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

SQLite support

Browse Source
This commit is contained in:
Mauro D
2023-04-02 17:10:59 +00:00
parent b4e392d1c2
commit 4d192de2fe
19 changed files with 1086 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/backend/foundationdb/bitmap.rs
View File

@@ -1,7 +1,7 @@
use ahash::AHashSet;
use roaring::RoaringBitmap;
const WORD_SIZE_BITS: u32 = 128;
const WORD_SIZE_BITS: u32 = (WORD_SIZE * 8) as u32;
const WORD_SIZE: usize = std::mem::size_of::<u128>();
const WORDS_PER_BLOCK: u32 = 8;
pub const BITS_PER_BLOCK: u32 = WORD_SIZE_BITS * WORDS_PER_BLOCK;

40
src/backend/foundationdb/read.rs
View File

@@ -5,7 +5,7 @@ use std::{
use foundationdb::{
options::{self, StreamingMode},
Database, KeySelector, RangeOption, Transaction,
KeySelector, RangeOption,
};
use futures::StreamExt;
use roaring::RoaringBitmap;
@@ -13,17 +13,11 @@ use roaring::RoaringBitmap;
use crate::{
query::Operator,
write::key::{DeserializeBigEndian, KeySerializer},
BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, Serialize, Store, ValueKey, BM_DOCUMENT_IDS,
BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, ReadTransaction, Serialize, Store, ValueKey,
};
use super::{bitmap::DeserializeBlock, 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>>
@@ -39,36 +33,6 @@ impl ReadTransaction<'_> {
}
}
#[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
}
async fn get_bitmap_<T: AsRef<[u8]>>(
&self,
mut key: BitmapKey<T>,

2
src/backend/foundationdb/write.rs
View File

@@ -310,7 +310,7 @@ impl Store {
key: &[],
}
.serialize();
trx.get(&key, false).await?;
trx.get(&key, false).await?; // Read to create conflict range
trx.set(&key, &now().serialize());
match trx.commit().await {

2
src/backend/mod.rs
View File

@@ -2,3 +2,5 @@
pub mod foundationdb;
#[cfg(feature = "rocks")]
pub mod rocksdb;
#[cfg(feature = "sqlite")]
pub mod sqlite;

92
src/backend/sqlite/main.rs Normal file
View File

@@ -0,0 +1,92 @@
use r2d2::Pool;
use tokio::sync::oneshot;
use crate::Store;
use super::pool::SqliteConnectionManager;
impl Store {
// TODO configure rayon thread pool
// TODO configure r2d2 pool
pub async fn open() -> crate::Result<Self> {
let db = Self {
conn_pool: Pool::new(
SqliteConnectionManager::file("/tmp/sqlite.db")
.with_init(|c| c.execute_batch("PRAGMA journal_mode=WAL;")),
)?,
worker_pool: rayon::ThreadPoolBuilder::new().build().map_err(|err| {
crate::Error::InternalError(format!("Failed to build worker pool: {}", err))
})?,
};
db.create_tables()?;
Ok(db)
}
pub(super) fn create_tables(&self) -> crate::Result<()> {
let conn = self.conn_pool.get()?;
for table in ["v", "l", "o", "c"] {
conn.execute(
&format!(
"CREATE TABLE IF NOT EXISTS {table} (
k BLOB PRIMARY KEY,
v BLOB NOT NULL
)"
),
[],
)?;
}
conn.execute(
"CREATE TABLE IF NOT EXISTS i (
k BLOB PRIMARY KEY
)",
[],
)?;
conn.execute(
"CREATE TABLE IF NOT EXISTS b (
z BLOB PRIMARY KEY,
a INTEGER NOT NULL DEFAULT 0,
b INTEGER NOT NULL DEFAULT 0,
c INTEGER NOT NULL DEFAULT 0,
d INTEGER NOT NULL DEFAULT 0,
e INTEGER NOT NULL DEFAULT 0,
f INTEGER NOT NULL DEFAULT 0,
g INTEGER NOT NULL DEFAULT 0,
h INTEGER NOT NULL DEFAULT 0,
i INTEGER NOT NULL DEFAULT 0,
j INTEGER NOT NULL DEFAULT 0,
k INTEGER NOT NULL DEFAULT 0,
l INTEGER NOT NULL DEFAULT 0,
m INTEGER NOT NULL DEFAULT 0,
n INTEGER NOT NULL DEFAULT 0,
o INTEGER NOT NULL DEFAULT 0,
p INTEGER NOT NULL DEFAULT 0
)",
[],
)?;
Ok(())
}
pub async fn spawn_worker<U, V>(&self, f: U) -> crate::Result<V>
where
U: FnOnce() -> crate::Result<V> + Send + 'static,
V: Sync + Send + 'static,
{
let (tx, rx) = oneshot::channel();
self.worker_pool.spawn(move || {
tx.send(f()).ok();
});
match rx.await {
Ok(result) => result,
Err(err) => Err(crate::Error::InternalError(format!(
"Worker thread failed: {}",
err
))),
}
}
}

126
src/backend/sqlite/mod.rs Normal file
View File

@@ -0,0 +1,126 @@
use crate::{
write::key::KeySerializer, AclKey, BitmapKey, BlobKey, IndexKey, IndexKeyPrefix, LogKey,
Serialize, ValueKey,
};
pub mod main;
pub mod pool;
pub mod read;
pub mod write;
const WORD_SIZE_BITS: u32 = (WORD_SIZE * 8) as u32;
const WORD_SIZE: usize = std::mem::size_of::<u64>();
const WORDS_PER_BLOCK: u32 = 16;
pub const BITS_PER_BLOCK: u32 = WORD_SIZE_BITS * WORDS_PER_BLOCK;
const BITS_MASK: u32 = BITS_PER_BLOCK - 1;
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(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(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_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_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(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(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_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(self.account_id)
.write(self.collection)
.write(self.change_id)
.finalize()
}
}
impl From<r2d2::Error> for crate::Error {
fn from(err: r2d2::Error) -> Self {
Self::InternalError(format!("Connection pool error: {}", err))
}
}
impl From<rusqlite::Error> for crate::Error {
fn from(err: rusqlite::Error) -> Self {
Self::InternalError(format!("SQLite error: {}", err))
}
}
impl From<rusqlite::types::FromSqlError> for crate::Error {
fn from(err: rusqlite::types::FromSqlError) -> Self {
Self::InternalError(format!("SQLite error: {}", err))
}
}

105
src/backend/sqlite/pool.rs Normal file
View File

@@ -0,0 +1,105 @@
use rusqlite::{Connection, Error, OpenFlags};
use std::fmt;
use std::path::{Path, PathBuf};
#[derive(Debug)]
enum Source {
File(PathBuf),
Memory,
}
type InitFn = dyn Fn(&mut Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static;
/// An `r2d2::ManageConnection` for `rusqlite::Connection`s.
pub struct SqliteConnectionManager {
source: Source,
flags: OpenFlags,
init: Option<Box<InitFn>>,
}
impl fmt::Debug for SqliteConnectionManager {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut builder = f.debug_struct("SqliteConnectionManager");
let _ = builder.field("source", &self.source);
let _ = builder.field("flags", &self.source);
let _ = builder.field("init", &self.init.as_ref().map(|_| "InitFn"));
builder.finish()
}
}
impl SqliteConnectionManager {
/// Creates a new `SqliteConnectionManager` from file.
///
/// See `rusqlite::Connection::open`
pub fn file<P: AsRef<Path>>(path: P) -> Self {
Self {
source: Source::File(path.as_ref().to_path_buf()),
flags: OpenFlags::default(),
init: None,
}
}
/// Creates a new `SqliteConnectionManager` from memory.
pub fn memory() -> Self {
Self {
source: Source::Memory,
flags: OpenFlags::default(),
init: None,
}
}
/// Converts `SqliteConnectionManager` into one that sets OpenFlags upon
/// connection creation.
///
/// See `rustqlite::OpenFlags` for a list of available flags.
pub fn with_flags(self, flags: OpenFlags) -> Self {
Self { flags, ..self }
}
/// Converts `SqliteConnectionManager` into one that calls an initialization
/// function upon connection creation. Could be used to set PRAGMAs, for
/// example.
///
/// ### Example
///
/// Make a `SqliteConnectionManager` that sets the `foreign_keys` pragma to
/// true for every connection.
///
/// ```rust,no_run
/// # use r2d2_sqlite::{SqliteConnectionManager};
/// let manager = SqliteConnectionManager::file("app.db")
/// .with_init(|c| c.execute_batch("PRAGMA foreign_keys=1;"));
/// ```
pub fn with_init<F>(self, init: F) -> Self
where
F: Fn(&mut Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static,
{
let init: Option<Box<InitFn>> = Some(Box::new(init));
Self { init, ..self }
}
}
impl r2d2::ManageConnection for SqliteConnectionManager {
type Connection = Connection;
type Error = rusqlite::Error;
fn connect(&self) -> Result<Connection, Error> {
match self.source {
Source::File(ref path) => Connection::open_with_flags(path, self.flags),
Source::Memory => Connection::open_in_memory_with_flags(self.flags),
}
.map_err(Into::into)
.and_then(|mut c| match self.init {
None => Ok(c),
Some(ref init) => init(&mut c).map(|_| c),
})
}
fn is_valid(&self, conn: &mut Connection) -> Result<(), Error> {
conn.execute_batch("").map_err(Into::into)
}
fn has_broken(&self, _: &mut Connection) -> bool {
false
}
}

247
src/backend/sqlite/read.rs Normal file
View File

@@ -0,0 +1,247 @@
use std::ops::BitAndAssign;
use roaring::RoaringBitmap;
use rusqlite::OptionalExtension;
use crate::{
query::Operator,
write::key::{DeserializeBigEndian, KeySerializer},
BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, ReadTransaction, Serialize, Store, ValueKey,
};
use super::{BITS_PER_BLOCK, WORDS_PER_BLOCK, WORD_SIZE_BITS};
impl ReadTransaction<'_> {
#[inline(always)]
#[maybe_async::maybe_async]
pub async fn get_value<U>(&self, key: ValueKey) -> crate::Result<Option<U>>
where
U: Deserialize,
{
let key = key.serialize();
self.conn
.prepare_cached("SELECT v FROM v WHERE k = ?")?
.query_row([&key], |row| {
U::deserialize(row.get_ref(0)?.as_bytes()?)
.map_err(|err| rusqlite::Error::ToSqlConversionFailure(err.into()))
})
.optional()
.map_err(Into::into)
}
#[maybe_async::maybe_async]
async fn get_bitmap_<T: AsRef<[u8]>>(
&self,
mut key: BitmapKey<T>,
bm: &mut RoaringBitmap,
) -> crate::Result<()> {
let begin = key.serialize();
key.block_num = u32::MAX;
let end = key.serialize();
let mut query = self
.conn
.prepare_cached("SELECT z, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p FROM b WHERE z >= ? AND z <= ?")?;
let mut rows = query.query([&begin, &end])?;
while let Some(row) = rows.next()? {
let key = row.get_ref(0)?.as_bytes()?;
let block_num = key.deserialize_be_u32(key.len() - std::mem::size_of::<u32>())?;
for word_num in 0..WORDS_PER_BLOCK {
match row.get::<_, i64>((word_num + 1) as usize)? as u64 {
0 => (),
u64::MAX => {
bm.insert_range(
block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS
..(block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS)
+ WORD_SIZE_BITS,
);
}
mut word => {
while word != 0 {
let trailing_zeros = word.trailing_zeros();
bm.insert(
block_num * BITS_PER_BLOCK
+ word_num * WORD_SIZE_BITS
+ trailing_zeros,
);
word ^= 1 << trailing_zeros;
}
}
}
}
}
Ok(())
}
#[maybe_async::maybe_async]
pub async fn get_bitmap<T: AsRef<[u8]>>(
&self,
key: BitmapKey<T>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut bm = RoaringBitmap::new();
self.get_bitmap_(key, &mut bm).await?;
Ok(if !bm.is_empty() { Some(bm) } else { None })
}
#[maybe_async::maybe_async]
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 key in keys {
if let Some(bitmap) = self.get_bitmap(key).await? {
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)
}
#[maybe_async::maybe_async]
pub(crate) async fn get_bitmaps_union<T: AsRef<[u8]>>(
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut bm = RoaringBitmap::new();
for key in keys {
self.get_bitmap_(key, &mut bm).await?;
}
Ok(if !bm.is_empty() { Some(bm) } else { None })
}
#[maybe_async::maybe_async]
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(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(account_id)
.write(collection)
.write(field + matches!(op, Operator::GreaterThan | Operator::GreaterEqualThan) as u8);
let (query, begin, end) = match op {
Operator::LowerThan => (
("SELECT k FROM i WHERE k >= ? AND k < ?"),
(k1.finalize()),
(k2.write(&value[..]).write(0u32).finalize()),
),
Operator::LowerEqualThan => (
("SELECT k FROM i WHERE k >= ? AND k <= ?"),
(k1.finalize()),
(k2.write(&value[..]).write(u32::MAX).finalize()),
),
Operator::GreaterThan => (
("SELECT k FROM i WHERE k > ? AND k <= ?"),
(k1.write(&value[..]).write(u32::MAX).finalize()),
(k2.finalize()),
),
Operator::GreaterEqualThan => (
("SELECT k FROM i WHERE k >= ? AND k <= ?"),
(k1.write(&value[..]).write(0u32).finalize()),
(k2.finalize()),
),
Operator::Equal => (
("SELECT k FROM i WHERE k >= ? AND k <= ?"),
(k1.write(&value[..]).write(0u32).finalize()),
(k2.write(&value[..]).write(u32::MAX).finalize()),
),
};
let mut bm = RoaringBitmap::new();
let mut query = self.conn.prepare_cached(query)?;
let mut rows = query.query([&begin, &end])?;
while let Some(row) = rows.next()? {
let key = row.get_ref(0)?.as_bytes()?;
bm.insert(key.deserialize_be_u32(key.len() - std::mem::size_of::<u32>())?);
}
Ok(Some(bm))
}
#[maybe_async::maybe_async]
pub(crate) async fn sort_index(
&self,
account_id: u32,
collection: u8,
field: u8,
ascending: bool,
mut cb: impl FnMut(&[u8], u32) -> bool,
) -> crate::Result<()> {
let begin = IndexKeyPrefix {
account_id,
collection,
field,
}
.serialize();
let end = IndexKeyPrefix {
account_id,
collection,
field: field + 1,
}
.serialize();
let prefix_len = begin.len();
let mut query = self.conn.prepare_cached(if ascending {
"SELECT k FROM i WHERE k >= ? AND k < ? ORDER BY k ASC"
} else {
"SELECT k FROM i WHERE k >= ? AND k < ? ORDER BY k DESC"
})?;
let mut rows = query.query([&begin, &end])?;
while let Some(row) = rows.next()? {
let key = row.get_ref(0)?.as_bytes()?;
let id_pos = key.len() - std::mem::size_of::<u32>();
debug_assert!(key.starts_with(&begin));
if !cb(
key.get(prefix_len..id_pos).ok_or_else(|| {
crate::Error::InternalError("Invalid key found in index".to_string())
})?,
key.deserialize_be_u32(id_pos)?,
) {
return Ok(());
}
}
Ok(())
}
#[maybe_async::maybe_async]
pub async fn refresh_if_old(&mut self) -> crate::Result<()> {
Ok(())
}
}
impl Store {
#[maybe_async::maybe_async]
pub async fn read_transaction(&self) -> crate::Result<ReadTransaction<'static>> {
Ok(ReadTransaction {
conn: self.conn_pool.get()?,
_p: std::marker::PhantomData,
})
}
}

242
src/backend/sqlite/write.rs Normal file
View File

@@ -0,0 +1,242 @@
use rusqlite::params;
use crate::{
write::{Batch, Operation},
AclKey, BitmapKey, BlobKey, IndexKey, LogKey, Serialize, Store, ValueKey,
};
use super::{BITS_MASK, BITS_PER_BLOCK};
const INSERT_QUERIES: &[&str] = &[
"INSERT INTO b (z, a) VALUES (?, ?)",
"INSERT INTO b (z, b) VALUES (?, ?)",
"INSERT INTO b (z, c) VALUES (?, ?)",
"INSERT INTO b (z, d) VALUES (?, ?)",
"INSERT INTO b (z, e) VALUES (?, ?)",
"INSERT INTO b (z, f) VALUES (?, ?)",
"INSERT INTO b (z, g) VALUES (?, ?)",
"INSERT INTO b (z, h) VALUES (?, ?)",
"INSERT INTO b (z, i) VALUES (?, ?)",
"INSERT INTO b (z, j) VALUES (?, ?)",
"INSERT INTO b (z, k) VALUES (?, ?)",
"INSERT INTO b (z, l) VALUES (?, ?)",
"INSERT INTO b (z, m) VALUES (?, ?)",
"INSERT INTO b (z, n) VALUES (?, ?)",
"INSERT INTO b (z, o) VALUES (?, ?)",
"INSERT INTO b (z, p) VALUES (?, ?)",
];
const SET_QUERIES: &[&str] = &[
"UPDATE b SET a = a | ? WHERE z = ?",
"UPDATE b SET b = b | ? WHERE z = ?",
"UPDATE b SET c = c | ? WHERE z = ?",
"UPDATE b SET d = d | ? WHERE z = ?",
"UPDATE b SET e = e | ? WHERE z = ?",
"UPDATE b SET f = f | ? WHERE z = ?",
"UPDATE b SET g = g | ? WHERE z = ?",
"UPDATE b SET h = h | ? WHERE z = ?",
"UPDATE b SET i = i | ? WHERE z = ?",
"UPDATE b SET j = j | ? WHERE z = ?",
"UPDATE b SET k = k | ? WHERE z = ?",
"UPDATE b SET l = l | ? WHERE z = ?",
"UPDATE b SET m = m | ? WHERE z = ?",
"UPDATE b SET n = n | ? WHERE z = ?",
"UPDATE b SET o = o | ? WHERE z = ?",
"UPDATE b SET p = p | ? WHERE z = ?",
];
const CLEAR_QUERIES: &[&str] = &[
"UPDATE b SET a = a & ? WHERE z = ?",
"UPDATE b SET b = b & ? WHERE z = ?",
"UPDATE b SET c = c & ? WHERE z = ?",
"UPDATE b SET d = d & ? WHERE z = ?",
"UPDATE b SET e = e & ? WHERE z = ?",
"UPDATE b SET f = f & ? WHERE z = ?",
"UPDATE b SET g = g & ? WHERE z = ?",
"UPDATE b SET h = h & ? WHERE z = ?",
"UPDATE b SET i = i & ? WHERE z = ?",
"UPDATE b SET j = j & ? WHERE z = ?",
"UPDATE b SET k = k & ? WHERE z = ?",
"UPDATE b SET l = l & ? WHERE z = ?",
"UPDATE b SET m = m & ? WHERE z = ?",
"UPDATE b SET n = n & ? WHERE z = ?",
"UPDATE b SET o = o & ? WHERE z = ?",
"UPDATE b SET p = p & ? WHERE z = ?",
];
impl Store {
pub async fn write(&self, batch: Batch) -> crate::Result<()> {
let mut conn = self.conn_pool.get()?;
self.spawn_worker(move || {
let mut account_id = u32::MAX;
let mut collection = u8::MAX;
let mut document_id = u32::MAX;
let mut bitmap_block_num = 0;
let mut bitmap_col_num = 0;
let mut bitmap_value_set = 0i64;
let mut bitmap_value_clear = 0i64;
let trx = conn.transaction()?;
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_,
} => {
document_id = *document_id_;
bitmap_block_num = document_id / BITS_PER_BLOCK;
let index = document_id & BITS_MASK;
bitmap_col_num = (index / 64) as usize;
bitmap_value_set = (1u64 << (index as u64 & 63)) as i64;
bitmap_value_clear = (!(1u64 << (index as u64 & 63))) as i64;
}
Operation::Value { family, field, set } => {
let key = ValueKey {
account_id,
collection,
document_id,
family: *family,
field: *field,
}
.serialize();
if let Some(value) = set {
trx.prepare_cached("INSERT OR REPLACE INTO v (k, v) VALUES (?, ?)")?
.execute([&key, value])?;
} else {
trx.prepare_cached("DELETE FROM v WHERE k = ?")?
.execute([&key])?;
}
}
Operation::Index { field, key, set } => {
let key = IndexKey {
account_id,
collection,
document_id,
field: *field,
key,
}
.serialize();
if *set {
trx.prepare_cached("INSERT OR REPLACE INTO i (k) VALUES (?)")?
.execute([&key])?;
} else {
trx.prepare_cached("DELETE FROM v WHERE k = ?")?
.execute([&key])?;
}
}
Operation::Bitmap {
family,
field,
key,
set,
} => {
let key = BitmapKey {
account_id,
collection,
family: *family,
field: *field,
block_num: bitmap_block_num,
key,
}
.serialize();
if *set {
//trx.prepare_cached("INSERT OR IGNORE INTO b (z) VALUES (?)")?
// .execute([&key])?;
trx.prepare_cached(SET_QUERIES[bitmap_col_num])?
.execute(params![bitmap_value_set, &key])?;
if trx.changes() == 0 {
trx.prepare_cached(INSERT_QUERIES[bitmap_col_num])?
.execute(params![&key, bitmap_value_set])?;
}
} else {
trx.prepare_cached(CLEAR_QUERIES[bitmap_col_num])?
.execute(params![bitmap_value_clear, &key])?;
};
}
Operation::Blob { key, set } => {
let key = BlobKey {
account_id,
collection,
document_id,
hash: key,
}
.serialize();
if *set {
trx.prepare_cached("INSERT OR REPLACE INTO b (k, v) VALUES (?, ?)")?
.execute([&key[..], &[]])?;
} else {
trx.prepare_cached("DELETE FROM b WHERE k = ?")?
.execute([&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.prepare_cached("INSERT OR REPLACE INTO a (k, v) VALUES (?, ?)")?
.execute([&key, value])?;
} else {
trx.prepare_cached("DELETE FROM a WHERE k = ?")?
.execute([&key])?;
}
}
Operation::Log {
collection,
change_id,
set,
} => {
let key = LogKey {
account_id,
collection: *collection,
change_id: *change_id,
}
.serialize();
trx.prepare_cached("INSERT OR REPLACE INTO l (k, v) VALUES (?, ?)")?
.execute([&key, set])?;
}
}
}
trx.commit().map_err(Into::into)
})
.await
}
pub async fn assign_document_id(&self, account_id: u32, collection: u8) -> crate::Result<u32> {
todo!()
}
pub async fn assign_change_id(&self, account_id: u32, collection: u8) -> crate::Result<u64> {
todo!()
}
#[cfg(test)]
pub async fn destroy(&self) {
let conn = self.conn_pool.get().unwrap();
for table in ["v", "l", "o", "c", "i", "b"] {
conn.execute(&format!("DROP TABLE {table}"), []).unwrap();
}
self.create_tables().unwrap();
}
}

4
src/fts/query.rs
View File

@@ -3,7 +3,6 @@ use std::time::Instant;
use roaring::RoaringBitmap;
use crate::{
backend::foundationdb::read::ReadTransaction,
fts::{
bloom::{BloomFilter, BloomHashGroup},
builder::MAX_TOKEN_LENGTH,
@@ -11,12 +10,13 @@ use crate::{
stemmer::Stemmer,
tokenizers::Tokenizer,
},
BitmapKey, ValueKey, BLOOM_BIGRAM, BLOOM_TRIGRAM, HASH_EXACT, HASH_STEMMED,
BitmapKey, ReadTransaction, ValueKey, BLOOM_BIGRAM, BLOOM_TRIGRAM, HASH_EXACT, HASH_STEMMED,
};
use super::Language;
impl ReadTransaction<'_> {
#[maybe_async::maybe_async]
pub(crate) async fn fts_query(
&mut self,
account_id: u32,

33
src/lib.rs
View File

@@ -1,3 +1,5 @@
use std::fmt::Display;
pub mod backend;
pub mod fts;
pub mod query;
@@ -18,6 +20,25 @@ pub struct Store {
guard: foundationdb::api::NetworkAutoStop,
}
#[cfg(feature = "foundation")]
pub struct ReadTransaction<'x> {
db: &'x foundationdb::Database,
pub trx: foundationdb::Transaction,
trx_age: std::time::Instant,
}
#[cfg(feature = "sqlite")]
pub struct Store {
conn_pool: r2d2::Pool<backend::sqlite::pool::SqliteConnectionManager>,
worker_pool: rayon::ThreadPool,
}
#[cfg(feature = "sqlite")]
pub struct ReadTransaction<'x> {
conn: r2d2::PooledConnection<backend::sqlite::pool::SqliteConnectionManager>,
_p: std::marker::PhantomData<&'x ()>,
}
pub trait Deserialize: Sized + Sync + Send {
fn deserialize(bytes: &[u8]) -> crate::Result<Self>;
}
@@ -32,7 +53,6 @@ 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,
}
@@ -93,6 +113,17 @@ pub enum Error {
InternalError(String),
}
impl std::error::Error for Error {}
impl Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Error::NotFound => write!(f, "not found"),
Error::InternalError(msg) => write!(f, "internal error: {}", msg),
}
}
}
pub const BM_DOCUMENT_IDS: u8 = 0;
pub const BM_KEYWORD: u8 = 1 << 5;
pub const BM_TAG: u8 = 1 << 6;

57
src/query/filter.rs
View File

@@ -2,7 +2,7 @@ use std::ops::{BitAndAssign, BitOrAssign, BitXorAssign};
use roaring::RoaringBitmap;
use crate::{write::Tokenize, BitmapKey, Store, BM_KEYWORD};
use crate::{write::Tokenize, BitmapKey, ReadTransaction, Store, BM_KEYWORD};
use super::{Filter, ResultSet};
@@ -11,22 +11,22 @@ struct State {
bm: Option<RoaringBitmap>,
}
impl Store {
impl ReadTransaction<'_> {
#[maybe_async::maybe_async]
pub async fn filter(
&self,
&mut self,
account_id: u32,
collection: u8,
filters: Vec<Filter>,
) -> crate::Result<ResultSet> {
let mut trx = self.read_transaction().await?;
let mut not_mask = RoaringBitmap::new();
let mut not_fetch = false;
if filters.is_empty() {
return Ok(ResultSet {
account_id,
collection,
results: trx
.get_document_ids(account_id, collection)
results: self
.get_bitmap(BitmapKey::new_document_ids(account_id, collection))
.await?
.unwrap_or_else(RoaringBitmap::new),
});
@@ -37,23 +37,22 @@ impl Store {
let mut filters = filters.into_iter().peekable();
while let Some(filter) = filters.next() {
trx.refresh_if_old().await?;
self.refresh_if_old().await?;
let result = match filter {
Filter::HasKeyword { field, value } => {
trx.get_bitmap(BitmapKey {
self.get_bitmap(BitmapKey {
account_id,
collection,
family: BM_KEYWORD,
field,
key: value.as_bytes(),
#[cfg(feature = "foundation")]
block_num: 0,
})
.await?
}
Filter::HasKeywords { field, value } => {
trx.get_bitmaps_intersection(
self.get_bitmaps_intersection(
value
.tokenize()
.into_iter()
@@ -63,7 +62,6 @@ impl Store {
family: BM_KEYWORD,
field,
key: key.into_bytes(),
#[cfg(feature = "foundation")]
block_num: 0,
})
.collect(),
@@ -71,7 +69,7 @@ impl Store {
.await?
}
Filter::MatchValue { field, op, value } => {
trx.range_to_bitmap(account_id, collection, field, value, op)
self.range_to_bitmap(account_id, collection, field, value, op)
.await?
}
Filter::HasText {
@@ -80,17 +78,16 @@ impl Store {
language,
match_phrase,
} => {
trx.fts_query(account_id, collection, field, &text, language, match_phrase)
self.fts_query(account_id, collection, field, &text, language, match_phrase)
.await?
}
Filter::InBitmap { family, field, key } => {
trx.get_bitmap(BitmapKey {
self.get_bitmap(BitmapKey {
account_id,
collection,
family,
field,
key: &key,
#[cfg(feature = "foundation")]
block_num: 0,
})
.await?
@@ -113,8 +110,8 @@ impl Store {
};
if matches!(state.op, Filter::Not) && !not_fetch {
not_mask = trx
.get_document_ids(account_id, collection)
not_mask = self
.get_bitmap(BitmapKey::new_document_ids(account_id, collection))
.await?
.unwrap_or_else(RoaringBitmap::new);
not_fetch = true;
@@ -122,8 +119,6 @@ impl Store {
state.op.apply(&mut state.bm, result, &not_mask);
//println!("{:?}: {:?}", state.op, state.bm);
if matches!(state.op, Filter::And) && state.bm.as_ref().unwrap().is_empty() {
while let Some(filter) = filters.peek() {
if matches!(filter, Filter::End) {
@@ -143,6 +138,30 @@ impl Store {
}
}
impl Store {
pub async fn filter(
&self,
account_id: u32,
collection: u8,
filters: Vec<Filter>,
) -> crate::Result<ResultSet> {
#[cfg(feature = "is_async")]
{
self.read_transaction()
.await?
.filter(account_id, collection, filters)
.await
}
#[cfg(feature = "is_sync")]
{
let mut trx = self.read_transaction()?;
self.spawn_worker(move || trx.filter(account_id, collection, filters))
.await
}
}
}
impl Filter {
#[inline(always)]
pub fn apply(

51
src/query/get.rs Normal file
View File

@@ -0,0 +1,51 @@
use crate::{Deserialize, Store, ValueKey};
impl Store {
pub async fn get_value<U>(&self, key: ValueKey) -> crate::Result<Option<U>>
where
U: Deserialize + 'static,
{
#[cfg(feature = "is_async")]
{
self.read_transaction().await?.get_value(key).await
}
#[cfg(feature = "is_sync")]
{
let trx = self.read_transaction()?;
self.spawn_worker(move || trx.get_value(key)).await
}
}
pub async fn get_values<U>(&self, key: Vec<ValueKey>) -> crate::Result<Vec<Option<U>>>
where
U: Deserialize + 'static,
{
#[cfg(feature = "is_async")]
{
let mut trx = self.read_transaction().await?;
let mut results = Vec::with_capacity(key.len());
for key in key {
trx.refresh_if_old().await?;
results.push(trx.get_value(key).await?);
}
Ok(results)
}
#[cfg(feature = "is_sync")]
{
let trx = self.read_transaction()?;
self.spawn_worker(move || {
let mut results = Vec::with_capacity(key.len());
for key in key {
results.push(trx.get_value(key)?);
}
Ok(results)
})
.await
}
}
}

16
src/query/mod.rs
View File

@@ -1,4 +1,5 @@
pub mod filter;
pub mod get;
pub mod log;
pub mod sort;
@@ -6,7 +7,7 @@ use roaring::RoaringBitmap;
use crate::{
fts::{lang::LanguageDetector, Language},
Serialize,
BitmapKey, Serialize, BM_DOCUMENT_IDS,
};
#[derive(Debug, Clone, Copy)]
@@ -185,3 +186,16 @@ impl Comparator {
}
}
}
impl BitmapKey<&'static [u8]> {
pub fn new_document_ids(account_id: u32, collection: u8) -> Self {
BitmapKey {
account_id,
collection,
family: BM_DOCUMENT_IDS,
field: u8::MAX,
key: b"",
block_num: 0,
}
}
}

100
src/query/sort.rs
View File

@@ -1,12 +1,24 @@
use ahash::AHashMap;
use crate::Store;
use crate::{ReadTransaction, Store};
use super::{Comparator, ResultSet, SortedResultRet};
impl Store {
pub struct Pagination {
requested_position: i32,
position: i32,
limit: usize,
anchor: u32,
anchor_offset: i32,
has_anchor: bool,
anchor_found: bool,
ids: Vec<u32>,
}
impl ReadTransaction<'_> {
#[maybe_async::maybe_async]
pub async fn sort(
&self,
&mut self,
result_set: ResultSet,
mut comparators: Vec<Comparator>,
limit: usize,
@@ -14,27 +26,14 @@ impl Store {
anchor: Option<u32>,
anchor_offset: i32,
) -> crate::Result<SortedResultRet> {
let limit = match (result_set.results.len(), limit) {
(0, _) => {
return Ok(SortedResultRet {
position,
ids: vec![],
found_anchor: true,
});
}
(_, 0) => result_set.results.len() as usize,
(a, b) => std::cmp::min(a as usize, b),
};
let mut paginate = Pagination::new(limit, position, anchor, anchor_offset);
if comparators.len() == 1 {
match comparators.pop().unwrap() {
Comparator::Field { field, ascending } => {
let trx = self.read_transaction().await?;
let mut results = result_set.results;
trx.sort_index(
self.sort_index(
result_set.account_id,
result_set.collection,
field,
@@ -70,7 +69,6 @@ impl Store {
}
}
} else {
let mut trx = self.read_transaction().await?;
let mut sorted_ids = AHashMap::with_capacity(paginate.limit);
for (pos, comparator) in comparators.into_iter().take(4).enumerate() {
@@ -81,8 +79,8 @@ impl Store {
let mut has_grouped_ids = false;
let mut idx = 0;
trx.refresh_if_old().await?;
trx.sort_index(
self.refresh_if_old().await?;
self.sort_index(
result_set.account_id,
result_set.collection,
field,
@@ -153,15 +151,59 @@ impl Store {
}
}
pub struct Pagination {
requested_position: i32,
position: i32,
limit: usize,
anchor: u32,
anchor_offset: i32,
has_anchor: bool,
anchor_found: bool,
ids: Vec<u32>,
impl Store {
pub async fn sort(
&self,
result_set: ResultSet,
comparators: Vec<Comparator>,
limit: usize,
position: i32,
anchor: Option<u32>,
anchor_offset: i32,
) -> crate::Result<SortedResultRet> {
let limit = match (result_set.results.len(), limit) {
(0, _) => {
return Ok(SortedResultRet {
position,
ids: vec![],
found_anchor: true,
});
}
(_, 0) => result_set.results.len() as usize,
(a, b) => std::cmp::min(a as usize, b),
};
#[cfg(feature = "is_async")]
{
self.read_transaction()
.await?
.sort(
result_set,
comparators,
limit,
position,
anchor,
anchor_offset,
)
.await
}
#[cfg(feature = "is_sync")]
{
let mut trx = self.read_transaction()?;
self.spawn_worker(move || {
trx.sort(
result_set,
comparators,
limit,
position,
anchor,
anchor_offset,
)
})
.await
}
}
}
impl Pagination {

6
src/tests/mod.rs
View File

@@ -8,13 +8,13 @@ use super::*;
#[tokio::test]
pub async fn store_test() {
let db = Arc::new(Store::open().await.unwrap());
let insert = true;
let insert = false;
if insert {
db.destroy().await;
}
assign_id::test(db).await;
//assign_id::test(db).await;
//query::test(db, insert).await;
query::test(db, insert).await;
}
pub fn deflate_artwork_data() -> Vec<u8> {

82
src/tests/query.rs
View File

@@ -178,6 +178,7 @@ pub async fn test(db: Arc<Store>, do_insert: bool) {
let mut chunk = Vec::new();
for batch in batches {
let chunk_instance = Instant::now();
chunk.push({
let db = db.clone();
tokio::spawn(async move { db.write(batch).await })
@@ -186,6 +187,10 @@ pub async fn test(db: Arc<Store>, do_insert: bool) {
for handle in chunk {
handle.await.unwrap().unwrap();
}
println!(
"Chunk insert took {} ms.",
chunk_instance.elapsed().as_millis()
);
chunk = Vec::new();
}
}
@@ -315,7 +320,6 @@ pub async fn test_filter(db: Arc<Store>) {
for (filter, expected_results) in tests {
//println!("Running test: {:?}", filter);
let mut results: Vec<String> = Vec::with_capacity(expected_results.len());
let docset = db.filter(0, COLLECTION_ID, filter).await.unwrap();
let sorted_docset = db
.sort(
@@ -329,22 +333,27 @@ pub async fn test_filter(db: Arc<Store>) {
.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(),
);
}
assert_eq!(results, expected_results);
assert_eq!(
db.get_values::<String>(
sorted_docset
.ids
.into_iter()
.map(|document_id| ValueKey {
account_id: 0,
collection: COLLECTION_ID,
document_id,
family: 0,
field: fields["accession_number"],
})
.collect()
)
.await
.unwrap()
.into_iter()
.flatten()
.collect::<Vec<_>>(),
expected_results
);
}
}
@@ -409,29 +418,32 @@ pub async fn test_sort(db: Arc<Store>) {
for (filter, sort, expected_results) in tests {
//println!("Running test: {:?}", sort);
let mut results: Vec<String> = Vec::with_capacity(expected_results.len());
let docset = db.filter(0, COLLECTION_ID, filter).await.unwrap();
let sorted_docset = db
.sort(docset, sort, expected_results.len(), 0, None, 0)
.await
.unwrap();
let mut db = db.read_transaction().await.unwrap();
for document_id in sorted_docset.ids {
db.refresh_if_old().await.unwrap();
results.push(
db.get_value(ValueKey {
account_id: 0,
collection: COLLECTION_ID,
document_id,
family: 0,
field: fields["accession_number"],
})
.await
.unwrap()
.unwrap(),
);
}
assert_eq!(results, expected_results);
assert_eq!(
db.get_values::<String>(
sorted_docset
.ids
.into_iter()
.map(|document_id| ValueKey {
account_id: 0,
collection: COLLECTION_ID,
document_id,
family: 0,
field: fields["accession_number"],
})
.collect()
)
.await
.unwrap()
.into_iter()
.flatten()
.collect::<Vec<_>>(),
expected_results
);
}
}