EFF-833 Backport cluster shard group fixes

.changeset/eff-833-cluster-shard-groups.md

@@ -0,0 +1,5 @@
+---
+"@effect/cluster": patch
+---
+
+Backport cluster shard group fixes: use available shard groups for SQL advisory lock numbering and route workflow durable clock/deferred messages through the owning workflow shard group.

packages/cluster/src/ClusterWorkflowEngine.ts

@@ -1,6 +1,7 @@
 /**
  * @since 1.0.0
  */
+import * as Headers from "@effect/platform/Headers"
 import * as Rpc from "@effect/rpc/Rpc"
 import * as RpcServer from "@effect/rpc/RpcServer"
 import { DurableDeferred } from "@effect/workflow"
@@ -33,6 +34,8 @@ import * as Entity from "./Entity.js"
 import { EntityAddress } from "./EntityAddress.js"
 import { EntityId } from "./EntityId.js"
 import { EntityType } from "./EntityType.js"
+import * as Envelope from "./Envelope.js"
+import * as Message from "./Message.js"
 import { MessageStorage } from "./MessageStorage.js"
 import type { WithExitEncoded } from "./Reply.js"
 import * as Reply from "./Reply.js"
@@ -130,24 +133,57 @@ export const make = Effect.gen(function*() {
     }),
     idleTimeToLive: "5 minutes"
   })
-  const clockClient = yield* ClockEntity.client
-
-  const requestIdFor = Effect.fnUntraced(function*(options: {
+  const entityAddressFor = (options: {
     readonly workflow: Workflow.Any
     readonly entityType: string
     readonly executionId: string
-    readonly tag: string
-    readonly id: string
-  }) {
+  }) => {
     const shardGroup = Context.get(options.workflow.annotations, ClusterSchema.ShardGroup)(
       options.executionId as EntityId
     )
     const entityId = EntityId.make(options.executionId)
-    const address = new EntityAddress({
+    return new EntityAddress({
       entityType: EntityType.make(options.entityType),
       entityId,
       shardId: sharding.getShardId(entityId, shardGroup)
     })
+  }
+
+  const sendDiscard = Effect.fnUntraced(function*(options: {
+    readonly rpc: Rpc.AnyWithProps
+    readonly address: EntityAddress
+    readonly payload: unknown
+  }) {
+    const payload = options.rpc.payloadSchema.make
+      ? options.rpc.payloadSchema.make(options.payload)
+      : options.payload
+    const envelope = Envelope.makeRequest<any>({
+      requestId: yield* sharding.getSnowflake,
+      address: options.address,
+      tag: options.rpc._tag as any,
+      payload,
+      headers: Headers.empty
+    })
+    yield* sharding.sendOutgoing(
+      new Message.OutgoingRequest({
+        envelope,
+        context: Context.empty() as Context.Context<any>,
+        lastReceivedReply: Option.none(),
+        rpc: options.rpc,
+        respond: () => Effect.void
+      }),
+      true
+    )
+  })
+
+  const requestIdFor = Effect.fnUntraced(function*(options: {
+    readonly workflow: Workflow.Any
+    readonly entityType: string
+    readonly executionId: string
+    readonly tag: string
+    readonly id: string
+  }) {
+    const address = entityAddressFor(options)
     return yield* storage.requestIdForPrimaryKey({ address, tag: options.tag, id: options.id })
   })
 
@@ -493,6 +529,21 @@ export const make = Effect.gen(function*() {
 
     deferredDone: Effect.fnUntraced(
       function*({ deferredName, executionId, exit, workflowName }) {
+        const workflow = workflows.get(workflowName)
+        if (workflow) {
+          return yield* Effect.orDie(sendDiscard({
+            rpc: DeferredRpc,
+            address: entityAddressFor({
+              workflow,
+              entityType: `Workflow/${workflowName}`,
+              executionId
+            }),
+            payload: {
+              name: deferredName,
+              exit
+            }
+          }))
+        }
         const client = yield* RcMap.get(clientsPartial, workflowName)
         return yield* Effect.orDie(
           client(executionId).deferred({
@@ -505,14 +556,21 @@ export const make = Effect.gen(function*() {
     ),
 
     scheduleClock(workflow, options) {
-      const client = clockClient(options.executionId)
       return DateTime.now.pipe(
         Effect.flatMap((now) =>
-          client.run({
-            name: options.clock.name,
-            workflowName: workflow.name,
-            wakeUp: DateTime.addDuration(now, options.clock.duration)
-          }, { discard: true })
+          sendDiscard({
+            rpc: ClockRpc,
+            address: entityAddressFor({
+              workflow,
+              entityType: ClockEntity.type,
+              executionId: options.executionId
+            }),
+            payload: {
+              name: options.clock.name,
+              workflowName: workflow.name,
+              wakeUp: DateTime.addDuration(now, options.clock.duration)
+            }
+          })
         ),
         Effect.orDie
       )
@@ -621,11 +679,11 @@ class ClockPayload extends Schema.Class<ClockPayload>(`Workflow/DurableClock/Run
   }
 }
 
-const ClockEntity = Entity.make("Workflow/-/DurableClock", [
-  Rpc.make("run", { payload: ClockPayload })
-    .annotate(ClusterSchema.Persisted, true)
-    .annotate(ClusterSchema.Uninterruptible, true)
-])
+const ClockRpc = Rpc.make("run", { payload: ClockPayload })
+  .annotate(ClusterSchema.Persisted, true)
+  .annotate(ClusterSchema.Uninterruptible, true)
+
+const ClockEntity = Entity.make("Workflow/-/DurableClock", [ClockRpc])
 
 const ClockEntityLayer = ClockEntity.toLayer(Effect.gen(function*() {
   const engine = yield* WorkflowEngine

packages/cluster/src/Sharding.ts

@@ -56,7 +56,7 @@ import { Runners } from "./Runners.js"
 import { RunnerStorage } from "./RunnerStorage.js"
 import type { ShardId } from "./ShardId.js"
 import { make as makeShardId } from "./ShardId.js"
-import { ShardingConfig } from "./ShardingConfig.js"
+import { shardGroupConfig, ShardingConfig } from "./ShardingConfig.js"
 import { EntityRegistered, type ShardingRegistrationEvent, SingletonRegistered } from "./ShardingRegistrationEvent.js"
 import { SingletonAddress } from "./SingletonAddress.js"
 import * as Snowflake from "./Snowflake.js"
@@ -199,6 +199,7 @@ interface EntityManagerState {
 
 const make = Effect.gen(function*() {
   const config = yield* ShardingConfig
+  const shardGroups = shardGroupConfig(config)
   const clock = yield* Effect.clock
 
   const runnersService = yield* Runners
@@ -871,7 +872,7 @@ const make = Effect.gen(function*() {
   const selfRunner = Option.isSome(config.runnerAddress) ?
     new Runner({
       address: config.runnerAddress.value,
-      groups: config.shardGroups,
+      groups: Array.from(shardGroups.assigned),
       weight: config.runnerShardWeight
     }) :
     undefined

packages/cluster/src/ShardingConfig.ts

@@ -43,11 +43,23 @@ export class ShardingConfig extends Context.Tag("@effect/cluster/ShardingConfig"
    * shards as a runner with a weight of `1`.
    */
   readonly runnerShardWeight: number
+  /**
+   * The shard groups available across all runners.
+   *
+   * Defaults to `["default"]`.
+   */
+  readonly availableShardGroups: ReadonlyArray<string>
   /**
    * The shard groups that are assigned to this runner.
    *
    * Defaults to `["default"]`.
    */
+  readonly assignedShardGroups: ReadonlyArray<string>
+  /**
+   * The shard groups that are assigned to this runner.
+   *
+   * @deprecated Use `assignedShardGroups` instead.
+   */
   readonly shardGroups: ReadonlyArray<string>
   /**
    * The number of shards to allocate per shard group.
@@ -134,6 +146,8 @@ export const defaults: ShardingConfig["Type"] = {
   runnerListenAddress: Option.none(),
   runnerShardWeight: 1,
   shardsPerGroup: 300,
+  availableShardGroups: ["default"],
+  assignedShardGroups: ["default"],
   shardGroups: ["default"],
   preemptiveShutdown: true,
   shardLockRefreshInterval: Duration.seconds(10),
@@ -156,7 +170,7 @@ export const defaults: ShardingConfig["Type"] = {
  * @category Layers
  */
 export const layer = (options?: Partial<ShardingConfig["Type"]>): Layer.Layer<ShardingConfig> =>
-  Layer.succeed(ShardingConfig, { ...defaults, ...options })
+  Layer.succeed(ShardingConfig, normalize({ ...defaults, ...options }, options))
 
 /**
  * @since 1.0.0
@@ -191,6 +205,14 @@ export const config: Config.Config<ShardingConfig["Type"]> = Config.all({
   runnerShardWeight: Config.integer("runnerShardWeight").pipe(
     Config.withDefault(defaults.runnerShardWeight)
   ),
+  availableShardGroups: Config.array(Config.string("availableShardGroups")).pipe(
+    Config.withDefault(["default"]),
+    Config.withDescription("The shard groups available across all runners.")
+  ),
+  assignedShardGroups: Config.array(Config.string("shardGroups")).pipe(
+    Config.withDefault(["default"]),
+    Config.withDescription("The shard groups that are assigned to this runner.")
+  ),
   shardGroups: Config.array(Config.string("shardGroups")).pipe(
     Config.withDefault(["default"]),
     Config.withDescription("The shard groups that are assigned to this runner.")
@@ -283,5 +305,36 @@ export const layerFromEnv = (options?: Partial<ShardingConfig["Type"]> | undefin
 > =>
   Layer.effect(
     ShardingConfig,
-    options ? Effect.map(configFromEnv, (config) => ({ ...config, ...options })) : configFromEnv
+    options ? Effect.map(configFromEnv, (config) => normalize({ ...config, ...options }, options)) : configFromEnv
   )
+
+function normalize(
+  config: ShardingConfig["Type"],
+  options: Partial<ShardingConfig["Type"]> | undefined
+): ShardingConfig["Type"] {
+  const assignedShardGroups = options?.assignedShardGroups ?? options?.shardGroups ?? config.assignedShardGroups
+  const availableShardGroups = options?.availableShardGroups ??
+    (options?.shardGroups && !options.assignedShardGroups ? assignedShardGroups : config.availableShardGroups)
+  return { ...config, availableShardGroups, assignedShardGroups, shardGroups: assignedShardGroups }
+}
+
+/**
+ * Normalizes the provided `ShardingConfig` to calculate the available and
+ * assigned shard groups.
+ *
+ * @since 1.0.0
+ * @category Shard groups
+ */
+export const shardGroupConfig = (config: ShardingConfig["Type"]): {
+  readonly available: ReadonlySet<string>
+  readonly assigned: ReadonlySet<string>
+} => {
+  const available = new Set(config.availableShardGroups.slice().sort())
+  const assigned = new Set<string>()
+  available.forEach((group) => {
+    if (config.assignedShardGroups.includes(group)) {
+      assigned.add(group)
+    }
+  })
+  return { available, assigned }
+}

packages/cluster/src/SqlRunnerStorage.ts

@@ -25,6 +25,8 @@ export const make = Effect.fnUntraced(function*(options: {
   readonly prefix?: string | undefined
 }) {
   const config = yield* ShardingConfig.ShardingConfig
+  const shardGroups = ShardingConfig.shardGroupConfig(config)
+  const availableShardGroups = Array.from(shardGroups.available)
   const disableAdvisoryLocks = config.shardLockDisableAdvisory
   const sql = (yield* SqlClient.SqlClient).withoutTransforms()
   const prefix = options?.prefix ?? "cluster"
@@ -397,8 +399,8 @@ export const make = Effect.fnUntraced(function*(options: {
 
   const lockNumbers = new Map<string, number>()
   const lockNumbersReverse = new Map<number, string>()
-  for (let i = 0; i < config.shardGroups.length; i++) {
-    const group = config.shardGroups[i]
+  for (let i = 0; i < availableShardGroups.length; i++) {
+    const group = availableShardGroups[i]
     const base = (i + 1) * 1000000
     for (let shard = 1; shard <= config.shardsPerGroup; shard++) {
       const shardId = ShardId.make(group, shard).toString()
@@ -413,8 +415,8 @@ export const make = Effect.fnUntraced(function*(options: {
   const lockNamesReverse = new Map<string, string>()
   {
     let index = 0
-    for (let i = 0; i < config.shardGroups.length; i++) {
-      const group = config.shardGroups[i]
+    for (let i = 0; i < availableShardGroups.length; i++) {
+      const group = availableShardGroups[i]
       for (let shard = 1; shard <= config.shardsPerGroup; shard++) {
         const shardId = ShardId.make(group, shard).toString()
         const lockName = `${prefix}.${shardId}`