effect 4.0.0-rc.118 → 4.0.1

package/src/cluster/ClusterWorkflowEngine.ts

@@ -299,7 +299,14 @@ export const make = Effect.gen(function*() {
     )
 
     if (Option.isNone(maybeSuspended)) return
-    yield* sharding.reset(Snowflake.Snowflake(maybeSuspended.value.requestId))
+    if (
+      !(yield* sharding.reset(
+        Snowflake.Snowflake(maybeSuspended.value.requestId),
+        { expectedReplyId: Snowflake.Snowflake(maybeSuspended.value.id) }
+      ))
+    ) {
+      return yield* Effect.fail(new Error("Failed to reset workflow run"))
+    }
     yield* sharding.pollStorage
   })
 
@@ -404,7 +411,7 @@ export const make = Effect.gen(function*() {
             // Replays reuse the request id but track a fresh set of awaits.
             let currentRun: {
               readonly request: Entity.Request<any>
-              readonly awaitedDeferreds: ReadonlySet<string>
+              readonly instance: WorkflowEngine.WorkflowInstance["Service"]
             } | undefined
             // Concurrent wakes share one wait for the current run to publish its reply.
             const resumeGate = Semaphore.makeUnsafe(1)
@@ -421,7 +428,7 @@ export const make = Effect.gen(function*() {
             return {
               run: (request: Entity.Request<any>) => {
                 const instance = WorkflowEngine.WorkflowInstance.initial(workflow, executionId)
-                currentRun = { request, awaitedDeferreds: instance.awaitedDeferreds }
+                currentRun = { request, instance }
                 const parent = (request.payload as any)[payloadParentKey] as
                   | { workflowName: string; executionId: string }
                   | undefined
@@ -515,9 +522,13 @@ export const make = Effect.gen(function*() {
                 // registered yet, a later await will see the result just stored.
                 // Check when the wake runs and skip preemption too: interrupting
                 // a later await would require waiting for its reply before reset.
+                // A run completing its own deferred also reads it from the cache
+                // when it awaits later, so that completion never wakes the run.
                 const wake = Effect.suspend(() =>
-                  currentRun && !currentRun.awaitedDeferreds.has(payload.name)
-                    ? ensureSuccess(resume(workflow, executionId))
+                  currentRun && !currentRun.instance.awaitedDeferreds.has(payload.name)
+                    ? currentRun.instance.completedDeferreds.has(payload.name)
+                      ? Effect.void
+                      : ensureSuccess(resume(workflow, executionId))
                     : deferredState.deferredDone(executionId, payload.name).pipe(Effect.andThen(resumeCurrentRun))
                 )
                 // An asynchronous reply releases the RPC concurrency permit while
@@ -696,6 +707,16 @@ export const make = Effect.gen(function*() {
 
     deferredDone: Effect.fnUntraced(
       function*({ deferredName, executionId, exit, workflowName }) {
+        const instance = yield* Effect.serviceOption(WorkflowEngine.WorkflowInstance)
+        // Activities run with their own instance, so their completions still
+        // wake a run that suspended before reaching the await. An interrupt
+        // must wake the run even when it signals itself.
+        if (
+          Option.isSome(instance) && instance.value.executionId === executionId &&
+          deferredName !== InterruptSignal.name
+        ) {
+          instance.value.completedDeferreds.add(deferredName)
+        }
         const workflow = workflows.get(workflowName)
         if (workflow) {
           return yield* Effect.orDie(sendDiscard({

package/src/cluster/MessageStorage.ts

@@ -74,9 +74,16 @@ export class MessageStorage extends Context.Service<MessageStorage, {
   ) => Effect.Effect<void, PersistenceError | MalformedMessage>
 
   /**
-   * Clear the `Reply`s for the given request id.
+   * Clear replies for the given request id. When options.expectedReplyId is
+   * provided, clear only if it is still the latest reply at the storage boundary.
+   * Without it, clear unconditionally.
+   * Custom storage implementations must honor this condition to protect
+   * workflows from stale concurrent resumes.
    */
-  readonly clearReplies: (requestId: Snowflake.Snowflake) => Effect.Effect<void, PersistenceError>
+  readonly clearReplies: (
+    requestId: Snowflake.Snowflake,
+    options?: { readonly expectedReplyId?: Snowflake.Snowflake | undefined }
+  ) => Effect.Effect<void, PersistenceError>
 
   /**
    * Retrieves the replies for the specified requests.
@@ -350,9 +357,16 @@ export type Encoded = {
   readonly saveReply: (reply: Reply.Encoded) => Effect.Effect<void, PersistenceError>
 
   /**
-   * Remove the replies for the specified request.
+   * Remove the replies for the specified request. If options.expectedReplyId
+   * is provided, compare it with the latest reply atomically before clearing.
+   * Without it, clear unconditionally.
+   * Custom storage implementations must honor this condition to protect
+   * workflows from stale concurrent resumes.
    */
-  readonly clearReplies: (requestId: Snowflake.Snowflake) => Effect.Effect<void, PersistenceError>
+  readonly clearReplies: (
+    requestId: Snowflake.Snowflake,
+    options?: { readonly expectedReplyId?: Snowflake.Snowflake | undefined }
+  ) => Effect.Effect<void, PersistenceError>
 
   /**
    * Retrieves the request id for the specified primary key.
@@ -1054,10 +1068,13 @@ export class MemoryDriver extends Context.Service<MemoryDriver>()("effect/cluste
           replyIds.add(reply.id)
           replyLatch.openUnsafe()
         }),
-      clearReplies: (id) =>
+      clearReplies: (id, options) =>
         Effect.sync(() => {
           const entry = requests.get(String(id))
           if (!entry) return
+          if (
+            options?.expectedReplyId !== undefined && entry.replies.at(-1)?.id !== String(options.expectedReplyId)
+          ) return
           entry.replies = []
           entry.lastReceivedChunk = undefined
           unprocessed.add(entry.envelope)

package/src/cluster/Sharding.ts

@@ -195,9 +195,13 @@ export class Sharding extends Context.Service<Sharding, {
   >
 
   /**
-   * Reset the state of a message
+   * Reset the state of a message. When options.expectedReplyId is provided,
+   * reset only if it is still the latest reply. Otherwise reset unconditionally.
    */
-  readonly reset: (requestId: Snowflake.Snowflake) => Effect.Effect<boolean>
+  readonly reset: (
+    requestId: Snowflake.Snowflake,
+    options?: { readonly expectedReplyId?: Snowflake.Snowflake | undefined }
+  ) => Effect.Effect<boolean>
 
   /**
    * Trigger a storage read, which will read all unprocessed messages.
@@ -545,30 +549,27 @@ const make = Effect.gen(function*() {
       return Effect.logInfo("Shard lock storage has recovered")
     })
 
-    const refreshShardLocks = Effect.suspend(() =>
-      runnerStorage.refresh(selfAddress, [
-        ...acquiredShards,
-        ...releasingShards
-      ])
-    ).pipe(
-      Effect.flatMap((acquired) => {
-        for (const shardId of acquiredShards) {
-          if (!acquired.includes(shardId)) {
-            MutableHashSet.remove(acquiredShards, shardId)
-            MutableHashSet.add(releasingShards, shardId)
-          }
+    const refreshShardLocks = Effect.gen(function*() {
+      const refreshed = [...acquiredShards, ...releasingShards]
+      const acquired = yield* runnerStorage.refresh(selfAddress, refreshed)
+      // Shards acquired during the request are absent from its response.
+      for (const shardId of refreshed) {
+        if (MutableHashSet.has(acquiredShards, shardId) && !acquired.includes(shardId)) {
+          MutableHashSet.remove(acquiredShards, shardId)
+          MutableHashSet.add(releasingShards, shardId)
         }
-        for (let i = 0; i < acquired.length; i++) {
-          const shardId = acquired[i]
-          if (!MutableHashSet.has(selfShards, shardId)) {
-            MutableHashSet.remove(acquiredShards, shardId)
-            MutableHashSet.add(releasingShards, shardId)
-          }
+      }
+      for (let i = 0; i < acquired.length; i++) {
+        const shardId = acquired[i]
+        if (!MutableHashSet.has(selfShards, shardId)) {
+          MutableHashSet.remove(acquiredShards, shardId)
+          MutableHashSet.add(releasingShards, shardId)
         }
-        return MutableHashSet.size(releasingShards) > 0
-          ? activeShardsLatch.open
-          : Effect.void
-      }),
+      }
+      if (MutableHashSet.size(releasingShards) > 0) {
+        yield* activeShardsLatch.open
+      }
+    }).pipe(
       Effect.retry({
         times: 5,
         schedule: Schedule.spaced(50)
@@ -1208,8 +1209,8 @@ const make = Effect.gen(function*() {
     )
   }
 
-  const reset: Sharding["Service"]["reset"] = (requestId) =>
-    Effect.matchCause(storage.clearReplies(requestId), {
+  const reset: Sharding["Service"]["reset"] = (requestId, options) =>
+    Effect.matchCause(storage.clearReplies(requestId, options), {
       onSuccess: () => true,
       onFailure: () => false
     })
@@ -1458,8 +1459,10 @@ const make = Effect.gen(function*() {
               // for durable messages, we ignore interrupts on shutdown or as a
               // result of a shard being resassigned
               const caller = Context.getOption(entry.context, CurrentAddress).valueOrUndefined
+              const singletonShard = Context.getOption(entry.context, SingletonShardTag).valueOrUndefined
               const isTransientInterrupt = MutableRef.get(isShutdown) ||
-                (caller !== undefined && ActiveTeardown.isActive(caller))
+                (caller !== undefined && ActiveTeardown.isActive(caller)) ||
+                (singletonShard !== undefined && ActiveTeardown.isShardActive(singletonShard))
               if (isTransientInterrupt && Context.get(entry.message.annotations, Persisted)) {
                 return Effect.void
               }
@@ -1588,7 +1591,7 @@ const make = Effect.gen(function*() {
         Effect.andThen(Effect.never),
         Effect.scoped,
         Effect.provideService(CurrentLogAnnotations, {}),
-        Effect.provideContext(services),
+        Effect.provideContext(Context.add(services, SingletonShardTag, address.shardId)),
         Effect.orDie,
         Effect.interruptible
       ) as Effect.Effect<never>
@@ -1851,3 +1854,4 @@ export const layer: Layer.Layer<
 // Utilities
 
 const ClientAddressTag = Context.Service<EntityAddress>("effect/cluster/Sharding/ClientAddress")
+const SingletonShardTag = Context.Service<ShardId>("effect/cluster/Sharding/SingletonShard")

package/src/cluster/SqlMessageStorage.ts

@@ -29,6 +29,7 @@ import * as Migrator from "../sql/Migrator.ts"
 import * as SqlClient from "../sql/SqlClient.ts"
 import type { Row } from "../sql/SqlConnection.ts"
 import { isSqlError, type SqlError } from "../sql/SqlError.ts"
+import type { Fragment } from "../sql/Statement.ts"
 import { PersistenceError } from "./ClusterError.ts"
 import type * as EntityAddress from "./EntityAddress.ts"
 import type * as Envelope from "./Envelope.ts"
@@ -494,6 +495,28 @@ export const makeEncoded: (options?: {
       ? Effect.succeed([])
       : getUnprocessedMessagesForDialect(shardIds, now, unprocessedFilters(options))
 
+  const resetUnprocessed = sql.onDialectOrElse({
+    // Match the claim query's lock order to avoid deadlocks during resets.
+    pg: () => (filter: Fragment) =>
+      sql`
+      WITH messages AS (
+        SELECT id FROM ${messagesTableSql}
+        WHERE processed = ${sqlFalse} AND ${filter}
+        ORDER BY rowid ASC
+        FOR UPDATE
+      )
+      UPDATE ${messagesTableSql}
+      SET last_read = NULL
+      WHERE id IN (SELECT id FROM messages)
+    `,
+    orElse: () => (filter: Fragment) =>
+      sql`
+      UPDATE ${messagesTableSql}
+      SET last_read = NULL
+      WHERE processed = ${sqlFalse} AND ${filter}
+    `
+  })
+
   const encoded: MessageStorage.Encoded = {
     saveEnvelope: ({ deliverAt, envelope, primaryKey }) =>
       Effect.suspend(() => {
@@ -576,14 +599,25 @@ export const makeEncoded: (options?: {
       ),
 
     clearReplies: Effect.fnUntraced(
-      function*(requestId) {
-        yield* sql`DELETE FROM ${repliesTableSql} WHERE request_id = ${String(requestId)} AND kind = 0`
-        yield* sql`DELETE FROM ${messagesTableSql} WHERE request_id = ${
-          String(requestId)
-        } AND kind = ${messageKindInterrupt}`
-        yield* sql`UPDATE ${messagesTableSql} SET processed = ${sqlFalse}, last_reply_id = NULL, last_read = NULL WHERE request_id = ${
-          String(requestId)
-        }`
+      function*(requestId, options) {
+        const expectedReplyId = options?.expectedReplyId
+        if (expectedReplyId === undefined) {
+          yield* sql`DELETE FROM ${repliesTableSql} WHERE request_id = ${String(requestId)} AND kind = 0`
+          yield* sql`DELETE FROM ${messagesTableSql} WHERE request_id = ${
+            String(requestId)
+          } AND kind = ${messageKindInterrupt}`
+          yield* sql`UPDATE ${messagesTableSql} SET processed = ${sqlFalse}, last_reply_id = NULL, last_read = NULL WHERE request_id = ${
+            String(requestId)
+          }`
+          return
+        }
+        const id = String(requestId)
+        const expected = String(expectedReplyId)
+        // Check at the storage boundary, not at the earlier workflow read.
+        yield* sql`DELETE FROM ${repliesTableSql} WHERE request_id = ${id} AND kind = 0
+          AND EXISTS (SELECT 1 FROM ${messagesTableSql} WHERE id = ${id} AND last_reply_id = ${expected})`
+        yield* sql`UPDATE ${messagesTableSql} SET processed = ${sqlFalse}, last_reply_id = NULL, last_read = NULL
+          WHERE request_id = ${id} AND last_reply_id = ${expected}`
       },
       sql.withTransaction,
       PersistenceError.refail,
@@ -705,11 +739,7 @@ export const makeEncoded: (options?: {
     resetAddresses: (addresses) =>
       addresses.length === 0
         ? Effect.void
-        : sql`
-        UPDATE ${messagesTableSql}
-        SET last_read = NULL
-        WHERE processed = ${sqlFalse}
-        AND (${
+        : resetUnprocessed(sql`(${
           sql.or(
             groupAddresses(addresses).map(
               (group) =>
@@ -720,8 +750,7 @@ export const makeEncoded: (options?: {
                 ])
             )
           )
-        })
-        `.pipe(
+        })`).pipe(
           Effect.asVoid,
           PersistenceError.refail,
           withTracerDisabled
@@ -750,12 +779,7 @@ export const makeEncoded: (options?: {
       ),
 
     resetShards: (shardIds) =>
-      sql`
-        UPDATE ${messagesTableSql}
-        SET last_read = NULL
-        WHERE processed = ${sqlFalse}
-        AND shard_id IN (${sql.literal(shardIds.map(wrapString).join(","))})
-      `.pipe(
+      resetUnprocessed(sql`shard_id IN (${sql.literal(shardIds.map(wrapString).join(","))})`).pipe(
         Effect.asVoid,
         PersistenceError.refail,
         withTracerDisabled

package/src/cluster/SqlRunnerStorage.ts

@@ -11,13 +11,14 @@
  * @since 4.0.0
  */
 import * as Arr from "../Array.ts"
+import * as Cause from "../Cause.ts"
 import * as Duration from "../Duration.ts"
 import * as Effect from "../Effect.ts"
-import * as Exit from "../Exit.ts"
 import * as Fiber from "../Fiber.ts"
 import * as Layer from "../Layer.ts"
 import * as Scope from "../Scope.ts"
 import * as SqlClient from "../sql/SqlClient.ts"
+import type { Connection } from "../sql/SqlConnection.ts"
 import type { SqlError } from "../sql/SqlError.ts"
 import type * as Statement from "../sql/Statement.ts"
 import { PersistenceError } from "./ClusterError.ts"
@@ -156,47 +157,28 @@ export const make = Effect.fnUntraced(function*(options: {
     )
   })
 
-  let lockConnRebuilding = false
-  // Incremented every time the reserved connection is replaced, so failures
-  // from operations that ran on an already replaced connection do not trigger
-  // another rebuild.
-  let lockConnGeneration = 0
-  const rebuildLockConn = (generation: number) => {
-    if (
-      !lockConn ||
-      lockConnRebuilding ||
-      generation !== lockConnGeneration ||
-      lockConn.state.current._tag === "Closed"
-    ) return Effect.void
-    lockConnRebuilding = true
-    // The rebuild starts by closing the previous scope, releasing the
-    // unresponsive connection back to the pool. Bound it with `withDeadline`
-    // so a release that never completes cannot leave `lockConnRebuilding` set
-    // forever, which would disable every subsequent rebuild.
-    return withDeadline(lockConn.rebuildUnsafe()).pipe(
+  const rebuildLockConn = (from?: readonly [Connection, number]) => {
+    const rebuild = lockConn?.rebuildUnsafe({ from })
+    if (!rebuild) return Effect.void
+    // Bound release so a stalled connection cannot block subsequent rebuilds.
+    return withDeadline(rebuild).pipe(
       Effect.exit,
-      Effect.tap((exit) =>
-        Effect.sync(() => {
-          if (Exit.isSuccess(exit) && lockConn.state.current._tag === "Acquired") {
-            lockConnGeneration++
-          }
-        })
-      ),
-      Effect.ensuring(Effect.sync(() => {
-        lockConnRebuilding = false
-      })),
       Effect.forkIn(layerScope, { startImmediately: true }),
       Effect.asVoid
     )
   }
-  // Rebuild the reserved connection when `effect` fails on it. Failures keep
-  // scheduling rebuilds, so a rebuilt connection that is also unresponsive is
-  // replaced again.
-  const onErrorRebuildLockConn = <A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<A, E, R> =>
-    Effect.suspend(() => {
-      const generation = lockConnGeneration
-      return Effect.onError(effect, () => rebuildLockConn(generation))
-    })
+  // SQL failures replace the held connection; deadlines are handled by withLockOperationDeadline.
+  const useLockConn = <A, E, R>(
+    f: (conn: Connection, pid: number) => Effect.Effect<A, E, R>
+  ): Effect.Effect<A, E, R> =>
+    Effect.flatMap(
+      lockConn!.await,
+      (held) =>
+        Effect.onError(
+          f(held[0], held[1]),
+          (cause) => Cause.hasInterruptsOnly(cause) ? Effect.void : rebuildLockConn(held)
+        )
+    )
 
   const runnersTable = table("runners")
   const runnersTableSql = sql(runnersTable)
@@ -378,30 +360,21 @@ export const make = Effect.fnUntraced(function*(options: {
   const execWithLockConn = <A>(effect: Statement.Statement<A>): Effect.Effect<unknown, SqlError> => {
     if (!lockConn) return effect
     const [query, params] = effect.compile()
-    return lockConn.await.pipe(
-      Effect.flatMap(([conn]) => conn.executeRaw(query, params)),
-      onErrorRebuildLockConn
-    )
+    return useLockConn((conn) => conn.executeRaw(query, params))
   }
   const execWithLockConnUnprepared = <A>(
     effect: Statement.Statement<A>
   ): Effect.Effect<ReadonlyArray<ReadonlyArray<any>>, SqlError> => {
     if (!lockConn) return effect.values
     const [query, params] = effect.compile()
-    return lockConn.await.pipe(
-      Effect.flatMap(([conn]) => conn.executeUnprepared(query, params, undefined)),
-      onErrorRebuildLockConn
-    )
+    return useLockConn((conn) => conn.executeUnprepared(query, params, undefined))
   }
   const execWithLockConnValues = <A>(
     effect: Statement.Statement<A>
   ): Effect.Effect<ReadonlyArray<ReadonlyArray<any>>, SqlError> => {
     if (!lockConn) return effect.values
     const [query, params] = effect.compile()
-    return lockConn.await.pipe(
-      Effect.flatMap(([conn]) => conn.executeValues(query, params)),
-      onErrorRebuildLockConn
-    )
+    return useLockConn((conn) => conn.executeValues(query, params))
   }
 
   const acquireLock = sql.onDialectOrElse({
@@ -425,33 +398,33 @@ export const make = Effect.fnUntraced(function*(options: {
           )
         }
       }
-      return Effect.fnUntraced(function*(_address: string, shardIds: ReadonlyArray<string>) {
-        const [conn, pid] = yield* lockConn!.await
-        const acquiredShardIds: Array<string> = []
-        const toAcquire = new Map(shardIds.map((shardId) => [lockNumbers.get(shardId)!, shardId]))
-        const takenLocks = yield* conn.executeValues(
-          `SELECT objid FROM pg_locks WHERE locktype = 'advisory' AND granted = true AND classid = ${pgLockNamespaceOid} AND objsubid = 2 AND pid = ${pid} ORDER BY objid`,
-          []
-        )
-        for (let i = 0; i < takenLocks.length; i++) {
-          const lockNum = takenLocks[i][0] as number
-          const shardId = toAcquire.get(lockNum)
-          if (shardId === undefined) continue
-          acquiredShardIds.push(shardId)
-          toAcquire.delete(lockNum)
-        }
-        if (toAcquire.size === 0) {
-          return acquiredShardIds
-        }
-        const rows = yield* conn.executeUnprepared(`SELECT ${pgLocks(toAcquire)}`, [], undefined)
-        const results = rows[0] as Record<string, boolean>
-        for (const shardId in results) {
-          if (results[shardId]) {
+      return (_address: string, shardIds: ReadonlyArray<string>) =>
+        useLockConn(Effect.fnUntraced(function*(conn, pid) {
+          const acquiredShardIds: Array<string> = []
+          const toAcquire = new Map(shardIds.map((shardId) => [lockNumbers.get(shardId)!, shardId]))
+          const takenLocks = yield* conn.executeValues(
+            `SELECT objid FROM pg_locks WHERE locktype = 'advisory' AND granted = true AND classid = ${pgLockNamespaceOid} AND objsubid = 2 AND pid = ${pid} ORDER BY objid`,
+            []
+          )
+          for (let i = 0; i < takenLocks.length; i++) {
+            const lockNum = takenLocks[i][0] as number
+            const shardId = toAcquire.get(lockNum)
+            if (shardId === undefined) continue
             acquiredShardIds.push(shardId)
+            toAcquire.delete(lockNum)
           }
-        }
-        return acquiredShardIds
-      }, onErrorRebuildLockConn)
+          if (toAcquire.size === 0) {
+            return acquiredShardIds
+          }
+          const rows = yield* conn.executeUnprepared(`SELECT ${pgLocks(toAcquire)}`, [], undefined)
+          const results = rows[0] as Record<string, boolean>
+          for (const shardId in results) {
+            if (results[shardId]) {
+              acquiredShardIds.push(shardId)
+            }
+          }
+          return acquiredShardIds
+        }))
     },
 
     mysql: () => {
@@ -471,31 +444,31 @@ export const make = Effect.fnUntraced(function*(options: {
           )
         }
       }
-      return Effect.fnUntraced(function*(_address: string, shardIds: ReadonlyArray<string>) {
-        const [conn, pid] = yield* lockConn!.await
-        const takenLocks = (yield* conn.executeValues(`SELECT ${allMySqlTakenLocks}`, []))[0] as Array<number | null>
-        const acquiredShardIds: Array<string> = []
-        const toAcquire: Array<string> = []
-        for (let i = 0; i < shardIds.length; i++) {
-          const shardId = shardIds[i]
-          const lockTakenBy = takenLocks[shardIdsIndex.get(shardId)!]
-          if (lockTakenBy === pid) {
-            acquiredShardIds.push(shardId)
-          } else if (shardIds.includes(shardId)) {
-            toAcquire.push(shardId)
+      return (_address: string, shardIds: ReadonlyArray<string>) =>
+        useLockConn(Effect.fnUntraced(function*(conn, pid) {
+          const takenLocks = (yield* conn.executeValues(`SELECT ${allMySqlTakenLocks}`, []))[0] as Array<number | null>
+          const acquiredShardIds: Array<string> = []
+          const toAcquire: Array<string> = []
+          for (let i = 0; i < shardIds.length; i++) {
+            const shardId = shardIds[i]
+            const lockTakenBy = takenLocks[shardIdsIndex.get(shardId)!]
+            if (lockTakenBy === pid) {
+              acquiredShardIds.push(shardId)
+            } else if (shardIds.includes(shardId)) {
+              toAcquire.push(shardId)
+            }
           }
-        }
-        if (toAcquire.length === 0) {
-          return acquiredShardIds
-        }
-        const results = (yield* conn.executeValues(`SELECT ${mysqlLocks(toAcquire)}`, []))[0] as Array<number>
-        for (let i = 0; i < results.length; i++) {
-          if (results[i] === 1) {
-            acquiredShardIds.push(toAcquire[i])
+          if (toAcquire.length === 0) {
+            return acquiredShardIds
           }
-        }
-        return acquiredShardIds
-      }, onErrorRebuildLockConn)
+          const results = (yield* conn.executeValues(`SELECT ${mysqlLocks(toAcquire)}`, []))[0] as Array<number>
+          for (let i = 0; i < results.length; i++) {
+            if (results[i] === 1) {
+              acquiredShardIds.push(toAcquire[i])
+            }
+          }
+          return acquiredShardIds
+        }))
     },
 
     mssql: () => (address: string, shardIds: ReadonlyArray<string>) => {
@@ -647,8 +620,16 @@ export const make = Effect.fnUntraced(function*(options: {
       `.pipe(execWithLockConnValues, Effect.map((rows) => rows.map((row) => row[0] as string)))
   })
 
+  // On failure, replace the connection ready at entry. If none was ready,
+  // retry only a failed rebuild, leaving pending or newly ready connections alone.
   const withLockOperationDeadline = <A, E, R>(operation: Effect.Effect<A, E, R>) =>
-    onErrorRebuildLockConn(withDeadline(operation))
+    Effect.suspend(() => {
+      const ready = lockConn?.getUnsafe()
+      return Effect.onError(withDeadline(operation), () => {
+        if (ready !== undefined) return rebuildLockConn(ready)
+        return lockConn?.getUnsafe() === undefined ? rebuildLockConn() : Effect.void
+      })
+    })
 
   const releaseShard = sql.onDialectOrElse({
     pg: () => {
@@ -656,11 +637,10 @@ export const make = Effect.fnUntraced(function*(options: {
         return (address: string, shardId: string) =>
           sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`.pipe(execWithLockConn)
       }
-      return Effect.fnUntraced(
-        function*(_address, shardId) {
+      return (_address: string, shardId: string) =>
+        useLockConn(Effect.fnUntraced(function*(conn) {
           const lockNum = lockNumbers.get(shardId)!
           for (let i = 0; i < 5; i++) {
-            const [conn] = yield* lockConn!.await
             yield* conn.executeRaw(`SELECT pg_advisory_unlock(${pgLockNamespace}, ${lockNum})`, [])
             const takenLocks = yield* conn.executeValues(
               `SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND granted = true AND classid = ${pgLockNamespaceOid} AND objid = ${lockNum} AND objsubid = 2 AND pid = pg_backend_pid()`,
@@ -668,23 +648,18 @@ export const make = Effect.fnUntraced(function*(options: {
             )
             if (takenLocks.length === 0) return
           }
-          const [conn] = yield* lockConn!.await
           yield* conn.executeRaw(`SELECT pg_advisory_unlock_all()`, [])
-        },
-        onErrorRebuildLockConn,
-        Effect.asVoid
-      )
+        }))
     },
     mysql: () => {
       if (disableAdvisoryLocks) {
         return (address: string, shardId: string) =>
           sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`.pipe(execWithLockConn)
       }
-      return Effect.fnUntraced(
-        function*(_address, shardId) {
+      return (_address: string, shardId: string) =>
+        useLockConn(Effect.fnUntraced(function*(conn, pid) {
           const lockName = lockNames.get(shardId)!
           while (true) {
-            const [conn, pid] = yield* lockConn!.await
             yield* conn.executeRaw(`SELECT RELEASE_LOCK('${lockName}')`, [])
             const takenLocks = yield* conn.executeValues(
               `SELECT IS_USED_LOCK('${lockName}')`,
@@ -692,10 +667,7 @@ export const make = Effect.fnUntraced(function*(options: {
             )
             if (takenLocks.length === 0 || takenLocks[0][0] !== pid) return
           }
-        },
-        onErrorRebuildLockConn,
-        Effect.asVoid
-      )
+        }))
     },
     orElse: () => (address: string, shardId: string) =>
       sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`