From 9c37706dbaa1da783556b47d69d45fa70ae98aa8 Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Mon, 12 Jul 2021 14:38:55 +0100 Subject: [PATCH] Routing table sort of works --- Cargo.lock | 1 + crates/dht/Cargo.toml | 1 + crates/dht/src/routing_table.rs | 215 ++++++++++++++++++++++++++++---- 3 files changed, 193 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8bda825..4fbb82d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -283,6 +283,7 @@ dependencies = [ "log", "parking_lot", "pretty_env_logger", + "rand 0.8.4", "serde", "tokio", "tokio-stream", diff --git a/crates/dht/Cargo.toml b/crates/dht/Cargo.toml index 9c578c8..7b12095 100644 --- a/crates/dht/Cargo.toml +++ b/crates/dht/Cargo.toml @@ -17,6 +17,7 @@ parking_lot = "0.11" log = "0.4" pretty_env_logger = "0.4" futures = "0.3" +rand = "0.8" librqbit_core = {path="../librqbit_core"} diff --git a/crates/dht/src/routing_table.rs b/crates/dht/src/routing_table.rs index f488e33..05e0957 100644 --- a/crates/dht/src/routing_table.rs +++ b/crates/dht/src/routing_table.rs @@ -1,14 +1,12 @@ -use std::{ - collections::BTreeMap, - net::SocketAddr, - time::{Duration, Instant}, -}; +use std::{net::SocketAddr, time::Instant}; +#[derive(Debug)] enum BucketTreeNode { Leaf(Vec), LeftRight(Box, Box), } +#[derive(Debug)] pub struct BucketTree { bits: u8, start: Id20, @@ -81,28 +79,63 @@ fn compute_split_start_end( c.set_bit(changing_bit, true); c }; + debug_assert!( + start < new_left_end, + "expected start({:?}) < new_left_end({:?}); start={:?}, end={:?}, bits={}", + start, + new_left_end, + start, + end_inclusive, + bits + ); + debug_assert!( + new_left_end < new_right_start, + "expected new_left_end({:?}) < new_right_start({:?}); start={:?}, end={:?}, bits={}", + new_left_end, + new_right_start, + start, + end_inclusive, + bits + ); + debug_assert!( + new_right_start < end_inclusive, + "expected new_right_start({:?}) < end_inclusive({:?}); start={:?}, end={:?}, bits={}", + new_right_start, + end_inclusive, + start, + end_inclusive, + bits + ); ((start, new_left_end), (new_right_start, end_inclusive)) } +#[derive(Debug)] +pub enum InsertResult { + WasExisting, + ReplacedBad(RoutingTableNode), + Added, + Ignored, +} + impl BucketTree { pub fn new() -> Self { BucketTree { bits: 160, start: Id20([0u8; 20]), end_inclusive: Id20([0xff; 20]), - data: BucketTreeNode::Leaf(Vec::new()), + data: BucketTreeNode::Leaf(Vec::with_capacity(8)), } } pub fn iter(&self) -> BucketTreeNodeIterator<'_> { BucketTreeNodeIterator::new(self) } - pub fn add_node(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) { + pub fn add_node(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) -> InsertResult { 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) + return tree.insert_into_leaf(self_id, id, addr); } BucketTreeNode::LeftRight(left, right) => { if id >= right.start { @@ -116,33 +149,76 @@ impl BucketTree { } } } - fn insert_into_leaf(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) { + fn insert_into_leaf(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) -> InsertResult { 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.iter().any(|r| r.id == id) { + return InsertResult::WasExisting; } + let mut new_node = RoutingTableNode { + id, + addr, + last_request: None, + last_response: None, + outstanding_queries_in_a_row: 0, + }; + if nodes.len() < 8 { - nodes.push(RoutingTableNode { - id, - addr, - last_request: None, - last_response: None, - outstanding_queries_in_a_row: 0, - }); - return; + nodes.push(new_node); + nodes.sort_by_key(|n| n.id); + return InsertResult::Added; + } + + // Try replace a bad node + if let Some(bad_node) = nodes + .iter_mut() + .find(|r| matches!(r.status(), NodeStatus::Bad)) + { + std::mem::swap(bad_node, &mut new_node); + return InsertResult::ReplacedBad(new_node); } // if our id is not inside, don't bother. if *self_id < self.start || *self_id > self.end_inclusive { - return; + return InsertResult::Ignored; } - todo!() + // Split + let ((ls, le), (rs, re)) = + compute_split_start_end(self.start, self.end_inclusive, self.bits); + let (mut ld, mut rd) = (Vec::with_capacity(8), Vec::with_capacity(8)); + for node in nodes.drain(0..) { + if node.id < rs { + ld.push(node); + } else { + rd.push(node) + } + } + let mut left = BucketTree { + bits: self.bits - 1, + start: ls, + end_inclusive: le, + data: BucketTreeNode::Leaf(ld), + }; + let mut right = BucketTree { + bits: self.bits - 1, + start: rs, + end_inclusive: re, + data: BucketTreeNode::Leaf(rd), + }; + + let result = if id < rs { + left.add_node(self_id, id, addr) + } else { + right.add_node(self_id, id, addr) + }; + + self.data = BucketTreeNode::LeftRight(Box::new(left), Box::new(right)); + result } } @@ -154,6 +230,7 @@ impl Default for BucketTree { use crate::id20::Id20; +#[derive(Debug)] pub struct RoutingTableNode { id: Id20, addr: SocketAddr, @@ -189,6 +266,7 @@ impl RoutingTableNode { } } +#[derive(Debug)] pub struct RoutingTable { id: Id20, size: usize, @@ -211,19 +289,35 @@ impl RoutingTable { result.sort_by_key(|n| id.distance(&n.id)); result } - pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> bool { - todo!() + pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> InsertResult { + let res = self.buckets.add_node(&self.id, id, addr); + let replaced = match &res { + InsertResult::WasExisting => false, + InsertResult::ReplacedBad(_) => true, + InsertResult::Added => true, + InsertResult::Ignored => false, + }; + if replaced { + self.size += 1; + } + res } } #[cfg(test)] mod tests { + use std::net::SocketAddrV4; + + use rand::Rng; + use crate::{id20::Id20, routing_table::compute_split_start_end}; + use super::RoutingTable; + #[test] fn compute_split_start_end_root() { let start = Id20([0u8; 20]); - let end = Id20([0xffu8; 20]); + let end = Id20([0xff; 20]); assert_eq!( compute_split_start_end(start, end, 160), ( @@ -244,4 +338,77 @@ mod tests { ) ) } + + #[test] + fn compute_split_start_end_second_split() { + let start = Id20([ + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ]); + let end = Id20([0xff; 20]); + assert_eq!( + compute_split_start_end(start, end, 159), + ( + ( + start, + Id20([ + 0xbf, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff + ]) + ), + ( + Id20([ + 0xc0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 + ]), + end + ) + ) + ) + } + + #[test] + fn compute_split_start_end_3() { + let start = Id20([ + 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ]); + let end = Id20([0xff; 20]); + assert_eq!( + compute_split_start_end(start, end, 159), + ( + ( + start, + Id20([ + 0xbf, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff + ]) + ), + ( + Id20([ + 0xc0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 + ]), + end + ) + ) + ) + } + + fn random_id_20() -> Id20 { + let mut id20 = [0u8; 20]; + rand::thread_rng().fill(&mut id20); + Id20(id20) + } + + #[test] + fn simulate_tree() { + let my_id = random_id_20(); + let mut rtable = RoutingTable::new(my_id); + for i in 0..u16::MAX { + let other_id = random_id_20(); + let addr = std::net::SocketAddr::V4(SocketAddrV4::new("0.0.0.0".parse().unwrap(), i)); + } + dbg!(rtable); + } }