Disable tcp_nodelay for reqwest::Client and add rate limiting for https requests (#458)

crates/binstalk/Cargo.toml

@@ -47,6 +47,7 @@ thiserror = "1.0.37"
 tinytemplate = "1.2.1"
 tokio = { version = "1.21.2", features = ["macros", "rt", "process", "sync", "signal"], default-features = false }
 toml_edit = { version = "0.14.4", features = ["easy"] }
+tower = { version = "0.4.13", features = ["limit", "util"] }
 trust-dns-resolver = { version = "0.21.2", optional = true, default-features = false, features = ["dnssec-ring"] }
 url = { version = "2.3.1", features = ["serde"] }
 xz2 = "0.1.7"

crates/binstalk/src/drivers/crates_io.rs

@@ -3,13 +3,14 @@ use std::path::PathBuf;
 use cargo_toml::Manifest;
 use crates_io_api::AsyncClient;
 use log::debug;
-use reqwest::Client;
 use semver::VersionReq;
-use url::Url;
 
 use crate::{
     errors::BinstallError,
-    helpers::download::Download,
+    helpers::{
+        download::Download,
+        remote::{Client, Url},
+    },
     manifests::cargo_toml_binstall::{Meta, TarBasedFmt},
 };
 

crates/binstalk/src/fetchers.rs

@@ -4,10 +4,10 @@ use compact_str::CompactString;
 pub use gh_crate_meta::*;
 pub use log::debug;
 pub use quickinstall::*;
-use reqwest::Client;
 
 use crate::{
     errors::BinstallError,
+    helpers::remote::Client,
     manifests::cargo_toml_binstall::{PkgFmt, PkgMeta},
 };
 

crates/binstalk/src/fetchers/gh_crate_meta.rs

@@ -4,7 +4,6 @@ use compact_str::{CompactString, ToCompactString};
 use futures_util::stream::{FuturesUnordered, StreamExt};
 use log::{debug, warn};
 use once_cell::sync::OnceCell;
-use reqwest::{Client, Method};
 use serde::Serialize;
 use strum::IntoEnumIterator;
 use tinytemplate::TinyTemplate;
@@ -14,7 +13,7 @@ use crate::{
     errors::BinstallError,
     helpers::{
         download::Download,
-        remote::{get_redirected_final_url, remote_exists},
+        remote::{Client, Method},
         tasks::AutoAbortJoinHandle,
     },
     manifests::cargo_toml_binstall::{PkgFmt, PkgMeta},
@@ -59,11 +58,9 @@ impl GhCrateMeta {
             AutoAbortJoinHandle::spawn(async move {
                 debug!("Checking for package at: '{url}'");
 
-                Ok(
-                    (remote_exists(client.clone(), url.clone(), Method::HEAD).await?
-                        || remote_exists(client, url.clone(), Method::GET).await?)
-                        .then_some((url, pkg_fmt)),
-                )
+                Ok((client.remote_exists(url.clone(), Method::HEAD).await?
+                    || client.remote_exists(url.clone(), Method::GET).await?)
+                    .then_some((url, pkg_fmt)))
             })
         })
     }
@@ -81,7 +78,11 @@ impl super::Fetcher for GhCrateMeta {
 
     async fn find(&self) -> Result<bool, BinstallError> {
         let repo = if let Some(repo) = self.data.repo.as_deref() {
-            Some(get_redirected_final_url(&self.client, Url::parse(repo)?).await?)
+            Some(
+                self.client
+                    .get_redirected_final_url(Url::parse(repo)?)
+                    .await?,
+            )
         } else {
             None
         };

crates/binstalk/src/fetchers/quickinstall.rs

@@ -2,14 +2,15 @@ use std::{path::Path, sync::Arc};
 
 use compact_str::CompactString;
 use log::debug;
-use reqwest::Client;
-use reqwest::Method;
 use tokio::task::JoinHandle;
 use url::Url;
 
 use crate::{
     errors::BinstallError,
-    helpers::{download::Download, remote::remote_exists},
+    helpers::{
+        download::Download,
+        remote::{Client, Method},
+    },
     manifests::cargo_toml_binstall::{PkgFmt, PkgMeta},
 };
 
@@ -43,7 +44,9 @@ impl super::Fetcher for QuickInstall {
         let url = self.package_url();
         self.report();
         debug!("Checking for package at: '{url}'");
-        remote_exists(self.client.clone(), Url::parse(&url)?, Method::HEAD).await
+        self.client
+            .remote_exists(Url::parse(&url)?, Method::HEAD)
+            .await
     }
 
     async fn fetch_and_extract(&self, dst: &Path) -> Result<(), BinstallError> {
@@ -112,15 +115,7 @@ impl QuickInstall {
             let url = Url::parse(&stats_url)?;
             debug!("Sending installation report to quickinstall ({url})");
 
-            client
-                .request(Method::HEAD, url.clone())
-                .send()
-                .await
-                .map_err(|err| BinstallError::Http {
-                    method: Method::HEAD,
-                    url,
-                    err,
-                })?;
+            client.remote_exists(url, Method::HEAD).await?;
 
             Ok(())
         })

crates/binstalk/src/helpers/download.rs

@@ -2,11 +2,10 @@ use std::{fmt::Debug, marker::PhantomData, path::Path};
 
 use digest::{Digest, FixedOutput, HashMarker, Output, OutputSizeUser, Update};
 use log::debug;
-use reqwest::{Client, Url};
 
 use crate::{
     errors::BinstallError,
-    helpers::remote::create_request,
+    helpers::remote::{Client, Url},
     manifests::cargo_toml_binstall::{PkgFmt, PkgFmtDecomposed, TarBasedFmt},
 };
 
@@ -48,7 +47,7 @@ impl Download {
         fmt: TarBasedFmt,
         visitor: V,
     ) -> Result<V::Target, BinstallError> {
-        let stream = create_request(self.client, self.url).await?;
+        let stream = self.client.create_request(self.url).await?;
 
         debug!("Downloading and extracting then in-memory processing");
 
@@ -65,7 +64,7 @@ impl Download {
         fmt: PkgFmt,
         path: impl AsRef<Path>,
     ) -> Result<(), BinstallError> {
-        let stream = create_request(self.client, self.url).await?;
+        let stream = self.client.create_request(self.url).await?;
 
         let path = path.as_ref();
         debug!("Downloading and extracting to: '{}'", path.display());

crates/binstalk/src/helpers/remote.rs

@@ -1,69 +1,109 @@
-use std::env;
+use std::{env, num::NonZeroU64, sync::Arc, time::Duration};
 
 use bytes::Bytes;
 use futures_util::stream::Stream;
 use log::debug;
-use reqwest::{tls, Client, ClientBuilder, Method, Response};
-use url::Url;
+use reqwest::{Request, Response};
+use tokio::sync::Mutex;
+use tower::{limit::rate::RateLimit, Service, ServiceBuilder, ServiceExt};
 
 use crate::errors::BinstallError;
 
-pub fn create_reqwest_client(min_tls: Option<tls::Version>) -> Result<Client, BinstallError> {
-    const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+pub use reqwest::{tls, Method};
+pub use url::Url;
 
-    let mut builder = ClientBuilder::new()
-        .user_agent(USER_AGENT)
-        .https_only(true)
-        .min_tls_version(tls::Version::TLS_1_2);
+#[derive(Clone, Debug)]
+pub struct Client {
+    client: reqwest::Client,
+    rate_limit: Arc<Mutex<RateLimit<reqwest::Client>>>,
+}
 
-    if let Some(ver) = min_tls {
-        builder = builder.min_tls_version(ver);
+impl Client {
+    /// * `per` - must not be 0.
+    pub fn new(
+        min_tls: Option<tls::Version>,
+        per: Duration,
+        num_request: NonZeroU64,
+    ) -> Result<Self, BinstallError> {
+        const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
+        let mut builder = reqwest::ClientBuilder::new()
+            .user_agent(USER_AGENT)
+            .https_only(true)
+            .min_tls_version(tls::Version::TLS_1_2)
+            .tcp_nodelay(false);
+
+        if let Some(ver) = min_tls {
+            builder = builder.min_tls_version(ver);
+        }
+
+        let client = builder.build()?;
+
+        Ok(Self {
+            client: client.clone(),
+            rate_limit: Arc::new(Mutex::new(
+                ServiceBuilder::new()
+                    .rate_limit(num_request.get(), per)
+                    .service(client),
+            )),
+        })
     }
 
-    Ok(builder.build()?)
-}
-
-pub async fn remote_exists(
-    client: Client,
-    url: Url,
-    method: Method,
-) -> Result<bool, BinstallError> {
-    let req = client
-        .request(method.clone(), url.clone())
-        .send()
-        .await
-        .map_err(|err| BinstallError::Http { method, url, err })?;
-    Ok(req.status().is_success())
-}
-
-pub async fn get_redirected_final_url(client: &Client, url: Url) -> Result<Url, BinstallError> {
-    let method = Method::HEAD;
-
-    let req = client
-        .request(method.clone(), url.clone())
-        .send()
-        .await
-        .and_then(Response::error_for_status)
-        .map_err(|err| BinstallError::Http { method, url, err })?;
-
-    Ok(req.url().clone())
-}
-
-pub(crate) async fn create_request(
-    client: Client,
-    url: Url,
-) -> Result<impl Stream<Item = reqwest::Result<Bytes>>, BinstallError> {
-    debug!("Downloading from: '{url}'");
-
-    client
-        .get(url.clone())
-        .send()
-        .await
-        .and_then(|r| r.error_for_status())
-        .map_err(|err| BinstallError::Http {
-            method: Method::GET,
-            url,
-            err,
-        })
-        .map(Response::bytes_stream)
+    pub fn get_inner(&self) -> &reqwest::Client {
+        &self.client
+    }
+
+    async fn send_request(
+        &self,
+        method: Method,
+        url: Url,
+        error_for_status: bool,
+    ) -> Result<Response, BinstallError> {
+        let request = Request::new(method.clone(), url.clone());
+
+        // Reduce critical section:
+        //  - Construct the request before locking
+        //  - Once the rate_limit is ready, call it and obtain
+        //    the future, then release the lock before
+        //    polling the future.
+        let future = self.rate_limit.lock().await.ready().await?.call(request);
+
+        future
+            .await
+            .and_then(|response| {
+                if error_for_status {
+                    response.error_for_status()
+                } else {
+                    Ok(response)
+                }
+            })
+            .map_err(|err| BinstallError::Http { method, url, err })
+    }
+
+    pub async fn remote_exists(&self, url: Url, method: Method) -> Result<bool, BinstallError> {
+        Ok(self
+            .send_request(method, url, false)
+            .await?
+            .status()
+            .is_success())
+    }
+
+    pub async fn get_redirected_final_url(&self, url: Url) -> Result<Url, BinstallError> {
+        Ok(self
+            .send_request(Method::HEAD, url, true)
+            .await?
+            .url()
+            .clone())
+    }
+
+    pub(crate) async fn create_request(
+        &self,
+        url: Url,
+    ) -> Result<impl Stream<Item = reqwest::Result<Bytes>>, BinstallError> {
+        debug!("Downloading from: '{url}'");
+
+        self.send_request(Method::GET, url, true)
+            .await
+            .map(Response::bytes_stream)
+    }
 }

crates/binstalk/src/ops/resolve.rs

@@ -10,7 +10,6 @@ use cargo_toml::{Manifest, Package, Product};
 use compact_str::{CompactString, ToCompactString};
 use itertools::Itertools;
 use log::{debug, info, warn};
-use reqwest::Client;
 use semver::{Version, VersionReq};
 use tokio::task::block_in_place;
 
@@ -20,7 +19,7 @@ use crate::{
     drivers::fetch_crate_cratesio,
     errors::BinstallError,
     fetchers::{Data, Fetcher, GhCrateMeta, QuickInstall},
-    helpers::tasks::AutoAbortJoinHandle,
+    helpers::{remote::Client, tasks::AutoAbortJoinHandle},
     manifests::cargo_toml_binstall::{Meta, PkgMeta},
 };