Write stats to cache_stats_file after compile

Signed-off-by: Zack Cerza <[email protected]>

Author: zmc

docs/Configuration.md

@@ -6,6 +6,9 @@
 # If specified, wait this long for the server to start up.
 server_startup_timeout_ms = 10000
 
+# If specified, write JSON-formatted stats to this file after each compile operation.
+cache_stats_file = "/home/user/.cache/sccache-stats.json"
+
 [dist]
 # where to find the scheduler
 scheduler_url = "http://1.2.3.4:10600"

src/config.rs

@@ -585,6 +585,7 @@ pub struct FileConfig {
     pub cache: CacheConfigs,
     pub dist: DistConfig,
     pub server_startup_timeout_ms: Option<u64>,
+    pub cache_stats_file: Option<PathBuf>,
 }
 
 // If the file doesn't exist or we can't read it, log the issue and proceed. If the
@@ -977,6 +978,7 @@ pub struct Config {
     pub fallback_cache: DiskCacheConfig,
     pub dist: DistConfig,
     pub server_startup_timeout: Option<std::time::Duration>,
+    pub cache_stats_file: Option<PathBuf>,
 }
 
 impl Config {
@@ -998,6 +1000,7 @@ impl Config {
             cache,
             dist,
             server_startup_timeout_ms,
+            cache_stats_file,
         } = file_conf;
         conf_caches.merge(cache);
 
@@ -1013,6 +1016,7 @@ impl Config {
             fallback_cache,
             dist,
             server_startup_timeout,
+            cache_stats_file,
         }
     }
 }
@@ -1312,6 +1316,7 @@ fn config_overrides() {
         },
         dist: Default::default(),
         server_startup_timeout_ms: None,
+        cache_stats_file: None,
     };
 
     assert_eq!(
@@ -1334,6 +1339,7 @@ fn config_overrides() {
             },
             dist: Default::default(),
             server_startup_timeout: None,
+            cache_stats_file: None,
         }
     );
 }
@@ -1463,6 +1469,7 @@ fn test_gcs_service_account() {
 fn full_toml_parse() {
     const CONFIG_STR: &str = r#"
 server_startup_timeout_ms = 10000
+cache_stats_file = "/home/user/.cache/sccache-stats.json"
 
 [dist]
 # where to find the scheduler
@@ -1623,6 +1630,7 @@ no_credentials = true
                 rewrite_includes_only: false,
             },
             server_startup_timeout_ms: Some(10000),
+            cache_stats_file: Some(PathBuf::from("/home/user/.cache/sccache-stats.json"))
         }
     )
 }
@@ -1715,6 +1723,7 @@ size = "7g"
                 ..Default::default()
             },
             server_startup_timeout_ms: None,
+            cache_stats_file: None,
         }
     );
 }

src/server.rs

@@ -31,6 +31,7 @@ use anyhow::Context as _;
 use bytes::{buf::BufMut, Bytes, BytesMut};
 use filetime::FileTime;
 use fs::metadata;
+use fs::File;
 use fs_err as fs;
 use futures::channel::mpsc;
 use futures::future::FutureExt;
@@ -42,15 +43,15 @@ use std::collections::{HashMap, HashSet};
 use std::env;
 use std::ffi::OsString;
 use std::future::Future;
-use std::io::{self, Write};
+use std::io::{self, BufWriter, Write};
 use std::marker::Unpin;
 #[cfg(feature = "dist-client")]
 use std::mem;
 #[cfg(target_os = "android")]
 use std::os::android::net::SocketAddrExt;
 #[cfg(target_os = "linux")]
 use std::os::linux::net::SocketAddrExt;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::process::{ExitStatus, Output};
 use std::sync::Arc;
@@ -488,8 +489,14 @@ pub fn start_server(config: &Config, addr: &crate::net::SocketAddr) -> Result<()
             crate::net::SocketAddr::Net(addr) => {
                 trace!("binding TCP {addr}");
                 let l = runtime.block_on(tokio::net::TcpListener::bind(addr))?;
-                let srv =
-                    SccacheServer::<_>::with_listener(l, runtime, client, dist_client, storage);
+                let srv = SccacheServer::<_>::with_listener(
+                    l,
+                    runtime,
+                    client,
+                    dist_client,
+                    storage,
+                    config.cache_stats_file.clone(),
+                );
                 Ok((
                     srv.local_addr().unwrap(),
                     Box::new(move |f| srv.run(f)) as Box<dyn FnOnce(_) -> _>,
@@ -504,8 +511,14 @@ pub fn start_server(config: &Config, addr: &crate::net::SocketAddr) -> Result<()
                     let _guard = runtime.enter();
                     tokio::net::UnixListener::bind(path)?
                 };
-                let srv =
-                    SccacheServer::<_>::with_listener(l, runtime, client, dist_client, storage);
+                let srv = SccacheServer::<_>::with_listener(
+                    l,
+                    runtime,
+                    client,
+                    dist_client,
+                    storage,
+                    config.cache_stats_file.clone(),
+                );
                 Ok((
                     srv.local_addr().unwrap(),
                     Box::new(move |f| srv.run(f)) as Box<dyn FnOnce(_) -> _>,
@@ -579,6 +592,7 @@ impl<C: CommandCreatorSync> SccacheServer<tokio::net::TcpListener, C> {
         client: Client,
         dist_client: DistClientContainer,
         storage: Arc<dyn Storage>,
+        cache_stats_file: Option<PathBuf>,
     ) -> Result<Self> {
         let addr = crate::net::SocketAddr::with_port(port);
         let listener = runtime.block_on(tokio::net::TcpListener::bind(addr.as_net().unwrap()))?;
@@ -589,6 +603,7 @@ impl<C: CommandCreatorSync> SccacheServer<tokio::net::TcpListener, C> {
             client,
             dist_client,
             storage,
+            cache_stats_file,
         ))
     }
 }
@@ -600,13 +615,22 @@ impl<A: crate::net::Acceptor, C: CommandCreatorSync> SccacheServer<A, C> {
         client: Client,
         dist_client: DistClientContainer,
         storage: Arc<dyn Storage>,
+        cache_stats_file: Option<PathBuf>,
     ) -> Self {
         // Prepare the service which we'll use to service all incoming TCP
         // connections.
         let (tx, rx) = mpsc::channel(1);
         let (wait, info) = WaitUntilZero::new();
         let pool = runtime.handle().clone();
-        let service = SccacheService::new(dist_client, storage, &client, pool, tx, info);
+        let service = SccacheService::new(
+            dist_client,
+            storage,
+            &client,
+            pool,
+            tx,
+            info,
+            cache_stats_file.clone(),
+        );
 
         SccacheServer {
             runtime,
@@ -818,6 +842,10 @@ where
     /// This field causes [WaitUntilZero] to wait until this struct drops.
     #[allow(dead_code)]
     info: ActiveInfo,
+
+    /// A file that will contain JSON-formatted stats output, written after
+    /// each compile operation.
+    cache_stats_file: Option<PathBuf>,
 }
 
 type SccacheRequest = Message<Request, Body<()>>;
@@ -857,7 +885,11 @@ where
                 Request::Compile(compile) => {
                     debug!("handle_client: compile");
                     me.stats.lock().await.compile_requests += 1;
-                    me.handle_compile(compile).await
+                    let resp = me.handle_compile(compile).await;
+                    if let Some(val) = &me.cache_stats_file {
+                        me.stats.lock().await.clone().write(val)?
+                    }
+                    resp
                 }
                 Request::GetStats => {
                     debug!("handle_client: get_stats");
@@ -916,6 +948,7 @@ where
         rt: tokio::runtime::Handle,
         tx: mpsc::Sender<ServerMessage>,
         info: ActiveInfo,
+        cache_stats_file: Option<PathBuf>,
     ) -> SccacheService<C> {
         SccacheService {
             stats: Arc::default(),
@@ -927,6 +960,7 @@ where
             creator: C::new(client),
             tx,
             info,
+            cache_stats_file,
         }
     }
 
@@ -938,6 +972,7 @@ where
         let (_, info) = WaitUntilZero::new();
         let client = Client::new_num(1);
         let dist_client = DistClientContainer::new_disabled();
+        let cache_stats_file = None;
         SccacheService {
             stats: Arc::default(),
             dist_client: Arc::new(dist_client),
@@ -948,6 +983,7 @@ where
             creator: C::new(&client),
             tx,
             info,
+            cache_stats_file,
         }
     }
 
@@ -960,6 +996,7 @@ where
         let (tx, _) = mpsc::channel(1);
         let (_, info) = WaitUntilZero::new();
         let client = Client::new_num(1);
+        let cache_stats_file = None;
         SccacheService {
             stats: Arc::default(),
             dist_client: Arc::new(DistClientContainer::new_with_state(DistClientState::Some(
@@ -981,6 +1018,7 @@ where
             creator: C::new(&client),
             tx,
             info,
+            cache_stats_file,
         }
     }
 
@@ -1909,6 +1947,19 @@ impl ServerStats {
             );
         }
     }
+
+    /// Write stats in JSON format to a file.
+    fn write(&self, path: &Path) -> Result<()> {
+        let file = match File::create(path) {
+            Ok(f) => f,
+            Err(e) => {
+                debug!("Couldn't open stats file for writing: {}", e);
+                return Ok(());
+            }
+        };
+        let mut writer = BufWriter::new(file);
+        Ok(serde_json::to_writer(&mut writer, self)?)
+    }
 }
 
 fn set_percentage_stat(

src/test/tests.rs

@@ -89,7 +89,7 @@ where
         ));
 
         let client = Client::new();
-        let srv = SccacheServer::new(0, runtime, client, dist_client, storage).unwrap();
+        let srv = SccacheServer::new(0, runtime, client, dist_client, storage, None).unwrap();
         let mut srv: SccacheServer<_, Arc<Mutex<MockCommandCreator>>> = srv;
         let addr = srv.local_addr().unwrap();
         assert!(matches!(addr, crate::net::SocketAddr::Net(a) if a.port() > 0));

tests/harness/mod.rs

@@ -189,6 +189,7 @@ pub fn sccache_client_cfg(
             rewrite_includes_only: false, // TODO
         },
         server_startup_timeout_ms: None,
+        cache_stats_file: None,
     }
 }
 

tests/oauth.rs

@@ -60,6 +60,7 @@ fn config_with_dist_auth(
             rewrite_includes_only: true,
         },
         server_startup_timeout_ms: None,
+        cache_stats_file: None,
     }
 }