Add tiered send/receive recovery to azure-core-amqp and azure-messaging-servicebus

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/AmqpChannelProcessor.java

@@ -140,15 +140,19 @@ public void onNext(T amqpChannel) {
                     logger.info("Channel is now active.");
                 }
             }, error -> {
-                setAndClearChannel();
-                onError(error);
+                if (setAndClearChannel(amqpChannel)) {
+                    onError(error);
+                } else {
+                    logger.verbose("Ignoring error from a replaced channel.", error);
+                }
             }, () -> {
                 if (isDisposed()) {
                     logger.info("Channel is disposed.");
                 } else {
-                    logger.info("Channel is closed. Requesting upstream.");
-                    setAndClearChannel();
-                    requestUpstream();
+                    if (setAndClearChannel(amqpChannel)) {
+                        logger.info("Channel is closed. Requesting upstream.");
+                        requestUpstream();
+                    }
                 }
             });
         }
@@ -347,6 +351,36 @@ private void setAndClearChannel() {
         close(oldChannel);
     }
 
+    /**
+     * Force-closes the current cached channel so that the next subscriber receives a fresh one.
+     * This is used for connection-level recovery when the current connection is stale
+     * but the processor has not detected it (e.g., heartbeats echoed by intermediate infrastructure).
+     */
+    public void forceCloseChannel() {
+        setAndClearChannel();
+    }
+
+    /**
+     * Closes the failed channel only if it is still the current channel.
+     *
+     * @param failedChannel The channel used by the failed operation.
+     */
+    public void forceCloseChannel(T failedChannel) {
+        Objects.requireNonNull(failedChannel, "'failedChannel' cannot be null.");
+        setAndClearChannel(failedChannel);
+    }
+
+    private boolean setAndClearChannel(T failedChannel) {
+        synchronized (lock) {
+            if (currentChannel != failedChannel) {
+                return false;
+            }
+            currentChannel = null;
+        }
+        close(failedChannel);
+        return true;
+    }
+
     /**
      * Checks the current state of the channel for this channel and returns true if the channel is null or if this
      * processor is disposed.

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorConnection.java

@@ -362,7 +362,7 @@ public Mono<AmqpSession> createSession(String sessionName) {
                     logger.atInfo()
                         .addKeyValue(SESSION_NAME_KEY, sessionName)
                         .log("Error occurred. Removing and disposing session", error);
-                    removeSession(key);
+                    removeSession(key, amqpSession);
                 }, () -> {
                     // If we were already disposing of the connection, the session would be removed.
                     if (isDisposed.get()) {
@@ -372,7 +372,7 @@ public Mono<AmqpSession> createSession(String sessionName) {
                     logger.atVerbose()
                         .addKeyValue(SESSION_NAME_KEY, sessionName)
                         .log("Complete. Removing and disposing session.");
-                    removeSession(key);
+                    removeSession(key, amqpSession);
                 });
 
                 return new SessionSubscription(amqpSession, subscription);
@@ -396,8 +396,7 @@ public Mono<AmqpSession> createSession(String sessionName) {
 
                     final AmqpException amqpException = (AmqpException) error;
                     if (amqpException.getErrorCondition() == TIMEOUT_ERROR) {
-                        final SessionSubscription removed = sessionMap.remove(sessionName);
-                        removed.dispose();
+                        removeSession(sessionName, sessionSubscription.getSession());
                     }
                 });
 
@@ -443,6 +442,28 @@ public boolean removeSession(String sessionName) {
         return removed != null;
     }
 
+    /**
+     * Removes a session only if it is still the cached instance.
+     *
+     * @param sessionName The name used to cache the session.
+     * @param expectedSession The session owned by the failed operation.
+     * @return Whether the expected session was removed.
+     */
+    protected boolean removeSession(String sessionName, AmqpSession expectedSession) {
+        if (sessionName == null || expectedSession == null) {
+            return false;
+        }
+        if (useSessionChannelCache) {
+            return sessionCache.evict(sessionName, expectedSession);
+        }
+        final SessionSubscription entry = sessionMap.get(sessionName);
+        if (entry == null || entry.getSession() != expectedSession || !sessionMap.remove(sessionName, entry)) {
+            return false;
+        }
+        entry.dispose();
+        return true;
+    }
+
     @Override
     public boolean isDisposed() {
         return isDisposed.get();

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorConnectionCache.java

@@ -48,6 +48,10 @@ public final class ReactorConnectionCache<T extends ReactorConnection> implement
     // any dependent type; instead, the dependent type must acquire Connection only through the cache route,
     // i.e., by subscribing to 'createOrGetCachedConnection' via 'get()' getter.
     private volatile T currentConnection;
+    // Holds the ID of the connection that invalidateConnection() asked to force-invalidate.
+    // Only the connection whose getId() matches this value will be invalidated by cacheInvalidateIf;
+    // a freshly created connection with a different ID is never accidentally invalidated.
+    private final AtomicReference<String> forceInvalidateConnectionId = new AtomicReference<>(null);
     private final State state = new State();
 
     /**
@@ -113,12 +117,25 @@ public ReactorConnectionCache(Supplier<T> connectionSupplier, String fullyQualif
             }
         }).cacheInvalidateIf(c -> {
             if (c.isDisposed()) {
+                // Connection disposed for any reason. Clean up the force-invalidate marker if it
+                // was targeting this connection so it is not accidentally consumed by a future
+                // connection that happens to have the same ID.
+                forceInvalidateConnectionId.compareAndSet(c.getId(), null);
                 withConnectionId(logger, c.getId()).log("The connection is closed, requesting a new connection.");
                 return true;
-            } else {
-                // Emit cached connection.
-                return false;
             }
+            final String targetId = forceInvalidateConnectionId.get();
+            if (targetId != null
+                && targetId.equals(c.getId())
+                && forceInvalidateConnectionId.compareAndSet(targetId, null)) {
+                // invalidateConnection() marked this connection dirty.
+                // Close it here — the cache owns the connection lifecycle.
+                withConnectionId(logger, c.getId()).log("Force-invalidating and closing connection for recovery.");
+                closeConnection(c, logger, "Force-invalidated for recovery.");
+                return true;
+            }
+            // No forced invalidation targeted this connection — emit it from cache.
+            return false;
         });
     }
 
@@ -172,6 +189,49 @@ public boolean isCurrentConnectionClosed() {
         return (currentConnection != null && currentConnection.isDisposed()) || terminated;
     }
 
+    /**
+     * Marks the current cached connection for invalidation so that the next {@link #get()} call
+     * closes it and creates a fresh connection. This is used for connection-level recovery when
+     * the current connection is in a stale state that the cache's normal error detection (via
+     * endpoint state signals) has not detected — for example, when intermediate infrastructure
+     * (load balancers, NAT gateways) is echoing AMQP heartbeats on behalf of a dead connection.
+     *
+     * <p>The actual close is deferred to the cache chain's {@code cacheInvalidateIf} predicate
+     * inside {@link #get()}, keeping the connection lifecycle local to the cache. This avoids
+     * the concurrency concern of two threads touching the cached connection simultaneously.</p>
+     *
+     * <p>This is modeled after the Go SDK's {@code Namespace.Recover()} which explicitly closes
+     * the old connection and increments the connection revision.</p>
+     *
+     * <p>This method is safe to call concurrently. If the connection is already closed or being
+     * closed, this is a no-op.</p>
+     */
+    public void invalidateConnection() {
+        final T connection = currentConnection;
+        if (connection != null && !connection.isDisposed()) {
+            withConnectionId(logger, connection.getId())
+                .log("Marking connection for force-invalidation. Next get() will close and replace it.");
+            forceInvalidateConnectionId.set(connection.getId());
+        }
+    }
+
+    /**
+     * Marks the failed connection for invalidation only if it is still cached.
+     * A delayed failure must not invalidate a healthy replacement connection.
+     *
+     * @param connectionId The identifier of the connection used by the failed operation.
+     */
+    public void invalidateConnection(String connectionId) {
+        Objects.requireNonNull(connectionId, "'connectionId' cannot be null.");
+        synchronized (lock) {
+            final T connection = currentConnection;
+            if (connection != null && connectionId.equals(connection.getId()) && !connection.isDisposed()) {
+                withConnectionId(logger, connectionId).log("Marking failed connection for force-invalidation.");
+                forceInvalidateConnectionId.set(connectionId);
+            }
+        }
+    }
+
     /**
      * Terminate so that consumers will no longer be able to request connection. If there is a current (cached)
      * connection then it will be closed.

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorSessionCache.java

@@ -3,6 +3,7 @@
 
 package com.azure.core.amqp.implementation;
 
+import com.azure.core.amqp.AmqpSession;
 import com.azure.core.util.logging.ClientLogger;
 import org.apache.qpid.proton.engine.Connection;
 import reactor.core.Disposable;
@@ -119,6 +120,18 @@ boolean evict(String name) {
         return removed != null;
     }
 
+    boolean evict(String name, AmqpSession expectedSession) {
+        if (name == null || expectedSession == null) {
+            return false;
+        }
+        final Entry entry = entries.get(name);
+        if (entry == null || entry.getSession() != expectedSession || !entries.remove(name, entry)) {
+            return false;
+        }
+        entry.dispose();
+        return true;
+    }
+
     /**
      * Signal that the owner ({@link ReactorConnection}) of the cache is disposed of.
      */
@@ -193,7 +206,7 @@ private void evict(ReactorSession session, String message, Throwable error) {
         } else {
             logger.atInfo().addKeyValue(SESSION_NAME_KEY, name).addKeyValue(SESSION_ID_KEY, id).log(message);
         }
-        evict(name);
+        evict(name, session);
     }
 
     /**