Fixed potential race condition issues with cancellation and exception handling.

Author: Oren (electricessence)

Open.ChannelExtensions.Tests/ExceptionTests.cs

@@ -40,29 +40,35 @@ await range
 		[Fact]
 		public static async Task ExceptionPropagationConcurrent()
 		{
-			const int testSize = 1000000;
+			const int testSize = 100000000;
 			int total = 0;
 			int count = 0;
 			var range = Enumerable.Range(0, testSize);
-			try
+			await Assert.ThrowsAsync<AggregateException>(async () =>
 			{
-				await range
-					.ToChannel()
-					.ReadAllConcurrently(8, i =>
-					{
-						Interlocked.Increment(ref total);
-						if (i == 500)
+				try
+				{
+
+					await range
+						.ToChannel()
+						.ReadAllConcurrently(8, i =>
 						{
-							Interlocked.Increment(ref count);
-							throw new TestException();
-						}
-					});
-			}
-			catch (Exception ex)
-			{
-				Assert.IsType<AggregateException>(ex);
-				Assert.IsType<TestException>(((AggregateException)ex).InnerException);
-			}
+							Interlocked.Increment(ref total);
+							if (i == 500)
+							{
+								Interlocked.Increment(ref count);
+								throw new TestException();
+							}
+						});
+
+				}
+				catch (Exception ex)
+				{
+					Assert.IsType<AggregateException>(ex);
+					Assert.IsType<TestException>(((AggregateException)ex).InnerException);
+					throw;
+				}
+			});
 
 			Assert.Equal(1, count);
 			Assert.NotEqual(testSize, total);

Open.ChannelExtensions/Extensions.Read.cs

@@ -87,22 +87,30 @@ public static async ValueTask<long> ReadUntilCancelledAsync<T>(this ChannelReade
 			if (deferredExecution)
 				await Task.Yield();
 
+
 			long index = 0;
-			do
+			try
 			{
-				var next = new ValueTask();
-				while (
-					!cancellationToken.IsCancellationRequested
-					&& reader.TryRead(out var item))
+				do
 				{
+					var next = new ValueTask();
+					while (
+						!cancellationToken.IsCancellationRequested
+						&& reader.TryRead(out var item))
+					{
+						await next.ConfigureAwait(false);
+						next = receiver(item, index++);
+					}
 					await next.ConfigureAwait(false);
-					next = receiver(item, index++);
 				}
-				await next.ConfigureAwait(false);
+				while (
+					!cancellationToken.IsCancellationRequested
+					&& await reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false));
+			}
+			catch (TaskCanceledException)
+			{
+				// In case WaitToReadAsync is cancelled.
 			}
-			while (
-				!cancellationToken.IsCancellationRequested
-				&& await reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false));
 
 			return index;
 		}

Open.ChannelExtensions/Extensions.ReadConcurrently.cs

@@ -35,6 +35,7 @@ public static Task<long> ReadAllConcurrentlyAsync<T>(this ChannelReader<T> reade
 				return reader.ReadAllAsync(receiver, cancellationToken, true).AsTask();
 
 			var tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+			var token = tokenSource.Token;
 			var readers = new Task<long>[maxConcurrency];
 			for (var r = 0; r < maxConcurrency; ++r)
 				readers[r] = Read();
@@ -45,18 +46,19 @@ public static Task<long> ReadAllConcurrentlyAsync<T>(this ChannelReader<T> reade
 					{
 						tokenSource.Dispose();
 						if (t.IsFaulted) return Task.FromException<long>(t.Exception);
+						if (t.IsCanceled) return Task.FromCanceled<long>(token);
 						return Task.FromResult(t.Result.Sum());
 					},
 					CancellationToken.None,
-					TaskContinuationOptions.NotOnCanceled | TaskContinuationOptions.ExecuteSynchronously,
+					TaskContinuationOptions.ExecuteSynchronously,
 					TaskScheduler.Current)
 				.Unwrap();
 
 			async Task<long> Read()
 			{
 				try
 				{
-					return await reader.ReadUntilCancelledAsync(tokenSource.Token, ParallelReceiver, true);
+					return await reader.ReadUntilCancelledAsync(token, ParallelReceiver, true);
 				}
 				catch
 				{