feat: send traffic info to prometheus (#8541)

Unleash · Oct 25, 2024 · 677ec9c · 677ec9c
1 parent 7c37de3
commit 677ec9c
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 4 deletions.
diff --git a/src/lib/features/instance-stats/createInstanceStatsService.ts b/src/lib/features/instance-stats/createInstanceStatsService.ts
@@ -42,6 +42,8 @@ import FakeStrategiesStore from '../../../test/fixtures/fake-strategies-store';
 import FakeFeatureStrategiesStore from '../feature-toggle/fakes/fake-feature-strategies-store';
 import { FeatureStrategiesReadModel } from '../feature-toggle/feature-strategies-read-model';
 import { FakeFeatureStrategiesReadModel } from '../feature-toggle/fake-feature-strategies-read-model';
+import { TrafficDataUsageStore } from '../traffic-data-usage/traffic-data-usage-store';
+import { FakeTrafficDataUsageStore } from '../traffic-data-usage/fake-traffic-data-usage-store';
 
 export const createInstanceStatsService = (db: Db, config: IUnleashConfig) => {
     const { eventBus, getLogger, flagResolver } = config;
@@ -93,6 +95,9 @@ export const createInstanceStatsService = (db: Db, config: IUnleashConfig) => {
     );
 
     const featureStrategiesReadModel = new FeatureStrategiesReadModel(db);
+
+    const trafficDataUsageStore = new TrafficDataUsageStore(db, getLogger);
+
     const instanceStatsServiceStores = {
         featureToggleStore,
         userStore,
@@ -109,6 +114,7 @@ export const createInstanceStatsService = (db: Db, config: IUnleashConfig) => {
         apiTokenStore,
         clientMetricsStoreV2,
         featureStrategiesReadModel,
+        trafficDataUsageStore,
     };
     const featureStrategiesStore = new FeatureStrategyStore(
         db,
@@ -157,6 +163,7 @@ export const createFakeInstanceStatsService = (config: IUnleashConfig) => {
     const apiTokenStore = new FakeApiTokenStore();
     const clientMetricsStoreV2 = new FakeClientMetricsStoreV2();
     const featureStrategiesReadModel = new FakeFeatureStrategiesReadModel();
+    const trafficDataUsageStore = new FakeTrafficDataUsageStore();
 
     const instanceStatsServiceStores = {
         featureToggleStore,
@@ -174,6 +181,7 @@ export const createFakeInstanceStatsService = (config: IUnleashConfig) => {
         apiTokenStore,
         clientMetricsStoreV2,
         featureStrategiesReadModel,
+        trafficDataUsageStore,
     };
     const featureStrategiesStore = new FakeFeatureStrategiesStore();
     const versionServiceStores = {

diff --git a/src/lib/features/instance-stats/instance-stats-service.ts b/src/lib/features/instance-stats/instance-stats-service.ts
@@ -6,6 +6,7 @@ import type {
     IClientMetricsStoreV2,
     IEventStore,
     IFeatureStrategiesReadModel,
+    ITrafficDataUsageStore,
     IUnleashStores,
 } from '../../types/stores';
 import type { IContextFieldStore } from '../../types/stores/context-field-store';
@@ -29,6 +30,7 @@ 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';
 
 export type TimeRange = 'allTime' | '30d' | '7d';
 
@@ -115,6 +117,8 @@ export class InstanceStatsService {
 
     private featureStrategiesReadModel: IFeatureStrategiesReadModel;
 
+    private trafficDataUsageStore: ITrafficDataUsageStore;
+
     constructor(
         {
             featureToggleStore,
@@ -132,6 +136,7 @@ export class InstanceStatsService {
             apiTokenStore,
             clientMetricsStoreV2,
             featureStrategiesReadModel,
+            trafficDataUsageStore,
         }: Pick<
             IUnleashStores,
             | 'featureToggleStore'
@@ -149,6 +154,7 @@ export class InstanceStatsService {
             | 'apiTokenStore'
             | 'clientMetricsStoreV2'
             | 'featureStrategiesReadModel'
+            | 'trafficDataUsageStore'
         >,
         {
             getLogger,
@@ -178,6 +184,7 @@ export class InstanceStatsService {
         this.clientMetricsStore = clientMetricsStoreV2;
         this.flagResolver = flagResolver;
         this.featureStrategiesReadModel = featureStrategiesReadModel;
+        this.trafficDataUsageStore = trafficDataUsageStore;
     }
 
     getProjectModeCount(): Promise<ProjectModeCount[]> {
@@ -361,6 +368,16 @@ export class InstanceStatsService {
         return this.userStore.countServiceAccounts();
     }
 
+    async getCurrentTrafficData(): Promise<number> {
+        const traffic =
+            await this.trafficDataUsageStore.getTrafficDataUsageForPeriod(
+                format(new Date(), 'yyyy-MM'),
+            );
+
+        const counts = traffic.map((item) => item.count);
+        return counts.reduce((total, current) => total + current, 0);
+    }
+
     async getLabeledAppCounts(): Promise<
         Partial<{ [key in TimeRange]: number }>
     > {

diff --git a/src/lib/features/traffic-data-usage/fake-traffic-data-usage-store.ts b/src/lib/features/traffic-data-usage/fake-traffic-data-usage-store.ts
@@ -3,8 +3,11 @@ import type {
     IStatTrafficUsage,
 } from './traffic-data-usage-store-type';
 import type { ITrafficDataUsageStore } from '../../types';
+import { isSameMonth, parse } from 'date-fns';
 
 export class FakeTrafficDataUsageStore implements ITrafficDataUsageStore {
+    private trafficData: IStatTrafficUsage[] = [];
+
     get(key: IStatTrafficUsageKey): Promise<IStatTrafficUsage> {
         throw new Error('Method not implemented.');
     }
@@ -23,10 +26,28 @@ export class FakeTrafficDataUsageStore implements ITrafficDataUsageStore {
     destroy(): void {
         throw new Error('Method not implemented.');
     }
-    upsert(trafficDataUsage: IStatTrafficUsage): Promise<void> {
-        throw new Error('Method not implemented.');
+    async upsert(trafficDataUsage: IStatTrafficUsage): Promise<void> {
+        const index = this.trafficData.findIndex(
+            (data) =>
+                data.day.getTime() === trafficDataUsage.day.getTime() &&
+                data.trafficGroup === trafficDataUsage.trafficGroup &&
+                data.statusCodeSeries === trafficDataUsage.statusCodeSeries,
+        );
+
+        if (index >= 0) {
+            this.trafficData[index].count += trafficDataUsage.count;
+        } else {
+            this.trafficData.push(trafficDataUsage);
+        }
     }
-    getTrafficDataUsageForPeriod(period: string): Promise<IStatTrafficUsage[]> {
-        throw new Error('Method not implemented.');
+
+    async getTrafficDataUsageForPeriod(
+        period: string,
+    ): Promise<IStatTrafficUsage[]> {
+        const periodDate = parse(period, 'yyyy-MM', new Date());
+
+        return this.trafficData.filter((data) =>
+            isSameMonth(data.day, periodDate),
+        );
     }
 }
diff --git a/src/lib/metrics.test.ts b/src/lib/metrics.test.ts
@@ -351,3 +351,9 @@ test('should collect limit exceeded metrics', async () => {
         /exceeds_limit_error{resource=\"feature flags\",limit=\"5000\"} 1/,
     );
 });
+
+test('should collect traffic_total metrics', async () => {
+    const recordedMetric =
+        await prometheusRegister.getSingleMetricAsString('traffic_total');
+    expect(recordedMetric).toMatch(/traffic_total 0/);
+});
diff --git a/src/lib/metrics.ts b/src/lib/metrics.ts
@@ -277,6 +277,10 @@ export function registerPrometheusMetrics(
         name: 'users_total',
         help: 'Number of users',
     });
+    const trafficTotal = createGauge({
+        name: 'traffic_total',
+        help: 'Traffic used current month',
+    });
     const serviceAccounts = createGauge({
         name: 'service_accounts_total',
         help: 'Number of service accounts',
@@ -957,6 +961,11 @@ export function registerPrometheusMetrics(
                     await instanceStatsService.countServiceAccounts(),
                 );
 
+                trafficTotal.reset();
+                trafficTotal.set(
+                    await instanceStatsService.getCurrentTrafficData(),
+                );
+
                 apiTokens.reset();
 
                 for (const [