1 files changed, 59 insertions, 0 deletions
diff --git a/lib/trace/metric/trace.ts b/lib/trace/metric/trace.ts
new file mode 100644
index 0000000..0c5fe37
--- /dev/null
+++ b/lib/trace/metric/trace.ts
@@ -0,0 +1,59 @@
+import { IMetric, isIMetric, isMetricValue, ITrace, ITraceWith, MetricValue, SideEffect } from '@emprespresso/pengueno';
+
+export type MetricsTraceSupplier =
+    | ITraceWith<IMetric | MetricValue | undefined>
+    | Array<ITraceWith<IMetric | MetricValue | undefined>>;
+export const isMetricsTraceSupplier = (t: unknown): t is MetricsTraceSupplier =>
+    isMetricValue(t) || isIMetric(t) || (Array.isArray(t) && t.every((_m) => isMetricValue(_m) || isIMetric(_m)));
+
+export class MetricsTrace implements ITrace<MetricsTraceSupplier> {
+    constructor(
+        private readonly metricConsumer: SideEffect<Array<MetricValue>>,
+        private readonly activeTraces: ReadonlyMap<IMetric, number> = new Map(),
+        private readonly completedTraces: ReadonlySet<IMetric> = new Set(),
+    ) {}
+
+    public traceScope(trace: MetricsTraceSupplier): MetricsTrace {
+        const now = Date.now();
+        const metricsToTrace = (Array.isArray(trace) ? trace : [trace]).filter(isIMetric);
+
+        const initialTraces = new Map(metricsToTrace.map((metric) => [metric, now]));
+
+        return new MetricsTrace(this.metricConsumer, initialTraces);
+    }
+
+    public trace(metrics: MetricsTraceSupplier): MetricsTrace {
+        if (!metrics || typeof metrics === 'string') {
+            return this;
+        }
+
+        const now = Date.now();
+        const allMetrics = Array.isArray(metrics) ? metrics : [metrics];
+
+        // partition the incoming metrics
+        const valuesToEmit = allMetrics.filter(isMetricValue);
+        const traceableMetrics = allMetrics.filter(isIMetric);
+
+        const metricsToStart = traceableMetrics.filter((m) => !this.activeTraces.has(m));
+        const metricsToEnd = traceableMetrics.filter((m) => this.activeTraces.has(m) && !this.completedTraces.has(m));
+
+        // the new metrics to emit based on traces ending *now*
+        const endedMetricValues = metricsToEnd.flatMap((metric) => [
+            metric.count.withValue(1.0),
+            metric.time.withValue(now - this.activeTraces.get(metric)!),
+        ]);
+
+        const allMetricsToEmit = [...valuesToEmit, ...endedMetricValues];
+        if (allMetricsToEmit.length > 0) {
+            this.metricConsumer(allMetricsToEmit);
+        }
+
+        // the next immutable state
+        const nextActiveTraces = new Map([
+            ...this.activeTraces,
+            ...metricsToStart.map((m): [IMetric, number] => [m, now]),
+        ]);
+        const nextCompletedTraces = new Set([...this.completedTraces, ...metricsToEnd]);
+        return new MetricsTrace(this.metricConsumer, nextActiveTraces, nextCompletedTraces);
+    }
+}