From 696f3b5e3d16ca28663c8507f4212df71cf9465b Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igor.doroxov@gmail.com>
Date: Fri, 14 Feb 2025 11:19:29 -0500
Subject: [PATCH 01/12] 9390 add tryAcquire and tryAcquireN to TSemaphore

---
 .../test/scala/zio/stm/TSemaphoreSpec.scala   | 47 +++++++++++++++++++
 .../src/main/scala/zio/stm/TSemaphore.scala   | 21 +++++++++
 2 files changed, 68 insertions(+)

diff --git a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
index c3c1d4cf01cb..f10ffa7defa8 100644
--- a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
@@ -137,6 +137,53 @@ object TSemaphoreSpec extends ZIOBaseSpec {
           assertTrue(remaining == 3L)
         }
       }
+    ),
+    suite("tryAcquire and tryAcquireN")(
+      test("tryAcquire should succeed when a permit is available") {
+        for {
+          sem <- TSemaphore.makeCommit(1L)
+          res <- sem.tryAcquire.commit
+        } yield assert(res)(isTrue)
+      },
+      test("tryAcquire should fail when no permits are available") {
+        for {
+          sem <- TSemaphore.makeCommit(0L)
+          res <- sem.tryAcquire.commit
+        } yield assert(res)(isFalse)
+      },
+      test("tryAcquire should decrease the permit count when successful") {
+        for {
+          sem    <- TSemaphore.makeCommit(1L)
+          _      <- sem.tryAcquire.commit
+          avail  <- sem.available.commit
+        } yield assert(avail)(equalTo(0L))
+      },
+      test("tryAcquireN should acquire permits if enough are available") {
+        for {
+          sem <- TSemaphore.makeCommit(5L)
+          res <- sem.tryAcquireN(3L).commit
+        } yield assert(res)(isTrue)
+      },
+      test("tryAcquireN should fail if not enough permits are available") {
+        for {
+          sem <- TSemaphore.makeCommit(2L)
+          res <- sem.tryAcquireN(3L).commit
+        } yield assert(res)(isFalse)
+      },
+      test("tryAcquireN should decrease the permit count when successful") {
+        for {
+          sem    <- TSemaphore.makeCommit(5L)
+          _      <- sem.tryAcquireN(3L).commit
+          avail  <- sem.available.commit
+        } yield assert(avail)(equalTo(2L))
+      },
+      test("tryAcquireN should not change permit count when unsuccessful") {
+        for {
+          sem    <- TSemaphore.makeCommit(2L)
+          _      <- sem.tryAcquireN(3L).commit
+          avail  <- sem.available.commit
+        } yield assert(avail)(equalTo(2L))
+      }
     )
   )
 
diff --git a/core/shared/src/main/scala/zio/stm/TSemaphore.scala b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
index 5731cf404e7b..975e5918f7f5 100644
--- a/core/shared/src/main/scala/zio/stm/TSemaphore.scala
+++ b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
@@ -60,6 +60,27 @@ final class TSemaphore private (val permits: TRef[Long]) extends Serializable {
   def acquireN(n: Long): USTM[Unit] =
     acquireBetween(n, n).unit
 
+  /**
+   * Tries to acquire a single permit in a transactional context.
+   * Returns `true` if the permit was acquired, otherwise `false`.
+   */
+  def tryAcquire: USTM[Boolean] = tryAcquireN(1L)
+
+  /**
+   * Tries to acquire the specified number of permits in a transactional context.
+   * Returns `true` if the permits were acquired, otherwise `false`.
+   */
+  def tryAcquireN(n: Long): USTM[Boolean] =
+    ZSTM.Effect{ (journal, _, _) =>
+      assertNonNegative(n)
+
+      val available: Long = permits.unsafeGet(journal)
+      if(available >= n){
+        permits.unsafeSet(journal, available - n)
+        true
+      } else false
+    }
+
   /**
    * Acquire at least `min` permits and at most `max` permits in a transactional
    * context.

From 994212276ec3360a501961bc33c4fbc5042eb26d Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igor.doroxov@gmail.com>
Date: Fri, 14 Feb 2025 11:42:37 -0500
Subject: [PATCH 02/12] run scalafmt

---
 .../test/scala/zio/stm/TSemaphoreSpec.scala    | 18 +++++++++---------
 .../src/main/scala/zio/stm/TSemaphore.scala    | 12 ++++++------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
index f10ffa7defa8..f868dde9d41c 100644
--- a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
@@ -153,9 +153,9 @@ object TSemaphoreSpec extends ZIOBaseSpec {
       },
       test("tryAcquire should decrease the permit count when successful") {
         for {
-          sem    <- TSemaphore.makeCommit(1L)
-          _      <- sem.tryAcquire.commit
-          avail  <- sem.available.commit
+          sem   <- TSemaphore.makeCommit(1L)
+          _     <- sem.tryAcquire.commit
+          avail <- sem.available.commit
         } yield assert(avail)(equalTo(0L))
       },
       test("tryAcquireN should acquire permits if enough are available") {
@@ -172,16 +172,16 @@ object TSemaphoreSpec extends ZIOBaseSpec {
       },
       test("tryAcquireN should decrease the permit count when successful") {
         for {
-          sem    <- TSemaphore.makeCommit(5L)
-          _      <- sem.tryAcquireN(3L).commit
-          avail  <- sem.available.commit
+          sem   <- TSemaphore.makeCommit(5L)
+          _     <- sem.tryAcquireN(3L).commit
+          avail <- sem.available.commit
         } yield assert(avail)(equalTo(2L))
       },
       test("tryAcquireN should not change permit count when unsuccessful") {
         for {
-          sem    <- TSemaphore.makeCommit(2L)
-          _      <- sem.tryAcquireN(3L).commit
-          avail  <- sem.available.commit
+          sem   <- TSemaphore.makeCommit(2L)
+          _     <- sem.tryAcquireN(3L).commit
+          avail <- sem.available.commit
         } yield assert(avail)(equalTo(2L))
       }
     )
diff --git a/core/shared/src/main/scala/zio/stm/TSemaphore.scala b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
index 975e5918f7f5..96b536c6f913 100644
--- a/core/shared/src/main/scala/zio/stm/TSemaphore.scala
+++ b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
@@ -61,21 +61,21 @@ final class TSemaphore private (val permits: TRef[Long]) extends Serializable {
     acquireBetween(n, n).unit
 
   /**
-   * Tries to acquire a single permit in a transactional context.
-   * Returns `true` if the permit was acquired, otherwise `false`.
+   * Tries to acquire a single permit in a transactional context. Returns `true`
+   * if the permit was acquired, otherwise `false`.
    */
   def tryAcquire: USTM[Boolean] = tryAcquireN(1L)
 
   /**
-   * Tries to acquire the specified number of permits in a transactional context.
-   * Returns `true` if the permits were acquired, otherwise `false`.
+   * Tries to acquire the specified number of permits in a transactional
+   * context. Returns `true` if the permits were acquired, otherwise `false`.
    */
   def tryAcquireN(n: Long): USTM[Boolean] =
-    ZSTM.Effect{ (journal, _, _) =>
+    ZSTM.Effect { (journal, _, _) =>
       assertNonNegative(n)
 
       val available: Long = permits.unsafeGet(journal)
-      if(available >= n){
+      if (available >= n) {
         permits.unsafeSet(journal, available - n)
         true
       } else false

From cae94b3e06011228b852b864c92d6e52f14318a7 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igor.doroxov@gmail.com>
Date: Fri, 14 Feb 2025 15:37:10 -0500
Subject: [PATCH 03/12] add tryAcquire to Semaphore

---
 .../src/test/scala/zio/SemaphoreSpec.scala    | 32 +++++++++++++++++++
 .../shared/src/main/scala/zio/Semaphore.scala | 21 ++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index 8fd5f58c48b0..375cfd4ac6b1 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -32,6 +32,38 @@ object SemaphoreSpec extends ZIOBaseSpec {
         permits   <- semaphore.available
       } yield assertTrue(permits == 2L)
     },
+    test("tryAcquire should succeed when a permit is available") {
+      for {
+        sem <- Semaphore.make(1L)
+        res <- sem.tryAcquire
+      } yield assert(res)(isTrue)
+    },
+    test("tryAcquireN should acquire permits if enough are available") {
+      for {
+        sem <- Semaphore.make(5L)
+        res <- sem.tryAcquireN(3L)
+      } yield assert(res)(isTrue)
+    },
+    test("tryAcquireN should fail if not enough permits are available") {
+      for {
+        sem <- Semaphore.make(2L)
+        res <- sem.tryAcquireN(3L)
+      } yield assert(res)(isFalse)
+    },
+    test("tryAcquireN should decrease the permit count when successful") {
+      for {
+        sem   <- Semaphore.make(5L)
+        _     <- sem.tryAcquireN(3L)
+        avail <- sem.available
+      } yield assert(avail)(equalTo(2L))
+    },
+    test("tryAcquireN should not change permit count when unsuccessful") {
+      for {
+        sem   <- Semaphore.make(2L)
+        _     <- sem.tryAcquireN(3L)
+        avail <- sem.available
+      } yield assert(avail)(equalTo(2L))
+    },
     test("awaiting returns the count of waiting fibers") {
       for {
         semaphore    <- Semaphore.make(1)
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index fbad74f692c1..bf9be7011b72 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -39,6 +39,19 @@ sealed trait Semaphore extends Serializable {
    */
   def available(implicit trace: Trace): UIO[Long]
 
+  /**
+   * Attempts to acquire a permit without blocking. Returns `true` if the permit
+   * was acquired, otherwise `false`.
+   */
+  def tryAcquire(implicit trace: Trace): UIO[Boolean] =
+    tryAcquireN(1L)
+
+  /**
+   * Attempts to acquire the specified number of permits without blocking.
+   * Returns `true` if the permits were acquired, otherwise `false`.
+   */
+  def tryAcquireN(n: Long)(implicit trace: Trace): UIO[Boolean] = ZIO.succeed(false)
+
   /**
    * Returns the number of tasks currently waiting for permits. The default
    * implementation returns 0.
@@ -98,6 +111,14 @@ object Semaphore {
             case Right(_)    => 0L
           }
 
+        override def tryAcquireN(n: Long)(implicit trace: Trace): UIO[Boolean] =
+          ref.modify {
+            case Right(permits) if permits >= n =>
+              true -> Right(permits - n)
+            case other =>
+              false -> other
+          }
+
         def withPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, A] =
           withPermits(1L)(zio)
 

From ddcf2d1d8fbbcdb5759b872916d938c7e2135ac4 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igor.doroxov@gmail.com>
Date: Mon, 17 Feb 2025 15:20:12 -0500
Subject: [PATCH 04/12] refactor implemetation

---
 .../src/test/scala/zio/SemaphoreSpec.scala    | 44 ++++++++--------
 .../test/scala/zio/stm/TSemaphoreSpec.scala   | 32 +++++++++++-
 .../shared/src/main/scala/zio/Semaphore.scala | 50 +++++++++++--------
 .../src/main/scala/zio/stm/TSemaphore.scala   | 19 +++++++
 4 files changed, 99 insertions(+), 46 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index 375cfd4ac6b1..1345de25ba6b 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -32,37 +32,33 @@ object SemaphoreSpec extends ZIOBaseSpec {
         permits   <- semaphore.available
       } yield assertTrue(permits == 2L)
     },
-    test("tryAcquire should succeed when a permit is available") {
+    test("tryWithPermits acquires and releases same number of permits") {
       for {
-        sem <- Semaphore.make(1L)
-        res <- sem.tryAcquire
-      } yield assert(res)(isTrue)
+        sem     <- Semaphore.make(3L)
+        ans     <- sem.tryWithPermits(2L)(ZIO.unit)
+        permits <- sem.available
+      } yield assertTrue(permits == 3L && ans.isDefined)
     },
-    test("tryAcquireN should acquire permits if enough are available") {
+    test("tryWithPermits returns None if no permits available") {
       for {
-        sem <- Semaphore.make(5L)
-        res <- sem.tryAcquireN(3L)
-      } yield assert(res)(isTrue)
+        sem     <- Semaphore.make(3L)
+        ans     <- sem.tryWithPermits(4L)(ZIO.unit)
+        permits <- sem.available
+      } yield assertTrue(permits == 3L && ans.isEmpty)
     },
-    test("tryAcquireN should fail if not enough permits are available") {
+    test("tryWithPermit acquires and releases same number of permits") {
       for {
-        sem <- Semaphore.make(2L)
-        res <- sem.tryAcquireN(3L)
-      } yield assert(res)(isFalse)
+        sem     <- Semaphore.make(3L)
+        ans     <- sem.tryWithPermit(ZIO.unit)
+        permits <- sem.available
+      } yield assertTrue(permits == 3L && ans.isDefined)
     },
-    test("tryAcquireN should decrease the permit count when successful") {
+    test("tryWithPermits returns None if requested permits in negative number") {
       for {
-        sem   <- Semaphore.make(5L)
-        _     <- sem.tryAcquireN(3L)
-        avail <- sem.available
-      } yield assert(avail)(equalTo(2L))
-    },
-    test("tryAcquireN should not change permit count when unsuccessful") {
-      for {
-        sem   <- Semaphore.make(2L)
-        _     <- sem.tryAcquireN(3L)
-        avail <- sem.available
-      } yield assert(avail)(equalTo(2L))
+        sem     <- Semaphore.make(3L)
+        ans     <- sem.tryWithPermits(-1L)(ZIO.unit)
+        permits <- sem.available
+      } yield assertTrue(permits == 3L && ans.isEmpty)
     },
     test("awaiting returns the count of waiting fibers") {
       for {
diff --git a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
index f868dde9d41c..1452a2c7bfb1 100644
--- a/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
@@ -138,7 +138,7 @@ object TSemaphoreSpec extends ZIOBaseSpec {
         }
       }
     ),
-    suite("tryAcquire and tryAcquireN")(
+    suite("tryAcquire, tryAcquireN, tryWithPermit and tryWithPermits")(
       test("tryAcquire should succeed when a permit is available") {
         for {
           sem <- TSemaphore.makeCommit(1L)
@@ -183,6 +183,36 @@ object TSemaphoreSpec extends ZIOBaseSpec {
           _     <- sem.tryAcquireN(3L).commit
           avail <- sem.available.commit
         } yield assert(avail)(equalTo(2L))
+      },
+      test("tryWithPermits should acquire a permit and release it") {
+        for {
+          sem    <- TSemaphore.makeCommit(2L)
+          result <- sem.tryWithPermits(1L)(ZIO.succeed(2))
+          avail  <- sem.available.commit
+        } yield assertTrue(result.contains(2) && avail == 2L)
+      },
+      test("tryWithPermits should return None if no permits available") {
+        for {
+          sem    <- TSemaphore.makeCommit(0L)
+          result <- sem.tryWithPermits(1L)(ZIO.succeed(2))
+          avail  <- sem.available.commit
+        } yield assertTrue(result.isEmpty && avail == 0L)
+      },
+      test(
+        "tryWithPermits should return None if requested amount of permits is greater than available amount of permits"
+      ) {
+        for {
+          sem    <- TSemaphore.makeCommit(3L)
+          result <- sem.tryWithPermits(5L)(ZIO.succeed(2))
+          avail  <- sem.available.commit
+        } yield assertTrue(result.isEmpty && avail == 3L)
+      },
+      test("tryWithPermit should acquire a permit and release it") {
+        for {
+          sem    <- TSemaphore.makeCommit(3L)
+          result <- sem.tryWithPermit(ZIO.succeed(2))
+          avail  <- sem.available.commit
+        } yield assertTrue(result.contains(2) && avail == 3L)
       }
     )
   )
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index bf9be7011b72..602a4b988a00 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -39,19 +39,6 @@ sealed trait Semaphore extends Serializable {
    */
   def available(implicit trace: Trace): UIO[Long]
 
-  /**
-   * Attempts to acquire a permit without blocking. Returns `true` if the permit
-   * was acquired, otherwise `false`.
-   */
-  def tryAcquire(implicit trace: Trace): UIO[Boolean] =
-    tryAcquireN(1L)
-
-  /**
-   * Attempts to acquire the specified number of permits without blocking.
-   * Returns `true` if the permits were acquired, otherwise `false`.
-   */
-  def tryAcquireN(n: Long)(implicit trace: Trace): UIO[Boolean] = ZIO.succeed(false)
-
   /**
    * Returns the number of tasks currently waiting for permits. The default
    * implementation returns 0.
@@ -84,6 +71,20 @@ sealed trait Semaphore extends Serializable {
    * permits and releasing them when the scope is closed.
    */
   def withPermitsScoped(n: Long)(implicit trace: Trace): ZIO[Scope, Nothing, Unit]
+
+  /**
+   * Executes the effect, acquiring a permit if available and releasing it after
+   * execution. Returns `None` if no permits were available.
+   */
+  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    tryWithPermits(1L)(zio)
+
+  /**
+   * Executes the effect, acquiring `n` permits if available and releasing them
+   * after execution. Returns `None` if no permits were available.
+   */
+  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    ZIO.succeed(None)
 }
 
 object Semaphore {
@@ -111,14 +112,6 @@ object Semaphore {
             case Right(_)    => 0L
           }
 
-        override def tryAcquireN(n: Long)(implicit trace: Trace): UIO[Boolean] =
-          ref.modify {
-            case Right(permits) if permits >= n =>
-              true -> Right(permits - n)
-            case other =>
-              false -> other
-          }
-
         def withPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, A] =
           withPermits(1L)(zio)
 
@@ -131,8 +124,23 @@ object Semaphore {
         def withPermitsScoped(n: Long)(implicit trace: Trace): ZIO[Scope, Nothing, Unit] =
           ZIO.acquireRelease(reserve(n))(_.release).flatMap(_.acquire)
 
+        override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+          tryReserve(n).flatMap {
+            case Some(reservation) => (reservation.acquire *> zio <* reservation.release).asSome
+            case None              => ZIO.succeed(None)
+          }
+
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])
 
+        def tryReserve(n: Long)(implicit trace: Trace): UIO[Option[Reservation]] =
+          if (n <= 0) ZIO.succeed(None)
+          else
+            ref.modify {
+              case Right(permits) if permits >= n =>
+                Some(Reservation(ZIO.unit, releaseN(n))) -> Right(permits - n)
+              case other => None -> other
+            }
+
         def reserve(n: Long)(implicit trace: Trace): UIO[Reservation] =
           if (n < 0)
             ZIO.die(new IllegalArgumentException(s"Unexpected negative `$n` permits requested."))
diff --git a/core/shared/src/main/scala/zio/stm/TSemaphore.scala b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
index 96b536c6f913..99fa35067c27 100644
--- a/core/shared/src/main/scala/zio/stm/TSemaphore.scala
+++ b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
@@ -81,6 +81,25 @@ final class TSemaphore private (val permits: TRef[Long]) extends Serializable {
       } else false
     }
 
+  /**
+   * Executes the specified effect, acquiring `1` permit if available and
+   * releasing them after execution. Returns `None` if no permits were
+   * available.
+   */
+  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    tryWithPermits(1L)(zio)
+
+  /**
+   * Executes the specified effect, acquiring `n` permits if available and
+   * releasing them after execution. Returns `None` if no permits were
+   * available.
+   */
+  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    tryAcquireN(n).commit.flatMap {
+      case true  => zio.onExit(_ => releaseN(n).commit).asSome
+      case false => ZIO.succeed(None)
+    }
+
   /**
    * Acquire at least `min` permits and at most `max` permits in a transactional
    * context.

From 3aae3b5cff5bbbc859cd141a1c01887b11d1ee0f Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igor.doroxov@gmail.com>
Date: Thu, 27 Feb 2025 10:15:16 -0500
Subject: [PATCH 05/12] make update to tryReserve

---
 core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala | 9 ++++-----
 core/shared/src/main/scala/zio/Semaphore.scala           | 3 ++-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index 1345de25ba6b..36fdb59dbdfa 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -53,12 +53,11 @@ object SemaphoreSpec extends ZIOBaseSpec {
         permits <- sem.available
       } yield assertTrue(permits == 3L && ans.isDefined)
     },
-    test("tryWithPermits returns None if requested permits in negative number") {
+    test("tryWithPermits fails if requested permits in negative number") {
       for {
-        sem     <- Semaphore.make(3L)
-        ans     <- sem.tryWithPermits(-1L)(ZIO.unit)
-        permits <- sem.available
-      } yield assertTrue(permits == 3L && ans.isEmpty)
+        sem <- Semaphore.make(3L)
+        ans <- sem.tryWithPermits(-1L)(ZIO.unit).exit
+      } yield assert(ans)(dies(isSubtype[IllegalArgumentException](anything)))
     },
     test("awaiting returns the count of waiting fibers") {
       for {
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 602a4b988a00..83974cebe71d 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -133,7 +133,8 @@ object Semaphore {
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])
 
         def tryReserve(n: Long)(implicit trace: Trace): UIO[Option[Reservation]] =
-          if (n <= 0) ZIO.succeed(None)
+          if (n < 0) ZIO.die(new IllegalArgumentException(s"Unexpected negative `$n` permits requested."))
+          else if (n == 0L) ZIO.succeed(Some(Reservation(ZIO.unit, ZIO.unit)))
           else
             ref.modify {
               case Right(permits) if permits >= n =>

From 3d1c5a71c7e72f37529b5e669a7692433b660fcc Mon Sep 17 00:00:00 2001
From: Adam Hearn <22334119+hearnadam@users.noreply.github.com>
Date: Mon, 17 Mar 2025 16:47:26 -0700
Subject: [PATCH 06/12] Avoid `None` match

Co-authored-by: kyri-petrou <67301607+kyri-petrou@users.noreply.github.com>
---
 core/shared/src/main/scala/zio/Semaphore.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 83974cebe71d..423bfdb89046 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -127,7 +127,7 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           tryReserve(n).flatMap {
             case Some(reservation) => (reservation.acquire *> zio <* reservation.release).asSome
-            case None              => ZIO.succeed(None)
+            case _                 => Exit.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])

From a97d87725345389b3bf19379b4e41362c8c32e78 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@Igors-MacBook-Pro.local>
Date: Thu, 27 Mar 2025 21:57:58 -0400
Subject: [PATCH 07/12] address comments

---
 .../shared/src/main/scala/zio/Semaphore.scala | 31 ++++----
 .../src/main/scala/zio/stm/TSemaphore.scala   | 77 +++++++++----------
 2 files changed, 53 insertions(+), 55 deletions(-)

diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 423bfdb89046..176c6cd43221 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -45,6 +45,20 @@ sealed trait Semaphore extends Serializable {
    */
   def awaiting(implicit trace: Trace): UIO[Long] = ZIO.succeed(0L)
 
+  /**
+   * Executes the effect, acquiring a permit if available and releasing it after
+   * execution. Returns `None` if no permits were available.
+   */
+  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    tryWithPermits(1L)(zio)
+
+  /**
+   * Executes the effect, acquiring `n` permits if available and releasing them
+   * after execution. Returns `None` if no permits were available.
+   */
+  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    ZIO.none
+
   /**
    * Executes the specified workflow, acquiring a permit immediately before the
    * workflow begins execution and releasing it immediately after the workflow
@@ -72,19 +86,6 @@ sealed trait Semaphore extends Serializable {
    */
   def withPermitsScoped(n: Long)(implicit trace: Trace): ZIO[Scope, Nothing, Unit]
 
-  /**
-   * Executes the effect, acquiring a permit if available and releasing it after
-   * execution. Returns `None` if no permits were available.
-   */
-  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
-    tryWithPermits(1L)(zio)
-
-  /**
-   * Executes the effect, acquiring `n` permits if available and releasing them
-   * after execution. Returns `None` if no permits were available.
-   */
-  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
-    ZIO.succeed(None)
 }
 
 object Semaphore {
@@ -127,7 +128,7 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           tryReserve(n).flatMap {
             case Some(reservation) => (reservation.acquire *> zio <* reservation.release).asSome
-            case _                 => Exit.none
+            case _                 => ZIO.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])
@@ -146,7 +147,7 @@ object Semaphore {
           if (n < 0)
             ZIO.die(new IllegalArgumentException(s"Unexpected negative `$n` permits requested."))
           else if (n == 0L)
-            ZIO.succeedNow(Reservation(ZIO.unit, ZIO.unit))
+            ZIO.succeed(Reservation(ZIO.unit, ZIO.unit))
           else
             Promise.make[Nothing, Unit].flatMap { promise =>
               ref.modify {
diff --git a/core/shared/src/main/scala/zio/stm/TSemaphore.scala b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
index 99fa35067c27..3ca0e810ad84 100644
--- a/core/shared/src/main/scala/zio/stm/TSemaphore.scala
+++ b/core/shared/src/main/scala/zio/stm/TSemaphore.scala
@@ -60,46 +60,6 @@ final class TSemaphore private (val permits: TRef[Long]) extends Serializable {
   def acquireN(n: Long): USTM[Unit] =
     acquireBetween(n, n).unit
 
-  /**
-   * Tries to acquire a single permit in a transactional context. Returns `true`
-   * if the permit was acquired, otherwise `false`.
-   */
-  def tryAcquire: USTM[Boolean] = tryAcquireN(1L)
-
-  /**
-   * Tries to acquire the specified number of permits in a transactional
-   * context. Returns `true` if the permits were acquired, otherwise `false`.
-   */
-  def tryAcquireN(n: Long): USTM[Boolean] =
-    ZSTM.Effect { (journal, _, _) =>
-      assertNonNegative(n)
-
-      val available: Long = permits.unsafeGet(journal)
-      if (available >= n) {
-        permits.unsafeSet(journal, available - n)
-        true
-      } else false
-    }
-
-  /**
-   * Executes the specified effect, acquiring `1` permit if available and
-   * releasing them after execution. Returns `None` if no permits were
-   * available.
-   */
-  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
-    tryWithPermits(1L)(zio)
-
-  /**
-   * Executes the specified effect, acquiring `n` permits if available and
-   * releasing them after execution. Returns `None` if no permits were
-   * available.
-   */
-  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
-    tryAcquireN(n).commit.flatMap {
-      case true  => zio.onExit(_ => releaseN(n).commit).asSome
-      case false => ZIO.succeed(None)
-    }
-
   /**
    * Acquire at least `min` permits and at most `max` permits in a transactional
    * context.
@@ -148,6 +108,43 @@ final class TSemaphore private (val permits: TRef[Long]) extends Serializable {
       permits.unsafeSet(journal, current + n)
     }
 
+  /**
+   * Tries to acquire a single permit in a transactional context. Returns `true`
+   * if the permit was acquired, otherwise `false`.
+   */
+  def tryAcquire: USTM[Boolean] = tryAcquireN(1L)
+
+  /**
+   * Tries to acquire the specified number of permits in a transactional
+   * context. Returns `true` if the permits were acquired, otherwise `false`.
+   */
+  def tryAcquireN(n: Long): USTM[Boolean] =
+    ZSTM.Effect { (journal, _, _) =>
+      assertNonNegative(n)
+
+      val available: Long = permits.unsafeGet(journal)
+      if (available >= n) {
+        permits.unsafeSet(journal, available - n)
+        true
+      } else false
+    }
+
+  /**
+   * Executes the specified effect, acquiring `1` permit if available and
+   * releasing them after execution. Returns `None` if no permits were
+   * available.
+   */
+  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    tryWithPermits(1L)(zio)
+
+  /**
+   * Executes the specified effect, acquiring `n` permits if available and
+   * releasing them after execution. Returns `None` if no permits were
+   * available.
+   */
+  def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+    ZSTM.acquireReleaseWith(tryAcquireN(n))(releaseN(n).commit.whenDiscard(_))(zio.when(_))
+
   /**
    * Executes the specified effect, acquiring a permit immediately before the
    * effect begins execution and releasing it immediately after the effect

From d1ef09bba8494a230c5a07f4cf09f66288c93283 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@Igors-MacBook-Pro.local>
Date: Thu, 27 Mar 2025 22:12:51 -0400
Subject: [PATCH 08/12] revert back to Exit.none

---
 core/shared/src/main/scala/zio/Semaphore.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 176c6cd43221..63f2154eeedf 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -128,7 +128,7 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           tryReserve(n).flatMap {
             case Some(reservation) => (reservation.acquire *> zio <* reservation.release).asSome
-            case _                 => ZIO.none
+            case _                 => Exit.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])

From 5a8de0e4091d772a72c2c1291f80d17d8263b439 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@mac.mynetworksettings.com>
Date: Mon, 31 Mar 2025 20:08:34 -0400
Subject: [PATCH 09/12] refactor

---
 .../src/test/scala/zio/SemaphoreSpec.scala       |  9 ++++++++-
 core/shared/src/main/scala/zio/Semaphore.scala   | 16 +++++++++++-----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index 36fdb59dbdfa..aede23c84303 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -39,10 +39,17 @@ object SemaphoreSpec extends ZIOBaseSpec {
         permits <- sem.available
       } yield assertTrue(permits == 3L && ans.isDefined)
     },
+    test("tryWithPermits if 0 permits requested") {
+      for {
+        sem     <- Semaphore.make(3L)
+        ans     <- sem.tryWithPermits(0L)(ZIO.succeed("I got executed"))
+        permits <- sem.available
+      } yield assertTrue(permits == 3L && ans.contains("I got executed"))
+    },
     test("tryWithPermits returns None if no permits available") {
       for {
         sem     <- Semaphore.make(3L)
-        ans     <- sem.tryWithPermits(4L)(ZIO.unit)
+        ans     <- sem.tryWithPermits(4L)(ZIO.succeed("Shouldn't get executed"))
         permits <- sem.available
       } yield assertTrue(permits == 3L && ans.isEmpty)
     },
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 63f2154eeedf..f273332e033f 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -126,16 +126,22 @@ object Semaphore {
           ZIO.acquireRelease(reserve(n))(_.release).flatMap(_.acquire)
 
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
-          tryReserve(n).flatMap {
-            case Some(reservation) => (reservation.acquire *> zio <* reservation.release).asSome
-            case _                 => Exit.none
+          ZIO.acquireReleaseWith(tryReserve(n)) {
+            case Some(reservation) => reservation.release
+            case None              => Exit.none
+          } {
+            case Some(_) => zio.asSome
+            case None    => Exit.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])
+        object Reservation {
+          private[zio] val zero = Reservation(ZIO.unit, ZIO.unit)
+        }
 
         def tryReserve(n: Long)(implicit trace: Trace): UIO[Option[Reservation]] =
           if (n < 0) ZIO.die(new IllegalArgumentException(s"Unexpected negative `$n` permits requested."))
-          else if (n == 0L) ZIO.succeed(Some(Reservation(ZIO.unit, ZIO.unit)))
+          else if (n == 0L) ZIO.succeed(Some(Reservation.zero))
           else
             ref.modify {
               case Right(permits) if permits >= n =>
@@ -147,7 +153,7 @@ object Semaphore {
           if (n < 0)
             ZIO.die(new IllegalArgumentException(s"Unexpected negative `$n` permits requested."))
           else if (n == 0L)
-            ZIO.succeed(Reservation(ZIO.unit, ZIO.unit))
+            ZIO.succeed(Reservation.zero)
           else
             Promise.make[Nothing, Unit].flatMap { promise =>
               ref.modify {

From ae5fce9e453c19174c94333628ac1e6563423a95 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@mac.mynetworksettings.com>
Date: Mon, 14 Apr 2025 20:35:51 -0400
Subject: [PATCH 10/12] address comments

---
 .../shared/src/test/scala/zio/SemaphoreSpec.scala    | 12 ++++++++++++
 core/shared/src/main/scala/zio/Semaphore.scala       |  6 +++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index aede23c84303..fd85b06d5797 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -66,6 +66,18 @@ object SemaphoreSpec extends ZIOBaseSpec {
         ans <- sem.tryWithPermits(-1L)(ZIO.unit).exit
       } yield assert(ans)(dies(isSubtype[IllegalArgumentException](anything)))
     },
+    test("tryWithPermits restores permits after failure") {
+      for {
+        sem     <- Semaphore.make(3L)
+        failure = ZIO.fail("exception")
+        result  <- sem.tryWithPermits(2L)(failure).exit
+        permits <- sem.available
+      } yield assertTrue(
+        permits == 3L,
+        result.isFailure,
+        result == Exit.fail("exception")
+      )
+    },
     test("awaiting returns the count of waiting fibers") {
       for {
         semaphore    <- Semaphore.make(1)
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index f273332e033f..957b9d86e228 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -128,10 +128,10 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           ZIO.acquireReleaseWith(tryReserve(n)) {
             case Some(reservation) => reservation.release
-            case None              => Exit.none
+            case _              => Exit.none
           } {
-            case Some(_) => zio.asSome
-            case None    => Exit.none
+            case _: Some[?] => zio.asSome
+            case _    => Exit.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])

From 2f0ef9c2a07bab443c11c2bdfe50b7831652f216 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@mac.mynetworksettings.com>
Date: Mon, 14 Apr 2025 20:56:24 -0400
Subject: [PATCH 11/12] fmt

---
 core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala | 2 +-
 core/shared/src/main/scala/zio/Semaphore.scala           | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
index fd85b06d5797..78076c00ef76 100644
--- a/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
@@ -69,7 +69,7 @@ object SemaphoreSpec extends ZIOBaseSpec {
     test("tryWithPermits restores permits after failure") {
       for {
         sem     <- Semaphore.make(3L)
-        failure = ZIO.fail("exception")
+        failure  = ZIO.fail("exception")
         result  <- sem.tryWithPermits(2L)(failure).exit
         permits <- sem.available
       } yield assertTrue(
diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 957b9d86e228..0f585d837265 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -128,10 +128,10 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           ZIO.acquireReleaseWith(tryReserve(n)) {
             case Some(reservation) => reservation.release
-            case _              => Exit.none
+            case _                 => Exit.none
           } {
             case _: Some[?] => zio.asSome
-            case _    => Exit.none
+            case _          => Exit.none
           }
 
         case class Reservation(acquire: UIO[Unit], release: UIO[Any])

From 74a8287e696a2a92e1e979041ffe571ea95388c4 Mon Sep 17 00:00:00 2001
From: Igor Dorokhov <igordorokhov@mac.mynetworksettings.com>
Date: Wed, 23 Apr 2025 17:04:42 -0400
Subject: [PATCH 12/12] address comments

---
 core/shared/src/main/scala/zio/Semaphore.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/shared/src/main/scala/zio/Semaphore.scala b/core/shared/src/main/scala/zio/Semaphore.scala
index 0f585d837265..1dacb46b0eb8 100644
--- a/core/shared/src/main/scala/zio/Semaphore.scala
+++ b/core/shared/src/main/scala/zio/Semaphore.scala
@@ -49,7 +49,7 @@ sealed trait Semaphore extends Serializable {
    * Executes the effect, acquiring a permit if available and releasing it after
    * execution. Returns `None` if no permits were available.
    */
-  def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
+  final def tryWithPermit[R, E, A](zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
     tryWithPermits(1L)(zio)
 
   /**
@@ -128,7 +128,7 @@ object Semaphore {
         override def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
           ZIO.acquireReleaseWith(tryReserve(n)) {
             case Some(reservation) => reservation.release
-            case _                 => Exit.none
+            case _                 => Exit.unit
           } {
             case _: Some[?] => zio.asSome
             case _          => Exit.none