@effect-app/infra 4.0.0-beta.263 → 4.0.0-beta.265

package/src/ClusterCosmos.ts

@@ -1,4 +1,4 @@
-import type { OperationInput } from "@azure/cosmos"
+import type { OperationInput, PatchRequestBody } from "@azure/cosmos"
 import * as Arr from "effect-app/Array"
 import * as Effect from "effect-app/Effect"
 import * as Layer from "effect-app/Layer"
@@ -84,6 +84,17 @@ interface RunnerDoc {
   runner: string
   healthy: boolean
   lastHeartbeat: number
+  /** Unique, stable machine id (mirrors SQL's auto-increment `machine_id`). */
+  readonly machineId: number
+  readonly _etag?: string
+}
+
+interface MachineIdCounterDoc {
+  readonly id: string
+  readonly _partitionKey: "runner"
+  readonly type: "machine-id-counter"
+  value: number
+  readonly _etag?: string
 }
 
 interface LockDoc {
@@ -331,13 +342,7 @@ const replyFromDoc = (doc: ReplyDoc): Reply.Encoded =>
 const shardIdFromString = (shardId: string): Envelope.Encoded["address"]["shardId"] =>
   ShardId.fromStringEncoded(shardId)
 
-const makeMachineId = (address: string) => {
-  let hash = 0
-  for (let i = 0; i < address.length; i++) {
-    hash = Math.imul(31, hash) + address.charCodeAt(i) | 0
-  }
-  return Math.abs(hash)
-}
+const machineIdCounterDocId = cosmosId("machine-id-counter")
 
 const createContainer = (prefix: string) =>
   Effect.fnUntraced(function*() {
@@ -395,29 +400,26 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         )
 
   const markReplyAcked = (requestId: string, replyId: string) =>
-    Effect.tryPromise(() => container.item(cosmosId(replyId), replyPartition(requestId)).read<ReplyDoc>()).pipe(
-      Effect.flatMap((resp) => {
-        const doc = resp.resource
-        if (!doc) return Effect.void
-        doc.acked = true
-        return Effect
-          .tryPromise(() =>
-            container.item(cosmosId(replyId), replyPartition(requestId)).replace(doc, {
-              accessCondition: { type: "IfMatch", condition: doc._etag ?? "" }
-            })
-          )
-          .pipe(Effect.tap(annotateItem), Effect.asVoid)
-      }),
-      Effect.catchIf(isNotFound, () => Effect.void),
-      Effect.catchIf(isPreconditionFailed, () => Effect.void)
-    )
+    Effect
+      .tryPromise(() =>
+        container.item(cosmosId(replyId), replyPartition(requestId)).patch<ReplyDoc>([
+          { op: "set", path: "/acked", value: true }
+        ])
+      )
+      .pipe(
+        Effect.tap(annotateItem),
+        Effect.asVoid,
+        Effect.catchIf(isNotFound, () => Effect.void),
+        Effect.catchIf(isPreconditionFailed, () => Effect.void)
+      )
 
-  const replaceMessage = (doc: MessageDoc) =>
+  const claimMessageRead = (doc: MessageDoc, now: number) =>
     Effect
       .tryPromise(() =>
-        container.item(doc.id, doc._partitionKey).replace(doc, {
-          accessCondition: { type: "IfMatch", condition: doc._etag ?? "" }
-        })
+        container.item(doc.id, doc._partitionKey).patch<MessageDoc>(
+          [{ op: "set", path: "/lastRead", value: now }],
+          { accessCondition: { type: "IfMatch", condition: doc._etag ?? "" } }
+        )
       )
       .pipe(
         Effect.tap(annotateItem),
@@ -425,28 +427,27 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         Effect.catchIf(isPreconditionFailed, () => Effect.succeed(false))
       )
 
-  const markAckMessagesProcessed = (docs: ReadonlyArray<MessageDoc>) =>
+  const batchDocs = <A extends { readonly id: string; readonly _partitionKey: string }>(
+    docs: ReadonlyArray<A>,
+    operation: (doc: A) => OperationInput,
+    fallback: (doc: A) => Effect.Effect<void, unknown, never>
+  ): Effect.Effect<void, unknown, never> =>
     Effect.forEach(
       Arr.groupByT(docs, (doc) => doc._partitionKey),
       ([partitionKey, partitionDocs]) =>
         Effect.forEach(
           Arr.chunksOf(partitionDocs, maxCosmosBatchOperations),
           (chunk) => {
-            const operations: Array<OperationInput> = chunk.map((doc) => ({
-              operationType: "Patch" as const,
-              id: doc.id,
-              resourceBody: [{ op: "set" as const, path: "/processed", value: true }]
-            }))
+            const operations: Array<OperationInput> = chunk.map(operation)
             return Effect
               .tryPromise(() => container.items.batch(operations, partitionKey))
               .pipe(
                 Effect.tap(annotateItem),
                 Effect.flatMap((resp) => {
                   const failed = resp.result?.find((result) => !isSuccessfulStatus(result.statusCode))
-                  return failed === undefined
-                    ? Effect.void
-                    : Effect.fail(new Error(`cluster cosmos ack batch failed: ${failed.statusCode}`))
-                })
+                  return failed === undefined ? Effect.void : Effect.fail(failed.statusCode)
+                }),
+                Effect.catchIf(() => true, () => Effect.forEach(chunk, fallback, { discard: true }))
               )
           },
           { discard: true }
@@ -454,27 +455,50 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
       { discard: true }
     )
 
-  const updateMessage = (
-    doc: MessageDoc,
-    update: (doc: MessageDoc) => boolean
-  ): Effect.Effect<void, unknown> =>
-    Effect.suspend(function loop(current = doc): Effect.Effect<void, unknown> {
-      if (!update(current)) return Effect.void
-      return replaceMessage(current).pipe(
-        Effect.flatMap((replaced) =>
-          replaced
-            ? Effect.void
-            : readMessage(current.id, current._partitionKey).pipe(
-              Effect.flatMap((found) =>
-                Option.match(found, {
-                  onNone: () => Effect.void,
-                  onSome: loop
-                })
-              )
-            )
-        )
-      )
-    })
+  const patchDoc = <A extends { readonly id: string; readonly _partitionKey: string }>(
+    doc: A,
+    resourceBody: PatchRequestBody
+  ): Effect.Effect<void, unknown, never> =>
+    Effect.tryPromise(() => container.item(doc.id, doc._partitionKey).patch(resourceBody)).pipe(
+      Effect.tap(annotateItem),
+      Effect.asVoid,
+      Effect.catchIf(isNotFound, () => Effect.void)
+    )
+
+  const patchDocs = <A extends { readonly id: string; readonly _partitionKey: string }>(
+    docs: ReadonlyArray<A>,
+    resourceBody: (doc: A) => PatchRequestBody
+  ): Effect.Effect<void, unknown, never> =>
+    batchDocs(
+      docs,
+      (doc) => ({
+        operationType: "Patch" as const,
+        id: doc.id,
+        resourceBody: resourceBody(doc)
+      }),
+      (doc) => patchDoc(doc, resourceBody(doc))
+    )
+
+  const deleteDoc = <A extends { readonly id: string; readonly _partitionKey: string }>(
+    doc: A
+  ): Effect.Effect<void, unknown, never> =>
+    Effect.tryPromise(() => container.item(doc.id, doc._partitionKey).delete()).pipe(
+      Effect.tap(annotateItem),
+      Effect.asVoid,
+      Effect.catchIf(isNotFound, () => Effect.void)
+    )
+
+  const deleteDocs = <A extends { readonly id: string; readonly _partitionKey: string }>(
+    docs: ReadonlyArray<A>
+  ): Effect.Effect<void, unknown, never> =>
+    batchDocs(
+      docs,
+      (doc) => ({
+        operationType: "Delete" as const,
+        id: doc.id
+      }),
+      deleteDoc
+    )
 
   return yield* MessageStorage.makeEncoded({
     saveEnvelope: ({ deliverAt, envelope, primaryKey }) =>
@@ -487,7 +511,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
               "SELECT * FROM c WHERE c.type = 'message' AND c.kind = 'AckChunk' AND c.processed = false AND c.requestId = @requestId",
               [{ name: "@requestId", value: envelope.requestId }]
             )
-            yield* markAckMessagesProcessed(pendingAcks)
+            yield* patchDocs(pendingAcks, () => [{ op: "set", path: "/processed", value: true }])
           }
           return yield* Effect.tryPromise(() => container.items.create(doc)).pipe(
             Effect.tap(annotateItem),
@@ -525,17 +549,16 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
             "SELECT * FROM c WHERE c.type = 'message' AND c.requestId = @requestId",
             [{ name: "@requestId", value: reply.requestId }]
           )
-          yield* Effect.forEach(messages, (message) => {
-            return updateMessage(message, (doc) => {
-              if (reply._tag === "WithExit") {
-                doc.processed = true
-              } else if (doc.id !== reply.requestId && doc.kind !== "Request") {
-                return false
-              }
-              doc.lastReplyId = reply.id
-              return true
-            })
-          }, { discard: true })
+          const updatedMessages = reply._tag === "WithExit"
+            ? messages
+            : messages.filter((message) => message.id === reply.requestId || message.kind === "Request")
+          yield* patchDocs(updatedMessages, () =>
+            reply._tag === "WithExit"
+              ? [
+                { op: "set", path: "/processed", value: true },
+                { op: "set", path: "/lastReplyId", value: reply.id }
+              ]
+              : [{ op: "set", path: "/lastReplyId", value: reply.id }])
         })
         .pipe(annotate("saveReply"), refailPersistence, withTracerDisabled),
 
@@ -549,29 +572,20 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
               { name: "@requestId", value: id }
             ]
           )
-          yield* Effect.forEach(replies, (reply) =>
-            Effect
-              .tryPromise(() => container.item(reply.id, reply._partitionKey).delete())
-              .pipe(
-                Effect.tap(annotateItem),
-                Effect.catchIf(isNotFound, () => Effect.void)
-              ), { discard: true })
+          yield* deleteDocs(replies)
           const messages = yield* queryMessages(
             "SELECT * FROM c WHERE c.type = 'message' AND c.requestId = @requestId",
             [{ name: "@requestId", value: id }]
           )
-          yield* Effect.forEach(messages, (message) => {
-            if (message.kind === "Interrupt") {
-              return Effect.tryPromise(() => container.item(message.id, message._partitionKey).delete()).pipe(
-                Effect.tap(annotateItem),
-                Effect.catchIf(isNotFound, () => Effect.void)
-              )
-            }
-            message.processed = false
-            message.lastReplyId = null
-            message.lastRead = null
-            return replaceMessage(message).pipe(Effect.asVoid)
-          }, { discard: true })
+          yield* deleteDocs(messages.filter((message) => message.kind === "Interrupt"))
+          yield* patchDocs(
+            messages.filter((message) => message.kind !== "Interrupt"),
+            () => [
+              { op: "set", path: "/processed", value: false },
+              { op: "set", path: "/lastReplyId", value: null },
+              { op: "set", path: "/lastRead", value: null }
+            ]
+          )
         })
         .pipe(annotate("clearReplies"), refailPersistence, withTracerDisabled),
 
@@ -620,7 +634,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         ]
       )
         .pipe(
-          Effect.flatMap((docs) => collectUnprocessed(docs, now, lastReply, replaceMessage, queryReplies)),
+          Effect.flatMap((docs) => collectUnprocessed(docs, now, claimMessageRead, queryReplies)),
           annotate("unprocessedMessages"),
           refailPersistence,
           withTracerDisabled
@@ -635,7 +649,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         ]
       )
         .pipe(
-          Effect.flatMap((docs) => collectUnprocessed(docs, now, lastReply, replaceMessage, queryReplies)),
+          Effect.flatMap((docs) => collectUnprocessedById(docs, queryReplies)),
           annotate("unprocessedMessagesById"),
           refailPersistence,
           withTracerDisabled
@@ -651,12 +665,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         ]
       )
         .pipe(
-          Effect.flatMap((docs) =>
-            Effect.forEach(docs, (doc) => {
-              doc.lastRead = null
-              return replaceMessage(doc).pipe(Effect.asVoid)
-            }, { discard: true })
-          ),
+          Effect.flatMap((docs) => patchDocs(docs, () => [{ op: "set", path: "/lastRead", value: null }])),
           annotate("resetAddress"),
           refailPersistence,
           withTracerDisabled
@@ -671,37 +680,18 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         ]
       )
         .pipe(
-          Effect.flatMap((messages) =>
-            Effect.forEach(messages, (message) =>
-              queryReplies("SELECT * FROM c WHERE c.type = 'reply' AND c.requestId = @requestId", [
-                {
-                  name: "@requestId",
-                  value: message
-                    .requestId
-                }
-              ])
-                .pipe(
-                  Effect
-                    .flatMap((replies) =>
-                      Effect
-                        .forEach(replies, (reply) =>
-                          Effect
-                            .tryPromise(() =>
-                              container.item(reply.id, reply._partitionKey).delete()
-                            )
-                            .pipe(
-                              Effect.tap(annotateItem),
-                              Effect.catchIf(isNotFound, () => Effect.void)
-                            ), { discard: true })
-                    ),
-                  Effect.andThen(
-                    Effect.tryPromise(() => container.item(message.id, message._partitionKey).delete()).pipe(
-                      Effect.tap(annotateItem),
-                      Effect.catchIf(isNotFound, () => Effect.void)
-                    )
-                  )
-                ), { discard: true })
-          ),
+          Effect.flatMap((messages) => {
+            if (!Arr.isArrayNonEmpty(messages)) return Effect.void
+            const requestIds = Array.from(new Set(messages.map((message) => message.requestId)))
+            return queryReplies(
+              "SELECT * FROM c WHERE c.type = 'reply' AND ARRAY_CONTAINS(@requestIds, c.requestId)",
+              [{ name: "@requestIds", value: requestIds }]
+            )
+              .pipe(
+                Effect.flatMap((replies) => deleteDocs(replies)),
+                Effect.andThen(deleteDocs(messages))
+              )
+          }),
           annotate("clearAddress"),
           refailPersistence,
           withTracerDisabled
@@ -713,12 +703,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
         [{ name: "@shardIds", value: Array.from(shardIds) }]
       )
         .pipe(
-          Effect.flatMap((docs) =>
-            Effect.forEach(docs, (doc) => {
-              doc.lastRead = null
-              return replaceMessage(doc).pipe(Effect.asVoid)
-            }, { discard: true })
-          ),
+          Effect.flatMap((docs) => patchDocs(docs, () => [{ op: "set", path: "/lastRead", value: null }])),
           annotate("resetShards"),
           refailPersistence,
           withTracerDisabled
@@ -731,8 +716,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: {
 const collectUnprocessed = <E>(
   docs: ReadonlyArray<MessageDoc>,
   now: number,
-  lastReply: (replyId: string | null) => Effect.Effect<Option.Option<Reply.Encoded>, E>,
-  replaceMessage: (doc: MessageDoc) => Effect.Effect<boolean, E>,
+  claimMessageRead: (doc: MessageDoc, now: number) => Effect.Effect<boolean, E>,
   queryReplies: (
     query: string,
     parameters: ReadonlyArray<CosmosParameter>
@@ -743,22 +727,77 @@ const collectUnprocessed = <E>(
       readonly envelope: Envelope.Encoded
       readonly lastSentReply: Option.Option<Reply.Encoded>
     }> = []
+    const activeRequestIds = yield* activeReplyRequestIds(docs, queryReplies)
+    const lastReplies = yield* lastRepliesById(docs, activeRequestIds, queryReplies)
     for (const doc of docs) {
-      const replies = yield* queryReplies(
-        "SELECT * FROM c WHERE c.type = 'reply' AND c.requestId = @requestId AND (c.kind = 'WithExit' OR (c.kind = 'Chunk' AND c.acked = false))",
-        [{ name: "@requestId", value: doc.requestId }]
-      )
-      if (Arr.isArrayNonEmpty(replies)) continue
-      const sentReply = yield* lastReply(doc.lastReplyId)
-      doc.lastRead = now
-      const replaced = yield* replaceMessage(doc)
-      if (replaced) {
-        messages.push(envelopeFromDoc(doc, sentReply))
+      if (activeRequestIds.has(doc.requestId)) continue
+      const sentReply = Option.fromNullishOr(doc.lastReplyId === null ? undefined : lastReplies.get(doc.lastReplyId))
+      const claimed = yield* claimMessageRead(doc, now)
+      if (claimed) {
+        messages.push(envelopeFromDoc({ ...doc, lastRead: now }, sentReply))
       }
     }
     return messages
   })
 
+const collectUnprocessedById = <E>(
+  docs: ReadonlyArray<MessageDoc>,
+  queryReplies: (
+    query: string,
+    parameters: ReadonlyArray<CosmosParameter>
+  ) => Effect.Effect<Array<ReplyDoc>, E>
+) =>
+  Effect.gen(function*() {
+    const messages: Array<{
+      readonly envelope: Envelope.Encoded
+      readonly lastSentReply: Option.Option<Reply.Encoded>
+    }> = []
+    const activeRequestIds = yield* activeReplyRequestIds(docs, queryReplies)
+    const lastReplies = yield* lastRepliesById(docs, activeRequestIds, queryReplies)
+    for (const doc of docs) {
+      if (activeRequestIds.has(doc.requestId)) continue
+      const sentReply = Option.fromNullishOr(doc.lastReplyId === null ? undefined : lastReplies.get(doc.lastReplyId))
+      messages.push(envelopeFromDoc(doc, sentReply))
+    }
+    return messages
+  })
+
+const activeReplyRequestIds = <E>(
+  docs: ReadonlyArray<MessageDoc>,
+  queryReplies: (
+    query: string,
+    parameters: ReadonlyArray<CosmosParameter>
+  ) => Effect.Effect<Array<ReplyDoc>, E>
+) => {
+  const requestIds = Array.from(new Set(docs.map((doc) => doc.requestId)))
+  if (!Arr.isArrayNonEmpty(requestIds)) return Effect.succeed(new Set<string>())
+  return queryReplies(
+    "SELECT * FROM c WHERE c.type = 'reply' AND ARRAY_CONTAINS(@requestIds, c.requestId) AND (c.kind = 'WithExit' OR (c.kind = 'Chunk' AND c.acked = false))",
+    [{ name: "@requestIds", value: requestIds }]
+  )
+    .pipe(Effect.map((replies) => new Set(replies.map((reply) => reply.requestId))))
+}
+
+const lastRepliesById = <E>(
+  docs: ReadonlyArray<MessageDoc>,
+  activeRequestIds: ReadonlySet<string>,
+  queryReplies: (
+    query: string,
+    parameters: ReadonlyArray<CosmosParameter>
+  ) => Effect.Effect<Array<ReplyDoc>, E>
+) => {
+  const replyIds = Array.from(
+    new Set(docs
+      .flatMap((doc) => activeRequestIds.has(doc.requestId) || doc.lastReplyId === null ? [] : [doc.lastReplyId]))
+  )
+  if (!Arr.isArrayNonEmpty(replyIds)) return Effect.succeed(new Map<string, Reply.Encoded>())
+  return queryReplies(
+    "SELECT * FROM c WHERE c.type = 'reply' AND ARRAY_CONTAINS(@replyIds, c.rowid)",
+    [{ name: "@replyIds", value: replyIds }]
+  )
+    .pipe(Effect.map((replies) => new Map(replies.map((reply) => [reply.rowid, replyFromDoc(reply)]))))
+}
+
 export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
   readonly prefix?: string | undefined
 }) {
@@ -780,6 +819,52 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
       )
       .pipe(Effect.tap(annotateFeed), Effect.map((resp) => resp.resources))
 
+  const deleteRunner = (doc: RunnerDoc) =>
+    Effect
+      .tryPromise(() =>
+        container.item(doc.id, "runner").delete({
+          accessCondition: { type: "IfMatch", condition: doc._etag ?? "" }
+        })
+      )
+      .pipe(
+        Effect.tap(annotateItem),
+        Effect.asVoid,
+        Effect.catchIf(isNotFound, () => Effect.void),
+        Effect.catchIf(isPreconditionFailed, () => Effect.void)
+      )
+
+  const readRunner = (address: string) =>
+    Effect.tryPromise(() => container.item(runnerDocId(address), "runner").read<RunnerDoc>()).pipe(
+      Effect.tap(annotateItem),
+      Effect.map((resp) => Option.fromNullishOr(resp.resource)),
+      Effect.catchIf(isNotFound, () => Effect.succeed(Option.none()))
+    )
+
+  // Allocate a unique, monotonically increasing machine id via an atomic
+  // server-side `incr` on a single counter doc — mirrors SQL's auto-increment
+  // `machine_id` PK. A per-address string hash (the previous approach) could
+  // collide mod 1024 across distinct runners, yielding duplicate Snowflake ids.
+  // The counter doc is seeded at init below, so it always exists here.
+  const allocateMachineId = Effect
+    .tryPromise(() =>
+      container.item(machineIdCounterDocId, "runner").patch<MachineIdCounterDoc>([
+        { op: "incr", path: "/value", value: 1 }
+      ])
+    )
+    .pipe(Effect.tap(annotateItem), Effect.map((resp) => resp.resource!.value))
+
+  // Seed the counter once (create-if-missing); a lost create race is fine.
+  yield* Effect
+    .tryPromise(() =>
+      container.items.create<MachineIdCounterDoc>({
+        id: machineIdCounterDocId,
+        _partitionKey: "runner",
+        type: "machine-id-counter",
+        value: 0
+      })
+    )
+    .pipe(Effect.catchIf(isConflict, () => Effect.void), Effect.orDie)
+
   const readLock = (shardId: string) =>
     Effect.tryPromise(() => container.item(lockDocId(shardId), "lock").read<LockDoc>()).pipe(
       Effect.tap(annotateItem),
@@ -800,6 +885,20 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
         Effect.catchIf(isPreconditionFailed, () => Effect.succeed(false))
       )
 
+  const deleteLock = (doc: LockDoc) =>
+    Effect
+      .tryPromise(() =>
+        container.item(doc.id, "lock").delete({
+          accessCondition: { type: "IfMatch", condition: doc._etag ?? "" }
+        })
+      )
+      .pipe(
+        Effect.tap(annotateItem),
+        Effect.asVoid,
+        Effect.catchIf(isNotFound, () => Effect.void),
+        Effect.catchIf(isPreconditionFailed, () => Effect.void)
+      )
+
   const createLock = (address: string, shardId: string, now: number) =>
     Effect
       .tryPromise(() =>
@@ -850,9 +949,17 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
     ),
 
     register: (address, runner, healthy) =>
-      Effect.sync(() => Date.now()).pipe(
-        Effect.flatMap((now) =>
-          Effect
+      Effect
+        .gen(function*() {
+          const now = Date.now()
+          // Reuse a previously-allocated machine id for this address (stable per
+          // address, like SQL's RETURNING on conflict); allocate a fresh unique
+          // one on first registration.
+          const existing = yield* readRunner(address)
+          const machineId = Option.isSome(existing) && existing.value.machineId !== undefined
+            ? existing.value.machineId
+            : yield* allocateMachineId
+          yield* Effect
             .tryPromise(() =>
               container.items.upsert<RunnerDoc>({
                 id: runnerDocId(address),
@@ -861,40 +968,47 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
                 address,
                 runner,
                 healthy,
-                lastHeartbeat: now
+                lastHeartbeat: now,
+                machineId
               })
             )
             .pipe(Effect.tap(annotateItem))
-        ),
-        Effect.as(makeMachineId(address)),
-        annotate("register"),
-        refailPersistence,
-        withTracerDisabled
-      ),
+          return machineId
+        })
+        .pipe(annotate("register"), refailPersistence, withTracerDisabled),
 
     unregister: (address) =>
-      Effect.tryPromise(() => container.item(runnerDocId(address), "runner").delete()).pipe(
-        Effect.tap(annotateItem),
-        Effect.catchIf(isNotFound, () => Effect.void),
+      Effect.sync(() => Date.now()).pipe(
+        Effect.flatMap((now) =>
+          queryRunners(
+            "SELECT * FROM c WHERE c.type = 'runner' AND (c.address = @address OR c.lastHeartbeat < @expiresAt)",
+            [
+              { name: "@address", value: address },
+              { name: "@expiresAt", value: now - expires }
+            ]
+          )
+        ),
+        Effect.flatMap((docs) => Effect.forEach(docs, deleteRunner, { discard: true })),
         annotate("unregister"),
         refailPersistence,
         withTracerDisabled
       ),
 
     setRunnerHealth: (address, healthy) =>
-      Effect.tryPromise(() => container.item(runnerDocId(address), "runner").read<RunnerDoc>()).pipe(
-        Effect.flatMap((resp) => {
-          const doc = resp.resource
-          if (!doc) return Effect.void
-          doc.healthy = healthy
-          return Effect.tryPromise(() => container.item(doc.id, "runner").replace(doc)).pipe(Effect.tap(annotateItem))
-        }),
-        Effect.asVoid,
-        Effect.catchIf(isNotFound, () => Effect.void),
-        annotate("setRunnerHealth"),
-        refailPersistence,
-        withTracerDisabled
-      ),
+      Effect
+        .tryPromise(() =>
+          container.item(runnerDocId(address), "runner").patch<RunnerDoc>([
+            { op: "set", path: "/healthy", value: healthy }
+          ])
+        )
+        .pipe(
+          Effect.tap(annotateItem),
+          Effect.asVoid,
+          Effect.catchIf(isNotFound, () => Effect.void),
+          annotate("setRunnerHealth"),
+          refailPersistence,
+          withTracerDisabled
+        ),
 
     acquire: (address, shardIds) =>
       Effect.sync(() => Date.now()).pipe(
@@ -909,17 +1023,17 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
       Effect
         .gen(function*() {
           const now = Date.now()
-          yield* Effect.tryPromise(() => container.item(runnerDocId(address), "runner").read<RunnerDoc>()).pipe(
-            Effect.flatMap((resp) => {
-              const doc = resp.resource
-              if (!doc) return Effect.void
-              doc.lastHeartbeat = now
-              return Effect.tryPromise(() => container.item(doc.id, "runner").replace(doc)).pipe(
-                Effect.tap(annotateItem)
-              )
-            }),
-            Effect.catchIf(isNotFound, () => Effect.void)
-          )
+          yield* Effect
+            .tryPromise(() =>
+              container.item(runnerDocId(address), "runner").patch<RunnerDoc>([
+                { op: "set", path: "/lastHeartbeat", value: now }
+              ])
+            )
+            .pipe(
+              Effect.tap(annotateItem),
+              Effect.asVoid,
+              Effect.catchIf(isNotFound, () => Effect.void)
+            )
           const refreshed = yield* Effect.forEach(shardIds, (shardId) =>
             readLock(shardId).pipe(
               Effect.flatMap((lock) =>
@@ -944,11 +1058,7 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
             onNone: () => Effect.void,
             onSome: (doc) =>
               doc.address === address
-                ? Effect.tryPromise(() => container.item(doc.id, "lock").delete()).pipe(
-                  Effect.tap(annotateItem),
-                  Effect.catchIf(isNotFound, () => Effect.void),
-                  Effect.asVoid
-                )
+                ? deleteLock(doc)
                 : Effect.void
           })
         ),
@@ -970,13 +1080,7 @@ export const makeRunnerStorage = Effect.fnUntraced(function*(options?: {
         )
         .pipe(
           Effect.tap(annotateFeed),
-          Effect.flatMap((resp) =>
-            Effect.forEach(resp.resources, (doc) =>
-              Effect.tryPromise(() => container.item(doc.id, "lock").delete()).pipe(
-                Effect.tap(annotateItem),
-                Effect.catchIf(isNotFound, () => Effect.void)
-              ), { discard: true })
-          ),
+          Effect.flatMap((resp) => Effect.forEach(resp.resources, deleteLock, { discard: true })),
           annotate("releaseAll"),
           refailPersistence,
           withTracerDisabled