Skip to content

Commit 90b6437

Browse files
nicktrnnicktrn
committed
add shutdown manager
1 parent a6e85f0 commit 90b6437

File tree

4 files changed

+268
-0
lines changed

4 files changed

+268
-0
lines changed

‎packages/core/src/v3/serverOnly/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
11
export * from "./checkpointClient.js";
22
export * from "./checkpointTest.js";
33
export * from "./httpServer.js";
4+
export * from "./singleton.js";
5+
export * from "./shutdownManager.js";

‎packages/core/src/v3/serverOnly/shutdownManager.test.ts‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
import { describe, test, expect, vi, beforeEach } from "vitest";
2+
import { shutdownManager } from "./shutdownManager.js";
3+
4+
// Type assertion to access private members for testing
5+
type PrivateShutdownManager = {
6+
handlers: Map<string, { handler: NodeJS.SignalsListener; signals: Array<"SIGTERM" | "SIGINT"> }>;
7+
shutdown: (signal: "SIGTERM" | "SIGINT") => Promise<void>;
8+
_reset: () => void;
9+
};
10+
11+
describe("ShutdownManager", { concurrent: false }, () => {
12+
const manager = shutdownManager as unknown as PrivateShutdownManager;
13+
// Mock process.exit to prevent actual exit
14+
const mockExit = vi.spyOn(process, "exit").mockImplementation(() => undefined as never);
15+
16+
beforeEach(() => {
17+
// Clear all mocks and reset the manager before each test
18+
vi.clearAllMocks();
19+
manager._reset();
20+
});
21+
22+
test("should successfully register a new handler", () => {
23+
const handler = vi.fn();
24+
shutdownManager.register("test-handler", handler);
25+
26+
expect(manager.handlers.has("test-handler")).toBe(true);
27+
const registeredHandler = manager.handlers.get("test-handler");
28+
expect(registeredHandler?.handler).toBe(handler);
29+
expect(registeredHandler?.signals).toEqual(["SIGTERM", "SIGINT"]);
30+
});
31+
32+
test("should throw error when registering duplicate handler name", () => {
33+
const handler = vi.fn();
34+
shutdownManager.register("duplicate-handler", handler);
35+
36+
expect(() => {
37+
shutdownManager.register("duplicate-handler", handler);
38+
}).toThrow('Shutdown handler "duplicate-handler" already registered');
39+
});
40+
41+
test("should register handler with custom signals", () => {
42+
const handler = vi.fn();
43+
shutdownManager.register("custom-signals", handler, ["SIGTERM"]);
44+
45+
const registeredHandler = manager.handlers.get("custom-signals");
46+
expect(registeredHandler?.signals).toEqual(["SIGTERM"]);
47+
});
48+
49+
test("should call registered handlers when shutdown is triggered", async () => {
50+
const handler1 = vi.fn();
51+
const handler2 = vi.fn();
52+
53+
shutdownManager.register("handler1", handler1);
54+
shutdownManager.register("handler2", handler2);
55+
56+
await manager.shutdown("SIGTERM");
57+
58+
expect(handler1).toHaveBeenCalledWith("SIGTERM");
59+
expect(handler2).toHaveBeenCalledWith("SIGTERM");
60+
expect(mockExit).toHaveBeenCalledWith(128 + 15); // SIGTERM number
61+
});
62+
63+
test("should only call handlers registered for specific signal", async () => {
64+
const handler1 = vi.fn();
65+
const handler2 = vi.fn();
66+
67+
shutdownManager.register("handler1", handler1, ["SIGTERM"]);
68+
shutdownManager.register("handler2", handler2, ["SIGINT"]);
69+
70+
await manager.shutdown("SIGTERM");
71+
72+
expect(handler1).toHaveBeenCalledWith("SIGTERM");
73+
expect(handler2).not.toHaveBeenCalled();
74+
expect(mockExit).toHaveBeenCalledWith(128 + 15);
75+
});
76+
77+
test("should handle errors in shutdown handlers gracefully", async () => {
78+
const successHandler = vi.fn();
79+
const errorHandler = vi.fn().mockRejectedValue(new Error("Handler failed"));
80+
81+
shutdownManager.register("success-handler", successHandler);
82+
shutdownManager.register("error-handler", errorHandler);
83+
84+
await manager.shutdown("SIGTERM");
85+
86+
expect(successHandler).toHaveBeenCalledWith("SIGTERM");
87+
expect(errorHandler).toHaveBeenCalledWith("SIGTERM");
88+
expect(mockExit).toHaveBeenCalledWith(128 + 15);
89+
});
90+
91+
test("should only run shutdown sequence once even if called multiple times", async () => {
92+
const handler = vi.fn();
93+
shutdownManager.register("test-handler", handler);
94+
95+
await Promise.all([manager.shutdown("SIGTERM"), manager.shutdown("SIGTERM")]);
96+
97+
expect(handler).toHaveBeenCalledTimes(1);
98+
expect(mockExit).toHaveBeenCalledTimes(1);
99+
expect(mockExit).toHaveBeenCalledWith(128 + 15);
100+
});
101+
102+
test("should exit with correct signal number", async () => {
103+
const handler = vi.fn();
104+
shutdownManager.register("test-handler", handler);
105+
106+
await manager.shutdown("SIGINT");
107+
expect(mockExit).toHaveBeenCalledWith(128 + 2); // SIGINT number
108+
109+
vi.clearAllMocks();
110+
manager._reset();
111+
shutdownManager.register("test-handler", handler);
112+
113+
await manager.shutdown("SIGTERM");
114+
expect(mockExit).toHaveBeenCalledWith(128 + 15); // SIGTERM number
115+
});
116+
117+
test("should only exit after all handlers have finished", async () => {
118+
const sequence: string[] = [];
119+
120+
const handler1 = vi.fn().mockImplementation(async () => {
121+
sequence.push("handler1 start");
122+
await new Promise((resolve) => setTimeout(resolve, 10));
123+
sequence.push("handler1 end");
124+
});
125+
126+
const handler2 = vi.fn().mockImplementation(async () => {
127+
sequence.push("handler2 start");
128+
await new Promise((resolve) => setTimeout(resolve, 20));
129+
sequence.push("handler2 end");
130+
});
131+
132+
const handler3 = vi.fn().mockImplementation(async () => {
133+
sequence.push("handler3 start");
134+
await new Promise((resolve) => setTimeout(resolve, 5));
135+
sequence.push("handler3 end");
136+
});
137+
138+
// Store the current mock implementation
139+
const currentExit = mockExit.getMockImplementation();
140+
141+
// Override with our sequence-tracking implementation
142+
mockExit.mockImplementation((code?: number | string | null) => {
143+
sequence.push("exit");
144+
return undefined as never;
145+
});
146+
147+
shutdownManager.register("handler1", handler1);
148+
shutdownManager.register("handler2", handler2);
149+
shutdownManager.register("handler3", handler3);
150+
151+
await manager.shutdown("SIGTERM");
152+
153+
// Verify the execution order
154+
expect(sequence).toEqual([
155+
"handler1 start",
156+
"handler2 start",
157+
"handler3 start",
158+
"handler3 end",
159+
"handler1 end",
160+
"handler2 end",
161+
"exit",
162+
]);
163+
164+
// Verify the handlers were called with correct arguments
165+
expect(handler1).toHaveBeenCalledWith("SIGTERM");
166+
expect(handler2).toHaveBeenCalledWith("SIGTERM");
167+
expect(handler3).toHaveBeenCalledWith("SIGTERM");
168+
expect(mockExit).toHaveBeenCalledWith(128 + 15);
169+
170+
// Restore original mock implementation
171+
if (currentExit) {
172+
mockExit.mockImplementation(currentExit);
173+
}
174+
});
175+
});

‎packages/core/src/v3/serverOnly/shutdownManager.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import { singleton } from "./singleton.js";
2+
3+
type ShutdownHandler = NodeJS.SignalsListener;
4+
// We intentionally keep these limited to avoid unexpected issues with signal handling
5+
type ShutdownSignal = Extract<NodeJS.Signals, "SIGTERM" | "SIGINT">;
6+
7+
class ShutdownManager {
8+
private isShuttingDown = false;
9+
private signalNumbers: Record<ShutdownSignal, number> = {
10+
SIGINT: 2,
11+
SIGTERM: 15,
12+
};
13+
14+
private handlers: Map<string, { handler: ShutdownHandler; signals: ShutdownSignal[] }> =
15+
new Map();
16+
17+
constructor() {
18+
process.on("SIGTERM", () => this.shutdown("SIGTERM"));
19+
process.on("SIGINT", () => this.shutdown("SIGINT"));
20+
}
21+
22+
register(
23+
name: string,
24+
handler: ShutdownHandler,
25+
signals: ShutdownSignal[] = ["SIGTERM", "SIGINT"]
26+
) {
27+
if (this.handlers.has(name)) {
28+
throw new Error(`Shutdown handler "${name}" already registered`);
29+
}
30+
this.handlers.set(name, { handler, signals });
31+
}
32+
33+
async shutdown(signal: ShutdownSignal) {
34+
if (this.isShuttingDown) return;
35+
this.isShuttingDown = true;
36+
37+
console.log(`\nReceived ${signal}. Starting graceful shutdown...`);
38+
39+
// Get handlers that are registered for this signal
40+
const handlersToRun = Array.from(this.handlers.entries()).filter(([_, { signals }]) =>
41+
signals.includes(signal)
42+
);
43+
44+
try {
45+
const results = await Promise.allSettled(
46+
handlersToRun.map(async ([name, { handler }]) => {
47+
try {
48+
console.log(`Running shutdown handler: ${name}`);
49+
await handler(signal);
50+
console.log(`Shutdown handler completed: ${name}`);
51+
} catch (error) {
52+
console.error(`Shutdown handler failed: ${name}`, error);
53+
throw error;
54+
}
55+
})
56+
);
57+
58+
// Log any failures
59+
results.forEach((result, index) => {
60+
if (result.status === "rejected") {
61+
const handlerEntry = handlersToRun[index];
62+
if (handlerEntry) {
63+
const [name] = handlerEntry;
64+
console.error(`Shutdown handler "${name}" failed:`, result.reason);
65+
}
66+
}
67+
});
68+
} catch (error) {
69+
console.error("Error during shutdown:", error);
70+
} finally {
71+
// Exit with the correct signal number
72+
process.exit(128 + this.signalNumbers[signal]);
73+
}
74+
}
75+
76+
// For testing purposes only - keep this
77+
private _reset() {
78+
this.isShuttingDown = false;
79+
this.handlers.clear();
80+
}
81+
}
82+
83+
export const shutdownManager = singleton("shutdownManager", () => new ShutdownManager());

‎packages/core/src/v3/serverOnly/singleton.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
export function singleton<T>(name: string, getValue: () => T): T {
2+
const thusly = globalThis as unknown as {
3+
__trigger_singletons: Record<string, T>;
4+
};
5+
thusly.__trigger_singletons ??= {};
6+
thusly.__trigger_singletons[name] ??= getValue();
7+
return thusly.__trigger_singletons[name];
8+
}

0 commit comments

Comments
 (0)