[SignalR] Seamless Reconnect

src/Servers/Connections.Abstractions/src/Features/IReconnectFeature.cs

@@ -0,0 +1,24 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.Connections.Abstractions;
+
+/// <summary>
+/// 
+/// </summary>
+public interface IReconnectFeature
+{
+    /// <summary>
+    /// 
+    /// </summary>
+    public Action NotifyOnReconnect { get; set; }
+
+    // TODO
+    // void DisableReconnect();
+}

src/Servers/Connections.Abstractions/src/PublicAPI/net462/PublicAPI.Unshipped.txt

@@ -1,4 +1,7 @@
 #nullable enable
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.get -> System.Action!
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.set -> void
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature.Tags.get -> System.Collections.Generic.ICollection<System.Collections.Generic.KeyValuePair<string!, object?>>!
 Microsoft.AspNetCore.Connections.Features.IConnectionNamedPipeFeature

src/Servers/Connections.Abstractions/src/PublicAPI/net8.0/PublicAPI.Unshipped.txt

@@ -1,4 +1,7 @@
 #nullable enable
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.get -> System.Action!
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.set -> void
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature.Tags.get -> System.Collections.Generic.ICollection<System.Collections.Generic.KeyValuePair<string!, object?>>!
 Microsoft.AspNetCore.Connections.Features.IConnectionNamedPipeFeature

src/Servers/Connections.Abstractions/src/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt

@@ -1,4 +1,7 @@
 #nullable enable
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.get -> System.Action!
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.set -> void
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature.Tags.get -> System.Collections.Generic.ICollection<System.Collections.Generic.KeyValuePair<string!, object?>>!
 Microsoft.AspNetCore.Connections.Features.IConnectionNamedPipeFeature

src/Servers/Connections.Abstractions/src/PublicAPI/netstandard2.1/PublicAPI.Unshipped.txt

@@ -1,4 +1,7 @@
 #nullable enable
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.get -> System.Action!
+Microsoft.AspNetCore.Connections.Abstractions.IReconnectFeature.NotifyOnReconnect.set -> void
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature
 Microsoft.AspNetCore.Connections.Features.IConnectionMetricsTagsFeature.Tags.get -> System.Collections.Generic.ICollection<System.Collections.Generic.KeyValuePair<string!, object?>>!
 Microsoft.AspNetCore.Connections.Features.IConnectionNamedPipeFeature

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

@@ -8,6 +8,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.IO;
+using System.IO.Pipelines;
 using System.Linq;
 using System.Net;
 using System.Reflection;
@@ -16,6 +17,7 @@
 using System.Threading.Channels;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Connections.Abstractions;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Internal;
 using Microsoft.AspNetCore.Shared;
@@ -946,11 +948,19 @@ private async Task InvokeStreamCore(ConnectionState connectionState, string meth
     private async Task SendHubMessage(ConnectionState connectionState, HubMessage hubMessage, CancellationToken cancellationToken = default)
     {
         _state.AssertConnectionValid();
-        _protocol.WriteMessage(hubMessage, connectionState.Connection.Transport.Output);
 
         Log.SendingMessage(_logger, hubMessage);
 
-        await connectionState.Connection.Transport.Output.FlushAsync(cancellationToken).ConfigureAwait(false);
+        if (connectionState.UsingAcks())
+        {
+            await connectionState.WriteAsync(new SerializedHubMessage(hubMessage), cancellationToken).ConfigureAwait(false);
+        }
+        else
+        {
+            _protocol.WriteMessage(hubMessage, connectionState.Connection.Transport.Output);
+
+            await connectionState.Connection.Transport.Output.FlushAsync(cancellationToken).ConfigureAwait(false);
+        }
         Log.MessageSent(_logger, hubMessage);
 
         // We've sent a message, so don't ping for a while
@@ -1004,6 +1014,15 @@ private async Task SendWithLock(ConnectionState expectedConnectionState, HubMess
         Log.ResettingKeepAliveTimer(_logger);
         connectionState.ResetTimeout();
 
+        if (connectionState.UsingAcks())
+        {
+            if (!connectionState.ShouldProcessMessage(message))
+            {
+                _logger.LogInformation($"Dropped {((HubInvocationMessage)message).GetType().Name}. ID: {((HubInvocationMessage)message).InvocationId}");
+                return null;
+            }
+        }
+
         InvocationRequest? irq;
         switch (message)
         {
@@ -1055,6 +1074,20 @@ private async Task SendWithLock(ConnectionState expectedConnectionState, HubMess
                 Log.ReceivedPing(_logger);
                 // timeout is reset above, on receiving any message
                 break;
+            case AckMessage ackMessage:
+                _logger.LogInformation("Received Ack with ID {id}", ackMessage.SequenceId);
+                if (connectionState.UsingAcks())
+                {
+                    connectionState.Ack(ackMessage);
+                }
+                break;
+            case SequenceMessage sequenceMessage:
+                _logger.LogInformation("Received SequenceMessage with ID {id}", sequenceMessage.SequenceId);
+                if (connectionState.UsingAcks())
+                {
+                    connectionState.ResetSequence(sequenceMessage);
+                }
+                break;
             default:
                 throw new InvalidOperationException($"Unexpected message type: {message.GetType().FullName}");
         }
@@ -1235,6 +1268,7 @@ private async Task HandshakeAsync(ConnectionState startingConnectionState, Cance
                                 }
 
                                 Log.HandshakeComplete(_logger);
+
                                 break;
                             }
                         }
@@ -1813,6 +1847,7 @@ private sealed class ConnectionState : IInvocationBinder
         private readonly HubConnection _hubConnection;
         private readonly ILogger _logger;
         private readonly bool _hasInherentKeepAlive;
+        private readonly MessageBuffer? _messageBuffer;
 
         private readonly object _lock = new object();
         private readonly Dictionary<string, InvocationRequest> _pendingCalls = new Dictionary<string, InvocationRequest>(StringComparer.Ordinal);
@@ -1850,6 +1885,13 @@ public ConnectionState(ConnectionContext connection, HubConnection hubConnection
 
             _logger = _hubConnection._logger;
             _hasInherentKeepAlive = connection.Features.Get<IConnectionInherentKeepAliveFeature>()?.HasInherentKeepAlive ?? false;
+
+            if (Connection.Features.Get<IReconnectFeature>() is IReconnectFeature feature)
+            {
+                _messageBuffer = new MessageBuffer(connection, hubConnection._protocol);
+
+                feature.NotifyOnReconnect = _messageBuffer.Resend;
+            }
         }
 
         public string GetNextId() => (++_nextInvocationId).ToString(CultureInfo.InvariantCulture);
@@ -1935,6 +1977,8 @@ private async Task StopAsyncCore()
         {
             Log.Stopping(_logger);
 
+            _messageBuffer?.Dispose();
+
             // Complete our write pipe, which should cause everything to shut down
             Log.TerminatingReceiveLoop(_logger);
             Connection.Transport.Input.CancelPendingRead();
@@ -1966,6 +2010,33 @@ public async Task TimerLoop(TimerAwaitable timer)
             }
         }
 
+        public ValueTask<FlushResult> WriteAsync(SerializedHubMessage message, CancellationToken cancellationToken)
+        {
+            Debug.Assert(_messageBuffer is not null);
+            return _messageBuffer.WriteAsync(message, cancellationToken);
+        }
+
+        public bool ShouldProcessMessage(HubMessage message)
+        {
+            Debug.Assert(_messageBuffer is not null);
+            return _messageBuffer.ShouldProcessMessage(message);
+        }
+
+        public void Ack(AckMessage ackMessage)
+        {
+            Debug.Assert(_messageBuffer is not null);
+            _messageBuffer.Ack(ackMessage);
+        }
+
+        public void ResetSequence(SequenceMessage sequenceMessage)
+        {
+            Debug.Assert(_messageBuffer is not null);
+            _messageBuffer.ResetSequence(sequenceMessage);
+        }
+
+        [MemberNotNullWhen(true, nameof(_messageBuffer))]
+        public bool UsingAcks() => _messageBuffer is not null;
+
         public void ResetSendPing()
         {
             Volatile.Write(ref _nextActivationSendPing, (DateTime.UtcNow + _hubConnection.KeepAliveInterval).Ticks);

src/SignalR/clients/csharp/Client.Core/src/Internal/SerializedHubMessage.cs

@@ -0,0 +1,191 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.SignalR.Protocol;
+
+namespace Microsoft.AspNetCore.SignalR.Internal;
+
+/// <summary>
+/// Represents a serialization cache for a single message.
+/// </summary>
+internal class SerializedHubMessage
+{
+    private SerializedMessage _cachedItem1;
+    private SerializedMessage _cachedItem2;
+    private List<SerializedMessage>? _cachedItems;
+    private readonly object _lock = new object();
+
+    /// <summary>
+    /// Gets the hub message for the serialization cache.
+    /// </summary>
+    public HubMessage? Message { get; }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SerializedHubMessage"/> class.
+    /// </summary>
+    /// <param name="messages">A collection of already serialized messages to cache.</param>
+    public SerializedHubMessage(IReadOnlyList<SerializedMessage> messages)
+    {
+        // A lock isn't needed here because nobody has access to this type until the constructor finishes.
+        for (var i = 0; i < messages.Count; i++)
+        {
+            var message = messages[i];
+            SetCacheUnsynchronized(message.ProtocolName, message.Serialized);
+        }
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SerializedHubMessage"/> class.
+    /// </summary>
+    /// <param name="message">The hub message for the cache. This will be serialized with an <see cref="IHubProtocol"/> in <see cref="GetSerializedMessage"/> to get the message's serialized representation.</param>
+    public SerializedHubMessage(HubMessage message)
+    {
+        Message = message;
+    }
+
+    /// <summary>
+    /// Gets the serialized representation of the <see cref="HubMessage"/> using the specified <see cref="IHubProtocol"/>.
+    /// </summary>
+    /// <param name="protocol">The protocol used to create the serialized representation.</param>
+    /// <returns>The serialized representation of the <see cref="HubMessage"/>.</returns>
+    public ReadOnlyMemory<byte> GetSerializedMessage(IHubProtocol protocol)
+    {
+        lock (_lock)
+        {
+            if (!TryGetCachedUnsynchronized(protocol.Name, out var serialized))
+            {
+                if (Message == null)
+                {
+                    throw new InvalidOperationException(
+                        "This message was received from another server that did not have the requested protocol available.");
+                }
+
+                serialized = protocol.GetMessageBytes(Message);
+                SetCacheUnsynchronized(protocol.Name, serialized);
+            }
+
+            return serialized;
+        }
+    }
+
+    // Used for unit testing.
+    internal IReadOnlyList<SerializedMessage> GetAllSerializations()
+    {
+        // Even if this is only used in tests, let's do it right.
+        lock (_lock)
+        {
+            if (_cachedItem1.ProtocolName == null)
+            {
+                return Array.Empty<SerializedMessage>();
+            }
+
+            var list = new List<SerializedMessage>(2);
+            list.Add(_cachedItem1);
+
+            if (_cachedItem2.ProtocolName != null)
+            {
+                list.Add(_cachedItem2);
+
+                if (_cachedItems != null)
+                {
+                    list.AddRange(_cachedItems);
+                }
+            }
+
+            return list;
+        }
+    }
+
+    private void SetCacheUnsynchronized(string protocolName, ReadOnlyMemory<byte> serialized)
+    {
+        // We set the fields before moving on to the list, if we need it to hold more than 2 items.
+        // We have to read/write these fields under the lock because the structs might tear and another
+        // thread might observe them half-assigned
+
+        if (_cachedItem1.ProtocolName == null)
+        {
+            _cachedItem1 = new SerializedMessage(protocolName, serialized);
+        }
+        else if (_cachedItem2.ProtocolName == null)
+        {
+            _cachedItem2 = new SerializedMessage(protocolName, serialized);
+        }
+        else
+        {
+            if (_cachedItems == null)
+            {
+                _cachedItems = new List<SerializedMessage>();
+            }
+
+            foreach (var item in _cachedItems)
+            {
+                if (string.Equals(item.ProtocolName, protocolName, StringComparison.Ordinal))
+                {
+                    // No need to add
+                    return;
+                }
+            }
+
+            _cachedItems.Add(new SerializedMessage(protocolName, serialized));
+        }
+    }
+
+    private bool TryGetCachedUnsynchronized(string protocolName, out ReadOnlyMemory<byte> result)
+    {
+        if (string.Equals(_cachedItem1.ProtocolName, protocolName, StringComparison.Ordinal))
+        {
+            result = _cachedItem1.Serialized;
+            return true;
+        }
+
+        if (string.Equals(_cachedItem2.ProtocolName, protocolName, StringComparison.Ordinal))
+        {
+            result = _cachedItem2.Serialized;
+            return true;
+        }
+
+        if (_cachedItems != null)
+        {
+            foreach (var serializedMessage in _cachedItems)
+            {
+                if (string.Equals(serializedMessage.ProtocolName, protocolName, StringComparison.Ordinal))
+                {
+                    result = serializedMessage.Serialized;
+                    return true;
+                }
+            }
+        }
+
+        result = default;
+        return false;
+    }
+}
+
+internal readonly struct SerializedMessage
+{
+    /// <summary>
+    /// Gets the protocol of the serialized message.
+    /// </summary>
+    public string ProtocolName { get; }
+
+    /// <summary>
+    /// Gets the serialized representation of the message.
+    /// </summary>
+    public ReadOnlyMemory<byte> Serialized { get; }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SerializedMessage"/> class.
+    /// </summary>
+    /// <param name="protocolName">The protocol of the serialized message.</param>
+    /// <param name="serialized">The serialized representation of the message.</param>
+    public SerializedMessage(string protocolName, ReadOnlyMemory<byte> serialized)
+    {
+        ProtocolName = protocolName;
+        Serialized = serialized;
+    }
+}