rqbit/crates/dht/src/routing_table.rs

use std::{
    net::SocketAddr,
    time::{Duration, Instant},
};

use librqbit_core::id20::Id20;
use log::debug;
use serde::Serialize;
use smallvec::SmallVec;

#[derive(Debug, Clone, Serialize)]
enum BucketTreeNodeData {
    // TODO: maybe replace that with SmallVec<8>?
    Leaf(SmallVec<[RoutingTableNode; 8]>),
    LeftRight(usize, usize),
}

#[derive(Debug, Clone, Serialize)]
struct BucketTreeNode {
    bits: u8,
    #[serde(serialize_with = "crate::utils::serialize_id20")]
    start: Id20,
    #[serde(serialize_with = "crate::utils::serialize_id20")]
    end_inclusive: Id20,
    data: BucketTreeNodeData,
}

#[derive(Debug, Clone, Serialize)]
pub struct BucketTree {
    data: Vec<BucketTreeNode>,
}

pub struct BucketTreeIterator<'a> {
    tree: &'a BucketTree,
    current: std::slice::Iter<'a, RoutingTableNode>,
    queue: Vec<usize>,
}

impl<'a> BucketTreeIterator<'a> {
    fn new(tree: &'a BucketTree) -> Self {
        let mut queue = Vec::new();
        let mut current = 0;
        let current_slice = loop {
            match &tree.data[current].data {
                BucketTreeNodeData::Leaf(nodes) => break nodes.iter(),
                BucketTreeNodeData::LeftRight(left, right) => {
                    queue.push(*right);
                    current = *left;
                }
            }
        };
        BucketTreeIterator {
            tree,
            current: current_slice,
            queue,
        }
    }
}

impl<'a> Iterator for BucketTreeIterator<'a> {
    type Item = &'a RoutingTableNode;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(v) = self.current.next() {
            return Some(v);
        };

        loop {
            let idx = self.queue.pop()?;
            match &self.tree.data[idx].data {
                BucketTreeNodeData::Leaf(nodes) => {
                    self.current = nodes.iter();
                    match self.current.next() {
                        Some(v) => return Some(v),
                        None => continue,
                    }
                }
                BucketTreeNodeData::LeftRight(left, right) => {
                    self.queue.push(*right);
                    self.queue.push(*left);
                    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
    };
    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 {
        let mut data = Vec::with_capacity(64);
        data.push(BucketTreeNode {
            bits: 160,
            start: Id20([0u8; 20]),
            end_inclusive: Id20([0xff; 20]),
            data: BucketTreeNodeData::Leaf(SmallVec::with_capacity(8)),
        });
        BucketTree { data }
    }
    pub fn iter(&self) -> BucketTreeIterator<'_> {
        BucketTreeIterator::new(self)
    }

    pub fn get_mut(&mut self, id: &Id20) -> Option<&mut RoutingTableNode> {
        let mut idx = 0;
        loop {
            let node = &self.data[idx];
            if !(*id >= node.start && *id <= node.end_inclusive) {
                return None;
            };
            match &node.data {
                BucketTreeNodeData::Leaf(_) => {
                    // re-borrow mutably
                    if let BucketTreeNodeData::Leaf(nodes) = &mut self.data[idx].data {
                        return nodes.iter_mut().find(|b| b.id == *id);
                    }
                    unreachable!()
                }
                BucketTreeNodeData::LeftRight(left_idx, right_idx) => {
                    let left_idx = *left_idx;
                    let right_idx = *right_idx;
                    let left = &self.data[left_idx];
                    if *id >= left.start && *id <= left.end_inclusive {
                        idx = left_idx;
                        continue;
                    };
                    idx = right_idx;
                }
            }
        }
    }

    pub fn add_node(&mut self, self_id: &Id20, id: Id20, addr: SocketAddr) -> InsertResult {
        let mut current = 0;
        loop {
            let node = &self.data[current];
            debug_assert!(id >= node.start && id <= node.end_inclusive);
            match &node.data {
                BucketTreeNodeData::Leaf(_) => {
                    return self.insert_into_leaf(current, self_id, id, addr);
                }
                BucketTreeNodeData::LeftRight(left_idx, right_idx) => {
                    let left = &self.data[*left_idx];
                    if id <= left.end_inclusive {
                        current = *left_idx;
                        continue;
                    }
                    current = *right_idx;
                }
            }
        }
    }
    fn insert_into_leaf(
        &mut self,
        mut idx: usize,
        self_id: &Id20,
        id: Id20,
        addr: SocketAddr,
    ) -> InsertResult {
        loop {
            let leaf = &mut self.data[idx];
            let nodes = match &mut leaf.data {
                BucketTreeNodeData::Leaf(nodes) => nodes,
                BucketTreeNodeData::LeftRight(_, _) => unreachable!(),
            };
            // if already found, quit
            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(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);
                nodes.sort_by_key(|n| n.id);
                debug!("replaced bad node {:?}", new_node);
                return InsertResult::ReplacedBad(new_node);
            }

            // if our id is not inside, don't bother.
            if *self_id < leaf.start || *self_id > leaf.end_inclusive {
                return InsertResult::Ignored;
            }

            // Split
            let ((ls, le), (rs, re)) =
                compute_split_start_end(leaf.start, leaf.end_inclusive, leaf.bits);
            let (mut ld, mut rd) = (SmallVec::with_capacity(8), SmallVec::with_capacity(8));
            for node in nodes.drain(0..) {
                if node.id < rs {
                    ld.push(node);
                } else {
                    rd.push(node)
                }
            }

            let left = BucketTreeNode {
                bits: leaf.bits - 1,
                start: ls,
                end_inclusive: le,
                data: BucketTreeNodeData::Leaf(ld),
            };
            let right = BucketTreeNode {
                bits: leaf.bits - 1,
                start: rs,
                end_inclusive: re,
                data: BucketTreeNodeData::Leaf(rd),
            };

            let left_idx = {
                let l = self.data.len();
                self.data.push(left);
                l
            };
            let right_idx = {
                let l = self.data.len();
                self.data.push(right);
                l
            };

            self.data[idx].data = BucketTreeNodeData::LeftRight(left_idx, right_idx);
            if id < rs {
                idx = left_idx
            } else {
                idx = right_idx
            }
        }
    }
}

impl Default for BucketTree {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct RoutingTableNode {
    #[serde(serialize_with = "crate::utils::serialize_id20")]
    id: Id20,
    addr: SocketAddr,
    #[serde(skip)]
    last_request: Option<Instant>,
    #[serde(skip)]
    last_response: Option<Instant>,
    #[serde(skip)]
    outstanding_queries_in_a_row: usize,
}

pub enum NodeStatus {
    Good,
    Questionable,
    Bad,
    Unknown,
}

impl RoutingTableNode {
    pub fn id(&self) -> Id20 {
        self.id
    }
    pub fn addr(&self) -> SocketAddr {
        self.addr
    }
    pub fn status(&self) -> NodeStatus {
        // TODO: this is just a stub with simpler logic
        let last_request = match self.last_request {
            Some(v) => v,
            None => return NodeStatus::Unknown,
        };
        if self.outstanding_queries_in_a_row > 0 && last_request.elapsed() > Duration::from_secs(10)
        {
            return NodeStatus::Bad;
        }
        if self.last_response.is_some() {
            return NodeStatus::Good;
        }
        NodeStatus::Questionable
    }

    pub fn mark_outgoing_request(&mut self) {
        self.last_request = Some(Instant::now());
        self.outstanding_queries_in_a_row += 1;
    }

    pub fn mark_response(&mut self) {
        let now = Instant::now();
        self.last_response = Some(now);
        if self.last_request.is_none() {
            self.last_request = Some(now);
        }
        self.outstanding_queries_in_a_row = 0;
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct RoutingTable {
    #[serde(serialize_with = "crate::utils::serialize_id20")]
    id: Id20,
    size: usize,
    buckets: BucketTree,
}

impl RoutingTable {
    pub fn new(id: Id20) -> Self {
        Self {
            id,
            buckets: BucketTree::new(),
            size: 0,
        }
    }
    pub fn len(&self) -> usize {
        self.size
    }
    pub fn sorted_by_distance_from(&self, id: Id20) -> Vec<&RoutingTableNode> {
        let mut result = Vec::with_capacity(self.size);
        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) -> 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
    }
    pub fn mark_outgoing_request(&mut self, id: &Id20) -> bool {
        let r = match self.buckets.get_mut(id) {
            Some(r) => r,
            None => return false,
        };
        r.mark_outgoing_request();
        true
    }

    pub fn mark_response(&mut self, id: &Id20) -> bool {
        let r = match self.buckets.get_mut(id) {
            Some(r) => r,
            None => return false,
        };
        r.mark_response();
        true
    }
}

#[cfg(test)]
mod tests {
    use std::net::SocketAddrV4;

    use librqbit_core::id20::Id20;
    use rand::Rng;

    use crate::routing_table::compute_split_start_end;

    use super::RoutingTable;

    #[test]
    fn compute_split_start_end_root() {
        let start = Id20([0u8; 20]);
        let end = Id20([0xff; 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
                )
            )
        )
    }

    #[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));
            rtable.add_node(other_id, addr);
        }
        dbg!(&rtable);
        assert_eq!(rtable.sorted_by_distance_from(my_id).len(), rtable.size);
    }
}