GroupBy pulls upstream when a substream materialization is waiting

Author: ismaelhamed

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

@@ -483,7 +483,7 @@ public void GroupBy_must_work_under_fuzzing_stress_test()
         }
 
         [Fact]
-        public void GroupBy_must_Work_if_pull_is_exercised_from_both_substream_and_main()
+        public void GroupBy_must_work_if_pull_is_exercised_from_both_substream_and_main()
         {
             this.AssertAllStagesStopped(() =>
             {
@@ -517,6 +517,63 @@ public void GroupBy_must_Work_if_pull_is_exercised_from_both_substream_and_main(
             }, Materializer);
         }
 
+        [Fact]
+        public void GroupBy_must_work_if_pull_is_exercised_from_multiple_substreams_while_downstream_is_backpressuring()
+        {
+            this.AssertAllStagesStopped(() =>
+            {
+                var upstream = this.CreatePublisherProbe<int>();
+                var downstreamMaster = this.CreateSubscriberProbe<Source<int, NotUsed>>();
+
+                Source.FromPublisher(upstream)
+                    .Via(new GroupBy<int, int>(10, element => element))
+                    .RunWith(Sink.FromSubscriber(downstreamMaster), Materializer);
+
+                var substream1 = this.CreateSubscriberProbe<int>();
+                downstreamMaster.Request(1);
+                upstream.SendNext(1);
+                downstreamMaster.ExpectNext().RunWith(Sink.FromSubscriber(substream1), Materializer);
+
+                var substream2 = this.CreateSubscriberProbe<int>();
+                downstreamMaster.Request(1);
+                upstream.SendNext(2);
+                downstreamMaster.ExpectNext().RunWith(Sink.FromSubscriber(substream2), Materializer);
+
+                substream1.Request(1);
+                substream1.ExpectNext(1);
+                substream2.Request(1);
+                substream2.ExpectNext(2);
+
+                // Both substreams pull
+                substream1.Request(1);
+                substream2.Request(1);
+
+                // Upstream sends new groups
+                upstream.SendNext(3);
+                upstream.SendNext(4);
+
+                var substream3 = this.CreateSubscriberProbe<int>();
+                var substream4 = this.CreateSubscriberProbe<int>();
+                downstreamMaster.Request(1);
+                downstreamMaster.ExpectNext().RunWith(Sink.FromSubscriber(substream3), Materializer);
+                downstreamMaster.Request(1);
+                downstreamMaster.ExpectNext().RunWith(Sink.FromSubscriber(substream4), Materializer);
+
+                substream3.Request(1);
+                substream3.ExpectNext(3);
+                substream4.Request(1);
+                substream4.ExpectNext(4);
+
+                // Cleanup, not part of the actual test
+                substream1.Cancel();
+                substream2.Cancel();
+                substream3.Cancel();
+                substream4.Cancel();
+                downstreamMaster.Cancel();
+                upstream.SendComplete();
+            }, Materializer);
+        }
+
         [Fact]
         public void GroupBy_must_work_with_random_demand()
         {
@@ -681,7 +738,7 @@ private void RandomDemand(Dictionary<int, SubFlowState> map, RandomDemandPropert
                         map[key] = new SubFlowState(state.Probe, false, null);
                     }
                     else if (props.BlockingNextElement == null)
-                     break;
+                        break;
                 }
             }
 

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

@@ -181,7 +181,7 @@ public FlattenMerge(int breadth)
     internal sealed class PrefixAndTail<T> : GraphStage<FlowShape<T, (IImmutableList<T>, Source<T, NotUsed>)>>
     {
         #region internal classes
-        
+
         private sealed class Logic : TimerGraphStageLogic, IInHandler, IOutHandler
         {
             private const string SubscriptionTimer = "SubstreamSubscriptionTimer";
@@ -205,11 +205,11 @@ public Logic(PrefixAndTail<T> stage) : base(stage.Shape)
                     Pull(_stage._in);
                     _tailSource.SetHandler(new LambdaOutHandler(onPull: () => Pull(_stage._in)));
                 });
-                
+
                 SetHandler(_stage._in, this);
                 SetHandler(_stage._out, this);
             }
-            
+
             protected internal override void OnTimer(object timerKey)
             {
                 var materializer = ActorMaterializerHelper.Downcast(Interpreter.Materializer);
@@ -360,7 +360,7 @@ public PrefixAndTail(int count)
     /// <typeparam name="TKey">TBD</typeparam>
     internal sealed class GroupBy<T, TKey> : GraphStage<FlowShape<T, Source<T, NotUsed>>>
     {
-        #region Loigc 
+        #region Logic 
 
         private sealed class Logic : TimerGraphStageLogic, IInHandler, IOutHandler
         {
@@ -370,7 +370,7 @@ private sealed class Logic : TimerGraphStageLogic, IInHandler, IOutHandler
             private readonly HashSet<SubstreamSource> _substreamsJustStarted = new HashSet<SubstreamSource>();
             private readonly Lazy<Decider> _decider;
             private TimeSpan _timeout;
-            private SubstreamSource _substreamWaitingToBePushed;
+            private Option<SubstreamSource> _substreamWaitingToBePushed = Option<SubstreamSource>.None;
             private Option<TKey> _nextElementKey = Option<TKey>.None;
             private Option<T> _nextElementValue = Option<T>.None;
             private long _nextId;
@@ -379,12 +379,12 @@ private sealed class Logic : TimerGraphStageLogic, IInHandler, IOutHandler
             public Logic(GroupBy<T, TKey> stage, Attributes inheritedAttributes) : base(stage.Shape)
             {
                 _stage = stage;
-                
+
                 _decider = new Lazy<Decider>(() =>
                 {
                     var attribute = inheritedAttributes.GetAttribute<ActorAttributes.SupervisionStrategy>(null);
                     return attribute != null ? attribute.Decider : Deciders.StoppingDecider;
-                }); 
+                });
 
                 SetHandler(_stage.In, this);
                 SetHandler(_stage.Out, this);
@@ -431,11 +431,12 @@ public void OnPush()
 
             public void OnPull()
             {
-                if (_substreamWaitingToBePushed != null)
+                if (_substreamWaitingToBePushed.HasValue)
                 {
-                    Push(_stage.Out, Source.FromGraph(_substreamWaitingToBePushed.Source));
-                    ScheduleOnce(_substreamWaitingToBePushed.Key.Value, _timeout);
-                    _substreamWaitingToBePushed = null;
+                    var substreamSource = _substreamWaitingToBePushed.Value;
+                    Push(_stage.Out, Source.FromGraph(substreamSource.Source));
+                    ScheduleOnce(substreamSource.Key.Value, _timeout);
+                    _substreamWaitingToBePushed = Option<SubstreamSource>.None;
                 }
                 else
                 {
@@ -500,7 +501,7 @@ private void Fail(Exception ex)
                 FailStage(ex);
             }
 
-            private bool NeedToPull => !(HasBeenPulled(_stage.In) || IsClosed(_stage.In) || HasNextElement);
+            private bool NeedToPull => !(HasBeenPulled(_stage.In) || IsClosed(_stage.In) || HasNextElement || _substreamWaitingToBePushed.HasValue);
 
             public override void PreStart()
             {
@@ -530,7 +531,7 @@ private void RunSubstream(TKey key, T value)
                 {
                     Push(_stage.Out, Source.FromGraph(substreamSource.Source));
                     ScheduleOnce(key, _timeout);
-                    _substreamWaitingToBePushed = null;
+                    _substreamWaitingToBePushed = Option<SubstreamSource>.None;
                 }
                 else
                 {
@@ -628,7 +629,7 @@ public GroupBy(int maxSubstreams, Func<T, TKey> keyFor)
         {
             _maxSubstreams = maxSubstreams;
             _keyFor = keyFor;
-            
+
             Shape = new FlowShape<T, Source<T, NotUsed>>(In, Out);
         }
 
@@ -778,7 +779,7 @@ public override void OnDownstreamFinish()
                     else
                         // Start draining
                         if (!_logic.HasBeenPulled(_inlet))
-                            _logic.Pull(_inlet);
+                        _logic.Pull(_inlet);
                 }
 
                 public override void OnPush()
@@ -1010,7 +1011,7 @@ protected CommandScheduledBeforeMaterialization(ICommand command)
         internal class RequestOneScheduledBeforeMaterialization : CommandScheduledBeforeMaterialization
         {
             public static readonly RequestOneScheduledBeforeMaterialization Instance = new RequestOneScheduledBeforeMaterialization(RequestOne.Instance);
-            
+
             private RequestOneScheduledBeforeMaterialization(ICommand command) : base(command)
             {
             }
@@ -1046,7 +1047,7 @@ private RequestOne()
             {
             }
         }
-        
+
         internal class Cancel : ICommand
         {
             public static readonly Cancel Instance = new Cancel();