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    async fn get(
45        &self,
46        url: &url::Url,
47        params: &[(&str, &str)],
48    ) -> Result<(u16, std::sync::Arc<String>), ClientError> {
49        let response = self
50            .try_get_or_wait(url, params, async {
51                match self.inner.get(url, params).await {
52                    Ok((status, data)) => Response::String(status, data),
53                    Err(err) => Response::Error(err),
54                }
55            })
56            .await;
57
58        match response {
59            Response::String(status, data) => Ok((status, data)),
60            Response::Binary(_, _) => panic!(),
61            Response::Error(client_error) => Err(client_error),
62        }
63    }
64
65    async fn get_bin(
66        &self,
67        url: &url::Url,
68        params: &[(&str, &str)],
69    ) -> Result<(u16, std::sync::Arc<Vec<u8>>), ClientError> {
70        let response = self
71            .try_get_or_wait(url, params, async {
72                match self.inner.get_bin(url, params).await {
73                    Ok((status, data)) => Response::Binary(status, data),
74                    Err(err) => Response::Error(err),
75                }
76            })
77            .await;
78
79        match response {
80            Response::String(_, _) => panic!(),
81            Response::Binary(status, data) => Ok((status, data)),
82            Response::Error(client_error) => Err(client_error),
83        }
84    }
85}
86
87impl<C: Client> RequestDeduplicationClient<C> {
88    async fn try_get_or_wait(
89        &self,
90        url: &url::Url,
91        params: &[(&str, &str)],
92        getter: impl Future<Output = Response>,
93    ) -> Response {
94        let key = DeduplicationKey {
95            url: url.clone(),
96            params: params
97                .iter()
98                .map(|(k, v)| (k.to_string(), v.to_string()))
99                .collect(),
100        };
101
102        let (rx, request) = {
103            let mut requests = self.requests.lock().unwrap();
104
105            let (tx, rx) = channel::<Response>();
106            match requests.get_mut(&key) {
107                Some(ongoing_requests) => {
108                    ongoing_requests.push(tx);
109
110                    tracing::trace!(
111                        "Deduplicated request to {}. Queue length: {}",
112                        key,
113                        ongoing_requests.len()
114                    );
115                    (rx, None)
116                }
117                None => {
118                    tracing::debug!("A fresh request to: {}", key);
119
120                    requests.insert(key.clone(), vec![tx]);
121
122                    let request = async move {
123                        let response = getter.await;
124                        let mut requests = self.requests.lock().unwrap();
125
126                        let senders = requests
127                            .remove(&key)
128                            .expect("Key no longer exist. This is a bug");
129
130                        tracing::debug!(
131                            "Sending response of {} to {} requesters",
132                            key,
133                            senders.len()
134                        );
135
136                        for tx in senders {
137                            tx.send(response.clone()).unwrap();
138                        }
139                    };
140
141                    (rx, Some(request))
142                }
143            }
144        };
145
146        // If we are making a request, wait on it
147        if let Some(request) = request {
148            request.await;
149        }
150
151        // Await on receiving the answer
152        rx.await
153            .expect("Dedup channel closed instead of answered. This is a bug")
154    }
155}
156
157#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
158struct DeduplicationKey {
159    url: url::Url,
160    params: Vec<(String, String)>,
161}
162
163impl fmt::Display for DeduplicationKey {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        write!(f, "DeduplicationKey: {}:{:?}", self.url, self.params)
166    }
167}
168
169#[derive(Debug, Clone)]
170enum Response {
171    String(u16, std::sync::Arc<String>),
172    Binary(u16, std::sync::Arc<Vec<u8>>),
173    Error(ClientError),
174}
luct_client/impls/deduplication.rs

luct_client/impls/
deduplication.rs
