Convert sync over async network handling to async

Version.props

@@ -1,6 +1,6 @@
 <Project>
 	<!-- Versioning property for builds and packages -->
 	<PropertyGroup>
-		<VersionPrefix>1.0.48</VersionPrefix>
+		<VersionPrefix>1.0.49</VersionPrefix>
 	</PropertyGroup>
 </Project>

benchmark/BDN.benchmark/Embedded/EmbeddedNetworkHandler.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Diagnostics;
+using System.Threading.Tasks;
 using Garnet.common;
 using Garnet.networking;
 using Microsoft.Extensions.Logging;
@@ -11,8 +12,11 @@ namespace Embedded.server
 {
     internal class EmbeddedNetworkHandler : NetworkHandler<GarnetServerEmbedded, EmbeddedNetworkSender>
     {
+        readonly bool useTLS;
+
         public EmbeddedNetworkHandler(GarnetServerEmbedded serverHook, EmbeddedNetworkSender networkSender, NetworkBufferSettings networkBufferSettings, LimitedFixedBufferPool networkPool, bool useTLS, IMessageConsumer messageConsumer = null, ILogger logger = null) : base(serverHook, networkSender, networkBufferSettings, networkPool, useTLS, messageConsumer, logger)
         {
+            this.useTLS = useTLS;
         }
 
         public override string RemoteEndpointName => throw new NotImplementedException();
@@ -24,14 +28,16 @@ public override void Dispose()
 
         public override bool TryClose() => throw new NotImplementedException();
 
-        public unsafe void Send(Request request)
+        public async ValueTask Send(Request request)
         {
             networkReceiveBuffer = request.buffer;
-            networkReceiveBufferPtr = request.bufferPtr;
+            unsafe { networkReceiveBufferPtr = request.bufferPtr; }
 
-            OnNetworkReceive(request.buffer.Length);
+            if (useTLS)
+                await OnNetworkReceiveWithTLSAsync(request.buffer.Length);
+            else
+                OnNetworkReceiveWithoutTLS(request.buffer.Length);
 
-            // We should have consumed the entire buffer
             Debug.Assert(networkBytesRead == 0);
             Debug.Assert(networkReadHead == 0);
         }

benchmark/BDN.benchmark/Network/BasicOperations.cs

@@ -10,7 +10,7 @@ namespace BDN.benchmark.Network
     /// Benchmark for BasicOperations
     /// </summary>
     [MemoryDiagnoser]
-    public unsafe class BasicOperations : NetworkBase
+    public class BasicOperations : NetworkBase
     {
         static ReadOnlySpan<byte> INLINE_PING => "PING\r\n"u8;
         Request ping;
@@ -22,9 +22,9 @@ public override void GlobalSetup()
         }
 
         [Benchmark]
-        public void InlinePing()
+        public async ValueTask InlinePing()
         {
-            Send(ping);
+            await Send(ping);
         }
     }
 }

benchmark/BDN.benchmark/Network/NetworkBase.cs

@@ -64,10 +64,7 @@ public virtual void GlobalCleanup()
             server.Dispose();
         }
 
-        protected void Send(Request request)
-        {
-            networkHandler.Send(request);
-        }
+        protected ValueTask Send(Request request) => networkHandler.Send(request);
 
         protected unsafe void SetupOperation(ref Request request, ReadOnlySpan<byte> operation, int batchSize = batchSize)
         {

benchmark/BDN.benchmark/Network/RawStringOperations.cs

@@ -10,7 +10,7 @@ namespace BDN.benchmark.Network
     /// Benchmark for RawStringOperations
     /// </summary>
     [MemoryDiagnoser]
-    public unsafe class RawStringOperations : NetworkBase
+    public class RawStringOperations : NetworkBase
     {
         static ReadOnlySpan<byte> SET => "*3\r\n$3\r\nSET\r\n$1\r\na\r\n$1\r\na\r\n"u8;
         Request set;
@@ -65,63 +65,63 @@ public override void GlobalSetup()
         }
 
         [Benchmark]
-        public void Set()
+        public async ValueTask Set()
         {
-            Send(set);
+            await Send(set);
         }
 
         [Benchmark]
-        public void SetEx()
+        public async ValueTask SetEx()
         {
-            Send(setex);
+            await Send(setex);
         }
 
         [Benchmark]
-        public void SetNx()
+        public async ValueTask SetNx()
         {
-            Send(setnx);
+            await Send(setnx);
         }
 
         [Benchmark]
-        public void SetXx()
+        public async ValueTask SetXx()
         {
-            Send(setxx);
+            await Send(setxx);
         }
 
         [Benchmark]
-        public void GetFound()
+        public async ValueTask GetFound()
         {
-            Send(getf);
+            await Send(getf);
         }
 
         [Benchmark]
-        public void GetNotFound()
+        public async ValueTask GetNotFound()
         {
-            Send(getnf);
+            await Send(getnf);
         }
 
         [Benchmark]
-        public void Increment()
+        public async ValueTask Increment()
         {
-            Send(incr);
+            await Send(incr);
         }
 
         [Benchmark]
-        public void Decrement()
+        public async ValueTask Decrement()
         {
-            Send(decr);
+            await Send(decr);
         }
 
         [Benchmark]
-        public void IncrementBy()
+        public async ValueTask IncrementBy()
         {
-            Send(incrby);
+            await Send(incrby);
         }
 
         [Benchmark]
-        public void DecrementBy()
+        public async ValueTask DecrementBy()
         {
-            Send(decrby);
+            await Send(decrby);
         }
     }
 }

libs/common/Networking/NetworkHandler.cs

@@ -267,39 +267,40 @@ async Task AuthenticateAsClientAsync(SslClientAuthenticationOptions sslClientOpt
             }
         }
 
+        public unsafe void OnNetworkReceiveWithoutTLS(int bytesTransferred)
+        {
+            networkBytesRead += bytesTransferred;
+            transportReceiveBuffer = networkReceiveBuffer;
+            transportReceiveBufferPtr = networkReceiveBufferPtr;
+            transportBytesRead = networkBytesRead;
+
+            // Process non-TLS code on the synchronous thread
+            Process();
+
+            EndTransformNetworkToTransport();
+            UpdateNetworkBuffers();
+        }
+
         /// <summary>
         /// On network receive
         /// </summary>
         /// <param name="bytesTransferred">Number of bytes transferred</param>
-        public unsafe void OnNetworkReceive(int bytesTransferred)
+        public async ValueTask OnNetworkReceiveWithTLSAsync(int bytesTransferred)
         {
             // Wait for SslStream async processing to complete, if any (e.g., authentication phase)
             while (readerStatus == TlsReaderStatus.Active)
-                expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false).GetAwaiter().GetResult();
+                await expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false);
 
             // Increment network bytes read
             networkBytesRead += bytesTransferred;
 
             switch (readerStatus)
             {
                 case TlsReaderStatus.Rest:
-                    // Synchronously try to process the received data
-                    if (sslStream == null)
-                    {
-                        transportReceiveBuffer = networkReceiveBuffer;
-                        transportReceiveBufferPtr = networkReceiveBufferPtr;
-                        transportBytesRead = networkBytesRead;
-
-                        // We do not have an active read task, so we will process on the network thread
-                        Process();
-                    }
-                    else
-                    {
-                        readerStatus = TlsReaderStatus.Active;
-                        Read();
-                        while (readerStatus == TlsReaderStatus.Active)
-                            expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false).GetAwaiter().GetResult();
-                    }
+                    readerStatus = TlsReaderStatus.Active;
+                    Read();
+                    while (readerStatus == TlsReaderStatus.Active)
+                        await expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false);
                     break;
                 case TlsReaderStatus.Waiting:
                     // We have a ReadAsync task waiting for new data, set it to active status
@@ -309,17 +310,19 @@ public unsafe void OnNetworkReceive(int bytesTransferred)
                     _ = receivedData.Release();
 
                     while (readerStatus == TlsReaderStatus.Active)
-                        expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false).GetAwaiter().GetResult();
+                        await expectingData.WaitAsync(cancellationTokenSource.Token).ConfigureAwait(false);
                     break;
                 default:
                     ThrowInvalidOperationException($"Unexpected reader status {readerStatus}");
                     break;
             }
 
             Debug.Assert(readerStatus != TlsReaderStatus.Active);
+            UpdateNetworkBuffers();
+        }
 
-            EndTransformNetworkToTransport();
-
+        void UpdateNetworkBuffers()
+        {
             // Shift network buffer after processing is done
             if (networkReadHead > 0)
                 ShiftNetworkReceiveBuffer();

libs/common/Networking/TcpNetworkHandlerBase.cs

@@ -39,7 +39,6 @@ public TcpNetworkHandlerBase(TServerHook serverHook, TNetworkSender networkSende
 
             remoteEndpoint = socket.RemoteEndPoint is IPEndPoint remote ? $"{remote.Address}:{remote.Port}" : "";
             localEndpoint = socket.LocalEndPoint is IPEndPoint local ? $"{local.Address}:{local.Port}" : "";
-
             AllocateNetworkReceiveBuffer();
         }
 
@@ -52,28 +51,28 @@ public TcpNetworkHandlerBase(TServerHook serverHook, TNetworkSender networkSende
         /// <inheritdoc />
         public override void Start(SslServerAuthenticationOptions tlsOptions = null, string remoteEndpointName = null, CancellationToken token = default)
         {
-            Start();
+            Start(tlsOptions != null);
             base.Start(tlsOptions, remoteEndpointName, token);
         }
 
         /// <inheritdoc />
         public override async Task StartAsync(SslServerAuthenticationOptions tlsOptions = null, string remoteEndpointName = null, CancellationToken token = default)
         {
-            Start();
+            Start(tlsOptions != null);
             await base.StartAsync(tlsOptions, remoteEndpointName, token).ConfigureAwait(false);
         }
 
         /// <inheritdoc />
         public override void Start(SslClientAuthenticationOptions tlsOptions, string remoteEndpointName = null, CancellationToken token = default)
         {
-            Start();
+            Start(tlsOptions != null);
             base.Start(tlsOptions, remoteEndpointName, token);
         }
 
         /// <inheritdoc />
         public override async Task StartAsync(SslClientAuthenticationOptions tlsOptions, string remoteEndpointName = null, CancellationToken token = default)
         {
-            Start();
+            Start(tlsOptions != null);
             await base.StartAsync(tlsOptions, remoteEndpointName, token).ConfigureAwait(false);
         }
 
@@ -102,17 +101,22 @@ public override bool TryClose()
             return true;
         }
 
-        void Start()
+        void Start(bool useTLS)
         {
             var receiveEventArgs = new SocketAsyncEventArgs { AcceptSocket = socket };
             receiveEventArgs.SetBuffer(networkReceiveBuffer, 0, networkReceiveBuffer.Length);
-            receiveEventArgs.Completed += RecvEventArg_Completed;
+            receiveEventArgs.Completed += useTLS ? RecvEventArgCompletedWithTLS : RecvEventArgCompletedWithoutTLS;
 
             // If the client already have packets, avoid handling it here on the handler so we don't block future accepts.
             try
             {
                 if (!socket.ReceiveAsync(receiveEventArgs))
-                    Task.Run(() => RecvEventArg_Completed(null, receiveEventArgs));
+                {
+                    if (useTLS)
+                        Task.Run(() => RecvEventArgCompletedWithTLS(null, receiveEventArgs));
+                    else
+                        Task.Run(() => RecvEventArgCompletedWithoutTLS(null, receiveEventArgs));
+                }
             }
             catch (Exception ex)
             {
@@ -134,7 +138,35 @@ void Dispose(SocketAsyncEventArgs e)
             e.Dispose();
         }
 
-        void RecvEventArg_Completed(object sender, SocketAsyncEventArgs e)
+        void RecvEventArgCompletedWithTLS(object sender, SocketAsyncEventArgs e) =>
+            _ = HandleReceiveWithTLSAsync(sender, e);
+
+        void RecvEventArgCompletedWithoutTLS(object sender, SocketAsyncEventArgs e) =>
+            HandleReceiveWithoutTLS(sender, e);
+
+        private void HandleReceiveWithoutTLS(object sender, SocketAsyncEventArgs e)
+        {
+            try
+            {
+                do
+                {
+                    if (e.BytesTransferred == 0 || e.SocketError != SocketError.Success || serverHook.Disposed)
+                    {
+                        // No more things to receive
+                        Dispose(e);
+                        break;
+                    }
+                    OnNetworkReceiveWithoutTLS(e.BytesTransferred);
+                    e.SetBuffer(networkReceiveBuffer, networkBytesRead, networkReceiveBuffer.Length - networkBytesRead);
+                } while (!e.AcceptSocket.ReceiveAsync(e));
+            }
+            catch (Exception ex)
+            {
+                HandleReceiveFailure(ex, e);
+            }
+        }
+
+        private async ValueTask HandleReceiveWithTLSAsync(object sender, SocketAsyncEventArgs e)
         {
             try
             {
@@ -146,20 +178,29 @@ void RecvEventArg_Completed(object sender, SocketAsyncEventArgs e)
                         Dispose(e);
                         break;
                     }
-                    OnNetworkReceive(e.BytesTransferred);
+                    var receiveTask = OnNetworkReceiveWithTLSAsync(e.BytesTransferred);
+                    if (!receiveTask.IsCompletedSuccessfully)
+                    {
+                        await receiveTask;
+                    }
                     e.SetBuffer(networkReceiveBuffer, networkBytesRead, networkReceiveBuffer.Length - networkBytesRead);
                 } while (!e.AcceptSocket.ReceiveAsync(e));
             }
             catch (Exception ex)
             {
-                if (ex is ObjectDisposedException ex2 && ex2.ObjectName == "System.Net.Sockets.Socket")
-                    logger?.LogTrace("Accept socket was disposed at RecvEventArg_Completed");
-                else
-                    logger?.LogError(ex, "An error occurred at RecvEventArg_Completed");
-                Dispose(e);
+                HandleReceiveFailure(ex, e);
             }
         }
 
+        void HandleReceiveFailure(Exception ex, SocketAsyncEventArgs e)
+        {
+            if (ex is ObjectDisposedException ex2 && ex2.ObjectName == "System.Net.Sockets.Socket")
+                logger?.LogTrace("Accept socket was disposed at RecvEventArg_Completed");
+            else
+                logger?.LogError(ex, "An error occurred at RecvEventArg_Completed");
+            Dispose(e);
+        }
+
         unsafe void AllocateNetworkReceiveBuffer()
         {
             networkReceiveBufferEntry = networkPool.Get(networkBufferSettings.initialReceiveBufferSize);