Fixed race conditions + unsafe struct assignment in SelectAsync (#7521)

* Fixed race conditions + unsafe struct assignment in `SelectAsync`

close #7518

* added nullability API approvals

* fixed `FlowSelectAsyncSpecs`

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

Co-authored-by: Michael Buck <[email protected]>

* relaxed nullability requirements for `Holder<TOut>`

`null` might be a perfectly acceptable value inside a `SelectAsync` stage

* fixed `FlowAskSpec`

* forgot to pass in innerException

* fix race condition in `FlowAskSpec`

---------

Co-authored-by: Michael Buck <[email protected]>

Author: Aaronontheweb

src/core/Akka.API.Tests/verify/CoreAPISpec.ApproveStreams.DotNet.verified.txt

@@ -4430,7 +4430,12 @@ namespace Akka.Streams.Implementation.Fusing
         public override string ToString() { }
     }
     [Akka.Annotations.InternalApiAttribute()]
-    public sealed class SelectAsyncUnordered<TIn, TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
+    [System.Runtime.CompilerServices.NullableAttribute(new byte[] {
+            0,
+            1,
+            1,
+            1})]
+    public sealed class SelectAsyncUnordered<[System.Runtime.CompilerServices.NullableAttribute(2)]  TIn, [System.Runtime.CompilerServices.NullableAttribute(2)]  TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
     {
         public readonly Akka.Streams.Inlet<TIn> In;
         public readonly Akka.Streams.Outlet<TOut> Out;
@@ -4440,7 +4445,12 @@ namespace Akka.Streams.Implementation.Fusing
         protected override Akka.Streams.Stage.GraphStageLogic CreateLogic(Akka.Streams.Attributes inheritedAttributes) { }
     }
     [Akka.Annotations.InternalApiAttribute()]
-    public sealed class SelectAsync<TIn, TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
+    [System.Runtime.CompilerServices.NullableAttribute(new byte[] {
+            0,
+            1,
+            1,
+            1})]
+    public sealed class SelectAsync<[System.Runtime.CompilerServices.NullableAttribute(2)]  TIn, [System.Runtime.CompilerServices.NullableAttribute(2)]  TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
     {
         public readonly Akka.Streams.Inlet<TIn> In;
         public readonly Akka.Streams.Outlet<TOut> Out;

src/core/Akka.API.Tests/verify/CoreAPISpec.ApproveStreams.Net.verified.txt

@@ -4404,7 +4404,12 @@ namespace Akka.Streams.Implementation.Fusing
         public override string ToString() { }
     }
     [Akka.Annotations.InternalApiAttribute()]
-    public sealed class SelectAsyncUnordered<TIn, TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
+    [System.Runtime.CompilerServices.NullableAttribute(new byte[] {
+            0,
+            1,
+            1,
+            1})]
+    public sealed class SelectAsyncUnordered<[System.Runtime.CompilerServices.NullableAttribute(2)]  TIn, [System.Runtime.CompilerServices.NullableAttribute(2)]  TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
     {
         public readonly Akka.Streams.Inlet<TIn> In;
         public readonly Akka.Streams.Outlet<TOut> Out;
@@ -4414,7 +4419,12 @@ namespace Akka.Streams.Implementation.Fusing
         protected override Akka.Streams.Stage.GraphStageLogic CreateLogic(Akka.Streams.Attributes inheritedAttributes) { }
     }
     [Akka.Annotations.InternalApiAttribute()]
-    public sealed class SelectAsync<TIn, TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
+    [System.Runtime.CompilerServices.NullableAttribute(new byte[] {
+            0,
+            1,
+            1,
+            1})]
+    public sealed class SelectAsync<[System.Runtime.CompilerServices.NullableAttribute(2)]  TIn, [System.Runtime.CompilerServices.NullableAttribute(2)]  TOut> : Akka.Streams.Stage.GraphStage<Akka.Streams.FlowShape<TIn, TOut>>
     {
         public readonly Akka.Streams.Inlet<TIn> In;
         public readonly Akka.Streams.Outlet<TOut> Out;

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

@@ -238,9 +238,7 @@ public async Task Flow_with_ask_must_signal_ask_timeout_failure() => await this.
 
             c.ExpectSubscription().Request(10);
             var error = c.ExpectError();
-            error.As<AggregateException>().Flatten()
-                .InnerException
-                .Should().BeOfType<AskTimeoutException>();
+            error.Should().BeOfType<AskTimeoutException>();
             return Task.CompletedTask;
         }, _materializer);
 
@@ -253,8 +251,17 @@ public async Task Flow_with_ask_must_signal_ask_failure() => await this.AssertAl
                 .Ask<Reply>(failsOn, _timeout, 1)
                 .RunWith(Sink.FromSubscriber(c), _materializer);
 
-            var error = (AggregateException)c.ExpectSubscriptionAndError();
-            error.InnerException.Message.Should().Be("Booming for 1!");
+            var error = c.ExpectSubscriptionAndError();
+            if (error is AggregateException aggregateException) // happens if we hit the fast path and don't await
+            {
+                aggregateException.Flatten()
+                    .InnerException!.Message.Should().Be("Booming for 1!");
+            }
+            else
+            {
+                error.Message.Should().Be("Booming for 1!");
+            }
+            
             return Task.CompletedTask;
         }, _materializer);
 

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

@@ -171,10 +171,10 @@ await this.AssertAllStagesStoppedAsync(async() => {
 
                 var exception = await probe.AsyncBuilder()
                     .Request(10)
-                    .ExpectNextN(new[]{1, 2})
+                    .ExpectNextN([1, 2])
                     .ExpectErrorAsync()
                     .ShouldCompleteWithin(RemainingOrDefault);
-                exception.InnerException!.Message.Should().Be("err1");
+                exception.Message.Should().Be("err1");
             }, Materializer);
         }
 
@@ -232,7 +232,7 @@ await this.AssertAllStagesStoppedAsync(async () => {
                     .RunWith(Sink.FromSubscriber(c), Materializer);
                 var sub = await c.ExpectSubscriptionAsync();
                 sub.Request(10);
-                c.ExpectError().Message.Should().Be("err2");
+                (await c.ExpectErrorAsync()).Message.Should().Be("err2");
             }, Materializer);
         }
 
@@ -258,7 +258,7 @@ await this.AssertAllStagesStoppedAsync(async () =>
 
                 await probe.AsyncBuilder()
                     .Request(10)
-                    .ExpectNextN(new[] { 1, 2 })
+                    .ExpectNextN([1, 2])
                     .ExpectErrorAsync();
 
                 invoked.Should().BeTrue();
@@ -358,21 +358,21 @@ public async Task A_Flow_with_SelectAsync_must_signal_NPE_when_task_is_completed
         {
             var c = this.CreateManualSubscriberProbe<string>();
 
-            Source.From(new[] {"a", "b"})
-                .SelectAsync(4, _ => Task.FromResult(null as string))
+            Source.From(["a", "b"])
+                .SelectAsync(4, _ => Task.FromResult<string>(null))
                 .To(Sink.FromSubscriber(c)).Run(Materializer);
 
             var sub = await c.ExpectSubscriptionAsync();
             sub.Request(10);
-            c.ExpectError().Message.Should().StartWith(ReactiveStreamsCompliance.ElementMustNotBeNullMsg);
+            (await c.ExpectErrorAsync()).Message.Should().StartWith(ReactiveStreamsCompliance.ElementMustNotBeNullMsg);
         }
 
         [Fact]
         public async Task A_Flow_with_SelectAsync_must_resume_when_task_is_completed_with_null()
         {
             var c = this.CreateManualSubscriberProbe<string>();
-            Source.From(new[] { "a", "b", "c" })
-                .SelectAsync(4, s => s.Equals("b") ? Task.FromResult(null as string) : Task.FromResult(s))
+            Source.From(["a", "b", "c"])
+                .SelectAsync(4, s => s.Equals("b") ? Task.FromResult<string>(null) : Task.FromResult(s))
                 .WithAttributes(ActorAttributes.CreateSupervisionStrategy(Deciders.ResumingDecider))
                 .To(Sink.FromSubscriber(c)).Run(Materializer);
             var sub = await c.ExpectSubscriptionAsync();
@@ -438,21 +438,6 @@ await this.AssertAllStagesStoppedAsync(async() =>
                 }, cancellation.Token);
 #pragma warning restore CS4014 // Because this call is not awaited, execution of the current method continues before the call is completed
 
-                Task<int> Deferred()
-                {
-                    var promise = new TaskCompletionSource<int>();
-                    if (counter.IncrementAndGet() > parallelism)
-                        promise.SetException(new Exception("parallelism exceeded"));
-                    else
-                    {
-                        var wrote = queue.Writer.TryWrite((promise, DateTime.Now.Ticks));
-                        if (!wrote)
-                            promise.SetException(new Exception("Failed to write to queue"));
-                    }
-                        
-                    return promise.Task;
-                }
-
                 try
                 {
                     const int n = 10000;
@@ -467,6 +452,23 @@ Task<int> Deferred()
                 {
                     cancellation.Cancel(false);
                 }
+
+                return;
+
+                Task<int> Deferred()
+                {
+                    var promise = new TaskCompletionSource<int>();
+                    if (counter.IncrementAndGet() > parallelism)
+                        promise.SetException(new Exception("parallelism exceeded"));
+                    else
+                    {
+                        var wrote = queue.Writer.TryWrite((promise, DateTime.Now.Ticks));
+                        if (!wrote)
+                            promise.SetException(new Exception("Failed to write to queue"));
+                    }
+                        
+                    return promise.Task;
+                }
             }, Materializer);
         }
 

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

@@ -22,6 +22,7 @@
 using Akka.Streams.Util;
 using Akka.Util;
 using Akka.Util.Internal;
+using Debug = System.Diagnostics.Debug;
 using Decider = Akka.Streams.Supervision.Decider;
 using Directive = Akka.Streams.Supervision.Directive;
 
@@ -2512,61 +2513,47 @@ public Expand(Func<TIn, IEnumerator<TOut>> extrapolate)
         /// </returns>
         public override string ToString() => "Expand";
     }
+    
+    #nullable enable
 
     /// <summary>
     /// INTERNAL API
     /// </summary>
-    /// <typeparam name="TIn">TBD</typeparam>
-    /// <typeparam name="TOut">TBD</typeparam>
     [InternalApi]
     public sealed class SelectAsync<TIn, TOut> : GraphStage<FlowShape<TIn, TOut>>
     {
         #region internal classes
 
         private sealed class Logic : InAndOutGraphStageLogic
         {
-            private class Holder<T>
+            private sealed class Holder<T>(object? message, Result<T> element)
             {
-                private readonly Action<Holder<T>> _callback;
-
-                public Holder(object message, Result<T> element, Action<Holder<T>> callback)
-                {
-                    _callback = callback;
-                    Message = message;
-                    Element = element;
-                }
-
-                public Result<T> Element { get; private set; }
-                public object Message { get; }
+                public object? Message { get; private set; } = message;
+                
+                public Result<T> Element { get; private set; } = element;
 
                 public void SetElement(Result<T> result)
                 {
                     Element = result.IsSuccess && result.Value == null
                         ? Result.Failure<T>(ReactiveStreamsCompliance.ElementMustNotBeNullException)
                         : result;
                 }
-
-                public void Invoke(Result<T> result)
-                {
-                    SetElement(result);
-                    _callback(this);
-                }
             }
 
             private static readonly Result<TOut> NotYetThere = Result.Failure<TOut>(new Exception());
 
             private readonly SelectAsync<TIn, TOut> _stage;
             private readonly Decider _decider;
             private IBuffer<Holder<TOut>> _buffer;
-            private readonly Action<Holder<TOut>> _taskCallback;
+            private readonly Action<(Holder<TOut>, Result<TOut>)> _taskCallback;
 
             public Logic(Attributes inheritedAttributes, SelectAsync<TIn, TOut> stage) : base(stage.Shape)
             {
                 _stage = stage;
-                var attr = inheritedAttributes.GetAttribute<ActorAttributes.SupervisionStrategy>(null);
+                var attr = inheritedAttributes.GetAttribute<ActorAttributes.SupervisionStrategy>();
                 _decider = attr != null ? attr.Decider : Deciders.StoppingDecider;
 
-                _taskCallback = GetAsyncCallback<Holder<TOut>>(HolderCompleted);
+                _taskCallback = GetAsyncCallback<(Holder<TOut> holder, Result<TOut> result)>(t => HolderCompleted(t.holder, t.result));
 
                 SetHandlers(stage.In, stage.Out, this);
             }
@@ -2577,19 +2564,33 @@ public override void OnPush()
                 try
                 {
                     var task = _stage._mapFunc(message);
-                    var holder = new Holder<TOut>(message, NotYetThere, _taskCallback);
+                    var holder = new Holder<TOut>(message, NotYetThere);
                     _buffer.Enqueue(holder);
 
                     // We dispatch the task if it's ready to optimize away
                     // scheduling it to an execution context
                     if (task.IsCompleted)
                     {
-                        holder.SetElement(Result.FromTask(task));
-                        HolderCompleted(holder);
+                        HolderCompleted(holder, Result.FromTask(task));
                     }
                     else
-                        task.ContinueWith(t => holder.Invoke(Result.FromTask(t)),
-                            TaskContinuationOptions.ExecuteSynchronously);
+                    {
+                        async Task WaitForTask()
+                        {
+                            try
+                            {
+                                var result = Result.Success(await task);
+                                _taskCallback((holder, result));
+                            }
+                            catch(Exception ex){
+                                var result = Result.Failure<TOut>(ex);
+                                _taskCallback((holder, result));
+                            }   
+                        }
+                        
+                        _ = WaitForTask();
+                    }
+                       
                 }
                 catch (Exception e)
                 {
@@ -2606,7 +2607,7 @@ public override void OnPush()
                             break;
 
                         default:
-                            throw new AggregateException($"Unknown SupervisionStrategy directive: {strategy}", e);
+                            throw new ArgumentOutOfRangeException($"Unknown SupervisionStrategy directive: {strategy}", e);
                     }
                 }
                 if (Todo < _stage._parallelism && !HasBeenPulled(_stage.In))
@@ -2663,12 +2664,12 @@ private void PushOne()
                                     break;
 
                                 default:
-                                    throw new AggregateException($"Unknown SupervisionStrategy directive: {strategy}", result.Exception);
+                                    throw new ArgumentOutOfRangeException($"Unknown SupervisionStrategy directive: {strategy}", result.Exception);
                             }
                             continue;
                         }
 
-                        Push(_stage.Out, result.Value);
+                        Push(_stage.Out!, result.Value);
                         if (Todo < _stage._parallelism && !HasBeenPulled(inlet))
                             TryPull(inlet);
                     }
@@ -2677,17 +2678,18 @@ private void PushOne()
                 }
             }
 
-            private void HolderCompleted(Holder<TOut> holder)
+            private void HolderCompleted(Holder<TOut> holder, Result<TOut> result)
             {
-                var element = holder.Element;
-                if (element.IsSuccess)
+                // we may not be at the front of the line right now, so save the result for later
+                holder.SetElement(result);
+                if (result.IsSuccess)
                 {
                     if (IsAvailable(_stage.Out))
                         PushOne();
                     return;
                 }
 
-                var exception = element.Exception;
+                var exception = result.Exception;
                 var strategy = _decider(exception);
                 Log.Error(exception, "An exception occured inside SelectAsync while executing Task. Supervision strategy: {0}", strategy);
                 switch (strategy)
@@ -2703,7 +2705,7 @@ private void HolderCompleted(Holder<TOut> holder)
                         break;
 
                     default:
-                        throw new AggregateException($"Unknown SupervisionStrategy directive: {strategy}", exception);
+                        throw new ArgumentOutOfRangeException($"Unknown SupervisionStrategy directive: {strategy}", exception);
                 }
             }
 
@@ -2758,8 +2760,6 @@ protected override GraphStageLogic CreateLogic(Attributes inheritedAttributes)
     /// <summary>
     /// INTERNAL API
     /// </summary>
-    /// <typeparam name="TIn">TBD</typeparam>
-    /// <typeparam name="TOut">TBD</typeparam>
     [InternalApi]
     public sealed class SelectAsyncUnordered<TIn, TOut> : GraphStage<FlowShape<TIn, TOut>>
     {
@@ -2904,6 +2904,8 @@ public SelectAsyncUnordered(int parallelism, Func<TIn, Task<TOut>> mapFunc)
         protected override GraphStageLogic CreateLogic(Attributes inheritedAttributes)
             => new Logic(inheritedAttributes, this);
     }
+    
+    #nullable disable
 
     /// <summary>
     /// INTERNAL API