Skip to content

Commit b23cfdd

Browse files
authored
Merge master into feature/hybridChat
2 parents 4ceed6e + 07a70bc commit b23cfdd

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)