@effect/cluster 0.37.2 → 0.38.1

package/src/SqlShardStorage.ts

@@ -66,21 +66,21 @@ export const make = Effect.fnUntraced(function*(options?: {
       sql`
         IF OBJECT_ID(N'${shardsTableSql}', N'U') IS NULL
         CREATE TABLE ${shardsTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255)
         )
       `,
     mysql: () =>
       sql`
         CREATE TABLE IF NOT EXISTS ${shardsTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255)
         )
       `,
     pg: () =>
       sql`
         CREATE TABLE IF NOT EXISTS ${shardsTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255)
         )
       `,
@@ -88,7 +88,7 @@ export const make = Effect.fnUntraced(function*(options?: {
       // sqlite
       sql`
         CREATE TABLE IF NOT EXISTS ${shardsTableSql} (
-          shard_id INTEGER PRIMARY KEY,
+          shard_id TEXT PRIMARY KEY,
           address TEXT
         )
       `
@@ -102,7 +102,7 @@ export const make = Effect.fnUntraced(function*(options?: {
       sql`
         IF OBJECT_ID(N'${locksTableSql}', N'U') IS NULL
         CREATE TABLE ${locksTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255) NOT NULL,
           acquired_at DATETIME NOT NULL
         )
@@ -110,7 +110,7 @@ export const make = Effect.fnUntraced(function*(options?: {
     mysql: () =>
       sql`
         CREATE TABLE IF NOT EXISTS ${locksTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255) NOT NULL,
           acquired_at DATETIME NOT NULL
         )
@@ -118,7 +118,7 @@ export const make = Effect.fnUntraced(function*(options?: {
     pg: () =>
       sql`
         CREATE TABLE IF NOT EXISTS ${locksTableSql} (
-          shard_id INT PRIMARY KEY,
+          shard_id VARCHAR(50) PRIMARY KEY,
           address VARCHAR(255) NOT NULL,
           acquired_at TIMESTAMP NOT NULL
         )
@@ -127,7 +127,7 @@ export const make = Effect.fnUntraced(function*(options?: {
       // sqlite
       sql`
         CREATE TABLE IF NOT EXISTS ${locksTableSql} (
-          shard_id INTEGER PRIMARY KEY,
+          shard_id TEXT PRIMARY KEY,
           address TEXT NOT NULL,
           acquired_at DATETIME NOT NULL
         )
@@ -143,7 +143,7 @@ export const make = Effect.fnUntraced(function*(options?: {
   const sqlNow = sql.literal(sqlNowString)
 
   const lockExpiresAt = sql.onDialectOrElse({
-    pg: () => sql`${sqlNow} - INTERVAL '5 seconds'`,
+    pg: () => sql`${sqlNow} - INTERVAL '10 seconds'`,
     mysql: () => sql`DATE_SUB(${sqlNow}, INTERVAL 5 SECOND)`,
     mssql: () => sql`DATEADD(SECOND, -5, ${sqlNow})`,
     orElse: () => sql`datetime(${sqlNow}, '-5 seconds')`
@@ -199,7 +199,7 @@ export const make = Effect.fnUntraced(function*(options?: {
     orElse: () => sql.literal("FOR UPDATE")
   })
 
-  return yield* ShardStorage.makeEncoded({
+  return ShardStorage.makeEncoded({
     getAssignments: sql`SELECT shard_id, address FROM ${shardsTableSql} ORDER BY shard_id`.values.pipe(
       PersistenceError.refail,
       withTracerDisabled
@@ -244,12 +244,12 @@ export const make = Effect.fnUntraced(function*(options?: {
       function*(address, shardIds) {
         const values = shardIds.map((shardId) => sql`(${shardId}, ${address}, ${sqlNow})`)
         yield* acquireLock(address, values)
-        const currentLocks = yield* sql<{ shard_id: number }>`
+        const currentLocks = yield* sql<{ shard_id: string }>`
           SELECT shard_id FROM ${sql(locksTable)}
-          WHERE address = ${address} AND ${sql.in("shard_id", shardIds)}
+          WHERE address = ${address} AND acquired_at >= ${lockExpiresAt}
           ${forUpdate}
-        `
-        return currentLocks.map((row) => row.shard_id)
+        `.values
+        return currentLocks.map((row) => row[0] as string)
       },
       sql.withTransaction,
       PersistenceError.refail,
@@ -264,7 +264,7 @@ export const make = Effect.fnUntraced(function*(options?: {
           sql`SELECT shard_id FROM ${locksTableSql} WHERE address = ${address} AND acquired_at >= ${lockExpiresAt} ${forUpdate}`
             .values
         ),
-        Effect.map((rows) => rows.map((row) => Number(row[0]))),
+        Effect.map((rows) => rows.map((row) => row[0] as string)),
         PersistenceError.refail,
         withTracerDisabled
       ),

package/src/index.ts

@@ -1,3 +1,8 @@
+/**
+ * @since 1.0.0
+ */
+export * as ClusterCron from "./ClusterCron.js"
+
 /**
  * @since 1.0.0
  */

package/src/internal/entityManager.ts

@@ -7,6 +7,7 @@ import * as Context from "effect/Context"
 import * as Duration from "effect/Duration"
 import type { DurationInput } from "effect/Duration"
 import * as Effect from "effect/Effect"
+import * as Equal from "effect/Equal"
 import * as Exit from "effect/Exit"
 import * as FiberRef from "effect/FiberRef"
 import { identity } from "effect/Function"
@@ -400,7 +401,7 @@ export const make = Effect.fnUntraced(function*<
     Effect.suspend(function loop(): Effect.Effect<void> {
       const toInterrupt = new Set<EntityState>()
       for (const state of activeServers.values()) {
-        if (shardId === state.address.shardId) {
+        if (shardId[Equal.symbol](state.address.shardId)) {
           toInterrupt.add(state)
         }
       }

package/src/internal/shardManager.ts

@@ -14,7 +14,9 @@ import { ShardStorage } from "../ShardStorage.js"
 
 /** @internal */
 export class State {
-  static fromStorage = Effect.fnUntraced(function*(numberOfShards: number) {
+  static fromStorage = Effect.fnUntraced(function*(
+    shardsPerGroup: number
+  ) {
     const storage = yield* ShardStorage
     const runnerHealth = yield* RunnerHealth
 
@@ -38,11 +40,11 @@ export class State {
     }
 
     // Determine which shards remain unassigned to a runner
-    const assignedShards = new Map<ShardId, RunnerAddress>()
+    const assignedShards = MutableHashMap.empty<ShardId, RunnerAddress>()
     const invalidAssignments = Arr.empty<[ShardId, RunnerAddress]>()
     for (const [shard, address] of storedAssignments) {
       if (Option.isSome(address) && MutableHashMap.has(aliveRunners, address.value)) {
-        assignedShards.set(shard, address.value)
+        MutableHashMap.set(assignedShards, shard, address.value)
       } else if (Option.isSome(address)) {
         invalidAssignments.push([shard, address.value])
       }
@@ -56,29 +58,96 @@ export class State {
 
     // Construct the initial state
     const now = yield* Clock.currentTimeMillis
-    const runnerState = MutableHashMap.empty<RunnerAddress, RunnerWithMetadata>()
+    const allRunners = MutableHashMap.empty<RunnerAddress, RunnerWithMetadata>()
+    const runnerState = new Map<string, MutableHashMap.MutableHashMap<RunnerAddress, RunnerWithMetadata>>()
+    // for (const group of groups) {
+    //   runnerState.set(group, MutableHashMap.empty<RunnerAddress, RunnerWithMetadata>())
+    // }
     for (const [address, runner] of aliveRunners) {
-      MutableHashMap.set(runnerState, address, RunnerWithMetadata({ runner, registeredAt: now }))
+      const withMetadata = RunnerWithMetadata({ runner, registeredAt: now })
+      MutableHashMap.set(allRunners, address, withMetadata)
+      for (const group of runner.groups) {
+        let groupMap = runnerState.get(group)
+        if (!groupMap) {
+          groupMap = MutableHashMap.empty<RunnerAddress, RunnerWithMetadata>()
+          runnerState.set(group, groupMap)
+        }
+        MutableHashMap.set(groupMap, address, withMetadata)
+      }
     }
 
-    const shardState = new Map<ShardId, Option.Option<RunnerAddress>>()
-    for (let n = 1; n <= numberOfShards; n++) {
-      const shardId = ShardId.make(n)
-      shardState.set(shardId, Option.fromNullable(assignedShards.get(shardId)))
+    const shardState = new Map<string, Map<number, Option.Option<RunnerAddress>>>()
+    for (const group of runnerState.keys()) {
+      const groupMap = new Map<number, Option.Option<RunnerAddress>>()
+      shardState.set(group, groupMap)
+      for (let n = 1; n <= shardsPerGroup; n++) {
+        const shardId = new ShardId({ group, id: n })
+        groupMap.set(n, MutableHashMap.get(assignedShards, shardId))
+      }
     }
 
-    return new State(runnerState, shardState)
+    return new State(allRunners, runnerState, shardState, shardsPerGroup)
   })
 
   constructor(
-    readonly runners: MutableHashMap.MutableHashMap<RunnerAddress, RunnerWithMetadata>,
-    readonly shards: Map<ShardId, Option.Option<RunnerAddress>>
-  ) {}
+    readonly allRunners: MutableHashMap.MutableHashMap<RunnerAddress, RunnerWithMetadata>,
+    readonly runners: Map<string, MutableHashMap.MutableHashMap<RunnerAddress, RunnerWithMetadata>>,
+    readonly shards: Map<string, Map<number, Option.Option<RunnerAddress>>>,
+    readonly shardsPerGroup: number
+  ) {
+    this.assignments = MutableHashMap.empty<ShardId, Option.Option<RunnerAddress>>()
+    for (const [group, groupMap] of this.shards) {
+      for (const [id, address] of groupMap) {
+        MutableHashMap.set(this.assignments, new ShardId({ group, id }), address)
+      }
+    }
+  }
+
+  readonly assignments: MutableHashMap.MutableHashMap<ShardId, Option.Option<RunnerAddress>>
+
+  addGroup(group: string): void {
+    this.runners.set(group, MutableHashMap.empty<RunnerAddress, RunnerWithMetadata>())
+    const shardMap = new Map<number, Option.Option<RunnerAddress>>()
+    for (let n = 1; n <= this.shardsPerGroup; n++) {
+      shardMap.set(n, Option.none())
+      MutableHashMap.set(this.assignments, new ShardId({ group, id: n }), Option.none())
+    }
+    this.shards.set(group, shardMap)
+  }
+
+  addAssignments(
+    shards: Iterable<ShardId>,
+    address: Option.Option<RunnerAddress>
+  ) {
+    for (const shardId of shards) {
+      MutableHashMap.set(this.assignments, shardId, address)
+      this.shards.get(shardId.group)?.set(shardId.id, address)
+    }
+  }
+
+  addRunner(runner: Runner, registeredAt: number): void {
+    const withMetadata = RunnerWithMetadata({ runner, registeredAt })
+    MutableHashMap.set(this.allRunners, runner.address, withMetadata)
+    for (const group of runner.groups) {
+      if (!this.runners.has(group)) {
+        this.addGroup(group)
+      }
+      const groupMap = this.runners.get(group)!
+      MutableHashMap.set(groupMap, runner.address, withMetadata)
+    }
+  }
+
+  removeRunner(address: RunnerAddress): void {
+    MutableHashMap.remove(this.allRunners, address)
+    for (const groupMap of this.runners.values()) {
+      MutableHashMap.remove(groupMap, address)
+    }
+  }
 
   get maxVersion(): Option.Option<number> {
-    if (MutableHashMap.size(this.runners) === 0) return Option.none()
+    if (MutableHashMap.size(this.allRunners) === 0) return Option.none()
     let version: number | undefined = undefined
-    for (const [, meta] of this.runners) {
+    for (const [, meta] of this.allRunners) {
       if (version === undefined || meta.runner.version > version) {
         version = meta.runner.version
       }
@@ -93,31 +162,45 @@ export class State {
     )
   }
 
-  get shardsPerRunner(): MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>> {
-    const shards = MutableHashMap.empty<RunnerAddress, Set<ShardId>>()
+  shardsPerRunner(group: string): MutableHashMap.MutableHashMap<RunnerAddress, Set<number>> {
+    const groupRunners = this.runners.get(group)
+    const shards = MutableHashMap.empty<RunnerAddress, Set<number>>()
 
-    if (MutableHashMap.isEmpty(this.runners)) return shards
-    MutableHashMap.forEach(this.runners, (_, address) => {
+    if (!groupRunners || MutableHashMap.isEmpty(groupRunners)) return shards
+    MutableHashMap.forEach(groupRunners, (_, address) => {
       MutableHashMap.set(shards, address, new Set())
     })
 
-    for (const [shard, address] of this.shards) {
+    const assignments = this.shards.get(group)!
+    for (const [id, address] of assignments) {
       if (Option.isNone(address)) continue
       const shardIds = Option.getOrUndefined(MutableHashMap.get(shards, address.value))!
-      shardIds.add(shard)
+      shardIds.add(id)
     }
 
     return shards
   }
 
-  get averageShardsPerRunner(): number {
-    const runnerCount = MutableHashMap.size(this.runners)
-    return runnerCount > 0 ? this.shards.size / runnerCount : 0
+  averageShardsPerRunner(group: string): number {
+    const runnerCount = MutableHashMap.size(this.runners.get(group) ?? MutableHashMap.empty())
+    const shardGroup = this.shards.get(group) ?? new Map()
+    return runnerCount > 0 ? shardGroup.size / runnerCount : 0
+  }
+
+  get allUnassignedShards(): Array<ShardId> {
+    const unassigned: Array<ShardId> = []
+    for (const [shardId, address] of this.assignments) {
+      if (Option.isNone(address)) {
+        unassigned.push(shardId)
+      }
+    }
+    return unassigned
   }
 
-  get unassignedShards(): Array<ShardId> {
-    const shardIds: Array<ShardId> = []
-    for (const [shard, address] of this.shards) {
+  unassignedShards(group: string): Array<number> {
+    const shardIds: Array<number> = []
+    const assignments = this.shards.get(group)!
+    for (const [shard, address] of assignments) {
       if (Option.isNone(address)) {
         shardIds.push(shard)
       }
@@ -127,7 +210,7 @@ export class State {
 
   private get runnerVersions(): Array<number> {
     const runnerVersions: Array<number> = []
-    for (const [, meta] of this.runners) {
+    for (const [, meta] of this.allRunners) {
       runnerVersions.push(meta.runner.version)
     }
     return runnerVersions
@@ -143,36 +226,39 @@ export interface RunnerWithMetadata {
 export const RunnerWithMetadata = (runner: RunnerWithMetadata): RunnerWithMetadata => runner
 
 /** @internal */
-export function decideAssignmentsForUnassignedShards(state: State): readonly [
-  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
-  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
+export function decideAssignmentsForUnassignedShards(state: State, group: string): readonly [
+  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
+  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
   changes: MutableHashSet.MutableHashSet<RunnerAddress>
 ] {
-  return pickNewRunners(state.unassignedShards, state, true, 1)
+  return pickNewRunners(state.unassignedShards(group), state, group, true, 1)
 }
 
-const allocationOrder: Order.Order<[ShardId, number, number]> = Order.combine(
+const allocationOrder: Order.Order<[number, number, number]> = Order.combine(
   Order.mapInput(Order.number, ([, shards]) => shards),
   Order.mapInput(Order.number, ([, , registeredAt]) => registeredAt)
 )
 
 /** @internal */
-export function decideAssignmentsForUnbalancedShards(state: State, rate: number): readonly [
-  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
-  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
+export function decideAssignmentsForUnbalancedShards(state: State, group: string, rate: number): readonly [
+  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
+  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
   changes: MutableHashSet.MutableHashSet<RunnerAddress>
 ] {
-  const shardsPerRunner = state.shardsPerRunner
+  const shardsPerRunner = state.shardsPerRunner(group)
   const maxVersion = state.maxVersion
-  const extraShardsToAllocate = Arr.empty<[ShardId, shardsInverse: number, registeredAt: number]>()
+  const extraShardsToAllocate = Arr.empty<[number, shardsInverse: number, registeredAt: number]>()
+
+  const runnerGroup = state.runners.get(group)!
+  const shardsGroup = state.shards.get(group)!
 
   if (state.allRunnersHaveVersion(maxVersion)) {
-    const averageShardsPerRunner = state.averageShardsPerRunner
+    const averageShardsPerRunner = state.averageShardsPerRunner(group)
     MutableHashMap.forEach(shardsPerRunner, (shards) => {
       // Count how many extra shards there are compared to the average
       const extraShards = Math.max(0, shards.size - averageShardsPerRunner)
       for (const shard of takeRandom(shards, extraShards)) {
-        const maybeAddress = state.shards.get(shard) ?? Option.none()
+        const maybeAddress = shardsGroup.get(shard) ?? Option.none()
         if (Option.isNone(maybeAddress)) {
           extraShardsToAllocate.push([shard, Number.MIN_SAFE_INTEGER, Number.MIN_SAFE_INTEGER])
           continue
@@ -184,7 +270,7 @@ export function decideAssignmentsForUnbalancedShards(state: State, rate: number)
             onNone: () => Number.MIN_SAFE_INTEGER,
             onSome: (shards) => -shards.size
           }),
-          Option.match(MutableHashMap.get(state.runners, address), {
+          Option.match(MutableHashMap.get(runnerGroup, address), {
             onNone: () => Number.MIN_SAFE_INTEGER,
             onSome: (meta) => meta.registeredAt
           })
@@ -195,23 +281,24 @@ export function decideAssignmentsForUnbalancedShards(state: State, rate: number)
 
   const sortedShardsToRebalance = extraShardsToAllocate.sort(allocationOrder).map(([shard]) => shard)
 
-  return pickNewRunners(sortedShardsToRebalance, state, false, rate, shardsPerRunner, maxVersion)
+  return pickNewRunners(sortedShardsToRebalance, state, group, false, rate, shardsPerRunner, maxVersion)
 }
 
 function pickNewRunners(
-  shardsToRebalance: ReadonlyArray<ShardId>,
+  shardsToRebalance: ReadonlyArray<number>,
   state: State,
+  group: string,
   immediate: boolean,
   rate: number,
-  shardsPerRunner = state.shardsPerRunner,
+  shardsPerRunner = state.shardsPerRunner(group),
   maybeMaxVersion = state.maxVersion
 ): readonly [
-  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
-  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<ShardId>>,
+  assignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
+  unassignments: MutableHashMap.MutableHashMap<RunnerAddress, Set<number>>,
   changes: MutableHashSet.MutableHashSet<RunnerAddress>
 ] {
-  const addressAssignments = MutableHashMap.empty<RunnerAddress, Set<ShardId>>()
-  const unassignments = MutableHashMap.empty<RunnerAddress, Set<ShardId>>()
+  const addressAssignments = MutableHashMap.empty<RunnerAddress, Set<number>>()
+  const unassignments = MutableHashMap.empty<RunnerAddress, Set<number>>()
   const changes = MutableHashSet.empty<RunnerAddress>()
 
   if (Option.isNone(maybeMaxVersion)) {
@@ -219,14 +306,17 @@ function pickNewRunners(
   }
   const maxVersion = maybeMaxVersion.value
 
+  const runnerGroup = state.runners.get(group)!
+  const shardsGroup = state.shards.get(group)!
+
   for (const shardId of shardsToRebalance) {
     // Find the runner with the fewest assigned shards
     let candidate: RunnerAddress | undefined
-    let candidateShards: Set<ShardId> | undefined
+    let candidateShards: Set<number> | undefined
 
     for (const [address, shards] of shardsPerRunner) {
       // Keep only runners with the maximum version
-      const maybeRunnerMeta = MutableHashMap.get(state.runners, address)
+      const maybeRunnerMeta = MutableHashMap.get(runnerGroup, address)
       if (Option.isNone(maybeRunnerMeta)) continue
       const runnerMeta = maybeRunnerMeta.value
       if (runnerMeta.runner.version !== maxVersion) continue
@@ -238,7 +328,7 @@ function pickNewRunners(
       // occur immediately
       if (!immediate) {
         const assignmentCount = Option.getOrUndefined(MutableHashMap.get(addressAssignments, address))?.size ?? 0
-        if (assignmentCount >= state.shards.size * rate) continue
+        if (assignmentCount >= shardsGroup.size * rate) continue
       }
 
       if (candidate === undefined || shards.size < candidateShards!.size) {
@@ -249,7 +339,7 @@ function pickNewRunners(
     if (!candidate || !candidateShards) break
 
     // If the old runner is the same as the new runner, do nothing
-    const oldRunner = Option.getOrUndefined(state.shards.get(shardId) ?? Option.none())
+    const oldRunner = Option.getOrUndefined(shardsGroup.get(shardId) ?? Option.none())
     if (oldRunner && oldRunner.toString() === candidate.toString()) {
       continue
     }
@@ -316,3 +406,19 @@ function swap<A>(array: Array<A>, i: number, j: number): ReadonlyArray<A> {
   array[j] = tmp
   return array
 }
+
+/** @internal */
+export const addAllNested = <K, V>(
+  self: MutableHashMap.MutableHashMap<K, MutableHashSet.MutableHashSet<V>>,
+  key: K,
+  values: Iterable<V>
+) => {
+  const oset = MutableHashMap.get(self, key)
+  if (Option.isSome(oset)) {
+    for (const value of values) {
+      MutableHashSet.add(oset.value, value)
+    }
+  } else {
+    MutableHashMap.set(self, key, MutableHashSet.fromIterable(values))
+  }
+}