dynamo-kv-router 1.3.0-dev.1

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Concurrent Radix Tree (compressed trie) implementation for KV cache routing.
//!
//! See `README.md` in this module for structure, removal, split, and concurrency
//! notes.

use std::sync::Arc;

use dashmap::DashMap;
use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
use std::collections::VecDeque;
#[cfg(feature = "bench")]
use std::sync::atomic::{AtomicU64, Ordering};

use super::{
    AnchorRef, AnchorTask, EventKind, EventWarningKind, KvIndexerMetrics, KvRouterError,
    MatchDetails, PreBoundEventCounters, SyncIndexer, WorkerLookupStats, WorkerTask,
};
use crate::cleanup::{CleanupGuard, CleanupState};
use crate::protocols::*;

mod node;
mod types;
use node::*;
use types::*;

mod dump;
mod matches;
mod remove;
mod repair;
mod store;
mod sync_impl;

#[cfg(test)]
mod tests;

/// Thread-safe radix tree (compressed trie) for concurrent KV cache lookups.
pub struct ConcurrentRadixTreeCompressed {
    /// The root of the radix tree. Has an empty edge and only contains children.
    root: SharedNode,

    anchor_nodes: DashMap<ExternalSequenceBlockHash, SharedNode, FxBuildHasher>,
    cleanup: CleanupState,
    #[cfg(feature = "bench")]
    bench_metrics: CrtcBenchMetrics,
}

#[cfg(feature = "bench")]
struct CrtcBenchMetrics {
    node_splits: AtomicU64,
    lookup_repair_scans: AtomicU64,
    lookup_repair_entries: AtomicU64,
}

#[cfg(feature = "bench")]
impl CrtcBenchMetrics {
    fn new() -> Self {
        Self {
            node_splits: AtomicU64::new(0),
            lookup_repair_scans: AtomicU64::new(0),
            lookup_repair_entries: AtomicU64::new(0),
        }
    }
}

#[cfg(test)]
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct EdgeTopologyForTest {
    pub(crate) edge: Vec<u64>,
    pub(crate) children: Vec<EdgeTopologyForTest>,
}

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

// Dropping nodes can cause a cascade of drops that overflow the stack.
// This custom drop uses an iterative approach.
impl Drop for ConcurrentRadixTreeCompressed {
    fn drop(&mut self) {
        self.anchor_nodes.clear();
        let mut stack = self.root.take_children();
        while let Some(node) = stack.pop() {
            stack.extend(node.take_children());
        }
    }
}

impl ConcurrentRadixTreeCompressed {
    pub fn new() -> Self {
        Self {
            root: Arc::new(Node::new()),
            anchor_nodes: DashMap::with_hasher(FxBuildHasher),
            cleanup: CleanupState::new(),
            #[cfg(feature = "bench")]
            bench_metrics: CrtcBenchMetrics::new(),
        }
    }

    #[cfg(test)]
    pub(crate) fn raw_child_edge_count(&self) -> usize {
        let mut queue = VecDeque::from([self.root.clone()]);
        let mut count = 0usize;

        while let Some(node) = queue.pop_front() {
            let children = node.children_snapshot();
            count += children.len();
            queue.extend(children);
        }

        count
    }

    #[cfg(test)]
    pub(crate) fn edge_lengths_for_test(&self) -> Vec<usize> {
        let mut queue = VecDeque::from([self.root.clone()]);
        let mut lengths = Vec::new();

        while let Some(node) = queue.pop_front() {
            let children = node.children_snapshot();
            for child in &children {
                lengths.push(child.edge_len_for_test());
            }
            queue.extend(children);
        }

        lengths.sort_unstable();
        lengths
    }

    #[cfg(test)]
    fn edge_topology_node_for_test(node: &SharedNode) -> EdgeTopologyForTest {
        let mut children: Vec<_> = node
            .children_snapshot()
            .iter()
            .map(Self::edge_topology_node_for_test)
            .collect();
        children.sort_by(|left, right| left.edge.cmp(&right.edge));

        EdgeTopologyForTest {
            edge: node.edge_local_hashes_for_test(),
            children,
        }
    }

    #[cfg(test)]
    pub(crate) fn edge_topology_for_test(&self) -> Vec<EdgeTopologyForTest> {
        let mut children: Vec<_> = self
            .root
            .children_snapshot()
            .iter()
            .map(Self::edge_topology_node_for_test)
            .collect();
        children.sort_by(|left, right| left.edge.cmp(&right.edge));
        children
    }

    fn resolve_anchor_lookup(
        &self,
        lookup: &mut FxHashMap<WorkerWithDpRank, WorkerLookup>,
        worker: WorkerWithDpRank,
        hash: ExternalSequenceBlockHash,
    ) -> Option<SharedNode> {
        let node = self.anchor_nodes.get(&hash)?.clone();
        node.promote_worker_to_full_edge(worker);
        lookup.entry(worker).or_default().insert(hash, node.clone());
        Some(node)
    }

    fn is_anchor_node(&self, hash: ExternalSequenceBlockHash, node: &SharedNode) -> bool {
        self.anchor_nodes
            .get(&hash)
            .is_some_and(|anchor| Arc::ptr_eq(anchor.value(), node))
    }

    // ------------------------------------------------------------------
    // Split helpers
    // ------------------------------------------------------------------

    /// Apply deferred lookup updates after `Node::split_at`.
    ///
    /// Updates worker lookup maps so entries for blocks that moved to the suffix now
    /// point to the suffix node. Must be called **after** the write guard is dropped.
    fn apply_split_lookup(
        &self,
        lookup: &mut FxHashMap<WorkerWithDpRank, WorkerLookup>,
        split: SplitLookupData,
    ) {
        #[cfg(feature = "bench")]
        self.bench_metrics
            .node_splits
            .fetch_add(1, Ordering::Relaxed);
        for (worker, hashes) in split.suffix.lookup_entries_by_worker() {
            if let Some(wl) = lookup.get_mut(&worker) {
                for hash in hashes {
                    wl.insert(hash, split.suffix.clone());
                }
            }
        }
    }

    fn update_lookup_for_blocks(
        worker_lookup: &mut WorkerLookup,
        blocks: &[KvCacheStoredBlockData],
        node: &SharedNode,
    ) -> bool {
        let mut changed = false;
        for block in blocks {
            match worker_lookup.insert(block.block_hash, node.clone()) {
                Some(existing) if Arc::ptr_eq(&existing, node) => {}
                Some(_) => {
                    changed = true;
                }
                None => {
                    changed = true;
                }
            }
        }
        changed
    }

    // ------------------------------------------------------------------
    // apply_event dispatch
    // ------------------------------------------------------------------

    #[cfg_attr(feature = "profile", inline(never))]
    fn apply_event(
        &self,
        lookup: &mut FxHashMap<WorkerWithDpRank, WorkerLookup>,
        event: RouterEvent,
        counters: Option<&PreBoundEventCounters>,
    ) -> Result<(), KvCacheEventError> {
        let (worker_id, kv_event) = (event.worker_id, event.event);
        let (id, op) = (kv_event.event_id, kv_event.data);
        let worker = WorkerWithDpRank::new(worker_id, kv_event.dp_rank);

        match op {
            KvCacheEventData::Stored(op) => self.apply_stored(lookup, worker, op, id, counters),
            KvCacheEventData::Removed(op) => self.apply_removed(lookup, worker, op, id),
            KvCacheEventData::Cleared => {
                lookup.entry(worker).or_default();
                self.clear_all_blocks(lookup, worker.worker_id);
                Ok(())
            }
        }
    }
}