wip cache resource implementation

Author: Will-ShaoHua

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,20 @@ 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, {
+                locked: false,
+                timestamp: 0,
+                result: undefined,
+            })
+        }
+
+        override resourceProvider(): Promise<RegionProfile[]> {
+            return this.profileProvider()
+        }
+    })(this.listRegionProfile.bind(this))
+
     get activeRegionProfile() {
         const conn = this.connectionProvider()
         if (isBuilderIdConnection(conn)) {
@@ -104,6 +119,14 @@ export class RegionProfileManager {
 
     constructor(private readonly connectionProvider: () => Connection | undefined) {}
 
+    async resetCache() {
+        await this.cache.releaseLock()
+    }
+
+    async getProfiles(): Promise<RegionProfile[]> {
+        return this.cache.getResource()
+    }
+
     async listRegionProfile(): Promise<RegionProfile[]> {
         this._profiles = []
 
@@ -238,7 +261,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 +323,7 @@ export class RegionProfileManager {
         const selected = this.activeRegionProfile
         let profiles: RegionProfile[] = []
         try {
-            profiles = await this.listRegionProfile()
+            profiles = await this.getProfiles()
         } catch (e) {
             return [
                 {

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/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/utilities/resourceCache.ts

@@ -0,0 +1,113 @@
+/*!
+ * 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'
+
+interface WithLock {
+    locked: boolean
+    timestamp: number
+}
+
+interface Resource<V> extends WithLock {
+    result: V | undefined
+}
+
+const logger = getLogger()
+
+function now() {
+    return globals.clock.Date.now()
+}
+
+export abstract class CachedResource<V> {
+    constructor(
+        readonly key: globalKey,
+        readonly expiration: number,
+        private readonly defaultValue: Resource<V>
+    ) {}
+
+    abstract resourceProvider(): Promise<V>
+
+    async getResource(): Promise<V> {
+        const resource = await this.readResourceAndLock()
+        // if cache is still fresh, return
+        if (resource && resource.result) {
+            if (now() - resource.timestamp < this.expiration) {
+                logger.info(`cache hit`)
+                // release the lock
+                await globals.globalState.update(this.key, {
+                    ...resource,
+                    locked: false,
+                })
+                return resource.result
+            } else {
+                logger.info(`cache hit but cached value is stale, invoking service API to pull the latest response`)
+            }
+        }
+
+        // catch and error case?
+        logger.info(`cache miss, invoking service API to pull the latest response`)
+        const latest = await this.resourceProvider()
+
+        // update resource cache and release the lock
+        const toUpdate: Resource<V> = {
+            locked: false,
+            timestamp: now(),
+            result: latest,
+        }
+        await globals.globalState.update(this.key, toUpdate)
+        return latest
+    }
+
+    async readResourceAndLock(): Promise<Resource<V> | undefined> {
+        const _acquireLock = async () => {
+            const cachedValue = this.readCacheOrDefault()
+
+            if (!cachedValue.locked) {
+                await globals.globalState.update(this.key, {
+                    ...cachedValue,
+                    locked: true,
+                })
+
+                return cachedValue
+            }
+
+            return undefined
+        }
+
+        const lock = await waitUntil(
+            async () => {
+                const lock = await _acquireLock()
+                logger.info(`try obtaining cache lock %s`, lock)
+                if (lock) {
+                    return lock
+                }
+            },
+            { timeout: 15000, interval: 1500, truthy: true } // TODO: pass via ctor
+        )
+
+        return lock
+    }
+
+    async releaseLock() {
+        await globals.globalState.update(this.key, {
+            ...this.readCacheOrDefault(),
+            locked: false,
+        })
+    }
+
+    private readCacheOrDefault(): Resource<V> {
+        const cachedValue = globals.globalState.tryGet<Resource<V>>(this.key, Object, {
+            ...this.defaultValue,
+            locked: false,
+            result: undefined,
+            timestamp: 0,
+        })
+
+        return cachedValue
+    }
+}