refactor(util): merge withRetries and waitUntil

nkomonen-amazon · nkomonen-amazon · commit f37ef957cb16 · 2025-01-17T13:20:54.000-05:00
Problem:

We had 2 different implementation of similar code, between
withRetries and waitUntil.

We couldn't easily merge them though since behavior was slightly different
when it came to expected return types, and since one threw and the other
did not.

- waitUntil returns undefined if it never properly resolves
- withRetries throws if it never properly resolves

Solution:

As a solution we keep both methods but they share the same underlying
code.

- waitUntil now has a backoff field to increase the delta between intervals
  - withRetries also has this
- withRetries will retry even on a thrown exeception
  - on timeout it will throw the last exception it encountered

Signed-off-by: nkomonen-amazon &lt;nkomonen@amazon.com&gt;
diff --git a/packages/core/src/notifications/controller.ts b/packages/core/src/notifications/controller.ts
@@ -24,11 +24,11 @@ import { NotificationsNode } from './panelNode'
 import { Commands } from '../shared/vscode/commands2'
 import { RuleEngine } from './rules'
 import { TreeNode } from '../shared/treeview/resourceTreeDataProvider'
-import { withRetries } from '../shared/utilities/functionUtils'
 import { FileResourceFetcher } from '../shared/resourcefetcher/fileResourceFetcher'
 import { isAmazonQ } from '../shared/extensionUtilities'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { randomUUID } from '../shared/crypto'
+import { waitUntil } from '../shared/utilities/timeoutUtils'
 
 const logger = getLogger('notifications')
 
@@ -266,8 +266,8 @@ export interface NotificationFetcher {
 }
 
 export class RemoteFetcher implements NotificationFetcher {
-    public static readonly retryNumber = 5
     public static readonly retryIntervalMs = 30000
+    public static readonly retryTimeout = RemoteFetcher.retryIntervalMs * 5
 
     private readonly startUpEndpoint: string =
         'https://idetoolkits-hostedfiles.amazonaws.com/Notifications/VSCode/startup/1.x.json'
@@ -286,7 +286,7 @@ export class RemoteFetcher implements NotificationFetcher {
         })
         logger.verbose('Attempting to fetch notifications for category: %s at endpoint: %s', category, endpoint)
 
-        return withRetries(
+        return waitUntil(
             async () => {
                 try {
                     return await fetcher.getNewETagContent(versionTag)
@@ -296,8 +296,9 @@ export class RemoteFetcher implements NotificationFetcher {
                 }
             },
             {
-                maxRetries: RemoteFetcher.retryNumber,
-                delay: RemoteFetcher.retryIntervalMs,
+                interval: RemoteFetcher.retryIntervalMs,
+                timeout: RemoteFetcher.retryTimeout,
+                retryOnFail: true,
                 // No exponential backoff - necessary?
             }
         )
diff --git a/packages/core/src/shared/crashMonitoring.ts b/packages/core/src/shared/crashMonitoring.ts
@@ -17,8 +17,8 @@ import fs from './fs/fs'
 import { getLogger } from './logger/logger'
 import { crashMonitoringDirName } from './constants'
 import { throwOnUnstableFileSystem } from './filesystemUtilities'
-import { withRetries } from './utilities/functionUtils'
 import { truncateUuid } from './crypto'
+import { waitUntil } from './utilities/timeoutUtils'
 
 const className = 'CrashMonitoring'
 
@@ -489,7 +489,12 @@ export class FileSystemState {
                 this.deps.devLogger?.debug(`HEARTBEAT sent for ${truncateUuid(this.ext.sessionId)}`)
             }
             const funcWithCtx = () => withFailCtx('sendHeartbeatState', func)
-            const funcWithRetries = withRetries(funcWithCtx, { maxRetries: 6, delay: 100, backoff: 2 })
+            const funcWithRetries = waitUntil(funcWithCtx, {
+                timeout: 15_000,
+                interval: 100,
+                backoff: 2,
+                retryOnFail: true,
+            })
 
             return funcWithRetries
         } catch (e) {
@@ -542,7 +547,12 @@ export class FileSystemState {
             await nodeFs.rm(filePath, { force: true })
         }
         const funcWithCtx = () => withFailCtx(ctx, func)
-        const funcWithRetries = withRetries(funcWithCtx, { maxRetries: 6, delay: 100, backoff: 2 })
+        const funcWithRetries = waitUntil(funcWithCtx, {
+            timeout: 15_000,
+            interval: 100,
+            backoff: 2,
+            retryOnFail: true,
+        })
         await funcWithRetries
     }
 
@@ -609,7 +619,7 @@ export class FileSystemState {
                 }
                 const funcWithIgnoreBadFile = () => ignoreBadFileError(loadExtFromDisk)
                 const funcWithRetries = () =>
-                    withRetries(funcWithIgnoreBadFile, { maxRetries: 6, delay: 100, backoff: 2 })
+                    waitUntil(funcWithIgnoreBadFile, { timeout: 15_000, interval: 100, backoff: 2, retryOnFail: true })
                 const funcWithCtx = () => withFailCtx('parseRunningExtFile', funcWithRetries)
                 const ext: ExtInstanceHeartbeat | undefined = await funcWithCtx()
 
diff --git a/packages/core/src/shared/utilities/functionUtils.ts b/packages/core/src/shared/utilities/functionUtils.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { sleep, Timeout } from './timeoutUtils'
+import { Timeout } from './timeoutUtils'
 
 /**
  * Creates a function that always returns a 'shared' Promise.
@@ -145,34 +145,3 @@ export function cancellableDebounce<T, U extends any[]>(
         cancel: cancel,
     }
 }
-
-/**
- * Executes the given function, retrying if it throws.
- *
- * @param opts - if no opts given, defaults are used
- */
-export async function withRetries<T>(
-    fn: () => Promise<T>,
-    opts?: { maxRetries?: number; delay?: number; backoff?: number }
-): Promise<T> {
-    const maxRetries = opts?.maxRetries ?? 3
-    const delay = opts?.delay ?? 0
-    const backoff = opts?.backoff ?? 1
-
-    let retryCount = 0
-    let latestDelay = delay
-    while (true) {
-        try {
-            return await fn()
-        } catch (err) {
-            retryCount++
-            if (retryCount >= maxRetries) {
-                throw err
-            }
-            if (latestDelay > 0) {
-                await sleep(latestDelay)
-                latestDelay = latestDelay * backoff
-            }
-        }
-    }
-}
diff --git a/packages/core/src/shared/utilities/timeoutUtils.ts b/packages/core/src/shared/utilities/timeoutUtils.ts
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { ToolkitError } from '../errors'
 import globals from '../extensionGlobals'
 import { CancellationToken, EventEmitter, Event } from 'vscode'
 
@@ -219,43 +220,111 @@ interface WaitUntilOptions {
     readonly interval?: number
     /** Wait for "truthy" result, else wait for any defined result including `false` (default: true) */
     readonly truthy?: boolean
+    /** A backoff multiplier for how long the next interval will be (default: None, i.e 1) */
+    readonly backoff?: number
+    /**
+     * Only retries when an error is thrown, otherwise returning the result regardless of truthiness.
+     * - Ignores 'truthy' arg
+     * - If the timeout is reached it will throw the last error
+     * - default: false
+     */
+    readonly retryOnFail?: boolean
 }
 
+export const waitUntilDefaultTimeout = 2000
+export const waitUntilDefaultInterval = 500
+
 /**
  * Invokes `fn()` until it returns a truthy value (or non-undefined if `truthy:false`).
  *
+ * Also look at {@link withRetries} to be able to retry on failures.
+ *
  * @param fn  Function whose result is checked
  * @param options  See {@link WaitUntilOptions}
  *
  * @returns Result of `fn()`, or `undefined` if timeout was reached.
  */
+export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptions & { retryOnFail: true }): Promise<T>
+export async function waitUntil<T>(
+    fn: () => Promise<T>,
+    options: WaitUntilOptions & { retryOnFail: false }
+): Promise<T | undefined>
+export async function waitUntil<T>(
+    fn: () => Promise<T>,
+    options: Omit<WaitUntilOptions, 'retryOnFail'>
+): Promise<T | undefined>
 export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptions): Promise<T | undefined> {
-    const opt = { timeout: 5000, interval: 500, truthy: true, ...options }
+    // set default opts
+    const opt = {
+        timeout: waitUntilDefaultTimeout,
+        interval: waitUntilDefaultInterval,
+        truthy: true,
+        backoff: 1,
+        retryOnFail: false,
+        ...options,
+    }
+
+    let interval = opt.interval
+    let lastError: Error | undefined
+    let elapsed: number = 0
+    let remaining = opt.timeout
+
     for (let i = 0; true; i++) {
         const start: number = globals.clock.Date.now()
         let result: T
 
-        // Needed in case a caller uses a 0 timeout (function is only called once)
-        if (opt.timeout > 0) {
-            result = await Promise.race([fn(), new Promise<T>((r) => globals.clock.setTimeout(r, opt.timeout))])
-        } else {
-            result = await fn()
+        try {
+            // Needed in case a caller uses a 0 timeout (function is only called once)
+            if (remaining > 0) {
+                result = await Promise.race([fn(), new Promise<T>((r) => globals.clock.setTimeout(r, remaining))])
+            } else {
+                result = await fn()
+            }
+
+            if (opt.retryOnFail || (opt.truthy && result) || (!opt.truthy && result !== undefined)) {
+                return result
+            }
+        } catch (e) {
+            if (!opt.retryOnFail) {
+                throw e
+            }
+
+            // Unlikely to hit this, but exists for typing
+            if (!(e instanceof Error)) {
+                throw e
+            }
+
+            lastError = e
         }
 
         // Ensures that we never overrun the timeout
-        opt.timeout -= globals.clock.Date.now() - start
+        remaining -= globals.clock.Date.now() - start
+
+        // If the sleep will exceed the timeout, abort early
+        if (elapsed + interval >= remaining) {
+            if (!opt.retryOnFail) {
+                return undefined
+            }
 
-        if ((opt.truthy && result) || (!opt.truthy && result !== undefined)) {
-            return result
+            throw lastError
         }
-        if (i * opt.interval >= opt.timeout) {
-            return undefined
+
+        // when testing, this saves the need to progress the stubbed clock
+        if (interval > 0) {
+            await sleep(interval)
         }
 
-        await sleep(opt.interval)
+        elapsed += interval
+        interval = interval * opt.backoff
     }
 }
 
+/**
+ * Implies the function failed due to an external timeout, the function
+ * could have been successful if given more time.
+ */
+export const RetryTimeoutError = ToolkitError.named('RetryTimeoutError')
+
 /**
  * @deprecated Prefer using event-driven timeout mechanisms over racing promises.
  *
diff --git a/packages/core/src/test/notifications/controller.test.ts b/packages/core/src/test/notifications/controller.test.ts
@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import * as FakeTimers from '@sinonjs/fake-timers'
 import assert from 'assert'
-import sinon from 'sinon'
+import sinon, { createSandbox } from 'sinon'
 import globals from '../../shared/extensionGlobals'
 import { randomUUID } from '../../shared/crypto'
 import { getContext } from '../../shared/vscode/setContext'
@@ -27,6 +27,8 @@ import {
 import { HttpResourceFetcher } from '../../shared/resourcefetcher/httpResourceFetcher'
 import { NotificationsNode } from '../../notifications/panelNode'
 import { RuleEngine } from '../../notifications/rules'
+import { RetryTimeoutError } from '../../shared/utilities/timeoutUtils'
+import Sinon from 'sinon'
 
 // one test node to use across different tests
 export const panelNode: NotificationsNode = NotificationsNode.instance
@@ -512,43 +514,46 @@ describe('Notifications Controller', function () {
 
 describe('RemoteFetcher', function () {
     let clock: FakeTimers.InstalledClock
+    let sandbox: Sinon.SinonSandbox
 
     before(function () {
         clock = installFakeClock()
+        sandbox = createSandbox()
     })
 
     afterEach(function () {
         clock.reset()
+        sandbox.restore()
     })
 
     after(function () {
         clock.uninstall()
     })
 
     it('retries and throws error', async function () {
-        const httpStub = sinon.stub(HttpResourceFetcher.prototype, 'getNewETagContent')
+        // Setup
+        const httpStub = sandbox.stub(HttpResourceFetcher.prototype, 'getNewETagContent')
         httpStub.throws(new Error('network error'))
 
-        const runClock = (async () => {
-            await clock.tickAsync(1)
-            for (let n = 1; n <= RemoteFetcher.retryNumber; n++) {
-                assert.equal(httpStub.callCount, n)
-                await clock.tickAsync(RemoteFetcher.retryIntervalMs)
-            }
-
-            // Stop trying
-            await clock.tickAsync(RemoteFetcher.retryNumber)
-            assert.equal(httpStub.callCount, RemoteFetcher.retryNumber)
-        })()
-
-        const fetcher = new RemoteFetcher()
-        await fetcher
-            .fetch('startUp', 'any')
-            .then(() => assert.ok(false, 'Did not throw exception.'))
-            .catch(() => assert.ok(true))
-        await runClock
-
-        httpStub.restore()
+        // Start function under test
+        const fetcher = new RemoteFetcher().fetch('startUp', 'any').then(() => assert.fail('Did not throw exception.'))
+
+        // Progresses the clock, allowing the fetcher logic to break out of sleep for each iteration of withRetries()
+        assert.strictEqual(httpStub.callCount, 1) // 0
+        await clock.tickAsync(RemoteFetcher.retryIntervalMs)
+        assert.strictEqual(httpStub.callCount, 2) // 30_000
+        await clock.tickAsync(RemoteFetcher.retryIntervalMs)
+        assert.strictEqual(httpStub.callCount, 3) // 60_000
+        await clock.tickAsync(RemoteFetcher.retryIntervalMs)
+        assert.strictEqual(httpStub.callCount, 4) // 120_000
+        httpStub.throws(new Error('last error'))
+        await clock.tickAsync(RemoteFetcher.retryIntervalMs)
+        assert.strictEqual(httpStub.callCount, 5) // 150_000
+
+        // We hit timeout so the last error will be thrown
+        await assert.rejects(fetcher, (e) => {
+            return e instanceof RetryTimeoutError && e.cause?.message === 'last error'
+        })
     })
 })
 
diff --git a/packages/core/src/test/shared/utilities/functionUtils.test.ts b/packages/core/src/test/shared/utilities/functionUtils.test.ts
diff --git a/packages/core/src/test/shared/utilities/timeoutUtils.test.ts b/packages/core/src/test/shared/utilities/timeoutUtils.test.ts
