Visitar URL original
fix(supervisor): keep run notification subscriptions across platform … · triggerdotdev/trigger.dev@a9e69ab · GitHub
Skip to content

Commit a9e69ab

Browse files
myftijaTrigger.dev RepoOps
myftija
authored and
Trigger.dev RepoOps
committed
fix(supervisor): keep run notification subscriptions across platform socket reconnects
## Problem The supervisor subscribes to run notifications over its socket to the webapp, one subscription per run it drives. On every reconnect, after a webapp deploy or restart, the webapp forgot those subscriptions and the supervisor never sent them again. If the webapp turned the handshake away while it was still starting up, or dropped the socket itself, socket.io took that as final and the supervisor had no socket at all until its next restart. ## Impact Cancels and completed waits reached the affected runs only through the runner's snapshot poll, every 5 seconds: up to 5 seconds of extra latency on each wait and cancel. In the worst case, a handshake turned away during a deploy, that covered every run on the supervisor until it restarted. Runs were not stuck, lost or run twice. Today's in-image controllers are affected. ## Fix The session keeps the set of runs it has subscribed to and sends the whole set on every connect. Subscribing or unsubscribing while the socket is down only updates the set; the connect handler sends it, so nothing is asked twice and an ask made before the socket exists is no longer dropped. After a refused handshake or a server-side disconnect, the session reconnects itself after a delay that starts at 1 second and doubles up to 30 seconds, resetting on connect. Stopping the session forgets the socket first, so a pending retry cannot revive it. Before: the supervisor connects and subscribes to runs A and B; the webapp restarts; the socket reconnects and sends nothing; cancels for A and B reach their runners only by the poll. After: the same, except that the reconnect sends the subscriptions for A and B again and push notifications resume at once. Mono-RevId: 3e920c637f0bfce37273e0a43987889c142ef3f7
1 parent 592c9bc commit a9e69ab

3 files changed

Lines changed: 198 additions & 6 deletions

File tree

‎.changeset/supervisor-run-notifications-reconnect.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@trigger.dev/core": patch
3+
---
4+
5+
The supervisor keeps its run notification subscriptions across reconnects to the webapp, and reconnects when the webapp turns its handshake away, so cancellations and resolved waits keep reaching running tasks after a webapp restart or deploy.

‎apps/supervisor/src/platformSession.test.ts‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import { once } from "node:events";
2+
import { createServer } from "node:http";
3+
import type { AddressInfo } from "node:net";
4+
import { SupervisorSession } from "@trigger.dev/core/v3/runEngineWorker";
5+
import { Server, type Socket } from "socket.io";
6+
import { afterEach, describe, expect, it } from "vitest";
7+
8+
describe("the platform session's run subscriptions", () => {
9+
const cleanups: Array<() => Promise<unknown> | unknown> = [];
10+
11+
afterEach(async () => {
12+
for (const cleanup of cleanups.splice(0).reverse()) {
13+
await cleanup();
14+
}
15+
});
16+
17+
async function platform() {
18+
const http = createServer((request, response) => {
19+
response.setHeader("Content-Type", "application/json");
20+
if (request.url === "/engine/v1/worker-actions/connect") {
21+
response.end(JSON.stringify({ ok: true, workerGroup: { type: "MANAGED", name: "test" } }));
22+
return;
23+
}
24+
response.end(JSON.stringify({ ok: true }));
25+
});
26+
const io = new Server(http);
27+
type Connection = { socket: Socket; subscribed: string[][]; unsubscribed: string[][] };
28+
const connections: Connection[] = [];
29+
let rejections = 0;
30+
let rejectNext = false;
31+
io.of("/worker").use((_socket, next) => {
32+
if (rejectNext) {
33+
rejectNext = false;
34+
rejections += 1;
35+
return next(new Error("not ready"));
36+
}
37+
next();
38+
});
39+
io.of("/worker").on("connection", (socket) => {
40+
const connection: Connection = { socket, subscribed: [], unsubscribed: [] };
41+
connections.push(connection);
42+
socket.on("run:subscribe", ({ runFriendlyIds }) =>
43+
connection.subscribed.push(runFriendlyIds)
44+
);
45+
socket.on("run:unsubscribe", ({ runFriendlyIds }) =>
46+
connection.unsubscribed.push(runFriendlyIds)
47+
);
48+
});
49+
http.listen(0, "127.0.0.1");
50+
await once(http, "listening");
51+
cleanups.push(() => io.close());
52+
return {
53+
url: `http://127.0.0.1:${(http.address() as AddressInfo).port}`,
54+
connections,
55+
/** The next handshake is turned away, as the platform does while its socket adapter is not ready. */
56+
rejectNextHandshake: () => (rejectNext = true),
57+
rejections: () => rejections,
58+
};
59+
}
60+
61+
async function until(check: () => boolean, what: string, ms = 10_000) {
62+
const deadline = performance.now() + ms;
63+
while (!check() && performance.now() < deadline) {
64+
await new Promise((resolve) => setTimeout(resolve, 10));
65+
}
66+
expect(check(), what).toBe(true);
67+
}
68+
69+
it("asks for the runs subscribed before the socket existed, and again on every reconnect", async () => {
70+
const { url, connections, rejectNextHandshake, rejections } = await platform();
71+
const session = new SupervisorSession({
72+
apiUrl: url,
73+
workerToken: "test-token",
74+
instanceName: "test-worker",
75+
queueConsumerEnabled: false,
76+
runNotificationsEnabled: true,
77+
heartbeatIntervalSeconds: 3600,
78+
dequeueIntervalMs: 10,
79+
dequeueIdleIntervalMs: 10,
80+
scaling: { strategy: "none", minConsumerCount: 1, maxConsumerCount: 1 },
81+
});
82+
cleanups.push(() => session.stop());
83+
84+
session.subscribeToRunNotifications(["run_1", "run_2"]);
85+
await session.start();
86+
await until(() => connections[0]?.subscribed.length === 1, "the first connection got no ask");
87+
expect(connections[0]!.subscribed).toEqual([["run_1", "run_2"]]);
88+
89+
session.subscribeToRunNotifications(["run_3"]);
90+
session.unsubscribeFromRunNotifications(["run_2"]);
91+
await until(() => connections[0]!.unsubscribed.length === 1, "the unsubscribe did not arrive");
92+
expect(connections[0]!.subscribed).toEqual([["run_1", "run_2"], ["run_3"]]);
93+
94+
// The platform drops the connection, as a deploy of its socket servers does; the client reconnects.
95+
connections[0]!.socket.conn.close();
96+
await until(() => connections.length === 2, "the session did not reconnect");
97+
await until(() => connections[1]!.subscribed.length === 1, "the reconnect got no ask");
98+
expect(connections[1]!.subscribed).toEqual([["run_1", "run_3"]]);
99+
100+
session.unsubscribeFromRunNotifications(["run_1", "run_3"]);
101+
await until(() => connections[1]!.unsubscribed.length === 1, "the unsubscribe did not arrive");
102+
connections[1]!.socket.conn.close();
103+
await until(() => connections.length === 3, "the session did not reconnect again");
104+
// Asks arrive in order, so a later one proves nothing was asked for the empty set on connect.
105+
session.subscribeToRunNotifications(["run_4"]);
106+
await until(() => connections[2]!.subscribed.length > 0, "the later ask did not arrive");
107+
expect(connections[2]!.subscribed).toEqual([["run_4"]]);
108+
109+
// The platform turns the reconnect away, which socket.io takes as final; the client is down for sure
110+
// once the rejection went out. What changes meanwhile reaches the next connection as one ask.
111+
rejectNextHandshake();
112+
connections[2]!.socket.conn.close();
113+
await until(() => rejections() === 1, "the handshake was not turned away");
114+
session.subscribeToRunNotifications(["run_5"]);
115+
session.unsubscribeFromRunNotifications(["run_4"]);
116+
await until(
117+
() => connections.length === 4,
118+
"the session did not connect again after the rejection"
119+
);
120+
await until(
121+
() => connections[3]!.subscribed.length === 1,
122+
"the connection after the rejection got no ask"
123+
);
124+
expect(connections[3]!.subscribed).toEqual([["run_5"]]);
125+
expect(connections[3]!.unsubscribed).toEqual([]);
126+
127+
// Disconnected by the platform itself, which socket.io takes as final too.
128+
connections[3]!.socket.disconnect();
129+
await until(
130+
() => connections.length === 5,
131+
"the session did not connect again after being disconnected"
132+
);
133+
await until(
134+
() => connections[4]!.subscribed.length === 1,
135+
"the connection after the disconnect got no ask"
136+
);
137+
expect(connections[4]!.subscribed).toEqual([["run_5"]]);
138+
}, 20_000);
139+
});

‎packages/core/src/v3/runEngineWorker/supervisor/session.ts‎

Lines changed: 54 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,13 +37,19 @@ type SupervisorSessionOptions = SupervisorClientCommonOptions & {
3737
metadata?: Record<string, unknown>;
3838
};
3939

40+
const RECONNECT_DELAY_MS = 1_000;
41+
const RECONNECT_DELAY_MAX_MS = 30_000;
42+
4043
export class SupervisorSession extends EventEmitter<WorkerEvents> {
4144
public readonly httpClient: SupervisorHttpClient;
4245

4346
private readonly logger = new SimpleStructuredLogger("supervisor-session");
4447

4548
private readonly runNotificationsEnabled: boolean;
4649
private runNotificationsSocket?: Socket<WorkerServerToClientEvents, WorkerClientToServerEvents>;
50+
/** The runs subscribed to, asked for again on every connection: the platform keeps subscriptions per connection. */
51+
private readonly runSubscriptions = new Set<string>();
52+
private reconnectDelayMs = RECONNECT_DELAY_MS;
4753

4854
private readonly queueConsumerEnabled: boolean;
4955
private readonly consumerPool: RunQueueConsumerPool;
@@ -111,12 +117,16 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
111117
subscribeToRunNotifications(runFriendlyIds: string[]) {
112118
this.logger.debug("Subscribing to run notifications", { runFriendlyIds });
113119

114-
if (!this.runNotificationsSocket) {
115-
this.logger.error("Socket not connected");
120+
for (const runFriendlyId of runFriendlyIds) {
121+
this.runSubscriptions.add(runFriendlyId);
122+
}
123+
124+
// Not connected, the connect handler asks for the whole set; asking here too would ask twice.
125+
if (!this.runNotificationsSocket?.connected) {
116126
return;
117127
}
118128

119-
this.runNotificationsSocket.emit("run:subscribe", { version: "1", runFriendlyIds });
129+
this.subscribe(this.runNotificationsSocket, runFriendlyIds);
120130

121131
Promise.allSettled(
122132
runFriendlyIds.map((runFriendlyId) =>
@@ -131,8 +141,11 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
131141
unsubscribeFromRunNotifications(runFriendlyIds: string[]) {
132142
this.logger.debug("Unsubscribing from run notifications", { runFriendlyIds });
133143

134-
if (!this.runNotificationsSocket) {
135-
this.logger.error("Socket not connected");
144+
for (const runFriendlyId of runFriendlyIds) {
145+
this.runSubscriptions.delete(runFriendlyId);
146+
}
147+
148+
if (!this.runNotificationsSocket?.connected) {
136149
return;
137150
}
138151

@@ -151,6 +164,13 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
151164
);
152165
}
153166

167+
private subscribe(
168+
socket: Socket<WorkerServerToClientEvents, WorkerClientToServerEvents>,
169+
runFriendlyIds: string[]
170+
) {
171+
socket.emit("run:subscribe", { version: "1", runFriendlyIds });
172+
}
173+
154174
private createRunNotificationsSocket() {
155175
const wsUrl = new URL(this.opts.apiUrl);
156176
wsUrl.pathname = "/worker";
@@ -170,17 +190,43 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
170190
});
171191
socket.on("connect", () => {
172192
this.logger.log("[WS] Connected to platform");
193+
this.reconnectDelayMs = RECONNECT_DELAY_MS;
194+
if (this.runSubscriptions.size > 0) {
195+
this.logger.log("[WS] Subscribing to run notifications on connect", {
196+
runs: this.runSubscriptions.size,
197+
});
198+
this.subscribe(socket, [...this.runSubscriptions]);
199+
}
173200
});
174201
socket.on("connect_error", (error) => {
175202
this.logger.error("[WS] Connection error", { error });
203+
this.reconnectIfGivenUp(socket);
176204
});
177205
socket.on("disconnect", (reason, description) => {
178206
this.logger.log("[WS] Disconnected from platform", { reason, description });
207+
this.reconnectIfGivenUp(socket);
179208
});
180209

181210
return socket;
182211
}
183212

213+
/** socket.io stops after the platform rejects a handshake or disconnects the socket; the session does not. */
214+
private reconnectIfGivenUp(
215+
socket: Socket<WorkerServerToClientEvents, WorkerClientToServerEvents>
216+
) {
217+
if (socket.active || this.runNotificationsSocket !== socket) {
218+
return;
219+
}
220+
const delayMs = this.reconnectDelayMs;
221+
this.reconnectDelayMs = Math.min(delayMs * 2, RECONNECT_DELAY_MAX_MS);
222+
this.logger.log("[WS] Connecting to platform again", { delayMs });
223+
setTimeout(() => {
224+
if (this.runNotificationsSocket === socket) {
225+
socket.connect();
226+
}
227+
}, delayMs).unref();
228+
}
229+
184230
async start() {
185231
const connect = await this.httpClient.connect({
186232
metadata: {
@@ -225,7 +271,9 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
225271
async stop() {
226272
await this.consumerPool.stop();
227273
this.heartbeat.stop();
228-
this.runNotificationsSocket?.disconnect();
274+
const socket = this.runNotificationsSocket;
275+
this.runNotificationsSocket = undefined;
276+
socket?.disconnect();
229277
}
230278

231279
private getHeartbeatBody(): WorkerApiHeartbeatRequestBody {

0 commit comments

Comments
 (0)