chore: improve the performance of our instance stats (#8766)

## About the changes
Our stats are used for many places and many times to publish prometheus
metrics and some other things.

Some of these queries are heavy, traversing all tables to calculate
aggregates.

This adds a feature flag to be able to memoize 1 minute (by default) how
long to keep the calculated values in memory.

We can use the key of the function to individually control which ones
are memoized or not and for how long using a numeric variant.

Initially, this will be disabled and we'll test in our instances first

src/lib/features/instance-stats/instance-stats-service.test.ts

@@ -6,11 +6,18 @@ import { createFakeGetActiveUsers } from './getActiveUsers';
 import { createFakeGetProductionChanges } from './getProductionChanges';
 import { registerPrometheusMetrics } from '../../metrics';
 import { register } from 'prom-client';
-import type { IClientInstanceStore } from '../../types';
+import type {
+    IClientInstanceStore,
+    IFlagResolver,
+    IUnleashStores,
+} from '../../types';
 import { createFakeGetLicensedUsers } from './getLicensedUsers';
 let instanceStatsService: InstanceStatsService;
 let versionService: VersionService;
 let clientInstanceStore: IClientInstanceStore;
+let stores: IUnleashStores;
+let flagResolver: IFlagResolver;
+
 let updateMetrics: () => Promise<void>;
 beforeEach(() => {
     jest.clearAllMocks();
@@ -18,7 +25,8 @@ beforeEach(() => {
     register.clear();
 
     const config = createTestConfig();
-    const stores = createStores();
+    flagResolver = config.flagResolver;
+    stores = createStores();
     versionService = new VersionService(
         stores,
         config,
@@ -74,3 +82,96 @@ test('get snapshot should not call getStats', async () => {
 test('before the snapshot is refreshed we can still get the appCount', async () => {
     expect(instanceStatsService.getAppCountSnapshot('7d')).toBeUndefined();
 });
+
+describe.each([true, false])(
+    'When feature enabled is %s',
+    (featureEnabled: boolean) => {
+        beforeEach(() => {
+            jest.spyOn(flagResolver, 'getVariant').mockReturnValue({
+                name: 'memorizeStats',
+                enabled: featureEnabled,
+                feature_enabled: featureEnabled,
+            });
+        });
+
+        test(`should${featureEnabled ? ' ' : ' not '}memoize query results`, async () => {
+            const segmentStore = stores.segmentStore;
+            jest.spyOn(segmentStore, 'count').mockReturnValue(
+                Promise.resolve(5),
+            );
+            expect(segmentStore.count).toHaveBeenCalledTimes(0);
+            expect(await instanceStatsService.segmentCount()).toBe(5);
+            expect(segmentStore.count).toHaveBeenCalledTimes(1);
+            expect(await instanceStatsService.segmentCount()).toBe(5);
+            expect(segmentStore.count).toHaveBeenCalledTimes(
+                featureEnabled ? 1 : 2,
+            );
+        });
+
+        test(`should${featureEnabled ? ' ' : ' not '}memoize async query results`, async () => {
+            const trafficDataUsageStore = stores.trafficDataUsageStore;
+            jest.spyOn(
+                trafficDataUsageStore,
+                'getTrafficDataUsageForPeriod',
+            ).mockReturnValue(
+                Promise.resolve([
+                    {
+                        day: new Date(),
+                        trafficGroup: 'default',
+                        statusCodeSeries: 200,
+                        count: 5,
+                    },
+                    {
+                        day: new Date(),
+                        trafficGroup: 'default',
+                        statusCodeSeries: 400,
+                        count: 2,
+                    },
+                ]),
+            );
+            expect(
+                trafficDataUsageStore.getTrafficDataUsageForPeriod,
+            ).toHaveBeenCalledTimes(0);
+            expect(await instanceStatsService.getCurrentTrafficData()).toBe(7);
+            expect(
+                trafficDataUsageStore.getTrafficDataUsageForPeriod,
+            ).toHaveBeenCalledTimes(1);
+            expect(await instanceStatsService.getCurrentTrafficData()).toBe(7);
+            expect(
+                trafficDataUsageStore.getTrafficDataUsageForPeriod,
+            ).toHaveBeenCalledTimes(featureEnabled ? 1 : 2);
+        });
+
+        test(`getStats should${featureEnabled ? ' ' : ' not '}be memorized`, async () => {
+            const featureStrategiesReadModel =
+                stores.featureStrategiesReadModel;
+            jest.spyOn(
+                featureStrategiesReadModel,
+                'getMaxFeatureEnvironmentStrategies',
+            ).mockReturnValue(
+                Promise.resolve({
+                    feature: 'x',
+                    environment: 'default',
+                    count: 3,
+                }),
+            );
+            expect(
+                featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies,
+            ).toHaveBeenCalledTimes(0);
+            expect(
+                (await instanceStatsService.getStats())
+                    .maxEnvironmentStrategies,
+            ).toBe(3);
+            expect(
+                featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies,
+            ).toHaveBeenCalledTimes(1);
+            expect(
+                (await instanceStatsService.getStats())
+                    .maxEnvironmentStrategies,
+            ).toBe(3);
+            expect(
+                featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies,
+            ).toHaveBeenCalledTimes(featureEnabled ? 1 : 2);
+        });
+    },
+);

src/lib/features/instance-stats/instance-stats-service.ts

@@ -30,7 +30,8 @@ import { CUSTOM_ROOT_ROLE_TYPE } from '../../util';
 import type { GetActiveUsers } from './getActiveUsers';
 import type { ProjectModeCount } from '../project/project-store';
 import type { GetProductionChanges } from './getProductionChanges';
-import { format } from 'date-fns';
+import { format, minutesToMilliseconds } from 'date-fns';
+import memoizee from 'memoizee';
 import type { GetLicensedUsers } from './getLicensedUsers';
 
 export type TimeRange = 'allTime' | '30d' | '7d';
@@ -58,6 +59,8 @@ export interface InstanceStats {
     strategies: number;
     SAMLenabled: boolean;
     OIDCenabled: boolean;
+    passwordAuthEnabled: boolean;
+    SCIMenabled: boolean;
     clientApps: { range: TimeRange; count: number }[];
     activeUsers: Awaited<ReturnType<GetActiveUsers>>;
     licensedUsers: Awaited<ReturnType<GetLicensedUsers>>;
@@ -193,39 +196,75 @@ export class InstanceStatsService {
         this.trafficDataUsageStore = trafficDataUsageStore;
     }
 
+    memory = new Map<string, () => Promise<any>>();
+    memorize<T>(key: string, fn: () => Promise<T>): Promise<T> {
+        const variant = this.flagResolver.getVariant('memorizeStats', {
+            memoryKey: key,
+        });
+        if (variant.feature_enabled) {
+            const minutes =
+                variant.payload?.type === 'number'
+                    ? Number(variant.payload.value)
+                    : 1;
+
+            let memoizedFunction = this.memory.get(key);
+            if (!memoizedFunction) {
+                memoizedFunction = memoizee(() => fn(), {
+                    promise: true,
+                    maxAge: minutesToMilliseconds(minutes),
+                });
+                this.memory.set(key, memoizedFunction);
+            }
+            return memoizedFunction();
+        } else {
+            return fn();
+        }
+    }
+
     getProjectModeCount(): Promise<ProjectModeCount[]> {
-        return this.projectStore.getProjectModeCounts();
+        return this.memorize('getProjectModeCount', () =>
+            this.projectStore.getProjectModeCounts(),
+        );
     }
 
     getToggleCount(): Promise<number> {
-        return this.featureToggleStore.count({
+        return this.memorize('getToggleCount', () =>
+            this.featureToggleStore.count({
                 archived: false,
-        });
+            }),
+        );
     }
 
     getArchivedToggleCount(): Promise<number> {
-        return this.featureToggleStore.count({
+        return this.memorize('hasOIDC', () =>
+            this.featureToggleStore.count({
                 archived: true,
-        });
+            }),
+        );
     }
 
     async hasOIDC(): Promise<boolean> {
+        return this.memorize('hasOIDC', async () => {
             const settings = await this.settingStore.get<{ enabled: boolean }>(
                 'unleash.enterprise.auth.oidc',
             );
 
             return settings?.enabled || false;
+        });
     }
 
     async hasSAML(): Promise<boolean> {
+        return this.memorize('hasSAML', async () => {
             const settings = await this.settingStore.get<{ enabled: boolean }>(
                 'unleash.enterprise.auth.saml',
             );
 
             return settings?.enabled || false;
+        });
     }
 
     async hasPasswordAuth(): Promise<boolean> {
+        return this.memorize('hasPasswordAuth', async () => {
             const settings = await this.settingStore.get<{ disabled: boolean }>(
                 'unleash.auth.simple',
             );
@@ -234,14 +273,17 @@ export class InstanceStatsService {
                 typeof settings?.disabled === 'undefined' ||
                 settings.disabled === false
             );
+        });
     }
 
     async hasSCIM(): Promise<boolean> {
+        return this.memorize('hasSCIM', async () => {
             const settings = await this.settingStore.get<{ enabled: boolean }>(
                 'scim',
             );
 
             return settings?.enabled || false;
+        });
     }
 
     async getStats(): Promise<InstanceStats> {
@@ -281,8 +323,8 @@ export class InstanceStatsService {
             this.getRegisteredUsers(),
             this.countServiceAccounts(),
             this.countApiTokensByType(),
-            this.getActiveUsers(),
+            this.memorize('getActiveUsers', this.getActiveUsers.bind(this)),
-            this.getLicencedUsers(),
+            this.memorize('getLicencedUsers', this.getLicencedUsers.bind(this)),
             this.getProjectModeCount(),
             this.contextFieldCount(),
             this.groupCount(),
@@ -297,17 +339,39 @@ export class InstanceStatsService {
             this.hasPasswordAuth(),
             this.hasSCIM(),
             this.appCount ? this.appCount : this.getLabeledAppCounts(),
+            this.memorize('deprecatedFilteredCountFeaturesExported', () =>
                 this.eventStore.deprecatedFilteredCount({
                     type: FEATURES_EXPORTED,
                 }),
+            ),
+            this.memorize('deprecatedFilteredCountFeaturesImported', () =>
                 this.eventStore.deprecatedFilteredCount({
                     type: FEATURES_IMPORTED,
                 }),
-            this.getProductionChanges(),
+            ),
+            this.memorize(
+                'getProductionChanges',
+                this.getProductionChanges.bind(this),
+            ),
             this.countPreviousDayHourlyMetricsBuckets(),
-            this.featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies(),
+            this.memorize(
-            this.featureStrategiesReadModel.getMaxConstraintValues(),
+                'maxFeatureEnvironmentStrategies',
-            this.featureStrategiesReadModel.getMaxConstraintsPerStrategy(),
+                this.featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies.bind(
+                    this.featureStrategiesReadModel,
+                ),
+            ),
+            this.memorize(
+                'maxConstraintValues',
+                this.featureStrategiesReadModel.getMaxConstraintValues.bind(
+                    this.featureStrategiesReadModel,
+                ),
+            ),
+            this.memorize(
+                'maxConstraintsPerStrategy',
+                this.featureStrategiesReadModel.getMaxConstraintsPerStrategy.bind(
+                    this.featureStrategiesReadModel,
+                ),
+            ),
         ]);
 
         return {
@@ -333,6 +397,8 @@ export class InstanceStatsService {
             strategies,
             SAMLenabled,
             OIDCenabled,
+            passwordAuthEnabled,
+            SCIMenabled,
             clientApps: Object.entries(clientApps).map(([range, count]) => ({
                 range: range as TimeRange,
                 count,
@@ -348,59 +414,78 @@ export class InstanceStatsService {
     }
 
     groupCount(): Promise<number> {
-        return this.groupStore.count();
+        return this.memorize('groupCount', () => this.groupStore.count());
     }
 
     roleCount(): Promise<number> {
-        return this.roleStore.count();
+        return this.memorize('roleCount', () => this.roleStore.count());
     }
 
     customRolesCount(): Promise<number> {
-        return this.roleStore.filteredCount({ type: CUSTOM_ROOT_ROLE_TYPE });
+        return this.memorize('customRolesCount', () =>
+            this.roleStore.filteredCount({ type: CUSTOM_ROOT_ROLE_TYPE }),
+        );
     }
 
     customRolesCountInUse(): Promise<number> {
-        return this.roleStore.filteredCountInUse({
+        return this.memorize('customRolesCountInUse', () =>
+            this.roleStore.filteredCountInUse({
                 type: CUSTOM_ROOT_ROLE_TYPE,
-        });
+            }),
+        );
     }
 
     segmentCount(): Promise<number> {
-        return this.segmentStore.count();
+        return this.memorize('segmentCount', () => this.segmentStore.count());
     }
 
     contextFieldCount(): Promise<number> {
-        return this.contextFieldStore.count();
+        return this.memorize('contextFieldCount', () =>
+            this.contextFieldStore.count(),
+        );
     }
 
     strategiesCount(): Promise<number> {
-        return this.strategyStore.count();
+        return this.memorize('strategiesCount', () =>
+            this.strategyStore.count(),
+        );
     }
 
     environmentCount(): Promise<number> {
-        return this.environmentStore.count();
+        return this.memorize('environmentCount', () =>
+            this.environmentStore.count(),
+        );
     }
 
     countPreviousDayHourlyMetricsBuckets(): Promise<{
         enabledCount: number;
         variantCount: number;
     }> {
-        return this.clientMetricsStore.countPreviousDayHourlyMetricsBuckets();
+        return this.memorize('countPreviousDayHourlyMetricsBuckets', () =>
+            this.clientMetricsStore.countPreviousDayHourlyMetricsBuckets(),
+        );
     }
 
     countApiTokensByType(): Promise<Map<string, number>> {
-        return this.apiTokenStore.countByType();
+        return this.memorize('countApiTokensByType', () =>
+            this.apiTokenStore.countByType(),
+        );
     }
 
     getRegisteredUsers(): Promise<number> {
-        return this.userStore.count();
+        return this.memorize('getRegisteredUsers', () =>
+            this.userStore.count(),
+        );
     }
 
     countServiceAccounts(): Promise<number> {
-        return this.userStore.countServiceAccounts();
+        return this.memorize('countServiceAccounts', () =>
+            this.userStore.countServiceAccounts(),
+        );
     }
 
     async getCurrentTrafficData(): Promise<number> {
+        return this.memorize('getCurrentTrafficData', async () => {
             const traffic =
                 await this.trafficDataUsageStore.getTrafficDataUsageForPeriod(
                     format(new Date(), 'yyyy-MM'),
@@ -408,11 +493,13 @@ export class InstanceStatsService {
 
             const counts = traffic.map((item) => item.count);
             return counts.reduce((total, current) => total + current, 0);
+        });
     }
 
     async getLabeledAppCounts(): Promise<
         Partial<{ [key in TimeRange]: number }>
     > {
+        return this.memorize('getLabeledAppCounts', async () => {
             const [t7d, t30d, allTime] = await Promise.all([
                 this.clientInstanceStore.getDistinctApplicationsCount(7),
                 this.clientInstanceStore.getDistinctApplicationsCount(30),
@@ -424,6 +511,7 @@ export class InstanceStatsService {
                 allTime,
             };
             return this.appCount;
+        });
     }
 
     getAppCountSnapshot(range: TimeRange): number | undefined {

src/lib/routes/admin-api/instance-admin.ts

@@ -85,6 +85,8 @@ class InstanceAdminController extends Controller {
         return {
             OIDCenabled: true,
             SAMLenabled: false,
+            passwordAuthEnabled: true,
+            SCIMenabled: false,
             clientApps: [
                 { range: 'allTime', count: 15 },
                 { range: '30d', count: 9 },

src/lib/types/experimental.ts

@@ -61,7 +61,8 @@ export type IFlagKey =
     | 'simplifyProjectOverview'
     | 'flagOverviewRedesign'
     | 'showUserDeviceCount'
-    | 'deleteStaleUserSessions';
+    | 'deleteStaleUserSessions'
+    | 'memorizeStats';
 
 export type IFlags = Partial<{ [key in IFlagKey]: boolean | Variant }>;