Add more sharding events + a way to check all pods health (#59)

Author: ghostdogpr

manager/src/main/scala/com/devsisters/shardcake/GraphQLApi.scala

@@ -15,7 +15,8 @@ object GraphQLApi extends GenericSchema[Has[ShardManager]] {
   case class Mutations(
     register: Pod => RIO[Has[ShardManager], Unit],
     unregister: Pod => RIO[Has[ShardManager], Unit],
-    notifyUnhealthyPod: PodAddressArgs => URIO[Has[ShardManager], Unit]
+    notifyUnhealthyPod: PodAddressArgs => URIO[Has[ShardManager], Unit],
+    checkAllPodsHealth: URIO[Has[ShardManager], Unit]
   )
   case class Subscriptions(events: ZStream[Has[ShardManager], Nothing, ShardingEvent])
 
@@ -26,7 +27,8 @@ object GraphQLApi extends GenericSchema[Has[ShardManager]] {
         Mutations(
           pod => ZIO.serviceWith(_.register(pod)),
           pod => ZIO.serviceWith(_.unregister(pod.address)),
-          args => ZIO.serviceWith(_.notifyUnhealthyPod(args.podAddress))
+          args => ZIO.serviceWith(_.notifyUnhealthyPod(args.podAddress)),
+          ZIO.serviceWith(_.checkAllPodsHealth)
         ),
         Subscriptions(ZStream.serviceWithStream(_.getShardingEvents))
       )

manager/src/main/scala/com/devsisters/shardcake/ShardManager.scala

@@ -33,24 +33,33 @@ class ShardManager(
     ZStream.fromHub(eventsHub)
 
   def register(pod: Pod): UIO[Unit] =
-    logger.logInfo(s"Registering $pod") *>
-      (stateRef
-        .updateAndGet(state =>
-          ZIO
-            .succeed(OffsetDateTime.now())
-            .map(cdt => state.copy(pods = state.pods.updated(pod.address, PodWithMetadata(pod, cdt))))
-        )
-        .flatMap(state => ZIO.when(state.unassignedShards.nonEmpty)(rebalance(false))) *>
-        persistPods.forkDaemon).unit
+    for {
+      _     <- logger.logInfo(s"Registering $pod")
+      state <- stateRef.updateAndGet(state =>
+                 ZIO
+                   .succeed(OffsetDateTime.now())
+                   .map(cdt => state.copy(pods = state.pods.updated(pod.address, PodWithMetadata(pod, cdt))))
+               )
+      _     <- eventsHub.publish(ShardingEvent.PodRegistered(pod.address))
+      _     <- ZIO.when(state.unassignedShards.nonEmpty)(rebalance(false))
+      _     <- persistPods.forkDaemon
+    } yield ()
 
   def notifyUnhealthyPod(podAddress: PodAddress): UIO[Unit] =
     ZIO
       .whenM(stateRef.get.map(_.pods.contains(podAddress))) {
-        ZIO.unlessM(healthApi.isAlive(podAddress))(
-          logger.logWarning(s"$podAddress is not alive, unregistering") *> unregister(podAddress)
-        )
+        eventsHub.publish(ShardingEvent.PodHealthChecked(podAddress)) *>
+          ZIO.unlessM(healthApi.isAlive(podAddress))(
+            logger.logWarning(s"$podAddress is not alive, unregistering") *> unregister(podAddress)
+          )
       }
 
+  def checkAllPodsHealth: UIO[Unit] =
+    for {
+      pods <- stateRef.get.map(_.pods.keySet)
+      _    <- ZIO.foreachParN_(4)(pods)(notifyUnhealthyPod)
+    } yield ()
+
   def unregister(podAddress: PodAddress): UIO[Unit] =
     ZIO
       .whenM(stateRef.get.map(_.pods.contains(podAddress))) {
@@ -68,6 +77,7 @@ class ShardManager(
                                )
                              )
                            }
+          _             <- eventsHub.publish(ShardingEvent.PodUnregistered(podAddress))
           _             <- eventsHub
                              .publish(ShardingEvent.ShardsUnassigned(podAddress, unassignments))
                              .when(unassignments.nonEmpty)
@@ -259,6 +269,9 @@ object ShardManager {
   object ShardingEvent {
     case class ShardsAssigned(pod: PodAddress, shards: Set[ShardId])   extends ShardingEvent
     case class ShardsUnassigned(pod: PodAddress, shards: Set[ShardId]) extends ShardingEvent
+    case class PodRegistered(pod: PodAddress)                          extends ShardingEvent
+    case class PodUnregistered(pod: PodAddress)                        extends ShardingEvent
+    case class PodHealthChecked(pod: PodAddress)                       extends ShardingEvent
   }
 
   def decideAssignmentsForUnassignedShards(