LCOV - code coverage report
Current view: top level - pageserver/src/tenant - disk_btree.rs (source / functions) Coverage Total Hit
Test: 32f4a56327bc9da697706839ed4836b2a00a408f.info Lines: 97.2 % 677 658
Test Date: 2024-02-07 07:37:29 Functions: 80.3 % 213 171

            Line data    Source code
       1              : //!
       2              : //! Simple on-disk B-tree implementation
       3              : //!
       4              : //! This is used as the index structure within image and delta layers
       5              : //!
       6              : //! Features:
       7              : //! - Fixed-width keys
       8              : //! - Fixed-width values (VALUE_SZ)
       9              : //! - The tree is created in a bulk operation. Insert/deletion after creation
      10              : //!   is not supported
      11              : //! - page-oriented
      12              : //!
      13              : //! TODO:
      14              : //! - maybe something like an Adaptive Radix Tree would be more efficient?
      15              : //! - the values stored by image and delta layers are offsets into the file,
      16              : //!   and they are in monotonically increasing order. Prefix compression would
      17              : //!   be very useful for them, too.
      18              : //! - An Iterator interface would be more convenient for the callers than the
      19              : //!   'visit' function
      20              : //!
      21              : use byteorder::{ReadBytesExt, BE};
      22              : use bytes::{BufMut, Bytes, BytesMut};
      23              : use either::Either;
      24              : use hex;
      25              : use std::{cmp::Ordering, io, result};
      26              : use thiserror::Error;
      27              : use tracing::error;
      28              : 
      29              : use crate::{
      30              :     context::{DownloadBehavior, RequestContext},
      31              :     task_mgr::TaskKind,
      32              :     tenant::block_io::{BlockReader, BlockWriter},
      33              : };
      34              : 
      35              : // The maximum size of a value stored in the B-tree. 5 bytes is enough currently.
      36              : pub const VALUE_SZ: usize = 5;
      37              : pub const MAX_VALUE: u64 = 0x007f_ffff_ffff;
      38              : 
      39              : #[allow(dead_code)]
      40              : pub const PAGE_SZ: usize = 8192;
      41              : 
      42            0 : #[derive(Clone, Copy, Debug)]
      43              : struct Value([u8; VALUE_SZ]);
      44              : 
      45              : impl Value {
      46     85996118 :     fn from_slice(slice: &[u8]) -> Value {
      47     85996118 :         let mut b = [0u8; VALUE_SZ];
      48     85996118 :         b.copy_from_slice(slice);
      49     85996118 :         Value(b)
      50     85996118 :     }
      51              : 
      52     52969499 :     fn from_u64(x: u64) -> Value {
      53     52969499 :         assert!(x <= 0x007f_ffff_ffff);
      54     52969499 :         Value([
      55     52969499 :             (x >> 32) as u8,
      56     52969499 :             (x >> 24) as u8,
      57     52969499 :             (x >> 16) as u8,
      58     52969499 :             (x >> 8) as u8,
      59     52969499 :             x as u8,
      60     52969499 :         ])
      61     52969499 :     }
      62              : 
      63       102730 :     fn from_blknum(x: u32) -> Value {
      64       102730 :         Value([
      65       102730 :             0x80,
      66       102730 :             (x >> 24) as u8,
      67       102730 :             (x >> 16) as u8,
      68       102730 :             (x >> 8) as u8,
      69       102730 :             x as u8,
      70       102730 :         ])
      71       102730 :     }
      72              : 
      73              :     #[allow(dead_code)]
      74            0 :     fn is_offset(self) -> bool {
      75            0 :         self.0[0] & 0x80 != 0
      76            0 :     }
      77              : 
      78     75930108 :     fn to_u64(self) -> u64 {
      79     75930108 :         let b = &self.0;
      80     75930108 :         (b[0] as u64) << 32
      81     75930108 :             | (b[1] as u64) << 24
      82     75930108 :             | (b[2] as u64) << 16
      83     75930108 :             | (b[3] as u64) << 8
      84     75930108 :             | b[4] as u64
      85     75930108 :     }
      86              : 
      87     10059986 :     fn to_blknum(self) -> u32 {
      88     10059986 :         let b = &self.0;
      89     10059986 :         assert!(b[0] == 0x80);
      90     10059986 :         (b[1] as u32) << 24 | (b[2] as u32) << 16 | (b[3] as u32) << 8 | b[4] as u32
      91     10059986 :     }
      92              : }
      93              : 
      94            0 : #[derive(Error, Debug)]
      95              : pub enum DiskBtreeError {
      96              :     #[error("Attempt to append a value that is too large {0} > {}", MAX_VALUE)]
      97              :     AppendOverflow(u64),
      98              : 
      99              :     #[error("Unsorted input: key {key:?} is <= last_key {last_key:?}")]
     100              :     UnsortedInput { key: Box<[u8]>, last_key: Box<[u8]> },
     101              : 
     102              :     #[error("Could not push to new leaf node")]
     103              :     FailedToPushToNewLeafNode,
     104              : 
     105              :     #[error("IoError: {0}")]
     106              :     Io(#[from] io::Error),
     107              : }
     108              : 
     109              : pub type Result<T> = result::Result<T, DiskBtreeError>;
     110              : 
     111              : /// This is the on-disk representation.
     112              : struct OnDiskNode<'a, const L: usize> {
     113              :     // Fixed-width fields
     114              :     num_children: u16,
     115              :     level: u8,
     116              :     prefix_len: u8,
     117              :     suffix_len: u8,
     118              : 
     119              :     // Variable-length fields. These are stored on-disk after the fixed-width
     120              :     // fields, in this order. In the in-memory representation, these point to
     121              :     // the right parts in the page buffer.
     122              :     prefix: &'a [u8],
     123              :     keys: &'a [u8],
     124              :     values: &'a [u8],
     125              : }
     126              : 
     127              : impl<'a, const L: usize> OnDiskNode<'a, L> {
     128              :     ///
     129              :     /// Interpret a PAGE_SZ page as a node.
     130              :     ///
     131     34719618 :     fn deparse(buf: &[u8]) -> Result<OnDiskNode<L>> {
     132     34719618 :         let mut cursor = std::io::Cursor::new(buf);
     133     34719618 :         let num_children = cursor.read_u16::<BE>()?;
     134     34719618 :         let level = cursor.read_u8()?;
     135     34719618 :         let prefix_len = cursor.read_u8()?;
     136     34719618 :         let suffix_len = cursor.read_u8()?;
     137              : 
     138     34719618 :         let mut off = cursor.position();
     139     34719618 :         let prefix_off = off as usize;
     140     34719618 :         off += prefix_len as u64;
     141     34719618 : 
     142     34719618 :         let keys_off = off as usize;
     143     34719618 :         let keys_len = num_children as usize * suffix_len as usize;
     144     34719618 :         off += keys_len as u64;
     145     34719618 : 
     146     34719618 :         let values_off = off as usize;
     147     34719618 :         let values_len = num_children as usize * VALUE_SZ;
     148     34719618 :         //off += values_len as u64;
     149     34719618 : 
     150     34719618 :         let prefix = &buf[prefix_off..prefix_off + prefix_len as usize];
     151     34719618 :         let keys = &buf[keys_off..keys_off + keys_len];
     152     34719618 :         let values = &buf[values_off..values_off + values_len];
     153     34719618 : 
     154     34719618 :         Ok(OnDiskNode {
     155     34719618 :             num_children,
     156     34719618 :             level,
     157     34719618 :             prefix_len,
     158     34719618 :             suffix_len,
     159     34719618 :             prefix,
     160     34719618 :             keys,
     161     34719618 :             values,
     162     34719618 :         })
     163     34719618 :     }
     164              : 
     165              :     ///
     166              :     /// Read a value at 'idx'
     167              :     ///
     168     85996132 :     fn value(&self, idx: usize) -> Value {
     169     85996132 :         let value_off = idx * VALUE_SZ;
     170     85996132 :         let value_slice = &self.values[value_off..value_off + VALUE_SZ];
     171     85996132 :         Value::from_slice(value_slice)
     172     85996132 :     }
     173              : 
     174     34409934 :     fn binary_search(
     175     34409934 :         &self,
     176     34409934 :         search_key: &[u8; L],
     177     34409934 :         keybuf: &mut [u8],
     178     34409934 :     ) -> result::Result<usize, usize> {
     179     34409934 :         let mut size = self.num_children as usize;
     180     34409934 :         let mut low = 0;
     181     34409934 :         let mut high = size;
     182    233694766 :         while low < high {
     183    200652796 :             let mid = low + size / 2;
     184    200652796 : 
     185    200652796 :             let key_off = mid * self.suffix_len as usize;
     186    200652796 :             let suffix = &self.keys[key_off..key_off + self.suffix_len as usize];
     187    200652796 :             // Does this match?
     188    200652796 :             keybuf[self.prefix_len as usize..].copy_from_slice(suffix);
     189    200652796 : 
     190    200652796 :             let cmp = keybuf[..].cmp(search_key);
     191    200652796 : 
     192    200652796 :             if cmp == Ordering::Less {
     193     63668611 :                 low = mid + 1;
     194    136984185 :             } else if cmp == Ordering::Greater {
     195    135616221 :                 high = mid;
     196    135616221 :             } else {
     197      1367964 :                 return Ok(mid);
     198              :             }
     199    199284832 :             size = high - low;
     200              :         }
     201     33041970 :         Err(low)
     202     34409934 :     }
     203              : }
     204              : 
     205              : ///
     206              : /// Public reader object, to search the tree.
     207              : ///
     208              : pub struct DiskBtreeReader<R, const L: usize>
     209              : where
     210              :     R: BlockReader,
     211              : {
     212              :     start_blk: u32,
     213              :     root_blk: u32,
     214              :     reader: R,
     215              : }
     216              : 
     217     34409920 : #[derive(Clone, Copy, Debug, PartialEq, Eq)]
     218              : pub enum VisitDirection {
     219              :     Forwards,
     220              :     Backwards,
     221              : }
     222              : 
     223              : impl<R, const L: usize> DiskBtreeReader<R, L>
     224              : where
     225              :     R: BlockReader,
     226              : {
     227     23938769 :     pub fn new(start_blk: u32, root_blk: u32, reader: R) -> Self {
     228     23938769 :         DiskBtreeReader {
     229     23938769 :             start_blk,
     230     23938769 :             root_blk,
     231     23938769 :             reader,
     232     23938769 :         }
     233     23938769 :     }
     234              : 
     235              :     ///
     236              :     /// Read the value for given key. Returns the value, or None if it doesn't exist.
     237              :     ///
     238      1538358 :     pub async fn get(&self, search_key: &[u8; L], ctx: &RequestContext) -> Result<Option<u64>> {
     239      1538358 :         let mut result: Option<u64> = None;
     240      1538358 :         self.visit(
     241      1538358 :             search_key,
     242      1538358 :             VisitDirection::Forwards,
     243      1538358 :             |key, value| {
     244      1338334 :                 if key == search_key {
     245      1336324 :                     result = Some(value);
     246      1336324 :                 }
     247      1338334 :                 false
     248      1538358 :             },
     249      1538358 :             ctx,
     250      1538358 :         )
     251         6280 :         .await?;
     252      1538358 :         Ok(result)
     253      1538358 :     }
     254              : 
     255              :     ///
     256              :     /// Scan the tree, starting from 'search_key', in the given direction. 'visitor'
     257              :     /// will be called for every key >= 'search_key' (or <= 'search_key', if scanning
     258              :     /// backwards)
     259              :     ///
     260     24349950 :     pub async fn visit<V>(
     261     24349950 :         &self,
     262     24349950 :         search_key: &[u8; L],
     263     24349950 :         dir: VisitDirection,
     264     24349950 :         mut visitor: V,
     265     24349950 :         ctx: &RequestContext,
     266     24349950 :     ) -> Result<bool>
     267     24349950 :     where
     268     24349950 :         V: FnMut(&[u8], u64) -> bool,
     269     24349950 :     {
     270     24349950 :         let mut stack = Vec::new();
     271     24349950 :         stack.push((self.root_blk, None));
     272     24349950 :         let block_cursor = self.reader.block_cursor();
     273     34920742 :         while let Some((node_blknum, opt_iter)) = stack.pop() {
     274              :             // Locate the node.
     275     34713586 :             let node_buf = block_cursor
     276     34713586 :                 .read_blk(self.start_blk + node_blknum, ctx)
     277       248496 :                 .await?;
     278              : 
     279     34713586 :             let node = OnDiskNode::deparse(node_buf.as_ref())?;
     280     34713586 :             let prefix_len = node.prefix_len as usize;
     281     34713586 :             let suffix_len = node.suffix_len as usize;
     282     34713586 : 
     283     34713586 :             assert!(node.num_children > 0);
     284              : 
     285     34713586 :             let mut keybuf = Vec::new();
     286     34713586 :             keybuf.extend(node.prefix);
     287     34713586 :             keybuf.resize(prefix_len + suffix_len, 0);
     288              : 
     289     34713586 :             let mut iter = if let Some(iter) = opt_iter {
     290       303652 :                 iter
     291     34409934 :             } else if dir == VisitDirection::Forwards {
     292              :                 // Locate the first match
     293      2054571 :                 let idx = match node.binary_search(search_key, keybuf.as_mut_slice()) {
     294      1338715 :                     Ok(idx) => idx,
     295       715856 :                     Err(idx) => {
     296       715856 :                         if node.level == 0 {
     297              :                             // Imagine that the node contains the following keys:
     298              :                             //
     299              :                             // 1
     300              :                             // 3  <-- idx
     301              :                             // 5
     302              :                             //
     303              :                             // If the search key is '2' and there is exact match,
     304              :                             // the binary search would return the index of key
     305              :                             // '3'. That's cool, '3' is the first key to return.
     306       239200 :                             idx
     307              :                         } else {
     308              :                             // This is an internal page, so each key represents a lower
     309              :                             // bound for what's in the child page. If there is no exact
     310              :                             // match, we have to return the *previous* entry.
     311              :                             //
     312              :                             // 1  <-- return this
     313              :                             // 3  <-- idx
     314              :                             // 5
     315       476656 :                             idx.saturating_sub(1)
     316              :                         }
     317              :                     }
     318              :                 };
     319      2054571 :                 Either::Left(idx..node.num_children.into())
     320              :             } else {
     321     32355363 :                 let idx = match node.binary_search(search_key, keybuf.as_mut_slice()) {
     322        29249 :                     Ok(idx) => {
     323        29249 :                         // Exact match. That's the first entry to return, and walk
     324        29249 :                         // backwards from there.
     325        29249 :                         idx
     326              :                     }
     327     32326114 :                     Err(idx) => {
     328              :                         // No exact match. The binary search returned the index of the
     329              :                         // first key that's > search_key. Back off by one, and walk
     330              :                         // backwards from there.
     331     32326114 :                         if let Some(idx) = idx.checked_sub(1) {
     332     19701169 :                             idx
     333              :                         } else {
     334     12624945 :                             return Ok(false);
     335              :                         }
     336              :                     }
     337              :                 };
     338     19730418 :                 Either::Right((0..=idx).rev())
     339              :             };
     340              : 
     341              :             // idx points to the first match now. Keep going from there
     342     86500908 :             while let Some(idx) = iter.next() {
     343     85990101 :                 let key_off = idx * suffix_len;
     344     85990101 :                 let suffix = &node.keys[key_off..key_off + suffix_len];
     345     85990101 :                 keybuf[prefix_len..].copy_from_slice(suffix);
     346     85990101 :                 let value = node.value(idx);
     347     85990101 :                 #[allow(clippy::collapsible_if)]
     348     85990101 :                 if node.level == 0 {
     349              :                     // leaf
     350     75930116 :                     if !visitor(&keybuf, value.to_u64()) {
     351     11517848 :                         return Ok(false);
     352     64412267 :                     }
     353              :                 } else {
     354     10059985 :                     stack.push((node_blknum, Some(iter)));
     355     10059985 :                     stack.push((value.to_blknum(), None));
     356     10059985 :                     break;
     357              :                 }
     358              :             }
     359              :         }
     360       207156 :         Ok(true)
     361     24349949 :     }
     362              : 
     363              :     #[allow(dead_code)]
     364           10 :     pub async fn dump(&self) -> Result<()> {
     365           10 :         let mut stack = Vec::new();
     366           10 :         let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
     367           10 : 
     368           10 :         stack.push((self.root_blk, String::new(), 0, 0, 0));
     369           10 : 
     370           10 :         let block_cursor = self.reader.block_cursor();
     371              : 
     372         6042 :         while let Some((blknum, path, depth, child_idx, key_off)) = stack.pop() {
     373         6032 :             let blk = block_cursor.read_blk(self.start_blk + blknum, &ctx).await?;
     374         6032 :             let buf: &[u8] = blk.as_ref();
     375         6032 :             let node = OnDiskNode::<L>::deparse(buf)?;
     376              : 
     377         6032 :             if child_idx == 0 {
     378           18 :                 print!("{:indent$}", "", indent = depth * 2);
     379           18 :                 let path_prefix = stack
     380           18 :                     .iter()
     381           18 :                     .map(|(_blknum, path, ..)| path.as_str())
     382           18 :                     .collect::<String>();
     383           18 :                 println!(
     384           18 :                     "blk #{blknum}: path {path_prefix}{path}: prefix {}, suffix_len {}",
     385           18 :                     hex::encode(node.prefix),
     386           18 :                     node.suffix_len
     387           18 :                 );
     388         6014 :             }
     389              : 
     390         6032 :             if child_idx + 1 < node.num_children {
     391         6014 :                 let key_off = key_off + node.suffix_len as usize;
     392         6014 :                 stack.push((blknum, path.clone(), depth, child_idx + 1, key_off));
     393         6014 :             }
     394         6032 :             let key = &node.keys[key_off..key_off + node.suffix_len as usize];
     395         6032 :             let val = node.value(child_idx as usize);
     396         6032 : 
     397         6032 :             print!("{:indent$}", "", indent = depth * 2 + 2);
     398         6032 :             println!("{}: {}", hex::encode(key), hex::encode(val.0));
     399         6032 : 
     400         6032 :             if node.level > 0 {
     401            8 :                 stack.push((val.to_blknum(), hex::encode(node.prefix), depth + 1, 0, 0));
     402         6024 :             }
     403              :         }
     404           10 :         Ok(())
     405           10 :     }
     406              : }
     407              : 
     408              : ///
     409              : /// Public builder object, for creating a new tree.
     410              : ///
     411              : /// Usage: Create a builder object by calling 'new', load all the data into the
     412              : /// tree by calling 'append' for each key-value pair, and then call 'finish'
     413              : ///
     414              : /// 'L' is the key length in bytes
     415              : pub struct DiskBtreeBuilder<W, const L: usize>
     416              : where
     417              :     W: BlockWriter,
     418              : {
     419              :     writer: W,
     420              : 
     421              :     ///
     422              :     /// `stack[0]` is the current root page, `stack.last()` is the leaf.
     423              :     ///
     424              :     /// We maintain the length of the stack to be always greater than zero.
     425              :     /// Two exceptions are:
     426              :     /// 1. `Self::flush_node`. The method will push the new node if it extracted the last one.
     427              :     ///   So because other methods cannot see the intermediate state invariant still holds.
     428              :     /// 2. `Self::finish`. It consumes self and does not return it back,
     429              :     ///  which means that this is where the structure is destroyed.
     430              :     ///  Thus stack of zero length cannot be observed by other methods.
     431              :     stack: Vec<BuildNode<L>>,
     432              : 
     433              :     /// Last key that was appended to the tree. Used to sanity check that append
     434              :     /// is called in increasing key order.
     435              :     last_key: Option<[u8; L]>,
     436              : }
     437              : 
     438              : impl<W, const L: usize> DiskBtreeBuilder<W, L>
     439              : where
     440              :     W: BlockWriter,
     441              : {
     442        21957 :     pub fn new(writer: W) -> Self {
     443        21957 :         DiskBtreeBuilder {
     444        21957 :             writer,
     445        21957 :             last_key: None,
     446        21957 :             stack: vec![BuildNode::new(0)],
     447        21957 :         }
     448        21957 :     }
     449              : 
     450     52969548 :     pub fn append(&mut self, key: &[u8; L], value: u64) -> Result<()> {
     451     52969548 :         if value > MAX_VALUE {
     452            0 :             return Err(DiskBtreeError::AppendOverflow(value));
     453     52969548 :         }
     454     52969548 :         if let Some(last_key) = &self.last_key {
     455     52947591 :             if key <= last_key {
     456            2 :                 return Err(DiskBtreeError::UnsortedInput {
     457            2 :                     key: key.as_slice().into(),
     458            2 :                     last_key: last_key.as_slice().into(),
     459            2 :                 });
     460     52947589 :             }
     461        21957 :         }
     462     52969546 :         self.last_key = Some(*key);
     463     52969546 : 
     464     52969546 :         self.append_internal(key, Value::from_u64(value))
     465     52969548 :     }
     466              : 
     467     53072276 :     fn append_internal(&mut self, key: &[u8; L], value: Value) -> Result<()> {
     468     53072276 :         // Try to append to the current leaf buffer
     469     53072276 :         let last = self
     470     53072276 :             .stack
     471     53072276 :             .last_mut()
     472     53072276 :             .expect("should always have at least one item");
     473     53072276 :         let level = last.level;
     474     53072276 :         if last.push(key, value) {
     475     52881043 :             return Ok(());
     476       191233 :         }
     477       191233 : 
     478       191233 :         // It did not fit. Try to compress, and if it succeeds to make
     479       191233 :         // some room on the node, try appending to it again.
     480       191233 :         #[allow(clippy::collapsible_if)]
     481       191233 :         if last.compress() {
     482        96904 :             if last.push(key, value) {
     483        96819 :                 return Ok(());
     484           85 :             }
     485        94329 :         }
     486              : 
     487              :         // Could not append to the current leaf. Flush it and create a new one.
     488        94414 :         self.flush_node()?;
     489              : 
     490              :         // Replace the node we flushed with an empty one and append the new
     491              :         // key to it.
     492        94414 :         let mut last = BuildNode::new(level);
     493        94414 :         if !last.push(key, value) {
     494            0 :             return Err(DiskBtreeError::FailedToPushToNewLeafNode);
     495        94414 :         }
     496        94414 : 
     497        94414 :         self.stack.push(last);
     498        94414 : 
     499        94414 :         Ok(())
     500     53072276 :     }
     501              : 
     502              :     /// Flush the bottommost node in the stack to disk. Appends a downlink to its parent,
     503              :     /// and recursively flushes the parent too, if it becomes full. If the root page becomes full,
     504              :     /// creates a new root page, increasing the height of the tree.
     505       102730 :     fn flush_node(&mut self) -> Result<()> {
     506       102730 :         // Get the current bottommost node in the stack and flush it to disk.
     507       102730 :         let last = self
     508       102730 :             .stack
     509       102730 :             .pop()
     510       102730 :             .expect("should always have at least one item");
     511       102730 :         let buf = last.pack();
     512       102730 :         let downlink_key = last.first_key();
     513       102730 :         let downlink_ptr = self.writer.write_blk(buf)?;
     514              : 
     515              :         // Append the downlink to the parent. If there is no parent, ie. this was the root page,
     516              :         // create a new root page, increasing the height of the tree.
     517       102730 :         if self.stack.is_empty() {
     518         8319 :             self.stack.push(BuildNode::new(last.level + 1));
     519        94411 :         }
     520       102730 :         self.append_internal(&downlink_key, Value::from_blknum(downlink_ptr))
     521       102730 :     }
     522              : 
     523              :     ///
     524              :     /// Flushes everything to disk, and returns the block number of the root page.
     525              :     /// The caller must store the root block number "out-of-band", and pass it
     526              :     /// to the DiskBtreeReader::new() when you want to read the tree again.
     527              :     /// (In the image and delta layers, it is stored in the beginning of the file,
     528              :     /// in the summary header)
     529              :     ///
     530        21952 :     pub fn finish(mut self) -> Result<(u32, W)> {
     531              :         // flush all levels, except the root.
     532        30268 :         while self.stack.len() > 1 {
     533         8316 :             self.flush_node()?;
     534              :         }
     535              : 
     536        21952 :         let root = self
     537        21952 :             .stack
     538        21952 :             .first()
     539        21952 :             .expect("by the check above we left one item there");
     540        21952 :         let buf = root.pack();
     541        21952 :         let root_blknum = self.writer.write_blk(buf)?;
     542              : 
     543        21952 :         Ok((root_blknum, self.writer))
     544        21952 :     }
     545              : 
     546      2373630 :     pub fn borrow_writer(&self) -> &W {
     547      2373630 :         &self.writer
     548      2373630 :     }
     549              : }
     550              : 
     551              : ///
     552              : /// BuildNode represesnts an incomplete page that we are appending to.
     553              : ///
     554            0 : #[derive(Clone, Debug)]
     555              : struct BuildNode<const L: usize> {
     556              :     num_children: u16,
     557              :     level: u8,
     558              :     prefix: Vec<u8>,
     559              :     suffix_len: usize,
     560              : 
     561              :     keys: Vec<u8>,
     562              :     values: Vec<u8>,
     563              : 
     564              :     size: usize, // physical size of this node, if it was written to disk like this
     565              : }
     566              : 
     567              : const NODE_SIZE: usize = PAGE_SZ;
     568              : 
     569              : const NODE_HDR_SIZE: usize = 2 + 1 + 1 + 1;
     570              : 
     571              : impl<const L: usize> BuildNode<L> {
     572       124690 :     fn new(level: u8) -> Self {
     573       124690 :         BuildNode {
     574       124690 :             num_children: 0,
     575       124690 :             level,
     576       124690 :             prefix: Vec::new(),
     577       124690 :             suffix_len: 0,
     578       124690 :             keys: Vec::new(),
     579       124690 :             values: Vec::new(),
     580       124690 :             size: NODE_HDR_SIZE,
     581       124690 :         }
     582       124690 :     }
     583              : 
     584              :     /// Try to append a key-value pair to this node. Returns 'true' on
     585              :     /// success, 'false' if the page was full or the key was
     586              :     /// incompatible with the prefix of the existing keys.
     587     53263594 :     fn push(&mut self, key: &[u8; L], value: Value) -> bool {
     588     53263594 :         // If we have already performed prefix-compression on the page,
     589     53263594 :         // check that the incoming key has the same prefix.
     590     53263594 :         if self.num_children > 0 {
     591              :             // does the prefix allow it?
     592     53138904 :             if !key.starts_with(&self.prefix) {
     593        24362 :                 return false;
     594     53114542 :             }
     595       124690 :         } else {
     596       124690 :             self.suffix_len = key.len();
     597       124690 :         }
     598              : 
     599              :         // Is the node too full?
     600     53239232 :         if self.size + self.suffix_len + VALUE_SZ >= NODE_SIZE {
     601       166956 :             return false;
     602     53072276 :         }
     603     53072276 : 
     604     53072276 :         // All clear
     605     53072276 :         self.num_children += 1;
     606     53072276 :         self.keys.extend(&key[self.prefix.len()..]);
     607     53072276 :         self.values.extend(value.0);
     608              : 
     609     53072276 :         assert!(self.keys.len() == self.num_children as usize * self.suffix_len);
     610     53072276 :         assert!(self.values.len() == self.num_children as usize * VALUE_SZ);
     611              : 
     612     53072276 :         self.size += self.suffix_len + VALUE_SZ;
     613     53072276 : 
     614     53072276 :         true
     615     53263594 :     }
     616              : 
     617              :     ///
     618              :     /// Perform prefix-compression.
     619              :     ///
     620              :     /// Returns 'true' on success, 'false' if no compression was possible.
     621              :     ///
     622       191233 :     fn compress(&mut self) -> bool {
     623       191233 :         let first_suffix = self.first_suffix();
     624       191233 :         let last_suffix = self.last_suffix();
     625       191233 : 
     626       191233 :         // Find the common prefix among all keys
     627       191233 :         let mut prefix_len = 0;
     628      1882999 :         while prefix_len < self.suffix_len {
     629      1882999 :             if first_suffix[prefix_len] != last_suffix[prefix_len] {
     630       191233 :                 break;
     631      1691766 :             }
     632      1691766 :             prefix_len += 1;
     633              :         }
     634       191233 :         if prefix_len == 0 {
     635        94329 :             return false;
     636        96904 :         }
     637        96904 : 
     638        96904 :         // Can compress. Rewrite the keys without the common prefix.
     639        96904 :         self.prefix.extend(&self.keys[..prefix_len]);
     640        96904 : 
     641        96904 :         let mut new_keys = Vec::new();
     642        96904 :         let mut key_off = 0;
     643     25720203 :         while key_off < self.keys.len() {
     644     25623299 :             let next_key_off = key_off + self.suffix_len;
     645     25623299 :             new_keys.extend(&self.keys[key_off + prefix_len..next_key_off]);
     646     25623299 :             key_off = next_key_off;
     647     25623299 :         }
     648        96904 :         self.keys = new_keys;
     649        96904 :         self.suffix_len -= prefix_len;
     650        96904 : 
     651        96904 :         self.size -= prefix_len * self.num_children as usize;
     652        96904 :         self.size += prefix_len;
     653              : 
     654        96904 :         assert!(self.keys.len() == self.num_children as usize * self.suffix_len);
     655        96904 :         assert!(self.values.len() == self.num_children as usize * VALUE_SZ);
     656              : 
     657        96904 :         true
     658       191233 :     }
     659              : 
     660              :     ///
     661              :     /// Serialize the node to on-disk format.
     662              :     ///
     663       124682 :     fn pack(&self) -> Bytes {
     664       124682 :         assert!(self.keys.len() == self.num_children as usize * self.suffix_len);
     665       124682 :         assert!(self.values.len() == self.num_children as usize * VALUE_SZ);
     666       124682 :         assert!(self.num_children > 0);
     667              : 
     668       124682 :         let mut buf = BytesMut::new();
     669       124682 : 
     670       124682 :         buf.put_u16(self.num_children);
     671       124682 :         buf.put_u8(self.level);
     672       124682 :         buf.put_u8(self.prefix.len() as u8);
     673       124682 :         buf.put_u8(self.suffix_len as u8);
     674       124682 :         buf.put(&self.prefix[..]);
     675       124682 :         buf.put(&self.keys[..]);
     676       124682 :         buf.put(&self.values[..]);
     677              : 
     678       124682 :         assert!(buf.len() == self.size);
     679              : 
     680       124682 :         assert!(buf.len() <= PAGE_SZ);
     681       124682 :         buf.resize(PAGE_SZ, 0);
     682       124682 :         buf.freeze()
     683       124682 :     }
     684              : 
     685       293963 :     fn first_suffix(&self) -> &[u8] {
     686       293963 :         &self.keys[..self.suffix_len]
     687       293963 :     }
     688       191233 :     fn last_suffix(&self) -> &[u8] {
     689       191233 :         &self.keys[self.keys.len() - self.suffix_len..]
     690       191233 :     }
     691              : 
     692              :     /// Return the full first key of the page, including the prefix
     693       102730 :     fn first_key(&self) -> [u8; L] {
     694       102730 :         let mut key = [0u8; L];
     695       102730 :         key[..self.prefix.len()].copy_from_slice(&self.prefix);
     696       102730 :         key[self.prefix.len()..].copy_from_slice(self.first_suffix());
     697       102730 :         key
     698       102730 :     }
     699              : }
     700              : 
     701              : #[cfg(test)]
     702              : pub(crate) mod tests {
     703              :     use super::*;
     704              :     use crate::context::DownloadBehavior;
     705              :     use crate::task_mgr::TaskKind;
     706              :     use crate::tenant::block_io::{BlockCursor, BlockLease, BlockReaderRef};
     707              :     use rand::Rng;
     708              :     use std::collections::BTreeMap;
     709              :     use std::sync::atomic::{AtomicUsize, Ordering};
     710              : 
     711           10 :     #[derive(Clone, Default)]
     712              :     pub(crate) struct TestDisk {
     713              :         blocks: Vec<Bytes>,
     714              :     }
     715              :     impl TestDisk {
     716           10 :         fn new() -> Self {
     717           10 :             Self::default()
     718           10 :         }
     719      1016332 :         pub(crate) fn read_blk(&self, blknum: u32) -> io::Result<BlockLease> {
     720      1016332 :             let mut buf = [0u8; PAGE_SZ];
     721      1016332 :             buf.copy_from_slice(&self.blocks[blknum as usize]);
     722      1016332 :             Ok(std::sync::Arc::new(buf).into())
     723      1016332 :         }
     724              :     }
     725              :     impl BlockReader for TestDisk {
     726       411199 :         fn block_cursor(&self) -> BlockCursor<'_> {
     727       411199 :             BlockCursor::new(BlockReaderRef::TestDisk(self))
     728       411199 :         }
     729              :     }
     730              :     impl BlockWriter for &mut TestDisk {
     731          214 :         fn write_blk(&mut self, buf: Bytes) -> io::Result<u32> {
     732          214 :             let blknum = self.blocks.len();
     733          214 :             self.blocks.push(buf);
     734          214 :             Ok(blknum as u32)
     735          214 :         }
     736              :     }
     737              : 
     738            2 :     #[tokio::test]
     739            2 :     async fn basic() -> Result<()> {
     740            2 :         let mut disk = TestDisk::new();
     741            2 :         let mut writer = DiskBtreeBuilder::<_, 6>::new(&mut disk);
     742            2 : 
     743            2 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
     744            2 : 
     745            2 :         let all_keys: Vec<&[u8; 6]> = vec![
     746            2 :             b"xaaaaa", b"xaaaba", b"xaaaca", b"xabaaa", b"xababa", b"xabaca", b"xabada", b"xabadb",
     747            2 :         ];
     748            2 :         let all_data: Vec<(&[u8; 6], u64)> = all_keys
     749            2 :             .iter()
     750            2 :             .enumerate()
     751           16 :             .map(|(idx, key)| (*key, idx as u64))
     752            2 :             .collect();
     753           16 :         for (key, val) in all_data.iter() {
     754           16 :             writer.append(key, *val)?;
     755              :         }
     756              : 
     757            2 :         let (root_offset, _writer) = writer.finish()?;
     758              : 
     759            2 :         let reader = DiskBtreeReader::new(0, root_offset, disk);
     760            2 : 
     761            2 :         reader.dump().await?;
     762              : 
     763              :         // Test the `get` function on all the keys.
     764           16 :         for (key, val) in all_data.iter() {
     765           16 :             assert_eq!(reader.get(key, &ctx).await?, Some(*val));
     766              :         }
     767              :         // And on some keys that don't exist
     768            2 :         assert_eq!(reader.get(b"aaaaaa", &ctx).await?, None);
     769            2 :         assert_eq!(reader.get(b"zzzzzz", &ctx).await?, None);
     770            2 :         assert_eq!(reader.get(b"xaaabx", &ctx).await?, None);
     771              : 
     772              :         // Test search with `visit` function
     773            2 :         let search_key = b"xabaaa";
     774            2 :         let expected: Vec<(Vec<u8>, u64)> = all_data
     775            2 :             .iter()
     776           16 :             .filter(|(key, _value)| key[..] >= search_key[..])
     777           10 :             .map(|(key, value)| (key.to_vec(), *value))
     778            2 :             .collect();
     779            2 : 
     780            2 :         let mut data = Vec::new();
     781            2 :         reader
     782            2 :             .visit(
     783            2 :                 search_key,
     784            2 :                 VisitDirection::Forwards,
     785           10 :                 |key, value| {
     786           10 :                     data.push((key.to_vec(), value));
     787           10 :                     true
     788           10 :                 },
     789            2 :                 &ctx,
     790            2 :             )
     791            0 :             .await?;
     792            2 :         assert_eq!(data, expected);
     793              : 
     794              :         // Test a backwards scan
     795            2 :         let mut expected: Vec<(Vec<u8>, u64)> = all_data
     796            2 :             .iter()
     797           16 :             .filter(|(key, _value)| key[..] <= search_key[..])
     798            8 :             .map(|(key, value)| (key.to_vec(), *value))
     799            2 :             .collect();
     800            2 :         expected.reverse();
     801            2 :         let mut data = Vec::new();
     802            2 :         reader
     803            2 :             .visit(
     804            2 :                 search_key,
     805            2 :                 VisitDirection::Backwards,
     806            8 :                 |key, value| {
     807            8 :                     data.push((key.to_vec(), value));
     808            8 :                     true
     809            8 :                 },
     810            2 :                 &ctx,
     811            2 :             )
     812            0 :             .await?;
     813            2 :         assert_eq!(data, expected);
     814              : 
     815              :         // Backward scan where nothing matches
     816            2 :         reader
     817            2 :             .visit(
     818            2 :                 b"aaaaaa",
     819            2 :                 VisitDirection::Backwards,
     820            2 :                 |key, value| {
     821            0 :                     panic!("found unexpected key {}: {}", hex::encode(key), value);
     822            2 :                 },
     823            2 :                 &ctx,
     824            2 :             )
     825            0 :             .await?;
     826              : 
     827              :         // Full scan
     828            2 :         let expected: Vec<(Vec<u8>, u64)> = all_data
     829            2 :             .iter()
     830           16 :             .map(|(key, value)| (key.to_vec(), *value))
     831            2 :             .collect();
     832            2 :         let mut data = Vec::new();
     833            2 :         reader
     834            2 :             .visit(
     835            2 :                 &[0u8; 6],
     836            2 :                 VisitDirection::Forwards,
     837           16 :                 |key, value| {
     838           16 :                     data.push((key.to_vec(), value));
     839           16 :                     true
     840           16 :                 },
     841            2 :                 &ctx,
     842            2 :             )
     843            0 :             .await?;
     844            2 :         assert_eq!(data, expected);
     845              : 
     846            2 :         Ok(())
     847              :     }
     848              : 
     849            2 :     #[tokio::test]
     850            2 :     async fn lots_of_keys() -> Result<()> {
     851            2 :         let mut disk = TestDisk::new();
     852            2 :         let mut writer = DiskBtreeBuilder::<_, 8>::new(&mut disk);
     853            2 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
     854            2 : 
     855            2 :         const NUM_KEYS: u64 = 1000;
     856            2 : 
     857            2 :         let mut all_data: BTreeMap<u64, u64> = BTreeMap::new();
     858              : 
     859         2002 :         for idx in 0..NUM_KEYS {
     860         2000 :             let key_int: u64 = 1 + idx * 2;
     861         2000 :             let key = u64::to_be_bytes(key_int);
     862         2000 :             writer.append(&key, idx)?;
     863              : 
     864         2000 :             all_data.insert(key_int, idx);
     865              :         }
     866              : 
     867            2 :         let (root_offset, _writer) = writer.finish()?;
     868              : 
     869            2 :         let reader = DiskBtreeReader::new(0, root_offset, disk);
     870            2 : 
     871            2 :         reader.dump().await?;
     872              : 
     873              :         use std::sync::Mutex;
     874              : 
     875            2 :         let result = Mutex::new(Vec::new());
     876            2 :         let limit: AtomicUsize = AtomicUsize::new(10);
     877        83820 :         let take_ten = |key: &[u8], value: u64| {
     878        83820 :             let mut keybuf = [0u8; 8];
     879        83820 :             keybuf.copy_from_slice(key);
     880        83820 :             let key_int = u64::from_be_bytes(keybuf);
     881        83820 : 
     882        83820 :             let mut result = result.lock().unwrap();
     883        83820 :             result.push((key_int, value));
     884        83820 : 
     885        83820 :             // keep going until we have 10 matches
     886        83820 :             result.len() < limit.load(Ordering::Relaxed)
     887        83820 :         };
     888              : 
     889         4020 :         for search_key_int in 0..(NUM_KEYS * 2 + 10) {
     890         4020 :             let search_key = u64::to_be_bytes(search_key_int);
     891         4020 :             assert_eq!(
     892         4020 :                 reader.get(&search_key, &ctx).await?,
     893         4020 :                 all_data.get(&search_key_int).cloned()
     894              :             );
     895              : 
     896              :             // Test a forward scan starting with this key
     897         4020 :             result.lock().unwrap().clear();
     898         4020 :             reader
     899         4020 :                 .visit(&search_key, VisitDirection::Forwards, take_ten, &ctx)
     900            0 :                 .await?;
     901         4020 :             let expected = all_data
     902         4020 :                 .range(search_key_int..)
     903         4020 :                 .take(10)
     904        39820 :                 .map(|(&key, &val)| (key, val))
     905         4020 :                 .collect::<Vec<(u64, u64)>>();
     906         4020 :             assert_eq!(*result.lock().unwrap(), expected);
     907              : 
     908              :             // And a backwards scan
     909         4020 :             result.lock().unwrap().clear();
     910         4020 :             reader
     911         4020 :                 .visit(&search_key, VisitDirection::Backwards, take_ten, &ctx)
     912            0 :                 .await?;
     913         4020 :             let expected = all_data
     914         4020 :                 .range(..=search_key_int)
     915         4020 :                 .rev()
     916         4020 :                 .take(10)
     917        40000 :                 .map(|(&key, &val)| (key, val))
     918         4020 :                 .collect::<Vec<(u64, u64)>>();
     919         4020 :             assert_eq!(*result.lock().unwrap(), expected);
     920              :         }
     921              : 
     922              :         // full scan
     923            2 :         let search_key = u64::to_be_bytes(0);
     924            2 :         limit.store(usize::MAX, Ordering::Relaxed);
     925            2 :         result.lock().unwrap().clear();
     926            2 :         reader
     927            2 :             .visit(&search_key, VisitDirection::Forwards, take_ten, &ctx)
     928            0 :             .await?;
     929            2 :         let expected = all_data
     930            2 :             .iter()
     931         2000 :             .map(|(&key, &val)| (key, val))
     932            2 :             .collect::<Vec<(u64, u64)>>();
     933            2 :         assert_eq!(*result.lock().unwrap(), expected);
     934              : 
     935              :         // full scan
     936            2 :         let search_key = u64::to_be_bytes(u64::MAX);
     937            2 :         limit.store(usize::MAX, Ordering::Relaxed);
     938            2 :         result.lock().unwrap().clear();
     939            2 :         reader
     940            2 :             .visit(&search_key, VisitDirection::Backwards, take_ten, &ctx)
     941            0 :             .await?;
     942            2 :         let expected = all_data
     943            2 :             .iter()
     944            2 :             .rev()
     945         2000 :             .map(|(&key, &val)| (key, val))
     946            2 :             .collect::<Vec<(u64, u64)>>();
     947            2 :         assert_eq!(*result.lock().unwrap(), expected);
     948              : 
     949            2 :         Ok(())
     950              :     }
     951              : 
     952            2 :     #[tokio::test]
     953            2 :     async fn random_data() -> Result<()> {
     954            2 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
     955            2 : 
     956            2 :         // Generate random keys with exponential distribution, to
     957            2 :         // exercise the prefix compression
     958            2 :         const NUM_KEYS: usize = 100000;
     959            2 :         let mut all_data: BTreeMap<u128, u64> = BTreeMap::new();
     960       200002 :         for idx in 0..NUM_KEYS {
     961       200000 :             let u: f64 = rand::thread_rng().gen_range(0.0..1.0);
     962       200000 :             let t = -(f64::ln(u));
     963       200000 :             let key_int = (t * 1000000.0) as u128;
     964       200000 : 
     965       200000 :             all_data.insert(key_int, idx as u64);
     966       200000 :         }
     967              : 
     968              :         // Build a tree from it
     969            2 :         let mut disk = TestDisk::new();
     970            2 :         let mut writer = DiskBtreeBuilder::<_, 16>::new(&mut disk);
     971              : 
     972       195093 :         for (&key, &val) in all_data.iter() {
     973       195093 :             writer.append(&u128::to_be_bytes(key), val)?;
     974              :         }
     975            2 :         let (root_offset, _writer) = writer.finish()?;
     976              : 
     977            2 :         let reader = DiskBtreeReader::new(0, root_offset, disk);
     978              : 
     979              :         // Test get() operation on all the keys
     980       195093 :         for (&key, &val) in all_data.iter() {
     981       195093 :             let search_key = u128::to_be_bytes(key);
     982       195093 :             assert_eq!(reader.get(&search_key, &ctx).await?, Some(val));
     983              :         }
     984              : 
     985              :         // Test get() operations on random keys, most of which will not exist
     986       200002 :         for _ in 0..100000 {
     987       200000 :             let key_int = rand::thread_rng().gen::<u128>();
     988       200000 :             let search_key = u128::to_be_bytes(key_int);
     989       200000 :             assert!(reader.get(&search_key, &ctx).await? == all_data.get(&key_int).cloned());
     990              :         }
     991              : 
     992              :         // Test boundary cases
     993            2 :         assert!(
     994            2 :             reader.get(&u128::to_be_bytes(u128::MIN), &ctx).await?
     995            2 :                 == all_data.get(&u128::MIN).cloned()
     996              :         );
     997            2 :         assert!(
     998            2 :             reader.get(&u128::to_be_bytes(u128::MAX), &ctx).await?
     999            2 :                 == all_data.get(&u128::MAX).cloned()
    1000              :         );
    1001              : 
    1002            2 :         Ok(())
    1003              :     }
    1004              : 
    1005            2 :     #[test]
    1006            2 :     fn unsorted_input() {
    1007            2 :         let mut disk = TestDisk::new();
    1008            2 :         let mut writer = DiskBtreeBuilder::<_, 2>::new(&mut disk);
    1009            2 : 
    1010            2 :         let _ = writer.append(b"ba", 1);
    1011            2 :         let _ = writer.append(b"bb", 2);
    1012            2 :         let err = writer.append(b"aa", 3).expect_err("should've failed");
    1013            2 :         match err {
    1014            2 :             DiskBtreeError::UnsortedInput { key, last_key } => {
    1015            2 :                 assert_eq!(key.as_ref(), b"aa".as_slice());
    1016            2 :                 assert_eq!(last_key.as_ref(), b"bb".as_slice());
    1017              :             }
    1018            0 :             _ => panic!("unexpected error variant, expected DiskBtreeError::UnsortedInput"),
    1019              :         }
    1020            2 :     }
    1021              : 
    1022              :     ///
    1023              :     /// This test contains a particular data set, see disk_btree_test_data.rs
    1024              :     ///
    1025            2 :     #[tokio::test]
    1026            2 :     async fn particular_data() -> Result<()> {
    1027            2 :         // Build a tree from it
    1028            2 :         let mut disk = TestDisk::new();
    1029            2 :         let mut writer = DiskBtreeBuilder::<_, 26>::new(&mut disk);
    1030            2 :         let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error);
    1031              : 
    1032         4002 :         for (key, val) in disk_btree_test_data::TEST_DATA {
    1033         4000 :             writer.append(&key, val)?;
    1034              :         }
    1035            2 :         let (root_offset, writer) = writer.finish()?;
    1036              : 
    1037            2 :         println!("SIZE: {} blocks", writer.blocks.len());
    1038            2 : 
    1039            2 :         let reader = DiskBtreeReader::new(0, root_offset, disk);
    1040              : 
    1041              :         // Test get() operation on all the keys
    1042         4002 :         for (key, val) in disk_btree_test_data::TEST_DATA {
    1043         4000 :             assert_eq!(reader.get(&key, &ctx).await?, Some(val));
    1044              :         }
    1045              : 
    1046              :         // Test full scan
    1047            2 :         let mut count = 0;
    1048            2 :         reader
    1049            2 :             .visit(
    1050            2 :                 &[0u8; 26],
    1051            2 :                 VisitDirection::Forwards,
    1052         4000 :                 |_key, _value| {
    1053         4000 :                     count += 1;
    1054         4000 :                     true
    1055         4000 :                 },
    1056            2 :                 &ctx,
    1057            2 :             )
    1058            0 :             .await?;
    1059            2 :         assert_eq!(count, disk_btree_test_data::TEST_DATA.len());
    1060              : 
    1061            2 :         reader.dump().await?;
    1062              : 
    1063            2 :         Ok(())
    1064              :     }
    1065              : }
    1066              : 
    1067              : #[cfg(test)]
    1068              : #[path = "disk_btree_test_data.rs"]
    1069              : mod disk_btree_test_data;
        

Generated by: LCOV version 2.1-beta