optimize ZStream's mapZIOPar and mapZIOParUnordered

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

@@ -86,14 +86,32 @@ final class FiberRuntime[E, A](fiberId: FiberId.Runtime, fiberRefs0: FiberRefs,
         )
     }
 
+  private def childrenChunk = if (_children eq null) Chunk.empty
+  else {
+    val bldr = Chunk.newBuilder[Fiber.Runtime[_, _]]
+    _children.forEach { child =>
+      if ((child ne null) && child.isAlive())
+        bldr.addOne(child)
+    }
+    bldr.result()
+  }
+
   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)
+    ZIO.withFiberRuntime[Any, Nothing, Chunk[Fiber.Runtime[_, _]]] { case (fib, _) =>
+      if (fib eq self)
+        ZIO.succeed(self.childrenChunk)
+      else {
+        ZIO.asyncZIO[Any, Nothing, Chunk[Fiber.Runtime[_, _]]] { k =>
+          ZIO.succeed {
+            self.tell {
+              FiberMessage.Stateful { case (fib, _) =>
+                val childs = fib.childrenChunk
+                k(Exit.succeed(childs))
+              }
+            }
+          }
         }
+      }
     }
 
   def fiberRefs(implicit trace: Trace): UIO[FiberRefs] = ZIO.succeed(_fiberRefs)
@@ -482,9 +500,9 @@ 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[_, _]] = {
     if (_children eq null) {
-      _children = Platform.newConcurrentWeakSet[Fiber.Runtime[_, _]]()(Unsafe.unsafe)
+      _children = Platform.newWeakSet[Fiber.Runtime[_, _]]()(Unsafe.unsafe)
     }
     _children
   }
@@ -651,18 +669,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 =
@@ -842,7 +871,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)
       ()
@@ -1169,16 +1198,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 =

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") {