Improve resilience in some failure cases (#25)

Author: ghostdogpr

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

@@ -37,7 +37,8 @@ class Sharding private (
       shardManager.register(address)
 
   val unregister: Task[Unit] =
-    ZIO.logDebug(s"Stopping local entities") *>
+    shardManager.getAssignments *> // ping the shard manager first to stop if it's not available
+      ZIO.logDebug(s"Stopping local entities") *>
       isShuttingDownRef.set(true) *>
       entityStates.get.flatMap(states => ZIO.foreachDiscard(states.values)(_.entityManager.terminateAllEntities)) *>
       ZIO.logDebug(s"Unregistering pod $address to Shard Manager") *>
@@ -108,23 +109,30 @@ class Sharding private (
   def getPods: UIO[Set[PodAddress]] =
     shardAssignments.get.map(_.values.toSet)
 
+  private def updateAssignments(
+    assignmentsOpt: Map[ShardId, Option[PodAddress]],
+    fromShardManager: Boolean
+  ): UIO[Unit] = {
+    val assignments = assignmentsOpt.flatMap { case (k, v) => v.map(k -> _) }
+    ZIO.logDebug("Received new shard assignments") *>
+      (if (fromShardManager) shardAssignments.update(map => if (map.isEmpty) assignments else map)
+       else
+         shardAssignments.update(map =>
+           // we keep self assignments (we don't override them with the new assignments
+           // because only the Shard Manager is able to change assignments of the current node, via assign/unassign
+           assignments.filter { case (_, pod) => pod != address } ++
+             map.filter { case (_, pod) => pod == address }
+         ))
+  }
+
   private[shardcake] val refreshAssignments: ZIO[Scope, Nothing, Unit] = {
     val assignmentStream =
       ZStream.fromZIO(
         shardManager.getAssignments.map(_ -> true) // first, get the assignments from the shard manager directly
       ) ++
         storage.assignmentsStream.map(_ -> false) // then, get assignments changes from Redis
     assignmentStream.mapZIO { case (assignmentsOpt, fromShardManager) =>
-      val assignments = assignmentsOpt.flatMap { case (k, v) => v.map(k -> _) }
-      ZIO.logDebug("Received new shard assignments") *>
-        (if (fromShardManager) shardAssignments.update(map => if (map.isEmpty) assignments else map)
-         else
-           shardAssignments.update(map =>
-             // we keep self assignments (we don't override them with the new assignments
-             // because only the Shard Manager is able to change assignments of the current node, via assign/unassign
-             assignments.filter { case (_, pod) => pod != address } ++
-               map.filter { case (_, pod) => pod == address }
-           ))
+      updateAssignments(assignmentsOpt, fromShardManager)
     }.runDrain
   }.retry(Schedule.fixed(config.refreshAssignmentsRetryInterval))
     .interruptible
@@ -230,7 +238,12 @@ class Sharding private (
                                               )
                                               .map(_ isEqual cdt)
                                           )
-                                          ZIO.whenZIO(notify)(shardManager.notifyUnhealthyPod(pod).forkDaemon)
+                                          ZIO.whenZIO(notify)(
+                                            (shardManager.notifyUnhealthyPod(pod) *>
+                                              // just in case we missed the update from the pubsub, refresh assignments
+                                              shardManager.getAssignments
+                                                .flatMap(updateAssignments(_, fromShardManager = true))).forkDaemon
+                                          )
                                         }
                                       }
                                       .flatMap(ZIO.foreach(_)(serialization.decode[Res]))