retry heartbeats on fail

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/amazonq/src/extensionNode.ts

@@ -32,7 +32,8 @@ export async function activate(context: vscode.ExtensionContext) {
  * the code compatible with web and move it to {@link activateAmazonQCommon}.
  */
 async function activateAmazonQNode(context: vscode.ExtensionContext) {
-    await (await CrashMonitoring.instance())?.start()
+    // Intentionally do not await since this is slow and non-critical
+    void (await CrashMonitoring.instance())?.start()
 
     const extContext = {
         extensionContext: context,

packages/core/src/extensionNode.ts

@@ -78,7 +78,8 @@ export async function activate(context: vscode.ExtensionContext) {
         // IMPORTANT: If you are doing setup that should also work in web mode (browser), it should be done in the function below
         const extContext = await activateCommon(context, contextPrefix, false)
 
-        await (await CrashMonitoring.instance())?.start()
+        // Intentionally do not await since this can be slow and non-critical
+        void (await CrashMonitoring.instance())?.start()
 
         initializeCredentialsProviderManager()
 

packages/core/src/shared/crashMonitoring.ts

@@ -17,6 +17,7 @@ import fs from './fs/fs'
 import { getLogger } from './logger/logger'
 import { crashMonitoringDirNames } from './constants'
 import { throwOnUnstableFileSystem } from './filesystemUtilities'
+import { withRetries } from './utilities/functionUtils'
 
 const className = 'CrashMonitoring'
 
@@ -427,17 +428,21 @@ export class FileSystemState {
 
     // ------------------ Heartbeat methods ------------------
     public async sendHeartbeat() {
-        await withFailCtx('sendHeartbeatState', async () => {
-            const dir = await this.runningExtsDir()
-            const extId = this.createExtId(this.ext)
-            await fs.writeFile(
-                path.join(dir, extId),
-                JSON.stringify({ ...this.ext, lastHeartbeat: this.deps.now() }, undefined, 4)
-            )
-            this.deps.devLogger?.debug(
-                `crashMonitoring: HEARTBEAT pid ${this.deps.pid} + sessionId: ${this.deps.sessionId.slice(0, 8)}-...`
-            )
-        })
+        const _sendHeartbeat = () =>
+            withFailCtx('sendHeartbeatState', async () => {
+                const dir = await this.runningExtsDir()
+                const extId = this.createExtId(this.ext)
+                await fs.writeFile(
+                    path.join(dir, extId),
+                    JSON.stringify({ ...this.ext, lastHeartbeat: this.deps.now() }, undefined, 4)
+                )
+                this.deps.devLogger?.debug(
+                    `crashMonitoring: HEARTBEAT pid ${this.deps.pid} + sessionId: ${this.deps.sessionId.slice(0, 8)}-...`
+                )
+            })
+        // retry a failed heartbeat to avoid incorrectly being reported as a crash
+        const _sendHeartbeatWithRetries = withRetries(_sendHeartbeat, { maxRetries: 8, delay: 100, backoff: 2 })
+        return _sendHeartbeatWithRetries
     }
 
     /**

packages/core/src/shared/utilities/functionUtils.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Timeout } from './timeoutUtils'
+import { sleep, Timeout } from './timeoutUtils'
 
 /**
  * Creates a function that always returns a 'shared' Promise.
@@ -145,3 +145,34 @@ 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/test/shared/utilities/functionUtils.test.ts

@@ -4,8 +4,10 @@
  */
 
 import assert from 'assert'
-import { once, onceChanged, debounce } from '../../../shared/utilities/functionUtils'
+import { once, onceChanged, debounce, withRetries } from '../../../shared/utilities/functionUtils'
 import { installFakeClock } from '../../testUtil'
+import { stub, SinonStub } from 'sinon'
+import { InstalledClock } from '@sinonjs/fake-timers'
 
 describe('functionUtils', function () {
     it('once()', function () {
@@ -107,3 +109,82 @@ describe('debounce', function () {
         })
     })
 })
+
+// function to test the withRetries method. It passes in a stub function as the argument and has different tests that throw on different iterations
+describe('withRetries', function () {
+    let clock: InstalledClock
+    let fn: SinonStub
+
+    beforeEach(function () {
+        fn = stub()
+        clock = installFakeClock()
+    })
+
+    afterEach(function () {
+        clock.uninstall()
+    })
+
+    it('retries the function until it succeeds, using defaults', async function () {
+        fn.onCall(0).throws()
+        fn.onCall(1).throws()
+        fn.onCall(2).resolves('success')
+        assert.strictEqual(await withRetries(fn), 'success')
+    })
+
+    it('retries the function until it succeeds at the final try', async function () {
+        fn.onCall(0).throws()
+        fn.onCall(1).throws()
+        fn.onCall(2).throws()
+        fn.onCall(3).resolves('success')
+        assert.strictEqual(await withRetries(fn, { maxRetries: 4 }), 'success')
+    })
+
+    it('throws the last error if the function always fails, using defaults', async function () {
+        fn.onCall(0).throws()
+        fn.onCall(1).throws()
+        fn.onCall(2).throws()
+        fn.onCall(3).resolves('unreachable')
+        await assert.rejects(async () => {
+            await withRetries(fn)
+        })
+    })
+
+    it('throws the last error if the function always fails', async function () {
+        fn.onCall(0).throws()
+        fn.onCall(1).throws()
+        fn.onCall(2).throws()
+        fn.onCall(3).throws()
+        fn.onCall(4).resolves('unreachable')
+        await assert.rejects(async () => {
+            await withRetries(fn, { maxRetries: 4 })
+        })
+    })
+
+    it('honors retry delay + backoff multiplier', async function () {
+        fn.onCall(0).throws() // 100ms
+        fn.onCall(1).throws() // 200ms
+        fn.onCall(2).throws() // 400ms
+        fn.onCall(3).resolves('success')
+
+        const res = withRetries(fn, { maxRetries: 4, delay: 100, backoff: 2 })
+
+        // Check the call count after each iteration, ensuring the function is called
+        // after the correct delay between retries.
+        await clock.tickAsync(99)
+        assert.strictEqual(fn.callCount, 1)
+        await clock.tickAsync(1)
+        assert.strictEqual(fn.callCount, 2)
+
+        await clock.tickAsync(199)
+        assert.strictEqual(fn.callCount, 2)
+        await clock.tickAsync(1)
+        assert.strictEqual(fn.callCount, 3)
+
+        await clock.tickAsync(399)
+        assert.strictEqual(fn.callCount, 3)
+        await clock.tickAsync(1)
+        assert.strictEqual(fn.callCount, 4)
+
+        assert.strictEqual(await res, 'success')
+    })
+})