Fix flaky HubConnectionHandler test

src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs

@@ -6,9 +6,11 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
+using System.IO.Pipelines;
 using System.Linq;
 using System.Security.Claims;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using MessagePack;
 using MessagePack.Formatters;
@@ -2797,6 +2799,78 @@ public async Task ReceivingMessagesPreventsConnectionTimeoutFromOccuring()
             }
         }
 
+        internal class PipeReaderWrapper : PipeReader
+        {
+            private readonly PipeReader _originalPipeReader;
+            private TaskCompletionSource<object> _waitForRead;
+            private object _lock = new object();
+
+            public PipeReaderWrapper(PipeReader pipeReader)
+            {
+                _originalPipeReader = pipeReader;
+                _waitForRead = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            }
+
+            public override void AdvanceTo(SequencePosition consumed) =>
+                _originalPipeReader.AdvanceTo(consumed);
+
+            public override void AdvanceTo(SequencePosition consumed, SequencePosition examined) =>
+                _originalPipeReader.AdvanceTo(consumed, examined);
+
+            public override void CancelPendingRead() =>
+                _originalPipeReader.CancelPendingRead();
+
+            public override void Complete(Exception exception = null) =>
+                _originalPipeReader.Complete(exception);
+
+            public override async ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
+            {
+                lock (_lock)
+                {
+                    _waitForRead.SetResult(null);
+                }
+
+                try
+                {
+                    return await _originalPipeReader.ReadAsync(cancellationToken);
+                }
+                finally
+                {
+                    lock (_lock)
+                    {
+                        _waitForRead = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    }
+                }
+            }
+
+            public override bool TryRead(out ReadResult result) =>
+                _originalPipeReader.TryRead(out result);
+
+            public Task WaitForReadStart()
+            {
+                lock (_lock)
+                {
+                    return _waitForRead.Task;
+                }
+            }
+        }
+
+        internal class CustomDuplex : IDuplexPipe
+        {
+            private readonly IDuplexPipe _originalDuplexPipe;
+            public readonly PipeReaderWrapper WrappedPipeReader;
+
+            public CustomDuplex(IDuplexPipe duplexPipe)
+            {
+                _originalDuplexPipe = duplexPipe;
+                WrappedPipeReader = new PipeReaderWrapper(_originalDuplexPipe.Input);
+            }
+
+            public PipeReader Input => WrappedPipeReader;
+
+            public PipeWriter Output => _originalDuplexPipe.Output;
+        }
+
         [Fact]
         public async Task HubMethodInvokeDoesNotCountTowardsClientTimeout()
         {
@@ -2813,6 +2887,9 @@ public async Task HubMethodInvokeDoesNotCountTowardsClientTimeout()
 
                 using (var client = new TestClient(new JsonHubProtocol()))
                 {
+                    var customDuplex = new CustomDuplex(client.Connection.Transport);
+                    client.Connection.Transport = customDuplex;
+
                     var connectionHandlerTask = await client.ConnectAsync(connectionHandler);
                     // This starts the timeout logic
                     await client.SendHubMessageAsync(PingMessage.Instance);
@@ -2829,6 +2906,11 @@ public async Task HubMethodInvokeDoesNotCountTowardsClientTimeout()
 
                     await hubMethodTask.OrTimeout();
 
+                    // There is a small window when the hub method finishes and the timer starts again
+                    // So we need to delay a little before ticking the heart beat.
+                    // We do this by waiting until we know the HubConnectionHandler code is in pipe.ReadAsync()
+                    await customDuplex.WrappedPipeReader.WaitForReadStart().OrTimeout();
+
                     // Tick heartbeat again now that we're outside of the hub method
                     client.TickHeartbeat();