Get closer to prom-client types

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();
+});

src/lib/metrics-gauge.ts

@@ -2,37 +2,40 @@ import type { Logger } from './logger';
 import type { IUnleashConfig } from './types';
 import { createGauge, type Gauge } from './util/metrics';
 
-type RestrictedRecord<T extends readonly string[]> = Record<T[number], string>;
 type Query<R> = () => Promise<R | undefined | null>;
-type MetricValue<R> = {
-    count: number;
-    labels: RestrictedRecord<GaugeDefinition<R>['labelNames']>;
+type MetricValue<L extends string> = {
+    value: number;
+    labels?: Record<L, string | number>;
 };
-type MapResult<R> = (result: R) => MetricValue<R> | MetricValue<R>[];
+type MapResult<R, L extends string> = (
+    result: R,
+) => MetricValue<L> | MetricValue<L>[];
 
-type GaugeDefinition<T> = {
+type GaugeDefinition<T, L extends string> = {
     name: string;
     help: string;
-    labelNames: string[];
+    labelNames: L[];
     query: Query<T>;
-    map: MapResult<T>;
+    map: MapResult<T, L>;
 };
 
 type Task = () => Promise<void>;
 export class DbMetricsMonitor {
     private tasks: Set<Task> = new Set();
     private gauges: Map<string, Gauge<string>> = new Map();
-    private logger: Logger;
+    private log: Logger;
 
-    constructor(config: IUnleashConfig) {
-        this.logger = config.getLogger('gauge-metrics');
+    constructor({ getLogger }: Pick<IUnleashConfig, 'getLogger'>) {
+        this.log = getLogger('gauge-metrics');
     }
 
     private asArray<T>(value: T | T[]): T[] {
         return Array.isArray(value) ? value : [value];
     }
 
-    registerGaugeDbMetric<T>(definition: GaugeDefinition<T>): Task {
+    registerGaugeDbMetric<T, L extends string>(
+        definition: GaugeDefinition<T, L>,
+    ): Task {
         const gauge = createGauge(definition);
         this.gauges.set(definition.name, gauge);
         const task = async () => {
@@ -42,11 +45,15 @@ export class DbMetricsMonitor {
                     const results = this.asArray(definition.map(result));
                     gauge.reset();
                     for (const r of results) {
-                        gauge.labels(r.labels).set(r.count);
+                        if (r.labels) {
+                            gauge.labels(r.labels).set(r.value);
+                        } else {
+                            gauge.set(r.value);
+                        }
                     }
                 }
             } catch (e) {
-                this.logger.warn(`Failed to refresh ${definition.name}`, e);
+                this.log.warn(`Failed to refresh ${definition.name}`, e);
             }
         };
         this.tasks.add(task);
@@ -59,10 +66,21 @@ export class DbMetricsMonitor {
         }
     };
 
-    async getLastValue(name: string): Promise<number | undefined> {
+    async findValue(
+        name: string,
+        labels?: Record<string, string | number>,
+    ): Promise<number | undefined> {
         const gauge = await this.gauges.get(name)?.gauge?.get();
         if (gauge && gauge.values.length > 0) {
-            return gauge.values[0].value;
+            const values = labels
+                ? gauge.values.filter(({ labels: l }) => {
+                      return Object.entries(labels).every(
+                          ([key, value]) => l[key] === value,
+                      );
+                  })
+                : gauge.values;
+            // return first value
+            return values.map(({ value }) => value).shift();
         }
         return undefined;
     }

src/lib/metrics.ts

@@ -124,7 +124,7 @@ export default class MetricsMonitor {
             help: 'Number of feature flags',
             labelNames: ['version'],
             query: () => instanceStatsService.getToggleCount(),
-            map: (count) => ({ count, labels: { version } }),
+            map: (value) => ({ value, labels: { version } }),
         })();
 
         dbMetrics.registerGaugeDbMetric({
@@ -134,7 +134,7 @@ export default class MetricsMonitor {
             query: () =>
                 stores.featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies(),
             map: (result) => ({
-                count: result.count,
+                value: result.count,
                 labels: {
                     environment: result.environment,
                     feature: result.feature,
@@ -149,7 +149,7 @@ export default class MetricsMonitor {
             query: () =>
                 stores.featureStrategiesReadModel.getMaxFeatureStrategies(),
             map: (result) => ({
-                count: result.count,
+                value: result.count,
                 labels: { feature: result.feature },
             }),
         });
@@ -268,7 +268,7 @@ export default class MetricsMonitor {
             query: () => instanceStatsService.getLabeledAppCounts(),
             map: (result) =>
                 Object.entries(result).map(([range, count]) => ({
-                    count,
+                    value: count,
                     labels: { range },
                 })),
         })();