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,102 @@ 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();
+                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);
+        }
+
+        private class TestAsyncEnumerable: IAsyncEnumerable<int>
+        {
+            private readonly AsyncEnumerator _enumerator = new();
+
+            public bool Disposed => _enumerator.Disposed;
+
+            public IAsyncEnumerator<int> GetAsyncEnumerator(CancellationToken token = default)
+            {
+                token.ThrowIfCancellationRequested();
+                return _enumerator;
+            }
+
+            private sealed class AsyncEnumerator: IAsyncEnumerator<int>
+            {
+                private int _current = -1;
+
+                public bool Disposed { get; private set; }
+
+                public async ValueTask DisposeAsync()
+                {
+                    await Task.Delay(100);
+                    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 +405,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,33 @@ 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.Error(ex, "AsyncEnumerable threw while cancelling CancellationTokenSource");
+                }
+                try
+                {
+                    var disposed = _enumerator.DisposeAsync().AsTask().Wait(TimeSpan.FromSeconds(10));
+                    if(!disposed)
+                        Log.Error("Failed to dispose underlying async enumerator. DisposeAsync timed out after 10 seconds.");
+                }
+                catch(Exception ex)
+                {
+                    Log.Error(ex, "Underlying async enumerator threw an exception while being disposed.");
+                }
+                base.PostStop();
+            }
 
             public override void OnPull()
             {
@@ -3859,26 +3884,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 +3908,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);
             }