From f095e5dfccda4d6ef09acce0a3e4f60d212626a7 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Thu, 21 Nov 2024 11:46:39 -0600 Subject: [PATCH] update --- bench/scan.rs | 3 - src/base_node.rs | 2 +- src/lib.rs | 7 +- src/stats.rs | 235 +++++++++++++++++++++++++++++----------------- src/tests/tree.rs | 12 ++- src/tree.rs | 122 ++++++++++++++++++------ 6 files changed, 257 insertions(+), 124 deletions(-) diff --git a/bench/scan.rs b/bench/scan.rs index 06965f4..f73e4b6 100644 --- a/bench/scan.rs +++ b/bench/scan.rs @@ -36,9 +36,6 @@ impl ShumaiBench for TestBench { unique_values.insert(k); } - #[cfg(feature = "stats")] - println!("{}", self.index.stats()); - Some(serde_json::json!({ "unique_values": unique_values.len(), })) diff --git a/src/base_node.rs b/src/base_node.rs index b9fcfa8..8dbffe7 100644 --- a/src/base_node.rs +++ b/src/base_node.rs @@ -299,7 +299,7 @@ impl BaseNode { (version & 0b10) == 0b10 } - pub(crate) fn get_count(&self) -> usize { + pub(crate) fn value_count(&self) -> usize { self.meta.count as usize } diff --git a/src/lib.rs b/src/lib.rs index a652a64..1a7623a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,7 +16,6 @@ mod utils; mod range_scan; -#[cfg(feature = "stats")] mod stats; #[cfg(test)] @@ -72,7 +71,7 @@ impl Allocator for DefaultAllocator { pub struct U64Congee< V: Clone + From + Into, - A: Allocator + Clone + 'static = DefaultAllocator, + A: Allocator + Clone + Send + 'static = DefaultAllocator, > { inner: RawCongee<8, A>, pt_val: PhantomData, @@ -130,7 +129,7 @@ impl + Into> U64Congee { pub struct Congee< K: Clone + From, V: Clone + From, - A: Allocator + Clone + 'static = DefaultAllocator, + A: Allocator + Clone + Send + 'static = DefaultAllocator, > where usize: From, usize: From, @@ -391,8 +390,6 @@ where } /// Display the internal node statistics - #[cfg(feature = "stats")] - #[cfg_attr(docsrs, doc(cfg(feature = "stats")))] pub fn stats(&self) -> stats::NodeStats { self.inner.stats() } diff --git a/src/stats.rs b/src/stats.rs index 339869d..c425da4 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -1,54 +1,74 @@ -use std::{fmt::Display, ptr::NonNull}; +use std::{collections::HashMap, fmt::Display, ptr::NonNull}; use crate::{ base_node::{BaseNode, NodeType}, - node_256::Node256, - node_ptr::PtrType, - tree::RawCongee, + tree::{CongeeVisitor, RawCongee}, Allocator, }; -#[derive(Default, Debug, serde::Serialize)] -pub struct NodeStats(Vec); +#[cfg_attr(feature = "stats", derive(serde::Serialize))] +#[derive(Default, Debug, Clone)] +pub struct NodeStats { + levels: HashMap, +} impl Display for NodeStats { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - fn calc_load_factor(n: (usize, usize), scale: usize) -> f64 { - if n.0 == 0 { + fn calc_load_factor(n: &NodeInfo, scale: usize) -> f64 { + if n.node_count == 0 { return 0.0; } - (n.1 as f64) / (n.0 as f64 * scale as f64) + (n.value_count as f64) / (n.node_count as f64 * scale as f64) } - let mut total_node = 0; + let mut levels = self.levels.values().collect::>(); + levels.sort_by_key(|l| l.level); + + let mut node_count = 0; let mut total_f = 0.0; let mut memory_size = 0; + let mut value_count = 0; - for l in self.0.iter() { - total_f += l.n4.1 as f64 / 4.0; - total_f += l.n16.1 as f64 / 16.0; - total_f += l.n48.1 as f64 / 48.0; - total_f += l.n256.1 as f64 / 256.0; + for l in levels.iter() { + total_f += l.n4.value_count as f64 / 4.0; + total_f += l.n16.value_count as f64 / 16.0; + total_f += l.n48.value_count as f64 / 48.0; + total_f += l.n256.value_count as f64 / 256.0; - total_node += l.total_nodes(); + node_count += l.node_count(); memory_size += l.memory_size(); + value_count += l.value_count(); writeln!( f, "Level: {} --- || N4: {:8}, {:8.2} || N16: {:8}, {:8.2} || N48: {:8}, {:8.2} || N256: {:8}, {:8.2} ||", l.level, - l.n4.0, - calc_load_factor(l.n4, 4), - l.n16.0, - calc_load_factor(l.n16, 16), - l.n48.0, - calc_load_factor(l.n48, 48), - l.n256.0, - calc_load_factor(l.n256, 256), + l.n4.node_count, + calc_load_factor(&l.n4, 4), + l.n16.node_count, + calc_load_factor(&l.n16, 16), + l.n48.node_count, + calc_load_factor(&l.n48, 48), + l.n256.node_count, + calc_load_factor(&l.n256, 256), )?; } - let load_factor = total_f / (total_node as f64); + writeln!( + f, + "Overall node count: {}, value count: {}", + node_count, value_count + )?; + + let last_level = levels.last().unwrap(); + writeln!( + f, + "Last level node: {}, value count: {}", + last_level.node_count(), + last_level.value_count(), + )?; + + let load_factor = total_f / (node_count as f64); if load_factor < 0.5 { writeln!(f, "Load factor: {load_factor:.2} (too low)")?; } else { @@ -61,88 +81,135 @@ impl Display for NodeStats { } } -#[derive(Debug, serde::Serialize, Clone)] +#[cfg_attr(feature = "stats", derive(serde::Serialize))] +#[derive(Debug, Clone, Default)] +struct NodeInfo { + node_count: usize, + value_count: usize, +} + +#[cfg_attr(feature = "stats", derive(serde::Serialize))] +#[derive(Debug, Clone)] pub struct LevelStats { level: usize, - n4: (usize, usize), // (node count, leaf count) - n16: (usize, usize), - n48: (usize, usize), - n256: (usize, usize), + n4: NodeInfo, // (node count, leaf count) + n16: NodeInfo, + n48: NodeInfo, + n256: NodeInfo, } impl LevelStats { fn new_level(level: usize) -> Self { Self { level, - n4: (0, 0), - n16: (0, 0), - n48: (0, 0), - n256: (0, 0), + n4: NodeInfo::default(), + n16: NodeInfo::default(), + n48: NodeInfo::default(), + n256: NodeInfo::default(), } } fn memory_size(&self) -> usize { - self.n4.0 * NodeType::N4.node_layout().size() - + self.n16.0 * NodeType::N16.node_layout().size() - + self.n48.0 * NodeType::N48.node_layout().size() - + self.n256.0 * NodeType::N256.node_layout().size() + self.n4.node_count * NodeType::N4.node_layout().size() + + self.n16.node_count * NodeType::N16.node_layout().size() + + self.n48.node_count * NodeType::N48.node_layout().size() + + self.n256.node_count * NodeType::N256.node_layout().size() + } + + fn node_count(&self) -> usize { + self.n4.node_count + self.n16.node_count + self.n48.node_count + self.n256.node_count } - fn total_nodes(&self) -> usize { - self.n4.0 + self.n16.0 + self.n48.0 + self.n256.0 + fn value_count(&self) -> usize { + self.n4.value_count + self.n16.value_count + self.n48.value_count + self.n256.value_count } } -impl RawCongee { - /// Returns the node stats for the tree. - pub fn stats(&self) -> NodeStats { - let mut node_stats = NodeStats::default(); +struct StatsVisitor { + node_stats: NodeStats, +} - let mut sub_nodes = vec![(0, 0, unsafe { - std::mem::transmute::, NonNull>(self.root) - })]; +impl CongeeVisitor for StatsVisitor { + fn pre_visit_sub_node(&mut self, node: NonNull) { + let node = BaseNode::read_lock(node).unwrap(); + let tree_level = node.as_ref().prefix().len(); - while let Some((level, key_level, node)) = sub_nodes.pop() { - let node = BaseNode::read_lock(node).unwrap(); + if !self.node_stats.levels.contains_key(&tree_level) { + self.node_stats + .levels + .insert(tree_level, LevelStats::new_level(tree_level)); + } - if node_stats.0.len() <= level { - node_stats.0.push(LevelStats::new_level(level)); + match node.as_ref().get_type() { + crate::base_node::NodeType::N4 => { + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n4 + .node_count += 1; + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n4 + .value_count += node.as_ref().value_count(); } - - match node.as_ref().get_type() { - crate::base_node::NodeType::N4 => { - node_stats.0[level].n4.0 += 1; - node_stats.0[level].n4.1 += node.as_ref().get_count(); - } - crate::base_node::NodeType::N16 => { - node_stats.0[level].n16.0 += 1; - node_stats.0[level].n16.1 += node.as_ref().get_count(); - } - crate::base_node::NodeType::N48 => { - node_stats.0[level].n48.0 += 1; - node_stats.0[level].n48.1 += node.as_ref().get_count(); - } - crate::base_node::NodeType::N256 => { - node_stats.0[level].n256.0 += 1; - node_stats.0[level].n256.1 += node.as_ref().get_count(); - } + crate::base_node::NodeType::N16 => { + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n16 + .node_count += 1; + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n16 + .value_count += node.as_ref().value_count(); } - - let children = node.as_ref().get_children(0, 255); - for (_k, n) in children { - match n.downcast::(level) { - PtrType::Payload(_) => {} - PtrType::SubNode(sub_node) => { - let child_node = BaseNode::read_lock(sub_node).unwrap(); - sub_nodes.push(( - level + 1, - key_level + 1 + child_node.as_ref().prefix().len(), - sub_node, - )); - } - } + crate::base_node::NodeType::N48 => { + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n48 + .node_count += 1; + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n48 + .value_count += node.as_ref().value_count(); + } + crate::base_node::NodeType::N256 => { + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n256 + .node_count += 1; + self.node_stats + .levels + .get_mut(&tree_level) + .unwrap() + .n256 + .value_count += node.as_ref().value_count(); } } - node_stats + } +} + +impl RawCongee { + /// Returns the node stats for the tree. + pub fn stats(&self) -> NodeStats { + let mut visitor = StatsVisitor { + node_stats: NodeStats::default(), + }; + + self.dfs_visitor_slow(&mut visitor).unwrap(); + + return visitor.node_stats; } } diff --git a/src/tests/tree.rs b/src/tests/tree.rs index 4081c97..8a6464f 100644 --- a/src/tests/tree.rs +++ b/src/tests/tree.rs @@ -23,13 +23,14 @@ fn small_insert() { #[test] fn test_sparse_keys() { - let key_cnt = 100_000; + let key_cnt = 30_000; let tree = RawCongee::default(); let mut keys = Vec::::with_capacity(key_cnt); let guard = crossbeam_epoch::pin(); + let mut rng = StdRng::seed_from_u64(12); for _i in 0..key_cnt { - let k = thread_rng().gen::() & 0x7fff_ffff_ffff_ffff; + let k = rng.gen::() & 0x7fff_ffff_ffff_ffff; keys.push(k); let key: [u8; 8] = k.to_be_bytes(); @@ -56,13 +57,12 @@ fn test_sparse_keys() { assert_eq!(v, *i); } - #[cfg(feature = "stats")] println!("{}", tree.stats()); } use rand::prelude::StdRng; use rand::seq::SliceRandom; -use rand::{thread_rng, Rng, SeedableRng}; +use rand::{Rng, SeedableRng}; #[test] fn test_concurrent_insert() { @@ -105,6 +105,8 @@ fn test_concurrent_insert() { let val = tree.get(&key, &guard).unwrap(); assert_eq!(val, *v); } + + assert_eq!(tree.value_count(&guard), key_space.len()); } #[cfg(all(feature = "shuttle", test))] @@ -192,6 +194,8 @@ fn test_concurrent_insert_read() { assert_eq!(val, *v); } + assert_eq!(tree.value_count(&guard), key_space.len()); + drop(guard); drop(tree); } diff --git a/src/tree.rs b/src/tree.rs index 342f026..62cd91e 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -16,14 +16,17 @@ use crate::{ /// Raw interface to the ART tree. /// The `Art` is a wrapper around the `RawArt` that provides a safe interface. -pub(crate) struct RawCongee { +pub(crate) struct RawCongee< + const K_LEN: usize, + A: Allocator + Clone + Send + 'static = DefaultAllocator, +> { pub(crate) root: NonNull, allocator: A, _pt_key: PhantomData<[u8; K_LEN]>, } -unsafe impl Send for RawCongee {} -unsafe impl Sync for RawCongee {} +unsafe impl Send for RawCongee {} +unsafe impl Sync for RawCongee {} impl Default for RawCongee { fn default() -> Self { @@ -31,29 +34,46 @@ impl Default for RawCongee { } } -impl Drop for RawCongee { - fn drop(&mut self) { - let mut sub_nodes = vec![( - unsafe { std::mem::transmute::, NonNull>(self.root) }, - 0, - )]; - - while let Some((node, level)) = sub_nodes.pop() { - let node_lock = BaseNode::read_lock(node).unwrap(); - let children = node_lock.as_ref().get_children(0, 255); - for (_k, n) in children { - match n.downcast::(level) { - PtrType::Payload(_) => {} - PtrType::SubNode(sub_sub_node) => { - let node_lock = BaseNode::read_lock(sub_sub_node).unwrap(); - sub_nodes.push((sub_sub_node, node_lock.as_ref().prefix().len())); - } - } - } - unsafe { - BaseNode::drop_node(node, self.allocator.clone()); - } +pub(crate) trait CongeeVisitor { + fn visit_payload(&mut self, _payload: usize) {} + fn pre_visit_sub_node(&mut self, _node: NonNull) {} + fn post_visit_sub_node(&mut self, _node: NonNull) {} +} + +struct DropVisitor { + allocator: A, +} + +impl CongeeVisitor + for DropVisitor +{ + fn visit_payload(&mut self, _payload: usize) {} + fn pre_visit_sub_node(&mut self, _node: NonNull) {} + fn post_visit_sub_node(&mut self, node: NonNull) { + unsafe { + BaseNode::drop_node(node, self.allocator.clone()); } + } +} + +#[cfg(test)] +struct ValueCountVisitor { + value_count: usize, +} + +#[cfg(test)] +impl CongeeVisitor for ValueCountVisitor { + fn visit_payload(&mut self, _payload: usize) { + self.value_count += 1; + } +} + +impl Drop for RawCongee { + fn drop(&mut self) { + let mut visitor = DropVisitor:: { + allocator: self.allocator.clone(), + }; + self.dfs_visitor_slow(&mut visitor).unwrap(); // see this: https://github.com/XiangpengHao/congee/issues/20 for _ in 0..128 { @@ -62,7 +82,7 @@ impl Drop for RawCongee { } } -impl RawCongee { +impl RawCongee { pub fn new(allocator: A) -> Self { let root = BaseNode::make_node::(&[], &allocator) .expect("Can't allocate memory for root node!"); @@ -135,6 +155,54 @@ impl RawCongee { } } + /// Depth-First Search visitor implemented recursively, use with caution + /// Will panic on any contention + pub(crate) fn dfs_visitor_slow>( + &self, + visitor: &mut V, + ) -> Result<(), ArtError> { + let first = PtrType::SubNode(unsafe { + std::mem::transmute::, NonNull>(self.root) + }); + self.recursive_dfs(first, visitor)?; + Ok(()) + } + + fn recursive_dfs>( + &self, + node: PtrType, + visitor: &mut V, + ) -> Result<(), ArtError> { + match node { + PtrType::Payload(v) => { + visitor.visit_payload(v); + } + PtrType::SubNode(node_ptr) => { + visitor.pre_visit_sub_node(node_ptr); + let node_lock = BaseNode::read_lock(node_ptr)?; + let children = node_lock.as_ref().get_children(0, 255); + for (_k, child_ptr) in children { + let next = child_ptr.downcast::(node_lock.as_ref().prefix().len()); + self.recursive_dfs(next, visitor)?; + } + visitor.post_visit_sub_node(node_ptr); + node_lock.check_version()?; + } + } + Ok(()) + } + + /// Returns the number of values in the tree. + #[cfg(test)] + pub(crate) fn value_count(&self, _guard: &Guard) -> usize { + loop { + let mut visitor = ValueCountVisitor:: { value_count: 0 }; + if let Ok(_) = self.dfs_visitor_slow(&mut visitor) { + return visitor.value_count; + } + } + } + #[inline] fn insert_inner( &self, @@ -448,7 +516,7 @@ impl RawCongee { None => { // new value is none, we need to delete this entry debug_assert!(parent.is_some()); // reaching leaf means we must have parent, bcs root can't be leaf - if node.as_ref().get_count() == 1 { + if node.as_ref().value_count() == 1 { let (parent_node, parent_key) = parent.unwrap(); let mut write_p = parent_node.upgrade().map_err(|(_n, v)| v)?;