From: Matt Corallo Date: Sat, 18 May 2019 20:03:09 +0000 (-0400) Subject: Track service bits X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=commitdiff_plain;h=a5e03b3a45858933dd0a6a9782f78e5270da6e5f;p=dnsseed-rust Track service bits --- diff --git a/src/datastore.rs b/src/datastore.rs index 1013a9e..4bab7ef 100644 --- a/src/datastore.rs +++ b/src/datastore.rs @@ -1,5 +1,5 @@ use std::{cmp, mem}; -use std::collections::{HashMap, hash_map}; +use std::collections::{HashSet, HashMap, hash_map}; use std::sync::RwLock; use std::net::SocketAddr; use std::time::{Duration, Instant}; @@ -38,11 +38,32 @@ pub enum StringSetting { SubverRegex, } +struct Node { + state: AddressState, + last_services: u64, + last_update: Instant, +} + struct Nodes { - good_node_services: HashMap>, - nodes_to_state: HashMap, + good_node_services: HashMap>, + nodes_to_state: HashMap, state_next_scan: HashMap>, } +struct NodesMutRef<'a> { + good_node_services: &'a mut HashMap>, + nodes_to_state: &'a mut HashMap, + state_next_scan: &'a mut HashMap>, + +} +impl Nodes { + fn borrow_mut<'a>(&'a mut self) -> NodesMutRef<'a> { + NodesMutRef { + good_node_services: &mut self.good_node_services, + nodes_to_state: &mut self.nodes_to_state, + state_next_scan: &mut self.state_next_scan, + } + } +} pub struct Store { u64_settings: RwLock>, @@ -82,7 +103,7 @@ impl Store { state_vecs.insert(AddressState::WasGood, Vec::new()); let mut good_node_services = HashMap::with_capacity(64); for i in 0..64 { - good_node_services.insert(i, Vec::new()); + good_node_services.insert(i, HashSet::new()); } Store { u64_settings: RwLock::new(u64s), @@ -114,7 +135,11 @@ impl Store { if let Ok(socketaddr) = addr.socket_addr() { match nodes.nodes_to_state.entry(socketaddr.clone()) { hash_map::Entry::Vacant(e) => { - e.insert(AddressState::Untested); + e.insert(Node { + state: AddressState::Untested, + last_services: 0, + last_update: cur_time, + }); nodes.state_next_scan.get_mut(&AddressState::Untested).unwrap().push((cur_time, socketaddr)); }, hash_map::Entry::Occupied(_) => {}, @@ -126,17 +151,31 @@ impl Store { } pub fn set_node_state(&self, addr: SocketAddr, state: AddressState, services: u64) { - let mut nodes = self.nodes.write().unwrap(); + let mut nodes_lock = self.nodes.write().unwrap(); + let nodes = nodes_lock.borrow_mut(); let state_ref = nodes.nodes_to_state.get_mut(&addr).unwrap(); - if *state_ref == AddressState::Good && state != AddressState::Good { - *state_ref = AddressState::WasGood; - nodes.state_next_scan.get_mut(&AddressState::WasGood).unwrap().push((Instant::now(), addr)); + state_ref.last_update = Instant::now(); + if state_ref.state == AddressState::Good && state != AddressState::Good { + state_ref.state = AddressState::WasGood; + for i in 0..64 { + if state_ref.last_services & (1 << i) != 0 { + nodes.good_node_services.get_mut(&i).unwrap().remove(&addr); + } + } + state_ref.last_services = 0; + nodes.state_next_scan.get_mut(&AddressState::WasGood).unwrap().push((state_ref.last_update, addr)); } else { - *state_ref = state; - nodes.state_next_scan.get_mut(&state).unwrap().push((Instant::now(), addr)); - } - if state == AddressState::Good { - + state_ref.state = state; + if state == AddressState::Good { + for i in 0..64 { + if services & (1 << i) != 0 && state_ref.last_services & (1 << i) == 0 { + nodes.good_node_services.get_mut(&i).unwrap().insert(addr); + } else if services & (1 << i) == 0 && state_ref.last_services & (1 << i) != 0 { + nodes.good_node_services.get_mut(&i).unwrap().remove(&addr); + } + } + } + nodes.state_next_scan.get_mut(&state).unwrap().push((state_ref.last_update, addr)); } }