Merge pull request #14 from dfinity/igornovg/improvements

add inflight req metric, cleanups, update some deps

Author: blind-oracle

.github/workflows/test.yml

@@ -16,7 +16,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - uses: Swatinem/rust-cache@23bce251a8cd2ffc3c1075eaa2367cf899916d84
+      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
 
       - name: Run tests
         run: cargo test

Cargo.toml

@@ -25,25 +25,26 @@ derive-new = "0.7.0"
 fqdn = "0.4.1"
 futures = "0.3.31"
 futures-util = "0.3.31"
-hickory-proto = "0.24.1"
-hickory-resolver = { version = "0.24.1", features = [
-    "dns-over-https-rustls",
-    "webpki-roots",
+hickory-proto = "0.25.1"
+hickory-resolver = { version = "0.25.1", features = [
+    "tls-ring",
+    "https-ring",
     "dnssec-ring",
+    "webpki-roots",
 ] }
 http = "1.3.1"
 http-body = "1.0.1"
 http-body-util = "0.1.2"
 humantime = "2.2.0"
-hyper = "1.5.0"
+hyper = "1.6.0"
 hyper-util = { version = "0.1.10", features = ["full"] }
 instant-acme = { version = "0.7.2", default-features = false, features = [
     "ring",
     "hyper-rustls",
 ] }
 moka = { version = "0.12.8", features = ["sync", "future"] }
 parse-size = { version = "1.1.0", features = ["std"] }
-prometheus = "0.13.4"
+prometheus = "0.14.0"
 prost = "0.13.3"
 prost-types = "0.13.3"
 rand = "0.8.5"
@@ -86,8 +87,8 @@ tower = { version = "0.5.1", features = ["util"] }
 tower-service = "0.3.3"
 tracing = "0.1.40"
 url = "2.5.3"
-uuid = { version = "1.15.0", features = ["v7"] }
-vrl = { version = "0.22.0", default-features = false, features = ["value"] }
+uuid = { version = "1.16.0", features = ["v7"] }
+vrl = { version = "0.23.0", default-features = false, features = ["value"] }
 webpki-root-certs = "0.26.6"
 x509-parser = "0.17.0"
 zeroize = { version = "1.8.1", features = ["derive"] }

src/http/client/mod.rs

@@ -3,17 +3,16 @@ pub mod cli;
 use std::{
     fmt,
     sync::{
+        Arc,
         atomic::{AtomicUsize, Ordering},
-        Arc, RwLock, RwLockWriteGuard,
     },
-    time::{Duration, Instant},
+    time::Duration,
 };
 
 use anyhow::Context;
 use async_trait::async_trait;
 use http::header::HeaderValue;
-use rand::{rngs::OsRng, seq::IteratorRandom};
-use reqwest::{dns::Resolve, Request, Response};
+use reqwest::{Request, Response, dns::Resolve};
 use scopeguard::defer;
 
 use super::Error;
@@ -160,6 +159,7 @@ impl ReqwestClientLeastLoaded {
 #[async_trait]
 impl Client for ReqwestClientLeastLoaded {
     async fn execute(&self, req: Request) -> Result<Response, reqwest::Error> {
+        // Select the client with least outstanding requests
         let cli = self
             .inner
             .iter()
@@ -176,14 +176,6 @@ impl Client for ReqwestClientLeastLoaded {
     }
 }
 
-pub trait GeneratesClients: Send + Sync + fmt::Debug + 'static {
-    fn generate(&self) -> Result<Arc<dyn Client>, Error>;
-}
-
-pub trait GeneratesClientsWithStats: Send + Sync + fmt::Debug + 'static {
-    fn generate(&self) -> Result<Arc<dyn ClientWithStats>, Error>;
-}
-
 #[derive(Debug, Clone)]
 pub struct ClientStats {
     pub pool_size: usize,
@@ -194,138 +186,6 @@ pub trait Stats {
     fn stats(&self) -> ClientStats;
 }
 
-#[derive(Debug)]
-pub struct ReqwestClientDynamic<G: GeneratesClients> {
-    generator: G,
-    min_clients: usize,
-    max_clients: usize,
-    max_outstanding: usize,
-    idle_timeout: Duration,
-    pool: RwLock<Vec<Arc<ReqwestClientDynamicInner>>>,
-}
-
-impl<G: GeneratesClients> ClientWithStats for ReqwestClientDynamic<G> {
-    fn to_client(self: Arc<Self>) -> Arc<dyn Client> {
-        self
-    }
-}
-
-#[derive(Debug)]
-struct ReqwestClientDynamicInner {
-    cli: Arc<dyn Client>,
-    outstanding: AtomicUsize,
-    last_request: RwLock<Instant>,
-}
-
-impl ReqwestClientDynamicInner {
-    fn new(cli: Arc<dyn Client>) -> Self {
-        Self {
-            cli,
-            outstanding: AtomicUsize::new(0),
-            last_request: RwLock::new(Instant::now()),
-        }
-    }
-}
-
-impl<G: GeneratesClients> ReqwestClientDynamic<G> {
-    pub fn new(
-        generator: G,
-        min_clients: usize,
-        max_clients: usize,
-        max_outstanding: usize,
-        idle_timeout: Duration,
-    ) -> Result<Self, Error> {
-        let mut pool = Vec::with_capacity(max_clients);
-
-        for _ in 0..min_clients {
-            let inner = Arc::new(ReqwestClientDynamicInner::new(generator.generate()?));
-            pool.push(inner);
-        }
-
-        Ok(Self {
-            generator,
-            min_clients,
-            max_clients,
-            max_outstanding,
-            idle_timeout,
-            pool: RwLock::new(pool),
-        })
-    }
-
-    /// Drop unused clients while leaving min_clients always available.
-    /// Algo mimics Vec::retain().
-    /// TODO find a better way?
-    fn cleanup(&self, pool: &mut RwLockWriteGuard<'_, Vec<Arc<ReqwestClientDynamicInner>>>) {
-        let mut j = self.min_clients;
-        for i in self.min_clients..pool.len() {
-            if !(pool[i].outstanding.load(Ordering::SeqCst) == 0
-                && pool[i].last_request.read().unwrap().elapsed() > self.idle_timeout)
-            {
-                pool.swap(i, j);
-                j += 1
-            }
-        }
-        pool.truncate(j);
-    }
-
-    fn get_client(&self) -> Arc<ReqwestClientDynamicInner> {
-        let mut pool = self.pool.write().unwrap();
-        self.cleanup(&mut pool);
-
-        pool.iter()
-            // First try to find an existing client with spare capacity
-            .find_map(|x| {
-                (x.outstanding.load(Ordering::SeqCst) < self.max_outstanding).then(|| x.clone())
-            })
-            // Otherwise see if we have spare space in the pool
-            .unwrap_or_else(|| {
-                // If not - just pick a random client
-                if pool.len() >= self.max_clients {
-                    pool.iter().choose(&mut OsRng).unwrap().clone()
-                } else {
-                    // Otherwise generate a new client and use it
-                    // The error is checked only in new() for now.
-                    let cli = self.generator.generate().unwrap();
-                    let inner = Arc::new(ReqwestClientDynamicInner::new(cli));
-                    pool.push(inner.clone());
-                    inner
-                }
-            })
-    }
-}
-
-impl<G: GeneratesClients> Stats for ReqwestClientDynamic<G> {
-    fn stats(&self) -> ClientStats {
-        let pool = self.pool.read().unwrap();
-
-        let outstanding: usize = pool
-            .iter()
-            .map(|x| x.outstanding.load(Ordering::SeqCst))
-            .sum();
-
-        ClientStats {
-            pool_size: pool.len(),
-            outstanding,
-        }
-    }
-}
-
-#[async_trait]
-impl<G: GeneratesClients> Client for ReqwestClientDynamic<G> {
-    async fn execute(&self, req: Request) -> Result<Response, reqwest::Error> {
-        let inner = self.get_client();
-
-        // The future can be cancelled so we have to use defer to make sure the counter is decreased
-        defer! {
-            inner.outstanding.fetch_sub(1, Ordering::SeqCst);
-        }
-
-        *inner.last_request.write().unwrap() = Instant::now();
-        inner.outstanding.fetch_add(1, Ordering::SeqCst);
-        inner.cli.execute(req).await
-    }
-}
-
 pub fn basic_auth<U, P>(username: U, password: Option<P>) -> HeaderValue
 where
     U: fmt::Display,
@@ -348,56 +208,3 @@ where
     header.set_sensitive(true);
     header
 }
-
-#[cfg(test)]
-mod test {
-    use futures::future::join_all;
-
-    use super::*;
-
-    #[derive(Debug)]
-    struct TestClient;
-
-    #[async_trait]
-    impl Client for TestClient {
-        async fn execute(&self, _req: Request) -> Result<Response, reqwest::Error> {
-            let resp = http::Response::new(vec![]);
-            tokio::time::sleep(Duration::from_millis(100)).await;
-            Ok(resp.into())
-        }
-    }
-
-    #[derive(Debug)]
-    struct TestClientGenerator;
-    impl GeneratesClients for TestClientGenerator {
-        fn generate(&self) -> Result<Arc<dyn Client>, Error> {
-            Ok(Arc::new(TestClient))
-        }
-    }
-
-    #[tokio::test]
-    async fn test_dynamic_client() {
-        let cli = Arc::new(
-            ReqwestClientDynamic::new(TestClientGenerator, 1, 10, 10, Duration::ZERO).unwrap(),
-        );
-
-        let mut futs = vec![];
-        for _ in 0..200 {
-            let req = Request::new(reqwest::Method::GET, url::Url::parse("http://foo").unwrap());
-
-            let cli = cli.clone();
-            futs.push(async move { cli.execute(req).await });
-        }
-
-        join_all(futs).await;
-        let mut pool = cli.pool.write().unwrap();
-        assert_eq!(pool.len(), 10);
-
-        for x in pool.iter() {
-            assert_eq!(x.outstanding.load(Ordering::SeqCst), 0);
-        }
-
-        cli.cleanup(&mut pool);
-        assert_eq!(pool.len(), 1);
-    }
-}

src/http/dns.rs

@@ -8,14 +8,15 @@ use anyhow::Context;
 use async_trait::async_trait;
 use hickory_proto::rr::RecordType;
 use hickory_resolver::{
-    config::{NameServerConfigGroup, ResolverConfig, ResolverOpts},
+    TokioResolver,
+    config::{NameServerConfigGroup, ResolveHosts, ResolverConfig, ResolverOpts},
     lookup_ip::LookupIpIntoIter,
-    TokioAsyncResolver,
+    name_server::TokioConnectionProvider,
 };
 use reqwest::dns::{Addrs, Name, Resolve, Resolving};
 use strum_macros::EnumString;
 
-use super::{client::CloneableDnsResolver, Error};
+use super::{Error, client::CloneableDnsResolver};
 
 #[derive(Clone, Copy, Debug, EnumString)]
 #[strum(serialize_all = "snake_case")]
@@ -39,7 +40,7 @@ pub struct Options {
 }
 
 #[derive(Debug, Clone)]
-pub struct Resolver(Arc<TokioAsyncResolver>);
+pub struct Resolver(Arc<TokioResolver>);
 impl CloneableDnsResolver for Resolver {}
 
 // new() must be called in Tokio context
@@ -56,14 +57,16 @@ impl Resolver {
         let cfg = ResolverConfig::from_parts(None, vec![], name_servers);
 
         let mut opts = ResolverOpts::default();
-        opts.rotate = true;
         opts.cache_size = o.cache_size;
-        opts.use_hosts_file = false;
+        opts.use_hosts_file = ResolveHosts::Never;
         opts.preserve_intermediates = false;
         opts.try_tcp_on_error = true;
 
-        let resolver = TokioAsyncResolver::tokio(cfg, opts);
-        Self(Arc::new(resolver))
+        let mut builder =
+            TokioResolver::builder_with_config(cfg, TokioConnectionProvider::default());
+        *builder.options_mut() = opts;
+
+        Self(Arc::new(builder.build()))
     }
 }