cherry-pick: feat(replication): gracefully drain in-flight blob sends before worker shutdown (#529 → v5.1)#552
Conversation
feat(replication): gracefully drain in-flight blob sends before worker shutdown
There was a problem hiding this comment.
Code Review
This pull request introduces a per-worker tracker (blobSendDrain.ts) to gracefully drain in-flight replication blob sends before a worker shuts down during a restart, preventing peer copy divergence. It integrates this tracking into the replication connection logic and registers the shutdown drain. The review feedback highlights two important issues: first, waitForDrainOrSocketEnd could hang indefinitely if the socket is already destroyed or not writable when called, leading to a leak of the drain token; second, a synchronous error in blob.stream() would bypass the finally block and permanently leak the blob ID in blobsBeingSent. Addressing these edge cases with defensive checks and error handling will make the implementation more robust.
| export function waitForDrainOrSocketEnd(socket: EventEmitter, ws: EventEmitter): Promise<void> { | ||
| return new Promise<void>((resolve) => { | ||
| const done = () => { | ||
| socket.off('drain', done); | ||
| socket.off('close', done); | ||
| socket.off('error', done); | ||
| ws.off('close', done); | ||
| resolve(); | ||
| }; | ||
| socket.once('drain', done); | ||
| socket.once('close', done); | ||
| socket.once('error', done); | ||
| ws.once('close', done); | ||
| }); | ||
| } |
There was a problem hiding this comment.
If the underlying socket is already destroyed or no longer writable when waitForDrainOrSocketEnd is called, none of the registered event listeners ('drain', 'close', 'error') will ever fire. This will cause the returned promise to hang indefinitely, preventing the finally block in sendBlobs from running and leading to a permanent leak of the drain token and outstandingBlobsBeingSent counter.
Adding a defensive check to resolve immediately if the socket is already destroyed or not writable prevents this potential hang.
export function waitForDrainOrSocketEnd(socket: EventEmitter, ws: EventEmitter): Promise<void> {
if ((socket as any).destroyed || !(socket as any).writable) {
return Promise.resolve();
}
return new Promise<void>((resolve) => {
const done = () => {
socket.off('drain', done);
socket.off('close', done);
socket.off('error', done);
ws.off('close', done);
resolve();
};
socket.once('drain', done);
socket.once('close', done);
socket.once('error', done);
ws.once('close', done);
});
}| const iterator = blob.stream()[Symbol.asyncIterator](); | ||
| // Track this send so a worker restart can gracefully drain it (finish it if it's still making | ||
| // progress) before shutting down, rather than tearing it down mid-stream. See blobSendDrain.ts. | ||
| const drainToken = registerBlobSend(); |
There was a problem hiding this comment.
If blob.stream() throws an error synchronously (for example, if the blob is invalid or has been deleted/evicted before the stream is opened), the error will propagate immediately before entering the try...catch block. Because blobsBeingSent.add(id) has already been called, this synchronous throw will bypass the finally block's cleanup, permanently leaking the blob id in the blobsBeingSent set and preventing this blob from ever being re-sent on this connection.
Wrapping the stream creation in a local try...catch block to clean up blobsBeingSent on failure prevents this resource leak.
let iterator;
try {
iterator = blob.stream()[Symbol.asyncIterator]();
} catch (error) {
blobsBeingSent.delete(id);
throw error;
}
// Track this send so a worker restart can gracefully drain it (finish it if it's still making
// progress) before shutting down, rather than tearing it down mid-stream. See blobSendDrain.ts.
const drainToken = registerBlobSend();|
Reviewed; no blockers found. |
Manual cherry-pick of #529 to
v5.1— the automated bot crashed on the submodule conflict (core pointer +replication/replicator.ts+replication/replicationConnection.tsdiverging fromv5.1) before it could push anything.Resolved by hand:
coresubmodule pointer set to harper's newv5.1tip (includes #1621, needed forcomponents/shutdownDrain.tswhich this PR'sregisterShutdownDrainimport depends on).replication/replicator.ts: dropped an unrelatedimport { redactOperationForLog } from './logRedaction.ts'that feat(replication): gracefully drain in-flight blob sends before worker shutdown #529's branch had absorbed from an unrelated main merge along the way — that file doesn't exist onv5.1and the import was unused in this file. Kept the two genuine feat(replication): gracefully drain in-flight blob sends before worker shutdown #529 imports (registerShutdownDrain,hasProgressingBlobSends/drainBlobSends).replication/replicationConnection.tsauto-merged cleanly, no manual resolution needed.Verified: full build (dist emits fine despite this repo's pre-existing, CI-tolerated tsc baseline errors in
analytics/profile.ts/replicationConnection.ts, unrelated to this change),unitTests/replication/**278 passing, lint/format clean on all touched files (one pre-existing lint warning + one pre-existing format issue inreplicationConnection.tsconfirmed present onv5.1before this cherry-pick, unrelated).