Routing table sort of works

Cargo.lock

@@ -283,6 +283,7 @@ dependencies = [
  "log",
  "parking_lot",
  "pretty_env_logger",
+ "rand 0.8.4",
  "serde",
  "tokio",
  "tokio-stream",

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"}
 

crates/dht/src/routing_table.rs

@@ -1,14 +1,12 @@
-use std::{
+use std::{net::SocketAddr, time::Instant};
-    collections::BTreeMap,
-    net::SocketAddr,
-    time::{Duration, Instant},
-};
 
+#[derive(Debug)]
 enum BucketTreeNode {
     Leaf(Vec<RoutingTableNode>),
     LeftRight(Box<BucketTree>, Box<BucketTree>),
 }
 
+#[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() {
+        if nodes.iter().any(|r| r.id == id) {
-            return;
+            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 {
+            nodes.push(new_node);
-                id,
+            nodes.sort_by_key(|n| n.id);
-                addr,
+            return InsertResult::Added;
-                last_request: None,
+        }
-                last_response: None,
+
-                outstanding_queries_in_a_row: 0,
+        // Try replace a bad node
-            });
+        if let Some(bad_node) = nodes
-            return;
+            .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 {
+    pub fn add_node(&mut self, id: Id20, addr: SocketAddr) -> InsertResult {
-        todo!()
+        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);
+    }
 }