maybe done refactoring but not tested
This commit is contained in:
parent
f51099b911
commit
e58cf2bbd0
3 changed files with 110 additions and 166 deletions
137
src/main.rs
137
src/main.rs
|
@ -1,17 +1,8 @@
|
||||||
#![feature(let_chains)]
|
#![feature(let_chains)]
|
||||||
#![feature(extend_one)]
|
#![feature(extend_one)]
|
||||||
|
|
||||||
use std::path::Path;
|
|
||||||
use std::sync::{
|
|
||||||
Arc, Mutex, RwLock,
|
|
||||||
atomic::{AtomicUsize, Ordering},
|
|
||||||
};
|
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use clap::{Parser, Subcommand};
|
use clap::{Parser, Subcommand};
|
||||||
use log::{debug, trace};
|
|
||||||
use tokio::process::Command;
|
|
||||||
use tokio::sync::{Semaphore, mpsc};
|
|
||||||
|
|
||||||
use nixcp::NixCp;
|
use nixcp::NixCp;
|
||||||
|
|
||||||
|
@ -60,138 +51,14 @@ enum Commands {
|
||||||
async fn main() -> Result<()> {
|
async fn main() -> Result<()> {
|
||||||
env_logger::init();
|
env_logger::init();
|
||||||
let cli = Cli::parse();
|
let cli = Cli::parse();
|
||||||
let mut nixcp = NixCp::new();
|
let nixcp = Box::leak(Box::new(NixCp::with_upstreams(&cli.upstream_caches)?));
|
||||||
nixcp.add_upstreams(&cli.upstream_caches)?;
|
|
||||||
|
|
||||||
match &cli.command {
|
match &cli.command {
|
||||||
Commands::Push { package } => {
|
Commands::Push { package } => {
|
||||||
nixcp.paths_from_package(package).await?;
|
nixcp.paths_from_package(package).await?;
|
||||||
|
nixcp.run().await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|
||||||
/*
|
|
||||||
let (cacheable_tx, mut cacheable_rx) = mpsc::channel(cli.uploader_concurrency.into());
|
|
||||||
|
|
||||||
println!("spawning check_upstream");
|
|
||||||
|
|
||||||
println!("spawning uploader");
|
|
||||||
handles.push(tokio::spawn(async move {
|
|
||||||
uploader(&mut cacheable_rx, binary_cache, cli.uploader_concurrency).await;
|
|
||||||
}));
|
|
||||||
|
|
||||||
// make sure all threads are done
|
|
||||||
for handle in handles {
|
|
||||||
handle.await.unwrap();
|
|
||||||
}
|
|
||||||
*/
|
|
||||||
}
|
|
||||||
|
|
||||||
// filter out store paths that exist in upstream caches
|
|
||||||
async fn check_upstream(
|
|
||||||
store_paths: Arc<RwLock<Vec<String>>>,
|
|
||||||
cacheable_tx: mpsc::Sender<String>,
|
|
||||||
concurrency: u8,
|
|
||||||
upstream_caches: Arc<Vec<String>>,
|
|
||||||
) {
|
|
||||||
let concurrency = Arc::new(Semaphore::new(concurrency.into()));
|
|
||||||
let c_store_paths = Arc::clone(&store_paths);
|
|
||||||
let store_paths = c_store_paths.read().unwrap().clone();
|
|
||||||
|
|
||||||
for store_path in store_paths {
|
|
||||||
let tx = cacheable_tx.clone();
|
|
||||||
let upstream_caches = Arc::clone(&upstream_caches);
|
|
||||||
let concurrency = Arc::clone(&concurrency);
|
|
||||||
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let _permit = concurrency.acquire().await.unwrap();
|
|
||||||
let basename = Path::new(&store_path)
|
|
||||||
.file_name()
|
|
||||||
.unwrap()
|
|
||||||
.to_str()
|
|
||||||
.unwrap()
|
|
||||||
.to_string();
|
|
||||||
let hash = basename.split("-").next().unwrap();
|
|
||||||
|
|
||||||
let mut hit = false;
|
|
||||||
for upstream in upstream_caches.as_ref() {
|
|
||||||
let mut uri = upstream.clone();
|
|
||||||
uri.push_str(format!("/{}.narinfo", hash).as_str());
|
|
||||||
|
|
||||||
let res_status = reqwest::Client::new()
|
|
||||||
.head(uri)
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
.map(|x| x.status());
|
|
||||||
|
|
||||||
if let Ok(res_status) = res_status
|
|
||||||
&& res_status.is_success()
|
|
||||||
{
|
|
||||||
debug!("{} was a hit upstream: {}", store_path, upstream);
|
|
||||||
hit = true;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if !hit {
|
|
||||||
trace!("sending {}", store_path);
|
|
||||||
tx.send(store_path).await.unwrap();
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn uploader(
|
|
||||||
cacheable_rx: &mut mpsc::Receiver<String>,
|
|
||||||
binary_cache: String,
|
|
||||||
concurrency: u8,
|
|
||||||
) {
|
|
||||||
let upload_count = Arc::new(AtomicUsize::new(0));
|
|
||||||
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
|
||||||
let concurrency = Arc::new(Semaphore::new(concurrency.into()));
|
|
||||||
let mut handles = Vec::new();
|
|
||||||
|
|
||||||
loop {
|
|
||||||
if let Some(path_to_upload) = cacheable_rx.recv().await {
|
|
||||||
let concurrency = Arc::clone(&concurrency);
|
|
||||||
let failures = Arc::clone(&failures);
|
|
||||||
let binary_cache = binary_cache.clone();
|
|
||||||
let upload_count = Arc::clone(&upload_count);
|
|
||||||
|
|
||||||
handles.push(tokio::spawn(async move {
|
|
||||||
let _permit = concurrency.acquire().await.unwrap();
|
|
||||||
println!("uploading: {}", path_to_upload);
|
|
||||||
if Command::new("nix")
|
|
||||||
.arg("copy")
|
|
||||||
.arg("--to")
|
|
||||||
.arg(&binary_cache)
|
|
||||||
.arg(&path_to_upload)
|
|
||||||
.output()
|
|
||||||
.await
|
|
||||||
.is_err()
|
|
||||||
{
|
|
||||||
println!("WARN: upload failed: {}", path_to_upload);
|
|
||||||
failures.lock().unwrap().push(path_to_upload);
|
|
||||||
} else {
|
|
||||||
upload_count.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}));
|
|
||||||
} else {
|
|
||||||
// make sure all threads are done
|
|
||||||
for handle in handles {
|
|
||||||
handle.await.unwrap();
|
|
||||||
}
|
|
||||||
println!("uploaded {} paths", upload_count.load(Ordering::Relaxed));
|
|
||||||
|
|
||||||
let failures = failures.lock().unwrap();
|
|
||||||
if !failures.is_empty() {
|
|
||||||
println!("failed to upload these paths: ");
|
|
||||||
for failure in failures.iter() {
|
|
||||||
print!("{}", failure);
|
|
||||||
}
|
|
||||||
println!();
|
|
||||||
}
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
123
src/nixcp.rs
123
src/nixcp.rs
|
@ -1,64 +1,135 @@
|
||||||
use std::sync::Arc;
|
use std::{
|
||||||
|
iter::once,
|
||||||
|
sync::{
|
||||||
|
Arc, Mutex,
|
||||||
|
atomic::{AtomicUsize, Ordering},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
use crate::path_info::PathInfo;
|
use crate::path_info::PathInfo;
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result};
|
||||||
use log::info;
|
use log::{info, warn};
|
||||||
use tokio::sync::{Semaphore, mpsc};
|
use tokio::{
|
||||||
|
process::Command,
|
||||||
|
sync::{RwLock, Semaphore, mpsc},
|
||||||
|
};
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
pub struct NixCp {
|
pub struct NixCp {
|
||||||
upstream_caches: Arc<Vec<Url>>,
|
upstream_caches: Arc<Vec<Url>>,
|
||||||
store_paths: Vec<PathInfo>,
|
store_paths: Arc<RwLock<Vec<PathInfo>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl NixCp {
|
impl NixCp {
|
||||||
pub fn new() -> Self {
|
pub fn with_upstreams(new_upstreams: &[String]) -> Result<Self> {
|
||||||
Self {
|
let mut upstreams = Vec::with_capacity(new_upstreams.len() + 1);
|
||||||
upstream_caches: vec![Url::parse("https://cache.nixos.org").unwrap()],
|
for upstream in new_upstreams
|
||||||
store_paths: Vec::new(),
|
.iter()
|
||||||
}
|
.chain(once(&"https://cache.nixos.org".to_string()))
|
||||||
}
|
{
|
||||||
|
upstreams
|
||||||
pub fn add_upstreams(&mut self, upstreams: &[String]) -> Result<()> {
|
|
||||||
self.upstream_caches.reserve(upstreams.len());
|
|
||||||
for upstream in upstreams {
|
|
||||||
self.upstream_caches
|
|
||||||
.push(Url::parse(upstream).context(format!("failed to parse {upstream} as url"))?);
|
.push(Url::parse(upstream).context(format!("failed to parse {upstream} as url"))?);
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(Self {
|
||||||
|
upstream_caches: Arc::new(upstreams),
|
||||||
|
store_paths: Arc::new(RwLock::new(Vec::new())),
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn paths_from_package(&mut self, package: &str) -> Result<()> {
|
pub async fn paths_from_package(&mut self, package: &str) -> Result<()> {
|
||||||
let path_info = PathInfo::from_path(package).await?;
|
let path_info = PathInfo::from_path(package).await?;
|
||||||
self.store_paths = path_info.get_closure().await?;
|
self.store_paths
|
||||||
info!("found {} store paths", self.store_paths.len());
|
.write()
|
||||||
|
.await
|
||||||
|
.extend(path_info.get_closure().await?);
|
||||||
|
info!("found {} store paths", self.store_paths.read().await.len());
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn run(&mut self) {}
|
pub async fn run(&'static self) {
|
||||||
|
let (tx, rx) = mpsc::channel(10);
|
||||||
|
let tx = Arc::new(tx);
|
||||||
|
tokio::spawn(self.filter_from_upstream(tx));
|
||||||
|
tokio::spawn(self.uploader("".to_string(), rx));
|
||||||
|
}
|
||||||
|
|
||||||
/// filter paths that are on upstream and send to `tx`
|
/// filter paths that are on upstream and send to `tx`
|
||||||
async fn filter_from_upstream(&self, tx: mpsc::Sender<&PathInfo>) {
|
async fn filter_from_upstream(&self, tx: Arc<mpsc::Sender<String>>) {
|
||||||
let permits = Arc::new(Semaphore::new(10));
|
let permits = Arc::new(Semaphore::new(10));
|
||||||
let mut handles = Vec::new();
|
let mut handles = Vec::with_capacity(10);
|
||||||
for path in &self.store_paths {
|
let store_paths = self.store_paths.read().await.clone();
|
||||||
|
|
||||||
|
for path in store_paths.into_iter() {
|
||||||
if path.check_upstream_signature(&self.upstream_caches) {
|
if path.check_upstream_signature(&self.upstream_caches) {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
handles.push({
|
||||||
let permits = permits.clone();
|
let permits = permits.clone();
|
||||||
let tx = tx.clone();
|
let tx = tx.clone();
|
||||||
handles.push(tokio::spawn(async move {
|
let upstream_caches = self.upstream_caches.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
let _permit = permits.acquire().await.unwrap();
|
let _permit = permits.acquire().await.unwrap();
|
||||||
|
|
||||||
if !path.check_upstream_hit(&self.upstream_caches).await {
|
if !path.check_upstream_hit(upstream_caches.as_slice()).await {
|
||||||
tx.send(path);
|
tx.send(path.to_string()).await.unwrap();
|
||||||
}
|
}
|
||||||
}));
|
})
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
for handle in handles {
|
for handle in handles {
|
||||||
handle.await.unwrap();
|
handle.await.unwrap();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn uploader(&self, cache: String, mut rx: mpsc::Receiver<String>) {
|
||||||
|
let upload_count = Arc::new(AtomicUsize::new(0));
|
||||||
|
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
||||||
|
let permits = Arc::new(Semaphore::new(10));
|
||||||
|
let mut handles = Vec::with_capacity(10);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
if let Some(path_to_upload) = rx.recv().await {
|
||||||
|
let permits = Arc::clone(&permits);
|
||||||
|
let failures = Arc::clone(&failures);
|
||||||
|
let binary_cache = cache.clone();
|
||||||
|
let upload_count = Arc::clone(&upload_count);
|
||||||
|
|
||||||
|
handles.push(tokio::spawn(async move {
|
||||||
|
let _permit = permits.acquire().await.unwrap();
|
||||||
|
info!("uploading: {}", path_to_upload.to_string());
|
||||||
|
if Command::new("nix")
|
||||||
|
.arg("copy")
|
||||||
|
.arg("--to")
|
||||||
|
.arg(&binary_cache)
|
||||||
|
.arg(&path_to_upload.to_string())
|
||||||
|
.output()
|
||||||
|
.await
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
warn!("upload failed: {}", path_to_upload);
|
||||||
|
failures.lock().unwrap().push(path_to_upload);
|
||||||
|
} else {
|
||||||
|
upload_count.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
} else {
|
||||||
|
// make sure all threads are done
|
||||||
|
for handle in handles {
|
||||||
|
handle.await.unwrap();
|
||||||
|
}
|
||||||
|
println!("uploaded {} paths", upload_count.load(Ordering::Relaxed));
|
||||||
|
|
||||||
|
let failures = failures.lock().unwrap();
|
||||||
|
if !failures.is_empty() {
|
||||||
|
warn!("failed to upload these paths: ");
|
||||||
|
for failure in failures.iter() {
|
||||||
|
warn!("{}", failure);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -8,7 +8,7 @@ use tokio::process::Command;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
// nix path-info --derivation --json
|
// nix path-info --derivation --json
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
#[serde(rename_all = "camelCase")]
|
#[serde(rename_all = "camelCase")]
|
||||||
pub struct PathInfo {
|
pub struct PathInfo {
|
||||||
deriver: String,
|
deriver: String,
|
||||||
|
@ -103,6 +103,12 @@ impl PathInfo {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl ToString for PathInfo {
|
||||||
|
fn to_string(&self) -> String {
|
||||||
|
self.path.clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue