From 7ffac90eb0b2236f389e311221135b8f32b4231e Mon Sep 17 00:00:00 2001
From: Kyri Petrou <kyri.petrou@outlook.com>
Date: Fri, 13 Dec 2024 21:02:16 +0530
Subject: [PATCH] Ensure that TestClock methods are thread-safe

---
 .../test/environment/TestClockSpecJVM.scala   |  9 +++++++--
 .../src/main/scala/zio/test/TestClock.scala   | 19 +++++++++++++------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/test-tests/jvm-native/src/test/scala/zio/test/environment/TestClockSpecJVM.scala b/test-tests/jvm-native/src/test/scala/zio/test/environment/TestClockSpecJVM.scala
index d044afbd43ff..df86fe01bf0c 100644
--- a/test-tests/jvm-native/src/test/scala/zio/test/environment/TestClockSpecJVM.scala
+++ b/test-tests/jvm-native/src/test/scala/zio/test/environment/TestClockSpecJVM.scala
@@ -2,8 +2,8 @@ package zio.test.environment
 
 import zio._
 import zio.test.Assertion._
-import zio.test.TestAspect.nonFlaky
-import zio.test.{TestClock, ZIOBaseSpec, assert}
+import zio.test.TestAspect.{nonFlaky, timeout}
+import zio.test.{TestClock, ZIOBaseSpec, assert, assertCompletes}
 
 import java.util.concurrent.TimeUnit
 
@@ -126,6 +126,11 @@ object TestClockSpecJVM extends ZIOBaseSpec {
             _      <- ZIO.logInfo(s"Values after interruption: $values")
           } yield assert(values.reverse)(equalTo(List(5L)))
         }
+      ),
+      suite("adjust")(
+        test("is thread safe") {
+          (TestClock.adjust(1.second) &> TestClock.adjust(1.second)).as(assertCompletes)
+        } @@ timeout(10.seconds)
       )
     ) @@ nonFlaky(20)
 }
diff --git a/test/shared/src/main/scala/zio/test/TestClock.scala b/test/shared/src/main/scala/zio/test/TestClock.scala
index 3b7ec133b112..9810094b82e4 100644
--- a/test/shared/src/main/scala/zio/test/TestClock.scala
+++ b/test/shared/src/main/scala/zio/test/TestClock.scala
@@ -104,6 +104,11 @@ object TestClock extends Serializable {
   ) extends TestClock
       with TestClockPlatformSpecific {
 
+    /**
+     * See https://github.com/zio/zio/issues/9385
+     */
+    private[this] val freezeLock = Semaphore.unsafe.make(1)(Unsafe)
+
     /**
      * Increments the current clock time by the specified duration. Any effects
      * that were scheduled to occur on or before the new time will be run in
@@ -294,12 +299,14 @@ object TestClock extends Serializable {
      * snapshot may not be fully consistent.
      */
     private def freeze(implicit trace: Trace): IO[Unit, Map[FiberId, Fiber.Status]] =
-      supervisedFibers.flatMap { fibers =>
-        ZIO.foldLeft(fibers)(Map.empty[FiberId, Fiber.Status]) { (map, fiber) =>
-          fiber.status.flatMap {
-            case done @ Fiber.Status.Done                    => ZIO.succeed(map.updated(fiber.id, done))
-            case suspended @ Fiber.Status.Suspended(_, _, _) => ZIO.succeed(map.updated(fiber.id, suspended))
-            case _                                           => ZIO.fail(())
+      freezeLock.withPermit {
+        supervisedFibers.flatMap { fibers =>
+          ZIO.foldLeft(fibers)(Map.empty[FiberId, Fiber.Status]) { (map, fiber) =>
+            fiber.status.flatMap {
+              case done @ Fiber.Status.Done                    => ZIO.succeed(map.updated(fiber.id, done))
+              case suspended @ Fiber.Status.Suspended(_, _, _) => ZIO.succeed(map.updated(fiber.id, suspended))
+              case _                                           => ZIO.fail(())
+            }
           }
         }
       }