ckb_network/peer_store/
addr_manager.rs

1//! Address manager
2use crate::peer_store::{base_addr, types::AddrInfo};
3use p2p::{multiaddr::Multiaddr, utils::multiaddr_to_socketaddr};
4use rand::Rng;
5use std::collections::{HashMap, HashSet};
6
7/// Address manager
8#[derive(Default)]
9pub struct AddrManager {
10    next_id: u64,
11    addr_to_id: HashMap<Multiaddr, u64>,
12    id_to_info: HashMap<u64, AddrInfo>,
13    random_ids: Vec<u64>,
14}
15
16impl AddrManager {
17    /// Add an address information to address manager
18    pub fn add(&mut self, mut addr_info: AddrInfo) {
19        if let Some(&id) = self.addr_to_id.get(&addr_info.addr) {
20            let (exist_last_connected_at_ms, random_id_pos) = {
21                let info = self.id_to_info.get(&id).expect("must exists");
22                (info.last_connected_at_ms, info.random_id_pos)
23            };
24            // Get time earlier than record time, return directly
25            if addr_info.last_connected_at_ms >= exist_last_connected_at_ms {
26                addr_info.random_id_pos = random_id_pos;
27                self.id_to_info.insert(id, addr_info);
28            }
29            return;
30        }
31
32        let id = self.next_id;
33        self.addr_to_id.insert(addr_info.addr.clone(), id);
34        addr_info.random_id_pos = self.random_ids.len();
35        self.id_to_info.insert(id, addr_info);
36        self.random_ids.push(id);
37        self.next_id += 1;
38    }
39
40    /// Randomly return addrs that worth to try or connect.
41    pub fn fetch_random<F>(&mut self, count: usize, filter: F) -> Vec<AddrInfo>
42    where
43        F: Fn(&AddrInfo) -> bool,
44    {
45        let mut duplicate_ips = HashSet::new();
46        let mut addr_infos = Vec::with_capacity(count);
47        let mut rng = rand::thread_rng();
48        let now_ms = ckb_systemtime::unix_time_as_millis();
49        for i in 0..self.random_ids.len() {
50            // reuse the for loop to shuffle random ids
51            // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
52            let j = rng.gen_range(i..self.random_ids.len());
53            self.swap_random_id(j, i);
54            let addr_info: AddrInfo = self.id_to_info[&self.random_ids[i]].to_owned();
55            match multiaddr_to_socketaddr(&addr_info.addr) {
56                Some(socket_addr) => {
57                    let ip = socket_addr.ip();
58                    let is_unique_ip = !duplicate_ips.contains(&ip);
59                    // A trick to make our tests work
60                    // TODO remove this after fix the network tests.
61                    let is_test_ip = ip.is_unspecified() || ip.is_loopback();
62                    if (is_test_ip || is_unique_ip)
63                        && addr_info.is_connectable(now_ms)
64                        && filter(&addr_info)
65                    {
66                        duplicate_ips.insert(ip);
67                        addr_infos.push(addr_info);
68                    }
69                }
70                None => {
71                    if addr_info.is_connectable(now_ms) && filter(&addr_info) {
72                        addr_infos.push(addr_info);
73                    }
74                }
75            }
76            if addr_infos.len() == count {
77                break;
78            }
79        }
80        addr_infos
81    }
82
83    /// The count of address in address manager
84    pub fn count(&self) -> usize {
85        self.addr_to_id.len()
86    }
87
88    /// Addresses iterator
89    pub fn addrs_iter(&self) -> impl Iterator<Item = &AddrInfo> {
90        self.id_to_info.values()
91    }
92
93    /// Remove an address by ip and port
94    pub fn remove(&mut self, addr: &Multiaddr) -> Option<AddrInfo> {
95        let base_addr = base_addr(addr);
96        self.addr_to_id.remove(&base_addr).and_then(|id| {
97            let random_id_pos = self.id_to_info.get(&id).expect("exists").random_id_pos;
98            // swap with last index, then remove the last index
99            self.swap_random_id(random_id_pos, self.random_ids.len() - 1);
100            self.random_ids.pop();
101            self.id_to_info.remove(&id)
102        })
103    }
104
105    /// Get an address information by ip and port
106    pub fn get(&self, addr: &Multiaddr) -> Option<&AddrInfo> {
107        let base_addr = base_addr(addr);
108        self.addr_to_id
109            .get(&base_addr)
110            .and_then(|id| self.id_to_info.get(id))
111    }
112
113    /// Get a mutable address information by ip and port
114    pub fn get_mut(&mut self, addr: &Multiaddr) -> Option<&mut AddrInfo> {
115        let base_addr = base_addr(addr);
116        if let Some(id) = self.addr_to_id.get(&base_addr) {
117            self.id_to_info.get_mut(id)
118        } else {
119            None
120        }
121    }
122
123    /// swap random_id i and j,
124    /// this function keep random_id_pos in consistency
125    fn swap_random_id(&mut self, i: usize, j: usize) {
126        if i == j {
127            return;
128        }
129        self.id_to_info
130            .get_mut(&self.random_ids[i])
131            .expect("exists")
132            .random_id_pos = j;
133        self.id_to_info
134            .get_mut(&self.random_ids[j])
135            .expect("exists")
136            .random_id_pos = i;
137        self.random_ids.swap(i, j);
138    }
139}
ckb_network/peer_store/addr_manager.rs

ckb_network/peer_store/
addr_manager.rs
