From 0a77ca1cf7e7a79c843b22adf7c82563ddf462a3 Mon Sep 17 00:00:00 2001
From: Donghyuk Kang <donghyuk1003+github@gmail.com>
Date: Tue, 24 Oct 2023 13:59:38 +0900
Subject: [PATCH] feat: implement startTimerWithExemplar

---
 lib/histogram.js      | 29 +++++++++++++++++++++++++++--
 test/exemplarsTest.js | 25 +++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/lib/histogram.js b/lib/histogram.js
index 18c99268..723e0192 100644
--- a/lib/histogram.js
+++ b/lib/histogram.js
@@ -24,8 +24,10 @@ class Histogram extends Metric {
 		this.type = 'histogram';
 		this.defaultLabels = {};
 		this.defaultExemplarLabelSet = {};
+		this.enableExemplars = false;
 
 		if (config.enableExemplars) {
+			this.enableExemplars = true;
 			this.observe = this.observeWithExemplar;
 		} else {
 			this.observe = this.observeWithoutExemplar;
@@ -143,6 +145,7 @@ class Histogram extends Metric {
 	/**
 	 * Start a timer that could be used to logging durations
 	 * @param {object} labels - Object with labels where key is the label key and value is label value. Can only be one level deep
+	 * @param {object} exemplarLabels - Object with labels for exemplar where key is the label key and value is label value. Can only be one level deep
 	 * @returns {function} - Function to invoke when you want to stop the timer and observe the duration in seconds
 	 * @example
 	 * var end = histogram.startTimer();
@@ -151,8 +154,10 @@ class Histogram extends Metric {
 	 * 	console.log('Duration', duration);
 	 * });
 	 */
-	startTimer(labels) {
-		return startTimer.call(this, labels)();
+	startTimer(labels, exemplarLabels) {
+		return this.enableExemplars
+			? startTimerWithExemplar.call(this, labels, exemplarLabels)()
+			: startTimer.call(this, labels)();
 	}
 
 	labels(...args) {
@@ -183,6 +188,26 @@ function startTimer(startLabels) {
 	};
 }
 
+function startTimerWithExemplar(startLabels, startExemplarLabels) {
+	return () => {
+		const start = process.hrtime();
+		return (endLabels, endExemplarLabels) => {
+			const delta = process.hrtime(start);
+			const value = delta[0] + delta[1] / 1e9;
+			this.observe({
+				labels: Object.assign({}, startLabels, endLabels),
+				value,
+				exemplarLabels: Object.assign(
+					{},
+					startExemplarLabels,
+					endExemplarLabels,
+				),
+			});
+			return value;
+		};
+	};
+}
+
 function setValuePair(labels, value, metricName, exemplar, sharedLabels = {}) {
 	return {
 		labels,
diff --git a/test/exemplarsTest.js b/test/exemplarsTest.js
index 13a38bc2..68a6046e 100644
--- a/test/exemplarsTest.js
+++ b/test/exemplarsTest.js
@@ -147,6 +147,31 @@ describe('Exemplars', () => {
 				jest.useRealTimers();
 			});
 
+			it('should allow exemplar labels before and after timers', async () => {
+				jest.useFakeTimers('modern');
+				jest.setSystemTime(0);
+				const histogramInstance = new Histogram({
+					name: 'histogram_start_timer_exemplar_label_test',
+					help: 'test',
+					labelNames: ['method', 'code'],
+					enableExemplars: true,
+				});
+				const end = histogramInstance.startTimer(
+					{ method: 'get' },
+					{ traceId: 'trace_id_test_1' },
+				);
+
+				jest.advanceTimersByTime(500);
+				end({ code: '200' }, { spanId: 'span_id_test_1' });
+
+				const vals = (await histogramInstance.get()).values;
+				expect(getValuesByLabel(0.5, vals)[0].value).toEqual(1);
+				expect(
+					getValuesByLabel(0.5, vals)[0].exemplar.labelSet.traceId,
+				).toEqual('trace_id_test_1');
+				jest.useRealTimers();
+			});
+
 			function getValueByLabel(label, values, key) {
 				return values.reduce((acc, val) => {
 					if (val.labels && val.labels[key || 'le'] === label) {