Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[9390] - add tryAcquire to Semaphore #9825

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 28, 2025
Merged

[9390] - add tryAcquire to Semaphore #9825

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions core-tests/shared/src/test/scala/zio/SemaphoreSpec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,52 @@ object SemaphoreSpec extends ZIOBaseSpec {
permits <- semaphore.available
} yield assertTrue(permits == 2L)
},
test("tryWithPermits acquires and releases same number of permits") {
for {
sem <- Semaphore.make(3L)
ans <- sem.tryWithPermits(2L)(ZIO.unit)
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.succeed("Shouldn't get executed"))
permits <- sem.available
} yield assertTrue(permits == 3L && ans.isEmpty)
},
test("tryWithPermit acquires and releases same number of permits") {
for {
sem <- Semaphore.make(3L)
ans <- sem.tryWithPermit(ZIO.unit)
permits <- sem.available
} yield assertTrue(permits == 3L && ans.isDefined)
},
test("tryWithPermits fails if requested permits in negative number") {
for {
sem <- Semaphore.make(3L)
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)
Expand Down
77 changes: 77 additions & 0 deletions core-tests/shared/src/test/scala/zio/stm/TSemaphoreSpec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,83 @@ object TSemaphoreSpec extends ZIOBaseSpec {
assertTrue(remaining == 3L)
}
}
),
suite("tryAcquire, tryAcquireN, tryWithPermit and tryWithPermits")(
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))
},
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)
}
)
)

Expand Down
39 changes: 38 additions & 1 deletion core/shared/src/main/scala/zio/Semaphore.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
final 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
Expand All @@ -71,6 +85,7 @@ 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]

}

object Semaphore {
Expand Down Expand Up @@ -110,13 +125,35 @@ 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]] =
ZIO.acquireReleaseWith(tryReserve(n)) {
case Some(reservation) => reservation.release
case _ => Exit.unit
} {
case _: Some[?] => zio.asSome
case _ => 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.zero))
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."))
else if (n == 0L)
ZIO.succeedNow(Reservation(ZIO.unit, ZIO.unit))
ZIO.succeed(Reservation.zero)
else
Promise.make[Nothing, Unit].flatMap { promise =>
ref.modify {
Expand Down
37 changes: 37 additions & 0 deletions core/shared/src/main/scala/zio/stm/TSemaphore.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def tryAcquire: USTM[Boolean] = tryAcquireN(1L)
final 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] =
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def tryAcquireN(n: Long): USTM[Boolean] =
final 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]] =
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
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]] =

Copy link
Contributor Author

@IgorDorokhov IgorDorokhov Apr 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@hearnadam TSemaphore is a final class, hence all member of the class are effectively final by default. that is why final modifier is redundant

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

zio/core/shared/src/main/scala/zio/stm/TSemaphore.scala

Line 49 in 66cfa25

final class TSemaphore private (val permits: TRef[Long]) extends Serializable {

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah right, thanks

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]] =
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def tryWithPermits[R, E, A](n: Long)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, Option[A]] =
final 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
Expand Down
Loading