optimize ZStream's mapZIOPar and mapZIOParUnordered

core/shared/src/main/scala/zio/Fiber.scala

@@ -530,6 +530,7 @@ object Fiber extends FiberPlatformSpecific {
      * '''NOTE''': This method must be invoked by the fiber itself.
      */
     private[zio] def addChild(child: Fiber.Runtime[_, _]): Unit
+    private[zio] def addChildren(children: Iterable[Fiber.Runtime[_, _]]): Unit
 
     /**
      * Deletes the specified fiber ref.
@@ -619,6 +620,7 @@ object Fiber extends FiberPlatformSpecific {
      * Adds a message to add a child to this fiber.
      */
     private[zio] def tellAddChild(child: Fiber.Runtime[_, _]): Unit
+    private[zio] def tellAddChildren(children: Iterable[Fiber.Runtime[_, _]]): Unit
 
     /**
      * Adds a message to interrupt this fiber.

core/shared/src/main/scala/zio/internal/FiberRuntime.scala

@@ -87,15 +87,22 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
         )
     }
 
-  def children(implicit trace: Trace): UIO[Chunk[Fiber.Runtime[_, _]]] =
-    ZIO.succeed {
-      val childs = _children
-      if (childs == null) Chunk.empty
-      else
-        zio.internal.Sync(childs) {
-          Chunk.fromJavaIterable(childs)
-        }
+  private def childrenChunk = {
+    //may be executed by a foreign fiber (under Sync), hence we're risking a race over the _children variable being set back to null by a concurrent transferChildren call
+    val childs = _children
+    if (childs eq null) Chunk.empty
+    else {
+      val bldr = Chunk.newBuilder[Fiber.Runtime[_, _]]
+      childs.forEach { child =>
+        if ((child ne null) && child.isAlive())
+          bldr.addOne(child)
+      }
+      bldr.result()
     }
+  }
+
+  def children(implicit trace: Trace): UIO[Chunk[Fiber.Runtime[_, _]]] =
+    ZIO.succeed(self.childrenChunk)
 
   def fiberRefs(implicit trace: Trace): UIO[FiberRefs] = ZIO.succeed(_fiberRefs)
 
@@ -151,14 +158,47 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
     }
 
   private[zio] def addChild(child: Fiber.Runtime[_, _]): Unit =
-    if (isAlive()) {
-      getChildren().add(child)
+    if (child.isAlive()) {
+      if (isAlive()) {
+        getChildren().add(child)
 
-      if (isInterrupted())
+        if (isInterrupted())
+          child.tellInterrupt(getInterruptedCause())
+      } else {
         child.tellInterrupt(getInterruptedCause())
+      }
+    }
+
+  private[zio] def addChildren(children: Iterable[Fiber.Runtime[_, _]]): Unit = {
+    val iter = children.iterator
+    if (isAlive()) {
+      val childs = getChildren()
+
+      if (isInterrupted()) {
+        val cause = getInterruptedCause()
+        while (iter.hasNext) {
+          val child = iter.next()
+          if (child.isAlive()) {
+            childs.add(child)
+            child.tellInterrupt(cause)
+          }
+        }
+      } else {
+        while (iter.hasNext) {
+          val child = iter.next()
+          if (child.isAlive())
+            childs.add(child)
+        }
+      }
     } else {
-      child.tellInterrupt(getInterruptedCause())
+      val cause = getInterruptedCause()
+      while (iter.hasNext) {
+        val child = iter.next()
+        if (child.isAlive())
+          child.tellInterrupt(cause)
+      }
     }
+  }
 
   /**
    * Adds an interruptor to the set of interruptors that are interrupting this
@@ -483,7 +523,8 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
    *
    * '''NOTE''': This method must be invoked by the fiber itself.
    */
-  private[zio] def getChildren(): JavaSet[Fiber.Runtime[_, _]] = {
+  private def getChildren(): JavaSet[Fiber.Runtime[_, _]] = {
+    //executed by the fiber itself, no risk of racing with transferChildren
     if (_children eq null) {
       _children = Platform.newConcurrentWeakSet[Fiber.Runtime[_, _]]()(Unsafe.unsafe)
     }
@@ -652,18 +693,29 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
   private def interruptAllChildren(): UIO[Any] =
     if (sendInterruptSignalToAllChildren(_children)) {
       val iterator = _children.iterator()
-
       _children = null
 
-      val body = () => {
-        val next = iterator.next()
+      var curr: Fiber.Runtime[_, _] = null
 
-        if (next != null) next.await(id.location) else Exit.unit
+      //this finds the next operable child fiber and stores it in the `curr` variable
+      def skip() = {
+        var next: Fiber.Runtime[_, _] = null
+        while (iterator.hasNext && (next eq null)) {
+          next = iterator.next()
+          if ((next ne null) && !next.isAlive())
+            next = null
+        }
+        curr = next
       }
 
-      // Now await all children to finish:
-      ZIO
-        .whileLoop(iterator.hasNext)(body())(_ => ())(id.location)
+      //find the first operable child fiber
+      //if there isn't any we can simply return null and save ourselves an effect evaluation
+      skip()
+
+      if (null ne curr) {
+        ZIO
+          .whileLoop(null ne curr)(curr.await(id.location))(_ => skip())(id.location)
+      } else null
     } else null
 
   private[zio] def isAlive(): Boolean =
@@ -843,7 +895,7 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
    *
    * '''NOTE''': This method must be invoked by the fiber itself.
    */
-  private[zio] def removeChild(child: FiberRuntime[_, _]): Unit =
+  private def removeChild(child: FiberRuntime[_, _]): Unit =
     if (_children ne null) {
       _children.remove(child)
       ()
@@ -1170,16 +1222,6 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
     done
   }
 
-  private def sendInterruptSignalToAllChildrenConcurrently(): Boolean = {
-    val childFibers = _children
-
-    if (childFibers ne null) {
-      internal.Sync(childFibers) {
-        sendInterruptSignalToAllChildren(childFibers)
-      }
-    } else false
-  }
-
   private def sendInterruptSignalToAllChildren(
     children: JavaSet[Fiber.Runtime[_, _]]
   ): Boolean =
@@ -1375,6 +1417,9 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
   private[zio] def tellAddChild(child: Fiber.Runtime[_, _]): Unit =
     tell(FiberMessage.Stateful((parentFiber, _) => parentFiber.addChild(child)))
 
+  private[zio] def tellAddChildren(children: Iterable[Fiber.Runtime[_, _]]): Unit =
+    tell(FiberMessage.Stateful((parentFiber, _) => parentFiber.addChildren(children)))
+
   private[zio] def tellInterrupt(cause: Cause[Nothing]): Unit =
     tell(FiberMessage.InterruptSignal(cause))
 
@@ -1385,24 +1430,16 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
    * '''NOTE''': This method must be invoked by the fiber itself after it has
    * evaluated the effects but prior to exiting
    */
-  private[zio] def transferChildren(scope: FiberScope): Unit = {
-    val children = _children
-    if ((children ne null) && !children.isEmpty) {
-      val iterator = children.iterator()
-      val flags    = _runtimeFlags
-
-      while (iterator.hasNext) {
-        val next = iterator.next()
-
-        // Only move alive children.
-        // Unless we forked fibers and didn't await them, we shouldn't have any alive children in the set.
-        if ((next ne null) && next.isAlive()) {
-          scope.add(self, flags, next)(location, Unsafe.unsafe)
-          iterator.remove()
-        }
-      }
+  private[zio] def transferChildren(scope: FiberScope): Unit =
+    if ((_children ne null) && !_children.isEmpty) {
+      val childs = childrenChunk
+      //we're effectively clearing this set, seems cheaper to 'drop' it and allocate a new one if we spawn more fibers
+      //a concurrent children call might get the stale set, but this method (and its primary usage for dumping fibers)
+      //is racy by definition
+      _children = null
+      val flags = _runtimeFlags
+      scope.addAll(self, flags, childs)(location, Unsafe.unsafe)
     }
-  }
 
   /**
    * Updates a fiber ref belonging to this fiber by using the provided update

core/shared/src/main/scala/zio/internal/FiberScope.scala

@@ -36,6 +36,15 @@ private[zio] sealed trait FiberScope {
     trace: Trace,
     unsafe: Unsafe
   ): Unit
+
+  private[zio] def addAll(
+    currentFiber: Fiber.Runtime[_, _],
+    runtimeFlags: RuntimeFlags,
+    children: Iterable[Fiber.Runtime[_, _]]
+  )(implicit
+    trace: Trace,
+    unsafe: Unsafe
+  ): Unit
 }
 
 private[zio] object FiberScope {
@@ -56,6 +65,20 @@ private[zio] object FiberScope {
       if (RuntimeFlags.fiberRoots(runtimeFlags)) {
         Fiber._roots.add(child)
       }
+
+    private[zio] def addAll(
+      currentFiber: Fiber.Runtime[_, _],
+      runtimeFlags: RuntimeFlags,
+      children: Iterable[Fiber.Runtime[_, _]]
+    )(implicit
+      trace: Trace,
+      unsafe: Unsafe
+    ): Unit =
+      if (RuntimeFlags.fiberRoots(runtimeFlags)) {
+        children.foreach {
+          Fiber._roots.add(_)
+        }
+      }
   }
 
   private final class Local(val fiberId: FiberId, parentRef: WeakReference[Fiber.Runtime[_, _]]) extends FiberScope {
@@ -85,6 +108,37 @@ private[zio] object FiberScope {
         child.tellInterrupt(Cause.interrupt(currentFiber.id))
       }
     }
+
+    private[zio] def addAll(
+      currentFiber: Fiber.Runtime[_, _],
+      runtimeFlags: RuntimeFlags,
+      children: Iterable[Fiber.Runtime[_, _]]
+    )(implicit
+      trace: Trace,
+      unsafe: Unsafe
+    ): Unit = if (children.nonEmpty) {
+      val parent = parentRef.get()
+
+      if (parent ne null) {
+        // Parent is not GC'd. Let's check to see if the parent is the current
+        // fiber:
+        if (currentFiber eq parent) {
+          // The parent is the current fiber so it is safe to directly add the
+          // child to the parent:
+          parent.addChildren(children)
+        } else {
+          // The parent is not the current fiber. So we need to send a message
+          // to the parent so it will add the child to itself:
+          parent.tellAddChildren(children)
+        }
+      } else {
+        // Parent was GC'd. We immediately interrupt the child fiber using the id
+        // of the current fiber (which is adding the child to the parent):
+        children.foreach(
+          _.tellInterrupt(Cause.interrupt(currentFiber.id))
+        )
+      }
+    }
   }
 
   private[zio] def make(fiber: FiberRuntime[_, _]): FiberScope =

streams-tests/shared/src/test/scala/zio/stream/ZStreamSpec.scala

@@ -2677,7 +2677,58 @@ object ZStreamSpec extends ZIOBaseSpec {
             val stream =
               ZStream.fromIterable(0 to 3).mapZIOParUnordered(10)(_ => ZIO.fail("fail"))
             assertZIO(stream.runDrain.exit)(fails(equalTo("fail")))
-          } @@ nonFlaky @@ TestAspect.diagnose(10.seconds)
+          } @@ nonFlaky @@ TestAspect.diagnose(10.seconds),
+          test("interruption propagation") {
+            for {
+              interrupted <- Ref.make(false)
+              latch       <- Promise.make[Nothing, Unit]
+              fib <-
+                ZStream(())
+                  .mapZIOParUnordered(1)(_ => (latch.succeed(()) *> ZIO.infinity).onInterrupt(interrupted.set(true)))
+                  .runDrain
+                  .fork
+              _      <- latch.await
+              _      <- fib.interrupt
+              result <- interrupted.get
+            } yield assert(result)(isTrue)
+          },
+          test("interrupts pending tasks when one of the tasks fails U") {
+            for {
+              interrupted <- Ref.make(0)
+              latch1      <- Promise.make[Nothing, Unit]
+              latch2      <- Promise.make[Nothing, Unit]
+              result <- ZStream(1, 2, 3)
+                          .mapZIOParUnordered(3) {
+                            case 1 => (latch1.succeed(()) *> ZIO.never).onInterrupt(interrupted.update(_ + 1))
+                            case 2 => (latch2.succeed(()) *> ZIO.never).onInterrupt(interrupted.update(_ + 1))
+                            case 3 => latch1.await *> latch2.await *> ZIO.fail("Boom")
+                          }
+                          .runDrain
+                          .exit
+              count <- interrupted.get
+            } yield assert(count)(equalTo(2)) && assert(result)(fails(equalTo("Boom")))
+          } @@ nonFlaky,
+          test("awaits children fibers properly") {
+            assertZIO(
+              ZStream
+                .fromIterable((0 to 100))
+                .interruptWhen(ZIO.never)
+                .mapZIOParUnordered(8)(_ => ZIO.succeed(1).repeatN(2000))
+                .runDrain
+                .exit
+                .map(_.isInterrupted)
+            )(equalTo(false))
+          },
+          test("propagates error of original stream") {
+            for {
+              fiber <- (ZStream(1, 2, 3, 4, 5, 6, 7, 8, 9, 10) ++ ZStream.fail(new Throwable("Boom")))
+                         .mapZIOParUnordered(2)(_ => ZIO.sleep(1.second))
+                         .runDrain
+                         .fork
+              _    <- TestClock.adjust(5.seconds)
+              exit <- fiber.await
+            } yield assert(exit)(fails(hasMessage(equalTo("Boom"))))
+          }
         ),
         suite("mergeLeft/Right")(
           test("mergeLeft with HaltStrategy.Right terminates as soon as the right stream terminates") {