diff --git a/core/src/main/scala/org/typelevel/keypool/internal/RequestSemaphore.scala b/core/src/main/scala/org/typelevel/keypool/internal/RequestSemaphore.scala index a98badfb..e2cb5fbd 100644 --- a/core/src/main/scala/org/typelevel/keypool/internal/RequestSemaphore.scala +++ b/core/src/main/scala/org/typelevel/keypool/internal/RequestSemaphore.scala @@ -36,7 +36,7 @@ import org.typelevel.keypool.Fairness * the order in which requests acquire a permit. * * Derived from cats-effect MiniSemaphore - * https://github.com/typelevel/cats-effect/blob/v3.5.4/kernel/shared/src/main/scala/cats/effect/kernel/MiniSemaphore.scala#L29 + * https://github.com/typelevel/cats-effect/blob/v3.7.1/kernel/shared/src/main/scala/cats/effect/kernel/MiniSemaphore.scala#L29 */ private[keypool] abstract class RequestSemaphore[F[_]] { def permit: Resource[F, Unit] @@ -95,10 +95,10 @@ private[keypool] object RequestSemaphore { new RequestSemaphore[F] { private def acquire: F[Unit] = F.deferred[Unit].flatMap { wait => - val cleanup = state.update { case s @ State(waiting, permits) => - if (B.nonEmpty(waiting)) - State(B.cleanup(waiting, wait), permits) - else s + val cleanup = state.flatModify { case State(waiting, permits) => + State(B.cleanup(waiting, wait), permits) -> wait.complete(()).flatMap { won => + if (won) F.unit else release + } } state.flatModifyFull { case (poll, State(waiting, permits)) => @@ -113,7 +113,9 @@ private[keypool] object RequestSemaphore { state.flatModify { case State(waiting, permits) => if (B.nonEmpty(waiting)) { val (rest, next) = B.take(waiting) - State(rest, permits) -> next.complete(()).void + State(rest, permits) -> next.complete(()).flatMap { granted => + if (granted) F.unit else release + } } else State(waiting, permits + 1) -> F.unit } diff --git a/core/src/test/scala/org/typelevel/keypool/PoolSpec.scala b/core/src/test/scala/org/typelevel/keypool/PoolSpec.scala index 12f69552..e81367d2 100644 --- a/core/src/test/scala/org/typelevel/keypool/PoolSpec.scala +++ b/core/src/test/scala/org/typelevel/keypool/PoolSpec.scala @@ -26,6 +26,7 @@ import cats.effect._ import cats.effect.std.CountDownLatch import cats.effect.testkit.TestControl import scala.concurrent.duration._ +import scala.concurrent.TimeoutException import munit.CatsEffectSuite class PoolSpec extends CatsEffectSuite { @@ -257,6 +258,30 @@ class PoolSpec extends CatsEffectSuite { } } + test("do not lose permits when requests time out while another request releases") { + val program = Pool + .Builder(Ref.of[IO, Int](1), nothing) + .withMaxTotal(1) + .build + .use { pool => + for { + holder <- pool.take.use(_ => IO.sleep(1.second)).start + _ <- IO.sleep(1.milli) + // These timeouts fire at the same instant the holder releases. A request that is + // cancelled while `release` hands it the permit takes the permit down with it. + _ <- (1 to 8).toList.parTraverse_(_ => pool.take.use_.timeout(999.millis).attempt) + _ <- holder.join + // Test whether we can acquire a permit, fail with timeout if we can't (no permits available) + _ <- pool.take.use_.timeout(1.minute) + } yield () + } + + TestControl + .executeEmbed(program) + .adaptErr { case e: TimeoutException => new AssertionError(s"permit lost!", e) } + .replicateA_(100) // repeat to increase the chance of hitting the race condition + } + private def reqAction( pool: Pool[IO, Ref[IO, Int]], ref: Ref[IO, List[Int]], diff --git a/core/src/test/scala/org/typelevel/keypool/internal/RequestSemaphoreSpec.scala b/core/src/test/scala/org/typelevel/keypool/internal/RequestSemaphoreSpec.scala index e57a2641..79873241 100644 --- a/core/src/test/scala/org/typelevel/keypool/internal/RequestSemaphoreSpec.scala +++ b/core/src/test/scala/org/typelevel/keypool/internal/RequestSemaphoreSpec.scala @@ -22,6 +22,7 @@ package org.typelevel.keypool.internal import munit.CatsEffectSuite +import cats.syntax.all._ import cats.effect._ import cats.effect.testkit.TestControl import scala.concurrent.duration._ @@ -122,6 +123,30 @@ class RequestSemaphoreSpec extends CatsEffectSuite { } } + List(Fifo, Lifo).foreach { fairness => + test( + s"$fairness: do not lose a permit when a waiter is cancelled while being handed the permit" + ) { + val program = for { + sem <- RequestSemaphore[IO](fairness, 1) + held <- sem.permit.allocated + waiter <- sem.permit.surround(IO.unit).start + _ <- IO.sleep(1.milli) // the waiter is now queued for the permit + // Cancelling only schedules the waiter's cleanup. If `release` dequeues the waiter first, + // it hands the permit to a fiber that is already being cancelled and never releases it. + _ <- IO.both(held._2, waiter.cancel) + _ <- sem.permit.surround(IO.unit) // never completes if the permit was lost + } yield () + + TestControl + .executeEmbed(program) + .adaptErr { case e: TestControl.NonTerminationException => + new AssertionError(s"permit lost!", e) + } + .replicateA_(100) // repeat to increase the chance of hitting the race condition + } + } + private def action( sem: RequestSemaphore[IO], ref: Ref[IO, List[Int]],