1import http from "node:http";
2
3const machine = process.env.FLY_MACHINE_ID;
4const version = process.env.VERSION;
5const mode = process.env.MODE || "drain";
6const afterSignal = Number(process.env.AFTER_SIGNAL_MS || "1000");
7const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
8const call = async (operation, args = []) => {
9 const response = await fetch(process.env.LEDGER_URL, {
10 method: "POST",
11 headers: {
12 authorization: `Bearer ${process.env.LEDGER_TOKEN}`,
13 "content-type": "application/json",
14 },
15 body: JSON.stringify({ operation, args }),
16 signal: AbortSignal.timeout(5000),
17 });
18 if (!response.ok) throw new Error(`ledger HTTP ${response.status}`);
19 return (await response.json()).result;
20};
21const event = (name, fields = {}) =>
22 call("event", [JSON.stringify({ event: name, machine, version, ...fields })]);
23let resolveStop;
24const stopping = new Promise((resolve) => {
25 resolveStop = resolve;
26});
27const jobs = new Set();
28let workerClientOpen = true;
29const workerCall = (operation, args) => {
30 if (!workerClientOpen)
31 throw new Error("worker queue client already released");
32 return call(operation, args);
33};
34const requestFinalizers = [];
35let accepting = true;
36let poll;
37let started = false;
38const server = http.createServer(async (request, response) => {
39 try {
40 if (request.url === "/health") return response.end("ready");
41 if (request.url === "/")
42 return response.end(JSON.stringify({ machine, version }));
43 await event("request-started");
44 response.once("finish", () =>
45 requestFinalizers.push(event("request-finalized")),
46 );
47 if (request.url === "/stream") response.write("first\n".repeat(32768));
48 await stopping;
49 await sleep(afterSignal);
50 await event("response-finished");
51 response.end(
52 request.url === "/stream"
53 ? "last\n".repeat(32768)
54 : JSON.stringify({ machine, version }),
55 );
56 } catch (error) {
57 console.error(error);
58 response.destroy();
59 }
60});
61const processJob = async (job) => {
62 if (!job.checkpoint && job.kind !== "quick") {
63 await stopping;
64 if (job.kind === "checkpoint") {
65 await workerCall("checkpoint", [job.id, job.job, machine]);
66 return;
67 }
68 await sleep(afterSignal);
69 } else await sleep(100);
70 await workerCall("finish", [job.id, job.job, machine]);
71};
72const acquire = async () => {
73 if (!accepting) return;
74 await workerCall("tick", [machine, version]);
75 const result = await workerCall("claim", [machine, version]);
76 if (!result) return;
77 const job = JSON.parse(result);
78 const running = processJob(job)
79 .catch(async (error) => {
80 process.exitCode = 1;
81 console.error(error);
82 await workerCall("checkpoint", [job.id, job.job, machine]);
83 })
84 .finally(() => {
85 jobs.delete(running);
86 });
87 jobs.add(running);
88};
89const timer = setInterval(() => {
90 if (poll || !accepting) return;
91 poll = acquire()
92 .catch((error) => {
93 console.error(error);
94 process.exitCode = 1;
95 })
96 .finally(() => {
97 poll = undefined;
98 });
99}, 250);
100const shutdown = async (signal) => {
101 if (started) return;
102 started = true;
103 setTimeout(
104 () => process.exit(1),
105 Number(process.env.SHUTDOWN_MS || "30000") * 0.9,
106 );
107 const httpClosed = new Promise((resolve) => server.close(resolve));
108 await event("stop-started", { worker: "a", signal });
109 if (mode === "stop-delay") await sleep(1000);
110 accepting = false;
111 clearInterval(timer);
112 await poll;
113 await event("stopped", { worker: "a" });
114 resolveStop();
115 await Promise.all([...jobs]);
116 await event("drained", { worker: "a" });
117 await event("work-closed", { worker: "a" });
118 await event("client-released", { worker: "a" });
119 workerClientOpen = false;
120 await httpClosed;
121 await Promise.all(requestFinalizers);
122 await event("shared-closed");
123 process.exit(process.exitCode || 0);
124};
125for (const signal of ["SIGTERM", "SIGINT", "SIGQUIT"])
126 process.on(signal, () => {
127 void shutdown(signal).catch((error) => {
128 console.error(error);
129 process.exit(1);
130 });
131 });
132await event("worker-ready", { worker: "a" });
133server.listen(3000, "::");