Fix behaviour around large buffers and early reader completion (#36)

* Fix behaviour around large buffers and early reader completion

* Cleanup pipe reader on behalf of the target config

* Tweak test for early completion/pipe blocking logic

Author: Turnerj

src/TurnerSoftware.Aqueduct/PipeBifurcation.cs

@@ -30,11 +30,26 @@ public BifurcationState(BifurcationTargetConfig<TResult> config)
 
 		public bool IsCompleted { get; private set; }
 
+		private async Task<TResult> RunReaderWithCleanup(CancellationToken cancellationToken)
+		{
+			try
+			{
+				var result = await Config.Reader(Pipe.Reader, cancellationToken);
+				await Pipe.Reader.CompleteAsync();
+				return result;
+			}
+			catch (Exception ex)
+			{
+				await Pipe.Reader.CompleteAsync(ex);
+				throw;
+			}
+		}
+
 		public void StartReader(CancellationToken cancellationToken)
 		{
 			try
 			{
-				ReaderTask = Config.Reader(Pipe.Reader, cancellationToken);
+				ReaderTask = RunReaderWithCleanup(cancellationToken);
 			}
 			catch (Exception ex)
 			{
@@ -55,10 +70,19 @@ public async ValueTask<bool> WriteAsync(ReadOnlySequence<byte> buffer, Cancellat
 				return false;
 			}
 
-			//Await faulted readers to correctly bubble exceptions
-			if (ReaderTask is not null && ReaderTask.IsFaulted)
+			if (ReaderTask is not null)
 			{
-				await ReaderTask;
+				//Await faulted readers to correctly bubble exceptions
+				if (ReaderTask.IsFaulted)
+				{
+					await ReaderTask;
+				}
+
+				//If the reader task finishes early for some other reason
+				if (ReaderTask.IsCompleted)
+				{
+					return false;
+				}
 			}
 
 			var bytesToRead = (int)buffer.Length;
@@ -68,7 +92,7 @@ public async ValueTask<bool> WriteAsync(ReadOnlySequence<byte> buffer, Cancellat
 			}
 
 			var destination = Pipe.Writer.GetMemory(bytesToRead);
-			buffer.CopyTo(destination.Span);
+			buffer.Slice(0, bytesToRead).CopyTo(destination.Span);
 
 			Pipe.Writer.Advance(bytesToRead);
 

tests/TurnerSoftware.Aqueduct.Tests/PipeBifurcationTests.cs

@@ -72,6 +72,37 @@ await PipeBifurcation.BifurcatedReadAsync(
 		targetReaderHasCompleted.Should().BeTrue();
 	}
 
+	[TestMethod]
+	public async Task SingleTarget_ReaderCompletesEarly_Success()
+	{
+		var source = new Pipe();
+		var buffer = new byte[16];
+		await source.Writer.WriteAsync(buffer);
+		var targetReaderHasCompleted = false;
+
+		var bifurcationTask = PipeBifurcation.BifurcatedReadAsync(
+			source.Reader,
+			new BifurcationSourceConfig(minReadBufferSize: -1),
+			new BifurcationTargetConfig(
+				async (Stream reader, CancellationToken cancellationToken) =>
+				{
+					var buffer = new byte[1];
+					await reader.ReadAsync(buffer, cancellationToken);
+					targetReaderHasCompleted = true;
+				},
+				//These are set to exaggerate the problem where exiting early
+				//still has the pipe being fed bytes till the point it blocks
+				blockAfter: 16,
+				resumeAfter: 8
+			)
+		);
+
+		await source.Writer.WriteAsync(buffer);
+		await source.Writer.CompleteAsync();
+		await bifurcationTask;
+		targetReaderHasCompleted.Should().BeTrue();
+	}
+
 	[TestMethod]
 	public async Task MultiTarget_DefaultConfig_Success()
 	{