From ea7d43c616c4cd3bb77f27eaf9f3becd32e41744 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20=C3=81ngel?= Date: Sat, 31 Oct 2026 17:08:22 +0000 Subject: [PATCH] fix(cli): stopping remove-background on a Mac no longer aborts the process (#6848) * fix(cli): stopping remove-background on a Mac no longer aborts the process * fix(cli): a quit after background removal released its session exits at once * Revert "fix(cli): a quit after background removal released its session exits at once" This reverts commit 3b9d6ffc247b8abe898b0312a0c395cd9ee8df38. --- .../background-removal/inference.stop.test.ts | 96 +++++++++++++++++++ .../cli/src/background-removal/inference.ts | 74 +++++++++----- packages/cli/src/cli.ts | 9 +- packages/cli/src/utils/commandResult.test.ts | 40 +++++++- packages/cli/src/utils/commandResult.ts | 19 ++++ 5 files changed, 211 insertions(+), 27 deletions(-) create mode 100644 packages/cli/src/background-removal/inference.stop.test.ts diff --git a/packages/cli/src/background-removal/inference.stop.test.ts b/packages/cli/src/background-removal/inference.stop.test.ts new file mode 100644 index 0000000000..b5cd8484d7 --- /dev/null +++ b/packages/cli/src/background-removal/inference.stop.test.ts @@ -0,0 +1,96 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const cleanups = vi.hoisted(() => new Set<() => Promise>()); +const native = vi.hoisted(() => ({ + frames: [] as (() => void)[], + released: 0, + releasedDuringFrame: 0, + inFrame: false, +})); + +vi.mock("../browser/launch.js", () => ({ + registerBrowserSessionCleanup: (cleanup: () => Promise) => { + cleanups.add(cleanup); + return () => cleanups.delete(cleanup); + }, +})); +vi.mock("./manager.js", () => ({ + ensureModel: async () => "model.onnx", + selectProviders: () => ({ providers: ["coreml", "cpu"], label: "CoreML" }), +})); +vi.mock("./coremlTemp.js", () => ({ + trackCoreMlTemp: () => ({ record: () => undefined, forget: () => undefined }), +})); +// A frame inside the native session ends when the test says so; releasing the session meanwhile is what aborts +// the process on a Mac ("mutex lock failed"). +vi.mock("../utils/optionalPackages.js", () => { + const chain: Record = {}; + for (const step of ["resize", "raw", "toColourspace"]) chain[step] = () => chain; + chain.toBuffer = async () => Buffer.alloc(320 * 320 * 3); + return { + loadSharp: async () => () => chain, + loadOptionalPackage: async () => ({ + Tensor: class {}, + InferenceSession: { + create: async () => ({ + inputNames: ["in"], + outputNames: ["out"], + run: () => + new Promise((resolve) => { + native.inFrame = true; + native.frames.push(() => { + native.inFrame = false; + resolve({ out: { data: new Float32Array(320 * 320) } }); + }); + }), + release: async () => { + native.released++; + if (native.inFrame) native.releasedDuringFrame++; + }, + }), + }, + }), + }; +}); + +import { createSession } from "./inference.js"; + +const frame = (session: Awaited>) => + session.process(Buffer.alloc(48), 4, 4); +const settle = () => new Promise((resolve) => setTimeout(resolve, 5)); + +beforeEach(() => { + cleanups.clear(); + Object.assign(native, { frames: [], released: 0, releasedDuringFrame: 0, inFrame: false }); +}); + +describe("a stop signal while a frame is inside the session", () => { + it("releases the session only after that frame has ended", async () => { + const session = await createSession(); + const running = frame(session); + await settle(); + const [stop] = [...cleanups]; + + const stopped = stop!(); + await settle(); + expect(native.released).toBe(0); + + native.frames.shift()!(); + await running.catch(() => undefined); + await stopped; + + expect(native.released).toBe(1); + expect(native.releasedDuringFrame).toBe(0); + }); + + it("starts no frame after the stop, and closing afterwards does not release a second time", async () => { + const session = await createSession(); + const [stop] = [...cleanups]; + await stop!(); + + await expect(frame(session)).rejects.toThrow(/stopped/i); + await session.close(); + + expect(native.released).toBe(1); + }); +}); diff --git a/packages/cli/src/background-removal/inference.ts b/packages/cli/src/background-removal/inference.ts index 1f124c45fa..a1fd33bbdc 100644 --- a/packages/cli/src/background-removal/inference.ts +++ b/packages/cli/src/background-removal/inference.ts @@ -9,6 +9,7 @@ import type { InferenceSession, Tensor } from "onnxruntime-node"; import type sharpType from "sharp"; import { loadOptionalPackage, loadSharp } from "../utils/optionalPackages.js"; import { registerBrowserSessionCleanup } from "../browser/launch.js"; +import { drainToExit } from "../utils/commandResult.js"; import { trackCoreMlTemp } from "./coremlTemp.js"; import { ensureModel, selectProviders, type Device, type ModelId } from "./manager.js"; @@ -75,7 +76,23 @@ export async function createSession(options: CreateSessionOptions = {}): Promise const coreMlTemp = choice.label === "CoreML" ? trackCoreMlTemp() : undefined; // A signal stops the run through the CLI's shutdown, which releases the session so CoreML removes its temp files. - const unregister = coreMlTemp && registerBrowserSessionCleanup(async () => session?.release()); + // Releasing it while a frame is inside it aborts the process, so the stop waits for that frame and refuses the next. + if (coreMlTemp) drainToExit(); + let stopped = false; + let inFlight: Promise | undefined; + let released = false; + const release = async () => { + if (released || !session) return; + released = true; + await session.release(); + }; + const unregister = + coreMlTemp && + registerBrowserSessionCleanup(async () => { + stopped = true; + await inFlight; + await release(); + }); let session: InferenceSession; let providerUsed = choice.label; try { @@ -125,38 +142,45 @@ export async function createSession(options: CreateSessionOptions = {}): Promise } }; + const convert = async (rgb: Buffer, width: number, height: number, withBackground: boolean) => { + const tensor = await preprocess(sharp, ort, rgb, width, height, inputData); + const outputs = await run(tensor); + const output = outputs[outputName]; + if (!output) throw new Error(`Model did not return output '${outputName}'`); + const expectedBytes = width * height * 4; + if (!rgbaBuf || rgbaBuf.length !== expectedBytes) { + rgbaBuf = Buffer.allocUnsafe(expectedBytes); + } + if (withBackground) { + if (!rgbaBgBuf || rgbaBgBuf.length !== expectedBytes) { + rgbaBgBuf = Buffer.allocUnsafe(expectedBytes); + } + } + return await postprocess( + sharp, + output, + rgb, + width, + height, + maskBuf, + rgbaBuf, + withBackground ? rgbaBgBuf : null, + ); + }; + return { get provider() { return providerUsed; }, async process(rgb, width, height, withBackground = false) { - const tensor = await preprocess(sharp, ort, rgb, width, height, inputData); - const outputs = await run(tensor); - const output = outputs[outputName]; - if (!output) throw new Error(`Model did not return output '${outputName}'`); - const expectedBytes = width * height * 4; - if (!rgbaBuf || rgbaBuf.length !== expectedBytes) { - rgbaBuf = Buffer.allocUnsafe(expectedBytes); - } - if (withBackground) { - if (!rgbaBgBuf || rgbaBgBuf.length !== expectedBytes) { - rgbaBgBuf = Buffer.allocUnsafe(expectedBytes); - } - } - return await postprocess( - sharp, - output, - rgb, - width, - height, - maskBuf, - rgbaBuf, - withBackground ? rgbaBgBuf : null, - ); + if (stopped) throw new Error("Background removal was stopped."); + const frame = convert(rgb, width, height, withBackground); + inFlight = frame.catch(() => undefined); + return await frame; }, async close() { unregister?.(); - await session.release(); + await release(); coreMlTemp?.forget(); }, }; diff --git a/packages/cli/src/cli.ts b/packages/cli/src/cli.ts index af7acf86d0..a2b784f690 100644 --- a/packages/cli/src/cli.ts +++ b/packages/cli/src/cli.ts @@ -97,6 +97,7 @@ import { CliUsageError, USAGE_EXIT_CODE, consumeCommandResult, + endProcess, outputWritten, registerRootExitCodeSanitizer, registerRootExitRequester, @@ -373,7 +374,13 @@ registerRootExitRequester((exitCode) => { presented: true, }).finally(async () => { await outputWritten(); - process.exit(exitCode); + endProcess(exitCode, { + exit: (code) => process.exit(code), + setExitCode: (code) => { + process.exitCode = code; + }, + later: setTimeout, + }); }); }); diff --git a/packages/cli/src/utils/commandResult.test.ts b/packages/cli/src/utils/commandResult.test.ts index c127462320..a42c5dbc01 100644 --- a/packages/cli/src/utils/commandResult.test.ts +++ b/packages/cli/src/utils/commandResult.test.ts @@ -1,9 +1,11 @@ -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { CliResultSignal, CliRuntimeError, CliUsageError, consumeCommandResult, + drainToExit, + endProcess, failCommand, failUsage, finishCommand, @@ -33,3 +35,39 @@ describe("command result contract", () => { expect(consumeCommandResult()).toMatchObject({ exitCode: 1, kind: "runtime_error" }); }); }); + +describe("ending the process", () => { + const host = () => { + const pending: (() => void)[] = []; + const calls = { exit: vi.fn(), setExitCode: vi.fn() }; + const later = ((run: () => void) => { + pending.push(run); + return { unref: () => undefined }; + }) as unknown as typeof setTimeout; + return { pending, calls, ...calls, later }; + }; + + it("exits at once when nothing native needs tearing down", async () => { + vi.resetModules(); + const fresh = await import("./commandResult.js"); + const ending = host(); + + fresh.endProcess(143, ending); + + expect(ending.exit).toHaveBeenCalledWith(143); + expect(ending.pending).toHaveLength(0); + }); + + it("after CoreML ran, leaves the exit to Node's own teardown and forces it only after a grace", () => { + const ending = host(); + drainToExit(); + + endProcess(143, ending); + + expect(ending.exit).not.toHaveBeenCalled(); + expect(ending.setExitCode).toHaveBeenCalledWith(143); + expect(ending.pending).toHaveLength(1); + ending.pending[0]!(); + expect(ending.exit).toHaveBeenCalledWith(143); + }); +}); diff --git a/packages/cli/src/utils/commandResult.ts b/packages/cli/src/utils/commandResult.ts index 92e150c53f..8f7ff1b8c9 100644 --- a/packages/cli/src/utils/commandResult.ts +++ b/packages/cli/src/utils/commandResult.ts @@ -120,6 +120,25 @@ export async function outputWritten(): Promise { await Promise.all([written(process.stdout), written(process.stderr)]); } +let drainsToExit = false; +const DRAIN_GRACE_MS = 8_000; + +export function drainToExit(): void { + drainsToExit = true; +} + +export interface ProcessEnd { + exit(exitCode: number): void; + setExitCode(exitCode: number): void; + later: typeof setTimeout; +} + +export function endProcess(exitCode: number, host: ProcessEnd): void { + if (!drainsToExit) return host.exit(exitCode); + host.setExitCode(exitCode); + host.later(() => host.exit(exitCode), DRAIN_GRACE_MS).unref(); +} + /** Ask cli.ts to finalize telemetry/output and then terminate the process. */ export function requestCliExit(exitCode = 0): void { if (!rootExitRequester) {