pipedash_core/infrastructure/
deduplication.rs

1use std::collections::hash_map::DefaultHasher;
2use std::future::Future;
3use std::hash::{
4    Hash,
5    Hasher,
6};
7use std::sync::Arc;
8
9use dashmap::DashMap;
10use sha2::{
11    Digest,
12    Sha256,
13};
14use tokio::sync::oneshot;
15
16type RequestId = u64;
17type InFlight<T> = DashMap<RequestId, Vec<oneshot::Sender<Arc<T>>>>;
18
19pub struct RequestDeduplicator<T: Clone> {
20    in_flight: Arc<InFlight<T>>,
21}
22
23struct InFlightGuard<'a, T> {
24    in_flight: &'a InFlight<T>,
25    request_id: RequestId,
26    armed: bool,
27}
28
29impl<T> InFlightGuard<'_, T> {
30    fn take_waiters(mut self) -> Vec<oneshot::Sender<Arc<T>>> {
31        self.armed = false;
32        self.in_flight
33            .remove(&self.request_id)
34            .map(|(_, waiters)| waiters)
35            .unwrap_or_default()
36    }
37}
38
39impl<T> Drop for InFlightGuard<'_, T> {
40    fn drop(&mut self) {
41        if self.armed {
42            self.in_flight.remove(&self.request_id);
43        }
44    }
45}
46
47impl<T: Clone> RequestDeduplicator<T> {
48    pub fn new() -> Self {
49        Self {
50            in_flight: Arc::new(DashMap::new()),
51        }
52    }
53
54    pub async fn deduplicate<F, Fut, E>(&self, request_id: RequestId, operation: F) -> Result<T, E>
55    where
56        F: FnOnce() -> Fut,
57        Fut: Future<Output = Result<T, E>>,
58    {
59        let waiter = match self.in_flight.entry(request_id) {
60            dashmap::Entry::Occupied(mut entry) => {
61                let (tx, rx) = oneshot::channel();
62                entry.get_mut().push(tx);
63                Some(rx)
64            }
65            dashmap::Entry::Vacant(entry) => {
66                entry.insert(Vec::new());
67                None
68            }
69        };
70
71        if let Some(rx) = waiter {
72            return match rx.await {
73                Ok(result) => Ok((*result).clone()),
74                Err(_) => operation().await,
75            };
76        }
77
78        let guard = InFlightGuard {
79            in_flight: &self.in_flight,
80            request_id,
81            armed: true,
82        };
83        let result = operation().await;
84        let waiters = guard.take_waiters();
85
86        if let Ok(value) = &result {
87            let shared_result = Arc::new(value.clone());
88            for waiter in waiters {
89                let _ = waiter.send(shared_result.clone());
90            }
91        }
92
93        result
94    }
95}
96
97impl<T: Clone> Default for RequestDeduplicator<T> {
98    fn default() -> Self {
99        Self::new()
100    }
101}
102
103pub fn hash_request(provider_id: i64, endpoint: &str) -> RequestId {
104    let mut hasher = DefaultHasher::new();
105    provider_id.hash(&mut hasher);
106    endpoint.hash(&mut hasher);
107    hasher.finish()
108}
109
110pub fn hash_pipeline_run(
111    run_number: i64, status: &str, branch: Option<&str>, started_at: &str,
112    duration_seconds: Option<i64>, commit_sha: Option<&str>,
113) -> String {
114    let mut hasher = Sha256::new();
115    hasher.update(run_number.to_string().as_bytes());
116    hasher.update(status.as_bytes());
117    hasher.update(branch.unwrap_or("").as_bytes());
118    hasher.update(started_at.as_bytes());
119    if let Some(duration) = duration_seconds {
120        hasher.update(duration.to_string().as_bytes());
121    }
122    hasher.update(commit_sha.unwrap_or("").as_bytes());
123
124    hasher
125        .finalize()
126        .iter()
127        .map(|b| format!("{b:02x}"))
128        .collect()
129}
130
131impl<T: Clone> Clone for RequestDeduplicator<T> {
132    fn clone(&self) -> Self {
133        Self {
134            in_flight: self.in_flight.clone(),
135        }
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use std::sync::atomic::{
142        AtomicUsize,
143        Ordering,
144    };
145
146    use super::*;
147
148    #[tokio::test]
149    async fn test_request_deduplication() {
150        let dedup = RequestDeduplicator::<i32>::new();
151        let call_count = Arc::new(AtomicUsize::new(0));
152
153        let request_id = hash_request(1, "test");
154
155        let dedup1 = dedup.clone();
156        let count1 = call_count.clone();
157        let handle1 = tokio::spawn(async move {
158            dedup1
159                .deduplicate(request_id, || async {
160                    count1.fetch_add(1, Ordering::SeqCst);
161                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
162                    Ok::<_, ()>(42)
163                })
164                .await
165        });
166
167        let dedup2 = dedup.clone();
168        let count2 = call_count.clone();
169        let handle2 = tokio::spawn(async move {
170            dedup2
171                .deduplicate(request_id, || async {
172                    count2.fetch_add(1, Ordering::SeqCst);
173                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
174                    Ok::<_, ()>(42)
175                })
176                .await
177        });
178
179        let (r1, r2) = tokio::join!(handle1, handle2);
180        assert_eq!(r1.unwrap().unwrap(), 42);
181        assert_eq!(r2.unwrap().unwrap(), 42);
182
183        assert_eq!(call_count.load(Ordering::SeqCst), 1);
184    }
185
186    #[tokio::test]
187    async fn failed_request_does_not_block_the_next_one() {
188        let dedup = RequestDeduplicator::<i32>::new();
189        let request_id = hash_request(1, "fetch_runs_x_100");
190
191        let first = dedup
192            .deduplicate(request_id, || async { Err::<i32, &str>("provider error") })
193            .await;
194        assert_eq!(first, Err("provider error"));
195
196        let second = tokio::time::timeout(
197            tokio::time::Duration::from_secs(2),
198            dedup.deduplicate(request_id, || async { Ok::<_, &str>(7) }),
199        )
200        .await;
201        assert_eq!(second, Ok(Ok(7)));
202    }
203
204    #[tokio::test]
205    async fn dropped_request_does_not_block_the_next_one() {
206        let dedup = RequestDeduplicator::<i32>::new();
207        let request_id = hash_request(1, "fetch_pipelines");
208
209        let abandoned = tokio::time::timeout(
210            tokio::time::Duration::from_millis(20),
211            dedup.deduplicate(request_id, || async {
212                tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
213                Ok::<_, &str>(1)
214            }),
215        )
216        .await;
217        assert!(abandoned.is_err());
218
219        let next = tokio::time::timeout(
220            tokio::time::Duration::from_secs(2),
221            dedup.deduplicate(request_id, || async { Ok::<_, &str>(2) }),
222        )
223        .await;
224        assert_eq!(next, Ok(Ok(2)));
225    }
226
227    #[tokio::test]
228    async fn waiter_runs_its_own_request_when_the_first_one_fails() {
229        let dedup = RequestDeduplicator::<i32>::new();
230        let request_id = hash_request(1, "fetch_pipelines");
231
232        let leader = dedup.clone();
233        let first = tokio::spawn(async move {
234            leader
235                .deduplicate(request_id, || async {
236                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
237                    Err::<i32, &str>("provider error")
238                })
239                .await
240        });
241        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
242
243        let waiter = tokio::time::timeout(
244            tokio::time::Duration::from_secs(2),
245            dedup.deduplicate(request_id, || async { Ok::<_, &str>(3) }),
246        )
247        .await;
248
249        assert_eq!(first.await.unwrap(), Err("provider error"));
250        assert_eq!(waiter, Ok(Ok(3)));
251    }
252
253    #[test]
254    fn test_hash_request() {
255        let h1 = hash_request(1, "pipelines");
256        let h2 = hash_request(1, "pipelines");
257        let h3 = hash_request(2, "pipelines");
258
259        assert_eq!(h1, h2);
260        assert_ne!(h1, h3);
261    }
262
263    #[test]
264    fn test_hash_pipeline_run() {
265        let h1 = hash_pipeline_run(
266            1,
267            "success",
268            Some("main"),
269            "2024-01-01T00:00:00Z",
270            Some(60),
271            None,
272        );
273        let h2 = hash_pipeline_run(
274            1,
275            "success",
276            Some("main"),
277            "2024-01-01T00:00:00Z",
278            Some(60),
279            None,
280        );
281        let h3 = hash_pipeline_run(
282            1,
283            "failed",
284            Some("main"),
285            "2024-01-01T00:00:00Z",
286            Some(60),
287            None,
288        );
289
290        assert_eq!(h1, h2);
291        assert_ne!(h1, h3);
292    }
293}
pipedash_core/infrastructure/deduplication.rs

pipedash_core/infrastructure/
deduplication.rs
