Merged PR 58375: [SignalR] Extra backpressure timeout

#### AI description  (iteration 1)
#### PR Classification
This PR is a feature enhancement that improves SignalR’s message buffering by adding an extra backpressure cancellation timeout.

#### PR Summary
The changes strengthen SignalR's handling of backpressure conditions by introducing a 5-second timeout for message buffering operations and updating the related tests to verify cancellation behavior.
- **`src/SignalR/server/SignalR/test/Internal/MessageBufferTests.cs`**: Added tests for both default cancellation and explicit cancellation during backpressure, and refined test naming and timeout assertions.
- **`src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs`**: Introduced a new backpressure timeout test scenario and updated the simulation of timeout conditions.
- **`src/SignalR/common/Shared/MessageBuffer.cs`**: Implemented the extra 5-second cancellation timeout using a TimeProvider and linked cancellation tokens to improve backpressure handling.
<!-- GitOpsUserAgent=GitOps.Apps.Server.pullrequestcopilot -->

Author: BrennanConroy

src/SignalR/common/Shared/MessageBuffer.cs

@@ -33,6 +33,7 @@ internal sealed class MessageBuffer : IDisposable
 
 #if NET8_0_OR_GREATER
     private readonly PeriodicTimer _timer;
+    private readonly TimeProvider _timeProvider;
 #else
     private readonly TimerAwaitable _timer = new(AckRate, AckRate);
 #endif
@@ -68,8 +69,8 @@ public MessageBuffer(ConnectionContext connection, IHubProtocol protocol, long b
     public MessageBuffer(ConnectionContext connection, IHubProtocol protocol, long bufferLimit, ILogger logger, TimeProvider timeProvider)
     {
 #if NET8_0_OR_GREATER
-        timeProvider ??= TimeProvider.System;
-        _timer = new(AckRate, timeProvider);
+        _timeProvider = timeProvider;
+        _timer = new(AckRate, _timeProvider);
 #endif
 
         _buffer = new LinkedBuffer();
@@ -132,14 +133,17 @@ public ValueTask<FlushResult> WriteAsync(HubMessage hubMessage, CancellationToke
 
     private async ValueTask<FlushResult> WriteAsyncCore(Type hubMessageType, ReadOnlyMemory<byte> messageBytes, CancellationToken cancellationToken)
     {
-        // If backpressure is being observed a cancelable token is needed to make sure we can break out of waiting when the connection is closed
-        Debug.Assert(cancellationToken.CanBeCanceled);
-
         // TODO: Add backpressure based on message count
         if (_bufferedByteCount > _bufferLimit)
         {
+#if NET
+            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5), _timeProvider);
+#else
+            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+#endif
+            using var _ = CancellationTokenUtils.CreateLinkedToken(cts.Token, cancellationToken, out var linkedToken);
             // primitive backpressure if buffer is full
-            while (await _waitForAck.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
+            while (await _waitForAck.Reader.WaitToReadAsync(linkedToken).ConfigureAwait(false))
             {
                 if (_waitForAck.Reader.TryRead(out var count) && count < _bufferLimit)
                 {

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

@@ -274,19 +274,14 @@ private ValueTask<FlushResult> WriteCore(HubMessage message, CancellationToken c
                 static async ValueTask<FlushResult> WriteAsync(MessageBuffer messageBuffer, HubConnectionContext hubConnectionContext,
                     HubMessage message, CancellationToken cancellationToken)
                 {
-                    CancellationTokenSource? cts = null;
                     var connectionToken = hubConnectionContext.ConnectionAborted;
                     if (message is CloseMessage)
                     {
                         // If it's a CloseMessage, we might already have triggered the ConnectionAborted token
-                        // We would like to successfully send the CloseMessage for graceful close which means we can't use the ConnectionAborted token,
-                        // but we need to make sure we don't get blocked by backpressure or anything, so we use a short-lived token.
-                        cts = new CancellationTokenSource(TimeSpan.FromSeconds(5), hubConnectionContext._timeProvider);
-                        connectionToken = cts.Token;
+                        // We would like to successfully send the CloseMessage for graceful close which means we can't use the ConnectionAborted token.
+                        connectionToken = CancellationToken.None;
                     }
 
-                    using var __ = cts;
-
                     // MessageBuffer can wait on things other than the PipeWriter (which is canceled by other means)
                     // So we need to make sure the cancellation token passed to it is also canceled when the connection is aborted
                     using var _ = CancellationTokenUtils.CreateLinkedToken(connectionToken, cancellationToken, out var linkedToken);

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

@@ -5272,11 +5272,13 @@ public enum CloseScenario
     {
         PingTimeout,
         Abort,
+        BackpressureTimeout,
     }
 
     [Theory]
     [InlineData(CloseScenario.PingTimeout)]
     [InlineData(CloseScenario.Abort)]
+    [InlineData(CloseScenario.BackpressureTimeout)]
     public async Task StatefulReconnectWithMessageBufferBackpressureIsCancelable(CloseScenario scenario)
     {
         using (StartVerifiableLog(write => write.EventId.Name == "FailedWritingMessage"))
@@ -5307,29 +5309,34 @@ public async Task StatefulReconnectWithMessageBufferBackpressureIsCancelable(Clo
 
             await client2.SendHubMessageAsync(new InvocationMessage(nameof(MethodHub.BroadcastMethod), [new string('a', 100)]));
 
+            Assert.IsType<InvocationMessage>(await client2.ReadAsync().DefaultTimeout());
+
+            await client2.SendHubMessageAsync(new InvocationMessage(nameof(MethodHub.BroadcastMethod), [new string('a', 100)]));
+
             switch (scenario)
             {
                 case CloseScenario.PingTimeout:
-                {
-                    // We go over the 100 ms timeout interval multiple times
-                    for (var i = 0; i < 3; i++)
                     {
-                        timeProvider.Advance(timeout + TimeSpan.FromMilliseconds(1));
-                        client1.TickHeartbeat();
+                        // We go over the 100 ms timeout interval multiple times
+                        for (var i = 0; i < 3; i++)
+                        {
+                            timeProvider.Advance(timeout + TimeSpan.FromMilliseconds(1));
+                            client1.TickHeartbeat();
+                        }
+                        break;
                     }
-                    break;
-                }
                 case CloseScenario.Abort:
-                {
-                    client1.Connection.Abort();
-                    break;
-                }
+                    {
+                        client1.Connection.Abort();
+                        break;
+                    }
+                case CloseScenario.BackpressureTimeout:
+                    {
+                        timeProvider.Advance(TimeSpan.FromSeconds(5) + TimeSpan.FromMilliseconds(1));
+                        break;
+                    }
             }
 
-            Assert.IsType<InvocationMessage>(await client2.ReadAsync().DefaultTimeout());
-
-            await client2.SendHubMessageAsync(new InvocationMessage(nameof(MethodHub.BroadcastMethod), [new string('a', 100)]));
-
             // This one might not be blocked on client1 if the server sends to client2 first during Broadcast
             Assert.IsType<InvocationMessage>(await client2.ReadAsync().DefaultTimeout());
 

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

@@ -563,7 +563,7 @@ public async Task SendingAckMessageDelayedDuringResend()
     }
 
     [Fact]
-    public async Task BackpressureWriteMessageCanBeCanceled()
+    public async Task PipeBackpressureWriteMessageCanBeCanceled()
     {
         var cts = new CancellationTokenSource();
         var protocol = new JsonHubProtocol();
@@ -581,7 +581,7 @@ public async Task BackpressureWriteMessageCanBeCanceled()
 
         cts.Cancel();
 
-        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await writeTask);
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await writeTask).DefaultTimeout();
 
         DuplexPipe.UpdateConnectionPair(ref pipes, connection, pipeOptions);
         var resendTask = messageBuffer.ResendAsync(pipes.Transport.Output);
@@ -607,6 +607,94 @@ public async Task BackpressureWriteMessageCanBeCanceled()
 
         await resendTask;
     }
+
+    [Fact]
+    public async Task BufferedBackpressureWriteMessageDefaultCancellation()
+    {
+        var cts = new CancellationTokenSource();
+        var protocol = new JsonHubProtocol();
+        var connection = new TestConnectionContext();
+        var pipeOptions = new PipeOptions(pauseWriterThreshold: 100, resumeWriterThreshold: 50);
+        var pipes = DuplexPipe.CreateConnectionPair(new PipeOptions(), pipeOptions);
+        connection.Transport = pipes.Transport;
+        var timeProvider = new FakeTimeProvider();
+        using var messageBuffer = new MessageBuffer(connection, protocol, bufferLimit: 50, NullLogger.Instance, timeProvider);
+
+        await messageBuffer.WriteAsync(new SerializedHubMessage(new InvocationMessage("t", new object[] { new byte[40] })), cts.Token);
+
+        // Write will hit pipe backpressure
+        var writeTask = messageBuffer.WriteAsync(new SerializedHubMessage(new InvocationMessage("t", new object[] { new byte[40] })), cts.Token);
+        Assert.False(writeTask.IsCompleted);
+
+        timeProvider.Advance(TimeSpan.FromSeconds(5) + TimeSpan.FromMilliseconds(1));
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await writeTask).DefaultTimeout();
+
+        DuplexPipe.UpdateConnectionPair(ref pipes, connection, pipeOptions);
+        var resendTask = messageBuffer.ResendAsync(pipes.Transport.Output);
+
+        var res = await pipes.Application.Input.ReadAsync();
+        var buffer = res.Buffer;
+        Assert.True(protocol.TryParseMessage(ref buffer, new TestBinder(), out var message));
+        Assert.IsType<SequenceMessage>(message);
+
+        pipes.Application.Input.AdvanceTo(buffer.Start);
+
+        res = await pipes.Application.Input.ReadAsync();
+        buffer = res.Buffer;
+        Assert.True(protocol.TryParseMessage(ref buffer, new TestBinder(), out message));
+        Assert.IsType<InvocationMessage>(message);
+
+        pipes.Application.Input.AdvanceTo(buffer.Start);
+
+        Assert.False(pipes.Application.Input.TryRead(out res));
+
+        await resendTask;
+    }
+
+    [Fact]
+    public async Task BufferedBackpressureWriteMessageCanBeCanceled()
+    {
+        var cts = new CancellationTokenSource();
+        var protocol = new JsonHubProtocol();
+        var connection = new TestConnectionContext();
+        var pipeOptions = new PipeOptions(pauseWriterThreshold: 100, resumeWriterThreshold: 50);
+        var pipes = DuplexPipe.CreateConnectionPair(new PipeOptions(), pipeOptions);
+        connection.Transport = pipes.Transport;
+        var timeProvider = new FakeTimeProvider();
+        using var messageBuffer = new MessageBuffer(connection, protocol, bufferLimit: 50, NullLogger.Instance, timeProvider);
+
+        await messageBuffer.WriteAsync(new SerializedHubMessage(new InvocationMessage("t", new object[] { new byte[40] })), cts.Token);
+
+        // Write will hit pipe backpressure
+        var writeTask = messageBuffer.WriteAsync(new SerializedHubMessage(new InvocationMessage("t", new object[] { new byte[40] })), cts.Token);
+        Assert.False(writeTask.IsCompleted);
+
+        cts.Cancel();
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await writeTask).DefaultTimeout();
+
+        DuplexPipe.UpdateConnectionPair(ref pipes, connection, pipeOptions);
+        var resendTask = messageBuffer.ResendAsync(pipes.Transport.Output);
+
+        var res = await pipes.Application.Input.ReadAsync();
+        var buffer = res.Buffer;
+        Assert.True(protocol.TryParseMessage(ref buffer, new TestBinder(), out var message));
+        Assert.IsType<SequenceMessage>(message);
+
+        pipes.Application.Input.AdvanceTo(buffer.Start);
+
+        res = await pipes.Application.Input.ReadAsync();
+        buffer = res.Buffer;
+        Assert.True(protocol.TryParseMessage(ref buffer, new TestBinder(), out message));
+        Assert.IsType<InvocationMessage>(message);
+
+        pipes.Application.Input.AdvanceTo(buffer.Start);
+
+        Assert.False(pipes.Application.Input.TryRead(out res));
+
+        await resendTask;
+    }
 }
 
 internal sealed class TestConnectionContext : ConnectionContext