ensure runner detects when it has been externally unregistered (#5508)

Author: tim-smart

.changeset/thick-falcons-brush.md

@@ -0,0 +1,5 @@
+---
+"@effect/cluster": patch
+---
+
+ensure runner detects when it has been externally unregistered

packages/cluster/src/Runners.ts

@@ -196,28 +196,27 @@ export const make: (options: Omit<Runners["Type"], "sendLocal" | "notifyLocal">)
   type StorageRequestEntry = {
     readonly latch: Effect.Latch
     doneLatch: Effect.Latch | undefined
-    readonly messages: Array<Message.OutgoingRequest<any>>
+    readonly messages: Set<Message.OutgoingRequest<any>>
     replies: Array<Reply.Reply<any>>
   }
   const storageRequests = new Map<Snowflake.Snowflake, StorageRequestEntry>()
   const waitingStorageRequests = new Map<Snowflake.Snowflake, Message.OutgoingRequest<any>>()
   const replyFromStorage = Effect.fnUntraced(
     function*(message: Message.OutgoingRequest<any>) {
       let entry = storageRequests.get(message.envelope.requestId)
-      if (!entry) {
+      if (entry) {
+        entry.messages.add(message)
+        entry.doneLatch ??= Effect.unsafeMakeLatch(false)
+        return yield* entry.doneLatch.await
+      } else {
         entry = {
           latch: Effect.unsafeMakeLatch(false),
           doneLatch: undefined,
           replies: [],
-          messages: []
+          messages: new Set([message])
         }
         storageRequests.set(message.envelope.requestId, entry)
       }
-      entry.messages.push(message)
-      if (entry.messages.length > 1) {
-        entry.doneLatch ??= Effect.unsafeMakeLatch(false)
-        return yield* entry.doneLatch.await
-      }
 
       while (true) {
         // wait for the storage loop to notify us
@@ -231,16 +230,16 @@ export const make: (options: Omit<Runners["Type"], "sendLocal" | "notifyLocal">)
           const reply = entry.replies[i]
           // we have reached the end
           if (reply._tag === "WithExit") {
-            for (let j = 0; j < entry.messages.length; j++) {
-              yield* entry.messages[j].respond(reply)
+            for (const message of entry.messages) {
+              yield* message.respond(reply)
             }
             entry.doneLatch?.unsafeOpen()
             return
           }
 
           entry.latch.unsafeClose()
-          for (let i = 0; i < entry.messages.length; i++) {
-            yield* entry.messages[i].respond(reply)
+          for (const message of entry.messages) {
+            yield* message.respond(reply)
           }
           yield* entry.latch.await
         }
@@ -251,6 +250,11 @@ export const make: (options: Omit<Runners["Type"], "sendLocal" | "notifyLocal">)
       Effect.ensuring(
         effect,
         Effect.sync(() => {
+          const entry = storageRequests.get(message.envelope.requestId)
+          if (!entry || entry.messages.size > 1) {
+            entry?.messages.delete(message)
+            return
+          }
           storageRequests.delete(message.envelope.requestId)
           waitingStorageRequests.delete(message.envelope.requestId)
         })

packages/cluster/src/ShardManager.ts

@@ -59,7 +59,7 @@ export class ShardManager extends Context.Tag("@effect/cluster/ShardManager")<Sh
    */
   readonly shardingEvents: (
     address: Option.Option<RunnerAddress>
-  ) => Effect.Effect<Queue.Dequeue<ShardingEvent>, never, Scope>
+  ) => Effect.Effect<Queue.Dequeue<ShardingEvent>, RunnerNotRegistered, Scope>
   /**
    * Register a new runner with the cluster.
    */
@@ -293,6 +293,7 @@ export class Rpcs extends RpcGroup.make(
   Rpc.make("ShardingEvents", {
     payload: { address: Schema.Option(RunnerAddress) },
     success: ShardingEventSchema,
+    error: RunnerNotRegistered,
     stream: true
   }),
   Rpc.make("GetTime", {
@@ -744,11 +745,16 @@ export const make = Effect.gen(function*() {
 
   return ShardManager.of({
     getAssignments,
-    shardingEvents: (address) =>
-      Effect.zipRight(
-        Option.isSome(address) ? runnerHealthApi.onConnection(address.value) : Effect.void,
-        PubSub.subscribe(events)
-      ),
+    shardingEvents: (address) => {
+      if (Option.isNone(address)) {
+        return PubSub.subscribe(events)
+      }
+      return Effect.tap(PubSub.subscribe(events), () => {
+        const isRegistered = MutableHashMap.has(state.allRunners, address.value)
+        if (isRegistered) return
+        return Effect.fail(new RunnerNotRegistered({ address: address.value }))
+      })
+    },
     register,
     unregister,
     rebalance,

packages/cluster/src/Sharding.ts

@@ -36,6 +36,7 @@ import {
   EntityNotManagedByRunner,
   RunnerUnavailable
 } from "./ClusterError.js"
+import * as ClusterError from "./ClusterError.js"
 import { Persisted, Uninterruptible } from "./ClusterSchema.js"
 import * as ClusterSchema from "./ClusterSchema.js"
 import type { CurrentAddress, CurrentRunnerAddress, Entity, HandlersFrom } from "./Entity.js"
@@ -830,11 +831,12 @@ const make = Effect.gen(function*() {
 
     yield* Effect.logDebug("Subscribing to sharding events")
     const mailbox = yield* shardManager.shardingEvents(config.runnerAddress)
-    const startedLatch = yield* Deferred.make<void>()
+    const startedLatch = yield* Deferred.make<void, ClusterError.RunnerNotRegistered>()
 
     const eventsFiber = yield* Effect.gen(function*() {
       while (true) {
-        const [events] = yield* mailbox.takeAll
+        const [events, done] = yield* mailbox.takeAll
+        if (done) return
         for (const event of events) {
           yield* Effect.logDebug("Received sharding event", event)
 
@@ -868,11 +870,16 @@ const make = Effect.gen(function*() {
               }
               break
             }
+            case "RunnerUnregistered": {
+              if (!isLocalRunner(event.address)) break
+              return yield* Effect.fail(new ClusterError.RunnerNotRegistered({ address: event.address }))
+            }
           }
         }
       }
     }).pipe(
       Effect.intoDeferred(startedLatch),
+      Effect.zipRight(Effect.dieMessage("Shard manager event stream down")),
       Effect.forkScoped
     )
 
@@ -886,7 +893,7 @@ const make = Effect.gen(function*() {
       Effect.forkScoped
     )
 
-    yield* Fiber.joinAll([eventsFiber, syncFiber])
+    return yield* Fiber.joinAll([eventsFiber, syncFiber])
   }).pipe(
     Effect.scoped,
     Effect.catchAllCause((cause) => Effect.logDebug(cause)),