Gracefully close sockets (nats-io#406)

* gracefully close sockets Signed-off-by: Caleb Lloyd <[email protected]> * MockServer handle socket closed Signed-off-by: Caleb Lloyd <[email protected]> --------- Signed-off-by: Caleb Lloyd <[email protected]>
rickdotnet · Feb 24, 2024 · 273a009 · 273a009
1 parent b16a182
commit 273a009
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 13 deletions.
diff --git a/src/NATS.Client.Core/Internal/SslStreamConnection.cs b/src/NATS.Client.Core/Internal/SslStreamConnection.cs
@@ -27,11 +27,7 @@ public SslStreamConnection(ILogger logger, SslStream sslStream, NatsTlsOpts tlsO
 
     public Task<Exception> WaitForClosed => _waitForClosedSource.Task;
 
-#if NET6_0
-    public ValueTask DisposeAsync()
-#else
     public async ValueTask DisposeAsync()
-#endif
     {
         if (Interlocked.Increment(ref _disposed) == 1)
         {
@@ -47,11 +43,9 @@ public async ValueTask DisposeAsync()
             catch
             {
             }
-        }
 
-#if NET6_0
-        return ValueTask.CompletedTask;
-#endif
+            await _sslStream.DisposeAsync().ConfigureAwait(false);
+        }
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]

diff --git a/src/NATS.Client.Core/Internal/TcpConnection.cs b/src/NATS.Client.Core/Internal/TcpConnection.cs
@@ -103,10 +103,18 @@ public ValueTask DisposeAsync()
             {
             }
 
+            try
+            {
+                _socket.Shutdown(SocketShutdown.Both);
+            }
+            catch
+            {
+            }
+
             _socket.Dispose();
         }
 
-        return default;
+        return ValueTask.CompletedTask;
     }
 
     // when catch SocketClosedException, call this method.

diff --git a/src/NATS.Client.Core/Internal/WebSocketConnection.cs b/src/NATS.Client.Core/Internal/WebSocketConnection.cs
@@ -8,6 +8,7 @@ internal sealed class WebSocketConnection : ISocketConnection
 {
     private readonly ClientWebSocket _socket;
     private readonly TaskCompletionSource<Exception> _waitForClosedSource = new();
+    private readonly TimeSpan _socketCloseTimeout = TimeSpan.FromSeconds(5); // matches _socketComponentDisposeTimeout in NatsConnection.cs
     private int _disposed;
 
     public WebSocketConnection()
@@ -78,7 +79,7 @@ public ValueTask AbortConnectionAsync(CancellationToken cancellationToken)
         return default;
     }
 
-    public ValueTask DisposeAsync()
+    public async ValueTask DisposeAsync()
     {
         if (Interlocked.Increment(ref _disposed) == 1)
         {
@@ -90,10 +91,17 @@ public ValueTask DisposeAsync()
             {
             }
 
+            try
+            {
+                var cts = new CancellationTokenSource(_socketCloseTimeout);
+                await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, default, cts.Token).ConfigureAwait(false);
+            }
+            catch
+            {
+            }
+
             _socket.Dispose();
         }
-
-        return default;
     }
 
     // when catch SocketClosedException, call this method.

diff --git a/tests/NATS.Client.TestUtilities/MockServer.cs b/tests/NATS.Client.TestUtilities/MockServer.cs
@@ -48,7 +48,12 @@ public MockServer(
                     {
                         while (!cancellationToken.IsCancellationRequested)
                         {
-                            var line = (await sr.ReadLineAsync())!;
+                            var line = await sr.ReadLineAsync();
+                            if (line == null)
+                            {
+                                // empty read, socket closed
+                                return;
+                            }
 
                             if (line.StartsWith("CONNECT"))
                             {