From e5183f85dc3a6b851429c183dfa1e21ed109ba6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Mon, 31 Jul 2023 12:52:29 +0200 Subject: [PATCH] Make DiskBtreeReader::dump async (#4838) ## Problem `DiskBtreeReader::dump` calls `read_blk` internally, which we want to make async in the future. As it is currently relying on recursion, and async doesn't like recursion, we want to find an alternative to that and instead traverse the tree using a loop and a manual stack. ## Summary of changes * Make `DiskBtreeReader::dump` and all the places calling it async * Make `DiskBtreeReader::dump` non-recursive internally and use a stack instead. It now deparses the node in each iteration, which isn't optimal, but on the other hand it's hard to store the node as it is referencing the buffer. Self referential data are hard in Rust. For a dumping function, speed isn't a priority so we deparse the node multiple times now (up to branching factor many times). Part of https://github.com/neondatabase/neon/issues/4743 I have verified that output is unchanged by comparing the output of this command both before and after this patch: ``` cargo test -p pageserver -- particular_data --nocapture ``` --- pageserver/src/tenant/disk_btree.rs | 67 ++++++++++--------- .../src/tenant/storage_layer/delta_layer.rs | 2 +- .../src/tenant/storage_layer/image_layer.rs | 2 +- 3 files changed, 37 insertions(+), 34 deletions(-) diff --git a/pageserver/src/tenant/disk_btree.rs b/pageserver/src/tenant/disk_btree.rs index 734409a619..518286ddb0 100644 --- a/pageserver/src/tenant/disk_btree.rs +++ b/pageserver/src/tenant/disk_btree.rs @@ -390,39 +390,42 @@ where } #[allow(dead_code)] - pub fn dump(&self) -> Result<()> { - self.dump_recurse(self.root_blk, &[], 0) - } + pub async fn dump(&self) -> Result<()> { + let mut stack = Vec::new(); - fn dump_recurse(&self, blknum: u32, path: &[u8], depth: usize) -> Result<()> { - let blk = self.reader.read_blk(self.start_blk + blknum)?; - let buf: &[u8] = blk.as_ref(); + stack.push((self.root_blk, String::new(), 0, 0, 0)); - let node = OnDiskNode::::deparse(buf)?; + while let Some((blknum, path, depth, child_idx, key_off)) = stack.pop() { + let blk = self.reader.read_blk(self.start_blk + blknum)?; + let buf: &[u8] = blk.as_ref(); + let node = OnDiskNode::::deparse(buf)?; - print!("{:indent$}", "", indent = depth * 2); - println!( - "blk #{}: path {}: prefix {}, suffix_len {}", - blknum, - hex::encode(path), - hex::encode(node.prefix), - node.suffix_len - ); + if child_idx == 0 { + print!("{:indent$}", "", indent = depth * 2); + let path_prefix = stack + .iter() + .map(|(_blknum, path, ..)| path.as_str()) + .collect::(); + println!( + "blk #{blknum}: path {path_prefix}{path}: prefix {}, suffix_len {}", + hex::encode(node.prefix), + node.suffix_len + ); + } - let mut idx = 0; - let mut key_off = 0; - while idx < node.num_children { + if child_idx + 1 < node.num_children { + let key_off = key_off + node.suffix_len as usize; + stack.push((blknum, path.clone(), depth, child_idx + 1, key_off)); + } let key = &node.keys[key_off..key_off + node.suffix_len as usize]; - let val = node.value(idx as usize); + let val = node.value(child_idx as usize); + print!("{:indent$}", "", indent = depth * 2 + 2); println!("{}: {}", hex::encode(key), hex::encode(val.0)); if node.level > 0 { - let child_path = [path, node.prefix].concat(); - self.dump_recurse(val.to_blknum(), &child_path, depth + 1)?; + stack.push((val.to_blknum(), hex::encode(node.prefix), depth + 1, 0, 0)); } - idx += 1; - key_off += node.suffix_len as usize; } Ok(()) } @@ -754,8 +757,8 @@ mod tests { } } - #[test] - fn basic() -> Result<()> { + #[tokio::test] + async fn basic() -> Result<()> { let mut disk = TestDisk::new(); let mut writer = DiskBtreeBuilder::<_, 6>::new(&mut disk); @@ -775,7 +778,7 @@ mod tests { let reader = DiskBtreeReader::new(0, root_offset, disk); - reader.dump()?; + reader.dump().await?; // Test the `get` function on all the keys. for (key, val) in all_data.iter() { @@ -835,8 +838,8 @@ mod tests { Ok(()) } - #[test] - fn lots_of_keys() -> Result<()> { + #[tokio::test] + async fn lots_of_keys() -> Result<()> { let mut disk = TestDisk::new(); let mut writer = DiskBtreeBuilder::<_, 8>::new(&mut disk); @@ -856,7 +859,7 @@ mod tests { let reader = DiskBtreeReader::new(0, root_offset, disk); - reader.dump()?; + reader.dump().await?; use std::sync::Mutex; @@ -994,8 +997,8 @@ mod tests { /// /// This test contains a particular data set, see disk_btree_test_data.rs /// - #[test] - fn particular_data() -> Result<()> { + #[tokio::test] + async fn particular_data() -> Result<()> { // Build a tree from it let mut disk = TestDisk::new(); let mut writer = DiskBtreeBuilder::<_, 26>::new(&mut disk); @@ -1022,7 +1025,7 @@ mod tests { })?; assert_eq!(count, disk_btree_test_data::TEST_DATA.len()); - reader.dump()?; + reader.dump().await?; Ok(()) } diff --git a/pageserver/src/tenant/storage_layer/delta_layer.rs b/pageserver/src/tenant/storage_layer/delta_layer.rs index 9585c04120..7574554b4e 100644 --- a/pageserver/src/tenant/storage_layer/delta_layer.rs +++ b/pageserver/src/tenant/storage_layer/delta_layer.rs @@ -256,7 +256,7 @@ impl Layer for DeltaLayer { file, ); - tree_reader.dump()?; + tree_reader.dump().await?; let mut cursor = file.block_cursor(); diff --git a/pageserver/src/tenant/storage_layer/image_layer.rs b/pageserver/src/tenant/storage_layer/image_layer.rs index f758347d9a..a5b9d8834e 100644 --- a/pageserver/src/tenant/storage_layer/image_layer.rs +++ b/pageserver/src/tenant/storage_layer/image_layer.rs @@ -175,7 +175,7 @@ impl Layer for ImageLayer { let tree_reader = DiskBtreeReader::<_, KEY_SIZE>::new(inner.index_start_blk, inner.index_root_blk, file); - tree_reader.dump()?; + tree_reader.dump().await?; tree_reader.visit(&[0u8; KEY_SIZE], VisitDirection::Forwards, |key, value| { println!("key: {} offset {}", hex::encode(key), value);