merge: 'issue-6' from email

Cargo.lock

@@ -1249,6 +1249,7 @@ dependencies = [
  "pretty_env_logger",
  "rayon",
  "reqwest",
+ "rlimit",
  "scraper",
  "serde",
  "serde_json",
@@ -1740,6 +1741,15 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "rlimit"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7043b63bd0cd1aaa628e476b80e6d4023a3b50eb32789f2728908107bd0c793a"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "roff"
 version = "0.2.2"

crates/nix-weather/Cargo.toml

@@ -28,6 +28,7 @@ openssl = { version = "0.10.63" }
 pretty_env_logger = "0.5.0"
 rayon = "1.9.0"
 reqwest = { version = "0.12", features = ["blocking"] }
+rlimit = "0.10.2"
 scraper = "0.18.1"
 serde = "1.0.197"
 serde_json = "1.0.114"

crates/nix-weather/src/main.rs

@@ -22,6 +22,9 @@ mod nix;
 /// The initial time to wait on http 104, in milliseconds
 const SLIDE: u64 = 100;
 
+// Open files limit to try to set
+const NOFILES_LIMIT: u64 = 16384;
+
 const DEFAULT_CACHE: &str = "cache.nixos.org";
 const DEFAULT_CONFIG_DIR: &str = "/etc/nixos";
 
@@ -109,6 +112,14 @@ async fn main() -> io::Result<()> {
 
   log::debug!("{:#?}", &ips);
 
+  // try to increase NOFILES runtime limit
+  if rlimit::increase_nofile_limit(NOFILES_LIMIT).is_err() {
+    log::warn!(
+      "Failed to increase NOFILES limit, still at {:#?}",
+      rlimit::Resource::NOFILE.get().unwrap_or_default()
+    );
+  }
+
   let domain_addr = SocketAddr::new(ips[0], 443);
 
   let client = reqwest::Client::builder()

crates/nix-weather/src/net.rs

@@ -63,10 +63,31 @@ pub async fn nar_exists(client: Client, domain: &str, hash: &str, slide: u64) ->
   match response {
     Ok(response) if response.status().as_u16() == 200 => 1,
     Ok(response) if response.status().as_u16() == 404 => 0,
-    _ => {
+    Ok(response) if response.status().as_u16() == 429 => {
       // We're so fast now we get rate limited.
-      //
-      // Writng an actual sliding window seems kinda hard,
+
+      let wait = if let Some(retry_after) = response.headers().get("Retry-After") {
+        if let Ok(seconds) = String::from_utf8_lossy(retry_after.as_bytes()).parse::<u64>() {
+          seconds * 1000
+        } else {
+          slide
+        }
+      } else {
+        slide
+      };
+
+      sleep(Duration::from_millis(wait)).await;
+      Box::pin(nar_exists(
+        client,
+        domain,
+        hash,
+        std::cmp::min(slide * 2, MAX_SLIDE),
+      ))
+      .await
+    }
+    _ => {
+      // We might also be out of sockets, so let's wait on that
+      // Writing an actual sliding window seems kinda hard,
       // so we do this instead.
       log::trace!("rate limited! {slide}");
       sleep(Duration::from_millis(slide)).await;