Streams ensure cts cancel fix statecheck

src/core/Akka.Streams.Tests/Dsl/AsyncEnumerableSpec.cs

@@ -9,6 +9,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Akka.Actor;
 using Akka.Pattern;
 using Akka.Streams.Dsl;
 using Akka.Streams.TestKit;
@@ -269,6 +270,150 @@ await EventFilter.Warning().ExpectAsync(0, async () =>
             });
         }
 
+        /// <summary>
+        /// Reproduction for https://github.com/akkadotnet/akka.net/issues/6903
+        /// </summary>
+        [Fact(DisplayName = "AsyncEnumerable Source should dispose underlying async enumerator on kill switch signal")]
+        public async Task AsyncEnumerableSource_Disposes_On_KillSwitch()
+        {
+            await this.AssertAllStagesStoppedAsync(async () =>
+            {
+                var probe = CreateTestProbe();
+                var enumerable = new TestAsyncEnumerable(500.Milliseconds());
+                var src = Source.From(() => enumerable)
+                    .ViaMaterialized(KillSwitches.Single<int>(), Keep.Right)
+                    .ToMaterialized(Sink.ActorRefWithAck<int>(probe, "init", "ack", "complete"), Keep.Left);
+                var killSwitch = src.Run(Materializer);
+
+                // assert init was sent
+                await probe.ExpectMsgAsync<string>(msg => msg == "init");
+                probe.Sender.Tell("ack");
+
+                // assert enumerator is working
+                foreach (var i in Enumerable.Range(0, 5))
+                {
+                    await probe.ExpectMsgAsync<int>(msg => msg == i);
+                    probe.Sender.Tell("ack");
+                }
+
+                // last message was not ack-ed
+                await probe.ExpectMsgAsync<int>(msg => msg == 5);
+
+                killSwitch.Shutdown();
+
+                // assert that enumerable resource was disposed
+                await AwaitConditionAsync(() => enumerable.Disposed);
+            }, Materializer);
+        }
+
+        [Fact(DisplayName = "AsyncEnumerable Source should dispose underlying async enumerator on kill switch signal even after ActorSystem termination")]
+        public async Task AsyncEnumerableSource_Disposes_On_KillSwitch2()
+        {
+            var probe = CreateTestProbe();
+            // A long disposing enumerable source
+            var enumerable = new TestAsyncEnumerable(2.Seconds());
+            var src = Source.From(() => enumerable)
+                .ViaMaterialized(KillSwitches.Single<int>(), Keep.Right)
+                .ToMaterialized(Sink.ActorRefWithAck<int>(probe, "init", "ack", "complete"), Keep.Left);
+            var killSwitch = src.Run(Materializer);
+
+            // assert init was sent
+            await probe.ExpectMsgAsync<string>(msg => msg == "init");
+            probe.Sender.Tell("ack");
+
+            // assert enumerator is working
+            foreach (var i in Enumerable.Range(0, 5))
+            {
+                await probe.ExpectMsgAsync<int>(msg => msg == i);
+                probe.Sender.Tell("ack");
+            }
+
+            // last message was not ack-ed
+            await probe.ExpectMsgAsync<int>(msg => msg == 5);
+
+            killSwitch.Shutdown();
+
+            await Sys.Terminate();
+
+            // enumerable was not disposed even after system termination
+            enumerable.Disposed.Should().BeFalse();
+
+            // assert that enumerable resource can still be disposed even after system termination
+            // (Not guaranteed if process was already killed)
+            await AwaitConditionAsync(() => enumerable.Disposed);
+        }
+
+        private class TestAsyncEnumerable: IAsyncEnumerable<int>
+        {
+            private readonly AsyncEnumerator _enumerator;
+
+            public bool Disposed => _enumerator.Disposed;
+
+            public TestAsyncEnumerable(TimeSpan shutdownDelay)
+            {
+                _enumerator = new AsyncEnumerator(shutdownDelay);
+            }
+
+            public IAsyncEnumerator<int> GetAsyncEnumerator(CancellationToken token = default)
+            {
+                token.ThrowIfCancellationRequested();
+                return _enumerator;
+            }
+
+            private sealed class AsyncEnumerator: IAsyncEnumerator<int>
+            {
+                private readonly TimeSpan _shutdownDelay;
+                private int _current = -1;
+
+                public AsyncEnumerator(TimeSpan shutdownDelay)
+                {
+                    _shutdownDelay = shutdownDelay;
+                }
+
+                public bool Disposed { get; private set; }
+
+                public async ValueTask DisposeAsync()
+                {
+                    await Task.Delay(_shutdownDelay);
+                    Disposed = true;
+                }
+
+                public async ValueTask<bool> MoveNextAsync()
+                {
+                    await Task.Delay(100);
+                    _current++;
+                    return true;
+                }
+
+                public int Current
+                {
+                    get
+                    {
+                        if (_current == -1)
+                            throw new IndexOutOfRangeException("MoveNextAsync has not been called");
+                        if (Disposed)
+                            throw new ObjectDisposedException("Enumerator already disposed");
+                        return _current;
+                    }
+                }
+            }
+        }
+
+        [Fact]
+        public async Task AsyncEnumerableSource_Disposes_OnCancel()
+        {
+            var resource = new Resource();
+            var tcs = new System.Threading.Tasks.TaskCompletionSource<NotUsed>(TaskCreationOptions
+                .RunContinuationsAsynchronously);
+            var src = Source.From(() =>
+                CancelTestGenerator(tcs, resource, default));
+            src.To(Sink.Ignore<int>()).Run(Materializer);
+            await tcs.Task;
+            Materializer.Shutdown();
+            await Task.Delay(500);
+            Assert.False(resource.IsActive);
+        }
+
         private static async IAsyncEnumerable<int> RangeAsync(int start, int count, 
             [EnumeratorCancellation] CancellationToken token = default)
         {
@@ -308,6 +453,39 @@ private static async IAsyncEnumerable<int> ProbeableRangeAsync(int start, int co
                 yield return i;
             }
         }
+
+        public static async IAsyncEnumerable<int> CancelTestGenerator(
+            TaskCompletionSource<NotUsed> tcs,
+            Resource resource,
+            [EnumeratorCancellation] CancellationToken token
+        )
+        {
+            await using var res = resource;
+            int i = 0;
+            bool isSet = false;
+            while (true)
+            {
+                await Task.Delay(1, token).ConfigureAwait(false);
+                yield return i++;
+                if (isSet == false)
+                {
+                    tcs.TrySetResult(NotUsed.Instance);
+                    isSet = true;
+                }
+            }
+            // ReSharper disable once IteratorNeverReturns
+        }
+
+        public class Resource : IAsyncDisposable
+        {
+            public bool IsActive = true;
+            public ValueTask DisposeAsync()
+            {
+                IsActive = false;
+                Console.WriteLine("Enumerator completed and resource disposed");
+                return new ValueTask();
+            }
+        }
     }
 #endif
 }

src/core/Akka.Streams/Implementation/Fusing/Ops.cs

@@ -3812,18 +3812,19 @@ private sealed class Logic : OutGraphStageLogic
             private readonly Action<T> _onSuccess;
             private readonly Action<Exception> _onFailure;
             private readonly Action _onComplete;
+            private readonly CancellationTokenSource _completionCts;
 
-            private CancellationTokenSource _completionCts;
             private IAsyncEnumerator<T> _enumerator;
 
             public Logic(SourceShape<T> shape, IAsyncEnumerable<T> enumerable) : base(shape)
             {
+
                 _enumerable = enumerable;
                 _outlet = shape.Outlet;
                 _onSuccess = GetAsyncCallback<T>(OnSuccess);
                 _onFailure = GetAsyncCallback<Exception>(OnFailure);
                 _onComplete = GetAsyncCallback(OnComplete);
-
+                _completionCts = new CancellationTokenSource();
                 SetHandler(_outlet, this);
             }
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -3838,9 +3839,50 @@ public Logic(SourceShape<T> shape, IAsyncEnumerable<T> enumerable) : base(shape)
             public override void PreStart()
             {
                 base.PreStart();
-                _completionCts = new CancellationTokenSource();
                 _enumerator = _enumerable.GetAsyncEnumerator(_completionCts.Token);
             }
+
+            public override void PostStop()
+            {
+                try
+                {
+                    _completionCts.Cancel();
+                    _completionCts.Dispose();
+                }
+                catch(Exception ex)
+                {
+                    // This should never happen
+                    Log.Debug(ex, "AsyncEnumerable threw while cancelling CancellationTokenSource");
+                }
+
+                try
+                {
+#pragma warning disable CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
+                    // Intentionally creating a detached dispose task
+                    DisposeEnumeratorAsync();
+#pragma warning restore CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
+                }
+                catch(Exception ex)
+                {
+                    Log.Debug(ex, "Underlying async enumerator threw an exception while being disposed.");
+                }
+                base.PostStop();
+                return;
+
+                async Task DisposeEnumeratorAsync()
+                {
+                    try
+                    {
+                        await _enumerator.DisposeAsync();
+                    }
+                    catch (Exception ex)
+                    {
+                        // This is best effort exception logging, this log will never appear if the ActorSystem
+                        // was shut down before we reach this code (BusEvent was not emitting new logs anymore)
+                        Log.Debug(ex, "Underlying async enumerator threw an exception while being disposed.");
+                    }
+                }
+            }
 
             public override void OnPull()
             {
@@ -3859,26 +3901,12 @@ public override void OnPull()
                         // if result is false, it means enumerator was closed. Complete stage in that case.
                         CompleteStage();
                     }
-                } 
-                else if (vtask.IsCompleted) // IsCompleted covers Faulted, Cancelled, and RanToCompletion async state
-                {
-                    // vtask will always contains an exception because we know we're not successful and always throws
-                    try
-                    {
-                        // This does not block because we know that the task already completed
-                        // Using GetAwaiter().GetResult() to automatically unwraps AggregateException inner exception
-                        vtask.GetAwaiter().GetResult();
-                    }
-                    catch (Exception ex)
-                    {
-                        FailStage(ex);
-                        return;
-                    }
-
-                    throw new InvalidOperationException("Should never reach this code");
                 }
                 else
                 {
+                    //We immediately fall into wait case.
+                    //Unlike Task, we don't have a 'status' Enum to switch off easily,
+                    //And Error cases can just live with the small cost of async callback.
                     async Task ProcessTask()
                     {
                         // Since this Action is used as task continuation, we cannot safely call corresponding
@@ -3897,16 +3925,12 @@ async Task ProcessTask()
                         }
                     }
 
-#pragma warning disable CS4014
-                    ProcessTask();
-#pragma warning restore CS4014
+                    _ = ProcessTask();
                 }
             }
 
             public override void OnDownstreamFinish(Exception cause)
             {
-                _completionCts.Cancel();
-                _completionCts.Dispose();
                 CompleteStage();
                 base.OnDownstreamFinish(cause);
             }