1import { NodeClusterSocket, NodeCrypto, NodeSocket, NodeSocketServer } from "@effect/platform-node"
2import { MysqlClient } from "@effect/sql-mysql2"
3import { PgClient } from "@effect/sql-pg"
4import { Clock, Context, Duration, Effect, Exit, Latch, Layer, Option, Redacted, Scope } from "effect"
5import {
6 ClusterWorkflowEngine,
7 type Entity,
8 EntityId,
9 type MessageStorage,
10 type Runner as RunnerModel,
11 RunnerAddress,
12 RunnerHealth,
13 Runners,
14 RunnerStorage,
15 ShardId,
16 Sharding,
17 ShardingConfig,
18 SocketRunner,
19 SqlMessageStorage,
20 SqlRunnerStorage
21} from "effect/cluster"
22import type { Rpc } from "effect/rpc"
23import { RpcClient, RpcSerialization } from "effect/rpc"
24import * as Socket from "effect/socket/Socket"
25import * as SocketServer from "effect/socket/SocketServer"
26import { SqlClient, type SqlConnection, SqlError } from "effect/sql"
27import { WorkflowEngine } from "effect/workflow"
28import * as Net from "node:net"
29import { inject } from "vitest"
30
31export type Backend = "mysql" | "pg"
32export type LockMode = "advisory" | "row"
33export type LockFaultMode = "blackhole" | "stuck" | "fail" | "hangRelease" | "clear"
34export type SocketDirection = "inbound" | "outbound"
35
36export interface ClusterRunner {
37 readonly address: RunnerAddress.RunnerAddress
38 readonly index: number
39 readonly shardGroups: ReadonlyArray<string>
40 readonly sharding: Sharding.Sharding["Service"]
41 readonly state: () => "frozen" | "killed" | "running" | "stopped"
42}
43
44/** A raw TCP caller whose transport can close without a Sharding retry. */
45export const makeRawRunnerClient = Effect.fnUntraced(function*(runner: ClusterRunner) {
46 const serialization = yield* RpcSerialization.RpcSerialization.pipe(
47 Effect.provide(RpcSerialization.layerSchemaBinary())
48 )
49 const transport = yield* Effect.acquireRelease(Scope.make(), (scope) => Scope.close(scope, Exit.void))
50 const socket = yield* NodeSocket.makeNet({ host: runner.address.host, port: runner.address.port })
51 const protocol = yield* RpcClient.makeProtocolSocket().pipe(
52 Effect.provideService(Socket.Socket, socket),
53 Effect.provideService(RpcSerialization.RpcSerialization, serialization),
54 Scope.provide(transport)
55 )
56 const client = yield* RpcClient.make(Runners.Rpcs).pipe(
57 Effect.provideService(RpcClient.Protocol, protocol),
58 Scope.provide(transport)
59 )
60 return { client, codecFor: serialization.codecFor, disconnect: Scope.close(transport, Exit.void) }
61})
62
63export interface MessageCounts {
64 readonly failed: number
65 readonly replied: number
66 readonly unprocessed: number
67}
68
69export interface MakeOptions {
70 readonly backend: Backend
71 readonly config?: Partial<ShardingConfig.ShardingConfig["Service"]> | undefined
72 readonly entities:
73 | RunnerEntities
74 | ((options: { readonly prefix: string }) => RunnerEntities)
75 readonly lockMode?: LockMode | undefined
76 readonly runnerLayer?: RunnerLayer | undefined
77}
78
79export interface StartOptions {
80 readonly assignedShardGroups?: ReadonlyArray<string> | undefined
81 readonly entities?: RunnerEntities | undefined
82 readonly runnerShardWeight?: number | undefined
83}
84
85export interface CutSocketOptions {
86 readonly direction: SocketDirection
87 readonly peer: ClusterRunner | "client"
88}
89
90export interface InsertMessageRow {
91 readonly deliver_at: number | bigint | null
92 readonly entity_id: string
93 readonly entity_type: string
94 readonly headers: string | null
95 readonly id: string | bigint
96 readonly kind: 0 | 1 | 2
97 readonly message_id: string | null
98 readonly payload: string | null
99 readonly reply_id: string | bigint | null
100 readonly request_id: string | bigint
101 readonly sampled: boolean | number | null
102 readonly shard_id: string
103 readonly span_id: string | null
104 readonly tag: string | null
105 readonly trace_id: string | null
106}
107
108type HarnessConfig = Partial<ShardingConfig.ShardingConfig["Service"]> & {
109 readonly shardLockDisableAdvisory: boolean
110}
111
112export type RunnerEntities = Layer.Layer<
113 never,
114 never,
115 Sharding.Sharding | MessageStorage.MessageStorage | SqlClient.SqlClient
116>
117
118interface RegistrationRow {
119 readonly address: string
120 readonly healthy: boolean | number
121 readonly last_heartbeat: Date | string
122 readonly runner: unknown
123}
124
125interface MessageRow {
126 readonly processed: boolean | number
127 readonly reply_payload: string | Record<string, unknown> | null
128}
129
130interface RunnerEntry extends ClusterRunner {
131 readonly controller: ReturnType<typeof makeRunnerStorageController>
132 readonly lockFault: ReturnType<typeof makeLockFaultController>
133 readonly scope: Scope.Closeable
134 setState(state: ReturnType<ClusterRunner["state"]>): void
135}
136
137const clusterConfig = {
138 entityMaxIdleTime: 3_000,
139 entityMessagePollInterval: 500,
140 refreshAssignmentsInterval: 150,
141 shardLockExpiration: 1_750,
142 shardLockRefreshInterval: 500
143} as const
144
145let nextCluster = 0
146
147interface LockSession {
148 fault: "blackhole" | "fail" | undefined
149}
150
151const makeLockFaultController = (sql: SqlClient.SqlClient) => {
152 const releaseGate = Latch.makeUnsafe(true)
153 let persistentFault: "stuck" | "hangRelease" | undefined
154 let currentSession: LockSession | undefined
155
156 const wrapConnection = (
157 connection: SqlConnection.Connection,
158 session: LockSession
159 ): SqlConnection.Connection => {
160 const execute = <A, E extends SqlError.SqlError, R>(sql: string, effect: Effect.Effect<A, E, R>) =>
161 Effect.suspend((): Effect.Effect<A, E | SqlError.SqlError, R> => {
162 // Let scope cleanup release advisory locks even when the simulated
163 // operational connection is stuck. `hangRelease` blocks teardown with
164 // the separate gate below, after the database lock has been dropped.
165 if (sql.includes("pg_advisory_unlock_all") || sql.includes("RELEASE_ALL_LOCKS")) return effect
166 const mode = session.fault ?? persistentFault
167 if (mode === "fail") {
168 return Effect.fail(
169 new SqlError.SqlError({
170 reason: new SqlError.ConnectionError({ cause: new Error("lock connection lost") })
171 })
172 )
173 }
174 if (mode === "blackhole") {
175 return Effect.never
176 }
177 if (mode === "stuck" || mode === "hangRelease") {
178 return Effect.never
179 }
180 return effect
181 })
182 return {
183 ...connection,
184 execute: (...args) => execute(args[0], connection.execute(...args)),
185 executeRaw: (...args) => execute(args[0], connection.executeRaw(...args)),
186 executeValues: (...args) => execute(args[0], connection.executeValues(...args)),
187 executeValuesUnprepared: (...args) => execute(args[0], connection.executeValuesUnprepared(...args)),
188 executeUnprepared: (...args) => execute(args[0], connection.executeUnprepared(...args))
189 }
190 }
191
192 const transformFreeSql = sql.withoutTransforms()
193 const client: SqlClient.SqlClient = new Proxy(transformFreeSql, {
194 get(target, property, receiver) {
195 if (property === "reserve") {
196 return Effect.gen(function*() {
197 const session: LockSession = { fault: undefined }
198 currentSession = session
199 yield* Effect.addFinalizer(() =>
200 Effect.uninterruptible(
201 Effect.suspend(() => persistentFault === "hangRelease" ? releaseGate.await : Effect.void)
202 )
203 )
204 return wrapConnection(yield* target.reserve, session)
205 })
206 }
207 if (property === "withoutTransforms") {
208 return () => client
209 }
210 return Reflect.get(target, property, receiver)
211 }
212 })
213
214 const set = (mode: LockFaultMode) =>
215 Effect.sync(() => {
216 if (mode === "clear") {
217 persistentFault = undefined
218 if (currentSession !== undefined) currentSession.fault = undefined
219 releaseGate.openUnsafe()
220 } else if (mode === "blackhole" || mode === "fail") {
221 if (currentSession === undefined) throw new Error("Runner has no reserved lock connection")
222 persistentFault = undefined
223 releaseGate.openUnsafe()
224 currentSession.fault = mode
225 } else {
226 persistentFault = mode
227 if (currentSession !== undefined) currentSession.fault = undefined
228 if (mode === "hangRelease") releaseGate.closeUnsafe()
229 else releaseGate.openUnsafe()
230 }
231 })
232
233 return { client, set }
234}
235
236const makeRunnerStorageController = (storage: RunnerStorage.RunnerStorage["Service"]) => {
237 const gate = Latch.makeUnsafe(true)
238 const refreshPaused = Latch.makeUnsafe()
239 const syncPaused = Latch.makeUnsafe()
240 let mode: "frozen" | "killed" | "running" = "running"
241 let lastRunners: Array<readonly [RunnerModel.Runner, boolean]> = []
242
243 const waitWhileFrozen = <A, E>(
244 paused: Latch.Latch,
245 effect: Effect.Effect<A, E>,
246 onKilled: () => A
247 ): Effect.Effect<A, E> =>
248 Effect.suspend(() => {
249 if (mode === "running") return effect
250 if (mode === "killed") return Effect.succeed(onKilled())
251 paused.openUnsafe()
252 return gate.await.pipe(
253 Effect.uninterruptible,
254 Effect.andThen(Effect.suspend(() => mode === "running" ? effect : Effect.succeed(onKilled())))
255 )
256 })
257
258 const controlled = RunnerStorage.RunnerStorage.of({
259 ...storage,
260 getRunners: waitWhileFrozen(
261 syncPaused,
262 storage.getRunners.pipe(Effect.tap((runners) => Effect.sync(() => lastRunners = runners))),
263 () => lastRunners
264 ),
265 refresh: (address, shardIds) => {
266 const shards = Array.from(shardIds)
267 return waitWhileFrozen(refreshPaused, storage.refresh(address, shards), () => shards)
268 },
269 release: (address, shardId) => mode === "killed" ? Effect.void : storage.release(address, shardId),
270 releaseAll: (address) => mode === "killed" ? Effect.void : storage.releaseAll(address),
271 unregister: (address) => mode === "killed" ? Effect.void : storage.unregister(address)
272 })
273
274 return {
275 controlled,
276 freeze: Effect.sync(() => {
277 mode = "frozen"
278 gate.closeUnsafe()
279 }).pipe(Effect.andThen(Effect.all([refreshPaused.await, syncPaused.await], { discard: true }))),
280 kill: Effect.sync(() => {
281 mode = "killed"
282 gate.openUnsafe()
283 }),
284 resume: Effect.sync(() => {
285 mode = "running"
286 gate.openUnsafe()
287 })
288 }
289}
290
291const clientPeer = "client"
292const addressKey = (address: RunnerAddress.RunnerAddress) => `${address.host}:${address.port}`
293
294const makeSocketController = () => {
295 const connections = new Map<string, Set<Net.Socket>>()
296 const linkKey = (source: string, target: string) => `${source}->${target}`
297
298 const add = (source: string, target: RunnerAddress.RunnerAddress, socket: Net.Socket) => {
299 const key = linkKey(source, addressKey(target))
300 let active = connections.get(key)
301 if (active === undefined) {
302 active = new Set()
303 connections.set(key, active)
304 }
305 active.add(socket)
306 }
307
308 const remove = (source: string, target: RunnerAddress.RunnerAddress, socket: Net.Socket) => {
309 const key = linkKey(source, addressKey(target))
310 const active = connections.get(key)
311 if (active === undefined) return
312 active.delete(socket)
313 if (active.size === 0) connections.delete(key)
314 }
315
316 const cut = (source: string, target: string) =>
317 Effect.sync(() => {
318 const active = connections.get(linkKey(source, target))
319 const socket = active?.values().next().value as Net.Socket | undefined
320 if (socket === undefined) {
321 throw new Error(`No active socket from ${source} to ${target}`)
322 }
323 socket.destroy(new Error(`Socket cut from ${source} to ${target}`))
324 })
325
326 return { add, cut, remove }
327}
328
329type SocketController = ReturnType<typeof makeSocketController>
330
331const trackedClientProtocolLayer = (
332 controller: SocketController,
333 source: string
334): Layer.Layer<Runners.RpcClientProtocol, never, RpcSerialization.RpcSerialization> =>
335 Layer.effect(Runners.RpcClientProtocol)(
336 Effect.gen(function*() {
337 const serialization = yield* RpcSerialization.RpcSerialization
338 return {
339 codecFor: serialization.codecFor,
340 make: Effect.fnUntraced(function*(address) {
341 const socket = yield* NodeSocket.fromDuplex(
342 Effect.acquireRelease(
343 Effect.callback<Net.Socket, Socket.SocketError>((resume) => {
344 const connection = Net.createConnection({
345 host: address.host,
346 port: address.port
347 })
348 connection.once("connect", () => {
349 controller.add(source, address, connection)
350 resume(Effect.succeed(connection))
351 })
352 connection.on("error", (cause) => {
353 resume(Effect.fail(
354 new Socket.SocketError({
355 reason: new Socket.SocketOpenError({ kind: "Unknown", cause })
356 })
357 ))
358 })
359 return Effect.sync(() => connection.destroy())
360 }),
361 (connection) =>
362 Effect.sync(() => {
363 controller.remove(source, address, connection)
364 if (!connection.closed) connection.destroySoon()
365 })
366 ),
367 { openTimeout: 1_000 }
368 )
369 return yield* RpcClient.makeProtocolSocket().pipe(
370 Effect.provideService(Socket.Socket, socket),
371 Effect.provideService(RpcSerialization.RpcSerialization, serialization)
372 )
373 }, Effect.orDie)
374 }
375 })
376 )
377
378const runnerHealthLayer = (
379 clientProtocol: Layer.Layer<Runners.RpcClientProtocol, never, RpcSerialization.RpcSerialization>
380) =>
381 RunnerHealth.layerPing.pipe(
382 Layer.provide(Runners.layerRpc),
383 Layer.provide(clientProtocol)
384 )
385
386export const socketRunnerLayer = (
387 address: RunnerAddress.RunnerAddress,
388 entities: RunnerEntities,
389 socketServer: SocketServer.SocketServer["Service"],
390 config: HarnessConfig,
391 clientProtocol = NodeClusterSocket.layerClientProtocol
392) =>
393 entities.pipe(
394 Layer.provideMerge(SocketRunner.layer),
395 Layer.provide(runnerHealthLayer(clientProtocol)),
396 Layer.provide(Layer.succeed(SocketServer.SocketServer, socketServer)),
397 Layer.provide(clientProtocol),
398 Layer.provide(ShardingConfig.layer({
399 ...config,
400 runnerAddress: Option.some(address)
401 })),
402 Layer.provide(RpcSerialization.layerSchemaBinary())
403 )
404
405export type RunnerLayer = typeof socketRunnerLayer
406
407const clientLayer = (
408 config: HarnessConfig,
409 clientProtocol: Layer.Layer<Runners.RpcClientProtocol, never, RpcSerialization.RpcSerialization>
410) =>
411 SocketRunner.layerClientOnly.pipe(
412 Layer.provide(clientProtocol),
413 Layer.provide(ShardingConfig.layer(config)),
414 Layer.provide(RpcSerialization.layerSchemaBinary())
415 )
416
417const parseReply = (payload: MessageRow["reply_payload"]): Record<string, unknown> | undefined => {
418 if (payload === null) return undefined
419 return typeof payload === "string" ? JSON.parse(payload) : payload
420}
421
422export const make = Effect.fnUntraced(function*(options: MakeOptions) {
423 const parentScope = yield* Effect.scope
424 const prefix = `cluster_${process.pid}_${nextCluster++}`
425 const config = {
426 ...clusterConfig,
427 ...options.config,
428 shardLockDisableAdvisory: options.lockMode === "row"
429 }
430 const databases = inject("clusterDatabases")
431 const databaseLayer = options.backend === "pg"
432 ? PgClient.layer({ url: Redacted.make(databases.pg), maxConnections: 32 })
433 : MysqlClient.layer({ url: Redacted.make(databases.mysql), maxConnections: 32 })
434 const database = yield* Layer.buildWithScope(databaseLayer, parentScope)
435 const sql = Context.get(database, SqlClient.SqlClient).withoutTransforms()
436 const messageStorage = yield* SqlMessageStorage.layerWith({ prefix }).pipe(
437 Layer.provide(NodeCrypto.layer),
438 Layer.provide(ShardingConfig.layer(config)),
439 Layer.buildWithScope(parentScope),
440 Effect.provide(database)
441 )
442 const shared = Context.merge(database, messageStorage)
443 const entities = typeof options.entities === "function" ? options.entities({ prefix }) : options.entities
444 const runners: Array<RunnerEntry> = []
445 const socketController = makeSocketController()
446
447 const makeRunnerStorage = Effect.fnUntraced(function*(
448 scope: Scope.Closeable,
449 storageConfig: HarnessConfig = config,
450 storageDatabase: Context.Context<SqlClient.SqlClient> = database
451 ) {
452 return yield* SqlRunnerStorage.layerWith({ prefix }).pipe(
453 Layer.provide(ShardingConfig.layer(storageConfig)),
454 Layer.orDie,
455 Layer.buildWithScope(scope),
456 Effect.provide(storageDatabase)
457 )
458 })
459
460 const clientScope = yield* Scope.fork(parentScope)
461 const clientStorage = yield* makeRunnerStorage(clientScope)
462 const clientBase = yield* clientLayer(config, trackedClientProtocolLayer(socketController, clientPeer)).pipe(
463 Layer.buildWithScope(clientScope),
464 Effect.provide(Context.merge(shared, clientStorage))
465 )
466 const workflowEngine = yield* ClusterWorkflowEngine.make.pipe(
467 Effect.provide(Context.merge(shared, clientBase))
468 )
469 const client = Context.add(clientBase, WorkflowEngine.WorkflowEngine, workflowEngine)
470 const clientSharding = Context.get(client, Sharding.Sharding)
471 let nextRunnerIndex = 0
472 const startRunner = Effect.fnUntraced(function*(index: number, startOptions?: StartOptions) {
473 const scope = yield* Scope.fork(parentScope)
474 const runnerConfig: HarnessConfig = {
475 ...config,
476 ...(startOptions?.assignedShardGroups === undefined
477 ? undefined
478 : { assignedShardGroups: startOptions.assignedShardGroups }),
479 ...(startOptions?.runnerShardWeight === undefined
480 ? undefined
481 : { runnerShardWeight: startOptions.runnerShardWeight })
482 }
483 const serverContext = yield* NodeSocketServer.layer({ host: "127.0.0.1", port: 0 }).pipe(
484 Layer.buildWithScope(scope)
485 )
486 const socketServer = Context.get(serverContext, SocketServer.SocketServer)
487 if (socketServer.address._tag === "UnixPathAddress") {
488 return yield* Effect.die("Expected a TCP socket server")
489 }
490 const address = RunnerAddress.make("127.0.0.1", socketServer.address.port)
491 const lockFault = makeLockFaultController(Context.get(database, SqlClient.SqlClient))
492 const runnerDatabase = Context.add(database, SqlClient.SqlClient, lockFault.client)
493 const rawStorageContext = yield* makeRunnerStorage(scope, runnerConfig, runnerDatabase)
494 yield* Scope.addFinalizer(scope, lockFault.set("clear"))
495 const controller = makeRunnerStorageController(Context.get(rawStorageContext, RunnerStorage.RunnerStorage))
496 const storage = Context.make(RunnerStorage.RunnerStorage, controller.controlled)
497 const context = yield* (options.runnerLayer ?? socketRunnerLayer)(
498 address,
499 startOptions?.entities ?? entities,
500 socketServer,
501 runnerConfig,
502 trackedClientProtocolLayer(socketController, addressKey(address))
503 ).pipe(
504 Layer.buildWithScope(scope),
505 Effect.provide(Context.mergeAll(shared, storage))
506 )
507 const sharding = Context.get(context, Sharding.Sharding)
508 let state: ReturnType<ClusterRunner["state"]> = "running"
509 const runner: RunnerEntry = {
510 address,
511 controller,
512 index,
513 lockFault,
514 shardGroups: runnerConfig.assignedShardGroups ?? ShardingConfig.defaults.assignedShardGroups,
515 scope,
516 setState(next) {
517 state = next
518 },
519 sharding,
520 state: () => state
521 }
522 runners.push(runner)
523 return runner as ClusterRunner
524 })
525
526 const start = Effect.fnUntraced(function*(runnerCount: number, startOptions?: StartOptions) {
527 return yield* Effect.forEach(
528 Array.from({ length: runnerCount }, () => nextRunnerIndex++),
529 (index) => startRunner(index, startOptions)
530 )
531 })
532
533 const entryFor = (runner: ClusterRunner) => runners.find((entry) => entry === runner)!
534
535 const stop = Effect.fnUntraced(function*(runner: ClusterRunner) {
536 const entry = entryFor(runner)
537 if (entry.state() === "stopped" || entry.state() === "killed") return
538 if (entry.state() === "frozen") yield* entry.controller.resume
539 entry.setState("stopped")
540 yield* Scope.close(entry.scope, Exit.void)
541 })
542
543 const kill = Effect.fnUntraced(function*(runner: ClusterRunner) {
544 const entry = entryFor(runner)
545 if (entry.state() === "stopped" || entry.state() === "killed") return
546 entry.setState("killed")
547 yield* entry.controller.kill
548 yield* Scope.close(entry.scope, Exit.void)
549 })
550
551 const freeze = Effect.fnUntraced(function*(runner: ClusterRunner) {
552 const entry = entryFor(runner)
553 if (entry.state() !== "running") return
554 entry.setState("frozen")
555 yield* entry.controller.freeze
556 })
557
558 const faultLock = (runner: ClusterRunner, mode: LockFaultMode) => entryFor(runner).lockFault.set(mode)
559
560 const cutSocket = (runner: ClusterRunner, options: CutSocketOptions) => {
561 const runnerAddress = addressKey(runner.address)
562 const peerAddress = options.peer === "client" ? clientPeer : addressKey(options.peer.address)
563 const [source, target] = options.direction === "inbound"
564 ? [peerAddress, runnerAddress]
565 : [runnerAddress, peerAddress]
566 return socketController.cut(source, target)
567 }
568
569 const insertMessage = (row: InsertMessageRow) => {
570 const messages = sql(`${prefix}_messages`)
571 return sql`INSERT INTO ${messages} ${sql.insert({ ...row })}`.unprepared.pipe(Effect.asVoid)
572 }
573
574 const assignmentMap = () => {
575 const assignments: Record<string, ReadonlyArray<string>> = {}
576 const groups = config.availableShardGroups ?? ShardingConfig.defaults.availableShardGroups
577 const shardsPerGroup = config.shardsPerGroup ?? ShardingConfig.defaults.shardsPerGroup
578 for (const group of groups) {
579 for (let id = 1; id <= shardsPerGroup; id++) {
580 const shard = ShardId.make(group, id)
581 assignments[shard.toString()] = runners
582 .filter((runner) => runner.state() === "running" && runner.sharding.hasShardId(shard))
583 .map((runner) => `${runner.address.host}:${runner.address.port}`)
584 }
585 }
586 return assignments
587 }
588
589 const ownersOfShard = (shard: ShardId.ShardId, includeInactive = false) =>
590 runners.filter((runner) => (includeInactive || runner.state() === "running") && runner.sharding.hasShardId(shard))
591
592 const messageCounts = Effect.fnUntraced(function*() {
593 const messages = sql(`${prefix}_messages`)
594 const replies = sql(`${prefix}_replies`)
595 const rows = yield* sql<MessageRow>`
596 SELECT m.processed, r.payload AS reply_payload
597 FROM ${messages} m
598 LEFT JOIN ${replies} r ON r.id = m.last_reply_id
599 WHERE m.kind = 0
600 `
601 let failed = 0
602 let replied = 0
603 let unprocessed = 0
604 for (const row of rows) {
605 if (!row.processed) {
606 unprocessed++
607 continue
608 }
609 const reply = parseReply(row.reply_payload)
610 if (reply?._tag === "Success") replied++
611 if (reply?._tag === "Failure") failed++
612 }
613 return { failed, replied, unprocessed } satisfies MessageCounts
614 })
615
616 const diagnostics = Effect.fnUntraced(function*() {
617 const table = sql(`${prefix}_runners`)
618 const registrations = yield* sql<RegistrationRow>`
619 SELECT address, runner, healthy, last_heartbeat
620 FROM ${table}
621 ORDER BY address
622 `
623 return {
624 assignments: assignmentMap(),
625 messageCounts: yield* messageCounts(),
626 registrations
627 }
628 })
629
630 const waitUntil = Effect.fnUntraced(function*<E, R>(
631 description: string,
632 condition: Effect.Effect<boolean, E, R>,
633 timeout: Duration.Input = "15 seconds"
634 ) {
635 const started = yield* Clock.currentTimeMillis
636 const deadline = started + Duration.toMillis(Duration.fromInputUnsafe(timeout))
637 if (yield* Effect.provide(condition, client)) return
638 while ((yield* Clock.currentTimeMillis) < deadline) {
639 yield* Effect.sleep(100)
640 if (yield* Effect.provide(condition, client)) return
641 }
642 const state = yield* diagnostics()
643 return yield* Effect.fail(new Error(`${description}\n${JSON.stringify(state, null, 2)}`))
644 })
645
646 const waitForStableAssignments = Effect.fnUntraced(function*(timeout?: Duration.Input) {
647 let previous = ""
648 let stablePolls = 0
649 yield* waitUntil(
650 "Shard assignments did not stabilize before the deadline",
651 Effect.sync(() => {
652 const current = assignmentMap()
653 const owners = new Set(Object.values(current).flat())
654 const complete = Object.values(current).every((owners) => owners.length === 1) &&
655 runners.every((runner) =>
656 runner.state() !== "running" || owners.has(`${runner.address.host}:${runner.address.port}`)
657 )
658 const encoded = complete ? JSON.stringify(current) : ""
659 stablePolls = encoded !== "" && encoded === previous ? stablePolls + 1 : 0
660 previous = encoded
661 return stablePolls >= 3
662 }),
663 timeout
664 )
665 return assignmentMap()
666 })
667
668 const shardOfEntity = <Type extends string, Rpcs extends Rpc.Any>(
669 entity: Entity.Entity<Type, Rpcs>,
670 entityId: string
671 ) => entity.getShardId(EntityId.make(entityId)).pipe(Effect.provide(client))
672
673 const ownerOfEntity = Effect.fnUntraced(function*<Type extends string, Rpcs extends Rpc.Any>(
674 entity: Entity.Entity<Type, Rpcs>,
675 entityId: string
676 ) {
677 const shardId = yield* shardOfEntity(entity, entityId)
678 return runners.find((runner) => runner.state() === "running" && runner.sharding.hasShardId(shardId)) as
679 | ClusterRunner
680 | undefined
681 })
682
683 const waitForEntityOwner = <Type extends string, Rpcs extends Rpc.Any>(
684 entity: Entity.Entity<Type, Rpcs>,
685 entityId: string,
686 runner: ClusterRunner,
687 timeout?: Duration.Input
688 ) =>
689 waitUntil(
690 `Entity ${entity.type}/${entityId} was not owned by runner ${runner.index} before the deadline`,
691 Effect.map(ownerOfEntity(entity, entityId), (owner) => owner === runner),
692 timeout
693 )
694
695 const getClient = <Type extends string, Rpcs extends Rpc.Any>(entity: Entity.Entity<Type, Rpcs>) =>
696 entity.client.pipe(Effect.provide(client))
697
698 return {
699 assignmentMap,
700 backend: options.backend,
701 clientSharding,
702 cutSocket,
703 diagnostics,
704 failedMessageCount: Effect.map(messageCounts(), (counts) => counts.failed),
705 faultLock,
706 freeze,
707 getClient,
708 insertMessage,
709 kill,
710 lockMode: options.lockMode ?? "advisory",
711 messageCounts,
712 ownerOfEntity,
713 ownersOfShard,
714 prefix,
715 repliedMessageCount: Effect.map(messageCounts(), (counts) => counts.replied),
716 runners: runners as ReadonlyArray<ClusterRunner>,
717 shardOfEntity,
718 start,
719 stop,
720 unprocessedMessageCount: Effect.map(messageCounts(), (counts) => counts.unprocessed),
721 waitForEntityOwner,
722 waitForStableAssignments,
723 waitUntil,
724 workflowEngine
725 } as const
726})