From 61ffa84f6cce362e0a83cd00f55bcda0a6b3df12 Mon Sep 17 00:00:00 2001 From: decfox Date: Tue, 21 Jul 2026 13:05:36 +0530 Subject: [PATCH 1/4] feat: introduce client cache with reuse abilities --- ooniprobe-ffi/src/client.rs | 57 ++++++++++++++----- ooniprobe-ffi/src/userauth.rs | 4 +- ooniprobe-services/src/client/reqwest_impl.rs | 8 ++- ooniprobe-services/src/client/wreq_impl.rs | 7 ++- 4 files changed, 57 insertions(+), 19 deletions(-) diff --git a/ooniprobe-ffi/src/client.rs b/ooniprobe-ffi/src/client.rs index c4309fd..c8b82dc 100644 --- a/ooniprobe-ffi/src/client.rs +++ b/ooniprobe-ffi/src/client.rs @@ -1,5 +1,8 @@ use crate::errors::OoniError; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; + use ooniprobe_services::client::{Client, ClientOptions, Response}; use serde::{Deserialize, Serialize}; @@ -51,23 +54,47 @@ impl From for HttpResponse { const DEFAULT_TIMEOUT_SECS: f32 = 30.0; const DEFAULT_USER_AGENT: &str = "ooniprobe"; +// The connection pool lives inside the client, so clients are cached and reused. +// Keyed by the options baked in at build time; +type ClientKey = (Option, String, u32); + +fn client_cache() -> &'static Mutex>> { + static CACHE: OnceLock>>> = OnceLock::new(); + CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + pub fn build_client( - url: &str, proxy: Option<&str>, timeout: Option, user_agent: Option<&str>, -) -> Result { - let mut options = ClientOptions::new(); +) -> Result, OoniError> { + let timeout = timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); + let user_agent = user_agent.unwrap_or(DEFAULT_USER_AGENT); + let key = ( + proxy.map(String::from), + user_agent.to_string(), + timeout.to_bits(), + ); + + let mut cache = client_cache().lock().unwrap_or_else(|e| e.into_inner()); + if let Some(client) = cache.get(&key) { + return Ok(Arc::clone(client)); + } + let mut options = ClientOptions::new(); options.set_proxy_url(proxy); - options.set_base_url(Some(url)); - options.set_timeout(Some(timeout.unwrap_or(DEFAULT_TIMEOUT_SECS))); - options.set_user_agent(Some(user_agent.unwrap_or(DEFAULT_USER_AGENT))); - - Client::builder() - .set_options(options) - .build() - .map_err(OoniError::from) + options.set_timeout(Some(timeout)); + options.set_user_agent(Some(user_agent)); + + let client = Arc::new( + Client::builder() + .set_options(options) + .build() + .map_err(OoniError::from)?, + ); + cache.insert(key, Arc::clone(&client)); + + Ok(client) } pub fn client_get( @@ -78,7 +105,7 @@ pub fn client_get( timeout: Option, user_agent: Option, ) -> Result { - let client = build_client(&url, proxy.as_deref(), timeout, user_agent.as_deref())?; + let client = build_client(proxy.as_deref(), timeout, user_agent.as_deref())?; let mut header_map = HeaderMap::new(); for kv in headers { @@ -111,7 +138,7 @@ pub fn client_post( timeout: Option, user_agent: Option, ) -> Result { - let client = build_client(&url, proxy.as_deref(), timeout, user_agent.as_deref())?; + let client = build_client(proxy.as_deref(), timeout, user_agent.as_deref())?; let mut header_map = HeaderMap::new(); for kv in headers { @@ -143,8 +170,8 @@ mod tests { #[test] fn get_manifest_returns_manifest_version_and_public_params() { let url = format!("{BASE_URL}/api/v1/manifest"); - let resp = client_get(url, vec![], vec![], None, None, None) - .expect("GET manifest should succeed"); + let resp = + client_get(url, vec![], vec![], None, None, None).expect("GET manifest should succeed"); assert_eq!(resp.status_code, 200, "incorrect status_code: {:?}", resp); diff --git a/ooniprobe-ffi/src/userauth.rs b/ooniprobe-ffi/src/userauth.rs index 577af8e..e26f994 100644 --- a/ooniprobe-ffi/src/userauth.rs +++ b/ooniprobe-ffi/src/userauth.rs @@ -141,7 +141,7 @@ pub fn userauth_register( let json_payload = serde_json::to_string(&payload)?; // make the API call - let client = build_client(&url, proxy.as_deref(), timeout, user_agent.as_deref())?; + let client = build_client(proxy.as_deref(), timeout, user_agent.as_deref())?; let request = client .request("POST", &url) .map(|b| b.body(json_payload)) @@ -254,7 +254,7 @@ pub fn userauth_submit( let json_payload = serde_json::to_string(&submit_payload)?; // make the API call - let client = build_client(&url, proxy.as_deref(), timeout, user_agent.as_deref())?; + let client = build_client(proxy.as_deref(), timeout, user_agent.as_deref())?; let request = client .request("POST", &url) .map(|b| b.body(json_payload)) diff --git a/ooniprobe-services/src/client/reqwest_impl.rs b/ooniprobe-services/src/client/reqwest_impl.rs index fe534be..a1b36cf 100644 --- a/ooniprobe-services/src/client/reqwest_impl.rs +++ b/ooniprobe-services/src/client/reqwest_impl.rs @@ -6,6 +6,9 @@ use tokio::runtime::Runtime; use super::{b64_encode, ClientOptions, Error, Response}; +const POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(30); +const POOL_MAX_IDLE_PER_HOST: usize = 4; + pub struct Client { http_client: reqwest::Client, rt: Runtime, @@ -117,7 +120,10 @@ impl ClientBuilder { } pub fn build(self) -> Result { - let mut client_builder = reqwest::Client::builder().use_rustls_tls(); + let mut client_builder = reqwest::Client::builder() + .use_rustls_tls() + .pool_idle_timeout(POOL_IDLE_TIMEOUT) + .pool_max_idle_per_host(POOL_MAX_IDLE_PER_HOST); if let Some(timeout) = self.client_options.timeout { client_builder = client_builder.timeout(Duration::from_secs_f32(timeout)); diff --git a/ooniprobe-services/src/client/wreq_impl.rs b/ooniprobe-services/src/client/wreq_impl.rs index 061764f..9326799 100644 --- a/ooniprobe-services/src/client/wreq_impl.rs +++ b/ooniprobe-services/src/client/wreq_impl.rs @@ -9,6 +9,9 @@ use wreq_util::Emulation; use super::{b64_encode, ClientOptions, Error, Response}; +const POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(30); +const POOL_MAX_IDLE_PER_HOST: usize = 4; + pub struct Client { inner: Arc, rt: Runtime, @@ -128,7 +131,9 @@ impl ClientBuilder { .cert_store(CertStore::from_der_certs( webpki_root_certs::TLS_SERVER_ROOT_CERTS, )?) - .emulation(Emulation::Chrome118); + .emulation(Emulation::Chrome118) + .pool_idle_timeout(POOL_IDLE_TIMEOUT) + .pool_max_idle_per_host(POOL_MAX_IDLE_PER_HOST); if let Some(timeout) = self.client_options.timeout { client_builder = client_builder.timeout(Duration::from_secs_f32(timeout)); From 1c0813d88a9ce54111a12244427e2282b11d3f68 Mon Sep 17 00:00:00 2001 From: decfox Date: Tue, 21 Jul 2026 14:16:43 +0530 Subject: [PATCH 2/4] chore: remove redundant Arc encapsulation on wreq client --- ooniprobe-services/src/client/wreq_impl.rs | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/ooniprobe-services/src/client/wreq_impl.rs b/ooniprobe-services/src/client/wreq_impl.rs index 9326799..25fdbd8 100644 --- a/ooniprobe-services/src/client/wreq_impl.rs +++ b/ooniprobe-services/src/client/wreq_impl.rs @@ -1,7 +1,6 @@ use bytes::Bytes; use encoding_rs::{Encoding, UTF_8}; use mime::Mime; -use std::sync::Arc; use std::time::Duration; use tokio::runtime::Runtime; use wreq::tls::CertStore; @@ -13,12 +12,8 @@ const POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(30); const POOL_MAX_IDLE_PER_HOST: usize = 4; pub struct Client { - inner: Arc, - rt: Runtime, -} - -struct ClientRef { http_client: wreq::Client, + rt: Runtime, } fn decode_to_text(bytes: &Bytes, headers: &wreq::header::HeaderMap) -> Result { @@ -46,7 +41,7 @@ impl Client { pub fn execute(&self, request: wreq::Request) -> Result { self.rt.block_on(async { - let wreq_resp: wreq::Response = self.inner.http_client.execute(request).await?; + let wreq_resp: wreq::Response = self.http_client.execute(request).await?; let status_code = wreq_resp.status().as_u16(); let version = match wreq_resp.version() { @@ -101,7 +96,7 @@ impl Client { "OPTIONS" => http::Method::OPTIONS, _ => return Err(Error::InvalidHttpMethod), }; - Ok(self.inner.http_client.request(m, url)) + Ok(self.http_client.request(m, url)) } } @@ -159,7 +154,7 @@ impl ClientBuilder { .expect("failed to build tokio runtime"); Ok(Client { - inner: Arc::new(ClientRef { http_client }), + http_client, rt, }) } From dd5a60da486e128f89435520cd2e7821ba0889a8 Mon Sep 17 00:00:00 2001 From: decfox Date: Thu, 30 Jul 2026 21:43:21 +0530 Subject: [PATCH 3/4] feat: add intergation tests for connection pool --- ooniprobe-ffi/tests/client.rs | 52 +++++++++++++++++- ooniprobe-ffi/tests/common/mod.rs | 90 +++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/ooniprobe-ffi/tests/client.rs b/ooniprobe-ffi/tests/client.rs index e01ba1e..7f64c51 100644 --- a/ooniprobe-ffi/tests/client.rs +++ b/ooniprobe-ffi/tests/client.rs @@ -4,7 +4,7 @@ mod common; use std::time::Duration; -use common::{start_server, start_server_with_delay}; +use common::{start_keepalive_server, start_server, start_server_with_delay}; use uniffi_ooniprobe::{client_get, client_post, KeyValue, OoniError}; fn kv(key: &str, value: &str) -> KeyValue { @@ -154,3 +154,53 @@ fn generous_timeout_tolerates_slow_response() { .expect("response within timeout should succeed"); assert_eq!(resp.status_code, 200); } + +#[test] +fn reuses_connection_across_calls() { + let server = start_keepalive_server("ok"); + let ua = Some("reuse-across-calls/1".to_string()); + + for _ in 0..3 { + client_get( + format!("{}/", server.url), + vec![], + vec![], + None, + None, + ua.clone(), + ) + .expect("GET should succeed"); + } + + assert_eq!(server.hits(), 3, "expected three requests"); + assert_eq!( + server.accepts(), + 1, + "all three requests should reuse one connection, got {} accepts", + server.accepts() + ); +} + +#[test] +fn distinct_user_agent_uses_distinct_client() { + let server = start_keepalive_server("ok"); + + for ua in ["distinct-a/1", "distinct-b/1"] { + client_get( + format!("{}/", server.url), + vec![], + vec![], + None, + None, + Some(ua.to_string()), + ) + .expect("GET should succeed"); + } + + assert_eq!(server.hits(), 2, "expected two requests"); + assert_eq!( + server.accepts(), + 2, + "differing user agents must use separate clients (separate pools)" + ); +} diff --git a/ooniprobe-ffi/tests/common/mod.rs b/ooniprobe-ffi/tests/common/mod.rs index f85e46c..67f4fde 100644 --- a/ooniprobe-ffi/tests/common/mod.rs +++ b/ooniprobe-ffi/tests/common/mod.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use std::io::{Read, Write}; use std::net::TcpListener; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -8,6 +10,7 @@ use std::time::Duration; pub struct MockServer { pub url: String, hits: Arc, + accepts: Arc, requests: Arc>>, } @@ -16,6 +19,10 @@ impl MockServer { self.hits.load(Ordering::SeqCst) } + pub fn accepts(&self) -> usize { + self.accepts.load(Ordering::SeqCst) + } + pub fn requests(&self) -> Vec { self.requests.lock().unwrap().clone() } @@ -45,13 +52,16 @@ pub fn start_server_with_delay(body: &'static str, delay: Duration) -> MockServe let addr = listener.local_addr().expect("local addr"); let hits = Arc::new(AtomicUsize::new(0)); + let accepts = Arc::new(AtomicUsize::new(0)); let requests = Arc::new(Mutex::new(Vec::new())); let hits_t = Arc::clone(&hits); + let accepts_t = Arc::clone(&accepts); let requests_t = Arc::clone(&requests); thread::spawn(move || { for stream in listener.incoming() { let Ok(mut stream) = stream else { continue }; + accepts_t.fetch_add(1, Ordering::SeqCst); let mut data = Vec::new(); let mut buf = [0u8; 8192]; @@ -102,6 +112,7 @@ pub fn start_server_with_delay(body: &'static str, delay: Duration) -> MockServe MockServer { url: format!("http://{addr}"), hits, + accepts, requests, } } @@ -109,3 +120,82 @@ pub fn start_server_with_delay(body: &'static str, delay: Duration) -> MockServe pub fn start_server(body: &'static str) -> MockServer { start_server_with_delay(body, Duration::ZERO) } + +pub fn start_keepalive_server(body: &'static str) -> MockServer { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock server"); + let addr = listener.local_addr().expect("local addr"); + + let hits = Arc::new(AtomicUsize::new(0)); + let accepts = Arc::new(AtomicUsize::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let hits_t = Arc::clone(&hits); + let accepts_t = Arc::clone(&accepts); + let requests_t = Arc::clone(&requests); + + thread::spawn(move || { + for stream in listener.incoming() { + let Ok(stream) = stream else { continue }; + accepts_t.fetch_add(1, Ordering::SeqCst); + + let hits = Arc::clone(&hits_t); + let requests = Arc::clone(&requests_t); + thread::spawn(move || serve_keepalive(stream, body, hits, requests)); + } + }); + + MockServer { + url: format!("http://{addr}"), + hits, + accepts, + requests, + } +} + +fn serve_keepalive( + mut stream: std::net::TcpStream, + body: &'static str, + hits: Arc, + requests: Arc>>, +) { + let mut buf = [0u8; 8192]; + let mut data: Vec = Vec::new(); + + loop { + let head_end = loop { + if let Some(p) = data.windows(4).position(|w| w == b"\r\n\r\n") { + break p + 4; + } + match stream.read(&mut buf) { + Ok(0) | Err(_) => return, + Ok(n) => data.extend_from_slice(&buf[..n]), + } + }; + + let headers = String::from_utf8_lossy(&data[..head_end]).to_string(); + let total = head_end + content_length(&headers); + while data.len() < total { + match stream.read(&mut buf) { + Ok(0) | Err(_) => return, + Ok(n) => data.extend_from_slice(&buf[..n]), + } + } + + requests + .lock() + .unwrap() + .push(String::from_utf8_lossy(&data[..total]).to_string()); + hits.fetch_add(1, Ordering::SeqCst); + + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ); + if stream.write_all(response.as_bytes()).is_err() { + return; + } + let _ = stream.flush(); + + data.drain(..total); + } +} From 15543d6c778741973f212e3943264af3b9c27fef Mon Sep 17 00:00:00 2001 From: decfox Date: Thu, 30 Jul 2026 22:01:51 +0530 Subject: [PATCH 4/4] chore: add unit tests for connection map --- ooniprobe-ffi/src/client.rs | 42 +++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/ooniprobe-ffi/src/client.rs b/ooniprobe-ffi/src/client.rs index c8b82dc..c539bd2 100644 --- a/ooniprobe-ffi/src/client.rs +++ b/ooniprobe-ffi/src/client.rs @@ -294,4 +294,46 @@ mod tests { body_text ); } + + #[test] + fn cache_same_key_returns_cached_client() { + let a = build_client(None, None, Some("cache-same/1")).unwrap(); + let b = build_client(None, None, Some("cache-same/1")).unwrap(); + assert!(Arc::ptr_eq(&a, &b), "same key must return the cached client"); + } + + #[test] + fn cache_distinct_user_agent_distinct_client() { + let a = build_client(None, None, Some("cache-ua-a/1")).unwrap(); + let b = build_client(None, None, Some("cache-ua-b/1")).unwrap(); + assert!(!Arc::ptr_eq(&a, &b), "different user agents must not share a client"); + } + + #[test] + fn cache_distinct_timeout_distinct_client() { + let a = build_client(None, Some(5.0), Some("cache-timeout/1")).unwrap(); + let b = build_client(None, Some(9.0), Some("cache-timeout/1")).unwrap(); + assert!(!Arc::ptr_eq(&a, &b), "different timeouts must not share a client"); + } + + #[test] + fn cache_distinct_proxy_distinct_client() { + let a = build_client(Some("http://127.0.0.1:8080"), None, Some("cache-proxy/1")).unwrap(); + let b = build_client(Some("http://127.0.0.1:9090"), None, Some("cache-proxy/1")).unwrap(); + assert!(!Arc::ptr_eq(&a, &b), "different proxies must not share a client"); + } + + #[test] + fn cache_none_timeout_normalizes_to_default() { + let a = build_client(None, None, Some("cache-norm-timeout/1")).unwrap(); + let b = build_client(None, Some(DEFAULT_TIMEOUT_SECS), Some("cache-norm-timeout/1")).unwrap(); + assert!(Arc::ptr_eq(&a, &b), "None timeout must map to the same key as the default"); + } + + #[test] + fn cache_none_user_agent_normalizes_to_default() { + let a = build_client(None, None, None).unwrap(); + let b = build_client(None, None, Some(DEFAULT_USER_AGENT)).unwrap(); + assert!(Arc::ptr_eq(&a, &b), "None user agent must map to the same key as the default"); + } }