Add support for the prometheus-client crate

.github/workflows/ci.yml

@@ -17,14 +17,15 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
 
-      # Build the packages using the different feature flags
+      # Build the packages
       - run: cargo build
-      - run: cargo build --features=metrics
-      - run: cargo build --features=prometheus
       - run: cargo build --features=custom-objective-percentile,custom-objective-latency
 
-      # Run the tests
+      # Run the tests with each of the different metrics libraries
       - run: cargo test --features=prometheus-exporter
+      - run: cargo test --no-default-features --features=prometheus-exporter,metrics --tests
+      - run: cargo test --no-default-features --features=prometheus-exporter,prometheus --tests
+      - run: cargo test --no-default-features --features=prometheus-exporter,prometheus-client --tests
 
       # Compile the examples
       - run: cargo build --package example-axum

autometrics/Cargo.toml

@@ -17,6 +17,7 @@ default = ["opentelemetry"]
 metrics = ["dep:metrics"]
 opentelemetry = ["opentelemetry_api"]
 prometheus = ["const_format", "dep:prometheus", "once_cell"]
+prometheus-client = ["dep:prometheus-client", "once_cell"]
 prometheus-exporter = [
   "metrics-exporter-prometheus",
   "once_cell",
@@ -47,6 +48,9 @@ prometheus = { version = "0.13", default-features = false, optional = true }
 # Used for prometheus feature
 const_format = { version = "0.2", features = ["rust_1_51"], optional = true }
 
+# Used for prometheus-client feature
+prometheus-client = { version = "0.21.1", optional = true }
+
 [dev-dependencies]
 axum = { version = "0.6", features = ["tokio"] }
 regex = "1.7"

autometrics/src/labels.rs

@@ -1,10 +1,16 @@
 use crate::{constants::*, objectives::*};
+#[cfg(feature = "prometheus-client")]
+use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue, LabelValueEncoder};
 use std::ops::Deref;
 
 pub(crate) type Label = (&'static str, &'static str);
 pub type ResultAndReturnTypeLabels = (&'static str, Option<&'static str>);
 
 /// These are the labels used for the `build_info` metric.
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(EncodeLabelSet, Debug, Clone, PartialEq, Eq, Hash)
+)]
 pub struct BuildInfoLabels {
     pub(crate) version: &'static str,
     pub(crate) commit: &'static str,
@@ -30,12 +36,47 @@ impl BuildInfoLabels {
 }
 
 /// These are the labels used for the `function.calls.count` metric.
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(EncodeLabelSet, Debug, Clone, PartialEq, Eq, Hash)
+)]
 pub struct CounterLabels {
     pub(crate) function: &'static str,
     pub(crate) module: &'static str,
     pub(crate) caller: &'static str,
-    pub(crate) result: Option<ResultAndReturnTypeLabels>,
-    pub(crate) objective: Option<(&'static str, ObjectivePercentile)>,
+    pub(crate) result: Option<ResultLabel>,
+    pub(crate) ok: Option<&'static str>,
+    pub(crate) error: Option<&'static str>,
+    pub(crate) objective_name: Option<&'static str>,
+    pub(crate) objective_percentile: Option<ObjectivePercentile>,
+}
+
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(Debug, Clone, PartialEq, Eq, Hash)
+)]
+pub(crate) enum ResultLabel {
+    Ok,
+    Error,
+}
+
+impl ResultLabel {
+    pub(crate) const fn as_str(&self) -> &'static str {
+        match self {
+            ResultLabel::Ok => OK_KEY,
+            ResultLabel::Error => ERROR_KEY,
+        }
+    }
+}
+
+#[cfg(feature = "prometheus-client")]
+impl EncodeLabelValue for ResultLabel {
+    fn encode(&self, encoder: &mut LabelValueEncoder) -> Result<(), std::fmt::Error> {
+        match self {
+            ResultLabel::Ok => EncodeLabelValue::encode(&OK_KEY, encoder),
+            ResultLabel::Error => EncodeLabelValue::encode(&ERROR_KEY, encoder),
+        }
+    }
 }
 
 impl CounterLabels {
@@ -46,21 +87,33 @@ impl CounterLabels {
         result: Option<ResultAndReturnTypeLabels>,
         objective: Option<Objective>,
     ) -> Self {
-        let objective = if let Some(objective) = objective {
+        let (objective_name, objective_percentile) = if let Some(objective) = objective {
             if let Some(success_rate) = objective.success_rate {
-                Some((objective.name, success_rate))
+                (Some(objective.name), Some(success_rate))
             } else {
-                None
+                (None, None)
             }
         } else {
-            None
+            (None, None)
+        };
+        let (result, ok, error) = if let Some((result, return_value_type)) = result {
+            match result {
+                OK_KEY => (Some(ResultLabel::Ok), return_value_type, None),
+                ERROR_KEY => (Some(ResultLabel::Error), None, return_value_type),
+                _ => (None, None, None),
+            }
+        } else {
+            (None, None, None)
         };
         Self {
             function,
             module,
             caller,
+            objective_name,
+            objective_percentile,
             result,
-            objective,
+            ok,
+            error,
         }
     }
 
@@ -70,62 +123,86 @@ impl CounterLabels {
             (MODULE_KEY, self.module),
             (CALLER_KEY, self.caller),
         ];
-        if let Some((result, return_value_type)) = self.result {
-            labels.push((RESULT_KEY, result));
-            if let Some(return_value_type) = return_value_type {
-                labels.push((result, return_value_type));
-            }
+        if let Some(result) = &self.result {
+            labels.push((RESULT_KEY, result.as_str()));
+        }
+        if let Some(ok) = self.ok {
+            labels.push((OK_KEY, ok));
         }
-        if let Some((name, percentile)) = &self.objective {
-            labels.push((OBJECTIVE_NAME, name));
-            labels.push((OBJECTIVE_PERCENTILE, percentile.as_str()));
+        if let Some(error) = self.error {
+            labels.push((ERROR_KEY, error));
+        }
+        if let Some(objective_name) = self.objective_name {
+            labels.push((OBJECTIVE_NAME, objective_name));
+        }
+        if let Some(objective_percentile) = &self.objective_percentile {
+            labels.push((OBJECTIVE_PERCENTILE, objective_percentile.as_str()));
         }
 
         labels
     }
 }
 
 /// These are the labels used for the `function.calls.duration` metric.
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(EncodeLabelSet, Debug, Clone, PartialEq, Eq, Hash)
+)]
 pub struct HistogramLabels {
     pub function: &'static str,
     pub module: &'static str,
-    /// The SLO name, objective percentile, and latency threshold
-    pub objective: Option<(&'static str, ObjectivePercentile, ObjectiveLatency)>,
+    pub objective_name: Option<&'static str>,
+    pub objective_percentile: Option<ObjectivePercentile>,
+    pub objective_latency_threshold: Option<ObjectiveLatency>,
 }
 
 impl HistogramLabels {
     pub fn new(function: &'static str, module: &'static str, objective: Option<Objective>) -> Self {
-        let objective = if let Some(objective) = objective {
-            if let Some((latency, percentile)) = objective.latency {
-                Some((objective.name, percentile, latency))
+        let (objective_name, objective_percentile, objective_latency_threshold) =
+            if let Some(objective) = objective {
+                if let Some((latency, percentile)) = objective.latency {
+                    (Some(objective.name), Some(percentile), Some(latency))
+                } else {
+                    (None, None, None)
+                }
             } else {
-                None
-            }
-        } else {
-            None
-        };
+                (None, None, None)
+            };
 
         Self {
             function,
             module,
-            objective,
+            objective_name,
+            objective_percentile,
+            objective_latency_threshold,
         }
     }
 
     pub fn to_vec(&self) -> Vec<Label> {
         let mut labels = vec![(FUNCTION_KEY, self.function), (MODULE_KEY, self.module)];
 
-        if let Some((name, percentile, latency)) = &self.objective {
-            labels.push((OBJECTIVE_NAME, name));
-            labels.push((OBJECTIVE_PERCENTILE, percentile.as_str()));
-            labels.push((OBJECTIVE_LATENCY_THRESHOLD, latency.as_str()));
+        if let Some(objective_name) = self.objective_name {
+            labels.push((OBJECTIVE_NAME, objective_name));
+        }
+        if let Some(objective_percentile) = &self.objective_percentile {
+            labels.push((OBJECTIVE_PERCENTILE, objective_percentile.as_str()));
+        }
+        if let Some(objective_latency_threshold) = &self.objective_latency_threshold {
+            labels.push((
+                OBJECTIVE_LATENCY_THRESHOLD,
+                objective_latency_threshold.as_str(),
+            ));
         }
 
         labels
     }
 }
 
 /// These are the labels used for the `function.calls.concurrent` metric.
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(EncodeLabelSet, Debug, Clone, PartialEq, Eq, Hash)
+)]
 pub struct GaugeLabels {
     pub function: &'static str,
     pub module: &'static str,
@@ -235,11 +312,12 @@ impl_trait_for_types!(GetStaticStr);
 /// The macro uses the autoref specialization trick through spez to get the labels for the type in a variety of circumstances.
 /// Specifically, if the value is a Result, it will add the ok or error label accordingly unless one or both of the types that
 /// the Result<T, E> is generic over implements the GetLabels trait. The label allows to override the inferred label, and the
-/// [`ResultLabels`](crate::result_labels) macro implements the GetLabels trait for the user using annotations.
+/// [`ResultLabels`](crate::ResultLabels) macro implements the GetLabels trait for the user using annotations.
 ///
 /// The macro is meant to be called with a reference as argument: `get_result_labels_for_value(&return_value)`
 ///
-/// Ref: https://github.com/dtolnay/case-studies/blob/master/autoref-specialization/README.md
+/// See: <https://github.com/dtolnay/case-studies/blob/master/autoref-specialization/README.md>
+#[doc(hidden)]
 #[macro_export]
 macro_rules! get_result_labels_for_value {
     ($e:expr) => {{

autometrics/src/lib.rs

@@ -186,9 +186,16 @@ pub use autometrics_macros::ResultLabels;
 #[cfg(feature = "prometheus-exporter")]
 pub use self::prometheus_exporter::*;
 
+#[cfg(feature = "prometheus-client")]
+pub use self::tracker::prometheus_client::PROMETHEUS_CLIENT_REGISTRY;
+
 /// We use the histogram buckets recommended by the OpenTelemetry specification
 /// https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/metrics/sdk.md#explicit-bucket-histogram-aggregation
-#[cfg(any(feature = "prometheus", feature = "prometheus-exporter"))]
+#[cfg(any(
+    feature = "prometheus",
+    feature = "prometheus-client",
+    feature = "prometheus-exporter"
+))]
 pub(crate) const HISTOGRAM_BUCKETS: [f64; 14] = [
     0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0,
 ];

autometrics/src/objectives.rs

@@ -24,6 +24,9 @@
 //! }
 //! ```
 
+#[cfg(feature = "prometheus-client")]
+use prometheus_client::encoding::{EncodeLabelValue, LabelValueEncoder};
+
 /// A Service-Level Objective (SLO) for a function or group of functions.
 ///
 /// The objective should be given a descriptive name and can represent
@@ -119,6 +122,10 @@ impl Objective {
 }
 
 /// The percentage of requests that must meet the given criteria (success rate or latency).
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(Clone, Debug, PartialEq, Eq, Hash)
+)]
 #[non_exhaustive]
 pub enum ObjectivePercentile {
     /// 90%
@@ -152,7 +159,18 @@ impl ObjectivePercentile {
     }
 }
 
+#[cfg(feature = "prometheus-client")]
+impl EncodeLabelValue for ObjectivePercentile {
+    fn encode(&self, encoder: &mut LabelValueEncoder) -> Result<(), std::fmt::Error> {
+        self.as_str().encode(encoder)
+    }
+}
+
 /// The latency threshold, in milliseoncds, for a given objective.
+#[cfg_attr(
+    feature = "prometheus-client",
+    derive(Clone, Debug, PartialEq, Eq, Hash)
+)]
 #[non_exhaustive]
 pub enum ObjectiveLatency {
     /// 5 milliseconds
@@ -234,3 +252,10 @@ impl ObjectiveLatency {
         }
     }
 }
+
+#[cfg(feature = "prometheus-client")]
+impl EncodeLabelValue for ObjectiveLatency {
+    fn encode(&self, encoder: &mut LabelValueEncoder) -> Result<(), std::fmt::Error> {
+        self.as_str().encode(encoder)
+    }
+}

autometrics/src/prometheus_exporter.rs

@@ -30,6 +30,21 @@ impl GlobalPrometheus {
             output.push_str(&self.handle.render());
         }
 
+        #[cfg(feature = "prometheus-client")]
+        {
+            output.push('\n');
+            prometheus_client::encoding::text::encode(
+                &mut output,
+                &crate::PROMETHEUS_CLIENT_REGISTRY,
+            )
+            .map_err(|err| {
+                Error::Msg(format!(
+                    "Failed to encode prometheus-client metrics: {}",
+                    err
+                ))
+            })?;
+        }
+
         Ok(output)
     }
 }

autometrics/src/tracker/metrics.rs

@@ -42,7 +42,7 @@ impl TrackMetrics for MetricsTracker {
         }
     }
 
-    fn finish<'a>(self, counter_labels: &CounterLabels, histogram_labels: &HistogramLabels) {
+    fn finish(self, counter_labels: &CounterLabels, histogram_labels: &HistogramLabels) {
         let duration = self.start.elapsed().as_secs_f64();
         register_counter!(COUNTER_NAME, &counter_labels.to_vec()).increment(1);
         register_histogram!(HISTOGRAM_NAME, &histogram_labels.to_vec()).record(duration);

autometrics/src/tracker/mod.rs

@@ -6,6 +6,8 @@ mod metrics;
 mod opentelemetry;
 #[cfg(feature = "prometheus")]
 mod prometheus;
+#[cfg(feature = "prometheus-client")]
+pub(crate) mod prometheus_client;
 
 // By default, use the opentelemetry crate
 #[cfg(all(