From 21a2c3d6f282f525b126669c673cef74d9d6cbbf Mon Sep 17 00:00:00 2001 From: Austin McKinley <54160+amckinley@users.noreply.github.com> Date: Thu, 20 Aug 2026 13:16:31 -0700 Subject: [PATCH] Add graceful shutdown to Blink (#11523) --- bskylink/src/bin.ts | 36 +++++-- bskylink/src/cache/safelinkClient.ts | 135 +++++++++++++++++++++----- bskylink/src/index.ts | 35 +++++-- bskylink/src/safelinkShutdown.test.ts | 118 ++++++++++++++++++++++ bskylink/src/shutdown.test.ts | 63 ++++++++++++ bskylink/src/shutdown.ts | 2 + 6 files changed, 350 insertions(+), 39 deletions(-) create mode 100644 bskylink/src/safelinkShutdown.test.ts create mode 100644 bskylink/src/shutdown.test.ts create mode 100644 bskylink/src/shutdown.ts diff --git a/bskylink/src/bin.ts b/bskylink/src/bin.ts index 0194471dc2..d0d4bcec5e 100644 --- a/bskylink/src/bin.ts +++ b/bskylink/src/bin.ts @@ -1,4 +1,11 @@ -import {Database, envToCfg, httpLogger, LinkService, readEnv} from './index.js' +import { + Database, + envToCfg, + FORCE_SHUTDOWN_TIMEOUT_MS, + httpLogger, + LinkService, + readEnv, +} from './index.js' async function main() { try { @@ -40,13 +47,30 @@ async function main() { await link.start() httpLogger.info('Link service is running') - process.on('SIGTERM', () => { + const shutdown = (signal: NodeJS.Signals) => { + const forceExitTimer = setTimeout(() => { + httpLogger.error( + {signal}, + 'Link service exceeded its shutdown deadline; forcing exit', + ) + process.exit(1) + }, FORCE_SHUTDOWN_TIMEOUT_MS) + forceExitTimer.unref() + void (async () => { - httpLogger.info('Link service is stopping') - await link.destroy() - httpLogger.info('Link service is stopped') + httpLogger.info({signal}, 'Link service is stopping') + try { + await link.destroy() + httpLogger.info({signal}, 'Link service is stopped') + } catch (err) { + process.exitCode = 1 + httpLogger.error({err, signal}, 'Failed to stop link service cleanly') + } })() - }) + } + + process.once('SIGTERM', shutdown) + process.once('SIGINT', shutdown) } catch (error) { httpLogger.error( { diff --git a/bskylink/src/cache/safelinkClient.ts b/bskylink/src/cache/safelinkClient.ts index c48b49ddc4..1883702468 100644 --- a/bskylink/src/cache/safelinkClient.ts +++ b/bskylink/src/cache/safelinkClient.ts @@ -26,6 +26,9 @@ export class SafelinkClient { private ozoneAgent: OzoneAgent private cursor?: string + private fetchEventsPromise?: Promise + private fetchEventsTimeout?: NodeJS.Timeout + private stopped = false constructor({cfg, db}: {cfg: ServiceConfig; db: Database}) { this.domainCache = new LRUCache({ @@ -124,7 +127,7 @@ export class SafelinkClient { return } - db.db + await db.db .insertInto('safelink_rule') .values({ id: rule.id, @@ -140,6 +143,7 @@ export class SafelinkClient { {error: err, rule}, 'failed to add rule to database', ) + throw err }) if (rule.pattern === 'domain') { @@ -174,6 +178,7 @@ export class SafelinkClient { {error: err, rule}, 'failed to remove rule from database', ) + throw err }) if (rule.pattern === 'domain') { @@ -183,13 +188,71 @@ export class SafelinkClient { } } - public async runFetchEvents() { + public runFetchEvents(): Promise { + if (this.stopped) { + return Promise.resolve() + } + + this.fetchEventsPromise ??= this.fetchEvents().finally(() => { + this.fetchEventsPromise = undefined + }) + return this.fetchEventsPromise + } + + public async stop(timeoutMs: number): Promise { + this.stopped = true + if (this.fetchEventsTimeout) { + clearTimeout(this.fetchEventsTimeout) + this.fetchEventsTimeout = undefined + } + const activePoll = this.fetchEventsPromise + if (!activePoll) { + return + } + + let timeout: NodeJS.Timeout | undefined + const stopWaiting = new Promise(resolve => { + timeout = setTimeout(() => { + redirectLogger.warn( + {timeoutMs}, + 'Safelink poll exceeded its shutdown deadline', + ) + resolve() + }, timeoutMs) + }) + try { + await Promise.race([activePoll, stopWaiting]) + } finally { + if (timeout) { + clearTimeout(timeout) + } + } + } + + private scheduleFetchEvents(delay: number) { + if (this.stopped) { + return + } + this.fetchEventsTimeout = setTimeout(() => { + this.fetchEventsTimeout = undefined + void this.runFetchEvents() + }, delay) + } + + private async fetchEvents() { let agent: AtpAgent try { agent = await this.ozoneAgent.getAgent() } catch (err) { + if (this.stopped) { + return + } redirectLogger.error({error: err}, 'error getting Ozone agent') - setTimeout(() => this.runFetchEvents(), SAFELINK_MAX_FETCH_INTERVAL) + this.scheduleFetchEvents(SAFELINK_MAX_FETCH_INTERVAL) + return + } + + if (this.stopped) { return } @@ -202,10 +265,13 @@ export class SafelinkClient { sortDirection: 'asc', }) } catch (err) { + if (this.stopped) { + return + } if (err instanceof ExpiredTokenError) { redirectLogger.info('ozone agent had expired session, refreshing...') await this.ozoneAgent.refreshSession() - setTimeout(() => this.runFetchEvents(), SAFELINK_MIN_FETCH_INTERVAL) + this.scheduleFetchEvents(SAFELINK_MIN_FETCH_INTERVAL) return } @@ -213,37 +279,56 @@ export class SafelinkClient { {error: err}, 'error fetching safelink events from Ozone', ) - setTimeout(() => this.runFetchEvents(), SAFELINK_MAX_FETCH_INTERVAL) + this.scheduleFetchEvents(SAFELINK_MAX_FETCH_INTERVAL) + return + } + + if (this.stopped) { return } if (res.data.events.length === 0) { redirectLogger.info('received no new safelink events from ozone') - setTimeout(() => this.runFetchEvents(), SAFELINK_MAX_FETCH_INTERVAL) + this.scheduleFetchEvents(SAFELINK_MAX_FETCH_INTERVAL) } else { - await this.db.transaction(async db => { - for (const rule of res.data.events) { - switch (rule.eventType) { - case 'removeRule': - await this.removeRule(db, rule) - break - case 'addRule': - case 'updateRule': - await this.addRule(db, rule) - break - default: - redirectLogger.warn({rule}, 'received unknown rule event type') + try { + await this.db.transaction(async db => { + for (const rule of res.data.events) { + switch (rule.eventType) { + case 'removeRule': + await this.removeRule(db, rule) + break + case 'addRule': + case 'updateRule': + await this.addRule(db, rule) + break + default: + redirectLogger.warn({rule}, 'received unknown rule event type') + } } + }) + if (this.stopped) { + return } - }) - if (res.data.cursor) { - redirectLogger.info( - {cursor: res.data.cursor}, - 'received new safelink events from Ozone', + if (res.data.cursor) { + redirectLogger.info( + {cursor: res.data.cursor}, + 'received new safelink events from Ozone', + ) + await this.setCursor(res.data.cursor) + } + } catch (err) { + if (this.stopped) { + return + } + redirectLogger.error( + {error: err}, + 'error applying safelink events from Ozone', ) - await this.setCursor(res.data.cursor) + this.scheduleFetchEvents(SAFELINK_MAX_FETCH_INTERVAL) + return } - setTimeout(() => this.runFetchEvents(), SAFELINK_MIN_FETCH_INTERVAL) + this.scheduleFetchEvents(SAFELINK_MIN_FETCH_INTERVAL) } } diff --git a/bskylink/src/index.ts b/bskylink/src/index.ts index 24e1a798a0..5f5bf33459 100644 --- a/bskylink/src/index.ts +++ b/bskylink/src/index.ts @@ -11,10 +11,12 @@ import {AppContext} from './context.js' import i18n from './i18n.js' import {createPrometheusRegistry} from './prometheus.js' import {default as routes, errorHandler} from './routes/index.js' +import {REQUEST_DRAIN_TIMEOUT_MS} from './shutdown.js' export * from './config.js' export * from './db/index.js' export * from './logger.js' +export * from './shutdown.js' export class LinkService { public server?: http.Server @@ -22,6 +24,7 @@ export class LinkService { private terminator?: HttpTerminator private metricsTerminator?: HttpTerminator private metricsRegistry: Registry + private destroyPromise?: Promise constructor( public app: express.Application, @@ -46,7 +49,10 @@ export class LinkService { this.ctx.metrics.start() this.server = this.app.listen(this.ctx.cfg.service.port) this.server.keepAliveTimeout = 90000 - this.terminator = createHttpTerminator({server: this.server}) + this.terminator = createHttpTerminator({ + server: this.server, + gracefulTerminationTimeout: REQUEST_DRAIN_TIMEOUT_MS, + }) await events.once(this.server, 'listening') const metricsApp = express() @@ -62,13 +68,26 @@ export class LinkService { await events.once(this.metricsServer, 'listening') } - async destroy() { + destroy(): Promise { + this.destroyPromise ??= this.destroyInternal() + + return this.destroyPromise + } + + private async destroyInternal() { this.ctx.abortController.abort() - await Promise.all([ - this.terminator?.terminate(), - this.metricsTerminator?.terminate(), - ]) - await this.ctx.db.close() - this.ctx.metrics.stop() + try { + await Promise.all([ + this.terminator?.terminate(), + this.metricsTerminator?.terminate(), + this.ctx.safelinkClient.stop(REQUEST_DRAIN_TIMEOUT_MS), + ]) + } finally { + try { + await this.ctx.db.close() + } finally { + this.ctx.metrics.stop() + } + } } } diff --git a/bskylink/src/safelinkShutdown.test.ts b/bskylink/src/safelinkShutdown.test.ts new file mode 100644 index 0000000000..63b99bd5e4 --- /dev/null +++ b/bskylink/src/safelinkShutdown.test.ts @@ -0,0 +1,118 @@ +import assert from 'node:assert' +import {describe, it} from 'node:test' + +import {SafelinkClient} from './cache/safelinkClient.js' + +const createClient = (getAgent: () => Promise) => { + const client: SafelinkClient = Object.create(SafelinkClient.prototype) + Reflect.set(client, 'stopped', false) + Reflect.set(client, 'ozoneAgent', {getAgent}) + return client +} + +describe('Safelink shutdown', () => { + it('clears a scheduled retry and cannot restart after stop', async () => { + const client = createClient(async () => { + throw new Error('Ozone unavailable') + }) + + await client.runFetchEvents() + assert.ok(Reflect.get(client, 'fetchEventsTimeout')) + + await client.stop(1_000) + assert.strictEqual(Reflect.get(client, 'fetchEventsTimeout'), undefined) + assert.strictEqual(Reflect.get(client, 'stopped'), true) + + await client.runFetchEvents() + assert.strictEqual(Reflect.get(client, 'fetchEventsTimeout'), undefined) + }) + + it('waits for an active poll to finish before stopping', async () => { + let pollStarted = () => {} + const started = new Promise(resolve => { + pollStarted = () => resolve(undefined) + }) + let finishPoll = () => {} + const releasePoll = new Promise(resolve => { + finishPoll = () => resolve(undefined) + }) + const client = createClient(async () => { + pollStarted() + await releasePoll + throw new Error('poll released during shutdown') + }) + + const polling = client.runFetchEvents() + await started + + let stopped = false + const stopping = client.stop(1_000).then(() => { + stopped = true + }) + await new Promise(resolve => setTimeout(resolve, 25)) + assert.strictEqual(stopped, false) + + finishPoll() + await Promise.all([polling, stopping]) + assert.strictEqual(stopped, true) + }) + + it( + 'bounds the wait for a poll that never finishes', + {timeout: 1_000}, + async () => { + const client = createClient(() => new Promise(() => {})) + void client.runFetchEvents() + + const startedAt = Date.now() + await client.stop(25) + assert.ok(Date.now() - startedAt >= 20) + }, + ) + + it('retries a failed rule write without advancing the cursor', async () => { + const client = createClient(async () => ({ + tools: { + ozone: { + safelink: { + queryEvents: async () => ({ + data: { + cursor: 'next', + events: [ + { + action: 'block', + createdAt: new Date().toISOString(), + eventType: 'addRule', + id: 1, + pattern: 'domain', + url: 'example.com', + }, + ], + }, + }), + }, + }, + }, + })) + Reflect.set(client, 'cursor', 'current') + Reflect.set(client, 'db', { + transaction: async (run: (db: unknown) => Promise) => + run({ + db: { + insertInto: () => ({ + values: () => ({ + execute: async () => { + throw new Error('database unavailable') + }, + }), + }), + }, + }), + }) + + await client.runFetchEvents() + assert.ok(Reflect.get(client, 'fetchEventsTimeout')) + assert.strictEqual(Reflect.get(client, 'cursor'), 'current') + await client.stop(1_000) + }) +}) diff --git a/bskylink/src/shutdown.test.ts b/bskylink/src/shutdown.test.ts new file mode 100644 index 0000000000..7d75ca3ab0 --- /dev/null +++ b/bskylink/src/shutdown.test.ts @@ -0,0 +1,63 @@ +import assert from 'node:assert' +import events from 'node:events' +import http from 'node:http' +import {describe, it} from 'node:test' + +import {createHttpTerminator} from 'http-terminator' + +import {REQUEST_DRAIN_TIMEOUT_MS} from './shutdown.js' + +describe('HTTP shutdown', () => { + it('allows in-flight requests to finish during the drain window', async () => { + let beginRequest = () => {} + const requestStarted = new Promise(resolve => { + beginRequest = () => resolve(undefined) + }) + let finishRequest = () => {} + const releaseRequest = new Promise(resolve => { + finishRequest = () => resolve(undefined) + }) + + const server = http.createServer(async (_req, res) => { + beginRequest() + await releaseRequest + res.end('finished') + }) + server.listen(0, '127.0.0.1') + await events.once(server, 'listening') + + const address = server.address() + assert.ok(address && typeof address !== 'string') + + const responsePromise = fetch(`http://127.0.0.1:${address.port}`) + await requestStarted + + const terminator = createHttpTerminator({ + server, + gracefulTerminationTimeout: REQUEST_DRAIN_TIMEOUT_MS, + }) + let termination: Promise | undefined + try { + let terminated = false + termination = terminator.terminate().then(() => { + terminated = true + }) + + await new Promise(resolve => setTimeout(resolve, 25)) + assert.strictEqual(terminated, false) + + finishRequest() + const response = await responsePromise + assert.strictEqual(await response.text(), 'finished') + await termination + assert.strictEqual(terminated, true) + } finally { + finishRequest() + await (termination ?? terminator.terminate()) + } + }) + + it('uses the shared 60 second request drain budget', () => { + assert.strictEqual(REQUEST_DRAIN_TIMEOUT_MS, 60_000) + }) +}) diff --git a/bskylink/src/shutdown.ts b/bskylink/src/shutdown.ts new file mode 100644 index 0000000000..d399e189c6 --- /dev/null +++ b/bskylink/src/shutdown.ts @@ -0,0 +1,2 @@ +export const REQUEST_DRAIN_TIMEOUT_MS = 60_000 +export const FORCE_SHUTDOWN_TIMEOUT_MS = REQUEST_DRAIN_TIMEOUT_MS + 4_000