[5752] FakeLogCollector waiting capabilities

src/Libraries/Microsoft.Extensions.Diagnostics.Testing/Logging/FakeLogCollector.LogEnumeration.cs

@@ -0,0 +1,189 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Extensions.Logging.Testing;
+
+public partial class FakeLogCollector
+{
+    private int _recordCollectionVersion;
+
+    private TaskCompletionSource<object?> _logEnumerationSharedWaiter =
+        new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+    private int _waitingEnumeratorCount;
+
+    public IAsyncEnumerable<FakeLogRecord> GetLogsAsync(CancellationToken cancellationToken = default)
+        => new LogAsyncEnumerable(this, cancellationToken);
+
+    private class LogAsyncEnumerable : IAsyncEnumerable<FakeLogRecord>
+    {
+        private readonly FakeLogCollector _collector;
+        private readonly CancellationToken _enumerableCancellationToken;
+
+        internal LogAsyncEnumerable(
+            FakeLogCollector collector,
+            CancellationToken enumerableCancellationToken)
+        {
+            _collector = collector;
+            _enumerableCancellationToken = enumerableCancellationToken;
+        }
+
+        public IAsyncEnumerator<FakeLogRecord> GetAsyncEnumerator(
+            CancellationToken enumeratorCancellationToken = default)
+                => new StreamEnumerator(_collector, _enumerableCancellationToken, enumeratorCancellationToken);
+    }
+
+    private sealed class StreamEnumerator : IAsyncEnumerator<FakeLogRecord>
+    {
+        private readonly FakeLogCollector _collector;
+        private readonly CancellationTokenSource _mainCts;
+
+        private FakeLogRecord? _current;
+        private int _index;
+        private bool _disposed;
+        private int _observedRecordCollectionVersion;
+
+        // Concurrent MoveNextAsync guard
+        private int _moveNextActive; // 0 = inactive, 1 = active (int type used for net462 compatibility)
+
+        public StreamEnumerator(
+            FakeLogCollector collector,
+            CancellationToken enumerableCancellationToken,
+            CancellationToken enumeratorCancellationToken)
+        {
+            _collector = collector;
+            _mainCts = enumerableCancellationToken.CanBeCanceled || enumeratorCancellationToken.CanBeCanceled
+                ? CancellationTokenSource.CreateLinkedTokenSource([enumerableCancellationToken, enumeratorCancellationToken])
+                : new CancellationTokenSource();
+            _observedRecordCollectionVersion = collector._recordCollectionVersion;
+        }
+
+        public FakeLogRecord Current => _current ?? throw new InvalidOperationException("Enumeration not started.");
+
+        public async ValueTask<bool> MoveNextAsync()
+        {
+            if (Interlocked.CompareExchange(ref _moveNextActive, 1, 0) == 1)
+            {
+                throw new InvalidOperationException("MoveNextAsync is already in progress. Concurrent calls are not allowed.");
+            }
+
+            try
+            {
+                ThrowIfDisposed();
+
+                var masterCancellationToken = _mainCts.Token;
+
+                masterCancellationToken.ThrowIfCancellationRequested();
+
+                while (true)
+                {
+                    TaskCompletionSource<object?>? waiter = null;
+
+                    try
+                    {
+                        masterCancellationToken.ThrowIfCancellationRequested();
+
+                        lock (_collector._records)
+                        {
+                            if (_observedRecordCollectionVersion != _collector._recordCollectionVersion)
+                            {
+                                _index = 0; // based on assumption that version changed on full collection clear
+                                _observedRecordCollectionVersion = _collector._recordCollectionVersion;
+                            }
+
+                            if (_index < _collector._records.Count)
+                            {
+                                _current = _collector._records[_index++];
+                                return true;
+                            }
+
+                            // waiter needs to be subscribed within records lock
+                            // if not: more records could be added in the meantime and the waiter could be stuck waiting even though the index is behind the actual count
+                            waiter = _collector._logEnumerationSharedWaiter;
+                            _collector._waitingEnumeratorCount++;
+                        }
+
+                        // Compatibility path for net462: emulate Task.WaitAsync(cancellationToken).
+                        // After the wait is complete in normal flow, no need to decrement because the shared waiter will be swapped and counter reset.
+                        await AwaitWithCancellationAsync(waiter.Task, masterCancellationToken).ConfigureAwait(false);
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        if (waiter is not null)
+                        {
+                            lock (_collector._records)
+                            {
+                                if (
+                                    _collector._waitingEnumeratorCount > 0 // counter can be zero during the cancellation path
+                                    && waiter == _collector._logEnumerationSharedWaiter // makes sure we adjust the counter for the same shared waiting session
+                                )
+                                {
+                                    _collector._waitingEnumeratorCount--;
+                                }
+                            }
+                        }
+
+                        throw;
+                    }
+                }
+
+            }
+            finally
+            {
+                Volatile.Write(ref _moveNextActive, 0);
+            }
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return default;
+            }
+
+            _disposed = true;
+
+            _mainCts.Cancel();
+            _mainCts.Dispose();
+
+            return default;
+        }
+
+        private static async Task AwaitWithCancellationAsync(Task task, CancellationToken cancellationToken)
+        {
+            if (!cancellationToken.CanBeCanceled || task.IsCompleted)
+            {
+                await task.ConfigureAwait(false);
+                return;
+            }
+
+            var cancelTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            CancellationTokenRegistration ctr = default;
+            try
+            {
+                ctr = cancellationToken.Register(static s =>
+                    ((TaskCompletionSource<bool>)s!).TrySetCanceled(), cancelTcs);
+
+                var completed = await Task.WhenAny(task, cancelTcs.Task).ConfigureAwait(false);
+                await completed.ConfigureAwait(false);
+            }
+            finally
+            {
+                ctr.Dispose();
+            }
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(nameof(StreamEnumerator));
+            }
+        }
+    }
+}

src/Libraries/Microsoft.Extensions.Diagnostics.Testing/Logging/FakeLogCollector.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Threading.Tasks;
 using Microsoft.Extensions.Options;
 using Microsoft.Shared.Diagnostics;
 
@@ -14,7 +15,7 @@ namespace Microsoft.Extensions.Logging.Testing;
 /// </summary>
 [DebuggerDisplay("Count = {Count}, LatestRecord = {LatestRecord}")]
 [DebuggerTypeProxy(typeof(FakeLogCollectorDebugView))]
-public class FakeLogCollector
+public partial class FakeLogCollector
 {
     private readonly List<FakeLogRecord> _records = [];
     private readonly FakeLogCollectorOptions _options;
@@ -54,6 +55,7 @@ public void Clear()
         lock (_records)
         {
             _records.Clear();
+            _recordCollectionVersion++;
         }
     }
 
@@ -136,11 +138,23 @@ internal void AddRecord(FakeLogRecord record)
             return;
         }
 
+        TaskCompletionSource<object?>? logEnumerationSharedWaiterToWake = null;
+
         lock (_records)
         {
             _records.Add(record);
+
+            if (_waitingEnumeratorCount > 0)
+            {
+                logEnumerationSharedWaiterToWake = _logEnumerationSharedWaiter;
+                _logEnumerationSharedWaiter = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+                _waitingEnumeratorCount = 0;
+            }
         }
 
+        // it is possible the task was already completed, but it does not matter and we can avoid locking
+        _ = logEnumerationSharedWaiterToWake?.TrySetResult(null);
+
         _options.OutputSink?.Invoke(_options.OutputFormatter(record));
     }