Merging internal commits for release/9.0

NuGet.config

@@ -6,8 +6,10 @@
     <!--  Begin: Package sources from dotnet-extensions -->
     <!--  End: Package sources from dotnet-extensions -->
     <!--  Begin: Package sources from dotnet-runtime -->
+    <add key="darc-int-dotnet-runtime-19c0782" value="https://pkgs.dev.azure.com/dnceng/internal/_packaging/darc-int-dotnet-runtime-19c07820/nuget/v3/index.json" />
     <!--  End: Package sources from dotnet-runtime -->
     <!--  Begin: Package sources from dotnet-efcore -->
+    <add key="darc-int-dotnet-efcore-1bea6ab" value="https://pkgs.dev.azure.com/dnceng/internal/_packaging/darc-int-dotnet-efcore-1bea6ab6/nuget/v3/index.json" />
     <!--  End: Package sources from dotnet-efcore -->
     <!--End: Package sources managed by Dependency Flow automation. Do not edit the sources above.-->
     <add key="dotnet-eng" value="https://pkgs.dev.azure.com/dnceng/public/_packaging/dotnet-eng/nuget/v3/index.json" />
@@ -30,8 +32,10 @@
     <clear />
     <!--Begin: Package sources managed by Dependency Flow automation. Do not edit the sources below.-->
     <!--  Begin: Package sources from dotnet-efcore -->
+    <add key="darc-int-dotnet-efcore-1bea6ab" value="true" />
     <!--  End: Package sources from dotnet-efcore -->
     <!--  Begin: Package sources from dotnet-runtime -->
+    <add key="darc-int-dotnet-runtime-19c0782" value="true" />
     <!--  End: Package sources from dotnet-runtime -->
     <!--End: Package sources managed by Dependency Flow automation. Do not edit the sources above.-->
   </disabledPackageSources>

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

@@ -2288,10 +2288,24 @@ public async Task TimerLoop(TimerAwaitable timer)
             }
         }
 
-        public ValueTask<FlushResult> WriteAsync(HubMessage message, CancellationToken cancellationToken)
+        public async ValueTask<FlushResult> WriteAsync(HubMessage message, CancellationToken cancellationToken)
         {
             Debug.Assert(_messageBuffer is not null);
-            return _messageBuffer.WriteAsync(message, cancellationToken);
+            CancellationTokenSource? cts = null;
+            var connectionToken = _hubConnection._state.StopCts.Token;
+            // StopAsync might have been called which would trigger the StopCts and thus prevent the CloseMessage from being sent.
+            // We'll use a short-lived token for CloseMessage specifically that isn't tied to the StopCts so we have a good chance of
+            // sending the CloseMessage to the server.
+            if (message is CloseMessage)
+            {
+                cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+                connectionToken = cts.Token;
+            }
+
+            using var __ = cts;
+
+            using var _ = CancellationTokenUtils.CreateLinkedToken(cancellationToken, connectionToken, out var linkedToken);
+            return await _messageBuffer.WriteAsync(message, linkedToken).ConfigureAwait(false);
         }
 
         public bool ShouldProcessMessage(HubMessage message)

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

@@ -941,6 +941,8 @@ public async Task DisableReconnectCalledWhenSendingCloseMessage()
         Assert.Null(exception);
 
         await reconnectFeature.DisableReconnectCalled.DefaultTimeout();
+
+        Assert.Equal("{\"type\":7}", await innerConnection.ReadSentTextMessageAsync().DefaultTimeout());
     }
 
     private class SampleObject

src/SignalR/common/Shared/MessageBuffer.cs

@@ -132,6 +132,9 @@ 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)
         {

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

@@ -205,7 +205,7 @@ internal ValueTask WriteAsync(HubMessage message, bool ignoreAbort, Cancellation
         // The write didn't complete synchronously so await completion
         if (!task.IsCompletedSuccessfully)
         {
-            return new ValueTask(CompleteWriteAsync(task));
+            return new ValueTask(CompleteWriteAsync(task, cancellationToken));
         }
         else
         {
@@ -250,7 +250,7 @@ public virtual ValueTask WriteAsync(SerializedHubMessage message, CancellationTo
         // The write didn't complete synchronously so await completion
         if (!task.IsCompletedSuccessfully)
         {
-            return new ValueTask(CompleteWriteAsync(task));
+            return new ValueTask(CompleteWriteAsync(task, cancellationToken));
         }
         else
         {
@@ -271,7 +271,30 @@ private ValueTask<FlushResult> WriteCore(HubMessage message, CancellationToken c
         {
             if (UsingStatefulReconnect())
             {
-                return _messageBuffer.WriteAsync(message, cancellationToken);
+                return WriteAsync(_messageBuffer, this, message, cancellationToken);
+
+                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;
+                    }
+
+                    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);
+                    var result = await messageBuffer.WriteAsync(message, linkedToken);
+                    return result;
+                }
             }
             else
             {
@@ -300,7 +323,17 @@ private ValueTask<FlushResult> WriteCore(SerializedHubMessage message, Cancellat
             if (UsingStatefulReconnect())
             {
                 Debug.Assert(_messageBuffer is not null);
-                return _messageBuffer.WriteAsync(message, cancellationToken);
+                return WriteAsync(_messageBuffer, this, message, cancellationToken);
+
+                static async ValueTask<FlushResult> WriteAsync(MessageBuffer messageBuffer, HubConnectionContext hubConnectionContext,
+                    SerializedHubMessage message, CancellationToken cancellationToken)
+                {
+                    // 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(hubConnectionContext.ConnectionAborted, cancellationToken, out var linkedToken);
+                    var result = await messageBuffer.WriteAsync(message, linkedToken);
+                    return result;
+                }
             }
             else
             {
@@ -321,13 +354,16 @@ private ValueTask<FlushResult> WriteCore(SerializedHubMessage message, Cancellat
         }
     }
 
-    private async Task CompleteWriteAsync(ValueTask<FlushResult> task)
+    private async Task CompleteWriteAsync(ValueTask<FlushResult> task, CancellationToken cancellationToken)
     {
         try
         {
             await task;
         }
-        catch (Exception ex)
+        // We care about errors while serializing to the PipeWriter as that will leave the Pipe
+        // in an invalid (for our scenario) state. OCE shouldn't occur while serializing bytes and
+        // writing to the Pipe. We assume that PipeWriter.WriteAsync(buffer) always writes the full message before calling FlushAsync
+        catch (Exception ex) when (ex is not OperationCanceledException || !cancellationToken.IsCancellationRequested)
         {
             CloseException = ex;
             Log.FailedWritingMessage(_logger, ex);
@@ -355,7 +391,10 @@ private async Task WriteSlowAsync(HubMessage message, bool ignoreAbort, Cancella
 
             await WriteCore(message, cancellationToken);
         }
-        catch (Exception ex)
+        // We care about errors while serializing to the PipeWriter as that will leave the Pipe
+        // in an invalid (for our scenario) state. OCE shouldn't occur while serializing bytes and
+        // writing to the Pipe. We assume that PipeWriter.WriteAsync(buffer) always writes the full message before calling FlushAsync
+        catch (Exception ex) when (ex is not OperationCanceledException || !cancellationToken.IsCancellationRequested)
         {
             CloseException = ex;
             Log.FailedWritingMessage(_logger, ex);
@@ -381,7 +420,10 @@ private async Task WriteSlowAsync(SerializedHubMessage message, CancellationToke
 
             await WriteCore(message, cancellationToken);
         }
-        catch (Exception ex)
+        // We care about errors while serializing to the PipeWriter as that will leave the Pipe
+        // in an invalid (for our scenario) state. OCE shouldn't occur while serializing bytes and
+        // writing to the Pipe. We assume that PipeWriter.WriteAsync(buffer) always writes the full message before calling FlushAsync
+        catch (Exception ex) when (ex is not OperationCanceledException || !cancellationToken.IsCancellationRequested)
         {
             CloseException = ex;
             Log.FailedWritingMessage(_logger, ex);

src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/DefaultHubLifetimeManagerTests.cs

@@ -35,18 +35,14 @@ public async Task SendAllAsyncWillCancelWithToken()
             var sendTask = manager.SendAllAsync("Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
             var message = Assert.IsType<InvocationMessage>(client1.TryRead());
             Assert.Equal("Hello", message.Target);
             Assert.Single(message.Arguments);
             Assert.Equal("World", (string)message.Arguments[0]);
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection2.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+
             Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
+            Assert.False(connection2.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -65,14 +61,10 @@ public async Task SendAllExceptAsyncWillCancelWithToken()
             var sendTask = manager.SendAllExceptAsync("Hello", new object[] { "World" }, new List<string> { connection1.ConnectionId }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection2.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
             Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
+            Assert.False(connection2.ConnectionAborted.IsCancellationRequested);
             Assert.Null(client1.TryRead());
         }
     }
@@ -89,13 +81,9 @@ public async Task SendConnectionAsyncWillCancelWithToken()
             var sendTask = manager.SendConnectionAsync(connection1.ConnectionId, "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection1.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
+            Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -111,13 +99,9 @@ public async Task SendConnectionsAsyncWillCancelWithToken()
             var sendTask = manager.SendConnectionsAsync(new List<string> { connection1.ConnectionId }, "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection1.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
+            Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -134,13 +118,9 @@ public async Task SendGroupAsyncWillCancelWithToken()
             var sendTask = manager.SendGroupAsync("group", "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection1.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
+            Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -161,14 +141,10 @@ public async Task SendGroupExceptAsyncWillCancelWithToken()
             var sendTask = manager.SendGroupExceptAsync("group", "Hello", new object[] { "World" }, new List<string> { connection1.ConnectionId }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection2.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
             Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
+            Assert.False(connection2.ConnectionAborted.IsCancellationRequested);
             Assert.Null(client1.TryRead());
         }
     }
@@ -186,13 +162,9 @@ public async Task SendGroupsAsyncWillCancelWithToken()
             var sendTask = manager.SendGroupsAsync(new List<string> { "group" }, "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection1.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
+
+            Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -211,18 +183,14 @@ public async Task SendUserAsyncWillCancelWithToken()
             var sendTask = manager.SendUserAsync("user", "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
             var message = Assert.IsType<InvocationMessage>(client1.TryRead());
             Assert.Equal("Hello", message.Target);
             Assert.Single(message.Arguments);
             Assert.Equal("World", (string)message.Arguments[0]);
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection2.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+
             Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
+            Assert.False(connection2.ConnectionAborted.IsCancellationRequested);
         }
     }
 
@@ -241,18 +209,14 @@ public async Task SendUsersAsyncWillCancelWithToken()
             var sendTask = manager.SendUsersAsync(new List<string> { "user1", "user2" }, "Hello", new object[] { "World" }, cts.Token).DefaultTimeout();
             Assert.False(sendTask.IsCompleted);
             cts.Cancel();
-            await sendTask.DefaultTimeout();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await sendTask.DefaultTimeout());
             var message = Assert.IsType<InvocationMessage>(client1.TryRead());
             Assert.Equal("Hello", message.Target);
             Assert.Single(message.Arguments);
             Assert.Equal("World", (string)message.Arguments[0]);
-            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            connection2.ConnectionAborted.Register(t =>
-            {
-                ((TaskCompletionSource)t).SetResult();
-            }, tcs);
-            await tcs.Task.DefaultTimeout();
+
             Assert.False(connection1.ConnectionAborted.IsCancellationRequested);
+            Assert.False(connection2.ConnectionAborted.IsCancellationRequested);
         }
     }
 }