Fix root causes behind recurring Garnet .NET CI flakes

libs/client/ClientSession/GarnetClientSession.cs

@@ -42,6 +42,8 @@ public sealed partial class GarnetClientSession : IServerHook, IMessageConsumer
         static readonly Exception disposeException = new GarnetClientDisposedException();
 
         Socket socket;
+
+        int closeRequested;
         int disposed;
 
         // Send        
@@ -147,6 +149,7 @@ public GarnetClientSession(
         public unsafe void Connect(CancellationToken token = default)
         {
             socket = ConnectSendSocket();
+            ThrowIfCloseRequested();
             networkHandler = new GarnetClientSessionTcpNetworkHandler(
                 this,
                 socket,
@@ -199,6 +202,7 @@ public unsafe void Connect(CancellationToken token = default)
         public async Task ConnectAsync(int timeoutMs = 0, CancellationToken token = default)
         {
             socket = await ConnectSendSocketAsync(timeoutMs, token).ConfigureAwait(false);
+            ThrowIfCloseRequested();
             networkHandler = new GarnetClientSessionTcpNetworkHandler(
                 this,
                 socket,
@@ -407,6 +411,44 @@ public Task ReconnectAsync(int timeoutMs = 0, CancellationToken token = default)
             return ConnectAsync(timeoutMs, token);
         }
 
+        /// <summary>
+        /// Closes the underlying connection so a thread blocked on network I/O through this session
+        /// fails fast, without disposing the session itself. The session expects mono-threaded access,
+        /// so a caller that is not the thread using the session must not dispose it: doing so races
+        /// that thread's use of the send buffer it has rented. Closing the connection is safe because
+        /// it only invalidates the socket, leaving the buffer owned by the thread that rented it.
+        ///
+        /// The session is not reusable afterwards: the request is sticky, so a later
+        /// <see cref="ReconnectAsync"/> or <see cref="Connect"/> throws.
+        /// </summary>
+        public void CloseConnection()
+        {
+            _ = Interlocked.Exchange(ref closeRequested, 1);
+            Volatile.Read(ref socket)?.Dispose();
+        }
+
+        /// <summary>
+        /// A close that lands while the connection is still being established would otherwise be lost:
+        /// the socket is not yet published for <see cref="CloseConnection"/> to find, and the exchanges
+        /// that follow await a reply with no cancellation token, so the connecting thread would park
+        /// indefinitely on a peer that is shutting down. Checking here after the socket is published
+        /// means either the closer observes the socket or the connecting thread observes the request.
+        /// </summary>
+        void ThrowIfCloseRequested()
+        {
+            // Publishing the socket and reading the request are a store followed by a load of a
+            // different location, so without a StoreLoad fence both sides can miss: this thread's
+            // store is still buffered while it reads a stale request, and the closer reads a stale
+            // socket. The closer's Interlocked.Exchange fences its own side; this fences ours.
+            Interlocked.MemoryBarrier();
+
+            if (Volatile.Read(ref closeRequested) == 0)
+                return;
+
+            socket?.Dispose();
+            throw new ObjectDisposedException(nameof(GarnetClientSession), "Connection was closed while connecting");
+        }
+
         /// <summary>
         /// Dispose instance
         /// </summary>

libs/cluster/Server/Replication/PrimaryOps/AofOperations/AofSyncDriver.cs

@@ -121,13 +121,21 @@ public void Dispose()
             // Cancel cts
             cts?.Cancel();
 
-            // Dispose sync tasks
+            // Break the connections so a task blocked on a send fails fast. The client sessions are
+            // deliberately not disposed here: they expect mono-threaded access, and a task may still
+            // be writing through one. Disposing it from this thread races that write and can strand
+            // the send buffer the task rented, which leaves the replication buffer pool's dispose
+            // spinning forever and the whole server stuck in shutdown holding its port.
             foreach (var aofSyncTask in aofSyncTasks)
-                aofSyncTask?.Dispose();
+                aofSyncTask?.CloseConnection();
 
-            // Wait for tasks to exit
+            // Wait for tasks to exit; each disposes its own client before leaving the monitor
             activeWorkerMonitor.Dispose();
 
+            // Dispose sync tasks, now that no task is running
+            foreach (var aofSyncTask in aofSyncTasks)
+                aofSyncTask?.Dispose();
+
             // Finally, dispose the cts
             cts?.Dispose();
         }

libs/cluster/Server/Replication/PrimaryOps/AofOperations/AofSyncTask.cs

@@ -143,6 +143,19 @@ public AofSyncTask(
                 this.logger = logger;
             }
 
+            /// <summary>
+            /// Closes the network connection so a task blocked on a send fails fast, without
+            /// disposing the client session the task may still be writing through.
+            /// </summary>
+            public void CloseConnection()
+            {
+                try
+                {
+                    garnetClient?.CloseConnection();
+                }
+                catch { }
+            }
+
             public void Dispose()
             {
                 try
@@ -344,9 +357,12 @@ await iter.BulkConsumeAllAsync(
                 }
                 finally
                 {
+                    // The client is disposed before leaving the monitor so that a drained monitor
+                    // means every client has already been torn down by the one thread that was
+                    // using it, and the send buffer it rented is back in the replication pool.
+                    garnetClient?.Dispose();
                     if (enteredMonitor)
                         _ = aofSyncDriver.activeWorkerMonitor.Exit();
-                    garnetClient?.Dispose();
                 }
 
                 [Conditional("DEBUG")]

libs/common/Testing/ExceptionInjectionHelper.cs

@@ -17,6 +17,12 @@ public static class ExceptionInjectionHelper
         static object @lock = new();
         static TaskCompletionSource<bool> update = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
+        /// <summary>
+        /// Non-zero while at least one <see cref="SuspendParking"/> scope is open. While it is non-zero
+        /// <see cref="ResetAndWaitAsync"/> refuses to park, because nothing is left to signal it.
+        /// </summary>
+        static int parkingSuspensions;
+
         /// <summary>
         /// Array of exception injection types
         /// </summary>
@@ -128,6 +134,51 @@ public static bool TriggerCondition(ExceptionInjectionType exceptionType)
 #endif
         }
 
+        /// <summary>
+        /// Stops <see cref="ResetAndWaitAsync"/> from parking, and releases anyone already parked.
+        ///
+        /// A parked waiter is only released by <see cref="EnableException"/>, but the cleanup a test runs on
+        /// its way out is <see cref="DisableException"/>. A test that leaves between a waiter arriving and
+        /// being re-enabled - an assertion failing, or a wait for the arrival timing out - therefore strands
+        /// that waiter permanently. The waiter is a server thread holding a pooled network buffer, so
+        /// <c>LimitedFixedBufferPool.Dispose</c> then spins forever waiting for a reference that is never
+        /// returned and the whole test process hangs rather than one test failing.
+        ///
+        /// Suspension covers arrivals as well as waiters already parked, so a request that reaches an
+        /// injection point after its owner has begun shutting down cannot re-create the same hang.
+        ///
+        /// Every call must be paired with <see cref="ResumeParking"/>. The count keeps concurrent shutdowns
+        /// independent, and ending every scope is what stops a suspension leaking into a later test sharing
+        /// this process-wide state.
+        /// </summary>
+        [Conditional("DEBUG")]
+        public static void SuspendParking()
+        {
+            TaskCompletionSource<bool> release;
+
+            lock (@lock)
+            {
+                parkingSuspensions++;
+                release = update;
+                update = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            }
+
+            _ = release.TrySetResult(true);
+        }
+
+        /// <summary>
+        /// Ends a <see cref="SuspendParking"/> scope. Parking resumes once every scope has ended.
+        /// </summary>
+        [Conditional("DEBUG")]
+        public static void ResumeParking()
+        {
+            lock (@lock)
+            {
+                Debug.Assert(parkingSuspensions > 0, "ResumeParking without a matching SuspendParking");
+                parkingSuspensions--;
+            }
+        }
+
         /// <summary>
         /// Wait on set condition
         /// </summary>
@@ -149,7 +200,10 @@ public static async Task ResetAndWaitAsync(ExceptionInjectionType exceptionType)
                     Task task;
                     lock (@lock)
                     {
-                        if (IsEnabled(exceptionType))
+                        // Parking is suspended while a server is shutting down, because whoever armed this
+                        // injection point is gone and will never re-enable it. Reading it under the lock is
+                        // what makes a suspension raised at any point before here take effect.
+                        if (IsEnabled(exceptionType) || parkingSuspensions > 0)
                             break;
                         task = update.Task;
                     }

libs/host/GarnetServer.cs

@@ -564,23 +564,35 @@ public void Dispose(bool deleteDir = true)
 
         private void InternalDispose()
         {
-            // Phase 1: Stop listening on all servers to free ports immediately.
-            for (var i = 0; i < servers.Length; i++)
-                servers[i]?.Close();
+            // A thread parked in a test injection point holds a pooled network buffer, and the drain below
+            // waits for every such buffer to come back. Suspending covers both waiters already parked and
+            // requests that reach an injection point while this runs, so shutdown cannot be blocked by an
+            // injection point whose owner is already gone. Compiled out in Release.
+            ExceptionInjectionHelper.SuspendParking();
+            try
+            {
+                // Phase 1: Stop listening on all servers to free ports immediately.
+                for (var i = 0; i < servers.Length; i++)
+                    servers[i]?.Close();
 
-            // Phase 2: Drain active handlers and clean up remaining resources.
-            for (var i = 0; i < servers.Length; i++)
-                servers[i]?.Dispose();
+                // Phase 2: Drain active handlers and clean up remaining resources.
+                for (var i = 0; i < servers.Length; i++)
+                    servers[i]?.Dispose();
 
-            // Phase 3: Dispose the provider (storage engine shutdown — may take time).
-            Provider?.Dispose();
+                // Phase 3: Dispose the provider (storage engine shutdown — may take time).
+                Provider?.Dispose();
 
-            subscribeBroker?.Dispose();
-            storeEpoch?.Dispose();
-            pubSubEpoch?.Dispose();
-            opts.AuthSettings?.Dispose();
-            if (disposeLoggerFactory)
-                loggerFactory?.Dispose();
+                subscribeBroker?.Dispose();
+                storeEpoch?.Dispose();
+                pubSubEpoch?.Dispose();
+                opts.AuthSettings?.Dispose();
+                if (disposeLoggerFactory)
+                    loggerFactory?.Dispose();
+            }
+            finally
+            {
+                ExceptionInjectionHelper.ResumeParking();
+            }
         }
 
         private static void DeleteDirectory(string path)

libs/server/Resp/Vector/VectorManager.ContextMetadata.cs

@@ -448,6 +448,10 @@ public readonly void Dispose()
                 manager.recoveredIndexes.Clear();
                 manager.recoveredMetadata.Clear();
 
+                // Migration remappings name contexts in the array just replaced, so a surviving entry would
+                // steer the rest of that migration into a context this flush has already handed back
+                manager.ClearMigratedContextRemap();
+
                 // Allow Vector Set operations again
                 manager.vectorSetLocks.ReleaseLock(lockToken);