From 677ec9cee8ed45a93393f4d180cc51bce281f4ae Mon Sep 17 00:00:00 2001 From: Jaanus Sellin Date: Fri, 25 Oct 2024 15:43:14 +0300 Subject: [PATCH] feat: send traffic info to prometheus (#8541) --- .../createInstanceStatsService.ts | 8 +++++ .../instance-stats/instance-stats-service.ts | 17 +++++++++++ .../fake-traffic-data-usage-store.ts | 29 ++++++++++++++++--- src/lib/metrics.test.ts | 6 ++++ src/lib/metrics.ts | 9 ++++++ 5 files changed, 65 insertions(+), 4 deletions(-) diff --git a/src/lib/features/instance-stats/createInstanceStatsService.ts b/src/lib/features/instance-stats/createInstanceStatsService.ts index 134e9f450151..f3dd4a0cd482 100644 --- 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 index 24e04ef3dfaa..e7d4a64f3272 100644 --- 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 { @@ -361,6 +368,16 @@ export class InstanceStatsService { return this.userStore.countServiceAccounts(); } + async getCurrentTrafficData(): Promise { + 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 index bb4f555e64ef..b3e6c0f2bf61 100644 --- 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 { 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 { - throw new Error('Method not implemented.'); + async upsert(trafficDataUsage: IStatTrafficUsage): Promise { + 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 { - throw new Error('Method not implemented.'); + + async getTrafficDataUsageForPeriod( + period: string, + ): Promise { + 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 index d007a2b64aee..4b53c8335217 100644 --- 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 index 2b082c011331..c2adb38be3df 100644 --- 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 [