Terminate entities when no termination message is provided (#30)

* Terminate entities when no termination message is provided

* Fix test

* Fix tests

Author: ghostdogpr

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

@@ -40,6 +40,7 @@ class Sharding private (
 
   val register: Task[Unit] =
     logger.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

@@ -55,11 +55,17 @@ private[shardcake] object EntityManager {
       entities.update(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).run)
-              .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).run.as(map.updated(entityId, (None, interruptionFiber)))
+                  case None      =>
+                    queue.shutdown.as(map - entityId)
+                }
+              }
           case _                                      =>
             // if no queue is found, do nothing
             ZIO.succeed(map)
@@ -145,7 +151,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

@@ -15,7 +15,7 @@ import zio.test.environment.TestEnvironment
 object ShardingSpec extends DefaultRunnableSpec {
 
   private val layer =
-    (Clock.live ++ Random.live ++ ZLayer.succeed(Config.default) >+>
+    (Clock.live ++ Random.live ++ ZLayer.succeed(Config.default.copy(entityMaxIdleTime = 2 seconds)) >+>
       ShardManagerClient.local ++ Logging.debug ++ Pods.noop ++ Storage.memory ++ Serialization.javaSerialization >+>
       Sharding.live).mapError(TestFailure.fail)
 
@@ -36,6 +36,16 @@ object ShardingSpec extends DefaultRunnableSpec {
           } yield assertTrue(c1 == 2) && assertTrue(c2 == 1)
         }
       },
+      testM("Entity termination") {
+        (Sharding.registerEntity(Counter, behavior) *> Sharding.registerManaged).use { _ =>
+          for {
+            counter <- Sharding.messenger(Counter)
+            _       <- counter.send("c3")(GetCounter.apply)
+            _       <- clock.sleep(3 seconds)
+            c       <- counter.send("c3")(GetCounter.apply)
+          } yield assertTrue(c == 0)
+        }
+      },
       testM("Cluster singleton") {
         Sharding.registerManaged.use(_ =>
           for {

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

@@ -2,7 +2,7 @@ package example.complex
 
 import com.devsisters.shardcake._
 import com.devsisters.shardcake.interfaces.Logging
-import example.complex.GuildBehavior.GuildMessage.Join
+import example.complex.GuildBehavior.GuildMessage.{ Join, Terminate }
 import example.complex.GuildBehavior._
 import zio._
 import zio.clock.Clock
@@ -16,7 +16,7 @@ object GuildApp extends zio.App {
       .toLayer
 
   val program =
-    (Sharding.registerEntity(Guild, behavior) *> Sharding.registerManaged).use { _ =>
+    (Sharding.registerEntity(Guild, behavior, p => Some(Terminate(p))) *> Sharding.registerManaged).use { _ =>
       Sharding.messenger(Guild).flatMap { guild =>
         for {
           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, Has, RIO, Task, ZIO }
+import zio.{ Dequeue, Has, 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")
@@ -47,5 +48,7 @@ object GuildBehavior {
           )
       case GuildMessage.Leave(userId)         =>
         redis.lRem(entityId, 1, userId).unit
+      case GuildMessage.Terminate(p)          =>
+        p.succeed(()) *> ZIO.interrupt
     }
 }