Support Supervision.restart for SubFlow's

apache · Jan 19, 2024 · e555122 · e555122
1 parent 19da736
commit e555122
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/StreamOfStreams.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/fusing/StreamOfStreams.scala
@@ -511,13 +511,6 @@ import pekko.util.ccompat.JavaConverters._
     private var substreamWaitingToBePushed = false
     private var substreamCancelled = false
 
-    def propagateSubstreamCancel(ex: Throwable): Boolean =
-      decider(ex) match {
-        case Supervision.Stop    => true
-        case Supervision.Resume  => false
-        case Supervision.Restart => false
-      }
-
     setHandler(
       out,
       new OutHandler {
@@ -614,14 +607,31 @@ import pekko.util.ccompat.JavaConverters._
 
       override def onDownstreamFinish(cause: Throwable): Unit = {
         substreamCancelled = true
-        if (isClosed(in) || propagateSubstreamCancel(cause)) {
-          cancelStage(cause)
-        } else {
-          // Start draining
-          if (!hasBeenPulled(in)) pull(in)
+        decider(cause) match {
+          case Supervision.Stop =>
+            cancelStage(cause)
+          case Supervision.Resume =>
+            if (isClosed(in)) cancelStage(cause)
+            else {
+              // Start draining
+              if (!hasBeenPulled(in)) pull(in)
+            }
+          case Supervision.Restart =>
+            if (isClosed(in)) completeStage()
+            else {
+              restartState()
+              // Start draining
+              if (!hasBeenPulled(in)) pull(in)
+            }
         }
       }
 
+      private def restartState(): Unit = {
+        substreamSource = null
+        substreamWaitingToBePushed = false
+        substreamCancelled = false
+      }
+
       override def onPush(): Unit = {
         val elem = grab(in)
         try {