chore: memoizee active tokens (#6135)

## About the changes
getAllActive from api-tokens store is the second most frequent query

![image](https://github.com/Unleash/unleash/assets/455064/63c5ae76-bb62-41b2-95b4-82aca59a7c16)

To prevent starving our db connections, we can cache this data that
rarely changes and clear the cache when we see changes. Because we will
only clear changes in the node receiving the change we're only caching
the data for 1 minute.

This should give us some room to test if this solution will work

---------

Co-authored-by: Nuno Góis <github@nunogois.com>

src/lib/__snapshots__/create-config.test.ts.snap

@@ -135,6 +135,7 @@ exports[`should create default config 1`] = `
       "showInactiveUsers": false,
       "strictSchemaValidation": false,
       "stripClientHeadersOn304": false,
+      "useMemoizedActiveTokens": false,
     },
     "externalResolver": {
       "getVariant": [Function],

src/lib/services/api-token-service.test.ts

@@ -29,6 +29,11 @@ test('Should init api token', async () => {
         authentication: {
             initApiTokens: [token],
         },
+        experimental: {
+            flags: {
+                useMemoizedActiveTokens: true,
+            },
+        },
     });
     const apiTokenStore = new FakeApiTokenStore();
     const environmentStore = new FakeEnvironmentStore();

src/lib/services/api-token-service.ts

@@ -24,6 +24,7 @@ import {
     ApiTokenCreatedEvent,
     ApiTokenDeletedEvent,
     ApiTokenUpdatedEvent,
+    IFlagResolver,
     IUser,
     SYSTEM_USER,
     SYSTEM_USER_ID,
@@ -60,23 +61,34 @@ export class ApiTokenService {
 
     private activeTokens: IApiToken[] = [];
 
+    private initialized = false;
+
     private eventService: EventService;
 
     private lastSeenSecrets: Set<string> = new Set<string>();
 
+    private flagResolver: IFlagResolver;
+
     constructor(
         {
             apiTokenStore,
             environmentStore,
         }: Pick<IUnleashStores, 'apiTokenStore' | 'environmentStore'>,
-        config: Pick<IUnleashConfig, 'getLogger' | 'authentication'>,
+        config: Pick<
+            IUnleashConfig,
+            'getLogger' | 'authentication' | 'flagResolver'
+        >,
         eventService: EventService,
     ) {
         this.store = apiTokenStore;
         this.eventService = eventService;
         this.environmentStore = environmentStore;
+        this.flagResolver = config.flagResolver;
         this.logger = config.getLogger('/services/api-token-service.ts');
-        this.fetchActiveTokens();
+        if (!this.flagResolver.isEnabled('useMemoizedActiveTokens')) {
+            // This is probably not needed because the scheduler will run it
+            this.fetchActiveTokens();
+        }
         this.updateLastSeen();
         if (config.authentication.initApiTokens.length > 0) {
             process.nextTick(async () =>
@@ -85,9 +97,13 @@ export class ApiTokenService {
         }
     }
 
+    /**
+     * Executed by a scheduler to refresh all active tokens
+     */
     async fetchActiveTokens(): Promise<void> {
         try {
-            this.activeTokens = await this.getAllActiveTokens();
+            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;
@@ -111,7 +127,16 @@ export class ApiTokenService {
     }
 
     public async getAllActiveTokens(): Promise<IApiToken[]> {
-        return this.store.getAllActive();
+        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[]) {

src/lib/services/edge-service.ts

@@ -1,25 +1,25 @@
-import { IUnleashStores, IUnleashConfig } from '../types';
+import { IUnleashConfig } from '../types';
 import { Logger } from '../logger';
-import { IApiTokenStore } from '../types/stores/api-token-store';
 import { EdgeTokenSchema } from '../openapi/spec/edge-token-schema';
 import { constantTimeCompare } from '../util/constantTimeCompare';
 import { ValidatedEdgeTokensSchema } from '../openapi/spec/validated-edge-tokens-schema';
+import { ApiTokenService } from './api-token-service';
 
 export default class EdgeService {
     private logger: Logger;
 
-    private apiTokenStore: IApiTokenStore;
+    private apiTokenService: ApiTokenService;
 
     constructor(
-        { apiTokenStore }: Pick<IUnleashStores, 'apiTokenStore'>,
+        { apiTokenService }: { apiTokenService: ApiTokenService },
         { getLogger }: Pick<IUnleashConfig, 'getLogger'>,
     ) {
         this.logger = getLogger('lib/services/edge-service.ts');
-        this.apiTokenStore = apiTokenStore;
+        this.apiTokenService = apiTokenService;
     }
 
     async getValidTokens(tokens: string[]): Promise<ValidatedEdgeTokensSchema> {
-        const activeTokens = await this.apiTokenStore.getAllActive();
+        const activeTokens = await this.apiTokenService.getAllActiveTokens();
         const edgeTokens = tokens.reduce((result: EdgeTokenSchema[], token) => {
             const dbToken = activeTokens.find((activeToken) =>
                 constantTimeCompare(activeToken.secret, token),

src/lib/services/index.ts

@@ -293,7 +293,7 @@ export const createServices = (
         configurationRevisionService,
     });
 
-    const edgeService = new EdgeService(stores, config);
+    const edgeService = new EdgeService({ apiTokenService }, config);
 
     const patService = new PatService(stores, config, eventService);
 

src/lib/types/experimental.ts

@@ -46,7 +46,8 @@ export type IFlagKey =
     | 'executiveDashboard'
     | 'feedbackComments'
     | 'createdByUserIdDataMigration'
-    | 'showInactiveUsers';
+    | 'showInactiveUsers'
+    | 'useMemoizedActiveTokens';
 
 export type IFlags = Partial<{ [key in IFlagKey]: boolean | Variant }>;
 
@@ -227,6 +228,10 @@ const flags: IFlags = {
         process.env.UNLEASH_EXPERIMENTAL_SHOW_INACTIVE_USERS,
         false,
     ),
+    useMemoizedActiveTokens: parseEnvVarBoolean(
+        process.env.UNLEASH_EXPERIMENTAL_MEMOIZED_ACTIVE_TOKENS,
+        false,
+    ),
 };
 
 export const defaultExperimentalOptions: IExperimentalOptions = {

src/test/e2e/api/admin/api-token.auth.e2e.test.ts

@@ -207,15 +207,19 @@ test('Token-admin should be allowed to create token', async () => {
 test('An admin token should be allowed to create a token', async () => {
     expect.assertions(2);
 
-    const adminToken = await db.stores.apiTokenStore.insert({
-        type: ApiTokenType.ADMIN,
-        secret: '12345',
-        environment: '',
-        projects: [],
-        tokenName: 'default-admin',
-    });
+    const { request, destroy, services } = await setupAppWithAuth(stores);
 
-    const { request, destroy } = await setupAppWithAuth(stores);
+    const { secret } =
+        await services.apiTokenService.createApiTokenWithProjects(
+            {
+                tokenName: 'default-admin',
+                type: ApiTokenType.ADMIN,
+                projects: ['*'],
+                environment: '*',
+            },
+            'test',
+            1,
+        );
 
     await request
         .post('/api/admin/api-tokens')
@@ -223,12 +227,12 @@ test('An admin token should be allowed to create a token', async () => {
             username: 'default-admin',
             type: 'admin',
         })
-        .set('Authorization', adminToken.secret)
+        .set('Authorization', secret)
         .set('Content-Type', 'application/json')
         .expect(201);
 
     const event = await getLastEvent();
-    expect(event.createdBy).toBe(adminToken.tokenName);
+    expect(event.createdBy).toBe('default-admin');
     expect(event.createdByUserId).toBe(ADMIN_TOKEN_USER.id);
     await destroy();
 });

src/test/e2e/services/api-token-service.e2e.test.ts

@@ -18,6 +18,11 @@ let projectService: ProjectService;
 beforeAll(async () => {
     const config = createTestConfig({
         server: { baseUriPath: '/test' },
+        experimental: {
+            flags: {
+                useMemoizedActiveTokens: true,
+            },
+        },
     });
     db = await dbInit('api_token_service_serial', getLogger);
     stores = db.stores;
@@ -179,7 +184,7 @@ test('should return user with multiple projects', async () => {
     const now = Date.now();
     const tomorrow = addDays(now, 1);
 
-    await apiTokenService.createApiToken({
+    const { secret: secret1 } = await apiTokenService.createApiToken({
         tokenName: 'default-valid',
         type: ApiTokenType.CLIENT,
         expiresAt: tomorrow,
@@ -187,7 +192,7 @@ test('should return user with multiple projects', async () => {
         environment: DEFAULT_ENV,
     });
 
-    await apiTokenService.createApiToken({
+    const { secret: secret2 } = await apiTokenService.createApiToken({
         tokenName: 'default-also-valid',
         type: ApiTokenType.CLIENT,
         expiresAt: tomorrow,
@@ -195,13 +200,8 @@ test('should return user with multiple projects', async () => {
         environment: DEFAULT_ENV,
     });
 
-    const tokens = await apiTokenService.getAllActiveTokens();
-    const multiProjectUser = await apiTokenService.getUserForToken(
-        tokens[0].secret,
-    );
-    const singleProjectUser = await apiTokenService.getUserForToken(
-        tokens[1].secret,
-    );
+    const multiProjectUser = apiTokenService.getUserForToken(secret1);
+    const singleProjectUser = apiTokenService.getUserForToken(secret2);
 
     expect(multiProjectUser!.projects).toStrictEqual([
         'test-project',