Fix permit loss

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
         }

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]],

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]],