1/**
2 * Runs shard ownership and message routing for Effect Cluster.
3 *
4 * `Sharding` decides which shard owns an entity id, tracks which shards belong
5 * to the local runner, and sends cluster messages to local handlers or remote
6 * runners. It also registers entities and singletons, creates clients for
7 * entity requests, polls stored messages, and tracks shutdown state. The main
8 * layer connects these responsibilities to runner communication, storage,
9 * health checks, configuration, and local resources.
10 *
11 * @stability unstable
12 * @since 4.0.0
13 */
14import * as Arr from "../Array.ts"
15import * as Cause from "../Cause.ts"
16import { Clock } from "../Clock.ts"
17import * as Context from "../Context.ts"
18import type { Input } from "../Duration.ts"
19import * as Duration from "../Duration.ts"
20import * as Effect from "../Effect.ts"
21import * as Equal from "../Equal.ts"
22import type * as Exit from "../Exit.ts"
23import * as Fiber from "../Fiber.ts"
24import * as FiberMap from "../FiberMap.ts"
25import { constant, flow } from "../Function.ts"
26import * as HashRing from "../HashRing.ts"
27import * as InternalRecord from "../internal/record.ts"
28import * as Latch from "../Latch.ts"
29import * as Layer from "../Layer.ts"
30import * as MutableHashMap from "../MutableHashMap.ts"
31import * as MutableHashSet from "../MutableHashSet.ts"
32import * as MutableRef from "../MutableRef.ts"
33import * as Option from "../Option.ts"
34import * as PubSub from "../PubSub.ts"
35import { CurrentLogAnnotations } from "../References.ts"
36import * as Result from "../Result.ts"
37import type * as Rpc from "../rpc/Rpc.ts"
38import * as RpcClient from "../rpc/RpcClient.ts"
39import type { FromServer } from "../rpc/RpcMessage.ts"
40import * as Schedule from "../Schedule.ts"
41import * as Scope from "../Scope.ts"
42import * as Semaphore from "../Semaphore.ts"
43import * as Stream from "../Stream.ts"
44import type { MailboxFull, PersistenceError } from "./ClusterError.ts"
45import { AlreadyProcessingMessage, EntityNotAssignedToRunner } from "./ClusterError.ts"
46import * as ClusterMetrics from "./ClusterMetrics.ts"
47import { Persisted } from "./ClusterSchema.ts"
48import * as ClusterSchema from "./ClusterSchema.ts"
49import { CurrentAddress, type CurrentRunnerAddress, type Entity, type HandlersFrom } from "./Entity.ts"
50import type { EntityAddress } from "./EntityAddress.ts"
51import { make as makeEntityAddress } from "./EntityAddress.ts"
52import type { EntityId } from "./EntityId.ts"
53import { make as makeEntityId } from "./EntityId.ts"
54import * as Envelope from "./Envelope.ts"
55import * as ClusterAbandon from "./internal/clusterAbandon.ts"
56import * as EntityManager from "./internal/entityManager.ts"
57import { EntityReaper } from "./internal/entityReaper.ts"
58import { hashString } from "./internal/hash.ts"
59import * as ActiveTeardown from "./internal/interruptors.ts"
60import { ResourceMap } from "./internal/resourceMap.ts"
61import { effectiveInterval } from "./internal/shardLock.ts"
62import * as Message from "./Message.ts"
63import * as MessageStorage from "./MessageStorage.ts"
64import * as Reply from "./Reply.ts"
65import { Runner } from "./Runner.ts"
66import type { RunnerAddress } from "./RunnerAddress.ts"
67import * as RunnerHealth from "./RunnerHealth.ts"
68import { Runners } from "./Runners.ts"
69import { RunnerStorage } from "./RunnerStorage.ts"
70import type { ShardId } from "./ShardId.ts"
71import { make as makeShardId } from "./ShardId.ts"
72import { shardGroupConfig, ShardingConfig } from "./ShardingConfig.ts"
73import { EntityRegistered, type ShardingRegistrationEvent, SingletonRegistered } from "./ShardingRegistrationEvent.ts"
74import { SingletonAddress } from "./SingletonAddress.ts"
75import * as Snowflake from "./Snowflake.ts"
76
77/**
78 * Service that registers entities and singletons, routes messages to owned
79 * shards, generates runner-local snowflake ids, and polls
80 * storage for persisted work.
81 *
82 * **When to use**
83 *
84 * Use to access or provide cluster routing, shard ownership, entity
85 * registration, singleton registration, and persisted-work polling.
86 *
87 * @stability unstable
88 * @category services
89 * @since 4.0.0
90 */
91export class Sharding extends Context.Service<Sharding, {
92 /**
93 * Returns a stream of events that occur when the runner registers entities or
94 * singletons.
95 */
96 readonly getRegistrationEvents: Stream.Stream<ShardingRegistrationEvent>
97
98 /**
99 * Returns the `ShardId` of the shard to which the entity at the specified
100 * `address` is assigned.
101 */
102 readonly getShardId: (entityId: EntityId, group: string) => ShardId
103
104 /**
105 * Returns `true` if the specified `shardId` is assigned to this runner.
106 */
107 readonly hasShardId: (shardId: ShardId) => boolean
108
109 /**
110 * Generate a Snowflake ID that is unique to this runner.
111 */
112 readonly getSnowflake: Effect.Effect<Snowflake.Snowflake>
113
114 /**
115 * Returns `true` if sharding is shutting down, `false` otherwise.
116 */
117 readonly isShutdown: Effect.Effect<boolean>
118
119 /**
120 * Constructs a `RpcClient` which can be used to send messages to the
121 * specified `Entity`.
122 */
123 readonly makeClient: <Type extends string, Rpcs extends Rpc.Any>(
124 entity: Entity<Type, Rpcs>
125 ) => Effect.Effect<
126 (
127 entityId: string
128 ) => RpcClient.RpcClient.From<
129 Rpcs,
130 MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
131 >
132 >
133
134 /**
135 * Registers a new entity with the runner.
136 */
137 readonly registerEntity: <Type extends string, Rpcs extends Rpc.Any, Handlers extends HandlersFrom<Rpcs>, RX>(
138 entity: Entity<Type, Rpcs>,
139 handlers: Effect.Effect<Handlers, never, RX>,
140 options?: {
141 readonly maxIdleTime?: Input | undefined
142 readonly concurrency?: number | "unbounded" | undefined
143 readonly mailboxCapacity?: number | "unbounded" | undefined
144 readonly disableFatalDefects?: boolean | undefined
145 readonly defectRetryPolicy?: Schedule.Schedule<any, unknown> | undefined
146 readonly spanAttributes?: Record<string, string> | undefined
147 }
148 ) => Effect.Effect<
149 void,
150 never,
151 | Scope.Scope
152 | Rpc.ServicesServer<Rpcs>
153 | Rpc.Middleware<Rpcs>
154 | Exclude<RX, Scope.Scope | CurrentAddress | CurrentRunnerAddress>
155 >
156
157 /**
158 * Registers a new singleton with the runner.
159 */
160 readonly registerSingleton: <E, R>(
161 name: string,
162 run: Effect.Effect<void, E, R>,
163 options?: {
164 readonly shardGroup?: string | undefined
165 }
166 ) => Effect.Effect<void, never, R | Scope.Scope>
167
168 /**
169 * Sends a message to the specified entity.
170 */
171 readonly send: (message: Message.Incoming<any>) => Effect.Effect<
172 void,
173 EntityNotAssignedToRunner | MailboxFull | AlreadyProcessingMessage
174 >
175
176 /**
177 * Sends an outgoing message
178 */
179 readonly sendOutgoing: (
180 message: Message.Outgoing<any>,
181 discard: boolean
182 ) => Effect.Effect<
183 void,
184 MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
185 >
186
187 /**
188 * Notify sharding that a message has been persisted to storage.
189 */
190 readonly notify: (message: Message.Incoming<any>, options?: {
191 readonly waitUntilRead?: boolean | undefined
192 }) => Effect.Effect<
193 void,
194 EntityNotAssignedToRunner | AlreadyProcessingMessage
195 >
196
197 /**
198 * Reset the state of a message. When options.expectedReplyId is provided,
199 * reset only if it is still the latest reply. Otherwise reset unconditionally.
200 */
201 readonly reset: (
202 requestId: Snowflake.Snowflake,
203 options?: { readonly expectedReplyId?: Snowflake.Snowflake | undefined }
204 ) => Effect.Effect<boolean>
205
206 /**
207 * Trigger a storage read, which will read all unprocessed messages.
208 */
209 readonly pollStorage: Effect.Effect<void>
210
211 /**
212 * Retrieves the active entity count for the current runner.
213 */
214 readonly activeEntityCount: Effect.Effect<number>
215}>()("effect/cluster/Sharding") {}
216
217// -----------------------------------------------------------------------------
218// Implementation
219// -----------------------------------------------------------------------------
220
221interface EntityManagerState {
222 readonly entity: Entity<any, any>
223 readonly manager: EntityManager.EntityManager
224 status: "alive" | "closing" | "closed"
225}
226
227const make = Effect.gen(function*() {
228 const config = yield* ShardingConfig
229 const shardLockInterval = effectiveInterval(config)
230 const shardGroups = shardGroupConfig(config)
231 const getRunnerAddress = () => Option.getOrUndefined(config.runnerAddress)
232 const clock = yield* Clock
233
234 const runnersService = yield* Runners
235 const runnerHealth = yield* RunnerHealth.RunnerHealth
236 const snowflakeGen = yield* Snowflake.Generator
237 const shardingScope = yield* Effect.scope
238 const isShutdown = MutableRef.make(false)
239 const services = Context.omit(Scope.Scope)(yield* Effect.context<ShardingConfig>())
240 const runFork = flow(
241 Effect.runForkWith(services),
242 Fiber.runIn(shardingScope)
243 )
244
245 const storage = yield* MessageStorage.MessageStorage
246 const storageEnabled = storage !== MessageStorage.noop
247 const runnerStorage = yield* RunnerStorage
248
249 const entityManagers = new Map<string, EntityManagerState>()
250 let entityRegistrationStartMillis: number | undefined
251 let entityRegistrationFallbackStartMillis: number | undefined
252 const entityRegistrationTimeoutMillis = Duration.toMillis(
253 Duration.fromInputUnsafe(config.entityRegistrationTimeout)
254 )
255 const entityRegistrationTimeRemaining = () => {
256 const now = clock.currentTimeMillisUnsafe()
257 const registrationStarted = entityRegistrationStartMillis !== undefined
258 // The first missing entity observed by either the storage read loop or a
259 // local send starts the shared fallback registration window.
260 const timeoutStartMillis = entityRegistrationStartMillis ??
261 (entityRegistrationFallbackStartMillis ??= now)
262 // If registration never starts, allow two intervals from the first missing
263 // entity before failing.
264 const timeoutMillis = registrationStarted
265 ? entityRegistrationTimeoutMillis
266 : entityRegistrationTimeoutMillis * 2
267 return Math.max(0, timeoutMillis - (now - timeoutStartMillis))
268 }
269 const entityNotRegistered = (entityType: string) =>
270 Effect.die(new Error(`Entity type '${entityType}' not registered`))
271
272 const shardAssignments = MutableHashMap.empty<ShardId, RunnerAddress>()
273 const selfShards = MutableHashSet.empty<ShardId>()
274 // open while shard lock storage is healthy
275 const shardLocksHealthyLatch = Latch.makeUnsafe(true)
276
277 // the active shards are the ones that we have acquired the lock for
278 const acquiredShards = MutableHashSet.empty<ShardId>()
279 const activeShardsLatch = yield* Latch.make(false)
280
281 const events = yield* PubSub.unbounded<ShardingRegistrationEvent>()
282 const getRegistrationEvents: Stream.Stream<ShardingRegistrationEvent> = Stream.fromPubSub(events)
283
284 const isLocalRunner = (address: RunnerAddress) => {
285 const runnerAddress = getRunnerAddress()
286 return runnerAddress !== undefined && Equal.equals(address, runnerAddress)
287 }
288
289 function getShardId(entityId: EntityId, group: string): ShardId {
290 const id = Math.abs(hashString(entityId) % config.shardsPerGroup) + 1
291 return makeShardId(group, id)
292 }
293
294 function isEntityOnLocalShards(address: EntityAddress): boolean {
295 return MutableHashSet.has(acquiredShards, address.shardId)
296 }
297
298 yield* Scope.addFinalizer(
299 shardingScope,
300 Effect.logDebug("Shutdown complete").pipe(Effect.annotateLogs({
301 module: "effect/cluster/Sharding"
302 }))
303 )
304
305 // --- Shard acquisition ---
306 //
307 // Responsible for acquiring and releasing shards from RunnerStorage.
308 //
309 // This should be shutdown last, when all entities have been shutdown, to
310 // allow them to move to another runner.
311
312 const releasingShards = MutableHashSet.empty<ShardId>()
313 // Shards whose entities must be force interrupted and locks fully released
314 // before normal reacquisition.
315 const forceReleasingShards = MutableHashSet.empty<ShardId>()
316 const initialRunnerAddress = getRunnerAddress()
317 if (initialRunnerAddress) {
318 const selfAddress = initialRunnerAddress
319 yield* Scope.addFinalizerExit(shardingScope, () => {
320 // the locks expire over time, so if this fails we ignore it
321 return Effect.ignore(runnerStorage.releaseAll(selfAddress))
322 })
323
324 const releaseShardsMap = yield* FiberMap.make<ShardId>()
325 let forcedShardReleaseRunning = false
326 // Interrupt the shards' entities, wait for lock health, run the storage
327 // release, then clear the shards' bookkeeping.
328 const runShardRelease = Effect.fnUntraced(function*<E>(
329 shardIds: ReadonlyArray<ShardId>,
330 force: boolean,
331 release: Effect.Effect<void, E>
332 ) {
333 const fibers = Arr.empty<Fiber.Fiber<void>>()
334 for (const shardId of shardIds) {
335 for (const state of entityManagers.values()) {
336 if (state.status === "closed") continue
337 fibers.push(yield* Effect.forkScoped(state.manager.interruptShard(shardId, { force })))
338 }
339 }
340 yield* Fiber.joinAll(fibers)
341 yield* shardLocksHealthyLatch.await
342 yield* release
343 for (const shardId of shardIds) {
344 MutableHashSet.remove(releasingShards, shardId)
345 MutableHashSet.remove(forceReleasingShards, shardId)
346 // During shutdown, parked waiters are interrupted: this runner will
347 // never observe the reply. During a rebalance the waiters are resumed
348 // with `EntityNotAssignedToRunner`, and the client falls back to
349 // waiting for the reply via storage while the entity moves.
350 yield* storage.unregisterShardReplyHandlers(shardId, {
351 interrupt: MutableRef.get(isShutdown)
352 })
353 }
354 })
355 const retryShardRelease =
356 (annotations: { readonly fiber: string; readonly shardId?: ShardId }) =>
357 <A, E, R>(effect: Effect.Effect<A, E, R>) =>
358 effect.pipe(
359 Effect.sandbox,
360 Effect.tapError((cause) =>
361 Effect.logDebug(`Could not release shards, retrying`, cause).pipe(
362 Effect.annotateLogs({
363 module: "effect/cluster/Sharding",
364 runner: selfAddress,
365 ...annotations
366 }),
367 // Effect.eventually retries immediately, so space failures to
368 // avoid hot-looping while storage is unavailable.
369 Effect.andThen(Effect.sleep(50))
370 )
371 ),
372 Effect.eventually
373 )
374 const releaseShard = Effect.fnUntraced(
375 function*(shardId: ShardId) {
376 yield* runShardRelease(
377 [shardId],
378 MutableHashSet.has(forceReleasingShards, shardId),
379 runnerStorage.release(selfAddress, shardId)
380 )
381 },
382 (effect, shardId) =>
383 effect.pipe(
384 retryShardRelease({ fiber: "releaseShard", shardId }),
385 FiberMap.run(releaseShardsMap, shardId, { onlyIfMissing: true })
386 )
387 )
388 // The forced release ends with `runnerStorage.releaseAll`, which drops
389 // every lock held by this runner. Shards must not be reacquired through the
390 // normal path while it is pending, otherwise the bulk release wipes a lock
391 // the runner already considers acquired.
392 const forcedShardReleasePending = () => forcedShardReleaseRunning || MutableHashSet.size(forceReleasingShards) > 0
393 const releaseForcedShards = Effect.suspend(() => {
394 if (forcedShardReleaseRunning || MutableHashSet.size(forceReleasingShards) === 0) {
395 return Effect.void
396 }
397 forcedShardReleaseRunning = true
398 const shardIds = [...forceReleasingShards]
399 return runShardRelease(shardIds, true, runnerStorage.releaseAll(selfAddress)).pipe(
400 retryShardRelease({ fiber: "releaseForcedShards" }),
401 Effect.ensuring(Effect.sync(() => {
402 forcedShardReleaseRunning = false
403 activeShardsLatch.openUnsafe()
404 })),
405 Effect.forkIn(shardingScope),
406 Effect.asVoid
407 )
408 })
409 const releaseShards = Effect.gen(function*() {
410 yield* releaseForcedShards
411 for (const shardId of releasingShards) {
412 if (MutableHashSet.has(forceReleasingShards, shardId)) continue
413 if (FiberMap.hasUnsafe(releaseShardsMap, shardId)) continue
414 yield* releaseShard(shardId)
415 }
416 })
417
418 yield* Effect.gen(function*() {
419 activeShardsLatch.openUnsafe()
420
421 while (true) {
422 yield* activeShardsLatch.await
423 activeShardsLatch.closeUnsafe()
424
425 // if a shard is no longer assigned to this runner, we release it
426 for (const shardId of acquiredShards) {
427 if (MutableHashSet.has(selfShards, shardId)) continue
428 MutableHashSet.remove(acquiredShards, shardId)
429 MutableHashSet.add(releasingShards, shardId)
430 }
431
432 if (MutableHashSet.size(releasingShards) > 0 || MutableHashSet.size(forceReleasingShards) > 0) {
433 yield* Effect.forkIn(syncSingletons, shardingScope)
434 yield* releaseShards
435 }
436
437 if (!shardLocksHealthyLatch.isOpen()) {
438 continue
439 }
440
441 // Wait for the pending bulk release before reacquiring, so it cannot
442 // drop a lock acquired here. `releaseForcedShards` reopens the latch.
443 if (forcedShardReleasePending()) {
444 continue
445 }
446
447 // if a shard has been assigned to this runner, we acquire it
448 const unacquiredShards = MutableHashSet.empty<ShardId>()
449 for (const shardId of selfShards) {
450 if (MutableHashSet.has(acquiredShards, shardId) || MutableHashSet.has(releasingShards, shardId)) continue
451 MutableHashSet.add(unacquiredShards, shardId)
452 }
453
454 if (MutableHashSet.size(unacquiredShards) === 0) {
455 continue
456 }
457
458 const oacquired = yield* runnerStorage.acquire(selfAddress, unacquiredShards).pipe(
459 Effect.timeoutOption(shardLockInterval)
460 )
461 if (Option.isNone(oacquired)) {
462 activeShardsLatch.openUnsafe()
463 continue
464 }
465
466 const acquired = oacquired.value
467 yield* storage.resetShards(acquired).pipe(
468 Effect.ignore,
469 Effect.timeoutOption(shardLockInterval)
470 )
471 // A forced release can start while `acquire` is in flight, so re-check
472 // it here as well as before acquiring.
473 const forcedReleasePending = forcedShardReleasePending()
474 for (const shardId of acquired) {
475 if (
476 !shardLocksHealthyLatch.isOpen() ||
477 forcedReleasePending ||
478 MutableHashSet.has(releasingShards, shardId) ||
479 !MutableHashSet.has(selfShards, shardId)
480 ) {
481 MutableHashSet.add(releasingShards, shardId)
482 continue
483 }
484 MutableHashSet.add(acquiredShards, shardId)
485 }
486 if (acquired.length > 0) {
487 yield* storageReadLatch.open
488 yield* Effect.forkIn(syncSingletons, shardingScope)
489
490 // update metrics
491 ClusterMetrics.shards.updateUnsafe(BigInt(MutableHashSet.size(acquiredShards)), Context.empty())
492 }
493 yield* Effect.sleep(1000)
494 activeShardsLatch.openUnsafe()
495 }
496 }).pipe(
497 Effect.catchCause((cause) => Effect.logWarning("Could not acquire/release shards", cause)),
498 Effect.repeat(Schedule.spaced(config.entityMessagePollInterval)),
499 Effect.annotateLogs({
500 module: "effect/cluster/Sharding",
501 fiber: "Shard acquisition loop",
502 runner: selfAddress
503 }),
504 Effect.forkIn(shardingScope)
505 )
506
507 const markShardLocksUnhealthy = (cause: Cause.Cause<unknown>) =>
508 Effect.suspend(() => {
509 if (!shardLocksHealthyLatch.closeUnsafe()) return Effect.void
510
511 const affectedShards = MutableHashSet.fromIterable([...acquiredShards, ...releasingShards])
512 MutableHashSet.clear(selfShards)
513 MutableHashSet.clear(acquiredShards)
514 for (const shardId of affectedShards) {
515 MutableHashSet.add(releasingShards, shardId)
516 MutableHashSet.add(forceReleasingShards, shardId)
517 }
518 ClusterMetrics.shards.updateUnsafe(BigInt(0), Context.empty())
519 activeShardsLatch.openUnsafe()
520
521 return Effect.gen(function*() {
522 yield* Effect.logError("Shard lock storage is unhealthy", cause)
523 yield* Effect.forkIn(syncSingletons, shardingScope, { startImmediately: true })
524
525 for (const shardId of affectedShards) {
526 for (const state of entityManagers.values()) {
527 if (state.status === "closed") continue
528 yield* Effect.forkIn(
529 state.manager.interruptShard(shardId, { force: true }),
530 shardingScope,
531 { startImmediately: true }
532 )
533 }
534 }
535 activeShardsLatch.openUnsafe()
536 })
537 })
538
539 const markShardLocksHealthy = Effect.suspend(() => {
540 if (!shardLocksHealthyLatch.openUnsafe()) return Effect.void
541
542 MutableHashSet.clear(selfShards)
543 MutableHashMap.forEach(shardAssignments, (runner, shardId) => {
544 if (isLocalRunner(runner)) {
545 MutableHashSet.add(selfShards, shardId)
546 }
547 })
548 activeShardsLatch.openUnsafe()
549 return Effect.logInfo("Shard lock storage has recovered")
550 })
551
552 const refreshShardLocks = Effect.gen(function*() {
553 const refreshed = [...acquiredShards, ...releasingShards]
554 const acquired = yield* runnerStorage.refresh(selfAddress, refreshed)
555 // Shards acquired during the request are absent from its response.
556 for (const shardId of refreshed) {
557 if (MutableHashSet.has(acquiredShards, shardId) && !acquired.includes(shardId)) {
558 MutableHashSet.remove(acquiredShards, shardId)
559 MutableHashSet.add(releasingShards, shardId)
560 }
561 }
562 for (let i = 0; i < acquired.length; i++) {
563 const shardId = acquired[i]
564 if (!MutableHashSet.has(selfShards, shardId)) {
565 MutableHashSet.remove(acquiredShards, shardId)
566 MutableHashSet.add(releasingShards, shardId)
567 }
568 }
569 if (MutableHashSet.size(releasingShards) > 0) {
570 yield* activeShardsLatch.open
571 }
572 }).pipe(
573 Effect.retry({
574 times: 5,
575 schedule: Schedule.spaced(50)
576 }),
577 Effect.timeout(shardLockInterval),
578 Effect.catchCause(markShardLocksUnhealthy)
579 )
580
581 const probeShardLocks = runnerStorage.refresh(selfAddress, []).pipe(
582 Effect.timeout(shardLockInterval),
583 Effect.andThen(markShardLocksHealthy),
584 Effect.catchCause((cause) => Effect.logWarning("Shard lock storage is still unhealthy", cause))
585 )
586
587 // Refresh shard locks at the lease-safe interval, or probe storage while
588 // lock ownership is uncertain.
589 yield* Effect.suspend(() => shardLocksHealthyLatch.isOpen() ? refreshShardLocks : probeShardLocks).pipe(
590 Effect.repeat(Schedule.fixed(shardLockInterval)),
591 Effect.forever,
592 Effect.forkIn(shardingScope)
593 )
594
595 // open the shard latch every poll interval
596 yield* activeShardsLatch.open.pipe(
597 Effect.delay(config.entityMessagePollInterval),
598 Effect.forever,
599 Effect.forkIn(shardingScope)
600 )
601 }
602
603 // --- Storage inbox ---
604 //
605 // Responsible for reading unprocessed messages from storage and sending them
606 // to the appropriate entity manager.
607 //
608 // This should be shutdown before shard acquisition, to ensure no messages are
609 // being processed before the shards are released.
610 //
611 // It should also be shutdown after the entity managers, to ensure interrupt
612 // & ack envelopes can still be processed.
613
614 const storageReadLatch = yield* Latch.make(true)
615 const openStorageReadLatch = constant(Effect.asVoid(storageReadLatch.open))
616
617 const storageReadLock = Semaphore.makeUnsafe(1)
618 const withStorageReadLock = storageReadLock.withPermits(1)
619
620 // --- Entity residency ---
621 //
622 // A runner-wide counter that bounds how many entities can be resident at
623 // the same time. Entity managers reserve a slot before spawning an entity,
624 // so the sequential storage read loop and concurrent volatile sends share
625 // the same cap.
626
627 const maxResidentEntities = config.maxResidentEntities
628 let residentEntityCount = 0
629 const residencyAtCapacityUnsafe = () =>
630 maxResidentEntities !== "unbounded" && residentEntityCount >= maxResidentEntities
631 const residency: EntityManager.Residency = {
632 admitUnsafe() {
633 if (residencyAtCapacityUnsafe()) return false
634 residentEntityCount++
635 return true
636 },
637 releaseUnsafe() {
638 residentEntityCount--
639 // a slot has freed up, so the storage read loop can admit messages for
640 // entities it previously had to skip
641 if (maxResidentEntities !== "unbounded") {
642 storageReadLatch.openUnsafe()
643 }
644 }
645 }
646
647 if (storageEnabled && initialRunnerAddress) {
648 const selfAddress = initialRunnerAddress
649
650 yield* Effect.gen(function*() {
651 yield* Effect.logDebug("Starting")
652 yield* Effect.addFinalizer(() => Effect.logDebug("Shutting down"))
653
654 let index = 0
655 let messages: Array<Message.Incoming<any>> = []
656 let deliveredThisRead = false
657 const removableNotifications = new Set<PendingNotification>()
658 const resetAddresses = MutableHashSet.empty<EntityAddress>()
659 const resetRequestIds: Array<Snowflake.Snowflake> = []
660 const cappedAddresses = MutableHashSet.empty<EntityAddress>()
661
662 const markDelivered = Effect.sync(() => {
663 deliveredThisRead = true
664 })
665
666 // A completed request may remain deduplicated after storage claims its reset.
667 // Release the claim so the reset can be redelivered immediately.
668 const releaseCompletedClaim = (message: Message.Incoming<any>) => {
669 const state = entityManagers.get(message.envelope.address.entityType)
670 if (!state?.manager.isProcessingFor(message, { excludeCompleted: true })) {
671 resetRequestIds.push(message.envelope.requestId)
672 }
673 }
674
675 const processMessages = Effect.whileLoop({
676 while: () => index < messages.length,
677 step: () => index++,
678 body: () => send
679 })
680
681 const readAndProcess = Effect.fnUntraced(function*(options?: {
682 readonly limit?: number | undefined
683 readonly addresses?: ReadonlyArray<EntityAddress> | undefined
684 }) {
685 messages = yield* storage.unprocessedMessages(acquiredShards, options)
686 index = 0
687 yield* processMessages
688 return messages.length
689 })
690
691 const send = Effect.catchCause(
692 Effect.suspend(() => {
693 const message = messages[index]
694 const address = message.envelope.address
695 if (!MutableHashSet.has(acquiredShards, address.shardId)) {
696 return Effect.void
697 }
698 const state = entityManagers.get(address.entityType)
699 if (!state) {
700 if (entityRegistrationTimeRemaining() > 0) {
701 // reset address in the case that the entity is slow to register
702 MutableHashSet.add(resetAddresses, address)
703 return Effect.void
704 }
705 // if the entity did not register in time, we save a defect reply
706 return entityNotRegistered(address.entityType)
707 } else if (state.status === "closed") {
708 return Effect.void
709 }
710
711 const isProcessing = state.manager.isProcessingFor(message)
712
713 if (message._tag === "IncomingEnvelope" && isProcessing) {
714 // If the message might affect a currently processing request, we
715 // send it to the entity manager to be processed.
716 return Effect.tap(state.manager.send(message), markDelivered)
717 } else if (isProcessing || state.status === "closing") {
718 // If the request is already processing, we skip it.
719 // Or if the entity is closing, we skip all incoming messages.
720 if (isProcessing) releaseCompletedClaim(message)
721 return Effect.void
722 } else if (message._tag === "IncomingRequest" && pendingNotifications.has(message.envelope.requestId)) {
723 const entry = pendingNotifications.get(message.envelope.requestId)!
724 pendingNotifications.delete(message.envelope.requestId)
725 removableNotifications.delete(entry)
726 entry.resume(Effect.void)
727 }
728
729 // The runner is at entity capacity, so messages for entities that
730 // are not already resident are skipped. Their claims are released
731 // after the read, so they stay eligible for once a slot frees up.
732 if (residencyAtCapacityUnsafe() && !state.manager.isResidentUnsafe(address)) {
733 MutableHashSet.add(cappedAddresses, address)
734 return Effect.void
735 }
736
737 // If the entity was resuming in another fiber, we add the message
738 // id to the unprocessed set.
739 const resumptionState = MutableHashMap.get(entityResumptionState, address)
740 if (Option.isSome(resumptionState)) {
741 resumptionState.value.unprocessed.add(message.envelope.requestId)
742 if (message.envelope._tag === "Interrupt") {
743 resumptionState.value.interrupts.set(message.envelope.requestId, message as Message.IncomingEnvelope)
744 }
745 return Effect.void
746 }
747 return Effect.tap(state.manager.send(message), markDelivered)
748 }),
749 (cause) => {
750 const message = messages[index]
751 const error = Cause.findError(cause)
752 // if we get a defect, then update storage
753 if (Result.isFailure(error)) {
754 if (Cause.hasInterrupts(cause)) {
755 return Effect.void
756 }
757 return storage.saveReply(Reply.ReplyWithContext.fromDefect({
758 id: snowflakeGen.nextUnsafe(),
759 requestId: message.envelope.requestId,
760 defect: Cause.squash(cause)
761 }))
762 }
763 if (error.success._tag === "AlreadyProcessingMessage") {
764 // Completion raced with decoding.
765 releaseCompletedClaim(message)
766 return Effect.void
767 }
768 if (error.success._tag === "MailboxFull") {
769 const address = message.envelope.address
770 const state = entityManagers.get(address.entityType)
771 // A resident entity has a full per-entity mailbox; its messages
772 // are resumed from storage once there is capacity again.
773 // MailboxFull can only happen for requests, so this cast is safe
774 if (message._tag === "IncomingRequest" && state?.manager.isResidentUnsafe(address)) {
775 return resumeEntityFromStorage(message as Message.IncomingRequest<any>)
776 }
777 // Otherwise the runner ran out of entity slots while the message
778 // was in flight; leave it in storage.
779 MutableHashSet.add(cappedAddresses, address)
780 return Effect.void
781 }
782 return Effect.void
783 }
784 )
785
786 const batchSize = Math.max(1, config.unprocessedMessageBatchSize)
787
788 while (true) {
789 // wait for the next poll interval, or if we get notified of a change
790 yield* storageReadLatch.await
791
792 // if we get notified of a change, ensure we start a read immediately
793 // next iteration
794 storageReadLatch.closeUnsafe()
795
796 // the lock is used to ensure resuming entities have a garantee that no
797 // more items are added to the unprocessed set while the semaphore is
798 // acquired.
799 yield* storageReadLock.take(1)
800
801 entityManagers.forEach((state) => state.manager.clearProcessed())
802 if (pendingNotifications.size > 0) {
803 pendingNotifications.forEach((entry) => removableNotifications.add(entry))
804 }
805
806 deliveredThisRead = false
807 let readCount = 0
808 let fullBatch = false
809 let exhaustive = false
810
811 // First deliver messages for entities that are already resident, so
812 // they keep making progress even when the runner is at entity
813 // capacity.
814 if (residencyAtCapacityUnsafe()) {
815 const residentAddresses: Array<EntityAddress> = []
816 for (const state of entityManagers.values()) {
817 if (state.status === "closed") continue
818 for (const address of state.manager.residentAddressesUnsafe()) {
819 residentAddresses.push(address)
820 }
821 }
822 if (residentAddresses.length > 0) {
823 readCount = yield* readAndProcess({ limit: batchSize, addresses: residentAddresses })
824 fullBatch = readCount >= batchSize
825 }
826 }
827
828 // Then walk the remaining messages, spawning new entities while the
829 // runner has entity slots left.
830 if (!residencyAtCapacityUnsafe() && readCount < batchSize) {
831 const limit = batchSize - readCount
832 const read = yield* readAndProcess({ limit })
833 if (read >= limit) {
834 fullBatch = true
835 } else if (MutableHashSet.size(cappedAddresses) === 0) {
836 exhaustive = true
837 }
838 }
839
840 if (removableNotifications.size > 0) {
841 // On an exhaustive read, a message that was not seen is no longer
842 // processable by this runner. At the entity cap, the messages are
843 // safely persisted and delivered once a slot frees up, so persisted
844 // senders succeed instead of failing. Otherwise the read was
845 // truncated by the batch size and the notifications stay registered
846 // for the immediately following read.
847 const capped = MutableHashSet.size(cappedAddresses) > 0 || residencyAtCapacityUnsafe()
848 if (exhaustive || capped) {
849 removableNotifications.forEach(({ message, resume }) => {
850 pendingNotifications.delete(message.envelope.requestId)
851 resume(
852 exhaustive
853 ? Effect.fail(new EntityNotAssignedToRunner({ address: message.envelope.address }))
854 : Effect.void
855 )
856 })
857 }
858 removableNotifications.clear()
859 }
860 if (MutableHashSet.size(resetAddresses) > 0) {
861 for (const address of resetAddresses) {
862 yield* Effect.logWarning("Could not find entity manager for address, retrying").pipe(
863 Effect.annotateLogs({ address })
864 )
865 yield* Effect.forkIn(storage.resetAddress(address), shardingScope)
866 }
867 MutableHashSet.clear(resetAddresses)
868 }
869 // Capture claims skipped because of the entity cap. They are reset in
870 // one storage operation after releasing the read lock.
871 let cappedAddressesToReset: Array<EntityAddress> | undefined
872 if (MutableHashSet.size(cappedAddresses) > 0) {
873 cappedAddressesToReset = Arr.fromIterable(cappedAddresses)
874 MutableHashSet.clear(cappedAddresses)
875 }
876
877 // let the resuming entities check if they are done
878 yield* storageReadLock.release(1)
879
880 if (resetRequestIds.length > 0) {
881 yield* Effect.ignore(storage.resetRequests(resetRequestIds))
882 resetRequestIds.length = 0
883 }
884
885 if (cappedAddressesToReset !== undefined) {
886 yield* Effect.ignore(storage.resetAddresses(cappedAddressesToReset))
887 }
888
889 // A full batch means more messages could be waiting; start the next
890 // read immediately, as long as this read made progress (a full batch
891 // of skipped messages must not spin the loop). When the runner is at
892 // capacity, the next read waits for an entity to be removed or the
893 // next poll interval instead.
894 if (fullBatch && deliveredThisRead) {
895 storageReadLatch.openUnsafe()
896 }
897 }
898 }).pipe(
899 Effect.scoped,
900 Effect.ensuring(storageReadLock.releaseAll),
901 Effect.catchCause((cause) => Effect.logWarning("Could not read messages from storage", cause)),
902 Effect.forever,
903 Effect.annotateLogs({
904 module: "effect/cluster/Sharding",
905 fiber: "Storage read loop",
906 runner: selfAddress
907 }),
908 Effect.forkIn(shardingScope)
909 )
910
911 // open the storage latch every poll interval
912 yield* storageReadLatch.open.pipe(
913 Effect.delay(config.entityMessagePollInterval),
914 Effect.forever,
915 Effect.forkIn(shardingScope)
916 )
917
918 // Resume unprocessed messages for entities that reached a full mailbox.
919 const entityResumptionState = MutableHashMap.empty<EntityAddress, {
920 unprocessed: Set<Snowflake.Snowflake>
921 interrupts: Map<Snowflake.Snowflake, Message.IncomingEnvelope>
922 }>()
923 const resumeEntityFromStorage = (lastReceivedMessage: Message.IncomingRequest<any>) => {
924 const address = lastReceivedMessage.envelope.address
925 const resumptionState = MutableHashMap.get(entityResumptionState, address)
926 if (Option.isSome(resumptionState)) {
927 resumptionState.value.unprocessed.add(lastReceivedMessage.envelope.requestId)
928 return Effect.void
929 }
930 MutableHashMap.set(entityResumptionState, address, {
931 unprocessed: new Set([lastReceivedMessage.envelope.requestId]),
932 interrupts: new Map()
933 })
934 return resumeEntityFromStorageImpl(address)
935 }
936 const resumeEntityFromStorageImpl = Effect.fnUntraced(
937 function*(address: EntityAddress) {
938 const state = entityManagers.get(address.entityType)
939 if (!state) {
940 MutableHashMap.remove(entityResumptionState, address)
941 return
942 }
943
944 const resumptionState = Option.getOrThrow(MutableHashMap.get(entityResumptionState, address))
945 let done = false
946
947 while (!done) { // oxlint-disable-line no-unmodified-loop-condition
948 // if the shard is no longer assigned to this runner, we stop
949 if (!MutableHashSet.has(acquiredShards, address.shardId)) {
950 return
951 }
952
953 // take a batch of unprocessed messages ids
954 const messageIds = Arr.empty<Snowflake.Snowflake>()
955 for (const id of resumptionState.unprocessed) {
956 if (messageIds.length === 1024) break
957 messageIds.push(id)
958 }
959
960 const messages = yield* storage.unprocessedMessagesById(messageIds)
961
962 // this should not happen, but we handle it just in case
963 if (messages.length === 0) {
964 yield* Effect.sleep(config.entityMessagePollInterval)
965 continue
966 }
967
968 let index = 0
969
970 const sendWithRetry: Effect.Effect<
971 void,
972 EntityNotAssignedToRunner
973 > = Effect.catchTags(
974 Effect.suspend(() => {
975 if (!MutableHashSet.has(acquiredShards, address.shardId)) {
976 return Effect.fail(new EntityNotAssignedToRunner({ address }))
977 }
978
979 const message = messages[index]
980 // check if this is a request that was interrupted
981 const interrupt = message._tag === "IncomingRequest" &&
982 resumptionState.interrupts.get(message.envelope.requestId)
983 return interrupt ?
984 Effect.flatMap(state.manager.send(message), () => {
985 resumptionState.interrupts.delete(message.envelope.requestId)
986 return state.manager.send(interrupt)
987 }) :
988 state.manager.send(message)
989 }),
990 {
991 MailboxFull: () => Effect.delay(sendWithRetry, config.sendRetryInterval),
992 AlreadyProcessingMessage: () => Effect.void
993 }
994 )
995
996 yield* Effect.whileLoop({
997 while: () => index < messages.length,
998 body: constant(sendWithRetry),
999 step: () => index++
1000 })
1001
1002 for (const id of messageIds) {
1003 resumptionState.unprocessed.delete(id)
1004 }
1005 if (resumptionState.unprocessed.size > 0) continue
1006
1007 // if we have caught up to the main storage loop, we let it take over
1008 yield* withStorageReadLock(Effect.sync(() => {
1009 if (resumptionState.unprocessed.size === 0) {
1010 MutableHashMap.remove(entityResumptionState, address)
1011 done = true
1012 }
1013 }))
1014 }
1015 },
1016 Effect.retry({
1017 while: (e) => e._tag === "PersistenceError",
1018 schedule: Schedule.spaced(config.entityMessagePollInterval)
1019 }),
1020 Effect.catchCause((cause) => Effect.logDebug("Could not resume unprocessed messages", cause)),
1021 (effect, address) =>
1022 Effect.annotateLogs(effect, {
1023 module: "effect/cluster/Sharding",
1024 fiber: "Resuming unprocessed messages",
1025 runner: selfAddress,
1026 entity: address
1027 }),
1028 (effect, address) =>
1029 Effect.ensuring(
1030 effect,
1031 Effect.sync(() => MutableHashMap.remove(entityResumptionState, address))
1032 ),
1033 Effect.forkIn(shardingScope)
1034 )
1035 }
1036
1037 // --- Sending messages ---
1038
1039 const sendLocal = <M extends Message.Outgoing<any> | Message.Incoming<any>>(message: M) =>
1040 Effect.suspend(function loop(): Effect.Effect<
1041 void,
1042 | EntityNotAssignedToRunner
1043 | MailboxFull
1044 | AlreadyProcessingMessage
1045 | (M extends Message.Incoming<any> ? never : PersistenceError)
1046 > {
1047 const address = message.envelope.address
1048 if (!isEntityOnLocalShards(address)) {
1049 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1050 }
1051 const state = entityManagers.get(address.entityType)
1052 if (!state) {
1053 return Effect.flatMap(waitForEntityManager(address.entityType), loop)
1054 } else if (state.status === "closed" || (state.status === "closing" && message._tag === "IncomingRequest")) {
1055 // if we are shutting down, we don't accept new requests
1056 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1057 }
1058
1059 return message._tag === "IncomingRequest" || message._tag === "IncomingEnvelope" ?
1060 state.manager.send(message) :
1061 runnersService.sendLocal({
1062 message,
1063 send: state.manager.sendLocal,
1064 simulateRemoteSerialization: config.simulateRemoteSerialization
1065 }) as any
1066 })
1067
1068 type PendingNotification = {
1069 resume: (_: Effect.Effect<void, EntityNotAssignedToRunner>) => void
1070 readonly message: Message.IncomingRequest<any>
1071 }
1072 const pendingNotifications = new Map<Snowflake.Snowflake, PendingNotification>()
1073 const notifyLocal = <M extends Message.Outgoing<any> | Message.Incoming<any>>(
1074 message: M,
1075 discard: boolean,
1076 options?: {
1077 readonly waitUntilRead?: boolean | undefined
1078 }
1079 ) =>
1080 Effect.suspend(function loop(): Effect.Effect<
1081 void,
1082 | EntityNotAssignedToRunner
1083 | AlreadyProcessingMessage
1084 | (M extends Message.Incoming<any> ? never : PersistenceError)
1085 > {
1086 const address = message.envelope.address
1087 const state = entityManagers.get(address.entityType)
1088 if (!state) {
1089 return Effect.flatMap(waitForEntityManager(address.entityType), loop)
1090 } else if (state.status === "closed" || !isEntityOnLocalShards(address)) {
1091 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1092 }
1093
1094 const isLocal = isEntityOnLocalShards(address)
1095 const notify = storageEnabled
1096 ? openStorageReadLatch
1097 : () => Effect.die("Sharding.notifyLocal: storage is disabled")
1098
1099 if (message._tag === "IncomingRequest" || message._tag === "IncomingEnvelope") {
1100 if (!isLocal) {
1101 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1102 } else if (
1103 message._tag === "IncomingRequest" && state.manager.isProcessingFor(message, { excludeReplies: true })
1104 ) {
1105 return Effect.fail(new AlreadyProcessingMessage({ address, envelopeId: message.envelope.requestId }))
1106 } else if (message._tag === "IncomingRequest" && options?.waitUntilRead) {
1107 if (!storageEnabled) return notify()
1108 return Effect.callback<void, EntityNotAssignedToRunner>((resume) => {
1109 let entry = pendingNotifications.get(message.envelope.requestId)
1110 if (entry) {
1111 const prevResume = entry.resume
1112 entry.resume = (effect) => {
1113 prevResume(effect)
1114 resume(effect)
1115 }
1116 return
1117 }
1118 entry = { resume, message }
1119 pendingNotifications.set(message.envelope.requestId, entry)
1120 storageReadLatch.openUnsafe()
1121 })
1122 }
1123 return notify()
1124 }
1125
1126 return runnersService.notifyLocal({ message, notify, discard, storageOnly: !isLocal }) as any
1127 })
1128
1129 function sendOutgoing(
1130 message: Message.Outgoing<any>,
1131 discard: boolean,
1132 retries?: number
1133 ): Effect.Effect<
1134 void,
1135 MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
1136 > {
1137 const isPersisted = Context.get(
1138 message._tag === "OutgoingRequest" ? message.annotations : message.rpc.annotations,
1139 Persisted
1140 )
1141 const shouldFail = !discard &&
1142 (message._tag === "OutgoingRequest" || message.envelope._tag === "AckChunk")
1143 const abandon = (error: EntityNotAssignedToRunner) => {
1144 if (!isPersisted) {
1145 return shouldFail
1146 ? Effect.fail(error)
1147 : Effect.logDebug("Abandoning outgoing message during shutdown", message.envelope.address)
1148 }
1149 const persist = message._tag === "OutgoingRequest"
1150 ? storage.saveRequest(message)
1151 : storage.saveEnvelope(message)
1152 return Effect.catchTag(persist, "MalformedMessage", Effect.die).pipe(
1153 Effect.andThen(
1154 shouldFail
1155 // The message is durable, so a transient routing state is never an
1156 // error for the caller. The runner is tearing down and can no
1157 // longer observe the reply, so interrupt the caller instead: the
1158 // request will be served under the next owner.
1159 ? message._tag === "OutgoingRequest"
1160 ? ClusterAbandon.interrupt
1161 : Effect.fail(error)
1162 : Effect.logWarning("Persisting outgoing message abandoned during shutdown", message.envelope.address)
1163 )
1164 )
1165 }
1166 return Effect.catchFilter(
1167 Effect.suspend(() => {
1168 const address = message.envelope.address
1169 if (isPersisted && !storageEnabled) {
1170 return Effect.die("Sharding.sendOutgoing: Persisted messages require MessageStorage")
1171 }
1172 if (shouldFail && MutableRef.get(isShutdown)) {
1173 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1174 }
1175 const maybeRunner = MutableHashMap.get(shardAssignments, address.shardId)
1176 const runnerIsLocal = Option.isSome(maybeRunner) && isLocalRunner(maybeRunner.value)
1177 if (isPersisted) {
1178 return runnerIsLocal
1179 ? notifyLocal(message, discard)
1180 : runnersService.notify({ address: maybeRunner, message, discard })
1181 } else if (Option.isNone(maybeRunner)) {
1182 return Effect.fail(new EntityNotAssignedToRunner({ address }))
1183 }
1184 return runnerIsLocal
1185 ? sendLocal(message)
1186 : discard
1187 ? runnersService.notify({ address: maybeRunner, message, discard })
1188 : runnersService.send({ address: maybeRunner.value, message })
1189 }),
1190 (error) =>
1191 error._tag === "EntityNotAssignedToRunner" || error._tag === "RunnerUnavailable"
1192 ? Result.succeed(error)
1193 : Result.fail(error),
1194 (error) => {
1195 // Abandon the message during teardown: retrying would loop forever once the runner is shutting down
1196 if (error._tag === "EntityNotAssignedToRunner") {
1197 const targetManager = entityManagers.get(message.envelope.address.entityType)
1198 const cannotRecover = MutableRef.get(isShutdown) ||
1199 (targetManager !== undefined && targetManager.status !== "alive")
1200 if (cannotRecover) {
1201 return abandon(error)
1202 }
1203 }
1204 if (retries === 0) {
1205 return Effect.die(error)
1206 }
1207 return Effect.delay(sendOutgoing(message, discard, retries && retries - 1), config.sendRetryInterval)
1208 }
1209 )
1210 }
1211
1212 const reset: Sharding["Service"]["reset"] = (requestId, options) =>
1213 Effect.matchCause(storage.clearReplies(requestId, options), {
1214 onSuccess: () => true,
1215 onFailure: () => false
1216 })
1217
1218 // --- RunnerStorage sync ---
1219 //
1220 // This is responsible for syncing the local view of runners and shard
1221 // assignments with RunnerStorage.
1222 //
1223 // It should be shutdown after the clients, so that they can still get correct
1224 // shard assignments for outgoing messages (they could still be in use by
1225 // entities that are shutting down).
1226
1227 const selfRunner = initialRunnerAddress ?
1228 new Runner({
1229 address: initialRunnerAddress,
1230 groups: Array.from(shardGroups.assigned),
1231 weight: config.runnerShardWeight
1232 }) :
1233 undefined
1234
1235 let allRunners = MutableHashMap.empty<Runner, boolean>()
1236 let healthyRunnerCount = 0
1237
1238 // update metrics
1239 if (selfRunner) {
1240 ClusterMetrics.runners.updateUnsafe(BigInt(1), Context.empty())
1241 ClusterMetrics.runnersHealthy.updateUnsafe(BigInt(1), Context.empty())
1242 }
1243
1244 yield* Effect.gen(function*() {
1245 const hashRings = new Map<string, HashRing.HashRing<RunnerAddress>>()
1246 let nextRunners = MutableHashMap.empty<Runner, boolean>()
1247 const healthyRunners = MutableHashSet.empty<Runner>()
1248 const withTimeout = Effect.timeout(Duration.seconds(5))
1249
1250 while (true) {
1251 // Ensure the current runner is registered
1252 if (selfRunner && !isShutdown.current && !MutableHashMap.has(allRunners, selfRunner)) {
1253 yield* Effect.logDebug("Registering runner", selfRunner)
1254 const machineId = yield* withTimeout(runnerStorage.register(selfRunner, true))
1255 yield* snowflakeGen.setMachineId(machineId)
1256 }
1257
1258 const runners = yield* withTimeout(runnerStorage.getRunners)
1259 let changed = false
1260 for (let i = 0; i < runners.length; i++) {
1261 const [runner, healthy] = runners[i]
1262 MutableHashMap.set(nextRunners, runner, healthy)
1263 const wasHealthy = MutableHashSet.has(healthyRunners, runner)
1264 if (!healthy || wasHealthy) {
1265 if (healthy === wasHealthy || !wasHealthy) {
1266 // no change
1267 MutableHashMap.remove(allRunners, runner)
1268 }
1269 continue
1270 }
1271 changed = true
1272 MutableHashSet.add(healthyRunners, runner)
1273 MutableHashMap.remove(allRunners, runner)
1274 for (let j = 0; j < runner.groups.length; j++) {
1275 const group = runner.groups[j]
1276 let ring = hashRings.get(group)
1277 if (!ring) {
1278 ring = HashRing.make()
1279 hashRings.set(group, ring)
1280 }
1281 HashRing.add(ring, runner.address, { weight: runner.weight })
1282 }
1283 }
1284
1285 // Remove runners that are no longer present or healthy
1286 MutableHashMap.forEach(allRunners, (_, runner) => {
1287 changed = true
1288 MutableHashMap.remove(allRunners, runner)
1289 MutableHashSet.remove(healthyRunners, runner)
1290 runFork(runnersService.onRunnerUnavailable(runner.address))
1291 for (let i = 0; i < runner.groups.length; i++) {
1292 HashRing.remove(hashRings.get(runner.groups[i])!, runner.address)
1293 }
1294 })
1295
1296 // swap allRunners and nextRunners
1297 const prevRunners = allRunners
1298 allRunners = nextRunners
1299 nextRunners = prevRunners
1300 healthyRunnerCount = MutableHashSet.size(healthyRunners)
1301
1302 // Ensure the current runner is registered
1303 if (selfRunner && !isShutdown.current && !MutableHashMap.has(allRunners, selfRunner)) {
1304 continue
1305 }
1306
1307 // Recompute shard assignments if the set of healthy runners has changed.
1308 if (changed) {
1309 MutableHashSet.clear(selfShards)
1310 hashRings.forEach((ring, group) => {
1311 const newAssignments = HashRing.getShards(ring, config.shardsPerGroup)
1312 for (let i = 0; i < config.shardsPerGroup; i++) {
1313 const shard = makeShardId(group, i + 1)
1314 if (newAssignments) {
1315 const runner = newAssignments[i]
1316 MutableHashMap.set(shardAssignments, shard, runner)
1317 if (shardLocksHealthyLatch.isOpen() && isLocalRunner(runner)) {
1318 MutableHashSet.add(selfShards, shard)
1319 }
1320 } else {
1321 MutableHashMap.remove(shardAssignments, shard)
1322 }
1323 }
1324 })
1325 yield* Effect.logDebug("New shard assignments", selfShards)
1326 activeShardsLatch.openUnsafe()
1327
1328 // update metrics
1329 if (selfRunner) {
1330 ClusterMetrics.runnersHealthy.updateUnsafe(
1331 BigInt(MutableHashSet.has(healthyRunners, selfRunner) ? 1 : 0),
1332 Context.empty()
1333 )
1334 }
1335 }
1336
1337 if (selfRunner && MutableHashSet.size(healthyRunners) === 0) {
1338 yield* Effect.logWarning("No healthy runners available")
1339 // to prevent a deadlock, we will mark the current node as healthy to
1340 // start the health check singleton again
1341 yield* withTimeout(runnerStorage.setRunnerHealth(selfRunner.address, true))
1342 }
1343
1344 yield* Effect.sleep(config.refreshAssignmentsInterval)
1345 }
1346 }).pipe(
1347 Effect.catchCause((cause) => Effect.logDebug(cause)),
1348 Effect.repeat(Schedule.spaced(1000)),
1349 Effect.annotateLogs({
1350 module: "effect/cluster/Sharding",
1351 fiber: "RunnerStorage sync",
1352 runner: initialRunnerAddress
1353 }),
1354 Effect.forkIn(shardingScope)
1355 )
1356
1357 // --- Clients ---
1358
1359 type ClientRequestEntry = {
1360 readonly rpc: Rpc.AnyWithProps
1361 readonly context: Context.Context<never>
1362 readonly message: Message.OutgoingRequest<any>
1363 lastChunkId?: Snowflake.Snowflake
1364 }
1365 const clientRequests = new Map<Snowflake.Snowflake, ClientRequestEntry>()
1366
1367 const clients: ResourceMap<
1368 Entity<any, any>,
1369 (entityId: string) => RpcClient.RpcClient<
1370 any,
1371 MailboxFull | AlreadyProcessingMessage | EntityNotAssignedToRunner
1372 >,
1373 never
1374 > = yield* ResourceMap.make(
1375 Effect.fnUntraced(function*(entity: Entity<string, any>) {
1376 const clientScope = yield* Effect.scope
1377 yield* Scope.addFinalizer(
1378 clientScope,
1379 Effect.sync(() => {
1380 ActiveTeardown.releaseEntityType(entity.type)
1381 })
1382 )
1383 const client = yield* RpcClient.makeNoSerialization(entity.protocol, {
1384 spanPrefix: `${entity.type}.client`,
1385 disableTracing: !Context.get(entity.protocol.annotations, ClusterSchema.ClientTracingEnabled),
1386 supportsAck: true,
1387 generateRequestId: () => snowflakeGen.nextUnsafe() as any,
1388 flatten: true,
1389 onFromClient(options): Effect.Effect<
1390 void,
1391 MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
1392 > {
1393 const address = Context.getUnsafe(options.context, ClientAddressTag)
1394 switch (options.message._tag) {
1395 case "Request": {
1396 const fiber = Fiber.getCurrent()!
1397 const id = Snowflake.Snowflake(options.message.id)
1398 const rpc = entity.protocol.requests.get(options.message.tag)!
1399 let respond: (reply: Reply.Reply<any>) => Effect.Effect<void>
1400 const envelope = Envelope.makeRequest<any>({
1401 requestId: id,
1402 address,
1403 tag: options.message.tag,
1404 payload: options.message.payload,
1405 headers: options.message.headers,
1406 traceId: options.message.traceId,
1407 spanId: options.message.spanId,
1408 sampled: options.message.sampled
1409 })
1410 const message = new Message.OutgoingRequest({
1411 envelope,
1412 lastReceivedReply: Option.none(),
1413 rpc,
1414 context: fiber.context as Context.Context<any>,
1415 respond: (reply) => respond(reply),
1416 annotations: Context.get(rpc.annotations, ClusterSchema.Dynamic)(
1417 rpc.annotations,
1418 envelope as any
1419 )
1420 })
1421 if (!options.discard) {
1422 const entry: ClientRequestEntry = {
1423 rpc: rpc as any,
1424 context: fiber.context,
1425 message
1426 }
1427 clientRequests.set(id, entry)
1428 respond = makeClientRespond(entry, client.write)
1429 } else {
1430 respond = clientRespondDiscard
1431 }
1432 return sendOutgoing(message, options.discard)
1433 }
1434 case "Ack": {
1435 const requestId = Snowflake.Snowflake(options.message.requestId)
1436 const entry = clientRequests.get(requestId)
1437 if (!entry) return Effect.void
1438 return sendOutgoing(
1439 new Message.OutgoingEnvelope({
1440 envelope: new Envelope.AckChunk({
1441 id: snowflakeGen.nextUnsafe(),
1442 address,
1443 requestId,
1444 replyId: entry.lastChunkId!
1445 }),
1446 rpc: entry.rpc
1447 }),
1448 false
1449 )
1450 }
1451 case "Interrupt": {
1452 const requestId = Snowflake.Snowflake(options.message.requestId)
1453 const entry = clientRequests.get(requestId)!
1454 if (!entry) return Effect.void
1455 clientRequests.delete(requestId)
1456 if (ClusterSchema.isUninterruptibleForClient(entry.message.annotations)) {
1457 return Effect.void
1458 }
1459 // for durable messages, we ignore interrupts on shutdown or as a
1460 // result of a shard being resassigned
1461 const caller = Context.getOption(entry.context, CurrentAddress).valueOrUndefined
1462 const singletonShard = Context.getOption(entry.context, SingletonShardTag).valueOrUndefined
1463 const isTransientInterrupt = MutableRef.get(isShutdown) ||
1464 (caller !== undefined && ActiveTeardown.isActive(caller)) ||
1465 (singletonShard !== undefined && ActiveTeardown.isShardActive(singletonShard))
1466 if (isTransientInterrupt && Context.get(entry.message.annotations, Persisted)) {
1467 return Effect.void
1468 }
1469 return Effect.ignore(sendOutgoing(
1470 new Message.OutgoingEnvelope({
1471 envelope: new Envelope.Interrupt({
1472 id: snowflakeGen.nextUnsafe(),
1473 address,
1474 requestId
1475 }),
1476 rpc: entry.rpc
1477 }),
1478 false,
1479 3
1480 ))
1481 }
1482 }
1483 return Effect.void
1484 }
1485 })
1486
1487 yield* Scope.addFinalizer(
1488 clientScope,
1489 Effect.sync(() => {
1490 ActiveTeardown.acquireEntityType(entity.type)
1491 })
1492 )
1493
1494 return (entityId: string) => {
1495 const id = makeEntityId(entityId)
1496 const address = ClientAddressTag.context(makeEntityAddress({
1497 shardId: getShardId(id, entity.getShardGroup(entityId as EntityId)),
1498 entityId: id,
1499 entityType: entity.type
1500 }))
1501 const clientFn = function(tag: string, payload: any, options?: {
1502 readonly context?: Context.Context<never>
1503 }) {
1504 const context = options?.context ? Context.merge(options.context, address) : address
1505 return client.client(tag, payload, {
1506 ...options,
1507 context
1508 })
1509 }
1510 const proxyClient: any = {}
1511 return new Proxy(proxyClient, {
1512 has(_, p) {
1513 return entity.protocol.requests.has(p as string)
1514 },
1515 get(target, p) {
1516 if (Object.hasOwn(target, p)) {
1517 return target[p]
1518 } else if (!entity.protocol.requests.has(p as string)) {
1519 return undefined
1520 }
1521 const method = (payload: any, options?: {}) => clientFn(p as string, payload, options)
1522 InternalRecord.assignProperty(target, p, method)
1523 return method
1524 }
1525 })
1526 }
1527 }),
1528 { referential: true }
1529 )
1530
1531 const makeClient = <Type extends string, Rpcs extends Rpc.Any>(entity: Entity<Type, Rpcs>): Effect.Effect<
1532 (
1533 entityId: string
1534 ) => RpcClient.RpcClient.From<Rpcs, MailboxFull | AlreadyProcessingMessage | EntityNotAssignedToRunner>
1535 > => clients.get(entity) as any
1536
1537 const clientRespondDiscard = (_reply: Reply.Reply<any>) => Effect.void
1538
1539 const makeClientRespond = (
1540 entry: ClientRequestEntry,
1541 write: (reply: FromServer<any>) => Effect.Effect<void>
1542 ) =>
1543 (reply: Reply.Reply<any>) => {
1544 switch (reply._tag) {
1545 case "Chunk": {
1546 entry.lastChunkId = reply.id
1547 return write({
1548 _tag: "Chunk",
1549 clientId: 0,
1550 requestId: reply.requestId as any,
1551 values: reply.values
1552 })
1553 }
1554 case "WithExit": {
1555 clientRequests.delete(reply.requestId)
1556 return write({
1557 _tag: "Exit",
1558 clientId: 0,
1559 requestId: reply.requestId as any,
1560 exit: reply.exit
1561 })
1562 }
1563 }
1564 }
1565
1566 // --- Singletons ---
1567
1568 const singletons = new Map<ShardId, MutableHashMap.MutableHashMap<SingletonAddress, Effect.Effect<void>>>()
1569 const singletonFibers = yield* FiberMap.make<SingletonAddress>()
1570 const withSingletonLock = Semaphore.makeUnsafe(1).withPermits(1)
1571
1572 const registerSingleton: Sharding["Service"]["registerSingleton"] = Effect.fnUntraced(
1573 function*(name, run, options) {
1574 const shardGroup = options?.shardGroup ?? "default"
1575 const address = new SingletonAddress({
1576 shardId: getShardId(makeEntityId(name), shardGroup),
1577 name
1578 })
1579
1580 let map = singletons.get(address.shardId)
1581 if (!map) {
1582 map = MutableHashMap.empty()
1583 singletons.set(address.shardId, map)
1584 }
1585 if (MutableHashMap.has(map, address)) {
1586 return yield* Effect.die(`Singleton '${name}' is already registered`)
1587 }
1588
1589 const services = yield* Effect.context<never>()
1590 const wrappedRun = run.pipe(
1591 Effect.andThen(Effect.never),
1592 Effect.scoped,
1593 Effect.provideService(CurrentLogAnnotations, {}),
1594 Effect.provideContext(Context.add(services, SingletonShardTag, address.shardId)),
1595 Effect.orDie,
1596 Effect.interruptible
1597 ) as Effect.Effect<never>
1598 MutableHashMap.set(map, address, wrappedRun)
1599
1600 yield* PubSub.publish(events, SingletonRegistered({ address }))
1601
1602 // start if we are on the right shard
1603 if (MutableHashSet.has(acquiredShards, address.shardId)) {
1604 yield* Effect.logDebug("Starting singleton", address)
1605 yield* FiberMap.run(singletonFibers, address, wrappedRun)
1606 }
1607
1608 yield* Effect.addFinalizer(() => {
1609 const map = singletons.get(address.shardId)!
1610 MutableHashMap.remove(map, address)
1611 return FiberMap.remove(singletonFibers, address)
1612 })
1613 },
1614 withSingletonLock
1615 )
1616
1617 const syncSingletons = withSingletonLock(Effect.gen(function*() {
1618 for (const [shardId, map] of singletons) {
1619 for (const [address, run] of map) {
1620 const running = FiberMap.hasUnsafe(singletonFibers, address)
1621 const shouldBeRunning = MutableHashSet.has(acquiredShards, shardId)
1622 if (running && !shouldBeRunning) {
1623 yield* Effect.logDebug("Stopping singleton", address)
1624 yield* ActiveTeardown.aroundShard(
1625 address.shardId,
1626 FiberMap.remove(singletonFibers, address)
1627 )
1628 } else if (!running && shouldBeRunning) {
1629 yield* Effect.logDebug("Starting singleton", address)
1630 yield* FiberMap.run(singletonFibers, address, run)
1631 }
1632 }
1633 }
1634 ClusterMetrics.singletons.updateUnsafe(
1635 BigInt(yield* FiberMap.size(singletonFibers)),
1636 Context.empty()
1637 )
1638 }))
1639
1640 // --- Entities ---
1641
1642 const reaper = yield* EntityReaper
1643 const entityManagerLatches = new Map<string, Latch.Latch>()
1644
1645 const registerEntity: Sharding["Service"]["registerEntity"] = Effect.fnUntraced(
1646 function*(entity, build, options) {
1647 const runnerAddress = getRunnerAddress()
1648 if (!runnerAddress || entityManagers.has(entity.type)) return
1649 const scope = yield* Effect.scope
1650 const registrationContext = yield* Effect.context<never>()
1651 yield* Scope.addFinalizer(
1652 scope,
1653 Effect.sync(() => {
1654 state.status = "closed"
1655 })
1656 )
1657 const manager = yield* EntityManager.make(entity, build, {
1658 ...options,
1659 storage,
1660 runnerAddress,
1661 residency,
1662 sharding
1663 }).pipe(
1664 Effect.provideContext(services.pipe(
1665 // Registration overrides construction fallbacks, except for runner-owned services.
1666 Context.merge(registrationContext),
1667 Context.add(ShardingConfig, config),
1668 Context.add(Clock, clock),
1669 Context.add(EntityReaper, reaper),
1670 Context.add(Scope.Scope, scope),
1671 Context.add(Snowflake.Generator, snowflakeGen)
1672 ))
1673 ) as Effect.Effect<EntityManager.EntityManager>
1674 const state: EntityManagerState = {
1675 entity,
1676 status: "alive",
1677 manager
1678 }
1679 yield* Scope.addFinalizer(
1680 scope,
1681 Effect.suspend(() => {
1682 state.status = "closing"
1683 return ActiveTeardown.aroundEntityType(
1684 entity.type,
1685 config.preemptiveShutdown ? shutdown() : Effect.void
1686 )
1687 })
1688 )
1689
1690 // register entities while storage is idle
1691 // this ensures message order is preserved
1692 yield* withStorageReadLock(Effect.sync(() => {
1693 entityRegistrationStartMillis ??= clock.currentTimeMillisUnsafe()
1694 entityManagers.set(entity.type, state)
1695 if (entityManagerLatches.has(entity.type)) {
1696 entityManagerLatches.get(entity.type)!.openUnsafe()
1697 entityManagerLatches.delete(entity.type)
1698 }
1699 }))
1700
1701 yield* PubSub.publish(events, EntityRegistered({ entity }))
1702 }
1703 )
1704
1705 // Sleeps until the registration deadline shared with the storage read loop,
1706 // re-checking once in case registration started and moved the deadline.
1707 const awaitRegistrationDeadline = (entityType: string): Effect.Effect<never> =>
1708 Effect.suspend(() => {
1709 const remaining = entityRegistrationTimeRemaining()
1710 return remaining > 0
1711 ? Effect.flatMap(Effect.sleep(remaining), () => awaitRegistrationDeadline(entityType))
1712 : entityNotRegistered(entityType)
1713 })
1714
1715 const waitForEntityManager = (entityType: string) => {
1716 let latch = entityManagerLatches.get(entityType)
1717 if (!latch) {
1718 latch = Latch.makeUnsafe()
1719 entityManagerLatches.set(entityType, latch)
1720 }
1721 return Effect.raceFirst(latch.await, awaitRegistrationDeadline(entityType))
1722 }
1723
1724 // --- Runner health checks ---
1725
1726 if (selfRunner) {
1727 const checkRunner = ([runner, healthy]: [Runner, boolean]) =>
1728 Effect.flatMap(runnerHealth.isAlive(runner.address), (isAlive) => {
1729 if (healthy === isAlive) return Effect.void
1730 if (isAlive) {
1731 healthyRunnerCount++
1732 return Effect.logDebug(`Runner is healthy`, runner).pipe(
1733 Effect.andThen(runnerStorage.setRunnerHealth(runner.address, isAlive))
1734 )
1735 }
1736 if (healthyRunnerCount <= 1) {
1737 // never mark the last runner as unhealthy, to prevent a deadlock
1738 return Effect.void
1739 }
1740 healthyRunnerCount--
1741 return Effect.logDebug(`Runner is unhealthy`, runner).pipe(
1742 Effect.andThen(runnerStorage.setRunnerHealth(runner.address, isAlive))
1743 )
1744 })
1745
1746 yield* registerSingleton(
1747 "effect/cluster/Sharding/RunnerHealth",
1748 Effect.gen(function*() {
1749 while (true) {
1750 // Skip health checks if we are the only runner
1751 if (MutableHashMap.size(allRunners) > 1) {
1752 yield* Effect.forEach(allRunners, checkRunner, { discard: true, concurrency: 10 })
1753 }
1754 yield* Effect.sleep(config.runnerHealthCheckInterval)
1755 }
1756 }).pipe(
1757 Effect.catchCause((cause) => Effect.logDebug("Runner health check failed", cause)),
1758 Effect.forever,
1759 Effect.annotateLogs({
1760 module: "effect/cluster/Sharding",
1761 fiber: "Runner health check"
1762 })
1763 )
1764 )
1765 }
1766
1767 // --- Finalization ---
1768
1769 const shutdown = Effect.fnUntraced(function*(exit?: Exit.Exit<unknown, unknown>) {
1770 if (exit) {
1771 yield* Effect.logDebug("Shutting down", exit._tag === "Failure" ? exit.cause : {}).pipe(
1772 Effect.annotateLogs({
1773 module: "effect/cluster/Sharding"
1774 })
1775 )
1776 }
1777
1778 if (isShutdown.current) return
1779 MutableRef.set(isShutdown, true)
1780 if (selfRunner) {
1781 yield* Effect.ignore(runnerStorage.unregister(selfRunner.address))
1782 }
1783 })
1784
1785 yield* Scope.addFinalizerExit(shardingScope, shutdown)
1786
1787 const activeEntityCount = Effect.gen(function*() {
1788 let count = 0
1789 for (const state of entityManagers.values()) {
1790 count += yield* state.manager.activeEntityCount
1791 }
1792 return count
1793 })
1794
1795 const sharding = Sharding.of({
1796 getRegistrationEvents,
1797 getShardId,
1798 hasShardId(shardId: ShardId) {
1799 if (isShutdown.current) return false
1800 return MutableHashSet.has(acquiredShards, shardId)
1801 },
1802 getSnowflake: Effect.sync(() => snowflakeGen.nextUnsafe()),
1803 isShutdown: Effect.sync(() => MutableRef.get(isShutdown)),
1804 registerEntity,
1805 registerSingleton,
1806 makeClient,
1807 send: sendLocal,
1808 sendOutgoing: (message, discard) => sendOutgoing(message, discard),
1809 notify: (message, options) => notifyLocal(message, false, options),
1810 activeEntityCount,
1811 pollStorage: storageReadLatch.open,
1812 reset
1813 })
1814
1815 return sharding
1816})
1817
1818/**
1819 * Layer that constructs the `Sharding` service from sharding configuration,
1820 * runner communication, message storage, runner storage, runner health, the
1821 * snowflake generator, and the entity reaper.
1822 *
1823 * **When to use**
1824 *
1825 * Use when you need to assemble a cluster sharding runtime from explicit
1826 * sharding configuration, runner communication, message storage, runner
1827 * storage, and runner health layers.
1828 *
1829 * **Details**
1830 *
1831 * The layer provides the `Sharding` service and installs its own snowflake
1832 * generator and entity reaper. Callers still provide `ShardingConfig`,
1833 * `Runners`, `MessageStorage`, `RunnerStorage`, and `RunnerHealth`.
1834 *
1835 * **Gotchas**
1836 *
1837 * Persisted messages require a non-no-op `MessageStorage`; if this layer is
1838 * provided with `MessageStorage.layerNoop`, persisted sends defect.
1839 *
1840 * @see {@link Sharding} for the service provided by this layer
1841 *
1842 * @stability unstable
1843 * @category layers
1844 * @since 4.0.0
1845 */
1846export const layer: Layer.Layer<
1847 Sharding,
1848 never,
1849 ShardingConfig | Runners | MessageStorage.MessageStorage | RunnerStorage | RunnerHealth.RunnerHealth
1850> = Layer.effect(Sharding)(make).pipe(
1851 Layer.provide([Snowflake.layerGenerator, EntityReaper.layer])
1852)
1853
1854// Utilities
1855
1856const ClientAddressTag = Context.Service<EntityAddress>("effect/cluster/Sharding/ClientAddress")
1857const SingletonShardTag = Context.Service<ShardId>("effect/cluster/Sharding/SingletonShard")