impl and working tests

Author: woutervanranst

src/WouterVanRanst.Utils.Tests/TaskCompletionBufferTests.cs

@@ -187,6 +187,136 @@ public async Task TaskSetTemporarilyEmptyButMoreTasksAreAdded()
     }
 
 
+
+
+    [Fact]
+    public async Task AddTaskAfterEnumeratorStarts()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        var task1 = SimulateTask("Task1", 200);
+        buffer.Add(task1);
+
+        await using var enumerator = buffer.GetConsumingEnumerable().GetAsyncEnumerator();
+        var moveNextTask = enumerator.MoveNextAsync();
+
+        var task2 = SimulateTask("Task2", 100);
+        buffer.Add(task2);
+        buffer.CompleteAdding();
+
+        Assert.True(await moveNextTask);
+        Assert.Equal("Task2", await enumerator.Current);
+
+        Assert.True(await enumerator.MoveNextAsync());
+        Assert.Equal("Task1", await enumerator.Current);
+
+        Assert.False(await enumerator.MoveNextAsync());
+    }
+
+    [Fact]
+    public async Task AllTasksCanceled()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        var cts = new CancellationTokenSource();
+        cts.Cancel();
+
+        var task1 = Task.FromCanceled<string>(cts.Token);
+        buffer.Add(task1);
+        buffer.CompleteAdding();
+
+        var processedTasks = new List<string>();
+        await foreach (var task in buffer.GetConsumingEnumerable())
+        {
+            try
+            {
+                processedTasks.Add(await task);
+            }
+            catch (TaskCanceledException)
+            {
+                processedTasks.Add("Canceled");
+            }
+        }
+
+        Assert.Equal(new[] { "Canceled" }, processedTasks);
+    }
+
+    [Fact]
+    public async Task MultipleProducersAddingConcurrently()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        var producer1 = Task.Run(() =>
+        {
+            buffer.Add(SimulateTask("P1T1", 200));
+            buffer.Add(SimulateTask("P1T2", 100));
+        });
+
+        var producer2 = Task.Run(() =>
+        {
+            buffer.Add(SimulateTask("P2T1", 150));
+            buffer.Add(SimulateTask("P2T2", 50));
+        });
+
+        await Task.WhenAll(producer1, producer2);
+        buffer.CompleteAdding();
+
+        var processedTasks = new List<string>();
+        await foreach (var task in buffer.GetConsumingEnumerable())
+        {
+            processedTasks.Add(await task);
+        }
+
+        Assert.Equal(new[] { "P2T2", "P1T2", "P2T1", "P1T1" }, processedTasks);
+    }
+
+    [Fact]
+    public async Task AddCompletedTaskImmediatelyProcessed()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        var task = Task.FromResult("Immediate");
+        buffer.Add(task);
+        buffer.CompleteAdding();
+
+        var processedTasks = new List<string>();
+        await foreach (var t in buffer.GetConsumingEnumerable())
+        {
+            processedTasks.Add(await t);
+        }
+
+        Assert.Equal(new[] { "Immediate" }, processedTasks);
+    }
+
+    [Fact]
+    public void CompleteAddingIsIdempotent()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        buffer.CompleteAdding();
+        buffer.CompleteAdding(); // Should not throw
+    }
+
+    [Fact]
+    public async Task EnumerateWithoutCompleteAddingBlocks()
+    {
+        var buffer = new TaskCompletionBuffer<string>();
+        var task = SimulateTask("Task", 100);
+        buffer.Add(task);
+
+        var processedTasks = new List<string>();
+        var cts = new CancellationTokenSource(500); // Cancel after 500ms if it blocks indefinitely
+
+        await Assert.ThrowsAsync<OperationCanceledException>(async () =>
+        {
+            await foreach (var t in buffer.GetConsumingEnumerable(cts.Token))
+            {
+                processedTasks.Add(await t);
+            }
+        });
+
+        Assert.Single(processedTasks); // The task should have completed within 500ms
+        Assert.Equal("Task", processedTasks[0]);
+    }
+
+
+
+
     private async Task<string> SimulateTask(string name, int delay)
     {
         await Task.Delay(delay);

src/WouterVanRanst.Utils/Collections/ConcurrentConsumingTaskCollection.cs

@@ -0,0 +1,104 @@
+using System.Threading.Channels;
+
+namespace WouterVanRanst.Utils.Collections;
+
+public class TaskCompletionBuffer<T>
+{
+    private readonly Channel<Task<T>> _channel = Channel.CreateUnbounded<Task<T>>();
+    private readonly object _lock = new object();
+    private volatile bool _completedAdding;
+    private int _pendingTaskCount;
+
+    public void Add(Task<T> task)
+    {
+        lock (_lock)
+        {
+            if (_completedAdding)
+                throw new InvalidOperationException("Cannot add tasks after CompleteAdding has been called.");
+            Interlocked.Increment(ref _pendingTaskCount);
+        }
+
+        task.ContinueWith(t =>
+        {
+            _channel.Writer.TryWrite(t);
+            int newCount = Interlocked.Decrement(ref _pendingTaskCount);
+            if (_completedAdding && newCount == 0)
+                _channel.Writer.TryComplete();
+        }, TaskContinuationOptions.ExecuteSynchronously);
+    }
+
+    public void CompleteAdding()
+    {
+        lock (_lock)
+        {
+            if (_completedAdding)
+                return;
+            _completedAdding = true;
+        }
+
+        if (Interlocked.CompareExchange(ref _pendingTaskCount, 0, 0) == 0)
+            _channel.Writer.TryComplete();
+    }
+
+    public IAsyncEnumerable<Task<T>> GetConsumingEnumerable(CancellationToken cancellationToken = default)
+    {
+        return _channel.Reader.ReadAllAsync(cancellationToken);
+    }
+}
+
+//public class TaskCompletionBuffer<T>
+//{
+//    private readonly Channel<Task<T>> _taskChannel = Channel.CreateUnbounded<Task<T>>();
+//    private readonly CancellationTokenSource _completionSignal = new();
+//    private bool _isAddingCompleted;
+
+//    public void Add(Task<T> task)
+//    {
+//        if (_isAddingCompleted)
+//            throw new InvalidOperationException("Adding tasks has been marked as complete.");
+
+//        _taskChannel.Writer.TryWrite(task);
+//    }
+
+//    public void CompleteAdding()
+//    {
+//        _isAddingCompleted = true;
+//        _taskChannel.Writer.Complete();
+//        _completionSignal.Cancel(); // Signal to exit enumeration when done
+//    }
+
+//    public async IAsyncEnumerable<Task<T>> GetConsumingEnumerable(
+//        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+//    {
+//        var pendingTasks = new List<Task<T>>();
+//        var combinedToken = CancellationTokenSource.CreateLinkedTokenSource(
+//            cancellationToken, _completionSignal.Token).Token;
+
+//        var reader = _taskChannel.Reader;
+
+//        while (!combinedToken.IsCancellationRequested || pendingTasks.Count > 0)
+//        {
+//            // Check for newly added tasks
+//            while (reader.TryRead(out var task))
+//                pendingTasks.Add(task);
+
+//            if (pendingTasks.Count == 0)
+//            {
+//                if (_isAddingCompleted) break; // No more tasks coming
+
+//                // Wait for new tasks or completion signal
+//                await reader.WaitToReadAsync(combinedToken).ConfigureAwait(false);
+//                continue;
+//            }
+
+//            // Wait for any task to complete or new tasks to arrive
+//            var completedTask = await Task.WhenAny([..pendingTasks, reader.WaitToReadAsync(combinedToken).AsTask()]).ConfigureAwait(false);
+
+//            if (completedTask is Task<T> resultTask)
+//            {
+//                pendingTasks.Remove(resultTask);
+//                yield return resultTask;
+//            }
+//        }
+//    }
+//}