feat: introduce connection pool

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<Response> 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>, String, u32);
+
+fn client_cache() -> &'static Mutex<HashMap<ClientKey, Arc<Client>>> {
+    static CACHE: OnceLock<Mutex<HashMap<ClientKey, Arc<Client>>>> = OnceLock::new();
+    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
 pub fn build_client(
-    url: &str,
     proxy: Option<&str>,
     timeout: Option<f32>,
     user_agent: Option<&str>,
-) -> Result<Client, OoniError> {
-    let mut options = ClientOptions::new();
+) -> Result<Arc<Client>, 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<f32>,
     user_agent: Option<String>,
 ) -> Result<HttpResponse, OoniError> {
-    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<f32>,
     user_agent: Option<String>,
 ) -> Result<HttpResponse, OoniError> {
-    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);
 
@@ -267,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");
+    }
 }

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))

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)"
+    );
+}

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<AtomicUsize>,
+    accepts: Arc<AtomicUsize>,
     requests: Arc<Mutex<Vec<String>>>,
 }
 
@@ -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<String> {
         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,10 +112,90 @@ pub fn start_server_with_delay(body: &'static str, delay: Duration) -> MockServe
     MockServer {
         url: format!("http://{addr}"),
         hits,
+        accepts,
         requests,
     }
 }
 
 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<AtomicUsize>,
+    requests: Arc<Mutex<Vec<String>>>,
+) {
+    let mut buf = [0u8; 8192];
+    let mut data: Vec<u8> = 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);
+    }
+}

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<Client, Error> {
-        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));