network/ip/
bucket.rs

1use crate::{
2    ip::sample::SampleHashSet,
3    node_database::NodeDatabase,
4    node_table::{NodeContact, NodeId},
5};
6use rand::{prelude::ThreadRng, Rng};
7use std::time::Duration;
8
9/// NodeBucket is used to manage the nodes that grouped by subnet,
10/// and support to sample any node from bucket in O(1) time complexity.
11#[derive(Default, Debug)]
12pub struct NodeBucket {
13    nodes: SampleHashSet<NodeId>,
14}
15
16impl NodeBucket {
17    #[inline]
18    pub fn count(&self) -> usize { self.nodes.len() }
19
20    #[inline]
21    pub fn add(&mut self, id: NodeId) -> bool { self.nodes.insert(id) }
22
23    #[inline]
24    pub fn remove(&mut self, id: &NodeId) -> bool { self.nodes.remove(id) }
25
26    #[inline]
27    pub fn sample(&self, rng: &mut ThreadRng) -> Option<NodeId> {
28        self.nodes.sample(rng)
29    }
30
31    /// Select a node to evict due to bucket is full. The basic priority is as
32    /// following:
33    /// - Do not evict connecting nodes.
34    /// - Evict nodes that have not been contacted for a long time.
35    /// - Randomly pick a node without "fresher" bias.
36    pub fn select_evictee(
37        &self, db: &NodeDatabase, evict_timeout: Duration,
38    ) -> Option<NodeId> {
39        let mut long_time_nodes = Vec::new();
40        let mut evictable_nodes = Vec::new();
41
42        for id in self.nodes.iter() {
43            if let Some(node) = db.get(id, false /* trusted_only */) {
44                // do not evict the connecting nodes
45                if let Some(NodeContact::Success(_)) = node.last_connected {
46                    continue;
47                }
48
49                match node.last_contact {
50                    Some(contact) => match contact.time().elapsed() {
51                        Ok(d) => {
52                            if d > evict_timeout {
53                                long_time_nodes.push(id);
54                            } else {
55                                evictable_nodes.push(id);
56                            }
57                        }
58                        Err(_) => long_time_nodes.push(id),
59                    },
60                    None => long_time_nodes.push(id),
61                }
62            }
63        }
64
65        let mut rng = rand::rng();
66
67        // evict out-of-date node with high priority
68        if !long_time_nodes.is_empty() {
69            let index = rng.random_range(0..long_time_nodes.len());
70            return Some(*long_time_nodes[index]);
71        }
72
73        // randomly evict one
74        if !evictable_nodes.is_empty() {
75            let index = rng.random_range(0..evictable_nodes.len());
76            return Some(*evictable_nodes[index]);
77        }
78
79        None
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::{NodeBucket, NodeId};
86
87    #[test]
88    fn test_add_remove() {
89        let mut bucket = NodeBucket::default();
90        assert_eq!(bucket.count(), 0);
91
92        // succeed to add n1
93        let n1 = NodeId::random();
94        assert_eq!(bucket.add(n1.clone()), true);
95        assert_eq!(bucket.count(), 1);
96
97        // cannot add n1 again
98        assert_eq!(bucket.add(n1.clone()), false);
99        assert_eq!(bucket.count(), 1);
100
101        // succeed to add n2
102        let n2 = NodeId::random();
103        assert_eq!(bucket.add(n2.clone()), true);
104        assert_eq!(bucket.count(), 2);
105
106        // failed to remove non-exist node n3
107        let n3 = NodeId::random();
108        assert_eq!(bucket.remove(&n3), false);
109
110        // succeed to remove existing n1/n2
111        assert_eq!(bucket.remove(&n1), true);
112        assert_eq!(bucket.count(), 1);
113
114        assert_eq!(bucket.remove(&n2), true);
115        assert_eq!(bucket.count(), 0);
116    }
117
118    #[test]
119    fn test_sample() {
120        let mut bucket = NodeBucket::default();
121        let mut rng = rand::rng();
122
123        // sample None if bucket is empty
124        assert_eq!(bucket.sample(&mut rng), None);
125
126        // sample any trusted node
127        let n1 = NodeId::random();
128        assert_eq!(bucket.add(n1.clone()), true);
129        assert_eq!(bucket.sample(&mut rng), Some(n1.clone()));
130    }
131}
network/ip/bucket.rs

network/ip/
bucket.rs
