Expose terminateLocalEntity (#103)

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

@@ -399,6 +399,12 @@ class Sharding private (
                            .catchAllCause(replyChannel.fail)
       _             <- entityStates.update(_.updated(recipientType.name, EntityState(entityManager, processBinary)))
     } yield ()
+
+  def terminateLocalEntity(entityType: EntityType[_], entityId: String): UIO[Unit] =
+    entityStates.get.flatMap(_.get(entityType.name) match {
+      case Some(state) => state.entityManager.terminateEntity(entityId)
+      case None        => ZIO.unit
+    })
 }
 
 object Sharding {
@@ -547,4 +553,13 @@ object Sharding {
    */
   def getPods: RIO[Sharding, Set[PodAddress]] =
     ZIO.serviceWithZIO[Sharding](_.getPods)
+
+  /**
+   * Terminate a given entity. If a termination message was provided, that message will be sent to the entity and
+   * no new message will be enqueued after that. If no termination message was provided, the entity will be stopped immediately.
+   * This method can only be used if the entity is hosted on the current pod (otherwise it will do nothing).
+   * Typically, you would use this method from inside the entity behavior to stop itself.
+   */
+  def terminateLocalEntity(entityType: EntityType[_], entityId: String): URIO[Sharding, Unit] =
+    ZIO.serviceWithZIO[Sharding](_.terminateLocalEntity(entityType, entityId))
 }

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

@@ -13,6 +13,7 @@ private[shardcake] trait EntityManager[-Req] {
     replyId: Option[String],
     replyChannel: ReplyChannel[Nothing]
   ): IO[EntityNotManagedByThisPod, Unit]
+  def terminateEntity(entityId: String): UIO[Unit]
   def terminateEntitiesOnShards(shards: Set[ShardId]): UIO[Unit]
   def terminateAllEntities: UIO[Unit]
 }
@@ -46,7 +47,7 @@ private[shardcake] object EntityManager {
   private val currentTimeInMilliseconds: UIO[EpochMillis] =
     Clock.currentTime(TimeUnit.MILLISECONDS)
 
-  class EntityManagerLive[Req](
+  private class EntityManagerLive[Req](
     recipientType: RecipientType[Req],
     behavior: (String, Queue[Req]) => Task[Nothing],
     terminateMessage: Signal => Option[Req],
@@ -75,7 +76,7 @@ private[shardcake] object EntityManager {
       } yield ()).interruptible.forkDaemon
     }
 
-    private def terminateEntity(entityId: String): UIO[Unit] =
+    def terminateEntity(entityId: String): UIO[Unit] =
       entities.updateZIO(map =>
         map.get(entityId) match {
           case Some(Left(queue)) =>

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

@@ -25,7 +25,7 @@ object ShardingSpec extends ZIOSpecDefault {
             _       <- Clock.sleep(1 second)
             c1      <- counter.send("c1")(GetCounter.apply)
             c2      <- counter.send("c2")(GetCounter.apply)
-          } yield assertTrue(c1 == 2) && assertTrue(c2 == 1)
+          } yield assertTrue(c1 == 2, c2 == 1)
         }
       },
       test("Streaming") {
@@ -34,7 +34,7 @@ object ShardingSpec extends ZIOSpecDefault {
             _       <- Sharding.registerEntity(Counter, behavior)
             _       <- Sharding.registerScoped
             counter <- Sharding.messenger(Counter)
-            stream  <- counter.sendStream("c1")(StreamingChanges(_))
+            stream  <- counter.sendStream("c1")(StreamingChanges.apply)
             latch   <- Promise.make[Nothing, Unit]
             fiber   <- stream.take(5).tap(_ => latch.succeed(())).runCollect.fork
             _       <- latch.await
@@ -52,7 +52,7 @@ object ShardingSpec extends ZIOSpecDefault {
             _       <- Sharding.registerEntity(Counter, behavior)
             _       <- Sharding.registerScoped
             counter <- Sharding.messenger(Counter)
-            stream  <- counter.sendStream("c1")(StreamingChanges(_))
+            stream  <- counter.sendStream("c1")(StreamingChanges.apply)
             latch   <- Promise.make[Nothing, Unit]
             fiber   <- stream.take(2).tap(_ => latch.succeed(())).runCollect.fork
             _       <- latch.await
@@ -75,6 +75,19 @@ object ShardingSpec extends ZIOSpecDefault {
           } yield assertTrue(c0 == 1, c1 == 0)
         }
       },
+      test("Entity manual termination") {
+        ZIO.scoped {
+          for {
+            _       <- Sharding.registerEntity(Counter, behavior)
+            _       <- Sharding.registerScoped
+            counter <- Sharding.messenger(Counter)
+            _       <- counter.sendDiscard("c3")(IncrementCounter)
+            c0      <- counter.send("c3")(GetCounter.apply)
+            _       <- counter.send("c3")(TriggerTerminate.apply)
+            c1 <- counter.send("c3")(GetCounter.apply) // counter should be restarted
+          } yield assertTrue(c0 == 1, c1 == 0)
+        }
+      },
       test("Entity termination with extension") {
         ZIO.scoped {
           for {
@@ -119,6 +132,7 @@ object CounterActor {
     case object IncrementCounter                             extends CounterMessage
     case object DecrementCounter                             extends CounterMessage
     case class StreamingChanges(replier: StreamReplier[Int]) extends CounterMessage
+    case class TriggerTerminate(replier: Replier[Unit])      extends CounterMessage
   }
 
   object Counter extends EntityType[CounterMessage]("counter")
@@ -134,6 +148,8 @@ object CounterActor {
             case CounterMessage.DecrementCounter          => state.update(_ - 1)
             case CounterMessage.StreamingChanges(replier) =>
               replier.replyStream(state.changes.ensuring(state.set(-1)))
+            case CounterMessage.TriggerTerminate(replier) =>
+              Sharding.terminateLocalEntity(Counter, entityId).ensuring(replier.reply(()))
           }.forever
         )
 }