filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

FoundationDB first impl

Browse Source
This commit is contained in:
Mauro D
2023-03-30 06:26:53 +00:00
parent 05d53da6fb
commit d74731f3c6
25 changed files with 1555 additions and 725 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
src/backend/foundationdb/bitmap.rs Normal file
View File

@@ -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::<u128>() * 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::<u128>() * 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::<u128>() * 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::<u128>() * WORD_SIZE as usize
);
for (word_num, word) in bytes.chunks_exact(std::mem::size_of::<u128>()).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);
}
}
}

21
src/backend/foundationdb/main.rs Normal file
View File

@@ -0,0 +1,21 @@
use foundationdb::Database;
use crate::Store;
impl Store {
pub async fn open() -> crate::Result<Self> {
Ok(Self {
guard: unsafe { foundationdb::boot() },
db: Database::default()?,
})
}
}
/*
impl Drop for Store {
fn drop(&mut self) {
self.guard.drop();
self.db.drop();
}
}
*/

125
src/backend/foundationdb/mod.rs
View File

@@ -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<T: AsRef<[u8]>> Serialize for &IndexKey<T> {
fn serialize(self) -> Vec<u8> {
let key = self.key.as_ref();
KeySerializer::new(std::mem::size_of::<IndexKey<T>>() + 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<u8> {
KeySerializer::new(std::mem::size_of::<IndexKeyPrefix>() + 1)
.write(SUBSPACE_INDEXES)
.write(self.account_id)
.write(self.collection)
.write(self.field)
.finalize()
}
}
impl Serialize for &ValueKey {
fn serialize(self) -> Vec<u8> {
if self.family == 0 {
KeySerializer::new(std::mem::size_of::<ValueKey>() + 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::<ValueKey>() + 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<T: AsRef<[u8]>> Serialize for &BitmapKey<T> {
fn serialize(self) -> Vec<u8> {
let key = self.key.as_ref();
KeySerializer::new(std::mem::size_of::<BitmapKey<T>>() + 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<T: AsRef<[u8]>> Serialize for &BlobKey<T> {
fn serialize(self) -> Vec<u8> {
let hash = self.hash.as_ref();
KeySerializer::new(std::mem::size_of::<BlobKey<T>>() + 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<u8> {
KeySerializer::new(std::mem::size_of::<AclKey>() + 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<u8> {
KeySerializer::new(std::mem::size_of::<LogKey>() + 1)
.write(SUBSPACE_LOGS)
.write(self.account_id)
.write(self.collection)
.write(self.change_id)
.finalize()
}
}
impl From<FdbError> for Error {
fn from(error: FdbError) -> Self {
Self::InternalError(format!("FoundationDB error: {}", error.message()))
}
}

312
src/backend/foundationdb/read.rs Normal file
View File

@@ -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<U>(&self, key: ValueKey) -> crate::Result<Option<U>>
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<U>(&self, keys: Vec<ValueKey>) -> crate::Result<Vec<Option<U>>>
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<Option<RoaringBitmap>> {
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<T: AsRef<[u8]>>(
&self,
mut key: BitmapKey<T>,
) -> crate::Result<Option<RoaringBitmap>> {
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::<u32>())?,
);
}
//println!("deserializing bitmap: {:?} {:?}", value.key(), bm.len());
}
Ok(if !bm.is_empty() { Some(bm) } else { None })
}
#[inline(always)]
async fn get_bitmaps<T: AsRef<[u8]>>(
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Vec<Option<RoaringBitmap>>> {
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<T: AsRef<[u8]>>(
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut result: Option<RoaringBitmap> = 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<T: AsRef<[u8]>>(
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut result: Option<RoaringBitmap> = 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<u8>,
op: Operator,
) -> crate::Result<Option<RoaringBitmap>> {
let k1 = KeySerializer::new(
std::mem::size_of::<IndexKey<&[u8]>>() + value.len() + 1 + std::mem::size_of::<u32>(),
)
.write(SUBSPACE_INDEXES)
.write(account_id)
.write(collection)
.write(field);
let k2 = KeySerializer::new(
std::mem::size_of::<IndexKey<&[u8]>>() + value.len() + 1 + std::mem::size_of::<u32>(),
)
.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::<u32>())?);
}
}
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<Vec<SortedId>> {
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::<u32>())
.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<ReadTransaction<'_>> {
Ok(ReadTransaction {
db: &self.db,
trx: self.db.create_trx()?,
trx_age: Instant::now(),
})
}
}

166
src/backend/foundationdb/write.rs Normal file
View File

@@ -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());
}
}
}
}
}
}

2
src/backend/mod.rs
View File

@@ -1,2 +1,4 @@
#[cfg(feature = "foundation")]
pub mod foundationdb;
#[cfg(feature = "rocks")]
pub mod rocksdb;

123
src/backend/rocksdb/log.rs Normal file
View File

@@ -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::<u32>() + std::mem::size_of::<u8>();
impl Store {
pub fn get_last_change_id(
&self,
account_id: u32,
collection: impl Into<u8>,
) -> crate::Result<Option<u64>> {
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<u8>,
query: Query,
) -> crate::Result<Option<Changes>> {
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))
}
}

36
src/backend/rocksdb/mod.rs
View File

@@ -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<T: AsRef<[u8]>> Serialize for IndexKey<T> {
impl Serialize for ValueKey {
fn serialize(self) -> Vec<u8> {
KeySerializer::new(std::mem::size_of::<ValueKey>())
.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::<ValueKey>())
.write_leb128(self.account_id)
.write(self.collection)
.write_leb128(self.document_id)
.write(self.field)
.finalize()
} else {
KeySerializer::new(std::mem::size_of::<ValueKey>() + 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<u8> {
KeySerializer::new(std::mem::size_of::<BitmapKey<&[u8]>>() + 2)
.write_leb128(account_id)
.write(collection)
.write(BM_BLOOM)
.write(field)
.write(self.as_high_rank_hash())
.finalize()
}
}
impl From<rocksdb::Error> for crate::Error {
fn from(value: rocksdb::Error) -> Self {
Self::InternalError(format!("RocksDB error: {}", value))

61
src/backend/rocksdb/write.rs
View File

@@ -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<WriteResult> {
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::<ValueKey>())
.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

24
src/fts/bloom.rs
View File

@@ -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<T: Hash + AsRef<[u8]> + ?Sized>(item: &T) -> Self;
}
impl BloomHash {
pub fn hash<T: Hash + AsRef<[u8]> + ?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<Self> {
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<Self> {
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))
}
}

24
src/fts/builder.rs
View File

@@ -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(),
});
}

4
src/fts/mod.rs
View File

@@ -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;

156
src/fts/query.rs
View File

@@ -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<Option<RoaringBitmap>> {
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,
<Vec<BloomHashGroup>>::to_ngrams(&tokens, 2),
BLOOM_BIGRAM,
BM_BLOOM | BLOOM_BIGRAM,
),
_ => (
bitmaps,
<Vec<BloomHashGroup>>::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::<ValueKey>())
.write_leb128(account_id)
.write(collection)
.write_leb128(document_id)
.write(u8::MAX)
.write(BM_BLOOM | family)
.write(field)
.finalize()
})
.collect::<Vec<_>>();
self.get_values::<BloomFilter>(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::<BloomFilter>(
bitmaps
.iter()
.map(|document_id| ValueKey {
account_id,
collection,
document_id,
family,
field,
})
.collect::<Vec<_>>(),
)
.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::<ValueKey>() + 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::<BloomFilter>(key)? {
if matched {
return Some(document_id);
}
}
None
})
.collect::<RoaringBitmap>();*/
let mut bm = RoaringBitmap::new();
for document_id in bitmaps {
trx.refresh_if_old().await?;
if let Some(bloom) = trx
.get_value::<BloomFilter>(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<u8> {
KeySerializer::new(std::mem::size_of::<BitmapKey<&[u8]>>() + 2)
.write_leb128(account_id)
.write(collection)
.write(BM_BLOOM)
.write(field)
.write(self.as_high_rank_hash(n))
.finalize()
) -> BitmapKey<Vec<u8>> {
BitmapKey {
account_id,
collection,
family: BM_BLOOM,
field,
block_num: 0,
key: self.as_high_rank_hash().serialize(),
}
}
}

26
src/lib.rs
View File

@@ -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<MultiThreaded>,
db: rocksdb::OptimisticTransactionDB<rocksdb::MultiThreaded>,
}
#[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<Self>;
fn deserialize(bytes: &[u8]) -> crate::Result<Self>;
}
pub trait Serialize {
@@ -26,6 +32,8 @@ pub struct BitmapKey<T: AsRef<[u8]>> {
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<T: AsRef<[u8]>> {
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;

143
src/query/filter.rs
View File

@@ -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<Filter>,
) -> crate::Result<ResultSet> {
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::<IndexKey<&[u8]>>() + 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,
})

109
src/query/log.rs Normal file
View File

@@ -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<Change>,
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(())
}
}

38
src/query/mod.rs
View File

@@ -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<u32>),
}
#[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<u8>, 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<u32> {
fn contains_id(&self, id: u32) -> bool {
self.iter().any(|&i| i == id)
}
fn len(&self) -> usize {
self.len()
}
}

225
src/query/sort.rs
View File

@@ -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<DBIteratorWithThreadMode<'x, OptimisticTransactionDB<MultiThreaded>>>,
prefix: Vec<u8>,
start_key: Vec<u8>,
from_key: Vec<u8>,
ascending: bool,
prev_item: Option<u32>,
prev_key: Option<Box<[u8]>>,
},
}
#[cfg(feature = "rocks")]
struct IndexIterator<'x> {
index: IndexType<'x>,
remaining: RoaringBitmap,
eof: bool,
}
#[cfg(feature = "foundation")]
enum IndexType<'x, T: Stream<Item = FdbResult<FdbValue>> + Unpin + 'x> {
DocumentSet {
set: RoaringBitmap,
it: Option<roaring::bitmap::IntoIter>,
},
DB {
it: Option<T>,
from_key: Vec<u8>,
to_key: Vec<u8>,
ascending: bool,
prev_item: Option<u32>,
prev_key: Option<Box<[u8]>>,
phantom: std::marker::PhantomData<&'x ()>,
},
}
#[cfg(feature = "foundation")]
struct IndexIterator<'x, T: Stream<Item = FdbResult<FdbValue>> + 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<Comparator>,
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::<Vec<_>>();
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::<u32>()..)
.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::<u32>())?;
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,
}
}
}

343
src/rocksdb.rs
View File

@@ -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 <http://www.gnu.org/licenses/>.
*
* 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<MultiThreaded>,
}
pub struct RocksDBIterator<'x> {
it: DBIteratorWithThreadMode<'x, OptimisticTransactionDB<MultiThreaded>>,
}
impl Iterator for RocksDBIterator<'_> {
type Item = (Box<[u8]>, Box<[u8]>);
#[allow(clippy::while_let_on_iterator)]
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
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<U>(&self, cf: crate::ColumnFamily, key: &[u8]) -> crate::Result<Option<U>>
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<WriteOperation>) -> 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<bool> {
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<T, U>(
&self,
cf: crate::ColumnFamily,
keys: Vec<U>,
) -> crate::Result<Vec<Option<T>>>
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::<Vec<_>>())
{
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<Self::Iterator<'x>> {
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<Self> {
// 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<Arc<BoundColumnFamily>> {
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<Vec<u8>> {
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::<i64>());
bytes.extend_from_slice(&value.to_le_bytes());
Some(bytes)
}

92
src/tests/query.rs
View File

@@ -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<Store>, 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::<Vec<_>>();
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>, u8) {
(vec![], BM_DOCUMENT_IDS)
}
}
pub fn test_filter(db: &Store) {
pub async fn test_filter(db: Arc<Store>) {
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<String> = 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<Store>) {
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<String> = 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(),
);

32
src/write/batch.rs
View File

@@ -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<u8>) -> &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 },
});
}

48
src/write/key.rs
View File

@@ -10,8 +10,8 @@ pub trait KeySerialize {
}
pub trait DeserializeBigEndian {
fn deserialize_be_u32(&self, index: usize) -> Option<u32>;
fn deserialize_be_u64(&self, index: usize) -> Option<u64>;
fn deserialize_be_u32(&self, index: usize) -> crate::Result<u32>;
fn deserialize_be_u64(&self, index: usize) -> crate::Result<u64>;
}
impl KeySerializer {
@@ -79,21 +79,37 @@ impl KeySerialize for u64 {
}
impl DeserializeBigEndian for &[u8] {
fn deserialize_be_u32(&self, index: usize) -> Option<u32> {
u32::from_be_bytes(
self.get(index..index + std::mem::size_of::<u32>())?
.try_into()
.ok()?,
)
.into()
fn deserialize_be_u32(&self, index: usize) -> crate::Result<u32> {
self.get(index..index + std::mem::size_of::<u32>())
.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> {
u64::from_be_bytes(
self.get(index..index + std::mem::size_of::<u64>())?
.try_into()
.ok()?,
)
.into()
fn deserialize_be_u64(&self, index: usize) -> crate::Result<u64> {
self.get(index..index + std::mem::size_of::<u64>())
.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)
}
}

16
src/write/log.rs
View File

@@ -5,13 +5,14 @@ use crate::Serialize;
use super::{IntoOperations, Operation};
#[derive(Default)]
pub struct ChangeLogBuilder {
pub change_id: u64,
pub changes: VecMap<u8, Change>,
pub changes: VecMap<u8, Changes>,
}
#[derive(Default)]
pub struct Change {
pub struct Changes {
pub inserts: AHashSet<u64>,
pub updates: AHashSet<u64>,
pub deletes: AHashSet<u64>,
@@ -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<u8> {
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);

15
src/write/mod.rs
View File

@@ -16,7 +16,6 @@ pub struct Batch {
}
pub struct BatchBuilder {
pub last_collection: u8,
pub ops: Vec<Operation>,
}
@@ -29,10 +28,10 @@ pub enum Operation {
},
DocumentId {
document_id: u32,
set: bool,
},
Value {
field: u8,
family: u8,
set: Option<Vec<u8>>,
},
Index {
@@ -46,11 +45,6 @@ pub enum Operation {
key: Vec<u8>,
set: bool,
},
Bloom {
field: u8,
family: u8,
set: Option<Vec<u8>>,
},
Blob {
key: Vec<u8>,
set: bool,
@@ -61,7 +55,8 @@ pub enum Operation {
},
Log {
change_id: u64,
changes: Vec<u8>,
collection: u8,
set: Vec<u8>,
},
}
@@ -108,8 +103,8 @@ impl Serialize for Vec<u8> {
}
impl Deserialize for String {
fn deserialize(bytes: &[u8]) -> Option<Self> {
String::from_utf8_lossy(bytes).into_owned().into()
fn deserialize(bytes: &[u8]) -> crate::Result<Self> {
Ok(String::from_utf8_lossy(bytes).into_owned())
}
}