Feat: add withLabels()

metrics/registry.go

@@ -1,6 +1,7 @@
 package metrics
 
 import (
+	"sort"
 	"strings"
 	"sync"
 
@@ -37,7 +38,7 @@ func (r *registry) sanitizeMetricName(name string) string {
 	return strings.ReplaceAll(name, "-", "_")
 }
 
-// Inc increments a counter for the given Series.
+// Inc increments a counter for the given Series, allowing additional custom labels.
 func (r *registry) Inc(name string, labels prometheus.Labels) {
 	r.metricsMu.Lock()
 	defer r.metricsMu.Unlock()
@@ -49,14 +50,14 @@ func (r *registry) Inc(name string, labels prometheus.Labels) {
 			Subsystem: r.Subsystem,
 			Namespace: r.Namespace,
 			Name:      sanitized,
-		}, []string{"series_type", "sub_type", "operation", "status", "error_code"})
+		}, labelKeys(labels))
 		r.PromRegistry.MustRegister(counter)
 		r.counters[sanitized] = counter
 	}
 	counter.With(labels).Inc()
 }
 
-// RecordDuration records a duration for the given Series.
+// RecordDuration records a duration for the given Series, allowing additional custom labels.
 func (r *registry) RecordDuration(name string, labels prometheus.Labels, duration float64) {
 	r.metricsMu.Lock()
 	defer r.metricsMu.Unlock()
@@ -69,7 +70,7 @@ func (r *registry) RecordDuration(name string, labels prometheus.Labels, duratio
 			Namespace: r.Namespace,
 			Name:      sanitized,
 			Buckets:   prometheus.DefBuckets,
-		}, []string{"series_type", "sub_type", "operation"})
+		}, labelKeys(labels))
 		r.PromRegistry.MustRegister(histogram)
 		r.histograms[sanitized] = histogram
 	}
@@ -88,3 +89,14 @@ func registerMetrics(registry *registry) {
 			collectors.WithGoCollectorRuntimeMetrics(collectors.MetricsScheduler),
 		))
 }
+
+// labelKeys returns the keys of the labels in a deterministic order.
+func labelKeys(labels prometheus.Labels) []string {
+	keys := make([]string, 0, len(labels))
+	for k := range labels {
+		keys = append(keys, k)
+	}
+	// Sort keys to ensure consistent label ordering
+	sort.Strings(keys)
+	return keys
+}

metrics/series.go

@@ -13,7 +13,8 @@
 		seriesType SeriesType
 		subType    string
 		operation  string
 		status     string
+		labels     prometheus.Labels
 	}
 
 	seriesContextKey struct{}
@@ -32,15 +33,24 @@
 	SeriesTypeDatabusConsumer SeriesType = "databus_consumer"
 )
 
-// NewSeries creates a new Series instance with the given type and name.
+// NewSeries creates a new Series instance with the given type and subType.
 func NewSeries(st SeriesType, subType string) Series {
 	return Series{
 		seriesType: st,
 		subType:    subType,
 		operation:  "undefined",
+		labels:     make(prometheus.Labels),
 	}
 }
 
+// WithLabels adds custom labels to the Series.
+func (s Series) WithLabels(labels prometheus.Labels) Series {
+	for k, v := range labels {
+		s.labels[k] = v
+	}
+	return s
+}
+
 // FromContext retrieves the Series from the context.
 func FromContext(ctx context.Context) Series {
 	series, ok := ctx.Value(seriesContextKey{}).(Series)
@@ -61,6 +71,7 @@
 			seriesType: s.seriesType,
 			subType:    s.subType,
 			operation:  series.appendOperation(operation),
+			labels:     series.labels,
 		}
 
 		return series.ToContext(ctx), series
@@ -70,6 +81,7 @@
 		seriesType: s.seriesType,
 		subType:    s.subType,
 		operation:  operation,
+		labels:     s.labels,
 	}
 
 	return series.ToContext(ctx), series
@@ -88,45 +100,58 @@
 
 // Info returns the metric name and labels for an informational event.
 func (s Series) Info(message string) (string, prometheus.Labels) {
-	return "operation_count", prometheus.Labels{
+	labels := prometheus.Labels{
 		"series_type":  s.seriesType.String(),
 		"sub_type":     s.subType,
 		"operation":    s.operation,
 		"status":       seriesTypeInfo,
 		"info_message": message,
 	}
+	return "operation_count", mergeLabels(labels, s.labels)
 }
 
 // Success returns the metric name and labels for a success event.
 func (s Series) Success() (string, prometheus.Labels) {
-	return "operation_count", prometheus.Labels{
+	labels := prometheus.Labels{
 		"series_type": s.seriesType.String(),
 		"sub_type":    s.subType,
 		"operation":   s.operation,
 		"status":      seriesTypeSuccess,
 	}
+	return "operation_count", mergeLabels(labels, s.labels)
 }
 
 // Error returns the metric name and labels for an error event.
 func (s Series) Error(message string) (string, prometheus.Labels) {
-	return "operation_count", prometheus.Labels{
+	labels := prometheus.Labels{
 		"series_type":   s.seriesType.String(),
 		"sub_type":      s.subType,
 		"operation":     s.operation,
 		"status":        seriesTypeError,
 		"error_message": message,
 	}
+	return "operation_count", mergeLabels(labels, s.labels)
 }
 
 // Duration returns the metric name and labels for recording a duration.
 func (s Series) Duration() (string, prometheus.Labels) {
-	return "operation_duration_seconds", prometheus.Labels{
+	labels := prometheus.Labels{
 		"series_type": s.seriesType.String(),
 		"sub_type":    s.subType,
 		"operation":   s.operation,
 	}
+	return "operation_duration_seconds", mergeLabels(labels, s.labels)
 }
 
+// appendOperation appends the operation to the Series operation string.
 func (s Series) appendOperation(operation string) string {
 	return s.operation + "_" + operation
 }
+
+// mergeLabels merges a set of additional labels into the base labels.
+func mergeLabels(base prometheus.Labels, additional prometheus.Labels) prometheus.Labels {
+	for k, v := range additional {
+		base[k] = v
+	}
+	return base
+}