refactor(util): merge withRetries and waitUntil

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 <[email protected]>

Author: nkomonen-amazon

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 { withRetries } 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'
@@ -296,8 +296,8 @@ export class RemoteFetcher implements NotificationFetcher {
                 }
             },
             {
-                maxRetries: RemoteFetcher.retryNumber,
-                delay: RemoteFetcher.retryIntervalMs,
+                interval: RemoteFetcher.retryIntervalMs,
+                timeout: RemoteFetcher.retryTimeout,
                 // No exponential backoff - necessary?
             }
         )

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 { withRetries } from './utilities/timeoutUtils'
 
 const className = 'CrashMonitoring'
 
@@ -489,7 +489,7 @@ 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 = withRetries(funcWithCtx, { timeout: 15_000, interval: 100, backoff: 2 })
 
             return funcWithRetries
         } catch (e) {
@@ -542,7 +542,7 @@ 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 = withRetries(funcWithCtx, { timeout: 15_000, interval: 100, backoff: 2 })
         await funcWithRetries
     }
 
@@ -609,7 +609,7 @@ export class FileSystemState {
                 }
                 const funcWithIgnoreBadFile = () => ignoreBadFileError(loadExtFromDisk)
                 const funcWithRetries = () =>
-                    withRetries(funcWithIgnoreBadFile, { maxRetries: 6, delay: 100, backoff: 2 })
+                    withRetries(funcWithIgnoreBadFile, { timeout: 15_000, interval: 100, backoff: 2 })
                 const funcWithCtx = () => withFailCtx('parseRunningExtFile', funcWithRetries)
                 const ext: ExtInstanceHeartbeat | undefined = await funcWithCtx()
 

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

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,116 @@ 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
+}
+
+interface WithRetriesOptions extends WaitUntilOptions {
+    /**
+     * Ignore intermediate errors thrown by fn(). If timeout it reached, it will throw the last error.
+     * (default: true)
+     */
+    readonly catch?: boolean
 }
 
 /**
  * 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): Promise<T | undefined> {
-    const opt = { timeout: 5000, interval: 500, truthy: true, ...options }
+    const opt: WithRetriesOptions = { timeout: 5000, interval: 500, truthy: true, backoff: 1, ...options }
+    try {
+        return await withRetries(fn, Object.assign(opt, { catch: false }))
+    } catch (e) {
+        if (e instanceof RetryTimeoutError) {
+            return undefined
+        }
+        throw e
+    }
+}
+
+export const withRetriesDefaultTimeout = 2000
+export const withRetriesInterval = 500
+
+/**
+ * Invokes `fn()` until it successfully returns a value, otherwise it throws.
+ *
+ * @param fn  Function whose result is checked
+ * @param options  See {@link WaitUntilOptions}
+ */
+export async function withRetries<T>(fn: () => Promise<T>, options: WithRetriesOptions = {}): Promise<T> {
+    const opt = {
+        timeout: withRetriesDefaultTimeout,
+        interval: withRetriesInterval,
+        truthy: true,
+        backoff: 1,
+        catch: true,
+        ...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.truthy && result) || (!opt.truthy && result !== undefined)) {
+                return result
+            }
+        } catch (e) {
+            if (!opt.catch) {
+                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 ((opt.truthy && result) || (!opt.truthy && result !== undefined)) {
-            return result
+        // If the sleep will exceed the timeout, abort early
+        if (elapsed + interval >= remaining) {
+            throw lastError
+                ? new RetryTimeoutError('Last error was:', { cause: lastError })
+                : new RetryTimeoutError('Function never returned a truthy result')
         }
-        if (i * opt.interval >= opt.timeout) {
-            return undefined
+
+        // when testing we can ignore progressing a 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.
  *

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