Open stream only when the client is connected

Author: sakno

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

@@ -248,7 +248,7 @@ public static async Task WaitForConnectionAsync()
         await using var server = new TcpMultiplexedListener(LocalEndPoint, new() { Timeout = DefaultTimeout });
         await server.StartAsync();
 
-        var task = client.EnsureConnectedAsync().AsTask();
+        var task = client.OpenStreamAsync().AsTask();
         False(task.IsCompleted);
 
         await client.StartAsync();
@@ -262,7 +262,7 @@ public static async Task WaitForDisposedConnectionAsync()
         Task task;
         await using (var client = new TcpMultiplexedClient(LocalEndPoint, new() { Timeout = DefaultTimeout }))
         {
-            task = client.EnsureConnectedAsync().AsTask();
+            task = client.OpenStreamAsync().AsTask();
         }
 
         await ThrowsAsync<ObjectDisposedException>(Func.Constant(task));

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

@@ -29,7 +29,14 @@ public bool TryAddStream(uint streamId, MultiplexedStream stream)
         ChangeStreamCount(Unsafe.BitCast<bool, byte>(result));
         return result;
     }
-    
+
+    public bool TryRemoveStream(uint streamId, MultiplexedStream stream)
+    {
+        var removed = streams.TryRemove(new(streamId, stream));
+        ChangeStreamCount(-Unsafe.BitCast<bool, byte>(removed));
+        return removed;
+    }
+
     public OutputMultiplexer CreateOutput(Memory<byte> framingBuffer, TimeSpan receiveTimeout)
         => new(streams, writeSignal, commands, framingBuffer, streamCounter, measurementTags, receiveTimeout, RootToken);
 
@@ -52,10 +59,9 @@ public async Task ProcessAsync(Func<bool> condition, Socket socket)
             {
                 var (streamId, stream) = enumerator.Current;
 
-                if (stream.IsCompleted && streams.TryRemove(streamId, out _))
+                if (stream.IsCompleted && TryRemoveStream(streamId, stream))
                 {
                     Protocol.WriteStreamClosed(framingBuffer, streamId);
-                    ChangeStreamCount(-1);
                 }
                 else
                 {

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

@@ -46,7 +46,7 @@ protected MultiplexedClient(Options configuration)
         output = input.CreateOutput(GC.AllocateArray<byte>(configuration.BufferCapacity, pinned: true), configuration.Timeout);
     }
 
-    private Task EnsureConnectedCoreAsync(CancellationToken token)
+    private Task EnsureConnectedAsync(CancellationToken token)
         => readiness?.Task.WaitAsync(token) ?? Task.CompletedTask;
 
     /// <summary>
@@ -57,7 +57,7 @@ private Task EnsureConnectedCoreAsync(CancellationToken token)
     /// <exception cref="ObjectDisposedException">The client is disposed.</exception>
     public ValueTask StartAsync(CancellationToken token = default)
     {
-        var task = EnsureConnectedCoreAsync(token);
+        var task = EnsureConnectedAsync(token);
         if (ReferenceEquals(dispatcher, Task.CompletedTask))
         {
             dispatcher = DispatchAsync();
@@ -66,18 +66,6 @@ public ValueTask StartAsync(CancellationToken token = default)
         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 EnsureConnectedAsync(CancellationToken token = default)
-        => new(EnsureConnectedCoreAsync(token));
-
     /// <summary>
     /// Creates a new multiplexed client stream.
     /// </summary>
@@ -99,29 +87,31 @@ public ValueTask EnsureConnectedAsync(CancellationToken token = default)
     /// </remarks>
     /// <returns>A duplex pipe for data input/output.</returns>
     /// <seealso cref="DotNext.IO.Pipelines.DuplexStream"/>
-    public ValueTask<IDuplexPipe> OpenStreamAsync(CancellationToken token = default)
-    {
-        var readinessCopy = readiness;
-        return readinessCopy is null or { Task.IsCompletedSuccessfully: true }
-            ? new(OpenStream())
-            : OpenStreamCoreAsync(readinessCopy.Task, token);
-    }
-
-    private async ValueTask<IDuplexPipe> OpenStreamCoreAsync(Task readinessTask, CancellationToken token)
+    public async ValueTask<IDuplexPipe> OpenStreamAsync(CancellationToken token = default)
     {
-        await readinessTask.WaitAsync(token).ConfigureAwait(false);
-        return OpenStream();
+        for (var stream = OpenStream(out var addedStreamId);; token.ThrowIfCancellationRequested())
+        {
+            try
+            {
+                await EnsureConnectedAsync(token).ConfigureAwait(false);
+                return stream;
+            }
+            catch (Exception e)
+            {
+                input.TryRemoveStream(addedStreamId, stream);
+                await stream.AbortAppSideAsync(e).ConfigureAwait(false);
+            }
+        }
     }
 
-    private MultiplexedStream OpenStream()
+    private MultiplexedStream OpenStream(out uint id)
     {
         var stream = new MultiplexedStream(options, writeSignal);
-        uint id;
         do
         {
             id = Interlocked.Increment(ref streamId);
         } while (!input.TryAddStream(id, stream));
-        
+
         return stream;
     }
 

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

@@ -85,10 +85,12 @@ private bool IsTransportInputCompleted(out bool appSideCompleted)
         return (stateCopy & TransportInputCompletedState) is not 0U;
     }
 
-    public async ValueTask AbortAppSideAsync()
+    public async ValueTask AbortAppSideAsync(Exception? e = null)
     {
-        var e = new ConnectionAbortedException();
-        ExceptionDispatchInfo.SetCurrentStackTrace(e);
+        if (e is null)
+        {
+            ExceptionDispatchInfo.SetCurrentStackTrace(e = new ConnectionAbortedException());
+        }
 
         await appWriter.CompleteAsync(e).ConfigureAwait(false);
         await appReader.CompleteAsync(e).ConfigureAwait(false);