feat: add retain() to Family to allow metrics filtering

Signed-off-by: John Howard <john.howard@solo.io>

Author: howardjohn

Cargo.toml

@@ -44,6 +44,9 @@ http-body-util = "0.1.1"
 [build-dependencies]
 prost-build = { version = "0.12.0", optional = true }
 
+[[example]]
+name = "prune"
+
 [[bench]]
 name = "baseline"
 harness = false

examples/prune.rs

@@ -0,0 +1,93 @@
+use prometheus_client::encoding::{text::encode, EncodeMetric, MetricEncoder};
+use prometheus_client::metrics::counter::Atomic;
+use prometheus_client::metrics::family::Family;
+use prometheus_client::metrics::{MetricType, TypedMetric};
+use prometheus_client::registry::Registry;
+use prometheus_client_derive_encode::EncodeLabelSet;
+use std::fmt::Error;
+use std::sync::atomic::AtomicU64;
+use std::sync::{Arc, Mutex};
+use std::thread;
+use std::time::Duration;
+use tokio::time::Instant;
+
+// The 'prune' example shows an advanced use case of a custom metric type that records it's last record date.
+// Then, label sets older than a certain time period are removed from the label set.
+fn main() {
+    let mut registry = Registry::default();
+
+    let metric: Family<Labels, MyCounter> = Family::default();
+    registry.register("my_custom_metric", "test", metric.clone());
+    // First we record two label sets, apple and banana.
+    metric
+        .get_or_create(&Labels {
+            name: "apple".to_string(),
+        })
+        .inc();
+    metric
+        .get_or_create(&Labels {
+            name: "banana".to_string(),
+        })
+        .inc();
+
+    let mut encoded = String::new();
+    encode(&mut encoded, &registry).unwrap();
+
+    println!("Scrape output:\n{}", encoded);
+    thread::sleep(Duration::from_secs(1));
+
+    // We update only 'banana'
+    metric
+        .get_or_create(&Labels {
+            name: "banana".to_string(),
+        })
+        .inc();
+    let now = Instant::now();
+    // Retain only metrics set within the last second.
+    metric.retain(|_labels, counter| {
+        let last = counter.last_access.lock().unwrap().unwrap();
+        now.saturating_duration_since(last) < Duration::from_secs(1)
+    });
+
+    // 'apple' should be removed now
+    let mut encoded = String::new();
+    encode(&mut encoded, &registry).unwrap();
+
+    println!("Scrape output:\n{}", encoded);
+}
+
+#[derive(Default, Debug)]
+struct MyCounter {
+    value: Arc<AtomicU64>,
+    last_access: Arc<Mutex<Option<Instant>>>,
+}
+
+impl TypedMetric for MyCounter {
+    const TYPE: MetricType = MetricType::Counter;
+}
+
+impl MyCounter {
+    pub fn get(&self) -> u64 {
+        self.value.get()
+    }
+    pub fn inc(&self) -> u64 {
+        let mut last = self.last_access.lock().unwrap();
+        *last = Some(Instant::now());
+        self.value.inc()
+    }
+}
+
+impl EncodeMetric for MyCounter {
+    fn encode(&self, mut encoder: MetricEncoder) -> Result<(), Error> {
+        encoder.encode_counter::<prometheus_client::encoding::NoLabelSet, _, u64>(&self.get(), None)
+    }
+
+    fn metric_type(&self) -> MetricType {
+        todo!()
+    }
+}
+
+#[derive(Clone, Hash, Default, Debug, PartialEq, Eq, EncodeLabelSet)]
+struct Labels {
+    name: String,
+}

src/metrics/family.rs

@@ -303,6 +303,33 @@ impl<S: Clone + std::hash::Hash + Eq, M, C: MetricConstructor<M>> Family<S, M, C
         self.metrics.write().clear()
     }
 
+    /// Retains only the label sets specified by the predicate.
+    ///
+    /// ```
+    /// # use prometheus_client::metrics::counter::{Atomic, Counter};
+    /// # use prometheus_client::metrics::family::Family;
+    /// #
+    /// let family = Family::<Vec<(String, String)>, Counter>::default();
+    ///
+    /// // Will create the metric with label `method="GET"` on first call and
+    /// // return a reference.
+    /// family.get_or_create(&vec![("method".to_owned(), "GET".to_owned())]).inc_by(10);
+    /// family.get_or_create(&vec![("method".to_owned(), "DELETE".to_owned())]).inc();
+    /// family.get_or_create(&vec![("method".to_owned(), "POST".to_owned())]).inc();
+    ///
+    /// // Retain only label sets where the counter is less than 10, or method="POST"
+    /// // This will leave the method="POST" and method="DELETE" label sets.
+    /// family.retain(|labels, counter| {
+    ///     counter.get() < 5 || labels.contains(&("method".to_owned(), "POST".to_owned()))
+    /// });
+    /// ```
+    pub fn retain<F>(&self, f: F)
+    where
+        F: FnMut(&S, &mut M) -> bool,
+    {
+        self.metrics.write().retain(f)
+    }
+
     pub(crate) fn read(&self) -> RwLockReadGuard<HashMap<S, M>> {
         self.metrics.read()
     }