chore: handle aggregation queries sequentially follow up

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

@@ -4,11 +4,18 @@ import createStores from '../../../test/fixtures/store';
 import VersionService from '../../services/version-service';
 import { createFakeGetActiveUsers } from './getActiveUsers';
 import { createFakeGetProductionChanges } from './getProductionChanges';
-
+import { registerPrometheusMetrics } from '../../metrics';
+import { register } from 'prom-client';
+import type { IClientInstanceStore } from '../../types';
 let instanceStatsService: InstanceStatsService;
 let versionService: VersionService;
-
+let clientInstanceStore: IClientInstanceStore;
+let updateMetrics: () => Promise<void>;
 beforeEach(() => {
+    jest.clearAllMocks();
+
+    register.clear();
+
     const config = createTestConfig();
     const stores = createStores();
     versionService = new VersionService(
@@ -17,6 +24,7 @@ beforeEach(() => {
         createFakeGetActiveUsers(),
         createFakeGetProductionChanges(),
     );
+    clientInstanceStore = stores.clientInstanceStore;
     instanceStatsService = new InstanceStatsService(
         stores,
         config,
@@ -25,23 +33,28 @@ beforeEach(() => {
         createFakeGetProductionChanges(),
     );
 
-    jest.spyOn(instanceStatsService, 'refreshAppCountSnapshot');
-    jest.spyOn(instanceStatsService, 'getLabeledAppCounts');
+    const { collectDbMetrics } = registerPrometheusMetrics(
+        config,
+        stores,
+        undefined as unknown as string,
+        config.eventBus,
+        instanceStatsService,
+    );
+    updateMetrics = collectDbMetrics;
+
+    jest.spyOn(clientInstanceStore, 'getDistinctApplicationsCount');
     jest.spyOn(instanceStatsService, 'getStats');
 
-    // validate initial state without calls to these methods
-    expect(instanceStatsService.refreshAppCountSnapshot).toHaveBeenCalledTimes(
-        0,
-    );
     expect(instanceStatsService.getStats).toHaveBeenCalledTimes(0);
 });
 
 test('get snapshot should not call getStats', async () => {
-    await instanceStatsService.refreshAppCountSnapshot();
-    expect(instanceStatsService.getLabeledAppCounts).toHaveBeenCalledTimes(1);
+    await updateMetrics();
+    expect(
+        clientInstanceStore.getDistinctApplicationsCount,
+    ).toHaveBeenCalledTimes(3);
     expect(instanceStatsService.getStats).toHaveBeenCalledTimes(0);
 
-    // subsequent calls to getStatsSnapshot don't call getStats
     for (let i = 0; i < 3; i++) {
         const { clientApps } = await instanceStatsService.getStats();
         expect(clientApps).toStrictEqual([
@@ -51,12 +64,11 @@ test('get snapshot should not call getStats', async () => {
         ]);
     }
     // after querying the stats snapshot no call to getStats should be issued
-    expect(instanceStatsService.getLabeledAppCounts).toHaveBeenCalledTimes(1);
+    expect(
+        clientInstanceStore.getDistinctApplicationsCount,
+    ).toHaveBeenCalledTimes(3);
 });
 
 test('before the snapshot is refreshed we can still get the appCount', async () => {
-    expect(instanceStatsService.refreshAppCountSnapshot).toHaveBeenCalledTimes(
-        0,
-    );
     expect(instanceStatsService.getAppCountSnapshot('7d')).toBeUndefined();
 });

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

@@ -109,9 +109,9 @@ export class InstanceStatsService {
 
     private appCount?: Partial<{ [key in TimeRange]: number }>;
 
-    private getActiveUsers: GetActiveUsers;
+    getActiveUsers: GetActiveUsers;
 
-    private getProductionChanges: GetProductionChanges;
+    getProductionChanges: GetProductionChanges;
 
     private featureStrategiesReadModel: IFeatureStrategiesReadModel;
 
@@ -180,25 +180,6 @@ export class InstanceStatsService {
         this.featureStrategiesReadModel = featureStrategiesReadModel;
     }
 
-    async refreshAppCountSnapshot(): Promise<
-        Partial<{ [key in TimeRange]: number }>
-    > {
-        try {
-            this.appCount = await this.getLabeledAppCounts();
-            return this.appCount;
-        } catch (error) {
-            this.logger.warn(
-                'Unable to retrieve statistics. This will be retried',
-                error,
-            );
-            return {
-                '7d': 0,
-                '30d': 0,
-                allTime: 0,
-            };
-        }
-    }
-
     getProjectModeCount(): Promise<ProjectModeCount[]> {
         return this.projectStore.getProjectModeCounts();
     }
@@ -231,9 +212,6 @@ export class InstanceStatsService {
         return settings?.enabled || false;
     }
 
-    /**
-     * use getStatsSnapshot for low latency, sacrificing data-freshness
-     */
     async getStats(): Promise<InstanceStats> {
         const versionInfo = await this.versionService.getVersionInfo();
         const [
@@ -265,30 +243,30 @@ export class InstanceStatsService {
         ] = await Promise.all([
             this.getToggleCount(),
             this.getArchivedToggleCount(),
-            this.userStore.count(),
-            this.userStore.countServiceAccounts(),
-            this.apiTokenStore.countByType(),
+            this.getRegisteredUsers(),
+            this.countServiceAccounts(),
+            this.countApiTokensByType(),
             this.getActiveUsers(),
             this.getProjectModeCount(),
-            this.contextFieldStore.count(),
-            this.groupStore.count(),
-            this.roleStore.count(),
-            this.roleStore.filteredCount({ type: CUSTOM_ROOT_ROLE_TYPE }),
-            this.roleStore.filteredCountInUse({ type: CUSTOM_ROOT_ROLE_TYPE }),
-            this.environmentStore.count(),
-            this.segmentStore.count(),
-            this.strategyStore.count(),
+            this.contextFieldCount(),
+            this.groupCount(),
+            this.roleCount(),
+            this.customRolesCount(),
+            this.customRolesCountInUse(),
+            this.environmentCount(),
+            this.segmentCount(),
+            this.strategiesCount(),
             this.hasSAML(),
             this.hasOIDC(),
-            this.appCount ? this.appCount : this.refreshAppCountSnapshot(),
+            this.appCount ? this.appCount : this.getLabeledAppCounts(),
             this.eventStore.deprecatedFilteredCount({
                 type: FEATURES_EXPORTED,
             }),
             this.eventStore.deprecatedFilteredCount({
                 type: FEATURES_IMPORTED,
             }),
             this.getProductionChanges(),
-            this.clientMetricsStore.countPreviousDayHourlyMetricsBuckets(),
+            this.countPreviousDayHourlyMetricsBuckets(),
             this.featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies(),
             this.featureStrategiesReadModel.getMaxConstraintValues(),
             this.featureStrategiesReadModel.getMaxConstraintsPerStrategy(),
@@ -330,6 +308,59 @@ export class InstanceStatsService {
         };
     }
 
+    groupCount(): Promise<number> {
+        return this.groupStore.count();
+    }
+
+    roleCount(): Promise<number> {
+        return this.roleStore.count();
+    }
+
+    customRolesCount(): Promise<number> {
+        return this.roleStore.filteredCount({ type: CUSTOM_ROOT_ROLE_TYPE });
+    }
+
+    customRolesCountInUse(): Promise<number> {
+        return this.roleStore.filteredCountInUse({
+            type: CUSTOM_ROOT_ROLE_TYPE,
+        });
+    }
+
+    segmentCount(): Promise<number> {
+        return this.segmentStore.count();
+    }
+
+    contextFieldCount(): Promise<number> {
+        return this.contextFieldStore.count();
+    }
+
+    strategiesCount(): Promise<number> {
+        return this.strategyStore.count();
+    }
+
+    environmentCount(): Promise<number> {
+        return this.environmentStore.count();
+    }
+
+    countPreviousDayHourlyMetricsBuckets(): Promise<{
+        enabledCount: number;
+        variantCount: number;
+    }> {
+        return this.clientMetricsStore.countPreviousDayHourlyMetricsBuckets();
+    }
+
+    countApiTokensByType(): Promise<Map<string, number>> {
+        return this.apiTokenStore.countByType();
+    }
+
+    getRegisteredUsers(): Promise<number> {
+        return this.userStore.count();
+    }
+
+    countServiceAccounts(): Promise<number> {
+        return this.userStore.countServiceAccounts();
+    }
+
     async getLabeledAppCounts(): Promise<
         Partial<{ [key in TimeRange]: number }>
     > {
@@ -338,11 +369,12 @@ export class InstanceStatsService {
             this.clientInstanceStore.getDistinctApplicationsCount(30),
             this.clientInstanceStore.getDistinctApplicationsCount(),
         ]);
-        return {
+        this.appCount = {
             '7d': t7d,
             '30d': t30d,
             allTime,
         };
+        return this.appCount;
     }
 
     getAppCountSnapshot(range: TimeRange): number | undefined {

src/lib/metrics-gauge.test.ts

@@ -0,0 +1,114 @@
+import { register } from 'prom-client';
+import { createTestConfig } from '../test/config/test-config';
+import type { IUnleashConfig } from './types';
+import { DbMetricsMonitor } from './metrics-gauge';
+
+const prometheusRegister = register;
+let config: IUnleashConfig;
+let dbMetrics: DbMetricsMonitor;
+
+beforeAll(async () => {
+    config = createTestConfig({
+        server: {
+            serverMetrics: true,
+        },
+    });
+});
+
+beforeEach(async () => {
+    dbMetrics = new DbMetricsMonitor(config);
+});
+
+test('should collect registered metrics', async () => {
+    dbMetrics.registerGaugeDbMetric({
+        name: 'my_metric',
+        help: 'This is the answer to life, the univers, and everything',
+        labelNames: [],
+        query: () => Promise.resolve(42),
+        map: (result) => ({ value: result }),
+    });
+
+    await dbMetrics.refreshDbMetrics();
+
+    const metrics = await prometheusRegister.metrics();
+    expect(metrics).toMatch(/my_metric 42/);
+});
+
+test('should collect registered metrics with labels', async () => {
+    dbMetrics.registerGaugeDbMetric({
+        name: 'life_the_universe_and_everything',
+        help: 'This is the answer to life, the univers, and everything',
+        labelNames: ['test'],
+        query: () => Promise.resolve(42),
+        map: (result) => ({ value: result, labels: { test: 'case' } }),
+    });
+
+    await dbMetrics.refreshDbMetrics();
+
+    const metrics = await prometheusRegister.metrics();
+    expect(metrics).toMatch(
+        /life_the_universe_and_everything\{test="case"\} 42/,
+    );
+});
+
+test('should collect multiple registered metrics with and without labels', async () => {
+    dbMetrics.registerGaugeDbMetric({
+        name: 'my_first_metric',
+        help: 'This is the answer to life, the univers, and everything',
+        labelNames: [],
+        query: () => Promise.resolve(42),
+        map: (result) => ({ value: result }),
+    });
+
+    dbMetrics.registerGaugeDbMetric({
+        name: 'my_other_metric',
+        help: 'This is Eulers number',
+        labelNames: ['euler'],
+        query: () => Promise.resolve(Math.E),
+        map: (result) => ({ value: result, labels: { euler: 'number' } }),
+    });
+
+    await dbMetrics.refreshDbMetrics();
+
+    const metrics = await prometheusRegister.metrics();
+    expect(metrics).toMatch(/my_first_metric 42/);
+    expect(metrics).toMatch(/my_other_metric\{euler="number"\} 2.71828/);
+});
+
+test('should support different label and value pairs', async () => {
+    dbMetrics.registerGaugeDbMetric({
+        name: 'multi_dimensional',
+        help: 'This metric has different values for different labels',
+        labelNames: ['version', 'range'],
+        query: () => Promise.resolve(2),
+        map: (result) => [
+            { value: result, labels: { version: '1', range: 'linear' } },
+            {
+                value: result * result,
+                labels: { version: '2', range: 'square' },
+            },
+            { value: result / 2, labels: { version: '3', range: 'half' } },
+        ],
+    });
+
+    await dbMetrics.refreshDbMetrics();
+
+    const metrics = await prometheusRegister.metrics();
+    expect(metrics).toMatch(
+        /multi_dimensional\{version="1",range="linear"\} 2\nmulti_dimensional\{version="2",range="square"\} 4\nmulti_dimensional\{version="3",range="half"\} 1/,
+    );
+    expect(
+        await dbMetrics.findValue('multi_dimensional', { range: 'linear' }),
+    ).toBe(2);
+    expect(
+        await dbMetrics.findValue('multi_dimensional', { range: 'half' }),
+    ).toBe(1);
+    expect(
+        await dbMetrics.findValue('multi_dimensional', { range: 'square' }),
+    ).toBe(4);
+    expect(
+        await dbMetrics.findValue('multi_dimensional', { range: 'x' }),
+    ).toBeUndefined();
+    expect(await dbMetrics.findValue('multi_dimensional')).toBe(2); // first match
+    expect(await dbMetrics.findValue('other')).toBeUndefined();
+});