From 59b59cfa3a7c7839ae247021222aa628a4ade9f0 Mon Sep 17 00:00:00 2001
From: Graham Christensen <graham@grahamc.com>
Date: Sun, 25 Jun 2023 21:45:20 -0400
Subject: [PATCH] Record some basic data about gets and puts

---
 magic-nix-cache/src/binary_cache.rs |  9 +++++++++
 magic-nix-cache/src/main.rs         | 28 ++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/magic-nix-cache/src/binary_cache.rs b/magic-nix-cache/src/binary_cache.rs
index 60cb958..5283684 100644
--- a/magic-nix-cache/src/binary_cache.rs
+++ b/magic-nix-cache/src/binary_cache.rs
@@ -56,16 +56,21 @@ async fn get_narinfo(
         .await
         .contains(&store_path_hash)
     {
+        state.metrics.narinfos_sent_upstream.incr();
+        state.metrics.narinfos_negative_cache_hits.incr();
         return pull_through(&state, &path);
     }
 
     if let Some(url) = state.api.get_file_url(&[&key]).await? {
+        state.metrics.narinfos_served.incr();
         return Ok(Redirect::temporary(&url));
     }
 
     let mut negative_cache = state.narinfo_nagative_cache.write().await;
     negative_cache.insert(store_path_hash);
 
+    state.metrics.narinfos_sent_upstream.incr();
+    state.metrics.narinfos_negative_cache_misses.incr();
     pull_through(&state, &path)
 }
 async fn put_narinfo(
@@ -90,6 +95,7 @@ async fn put_narinfo(
         body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))),
     );
     state.api.upload_file(allocation, stream).await?;
+    state.metrics.narinfos_uploaded.incr();
 
     state
         .narinfo_nagative_cache
@@ -102,10 +108,12 @@ async fn put_narinfo(
 
 async fn get_nar(Extension(state): Extension<State>, Path(path): Path<String>) -> Result<Redirect> {
     if let Some(url) = state.api.get_file_url(&[&path]).await? {
+        state.metrics.nars_served.incr();
         return Ok(Redirect::temporary(&url));
     }
 
     if let Some(upstream) = &state.upstream {
+        state.metrics.nars_sent_upstream.incr();
         Ok(Redirect::temporary(&format!("{}/nar/{}", upstream, path)))
     } else {
         Err(Error::NotFound)
@@ -121,6 +129,7 @@ async fn put_nar(
         body.map(|r| r.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))),
     );
     state.api.upload_file(allocation, stream).await?;
+    state.metrics.nars_uploaded.incr();
 
     Ok(())
 }
diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs
index 9b215fa..556c3b9 100644
--- a/magic-nix-cache/src/main.rs
+++ b/magic-nix-cache/src/main.rs
@@ -107,6 +107,31 @@ struct StateInner {
     ///
     /// This is used by our Action API to invoke `nix copy` to upload new paths.
     self_endpoint: SocketAddr,
+
+    /// Metrics for sending to perf at shutdown
+    metrics: Metrics,
+}
+
+/// The global server state.
+#[derive(Debug, Default)]
+struct Metrics {
+    narinfos_served: Metric,
+    narinfos_sent_upstream: Metric,
+    narinfos_negative_cache_hits: Metric,
+    narinfos_negative_cache_misses: Metric,
+    narinfos_uploaded: Metric,
+
+    nars_served: Metric,
+    nars_sent_upstream: Metric,
+    nars_uploaded: Metric,
+}
+
+#[derive(Debug, Default)]
+struct Metric(std::sync::atomic::AtomicUsize);
+impl Metric {
+    fn incr(&self) -> () {
+        self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+    }
 }
 
 fn main() {
@@ -149,6 +174,9 @@ fn main() {
         original_paths: Mutex::new(HashSet::new()),
         narinfo_nagative_cache: RwLock::new(HashSet::new()),
         self_endpoint: args.listen.to_owned(),
+        metrics: Metrics {
+            ..Default::default()
+        },
     });
 
     let app = Router::new()