Destroy borrowed resources during cleanup

core/src/main/scala/org/typelevel/keypool/KeyPool.scala

@@ -124,15 +124,18 @@ object KeyPool {
       m <- kpVar.getAndSet(PoolMap.closed[A, (B, ExitCase => F[Unit])])
       _ <- m match {
         case PoolClosed() => Applicative[F].unit
-        case PoolOpen(_, m2) =>
-          m2.toList.traverse_ { case (_, pl) =>
-            pl.toList
-              .traverse_ { case (_, r) =>
+        case PoolOpen(_, borrowed, m2) =>
+          borrowed.values.toList.traverse_ { case (_, release) =>
+            metrics.resourceDestroyed(Metrics.DestructionReason.PoolClosed) >>
+              release(exit).attempt.void
+          } >>
+            m2.toList.traverse_ { case (_, pl) =>
+              pl.toList.traverse_ { case (_, (_, release)) =>
                 metrics.idleDec >>
                   metrics.resourceDestroyed(Metrics.DestructionReason.PoolClosed) >>
-                  r._2(exit).attempt.void
+                  release(exit).attempt.void
               }
-          }
+            }
       }
     } yield ()
 
@@ -153,7 +156,7 @@ object KeyPool {
         now: FiniteDuration,
         idleCount: Int,
         m: Map[A, PoolList[(B, ExitCase => F[Unit])]]
-    ): (PoolMap[A, (B, ExitCase => F[Unit])], List[(A, (B, ExitCase => F[Unit]))]) = {
+    ): (Int, Map[A, PoolList[(B, ExitCase => F[Unit])]], List[(A, (B, ExitCase => F[Unit]))]) = {
       val isNotStale: FiniteDuration => Boolean =
         time =>
           time + idleTimeAllowedInPoolNanos >= now // Time value is alright inside the KeyPool in nanos.
@@ -194,7 +197,7 @@ object KeyPool {
       // May be able to use Span eventually
       val (toKeep, toDestroy) = findStale_(identity, identity, m.toList)
       val idleCount_ = idleCount - toDestroy.length
-      (PoolMap.open(idleCount_, toKeep), toDestroy)
+      (idleCount_, toKeep, toDestroy)
     }
 
     val sleep = Temporal[F].sleep(durationBetweenEvictionRuns).void
@@ -204,13 +207,13 @@ object KeyPool {
       .flatMap { now =>
         kpVar.tryModify {
           case p @ PoolClosed() => (p, F.unit)
-          case p @ PoolOpen(idleCount, m) =>
+          case p @ PoolOpen(idleCount, borrowed, m) =>
             if (m.isEmpty)
               (p, F.unit) // Not worth it to introduce deadlock concerns when hot loop is 5 seconds
             else {
-              val (m_, toDestroy) = findStale(now, idleCount, m)
+              val (idleCount_, toKeep, toDestroy) = findStale(now, idleCount, m)
               (
-                m_,
+                PoolMap.open(idleCount_, borrowed, toKeep),
                 // In this context, we're closing the resource due to it not being used for a while - hence a Succeeded exit case.
                 toDestroy
                   .traverse_(r =>
@@ -243,7 +246,7 @@ object KeyPool {
       case PoolClosed() =>
         (0, Map.empty)
 
-      case PoolOpen(idleCount, m) =>
+      case PoolOpen(idleCount, _, m) =>
         val modified = m.map { case (k, pl) =>
           pl match {
             case One(_, _) => (k, 1)
@@ -257,7 +260,8 @@ object KeyPool {
       kp: KeyPoolConcrete[F, A, B],
       k: A,
       r: B,
-      destroy: ExitCase => F[Unit]
+      destroy: ExitCase => F[Unit],
+      token: Unique.Token
   ): ExitCase => F[Unit] = {
     def destroyFor(reason: Metrics.DestructionReason)(ec: ExitCase): F[Unit] =
       kp.kpMetrics.resourceDestroyed(reason) >> destroy(ec)
@@ -283,19 +287,20 @@ object KeyPool {
     ): (PoolMap[A, (B, ExitCase => F[Unit])], ExitCase => F[Unit]) =
       pc match {
         case p @ PoolClosed() => (p, destroyFor(Metrics.DestructionReason.PoolClosed))
-        case p @ PoolOpen(idleCount, m) =>
+        case p @ PoolOpen(idleCount, borrowed, m) =>
           if (kp.kpMaxIdle == 0 || idleCount >= kp.kpMaxIdle)
             (p, destroyFor(Metrics.DestructionReason.MaxIdle))
           else
             m.get(k) match {
               case None =>
                 val cnt_ = idleCount + 1
-                val m_ = PoolMap.open(cnt_, m + (k -> One((r, destroy), now)))
+                val m_ = PoolMap.open(cnt_, borrowed - token, m + (k -> One((r, destroy), now)))
                 (m_, Function.const[F[Unit], ExitCase](kp.kpMetrics.idleInc))
               case Some(l) =>
                 val (l_, mx) = addToList(now, kp.kpMaxPerKey(k), (r, destroy), l)
                 val cnt_ = idleCount + mx.fold(1)(_ => 0)
-                val m_ = PoolMap.open(cnt_, m + (k -> l_))
+                val borrowed_ = mx.fold(borrowed - token)(_ => borrowed)
+                val m_ = PoolMap.open(cnt_, borrowed_, m + (k -> l_))
                 (
                   m_,
                   mx.fold((_: ExitCase) => kp.kpMetrics.idleInc)(_ =>
@@ -316,42 +321,67 @@ object KeyPool {
       k: A
   ): Resource[F, Managed[F, B]] = {
     def go(
+        token: Unique.Token,
         pm: PoolMap[A, (B, ExitCase => F[Unit])]
     ): (PoolMap[A, (B, ExitCase => F[Unit])], Option[(B, ExitCase => F[Unit])]) =
       pm match {
         case p @ PoolClosed() => (p, None)
-        case pOrig @ PoolOpen(idleCount, m) =>
+        case pOrig @ PoolOpen(idleCount, borrowed, m) =>
           m.get(k) match {
             case None => (pOrig, None)
             case Some(One(a, _)) =>
-              (PoolMap.open(idleCount - 1, m - k), Some(a))
+              (PoolMap.open(idleCount - 1, borrowed.updated(token, a), m - k), Some(a))
             case Some(Cons(a, _, _, rest)) =>
-              (PoolMap.open(idleCount - 1, m + (k -> rest)), Some(a))
+              (
+                PoolMap.open(idleCount - 1, borrowed.updated(token, a), m + (k -> rest)),
+                Some(a)
+              )
           }
       }
 
     def allocateNew: F[(B, ExitCase => F[Unit])] =
       kp.kpMetrics.createDuration.surround(kp.kpRes(k).allocatedCase)
 
+    def updateBorrowed(
+        pm: PoolMap[A, (B, ExitCase => F[Unit])],
+        update: Map[Unique.Token, (B, ExitCase => F[Unit])] => Map[
+          Unique.Token,
+          (B, ExitCase => F[Unit])
+        ]
+    ): PoolMap[A, (B, ExitCase => F[Unit])] =
+      pm match {
+        case p @ PoolClosed() => p
+        case PoolOpen(idleCount, borrowed, m) =>
+          PoolMap.open(idleCount, update(borrowed), m)
+      }
+
     for {
       acquisition <- kp.kpMetrics.acquire
       _ <- kp.kpMaxTotalSem.permit
       releasedState <- Resource.eval(Ref[F].of[Reusable](kp.kpDefaultReuseState))
+      token <- Resource.eval(Temporal[F].unique)
       taken <- Resource.makeCaseFull[F, (Boolean, (B, ExitCase => F[Unit]))] { poll =>
-        kp.kpVar.modify(go).flatMap {
-          case Some(r) => kp.kpMetrics.idleDec.as((true, r))
-          case None => poll(allocateNew).map((false, _))
-        }
+        kp.kpVar
+          .modify(go(token, _))
+          .flatMap {
+            case Some(r) => kp.kpMetrics.idleDec.as((true, r))
+            case None => poll(allocateNew).map((false, _))
+          }
+          .flatTap { case (reused, r) =>
+            if (reused) Applicative[F].unit
+            else kp.kpVar.update(pm => updateBorrowed(pm, _ + (token -> r)))
+          }
       } { case ((_, (b, destroy)), exitCase) =>
         for {
           reusable <- releasedState.get
           out <- reusable match {
-            case Reusable.Reuse => put(kp, k, b, destroy).apply(exitCase).attempt.void
+            case Reusable.Reuse => put(kp, k, b, destroy, token).apply(exitCase).attempt.void
             case Reusable.DontReuse =>
               kp.kpMetrics.resourceDestroyed(
                 Metrics.DestructionReason.NotReusable
               ) >> destroy(exitCase).attempt.void
           }
+          _ <- kp.kpVar.update(pm => updateBorrowed(pm, _ - token))
         } yield out
       }
       _ <- Resource.eval(acquisition.complete)
@@ -441,7 +471,7 @@ object KeyPool {
         kpMetrics <- Resource.eval(metricsProvider.get)
         kpVar <- Resource.makeCase(
           Ref[F].of[PoolMap[A, (B, ExitCase => F[Unit])]](
-            PoolMap.open(0, Map.empty[A, PoolList[(B, ExitCase => F[Unit])]])
+            PoolMap.open(0, Map.empty, Map.empty[A, PoolList[(B, ExitCase => F[Unit])]])
           )
         )((kpVar, exit) => KeyPool.destroy(kpVar, exit, kpMetrics))
         kpMaxTotalSem <- Resource.eval(RequestSemaphore[F](fairness, kpMaxTotal))

core/src/main/scala/org/typelevel/keypool/KeyPoolBuilder.scala

@@ -92,7 +92,7 @@ final class KeyPoolBuilder[F[_]: Temporal, A, B] private (
       kpMetrics <- Resource.pure(Metrics.noop)
       kpVar <- Resource.makeCase(
         Ref[F].of[PoolMap[A, (B, ExitCase => F[Unit])]](
-          PoolMap.open(0, Map.empty[A, PoolList[(B, ExitCase => F[Unit])]])
+          PoolMap.open(0, Map.empty, Map.empty[A, PoolList[(B, ExitCase => F[Unit])]])
         )
       )((kpVar, exit) => KeyPool.destroy(kpVar, exit, kpMetrics))
       kpMaxTotalSem <- Resource.eval(RequestSemaphore[F](Fairness.Fifo, kpMaxTotal))

core/src/main/scala/org/typelevel/keypool/internal/PoolMap.scala

@@ -22,16 +22,17 @@
 package org.typelevel.keypool.internal
 
 import cats._
+import cats.effect.kernel.Unique
 import cats.syntax.all._
 
 private[keypool] sealed trait PoolMap[Key, Rezource] extends Product with Serializable {
   def foldLeft[B](b: B)(f: (B, Rezource) => B): B = this match {
     case PoolClosed() => b
-    case PoolOpen(_, m) => m.foldLeft(b) { case (b, (_, pl)) => pl.foldLeft(b)(f) }
+    case PoolOpen(_, _, m) => m.foldLeft(b) { case (b, (_, pl)) => pl.foldLeft(b)(f) }
   }
   def foldRight[B](lb: Eval[B])(f: (Rezource, Eval[B]) => Eval[B]): Eval[B] = this match {
     case PoolClosed() => lb
-    case PoolOpen(_, m) =>
+    case PoolOpen(_, _, m) =>
       Foldable.iterateRight(m.values, lb) { case (pl, b) => pl.foldRight(b)(f) }
   }
 }
@@ -42,11 +43,13 @@ private[keypool] object PoolMap {
       fa.foldRight(lb)(f)
   }
   def closed[K, R]: PoolMap[K, R] = PoolClosed()
-  def open[K, R](n: Int, m: Map[K, PoolList[R]]): PoolMap[K, R] = PoolOpen(n, m)
+  def open[K, R](n: Int, borrowed: Map[Unique.Token, R], m: Map[K, PoolList[R]]): PoolMap[K, R] =
+    PoolOpen(n, borrowed, m)
 }
 
 private[keypool] final case class PoolClosed[Key, Rezource]() extends PoolMap[Key, Rezource]
 private[keypool] final case class PoolOpen[Key, Rezource](
     idleCount: Int,
+    borrowed: Map[Unique.Token, Rezource],
     m: Map[Key, PoolList[Rezource]]
 ) extends PoolMap[Key, Rezource]

core/src/test/scala/org/typelevel/keypool/KeyPoolSpec.scala

@@ -298,4 +298,43 @@ class KeyPoolSpec extends CatsEffectSuite {
     }
   }
 
+  // see https://github.com/typelevel/keypool/issues/291
+  test("Borrowed Resource destroyed during cleanup") {
+    val pool = KeyPool
+      .Builder(
+        (_: Int) => IO.ref(true),
+        (r: Ref[IO, Boolean]) => r.set(false)
+      )
+      .build
+
+    def escapedResource(outerAwait: Deferred[IO, Unit]): IO[FiberIO[(Boolean, Boolean)]] =
+      pool.use { p =>
+        def job(innerAwait: Deferred[IO, Unit]): IO[(Boolean, Boolean)] =
+          p.take(1).use { managed =>
+            val value = managed.value
+
+            for {
+              status1 <- value.get
+              _ <- innerAwait.complete(())
+              _ <- outerAwait.get
+              status2 <- value.get
+            } yield (status1, status2)
+          }
+
+        for {
+          await <- IO.deferred[Unit]
+          fiber <- job(await).start
+          _ <- await.get // make sure the first part happens inside of the open pool
+        } yield fiber
+      }
+
+    for {
+      await <- IO.deferred[Unit]
+      fiber <- escapedResource(await)
+      _ <- await.complete(()) // pool is closed and a fiber can continue its execution
+      outcome <- fiber.join
+      result <- outcome.embedNever
+      (status1, status2) = result
+    } yield assert(status1 === true && status2 === false)
+  }
 }

core/src/test/scala/org/typelevel/keypool/PoolSpec.scala

@@ -292,6 +292,46 @@ class PoolSpec extends CatsEffectSuite {
     else
       pool.take.use(_ => ref.update(l => l :+ id))
 
+  // see https://github.com/typelevel/keypool/issues/291
+  test("Borrowed Resource destroyed during cleanup") {
+    val pool = Pool
+      .Builder(
+        IO.ref(true),
+        (r: Ref[IO, Boolean]) => r.set(false)
+      )
+      .build
+
+    def escapedResource(outerAwait: Deferred[IO, Unit]): IO[FiberIO[(Boolean, Boolean)]] =
+      pool.use { p =>
+        def job(innerAwait: Deferred[IO, Unit]): IO[(Boolean, Boolean)] =
+          p.take.use { managed =>
+            val value = managed.value
+
+            for {
+              status1 <- value.get
+              _ <- innerAwait.complete(())
+              _ <- outerAwait.get
+              status2 <- value.get
+            } yield (status1, status2)
+          }
+
+        for {
+          await <- IO.deferred[Unit]
+          fiber <- job(await).start
+          _ <- await.get // make sure the first part happens inside of the open pool
+        } yield fiber
+      }
+
+    for {
+      await <- IO.deferred[Unit]
+      fiber <- escapedResource(await)
+      _ <- await.complete(()) // pool is closed and a fiber can continue its execution
+      outcome <- fiber.join
+      result <- outcome.embedNever
+      (status1, status2) = result
+    } yield assert(status1 === true && status2 === false)
+  }
+
   private def nothing(ref: Ref[IO, Int]): IO[Unit] =
     ref.get.void
 }