Fix possible missed completion if cancellation occured before WriteAll.

Author: electricessence

Open.ChannelExtensions.Tests/ExceptionTests.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Linq;
 using System.Threading;
+using System.Threading.Channels;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -71,4 +72,35 @@ await range
 		Assert.Equal(1, count);
 		Assert.NotEqual(testSize, total);
 	}
+
+	[Fact]
+	public static void ChannelClosed()
+	{
+		var channel = Channel.CreateBounded<int>(new BoundedChannelOptions(1000)
+		{
+			SingleWriter = true,
+			SingleReader = true,
+		});
+		channel.Writer.Complete();
+		// Needs to throw immediately if true.
+		Assert.Throws<ChannelClosedException>(() => channel.Writer.WaitToWriteAndThrowIfClosedAsync());
+	}
+
+	[Fact]
+	public static async Task WriteAllThrowIfClosed()
+	{
+		var channel = Channel.CreateBounded<int>(new BoundedChannelOptions(1000)
+		{
+			SingleWriter = true,
+			SingleReader = true,
+		});
+		var reader = channel.Source(Enumerable.Range(0, 10_000));
+		await reader.ReadAll(_ => { });
+
+		await Assert.ThrowsAsync<ChannelClosedException>(async ()=>
+		{
+			channel.Source(Enumerable.Range(0, 10_000), out var completion);
+			await completion;
+		});
+	}
 }

Open.ChannelExtensions.Tests/SourceTests.cs

@@ -0,0 +1,39 @@
+using System;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Open.ChannelExtensions.Tests;
+public static class SourceTests
+{
+	[Fact]
+	public static async Task ToChannelCancelledAfterwriteStarts()
+	{
+		var cts = new CancellationTokenSource();
+		var reader = Enumerable.Range(0, 10_000).ToChannel(10, true, cts.Token);
+		cts.Cancel();
+
+		try
+		{
+			await reader.ReadAll(_ => { }, cts.Token);
+		}
+		catch (OperationCanceledException)
+		{ }
+
+		await reader.ReadAll(_ => { });
+		await Assert.ThrowsAsync<TaskCanceledException>(()=>reader.Completion);
+	}
+
+	[Fact]
+	public static async Task ToChannelCancelledBeforeWriteStarts()
+	{
+		var cts = new CancellationTokenSource();
+		cts.Cancel();
+		var reader = Enumerable.Range(0, 10_000).ToChannel(10, true, cts.Token);
+
+		var count = await reader.ReadAll(_ => { });
+		Assert.Equal(0, count);
+		await Assert.ThrowsAsync<TaskCanceledException>(() => reader.Completion);
+	}
+}

Open.ChannelExtensions/Extensions.Source.cs

@@ -20,6 +20,7 @@ public static partial class Extensions
 	/// <param name="completion">The underlying ValueTask used to pass the data from the source to the channel.</param>
 	/// <param name="deferredExecution">If true, calls await Task.Yield() before writing to the channel.</param>
 	/// <param name="cancellationToken">An optional cancellation token.</param>
+	/// <remarks>Calling this method does not throw if the channel is already closed.</remarks>
 	/// <returns>The channel reader.</returns>
 	public static ChannelReader<TRead> SourceAsync<TWrite, TRead>(
 		this Channel<TWrite, TRead> target,

Open.ChannelExtensions/Extensions.Write.cs

@@ -34,10 +34,12 @@ public static async ValueTask<long> WriteAllAsync<T>(
 		if (source is null) throw new ArgumentNullException(nameof(source));
 		Contract.EndContractBlock();
 
-		await target.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken).ConfigureAwait(false);
-
 		try
 		{
+			await target
+				.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken)
+				.ConfigureAwait(false);
+
 			long count = 0;
 			var next = new ValueTask();
 			foreach (ValueTask<T> e in source)
@@ -50,19 +52,20 @@ public static async ValueTask<long> WriteAllAsync<T>(
 			await next.ConfigureAwait(false);
 			return count;
 		}
+		catch(ChannelClosedException) { throw; }
 		catch (Exception ex)
 		{
 			if (complete)
 			{
-				target.Complete(ex);
+				target.TryComplete(ex);
 				complete = false;
 			}
 			throw;
 		}
 		finally
 		{
 			if (complete)
-				target.Complete();
+				target.TryComplete();
 		}
 	}
 
@@ -246,11 +249,11 @@ public static async ValueTask<long> WriteAllLines(
 		if (source is null) throw new ArgumentNullException(nameof(source));
 		Contract.EndContractBlock();
 
-		ValueTask next = target.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken);
-		await next.ConfigureAwait(false);
-
 		try
 		{
+			ValueTask next = target.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken);
+			await next.ConfigureAwait(false);
+
 			long count = 0;
 			bool more = false; // if it completed and actually returned false, no need to bubble the cancellation since it actually completed.
 			while (!cancellationToken.IsCancellationRequested)
@@ -275,19 +278,20 @@ public static async ValueTask<long> WriteAllLines(
 			if (more) cancellationToken.ThrowIfCancellationRequested();
 			return count;
 		}
+		catch (ChannelClosedException) { throw; }
 		catch (Exception ex)
 		{
 			if (complete)
 			{
-				target.Complete(ex);
+				target.TryComplete(ex);
 				complete = false;
 			}
 			throw;
 		}
 		finally
 		{
 			if (complete)
-				target.Complete();
+				target.TryComplete();
 		}
 	}
 
@@ -330,12 +334,12 @@ public static async ValueTask<long> WriteAllAsync<T>(
 		if (source is null) throw new ArgumentNullException(nameof(source));
 		Contract.EndContractBlock();
 
-		await target
-			.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken)
-			.ConfigureAwait(false);
-
 		try
 		{
+			await target
+				.WaitToWriteAndThrowIfClosedAsync(ChannelClosedMessage, deferredExecution, cancellationToken)
+				.ConfigureAwait(false);
+
 			long count = 0;
 			var next = new ValueTask();
 			await foreach (T? value in source)
@@ -347,19 +351,20 @@ await target
 			await next.ConfigureAwait(false);
 			return count;
 		}
+		catch (ChannelClosedException) { throw; }
 		catch (Exception ex)
 		{
 			if (complete)
 			{
-				target.Complete(ex);
+				target.TryComplete(ex);
 				complete = false;
 			}
 			throw;
 		}
 		finally
 		{
 			if (complete)
-				target.Complete();
+				target.TryComplete();
 		}
 	}
 

Open.ChannelExtensions/Extensions._.cs

@@ -66,6 +66,27 @@ internal static IEnumerable<ValueTask<T>> WrapValueTask<T>(this IEnumerable<Func
 			yield return new ValueTask<T>(e());
 	}
 
+
+	/// <summary>
+	/// Uses <see cref="ChannelWriter{T}.WaitToWriteAsync(CancellationToken)"/> to peek and see if the channel can still be written to.
+	/// </summary>
+	/// <typeparam name="T">The type being written to the channel</typeparam>
+	/// <param name="writer">The channel writer.</param>
+	/// <param name="ifClosedMessage">The message to include with the ChannelClosedException if thrown.</param>
+	/// <exception cref="ChannelClosedException">If the channel writer will no longer accept messages.</exception>
+	public static void ThrowIfClosed<T>(this ChannelWriter<T> writer, string? ifClosedMessage = null)
+	{
+		if (writer is null) throw new ArgumentNullException(nameof(writer));
+		ValueTask<bool> waitForWrite = writer.WaitToWriteAsync();
+		if (!waitForWrite.IsCompletedSuccessfully || waitForWrite.Result)
+			return;
+
+		if (string.IsNullOrWhiteSpace(ifClosedMessage))
+			throw new ChannelClosedException();
+
+		throw new ChannelClosedException(ifClosedMessage);
+	}
+
 	/// <summary>
 	/// Waits for opportunity to write to a channel and throws a ChannelClosedException if the channel is closed.
 	/// </summary>

Open.ChannelExtensions/Open.ChannelExtensions.csproj

@@ -22,7 +22,7 @@
 		<RepositoryType>git</RepositoryType>
 		<GeneratePackageOnBuild>true</GeneratePackageOnBuild>
 		<GenerateDocumentationFile>true</GenerateDocumentationFile>
-		<Version>6.2.1</Version>
+		<Version>6.2.2</Version>
 		<PackageReleaseNotes></PackageReleaseNotes>
 		<PackageLicenseExpression>MIT</PackageLicenseExpression>
 		<PublishRepositoryUrl>true</PublishRepositoryUrl>