From 44a72b10880f7011a483ecdea9286b99853e63ad Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Mon, 12 Jul 2021 13:59:58 +0100 Subject: [PATCH] Saving --- crates/dht/src/id20.rs | 29 +++++ crates/dht/src/main.rs | 50 ++++---- crates/dht/src/routing_table.rs | 207 ++++++++++++++++++++++++++++---- 3 files changed, 239 insertions(+), 47 deletions(-) diff --git a/crates/dht/src/id20.rs b/crates/dht/src/id20.rs index 61befdc..37c8cd0 100644 --- a/crates/dht/src/id20.rs +++ b/crates/dht/src/id20.rs @@ -67,6 +67,20 @@ impl Id20 { } Id20(xor) } + pub fn set_bit(&mut self, bit: u8, value: bool) { + let n = &mut self.0[(bit / 8) as usize]; + if value { + *n |= 1 << (7 - bit % 8) + } else { + let mask = !(1 << (7 - bit % 8)); + *n &= mask; + } + } + pub fn set_bits_range(&mut self, r: std::ops::Range, value: bool) { + for bit in r { + self.set_bit(bit, value) + } + } } impl Ord for Id20 { @@ -87,3 +101,18 @@ impl PartialOrd for Id20 { Some(self.cmp(other)) } } + +#[cfg(test)] +mod tests { + use super::Id20; + + #[test] + fn test_set_bit_range() { + let mut id = Id20([0u8; 20]); + id.set_bits_range(9..17, true); + assert_eq!( + id, + Id20([0, 127, 128, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) + ) + } +} diff --git a/crates/dht/src/main.rs b/crates/dht/src/main.rs index d5f8fb6..ae6606c 100644 --- a/crates/dht/src/main.rs +++ b/crates/dht/src/main.rs @@ -130,12 +130,12 @@ impl DhtState { let nodes = response .nodes .ok_or_else(|| anyhow::anyhow!("expected nodes for find_node requests"))?; - let pn = match (response.nodes, response.values) { - (Some(nodes), None) => PeersOrNodes::Nodes(nodes), - (None, Some(peers)) => PeersOrNodes::Peers(peers), - _ => anyhow::bail!("expected nodes or values to be set in find_peers response"), - }; - self.on_found_peers_or_nodes(id, pn) + // let pn = match (response.nodes, response.values) { + // (Some(nodes), None) => PeersOrNodes::Nodes(nodes), + // (None, Some(peers)) => PeersOrNodes::Peers(peers), + // _ => anyhow::bail!("expected nodes or values to be set in find_peers response"), + // }; + // self.on_found_peers_or_nodes(id, pn) } }; Ok(()) @@ -222,31 +222,31 @@ impl DhtWorker { // bootstrap for addr in bootstrap_addrs { for addr in tokio::net::lookup_host(addr).await.unwrap() { - let msg = MessageKind::FindNodeRequest(FindNodeRequest { - id: self.peer_id, - target: self.peer_id, - }); - in_tx.send((msg, addr)).await.unwrap(); + // let msg = MessageKind::FindNodeRequest(FindNodeRequest { + // id: self.peer_id, + // target: self.peer_id, + // }); + // in_tx.send((msg, addr)).await.unwrap(); } } }; let mut bootstrap_done = false; - let request_reader = async { - while let Some((request, sender)) = self.request_rx.recv().await { - self.on_request(request, sender) - } - }; + // let request_reader = async { + // while let Some((request, sender)) = self.request_rx.recv().await { + // self.on_request(request, sender) + // } + // }; - tokio::select! { - _ = framer => { - anyhow::bail!("framer quit") - }, - _ = bootstrap, if !bootstrap_done => { - bootstrap_done = true - }, - _ = request_reader => {} - } + // tokio::select! { + // _ = framer => { + // anyhow::bail!("framer quit") + // }, + // _ = bootstrap, if !bootstrap_done => { + // bootstrap_done = true + // }, + // _ = request_reader => {} + // } todo!() } diff --git a/crates/dht/src/routing_table.rs b/crates/dht/src/routing_table.rs index ff51bcd..f488e33 100644 --- a/crates/dht/src/routing_table.rs +++ b/crates/dht/src/routing_table.rs @@ -4,6 +4,154 @@ use std::{ time::{Duration, Instant}, }; +enum BucketTreeNode { + Leaf(Vec), + LeftRight(Box, Box), +} + +pub struct BucketTree { + bits: u8, + start: Id20, + end_inclusive: Id20, + data: BucketTreeNode, +} + +pub struct BucketTreeNodeIterator<'a> { + current: std::slice::Iter<'a, RoutingTableNode>, + queue: Vec<&'a BucketTree>, +} + +impl<'a> BucketTreeNodeIterator<'a> { + fn new(mut tree: &'a BucketTree) -> Self { + let mut queue = Vec::new(); + let current = loop { + match &tree.data { + BucketTreeNode::Leaf(nodes) => break nodes.iter(), + BucketTreeNode::LeftRight(left, right) => { + queue.push(right.as_ref()); + tree = left.as_ref() + } + } + }; + BucketTreeNodeIterator { current, queue } + } +} + +impl<'a> Iterator for BucketTreeNodeIterator<'a> { + type Item = &'a RoutingTableNode; + + fn next(&mut self) -> Option { + if let Some(v) = self.current.next() { + return Some(v); + }; + + loop { + let tree = self.queue.pop()?; + match &tree.data { + BucketTreeNode::Leaf(nodes) => { + self.current = nodes.iter(); + match self.current.next() { + Some(v) => return Some(v), + None => continue, + } + } + BucketTreeNode::LeftRight(left, right) => { + self.queue.push(right.as_ref()); + self.queue.push(left.as_ref()); + continue; + } + } + } + } +} + +fn compute_split_start_end( + start: Id20, + end_inclusive: Id20, + bits: u8, +) -> ((Id20, Id20), (Id20, Id20)) { + let changing_bit = 160 - bits; + let new_left_end = { + let mut c = end_inclusive; + c.set_bit(changing_bit, false); + c + }; + let new_right_start = { + let mut c = start; + c.set_bit(changing_bit, true); + c + }; + ((start, new_left_end), (new_right_start, end_inclusive)) +} + +impl BucketTree { + pub fn new() -> Self { + BucketTree { + bits: 160, + start: Id20([0u8; 20]), + end_inclusive: Id20([0xff; 20]), + data: BucketTreeNode::Leaf(Vec::new()), + } + } + pub fn iter(&self) -> BucketTreeNodeIterator<'_> { + BucketTreeNodeIterator::new(self) + } + pub fn add_node(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) { + let mut tree = self; + loop { + match &mut tree.data { + BucketTreeNode::Leaf(_) => { + assert!(id >= tree.start && id <= tree.end_inclusive); + tree.insert_into_leaf(self_id, id, addr) + } + BucketTreeNode::LeftRight(left, right) => { + if id >= right.start { + // Erase lifetime. + // Safety: this is safe as it's a tree, not a DAG or Graph. + tree = unsafe { &mut *(right.as_mut() as *mut _) }; + continue; + } + tree = unsafe { &mut *(left.as_mut() as *mut _) }; + } + } + } + } + fn insert_into_leaf(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) { + let nodes = match &mut self.data { + BucketTreeNode::Leaf(nodes) => nodes, + BucketTreeNode::LeftRight(_, _) => unreachable!(), + }; + // if already found, quit + if nodes.iter().find(|r| r.id == id).is_some() { + return; + } + + if nodes.len() < 8 { + nodes.push(RoutingTableNode { + id, + addr, + last_request: None, + last_response: None, + outstanding_queries_in_a_row: 0, + }); + return; + } + + // if our id is not inside, don't bother. + if *self_id < self.start || *self_id > self.end_inclusive { + return; + } + + todo!() + } +} + +impl Default for BucketTree { + fn default() -> Self { + Self::new() + } +} + use crate::id20::Id20; pub struct RoutingTableNode { @@ -41,44 +189,59 @@ impl RoutingTableNode { } } -struct Bucket { - bits: u8, - nodes: Vec, - end: Id20, -} - pub struct RoutingTable { id: Id20, size: usize, - buckets: BTreeMap, + buckets: BucketTree, } impl RoutingTable { pub fn new(id: Id20) -> Self { - let initial_bucket = Id20([0u8; 20]); - let mut buckets = BTreeMap::new(); - buckets.insert( - initial_bucket, - Bucket { - bits: 160, - nodes: Vec::new(), - }, - ); Self { id, - buckets, + buckets: BucketTree::new(), size: 0, } } pub fn sorted_by_distance_from(&self, id: Id20) -> Vec<&RoutingTableNode> { let mut result = Vec::with_capacity(self.size); - for bucket in self.buckets.values() { - for node in bucket.nodes.iter() { - result.push(node); - } + for node in self.buckets.iter() { + result.push(node); } result.sort_by_key(|n| id.distance(&n.id)); result } - pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> bool {} + pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> bool { + todo!() + } +} + +#[cfg(test)] +mod tests { + use crate::{id20::Id20, routing_table::compute_split_start_end}; + + #[test] + fn compute_split_start_end_root() { + let start = Id20([0u8; 20]); + let end = Id20([0xffu8; 20]); + assert_eq!( + compute_split_start_end(start, end, 160), + ( + ( + start, + Id20([ + 0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff + ]) + ), + ( + Id20([ + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 + ]), + end + ) + ) + ) + } }