mirror of
https://github.com/zhaofengli/attic.git
synced 2024-12-14 11:57:30 +00:00
client: Implement watch-store
This commit is contained in:
parent
a2bc969594
commit
d540cc6888
7 changed files with 566 additions and 178 deletions
89
Cargo.lock
generated
89
Cargo.lock
generated
|
@ -194,6 +194,7 @@ dependencies = [
|
|||
"humantime",
|
||||
"indicatif",
|
||||
"lazy_static",
|
||||
"notify",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"serde",
|
||||
|
@ -1534,6 +1535,18 @@ dependencies = [
|
|||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "filetime"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4e884668cd0c7480504233e951174ddc3b382f7c2666e3b7310b5c4e7b0c37f9"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"redox_syscall",
|
||||
"windows-sys 0.42.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "flate2"
|
||||
version = "1.0.25"
|
||||
|
@ -2023,6 +2036,26 @@ dependencies = [
|
|||
"unicode-width",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inotify"
|
||||
version = "0.9.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8069d3ec154eb856955c1c0fbffefbf5f3c40a104ec912d4797314c1801abff"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"inotify-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inotify-sys"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e05c02b5e89bff3b946cedeca278abc628fe811e604f027c45a8aa3cf793d0eb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "instant"
|
||||
version = "0.1.12"
|
||||
|
@ -2131,6 +2164,26 @@ dependencies = [
|
|||
"sha2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kqueue"
|
||||
version = "1.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2c8fc60ba15bf51257aa9807a48a61013db043fcf3a78cb0d916e8e396dcad98"
|
||||
dependencies = [
|
||||
"kqueue-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "kqueue-sys"
|
||||
version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8367585489f01bc55dd27404dcf56b95e6da061a256a666ab23be9ba96a2e587"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.4.0"
|
||||
|
@ -2308,6 +2361,22 @@ dependencies = [
|
|||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "notify"
|
||||
version = "5.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "58ea850aa68a06e48fdb069c0ec44d0d64c8dbffa49bf3b6f7f0a901fdea1ba9"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"filetime",
|
||||
"inotify",
|
||||
"kqueue",
|
||||
"libc",
|
||||
"mio",
|
||||
"walkdir",
|
||||
"windows-sys 0.42.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nu-ansi-term"
|
||||
version = "0.46.0"
|
||||
|
@ -2995,6 +3064,15 @@ version = "1.0.12"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b4b9743ed687d4b4bcedf9ff5eaa7398495ae14e61cba0a295704edbc7decde"
|
||||
|
||||
[[package]]
|
||||
name = "same-file"
|
||||
version = "1.0.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
|
||||
dependencies = [
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schannel"
|
||||
version = "0.1.20"
|
||||
|
@ -4106,6 +4184,17 @@ version = "0.9.4"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "808cf2735cd4b6866113f648b791c6adc5714537bc222d9347bb203386ffda56"
|
||||
dependencies = [
|
||||
"same-file",
|
||||
"winapi",
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "want"
|
||||
version = "0.3.0"
|
||||
|
|
|
@ -24,6 +24,7 @@ futures = "0.3.25"
|
|||
humantime = "2.1.0"
|
||||
indicatif = "0.17.2"
|
||||
lazy_static = "1.4.0"
|
||||
notify = { version = "5.1.0", default-features = false, features = ["macos_kqueue"] }
|
||||
regex = "1.7.0"
|
||||
reqwest = { version = "0.11.13", default-features = false, features = ["json", "rustls-tls", "rustls-tls-native-roots", "stream"] }
|
||||
serde = { version = "1.0.151", features = ["derive"] }
|
||||
|
|
|
@ -12,6 +12,7 @@ use crate::command::get_closure::{self, GetClosure};
|
|||
use crate::command::login::{self, Login};
|
||||
use crate::command::push::{self, Push};
|
||||
use crate::command::r#use::{self, Use};
|
||||
use crate::command::watch_store::{self, WatchStore};
|
||||
|
||||
/// Attic binary cache client.
|
||||
#[derive(Debug, Parser)]
|
||||
|
@ -28,6 +29,7 @@ pub enum Command {
|
|||
Use(Use),
|
||||
Push(Push),
|
||||
Cache(Cache),
|
||||
WatchStore(WatchStore),
|
||||
|
||||
#[clap(hide = true)]
|
||||
GetClosure(GetClosure),
|
||||
|
@ -53,6 +55,7 @@ pub async fn run() -> Result<()> {
|
|||
Command::Use(_) => r#use::run(opts).await,
|
||||
Command::Push(_) => push::run(opts).await,
|
||||
Command::Cache(_) => cache::run(opts).await,
|
||||
Command::WatchStore(_) => watch_store::run(opts).await,
|
||||
Command::GetClosure(_) => get_closure::run(opts).await,
|
||||
}
|
||||
}
|
||||
|
|
|
@ -3,3 +3,4 @@ pub mod get_closure;
|
|||
pub mod login;
|
||||
pub mod push;
|
||||
pub mod r#use;
|
||||
pub mod watch_store;
|
||||
|
|
|
@ -1,19 +1,16 @@
|
|||
use std::collections::{HashMap, HashSet};
|
||||
use std::cmp;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use clap::Parser;
|
||||
use futures::future::join_all;
|
||||
use indicatif::MultiProgress;
|
||||
|
||||
use crate::api::ApiClient;
|
||||
use crate::cache::{CacheName, CacheRef};
|
||||
use crate::cache::CacheRef;
|
||||
use crate::cli::Opts;
|
||||
use crate::config::Config;
|
||||
use crate::push::{Pusher, PushConfig};
|
||||
use attic::nix_store::{NixStore, StorePath, StorePathHash, ValidPathInfo};
|
||||
use attic::nix_store::NixStore;
|
||||
|
||||
/// Push closures to a binary cache.
|
||||
#[derive(Debug, Parser)]
|
||||
|
@ -41,20 +38,6 @@ pub struct Push {
|
|||
force_preamble: bool,
|
||||
}
|
||||
|
||||
struct PushPlan {
|
||||
/// Store paths to push.
|
||||
store_path_map: HashMap<StorePathHash, ValidPathInfo>,
|
||||
|
||||
/// The number of paths in the original full closure.
|
||||
num_all_paths: usize,
|
||||
|
||||
/// Number of paths that have been filtered out because they are already cached.
|
||||
num_already_cached: usize,
|
||||
|
||||
/// Number of paths that have been filtered out because they are signed by an upstream cache.
|
||||
num_upstream: usize,
|
||||
}
|
||||
|
||||
pub async fn run(opts: Opts) -> Result<()> {
|
||||
let sub = opts.command.as_push().unwrap();
|
||||
if sub.jobs == 0 {
|
||||
|
@ -74,15 +57,24 @@ pub async fn run(opts: Opts) -> Result<()> {
|
|||
let (server_name, server, cache) = config.resolve_cache(&sub.cache)?;
|
||||
|
||||
let mut api = ApiClient::from_server_config(server.clone())?;
|
||||
let plan = PushPlan::plan(
|
||||
store.clone(),
|
||||
&mut api,
|
||||
cache,
|
||||
roots,
|
||||
sub.no_closure,
|
||||
sub.ignore_upstream_cache_filter,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Confirm remote cache validity, query cache config
|
||||
let cache_config = api.get_cache_config(cache).await?;
|
||||
|
||||
if let Some(api_endpoint) = &cache_config.api_endpoint {
|
||||
// Use delegated API endpoint
|
||||
api.set_endpoint(api_endpoint)?;
|
||||
}
|
||||
|
||||
let push_config = PushConfig {
|
||||
num_workers: sub.jobs,
|
||||
force_preamble: sub.force_preamble,
|
||||
};
|
||||
|
||||
let mp = MultiProgress::new();
|
||||
|
||||
let pusher = Pusher::new(store, api, cache.to_owned(), cache_config, mp, push_config);
|
||||
let plan = pusher.plan(roots, sub.no_closure, sub.ignore_upstream_cache_filter).await?;
|
||||
|
||||
if plan.store_path_map.is_empty() {
|
||||
if plan.num_all_paths == 0 {
|
||||
|
@ -106,16 +98,8 @@ pub async fn run(opts: Opts) -> Result<()> {
|
|||
);
|
||||
}
|
||||
|
||||
let push_config = PushConfig {
|
||||
num_workers: cmp::min(sub.jobs, plan.store_path_map.len()),
|
||||
force_preamble: sub.force_preamble,
|
||||
};
|
||||
|
||||
let mp = MultiProgress::new();
|
||||
|
||||
let pusher = Pusher::new(store, api, cache.to_owned(), mp, push_config);
|
||||
for (_, path_info) in plan.store_path_map {
|
||||
pusher.push(path_info).await?;
|
||||
pusher.queue(path_info).await?;
|
||||
}
|
||||
|
||||
let results = pusher.wait().await;
|
||||
|
@ -123,103 +107,3 @@ pub async fn run(opts: Opts) -> Result<()> {
|
|||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl PushPlan {
|
||||
/// Creates a plan.
|
||||
async fn plan(
|
||||
store: Arc<NixStore>,
|
||||
api: &mut ApiClient,
|
||||
cache: &CacheName,
|
||||
roots: Vec<StorePath>,
|
||||
no_closure: bool,
|
||||
ignore_upstream_filter: bool,
|
||||
) -> Result<Self> {
|
||||
// Compute closure
|
||||
let closure = if no_closure {
|
||||
roots
|
||||
} else {
|
||||
store
|
||||
.compute_fs_closure_multi(roots, false, false, false)
|
||||
.await?
|
||||
};
|
||||
|
||||
let mut store_path_map: HashMap<StorePathHash, ValidPathInfo> = {
|
||||
let futures = closure
|
||||
.iter()
|
||||
.map(|path| {
|
||||
let store = store.clone();
|
||||
let path = path.clone();
|
||||
let path_hash = path.to_hash();
|
||||
|
||||
async move {
|
||||
let path_info = store.query_path_info(path).await?;
|
||||
Ok((path_hash, path_info))
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
join_all(futures).await.into_iter().collect::<Result<_>>()?
|
||||
};
|
||||
|
||||
let num_all_paths = store_path_map.len();
|
||||
if store_path_map.is_empty() {
|
||||
return Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: 0,
|
||||
num_upstream: 0,
|
||||
});
|
||||
}
|
||||
|
||||
// Confirm remote cache validity, query cache config
|
||||
let cache_config = api.get_cache_config(cache).await?;
|
||||
|
||||
if let Some(api_endpoint) = &cache_config.api_endpoint {
|
||||
// Use delegated API endpoint
|
||||
api.set_endpoint(api_endpoint)?;
|
||||
}
|
||||
|
||||
if !ignore_upstream_filter {
|
||||
// Filter out paths signed by upstream caches
|
||||
let upstream_cache_key_names =
|
||||
cache_config.upstream_cache_key_names.unwrap_or_default();
|
||||
store_path_map.retain(|_, pi| {
|
||||
for sig in &pi.sigs {
|
||||
if let Some((name, _)) = sig.split_once(':') {
|
||||
if upstream_cache_key_names.iter().any(|u| name == u) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
});
|
||||
}
|
||||
|
||||
let num_filtered_paths = store_path_map.len();
|
||||
if store_path_map.is_empty() {
|
||||
return Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: 0,
|
||||
num_upstream: num_all_paths - num_filtered_paths,
|
||||
});
|
||||
}
|
||||
|
||||
// Query missing paths
|
||||
let missing_path_hashes: HashSet<StorePathHash> = {
|
||||
let store_path_hashes = store_path_map.keys().map(|sph| sph.to_owned()).collect();
|
||||
let res = api.get_missing_paths(cache, store_path_hashes).await?;
|
||||
res.missing_paths.into_iter().collect()
|
||||
};
|
||||
store_path_map.retain(|sph, _| missing_path_hashes.contains(sph));
|
||||
let num_missing_paths = store_path_map.len();
|
||||
|
||||
Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: num_filtered_paths - num_missing_paths,
|
||||
num_upstream: num_all_paths - num_filtered_paths,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
114
client/src/command/watch_store.rs
Normal file
114
client/src/command/watch_store.rs
Normal file
|
@ -0,0 +1,114 @@
|
|||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use clap::Parser;
|
||||
use indicatif::MultiProgress;
|
||||
use notify::{RecursiveMode, Watcher, EventKind};
|
||||
|
||||
use crate::api::ApiClient;
|
||||
use crate::cache::CacheRef;
|
||||
use crate::cli::Opts;
|
||||
use crate::config::Config;
|
||||
use crate::push::{Pusher, PushConfig, PushSessionConfig};
|
||||
use attic::nix_store::{NixStore, StorePath};
|
||||
|
||||
/// Watch the Nix Store for new paths and upload them to a binary cache.
|
||||
#[derive(Debug, Parser)]
|
||||
pub struct WatchStore {
|
||||
/// The cache to push to.
|
||||
cache: CacheRef,
|
||||
|
||||
/// Push the new paths only and do not compute closures.
|
||||
#[clap(long)]
|
||||
no_closure: bool,
|
||||
|
||||
/// Ignore the upstream cache filter.
|
||||
#[clap(long)]
|
||||
ignore_upstream_cache_filter: bool,
|
||||
|
||||
/// The maximum number of parallel upload processes.
|
||||
#[clap(short = 'j', long, default_value = "5")]
|
||||
jobs: usize,
|
||||
|
||||
/// Always send the upload info as part of the payload.
|
||||
#[clap(long, hide = true)]
|
||||
force_preamble: bool,
|
||||
}
|
||||
|
||||
pub async fn run(opts: Opts) -> Result<()> {
|
||||
let sub = opts.command.as_watch_store().unwrap();
|
||||
if sub.jobs == 0 {
|
||||
return Err(anyhow!("The number of jobs cannot be 0"));
|
||||
}
|
||||
|
||||
let config = Config::load()?;
|
||||
|
||||
let store = Arc::new(NixStore::connect()?);
|
||||
let store_dir = store.store_dir().to_owned();
|
||||
|
||||
let (server_name, server, cache) = config.resolve_cache(&sub.cache)?;
|
||||
let mut api = ApiClient::from_server_config(server.clone())?;
|
||||
|
||||
// Confirm remote cache validity, query cache config
|
||||
let cache_config = api.get_cache_config(cache).await?;
|
||||
|
||||
if let Some(api_endpoint) = &cache_config.api_endpoint {
|
||||
// Use delegated API endpoint
|
||||
api.set_endpoint(api_endpoint)?;
|
||||
}
|
||||
|
||||
let push_config = PushConfig {
|
||||
num_workers: sub.jobs,
|
||||
force_preamble: sub.force_preamble,
|
||||
};
|
||||
|
||||
let push_session_config = PushSessionConfig {
|
||||
no_closure: sub.no_closure,
|
||||
ignore_upstream_cache_filter: sub.ignore_upstream_cache_filter,
|
||||
};
|
||||
|
||||
let mp = MultiProgress::new();
|
||||
let session = Pusher::new(store.clone(), api, cache.to_owned(), cache_config, mp, push_config)
|
||||
.into_push_session(push_session_config);
|
||||
|
||||
let mut watcher = notify::recommended_watcher(move |res: notify::Result<notify::Event>| {
|
||||
match res {
|
||||
Ok(event) => {
|
||||
// We watch the removals of lock files which signify
|
||||
// store paths becoming valid
|
||||
if let EventKind::Remove(_) = event.kind {
|
||||
let paths = event.paths
|
||||
.iter()
|
||||
.filter_map(|p| {
|
||||
let base = strip_lock_file(&p)?;
|
||||
store.parse_store_path(base).ok()
|
||||
})
|
||||
.collect::<Vec<StorePath>>();
|
||||
|
||||
if !paths.is_empty() {
|
||||
session.queue_many(paths).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => eprintln!("Error during watch: {:?}", e),
|
||||
}
|
||||
})?;
|
||||
|
||||
watcher.watch(&store_dir, RecursiveMode::NonRecursive)?;
|
||||
|
||||
eprintln!("👀 Pushing new store paths to \"{cache}\" on \"{server}\"",
|
||||
cache = cache.as_str(),
|
||||
server = server_name.as_str(),
|
||||
);
|
||||
|
||||
loop {
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_lock_file(p: &Path) -> Option<PathBuf> {
|
||||
p.to_str()
|
||||
.and_then(|p| p.strip_suffix(".lock"))
|
||||
.filter(|t| !t.ends_with(".drv"))
|
||||
.map(PathBuf::from)
|
||||
}
|
|
@ -1,10 +1,21 @@
|
|||
//! Store path uploader.
|
||||
//!
|
||||
//! Multiple workers are spawned to upload store paths concurrently.
|
||||
//! There are two APIs: `Pusher` and `PushSession`.
|
||||
//!
|
||||
//! A `Pusher` simply dispatches `ValidPathInfo`s for workers to push. Use this
|
||||
//! when you know all store paths to push beforehand. The push plan (closure, missing
|
||||
//! paths, all path metadata) should be computed prior to pushing.
|
||||
//!
|
||||
//! A `PushSession`, on the other hand, accepts a stream of `StorePath`s and
|
||||
//! takes care of retrieving the closure and path metadata. It automatically
|
||||
//! batches expensive operations (closure computation, querying missing paths).
|
||||
//! Use this when the list of store paths is streamed from some external
|
||||
//! source (e.g., FS watcher, Unix Domain Socket) and a push plan cannot be
|
||||
//! created statically.
|
||||
//!
|
||||
//! TODO: Refactor out progress reporting and support a simple output style without progress bars
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fmt::Write;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
@ -18,11 +29,14 @@ use futures::stream::{Stream, TryStreamExt};
|
|||
use futures::future::join_all;
|
||||
use indicatif::{HumanBytes, MultiProgress, ProgressBar, ProgressState, ProgressStyle};
|
||||
use tokio::task::{JoinHandle, spawn};
|
||||
use tokio::time;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use attic::api::v1::cache_config::CacheConfig;
|
||||
use attic::api::v1::upload_path::{UploadPathNarInfo, UploadPathResult, UploadPathResultKind};
|
||||
use attic::cache::CacheName;
|
||||
use attic::error::AtticResult;
|
||||
use attic::nix_store::{NixStore, StorePath, ValidPathInfo};
|
||||
use attic::nix_store::{NixStore, StorePath, StorePathHash, ValidPathInfo};
|
||||
use crate::api::ApiClient;
|
||||
|
||||
type JobSender = channel::Sender<ValidPathInfo>;
|
||||
|
@ -38,16 +52,78 @@ pub struct PushConfig {
|
|||
pub force_preamble: bool,
|
||||
}
|
||||
|
||||
/// Configuration for a push session.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PushSessionConfig {
|
||||
/// Push the specified paths only and do not compute closures.
|
||||
pub no_closure: bool,
|
||||
|
||||
/// Ignore the upstream cache filter.
|
||||
pub ignore_upstream_cache_filter: bool,
|
||||
}
|
||||
|
||||
/// A handle to push store paths to a cache.
|
||||
///
|
||||
/// The caller is responsible for computing closures and
|
||||
/// checking for paths that already exist on the remote
|
||||
/// cache.
|
||||
pub struct Pusher {
|
||||
api: ApiClient,
|
||||
store: Arc<NixStore>,
|
||||
cache: CacheName,
|
||||
cache_config: CacheConfig,
|
||||
workers: Vec<JoinHandle<HashMap<StorePath, Result<()>>>>,
|
||||
sender: JobSender,
|
||||
}
|
||||
|
||||
/// A wrapper over a `Pusher` that accepts a stream of `StorePath`s.
|
||||
///
|
||||
/// Unlike a `Pusher`, a `PushSession` takes a stream of `StorePath`s
|
||||
/// instead of `ValidPathInfo`s, taking care of retrieving the closure
|
||||
/// and path metadata.
|
||||
///
|
||||
/// This is useful when the list of store paths is streamed from some
|
||||
/// external source (e.g., FS watcher, Unix Domain Socket) and a push
|
||||
/// plan cannot be computed statically.
|
||||
///
|
||||
/// ## Batching
|
||||
///
|
||||
/// Many store paths can be built in a short period of time, with each
|
||||
/// having a big closure. It can be very inefficient if we were to compute
|
||||
/// closure and query for missing paths for each individual path. This is
|
||||
/// especially true if we have a lot of remote builders (e.g., `attic watch-store`
|
||||
/// running alongside a beefy Hydra instance).
|
||||
///
|
||||
/// `PushSession` batches operations in order to minimize the number of
|
||||
/// closure computations and API calls. It also remembers which paths already
|
||||
/// exist on the remote cache. By default, it submits a batch if it's been 2
|
||||
/// seconds since the last path is queued or it's been 10 seconds in total.
|
||||
pub struct PushSession {
|
||||
/// Sender to the batching future.
|
||||
sender: channel::Sender<Vec<StorePath>>,
|
||||
}
|
||||
|
||||
enum SessionQueuePoll {
|
||||
Paths(Vec<StorePath>),
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct PushPlan {
|
||||
/// Store paths to push.
|
||||
pub store_path_map: HashMap<StorePathHash, ValidPathInfo>,
|
||||
|
||||
/// The number of paths in the original full closure.
|
||||
pub num_all_paths: usize,
|
||||
|
||||
/// Number of paths that have been filtered out because they are already cached.
|
||||
pub num_already_cached: usize,
|
||||
|
||||
/// Number of paths that have been filtered out because they are signed by an upstream cache.
|
||||
pub num_upstream: usize,
|
||||
}
|
||||
|
||||
/// Wrapper to update a progress bar as a NAR is streamed.
|
||||
struct NarStreamProgress<S> {
|
||||
stream: S,
|
||||
|
@ -55,12 +131,12 @@ struct NarStreamProgress<S> {
|
|||
}
|
||||
|
||||
impl Pusher {
|
||||
pub fn new(store: Arc<NixStore>, api: ApiClient, cache: CacheName, mp: MultiProgress, config: PushConfig) -> Self {
|
||||
pub fn new(store: Arc<NixStore>, api: ApiClient, cache: CacheName, cache_config: CacheConfig, mp: MultiProgress, config: PushConfig) -> Self {
|
||||
let (sender, receiver) = channel::unbounded();
|
||||
let mut workers = Vec::new();
|
||||
|
||||
for _ in 0..config.num_workers {
|
||||
workers.push(spawn(worker(
|
||||
workers.push(spawn(Self::worker(
|
||||
receiver.clone(),
|
||||
store.clone(),
|
||||
api.clone(),
|
||||
|
@ -70,11 +146,11 @@ impl Pusher {
|
|||
)));
|
||||
}
|
||||
|
||||
Self { workers, sender }
|
||||
Self { api, store, cache, cache_config, workers, sender }
|
||||
}
|
||||
|
||||
/// Sends a path to be pushed.
|
||||
pub async fn push(&self, path_info: ValidPathInfo) -> Result<()> {
|
||||
/// Queues a store path to be pushed.
|
||||
pub async fn queue(&self, path_info: ValidPathInfo) -> Result<()> {
|
||||
self.sender.send(path_info).await
|
||||
.map_err(|e| anyhow!(e))
|
||||
}
|
||||
|
@ -96,42 +172,262 @@ impl Pusher {
|
|||
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
async fn worker(
|
||||
receiver: JobReceiver,
|
||||
store: Arc<NixStore>,
|
||||
api: ApiClient,
|
||||
cache: CacheName,
|
||||
mp: MultiProgress,
|
||||
config: PushConfig,
|
||||
) -> HashMap<StorePath, Result<()>> {
|
||||
let mut results = HashMap::new();
|
||||
|
||||
loop {
|
||||
let path_info = match receiver.recv().await {
|
||||
Ok(path_info) => path_info,
|
||||
Err(_) => {
|
||||
// channel is closed - we are done
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let store_path = path_info.path.clone();
|
||||
|
||||
let r = upload_path(
|
||||
path_info,
|
||||
store.clone(),
|
||||
api.clone(),
|
||||
&cache,
|
||||
mp.clone(),
|
||||
config.force_preamble,
|
||||
).await;
|
||||
|
||||
results.insert(store_path, r);
|
||||
/// Creates a push plan.
|
||||
pub async fn plan(&self, roots: Vec<StorePath>, no_closure: bool, ignore_upstream_filter: bool) -> Result<PushPlan> {
|
||||
PushPlan::plan(
|
||||
self.store.clone(),
|
||||
&self.api,
|
||||
&self.cache,
|
||||
&self.cache_config,
|
||||
roots,
|
||||
no_closure,
|
||||
ignore_upstream_filter,
|
||||
).await
|
||||
}
|
||||
|
||||
results
|
||||
/// Converts the pusher into a `PushSession`.
|
||||
///
|
||||
/// This is useful when the list of store paths is streamed from some
|
||||
/// external source (e.g., FS watcher, Unix Domain Socket) and a push
|
||||
/// plan cannot be computed statically.
|
||||
pub fn into_push_session(self, config: PushSessionConfig) -> PushSession {
|
||||
PushSession::with_pusher(self, config)
|
||||
}
|
||||
|
||||
async fn worker(
|
||||
receiver: JobReceiver,
|
||||
store: Arc<NixStore>,
|
||||
api: ApiClient,
|
||||
cache: CacheName,
|
||||
mp: MultiProgress,
|
||||
config: PushConfig,
|
||||
) -> HashMap<StorePath, Result<()>> {
|
||||
let mut results = HashMap::new();
|
||||
|
||||
loop {
|
||||
let path_info = match receiver.recv().await {
|
||||
Ok(path_info) => path_info,
|
||||
Err(_) => {
|
||||
// channel is closed - we are done
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
let store_path = path_info.path.clone();
|
||||
|
||||
let r = upload_path(
|
||||
path_info,
|
||||
store.clone(),
|
||||
api.clone(),
|
||||
&cache,
|
||||
mp.clone(),
|
||||
config.force_preamble,
|
||||
).await;
|
||||
|
||||
results.insert(store_path, r);
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
impl PushSession {
|
||||
pub fn with_pusher(pusher: Pusher, config: PushSessionConfig) -> Self {
|
||||
let (sender, receiver) = channel::unbounded();
|
||||
|
||||
let known_paths_mutex = Arc::new(Mutex::new(HashSet::new()));
|
||||
|
||||
// FIXME
|
||||
spawn(async move {
|
||||
let pusher = Arc::new(pusher);
|
||||
loop {
|
||||
if let Err(e) = Self::worker(
|
||||
pusher.clone(),
|
||||
config.clone(),
|
||||
known_paths_mutex.clone(),
|
||||
receiver.clone(),
|
||||
).await {
|
||||
eprintln!("Worker exited: {:?}", e);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self {
|
||||
sender,
|
||||
}
|
||||
}
|
||||
|
||||
async fn worker(
|
||||
pusher: Arc<Pusher>,
|
||||
config: PushSessionConfig,
|
||||
known_paths_mutex: Arc<Mutex<HashSet<StorePathHash>>>,
|
||||
receiver: channel::Receiver<Vec<StorePath>>,
|
||||
) -> Result<()> {
|
||||
let mut roots = HashSet::new();
|
||||
|
||||
loop {
|
||||
// Get outstanding paths in queue
|
||||
let done = tokio::select! {
|
||||
// 2 seconds since last queued path
|
||||
done = async {
|
||||
loop {
|
||||
let poll = tokio::select! {
|
||||
r = receiver.recv() => match r {
|
||||
Ok(paths) => SessionQueuePoll::Paths(paths),
|
||||
_ => SessionQueuePoll::Closed,
|
||||
},
|
||||
_ = time::sleep(Duration::from_secs(2)) => SessionQueuePoll::TimedOut,
|
||||
};
|
||||
|
||||
match poll {
|
||||
SessionQueuePoll::Paths(store_paths) => {
|
||||
roots.extend(store_paths.into_iter());
|
||||
}
|
||||
SessionQueuePoll::Closed => {
|
||||
break true;
|
||||
}
|
||||
SessionQueuePoll::TimedOut => {
|
||||
break false;
|
||||
}
|
||||
}
|
||||
}
|
||||
} => done,
|
||||
|
||||
// 10 seconds
|
||||
_ = time::sleep(Duration::from_secs(10)) => {
|
||||
false
|
||||
},
|
||||
};
|
||||
|
||||
// Compute push plan
|
||||
let roots_vec: Vec<StorePath> = {
|
||||
let known_paths = known_paths_mutex.lock().await;
|
||||
roots.drain()
|
||||
.filter(|root| !known_paths.contains(&root.to_hash()))
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut plan = pusher.plan(roots_vec, config.no_closure, config.ignore_upstream_cache_filter).await?;
|
||||
|
||||
let mut known_paths = known_paths_mutex.lock().await;
|
||||
plan.store_path_map
|
||||
.retain(|sph, _| !known_paths.contains(&sph));
|
||||
|
||||
// Push everything
|
||||
for (store_path_hash, path_info) in plan.store_path_map.into_iter() {
|
||||
pusher.queue(path_info).await?;
|
||||
known_paths.insert(store_path_hash);
|
||||
}
|
||||
|
||||
drop(known_paths);
|
||||
|
||||
if done {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Queues multiple store paths to be pushed.
|
||||
pub fn queue_many(&self, store_paths: Vec<StorePath>) -> Result<()> {
|
||||
self.sender.send_blocking(store_paths)
|
||||
.map_err(|e| anyhow!(e))
|
||||
}
|
||||
}
|
||||
|
||||
impl PushPlan {
|
||||
/// Creates a plan.
|
||||
async fn plan(
|
||||
store: Arc<NixStore>,
|
||||
api: &ApiClient,
|
||||
cache: &CacheName,
|
||||
cache_config: &CacheConfig,
|
||||
roots: Vec<StorePath>,
|
||||
no_closure: bool,
|
||||
ignore_upstream_filter: bool,
|
||||
) -> Result<Self> {
|
||||
// Compute closure
|
||||
let closure = if no_closure {
|
||||
roots
|
||||
} else {
|
||||
store
|
||||
.compute_fs_closure_multi(roots, false, false, false)
|
||||
.await?
|
||||
};
|
||||
|
||||
let mut store_path_map: HashMap<StorePathHash, ValidPathInfo> = {
|
||||
let futures = closure
|
||||
.iter()
|
||||
.map(|path| {
|
||||
let store = store.clone();
|
||||
let path = path.clone();
|
||||
let path_hash = path.to_hash();
|
||||
|
||||
async move {
|
||||
let path_info = store.query_path_info(path).await?;
|
||||
Ok((path_hash, path_info))
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
join_all(futures).await.into_iter().collect::<Result<_>>()?
|
||||
};
|
||||
|
||||
let num_all_paths = store_path_map.len();
|
||||
if store_path_map.is_empty() {
|
||||
return Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: 0,
|
||||
num_upstream: 0,
|
||||
});
|
||||
}
|
||||
|
||||
if !ignore_upstream_filter {
|
||||
// Filter out paths signed by upstream caches
|
||||
let upstream_cache_key_names =
|
||||
cache_config.upstream_cache_key_names.as_ref().map_or([].as_slice(), |v| v.as_slice());
|
||||
store_path_map.retain(|_, pi| {
|
||||
for sig in &pi.sigs {
|
||||
if let Some((name, _)) = sig.split_once(':') {
|
||||
if upstream_cache_key_names.iter().any(|u| name == u) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
});
|
||||
}
|
||||
|
||||
let num_filtered_paths = store_path_map.len();
|
||||
if store_path_map.is_empty() {
|
||||
return Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: 0,
|
||||
num_upstream: num_all_paths - num_filtered_paths,
|
||||
});
|
||||
}
|
||||
|
||||
// Query missing paths
|
||||
let missing_path_hashes: HashSet<StorePathHash> = {
|
||||
let store_path_hashes = store_path_map.keys().map(|sph| sph.to_owned()).collect();
|
||||
let res = api.get_missing_paths(cache, store_path_hashes).await?;
|
||||
res.missing_paths.into_iter().collect()
|
||||
};
|
||||
store_path_map.retain(|sph, _| missing_path_hashes.contains(sph));
|
||||
let num_missing_paths = store_path_map.len();
|
||||
|
||||
Ok(Self {
|
||||
store_path_map,
|
||||
num_all_paths,
|
||||
num_already_cached: num_filtered_paths - num_missing_paths,
|
||||
num_upstream: num_all_paths - num_filtered_paths,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Uploads a single path to a cache.
|
||||
|
|
Loading…
Reference in a new issue