make things faster

src/main.rs

@@ -1,3 +1,5 @@
+#![feature(let_chains)]
+
 use std::process::{Command, Stdio};
 use std::sync::mpsc;
 use std::{env, path::Path};
@@ -5,7 +7,7 @@ use std::{env, path::Path};
 use log::{debug, trace};
 use serde::{Deserialize, Serialize};
 use serde_json;
-use tokio;
+use tokio::sync::Semaphore;
 
 const UPSTREAM_CACHES: &'static [&'static str] = &[
     "https://cache.nixos.org",
@@ -75,20 +77,31 @@ async fn main() {
     let store_paths = path_infos[0].get_store_paths();
     let (cacheable_tx, cacheable_rx) = mpsc::channel();
 
+    let mut handles = Vec::new();
+    
     println!("spawning check_upstream");
-    tokio::spawn(async move {
+    handles.push(tokio::spawn(async move {
         check_upstream(store_paths, cacheable_tx).await;
-    });
+    }));
 
     println!("spawning uploader");
-    tokio::spawn(async move {
+    handles.push(tokio::spawn(async move {
         uploader(cacheable_rx).await;
-    }).await.unwrap();
+    }));
+
+    // 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: Vec<String>, cacheable_tx: mpsc::Sender<String>) {
+    let concurrent = Semaphore::new(50);
     for store_path in store_paths {
+        let _ = concurrent.acquire().await.unwrap();
+        let tx = cacheable_tx.clone();
+        tokio::spawn(async move {
             let basename = Path::new(&store_path)
                 .file_name()
                 .unwrap()
@@ -106,27 +119,31 @@ async fn check_upstream(store_paths: Vec<String>, cacheable_tx: mpsc::Sender<Str
                     .head(uri)
                     .send()
                     .await
-                .unwrap()
+                    .map(|x| x.status());
-                .status();
 
-            if res_status.is_success() {
+                if let Ok(res_status) = res_status && res_status.is_success() {
-                debug!("{} was a hit upstream: {}", store_path, upstream);
+                    println!("{} was a hit upstream: {}", store_path, upstream);
                     hit = true;
                     break;
                 }
             }
             if !hit {
                 trace!("sending {}", store_path);
-            cacheable_tx.send(store_path).unwrap();
+                tx.send(store_path).unwrap();
             }
+        });
     }
 }
 
 async fn uploader(cacheable_rx: mpsc::Receiver<String>) {
     let mut count = 0;
+    let concurrent = Semaphore::new(10);
+    let mut handles = Vec::new();
     loop {
         if let Ok(path_to_upload) = cacheable_rx.recv() {
-            trace!("to upload: {}", path_to_upload);
+            let _ = concurrent.acquire().await.unwrap();
+            handles.push(tokio::spawn(async move {
+                println!("uploading: {}", path_to_upload);
                 if Command::new("nix")
                     .arg("copy")
                     .arg("--to")
@@ -139,7 +156,12 @@ async fn uploader(cacheable_rx: mpsc::Receiver<String>) {
                 } else {
                     count += 1;
                 }
+            }));
         } else {
+            // make sure all threads are done
+            for handle in handles {
+                handle.await.unwrap();
+            }
             println!("uploaded {} paths", count);
             break;
         }