feat(amazonq): allow child processes to define custom thresholds

packages/amazonq/src/lsp/client.ts

@@ -35,6 +35,7 @@ import {
     isAmazonInternalOs,
     fs,
 } from 'aws-core-vscode/shared'
+import { processUtils } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
@@ -80,11 +81,13 @@ export async function startLanguageServer(
         executable = [resourcePaths.node]
     }
 
+    const memoryWarnThreshold = 1024 * processUtils.oneMB
     const serverOptions = createServerOptions({
         encryptionKey,
         executable: executable,
         serverModule,
         execArgv: argv,
+        warnThresholds: { memory: memoryWarnThreshold },
     })
 
     await validateNodeExe(executable, resourcePaths.lsp, argv, logger)

packages/core/src/amazonq/lsp/lspClient.ts

@@ -8,6 +8,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  */
 import * as vscode from 'vscode'
+import { oneMB } from '../../shared/utilities/processUtils'
 import * as path from 'path'
 import * as nls from 'vscode-nls'
 import * as crypto from 'crypto'
@@ -252,13 +253,15 @@ export async function activate(extensionContext: ExtensionContext, resourcePaths
     }
 
     const serverModule = resourcePaths.lsp
+    const memoryWarnThreshold = 800 * oneMB
 
     const serverOptions = createServerOptions({
         encryptionKey: key,
         executable: [resourcePaths.node],
         serverModule,
         // TODO(jmkeyes): we always use the debug options...?
         execArgv: debugOptions.execArgv,
+        warnThresholds: { memory: memoryWarnThreshold },
     })
 
     const documentSelector = [{ scheme: 'file', language: '*' }]

packages/core/src/shared/lsp/utils/platform.ts

@@ -81,19 +81,21 @@ export function createServerOptions({
     executable,
     serverModule,
     execArgv,
+    warnThresholds,
 }: {
     encryptionKey: Buffer
     executable: string[]
     serverModule: string
     execArgv: string[]
+    warnThresholds?: { cpu?: number; memory?: number }
 }) {
     return async () => {
         const bin = executable[0]
         const args = [...executable.slice(1), serverModule, ...execArgv]
         if (isDebugInstance()) {
             args.unshift('--inspect=6080')
         }
-        const lspProcess = new ChildProcess(bin, args)
+        const lspProcess = new ChildProcess(bin, args, { warnThresholds })
 
         // this is a long running process, awaiting it will never resolve
         void lspProcess.run()

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

@@ -44,6 +44,13 @@ export interface ChildProcessOptions {
     onStdout?: (text: string, context: RunParameterContext) => void
     /** Callback for intercepting text from the stderr stream. */
     onStderr?: (text: string, context: RunParameterContext) => void
+    /** Thresholds to configure warning logs */
+    warnThresholds?: {
+        /** Threshold for memory usage in bytes */
+        memory?: number
+        /** Threshold for CPU usage by percentage */
+        cpu?: number
+    }
 }
 
 export interface ChildProcessRunOptions extends Omit<ChildProcessOptions, 'logging'> {
@@ -60,19 +67,19 @@ export interface ChildProcessResult {
     stderr: string
     signal?: string
 }
-
+export const oneMB = 1024 * 1024
 export const eof = Symbol('EOF')
+export const defaultProcessWarnThresholds = {
+    memory: 100 * oneMB,
+    cpu: 50,
+}
 
 export interface ProcessStats {
     memory: number
     cpu: number
 }
 export class ChildProcessTracker {
     static readonly pollingInterval: number = 10000 // Check usage every 10 seconds
-    static readonly thresholds: ProcessStats = {
-        memory: 100 * 1024 * 1024, // 100 MB
-        cpu: 50,
-    }
     static readonly logger = logger.getLogger('childProcess')
     static readonly loggedPids = new CircularBuffer(1000)
     #processByPid: Map<number, ChildProcess> = new Map<number, ChildProcess>()
@@ -82,6 +89,15 @@ export class ChildProcessTracker {
         this.#pids = new PollingSet(ChildProcessTracker.pollingInterval, () => this.monitor())
     }
 
+    private getThreshold(pid: number): ProcessStats {
+        if (!this.#processByPid.has(pid)) {
+            ChildProcessTracker.logOnce(pid, `Missing process with id ${pid}, returning default threshold`)
+            return defaultProcessWarnThresholds
+        }
+        // Safe to assert since it exists from check above.
+        return this.#processByPid.get(pid)!.getWarnThresholds()
+    }
+
     private cleanUp() {
         const terminatedProcesses = Array.from(this.#pids.values()).filter(
             (pid: number) => this.#processByPid.get(pid)?.stopped
@@ -106,13 +122,17 @@ export class ChildProcessTracker {
             return
         }
         const stats = this.getUsage(pid)
+        const threshold = this.getThreshold(pid)
         if (stats) {
             ChildProcessTracker.logger.debug(`Process ${pid} usage: %O`, stats)
-            if (stats.memory > ChildProcessTracker.thresholds.memory) {
-                ChildProcessTracker.logOnce(pid, `Process ${pid} exceeded memory threshold: ${stats.memory}`)
+            if (stats.memory > threshold.memory) {
+                ChildProcessTracker.logOnce(
+                    pid,
+                    `Process ${pid} exceeded memory threshold: ${(stats.memory / oneMB).toFixed(2)} MB`
+                )
             }
-            if (stats.cpu > ChildProcessTracker.thresholds.cpu) {
-                ChildProcessTracker.logOnce(pid, `Process ${pid} exceeded cpu threshold: ${stats.cpu}`)
+            if (stats.cpu > threshold.cpu) {
+                ChildProcessTracker.logOnce(pid, `Process ${pid} exceeded cpu threshold: ${stats.cpu}%`)
             }
         }
     }
@@ -248,6 +268,10 @@ export class ChildProcess {
         return await new ChildProcess(command, args, options).run()
     }
 
+    public getWarnThresholds(): ProcessStats {
+        return { ...defaultProcessWarnThresholds, ...this.#baseOptions.warnThresholds }
+    }
+
     // Inspired by 'got'
     /**
      * Creates a one-off {@link ChildProcess} class that always uses the specified options.

packages/core/src/test/shared/utilities/processUtils.test.ts

@@ -10,8 +10,10 @@ import * as sinon from 'sinon'
 import { makeTemporaryToolkitFolder, tryRemoveFolder } from '../../../shared/filesystemUtilities'
 import {
     ChildProcess,
+    ChildProcessOptions,
     ChildProcessResult,
     ChildProcessTracker,
+    defaultProcessWarnThresholds,
     eof,
     ProcessStats,
 } from '../../../shared/utilities/processUtils'
@@ -376,8 +378,8 @@ async function stopAndWait(runningProcess: RunningProcess): Promise<void> {
     await runningProcess.result
 }
 
-function startSleepProcess(timeout: number = 90): RunningProcess {
-    const childProcess = new ChildProcess(getSleepCmd(), [timeout.toString()])
+function startSleepProcess(options?: ChildProcessOptions, timeout: number = 90): RunningProcess {
+    const childProcess = new ChildProcess(getSleepCmd(), [timeout.toString()], options)
     const result = childProcess.run().catch(() => assert.fail('sleep command threw an error'))
     return { childProcess, result }
 }
@@ -454,12 +456,12 @@ describe('ChildProcessTracker', function () {
         tracker.add(runningProcess.childProcess)
 
         const highCpu: ProcessStats = {
-            cpu: ChildProcessTracker.thresholds.cpu + 1,
+            cpu: defaultProcessWarnThresholds.cpu + 1,
             memory: 0,
         }
         const highMemory: ProcessStats = {
             cpu: 0,
-            memory: ChildProcessTracker.thresholds.memory + 1,
+            memory: defaultProcessWarnThresholds.memory + 1,
         }
 
         usageMock.returns(highCpu)
@@ -480,7 +482,7 @@ describe('ChildProcessTracker', function () {
         tracker.add(runningProcess.childProcess)
 
         usageMock.returns({
-            cpu: ChildProcessTracker.thresholds.cpu + 1,
+            cpu: defaultProcessWarnThresholds.cpu + 1,
             memory: 0,
         })
 
@@ -492,10 +494,11 @@ describe('ChildProcessTracker', function () {
 
     it('does not log for processes within threshold', async function () {
         const runningProcess = startSleepProcess()
+        tracker.add(runningProcess.childProcess)
 
         usageMock.returns({
-            cpu: ChildProcessTracker.thresholds.cpu - 1,
-            memory: ChildProcessTracker.thresholds.memory - 1,
+            cpu: defaultProcessWarnThresholds.cpu - 1,
+            memory: defaultProcessWarnThresholds.memory - 1,
         })
 
         await clock.tickAsync(ChildProcessTracker.pollingInterval)
@@ -504,4 +507,51 @@ describe('ChildProcessTracker', function () {
 
         await stopAndWait(runningProcess)
     })
+
+    it('respects custom thresholds', async function () {
+        const largeRunningProcess = startSleepProcess({
+            warnThresholds: {
+                cpu: defaultProcessWarnThresholds.cpu + 10,
+                memory: defaultProcessWarnThresholds.memory + 10,
+            },
+        })
+        tracker.add(largeRunningProcess.childProcess)
+        const smallRunningProcess = startSleepProcess({
+            warnThresholds: {
+                cpu: defaultProcessWarnThresholds.cpu - 10,
+                memory: defaultProcessWarnThresholds.memory - 10,
+            },
+        })
+        tracker.add(smallRunningProcess.childProcess)
+
+        usageMock.returns({
+            cpu: defaultProcessWarnThresholds.cpu + 5,
+            memory: defaultProcessWarnThresholds.memory + 5,
+        })
+
+        await clock.tickAsync(ChildProcessTracker.pollingInterval)
+        assert.throws(() => assertLogsContain(largeRunningProcess.childProcess.pid().toString(), false, 'warn'))
+        assertLogsContain(smallRunningProcess.childProcess.pid().toString(), false, 'warn')
+
+        await stopAndWait(largeRunningProcess)
+        await stopAndWait(smallRunningProcess)
+    })
+
+    it('fills custom thresholds with default', async function () {
+        const runningProcess = startSleepProcess({
+            warnThresholds: {
+                cpu: defaultProcessWarnThresholds.cpu + 10,
+            },
+        })
+        tracker.add(runningProcess.childProcess)
+
+        usageMock.returns({
+            memory: defaultProcessWarnThresholds.memory + 1,
+        })
+
+        await clock.tickAsync(ChildProcessTracker.pollingInterval)
+        assertLogsContain(runningProcess.childProcess.pid().toString(), false, 'warn')
+
+        await stopAndWait(runningProcess)
+    })
 })