Increment HubProtocol version for Stateful Reconnect (#50442)

* Increment HubProtocol version for Stateful Reconnect

* another

* fb

* fb

Author: BrennanConroy

src/Components/Server/src/BlazorPack/BlazorPackHubProtocol.cs

@@ -16,7 +16,7 @@ namespace Microsoft.AspNetCore.Components.Server.BlazorPack;
 internal sealed class BlazorPackHubProtocol : IHubProtocol
 {
     internal const string ProtocolName = "blazorpack";
-    private const int ProtocolVersion = 1;
+    private const int ProtocolVersion = 2;
 
     private readonly BlazorPackHubProtocolWorker _worker = new BlazorPackHubProtocolWorker();
 
@@ -32,7 +32,7 @@ internal sealed class BlazorPackHubProtocol : IHubProtocol
     /// <inheritdoc />
     public bool IsVersionSupported(int version)
     {
-        return version == Version;
+        return version <= Version;
     }
 
     /// <inheritdoc />

src/SignalR/clients/csharp/Client.Core/src/HubConnection.Log.cs

@@ -334,5 +334,8 @@ public static void ErrorHandshakeTimedOut(ILogger logger, TimeSpan handshakeTime
 
         [LoggerMessage(92, LogLevel.Trace, "Received SequenceMessage with Sequence ID '{SequenceId}'.", EventName = "ReceivedSequenceMessage")]
         public static partial void ReceivedSequenceMessage(ILogger logger, long sequenceId);
+
+        [LoggerMessage(93, LogLevel.Debug, "HubProtocol '{Protocol} v{Version}' does not support Stateful Reconnect. Disabling the feature.", EventName = "DisablingReconnect")]
+        public static partial void DisablingReconnect(ILogger logger, string protocol, int version);
     }
 }

src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs

@@ -481,12 +481,35 @@ private async Task StartAsyncCore(CancellationToken cancellationToken)
         var connection = await _connectionFactory.ConnectAsync(_endPoint, cancellationToken).ConfigureAwait(false);
         var startingConnectionState = new ConnectionState(connection, this);
 
+#pragma warning disable CA2252 // This API requires opting into preview features
+        var statefulReconnectFeature = connection.Features.Get<IStatefulReconnectFeature>();
+#pragma warning restore CA2252 // This API requires opting into preview features
+
         // From here on, if an error occurs we need to shut down the connection because
         // we still own it.
         try
         {
-            Log.HubProtocol(_logger, _protocol.Name, _protocol.Version);
-            await HandshakeAsync(startingConnectionState, cancellationToken).ConfigureAwait(false);
+            var usedProtocolVersion = _protocol.Version;
+            if (statefulReconnectFeature is null && _protocol.IsVersionSupported(1))
+            {
+                // Stateful Reconnect starts with HubProtocol version 2, newer clients connecting to older servers will fail to connect due to
+                // the handshake only supporting version 1, so we will try to send version 1 during the handshake to keep old servers working
+                // if the client is not attempting to enable stateful reconnect and therefore does not require a newer HubProtocol.
+                usedProtocolVersion = 1;
+            }
+            else if (_protocol.Version < 2)
+            {
+                if (statefulReconnectFeature is not null)
+                {
+                    Log.DisablingReconnect(_logger, _protocol.Name, _protocol.Version);
+#pragma warning disable CA2252 // This API requires opting into preview features
+                    statefulReconnectFeature.DisableReconnect();
+#pragma warning restore CA2252 // This API requires opting into preview features
+                }
+            }
+
+            Log.HubProtocol(_logger, _protocol.Name, usedProtocolVersion);
+            await HandshakeAsync(startingConnectionState, usedProtocolVersion, cancellationToken).ConfigureAwait(false);
         }
         catch (Exception ex)
         {
@@ -1240,12 +1263,12 @@ private void CheckDisposed()
         ObjectDisposedThrowHelper.ThrowIf(_disposed, this);
     }
 
-    private async Task HandshakeAsync(ConnectionState startingConnectionState, CancellationToken cancellationToken)
+    private async Task HandshakeAsync(ConnectionState startingConnectionState, int protocolVersion, CancellationToken cancellationToken)
     {
         // Send the Handshake request
         Log.SendingHubHandshake(_logger);
 
-        var handshakeRequest = new HandshakeRequestMessage(_protocol.Name, _protocol.Version);
+        var handshakeRequest = new HandshakeRequestMessage(_protocol.Name, protocolVersion);
         HandshakeProtocol.WriteRequestMessage(handshakeRequest, startingConnectionState.Connection.Transport.Output);
 
         var sendHandshakeResult = await startingConnectionState.Connection.Transport.Output.FlushAsync(CancellationToken.None).ConfigureAwait(false);

src/SignalR/clients/csharp/Client/test/FunctionalTests/HubConnectionTests.cs

@@ -2833,6 +2833,83 @@ public async Task CanSetMessageBufferSizeOnClient()
         }
     }
 
+    [Fact]
+    public async Task ServerWithOldProtocolVersionClientWithNewProtocolVersionWorksDoesNotAllowStatefulReconnect()
+    {
+        bool ExpectedErrors(WriteContext writeContext)
+        {
+            return writeContext.LoggerName == typeof(HubConnection).FullName &&
+                   (writeContext.EventId.Name == "ShutdownWithError" ||
+                   writeContext.EventId.Name == "ServerDisconnectedWithError");
+        }
+
+        var protocol = HubProtocols["json"];
+        await using (var server = await StartServer<Startup>(ExpectedErrors))
+        {
+            var websocket = new ClientWebSocket();
+            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            tcs.SetResult();
+
+            const string originalMessage = "SignalR";
+            var connectionBuilder = new HubConnectionBuilder()
+                .WithLoggerFactory(LoggerFactory)
+                .WithUrl(server.Url + "/default", HttpTransportType.WebSockets, o =>
+                {
+                    o.WebSocketFactory = async (context, token) =>
+                    {
+                        await tcs.Task;
+                        await websocket.ConnectAsync(context.Uri, token);
+                        tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+                        return websocket;
+                    };
+                    o.UseStatefulReconnect = true;
+                });
+            // Force version 1 on the server so it turns off Stateful Reconnects
+            connectionBuilder.Services.AddSingleton<IHubProtocol>(new HubProtocolVersionTests.SingleVersionHubProtocol(HubProtocols["json"], 1));
+            var connection = connectionBuilder.Build();
+
+            var closedTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            connection.Closed += (_) =>
+            {
+                closedTcs.SetResult();
+                return Task.CompletedTask;
+            };
+
+            try
+            {
+                await connection.StartAsync().DefaultTimeout();
+                var originalConnectionId = connection.ConnectionId;
+
+                var result = await connection.InvokeAsync<string>(nameof(TestHub.Echo), originalMessage).DefaultTimeout();
+
+                Assert.Equal(originalMessage, result);
+
+                var originalWebsocket = websocket;
+                websocket = new ClientWebSocket();
+                originalWebsocket.Dispose();
+
+                var resultTask = connection.InvokeAsync<string>(nameof(TestHub.Echo), originalMessage).DefaultTimeout();
+                tcs.SetResult();
+
+                // In-progress send canceled when connection closes
+                var ex = await Assert.ThrowsAnyAsync<Exception>(() => resultTask);
+                Assert.True(ex is TaskCanceledException or WebSocketException);
+                await closedTcs.Task;
+
+                Assert.Equal(HubConnectionState.Disconnected, connection.State);
+            }
+            catch (Exception ex)
+            {
+                LoggerFactory.CreateLogger<HubConnectionTests>().LogError(ex, "{ExceptionType} from test", ex.GetType().FullName);
+                throw;
+            }
+            finally
+            {
+                await connection.DisposeAsync().DefaultTimeout();
+            }
+        }
+    }
+
     private class OneAtATimeSynchronizationContext : SynchronizationContext, IAsyncDisposable
     {
         private readonly Channel<(SendOrPostCallback, object)> _taskQueue = Channel.CreateUnbounded<(SendOrPostCallback, object)>();

src/SignalR/clients/csharp/Client/test/FunctionalTests/HubProtocolVersionTests.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Net;
 using System.Text;
 using System.Threading;
@@ -171,7 +172,7 @@ bool ExpectedErrors(WriteContext writeContext)
             var connectionBuilder = new HubConnectionBuilder()
                 .WithLoggerFactory(LoggerFactory)
                 .WithUrl(server.Url + "/version", transportType);
-            connectionBuilder.Services.AddSingleton<IHubProtocol>(new VersionedJsonHubProtocol(int.MaxValue));
+            connectionBuilder.Services.AddSingleton<IHubProtocol>(new SingleVersionHubProtocol(new VersionedJsonHubProtocol(int.MaxValue), int.MaxValue));
 
             var connection = connectionBuilder.Build();
 
@@ -193,6 +194,41 @@ await ExceptionAssert.ThrowsAsync<HubException>(
         }
     }
 
+    public class SingleVersionHubProtocol : IHubProtocol
+    {
+        private readonly IHubProtocol _protocol;
+        private readonly int _version;
+
+        public SingleVersionHubProtocol(IHubProtocol inner, int version)
+        {
+            _protocol = inner;
+            _version = version;
+        }
+
+        public string Name => _protocol.Name;
+
+        public int Version => _version;
+
+        public TransferFormat TransferFormat => _protocol.TransferFormat;
+
+        public ReadOnlyMemory<byte> GetMessageBytes(HubMessage message) => _protocol.GetMessageBytes(message);
+
+        public bool IsVersionSupported(int version)
+        {
+            return version == _version;
+        }
+
+        public bool TryParseMessage(ref ReadOnlySequence<byte> input, IInvocationBinder binder, [NotNullWhen(true)] out HubMessage message)
+        {
+            return _protocol.TryParseMessage(ref input, binder, out message);
+        }
+
+        public void WriteMessage(HubMessage message, IBufferWriter<byte> output)
+        {
+            _protocol.WriteMessage(message, output);
+        }
+    }
+
     private class ProxyConnectionFactory : IConnectionFactory
     {
         private readonly IConnectionFactory _innerFactory;

src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/WebSocketsTransport.cs

@@ -372,6 +372,9 @@ private async Task ProcessSocketAsync(WebSocket socket, Uri url, bool isReconnec
 
                     // Abort the websocket if we're stuck in a pending send to the client
                     socket.Abort();
+
+                    // Should not throw
+                    await sending.ConfigureAwait(false);
                 }
             }
             else
@@ -390,18 +393,40 @@ private async Task ProcessSocketAsync(WebSocket socket, Uri url, bool isReconnec
                 {
                     _application.Output.CancelPendingFlush();
                 }
+
+                // Should not throw
+                await receiving.ConfigureAwait(false);
             }
         }
 
-        if (_useStatefulReconnect && !_gracefulClose)
+        var cleanup = true;
+        try
         {
-            if (!UpdateConnectionPair())
+            if (!_gracefulClose && UpdateConnectionPair())
             {
-                return;
+                try
+                {
+                    await StartAsync(url, _webSocketMessageType == WebSocketMessageType.Binary ? TransferFormat.Binary : TransferFormat.Text,
+                        cancellationToken: default).ConfigureAwait(false);
+                    cleanup = false;
+                }
+                catch (Exception ex)
+                {
+                    throw new InvalidOperationException("Reconnect attempt failed.", innerException: ex);
+                }
+            }
+        }
+        finally
+        {
+            if (cleanup)
+            {
+                // Pipes will usually already be completed.
+                // If stateful reconnect fails we want to make sure the Pipes are cleaned up.
+                // And in rare cases where the websocket is closing at the same time StopAsync is called
+                // It's possible a Pipe won't be completed so let's be safe and call Complete again.
+                _application.Output.Complete();
+                _application.Input.Complete();
             }
-
-            await StartAsync(url, _webSocketMessageType == WebSocketMessageType.Binary ? TransferFormat.Binary : TransferFormat.Text,
-                cancellationToken: default).ConfigureAwait(false);
         }
     }
 
@@ -479,7 +504,7 @@ private async Task StartReceiving(WebSocket socket)
         {
             if (!_aborted)
             {
-                if (_gracefulClose)
+                if (_gracefulClose || !_useStatefulReconnect)
                 {
                     _application.Output.Complete(ex);
                 }
@@ -493,7 +518,7 @@ private async Task StartReceiving(WebSocket socket)
         finally
         {
             // We're done writing
-            if (_gracefulClose)
+            if (_gracefulClose || !_useStatefulReconnect)
             {
                 _application.Output.Complete();
             }
@@ -589,7 +614,7 @@ private async Task StartSending(WebSocket socket, bool ignoreFirstCanceled)
                 }
             }
 
-            if (_gracefulClose)
+            if (_gracefulClose || !_useStatefulReconnect)
             {
                 _application.Input.Complete(error);
             }

src/SignalR/clients/ts/signalr-protocol-msgpack/src/MessagePackHubProtocol.ts

@@ -26,7 +26,7 @@ export class MessagePackHubProtocol implements IHubProtocol {
     /** The name of the protocol. This is used by SignalR to resolve the protocol between the client and server. */
     public readonly name: string = "messagepack";
     /** The version of the protocol. */
-    public readonly version: number = 1;
+    public readonly version: number = 2;
     /** The TransferFormat of the protocol. */
     public readonly transferFormat: TransferFormat = TransferFormat.Binary;
 

src/SignalR/clients/ts/signalr/src/HubConnection.ts

@@ -228,9 +228,16 @@ export class HubConnection {
         await this.connection.start(this._protocol.transferFormat);
 
         try {
+            let version = this._protocol.version;
+            if (!this.connection.features.reconnect) {
+                // Stateful Reconnect starts with HubProtocol version 2, newer clients connecting to older servers will fail to connect due to
+                // the handshake only supporting version 1, so we will try to send version 1 during the handshake to keep old servers working.
+                version = 1;
+            }
+
             const handshakeRequest: HandshakeRequestMessage = {
                 protocol: this._protocol.name,
-                version: this._protocol.version,
+                version,
             };
 
             this._logger.log(LogLevel.Debug, "Sending handshake request.");
@@ -257,8 +264,8 @@ export class HubConnection {
                 throw this._stopDuringStartError;
             }
 
-            const useAck = this.connection.features.reconnect || false;
-            if (useAck) {
+            const useStatefulReconnect = this.connection.features.reconnect || false;
+            if (useStatefulReconnect) {
                 this._messageBuffer = new MessageBuffer(this._protocol, this.connection, this._statefulReconnectBufferSize);
                 this.connection.features.disconnected = this._messageBuffer._disconnected.bind(this._messageBuffer);
                 this.connection.features.resend = () => {

src/SignalR/clients/ts/signalr/src/JsonHubProtocol.ts

@@ -15,7 +15,7 @@ export class JsonHubProtocol implements IHubProtocol {
     /** @inheritDoc */
     public readonly name: string = JSON_HUB_PROTOCOL_NAME;
     /** @inheritDoc */
-    public readonly version: number = 1;
+    public readonly version: number = 2;
 
     /** @inheritDoc */
     public readonly transferFormat: TransferFormat = TransferFormat.Text;

src/SignalR/common/Protocols.Json/src/Protocol/JsonHubProtocol.cs

@@ -47,7 +47,7 @@ public sealed class JsonHubProtocol : IHubProtocol
     private static readonly JsonEncodedText SequenceIdPropertyNameBytes = JsonEncodedText.Encode(SequenceIdPropertyName);
 
     private const string ProtocolName = "json";
-    private const int ProtocolVersion = 1;
+    private const int ProtocolVersion = 2;
 
     /// <summary>
     /// Gets the serializer used to serialize invocation arguments and return values.
@@ -82,7 +82,7 @@ public JsonHubProtocol(IOptions<JsonHubProtocolOptions> options)
     /// <inheritdoc />
     public bool IsVersionSupported(int version)
     {
-        return version == Version;
+        return version <= Version;
     }
 
     /// <inheritdoc />