millipede_http/
coalesce.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    sync::{Arc, Mutex},
5};
6
7use anyhow::anyhow;
8use async_trait::async_trait;
9use millipede_core::{
10    http_client::{HttpClient, HttpClientError, HttpRequest, HttpResponse, StreamingResponse},
11    request::{Method, Request},
12};
13use tokio::sync::OnceCell;
14
15#[derive(Hash, Eq, PartialEq, Clone)]
16struct CoalesceKey {
17    unique_key: String,
18    proxy: Option<String>,
19    jar: Option<usize>,
20}
21
22type SharedResponse = Arc<OnceCell<Result<HttpResponse, String>>>;
23type InFlightRequests = HashMap<CoalesceKey, SharedResponse>;
24
25struct InFlightGuard<'a> {
26    in_flight: &'a Mutex<InFlightRequests>,
27    key: CoalesceKey,
28    cell: SharedResponse,
29    active: bool,
30}
31
32impl Drop for InFlightGuard<'_> {
33    fn drop(&mut self) {
34        // `OnceCell` cancels an initializer when its caller is dropped and lets
35        // another waiter take over. Keep the entry discoverable during that
36        // handoff; only the caller that observes a settled cell may remove it.
37        if !self.active || self.cell.get().is_none() {
38            return;
39        }
40        let mut in_flight = self
41            .in_flight
42            .lock()
43            .unwrap_or_else(|error| error.into_inner());
44        if in_flight
45            .get(&self.key)
46            .is_some_and(|current| Arc::ptr_eq(current, &self.cell))
47        {
48            in_flight.remove(&self.key);
49        }
50    }
51}
52
53/// An HTTP client decorator that joins safe, identical requests already in flight.
54///
55/// Only bodyless `GET` and `HEAD` requests are coalesced. Cookie-jar identity and
56/// proxy selection are part of the key, so distinct sessions and routes never share
57/// a fetch. Joined callers of a failed request receive an `Other` error because
58/// [`HttpClientError`] is not cloneable; this intentionally loses connect/timeout
59/// variant fidelity.
60///
61/// # Examples
62///
63/// ```
64/// use std::sync::Arc;
65/// use millipede_http::{CoalescingClient, ReqwestClient};
66///
67/// let inner = Arc::new(ReqwestClient::new()?);
68/// let client = CoalescingClient::new(inner);
69/// # Ok::<(), millipede_core::http_client::HttpClientError>(())
70/// ```
71pub struct CoalescingClient {
72    inner: Arc<dyn HttpClient>,
73    in_flight: Mutex<InFlightRequests>,
74}
75
76impl CoalescingClient {
77    /// Wraps an HTTP client with in-flight request coalescing.
78    pub fn new(inner: Arc<dyn HttpClient>) -> Self {
79        Self {
80            inner,
81            in_flight: Mutex::new(HashMap::new()),
82        }
83    }
84}
85
86#[async_trait]
87impl HttpClient for CoalescingClient {
88    async fn send(&self, request: HttpRequest) -> Result<HttpResponse, HttpClientError> {
89        if (request.method != Method::GET && request.method != Method::HEAD)
90            || request.body.is_some()
91        {
92            return self.inner.send(request).await;
93        }
94
95        let key = CoalesceKey {
96            unique_key: Request::compute_unique_key(&request.url, &request.method, None),
97            proxy: request.proxy.as_ref().map(ToString::to_string),
98            jar: request
99                .cookie_jar
100                .as_ref()
101                .map(|jar| Arc::as_ptr(jar) as usize),
102        };
103        let cell = {
104            let mut in_flight = self
105                .in_flight
106                .lock()
107                .unwrap_or_else(|error| error.into_inner());
108            if let Some(cell) = in_flight.get(&key) {
109                Arc::clone(cell)
110            } else {
111                let cell = Arc::new(OnceCell::new());
112                in_flight.insert(key.clone(), Arc::clone(&cell));
113                cell
114            }
115        };
116
117        let mut cleanup = InFlightGuard {
118            in_flight: &self.in_flight,
119            key,
120            cell: Arc::clone(&cell),
121            active: false,
122        };
123        let mut leader_result = None;
124        let shared_result = cell
125            .get_or_init(|| async {
126                cleanup.active = true;
127                let result = self.inner.send(request).await;
128                let follower_result = result
129                    .as_ref()
130                    .map(Clone::clone)
131                    .map_err(ToString::to_string);
132                leader_result = Some(result);
133                follower_result
134            })
135            .await;
136
137        leader_result.unwrap_or_else(|| {
138            shared_result.clone().map_err(|message| {
139                HttpClientError::other(anyhow!("coalesced request failed: {message}"))
140            })
141        })
142    }
143
144    async fn stream(&self, request: HttpRequest) -> Result<StreamingResponse, HttpClientError> {
145        self.inner.stream(request).await
146    }
147}
148
149impl fmt::Debug for CoalescingClient {
150    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
151        let in_flight = self
152            .in_flight
153            .lock()
154            .unwrap_or_else(|error| error.into_inner())
155            .len();
156        formatter
157            .debug_struct("CoalescingClient")
158            .field("inner", &"dyn HttpClient")
159            .field("in_flight", &in_flight)
160            .finish()
161    }
162}
millipede_http/coalesce.rs

millipede_http/
coalesce.rs
