Added a way to avoid stream ID inflation on the client side in the case of unstable TCP connection

Author: sakno

src/DotNext.Tests/Net/Multiplexing/TcpMultiplexerTests.cs

@@ -239,4 +239,32 @@ protected override void Record(int value)
 
         public Task WaitForZero(TimeSpan timeout) => zeroReached.Task.WaitAsync(timeout);
     }
+
+    [Fact]
+    public static async Task WaitForConnectionAsync()
+    {
+        await using var client = new TcpMultiplexedClient(LocalEndPoint, new() { Timeout = DefaultTimeout });
+        
+        await using var server = new TcpMultiplexedListener(LocalEndPoint, new() { Timeout = DefaultTimeout });
+        await server.StartAsync();
+
+        var task = client.WaitForConnectionAsync().AsTask();
+        False(task.IsCompleted);
+
+        await client.StartAsync();
+        await task;
+        True(task.IsCompletedSuccessfully);
+    }
+
+    [Fact]
+    public static async Task WaitForDisposedConnectionAsync()
+    {
+        Task task;
+        await using (var client = new TcpMultiplexedClient(LocalEndPoint, new() { Timeout = DefaultTimeout }))
+        {
+            task = client.WaitForConnectionAsync().AsTask();
+        }
+
+        await ThrowsAsync<ObjectDisposedException>(Func.Constant(task));
+    }
 }

src/cluster/DotNext.Net.Cluster/Net/Multiplexing/MultiplexedClient.Dispatcher.cs

@@ -1,6 +1,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.IO.Pipelines;
 using System.Net.Sockets;
+using System.Runtime.ExceptionServices;
 using Microsoft.AspNetCore.Connections;
 
 namespace DotNext.Net.Multiplexing;
@@ -24,6 +25,30 @@ partial class MultiplexedClient
     private readonly OutputMultiplexer output;
     private uint streamId;
 
+    private void ReportConnected()
+        => Interlocked.Exchange(ref readiness, null)?.TrySetResult();
+
+    private void ReportDisconnected()
+    {
+        if (readiness is null)
+        {
+            var source = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            Interlocked.CompareExchange(ref readiness, source, null);
+        }
+    }
+
+    private void ReportDisposed()
+    {
+        var e = new ObjectDisposedException(GetType().Name);
+        ExceptionDispatchInfo.SetCurrentStackTrace(e);
+        if (readiness?.TrySetException(e) is null)
+        {
+            var source = new TaskCompletionSource();
+            source.SetException(e);
+            Interlocked.CompareExchange(ref readiness, source, null);
+        }
+    }
+
     private async Task DispatchAsync()
     {
         var socket = default(Socket);
@@ -43,7 +68,7 @@ private async Task DispatchAsync()
                 {
                     socket = await ConnectAsync(input.RootToken).ConfigureAwait(false);
                     receiveLoop = output.ProcessAsync(socket);
-                    readiness.TrySetResult();
+                    ReportConnected();
                 }
 
                 // send data
@@ -61,6 +86,8 @@ await input.CompleteAllAsync(new ConnectionResetException(ExceptionMessages.Conn
             {
                 socket?.Dispose();
                 await receiveLoop.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+                
+                ReportDisconnected();
                 await input.CompleteAllAsync(e).ConfigureAwait(false);
             }
         }

src/cluster/DotNext.Net.Cluster/Net/Multiplexing/MultiplexedClient.cs

@@ -13,7 +13,7 @@ namespace DotNext.Net.Multiplexing;
 [Experimental("DOTNEXT001")]
 public abstract partial class MultiplexedClient : Disposable, IAsyncDisposable
 {
-    private readonly TaskCompletionSource readiness;
+    private volatile TaskCompletionSource? readiness;
     private Task dispatcher;
 
     [SuppressMessage("Usage", "CA2213", Justification = "False positive")]
@@ -46,6 +46,9 @@ protected MultiplexedClient(Options configuration)
         output = input.CreateOutput(GC.AllocateArray<byte>(configuration.BufferCapacity, pinned: true), configuration.Timeout);
     }
 
+    private Task WaitForConnectionCoreAsync(CancellationToken token)
+        => readiness?.Task.WaitAsync(token) ?? Task.CompletedTask;
+
     /// <summary>
     /// Connects to the server and starts the dispatching loop.
     /// </summary>
@@ -54,14 +57,27 @@ protected MultiplexedClient(Options configuration)
     /// <exception cref="ObjectDisposedException">The client is disposed.</exception>
     public ValueTask StartAsync(CancellationToken token = default)
     {
+        var task = WaitForConnectionCoreAsync(token);
         if (ReferenceEquals(dispatcher, Task.CompletedTask))
         {
             dispatcher = DispatchAsync();
         }
 
-        return new(readiness.Task.WaitAsync(token));
+        return new(task);
     }
 
+    /// <summary>
+    /// Waits for the connection to be established.
+    /// </summary>
+    /// <remarks>
+    /// The method can be called to ensure that the connection to the server is established successfully.
+    /// This is useful when the underlying connection is lost to prevent inflation of the stream IDs.
+    /// </remarks>
+    /// <param name="token">The token that can be used to cancel the operation.</param>
+    /// <returns>The task representing connection state.</returns>
+    public ValueTask WaitForConnectionAsync(CancellationToken token = default)
+        => new(WaitForConnectionCoreAsync(token));
+
     /// <summary>
     /// Creates a new multiplexed client stream.
     /// </summary>
@@ -84,11 +100,16 @@ public ValueTask StartAsync(CancellationToken token = default)
     /// <returns>A duplex pipe for data input/output.</returns>
     /// <seealso cref="DotNext.IO.Pipelines.DuplexStream"/>
     public ValueTask<IDuplexPipe> OpenStreamAsync(CancellationToken token = default)
-        => readiness.Task.IsCompletedSuccessfully ? new(OpenStream()) : OpenStreamCoreAsync(token);
+    {
+        var readinessCopy = readiness;
+        return readinessCopy is null or { Task.IsCompletedSuccessfully: true }
+            ? new(OpenStream())
+            : OpenStreamCoreAsync(readinessCopy.Task, token);
+    }
 
-    private async ValueTask<IDuplexPipe> OpenStreamCoreAsync(CancellationToken token)
+    private async ValueTask<IDuplexPipe> OpenStreamCoreAsync(Task readinessTask, CancellationToken token)
     {
-        await readiness.Task.WaitAsync(token).ConfigureAwait(false);
+        await readinessTask.WaitAsync(token).ConfigureAwait(false);
         return OpenStream();
     }
 
@@ -130,7 +151,6 @@ private void Cancel()
     {
         if (Interlocked.Exchange(ref lifetimeTokenSource, null) is { } cts)
         {
-            readiness.TrySetException(new ObjectDisposedException(GetType().Name));
             using (cts)
             {
                 cts.Cancel();
@@ -143,6 +163,8 @@ protected override async ValueTask DisposeAsyncCore()
     {
         Cancel();
         await dispatcher.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
+        ReportDisposed();
+        
         await input.DisposeAsync().ConfigureAwait(false);
         await output.DisposeAsync().ConfigureAwait(false);
         writeSignal.Dispose();