refactor(util): merge withRetries into waitUntil aws#6387

## Problem:
`withRetries` and `waitUntil` are two different interfaces for very
similar concepts.

## Solution:
- Add `backoff` and `retryOnFail` flags to `waitUntil()`.
    - `waitUntil(…, {retryOnFail:true})` returns `undefined` if it never
      properly resolves.
    - `waitUntil(…, {retryOnFail:false})` throws if it never properly
      resolves.
- Delete `withRetries()`.

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 { 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?
             }
         )

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

packages/core/src/shared/resourcefetcher/httpResourceFetcher.ts

@@ -6,9 +6,8 @@
 import { VSCODE_EXTENSION_ID } from '../extensions'
 import { getLogger, Logger } from '../logger'
 import { ResourceFetcher } from './resourcefetcher'
-import { Timeout, CancelEvent } from '../utilities/timeoutUtils'
+import { Timeout, CancelEvent, waitUntil } from '../utilities/timeoutUtils'
 import request, { RequestError } from '../request'
-import { withRetries } from '../utilities/functionUtils'
 
 type RequestHeaders = { eTag?: string; gZip?: boolean }
 
@@ -30,7 +29,6 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
             showUrl: boolean
             friendlyName?: string
             timeout?: Timeout
-            retries?: number
         }
     ) {}
 
@@ -97,7 +95,7 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
     }
 
     private async getResponseFromGetRequest(timeout?: Timeout, headers?: RequestHeaders): Promise<Response> {
-        return withRetries(
+        return waitUntil(
             () => {
                 const req = request.fetch('GET', this.url, {
                     headers: this.buildRequestHeaders(headers),
@@ -111,7 +109,10 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
                 return req.response.finally(() => cancelListener?.dispose())
             },
             {
-                maxRetries: this.params.retries ?? 1,
+                timeout: 3000,
+                interval: 100,
+                backoff: 2,
+                retryOnFail: true,
             }
         )
     }

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

@@ -219,40 +219,101 @@ 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 immediate result.
+     * - 'truthy' arg is ignored
+     * - If the timeout is reached it throws 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`).
+ * Invokes `fn()` on an interval based on the given arguments. This can be used for retries, or until
+ * an expected result is given. Read {@link WaitUntilOptions} carefully.
  *
  * @param fn  Function whose result is checked
  * @param options  See {@link WaitUntilOptions}
  *
- * @returns Result of `fn()`, or `undefined` if timeout was reached.
+ * @returns Result of `fn()`, or possibly `undefined` depending on the arguments.
  */
+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 avoids the need to progress the stubbed clock
+        if (interval > 0) {
+            await sleep(interval)
         }
 
-        await sleep(opt.interval)
+        elapsed += interval
+        interval = interval * opt.backoff
     }
 }
 

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,7 @@ import {
 import { HttpResourceFetcher } from '../../shared/resourcefetcher/httpResourceFetcher'
 import { NotificationsNode } from '../../notifications/panelNode'
 import { RuleEngine } from '../../notifications/rules'
+import Sinon from 'sinon'
 
 // one test node to use across different tests
 export const panelNode: NotificationsNode = NotificationsNode.instance
@@ -512,43 +513,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)
-        })()
+        // Start function under test
+        const fetcher = assert.rejects(new RemoteFetcher().fetch('startUp', 'any'), (e) => {
+            return e instanceof Error && e.message === 'last error'
+        })
 
-        const fetcher = new RemoteFetcher()
+        // 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 fetcher
-            .fetch('startUp', 'any')
-            .then(() => assert.ok(false, 'Did not throw exception.'))
-            .catch(() => assert.ok(true))
-        await runClock
-
-        httpStub.restore()
     })
 })