fix(amazonq): fix enterprise users not able to sign in correctly if they have 2+ vscode instances open

packages/amazonq/.changes/next-release/Bug Fix-8290a06f-ee3f-4235-b8af-faedb1bdfbe4.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix users can not log in successfully with 2+ IDE instnaces open due to throttle error throw by the service"
+}

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

@@ -29,6 +29,7 @@ import { isAwsError, ToolkitError } from '../../shared/errors'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { localize } from '../../shared/utilities/vsCodeUtils'
 import { Commands } from '../../shared/vscode/commands2'
+import { CachedResource } from '../../shared/utilities/resourceCache'
 
 // TODO: is there a better way to manage all endpoint strings in one place?
 export const defaultServiceConfig: CodeWhispererConfig = {
@@ -59,6 +60,27 @@ export class RegionProfileManager {
     // Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
     private _profiles: RegionProfile[] = []
 
+    private readonly cache = new (class extends CachedResource<RegionProfile[]> {
+        constructor(private readonly profileProvider: () => Promise<RegionProfile[]>) {
+            super(
+                'aws.amazonq.regionProfiles.cache',
+                60000,
+                {
+                    resource: {
+                        locked: false,
+                        timestamp: 0,
+                        result: undefined,
+                    },
+                },
+                { timeout: 15000, interval: 1500, truthy: true }
+            )
+        }
+
+        override resourceProvider(): Promise<RegionProfile[]> {
+            return this.profileProvider()
+        }
+    })(this.listRegionProfile.bind(this))
+
     get activeRegionProfile() {
         const conn = this.connectionProvider()
         if (isBuilderIdConnection(conn)) {
@@ -104,6 +126,10 @@ export class RegionProfileManager {
 
     constructor(private readonly connectionProvider: () => Connection | undefined) {}
 
+    async getProfiles(): Promise<RegionProfile[]> {
+        return this.cache.getResource()
+    }
+
     async listRegionProfile(): Promise<RegionProfile[]> {
         this._profiles = []
 
@@ -238,7 +264,7 @@ export class RegionProfileManager {
             return
         }
         // cross-validation
-        this.listRegionProfile()
+        this.getProfiles()
             .then(async (profiles) => {
                 const r = profiles.find((it) => it.arn === previousSelected.arn)
                 if (!r) {
@@ -300,7 +326,7 @@ export class RegionProfileManager {
         const selected = this.activeRegionProfile
         let profiles: RegionProfile[] = []
         try {
-            profiles = await this.listRegionProfile()
+            profiles = await this.getProfiles()
         } catch (e) {
             return [
                 {
@@ -347,6 +373,11 @@ export class RegionProfileManager {
         }
     }
 
+    // Should be called on connection changed in case users change to a differnet connection and use the wrong resultset.
+    async clearCache() {
+        await this.cache.clearCache()
+    }
+
     async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
         const token = (await conn.getToken()).accessToken
         const serviceOption: ServiceOptions = {

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

@@ -143,6 +143,8 @@ export class AuthUtil {
             if (!this.isConnected()) {
                 await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
             }
+
+            await this.regionProfileManager.clearCache()
         })
 
         this.regionProfileManager.onDidChangeRegionProfile(async () => {

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -223,7 +223,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
      */
     override async listRegionProfiles(): Promise<RegionProfile[] | string> {
         try {
-            return await AuthUtil.instance.regionProfileManager.listRegionProfile()
+            return await AuthUtil.instance.regionProfileManager.getProfiles()
         } catch (e) {
             const conn = AuthUtil.instance.conn as SsoConnection | undefined
             telemetry.amazonq_didSelectProfile.emit({

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.cache'
     // Deprecated/legacy names. New keys should start with "aws.".
     | '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
     | 'CODECATALYST_RECONNECT'

packages/core/src/shared/logger/logger.ts

@@ -18,6 +18,7 @@ export type LogTopic =
     | 'chat'
     | 'stepfunctions'
     | 'unknown'
+    | 'resourceCache'
 
 class ErrorLog {
     constructor(

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

@@ -0,0 +1,193 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import globals from '../extensionGlobals'
+import { globalKey } from '../globalState'
+import { getLogger } from '../logger/logger'
+import { waitUntil } from '../utilities/timeoutUtils'
+
+/**
+ * args:
+ *  @member result: the actual resource type callers want to use
+ *  @member locked: readWriteLock, while the lock is acquired by one process, the other can't access to it until it's released by the previous
+ *  @member timestamp: used for determining the resource is stale or not
+ */
+interface Resource<V> {
+    result: V | undefined
+    locked: boolean
+    timestamp: number
+}
+
+/**
+ * GlobalStates schema, which is used for vscode global states deserialization, [globals.globalState#tryGet<T>] etc.
+ * The purpose of it is to allow devs to overload the resource into existing global key and no need to create a specific key for only this purpose.
+ */
+export interface GlobalStateSchema<V> {
+    resource: Resource<V>
+}
+
+const logger = getLogger('resourceCache')
+
+function now() {
+    return globals.clock.Date.now()
+}
+
+/**
+ * CacheResource utilizes VSCode global states API to cache resources which are expensive to get so that the result can be shared across multiple VSCode instances.
+ *  The first VSCode instance invoking #getResource will hold a lock and make the actual network call/FS read to pull the real response.
+ *  When the pull is done, the lock will be released and it then caches the result in the global states. Then the rest of instances can now acquire the lock 1 by 1 and read the resource from the cache.
+ *
+ * constructor:
+ *  @param key: global state key, which is used for globals.globalState#update, #tryGet etc.
+ *  @param expirationInMilli: cache expiration time in milli seconds
+ *  @param defaultValue: default value for the cache if the cache doesn't pre-exist in users' FS
+ *  @param waitUntilOption: waitUntil option for acquire lock
+ *
+ * methods:
+ *  @method resourceProvider: implementation needs to implement this method to obtain the latest resource either via network calls or FS read
+ *  @method getResource: obtain the resource from cache or pull the latest from the service if the cache either expires or doesn't exist
+ */
+export abstract class CachedResource<V> {
+    constructor(
+        private readonly key: globalKey,
+        private readonly expirationInMilli: number,
+        private readonly defaultValue: GlobalStateSchema<V>,
+        private readonly waitUntilOption: { timeout: number; interval: number; truthy: boolean }
+    ) {}
+
+    abstract resourceProvider(): Promise<V>
+
+    async getResource(): Promise<V> {
+        const cachedValue = await this.tryLoadResourceAndLock()
+        const resource = cachedValue?.resource
+
+        // If cache is still fresh, return cached result, otherwise pull latest from the service
+        if (cachedValue && resource && resource.result) {
+            const duration = now() - resource.timestamp
+            if (duration < this.expirationInMilli) {
+                logger.debug(
+                    `cache hit, duration(%sms) is less than expiration(%sms), returning cached value %s`,
+                    duration,
+                    this.expirationInMilli,
+                    this.key
+                )
+                // release the lock
+                await this.releaseLock(resource, cachedValue)
+                return resource.result
+            } else {
+                logger.debug(
+                    `cache is stale, duration(%sms) is older than expiration(%sms), pulling latest resource %s`,
+                    duration,
+                    this.expirationInMilli,
+                    this.key
+                )
+            }
+        } else {
+            logger.debug(`cache miss, pulling latest resource %s`, this.key)
+        }
+
+        /**
+         * Possible paths here
+         *  1. cache doesn't exist.
+         *  2. cache exists but expired.
+         *  3. lock is held by other process and the waiting time is greater than the specified waiting time
+         */
+        try {
+            // Make the real network call / FS read to pull the resource
+            const latest = await this.resourceProvider()
+
+            // Update resource cache and release the lock
+            const r: Resource<V> = {
+                locked: false,
+                timestamp: now(),
+                result: latest,
+            }
+            logger.debug(`doen loading latest resource, updating resource cache: %s`, this.key)
+            await this.releaseLock(r)
+            return latest
+        } catch (e) {
+            logger.debug(
+                `encountered unexpected error while loading the latest of resource(%s), releasing resource lock`,
+                this.key
+            )
+            await this.releaseLock()
+            throw e
+        }
+    }
+
+    // This method will lock the resource so other callers have to wait until the lock is released, otherwise will return undefined if it times out
+    private async tryLoadResourceAndLock(): Promise<GlobalStateSchema<V> | undefined> {
+        const _acquireLock = async () => {
+            const cachedValue = this.readCacheOrDefault()
+
+            if (!cachedValue.resource.locked) {
+                await this.lockResource(cachedValue)
+                return cachedValue
+            }
+
+            return undefined
+        }
+
+        const lock = await waitUntil(async () => {
+            const lock = await _acquireLock()
+            logger.debug(`try obtain resource cache read write lock for resource %s`, this.key)
+            if (lock) {
+                return lock
+            }
+        }, this.waitUntilOption)
+
+        return lock
+    }
+
+    async lockResource(baseCache: GlobalStateSchema<V>): Promise<void> {
+        await this.updateResourceCache({ locked: true }, baseCache)
+    }
+
+    async releaseLock(): Promise<void>
+    async releaseLock(resource: Partial<Resource<V>>): Promise<void>
+    async releaseLock(resource: Partial<Resource<V>>, baseCache: GlobalStateSchema<V>): Promise<void>
+    async releaseLock(resource?: Partial<Resource<V>>, baseCache?: GlobalStateSchema<V>): Promise<void> {
+        if (!resource) {
+            await this.updateResourceCache({ locked: false }, undefined)
+        } else if (baseCache) {
+            await this.updateResourceCache(resource, baseCache)
+        } else {
+            await this.updateResourceCache(resource, undefined)
+        }
+    }
+
+    async clearCache() {
+        const baseCache = this.readCacheOrDefault()
+        await this.updateResourceCache({ result: undefined, timestamp: 0, locked: false }, baseCache)
+    }
+
+    private async updateResourceCache(resource: Partial<Resource<any>>, cache: GlobalStateSchema<any> | undefined) {
+        const baseCache = cache ?? this.readCacheOrDefault()
+
+        const toUpdate: GlobalStateSchema<V> = {
+            ...baseCache,
+            resource: {
+                ...baseCache.resource,
+                ...resource,
+            },
+        }
+
+        await globals.globalState.update(this.key, toUpdate)
+    }
+
+    private readCacheOrDefault(): GlobalStateSchema<V> {
+        const cachedValue = globals.globalState.tryGet<GlobalStateSchema<V>>(this.key, Object, {
+            ...this.defaultValue,
+            resource: {
+                ...this.defaultValue.resource,
+                locked: false,
+                result: undefined,
+                timestamp: 0,
+            },
+        })
+
+        return cachedValue
+    }
+}