Fix dropping messages on reconnect (#50409)

* Fix dropping messages on reconnect

* fb

* Update src/SignalR/common/Shared/MessageBuffer.cs

Co-authored-by: Stephen Halter <[email protected]>

---------

Co-authored-by: Stephen Halter <[email protected]>

Author: BrennanConroy

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

@@ -1137,7 +1137,6 @@ private async Task SendWithLock(ConnectionState expectedConnectionState, HubMess
                 break;
             case SequenceMessage sequenceMessage:
                 Log.ReceivedSequenceMessage(_logger, sequenceMessage.SequenceId);
-                connectionState.ResetSequence(sequenceMessage);
                 break;
             default:
                 throw new InvalidOperationException($"Unexpected message type: {message.GetType().FullName}");
@@ -2100,14 +2099,6 @@ public Task AckAsync(AckMessage ackMessage)
             return Task.CompletedTask;
         }
 
-        public void ResetSequence(SequenceMessage sequenceMessage)
-        {
-            if (UsingAcks())
-            {
-                _messageBuffer.ResetSequence(sequenceMessage);
-            }
-        }
-
         [MemberNotNullWhen(true, nameof(_messageBuffer))]
         public bool UsingAcks() => _messageBuffer is not null;
 

src/SignalR/common/Shared/MessageBuffer.cs

@@ -41,7 +41,6 @@ internal sealed class MessageBuffer : IDisposable
     private PipeWriter _writer;
 
     private long _totalMessageCount;
-    private bool _waitForSequenceMessage;
 
     // Message IDs start at 1 and always increment by 1
     private long _currentReceivingSequenceId = 1;
@@ -210,19 +209,20 @@ public async Task AckAsync(AckMessage ackMessage)
 
     internal bool ShouldProcessMessage(HubMessage message)
     {
-        // TODO: if we're expecting a sequence message but get here should we error or ignore or maybe even continue to process them?
-        if (_waitForSequenceMessage)
+        if (message is SequenceMessage sequenceMessage)
         {
-            if (message is SequenceMessage)
-            {
-                _waitForSequenceMessage = false;
-                return true;
-            }
-            else
+            // TODO: is a sequence message expected right now?
+
+            if (sequenceMessage.SequenceId > _currentReceivingSequenceId)
             {
-                // ignore messages received while waiting for sequence message
-                return false;
+                throw new InvalidOperationException("Sequence ID greater than amount of messages we've received.");
             }
+
+            _currentReceivingSequenceId = sequenceMessage.SequenceId;
+
+            // Technically handled by the 'is not HubInvocationMessage' check, but this is future proofing in case that check changes
+            // SequenceMessage should not be counted towards ackable messages
+            return true;
         }
 
         // Only care about messages implementing HubInvocationMessage currently (e.g. ignore ping, close, ack, sequence)
@@ -233,6 +233,7 @@ internal bool ShouldProcessMessage(HubMessage message)
         }
 
         var currentId = _currentReceivingSequenceId;
+        // ShouldProcessMessage is never called in parallel and is the only method referencing _currentReceivingSequenceId
         _currentReceivingSequenceId++;
         if (currentId <= _latestReceivedSequenceId)
         {
@@ -244,21 +245,8 @@ internal bool ShouldProcessMessage(HubMessage message)
         return true;
     }
 
-    internal void ResetSequence(SequenceMessage sequenceMessage)
-    {
-        // TODO: is a sequence message expected right now?
-
-        if (sequenceMessage.SequenceId > _currentReceivingSequenceId)
-        {
-            throw new InvalidOperationException("Sequence ID greater than amount of messages we've received.");
-        }
-        _currentReceivingSequenceId = sequenceMessage.SequenceId;
-    }
-
     internal async Task ResendAsync(PipeWriter writer)
     {
-        _waitForSequenceMessage = true;
-
         var tcs = new TaskCompletionSource<FlushResult>(TaskCreationOptions.RunContinuationsAsynchronously);
         _resend = tcs;
 

src/SignalR/server/Core/src/HubConnectionContext.cs

@@ -802,12 +802,4 @@ internal bool ShouldProcessMessage(HubMessage message)
         }
         return true;
     }
-
-    internal void ResetSequence(SequenceMessage sequenceMessage)
-    {
-        if (UsingStatefulReconnect())
-        {
-            _messageBuffer.ResetSequence(sequenceMessage);
-        }
-    }
 }

src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs

@@ -198,7 +198,6 @@ public override Task DispatchMessageAsync(HubConnectionContext connection, HubMe
 
             case SequenceMessage sequenceMessage:
                 Log.ReceivedSequenceMessage(_logger, sequenceMessage.SequenceId);
-                connection.ResetSequence(sequenceMessage);
                 break;
 
             case CloseMessage closeMessage:

src/SignalR/server/SignalR/test/Internal/MessageBufferTests.cs

@@ -142,9 +142,8 @@ public async Task UnAckedMessageResentOnReconnect()
         DuplexPipe.UpdateConnectionPair(ref pipes, connection);
         await messageBuffer.ResendAsync(pipes.Transport.Output);
 
-        // Any message except SequenceMessage will be ignored until a SequenceMessage is received
-        Assert.False(messageBuffer.ShouldProcessMessage(PingMessage.Instance));
-        Assert.False(messageBuffer.ShouldProcessMessage(CompletionMessage.WithResult("1", null)));
+        Assert.True(messageBuffer.ShouldProcessMessage(PingMessage.Instance));
+        Assert.True(messageBuffer.ShouldProcessMessage(CompletionMessage.WithResult("1", null)));
         Assert.True(messageBuffer.ShouldProcessMessage(new SequenceMessage(1)));
 
         res = await pipes.Application.Input.ReadAsync();
@@ -164,10 +163,10 @@ public async Task UnAckedMessageResentOnReconnect()
 
         pipes.Application.Input.AdvanceTo(buffer.Start);
 
-        messageBuffer.ResetSequence(new SequenceMessage(1));
+        messageBuffer.ShouldProcessMessage(new SequenceMessage(1));
 
         Assert.True(messageBuffer.ShouldProcessMessage(PingMessage.Instance));
-        Assert.True(messageBuffer.ShouldProcessMessage(CompletionMessage.WithResult("1", null)));
+        Assert.False(messageBuffer.ShouldProcessMessage(CompletionMessage.WithResult("1", null)));
     }
 
     [Fact]
@@ -235,7 +234,7 @@ public async Task ReceiveSequenceMessageWithLargerIDThanMessagesReceived()
 
         pipes.Application.Input.AdvanceTo(buffer.Start);
 
-        Assert.Throws<InvalidOperationException>(() => messageBuffer.ResetSequence(new SequenceMessage(2)));
+        Assert.Throws<InvalidOperationException>(() => messageBuffer.ShouldProcessMessage(new SequenceMessage(2)));
     }
 
     [Fact]