diff --git a/magic-nix-cache/src/api.rs b/magic-nix-cache/src/api.rs index 02e4bc3..0ebc04b 100644 --- a/magic-nix-cache/src/api.rs +++ b/magic-nix-cache/src/api.rs @@ -2,6 +2,7 @@ //! //! This API is intended to be used by nix-installer-action. +use attic::nix_store::StorePath; use axum::{extract::Extension, routing::post, Json, Router}; use axum_macros::debug_handler; use serde::{Deserialize, Serialize}; @@ -10,28 +11,41 @@ use super::State; use crate::error::{Error, Result}; #[derive(Debug, Clone, Serialize)] -struct WorkflowStartResponse {} +struct WorkflowStartResponse { + num_original_paths: usize, +} #[derive(Debug, Clone, Serialize)] struct WorkflowFinishResponse { - //num_new_paths: usize, + num_original_paths: usize, + num_final_paths: usize, + num_new_paths: usize, } pub fn get_router() -> Router { Router::new() .route("/api/workflow-start", post(workflow_start)) .route("/api/workflow-finish", post(workflow_finish)) - .route("/api/enqueue-paths", post(enqueue_paths)) + .route("/api/enqueue-paths", post(post_enqueue_paths)) } /// Record existing paths. #[debug_handler] -async fn workflow_start( - Extension(_state): Extension, -) -> Result> { +async fn workflow_start(Extension(state): Extension) -> Result> { tracing::info!("Workflow started"); + let mut original_paths = state.original_paths.lock().await; + *original_paths = crate::util::get_store_paths(&state.store).await?; - Ok(Json(WorkflowStartResponse {})) + let reply = WorkflowStartResponse { + num_original_paths: original_paths.len(), + }; + + state + .metrics + .num_original_paths + .set(reply.num_original_paths); + + Ok(Json(reply)) } /// Push new paths and shut down. @@ -40,6 +54,23 @@ async fn workflow_finish( ) -> Result> { tracing::info!("Workflow finished"); + let original_paths = state.original_paths.lock().await; + let final_paths = crate::util::get_store_paths(&state.store).await?; + let new_paths = final_paths + .difference(&original_paths) + .cloned() + .map(|path| state.store.follow_store_path(path).map_err(Error::Attic)) + .collect::>>()?; + + let num_original_paths = original_paths.len(); + let num_final_paths = final_paths.len(); + let num_new_paths = new_paths.len(); + + // NOTE(cole-h): If we're substituting from an upstream cache, those paths won't have the + // post-build-hook run on it, so we diff the store to ensure we cache everything we can. + tracing::info!("Diffing the store and uploading any new paths before we shut down"); + enqueue_paths(&state, new_paths).await?; + if let Some(gha_cache) = &state.gha_cache { tracing::info!("Waiting for GitHub action cache uploads to finish"); gha_cache.shutdown().await?; @@ -63,9 +94,18 @@ async fn workflow_finish( println!("Every log line throughout the lifetime of the program:"); println!("\n{logfile_contents}\n"); - let reply = WorkflowFinishResponse {}; + let reply = WorkflowFinishResponse { + num_original_paths, + num_final_paths, + num_new_paths, + }; - //state.metrics.num_new_paths.set(num_new_paths); + state + .metrics + .num_original_paths + .set(reply.num_original_paths); + state.metrics.num_final_paths.set(reply.num_final_paths); + state.metrics.num_new_paths.set(reply.num_new_paths); Ok(Json(reply)) } @@ -80,7 +120,7 @@ pub struct EnqueuePathsResponse {} /// Schedule paths in the local Nix store for uploading. #[tracing::instrument(skip_all)] -async fn enqueue_paths( +async fn post_enqueue_paths( Extension(state): Extension, Json(req): Json, ) -> Result> { @@ -92,6 +132,12 @@ async fn enqueue_paths( .map(|path| state.store.follow_store_path(path).map_err(Error::Attic)) .collect::>>()?; + enqueue_paths(&state, store_paths).await?; + + Ok(Json(EnqueuePathsResponse {})) +} + +async fn enqueue_paths(state: &State, store_paths: Vec) -> Result<()> { if let Some(gha_cache) = &state.gha_cache { gha_cache .enqueue_paths(state.store.clone(), store_paths.clone()) @@ -103,5 +149,5 @@ async fn enqueue_paths( crate::flakehub::enqueue_paths(flakehub_state, store_paths).await?; } - Ok(Json(EnqueuePathsResponse {})) + Ok(()) } diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs index 9f94e54..2113606 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -18,6 +18,7 @@ mod error; mod flakehub; mod gha; mod telemetry; +mod util; use std::collections::HashSet; use std::fs::{self, create_dir_all}; @@ -136,6 +137,9 @@ struct StateInner { /// FlakeHub cache state. flakehub_state: RwLock>, + + /// The paths in the Nix store when Magic Nix Cache started. + original_paths: Mutex>, } async fn main_cli() -> Result<()> { @@ -324,6 +328,7 @@ async fn main_cli() -> Result<()> { metrics, store, flakehub_state: RwLock::new(flakehub_state), + original_paths: Mutex::new(HashSet::new()), }); let app = Router::new() diff --git a/magic-nix-cache/src/util.rs b/magic-nix-cache/src/util.rs new file mode 100644 index 0000000..e22d531 --- /dev/null +++ b/magic-nix-cache/src/util.rs @@ -0,0 +1,47 @@ +//! Utilities. + +use std::collections::HashSet; +use std::path::{Path, PathBuf}; + +use attic::nix_store::NixStore; + +use crate::error::Result; + +/// Returns the list of store paths that are currently present. +pub async fn get_store_paths(store: &NixStore) -> Result> { + // FIXME: use the Nix API. + let store_dir = store.store_dir(); + let mut listing = tokio::fs::read_dir(store_dir).await?; + let mut paths = HashSet::new(); + while let Some(entry) = listing.next_entry().await? { + let file_name = entry.file_name(); + let file_name = Path::new(&file_name); + + if let Some(extension) = file_name.extension() { + match extension.to_str() { + None | Some("drv") | Some("chroot") => { + tracing::debug!( + "skipping file with weird or uninteresting extension {extension:?}" + ); + continue; + } + _ => {} + } + } + + if let Some(s) = file_name.to_str() { + // Let's not push any sources + if s.ends_with("-source") { + continue; + } + + // Special paths (so far only `.links`) + if s.starts_with(".links") { + continue; + } + } + + paths.insert(store_dir.join(file_name)); + } + Ok(paths) +}