oximedia_distributed/
node_topology.rs

1#![allow(dead_code)]
2//! Network topology awareness for data-local task scheduling.
3//!
4//! Models the physical and logical topology of a distributed cluster so that
5//! the scheduler can prefer nodes that are close to the data, minimising
6//! network transfer costs.
7
8use std::collections::{HashMap, HashSet};
9use std::fmt;
10
11/// Identifies a location tier in the topology hierarchy.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub enum LocationTier {
14    /// Same physical host (local).
15    Host,
16    /// Same rack / network switch.
17    Rack,
18    /// Same data-centre / availability zone.
19    DataCenter,
20    /// Same geographic region.
21    Region,
22    /// Different region (cross-region).
23    Remote,
24}
25
26impl LocationTier {
27    /// Relative cost weight (lower is better).
28    #[must_use]
29    pub fn cost_weight(self) -> u32 {
30        match self {
31            Self::Host => 0,
32            Self::Rack => 1,
33            Self::DataCenter => 5,
34            Self::Region => 20,
35            Self::Remote => 100,
36        }
37    }
38}
39
40impl fmt::Display for LocationTier {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        match self {
43            Self::Host => write!(f, "Host"),
44            Self::Rack => write!(f, "Rack"),
45            Self::DataCenter => write!(f, "DataCenter"),
46            Self::Region => write!(f, "Region"),
47            Self::Remote => write!(f, "Remote"),
48        }
49    }
50}
51
52/// Physical location descriptor for a node.
53#[derive(Debug, Clone, PartialEq, Eq, Hash)]
54pub struct NodeLocation {
55    /// Region label (e.g. "us-east-1").
56    pub region: String,
57    /// Data-centre / availability zone.
58    pub data_center: String,
59    /// Rack identifier.
60    pub rack: String,
61    /// Hostname.
62    pub host: String,
63}
64
65impl NodeLocation {
66    /// Create a new location.
67    pub fn new(
68        region: impl Into<String>,
69        data_center: impl Into<String>,
70        rack: impl Into<String>,
71        host: impl Into<String>,
72    ) -> Self {
73        Self {
74            region: region.into(),
75            data_center: data_center.into(),
76            rack: rack.into(),
77            host: host.into(),
78        }
79    }
80
81    /// Determine the tier of proximity between two locations.
82    #[must_use]
83    pub fn tier_to(&self, other: &Self) -> LocationTier {
84        if self.host == other.host
85            && self.rack == other.rack
86            && self.data_center == other.data_center
87            && self.region == other.region
88        {
89            LocationTier::Host
90        } else if self.rack == other.rack
91            && self.data_center == other.data_center
92            && self.region == other.region
93        {
94            LocationTier::Rack
95        } else if self.data_center == other.data_center {
96            LocationTier::DataCenter
97        } else if self.region == other.region {
98            LocationTier::Region
99        } else {
100            LocationTier::Remote
101        }
102    }
103
104    /// Cost weight of transferring data to another location.
105    #[must_use]
106    pub fn cost_to(&self, other: &Self) -> u32 {
107        self.tier_to(other).cost_weight()
108    }
109}
110
111/// A node registered in the topology.
112#[derive(Debug, Clone)]
113pub struct TopologyNode {
114    /// Unique node identifier.
115    pub node_id: String,
116    /// Physical location.
117    pub location: NodeLocation,
118    /// Whether the node is currently available.
119    pub available: bool,
120    /// Set of data block IDs this node holds locally.
121    pub local_data: HashSet<String>,
122}
123
124impl TopologyNode {
125    /// Create a new topology node.
126    pub fn new(node_id: impl Into<String>, location: NodeLocation) -> Self {
127        Self {
128            node_id: node_id.into(),
129            location,
130            available: true,
131            local_data: HashSet::new(),
132        }
133    }
134
135    /// Mark data as locally available on this node.
136    pub fn add_local_data(&mut self, data_id: impl Into<String>) {
137        self.local_data.insert(data_id.into());
138    }
139
140    /// Remove a data block reference.
141    pub fn remove_local_data(&mut self, data_id: &str) {
142        self.local_data.remove(data_id);
143    }
144
145    /// Check whether this node has a specific data block.
146    #[must_use]
147    pub fn has_data(&self, data_id: &str) -> bool {
148        self.local_data.contains(data_id)
149    }
150}
151
152/// The cluster topology manager.
153#[derive(Debug, Clone)]
154pub struct TopologyManager {
155    /// All registered nodes.
156    nodes: HashMap<String, TopologyNode>,
157}
158
159impl TopologyManager {
160    /// Create an empty topology.
161    #[must_use]
162    pub fn new() -> Self {
163        Self {
164            nodes: HashMap::new(),
165        }
166    }
167
168    /// Register a node.
169    pub fn add_node(&mut self, node: TopologyNode) {
170        self.nodes.insert(node.node_id.clone(), node);
171    }
172
173    /// Remove a node by ID.
174    pub fn remove_node(&mut self, node_id: &str) -> Option<TopologyNode> {
175        self.nodes.remove(node_id)
176    }
177
178    /// Get a node by ID.
179    #[must_use]
180    pub fn get_node(&self, node_id: &str) -> Option<&TopologyNode> {
181        self.nodes.get(node_id)
182    }
183
184    /// Get a mutable reference to a node.
185    pub fn get_node_mut(&mut self, node_id: &str) -> Option<&mut TopologyNode> {
186        self.nodes.get_mut(node_id)
187    }
188
189    /// Number of nodes.
190    #[must_use]
191    pub fn node_count(&self) -> usize {
192        self.nodes.len()
193    }
194
195    /// List all node IDs.
196    #[must_use]
197    pub fn node_ids(&self) -> Vec<&str> {
198        self.nodes.keys().map(std::string::String::as_str).collect()
199    }
200
201    /// Find nodes that have a specific data block locally, sorted by availability.
202    #[must_use]
203    pub fn nodes_with_data(&self, data_id: &str) -> Vec<&TopologyNode> {
204        let mut nodes: Vec<&TopologyNode> = self
205            .nodes
206            .values()
207            .filter(|n| n.has_data(data_id) && n.available)
208            .collect();
209        // Stable sort: available nodes first (already filtered), then by node_id for determinism
210        nodes.sort_by(|a, b| a.node_id.cmp(&b.node_id));
211        nodes
212    }
213
214    /// Rank candidate nodes for a task that needs `data_id`, preferring
215    /// nodes closest to `reference_location`.
216    ///
217    /// Returns node IDs sorted by ascending transfer cost.
218    #[must_use]
219    pub fn rank_by_locality(
220        &self,
221        data_id: &str,
222        reference_location: &NodeLocation,
223    ) -> Vec<(String, u32)> {
224        let mut candidates: Vec<(String, u32)> = self
225            .nodes
226            .values()
227            .filter(|n| n.available)
228            .map(|n| {
229                let mut cost = reference_location.cost_to(&n.location);
230                // Bonus: if the node already has the data, cost is even lower
231                if n.has_data(data_id) {
232                    cost = cost.saturating_sub(1);
233                }
234                (n.node_id.clone(), cost)
235            })
236            .collect();
237        candidates.sort_by_key(|&(_, cost)| cost);
238        candidates
239    }
240
241    /// Get all available nodes in a given region.
242    #[must_use]
243    pub fn nodes_in_region(&self, region: &str) -> Vec<&TopologyNode> {
244        self.nodes
245            .values()
246            .filter(|n| n.location.region == region && n.available)
247            .collect()
248    }
249
250    /// Get all available nodes in a given data centre.
251    #[must_use]
252    pub fn nodes_in_data_center(&self, dc: &str) -> Vec<&TopologyNode> {
253        self.nodes
254            .values()
255            .filter(|n| n.location.data_center == dc && n.available)
256            .collect()
257    }
258
259    /// Get all available nodes on a given rack.
260    #[must_use]
261    pub fn nodes_in_rack(&self, rack: &str) -> Vec<&TopologyNode> {
262        self.nodes
263            .values()
264            .filter(|n| n.location.rack == rack && n.available)
265            .collect()
266    }
267
268    /// Set a node's availability.
269    pub fn set_available(&mut self, node_id: &str, available: bool) {
270        if let Some(node) = self.nodes.get_mut(node_id) {
271            node.available = available;
272        }
273    }
274
275    /// List distinct regions.
276    #[must_use]
277    pub fn regions(&self) -> Vec<String> {
278        let mut set: HashSet<String> = HashSet::new();
279        for n in self.nodes.values() {
280            set.insert(n.location.region.clone());
281        }
282        let mut regions: Vec<String> = set.into_iter().collect();
283        regions.sort();
284        regions
285    }
286}
287
288impl Default for TopologyManager {
289    fn default() -> Self {
290        Self::new()
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    fn loc(region: &str, dc: &str, rack: &str, host: &str) -> NodeLocation {
299        NodeLocation::new(region, dc, rack, host)
300    }
301
302    #[test]
303    fn test_tier_same_host() {
304        let a = loc("us", "dc1", "r1", "h1");
305        let b = loc("us", "dc1", "r1", "h1");
306        assert_eq!(a.tier_to(&b), LocationTier::Host);
307    }
308
309    #[test]
310    fn test_tier_same_rack() {
311        let a = loc("us", "dc1", "r1", "h1");
312        let b = loc("us", "dc1", "r1", "h2");
313        assert_eq!(a.tier_to(&b), LocationTier::Rack);
314    }
315
316    #[test]
317    fn test_tier_same_dc() {
318        let a = loc("us", "dc1", "r1", "h1");
319        let b = loc("us", "dc1", "r2", "h3");
320        assert_eq!(a.tier_to(&b), LocationTier::DataCenter);
321    }
322
323    #[test]
324    fn test_tier_same_region() {
325        let a = loc("us", "dc1", "r1", "h1");
326        let b = loc("us", "dc2", "r5", "h9");
327        assert_eq!(a.tier_to(&b), LocationTier::Region);
328    }
329
330    #[test]
331    fn test_tier_remote() {
332        let a = loc("us", "dc1", "r1", "h1");
333        let b = loc("eu", "dc3", "r1", "h1");
334        assert_eq!(a.tier_to(&b), LocationTier::Remote);
335    }
336
337    #[test]
338    fn test_cost_ordering() {
339        assert!(LocationTier::Host.cost_weight() < LocationTier::Rack.cost_weight());
340        assert!(LocationTier::Rack.cost_weight() < LocationTier::DataCenter.cost_weight());
341        assert!(LocationTier::DataCenter.cost_weight() < LocationTier::Region.cost_weight());
342        assert!(LocationTier::Region.cost_weight() < LocationTier::Remote.cost_weight());
343    }
344
345    #[test]
346    fn test_topology_manager_add_remove() {
347        let mut mgr = TopologyManager::new();
348        let node = TopologyNode::new("n1", loc("us", "dc1", "r1", "h1"));
349        mgr.add_node(node);
350        assert_eq!(mgr.node_count(), 1);
351        mgr.remove_node("n1");
352        assert_eq!(mgr.node_count(), 0);
353    }
354
355    #[test]
356    fn test_nodes_with_data() {
357        let mut mgr = TopologyManager::new();
358        let mut n1 = TopologyNode::new("n1", loc("us", "dc1", "r1", "h1"));
359        n1.add_local_data("block-42");
360        let n2 = TopologyNode::new("n2", loc("us", "dc1", "r1", "h2"));
361        mgr.add_node(n1);
362        mgr.add_node(n2);
363        let holders = mgr.nodes_with_data("block-42");
364        assert_eq!(holders.len(), 1);
365        assert_eq!(holders[0].node_id, "n1");
366    }
367
368    #[test]
369    fn test_rank_by_locality() {
370        let mut mgr = TopologyManager::new();
371        let mut n_local = TopologyNode::new("local", loc("us", "dc1", "r1", "h1"));
372        n_local.add_local_data("data-1");
373        let n_remote = TopologyNode::new("remote", loc("eu", "dc3", "r1", "h9"));
374        mgr.add_node(n_local);
375        mgr.add_node(n_remote);
376        let ref_loc = loc("us", "dc1", "r1", "h1");
377        let ranked = mgr.rank_by_locality("data-1", &ref_loc);
378        assert_eq!(ranked[0].0, "local");
379        assert!(ranked[0].1 < ranked[1].1);
380    }
381
382    #[test]
383    fn test_set_available() {
384        let mut mgr = TopologyManager::new();
385        mgr.add_node(TopologyNode::new("n1", loc("us", "dc1", "r1", "h1")));
386        mgr.set_available("n1", false);
387        assert!(!mgr.get_node("n1").expect("node should exist").available);
388        mgr.set_available("n1", true);
389        assert!(mgr.get_node("n1").expect("node should exist").available);
390    }
391
392    #[test]
393    fn test_nodes_in_region() {
394        let mut mgr = TopologyManager::new();
395        mgr.add_node(TopologyNode::new("n1", loc("us", "dc1", "r1", "h1")));
396        mgr.add_node(TopologyNode::new("n2", loc("eu", "dc2", "r1", "h1")));
397        assert_eq!(mgr.nodes_in_region("us").len(), 1);
398        assert_eq!(mgr.nodes_in_region("eu").len(), 1);
399    }
400
401    #[test]
402    fn test_nodes_in_data_center() {
403        let mut mgr = TopologyManager::new();
404        mgr.add_node(TopologyNode::new("n1", loc("us", "dc1", "r1", "h1")));
405        mgr.add_node(TopologyNode::new("n2", loc("us", "dc1", "r2", "h2")));
406        mgr.add_node(TopologyNode::new("n3", loc("us", "dc2", "r1", "h3")));
407        assert_eq!(mgr.nodes_in_data_center("dc1").len(), 2);
408    }
409
410    #[test]
411    fn test_regions_list() {
412        let mut mgr = TopologyManager::new();
413        mgr.add_node(TopologyNode::new("n1", loc("us", "dc1", "r1", "h1")));
414        mgr.add_node(TopologyNode::new("n2", loc("eu", "dc2", "r1", "h1")));
415        mgr.add_node(TopologyNode::new("n3", loc("us", "dc3", "r1", "h2")));
416        let regions = mgr.regions();
417        assert_eq!(regions, vec!["eu", "us"]);
418    }
419
420    #[test]
421    fn test_location_tier_display() {
422        assert_eq!(LocationTier::Host.to_string(), "Host");
423        assert_eq!(LocationTier::Remote.to_string(), "Remote");
424    }
425
426    #[test]
427    fn test_node_local_data_operations() {
428        let mut node = TopologyNode::new("n1", loc("us", "dc1", "r1", "h1"));
429        node.add_local_data("d1");
430        node.add_local_data("d2");
431        assert!(node.has_data("d1"));
432        assert!(node.has_data("d2"));
433        node.remove_local_data("d1");
434        assert!(!node.has_data("d1"));
435        assert!(node.has_data("d2"));
436    }
437}
oximedia_distributed/node_topology.rs

oximedia_distributed/
node_topology.rs
