diff --git a/src/lib/features/metrics/impact/batch-histogram.test.ts b/src/lib/features/metrics/impact/batch-histogram.test.ts new file mode 100644 index 0000000000..6e7e0e7019 --- /dev/null +++ b/src/lib/features/metrics/impact/batch-histogram.test.ts @@ -0,0 +1,196 @@ +import { Registry } from 'prom-client'; +import { BatchHistogram } from './batch-histogram.js'; + +describe('BatchHistogram', () => { + let registry: Registry; + let histogram: BatchHistogram; + + beforeEach(() => { + registry = new Registry(); + histogram = new BatchHistogram({ + name: 'test_histogram', + help: 'Test histogram', + registry: registry, + }); + }); + + test('should record batch data and preserve bucket distribution', async () => { + histogram.recordBatch( + { label1: 'value1' }, + { + count: 10, + sum: 15.5, + buckets: [ + { le: 0.1, count: 2 }, + { le: 0.5, count: 5 }, + { le: 1, count: 7 }, + { le: 2.5, count: 9 }, + { le: 5, count: 10 }, + { le: '+Inf', count: 10 }, + ], + }, + ); + + const metrics = await registry.metrics(); + + expect(metrics).toContain('# HELP test_histogram Test histogram'); + expect(metrics).toContain('# TYPE test_histogram histogram'); + + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="0.1"} 2/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="0.5"} 5/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="1"} 7/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="2.5"} 9/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="5"} 10/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="\+Inf"} 10/, + ); + + expect(metrics).toMatch(/test_histogram_sum{label1="value1"} 15\.5/); + expect(metrics).toMatch(/test_histogram_count{label1="value1"} 10/); + }); + + test('should aggregate multiple batches correctly', async () => { + histogram.recordBatch( + { label1: 'value1' }, + { + count: 5, + sum: 7.5, + buckets: [ + { le: 0.1, count: 1 }, + { le: 0.5, count: 3 }, + { le: 1, count: 4 }, + { le: 2.5, count: 5 }, + { le: 5, count: 5 }, + { le: '+Inf', count: 5 }, + ], + }, + ); + + histogram.recordBatch( + { label1: 'value1' }, + { + count: 3, + sum: 4.2, + buckets: [ + { le: 0.1, count: 0 }, + { le: 0.5, count: 1 }, + { le: 1, count: 2 }, + { le: 2.5, count: 3 }, + { le: 5, count: 3 }, + { le: '+Inf', count: 3 }, + ], + }, + ); + + const metrics = await registry.metrics(); + + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="0.1"} 1/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="0.5"} 4/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="1"} 6/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="2.5"} 8/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="5"} 8/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{label1="value1",le="\+Inf"} 8/, + ); + + expect(metrics).toMatch(/test_histogram_sum{label1="value1"} 11\.7/); + expect(metrics).toMatch(/test_histogram_count{label1="value1"} 8/); + }); + + test('should record different labels separately', async () => { + histogram.recordBatch( + { service: 'api', app: 'my_app' }, + { + count: 3, + sum: 1.5, + buckets: [ + { le: 1, count: 2 }, + { le: '+Inf', count: 3 }, + ], + }, + ); + + histogram.recordBatch( + { service: 'web', app: 'my_app' }, + { + count: 2, + sum: 3.0, + buckets: [ + { le: 1, count: 1 }, + { le: '+Inf', count: 2 }, + ], + }, + ); + + const metrics = await registry.metrics(); + + expect(metrics).toMatch( + /test_histogram_bucket{app="my_app",service="api",le="1"} 2/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{app="my_app",service="api",le="\+Inf"} 3/, + ); + expect(metrics).toMatch( + /test_histogram_sum{app="my_app",service="api"} 1\.5/, + ); + expect(metrics).toMatch( + /test_histogram_count{app="my_app",service="api"} 3/, + ); + + expect(metrics).toMatch( + /test_histogram_bucket{app="my_app",service="web",le="1"} 1/, + ); + expect(metrics).toMatch( + /test_histogram_bucket{app="my_app",service="web",le="\+Inf"} 2/, + ); + expect(metrics).toMatch( + /test_histogram_sum{app="my_app",service="web"} 3/, + ); + expect(metrics).toMatch( + /test_histogram_count{app="my_app",service="web"} 2/, + ); + }); + + test('should handle "+Inf" string from SDK serialization', async () => { + histogram.recordBatch( + { client: 'sdk' }, + { + count: 5, + sum: 12.3, + buckets: [ + { le: 1, count: 3 }, + { le: '+Inf', count: 5 }, // String instead of Infinity + ], + }, + ); + + const metrics = await registry.metrics(); + + expect(metrics).toMatch(/test_histogram_bucket{client="sdk",le="1"} 3/); + expect(metrics).toMatch( + /test_histogram_bucket{client="sdk",le="\+Inf"} 5/, + ); + expect(metrics).toMatch(/test_histogram_sum{client="sdk"} 12\.3/); + expect(metrics).toMatch(/test_histogram_count{client="sdk"} 5/); + }); +}); diff --git a/src/lib/features/metrics/impact/batch-histogram.ts b/src/lib/features/metrics/impact/batch-histogram.ts new file mode 100644 index 0000000000..4b604c294b --- /dev/null +++ b/src/lib/features/metrics/impact/batch-histogram.ts @@ -0,0 +1,126 @@ +import type { Registry } from 'prom-client'; + +interface BucketData { + le: number | '+Inf'; + count: number; +} + +interface BatchData { + count: number; + sum: number; + buckets: BucketData[]; +} + +export class BatchHistogram { + private name: string; + private help: string; + private registry: Registry; + + // Store accumulated data for each label combination + private store: Map< + string, + { + count: number; + sum: number; + buckets: Map; + } + > = new Map(); + + constructor(config: { + name: string; + help: string; + registry: Registry; + }) { + this.name = config.name; + this.help = config.help; + this.registry = config.registry; + + this.registry.registerMetric(this as any); + } + + recordBatch( + labels: Record, + data: BatchData, + ): void { + const labelKey = this.createLabelKey(labels); + + let entry = this.store.get(labelKey); + if (!entry) { + entry = { + count: 0, + sum: 0, + buckets: new Map(), + }; + this.store.set(labelKey, entry); + } + + entry.count += data.count; + entry.sum += data.sum; + + for (const bucket of data.buckets) { + const current = entry.buckets.get(bucket.le) || 0; + entry.buckets.set(bucket.le, current + bucket.count); + } + } + + private createLabelKey(labels: Record): string { + const sortedKeys = Object.keys(labels).sort(); + return sortedKeys.map((key) => `${key}:${labels[key]}`).join(','); + } + + reset(): void { + this.store.clear(); + } + + get() { + const values: any[] = []; + + for (const [labelKey, data] of this.store) { + const labels: Record = {}; + if (labelKey) { + labelKey.split(',').forEach((pair) => { + const [key, value] = pair.split(':'); + labels[key] = value; + }); + } + + for (const [le, cumulativeCount] of Array.from( + data.buckets.entries(), + ).sort((a, b) => { + // Sort buckets: numbers first (ascending), then '+Inf' last + if (a[0] === '+Inf') return 1; + if (b[0] === '+Inf') return -1; + return (a[0] as number) - (b[0] as number); + })) { + values.push({ + value: cumulativeCount, + labels: { + ...labels, + le: le.toString(), + }, + metricName: `${this.name}_bucket`, + }); + } + + values.push({ + value: data.sum, + labels, + metricName: `${this.name}_sum`, + }); + + values.push({ + value: data.count, + labels, + metricName: `${this.name}_count`, + }); + } + + return { + name: this.name, + help: this.help, + type: 'histogram', + values, + aggregator: 'sum', + }; + } +}