Skip to main content

confium_coordinator/
request_coalescing.rs

1//! Request coalescing — merge duplicate concurrent requests.
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6pub struct RequestCoalescer<T: Clone + Send + Sync + 'static> {
7    pending: Mutex<HashMap<String, CoalescedRequest<T>>>,
8}
9
10struct CoalescedRequest<T> {
11    waiters: usize,
12    result: Option<T>,
13}
14
15impl<T: Clone + Send + Sync + 'static> RequestCoalescer<T> {
16    pub fn new() -> Self {
17        Self {
18            pending: Mutex::new(HashMap::new()),
19        }
20    }
21
22    pub fn begin(&self, key: &str) -> bool {
23        let mut pending = self.pending.lock().unwrap();
24        if let Some(req) = pending.get_mut(key) {
25            req.waiters += 1;
26            false // already in progress
27        } else {
28            pending.insert(
29                key.into(),
30                CoalescedRequest {
31                    waiters: 1,
32                    result: None,
33                },
34            );
35            true // this caller should execute
36        }
37    }
38
39    pub fn complete(&self, key: &str, result: T) {
40        let mut pending = self.pending.lock().unwrap();
41        if let Some(req) = pending.get_mut(key) {
42            req.result = Some(result);
43        }
44    }
45
46    pub fn collect_result(&self, key: &str) -> Option<T> {
47        let mut pending = self.pending.lock().unwrap();
48        if let Some(req) = pending.get_mut(key) {
49            if let Some(ref result) = req.result {
50                req.waiters -= 1;
51                if req.waiters == 0 {
52                    return pending.remove(key).and_then(|r| r.result);
53                }
54                return Some(result.clone());
55            }
56        }
57        None
58    }
59
60    pub fn pending_count(&self) -> usize {
61        self.pending.lock().unwrap().len()
62    }
63
64    pub fn waiters_for(&self, key: &str) -> usize {
65        self.pending
66            .lock()
67            .unwrap()
68            .get(key)
69            .map(|r| r.waiters)
70            .unwrap_or(0)
71    }
72}
73
74impl<T: Clone + Send + Sync + 'static> Default for RequestCoalescer<T> {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn first_call_begins() {
86        let coal = RequestCoalescer::<String>::new();
87        assert!(coal.begin("key1"));
88    }
89
90    #[test]
91    fn second_call_coalesced() {
92        let coal = RequestCoalescer::<String>::new();
93        assert!(coal.begin("key1"));
94        assert!(!coal.begin("key1"));
95    }
96
97    #[test]
98    fn different_keys_independent() {
99        let coal = RequestCoalescer::<String>::new();
100        assert!(coal.begin("key1"));
101        assert!(coal.begin("key2"));
102        assert_eq!(coal.pending_count(), 2);
103    }
104
105    #[test]
106    fn complete_and_collect() {
107        let coal = RequestCoalescer::<String>::new();
108        coal.begin("k1");
109        coal.complete("k1", "result".into());
110        let result = coal.collect_result("k1");
111        assert_eq!(result, Some("result".into()));
112    }
113
114    #[test]
115    fn multiple_waiters_share_result() {
116        let coal = RequestCoalescer::<String>::new();
117        coal.begin("k1");
118        coal.begin("k1"); // waiter 2
119        coal.begin("k1"); // waiter 3
120        assert_eq!(coal.waiters_for("k1"), 3);
121        coal.complete("k1", "shared".into());
122        let r1 = coal.collect_result("k1");
123        let r2 = coal.collect_result("k1");
124        let r3 = coal.collect_result("k1");
125        assert_eq!(r1, Some("shared".into()));
126        assert_eq!(r2, Some("shared".into()));
127        assert_eq!(r3, Some("shared".into()));
128        assert_eq!(coal.pending_count(), 0);
129    }
130
131    #[test]
132    fn collect_without_complete_returns_none() {
133        let coal = RequestCoalescer::<String>::new();
134        coal.begin("k1");
135        assert!(coal.collect_result("k1").is_none());
136    }
137}