Allow CancellationToken in streaming hub methods

test/Microsoft.AspNetCore.SignalR.Tests.Utils/CancellationTokenExtensions.cs

@@ -0,0 +1,21 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.SignalR.Tests
+{
+    public static class CancellationTokenExtensions
+    {
+        public static Task WaitForCancellationAsync(this CancellationToken token)
+        {
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            token.Register((t) =>
+            {
+                ((TaskCompletionSource<object>)t).SetResult(null);
+            }, tcs);
+            return tcs.Task;
+        }
+    }
+}

test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs

@@ -556,10 +556,10 @@ public ChannelReader<int> CancelableStream(CancellationToken token)
         {
             var channel = Channel.CreateBounded<int>(10);
 
-            Task.Run(() =>
+            Task.Run(async () =>
             {
                 _tcsService.StartedMethod.SetResult(null);
-                token.WaitHandle.WaitOne();
+                await token.WaitForCancellationAsync();
                 channel.Writer.TryComplete();
                 _tcsService.EndMethod.SetResult(null);
             });
@@ -571,10 +571,10 @@ public ChannelReader<int> CancelableStream2(int ignore, int ignore2, Cancellatio
         {
             var channel = Channel.CreateBounded<int>(10);
 
-            Task.Run(() =>
+            Task.Run(async () =>
             {
                 _tcsService.StartedMethod.SetResult(null);
-                token.WaitHandle.WaitOne();
+                await token.WaitForCancellationAsync();
                 channel.Writer.TryComplete();
                 _tcsService.EndMethod.SetResult(null);
             });
@@ -586,10 +586,10 @@ public ChannelReader<int> CancelableStreamMiddle(int ignore, CancellationToken t
         {
             var channel = Channel.CreateBounded<int>(10);
 
-            Task.Run(() =>
+            Task.Run(async () =>
             {
                 _tcsService.StartedMethod.SetResult(null);
-                token.WaitHandle.WaitOne();
+                await token.WaitForCancellationAsync();
                 channel.Writer.TryComplete();
                 _tcsService.EndMethod.SetResult(null);
             });