From 1af896f31b354bd44966ae88d50a6b5ee73da0f9 Mon Sep 17 00:00:00 2001 From: Devin Ivy Date: Mon, 20 Jul 2026 00:24:46 -0700 Subject: [PATCH] fix: require initialState for stateful balancers, strengthen tests --- examples/worker-affinity/key-affinity.ts | 8 ++---- src/balancers.ts | 14 +++++++--- src/index.ts | 1 + src/runner.ts | 20 ++++++++------ test/balancers.test.ts | 6 ++-- test/fixtures/load-balancing.ts | 3 ++ test/load-balancing.test.ts | 35 ++++++++++++++++++++++-- 7 files changed, 63 insertions(+), 24 deletions(-) diff --git a/examples/worker-affinity/key-affinity.ts b/examples/worker-affinity/key-affinity.ts index 5311315..cb875d4 100644 --- a/examples/worker-affinity/key-affinity.ts +++ b/examples/worker-affinity/key-affinity.ts @@ -1,11 +1,7 @@ import { isTask, roundRobin } from '../../src/index.ts'; -import type { Balancer, Task, WorkerHandle } from '../../src/index.ts'; +import type { Balancer, RoundRobinState, Task, WorkerHandle } from '../../src/index.ts'; import { increment, read } from './counter.ts'; -// The state type produced by roundRobin()'s initialState(), derived rather -// than hardcoded so this example doesn't need to know its shared-memory shape. -type RoundRobinState = ReturnType['initialState']>>; - // Routes tasks to workers based on a key in the task's args, so that // per-worker state (caches, connections, compiled resources) can be reused // across successive calls for the same key. Falls back to round-robin for @@ -19,7 +15,7 @@ type RoundRobinState = ReturnType['ini export function keyAffinity(): Balancer<{ fallbackState: RoundRobinState }> { const fallback = roundRobin(); return { - initialState: () => ({ fallbackState: fallback.initialState!() }), + initialState: () => ({ fallbackState: fallback.initialState() }), select(workers: readonly WorkerHandle[], task: Task, { fallbackState }): WorkerHandle { let key: string | undefined; if (isTask(increment, task)) key = task.args[0]; diff --git a/src/balancers.ts b/src/balancers.ts index bb31c30..f2eaa70 100644 --- a/src/balancers.ts +++ b/src/balancers.ts @@ -2,14 +2,18 @@ import { uint32atomic } from './shared/descriptors.ts'; import type { Uint32Atomic } from './shared/uint32-atomic.ts'; import type { Balancer, WorkerHandle } from './runner.ts'; +/** Shared state for {@link roundRobin}: a cursor advanced atomically on every select. */ +export type RoundRobinState = { cursor: Uint32Atomic }; + /** * Creates a round-robin balancer that cycles through workers in order. * The cursor lives in shared memory so all scheduling threads advance the - * same sequence. (The uint32 cursor wraps at 2^32; a wrap may skip one - * position, which round-robin semantics tolerate.) + * same sequence. (The uint32 cursor wraps at 2^32; a wrap resets the cycle + * phase, which may skip or repeat positions once per 2^32 dispatches — + * round-robin semantics tolerate it.) * @returns A fresh Balancer instance. */ -export function roundRobin(): Balancer<{ cursor: Uint32Atomic }> { +export function roundRobin(): Balancer { return { initialState: () => ({ cursor: uint32atomic() }), select(workers: readonly WorkerHandle[], _task, { cursor }): WorkerHandle { @@ -21,7 +25,9 @@ export function roundRobin(): Balancer<{ cursor: Uint32Atomic }> { /** * Creates a least-busy balancer that picks the worker with the lowest activeCount. * Ties are broken by index (first wins). Reads the pool's shared active counts - * via the handles, so it needs no state of its own. + * via the handles, so it needs no state of its own. Under multi-threaded + * scheduling, concurrent selects may briefly herd onto the same worker before + * counts update; this is momentary imbalance, not corruption. * @returns A fresh Balancer instance. */ export function leastBusy(): Balancer { diff --git a/src/index.ts b/src/index.ts index 5963eb1..368c737 100644 --- a/src/index.ts +++ b/src/index.ts @@ -13,6 +13,7 @@ export { isTask } from './is-task.ts'; export { map } from './map.ts'; export type { MapOptions } from './map.ts'; export { roundRobin, leastBusy } from './balancers.ts'; +export type { RoundRobinState } from './balancers.ts'; export type { Task, RunResult, Balancer, Runner, WorkerHandle, WorkerOptions } from './runner.ts'; export { shared, diff --git a/src/runner.ts b/src/runner.ts index e6c3b45..1a52138 100644 --- a/src/runner.ts +++ b/src/runner.ts @@ -21,20 +21,24 @@ type TaskResults[]> = { [K in keyof T]: T[K] extends Task ? R : never; }; -/** A load balancing strategy for choosing which worker runs a task. */ -export interface Balancer { - /** Produces the balancer's state once at pool/runtime boot. May contain - * shared memory (which crosses threads by reference); plain values ride - * along by serialization. Closure state is safe only where scheduling is - * single-threaded (a `workers()` pool); shared, declared state works everywhere. */ - initialState?(): S; +/** + * A load balancing strategy for choosing which worker runs a task. + * + * `initialState()` produces the balancer's state once at pool/runtime boot. May + * contain shared memory (which crosses threads by reference); plain values ride + * along by serialization. Closure state is safe only where scheduling is + * single-threaded (a `workers()` pool); shared, declared state works everywhere. + * Required whenever `S` is non-`void`, so a stateful balancer can't omit it and + * crash at first dispatch. + */ +export type Balancer = { /** Choose a worker for the given task. Called synchronously on every dispatch. */ select(workers: readonly WorkerHandle[], task: Task, state: S): WorkerHandle; /** Optional cleanup on sync dispose. */ [Symbol.dispose]?(): void; /** Optional cleanup on async dispose. */ [Symbol.asyncDispose]?(): Promise; -} +} & ([S] extends [void] ? { initialState?(): S } : { initialState(): S }); /** Options for configuring a worker pool. */ export interface WorkerOptions { diff --git a/test/balancers.test.ts b/test/balancers.test.ts index beeb5a8..a9c7a04 100644 --- a/test/balancers.test.ts +++ b/test/balancers.test.ts @@ -10,7 +10,7 @@ function mockHandle(activeCount: number, index = 0): WorkerHandle { describe('roundRobin()', () => { it('cycles through workers in order', () => { const b = roundRobin(); - const state = b.initialState!(); + const state = b.initialState(); const handles = [mockHandle(0, 0), mockHandle(0, 1), mockHandle(0, 2)]; const task = { id: 'test', args: [], uid: 0 } as any; assert.equal(b.select(handles, task, state), handles[0]); @@ -19,9 +19,9 @@ describe('roundRobin()', () => { assert.equal(b.select(handles, task, state), handles[0]); }); - it('state is shared-memory backed (two instances over the same state agree)', () => { + it('state is externalized (two instances over the same state agree)', () => { const b1 = roundRobin(); - const state = b1.initialState!(); + const state = b1.initialState(); const b2 = roundRobin(); const handles = [mockHandle(0, 0), mockHandle(0, 1)]; const task = { id: 'test', args: [], uid: 0 } as any; diff --git a/test/fixtures/load-balancing.ts b/test/fixtures/load-balancing.ts index df4dda0..5c1cf40 100644 --- a/test/fixtures/load-balancing.ts +++ b/test/fixtures/load-balancing.ts @@ -1,8 +1,11 @@ import { setTimeout } from 'node:timers/promises'; +import { threadId } from 'node:worker_threads'; import { mo } from 'moroutine'; export const identity = mo(import.meta, (n: number): number => n); +export const whichThread = mo(import.meta, (): number => threadId); + export const slow = mo(import.meta, async (ms: number): Promise => { await setTimeout(ms); return 'done'; diff --git a/test/load-balancing.test.ts b/test/load-balancing.test.ts index e4dec09..b3d1944 100644 --- a/test/load-balancing.test.ts +++ b/test/load-balancing.test.ts @@ -2,14 +2,43 @@ import { describe, it } from 'node:test'; import assert from 'node:assert/strict'; import { workers, assign, leastBusy } from 'moroutine'; import type { Balancer } from 'moroutine'; -import { identity, slow } from './fixtures/load-balancing.ts'; +import { identity, slow, whichThread } from './fixtures/load-balancing.ts'; describe('load balancing', () => { it('defaults to least-busy', async () => { const run = workers(2); try { - const result = await run(identity(42)); - assert.equal(result, 42); + const worker1Thread = await run(assign(run.workers[1], whichThread())); + const busy = run(assign(run.workers[0], slow(200))); + assert.equal(await run(whichThread()), worker1Thread); + await busy; + } finally { + run[Symbol.dispose](); + } + }); + + it('threads state into select, and initialState runs exactly once', async () => { + let initCount = 0; + const sentinel = { tag: 'sentinel' }; + const seenStates: unknown[] = []; + const custom: Balancer = { + initialState() { + initCount++; + return sentinel; + }, + select(workers, _task, state) { + seenStates.push(state); + return workers[0]; + }, + }; + const run = workers(1, { balance: custom }); + try { + await run(identity(1)); + await run(identity(2)); + assert.equal(initCount, 1); + assert.equal(seenStates.length, 2); + assert.equal(seenStates[0], sentinel); + assert.equal(seenStates[1], sentinel); } finally { run[Symbol.dispose](); } -- 2.51.2