network/ip/
bucket.rs

1use crate::{
2    ip::sample::SampleHashSet,
3    node_database::NodeDatabase,
4    node_table::{NodeContact, NodeId},
5};
6use rand::{prelude::ThreadRng, thread_rng, 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 = thread_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].clone());
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].clone());
77        }
78
79        None
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::{NodeBucket, NodeId};
86    use rand::thread_rng;
87
88    #[test]
89    fn test_add_remove() {
90        let mut bucket = NodeBucket::default();
91        assert_eq!(bucket.count(), 0);
92
93        // succeed to add n1
94        let n1 = NodeId::random();
95        assert_eq!(bucket.add(n1.clone()), true);
96        assert_eq!(bucket.count(), 1);
97
98        // cannot add n1 again
99        assert_eq!(bucket.add(n1.clone()), false);
100        assert_eq!(bucket.count(), 1);
101
102        // succeed to add n2
103        let n2 = NodeId::random();
104        assert_eq!(bucket.add(n2.clone()), true);
105        assert_eq!(bucket.count(), 2);
106
107        // failed to remove non-exist node n3
108        let n3 = NodeId::random();
109        assert_eq!(bucket.remove(&n3), false);
110
111        // succeed to remove existing n1/n2
112        assert_eq!(bucket.remove(&n1), true);
113        assert_eq!(bucket.count(), 1);
114
115        assert_eq!(bucket.remove(&n2), true);
116        assert_eq!(bucket.count(), 0);
117    }
118
119    #[test]
120    fn test_sample() {
121        let mut bucket = NodeBucket::default();
122        let mut rng = thread_rng();
123
124        // sample None if bucket is empty
125        assert_eq!(bucket.sample(&mut rng), None);
126
127        // sample any trusted node
128        let n1 = NodeId::random();
129        assert_eq!(bucket.add(n1.clone()), true);
130        assert_eq!(bucket.sample(&mut rng), Some(n1.clone()));
131    }
132}
network/ip/bucket.rs

network/ip/
bucket.rs
