many changes

src/nixcp.rs

@@ -1,4 +1,5 @@
 use std::{
+    fs,
     iter::once,
     sync::{
         Arc, Mutex,
@@ -6,56 +7,86 @@ use std::{
     },
 };
 
-use crate::path_info::PathInfo;
 use anyhow::{Context, Result};
-use log::{info, warn};
-use tokio::{
-    process::Command,
-    sync::{RwLock, Semaphore, mpsc},
-};
+use aws_config::Region;
+use aws_sdk_s3 as s3;
+use futures::future::join_all;
+use log::{debug, info, warn};
+use nix_compat::narinfo::{self, SigningKey};
+use tokio::sync::{RwLock, Semaphore, mpsc};
 use url::Url;
 
+use crate::{Cli, path_info::PathInfo, uploader::Uploader};
+
 pub struct NixCp {
     upstream_caches: Arc<Vec<Url>>,
     store_paths: Arc<RwLock<Vec<PathInfo>>>,
+    s3_client: s3::Client,
+    signing_key: SigningKey<ed25519_dalek::SigningKey>,
+    bucket: String,
 }
 
 impl NixCp {
-    pub fn with_upstreams(new_upstreams: &[String]) -> Result<Self> {
-        let mut upstreams = Vec::with_capacity(new_upstreams.len() + 1);
-        for upstream in new_upstreams
+    pub async fn new(cli: &Cli) -> Result<Self> {
+        let mut upstreams = Vec::with_capacity(cli.upstreams.len() + 1);
+        for upstream in cli
+            .upstreams
             .iter()
             .chain(once(&"https://cache.nixos.org".to_string()))
         {
             upstreams
                 .push(Url::parse(upstream).context(format!("failed to parse {upstream} as url"))?);
         }
+
+        let key = fs::read_to_string(&cli.signing_key)?;
+        let signing_key = narinfo::parse_keypair(key.as_str())?.0;
+
+        let mut s3_config = aws_config::from_env();
+        if let Some(region) = &cli.region {
+            s3_config = s3_config.region(Region::new(region.clone()));
+        }
+        if let Some(endpoint) = &cli.endpoint {
+            s3_config = s3_config.endpoint_url(endpoint);
+        }
+        if let Some(profile) = &cli.profile {
+            s3_config = s3_config.profile_name(profile);
+        }
+
+        let s3_client = s3::Client::new(&s3_config.load().await);
         Ok(Self {
             upstream_caches: Arc::new(upstreams),
             store_paths: Arc::new(RwLock::new(Vec::new())),
+            s3_client,
+            signing_key,
+            bucket: cli.bucket.clone(),
         })
     }
 
     pub async fn paths_from_package(&mut self, package: &str) -> Result<()> {
-        let path_info = PathInfo::from_path(package).await?;
-        self.store_paths
-            .write()
+        let path_info = PathInfo::from_path(package)
             .await
-            .extend(path_info.get_closure().await?);
+            .context("get path info for package")?;
+        debug!("path-info for {package}: {:?}", path_info);
+        self.store_paths.write().await.extend(
+            path_info
+                .get_closure()
+                .await
+                .context("closure from path info")?,
+        );
         info!("found {} store paths", self.store_paths.read().await.len());
 
         Ok(())
     }
 
-    pub async fn run(&'static self) {
+    pub async fn run(&'static self) -> Result<()> {
         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));
+        self.upload(rx).await
     }
 
     /// filter paths that are on upstream and send to `tx`
-    async fn filter_from_upstream(&self, tx: Arc<mpsc::Sender<String>>) {
+    async fn filter_from_upstream(&self, tx: Arc<mpsc::Sender<PathInfo>>) {
         let permits = Arc::new(Semaphore::new(10));
         let mut handles = Vec::with_capacity(10);
         let store_paths = self.store_paths.read().await.clone();
@@ -72,7 +103,7 @@ impl NixCp {
                     let _permit = permits.acquire().await.unwrap();
 
                     if !path.check_upstream_hit(upstream_caches.as_slice()).await {
-                        tx.send(path.absolute_path()).await.unwrap();
+                        tx.send(path).await.unwrap();
                     }
                 })
             });
@@ -83,42 +114,32 @@ impl NixCp {
         }
     }
 
-    async fn uploader(&self, cache: String, mut rx: mpsc::Receiver<String>) {
-        let upload_count = Arc::new(AtomicUsize::new(0));
+    async fn upload(&'static self, mut rx: mpsc::Receiver<PathInfo>) -> Result<()> {
+        let upload_count = 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);
+        let mut uploads = 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);
+                let absolute_path = path_to_upload.absolute_path();
 
-                handles.push(tokio::spawn(async move {
+                info!("uploading: {}", absolute_path);
+                let uploader = Uploader::new(
+                    &self.signing_key,
+                    path_to_upload,
+                    &self.s3_client,
+                    self.bucket.clone(),
+                )?;
+
+                let fut = tokio::spawn({
                     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);
-                    }
-                }));
+                    async move { uploader.upload().await }
+                });
+                uploads.push(fut);
             } else {
-                // make sure all threads are done
-                for handle in handles {
-                    handle.await.unwrap();
-                }
+                join_all(uploads).await;
                 println!("uploaded {} paths", upload_count.load(Ordering::Relaxed));
 
                 let failures = failures.lock().unwrap();
@@ -131,5 +152,6 @@ impl NixCp {
                 break;
             }
         }
+        Ok(())
     }
 }