Improve network connection handling (#2268)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: badrishc

libs/common/AsyncUtils.cs

@@ -22,6 +22,18 @@ public static class AsyncUtils
         public static void BlockingWait(Task t)
         => t.GetAwaiter().GetResult();
 
+        /// <summary>
+        /// Block on a given <see cref="Task"/> for at most <paramref name="millisecondsTimeout"/> milliseconds.
+        /// 
+        /// If at all possible, you should instead await your tasks, but for cases where this is not possible use this helper.
+        /// </summary>
+        /// <param name="t">The task to wait for.</param>
+        /// <param name="millisecondsTimeout">The maximum time to wait, or <see cref="System.Threading.Timeout.Infinite"/>.</param>
+        /// <returns>True if the task completed within the timeout; otherwise false, and the task keeps running.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static bool BlockingWait(Task t, int millisecondsTimeout)
+        => t.Wait(millisecondsTimeout);
+
         /// <summary>
         /// Block on a given <see cref="Task{TResult}"/>.
         /// 

libs/common/Networking/GarnetTcpNetworkSender.cs

@@ -7,7 +7,9 @@
 using System.Net.Sockets;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 using Garnet.networking;
+using Microsoft.Extensions.Logging;
 
 namespace Garnet.common
 {
@@ -46,6 +48,23 @@ public class GarnetTcpNetworkSender : NetworkSenderBase
         /// </summary>
         protected readonly int ThrottleMax = 8;
 
+        /// <summary>
+        /// Longest a send waits for an earlier send to complete, in milliseconds, before the connection is
+        /// closed as stalled; <see cref="Timeout.Infinite"/> waits indefinitely.
+        /// </summary>
+        readonly int sendTimeoutMs;
+
+        /// <summary>
+        /// Progressive spins a draining close performs before it blocks.
+        /// </summary>
+        const int DrainSpinIterations = 32;
+
+        /// <summary>
+        /// Progressive spins a send performs while waiting for a send slot before it blocks; this matches the spin
+        /// <see cref="SemaphoreSlim.Wait()"/> performs before it blocks.
+        /// </summary>
+        const int CapacitySpinIterations = 140;
+
         readonly string remoteEndpoint;
         readonly string localEndpoint;
 
@@ -64,12 +83,18 @@ public class GarnetTcpNetworkSender : NetworkSenderBase
         /// </summary>
         /// <param name="socket"></param>
         /// <param name="networkBufferSettings"></param>
+        /// <param name="networkPool"></param>
         /// <param name="throttleMax"></param>
+        /// <param name="sendTimeout">
+        /// Longest a send waits for an earlier send to complete before the connection is closed as stalled.
+        /// Zero or negative waits indefinitely.
+        /// </param>
         public GarnetTcpNetworkSender(
             Socket socket,
             NetworkBufferSettings networkBufferSettings,
             LimitedFixedBufferPool networkPool,
-            int throttleMax = 8)
+            int throttleMax = 8,
+            TimeSpan sendTimeout = default)
             : base(networkBufferSettings.sendBufferSize)
         {
             this.networkBufferSettings = networkBufferSettings;
@@ -78,6 +103,7 @@ public GarnetTcpNetworkSender(
             this.saeaStack = new(2 * ThrottleMax);
             this.responseObject = null;
             this.ThrottleMax = throttleMax;
+            this.sendTimeoutMs = sendTimeout > TimeSpan.Zero ? (int)Math.Min(sendTimeout.TotalMilliseconds, int.MaxValue) : Timeout.Infinite;
             this.spinLock = new();
             this.closeRequested = 0;
 
@@ -258,11 +284,19 @@ public override void SendCallback(object context)
         /// <inheritdoc />
         public override void DisposeNetworkSender(bool waitForSendCompletion)
         {
-            if (!waitForSendCompletion)
+            // Invariant: every in-flight send completes before the buffer stack and throttle are disposed, since a
+            // send's completion returns its buffer to the stack and may release the throttle. Draining also closes
+            // the sender, so any later send completes immediately without reaching the socket.
+            //
+            // A graceful close first lets in-flight sends reach the peer, giving up once no send has completed for the
+            // send timeout (never, when the timeout is infinite). An abortive close, or a graceful one that gives up,
+            // disposes the socket instead. That aborts the in-flight sends, so the unbounded drain after it is prompt
+            // and still upholds the invariant.
+            if (!waitForSendCompletion || !TryDrainSends(sendTimeoutMs))
+            {
                 socket.Dispose();
-
-            // Wait for ongoing sends to complete
-            while (throttleCount >= 0 && Interlocked.CompareExchange(ref throttleCount, int.MinValue, 0) != 0) Thread.Yield();
+                _ = TryDrainSends(Timeout.Infinite);
+            }
 
             // Empty and dispose the stack
             saeaStack.Dispose();
@@ -272,6 +306,56 @@ public override void DisposeNetworkSender(bool waitForSendCompletion)
                 socket.Dispose();
         }
 
+        /// <summary>
+        /// Waits for in-flight sends to complete, then closes the sender to further sends.
+        /// </summary>
+        /// <param name="timeoutMs">Longest to wait without any send completing, or <see cref="Timeout.Infinite"/>.</param>
+        /// <returns>False if no send completed within <paramref name="timeoutMs"/>; the sender is then still open.</returns>
+        bool TryDrainSends(int timeoutMs)
+        {
+            // Sends complete within microseconds while the peer reads, so spin progressively first
+            var spinner = ConfiguredSpinWait.BoundedProgressive(DrainSpinIterations);
+            do
+            {
+                if (TryCloseIfDrained())
+                    return true;
+            }
+            while (spinner.TryWait());
+
+            // A peer that reads slowly, or not at all, can take far longer. Blocking on a task, rather than sleeping,
+            // lets the thread pool inject a replacement for this thread while it waits.
+            return AsyncUtils.BlockingWait(DrainSendsAsync(timeoutMs));
+        }
+
+        async Task<bool> DrainSendsAsync(int timeoutMs)
+        {
+            using var poll = new PeriodicTimer(TimeSpan.FromMilliseconds(1));
+            var pending = Volatile.Read(ref throttleCount);
+            var progressAt = Environment.TickCount64;
+            while (!TryCloseIfDrained())
+            {
+                var count = Volatile.Read(ref throttleCount);
+                if (count != pending)
+                {
+                    pending = count;
+                    progressAt = Environment.TickCount64;
+                }
+                else if (timeoutMs != Timeout.Infinite && Environment.TickCount64 - progressAt >= timeoutMs)
+                {
+                    return false;
+                }
+
+                _ = await poll.WaitForNextTickAsync().ConfigureAwait(false);
+            }
+            return true;
+        }
+
+        /// <summary>
+        /// Closes the sender to further sends if no send is in flight, or reports that it is already closed.
+        /// </summary>
+        bool TryCloseIfDrained()
+            => Volatile.Read(ref throttleCount) < 0 || Interlocked.CompareExchange(ref throttleCount, int.MinValue, 0) == 0;
+
         /// <inheritdoc />
         public override void Throttle()
         {
@@ -329,7 +413,7 @@ private unsafe void Send(Socket socket, GarnetSaeaBuffer sendObject, int offset,
                 return;
             }
             if (cnt > ThrottleMax)
-                throttle.Wait();
+                WaitForSendCapacity();
 
             // Debug.WriteLine($"SEND: [{size}][{System.Text.Encoding.UTF8.GetString(new Span<byte>(sendObject.socketEventAsyncArgs.Buffer, offset, size)).Replace("\n", "|").Replace("\r", "")}]");
 
@@ -341,12 +425,62 @@ private unsafe void Send(Socket socket, GarnetSaeaBuffer sendObject, int offset,
                 SeaaBuffer_Completed(null, sendObject.socketEventAsyncArgs);
         }
 
+        /// <summary>
+        /// Waits for an in-flight send to complete, which frees a send slot for the caller.
+        /// </summary>
+        /// <remarks>
+        /// The caller is usually a thread-pool thread that also serves other connections. Blocking on a task lets
+        /// the thread pool inject a replacement thread, whereas blocking on the semaphore does not, so other
+        /// connections continue to be served while this one waits. A peer that accepts no data for
+        /// the send timeout has its connection closed, which also completes its in-flight sends.
+        /// <para>
+        /// A completion that frees a slot for a waiter decrements the send count before it releases the semaphore,
+        /// so the count can reach zero, letting the sender be disposed, while that release is still to come. A waiter
+        /// therefore leaves only once it has taken its release, even after a timeout; otherwise the release could
+        /// land on a disposed semaphore.
+        /// </para>
+        /// <para>
+        /// The timeout is applied to the wait on the untimed semaphore task rather than passed to the semaphore,
+        /// because a timed semaphore wait allocates a timer and async state on every call. The untimed task also
+        /// stays queued after a timeout, so it is the one that takes the owed release.
+        /// </para>
+        /// </remarks>
+        /// <exception cref="GarnetException">No send completed within the send timeout.</exception>
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void WaitForSendCapacity()
+        {
+            // A send usually completes shortly while the peer reads, so spin progressively, as the semaphore's own
+            // blocking wait does, before paying for a task and a blocking wait.
+            var spinner = ConfiguredSpinWait.BoundedProgressive(CapacitySpinIterations, sleep1Threshold: -1);
+            do
+            {
+                if (throttle.Wait(0))
+                    return;
+            }
+            while (spinner.TryWait());
+
+            var slot = throttle.WaitAsync();
+            if (AsyncUtils.BlockingWait(slot, sendTimeoutMs))
+                return;
+
+            // Closing the socket completes every in-flight send, so the release this waiter is owed arrives promptly
+            _ = TryClose();
+            AsyncUtils.BlockingWait(slot);
+            throw new GarnetException($"Closed connection to {remoteEndpoint}: no send completed within {sendTimeoutMs} ms",
+                LogLevel.Warning, clientResponse: false, disposeSession: true);
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void SeaaBuffer_Completed(object sender, SocketAsyncEventArgs e)
         {
             ReturnBuffer((GarnetSaeaBuffer)e.UserToken);
             if (Interlocked.Decrement(ref throttleCount) >= ThrottleMax)
+            {
+#if DEBUG
+                ExceptionInjectionHelper.ResetAndWait(ExceptionInjectionType.Sender_Pause_Before_Slot_Release);
+#endif
                 throttle.Release();
+            }
         }
     }
 }

libs/common/Networking/NetworkHandler.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Diagnostics;
+using System.IO;
 using System.Net.Security;
 using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
@@ -161,13 +163,26 @@ protected int BaseReceiveBufferSize
         // Number of times Dispose has been called
         int disposeCount;
 
+        // Set once a TLS write fails; a failed write can leave a partial record on the connection, so nothing
+        // further can be written to it
+        bool writeFailed;
+
+        // Longest a server-side TLS handshake may take; zero waits indefinitely
+        readonly TimeSpan tlsHandshakeTimeout;
+
+        // Longest delay a cancellation timer supports, about 49.7 days
+        static readonly TimeSpan MaxTlsHandshakeTimeout = TimeSpan.FromMilliseconds(uint.MaxValue - 1);
+
         /// <summary>
         /// Constructor
         /// </summary>
-        public unsafe NetworkHandler(TServerHook serverHook, TNetworkSender networkSender, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, ILogger logger = null)
+        /// <param name="tlsHandshakeTimeout">Longest a server-side TLS handshake may take before the connection is closed. Zero or negative waits indefinitely; a timeout beyond the longest a timer supports, about 49.7 days, is limited to it.</param>
+        public unsafe NetworkHandler(TServerHook serverHook, TNetworkSender networkSender, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, ILogger logger = null, TimeSpan tlsHandshakeTimeout = default)
             : base(networkPool.MinAllocationSize)
         {
             this.logger = logger;
+            this.tlsHandshakeTimeout = tlsHandshakeTimeout <= TimeSpan.Zero ? TimeSpan.Zero
+                : tlsHandshakeTimeout < MaxTlsHandshakeTimeout ? tlsHandshakeTimeout : MaxTlsHandshakeTimeout;
             this.serverHook = serverHook;
             this.networkSender = networkSender;
             this.session = messageConsumer;
@@ -246,9 +261,13 @@ public virtual async Task StartAsync(SslServerAuthenticationOptions tlsOptions =
         async Task AuthenticateAsServerAsync(SslServerAuthenticationOptions tlsOptions, string remoteEndpointName, CancellationToken token = default)
         {
             Debug.Assert(readerStatus == TlsReaderStatus.Active);
+
+            // The deadline is absolute; handshake bytes that arrive before it do not extend it
+            using var handshakeDeadline = tlsHandshakeTimeout > TimeSpan.Zero ? CancellationTokenSource.CreateLinkedTokenSource(token) : null;
+            handshakeDeadline?.CancelAfter(tlsHandshakeTimeout);
             try
             {
-                await sslStream.AuthenticateAsServerAsync(tlsOptions, token).ConfigureAwait(false);
+                await sslStream.AuthenticateAsServerAsync(tlsOptions, handshakeDeadline?.Token ?? token).ConfigureAwait(false);
 
                 if (token.IsCancellationRequested) throw new TaskCanceledException("AuthenticateAsServerAsync was cancelled");
 
@@ -263,10 +282,16 @@ async Task AuthenticateAsServerAsync(SslServerAuthenticationOptions tlsOptions,
             }
             catch (Exception ex)
             {
-                logger?.LogWarning(ex, "An error has occurred");
+                var timedOut = handshakeDeadline?.IsCancellationRequested == true && !token.IsCancellationRequested;
+                if (timedOut)
+                    logger?.LogInformation("TLS handshake with {remoteEndpoint} did not complete within {timeout}", remoteEndpointName, tlsHandshakeTimeout);
+                else
+                    logger?.LogWarning(ex, "An error has occurred");
                 readerStatus = TlsReaderStatus.Rest;
                 if (expectingData.CurrentCount == 0) expectingData.Release();
                 Dispose();
+                if (timedOut)
+                    throw new TimeoutException($"TLS handshake did not complete within {tlsHandshakeTimeout}", ex);
                 throw;
             }
         }
@@ -926,8 +951,24 @@ public override unsafe bool SendResponse(int offset, int size)
             logger?.LogTrace("SEND: [{send}]", System.Text.Encoding.UTF8.GetString(
                 new Span<byte>(transportSendBuffer).Slice(offset, size)).Replace("\n", "|").Replace("\r", ""));
 #endif
-            sslStream.Write(transportSendBuffer, offset, size);
-            sslStream.Flush();
+            if (writeFailed)
+                return false;
+
+            try
+            {
+                sslStream.Write(transportSendBuffer, offset, size);
+                sslStream.Flush();
+            }
+            catch (Exception ex)
+            {
+                writeFailed = true;
+
+                // SslStream wraps what the underlying sender throws; surface the sender's own failure, such as a send
+                // timeout, as the session would see it without TLS
+                if (ex is IOException { InnerException: GarnetException inner })
+                    ExceptionDispatchInfo.Throw(inner);
+                throw;
+            }
             return true;
         }
 

libs/common/Networking/TcpNetworkHandler.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+using System;
 using System.Net.Sockets;
 using Garnet.networking;
 using Microsoft.Extensions.Logging;
@@ -17,8 +18,10 @@ public abstract class TcpNetworkHandler<TServerHook> : TcpNetworkHandlerBase<TSe
         /// <summary>
         /// Constructor
         /// </summary>
-        public TcpNetworkHandler(TServerHook serverHook, Socket socket, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, int networkSendThrottleMax = 8, ILogger logger = null)
-            : base(serverHook, new GarnetTcpNetworkSender(socket, networkBufferSettings, networkPool, networkSendThrottleMax), socket, networkBufferSettings, networkPool, useTLS, messageConsumer: messageConsumer, logger: logger)
+        /// <param name="networkSendTimeout">Longest a send waits for an earlier send to complete before the connection is closed. Zero or negative waits indefinitely.</param>
+        /// <param name="tlsHandshakeTimeout">Longest a server-side TLS handshake may take before the connection is closed. Zero or negative waits indefinitely.</param>
+        public TcpNetworkHandler(TServerHook serverHook, Socket socket, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, int networkSendThrottleMax = 8, ILogger logger = null, TimeSpan networkSendTimeout = default, TimeSpan tlsHandshakeTimeout = default)
+            : base(serverHook, new GarnetTcpNetworkSender(socket, networkBufferSettings, networkPool, networkSendThrottleMax, networkSendTimeout), socket, networkBufferSettings, networkPool, useTLS, messageConsumer: messageConsumer, logger: logger, tlsHandshakeTimeout: tlsHandshakeTimeout)
         {
         }
     }

libs/common/Networking/TcpNetworkHandlerBase.cs

@@ -31,8 +31,9 @@ public abstract class TcpNetworkHandlerBase<TServerHook, TNetworkSender> : Netwo
         /// <summary>
         /// Constructor
         /// </summary>
-        public TcpNetworkHandlerBase(TServerHook serverHook, TNetworkSender networkSender, Socket socket, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, ILogger logger = null)
-            : base(serverHook, networkSender, networkBufferSettings, networkPool, useTLS, messageConsumer: messageConsumer, logger: logger)
+        /// <param name="tlsHandshakeTimeout">Longest a server-side TLS handshake may take before the connection is closed. Zero or negative waits indefinitely.</param>
+        public TcpNetworkHandlerBase(TServerHook serverHook, TNetworkSender networkSender, Socket socket, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, ILogger logger = null, TimeSpan tlsHandshakeTimeout = default)
+            : base(serverHook, networkSender, networkBufferSettings, networkPool, useTLS, messageConsumer: messageConsumer, logger: logger, tlsHandshakeTimeout: tlsHandshakeTimeout)
         {
             this.logger = logger;
             this.socket = socket;

libs/common/Testing/ExceptionInjectionType.cs

@@ -160,5 +160,10 @@ public enum ExceptionInjectionType
         /// Pause after receiving the full diskless snapshot, before recovery resets the AOF addresses.
         /// </summary>
         Replication_Pause_Before_Diskless_Snapshot_Recovery,
+        /// <summary>
+        /// Server network sender: pause a send completion that frees a slot for a waiting send, after it has
+        /// decremented the send count and before it releases the waiter.
+        /// </summary>
+        Sender_Pause_Before_Slot_Release,
     }
 }