call resource with unique id

evolution-gaming · Sep 21, 2023 · 266a2b6 · 266a2b6
1 parent 74e12ac
commit 266a2b6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/src/main/scala/com/evolution/resourcepool/ResourcePool.scala b/src/main/scala/com/evolution/resourcepool/ResourcePool.scala
@@ -57,21 +57,22 @@ object ResourcePool {
 
     def apply(maxSize: Int, partitions: Int) = {
 
-      def of(maxSize: Int) = {
+      def of(maxSize: Int)(resource: Id => Resource[F, A]) = {
         of0(
           maxSize,
           expireAfter,
           resource)
       }
 
       if (partitions <= 1) {
-        of(maxSize)
+        of(maxSize)(resource)
       } else {
         for {
           ref    <- Ref[F].of(0).toResource
           values <- maxSize
             .divide(partitions)
-            .traverse { maxSize => of(maxSize) }
+            .zipWithIndex
+            .traverse { case (maxSize, idx) => of(maxSize) { id => resource(s"$idx-$id") } }
           values <- values
             .toVector
             .pure[Resource[F, *]]

diff --git a/src/test/scala/com/evolution/resourcepool/ResourcePoolTest.scala b/src/test/scala/com/evolution/resourcepool/ResourcePoolTest.scala
@@ -288,24 +288,25 @@ class ResourcePoolTest extends AsyncFunSuite with Matchers {
   }
 
   test("not exceed `maxSize` with partitioned pool") {
-    val maxSize = 100
+    val maxSize = 10
     val resource = for {
-      ref  <- Ref[IO].of(0).toResource
+      ref  <- Ref[IO].of(List.empty[String]).toResource
       pool <- ResourcePool.of(
         maxSize = maxSize,
-        partitions = 10,
+        partitions = 3,
         expireAfter = 1.day,
-        resource = _ => ref.update { _ + 1 }.toResource
+        resource = id => ref.update { id :: _ }.toResource
       )
     } yield {
       for {
-        _ <- pool
+        _      <- pool
           .resource
           .use { _ => Temporal[IO].sleep(1.millis) }
           .parReplicateA(maxSize * 100)
           .map { _.combineAll }
         result <- ref.get
-        _ <- IO { result shouldEqual maxSize }
+        _      <- IO { result.size shouldEqual maxSize }
+        _      <- IO { result.toSet shouldEqual Set("2-3", "1-2", "1-0", "0-2", "1-1", "0-0", "2-2", "0-1", "2-0", "2-1") }
       } yield {}
     }