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

pipedash_core/infrastructure/
deduplication.rs
