zio · jdegoes · May 15, 2025 · Feb 10, 2025 · Feb 10, 2025 · Mar 5, 2025
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -7,5 +7,5 @@
   ],
   "files.watcherExclude": {
     "**/target": true
-  }
+}
 }
diff --git a/benchmarks/src/main/scala/zio/BenchmarkUtil.scala b/benchmarks/src/main/scala/zio/BenchmarkUtil.scala
@@ -39,9 +39,12 @@ object BenchmarkUtil extends Runtime[Any] { self =>
     Unsafe.unsafe(implicit unsafe => rt.unsafe.run(zio).getOrThrowFiberFailure())
   }
 
+  override val unsafe = super.unsafe
+
   private object NoFiberRootsRuntime extends Runtime[Any] {
-    val environment  = Runtime.default.environment
-    val fiberRefs    = Runtime.default.fiberRefs
-    val runtimeFlags = RuntimeFlags(RuntimeFlag.CooperativeYielding, RuntimeFlag.Interruption)
+    override val unsafe = super.unsafe
+    val environment     = Runtime.default.environment
+    val fiberRefs       = Runtime.default.fiberRefs
+    val runtimeFlags    = RuntimeFlags(RuntimeFlag.CooperativeYielding, RuntimeFlag.Interruption)
   }
 }
diff --git a/benchmarks/src/main/scala/zio/PromiseBenchmarks.scala b/benchmarks/src/main/scala/zio/PromiseBenchmarks.scala
@@ -1,8 +1,11 @@
 package zio
 
 import cats.effect.kernel.Deferred
+import cats.syntax.traverse._
+import cats.instances.list._
 import cats.effect.unsafe.implicits.global
 import cats.effect.{IO => CIO}
+import cats.syntax.foldable._
 import org.openjdk.jmh.annotations.{Scope => JScope, _}
 import zio.BenchmarkUtil._
 
@@ -11,35 +14,109 @@ import java.util.concurrent.TimeUnit
 @State(JScope.Thread)
 @BenchmarkMode(Array(Mode.Throughput))
 @OutputTimeUnit(TimeUnit.SECONDS)
-@Measurement(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 3)
-@Warmup(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 3)
+@Measurement(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 10)
+@Warmup(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 10)
 @Fork(value = 3)
 class PromiseBenchmarks {
 
-  val size = 100000
+  val n            = 100000
+  val waiters: Int = 8
 
-  val ints: List[Int] = List.range(0, size)
+  def createWaitersZIO(promise: Promise[Nothing, Unit]): ZIO[Any, Nothing, Seq[Fiber[Nothing, Unit]]] =
+    ZIO.foreach(Vector.range(0, waiters))(_ => promise.await.forkDaemon)
+
+  def createWaitersCats(promise: Deferred[CIO, Unit]) =
+    List.range(0, waiters).traverse(_ => promise.get.start)
 
   @Benchmark
-  def zioPromiseAwaitDone(): Unit = {
+  def zioPromiseDoneAwait(): Unit = {
 
-    val io = ZIO.foreachDiscard(ints) { _ =>
-      Promise.make[Nothing, Unit].flatMap { promise =>
-        promise.succeed(()) *> promise.await
-      }
-    }
+    val io =
+      Promise
+        .make[Nothing, Unit]
+        .flatMap { promise =>
+          promise.done(Exit.unit) *> promise.await
+        }
+        .repeatN(n)
 
     unsafeRun(io)
   }
 
   @Benchmark
-  def catsPromiseAwaitDone(): Unit = {
+  def catsPromiseDoneAwait(): Unit = {
 
-    val io = catsForeachDiscard(List.range(1, size)) { _ =>
+    val io =
       Deferred[CIO, Unit].flatMap { promise =>
         promise.complete(()).flatMap(_ => promise.get)
+      }.replicateA_(n)
+
+    io.unsafeRunSync()
+  }
+
+  @Benchmark
+  def zioPromiseMultiAwaitDone(): Unit = {
+    val io = Promise
+      .make[Nothing, Unit]
+      .flatMap { promise =>
+        for {
+          fibers <- createWaitersZIO(promise)
+          _      <- promise.done(Exit.unit)
+          _      <- ZIO.foreachDiscard(fibers)(_.await)
+        } yield ()
       }
-    }
+      .repeatN(1023)
+
+    unsafeRun(io)
+  }
+
+  @Benchmark
+  def catsPromiseMultiAwaitDone(): Unit = {
+    val io =
+      Deferred[CIO, Unit].flatMap { promise =>
+        for {
+          fibers <- createWaitersCats(promise)
+          _      <- promise.complete(())
+          _      <- fibers.traverse_(_.join)
+        } yield ()
+      }.replicateA_(1023)
+
+    io.unsafeRunSync()
+  }
+
+  @Benchmark
+  def zioPromiseMultiAwaitMultiDone(): Unit = {
+    def createCompleters(promise: Promise[Nothing, Unit], latch: Promise[Nothing, Unit]) =
+      ZIO.foreach(Vector.range(0, waiters))(_ => (latch.await *> promise.done(Exit.unit)).forkDaemon)
+
+    val io = {
+      for {
+        latch   <- Promise.make[Nothing, Unit]
+        promise <- Promise.make[Nothing, Unit]
+        waiters <- createWaitersZIO(promise)
+        fibers  <- createCompleters(promise, latch)
+        _       <- latch.done(Exit.unit)
+        result  <- promise.await
+      } yield result
+    }.repeatN(1023)
+
+    unsafeRun(io)
+  }
+
+  @Benchmark
+  def catsPromiseMultiAwaitMultiDone(): Unit = {
+    def createCompleters(promise: Deferred[CIO, Unit], latch: Deferred[CIO, Unit]) =
+      List.range(0, waiters).traverse(_ => (latch.get *> promise.complete(())).start)
+
+    val io = {
+      for {
+        latch   <- Deferred[CIO, Unit]
+        promise <- Deferred[CIO, Unit]
+        waiters <- createWaitersCats(promise)
+        fibers  <- createCompleters(promise, latch)
+        _       <- latch.complete(())
+        result  <- promise.get
+      } yield result
+    }.replicateA_(1023)
 
     io.unsafeRunSync()
   }

diff --git a/core-tests/shared/src/test/scala/zio/FiberRuntimeSpec.scala b/core-tests/shared/src/test/scala/zio/FiberRuntimeSpec.scala
@@ -66,6 +66,40 @@ object FiberRuntimeSpec extends ZIOBaseSpec {
         }
       }
     ),
+    suite("async")(
+      test("async callback after interruption is ignored") {
+        ZIO.suspendSucceed {
+          val executed = Ref.unsafe.make(0)
+          val cb       = Ref.unsafe.make[Option[ZIO[Any, Nothing, Unit] => Unit]](None)
+          val latch    = Promise.unsafe.make[Nothing, Unit](FiberId.None)
+          val async = ZIO.async[Any, Nothing, Unit] { k =>
+            cb.unsafe.set(Some(k))
+            latch.unsafe.done(Exit.unit)
+          }
+          val increment = executed.update(_ + 1)
+          for {
+            fiber          <- async.fork
+            _              <- latch.await
+            exit           <- fiber.interrupt
+            callback       <- cb.get.some
+            state1         <- fiber.poll
+            _              <- ZIO.succeed(callback(increment))
+            state2         <- fiber.poll
+            executedBefore <- executed.get
+            _              <- ZIO.succeed(callback(increment))
+            state3         <- fiber.poll
+            executedAfter  <- executed.get
+          } yield assertTrue(
+            state1 == Some(exit),
+            state2 == Some(exit),
+            state3 == Some(exit),
+            executedBefore == 0,
+            executedAfter == 0,
+            exit.isInterrupted
+          )
+        }
+      } @@ TestAspect.nonFlaky(10)
+    ),
     suite("runtime metrics")(
       test("Failures are counted once for the fiber that caused them and exits are not") {
         val nullErrors = ZIO.foreachParDiscard(1 to 2)(_ => ZIO.attempt(throw new NullPointerException))

diff --git a/core-tests/shared/src/test/scala/zio/PromiseSpec.scala b/core-tests/shared/src/test/scala/zio/PromiseSpec.scala
@@ -7,6 +7,11 @@ object PromiseSpec extends ZIOBaseSpec {
 
   import ZIOTag._
 
+  private def empty[E, A]: Promise.internal.Pending[E, A] =
+    Promise.internal.State.empty[E, A].asInstanceOf[Promise.internal.Pending[E, A]]
+
+  val n = 10000
+
   def spec: Spec[Any, TestFailure[Any]] = suite("PromiseSpec")(
     test("complete a promise using succeed") {
       for {
@@ -120,6 +125,72 @@ object PromiseSpec extends ZIOBaseSpec {
         _ <- p.fail("failure")
         d <- p.isDone
       } yield assert(d)(isTrue)
-    } @@ zioTag(errors)
+    } @@ zioTag(errors),
+    test("waiter stack safety") {
+      for {
+        p      <- Promise.make[Nothing, Unit]
+        fibers <- ZIO.foreach(1 to n)(_ => p.await.forkDaemon)
+        _      <- p.complete(Exit.unit)
+        _      <- ZIO.foreach(fibers)(_.await)
+      } yield assertCompletes
+    },
+    suite("State")(
+      suite("add")(
+        test("stack safety") {
+          (0 to 100000).foldLeft(empty[Nothing, Unit])((acc, _) => acc.add(_ => ()))
+          assertCompletes
+        }
+      ),
+      suite("complete")(
+        test("one") {
+          var increment = 0
+          val state     = empty[Nothing, Unit].add(_ => increment += 1)
+          state.complete(ZIO.unit)
+          assert(increment)(equalTo(1))
+        },
+        test("multiple") {
+          var increment = 0
+          val state     = (0 until n).foldLeft(empty[Nothing, Unit])((acc, _) => acc.add(_ => increment += 1))
+          state.complete(ZIO.unit)
+          assert(increment)(equalTo(n))
+        }
+      ),
+      suite("remove")(
+        test("one") {
+          var increment = 0
+          val cb        = (_: IO[Nothing, Unit]) => increment += 1
+          val state     = empty[Nothing, Unit].add(cb)
+          val removed   = state.remove(cb)
+          removed.complete(ZIO.unit)
+          assert(removed)(equalTo(empty[Nothing, Unit])) &&
+          assert(increment)(equalTo(0))
+        },
+        test("multiple") {
+          var fired    = 0
+          val cb       = (_: IO[Nothing, Unit]) => ()
+          val toRemove = (_: IO[Nothing, Unit]) => fired += 1
+          val state =
+            (0 until n).foldLeft(empty[Nothing, Unit])((acc, i) => if (i < 5) acc.add(cb) else acc.add(toRemove))
+          val removed = state.remove(toRemove)
+          removed.complete(ZIO.unit)
+          assert(removed.size)(equalTo(5)) &&
+          assert(fired)(equalTo(0))
+        }
+      ),
+      suite("complete")(
+        test("one") {
+          var completed = 0
+          val state     = empty[Nothing, Unit].add(_ => completed += 1)
+          state.complete(ZIO.unit)
+          assert(completed)(equalTo(1))
+        },
+        test("multiple") {
+          var completed = List.empty[Int]
+          val state     = (0 until n).foldLeft(empty[Nothing, Unit])((acc, i) => acc.add(_ => completed = i :: completed))
+          state.complete(ZIO.unit)
+          assert(completed)(equalTo(List.range(0, n)))
+        }
+      )
+    )
   )
 }
diff --git a/core/shared/src/main/scala-2/zio/ZIOCompanionVersionSpecific.scala b/core/shared/src/main/scala-2/zio/ZIOCompanionVersionSpecific.scala
@@ -4,6 +4,7 @@ import zio.ZIO.Async
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 import java.io.IOException
+import java.util.concurrent.atomic.AtomicReference
 
 private[zio] trait ZIOCompanionVersionSpecific {
 
@@ -44,22 +45,15 @@ private[zio] trait ZIOCompanionVersionSpecific {
     blockingOn: => FiberId = FiberId.None
   )(implicit trace: Trace): ZIO[R, E, A] =
     ZIO.suspendSucceed {
-      val cancelerRef = new java.util.concurrent.atomic.AtomicReference[URIO[R, Any]](ZIO.unit)
-
-      ZIO
-        .Async[R, E, A](
-          trace,
-          { k =>
-            val result = register(k(_))
-
-            result match {
-              case Left(canceler) => cancelerRef.set(canceler); null.asInstanceOf[ZIO[R, E, A]]
-              case Right(done)    => done
-            }
-          },
-          () => blockingOn
-        )
-        .onInterrupt(cancelerRef.get())
+      val state = new AtomicReference[URIO[R, Any]](Exit.unit) with ((ZIO[R, E, A] => Unit) => ZIO[R, E, A]) {
+        def apply(k: ZIO[R, E, A] => Unit): ZIO[R, E, A] =
+          register(k(_)) match {
+            case Left(canceler) => set(canceler); null.asInstanceOf[ZIO[R, E, A]]
+            case Right(done)    => done
+          }
+      }
+
+      ZIO.Async[R, E, A](trace, state, () => blockingOn).onInterrupt(state.get())
     }
 
   /**

diff --git a/core/shared/src/main/scala-3/zio/ZIOCompanionVersionSpecific.scala b/core/shared/src/main/scala-3/zio/ZIOCompanionVersionSpecific.scala
@@ -4,6 +4,7 @@ import zio.ZIO.Async
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 import java.io.IOException
+import java.util.concurrent.atomic.AtomicReference
 import scala.annotation.targetName
 
 private[zio] transparent trait ZIOCompanionVersionSpecific {
@@ -51,22 +52,15 @@ private[zio] transparent trait ZIOCompanionVersionSpecific {
     blockingOn: => FiberId = FiberId.None
   )(implicit trace: Trace): ZIO[R, E, A] =
     ZIO.suspendSucceed {
-      val cancelerRef = new java.util.concurrent.atomic.AtomicReference[URIO[R, Any]](ZIO.unit)
-
-      ZIO
-        .Async[R, E, A](
-          trace,
-          { k =>
-            val result = register(using Unsafe)(k(_))
-
-            result match {
-              case Left(canceler) => cancelerRef.set(canceler); null.asInstanceOf[ZIO[R, E, A]]
-              case Right(done)    => done
-            }
-          },
-          () => blockingOn
-        )
-        .onInterrupt(cancelerRef.get())
+      val state = new AtomicReference[URIO[R, Any]](Exit.unit) with ((ZIO[R, E, A] => Unit) => ZIO[R, E, A]) {
+        def apply(k: ZIO[R, E, A] => Unit): ZIO[R, E, A] =
+          register(using Unsafe)(k(_)) match {
+            case Left(canceler) => set(canceler); null.asInstanceOf[ZIO[R, E, A]]
+            case Right(done)    => done
+          }
+      }
+
+      ZIO.Async[R, E, A](trace, state, () => blockingOn).onInterrupt(state.get())
     }
 
   /**
-Original file line number
+Diff line change
@@ Expand Up / @@ -7,5 +7,5 @@ @@
       ],
       "files.watcherExclude": {
         "**/target": true
-      }
+    }
     }