Making HttpConnection restartable (C#)

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

Author: Pawel Kadluczka

SignalR.sln

@@ -9,6 +9,7 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "src", "src", "{DA69F624-539
 EndProject
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Solution Items", "Solution Items", "{83B2C3EB-A3D8-4E6F-9A3C-A380B005EF31}"
 	ProjectSection(SolutionItems) = preProject
+		build\dependencies.props = build\dependencies.props
 		Directory.Build.props = Directory.Build.props
 		Directory.Build.targets = Directory.Build.targets
 		build\Key.snk = build\Key.snk
@@ -85,7 +86,7 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "JwtSample", "samples\JwtSam
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "JwtClientSample", "samples\JwtClientSample\JwtClientSample.csproj", "{1A953296-E869-4DE2-A693-FD5FCDE27057}"
 EndProject
-Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Microsoft.AspNetCore.SignalR.Tests.Utils", "test\Microsoft.AspNetCore.SignalR.Tests.Utils\Microsoft.AspNetCore.SignalR.Tests.Utils.csproj", "{0A0A6135-EA24-4307-95C2-CE1B7E164A5E}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.AspNetCore.SignalR.Tests.Utils", "test\Microsoft.AspNetCore.SignalR.Tests.Utils\Microsoft.AspNetCore.SignalR.Tests.Utils.csproj", "{0A0A6135-EA24-4307-95C2-CE1B7E164A5E}"
 EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution

client-ts/package-lock.json

@@ -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/HubSample.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/ClientSample/RawSample.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++;

samples/JwtClientSample/Program.cs

@@ -34,16 +34,16 @@ 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;
         private Timer _timeoutTimer;
         private bool _needKeepAlive;
 
-        public Task Closed { get; }
+        public event Action<Exception> Closed;
 
         /// <summary>
         /// Gets or sets the server timeout interval for the connection. Changes to this value
@@ -69,11 +69,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);
 
             // Create the timer for timeout, but disabled by default (we enable it when started).
             _timeoutTimer = new Timer(state => ((HubConnection)state).TimeoutElapsed(), this, Timeout.Infinite, Timeout.Infinite);
@@ -122,12 +118,14 @@ private async Task StartAsyncCore()
             transferModeFeature.TransferMode = requestedTransferMode;
             await _connection.StartAsync();
             _needKeepAlive = _connection.Features.Get<IConnectionInherentKeepAliveFeature>() == null;
+
             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);
@@ -151,13 +149,16 @@ 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()
         {
             _timeoutTimer.Dispose();
             await _connection.DisposeAsync();
-            await Closed;
         }
 
         // TODO: Client return values/tasks?
@@ -370,12 +371,12 @@ private async Task OnDataReceivedAsync(byte[] data)
             }
         }
 
-        private void Shutdown(Exception ex = null)
+        private void Shutdown(Exception exception = null)
         {
             _logger.ShutdownConnection();
-            if (ex != null)
+            if (exception != null)
             {
-                _logger.ShutdownWithError(ex);
+                _logger.ShutdownWithError(exception);
             }
 
             lock (_pendingCallsLock)
@@ -388,14 +389,23 @@ private void Shutdown(Exception ex = null)
                 foreach (var outstandingCall in _pendingCalls.Values)
                 {
                     _logger.RemoveInvocation(outstandingCall.InvocationId);
-                    if (ex != null)
+                    if (exception != null)
                     {
-                        outstandingCall.Fail(ex);
+                        outstandingCall.Fail(exception);
                     }
                     outstandingCall.Dispose();
                 }
                 _pendingCalls.Clear();
             }
+
+            try
+            {
+                Closed?.Invoke(exception);
+            }
+            catch (Exception ex)
+            {
+                _logger.ErrorDuringClosedEvent(ex);
+            }
         }
 
         private async Task DispatchInvocationAsync(InvocationMessage invocation, CancellationToken cancellationToken)

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

@@ -88,6 +88,9 @@ internal static class SignalRClientLoggerExtensions
         private static readonly Action<ILogger, Exception> _resettingKeepAliveTimer =
             LoggerMessage.Define(LogLevel.Trace, new EventId(25, nameof(ResettingKeepAliveTimer)), "Resetting keep-alive timer, received a message from the server.");
 
+        private static readonly Action<ILogger, Exception> _errorDuringClosedEvent =
+            LoggerMessage.Define(LogLevel.Error, new EventId(26, nameof(ErrorDuringClosedEvent)), "An exception was thrown in the handler for the Closed event.");
+
         // Category: Streaming and NonStreaming
         private static readonly Action<ILogger, string, Exception> _invocationCreated =
             LoggerMessage.Define<string>(LogLevel.Trace, new EventId(0, nameof(InvocationCreated)), "Invocation {invocationId} created.");
@@ -292,5 +295,10 @@ public static void ResettingKeepAliveTimer(this ILogger logger)
         {
             _resettingKeepAliveTimer(logger, null);
         }
+
+        public static void ErrorDuringClosedEvent(this ILogger logger, Exception exception)
+        {
+            _errorDuringClosedEvent(logger, exception);
+        }
     }
 }

src/Microsoft.AspNetCore.SignalR.Client.Core/Internal/SignalRClientLoggerExtensions.cs

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