luct_client/impls/
deduplication.rs

1//! [`Client`] wrapper to deduplicate requests to the same [`Url`](url::Url)
2
3use crate::{Client, ClientError};
4use futures::channel::oneshot::{Sender, channel};
5use std::{
6    collections::BTreeMap,
7    fmt,
8    sync::{Arc, Mutex},
9};
10
11/// Wraps an inner [`Client`] and deduplicates running requests.
12///
13/// The endpoint must be idempotent.
14/// In particular, the following things must be guaranteed:
15///
16/// - The endpoint must return the same response on the same request
17/// - The deduplication may fail due to TOCTOU, and sending the same request
18///   twice must not change the servers behavioru
19#[derive(Clone, Default)]
20pub struct RequestDeduplicationClient<C> {
21    inner: C,
22    requests: Arc<Mutex<BTreeMap<DeduplicationKey, Vec<Sender<Response>>>>>,
23}
24
25impl<C: fmt::Debug> fmt::Debug for RequestDeduplicationClient<C> {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("RequestDeduplicationClient")
28            .field("inner", &self.inner)
29            .field("requests", &self.requests.lock().unwrap().len())
30            .finish()
31    }
32}
33
34impl<C> RequestDeduplicationClient<C> {
35    pub fn new(inner: C) -> Self {
36        Self {
37            inner,
38            requests: Arc::new(Mutex::new(BTreeMap::new())),
39        }
40    }
41}
42
43impl<C: Client> Client for RequestDeduplicationClient<C> {
44    #[tracing::instrument(level = "trace")]
45    async fn get(
46        &self,
47        url: &url::Url,
48        params: &[(&str, &str)],
49    ) -> Result<(u16, std::sync::Arc<String>), ClientError> {
50        let response = self
51            .try_get_or_wait(url, params, async {
52                match self.inner.get(url, params).await {
53                    Ok((status, data)) => Response::String(status, data),
54                    Err(err) => Response::Error(err),
55                }
56            })
57            .await;
58
59        match response {
60            Response::String(status, data) => Ok((status, data)),
61            Response::Binary(_, _) => panic!(),
62            Response::Error(client_error) => Err(client_error),
63        }
64    }
65
66    #[tracing::instrument(level = "trace")]
67    async fn get_bin(
68        &self,
69        url: &url::Url,
70        params: &[(&str, &str)],
71    ) -> Result<(u16, std::sync::Arc<Vec<u8>>), ClientError> {
72        let response = self
73            .try_get_or_wait(url, params, async {
74                match self.inner.get_bin(url, params).await {
75                    Ok((status, data)) => Response::Binary(status, data),
76                    Err(err) => Response::Error(err),
77                }
78            })
79            .await;
80
81        match response {
82            Response::String(_, _) => panic!(),
83            Response::Binary(status, data) => Ok((status, data)),
84            Response::Error(client_error) => Err(client_error),
85        }
86    }
87}
88
89impl<C: Client> RequestDeduplicationClient<C> {
90    async fn try_get_or_wait(
91        &self,
92        url: &url::Url,
93        params: &[(&str, &str)],
94        getter: impl Future<Output = Response>,
95    ) -> Response {
96        let key = DeduplicationKey {
97            url: url.clone(),
98            params: params
99                .iter()
100                .map(|(k, v)| (k.to_string(), v.to_string()))
101                .collect(),
102        };
103
104        let (rx, request) = {
105            let mut requests = self.requests.lock().unwrap();
106
107            let (tx, rx) = channel::<Response>();
108            match requests.get_mut(&key) {
109                Some(ongoing_requests) => {
110                    ongoing_requests.push(tx);
111
112                    tracing::trace!(
113                        "Deduplicated request to {}. Queue length: {}",
114                        key,
115                        ongoing_requests.len()
116                    );
117                    (rx, None)
118                }
119                None => {
120                    tracing::debug!("A fresh request to: {}", key);
121
122                    requests.insert(key.clone(), vec![tx]);
123
124                    let request = async move {
125                        let response = getter.await;
126                        let mut requests = self.requests.lock().unwrap();
127
128                        let senders = requests
129                            .remove(&key)
130                            .expect("Key no longer exist. This is a bug");
131
132                        tracing::debug!(
133                            "Sending response of {} to {} requesters",
134                            key,
135                            senders.len()
136                        );
137
138                        for tx in senders {
139                            tx.send(response.clone()).unwrap();
140                        }
141                    };
142
143                    (rx, Some(request))
144                }
145            }
146        };
147
148        // If we are making a request, wait on it
149        if let Some(request) = request {
150            request.await;
151        }
152
153        // Await on receiving the answer
154        rx.await
155            .expect("Dedup channel closed instead of answered. This is a bug")
156    }
157}
158
159#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
160struct DeduplicationKey {
161    url: url::Url,
162    params: Vec<(String, String)>,
163}
164
165impl fmt::Display for DeduplicationKey {
166    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167        write!(f, "DeduplicationKey: {}:{:?}", self.url, self.params)
168    }
169}
170
171#[derive(Debug, Clone)]
172enum Response {
173    String(u16, std::sync::Arc<String>),
174    Binary(u16, std::sync::Arc<Vec<u8>>),
175    Error(ClientError),
176}
luct_client/impls/deduplication.rs

luct_client/impls/
deduplication.rs
