diff --git a/packages/stage-ui/package.json b/packages/stage-ui/package.json index 96ab22460..9fc4c61b9 100644 --- a/packages/stage-ui/package.json +++ b/packages/stage-ui/package.json @@ -103,6 +103,7 @@ "@xsai/tool": "catalog:", "@xsai/utils-chat": "catalog:", "animejs": "^4.3.6", + "async-mutex": "catalog:", "better-auth": "catalog:", "culori": "^4.0.2", "d3": "catalog:", diff --git a/packages/stage-ui/src/libs/inference/adapters/background-removal.ts b/packages/stage-ui/src/libs/inference/adapters/background-removal.ts index 542681e84..2403e34cb 100644 --- a/packages/stage-ui/src/libs/inference/adapters/background-removal.ts +++ b/packages/stage-ui/src/libs/inference/adapters/background-removal.ts @@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator' import type { ProgressPayload } from '../protocol' import { defaultPerfTracer } from '@proj-airi/stage-shared' +import { Mutex } from 'async-mutex' import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status' -import { AsyncMutex } from '../async-mutex' import { MODEL_IDS, MODEL_NAMES, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' @@ -58,7 +58,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { let state: BackgroundRemovalAdapter['state'] = 'idle' let allocationToken: AllocationToken | null = null - const operationMutex = new AsyncMutex() + const operationMutex = new Mutex() function ensureWorker(): Worker { if (!worker) { @@ -66,9 +66,9 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { new URL('../../../workers/background-removal/worker.ts', import.meta.url), { type: 'module' }, ) - worker.addEventListener('error', (event) => { + worker.addEventListener('error', (_event) => { state = 'error' - operationMutex.reset(new Error(event.message ?? 'Worker error')) + operationMutex.cancel() }) } return worker @@ -119,7 +119,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { } async function load(onProgress?: (p: ProgressPayload) => void): Promise { - return operationMutex.run(async () => { + return operationMutex.runExclusive(async () => { state = 'loading' updateInferenceStatus(MODEL_NAMES.BG_REMOVAL, { state: 'downloading', device: 'webgpu' }) @@ -172,7 +172,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { } async function processImage(imageData: ImageData): Promise { - return defaultPerfTracer.withMeasure('inference', 'bg-removal-process', () => operationMutex.run(async () => { + return defaultPerfTracer.withMeasure('inference', 'bg-removal-process', () => operationMutex.runExclusive(async () => { if (!worker || (state !== 'ready' && state !== 'processing')) throw new Error('Model not loaded. Call load() first.') @@ -222,7 +222,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { } function terminateAdapter(): void { - operationMutex.reset(new Error('Adapter terminated')) + operationMutex.cancel() if (worker) { worker.terminate() worker = null diff --git a/packages/stage-ui/src/libs/inference/adapters/kokoro.ts b/packages/stage-ui/src/libs/inference/adapters/kokoro.ts index 0a4978651..1495494b1 100644 --- a/packages/stage-ui/src/libs/inference/adapters/kokoro.ts +++ b/packages/stage-ui/src/libs/inference/adapters/kokoro.ts @@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator' import type { ProgressPayload } from '../protocol' import { defaultPerfTracer } from '@proj-airi/stage-shared' +import { Mutex } from 'async-mutex' import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status' -import { AsyncMutex } from '../async-mutex' import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' @@ -161,8 +161,8 @@ export function createKokoroAdapter(): KokoroAdapter { let allocationToken: AllocationToken | null = null let currentModelStatusId: string | null = null - const operationMutex = new AsyncMutex() - const lifecycleMutex = new AsyncMutex() + const operationMutex = new Mutex() + const lifecycleMutex = new Mutex() function initializeWorker(): void { worker = new Worker( @@ -172,13 +172,9 @@ export function createKokoroAdapter(): KokoroAdapter { worker.addEventListener('error', handleWorkerError) } - function handleWorkerError(event: ErrorEvent | Error): void { - const message = event instanceof Error - ? event.message - : (event as ErrorEvent).message ?? 'Unknown worker error' - + function handleWorkerError(_event: ErrorEvent | Error): void { state = 'error' - operationMutex.reset(new Error(message)) + operationMutex.cancel() destroyWorker() scheduleRestart() } @@ -218,7 +214,7 @@ export function createKokoroAdapter(): KokoroAdapter { } async function ensureStarted(): Promise { - await lifecycleMutex.run(async () => { + await lifecycleMutex.runExclusive(async () => { if (!worker) { initializeWorker() state = 'idle' @@ -235,7 +231,7 @@ export function createKokoroAdapter(): KokoroAdapter { ): Promise { await ensureStarted() - return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.run(async () => { + return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.runExclusive(async () => { state = 'loading' const modelStatusId = `kokoro-${quantization}` @@ -298,7 +294,7 @@ export function createKokoroAdapter(): KokoroAdapter { } async function generate(text: string, voice: VoiceKey): Promise { - return defaultPerfTracer.withMeasure('inference', 'kokoro-generate', () => operationMutex.run(async () => { + return defaultPerfTracer.withMeasure('inference', 'kokoro-generate', () => operationMutex.runExclusive(async () => { if (!worker) throw new Error('Worker not initialized. Call loadModel() first.') @@ -341,7 +337,7 @@ export function createKokoroAdapter(): KokoroAdapter { } function terminateAdapter(): void { - operationMutex.reset(new Error('Adapter terminated')) + operationMutex.cancel() destroyWorker() if (allocationToken) { removeInferenceStatus(allocationToken.modelId) @@ -366,14 +362,14 @@ export function createKokoroAdapter(): KokoroAdapter { // --------------------------------------------------------------------------- let globalAdapter: KokoroAdapter | null = null -const singletonMutex = new AsyncMutex() +const singletonMutex = new Mutex() /** * Get the global Kokoro adapter instance. * Creates and starts the worker on first call. */ export async function getKokoroAdapter(): Promise { - return singletonMutex.run(async () => { + return singletonMutex.runExclusive(async () => { if (!globalAdapter) globalAdapter = createKokoroAdapter() return globalAdapter diff --git a/packages/stage-ui/src/libs/inference/adapters/whisper.ts b/packages/stage-ui/src/libs/inference/adapters/whisper.ts index 29a0ae370..9808e9046 100644 --- a/packages/stage-ui/src/libs/inference/adapters/whisper.ts +++ b/packages/stage-ui/src/libs/inference/adapters/whisper.ts @@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator' import type { ProgressPayload } from '../protocol' import { defaultPerfTracer } from '@proj-airi/stage-shared' +import { Mutex } from 'async-mutex' import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status' -import { AsyncMutex } from '../async-mutex' import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' @@ -84,15 +84,11 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { let restartAttempts = 0 const messageHandlers = new Set<(event: WhisperEvent) => void>() - const operationMutex = new AsyncMutex() - - function handleWorkerError(event: ErrorEvent | Error): void { - const message = event instanceof Error - ? event.message - : (event as ErrorEvent).message ?? 'Unknown worker error' + const operationMutex = new Mutex() + function handleWorkerError(_event: ErrorEvent | Error): void { state = 'error' - operationMutex.reset(new Error(message)) + operationMutex.cancel() destroyWorker() scheduleRestart() } @@ -199,7 +195,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { async function load( onProgress?: (p: ProgressPayload) => void, ): Promise { - return operationMutex.run(async () => { + return operationMutex.runExclusive(async () => { state = 'loading' updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'downloading', device: 'webgpu' }) @@ -253,7 +249,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { } async function transcribe(input: WhisperTranscribeInput): Promise { - return defaultPerfTracer.withMeasure('inference', 'whisper-transcribe', () => operationMutex.run(async () => { + return defaultPerfTracer.withMeasure('inference', 'whisper-transcribe', () => operationMutex.runExclusive(async () => { if (!worker || state !== 'ready') throw new Error('Model not loaded. Call load() first.') @@ -286,7 +282,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { } function terminateAdapter(): void { - operationMutex.reset(new Error('Adapter terminated')) + operationMutex.cancel() if (worker) { worker.terminate() worker = null diff --git a/packages/stage-ui/src/libs/inference/async-mutex.test.ts b/packages/stage-ui/src/libs/inference/async-mutex.test.ts deleted file mode 100644 index d68355f0a..000000000 --- a/packages/stage-ui/src/libs/inference/async-mutex.test.ts +++ /dev/null @@ -1,68 +0,0 @@ -import { describe, expect, it } from 'vitest' - -import { AsyncMutex } from './async-mutex' - -describe('asyncMutex', () => { - it('should serialize concurrent operations', async () => { - const mutex = new AsyncMutex() - const order: number[] = [] - - const p1 = mutex.run(async () => { - await new Promise(r => setTimeout(r, 50)) - order.push(1) - return 'a' - }) - - const p2 = mutex.run(async () => { - order.push(2) - return 'b' - }) - - const [r1, r2] = await Promise.all([p1, p2]) - - expect(r1).toBe('a') - expect(r2).toBe('b') - expect(order).toEqual([1, 2]) - }) - - it('should handle errors without blocking subsequent operations', async () => { - const mutex = new AsyncMutex() - - await expect(mutex.run(async () => { - throw new Error('fail') - })).rejects.toThrow('fail') - - // Should still work after error - const result = await mutex.run(async () => 'ok') - expect(result).toBe('ok') - }) - - it('should reset and reject waiting operations', async () => { - const mutex = new AsyncMutex() - - // Queue a waiting operation that will never resolve on its own - let resolveHold!: () => void - const holdPromise = new Promise(r => resolveHold = r) - - const hold = mutex.run(() => holdPromise).catch(() => 'hold-rejected') - const waiting = mutex.run(async () => 'waited').catch(() => 'waiting-rejected') - - // Reset — should reject waiting operations - await new Promise(r => setTimeout(r, 10)) - mutex.reset(new Error('reset')) - resolveHold() - - const [holdResult, waitingResult] = await Promise.all([hold, waiting]) - - // Waiting should be rejected - expect(waitingResult).toBe('waiting-rejected') - // Hold may or may not be rejected (it was active) - expect(['hold-rejected', undefined]).toContain(holdResult) - }) - - it('should support synchronous callbacks', async () => { - const mutex = new AsyncMutex() - const result = await mutex.run(() => 42) - expect(result).toBe(42) - }) -}) diff --git a/packages/stage-ui/src/libs/inference/async-mutex.ts b/packages/stage-ui/src/libs/inference/async-mutex.ts deleted file mode 100644 index 28d4c0bd9..000000000 --- a/packages/stage-ui/src/libs/inference/async-mutex.ts +++ /dev/null @@ -1,62 +0,0 @@ -/** - * An async mutex that ensures only one callback runs at a time. - * Waiters queue up and are processed in FIFO order. - * - * Extracted from KokoroWorkerManager for reuse across inference workers. - */ -export class AsyncMutex { - private locked = false - private waiters: { resolve: () => void, reject: (error: Error) => void }[] = [] - - // Incremented on reset() to invalidate stale lock holders - private generation = 0 - - /** - * Executes the callback with exclusive access to the mutex. - * If the mutex is locked, waits in queue until it's our turn. - */ - async run(callback: () => Promise | T): Promise { - const myGeneration = this.generation - - if (this.locked) { - await new Promise((resolve, reject) => { - this.waiters.push({ resolve, reject }) - }) - if (myGeneration !== this.generation) { - throw new Error('Mutex was reset') - } - } - this.locked = true - - try { - return await callback() - } - finally { - if (myGeneration === this.generation) { - const next = this.waiters.shift() - if (next) { - next.resolve() - } - else { - this.locked = false - } - } - } - } - - /** - * Cancels all waiting tasks and releases the lock. - * Atomic thanks to JavaScript's Run-To-Completion semantics. - */ - reset(error: Error = new Error('Mutex reset')): void { - this.generation++ - this.locked = false - - const waitersToReject = this.waiters - this.waiters = [] - - for (const waiter of waitersToReject) { - waiter.reject(error) - } - } -} diff --git a/packages/stage-ui/src/libs/inference/index.ts b/packages/stage-ui/src/libs/inference/index.ts index d6c7d9a89..532c567d2 100644 --- a/packages/stage-ui/src/libs/inference/index.ts +++ b/packages/stage-ui/src/libs/inference/index.ts @@ -1,5 +1,3 @@ -// Core -export { AsyncMutex } from './async-mutex' // Cache utilities export { clearModelCache, diff --git a/packages/stage-ui/src/libs/inference/worker-manager.ts b/packages/stage-ui/src/libs/inference/worker-manager.ts index 6b0502601..3e0f881e8 100644 --- a/packages/stage-ui/src/libs/inference/worker-manager.ts +++ b/packages/stage-ui/src/libs/inference/worker-manager.ts @@ -2,7 +2,7 @@ * Generic inference worker manager. * * Provides lifecycle management (start / restart / terminate), request - * serialisation via AsyncMutex, timeout handling, and a unified + * serialisation via async-mutex, timeout handling, and a unified * message protocol for any inference worker. * * NOTICE: Currently not consumed by any adapter. Adapters implement their @@ -20,8 +20,8 @@ import type { } from './protocol' import { errorMessageFrom } from '@moeru/std' +import { Mutex } from 'async-mutex' -import { AsyncMutex } from './async-mutex' import { createRequestId } from './protocol' // --------------------------------------------------------------------------- @@ -158,8 +158,8 @@ export function createInferenceWorkerManager( let lastError: ErrorPayload | null = null let restartAttempts = 0 - const operationMutex = new AsyncMutex() - const lifecycleMutex = new AsyncMutex() + const operationMutex = new Mutex() + const lifecycleMutex = new Mutex() // -- Worker lifecycle ----------------------------------------------------- @@ -181,7 +181,7 @@ export function createInferenceWorkerManager( state = 'error' // Reject all pending operations - operationMutex.reset(new Error(message)) + operationMutex.cancel() // Clean up and try to restart destroyWorker() @@ -223,7 +223,7 @@ export function createInferenceWorkerManager( } async function ensureStarted(): Promise { - await lifecycleMutex.run(async () => { + await lifecycleMutex.runExclusive(async () => { if (!worker) { initializeWorker() state = 'idle' @@ -239,7 +239,7 @@ export function createInferenceWorkerManager( ): Promise { await ensureStarted() - return operationMutex.run(async () => { + return operationMutex.runExclusive(async () => { state = 'loading' const requestId = createRequestId() @@ -278,7 +278,7 @@ export function createInferenceWorkerManager( input: TInput, onProgress?: (p: ProgressPayload) => void, ): Promise { - return operationMutex.run(async () => { + return operationMutex.runExclusive(async () => { if (!worker) throw new Error('Worker not initialized. Call loadModel() first.') @@ -319,7 +319,7 @@ export function createInferenceWorkerManager( } async function unloadModel(): Promise { - return operationMutex.run(async () => { + return operationMutex.runExclusive(async () => { if (!worker) return @@ -341,7 +341,7 @@ export function createInferenceWorkerManager( } function terminateManager(): void { - operationMutex.reset(new Error('Manager terminated')) + operationMutex.cancel() destroyWorker() state = 'terminated' } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index a5a7f498e..bc1f69fe6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3065,6 +3065,9 @@ importers: animejs: specifier: ^4.3.6 version: 4.3.6 + async-mutex: + specifier: 'catalog:' + version: 0.5.0 better-auth: specifier: 'catalog:' version: 1.6.2(@opentelemetry/api@1.9.1)(@prisma/client@5.22.0)(better-sqlite3@12.5.0)(drizzle-kit@0.31.10)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.4)(@opentelemetry/api@1.9.1)(@prisma/client@5.22.0)(@types/pg@8.20.0)(better-sqlite3@12.5.0)(kysely@0.28.14)(pg@8.20.0)(postgres@3.4.9))(pg@8.20.0)(react@19.2.3)(vitest@4.1.4)(vue@3.5.32(typescript@5.9.3)) @@ -8733,6 +8736,7 @@ packages: '@simple-git/argv-parser@1.1.0': resolution: {integrity: sha512-sUKOu2lb5vGIWADNNLpscyj07DAeQZU3KLbnE2Tj53tW6BbDQKMly2CCfnR4oYzqtRELCPWfwaPg+Q0T8qfKBg==} + deprecated: Contains a breaking change that should be a major version bump '@sindresorhus/base62@1.0.0': resolution: {integrity: sha512-TeheYy0ILzBEI/CO55CP6zJCSdSWeRtGnHy8U8dWSUH4I68iqTsy7HkMktR4xakThc9jotkPQUXT4ITdbV7cHA==}