fix(relay): order replacement sockets explicitly

Wall-clock timestamps can tie during rapid daemon socket replacement. Persist a per-connection generation so forwarding ownership follows acceptance order, including when replacing pre-generation sockets during rollout.
This commit is contained in:
Mohamed Boudra
2026-08-12 10:04:16 +00:00
parent c4cb6aed30
commit ca0e14e549
3 changed files with 52 additions and 7 deletions
+36 -1
View File
@@ -246,7 +246,8 @@ describe("RelayDurableObject data forwarding", () => {
role: "server",
connectionId,
serverId: "srv_test",
createdAt: index + 1,
generation: index + 1,
createdAt: 1,
}),
);
const { state, setTagSockets } = createMockState();
@@ -261,6 +262,40 @@ describe("RelayDurableObject data forwarding", () => {
expect(daemonSockets.at(-1)?.send).toHaveBeenCalledOnce();
expect(daemonSockets.at(-1)?.send).toHaveBeenCalledWith("create_agent_request");
});
it("prefers a generated replacement over rollout sockets without generations", () => {
const connectionId = "clt_rollout";
const client = createMockSocket({
version: "2",
role: "client",
connectionId,
serverId: "srv_test",
createdAt: 1,
});
const staleSocket = createMockSocket({
version: "2",
role: "server",
connectionId,
serverId: "srv_test",
createdAt: 2,
});
const replacementSocket = createMockSocket({
version: "2",
role: "server",
connectionId,
serverId: "srv_test",
generation: 1,
createdAt: 2,
});
const { state, setTagSockets } = createMockState();
setTagSockets(`server:${connectionId}`, [staleSocket, replacementSocket]);
const relay = new RelayDurableObject(state as unknown as DurableObjectStateArg);
relay.webSocketMessage(client as unknown as WebSocket, "create_agent_request");
expect(staleSocket.send).not.toHaveBeenCalled();
expect(replacementSocket.send).toHaveBeenCalledWith("create_agent_request");
});
});
describe("relay worker endpoint routing", () => {
+14 -6
View File
@@ -86,25 +86,29 @@ function getString(record: Record<string, unknown>, key: string): string | undef
return typeof value === "string" ? value : undefined;
}
function getSocketCreatedAt(ws: WebSocket): number {
function getSocketGeneration(ws: WebSocket): number {
const attachment = deserializeAttachment(ws);
if (!isRecord(attachment)) return 0;
return typeof attachment.createdAt === "number" ? attachment.createdAt : 0;
return typeof attachment.generation === "number" ? attachment.generation : 0;
}
function selectNewestSocket(sockets: WebSocket[]): WebSocket | null {
let newest: WebSocket | null = null;
let newestCreatedAt = Number.NEGATIVE_INFINITY;
let newestGeneration = Number.NEGATIVE_INFINITY;
for (const socket of sockets) {
const createdAt = getSocketCreatedAt(socket);
if (createdAt >= newestCreatedAt) {
const generation = getSocketGeneration(socket);
if (generation > newestGeneration) {
newest = socket;
newestCreatedAt = createdAt;
newestGeneration = generation;
}
}
return newest;
}
function nextSocketGeneration(sockets: WebSocket[]): number {
return sockets.reduce((highest, socket) => Math.max(highest, getSocketGeneration(socket)), 0) + 1;
}
function getGlobalWebSocketPair(): (new () => WebSocketPair) | undefined {
// Access WebSocketPair from global scope (Cloudflare Workers runtime)
// Use Reflect to access global property without type assertions
@@ -371,6 +375,9 @@ export class RelayDurableObject {
const isServerControl = role === "server" && !resolvedConnectionId;
const isServerData = role === "server" && !!resolvedConnectionId;
const serverDataGeneration = isServerData
? nextSocketGeneration(this.state.getWebSockets(`server:${resolvedConnectionId}`))
: undefined;
// Close any existing server-side connection with the same identity.
// - server-control: single per serverId
@@ -396,6 +403,7 @@ export class RelayDurableObject {
role,
version: CURRENT_RELAY_VERSION,
connectionId: resolvedConnectionId || null,
...(serverDataGeneration === undefined ? {} : { generation: serverDataGeneration }),
createdAt: Date.now(),
};
serializeAttachment(server, attachment);
+2
View File
@@ -24,5 +24,7 @@ export interface RelaySessionAttachment {
* independent socket + E2EE channel per connected connection.
*/
connectionId?: string | null;
/** Monotonic acceptance order for v2 daemon data sockets sharing a connectionId. */
generation?: number;
createdAt: number;
}