zio · kyri-petrou · Nov 16, 2024 · Nov 15, 2024 · Nov 15, 2024 · Nov 15, 2024
diff --git a/core-tests/shared/src/test/scala/zio/ZKeyedPoolSpec.scala b/core-tests/shared/src/test/scala/zio/ZKeyedPoolSpec.scala
@@ -40,6 +40,24 @@ object ZKeyedPoolSpec extends ZIOBaseSpec {
           _ <- TestClock.adjust((15 * 400).millis)
           _ <- fiber.join
         } yield assertCompletes
-      }
+      },
+      test("invalidate does not cause memory leaks (i9306)") {
+        ZKeyedPool
+          .make[String, Any, Nothing, Array[Int]]((_: String) => ZIO.succeed(Array.ofDim[Int](1000000)), size = 1)
+          .flatMap { pool =>
+            ZIO
+              .foreachDiscard(1 to 10000)(_ =>
+                ZIO.scoped {
+                  for {
+                    item1 <- pool.get("key0")
+                    _     <- ZIO.foreachDiscard(1 to 5)(i => pool.get(s"key$i"))
+                    _     <- pool.invalidate(item1)
+                  } yield ()
+                }
+              )
+          }
+          .as(assertCompletes)
+      } @@ jvmOnly
     ) @@ exceptJS
+
 }
diff --git a/core-tests/shared/src/test/scala/zio/ZPoolSpec.scala b/core-tests/shared/src/test/scala/zio/ZPoolSpec.scala
@@ -206,6 +206,12 @@ object ZPoolSpec extends ZIOBaseSpec {
             _         <- ZIO.scoped(ZPool.make(incCounter <* ZIO.fail("oh no"), 10))
             _         <- latch.await
           } yield assertCompletes
-        } @@ exceptJS(nonFlaky(1000))
+        } @@ exceptJS(nonFlaky(1000)) +
+        test("calling invalidate with items not in the pool doesn't cause memory leaks") {
+          for {
+            pool <- ZPool.make(ZIO.succeed(Array.empty[Int]), 1)
+            _    <- ZIO.foreachDiscard(1 to 1000)(_ => pool.invalidate(Array.ofDim[Int](10000000)))
+          } yield assertCompletes
+        } @@ jvmOnly
     }.provideLayer(Scope.default) @@ timeout(30.seconds)
 }
diff --git a/core/shared/src/main/scala/zio/ZPool.scala b/core/shared/src/main/scala/zio/ZPool.scala
@@ -16,6 +16,7 @@
 
 package zio
 
+import zio.ZIO.InterruptibilityRestorer
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 /**
@@ -115,15 +116,15 @@ object ZPool {
         down     <- Ref.make(false)
         state    <- Ref.make(State(0, 0))
         items    <- Queue.bounded[Attempted[E, A]](range.end)
-        inv      <- Ref.make(Set.empty[A])
+        alloc    <- Ref.make(Set.empty[A])
         initial  <- strategy.initial
         pool = DefaultPool(
                  get.provideSomeEnvironment[Scope](env.union[Scope](_)),
                  range,
                  down,
                  state,
                  items,
-                 inv,
+                 alloc,
                  strategy.track(initial)
                )
         _ <- restore(pool.initialize).foldCauseZIO(
@@ -148,24 +149,35 @@ object ZPool {
       ZIO.done(result)
   }
 
-  private case class DefaultPool[R, E, A](
+  private final case class DefaultPool[E, A](
     creator: ZIO[Scope, E, A],
     range: Range,
     isShuttingDown: Ref[Boolean],
     state: Ref[State],
     items: Queue[Attempted[E, A]],
-    invalidated: Ref[Set[A]],
+    allocated: Ref[Set[A]],
     track: Exit[E, A] => UIO[Any]
   ) extends ZPool[E, A] {
 
+    private def allocate(implicit restore: InterruptibilityRestorer, trace: Trace): UIO[Any] =
+      for {
+        scope     <- Scope.make
+        exit      <- scope.extend(restore(creator)).exit
+        attempted <- ZIO.succeed(Attempted(exit, scope.close(exit)))
+        _         <- attempted.forEach(a => allocated.update(_ + a))
+        _         <- items.offer(attempted)
+        _         <- track(attempted.result)
+        _         <- getAndShutdown.whenZIO(isShuttingDown.get)
+      } yield attempted
+
     /**
      * Returns the number of items in the pool in excess of the minimum size.
      */
     def excess(implicit trace: Trace): UIO[Int] =
       state.get.map { case State(free, size) => size - range.start min free }
 
     def get(implicit trace: Trace): ZIO[Scope, E, A] =
-      ZIO.InterruptibilityRestorer.make.flatMap { restore =>
+      ZIO.InterruptibilityRestorer.make.flatMap { implicit restore =>
         def acquire: UIO[Attempted[E, A]] =
           isShuttingDown.get.flatMap { down =>
             if (down) ZIO.interrupt
@@ -176,9 +188,9 @@ object ZPool {
                     items.take.flatMap { attempted =>
                       attempted.result match {
                         case Exit.Success(item) =>
-                          invalidated.get.flatMap { set =>
-                            if (set.contains(item)) finalizeInvalid(attempted) *> acquire
-                            else Exit.succeed(attempted)
+                          allocated.get.flatMap { set =>
+                            if (set.contains(item)) Exit.succeed(attempted)
+                            else finalizeInvalid(attempted) *> acquire
                           }
                         case _ =>
                           state.modify { case State(size, free) =>
@@ -201,68 +213,47 @@ object ZPool {
         def release(attempted: Attempted[E, A]): UIO[Any] =
           attempted.result match {
             case Exit.Success(item) =>
-              invalidated.get.flatMap { set =>
-                if (set.contains(item)) finalizeInvalid(attempted)
-                else
+              allocated.get.flatMap { set =>
+                if (set.contains(item))
                   state.update(state => state.copy(free = state.free + 1)) *>
                     items.offer(attempted) *>
                     track(attempted.result) *>
                     getAndShutdown.whenZIO(isShuttingDown.get)
+                else finalizeInvalid(attempted)
               }
             case _ =>
               Exit.unit // Handled during acquire
           }
 
         def finalizeInvalid(attempted: Attempted[E, A]): UIO[Any] =
-          attempted.forEach(a => invalidated.update(_ - a)) *>
-            attempted.finalizer *>
+          attempted.finalizer *>
             state.modify { case State(size, free) =>
               if (size <= range.start || free < 0)
                 allocate -> State(size, free + 1)
               else
                 ZIO.unit -> State(size - 1, free)
             }.flatten
 
-        def allocate: UIO[Any] =
-          for {
-            scope     <- Scope.make
-            exit      <- scope.extend(restore(creator)).exit
-            attempted <- ZIO.succeed(Attempted(exit, scope.close(exit)))
-            _         <- items.offer(attempted)
-            _         <- track(attempted.result)
-            _         <- getAndShutdown.whenZIO(isShuttingDown.get)
-          } yield attempted
-
         ZIO.acquireRelease(acquire)(release).flatMap(_.result).disconnect
       }
 
     /**
      * Begins pre-allocating pool entries based on minimum pool size.
      */
-    final def initialize(implicit trace: Trace): UIO[Unit] =
-      ZIO.uninterruptibleMask { restore =>
+    def initialize(implicit trace: Trace): UIO[Unit] =
+      ZIO.uninterruptibleMask { implicit restore =>
         ZIO.replicateZIODiscard(range.start) {
           state.modify { case State(size, free) =>
             if (size < range.start && size >= 0)
-              (
-                for {
-                  scope     <- Scope.make
-                  exit      <- scope.extend(restore(creator)).exit
-                  attempted <- ZIO.succeed(Attempted(exit, scope.close(exit)))
-                  _         <- items.offer(attempted)
-                  _         <- track(attempted.result)
-                  _         <- getAndShutdown.whenZIO(isShuttingDown.get)
-                } yield attempted,
-                State(size + 1, free + 1)
-              )
+              allocate -> State(size + 1, free + 1)
             else
               ZIO.unit -> State(size, free)
           }.flatten
         }
       }
 
     def invalidate(item: A)(implicit trace: zio.Trace): UIO[Unit] =
-      invalidated.update(_ + item)
+      allocated.update(_ - item)
 
     /**
      * Shrinks the pool down, but never to less than the minimum size.
@@ -273,7 +264,7 @@ object ZPool {
           if (size > range.start && free > 0)
             (
               items.take.flatMap { attempted =>
-                attempted.forEach(a => invalidated.update(_ - a)) *>
+                attempted.forEach(a => allocated.update(_ - a)) *>
                   attempted.finalizer *>
                   state.update(state => state.copy(size = state.size - 1))
               },
@@ -295,7 +286,7 @@ object ZPool {
             items.take.foldCauseZIO(
               _ => ZIO.unit,
               attempted =>
-                attempted.forEach(a => invalidated.update(_ - a)) *>
+                attempted.forEach(a => allocated.update(_ - a)) *>
                   attempted.finalizer *>
                   state.update(state => state.copy(size = state.size - 1)) *>
                   getAndShutdown

diff --git a/project/MimaSettings.scala b/project/MimaSettings.scala
@@ -17,7 +17,9 @@ object MimaSettings {
         exclude[IncompatibleResultTypeProblem]("zio.stm.TRef.todo"),
         exclude[DirectMissingMethodProblem]("zio.stm.TRef.versioned_="),
         exclude[IncompatibleResultTypeProblem]("zio.stm.TRef.versioned"),
-        exclude[ReversedMissingMethodProblem]("zio.Fiber#Runtime#UnsafeAPI.zio$Fiber$Runtime$UnsafeAPI$$$outer")
+        exclude[ReversedMissingMethodProblem]("zio.Fiber#Runtime#UnsafeAPI.zio$Fiber$Runtime$UnsafeAPI$$$outer"),
+        exclude[FinalClassProblem]("zio.ZPool$DefaultPool"),
+        exclude[DirectMissingMethodProblem]("zio.ZPool#DefaultPool.invalidated")
       ),
       mimaFailOnProblem := failOnProblem
     )