llkv_btree/bplus_tree/

shared_bplus_tree.rs

use std::sync::{Mutex as StdMutex, RwLock};

use crossbeam_channel as xchan;

use crate::bplus_tree::BPlusTree;
use crate::codecs::{IdCodec, KeyCodec};
use crate::errors::Error;
use crate::iter::{BPlusTreeIter, ScanOpts};
use crate::node_cache::NodeCache;
use crate::pager::Pager;
use crate::traits::BTree;
use std::sync::Arc;

/// A façade that provides a serialized writer and snapshot readers,
/// while implementing the same `BTree` trait.
pub struct SharedBPlusTree<P, KC, IC>
where
    P: Pager + Clone,
    P::Page: Send + Sync + 'static,
    P::Id: Send + Sync + 'static,
    KC: KeyCodec,
    IC: IdCodec<Id = P::Id>,
{
    pager: P,
    writer: StdMutex<BPlusTree<P, KC, IC>>, // single-threaded writer
    latest_root: RwLock<P::Id>,             // published root for readers
    shared_node_cache: Arc<RwLock<NodeCache<P>>>,
    // Cache page size so readers do not lock the writer.
    page_size: usize,
}

impl<P, KC, IC> SharedBPlusTree<P, KC, IC>
where
    P: Pager + Clone,
    P::Page: Send + Sync + 'static,
    P::Id: Send + Sync + 'static,
    KC: KeyCodec,
    IC: IdCodec<Id = P::Id>,
{
    /// Create an empty tree + shared writer.
    pub fn create_empty(
        pager: P,
        opt_shared_node_cache: Option<Arc<RwLock<NodeCache<P>>>>,
    ) -> Result<Self, Error> {
        let shared_node_cache = BPlusTree::<P, KC, IC>::init_node_cache(opt_shared_node_cache);

        let writer = BPlusTree::<P, KC, IC>::create_empty(
            pager.clone(),
            Some(Arc::clone(&shared_node_cache)),
        )?;
        let root = writer.root_id();
        let page_size = writer.page_size();

        Ok(Self {
            pager,
            writer: StdMutex::new(writer),
            latest_root: RwLock::new(root),
            shared_node_cache,
            page_size,
        })
    }

    /// Open a shared B+Tree backed by an existing root.
    pub fn open(
        pager: P,
        root: P::Id,
        opt_shared_node_cache: Option<Arc<RwLock<NodeCache<P>>>>,
    ) -> Self {
        let shared_node_cache = BPlusTree::<P, KC, IC>::init_node_cache(opt_shared_node_cache);

        let writer = BPlusTree::<P, KC, IC>::new(
            pager.clone(),
            root.clone(),
            Some(Arc::clone(&shared_node_cache)),
        );
        let page_size = writer.page_size();

        Self {
            pager,
            writer: StdMutex::new(writer),
            latest_root: RwLock::new(root),
            shared_node_cache,
            page_size,
        }
    }

    /// Produce a **snapshot** tree at the current published root.
    /// This is a normal `BPlusTree` you can iterate with `ScanOpts`.
    pub fn snapshot(&self) -> BPlusTree<P, KC, IC> {
        let root = { self.latest_root.read().unwrap().clone() };
        let shared_node_cache = Arc::clone(&self.shared_node_cache);

        BPlusTree::<P, KC, IC>::with_readonly_pager_state(
            self.pager.clone(),
            root,
            shared_node_cache,
            self.page_size,
        )
    }

    // TODO: Consider putting this behind a feature flag and only using it for
    // local testing purposes. This is good for some integration and benchmark
    // testing, but callers wil likely benefit from manually snapshotting on the
    // thread they choose to run it on.
    //
    /// Start a streaming scan at the current root, using any `ScanOpts`.
    /// Returns a crossbeam MPMC receiver; read it from as many threads as you
    /// like. Keys/values are owned (no borrows across threads).
    pub fn start_stream_with_opts(
        &self,
        opts: ScanOpts<'static, KC>,
    ) -> xchan::Receiver<(KC::Key, Vec<u8>)>
    where
        // required because we move these into a rayon worker thread
        P: Send + 'static,
        P::Id: Send + 'static,
        KC: KeyCodec + 'static,
        KC::Key: Send + Sync + 'static,
        IC: IdCodec<Id = P::Id> + 'static,
    {
        let root = { self.latest_root.read().unwrap().clone() };
        let pager = self.pager.clone();

        // TODO: Make the capacity configurable
        let (tx, rx) = xchan::bounded(1);

        let shared_node_cache = Arc::clone(&self.shared_node_cache);
        let page_size = self.page_size;

        rayon::spawn(move || {
            let tree = BPlusTree::<P, KC, IC>::with_readonly_pager_state(
                pager,
                root,
                shared_node_cache,
                page_size,
            );
            if let Ok(it) = BPlusTreeIter::with_opts(&tree, opts) {
                for (kref, vref) in it {
                    let key = KC::decode_from(kref.as_ref()).expect("decode key");
                    let val = vref.as_ref().to_vec();
                    if tx.send((key, val)).is_err() {
                        break;
                    }
                    // std::thread::yield_now(); // encourages interleaving
                }
            }
            // drop(tx) ends the stream
        });

        rx
    }

    // Optionally expose a publish step (useful if you batch).
    // pub fn flush(&self) -> Result<(), Error> {
    //     let mut w = self.writer.lock().unwrap();
    //     w.flush()?;
    //     *self.latest_root.write().unwrap() = w.root_id();
    //     Ok(())
    // }
}

// --- Implement the public BTree trait on SharedBPlusTree ---

impl<'a, P, KC, IC> BTree<'a, KC::Key> for SharedBPlusTree<P, KC, IC>
where
    P: Pager + Clone + 'a,
    P::Page: Send + Sync + 'static,
    P::Id: Send + Sync + 'static,
    KC: KeyCodec + 'a,
    IC: IdCodec<Id = P::Id> + 'a,
{
    type Value = <BPlusTree<P, KC, IC> as BTree<'a, KC::Key>>::Value;

    fn get_many(&'a self, keys: &[KC::Key]) -> Result<Vec<Option<Self::Value>>, Error> {
        self.snapshot().get_many(keys)
    }

    fn contains_key(&'a self, key: &KC::Key) -> Result<bool, Error> {
        self.snapshot().contains_key(key)
    }

    fn insert_many(&self, items: &[(KC::Key, &[u8])]) -> Result<(), Error>
    where
        KC::Key: Clone,
    {
        // Perform write work under the writer lock, then release it
        // before publishing the new root to avoid cross-lock waits.
        let new_root = {
            let w = self.writer.lock().unwrap();
            w.insert_many(items)?;
            w.flush()?;
            w.root_id()
        };
        *self.latest_root.write().unwrap() = new_root;
        Ok(())
    }

    fn delete_many(&self, keys: &[KC::Key]) -> Result<(), Error>
    where
        KC::Key: Clone,
    {
        let new_root = {
            let w = self.writer.lock().unwrap();
            w.delete_many(keys)?;
            w.flush()?;
            w.root_id()
        };
        *self.latest_root.write().unwrap() = new_root;
        Ok(())
    }
}