diff --git a/.gitignore b/.gitignore index 86c77cc0f..d34adcacf 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,8 @@ backend/test-store /proxy/ /gradjoeng/ -/gradient-hypervisor/ +/hypervisor/ +/hosted/ # ---> Python # Byte-compiled / optimized / DLL files diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 9f9f33ecd..ede78b379 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -2705,6 +2705,7 @@ name = "gradient-entity" version = "2.0.0" dependencies = [ "chrono", + "futures", "num_enum", "sea-orm", "serde", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 1aa69e44c..de96cfbb2 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -242,6 +242,10 @@ publish.workspace = true [lints] workspace = true +[lib] +name = "gradient_server" +path = "src/lib.rs" + [[bin]] name = "gradient-server" path = "src/main.rs" diff --git a/backend/gradient-cache/src/eval_cache.rs b/backend/gradient-cache/src/eval_cache.rs index 8d8924d0e..a6bcdc878 100644 --- a/backend/gradient-cache/src/eval_cache.rs +++ b/backend/gradient-cache/src/eval_cache.rs @@ -110,8 +110,8 @@ impl EvalCacheSweep { let mut freed: u64 = 0; let mut evicted = 0usize; for row in models { - if let Err(e) = self.nar_storage.delete_eval_cache(&row.fingerprint).await { - warn!(fingerprint = %row.fingerprint, error = ?e, "eval-cache sweep: blob delete failed"); + if let Err(e) = self.nar_storage.delete_eval_cache(&row.storage_path).await { + warn!(key = %row.storage_path, error = ?e, "eval-cache sweep: blob delete failed"); } match EEvalCacheStore::delete_by_id(row.id).exec(&self.db).await { diff --git a/backend/gradient-cache/src/schedule.rs b/backend/gradient-cache/src/schedule.rs index ec9d03e69..20684b7a7 100644 --- a/backend/gradient-cache/src/schedule.rs +++ b/backend/gradient-cache/src/schedule.rs @@ -35,6 +35,7 @@ pub(crate) fn child_specs(state: &ServerState) -> Vec { ChildSpec::pass(SignSweep { db: state.worker_db.clone(), config: Arc::clone(config), + store_path_mapping: Arc::clone(&state.extensions.store_path_mapping), period: secs(config.cache.sign_sweep_interval_secs), }), ChildSpec::pass(DebugIndex { diff --git a/backend/gradient-cache/src/signatures.rs b/backend/gradient-cache/src/signatures.rs index ddeb84cc8..dce573e56 100644 --- a/backend/gradient-cache/src/signatures.rs +++ b/backend/gradient-cache/src/signatures.rs @@ -55,6 +55,7 @@ ON CONFLICT (cached_path, cache) DO NOTHING pub(crate) struct SignSweep { pub(crate) db: WorkerDb, pub(crate) config: Arc, + pub(crate) store_path_mapping: Arc, pub(crate) period: Duration, } @@ -185,8 +186,14 @@ impl SignSweep { let nar_hash_nix32 = normalize_nar_hash(nar_hash); - let sig_bytes = - signer.sign_narinfo_raw(&store_path, &nar_hash_nix32, nar_size as u64, &refs); + let mapper = StorePathMapper::for_cache(Arc::clone(&self.store_path_mapping), cache.id); + let sig_bytes = signer.sign_stored_narinfo( + &mapper, + &store_path, + &nar_hash_nix32, + nar_size as u64, + &refs, + ); let mut am = row.into_active_model(); am.signature = Set(Some(sig_bytes)); diff --git a/backend/gradient-cache/src/storage_migrations/eval_cache_per_cache.rs b/backend/gradient-cache/src/storage_migrations/eval_cache_per_cache.rs new file mode 100644 index 000000000..0d26b529f --- /dev/null +++ b/backend/gradient-cache/src/storage_migrations/eval_cache_per_cache.rs @@ -0,0 +1,80 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use anyhow::Result; +use futures::FutureExt as _; +use futures::future::BoxFuture; +use gradient_db::DbContext; +use tracing::warn; + +use super::StorageMigration; + +const FLAT_KEYS: &str = "flat-keys"; + +pub(super) struct EvalCachePerCache; + +impl StorageMigration for EvalCachePerCache { + fn name(&self) -> &'static str { + "m20261019_000000_eval_cache_per_cache" + } + + fn units(&self) -> Vec { + vec![FLAT_KEYS.to_owned()] + } + + fn migrate<'a>(&'a self, db: &'a DbContext, _unit: &'a str) -> BoxFuture<'a, Result<()>> { + delete_flat_keys(db).boxed() + } +} + +fn is_flat(key: &str) -> bool { + key.strip_prefix("eval-cache/") + .is_some_and(|rest| !rest.is_empty() && !rest.contains('/')) +} + +async fn delete_flat_keys(db: &DbContext) -> Result<()> { + let store = &db.storage.nar_storage; + for key in store.list_eval_cache_keys().await? { + if is_flat(&key) + && let Err(e) = store.delete_eval_cache(&key).await + { + warn!(%key, error = %e, "flat eval-cache blob delete failed"); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::test_server_state; + use gradient_fixtures::mock::empty_db; + use gradient_storage::NarStore; + use gradient_types::ids::CacheId; + + #[tokio::test] + async fn only_eval_cache_blobs_without_a_cache_are_deleted() { + let tmp = tempfile::tempdir().unwrap(); + let nar = NarStore::local(tmp.path()).unwrap(); + let per_cache = format!("eval-cache/{}/fp2", CacheId::now_v7()); + nar.put_eval_cache("eval-cache/fp1", b"flat".to_vec()) + .await + .unwrap(); + nar.put_eval_cache(&per_cache, b"kept".to_vec()) + .await + .unwrap(); + let state = test_server_state(nar, empty_db(), |_| {}); + + EvalCachePerCache + .migrate(&state.db(), "flat-keys") + .await + .unwrap(); + + let store = &state.db().storage.nar_storage; + assert_eq!(store.get_eval_cache("eval-cache/fp1").await.unwrap(), None); + assert!(store.get_eval_cache(&per_cache).await.unwrap().is_some()); + } +} diff --git a/backend/gradient-cache/src/storage_migrations/mod.rs b/backend/gradient-cache/src/storage_migrations/mod.rs index 540f18a54..37e1f5b50 100644 --- a/backend/gradient-cache/src/storage_migrations/mod.rs +++ b/backend/gradient-cache/src/storage_migrations/mod.rs @@ -7,6 +7,8 @@ //! Storage migrations are to the NAR, log and blob layout what database migrations are to the //! schema. Apply them in order and before any deep GC unit. +mod eval_cache_per_cache; + use crate::units::{Step, next_unit}; use anyhow::{Context, Result}; use futures::future::BoxFuture; @@ -25,7 +27,7 @@ pub(crate) trait StorageMigration: Send + Sync { } fn migrations() -> Vec> { - Vec::new() + vec![Box::new(eval_cache_per_cache::EvalCachePerCache)] } pub(super) async fn step(db: &DbContext) -> Result { diff --git a/backend/gradient-ci/src/actions/tests/fixtures.rs b/backend/gradient-ci/src/actions/tests/fixtures.rs index 514cfb068..81d0f3cdb 100644 --- a/backend/gradient-ci/src/actions/tests/fixtures.rs +++ b/backend/gradient-ci/src/actions/tests/fixtures.rs @@ -174,6 +174,7 @@ pub fn make_ctx_with(worker: sea_orm::MockDatabase) -> crate::CiContext { events: gradient_types::EventBus::default(), wake: Default::default(), memory: Default::default(), + store_path_mapping: std::sync::Arc::new(gradient_entity::IdentityMapping), }; CiContext { db, diff --git a/backend/gradient-core/src/extensions.rs b/backend/gradient-core/src/extensions.rs new file mode 100644 index 000000000..c06e471a8 --- /dev/null +++ b/backend/gradient-core/src/extensions.rs @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::sync::Arc; + +use gradient_entity::{IdentityMapping, StorePathMapping}; + +#[derive(Clone, Debug)] +pub struct ServerExtensions { + pub store_path_mapping: Arc, +} + +impl Default for ServerExtensions { + fn default() -> Self { + Self { + store_path_mapping: Arc::new(IdentityMapping), + } + } +} diff --git a/backend/gradient-core/src/lib.rs b/backend/gradient-core/src/lib.rs index 9766eedbc..cc82557ce 100644 --- a/backend/gradient-core/src/lib.rs +++ b/backend/gradient-core/src/lib.rs @@ -4,10 +4,12 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +pub mod extensions; pub mod state_root; pub mod upstream; pub mod upstream_source; +pub use extensions::ServerExtensions; pub use gradient_graph::Graph; pub use state_root::{ AppState, BUILD_PROGRESS_TTL, EVAL_PROGRESS_TTL, InstanceMemory, LAST_USED_STAMP_INTERVAL, @@ -57,7 +59,10 @@ pub enum InitError { LocalStorage(String), } -pub async fn init_state(cli: Cli) -> Result, InitError> { +pub async fn init_state( + cli: Cli, + extensions: ServerExtensions, +) -> Result, InitError> { if cli.secrets.crypt_file.is_empty() || cli.secrets.jwt_file.is_empty() { return Err(InitError::MissingSecrets); } @@ -219,6 +224,7 @@ pub async fn init_state(cli: Cli) -> Result, InitError> { events: gradient_types::EventBus::default(), wake: Default::default(), graph: Graph::new(), + extensions, memory: InstanceMemory { db: DbMemory { probe_requests: gradient_db::ProbeRequests::channel(), diff --git a/backend/gradient-core/src/state_root.rs b/backend/gradient-core/src/state_root.rs index 49ef5ccc4..4b8f943c8 100644 --- a/backend/gradient-core/src/state_root.rs +++ b/backend/gradient-core/src/state_root.rs @@ -18,6 +18,8 @@ use gradient_ci::CiContext; use gradient_ci::manifest_state::{ManifestStateStore, PendingCredentialsStore}; use gradient_db::metrics::cache_traffic::CacheTraffic; use gradient_db::{CacheDb, DbContext, DbMemory, Wakeups, WebDb, WorkerDb}; +use gradient_entity::ids::{CacheId, ProjectId}; +use gradient_entity::{MappingError, StorePathMapper}; use gradient_git_host::GitHostRegistry; use gradient_graph::Graph; use gradient_notify::EmailSender; @@ -30,6 +32,8 @@ use gradient_util::debounce::Debounce; use gradient_util::latest::Latest; use gradient_util::shutdown::Shutdown; +use crate::ServerExtensions; + #[derive(Debug)] pub struct InstanceMemory { pub db: DbMemory, @@ -66,6 +70,7 @@ pub struct AppState { pub graph: Arc, pub memory: InstanceMemory, pub wake: Wakeups, + pub extensions: ServerExtensions, } pub type ServerState = AppState; @@ -106,9 +111,21 @@ impl AppState { events: self.events.clone(), wake: self.wake.clone(), memory: self.memory.db.clone(), + store_path_mapping: Arc::clone(&self.extensions.store_path_mapping), } } + pub fn cache_mapper(&self, cache: CacheId) -> StorePathMapper { + StorePathMapper::for_cache(Arc::clone(&self.extensions.store_path_mapping), cache) + } + + pub async fn project_mapper( + &self, + project: ProjectId, + ) -> Result { + StorePathMapper::for_project(Arc::clone(&self.extensions.store_path_mapping), project).await + } + pub fn ci(&self) -> CiContext { CiContext { db: self.db(), @@ -126,3 +143,23 @@ impl AppState { self.db().record_evaluation_created(eval).await; } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use gradient_entity::ids::ProjectId; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + + #[tokio::test] + async fn the_server_mapping_reaches_the_database_context() { + let state = gradient_test_support::state::TestState::new() + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + let mapper = state.db().project_mapper(ProjectId::nil()).await.unwrap(); + assert_eq!( + mapper.outbound_text(EXAMPLE_STORED_HASH).unwrap(), + EXAMPLE_ORIGINAL_HASH + ); + } +} diff --git a/backend/gradient-core/src/upstream.rs b/backend/gradient-core/src/upstream.rs index 0ff2d737a..bfc18c0bc 100644 --- a/backend/gradient-core/src/upstream.rs +++ b/backend/gradient-core/src/upstream.rs @@ -13,7 +13,9 @@ use gradient_util::sync::Mutex; use tokio::sync::Semaphore; use gradient_db::caches::upstream::{UpstreamAccum, UpstreamEndpoint}; +use gradient_entity::StorePathError; use gradient_types::ids::CacheUpstreamId; +use gradient_types::{StorePath, StorePathMapper}; use gradient_wire::types::CachedPath; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -226,6 +228,7 @@ impl ProbeResult { async fn probe_one( pool: &Arc, ep: &UpstreamEndpoint, + mapper: &StorePathMapper, hash: &str, store_path: &str, ) -> (f64, SampleKind, Option) { @@ -253,7 +256,9 @@ async fn probe_one( let out = match resp { Ok(r) if r.status().is_success() => match r.text().await { - Ok(body) => match verified_narinfo(ep, hash, store_path, &body) { + Ok(body) => match verified_narinfo(ep, hash, store_path, &body) + .and_then(|cp| inbound_cached_path(mapper, cp).ok()) + { Some(cp) => (latency_ms, SampleKind::Hit, Some(cp)), None => (latency_ms, SampleKind::Miss, None), }, @@ -334,12 +339,30 @@ pub async fn fetch_narinfo_body( None } +pub fn inbound_cached_path( + mapper: &StorePathMapper, + cp: CachedPath, +) -> Result { + Ok(CachedPath { + path: mapper.inbound_text(&cp.path)?, + references: cp + .references + .map(|refs| refs.iter().map(|r| mapper.inbound_text(r)).collect()) + .transpose()?, + deriver: cp.deriver.map(|d| mapper.inbound_text(&d)).transpose()?, + ..cp + }) +} + pub async fn lookup_upstream_narinfo( endpoints: Arc>, pool: Arc, - hash: String, - store_path: String, + mapper: StorePathMapper, + target: StorePath, ) -> ProbeResult { + let original = mapper.outbound(&target); + let hash = original.hash().to_string(); + let store_path = original.full(); let mut samples = Vec::new(); // A tripped upstream is skipped rather than probed and recorded. Folding a sample never taken @@ -367,10 +390,9 @@ pub async fn lookup_upstream_narinfo( .iter() .map(|ep| { let pool = Arc::clone(&pool); - let hash = hash.clone(); - let path = store_path.clone(); + let (mapper, hash, path) = (&mapper, &hash, &store_path); async move { - let (latency, kind, cp) = probe_one(&pool, ep, &hash, &path).await; + let (latency, kind, cp) = probe_one(&pool, ep, mapper, hash, path).await; (ep.id, latency, kind, cp) } }) @@ -390,7 +412,7 @@ pub async fn lookup_upstream_narinfo( } for ep in endpoints.iter() { - let (latency, kind, cp) = probe_one(&pool, ep, &hash, &store_path).await; + let (latency, kind, cp) = probe_one(&pool, ep, &mapper, &hash, &store_path).await; let is_hit = matches!(kind, SampleKind::Hit); samples.push(ProbeSample { upstream: ep.id, @@ -415,7 +437,8 @@ pub struct BatchProbe { pub async fn probe_batch( mut endpoints: Vec, pool: Arc, - targets: Vec<(String, String)>, + mapper: StorePathMapper, + targets: Vec, ) -> BatchProbe { use futures::stream::{FuturesUnordered, StreamExt as _}; @@ -429,18 +452,20 @@ pub async fn probe_batch( let mut futs = FuturesUnordered::new(); let mut iter = targets.into_iter(); - let push = |futs: &mut FuturesUnordered<_>, hash: String, path: String| { + let push = |futs: &mut FuturesUnordered<_>, target: StorePath| { let eps = Arc::clone(&endpoints); let pool = Arc::clone(&pool); + let mapper = mapper.clone(); futs.push(async move { - let res = lookup_upstream_narinfo(eps, pool, hash.clone(), path).await; + let hash = target.hash().to_owned(); + let res = lookup_upstream_narinfo(eps, pool, mapper, target).await; (hash, res) }); }; for _ in 0..BATCH_WINDOW { match iter.next() { - Some((hash, path)) => push(&mut futs, hash, path), + Some(target) => push(&mut futs, target), None => break, } } @@ -453,8 +478,8 @@ pub async fn probe_batch( batch.unanswered.insert(hash, res.unanswered); } - if let Some((hash, path)) = iter.next() { - push(&mut futs, hash, path); + if let Some(target) = iter.next() { + push(&mut futs, target); } } @@ -490,8 +515,8 @@ fn names_path_hash(body: &str, path_hash: &str) -> bool { body.lines() .filter_map(|l| l.strip_prefix("StorePath: ")) .next_back() - .and_then(|p| gradient_sources::get_hash_from_path(p.trim().to_owned()).ok()) - .is_some_and(|(hash, _)| hash == path_hash) + .and_then(|p| StorePath::parse_stored(p.trim()).ok()) + .is_some_and(|p| p.hash() == path_hash) } pub fn parse_upstream_narinfo(base_url: &str, store_path: &str, body: &str) -> Option { @@ -737,8 +762,11 @@ mod tests { lookup_upstream_narinfo( Arc::new(endpoints), Arc::new(Semaphore::new(4)), - SIGNED_HASH.to_owned(), - SIGNED_PATH.to_owned(), + gradient_types::StorePathMapper::for_cache( + Arc::new(gradient_types::IdentityMapping), + gradient_types::ids::CacheId::nil(), + ), + StorePath::parse_stored(SIGNED_PATH).expect("path"), ) .await } @@ -873,18 +901,64 @@ mod tests { const SIGNED_PATH: &str = "/nix/store/brj5bb4pny8pnngq3qdymkllwql6z29j-hello-2.12"; fn signed_narinfo(store_path: &str) -> (String, String) { + signed_narinfo_referencing(store_path, None) + } + + fn signed_narinfo_referencing(store_path: &str, reference: Option<&str>) -> (String, String) { use base64::{Engine as _, engine::general_purpose::STANDARD}; let keypair = ed25519_compact::KeyPair::generate(); let nar_hash = "sha256:1bnnhb0pfx49mg15fmk3jx34wj8j24ygqcq7xww9g8qcyaf23rkf"; - let fingerprint = format!("1;{store_path};{nar_hash};120;"); + let full_reference = reference.map_or(String::new(), |r| format!("/nix/store/{r}")); + let fingerprint = format!("1;{store_path};{nar_hash};120;{full_reference}"); let sig = STANDARD.encode(*keypair.sk.sign(fingerprint, None)); let body = format!( "StorePath: {store_path}\nURL: nar/x.nar.xz\nNarHash: {nar_hash}\nNarSize: 120\n\ - References: \nSig: upstream-1:{sig}\n" + References: {}\nSig: upstream-1:{sig}\n", + reference.unwrap_or_default() ); (body, format!("upstream-1:{}", STANDARD.encode(*keypair.pk))) } + #[tokio::test] + async fn an_upstream_is_asked_for_the_original_hash_and_answers_in_stored_form() { + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + let (original, stored) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let (body, key) = signed_narinfo_referencing( + &format!("/nix/store/{original}-hello"), + Some(&format!("{original}-dep")), + ); + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path(format!("/{original}.narinfo"))) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_string(body)) + .mount(&server) + .await; + let endpoint = UpstreamEndpoint { + url: server.uri(), + public_key: Some(key), + ..ep(None, None) + }; + let mapper = gradient_types::StorePathMapper::for_cache( + Arc::new(RotatedHashMapping), + gradient_types::ids::CacheId::nil(), + ); + + let batch = probe_batch( + vec![endpoint], + Arc::new(Semaphore::new(4)), + mapper, + vec![gradient_types::StorePath::from_stored(stored, "hello")], + ) + .await; + + let found = &batch.found[stored]; + assert_eq!(found.path, format!("/nix/store/{stored}-hello")); + assert_eq!( + found.references, + Some(vec![format!("/nix/store/{stored}-dep")]) + ); + } + fn ep_with_key(public_key: Option) -> UpstreamEndpoint { UpstreamEndpoint { public_key, diff --git a/backend/gradient-db/src/cache_paths.rs b/backend/gradient-db/src/cache_paths.rs index ccf981833..e1e999726 100644 --- a/backend/gradient-db/src/cache_paths.rs +++ b/backend/gradient-db/src/cache_paths.rs @@ -82,6 +82,37 @@ pub async fn project_read_caches( .await } +pub async fn project_write_caches( + db: &C, + project: ProjectId, +) -> Result, DbErr> { + let subscribed: Vec = EProjectCache::find() + .filter(CProjectCache::Project.eq(project)) + .filter(CProjectCache::Mode.is_in([ + CacheSubscriptionMode::ReadWrite, + CacheSubscriptionMode::WriteOnly, + ])) + .all(db) + .await? + .into_iter() + .map(|link| link.cache) + .collect(); + if subscribed.is_empty() { + return Ok(Vec::new()); + } + + let mut active: Vec = ECache::find() + .filter(CCache::Id.is_in(subscribed)) + .filter(CCache::Active.eq(true)) + .all(db) + .await? + .into_iter() + .map(|cache| cache.id) + .collect(); + active.sort(); + Ok(active) +} + pub async fn served_path( db: &C, caches: &[CacheId], diff --git a/backend/gradient-db/src/context.rs b/backend/gradient-db/src/context.rs index 163b7f3d5..70b69aee9 100644 --- a/backend/gradient-db/src/context.rs +++ b/backend/gradient-db/src/context.rs @@ -13,6 +13,8 @@ use tokio::sync::{Notify, mpsc}; use crate::pool::{WebDb, WorkerDb}; use crate::scheduling::startable_set::StartableSet; +use gradient_entity::ids::{CacheId, ProjectId}; +use gradient_entity::{MappingError, StorePathMapper, StorePathMapping}; use gradient_storage::StorageCtx; use gradient_types::{DerivationId, EvaluationId, RuntimeConfig}; use gradient_util::shutdown::Shutdown; @@ -182,9 +184,21 @@ pub struct DbContext { pub events: gradient_types::EventBus, pub wake: Wakeups, pub memory: DbMemory, + pub store_path_mapping: Arc, } impl DbContext { + pub fn cache_mapper(&self, cache: CacheId) -> StorePathMapper { + StorePathMapper::for_cache(Arc::clone(&self.store_path_mapping), cache) + } + + pub async fn project_mapper( + &self, + project: ProjectId, + ) -> Result { + StorePathMapper::for_project(Arc::clone(&self.store_path_mapping), project).await + } + pub fn in_transaction(&self, tx: Arc) -> DbContext { DbContext { worker_db: self.worker_db.in_transaction(tx), diff --git a/backend/gradient-db/src/dashboard/mod.rs b/backend/gradient-db/src/dashboard/mod.rs index ecaa056c3..8537b991e 100644 --- a/backend/gradient-db/src/dashboard/mod.rs +++ b/backend/gradient-db/src/dashboard/mod.rs @@ -31,8 +31,8 @@ mod totals; pub use rail::{RailCacheRow, RailProjectRow, RailTaskRow, rail_caches, rail_projects, rail_tasks}; pub use search::{ - CommitHitRow, NameHitRow, NameKind, NarHitRow, ilike_contains, search_commits, search_names, - search_nars, + CommitHitRow, NameHitRow, NameKind, NarHitRow, ilike_contains, readable_cache_ids, + search_commits, search_names, search_nars, }; pub use stars::{StarKind, StarredNames, StarredTask, star, starred, starred_names, unstar}; pub use tasks::{HistoryRow, TaskFactsRow, entry_point_outcomes, histories, task_facts}; diff --git a/backend/gradient-db/src/dashboard/search.rs b/backend/gradient-db/src/dashboard/search.rs index b15f08e9a..c26378faa 100644 --- a/backend/gradient-db/src/dashboard/search.rs +++ b/backend/gradient-db/src/dashboard/search.rs @@ -9,10 +9,15 @@ use gradient_types::*; use sea_orm::{ConnectionTrait, DbErr, QueryResult, Value}; crate::sql! { - SEARCH_NARS = concat!("SELECT c.name AS cache, cp.hash, cp.package FROM cached_path cp \ + READABLE_CACHE_IDS = concat!("SELECT c.id FROM cache c WHERE ", cache_readable!()), + params = [UserId], + tier = Bulk; + + SEARCH_NARS = concat!("SELECT c.id AS cache_id, c.name AS cache, cp.hash, cp.package FROM cached_path cp \ JOIN cached_path_signature s ON s.cached_path = cp.id JOIN cache c ON c.id = s.cache \ - WHERE cp.hash = $2 AND ", cache_readable!(), " ORDER BY c.name LIMIT 10"), - params = [UserId, CachedPathHash], + WHERE (c.id, cp.hash) IN (SELECT * FROM unnest($2::uuid[], $3::text[])) AND ", cache_readable!(), " \ + ORDER BY c.name LIMIT 10"), + params = [UserId, CacheIds(4), CachedPathHashes(4)], tier = Bulk; SEARCH_COMMITS = concat!("SELECT p.name AS project, t.name AS task, e.id AS evaluation, \ @@ -42,6 +47,7 @@ crate::sql! { #[derive(Clone, Debug, PartialEq, Eq)] pub struct NarHitRow { + pub cache_id: CacheId, pub cache: String, pub hash: String, pub package: String, @@ -90,6 +96,7 @@ fn name_kind(raw: &str, project: Option) -> Result { fn nar_hit_row(r: &QueryResult) -> Result { Ok(NarHitRow { + cache_id: r.try_get("", "cache_id")?, cache: r.try_get("", "cache")?, hash: r.try_get("", "hash")?, package: r.try_get("", "package")?, @@ -114,12 +121,27 @@ pub(super) fn name_hit_row(r: &QueryResult) -> Result { }) } +pub async fn readable_cache_ids( + db: &C, + user: UserId, +) -> Result, DbErr> { + db.query_all_raw(READABLE_CACHE_IDS.bind([user_value(user)])) + .await? + .iter() + .map(|r| r.try_get("", "id")) + .collect() +} + pub async fn search_nars( db: &C, - hash: &str, + lookups: &[(CacheId, StoreHash)], user: UserId, ) -> Result, DbErr> { - db.query_all_raw(SEARCH_NARS.bind([user_value(user), hash.into()])) + let (caches, hashes): (Vec, Vec) = lookups + .iter() + .map(|(cache, hash)| (cache.into_inner(), hash.to_string())) + .unzip(); + db.query_all_raw(SEARCH_NARS.bind([user_value(user), Value::from(caches), Value::from(hashes)])) .await? .iter() .map(nar_hit_row) diff --git a/backend/gradient-db/src/graph/runtime_closure.rs b/backend/gradient-db/src/graph/runtime_closure.rs index b60072abd..1ae49af48 100644 --- a/backend/gradient-db/src/graph/runtime_closure.rs +++ b/backend/gradient-db/src/graph/runtime_closure.rs @@ -19,8 +19,9 @@ struct ReferenceTokens { } pub fn parse_reference_hash(reference: &str) -> Option { - let hash = reference.split('-').next().unwrap_or_default(); - (!hash.is_empty()).then(|| hash.to_string()) + StorePath::parse_stored(reference) + .ok() + .map(|path| path.hash().to_owned()) } fn tokens(references: Option) -> Vec { @@ -178,7 +179,7 @@ mod tests { parse_reference_hash("abc123-hello-2.10").as_deref(), Some("abc123") ); - assert_eq!(parse_reference_hash("abc123").as_deref(), Some("abc123")); + assert_eq!(parse_reference_hash("abc123"), None); assert_eq!(parse_reference_hash(""), None); } diff --git a/backend/gradient-db/src/graph/runtime_dependencies.rs b/backend/gradient-db/src/graph/runtime_dependencies.rs index e2f43359a..5f88edc23 100644 --- a/backend/gradient-db/src/graph/runtime_dependencies.rs +++ b/backend/gradient-db/src/graph/runtime_dependencies.rs @@ -89,9 +89,8 @@ crate::sql_lazy! { } pub(crate) fn hash_of_token(token: &str) -> Option { - let base = token.trim_start_matches("/nix/store/"); - let hash = base.split('-').next()?; - gradient_util::nix_hash::is_nix32_hash(hash).then(|| hash.to_owned()) + let path = gradient_types::StorePath::parse_stored(token).ok()?; + gradient_util::nix_hash::is_nix32_hash(path.hash()).then(|| path.hash().to_owned()) } pub async fn producers_of_tokens( diff --git a/backend/gradient-db/src/projects/caches.rs b/backend/gradient-db/src/projects/caches.rs index 2b2c7b280..c9d6c3354 100644 --- a/backend/gradient-db/src/projects/caches.rs +++ b/backend/gradient-db/src/projects/caches.rs @@ -4,7 +4,6 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use gradient_entity::project_cache::CacheSubscriptionMode; use gradient_types::ids::{CacheId, ProjectId}; use gradient_types::{CProjectCache, EProjectCache, MProjectCache}; use sea_orm::{ColumnTrait, Condition, ConnectionTrait, DbErr, EntityTrait, QueryFilter}; @@ -13,31 +12,9 @@ pub async fn project_has_writable_cache( db: &C, project: ProjectId, ) -> Result { - use gradient_entity::cache::{Column as CCache, Entity as ECache}; - use gradient_entity::project_cache::{Column as COC, Entity as EOC}; - - let cache_ids: Vec = EOC::find() - .filter(COC::Project.eq(project)) - .filter(COC::Mode.is_in([ - CacheSubscriptionMode::ReadWrite, - CacheSubscriptionMode::WriteOnly, - ])) - .all(db) + Ok(!crate::cache_paths::project_write_caches(db, project) .await? - .into_iter() - .map(|r| r.cache) - .collect(); - - if cache_ids.is_empty() { - return Ok(false); - } - - let row = ECache::find() - .filter(CCache::Id.is_in(cache_ids)) - .filter(CCache::Active.eq(true)) - .one(db) - .await?; - Ok(row.is_some()) + .is_empty()) } pub async fn subscriptions_of( diff --git a/backend/gradient-db/src/test_ctx.rs b/backend/gradient-db/src/test_ctx.rs index 0e755697a..d53182229 100644 --- a/backend/gradient-db/src/test_ctx.rs +++ b/backend/gradient-db/src/test_ctx.rs @@ -90,6 +90,7 @@ pub(crate) async fn ctx_at(db: DatabaseConnection, dir: &std::path::Path) -> (Db events: gradient_types::EventBus::new(16), wake: Default::default(), memory: Default::default(), + store_path_mapping: Arc::new(gradient_entity::IdentityMapping), }; (ctx, worker_db) diff --git a/backend/gradient-effects/src/enrich.rs b/backend/gradient-effects/src/enrich.rs index 296d85300..0a7996fc0 100644 --- a/backend/gradient-effects/src/enrich.rs +++ b/backend/gradient-effects/src/enrich.rs @@ -13,15 +13,20 @@ use gradient_db::DbContext; use gradient_types::events::{Event, build, evaluation}; use gradient_types::*; use sea_orm::{ConnectionTrait, EntityTrait}; +use std::sync::Arc; pub async fn enrich(ctx: &DbContext, event: Event) -> Result> { - enrich_with(&ctx.worker_db, event).await + enrich_with(&ctx.worker_db, &ctx.store_path_mapping, event).await } -pub async fn enrich_with(db: &C, event: Event) -> Result> { +pub async fn enrich_with( + db: &C, + mapping: &Arc, + event: Event, +) -> Result> { Ok(match event { - Event::BuildReported(b) => enrich_build(db, b).await?.map(Event::from), - Event::EvaluationReported(e) => enrich_evaluation(db, e).await?.map(Event::from), + Event::BuildReported(b) => enrich_build(db, mapping, b).await?.map(Event::from), + Event::EvaluationReported(e) => enrich_evaluation(db, mapping, e).await?.map(Event::from), other => Some(other), }) } @@ -59,6 +64,7 @@ async fn evaluation_scope(db: &C, id: EvaluationId) -> Resul async fn enrich_build( db: &C, + mapping: &Arc, mut b: build::Reported, ) -> Result> { let Some(scope) = evaluation_scope(db, b.evaluation_id).await? else { @@ -67,16 +73,19 @@ async fn enrich_build( b.task = Some(scope.task); b.project = Some(scope.project); b.evaluation_kind = Some(eval_kind_str(scope.evaluation.kind).to_owned()); - b.derivation_path = EDerivation::find_by_id(b.derivation) - .one(db) - .await - .context("look up the derivation of a build report")? - .map(|d| d.store_path()); + if let Some(mapper) = project_mapper(mapping, scope.project).await { + b.derivation_path = EDerivation::find_by_id(b.derivation) + .one(db) + .await + .context("look up the derivation of a build report")? + .map(|d| mapper.outbound(&d.as_store_path()).full()); + } Ok(Some(b)) } async fn enrich_evaluation( db: &C, + mapping: &Arc, mut e: evaluation::Reported, ) -> Result> { let Some(scope) = evaluation_scope(db, e.evaluation_id).await? else { @@ -84,15 +93,30 @@ async fn enrich_evaluation( }; e.task = Some(scope.task); e.project = Some(scope.project); - e.repository = Some(scope.evaluation.repository); + e.repository = project_mapper(mapping, scope.project) + .await + .map(|mapper| mapper.outbound_if_store_path(&scope.evaluation.repository)); e.evaluation_kind = Some(eval_kind_str(scope.evaluation.kind).to_owned()); Ok(Some(e)) } +async fn project_mapper( + mapping: &Arc, + project: ProjectId, +) -> Option { + StorePathMapper::for_project(Arc::clone(mapping), project) + .await + .inspect_err(|error| tracing::warn!(%project, %error, "event leaves out its store paths")) + .ok() +} + #[cfg(test)] mod tests { use super::*; use gradient_entity::build::BuildStatus; + use gradient_entity::{ + EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping, UnresolvableProjectMapping, + }; use gradient_fixtures::mock::{empty_db, mock_db}; #[tokio::test] @@ -106,7 +130,80 @@ mod tests { } .into(); - assert!(enrich_with(&db, event).await.unwrap().is_none()); + let mapping: Arc = Arc::new(gradient_entity::IdentityMapping); + assert!(enrich_with(&db, &mapping, event).await.unwrap().is_none()); + } + + #[tokio::test] + async fn a_reported_build_names_its_original_derivation_path() { + let evaluation = MEvaluation { + id: EvaluationId::now_v7(), + task: Some(TaskId::now_v7()), + ..Default::default() + }; + let task = MTask { + id: evaluation.task.unwrap(), + project: ProjectId::now_v7(), + ..Default::default() + }; + let derivation = MDerivation { + id: DerivationId::now_v7(), + hash: EXAMPLE_STORED_HASH.into(), + name: "hello".into(), + ..Default::default() + }; + let db = mock_db() + .append_query_results([vec![evaluation.clone()]]) + .append_query_results([vec![task]]) + .append_query_results([vec![derivation.clone()]]) + .into_connection(); + let event: Event = build::Reported { + evaluation_id: evaluation.id, + derivation: derivation.id, + ..Default::default() + } + .into(); + let mapping: Arc = Arc::new(RotatedHashMapping); + + let Some(Event::BuildReported(reported)) = enrich_with(&db, &mapping, event).await.unwrap() + else { + panic!("the build is reported"); + }; + assert_eq!( + reported.derivation_path, + Some(format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-hello.drv")) + ); + } + + #[tokio::test] + async fn a_build_of_a_project_that_does_not_resolve_names_no_derivation_path() { + let evaluation = MEvaluation { + id: EvaluationId::now_v7(), + task: Some(TaskId::now_v7()), + ..Default::default() + }; + let task = MTask { + id: evaluation.task.unwrap(), + project: ProjectId::now_v7(), + ..Default::default() + }; + let db = mock_db() + .append_query_results([vec![evaluation.clone()]]) + .append_query_results([vec![task]]) + .into_connection(); + let event: Event = build::Reported { + evaluation_id: evaluation.id, + derivation: DerivationId::now_v7(), + ..Default::default() + } + .into(); + let mapping: Arc = Arc::new(UnresolvableProjectMapping); + + let Some(Event::BuildReported(reported)) = enrich_with(&db, &mapping, event).await.unwrap() + else { + panic!("the build is reported"); + }; + assert_eq!(reported.derivation_path, None); } #[tokio::test] @@ -118,7 +215,11 @@ mod tests { } .into(); - assert_eq!(enrich_with(&db, event.clone()).await.unwrap(), Some(event)); + let mapping: Arc = Arc::new(gradient_entity::IdentityMapping); + assert_eq!( + enrich_with(&db, &mapping, event.clone()).await.unwrap(), + Some(event) + ); assert!(db.into_transaction_log().is_empty()); } } diff --git a/backend/gradient-entity/Cargo.toml b/backend/gradient-entity/Cargo.toml index 408c51abe..f6089fea7 100644 --- a/backend/gradient-entity/Cargo.toml +++ b/backend/gradient-entity/Cargo.toml @@ -14,8 +14,12 @@ path = "src/lib.rs" [lints] workspace = true +[features] +testing = [] + [dependencies] chrono = { workspace = true } +futures = { workspace = true } num_enum = { workspace = true } sea-orm = { workspace = true } serde = { workspace = true } diff --git a/backend/gradient-entity/src/cached_path.rs b/backend/gradient-entity/src/cached_path.rs index f549fe7ef..e99a27c40 100644 --- a/backend/gradient-entity/src/cached_path.rs +++ b/backend/gradient-entity/src/cached_path.rs @@ -41,7 +41,7 @@ impl ActiveModelBehavior for ActiveModel {} impl Model { pub fn as_store_path(&self) -> crate::StorePath { - crate::StorePath::from_parts(self.hash.clone(), self.package.clone()) + crate::StorePath::from_stored(self.hash.clone(), self.package.clone()) } pub fn store_path(&self) -> String { diff --git a/backend/gradient-entity/src/derivation.rs b/backend/gradient-entity/src/derivation.rs index 859a2817d..3ae85e470 100644 --- a/backend/gradient-entity/src/derivation.rs +++ b/backend/gradient-entity/src/derivation.rs @@ -35,7 +35,7 @@ pub struct Model { impl Model { pub fn as_store_path(&self) -> crate::StorePath { - crate::StorePath::from_parts(self.hash.clone(), format!("{}.drv", self.name)) + crate::StorePath::from_stored(self.hash.clone(), format!("{}.drv", self.name)) } pub fn drv_path(&self) -> String { diff --git a/backend/gradient-entity/src/eval_cache_store.rs b/backend/gradient-entity/src/eval_cache_store.rs index d0de64308..f894d46b3 100644 --- a/backend/gradient-entity/src/eval_cache_store.rs +++ b/backend/gradient-entity/src/eval_cache_store.rs @@ -8,14 +8,14 @@ use chrono::NaiveDateTime; use sea_orm::entity::prelude::*; use serde::{Deserialize, Serialize}; -use crate::ids::EvalCacheStoreId; +use crate::ids::{CacheId, EvalCacheStoreId}; #[derive(Clone, Debug, Default, PartialEq, DeriveEntityModel, Deserialize, Serialize)] #[sea_orm(table_name = "eval_cache_store")] pub struct Model { #[sea_orm(primary_key)] pub id: EvalCacheStoreId, - #[sea_orm(unique)] + pub cache: CacheId, pub fingerprint: String, pub storage_path: String, pub size_bytes: i64, diff --git a/backend/gradient-entity/src/lib.rs b/backend/gradient-entity/src/lib.rs index 1dd58c8a3..e41bb0608 100644 --- a/backend/gradient-entity/src/lib.rs +++ b/backend/gradient-entity/src/lib.rs @@ -6,9 +6,17 @@ pub mod ids; pub mod store_path; +pub mod store_path_mapping; pub mod transition; -pub use store_path::{STORE_DIR, StorePath, StorePathError}; +pub use store_path::{OriginalStorePath, STORE_DIR, StoreHash, StorePath, StorePathError}; +#[cfg(feature = "testing")] +pub use store_path_mapping::{ + EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping, UnresolvableProjectMapping, +}; +pub use store_path_mapping::{ + IdentityMapping, MappingContext, MappingError, StorePathMapper, StorePathMapping, +}; pub mod admin_task; pub mod api; diff --git a/backend/gradient-entity/src/store_path.rs b/backend/gradient-entity/src/store_path.rs index f12ba4e8e..0ddaafa7f 100644 --- a/backend/gradient-entity/src/store_path.rs +++ b/backend/gradient-entity/src/store_path.rs @@ -11,48 +11,138 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; pub const STORE_DIR: &str = "/nix/store"; -/// A store path is stored and compared without the `/nix/store/` prefix. [`StorePath::full`] is -/// adding it only where a real filesystem path is needed. `name` is keeping a trailing `.drv` for -/// derivations. +pub(crate) const NIX32_ALPHABET: &str = "0123456789abcdfghijklmnpqrsvwxyz"; +const STORE_HASH_LENGTH: usize = 32; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct StoreHash(String); + +impl StoreHash { + pub fn parse(hash: &str) -> Result { + let is_nix32 = + hash.len() == STORE_HASH_LENGTH && hash.chars().all(|c| NIX32_ALPHABET.contains(c)); + + if !is_nix32 { + return Err(StorePathError::Malformed(hash.to_owned())); + } + + Ok(Self(hash.to_owned())) + } + + pub(crate) fn stored(hash: String) -> Self { + Self(hash) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Display for StoreHash { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +fn split_base(input: &str) -> Result<(&str, &str), StorePathError> { + let base = input + .strip_prefix(STORE_DIR) + .map(|rest| rest.trim_start_matches('/')) + .unwrap_or(input); + + match base.split_once('-') { + Some((hash, name)) if !hash.is_empty() && !name.is_empty() => Ok((hash, name)), + _ => Err(StorePathError::Malformed(input.to_owned())), + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct OriginalStorePath { + hash: StoreHash, + name: String, +} + +impl OriginalStorePath { + pub(crate) fn from_hash(hash: StoreHash, name: &str) -> Self { + Self { + hash, + name: name.to_owned(), + } + } + + pub fn parse(input: &str) -> Result { + let (hash, name) = split_base(input)?; + let hash = + StoreHash::parse(hash).map_err(|_| StorePathError::Malformed(input.to_owned()))?; + + Ok(Self { + hash, + name: name.to_owned(), + }) + } + + pub fn hash(&self) -> &StoreHash { + &self.hash + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn base(&self) -> String { + format!("{}-{}", self.hash, self.name) + } + + pub fn full(&self) -> String { + format!("{}/{}", STORE_DIR, self.base()) + } +} + +impl fmt::Display for OriginalStorePath { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}-{}", self.hash, self.name) + } +} + +// Stored without the `/nix/store/` prefix; `name` keeps a trailing `.drv` for derivations. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct StorePath { - hash: String, + hash: StoreHash, name: String, } impl From for StorePath { fn from(output: crate::derivation_output::Model) -> Self { - StorePath::from_parts(output.hash, output.package) + StorePath::from_stored(output.hash, output.package) } } impl StorePath { - pub fn from_parts(hash: impl Into, name: impl Into) -> Self { + pub fn from_stored(hash: impl Into, name: impl Into) -> Self { Self { - hash: hash.into(), + hash: StoreHash::stored(hash.into()), name: name.into(), } } - pub fn parse(input: &str) -> Result { - let base = input - .strip_prefix(STORE_DIR) - .map(|rest| rest.trim_start_matches('/')) - .unwrap_or(input); - - let (hash, name) = base - .split_once('-') - .ok_or_else(|| StorePathError::Malformed(input.to_owned()))?; - - if hash.is_empty() || name.is_empty() { - return Err(StorePathError::Malformed(input.to_owned())); + pub(crate) fn from_hash(hash: StoreHash, name: &str) -> Self { + Self { + hash, + name: name.to_owned(), } + } + + pub(crate) fn store_hash(&self) -> &StoreHash { + &self.hash + } - Ok(Self::from_parts(hash, name)) + pub fn parse_stored(input: &str) -> Result { + let (hash, name) = split_base(input)?; + Ok(Self::from_stored(hash, name)) } pub fn hash(&self) -> &str { - &self.hash + self.hash.as_str() } pub fn name(&self) -> &str { @@ -82,7 +172,7 @@ impl FromStr for StorePath { type Err = StorePathError; fn from_str(s: &str) -> Result { - Self::parse(s) + Self::parse_stored(s) } } @@ -95,7 +185,7 @@ impl Serialize for StorePath { impl<'de> Deserialize<'de> for StorePath { fn deserialize>(deserializer: D) -> Result { let raw = String::deserialize(deserializer)?; - StorePath::parse(&raw).map_err(serde::de::Error::custom) + StorePath::parse_stored(&raw).map_err(serde::de::Error::custom) } } @@ -111,7 +201,7 @@ impl sea_orm::TryGetable for StorePath { idx: I, ) -> Result { let raw = String::try_get_by(res, idx)?; - StorePath::parse(&raw) + StorePath::parse_stored(&raw) .map_err(|e| sea_orm::TryGetError::DbErr(sea_orm::DbErr::Type(e.to_string()))) } } @@ -120,7 +210,7 @@ impl sea_orm::sea_query::ValueType for StorePath { fn try_from(v: sea_orm::Value) -> Result { match v { sea_orm::Value::String(Some(s)) => { - StorePath::parse(&s).map_err(|_| sea_orm::sea_query::ValueTypeErr) + StorePath::parse_stored(&s).map_err(|_| sea_orm::sea_query::ValueTypeErr) } _ => Err(sea_orm::sea_query::ValueTypeErr), } @@ -164,14 +254,32 @@ impl std::error::Error for StorePathError {} mod tests { use super::*; + #[test] + fn an_original_path_needs_a_full_nix32_hash() { + let hash = "0123456789abcdfghijklmnpqrsvwxyz"; + let path = OriginalStorePath::parse(&format!("/nix/store/{hash}-hello-2.12.1")).unwrap(); + assert_eq!(path.hash().as_str(), hash); + assert_eq!(path.name(), "hello-2.12.1"); + assert_eq!(path.full(), format!("/nix/store/{hash}-hello-2.12.1")); + assert_eq!(OriginalStorePath::parse(&path.base()).unwrap(), path); + + assert!(OriginalStorePath::parse("abc123-hello").is_err()); + assert!(OriginalStorePath::parse("0123456789abcdfghijklmnpqrsvwxye-hello").is_err()); + assert!(OriginalStorePath::parse(&format!("{hash}-")).is_err()); + assert_eq!( + StorePath::parse_stored("abc123-hello").unwrap().hash(), + "abc123" + ); + } + #[test] fn parse_reads_hash_name_and_derivation_flag() { - let p = StorePath::parse("/nix/store/abc123-hello-2.12.1").unwrap(); + let p = StorePath::parse_stored("/nix/store/abc123-hello-2.12.1").unwrap(); assert_eq!(p.hash(), "abc123"); assert_eq!(p.name(), "hello-2.12.1"); assert!(!p.is_derivation()); - let p = StorePath::parse("abc123-hello-2.12.1.drv").unwrap(); + let p = StorePath::parse_stored("abc123-hello-2.12.1.drv").unwrap(); assert_eq!(p.hash(), "abc123"); assert_eq!(p.name(), "hello-2.12.1.drv"); assert!(p.is_derivation()); @@ -179,19 +287,19 @@ mod tests { #[test] fn full_and_base_round_trip_and_display_blanks_the_prefix() { - let p = StorePath::from_parts("abc123", "hello-2.12.1"); + let p = StorePath::from_stored("abc123", "hello-2.12.1"); assert_eq!(p.base(), "abc123-hello-2.12.1"); assert_eq!(p.full(), "/nix/store/abc123-hello-2.12.1"); - assert_eq!(StorePath::parse(&p.full()).unwrap(), p); - assert_eq!(StorePath::parse(&p.base()).unwrap(), p); + assert_eq!(StorePath::parse_stored(&p.full()).unwrap(), p); + assert_eq!(StorePath::parse_stored(&p.base()).unwrap(), p); - let drv = StorePath::from_parts("abc123", "hello.drv"); + let drv = StorePath::from_stored("abc123", "hello.drv"); assert_eq!(drv.to_string(), "abc123-hello.drv"); } #[test] fn serde_uses_base_form() { - let p = StorePath::from_parts("abc123", "hello-2.12.1"); + let p = StorePath::from_stored("abc123", "hello-2.12.1"); let json = serde_json::to_string(&p).unwrap(); assert_eq!(json, "\"abc123-hello-2.12.1\""); @@ -203,16 +311,16 @@ mod tests { #[test] fn rejects_malformed() { - assert!(StorePath::parse("nodash").is_err()); - assert!(StorePath::parse("-noname").is_err()); - assert!(StorePath::parse("nohash-").is_err()); - assert!(StorePath::parse("/nix/store/").is_err()); + assert!(StorePath::parse_stored("nodash").is_err()); + assert!(StorePath::parse_stored("-noname").is_err()); + assert!(StorePath::parse_stored("nohash-").is_err()); + assert!(StorePath::parse_stored("/nix/store/").is_err()); } #[test] fn sea_orm_value_round_trip_is_prefix_free() { use sea_orm::sea_query::ValueType; - let p = StorePath::from_parts("abc123", "hello-2.12.1"); + let p = StorePath::from_stored("abc123", "hello-2.12.1"); let v = sea_orm::Value::from(p.clone()); assert_eq!(v, sea_orm::Value::String(Some(p.base()))); assert_eq!(::try_from(v).unwrap(), p); diff --git a/backend/gradient-entity/src/store_path_mapping.rs b/backend/gradient-entity/src/store_path_mapping.rs new file mode 100644 index 000000000..ebe4613ae --- /dev/null +++ b/backend/gradient-entity/src/store_path_mapping.rs @@ -0,0 +1,373 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::fmt; +use std::sync::Arc; + +use futures::future::BoxFuture; + +use crate::ids::{CacheId, ProjectId}; +use crate::store_path::{ + NIX32_ALPHABET, OriginalStorePath, STORE_DIR, StoreHash, StorePath, StorePathError, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum MappingContext { + Project(ProjectId), + Cache(CacheId), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MappingError(String); + +impl MappingError { + pub fn new(message: impl Into) -> Self { + Self(message.into()) + } +} + +impl fmt::Display for MappingError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for MappingError {} + +pub trait StorePathMapping: Send + Sync { + fn resolve( + &self, + context: MappingContext, + ) -> BoxFuture<'_, Result> { + Box::pin(std::future::ready(Ok(context))) + } + + fn inbound(&self, context: MappingContext, hash: StoreHash) -> StoreHash; + fn outbound(&self, context: MappingContext, hash: StoreHash) -> StoreHash; + fn inbound_prefix(&self, context: MappingContext, prefix: &str) -> Option; +} + +impl fmt::Debug for dyn StorePathMapping { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("StorePathMapping") + } +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct IdentityMapping; + +impl StorePathMapping for IdentityMapping { + fn inbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn outbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn inbound_prefix(&self, _context: MappingContext, prefix: &str) -> Option { + Some(prefix.to_owned()) + } +} + +#[derive(Clone, Debug)] +pub struct StorePathMapper { + mapping: Arc, + context: MappingContext, +} + +impl StorePathMapper { + pub fn for_cache(mapping: Arc, cache: CacheId) -> Self { + Self { + mapping, + context: MappingContext::Cache(cache), + } + } + + pub async fn for_project( + mapping: Arc, + project: ProjectId, + ) -> Result { + let context = mapping.resolve(MappingContext::Project(project)).await?; + Ok(Self { mapping, context }) + } + + pub fn context(&self) -> MappingContext { + self.context + } + + pub fn inbound(&self, path: &OriginalStorePath) -> StorePath { + let hash = self.mapping.inbound(self.context, path.hash().clone()); + StorePath::from_hash(hash, path.name()) + } + + pub fn outbound(&self, path: &StorePath) -> OriginalStorePath { + let hash = self + .mapping + .outbound(self.context, path.store_hash().clone()); + OriginalStorePath::from_hash(hash, path.name()) + } + + pub fn inbound_hash(&self, hash: StoreHash) -> StoreHash { + self.mapping.inbound(self.context, hash) + } + + pub fn inbound_text(&self, text: &str) -> Result { + let parts = TextParts::split(text); + let hash = + StoreHash::parse(parts.hash).map_err(|_| StorePathError::Malformed(text.to_owned()))?; + Ok(parts.join(&self.mapping.inbound(self.context, hash))) + } + + pub fn outbound_text(&self, text: &str) -> Result { + let parts = TextParts::split(text); + if parts.hash.is_empty() { + return Err(StorePathError::Malformed(text.to_owned())); + } + let hash = StoreHash::stored(parts.hash.to_owned()); + Ok(parts.join(&self.mapping.outbound(self.context, hash))) + } + + pub fn inbound_if_store_path(&self, text: &str) -> String { + self.if_store_path(text, Self::inbound_text) + } + + pub fn outbound_if_store_path(&self, text: &str) -> String { + self.if_store_path(text, Self::outbound_text) + } + + fn if_store_path( + &self, + text: &str, + map: fn(&Self, &str) -> Result, + ) -> String { + match text.strip_prefix(STORE_DIR) { + Some(_) => map(self, text).unwrap_or_else(|_| text.to_owned()), + None => text.to_owned(), + } + } + + pub fn inbound_prefix(&self, prefix: &str) -> Option { + self.mapping.inbound_prefix(self.context, prefix) + } +} + +struct TextParts<'a> { + store_dir: &'a str, + hash: &'a str, + rest: &'a str, +} + +impl<'a> TextParts<'a> { + fn split(text: &'a str) -> Self { + let store_dir_length = text + .strip_prefix(STORE_DIR) + .and_then(|rest| rest.strip_prefix('/')) + .map_or(0, |_| STORE_DIR.len() + 1); + let (store_dir, base) = text.split_at(store_dir_length); + let hash_length = base.find('-').unwrap_or(base.len()); + let (hash, rest) = base.split_at(hash_length); + Self { + store_dir, + hash, + rest, + } + } + + fn join(&self, hash: &StoreHash) -> String { + format!("{}{}{}", self.store_dir, hash, self.rest) + } +} + +#[cfg(feature = "testing")] +pub const EXAMPLE_ORIGINAL_HASH: &str = "0123456789abcdfghijklmnpqrsvwxyz"; +#[cfg(feature = "testing")] +pub const EXAMPLE_STORED_HASH: &str = "123456789abcdfghijklmnpqrsvwxyz0"; + +#[cfg(feature = "testing")] +#[derive(Clone, Copy, Debug, Default)] +pub struct RotatedHashMapping; + +#[cfg(feature = "testing")] +impl RotatedHashMapping { + fn rotate(text: &str, steps: usize) -> String { + let alphabet: Vec = NIX32_ALPHABET.chars().collect(); + text.chars() + .map(|c| match alphabet.iter().position(|&a| a == c) { + Some(index) => alphabet[(index + steps) % alphabet.len()], + None => c, + }) + .collect() + } +} + +#[cfg(feature = "testing")] +impl StorePathMapping for RotatedHashMapping { + fn inbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + StoreHash::stored(Self::rotate(hash.as_str(), 1)) + } + + fn outbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + StoreHash::stored(Self::rotate(hash.as_str(), NIX32_ALPHABET.len() - 1)) + } + + fn inbound_prefix(&self, _context: MappingContext, prefix: &str) -> Option { + Some(Self::rotate(prefix, 1)) + } +} + +#[cfg(feature = "testing")] +#[derive(Clone, Copy, Debug, Default)] +pub struct UnresolvableProjectMapping; + +#[cfg(feature = "testing")] +impl StorePathMapping for UnresolvableProjectMapping { + fn resolve( + &self, + context: MappingContext, + ) -> BoxFuture<'_, Result> { + let resolved = match context { + MappingContext::Project(project) => Err(MappingError::new(format!( + "project {project} has no mapping context" + ))), + MappingContext::Cache(_) => Ok(context), + }; + Box::pin(std::future::ready(resolved)) + } + + fn inbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn outbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn inbound_prefix(&self, _context: MappingContext, prefix: &str) -> Option { + Some(prefix.to_owned()) + } +} + +#[cfg(all(test, feature = "testing"))] +mod tests { + use super::*; + + fn rotated() -> StorePathMapper { + StorePathMapper::for_cache(Arc::new(RotatedHashMapping), CacheId::nil()) + } + + struct ProjectCacheMapping(CacheId); + + impl StorePathMapping for ProjectCacheMapping { + fn resolve( + &self, + context: MappingContext, + ) -> BoxFuture<'_, Result> { + let resolved = match context { + MappingContext::Project(_) => MappingContext::Cache(self.0), + MappingContext::Cache(_) => context, + }; + Box::pin(std::future::ready(Ok(resolved))) + } + + fn inbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn outbound(&self, _context: MappingContext, hash: StoreHash) -> StoreHash { + hash + } + + fn inbound_prefix(&self, _context: MappingContext, prefix: &str) -> Option { + Some(prefix.to_owned()) + } + } + + #[tokio::test] + async fn a_project_mapper_maps_with_the_resolved_context() { + let cache = CacheId::new(uuid::Uuid::from_u128(0x5c4e_d2a1)); + let project = ProjectId::new(uuid::Uuid::from_u128(0x9b07_e113)); + let mapper = StorePathMapper::for_project(Arc::new(ProjectCacheMapping(cache)), project) + .await + .unwrap(); + assert_eq!(mapper.context(), MappingContext::Cache(cache)); + + let identity = StorePathMapper::for_project(Arc::new(IdentityMapping), project) + .await + .unwrap(); + assert_eq!(identity.context(), MappingContext::Project(project)); + } + + #[test] + fn both_mappings_return_the_hash_they_were_given() { + let hash = StoreHash::parse(EXAMPLE_ORIGINAL_HASH).unwrap(); + let mappings: [Arc; 2] = + [Arc::new(IdentityMapping), Arc::new(RotatedHashMapping)]; + for mapping in mappings { + for context in [ + MappingContext::Project(ProjectId::nil()), + MappingContext::Cache(CacheId::nil()), + ] { + let stored = mapping.inbound(context, hash.clone()); + assert_eq!(mapping.outbound(context, stored), hash); + } + } + } + + #[test] + fn the_rotated_mapping_moves_every_character_one_step() { + let hash = StoreHash::parse(EXAMPLE_ORIGINAL_HASH).unwrap(); + let context = MappingContext::Cache(CacheId::nil()); + assert_eq!( + RotatedHashMapping.inbound(context, hash).as_str(), + EXAMPLE_STORED_HASH + ); + assert_eq!( + RotatedHashMapping.inbound_prefix(context, "012").as_deref(), + Some("123") + ); + assert_eq!( + IdentityMapping.inbound_prefix(context, "012").as_deref(), + Some("012") + ); + } + + #[test] + fn text_mapping_swaps_only_the_hash() { + let (o, s) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let mapper = rotated(); + assert_eq!( + mapper + .inbound_text(&format!("/nix/store/{o}-hello/share/doc")) + .unwrap(), + format!("/nix/store/{s}-hello/share/doc") + ); + assert_eq!( + mapper.inbound_text(&format!("{o}-hello.drv")).unwrap(), + format!("{s}-hello.drv") + ); + assert_eq!(mapper.inbound_text(o).unwrap(), s); + assert_eq!( + mapper + .outbound_text(&format!("/nix/store/{s}-hello")) + .unwrap(), + format!("/nix/store/{o}-hello") + ); + assert!(mapper.inbound_text("abc-hello").is_err()); + } + + #[test] + fn only_a_store_path_repository_is_mapped_outbound() { + let (o, s) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let mapper = rotated(); + assert_eq!( + mapper.outbound_if_store_path(&format!("/nix/store/{s}-source")), + format!("/nix/store/{o}-source") + ); + let url = "https://example.invalid/my-repo"; + assert_eq!(mapper.outbound_if_store_path(url), url); + } +} diff --git a/backend/gradient-fixtures/Cargo.toml b/backend/gradient-fixtures/Cargo.toml index 9c84b9888..aacedde0f 100644 --- a/backend/gradient-fixtures/Cargo.toml +++ b/backend/gradient-fixtures/Cargo.toml @@ -15,7 +15,7 @@ path = "src/lib.rs" workspace = true [dependencies] -gradient-entity = { workspace = true } +gradient-entity = { workspace = true, features = ["testing"] } gradient-migration = { workspace = true } gradient-types = { workspace = true } diff --git a/backend/gradient-fixtures/src/database.rs b/backend/gradient-fixtures/src/database.rs index 035f750ad..aabfd7502 100644 --- a/backend/gradient-fixtures/src/database.rs +++ b/backend/gradient-fixtures/src/database.rs @@ -101,6 +101,33 @@ impl TestDatabase { .collect() } + pub async fn columns_containing(&self, needle: &str) -> Vec { + assert!( + needle.chars().all(|c| c.is_ascii_alphanumeric()), + "{needle} is no plain hash" + ); + let columns = self + .texts( + "SELECT table_name || '.' || column_name FROM information_schema.columns \ + WHERE table_schema = 'public' \ + AND (data_type IN ('text', 'character varying', 'jsonb') OR udt_name = '_text') \ + ORDER BY 1", + ) + .await; + let mut containing = Vec::new(); + for column in columns { + let (table, name) = column.split_once('.').expect("table.column"); + let hits = self + .texts(&format!( + "SELECT '{column}'::text FROM public.\"{table}\" \ + WHERE \"{name}\"::text LIKE '%{needle}%' LIMIT 1" + )) + .await; + containing.extend(hits); + } + containing + } + pub async fn find_derivation(&self, path: &str) -> Option { let (hash, name) = crate::shape::drv_hash_name(path); derivation::Entity::find() diff --git a/backend/gradient-fixtures/src/shape/mod.rs b/backend/gradient-fixtures/src/shape/mod.rs index 04678c25a..b259d50aa 100644 --- a/backend/gradient-fixtures/src/shape/mod.rs +++ b/backend/gradient-fixtures/src/shape/mod.rs @@ -623,7 +623,8 @@ fn known_upstream(output: &mut derivation_output::Model) { } pub(crate) fn drv_hash_name(path: &str) -> (String, String) { - let parsed = StorePath::parse(path).unwrap_or_else(|_| panic!("{path} is no store path")); + let parsed = + StorePath::parse_stored(path).unwrap_or_else(|_| panic!("{path} is no store path")); let name = parsed .name() .strip_suffix(".drv") diff --git a/backend/gradient-fixtures/src/shape/scheduling.rs b/backend/gradient-fixtures/src/shape/scheduling.rs index dcd6325c6..b47f8c5c3 100644 --- a/backend/gradient-fixtures/src/shape/scheduling.rs +++ b/backend/gradient-fixtures/src/shape/scheduling.rs @@ -296,10 +296,12 @@ impl Shape { } pub fn eval_cache(mut self, fingerprint: &str) -> Self { + let cache = self.caches[self.cache.expect("name a cache first")].id; self.eval_caches.push(eval_cache_store::Model { id: EvalCacheStoreId::now_v7(), + cache, fingerprint: fingerprint.into(), - storage_path: format!("eval-cache/{fingerprint}"), + storage_path: format!("eval-cache/{cache}/{fingerprint}"), size_bytes: 1, created_at: test_date(), updated_at: test_date(), diff --git a/backend/gradient-graph/src/cases/mod.rs b/backend/gradient-graph/src/cases/mod.rs index 55682aa5a..b3bdab766 100644 --- a/backend/gradient-graph/src/cases/mod.rs +++ b/backend/gradient-graph/src/cases/mod.rs @@ -29,6 +29,7 @@ use std::time::Duration; use gradient_db::evaluations::counters::eval_counters; use gradient_db::graph::consistency::ConsistencyReport; use gradient_db::{DbContext, DbMemory, ProbeRequests}; +use gradient_entity::StorePathMapping; use gradient_entity::build::BuildStatus; use gradient_entity::derivation_dependency::EdgeKind; use gradient_entity::dispatched_job::DispatchedJobOutcome; @@ -95,7 +96,20 @@ impl Case { case } + pub async fn start_with_mapping(shape: Shape, mapping: Arc) -> Self { + let (case, report) = Self::start_mapped(shape, Some(mapping)).await; + assert_eq!(report.total(), 0, "drift after the seed: {report:?}"); + case + } + pub async fn start_with_drift(shape: Shape) -> (Self, ConsistencyReport) { + Self::start_mapped(shape, None).await + } + + async fn start_mapped( + shape: Shape, + mapping: Option>, + ) -> (Self, ConsistencyReport) { let declared = shape.declared_statuses(); let secret = SECRET.with(|(_, path)| path.clone()); let db = shape.into_database().await; @@ -109,6 +123,7 @@ impl Case { probe_requests, ..ctx.memory.clone() }, + store_path_mapping: mapping.unwrap_or_else(|| Arc::clone(&ctx.store_path_mapping)), ..ctx }; diff --git a/backend/gradient-graph/src/cases/signing.rs b/backend/gradient-graph/src/cases/signing.rs index 40be6ad03..64beded4d 100644 --- a/backend/gradient-graph/src/cases/signing.rs +++ b/backend/gradient-graph/src/cases/signing.rs @@ -4,6 +4,9 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +use std::sync::Arc; + +use gradient_entity::RotatedHashMapping; use gradient_entity::build::BuildStatus; use gradient_entity::cached_path_signature; use gradient_entity::ids::{CacheId, EvaluationId}; @@ -410,3 +413,32 @@ async fn k09_a_drv_path_without_a_producer_is_signed_for_the_uploading_project() case.zero_drift("step 1").await; case.finish().await; } + +#[tokio::test] +async fn k10_a_signature_verifies_over_the_outbound_paths() { + let (e1, c1) = (EvaluationId::now_v7(), CacheId::now_v7()); + let case = Case::start_with_mapping( + chain(e1).into_cache(c1).keyed(Case::key()), + Arc::new(RotatedHashMapping), + ) + .await; + let stored = commit_of_b(case.project_caches(e1).await); + case.commit_nar(stored.clone()).await; + + let mapper = case.ctx.cache_mapper(c1); + let seen = NarCommit { + store_path: mapper + .outbound_text(&stored.store_path) + .expect("the stored path maps"), + references: stored + .references + .iter() + .map(|reference| mapper.outbound_text(reference).expect("a reference maps")) + .collect(), + ..stored.clone() + }; + let row = row(&case, B_OUT, c1, "k10").await; + assert!(case.verifies(&row, &seen).await, "{row:?}"); + assert!(!case.verifies(&row, &stored).await, "{row:?}"); + case.finish().await; +} diff --git a/backend/gradient-graph/src/claims.rs b/backend/gradient-graph/src/claims.rs index 78c3facd4..01b10c422 100644 --- a/backend/gradient-graph/src/claims.rs +++ b/backend/gradient-graph/src/claims.rs @@ -73,7 +73,7 @@ struct Claim { } impl Claim { - fn signature(&self, signer: &CacheSigner) -> Option> { + fn signature(&self, signer: &CacheSigner, mapper: &StorePathMapper) -> Option> { let (nar_hash, nar_size) = (self.nar_hash.as_deref()?, self.nar_size?); let references: Vec = self .references @@ -83,8 +83,9 @@ impl Claim { .map(str::to_owned) .collect(); - Some(signer.sign_narinfo_raw( - &StorePath::from_parts(&self.hash, &self.package).full(), + Some(signer.sign_stored_narinfo( + mapper, + &StorePath::from_stored(&self.hash, &self.package).full(), &normalize_nar_hash(nar_hash), nar_size as u64, &references, @@ -143,7 +144,7 @@ async fn sign_claims(ctx: &DbContext, claims: Statement) -> anyhow::Result<()> { signature: signers .get(&claim.cache) .and_then(Option::as_ref) - .and_then(|signer| claim.signature(signer)), + .and_then(|signer| claim.signature(signer, &ctx.cache_mapper(claim.cache))), created_at, ..Default::default() } diff --git a/backend/gradient-graph/src/nar.rs b/backend/gradient-graph/src/nar.rs index 82e23e146..fc9570a52 100644 --- a/backend/gradient-graph/src/nar.rs +++ b/backend/gradient-graph/src/nar.rs @@ -25,7 +25,7 @@ use crate::messages::{NarCommit, NarCommitted, SignTargets}; /// The shared build locks are always taken after the `cached_path` lock, never before. pub(crate) async fn commit(ctx: &DbContext, c: &NarCommit) -> anyhow::Result { let db = &ctx.worker_db; - let sp = StorePath::parse(&c.store_path)?; + let sp = StorePath::parse_stored(&c.store_path)?; if !is_nix32_hash(sp.hash()) { anyhow::bail!(StorePathError::Malformed(c.store_path.clone())); } @@ -107,7 +107,7 @@ pub(crate) async fn commit_batch( let paths = commits .iter() .map(|c| { - let sp = StorePath::parse(&c.store_path)?; + let sp = StorePath::parse_stored(&c.store_path)?; if !is_nix32_hash(sp.hash()) { anyhow::bail!(StorePathError::Malformed(c.store_path.clone())); } @@ -308,7 +308,7 @@ fn union(a: &[DerivationId], b: &[DerivationId]) -> Vec { } pub(crate) fn after_commit(ctx: &DbContext, committed: &NarCommitted, store_path: &str) { - let Ok(sp) = StorePath::parse(store_path) else { + let Ok(sp) = StorePath::parse_stored(store_path) else { return; }; if !gradient_db::caches::debug_info::carries_debug_info(sp.name()) { @@ -496,7 +496,8 @@ async fn sign_into_caches( (!private.contains(sp.hash())).then(|| (sp.full(), normalize_nar_hash(&c.nar_hash))); for (cache, signer) in &signers[&c.targets] { let signature = match (&fingerprint, signer) { - (Some((store_path, nar_hash)), Some(signer)) => Some(signer.sign_narinfo_raw( + (Some((store_path, nar_hash)), Some(signer)) => Some(signer.sign_stored_narinfo( + &ctx.cache_mapper(*cache), store_path, nar_hash, c.nar_size as u64, diff --git a/backend/gradient-graph/src/record.rs b/backend/gradient-graph/src/record.rs index 26abf163f..f4a19cfd9 100644 --- a/backend/gradient-graph/src/record.rs +++ b/backend/gradient-graph/src/record.rs @@ -454,7 +454,7 @@ impl BatchWriter<'_> { let hashes: Vec = d .outputs .iter() - .filter_map(|o| StorePath::parse(&o.path).ok()) + .filter_map(|o| StorePath::parse_stored(&o.path).ok()) .map(|sp| sp.hash().to_owned()) .collect(); ids.get(&d.drv_path) @@ -771,7 +771,7 @@ impl BatchWriter<'_> { }; for src in &d.input_sources { - let Ok(store_path) = StorePath::parse(src) else { + let Ok(store_path) = StorePath::parse_stored(src) else { continue; }; @@ -1309,13 +1309,13 @@ pub(crate) async fn fail_evaluation(ctx: &DbContext, evaluation: EvaluationId, m } fn drv_hash_name(drv_path: &str) -> Option<(String, String)> { - let sp = StorePath::parse(drv_path).ok()?; + let sp = StorePath::parse_stored(drv_path).ok()?; let name = sp.name().strip_suffix(".drv")?; Some((sp.hash().to_owned(), name.to_owned())) } fn output_hash_name(path: &str) -> Option<(String, String)> { - let sp = StorePath::parse(path).ok()?; + let sp = StorePath::parse_stored(path).ok()?; Some((sp.hash().to_owned(), sp.name().to_owned())) } diff --git a/backend/gradient-graph/src/self_heal.rs b/backend/gradient-graph/src/self_heal.rs index d3dbc00e6..3cec961bc 100644 --- a/backend/gradient-graph/src/self_heal.rs +++ b/backend/gradient-graph/src/self_heal.rs @@ -9,13 +9,6 @@ use gradient_db::DbContext; use gradient_types::*; use tracing::{info, warn}; -fn store_path_hash(store_path: &str) -> Option<&str> { - store_path - .strip_prefix("/nix/store/") - .and_then(|s| s.split('-').next()) - .filter(|h| !h.is_empty()) -} - async fn any_reachable(db: &C, derivations: &[DerivationId]) -> bool { for d in derivations { if gradient_db::graph::reachability::derivation_is_reachable(db, *d) @@ -45,9 +38,10 @@ pub(crate) async fn repair_missing_inputs( let mut kept: Vec<&str> = Vec::new(); let mut demoted_producers: Vec = Vec::new(); for path in missing_paths { - let Some(hash) = store_path_hash(path) else { + let Ok(stored) = StorePath::parse_stored(path) else { continue; }; + let hash = stored.hash(); match gradient_db::caches::demotion::diagnose_missing_input( db, @@ -124,18 +118,3 @@ pub(crate) async fn repair_missing_inputs( ); Ok(()) } - -#[cfg(test)] -mod tests { - use super::store_path_hash; - - #[test] - fn store_path_hash_extracts_32_char_hash() { - assert_eq!( - store_path_hash("/nix/store/g9y0fvqh2c991vjprgz9mvdm0zj7ggij-python3-static-3.13"), - Some("g9y0fvqh2c991vjprgz9mvdm0zj7ggij") - ); - assert_eq!(store_path_hash("not-a-store-path"), None); - assert_eq!(store_path_hash("/nix/store/"), None); - } -} diff --git a/backend/gradient-graph/src/testing.rs b/backend/gradient-graph/src/testing.rs index 50a9ca1bd..bf5479e1b 100644 --- a/backend/gradient-graph/src/testing.rs +++ b/backend/gradient-graph/src/testing.rs @@ -51,6 +51,7 @@ pub async fn ctx_with_crypt_file(db: DatabaseConnection, crypt_file: &str) -> Db events: gradient_types::EventBus::new(16), wake: Default::default(), memory: Default::default(), + store_path_mapping: Arc::new(gradient_entity::IdentityMapping), } } diff --git a/backend/gradient-graph/src/transition.rs b/backend/gradient-graph/src/transition.rs index 6c1ca1959..410ece892 100644 --- a/backend/gradient-graph/src/transition.rs +++ b/backend/gradient-graph/src/transition.rs @@ -311,18 +311,28 @@ async fn heal_corrupt_eval_cache( evaluation_id: EvaluationId, fingerprint: &str, ) -> Result { - let purged = EEvalCacheStore::delete_many() + let rows = EEvalCacheStore::find() .filter(CEvalCacheStore::Fingerprint.eq(fingerprint)) - .exec(&ctx.worker_db) - .await? - .rows_affected; - if purged == 0 { + .all(&ctx.worker_db) + .await?; + if rows.is_empty() { warn!(%evaluation_id, %fingerprint, "corrupt eval-cache recurred with no shared blob to purge; failing eval"); return Ok(false); } - if let Err(e) = ctx.storage.nar_storage.delete_eval_cache(fingerprint).await { - warn!(%fingerprint, error = %e, "failed to delete corrupt eval-cache object"); + EEvalCacheStore::delete_many() + .filter(CEvalCacheStore::Fingerprint.eq(fingerprint)) + .exec(&ctx.worker_db) + .await?; + for row in rows { + if let Err(e) = ctx + .storage + .nar_storage + .delete_eval_cache(&row.storage_path) + .await + { + warn!(key = %row.storage_path, error = %e, "failed to delete corrupt eval-cache object"); + } } requeue_evaluation(ctx, evaluation_id).await?; diff --git a/backend/gradient-graph/src/writer.rs b/backend/gradient-graph/src/writer.rs index c68eb5f5b..5b4f219eb 100644 --- a/backend/gradient-graph/src/writer.rs +++ b/backend/gradient-graph/src/writer.rs @@ -325,7 +325,7 @@ async fn claim_landed_paths( .iter() .zip(committed) .filter(|(_, outcome)| outcome.is_ok()) - .filter_map(|(commit, _)| StorePath::parse(&commit.store_path).ok()) + .filter_map(|(commit, _)| StorePath::parse_stored(&commit.store_path).ok()) .map(|path| path.hash().to_owned()) .collect(); if hashes.is_empty() { diff --git a/backend/gradient-migration/src/lib.rs b/backend/gradient-migration/src/lib.rs index 54fd4ce11..6292f92c5 100644 --- a/backend/gradient-migration/src/lib.rs +++ b/backend/gradient-migration/src/lib.rs @@ -21,6 +21,7 @@ mod m20261016_000000_metric_input_pnames; mod m20261016_000001_metric_pname_without_version; mod m20261016_000002_team_webhooks; mod m20261018_000000_worker_network_capability; +mod m20261019_000000_eval_cache_per_cache; pub struct UpgradeFloor { pub release: &'static str, @@ -59,6 +60,7 @@ impl MigratorTrait for Migrator { Box::new(m20261016_000001_metric_pname_without_version::Migration), Box::new(m20261016_000002_team_webhooks::Migration), Box::new(m20261018_000000_worker_network_capability::Migration), + Box::new(m20261019_000000_eval_cache_per_cache::Migration), ] } } diff --git a/backend/gradient-migration/src/m20261019_000000_eval_cache_per_cache.rs b/backend/gradient-migration/src/m20261019_000000_eval_cache_per_cache.rs new file mode 100644 index 000000000..a8b906d2e --- /dev/null +++ b/backend/gradient-migration/src/m20261019_000000_eval_cache_per_cache.rs @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use sea_orm_migration::prelude::*; +use sea_orm_migration::sea_orm::ConnectionTrait; + +const UP: &str = r#" +DELETE FROM public.eval_cache_store; + +ALTER TABLE public.eval_cache_store + ADD COLUMN cache uuid NOT NULL REFERENCES public.cache(id) ON DELETE CASCADE; + +DROP INDEX public."idx-eval_cache_store-fingerprint"; + +CREATE UNIQUE INDEX "idx-eval_cache_store-cache-fingerprint" ON public.eval_cache_store (cache, fingerprint); +"#; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager.get_connection().execute_unprepared(UP).await?; + Ok(()) + } + + async fn down(&self, _: &SchemaManager) -> Result<(), DbErr> { + Err(DbErr::Migration( + "m20261019_000000_eval_cache_per_cache is irreversible".into(), + )) + } +} diff --git a/backend/gradient-proto/Cargo.toml b/backend/gradient-proto/Cargo.toml index 43ea2ffaf..0abd57783 100644 --- a/backend/gradient-proto/Cargo.toml +++ b/backend/gradient-proto/Cargo.toml @@ -37,6 +37,7 @@ ractor = { workspace = true } reqwest = { workspace = true } sea-orm = { workspace = true } serde = { workspace = true } +serde_json = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["fs", "macros", "sync"] } tokio-tungstenite = { workspace = true } @@ -49,8 +50,9 @@ axum = { workspace = true, features = ["ws", "http1", "tokio"] } [dev-dependencies] gradient-wire = { workspace = true, features = ["testing"] } -serde_json = { workspace = true } +gradient-entity = { workspace = true, features = ["testing"] } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "sync", "net", "rt-multi-thread"] } gradient-fixtures = { workspace = true } +gradient-graph = { workspace = true, features = ["testing"] } gradient-test-support = { workspace = true } diff --git a/backend/gradient-proto/src/handler/cache.rs b/backend/gradient-proto/src/handler/cache.rs index af3ac5c4b..9c7341fb4 100644 --- a/backend/gradient-proto/src/handler/cache.rs +++ b/backend/gradient-proto/src/handler/cache.rs @@ -20,6 +20,7 @@ use tokio::sync::Semaphore; use tracing::warn; use super::SessionHandles; +use super::store_paths::project_mapper; fn stored_sizes( cached_paths: &[gradient_entity::cached_path::Model], @@ -346,7 +347,8 @@ async fn extend_with_upstream_results( db: &CacheDb, upstream_query: &Arc, project_id: ProjectId, - uncached_pairs: Vec<(String, String)>, + mapper: &StorePathMapper, + uncached_pairs: &[(String, String)], result: &mut Vec, ) { const UPSTREAM_WINDOW_MINUTES: i64 = 60; @@ -370,11 +372,30 @@ async fn extend_with_upstream_results( let id_to_url: HashMap<_, String> = endpoints.iter().map(|e| (e.id, e.url.clone())).collect(); - let batch = - gradient_core::upstream::probe_batch(endpoints, Arc::clone(upstream_query), uncached_pairs) - .await; - - result.extend(batch.found.into_values()); + let requested: HashMap<&str, &str> = uncached_pairs + .iter() + .map(|(hash, path)| (hash.as_str(), path.as_str())) + .collect(); + let targets = uncached_pairs + .iter() + .filter_map(|(_, path)| StorePath::parse_stored(path).ok()) + .collect(); + let batch = gradient_core::upstream::probe_batch( + endpoints, + Arc::clone(upstream_query), + mapper.clone(), + targets, + ) + .await; + + result.extend(batch.found.into_iter().map(|(hash, cp)| { + gradient_wire::types::CachedPath { + path: requested + .get(hash.as_str()) + .map_or(cp.path.clone(), |p| (*p).to_owned()), + ..cp + } + })); let mut by_url: HashMap = HashMap::new(); for (id, accum) in &batch.stats { @@ -401,6 +422,7 @@ async fn extend_with_gradient_proto_results( db: &CacheDb, config: &RuntimeConfig, project_id: ProjectId, + mapper: &StorePathMapper, uncached_pairs: &[(String, String)], result: &mut Vec, ) { @@ -421,10 +443,10 @@ async fn extend_with_gradient_proto_results( } let already: HashSet = result.iter().map(|c| c.path.clone()).collect(); - let want: Vec = uncached_pairs + let want: Vec = uncached_pairs .iter() - .map(|(_, p)| p.clone()) - .filter(|p| !already.contains(p)) + .filter(|(_, p)| !already.contains(p)) + .filter_map(|(_, p)| StorePath::parse_stored(p).ok()) .collect(); if want.is_empty() { return; @@ -435,9 +457,14 @@ async fn extend_with_gradient_proto_results( .api_key_enc .as_deref() .and_then(|enc| gradient_sources::decrypt_secret(&config.secrets.crypt_file, enc).ok()); - let found = - super::cache_consumer::pull_paths(&up.url, &up.remote_cache, api_key.as_deref(), &want) - .await; + let found = super::cache_consumer::pull_paths( + &up.url, + &up.remote_cache, + api_key.as_deref(), + mapper, + &want, + ) + .await; for cp in found { if !result.iter().any(|r| r.path == cp.path) { result.push(cp); @@ -448,7 +475,7 @@ async fn extend_with_gradient_proto_results( async fn query( handles: &SessionHandles, - project_id: Option, + project_id: ProjectId, paths: &[String], nar_sizes: &[Option], mode: gradient_wire::types::QueryMode, @@ -470,17 +497,10 @@ async fn query( )); } - let hash_path_pairs: Vec<(&str, &str)> = paths + let parsed = parse_stored_paths(paths); + let hash_path_pairs: Vec<(&str, &str)> = parsed .iter() - .filter_map(|p| { - let base = p.strip_prefix("/nix/store/").unwrap_or(p); - let hash = base.split('-').next()?; - if hash.len() == 32 { - Some((hash, p.as_str())) - } else { - None - } - }) + .map(|(stored, path)| (stored.hash(), *path)) .collect(); if hash_path_pairs.is_empty() { @@ -492,10 +512,8 @@ async fn query( let cached_path_rows = load_cached_path_rows(&handles.cache_db, &hashes).await?; let cached_map = stored_sizes(&cached_path_rows); - if matches!(mode, QueryMode::Push) - && let Some(oid) = project_id - { - ensure_push_signatures(&handles.cache_db, oid, &cached_path_rows).await; + if matches!(mode, QueryMode::Push) { + ensure_push_signatures(&handles.cache_db, project_id, &cached_path_rows).await; } let rows_by_hash: HashMap<&str, &gradient_entity::cached_path::Model> = cached_path_rows @@ -553,21 +571,24 @@ async fn query( .filter(|(h, _)| !resolved.contains(h)) .collect(); - if !uncached_pairs.is_empty() - && let Some(oid) = project_id - { + if uncached_pairs.is_empty() { + return Ok(result); + } + if let Some(mapper) = project_mapper(&handles.db, project_id).await { extend_with_upstream_results( &handles.cache_db, &handles.upstream_query, - oid, - uncached_pairs.clone(), + project_id, + &mapper, + &uncached_pairs, &mut result, ) .await; extend_with_gradient_proto_results( &handles.cache_db, &handles.db.config, - oid, + project_id, + &mapper, &uncached_pairs, &mut result, ) @@ -637,17 +658,10 @@ pub(super) async fn query_for_cache( return vec![]; } - let hash_path_pairs: Vec<(&str, &str)> = paths + let parsed = parse_stored_paths(paths); + let hash_path_pairs: Vec<(&str, &str)> = parsed .iter() - .filter_map(|p| { - let base = p.strip_prefix("/nix/store/").unwrap_or(p); - let hash = base.split('-').next()?; - if hash.len() == 32 { - Some((hash, p.as_str())) - } else { - None - } - }) + .map(|(stored, path)| (stored.hash(), *path)) .collect(); if hash_path_pairs.is_empty() { @@ -704,20 +718,22 @@ pub(super) async fn query_for_cache( result } -pub(super) async fn path_in_cache(db: &CacheDb, cache_id: CacheId, store_path: &str) -> bool { - let base = store_path.strip_prefix("/nix/store/").unwrap_or(store_path); - let Some(hash) = base.split('-').next() else { - return false; - }; - if hash.len() != 32 { - return false; - } - !hashes_in_cache(db, cache_id, &[hash]).await.is_empty() +pub(super) async fn path_in_cache(db: &CacheDb, cache_id: CacheId, path: &StorePath) -> bool { + !hashes_in_cache(db, cache_id, &[path.hash()]) + .await + .is_empty() +} + +fn parse_stored_paths(paths: &[String]) -> Vec<(StorePath, &str)> { + paths + .iter() + .filter_map(|path| Some((StorePath::parse_stored(path).ok()?, path.as_str()))) + .collect() } pub(super) async fn handle_cache_query( handles: &SessionHandles, - project_id: Option, + project_id: ProjectId, paths: &[String], nar_sizes: &[Option], mode: gradient_wire::types::QueryMode, @@ -770,7 +786,7 @@ mod tests { assert!( query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &[], QueryMode::Push, @@ -807,7 +823,7 @@ mod tests { let sizes = [Some(1024), Some(8 * 1024 * 1024 * 1024)]; let result = query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &sizes, QueryMode::Push, @@ -832,7 +848,7 @@ mod tests { assert!( query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &[], QueryMode::Pull, @@ -849,7 +865,7 @@ mod tests { let state = make_state(); let err = handle_cache_query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &["/nix/store/a-a".to_owned(), "/nix/store/b-b".to_owned()], &[], QueryMode::Pull, @@ -867,7 +883,7 @@ mod tests { assert!( query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &[], &[], QueryMode::Normal, @@ -880,7 +896,7 @@ mod tests { assert!( query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &[], &[], QueryMode::Push, @@ -893,7 +909,7 @@ mod tests { assert!( query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &[], &[], QueryMode::Pull, @@ -905,62 +921,13 @@ mod tests { ); } - #[tokio::test] - async fn cache_query_invalid_store_paths_skipped() { - let state = make_state(); - - let paths = vec![ - "not-a-store-path".to_string(), - "/nix/store/short-name".to_string(), - ]; - assert!( - query( - &SessionHandles::from_state(&state), - None, - &paths, - &[], - QueryMode::Normal, - false - ) - .await - .unwrap() - .is_empty() - ); - assert!( - query( - &SessionHandles::from_state(&state), - None, - &paths, - &[None; 2], - QueryMode::Push, - false - ) - .await - .unwrap() - .is_empty() - ); - assert!( - query( - &SessionHandles::from_state(&state), - None, - &paths, - &[], - QueryMode::Pull, - false - ) - .await - .unwrap() - .is_empty() - ); - } - #[tokio::test] async fn cache_query_normal_uncached_returns_empty() { let state = make_state(); let paths = vec!["/nix/store/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-hello".to_string()]; let result = query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &[], QueryMode::Normal, @@ -984,7 +951,7 @@ mod tests { let result = query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &[], QueryMode::Pull, @@ -1008,7 +975,7 @@ mod tests { ]; let result = query( &SessionHandles::from_state(&state), - None, + ProjectId::nil(), &paths, &[None; 2], QueryMode::Push, @@ -1033,30 +1000,6 @@ mod tests { assert!(returned_paths.contains(&paths[1].as_str())); } - #[tokio::test] - async fn cache_query_rejects_overlong_hash() { - let state = make_state(); - let paths = vec!["/nix/store/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-foo".to_string()]; - for mode in [QueryMode::Normal, QueryMode::Pull, QueryMode::Push] { - let sizes: &[Option] = if matches!(mode, QueryMode::Push) { - &[None] - } else { - &[] - }; - let result = query( - &SessionHandles::from_state(&state), - None, - &paths, - sizes, - mode, - false, - ) - .await - .unwrap(); - assert!(result.is_empty(), "33-char hash must be filtered out"); - } - } - #[tokio::test] async fn cache_scoped_query_rejects_push() { let state = make_state(); diff --git a/backend/gradient-proto/src/handler/cache_consumer.rs b/backend/gradient-proto/src/handler/cache_consumer.rs index b0d5be446..fcbd72163 100644 --- a/backend/gradient-proto/src/handler/cache_consumer.rs +++ b/backend/gradient-proto/src/handler/cache_consumer.rs @@ -6,6 +6,8 @@ use std::time::Duration; +use gradient_core::upstream::inbound_cached_path; +use gradient_types::{StorePath, StorePathMapper}; use gradient_wire::types::{CachedPath, GradientCapabilities, QueryMode}; use gradient_wire::messages::{ClientMessage, ServerMessage}; @@ -41,7 +43,8 @@ pub(crate) async fn pull_paths( base_url: &str, remote_cache: &str, api_key: Option<&str>, - paths: &[String], + mapper: &StorePathMapper, + paths: &[StorePath], ) -> Vec { let url = proto_ws_url(base_url, remote_cache); let mut socket = match gradient_wire::client::dial_with_auth(&url, api_key).await { @@ -67,7 +70,7 @@ pub(crate) async fn pull_paths( let query = ClientMessage::CacheQuery { job_id: uuid::Uuid::now_v7().to_string(), query_id: uuid::Uuid::now_v7().to_string(), - paths: paths.to_vec(), + paths: paths.iter().map(|p| mapper.outbound(p).full()).collect(), mode: QueryMode::Pull, nar_sizes: Vec::new(), external: false, @@ -77,9 +80,11 @@ pub(crate) async fn pull_paths( } match tokio::time::timeout(Duration::from_secs(30), reader.recv()).await { - Ok(Some(ServerMessage::CacheStatus { cached, .. })) => { - cached.into_iter().filter(|c| c.cached).collect() - } + Ok(Some(ServerMessage::CacheStatus { cached, .. })) => cached + .into_iter() + .filter(|c| c.cached) + .filter_map(|c| inbound_cached_path(mapper, c).ok()) + .collect(), _ => vec![], } } diff --git a/backend/gradient-proto/src/handler/cache_session.rs b/backend/gradient-proto/src/handler/cache_session.rs index 568fa3d96..3c04e1296 100644 --- a/backend/gradient-proto/src/handler/cache_session.rs +++ b/backend/gradient-proto/src/handler/cache_session.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; use gradient_types::ids::CacheId; +use gradient_types::{StorePath, StorePathMapper}; use tokio::sync::Semaphore; use tracing::{debug, info, warn}; @@ -16,7 +17,12 @@ use gradient_wire::messages::{ClientMessage, GradientCapabilities, ServerMessage use gradient_wire::session::handshake as handshake_fsm; use super::SessionHandles; -use super::socket::{HANDSHAKE_TIMEOUT, ProtoSocket, recv_client_msg, send_server_msg}; +use super::cache::{path_in_cache, query_for_cache}; +use super::nar_serve::{ServeSlot, serve_nar_request}; +use super::socket::{ + HANDSHAKE_TIMEOUT, ProtoSocket, ProtoWriter, recv_client_msg, send_server_msg, +}; +use super::store_paths::{inbound_or_refuse, outbound_message}; const CACHE_SESSION_IDLE_TIMEOUT_SECS: u64 = 120; @@ -45,6 +51,103 @@ fn readonly_capabilities() -> GradientCapabilities { } } +struct CacheSession { + handles: SessionHandles, + writer: ProtoWriter, + cache_id: CacheId, + mapper: StorePathMapper, + nar_serve_semaphore: Arc, + max_serves: usize, +} + +impl CacheSession { + fn new(handles: SessionHandles, writer: ProtoWriter, cache_id: CacheId) -> Self { + let max_serves = handles.db.config.nar.max_concurrent_serves.get(); + let mapper = handles.db.cache_mapper(cache_id); + Self { + handles, + writer, + cache_id, + mapper, + nar_serve_semaphore: Arc::new(Semaphore::new(max_serves)), + max_serves, + } + } + + fn serving_nothing(&self) -> bool { + self.nar_serve_semaphore.available_permits() == self.max_serves + } + + async fn serve(&self, msg: ClientMessage) -> bool { + let cache_id = self.cache_id; + if let Some(rejection) = rejection_for(&msg) { + warn!(%cache_id, variant = msg.variant_name(), "rejecting message on cache session"); + return send_server_msg(&self.writer, &ServerMessage::from(rejection)) + .await + .is_ok(); + } + let Some(msg) = inbound_or_refuse(&self.writer, &self.mapper, msg).await else { + return true; + }; + + match msg { + ClientMessage::CacheQuery { + query_id, + paths, + mode, + .. + } => { + let cached = query_for_cache(&self.handles, cache_id, &paths, mode).await; + let status = ServerMessage::CacheStatus { query_id, cached }; + send_server_msg(&self.writer, &outbound_message(status, &self.mapper)) + .await + .is_ok() + } + ClientMessage::NarRequest { job_id, paths } => self.serve_nars(job_id, paths).await, + _ => true, + } + } + + async fn serve_nars(&self, job_id: String, paths: Vec) -> bool { + let cache_id = self.cache_id; + for store_path in paths { + let in_cache = match StorePath::parse_stored(&store_path) { + Ok(path) => path_in_cache(&self.handles.cache_db, cache_id, &path).await, + Err(_) => false, + }; + if !in_cache { + debug!(%cache_id, %store_path, "skipping path not in this cache"); + continue; + } + let Some(slot) = ServeSlot::acquire( + Arc::clone(&self.nar_serve_semaphore), + Arc::clone(&self.handles.nar_downloads), + ) + .await + else { + warn!("nar serve semaphore closed"); + return false; + }; + + let handles = self.handles.clone(); + let writer = self.writer.clone(); + let mapper = self.mapper.clone(); + let job_id = job_id.clone(); + let shutdown = handles.db.shutdown.clone(); + shutdown.spawn(async move { + let _slot = slot; + if let Err(e) = + serve_nar_request(&handles, &writer, &mapper, &job_id, &store_path, 0, None) + .await + { + debug!(%store_path, error = %e, "cache NAR serve task failed"); + } + }); + } + true + } +} + pub async fn handle_cache_socket( mut socket: ProtoSocket, handles: SessionHandles, @@ -83,16 +186,15 @@ pub async fn handle_cache_socket( let send_chunk_timeout = Duration::from_secs(handles.db.config.nar.send_chunk_timeout_secs.get()); let (mut reader, writer) = socket.split(send_chunk_timeout, &handles.db.shutdown); - let max_serves = handles.db.config.nar.max_concurrent_serves.get(); - let nar_serve_semaphore = Arc::new(Semaphore::new(max_serves)); let idle = Duration::from_secs(CACHE_SESSION_IDLE_TIMEOUT_SECS); let cancel = handles.db.shutdown.token(); + let session = CacheSession::new(handles, writer, cache_id); loop { // The idle timeout is only enforced while no NAR transfer is in flight. A client quietly // receiving a large download must not be disconnected. let next = async { - if nar_serve_semaphore.available_permits() == max_serves { + if session.serving_nothing() { match tokio::time::timeout(idle, recv_client_msg(&mut reader)).await { Ok(m) => m, Err(_) => { @@ -112,71 +214,8 @@ pub async fn handle_cache_socket( }, }; - if let Some(rejection) = rejection_for(&msg) { - warn!(%cache_id, variant = msg.variant_name(), "rejecting message on cache session"); - if send_server_msg(&writer, &ServerMessage::from(rejection)) - .await - .is_err() - { - break; - } - continue; - } - - match msg { - ClientMessage::CacheQuery { - query_id, - paths, - mode, - .. - } => { - let cached = super::cache::query_for_cache(&handles, cache_id, &paths, mode).await; - if send_server_msg(&writer, &ServerMessage::CacheStatus { query_id, cached }) - .await - .is_err() - { - break; - } - } - ClientMessage::NarRequest { job_id, paths } => { - for store_path in paths { - if !super::cache::path_in_cache(&handles.cache_db, cache_id, &store_path).await - { - debug!(%cache_id, %store_path, "skipping path not in this cache"); - continue; - } - let Some(slot) = super::nar_serve::ServeSlot::acquire( - Arc::clone(&nar_serve_semaphore), - Arc::clone(&handles.nar_downloads), - ) - .await - else { - warn!("nar serve semaphore closed"); - return; - }; - - let handles = handles.clone(); - let writer = writer.clone(); - let job_id = job_id.clone(); - let shutdown = handles.db.shutdown.clone(); - shutdown.spawn(async move { - let _slot = slot; - if let Err(e) = super::nar_serve::serve_nar_request( - &handles, - &writer, - &job_id, - &store_path, - 0, - None, - ) - .await - { - debug!(%store_path, error = %e, "cache NAR serve task failed"); - } - }); - } - } - _ => {} + if !session.serve(msg).await { + break; } } @@ -238,4 +277,61 @@ mod tests { .is_some() ); } + + #[tokio::test] + async fn a_cache_socket_query_answers_with_the_paths_asked_for() { + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + use gradient_test_support::state::TestState; + + let cache_id = CacheId::nil(); + let stored = gradient_entity::cached_path::Model { + hash: EXAMPLE_STORED_HASH.into(), + file_hash: Some("sha256:f".into()), + confirmed: true, + ..Default::default() + }; + let cache_db = gradient_fixtures::mock::mock_db() + .append_query_results([vec![stored.clone()]]) + .append_query_results([vec![gradient_entity::cached_path_signature::Model { + cached_path: stored.id, + cache: cache_id, + ..Default::default() + }]]) + .append_query_results([vec![stored]]) + .into_connection(); + let state = TestState::new() + .cache_db(cache_db) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + let (writer, mut sent) = ProtoWriter::spy(Duration::from_secs(5)); + let session = CacheSession::new(SessionHandles::from_state(&state), writer, cache_id); + let query = |query_id: &str, path: String| ClientMessage::CacheQuery { + job_id: "job".into(), + query_id: query_id.into(), + paths: vec![path], + mode: QueryMode::Normal, + nar_sizes: vec![], + external: false, + }; + let decode = |bytes| { + gradient_wire::codec::from_bytes::( + bytes, + *gradient_wire::PROTO_VERSIONS.end(), + ) + .expect("a server message") + }; + let asked = format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-hello"); + + assert!(session.serve(query("q1", asked.clone())).await); + let ServerMessage::CacheStatus { cached, .. } = decode(sent.try_recv().unwrap()) else { + panic!("a CacheStatus") + }; + assert_eq!(cached[0].path, asked); + + assert!(session.serve(query("q2", "abc-hello".into())).await); + assert!(matches!( + decode(sent.try_recv().unwrap()), + ServerMessage::CacheStatus { query_id, cached } if query_id == "q2" && cached.is_empty() + )); + } } diff --git a/backend/gradient-proto/src/handler/eval_cache.rs b/backend/gradient-proto/src/handler/eval_cache.rs index 6d41ff0e7..d2db021c8 100644 --- a/backend/gradient-proto/src/handler/eval_cache.rs +++ b/backend/gradient-proto/src/handler/eval_cache.rs @@ -7,7 +7,7 @@ use gradient_db::WorkerDb; use gradient_entity::eval_cache_store; use gradient_storage::NarStore; -use gradient_types::ids::EvalCacheStoreId; +use gradient_types::ids::{CacheId, EvalCacheStoreId}; use gradient_types::*; use gradient_wire::types::EvalCachePullOutcome; use sea_orm::sea_query::OnConflict; @@ -17,8 +17,8 @@ use tracing::{debug, warn}; use super::socket::{BULK_CHUNK_SIZE, ProtoWriter, send_server_msg}; use gradient_wire::messages::{PRESIGN_TTL, ServerMessage}; -fn storage_key(fingerprint: &str) -> String { - format!("eval-cache/{fingerprint}") +pub(super) fn eval_cache_key(cache: CacheId, fingerprint: &str) -> String { + format!("eval-cache/{cache}/{fingerprint}") } fn should_accept_push(existing: Option, incoming: u64) -> bool { @@ -49,15 +49,15 @@ pub(super) async fn handle_eval_cache_pull( db: &WorkerDb, nar_storage: &NarStore, writer: &ProtoWriter, + readable: &[CacheId], job_id: String, fingerprint: String, ) { - let row = lookup_row(db, &fingerprint).await; - let key = storage_key(&fingerprint); + let row = newest_row(db, readable, &fingerprint).await; let presigned = match &row { - Some(_) => nar_storage - .presigned_eval_cache_get_url(&fingerprint, PRESIGN_TTL) + Some(row) => nar_storage + .presigned_eval_cache_get_url(&row.storage_path, PRESIGN_TTL) .await .unwrap_or_else(|e| { warn!(%fingerprint, error = %e, "presigned eval-cache GET failed; falling back to inline"); @@ -79,36 +79,81 @@ pub(super) async fn handle_eval_cache_pull( ) .await; - if inline && let Err(e) = stream_blob_inline(nar_storage, writer, &job_id, &key).await { + if inline + && let Some(row) = row + && let Err(e) = stream_blob_inline(nar_storage, writer, &job_id, &row.storage_path).await + { warn!(%fingerprint, error = %e, "inline eval-cache stream failed"); } } -pub(super) async fn accepts_push(db: &WorkerDb, fingerprint: &str, size_bytes: u64) -> bool { - let existing = lookup_row(db, fingerprint).await.map(|r| r.size_bytes); +pub(super) async fn accepts_push( + db: &WorkerDb, + writable: &[CacheId], + fingerprint: &str, + size_bytes: u64, +) -> bool { + let rows = rows_of(db, writable, fingerprint).await; + let existing = (rows.len() == writable.len()) + .then(|| rows.iter().map(|r| r.size_bytes).min()) + .flatten(); should_accept_push(existing, size_bytes) } -async fn lookup_row(db: &WorkerDb, fingerprint: &str) -> Option { - match EEvalCacheStore::find() +async fn newest_row( + db: &WorkerDb, + caches: &[CacheId], + fingerprint: &str, +) -> Option { + rows_of(db, caches, fingerprint) + .await + .into_iter() + .max_by_key(|r| r.updated_at) +} + +async fn rows_of( + db: &WorkerDb, + caches: &[CacheId], + fingerprint: &str, +) -> Vec { + if caches.is_empty() { + return Vec::new(); + } + EEvalCacheStore::find() .filter(eval_cache_store::Column::Fingerprint.eq(fingerprint)) - .one(db) + .filter(eval_cache_store::Column::Cache.is_in(caches.iter().copied())) + .all(db) .await - { - Ok(row) => row, - Err(e) => { + .unwrap_or_else(|e| { warn!(%fingerprint, error = %e, "eval_cache_store lookup failed"); - None - } - } + Vec::new() + }) } -pub(super) async fn record_eval_cache(db: &WorkerDb, fingerprint: &str, size_bytes: u64) { - upsert_eval_cache_row(db, fingerprint, &storage_key(fingerprint), size_bytes).await; +pub(super) async fn record_eval_cache( + db: &WorkerDb, + nar_storage: &NarStore, + writable: &[CacheId], + fingerprint: &str, + size_bytes: u64, +) { + let Some((&uploaded, copies)) = writable.split_first() else { + return; + }; + let uploaded_key = eval_cache_key(uploaded, fingerprint); + upsert_eval_cache_row(db, uploaded, fingerprint, &uploaded_key, size_bytes).await; + for &cache in copies { + let key = eval_cache_key(cache, fingerprint); + match nar_storage.copy_eval_cache(&uploaded_key, &key).await { + Ok(()) => upsert_eval_cache_row(db, cache, fingerprint, &key, size_bytes).await, + Err(e) => warn!(%fingerprint, %cache, error = %e, "eval-cache copy failed"), + } + } } async fn upsert_eval_cache_row( db: &WorkerDb, + cache: CacheId, fingerprint: &str, storage_path: &str, size_bytes: u64, @@ -116,6 +161,7 @@ async fn upsert_eval_cache_row( let now = gradient_types::now(); let model = MEvalCacheStore { id: EvalCacheStoreId::now_v7(), + cache, fingerprint: fingerprint.to_owned(), storage_path: storage_path.to_owned(), size_bytes: size_bytes as i64, @@ -125,20 +171,23 @@ async fn upsert_eval_cache_row( let result = EEvalCacheStore::insert(model.into_active_model()) .on_conflict( - OnConflict::column(eval_cache_store::Column::Fingerprint) - .update_columns([ - eval_cache_store::Column::StoragePath, - eval_cache_store::Column::SizeBytes, - eval_cache_store::Column::UpdatedAt, - ]) - .to_owned(), + OnConflict::columns([ + eval_cache_store::Column::Cache, + eval_cache_store::Column::Fingerprint, + ]) + .update_columns([ + eval_cache_store::Column::StoragePath, + eval_cache_store::Column::SizeBytes, + eval_cache_store::Column::UpdatedAt, + ]) + .to_owned(), ) - .exec(db) + .exec_without_returning(db) .await; if let Err(e) = result { - warn!(%fingerprint, error = %e, "eval_cache_store upsert failed"); + warn!(%fingerprint, %cache, error = %e, "eval_cache_store upsert failed"); } else { - debug!(%fingerprint, size_bytes, "eval_cache_store upserted"); + debug!(%fingerprint, %cache, size_bytes, "eval_cache_store upserted"); } } @@ -150,8 +199,7 @@ async fn stream_blob_inline( ) -> anyhow::Result<()> { use futures::StreamExt as _; - let fingerprint = key.strip_prefix("eval-cache/").unwrap_or(key); - let Some((_size, mut stream)) = nar_storage.get_eval_cache_stream(fingerprint).await? else { + let Some((_size, mut stream)) = nar_storage.get_eval_cache_stream(key).await? else { return Err(anyhow::anyhow!( "eval-cache blob {key} vanished before stream" )); @@ -214,12 +262,14 @@ fn stream_token(fingerprint: &str) -> String { #[cfg(test)] mod tests { use super::*; + use gradient_fixtures::mock::{exec_ok, mock_db, raw_statements}; fn row(size: i64) -> eval_cache_store::Model { eval_cache_store::Model { id: EvalCacheStoreId::now_v7(), + cache: CacheId::nil(), fingerprint: "fp".into(), - storage_path: storage_key("fp"), + storage_path: eval_cache_key(CacheId::nil(), "fp"), size_bytes: size, created_at: gradient_types::now(), updated_at: gradient_types::now(), @@ -257,4 +307,70 @@ mod tests { } ); } + + #[tokio::test] + async fn a_push_is_recorded_for_every_writable_cache() { + let (first, second) = (CacheId::now_v7(), CacheId::now_v7()); + let dir = tempfile::TempDir::new().unwrap(); + let store = NarStore::local(dir.path()).unwrap(); + store + .put_eval_cache(&eval_cache_key(first, "fp"), b"blob".to_vec()) + .await + .unwrap(); + let db = WorkerDb::new( + mock_db() + .append_exec_results([exec_ok(), exec_ok()]) + .into_connection(), + ); + + record_eval_cache(&db, &store, &[first, second], "fp", 10).await; + + assert_eq!( + store + .get_eval_cache(&eval_cache_key(second, "fp")) + .await + .unwrap(), + Some(b"blob".to_vec()) + ); + let inserts = raw_statements(db.into_transaction_log()); + assert_eq!(inserts.len(), 2, "{inserts:?}"); + for (insert, cache) in inserts.iter().zip([first, second]) { + assert!( + insert + .sql + .contains(r#"ON CONFLICT ("cache", "fingerprint")"#), + "{}", + insert.sql + ); + let values = format!("{:?}", insert.values); + assert!( + values.contains(&format!("{:?}", eval_cache_key(cache, "fp"))), + "{values}" + ); + } + } + + #[tokio::test] + async fn a_project_without_readable_caches_misses_without_a_lookup() { + let dir = tempfile::TempDir::new().unwrap(); + let store = NarStore::local(dir.path()).unwrap(); + let db = WorkerDb::new(mock_db().into_connection()); + let (writer, mut sent) = ProtoWriter::spy(std::time::Duration::from_secs(5)); + + handle_eval_cache_pull(&db, &store, &writer, &[], "j1".into(), "fp".into()).await; + + let reply = gradient_wire::codec::from_bytes::( + sent.try_recv().unwrap(), + *gradient_wire::PROTO_VERSIONS.end(), + ) + .unwrap(); + assert!(matches!( + reply, + ServerMessage::EvalCachePullResult { + outcome: EvalCachePullOutcome::Miss, + .. + } + )); + assert!(db.into_transaction_log().is_empty()); + } } diff --git a/backend/gradient-proto/src/handler/inbound.rs b/backend/gradient-proto/src/handler/inbound.rs index b0c9dbbcd..1ec9f08a6 100644 --- a/backend/gradient-proto/src/handler/inbound.rs +++ b/backend/gradient-proto/src/handler/inbound.rs @@ -10,8 +10,9 @@ use std::time::Instant; use gradient_types::events::{build, evaluation}; use gradient_types::ids::{DerivationBuildId, DispatchedJobId, EvaluationId, ProjectId, TaskId}; -use gradient_types::{BuildProgress, BuildProgressPhase, EvaluationProgress}; -use gradient_util::store_path::strip_nix_store_prefix; +use gradient_types::{ + BuildProgress, BuildProgressPhase, EvaluationProgress, StorePath, StorePathMapper, +}; use tokio::sync::Semaphore; use tracing::{Instrument as _, debug, debug_span, info, trace, warn}; @@ -20,8 +21,8 @@ use gradient_scheduler::jobs::{Assignment, PendingJob}; use gradient_scheduler::{ReportedTimeline, Scheduler}; use gradient_wire::Rejection; use gradient_wire::messages::{ - CACHE_QUERY_BUDGET, CandidateScore, ClientMessage, ClusterMembership, JobKind, JobUpdateKind, - ServerMessage, + BuildFailureKind, CACHE_QUERY_BUDGET, CandidateScore, ClientMessage, ClusterMembership, + JobKind, JobUpdateKind, ServerMessage, }; use gradient_wire::types::{ BuildProgressPhase as WireBuildProgressPhase, EvalProgress as WireEvalProgress, @@ -35,10 +36,15 @@ use super::eval_cache::handle_eval_cache_pull; use super::job_events::{JobEvent, JobEvents}; use super::log_lane::LogLane; use super::nar_serve::{ServeSlot, serve_nar_request}; +use super::refusal::refuse; use super::socket::{ ProtoWriter, job_offer_chunks, send_credentials_for_job, send_error, send_server_msg, }; +use super::store_paths::{ + inbound_or_refuse, outbound_message, outbound_offer, project_mapper, unmapped_project, +}; use super::upload::UploadSession; +use gradient_db::cache_paths::project_read_caches; #[derive(Clone)] pub(crate) struct ActiveJob { @@ -86,6 +92,17 @@ impl ActiveJobs { self.0.lock().get(job_id).map(|a| a.pending.project_id()) } + pub(crate) fn session_project(&self, job_id: &str) -> Option { + self.project(job_id).or_else(|| self.fetch_project(job_id)) + } + + async fn job_project(&self, scheduler: &Scheduler, job_id: &str) -> Option { + match self.session_project(job_id) { + Some(project) => Some(project), + None => scheduler.project_for_job(job_id).await, + } + } + pub(crate) fn build(&self, job_id: &str, task_index: u32) -> Option { match &self.0.lock().get(job_id)?.pending { PendingJob::Build(j) => j.job.builds.get(task_index as usize)?.build_id.parse().ok(), @@ -154,6 +171,9 @@ pub(super) struct InboundContext<'a> { impl<'a> InboundContext<'a> { pub async fn handle(&mut self, msg: ClientMessage, uploads: &mut UploadSession) -> bool { trace!(variant = msg.variant_name(), "received client message"); + let Some((msg, mapper)) = self.admit(msg).await else { + return true; + }; match msg { ClientMessage::InitConnection { .. } => { send_error(self.writer, 400, "unexpected InitConnection".into()).await; @@ -338,7 +358,9 @@ impl<'a> InboundContext<'a> { true } ClientMessage::NarRequest { job_id, paths } => { - self.on_nar_request(job_id, paths).await; + if let Some(mapper) = mapper { + self.on_nar_request(job_id, paths, mapper).await; + } true } ClientMessage::NarRequestResume { @@ -347,8 +369,16 @@ impl<'a> InboundContext<'a> { received_bytes, stream_token, } => { - self.on_nar_request_resume(job_id, store_path, received_bytes, stream_token) + if let Some(mapper) = mapper { + self.on_nar_request_resume( + job_id, + store_path, + received_bytes, + stream_token, + mapper, + ) .await; + } true } ClientMessage::EvalCachePull { @@ -486,6 +516,40 @@ impl<'a> InboundContext<'a> { } } + async fn admit(&self, msg: ClientMessage) -> Option<(ClientMessage, Option)> { + let project = match &msg { + ClientMessage::InputFetched { fetch_id, .. } => self.active.fetch_project(fetch_id), + ClientMessage::JobUpdate { job_id, .. } + | ClientMessage::JobFailed { job_id, .. } + | ClientMessage::NarRequest { job_id, .. } + | ClientMessage::NarRequestResume { job_id, .. } + | ClientMessage::EvalCachePull { job_id, .. } + | ClientMessage::ImportRequest { job_id, .. } + | ClientMessage::InputRequest { job_id, .. } + | ClientMessage::UploadRequest { job_id, .. } => { + let project = self.active.job_project(self.scheduler, job_id).await; + if project.is_none() { + refuse(self.writer, &msg, "the job does not resolve to a project").await; + return None; + } + project + } + _ => None, + }; + let Some(project) = project else { + return Some((msg, None)); + }; + let mapper = match self.handles.db.project_mapper(project).await { + Ok(mapper) => mapper, + Err(e) => { + refuse(self.writer, &msg, &unmapped_project(&e)).await; + return None; + } + }; + let mapped = inbound_or_refuse(self.writer, &mapper, msg).await?; + Some((mapped, Some(mapper))) + } + async fn on_job_update(&self, job_id: String, update: JobUpdateKind) { let starts_build = matches!(update, JobUpdateKind::Building { .. }); let event = JobEvent::Update { job_id, update }; @@ -544,10 +608,19 @@ impl<'a> InboundContext<'a> { } async fn on_eval_cache_pull(&mut self, job_id: String, fingerprint: String) { + let db = &self.handles.db.worker_db; + let readable = match self.active.job_project(self.scheduler, &job_id).await { + Some(project) => project_read_caches(db, project).await.unwrap_or_else(|e| { + warn!(peer_id = %self.peer_id, %job_id, error = %e, "readable caches lookup failed"); + Vec::new() + }), + None => Vec::new(), + }; handle_eval_cache_pull( - &self.handles.db.worker_db, + db, &self.handles.db.storage.nar_storage, self.writer, + &readable, job_id, fingerprint, ) @@ -638,8 +711,9 @@ impl<'a> InboundContext<'a> { async fn on_request_job_list(&mut self) -> bool { debug!(peer_id = %self.peer_id, "RequestJobList"); - let candidates = self.scheduler.get_job_candidates(self.peer_id).await; - self.send_job_list_chunks(candidates).await + let offer = self.scheduler.get_job_candidates(self.peer_id).await; + self.send_job_list_chunks(outbound_offer(offer, &self.handles.db).await) + .await } async fn send_job_list_chunks( @@ -692,6 +766,9 @@ impl<'a> InboundContext<'a> { cluster: Option, ) -> bool { tracing::Span::current().record("job_id", assignment.job_id()); + let Some(mapper) = project_mapper(&self.handles.db, assignment.project_id).await else { + return self.fail_unmappable(assignment.job_id()).await; + }; self.active.insert( assignment.job_id().to_owned(), ActiveJob { @@ -712,18 +789,34 @@ impl<'a> InboundContext<'a> { .await; let job_id = assignment.job_id().to_owned(); let assigned = debug_span!("assign_job", %job_id); - send_server_msg( - self.writer, - &ServerMessage::AssignJob { + let assign = ServerMessage::AssignJob { + job_id, + assignment_id: assignment.assignment_id().to_string(), + job: assignment.job, + cluster: cluster.map(Box::new), + }; + send_server_msg(self.writer, &outbound_message(assign, &mapper)) + .instrument(assigned) + .await + .is_ok() + } + + async fn fail_unmappable(&self, job_id: &str) -> bool { + let failed = self + .scheduler + .handle_job_failed( + self.peer_id, job_id, - assignment_id: assignment.assignment_id().to_string(), - job: assignment.job, - cluster: cluster.map(Box::new), - }, - ) - .instrument(assigned) - .await - .is_ok() + "the project's store paths cannot be mapped", + BuildFailureKind::Transient, + &[], + None, + ) + .await; + if let Err(e) = failed { + warn!(peer_id = %self.peer_id, %job_id, error = %e, "releasing an unmappable job failed"); + } + true } #[tracing::instrument(level = "debug", skip_all, fields(scores = scores.len(), is_final))] @@ -809,7 +902,12 @@ impl<'a> InboundContext<'a> { }); } - async fn on_nar_request(&mut self, job_id: String, paths: Vec) { + async fn on_nar_request( + &mut self, + job_id: String, + paths: Vec, + mapper: StorePathMapper, + ) { debug!(peer_id = %self.peer_id, %job_id, count = paths.len(), "NarRequest"); let shutdown = self.handles.db.shutdown.clone(); for store_path in paths { @@ -818,6 +916,7 @@ impl<'a> InboundContext<'a> { let permit = Arc::clone(self.nar_serve_semaphore); let peer_id = self.peer_id.to_owned(); let job_id = job_id.clone(); + let mapper = mapper.clone(); shutdown.spawn(async move { let server = Arc::clone(&handles.nar_downloads); let Some(_slot) = ServeSlot::acquire(permit, server).await else { @@ -825,7 +924,8 @@ impl<'a> InboundContext<'a> { }; if let Err(e) = - serve_nar_request(&handles, &writer, &job_id, &store_path, 0, None).await + serve_nar_request(&handles, &writer, &mapper, &job_id, &store_path, 0, None) + .await { warn!(%peer_id, %job_id, %store_path, error = %e, "NarRequest serve failed"); } @@ -839,6 +939,7 @@ impl<'a> InboundContext<'a> { store_path: String, received_bytes: u64, stream_token: String, + mapper: StorePathMapper, ) { debug!(peer_id = %self.peer_id, %job_id, %store_path, received_bytes, "NarRequestResume"); let handles = self.handles.clone(); @@ -855,6 +956,7 @@ impl<'a> InboundContext<'a> { if let Err(e) = serve_nar_request( &handles, &writer, + &mapper, &job_id, &store_path, received_bytes, @@ -895,8 +997,40 @@ impl RpcContext { } pub(super) fn serve(&self, msg: ClientMessage) -> Option { + if !matches!( + msg, + ClientMessage::CacheQuery { .. } | ClientMessage::QueryKnownDerivations { .. } + ) { + return Some(msg); + } let rpc = self.clone(); - match msg { + self.handles + .db + .shutdown + .spawn(async move { rpc.answer(msg).await }); + None + } + + pub(super) async fn job_project(&self, job_id: &str) -> Option { + self.active.job_project(&self.scheduler, job_id).await + } + + async fn answer(&self, msg: ClientMessage) { + let Some(project_id) = self.job_project(msg.job_id().unwrap_or_default()).await else { + refuse(&self.writer, &msg, "the job does not resolve to a project").await; + return; + }; + let mapper = match self.handles.db.project_mapper(project_id).await { + Ok(mapper) => mapper, + Err(e) => { + refuse(&self.writer, &msg, &unmapped_project(&e)).await; + return; + } + }; + let Some(msg) = inbound_or_refuse(&self.writer, &mapper, msg).await else { + return; + }; + let reply = match msg { ClientMessage::CacheQuery { job_id, query_id, @@ -905,24 +1039,27 @@ impl RpcContext { mode, external, } => { - self.handles.db.shutdown.spawn(async move { - rpc.on_cache_query(job_id, query_id, paths, nar_sizes, mode, external) - .await - }); - None + self.on_cache_query( + project_id, job_id, query_id, paths, nar_sizes, mode, external, + ) + .await } ClientMessage::QueryKnownDerivations { job_id, query_id, drv_paths, } => { - self.handles.db.shutdown.spawn(async move { - rpc.on_query_known_derivations(job_id, query_id, drv_paths) - .await - }); - None + self.on_query_known_derivations(job_id, query_id, drv_paths) + .await } - other => Some(other), + _ => return, + }; + let variant = reply.variant_name(); + if send_server_msg(&self.writer, &outbound_message(reply, &mapper)) + .await + .is_err() + { + debug!(peer_id = %self.peer_id, variant, "reply send failed; connection closing"); } } @@ -932,37 +1069,27 @@ impl RpcContext { )] async fn on_cache_query( &self, + project_id: ProjectId, job_id: String, query_id: String, paths: Vec, nar_sizes: Vec>, mode: gradient_wire::types::QueryMode, external: bool, - ) { + ) -> ServerMessage { debug!(peer_id = %self.peer_id, %job_id, %query_id, count = paths.len(), ?mode, external, "CacheQuery"); - let answer = async { - let project_id = match self - .active - .project(&job_id) - .or_else(|| self.active.fetch_project(&job_id)) - { - Some(project_id) => Some(project_id), - None => self.scheduler.project_for_job(&job_id).await, - }; - handle_cache_query( - &self.handles, - project_id, - &paths, - &nar_sizes, - mode, - external, - ) - .await - }; + let answer = handle_cache_query( + &self.handles, + project_id, + &paths, + &nar_sizes, + mode, + external, + ); // A DB error or an over-budget handler is indeterminate, never "absent". `CacheError` is // making the worker retry instead of failing the eval with `InputsUnavailable`. - let reply = match tokio::time::timeout(CACHE_QUERY_BUDGET, answer).await { + match tokio::time::timeout(CACHE_QUERY_BUDGET, answer).await { Ok(Ok(cached)) => { debug!(peer_id = %self.peer_id, %job_id, %query_id, entries = cached.len(), "CacheStatus"); ServerMessage::CacheStatus { query_id, cached } @@ -981,10 +1108,6 @@ impl RpcContext { message: "cache query exceeded server budget".to_string(), } } - }; - - if send_server_msg(&self.writer, &reply).await.is_err() { - debug!(peer_id = %self.peer_id, "CacheStatus/CacheError send failed; connection closing"); } } @@ -993,16 +1116,13 @@ impl RpcContext { job_id: String, query_id: String, drv_paths: Vec, - ) { + ) -> ServerMessage { debug!(peer_id = %self.peer_id, %job_id, %query_id, count = drv_paths.len(), "QueryKnownDerivations"); let hashes: Vec = drv_paths .iter() - .map(|p| strip_nix_store_prefix(p)) - .filter_map(|p| { - gradient_sources::parse_drv_hash_name(&p) - .ok() - .map(|(h, _)| h) - }) + .filter_map(|p| StorePath::parse_stored(p).ok()) + .filter(StorePath::is_derivation) + .map(|p| p.hash().to_owned()) .collect::>() .into_iter() .collect(); @@ -1023,15 +1143,7 @@ impl RpcContext { } }; debug!(peer_id = %self.peer_id, %job_id, %query_id, known = known.len(), "KnownDerivations"); - if send_server_msg( - &self.writer, - &ServerMessage::KnownDerivations { query_id, known }, - ) - .await - .is_err() - { - debug!(peer_id = %self.peer_id, "KnownDerivations send failed; connection closing"); - } + ServerMessage::KnownDerivations { query_id, known } } async fn on_worker_metrics(&self, metrics: WorkerMetrics) { @@ -1544,4 +1656,185 @@ mod assignment_response_tests { "{reply:?}" ); } + + #[tokio::test] + async fn a_cache_query_for_a_job_no_session_runs_is_refused_and_the_next_is_answered() { + use super::fixture::{JOB, TestSession, decode}; + use gradient_entity::EXAMPLE_ORIGINAL_HASH; + + let cache_db = gradient_fixtures::mock::mock_db() + .append_query_results([Vec::::new()]) + .into_connection(); + let state = TestState::new().cache_db(cache_db).build(); + let (mut session, mut sent, _admitted) = TestSession::new(&state).await; + let (ctx, _uploads) = session.split(); + let query = |job_id: &str, query_id: &str| ClientMessage::CacheQuery { + job_id: job_id.into(), + query_id: query_id.into(), + paths: vec![format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-a")], + nar_sizes: vec![None], + mode: gradient_wire::types::QueryMode::Normal, + external: false, + }; + let mut reply = async || { + let bytes = tokio::time::timeout(Duration::from_secs(5), sent.recv()) + .await + .expect("answered") + .expect("reply sent"); + decode(bytes) + }; + + ctx.rpc().serve(query("unknown", "q1")); + let refused = reply().await; + assert!( + matches!(&refused, ServerMessage::CacheError { query_id, .. } if query_id == "q1"), + "{refused:?}" + ); + + ctx.rpc().serve(query(JOB, "q2")); + let answered = reply().await; + assert!( + matches!(&answered, ServerMessage::CacheStatus { query_id, .. } if query_id == "q2"), + "{answered:?}" + ); + } + + #[tokio::test] + async fn a_malformed_path_refuses_only_its_message() { + use super::fixture::{JOB, TestSession, decode}; + + let state = test_state(empty_db()); + let (mut session, mut sent, _admitted) = TestSession::new(&state).await; + let (mut ctx, uploads) = session.split(); + let malformed = ClientMessage::JobUpdate { + job_id: JOB.into(), + assignment_id: DispatchedJobId::now_v7().to_string(), + update: JobUpdateKind::EvalResult { + derivations: vec![gradient_wire::types::DiscoveredDerivation { + attr: "a".into(), + drv_path: "abc-a.drv".into(), + ..Default::default() + }], + warnings: vec![], + errors: vec![], + }, + }; + + assert!(ctx.handle(malformed, uploads).await); + assert!(matches!( + decode(sent.try_recv().expect("a refusal")), + ServerMessage::Error { code: 400, .. } + )); + } + + #[tokio::test] + async fn a_message_whose_project_does_not_resolve_is_refused_in_the_reply_it_waits_for() { + use super::fixture::{JOB, TestSession, decode}; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, UnresolvableProjectMapping}; + use gradient_wire::types::ImportOutcome; + + let state = TestState::new() + .worker_db(empty_db()) + .store_path_mapping(Arc::new(UnresolvableProjectMapping)) + .build(); + let (mut session, mut sent, _admitted) = TestSession::new(&state).await; + let (mut ctx, uploads) = session.split(); + let asked = format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-hello.drv"); + + let import = ClientMessage::ImportRequest { + job_id: JOB.into(), + request_id: "r1".into(), + drv_paths: vec![asked.clone()], + }; + assert!(ctx.handle(import, uploads).await); + assert!(matches!( + decode(sent.try_recv().expect("a refusal")), + ServerMessage::ImportResult { + request_id, + outcome: ImportOutcome::Unknown { drv_path }, + .. + } if request_id == "r1" && drv_path == asked + )); + } + + #[tokio::test] + async fn a_cache_status_names_the_paths_the_worker_asked_for() { + use super::fixture::{JOB, TestSession, decode}; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + + let stored = gradient_entity::cached_path::Model { + hash: EXAMPLE_STORED_HASH.into(), + file_hash: Some("sha256:f".into()), + confirmed: true, + ..Default::default() + }; + let cache_db = gradient_fixtures::mock::mock_db() + .append_query_results([vec![stored]]) + .into_connection(); + let state = TestState::new() + .cache_db(cache_db) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + let (mut session, mut sent, _admitted) = TestSession::new(&state).await; + let (ctx, _uploads) = session.split(); + let asked = format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-hello"); + + ctx.rpc().serve(ClientMessage::CacheQuery { + job_id: JOB.into(), + query_id: "q1".into(), + paths: vec![asked.clone()], + nar_sizes: vec![None], + mode: gradient_wire::types::QueryMode::Normal, + external: false, + }); + let reply = tokio::time::timeout(Duration::from_secs(5), sent.recv()) + .await + .expect("answered") + .expect("reply sent"); + let ServerMessage::CacheStatus { cached, .. } = decode(reply) else { + panic!("a CacheStatus") + }; + assert_eq!(cached[0].path, asked); + } + + #[tokio::test] + async fn a_malformed_cache_query_path_is_answered_as_not_cached() { + use super::fixture::{JOB, TestSession, decode}; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + + let stored = gradient_entity::cached_path::Model { + hash: EXAMPLE_STORED_HASH.into(), + file_hash: Some("sha256:f".into()), + confirmed: true, + ..Default::default() + }; + let cache_db = gradient_fixtures::mock::mock_db() + .append_query_results([vec![stored]]) + .into_connection(); + let state = TestState::new() + .cache_db(cache_db) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + let (mut session, mut sent, _admitted) = TestSession::new(&state).await; + let (ctx, _uploads) = session.split(); + let asked = format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-hello"); + + ctx.rpc().serve(ClientMessage::CacheQuery { + job_id: JOB.into(), + query_id: "q1".into(), + paths: vec!["/nix/store/abc-malformed".into(), asked.clone()], + nar_sizes: vec![None, None], + mode: gradient_wire::types::QueryMode::Normal, + external: false, + }); + let reply = tokio::time::timeout(Duration::from_secs(5), sent.recv()) + .await + .expect("answered") + .expect("reply sent"); + let ServerMessage::CacheStatus { cached, .. } = decode(reply) else { + panic!("a CacheStatus") + }; + let cached: Vec = cached.into_iter().map(|hit| hit.path).collect(); + assert_eq!(cached, vec![asked]); + } } diff --git a/backend/gradient-proto/src/handler/mod.rs b/backend/gradient-proto/src/handler/mod.rs index 714bbdf17..10ab30292 100644 --- a/backend/gradient-proto/src/handler/mod.rs +++ b/backend/gradient-proto/src/handler/mod.rs @@ -17,10 +17,14 @@ mod job_events; mod log_lane; mod nar; mod nar_serve; +mod refusal; mod session; mod session_actor; mod sessions; mod socket; +mod store_paths; +#[cfg(test)] +mod stored_paths; mod tap; mod upload; diff --git a/backend/gradient-proto/src/handler/nar.rs b/backend/gradient-proto/src/handler/nar.rs index dfd5b7a89..63cff2203 100644 --- a/backend/gradient-proto/src/handler/nar.rs +++ b/backend/gradient-proto/src/handler/nar.rs @@ -4,15 +4,14 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use gradient_db::WorkerDb; use gradient_db::metrics::cache_traffic::CacheTraffic; -use gradient_entity::dispatched_job::{Column as CDispatchedJob, Entity as EDispatchedJob}; +use gradient_db::{DbContext, WorkerDb}; use gradient_graph::{Graph, NarCommit, SignTargets}; use gradient_types::events::cache::NarSigned; use gradient_types::ids::ProjectId; use gradient_types::*; -use sea_orm::{ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QueryOrder, Select}; -use tracing::warn; +use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; +use std::sync::Arc; pub(super) struct NarUploadRecord<'a> { pub file_hash: &'a str, @@ -28,13 +27,9 @@ pub(super) struct NarUploadRecord<'a> { pub(super) async fn record_nar_push_metric( db: &WorkerDb, cache_traffic: &CacheTraffic, - project_id: Option, + project_id: ProjectId, bytes: i64, ) -> anyhow::Result<()> { - let Some(project_id) = project_id else { - return Ok(()); - }; - let project_cache = EProjectCache::find() .filter(CProjectCache::Project.eq(project_id)) .one(db) @@ -49,22 +44,15 @@ pub(super) async fn record_nar_push_metric( pub(super) async fn mark_nar_stored( graph: &Graph, - events: &EventBus, - project_id: Option, + db: &DbContext, + project_id: ProjectId, store_path: &str, record: &NarUploadRecord<'_>, ) -> anyhow::Result<()> { - let hash_name = store_path.strip_prefix("/nix/store/").unwrap_or(store_path); - let hash = hash_name.split('-').next().unwrap_or(""); - - if hash.is_empty() { + let Ok(path) = StorePath::parse_stored(store_path) else { return Ok(()); - } - - let targets = match project_id { - Some(project_id) => SignTargets::ProjectCaches(project_id), - None => SignTargets::None, }; + let committed = graph .commit_nar(NarCommit { store_path: store_path.to_owned(), @@ -75,103 +63,54 @@ pub(super) async fn mark_nar_stored( references: record.references.to_vec(), deriver: record.deriver.map(str::to_owned), ca: record.ca.map(str::to_owned), - targets, + targets: SignTargets::ProjectCaches(project_id), confirmed: record.confirmed, built_by_worker: true, }) .await?; - for cache in committed.signed { - events.publish(NarSigned { - cache, - hash: hash.to_owned(), - }); + for event in signed_events(&db.store_path_mapping, &committed.signed, &path) { + db.events.publish(event); } Ok(()) } -fn dispatched_job_query(worker_id: &str, job_id: &str) -> Select { - EDispatchedJob::find() - .filter(CDispatchedJob::JobId.eq(job_id)) - .filter(CDispatchedJob::WorkerId.eq(worker_id)) - .order_by_desc(CDispatchedJob::DispatchedAt) -} - -/// `JobCompleted` is overtaking its own trailing `NarUploaded` because the control lane is drained -/// first. The scheduler has then evicted the job, and the dispatch row is the only source for the -/// cache claim. The worker filter is keeping a concurrent dispatch elsewhere from answering for -/// this upload. -pub(super) async fn project_for_dispatched_job( - db: &C, - worker_id: &str, - job_id: &str, -) -> Option { - dispatched_job_query(worker_id, job_id) - .one(db) - .await - .map_err(|e| warn!(%job_id, %worker_id, error = %e, "dispatched_job lookup for a late NAR failed")) - .ok() - .flatten() - .map(|row| row.project) +fn signed_events( + mapping: &Arc, + caches: &[CacheId], + path: &StorePath, +) -> Vec { + caches + .iter() + .map(|&cache| NarSigned { + cache, + hash: StorePathMapper::for_cache(Arc::clone(mapping), cache) + .outbound(path) + .hash() + .to_string(), + }) + .collect() } #[cfg(test)] mod tests { use super::*; - use gradient_fixtures::mock::mock_db; - use sea_orm::{DatabaseBackend, QueryTrait}; - use uuid::Uuid; - - const JOB: &str = "build:01a07af6-9d9f-7300-8aa1-fb82b0c18446"; - const WORKER: &str = "builder-1"; - - fn project() -> ProjectId { - ProjectId::new( - Uuid::parse_str("019e8549-eef3-7b92-96a2-f3842f585407").expect("project uuid"), - ) - } - - fn row() -> gradient_entity::dispatched_job::Model { - gradient_entity::dispatched_job::Model { - id: gradient_entity::ids::DispatchedJobId::now_v7(), - project: project(), - worker_id: WORKER.to_owned(), - job_id: Some(JOB.to_owned()), - ..Default::default() - } - } - - #[tokio::test] - async fn a_late_nar_still_resolves_its_project_from_the_assignment_row() { - let db = mock_db() - .append_query_results([vec![row()]]) - .into_connection(); - - assert_eq!( - project_for_dispatched_job(&db, WORKER, JOB).await, - Some(project()) - ); - } - - #[tokio::test] - async fn no_assignment_row_yields_no_project() { - let db = mock_db() - .append_query_results([Vec::::new()]) - .into_connection(); - - assert_eq!(project_for_dispatched_job(&db, WORKER, JOB).await, None); - } + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; #[test] - fn the_lookup_is_pinned_to_this_worker_and_the_newest_assignment() { - let sql = dispatched_job_query(WORKER, JOB) - .build(DatabaseBackend::Postgres) - .to_string(); - - assert!(sql.contains("\"job_id\" = "), "{sql}"); - assert!(sql.contains("\"worker_id\" = "), "{sql}"); - assert!(sql.contains("ORDER BY"), "{sql}"); - assert!(sql.contains("\"dispatched_at\" DESC"), "{sql}"); + fn a_signed_nar_is_announced_with_its_original_hash() { + let mapping: Arc = Arc::new(RotatedHashMapping); + let path = StorePath::parse_stored(&format!("/nix/store/{EXAMPLE_STORED_HASH}-hello")) + .expect("a stored path"); + let caches = vec![CacheId::now_v7(), CacheId::now_v7()]; + + let events = signed_events(&mapping, &caches, &path); + + assert_eq!(events.len(), 2); + for (event, cache) in events.iter().zip(&caches) { + assert_eq!(event.cache, *cache); + assert_eq!(event.hash, EXAMPLE_ORIGINAL_HASH); + } } } diff --git a/backend/gradient-proto/src/handler/nar_serve.rs b/backend/gradient-proto/src/handler/nar_serve.rs index 11a77c902..90424c779 100644 --- a/backend/gradient-proto/src/handler/nar_serve.rs +++ b/backend/gradient-proto/src/handler/nar_serve.rs @@ -12,7 +12,7 @@ use gradient_graph::Demotion; use gradient_storage::passthrough::{ PassthroughError, PassthroughLimits, PassthroughRequest, send_nar, }; -use gradient_types::RuntimeConfig; +use gradient_types::{RuntimeConfig, StorePath, StorePathMapper}; use gradient_util::telemetry::{GAUGES, Gauges}; use gradient_wire::messages::ServerMessage; use gradient_wire::session::frame::send_server_msg; @@ -75,6 +75,7 @@ impl Drop for ServeSlot { pub(super) async fn serve_nar_request( handles: &SessionHandles, writer: &ProtoWriter, + mapper: &StorePathMapper, job_id: &str, store_path: &str, resume_from: u64, @@ -86,14 +87,18 @@ pub(super) async fn serve_nar_request( chunk_read: Duration::from_secs(nar_cfg.send_chunk_timeout_secs.get()), chunk_bytes: usize::try_from(nar_cfg.chunk_bytes).unwrap_or(usize::MAX), }; - let Some(key) = store_hash(store_path) else { - let reason = format!("invalid store path: {store_path}"); - unavailable(writer, job_id, store_path, &reason).await; + let original = mapper + .outbound_text(store_path) + .unwrap_or_else(|_| store_path.to_owned()); + let Ok(path) = StorePath::parse_stored(store_path) else { + let reason = format!("invalid store path: {original}"); + unavailable(writer, job_id, &original, &reason).await; return Err(anyhow::anyhow!(reason)); }; + let key = path.hash(); let req = PassthroughRequest { job_id, - store_path, + store_path: &original, key, resume_from, client_token, @@ -101,7 +106,7 @@ pub(super) async fn serve_nar_request( match send_nar(&handles.db.storage.nar_storage, writer, req, limits).await { Ok(_) => Ok(()), Err(PassthroughError::NotFound(reason)) => { - unavailable(writer, job_id, store_path, &reason).await; + unavailable(writer, job_id, &original, &reason).await; invalidate_cached_path(handles, key, store_path).await; Err(anyhow::anyhow!(reason)) } @@ -159,15 +164,6 @@ async fn awaiting_upload(db: &WorkerDb, config: &RuntimeConfig, hash: &str) -> b } } -fn store_hash(store_path: &str) -> Option<&str> { - let hash = store_path - .strip_prefix("/nix/store/") - .unwrap_or(store_path) - .split('-') - .next()?; - (hash.len() == 32 && hash.bytes().all(|b| b.is_ascii_alphanumeric())).then_some(hash) -} - #[cfg(test)] mod serve_nar_tests { use super::*; @@ -178,6 +174,13 @@ mod serve_nar_tests { use gradient_wire::messages::ServerMessage; use tokio::sync::mpsc; + fn identity() -> StorePathMapper { + StorePathMapper::for_cache( + Arc::new(gradient_types::IdentityMapping), + gradient_types::ids::CacheId::nil(), + ) + } + fn spy_writer(timeout: Duration) -> (ProtoWriter, mpsc::Receiver) { ProtoWriter::spy(timeout) } @@ -206,6 +209,7 @@ mod serve_nar_tests { serve_nar_request( &SessionHandles::from_state(&state), &writer, + &identity(), "job-1", &store_path, 0, @@ -259,6 +263,7 @@ mod serve_nar_tests { let res = serve_nar_request( &SessionHandles::from_state(&state), &writer, + &identity(), "job-1", "/nix/store/zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-missing", 0, @@ -299,6 +304,7 @@ mod serve_nar_tests { serve_nar_request( &SessionHandles::from_state(&state), &writer, + &identity(), "job-1", &store_path, 0, @@ -341,8 +347,9 @@ mod serve_nar_tests { let res = serve_nar_request( &SessionHandles::from_state(&state), &writer, + &identity(), "job-1", - "not-a-store-path", + "not_a_store_path", 0, None, ) diff --git a/backend/gradient-proto/src/handler/refusal.rs b/backend/gradient-proto/src/handler/refusal.rs new file mode 100644 index 000000000..eb78dd36a --- /dev/null +++ b/backend/gradient-proto/src/handler/refusal.rs @@ -0,0 +1,154 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use gradient_wire::messages::{ClientMessage, ServerMessage}; +use gradient_wire::types::{EvalCachePullOutcome, ImportOutcome, InputOutcome, UploadOutcome}; +use tracing::warn; + +use super::socket::{ProtoWriter, send_server_msg}; + +pub(super) fn refusal(msg: &ClientMessage, reason: &str) -> Vec { + let message = format!("{} refused: {reason}", msg.variant_name()); + match msg { + ClientMessage::CacheQuery { query_id, .. } + | ClientMessage::QueryKnownDerivations { query_id, .. } => { + vec![ServerMessage::CacheError { + query_id: query_id.clone(), + message, + }] + } + ClientMessage::NarRequest { job_id, paths } => paths + .iter() + .map(|store_path| nar_unavailable(job_id, store_path, &message)) + .collect(), + ClientMessage::NarRequestResume { + job_id, store_path, .. + } => vec![nar_unavailable(job_id, store_path, &message)], + ClientMessage::UploadRequest { request_id, .. } + | ClientMessage::UploadFinished { request_id, .. } => { + vec![ServerMessage::UploadCommitted { + request_id: *request_id, + outcome: UploadOutcome::Rejected { reason: message }, + }] + } + ClientMessage::EvalCachePull { job_id, .. } => vec![ServerMessage::EvalCachePullResult { + job_id: job_id.clone(), + outcome: EvalCachePullOutcome::Miss, + }], + ClientMessage::ImportRequest { + job_id, + request_id, + drv_paths, + } => vec![ServerMessage::ImportResult { + job_id: job_id.clone(), + request_id: request_id.clone(), + outcome: ImportOutcome::Unknown { + drv_path: drv_paths.first().cloned().unwrap_or_default(), + }, + }], + ClientMessage::InputRequest { + job_id, request_id, .. + } => vec![ServerMessage::InputResult { + job_id: job_id.clone(), + request_id: request_id.clone(), + outcome: InputOutcome::Failed { message }, + }], + _ => vec![ServerMessage::Error { code: 400, message }], + } +} + +pub(super) async fn refuse(writer: &ProtoWriter, msg: &ClientMessage, reason: &str) { + warn!(variant = msg.variant_name(), job_id = msg.job_id(), %reason, "message refused"); + for reply in refusal(msg, reason) { + if send_server_msg(writer, &reply).await.is_err() { + return; + } + } +} + +fn nar_unavailable(job_id: &str, store_path: &str, reason: &str) -> ServerMessage { + ServerMessage::NarUnavailable { + job_id: job_id.to_owned(), + store_path: store_path.to_owned(), + reason: reason.to_owned(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_nar_request_is_refused_path_by_path() { + let msg = ClientMessage::NarRequest { + job_id: "j1".into(), + paths: vec!["/nix/store/a".into(), "/nix/store/b".into()], + }; + let refused: Vec = refusal(&msg, "no job") + .into_iter() + .map(|reply| match reply { + ServerMessage::NarUnavailable { store_path, .. } => store_path, + other => panic!("{other:?}"), + }) + .collect(); + assert_eq!(refused, vec!["/nix/store/a", "/nix/store/b"]); + } + + #[test] + fn each_request_is_refused_in_the_reply_its_sender_waits_for() { + let pull = ClientMessage::EvalCachePull { + job_id: "j1".into(), + fingerprint: "f".into(), + }; + assert!(matches!( + refusal(&pull, "no job").as_slice(), + [ServerMessage::EvalCachePullResult { + outcome: EvalCachePullOutcome::Miss, + .. + }] + )); + + let finished = ClientMessage::UploadFinished { + request_id: 7, + metadata: gradient_wire::types::UploadMetadata::EvalCache { size_bytes: 1 }, + }; + assert!(matches!( + refusal(&finished, "no job").as_slice(), + [ServerMessage::UploadCommitted { + request_id: 7, + outcome: UploadOutcome::Rejected { .. } + }] + )); + + let import = ClientMessage::ImportRequest { + job_id: "j1".into(), + request_id: "r".into(), + drv_paths: vec!["/nix/store/a.drv".into(), "/nix/store/b.drv".into()], + }; + assert!(matches!( + refusal(&import, "no job").as_slice(), + [ServerMessage::ImportResult { + request_id, + outcome: ImportOutcome::Unknown { drv_path }, + .. + }] if request_id == "r" && drv_path == "/nix/store/a.drv" + )); + + let input = ClientMessage::InputRequest { + job_id: "j1".into(), + request_id: "r".into(), + request: "{}".into(), + }; + assert!(matches!( + refusal(&input, "no job").as_slice(), + [ServerMessage::InputResult { + request_id, + outcome: InputOutcome::Failed { .. }, + .. + }] if request_id == "r" + )); + } +} diff --git a/backend/gradient-proto/src/handler/session_actor.rs b/backend/gradient-proto/src/handler/session_actor.rs index 77f37f262..631319937 100644 --- a/backend/gradient-proto/src/handler/session_actor.rs +++ b/backend/gradient-proto/src/handler/session_actor.rs @@ -34,6 +34,7 @@ use super::socket::{ InputFetch, ProtoSocket, ProtoWriter, job_offer_chunks, recv_client_msg, send_input_fetch, send_server_msg, }; +use super::store_paths::{outbound_message, outbound_offer, project_mapper}; use super::upload::{UploadSession, UploadTable, abandon_transfer}; use gradient_wire::Rejection; use gradient_wire::messages::{ClientMessage, GradientCapabilities, ServerMessage}; @@ -272,16 +273,19 @@ impl Actor for SessionActor { request_id, outcome, }) => { - if !st.active.contains(&job_id) { + let Some(project) = st.active.project(&job_id) else { debug!(peer_id = %st.peer_id, %job_id, "import result for a job this session no longer runs"); return Ok(()); - } + }; let msg = ServerMessage::ImportResult { job_id, request_id, outcome, }; + let Some(msg) = outbound_to(st, project, msg).await else { + return Ok(()); + }; if send_server_msg(&st.writer, &msg).await.is_err() { myself.stop(Some("write failed".into())); } @@ -291,16 +295,19 @@ impl Actor for SessionActor { request_id, outcome, }) => { - if !st.active.contains(&job_id) { + let Some(project) = st.active.project(&job_id) else { debug!(peer_id = %st.peer_id, %job_id, "input result for a job this session no longer runs"); return Ok(()); - } + }; let msg = ServerMessage::InputResult { job_id, request_id, outcome, }; + let Some(msg) = outbound_to(st, project, msg).await else { + return Ok(()); + }; if send_server_msg(&st.writer, &msg).await.is_err() { myself.stop(Some("write failed".into())); } @@ -443,17 +450,27 @@ impl Actor for SessionActor { } } +async fn outbound_to( + st: &SessionState, + project: ProjectId, + msg: ServerMessage, +) -> Option { + let mapper = project_mapper(&st.handles.db, project).await?; + Some(outbound_message(msg, &mapper)) +} + #[tracing::instrument(level = "debug", skip_all, fields(candidates = tracing::field::Empty))] async fn offer_jobs(st: &mut SessionState) -> bool { let offer = st.scheduler.get_new_job_candidates(&st.peer_id).await; - tracing::Span::current().record("candidates", offer.candidates.len()); st.offers_seen = offer.generation; - if offer.candidates.is_empty() { + let candidates = outbound_offer(offer, &st.handles.db).await; + tracing::Span::current().record("candidates", candidates.len()); + if candidates.is_empty() { return true; } - debug!(peer_id = %st.peer_id, count = offer.candidates.len(), "pushing job offer (delta)"); - for chunk in job_offer_chunks(&offer.candidates) { + debug!(peer_id = %st.peer_id, count = candidates.len(), "pushing job offer (delta)"); + for chunk in job_offer_chunks(&candidates) { if send_server_msg( &st.writer, &ServerMessage::JobOffer { @@ -569,7 +586,7 @@ fn split_uploads(st: &mut SessionState) -> (InboundContext<'_>, &mut UploadSessi } #[cfg(test)] -mod tests { +pub(in crate::handler) mod tests { use super::super::sessions::{Sessions, SessionsArgs, SessionsMsg}; use super::*; use futures::{SinkExt, StreamExt}; @@ -650,7 +667,8 @@ mod tests { } } - async fn connected_pair() -> (ProtoSocket, WebSocketStream>) { + pub(in crate::handler) async fn connected_pair() + -> (ProtoSocket, WebSocketStream>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let dial = tokio_tungstenite::connect_async(format!("ws://{addr}/proto")); diff --git a/backend/gradient-proto/src/handler/socket.rs b/backend/gradient-proto/src/handler/socket.rs index adfab8862..cdc62a0b9 100644 --- a/backend/gradient-proto/src/handler/socket.rs +++ b/backend/gradient-proto/src/handler/socket.rs @@ -13,6 +13,8 @@ use tracing::{debug, warn}; use gradient_scheduler::Scheduler; use gradient_wire::messages::ServerMessage; +use super::store_paths::{outbound_message, outbound_offer, project_mapper}; + pub use gradient_wire::session::frame::{ BULK_CHUNK_SIZE, HANDSHAKE_TIMEOUT, ProtoSocket, ProtoWriter, recv_client_msg, send_error, send_server_msg, @@ -25,11 +27,12 @@ pub(super) async fn push_pending_candidates( peer_id: &str, ) { let offer = scheduler.get_new_job_candidates(peer_id).await; - if offer.candidates.is_empty() { + let candidates = outbound_offer(offer, &scheduler.db).await; + if candidates.is_empty() { return; } - debug!(%peer_id, count = offer.candidates.len(), "pushing job offer (delta) after message processing"); - for chunk in job_offer_chunks(&offer.candidates) { + debug!(%peer_id, count = candidates.len(), "pushing job offer (delta) after message processing"); + for chunk in job_offer_chunks(&candidates) { let _ = send_server_msg( writer, &ServerMessage::JobOffer { @@ -71,22 +74,22 @@ pub(super) async fn send_input_fetch( worker_id: &str, fetch: InputFetch, ) -> Result<(), gradient_wire::session::frame::SendError> { + let Some(mapper) = project_mapper(db, fetch.project).await else { + return Ok(()); + }; let caps = scheduler.worker_gradient_caps(worker_id).await; let ssh_key = if sends_input_ssh_key(caps.as_ref(), &fetch.request) { project_ssh_key(db, fetch.project).await } else { None }; - send_server_msg( - writer, - &ServerMessage::FetchInput { - fetch_id: fetch.fetch_id, - project: fetch.project.to_string(), - request: fetch.request, - ssh_key, - }, - ) - .await + let msg = ServerMessage::FetchInput { + fetch_id: fetch.fetch_id, + project: fetch.project.to_string(), + request: fetch.request, + ssh_key, + }; + send_server_msg(writer, &outbound_message(msg, &mapper)).await } pub(super) struct InputFetch { diff --git a/backend/gradient-proto/src/handler/store_paths.rs b/backend/gradient-proto/src/handler/store_paths.rs new file mode 100644 index 000000000..16d73a214 --- /dev/null +++ b/backend/gradient-proto/src/handler/store_paths.rs @@ -0,0 +1,587 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::convert::Infallible; + +use gradient_db::DbContext; +use gradient_entity::StorePathError; +use gradient_scheduler::actor::Offer; +use gradient_types::ids::ProjectId; +use gradient_types::{MappingError, StorePathMapper}; +use gradient_wire::messages::{BuildFailureKind, ClientMessage, ServerMessage}; +use gradient_wire::types::{ + FlakeSource, ImportOutcome, InputOutcome, Job, JobCandidate, JobUpdateKind, UploadMetadata, + UploadObject, +}; +use serde_json::Value; + +use super::refusal::refuse; +use super::socket::ProtoWriter; + +struct TextMapper<'a, E> { + mapper: &'a StorePathMapper, + map: fn(&StorePathMapper, &str) -> Result, +} + +impl TextMapper<'_, E> { + fn path(&self, text: &mut String) -> Result<(), E> { + if !text.is_empty() { + *text = (self.map)(self.mapper, text)?; + } + Ok(()) + } + + fn paths(&self, texts: &mut [String]) -> Result<(), E> { + texts.iter_mut().try_for_each(|text| self.path(text)) + } + + fn json_store_path(&self, json: &mut String) -> Result<(), E> { + let Ok(mut value) = serde_json::from_str::(json) else { + return Ok(()); + }; + let Some(Value::String(path)) = value.get_mut("storePath") else { + return Ok(()); + }; + self.path(path)?; + *json = value.to_string(); + Ok(()) + } +} + +fn inbound(mapper: &StorePathMapper) -> TextMapper<'_, StorePathError> { + TextMapper { + mapper, + map: StorePathMapper::inbound_text, + } +} + +fn outbound(mapper: &StorePathMapper) -> TextMapper<'_, Infallible> { + TextMapper { + mapper, + map: |mapper, text| { + Ok(mapper + .outbound_text(text) + .unwrap_or_else(|_| text.to_owned())) + }, + } +} + +pub(super) fn inbound_message( + mut msg: ClientMessage, + mapper: &StorePathMapper, +) -> Result { + let paths = inbound(mapper); + match &mut msg { + ClientMessage::JobUpdate { update, .. } => inbound_update(update, &paths)?, + ClientMessage::JobFailed { + kind: BuildFailureKind::CorruptEvalCache, + .. + } => {} + ClientMessage::JobFailed { missing_paths, .. } => paths.paths(missing_paths)?, + ClientMessage::CacheQuery { + paths: queried, + nar_sizes, + .. + } => inbound_queried(queried, nar_sizes, &paths), + ClientMessage::QueryKnownDerivations { drv_paths, .. } + | ClientMessage::ImportRequest { drv_paths, .. } => paths.paths(drv_paths)?, + ClientMessage::NarRequest { + paths: requested, .. + } => paths.paths(requested)?, + ClientMessage::NarRequestResume { store_path, .. } + | ClientMessage::UploadRequest { + object: UploadObject::Nar { store_path }, + .. + } => paths.path(store_path)?, + ClientMessage::InputRequest { request, .. } => paths.json_store_path(request)?, + ClientMessage::InputFetched { + outcome: InputOutcome::Fetched { response }, + .. + } => paths.json_store_path(response)?, + _ => {} + } + Ok(msg) +} + +fn inbound_queried( + queried: &mut Vec, + nar_sizes: &mut Vec>, + paths: &TextMapper, +) { + let mapped: Vec = queried + .iter_mut() + .map(|path| paths.path(path).is_ok()) + .collect(); + if nar_sizes.len() == mapped.len() { + retain_mapped(nar_sizes, &mapped); + } + retain_mapped(queried, &mapped); +} + +fn retain_mapped(items: &mut Vec, mapped: &[bool]) { + let mut mapped = mapped.iter(); + items.retain(|_| mapped.next().copied().unwrap_or(false)); +} + +fn inbound_update( + update: &mut JobUpdateKind, + paths: &TextMapper, +) -> Result<(), StorePathError> { + match update { + JobUpdateKind::EvalResult { derivations, .. } => { + for derivation in derivations { + paths.path(&mut derivation.drv_path)?; + for output in &mut derivation.outputs { + paths.path(&mut output.path)?; + } + paths.paths(&mut derivation.dependencies)?; + paths.paths(&mut derivation.input_sources)?; + } + } + JobUpdateKind::EvalStats(report) => { + for node in &mut report.flake_nodes { + if let Some(drv_path) = &mut node.drv_path { + paths.path(drv_path)?; + } + } + } + JobUpdateKind::FetchResult { + flake_source: Some(flake_source), + } => paths.path(flake_source)?, + JobUpdateKind::BuildOutput { outputs, .. } => { + for output in outputs { + paths.path(&mut output.store_path)?; + paths.path(&mut output.hash)?; + for product in &mut output.products { + product.path = paths.mapper.inbound_if_store_path(&product.path); + } + } + } + _ => {} + } + Ok(()) +} + +pub(super) async fn inbound_or_refuse( + writer: &ProtoWriter, + mapper: &StorePathMapper, + msg: ClientMessage, +) -> Option { + match inbound_message(msg.clone(), mapper) { + Ok(mapped) => Some(mapped), + Err(e) => { + refuse(writer, &msg, &e.to_string()).await; + None + } + } +} + +pub(super) fn inbound_upload_metadata( + mut metadata: UploadMetadata, + mapper: &StorePathMapper, +) -> Result { + if let UploadMetadata::Nar(meta) = &mut metadata { + let paths = inbound(mapper); + paths.paths(&mut meta.references)?; + if let Some(deriver) = &mut meta.deriver { + paths.path(deriver)?; + } + } + Ok(metadata) +} + +pub(super) fn outbound_message(mut msg: ServerMessage, mapper: &StorePathMapper) -> ServerMessage { + let paths = outbound(mapper); + let Ok(()) = outbound_fields(&mut msg, &paths); + msg +} + +fn outbound_fields( + msg: &mut ServerMessage, + paths: &TextMapper, +) -> Result<(), Infallible> { + match msg { + ServerMessage::AssignJob { job, .. } => match job { + Job::Build(build) => { + for spec in &mut build.builds { + paths.path(&mut spec.drv_path)?; + for output in &mut spec.outputs { + paths.path(&mut output.path)?; + } + } + } + Job::Flake(flake) => { + if let FlakeSource::Cached { store_path } = &mut flake.source { + paths.path(store_path)?; + } + } + }, + ServerMessage::JobOffer { candidates } | ServerMessage::JobListChunk { candidates, .. } => { + for candidate in candidates { + outbound_candidate_fields(candidate, paths)?; + } + } + ServerMessage::CacheStatus { cached, .. } => { + for entry in cached { + paths.path(&mut entry.path)?; + if let Some(references) = &mut entry.references { + paths.paths(references)?; + } + if let Some(deriver) = &mut entry.deriver { + paths.path(deriver)?; + } + } + } + ServerMessage::KnownDerivations { known, .. } => paths.paths(known)?, + ServerMessage::ImportResult { + outcome: ImportOutcome::Failed { drv_path, .. } | ImportOutcome::Unknown { drv_path }, + .. + } => paths.path(drv_path)?, + ServerMessage::FetchInput { request, .. } => paths.json_store_path(request)?, + ServerMessage::InputResult { + outcome: InputOutcome::Fetched { response }, + .. + } => paths.json_store_path(response)?, + _ => {} + } + Ok(()) +} + +pub(super) fn outbound_candidate( + mut candidate: JobCandidate, + mapper: &StorePathMapper, +) -> JobCandidate { + let Ok(()) = outbound_candidate_fields(&mut candidate, &outbound(mapper)); + candidate +} + +fn outbound_candidate_fields( + candidate: &mut JobCandidate, + paths: &TextMapper, +) -> Result<(), Infallible> { + for required in &mut candidate.required_paths { + paths.path(&mut required.path)?; + } + paths.paths(&mut candidate.drv_paths)?; + paths.paths(&mut candidate.output_paths) +} + +pub(super) async fn outbound_offer(offer: Offer, db: &DbContext) -> Vec { + let mut candidates = Vec::with_capacity(offer.candidates.len()); + for candidate in offer.candidates { + let Some(&project) = offer.projects.get(&candidate.job_id) else { + continue; + }; + if let Some(mapper) = project_mapper(db, project).await { + candidates.push(outbound_candidate(candidate, &mapper)); + } + } + candidates +} + +pub(super) fn unmapped_project(error: &MappingError) -> String { + format!("the project's store paths cannot be mapped: {error}") +} + +pub(super) async fn project_mapper(db: &DbContext, project: ProjectId) -> Option { + db.project_mapper(project) + .await + .inspect_err(|error| tracing::warn!(%project, "{}", unmapped_project(error))) + .ok() +} + +#[cfg(test)] +mod tests { + use super::*; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + use gradient_types::StorePathMapper; + use gradient_types::ids::CacheId; + use gradient_wire::messages::ServerMessage; + use gradient_wire::types::{ + BuildJob, BuildOutput, BuildProduct, BuildSpec, DerivationOutput, DiscoveredDerivation, + Job, JobCandidate, JobUpdateKind, NarUploadMetadata, RequiredPath, UploadMetadata, + }; + use std::sync::Arc; + + const O: &str = EXAMPLE_ORIGINAL_HASH; + const S: &str = EXAMPLE_STORED_HASH; + + fn mapper() -> StorePathMapper { + StorePathMapper::for_cache(Arc::new(RotatedHashMapping), CacheId::nil()) + } + + #[test] + fn every_path_of_an_eval_result_is_mapped_inbound() { + let derivation = DiscoveredDerivation { + attr: "a".into(), + drv_path: format!("/nix/store/{O}-a.drv"), + outputs: vec![ + DerivationOutput { + name: "out".into(), + path: format!("/nix/store/{O}-a"), + }, + DerivationOutput { + name: "dev".into(), + path: String::new(), + }, + ], + dependencies: vec![format!("/nix/store/{O}-b.drv")], + input_sources: vec![format!("/nix/store/{O}-src")], + ..Default::default() + }; + let msg = ClientMessage::JobUpdate { + job_id: "j1".into(), + assignment_id: "a1".into(), + update: JobUpdateKind::EvalResult { + derivations: vec![derivation], + warnings: vec![], + errors: vec![], + }, + }; + let ClientMessage::JobUpdate { + update: JobUpdateKind::EvalResult { derivations, .. }, + .. + } = inbound_message(msg, &mapper()).unwrap() + else { + panic!("an EvalResult stays an EvalResult") + }; + assert_eq!(derivations[0].attr, "a"); + assert_eq!(derivations[0].drv_path, format!("/nix/store/{S}-a.drv")); + assert_eq!(derivations[0].outputs[0].path, format!("/nix/store/{S}-a")); + assert_eq!(derivations[0].outputs[1].path, ""); + assert_eq!( + derivations[0].dependencies, + vec![format!("/nix/store/{S}-b.drv")] + ); + assert_eq!( + derivations[0].input_sources, + vec![format!("/nix/store/{S}-src")] + ); + } + + #[test] + fn a_build_product_keeps_its_path_inside_the_output() { + let msg = ClientMessage::JobUpdate { + job_id: "j1".into(), + assignment_id: "a1".into(), + update: JobUpdateKind::BuildOutput { + build_id: "b1".into(), + outputs: vec![BuildOutput { + name: "out".into(), + store_path: format!("/nix/store/{O}-a"), + hash: O.into(), + nar_size: None, + nar_hash: None, + products: vec![BuildProduct { + file_type: "doc".into(), + subtype: "pdf".into(), + name: "doc.pdf".into(), + path: format!("/nix/store/{O}-a/doc.pdf"), + size: None, + }], + }], + metrics: None, + substituted: false, + }, + }; + let ClientMessage::JobUpdate { + update: JobUpdateKind::BuildOutput { outputs, .. }, + .. + } = inbound_message(msg, &mapper()).unwrap() + else { + panic!("a BuildOutput stays a BuildOutput") + }; + assert_eq!(outputs[0].store_path, format!("/nix/store/{S}-a")); + assert_eq!(outputs[0].hash, S); + assert_eq!( + outputs[0].products[0].path, + format!("/nix/store/{S}-a/doc.pdf") + ); + } + + #[test] + fn a_build_product_outside_the_store_is_admitted_unchanged() { + let msg = ClientMessage::JobUpdate { + job_id: "j1".into(), + assignment_id: "a1".into(), + update: JobUpdateKind::BuildOutput { + build_id: "b1".into(), + outputs: vec![BuildOutput { + name: "out".into(), + store_path: format!("/nix/store/{O}-a"), + hash: O.into(), + nar_size: None, + nar_hash: None, + products: vec![BuildProduct { + file_type: "doc".into(), + subtype: "pdf".into(), + name: "manual.pdf".into(), + path: "share/doc/manual.pdf".into(), + size: None, + }], + }], + metrics: None, + substituted: false, + }, + }; + let ClientMessage::JobUpdate { + update: JobUpdateKind::BuildOutput { outputs, .. }, + .. + } = inbound_message(msg, &mapper()).unwrap() + else { + panic!("a BuildOutput stays a BuildOutput") + }; + assert_eq!(outputs[0].products[0].path, "share/doc/manual.pdf"); + } + + #[test] + fn an_input_request_maps_its_store_path_key() { + let request = |hash: &str| { + format!( + r#"{{"builtin":"fetchTree","args":{{"type":"tarball","url":"https://example.com/x.tar.gz"}},"storePath":"/nix/store/{hash}-input"}}"# + ) + }; + let msg = ClientMessage::InputRequest { + job_id: "j1".into(), + request_id: "r1".into(), + request: request(O), + }; + let ClientMessage::InputRequest { + request: mapped, .. + } = inbound_message(msg, &mapper()).unwrap() + else { + panic!("an InputRequest stays an InputRequest") + }; + let json = |text: &str| serde_json::from_str::(text).unwrap(); + assert_eq!(json(&mapped), json(&request(S))); + } + + #[test] + fn upload_metadata_maps_its_references_and_deriver() { + let metadata = UploadMetadata::Nar(Box::new(NarUploadMetadata { + file_hash: "sha256:f".into(), + file_size: 1, + nar_size: 1, + nar_hash: "sha256:n".into(), + references: vec![format!("{O}-dep")], + deriver: Some(format!("{O}-a.drv")), + ca: None, + multipart: None, + })); + let UploadMetadata::Nar(meta) = inbound_upload_metadata(metadata, &mapper()).unwrap() + else { + panic!("NAR metadata stays NAR metadata") + }; + assert_eq!(meta.references, vec![format!("{S}-dep")]); + assert_eq!(meta.deriver, Some(format!("{S}-a.drv"))); + assert_eq!(meta.file_hash, "sha256:f"); + } + + #[test] + fn a_short_hash_is_malformed() { + let msg = ClientMessage::NarRequest { + job_id: "j1".into(), + paths: vec!["/nix/store/abc-a".into()], + }; + assert!(inbound_message(msg, &mapper()).is_err()); + } + + #[test] + fn a_cache_query_drops_only_its_malformed_paths() { + let msg = ClientMessage::CacheQuery { + job_id: "j1".into(), + query_id: "q1".into(), + paths: vec!["/nix/store/abc-a".into(), format!("/nix/store/{O}-b")], + mode: gradient_wire::types::QueryMode::Push, + nar_sizes: vec![Some(1), Some(2)], + external: false, + }; + let ClientMessage::CacheQuery { + paths, nar_sizes, .. + } = inbound_message(msg, &mapper()).unwrap() + else { + panic!("a CacheQuery stays a CacheQuery") + }; + assert_eq!(paths, vec![format!("/nix/store/{S}-b")]); + assert_eq!(nar_sizes, vec![Some(2)]); + } + + #[test] + fn a_corrupt_eval_cache_failure_keeps_its_fingerprint() { + let fingerprint = "a1b2c3d4e5f60718293a4b5c6d7e8f90a1b2c3d4e5f60718293a4b5c6d7e8f90"; + let msg = ClientMessage::JobFailed { + job_id: "j1".into(), + assignment_id: "a1".into(), + error: "corrupt".into(), + kind: BuildFailureKind::CorruptEvalCache, + missing_paths: vec![fingerprint.into()], + spans: Vec::new(), + elapsed_ms: 0, + metrics: None, + }; + let ClientMessage::JobFailed { missing_paths, .. } = + inbound_message(msg, &mapper()).unwrap() + else { + unreachable!() + }; + assert_eq!(missing_paths, vec![fingerprint]); + } + + #[test] + fn an_assigned_build_names_original_paths() { + let msg = ServerMessage::AssignJob { + job_id: "j1".into(), + assignment_id: "a1".into(), + job: Job::Build(BuildJob { + builds: vec![BuildSpec { + build_id: "b1".into(), + drv_path: format!("{S}-a.drv"), + kind: Default::default(), + is_fixed_output: false, + outputs: vec![DerivationOutput { + name: "out".into(), + path: format!("/nix/store/{S}-a"), + }], + timeout_secs: None, + max_silent_secs: None, + build_cores: None, + }], + requirement: Default::default(), + }), + cluster: None, + }; + let ServerMessage::AssignJob { + job: Job::Build(job), + .. + } = outbound_message(msg, &mapper()) + else { + panic!("an AssignJob stays an AssignJob") + }; + assert_eq!(job.builds[0].drv_path, format!("{O}-a.drv")); + assert_eq!(job.builds[0].outputs[0].path, format!("/nix/store/{O}-a")); + } + + #[test] + fn an_offered_candidate_names_original_paths() { + let candidate = JobCandidate { + job_id: "j1".into(), + required_paths: vec![RequiredPath { + path: format!("/nix/store/{S}-dep"), + cache_info: None, + }], + drv_paths: vec![format!("/nix/store/{S}-a.drv")], + output_paths: vec![format!("/nix/store/{S}-a")], + requirement: None, + }; + let candidate = outbound_candidate(candidate, &mapper()); + assert_eq!( + candidate.required_paths[0].path, + format!("/nix/store/{O}-dep") + ); + assert_eq!(candidate.drv_paths, vec![format!("/nix/store/{O}-a.drv")]); + assert_eq!(candidate.output_paths, vec![format!("/nix/store/{O}-a")]); + } +} diff --git a/backend/gradient-proto/src/handler/stored_paths.rs b/backend/gradient-proto/src/handler/stored_paths.rs new file mode 100644 index 000000000..1f28a984e --- /dev/null +++ b/backend/gradient-proto/src/handler/stored_paths.rs @@ -0,0 +1,547 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::collections::HashSet; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use futures::{SinkExt, StreamExt}; +use gradient_entity::RotatedHashMapping; +use gradient_entity::evaluation::EvaluationStatus; +use gradient_fixtures::database::TestDatabase; +use gradient_fixtures::shape::shape; +use gradient_graph::{Graph, testing}; +use gradient_scheduler::Scheduler; +use gradient_scheduler::jobs::eval_job_key; +use gradient_test_support::state::TestState; +use gradient_types::events::EventRx; +use gradient_types::ids::{CacheId, EvaluationId}; +use gradient_types::{BuildArgs, MappingContext, StoreHash, StorePathMapping}; +use gradient_wire::PROTO_VERSIONS; +use gradient_wire::codec::{from_bytes, to_bytes}; +use gradient_wire::messages::{ClientMessage, GradientCapabilities, ServerMessage}; +use gradient_wire::types::{ + BuildOutput, CandidateScore, DerivationOutput, DiscoveredDerivation, GrantTarget, Job, + JobCandidate, JobKind, JobUpdateKind, NarUploadMetadata, QueryMode, UploadMetadata, + UploadObject, UploadOutcome, +}; +use ractor::Actor; +use ractor::rpc::CallResult; +use tokio::net::TcpStream; +use tokio::sync::broadcast::error::RecvError; +use tokio::sync::oneshot; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use super::SessionHandles; +use super::session_actor::SessionArgs; +use super::session_actor::tests::connected_pair; +use super::sessions::{Sessions, SessionsArgs, SessionsMsg}; + +const A_DRV: &str = "assjydybg00pv49ff2vsw2mz7p19571x"; +const A_OUT: &str = "amrz9b2yg4cd47yqcm9sjpwqlzqi15az"; +const B_DRV: &str = "0z4fl7banmnkg6a4bvpa2gsa0nzqgga4"; +const B_OUT: &str = "wsy3j3hj2sgar4m99y3h8y1dp9pvzyk8"; +const SRC: &str = "i80q81hlgshr8896r6b1w947j8klrzfa"; +const ORIGINALS: [&str; 5] = [A_DRV, A_OUT, B_DRV, B_OUT, SRC]; +const WAIT: Duration = Duration::from_secs(60); + +fn rotate(original: &str) -> String { + RotatedHashMapping + .inbound( + MappingContext::Cache(CacheId::nil()), + StoreHash::parse(original).unwrap(), + ) + .to_string() +} + +fn store_path(hash: &str, name: &str) -> String { + format!("/nix/store/{hash}-{name}") +} + +struct Worker { + socket: WebSocketStream>, + received: Vec, +} + +struct Claimed { + job_id: String, + assignment_id: String, + job: Job, +} + +impl Worker { + async fn send(&mut self, msg: ClientMessage) { + let bytes = to_bytes(&msg, *PROTO_VERSIONS.end()).unwrap(); + self.socket.send(Message::Binary(bytes)).await.unwrap(); + } + + async fn next_within( + &mut self, + limit: Duration, + mut pick: impl FnMut(&ServerMessage) -> Option, + ) -> Option { + tokio::time::timeout(limit, async { + loop { + let Some(Ok(frame)) = self.socket.next().await else { + panic!("the server closed the session"); + }; + let Message::Binary(bytes) = frame else { + continue; + }; + let msg = from_bytes::(bytes, *PROTO_VERSIONS.end()).unwrap(); + let picked = pick(&msg); + self.received.push(msg); + if let Some(picked) = picked { + return picked; + } + } + }) + .await + .ok() + } + + async fn next(&mut self, pick: impl FnMut(&ServerMessage) -> Option) -> T { + self.next_within(WAIT, pick) + .await + .expect("the server answers in time") + } + + async fn candidates(&mut self) -> Vec { + self.send(ClientMessage::RequestJobList).await; + let mut all = Vec::new(); + loop { + let (candidates, is_final) = self + .next(|msg| match msg { + ServerMessage::JobListChunk { + candidates, + is_final, + } => Some((candidates.clone(), *is_final)), + _ => None, + }) + .await; + all.extend(candidates); + if is_final { + return all; + } + } + } + + async fn claim(&mut self, kind: JobKind, wanted: impl Fn(&JobCandidate) -> bool) -> Claimed { + let deadline = tokio::time::Instant::now() + WAIT; + while tokio::time::Instant::now() < deadline { + let candidates = self.candidates().await; + if let Some(candidate) = candidates.iter().find(|c| wanted(c)) { + let scores = vec![CandidateScore { + job_id: candidate.job_id.clone(), + missing_count: 0, + missing_nar_size: 0, + outputs_present: false, + }]; + self.send(ClientMessage::RequestJobChunk { + scores, + is_final: true, + }) + .await; + self.send(ClientMessage::RequestJob { kind: kind.clone() }) + .await; + let assigned = self + .next_within(Duration::from_secs(2), |msg| match msg { + ServerMessage::AssignJob { + job_id, + assignment_id, + job, + .. + } => Some(Claimed { + job_id: job_id.clone(), + assignment_id: assignment_id.clone(), + job: job.clone(), + }), + _ => None, + }) + .await; + if let Some(claimed) = assigned { + self.send(ClientMessage::AssignJobResponse { + job_id: claimed.job_id.clone(), + accepted: true, + reason: None, + }) + .await; + return claimed; + } + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + panic!("no {kind:?} job was handed out"); + } + + async fn update(&mut self, claimed: &Claimed, update: JobUpdateKind) { + self.send(ClientMessage::JobUpdate { + job_id: claimed.job_id.clone(), + assignment_id: claimed.assignment_id.clone(), + update, + }) + .await; + } + + async fn complete(&mut self, claimed: &Claimed) { + self.send(ClientMessage::JobCompleted { + job_id: claimed.job_id.clone(), + assignment_id: claimed.assignment_id.clone(), + spans: Vec::new(), + elapsed_ms: 0, + }) + .await; + } + + async fn upload( + &mut self, + claimed: &Claimed, + request_id: u64, + path: String, + references: Vec, + deriver: Option, + ) { + let nar = format!("nar of {path}").into_bytes(); + self.send(ClientMessage::UploadRequest { + job_id: claimed.job_id.clone(), + request_id, + object: UploadObject::Nar { store_path: path }, + size: nar.len() as u64, + }) + .await; + let target = self + .next(|msg| match msg { + ServerMessage::UploadGrant { + request_id: granted, + target, + } if *granted == request_id => Some(target.clone()), + _ => None, + }) + .await; + assert_eq!(target, GrantTarget::Passthrough { resume_offset: 0 }); + self.send(ClientMessage::UploadChunk { + request_id, + data: nar.clone().into(), + offset: 0, + is_final: true, + }) + .await; + self.send(ClientMessage::UploadFinished { + request_id, + metadata: UploadMetadata::Nar(Box::new(NarUploadMetadata { + file_hash: gradient_storage::file_hash_sri(&nar), + file_size: nar.len() as u64, + nar_size: nar.len() as u64, + nar_hash: "sha256:1111".into(), + references, + deriver, + ca: None, + multipart: None, + })), + }) + .await; + let outcome = self + .next(|msg| match msg { + ServerMessage::UploadCommitted { + request_id: committed, + outcome, + } if *committed == request_id => Some(outcome.clone()), + _ => None, + }) + .await; + assert_eq!(outcome, UploadOutcome::Ok); + } +} + +async fn eventually(db: &TestDatabase, ready: impl AsyncFn(&TestDatabase) -> bool) { + let deadline = tokio::time::Instant::now() + WAIT; + while !ready(db).await { + assert!(tokio::time::Instant::now() < deadline, "timed out"); + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +async fn collect_event_lines( + mut events: EventRx, + mut flow_ended: oneshot::Receiver<()>, +) -> Vec { + let mut lines = Vec::new(); + loop { + tokio::select! { + biased; + received = events.recv() => match received { + Ok(envelope) => lines.push(envelope.to_line()), + Err(RecvError::Lagged(missed)) => panic!("missed {missed} events"), + Err(RecvError::Closed) => break, + }, + _ = &mut flow_ended => break, + } + } + while let Ok(envelope) = events.try_recv() { + lines.push(envelope.to_line()); + } + lines +} + +fn object_keys(dir: &Path) -> Vec { + std::fs::read_dir(dir) + .unwrap() + .flat_map(|entry| { + let path = entry.unwrap().path(); + if path.is_dir() { + object_keys(&path) + } else { + vec![path.to_string_lossy().into_owned()] + } + }) + .collect() +} + +fn discovered(drv: &str, out: &str, name: &str, attr: &str) -> DiscoveredDerivation { + DiscoveredDerivation { + attr: attr.into(), + drv_path: store_path(drv, &format!("{name}.drv")), + outputs: vec![DerivationOutput { + name: "out".into(), + path: store_path(out, name), + }], + input_sources: vec![store_path(SRC, "source")], + architecture: "x86_64-linux".into(), + allow_substitutes: true, + ..Default::default() + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn after_an_evaluation_build_and_upload_no_original_hash_is_stored() { + let e1 = EvaluationId::now_v7(); + let (secret, _) = testing::secret_file(); + let db = shape() + .evaluation(e1) + .evaluation_status(EvaluationStatus::Queued) + .repository("https://example.invalid/repo") + .into_cache(CacheId::now_v7()) + .keyed(testing::signing_key(&secret.path().to_string_lossy())) + .into_database() + .await; + let project = db.project_of(e1).await; + let storage = tempfile::TempDir::new().unwrap(); + let crypt_file = secret.path().to_string_lossy().into_owned(); + let graph = Graph::new(); + let state = TestState::new() + .worker_db(db.connection()) + .web_db(db.connection()) + .cache_db(db.connection()) + .cli(|cli| { + cli.secrets.crypt_file = crypt_file; + cli.build = BuildArgs::default(); + }) + .nar_storage(gradient_storage::NarStore::local(storage.path()).unwrap()) + .graph(Arc::clone(&graph)) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + let (end_flow, flow_ended) = oneshot::channel(); + let event_lines = state + .shutdown + .spawn(collect_event_lines(state.events.subscribe(), flow_ended)); + graph.spawn(state.db(), None, None).await.unwrap(); + let scheduler = Arc::new(Scheduler::new(Arc::clone(&state))); + scheduler.start(); + let (sessions, _) = Actor::spawn( + None, + Sessions, + SessionsArgs { + scheduler: Arc::clone(&scheduler), + shutdown: state.shutdown.clone(), + }, + ) + .await + .unwrap(); + let (socket, client) = connected_pair().await; + let args = SessionArgs { + peer_id: "w1".into(), + handles: SessionHandles::from_state(&state), + scheduler: Arc::clone(&scheduler), + socket, + capabilities: GradientCapabilities { + eval: true, + build: true, + network: true, + ..Default::default() + }, + authorized_peers: HashSet::from([project]), + dialed: None, + }; + let CallResult::Success(Ok(_)) = sessions + .call(|reply| SessionsMsg::Attach(Box::new(args), reply), None) + .await + .unwrap() + else { + panic!("the session attaches"); + }; + let mut worker = Worker { + socket: client, + received: Vec::new(), + }; + worker + .send(ClientMessage::WorkerCapabilities { + architectures: vec!["x86_64-linux".into()], + system_features: Vec::new(), + max_concurrent_builds: 1, + cpu_count: 4, + ram_total_mb: 8192, + cpu_core_score: 1, + zone: None, + endpoint: None, + build_cores: None, + }) + .await; + + let eval = worker + .claim(JobKind::Flake, |c| c.job_id == eval_job_key(e1)) + .await; + worker.update(&eval, JobUpdateKind::Fetching).await; + worker + .update( + &eval, + JobUpdateKind::FetchResult { + flake_source: Some(store_path(SRC, "source")), + }, + ) + .await; + worker.update(&eval, JobUpdateKind::EvaluatingFlake).await; + worker + .update(&eval, JobUpdateKind::EvaluatingDerivations) + .await; + worker + .update( + &eval, + JobUpdateKind::EvalResult { + derivations: vec![ + discovered(A_DRV, A_OUT, "a", "packages.x86_64-linux.a"), + discovered(B_DRV, B_OUT, "b", ""), + ], + warnings: Vec::new(), + errors: Vec::new(), + }, + ) + .await; + let stored_a_drv = format!("{}-a.drv", rotate(A_DRV)); + eventually(&db, async |db| { + db.find_derivation(&stored_a_drv).await.is_some() + }) + .await; + worker + .upload(&eval, 1, store_path(SRC, "source"), Vec::new(), None) + .await; + worker + .upload( + &eval, + 2, + store_path(A_DRV, "a.drv"), + vec![format!("{SRC}-source")], + None, + ) + .await; + worker + .send(ClientMessage::CacheQuery { + job_id: eval.job_id.clone(), + query_id: "q1".into(), + paths: vec![store_path(A_OUT, "a")], + mode: QueryMode::Normal, + nar_sizes: vec![None], + external: false, + }) + .await; + worker + .next(|msg| matches!(msg, ServerMessage::CacheStatus { .. }).then_some(())) + .await; + worker.complete(&eval).await; + + let original_a_drv = store_path(A_DRV, "a.drv"); + let build = worker + .claim(JobKind::Build, |c| c.drv_paths.contains(&original_a_drv)) + .await; + let Job::Build(job) = &build.job else { + panic!("a build job, got {:?}", build.job); + }; + let build_id = job.builds[0].build_id.clone(); + worker + .update( + &build, + JobUpdateKind::Building { + build_id: build_id.clone(), + }, + ) + .await; + worker + .update( + &build, + JobUpdateKind::BuildOutput { + build_id, + outputs: vec![BuildOutput { + name: "out".into(), + store_path: store_path(A_OUT, "a"), + hash: A_OUT.into(), + nar_size: Some(1), + nar_hash: Some("sha256:1111".into()), + products: Vec::new(), + }], + metrics: None, + substituted: false, + }, + ) + .await; + worker + .upload( + &build, + 3, + store_path(A_OUT, "a"), + vec![format!("{B_OUT}-b")], + Some(format!("{A_DRV}-a.drv")), + ) + .await; + worker.complete(&build).await; + eventually(&db, async |db| db.output_stored(&stored_a_drv, "out").await).await; + + for original in ORIGINALS { + assert_eq!( + db.columns_containing(original).await, + Vec::::new(), + "{original}" + ); + } + assert!(!db.columns_containing(&rotate(A_DRV)).await.is_empty()); + let frames = format!("{:?}", worker.received); + for original in ORIGINALS { + assert!(!frames.contains(&rotate(original)), "{original}"); + } + let _ = end_flow.send(()); + let event_lines = event_lines.await.unwrap(); + assert!(!event_lines.is_empty()); + for line in &event_lines { + for original in ORIGINALS { + assert!(!line.contains(&rotate(original)), "{original} in {line}"); + } + } + let keys = object_keys(storage.path()); + assert!(!keys.is_empty()); + let root = storage.path().to_string_lossy().into_owned(); + for key in keys.iter().map(|key| key.trim_start_matches(&root)) { + for original in ORIGINALS { + assert!(!key.contains(original), "{original} in {key}"); + } + } + + sessions.stop(None); + state + .shutdown + .cancel_and_drain(Duration::from_secs(10)) + .await; + db.finish().await; +} diff --git a/backend/gradient-proto/src/handler/upload/commit.rs b/backend/gradient-proto/src/handler/upload/commit.rs index 826026c49..e15beea49 100644 --- a/backend/gradient-proto/src/handler/upload/commit.rs +++ b/backend/gradient-proto/src/handler/upload/commit.rs @@ -4,19 +4,21 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use futures::future::BoxFuture; +use gradient_db::cache_paths::project_write_caches; use gradient_db::{DbContext, WorkerDb}; use gradient_storage::NarStore; -use gradient_storage::admission::{ObjectKey, UploadPermit}; +use gradient_storage::admission::UploadPermit; use gradient_types::RuntimeConfig; -use gradient_types::ids::ProjectId; +use gradient_types::ids::{CacheId, ProjectId}; use gradient_wire::messages::ServerMessage; use gradient_wire::types::{NarUploadMetadata, UploadMetadata, UploadObject, UploadOutcome}; use tracing::warn; use super::super::SessionHandles; +use super::super::eval_cache::eval_cache_key; use super::super::nar::{NarUploadRecord, mark_nar_stored, record_nar_push_metric}; use super::super::socket::{ProtoWriter, send_server_msg}; +use super::super::store_paths::{inbound_upload_metadata, unmapped_project}; use super::table::Transfer; use crate::import::{WriteNeeded, nar_write_needed}; @@ -25,9 +27,7 @@ pub(super) struct Commit { pub handles: SessionHandles, pub peer_id: String, pub request_id: u64, - /// This is resolved off the session's reader loop. The scheduler is answering from a mailbox - /// that a claim in flight can hold for its lock waits. - pub project: BoxFuture<'static, Option>, + pub project: ProjectId, pub object: UploadObject, pub transfer: Transfer, pub metadata: UploadMetadata, @@ -38,12 +38,30 @@ pub(super) struct Commit { /// recorded it, and a second upload of the path is waiting instead of overwriting the object. #[tracing::instrument(name = "upload_commit", level = "debug", skip_all)] pub(super) async fn run(c: Commit) { - let outcome = match place(&c.handles.db, &c.object, c.transfer, &c.metadata).await { - Ok(()) => { + let metadata = match c.handles.db.project_mapper(c.project).await { + Ok(mapper) => inbound_upload_metadata(c.metadata, &mapper).map_err(|e| e.to_string()), + Err(e) => Err(unmapped_project(&e)), + }; + let writable = match &c.object { + UploadObject::EvalCache { .. } => { + project_write_caches(&c.handles.db.worker_db, c.project).await + } + UploadObject::Nar { .. } => Ok(Vec::new()), + }; + let placed = match (metadata, writable) { + (Ok(metadata), Ok(writable)) => { + place(&c.handles.db, &writable, &c.object, c.transfer, &metadata) + .await + .map(|()| (metadata, writable)) + } + (Err(reason), _) => Err(rejected(reason)), + (_, Err(e)) => Err(retry(format!("looking up the writable caches: {e}"))), + }; + let outcome = match placed { + Ok((metadata, writable)) => { let mut permit = c.permit; permit.release_budget(); - let project_id = c.project.await; - let outcome = record(&c.handles, project_id, &c.object, &c.metadata).await; + let outcome = record(&c.handles, c.project, &writable, &c.object, &metadata).await; if matches!(outcome, UploadOutcome::Ok) { permit.committed(); } @@ -81,13 +99,14 @@ fn rejected(reason: String) -> UploadOutcome { #[tracing::instrument(name = "place_upload", level = "debug", skip_all)] async fn place( db: &DbContext, + writable: &[CacheId], object: &UploadObject, transfer: Transfer, metadata: &UploadMetadata, ) -> Result<(), UploadOutcome> { match (object, metadata) { (UploadObject::Nar { store_path }, UploadMetadata::Nar(meta)) => { - let Some(ObjectKey::Nar(hash)) = super::object_key(object) else { + let Some(hash) = super::nar_key(store_path) else { return Err(rejected(format!("malformed store path {store_path}"))); }; match transfer { @@ -107,7 +126,11 @@ async fn place( } } (UploadObject::EvalCache { fingerprint }, UploadMetadata::EvalCache { size_bytes }) => { - place_eval_cache(&db.storage.nar_storage, fingerprint, transfer, *size_bytes).await + let Some(&target) = writable.first() else { + return Err(rejected("the project has no writable cache".into())); + }; + let key = eval_cache_key(target, fingerprint); + place_eval_cache(&db.storage.nar_storage, &key, transfer, *size_bytes).await } _ => Err(rejected("metadata does not match the upload object".into())), } @@ -116,7 +139,8 @@ async fn place( #[tracing::instrument(name = "upload_graph_commit", level = "debug", skip_all)] async fn record( handles: &SessionHandles, - project_id: Option, + project_id: ProjectId, + writable: &[CacheId], object: &UploadObject, metadata: &UploadMetadata, ) -> UploadOutcome { @@ -132,14 +156,8 @@ async fn record( ca: meta.ca.as_deref(), confirmed: true, }; - if let Err(e) = mark_nar_stored( - &handles.graph, - &handles.db.events, - project_id, - store_path, - &record, - ) - .await + if let Err(e) = + mark_nar_stored(&handles.graph, &handles.db, project_id, store_path, &record).await { return retry(format!("recording {store_path} in the cache index: {e:#}")); } @@ -154,6 +172,8 @@ async fn record( (UploadObject::EvalCache { fingerprint }, UploadMetadata::EvalCache { size_bytes }) => { super::super::eval_cache::record_eval_cache( &handles.db.worker_db, + &handles.db.storage.nar_storage, + writable, fingerprint, *size_bytes, ) @@ -237,7 +257,7 @@ async fn place_presigned( async fn place_eval_cache( nar_storage: &NarStore, - fingerprint: &str, + key: &str, transfer: Transfer, size_bytes: u64, ) -> Result<(), UploadOutcome> { @@ -252,13 +272,10 @@ async fn place_eval_cache( } let bytes = tokio::fs::read(&staged.path).await.map_err(retry)?; let _ = tokio::fs::remove_file(&staged.path).await; - nar_storage - .put_eval_cache(fingerprint, bytes) - .await - .map_err(retry) + nar_storage.put_eval_cache(key, bytes).await.map_err(retry) } Transfer::Put | Transfer::Multipart { .. } => nar_storage - .verify_eval_cache(fingerprint, size_bytes) + .verify_eval_cache(key, size_bytes) .await .map_err(retry), } @@ -302,7 +319,7 @@ mod tests { handles: SessionHandles::from_state(&state), peer_id: "w1".into(), request_id: 1, - project: Box::pin(std::future::ready(None)), + project: ProjectId::nil(), object: UploadObject::Nar { store_path: format!("/nix/store/{}-p", "c".repeat(32)), }, @@ -347,7 +364,7 @@ mod tests { handles: SessionHandles::from_state(&state), peer_id: "w1".into(), request_id: 1, - project: Box::pin(std::future::ready(None)), + project: ProjectId::nil(), object: UploadObject::Nar { store_path: format!("/nix/store/{}-p", "c".repeat(32)), }, @@ -402,6 +419,7 @@ mod tests { place( &state.db(), + &[], &UploadObject::Nar { store_path: format!("/nix/store/{hash}-p"), }, @@ -446,7 +464,7 @@ mod tests { handles: SessionHandles::from_state(&state), peer_id: "w1".into(), request_id: 1, - project: Box::pin(std::future::ready(None)), + project: ProjectId::nil(), object: UploadObject::Nar { store_path: format!("/nix/store/{}-p", "c".repeat(32)), }, @@ -505,7 +523,7 @@ mod tests { handles: SessionHandles::from_state(&state), peer_id: "w1".into(), request_id: 1, - project: Box::pin(std::future::ready(None)), + project: ProjectId::nil(), object: UploadObject::Nar { store_path: format!("/nix/store/{}-p", "c".repeat(32)), }, diff --git a/backend/gradient-proto/src/handler/upload/mod.rs b/backend/gradient-proto/src/handler/upload/mod.rs index 3fec4019b..4c5aef6e9 100644 --- a/backend/gradient-proto/src/handler/upload/mod.rs +++ b/backend/gradient-proto/src/handler/upload/mod.rs @@ -7,11 +7,12 @@ mod commit; mod table; -use std::sync::Arc; use std::time::{Duration, Instant}; use gradient_storage::admission::{AdmissionSession, Admitted, ObjectKey}; use gradient_storage::{NarStore, PartialStore, StorageTarget, upload_lease}; +use gradient_types::StorePath; +use gradient_types::ids::ProjectId; use gradient_wire::messages::{ServerMessage, is_small_upload}; use gradient_wire::types::{GrantTarget, UploadMetadata, UploadObject, UploadOutcome}; use tracing::warn; @@ -28,15 +29,10 @@ pub(super) struct UploadSession { pub idle_lease: Duration, } -pub(super) fn object_key(object: &UploadObject) -> Option { - match object { - UploadObject::Nar { store_path } => { - let name = store_path.strip_prefix("/nix/store/")?; - let hash = name.split('-').next()?; - (hash.len() == 32).then(|| ObjectKey::Nar(hash.to_owned())) - } - UploadObject::EvalCache { fingerprint } => Some(ObjectKey::EvalCache(fingerprint.clone())), - } +pub(super) fn nar_key(store_path: &str) -> Option { + StorePath::parse_stored(store_path) + .ok() + .map(|path| path.hash().to_owned()) } fn key_name(key: &ObjectKey) -> &str { @@ -55,17 +51,7 @@ impl InboundContext<'_> { size: u64, uploads: &mut UploadSession, ) { - let Some(key) = object_key(&object) else { - return self - .settle( - request_id, - UploadOutcome::Rejected { - reason: format!("malformed upload object {object:?}"), - }, - ) - .await; - }; - if !self.active.contains(&job_id) && self.active.fetch_project(&job_id).is_none() { + let Some(project) = self.active.session_project(&job_id) else { return self .settle( request_id, @@ -74,14 +60,33 @@ impl InboundContext<'_> { }, ) .await; - } - if let UploadObject::EvalCache { fingerprint } = &object - && !super::eval_cache::accepts_push(&self.handles.db.worker_db, fingerprint, size).await - { - return self.grant(request_id, GrantTarget::Skip).await; - } + }; + let key = match &object { + UploadObject::Nar { store_path } => match nar_key(store_path) { + Some(hash) => ObjectKey::Nar(hash), + None => { + return self + .settle( + request_id, + UploadOutcome::Rejected { + reason: format!("malformed upload object {object:?}"), + }, + ) + .await; + } + }, + UploadObject::EvalCache { fingerprint } => { + match self.eval_cache_target(project, fingerprint, size).await { + Some(key) => ObjectKey::EvalCache(key), + None => return self.grant(request_id, GrantTarget::Skip).await, + } + } + }; let priority = gradient_wire::messages::is_small_upload(size); - if !uploads.table.queue(request_id, job_id, object, size) { + if !uploads + .table + .queue(request_id, job_id, project, object, size) + { return self .settle( request_id, @@ -137,6 +142,7 @@ impl InboundContext<'_> { id, Granted { job_id: queued.job_id, + project: queued.project, object: queued.object, size: queued.size, permit, @@ -164,6 +170,25 @@ impl InboundContext<'_> { } } + async fn eval_cache_target( + &self, + project: ProjectId, + fingerprint: &str, + size: u64, + ) -> Option { + let db = &self.handles.db.worker_db; + let writable = gradient_db::cache_paths::project_write_caches(db, project) + .await + .unwrap_or_else(|e| { + warn!(peer_id = %self.peer_id, %project, error = %e, "writable caches lookup failed"); + Vec::new() + }); + let &target = writable.first()?; + super::eval_cache::accepts_push(db, &writable, fingerprint, size) + .await + .then(|| super::eval_cache::eval_cache_key(target, fingerprint)) + } + #[tracing::instrument(level = "debug", skip_all)] async fn already_stored(&self, key: &ObjectKey, size: u64) -> Result { let ObjectKey::Nar(hash) = key else { @@ -310,28 +335,12 @@ impl InboundContext<'_> { let Some(metadata) = granted.finished else { return; }; - let project = { - let scheduler = Arc::clone(self.scheduler); - let db = self.handles.db.worker_db.clone(); - let peer_id = self.peer_id.to_owned(); - let job_id = granted.job_id.clone(); - let fetch_project = self.active.fetch_project(&job_id); - Box::pin(async move { - if fetch_project.is_some() { - return fetch_project; - } - match scheduler.project_for_job(&job_id).await { - Some(id) => Some(id), - None => super::nar::project_for_dispatched_job(&db, &peer_id, &job_id).await, - } - }) - }; let c = commit::Commit { writer: self.writer.clone(), handles: self.handles.clone(), peer_id: self.peer_id.to_owned(), request_id, - project, + project: granted.project, object: granted.object, transfer: granted.transfer, metadata, @@ -375,8 +384,9 @@ impl InboundContext<'_> { } pub(super) async fn abandon_transfer(nar_storage: &NarStore, granted: Granted) { - if let (Transfer::Multipart { upload_id }, Some(ObjectKey::Nar(hash))) = - (&granted.transfer, object_key(&granted.object)) + if let (Transfer::Multipart { upload_id }, UploadObject::Nar { store_path }) = + (&granted.transfer, &granted.object) + && let Some(hash) = nar_key(store_path) { nar_storage.abort_multipart(&hash, upload_id).await; } diff --git a/backend/gradient-proto/src/handler/upload/table.rs b/backend/gradient-proto/src/handler/upload/table.rs index a8f3a8dfb..6df557876 100644 --- a/backend/gradient-proto/src/handler/upload/table.rs +++ b/backend/gradient-proto/src/handler/upload/table.rs @@ -9,6 +9,7 @@ use std::time::{Duration, Instant}; use gradient_storage::PartialWriter; use gradient_storage::admission::UploadPermit; +use gradient_types::ids::ProjectId; use gradient_wire::types::{UploadMetadata, UploadObject}; pub(in crate::handler) enum Transfer { @@ -39,12 +40,14 @@ impl Lease { pub(in crate::handler) struct Queued { pub job_id: String, + pub project: ProjectId, pub object: UploadObject, pub size: u64, } pub(in crate::handler) struct Granted { pub job_id: String, + pub project: ProjectId, pub object: UploadObject, pub size: u64, pub permit: UploadPermit, @@ -79,6 +82,7 @@ impl UploadTable { &mut self, id: u64, job_id: String, + project: ProjectId, object: UploadObject, size: u64, ) -> bool { @@ -89,6 +93,7 @@ impl UploadTable { id, Entry::Queued(Queued { job_id, + project, object, size, }), @@ -221,11 +226,12 @@ mod tests { let (_session, permit) = permit(&admission, 1).await; let mut table = UploadTable::default(); let start = Instant::now(); - assert!(table.queue(1, "build:1".into(), nar(0), 10)); + assert!(table.queue(1, "build:1".into(), ProjectId::nil(), nar(0), 10)); table.grant( 1, Granted { job_id: "build:1".into(), + project: ProjectId::nil(), object: nar(0), size: 10, permit, @@ -262,11 +268,12 @@ mod tests { let (_session, permit) = permit(&admission, 1).await; let mut table = UploadTable::default(); let start = Instant::now(); - table.queue(1, "build:1".into(), nar(0), 10); + table.queue(1, "build:1".into(), ProjectId::nil(), nar(0), 10); table.grant( 1, Granted { job_id: "build:1".into(), + project: ProjectId::nil(), object: nar(0), size: 10, permit, @@ -290,16 +297,16 @@ mod tests { #[test] fn a_duplicate_request_id_is_refused() { let mut table = UploadTable::default(); - assert!(table.queue(1, "build:1".into(), nar(0), 1)); - assert!(!table.queue(1, "build:1".into(), nar(1), 1)); + assert!(table.queue(1, "build:1".into(), ProjectId::nil(), nar(0), 1)); + assert!(!table.queue(1, "build:1".into(), ProjectId::nil(), nar(1), 1)); } #[test] fn forgetting_a_job_returns_only_its_requests() { let mut table = UploadTable::default(); - table.queue(1, "build:1".into(), nar(0), 1); - table.queue(2, "build:2".into(), nar(1), 1); - table.queue(3, "build:1".into(), nar(2), 1); + table.queue(1, "build:1".into(), ProjectId::nil(), nar(0), 1); + table.queue(2, "build:2".into(), ProjectId::nil(), nar(1), 1); + table.queue(3, "build:1".into(), ProjectId::nil(), nar(2), 1); let mut ids: Vec = table .forget_job("build:1") .into_iter() diff --git a/backend/gradient-proto/src/import.rs b/backend/gradient-proto/src/import.rs index b8c320e6f..00cfa8879 100644 --- a/backend/gradient-proto/src/import.rs +++ b/backend/gradient-proto/src/import.rs @@ -45,7 +45,7 @@ impl ImportInput<'_> { } fn parse_store_path(store_path: &str) -> anyhow::Result { - let sp = StorePath::parse(store_path)?; + let sp = StorePath::parse_stored(store_path)?; if !is_nix32_hash(sp.hash()) { anyhow::bail!(StorePathError::Malformed(store_path.to_owned())); } diff --git a/backend/gradient-scheduler/src/actor.rs b/backend/gradient-scheduler/src/actor.rs index e48f1f21a..8f7b484c6 100644 --- a/backend/gradient-scheduler/src/actor.rs +++ b/backend/gradient-scheduler/src/actor.rs @@ -4,7 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::time::{Duration, Instant}; @@ -55,6 +55,7 @@ pub struct WorkerMetrics { #[derive(Debug, Clone, Default)] pub struct Offer { pub candidates: Vec, + pub projects: HashMap, pub generation: u64, } @@ -477,16 +478,21 @@ impl SchedulerCore { fn candidates(&mut self, worker: &str, only_new: bool) -> Offer { let (authorized, caps) = self.auth_and_caps(worker); - let mut candidates = self + let mut offered = self .tracker .candidates_for_worker(authorized.as_ref(), caps.as_ref()); - let visible: HashSet = candidates.iter().map(|c| c.job_id.clone()).collect(); + let visible: HashSet = offered.iter().map(|(c, _)| c.job_id.clone()).collect(); if only_new && let Some(sent) = self.pool.sent_candidates_for(worker) { - candidates.retain(|c| !sent.contains(&c.job_id)); + offered.retain(|(c, _)| !sent.contains(&c.job_id)); } self.pool.set_sent_candidates(worker, visible); + let projects = offered + .iter() + .map(|(c, project)| (c.job_id.clone(), *project)) + .collect(); Offer { - candidates, + candidates: offered.into_iter().map(|(c, _)| c).collect(), + projects, generation: self.offers, } } diff --git a/backend/gradient-scheduler/src/cases/readers.rs b/backend/gradient-scheduler/src/cases/readers.rs index 582cf2b62..9b8e4f025 100644 --- a/backend/gradient-scheduler/src/cases/readers.rs +++ b/backend/gradient-scheduler/src/cases/readers.rs @@ -46,7 +46,7 @@ pub(crate) fn planned_targets(plan: &ProbePlan) -> BTreeMap, upstream_query: &Arc, evaluation: &MEvaluation, - to_probe: Vec<(String, String)>, + to_probe: Vec, ) -> ProbeAnswer { let mut answer = ProbeAnswer::default(); if to_probe.is_empty() { @@ -48,7 +49,7 @@ pub async fn probe_outputs( Ok(endpoints) => endpoints, Err(e) => { warn!(error = %e, evaluation = %evaluation.id, "upstream probe: could not read the upstream caches; asking again later"); - answer.unanswered = to_probe.into_iter().map(|(hash, _)| hash).collect(); + answer.unanswered = to_probe.iter().map(|p| p.hash().to_owned()).collect(); return answer; } }; @@ -56,9 +57,21 @@ pub async fn probe_outputs( return answer; } + let mapper = match StorePathMapper::for_project(Arc::clone(mapping), project_id).await { + Ok(mapper) => mapper, + Err(e) => { + warn!(error = %e, evaluation = %evaluation.id, "upstream probe: the project's store paths cannot be mapped; treating every path as a miss"); + return answer; + } + }; let id_to_url: HashMap<_, String> = endpoints.iter().map(|e| (e.id, e.url.clone())).collect(); - let batch = - gradient_core::upstream::probe_batch(endpoints, Arc::clone(upstream_query), to_probe).await; + let batch = gradient_core::upstream::probe_batch( + endpoints, + Arc::clone(upstream_query), + mapper, + to_probe, + ) + .await; warn_unanswered(db, evaluation, &batch.unanswered, &id_to_url).await; // The same URL under different upstream ids is folding into one metric series (#417). diff --git a/backend/gradient-scheduler/src/job_handlers/eval_status.rs b/backend/gradient-scheduler/src/job_handlers/eval_status.rs index 1d14e038b..6b39b4dc6 100644 --- a/backend/gradient-scheduler/src/job_handlers/eval_status.rs +++ b/backend/gradient-scheduler/src/job_handlers/eval_status.rs @@ -11,7 +11,6 @@ use tracing::{debug, warn}; use gradient_graph::RecordBatch; use gradient_types::*; -use gradient_util::store_path::strip_nix_store_prefix; use gradient_wire::types::DiscoveredDerivation; use crate::Scheduler; @@ -187,9 +186,9 @@ impl Scheduler { // Store paths are canonicalised to the bare `-` form before reaching the graph // writer. The `/nix/store/` prefix is added back only at the worker and API boundary. for d in &mut derivations { - d.drv_path = strip_nix_store_prefix(&d.drv_path); + d.drv_path = StorePath::parse_stored(&d.drv_path)?.base(); for dep in &mut d.dependencies { - *dep = strip_nix_store_prefix(dep); + *dep = StorePath::parse_stored(dep)?.base(); } } diff --git a/backend/gradient-scheduler/src/job_handlers/import.rs b/backend/gradient-scheduler/src/job_handlers/import.rs index 1ded9fc27..ca205ba95 100644 --- a/backend/gradient-scheduler/src/job_handlers/import.rs +++ b/backend/gradient-scheduler/src/job_handlers/import.rs @@ -107,7 +107,11 @@ impl Scheduler { ) -> Result, ImportOutcome>> { let hashes: Vec> = drv_paths .iter() - .map(|p| StorePath::parse(p).ok().map(|sp| sp.hash().to_owned())) + .map(|p| { + StorePath::parse_stored(p) + .ok() + .map(|sp| sp.hash().to_owned()) + }) .collect(); let rows = gradient_db::evaluations::imports::import_builds( &self.db.worker_db, diff --git a/backend/gradient-scheduler/src/job_handlers/queue.rs b/backend/gradient-scheduler/src/job_handlers/queue.rs index 8a1e1bfa0..3ebb66932 100644 --- a/backend/gradient-scheduler/src/job_handlers/queue.rs +++ b/backend/gradient-scheduler/src/job_handlers/queue.rs @@ -8,7 +8,7 @@ use crate::Scheduler; use crate::actor::{Offer, SchedulerMsg}; use crate::jobs::{PendingBuildJob, PendingEvalJob, PendingJob}; use gradient_pool::WorkerInfo; -use gradient_wire::types::{GradientCapabilities, JobCandidate}; +use gradient_wire::types::GradientCapabilities; impl Scheduler { pub async fn enqueue_eval_job( @@ -86,7 +86,7 @@ impl Scheduler { .await } - pub async fn get_job_candidates(&self, worker_id: &str) -> Vec { + pub async fn get_job_candidates(&self, worker_id: &str) -> Offer { let worker = worker_id.to_owned(); self.call(|reply| SchedulerMsg::Candidates { worker, @@ -94,7 +94,6 @@ impl Scheduler { reply, }) .await - .map(|o| o.candidates) .unwrap_or_default() } diff --git a/backend/gradient-scheduler/src/jobs.rs b/backend/gradient-scheduler/src/jobs.rs index 22c7b66aa..5d7d7a340 100644 --- a/backend/gradient-scheduler/src/jobs.rs +++ b/backend/gradient-scheduler/src/jobs.rs @@ -715,14 +715,14 @@ impl JobTracker { &self, authorized: Option<&HashSet>, caps: Option<&WorkerCaps>, - ) -> Vec { + ) -> Vec<(JobCandidate, ProjectId)> { let members = self .clusters .jobs() .filter(|(_, job)| visible_to(job, authorized, caps)); self.eligible_for_worker(authorized, caps) .chain(members) - .map(|(id, job)| job.as_candidate(id)) + .map(|(id, job)| (job.as_candidate(id), job.project_id())) .collect() } @@ -2333,7 +2333,7 @@ mod tests { let candidates = tracker.candidates_for_worker(Some(&authorized), None); assert_eq!(candidates.len(), 1); - assert_eq!(candidates[0].job_id, "ja"); + assert_eq!(candidates[0].0.job_id, "ja"); } #[test] @@ -2355,7 +2355,7 @@ mod tests { let x86_caps = build_caps(&["x86_64-linux"], &[]); let candidates = tracker.candidates_for_worker(None, Some(&x86_caps)); - let mut ids: Vec<_> = candidates.iter().map(|c| c.job_id.clone()).collect(); + let mut ids: Vec<_> = candidates.iter().map(|(c, _)| c.job_id.clone()).collect(); ids.sort(); assert_eq!(ids, vec!["any".to_string(), "x86".to_string()]); } @@ -3086,7 +3086,7 @@ mod tests { let offered = tracker.candidates_for_worker(None, None); assert_eq!(offered.len(), 1); - assert_eq!(offered[0].job_id, "m1"); + assert_eq!(offered[0].0.job_id, "m1"); assert!(tracker.contains_job("m1")); let p = gradient_pool::score::policy_for(gradient_types::ScoringPolicy::Simple); diff --git a/backend/gradient-scheduler/src/log_substitution.rs b/backend/gradient-scheduler/src/log_substitution.rs index 0c60465d6..21dbd9ad6 100644 --- a/backend/gradient-scheduler/src/log_substitution.rs +++ b/backend/gradient-scheduler/src/log_substitution.rs @@ -72,6 +72,16 @@ pub async fn substitute_log( } }; + let mapper = match db.project_mapper(project_id).await { + Ok(mapper) => mapper, + Err(e) => { + warn!(error = %e, "substitute_log: the project's store paths cannot be mapped"); + return Ok(()); + } + }; + let Ok(drv_path) = mapper.outbound_text(&drv_path) else { + return Ok(()); + }; let Some(drv_basename) = std::path::Path::new(&drv_path) .file_name() .and_then(|n| n.to_str()) diff --git a/backend/gradient-scheduler/src/probe.rs b/backend/gradient-scheduler/src/probe.rs index d0ddeb0a9..b0d21782e 100644 --- a/backend/gradient-scheduler/src/probe.rs +++ b/backend/gradient-scheduler/src/probe.rs @@ -32,6 +32,7 @@ type Inbox = Mutex>>>; pub(crate) struct UpstreamProbe { db: WorkerDb, + store_path_mapping: Arc, graph: Arc, upstream_query: Arc, wake: Arc, @@ -43,6 +44,7 @@ impl UpstreamProbe { pub(crate) fn new(db: &DbContext, graph: &Arc, upstream_query: &Arc) -> Self { Self { db: db.worker_db.clone(), + store_path_mapping: Arc::clone(&db.store_path_mapping), graph: Arc::clone(graph), upstream_query: Arc::clone(upstream_query), wake: db.memory.probe_requests.wake(), @@ -81,6 +83,7 @@ impl UpstreamProbe { for chunk in targets.chunks(PROBE_BATCH) { let answer = crate::eval::probe_outputs( &self.db, + &self.store_path_mapping, &self.upstream_query, evaluation, chunk.to_vec(), @@ -178,7 +181,7 @@ fn unanswered_need_query() -> sea_orm::Select { #[derive(Debug, Default)] pub(crate) struct ProbePlan { - pub rounds: Vec<(MEvaluation, Vec<(String, String)>)>, + pub rounds: Vec<(MEvaluation, Vec)>, pub answered: Vec, pub owners: HashMap>, } @@ -300,7 +303,7 @@ pub(crate) async fn plan_probes( .or_insert(job.evaluation); } - let mut by_evaluation: HashMap> = HashMap::new(); + let mut by_evaluation: HashMap> = HashMap::new(); let mut owners: HashMap> = HashMap::new(); let cached = gradient_db::derivation_outputs::cached_anywhere(db, &outputs) .await @@ -309,11 +312,10 @@ pub(crate) async fn plan_probes( let Some(evaluation) = evaluation_of.get(&o.derivation) else { continue; }; - let sp = StorePath::from_parts(o.hash.clone(), o.package.clone()); by_evaluation .entry(*evaluation) .or_default() - .insert(o.hash.clone(), sp.full()); + .insert(o.hash.clone(), StorePath::from(o.clone())); owners.entry(o.hash.clone()).or_default().push(o.derivation); } if by_evaluation.is_empty() { @@ -339,7 +341,7 @@ pub(crate) async fn plan_probes( .filter_map(|e| { by_evaluation .remove(&e.id) - .map(|targets| (e, targets.into_iter().collect())) + .map(|targets| (e, targets.into_values().collect())) }) .collect(); diff --git a/backend/gradient-scheduler/src/views.rs b/backend/gradient-scheduler/src/views.rs index 1f2fd6920..7f7a0f45f 100644 --- a/backend/gradient-scheduler/src/views.rs +++ b/backend/gradient-scheduler/src/views.rs @@ -156,7 +156,7 @@ impl JobContextView { .iter() .map(|t| DerivationRef { build_id: t.build_id.clone(), - drv_path: gradient_types::StorePath::parse(&t.drv_path) + drv_path: gradient_types::StorePath::parse_stored(&t.drv_path) .map(|sp| sp.base()) .unwrap_or_else(|_| t.drv_path.clone()), pname: b.pname.clone(), diff --git a/backend/gradient-sources/src/cache_key/signing.rs b/backend/gradient-sources/src/cache_key/signing.rs index 996d22140..cc99c9093 100644 --- a/backend/gradient-sources/src/cache_key/signing.rs +++ b/backend/gradient-sources/src/cache_key/signing.rs @@ -46,12 +46,29 @@ impl CacheSigner { nar_size: u64, references: &[String], ) -> String { - let raw = self.sign_narinfo_raw(store_path, nar_hash, nar_size, references); + let raw = self.sign_fingerprint(store_path, nar_hash, nar_size, references); let sig_b64 = general_purpose::STANDARD.encode(raw); format!("{}-{}:{}", self.base_url, self.cache_name, sig_b64) } - pub fn sign_narinfo_raw( + pub fn sign_stored_narinfo( + &self, + mapper: &StorePathMapper, + store_path: &str, + nar_hash: &str, + nar_size: u64, + references: &[String], + ) -> Vec { + let original = |text: &str| { + mapper + .outbound_text(text) + .unwrap_or_else(|_| text.to_owned()) + }; + let references: Vec = references.iter().map(|r| original(r)).collect(); + self.sign_fingerprint(&original(store_path), nar_hash, nar_size, &references) + } + + fn sign_fingerprint( &self, store_path: &str, nar_hash: &str, diff --git a/backend/gradient-sources/src/store.rs b/backend/gradient-sources/src/store.rs index adf319982..31dbdf655 100644 --- a/backend/gradient-sources/src/store.rs +++ b/backend/gradient-sources/src/store.rs @@ -50,15 +50,6 @@ pub fn get_hash_from_path(path: String) -> Result<(String, String), SourceError> Ok((hash, package)) } -pub fn parse_drv_hash_name(drv_path: &str) -> Result<(String, String), SourceError> { - let (hash, rest) = drv_path.split_once('-').ok_or(SourceError::InvalidPath)?; - let name = rest.strip_suffix(".drv").ok_or(SourceError::InvalidPath)?; - if hash.is_empty() || name.is_empty() { - return Err(SourceError::InvalidPath); - } - Ok((hash.to_string(), name.to_string())) -} - pub fn get_cache_nar_location(base_path: String, hash: String) -> Result { let hash_hex = hash.as_str(); std::fs::create_dir_all(format!("{}/nars/{}", base_path, &hash_hex[0..2])).map_err(|e| { @@ -97,32 +88,3 @@ pub fn get_cache_nar_compressed_location( pub fn get_derivation_paths(derivations: &[MDerivation]) -> Vec { derivations.iter().map(|d| d.store_path()).collect() } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_canonical_drv_form() { - let (h, n) = - parse_drv_hash_name("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-hello-2.12.1.drv").unwrap(); - assert_eq!(h, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"); - assert_eq!(n, "hello-2.12.1"); - } - - #[test] - fn rejects_missing_drv_suffix() { - assert!(parse_drv_hash_name("abc-foo").is_err()); - } - - #[test] - fn rejects_missing_hash_or_name() { - assert!(parse_drv_hash_name("-foo.drv").is_err()); - assert!(parse_drv_hash_name("abc-.drv").is_err()); - } - - #[test] - fn rejects_missing_dash() { - assert!(parse_drv_hash_name("abcfoo.drv").is_err()); - } -} diff --git a/backend/gradient-ssh/src/auth.rs b/backend/gradient-ssh/src/auth.rs index 5e1078aee..1a70d0208 100644 --- a/backend/gradient-ssh/src/auth.rs +++ b/backend/gradient-ssh/src/auth.rs @@ -74,14 +74,9 @@ async fn load_session( used.last_used_at = Set(Some(now())); used.update(&handles.db.web_db).await?; - Ok(Some(Session { - handles: handles.clone(), - user, - project, - permissions, - caches, - closed: handles.db.shutdown.child_token(), - })) + Ok(Some( + Session::open(handles.clone(), user, project, permissions, caches).await?, + )) } #[cfg(test)] diff --git a/backend/gradient-ssh/src/basic_derivation.rs b/backend/gradient-ssh/src/basic_derivation.rs index 8ced72b4f..b0d54cea6 100644 --- a/backend/gradient-ssh/src/basic_derivation.rs +++ b/backend/gradient-ssh/src/basic_derivation.rs @@ -5,6 +5,8 @@ */ use crate::session::Session; +use crate::store_paths::StoredPathInfo; +use gradient_types::StorePathMapper; use gradient_util::nar::single_file_nar; use harmonia_protocol::valid_path_info::{UnkeyedValidPathInfo, ValidPathInfo}; use harmonia_store_aterm::print_derivation_aterm; @@ -51,9 +53,17 @@ fn drv_file(drv: &BasicDerivation) -> anyhow::Result { } pub async fn import(session: &Session, drv: &BasicDerivation) -> anyhow::Result { + let (info, nar) = stored_drv(drv, &session.mapper)?; + crate::ingest::import(session, &info, nar.as_slice()).await?; + Ok(info.path.full()) +} + +fn stored_drv( + drv: &BasicDerivation, + mapper: &StorePathMapper, +) -> anyhow::Result<(StoredPathInfo, Vec)> { let file = drv_file(drv)?; let nar = single_file_nar(&file.aterm, false); - let store_dir = StoreDir::default(); let info = ValidPathInfo { path: file.path.clone(), info: UnkeyedValidPathInfo { @@ -65,12 +75,11 @@ pub async fn import(session: &Session, drv: &BasicDerivation) -> anyhow::Result< ultimate: false, signatures: Default::default(), ca: Some(ContentAddress::Text(file.digest)), - store_dir: store_dir.clone(), + store_dir: StoreDir::default(), }, }; - crate::ingest::import(session, &info, nar.as_slice()).await?; - Ok(store_dir.display(&file.path).to_string()) + Ok((StoredPathInfo::inbound(mapper, &info)?, nar)) } #[cfg(test)] @@ -109,6 +118,27 @@ mod tests { } } + #[test] + fn a_derivation_path_is_computed_over_original_inputs() { + let mapper = gradient_types::StorePathMapper::for_cache( + std::sync::Arc::new(gradient_entity::RotatedHashMapping), + gradient_types::ids::CacheId::nil(), + ); + let drv = basic(ATERM); + + let (info, nar) = stored_drv(&drv, &mapper).expect("drv"); + + let original = gradient_types::OriginalStorePath::parse(PATH.trim()).expect("path"); + assert_eq!(info.path, mapper.inbound(&original)); + assert_eq!(info.original.to_string(), original.base()); + assert_eq!(info.nar_size, nar.len() as u64); + for input in &drv.inputs { + let input = + gradient_types::OriginalStorePath::parse(&input.to_string()).expect("input"); + assert!(info.references.contains(&mapper.inbound(&input).base())); + } + } + #[test] fn a_fixed_output_derivation_lands_at_the_path_nix_gave_it() { let file = drv_file(&basic(ATERM)).expect("drv"); diff --git a/backend/gradient-ssh/src/build_log.rs b/backend/gradient-ssh/src/build_log.rs index 57e010a96..54b4bce49 100644 --- a/backend/gradient-ssh/src/build_log.rs +++ b/backend/gradient-ssh/src/build_log.rs @@ -198,7 +198,7 @@ impl PlainLog { LogMessage::Message(m) => Some(String::from_utf8_lossy(&m.text).into_owned()), LogMessage::StartActivity(a) => { if let Some(Field::String(path)) = a.fields.first() { - let name = drv_name(&String::from_utf8_lossy(path)).to_owned(); + let name = drv_name(&String::from_utf8_lossy(path)); self.names.insert(a.id, name); } @@ -326,14 +326,15 @@ mod tests { .append_query_results([vec![attempt]]) .append_query_results([Vec::::new()]) .into_connection(); - let session = Session { - handles: SshHandles::from_state(&TestState::new().web_db(db).build()), - user: gradient_fixtures::rows::user(), - project: gradient_fixtures::rows::project(), - permissions: 0, - caches: vec![], - closed: Default::default(), - }; + let session = Session::open( + SshHandles::from_state(&TestState::new().web_db(db).build()), + gradient_fixtures::rows::user(), + gradient_fixtures::rows::project(), + 0, + vec![], + ) + .await + .unwrap(); let mut logs = BuildLogs { drv_paths: HashMap::from([(build.id, DRV.to_owned())]), open: vec![build.id], diff --git a/backend/gradient-ssh/src/build_request.rs b/backend/gradient-ssh/src/build_request.rs index ae06076ff..6cdb5b03d 100644 --- a/backend/gradient-ssh/src/build_request.rs +++ b/backend/gradient-ssh/src/build_request.rs @@ -6,6 +6,7 @@ use crate::build_wait::{BuildOutcome, wait}; use crate::session::Session; +use crate::store_paths::stored_derivation; use futures::StreamExt as _; use gradient_db::Hold; use gradient_db::build_request_task::ensure_build_request_task; @@ -15,7 +16,6 @@ use gradient_entity::evaluation::{EvaluationKind, EvaluationStatus}; use gradient_graph::{RecordBatch, Transition}; use gradient_types::permissions::Permission; use gradient_types::*; -use gradient_util::store_path::strip_nix_store_prefix; use gradient_wire::types::DiscoveredDerivation; use harmonia_protocol::log::LogMessage; use sea_orm::{ @@ -61,8 +61,8 @@ struct CacheDrvSource<'a> { impl DrvSource for CacheDrvSource<'_> { async fn read(&self, drv_path: &str) -> anyhow::Result> { - let base = strip_nix_store_prefix(drv_path); - let hash = base.split('-').next().unwrap_or_default(); + let path = StorePath::parse_stored(drv_path)?; + let hash = path.hash(); let handles = &self.session.handles; if gradient_db::cache_paths::served_path(&handles.db.web_db, &self.session.caches, hash) .await? @@ -81,7 +81,8 @@ impl DrvSource for CacheDrvSource<'_> { if let harmonia_file_nar::NarEvent::File { mut reader, .. } = event? { let mut contents = Vec::new(); reader.read_to_end(&mut contents).await?; - return parse_drv(&contents).map(Some); + let drv = parse_drv(&contents)?; + return stored_derivation(&self.session.mapper, drv).map(Some); } } @@ -89,10 +90,13 @@ impl DrvSource for CacheDrvSource<'_> { } } -pub fn drv_name(drv_path: &str) -> &str { - let base = drv_path.strip_prefix("/nix/store/").unwrap_or(drv_path); - let name = base.split_once('-').map_or(base, |(_, name)| name); - name.strip_suffix(".drv").unwrap_or(name) +pub fn drv_name(drv_path: &str) -> String { + let name = StorePath::parse_stored(drv_path) + .map_or_else(|_| drv_path.to_owned(), |path| path.name().to_owned()); + match name.strip_suffix(".drv") { + Some(name) => name.to_owned(), + None => name, + } } pub async fn closure( @@ -224,7 +228,7 @@ async fn create_evaluation( ) -> anyhow::Result { let handles = &session.handles; - let names: Vec<&str> = drv_paths.iter().map(|p| drv_name(p)).collect(); + let names: Vec = drv_paths.iter().map(|p| drv_name(p)).collect(); let tx = handles.db.web_db.inner().begin().await?; let commit = MCommit { id: CommitId::now_v7(), @@ -310,17 +314,16 @@ async fn add_batches( let derivations: Vec = closure .iter() .map(|(path, drv)| { - let attr = requested - .contains(path.as_str()) - .then(|| drv_name(path).to_string()); - let mut discovered = discovered_derivation(attr, strip_nix_store_prefix(path), drv); + let attr = requested.contains(path.as_str()).then(|| drv_name(path)); + let mut discovered = + discovered_derivation(attr, StorePath::parse_stored(path)?.base(), drv); for dependency in &mut discovered.dependencies { - *dependency = strip_nix_store_prefix(dependency); + *dependency = StorePath::parse_stored(dependency)?.base(); } - discovered + Ok(discovered) }) - .collect(); + .collect::>()?; for batch in derivations.chunks(RECORD_BATCH_SIZE) { let report = handles @@ -442,29 +445,23 @@ mod tests { #[tokio::test] async fn building_without_trigger_evaluation_is_refused() { let db = empty_db(); - let session = Session { - handles: SshHandles::from_state(&TestState::new().web_db(db.clone()).build()), - user: gradient_fixtures::rows::user(), - project: gradient_fixtures::rows::project(), - permissions: 0, - caches: vec![], - closed: Default::default(), - }; + let session = session(db.clone()).await; let e = start(&session, &[path("a")]).await.expect_err("refused"); assert!(e.to_string().contains("TriggerEvaluation"), "{e}"); assert!(db.into_transaction_log().is_empty()); } - fn session(db: sea_orm::DatabaseConnection) -> Session { - Session { - handles: SshHandles::from_state(&TestState::new().web_db(db).build()), - user: gradient_fixtures::rows::user(), - project: gradient_fixtures::rows::project(), - permissions: 0, - caches: vec![], - closed: Default::default(), - } + async fn session(db: sea_orm::DatabaseConnection) -> Session { + Session::open( + SshHandles::from_state(&TestState::new().web_db(db).build()), + gradient_fixtures::rows::user(), + gradient_fixtures::rows::project(), + 0, + vec![], + ) + .await + .unwrap() } fn build_request_task() -> MTask { @@ -503,7 +500,7 @@ mod tests { .append_query_results([vec![task.clone()]]) .append_query_results([vec![running.clone()]]) .into_connection(); - let session = session(db); + let session = session(db).await; let target = evaluation_for(&session, &[path("a")]) .await @@ -541,7 +538,7 @@ mod tests { .append_query_results([vec![commit()]]) .append_query_results([vec![created.clone()]]) .into_connection(); - let session = session(db); + let session = session(db).await; let target = evaluation_for(&session, &[path("a")]) .await diff --git a/backend/gradient-ssh/src/build_wait.rs b/backend/gradient-ssh/src/build_wait.rs index 5f656154a..c092acd4c 100644 --- a/backend/gradient-ssh/src/build_wait.rs +++ b/backend/gradient-ssh/src/build_wait.rs @@ -12,7 +12,6 @@ use gradient_derivation::Derivation; use gradient_entity::build::BuildStatus; use gradient_entity::evaluation::EvaluationStatus; use gradient_types::*; -use gradient_util::store_path::strip_nix_store_prefix; use harmonia_protocol::daemon::wire::types2::FailureStatus; use harmonia_protocol::log::LogMessage; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; @@ -173,11 +172,8 @@ async fn closure_derivations( let paths: HashSet<&str> = closure.iter().map(String::as_str).collect(); let hashes: Vec = closure .iter() - .filter_map(|p| { - strip_nix_store_prefix(p) - .split_once('-') - .map(|(h, _)| h.to_string()) - }) + .filter_map(|p| StorePath::parse_stored(p).ok()) + .map(|p| p.hash().to_owned()) .collect(); Ok(gradient_db::fetch_in_chunks(&hashes, |chunk| async { EDerivation::find() diff --git a/backend/gradient-ssh/src/daemon_build.rs b/backend/gradient-ssh/src/daemon_build.rs index d21605043..f5237a3d8 100644 --- a/backend/gradient-ssh/src/daemon_build.rs +++ b/backend/gradient-ssh/src/daemon_build.rs @@ -7,7 +7,9 @@ use crate::build_wait::{BuildOutcome, DrvResult}; use crate::session::Session; use crate::store::err; +use crate::store_paths::{original, stored}; use gradient_db::cache_paths::served_hashes; +use gradient_types::StorePathMapper; use harmonia_protocol::daemon::wire::types2::{ BuildResult, BuildResultFailure, BuildResultInner, BuildResultSuccess, KeyedBuildResult, SuccessStatus, @@ -52,7 +54,8 @@ pub fn build_paths( paths: Vec, ) -> impl ResultLog>> + Send + 'static { streamed(session, move |session, log| async move { - let derivations = requested_derivations(&session, &paths).await?; + let mapper = session.mapper.clone(); + let derivations = requested_derivations(&session, &mapper, &paths).await?; let outcome = if derivations.is_empty() { None } else { @@ -63,7 +66,7 @@ pub fn build_paths( ) }; - Ok(keyed_results(&paths, outcome)) + Ok(keyed_results(&mapper, &paths, outcome)) }) } @@ -79,12 +82,16 @@ pub fn build_derivation( .await .map_err(err)?; let result = outcome.results.iter().find(|(path, _)| *path == drv_path); - Ok(build_result(result_inner(result.map(|(_, r)| r)))) + Ok(build_result(result_inner( + &session.mapper, + result.map(|(_, r)| r), + ))) }) } async fn requested_derivations( session: &Session, + mapper: &StorePathMapper, paths: &[DerivedPath], ) -> DaemonResult> { let mut derivations = Vec::new(); @@ -92,7 +99,9 @@ async fn requested_derivations( for path in paths { match path { DerivedPath::Built { drv_path, .. } => match drv_path.as_ref() { - SingleDerivedPath::Opaque(drv) => derivations.push(format!("/nix/store/{drv}")), + SingleDerivedPath::Opaque(drv) => { + derivations.push(stored(mapper, drv).map_err(err)?.full()) + } SingleDerivedPath::Built { .. } => { return Err(err("dynamic derivations are not supported over SSH")); } @@ -101,13 +110,18 @@ async fn requested_derivations( } } - let hashes: Vec = opaque.iter().map(|p| p.hash().to_string()).collect(); + let hashes: Vec = opaque + .iter() + .map(|p| stored(mapper, p).map(|s| s.hash().to_owned())) + .collect::>() + .map_err(err)?; let served = served_hashes(&session.handles.db.web_db, &session.caches, &hashes) .await .map_err(err)?; - if let Some(missing) = opaque + if let Some((missing, _)) = opaque .iter() - .find(|p| !served.contains(&p.hash().to_string())) + .zip(&hashes) + .find(|(_, hash)| !served.contains(*hash)) { return Err(err(format!("{missing} is not in the project caches"))); } @@ -115,7 +129,11 @@ async fn requested_derivations( Ok(derivations) } -fn keyed_results(paths: &[DerivedPath], outcome: Option) -> Vec { +fn keyed_results( + mapper: &StorePathMapper, + paths: &[DerivedPath], + outcome: Option, +) -> Vec { let by_drv: HashMap = outcome .map(|o| o.results.into_iter().collect()) .unwrap_or_default(); @@ -125,9 +143,12 @@ fn keyed_results(paths: &[DerivedPath], outcome: Option) -> Vec match drv_path.as_ref() { - SingleDerivedPath::Opaque(drv) => { - result_inner(by_drv.get(&format!("/nix/store/{drv}"))) - } + SingleDerivedPath::Opaque(drv) => result_inner( + mapper, + stored(mapper, drv) + .ok() + .and_then(|drv| by_drv.get(&drv.full())), + ), SingleDerivedPath::Built { .. } => already_valid(), }, DerivedPath::Opaque(_) => already_valid(), @@ -141,9 +162,9 @@ fn keyed_results(paths: &[DerivedPath], outcome: Option) -> Vec) -> BuildResultInner { +fn result_inner(mapper: &StorePathMapper, result: Option<&DrvResult>) -> BuildResultInner { match result { - Some(Ok(outputs)) => built(outputs), + Some(Ok(outputs)) => built(mapper, outputs), Some(Err(failure)) => BuildResultInner::Failure(BuildResultFailure { status: failure.status, error_msg: failure.message.clone().into_bytes(), @@ -168,14 +189,11 @@ fn build_result(inner: BuildResultInner) -> BuildResult { } } -fn built(outputs: &BTreeMap) -> BuildResultInner { +fn built(mapper: &StorePathMapper, outputs: &BTreeMap) -> BuildResultInner { let built_outputs = outputs .iter() .filter_map(|(name, path)| { - let out_path = harmonia_store_path::StorePath::from_base_path( - path.strip_prefix("/nix/store/").unwrap_or(path), - ) - .ok()?; + let out_path = original(mapper, path).ok()?; Some(( name.parse().ok()?, UnkeyedRealisation { diff --git a/backend/gradient-ssh/src/exec.rs b/backend/gradient-ssh/src/exec.rs index 548b0e456..5c0e40c90 100644 --- a/backend/gradient-ssh/src/exec.rs +++ b/backend/gradient-ssh/src/exec.rs @@ -9,6 +9,7 @@ use crate::commands::Command; use crate::roots::Roots; use crate::session::Session; use crate::store::SshBackend; +use crate::store_paths::{original_text, stored_full}; use harmonia_protocol::log::LogMessage; use russh::ChannelId; use russh::server::{Handle, Msg}; @@ -152,9 +153,21 @@ async fn realise( derivations: &[String], add_root: Option, ) -> u32 { + let mapper = session.mapper.clone(); + let derivations = match derivations + .iter() + .map(|drv| stored_full(&mapper, drv)) + .collect::>>() + { + Ok(derivations) => derivations, + Err(error) => { + output.stderr(format!("error: {error}")).await; + return 1; + } + }; let (tx, mut rx) = mpsc::unbounded_channel::(); let build = async move { - crate::build_request::run(session, derivations, move |message| { + crate::build_request::run(session, &derivations, move |message| { let _ = tx.send(message); }) .await @@ -181,7 +194,11 @@ async fn realise( let mut paths = Vec::new(); for (_, result) in outcome.results { match result { - Ok(outputs) => paths.extend(outputs.into_values()), + Ok(outputs) => paths.extend( + outputs + .values() + .filter_map(|path| original_text(&mapper, path).ok()), + ), Err(failure) => { failed = true; output.stderr(failure.message).await; diff --git a/backend/gradient-ssh/src/ingest.rs b/backend/gradient-ssh/src/ingest.rs index ee77f699e..edf317331 100644 --- a/backend/gradient-ssh/src/ingest.rs +++ b/backend/gradient-ssh/src/ingest.rs @@ -5,6 +5,7 @@ */ use crate::session::Session; +use crate::store_paths::StoredPathInfo; use anyhow::Context as _; use async_compression::Level; use async_compression::tokio::write::ZstdEncoder; @@ -15,7 +16,6 @@ use gradient_types::events::cache::NarSigned; use gradient_types::permissions::Permission; use gradient_types::*; use gradient_util::nix_hash::normalize_nar_hash; -use harmonia_protocol::valid_path_info::ValidPathInfo; use harmonia_store_path_info::NarHash; use harmonia_utils_hash::{Algorithm, Context, HashFormat as _, Sha256}; use std::path::Path; @@ -41,7 +41,7 @@ pub struct Staged { pub async fn import( session: &Session, - info: &ValidPathInfo, + info: &StoredPathInfo, reader: impl AsyncRead + Unpin + Send, ) -> anyhow::Result<()> { let pending = stage(session, info, reader).await?; @@ -50,7 +50,7 @@ pub async fn import( pub async fn stage( session: &Session, - info: &ValidPathInfo, + info: &StoredPathInfo, mut reader: impl AsyncRead + Unpin + Send, ) -> anyhow::Result { anyhow::ensure!( @@ -60,49 +60,49 @@ pub async fn stage( ); let handles = &session.handles; - let hash = info.path.hash().to_string(); + let hash = info.path.hash(); let claim = claim_path( &handles.db.config, &handles.upload_admission, - &hash, - info.info.nar_size, + hash, + info.nar_size, ) .await?; - if let Some(known) = stored_path(&handles.db.web_db, &hash).await? { + if let Some(known) = stored_path(&handles.db.web_db, hash).await? { if let Some(claim) = claim { claim.committed(); } tokio::io::copy(&mut reader, &mut tokio::io::sink()).await?; - let announced = normalize_nar_hash(&info.info.nar_hash.as_sri().to_string()); + let announced = normalize_nar_hash(&info.nar_hash.as_sri().to_string()); if known.nar_hash.as_deref().map(normalize_nar_hash) == Some(announced) { return Ok(Pending::Known(known)); } - debug!(path = %info.path, "other content is stored for this path; dropping the upload"); + debug!(path = %info.original, "other content is stored for this path; dropping the upload"); return Ok(Pending::OtherContentStored); } let permit = claim.with_context(|| { format!( "{} was stored by another upload and removed again, retry", - info.path + info.original ) })?; let dir = handles.db.config.server.nar_upload_partial_dir(); tokio::fs::create_dir_all(&dir).await?; let file = tempfile::NamedTempFile::new_in(&dir)?; let (nar_hash, nar_size, file_hash, file_size) = - compress_into(reader, file.path(), info.info.nar_size).await?; + compress_into(reader, file.path(), info.nar_size).await?; anyhow::ensure!( - nar_size == info.info.nar_size, + nar_size == info.nar_size, "NAR size of {} does not match", - info.path + info.original ); anyhow::ensure!( - nar_hash == info.info.nar_hash, + nar_hash == info.nar_hash, "NAR hash of {} does not match", - info.path + info.original ); Ok(Pending::New(Staged { @@ -117,7 +117,7 @@ pub async fn stage( pub async fn commit( session: &Session, - info: &ValidPathInfo, + info: &StoredPathInfo, pending: Pending, ) -> anyhow::Result<()> { let handles = &session.handles; @@ -132,15 +132,7 @@ pub async fn commit( } Pending::OtherContentStored => Vec::new(), Pending::New(staged) => { - let store_path = format!("/nix/store/{}", info.path); - let references: Vec = - info.info.references.iter().map(|r| r.to_string()).collect(); - let deriver = info - .info - .deriver - .as_ref() - .map(|d| format!("/nix/store/{d}")); - let ca = info.info.ca.as_ref().map(|c| c.to_string()); + let store_path = info.path.full(); let nar_hash = staged.nar_hash.as_sri().to_string(); let committed = import_nar_reader( &handles.db.web_db, @@ -153,9 +145,9 @@ pub async fn commit( file_size: staged.file_size as i64, nar_size: staged.nar_size as i64, nar_hash: &nar_hash, - references: &references, - deriver: deriver.as_deref(), - ca: ca.as_deref(), + references: &info.references, + deriver: info.deriver.as_deref(), + ca: info.ca.as_deref(), }, targets, ) @@ -168,7 +160,7 @@ pub async fn commit( for cache in signed { handles.db.events.publish(NarSigned { cache, - hash: info.path.hash().to_string(), + hash: info.original.hash().to_string(), }); } @@ -248,26 +240,29 @@ mod tests { use gradient_test_support::state::TestState; use gradient_types::permissions::{Permission, mask_from}; use gradient_types::{CacheId, MCachedPath}; - use harmonia_protocol::valid_path_info::UnkeyedValidPathInfo; + use harmonia_protocol::valid_path_info::{UnkeyedValidPathInfo, ValidPathInfo}; use harmonia_store_path::{StoreDir, StorePath}; use sea_orm::DatabaseConnection; use std::sync::Arc; const NAR: &[u8] = b"nix-archive-1 stand-in bytes"; - fn session(db: DatabaseConnection, permissions: i64) -> Arc { - Arc::new(Session { - handles: SshHandles::from_state(&TestState::new().web_db(db).build()), - user: gradient_fixtures::rows::user(), - project: gradient_fixtures::rows::project(), + async fn session(db: DatabaseConnection, permissions: i64) -> Arc { + let session = Session::open( + SshHandles::from_state(&TestState::new().web_db(db).build()), + gradient_fixtures::rows::user(), + gradient_fixtures::rows::project(), permissions, - caches: vec![CacheId::now_v7()], - closed: Default::default(), - }) + vec![CacheId::now_v7()], + ) + .await + .unwrap(); + Arc::new(session) } - fn info(nar: &[u8]) -> ValidPathInfo { - ValidPathInfo { + fn info(nar: &[u8]) -> StoredPathInfo { + let mapper = StorePathMapper::for_cache(Arc::new(IdentityMapping), CacheId::nil()); + let client = ValidPathInfo { path: StorePath::from_base_path("0123456789abcdfghijklmnpqrsvwxyz-hello") .expect("path"), info: UnkeyedValidPathInfo { @@ -281,7 +276,8 @@ mod tests { ca: None, store_dir: StoreDir::default(), }, - } + }; + StoredPathInfo::inbound(&mapper, &client).expect("stored info") } #[tokio::test] @@ -289,9 +285,9 @@ mod tests { let db = mock_db() .append_query_results([Vec::::new()]) .into_connection(); - let session = session(db.clone(), mask_from(&[Permission::TriggerEvaluation])); + let session = session(db.clone(), mask_from(&[Permission::TriggerEvaluation])).await; let mut info = info(NAR); - info.info.nar_hash = NarHash::digest(b"other"); + info.nar_hash = NarHash::digest(b"other"); let e = import(&session, &info, NAR) .await @@ -306,9 +302,9 @@ mod tests { let db = mock_db() .append_query_results([Vec::::new()]) .into_connection(); - let session = session(db, mask_from(&[Permission::TriggerEvaluation])); + let session = session(db, mask_from(&[Permission::TriggerEvaluation])).await; let mut info = info(NAR); - info.info.nar_size += 1; + info.nar_size += 1; let e = import(&session, &info, NAR) .await @@ -319,7 +315,7 @@ mod tests { #[tokio::test] async fn copying_in_without_trigger_evaluation_is_refused() { let db = empty_db(); - let session = session(db, 0); + let session = session(db, 0).await; let e = import(&session, &info(NAR), NAR) .await @@ -332,9 +328,9 @@ mod tests { let db = mock_db() .append_query_results([Vec::::new()]) .into_connection(); - let session = session(db, mask_from(&[Permission::TriggerEvaluation])); + let session = session(db, mask_from(&[Permission::TriggerEvaluation])).await; let mut info = info(NAR); - info.info.nar_size = 4; + info.nar_size = 4; let e = import(&session, &info, NAR).await.expect_err("too long"); assert!(e.to_string().contains("longer than announced"), "{e}"); @@ -353,7 +349,7 @@ mod tests { let db = mock_db() .append_query_results([vec![existing]]) .into_connection(); - let session = session(db.clone(), mask_from(&[Permission::TriggerEvaluation])); + let session = session(db.clone(), mask_from(&[Permission::TriggerEvaluation])).await; let mut upload = NAR; import(&session, &info(NAR), &mut upload) @@ -378,10 +374,10 @@ mod tests { let db = mock_db() .append_query_results([vec![existing]]) .into_connection(); - let session = session(db, mask_from(&[Permission::TriggerEvaluation])); + let session = session(db, mask_from(&[Permission::TriggerEvaluation])).await; let info = info(NAR); let (holder, mut admitted) = session.handles.upload_admission.open_session("holder"); - holder.request(1, ObjectKey::Nar(info.path.hash().to_string()), 1, false); + holder.request(1, ObjectKey::Nar(info.path.hash().to_owned()), 1, false); let Some(Admitted::Granted { permit, .. }) = admitted.recv().await else { panic!("the holder leads the path"); }; diff --git a/backend/gradient-ssh/src/lib.rs b/backend/gradient-ssh/src/lib.rs index 41d097f78..b5b35e100 100644 --- a/backend/gradient-ssh/src/lib.rs +++ b/backend/gradient-ssh/src/lib.rs @@ -20,6 +20,7 @@ mod roots; mod server; mod session; mod store; +mod store_paths; use gradient_core::ServerState; use russh::server::Server as _; diff --git a/backend/gradient-ssh/src/session.rs b/backend/gradient-ssh/src/session.rs index b5ceaae04..450827459 100644 --- a/backend/gradient-ssh/src/session.rs +++ b/backend/gradient-ssh/src/session.rs @@ -36,10 +36,30 @@ pub struct Session { pub project: MProject, pub permissions: PermissionMask, pub caches: Vec, + pub mapper: StorePathMapper, pub closed: CancellationToken, } impl Session { + pub async fn open( + handles: SshHandles, + user: MUser, + project: MProject, + permissions: PermissionMask, + caches: Vec, + ) -> Result { + let mapper = handles.db.project_mapper(project.id).await?; + Ok(Self { + closed: handles.db.shutdown.child_token(), + handles, + user, + project, + permissions, + caches, + mapper, + }) + } + pub fn may(&self, permission: Permission) -> bool { mask_grants(self.permissions, permission) } diff --git a/backend/gradient-ssh/src/store.rs b/backend/gradient-ssh/src/store.rs index 118b823cb..da0672106 100644 --- a/backend/gradient-ssh/src/store.rs +++ b/backend/gradient-ssh/src/store.rs @@ -5,10 +5,12 @@ */ use crate::session::Session; +use crate::store_paths::{StoredPathInfo, original, stored, stored_hash}; use futures::{Stream, StreamExt as _}; use gradient_daemon::backend::{Backend, ConnInfo}; use gradient_daemon::journal::Journal; use gradient_db::cache_paths::{ServedPath, served_hashes, served_path}; +use gradient_types::StorePathMapper; use gradient_util::nix_hash::normalize_nar_hash; use harmonia_protocol::daemon::wire::types2::{ BuildMode, BuildResult, KeyedBuildResult, QueryMissingResult, @@ -54,6 +56,7 @@ impl Backend for SshBackend { fn handler(self: &Arc, _conn: ConnInfo) -> CacheStore { CacheStore { + mapper: self.session.mapper.clone(), session: self.session.clone(), } } @@ -69,6 +72,7 @@ impl Backend for SshBackend { pub struct CacheStore { session: Arc, + mapper: StorePathMapper, } const CONCURRENT_COMMITS: usize = 8; @@ -85,20 +89,20 @@ pub(crate) fn err(e: impl Display) -> DaemonError { DaemonError::custom(format!("{e:#}")) } -fn parse_store_path(path: &str) -> anyhow::Result { - let base = path.strip_prefix("/nix/store/").unwrap_or(path); - StorePath::from_base_path(base).map_err(|e| anyhow::anyhow!("{base}: {e}")) -} - impl CacheStore { + fn stored_info(&self, info: &ValidPathInfo) -> DaemonResult { + StoredPathInfo::inbound(&self.mapper, info).map_err(err) + } + async fn served(&self, path: &StorePath) -> DaemonResult> { - served_path( - &self.session.handles.db.web_db, - &self.session.caches, - &path.hash().to_string(), - ) - .await - .map_err(err) + let path = stored(&self.mapper, path).map_err(err)?; + self.served_stored(path.hash()).await + } + + async fn served_stored(&self, hash: &str) -> DaemonResult> { + served_path(&self.session.handles.db.web_db, &self.session.caches, hash) + .await + .map_err(err) } async fn copy_in( @@ -106,12 +110,13 @@ impl CacheStore { info: &ValidPathInfo, reader: impl AsyncBufRead + Send + Unpin, ) -> DaemonResult<()> { - if self.served(&info.path).await?.is_some() { + let info = self.stored_info(info)?; + if self.served_stored(info.path.hash()).await?.is_some() { tokio::io::copy(&mut std::pin::pin!(reader), &mut tokio::io::sink()).await?; return Ok(()); } - crate::ingest::import(&self.session, info, reader) + crate::ingest::import(&self.session, &info, reader) .await .map_err(err) } @@ -131,7 +136,7 @@ impl CacheStore { gradient_db::graph::runtime_closure::references_for_hash(&handles.db.web_db, &row.hash) .await? .iter() - .map(|r| parse_store_path(r)) + .map(|r| original(&self.mapper, r)) .collect::>()?; let signatures = row @@ -149,7 +154,11 @@ impl CacheStore { .collect(); Ok(UnkeyedValidPathInfo { - deriver: row.deriver.as_deref().map(parse_store_path).transpose()?, + deriver: row + .deriver + .as_deref() + .map(|d| original(&self.mapper, d)) + .transpose()?, nar_hash, references, registration_time: None, @@ -188,7 +197,11 @@ impl DaemonStore for CacheStore { _substitute: bool, ) -> impl ResultLog> + Send + 'a { async move { - let hashes: Vec = paths.iter().map(|p| p.hash().to_string()).collect(); + let hashes: Vec = paths + .iter() + .map(|p| stored(&self.mapper, p).map(|s| s.hash().to_owned())) + .collect::>() + .map_err(err)?; let served = served_hashes( &self.session.handles.db.web_db, &self.session.caches, @@ -199,8 +212,9 @@ impl DaemonStore for CacheStore { Ok(paths .iter() - .filter(|p| served.contains(&p.hash().to_string())) - .cloned() + .zip(&hashes) + .filter(|(_, hash)| served.contains(*hash)) + .map(|(p, _)| p.clone()) .collect()) } .empty_logs() @@ -225,15 +239,10 @@ impl DaemonStore for CacheStore { hash: &'a StorePathHash, ) -> impl ResultLog>> + Send + 'a { async move { - let row = served_path( - &self.session.handles.db.web_db, - &self.session.caches, - &hash.to_string(), - ) - .await - .map_err(err)?; + let hash = stored_hash(&self.mapper, hash).map_err(err)?; + let row = self.served_stored(&hash).await?; - row.map(|r| parse_store_path(&format!("{}-{}", r.hash, r.package)).map_err(err)) + row.map(|r| original(&self.mapper, &format!("{}-{}", r.hash, r.package)).map_err(err)) .transpose() } .empty_logs() @@ -244,17 +253,15 @@ impl DaemonStore for CacheStore { path: &'s StorePath, ) -> impl ResultLog>> + Send + 's { async move { - if self.served(path).await?.is_none() { + let stored = stored(&self.mapper, path).map_err(err)?; + if self.served_stored(stored.hash()).await?.is_none() { return Err(err(format!("{path} is not in the project caches"))); } - crate::nar::open_raw( - &self.session.handles.db.storage.nar_storage, - &path.hash().to_string(), - ) - .await - .map_err(err)? - .ok_or_else(|| err(format!("{path} has no stored NAR"))) + crate::nar::open_raw(&self.session.handles.db.storage.nar_storage, stored.hash()) + .await + .map_err(err)? + .ok_or_else(|| err(format!("{path} has no stored NAR"))) } .empty_logs() } @@ -293,7 +300,8 @@ impl DaemonStore for CacheStore { let mut commits = JoinSet::new(); while let Some(item) = stream.next().await { let AddToStoreItem { info, reader } = item?; - if self.served(&info.path).await?.is_some() { + let info = self.stored_info(&info)?; + if self.served_stored(info.path.hash()).await?.is_some() { tokio::io::copy(&mut std::pin::pin!(reader), &mut tokio::io::sink()).await?; continue; } @@ -415,6 +423,7 @@ mod tests { use super::*; use crate::session::SshHandles; use gradient_daemon::server::{next_conn, serve_stream}; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; use gradient_fixtures::mock::{empty_db, mock_db}; use gradient_test_support::state::TestState; use gradient_types::CacheId; @@ -431,15 +440,21 @@ mod tests { tokio::io::WriteHalf, >; - fn session_over(db: DatabaseConnection) -> Arc { - Arc::new(Session { - handles: SshHandles::from_state(&TestState::new().web_db(db).build()), - user: gradient_fixtures::rows::user(), - project: gradient_fixtures::rows::project(), - permissions: 0, - caches: vec![CacheId::now_v7()], - closed: Default::default(), - }) + async fn session_over(db: DatabaseConnection) -> Arc { + session_in(TestState::new().web_db(db)).await + } + + async fn session_in(state: TestState) -> Arc { + let session = Session::open( + SshHandles::from_state(&state.build()), + gradient_fixtures::rows::user(), + gradient_fixtures::rows::project(), + 0, + vec![CacheId::now_v7()], + ) + .await + .unwrap(); + Arc::new(session) } async fn connect(session: Arc) -> (JoinSet<()>, TestClient) { @@ -502,7 +517,7 @@ mod tests { .append_query_results([vec![served_row()]]) .append_query_results([vec![references_row()]]) .into_connection(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; let info = client .query_path_info(&hello()) @@ -514,13 +529,49 @@ mod tests { assert!(info.references.contains(&hello())); } + #[tokio::test] + async fn path_info_names_the_path_the_client_asked_for() { + let (original, stored) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let mut row = served_row(); + row.insert("hash".to_owned(), Value::from(stored)); + let references = BTreeMap::from([ + ("hash".to_owned(), Value::from(stored)), + ( + "references".to_owned(), + Value::from(Some(format!("{stored}-hello"))), + ), + ]); + let db = mock_db() + .append_query_results([vec![row]]) + .append_query_results([vec![references]]) + .into_connection(); + let session = session_in( + TestState::new() + .web_db(db.clone()) + .store_path_mapping(Arc::new(RotatedHashMapping)), + ) + .await; + let (_server, mut client) = connect(session).await; + let asked = StorePath::from_base_path(&format!("{original}-hello")).expect("path"); + + let info = client + .query_path_info(&asked) + .await + .expect("info") + .expect("served"); + assert!(info.references.contains(&asked)); + let log = format!("{:?}", db.into_transaction_log()); + assert!(log.contains(stored), "{log}"); + assert!(!log.contains(original), "{log}"); + } + #[tokio::test] async fn a_path_outside_the_project_caches_is_invalid() { let db = mock_db() .append_query_results([Vec::>::new()]) .append_query_results([Vec::>::new()]) .into_connection(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; assert!(!client.is_valid_path(&hello()).await.expect("valid")); assert!( @@ -541,7 +592,7 @@ mod tests { #[tokio::test] async fn the_client_is_not_trusted() { - let session = session_over(empty_db()); + let session = session_over(empty_db()).await; let backend = Arc::new(SshBackend::new(session)); let store = backend.handler(next_conn(None)); assert_eq!(store.trust_level(), Some(TrustLevel::NotTrusted)); @@ -550,7 +601,7 @@ mod tests { #[tokio::test] async fn garbage_collection_is_refused() { let db = empty_db(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; let err = client.find_roots().await.expect_err("refused"); assert!(err.to_string().contains("unimplemented"), "{err}"); @@ -561,7 +612,7 @@ mod tests { let db = mock_db() .append_query_results([vec![served_row()]]) .into_connection(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; let nar = nar_of(b"already there").await; let info = ValidPathInfo { path: hello(), @@ -596,7 +647,7 @@ mod tests { let db = mock_db() .append_query_results([vec![served_row()]]) .into_connection(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; let results = client .build_paths_with_results(&[DerivedPath::Opaque(hello())], BuildMode::Normal) @@ -614,7 +665,7 @@ mod tests { let db = mock_db() .append_query_results([Vec::>::new()]) .into_connection(); - let (_server, mut client) = connect(session_over(db)).await; + let (_server, mut client) = connect(session_over(db).await).await; let err = client .build_paths(&[DerivedPath::Opaque(hello())], BuildMode::Normal) diff --git a/backend/gradient-ssh/src/store_paths.rs b/backend/gradient-ssh/src/store_paths.rs new file mode 100644 index 000000000..4f01a9911 --- /dev/null +++ b/backend/gradient-ssh/src/store_paths.rs @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use gradient_derivation::Derivation; +use gradient_types::{OriginalStorePath, StoreHash, StorePath, StorePathMapper}; +use harmonia_protocol::valid_path_info::ValidPathInfo; +use harmonia_store_path::{StorePath as ClientPath, StorePathHash}; +use harmonia_store_path_info::NarHash; + +pub fn stored(mapper: &StorePathMapper, path: &ClientPath) -> anyhow::Result { + Ok(mapper.inbound(&OriginalStorePath::parse(&path.to_string())?)) +} + +pub fn stored_hash(mapper: &StorePathMapper, hash: &StorePathHash) -> anyhow::Result { + Ok(mapper + .inbound_hash(StoreHash::parse(&hash.to_string())?) + .to_string()) +} + +pub fn original(mapper: &StorePathMapper, stored: &str) -> anyhow::Result { + let original = mapper.outbound_text(stored)?; + let base = original + .strip_prefix(gradient_types::STORE_DIR) + .map_or(original.as_str(), |rest| rest.trim_start_matches('/')); + ClientPath::from_base_path(base).map_err(|e| anyhow::anyhow!("{base}: {e}")) +} + +pub fn original_text(mapper: &StorePathMapper, stored: &str) -> anyhow::Result { + Ok(mapper.outbound_text(stored)?) +} + +pub fn stored_full(mapper: &StorePathMapper, original: &str) -> anyhow::Result { + Ok(mapper.inbound(&OriginalStorePath::parse(original)?).full()) +} + +pub fn stored_derivation( + mapper: &StorePathMapper, + mut drv: Derivation, +) -> anyhow::Result { + for output in drv.outputs.iter_mut().filter(|o| !o.path.is_empty()) { + output.path = mapper.inbound_text(&output.path)?; + } + for (input, _) in &mut drv.input_derivations { + *input = mapper.inbound_text(input)?; + } + for source in &mut drv.input_sources { + *source = mapper.inbound_text(source)?; + } + Ok(drv) +} + +pub struct StoredPathInfo { + pub path: StorePath, + pub original: ClientPath, + pub references: Vec, + pub deriver: Option, + pub ca: Option, + pub nar_hash: NarHash, + pub nar_size: u64, +} + +impl StoredPathInfo { + pub fn inbound(mapper: &StorePathMapper, info: &ValidPathInfo) -> anyhow::Result { + Ok(Self { + path: stored(mapper, &info.path)?, + original: info.path.clone(), + references: info + .info + .references + .iter() + .map(|r| stored(mapper, r).map(|p| p.base())) + .collect::>()?, + deriver: info + .info + .deriver + .as_ref() + .map(|d| stored(mapper, d).map(|p| p.full())) + .transpose()?, + ca: info.info.ca.as_ref().map(|c| c.to_string()), + nar_hash: info.info.nar_hash, + nar_size: info.info.nar_size, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use gradient_derivation::DerivationOutput; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + use gradient_types::CacheId; + use std::sync::Arc; + + #[test] + fn a_derivation_read_from_the_cache_names_stored_paths() { + let (o, s) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let mapper = StorePathMapper::for_cache(Arc::new(RotatedHashMapping), CacheId::nil()); + let output = |path: String| DerivationOutput { + name: "out".into(), + path, + hash_algo: String::new(), + hash: String::new(), + }; + let drv = Derivation { + outputs: vec![output(format!("/nix/store/{o}-hello"))], + input_derivations: vec![(format!("/nix/store/{o}-dep.drv"), vec!["out".into()])], + input_sources: vec![format!("/nix/store/{o}-src")], + system: "x86_64-linux".into(), + builder: format!("/nix/store/{o}-bash/bin/bash"), + args: Vec::new(), + environment: Default::default(), + }; + + let stored = stored_derivation(&mapper, drv).expect("stored"); + + assert_eq!( + stored.outputs, + vec![output(format!("/nix/store/{s}-hello"))] + ); + assert_eq!( + stored.input_derivations, + vec![(format!("/nix/store/{s}-dep.drv"), vec!["out".to_owned()])] + ); + assert_eq!(stored.input_sources, vec![format!("/nix/store/{s}-src")]); + assert_eq!(stored.builder, format!("/nix/store/{o}-bash/bin/bash")); + } +} diff --git a/backend/gradient-storage/src/nar.rs b/backend/gradient-storage/src/nar.rs index a5a63bb98..69d66eab5 100644 --- a/backend/gradient-storage/src/nar.rs +++ b/backend/gradient-storage/src/nar.rs @@ -529,55 +529,56 @@ impl NarStore { } } - fn eval_cache_path(&self, fingerprint: &str) -> Path { - Path::from(format!("{}eval-cache/{}", self.prefix, fingerprint)) + fn eval_cache_path(&self, key: &str) -> Path { + Path::from(format!("{}{}", self.prefix, key)) } - pub async fn put_eval_cache(&self, fingerprint: &str, data: Vec) -> Result<()> { - self.put_object(self.eval_cache_path(fingerprint), data) - .await + pub async fn put_eval_cache(&self, key: &str, data: Vec) -> Result<()> { + self.put_object(self.eval_cache_path(key), data).await } - pub async fn get_eval_cache(&self, fingerprint: &str) -> Result>> { - self.get_object(&self.eval_cache_path(fingerprint)).await + pub async fn get_eval_cache(&self, key: &str) -> Result>> { + self.get_object(&self.eval_cache_path(key)).await } pub async fn get_eval_cache_stream( &self, - fingerprint: &str, + key: &str, ) -> Result>)>> { - self.get_stream_object(&self.eval_cache_path(fingerprint), None) + self.get_stream_object(&self.eval_cache_path(key), None) .await } pub async fn presigned_eval_cache_get_url( &self, - fingerprint: &str, + key: &str, expires_in: std::time::Duration, ) -> Result> { - self.presign_object( - reqwest::Method::GET, - &self.eval_cache_path(fingerprint), - expires_in, - ) - .await + self.presign_object(reqwest::Method::GET, &self.eval_cache_path(key), expires_in) + .await } pub async fn presigned_eval_cache_put_url( &self, - fingerprint: &str, + key: &str, expires_in: std::time::Duration, ) -> Result> { - self.presign_object( - reqwest::Method::PUT, - &self.eval_cache_path(fingerprint), - expires_in, - ) - .await + self.presign_object(reqwest::Method::PUT, &self.eval_cache_path(key), expires_in) + .await } - pub async fn delete_eval_cache(&self, fingerprint: &str) -> Result<()> { - self.delete_object(&self.eval_cache_path(fingerprint)).await + pub async fn copy_eval_cache(&self, from_key: &str, to_key: &str) -> Result<()> { + self.inner + .copy( + &self.eval_cache_path(from_key), + &self.eval_cache_path(to_key), + ) + .await + .with_context(|| format!("copy eval-cache object {from_key} to {to_key}")) + } + + pub async fn delete_eval_cache(&self, key: &str) -> Result<()> { + self.delete_object(&self.eval_cache_path(key)).await } fn blob_path(&self, project: uuid::Uuid, hash: &[u8; 32]) -> Path { @@ -638,6 +639,19 @@ impl NarStore { Ok(hashes) } + pub async fn list_eval_cache_keys(&self) -> Result> { + let listed = Path::from(format!("{}eval-cache", self.prefix)); + let mut stream = self.inner.list(Some(&listed)); + let mut keys = Vec::new(); + while let Some(item) = stream.next().await { + let location = item.context("list eval-cache")?.location.to_string(); + if let Some(key) = location.strip_prefix(self.prefix.as_str()) { + keys.push(key.to_owned()); + } + } + Ok(keys) + } + pub async fn list_blobs(&self) -> Result> { let prefix = Path::from(format!("{}build-request-blobs", self.prefix)); let mut stream = self.inner.list(Some(&prefix)); @@ -786,22 +800,22 @@ impl NarStore { .await? .map(|url| StorageTarget::Put { url }) .ok_or_else(missing), - ObjectKey::EvalCache(fingerprint) => self - .presigned_eval_cache_put_url(fingerprint, PRESIGN_TTL) + ObjectKey::EvalCache(key) => self + .presigned_eval_cache_put_url(key, PRESIGN_TTL) .await? .map(|url| StorageTarget::Put { url }) .ok_or_else(missing), } } - pub async fn verify_eval_cache(&self, fingerprint: &str, size: u64) -> Result<()> { + pub async fn verify_eval_cache(&self, key: &str, size: u64) -> Result<()> { let meta = self - .head_object(&self.eval_cache_path(fingerprint)) + .head_object(&self.eval_cache_path(key)) .await? - .with_context(|| format!("eval-cache object for {fingerprint} is missing"))?; + .with_context(|| format!("eval-cache object {key} is missing"))?; anyhow::ensure!( meta.size == size, - "eval-cache object for {fingerprint} is {} bytes, expected {size}", + "eval-cache object {key} is {} bytes, expected {size}", meta.size ); Ok(()) @@ -840,7 +854,7 @@ mod tests { let store = NarStore::local(dir.path()).unwrap(); for object in [ ObjectKey::Nar("a".repeat(32)), - ObjectKey::EvalCache("fp".into()), + ObjectKey::EvalCache("eval-cache/c/fp".into()), ] { assert_eq!( store @@ -871,7 +885,7 @@ mod tests { .unwrap(); assert!(matches!(nar, StorageTarget::Put { .. }), "{nar:?}"); let blob = store - .upload_target(&ObjectKey::EvalCache("fp".into()), 1024) + .upload_target(&ObjectKey::EvalCache("eval-cache/c/fp".into()), 1024) .await .unwrap(); assert!(matches!(blob, StorageTarget::Put { .. }), "{blob:?}"); @@ -1012,6 +1026,23 @@ mod tests { assert_eq!(got, payload); } + #[tokio::test] + async fn a_copied_eval_cache_blob_reads_back_under_its_new_key() { + let (_d, store) = local_store(); + store + .put_eval_cache("eval-cache/a/fp", b"blob".to_vec()) + .await + .expect("put"); + store + .copy_eval_cache("eval-cache/a/fp", "eval-cache/b/fp") + .await + .expect("copy"); + assert_eq!( + store.get_eval_cache("eval-cache/b/fp").await.expect("get"), + Some(b"blob".to_vec()) + ); + } + #[tokio::test] async fn exists_reflects_presence() { let (_d, store) = local_store(); diff --git a/backend/gradient-storage/src/source_nar.rs b/backend/gradient-storage/src/source_nar.rs index e1102539a..74e82273c 100644 --- a/backend/gradient-storage/src/source_nar.rs +++ b/backend/gradient-storage/src/source_nar.rs @@ -56,7 +56,7 @@ pub async fn source_nar_from_bytes(nar_bytes: Vec) -> Result { let store_path_obj = make_store_path_from_ca(&store_dir, name, ca); let store_hash = store_path_obj.hash().to_string(); - let store_path = gradient_types::StorePath::from_parts(store_hash.clone(), "source").full(); + let store_path = gradient_types::StorePath::from_stored(store_hash.clone(), "source").full(); let compressed_bytes = zstd::encode_all( std::io::Cursor::new(&nar_bytes), diff --git a/backend/gradient-test-support/src/cache_fixture.rs b/backend/gradient-test-support/src/cache_fixture.rs index c153534e0..97f4e5b14 100644 --- a/backend/gradient-test-support/src/cache_fixture.rs +++ b/backend/gradient-test-support/src/cache_fixture.rs @@ -142,7 +142,6 @@ async fn public_cache_storing_nar(served: bool) -> Arc { let db = mock_db() .append_query_results([vec![cache_row]]) - .append_query_results([Vec::::new()]) .append_query_results([served_gate(served)]) .into_connection(); diff --git a/backend/gradient-test-support/src/state.rs b/backend/gradient-test-support/src/state.rs index 505827df3..723232a01 100644 --- a/backend/gradient-test-support/src/state.rs +++ b/backend/gradient-test-support/src/state.rs @@ -9,8 +9,9 @@ use crate::fakes::email::InMemoryEmailSender; use crate::log_storage::NoopLogStorage; use axum_test::TestServer; use chrono::{DateTime, Utc}; -use gradient_core::{Graph, InstanceMemory, ServerState}; +use gradient_core::{Graph, InstanceMemory, ServerExtensions, ServerState}; use gradient_db::{CacheDb, WebDb, WorkerDb}; +use gradient_entity::StorePathMapping; use gradient_fixtures::mock::empty_db; use gradient_notify::EmailSender; use gradient_storage::{LogStorage, NarStore}; @@ -33,6 +34,7 @@ pub struct TestState { email: Arc, graph: Option>, started_at: Option>, + extensions: ServerExtensions, } impl Default for TestState { @@ -48,6 +50,7 @@ impl Default for TestState { email: Arc::new(InMemoryEmailSender::new()), graph: None, started_at: None, + extensions: ServerExtensions::default(), } } } @@ -107,6 +110,11 @@ impl TestState { self } + pub fn store_path_mapping(mut self, mapping: Arc) -> Self { + self.extensions.store_path_mapping = mapping; + self + } + pub fn build(self) -> Arc { let mut config = RuntimeConfig::from_cli(&self.cli).expect("valid test config"); for change in self.config_changes { @@ -132,6 +140,7 @@ impl TestState { events: gradient_types::EventBus::default(), wake: Default::default(), graph: self.graph.unwrap_or_else(Graph::stub), + extensions: self.extensions, memory: InstanceMemory { db: Default::default(), last_used_stamps: gradient_core::last_used_stamps(), diff --git a/backend/gradient-types/src/entity_aliases.rs b/backend/gradient-types/src/entity_aliases.rs index a9a61d282..381f8fb3e 100644 --- a/backend/gradient-types/src/entity_aliases.rs +++ b/backend/gradient-types/src/entity_aliases.rs @@ -12,7 +12,10 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; pub use gradient_entity::team_user::{TeamMemberSource, TeamRole}; -pub use gradient_entity::{STORE_DIR, StorePath}; +pub use gradient_entity::{ + IdentityMapping, MappingContext, MappingError, OriginalStorePath, STORE_DIR, StoreHash, + StorePath, StorePathMapper, StorePathMapping, +}; #[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)] pub struct ListItem { diff --git a/backend/gradient-util/src/store_path.rs b/backend/gradient-util/src/store_path.rs index abf388a43..b7a335cb8 100644 --- a/backend/gradient-util/src/store_path.rs +++ b/backend/gradient-util/src/store_path.rs @@ -12,10 +12,6 @@ pub fn nix_store_path(hash_name: &str) -> String { } } -pub fn strip_nix_store_prefix(path: &str) -> String { - path.strip_prefix("/nix/store/").unwrap_or(path).to_string() -} - pub fn strip_store_prefix(path: &str) -> &str { path.strip_prefix("/nix/store/").unwrap_or(path) } @@ -25,7 +21,7 @@ mod tests { use super::*; #[test] - fn nix_store_path_and_strip_round_trip() { + fn nix_store_path_adds_the_store_prefix() { for (input, expected) in [ ("abc123-foo", "/nix/store/abc123-foo"), ("/nix/store/abc123-foo", "/nix/store/abc123-foo"), @@ -34,14 +30,6 @@ mod tests { ] { assert_eq!(nix_store_path(input), expected, "{input}"); } - for (input, expected) in [ - ("/nix/store/abc123-foo", "abc123-foo"), - ("abc123-foo", "abc123-foo"), - ("/tmp/foo", "/tmp/foo"), - ("", ""), - ] { - assert_eq!(strip_nix_store_prefix(input), expected, "{input}"); - } } #[test] diff --git a/backend/gradient-web/src/endpoints/board.rs b/backend/gradient-web/src/endpoints/board.rs index edba209f4..070d18d66 100644 --- a/backend/gradient-web/src/endpoints/board.rs +++ b/backend/gradient-web/src/endpoints/board.rs @@ -14,6 +14,7 @@ use crate::helpers::ok_json; use crate::lookup::require; use crate::metrics_scope::MetricsScope; use crate::resource::Resource; +use crate::store_paths::original_text; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{Path, Query, State}; use axum::response::Response; @@ -376,6 +377,27 @@ pub struct JobDerivationView { pub pname: Option, } +fn outbound_drv_paths(value: &mut serde_json::Value, mapper: &StorePathMapper) { + match value { + serde_json::Value::Object(fields) => { + for (key, field) in fields.iter_mut() { + match field { + serde_json::Value::String(path) if key == "drv_path" => { + *path = original_text(mapper, std::mem::take(path)); + } + _ => outbound_drv_paths(field, mapper), + } + } + } + serde_json::Value::Array(items) => { + for item in items { + outbound_drv_paths(item, mapper); + } + } + _ => {} + } +} + fn snapshot_derivations( job_context: &serde_json::Value, ) -> Vec<(DerivationBuildId, String, Option)> { @@ -494,11 +516,13 @@ pub async fn get_dispatched_job( let scope = MetricsScope::resolve(&state.web_db, &maybe_user).await?; - if let Some(c) = scheduler.candidate_detail(DispatchedJobId::from(id)).await { + if let Some(mut c) = scheduler.candidate_detail(DispatchedJobId::from(id)).await { if !scope.allows(&Uuid::from(c.project)) { return Err(WebError::not_found(Resource::Job)); } + outbound_drv_paths(&mut c.job_context, &state.project_mapper(c.project).await?); + let (project, worker_name) = job_names(&state.web_db, &scope, c.project.into(), &c.worker_id).await?; @@ -534,7 +558,13 @@ pub async fn get_dispatched_job( derivation_build_id: c.derivation_build.map(Into::into), derivations, evaluation_id: c.evaluation_id.into(), - evaluation: eval_job_evaluation(&state.web_db, c.kind, c.evaluation_id).await?, + evaluation: eval_job_evaluation( + &state.web_db, + &state.project_mapper(c.project).await?, + c.kind, + c.evaluation_id, + ) + .await?, pname: c.pname, score_breakdown: c.score_breakdown, worker_context: c.worker_context, @@ -546,7 +576,7 @@ pub async fn get_dispatched_job( })); } - let j = require::( + let mut j = require::( &state.web_db, DispatchedJobId::from(id), ) @@ -556,6 +586,12 @@ pub async fn get_dispatched_job( return Err(WebError::not_found(Resource::Job)); } + let mapper = state.project_mapper(j.project).await?; + outbound_drv_paths(&mut j.job_context, &mapper); + if let Some(candidates) = &mut j.candidates { + outbound_drv_paths(candidates, &mapper); + } + let (project, worker_name) = job_names(&state.web_db, &scope, j.project.into(), &j.worker_id).await?; @@ -628,7 +664,8 @@ pub async fn get_dispatched_job( derivation_build_id: shared_build_id.map(Into::into), derivations, evaluation_id: j.evaluation_id.into(), - evaluation: eval_job_evaluation(&state.web_db, i16::from(j.kind), j.evaluation_id).await?, + evaluation: eval_job_evaluation(&state.web_db, &mapper, i16::from(j.kind), j.evaluation_id) + .await?, pname, score_breakdown: j.score_breakdown, worker_context: j.worker_context, @@ -665,6 +702,7 @@ async fn job_names( async fn eval_job_evaluation( db: &C, + mapper: &StorePathMapper, kind: i16, evaluation: EvaluationId, ) -> Result, sea_orm::DbErr> { @@ -672,7 +710,10 @@ async fn eval_job_evaluation( return Ok(None); } - job_evaluation(db, evaluation).await + Ok(job_evaluation(db, evaluation).await?.map(|mut view| { + view.repository = mapper.outbound_if_store_path(&view.repository); + view + })) } #[derive(Serialize, utoipa::ToSchema)] @@ -1331,7 +1372,7 @@ impl From for FlakeGraphNode { kind: n.kind, is_derivation: n.is_derivation, drv_path: n.drv_path.map(|p| { - gradient_entity::StorePath::parse(&p) + gradient_entity::StorePath::parse_stored(&p) .map(|sp| sp.base()) .unwrap_or(p) }), @@ -1355,11 +1396,19 @@ pub async fn get_eval_flake_graph( Extension(api_key): Extension, Path(evaluation_id): Path, ) -> WebResult>>> { - let _ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) + let ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) .await?; let rows = gradient_db::flake_output_nodes::of_evaluation(&state.web_db, evaluation_id).await?; - let out = rows.into_iter().map(FlakeGraphNode::from).collect(); + let mapper = state.project_mapper(ctx.project_id).await?; + let out = rows + .into_iter() + .map(FlakeGraphNode::from) + .map(|node| FlakeGraphNode { + drv_path: node.drv_path.map(|p| original_text(&mapper, p)), + ..node + }) + .collect(); Ok(ok_json(out)) } @@ -1616,6 +1665,30 @@ mod tests { ); } + #[test] + fn drv_paths_inside_job_context_are_mapped_outbound() { + let (original, stored) = ( + gradient_entity::EXAMPLE_ORIGINAL_HASH, + gradient_entity::EXAMPLE_STORED_HASH, + ); + let mapper = StorePathMapper::for_cache( + Arc::new(gradient_entity::RotatedHashMapping), + CacheId::nil(), + ); + let mut value = serde_json::json!({ + "derivations": [{ "drv_path": format!("{stored}-a.drv"), "name": "a" }] + }); + + outbound_drv_paths(&mut value, &mapper); + + assert_eq!( + value, + serde_json::json!({ + "derivations": [{ "drv_path": format!("{original}-a.drv"), "name": "a" }] + }) + ); + } + #[test] fn snapshot_derivations_of_an_eval_job_is_empty() { assert!(snapshot_derivations(&serde_json::json!({ "kind": "Eval" })).is_empty()); diff --git a/backend/gradient-web/src/endpoints/build_requests/dispatch.rs b/backend/gradient-web/src/endpoints/build_requests/dispatch.rs index 24ee5e106..29a8b0e74 100644 --- a/backend/gradient-web/src/endpoints/build_requests/dispatch.rs +++ b/backend/gradient-web/src/endpoints/build_requests/dispatch.rs @@ -31,7 +31,7 @@ use gradient_types::ids::{ use gradient_types::input::{check_flake_input_override, check_manifest_path, decode_blake3_hex}; use gradient_types::{ BaseResponse, ECache, EUploadSession, MCommit, MEvaluation, MEvaluationFlakeInputOverride, - MUploadSession, MUser, now, + MUploadSession, MUser, OriginalStorePath, StorePath, StorePathMapper, now, }; use sea_orm::{ConnectionTrait, TransactionTrait}; use serde::{Deserialize, Serialize}; @@ -171,15 +171,17 @@ pub(super) async fn finalize_build_request( check_flake_input_override(input_name, url)?; } + let mapper = db.project_mapper(project).await?; + let stored = stored_source_path(&mapper, &nar.store_path)?; db.storage .nar_storage - .put(&nar.store_hash, nar.compressed_bytes.clone()) + .put(stored.hash(), nar.compressed_bytes.clone()) .await .context("store source NAR")?; graph .commit_nar(NarCommit { - store_path: format!("{}-source", nar.store_hash), + store_path: stored.base(), file_hash: nar.file_hash_sri.clone(), file_size: nar.file_size as i64, nar_size: nar.nar_size as i64, @@ -198,7 +200,7 @@ pub(super) async fn finalize_build_request( project, user, BuildRequestSource { - repository: nar.store_path.clone(), + repository: stored.full(), commit_hash: vec![0; 20], commit_message: format!("Build request {}", nar.store_hash), target: options.target, @@ -220,6 +222,11 @@ pub(super) async fn finalize_build_request( Ok(response) } +fn stored_source_path(mapper: &StorePathMapper, store_path: &str) -> anyhow::Result { + let original = OriginalStorePath::parse(store_path).context("source NAR store path")?; + Ok(mapper.inbound(&original)) +} + pub(super) struct BuildRequestSource { pub repository: String, pub commit_hash: Vec, @@ -348,3 +355,26 @@ async fn materialise_staging( } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + use gradient_types::CacheId; + + #[test] + fn a_source_nar_is_stored_under_its_mapped_path() { + let mapper = StorePathMapper::for_cache(Arc::new(RotatedHashMapping), CacheId::nil()); + + let stored = stored_source_path( + &mapper, + &format!("/nix/store/{EXAMPLE_ORIGINAL_HASH}-source"), + ) + .expect("stored"); + + assert_eq!( + stored.full(), + format!("/nix/store/{EXAMPLE_STORED_HASH}-source") + ); + } +} diff --git a/backend/gradient-web/src/endpoints/builds/closure.rs b/backend/gradient-web/src/endpoints/builds/closure.rs index a2fa938a7..c327f9520 100644 --- a/backend/gradient-web/src/endpoints/builds/closure.rs +++ b/backend/gradient-web/src/endpoints/builds/closure.rs @@ -8,6 +8,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::endpoints::evals::EvalAccessContext; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; +use crate::store_paths::original_text; use axum::extract::{Path, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -76,6 +77,28 @@ pub struct ClosureGraph { pub edges: Vec, } +impl ClosureGraph { + fn with_original_paths(mut self, mapper: &StorePathMapper) -> Self { + for node in &mut self.nodes { + node.path = original_text(mapper, std::mem::take(&mut node.path)); + } + self + } + + fn with_original_hashes(mut self, mapper: &StorePathMapper) -> Self { + let original = |hash: &mut String| *hash = original_text(mapper, std::mem::take(hash)); + self.roots.iter_mut().for_each(original); + for node in &mut self.nodes { + original(&mut node.id); + } + for edge in &mut self.edges { + original(&mut edge.source); + original(&mut edge.target); + } + self.with_original_paths(mapper) + } +} + pub async fn derivation_closure_reachable( db: &C, seed_drv_ids: Vec, @@ -253,8 +276,9 @@ pub async fn get_build_closure( ) -> WebResult>> { let ctx = BuildAccessContext::load(&state.web_db, build_id, &maybe_user, api_key.as_ref()).await?; + let mapper = state.project_mapper(ctx.project.id).await?; let graph = build_closure_graph(&state.web_db, vec![ctx.build_job.derivation]).await?; - Ok(ok_json(graph)) + Ok(ok_json(graph.with_original_paths(&mapper))) } #[utoipa::path( @@ -273,7 +297,7 @@ pub async fn get_eval_closure( Extension(api_key): Extension, Path(evaluation_id): Path, ) -> WebResult>> { - let _ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) + let ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) .await?; let roots: Vec = @@ -283,8 +307,9 @@ pub async fn get_eval_closure( .map(|ep| ep.derivation) .collect(); + let mapper = state.project_mapper(ctx.project_id).await?; let graph = build_closure_graph(&state.web_db, roots).await?; - Ok(ok_json(graph)) + Ok(ok_json(graph.with_original_paths(&mapper))) } pub async fn build_runtime_closure_graph( @@ -351,8 +376,9 @@ pub async fn get_build_runtime_closure( ) .await?; + let mapper = state.project_mapper(ctx.project.id).await?; let graph = build_runtime_closure_graph(&state.web_db, seeds).await?; - Ok(ok_json(graph)) + Ok(ok_json(graph.with_original_hashes(&mapper))) } #[utoipa::path( @@ -371,7 +397,7 @@ pub async fn get_eval_runtime_closure( Extension(api_key): Extension, Path(evaluation_id): Path, ) -> WebResult>> { - let _ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) + let ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) .await?; let roots: Vec = @@ -383,8 +409,9 @@ pub async fn get_eval_runtime_closure( let seeds = gradient_db::graph::runtime_closure::output_hashes_for_drvs(&state.web_db, &roots).await?; + let mapper = state.project_mapper(ctx.project_id).await?; let graph = build_runtime_closure_graph(&state.web_db, seeds).await?; - Ok(ok_json(graph)) + Ok(ok_json(graph.with_original_hashes(&mapper))) } #[cfg(test)] diff --git a/backend/gradient-web/src/endpoints/builds/downloads.rs b/backend/gradient-web/src/endpoints/builds/downloads.rs index 506619eac..fe9691947 100644 --- a/backend/gradient-web/src/endpoints/builds/downloads.rs +++ b/backend/gradient-web/src/endpoints/builds/downloads.rs @@ -10,6 +10,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser, decode_download_token, encode use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; use crate::resource::Resource; +use crate::store_paths::original_text; use axum::extract::{Path, Query, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -204,7 +205,15 @@ pub async fn get_build_downloads( gradient_db::derivation_outputs::of_derivation(&state.web_db, ctx.build_job.derivation) .await?; - let products = collect_build_products(&state.web_db, build_id, build_outputs).await; + let mapper = state.project_mapper(ctx.project.id).await?; + let products = collect_build_products(&state.web_db, build_id, build_outputs) + .await + .into_iter() + .map(|product| BuildProduct { + path: original_text(&mapper, product.path), + ..product + }) + .collect::>(); Ok(ok_json(products)) } diff --git a/backend/gradient-web/src/endpoints/builds/graph.rs b/backend/gradient-web/src/endpoints/builds/graph.rs index 37c0e30d2..2ab870da7 100644 --- a/backend/gradient-web/src/endpoints/builds/graph.rs +++ b/backend/gradient-web/src/endpoints/builds/graph.rs @@ -7,6 +7,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; +use crate::store_paths::original_text; use axum::extract::{Path, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -150,9 +151,13 @@ pub async fn get_build_graph( ) -> WebResult>> { let ctx = BuildAccessContext::load(&state.web_db, build_id, &maybe_user, api_key.as_ref()).await?; - Ok(ok_json( - dependency_graph(&state.web_db, &ctx.build_job).await?, - )) + let mapper = state.project_mapper(ctx.project.id).await?; + let mut graph = dependency_graph(&state.web_db, &ctx.build_job).await?; + for node in &mut graph.nodes { + node.path = original_text(&mapper, std::mem::take(&mut node.path)); + } + + Ok(ok_json(graph)) } #[cfg(test)] diff --git a/backend/gradient-web/src/endpoints/builds/query.rs b/backend/gradient-web/src/endpoints/builds/query.rs index 94a3054b9..a5817118b 100644 --- a/backend/gradient-web/src/endpoints/builds/query.rs +++ b/backend/gradient-web/src/endpoints/builds/query.rs @@ -8,6 +8,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; use crate::resource::Resource; +use crate::store_paths::original_text; use axum::extract::{Path, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -86,10 +87,11 @@ pub async fn get_build( let derivation_outputs = gradient_db::derivation_outputs::of_derivation(&state.web_db, derivation.id).await?; + let mapper = state.project_mapper(ctx.project.id).await?; let mut outputs = HashMap::new(); for output in derivation_outputs { let path = StorePath::from(output.clone()).base(); - outputs.insert(output.name, path); + outputs.insert(output.name, original_text(&mapper, path)); } let attempt = latest_attempt(&state.web_db, shared_build.id) @@ -116,7 +118,7 @@ pub async fn get_build( id: build_job.id, evaluation: build_job.evaluation, status: shared_build.status.for_api(), - derivation_path: derivation.drv_path(), + derivation_path: original_text(&mapper, derivation.drv_path()), architecture: derivation.architecture, worker, dispatched_job, diff --git a/backend/gradient-web/src/endpoints/caches/build_log.rs b/backend/gradient-web/src/endpoints/caches/build_log.rs index 3b2d95b4d..cffe8c1f2 100644 --- a/backend/gradient-web/src/endpoints/caches/build_log.rs +++ b/backend/gradient-web/src/endpoints/caches/build_log.rs @@ -16,8 +16,8 @@ use gradient_core::ServerState; use gradient_core::upstream_source::{fetch_upstream_log, substitution_sources}; use gradient_db::WebDb; use gradient_db::caches::upstream::pull_through_upstream_caches; -use gradient_sources::parse_drv_hash_name; use gradient_storage::LogStorage; +use gradient_types::{OriginalStorePath, StorePathMapper}; use std::sync::Arc; /// A pull-through cache is substituting paths it never built (#547). The upstream caches are asked @@ -49,7 +49,8 @@ pub async fn log( ) .await?; - if let Some(body) = local_log(&state.web_db, &state.log_storage, &ctx, &drv).await? { + let mapper = state.cache_mapper(ctx.cache.id); + if let Some(body) = local_log(&state.web_db, &state.log_storage, &ctx, &mapper, &drv).await? { return log_response(body, "HIT"); } @@ -66,17 +67,26 @@ async fn local_log( db: &WebDb, log_storage: &Arc, ctx: &CacheContext, + mapper: &StorePathMapper, drv: &str, ) -> WebResult> { - let Ok((drv_hash, drv_name)) = parse_drv_hash_name(drv) else { + let Ok(original) = OriginalStorePath::parse(drv) else { + return Ok(None); + }; + let stored = mapper.inbound(&original); + let Some(drv_name) = stored + .name() + .strip_suffix(".drv") + .filter(|name| !name.is_empty()) + else { return Ok(None); }; let Some(derivation) = gradient_db::cached_paths::derivation_signed_in_cache( db, ctx.cache.id, - &drv_hash, - &drv_name, + stored.hash(), + drv_name, ) .await? else { diff --git a/backend/gradient-web/src/endpoints/caches/helpers.rs b/backend/gradient-web/src/endpoints/caches/helpers.rs index db3ab26d8..c4b3619db 100644 --- a/backend/gradient-web/src/endpoints/caches/helpers.rs +++ b/backend/gradient-web/src/endpoints/caches/helpers.rs @@ -12,6 +12,7 @@ use crate::helpers::OptionExt; use crate::ip_allowlist::is_allowed as ip_allowed; use crate::refusal::SharedRefusal; use crate::resource::Resource; +use crate::store_paths::original_text; use anyhow::Context; use axum::http::HeaderMap; use base64::Engine; @@ -123,9 +124,11 @@ pub(super) async fn get_nar_by_hash( db: &WebDb, config: &RuntimeConfig, cache: MCache, - hash: String, + mapper: &StorePathMapper, + hash: StoreHash, ) -> Result { - let cached_path_row = gradient_db::cached_paths::by_hash(db, &hash) + let hash = mapper.inbound_hash(hash); + let cached_path_row = gradient_db::cached_paths::by_hash(db, hash.as_str()) .await .map_err(WebError::from)? .filter(|row| row.is_fully_cached()) @@ -169,17 +172,22 @@ pub(super) async fn get_nar_by_hash( let file_hash_nix32 = strip_hash_algo(&file_hash).to_string(); Ok(NixPathInfo { - store_path: cached_path_row.store_path(), + store_path: original_text(mapper, cached_path_row.store_path()), url: format!("nar/{}.nar.zst", file_hash_nix32), compression: "zstd".to_string(), file_hash, file_size, nar_hash, nar_size, - references, - deriver: cached_path_row.deriver.clone(), + references: references + .into_iter() + .map(|reference| original_text(mapper, reference)) + .collect(), + deriver: cached_path_row + .deriver + .map(|deriver| original_text(mapper, deriver)), sig, - ca: cached_path_row.ca.clone(), + ca: cached_path_row.ca, }) } @@ -237,15 +245,13 @@ pub async fn fetch_nar_stream( db: &WebDb, nar_storage: &NarStore, cache: CacheId, - path_hash: &str, -) -> WebResult<(String, u64, BoxStream<'static, anyhow::Result>)> { - let effective_hash = - crate::endpoints::caches::nar::resolve_effective_hash_db(db, path_hash).await?; - if !gradient_db::cached_paths::signed_in_cache(db, cache, &effective_hash).await? { + stored_hash: &str, +) -> WebResult<(u64, BoxStream<'static, anyhow::Result>)> { + if !gradient_db::cached_paths::signed_in_cache(db, cache, stored_hash).await? { return Err(WebError::not_found(Resource::Path)); } let source = nar_storage - .open(&effective_hash, 0) + .open(stored_hash, 0) .await .context("read NAR")? .or_not_found(Resource::Path)?; @@ -259,7 +265,7 @@ pub async fn fetch_nar_stream( } NarSource::Stream { size, stream } => (size, stream), }; - Ok((effective_hash, size, stream)) + Ok((size, stream)) } #[derive(Debug, Clone, Copy)] diff --git a/backend/gradient-web/src/endpoints/caches/nar.rs b/backend/gradient-web/src/endpoints/caches/nar.rs index 9979a886f..33ff36bde 100644 --- a/backend/gradient-web/src/endpoints/caches/nar.rs +++ b/backend/gradient-web/src/endpoints/caches/nar.rs @@ -10,6 +10,7 @@ use crate::api_spec::Binary; use crate::client_ip::OptionalPeer; use crate::error::{WebError, WebResult}; use crate::resource::Resource; +use crate::store_paths::{original_text, stored_hash}; use anyhow::Context; use axum::body::Body; use axum::extract::{Path, RawQuery, State}; @@ -63,18 +64,20 @@ pub async fn nar( ) .await?; - let (effective_hash, size, stream) = super::helpers::fetch_nar_stream( + let mapper = state.cache_mapper(ctx.cache.id); + let effective_hash = resolve_effective_hash_db(&state.web_db, &mapper, &path_hash).await?; + let (size, stream) = super::helpers::fetch_nar_stream( &state.web_db, &state.nar_storage, ctx.cache.id, - &path_hash, + &effective_hash, ) .await?; super::super::stats::record_nar_traffic(&state.memory.cache_traffic, ctx.cache.id, size as i64); state.events.publish(NarFetched { cache: ctx.cache.id, - hash: effective_hash.clone(), + hash: original_text(&mapper, effective_hash.clone()), size, }); spawn_fetch_stamp( @@ -167,6 +170,7 @@ fn named_first(sources: Vec, named: Uuid) -> Vec pub(crate) async fn resolve_effective_hash_db( db: &C, + mapper: &StorePathMapper, path_hash: &str, ) -> WebResult { let candidates = vec![format!("blake3:{path_hash}"), format!("sha256:{path_hash}")]; @@ -177,7 +181,7 @@ pub(crate) async fn resolve_effective_hash_db( return Ok(row.hash); } - Ok(path_hash.to_string()) + stored_hash(mapper, path_hash) } fn spawn_fetch_stamp(shutdown: &Shutdown, worker_db: &WorkerDb, cache_id: CacheId, hash: String) { @@ -300,6 +304,10 @@ mod tests { } } + fn identity() -> StorePathMapper { + StorePathMapper::for_cache(Arc::new(IdentityMapping), CacheId::nil()) + } + fn runtime() -> tokio::runtime::Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() @@ -314,21 +322,40 @@ mod tests { .into_connection(); let effective = runtime() - .block_on(resolve_effective_hash_db(&db, FILE_HASH_NIX32)) + .block_on(resolve_effective_hash_db(&db, &identity(), FILE_HASH_NIX32)) .expect("resolve should succeed"); assert_eq!(effective, STORE_HASH); } #[test] - fn resolve_falls_back_to_url_hash_when_no_match() { + fn an_unknown_file_hash_is_not_found() { + let db = mock_db() + .append_query_results([Vec::::new()]) + .into_connection(); + + let result = + runtime().block_on(resolve_effective_hash_db(&db, &identity(), FILE_HASH_NIX32)); + assert!(result.is_err()); + } + + #[test] + fn a_store_hash_in_a_nar_url_is_mapped_inbound() { let db = mock_db() .append_query_results([Vec::::new()]) .into_connection(); + let mapper = StorePathMapper::for_cache( + Arc::new(gradient_entity::RotatedHashMapping), + CacheId::nil(), + ); let effective = runtime() - .block_on(resolve_effective_hash_db(&db, FILE_HASH_NIX32)) + .block_on(resolve_effective_hash_db( + &db, + &mapper, + gradient_entity::EXAMPLE_ORIGINAL_HASH, + )) .expect("resolve should succeed"); - assert_eq!(effective, FILE_HASH_NIX32); + assert_eq!(effective, gradient_entity::EXAMPLE_STORED_HASH); } #[test] @@ -351,7 +378,7 @@ mod tests { .into_connection(); let effective = runtime() - .block_on(resolve_effective_hash_db(&db, FILE_HASH_NIX32)) + .block_on(resolve_effective_hash_db(&db, &identity(), FILE_HASH_NIX32)) .expect("resolve should succeed"); assert_eq!(effective, STORE_HASH); } diff --git a/backend/gradient-web/src/endpoints/caches/narinfo.rs b/backend/gradient-web/src/endpoints/caches/narinfo.rs index 21cc715fa..35db5ab4c 100644 --- a/backend/gradient-web/src/endpoints/caches/narinfo.rs +++ b/backend/gradient-web/src/endpoints/caches/narinfo.rs @@ -187,6 +187,8 @@ pub async fn path( let path_hash = get_hash_from_url(path.clone()).map_err(|_| WebError::not_found(Resource::Path))?; + let store_hash = + StoreHash::parse(&path_hash).map_err(|_| WebError::not_found(Resource::Path))?; let client_ip = cache_client_ip(&state.config, &headers, peer); let ctx = CacheContext::load( @@ -199,11 +201,13 @@ pub async fn path( ) .await?; + let mapper = state.cache_mapper(ctx.cache.id); if let Ok(path_info) = get_nar_by_hash( &state.web_db, &state.config, ctx.cache.clone(), - path_hash.clone(), + &mapper, + store_hash, ) .await { diff --git a/backend/gradient-web/src/endpoints/caches/narlist.rs b/backend/gradient-web/src/endpoints/caches/narlist.rs index 1fc821b0b..33fc961cd 100644 --- a/backend/gradient-web/src/endpoints/caches/narlist.rs +++ b/backend/gradient-web/src/endpoints/caches/narlist.rs @@ -6,6 +6,7 @@ use super::helpers::{CacheContext, cache_client_ip, fetch_nar_stream}; use crate::client_ip::OptionalPeer; use crate::error::{WebError, WebResult}; +use crate::store_paths::stored_hash; use anyhow::Context; use axum::Json; use axum::extract::{Path, State}; @@ -73,8 +74,10 @@ pub async fn ls( cache, ) .await?; - let (_effective_hash, _size, stream) = - fetch_nar_stream(&state.web_db, &state.nar_storage, ctx.cache.id, &hash).await?; + let mapper = state.cache_mapper(ctx.cache.id); + let stored = stored_hash(&mapper, &hash)?; + let (_size, stream) = + fetch_nar_stream(&state.web_db, &state.nar_storage, ctx.cache.id, &stored).await?; let reader = nar_reader_from_stream(stream); let root = walk_nar(reader).await.context("walk the NAR")?; diff --git a/backend/gradient-web/src/endpoints/caches/nars.rs b/backend/gradient-web/src/endpoints/caches/nars.rs index 3f79b49fa..4a4e240ec 100644 --- a/backend/gradient-web/src/endpoints/caches/nars.rs +++ b/backend/gradient-web/src/endpoints/caches/nars.rs @@ -11,6 +11,7 @@ use crate::error::{WebError, WebResult}; use crate::helpers::{OptionExt, ok_json}; use crate::resource::Resource; use crate::scope::{CacheScope, Readable, Requires, can}; +use crate::store_paths::{original_text, stored_hash}; use anyhow::Context; use axum::Json; use axum::extract::{Query, State}; @@ -57,17 +58,25 @@ pub struct NarDetail { } impl NarDetail { - fn new(path: MCachedPath, signature: MCachedPathSignature, references: Vec) -> Self { + fn new( + path: MCachedPath, + signature: MCachedPathSignature, + references: Vec, + mapper: &StorePathMapper, + ) -> Self { Self { - store_path: path.as_store_path().base(), - hash: path.hash, + store_path: original_text(mapper, path.as_store_path().base()), + hash: original_text(mapper, path.hash), package: path.package, nar_size: path.nar_size, file_size: path.file_size, file_hash: path.file_hash, nar_hash: path.nar_hash, - references, - deriver: path.deriver, + references: references + .into_iter() + .map(|reference| original_text(mapper, reference)) + .collect(), + deriver: path.deriver.map(|deriver| original_text(mapper, deriver)), ca: path.ca, created_at: path.created_at, last_fetched_at: signature.last_fetched_at, @@ -136,9 +145,24 @@ pub async fn list( "last_fetched_at" => NarOrder::LastFetchedAt, _ => NarOrder::CreatedAt, }; + let mapper = db.cache_mapper(cache.id); + let hash_prefix = match q.hash.as_deref().filter(|s| !s.is_empty()) { + Some(prefix) => match mapper.inbound_prefix(prefix) { + Some(stored) => Some(stored), + None => { + return Ok(ok_json(Paginated { + items: Vec::new(), + total: 0, + page, + per_page, + })); + } + }, + None => None, + }; let filter = NarFilter { cache: cache.id, - hash_prefix: q.hash.as_deref().filter(|s| !s.is_empty()), + hash_prefix: hash_prefix.as_deref(), package: q.package.as_deref().filter(|s| !s.is_empty()), }; @@ -146,7 +170,14 @@ pub async fn list( let items = gradient_db::cached_paths::page_in_cache( &db.web_db, filter, order, ascending, per_page, offset, ) - .await?; + .await? + .into_iter() + .map(|item| NarSummary { + hash: original_text(&mapper, item.hash), + store_path: original_text(&mapper, item.store_path), + ..item + }) + .collect(); Ok(ok_json(Paginated { items, @@ -172,6 +203,8 @@ pub async fn show( }: CacheScope, ) -> WebResult>> { let hash: String = params.param("hash")?; + let mapper = db.cache_mapper(cache.id); + let hash = stored_hash(&mapper, &hash)?; let cp = gradient_db::cached_paths::by_hash(&db.web_db, &hash) .await? @@ -181,7 +214,7 @@ pub async fn show( .or_not_found(Resource::Signature)?; let references = gradient_db::graph::runtime_closure::references_for_hash(&db.web_db, &hash).await?; - Ok(ok_json(NarDetail::new(cp, sig, references))) + Ok(ok_json(NarDetail::new(cp, sig, references, &mapper))) } #[utoipa::path( @@ -222,6 +255,7 @@ pub async fn available( if hash.is_empty() { return Err(CacheRefusal::MissingHashQuery.into()); } + let hash = stored_hash(&db.cache_mapper(cache.id), &hash)?; let cp = gradient_db::cached_paths::by_hash(&db.web_db, &hash).await?; let exists = if let Some(cp) = cp { gradient_db::cached_paths::signature_in_cache(&db.web_db, cp.id, cache.id) @@ -244,15 +278,19 @@ pub async fn available( responses((status = NO_CONTENT, description = "NAR removed from this cache"), WebError) )] pub async fn delete( - CacheScope { cache, params, .. }: CacheScope>, + CacheScope { + cache, params, db, .. + }: CacheScope>, state: State>, audit: Audit, ) -> WebResult { let hash: String = params.param("hash")?; + let mapper = db.cache_mapper(cache.id); + let stored = stored_hash(&mapper, &hash)?; let (cp, outcome) = - delete_nar_from_cache(&state.nar_storage, &state.graph, cache.id, &hash).await?; - audit.write(Action::CacheNarDelete, EventOwner::cache(cache.id), serde_json::json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "hash": hash, "store_path": cp.store_path(), "ref_counted_others": outcome.ref_counted_others, })) + delete_nar_from_cache(&state.nar_storage, &state.graph, cache.id, &stored).await?; + audit.write(Action::CacheNarDelete, EventOwner::cache(cache.id), serde_json::json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "hash": hash, "store_path": original_text(&mapper, cp.store_path()), "ref_counted_others": outcome.ref_counted_others, })) .await; Ok(StatusCode::NO_CONTENT) } diff --git a/backend/gradient-web/src/endpoints/caches/serve.rs b/backend/gradient-web/src/endpoints/caches/serve.rs index 82685866f..3e6ebd867 100644 --- a/backend/gradient-web/src/endpoints/caches/serve.rs +++ b/backend/gradient-web/src/endpoints/caches/serve.rs @@ -8,6 +8,7 @@ use crate::api_spec::Binary; use crate::client_ip::OptionalPeer; use crate::error::{WebError, WebResult}; use crate::resource::Resource; +use crate::store_paths::stored_hash; use anyhow::Context; use axum::body::Body; use axum::extract::{Path, State}; @@ -51,8 +52,10 @@ pub async fn serve( cache, ) .await?; - let (_effective_hash, _size, stream) = - fetch_nar_stream(&state.web_db, &state.nar_storage, ctx.cache.id, &hash).await?; + let mapper = state.cache_mapper(ctx.cache.id); + let stored = stored_hash(&mapper, &hash)?; + let (_size, stream) = + fetch_nar_stream(&state.web_db, &state.nar_storage, ctx.cache.id, &stored).await?; let reader = nar_reader_from_stream(stream); match extract_path_from_reader(reader, &rel_path).await { diff --git a/backend/gradient-web/src/endpoints/caches/upload.rs b/backend/gradient-web/src/endpoints/caches/upload.rs index e5487d791..811133547 100644 --- a/backend/gradient-web/src/endpoints/caches/upload.rs +++ b/backend/gradient-web/src/endpoints/caches/upload.rs @@ -17,6 +17,7 @@ use axum::extract::{Multipart, Query, State}; use axum::http::{HeaderMap, StatusCode, header}; use axum::response::IntoResponse; use gradient_core::ServerState; +use gradient_entity::StorePathError; use gradient_proto::import::{ImportInput, SignTargets, import_nar_reader}; use gradient_storage::PartialStore; use gradient_types::events::EventOwner; @@ -120,6 +121,14 @@ pub async fn nars_upload( return Err(CacheRefusal::MissingPart("narinfo").into()); } }; + let upload = match MappedUpload::new(narinfo, &state.cache_mapper(cache.id)) { + Ok(upload) => upload, + Err(e) => { + let _ = upload_store.discard(&stage_key).await; + return Err(e); + } + }; + let narinfo = &upload.narinfo; let staged_size = staged_size.ok_or(CacheRefusal::MissingPart("nar"))?; if staged_size as i64 != narinfo.file_size { @@ -153,7 +162,7 @@ pub async fn nars_upload( let claim = admit_nar( &state.config, &state.memory.upload_admission, - &narinfo_hash(&narinfo.store_path), + &upload.stored_hash, narinfo.file_size as u64, ) .await?; @@ -166,16 +175,7 @@ pub async fn nars_upload( &state.nar_storage, &state.graph, nar_reader, - ImportInput { - store_path: &narinfo.store_path, - file_hash: &narinfo.file_hash, - file_size: narinfo.file_size, - nar_size: narinfo.nar_size, - nar_hash: &narinfo.nar_hash, - references: &narinfo.references, - deriver: narinfo.deriver.as_deref(), - ca: None, - }, + upload.import_input(), SignTargets::Cache(cache.id), ) .await @@ -189,30 +189,76 @@ pub async fn nars_upload( for cache in outcome.signed { state.events.publish(NarSigned { cache, - hash: narinfo_hash(&narinfo.store_path), + hash: upload.original_hash.clone(), }); } - audit.write(Action::CacheNarUpload, EventOwner::cache(cache.id), json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "store_path": narinfo.store_path, "created": outcome.created, })) + audit.write(Action::CacheNarUpload, EventOwner::cache(cache.id), json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "store_path": upload.client_path, "created": outcome.created, })) .await; Ok(( StatusCode::CREATED, ok_json(StoredNar { - store_path: narinfo.store_path, + store_path: upload.client_path, created: outcome.created, }), )) } -fn narinfo_hash(store_path: &str) -> String { - store_path - .strip_prefix("/nix/store/") - .unwrap_or(store_path) - .split('-') - .next() - .unwrap_or("") - .to_owned() +fn stored_narinfo( + part: NarinfoPart, + mapper: &StorePathMapper, +) -> Result { + Ok(NarinfoPart { + store_path: mapper.inbound_text(&part.store_path)?, + references: part + .references + .iter() + .map(|reference| mapper.inbound_text(reference)) + .collect::>()?, + deriver: part + .deriver + .as_deref() + .map(|deriver| mapper.inbound_text(deriver)) + .transpose()?, + ..part + }) +} + +struct MappedUpload { + client_path: String, + original_hash: String, + stored_hash: String, + narinfo: NarinfoPart, +} + +impl MappedUpload { + fn new(narinfo: NarinfoPart, mapper: &StorePathMapper) -> WebResult { + let original = OriginalStorePath::parse(&narinfo.store_path) + .map_err(|_| CacheRefusal::InvalidStoreHash)?; + let client_path = narinfo.store_path.clone(); + let narinfo = + stored_narinfo(narinfo, mapper).map_err(|_| CacheRefusal::InvalidStoreHash)?; + Ok(Self { + client_path, + original_hash: original.hash().to_string(), + stored_hash: mapper.inbound_hash(original.hash().clone()).to_string(), + narinfo, + }) + } + + fn import_input(&self) -> ImportInput<'_> { + ImportInput { + store_path: &self.narinfo.store_path, + file_hash: &self.narinfo.file_hash, + file_size: self.narinfo.file_size, + nar_size: self.narinfo.nar_size, + nar_hash: &self.narinfo.nar_hash, + references: &self.narinfo.references, + deriver: self.narinfo.deriver.as_deref(), + ca: None, + } + } } #[derive(Deserialize, utoipa::IntoParams)] @@ -251,6 +297,14 @@ async fn stage_nar_field( Ok(offset) } +fn stored_segment(segment: &str, mapper: &StorePathMapper) -> WebResult { + let stored = mapper + .inbound_text(segment) + .map_err(|_| CacheRefusal::InvalidStoreHash)?; + require_safe_hash(&stored)?; + Ok(stored) +} + fn require_safe_hash(store_hash: &str) -> WebResult<()> { if store_hash.is_empty() || store_hash.contains('/') || store_hash.contains("..") { return Err(CacheRefusal::InvalidStoreHash.into()); @@ -276,9 +330,10 @@ pub async fn nar_chunk( Query(NarChunkQuery { offset }): Query, body: Bytes, ) -> WebResult { - let store_hash: String = params.param("hash")?; - - require_safe_hash(&store_hash)?; + let store_hash = stored_segment( + ¶ms.param::("hash")?, + &state.cache_mapper(cache.id), + )?; let max = u64::try_from(state.config.nar.max_upload_size.get()).unwrap_or(u64::MAX); if offset + body.len() as u64 > max { @@ -327,9 +382,10 @@ pub async fn nar_finalize( audit: Audit, axum::Json(narinfo): axum::Json, ) -> WebResult { - let store_hash: String = params.param("hash")?; - - require_safe_hash(&store_hash)?; + let mapper = state.cache_mapper(cache.id); + let store_hash = stored_segment(¶ms.param::("hash")?, &mapper)?; + let upload = MappedUpload::new(narinfo, &mapper)?; + let narinfo = &upload.narinfo; let store = upload_partial_store(&state.config)?; let key = format!("{}/{store_hash}", cache.id); @@ -345,7 +401,7 @@ pub async fn nar_finalize( let claim = admit_nar( &state.config, &state.memory.upload_admission, - &narinfo_hash(&narinfo.store_path), + &upload.stored_hash, narinfo.file_size as u64, ) .await?; @@ -368,16 +424,7 @@ pub async fn nar_finalize( &state.nar_storage, &state.graph, nar_reader, - ImportInput { - store_path: &narinfo.store_path, - file_hash: &narinfo.file_hash, - file_size: narinfo.file_size, - nar_size: narinfo.nar_size, - nar_hash: &narinfo.nar_hash, - references: &narinfo.references, - deriver: narinfo.deriver.as_deref(), - ca: None, - }, + upload.import_input(), SignTargets::Cache(cache.id), ) .await @@ -391,17 +438,17 @@ pub async fn nar_finalize( for cache in outcome.signed { state.events.publish(NarSigned { cache, - hash: narinfo_hash(&narinfo.store_path), + hash: upload.original_hash.clone(), }); } - audit.write(Action::CacheNarUpload, EventOwner::cache(cache.id), json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "store_path": narinfo.store_path, "created": outcome.created, "chunked": true, })) + audit.write(Action::CacheNarUpload, EventOwner::cache(cache.id), json!({ "cache_id": cache.id.to_string(), "cache_name": cache.name, "store_path": upload.client_path, "created": outcome.created, "chunked": true, })) .await; Ok(( StatusCode::CREATED, ok_json(StoredNar { - store_path: narinfo.store_path, + store_path: upload.client_path, created: outcome.created, }), )) @@ -409,7 +456,58 @@ pub async fn nar_finalize( #[cfg(test)] mod tests { - use super::require_safe_hash; + use super::{NarinfoPart, require_safe_hash, stored_narinfo}; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; + use gradient_types::{CacheId, StorePathMapper}; + use std::sync::Arc; + + fn part(store_path: String, references: Vec) -> NarinfoPart { + NarinfoPart { + store_path, + file_hash: "sha256:0mdqa9w1p6cmli6976v4wi0sw9r4p5prkj7lzfd1877wk11c9c73".into(), + file_size: 5, + nar_size: 5, + nar_hash: "sha256:0mdqa9w1p6cmli6976v4wi0sw9r4p5prkj7lzfd1877wk11c9c73".into(), + references, + deriver: None, + } + } + + fn rotated() -> StorePathMapper { + StorePathMapper::for_cache(Arc::new(RotatedHashMapping), CacheId::nil()) + } + + #[test] + fn an_uploaded_narinfo_is_stored_under_its_mapped_hash() { + let (o, s) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let uploaded = NarinfoPart { + deriver: Some(format!("{o}-hello.drv")), + ..part( + format!("/nix/store/{o}-hello"), + vec![format!("/nix/store/{o}-hello")], + ) + }; + + let stored = stored_narinfo(uploaded, &rotated()).unwrap(); + + assert_eq!(stored.store_path, format!("/nix/store/{s}-hello")); + assert_eq!(stored.references, vec![format!("/nix/store/{s}-hello")]); + assert_eq!(stored.deriver, Some(format!("{s}-hello.drv"))); + assert_eq!(stored.file_size, 5); + } + + #[test] + fn an_uploaded_narinfo_with_a_short_hash_is_refused() { + let o = EXAMPLE_ORIGINAL_HASH; + assert!(stored_narinfo(part("abc-hello".into(), Vec::new()), &rotated()).is_err()); + assert!( + stored_narinfo( + part(format!("{o}-hello"), vec!["abc-bash".into()]), + &rotated() + ) + .is_err() + ); + } #[test] fn accepts_a_normal_store_basename() { diff --git a/backend/gradient-web/src/endpoints/evals/artefacts.rs b/backend/gradient-web/src/endpoints/evals/artefacts.rs index 941ef4829..91b5b8330 100644 --- a/backend/gradient-web/src/endpoints/evals/artefacts.rs +++ b/backend/gradient-web/src/endpoints/evals/artefacts.rs @@ -7,6 +7,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; +use crate::store_paths::original_text; use axum::extract::{Path, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -52,9 +53,13 @@ pub struct OutputArtefacts { } impl OutputArtefacts { - fn new(output: MDerivationOutput, products: Vec) -> Self { + fn new( + mapper: &StorePathMapper, + output: MDerivationOutput, + products: Vec, + ) -> Self { Self { - store_path: StorePath::from(output.clone()).base(), + store_path: original_text(mapper, StorePath::from(output.clone()).base()), name: output.name, products, } @@ -104,6 +109,7 @@ pub async fn get_artefacts( let ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) .await?; let evaluation = ctx.evaluation; + let mapper = state.project_mapper(ctx.project_id).await?; let entry_points = gradient_db::entry_points::of_evaluation(&state.web_db, evaluation.id).await?; @@ -178,14 +184,17 @@ pub async fn get_artefacts( .cloned() .unwrap_or_default() .into_iter() - .map(ProductArtefact::from) + .map(|p| ProductArtefact { + path: original_text(&mapper, p.path.clone()), + ..ProductArtefact::from(p) + }) .collect(); - OutputArtefacts::new(o, products) + OutputArtefacts::new(&mapper, o, products) }) .collect(); Some(EntryPointArtefacts { attr: ep.eval, - derivation: drv.drv_path(), + derivation: original_text(&mapper, drv.drv_path()), build_id, outputs, }) diff --git a/backend/gradient-web/src/endpoints/evals/query.rs b/backend/gradient-web/src/endpoints/evals/query.rs index 5f7c95fd9..aa1ab36d8 100644 --- a/backend/gradient-web/src/endpoints/evals/query.rs +++ b/backend/gradient-web/src/endpoints/evals/query.rs @@ -8,6 +8,7 @@ use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; use crate::resource::Resource; +use crate::store_paths::original_text; use axum::extract::{Path, Query, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -150,7 +151,10 @@ pub async fn get_evaluation( task: evaluation.task, task_name: ctx.task_name, task_display_name: ctx.task_display_name, - repository: evaluation.repository, + repository: state + .project_mapper(ctx.project_id) + .await? + .outbound_if_store_path(&evaluation.repository), commit: commit_hash, wildcard: evaluation.wildcard, status: evaluation.status, @@ -213,6 +217,7 @@ pub async fn get_evaluation_builds( let ctx = EvalAccessContext::load(&state.web_db, evaluation_id, &maybe_user, api_key.as_ref()) .await?; let evaluation = ctx.evaluation; + let mapper = state.project_mapper(ctx.project_id).await?; let jobs = gradient_db::build_jobs::of_evaluation(&state.web_db, evaluation.id).await?; @@ -335,7 +340,7 @@ pub async fn get_evaluation_builds( page.push(BuildItem { id: j.id, - name: drv.drv_path(), + name: original_text(&mapper, drv.drv_path()), status: format!("{:?}", status), has_artefacts: has_artefacts.contains(&j.derivation), updated_at: shared_build.updated_at, diff --git a/backend/gradient-web/src/endpoints/evals/summary/mod.rs b/backend/gradient-web/src/endpoints/evals/summary/mod.rs index 98006e321..485603efb 100644 --- a/backend/gradient-web/src/endpoints/evals/summary/mod.rs +++ b/backend/gradient-web/src/endpoints/evals/summary/mod.rs @@ -9,6 +9,7 @@ mod analysis; use crate::authorization::{MaybeApiKey, MaybeUser}; use crate::error::{WebError, WebResult}; use crate::helpers::ok_json; +use crate::store_paths::original_text; use axum::extract::{Path, State}; use axum::{Extension, Json}; use gradient_core::ServerState; @@ -115,7 +116,13 @@ pub async fn get_evaluation_summary( Ok(ok_json(EvaluationFailureSummary { packages: analysis::package_counts(&loaded.graph), failures_total: failures.len(), - failures: failed_builds(db, &loaded, failures).await?, + failures: failed_builds( + db, + &state.project_mapper(ctx.project_id).await?, + &loaded, + failures, + ) + .await?, failed_attributes: failed_attribute_reports(loaded.failed_attributes, baseline), fixed_total: fixed.len(), fixed: fixed.into_iter().take(FIXED_SHOWN).collect(), @@ -194,6 +201,7 @@ async fn compared_evaluation( async fn failed_builds( db: &C, + mapper: &StorePathMapper, loaded: &LoadedEvaluation, failures: Vec, ) -> Result, DbErr> { @@ -210,7 +218,7 @@ async fn failed_builds( Some(FailedBuildSummary { build_id: *loaded.build_jobs.get(&failure.derivation)?, name: derivation.name.clone(), - derivation_path: derivation.drv_path(), + derivation_path: original_text(mapper, derivation.drv_path()), architecture: derivation.architecture.clone(), status: failure.status, attributes: failure.attributes, diff --git a/backend/gradient-web/src/endpoints/search/classify.rs b/backend/gradient-web/src/endpoints/search/classify.rs index ba44e2b00..1c7d7f39c 100644 --- a/backend/gradient-web/src/endpoints/search/classify.rs +++ b/backend/gradient-web/src/endpoints/search/classify.rs @@ -4,7 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -const NIX32: &str = "0123456789abcdfghijklmnpqrsvwxyz"; +use gradient_types::{OriginalStorePath, StoreHash}; #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct CommitRange { @@ -14,17 +14,15 @@ pub struct CommitRange { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Lookups { - pub nar: Option, + pub nar: Option, pub commit: Option, pub text: Option, } -fn nar_hash(q: &str) -> Option { - let base = q.strip_prefix("/nix/store/").unwrap_or(q); - let hash = base.get(..32)?; - let rest = &base[32..]; - (hash.chars().all(|c| NIX32.contains(c)) && (rest.is_empty() || rest.starts_with('-'))) - .then(|| hash.to_string()) +fn nar_hash(q: &str) -> Option { + StoreHash::parse(q) + .or_else(|_| OriginalStorePath::parse(q).map(|path| path.hash().clone())) + .ok() } fn commit_range(q: &str) -> Option { @@ -58,14 +56,17 @@ mod tests { #[test] fn a_store_path_is_a_nar_and_a_name() { let l = classify(&format!("/nix/store/{HASH}-hello-2.12.1")); - assert_eq!(l.nar.as_deref(), Some(HASH)); + assert_eq!(l.nar.as_ref().map(StoreHash::as_str), Some(HASH)); assert!(l.commit.is_none()); assert!(l.text.is_some()); } #[test] fn a_bare_hash_is_a_nar() { - assert_eq!(classify(HASH).nar.as_deref(), Some(HASH)); + assert_eq!( + classify(HASH).nar.as_ref().map(StoreHash::as_str), + Some(HASH) + ); } #[test] diff --git a/backend/gradient-web/src/endpoints/search/mod.rs b/backend/gradient-web/src/endpoints/search/mod.rs index 8e01615eb..1553a7a19 100644 --- a/backend/gradient-web/src/endpoints/search/mod.rs +++ b/backend/gradient-web/src/endpoints/search/mod.rs @@ -13,8 +13,8 @@ use axum::{Extension, Json}; use classify::{Lookups, classify}; use gradient_core::ServerState; use gradient_db::dashboard::{ - CommitHitRow, NameHitRow, NameKind, NarHitRow, search_commits, search_names, search_nars, - starred, + CommitHitRow, NameHitRow, NameKind, NarHitRow, readable_cache_ids, search_commits, + search_names, search_nars, starred, }; use gradient_types::*; use sea_orm::{ConnectionTrait, DbErr}; @@ -75,18 +75,30 @@ impl From for SearchHit { } } -impl From for SearchHit { - fn from(n: NarHitRow) -> Self { - Self { - kind: HitKind::Nar, - route: format!("/caches/{}/nars?hash={}", n.cache, n.hash), - label: n.package, - sublabel: n.cache, - starred: false, - } +fn nar_hit(n: NarHitRow, typed: &StoreHash) -> SearchHit { + SearchHit { + kind: HitKind::Nar, + route: format!("/caches/{}/nars?hash={typed}", n.cache), + label: n.package, + sublabel: n.cache, + starred: false, } } +fn nar_lookups( + hash: &StoreHash, + caches: &[CacheId], + mapping: &Arc, +) -> Vec<(CacheId, StoreHash)> { + caches + .iter() + .map(|&cache| { + let stored = mapping.inbound(MappingContext::Cache(cache), hash.clone()); + (cache, stored) + }) + .collect() +} + fn per_kind(limit: Option) -> u64 { limit.unwrap_or(5).clamp(1, 20) } @@ -123,6 +135,7 @@ fn cache_hit(name: String, starred: bool) -> SearchHit { pub async fn search( db: &C, + mapping: &Arc, user: UserId, lookups: Lookups, per_kind: u64, @@ -136,8 +149,9 @@ pub async fn search( }; let mut hits = Vec::new(); if let Some(hash) = &lookups.nar { - let nars = search_nars(db, hash, user).await?; - hits.extend(nars.into_iter().map(SearchHit::from)); + let caches = readable_cache_ids(db, user).await?; + let nars = search_nars(db, &nar_lookups(hash, &caches, mapping), user).await?; + hits.extend(nars.into_iter().map(|n| nar_hit(n, hash))); } if let Some(range) = &lookups.commit { let commits = search_commits(db, &range.low, &range.high, user).await?; @@ -163,13 +177,21 @@ pub async fn get_search( Extension(user): Extension, Query(q): Query, ) -> WebResult>>> { - let hits = search(&state.web_db, user.id, classify(&q.q), per_kind(q.limit)).await?; + let hits = search( + &state.web_db, + &state.extensions.store_path_mapping, + user.id, + classify(&q.q), + per_kind(q.limit), + ) + .await?; Ok(ok_json(hits)) } #[cfg(test)] mod tests { use super::*; + use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; use gradient_fixtures::mock::mock_db; use sea_orm::Value; use std::collections::BTreeMap; @@ -194,10 +216,25 @@ mod tests { async fn run(results: Vec>, q: &str) -> (Vec, usize) { let db = mock_db().append_query_results(results).into_connection(); - let hits = search(&db, UserId::now_v7(), classify(q), 5).await.unwrap(); + let mapping: Arc = Arc::new(IdentityMapping); + let hits = search(&db, &mapping, UserId::now_v7(), classify(q), 5) + .await + .unwrap(); (hits, db.into_transaction_log().len()) } + #[test] + fn a_nar_query_is_mapped_once_per_readable_cache() { + let (c1, c2) = (CacheId::now_v7(), CacheId::now_v7()); + let hash = StoreHash::parse(EXAMPLE_ORIGINAL_HASH).unwrap(); + let stored = StoreHash::parse(EXAMPLE_STORED_HASH).unwrap(); + let mapping: Arc = Arc::new(RotatedHashMapping); + assert_eq!( + nar_lookups(&hash, &[c1, c2], &mapping), + vec![(c1, stored.clone()), (c2, stored)] + ); + } + #[test] fn limit_is_clamped() { assert_eq!(per_kind(None), 5); @@ -236,7 +273,8 @@ mod tests { let db = mock_db() .append_query_results([vec![name_row("task", None, "hosts", false)]]) .into_connection(); - let result = search(&db, UserId::now_v7(), classify("hosts"), 5).await; + let mapping: Arc = Arc::new(IdentityMapping); + let result = search(&db, &mapping, UserId::now_v7(), classify("hosts"), 5).await; assert!(result.is_err()); } @@ -268,18 +306,21 @@ mod tests { #[tokio::test] async fn a_nar_hash_routes_to_the_cache_nar_listing() { + let cache = CacheId::now_v7(); + let readable = BTreeMap::from([("id", Value::from(cache.into_inner()))]); let nar = BTreeMap::from([ + ("cache_id", Value::from(cache.into_inner())), ("cache", Value::from("main")), ("hash", Value::from(HASH)), ("package", Value::from("hello-2.12.1")), ]); - let (hits, statements) = run(vec![vec![nar], vec![]], HASH).await; + let (hits, statements) = run(vec![vec![readable], vec![nar], vec![]], HASH).await; assert_eq!(hits.len(), 1); assert_eq!(hits[0].kind, HitKind::Nar); assert_eq!(hits[0].route, format!("/caches/main/nars?hash={HASH}")); assert_eq!(hits[0].label, "hello-2.12.1"); - assert_eq!(statements, 2); + assert_eq!(statements, 3); } #[tokio::test] diff --git a/backend/gradient-web/src/endpoints/tasks/evaluations.rs b/backend/gradient-web/src/endpoints/tasks/evaluations.rs index bfeb1e6e1..e6aaf0082 100644 --- a/backend/gradient-web/src/endpoints/tasks/evaluations.rs +++ b/backend/gradient-web/src/endpoints/tasks/evaluations.rs @@ -21,6 +21,7 @@ use crate::permissions::Permission; use crate::refusal::SharedRefusal; use crate::resource::Resource; use crate::scope::{Readable, Requires, TaskScope, can}; +use crate::store_paths::original_text; use axum::extract::{Path, Query, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -548,7 +549,7 @@ pub async fn get_task_entry_points( } else { EntryPointRelatedData::load(&db.web_db, &evaluation, &entry_points) .await? - .build_summaries(&entry_points) + .build_summaries(&db.project_mapper(task.project).await?, &entry_points) }; Ok(ok_json(PaginatedEntryPoints { @@ -685,7 +686,11 @@ impl EntryPointRelatedData { }) } - fn build_summaries(&self, entry_points: &[MEntryPoint]) -> Vec { + fn build_summaries( + &self, + mapper: &StorePathMapper, + entry_points: &[MEntryPoint], + ) -> Vec { let mut summaries = Vec::new(); for ep in entry_points { let Some(&build_id) = self.build_jobs.get(&ep.derivation) else { @@ -704,7 +709,7 @@ impl EntryPointRelatedData { summaries.push(EntryPointSummary { id: ep.id, build_id, - derivation_path: drv.drv_path(), + derivation_path: original_text(mapper, drv.drv_path()), eval: ep.eval.clone(), build_status, has_artefacts: *self.has_products.get(&ep.derivation).unwrap_or(&false), @@ -712,7 +717,10 @@ impl EntryPointRelatedData { .outputs .get(&ep.derivation) .cloned() - .unwrap_or_default(), + .unwrap_or_default() + .into_iter() + .map(|(name, path)| (name, original_text(mapper, path))) + .collect(), architecture: drv.architecture.clone(), build_time_ms: attempt.and_then(|a| a.duration_ms()), build_started_at: attempt.and_then(|a| a.build_started_at), diff --git a/backend/gradient-web/src/error.rs b/backend/gradient-web/src/error.rs index 517e6bbe0..ccca33adf 100644 --- a/backend/gradient-web/src/error.rs +++ b/backend/gradient-web/src/error.rs @@ -193,6 +193,15 @@ impl From for WebError { } } +impl From for WebError { + fn from(err: gradient_types::MappingError) -> Self { + Self::ServiceUnavailable( + ErrorCode::SERVICE_UNAVAILABLE, + format!("store paths of this project cannot be mapped: {err}"), + ) + } +} + impl From for WebError { fn from(err: JsonRejection) -> Self { let text = err.to_string(); diff --git a/backend/gradient-web/src/lib.rs b/backend/gradient-web/src/lib.rs index 8d7a8c2cd..784be5fd9 100644 --- a/backend/gradient-web/src/lib.rs +++ b/backend/gradient-web/src/lib.rs @@ -29,6 +29,7 @@ pub mod resource; mod routes; pub mod scim; pub mod scope; +mod store_paths; use axum::body::Body; use axum::extract::{DefaultBodyLimit, MatchedPath}; diff --git a/backend/gradient-web/src/store_paths.rs b/backend/gradient-web/src/store_paths.rs new file mode 100644 index 000000000..dbca02b4a --- /dev/null +++ b/backend/gradient-web/src/store_paths.rs @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use crate::error::{WebError, WebResult}; +use crate::resource::Resource; +use gradient_types::{StoreHash, StorePathMapper}; + +pub(crate) fn stored_hash(mapper: &StorePathMapper, hash: &str) -> WebResult { + let hash = StoreHash::parse(hash).map_err(|_| WebError::not_found(Resource::Path))?; + Ok(mapper.inbound_hash(hash).to_string()) +} + +pub(crate) fn original_text(mapper: &StorePathMapper, text: String) -> String { + mapper.outbound_text(&text).unwrap_or(text) +} diff --git a/backend/gradient-web/tests/build_paths.rs b/backend/gradient-web/tests/build_paths.rs new file mode 100644 index 000000000..25cbda4a7 --- /dev/null +++ b/backend/gradient-web/tests/build_paths.rs @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use axum::http::StatusCode; +use gradient_entity::ids::*; +use gradient_entity::{ + EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping, UnresolvableProjectMapping, +}; +use gradient_fixtures::ids::{project_id, task_id}; +use gradient_fixtures::mock::mock_db; +use gradient_fixtures::rows::{eval_at, public_project}; +use gradient_test_support::state::TestState; +use sea_orm::{DatabaseConnection, Value}; +use std::collections::BTreeMap; +use std::sync::Arc; +use uuid::Uuid; + +fn build_job_id() -> BuildJobId { + BuildJobId::new(Uuid::from_u128(0xb1)) +} +fn derivation_id() -> DerivationId { + DerivationId::new(Uuid::from_u128(0xd1)) +} +fn shared_build_id() -> DerivationBuildId { + DerivationBuildId::new(Uuid::from_u128(0xa1)) +} +fn eval_id() -> EvaluationId { + EvaluationId::new(Uuid::from_u128(0xe1)) +} + +fn build_rows() -> DatabaseConnection { + mock_db() + .append_query_results([vec![gradient_entity::build_job::Model { + id: build_job_id(), + evaluation: eval_id(), + derivation: derivation_id(), + derivation_build: shared_build_id(), + ..Default::default() + }]]) + .append_query_results([vec![gradient_entity::derivation_build::Model { + id: shared_build_id(), + derivation: derivation_id(), + ..Default::default() + }]]) + .append_query_results([vec![eval_at(eval_id(), 0)]]) + .append_query_results([vec![gradient_entity::task::Model { + id: task_id(), + project: project_id(), + ..Default::default() + }]]) + .append_query_results([vec![public_project()]]) + .append_query_results([vec![gradient_entity::derivation::Model { + id: derivation_id(), + hash: EXAMPLE_STORED_HASH.into(), + name: "hello".into(), + ..Default::default() + }]]) + .append_query_results([vec![gradient_entity::derivation_output::Model { + derivation: derivation_id(), + name: "out".into(), + hash: EXAMPLE_STORED_HASH.into(), + package: "hello".into(), + ..Default::default() + }]]) + .append_query_results([Vec::::new()]) + .append_query_results([Vec::>::new()]) + .into_connection() +} + +#[tokio::test] +async fn a_build_names_its_original_derivation_path() { + let original = EXAMPLE_ORIGINAL_HASH; + let db = build_rows(); + let server = TestState::new() + .web_db(db) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .server(); + + let response = server + .get(&format!("/api/v1/builds/{}", build_job_id())) + .await; + + response.assert_status_ok(); + let body: serde_json::Value = response.json(); + assert_eq!( + body["message"]["derivation_path"], + format!("{original}-hello.drv") + ); + assert_eq!( + body["message"]["output"]["out"], + format!("{original}-hello") + ); +} + +#[tokio::test] +async fn a_build_of_a_project_that_does_not_resolve_is_unavailable() { + let server = TestState::new() + .web_db(build_rows()) + .store_path_mapping(Arc::new(UnresolvableProjectMapping)) + .server(); + + let response = server + .get(&format!("/api/v1/builds/{}", build_job_id())) + .await; + + response.assert_status(StatusCode::SERVICE_UNAVAILABLE); + assert!(!response.text().contains(EXAMPLE_STORED_HASH)); +} diff --git a/backend/gradient-web/tests/nar_serve.rs b/backend/gradient-web/tests/nar_serve.rs index b1c60d275..4790e284a 100644 --- a/backend/gradient-web/tests/nar_serve.rs +++ b/backend/gradient-web/tests/nar_serve.rs @@ -11,12 +11,14 @@ use axum::extract::connect_info::MockConnectInfo; use axum_test::TestServer; +use gradient_entity::{MappingContext, RotatedHashMapping, StorePathMapping}; use gradient_fixtures::mock::mock_db; use gradient_storage::NarStore; use gradient_test_support::prelude::test_cli; use gradient_test_support::state::TestState; use gradient_types::ids::*; use std::net::SocketAddr; +use std::sync::Arc; use uuid::Uuid; const STORE_HASH: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; @@ -205,3 +207,40 @@ async fn a_nar_this_cache_holds_no_claim_on_is_not_served() { resp.assert_status_not_found(); } + +#[tokio::test] +async fn a_nar_url_file_hash_is_looked_up_unmapped() { + let cli = test_cli(); + let db = mock_db() + .append_query_results([vec![cache_row()]]) + .append_query_results([vec![cached_path_row()]]) + .append_query_results([vec![served()]]) + .into_connection(); + let nar_storage = NarStore::local(&cli.server.base_dir).expect("create test NarStore"); + nar_storage + .put(STORE_HASH, blob()) + .await + .expect("seed NAR blob"); + let state = TestState::new() + .web_db(db.clone()) + .nar_storage(nar_storage) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + + let router = gradient_web::create_router(state) + .expect("router") + .layer(MockConnectInfo( + "127.0.0.1:0".parse::().unwrap(), + )); + TestServer::new(router) + .get(&format!("/cache/test-cache/nar/{FILE_HASH_NIX32}.nar.zst")) + .await + .assert_status_ok(); + + let rotated = RotatedHashMapping + .inbound_prefix(MappingContext::Cache(cache_id()), FILE_HASH_NIX32) + .unwrap(); + let log = format!("{:?}", db.into_transaction_log()); + assert!(log.contains(&format!("sha256:{FILE_HASH_NIX32}")), "{log}"); + assert!(!log.contains(&rotated), "{log}"); +} diff --git a/backend/gradient-web/tests/narinfo.rs b/backend/gradient-web/tests/narinfo.rs index b9d377914..c0f8c19df 100644 --- a/backend/gradient-web/tests/narinfo.rs +++ b/backend/gradient-web/tests/narinfo.rs @@ -10,10 +10,13 @@ )] use axum_test::TestServer; +use gradient_entity::{EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH, RotatedHashMapping}; use gradient_fixtures::mock::mock_db; use gradient_test_support::state::TestState; use gradient_types::ids::*; use gradient_web::create_router; +use std::collections::BTreeMap; +use std::sync::Arc; use uuid::Uuid; fn cache_id() -> CacheId { @@ -185,3 +188,80 @@ async fn narinfo_returns_404_when_signature_null() { response.assert_status_not_found(); } + +#[tokio::test] +async fn a_narinfo_names_the_path_the_client_asked_for() { + let (original, stored) = (EXAMPLE_ORIGINAL_HASH, EXAMPLE_STORED_HASH); + let cache_row = gradient_entity::cache::Model { + id: cache_id(), + name: "test-cache".into(), + active: true, + public: true, + created_by: UserId::new(project_id().into_inner()), + created_at: test_date(), + ..Default::default() + }; + let cached_path_row = gradient_entity::cached_path::Model { + id: cached_path_id(), + hash: stored.into(), + package: "hello".into(), + file_hash: Some( + "sha256:0000000000000000000000000000000000000000000000000000000000000000".into(), + ), + file_size: Some(12345), + nar_size: Some(67890), + nar_hash: Some("sha256:0mdqa9w1p6cmli6976v4wi0sw9r4p5prkj7lzfd1877wk11c9c73".into()), + deriver: Some(format!("/nix/store/{stored}-hello.drv")), + created_at: test_date(), + ..Default::default() + }; + let signature_row = gradient_entity::cached_path_signature::Model { + id: cached_path_sig_id(), + cached_path: cached_path_id(), + cache: cache_id(), + signature: Some(vec![0x42; 64]), + created_at: test_date(), + ..Default::default() + }; + let references_row = BTreeMap::from([ + ("hash", sea_orm::Value::from(stored)), + ( + "references", + sea_orm::Value::from(format!("{stored}-hello")), + ), + ]); + let db = mock_db() + .append_query_results([vec![cache_row]]) + .append_query_results([vec![cached_path_row]]) + .append_query_results([vec![signature_row]]) + .append_query_results([vec![references_row]]) + .into_connection(); + let state = TestState::new() + .web_db(db.clone()) + .store_path_mapping(Arc::new(RotatedHashMapping)) + .build(); + + let response = TestServer::new(create_router(state).expect("router")) + .get(&format!("/cache/test-cache/{original}.narinfo")) + .await; + + response.assert_status_ok(); + let body = response.text(); + assert!( + body.contains(&format!("StorePath: /nix/store/{original}-hello\n")), + "{body}" + ); + assert!( + body.contains(&format!("References: {original}-hello\n")), + "{body}" + ); + assert!( + body.contains(&format!("Deriver: {original}-hello.drv\n")), + "{body}" + ); + let log = format!("{:?}", db.into_transaction_log()); + assert!( + log.contains(stored), + "the lookup uses the stored hash: {log}" + ); +} diff --git a/backend/src/lib.rs b/backend/src/lib.rs new file mode 100644 index 000000000..ae17210ae --- /dev/null +++ b/backend/src/lib.rs @@ -0,0 +1,181 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use gradient_core::{ServerExtensions, init_state}; +use gradient_types::Cli; +use gradient_types::cli::LogArgs; +use gradient_util::logging::{LogSetup, LogWriter, NOISY_DEPS, TraceSetup}; +use std::sync::Arc; +use tracing::info; + +fn log_overrides(log: &LogArgs) -> [(&'static str, Option<&str>); 7] { + [ + ("gradient_web", log.level_web.as_deref()), + ("gradient_cache", log.level_cache.as_deref()), + ("gradient_proto", log.level_proto.as_deref()), + ("gradient_wire", log.level_proto.as_deref()), + ("gradient_storage::passthrough", log.level_proto.as_deref()), + ("gradient_scheduler", log.level_scheduler.as_deref()), + ("gradient_pool", log.level_scheduler.as_deref()), + ] +} + +fn init_logging(log: &LogArgs) { + let overrides = log_overrides(log); + gradient_util::logging::init(&LogSetup { + level: &log.level_default, + overrides: &overrides, + quiet: NOISY_DEPS, + honor_rust_log: true, + writer: LogWriter::Stdout, + trace: log.trace_dir.as_deref().map(|dir| TraceSetup { + dir, + process: "server", + }), + }); +} + +pub async fn run(cli: Cli, extensions: ServerExtensions) -> std::io::Result<()> { + init_logging(&cli.log); + + if cli.state.validate { + return validate_state_and_exit(cli.state.file.as_deref()); + } + + let state = match init_state(cli, extensions).await { + Ok(s) => s, + Err(e) => { + tracing::error!(error = %e, "server bootstrap failed"); + std::process::exit(1); + } + }; + + info!( + version = env!("CARGO_PKG_VERSION"), + ip = %state.config.server.listen_addr, + port = state.config.server.port, + log_level = %state.config.log.level_default, + web_log_level = state.config.log.level_web.as_deref().unwrap_or("(default)"), + cache_log_level = state.config.log.level_cache.as_deref().unwrap_or("(default)"), + proto_log_level = state.config.log.level_proto.as_deref().unwrap_or("(default)"), + scheduler_log_level = state.config.log.level_scheduler.as_deref().unwrap_or("(default)"), + "Starting Gradient server" + ); + + let _guard = if state.config.sentry.enable { + let dsn = gradient_types::cli::effective_sentry_dsn(&state.config.sentry); + info!(dsn, "Error reporting enabled - initializing Sentry"); + Some(sentry::init(dsn.to_string())) + } else { + info!("Error reporting disabled"); + None + }; + + state + .shutdown + .supervise_now(state.graph.child_spec(state.db())) + .await + .map_err(std::io::Error::other)?; + let scheduler = Arc::new(gradient_scheduler::Scheduler::new(Arc::clone(&state))); + state + .shutdown + .supervise_now(gradient_effects::child_spec( + Arc::clone(&state), + Arc::clone(&scheduler), + )) + .await + .map_err(std::io::Error::other)?; + + info!("Starting cache service"); + gradient_cache::start_cache(Arc::clone(&state)).await?; + + if state.config.ssh.enable { + info!("Starting SSH service"); + gradient_ssh::start(Arc::clone(&state)).await?; + } + + info!("Starting web service"); + gradient_web::serve_web(Arc::clone(&state), scheduler).await?; + + Ok(()) +} + +fn validate_state_and_exit(state_file: Option<&str>) -> std::io::Result<()> { + let Some(path) = state_file else { + eprintln!("--state-validate requires --state-file"); + std::process::exit(2); + }; + match gradient_state::validate_state_file(path) { + Ok(errors) if errors.is_empty() => { + println!("State configuration '{path}' is valid"); + Ok(()) + } + Ok(errors) => { + eprintln!("State configuration '{path}' is invalid:"); + for e in &errors { + eprintln!(" - {e}"); + } + std::process::exit(1); + } + Err(e) => { + eprintln!("Failed to load state file '{path}': {e}"); + std::process::exit(1); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn server_directive(log: &LogArgs) -> String { + let overrides = log_overrides(log); + gradient_util::logging::directive(&LogSetup { + level: &log.level_default, + overrides: &overrides, + quiet: NOISY_DEPS, + honor_rust_log: true, + writer: LogWriter::Stdout, + trace: None, + }) + } + + #[test] + fn directive_targets_gradient_crates_and_suppresses_noise() { + let logging = LogArgs { + level_default: "info".into(), + level_web: Some("debug".into()), + level_scheduler: Some("trace".into()), + ..Default::default() + }; + let d = server_directive(&logging); + assert!(d.starts_with("info,")); + assert!(d.contains("hyper=warn")); + assert!(d.contains("sqlx=warn")); + assert!(d.contains("gradient_web=debug")); + assert!(d.contains("gradient_scheduler=trace")); + assert!(d.contains("gradient_pool=trace")); + assert!(!d.contains("builder=")); + } + + #[test] + fn the_proto_level_reaches_the_passed_through_nar_sender() { + let logging = LogArgs { + level_proto: Some("trace".into()), + ..Default::default() + }; + let d = server_directive(&logging); + assert!(d.contains("gradient_storage::passthrough=trace"), "{d}"); + } + + #[test] + fn directive_without_overrides_is_global_plus_noise() { + let d = server_directive(&LogArgs::default()); + assert!(d.starts_with("info,")); + assert!(d.contains("rustls=warn")); + assert!(!d.contains("gradient_web=")); + } +} diff --git a/backend/src/main.rs b/backend/src/main.rs index 8204a44f4..15a7cd730 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -5,39 +5,8 @@ */ use clap::Parser; -use gradient_core::init_state; +use gradient_core::ServerExtensions; use gradient_types::Cli; -use gradient_types::cli::LogArgs; -use gradient_util::logging::{LogSetup, LogWriter, NOISY_DEPS, TraceSetup}; -use std::sync::Arc; -use tracing::info; - -fn log_overrides(log: &LogArgs) -> [(&'static str, Option<&str>); 7] { - [ - ("gradient_web", log.level_web.as_deref()), - ("gradient_cache", log.level_cache.as_deref()), - ("gradient_proto", log.level_proto.as_deref()), - ("gradient_wire", log.level_proto.as_deref()), - ("gradient_storage::passthrough", log.level_proto.as_deref()), - ("gradient_scheduler", log.level_scheduler.as_deref()), - ("gradient_pool", log.level_scheduler.as_deref()), - ] -} - -fn init_logging(log: &LogArgs) { - let overrides = log_overrides(log); - gradient_util::logging::init(&LogSetup { - level: &log.level_default, - overrides: &overrides, - quiet: NOISY_DEPS, - honor_rust_log: true, - writer: LogWriter::Stdout, - trace: log.trace_dir.as_deref().map(|dir| TraceSetup { - dir, - process: "server", - }), - }); -} pub fn main() -> std::io::Result<()> { // rustls 0.23 is panicking on a TLS handshake without an installed provider. See issue #232. @@ -48,149 +17,8 @@ pub fn main() -> std::io::Result<()> { .thread_stack_size(8 * 1024 * 1024) .build() .expect("Failed to build tokio runtime") - .block_on(run()) -} - -async fn run() -> std::io::Result<()> { - let cli = Cli::parse(); - init_logging(&cli.log); - - if cli.state.validate { - return validate_state_and_exit(cli.state.file.as_deref()); - } - - let state = match init_state(cli).await { - Ok(s) => s, - Err(e) => { - tracing::error!(error = %e, "server bootstrap failed"); - std::process::exit(1); - } - }; - - info!( - version = env!("CARGO_PKG_VERSION"), - ip = %state.config.server.listen_addr, - port = state.config.server.port, - log_level = %state.config.log.level_default, - web_log_level = state.config.log.level_web.as_deref().unwrap_or("(default)"), - cache_log_level = state.config.log.level_cache.as_deref().unwrap_or("(default)"), - proto_log_level = state.config.log.level_proto.as_deref().unwrap_or("(default)"), - scheduler_log_level = state.config.log.level_scheduler.as_deref().unwrap_or("(default)"), - "Starting Gradient server" - ); - - let _guard = if state.config.sentry.enable { - let dsn = gradient_types::cli::effective_sentry_dsn(&state.config.sentry); - info!(dsn, "Error reporting enabled - initializing Sentry"); - Some(sentry::init(dsn.to_string())) - } else { - info!("Error reporting disabled"); - None - }; - - state - .shutdown - .supervise_now(state.graph.child_spec(state.db())) - .await - .map_err(std::io::Error::other)?; - let scheduler = Arc::new(gradient_scheduler::Scheduler::new(Arc::clone(&state))); - state - .shutdown - .supervise_now(gradient_effects::child_spec( - Arc::clone(&state), - Arc::clone(&scheduler), + .block_on(gradient_server::run( + Cli::parse(), + ServerExtensions::default(), )) - .await - .map_err(std::io::Error::other)?; - - info!("Starting cache service"); - gradient_cache::start_cache(Arc::clone(&state)).await?; - - if state.config.ssh.enable { - info!("Starting SSH service"); - gradient_ssh::start(Arc::clone(&state)).await?; - } - - info!("Starting web service"); - gradient_web::serve_web(Arc::clone(&state), scheduler).await?; - - Ok(()) -} - -/// A non-zero exit is failing a NixOS build or CI run fast on an invalid state file. -fn validate_state_and_exit(state_file: Option<&str>) -> std::io::Result<()> { - let Some(path) = state_file else { - eprintln!("--state-validate requires --state-file"); - std::process::exit(2); - }; - match gradient_state::validate_state_file(path) { - Ok(errors) if errors.is_empty() => { - println!("State configuration '{path}' is valid"); - Ok(()) - } - Ok(errors) => { - eprintln!("State configuration '{path}' is invalid:"); - for e in &errors { - eprintln!(" - {e}"); - } - std::process::exit(1); - } - Err(e) => { - eprintln!("Failed to load state file '{path}': {e}"); - std::process::exit(1); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn server_directive(log: &LogArgs) -> String { - let overrides = log_overrides(log); - gradient_util::logging::directive(&LogSetup { - level: &log.level_default, - overrides: &overrides, - quiet: NOISY_DEPS, - honor_rust_log: true, - writer: LogWriter::Stdout, - trace: None, - }) - } - - #[test] - fn directive_targets_gradient_crates_and_suppresses_noise() { - let logging = LogArgs { - level_default: "info".into(), - level_web: Some("debug".into()), - level_scheduler: Some("trace".into()), - ..Default::default() - }; - let d = server_directive(&logging); - assert!(d.starts_with("info,")); - assert!(d.contains("hyper=warn")); - assert!(d.contains("sqlx=warn")); - assert!(d.contains("gradient_web=debug")); - assert!(d.contains("gradient_scheduler=trace")); - assert!(d.contains("gradient_pool=trace")); - assert!(!d.contains("builder=")); - } - - #[test] - fn the_proto_level_reaches_the_passed_through_nar_sender() { - let logging = LogArgs { - level_proto: Some("trace".into()), - ..Default::default() - }; - let d = server_directive(&logging); - assert!(d.contains("gradient_storage::passthrough=trace"), "{d}"); - } - - #[test] - fn directive_without_overrides_is_global_plus_noise() { - let d = server_directive(&LogArgs::default()); - assert!(d.starts_with("info,")); - assert!(d.contains("rustls=warn")); - assert!(!d.contains("gradient_web=")); - } } diff --git a/docs/src/contributors/proto/jobs.md b/docs/src/contributors/proto/jobs.md index bf9d757f1..a262c97a7 100644 --- a/docs/src/contributors/proto/jobs.md +++ b/docs/src/contributors/proto/jobs.md @@ -50,6 +50,7 @@ Flake jobs fetch and evaluate a flake. The server can pick the steps at queue ti | `EvaluateDerivations` | Eval cache pull, attribute discovery, closure walk, eval cache push | - **Eval Cache:** Workers pull the evaluation cache before discovery and push the cache after a walk with derivations. Empty or schemaless caches stay local. +- **Eval Cache Entries:** Entries belong to a single cache. Pulls read the newest entry among the readable caches of the project. Pushes write one entry into every writable cache of the project. - **Attribute Failures:** Workers send them as `EvalMessage { level: Error, source: "nix-eval:" }` before the last `EvalResult`. - **Import From Derivation:** Workers send `ImportRequest { drv_paths }` during discovery and wait for `ImportResult`. - **Missing Inputs:** Workers send `InputRequest { request }` for an input in neither the store nor the cache, and wait for `InputResult`. The server can send the download as `FetchInput` to a worker with `network`. That worker can push the input under the `fetch_id` and answer with `InputFetched`.