fix(db): measure the migration lock budget on a monotonic clock and never start an attempt past it

Author: waleedlatif1

packages/db/scripts/lock-timeout-retry.test.ts

@@ -44,10 +44,12 @@ describe('retryOnLockTimeout', () => {
     expect(clock.sleep).toHaveBeenCalledTimes(19)
   })
 
-  it('throws the last lock timeout once the next retry would end past the budget', async () => {
+  it('starts no attempt after the budget and throws the last lock timeout', async () => {
     const clock = fakeClock()
     const onRetry = vi.fn()
+    const startedAt: number[] = []
     const attempt = vi.fn(async () => {
+      startedAt.push(clock.now())
       clock.advance(5_000)
       throw pgError('55P03')
     })
@@ -62,13 +64,33 @@ describe('retryOnLockTimeout', () => {
       })
     ).rejects.toMatchObject({ code: '55P03' })
 
-    expect(clock.now()).toBeLessThan(2 * 60_000)
-    for (const [{ elapsedMs, delayMs }] of onRetry.mock.calls) {
-      expect(elapsedMs + delayMs).toBeLessThan(2 * 60_000)
-    }
+    for (const start of startedAt) expect(start).toBeLessThan(2 * 60_000)
+    /** The last attempt may run one lock timeout past the budget, never more. */
+    expect(clock.now()).toBeLessThan(2 * 60_000 + 5_000)
     expect(attempt).toHaveBeenCalledTimes(onRetry.mock.calls.length + 1)
   })
 
+  it('does not start an attempt when a timer resolves after the budget', async () => {
+    let nowMs = 0
+    const attempt = vi.fn(async () => {
+      nowMs += 1_000
+      throw pgError('55P03')
+    })
+
+    await expect(
+      retryOnLockTimeout(attempt, {
+        budgetMs: 60_000,
+        backoff: BACKOFF,
+        now: () => nowMs,
+        /** The process stalls: the timer fires long after its delay. */
+        sleep: async () => {
+          nowMs += 120_000
+        },
+      })
+    ).rejects.toMatchObject({ code: '55P03' })
+    expect(attempt).toHaveBeenCalledOnce()
+  })
+
   it('finds a lock timeout wrapped in a cause chain', async () => {
     const clock = fakeClock()
     let calls = 0

packages/db/scripts/lock-timeout-retry.ts

@@ -15,13 +15,16 @@ export interface LockTimeoutRetryAttempt {
 
 export interface LockTimeoutRetryOptions {
   /**
-   * Total wall-clock time, measured from the first attempt, during which a lock
-   * timeout is retried. A retry whose delay would end past the budget is not
-   * started; the last lock timeout is thrown instead.
+   * Time, measured on a monotonic clock from the first attempt, within which
+   * attempts may start. No attempt starts once the budget has elapsed, whether
+   * the backoff delay would end past it or a timer resolved late; the last lock
+   * timeout is thrown instead. An attempt that starts in time can still run for
+   * up to one `lock_timeout` past the budget.
    */
   budgetMs: number
   backoff: { baseMs: number; maxMs: number }
   onRetry?: (attempt: LockTimeoutRetryAttempt) => void
+  /** Monotonic milliseconds; defaults to `performance.now`, immune to wall-clock corrections. */
   now?: () => number
   sleep?: (ms: number) => Promise<void>
 }
@@ -43,7 +46,7 @@ export async function retryOnLockTimeout<T>(
   attempt: (attemptNumber: number) => Promise<T>,
   options: LockTimeoutRetryOptions
 ): Promise<T> {
-  const now = options.now ?? Date.now
+  const now = options.now ?? (() => performance.now())
   const sleep = options.sleep ?? defaultSleep
   const startedAt = now()
   for (let attemptNumber = 1; ; attemptNumber++) {
@@ -61,6 +64,7 @@ export async function retryOnLockTimeout<T>(
         budgetMs: options.budgetMs,
       })
       await sleep(delayMs)
+      if (now() - startedAt >= options.budgetMs) throw error
     }
   }
 }