diff --git a/Cargo.lock b/Cargo.lock index 8a0ab98..9ccd545 100644 --- a/Cargo.lock +++ b/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" diff --git a/crates/nix-weather/Cargo.toml b/crates/nix-weather/Cargo.toml index 325ef70..86a9382 100644 --- a/crates/nix-weather/Cargo.toml +++ b/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" diff --git a/crates/nix-weather/src/main.rs b/crates/nix-weather/src/main.rs index ff850e9..97e9381 100644 --- a/crates/nix-weather/src/main.rs +++ b/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() diff --git a/crates/nix-weather/src/net.rs b/crates/nix-weather/src/net.rs index b7616aa..b01e5dd 100644 --- a/crates/nix-weather/src/net.rs +++ b/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::() { + 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;