Preserve OCE.CancellationToken in ReadAtLeastAsync of StreamPipeReader (#83926)

Author: manandre

src/libraries/System.IO.Pipelines/src/System/IO/Pipelines/StreamPipeReader.cs

@@ -200,94 +200,15 @@ private bool CompleteAndGetNeedsDispose()
         /// <inheritdoc />
         public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
         {
-            // TODO ReadyAsync needs to throw if there are overlapping reads.
-            ThrowIfCompleted();
-
-            if (cancellationToken.IsCancellationRequested)
-            {
-                return new ValueTask<ReadResult>(Task.FromCanceled<ReadResult>(cancellationToken));
-            }
-
-            // PERF: store InternalTokenSource locally to avoid querying it twice (which acquires a lock)
-            CancellationTokenSource tokenSource = InternalTokenSource;
-            if (TryReadInternal(tokenSource, out ReadResult readResult))
-            {
-                return new ValueTask<ReadResult>(readResult);
-            }
-
-            if (_isStreamCompleted)
-            {
-                ReadResult completedResult = new ReadResult(buffer: default, isCanceled: false, isCompleted: true);
-                return new ValueTask<ReadResult>(completedResult);
-            }
-
-            return Core(this, tokenSource, cancellationToken);
-
-#if NETCOREAPP
-            [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
-#endif
-            static async ValueTask<ReadResult> Core(StreamPipeReader reader, CancellationTokenSource tokenSource, CancellationToken cancellationToken)
-            {
-                CancellationTokenRegistration reg = default;
-                if (cancellationToken.CanBeCanceled)
-                {
-                    reg = cancellationToken.UnsafeRegister(state => ((StreamPipeReader)state!).Cancel(), reader);
-                }
-
-                using (reg)
-                {
-                    var isCanceled = false;
-                    try
-                    {
-                        // This optimization only makes sense if we don't have anything buffered
-                        if (reader.UseZeroByteReads && reader._bufferedBytes == 0)
-                        {
-                            // Wait for data by doing 0 byte read before
-                            await reader.InnerStream.ReadAsync(Memory<byte>.Empty, tokenSource.Token).ConfigureAwait(false);
-                        }
-
-                        reader.AllocateReadTail();
-
-                        Memory<byte> buffer = reader._readTail!.AvailableMemory.Slice(reader._readTail.End);
-
-                        int length = await reader.InnerStream.ReadAsync(buffer, tokenSource.Token).ConfigureAwait(false);
-
-                        Debug.Assert(length + reader._readTail.End <= reader._readTail.AvailableMemory.Length);
-
-                        reader._readTail.End += length;
-                        reader._bufferedBytes += length;
-
-                        if (length == 0)
-                        {
-                            reader._isStreamCompleted = true;
-                        }
-                    }
-                    catch (OperationCanceledException ex)
-                    {
-                        reader.ClearCancellationToken();
-
-                        if (cancellationToken.IsCancellationRequested)
-                        {
-                            // Simulate an OCE triggered directly by the cancellationToken rather than the InternalTokenSource
-                            throw new OperationCanceledException(ex.Message, ex, cancellationToken);
-                        }
-                        else if (tokenSource.IsCancellationRequested)
-                        {
-                            // Catch cancellation and translate it into setting isCanceled = true
-                            isCanceled = true;
-                        }
-                        else
-                        {
-                            throw;
-                        }
-                    }
-
-                    return new ReadResult(reader.GetCurrentReadOnlySequence(), isCanceled, reader._isStreamCompleted);
-                }
-            }
+            return ReadInternalAsync(null, cancellationToken);
         }
 
         protected override ValueTask<ReadResult> ReadAtLeastAsyncCore(int minimumSize, CancellationToken cancellationToken)
+        {
+            return ReadInternalAsync(minimumSize, cancellationToken);
+        }
+
+        private ValueTask<ReadResult> ReadInternalAsync(int? minimumSize, CancellationToken cancellationToken)
         {
             // TODO ReadyAsync needs to throw if there are overlapping reads.
             ThrowIfCompleted();
@@ -301,7 +222,10 @@ protected override ValueTask<ReadResult> ReadAtLeastAsyncCore(int minimumSize, C
             CancellationTokenSource tokenSource = InternalTokenSource;
             if (TryReadInternal(tokenSource, out ReadResult readResult))
             {
-                if (readResult.Buffer.Length >= minimumSize || readResult.IsCompleted || readResult.IsCanceled)
+                if (minimumSize is null
+                    || readResult.Buffer.Length >= minimumSize
+                    || readResult.IsCompleted
+                    || readResult.IsCanceled)
                 {
                     return new ValueTask<ReadResult>(readResult);
                 }
@@ -318,7 +242,7 @@ protected override ValueTask<ReadResult> ReadAtLeastAsyncCore(int minimumSize, C
 #if NETCOREAPP
             [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
 #endif
-            static async ValueTask<ReadResult> Core(StreamPipeReader reader, int minimumSize, CancellationTokenSource tokenSource, CancellationToken cancellationToken)
+            static async ValueTask<ReadResult> Core(StreamPipeReader reader, int? minimumSize, CancellationTokenSource tokenSource, CancellationToken cancellationToken)
             {
                 CancellationTokenRegistration reg = default;
                 if (cancellationToken.CanBeCanceled)
@@ -356,13 +280,18 @@ static async ValueTask<ReadResult> Core(StreamPipeReader reader, int minimumSize
                                 reader._isStreamCompleted = true;
                                 break;
                             }
-                        } while (reader._bufferedBytes < minimumSize);
+                        } while (minimumSize != null && reader._bufferedBytes < minimumSize);
                     }
-                    catch (OperationCanceledException)
+                    catch (OperationCanceledException ex)
                     {
                         reader.ClearCancellationToken();
 
-                        if (tokenSource.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+                        if (cancellationToken.IsCancellationRequested)
+                        {
+                            // Simulate an OCE triggered directly by the cancellationToken rather than the InternalTokenSource
+                            throw new OperationCanceledException(ex.Message, ex, cancellationToken);
+                        }
+                        else if (tokenSource.IsCancellationRequested)
                         {
                             // Catch cancellation and translate it into setting isCanceled = true
                             isCanceled = true;
@@ -371,7 +300,6 @@ static async ValueTask<ReadResult> Core(StreamPipeReader reader, int minimumSize
                         {
                             throw;
                         }
-
                     }
 
                     return new ReadResult(reader.GetCurrentReadOnlySequence(), isCanceled, reader._isStreamCompleted);

src/libraries/System.IO.Pipelines/tests/PipeReaderReadAtLeastAsyncTests.cs

@@ -139,10 +139,13 @@ public async Task CanReadAtLeast(int bufferSize, bool bufferedRead)
         }
 
         [Fact]
-        public Task ReadAtLeastAsyncThrowsIfPassedCanceledCancellationToken()
+        public async Task ReadAtLeastAsyncThrowsIfPassedCanceledCancellationToken()
         {
-            ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(0, new CancellationToken(canceled: true));
-            return Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
+            CancellationToken token = new CancellationToken(canceled: true);
+            ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(0, token);
+            TaskCanceledException tce = await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
+            Assert.Equal(token, tce.CancellationToken);
+            Assert.Null(tce.InnerException);
         }
 
         [Fact]
@@ -164,12 +167,13 @@ public async Task WriteAndCancellingPendingReadBeforeReadAtLeastAsync()
         }
 
         [Fact]
-        public Task ReadAtLeastAsyncCancelableWhenWaitingForMoreData()
+        public async Task ReadAtLeastAsyncCancelableWhenWaitingForMoreData()
         {
             CancellationTokenSource cts = new CancellationTokenSource();
             ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(1, cts.Token);
             cts.Cancel();
-            return Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+            var oce = await Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+            Assert.Equal(cts.Token, oce.CancellationToken);
         }
 
         [Fact]
@@ -179,7 +183,8 @@ public async Task ReadAtLeastAsyncCancelableAfterReadingSome()
             await Pipe.WriteAsync(new byte[10], default);
             ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(11, cts.Token);
             cts.Cancel();
-            await Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+            var oce = await Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+            Assert.Equal(cts.Token, oce.CancellationToken);
         }
 
         [Fact]
@@ -191,7 +196,8 @@ public async Task ReadAtLeastAsyncCancelableAfterReadingSomeAndWritingAfterStart
             // Write, but not enough to unblock ReadAtLeastAsync
             await Pipe.WriteAsync(new byte[1], default);
             cts.Cancel();
-            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await task);
+            var oce = await Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+            Assert.Equal(cts.Token, oce.CancellationToken);
         }
     }
 }