Fix promise leak when using streaming with long-running actors (#136)

devsisters · Jul 23, 2024 · 2a5d6c5 · 2a5d6c5
1 parent 9f1eb4a
commit 2a5d6c5
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/entities/src/main/scala/com/devsisters/shardcake/Sharding.scala b/entities/src/main/scala/com/devsisters/shardcake/Sharding.scala
@@ -214,8 +214,12 @@ class Sharding private (
     })
 
   private[shardcake] def initReply(id: String, replyChannel: ReplyChannel[Nothing]): UIO[Unit] =
-    replyChannels.update(_.updated(id, replyChannel)) <*
-      replyChannel.await.ensuring(replyChannels.update(_ - id)).forkDaemon
+    replyChannels
+      .getAndUpdate(_.updated(id, replyChannel))
+      .flatMap(beforeReplyChannels =>
+        replyChannel.await.ensuring(replyChannels.update(_ - id)).forkDaemon.unless(beforeReplyChannels.contains(id))
+      )
+      .unit
 
   def reply[Reply](reply: Reply, replier: Replier[Reply]): UIO[Unit] =
     replyChannels