armour/database/

raw_tree.rs

use std::{
    ops::Bound,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
};

use async_broadcast::{Receiver, broadcast};
use derive_more::Debug;
use xxhash_rust::xxh3::Xxh3Default;

use super::{ByteValue, KeyValue};
use crate::{
    error::{DbError, DbResult},
    replication::{ChangeEvent, ReplicationEvent, ReplicationEventHandler, Subscribers},
    types::{ArmourError, attribute::EntityAttribute, num_ops::g4bits},
    utils::{CheckSumVec, CollectionInfo, GroupVal, HashPoints},
};

#[derive(Debug)]
pub(crate) struct InnerFields {
    pub(crate) info: CollectionInfo,
    /// map of groups hashes
    pub(crate) hashpoints: HashPoints,
    /// approximate number of elements in the collection
    pub(crate) seq: AtomicU64,
    #[debug("{}", self.replication_handler.is_some())]
    pub(crate) replication_handler: ReplicationEventHandler<ByteValue, ByteValue>,
}

impl InnerFields {
    pub(crate) fn invalidate_hash(&self, group_id: u32) {
        self.hashpoints.insert(
            group_id,
            GroupVal {
                hash: 0,
                changed: true,
            },
        );
    }
}

#[derive(Clone, Debug)]
pub struct RawTree {
    /// custom name (with_name fn) or Record::NAME
    pub name: String,
    pub attributes: &'static EntityAttribute,
    #[debug(skip)]
    pub(crate) seq_tree: super::Tree,
    #[debug(skip)]
    pub(crate) tree: super::Tree,
    #[debug(skip)]
    pub(crate) removed: super::Tree,
    pub(crate) inner: Arc<InnerFields>,
    pub(crate) subscribers: Arc<Subscribers<ByteValue, ByteValue>>,
}

impl RawTree {
    #[instrument(level = "info", skip(self), fields(name = self.name, ret))]
    pub fn checksum(&self) -> u32 {
        let start = std::time::Instant::now();
        let checksum = self.tree.checksum().expect("cannot calculate checksum");
        if checksum != 0 {
            debug!("checksum: {checksum:#X}");
        }

        histogram!("armour_sled_rawtree_checksum_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_checksum_total", "name" => self.name.clone()).increment(1);

        checksum
    }

    pub fn is_empty(&self) -> bool {
        self.tree.is_empty().expect("cannot check if tree is empty")
    }

    #[instrument(level = "trace", skip(self), fields(name = self.name), ret)]
    pub fn count(&self) -> u64 {
        self.inner.seq.load(Ordering::Relaxed)
    }

    #[instrument(level = "info", skip(self), fields(name = self.name))]
    pub fn subscribe(&self) -> Receiver<ChangeEvent<ByteValue, ByteValue>> {
        // let subscriber = self.tree.watch_prefix([]);
        // let f = futures::stream::unfold(subscriber, |mut sub| async move {
        //     (&mut sub).await.map(|event| {
        //         let event_batch: Box<dyn EventBatch<Bytes = ByteValue>> =
        //             Box::new(SledEvent(event));
        //         (event_batch, sub)
        //     })
        // });
        // f.boxed_local()

        let (_, r) = self.subscribers.get_or_init(|| {
            let (mut s, r) = broadcast(1024);
            let r = r.deactivate();
            s.set_await_active(false);
            s.set_overflow(true);
            (s, r)
        });
        r.activate_cloned()
    }

    #[instrument(level = "info", skip(self), fields(name = self.name, ret))]
    pub fn hashpoints(&self) -> CheckSumVec {
        let start = std::time::Instant::now();
        let res = self
            .inner
            .hashpoints
            .iter()
            .map(|entry| {
                let key = *entry.key();
                (key, entry.value().hash)
            })
            .collect();

        histogram!("armour_sled_rawtree_hashpoints_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_hashpoints_total", "name" => self.name.clone()).increment(1);

        res
    }

    #[instrument(level = "trace", skip(self), fields(name = self.name))]
    pub fn scan_group(&self, group: u32) -> impl DoubleEndedIterator<Item = KeyValue> + use<> {
        // for ID: 0x0
        // for ID: 0x0
        // for Fuid: 0x0 - April 09 2023 10:13:00 CET?
        let start_bytes = group.to_be_bytes();

        // for ID: 24 bits
        // for Id: 32 bits
        // for Fuid: 10 bits
        // for LowId: 8 bits
        let group_bits = self.attributes.group_bits;
        // for ID: 32 - 24 = 8 bits
        // for Id: 32 - 32 = 0 bits
        // for Fuid: 32 - 10 = 22 bits
        let bits_sub = u32::BITS - group_bits;
        // for ID: 2^8 = 256 = 0x100
        // for Id: 2^0 = 1 = 0x1
        // for Fuid: 2^22 = 4_194_304 = 0x40_0000
        let bits_pow_of_two = 2u32.pow(bits_sub);
        // for ID: 0x0 + 2^8 = 0x100
        // for Id: 0x0 + 2^0 = 0x1
        // for Fuid: 0x0 + 2^22 = 0x40_0000
        let end = group + bits_pow_of_two;
        let end_bytes = end.to_be_bytes();

        let start_bound = Bound::Included(start_bytes.as_ref());
        let end_bound = Bound::Excluded(end_bytes.as_ref());

        let res = self
            .tree
            .range::<&[u8], _>((start_bound, end_bound))
            .filter_map(|item| match item {
                Ok((key, value)) => {
                    // ...
                    Some((key, value))
                }
                Err(e) => {
                    error!(%e);
                    None
                }
            });

        counter!("armour_sled_rawtree_scan_group_total", "name" => self.name.clone()).increment(1);
        res
    }

    #[instrument(level = "info", skip(self), fields(name = self.name, ret))]
    pub fn recalcucate_hash(&self) -> u64 {
        let start = std::time::Instant::now();
        let hash = self
            .inner
            .hashpoints
            .iter()
            .map(|item| {
                let group_val = item.value();

                if group_val.changed {
                    let group = *item.key();
                    drop(item);
                    let mut hash_val = Xxh3Default::new();

                    for (key, value) in self.scan_group(group) {
                        hash_val.update(&key);
                        hash_val.update(&value);
                    }
                    let hash = hash_val.digest();
                    // TODO: deadlock? race condition?
                    self.inner.hashpoints.insert(
                        group,
                        GroupVal {
                            hash,
                            changed: false,
                        },
                    );
                    hash
                } else {
                    item.value().hash
                }
            })
            .fold(Xxh3Default::new(), |mut hasher, item| {
                hasher.update(&item.to_le_bytes());
                hasher
            });

        histogram!("armour_sled_rawtree_recalcucate_hash_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_recalcucate_hash_total", "name" => self.name.clone())
            .increment(1);

        hash.digest()
    }

    #[instrument(skip(self), fields(name = self.name))]
    pub fn tree_info(&self) -> CollectionInfo {
        let seq = self.inner.seq.load(Ordering::SeqCst);

        if seq != 0 {
            debug!(seq, "close seq");
        }

        let typ_hash = self.attributes.ty.h();
        let version = self.attributes.version;

        CollectionInfo {
            seq,
            typ_hash,
            version,
        }
    }

    /// return slice with [key + value] bytes
    #[instrument(level = "trace", skip(self), fields(name = self.name, ret))]
    pub fn get(&self, id: &[u8]) -> DbResult<Option<Vec<u8>>> {
        let start = std::time::Instant::now();
        let res = self
            .tree
            .get(id)
            .map(|item| {
                item.map(|item| {
                    let len = id.len() + item.len();
                    let mut v = vec![0; len];
                    v[..id.len()].copy_from_slice(id);
                    v[id.len()..].copy_from_slice(&item);
                    v
                })
            })
            .map_err(DbError::from);

        histogram!("armour_sled_rawtree_get_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_get_total", "name" => self.name.clone()).increment(1);

        res
    }

    #[instrument(level = "trace", skip(self), fields(name = self.name))]
    pub fn range(
        &self,
        range: (Bound<&[u8]>, Bound<&[u8]>),
    ) -> impl DoubleEndedIterator<Item = KeyValue> {
        let iter = self.tree.range::<&[u8], _>(range);
        let iter = iter.filter_map(|item| match item {
            Ok((key, value)) => Some((key, value)),
            Err(e) => {
                error!(%e);
                None
            }
        });

        counter!("armour_sled_rawtree_range_total", "name" => self.name.clone()).increment(1);

        iter
    }

    #[instrument(level = "trace", skip(self), fields(name = self.name))]
    fn invalidate_hash(&self, id: &[u8]) {
        let start = std::time::Instant::now();
        let mut bytes = [0; 4];
        bytes.copy_from_slice(&id[..4]);
        let group = u32::from_be_bytes(bytes);
        let group = g4bits(group, self.attributes.group_bits);
        self.inner.invalidate_hash(group);

        histogram!("armour_sled_rawtree_invalidate_hash_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_invalidate_hash_total", "name" => self.name.clone())
            .increment(1);
    }

    #[instrument(level = "debug", skip(self), fields(name = self.name), err, ret)]
    pub fn next_id(&self) -> DbResult<u64> {
        let start = std::time::Instant::now();
        let next_id_key = format!("next_id-{}", self.name);

        let key_ref = next_id_key.as_bytes();
        let mut current = self.seq_tree.get(key_ref)?;

        let res = (|| loop {
            let tmp = current.as_ref().map(AsRef::as_ref);

            let mut id = 1u64;

            let next = match tmp {
                Some(bytes) => {
                    let bytes = bytes
                        .try_into()
                        .map_err(|err| DbError::Armour(ArmourError::from(err)))?;
                    let old = u64::from_le_bytes(bytes);
                    id = old + 1;
                    id.to_le_bytes().to_vec()
                }
                None => id.to_le_bytes().to_vec(),
            };

            match self.seq_tree.compare_and_swap(key_ref, tmp, Some(next))? {
                Ok(_) => return Ok(id),
                Err(sled::CompareAndSwapError { current: cur, .. }) => {
                    current = cur;
                }
            }
        })();

        histogram!("armour_sled_rawtree_next_id_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_next_id_total", "name" => self.name.clone()).increment(1);

        res
    }

    #[instrument(level = "debug", skip(self, event), fields(name = self.name, event = event.variant()), err)]
    pub fn apply_event(&self, event: ChangeEvent<ByteValue, ByteValue>) -> DbResult<()> {
        let start = std::time::Instant::now();
        match &event {
            ChangeEvent::Upsert((key, val)) => {
                let old = self.tree.insert(key, val.clone())?;
                if old.is_none() {
                    self.inner.seq.fetch_add(1, Ordering::Relaxed);
                }
                if let Some(f) = self.inner.replication_handler.as_ref() {
                    f(ReplicationEvent::Upsert {
                        key,
                        val,
                        old_val: old.as_ref(),
                    });
                }
            }
            ChangeEvent::Delete(key) => {
                match self.tree.remove(key)? {
                    Some(val) => {
                        self.inner.seq.fetch_sub(1, Ordering::AcqRel);
                        if let Some(f) = self.inner.replication_handler.as_ref() {
                            f(ReplicationEvent::Delete { key, val: &val });
                        }
                    }
                    _ => {
                        error!(?key, "delete not found");
                    }
                };
            }
            ChangeEvent::ChangeId(old, new) => {
                let val = self.tree.remove(old)?.ok_or(DbError::NotFound)?;
                if self.tree.insert(new, val.clone())?.is_some() {
                    error!(?new, "change_id already exists");
                };
                if let Some(f) = self.inner.replication_handler.as_ref() {
                    f(ReplicationEvent::IdChange {
                        old_key: old,
                        new_key: new,
                        val: &val,
                    });
                }
                self.invalidate_hash(new);
            }
        }
        self.invalidate_hash(event.key());
        if let Some((sender, _)) = self.subscribers.get() {
            match sender.try_broadcast(event) {
                Ok(None) => {}
                Ok(Some(item)) => {
                    error!(?item, "broadcast is full");
                }
                Err(e) => {
                    error!(%e);
                }
            }
        }

        histogram!("armour_sled_rawtree_apply_event_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_apply_event_total", "name" => self.name.clone()).increment(1);

        Ok(())
    }

    /// doesn't change indexes
    #[instrument(level = "info", skip(self, iter), fields(name = self.name), err)]
    pub fn apply_batch(
        &self,
        iter: impl Iterator<Item = (ByteValue, Option<ByteValue>)>,
    ) -> DbResult<()> {
        let start = std::time::Instant::now();
        let mut batch = sled::Batch::default();

        // TODO
        for (key, val) in iter {
            self.invalidate_hash(&key);

            match &val {
                Some(val) => {
                    batch.insert(key.clone(), val.clone());

                    if let Some(f) = self.inner.replication_handler.as_ref() {
                        let old = self.tree.get(&key)?;
                        f(ReplicationEvent::Upsert {
                            key: &key,
                            val,
                            old_val: old.as_ref(),
                        });
                    }
                }
                None => {
                    batch.remove(key.clone());
                    if let Some(f) = self.inner.replication_handler.as_ref() {
                        match self.tree.get(&key)? {
                            Some(item) => {
                                f(ReplicationEvent::Delete {
                                    key: &key,
                                    val: &item,
                                });
                            }
                            _ => {
                                error!(?key, "delete not found");
                            }
                        }
                    }
                }
            }
            if let Some((sender, _)) = self.subscribers.get() {
                let msg = ChangeEvent::from_kv(key, val);
                match sender.try_broadcast(msg) {
                    Ok(None) => {}
                    Ok(Some(item)) => {
                        error!(?item, "broadcast is full");
                    }
                    Err(e) => {
                        error!(%e);
                    }
                }
            }
        }

        self.tree.apply_batch(batch)?;

        histogram!("armour_sled_rawtree_apply_batch_duration", "name" => self.name.clone())
            .record(start.elapsed().as_secs_f64());
        counter!("armour_sled_rawtree_apply_batch_total", "name" => self.name.clone()).increment(1);

        Ok(())
    }

    #[doc(hidden)]
    pub fn inner(&self) -> &super::Tree {
        &self.tree
    }
}