Skip to content

Commit e651809

Browse files
steipetesteipete
committed
perf: slim gateway startup imports
1 parent b6a9018 commit e651809

6 files changed

Lines changed: 297 additions & 131 deletions

File tree

‎src/cli/gateway-cli/run-loop.test.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,20 @@ async function waitForStart(started: Promise<void>) {
241241
await new Promise<void>((resolve) => setImmediate(resolve));
242242
}
243243

244+
async function waitForAssertion(assertion: () => void, maxTicks = 20) {
245+
for (let tick = 0; tick < maxTicks; tick += 1) {
246+
try {
247+
assertion();
248+
return;
249+
} catch (err) {
250+
if (tick === maxTicks - 1) {
251+
throw err;
252+
}
253+
await new Promise<void>((resolve) => setImmediate(resolve));
254+
}
255+
}
256+
}
257+
244258
async function createSignaledLoopHarness(exitCallOrder?: string[]) {
245259
const close = vi.fn(async () => {});
246260
const { start, started } = createSignaledStart(close);
@@ -302,7 +316,7 @@ describe("runGatewayLoop", () => {
302316
reason: "gateway restarting",
303317
restartExpectedMs: 1500,
304318
});
305-
expect(start).toHaveBeenCalledTimes(2);
319+
await waitForAssertion(() => expect(start).toHaveBeenCalledTimes(2));
306320

307321
sigint();
308322
await expect(exited).resolves.toBe(0);

‎src/cli/gateway-cli/run-loop.ts‎

Lines changed: 152 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,9 @@
11
import net from "node:net";
2-
import {
3-
abortEmbeddedPiRun,
4-
getActiveEmbeddedRunCount,
5-
waitForActiveEmbeddedRuns,
6-
} from "../../agents/pi-embedded-runner/runs.js";
7-
import { getRuntimeConfig } from "../../config/config.js";
82
import type { startGatewayServer } from "../../gateway/server.js";
93
import { formatErrorMessage } from "../../infra/errors.js";
104
import { acquireGatewayLock } from "../../infra/gateway-lock.js";
11-
import {
12-
respawnGatewayProcessForUpdate,
13-
restartGatewayProcessWithFreshPid,
14-
} from "../../infra/process-respawn.js";
15-
import { markUpdateRestartSentinelFailure } from "../../infra/restart-sentinel.js";
16-
import {
17-
consumeGatewaySigusr1RestartAuthorization,
18-
consumeGatewayRestartIntentSync,
19-
isGatewaySigusr1RestartExternallyAllowed,
20-
markGatewaySigusr1RestartHandled,
21-
peekGatewaySigusr1RestartReason,
22-
resetGatewayRestartStateForInProcessRestart,
23-
scheduleGatewaySigusr1Restart,
24-
} from "../../infra/restart.js";
25-
import { detectRespawnSupervisor } from "../../infra/supervisor-markers.js";
26-
import { writeDiagnosticStabilityBundleForFailureSync } from "../../logging/diagnostic-stability-bundle.js";
275
import { createSubsystemLogger } from "../../logging/subsystem.js";
28-
import {
29-
getActiveBundledRuntimeDepsInstallCount,
30-
waitForBundledRuntimeDepsInstallIdle,
31-
} from "../../plugins/bundled-runtime-deps-activity.js";
32-
import {
33-
getActiveTaskCount,
34-
markGatewayDraining,
35-
resetAllLanes,
36-
waitForActiveTasks,
37-
} from "../../process/command-queue.js";
38-
import { createRestartIterationHook } from "../../process/restart-recovery.js";
396
import type { RuntimeEnv } from "../../runtime.js";
40-
import { reloadTaskRegistryFromStore } from "../../tasks/runtime-internal.js";
417

428
const gatewayLog = createSubsystemLogger("gateway");
439
const LAUNCHD_SUPERVISED_RESTART_EXIT_DELAY_MS = 1500;
@@ -49,6 +15,61 @@ const UPDATE_RESPAWN_HEALTH_POLL_MS = 200;
4915
type GatewayRunSignalAction = "stop" | "restart";
5016
type RestartDrainTimeoutMs = number | undefined;
5117

18+
type EmbeddedRunsModule = typeof import("../../agents/pi-embedded-runner/runs.js");
19+
type RuntimeConfigModule = typeof import("../../config/config.js");
20+
type ProcessRespawnModule = typeof import("../../infra/process-respawn.js");
21+
type RestartSentinelModule = typeof import("../../infra/restart-sentinel.js");
22+
type RestartModule = typeof import("../../infra/restart.js");
23+
type SupervisorMarkersModule = typeof import("../../infra/supervisor-markers.js");
24+
type DiagnosticStabilityBundleModule =
25+
typeof import("../../logging/diagnostic-stability-bundle.js");
26+
type BundledRuntimeDepsActivityModule =
27+
typeof import("../../plugins/bundled-runtime-deps-activity.js");
28+
type CommandQueueModule = typeof import("../../process/command-queue.js");
29+
type RuntimeInternalModule = typeof import("../../tasks/runtime-internal.js");
30+
31+
let embeddedRunsModule: Promise<EmbeddedRunsModule> | undefined;
32+
let runtimeConfigModule: Promise<RuntimeConfigModule> | undefined;
33+
let processRespawnModule: Promise<ProcessRespawnModule> | undefined;
34+
let restartSentinelModule: Promise<RestartSentinelModule> | undefined;
35+
let restartModule: Promise<RestartModule> | undefined;
36+
let supervisorMarkersModule: Promise<SupervisorMarkersModule> | undefined;
37+
let diagnosticStabilityBundleModule: Promise<DiagnosticStabilityBundleModule> | undefined;
38+
let bundledRuntimeDepsActivityModule: Promise<BundledRuntimeDepsActivityModule> | undefined;
39+
let commandQueueModule: Promise<CommandQueueModule> | undefined;
40+
let runtimeInternalModule: Promise<RuntimeInternalModule> | undefined;
41+
42+
const loadEmbeddedRunsModule = () =>
43+
(embeddedRunsModule ??= import("../../agents/pi-embedded-runner/runs.js"));
44+
const loadRuntimeConfigModule = () => (runtimeConfigModule ??= import("../../config/config.js"));
45+
const loadProcessRespawnModule = () =>
46+
(processRespawnModule ??= import("../../infra/process-respawn.js"));
47+
const loadRestartSentinelModule = () =>
48+
(restartSentinelModule ??= import("../../infra/restart-sentinel.js"));
49+
const loadRestartModule = () => (restartModule ??= import("../../infra/restart.js"));
50+
const loadSupervisorMarkersModule = () =>
51+
(supervisorMarkersModule ??= import("../../infra/supervisor-markers.js"));
52+
const loadDiagnosticStabilityBundleModule = () =>
53+
(diagnosticStabilityBundleModule ??= import("../../logging/diagnostic-stability-bundle.js"));
54+
const loadBundledRuntimeDepsActivityModule = () =>
55+
(bundledRuntimeDepsActivityModule ??= import("../../plugins/bundled-runtime-deps-activity.js"));
56+
const loadCommandQueueModule = () =>
57+
(commandQueueModule ??= import("../../process/command-queue.js"));
58+
const loadRuntimeInternalModule = () =>
59+
(runtimeInternalModule ??= import("../../tasks/runtime-internal.js"));
60+
61+
function createRestartIterationHook(onRestart: () => Promise<void> | void): () => Promise<boolean> {
62+
let isFirstIteration = true;
63+
return async () => {
64+
if (isFirstIteration) {
65+
isFirstIteration = false;
66+
return false;
67+
}
68+
await onRestart();
69+
return true;
70+
};
71+
}
72+
5273
async function waitForGatewayPortReady(host: string, port: number): Promise<boolean> {
5374
return await new Promise<boolean>((resolve) => {
5475
const socket = net.createConnection({ host, port });
@@ -114,7 +135,9 @@ export async function runGatewayLoop(params: {
114135
cleanupSignals();
115136
params.runtime.exit(code);
116137
};
117-
const writeStabilityBundle = (reason: string, error?: unknown) => {
138+
const writeStabilityBundle = async (reason: string, error?: unknown) => {
139+
const { writeDiagnosticStabilityBundleForFailureSync } =
140+
await loadDiagnosticStabilityBundleModule();
118141
const result = writeDiagnosticStabilityBundleForFailureSync(reason, error);
119142
if ("message" in result) {
120143
gatewayLog.warn(result.message);
@@ -142,6 +165,10 @@ export async function runGatewayLoop(params: {
142165
const handleRestartAfterServerClose = async (restartReason?: string) => {
143166
const hadLock = await releaseLockIfHeld();
144167
const isUpdateRestart = restartReason === "update.run";
168+
const { respawnGatewayProcessForUpdate, restartGatewayProcessWithFreshPid } =
169+
await loadProcessRespawnModule();
170+
const { detectRespawnSupervisor } = await loadSupervisorMarkersModule();
171+
const { markUpdateRestartSentinelFailure } = await loadRestartSentinelModule();
145172

146173
if (isUpdateRestart) {
147174
const respawn = respawnGatewayProcessForUpdate();
@@ -228,7 +255,7 @@ export async function runGatewayLoop(params: {
228255
return;
229256
}
230257
if (respawn.mode === "failed") {
231-
writeStabilityBundle("gateway.restart_respawn_failed");
258+
await writeStabilityBundle("gateway.restart_respawn_failed");
232259
gatewayLog.warn(
233260
`full process restart failed (${respawn.detail ?? "unknown error"}); falling back to in-process restart`,
234261
);
@@ -250,8 +277,9 @@ export async function runGatewayLoop(params: {
250277

251278
const SUPERVISOR_STOP_TIMEOUT_MS = 30_000;
252279
const SHUTDOWN_TIMEOUT_MS = SUPERVISOR_STOP_TIMEOUT_MS - 5_000;
253-
const resolveRestartDrainTimeoutMs = (): RestartDrainTimeoutMs => {
280+
const resolveRestartDrainTimeoutMs = async (): Promise<RestartDrainTimeoutMs> => {
254281
try {
282+
const { getRuntimeConfig } = await loadRuntimeConfigModule();
255283
const timeoutMs = getRuntimeConfig().gateway?.reload?.deferralTimeoutMs;
256284
return typeof timeoutMs === "number" && Number.isFinite(timeoutMs) && timeoutMs > 0
257285
? timeoutMs
@@ -268,7 +296,6 @@ export async function runGatewayLoop(params: {
268296
}
269297
shuttingDown = true;
270298
const isRestart = action === "restart";
271-
const restartDrainTimeoutMs = isRestart ? resolveRestartDrainTimeoutMs() : 0;
272299
gatewayLog.info(`received ${signal}; ${isRestart ? "restarting" : "shutting down"}`);
273300

274301
let forceExitTimer: ReturnType<typeof setTimeout> | null = null;
@@ -278,13 +305,18 @@ export async function runGatewayLoop(params: {
278305
}
279306
forceExitTimer = setTimeout(() => {
280307
gatewayLog.error("shutdown timed out; exiting without full cleanup");
281-
writeStabilityBundle(
282-
isRestart ? "gateway.restart_shutdown_timeout" : "gateway.stop_shutdown_timeout",
283-
);
284-
// Keep the in-process watchdog below the supervisor stop budget so this
285-
// path wins before launchd/systemd escalates to a hard kill. Exit
286-
// non-zero on any timeout so supervised installs restart cleanly.
287-
exitProcess(1);
308+
void (async () => {
309+
try {
310+
await writeStabilityBundle(
311+
isRestart ? "gateway.restart_shutdown_timeout" : "gateway.stop_shutdown_timeout",
312+
);
313+
} finally {
314+
// Keep the in-process watchdog below the supervisor stop budget so this
315+
// path wins before launchd/systemd escalates to a hard kill. Exit
316+
// non-zero on any timeout so supervised installs restart cleanly.
317+
exitProcess(1);
318+
}
319+
})();
288320
}, forceExitMs);
289321
};
290322
const clearForceExitTimer = () => {
@@ -295,35 +327,45 @@ export async function runGatewayLoop(params: {
295327
forceExitTimer = null;
296328
};
297329

298-
if (!isRestart) {
299-
armForceExitTimer(SHUTDOWN_TIMEOUT_MS);
300-
} else if (restartDrainTimeoutMs !== undefined) {
301-
// Allow extra time for draining active turns on explicitly capped restarts.
302-
armForceExitTimer(restartDrainTimeoutMs + SHUTDOWN_TIMEOUT_MS);
303-
}
304-
305-
const formatRestartDrainBudget = () =>
306-
restartDrainTimeoutMs === undefined
307-
? "without a timeout"
308-
: `with timeout ${restartDrainTimeoutMs}ms`;
309-
const createStillPendingDrainLogger = () =>
310-
setInterval(() => {
311-
gatewayLog.warn(
312-
`still draining ${getActiveTaskCount()} active task(s), ${getActiveEmbeddedRunCount()} active embedded run(s), and ${getActiveBundledRuntimeDepsInstallCount()} runtime deps install(s) before restart`,
313-
);
314-
}, RESTART_DRAIN_STILL_PENDING_WARN_MS);
315-
316-
const armCloseForceExitTimerForIndefiniteRestart = () => {
317-
if (isRestart && restartDrainTimeoutMs === undefined) {
330+
void (async () => {
331+
const restartDrainTimeoutMs = isRestart ? await resolveRestartDrainTimeoutMs() : 0;
332+
if (!isRestart) {
318333
armForceExitTimer(SHUTDOWN_TIMEOUT_MS);
334+
} else if (restartDrainTimeoutMs !== undefined) {
335+
// Allow extra time for draining active turns on explicitly capped restarts.
336+
armForceExitTimer(restartDrainTimeoutMs + SHUTDOWN_TIMEOUT_MS);
319337
}
320-
};
321338

322-
void (async () => {
339+
const formatRestartDrainBudget = () =>
340+
restartDrainTimeoutMs === undefined
341+
? "without a timeout"
342+
: `with timeout ${restartDrainTimeoutMs}ms`;
343+
const armCloseForceExitTimerForIndefiniteRestart = () => {
344+
if (isRestart && restartDrainTimeoutMs === undefined) {
345+
armForceExitTimer(SHUTDOWN_TIMEOUT_MS);
346+
}
347+
};
348+
323349
try {
324350
// On restart, wait for in-flight agent turns to finish before
325351
// tearing down the server so buffered messages are delivered.
326352
if (isRestart) {
353+
const [
354+
{ abortEmbeddedPiRun, getActiveEmbeddedRunCount, waitForActiveEmbeddedRuns },
355+
{ getActiveBundledRuntimeDepsInstallCount, waitForBundledRuntimeDepsInstallIdle },
356+
{ getActiveTaskCount, markGatewayDraining, waitForActiveTasks },
357+
] = await Promise.all([
358+
loadEmbeddedRunsModule(),
359+
loadBundledRuntimeDepsActivityModule(),
360+
loadCommandQueueModule(),
361+
]);
362+
const createStillPendingDrainLogger = () =>
363+
setInterval(() => {
364+
gatewayLog.warn(
365+
`still draining ${getActiveTaskCount()} active task(s), ${getActiveEmbeddedRunCount()} active embedded run(s), and ${getActiveBundledRuntimeDepsInstallCount()} runtime deps install(s) before restart`,
366+
);
367+
}, RESTART_DRAIN_STILL_PENDING_WARN_MS);
368+
327369
// Reject new enqueues immediately during the drain window so
328370
// sessions get an explicit restart error instead of silent task loss.
329371
markGatewayDraining();
@@ -385,48 +427,69 @@ export async function runGatewayLoop(params: {
385427

386428
const onSigterm = () => {
387429
gatewayLog.info("signal SIGTERM received");
388-
request(consumeGatewayRestartIntentSync() ? "restart" : "stop", "SIGTERM");
430+
void (async () => {
431+
const { consumeGatewayRestartIntentSync } = await loadRestartModule();
432+
request(consumeGatewayRestartIntentSync() ? "restart" : "stop", "SIGTERM");
433+
})();
389434
};
390435
const onSigint = () => {
391436
gatewayLog.info("signal SIGINT received");
392437
request("stop", "SIGINT");
393438
};
394439
const onSigusr1 = () => {
395440
gatewayLog.info("signal SIGUSR1 received");
396-
const authorized = consumeGatewaySigusr1RestartAuthorization();
397-
if (!authorized) {
398-
if (!isGatewaySigusr1RestartExternallyAllowed()) {
399-
gatewayLog.warn(
400-
"SIGUSR1 restart ignored (not authorized; commands.restart=false or use gateway tool).",
401-
);
402-
return;
403-
}
404-
if (shuttingDown) {
405-
gatewayLog.info("received SIGUSR1 during shutdown; ignoring");
441+
void (async () => {
442+
const {
443+
consumeGatewaySigusr1RestartAuthorization,
444+
isGatewaySigusr1RestartExternallyAllowed,
445+
markGatewaySigusr1RestartHandled,
446+
peekGatewaySigusr1RestartReason,
447+
scheduleGatewaySigusr1Restart,
448+
} = await loadRestartModule();
449+
const authorized = consumeGatewaySigusr1RestartAuthorization();
450+
if (!authorized) {
451+
if (!isGatewaySigusr1RestartExternallyAllowed()) {
452+
gatewayLog.warn(
453+
"SIGUSR1 restart ignored (not authorized; commands.restart=false or use gateway tool).",
454+
);
455+
return;
456+
}
457+
if (shuttingDown) {
458+
gatewayLog.info("received SIGUSR1 during shutdown; ignoring");
459+
return;
460+
}
461+
// External SIGUSR1 requests should still reuse the in-process restart
462+
// scheduler so idle drain and restart coalescing stay consistent.
463+
scheduleGatewaySigusr1Restart({ delayMs: 0, reason: "SIGUSR1" });
406464
return;
407465
}
408-
// External SIGUSR1 requests should still reuse the in-process restart
409-
// scheduler so idle drain and restart coalescing stay consistent.
410-
scheduleGatewaySigusr1Restart({ delayMs: 0, reason: "SIGUSR1" });
411-
return;
412-
}
413-
const restartReason = peekGatewaySigusr1RestartReason();
414-
markGatewaySigusr1RestartHandled();
415-
request("restart", "SIGUSR1", restartReason);
466+
const restartReason = peekGatewaySigusr1RestartReason();
467+
markGatewaySigusr1RestartHandled();
468+
request("restart", "SIGUSR1", restartReason);
469+
})();
416470
};
417471

418472
process.on("SIGTERM", onSigterm);
419473
process.on("SIGINT", onSigint);
420474
process.on("SIGUSR1", onSigusr1);
421475

422476
try {
423-
const onIteration = createRestartIterationHook(() => {
477+
const onIteration = createRestartIterationHook(async () => {
424478
// After an in-process restart (SIGUSR1), reset command-queue lane state.
425479
// Interrupted tasks from the previous lifecycle may have left `active`
426480
// counts elevated (their finally blocks never ran), permanently blocking
427481
// new work from draining. The same boundary also discards stale restart
428482
// deferral timers and reloads the task registry from durable state so
429483
// cancelled/completed work is not kept alive by old in-memory maps.
484+
const [
485+
{ resetAllLanes },
486+
{ resetGatewayRestartStateForInProcessRestart },
487+
{ reloadTaskRegistryFromStore },
488+
] = await Promise.all([
489+
loadCommandQueueModule(),
490+
loadRestartModule(),
491+
loadRuntimeInternalModule(),
492+
]);
430493
resetAllLanes();
431494
resetGatewayRestartStateForInProcessRestart();
432495
reloadTaskRegistryFromStore();
@@ -436,7 +499,7 @@ export async function runGatewayLoop(params: {
436499
// SIGTERM/SIGINT still exit after a graceful shutdown.
437500
let isFirstStart = true;
438501
for (;;) {
439-
onIteration();
502+
await onIteration();
440503
try {
441504
server = await params.start({ startupStartedAt });
442505
isFirstStart = false;
@@ -456,7 +519,7 @@ export async function runGatewayLoop(params: {
456519
await releaseLockIfHeld();
457520
const errMsg = formatErrorMessage(err);
458521
const errStack = err instanceof Error && err.stack ? `\n${err.stack}` : "";
459-
writeStabilityBundle("gateway.restart_startup_failed", err);
522+
await writeStabilityBundle("gateway.restart_startup_failed", err);
460523
gatewayLog.error(
461524
`gateway startup failed: ${errMsg}. ` +
462525
`Process will stay alive; fix the issue and restart.${errStack}`,

0 commit comments

Comments
 (0)