Ensured forced update is working.

Author: Oren (electricessence)

Open.ChannelExtensions/BatchingChannelReader.cs

@@ -1,7 +1,9 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics.Contracts;
+using System.Threading;
 using System.Threading.Channels;
+using System.Threading.Tasks;
 
 namespace Open.ChannelExtensions
 {
@@ -12,7 +14,7 @@ namespace Open.ChannelExtensions
 	public class BatchingChannelReader<T> : BufferingChannelReader<T, List<T>>
 	{
 		private readonly int _batchSize;
-		private List<T>? _current;
+		private List<T>? _batch;
 
 		/// <summary>
 		/// Constructs a BatchingChannelReader.
@@ -24,7 +26,6 @@ public BatchingChannelReader(ChannelReader<T> source, int batchSize, bool single
 			Contract.EndContractBlock();
 
 			_batchSize = batchSize;
-			_current = source.Completion.IsCompleted ? null : new List<T>(batchSize);
 		}
 
 		/// <summary>
@@ -35,54 +36,53 @@ public bool ForceBatch()
 		{
 			if (Buffer == null || Buffer.Reader.Completion.IsCompleted) return false;
 			if (TryPipeItems()) return true;
+			if (_batch == null) return false;
 
 			lock (Buffer)
 			{
 				if (Buffer.Reader.Completion.IsCompleted) return false;
 				if (TryPipeItems()) return true;
-				var c = _current;
-				if (c == null || c.Count == 0 || Buffer.Reader.Completion.IsCompleted)
+				var c = _batch;
+				if (c == null || Buffer.Reader.Completion.IsCompleted)
 					return false;
 				c.TrimExcess();
-				_current = new List<T>(_batchSize);
-				Buffer.Writer.TryWrite(c);
+				_batch = null;
+				return Buffer.Writer.TryWrite(c); // Should always be true at this point.
 			}
-
-			return true;
 		}
 
 		/// <inheritdoc />
 		protected override bool TryPipeItems()
 		{
-			if (_current == null || Buffer == null || Buffer.Reader.Completion.IsCompleted)
+			if (Buffer == null || Buffer.Reader.Completion.IsCompleted)
 				return false;
 
 			lock (Buffer)
 			{
-				var c = _current;
-				if (c == null || Buffer.Reader.Completion.IsCompleted)
-					return false;
+				if (Buffer.Reader.Completion.IsCompleted) return false;
 
+				var c = _batch;
 				var source = Source;
 				if (source == null || source.Completion.IsCompleted)
 				{
 					// All finished, release the last batch to the buffer.
+					if (c == null) return false;
+
 					c.TrimExcess();
-					_current = null;
-					if (c.Count == 0)
-						return false;
+					_batch = null;
 
 					Buffer.Writer.TryWrite(c);
 					return true;
 				}
 
 				while (source.TryRead(out T item))
 				{
-					c.Add(item);
+					if (c == null) _batch = c = new List<T>(_batchSize) { item };
+					else c.Add(item);
 
 					if (c.Count == _batchSize)
 					{
-						_current = new List<T>(_batchSize);
+						_batch = null;
 						Buffer.Writer.TryWrite(c);
 						return true;
 					}
@@ -91,5 +91,39 @@ protected override bool TryPipeItems()
 				return false;
 			}
 		}
+
+		/// <inheritdoc />
+		protected override async ValueTask<bool> WaitToReadAsyncCore(ValueTask<bool> bufferWait, CancellationToken cancellationToken)
+		{
+
+			var source = Source;
+			if (source == null) return await bufferWait;
+
+			var b = bufferWait.AsTask();
+			using var tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+			var token = tokenSource.Token;
+
+		start:
+
+			if (b.IsCompleted) return await b.ConfigureAwait(false);
+
+			var s = source.WaitToReadAsync(token);
+			if (s.IsCompleted && !b.IsCompleted) TryPipeItems();
+
+			if (b.IsCompleted)
+			{
+				tokenSource.Cancel();
+				return await b.ConfigureAwait(false);
+			}
+			await Task.WhenAny(s.AsTask(), b).ConfigureAwait(false);
+			if (b.IsCompleted)
+			{
+				tokenSource.Cancel();
+				return await b.ConfigureAwait(false);
+			}
+
+			TryPipeItems();
+			goto start;
+		}
 	}
 }

Open.ChannelExtensions/BufferingChannelReader.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Diagnostics;
 using System.Diagnostics.Contracts;
 using System.Threading;
 using System.Threading.Channels;
@@ -35,10 +36,12 @@ protected BufferingChannelReader(ChannelReader<TIn> source, bool singleReader, b
 			if (source.Completion.IsCompleted)
 			{
 				Buffer = null;
+				_completion = Task.CompletedTask;
 			}
 			else
 			{
 				Buffer = Extensions.CreateChannel<TOut>(-1, singleReader, syncCont);
+				_completion = Buffer.Reader.Completion;
 
 				source.Completion.ContinueWith(t =>
 				{
@@ -54,8 +57,9 @@ protected BufferingChannelReader(ChannelReader<TIn> source, bool singleReader, b
 			}
 		}
 
+		private readonly Task _completion;
 		/// <inheritdoc />
-		public override Task Completion => Buffer?.Reader.Completion ?? Task.CompletedTask;
+		public override Task Completion => _completion;
 
 		/// <summary>
 		/// The method that triggers adding entries to the buffer.
@@ -87,33 +91,42 @@ public override ValueTask<bool> WaitToReadAsync(CancellationToken cancellationTo
 				return new ValueTask<bool>(Task.FromCanceled<bool>(cancellationToken));
 
 			var b = Buffer.Reader.WaitToReadAsync(cancellationToken);
-			if (b.IsCompleted)
-				return b;
+			return b.IsCompleted ? b : WaitToReadAsyncCore(b, cancellationToken);
+		}
 
+		/// <summary>
+		/// Implementation for waiting.
+		/// Can be overridden.
+		/// </summary>
+		protected virtual async ValueTask<bool> WaitToReadAsyncCore(ValueTask<bool> bufferWait, CancellationToken cancellationToken)
+		{
 			var source = Source;
-			if (source == null)
-				return b;
-
-			return WaitCore();
+			if (source == null) return await bufferWait;
 
-			async ValueTask<bool> WaitCore()
-			{
+			using var tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+			var token = tokenSource.Token;
 
-			start:
+		start:
 
-				if (b.IsCompleted) return await b;
+			if (bufferWait.IsCompleted) return await bufferWait;
 
-				var s = source!.WaitToReadAsync(cancellationToken);
-				if (s.IsCompleted && !b.IsCompleted)
-					TryPipeItems();
+			var s = source.WaitToReadAsync(token);
+			if (s.IsCompleted && !bufferWait.IsCompleted) TryPipeItems();
 
-				if (b.IsCompleted) return await b;
-				await s;
-				if (b.IsCompleted) return await b;
-				TryPipeItems();
-
-				goto start;
+			if (bufferWait.IsCompleted)
+			{
+				tokenSource.Cancel();
+				return await bufferWait.ConfigureAwait(false);
 			}
+			await s;
+			if (bufferWait.IsCompleted)
+			{
+				tokenSource.Cancel();
+				return await bufferWait.ConfigureAwait(false);
+			}
+			TryPipeItems();
+
+			goto start;
 		}
 	}
 }

Open.ChannelExtensions/Documentation.xml

@@ -18,7 +18,7 @@
 		<RepositoryType>git</RepositoryType>
 		<PackageTags>channels dotnet threading tasks extensions</PackageTags>
 		<GeneratePackageOnBuild>true</GeneratePackageOnBuild>
-		<Version>3.3.2</Version>
+		<Version>3.4.0</Version>
 		<PackageLicenseExpression>MIT</PackageLicenseExpression>
 		<PackageReleaseNotes></PackageReleaseNotes>
 		<RepositoryUrl>https://github.com/electricessence/Open.ChannelExtensions</RepositoryUrl>