Fix ZPool hangs during parallel item acquisition after invalidation (#…

…8820) (#8833)
zio · May 21, 2024 · 83940e8 · 83940e8
1 parent ee8c132
commit 83940e8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/core-tests/shared/src/test/scala/zio/ZPoolSpec.scala b/core-tests/shared/src/test/scala/zio/ZPoolSpec.scala
@@ -74,6 +74,20 @@ object ZPoolSpec extends ZIOBaseSpec {
             value  <- count.get
           } yield assertTrue(result == 2 && value == 10)
         } +
+        test("reallocate invalidated item on concurrent demand and maxed out pool size") {
+          ZIO.scoped {
+            for {
+              allocated      <- Ref.make(0)
+              finalized      <- Ref.make(0)
+              get             = ZIO.acquireRelease(allocated.updateAndGet(_ + 1))(_ => finalized.update(_ + 1))
+              pool           <- ZPool.make(get, 0 to 1, Duration.fromSeconds(100))
+              getZIO          = ZIO.scoped(pool.get.flatMap(pool.invalidate))
+              _              <- Live.live(ZIO.foreachPar(0 to 9)(_ => getZIO).disconnect.timeout(1.second))
+              allocatedCount <- allocated.get
+              finalizedCount <- finalized.get
+            } yield assertTrue(allocatedCount == 10 && finalizedCount == 10)
+          }
+        } +
         test("invalidate all items in pool and check that pool.get doesn't hang forever") {
           for {
             allocated      <- Ref.make(0)

diff --git a/core/shared/src/main/scala/zio/ZPool.scala b/core/shared/src/main/scala/zio/ZPool.scala
@@ -215,7 +215,7 @@ object ZPool {
           attempted.forEach(a => invalidated.update(_ - a)) *>
             attempted.finalizer *>
             state.modify { case State(size, free) =>
-              if (size <= range.start)
+              if (size <= range.start || free < 0)
                 allocate -> State(size, free + 1)
               else
                 ZIO.unit -> State(size - 1, free)