Add graceful shutdown to Blink (#11523)

2026-08-20 13:16:31 -07:00
parent a892057e29
commit 21a2c3d6f2
6 changed files with 350 additions and 39 deletions
@@ -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(
      {
@@ -26,6 +26,9 @@ export class SafelinkClient {
  private ozoneAgent: OzoneAgent

  private cursor?: string
+  private fetchEventsPromise?: Promise<void>
+  private fetchEventsTimeout?: NodeJS.Timeout
+  private stopped = false

  constructor({cfg, db}: {cfg: ServiceConfig; db: Database}) {
    this.domainCache = new LRUCache<string, SafelinkRule | 'ok'>({
@@ -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<void> {
+    if (this.stopped) {
+      return Promise.resolve()
+    }
+
+    this.fetchEventsPromise ??= this.fetchEvents().finally(() => {
+      this.fetchEventsPromise = undefined
+    })
+    return this.fetchEventsPromise
+  }
+
+  public async stop(timeoutMs: number): Promise<void> {
+    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<void>(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)
    }
  }

@@ -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<void>

  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<void> {
+    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()
+      }
+    }
  }
 }
@@ -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<unknown>) => {
+  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<void>(resolve => {
+      pollStarted = () => resolve(undefined)
+    })
+    let finishPoll = () => {}
+    const releasePoll = new Promise<void>(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<never>(() => {}))
+      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<void>) =>
+        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)
+  })
+})
@@ -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<void>(resolve => {
+      beginRequest = () => resolve(undefined)
+    })
+    let finishRequest = () => {}
+    const releaseRequest = new Promise<void>(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<void> | 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)
+  })
+})
@@ -0,0 +1,2 @@
+export const REQUEST_DRAIN_TIMEOUT_MS = 60_000
+export const FORCE_SHUTDOWN_TIMEOUT_MS = REQUEST_DRAIN_TIMEOUT_MS + 4_000