Terminate entities when no termination message is provided

Author: ghostdogpr

entities/src/main/scala/com/devsisters/shardcake/Sharding.scala

@@ -34,6 +34,7 @@ class Sharding private (
 
   val register: Task[Unit] =
     ZIO.logDebug(s"Registering pod $address to Shard Manager") *>
+      isShuttingDownRef.set(false) *>
       shardManager.register(address)
 
   val unregister: Task[Unit] =

entities/src/main/scala/com/devsisters/shardcake/internal/EntityManager.scala

@@ -50,11 +50,17 @@ private[shardcake] object EntityManager {
       entities.updateZIO(map =>
         map.get(entityId) match {
           case Some((Some(queue), interruptionFiber)) =>
-            // if a queue is found, offer the termination message, and set the queue to None so that no new message is enqueued
             Promise
               .make[Nothing, Unit]
-              .flatMap(p => ZIO.foreach(terminateMessage(p))(queue.offer).exit)
-              .as(map.updated(entityId, (None, interruptionFiber)))
+              .flatMap { p =>
+                terminateMessage(p) match {
+                  case Some(msg) =>
+                    // if a queue is found, offer the termination message, and set the queue to None so that no new message is enqueued
+                    queue.offer(msg).exit.as(map.updated(entityId, (None, interruptionFiber)))
+                  case None      =>
+                    queue.shutdown.as(map - entityId)
+                }
+              }
           case _                                      =>
             // if no queue is found, do nothing
             ZIO.succeed(map)
@@ -139,7 +145,7 @@ private[shardcake] object EntityManager {
                             case Some(queue) =>
                               terminateMessage(p) match {
                                 case Some(terminate) => queue.offer(terminate).catchAllCause(_ => p.succeed(()))
-                                case None            => p.succeed(())
+                                case None            => queue.shutdown *> p.succeed(())
                               }
                             case None        => p.succeed(())
                           }

entities/src/test/scala/com/devsisters/shardcake/ShardingSpec.scala

@@ -28,6 +28,18 @@ object ShardingSpec extends ZIOSpecDefault {
           } yield assertTrue(c1 == 2) && assertTrue(c2 == 1)
         }
       },
+      test("Entity termination") {
+        ZIO.scoped {
+          for {
+            _       <- Sharding.registerEntity(Counter, behavior)
+            _       <- Sharding.registerScoped
+            counter <- Sharding.messenger(Counter)
+            _       <- counter.send("c3")(GetCounter.apply)
+            _       <- Clock.sleep(3 seconds)
+            c       <- counter.send("c3")(GetCounter.apply)
+          } yield assertTrue(c == 0)
+        }
+      },
       test("Cluster singleton") {
         ZIO.scoped {
           for {

examples/src/main/scala/example/complex/GuildApp.scala

@@ -3,7 +3,7 @@ package example.complex
 import com.devsisters.shardcake._
 import com.devsisters.shardcake.interfaces.Serialization
 import dev.profunktor.redis4cats.RedisCommands
-import example.complex.GuildBehavior.GuildMessage.Join
+import example.complex.GuildBehavior.GuildMessage.{ Join, Terminate }
 import example.complex.GuildBehavior._
 import zio._
 
@@ -17,7 +17,7 @@ object GuildApp extends ZIOAppDefault {
 
   val program: ZIO[Sharding with Scope with Serialization with RedisCommands[Task, String, String], Throwable, Unit] =
     for {
-      _     <- Sharding.registerEntity(Guild, behavior)
+      _     <- Sharding.registerEntity(Guild, behavior, p => Some(Terminate(p)))
       _     <- Sharding.registerScoped
       guild <- Sharding.messenger(Guild)
       user1 <- Random.nextUUID.map(_.toString)

examples/src/main/scala/example/complex/GuildBehavior.scala

@@ -3,7 +3,7 @@ package example.complex
 import com.devsisters.shardcake.Messenger.Replier
 import com.devsisters.shardcake.{ EntityType, Sharding }
 import dev.profunktor.redis4cats.RedisCommands
-import zio.{ Dequeue, RIO, Task, ZIO }
+import zio.{ Dequeue, Promise, RIO, Task, ZIO }
 
 import scala.util.{ Failure, Success, Try }
 
@@ -13,6 +13,7 @@ object GuildBehavior {
   object GuildMessage {
     case class Join(userId: String, replier: Replier[Try[Set[String]]]) extends GuildMessage
     case class Leave(userId: String)                                    extends GuildMessage
+    case class Terminate(p: Promise[Nothing, Unit])                     extends GuildMessage
   }
 
   object Guild extends EntityType[GuildMessage]("guild")
@@ -45,5 +46,7 @@ object GuildBehavior {
           )
       case GuildMessage.Leave(userId)         =>
         redis.lRem(entityId, 1, userId).unit
+      case GuildMessage.Terminate(p)          =>
+        p.succeed(()) *> ZIO.interrupt
     }
 }