From 3cbbefff1ddd57b25c97c0730b1b3512eca0a08f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gast=C3=B3n=20Fournier?= <gaston@getunleash.io>
Date: Mon, 14 Oct 2024 16:13:42 +0200
Subject: [PATCH] chore: add a class to handle aggreggation queries

---
 src/lib/metrics-gauge.ts | 50 ++++++++++++++++++++++++++++++++++++++++
 src/lib/metrics.ts       | 23 ++++++++++--------
 2 files changed, 64 insertions(+), 9 deletions(-)
 create mode 100644 src/lib/metrics-gauge.ts

diff --git a/src/lib/metrics-gauge.ts b/src/lib/metrics-gauge.ts
new file mode 100644
index 0000000000..6d1d6fe177
--- /dev/null
+++ b/src/lib/metrics-gauge.ts
@@ -0,0 +1,50 @@
+import { createGauge, type Gauge } from './util/metrics';
+
+type RestrictedRecord<T extends string[]> = Record<T[number], string>;
+type Query<R> = () => Promise<R | undefined | null>;
+type MapResult<R> = (result: R) => {
+    count: number;
+    labels: RestrictedRecord<GaugeDefinition<R>['labelNames']>;
+};
+
+type GaugeDefinition<T> = {
+    name: string;
+    help: string;
+    labelNames: string[];
+    query: Query<T>;
+    map: MapResult<T>;
+};
+
+export class DbMetricsMonitor {
+    private tasks: (() => Promise<void>)[] = [];
+    private gauges: Record<string, Gauge<string>> = {};
+
+    constructor() {}
+
+    registerGaugeDbMetric<T>(definition: GaugeDefinition<T>) {
+        const gauge = createGauge(definition);
+        this.gauges[definition.name] = gauge;
+        this.tasks.push(async () => {
+            const result = await definition.query();
+            if (result) {
+                const { count, labels } = definition.map(result);
+                gauge.reset();
+                gauge.labels(labels).set(count);
+            }
+        });
+    }
+
+    refreshDbMetrics = async () => {
+        for (const task of this.tasks) {
+            await task();
+        }
+    };
+
+    async getLastValue(name: string): Promise<number | undefined> {
+        try {
+            return (await this.gauges[name].gauge.get()).values[0].value;
+        } catch (e) {
+            return undefined;
+        }
+    }
+}
diff --git a/src/lib/metrics.ts b/src/lib/metrics.ts
index 6a8c973073..6d08ffde5c 100644
--- a/src/lib/metrics.ts
+++ b/src/lib/metrics.ts
@@ -37,6 +37,7 @@ import {
 } from './util/metrics';
 import type { SchedulerService } from './services';
 import type { IClientMetricsEnv } from './features/metrics/client-metrics/client-metrics-store-v2-type';
+import { DbMetricsMonitor } from './metrics-gauge';
 
 export default class MetricsMonitor {
     constructor() {}
@@ -56,6 +57,7 @@ export default class MetricsMonitor {
 
         const { eventStore, environmentStore } = stores;
         const { flagResolver } = config;
+        const dbMetrics = new DbMetricsMonitor();
 
         const cachedEnvironments: () => Promise<IEnvironment[]> = memoizee(
             async () => environmentStore.getAll(),
@@ -125,11 +127,19 @@ export default class MetricsMonitor {
             help: 'Maximum number of environment strategies in one feature',
             labelNames: ['feature', 'environment'],
         });
-        const maxFeatureStrategies = createGauge({
+
+        dbMetrics.registerGaugeDbMetric({
             name: 'max_feature_strategies',
             help: 'Maximum number of strategies in one feature',
             labelNames: ['feature'],
+            query: () =>
+                stores.featureStrategiesReadModel.getMaxFeatureStrategies(),
+            map: (result) => ({
+                count: result.count,
+                labels: { feature: result.feature },
+            }),
         });
+
         const maxConstraintValues = createGauge({
             name: 'max_constraint_values',
             help: 'Maximum number of constraint values used in a single constraint',
@@ -394,9 +404,10 @@ export default class MetricsMonitor {
 
         async function collectStaticCounters() {
             try {
+                dbMetrics.refreshDbMetrics();
+
                 const stats = await instanceStatsService.getStats();
                 const [
-                    maxStrategies,
                     maxEnvironmentStrategies,
                     maxConstraintValuesResult,
                     maxConstraintsPerStrategyResult,
@@ -408,7 +419,6 @@ export default class MetricsMonitor {
                     instanceOnboardingMetrics,
                     projectsOnboardingMetrics,
                 ] = await Promise.all([
-                    stores.featureStrategiesReadModel.getMaxFeatureStrategies(),
                     stores.featureStrategiesReadModel.getMaxFeatureEnvironmentStrategies(),
                     stores.featureStrategiesReadModel.getMaxConstraintValues(),
                     stores.featureStrategiesReadModel.getMaxConstraintsPerStrategy(),
@@ -511,12 +521,7 @@ export default class MetricsMonitor {
                         })
                         .set(maxEnvironmentStrategies.count);
                 }
-                if (maxStrategies) {
-                    maxFeatureStrategies.reset();
-                    maxFeatureStrategies
-                        .labels({ feature: maxStrategies.feature })
-                        .set(maxStrategies.count);
-                }
+
                 if (maxConstraintValuesResult) {
                     maxConstraintValues.reset();
                     maxConstraintValues