Skip to content

Commit 3549e33

Browse files
authored
fix(amazonq): fix enterprise users not able to sign in correctly if they have 2+ vscode instances open (#7151)
## Problem revision of #7134, this pr aims to address the comment from the previous PR #7134 (comment) to extract the logic to a shared module so the diff against 7134 is expected to be large. after deployment 4/21, service has a strict 1 tps throttling policy and there was no caching for the API call previously. It will impact users as long as they have multiple ide instances opened as all of them will make list profile call while users attempt to sign in. ## Solution 1. cache the api result for 60 seconds for reuse 2. use lock to prevent multiple instances trying to call at the same time. Only 1 will make the service call and the rest will wait until it's done and read from the cache. --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 638778b commit 3549e33

File tree

7 files changed

+232
-3
lines changed

7 files changed

+232
-3
lines changed
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Bug Fix",
3+
"description": "Fix users can not log in successfully with 2+ IDE instnaces open due to throttle error throw by the service"
4+
}

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

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import { isAwsError, ToolkitError } from '../../shared/errors'
2929
import { telemetry } from '../../shared/telemetry/telemetry'
3030
import { localize } from '../../shared/utilities/vsCodeUtils'
3131
import { Commands } from '../../shared/vscode/commands2'
32+
import { CachedResource } from '../../shared/utilities/resourceCache'
3233

3334
// TODO: is there a better way to manage all endpoint strings in one place?
3435
export const defaultServiceConfig: CodeWhispererConfig = {
@@ -59,6 +60,27 @@ export class RegionProfileManager {
5960
// Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
6061
private _profiles: RegionProfile[] = []
6162

63+
private readonly cache = new (class extends CachedResource<RegionProfile[]> {
64+
constructor(private readonly profileProvider: () => Promise<RegionProfile[]>) {
65+
super(
66+
'aws.amazonq.regionProfiles.cache',
67+
60000,
68+
{
69+
resource: {
70+
locked: false,
71+
timestamp: 0,
72+
result: undefined,
73+
},
74+
},
75+
{ timeout: 15000, interval: 1500, truthy: true }
76+
)
77+
}
78+
79+
override resourceProvider(): Promise<RegionProfile[]> {
80+
return this.profileProvider()
81+
}
82+
})(this.listRegionProfile.bind(this))
83+
6284
get activeRegionProfile() {
6385
const conn = this.connectionProvider()
6486
if (isBuilderIdConnection(conn)) {
@@ -104,6 +126,10 @@ export class RegionProfileManager {
104126

105127
constructor(private readonly connectionProvider: () => Connection | undefined) {}
106128

129+
async getProfiles(): Promise<RegionProfile[]> {
130+
return this.cache.getResource()
131+
}
132+
107133
async listRegionProfile(): Promise<RegionProfile[]> {
108134
this._profiles = []
109135

@@ -238,7 +264,7 @@ export class RegionProfileManager {
238264
return
239265
}
240266
// cross-validation
241-
this.listRegionProfile()
267+
this.getProfiles()
242268
.then(async (profiles) => {
243269
const r = profiles.find((it) => it.arn === previousSelected.arn)
244270
if (!r) {
@@ -300,7 +326,7 @@ export class RegionProfileManager {
300326
const selected = this.activeRegionProfile
301327
let profiles: RegionProfile[] = []
302328
try {
303-
profiles = await this.listRegionProfile()
329+
profiles = await this.getProfiles()
304330
} catch (e) {
305331
return [
306332
{
@@ -347,6 +373,11 @@ export class RegionProfileManager {
347373
}
348374
}
349375

376+
// Should be called on connection changed in case users change to a differnet connection and use the wrong resultset.
377+
async clearCache() {
378+
await this.cache.clearCache()
379+
}
380+
350381
async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
351382
const token = (await conn.getToken()).accessToken
352383
const serviceOption: ServiceOptions = {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,8 @@ export class AuthUtil {
144144
if (!this.isConnected()) {
145145
await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn)
146146
}
147+
148+
await this.regionProfileManager.clearCache()
147149
})
148150

149151
this.regionProfileManager.onDidChangeRegionProfile(async () => {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
223223
*/
224224
override async listRegionProfiles(): Promise<RegionProfile[] | string> {
225225
try {
226-
return await AuthUtil.instance.regionProfileManager.listRegionProfile()
226+
return await AuthUtil.instance.regionProfileManager.getProfiles()
227227
} catch (e) {
228228
const conn = AuthUtil.instance.conn as SsoConnection | undefined
229229
telemetry.amazonq_didSelectProfile.emit({

packages/core/src/shared/globalState.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ export type globalKey =
4949
| 'aws.toolkit.lsp.manifest'
5050
| 'aws.amazonq.customization.overrideV2'
5151
| 'aws.amazonq.regionProfiles'
52+
| 'aws.amazonq.regionProfiles.cache'
5253
// Deprecated/legacy names. New keys should start with "aws.".
5354
| '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
5455
| 'CODECATALYST_RECONNECT'

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ export type LogTopic =
1818
| 'chat'
1919
| 'stepfunctions'
2020
| 'unknown'
21+
| 'resourceCache'
2122

2223
class ErrorLog {
2324
constructor(
Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import globals from '../extensionGlobals'
7+
import { globalKey } from '../globalState'
8+
import { getLogger } from '../logger/logger'
9+
import { waitUntil } from '../utilities/timeoutUtils'
10+
11+
/**
12+
* args:
13+
* @member result: the actual resource type callers want to use
14+
* @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
15+
* @member timestamp: used for determining the resource is stale or not
16+
*/
17+
interface Resource<V> {
18+
result: V | undefined
19+
locked: boolean
20+
timestamp: number
21+
}
22+
23+
/**
24+
* GlobalStates schema, which is used for vscode global states deserialization, [globals.globalState#tryGet<T>] etc.
25+
* 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.
26+
*/
27+
export interface GlobalStateSchema<V> {
28+
resource: Resource<V>
29+
}
30+
31+
const logger = getLogger('resourceCache')
32+
33+
function now() {
34+
return globals.clock.Date.now()
35+
}
36+
37+
/**
38+
* 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.
39+
* The first VSCode instance invoking #getResource will hold a lock and make the actual network call/FS read to pull the real response.
40+
* 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.
41+
*
42+
* constructor:
43+
* @param key: global state key, which is used for globals.globalState#update, #tryGet etc.
44+
* @param expirationInMilli: cache expiration time in milli seconds
45+
* @param defaultValue: default value for the cache if the cache doesn't pre-exist in users' FS
46+
* @param waitUntilOption: waitUntil option for acquire lock
47+
*
48+
* methods:
49+
* @method resourceProvider: implementation needs to implement this method to obtain the latest resource either via network calls or FS read
50+
* @method getResource: obtain the resource from cache or pull the latest from the service if the cache either expires or doesn't exist
51+
*/
52+
export abstract class CachedResource<V> {
53+
constructor(
54+
private readonly key: globalKey,
55+
private readonly expirationInMilli: number,
56+
private readonly defaultValue: GlobalStateSchema<V>,
57+
private readonly waitUntilOption: { timeout: number; interval: number; truthy: boolean }
58+
) {}
59+
60+
abstract resourceProvider(): Promise<V>
61+
62+
async getResource(): Promise<V> {
63+
const cachedValue = await this.tryLoadResourceAndLock()
64+
const resource = cachedValue?.resource
65+
66+
// If cache is still fresh, return cached result, otherwise pull latest from the service
67+
if (cachedValue && resource && resource.result) {
68+
const duration = now() - resource.timestamp
69+
if (duration < this.expirationInMilli) {
70+
logger.debug(
71+
`cache hit, duration(%sms) is less than expiration(%sms), returning cached value %s`,
72+
duration,
73+
this.expirationInMilli,
74+
this.key
75+
)
76+
// release the lock
77+
await this.releaseLock(resource, cachedValue)
78+
return resource.result
79+
} else {
80+
logger.debug(
81+
`cache is stale, duration(%sms) is older than expiration(%sms), pulling latest resource %s`,
82+
duration,
83+
this.expirationInMilli,
84+
this.key
85+
)
86+
}
87+
} else {
88+
logger.info(`cache miss, pulling latest resource %s`, this.key)
89+
}
90+
91+
/**
92+
* Possible paths here
93+
* 1. cache doesn't exist.
94+
* 2. cache exists but expired.
95+
* 3. lock is held by other process and the waiting time is greater than the specified waiting time
96+
*/
97+
try {
98+
// Make the real network call / FS read to pull the resource
99+
const latest = await this.resourceProvider()
100+
101+
// Update resource cache and release the lock
102+
const r: Resource<V> = {
103+
locked: false,
104+
timestamp: now(),
105+
result: latest,
106+
}
107+
logger.info(`doen loading latest resource, updating resource cache: %s`, this.key)
108+
await this.releaseLock(r)
109+
return latest
110+
} catch (e) {
111+
logger.error(`failed to load latest resource, releasing lock: %s`, this.key)
112+
await this.releaseLock()
113+
throw e
114+
}
115+
}
116+
117+
// 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
118+
private async tryLoadResourceAndLock(): Promise<GlobalStateSchema<V> | undefined> {
119+
const _acquireLock = async () => {
120+
const cachedValue = this.readCacheOrDefault()
121+
122+
if (!cachedValue.resource.locked) {
123+
await this.lockResource(cachedValue)
124+
return cachedValue
125+
}
126+
127+
return undefined
128+
}
129+
130+
const lock = await waitUntil(async () => {
131+
const lock = await _acquireLock()
132+
logger.debug(`try obtain resource cache read write lock for resource %s`, this.key)
133+
if (lock) {
134+
return lock
135+
}
136+
}, this.waitUntilOption)
137+
138+
return lock
139+
}
140+
141+
async lockResource(baseCache: GlobalStateSchema<V>): Promise<void> {
142+
await this.updateResourceCache({ locked: true }, baseCache)
143+
}
144+
145+
async releaseLock(): Promise<void>
146+
async releaseLock(resource: Partial<Resource<V>>): Promise<void>
147+
async releaseLock(resource: Partial<Resource<V>>, baseCache: GlobalStateSchema<V>): Promise<void>
148+
async releaseLock(resource?: Partial<Resource<V>>, baseCache?: GlobalStateSchema<V>): Promise<void> {
149+
if (!resource) {
150+
await this.updateResourceCache({ locked: false }, undefined)
151+
} else if (baseCache) {
152+
await this.updateResourceCache(resource, baseCache)
153+
} else {
154+
await this.updateResourceCache(resource, undefined)
155+
}
156+
}
157+
158+
async clearCache() {
159+
const baseCache = this.readCacheOrDefault()
160+
await this.updateResourceCache({ result: undefined, timestamp: 0, locked: false }, baseCache)
161+
}
162+
163+
private async updateResourceCache(resource: Partial<Resource<any>>, cache: GlobalStateSchema<any> | undefined) {
164+
const baseCache = cache ?? this.readCacheOrDefault()
165+
166+
const toUpdate: GlobalStateSchema<V> = {
167+
...baseCache,
168+
resource: {
169+
...baseCache.resource,
170+
...resource,
171+
},
172+
}
173+
174+
await globals.globalState.update(this.key, toUpdate)
175+
}
176+
177+
private readCacheOrDefault(): GlobalStateSchema<V> {
178+
const cachedValue = globals.globalState.tryGet<GlobalStateSchema<V>>(this.key, Object, {
179+
...this.defaultValue,
180+
resource: {
181+
...this.defaultValue.resource,
182+
locked: false,
183+
result: undefined,
184+
timestamp: 0,
185+
},
186+
})
187+
188+
return cachedValue
189+
}
190+
}

0 commit comments

Comments
 (0)