fix: accepted network clients are disposed twice causing the server t…

…o throw already disposed exceptions when disconnecting
Fresa · Aug 16, 2021 · fdc237b · fdc237b
1 parent edd8903
commit fdc237b
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 27 deletions.
diff --git a/src/Server/AsyncDisposables.cs b/src/Server/AsyncDisposables.cs
@@ -1,6 +1,6 @@
 using System;
-using System.Linq;
 using System.Threading.Tasks;
+using Port.Server.Framework;
 
 namespace Port.Server
 {
@@ -15,11 +15,8 @@ public AsyncDisposables(params IAsyncDisposable[] disposables)
 
         public async ValueTask DisposeAsync()
         {
-            await Task.WhenAll(
-                _disposables.Select(forwarder => forwarder.DisposeAsync())
-                            .Where(
-                                valueTask => !valueTask.IsCompletedSuccessfully)
-                            .Select(valueTask => valueTask.AsTask()));
+            await _disposables.DisposeAllAsync()
+                              .ConfigureAwait(false);
         }
     }
 }
diff --git a/src/Server/Framework/AsyncDisposableExtensions.cs b/src/Server/Framework/AsyncDisposableExtensions.cs
@@ -0,0 +1,15 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace Port.Server.Framework
+{
+    internal static class AsyncDisposableExtensions
+    {
+        internal static Task DisposeAllAsync(
+            this IEnumerable<IAsyncDisposable> disposables)
+            => disposables.Select(client => client.DisposeAsync())
+                          .WhenAllAsync();
+    }
+}
diff --git a/src/Server/Framework/ValueTaskExtensions.cs b/src/Server/Framework/ValueTaskExtensions.cs
@@ -0,0 +1,17 @@
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace Port.Server.Framework
+{
+    internal static class ValueTaskExtensions
+    {
+        internal static Task WhenAllAsync(
+            this IEnumerable<ValueTask> tasks)
+            => Task.WhenAll(
+                tasks.Where(
+                         valueTask
+                             => !valueTask.IsCompletedSuccessfully)
+                     .Select(valueTask => valueTask.AsTask()));
+    }
+}
diff --git a/src/Server/Services/PortForwardService.cs b/src/Server/Services/PortForwardService.cs
@@ -4,6 +4,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Grpc.Core;
+using Port.Server.Framework;
 using Port.Server.Messages;
 
 namespace Port.Server.Services
@@ -86,18 +87,14 @@ public async ValueTask DisposeAsync()
 
                 while (!_portForwardHandlers.IsEmpty)
                 {
-                    await Task.WhenAll(
-                        _portForwardHandlers
-                            .Keys
-                            .Select(
-                                key => _portForwardHandlers
-                                    .TryRemove(key, out var handler)
-                                    ? handler.DisposeAsync()
-                                    : ValueTask.CompletedTask)
-                            .Where(
-                                valueTask => !valueTask
-                                    .IsCompletedSuccessfully)
-                            .Select(valueTask => valueTask.AsTask()));
+                    await _portForwardHandlers
+                          .Keys
+                          .Select(
+                              key => _portForwardHandlers
+                                  .TryRemove(key, out var handler)
+                                  ? handler.DisposeAsync()
+                                  : ValueTask.CompletedTask)
+                          .WhenAllAsync();
                 }
             }
         }

diff --git a/src/Server/SocketServer.cs b/src/Server/SocketServer.cs
@@ -1,19 +1,16 @@
 using System;
-using System.Collections.Concurrent;
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
 using Log.It;
+using Port.Server.Framework;
 
 namespace Port.Server
 {
     internal class SocketServer : INetworkServer
     {
-        private readonly ConcurrentQueue<INetworkClient> _clients =
-            new();
-
         private readonly BufferBlock<INetworkClient> _waitingClients =
             new();
 
@@ -36,7 +33,6 @@ public async Task<INetworkClient> WaitForConnectedClientAsync(
                 .ReceiveAsync(cancellationToken)
                 .ConfigureAwait(false);
             Logger.Debug("Client accepted {@client}", client);
-            _clients.Enqueue(client);
             return client;
         }
 
@@ -128,11 +124,11 @@ public async ValueTask DisposeAsync()
 
             await _acceptingClientsBackgroundTask
                 .ConfigureAwait(false);
-            while (_clients.TryDequeue(out var client))
+
+            if (_waitingClients.TryReceiveAll(out var clients))
             {
-                await client
-                    .DisposeAsync()
-                    .ConfigureAwait(false);
+                await clients.DisposeAllAsync()
+                             .ConfigureAwait(false);
             }
             Logger.Trace("Disposed");
         }