Suggestions from code review

Author: MackinnonBuck

src/Components/Server/src/Circuits/RemoteJSDataStream.cs

@@ -181,18 +181,14 @@ public override void Write(byte[] buffer, int offset, int count)
 
     public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
     {
-        using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_streamCancellationToken, cancellationToken))
-        {
-            return await _pipeReaderStream.ReadAsync(buffer.AsMemory(offset, count), linkedCts.Token);
-        }
+        using var linkedCts = ValueLinkedCancellationTokenSource.Create(_streamCancellationToken, cancellationToken);
+        return await _pipeReaderStream.ReadAsync(buffer.AsMemory(offset, count), linkedCts.Token);
     }
 
     public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
     {
-        using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_streamCancellationToken, cancellationToken))
-        {
-            return await _pipeReaderStream.ReadAsync(buffer, linkedCts.Token);
-        }
+        using var linkedCts = ValueLinkedCancellationTokenSource.Create(_streamCancellationToken, cancellationToken);
+        return await _pipeReaderStream.ReadAsync(buffer, linkedCts.Token);
     }
 
     private async Task ThrowOnTimeout()
@@ -233,4 +229,45 @@ protected override void Dispose(bool disposing)
 
         _disposed = true;
     }
+
+    // A helper for creating and disposing linked CancellationTokenSources
+    // without allocating, when possible.
+    // Internal for testing.
+    internal readonly struct ValueLinkedCancellationTokenSource : IDisposable
+    {
+        private readonly CancellationTokenSource? _linkedCts;
+
+        public readonly CancellationToken Token;
+
+        // For testing.
+        internal bool HasLinkedCancellationTokenSource => _linkedCts is not null;
+
+        public static ValueLinkedCancellationTokenSource Create(
+            CancellationToken token1, CancellationToken token2)
+        {
+            if (!token1.CanBeCanceled)
+            {
+                return new(linkedCts: null, token2);
+            }
+
+            if (!token2.CanBeCanceled)
+            {
+                return new(linkedCts: null, token1);
+            }
+
+            var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(token1, token2);
+            return new(linkedCts, linkedCts.Token);
+        }
+
+        private ValueLinkedCancellationTokenSource(CancellationTokenSource? linkedCts, CancellationToken token)
+        {
+            _linkedCts = linkedCts;
+            Token = token;
+        }
+
+        public void Dispose()
+        {
+            _linkedCts?.Dispose();
+        }
+    }
 }

src/Components/Server/test/Circuits/RemoteJSDataStreamTest.cs

@@ -287,54 +287,59 @@ public async Task ReceiveData_ReceivesDataThenTimesout_StreamDisposed()
         Assert.False(success);
     }
 
-    [Fact]
-    public async Task ReadAsync_ByteArray_DisposesCancellationTokenSource()
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public void ValueLinkedCts_Works_WhenOneTokenCannotBeCanceled(bool isToken1Cancelable)
     {
-        // Arrange
-        var jsStreamReference = Mock.Of<IJSStreamReference>();
-        var remoteJSDataStream = await RemoteJSDataStream.CreateRemoteJSDataStreamAsync(_jsRuntime, jsStreamReference, totalLength: 100, signalRMaximumIncomingBytes: 10_000, jsInteropDefaultCallTimeout: TimeSpan.FromMinutes(1), cancellationToken: CancellationToken.None).DefaultTimeout();
-        var buffer = new byte[100];
-        var chunk = new byte[100];
-        new Random().NextBytes(chunk);
-
-        // Act
         var cts = new CancellationTokenSource();
-        var sendDataTask = Task.Run(async () =>
-        {
-            await RemoteJSDataStream.ReceiveData(_jsRuntime, GetStreamId(remoteJSDataStream, _jsRuntime), chunkId: 0, chunk, error: null);
-        });
+        var token1 = isToken1Cancelable ? cts.Token : CancellationToken.None;
+        var token2 = isToken1Cancelable ? CancellationToken.None : cts.Token;
 
-        await sendDataTask;
-        var result = await remoteJSDataStream.ReadAsync(buffer, 0, buffer.Length, cts.Token);
+        using var linkedCts = RemoteJSDataStream.ValueLinkedCancellationTokenSource.Create(token1, token2);
 
-        // Assert
-        Assert.True(cts.Token.CanBeCanceled);
-        Assert.False(cts.IsCancellationRequested);
+        Assert.False(linkedCts.HasLinkedCancellationTokenSource);
+        Assert.False(linkedCts.Token.IsCancellationRequested);
+
+        cts.Cancel();
+
+        Assert.True(linkedCts.Token.IsCancellationRequested);
     }
 
     [Fact]
-    public async Task ReadAsync_Memory_DisposesCancellationTokenSource()
+    public void ValueLinkedCts_Works_WhenBothTokensCannotBeCanceled()
     {
-        // Arrange
-        var jsStreamReference = Mock.Of<IJSStreamReference>();
-        var remoteJSDataStream = await RemoteJSDataStream.CreateRemoteJSDataStreamAsync(_jsRuntime, jsStreamReference, totalLength: 100, signalRMaximumIncomingBytes: 10_000, jsInteropDefaultCallTimeout: TimeSpan.FromMinutes(1), cancellationToken: CancellationToken.None).DefaultTimeout();
-        var buffer = new Memory<byte>(new byte[100]);
-        var chunk = new byte[100];
-        new Random().NextBytes(chunk);
+        using var linkedCts = RemoteJSDataStream.ValueLinkedCancellationTokenSource.Create(
+            CancellationToken.None,
+            CancellationToken.None);
 
-        // Act
-        var cts = new CancellationTokenSource();
-        var sendDataTask = Task.Run(async () =>
-        {
-            await RemoteJSDataStream.ReceiveData(_jsRuntime, GetStreamId(remoteJSDataStream, _jsRuntime), chunkId: 0, chunk, error: null);
-        });
+        Assert.False(linkedCts.HasLinkedCancellationTokenSource);
+        Assert.False(linkedCts.Token.IsCancellationRequested);
+    }
 
-        await sendDataTask;
-        var result = await remoteJSDataStream.ReadAsync(buffer, cts.Token);
+    [Theory]
+    [InlineData(false, true)]
+    [InlineData(true, false)]
+    [InlineData(true, true)]
+    public void ValueLinkedCts_Works_WhenBothTokensCanBeCanceled(bool shouldCancelToken1, bool shouldCancelToken2)
+    {
+        var cts1 = new CancellationTokenSource();
+        var cts2 = new CancellationTokenSource();
+        using var linkedCts = RemoteJSDataStream.ValueLinkedCancellationTokenSource.Create(cts1.Token, cts2.Token);
 
-        // Assert
-        Assert.True(cts.Token.CanBeCanceled);
-        Assert.False(cts.IsCancellationRequested);
+        Assert.True(linkedCts.HasLinkedCancellationTokenSource);
+        Assert.False(linkedCts.Token.IsCancellationRequested);
+
+        if (shouldCancelToken1)
+        {
+            cts1.Cancel();
+        }
+        if (shouldCancelToken2)
+        {
+            cts2.Cancel();
+        }
+
+        Assert.True(linkedCts.Token.IsCancellationRequested);
     }
 
     private static async Task<RemoteJSDataStream> CreateRemoteJSDataStreamAsync(TestRemoteJSRuntime jsRuntime = null)