diff --git a/Cargo.toml b/Cargo.toml index 4491e6f2..686c6c44 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,9 +6,12 @@ edition = "2021" [dependencies] utils = { path = "../utils" } +rocksdb = { version = "0.20.1", optional = true } +foundationdb = { version = "0.7.0", optional = true } +futures = { version = "0.3", optional = true } rand = "0.8.5" roaring = "0.10.1" -rocksdb = "0.20.1" +rayon = { version = "1.5.1", optional = true } serde = { version = "1.0", features = ["derive"]} ahash = { version = "0.8.0", features = ["serde"] } bitpacking = "0.8.4" @@ -21,7 +24,13 @@ xxhash-rust = { version = "0.8.5", features = ["xxh3"] } farmhash = "1.1.5" siphasher = "0.3" +[features] +default = ["foundation"] +rocks = ["rocksdb", "rayon"] +foundation = ["foundationdb", "futures"] + [dev-dependencies] +tokio = { version = "1.23", features = ["full"] } csv = "1.1" +rayon = { version = "1.5.1" } flate2 = { version = "1.0.17", features = ["zlib"], default-features = false } -rayon = "1.5.1" diff --git a/src/backend/foundationdb/bitmap.rs b/src/backend/foundationdb/bitmap.rs new file mode 100644 index 00000000..b815ce29 --- /dev/null +++ b/src/backend/foundationdb/bitmap.rs @@ -0,0 +1,126 @@ +use roaring::RoaringBitmap; + +const BITS: u32 = 128; +const WORD_SIZE: u32 = 8; +pub const BITS_PER_BLOCK: u32 = BITS * WORD_SIZE; +const BITS_MASK: u32 = BITS_PER_BLOCK - 1; + +pub struct DenseBitmap { + restore_value: u8, + restore_pos: usize, + pub block_num: u32, + pub bitmap: [u8; std::mem::size_of::() * WORD_SIZE as usize], +} + +impl DenseBitmap { + pub fn empty() -> Self { + Self { + block_num: 0, + restore_pos: 0, + restore_value: 0, + bitmap: [0; std::mem::size_of::() * WORD_SIZE as usize], + } + } + + pub fn full() -> Self { + Self { + block_num: 0, + restore_pos: 0, + restore_value: u8::MAX, + bitmap: [u8::MAX; std::mem::size_of::() * WORD_SIZE as usize], + } + } + + pub fn set(&mut self, index: u32) { + self.block_num = index / BITS_PER_BLOCK; + let index = index & BITS_MASK; + self.restore_pos = (index / 8) as usize; + self.bitmap[self.restore_pos] = 1 << (index & 7); + } + + #[cfg(test)] + pub fn set_or(&mut self, index: u32) { + let _index = index; + self.block_num = index / BITS_PER_BLOCK; + let index = index & BITS_MASK; + self.restore_pos = (index / 8) as usize; + self.bitmap[self.restore_pos] |= 1 << (index & 7); + } + + pub fn clear(&mut self, index: u32) { + self.block_num = index / BITS_PER_BLOCK; + let index = BITS_MASK - (index & BITS_MASK); + self.restore_pos = (index / 8) as usize; + self.bitmap[self.restore_pos] = !(1 << (index & 7)); + } + + pub fn reset(&mut self) { + self.bitmap[self.restore_pos] = self.restore_value; + } +} + +pub trait DeserializeBlock { + fn deserialize_block(&mut self, bytes: &[u8], block_num: u32); +} + +impl DeserializeBlock for RoaringBitmap { + fn deserialize_block(&mut self, bytes: &[u8], block_num: u32) { + debug_assert_eq!( + bytes.len(), + std::mem::size_of::() * WORD_SIZE as usize + ); + + for (word_num, word) in bytes.chunks_exact(std::mem::size_of::()).enumerate() { + match u128::from_le_bytes(word.try_into().unwrap()) { + 0 => continue, + u128::MAX => { + self.insert_range( + block_num * BITS_PER_BLOCK + word_num as u32 * 128 + ..(block_num * BITS_PER_BLOCK + word_num as u32 * 128) + 128, + ); + } + mut word => { + while word != 0 { + let trailing_zeros = word.trailing_zeros(); + self.insert( + block_num * BITS_PER_BLOCK + word_num as u32 * 128 + trailing_zeros, + ); + word ^= 1 << trailing_zeros; + } + } + } + } + + //println!("deserializing block {} {}", block_num, self.len()); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use roaring::RoaringBitmap; + + use crate::backend::foundationdb::bitmap::{DenseBitmap, DeserializeBlock, BITS_PER_BLOCK}; + + #[test] + fn serialize_bitmap_block() { + for range in [(0..128), (128..256), (5076..5093), (1762..19342)] { + let mut blocks = HashMap::new(); + let mut bitmap = RoaringBitmap::new(); + for item in range { + bitmap.insert(item); + blocks + .entry(item / BITS_PER_BLOCK) + .or_insert_with(DenseBitmap::empty) + .set_or(item); + } + let mut bitmap_blocks = RoaringBitmap::new(); + for (block_num, dense_bitmap) in blocks { + bitmap_blocks.deserialize_block(&dense_bitmap.bitmap, block_num); + } + + assert_eq!(bitmap, bitmap_blocks); + } + } +} diff --git a/src/backend/foundationdb/main.rs b/src/backend/foundationdb/main.rs new file mode 100644 index 00000000..b4114a75 --- /dev/null +++ b/src/backend/foundationdb/main.rs @@ -0,0 +1,21 @@ +use foundationdb::Database; + +use crate::Store; + +impl Store { + pub async fn open() -> crate::Result { + Ok(Self { + guard: unsafe { foundationdb::boot() }, + db: Database::default()?, + }) + } +} + +/* +impl Drop for Store { + fn drop(&mut self) { + self.guard.drop(); + self.db.drop(); + } +} +*/ diff --git a/src/backend/foundationdb/mod.rs b/src/backend/foundationdb/mod.rs index e69de29b..49b5b565 100644 --- a/src/backend/foundationdb/mod.rs +++ b/src/backend/foundationdb/mod.rs @@ -0,0 +1,125 @@ +use foundationdb::FdbError; + +use crate::{ + write::key::KeySerializer, AclKey, BitmapKey, BlobKey, Error, IndexKey, IndexKeyPrefix, LogKey, + Serialize, ValueKey, +}; + +pub mod bitmap; +pub mod main; +pub mod read; +pub mod write; + +pub const SUBSPACE_BITMAPS: u8 = b'b'; +pub const SUBSPACE_VALUES: u8 = b'v'; +pub const SUBSPACE_LOGS: u8 = b'l'; +pub const SUBSPACE_BLOBS: u8 = b'o'; +pub const SUBSPACE_INDEXES: u8 = b'i'; +pub const SUBSPACE_ACLS: u8 = b'c'; + +impl> Serialize for &IndexKey { + fn serialize(self) -> Vec { + let key = self.key.as_ref(); + KeySerializer::new(std::mem::size_of::>() + key.len() + 1) + .write(SUBSPACE_INDEXES) + .write(self.account_id) + .write(self.collection) + .write(self.field) + .write(key) + .write(self.document_id) + .finalize() + } +} + +impl Serialize for &IndexKeyPrefix { + fn serialize(self) -> Vec { + KeySerializer::new(std::mem::size_of::() + 1) + .write(SUBSPACE_INDEXES) + .write(self.account_id) + .write(self.collection) + .write(self.field) + .finalize() + } +} + +impl Serialize for &ValueKey { + fn serialize(self) -> Vec { + if self.family == 0 { + KeySerializer::new(std::mem::size_of::() + 1) + .write(SUBSPACE_VALUES) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(self.field) + .finalize() + } else { + KeySerializer::new(std::mem::size_of::() + 2) + .write(SUBSPACE_VALUES) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(u8::MAX) + .write(self.family) + .write(self.field) + .finalize() + } + } +} + +impl> Serialize for &BitmapKey { + fn serialize(self) -> Vec { + let key = self.key.as_ref(); + KeySerializer::new(std::mem::size_of::>() + key.len() + 1) + .write(SUBSPACE_BITMAPS) + .write(self.account_id) + .write(self.collection) + .write(self.family) + .write(self.field) + .write(key) + .write(self.block_num) + .finalize() + } +} + +impl> Serialize for &BlobKey { + fn serialize(self) -> Vec { + let hash = self.hash.as_ref(); + KeySerializer::new(std::mem::size_of::>() + hash.len() + 1) + .write(SUBSPACE_BLOBS) + .write(hash) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .finalize() + } +} + +impl Serialize for &AclKey { + fn serialize(self) -> Vec { + KeySerializer::new(std::mem::size_of::() + 1) + .write(SUBSPACE_ACLS) + .write_leb128(self.grant_account_id) + .write(u8::MAX) + .write_leb128(self.to_account_id) + .write(self.to_collection) + .write_leb128(self.to_document_id) + .finalize() + } +} + +impl Serialize for &LogKey { + fn serialize(self) -> Vec { + KeySerializer::new(std::mem::size_of::() + 1) + .write(SUBSPACE_LOGS) + .write(self.account_id) + .write(self.collection) + .write(self.change_id) + .finalize() + } +} + +impl From for Error { + fn from(error: FdbError) -> Self { + Self::InternalError(format!("FoundationDB error: {}", error.message())) + } +} diff --git a/src/backend/foundationdb/read.rs b/src/backend/foundationdb/read.rs new file mode 100644 index 00000000..6f2542e5 --- /dev/null +++ b/src/backend/foundationdb/read.rs @@ -0,0 +1,312 @@ +use std::{ + ops::{BitAndAssign, BitOrAssign}, + time::{Duration, Instant}, +}; + +use foundationdb::{ + options::{self, StreamingMode}, + Database, KeySelector, RangeOption, Transaction, +}; +use futures::StreamExt; +use roaring::RoaringBitmap; + +use crate::{ + query::{Operator, SortedId, UnsortedIds}, + write::key::{DeserializeBigEndian, KeySerializer}, + BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, Serialize, Store, ValueKey, BM_DOCUMENT_IDS, +}; + +use super::{ + bitmap::{DeserializeBlock, BITS_PER_BLOCK}, + SUBSPACE_INDEXES, +}; + +pub struct ReadTransaction<'x> { + db: &'x Database, + pub trx: Transaction, + trx_age: Instant, +} + +impl ReadTransaction<'_> { + #[inline(always)] + pub async fn get_value(&self, key: ValueKey) -> crate::Result> + where + U: Deserialize, + { + let key = key.serialize(); + + if let Some(bytes) = self.trx.get(&key, true).await? { + U::deserialize(&bytes).map(Some) + } else { + Ok(None) + } + } + + #[inline(always)] + pub async fn get_values(&self, keys: Vec) -> crate::Result>> + where + U: Deserialize, + { + let mut results = Vec::with_capacity(keys.len()); + + for key in keys { + results.push(self.get_value(key).await?); + } + + Ok(results) + } + + pub async fn get_document_ids( + &self, + account_id: u32, + collection: u8, + ) -> crate::Result> { + self.get_bitmap(BitmapKey { + account_id, + collection, + family: BM_DOCUMENT_IDS, + field: u8::MAX, + key: b"", + block_num: 0, + }) + .await + } + + #[inline(always)] + pub async fn get_bitmap>( + &self, + mut key: BitmapKey, + ) -> crate::Result> { + let from_key = key.serialize(); + key.block_num = u32::MAX; + let to_key = key.serialize(); + let opt = RangeOption { + mode: StreamingMode::WantAll, + reverse: false, + ..RangeOption::from((from_key.as_ref(), to_key.as_ref())) + }; + //println!("deserializing bitmap: {:?} {:?}", from_key, to_key); + let mut bm = RoaringBitmap::new(); + let mut values = self.trx.get_ranges(opt, true); + while let Some(values) = values.next().await { + for value in values? { + let key = value.key(); + bm.deserialize_block( + value.value(), + value + .key() + .deserialize_be_u32(key.len() - std::mem::size_of::())?, + ); + } + //println!("deserializing bitmap: {:?} {:?}", value.key(), bm.len()); + } + + Ok(if !bm.is_empty() { Some(bm) } else { None }) + } + + #[inline(always)] + async fn get_bitmaps>( + &self, + keys: Vec>, + ) -> crate::Result>> { + let mut results = Vec::with_capacity(keys.len()); + for key in keys { + results.push(self.get_bitmap(key).await?); + } + + Ok(results) + } + + pub(crate) async fn get_bitmaps_intersection>( + &self, + keys: Vec>, + ) -> crate::Result> { + let mut result: Option = None; + for bitmap in self.get_bitmaps(keys).await? { + if let Some(bitmap) = bitmap { + if let Some(result) = &mut result { + result.bitand_assign(&bitmap); + if result.is_empty() { + break; + } + } else { + result = Some(bitmap); + } + } else { + return Ok(None); + } + } + Ok(result) + } + + pub(crate) async fn get_bitmaps_union>( + &self, + keys: Vec>, + ) -> crate::Result> { + let mut result: Option = None; + for bitmap in (self.get_bitmaps(keys).await?).into_iter().flatten() { + if let Some(result) = &mut result { + result.bitor_assign(&bitmap); + } else { + result = Some(bitmap); + } + } + Ok(result) + } + + pub(crate) async fn range_to_bitmap( + &self, + account_id: u32, + collection: u8, + field: u8, + value: Vec, + op: Operator, + ) -> crate::Result> { + let k1 = KeySerializer::new( + std::mem::size_of::>() + value.len() + 1 + std::mem::size_of::(), + ) + .write(SUBSPACE_INDEXES) + .write(account_id) + .write(collection) + .write(field); + let k2 = KeySerializer::new( + std::mem::size_of::>() + value.len() + 1 + std::mem::size_of::(), + ) + .write(SUBSPACE_INDEXES) + .write(account_id) + .write(collection) + .write(field + matches!(op, Operator::GreaterThan | Operator::GreaterEqualThan) as u8); + + let (begin, end) = match op { + Operator::LowerThan => ( + KeySelector::first_greater_or_equal(k1.finalize()), + KeySelector::last_less_than(k2.write(&value[..]).write(0u32).finalize()), + ), + Operator::LowerEqualThan => ( + KeySelector::first_greater_or_equal(k1.finalize()), + KeySelector::last_less_or_equal(k2.write(&value[..]).write(u32::MAX).finalize()), + ), + Operator::GreaterThan => ( + KeySelector::first_greater_than(k1.write(&value[..]).write(u32::MAX).finalize()), + KeySelector::last_less_than(k2.finalize()), + ), + Operator::GreaterEqualThan => ( + KeySelector::first_greater_or_equal(k1.write(&value[..]).write(0u32).finalize()), + KeySelector::last_less_than(k2.finalize()), + ), + Operator::Equal => ( + KeySelector::first_greater_or_equal(k1.write(&value[..]).write(0u32).finalize()), + KeySelector::last_less_or_equal(k2.write(&value[..]).write(u32::MAX).finalize()), + ), + }; + + let opt = RangeOption { + begin, + end, + mode: StreamingMode::WantAll, + reverse: false, + ..RangeOption::default() + }; + + let mut bm = RoaringBitmap::new(); + let mut range_stream = self.trx.get_ranges(opt, true); + + while let Some(values) = range_stream.next().await { + for value in values? { + let key = value.key(); + bm.insert(key.deserialize_be_u32(key.len() - std::mem::size_of::())?); + } + } + + Ok(Some(bm)) + } + + pub(crate) async fn sort_bitmap( + &self, + account_id: u32, + collection: u8, + field: u8, + documents: &impl UnsortedIds, + limit: usize, + ascending: bool, + ) -> crate::Result> { + let from_key = IndexKeyPrefix { + account_id, + collection, + field, + } + .serialize(); + let to_key = IndexKeyPrefix { + account_id, + collection, + field: field + 1, + } + .serialize(); + let mut results = Vec::with_capacity(documents.len()); + let mut sorted_iter = self.trx.get_ranges( + RangeOption { + begin: KeySelector::first_greater_or_equal(&from_key), + end: KeySelector::last_less_than(&to_key), + mode: options::StreamingMode::Iterator, + reverse: !ascending, + ..Default::default() + }, + true, + ); + + let mut prev_prefix = vec![]; + while let Some(values) = sorted_iter.next().await { + for value in values? { + let key = value.key(); + let document_id = key.deserialize_be_u32(value.key().len() - 4)?; + + if documents.contains_id(document_id) { + let prefix = key + .get(..key.len() - std::mem::size_of::()) + .ok_or_else(|| { + crate::Error::InternalError("Invalid key found in index".to_string()) + })?; + + if prefix == prev_prefix { + let last = results.last_mut().unwrap(); + match last { + SortedId::Id(id) => { + *last = SortedId::GroupedId(vec![*id, document_id]); + } + SortedId::GroupedId(ids) => { + ids.push(document_id); + } + } + } else { + results.push(SortedId::Id(document_id)); + prev_prefix = prefix.to_vec(); + } + + if results.len() == limit { + return Ok(results); + } + } + } + } + + Ok(results) + } + + pub async fn refresh_if_old(&mut self) -> crate::Result<()> { + if self.trx_age.elapsed() > Duration::from_millis(2000) { + self.trx = self.db.create_trx()?; + self.trx_age = Instant::now(); + } + Ok(()) + } +} + +impl Store { + pub async fn read_transaction(&self) -> crate::Result> { + Ok(ReadTransaction { + db: &self.db, + trx: self.db.create_trx()?, + trx_age: Instant::now(), + }) + } +} diff --git a/src/backend/foundationdb/write.rs b/src/backend/foundationdb/write.rs new file mode 100644 index 00000000..b185f711 --- /dev/null +++ b/src/backend/foundationdb/write.rs @@ -0,0 +1,166 @@ +use std::time::Instant; + +use foundationdb::{options::MutationType, FdbError}; + +use crate::{ + write::{Batch, Operation}, + AclKey, BitmapKey, BlobKey, IndexKey, LogKey, Serialize, Store, ValueKey, +}; + +use super::bitmap::DenseBitmap; + +impl Store { + pub async fn write(&self, batch: Batch) -> crate::Result<()> { + let start = Instant::now(); + let mut or_bitmap = DenseBitmap::empty(); + let mut and_bitmap = DenseBitmap::full(); + let mut block_num = u32::MAX; + let mut retry_count = 0; + + loop { + let mut account_id = u32::MAX; + let mut collection = u8::MAX; + let mut document_id = u32::MAX; + let trx = self.db.create_trx()?; + + for op in &batch.ops { + match op { + Operation::AccountId { + account_id: account_id_, + } => { + account_id = *account_id_; + } + Operation::Collection { + collection: collection_, + } => { + collection = *collection_; + } + Operation::DocumentId { + document_id: document_id_, + } => { + if block_num != u32::MAX { + or_bitmap.reset(); + and_bitmap.reset(); + } + document_id = *document_id_; + or_bitmap.set(document_id); + and_bitmap.clear(document_id); + block_num = or_bitmap.block_num; + } + Operation::Value { family, field, set } => { + let key = ValueKey { + account_id, + collection, + document_id, + family: *family, + field: *field, + } + .serialize(); + if let Some(value) = set { + trx.set(&key, value); + } else { + trx.clear(&key); + } + } + Operation::Index { field, key, set } => { + let key = IndexKey { + account_id, + collection, + document_id, + field: *field, + key, + } + .serialize(); + if *set { + trx.set(&key, &[]); + } else { + trx.clear(&key); + } + } + Operation::Bitmap { + family, + field, + key, + set, + } => { + let key = BitmapKey { + account_id, + collection, + family: *family, + field: *field, + block_num, + key, + } + .serialize(); + if *set { + trx.atomic_op(&key, &or_bitmap.bitmap, MutationType::BitOr); + } else { + trx.atomic_op(&key, &and_bitmap.bitmap, MutationType::BitAnd); + }; + } + Operation::Blob { key, set } => { + let key = BlobKey { + account_id, + collection, + document_id, + hash: key, + } + .serialize(); + if *set { + trx.set(&key, &[]); + } else { + trx.clear(&key); + } + } + Operation::Acl { + grant_account_id, + set, + } => { + let key = AclKey { + grant_account_id: *grant_account_id, + to_account_id: account_id, + to_collection: collection, + to_document_id: document_id, + } + .serialize(); + if let Some(value) = set { + trx.set(&key, value); + } else { + trx.clear(&key); + } + } + Operation::Log { + collection, + change_id, + set, + } => { + let key = LogKey { + account_id, + collection: *collection, + change_id: *change_id, + } + .serialize(); + trx.set(&key, set); + } + } + } + + match trx.commit().await { + Ok(_) => { + //println!("Success with id {} block {block_num}", document_id); + return Ok(()); + } + Err(err) => { + if retry_count < 10 && start.elapsed().as_secs() < 5 { + println!("Retrying with id {}", document_id); + err.on_error().await?; + retry_count += 1; + } else { + println!("Error with id {}", document_id); + return Err(FdbError::from(err).into()); + } + } + } + } + } +} diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 54822ba1..9808078e 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -1,2 +1,4 @@ +#[cfg(feature = "foundation")] pub mod foundationdb; +#[cfg(feature = "rocks")] pub mod rocksdb; diff --git a/src/backend/rocksdb/log.rs b/src/backend/rocksdb/log.rs new file mode 100644 index 00000000..4afaeb75 --- /dev/null +++ b/src/backend/rocksdb/log.rs @@ -0,0 +1,123 @@ +use rocksdb::{Direction, IteratorMode}; + +use crate::{ + query::log::{Changes, Query}, + write::key::DeserializeBigEndian, + Error, LogKey, Serialize, Store, +}; + +use super::CF_LOGS; + +const CHANGE_ID_POS: usize = std::mem::size_of::() + std::mem::size_of::(); + +impl Store { + pub fn get_last_change_id( + &self, + account_id: u32, + collection: impl Into, + ) -> crate::Result> { + let collection = collection.into(); + let match_key = LogKey { + account_id, + collection, + change_id: u64::MAX, + } + .serialize(); + + if let Some(Ok((key, _))) = self + .db + .iterator_cf( + &self.db.cf_handle(CF_LOGS).unwrap(), + IteratorMode::From(&match_key, Direction::Reverse), + ) + .next() + { + if key.starts_with(&match_key[0..CHANGE_ID_POS]) { + return Ok(Some( + key.as_ref() + .deserialize_be_u64(CHANGE_ID_POS) + .ok_or_else(|| { + Error::InternalError(format!( + "Failed to deserialize changelog key for [{}/{:?}]: [{:?}]", + account_id, collection, key + )) + })?, + )); + } + } + Ok(None) + } + + pub fn get_changes( + &self, + account_id: u32, + collection: impl Into, + query: Query, + ) -> crate::Result> { + let collection = collection.into(); + let mut changelog = Changes::default(); + let (is_inclusive, from_change_id, to_change_id) = match query { + Query::All => (true, 0, 0), + Query::Since(change_id) => (false, change_id, 0), + Query::SinceInclusive(change_id) => (true, change_id, 0), + Query::RangeInclusive(from_change_id, to_change_id) => { + (true, from_change_id, to_change_id) + } + }; + let key = LogKey { + account_id, + collection, + change_id: from_change_id, + } + .serialize(); + let prefix = &key[0..CHANGE_ID_POS]; + let mut is_first = true; + + for entry in self.db.iterator_cf( + &self.db.cf_handle(CF_LOGS).unwrap(), + IteratorMode::From(&key, Direction::Forward), + ) { + let (key, value) = entry?; + if !key.starts_with(prefix) { + break; + } + let change_id = key + .as_ref() + .deserialize_be_u64(CHANGE_ID_POS) + .ok_or_else(|| { + Error::InternalError(format!( + "Failed to deserialize changelog key for [{}/{:?}]: [{:?}]", + account_id, collection, key + )) + })?; + + if change_id > from_change_id || (is_inclusive && change_id == from_change_id) { + if to_change_id > 0 && change_id > to_change_id { + break; + } + if is_first { + changelog.from_change_id = change_id; + is_first = false; + } + changelog.to_change_id = change_id; + changelog.deserialize(&value).ok_or_else(|| { + Error::InternalError(format!( + "Failed to deserialize changelog for [{}/{:?}]: [{:?}]", + account_id, collection, query + )) + })?; + } + } + + if is_first { + changelog.from_change_id = from_change_id; + changelog.to_change_id = if to_change_id > 0 { + to_change_id + } else { + from_change_id + }; + } + + Ok(Some(changelog)) + } +} diff --git a/src/backend/rocksdb/mod.rs b/src/backend/rocksdb/mod.rs index d156208c..50f3f6ba 100644 --- a/src/backend/rocksdb/mod.rs +++ b/src/backend/rocksdb/mod.rs @@ -3,6 +3,7 @@ use crate::{ }; pub mod bitmap; +pub mod log; pub mod main; pub mod read; pub mod write; @@ -33,12 +34,23 @@ impl> Serialize for IndexKey { impl Serialize for ValueKey { fn serialize(self) -> Vec { - KeySerializer::new(std::mem::size_of::()) - .write_leb128(self.account_id) - .write(self.collection) - .write_leb128(self.document_id) - .write(self.field) - .finalize() + if self.family == 0 { + KeySerializer::new(std::mem::size_of::()) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(self.field) + .finalize() + } else { + KeySerializer::new(std::mem::size_of::() + 1) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(u8::MAX) + .write(self.family) + .write(self.field) + .finalize() + } } } @@ -89,6 +101,18 @@ impl Serialize for LogKey { } } +impl BloomHash { + pub fn to_high_rank_key(&self, account_id: u32, collection: u8, field: u8) -> Vec { + KeySerializer::new(std::mem::size_of::>() + 2) + .write_leb128(account_id) + .write(collection) + .write(BM_BLOOM) + .write(field) + .write(self.as_high_rank_hash()) + .finalize() + } +} + impl From for crate::Error { fn from(value: rocksdb::Error) -> Self { Self::InternalError(format!("RocksDB error: {}", value)) diff --git a/src/backend/rocksdb/write.rs b/src/backend/rocksdb/write.rs index 1501963b..23e2375c 100644 --- a/src/backend/rocksdb/write.rs +++ b/src/backend/rocksdb/write.rs @@ -2,11 +2,12 @@ use std::time::Instant; use roaring::RoaringBitmap; use rocksdb::ErrorKind; +use utils::map::vec_map::VecMap; use crate::{ - write::{key::KeySerializer, Batch, Operation}, + write::{AccountCollection, Batch, Operation, WriteResult}, AclKey, BitmapKey, BlobKey, Deserialize, Error, IndexKey, LogKey, Serialize, Store, ValueKey, - BM_BLOOM, BM_DOCUMENT_IDS, + BM_DOCUMENT_IDS, UNASSIGNED_ID, }; use super::{ @@ -15,7 +16,7 @@ use super::{ }; impl Store { - pub fn write(&self, batch: Batch) -> crate::Result<()> { + pub fn write(&self, batch: Batch) -> crate::Result { let cf_values = self.db.cf_handle(CF_VALUES).unwrap(); let cf_bitmaps = self.db.cf_handle(CF_BITMAPS).unwrap(); let cf_indexes = self.db.cf_handle(CF_INDEXES).unwrap(); @@ -27,6 +28,10 @@ impl Store { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; + let mut result = WriteResult { + change_ids: VecMap::new(), + assigned_ids: VecMap::new(), + }; let txn = self.db.transaction(); let mut wb = txn.get_writebatch(); @@ -46,7 +51,7 @@ impl Store { document_id: document_id_, set, } => { - if *document_id_ == u32::MAX { + if *document_id_ == UNASSIGNED_ID { let key = BitmapKey { account_id, collection, @@ -80,6 +85,9 @@ impl Store { } else { 0 }; + result + .assigned_ids + .append((account_id, collection).into(), document_id); wb.merge_cf(&cf_bitmaps, key, set_bit(document_id)); } else { document_id = *document_id_; @@ -99,11 +107,12 @@ impl Store { } } } - Operation::Value { field, set } => { + Operation::Value { family, field, set } => { let key = ValueKey { account_id, collection, document_id, + family: *family, field: *field, } .serialize(); @@ -149,21 +158,6 @@ impl Store { }; wb.merge_cf(&cf_bitmaps, key, value); } - Operation::Bloom { family, field, set } => { - let key = KeySerializer::new(std::mem::size_of::()) - .write_leb128(account_id) - .write(collection) - .write_leb128(document_id) - .write(u8::MAX) - .write(BM_BLOOM | *family) - .write(*field) - .finalize(); - if let Some(value) = set { - wb.put_cf(&cf_values, key, value); - } else { - wb.delete_cf(&cf_values, key); - } - } Operation::Blob { key, set } => { let key = BlobKey { account_id, @@ -195,15 +189,30 @@ impl Store { wb.delete_cf(&cf_values, key); } } - Operation::Log { change_id, changes } => { - let coco = "_"; + Operation::Log { + collection, + changes, + } => { + let ac: AccountCollection = (account_id, *collection).into(); + let coco = "read for write"; + let change_id = self + .get_last_change_id(account_id, *collection)? + .map(|id| id + 1) + .unwrap_or(0); let key = LogKey { account_id, - collection, - change_id: *change_id, + collection: *collection, + change_id, } .serialize(); - wb.put_cf(&cf_logs, key, changes); + wb.put_cf( + &cf_logs, + key, + changes.serialize( + result.assigned_ids.get(&ac).copied().unwrap_or_default(), + ), + ); + result.change_ids.append(ac, change_id); } } } @@ -211,7 +220,7 @@ impl Store { match self.db.write(wb) { Ok(_) => { //println!("Success with id {}", document_id); - return Ok(()); + return Ok(result); } Err(err) => match err.kind() { ErrorKind::Busy | ErrorKind::MergeInProgress | ErrorKind::TryAgain diff --git a/src/fts/bloom.rs b/src/fts/bloom.rs index bb613bc3..60dfb397 100644 --- a/src/fts/bloom.rs +++ b/src/fts/bloom.rs @@ -7,7 +7,7 @@ use std::{ use roaring::RoaringBitmap; use utils::codec::leb128::{Leb128Reader, Leb128Vec}; -use crate::{Deserialize, Serialize}; +use crate::{Deserialize, Error, Serialize}; use super::{stemmer::StemmedToken, tokenizers::Token}; @@ -94,6 +94,10 @@ impl BloomFilter { } } +pub trait BloomHasher { + fn hash + ?Sized>(item: &T) -> Self; +} + impl BloomHash { pub fn hash + ?Sized>(item: &T) -> Self { let h1 = xxhash_rust::xxh3::xxh3_64(item.as_ref()); @@ -175,10 +179,18 @@ impl Serialize for BloomFilter { } impl Deserialize for BloomFilter { - fn deserialize(bytes: &[u8]) -> Option { - let (m, pos) = bytes.read_leb128()?; - let b = RoaringBitmap::deserialize_unchecked_from(bytes.get(pos..)?).ok()?; - - Some(Self::from_params(m, b)) + fn deserialize(bytes: &[u8]) -> crate::Result { + let (m, pos) = bytes.read_leb128().ok_or_else(|| { + Error::InternalError( + "Failed to read 'm' value while deserializing bloom filter.".to_string(), + ) + })?; + RoaringBitmap::deserialize_unchecked_from(bytes.get(pos..).ok_or_else(|| { + Error::InternalError( + "Failed to read bitmap while deserializing bloom filter.".to_string(), + ) + })?) + .map_err(|err| Error::InternalError(format!("Failed to deserialize bloom filter: {err}."))) + .map(|b| Self::from_params(m, b)) } } diff --git a/src/fts/builder.rs b/src/fts/builder.rs index 64aabd4a..f6487f53 100644 --- a/src/fts/builder.rs +++ b/src/fts/builder.rs @@ -4,7 +4,7 @@ use ahash::AHashSet; use crate::{ write::{BatchBuilder, IntoOperations, Operation}, - Serialize, BLOOM_BIGRAM, BLOOM_STEMMED, BLOOM_TRIGRAM, BM_BLOOM, + Serialize, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BM_BLOOM, }; use super::{ @@ -80,36 +80,34 @@ impl<'x> IntoOperations for FtsIndexBuilder<'x> { phrase_words.push(token.word); } - let mut bloom_stemmed = BloomFilter::new(unique_words.len()); + let mut bloom_unigram = BloomFilter::new(unique_words.len()); for word in unique_words { let hash = BloomHash::from(word); - bloom_stemmed.insert(&hash); - //for h in [0, 1] { + bloom_unigram.insert(&hash); batch.ops.push(Operation::Bitmap { family: BM_BLOOM, field: part.field, - key: hash.as_high_rank_hash(0).serialize(), + key: hash.as_high_rank_hash().serialize(), set: true, }); - //} } - batch.ops.push(Operation::Bloom { + batch.ops.push(Operation::Value { field: part.field, - family: BLOOM_STEMMED, - set: bloom_stemmed.serialize().into(), + family: BM_BLOOM | BLOOM_UNIGRAM, + set: bloom_unigram.serialize().into(), }); if phrase_words.len() > 1 { - batch.ops.push(Operation::Bloom { + batch.ops.push(Operation::Value { field: part.field, - family: BLOOM_BIGRAM, + family: BM_BLOOM | BLOOM_BIGRAM, set: BloomFilter::to_ngrams(&phrase_words, 2).serialize().into(), }); if phrase_words.len() > 2 { - batch.ops.push(Operation::Bloom { + batch.ops.push(Operation::Value { field: part.field, - family: BLOOM_TRIGRAM, + family: BM_BLOOM | BLOOM_TRIGRAM, set: BloomFilter::to_ngrams(&phrase_words, 3).serialize().into(), }); } diff --git a/src/fts/mod.rs b/src/fts/mod.rs index 6c14e863..4bc63421 100644 --- a/src/fts/mod.rs +++ b/src/fts/mod.rs @@ -27,9 +27,9 @@ pub mod bloom; pub mod builder; pub mod ngram; pub mod query; -pub mod search_snippet; +//pub mod search_snippet; pub mod stemmer; -pub mod term_index; +//pub mod term_index; pub mod tokenizers; pub const HIGH_RANK_MOD: u64 = 10_240; diff --git a/src/fts/query.rs b/src/fts/query.rs index bfc51924..164a0082 100644 --- a/src/fts/query.rs +++ b/src/fts/query.rs @@ -10,14 +10,13 @@ use crate::{ stemmer::Stemmer, tokenizers::Tokenizer, }, - write::key::KeySerializer, - BitmapKey, Store, ValueKey, BLOOM_BIGRAM, BLOOM_STEMMED, BLOOM_TRIGRAM, BM_BLOOM, + BitmapKey, Serialize, Store, ValueKey, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BM_BLOOM, }; use super::{Language, HIGH_RANK_MOD}; impl Store { - pub(crate) fn fts_query( + pub(crate) async fn fts_query( &self, account_id: u32, collection: u8, @@ -27,40 +26,41 @@ impl Store { match_phrase: bool, ) -> crate::Result> { let real_now = Instant::now(); + let mut trx = self.read_transaction().await?; let (bitmaps, hashes, family) = if match_phrase { let mut tokens = Vec::new(); let mut bit_keys = Vec::new(); for token in Tokenizer::new(text, language, MAX_TOKEN_LENGTH) { let hash = BloomHash::from(token.word.as_ref()); - let key = hash.to_high_rank_key(account_id, collection, field, 0); + let key = hash.to_high_rank_key(account_id, collection, field); if !bit_keys.contains(&key) { bit_keys.push(key); } tokens.push(token.word); } - let bitmaps = match self.get_bitmaps_intersection(bit_keys)? { + let bitmaps = match trx.get_bitmaps_intersection(bit_keys).await? { Some(b) if !b.is_empty() => b, _ => return Ok(None), }; match tokens.len() { - 0 => (bitmaps, vec![], BLOOM_STEMMED), + 0 => return Ok(None), 1 => ( bitmaps, vec![tokens.into_iter().next().unwrap().into()], - BLOOM_STEMMED, + BM_BLOOM | BLOOM_UNIGRAM, ), 2 => ( bitmaps, >::to_ngrams(&tokens, 2), - BLOOM_BIGRAM, + BM_BLOOM | BLOOM_BIGRAM, ), _ => ( bitmaps, >::to_ngrams(&tokens, 3), - BLOOM_TRIGRAM, + BM_BLOOM | BLOOM_TRIGRAM, ), } } else { @@ -76,14 +76,18 @@ impl Store { }, h1: token.word.into(), }; + trx.refresh_if_old().await?; - match self.get_bitmaps_union(vec![ - hash.h1.to_high_rank_key(account_id, collection, field, 0), - hash.h2 - .as_ref() - .unwrap() - .to_high_rank_key(account_id, collection, field, 0), - ])? { + match trx + .get_bitmaps_union(vec![ + hash.h1.to_high_rank_key(account_id, collection, field), + hash.h2 + .as_ref() + .unwrap() + .to_high_rank_key(account_id, collection, field), + ]) + .await? + { Some(b) if !b.is_empty() => { if !bitmaps.is_empty() { bitmaps &= b; @@ -100,59 +104,63 @@ impl Store { hashes.push(hash); } - (bitmaps, hashes, BLOOM_STEMMED) + (bitmaps, hashes, BM_BLOOM | BLOOM_UNIGRAM) }; let b_count = bitmaps.len(); - let mut bm = RoaringBitmap::new(); - /*let keys = bitmaps - .iter() - .map(|document_id| { - KeySerializer::new(std::mem::size_of::()) - .write_leb128(account_id) - .write(collection) - .write_leb128(document_id) - .write(u8::MAX) - .write(BM_BLOOM | family) - .write(field) - .finalize() - }) - .collect::>(); - - self.get_values::(keys)? - .into_iter() - .zip(bitmaps) - .for_each(|(bloom, document_id)| { - if let Some(bloom) = bloom { - if !bloom.is_empty() { - let mut matched = true; - for hash in &hashes { - if !(bloom.contains(&hash.h1) - || hash.h2.as_ref().map_or(false, |h2| bloom.contains(h2))) - { - matched = false; - break; - } - } - - if matched { - bm.insert(document_id); - } + /*let bm = self + .get_values::( + bitmaps + .iter() + .map(|document_id| ValueKey { + account_id, + collection, + document_id, + family, + field, + }) + .collect::>(), + ) + .await? + .into_iter() + .zip(bitmaps) + .filter_map(|(bloom, document_id)| { + let bloom = bloom?; + if !bloom.is_empty() { + let mut matched = true; + for hash in &hashes { + if !(bloom.contains(&hash.h1) + || hash.h2.as_ref().map_or(false, |h2| bloom.contains(h2))) + { + matched = false; + break; } } - });*/ - for document_id in bitmaps { - let key = KeySerializer::new(std::mem::size_of::() + 2) - .write_leb128(account_id) - .write(collection) - .write_leb128(document_id) - .write(u8::MAX) - .write(BM_BLOOM | family) - .write(field) - .finalize(); - if let Some(bloom) = self.get_value::(key)? { + if matched { + return Some(document_id); + } + } + + None + }) + .collect::();*/ + + let mut bm = RoaringBitmap::new(); + for document_id in bitmaps { + trx.refresh_if_old().await?; + + if let Some(bloom) = trx + .get_value::(ValueKey { + account_id, + collection, + document_id, + family, + field, + }) + .await? + { if !bloom.is_empty() { let mut matched = true; for hash in &hashes { @@ -172,7 +180,7 @@ impl Store { } println!( - "bloom_match {b_count} items in {:?}ms", + "bloom_match {text:?} {b_count} items in {:?}ms", real_now.elapsed().as_millis() ); @@ -182,8 +190,8 @@ impl Store { impl BloomHash { #[inline(always)] - pub fn as_high_rank_hash(&self, n: usize) -> u16 { - (self.h[n] % HIGH_RANK_MOD) as u16 + pub fn as_high_rank_hash(&self) -> u16 { + (self.h[0] % HIGH_RANK_MOD) as u16 } pub fn to_high_rank_key( @@ -191,14 +199,14 @@ impl BloomHash { account_id: u32, collection: u8, field: u8, - n: usize, - ) -> Vec { - KeySerializer::new(std::mem::size_of::>() + 2) - .write_leb128(account_id) - .write(collection) - .write(BM_BLOOM) - .write(field) - .write(self.as_high_rank_hash(n)) - .finalize() + ) -> BitmapKey> { + BitmapKey { + account_id, + collection, + family: BM_BLOOM, + field, + block_num: 0, + key: self.as_high_rank_hash().serialize(), + } } } diff --git a/src/lib.rs b/src/lib.rs index c3f52dc8..a11c6f1d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,3 @@ -use rocksdb::{MultiThreaded, OptimisticTransactionDB}; - pub mod backend; pub mod fts; pub mod query; @@ -8,12 +6,20 @@ pub mod write; #[cfg(test)] pub mod tests; +#[cfg(feature = "rocks")] pub struct Store { - db: OptimisticTransactionDB, + db: rocksdb::OptimisticTransactionDB, +} + +#[cfg(feature = "foundation")] +#[allow(dead_code)] +pub struct Store { + db: foundationdb::Database, + guard: foundationdb::api::NetworkAutoStop, } pub trait Deserialize: Sized + Sync + Send { - fn deserialize(bytes: &[u8]) -> Option; + fn deserialize(bytes: &[u8]) -> crate::Result; } pub trait Serialize { @@ -26,6 +32,8 @@ pub struct BitmapKey> { pub collection: u8, pub family: u8, pub field: u8, + #[cfg(feature = "foundation")] + pub block_num: u32, pub key: T, } @@ -38,11 +46,19 @@ pub struct IndexKey> { pub key: T, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct IndexKeyPrefix { + pub account_id: u32, + pub collection: u8, + pub field: u8, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct ValueKey { pub account_id: u32, pub collection: u8, pub document_id: u32, + pub family: u8, pub field: u8, } @@ -82,7 +98,7 @@ pub const BM_TERM: u8 = 0x10; pub const BM_TAG: u8 = 0x20; pub const BM_BLOOM: u8 = 0x40; -pub const BLOOM_STEMMED: u8 = 0x00; +pub const BLOOM_UNIGRAM: u8 = 0x00; pub const BLOOM_BIGRAM: u8 = 0x01; pub const BLOOM_TRIGRAM: u8 = 0x02; diff --git a/src/query/filter.rs b/src/query/filter.rs index 9702f34a..58d57cc1 100644 --- a/src/query/filter.rs +++ b/src/query/filter.rs @@ -1,10 +1,13 @@ -use std::ops::{BitAndAssign, BitOrAssign, BitXorAssign}; +use std::{ + ops::{BitAndAssign, BitOrAssign, BitXorAssign}, + time::Instant, +}; use roaring::RoaringBitmap; use crate::{ - write::{key::KeySerializer, Tokenize}, - BitmapKey, IndexKey, Store, BM_TERM, TERM_EXACT, + backend::foundationdb::read::ReadTransaction, write::Tokenize, BitmapKey, Store, BM_TERM, + TERM_EXACT, }; use super::{Filter, ResultSet}; @@ -15,17 +18,21 @@ struct State { } impl Store { - pub fn filter( + pub async fn filter( &self, account_id: u32, collection: u8, filters: Vec, ) -> crate::Result { - let document_ids = self - .get_document_ids(account_id, collection)? + let mut trx = self.read_transaction().await?; + let document_ids = trx + .get_document_ids(account_id, collection) + .await? .unwrap_or_else(RoaringBitmap::new); if filters.is_empty() { return Ok(ResultSet { + account_id, + collection, results: document_ids.clone(), document_ids, }); @@ -36,52 +43,42 @@ impl Store { let mut filters = filters.into_iter().peekable(); while let Some(filter) = filters.next() { - match filter { + trx.refresh_if_old().await?; + + let result = match filter { Filter::HasKeyword { field, value } => { - state.op.apply( - &mut state.bm, - self.get_bitmap(BitmapKey { - account_id, - collection, - family: BM_TERM | TERM_EXACT, - field, - key: value.as_bytes(), - })?, - &document_ids, - ); + trx.get_bitmap(BitmapKey { + account_id, + collection, + family: BM_TERM | TERM_EXACT, + field, + key: value.as_bytes(), + #[cfg(feature = "foundation")] + block_num: 0, + }) + .await? } Filter::HasKeywords { field, value } => { - let tokens = value.tokenize(); - state.op.apply( - &mut state.bm, - self.get_bitmaps_intersection( - tokens - .iter() - .map(|key| BitmapKey { - account_id, - collection, - family: BM_TERM | TERM_EXACT, - field, - key: key.as_bytes(), - }) - .collect(), - )?, - &document_ids, - ); + trx.get_bitmaps_intersection( + value + .tokenize() + .into_iter() + .map(|key| BitmapKey { + account_id, + collection, + family: BM_TERM | TERM_EXACT, + field, + key: key.into_bytes(), + #[cfg(feature = "foundation")] + block_num: 0, + }) + .collect(), + ) + .await? } Filter::MatchValue { field, op, value } => { - let key = - KeySerializer::new(std::mem::size_of::>() + value.len()) - .write(account_id) - .write(collection) - .write(field) - .write(&value[..]) - .finalize(); - state.op.apply( - &mut state.bm, - self.range_to_bitmap(&key, &value, op)?, - &document_ids, - ); + trx.range_to_bitmap(account_id, collection, field, value, op) + .await? } Filter::HasText { field, @@ -89,51 +86,39 @@ impl Store { language, match_phrase, } => { - state.op.apply( - &mut state.bm, - self.fts_query( - account_id, - collection, - field, - &text, - language, - match_phrase, - )?, - &document_ids, - ); + self.fts_query(account_id, collection, field, &text, language, match_phrase) + .await? } Filter::InBitmap { family, field, key } => { - state.op.apply( - &mut state.bm, - self.get_bitmap(BitmapKey { - account_id, - collection, - family, - field, - key: &key, - })?, - &document_ids, - ); - } - Filter::DocumentSet(set) => { - state.op.apply(&mut state.bm, Some(set), &document_ids); + trx.get_bitmap(BitmapKey { + account_id, + collection, + family, + field, + key: &key, + #[cfg(feature = "foundation")] + block_num: 0, + }) + .await? } + Filter::DocumentSet(set) => Some(set), op @ (Filter::And | Filter::Or | Filter::Not) => { stack.push(state); state = op.into(); continue; } Filter::End => { - if let Some(mut prev_state) = stack.pop() { - prev_state - .op - .apply(&mut prev_state.bm, state.bm, &document_ids); + if let Some(prev_state) = stack.pop() { + let bm = state.bm; state = prev_state; + bm } else { break; } } - } + }; + + state.op.apply(&mut state.bm, result, &document_ids); //println!("{:?}: {:?}", state.op, state.bm); @@ -149,6 +134,8 @@ impl Store { } Ok(ResultSet { + account_id, + collection, results: state.bm.unwrap_or_else(RoaringBitmap::new), document_ids, }) diff --git a/src/query/log.rs b/src/query/log.rs new file mode 100644 index 00000000..53dafcdb --- /dev/null +++ b/src/query/log.rs @@ -0,0 +1,109 @@ +use utils::codec::leb128::Leb128Iterator; + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum Change { + Insert(u64), + Update(u64), + ChildUpdate(u64), + Delete(u64), +} + +pub struct Changes { + pub changes: Vec, + pub from_change_id: u64, + pub to_change_id: u64, +} + +#[derive(Debug)] +pub enum Query { + All, + Since(u64), + SinceInclusive(u64), + RangeInclusive(u64, u64), +} + +impl Default for Changes { + fn default() -> Self { + Self { + changes: Vec::with_capacity(10), + from_change_id: 0, + to_change_id: 0, + } + } +} + +impl Changes { + pub fn deserialize(&mut self, bytes: &[u8]) -> Option<()> { + let mut bytes_it = bytes.iter(); + let total_inserts: usize = bytes_it.next_leb128()?; + let total_updates: usize = bytes_it.next_leb128()?; + let total_child_updates: usize = bytes_it.next_leb128()?; + let total_deletes: usize = bytes_it.next_leb128()?; + + if total_inserts > 0 { + for _ in 0..total_inserts { + self.changes.push(Change::Insert(bytes_it.next_leb128()?)); + } + } + + if total_updates > 0 || total_child_updates > 0 { + 'update_outer: for change_pos in 0..(total_updates + total_child_updates) { + let id = bytes_it.next_leb128()?; + let mut is_child_update = change_pos >= total_updates; + + for (idx, change) in self.changes.iter().enumerate() { + match change { + Change::Insert(insert_id) if *insert_id == id => { + // Item updated after inserted, no need to count this change. + continue 'update_outer; + } + Change::Update(update_id) if *update_id == id => { + // Move update to the front + is_child_update = false; + self.changes.remove(idx); + break; + } + Change::ChildUpdate(update_id) if *update_id == id => { + // Move update to the front + self.changes.remove(idx); + break; + } + _ => (), + } + } + + self.changes.push(if !is_child_update { + Change::Update(id) + } else { + Change::ChildUpdate(id) + }); + } + } + + if total_deletes > 0 { + 'delete_outer: for _ in 0..total_deletes { + let id = bytes_it.next_leb128()?; + + 'delete_inner: for (idx, change) in self.changes.iter().enumerate() { + match change { + Change::Insert(insert_id) if *insert_id == id => { + self.changes.remove(idx); + continue 'delete_outer; + } + Change::Update(update_id) | Change::ChildUpdate(update_id) + if *update_id == id => + { + self.changes.remove(idx); + break 'delete_inner; + } + _ => (), + } + } + + self.changes.push(Change::Delete(id)); + } + } + + Some(()) + } +} diff --git a/src/query/mod.rs b/src/query/mod.rs index 50156244..1117c40b 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -1,4 +1,5 @@ pub mod filter; +pub mod log; pub mod sort; use roaring::RoaringBitmap; @@ -58,8 +59,10 @@ pub enum Comparator { #[derive(Debug)] pub struct ResultSet { - results: RoaringBitmap, - document_ids: RoaringBitmap, + account_id: u32, + collection: u8, + pub results: RoaringBitmap, + pub document_ids: RoaringBitmap, } pub struct SortedResultRet { @@ -68,6 +71,17 @@ pub struct SortedResultRet { pub found_anchor: bool, } +pub enum SortedId { + Id(u32), + GroupedId(Vec), +} + +#[allow(clippy::len_without_is_empty)] +pub trait UnsortedIds { + fn contains_id(&self, id: u32) -> bool; + fn len(&self) -> usize; +} + impl Filter { pub fn cond(field: impl Into, op: Operator, value: impl Serialize) -> Self { Filter::MatchValue { @@ -183,3 +197,23 @@ impl Comparator { } } } + +impl UnsortedIds for RoaringBitmap { + fn contains_id(&self, id: u32) -> bool { + self.contains(id) + } + + fn len(&self) -> usize { + self.len() as usize + } +} + +impl UnsortedIds for Vec { + fn contains_id(&self, id: u32) -> bool { + self.iter().any(|&i| i == id) + } + + fn len(&self) -> usize { + self.len() + } +} diff --git a/src/query/sort.rs b/src/query/sort.rs index 823c11b2..3eeb3bc8 100644 --- a/src/query/sort.rs +++ b/src/query/sort.rs @@ -1,18 +1,24 @@ use std::ops::{BitAndAssign, BitXorAssign}; +use foundationdb::{future::FdbValue, options, FdbResult, KeySelector, RangeOption}; +use futures::{Stream, StreamExt}; use roaring::RoaringBitmap; +#[cfg(feature = "rocks")] use rocksdb::{ DBIteratorWithThreadMode, Direction, IteratorMode, MultiThreaded, OptimisticTransactionDB, }; +#[cfg(feature = "rocks")] +use crate::backend::rocksdb::{ACCOUNT_KEY_LEN, CF_INDEXES}; + use crate::{ - backend::rocksdb::{ACCOUNT_KEY_LEN, CF_INDEXES}, - write::key::KeySerializer, - Error, Store, + backend::foundationdb::read::ReadTransaction, write::key::DeserializeBigEndian, Error, + IndexKeyPrefix, Serialize, Store, }; use super::{Comparator, ResultSet, SortedResultRet}; +#[cfg(feature = "rocks")] enum IndexType<'x> { DocumentSet { set: RoaringBitmap, @@ -21,25 +27,47 @@ enum IndexType<'x> { DB { it: Option>>, prefix: Vec, - start_key: Vec, + from_key: Vec, ascending: bool, prev_item: Option, prev_key: Option>, }, } +#[cfg(feature = "rocks")] struct IndexIterator<'x> { index: IndexType<'x>, remaining: RoaringBitmap, eof: bool, } +#[cfg(feature = "foundation")] +enum IndexType<'x, T: Stream> + Unpin + 'x> { + DocumentSet { + set: RoaringBitmap, + it: Option, + }, + DB { + it: Option, + from_key: Vec, + to_key: Vec, + ascending: bool, + prev_item: Option, + prev_key: Option>, + phantom: std::marker::PhantomData<&'x ()>, + }, +} + +#[cfg(feature = "foundation")] +struct IndexIterator<'x, T: Stream> + Unpin + 'x> { + index: IndexType<'x, T>, + remaining: RoaringBitmap, + eof: bool, +} + impl Store { - #[allow(clippy::too_many_arguments)] - pub fn sort( + pub async fn sort( &self, - account_id: u32, - collection: u8, mut result_set: ResultSet, comparators: Vec, limit: usize, @@ -50,6 +78,7 @@ impl Store { let has_anchor = anchor.is_some(); let mut anchor_found = false; let requested_position = position; + let trx = self.read_transaction().await?; let mut result = SortedResultRet { position, @@ -60,37 +89,23 @@ impl Store { .into_iter() .map(|comp| IndexIterator { index: match comp { - Comparator::Field { field, ascending } => { - let prefix = KeySerializer::new(ACCOUNT_KEY_LEN) - .write(account_id) - .write(collection) - .write(field) - .finalize(); - IndexType::DB { - it: None, - start_key: if !ascending { - let (key_account_id, key_collection, key_field) = if field < u8::MAX - { - (account_id, collection, field + 1) - } else if (collection) < u8::MAX { - (account_id, (collection) + 1, field) - } else { - (account_id + 1, collection, field) - }; - KeySerializer::new(ACCOUNT_KEY_LEN) - .write(key_account_id) - .write(key_collection) - .write(key_field) - .finalize() - } else { - prefix.clone() - }, - prefix, - ascending, - prev_item: None, - prev_key: None, - } - } + Comparator::Field { field, ascending } => IndexType::DB { + it: None, + from_key: if !ascending { + result_set.from_key(field).serialize() + } else { + result_set.to_key(field).serialize() + }, + to_key: if !ascending { + result_set.to_key(field).serialize() + } else { + result_set.from_key(field).serialize() + }, + ascending, + prev_item: None, + prev_key: None, + phantom: std::marker::PhantomData, + }, Comparator::DocumentSet { mut set, ascending } => IndexType::DocumentSet { set: if !ascending { if !set.is_empty() { @@ -111,12 +126,13 @@ impl Store { .collect::>(); let mut current = 0; + let iter_len = iterators.len() - 1; 'outer: loop { let mut doc_id; 'inner: loop { - let (it_opts, mut next_it_opts) = if current < iterators.len() - 1 { + let (it_opts, mut next_it_opts) = if current < iter_len { let (iterators_first, iterators_last) = iterators.split_at_mut(current + 1); ( iterators_first.last_mut().unwrap(), @@ -145,26 +161,46 @@ impl Store { match &mut it_opts.index { IndexType::DB { it, - prefix, - start_key, + from_key, + to_key, ascending, prev_item, prev_key, + .. } => { let it = if let Some(it) = it { it } else { - *it = Some(self.db.iterator_cf( - &self.db.cf_handle(CF_INDEXES).unwrap(), - IteratorMode::From( - start_key, - if *ascending { - Direction::Forward - } else { - Direction::Reverse + #[cfg(feature = "foundation")] + { + *it = Some(trx.trx.get_ranges_keyvalues( + RangeOption { + begin: KeySelector::first_greater_or_equal( + from_key.clone(), + ), + end: KeySelector::last_less_than(to_key.clone()), + mode: options::StreamingMode::Iterator, + reverse: !*ascending, + ..Default::default() }, - ), - )); + true, + )); + } + + #[cfg(feature = "rocks")] + { + *it = Some(self.db.iterator_cf( + &self.db.cf_handle(CF_INDEXES).unwrap(), + IteratorMode::From( + from_key, + if *ascending { + Direction::Forward + } else { + Direction::Reverse + }, + ), + )); + } it.as_mut().unwrap() }; @@ -184,24 +220,20 @@ impl Store { let mut is_eof = false; loop { - if let Some(result) = it.next() { - let (key, _) = result.map_err(|e| { + if let Some(result) = it.next().await { + let key = result?.key().to_vec().into_boxed_slice(); + /*let (key, _) = result.map_err(|e| { Error::InternalError(format!("Iterator error: {}", e)) })?; if !key.starts_with(prefix) { *prev_key = None; is_eof = true; break; - } + }*/ - doc_id = u32::from_be_bytes( - key.get(key.len() - std::mem::size_of::()..) - .ok_or_else(|| { - Error::InternalError("Invalid index entry".to_string()) - })? - .try_into() - .unwrap(), - ); + doc_id = key + .as_ref() + .deserialize_be_u32(key.len() - std::mem::size_of::())?; if it_opts.remaining.contains(doc_id) { it_opts.remaining.remove(doc_id); @@ -305,24 +337,45 @@ impl Store { match &mut next_it_opts.index { IndexType::DB { it, - start_key, + from_key, + to_key, ascending, prev_item, prev_key, .. } => { if let Some(it) = it { - *it = self.db.iterator_cf( - &self.db.cf_handle(CF_INDEXES).unwrap(), - IteratorMode::From( - start_key, - if *ascending { - Direction::Forward - } else { - Direction::Reverse + #[cfg(feature = "rocks")] + { + *it = self.db.iterator_cf( + &self.db.cf_handle(CF_INDEXES).unwrap(), + IteratorMode::From( + from_key, + if *ascending { + Direction::Forward + } else { + Direction::Reverse + }, + ), + ); + } + #[cfg(feature = "foundation")] + { + *it = trx.trx.get_ranges_keyvalues( + RangeOption { + begin: KeySelector::first_greater_or_equal( + from_key.clone(), + ), + end: KeySelector::last_less_than( + to_key.clone(), + ), + mode: options::StreamingMode::Iterator, + reverse: !*ascending, + ..Default::default() }, - ), - ); + true, + ); + } } *prev_item = None; *prev_key = None; @@ -422,3 +475,27 @@ impl Store { Ok(result) } } + +impl ResultSet { + pub fn from_key(&self, field: u8) -> IndexKeyPrefix { + IndexKeyPrefix { + account_id: self.account_id, + collection: self.collection, + field, + } + } + pub fn to_key(&self, field: u8) -> IndexKeyPrefix { + let (account_id, collection, field) = if field < u8::MAX { + (self.account_id, self.collection, field + 1) + } else if (self.collection) < u8::MAX { + (self.account_id, (self.collection) + 1, field) + } else { + (self.account_id + 1, self.collection, field) + }; + IndexKeyPrefix { + account_id, + collection, + field, + } + } +} diff --git a/src/rocksdb.rs b/src/rocksdb.rs deleted file mode 100644 index 9b7f3b42..00000000 --- a/src/rocksdb.rs +++ /dev/null @@ -1,343 +0,0 @@ -/* - * Copyright (c) 2020-2022, Stalwart Labs Ltd. - * - * This file is part of the Stalwart JMAP Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use std::{convert::TryInto, path::PathBuf, sync::Arc}; - -use rocksdb::{ - BoundColumnFamily, ColumnFamilyDescriptor, DBIteratorWithThreadMode, MergeOperands, - MultiThreaded, OptimisticTransactionDB, Options, -}; - -use crate::{Deserialize, Error, InnerStore}; - -pub struct RocksDB { - db: OptimisticTransactionDB, -} - -pub struct RocksDBIterator<'x> { - it: DBIteratorWithThreadMode<'x, OptimisticTransactionDB>, -} - -impl Iterator for RocksDBIterator<'_> { - type Item = (Box<[u8]>, Box<[u8]>); - - #[allow(clippy::while_let_on_iterator)] - #[inline(always)] - fn next(&mut self) -> Option { - while let Some(result) = self.it.next() { - if let Ok(item) = result { - return Some(item); - } - } - None - } -} - -impl InnerStore for RocksDB { - type Iterator<'x> = RocksDBIterator<'x>; - - #[inline(always)] - fn delete(&self, cf: crate::ColumnFamily, key: &[u8]) -> crate::Result<()> { - self.db - .delete_cf(&self.cf_handle(cf)?, key) - .map_err(|err| Error::InternalError(format!("delete_cf failed: {}", err))) - } - - #[inline(always)] - fn set(&self, cf: crate::ColumnFamily, key: &[u8], value: &[u8]) -> crate::Result<()> { - self.db - .put_cf(&self.cf_handle(cf)?, key, value) - .map_err(|err| Error::InternalError(format!("put_cf failed: {}", err))) - } - - #[inline(always)] - fn get(&self, cf: crate::ColumnFamily, key: &[u8]) -> crate::Result> - where - U: Deserialize, - { - if let Some(bytes) = self - .db - .get_pinned_cf(&self.cf_handle(cf)?, key) - .map_err(|err| Error::InternalError(format!("get_cf failed: {}", err)))? - { - Ok(Some(U::deserialize(&bytes).ok_or_else(|| { - Error::DeserializeError(format!("Failed to deserialize key: {:?}", key)) - })?)) - } else { - Ok(None) - } - } - - #[inline(always)] - fn merge(&self, cf: crate::ColumnFamily, key: &[u8], value: &[u8]) -> crate::Result<()> { - self.db - .merge_cf(&self.cf_handle(cf)?, key, value) - .map_err(|err| Error::InternalError(format!("merge_cf failed: {}", err))) - } - - /* - #[inline(always)] - fn write(&self, batch: Vec) -> crate::Result<()> { - let mut rocks_batch = rocksdb::WriteBatch::default(); - let cf_bitmaps = self.cf_handle(crate::ColumnFamily::Bitmaps)?; - let cf_values = self.cf_handle(crate::ColumnFamily::Values)?; - let cf_indexes = self.cf_handle(crate::ColumnFamily::Indexes)?; - let cf_blobs = self.cf_handle(crate::ColumnFamily::Blobs)?; - let cf_logs = self.cf_handle(crate::ColumnFamily::Logs)?; - - for op in batch { - match op { - WriteOperation::Set { cf, key, value } => { - rocks_batch.put_cf( - match cf { - crate::ColumnFamily::Bitmaps => &cf_bitmaps, - crate::ColumnFamily::Values => &cf_values, - crate::ColumnFamily::Indexes => &cf_indexes, - crate::ColumnFamily::Blobs => &cf_blobs, - crate::ColumnFamily::Logs => &cf_logs, - }, - key, - value, - ); - } - WriteOperation::Delete { cf, key } => { - rocks_batch.delete_cf( - match cf { - crate::ColumnFamily::Bitmaps => &cf_bitmaps, - crate::ColumnFamily::Values => &cf_values, - crate::ColumnFamily::Indexes => &cf_indexes, - crate::ColumnFamily::Blobs => &cf_blobs, - crate::ColumnFamily::Logs => &cf_logs, - }, - key, - ); - } - WriteOperation::Merge { cf, key, value } => { - rocks_batch.merge_cf( - match cf { - crate::ColumnFamily::Bitmaps => &cf_bitmaps, - crate::ColumnFamily::Values => &cf_values, - crate::ColumnFamily::Indexes => &cf_indexes, - crate::ColumnFamily::Blobs => &cf_blobs, - crate::ColumnFamily::Logs => &cf_logs, - }, - key, - value, - ); - } - } - } - self.db - .write(rocks_batch) - .map_err(|err| Error::InternalError(format!("batch write failed: {}", err))) - } - - */ - - #[inline(always)] - fn exists(&self, cf: crate::ColumnFamily, key: &[u8]) -> crate::Result { - Ok(self - .db - .get_pinned_cf(&self.cf_handle(cf)?, key) - .map_err(|err| Error::InternalError(format!("get_cf failed: {}", err)))? - .is_some()) - } - - #[inline(always)] - fn multi_get( - &self, - cf: crate::ColumnFamily, - keys: Vec, - ) -> crate::Result>> - where - T: Deserialize, - U: AsRef<[u8]>, - { - let cf_handle = self.cf_handle(cf)?; - let mut results = Vec::with_capacity(keys.len()); - for value in self - .db - .multi_get_cf(keys.iter().map(|key| (&cf_handle, key)).collect::>()) - { - results.push( - if let Some(bytes) = value - .map_err(|err| Error::InternalError(format!("multi_get_cf failed: {}", err)))? - { - T::deserialize(&bytes) - .ok_or_else(|| { - Error::DeserializeError("Failed to deserialize keys.".to_string()) - })? - .into() - } else { - None - }, - ); - } - - Ok(results) - } - - #[inline(always)] - fn iterator<'x>( - &'x self, - cf: crate::ColumnFamily, - start: &[u8], - direction: crate::Direction, - ) -> crate::Result> { - Ok(RocksDBIterator { - it: self.db.iterator_cf( - &self.cf_handle(cf)?, - rocksdb::IteratorMode::From( - start, - match direction { - crate::Direction::Forward => rocksdb::Direction::Forward, - crate::Direction::Backward => rocksdb::Direction::Reverse, - }, - ), - ), - }) - } - - fn compact(&self, cf: crate::ColumnFamily) -> crate::Result<()> { - self.db - .compact_range_cf(&self.cf_handle(cf)?, None::<&[u8]>, None::<&[u8]>); - Ok(()) - } - - fn open() -> crate::Result { - // Create the database directory if it doesn't exist - let path = PathBuf::from( - "/tmp/rocksdb.test", /*&settings - .get("db-path") - .unwrap_or_else(|| "/usr/local/stalwart-jmap/data".to_string())*/ - ); - let mut idx_path = path; - idx_path.push("idx"); - std::fs::create_dir_all(&idx_path).map_err(|err| { - Error::InternalError(format!( - "Failed to create index directory {}: {:?}", - idx_path.display(), - err - )) - })?; - - // Bitmaps - let cf_bitmaps = { - let mut cf_opts = Options::default(); - //cf_opts.set_max_write_buffer_number(16); - //cf_opts.set_merge_operator("merge", bitmap_merge, bitmap_partial_merge); - //cf_opts.set_compaction_filter("compact", bitmap_compact); - ColumnFamilyDescriptor::new("bitmaps", cf_opts) - }; - - // Stored values - let cf_values = { - let mut cf_opts = Options::default(); - cf_opts.set_merge_operator_associative("merge", numeric_value_merge); - ColumnFamilyDescriptor::new("values", cf_opts) - }; - - // Secondary indexes - let cf_indexes = { - let cf_opts = Options::default(); - ColumnFamilyDescriptor::new("indexes", cf_opts) - }; - - // Blobs - let cf_blobs = { - let mut cf_opts = Options::default(); - cf_opts.set_enable_blob_files(true); - cf_opts.set_min_blob_size( - 16834, /*settings.parse("blob-min-size").unwrap_or(16384) */ - ); - ColumnFamilyDescriptor::new("blobs", cf_opts) - }; - - // Raft log and change log - let cf_log = { - let cf_opts = Options::default(); - ColumnFamilyDescriptor::new("logs", cf_opts) - }; - - let mut db_opts = Options::default(); - db_opts.create_missing_column_families(true); - db_opts.create_if_missing(true); - - Ok(RocksDB { - db: OptimisticTransactionDB::open_cf_descriptors( - &db_opts, - idx_path, - vec![cf_bitmaps, cf_values, cf_indexes, cf_blobs, cf_log], - ) - .map_err(|e| Error::InternalError(e.into_string()))?, - }) - } - - fn close(&self) -> crate::Result<()> { - self.db - .flush() - .map_err(|e| Error::InternalError(e.to_string()))?; - self.db.cancel_all_background_work(true); - Ok(()) - } -} - -impl RocksDB { - #[inline(always)] - fn cf_handle(&self, cf: crate::ColumnFamily) -> crate::Result> { - self.db - .cf_handle(match cf { - crate::ColumnFamily::Bitmaps => "bitmaps", - crate::ColumnFamily::Values => "values", - crate::ColumnFamily::Indexes => "indexes", - crate::ColumnFamily::Blobs => "blobs", - crate::ColumnFamily::Logs => "logs", - }) - .ok_or_else(|| { - Error::InternalError(format!( - "Failed to get handle for '{:?}' column family.", - cf - )) - }) - } -} - -pub fn numeric_value_merge( - _key: &[u8], - value: Option<&[u8]>, - operands: &MergeOperands, -) -> Option> { - let mut value = if let Some(value) = value { - i64::from_le_bytes(value.try_into().ok()?) - } else { - 0 - }; - - for op in operands.iter() { - value += i64::from_le_bytes(op.try_into().ok()?); - } - - let mut bytes = Vec::with_capacity(std::mem::size_of::()); - bytes.extend_from_slice(&value.to_le_bytes()); - Some(bytes) -} diff --git a/src/tests/query.rs b/src/tests/query.rs index ac53b441..a1f526cc 100644 --- a/src/tests/query.rs +++ b/src/tests/query.rs @@ -91,14 +91,21 @@ const FIELDS_OPTIONS: [FieldType; 20] = [ FieldType::Text, // "url", ]; -#[test] -pub fn db_test() { - let db = Store::open().unwrap(); - test(&db, false); +#[tokio::test] +pub async fn db_test() { + let db = Arc::new(Store::open().await.unwrap()); + let insert = false; + if insert { + let trx = db.db.create_trx().unwrap(); + trx.clear_range(&[0u8], &[u8::MAX]); + trx.commit().await.unwrap(); + } + + test(db, insert).await; } #[allow(clippy::mutex_atomic)] -pub fn test(db: &Store, do_insert: bool) { +pub async fn test(db: Arc, do_insert: bool) { let pool = rayon::ThreadPoolBuilder::new() .num_threads(8) .build() @@ -123,8 +130,7 @@ pub fn test(db: &Store, do_insert: bool) { builder .with_account_id(0) .with_collection(COLLECTION_ID) - .update_document(document_id as u32) // Speed up insertion by manually assigning id - .bitmap(u8::MAX, (), 0); + .create_document(document_id as u32); for (pos, field) in record.iter().enumerate() { let field_id = pos as u8; match FIELDS_OPTIONS[pos] { @@ -179,37 +185,40 @@ pub fn test(db: &Store, do_insert: bool) { now.elapsed().as_millis() ); - let db_ = Arc::new(db); let now = Instant::now(); + let batches = documents.lock().unwrap().drain(..).collect::>(); + let mut chunk = Vec::new(); - pool.scope_fifo(|s| { - let mut documents = documents.lock().unwrap(); - - for document in documents.drain(..) { - let db = db_.clone(); - s.spawn_fifo(move |_| { - db.write(document).unwrap(); - }); + for batch in batches { + chunk.push({ + let db = db.clone(); + tokio::spawn(async move { db.write(batch).await }) + }); + if chunk.len() == 1000 { + for handle in chunk { + handle.await.unwrap().unwrap(); + } + chunk = Vec::new(); } - }); + } + + if !chunk.is_empty() { + for handle in chunk { + handle.await.unwrap().unwrap(); + } + } println!("Insert took {} ms.", now.elapsed().as_millis()); } - println!("Running filter tests..."); - test_filter(db); + //println!("Running filter tests..."); + //test_filter(db.clone()).await; println!("Running sort tests..."); - test_sort(db); + test_sort(db).await; } -impl IntoBitmap for () { - fn into_bitmap(self) -> (Vec, u8) { - (vec![], BM_DOCUMENT_IDS) - } -} - -pub fn test_filter(db: &Store) { +pub async fn test_filter(db: Arc) { let mut fields = AHashMap::default(); for (field_num, field) in FIELDS.iter().enumerate() { fields.insert(field.to_string(), field_num as u8); @@ -317,13 +326,11 @@ pub fn test_filter(db: &Store) { ]; for (filter, expected_results) in tests { - println!("Running test: {:?}", filter); + //println!("Running test: {:?}", filter); let mut results: Vec = Vec::with_capacity(expected_results.len()); - let docset = db.filter(0, COLLECTION_ID, filter).unwrap(); + let docset = db.filter(0, COLLECTION_ID, filter).await.unwrap(); let sorted_docset = db .sort( - 0, - COLLECTION_ID, docset, vec![Comparator::ascending(fields["accession_number"])], 0, @@ -331,16 +338,20 @@ pub fn test_filter(db: &Store) { None, 0, ) + .await .unwrap(); + let db = db.read_transaction().await.unwrap(); for document_id in sorted_docset.ids { results.push( db.get_value(ValueKey { account_id: 0, collection: COLLECTION_ID, document_id, + family: 0, field: fields["accession_number"], }) + .await .unwrap() .unwrap(), ); @@ -349,7 +360,7 @@ pub fn test_filter(db: &Store) { } } -pub fn test_sort(db: &Store) { +pub async fn test_sort(db: Arc) { let mut fields = AHashMap::default(); for (field_num, field) in FIELDS.iter().enumerate() { fields.insert(field.to_string(), field_num as u8); @@ -410,28 +421,23 @@ pub fn test_sort(db: &Store) { for (filter, sort, expected_results) in tests { let mut results: Vec = Vec::with_capacity(expected_results.len()); - let docset = db.filter(0, COLLECTION_ID, filter).unwrap(); + let docset = db.filter(0, COLLECTION_ID, filter).await.unwrap(); let sorted_docset = db - .sort( - 0, - COLLECTION_ID, - docset, - sort, - expected_results.len(), - 0, - None, - 0, - ) + .sort(docset, sort, expected_results.len(), 0, None, 0) + .await .unwrap(); + let db = db.read_transaction().await.unwrap(); for document_id in sorted_docset.ids { results.push( db.get_value(ValueKey { account_id: 0, collection: COLLECTION_ID, document_id, + family: 0, field: fields["accession_number"], }) + .await .unwrap() .unwrap(), ); diff --git a/src/write/batch.rs b/src/write/batch.rs index e4590c4e..907eeefb 100644 --- a/src/write/batch.rs +++ b/src/write/batch.rs @@ -1,4 +1,4 @@ -use crate::{BM_TERM, TERM_EXACT}; +use crate::{BM_DOCUMENT_IDS, BM_TERM, TERM_EXACT}; use super::{ Batch, BatchBuilder, HasFlag, IntoBitmap, IntoOperations, Operation, Serialize, Tokenize, @@ -7,10 +7,7 @@ use super::{ impl BatchBuilder { pub fn new() -> Self { - Self { - ops: Vec::new(), - last_collection: 0, - } + Self { ops: Vec::new() } } pub fn with_account_id(&mut self, account_id: u32) -> &mut Self { @@ -19,32 +16,34 @@ impl BatchBuilder { } pub fn with_collection(&mut self, collection: impl Into) -> &mut Self { - self.last_collection = collection.into(); self.ops.push(Operation::Collection { - collection: self.last_collection, + collection: collection.into(), }); self } - pub fn create_document(&mut self) -> &mut Self { - self.ops.push(Operation::DocumentId { - document_id: u32::MAX, + pub fn create_document(&mut self, document_id: u32) -> &mut Self { + self.ops.push(Operation::DocumentId { document_id }); + self.ops.push(Operation::Bitmap { + family: BM_DOCUMENT_IDS, + field: u8::MAX, + key: vec![], set: true, }); self } pub fn update_document(&mut self, document_id: u32) -> &mut Self { - self.ops.push(Operation::DocumentId { - document_id, - set: true, - }); + self.ops.push(Operation::DocumentId { document_id }); self } pub fn delete_document(&mut self, document_id: u32) -> &mut Self { - self.ops.push(Operation::DocumentId { - document_id, + self.ops.push(Operation::DocumentId { document_id }); + self.ops.push(Operation::Bitmap { + family: BM_DOCUMENT_IDS, + field: u8::MAX, + key: vec![], set: false, }); self @@ -83,6 +82,7 @@ impl BatchBuilder { if options.has_flag(F_VALUE) { self.ops.push(Operation::Value { field, + family: 0, set: if is_set { Some(value) } else { None }, }); } diff --git a/src/write/key.rs b/src/write/key.rs index e844d50c..256064db 100644 --- a/src/write/key.rs +++ b/src/write/key.rs @@ -10,8 +10,8 @@ pub trait KeySerialize { } pub trait DeserializeBigEndian { - fn deserialize_be_u32(&self, index: usize) -> Option; - fn deserialize_be_u64(&self, index: usize) -> Option; + fn deserialize_be_u32(&self, index: usize) -> crate::Result; + fn deserialize_be_u64(&self, index: usize) -> crate::Result; } impl KeySerializer { @@ -79,21 +79,37 @@ impl KeySerialize for u64 { } impl DeserializeBigEndian for &[u8] { - fn deserialize_be_u32(&self, index: usize) -> Option { - u32::from_be_bytes( - self.get(index..index + std::mem::size_of::())? - .try_into() - .ok()?, - ) - .into() + fn deserialize_be_u32(&self, index: usize) -> crate::Result { + self.get(index..index + std::mem::size_of::()) + .ok_or_else(|| { + crate::Error::InternalError( + "Index out of range while deserializing u32.".to_string(), + ) + }) + .and_then(|bytes| { + bytes.try_into().map_err(|_| { + crate::Error::InternalError( + "Index out of range while deserializing u32.".to_string(), + ) + }) + }) + .map(u32::from_be_bytes) } - fn deserialize_be_u64(&self, index: usize) -> Option { - u64::from_be_bytes( - self.get(index..index + std::mem::size_of::())? - .try_into() - .ok()?, - ) - .into() + fn deserialize_be_u64(&self, index: usize) -> crate::Result { + self.get(index..index + std::mem::size_of::()) + .ok_or_else(|| { + crate::Error::InternalError( + "Index out of range while deserializing u64.".to_string(), + ) + }) + .and_then(|bytes| { + bytes.try_into().map_err(|_| { + crate::Error::InternalError( + "Index out of range while deserializing u64.".to_string(), + ) + }) + }) + .map(u64::from_be_bytes) } } diff --git a/src/write/log.rs b/src/write/log.rs index 4b628056..7d09ca25 100644 --- a/src/write/log.rs +++ b/src/write/log.rs @@ -5,13 +5,14 @@ use crate::Serialize; use super::{IntoOperations, Operation}; +#[derive(Default)] pub struct ChangeLogBuilder { pub change_id: u64, - pub changes: VecMap, + pub changes: VecMap, } #[derive(Default)] -pub struct Change { +pub struct Changes { pub inserts: AHashSet, pub updates: AHashSet, pub deletes: AHashSet, @@ -69,14 +70,10 @@ impl ChangeLogBuilder { impl IntoOperations for ChangeLogBuilder { fn build(self, batch: &mut super::BatchBuilder) -> crate::Result<()> { for (collection, changes) in self.changes { - if collection != batch.last_collection { - batch.last_collection = collection; - batch.ops.push(Operation::Collection { collection }); - } - batch.ops.push(Operation::Log { change_id: self.change_id, - changes: changes.serialize(), + collection, + set: changes.serialize(), }); } @@ -84,7 +81,7 @@ impl IntoOperations for ChangeLogBuilder { } } -impl Serialize for Change { +impl Serialize for Changes { fn serialize(self) -> Vec { let mut buf = Vec::with_capacity( 1 + (self.inserts.len() @@ -99,6 +96,7 @@ impl Serialize for Change { buf.push_leb128(self.updates.len()); buf.push_leb128(self.child_updates.len()); buf.push_leb128(self.deletes.len()); + for list in [self.inserts, self.updates, self.child_updates, self.deletes] { for id in list { buf.push_leb128(id); diff --git a/src/write/mod.rs b/src/write/mod.rs index 2fd2b63f..56fb7600 100644 --- a/src/write/mod.rs +++ b/src/write/mod.rs @@ -16,7 +16,6 @@ pub struct Batch { } pub struct BatchBuilder { - pub last_collection: u8, pub ops: Vec, } @@ -29,10 +28,10 @@ pub enum Operation { }, DocumentId { document_id: u32, - set: bool, }, Value { field: u8, + family: u8, set: Option>, }, Index { @@ -46,11 +45,6 @@ pub enum Operation { key: Vec, set: bool, }, - Bloom { - field: u8, - family: u8, - set: Option>, - }, Blob { key: Vec, set: bool, @@ -61,7 +55,8 @@ pub enum Operation { }, Log { change_id: u64, - changes: Vec, + collection: u8, + set: Vec, }, } @@ -108,8 +103,8 @@ impl Serialize for Vec { } impl Deserialize for String { - fn deserialize(bytes: &[u8]) -> Option { - String::from_utf8_lossy(bytes).into_owned().into() + fn deserialize(bytes: &[u8]) -> crate::Result { + Ok(String::from_utf8_lossy(bytes).into_owned()) } }