Fix cancel in PipeReader.ReadAtLeastAsync (#66752)

Author: BrennanConroy

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

@@ -347,15 +347,15 @@ internal ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken)
             ValueTask<FlushResult> result;
             lock (SyncObj)
             {
-                PrepareFlush(out completionData, out result, cancellationToken);
+                PrepareFlushUnsynchronized(out completionData, out result, cancellationToken);
             }
 
             TrySchedule(ReaderScheduler, completionData);
 
             return result;
         }
 
-        private void PrepareFlush(out CompletionData completionData, out ValueTask<FlushResult> result, CancellationToken cancellationToken)
+        private void PrepareFlushUnsynchronized(out CompletionData completionData, out ValueTask<FlushResult> result, CancellationToken cancellationToken)
         {
             var completeReader = CommitUnsynchronized();
 
@@ -691,6 +691,9 @@ internal ValueTask<ReadResult> ReadAtLeastAsync(int minimumBytes, CancellationTo
 
                     // We also need to flip the reading state off
                     _operationState.EndRead();
+
+                    // Begin read again to wire up cancellation token
+                    _readerAwaitable.BeginOperation(token, s_signalReaderAwaitable, this);
                 }
 
                 // If the writer is currently paused and we are about the wait for more data then this would deadlock.
@@ -1057,7 +1060,7 @@ internal ValueTask<FlushResult> WriteAsync(ReadOnlyMemory<byte> source, Cancella
                     WriteMultiSegment(source.Span);
                 }
 
-                PrepareFlush(out completionData, out result, cancellationToken);
+                PrepareFlushUnsynchronized(out completionData, out result, cancellationToken);
             }
 
             TrySchedule(ReaderScheduler, completionData);

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

@@ -162,5 +162,36 @@ public async Task WriteAndCancellingPendingReadBeforeReadAtLeastAsync()
             Assert.True(result.IsCanceled);
             PipeReader.AdvanceTo(buffer.End);
         }
+
+        [Fact]
+        public Task ReadAtLeastAsyncCancelableWhenWaitingForMoreData()
+        {
+            CancellationTokenSource cts = new CancellationTokenSource();
+            ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(1, cts.Token);
+            cts.Cancel();
+            return Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+        }
+
+        [Fact]
+        public async Task ReadAtLeastAsyncCancelableAfterReadingSome()
+        {
+            CancellationTokenSource cts = new CancellationTokenSource();
+            await Pipe.WriteAsync(new byte[10], default);
+            ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(11, cts.Token);
+            cts.Cancel();
+            await Assert.ThrowsAsync<OperationCanceledException>(async () => await task);
+        }
+
+        [Fact]
+        public async Task ReadAtLeastAsyncCancelableAfterReadingSomeAndWritingAfterStartingRead()
+        {
+            CancellationTokenSource cts = new CancellationTokenSource();
+            await Pipe.WriteAsync(new byte[10], default);
+            ValueTask<ReadResult> task = PipeReader.ReadAtLeastAsync(12, cts.Token);
+            // Write, but not enough to unblock ReadAtLeastAsync
+            await Pipe.WriteAsync(new byte[1], default);
+            cts.Cancel();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await task);
+        }
     }
 }