From a2a55592b2e40b374d0bf15f4aa9d2255bbb844b Mon Sep 17 00:00:00 2001
From: Adam Fraser <adam.fraser@gmail.com>
Date: Thu, 22 Jun 2023 16:58:22 -0700
Subject: [PATCH] join fiber refs in order

---
 core-tests/shared/src/test/scala/zio/ZIOSpec.scala | 10 ++++++++++
 core/shared/src/main/scala/zio/ZIO.scala           | 12 +++++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/ZIOSpec.scala b/core-tests/shared/src/test/scala/zio/ZIOSpec.scala
index f707564a9a29..0e4e24d4b163 100644
--- a/core-tests/shared/src/test/scala/zio/ZIOSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/ZIOSpec.scala
@@ -4193,6 +4193,16 @@ object ZIOSpec extends ZIOBaseSpec {
       test("propagates interruption") {
         val zio = ZIO.never <&> ZIO.never <&> ZIO.fail("fail")
         assertZIO(zio.exit)(fails(equalTo("fail")))
+      },
+      test("propagates FiberRef values") {
+        for {
+          fiberRef <- FiberRef.make(5)
+          workflow  = fiberRef.set(10).delay(2.seconds) <&> fiberRef.set(15)
+          fiber    <- workflow.fork
+          _        <- TestClock.adjust(2.seconds)
+          _        <- fiber.join
+          value    <- fiberRef.get
+        } yield assertTrue(value == 10)
       }
     ),
     suite("toFuture")(
diff --git a/core/shared/src/main/scala/zio/ZIO.scala b/core/shared/src/main/scala/zio/ZIO.scala
index 693a1980b39d..b5eff89c407b 100644
--- a/core/shared/src/main/scala/zio/ZIO.scala
+++ b/core/shared/src/main/scala/zio/ZIO.scala
@@ -2521,16 +2521,16 @@ sealed trait ZIO[-R, +E, +A]
   )(f: (A, B) => C)(implicit trace: Trace): ZIO[R1, E1, C] =
     ZIO.uninterruptibleMask { restore =>
       ZIO.transplant { graft =>
-        val promise = Promise.unsafe.make[Unit, Unit](FiberId.None)(Unsafe.unsafe)
+        val promise = Promise.unsafe.make[Unit, Boolean](FiberId.None)(Unsafe.unsafe)
         val ref     = new java.util.concurrent.atomic.AtomicBoolean(false)
 
-        def fork[R, E, A](zio: => ZIO[R, E, A]): ZIO[R, Nothing, Fiber[E, A]] =
+        def fork[R, E, A](zio: => ZIO[R, E, A], side: Boolean): ZIO[R, Nothing, Fiber[E, A]] =
           graft(restore(zio))
             .foldCauseZIO(
               cause => promise.fail(()) *> ZIO.refailCause(cause),
               a =>
                 if (ref.getAndSet(true)) {
-                  promise.unsafe.done(ZIO.unit)(Unsafe.unsafe)
+                  promise.unsafe.done(ZIO.succeedNow(side))(Unsafe.unsafe)
                   ZIO.succeed(a)
                 } else {
                   ZIO.succeed(a)
@@ -2538,7 +2538,7 @@ sealed trait ZIO[-R, +E, +A]
             )
             .forkDaemon
 
-        fork(self).zip(fork(that)).flatMap { case (left, right) =>
+        fork(self, false).zip(fork(that, true)).flatMap { case (left, right) =>
           restore(promise.await).foldCauseZIO(
             cause =>
               left.interruptFork *> right.interruptFork *>
@@ -2548,7 +2548,9 @@ sealed trait ZIO[-R, +E, +A]
                     case _                    => ZIO.refailCause(cause.stripFailures)
                   }
                 },
-            _ => left.join.zipWith(right.join)(f)
+            leftWins =>
+              if (leftWins) left.join.zipWith(right.join)((a, b) => f(a, b))
+              else right.join.zipWith(left.join)((b, a) => f(a, b))
           )
         }
       }