effect 4.0.0 → 4.0.2

package/src/cluster/Sharding.ts

@@ -549,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)
@@ -1462,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
               }
@@ -1592,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>
@@ -1855,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(() => {
@@ -716,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) =>
@@ -731,8 +750,7 @@ export const makeEncoded: (options?: {
                 ])
             )
           )
-        })
-        `.pipe(
+        })`).pipe(
           Effect.asVoid,
           PersistenceError.refail,
           withTracerDisabled
@@ -761,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}`