Add graceful shutdown to Blink (#11523)

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')
+        httpLogger.info({signal}, 'Link service is stopping')
+        try {
           await link.destroy()
-        httpLogger.info('Link service is stopped')
+          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(
       {

bskylink/src/cache/safelinkClient.ts

@@ -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,14 +279,19 @@ 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 {
+      try {
         await this.db.transaction(async db => {
           for (const rule of res.data.events) {
             switch (rule.eventType) {
@@ -236,6 +307,9 @@ export class SafelinkClient {
             }
           }
         })
+        if (this.stopped) {
+          return
+        }
         if (res.data.cursor) {
           redirectLogger.info(
             {cursor: res.data.cursor},
@@ -243,7 +317,18 @@ export class SafelinkClient {
           )
           await this.setCursor(res.data.cursor)
         }
-      setTimeout(() => this.runFetchEvents(), SAFELINK_MIN_FETCH_INTERVAL)
+      } catch (err) {
+        if (this.stopped) {
+          return
+        }
+        redirectLogger.error(
+          {error: err},
+          'error applying safelink events from Ozone',
+        )
+        this.scheduleFetchEvents(SAFELINK_MAX_FETCH_INTERVAL)
+        return
+      }
+      this.scheduleFetchEvents(SAFELINK_MIN_FETCH_INTERVAL)
     }
   }
 

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<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()
+    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()
       }
+    }
+  }
 }

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<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)
+  })
+})

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<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)
+  })
+})

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