From 295c574f52857e595132d6a7635d2add54953882 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Fri, 20 Dec 2024 12:02:45 -0500 Subject: [PATCH 1/5] fix(auth): debounce getToken() function Problem: The Identity team noticed a large spike in token refreshes for specific users. One user would trigger refresh over 50 times within a few seconds. Solution: The telemetry showed that `getChatAuthState()` was being called many times in a short period. This eventually triggered the token refresh logic many times, if the token was expired. The solution is to add a debounce to `getToken()` which calls the refresh logic. - `debounce()` only accepts functions without any args, the refresh logic requires args - `getToken()` will also load from disk is the token is not expired, so debouncing here saves disk I/O as well. The current debounce interval is 100 milliseconds, which based on telemetry should be enough to capture the barrage of calls. Signed-off-by: nkomonen-amazon --- packages/core/src/auth/sso/ssoAccessTokenProvider.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts index 6639e9f3832..aeb6f50010a 100644 --- a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts +++ b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts @@ -33,7 +33,7 @@ import { randomUUID } from '../../shared/crypto' import { getExtRuntimeContext } from '../../shared/vscode/env' import { showInputBox } from '../../shared/ui/inputPrompter' import { AmazonQPromptSettings, DevSettings, PromptSettings, ToolkitPromptSettings } from '../../shared/settings' -import { onceChanged } from '../../shared/utilities/functionUtils' +import { debounce, onceChanged } from '../../shared/utilities/functionUtils' import { NestedMap } from '../../shared/utilities/map' import { asStringifiedStack } from '../../shared/telemetry/spans' import { showViewLogsMessage } from '../../shared/utilities/messages' @@ -97,7 +97,13 @@ export abstract class SsoAccessTokenProvider { this.reAuthState.set(this.profile, { reAuthReason: `invalidate():${reason}` }) } - public async getToken(): Promise { + /** + * Sometimes we get many calls at once and this + * can trigger redundant disk reads, or token refreshes. + * We debounce to avoid this. + */ + public getToken = debounce(this._getToken.bind(this), 100) + public async _getToken(): Promise { const data = await this.cache.token.load(this.tokenCacheKey) SsoAccessTokenProvider.logIfChanged( indent( From 471739c6f7f6c6d6fa2e6ab0c62f0523b54dfd80 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Tue, 7 Jan 2025 18:48:25 -0500 Subject: [PATCH 2/5] fix failing tests Problem: By default a sinon fake clock was installed on all tests. This caused the new debounce functionality on getToken to freeze since the clock was not progressed, multiple tests were failing. Solution: Only use the fake clock in tests that need it Signed-off-by: nkomonen-amazon --- .../sso/ssoAccessTokenProvider.test.ts | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts index f552e8d08c8..2cb98193224 100644 --- a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts +++ b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts @@ -27,6 +27,7 @@ import { ToolkitError } from '../../../shared/errors' import * as fs from 'fs' // eslint-disable-line no-restricted-imports import * as path from 'path' import { Stub, stub } from '../../utilities/stubber' +import { globals } from '../../../shared' const hourInMs = 3600000 @@ -37,14 +38,14 @@ describe('SsoAccessTokenProvider', function () { let oidcClient: Stub let sut: SsoAccessTokenProvider let cache: ReturnType - let clock: FakeTimers.InstalledClock + let clock: FakeTimers.InstalledClock | undefined let tempDir: string let reAuthState: TestReAuthState function createToken(timeDelta: number, extras: Partial = {}) { return { accessToken: 'dummyAccessToken', - expiresAt: new clock.Date(clock.Date.now() + timeDelta), + expiresAt: new globals.clock.Date(globals.clock.Date.now() + timeDelta), ...extras, } } @@ -54,7 +55,7 @@ describe('SsoAccessTokenProvider', function () { scopes: [], clientId: 'dummyClientId', clientSecret: 'dummyClientSecret', - expiresAt: new clock.Date(clock.Date.now() + timeDelta), + expiresAt: new globals.clock.Date(globals.clock.Date.now() + timeDelta), startUrl, ...extras, } @@ -66,7 +67,7 @@ describe('SsoAccessTokenProvider', function () { deviceCode: 'dummyCode', userCode: 'dummyUserCode', verificationUri: 'dummyLink', - expiresAt: new clock.Date(clock.Date.now() + timeDelta), + expiresAt: new globals.clock.Date(globals.clock.Date.now() + timeDelta), } } @@ -77,14 +78,6 @@ describe('SsoAccessTokenProvider', function () { return cacheDir } - before(function () { - clock = installFakeClock() - }) - - after(function () { - clock.uninstall() - }) - beforeEach(async function () { oidcClient = stub(OidcClient) tempDir = await makeTemporaryTokenCacheFolder() @@ -95,7 +88,7 @@ describe('SsoAccessTokenProvider', function () { afterEach(async function () { sinon.restore() - clock.reset() + clock?.uninstall() await tryRemoveFolder(tempDir) }) @@ -267,6 +260,7 @@ describe('SsoAccessTokenProvider', function () { }) it(`emits session duration between logins of the same startUrl`, async function () { + clock = installFakeClock() setupFlow() stubOpen() @@ -311,6 +305,7 @@ describe('SsoAccessTokenProvider', function () { }) it('respects the device authorization expiration time', async function () { + clock = installFakeClock() setupFlow() stubOpen() const exception = new AuthorizationPendingException({ message: '', $metadata: {} }) @@ -352,7 +347,7 @@ describe('SsoAccessTokenProvider', function () { const registration = { clientId: 'myExpiredClientId', clientSecret: 'myExpiredClientSecret', - expiresAt: new clock.Date(clock.Date.now() - 1), // expired date + expiresAt: new globals.clock.Date(globals.clock.Date.now() - 1), // expired date startUrl: key.startUrl, } await cache.registration.save(key, registration) From 6bcf269d8849b2477090d041390b6ad0fb248593 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Tue, 7 Jan 2025 18:55:41 -0500 Subject: [PATCH 3/5] add debounce sanity check test Signed-off-by: nkomonen-amazon --- .../credentials/sso/ssoAccessTokenProvider.test.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts index 2cb98193224..ffabab68691 100644 --- a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts +++ b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts @@ -156,6 +156,18 @@ describe('SsoAccessTokenProvider', function () { assert.strictEqual(cachedToken, undefined) }) + it('concurrent calls resolve successfully', async function () { + // This test verifies debounced getToken() does not break things. + // But this test is not perfect since we are unable to spy _getToken() for some reason. + const validToken = createToken(hourInMs) + await cache.token.save(startUrl, { region, startUrl, token: validToken }) + + const result = await Promise.all([sut.getToken(), sut.getToken(), sut.getToken()]) + for (const r of result) { + assert.deepStrictEqual(r, validToken) + } + }) + describe('Exceptions', function () { it('drops expired tokens if failure was a client-fault', async function () { const exception = new UnauthorizedClientException({ message: '', $metadata: {} }) From 2efa3ece31f43e609a450c0c5d4ad583ee754451 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 8 Jan 2025 11:39:18 -0500 Subject: [PATCH 4/5] fix another broken test - Removed unnecessary fake clock. I guess this existed for historical reasons but is not needed anymore - Refactored one of the methods so that it could be stubbed Signed-off-by: nkomonen-amazon --- packages/core/src/auth/sso/ssoAccessTokenProvider.ts | 10 ++++++++-- .../provider/sharedCredentialsProvider.test.ts | 9 --------- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts index aeb6f50010a..b95abca3c4f 100644 --- a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts +++ b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts @@ -101,9 +101,15 @@ export abstract class SsoAccessTokenProvider { * Sometimes we get many calls at once and this * can trigger redundant disk reads, or token refreshes. * We debounce to avoid this. + * + * NOTE: The property {@link getTokenDebounced()} does not work with being stubbed for tests, so + * this redundant function was created to work around that. */ - public getToken = debounce(this._getToken.bind(this), 100) - public async _getToken(): Promise { + public async getToken(): Promise { + return this.getTokenDebounced() + } + private getTokenDebounced = debounce(this._getToken.bind(this), 100) + private async _getToken(): Promise { const data = await this.cache.token.load(this.tokenCacheKey) SsoAccessTokenProvider.logIfChanged( indent( diff --git a/packages/core/src/test/credentials/provider/sharedCredentialsProvider.test.ts b/packages/core/src/test/credentials/provider/sharedCredentialsProvider.test.ts index 8e266cc11d3..9f12d7c6c64 100644 --- a/packages/core/src/test/credentials/provider/sharedCredentialsProvider.test.ts +++ b/packages/core/src/test/credentials/provider/sharedCredentialsProvider.test.ts @@ -4,13 +4,11 @@ */ import assert from 'assert' -import * as FakeTimers from '@sinonjs/fake-timers' import * as sinon from 'sinon' import { SharedCredentialsProvider } from '../../../auth/providers/sharedCredentialsProvider' import { stripUndefined } from '../../../shared/utilities/collectionUtils' import * as process from '@aws-sdk/credential-provider-process' import { ParsedIniData } from '@smithy/shared-ini-file-loader' -import { installFakeClock } from '../../testUtil' import { SsoClient } from '../../../auth/sso/clients' import { stub } from '../../utilities/stubber' import { SsoAccessTokenProvider } from '../../../auth/sso/ssoAccessTokenProvider' @@ -19,20 +17,13 @@ import { createTestSections } from '../testUtil' const missingPropertiesFragment = 'missing properties' describe('SharedCredentialsProvider', async function () { - let clock: FakeTimers.InstalledClock let sandbox: sinon.SinonSandbox before(function () { sandbox = sinon.createSandbox() - clock = installFakeClock() - }) - - after(function () { - clock.uninstall() }) afterEach(function () { - clock.reset() sandbox.restore() }) From 5c56b23a7972088fbc58805102b73defc8133db1 Mon Sep 17 00:00:00 2001 From: nkomonen-amazon Date: Wed, 8 Jan 2025 17:02:04 -0500 Subject: [PATCH 5/5] refactor to spy on debounced funcitionality We can now spy on the underlying getToken method to verify it is being debounced as expected. Also reduce the debounce interval to 50ms to reduce the delay but still catch a barrage of calls Signed-off-by: nkomonen-amazon --- packages/core/src/auth/sso/ssoAccessTokenProvider.ts | 5 +++-- .../test/credentials/sso/ssoAccessTokenProvider.test.ts | 8 +++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts index b95abca3c4f..d0c8af56c3f 100644 --- a/packages/core/src/auth/sso/ssoAccessTokenProvider.ts +++ b/packages/core/src/auth/sso/ssoAccessTokenProvider.ts @@ -108,8 +108,9 @@ export abstract class SsoAccessTokenProvider { public async getToken(): Promise { return this.getTokenDebounced() } - private getTokenDebounced = debounce(this._getToken.bind(this), 100) - private async _getToken(): Promise { + private getTokenDebounced = debounce(() => this._getToken(), 50) + /** Exposed for testing purposes only */ + public async _getToken(): Promise { const data = await this.cache.token.load(this.tokenCacheKey) SsoAccessTokenProvider.logIfChanged( indent( diff --git a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts index ffabab68691..6536095cd59 100644 --- a/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts +++ b/packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts @@ -156,13 +156,15 @@ describe('SsoAccessTokenProvider', function () { assert.strictEqual(cachedToken, undefined) }) - it('concurrent calls resolve successfully', async function () { - // This test verifies debounced getToken() does not break things. - // But this test is not perfect since we are unable to spy _getToken() for some reason. + it('concurrent calls are debounced', async function () { const validToken = createToken(hourInMs) await cache.token.save(startUrl, { region, startUrl, token: validToken }) + const actualGetToken = sinon.spy(sut, '_getToken') const result = await Promise.all([sut.getToken(), sut.getToken(), sut.getToken()]) + + // Subsequent other calls were debounced so this was only called once + assert.strictEqual(actualGetToken.callCount, 1) for (const r of result) { assert.deepStrictEqual(r, validToken) }