Fix #144: Add sub_registry_with_labels method to Registry

prometheus · May 28, 2023 · 89f8c3c · 89f8c3c
1 parent d001ba0
commit 89f8c3c
Showing 1 changed file with 39 additions and 6 deletions.
diff --git a/src/registry.rs b/src/registry.rs
@@ -250,6 +250,23 @@ impl Registry {
         self.priv_sub_registry(sub_registry)
     }
 
+    /// Like [`Registry::sub_registry_with_prefix`] but with multiple labels instead.
+    pub fn sub_registry_with_labels<I>(&mut self, labels: I) -> &mut Self
+    where
+        I: Iterator<Item = (Cow<'static, str>, Cow<'static, str>)>,
+    {
+        let mut new_labels = self.labels.clone();
+        new_labels.extend(labels);
+
+        let sub_registry = Registry {
+            prefix: self.prefix.clone(),
+            labels: new_labels,
+            ..Default::default()
+        };
+
+        self.priv_sub_registry(sub_registry)
+    }
+
     fn priv_sub_registry(&mut self, sub_registry: Self) -> &mut Self {
         self.sub_registries.push(sub_registry);
 
@@ -549,10 +566,19 @@ mod tests {
         let sub_sub_registry = sub_registry.sub_registry_with_label(label_1_2.clone());
         sub_sub_registry.register(prefix_1_2_metric_name, "some help", counter.clone());
 
-        let prefix_1_2_1 = "prefix_1_2_1";
-        let prefix_1_2_1_metric_name = "my_prefix_1_2_1_metric";
-        let sub_sub_sub_registry = sub_sub_registry.sub_registry_with_prefix(prefix_1_2_1);
-        sub_sub_sub_registry.register(prefix_1_2_1_metric_name, "some help", counter.clone());
+        let labels_1_3 = vec![
+            (Cow::Borrowed("registry"), Cow::Borrowed("1_3_1")),
+            (Cow::Borrowed("registry"), Cow::Borrowed("1_3_2")),
+        ];
+        let prefix_1_3_metric_name = "my_prefix_1_3_metric";
+        let sub_sub_registry =
+            sub_registry.sub_registry_with_labels(labels_1_3.clone().into_iter());
+        sub_sub_registry.register(prefix_1_3_metric_name, "some help", counter.clone());
+
+        let prefix_1_3_1 = "prefix_1_3_1";
+        let prefix_1_3_1_metric_name = "my_prefix_1_3_1_metric";
+        let sub_sub_sub_registry = sub_sub_registry.sub_registry_with_prefix(prefix_1_3_1);
+        sub_sub_sub_registry.register(prefix_1_3_1_metric_name, "some help", counter.clone());
 
         let prefix_2 = "prefix_2";
         let _ = registry.sub_registry_with_prefix(prefix_2);
@@ -589,8 +615,15 @@ mod tests {
         );
         assert_eq!(
             Some((
-                prefix_1.to_string() + "_" + prefix_1_2_1 + "_" + prefix_1_2_1_metric_name,
-                vec![label_1_2]
+                prefix_1.to_string() + "_" + prefix_1_3_metric_name,
+                labels_1_3.clone()
+            )),
+            metric_iter.next()
+        );
+        assert_eq!(
+            Some((
+                prefix_1.to_string() + "_" + prefix_1_3_1 + "_" + prefix_1_3_1_metric_name,
+                labels_1_3.clone()
             )),
             metric_iter.next()
         );