fix(platform-cloudflare): preserve delayed request waiters

Author: tim-smart

packages/platform/cloudflare/src/CloudflareCluster.ts

@@ -119,16 +119,16 @@ const notImplemented = (method: string) =>
 
 interface EntityStub {
   readonly invoke: (envelope: string, discard: boolean, delivery?: {
-    readonly deliverAt: number
-    readonly primaryKey: string | null
+    readonly deliverAt?: number | undefined
+    readonly primaryKey?: string | null | undefined
     readonly replyTo?: string | undefined
   }) => Promise<{
     readonly requestId: string
     readonly replies: ReadonlyArray<string>
     readonly error?: "MailboxFull" | "EncodedMessageTooLarge" | undefined
   }>
   readonly acknowledge: (requestId: string, replyId: string) => Promise<ReadonlyArray<string>>
-  readonly interrupt?: (requestId: string) => Promise<void>
+  readonly interrupt?: (storageRequestId: string, clientRequestId?: string) => Promise<void>
   readonly reset?: (requestId: string) => Promise<void>
 }
 
@@ -192,6 +192,7 @@ const make = Effect.fnUntraced(function*(options: LayerOptions) {
       readonly clientRequestId: string
       storageRequestId: string
       lastChunkId?: string
+      replyHandler?: (reply: string) => Promise<void>
     }
     const entries = new Map<string, ClientEntry>()
     let client!: Effect.Success<ReturnType<typeof RpcClient.makeNoSerialization<any, MailboxFull | PersistenceError>>>
@@ -292,14 +293,17 @@ const make = Effect.fnUntraced(function*(options: LayerOptions) {
               }
               const delivery = delayed
                 ? { deliverAt: deliverAt!, primaryKey, ...(replyTo === undefined ? undefined : { replyTo }) }
-                : undefined
+                : replyTo === undefined
+                ? undefined
+                : { replyTo }
               let replyHandler: ((reply: string) => Promise<void>) | undefined
               if (delivery?.replyTo !== undefined) {
                 replyHandler = async (reply) => {
-                  unregisterReplyHandler(clientRequestId)
-                  unregisterReplyHandler(entry.storageRequestId)
+                  unregisterReplyHandler(clientRequestId, replyHandler)
+                  unregisterReplyHandler(entry.storageRequestId, replyHandler)
                   await Effect.runPromise(deliverReplies(entry, [reply]))
                 }
+                entry.replyHandler = replyHandler
                 registerReplyHandler(clientRequestId, replyHandler)
               }
               return Effect.promise(() => target.stub.invoke(envelope, discard, delivery)).pipe(
@@ -328,15 +332,15 @@ const make = Effect.fnUntraced(function*(options: LayerOptions) {
                   return Effect.ensuring(
                     deliver,
                     Effect.sync(() => {
-                      unregisterReplyHandler(clientRequestId)
-                      unregisterReplyHandler(entry.storageRequestId)
+                      unregisterReplyHandler(clientRequestId, replyHandler)
+                      unregisterReplyHandler(entry.storageRequestId, replyHandler)
                     })
                   )
                 }),
                 Effect.tapCause(() =>
                   Effect.sync(() => {
-                    unregisterReplyHandler(clientRequestId)
-                    unregisterReplyHandler(entry.storageRequestId)
+                    unregisterReplyHandler(clientRequestId, replyHandler)
+                    unregisterReplyHandler(entry.storageRequestId, replyHandler)
                   })
                 )
               )
@@ -354,10 +358,12 @@ const make = Effect.fnUntraced(function*(options: LayerOptions) {
             const entry = entries.get(clientRequestId)
             entries.delete(clientRequestId)
             requestTargets.delete(clientRequestId)
-            unregisterReplyHandler(clientRequestId)
-            if (entry !== undefined) unregisterReplyHandler(entry.storageRequestId)
+            unregisterReplyHandler(clientRequestId, entry?.replyHandler)
+            if (entry?.replyHandler !== undefined) {
+              unregisterReplyHandler(entry.storageRequestId, entry.replyHandler)
+            }
             if (entry === undefined || target.stub.interrupt === undefined) return Effect.void
-            return Effect.promise(() => target.stub.interrupt!(entry.storageRequestId))
+            return Effect.promise(() => target.stub.interrupt!(entry.storageRequestId, clientRequestId))
           }
           default:
             return Effect.void

packages/platform/cloudflare/src/CloudflareDurableObjects.ts

@@ -74,7 +74,7 @@ interface ReplayMessage {
   readonly lastSentChunk: string | undefined
   readonly discard: boolean
   readonly deliverAt?: number | undefined
-  readonly replyTo?: string | undefined
+  readonly replyTos?: ReadonlyArray<string> | undefined
 }
 
 interface InvokeResult {
@@ -89,8 +89,8 @@ interface InvokeOutcome {
 }
 
 interface DeliveryOptions {
-  readonly deliverAt: number
-  readonly primaryKey: string | null
+  readonly deliverAt?: number | undefined
+  readonly primaryKey?: string | null | undefined
   readonly replyTo?: string | undefined
 }
 
@@ -121,7 +121,6 @@ export class ClusterEntity extends DurableObject<unknown> {
   #serial: Promise<void> = Promise.resolve()
   readonly #sessions = new Map<string, ReplySession>()
   readonly #workerWaiters = new Map<string, Array<WorkerWaiter>>()
-  readonly #workerWaiterTargets = new Map<string, string>()
 
   constructor(ctx: DurableObjectState, env: unknown) {
     super(ctx, env)
@@ -176,7 +175,7 @@ export class ClusterEntity extends DurableObject<unknown> {
               ? undefined
               : {
                 scheduled: true,
-                ...(row.replyTo === undefined ? undefined : { replyTo: row.replyTo })
+                ...(row.replyTos === undefined ? undefined : { replyTos: row.replyTos })
               }
           ).pipe(
             Effect.catchCause((cause) => this.#completeReplayFailure(registration, row, cause))
@@ -223,7 +222,7 @@ export class ClusterEntity extends DurableObject<unknown> {
         if (persisted.processed) {
           return { result: { requestId: persisted.originalId, replies: [] } }
         }
-        if (delivery !== undefined) {
+        if (delivery?.deliverAt !== undefined) {
           yield* this.#armEarliestAlarm()
           return this.#delayedOutcome(
             persisted.originalId,
@@ -234,6 +233,15 @@ export class ClusterEntity extends DurableObject<unknown> {
         }
         const original = loadMessage(storage.sql, persisted.originalId)
         if (original === undefined) return yield* Effect.die("Duplicate mailbox row disappeared")
+        if (original.deliverAt !== undefined && original.deliverAt > Date.now()) {
+          yield* this.#armEarliestAlarm()
+          return this.#delayedOutcome(
+            persisted.originalId,
+            discard,
+            delivery?.replyTo,
+            String(envelope.requestId)
+          )
+        }
         const replies = yield* this.#runStored(
           registration,
           runtime,
@@ -243,7 +251,7 @@ export class ClusterEntity extends DurableObject<unknown> {
         )
         return { result: { requestId: persisted.originalId, replies } }
       }
-      if (delivery !== undefined) {
+      if (delivery?.deliverAt !== undefined) {
         yield* this.#armEarliestAlarm()
         return this.#delayedOutcome(String(envelope.requestId), discard, delivery.replyTo)
       }
@@ -264,7 +272,6 @@ export class ClusterEntity extends DurableObject<unknown> {
       const waiters = this.#workerWaiters.get(requestId) ?? []
       waiters.push({ clientRequestId, resolve, reject })
       this.#workerWaiters.set(requestId, waiters)
-      this.#workerWaiterTargets.set(clientRequestId, requestId)
     })
     return { result, deferred }
   }
@@ -287,22 +294,20 @@ export class ClusterEntity extends DurableObject<unknown> {
   }
 
   /** @internal Interrupts an in-memory handler execution. Persisted rows remain replayable. */
-  interrupt(requestId: string): Promise<void> {
-    const storageRequestId = this.#workerWaiterTargets.get(requestId) ?? requestId
+  interrupt(storageRequestId: string, clientRequestId = storageRequestId): Promise<void> {
     const waiters = this.#workerWaiters.get(storageRequestId)
     if (waiters !== undefined) {
       const remaining = waiters.filter((waiter) => {
-        if (waiter.clientRequestId !== requestId) return true
+        if (waiter.clientRequestId !== clientRequestId) return true
         waiter.reject(new Error("Delayed entity request interrupted"))
         return false
       })
-      this.#workerWaiterTargets.delete(requestId)
       if (remaining.length === 0) this.#workerWaiters.delete(storageRequestId)
       else this.#workerWaiters.set(storageRequestId, remaining)
     }
-    const session = this.#sessions.get(requestId)
+    const session = this.#sessions.get(storageRequestId)
     if (session === undefined) return Promise.resolve()
-    this.#sessions.delete(requestId)
+    this.#sessions.delete(storageRequestId)
     session.ack?.resolve()
     session.ack = undefined
     session.done = true
@@ -338,7 +343,7 @@ export class ClusterEntity extends DurableObject<unknown> {
     envelopeText: string,
     lastSentChunk: string | undefined,
     discard: boolean,
-    options?: { readonly scheduled?: boolean; readonly replyTo?: string | undefined }
+    options?: { readonly scheduled?: boolean; readonly replyTos?: ReadonlyArray<string> | undefined }
   ) {
     return Effect.flatMap(
       decodeRequest(registration, envelopeText),
@@ -358,7 +363,7 @@ export class ClusterEntity extends DurableObject<unknown> {
         (row) =>
           this.#runStored(registration, runtime, row.envelope, row.lastSentChunk, row.discard, {
             scheduled: true,
-            ...(row.replyTo === undefined ? undefined : { replyTo: row.replyTo })
+            ...(row.replyTos === undefined ? undefined : { replyTos: row.replyTos })
           }).pipe(
             Effect.catchCause((cause) => this.#completeReplayFailure(registration, row, cause))
           ),
@@ -403,7 +408,7 @@ export class ClusterEntity extends DurableObject<unknown> {
         ),
         (reply) =>
           Effect.sync(() => storage.transactionSync(() => saveReply(storage.sql, reply))).pipe(
-            Effect.andThen(this.#deliverScheduledReply(encoded.requestId as string, reply, row.replyTo))
+            Effect.andThen(this.#deliverScheduledReply(encoded.requestId as string, reply, row.replyTos))
           )
       ).pipe(
         Effect.catchCause(() =>
@@ -422,7 +427,7 @@ export class ClusterEntity extends DurableObject<unknown> {
     lastSentChunkText: string | undefined,
     discard: boolean,
     persisted: boolean,
-    options?: { readonly scheduled?: boolean; readonly replyTo?: string | undefined }
+    options?: { readonly scheduled?: boolean; readonly replyTos?: ReadonlyArray<string> | undefined }
   ): Effect.Effect<ReadonlyArray<string>> {
     const rpc = registration.entity.protocol.requests.get(envelope.tag) as Rpc.AnyWithProps
     const storage = this.#state.storage
@@ -457,7 +462,7 @@ export class ClusterEntity extends DurableObject<unknown> {
               yield* Effect.promise(() => this.#offerReply(session, encoded))
             }
             if (scheduled && reply._tag === "WithExit") {
-              yield* this.#deliverScheduledReply(requestId, encoded, options?.replyTo)
+              yield* this.#deliverScheduledReply(requestId, encoded, options?.replyTos)
             }
           })
       )
@@ -477,25 +482,48 @@ export class ClusterEntity extends DurableObject<unknown> {
     })
   }
 
-  #deliverScheduledReply(requestId: string, reply: string, replyTo: string | undefined): Effect.Effect<void> {
+  #deliverScheduledReply(
+    requestId: string,
+    reply: string,
+    replyTos: ReadonlyArray<string> | undefined
+  ): Effect.Effect<void> {
     const waiters = this.#workerWaiters.get(requestId)
     if (waiters !== undefined) {
       this.#workerWaiters.delete(requestId)
       for (const waiter of waiters) {
-        this.#workerWaiterTargets.delete(waiter.clientRequestId)
         waiter.resolve({ requestId, replies: [reply] })
       }
     }
-    if (replyTo === undefined) return Effect.void
+    if (replyTos === undefined) return Effect.void
     const namespace = (this.#state.exports as Record<string, unknown>).ClusterEntity as
       | {
         readonly getByName: (
           name: string
         ) => { readonly deliverReply: (requestId: string, reply: string) => Promise<boolean> }
       }
       | undefined
-    if (namespace === undefined) return Effect.void
-    return Effect.promise(() => namespace.getByName(replyTo).deliverReply(requestId, reply)).pipe(Effect.ignore)
+    if (namespace === undefined) {
+      return Effect.logError(
+        "Scheduled entity reply delivery failed",
+        new Error("CloudflareCluster: ClusterEntity export is unavailable for scheduled reply delivery")
+      )
+    }
+    return Effect.forEach(
+      replyTos,
+      (replyTo) =>
+        Effect.promise(() => namespace.getByName(replyTo).deliverReply(requestId, reply)).pipe(
+          Effect.flatMap((delivered) =>
+            delivered
+              ? Effect.void
+              : Effect.logError(
+                "Scheduled entity reply delivery failed",
+                new Error(`Scheduled entity reply target is unavailable: ${replyTo}`)
+              )
+          ),
+          Effect.catchCause((cause) => Effect.logError("Scheduled entity reply delivery failed", cause))
+        ),
+      { discard: true }
+    )
   }
 
   /** @internal Completes an in-memory delayed ask owned by this entity object. */

packages/platform/cloudflare/src/internal/entityMailbox.ts

@@ -53,7 +53,7 @@ export const persistRequest = (
   }
 
   const existing = sql.exec(
-    `SELECT m.request_id, m.processed, r.reply AS last_reply
+    `SELECT m.request_id, m.processed, m.reply_to, r.reply AS last_reply
      FROM cluster_messages m
      LEFT JOIN cluster_replies r ON r.reply_id = m.last_reply_id
      WHERE m.request_id = ? OR (? IS NOT NULL AND m.message_id = ?)
@@ -64,9 +64,11 @@ export const persistRequest = (
   ).toArray()[0]
   if (existing !== undefined) {
     if (replyTo !== null && Number(existing.processed) === 0) {
+      const replyTos = decodeReplyTargets(existing.reply_to)
+      if (!replyTos.includes(replyTo)) replyTos.push(replyTo)
       sql.exec(
         "UPDATE cluster_messages SET reply_to = ? WHERE request_id = ?",
-        replyTo,
+        JSON.stringify(replyTos),
         String(existing.request_id)
       )
     }
@@ -99,7 +101,7 @@ export const persistRequest = (
     envelopeText,
     discard ? 1 : 0,
     deliverAt,
-    replyTo
+    replyTo === null ? null : JSON.stringify([replyTo])
   )
   return { _tag: "Success" }
 }
@@ -148,16 +150,28 @@ export interface StoredMessage {
   readonly lastSentChunk: string | undefined
   readonly discard: boolean
   readonly deliverAt?: number | undefined
-  readonly replyTo?: string | undefined
+  readonly replyTos?: ReadonlyArray<string> | undefined
+}
+
+const decodeReplyTargets = (value: unknown): Array<string> => {
+  if (typeof value !== "string") return []
+  try {
+    const decoded = JSON.parse(value)
+    if (Array.isArray(decoded) && decoded.every((item) => typeof item === "string")) return decoded
+  } catch {
+    // Rows written before reply targets became a collection contain one plain name.
+  }
+  return [value]
 }
 
 const rowToMessage = (row: Record<string, unknown>): StoredMessage => {
+  const replyTos = decodeReplyTargets(row.reply_to)
   const message: StoredMessage = {
     envelope: String(row.envelope),
     lastSentChunk: typeof row.last_reply === "string" ? row.last_reply : undefined,
     discard: Number(row.discard) === 1,
     ...(typeof row.deliver_at === "number" ? { deliverAt: row.deliver_at } : undefined),
-    ...(typeof row.reply_to === "string" ? { replyTo: row.reply_to } : undefined)
+    ...(replyTos.length === 0 ? undefined : { replyTos })
   }
   return message
 }

packages/platform/cloudflare/src/internal/entityReply.ts

@@ -9,23 +9,32 @@ export const CurrentEntityName = Context.Reference<string | undefined>(
 
 type ReplyHandler = (reply: string) => Promise<void>
 
-const handlers = new Map<string, ReplyHandler>()
+const handlers = new Map<string, Set<ReplyHandler>>()
 
 /** @internal */
 export const registerReplyHandler = (requestId: string, handler: ReplyHandler): void => {
-  handlers.set(requestId, handler)
+  const registered = handlers.get(requestId) ?? new Set()
+  registered.add(handler)
+  handlers.set(requestId, registered)
 }
 
 /** @internal */
-export const unregisterReplyHandler = (requestId: string): void => {
-  handlers.delete(requestId)
+export const unregisterReplyHandler = (requestId: string, handler?: ReplyHandler): void => {
+  if (handler === undefined) {
+    handlers.delete(requestId)
+    return
+  }
+  const registered = handlers.get(requestId)
+  if (registered === undefined) return
+  registered.delete(handler)
+  if (registered.size === 0) handlers.delete(requestId)
 }
 
 /** @internal */
 export const deliverReply = async (requestId: string, reply: string): Promise<boolean> => {
-  const handler = handlers.get(requestId)
-  if (handler === undefined) return false
+  const registered = handlers.get(requestId)
+  if (registered === undefined) return false
   handlers.delete(requestId)
-  await handler(reply)
+  await Promise.all(Array.from(registered, (handler) => handler(reply)))
   return true
 }