fix(amazonq): listAvailableProfile calls are throttled when users have 2+ IDE instances

packages/core/src/codewhisperer/region/regionProfileManager.ts

@@ -28,6 +28,7 @@ import { parse } from '@aws-sdk/util-arn-parser'
 import { isAwsError, ToolkitError } from '../../shared/errors'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { localize } from '../../shared/utilities/vsCodeUtils'
+import { waitUntil } from '../../shared/utilities/timeoutUtils'
 
 // TODO: is there a better way to manage all endpoint strings in one place?
 export const defaultServiceConfig: CodeWhispererConfig = {
@@ -42,13 +43,41 @@ const endpoints = createConstantMap({
 })
 
 /**
- * 'user' -> users change the profile through Q menu
+ * 'user' -> users change the profile through Q menu`
  * 'auth' -> users change the profile through webview profile selector page
  * 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile
  * 'reload' -> on plugin restart, plugin will try to reload previous selected profile
  */
 export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
 
+// Only "valid" state will have non null profiles
+type CachedApiResultWithLock = {
+    // Lock
+    isAcquired: boolean
+    // Metadata
+    result: SuccessResult | FailureResult | undefined
+}
+
+interface ApiResult {
+    type: 'success' | 'failure'
+    timestamp: number
+}
+
+interface SuccessResult extends ApiResult {
+    type: 'success'
+    profiles: RegionProfile[]
+}
+
+interface FailureResult extends ApiResult {
+    type: 'failure'
+    // Error type can't be serialized
+    errorMsg: string
+}
+
+function now(): number {
+    return globals.clock.Date.now()
+}
+
 export class RegionProfileManager {
     private static logger = getLogger()
     private _activeRegionProfile: RegionProfile | undefined
@@ -110,7 +139,56 @@ export class RegionProfileManager {
         if (conn === undefined || !isSsoConnection(conn)) {
             return []
         }
+
+        // WaitUntil lock is acquired, the API tps is low, server can't afford multiple IDE instances making the calls simultaneiously
+        const cached = await waitUntil(
+            async () => {
+                const lock = await this.tryAcquireLock()
+                RegionProfileManager.logger.info(`try obtaining cache lock %s`, lock)
+                if (lock) {
+                    return lock
+                }
+            },
+            { timeout: 15000, interval: 1500, truthy: true }
+        )
+
+        RegionProfileManager.logger.info(`obtained cache lock %s`, cached)
+
         const availableProfiles: RegionProfile[] = []
+        // Load cache and use it directly if it's in "valid" state
+        // If it's pending, keep waiting until the ongoing pull finishes its job
+
+        const now = globals.clock.Date.now()
+        if (cached?.result && now - cached.result.timestamp < 60000) {
+            RegionProfileManager.logger.info(
+                `cache hit, duration=%s, previous listAvailableProfile result=%s`,
+                now - cached.result.timestamp,
+                cached
+            )
+
+            // Release lock so that other ide instances can continue
+            await this.releaeLock(cached)
+
+            if (cached.result.type === 'success') {
+                availableProfiles.push(...(cached.result as SuccessResult).profiles)
+                this._profiles = availableProfiles
+                return availableProfiles
+            } else if (cached.result.type === 'failure') {
+                const errorMsg = (cached.result as FailureResult).errorMsg
+                throw new Error(errorMsg)
+            }
+        }
+
+        if (cached?.result && now - cached.result.timestamp > 60000) {
+            RegionProfileManager.logger.info(
+                `listAvailableProfile cache hit but cached value is stale, invoking service API to pull the latest response`
+            )
+        } else {
+            RegionProfileManager.logger.info(
+                `listAvailableProfile cache miss, invoking service API to pull the latest response`
+            )
+        }
+
         for (const [region, endpoint] of endpoints.entries()) {
             const client = await this.createQClient(region, endpoint, conn as SsoConnection)
             const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
@@ -138,12 +216,17 @@ export class RegionProfileManager {
             } catch (e) {
                 const logMsg = isAwsError(e) ? `requestId=${e.requestId}; message=${e.message}` : (e as Error).message
                 RegionProfileManager.logger.error(`failed to listRegionProfile: ${logMsg}`)
+                // Should update cache even for failure path
+                await this.updateCacheWithError(e as Error)
                 throw e
             }
 
             RegionProfileManager.logger.info(`available amazonq profiles: ${availableProfiles.length}`)
         }
 
+        // Update cache and save the result
+        await this.updateCacheWithResult(availableProfiles)
+
         this._profiles = availableProfiles
         return availableProfiles
     }
@@ -367,4 +450,63 @@ export class RegionProfileManager {
 
         return c
     }
+
+    // Should reset cache when users signout in case users want to change to a different connection
+    async resetCache() {
+        await globals.globalState.update('aws.amazonq.regionProfiles.cachedResult', {
+            isAcquired: false,
+            result: undefined,
+        })
+    }
+
+    private async tryAcquireLock(): Promise<CachedApiResultWithLock | undefined> {
+        const cachedValue = globals.globalState.tryGet<CachedApiResultWithLock>(
+            'aws.amazonq.regionProfiles.cachedResult',
+            Object,
+            { isAcquired: false, result: undefined }
+        )
+
+        if (!cachedValue.isAcquired) {
+            await globals.globalState.update('aws.amazonq.regionProfiles.cachedResult', {
+                ...cachedValue,
+                isAcquired: true,
+            })
+
+            return cachedValue
+        }
+        return undefined
+    }
+
+    private async releaeLock(cached: CachedApiResultWithLock) {
+        await globals.globalState.update('aws.amazonq.regionProfiles.cachedResult', {
+            ...cached,
+            isAcquired: false,
+        })
+    }
+
+    private async updateCacheWithResult(r: RegionProfile[]) {
+        const result: SuccessResult = {
+            type: 'success',
+            profiles: r,
+            timestamp: now(),
+        }
+        const pojo: CachedApiResultWithLock = {
+            isAcquired: false, // release the lock
+            result: result,
+        }
+        await globals.globalState.update('aws.amazonq.regionProfiles.cachedResult', pojo)
+    }
+
+    private async updateCacheWithError(e: Error) {
+        const result: FailureResult = {
+            type: 'failure',
+            errorMsg: e.message,
+            timestamp: now(),
+        }
+        const pojo: CachedApiResultWithLock = {
+            isAcquired: false, // release the lock
+            result: result,
+        }
+        await globals.globalState.update('aws.amazonq.regionProfiles.cachedResult', pojo)
+    }
 }

packages/core/src/codewhisperer/util/authUtil.ts

@@ -142,6 +142,7 @@ export class AuthUtil {
 
             if (!this.isConnected()) {
                 await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
+                await this.regionProfileManager.resetCache()
             }
         })
 

packages/core/src/shared/globalState.ts

@@ -49,6 +49,7 @@ export type globalKey =
     | 'aws.toolkit.lsp.manifest'
     | 'aws.amazonq.customization.overrideV2'
     | 'aws.amazonq.regionProfiles'
+    | 'aws.amazonq.regionProfiles.cachedResult'
     // Deprecated/legacy names. New keys should start with "aws.".
     | '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
     | 'CODECATALYST_RECONNECT'