Making HttpConnection restartable (C#)

~~(a.k.a. Opening Pandora box)~~

Author: Pawel Kadluczka

samples/ClientSample/HubSample.cs

@@ -30,11 +30,22 @@ public static async Task<int> ExecuteAsync(string baseUrl)
             baseUrl = string.IsNullOrEmpty(baseUrl) ? "http://localhost:5000/default" : baseUrl;
 
             Console.WriteLine("Connecting to {0}", baseUrl);
-            HubConnection connection = await ConnectAsync(baseUrl);
-            Console.WriteLine("Connected to {0}", baseUrl);
+            var connection = new HubConnectionBuilder()
+                .WithUrl(baseUrl)
+                .WithConsoleLogger(LogLevel.Trace)
+                .Build();
 
             try
             {
+                var closeTcs = new TaskCompletionSource<object>();
+                connection.Closed += e => closeTcs.SetResult(null);
+                // Set up handler
+                connection.On<string>("Send", Console.WriteLine);
+
+                await ConnectAsync(connection);
+
+                Console.WriteLine("Connected to {0}", baseUrl);
+
                 var sendCts = new CancellationTokenSource();
 
                 Console.CancelKeyPress += async (sender, a) =>
@@ -45,13 +56,10 @@ public static async Task<int> ExecuteAsync(string baseUrl)
                     await connection.DisposeAsync();
                 };
 
-                // Set up handler
-                connection.On<string>("Send", Console.WriteLine);
-
-                while (!connection.Closed.IsCompleted)
+                while (!closeTcs.Task.IsCompleted)
                 {
-                    var completedTask = await Task.WhenAny(Task.Run(() => Console.ReadLine()), connection.Closed);
-                    if (completedTask == connection.Closed)
+                    var completedTask = await Task.WhenAny(Task.Run(() => Console.ReadLine()), closeTcs.Task);
+                    if (completedTask == closeTcs.Task)
                     {
                         break;
                     }
@@ -79,19 +87,15 @@ public static async Task<int> ExecuteAsync(string baseUrl)
             return 0;
         }
 
-        private static async Task<HubConnection> ConnectAsync(string baseUrl)
+        private static async Task ConnectAsync(HubConnection connection)
         {
             // Keep trying to until we can start
             while (true)
             {
-                var connection = new HubConnectionBuilder()
-                                .WithUrl(baseUrl)
-                                .WithConsoleLogger(LogLevel.Trace)
-                                .Build();
+
                 try
                 {
                     await connection.StartAsync();
-                    return connection;
                 }
                 catch (Exception)
                 {

samples/ClientSample/RawSample.cs

@@ -39,8 +39,9 @@ public static async Task<int> ExecuteAsync(string baseUrl)
             var connection = new HttpConnection(new Uri(baseUrl), loggerFactory);
             try
             {
+                var closeTcs = new TaskCompletionSource<object>();
+                connection.Closed += e => closeTcs.SetResult(null);
                 connection.OnReceived(data => Console.Out.WriteLineAsync($"{Encoding.UTF8.GetString(data)}"));
-
                 await connection.StartAsync();
 
                 Console.WriteLine($"Connected to {baseUrl}");
@@ -51,7 +52,7 @@ public static async Task<int> ExecuteAsync(string baseUrl)
                     await connection.DisposeAsync();
                 };
 
-                while (!connection.Closed.IsCompleted)
+                while (!closeTcs.Task.IsCompleted)
                 {
                     var line = await Task.Run(() => Console.ReadLine(), cts.Token);
 

samples/JwtClientSample/Program.cs

@@ -34,6 +34,9 @@ private async Task RunConnection(TransportType transportType)
                 .WithJwtBearer(() => _tokens[userId])
                 .Build();
 
+            var closedTcs = new TaskCompletionSource<object>();
+            hubConnection.Closed += e => closedTcs.SetResult(null);
+
             hubConnection.On<string, string>("Message", (sender, message) => Console.WriteLine($"[{userId}] {sender}: {message}"));
             await hubConnection.StartAsync();
             Console.WriteLine($"[{userId}] Connection Started");
@@ -43,7 +46,7 @@ private async Task RunConnection(TransportType transportType)
 
             try
             {
-                while (!hubConnection.Closed.IsCompleted)
+                while (!closedTcs.Task.IsCompleted)
                 {
                     await Task.Delay(1000);
                     ticks++;

src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs

@@ -32,14 +32,14 @@ public class HubConnection
         private HubProtocolReaderWriter _protocolReaderWriter;
 
         private readonly object _pendingCallsLock = new object();
-        private readonly CancellationTokenSource _connectionActive = new CancellationTokenSource();
         private readonly Dictionary<string, InvocationRequest> _pendingCalls = new Dictionary<string, InvocationRequest>();
         private readonly ConcurrentDictionary<string, List<InvocationHandler>> _handlers = new ConcurrentDictionary<string, List<InvocationHandler>>();
+        private CancellationTokenSource _connectionActive;
 
         private int _nextId = 0;
         private volatile bool _startCalled;
 
-        public Task Closed { get; }
+        public event Action<Exception> Closed;
 
         public HubConnection(IConnection connection, IHubProtocol protocol, ILoggerFactory loggerFactory)
         {
@@ -59,11 +59,7 @@ public HubConnection(IConnection connection, IHubProtocol protocol, ILoggerFacto
             _loggerFactory = loggerFactory ?? NullLoggerFactory.Instance;
             _logger = _loggerFactory.CreateLogger<HubConnection>();
             _connection.OnReceived((data, state) => ((HubConnection)state).OnDataReceivedAsync(data), this);
-            Closed = _connection.Closed.ContinueWith(task =>
-            {
-                Shutdown(task.Exception);
-                return task;
-            }).Unwrap();
+            _connection.Closed += e => Shutdown(e);
         }
 
         public async Task StartAsync()
@@ -94,12 +90,14 @@ private async Task StartAsyncCore()
 
             transferModeFeature.TransferMode = requestedTransferMode;
             await _connection.StartAsync();
+
             var actualTransferMode = transferModeFeature.TransferMode;
 
             _protocolReaderWriter = new HubProtocolReaderWriter(_protocol, GetDataEncoder(requestedTransferMode, actualTransferMode));
 
             _logger.HubProtocol(_protocol.Name);
 
+            _connectionActive = new CancellationTokenSource();
             using (var memoryStream = new MemoryStream())
             {
                 NegotiationProtocol.WriteMessage(new NegotiationMessage(_protocol.Name), memoryStream);
@@ -121,12 +119,15 @@ private IDataEncoder GetDataEncoder(TransferMode requestedTransferMode, Transfer
             return new PassThroughEncoder();
         }
 
+        public async Task StopAsync() => await StopAsyncCore().ForceAsync();
+
+        private Task StopAsyncCore() => _connection.StopAsync();
+
         public async Task DisposeAsync() => await DisposeAsyncCore().ForceAsync();
 
         private async Task DisposeAsyncCore()
         {
             await _connection.DisposeAsync();
-            await Closed;
         }
 
         // TODO: Client return values/tasks?
@@ -364,6 +365,8 @@ private void Shutdown(Exception ex = null)
                 }
                 _pendingCalls.Clear();
             }
+
+            Closed?.Invoke(ex);
         }
 
         private async Task DispatchInvocationAsync(InvocationMessage invocation, CancellationToken cancellationToken)

src/Microsoft.AspNetCore.Sockets.Abstractions/IConnection.cs

@@ -12,11 +12,12 @@ public interface IConnection
     {
         Task StartAsync();
         Task SendAsync(byte[] data, CancellationToken cancellationToken);
+        Task StopAsync();
         Task DisposeAsync();
 
         IDisposable OnReceived(Func<byte[], object, Task> callback, object state);
 
-        Task Closed { get; }
+        event Action<Exception> Closed;
 
         IFeatureCollection Features { get; }
     }