infinitree 0.11.0

Embedded, encrypted database with tiered cache
Documentation
//! Store an encrypted, indexed stream of arbitrary data in the object store
use super::{store, Action, RawAction};
use crate::{
    fields::{
        depth::Incremental, Collection, Intent, List, Load, LocalField, SparseField, Store,
        Strategy, Value,
    },
    index::{reader, writer, FieldReader, FieldWriter},
    object::{self, ObjectError, StreamChunks},
    ChunkPointer, Digest,
};
use itertools::Itertools;
use parking_lot::RwLock;
use scc::ebr::{Arc, AtomicArc, Barrier, Tag};
use serde::{de::DeserializeOwned, Serialize};
use std::{
    collections::BTreeMap,
    io,
    ops::{Bound, RangeBounds},
    sync::atomic::Ordering,
};

enum RangeAction {
    Skip,
    Take,
    End,
}

struct StreamInner<Index: 'static + Ord + Sync + Send + Clone> {
    current: RwLock<BTreeMap<Index, Arc<StreamChunks>>>,
    base: RwLock<BTreeMap<Index, Arc<StreamChunks>>>,
    max: AtomicArc<Index>,
    min: AtomicArc<Index>,
}

impl<Index> Default for StreamInner<Index>
where
    Index: 'static + Ord + Send + Sync + Clone,
{
    fn default() -> Self {
        Self {
            current: RwLock::default(),
            base: RwLock::default(),
            max: AtomicArc::default(),
            min: AtomicArc::default(),
        }
    }
}

#[derive(Clone)]
pub struct Stream<Index>
where
    Index: 'static + Ord + Send + Sync + Clone,
{
    inner: Arc<StreamInner<Index>>,
}

impl<Index> Default for Stream<Index>
where
    Index: 'static + Ord + Send + Sync + Clone,
{
    fn default() -> Self {
        Self {
            inner: Arc::new(StreamInner::default()),
        }
    }
}

impl<Index: Send + Sync + Ord + Clone> Stream<Index> {
    pub fn commit(&self) {
        let barrier = Barrier::new();

        for (k, v) in self.inner.current.read().iter() {
            self.inner.base.write().insert(k.clone(), v.clone());
        }

        self.inner.current.write().clear();
    }

    pub fn writer<W: object::Writer>(&self, index: W) -> StreamWriter<'_, Index, W> {
        StreamWriter {
            stream: &self.inner,
            barrier: Barrier::new(),
            writer,
        }
    }

    pub fn reader<R: object::Reader>(&self, reader: R) -> StreamReader<'_, Index, R, Box<[u8]>> {
        self.reader_with_buffer(reader, vec![0; crate::BLOCK_SIZE].into_boxed_slice())
    }

    pub fn reader_with_buffer<R: object::Reader, B: AsRef<[u8]> + AsMut<[u8]> + 'static>(
        &self,
        reader: R,
        buffer: B,
    ) -> StreamReader<'_, Index, R, B> {
        StreamReader {
            stream: &self.inner,
            barrier: Barrier::new(),
            reader,
            buffer,
        }
    }

    pub fn clear(&self) {
        self.inner.current.write().clear();
        self.inner.base.write().clear();
    }
}

// impl<'stream, Index, R, Buf> StreamReader<'stream, Index, R, Buf>
// where
//     R: object::Reader,
//     Index: 'static + Ord + Sync + Send + Clone,
//     Buf: AsRef<[u8]> + AsMut<[u8]>,
// {
//     /// Query a range from the stream
//     ///
//     /// # Examples
//     /// ```
//     /// println!("TODO");
//     /// // TODO
//     /// ```
//     pub fn range<B>(&mut self, range: B) -> impl Iterator<Item = u8> + '_
//     where
//         B: RangeBounds<Index> + Clone + 'static,
//     {
//         itertools::unfold(
//             vec![
//                 self.stream.base.read().iter(),
//                 self.stream.current.read().iter(),
//             ]
//             .into_iter()
//             .kmerge_by(|a, b| a.0 < b.0)
//             .peekable(),
//             move |iter| -> Option<(
//                 &Index,
//                 &Arc<parking_lot::lock_api::RwLock<parking_lot::RawRwLock, Vec<ChunkPointer>>>,
//             )> {
//                 let start_bound = range.start_bound();
//                 let end_bound = range.end_bound();

//                 loop {
//                     let next = iter.next();
//                     let peek = iter.peek();

//                     if next.is_none() {
//                         return None;
//                     }

//                     use {Bound::*, RangeAction::*};
//                     match (
//                         end_bound,
//                         next,
//                         match (next, peek, start_bound) {
//                             (_, _, Unbounded) => Take,
//                             (None, _, _) => End,
//                             (Some(next), Some(peek), Included(v)) => {
//                                 if v >= next.0 && v < peek.0 {
//                                     Take
//                                 } else if v <= next.0 {
//                                     Take
//                                 } else {
//                                     Skip
//                                 }
//                             }
//                             (Some(next), Some(peek), Excluded(v)) => {
//                                 if v > next.0 && v < peek.0 {
//                                     Take
//                                 } else if v <= next.0 {
//                                     Take
//                                 } else {
//                                     Skip
//                                 }
//                             }
//                             (Some(next), None, Included(v)) => {
//                                 if v >= next.0 {
//                                     Take
//                                 } else {
//                                     Skip
//                                 }
//                             }
//                             (Some(next), None, Excluded(v)) => {
//                                 if v > next.0 {
//                                     Take
//                                 } else {
//                                     Skip
//                                 }
//                             }
//                         },
//                     ) {
//                         (_, _, End) => return None,
//                         (_, None, _) => return None,
//                         (_, _, Skip) => continue,

//                         (Unbounded, Some(next), Take) => return Some(next),
//                         (Included(v), Some(next), Take) => {
//                             if v <= next.0 {
//                                 return Some(next);
//                             } else {
//                                 return None;
//                             }
//                         }
//                         (Excluded(v), Some(next), Take) => {
//                             if v < next.0 {
//                                 return Some(next);
//                             } else {
//                                 return None;
//                             }
//                         }
//                     }
//                 }
//             },
//         )
//         .flat_map(|(_, ptrs)| ptrs.read().to_vec())
//         .flat_map(|ref ptr| {
//             self.reader
//                 .read_chunk(ptr, self.buffer.as_mut())
//                 .map(Vec::from)
//         })
//         .flatten()
//     }

//     pub fn min(&self) -> Arc<Index> {
//         self.stream
//             .min
//             .load(Ordering::Relaxed, &self.barrier)
//             .get_arc()
//             .unwrap()
//     }

//     pub fn max(&self) -> Arc<Index> {
//         self.stream
//             .max
//             .load(Ordering::Relaxed, &self.barrier)
//             .get_arc()
//             .unwrap()
//     }
// }

impl<Idx> crate::Index for Stream<Idx>
where
    Idx: 'static + Ord + Send + Sync + Clone + DeserializeOwned + Serialize,
{
    fn store_all(&mut self) -> anyhow::Result<Vec<Intent<Box<dyn Store>>>> {
        Ok(vec![Intent::new(
            "root",
            Box::new(LocalField::for_field(self)),
        )])
    }

    fn load_all(&mut self) -> anyhow::Result<Vec<Intent<Box<dyn Load>>>> {
        Ok(vec![Intent::new(
            "root",
            Box::new(LocalField::for_field(self)),
        )])
    }
}

impl<Index> Store for LocalField<Stream<Index>>
where
    Index: 'static + Ord + Send + Sync + Clone + Serialize,
{
    #[inline(always)]
    fn store(
        &mut self,
        transaction: &mut writer::Transaction<'_>,
        _object: &mut dyn object::Writer,
    ) {
        let barrier = Barrier::new();
        for (k, v) in self.field.inner.current.read().iter() {
            transaction.write_next((k, v.read().to_vec()));
        }

        self.field.commit();
    }
}

impl<Index> Collection for LocalField<Stream<Index>>
where
    Index: 'static + Ord + Send + Sync + Clone + DeserializeOwned,
{
    type Depth = Incremental;
    type Key = Index;
    type Serialized = (Index, Vec<ChunkPointer>);
    type Item = (Index, Vec<ChunkPointer>);

    #[inline(always)]
    fn key(from: &Self::Serialized) -> &Self::Key {
        &from.0
    }

    #[inline(always)]
    fn load(from: Self::Serialized, _object: &mut dyn crate::object::Reader) -> Self::Item {
        from
    }

    #[inline(always)]
    fn insert(&mut self, record: Self::Item) {
        self.field
            .inner
            .base
            .write()
            .insert(record.0, Arc::new(RwLock::new(record.1)));
    }
}

/// This is the inverse of [`lz4_flex::compress::get_maximum_output_size`]
fn lz4_max_input_size_for(output_len: usize) -> usize {
    ((output_len - 16 - 4) as f64 / 1.1) as usize
}

#[cfg(test)]
mod test {
    #[test]
    fn faszom() {
        use crate::{backends::Directory, fields::Stream, Infinitree, Key};
        use std::io::Write;

        let tree = Infinitree::<Stream<usize>>::empty(
            Directory::new("test_data").unwrap(),
            Key::from_credentials("username", "password").unwrap(),
        );

        let index = tree.index();

        // Writer from this point is thread-locked.
        let mut writer = index.writer(tree.object_writer().unwrap());

        // // generate 8MiB of data
        let buffer = vec![123; 8 * 1024 * 1024];
        assert_eq!(writer.write(&buffer).unwrap(), 8 * 1024 * 1024);

        drop(writer);

        let mut reader = index.reader(tree.object_reader().unwrap());
        // let output = reader.range(..).collect::<Vec<u8>>();

        // assert_eq!(output.len(), buffer.len());
    }
}