From 126b78896e59f37f54741484a0ab65162e6b6b96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gast=C3=B3n=20Fournier?= Date: Fri, 19 Apr 2024 15:40:15 +0200 Subject: [PATCH] feat: make edge use token's cache (#6893) ## About the changes This PR removes the feature flag `queryMissingTokens` that was fully rolled out. It introduces a new way of checking edgeValidTokens controlled by the flag `checkEdgeValidTokensFromCache` that relies in the cached data but hits the DB if needed. The assumption is that most of the times edge will find tokens in the cache, except for a few cases in which a new token is queried. From all tokens we expect at most one to hit the DB and in this case querying a single token should be better than querying all the tokens. --- src/lib/db/api-token-store.ts | 2 + .../features/scheduler/schedule-services.ts | 2 +- src/lib/middleware/api-token-middleware.ts | 20 +--- src/lib/services/api-token-service.test.ts | 69 ++++------- src/lib/services/api-token-service.ts | 113 +++++++++--------- src/lib/services/edge-service.ts | 76 +++++++++--- src/lib/types/experimental.ts | 1 + 7 files changed, 146 insertions(+), 137 deletions(-) diff --git a/src/lib/db/api-token-store.ts b/src/lib/db/api-token-store.ts index 48f4fe27f82f..57838bc5e042 100644 --- a/src/lib/db/api-token-store.ts +++ b/src/lib/db/api-token-store.ts @@ -194,10 +194,12 @@ export class ApiTokenStore implements IApiTokenStore { } async get(key: string): Promise { + const stopTimer = this.timer('get-by-secret'); const row = await this.makeTokenProjectQuery().where( 'tokens.secret', key, ); + stopTimer(); return toTokens(row)[0]; } diff --git a/src/lib/features/scheduler/schedule-services.ts b/src/lib/features/scheduler/schedule-services.ts index 41636e58f5c3..a096c237a9b0 100644 --- a/src/lib/features/scheduler/schedule-services.ts +++ b/src/lib/features/scheduler/schedule-services.ts @@ -49,7 +49,7 @@ export const scheduleServices = async ( apiTokenService.fetchActiveTokens.bind(apiTokenService), minutesToMilliseconds(1), 'fetchActiveTokens', - 0, + 0, // no jitter, we need tokens at startup ); schedulerService.schedule( diff --git a/src/lib/middleware/api-token-middleware.ts b/src/lib/middleware/api-token-middleware.ts index 1a3e95390a3c..bde8fa031545 100644 --- a/src/lib/middleware/api-token-middleware.ts +++ b/src/lib/middleware/api-token-middleware.ts @@ -2,7 +2,6 @@ import { ApiTokenType } from '../types/models/api-token'; import type { IUnleashConfig } from '../types/option'; import type { IApiRequest, IAuthRequest } from '../routes/unleash-types'; import type { IUnleashServices } from '../server-impl'; -import type { IFlagContext } from '../types'; const isClientApi = ({ path }) => { return path && path.indexOf('/api/client') > -1; @@ -27,20 +26,6 @@ const isProxyApi = ({ path }) => { ); }; -const contextFrom = ( - req: IAuthRequest | IApiRequest, -): IFlagContext | undefined => { - // this is what we'd get from edge: - // req_path: '/api/client/features', - // req_user_agent: 'unleash-edge-16.0.4' - return { - reqPath: req.path, - reqUserAgent: req.get ? req.get('User-Agent') ?? '' : '', - reqAppName: - req.headers?.['unleash-appname'] ?? req.query?.appName ?? '', - }; -}; - export const TOKEN_TYPE_ERROR_MESSAGE = 'invalid token: expected a different token type for this endpoint'; @@ -70,10 +55,7 @@ const apiAccessMiddleware = ( const apiToken = req.header('authorization'); if (!apiToken?.startsWith('user:')) { const apiUser = apiToken - ? await apiTokenService.getUserForToken( - apiToken, - contextFrom(req), - ) + ? await apiTokenService.getUserForToken(apiToken) : undefined; const { CLIENT, FRONTEND } = ApiTokenType; diff --git a/src/lib/services/api-token-service.test.ts b/src/lib/services/api-token-service.test.ts index 95d18314ed35..83f3ab13fcc2 100644 --- a/src/lib/services/api-token-service.test.ts +++ b/src/lib/services/api-token-service.test.ts @@ -197,17 +197,17 @@ test('getUserForToken should get a user with admin token user id and token name' expect(user!.internalAdminTokenUserId).toBe(ADMIN_TOKEN_USER.id); }); -describe('When token is added by another instance', () => { - const setup = (options?: IUnleashOptions) => { - const token: IApiTokenCreate = { - environment: 'default', - projects: ['*'], - secret: '*:*:some-random-string', - type: ApiTokenType.CLIENT, - tokenName: 'new-token-by-another-instance', - expiresAt: undefined, - }; +describe('API token getTokenWithCache', () => { + const token: IApiTokenCreate = { + environment: 'default', + projects: ['*'], + secret: '*:*:some-random-string', + type: ApiTokenType.CLIENT, + tokenName: 'new-token-by-another-instance', + expiresAt: undefined, + }; + const setup = (options?: IUnleashOptions) => { const config: IUnleashConfig = createTestConfig(options); const apiTokenStore = new FakeApiTokenStore(); const environmentStore = new FakeEnvironmentStore(); @@ -220,60 +220,43 @@ describe('When token is added by another instance', () => { return { apiTokenService, apiTokenStore, - token, }; }; - test('should not return the token when query db flag is disabled', async () => { - const { apiTokenService, apiTokenStore, token } = setup(); - // simulate this token being inserted by another instance (apiTokenService does not know about it) - apiTokenStore.insert(token); - - const found = await apiTokenService.getUserForToken(token.secret); - expect(found).toBeUndefined(); - }); - - test('should return the token when query db flag is enabled', async () => { - const { apiTokenService, apiTokenStore, token } = setup({ - experimental: { - flags: { - queryMissingTokens: true, - }, - }, - }); + test('should return the token and perform only one db query', async () => { + const { apiTokenService, apiTokenStore } = setup(); + const apiTokenStoreGet = jest.spyOn(apiTokenStore, 'get'); - // simulate this token being inserted by another instance (apiTokenService does not know about it) + // valid token not present in cache (could be inserted by another instance) apiTokenStore.insert(token); - const found = await apiTokenService.getUserForToken(token.secret); - expect(found).toBeDefined(); - expect(found?.username).toBe(token.tokenName); + for (let i = 0; i < 5; i++) { + const found = await apiTokenService.getTokenWithCache(token.secret); + expect(found).toBeDefined(); + expect(found?.tokenName).toBe(token.tokenName); + expect(found?.createdAt).toBeDefined(); + } + expect(apiTokenStoreGet).toHaveBeenCalledTimes(1); }); test('should query the db only once for invalid tokens', async () => { jest.useFakeTimers(); - const { apiTokenService, apiTokenStore } = setup({ - experimental: { - flags: { - queryMissingTokens: true, - }, - }, - }); + const { apiTokenService, apiTokenStore } = setup(); const apiTokenStoreGet = jest.spyOn(apiTokenStore, 'get'); const invalidToken = 'invalid-token'; - for (let i = 0; i < 10; i++) { + for (let i = 0; i < 5; i++) { expect( - await apiTokenService.getUserForToken(invalidToken), + await apiTokenService.getTokenWithCache(invalidToken), ).toBeUndefined(); } expect(apiTokenStoreGet).toHaveBeenCalledTimes(1); // after more than 5 minutes we should be able to query again jest.advanceTimersByTime(minutesToMilliseconds(6)); - for (let i = 0; i < 10; i++) { + for (let i = 0; i < 5; i++) { expect( - await apiTokenService.getUserForToken(invalidToken), + await apiTokenService.getTokenWithCache(invalidToken), ).toBeUndefined(); } expect(apiTokenStoreGet).toHaveBeenCalledTimes(2); diff --git a/src/lib/services/api-token-service.ts b/src/lib/services/api-token-service.ts index 211371698448..c128074e572e 100644 --- a/src/lib/services/api-token-service.ts +++ b/src/lib/services/api-token-service.ts @@ -25,13 +25,14 @@ import { ApiTokenDeletedEvent, ApiTokenUpdatedEvent, type IAuditUser, - type IFlagContext, type IFlagResolver, SYSTEM_USER_AUDIT, } from '../types'; import { omitKeys } from '../util'; import type EventService from '../features/events/event-service'; import { addMinutes, isPast } from 'date-fns'; +import metricsHelper from '../util/metrics-helper'; +import { FUNCTION_TIME } from '../metric-events'; const resolveTokenPermissions = (tokenType: string) => { if (tokenType === ApiTokenType.ADMIN) { @@ -60,14 +61,14 @@ export class ApiTokenService { private queryAfter = new Map(); - private initialized = false; - private eventService: EventService; private lastSeenSecrets: Set = new Set(); private flagResolver: IFlagResolver; + private timer: Function; + constructor( { apiTokenStore, @@ -75,7 +76,7 @@ export class ApiTokenService { }: Pick, config: Pick< IUnleashConfig, - 'getLogger' | 'authentication' | 'flagResolver' + 'getLogger' | 'authentication' | 'flagResolver' | 'eventBus' >, eventService: EventService, ) { @@ -94,18 +95,21 @@ export class ApiTokenService { this.initApiTokens(config.authentication.initApiTokens), ); } + this.timer = (functionName: string) => + metricsHelper.wrapTimer(config.eventBus, FUNCTION_TIME, { + className: 'ApiTokenService', + functionName, + }); } /** - * Executed by a scheduler to refresh all active tokens + * Called by a scheduler without jitter to refresh all active tokens */ async fetchActiveTokens(): Promise { try { this.activeTokens = await this.store.getAllActive(); - this.initialized = true; - } finally { - // biome-ignore lint/correctness/noUnsafeFinally: We ignored this for eslint. Leaving this here for now, server-impl test fails without it - return; + } catch (e) { + this.logger.warn('Failed to fetch active tokens', e); } } @@ -113,50 +117,7 @@ export class ApiTokenService { return this.store.get(secret); } - async updateLastSeen(): Promise { - if (this.lastSeenSecrets.size > 0) { - const toStore = [...this.lastSeenSecrets]; - this.lastSeenSecrets = new Set(); - await this.store.markSeenAt(toStore); - } - } - - public async getAllTokens(): Promise { - return this.store.getAll(); - } - - public async getAllActiveTokens(): Promise { - if (this.flagResolver.isEnabled('useMemoizedActiveTokens')) { - if (!this.initialized) { - // unlikely this will happen but nice to have a fail safe - this.logger.info('Fetching active tokens before initialized'); - await this.fetchActiveTokens(); - } - return this.activeTokens; - } else { - return this.store.getAllActive(); - } - } - - private async initApiTokens(tokens: ILegacyApiTokenCreate[]) { - const tokenCount = await this.store.count(); - if (tokenCount > 0) { - return; - } - try { - const createAll = tokens - .map(mapLegacyTokenWithSecret) - .map((t) => this.insertNewApiToken(t, SYSTEM_USER_AUDIT)); - await Promise.all(createAll); - } catch (e) { - this.logger.error('Unable to create initial Admin API tokens'); - } - } - - public async getUserForToken( - secret: string, - flagContext?: IFlagContext, // temporarily added, expected from the middleware - ): Promise { + async getTokenWithCache(secret: string): Promise { if (!secret) { return undefined; } @@ -178,11 +139,7 @@ export class ApiTokenService { } const nextAllowedQuery = this.queryAfter.get(secret) ?? 0; - if ( - !token && - isPast(nextAllowedQuery) && - this.flagResolver.isEnabled('queryMissingTokens', flagContext) - ) { + if (!token && isPast(nextAllowedQuery)) { if (this.queryAfter.size > 1000) { // establish a max limit for queryAfter size to prevent memory leak this.queryAfter.clear(); @@ -190,12 +147,52 @@ export class ApiTokenService { // prevent querying the same invalid secret multiple times. Expire after 5 minutes this.queryAfter.set(secret, addMinutes(new Date(), 5)); + const stopCacheTimer = this.timer('getTokenWithCache.query'); token = await this.store.get(secret); if (token) { this.activeTokens.push(token); } + stopCacheTimer(); + } + + return token; + } + + async updateLastSeen(): Promise { + if (this.lastSeenSecrets.size > 0) { + const toStore = [...this.lastSeenSecrets]; + this.lastSeenSecrets = new Set(); + await this.store.markSeenAt(toStore); + } + } + + public async getAllTokens(): Promise { + return this.store.getAll(); + } + + public async getAllActiveTokens(): Promise { + return this.store.getAllActive(); + } + + private async initApiTokens(tokens: ILegacyApiTokenCreate[]) { + const tokenCount = await this.store.count(); + if (tokenCount > 0) { + return; } + try { + const createAll = tokens + .map(mapLegacyTokenWithSecret) + .map((t) => this.insertNewApiToken(t, SYSTEM_USER_AUDIT)); + await Promise.all(createAll); + } catch (e) { + this.logger.error('Unable to create initial Admin API tokens'); + } + } + public async getUserForToken( + secret: string, + ): Promise { + const token = await this.getTokenWithCache(secret); if (token) { this.lastSeenSecrets.add(token.secret); const apiUser: IApiUser = new ApiUser({ diff --git a/src/lib/services/edge-service.ts b/src/lib/services/edge-service.ts index b52dba4bf30e..4cc507dd0f8f 100644 --- a/src/lib/services/edge-service.ts +++ b/src/lib/services/edge-service.ts @@ -1,39 +1,83 @@ -import type { IUnleashConfig } from '../types'; +import type { IFlagResolver, IUnleashConfig } from '../types'; import type { Logger } from '../logger'; import type { EdgeTokenSchema } from '../openapi/spec/edge-token-schema'; import { constantTimeCompare } from '../util/constantTimeCompare'; import type { ValidatedEdgeTokensSchema } from '../openapi/spec/validated-edge-tokens-schema'; import type { ApiTokenService } from './api-token-service'; +import metricsHelper from '../util/metrics-helper'; +import { FUNCTION_TIME } from '../metric-events'; export default class EdgeService { private logger: Logger; private apiTokenService: ApiTokenService; + private flagResolver: IFlagResolver; + + private timer: Function; + constructor( { apiTokenService }: { apiTokenService: ApiTokenService }, - { getLogger }: Pick, + { + getLogger, + flagResolver, + eventBus, + }: Pick, ) { this.logger = getLogger('lib/services/edge-service.ts'); this.apiTokenService = apiTokenService; + this.flagResolver = flagResolver; + this.timer = (functionName: string) => + metricsHelper.wrapTimer(eventBus, FUNCTION_TIME, { + className: 'EdgeService', + functionName, + }); } async getValidTokens(tokens: string[]): Promise { - const activeTokens = await this.apiTokenService.getAllActiveTokens(); - const edgeTokens = tokens.reduce((result: EdgeTokenSchema[], token) => { - const dbToken = activeTokens.find((activeToken) => - constantTimeCompare(activeToken.secret, token), - ); - if (dbToken) { - result.push({ - token: token, - type: dbToken.type, - projects: dbToken.projects, - }); + if (this.flagResolver.isEnabled('checkEdgeValidTokensFromCache')) { + const stopTimer = this.timer('validateTokensWithCache'); + // new behavior: use cached tokens when possible + // use the db to fetch the missing ones + // cache stores both missing and active so we don't hammer the db + const validatedTokens: EdgeTokenSchema[] = []; + for (const token of tokens) { + const found = + await this.apiTokenService.getTokenWithCache(token); + if (found) { + validatedTokens.push({ + token: token, + type: found.type, + projects: found.projects, + }); + } } - return result; - }, []); - return { tokens: edgeTokens }; + stopTimer(); + return { tokens: validatedTokens }; + } else { + // old behavior: go to the db to fetch all tokens and then filter in memory + const stopTimer = this.timer('validateTokensWithoutCache'); + const activeTokens = + await this.apiTokenService.getAllActiveTokens(); + const edgeTokens = tokens.reduce( + (result: EdgeTokenSchema[], token) => { + const dbToken = activeTokens.find((activeToken) => + constantTimeCompare(activeToken.secret, token), + ); + if (dbToken) { + result.push({ + token: token, + type: dbToken.type, + projects: dbToken.projects, + }); + } + return result; + }, + [], + ); + stopTimer(); + return { tokens: edgeTokens }; + } } } diff --git a/src/lib/types/experimental.ts b/src/lib/types/experimental.ts index c5be30ed3045..1263aefa90fe 100644 --- a/src/lib/types/experimental.ts +++ b/src/lib/types/experimental.ts @@ -43,6 +43,7 @@ export type IFlagKey = | 'displayTrafficDataUsage' | 'useMemoizedActiveTokens' | 'queryMissingTokens' + | 'checkEdgeValidTokensFromCache' | 'userAccessUIEnabled' | 'disableUpdateMaxRevisionId' | 'disablePublishUnannouncedEvents'