1/**
2 * Stores cluster runner registration and shard ownership in SQL.
3 *
4 * The SQL-backed `RunnerStorage` records runners, health flags, machine ids,
5 * and shard locks so multiple processes can coordinate which runner owns each
6 * shard. This module creates the required runner and lock tables, supports an
7 * optional table prefix, uses advisory locks for PostgreSQL and MySQL when
8 * enabled, and provides constructors and layers for the storage service.
9 *
10 * @stability unstable
11 * @since 4.0.0
12 */
13import * as Arr from "../Array.ts"
14import * as Cause from "../Cause.ts"
15import * as Duration from "../Duration.ts"
16import * as Effect from "../Effect.ts"
17import * as Fiber from "../Fiber.ts"
18import * as Layer from "../Layer.ts"
19import * as Scope from "../Scope.ts"
20import * as SqlClient from "../sql/SqlClient.ts"
21import type { Connection } from "../sql/SqlConnection.ts"
22import type { SqlError } from "../sql/SqlError.ts"
23import type * as Statement from "../sql/Statement.ts"
24import { PersistenceError } from "./ClusterError.ts"
25import { ResourceRef } from "./internal/resourceRef.ts"
26import { effectiveInterval } from "./internal/shardLock.ts"
27import * as RunnerStorage from "./RunnerStorage.ts"
28import * as ShardId from "./ShardId.ts"
29import * as ShardingConfig from "./ShardingConfig.ts"
30
31const withTracerDisabled = Effect.withTracerEnabled(false)
32
33// This exact FNV-1a hash, including its tag and UTF-8 encoding, is a persistent
34// advisory-lock wire format and must never change.
35const postgresLockNamespace = (prefix: string): number => {
36 const bytes = new TextEncoder().encode(`effect-cluster:${prefix}`)
37 let hash = 0x811c9dc5
38 for (let i = 0; i < bytes.length; i++) {
39 hash = Math.imul(hash ^ bytes[i], 0x01000193)
40 }
41 return hash | 0
42}
43
44/**
45 * Creates a SQL-backed `RunnerStorage` implementation for registered runners and
46 * shard locks, using the configured table prefix and advisory locks where
47 * supported and enabled.
48 *
49 * **When to use**
50 *
51 * Use to create a SQL-backed `RunnerStorage` value directly when building
52 * custom service or layer composition around the storage implementation.
53 *
54 * **Details**
55 *
56 * When `prefix` is omitted, `make` uses the `cluster` prefix, creating
57 * `cluster_runners` and `cluster_locks`. PostgreSQL and MySQL use advisory
58 * locks unless `ShardingConfig.shardLockDisableAdvisory` is enabled; other
59 * dialects use rows in the locks table.
60 *
61 * **Gotchas**
62 *
63 * Changing `prefix` changes both generated table names, so runners using
64 * different prefixes do not share registrations or shard locks.
65 *
66 * @see {@link layer} for the default SQL-backed storage layer
67 * @see {@link layerWith} for a SQL-backed storage layer with a custom table prefix
68 *
69 * @stability unstable
70 * @category constructors
71 * @since 4.0.0
72 */
73export const make = Effect.fnUntraced(function*(options: {
74 readonly prefix?: string | undefined
75}) {
76 const config = yield* ShardingConfig.ShardingConfig
77 const shardGroups = ShardingConfig.shardGroupConfig(config)
78 const availableShardGroups = Array.from(shardGroups.available)
79 const disableAdvisoryLocks = config.shardLockDisableAdvisory
80 const lockOperationInterval = effectiveInterval(config)
81 const sql = (yield* SqlClient.SqlClient).withoutTransforms()
82 const layerScope = yield* Effect.scope
83 const prefix = options?.prefix ?? "cluster"
84 const table = (name: string) => `${prefix}_${name}`
85 const pgLockNamespace = postgresLockNamespace(prefix)
86 // PostgreSQL exposes the signed int4 lock key through pg_locks as an unsigned oid.
87 const pgLockNamespaceOid = pgLockNamespace >>> 0
88
89 // Keep all PostgreSQL and MySQL shard-lock operations on a rebuildable
90 // reserved connection, including when advisory locks are disabled.
91 const acquireLockConn = sql.onDialectOrElse({
92 pg: () =>
93 Effect.fnUntraced(function*(scope: Scope.Scope) {
94 const conn = yield* Effect.orDie(sql.reserve).pipe(
95 Scope.provide(scope)
96 )
97 const pid = disableAdvisoryLocks
98 ? 0
99 : (yield* conn.executeValues("SELECT pg_backend_pid()", []))[0][0] as number
100 if (!disableAdvisoryLocks) {
101 yield* Scope.addFinalizerExit(scope, () =>
102 conn.executeRaw("SELECT pg_advisory_unlock_all()", []).pipe(
103 Effect.timeout(lockOperationInterval),
104 Effect.interruptible,
105 Effect.ignoreCause
106 ))
107 }
108 return [conn, pid] as const
109 }, Effect.orDie),
110 mysql: () =>
111 Effect.fnUntraced(function*(scope: Scope.Scope) {
112 const conn = yield* Effect.orDie(sql.reserve).pipe(
113 Scope.provide(scope)
114 )
115 if (disableAdvisoryLocks) return [conn, 0] as const
116 // we need to get the connection id using IS_USED_LOCK to properly
117 // support vitess
118 let pid: number | undefined = undefined
119 while (pid === undefined) {
120 const address = `cluster:pid:${(Math.random() * Number.MAX_SAFE_INTEGER) | 0}`
121 const taken = (yield* conn.executeValues(
122 `SELECT GET_LOCK('${address}', 10), IS_USED_LOCK('${address}')`,
123 []
124 ))[0] as [1 | null, number]
125 if (taken[0] === null) continue
126 pid = taken[1]
127 }
128 if (!disableAdvisoryLocks) {
129 yield* Scope.addFinalizerExit(scope, () =>
130 conn.executeRaw("SELECT RELEASE_ALL_LOCKS()", []).pipe(
131 Effect.timeout(lockOperationInterval),
132 Effect.interruptible,
133 Effect.ignoreCause
134 ))
135 }
136 return [conn, pid] as const
137 }, Effect.orDie),
138 orElse: () => undefined
139 })
140 const lockConn = acquireLockConn && (yield* ResourceRef.from(layerScope, acquireLockConn))
141
142 // `Effect.timeout` waits for the timed-out effect to finish interrupting, so
143 // an operation stuck in an uninterruptible region (such as a scope finalizer
144 // releasing an unresponsive connection) can outlive its deadline. Fork the
145 // operation and timeout the join instead, leaving stalled cleanup to finish
146 // detached in the layer scope.
147 const withDeadline = Effect.fnUntraced(function*<A, E, R>(operation: Effect.Effect<A, E, R>) {
148 const fiber = yield* Effect.forkIn(operation, layerScope, { startImmediately: true })
149 return yield* Fiber.join(fiber).pipe(
150 Effect.timeout(lockOperationInterval),
151 Effect.ensuring(Effect.suspend(() =>
152 fiber.pollUnsafe() !== undefined ? Effect.void : Fiber.interrupt(fiber).pipe(
153 Effect.forkIn(layerScope, { startImmediately: true }),
154 Effect.asVoid
155 )
156 ))
157 )
158 })
159
160 const rebuildLockConn = (from?: readonly [Connection, number]) => {
161 const rebuild = lockConn?.rebuildUnsafe({ from })
162 if (!rebuild) return Effect.void
163 // Bound release so a stalled connection cannot block subsequent rebuilds.
164 return withDeadline(rebuild).pipe(
165 Effect.exit,
166 Effect.forkIn(layerScope, { startImmediately: true }),
167 Effect.asVoid
168 )
169 }
170 // SQL failures replace the held connection; deadlines are handled by withLockOperationDeadline.
171 const useLockConn = <A, E, R>(
172 f: (conn: Connection, pid: number) => Effect.Effect<A, E, R>
173 ): Effect.Effect<A, E, R> =>
174 Effect.flatMap(
175 lockConn!.await,
176 (held) =>
177 Effect.onError(
178 f(held[0], held[1]),
179 (cause) => Cause.hasInterruptsOnly(cause) ? Effect.void : rebuildLockConn(held)
180 )
181 )
182
183 const runnersTable = table("runners")
184 const runnersTableSql = sql(runnersTable)
185
186 yield* sql.onDialectOrElse({
187 mssql: () =>
188 sql`
189 IF OBJECT_ID(N'${runnersTableSql}', N'U') IS NULL
190 CREATE TABLE ${runnersTableSql} (
191 machine_id INT IDENTITY PRIMARY KEY,
192 address VARCHAR(255) NOT NULL,
193 runner TEXT NOT NULL,
194 healthy BIT NOT NULL DEFAULT 1,
195 last_heartbeat DATETIME NOT NULL DEFAULT GETDATE(),
196 UNIQUE(address)
197 )
198 `,
199 mysql: () =>
200 sql`
201 CREATE TABLE IF NOT EXISTS ${runnersTableSql} (
202 machine_id INT AUTO_INCREMENT PRIMARY KEY,
203 address VARCHAR(255) NOT NULL,
204 runner TEXT NOT NULL,
205 healthy BOOLEAN NOT NULL DEFAULT TRUE,
206 last_heartbeat DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
207 UNIQUE(address)
208 )
209 `,
210 pg: () =>
211 sql`
212 CREATE TABLE IF NOT EXISTS ${runnersTableSql} (
213 machine_id SERIAL PRIMARY KEY,
214 address VARCHAR(255) NOT NULL,
215 runner TEXT NOT NULL,
216 healthy BOOLEAN NOT NULL DEFAULT TRUE,
217 last_heartbeat TIMESTAMP NOT NULL DEFAULT NOW(),
218 UNIQUE(address)
219 )
220 `,
221 orElse: () =>
222 // sqlite
223 sql`
224 CREATE TABLE IF NOT EXISTS ${runnersTableSql} (
225 machine_id INTEGER PRIMARY KEY AUTOINCREMENT,
226 address TEXT NOT NULL,
227 runner TEXT NOT NULL,
228 healthy INTEGER NOT NULL DEFAULT 1,
229 last_heartbeat DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP),
230 UNIQUE(address)
231 )
232 `
233 })
234
235 const locksTable = table("locks")
236 const locksTableSql = sql(locksTable)
237
238 yield* sql.onDialectOrElse({
239 mssql: () =>
240 sql`
241 IF OBJECT_ID(N'${locksTableSql}', N'U') IS NULL
242 CREATE TABLE ${locksTableSql} (
243 shard_id VARCHAR(50) PRIMARY KEY,
244 address VARCHAR(255) NOT NULL,
245 acquired_at DATETIME NOT NULL
246 )
247 `,
248 mysql: () =>
249 disableAdvisoryLocks ?
250 sql`
251 CREATE TABLE IF NOT EXISTS ${locksTableSql} (
252 shard_id VARCHAR(50) PRIMARY KEY,
253 address VARCHAR(255) NOT NULL,
254 acquired_at DATETIME NOT NULL
255 )
256 ` :
257 Effect.void,
258 pg: () =>
259 disableAdvisoryLocks ?
260 sql`
261 CREATE TABLE IF NOT EXISTS ${locksTableSql} (
262 shard_id VARCHAR(50) PRIMARY KEY,
263 address VARCHAR(255) NOT NULL,
264 acquired_at TIMESTAMP NOT NULL
265 )
266 ` :
267 Effect.void,
268 orElse: () =>
269 // sqlite
270 sql`
271 CREATE TABLE IF NOT EXISTS ${locksTableSql} (
272 shard_id TEXT PRIMARY KEY,
273 address TEXT NOT NULL,
274 acquired_at DATETIME NOT NULL
275 )
276 `
277 })
278
279 const sqlNowString = sql.onDialectOrElse({
280 pg: () => "NOW()",
281 mysql: () => "NOW()",
282 mssql: () => "GETDATE()",
283 orElse: () => "CURRENT_TIMESTAMP"
284 })
285 const sqlNow = sql.literal(sqlNowString)
286
287 // Use the same shard ID order for all multi-row PostgreSQL lease locks to
288 // avoid deadlocks between acquisition, refresh and bulk release.
289 const pgLockOrder = sql.literal(`ORDER BY shard_id COLLATE "C"`)
290
291 const expiresSeconds = sql.literal(
292 Math.ceil(Duration.toSeconds(
293 Duration.fromInputUnsafe(config.shardLockExpiration)
294 )).toString()
295 )
296 const lockExpiresAt = sql.onDialectOrElse({
297 pg: () => sql`${sqlNow} - INTERVAL '${expiresSeconds} seconds'`,
298 mysql: () => sql`DATE_SUB(${sqlNow}, INTERVAL ${expiresSeconds} SECOND)`,
299 mssql: () => sql`DATEADD(SECOND, -${expiresSeconds}, ${sqlNow})`,
300 orElse: () => sql`datetime(${sqlNow}, '-${expiresSeconds} seconds')`
301 })
302
303 const encodeBoolean = sql.onDialectOrElse({
304 mssql: () => (b: boolean) => (b ? 1 : 0),
305 sqlite: () => (b: boolean) => (b ? 1 : 0),
306 orElse: () => (b: boolean) => b
307 })
308
309 // Upsert runner and return machine_id
310 const insertRunner = sql.onDialectOrElse({
311 mssql: () => (address: string, runner: string, healthy: boolean) =>
312 sql`
313 MERGE ${runnersTableSql} AS target
314 USING (SELECT ${address} AS address, ${runner} AS runner, ${sqlNow} AS last_heartbeat, ${
315 encodeBoolean(healthy)
316 } AS healthy) AS source
317 ON target.address = source.address
318 WHEN MATCHED THEN
319 UPDATE SET runner = source.runner, last_heartbeat = source.last_heartbeat, healthy = source.healthy
320 WHEN NOT MATCHED THEN
321 INSERT (address, runner, last_heartbeat, healthy)
322 VALUES (source.address, source.runner, source.last_heartbeat, source.healthy)
323 OUTPUT INSERTED.machine_id;
324 `.values,
325 mysql: () => (address: string, runner: string, healthy: boolean) =>
326 sql<{ machine_id: number }>`
327 INSERT INTO ${runnersTableSql} (address, runner, last_heartbeat, healthy)
328 VALUES (${address}, ${runner}, ${sqlNow}, ${healthy})
329 ON DUPLICATE KEY UPDATE
330 runner = VALUES(runner),
331 last_heartbeat = VALUES(last_heartbeat),
332 healthy = VALUES(healthy);
333 SELECT machine_id FROM ${runnersTableSql} WHERE address = ${address};
334 `.unprepared.pipe(
335 Effect.map((results: any) => [[results[1][0].machine_id]])
336 ),
337 pg: () => (address: string, runner: string, healthy: boolean) =>
338 sql`
339 INSERT INTO ${runnersTableSql} (address, runner, last_heartbeat, healthy)
340 VALUES (${address}, ${runner}, ${sqlNow}, ${healthy})
341 ON CONFLICT (address) DO UPDATE
342 SET runner = EXCLUDED.runner,
343 last_heartbeat = EXCLUDED.last_heartbeat,
344 healthy = EXCLUDED.healthy
345 RETURNING machine_id
346 `.values,
347 orElse: () => (address: string, runner: string, healthy: boolean) =>
348 // sqlite
349 sql`
350 INSERT INTO ${runnersTableSql} (address, runner, last_heartbeat, healthy)
351 VALUES (${address}, ${runner}, ${sqlNow}, ${encodeBoolean(healthy)})
352 ON CONFLICT(address) DO UPDATE SET
353 runner = excluded.runner,
354 last_heartbeat = excluded.last_heartbeat,
355 healthy = excluded.healthy
356 RETURNING machine_id;
357 `.values
358 })
359
360 const execWithLockConn = <A>(effect: Statement.Statement<A>): Effect.Effect<unknown, SqlError> => {
361 if (!lockConn) return effect
362 const [query, params] = effect.compile()
363 return useLockConn((conn) => conn.executeRaw(query, params))
364 }
365 const execWithLockConnUnprepared = <A>(
366 effect: Statement.Statement<A>
367 ): Effect.Effect<ReadonlyArray<ReadonlyArray<any>>, SqlError> => {
368 if (!lockConn) return effect.values
369 const [query, params] = effect.compile()
370 return useLockConn((conn) => conn.executeUnprepared(query, params, undefined))
371 }
372 const execWithLockConnValues = <A>(
373 effect: Statement.Statement<A>
374 ): Effect.Effect<ReadonlyArray<ReadonlyArray<any>>, SqlError> => {
375 if (!lockConn) return effect.values
376 const [query, params] = effect.compile()
377 return useLockConn((conn) => conn.executeValues(query, params))
378 }
379
380 const acquireLock = sql.onDialectOrElse({
381 pg: () => {
382 if (disableAdvisoryLocks) {
383 return (address: string, shardIds: ReadonlyArray<string>) => {
384 const values = shardIds.map((shardId) =>
385 sql`(${stringLiteral(shardId)}, ${stringLiteral(address)}, ${sqlNow})`
386 )
387 return sql`
388 INSERT INTO ${locksTableSql} (shard_id, address, acquired_at)
389 SELECT * FROM (VALUES ${sql.csv(values)}) AS requested (shard_id, address, acquired_at)
390 ${pgLockOrder}
391 ON CONFLICT (shard_id) DO UPDATE
392 SET address = ${address}, acquired_at = ${sqlNow}
393 WHERE ${locksTableSql}.address = ${address}
394 OR ${locksTableSql}.acquired_at < ${lockExpiresAt}
395`.pipe(
396 execWithLockConn,
397 Effect.andThen(acquiredLocks(address, shardIds))
398 )
399 }
400 }
401 return (_address: string, shardIds: ReadonlyArray<string>) =>
402 useLockConn(Effect.fnUntraced(function*(conn, pid) {
403 const acquiredShardIds: Array<string> = []
404 const toAcquire = new Map(shardIds.map((shardId) => [lockNumbers.get(shardId)!, shardId]))
405 const takenLocks = yield* conn.executeValues(
406 `SELECT objid FROM pg_locks WHERE locktype = 'advisory' AND granted = true AND classid = ${pgLockNamespaceOid} AND objsubid = 2 AND pid = ${pid} ORDER BY objid`,
407 []
408 )
409 for (let i = 0; i < takenLocks.length; i++) {
410 const lockNum = takenLocks[i][0] as number
411 const shardId = toAcquire.get(lockNum)
412 if (shardId === undefined) continue
413 acquiredShardIds.push(shardId)
414 toAcquire.delete(lockNum)
415 }
416 if (toAcquire.size === 0) {
417 return acquiredShardIds
418 }
419 const rows = yield* conn.executeUnprepared(`SELECT ${pgLocks(toAcquire)}`, [], undefined)
420 const results = rows[0] as Record<string, boolean>
421 for (const shardId in results) {
422 if (results[shardId]) {
423 acquiredShardIds.push(shardId)
424 }
425 }
426 return acquiredShardIds
427 }))
428 },
429
430 mysql: () => {
431 if (disableAdvisoryLocks) {
432 return (address: string, shardIds: ReadonlyArray<string>) => {
433 const values = shardIds.map((shardId) =>
434 sql`(${stringLiteral(shardId)}, ${stringLiteral(address)}, ${sqlNow})`
435 )
436 return sql`
437 INSERT INTO ${locksTableSql} (shard_id, address, acquired_at) VALUES ${sql.csv(values)}
438 ON DUPLICATE KEY UPDATE
439 address = IF(address = VALUES(address) OR acquired_at < ${lockExpiresAt}, VALUES(address), address),
440 acquired_at = IF(address = VALUES(address) OR acquired_at < ${lockExpiresAt}, VALUES(acquired_at), acquired_at)
441`.pipe(
442 execWithLockConnUnprepared,
443 Effect.andThen(acquiredLocks(address, shardIds))
444 )
445 }
446 }
447 return (_address: string, shardIds: ReadonlyArray<string>) =>
448 useLockConn(Effect.fnUntraced(function*(conn, pid) {
449 const takenLocks = (yield* conn.executeValues(`SELECT ${allMySqlTakenLocks}`, []))[0] as Array<number | null>
450 const acquiredShardIds: Array<string> = []
451 const toAcquire: Array<string> = []
452 for (let i = 0; i < shardIds.length; i++) {
453 const shardId = shardIds[i]
454 const lockTakenBy = takenLocks[shardIdsIndex.get(shardId)!]
455 if (lockTakenBy === pid) {
456 acquiredShardIds.push(shardId)
457 } else if (shardIds.includes(shardId)) {
458 toAcquire.push(shardId)
459 }
460 }
461 if (toAcquire.length === 0) {
462 return acquiredShardIds
463 }
464 const results = (yield* conn.executeValues(`SELECT ${mysqlLocks(toAcquire)}`, []))[0] as Array<number>
465 for (let i = 0; i < results.length; i++) {
466 if (results[i] === 1) {
467 acquiredShardIds.push(toAcquire[i])
468 }
469 }
470 return acquiredShardIds
471 }))
472 },
473
474 mssql: () => (address: string, shardIds: ReadonlyArray<string>) => {
475 const values = shardIds.map((shardId) => sql`(${stringLiteral(shardId)}, ${stringLiteral(address)}, ${sqlNow})`)
476 return sql`
477 MERGE ${locksTableSql} WITH (HOLDLOCK) AS target
478 USING (SELECT * FROM (VALUES ${sql.csv(values)})) AS source (shard_id, address, acquired_at)
479 ON target.shard_id = source.shard_id
480 WHEN MATCHED AND (target.address = source.address OR DATEDIFF(SECOND, target.acquired_at, ${sqlNow}) > ${expiresSeconds}) THEN
481 UPDATE SET address = source.address, acquired_at = source.acquired_at
482 WHEN NOT MATCHED THEN
483 INSERT (shard_id, address, acquired_at)
484 VALUES (source.shard_id, source.address, source.acquired_at);
485 `.pipe(
486 Effect.andThen(acquiredLocks(address, shardIds)),
487 sql.withTransaction
488 )
489 },
490
491 orElse: () => (address: string, shardIds: ReadonlyArray<string>) => {
492 const values = shardIds.map((shardId) => sql`(${stringLiteral(shardId)}, ${stringLiteral(address)}, ${sqlNow})`)
493 return sql`
494 WITH source(shard_id, address, acquired_at) AS (VALUES ${sql.csv(values)})
495 INSERT INTO ${locksTableSql} (shard_id, address, acquired_at)
496 SELECT source.shard_id, source.address, source.acquired_at
497 FROM source
498 WHERE NOT EXISTS (
499 SELECT 1 FROM ${locksTableSql}
500 WHERE shard_id = source.shard_id
501 AND address != ${address}
502 AND (strftime('%s', ${sqlNow}) - strftime('%s', acquired_at)) <= ${expiresSeconds}
503 )
504 ON CONFLICT(shard_id) DO UPDATE
505 SET address = ${address}, acquired_at = ${sqlNow}
506 `.pipe(
507 Effect.andThen(acquiredLocks(address, shardIds)),
508 sql.withTransaction
509 )
510 }
511 })
512
513 const lockNumbers = new Map<string, number>()
514 for (let i = 0; i < availableShardGroups.length; i++) {
515 const group = availableShardGroups[i]
516 const base = (i + 1) * 1000000
517 for (let shard = 1; shard <= config.shardsPerGroup; shard++) {
518 const shardId = ShardId.make(group, shard).toString()
519 const lockNum = base + shard
520 lockNumbers.set(shardId, lockNum)
521 }
522 }
523
524 const shardIdsIndex = new Map<string, number>()
525 const lockNames = new Map<string, string>()
526 const lockNamesReverse = new Map<string, string>()
527 {
528 let index = 0
529 for (let i = 0; i < availableShardGroups.length; i++) {
530 const group = availableShardGroups[i]
531 for (let shard = 1; shard <= config.shardsPerGroup; shard++) {
532 const shardId = ShardId.make(group, shard).toString()
533 const lockName = `${prefix}.${shardId}`
534 shardIdsIndex.set(shardId, index++)
535 lockNames.set(shardId, lockName)
536 lockNamesReverse.set(lockName, shardId)
537 }
538 }
539 }
540
541 const pgLocks = (shardIdsMap: Map<number, string>) =>
542 Array.from(
543 shardIdsMap.entries(),
544 ([lockNum, shardId]) => `pg_try_advisory_lock(${pgLockNamespace}, ${lockNum}) AS "${shardId}"`
545 ).join(", ")
546
547 const mysqlLocks = (shardIds: ReadonlyArray<string>) =>
548 shardIds.map((shardId) => `GET_LOCK('${lockNames.get(shardId)!}', 0) AS "${shardId}"`).join(", ")
549
550 const allMySqlTakenLocks = Array.from(
551 lockNames.entries(),
552 ([shardId, lockName]) => `IS_USED_LOCK('${lockName}') AS "${shardId}"`
553 ).join(", ")
554
555 const acquiredLocks = (address: string, shardIds: ReadonlyArray<string>) =>
556 sql<{ shard_id: string }>`
557 SELECT shard_id FROM ${sql(locksTable)}
558 WHERE address = ${address}
559 AND acquired_at >= ${lockExpiresAt}
560 AND shard_id IN ${stringLiteralArr(shardIds)}
561 `.pipe(
562 execWithLockConnValues,
563 Effect.map((rows) => rows.map((row) => row[0] as string))
564 )
565
566 const wrapString = sql.onDialectOrElse({
567 mssql: () => (s: string) => `N'${s}'`,
568 orElse: () => (s: string) => `'${s}'`
569 })
570 const stringLiteral = (s: string) => sql.literal(wrapString(s))
571 const stringLiteralArr = (arr: ReadonlyArray<string>) => sql.literal(`(${arr.map(wrapString).join(",")})`)
572
573 const refreshShards = sql.onDialectOrElse({
574 pg: () => {
575 if (!disableAdvisoryLocks) return acquireLock
576 return (address: string, shardIds: ReadonlyArray<string>) =>
577 sql`
578 WITH locked AS (
579 SELECT shard_id FROM ${locksTableSql}
580 WHERE address = ${address} AND shard_id IN ${stringLiteralArr(shardIds)}
581 ${pgLockOrder} FOR UPDATE
582 )
583 UPDATE ${locksTableSql}
584 SET acquired_at = ${sqlNow}
585 WHERE address = ${address} AND shard_id IN (SELECT shard_id FROM locked)
586 RETURNING shard_id
587 `.pipe(
588 execWithLockConnValues,
589 Effect.map((rows) => rows.map((row) => row[0] as string))
590 )
591 },
592 mysql: () => {
593 if (!disableAdvisoryLocks) return acquireLock
594 return (address: string, shardIds: ReadonlyArray<string>) => {
595 const shardIdsStr = stringLiteralArr(shardIds)
596 return sql<Array<{ shard_id: string }>>`
597 UPDATE ${locksTableSql}
598 SET acquired_at = ${sqlNow}
599 WHERE address = ${address} AND shard_id IN ${shardIdsStr};
600 SELECT shard_id FROM ${locksTableSql} WHERE address = ${address} AND shard_id IN ${shardIdsStr}
601 `.pipe(
602 execWithLockConnUnprepared,
603 Effect.map((rows) => rows[1].map((row) => row.shard_id))
604 )
605 }
606 },
607 mssql: () => (address: string, shardIds: ReadonlyArray<string>) =>
608 sql`
609 UPDATE ${locksTableSql}
610 SET acquired_at = ${sqlNow}
611 OUTPUT inserted.shard_id
612 WHERE address = ${address} AND shard_id IN ${stringLiteralArr(shardIds)}
613 `.pipe(execWithLockConnValues, Effect.map((rows) => rows.map((row) => row[0] as string))),
614 orElse: () => (address: string, shardIds: ReadonlyArray<string>) =>
615 sql`
616 UPDATE ${locksTableSql}
617 SET acquired_at = ${sqlNow}
618 WHERE address = ${address} AND shard_id IN ${stringLiteralArr(shardIds)}
619 RETURNING shard_id
620 `.pipe(execWithLockConnValues, Effect.map((rows) => rows.map((row) => row[0] as string)))
621 })
622
623 // On failure, replace the connection ready at entry. If none was ready,
624 // retry only a failed rebuild, leaving pending or newly ready connections alone.
625 const withLockOperationDeadline = <A, E, R>(operation: Effect.Effect<A, E, R>) =>
626 Effect.suspend(() => {
627 const ready = lockConn?.getUnsafe()
628 return Effect.onError(withDeadline(operation), () => {
629 if (ready !== undefined) return rebuildLockConn(ready)
630 return lockConn?.getUnsafe() === undefined ? rebuildLockConn() : Effect.void
631 })
632 })
633
634 const releaseShard = sql.onDialectOrElse({
635 pg: () => {
636 if (disableAdvisoryLocks) {
637 return (address: string, shardId: string) =>
638 sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`.pipe(execWithLockConn)
639 }
640 return (_address: string, shardId: string) =>
641 useLockConn(Effect.fnUntraced(function*(conn) {
642 const lockNum = lockNumbers.get(shardId)!
643 for (let i = 0; i < 5; i++) {
644 yield* conn.executeRaw(`SELECT pg_advisory_unlock(${pgLockNamespace}, ${lockNum})`, [])
645 const takenLocks = yield* conn.executeValues(
646 `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()`,
647 []
648 )
649 if (takenLocks.length === 0) return
650 }
651 yield* conn.executeRaw(`SELECT pg_advisory_unlock_all()`, [])
652 }))
653 },
654 mysql: () => {
655 if (disableAdvisoryLocks) {
656 return (address: string, shardId: string) =>
657 sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`.pipe(execWithLockConn)
658 }
659 return (_address: string, shardId: string) =>
660 useLockConn(Effect.fnUntraced(function*(conn, pid) {
661 const lockName = lockNames.get(shardId)!
662 while (true) {
663 yield* conn.executeRaw(`SELECT RELEASE_LOCK('${lockName}')`, [])
664 const takenLocks = yield* conn.executeValues(
665 `SELECT IS_USED_LOCK('${lockName}')`,
666 []
667 )
668 if (takenLocks.length === 0 || takenLocks[0][0] !== pid) return
669 }
670 }))
671 },
672 orElse: () => (address: string, shardId: string) =>
673 sql`DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id = ${shardId}`
674 })
675
676 const releaseAllShards = sql.onDialectOrElse({
677 pg: () => (address: string) =>
678 disableAdvisoryLocks
679 ? sql`
680 WITH locked AS (
681 SELECT shard_id FROM ${locksTableSql} WHERE address = ${address} ${pgLockOrder} FOR UPDATE
682 )
683 DELETE FROM ${locksTableSql} WHERE address = ${address} AND shard_id IN (SELECT shard_id FROM locked)
684 `.pipe(execWithLockConn)
685 : sql`SELECT pg_advisory_unlock_all()`.pipe(execWithLockConn, Effect.asVoid),
686 mysql: () => (address: string) =>
687 disableAdvisoryLocks
688 ? sql`DELETE FROM ${locksTableSql} WHERE address = ${address}`.pipe(execWithLockConn)
689 : sql`SELECT RELEASE_ALL_LOCKS()`.pipe(execWithLockConn, Effect.asVoid),
690 orElse: () => (address: string) => sql`DELETE FROM ${locksTableSql} WHERE address = ${address}`
691 })
692
693 return RunnerStorage.makeEncoded({
694 getRunners: sql`SELECT runner, healthy FROM ${runnersTableSql} WHERE last_heartbeat > ${lockExpiresAt}`.values.pipe(
695 PersistenceError.refail,
696 Effect.map(Arr.map(([runner, healthy]) => [String(runner), Boolean(healthy)] as const)),
697 withTracerDisabled
698 ),
699
700 register: (address, runner, healthy) =>
701 insertRunner(address, runner, healthy).pipe(
702 Effect.map((rows: any) => Number(rows[0][0])),
703 PersistenceError.refail,
704 withTracerDisabled
705 ),
706
707 unregister: (address) =>
708 sql`DELETE FROM ${runnersTableSql} WHERE address = ${address} OR last_heartbeat < ${lockExpiresAt}`.pipe(
709 Effect.asVoid,
710 PersistenceError.refail,
711 withTracerDisabled
712 ),
713
714 setRunnerHealth: (address, healthy) =>
715 sql`UPDATE ${runnersTableSql} SET healthy = ${encodeBoolean(healthy)} WHERE address = ${address}`
716 .pipe(
717 Effect.asVoid,
718 PersistenceError.refail,
719 withTracerDisabled
720 ),
721
722 acquire: (address, shardIds) =>
723 withLockOperationDeadline(acquireLock(address, shardIds)).pipe(
724 PersistenceError.refail,
725 withTracerDisabled
726 ),
727
728 refresh: (address, shardIds) => {
729 const heartbeat = sql`UPDATE ${runnersTableSql} SET last_heartbeat = ${sqlNow} WHERE address = ${address}`
730 // An empty refresh is the liveness probe used while lock storage is
731 // unhealthy. Run it on the shared pool, so a wedged reserved lock
732 // connection cannot block recovery.
733 if (shardIds.length === 0) {
734 return withDeadline(heartbeat).pipe(
735 Effect.as([]),
736 PersistenceError.refail,
737 withTracerDisabled
738 )
739 }
740 return withLockOperationDeadline(
741 heartbeat.pipe(
742 execWithLockConn,
743 Effect.andThen(refreshShards(address, shardIds))
744 )
745 ).pipe(
746 PersistenceError.refail,
747 withTracerDisabled
748 )
749 },
750
751 release: (address, shardId) =>
752 withLockOperationDeadline(releaseShard(address, shardId)).pipe(
753 Effect.asVoid,
754 PersistenceError.refail,
755 withTracerDisabled
756 ),
757
758 releaseAll: (address) =>
759 withLockOperationDeadline(releaseAllShards(address)).pipe(
760 Effect.asVoid,
761 PersistenceError.refail,
762 withTracerDisabled
763 )
764 })
765}, withTracerDisabled)
766
767/**
768 * Layer that provides SQL-backed `RunnerStorage` using the default table prefix.
769 *
770 * @stability unstable
771 * @category layers
772 * @since 4.0.0
773 */
774export const layer: Layer.Layer<
775 RunnerStorage.RunnerStorage,
776 SqlError,
777 SqlClient.SqlClient | ShardingConfig.ShardingConfig
778> = Layer.effect(RunnerStorage.RunnerStorage)(make({}))
779
780/**
781 * Layer that provides SQL-backed `RunnerStorage` using a custom table prefix.
782 *
783 * @stability unstable
784 * @category layers
785 * @since 4.0.0
786 */
787export const layerWith = (options: {
788 readonly prefix?: string | undefined
789}): Layer.Layer<RunnerStorage.RunnerStorage, SqlError, SqlClient.SqlClient | ShardingConfig.ShardingConfig> =>
790 Layer.effect(RunnerStorage.RunnerStorage)(make(options))