feat(fmt): support custom parallel worker count

Author: chenjiahan

packages/rstack/src/fmt/cli.ts

@@ -11,6 +11,7 @@ interface ParsedFmtCLIArgs {
   mode: FmtMode;
   patterns: string[];
   parallel: boolean;
+  parallelWorkers?: number;
   help: boolean;
 }
 
@@ -26,8 +27,22 @@ ${color.cyan('Options')}:
   --check             Check whether files are formatted
   --list-different    Print paths of unformatted files
   --no-parallel       Disable worker parallelism
+  --parallel-workers <count>  Number of parallel workers
   -h, --help          Display this help message`;
 
+const parseParallelWorkers = (value: string | undefined): number | undefined => {
+  if (value === undefined) {
+    return undefined;
+  }
+
+  const parallelWorkers = Number(value);
+  if (!/^\d+$/.test(value) || !Number.isSafeInteger(parallelWorkers) || parallelWorkers < 1) {
+    throw new Error('The --parallel-workers option must be a positive integer.');
+  }
+
+  return parallelWorkers;
+};
+
 const parseFmtCLIArgs = (args: string[]): ParsedFmtCLIArgs => {
   const { values, positionals } = parseArgs({
     args,
@@ -38,6 +53,8 @@ const parseFmtCLIArgs = (args: string[]): ParsedFmtCLIArgs => {
       listDifferent: { type: 'boolean' },
       'no-parallel': { type: 'boolean' },
       noParallel: { type: 'boolean' },
+      'parallel-workers': { type: 'string' },
+      parallelWorkers: { type: 'string' },
       help: { type: 'boolean', short: 'h' },
     },
     allowPositionals: true,
@@ -51,11 +68,27 @@ const parseFmtCLIArgs = (args: string[]): ParsedFmtCLIArgs => {
   }
 
   const mode = values.check ? 'check' : listDifferent ? 'list-different' : 'write';
+  const noParallel = values['no-parallel'] || values.noParallel;
+  const kebabParallelWorkers = values['parallel-workers'];
+  const camelParallelWorkers = values.parallelWorkers;
+
+  if (kebabParallelWorkers !== undefined && camelParallelWorkers !== undefined) {
+    throw new Error(
+      'The --parallel-workers and --parallelWorkers options cannot be used together.',
+    );
+  }
+
+  const parallelWorkers = parseParallelWorkers(kebabParallelWorkers ?? camelParallelWorkers);
+
+  if (noParallel && parallelWorkers !== undefined) {
+    throw new Error('The --parallel-workers and --no-parallel options cannot be used together.');
+  }
 
   return {
     mode,
     patterns: positionals,
-    parallel: !(values['no-parallel'] || values.noParallel),
+    parallel: !noParallel,
+    parallelWorkers,
     help: values.help ?? false,
   };
 };
@@ -98,7 +131,7 @@ const logFmtResult = (result: FmtRunResult, mode: FmtMode, cwd: string): void =>
 };
 
 const runFmtCLI = async (args: string[]): Promise<void> => {
-  const { help, mode, parallel, patterns } = parseFmtCLIArgs(args);
+  const { help, mode, parallel, parallelWorkers, patterns } = parseFmtCLIArgs(args);
   if (help) {
     console.log(fmtHelpMessage);
     return;
@@ -124,6 +157,7 @@ const runFmtCLI = async (args: string[]): Promise<void> => {
       mode,
       cache: false,
       parallel,
+      parallelWorkers,
     });
 
     logFmtResult(result, mode, cwd);

packages/rstack/src/fmt/parallel.ts

@@ -10,8 +10,8 @@ interface FmtWorker {
   terminate: () => void;
 }
 
-const getFmtWorkerCount = (fileCount: number): number =>
-  Math.min(fileCount, Math.max(1, availableParallelism() - 1));
+const getFmtWorkerCount = (fileCount: number, parallelWorkers?: number): number =>
+  Math.min(fileCount, parallelWorkers ?? Math.max(1, availableParallelism() - 1));
 
 const getFmtWorkerUrl = (): URL => {
   // Source tests run after build and exercise the same worker artifact as the CLI.
@@ -22,8 +22,8 @@ const getFmtWorkerUrl = (): URL => {
 };
 
 /** Creates and starts every worker before formatting can begin. */
-const createFmtWorker = async (fileCount: number): Promise<FmtWorker> => {
-  const workerCount = getFmtWorkerCount(fileCount);
+const createFmtWorker = async (fileCount: number, parallelWorkers?: number): Promise<FmtWorker> => {
+  const workerCount = getFmtWorkerCount(fileCount, parallelWorkers);
   const pool = new WorkTank<FmtWorkerMethods>({
     pool: {
       name: 'rstack-fmt',
@@ -51,4 +51,4 @@ const createFmtWorker = async (fileCount: number): Promise<FmtWorker> => {
   };
 };
 
-export { createFmtWorker };
+export { createFmtWorker, getFmtWorkerCount };

packages/rstack/src/fmt/runner.ts

@@ -54,9 +54,10 @@ const runFmtFilesSerial = async (
 const runFmtFilesParallel = async (
   files: FmtFileRequest[],
   shouldWrite: boolean,
+  parallelWorkers?: number,
 ): Promise<FmtFileResult[]> => {
   const { createFmtWorker } = await import('./parallel.ts');
-  const worker = await createFmtWorker(files.length);
+  const worker = await createFmtWorker(files.length, parallelWorkers);
 
   try {
     return await Promise.all(files.map((file) => runFmtFile(file, shouldWrite, worker.formatFile)));
@@ -96,12 +97,13 @@ const runFmtFiles = async ({
   files,
   mode,
   parallel,
+  parallelWorkers,
 }: RunFmtFilesOptions): Promise<FmtRunResult> => {
   const startTime = performance.now();
   const shouldWrite = mode === 'write';
   const results =
     parallel && files.length > 1 && canRunFmtFilesParallel(files)
-      ? await runFmtFilesParallel(files, shouldWrite)
+      ? await runFmtFilesParallel(files, shouldWrite, parallelWorkers)
       : await runFmtFilesSerial(files, shouldWrite);
 
   return {

packages/rstack/src/fmt/types.ts

@@ -56,6 +56,8 @@ interface RunFmtFilesOptions {
   cache: false;
   /** Whether cloneable file requests should run in worker threads. */
   parallel: boolean;
+  /** Maximum worker count when parallel execution is enabled. */
+  parallelWorkers?: number;
 }
 
 interface SuccessfulFmtFileResult {

packages/rstack/tests/cli/fmt/index.test.ts

@@ -84,11 +84,14 @@ test('formats the current directory with Prettier defaults', () => {
   expect(readProjectFile('index.ts')).toBe('const message = "hello";\n');
 });
 
-test('supports disabling parallel execution', () => {
+test.each([
+  ['disabling parallel execution', ['--no-parallel']],
+  ['configuring parallel worker count', ['--parallel-workers', '1']],
+] as const)('supports %s', (_, options) => {
   writeProjectFile('first.ts', 'const first="first"');
   writeProjectFile('second.ts', 'const second="second"');
 
-  const result = runFmt(['--no-parallel', 'first.ts', 'second.ts']);
+  const result = runFmt([...options, 'first.ts', 'second.ts']);
 
   expect(result.status).toBe(0);
   expect(result.stdout).toBe('first.ts\nsecond.ts\n');

packages/rstack/tests/fmt/cli.test.ts

@@ -6,6 +6,7 @@ test('uses write mode by default', () => {
     mode: 'write',
     patterns: [],
     parallel: true,
+    parallelWorkers: undefined,
     help: false,
   });
 });
@@ -20,6 +21,7 @@ test.each([
     mode,
     patterns: [],
     parallel: true,
+    parallelWorkers: undefined,
     help: false,
   });
 });
@@ -29,17 +31,56 @@ test.each(['--no-parallel', '--noParallel'])('disables parallel execution with %
     mode: 'write',
     patterns: [],
     parallel: false,
+    parallelWorkers: undefined,
     help: false,
   });
 });
 
+test.each(['--parallel-workers', '--parallelWorkers'])(
+  'configures parallel worker count with %s',
+  (option) => {
+    expect(parseFmtCLIArgs([option, '3'])).toEqual({
+      mode: 'write',
+      patterns: [],
+      parallel: true,
+      parallelWorkers: 3,
+      help: false,
+    });
+  },
+);
+
+test.each(['0', '-1', '1.5', 'invalid', '9007199254740992'])(
+  'rejects invalid parallel worker count %s',
+  (count) => {
+    expect(() => parseFmtCLIArgs([`--parallel-workers=${count}`])).toThrow(
+      'The --parallel-workers option must be a positive integer.',
+    );
+  },
+);
+
+test('rejects using both parallel worker aliases', () => {
+  expect(() => parseFmtCLIArgs(['--parallel-workers', '2', '--parallelWorkers', '3'])).toThrow(
+    'The --parallel-workers and --parallelWorkers options cannot be used together.',
+  );
+});
+
+test.each([
+  ['--no-parallel', '--parallel-workers'],
+  ['--noParallel', '--parallelWorkers'],
+])('rejects conflicting parallel options: %s and %s', (noParallel, parallelWorkers) => {
+  expect(() => parseFmtCLIArgs([noParallel, parallelWorkers, '2'])).toThrow(
+    'The --parallel-workers and --no-parallel options cannot be used together.',
+  );
+});
+
 test('preserves file paths and globs', () => {
   const patterns = ['src/file with spaces.ts', 'src/**/*.{js,ts}', '!src/generated/**'];
 
   expect(parseFmtCLIArgs([patterns[0], '--check', ...patterns.slice(1)])).toEqual({
     mode: 'check',
     patterns,
     parallel: true,
+    parallelWorkers: undefined,
     help: false,
   });
 });
@@ -49,6 +90,7 @@ test('treats arguments after the terminator as paths', () => {
     mode: 'check',
     patterns: ['--write', '--help'],
     parallel: true,
+    parallelWorkers: undefined,
     help: false,
   });
 });
@@ -63,6 +105,7 @@ test('provides command help', () => {
   expect(fmtHelpMessage).toContain('--check');
   expect(fmtHelpMessage).toContain('--list-different');
   expect(fmtHelpMessage).toContain('--no-parallel');
+  expect(fmtHelpMessage).toContain('--parallel-workers <count>');
   expect(fmtHelpMessage).toContain('-h, --help');
 });
 
@@ -78,9 +121,6 @@ test.each([
   );
 });
 
-test.each(['--unknown', '--no-cache', '--parallel-workers'])(
-  'rejects unsupported option %s',
-  (option) => {
-    expect(() => parseFmtCLIArgs([option])).toThrow();
-  },
-);
+test.each(['--unknown', '--no-cache'])('rejects unsupported option %s', (option) => {
+  expect(() => parseFmtCLIArgs([option])).toThrow();
+});

packages/rstack/tests/fmt/parallel.test.ts

@@ -0,0 +1,17 @@
+import { availableParallelism } from 'node:os';
+import { expect, test } from 'rstack/test';
+import { getFmtWorkerCount } from '../../src/fmt/parallel.ts';
+
+test('uses one fewer worker than the available parallelism by default', () => {
+  const defaultWorkerCount = Math.max(1, availableParallelism() - 1);
+
+  expect(getFmtWorkerCount(defaultWorkerCount + 1)).toBe(defaultWorkerCount);
+});
+
+test.each([
+  [4, 1, 1],
+  [4, 2, 2],
+  [2, 4, 2],
+])('uses %s files and %s configured workers as %s workers', (files, workers, expected) => {
+  expect(getFmtWorkerCount(files, workers)).toBe(expected);
+});

packages/rstack/tests/fmt/runnerParallelPreflight.test.ts

@@ -1,13 +1,24 @@
 import { readFileSync } from 'node:fs';
-import { expect, rs, test } from 'rstack/test';
+import { beforeEach, expect, rs, test } from 'rstack/test';
 import { runFmtFiles } from '../../src/fmt/runner.ts';
 import type { FmtFileRequest } from '../../src/fmt/types.ts';
 import { withTempProject, writeProjectFile } from './helpers.ts';
 
+const mocks = rs.hoisted(() => ({
+  createFmtWorkerCalls: [] as [number, number | undefined][],
+}));
+
 rs.mock('../../src/fmt/parallel.ts', () => ({
-  createFmtWorker: () => Promise.reject(new Error('worker startup failed')),
+  createFmtWorker: (fileCount: number, parallelWorkers?: number) => {
+    mocks.createFmtWorkerCalls.push([fileCount, parallelWorkers]);
+    return Promise.reject(new Error('worker startup failed'));
+  },
 }));
 
+beforeEach(() => {
+  mocks.createFmtWorkerCalls.length = 0;
+});
+
 const createRequest = (
   filePath: string,
   plugins?: FmtFileRequest['options']['plugins'],
@@ -32,9 +43,12 @@ test('does not write files when worker startup fails', async () => {
         mode: 'write',
         cache: false,
         parallel: true,
+        parallelWorkers: 3,
       }),
     ).rejects.toThrow('worker startup failed');
 
+    expect(mocks.createFmtWorkerCalls).toEqual([[2, 3]]);
+
     for (const filePath of filePaths) {
       expect(readFileSync(filePath, 'utf8')).toBe('const value=1');
     }