From 4d192de2fee0a1d342072e717d8918631f2b6218 Mon Sep 17 00:00:00 2001 From: Mauro D Date: Sun, 2 Apr 2023 17:10:59 +0000 Subject: [PATCH] SQLite support --- Cargo.toml | 11 +- coco.txt | 0 src/backend/foundationdb/bitmap.rs | 2 +- src/backend/foundationdb/read.rs | 40 +---- src/backend/foundationdb/write.rs | 2 +- src/backend/mod.rs | 2 + src/backend/sqlite/main.rs | 92 +++++++++++ src/backend/sqlite/mod.rs | 126 +++++++++++++++ src/backend/sqlite/pool.rs | 105 ++++++++++++ src/backend/sqlite/read.rs | 247 +++++++++++++++++++++++++++++ src/backend/sqlite/write.rs | 242 ++++++++++++++++++++++++++++ src/fts/query.rs | 4 +- src/lib.rs | 33 +++- src/query/filter.rs | 57 ++++--- src/query/get.rs | 51 ++++++ src/query/mod.rs | 16 +- src/query/sort.rs | 100 ++++++++---- src/tests/mod.rs | 6 +- src/tests/query.rs | 82 ++++++---- 19 files changed, 1086 insertions(+), 132 deletions(-) delete mode 100644 coco.txt create mode 100644 src/backend/sqlite/main.rs create mode 100644 src/backend/sqlite/mod.rs create mode 100644 src/backend/sqlite/pool.rs create mode 100644 src/backend/sqlite/read.rs create mode 100644 src/backend/sqlite/write.rs create mode 100644 src/query/get.rs diff --git a/Cargo.toml b/Cargo.toml index 686c6c44..3cf2bda8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,9 @@ edition = "2021" utils = { path = "../utils" } rocksdb = { version = "0.20.1", optional = true } foundationdb = { version = "0.7.0", optional = true } +rusqlite = { version = "0.29.0", features = ["bundled"], optional = true } +tokio = { version = "1.23", features = ["sync"], optional = true } +r2d2 = { version = "0.8.10", optional = true } futures = { version = "0.3", optional = true } rand = "0.8.5" roaring = "0.10.1" @@ -23,11 +26,15 @@ jieba-rs = "0.6" # Chinese stemmer xxhash-rust = { version = "0.8.5", features = ["xxh3"] } farmhash = "1.1.5" siphasher = "0.3" +maybe-async = "0.2" [features] default = ["foundation"] -rocks = ["rocksdb", "rayon"] -foundation = ["foundationdb", "futures"] +rocks = ["rocksdb", "rayon", "is_sync"] +sqlite = ["rusqlite", "rayon", "r2d2", "tokio", "is_sync"] +foundation = ["foundationdb", "futures", "is_async"] +is_sync = ["maybe-async/is_sync"] +is_async = [] [dev-dependencies] tokio = { version = "1.23", features = ["full"] } diff --git a/coco.txt b/coco.txt deleted file mode 100644 index e69de29b..00000000 diff --git a/src/backend/foundationdb/bitmap.rs b/src/backend/foundationdb/bitmap.rs index bfb19f43..defadbe5 100644 --- a/src/backend/foundationdb/bitmap.rs +++ b/src/backend/foundationdb/bitmap.rs @@ -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::(); const WORDS_PER_BLOCK: u32 = 8; pub const BITS_PER_BLOCK: u32 = WORD_SIZE_BITS * WORDS_PER_BLOCK; diff --git a/src/backend/foundationdb/read.rs b/src/backend/foundationdb/read.rs index 33662e01..f78f22b7 100644 --- a/src/backend/foundationdb/read.rs +++ b/src/backend/foundationdb/read.rs @@ -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(&self, key: ValueKey) -> crate::Result> @@ -39,36 +33,6 @@ impl ReadTransaction<'_> { } } - #[inline(always)] - pub async fn get_values(&self, keys: Vec) -> crate::Result>> - where - U: Deserialize, - { - let mut results = Vec::with_capacity(keys.len()); - - for key in keys { - results.push(self.get_value(key).await?); - } - - Ok(results) - } - - pub async fn get_document_ids( - &self, - account_id: u32, - collection: u8, - ) -> crate::Result> { - self.get_bitmap(BitmapKey { - account_id, - collection, - family: BM_DOCUMENT_IDS, - field: u8::MAX, - key: b"", - block_num: 0, - }) - .await - } - async fn get_bitmap_>( &self, mut key: BitmapKey, diff --git a/src/backend/foundationdb/write.rs b/src/backend/foundationdb/write.rs index ca9ff891..f2cef164 100644 --- a/src/backend/foundationdb/write.rs +++ b/src/backend/foundationdb/write.rs @@ -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 { diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 9808078e..45333889 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -2,3 +2,5 @@ pub mod foundationdb; #[cfg(feature = "rocks")] pub mod rocksdb; +#[cfg(feature = "sqlite")] +pub mod sqlite; diff --git a/src/backend/sqlite/main.rs b/src/backend/sqlite/main.rs new file mode 100644 index 00000000..207d0327 --- /dev/null +++ b/src/backend/sqlite/main.rs @@ -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 { + 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(&self, f: U) -> crate::Result + where + U: FnOnce() -> crate::Result + 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 + ))), + } + } +} diff --git a/src/backend/sqlite/mod.rs b/src/backend/sqlite/mod.rs new file mode 100644 index 00000000..e19bb204 --- /dev/null +++ b/src/backend/sqlite/mod.rs @@ -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::(); +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> Serialize for &IndexKey { + fn serialize(self) -> Vec { + let key = self.key.as_ref(); + KeySerializer::new(std::mem::size_of::>() + 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 { + KeySerializer::new(std::mem::size_of::() + 1) + .write(self.account_id) + .write(self.collection) + .write(self.field) + .finalize() + } +} + +impl Serialize for &ValueKey { + fn serialize(self) -> Vec { + if self.family == 0 { + KeySerializer::new(std::mem::size_of::() + 1) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(self.field) + .finalize() + } else { + KeySerializer::new(std::mem::size_of::() + 2) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .write(u8::MAX) + .write(self.family) + .write(self.field) + .finalize() + } + } +} + +impl> Serialize for &BitmapKey { + fn serialize(self) -> Vec { + let key = self.key.as_ref(); + KeySerializer::new(std::mem::size_of::>() + key.len() + 1) + .write(self.account_id) + .write(self.collection) + .write(self.family) + .write(self.field) + .write(key) + .write(self.block_num) + .finalize() + } +} + +impl> Serialize for &BlobKey { + fn serialize(self) -> Vec { + let hash = self.hash.as_ref(); + KeySerializer::new(std::mem::size_of::>() + hash.len() + 1) + .write(hash) + .write_leb128(self.account_id) + .write(self.collection) + .write_leb128(self.document_id) + .finalize() + } +} + +impl Serialize for &AclKey { + fn serialize(self) -> Vec { + KeySerializer::new(std::mem::size_of::() + 1) + .write_leb128(self.grant_account_id) + .write(u8::MAX) + .write_leb128(self.to_account_id) + .write(self.to_collection) + .write_leb128(self.to_document_id) + .finalize() + } +} + +impl Serialize for &LogKey { + fn serialize(self) -> Vec { + KeySerializer::new(std::mem::size_of::() + 1) + .write(self.account_id) + .write(self.collection) + .write(self.change_id) + .finalize() + } +} + +impl From for crate::Error { + fn from(err: r2d2::Error) -> Self { + Self::InternalError(format!("Connection pool error: {}", err)) + } +} + +impl From for crate::Error { + fn from(err: rusqlite::Error) -> Self { + Self::InternalError(format!("SQLite error: {}", err)) + } +} + +impl From for crate::Error { + fn from(err: rusqlite::types::FromSqlError) -> Self { + Self::InternalError(format!("SQLite error: {}", err)) + } +} diff --git a/src/backend/sqlite/pool.rs b/src/backend/sqlite/pool.rs new file mode 100644 index 00000000..1064c7d4 --- /dev/null +++ b/src/backend/sqlite/pool.rs @@ -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>, +} + +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>(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(self, init: F) -> Self + where + F: Fn(&mut Connection) -> Result<(), rusqlite::Error> + Send + Sync + 'static, + { + let init: Option> = Some(Box::new(init)); + Self { init, ..self } + } +} + +impl r2d2::ManageConnection for SqliteConnectionManager { + type Connection = Connection; + type Error = rusqlite::Error; + + fn connect(&self) -> Result { + 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 + } +} diff --git a/src/backend/sqlite/read.rs b/src/backend/sqlite/read.rs new file mode 100644 index 00000000..3c39556a --- /dev/null +++ b/src/backend/sqlite/read.rs @@ -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(&self, key: ValueKey) -> crate::Result> + 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_>( + &self, + mut key: BitmapKey, + 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::())?; + + 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>( + &self, + key: BitmapKey, + ) -> crate::Result> { + 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>( + &self, + keys: Vec>, + ) -> crate::Result> { + let mut result: Option = 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>( + &self, + keys: Vec>, + ) -> crate::Result> { + 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, + op: Operator, + ) -> crate::Result> { + let k1 = KeySerializer::new( + std::mem::size_of::>() + value.len() + 1 + std::mem::size_of::(), + ) + .write(account_id) + .write(collection) + .write(field); + let k2 = KeySerializer::new( + std::mem::size_of::>() + value.len() + 1 + std::mem::size_of::(), + ) + .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::())?); + } + + 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::(); + 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> { + Ok(ReadTransaction { + conn: self.conn_pool.get()?, + _p: std::marker::PhantomData, + }) + } +} diff --git a/src/backend/sqlite/write.rs b/src/backend/sqlite/write.rs new file mode 100644 index 00000000..dc304401 --- /dev/null +++ b/src/backend/sqlite/write.rs @@ -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 { + todo!() + } + + pub async fn assign_change_id(&self, account_id: u32, collection: u8) -> crate::Result { + 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(); + } +} diff --git a/src/fts/query.rs b/src/fts/query.rs index 9d57c85a..3140424c 100644 --- a/src/fts/query.rs +++ b/src/fts/query.rs @@ -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, diff --git a/src/lib.rs b/src/lib.rs index 28a4e065..2ab4c47a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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, + worker_pool: rayon::ThreadPool, +} + +#[cfg(feature = "sqlite")] +pub struct ReadTransaction<'x> { + conn: r2d2::PooledConnection, + _p: std::marker::PhantomData<&'x ()>, +} + pub trait Deserialize: Sized + Sync + Send { fn deserialize(bytes: &[u8]) -> crate::Result; } @@ -32,7 +53,6 @@ pub struct BitmapKey> { 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; diff --git a/src/query/filter.rs b/src/query/filter.rs index 4b514b93..26a80a62 100644 --- a/src/query/filter.rs +++ b/src/query/filter.rs @@ -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, } -impl Store { +impl ReadTransaction<'_> { + #[maybe_async::maybe_async] pub async fn filter( - &self, + &mut self, account_id: u32, collection: u8, filters: Vec, ) -> crate::Result { - 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, ¬_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, + ) -> crate::Result { + #[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( diff --git a/src/query/get.rs b/src/query/get.rs new file mode 100644 index 00000000..a8bd4a59 --- /dev/null +++ b/src/query/get.rs @@ -0,0 +1,51 @@ +use crate::{Deserialize, Store, ValueKey}; + +impl Store { + pub async fn get_value(&self, key: ValueKey) -> crate::Result> + 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(&self, key: Vec) -> crate::Result>> + 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 + } + } +} diff --git a/src/query/mod.rs b/src/query/mod.rs index 09571c00..4072f3f1 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -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, + } + } +} diff --git a/src/query/sort.rs b/src/query/sort.rs index 5563e08b..2eb13cd0 100644 --- a/src/query/sort.rs +++ b/src/query/sort.rs @@ -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, +} + +impl ReadTransaction<'_> { + #[maybe_async::maybe_async] pub async fn sort( - &self, + &mut self, result_set: ResultSet, mut comparators: Vec, limit: usize, @@ -14,27 +26,14 @@ impl Store { anchor: Option, anchor_offset: i32, ) -> crate::Result { - 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, +impl Store { + pub async fn sort( + &self, + result_set: ResultSet, + comparators: Vec, + limit: usize, + position: i32, + anchor: Option, + anchor_offset: i32, + ) -> crate::Result { + 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 { diff --git a/src/tests/mod.rs b/src/tests/mod.rs index cb733dae..eb712e1f 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -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 { diff --git a/src/tests/query.rs b/src/tests/query.rs index 1c6851de..1283a697 100644 --- a/src/tests/query.rs +++ b/src/tests/query.rs @@ -178,6 +178,7 @@ pub async fn test(db: Arc, 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, 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) { for (filter, expected_results) in tests { //println!("Running test: {:?}", filter); - let mut results: Vec = 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) { .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::( + 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::>(), + expected_results + ); } } @@ -409,29 +418,32 @@ pub async fn test_sort(db: Arc) { for (filter, sort, expected_results) in tests { //println!("Running test: {:?}", sort); - let mut results: Vec = 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::( + 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::>(), + expected_results + ); } }