From cb64acedd3f77a0098c42a35cdf5cf189f780ac3 Mon Sep 17 00:00:00 2001
From: John Howard <john.howard@solo.io>
Date: Mon, 29 Apr 2024 13:09:57 -0700
Subject: [PATCH] Prototype of metrics expiry

Signed-off-by: John Howard <john.howard@solo.io>
---
 Cargo.toml            |  3 ++
 examples/prune.rs     | 84 +++++++++++++++++++++++++++++++++++++++++++
 src/metrics/family.rs |  8 +++++
 3 files changed, 95 insertions(+)
 create mode 100644 examples/prune.rs

diff --git a/Cargo.toml b/Cargo.toml
index 5514ab68..7f23ca0b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -43,6 +43,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
diff --git a/examples/prune.rs b/examples/prune.rs
new file mode 100644
index 00000000..f6980ba7
--- /dev/null
+++ b/examples/prune.rs
@@ -0,0 +1,84 @@
+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;
+
+#[derive(Default, Debug)]
+struct MyCounter {
+    value: Arc<AtomicU64>,
+    last_access: Arc<Mutex<Option<Instant>>>,
+}
+
+impl TypedMetric for MyCounter {}
+
+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::<(), _, u64>(&self.get(), None)
+    }
+
+    fn metric_type(&self) -> MetricType {
+        todo!()
+    }
+}
+
+#[derive(Clone, Hash, Default, Debug, PartialEq, Eq, EncodeLabelSet)]
+struct Labels {
+    name: String,
+}
+
+fn main() {
+    let mut registry = Registry::default();
+
+    let metric: Family<Labels, MyCounter> = Family::default();
+    registry.register("my_custom_metric", "test", metric.clone());
+    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));
+    metric
+        .get_or_create(&Labels {
+            name: "banana".to_string(),
+        })
+        .inc();
+    let now = Instant::now();
+    metric.retain(|a, b| {
+        let last = b.last_access.lock().unwrap().unwrap();
+        now.saturating_duration_since(last) > Duration::from_secs(1)
+    });
+
+    let mut encoded = String::new();
+    encode(&mut encoded, &registry).unwrap();
+
+    println!("Scrape output:\n{}", encoded);
+}
diff --git a/src/metrics/family.rs b/src/metrics/family.rs
index 2f23b198..6e6bb055 100644
--- a/src/metrics/family.rs
+++ b/src/metrics/family.rs
@@ -288,6 +288,14 @@ impl<S: Clone + std::hash::Hash + Eq, M, C: MetricConstructor<M>> Family<S, M, C
         self.metrics.write().clear()
     }
 
+    /// retain
+    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()
     }