diff --git a/Cargo.toml b/Cargo.toml index 3cf2bda8..1563d8e5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,13 +27,16 @@ xxhash-rust = { version = "0.8.5", features = ["xxh3"] } farmhash = "1.1.5" siphasher = "0.3" maybe-async = "0.2" +parking_lot = { version = "0.12.1", optional = true } +lru-cache = { version = "0.1.2", optional = true } +blake3 = "1.3.3" [features] -default = ["foundation"] +default = ["sqlite"] rocks = ["rocksdb", "rayon", "is_sync"] sqlite = ["rusqlite", "rayon", "r2d2", "tokio", "is_sync"] foundation = ["foundationdb", "futures", "is_async"] -is_sync = ["maybe-async/is_sync"] +is_sync = ["maybe-async/is_sync", "parking_lot", "lru-cache"] is_async = [] [dev-dependencies] diff --git a/coco.sh b/coco.sh deleted file mode 100644 index a5d40801..00000000 --- a/coco.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash -while true; do - cargo test store_test -- --nocapture - exit_code=$? - if [ $exit_code -ne 0 ]; then - break - fi -done diff --git a/src/backend/foundationdb/main.rs b/src/backend/foundationdb/main.rs index b4114a75..bf83468e 100644 --- a/src/backend/foundationdb/main.rs +++ b/src/backend/foundationdb/main.rs @@ -3,19 +3,11 @@ use foundationdb::Database; use crate::Store; impl Store { - pub async fn open() -> crate::Result { + pub async fn open(config: &Config) -> crate::Result { Ok(Self { guard: unsafe { foundationdb::boot() }, db: Database::default()?, + blob: BlobStore::new(config).await?, }) } } - -/* -impl Drop for Store { - fn drop(&mut self) { - self.guard.drop(); - self.db.drop(); - } -} -*/ diff --git a/src/backend/foundationdb/mod.rs b/src/backend/foundationdb/mod.rs index 49b5b565..adcd7c51 100644 --- a/src/backend/foundationdb/mod.rs +++ b/src/backend/foundationdb/mod.rs @@ -10,13 +10,6 @@ pub mod main; pub mod read; pub mod write; -pub const SUBSPACE_BITMAPS: u8 = b'b'; -pub const SUBSPACE_VALUES: u8 = b'v'; -pub const SUBSPACE_LOGS: u8 = b'l'; -pub const SUBSPACE_BLOBS: u8 = b'o'; -pub const SUBSPACE_INDEXES: u8 = b'i'; -pub const SUBSPACE_ACLS: u8 = b'c'; - impl> Serialize for &IndexKey { fn serialize(self) -> Vec { let key = self.key.as_ref(); @@ -84,7 +77,7 @@ impl> Serialize for &BitmapKey { impl> Serialize for &BlobKey { fn serialize(self) -> Vec { let hash = self.hash.as_ref(); - KeySerializer::new(std::mem::size_of::>() + hash.len() + 1) + KeySerializer::new(std::mem::size_of::>() + BLOB_HASH_LEN + 1) .write(SUBSPACE_BLOBS) .write(hash) .write_leb128(self.account_id) diff --git a/src/backend/foundationdb/write.rs b/src/backend/foundationdb/write.rs index f2cef164..824671cf 100644 --- a/src/backend/foundationdb/write.rs +++ b/src/backend/foundationdb/write.rs @@ -128,7 +128,15 @@ impl Store { } .serialize(); if *set { - trx.set(&key, &[]); + let now_; + let value = if document_id != u32::MAX { + &[] + } else { + now_ = now().to_be_bytes(); + &now_[..] + }; + + trx.set(&key, value); } else { trx.clear(&key); } @@ -370,10 +378,3 @@ impl Store { trx.commit().await.unwrap(); } } - -#[inline(always)] -fn now() -> u64 { - SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .map_or(0, |d| d.as_secs()) -} diff --git a/src/backend/sqlite/id_assign.rs b/src/backend/sqlite/id_assign.rs new file mode 100644 index 00000000..591ce8ec --- /dev/null +++ b/src/backend/sqlite/id_assign.rs @@ -0,0 +1,151 @@ +/* + * Copyright (c) 2020-2022, Stalwart Labs Ltd. + * + * This file is part of the Stalwart JMAP Server. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * in the LICENSE file at the top-level directory of this distribution. + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + * + * You can be released from the requirements of the AGPLv3 license by + * purchasing a commercial license. Please contact licensing@stalw.art + * for more details. +*/ + +use roaring::RoaringBitmap; + +use crate::{BitmapKey, Store}; + +#[derive(Clone, Copy, Hash, PartialEq, Eq)] +pub struct IdCacheKey { + pub account_id: u32, + pub collection: u8, +} + +impl IdCacheKey { + pub fn new(account_id: u32, collection: impl Into) -> Self { + Self { + account_id, + collection: collection.into(), + } + } +} + +#[derive(Clone)] +pub struct IdAssigner { + pub available_document_ids: RoaringBitmap, + pub next_change_id: u64, +} + +impl IdAssigner { + pub fn new(used_ids: Option, next_change_id: u64) -> Self { + let mut assigner = IdAssigner { + available_document_ids: RoaringBitmap::full(), + next_change_id, + }; + + if let Some(used_ids) = used_ids { + assigner.available_document_ids ^= &used_ids; + } + + assigner + } + + pub fn assign_document_id(&mut self) -> u32 { + let id = self.available_document_ids.min().unwrap(); + self.available_document_ids.remove(id); + id + } + + pub fn assign_change_id(&mut self) -> u64 { + let id = self.next_change_id; + self.next_change_id += 1; + id + } +} + +impl Store { + pub async fn assign_document_id(&self, account_id: u32, collection: u8) -> crate::Result { + let key = IdCacheKey::new(account_id, collection); + for _ in 0..2 { + if let Some(assigner) = self.id_assigner.lock().get_mut(&key) { + return Ok(assigner.assign_document_id()); + } + self.build_id_assigner(key).await?; + } + + unreachable!() + } + + pub async fn assign_change_id(&self, account_id: u32, collection: u8) -> crate::Result { + let key = IdCacheKey::new(account_id, collection); + for _ in 0..2 { + if let Some(assigner) = self.id_assigner.lock().get_mut(&key) { + return Ok(assigner.assign_change_id()); + } + self.build_id_assigner(key).await?; + } + + unreachable!() + } + + async fn build_id_assigner(&self, key: IdCacheKey) -> crate::Result<()> { + let conn = self.read_transaction()?; + let id_assigner = self.id_assigner.clone(); + self.spawn_worker(move || { + let mut id_assigner = id_assigner.lock(); + // Make sure id assigner was not added by another thread + if id_assigner.get_mut(&key).is_some() { + return Ok(()); + } + + // Obtain used ids + let used_ids = + conn.get_bitmap(BitmapKey::new_document_ids(key.account_id, key.collection))?; + let next_change_id = conn + .get_last_change_id(key.account_id, key.collection)? + .map(|id| id + 1) + .unwrap_or(0); + id_assigner.insert(key, IdAssigner::new(used_ids, next_change_id)); + + Ok(()) + }) + .await + } +} + +#[cfg(test)] +mod tests { + use roaring::RoaringBitmap; + + use super::IdAssigner; + + #[test] + fn id_assigner() { + let mut assigner = IdAssigner::new(None, 0); + assert_eq!(assigner.assign_document_id(), 0); + assert_eq!(assigner.assign_document_id(), 1); + assert_eq!(assigner.assign_document_id(), 2); + + let mut assigner = IdAssigner::new( + RoaringBitmap::from_sorted_iter([0, 2, 4, 6]) + .unwrap() + .into(), + 0, + ); + assert_eq!(assigner.assign_document_id(), 1); + assert_eq!(assigner.assign_document_id(), 3); + assert_eq!(assigner.assign_document_id(), 5); + assert_eq!(assigner.assign_document_id(), 7); + assert_eq!(assigner.assign_document_id(), 8); + } +} diff --git a/src/backend/sqlite/main.rs b/src/backend/sqlite/main.rs index 207d0327..d0cd24b4 100644 --- a/src/backend/sqlite/main.rs +++ b/src/backend/sqlite/main.rs @@ -1,22 +1,38 @@ +use std::sync::Arc; + +use lru_cache::LruCache; +use parking_lot::Mutex; use r2d2::Pool; use tokio::sync::oneshot; +use utils::config::Config; -use crate::Store; +use crate::{ + blob::BlobStore, Store, SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_INDEXES, + SUBSPACE_LOGS, SUBSPACE_VALUES, +}; use super::pool::SqliteConnectionManager; impl Store { // TODO configure rayon thread pool // TODO configure r2d2 pool - pub async fn open() -> crate::Result { + // TODO configure id assigner size + pub async fn open(config: &Config) -> crate::Result { let db = Self { conn_pool: Pool::new( - SqliteConnectionManager::file("/tmp/sqlite.db") - .with_init(|c| c.execute_batch("PRAGMA journal_mode=WAL;")), + SqliteConnectionManager::file("/tmp/sqlite.db").with_init(|c| { + c.execute_batch(concat!( + "PRAGMA journal_mode = WAL; ", + "PRAGMA synchronous = normal; ", + "PRAGMA temp_store = memory;" + )) + }), )?, worker_pool: rayon::ThreadPoolBuilder::new().build().map_err(|err| { crate::Error::InternalError(format!("Failed to build worker pool: {}", err)) })?, + id_assigner: Arc::new(Mutex::new(LruCache::new(1000))), + blob: BlobStore::new(config).await?, }; db.create_tables()?; Ok(db) @@ -25,7 +41,13 @@ impl Store { pub(super) fn create_tables(&self) -> crate::Result<()> { let conn = self.conn_pool.get()?; - for table in ["v", "l", "o", "c"] { + for table in [ + SUBSPACE_VALUES, + SUBSPACE_LOGS, + SUBSPACE_BLOBS, + SUBSPACE_ACLS, + ] { + let table = char::from(table); conn.execute( &format!( "CREATE TABLE IF NOT EXISTS {table} ( @@ -38,14 +60,18 @@ impl Store { } conn.execute( - "CREATE TABLE IF NOT EXISTS i ( + &format!( + "CREATE TABLE IF NOT EXISTS {} ( k BLOB PRIMARY KEY )", + char::from(SUBSPACE_INDEXES) + ), [], )?; conn.execute( - "CREATE TABLE IF NOT EXISTS b ( + &format!( + "CREATE TABLE IF NOT EXISTS {} ( z BLOB PRIMARY KEY, a INTEGER NOT NULL DEFAULT 0, b INTEGER NOT NULL DEFAULT 0, @@ -64,6 +90,8 @@ impl Store { o INTEGER NOT NULL DEFAULT 0, p INTEGER NOT NULL DEFAULT 0 )", + char::from(SUBSPACE_BITMAPS) + ), [], )?; diff --git a/src/backend/sqlite/mod.rs b/src/backend/sqlite/mod.rs index e19bb204..35c4f17d 100644 --- a/src/backend/sqlite/mod.rs +++ b/src/backend/sqlite/mod.rs @@ -1,8 +1,9 @@ use crate::{ write::key::KeySerializer, AclKey, BitmapKey, BlobKey, IndexKey, IndexKeyPrefix, LogKey, - Serialize, ValueKey, + Serialize, ValueKey, BLOB_HASH_LEN, }; +pub mod id_assign; pub mod main; pub mod pool; pub mod read; @@ -76,7 +77,7 @@ impl> Serialize for &BitmapKey { impl> Serialize for &BlobKey { fn serialize(self) -> Vec { let hash = self.hash.as_ref(); - KeySerializer::new(std::mem::size_of::>() + hash.len() + 1) + KeySerializer::new(std::mem::size_of::>() + BLOB_HASH_LEN + 1) .write(hash) .write_leb128(self.account_id) .write(self.collection) @@ -85,6 +86,12 @@ impl> Serialize for &BlobKey { } } +impl> Serialize for BlobKey { + fn serialize(self) -> Vec { + (&self).serialize() + } +} + impl Serialize for &AclKey { fn serialize(self) -> Vec { KeySerializer::new(std::mem::size_of::() + 1) diff --git a/src/backend/sqlite/read.rs b/src/backend/sqlite/read.rs index 3c39556a..404e719c 100644 --- a/src/backend/sqlite/read.rs +++ b/src/backend/sqlite/read.rs @@ -6,7 +6,8 @@ use rusqlite::OptionalExtension; use crate::{ query::Operator, write::key::{DeserializeBigEndian, KeySerializer}, - BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, ReadTransaction, Serialize, Store, ValueKey, + BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, Key, LogKey, ReadTransaction, Serialize, + Store, ValueKey, }; use super::{BITS_PER_BLOCK, WORDS_PER_BLOCK, WORD_SIZE_BITS}; @@ -230,6 +231,73 @@ impl ReadTransaction<'_> { Ok(()) } + #[maybe_async::maybe_async] + pub(crate) async fn iterate( + &self, + mut acc: T, + begin: impl Key, + end: impl Key, + first: bool, + ascending: bool, + cb: impl Fn(&mut T, &[u8], &[u8]) -> crate::Result + Sync + Send + 'static, + ) -> crate::Result { + let table = char::from(begin.subspace()); + let begin = begin.serialize(); + let end = end.serialize(); + + let mut query = self.conn.prepare_cached(&match (first, ascending) { + (true, true) => { + format!("SELECT k, v FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC LIMIT 1") + } + (true, false) => { + format!("SELECT k, v FROM {table} WHERE k >= ? AND k <= ? ORDER BY k DESC LIMIT 1") + } + (false, true) => { + format!("SELECT k, v FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC") + } + (false, false) => { + format!("SELECT k, v FROM {table} 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 value = row.get_ref(1)?.as_bytes()?; + + if !cb(&mut acc, key, value)? { + return Ok(acc); + } + } + + Ok(acc) + } + + #[maybe_async::maybe_async] + pub(crate) async fn get_last_change_id( + &self, + account_id: u32, + collection: u8, + ) -> crate::Result> { + let key = LogKey { + account_id, + collection, + change_id: u64::MAX, + } + .serialize(); + + self.conn + .prepare_cached("SELECT k FROM l WHERE k < ? ORDER BY k DESC LIMIT 1")? + .query_row([&key], |row| { + let key = row.get_ref(0)?.as_bytes()?; + + key.deserialize_be_u64(key.len() - std::mem::size_of::()) + .map_err(|err| rusqlite::Error::ToSqlConversionFailure(err.into())) + }) + .optional() + .map_err(Into::into) + } + #[maybe_async::maybe_async] pub async fn refresh_if_old(&mut self) -> crate::Result<()> { Ok(()) diff --git a/src/backend/sqlite/write.rs b/src/backend/sqlite/write.rs index dc304401..f7b04d16 100644 --- a/src/backend/sqlite/write.rs +++ b/src/backend/sqlite/write.rs @@ -1,7 +1,7 @@ use rusqlite::params; use crate::{ - write::{Batch, Operation}, + write::{now, Batch, Operation}, AclKey, BitmapKey, BlobKey, IndexKey, LogKey, Serialize, Store, ValueKey, }; @@ -173,10 +173,17 @@ impl Store { .serialize(); if *set { - trx.prepare_cached("INSERT OR REPLACE INTO b (k, v) VALUES (?, ?)")? - .execute([&key[..], &[]])?; + let now_; + let value = if document_id != u32::MAX { + &[] + } else { + now_ = now().to_be_bytes(); + &now_[..] + }; + trx.prepare_cached("INSERT OR REPLACE INTO o (k, v) VALUES (?, ?)")? + .execute([&key[..], value])?; } else { - trx.prepare_cached("DELETE FROM b WHERE k = ?")? + trx.prepare_cached("DELETE FROM o WHERE k = ?")? .execute([&key])?; } } @@ -223,19 +230,24 @@ impl Store { .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) { + use crate::{ + SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_INDEXES, SUBSPACE_LOGS, + SUBSPACE_VALUES, + }; + let conn = self.conn_pool.get().unwrap(); - for table in ["v", "l", "o", "c", "i", "b"] { - conn.execute(&format!("DROP TABLE {table}"), []).unwrap(); + for table in [ + SUBSPACE_VALUES, + SUBSPACE_LOGS, + SUBSPACE_BLOBS, + SUBSPACE_ACLS, + SUBSPACE_BITMAPS, + SUBSPACE_INDEXES, + ] { + conn.execute(&format!("DROP TABLE {}", char::from(table)), []) + .unwrap(); } self.create_tables().unwrap(); } diff --git a/src/blob/mod.rs b/src/blob/mod.rs new file mode 100644 index 00000000..e363db1b --- /dev/null +++ b/src/blob/mod.rs @@ -0,0 +1,56 @@ +pub mod purge; +pub mod read; +pub mod write; + +use std::{ + io::Write, + path::{Path, PathBuf}, +}; + +use utils::{codec::base32_custom::Base32Writer, config::Config}; + +use crate::{BlobId, Serialize}; + +pub enum BlobStore { + Local { + base_path: PathBuf, + hash_levels: usize, + }, + Remote(String), +} + +impl BlobStore { + pub async fn new(config: &Config) -> crate::Result { + Ok(BlobStore::Local { + base_path: config.value_require("blob.store.path")?.into(), + hash_levels: config.property("blob.store.hash")?.unwrap_or(1), + }) + } +} + +impl Serialize for &BlobId { + fn serialize(self) -> Vec { + self.hash.to_vec() + } +} + +impl From for crate::Error { + fn from(err: std::io::Error) -> Self { + Self::InternalError(format!("IO error: {}", err)) + } +} + +fn get_path(base_path: &Path, hash_levels: usize, blob_id: &BlobId) -> crate::Result { + let mut path = base_path.to_path_buf(); + let hash = &blob_id.hash; + for byte in hash.iter().take(hash_levels) { + path.push(format!("{:x}", byte)); + } + + // Base32 encode the hash + let mut writer = Base32Writer::with_capacity(hash.len()); + writer.write_all(hash).unwrap(); + path.push(&writer.finalize()); + + Ok(path) +} diff --git a/src/blob/purge.rs b/src/blob/purge.rs new file mode 100644 index 00000000..b2bfa26d --- /dev/null +++ b/src/blob/purge.rs @@ -0,0 +1,90 @@ +use utils::codec::leb128::Leb128Iterator; + +use crate::{ + write::{now, BatchBuilder, F_CLEAR}, + BlobKey, Deserialize, Store, BLOB_HASH_LEN, +}; + +struct BlobPurge { + batch: BatchBuilder, + id: [u8; BLOB_HASH_LEN], + link_count: u32, + delete: Vec<[u8; BLOB_HASH_LEN]>, +} + +impl Store { + pub async fn purge_blobs(&self, ttl: u64) -> crate::Result<()> { + let now = now(); + + let results = BlobPurge { + batch: BatchBuilder::new(), + id: [0u8; BLOB_HASH_LEN], + link_count: u32::MAX, + delete: vec![], + }; + + let from_key = BlobKey { + account_id: 0, + collection: 0, + document_id: 0, + hash: [0; BLOB_HASH_LEN], + }; + let to_key = BlobKey { + account_id: u32::MAX, + collection: u8::MAX, + document_id: u32::MAX, + hash: [u8::MAX; BLOB_HASH_LEN], + }; + + let mut results = self + .iterate(results, from_key, to_key, false, true, move |b, k, v| { + if !k.starts_with(&b.id) { + if b.link_count == 0 { + b.delete.push(b.id); + } + b.link_count = 0; + b.id.copy_from_slice(&k[..BLOB_HASH_LEN]); + } + + if !v.is_empty() { + let timestamp = u64::deserialize(v)?; + + if (now >= timestamp && now - timestamp >= ttl) + || (now < timestamp && timestamp - now >= ttl) + { + let mut iter = k[BLOB_HASH_LEN..].iter(); + if let (Some(account_id), Some(collection), Some(document_id)) = + (iter.next_leb128(), iter.next(), iter.next_leb128()) + { + b.batch + .with_account_id(account_id) + .with_collection(*collection) + .update_document(document_id) + .blob(b.id.to_vec(), F_CLEAR); + } + } else { + b.link_count += 1; + } + } else { + b.link_count += 1; + } + + Ok(true) + }) + .await?; + + if results.link_count == 0 { + results.delete.push(results.id); + } + + if !results.batch.is_empty() { + self.write(results.batch.build()).await?; + } + + for hash in results.delete { + self.blob.delete(&crate::BlobId { hash }).await?; + } + + Ok(()) + } +} diff --git a/src/blob/read.rs b/src/blob/read.rs new file mode 100644 index 00000000..345b87cf --- /dev/null +++ b/src/blob/read.rs @@ -0,0 +1,52 @@ +use std::{io::SeekFrom, ops::Range}; + +use tokio::{ + fs::{self, File}, + io::{AsyncReadExt, AsyncSeekExt}, +}; + +use crate::BlobId; + +use super::{get_path, BlobStore}; + +impl BlobStore { + pub async fn get(&self, id: &BlobId, range: Range) -> crate::Result>> { + match self { + BlobStore::Local { + base_path, + hash_levels, + } => { + let blob_path = get_path(base_path, *hash_levels, id)?; + let blob_size = match fs::metadata(&blob_path).await { + Ok(m) => m.len(), + Err(_) => return Ok(None), + }; + let mut blob = File::open(&blob_path).await?; + + Ok(Some(if range.start != 0 || range.end != u32::MAX { + let from_offset = if range.start < blob_size as u32 { + range.start + } else { + 0 + }; + let mut buf = vec![ + 0; + (std::cmp::min(range.end, blob_size as u32) - from_offset) + as usize + ]; + + if from_offset > 0 { + blob.seek(SeekFrom::Start(from_offset as u64)).await?; + } + blob.read_exact(&mut buf).await?; + buf + } else { + let mut buf = Vec::with_capacity(blob_size as usize); + blob.read_to_end(&mut buf).await?; + buf + })) + } + BlobStore::Remote(_) => todo!(), + } + } +} diff --git a/src/blob/write.rs b/src/blob/write.rs new file mode 100644 index 00000000..4388ac05 --- /dev/null +++ b/src/blob/write.rs @@ -0,0 +1,110 @@ +use tokio::{ + fs::{self, File}, + io::AsyncWriteExt, +}; + +use crate::{write::BatchBuilder, BlobId, BlobKey, Store, BLOB_HASH_LEN}; + +use super::{get_path, BlobStore}; + +impl Store { + pub async fn write_blob(&self, account_id: u32, data: &[u8]) -> crate::Result { + let id = BlobId::from(data); + + // Check if the blob already exists + let from_key = BlobKey { + account_id: 0, + collection: 0, + document_id: 0, + hash: [0; BLOB_HASH_LEN], + }; + let to_key = BlobKey { + account_id: u32::MAX, + collection: u8::MAX, + document_id: u32::MAX, + hash: id.hash, + }; + + let found = self + .iterate(false, from_key, to_key, true, false, |acc, _, _| { + *acc = true; + Ok(false) + }) + .await?; + + if !found { + // Write the blob + self.blob.put(&id, data).await?; + + // Write a temporary link to the blob + self.write( + BatchBuilder::new() + .with_account_id(account_id) + .with_collection(u8::MAX) + .update_document(u32::MAX) + .blob(&id, 0) + .build_batch(), + ) + .await?; + } + + Ok(id) + } +} + +impl BlobStore { + pub async fn put(&self, id: &BlobId, data: &[u8]) -> crate::Result { + match self { + BlobStore::Local { + base_path, + hash_levels, + } => { + let blob_path = get_path(base_path, *hash_levels, id)?; + + if blob_path.exists() { + let metadata = fs::metadata(&blob_path).await?; + if metadata.len() as usize == data.len() { + return Ok(false); + } + } + + fs::create_dir_all(blob_path.parent().unwrap()).await?; + let mut blob_file = File::create(&blob_path).await?; + blob_file.write_all(data).await?; + blob_file.flush().await?; + + Ok(true) + } + BlobStore::Remote(_) => todo!(), + } + } + + pub async fn delete(&self, id: &BlobId) -> crate::Result { + match self { + BlobStore::Local { + base_path, + hash_levels, + } => { + let blob_path = get_path(base_path, *hash_levels, id)?; + + if blob_path.exists() { + fs::remove_file(&blob_path).await?; + Ok(true) + } else { + Ok(false) + } + } + BlobStore::Remote(_) => todo!(), + } + } +} + +impl From<&[u8]> for BlobId { + fn from(data: &[u8]) -> Self { + let mut hasher = blake3::Hasher::new(); + hasher.update(data); + Self { + hash: hasher.finalize().into(), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 2ab4c47a..40d78cb0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,9 @@ use std::fmt::Display; +use blob::BlobStore; + pub mod backend; +pub mod blob; pub mod fts; pub mod query; pub mod write; @@ -18,6 +21,7 @@ pub struct Store { pub struct Store { db: foundationdb::Database, guard: foundationdb::api::NetworkAutoStop, + blob: BlobStore, } #[cfg(feature = "foundation")] @@ -30,7 +34,16 @@ pub struct ReadTransaction<'x> { #[cfg(feature = "sqlite")] pub struct Store { conn_pool: r2d2::Pool, + id_assigner: std::sync::Arc< + parking_lot::Mutex< + lru_cache::LruCache< + backend::sqlite::id_assign::IdCacheKey, + backend::sqlite::id_assign::IdAssigner, + >, + >, + >, worker_pool: rayon::ThreadPool, + blob: BlobStore, } #[cfg(feature = "sqlite")] @@ -47,6 +60,10 @@ pub trait Serialize { fn serialize(self) -> Vec; } +pub trait Key: Serialize + Sync + Send + 'static { + fn subspace(&self) -> u8; +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct BitmapKey> { pub account_id: u32, @@ -105,6 +122,11 @@ pub struct LogKey { pub change_id: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct BlobId { + pub hash: [u8; BLOB_HASH_LEN], +} + pub type Result = std::result::Result; #[derive(Debug)] @@ -124,6 +146,12 @@ impl Display for Error { } } +impl From for Error { + fn from(msg: String) -> Self { + Error::InternalError(msg) + } +} + pub const BM_DOCUMENT_IDS: u8 = 0; pub const BM_KEYWORD: u8 = 1 << 5; pub const BM_TAG: u8 = 1 << 6; @@ -138,3 +166,12 @@ pub const BLOOM_TRIGRAM: u8 = 1 << 1; pub const TAG_ID: u8 = 0; pub const TAG_TEXT: u8 = 1 << 0; pub const TAG_STATIC: u8 = 1 << 1; + +pub const BLOB_HASH_LEN: usize = 32; + +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'; diff --git a/src/query/get.rs b/src/query/get.rs index a8bd4a59..a7e8ce94 100644 --- a/src/query/get.rs +++ b/src/query/get.rs @@ -1,4 +1,4 @@ -use crate::{Deserialize, Store, ValueKey}; +use crate::{Deserialize, Key, Store, ValueKey}; impl Store { pub async fn get_value(&self, key: ValueKey) -> crate::Result> @@ -48,4 +48,29 @@ impl Store { .await } } + + pub async fn iterate( + &self, + acc: T, + begin: impl Key, + end: impl Key, + first: bool, + ascending: bool, + cb: impl Fn(&mut T, &[u8], &[u8]) -> crate::Result + Sync + Send + 'static, + ) -> crate::Result { + #[cfg(feature = "is_async")] + { + self.read_transaction() + .await? + .iterate(acc, begin, end, first, ascending, cb) + .await + } + + #[cfg(feature = "is_sync")] + { + let trx = self.read_transaction()?; + self.spawn_worker(move || trx.iterate(acc, begin, end, first, ascending, cb)) + .await + } + } } diff --git a/src/tests/blobs.rs b/src/tests/blobs.rs new file mode 100644 index 00000000..262e5305 --- /dev/null +++ b/src/tests/blobs.rs @@ -0,0 +1,152 @@ +use std::{sync::Arc, time::Duration}; + +use ahash::AHashMap; + +use crate::{ + write::{BatchBuilder, F_CLEAR}, + BlobId, BlobKey, Store, BLOB_HASH_LEN, +}; + +pub async fn test(db: Arc) { + let ttl = 1_u64; + + let blob_1 = vec![b'a'; 1024]; + let blob_2 = vec![b'b'; 1024]; + + let blob_id_1 = BlobId::from(&blob_1[..]); + let blob_id_2 = BlobId::from(&blob_2[..]); + + // Insert the same blobs concurrently + let handles = (1..=100) + .map(|_| { + let db = db.clone(); + let blob_1 = blob_1.clone(); + let blob_2 = blob_2.clone(); + tokio::spawn(async move { + db.write_blob(u32::MAX, &blob_1).await.unwrap(); + db.write_blob(u32::MAX, &blob_2).await.unwrap(); + db.write( + BatchBuilder::new() + .with_account_id(0) + .with_collection(u8::MAX) + .update_document(u32::MAX) + .blob(&blob_id_2, 0) + .build_batch(), + ) + .await + .unwrap(); + }) + }) + .collect::>(); + + for handle in handles { + handle.await.unwrap(); + } + + // Count number of blobs + let mut expected_count = AHashMap::from_iter([(blob_id_1, (0, 1)), (blob_id_2, (0, 2))]); + assert_eq!(expected_count, db.get_all_blobs().await); + + // Purgimg should not delete any blobs at this point + db.purge_blobs(ttl).await.unwrap(); + assert_eq!(expected_count, db.get_all_blobs().await); + + // Link blob to an account + db.write( + BatchBuilder::new() + .with_account_id(2) + .with_collection(u8::MAX) + .update_document(2) + .blob(&blob_id_1, 0) + .build_batch(), + ) + .await + .unwrap(); + + // Check expected count + expected_count.insert(blob_id_1, (1, 1)); + assert_eq!(expected_count, db.get_all_blobs().await); + + // Wait 1 second until the blob reaches its TTL + tokio::time::sleep(Duration::from_millis(1100)).await; + db.purge_blobs(ttl).await.unwrap(); + expected_count.insert(blob_id_1, (1, 0)); + expected_count.remove(&blob_id_2); + assert_eq!(expected_count, db.get_all_blobs().await); + + // Unlink blob, purge and make sure it is removed. + db.write( + BatchBuilder::new() + .with_account_id(2) + .with_collection(u8::MAX) + .update_document(2) + .blob(&blob_id_1, F_CLEAR) + .build_batch(), + ) + .await + .unwrap(); + db.purge_blobs(ttl).await.unwrap(); + expected_count.remove(&blob_id_1); + assert_eq!(expected_count, db.get_all_blobs().await); +} + +struct BlobPurge { + result: AHashMap, + link_count: u32, + ephemeral_count: u32, + id: [u8; BLOB_HASH_LEN], +} + +impl Store { + async fn get_all_blobs(&self) -> AHashMap { + let results = BlobPurge { + result: AHashMap::new(), + id: [0u8; BLOB_HASH_LEN], + link_count: u32::MAX, + ephemeral_count: u32::MAX, + }; + + let from_key = BlobKey { + account_id: 0, + collection: 0, + document_id: 0, + hash: [0; BLOB_HASH_LEN], + }; + let to_key = BlobKey { + account_id: u32::MAX, + collection: u8::MAX, + document_id: u32::MAX, + hash: [u8::MAX; BLOB_HASH_LEN], + }; + + let mut b = self + .iterate(results, from_key, to_key, false, true, move |b, k, v| { + if !k.starts_with(&b.id) { + if b.link_count != u32::MAX { + let id = BlobId { hash: b.id }; + b.result.insert(id, (b.link_count, b.ephemeral_count)); + } + b.link_count = 0; + b.ephemeral_count = 0; + b.id.copy_from_slice(&k[..BLOB_HASH_LEN]); + } + + if v.is_empty() { + b.link_count += 1; + } else { + b.ephemeral_count += 1; + } + + Ok(true) + }) + .await + .unwrap(); + + if b.link_count != u32::MAX { + let id = BlobId { hash: b.id }; + b.result.insert(id, (b.link_count, b.ephemeral_count)); + } + + b.result + } +} diff --git a/src/tests/mod.rs b/src/tests/mod.rs index eb712e1f..d474e541 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,20 +1,37 @@ pub mod assign_id; +pub mod blobs; pub mod query; use std::{io::Read, sync::Arc}; +use utils::config::Config; + use super::*; +struct TempDir { + path: std::path::PathBuf, +} + #[tokio::test] pub async fn store_test() { - let db = Arc::new(Store::open().await.unwrap()); - let insert = false; + let temp_dir = TempDir::new("store_tests", true); + let config_file = format!( + concat!("[blob.store]\n", "path = \"{}\"\n", "hash = 1\n"), + temp_dir.path.display() + ); + let db = Arc::new( + Store::open(&Config::parse(&config_file).unwrap()) + .await + .unwrap(), + ); + let insert = true; if insert { db.destroy().await; } //assign_id::test(db).await; - - query::test(db, insert).await; + blobs::test(db).await; + //query::test(db, insert).await; + temp_dir.delete(); } pub fn deflate_artwork_data() -> Vec { @@ -32,59 +49,18 @@ pub fn deflate_artwork_data() -> Vec { result } -/* -#[test] -fn it_works() { - for n in [10, 100, 1000, 5000, 10000, 100000] { - let mut rb1 = RoaringBitmap::new(); - let mut h = BTreeSet::new(); - let m = (((n as f64) * f64::ln(0.01) / (-8.0 * LN_2.powi(2))).ceil() as u64) * 8; - - for pos in 0..(n * 7_usize) { - let num = rand::thread_rng().gen_range(0..m as u32); - rb1.insert(num); - h.insert(num); +impl TempDir { + pub fn new(name: &str, delete_if_exists: bool) -> Self { + let mut path = std::env::temp_dir(); + path.push(name); + if delete_if_exists && path.exists() { + std::fs::remove_dir_all(&path).unwrap(); } - - let mut compressed = vec![0u8; 4 * BitPacker8x::BLOCK_LEN]; - let mut bitpacker = BitPacker8x::new(); - let mut initial_value = 0; - let mut bytes = vec![]; - for chunk in h - .into_iter() - .collect::>() - .chunks_exact(BitPacker8x::BLOCK_LEN) - { - let num_bits: u8 = bitpacker.num_bits_sorted(initial_value, chunk); - let compressed_len = - bitpacker.compress_sorted(initial_value, chunk, &mut compressed[..], num_bits); - initial_value = chunk[chunk.len() - 1]; - //println!("{:?} {}", compressed_len, num_bits); - bytes.push(num_bits); - bytes.extend_from_slice(&compressed[..compressed_len]); - } - - let rb_size = rb1.serialized_size(); - let bp_size = bytes.len(); - if rb_size < bp_size { - println!("For {} Roaring is better {} vs {}", n, rb_size, bp_size); - } else { - println!("For {} BitPack is better {} vs {}", n, bp_size, rb_size); - } - let now = Instant::now(); - let mut ser = Vec::with_capacity(rb_size); - rb1.serialize_into(&mut ser).unwrap(); - println!("Roaring serialization took {:?}", now.elapsed().as_millis()); - let now = Instant::now(); - let deser = RoaringBitmap::deserialize_unchecked_from(&ser[..]).unwrap(); - println!( - "Roaring deserialization took {:?}", - now.elapsed().as_millis() - ); + std::fs::create_dir_all(&path).unwrap(); + Self { path } + } + + pub fn delete(&self) { + std::fs::remove_dir_all(&self.path).unwrap(); } - /*println!( - "ratio: {}", - rb1.serialized_size() as f64 / rb2.serialized_size() as f64 - );*/ } -*/ diff --git a/src/write/batch.rs b/src/write/batch.rs index 7307279c..9b4a1d43 100644 --- a/src/write/batch.rs +++ b/src/write/batch.rs @@ -99,7 +99,12 @@ impl BatchBuilder { self } - pub fn bitmap(&mut self, field: impl Into, value: impl IntoBitmap, options: u32) { + pub fn bitmap( + &mut self, + field: impl Into, + value: impl IntoBitmap, + options: u32, + ) -> &mut Self { let (key, family) = value.into_bitmap(); self.ops.push(Operation::Bitmap { family, @@ -107,20 +112,23 @@ impl BatchBuilder { key, set: !options.has_flag(F_CLEAR), }); + self } - pub fn acl(&mut self, grant_account_id: u32, acl: Option) { + pub fn acl(&mut self, grant_account_id: u32, acl: Option) -> &mut Self { self.ops.push(Operation::Acl { grant_account_id, set: acl.map(|acl| acl.serialize()), - }) + }); + self } - pub fn blob(&mut self, blob_id: impl Serialize, options: u32) { + pub fn blob(&mut self, blob_id: impl Serialize, options: u32) -> &mut Self { self.ops.push(Operation::Blob { key: blob_id.serialize(), set: !options.has_flag(F_CLEAR), }); + self } pub fn custom(&mut self, value: impl IntoOperations) -> crate::Result<()> { @@ -130,6 +138,16 @@ impl BatchBuilder { pub fn build(self) -> Batch { Batch { ops: self.ops } } + + pub fn build_batch(&mut self) -> Batch { + Batch { + ops: std::mem::take(&mut self.ops), + } + } + + pub fn is_empty(&self) -> bool { + self.ops.is_empty() + } } impl Default for BatchBuilder { diff --git a/src/write/key.rs b/src/write/key.rs index 256064db..5351e438 100644 --- a/src/write/key.rs +++ b/src/write/key.rs @@ -1,6 +1,8 @@ use std::convert::TryInto; use utils::codec::leb128::Leb128_; +use crate::{BlobKey, Key, SUBSPACE_BLOBS}; + pub struct KeySerializer { buf: Vec, } @@ -113,3 +115,9 @@ impl DeserializeBigEndian for &[u8] { .map(u64::from_be_bytes) } } + +impl + Sync + Send + 'static> Key for BlobKey { + fn subspace(&self) -> u8 { + SUBSPACE_BLOBS + } +} diff --git a/src/write/mod.rs b/src/write/mod.rs index adb93757..446c3488 100644 --- a/src/write/mod.rs +++ b/src/write/mod.rs @@ -1,4 +1,4 @@ -use std::collections::HashSet; +use std::{collections::HashSet, time::SystemTime}; use crate::{Deserialize, Serialize}; @@ -190,3 +190,10 @@ pub trait IntoBitmap { pub trait IntoOperations { fn build(self, batch: &mut BatchBuilder) -> crate::Result<()>; } + +#[inline(always)] +pub fn now() -> u64 { + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map_or(0, |d| d.as_secs()) +}