Skip to main content

reddb_client/
routing_cache.rs

1//! Driver routing cache updates from MOVED redirects.
2
3use std::collections::BTreeMap;
4
5use reddb_wire::MovedRedirect;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct CachedOwner {
9    pub owner_addr: String,
10    pub ownership_epoch: u64,
11    pub catalog_version: u64,
12}
13
14#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
15struct RangeKey {
16    collection: String,
17    range_id: u64,
18    slot: Option<u64>,
19}
20
21#[derive(Debug, Default, Clone)]
22pub struct RoutingCache {
23    ranges: BTreeMap<RangeKey, CachedOwner>,
24}
25
26impl RoutingCache {
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    pub fn apply_moved(&mut self, moved: &MovedRedirect) -> bool {
32        let key = RangeKey {
33            collection: moved.collection.clone(),
34            range_id: moved.range_id,
35            slot: moved.slot,
36        };
37        let incoming = CachedOwner {
38            owner_addr: moved.owner_addr.clone(),
39            ownership_epoch: moved.ownership_epoch,
40            catalog_version: moved.catalog_version,
41        };
42
43        match self.ranges.get(&key) {
44            Some(existing)
45                if existing.catalog_version > incoming.catalog_version
46                    || existing.ownership_epoch > incoming.ownership_epoch =>
47            {
48                false
49            }
50            _ => {
51                self.ranges.insert(key, incoming);
52                true
53            }
54        }
55    }
56
57    pub fn owner_for(
58        &self,
59        collection: &str,
60        range_id: u64,
61        slot: Option<u64>,
62    ) -> Option<&CachedOwner> {
63        self.ranges.get(&RangeKey {
64            collection: collection.to_string(),
65            range_id,
66            slot,
67        })
68    }
69
70    pub fn retry_owner_after_moved(&mut self, moved: &MovedRedirect) -> &str {
71        self.apply_moved(moved);
72        &self
73            .owner_for(&moved.collection, moved.range_id, moved.slot)
74            .expect("MOVED update must populate retry owner")
75            .owner_addr
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    fn moved(owner_addr: &str, epoch: u64, version: u64) -> MovedRedirect {
84        MovedRedirect {
85            slot: Some(12),
86            collection: "orders".to_string(),
87            range_id: 3,
88            owner_addr: owner_addr.to_string(),
89            ownership_epoch: epoch,
90            catalog_version: version,
91            reason: "transaction".to_string(),
92        }
93    }
94
95    #[test]
96    fn moved_redirect_updates_cache_and_retry_targets_current_owner() {
97        let mut cache = RoutingCache::new();
98        cache.apply_moved(&moved("node-a:5050", 1, 1));
99
100        let retry_owner = cache.retry_owner_after_moved(&moved("node-b:5050", 2, 2));
101
102        assert_eq!(retry_owner, "node-b:5050");
103        assert_eq!(
104            cache.owner_for("orders", 3, Some(12)),
105            Some(&CachedOwner {
106                owner_addr: "node-b:5050".to_string(),
107                ownership_epoch: 2,
108                catalog_version: 2,
109            })
110        );
111    }
112
113    #[test]
114    fn stale_moved_payload_does_not_replace_newer_route() {
115        let mut cache = RoutingCache::new();
116        assert!(cache.apply_moved(&moved("node-b:5050", 3, 3)));
117
118        assert!(!cache.apply_moved(&moved("node-a:5050", 2, 2)));
119
120        assert_eq!(
121            cache.owner_for("orders", 3, Some(12)).unwrap().owner_addr,
122            "node-b:5050"
123        );
124    }
125}