Fix cluster request resets stranded by in-memory deduplication

.changeset/cluster-reset-claims.md

@@ -0,0 +1,5 @@
+---
+"effect": patch
+---
+
+Fix cluster workflows stalling after request resets.

packages/effect/src/unstable/cluster/MessageStorage.ts

@@ -178,6 +178,11 @@ export class MessageStorage extends Context.Service<MessageStorage, {
     address: EntityAddress
   ) => Effect.Effect<void, PersistenceError>
 
+  /** Release request claims without changing replies or processed state. */
+  readonly resetRequests: (
+    requestIds: ReadonlyArray<Snowflake.Snowflake>
+  ) => Effect.Effect<void, PersistenceError>
+
   /**
    * Reset the mailbox state for the provided addresses.
    */
@@ -412,6 +417,11 @@ export type Encoded = {
     PersistenceError
   >
 
+  /** Release request claims without changing replies or processed state. */
+  readonly resetRequests: (
+    requestIds: ReadonlyArray<Snowflake.Snowflake>
+  ) => Effect.Effect<void, PersistenceError>
+
   /**
    * Reset the mailbox state for the provided addresses.
    */
@@ -716,6 +726,7 @@ export const makeEncoded: (encoded: Encoded) => Effect.Effect<
       )
     },
     resetAddress: (address) => encoded.resetAddresses([address]),
+    resetRequests: (requestIds) => requestIds.length === 0 ? Effect.void : encoded.resetRequests(requestIds),
     resetAddresses: (addresses) => addresses.length === 0 ? Effect.void : encoded.resetAddresses(addresses),
     clearAddress: encoded.clearAddress,
     resetShards: (shardIds) => {
@@ -846,6 +857,7 @@ export const noop: MessageStorage["Service"] = Effect.runSync(make({
   unprocessedMessages: () => Effect.succeed([]),
   unprocessedMessagesById: () => Effect.succeed([]),
   resetAddress: () => Effect.void,
+  resetRequests: () => Effect.void,
   resetAddresses: () => Effect.void,
   clearAddress: () => Effect.void,
   resetShards: () => Effect.void,
@@ -1067,6 +1079,7 @@ export class MemoryDriver extends Context.Service<MemoryDriver>()("effect/cluste
               if (claimedAt !== undefined && claimedAt > now - claimExpirationMillis) {
                 continue
               }
+              // Memory storage models claim expiry, but not SQL reply filtering or transactions.
               messages.push({
                 envelope,
                 lastSentReply: Option.fromNullishOr(entry.replies[entry.replies.length - 1])
@@ -1091,6 +1104,13 @@ export class MemoryDriver extends Context.Service<MemoryDriver>()("effect/cluste
           return unprocessedWith((envelope) => envelopeIds.has(envelope.requestId))
         }),
       resetAddresses,
+      resetRequests: (requestIds) =>
+        Effect.sync(() => {
+          for (const id of requestIds) {
+            const entry = requests.get(String(id))
+            if (entry) lastRead.delete(entry.envelope)
+          }
+        }),
       clearAddress: (address) =>
         Effect.sync(() => {
           for (const [primaryKey, entry] of requestsByPrimaryKey) {

packages/effect/src/unstable/cluster/Sharding.ts

@@ -637,12 +637,22 @@ const make = Effect.gen(function*() {
       let deliveredThisRead = false
       const removableNotifications = new Set<PendingNotification>()
       const resetAddresses = MutableHashSet.empty<EntityAddress>()
+      const resetRequestIds: Array<Snowflake.Snowflake> = []
       const cappedAddresses = MutableHashSet.empty<EntityAddress>()
 
       const markDelivered = Effect.sync(() => {
         deliveredThisRead = true
       })
 
+      // A completed request may remain deduplicated after storage claims its reset.
+      // Release the claim so the reset can be redelivered immediately.
+      const releaseCompletedClaim = (message: Message.Incoming<any>) => {
+        const state = entityManagers.get(message.envelope.address.entityType)
+        if (!state?.manager.isProcessingFor(message, { excludeCompleted: true })) {
+          resetRequestIds.push(message.envelope.requestId)
+        }
+      }
+
       const processMessages = Effect.whileLoop({
         while: () => index < messages.length,
         step: () => index++,
@@ -696,6 +706,7 @@ const make = Effect.gen(function*() {
           } else if (isProcessing || state.status === "closing") {
             // If the request is already processing, we skip it.
             // Or if the entity is closing, we skip all incoming messages.
+            if (isProcessing) releaseCompletedClaim(message)
             return Effect.void
           } else if (message._tag === "IncomingRequest" && pendingNotifications.has(message.envelope.requestId)) {
             const entry = pendingNotifications.get(message.envelope.requestId)!
@@ -738,6 +749,11 @@ const make = Effect.gen(function*() {
               defect: Cause.squash(cause)
             }))
           }
+          if (error.success._tag === "AlreadyProcessingMessage") {
+            // Completion raced with decoding.
+            releaseCompletedClaim(message)
+            return Effect.void
+          }
           if (error.success._tag === "MailboxFull") {
             const address = message.envelope.address
             const state = entityManagers.get(address.entityType)
@@ -850,6 +866,11 @@ const make = Effect.gen(function*() {
         // let the resuming entities check if they are done
         yield* storageReadLock.release(1)
 
+        if (resetRequestIds.length > 0) {
+          yield* Effect.ignore(storage.resetRequests(resetRequestIds))
+          resetRequestIds.length = 0
+        }
+
         if (cappedAddressesToReset !== undefined) {
           yield* Effect.ignore(storage.resetAddresses(cappedAddressesToReset))
         }

packages/effect/src/unstable/cluster/SqlMessageStorage.ts

@@ -686,6 +686,20 @@ export const makeEncoded: (options?: {
       )
     },
 
+    resetRequests: (requestIds) =>
+      requestIds.length === 0
+        ? Effect.void
+        : sql`
+        UPDATE ${messagesTableSql}
+        SET last_read = NULL
+        WHERE processed = ${sqlFalse}
+        AND ${sql.in("id", requestIds.map(String))}
+        `.pipe(
+          Effect.asVoid,
+          PersistenceError.refail,
+          withTracerDisabled
+        ),
+
     resetAddresses: (addresses) =>
       addresses.length === 0
         ? Effect.void

packages/effect/src/unstable/cluster/internal/entityManager.ts

@@ -55,6 +55,7 @@ export interface EntityManager {
 
   readonly isProcessingFor: (message: Message.Incoming<any>, options?: {
     readonly excludeReplies?: boolean
+    readonly excludeCompleted?: boolean
   }) => boolean
   readonly clearProcessed: () => void
 
@@ -90,6 +91,7 @@ export type EntityState = {
     readonly rpc: Rpc.AnyWithProps
     readonly message: Message.IncomingRequestLocal<any>
     sentReply: boolean
+    sentExit: boolean
     lastSentChunk: Option.Option<Reply.Chunk<Rpc.Any>>
     sequence: number
     /** Set when the request should not outlive its caller. */
@@ -233,6 +235,7 @@ export const make = Effect.fnUntraced(function*<
                 if (!request) return Effect.void
 
                 request.sentReply = true
+                request.sentExit = true
 
                 if (
                   isShuttingDown &&
@@ -253,6 +256,7 @@ export const make = Effect.fnUntraced(function*<
                   (isShuttingDown || isUninterruptibleForServer(request.message.annotations))
                 ) {
                   if (!isShuttingDown) {
+                    request.sentExit = false
                     return server.write(
                       0,
                       {
@@ -346,6 +350,7 @@ export const make = Effect.fnUntraced(function*<
             const request = activeRequests.get(id)
             if (!request) continue
             const { lastSentChunk, message } = request
+            request.sentExit = false
             yield* server.write(
               0,
               {
@@ -530,6 +535,7 @@ export const make = Effect.fnUntraced(function*<
                 rpc,
                 message,
                 sentReply: false,
+                sentExit: false,
                 lastSentChunk: Option.filter(
                   message.lastSentReply,
                   (reply): reply is Reply.Chunk<Rpc.Any> => reply._tag === "Chunk"
@@ -645,7 +651,10 @@ export const make = Effect.fnUntraced(function*<
         return Effect.flatMap(Fiber.joinAll(fibers), loop)
       }),
     isProcessingFor(message, options) {
-      if (options?.excludeReplies !== true && processedRequestIds.has(message.envelope.requestId)) {
+      if (
+        options?.excludeReplies !== true && options?.excludeCompleted !== true &&
+        processedRequestIds.has(message.envelope.requestId)
+      ) {
         return true
       }
       const state = activeServers.get(message.envelope.address.entityId)
@@ -655,6 +664,8 @@ export const make = Effect.fnUntraced(function*<
         return false
       } else if (options?.excludeReplies && request.sentReply) {
         return false
+      } else if (options?.excludeCompleted && request.sentExit) {
+        return false
       }
       return true
     },

packages/effect/test/cluster/MessageStorage.test.ts

@@ -104,6 +104,52 @@ describe("MessageStorage", () => {
         expect(messages).toHaveLength(0)
       }).pipe(Effect.provide(MessageStorage.MemoryDriver.layer)))
 
+    it.effect("resetRequests with no IDs leaves existing claims untouched", () =>
+      Effect.gen(function*() {
+        const storage = yield* MessageStorage.MessageStorage
+        const request = yield* makeRequest()
+        yield* storage.saveRequest(request)
+        expect(yield* storage.unprocessedMessages([request.envelope.address.shardId])).toHaveLength(1)
+        yield* storage.resetRequests([])
+        const driver = yield* MessageStorage.MemoryDriver
+        yield* driver.encoded.resetRequests([])
+        expect(yield* storage.unprocessedMessages([request.envelope.address.shardId])).toHaveLength(0)
+      }).pipe(Effect.provide(MemoryLayer)))
+
+    it.effect("resetRequests releases only the selected request at a shared address", () =>
+      Effect.gen(function*() {
+        const storage = yield* MessageStorage.MessageStorage
+        const selected = yield* makeRequest()
+        const unrelated = yield* makeRequest()
+        yield* storage.saveRequest(selected)
+        yield* storage.saveRequest(unrelated)
+        const shards = [selected.envelope.address.shardId]
+        expect(yield* storage.unprocessedMessages(shards)).toHaveLength(2)
+        yield* storage.resetRequests([selected.envelope.requestId])
+        const messages = yield* storage.unprocessedMessages(shards)
+        expect(messages.map((message) => message.envelope.requestId)).toEqual([selected.envelope.requestId])
+        expect(yield* storage.unprocessedMessages(shards)).toHaveLength(0)
+      }).pipe(Effect.provide(MemoryLayer)))
+
+    it.effect("resetRequests preserves chunk replies, exit replies, and completed state", () =>
+      Effect.gen(function*() {
+        const storage = yield* MessageStorage.MessageStorage
+        const streaming = yield* makeRequest({ rpc: StreamRpc, payload: StreamRpc.payloadSchema.make({ id: 123 }) })
+        const completed = yield* makeRequest()
+        yield* storage.saveRequest(streaming)
+        yield* storage.saveRequest(completed)
+        yield* storage.saveReply(yield* makeChunkReply(streaming))
+        yield* storage.saveReply(yield* makeReply(completed))
+        const replies = yield* storage.repliesFor([streaming, completed])
+        expect(replies).toHaveLength(2)
+        const shards = [streaming.envelope.address.shardId]
+        expect(yield* storage.unprocessedMessages(shards)).toHaveLength(1)
+        yield* storage.resetRequests([streaming.envelope.requestId, completed.envelope.requestId])
+        const messages = yield* storage.unprocessedMessages(shards)
+        expect(messages.map((message) => message.envelope.requestId)).toEqual([streaming.envelope.requestId])
+        expect(yield* storage.repliesFor([streaming, completed])).toEqual(replies)
+      }).pipe(Effect.provide(MemoryLayer)))
+
     it.effect("saves a request", () =>
       Effect.gen(function*() {
         const storage = yield* MessageStorage.MessageStorage