diff --git a/Cargo.lock b/Cargo.lock index 5fd6d148..c38d1d7f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -535,15 +535,6 @@ dependencies = [ "serde", ] -[[package]] -name = "bitpacking" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8c7d2ac73c167c06af4a5f37e6e59d84148d57ccbe4480b76f0273eefea82d7" -dependencies = [ - "crunchy", -] - [[package]] name = "bitvec" version = "1.0.1" @@ -1140,6 +1131,39 @@ dependencies = [ "generic-array", ] +[[package]] +name = "deadpool" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb84100978c1c7b37f09ed3ce3e5f843af02c2a2c431bae5b19230dad2c1b490" +dependencies = [ + "async-trait", + "deadpool-runtime", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-postgres" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40b75ba49590d27f677d3bebaf76cd15889ca8b308bc7ba99bfa25f1d7269c13" +dependencies = [ + "deadpool", + "tokio", + "tokio-postgres", + "tracing", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63dfa964fe2a66f3fde91fc70b267fe193d822c7e603e2a675a49a7f46ad3f49" +dependencies = [ + "tokio", +] + [[package]] name = "decancer" version = "1.6.5" @@ -3309,40 +3333,46 @@ dependencies = [ [[package]] name = "opentelemetry" -version = "0.20.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9591d937bc0e6d2feb6f71a559540ab300ea49955229c347a517a28d27784c54" +checksum = "1e32339a5dc40459130b3bd269e9892439f55b33e772d2a9d402a789baaf4e8a" dependencies = [ - "opentelemetry_api", - "opentelemetry_sdk", + "futures-core", + "futures-sink", + "indexmap 2.1.0", + "js-sys", + "once_cell", + "pin-project-lite", + "thiserror", + "urlencoding", ] [[package]] name = "opentelemetry-http" -version = "0.9.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7594ec0e11d8e33faf03530a4c49af7064ebba81c1480e01be67d90b356508b" +checksum = "7f51189ce8be654f9b5f7e70e49967ed894e84a06fc35c6c042e64ac1fc5399e" dependencies = [ "async-trait", "bytes", "http 0.2.11", - "opentelemetry_api", + "opentelemetry", "reqwest", ] [[package]] name = "opentelemetry-otlp" -version = "0.13.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e5e5a5c4135864099f3faafbe939eb4d7f9b80ebf68a8448da961b32a7c1275" +checksum = "f24cda83b20ed2433c68241f918d0f6fdec8b1d43b7a9590ab4420c5095ca930" dependencies = [ "async-trait", "futures-core", "http 0.2.11", + "opentelemetry", "opentelemetry-http", "opentelemetry-proto", "opentelemetry-semantic-conventions", - "opentelemetry_api", "opentelemetry_sdk", "prost", "reqwest", @@ -3353,11 +3383,11 @@ dependencies = [ [[package]] name = "opentelemetry-proto" -version = "0.3.0" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1e3f814aa9f8c905d0ee4bde026afd3b2577a97c10e1699912e3e44f0c4cbeb" +checksum = "a2e155ce5cc812ea3d1dffbd1539aed653de4bf4882d60e6e04dcf0901d674e1" dependencies = [ - "opentelemetry_api", + "opentelemetry", "opentelemetry_sdk", "prost", "tonic", @@ -3365,47 +3395,30 @@ dependencies = [ [[package]] name = "opentelemetry-semantic-conventions" -version = "0.12.0" +version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73c9f9340ad135068800e7f1b24e9e09ed9e7143f5bf8518ded3d3ec69789269" +checksum = "f5774f1ef1f982ef2a447f6ee04ec383981a3ab99c8e77a1a7b30182e65bbc84" dependencies = [ "opentelemetry", ] -[[package]] -name = "opentelemetry_api" -version = "0.20.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a81f725323db1b1206ca3da8bb19874bbd3f57c3bcd59471bfb04525b265b9b" -dependencies = [ - "futures-channel", - "futures-util", - "indexmap 1.9.3", - "js-sys", - "once_cell", - "pin-project-lite", - "thiserror", - "urlencoding", -] - [[package]] name = "opentelemetry_sdk" -version = "0.20.0" +version = "0.21.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa8e705a0612d48139799fcbaba0d4a90f06277153e43dd2bdc16c6f0edd8026" +checksum = "968ba3f2ca03e90e5187f5e4f46c791ef7f2c163ae87789c8ce5f5ca3b7b7de5" dependencies = [ "async-trait", "crossbeam-channel", "futures-channel", "futures-executor", "futures-util", + "glob", "once_cell", - "opentelemetry_api", + "opentelemetry", "ordered-float", "percent-encoding", "rand", - "regex", - "serde_json", "thiserror", "tokio", "tokio-stream", @@ -3413,9 +3426,9 @@ dependencies = [ [[package]] name = "ordered-float" -version = "3.9.2" +version = "4.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1e1c390732d15f1d48471625cd92d154e66db2c56645e29a9cd26f4699f72dc" +checksum = "536900a8093134cf9ccf00a27deb3532421099e958d9dd431135d0c7543ca1e8" dependencies = [ "num-traits", ] @@ -3698,6 +3711,35 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3bccab0e7fd7cc19f820a1c8c91720af652d0c88dc9664dd72aef2614f04af3b" +[[package]] +name = "postgres-protocol" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49b6c5ef183cd3ab4ba005f1ca64c21e8bd97ce4699cfea9e8d9a2c4958ca520" +dependencies = [ + "base64 0.21.5", + "byteorder", + "bytes", + "fallible-iterator", + "hmac 0.12.1", + "md-5 0.10.6", + "memchr", + "rand", + "sha2 0.10.8", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d2234cdee9408b523530a9b6d2d6b373d1db34f6a8e51dc03ded1828d7fb67c" +dependencies = [ + "bytes", + "fallible-iterator", + "postgres-protocol", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -5141,8 +5183,8 @@ version = "0.1.0" dependencies = [ "ahash 0.8.6", "async-trait", - "bitpacking", "blake3", + "deadpool-postgres", "farmhash", "foundationdb", "futures", @@ -5155,13 +5197,16 @@ dependencies = [ "r2d2", "rand", "rayon", + "ring 0.17.5", "roaring", "rocksdb", "rusqlite", "rust-s3", + "rustls 0.21.9", "serde", - "siphasher 1.0.0", "tokio", + "tokio-postgres", + "tokio-rustls", "tracing", "utils", "xxhash-rust", @@ -5487,6 +5532,32 @@ dependencies = [ "syn 2.0.39", ] +[[package]] +name = "tokio-postgres" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d340244b32d920260ae7448cb72b6e238bddc3d4f7603394e7dd46ed8e48f5b8" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand", + "socket2 0.5.5", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.24.1" @@ -5653,17 +5724,6 @@ dependencies = [ "tracing-subscriber", ] -[[package]] -name = "tracing-log" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f751112709b4e791d8ce53e32c4ed2d353565a795ce84da2285393f41557bdf2" -dependencies = [ - "log", - "once_cell", - "tracing-core", -] - [[package]] name = "tracing-log" version = "0.2.0" @@ -5677,18 +5737,20 @@ dependencies = [ [[package]] name = "tracing-opentelemetry" -version = "0.21.0" +version = "0.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75327c6b667828ddc28f5e3f169036cb793c3f588d83bf0f262a7f062ffed3c8" +checksum = "c67ac25c5407e7b961fafc6f7e9aa5958fd297aada2d20fa2ae1737357e55596" dependencies = [ + "js-sys", "once_cell", "opentelemetry", "opentelemetry_sdk", "smallvec", "tracing", "tracing-core", - "tracing-log 0.1.4", + "tracing-log", "tracing-subscriber", + "web-time", ] [[package]] @@ -5706,7 +5768,7 @@ dependencies = [ "thread_local", "tracing", "tracing-core", - "tracing-log 0.2.0", + "tracing-log", ] [[package]] @@ -5914,6 +5976,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions", + "opentelemetry_sdk", "privdrop", "rand", "rustls 0.21.9", @@ -5927,6 +5990,7 @@ dependencies = [ "tracing-journald", "tracing-opentelemetry", "tracing-subscriber", + "webpki-roots 0.25.3", ] [[package]] @@ -6060,6 +6124,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57099a701fb3a8043f993e8228dc24229c7b942e2b009a1b962e54489ba1d3bf" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki" version = "0.22.4" @@ -6121,6 +6195,10 @@ name = "whoami" version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22fc3756b8a9133049b26c7f61ab35416c130e8c09b660f5b3958b446f52cc50" +dependencies = [ + "wasm-bindgen", + "web-sys", +] [[package]] name = "widestring" diff --git a/crates/imap/src/core/client.rs b/crates/imap/src/core/client.rs index acb1a601..9fbeb4ef 100644 --- a/crates/imap/src/core/client.rs +++ b/crates/imap/src/core/client.rs @@ -37,7 +37,6 @@ use super::{SelectedMailbox, Session, SessionData, State, IMAP}; impl Session { pub async fn ingest(&mut self, bytes: &[u8]) -> crate::Result { /*for line in String::from_utf8_lossy(bytes).split("\r\n") { - //let c = println!("<- {:?}", &line[..std::cmp::min(line.len(), 100)]); let c = println!("{}", line); }*/ diff --git a/crates/imap/src/core/writer.rs b/crates/imap/src/core/writer.rs index 86ed287b..68ac1404 100644 --- a/crates/imap/src/core/writer.rs +++ b/crates/imap/src/core/writer.rs @@ -131,11 +131,9 @@ pub fn spawn_writer(mut stream: Event, span: tracing::Span) -> mpsc::Sender Session { pub async fn write_bytes(&self, bytes: impl Into>) -> crate::OpResult { let bytes = bytes.into(); - - /*let c = println!( - "-> {:?}", - String::from_utf8_lossy(&bytes[..std::cmp::min(bytes.len(), 100)]) - );*/ + /*for line in String::from_utf8_lossy(bytes.as_ref()).split("\r\n") { + let c = println!("{}", line); + }*/ if let Err(err) = self.writer.send(Event::Bytes(bytes)).await { debug!("Failed to send bytes: {}", err); @@ -149,10 +147,9 @@ impl Session { impl SessionData { pub async fn write_bytes(&self, bytes: impl Into>) -> bool { let bytes = bytes.into(); - /*let c = println!( - "-> {:?}", - String::from_utf8_lossy(&bytes[..std::cmp::min(bytes.len(), 100)]) - );*/ + /*for line in String::from_utf8_lossy(bytes.as_ref()).split("\r\n") { + let c = println!("{}", line); + }*/ if let Err(err) = self.writer.send(Event::Bytes(bytes)).await { debug!("Failed to send bytes: {}", err); diff --git a/crates/jmap/src/email/crypto.rs b/crates/jmap/src/email/crypto.rs index 8cd57a83..5218668b 100644 --- a/crates/jmap/src/email/crypto.rs +++ b/crates/jmap/src/email/crypto.rs @@ -194,7 +194,7 @@ impl EncryptMessage for Message<'_> { .map_err(|err| { EncryptMessageError::Error(format!("Failed to create armorer: {}", err)) })?; - let message = stream::Encryptor::for_recipients(message, keys) + let message = stream::Encryptor2::for_recipients(message, keys) .symmetric_algo(match algo { Algorithm::Aes128 => SymmetricAlgorithm::AES128, Algorithm::Aes256 => SymmetricAlgorithm::AES256, diff --git a/crates/jmap/src/lib.rs b/crates/jmap/src/lib.rs index dfc803e2..027691ae 100644 --- a/crates/jmap/src/lib.rs +++ b/crates/jmap/src/lib.rs @@ -48,7 +48,7 @@ use services::{ }; use smtp::core::SMTP; use store::{ - backend::{fs::FsStore, sqlite::SqliteStore}, + backend::{foundationdb::FdbStore, fs::FsStore, postgres::PostgresStore, sqlite::SqliteStore}, fts::FtsFilter, parking_lot::Mutex, query::{sort::Pagination, Comparator, Filter, ResultSet, SortedResultSet}, @@ -196,11 +196,27 @@ impl JMAP { .property::("global.shared-map.shard")? .unwrap_or(32) .next_power_of_two() as usize; - let store = Store::SQLite(Arc::new( - SqliteStore::open(config) + let store = Store::PostgreSQL(Arc::new( + PostgresStore::open(config) .await .failed("Unable to open database"), )); + /*let store = Store::SQLite(Arc::new( + SqliteStore::open(config) + .await + .failed("Unable to open database"), + ));*/ + /*let store = Store::FoundationDb(Arc::new( + FdbStore::open(config) + .await + .failed("Unable to open database"), + ));*/ + let blob_store = store.clone().into(); + /*let blob_store = BlobStore::Fs(Arc::new( + FsStore::open(config) + .await + .failed("Unable to open blob store"), + ));*/ let jmap_server = Arc::new(JMAP { directory: directory_config @@ -217,11 +233,7 @@ impl JMAP { .unwrap_or_else(SnowflakeIdGenerator::new), fts_store: FtsStore::Store(store.clone()), store, - blob_store: BlobStore::Fs(Arc::new( - FsStore::open(config) - .await - .failed("Unable to open blob store"), - )), + blob_store, config: Config::new(config).failed("Invalid configuration file"), sessions: TtlDashMap::with_capacity( config.property("jmap.session.cache.size")?.unwrap_or(100), diff --git a/crates/store/Cargo.toml b/crates/store/Cargo.toml index d67c8df3..3ac40bec 100644 --- a/crates/store/Cargo.toml +++ b/crates/store/Cargo.toml @@ -19,11 +19,9 @@ roaring = "0.10.1" rayon = { version = "1.5.1", optional = true } serde = { version = "1.0", features = ["derive"]} ahash = { version = "0.8.0", features = ["serde"] } -bitpacking = "0.8.4" lazy_static = "1.4" xxhash-rust = { version = "0.8.5", features = ["xxh3"] } farmhash = "1.1.5" -siphasher = "1.0" parking_lot = "0.12.1" lru-cache = { version = "0.1.2", optional = true } num_cpus = { version = "1.15.0", optional = true } @@ -31,6 +29,11 @@ blake3 = "1.3.3" tracing = "0.1" async-trait = "0.1.68" lz4_flex = { version = "0.11" } +deadpool-postgres = "0.11.0" +tokio-postgres = "0.7.10" +tokio-rustls = { version = "0.24.0"} +rustls = "0.21.0" +ring = "0.17" [dev-dependencies] tokio = { version = "1.23", features = ["full"] } diff --git a/crates/store/src/backend/foundationdb/id_assign.rs b/crates/store/src/backend/foundationdb/id_assign.rs deleted file mode 100644 index 5f873b98..00000000 --- a/crates/store/src/backend/foundationdb/id_assign.rs +++ /dev/null @@ -1,185 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail 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 crate::{ - backend::ID_ASSIGNMENT_EXPIRY, write::key::DeserializeBigEndian, Deserialize, Key, Serialize, - U32_LEN, -}; -use ahash::AHashSet; -use foundationdb::{options::StreamingMode, FdbError, KeySelector, RangeOption}; -use futures::StreamExt; -use rand::Rng; -use std::time::Instant; - -use crate::{write::now, BitmapKey, IndexKey}; - -use super::{ - bitmap::{next_available_index, BITS_PER_BLOCK}, - write::MAX_COMMIT_TIME, - FdbStore, -}; - -impl FdbStore { - pub(crate) async fn assign_document_id( - &self, - account_id: u32, - collection: impl Into + Sync + Send, - ) -> crate::Result { - let start = Instant::now(); - let collection = collection.into(); - - loop { - // First try to reuse an expired assigned id - let trx = self.db.create_trx()?; - let mut reserved_ids = AHashSet::new(); - let mut expired_ids = Vec::new(); - { - let begin = IndexKey { - account_id, - collection, - document_id: 0, - field: u8::MAX, - key: &[], - } - .serialize(true); - let end = IndexKey { - account_id, - collection, - document_id: u32::MAX, - field: u8::MAX, - key: &[], - } - .serialize(true); - - let mut values = trx.get_ranges( - RangeOption { - begin: KeySelector::first_greater_or_equal(begin), - end: KeySelector::first_greater_or_equal(end), - mode: StreamingMode::Iterator, - reverse: false, - ..RangeOption::default() - }, - true, - ); - - #[cfg(not(feature = "test_mode"))] - let expired_timestamp = now() - ID_ASSIGNMENT_EXPIRY; - #[cfg(feature = "test_mode")] - let expired_timestamp = - now() - ID_ASSIGNMENT_EXPIRY.load(std::sync::atomic::Ordering::Relaxed); - while let Some(values) = values.next().await { - for value in values? { - let key = value.key(); - let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?; - if u64::deserialize(value.value())? <= expired_timestamp { - // Found an expired id, reuse it - expired_ids.push(document_id); - } else { - // Keep track of all reserved ids - reserved_ids.insert(document_id); - } - } - } - } - - let mut document_id = u32::MAX; - - if !expired_ids.is_empty() { - // Obtain a random id from the expired ids - if expired_ids.len() > 1 { - document_id = expired_ids[rand::thread_rng().gen_range(0..expired_ids.len())]; - } else { - document_id = expired_ids[0]; - } - } else { - // Find the next available id - let mut key = BitmapKey::document_ids(account_id, collection); - let begin = key.serialize(true); - key.block_num = u32::MAX; - let end = key.serialize(true); - let mut values = trx.get_ranges( - RangeOption { - begin: KeySelector::first_greater_or_equal(begin), - end: KeySelector::first_greater_or_equal(end), - mode: StreamingMode::Iterator, - reverse: false, - ..RangeOption::default() - }, - true, - ); - - 'outer: while let Some(values) = values.next().await { - for value in values? { - let key = value.key(); - if let Some(next_id) = next_available_index( - value.value(), - key.deserialize_be_u32(key.len() - U32_LEN)?, - &reserved_ids, - ) { - document_id = next_id; - //assign_source = 3; - - break 'outer; - } - } - } - } - - // If no ids were found, assign the first available id that is not reserved - if document_id == u32::MAX { - document_id = 1024; - for document_id_ in 0..BITS_PER_BLOCK { - if !reserved_ids.contains(&document_id_) { - document_id = document_id_; - break; - } - } - } - - // Reserve the id - let key = IndexKey { - account_id, - collection, - document_id, - field: u8::MAX, - key: &[], - } - .serialize(true); - trx.get(&key, false).await?; // Read to create conflict range - trx.set(&key, &now().serialize()); - - match trx.commit().await { - Ok(_) => { - return Ok(document_id); - } - Err(err) => { - if start.elapsed() < MAX_COMMIT_TIME { - err.on_error().await?; - } else { - return Err(FdbError::from(err).into()); - } - } - } - } - } -} diff --git a/crates/store/src/backend/foundationdb/mod.rs b/crates/store/src/backend/foundationdb/mod.rs index 48236d3d..82a11b4e 100644 --- a/crates/store/src/backend/foundationdb/mod.rs +++ b/crates/store/src/backend/foundationdb/mod.rs @@ -25,9 +25,7 @@ use foundationdb::{api::NetworkAutoStop, Database, FdbError}; use crate::Error; -pub mod bitmap; pub mod blob; -pub mod id_assign; pub mod main; pub mod purge; pub mod read; diff --git a/crates/store/src/backend/foundationdb/purge.rs b/crates/store/src/backend/foundationdb/purge.rs index c833aa5c..4a5590ca 100644 --- a/crates/store/src/backend/foundationdb/purge.rs +++ b/crates/store/src/backend/foundationdb/purge.rs @@ -28,11 +28,12 @@ use foundationdb::{ use futures::StreamExt; use crate::{ - write::key::KeySerializer, SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_VALUES, + write::{bitmap::DenseBitmap, key::KeySerializer}, + SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES, U32_LEN, }; -use super::{bitmap::DenseBitmap, FdbStore}; +use super::FdbStore; const MAX_COMMIT_ATTEMPTS: u8 = 25; @@ -97,6 +98,7 @@ impl FdbStore { SUBSPACE_VALUES, SUBSPACE_LOGS, SUBSPACE_INDEXES, + SUBSPACE_INDEX_VALUES, ] { let from_key = KeySerializer::new(U32_LEN + 2) .write(subspace) diff --git a/crates/store/src/backend/foundationdb/read.rs b/crates/store/src/backend/foundationdb/read.rs index 8dffd31f..8d037c79 100644 --- a/crates/store/src/backend/foundationdb/read.rs +++ b/crates/store/src/backend/foundationdb/read.rs @@ -31,6 +31,7 @@ use roaring::RoaringBitmap; use crate::{ query::{self, Operator}, write::{ + bitmap::DeserializeBlock, key::{DeserializeBigEndian, KeySerializer}, BitmapClass, ValueClass, }, @@ -38,7 +39,7 @@ use crate::{ SUBSPACE_INDEXES, U32_LEN, }; -use super::{bitmap::DeserializeBlock, FdbStore}; +use super::FdbStore; impl FdbStore { pub(crate) async fn get_value(&self, key: impl Key) -> crate::Result> @@ -279,7 +280,7 @@ impl FdbStore { #[cfg(feature = "test_mode")] pub(crate) async fn assert_is_empty(&self) { - use crate::{SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_LOGS, SUBSPACE_VALUES}; + use crate::{SUBSPACE_BITMAPS, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES}; let conn = self.db.create_trx().unwrap(); @@ -344,7 +345,7 @@ impl FdbStore { ); } } - SUBSPACE_BLOBS | SUBSPACE_ACLS => { + SUBSPACE_BLOBS | SUBSPACE_INDEX_VALUES => { panic!( "Subspace {:?} is not empty: {key:?} {value:?}", char::from(subspace) diff --git a/crates/store/src/backend/foundationdb/write.rs b/crates/store/src/backend/foundationdb/write.rs index e72cc03e..89eb776c 100644 --- a/crates/store/src/backend/foundationdb/write.rs +++ b/crates/store/src/backend/foundationdb/write.rs @@ -25,25 +25,14 @@ use std::time::{Duration, Instant}; use ahash::AHashMap; use foundationdb::{options::MutationType, FdbError}; +use rand::Rng; use crate::{ - write::{Batch, Operation, ValueOp}, + write::{bitmap::DenseBitmap, Batch, Operation, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME}, BitmapKey, BlobKey, IndexKey, Key, LogKey, ValueKey, }; -use super::{bitmap::DenseBitmap, FdbStore}; - -#[cfg(not(feature = "test_mode"))] -pub const ID_ASSIGNMENT_EXPIRY: u64 = 60 * 60; // seconds -#[cfg(not(feature = "test_mode"))] -pub const MAX_COMMIT_ATTEMPTS: u32 = 10; -#[cfg(not(feature = "test_mode"))] -pub const MAX_COMMIT_TIME: Duration = Duration::from_secs(10); - -#[cfg(feature = "test_mode")] -pub const MAX_COMMIT_ATTEMPTS: u32 = 1000; -#[cfg(feature = "test_mode")] -pub const MAX_COMMIT_TIME: Duration = Duration::from_secs(3600); +use super::FdbStore; #[cfg(feature = "test_mode")] lazy_static::lazy_static! { @@ -273,6 +262,8 @@ impl FdbStore { Err(err) => { if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME { err.on_error().await?; + let backoff = rand::thread_rng().gen_range(50..=300); + tokio::time::sleep(Duration::from_millis(backoff)).await; retry_count += 1; } else { return Err(FdbError::from(err).into()); diff --git a/crates/store/src/backend/mod.rs b/crates/store/src/backend/mod.rs index 60202d0d..850cbc02 100644 --- a/crates/store/src/backend/mod.rs +++ b/crates/store/src/backend/mod.rs @@ -24,6 +24,7 @@ #[cfg(feature = "foundation")] pub mod foundationdb; pub mod fs; +pub mod postgres; #[cfg(feature = "rocks")] pub mod rocksdb; pub mod s3; diff --git a/crates/store/src/backend/postgres/blob.rs b/crates/store/src/backend/postgres/blob.rs new file mode 100644 index 00000000..cf2902ea --- /dev/null +++ b/crates/store/src/backend/postgres/blob.rs @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail Server. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * in the LICENSE file at the top-level directory of this distribution. + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + * + * You can be released from the requirements of the AGPLv3 license by + * purchasing a commercial license. Please contact licensing@stalw.art + * for more details. +*/ + +use std::ops::Range; + +use super::PostgresStore; + +impl PostgresStore { + pub(crate) async fn get_blob( + &self, + key: &[u8], + range: Range, + ) -> crate::Result>> { + let conn = self.conn_pool.get().await?; + let s = conn.prepare_cached("SELECT v FROM t WHERE k = $1").await?; + conn.query_opt(&s, &[&key]) + .await + .and_then(|row| { + if let Some(row) = row { + Ok(Some(if range.start == 0 && range.end == u32::MAX { + row.try_get::<_, Vec>(0)? + } else { + let bytes = row.try_get::<_, &[u8]>(0)?; + bytes + .get( + range.start as usize + ..std::cmp::min(bytes.len(), range.end as usize), + ) + .unwrap_or_default() + .to_vec() + })) + } else { + Ok(None) + } + }) + .map_err(Into::into) + } + + pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> crate::Result<()> { + let conn = self.conn_pool.get().await?; + let s = conn + .prepare_cached( + "INSERT INTO t (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v", + ) + .await?; + conn.execute(&s, &[&key, &data]) + .await + .map_err(|e| crate::Error::InternalError(format!("Failed to insert blob: {}", e))) + .map(|_| ()) + } + + pub(crate) async fn delete_blob(&self, key: &[u8]) -> crate::Result { + let conn = self.conn_pool.get().await?; + let s = conn.prepare_cached("DELETE FROM t WHERE k = $1").await?; + conn.execute(&s, &[&key]) + .await + .map_err(|e| crate::Error::InternalError(format!("Failed to delete blob: {}", e))) + .map(|hits| hits > 0) + } +} diff --git a/crates/store/src/backend/postgres/main.rs b/crates/store/src/backend/postgres/main.rs new file mode 100644 index 00000000..d12e31d7 --- /dev/null +++ b/crates/store/src/backend/postgres/main.rs @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail 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 crate::{ + backend::postgres::tls::MakeRustlsConnect, SUBSPACE_BITMAPS, SUBSPACE_BLOBS, + SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, + SUBSPACE_VALUES, +}; + +use super::PostgresStore; + +use deadpool_postgres::{ + Config, CreatePoolError, ManagerConfig, PoolConfig, RecyclingMethod, Runtime, +}; +use tokio_postgres::NoTls; +use utils::rustls_client_config; + +impl PostgresStore { + pub async fn open(config: &utils::config::Config) -> crate::Result { + let mut cfg = Config::new(); + cfg.dbname = config + .value_require("store.db.database")? + .to_string() + .into(); + cfg.host = config.value("store.db.host").map(|s| s.to_string()); + cfg.user = config.value("store.db.user").map(|s| s.to_string()); + cfg.password = config.value("store.db.password").map(|s| s.to_string()); + cfg.port = config.property("store.db.port")?; + cfg.connect_timeout = config.property("store.db.timeout")?; + cfg.manager = Some(ManagerConfig { + recycling_method: RecyclingMethod::Fast, + }); + if let Some(max_conn) = config.property::("store.db.pool.max-connections")? { + cfg.pool = PoolConfig::new(max_conn).into(); + } + let db = Self { + conn_pool: if config.property_or_static::("store.db.tls.enable", "false")? { + cfg.create_pool( + Some(Runtime::Tokio1), + MakeRustlsConnect::new(rustls_client_config( + config.property_or_static("store.db.tls.allow-invalid-certs", "false")?, + )), + )? + } else { + cfg.create_pool(Some(Runtime::Tokio1), NoTls)? + }, + }; + + db.create_tables().await?; + + Ok(db) + } + + pub(super) async fn create_tables(&self) -> crate::Result<()> { + let conn = self.conn_pool.get().await?; + + for table in [ + SUBSPACE_VALUES, + SUBSPACE_LOGS, + SUBSPACE_INDEX_VALUES, + SUBSPACE_BLOB_DATA, + ] { + let table = char::from(table); + conn.execute( + &format!( + "CREATE TABLE IF NOT EXISTS {table} ( + k BYTEA PRIMARY KEY, + v BYTEA NOT NULL + )" + ), + &[], + ) + .await?; + } + + for table in [SUBSPACE_INDEXES, SUBSPACE_BITMAPS, SUBSPACE_BLOBS] { + let table = char::from(table); + conn.execute( + &format!( + "CREATE TABLE IF NOT EXISTS {table} ( + k BYTEA PRIMARY KEY + )" + ), + &[], + ) + .await?; + } + + conn.execute( + &format!( + "CREATE TABLE IF NOT EXISTS {} ( + k BYTEA PRIMARY KEY, + v BIGINT NOT NULL DEFAULT 0 + )", + char::from(SUBSPACE_COUNTERS) + ), + &[], + ) + .await?; + + Ok(()) + } +} + +impl From for crate::Error { + fn from(err: CreatePoolError) -> Self { + crate::Error::InternalError(format!("Failed to create connection pool: {}", err)) + } +} diff --git a/crates/store/src/backend/postgres/mod.rs b/crates/store/src/backend/postgres/mod.rs new file mode 100644 index 00000000..a9487f06 --- /dev/null +++ b/crates/store/src/backend/postgres/mod.rs @@ -0,0 +1,47 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail 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 deadpool_postgres::{Pool, PoolError}; + +pub mod blob; +pub mod main; +pub mod purge; +pub mod read; +pub mod tls; +pub mod write; + +pub struct PostgresStore { + pub(crate) conn_pool: Pool, +} + +impl From for crate::Error { + fn from(err: PoolError) -> Self { + Self::InternalError(format!("Connection pool error: {}", err)) + } +} + +impl From for crate::Error { + fn from(err: tokio_postgres::Error) -> Self { + Self::InternalError(format!("PostgreSQL error: {}", err)) + } +} diff --git a/crates/store/src/backend/postgres/purge.rs b/crates/store/src/backend/postgres/purge.rs new file mode 100644 index 00000000..73dbb838 --- /dev/null +++ b/crates/store/src/backend/postgres/purge.rs @@ -0,0 +1,62 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail 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 crate::{ + write::key::KeySerializer, SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, + SUBSPACE_LOGS, SUBSPACE_VALUES, U32_LEN, +}; + +use super::PostgresStore; + +impl PostgresStore { + pub(crate) async fn purge_bitmaps(&self) -> crate::Result<()> { + // Not needed for PostgreSQL + Ok(()) + } + + pub(crate) async fn purge_account(&self, account_id: u32) -> crate::Result<()> { + let conn = self.conn_pool.get().await?; + let from_key = KeySerializer::new(U32_LEN).write(account_id).finalize(); + let to_key = KeySerializer::new(U32_LEN).write(account_id + 1).finalize(); + + for (table, i) in [ + (SUBSPACE_BITMAPS, 'z'), + (SUBSPACE_VALUES, 'k'), + (SUBSPACE_LOGS, 'k'), + (SUBSPACE_INDEXES, 'k'), + (SUBSPACE_INDEX_VALUES, 'k'), + ] { + let s = conn + .prepare_cached(&format!( + "DELETE FROM {} WHERE {} >= ? AND {} < ?", + char::from(table), + i, + i + )) + .await?; + conn.execute(&s, &[&from_key, &to_key]).await?; + } + + Ok(()) + } +} diff --git a/crates/store/src/backend/postgres/read.rs b/crates/store/src/backend/postgres/read.rs new file mode 100644 index 00000000..8304bd6a --- /dev/null +++ b/crates/store/src/backend/postgres/read.rs @@ -0,0 +1,378 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail 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 futures::{pin_mut, TryStreamExt}; +use roaring::RoaringBitmap; + +use crate::{ + query::{self, Operator}, + write::{ + key::{DeserializeBigEndian, KeySerializer}, + BitmapClass, ValueClass, + }, + BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, IterateParams, Key, ValueKey, U32_LEN, +}; + +use super::PostgresStore; + +impl PostgresStore { + pub(crate) async fn get_value(&self, key: impl Key) -> crate::Result> + where + U: Deserialize + 'static, + { + let key = key.serialize(false); + let conn = self.conn_pool.get().await?; + let s = conn.prepare_cached("SELECT v FROM v WHERE k = $1").await?; + conn.query_opt(&s, &[&key]) + .await + .map_err(Into::into) + .and_then(|r| { + if let Some(r) = r { + Ok(Some(U::deserialize(r.get(0))?)) + } else { + Ok(None) + } + }) + } + + pub(crate) async fn get_bitmap( + &self, + mut key: BitmapKey, + ) -> crate::Result> { + let begin = key.serialize(false); + key.block_num = u32::MAX; + let key_len = begin.len(); + let end = key.serialize(false); + let conn = self.conn_pool.get().await?; + + let mut bm = RoaringBitmap::new(); + let s = conn + .prepare_cached("SELECT k FROM b WHERE k >= $1 AND k <= $2") + .await?; + let rows = conn.query_raw(&s, &[&begin, &end]).await?; + + pin_mut!(rows); + + while let Some(row) = rows.try_next().await? { + let key: &[u8] = row.try_get(0)?; + if key.len() == key_len { + bm.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); + } + } + Ok(if !bm.is_empty() { Some(bm) } else { None }) + } + + pub(crate) async fn range_to_bitmap( + &self, + account_id: u32, + collection: u8, + field: u8, + value: &[u8], + op: query::Operator, + ) -> crate::Result> { + let conn = self.conn_pool.get().await?; + let k1 = + KeySerializer::new(std::mem::size_of::>() + value.len() + 1 + U32_LEN) + .write(account_id) + .write(collection) + .write(field); + let k2 = + KeySerializer::new(std::mem::size_of::>() + value.len() + 1 + U32_LEN) + .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 >= $1 AND k < $2"), + (k1.finalize()), + (k2.write(value).write(0u32).finalize()), + ), + Operator::LowerEqualThan => ( + ("SELECT k FROM i WHERE k >= $1 AND k <= $2"), + (k1.finalize()), + (k2.write(value).write(u32::MAX).finalize()), + ), + Operator::GreaterThan => ( + ("SELECT k FROM i WHERE k > $1 AND k <= $2"), + (k1.write(value).write(u32::MAX).finalize()), + (k2.finalize()), + ), + Operator::GreaterEqualThan => ( + ("SELECT k FROM i WHERE k >= $1 AND k <= $2"), + (k1.write(value).write(0u32).finalize()), + (k2.finalize()), + ), + Operator::Equal => ( + ("SELECT k FROM i WHERE k >= $1 AND k <= $2"), + (k1.write(value).write(0u32).finalize()), + (k2.write(value).write(u32::MAX).finalize()), + ), + }; + + let mut bm = RoaringBitmap::new(); + let s = conn.prepare_cached(query).await?; + let rows = conn.query_raw(&s, &[&begin, &end]).await?; + + pin_mut!(rows); + + if op != Operator::Equal { + while let Some(row) = rows.try_next().await? { + let key = row.try_get::<_, &[u8]>(0)?; + bm.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); + } + } else { + let key_len = begin.len(); + while let Some(row) = rows.try_next().await? { + let key = row.try_get::<_, &[u8]>(0)?; + if key.len() == key_len { + bm.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); + } + } + } + + Ok(Some(bm)) + } + + pub(crate) async fn sort_index( + &self, + account_id: u32, + collection: impl Into + Sync + Send, + field: impl Into + Sync + Send, + ascending: bool, + mut cb: impl for<'x> FnMut(&'x [u8], u32) -> crate::Result + Sync + Send, + ) -> crate::Result<()> { + let collection = collection.into(); + let field = field.into(); + + let conn = self.conn_pool.get().await?; + let begin = IndexKeyPrefix { + account_id, + collection, + field, + } + .serialize(false); + let end = IndexKeyPrefix { + account_id, + collection, + field: field + 1, + } + .serialize(false); + let prefix_len = begin.len(); + let s = conn + .prepare_cached(if ascending { + "SELECT k FROM i WHERE k >= $1 AND k < $2 ORDER BY k ASC" + } else { + "SELECT k FROM i WHERE k >= $1 AND k < $2 ORDER BY k DESC" + }) + .await?; + let rows = conn.query_raw(&s, &[&begin, &end]).await?; + + pin_mut!(rows); + + while let Some(row) = rows.try_next().await? { + let key = row.try_get::<_, &[u8]>(0)?; + let id_pos = key.len() - U32_LEN; + 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(()) + } + + pub(crate) async fn iterate( + &self, + params: IterateParams, + mut cb: impl for<'x> FnMut(&'x [u8], &'x [u8]) -> crate::Result + Sync + Send, + ) -> crate::Result<()> { + let conn = self.conn_pool.get().await?; + let table = char::from(params.begin.subspace()); + let begin = params.begin.serialize(false); + let end = params.end.serialize(false); + let keys = if params.values { "k, v" } else { "k" }; + + let s = conn + .prepare_cached(&match (params.first, params.ascending) { + (true, true) => { + format!( + "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1" + ) + } + (true, false) => { + format!( + "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC LIMIT 1" + ) + } + (false, true) => { + format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC") + } + (false, false) => { + format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC") + } + }) + .await?; + let rows = conn.query_raw(&s, &[&begin, &end]).await?; + + pin_mut!(rows); + + if params.values { + while let Some(row) = rows.try_next().await? { + let key = row.try_get::<_, &[u8]>(0)?; + let value = row.try_get::<_, &[u8]>(1)?; + + if !cb(key, value)? { + break; + } + } + } else { + while let Some(row) = rows.try_next().await? { + if !cb(row.try_get::<_, &[u8]>(0)?, b"")? { + break; + } + } + } + + Ok(()) + } + + pub(crate) async fn get_counter( + &self, + key: impl Into> + Sync + Send, + ) -> crate::Result { + let key = key.into().serialize(false); + let conn = self.conn_pool.get().await?; + let s = conn.prepare_cached("SELECT v FROM c WHERE k = $1").await?; + match conn.query_opt(&s, &[&key]).await { + Ok(Some(row)) => row.try_get(0).map_err(Into::into), + Ok(None) => Ok(0), + Err(e) => Err(e.into()), + } + } + + #[cfg(feature = "test_mode")] + pub(crate) async fn assert_is_empty(&self) { + let conn = self.conn_pool.get().await.unwrap(); + + // Values + let mut has_errors = false; + for table in [ + crate::SUBSPACE_VALUES, + crate::SUBSPACE_INDEX_VALUES, + crate::SUBSPACE_COUNTERS, + crate::SUBSPACE_BLOB_DATA, + ] { + let table = char::from(table); + let s = conn + .prepare_cached(&format!("SELECT k, v FROM {table}")) + .await + .unwrap(); + let rows = conn.query_raw(&s, Vec::<&[u8]>::new()).await.unwrap(); + + pin_mut!(rows); + + while let Some(row) = rows.try_next().await.unwrap() { + let key = row.get::<_, &[u8]>(0); + if table != 'c' { + let value = row.get::<_, &[u8]>(1); + + if key[0..4] != u32::MAX.to_be_bytes() { + eprintln!("Table {table:?} is not empty: {key:?} {value:?}"); + has_errors = true; + } + } else { + let value = row.get::<_, i64>(1); + if value != 0 { + eprintln!( + "Table counter is not empty, account {:?}, quota: {}", + key, value, + ); + has_errors = true; + } + } + } + } + + // Indexes + for table in [crate::SUBSPACE_INDEXES, crate::SUBSPACE_BLOBS] { + let table = char::from(table); + let s = conn + .prepare_cached(&format!("SELECT k FROM {table}")) + .await + .unwrap(); + let rows = conn.query_raw(&s, Vec::<&[u8]>::new()).await.unwrap(); + pin_mut!(rows); + while let Some(row) = rows.try_next().await.unwrap() { + let key = row.get::<_, &[u8]>(0); + + if table == 'i' { + eprintln!( + "Table index is not empty, account {}, collection {}, document {}, property {}, value {:?}: {:?}", + u32::from_be_bytes(key[0..4].try_into().unwrap()), + key[4], + u32::from_be_bytes(key[key.len()-4..].try_into().unwrap()), + key[5], + String::from_utf8_lossy(&key[6..key.len()-4]), + key + ); + } else { + eprintln!("Table {table:?} is not empty: {key:?}"); + } + has_errors = true; + } + } + + // Bitmaps + let s = conn + .prepare_cached(&format!( + "SELECT k FROM {}", + char::from(crate::SUBSPACE_BITMAPS) + )) + .await + .unwrap(); + let rows = conn.query_raw(&s, Vec::<&[u8]>::new()).await.unwrap(); + pin_mut!(rows); + while let Some(row) = rows.try_next().await.unwrap() { + let key = row.get::<_, &[u8]>(0); + if key[0..4] != u32::MAX.to_be_bytes() { + eprintln!("Table bitmaps failed to purge, found key: {key:?}"); + has_errors = true; + } + } + + // Delete logs + conn.execute("DELETE FROM l", &[]).await.unwrap(); + + if has_errors { + panic!("Database is not empty"); + } + } +} diff --git a/crates/store/src/backend/postgres/tls.rs b/crates/store/src/backend/postgres/tls.rs new file mode 100644 index 00000000..76805aeb --- /dev/null +++ b/crates/store/src/backend/postgres/tls.rs @@ -0,0 +1,152 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail 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. +*/ + +// Credits: https://github.com/jbg/tokio-postgres-rustls + +use std::{ + convert::TryFrom, + future::Future, + io, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use futures::future::{FutureExt, TryFutureExt}; +use ring::digest; +use rustls::{ClientConfig, ServerName}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio_postgres::tls::{ChannelBinding, MakeTlsConnect, TlsConnect}; +use tokio_rustls::{client::TlsStream, TlsConnector}; + +#[derive(Clone)] +pub struct MakeRustlsConnect { + config: Arc, +} + +impl MakeRustlsConnect { + pub fn new(config: ClientConfig) -> Self { + Self { + config: Arc::new(config), + } + } +} + +impl MakeTlsConnect for MakeRustlsConnect +where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + type Stream = RustlsStream; + type TlsConnect = RustlsConnect; + type Error = io::Error; + + fn make_tls_connect(&mut self, hostname: &str) -> io::Result { + ServerName::try_from(hostname) + .map(|dns_name| { + RustlsConnect(Some(RustlsConnectData { + hostname: dns_name, + connector: Arc::clone(&self.config).into(), + })) + }) + .or(Ok(RustlsConnect(None))) + } +} + +pub struct RustlsConnect(Option); + +struct RustlsConnectData { + hostname: ServerName, + connector: TlsConnector, +} + +impl TlsConnect for RustlsConnect +where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + type Stream = RustlsStream; + type Error = io::Error; + type Future = Pin>> + Send>>; + + fn connect(self, stream: S) -> Self::Future { + match self.0 { + None => Box::pin(core::future::ready(Err(io::ErrorKind::InvalidInput.into()))), + Some(c) => c + .connector + .connect(c.hostname, stream) + .map_ok(|s| RustlsStream(Box::pin(s))) + .boxed(), + } + } +} + +pub struct RustlsStream(Pin>>); + +impl tokio_postgres::tls::TlsStream for RustlsStream +where + S: AsyncRead + AsyncWrite + Unpin, +{ + fn channel_binding(&self) -> ChannelBinding { + let (_, session) = self.0.get_ref(); + match session.peer_certificates() { + Some(certs) if !certs.is_empty() => { + let sha256 = digest::digest(&digest::SHA256, certs[0].as_ref()); + ChannelBinding::tls_server_end_point(sha256.as_ref().into()) + } + _ => ChannelBinding::none(), + } + } +} + +impl AsyncRead for RustlsStream +where + S: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + self.0.as_mut().poll_read(cx, buf) + } +} + +impl AsyncWrite for RustlsStream +where + S: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context, + buf: &[u8], + ) -> Poll> { + self.0.as_mut().poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + self.0.as_mut().poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + self.0.as_mut().poll_shutdown(cx) + } +} diff --git a/crates/store/src/backend/postgres/write.rs b/crates/store/src/backend/postgres/write.rs new file mode 100644 index 00000000..7716c79d --- /dev/null +++ b/crates/store/src/backend/postgres/write.rs @@ -0,0 +1,328 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail Server. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * in the LICENSE file at the top-level directory of this distribution. + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + * + * You can be released from the requirements of the AGPLv3 license by + * purchasing a commercial license. Please contact licensing@stalw.art + * for more details. +*/ + +use std::time::{Duration, Instant}; + +use ahash::AHashMap; +use deadpool_postgres::Object; +use rand::Rng; +use tokio_postgres::{error::SqlState, IsolationLevel}; + +use crate::{ + write::{ + Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, + }, + BitmapKey, BlobKey, IndexKey, Key, LogKey, ValueKey, +}; + +use super::PostgresStore; + +impl PostgresStore { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { + let mut conn = self.conn_pool.get().await?; + let start = Instant::now(); + let mut retry_count = 0; + + loop { + match self.write_trx(&mut conn, &batch).await { + Ok(success) => { + return if success { + Ok(()) + } else { + Err(crate::Error::AssertValueFailed) + }; + } + Err(err) => match err.code() { + Some( + &SqlState::T_R_SERIALIZATION_FAILURE | &SqlState::T_R_DEADLOCK_DETECTED, + ) if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => { + let backoff = rand::thread_rng().gen_range(50..=300); + tokio::time::sleep(Duration::from_millis(backoff)).await; + retry_count += 1; + } + Some(&SqlState::UNIQUE_VIOLATION) => { + return Err(crate::Error::AssertValueFailed); + } + _ => return Err(err.into()), + }, + } + } + } + + async fn write_trx( + &self, + conn: &mut Object, + batch: &Batch, + ) -> Result { + let mut account_id = u32::MAX; + let mut collection = u8::MAX; + let mut document_id = u32::MAX; + let mut asserted_values = AHashMap::new(); + let trx = conn + .build_transaction() + .isolation_level(IsolationLevel::ReadCommitted) + .start() + .await?; + + 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_; + } + Operation::Value { + class, + op: ValueOp::Add(by), + } => { + let key = ValueKey { + account_id, + collection, + document_id, + class, + } + .serialize(false); + + if *by >= 0 { + let s = trx + .prepare_cached(concat!( + "INSERT INTO c (k, v) VALUES ($1, $2) ", + "ON CONFLICT(k) DO UPDATE SET v = c.v + EXCLUDED.v" + )) + .await?; + trx.execute(&s, &[&key, &by]).await?; + } else { + let s = trx + .prepare_cached("UPDATE c SET v = v + $1 WHERE k = $2") + .await?; + trx.execute(&s, &[&by, &key]).await?; + } + } + Operation::Value { class, op } => { + let key = ValueKey { + account_id, + collection, + document_id, + class, + }; + let table = char::from(key.subspace()); + let key = key.serialize(false); + + if let ValueOp::Set(value) = op { + let s = if let Some(exists) = asserted_values.get(&key) { + if *exists { + trx.prepare_cached(&format!( + "UPDATE {} SET v = $2 WHERE k = $1", + table + )) + .await? + } else { + trx.prepare_cached(&format!( + "INSERT INTO {} (k, v) VALUES ($1, $2)", + table + )) + .await? + } + } else { + trx + .prepare_cached( + &format!("INSERT INTO {} (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v", table), + ) + .await? + }; + + if trx.execute(&s, &[&key, value]).await? == 0 { + return Ok(false); + } + + if matches!(class, ValueClass::ReservedId) { + // Make sure the reserved id is not already in use + let s = trx.prepare_cached("SELECT 1 FROM b WHERE k = $1").await?; + let key = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + block_num: document_id, + } + .serialize(false); + if trx.query_opt(&s, &[&key]).await?.is_some() { + return Ok(false); + } + } + } else { + let s = trx + .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) + .await?; + trx.execute(&s, &[&key]).await?; + } + } + Operation::Index { field, key, set } => { + let key = IndexKey { + account_id, + collection, + document_id, + field: *field, + key, + } + .serialize(false); + + let s = if *set { + trx.prepare_cached( + "INSERT INTO i (k) VALUES ($1) ON CONFLICT (k) DO NOTHING", + ) + .await? + } else { + trx.prepare_cached("DELETE FROM i WHERE k = $1").await? + }; + trx.execute(&s, &[&key]).await?; + } + Operation::Bitmap { class, set } => { + let key = BitmapKey { + account_id, + collection, + class, + block_num: document_id, + } + .serialize(false); + + let s = if *set { + if matches!(class, BitmapClass::DocumentIds) { + trx.prepare_cached("INSERT INTO b (k) VALUES ($1)").await? + } else { + trx.prepare_cached( + "INSERT INTO b (k) VALUES ($1) ON CONFLICT (k) DO NOTHING", + ) + .await? + } + } else { + trx.prepare_cached("DELETE FROM b WHERE k = $1").await? + }; + trx.execute(&s, &[&key]).await?; + } + Operation::Blob { hash, op, set } => { + let key = BlobKey { + account_id, + collection, + document_id, + hash, + op: *op, + } + .serialize(false); + + let s = if *set { + trx.prepare_cached( + "INSERT INTO o (k) VALUES ($1) ON CONFLICT (k) DO NOTHING", + ) + .await? + } else { + trx.prepare_cached("DELETE FROM o WHERE k = $1").await? + }; + trx.execute(&s, &[&key]).await?; + } + Operation::Log { + collection, + change_id, + set, + } => { + let key = LogKey { + account_id, + collection: *collection, + change_id: *change_id, + } + .serialize(false); + + let s = trx + .prepare_cached("INSERT INTO l (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v") + .await?; + trx.execute(&s, &[&key, set]).await?; + } + Operation::AssertValue { + class, + assert_value, + } => { + let key = ValueKey { + account_id, + collection, + document_id, + class, + }; + let table = char::from(key.subspace()); + let key = key.serialize(false); + + let s = trx + .prepare_cached(&format!("SELECT v FROM {} WHERE k = $1 FOR UPDATE", table)) + .await?; + let (exists, matches) = trx + .query_opt(&s, &[&key]) + .await? + .map(|row| { + row.try_get::<_, &[u8]>(0) + .map_or((true, false), |v| (true, assert_value.matches(v))) + }) + .unwrap_or_else(|| (false, assert_value.is_none())); + if !matches { + return Ok(false); + } + asserted_values.insert(key, exists); + } + } + } + + trx.commit().await.map(|_| true) + } + + #[cfg(feature = "test_mode")] + pub(crate) async fn destroy(&self) { + use crate::{ + SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, + SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES, + }; + + let conn = self.conn_pool.get().await.unwrap(); + for table in [ + SUBSPACE_VALUES, + SUBSPACE_LOGS, + SUBSPACE_BITMAPS, + SUBSPACE_INDEXES, + SUBSPACE_BLOBS, + SUBSPACE_INDEX_VALUES, + SUBSPACE_COUNTERS, + SUBSPACE_BLOB_DATA, + ] { + conn.execute(&format!("DROP TABLE {}", char::from(table)), &[]) + .await + .unwrap(); + } + self.create_tables().await.unwrap(); + } +} diff --git a/crates/store/src/backend/sqlite/id_assign.rs b/crates/store/src/backend/sqlite/id_assign.rs deleted file mode 100644 index 037b5645..00000000 --- a/crates/store/src/backend/sqlite/id_assign.rs +++ /dev/null @@ -1,146 +0,0 @@ -/* - * Copyright (c) 2023, Stalwart Labs Ltd. - * - * This file is part of Stalwart Mail 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; - -use super::SqliteStore; - -#[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 freed_document_ids: Option, - pub next_document_id: u32, -} - -impl IdAssigner { - pub fn new(used_ids: Option) -> Self { - let mut assigner = IdAssigner { - freed_document_ids: None, - next_document_id: 0, - }; - if let Some(used_ids) = used_ids { - if let Some(max) = used_ids.max() { - assigner.next_document_id = max + 1; - let mut freed_ids = - RoaringBitmap::from_sorted_iter(0..assigner.next_document_id).unwrap(); - freed_ids ^= used_ids; - if !freed_ids.is_empty() { - assigner.freed_document_ids = Some(freed_ids); - } - } - } - - assigner - } - - pub fn assign_document_id(&mut self) -> u32 { - if let Some(freed_ids) = &mut self.freed_document_ids { - let id = freed_ids.min().unwrap(); - freed_ids.remove(id); - if freed_ids.is_empty() { - self.freed_document_ids = None; - } - id - } else { - let id = self.next_document_id; - self.next_document_id += 1; - id - } - } -} - -impl SqliteStore { - pub(crate) async fn assign_document_id( - &self, - account_id: u32, - collection: impl Into + Sync + Send, - ) -> crate::Result { - let key = IdCacheKey::new(account_id, collection.into()); - 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(crate) async fn build_id_assigner(&self, key: IdCacheKey) -> crate::Result<()> { - // Obtain used ids - let used_ids = self - .get_bitmap(BitmapKey::document_ids(key.account_id, key.collection)) - .await?; - - let id_assigner = self.id_assigner.clone(); - let mut id_assigner = id_assigner.lock(); - // Make sure id assigner was not added by another thread - if id_assigner.get_mut(&key).is_none() { - id_assigner.insert(key, IdAssigner::new(used_ids)); - } - - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use roaring::RoaringBitmap; - - use super::IdAssigner; - - #[test] - fn id_assigner() { - let mut assigner = IdAssigner::new(None); - 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(), - ); - 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/crates/store/src/backend/sqlite/main.rs b/crates/store/src/backend/sqlite/main.rs index b308a408..e4470afe 100644 --- a/crates/store/src/backend/sqlite/main.rs +++ b/crates/store/src/backend/sqlite/main.rs @@ -21,17 +21,13 @@ * for more details. */ -use std::sync::Arc; - -use lru_cache::LruCache; -use parking_lot::Mutex; use r2d2::Pool; use tokio::sync::oneshot; use utils::{config::Config, UnwrapFailure}; use crate::{ - SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, - SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_VALUES, + SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, SUBSPACE_INDEXES, + SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES, }; use super::{pool::SqliteConnectionManager, SqliteStore}; @@ -67,9 +63,6 @@ impl SqliteStore { .map_err(|err| { crate::Error::InternalError(format!("Failed to build worker pool: {}", err)) })?, - id_assigner: Arc::new(Mutex::new(LruCache::new( - config.property_or_static("store.db.cache.size", "1000")?, - ))), }; db.create_tables()?; Ok(db) @@ -81,7 +74,7 @@ impl SqliteStore { for table in [ SUBSPACE_VALUES, SUBSPACE_LOGS, - SUBSPACE_ACLS, + SUBSPACE_INDEX_VALUES, SUBSPACE_BLOB_DATA, ] { let table = char::from(table); diff --git a/crates/store/src/backend/sqlite/mod.rs b/crates/store/src/backend/sqlite/mod.rs index f11f3ab0..0d25582f 100644 --- a/crates/store/src/backend/sqlite/mod.rs +++ b/crates/store/src/backend/sqlite/mod.rs @@ -21,33 +21,17 @@ * for more details. */ -use std::sync::Arc; - -use lru_cache::LruCache; -use parking_lot::Mutex; use r2d2::Pool; -use crate::U64_LEN; - -use self::{ - id_assign::{IdAssigner, IdCacheKey}, - pool::SqliteConnectionManager, -}; +use self::pool::SqliteConnectionManager; pub mod blob; -pub mod id_assign; pub mod main; pub mod pool; pub mod purge; pub mod read; pub mod write; -const WORD_SIZE_BITS: u32 = (WORD_SIZE * 8) as u32; -const WORD_SIZE: usize = U64_LEN; -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 From for crate::Error { fn from(err: r2d2::Error) -> Self { Self::InternalError(format!("Connection pool error: {}", err)) @@ -68,6 +52,5 @@ impl From for crate::Error { pub struct SqliteStore { pub(crate) conn_pool: Pool, - pub(crate) id_assigner: Arc>>, pub(crate) worker_pool: rayon::ThreadPool, } diff --git a/crates/store/src/backend/sqlite/read.rs b/crates/store/src/backend/sqlite/read.rs index 8614d03c..4648b90a 100644 --- a/crates/store/src/backend/sqlite/read.rs +++ b/crates/store/src/backend/sqlite/read.rs @@ -27,13 +27,14 @@ use rusqlite::OptionalExtension; use crate::{ query::{self, Operator}, write::{ + bitmap::{BITS_PER_BLOCK_S, WORDS_PER_BLOCK_S, WORD_SIZE_BITS_S}, key::{DeserializeBigEndian, KeySerializer}, BitmapClass, ValueClass, }, BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, IterateParams, Key, ValueKey, U32_LEN, }; -use super::{SqliteStore, BITS_PER_BLOCK, WORDS_PER_BLOCK, WORD_SIZE_BITS}; +use super::SqliteStore; impl SqliteStore { pub(crate) async fn get_value(&self, key: impl Key) -> crate::Result> @@ -76,22 +77,22 @@ impl SqliteStore { if key.len() == key_len { let block_num = key.deserialize_be_u32(key.len() - U32_LEN)?; - for word_num in 0..WORDS_PER_BLOCK { + for word_num in 0..WORDS_PER_BLOCK_S { 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, + block_num * BITS_PER_BLOCK_S + word_num * WORD_SIZE_BITS_S + ..(block_num * BITS_PER_BLOCK_S + word_num * WORD_SIZE_BITS_S) + + WORD_SIZE_BITS_S, ); } mut word => { while word != 0 { let trailing_zeros = word.trailing_zeros(); bm.insert( - block_num * BITS_PER_BLOCK - + word_num * WORD_SIZE_BITS + block_num * BITS_PER_BLOCK_S + + word_num * WORD_SIZE_BITS_S + trailing_zeros, ); word ^= 1 << trailing_zeros; @@ -314,7 +315,7 @@ impl SqliteStore { // Values let mut has_errors = false; - for table in [crate::SUBSPACE_VALUES, crate::SUBSPACE_ACLS, crate::SUBSPACE_COUNTERS, crate::SUBSPACE_BLOB_DATA] { + for table in [crate::SUBSPACE_VALUES, crate::SUBSPACE_INDEX_VALUES, crate::SUBSPACE_COUNTERS, crate::SUBSPACE_BLOB_DATA] { let table = char::from(table); let mut query = conn.prepare_cached(&format!("SELECT k, v FROM {table}")).unwrap(); let mut rows = query.query([]).unwrap(); @@ -400,6 +401,5 @@ impl SqliteStore { Ok(()) }).await.unwrap(); - self.id_assigner.lock().clear(); } } diff --git a/crates/store/src/backend/sqlite/write.rs b/crates/store/src/backend/sqlite/write.rs index 4c759045..54e67d4f 100644 --- a/crates/store/src/backend/sqlite/write.rs +++ b/crates/store/src/backend/sqlite/write.rs @@ -24,11 +24,14 @@ use rusqlite::{params, OptionalExtension, TransactionBehavior}; use crate::{ - write::{Batch, Operation, ValueClass, ValueOp}, + write::{ + bitmap::{BITS_MASK_S, BITS_PER_BLOCK_S}, + Batch, Operation, ValueOp, + }, BitmapKey, BlobKey, IndexKey, Key, LogKey, ValueKey, }; -use super::{SqliteStore, BITS_MASK, BITS_PER_BLOCK}; +use super::SqliteStore; const INSERT_QUERIES: &[&str] = &[ "INSERT INTO b (z, a) VALUES (?, ?)", @@ -114,8 +117,8 @@ impl SqliteStore { document_id: document_id_, } => { document_id = *document_id_; - bitmap_block_num = document_id / BITS_PER_BLOCK; - let index = document_id & BITS_MASK; + bitmap_block_num = document_id / BITS_PER_BLOCK_S; + let index = document_id & BITS_MASK_S; 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; @@ -149,23 +152,19 @@ impl SqliteStore { collection, document_id, class, - } - .serialize(false); + }; + let table = char::from(key.subspace()); + let key = key.serialize(false); if let ValueOp::Set(value) = op { - trx.prepare_cached(if !matches!(class, ValueClass::Acl(_)) { - "INSERT OR REPLACE INTO v (k, v) VALUES (?, ?)" - } else { - "INSERT OR REPLACE INTO a (k, v) VALUES (?, ?)" - })? + trx.prepare_cached(&format!( + "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", + table + ))? .execute([&key, value])?; } else { - trx.prepare_cached(if !matches!(class, ValueClass::Acl(_)) { - "DELETE FROM v WHERE k = ?" - } else { - "DELETE FROM a WHERE k = ?" - })? - .execute([&key])?; + trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table))? + .execute([&key])?; } } Operation::Index { field, key, set } => { @@ -249,11 +248,12 @@ impl SqliteStore { collection, document_id, class, - } - .serialize(false); + }; + let table = char::from(key.subspace()); + let key = key.serialize(false); let matches = trx - .prepare_cached("SELECT v FROM v WHERE k = ?")? + .prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table))? .query_row([&key], |row| { Ok(assert_value.matches(row.get_ref(0)?.as_bytes()?)) }) @@ -274,8 +274,8 @@ impl SqliteStore { #[cfg(feature = "test_mode")] pub(crate) async fn destroy(&self) { use crate::{ - SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, - SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_VALUES, + SUBSPACE_BITMAPS, SUBSPACE_BLOBS, SUBSPACE_BLOB_DATA, SUBSPACE_COUNTERS, + SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES, }; let conn = self.conn_pool.get().unwrap(); @@ -285,7 +285,7 @@ impl SqliteStore { SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_BLOBS, - SUBSPACE_ACLS, + SUBSPACE_INDEX_VALUES, SUBSPACE_COUNTERS, SUBSPACE_BLOB_DATA, ] { diff --git a/crates/store/src/dispatch.rs b/crates/store/src/dispatch.rs index a5a69394..e035f6eb 100644 --- a/crates/store/src/dispatch.rs +++ b/crates/store/src/dispatch.rs @@ -36,24 +36,6 @@ use crate::{ }; impl Store { - /*pub async fn assign_change_id(&self, account_id: u32) -> crate::Result { - match self { - Self::SQLite(store) => store.assign_change_id(account_id).await, - Self::FoundationDb(store) => store.assign_change_id(account_id).await, - } - }*/ - - pub async fn assign_document_id( - &self, - account_id: u32, - collection: impl Into + Sync + Send, - ) -> crate::Result { - match self { - Self::SQLite(store) => store.assign_document_id(account_id, collection).await, - Self::FoundationDb(store) => store.assign_document_id(account_id, collection).await, - } - } - pub async fn get_value(&self, key: impl Key) -> crate::Result> where U: Deserialize + 'static, @@ -61,6 +43,7 @@ impl Store { match self { Self::SQLite(store) => store.get_value(key).await, Self::FoundationDb(store) => store.get_value(key).await, + Self::PostgreSQL(store) => store.get_value(key).await, } } @@ -84,6 +67,7 @@ impl Store { match self { Self::SQLite(store) => store.get_bitmap(key).await, Self::FoundationDb(store) => store.get_bitmap(key).await, + Self::PostgreSQL(store) => store.get_bitmap(key).await, } } @@ -128,6 +112,11 @@ impl Store { .range_to_bitmap(account_id, collection, field, value, op) .await } + Self::PostgreSQL(store) => { + store + .range_to_bitmap(account_id, collection, field, value, op) + .await + } } } @@ -150,6 +139,11 @@ impl Store { .sort_index(account_id, collection, field, ascending, cb) .await } + Self::PostgreSQL(store) => { + store + .sort_index(account_id, collection, field, ascending, cb) + .await + } } } @@ -161,6 +155,7 @@ impl Store { match self { Self::SQLite(store) => store.iterate(params, cb).await, Self::FoundationDb(store) => store.iterate(params, cb).await, + Self::PostgreSQL(store) => store.iterate(params, cb).await, } } @@ -171,6 +166,7 @@ impl Store { match self { Self::SQLite(store) => store.get_counter(key).await, Self::FoundationDb(store) => store.get_counter(key).await, + Self::PostgreSQL(store) => store.get_counter(key).await, } } @@ -178,6 +174,7 @@ impl Store { match self { Self::SQLite(store) => store.write(batch).await, Self::FoundationDb(store) => store.write(batch).await, + Self::PostgreSQL(store) => store.write(batch).await, } } @@ -185,12 +182,14 @@ impl Store { match self { Self::SQLite(store) => store.purge_bitmaps().await, Self::FoundationDb(store) => store.purge_bitmaps().await, + Self::PostgreSQL(store) => store.purge_bitmaps().await, } } pub async fn purge_account(&self, account_id: u32) -> crate::Result<()> { match self { Self::SQLite(store) => store.purge_account(account_id).await, Self::FoundationDb(store) => store.purge_account(account_id).await, + Self::PostgreSQL(store) => store.purge_account(account_id).await, } } @@ -198,6 +197,7 @@ impl Store { match self { Self::SQLite(store) => store.get_blob(key, range).await, Self::FoundationDb(store) => store.get_blob(key, range).await, + Self::PostgreSQL(store) => store.get_blob(key, range).await, } } @@ -205,6 +205,7 @@ impl Store { match self { Self::SQLite(store) => store.put_blob(key, data).await, Self::FoundationDb(store) => store.put_blob(key, data).await, + Self::PostgreSQL(store) => store.put_blob(key, data).await, } } @@ -212,6 +213,7 @@ impl Store { match self { Self::SQLite(store) => store.delete_blob(key).await, Self::FoundationDb(store) => store.delete_blob(key).await, + Self::PostgreSQL(store) => store.delete_blob(key).await, } } @@ -220,6 +222,7 @@ impl Store { match self { Self::SQLite(store) => store.destroy().await, Self::FoundationDb(store) => store.destroy().await, + Self::PostgreSQL(store) => store.destroy().await, } } @@ -285,6 +288,7 @@ impl Store { match self { Self::SQLite(store) => store.assert_is_empty().await, Self::FoundationDb(store) => store.assert_is_empty().await, + Self::PostgreSQL(store) => store.assert_is_empty().await, } } } @@ -296,6 +300,7 @@ impl BlobStore { Self::S3(store) => store.get_blob(key, range).await, Self::Sqlite(store) => store.get_blob(key, range).await, Self::FoundationDb(store) => store.get_blob(key, range).await, + Self::PostgreSQL(store) => store.get_blob(key, range).await, } } @@ -305,6 +310,7 @@ impl BlobStore { Self::S3(store) => store.put_blob(key, data).await, Self::Sqlite(store) => store.put_blob(key, data).await, Self::FoundationDb(store) => store.put_blob(key, data).await, + Self::PostgreSQL(store) => store.put_blob(key, data).await, } } @@ -314,6 +320,7 @@ impl BlobStore { Self::S3(store) => store.delete_blob(key).await, Self::Sqlite(store) => store.delete_blob(key).await, Self::FoundationDb(store) => store.delete_blob(key).await, + Self::PostgreSQL(store) => store.delete_blob(key).await, } } } diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 2d848be3..91b75575 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -30,7 +30,9 @@ pub mod query; pub mod write; pub use ahash; -use backend::{foundationdb::FdbStore, fs::FsStore, s3::S3Store, sqlite::SqliteStore}; +use backend::{ + foundationdb::FdbStore, fs::FsStore, postgres::PostgresStore, s3::S3Store, sqlite::SqliteStore, +}; pub use blake3; pub use parking_lot; pub use rand; @@ -154,7 +156,7 @@ pub const SUBSPACE_LOGS: u8 = b'l'; pub const SUBSPACE_INDEXES: u8 = b'i'; pub const SUBSPACE_BLOBS: u8 = b'o'; pub const SUBSPACE_BLOB_DATA: u8 = b't'; -pub const SUBSPACE_ACLS: u8 = b'a'; +pub const SUBSPACE_INDEX_VALUES: u8 = b'a'; pub const SUBSPACE_COUNTERS: u8 = b'c'; pub struct IterateParams { @@ -169,6 +171,7 @@ pub struct IterateParams { pub enum Store { SQLite(Arc), FoundationDb(Arc), + PostgreSQL(Arc), } #[derive(Clone)] @@ -177,6 +180,7 @@ pub enum BlobStore { S3(Arc), Sqlite(Arc), FoundationDb(Arc), + PostgreSQL(Arc), } #[derive(Clone)] @@ -196,6 +200,12 @@ impl From for Store { } } +impl From for Store { + fn from(store: PostgresStore) -> Self { + Self::PostgreSQL(Arc::new(store)) + } +} + impl From for BlobStore { fn from(store: FsStore) -> Self { Self::Fs(Arc::new(store)) @@ -219,6 +229,7 @@ impl From for BlobStore { match store { Store::SQLite(store) => Self::Sqlite(store), Store::FoundationDb(store) => Self::FoundationDb(store), + Store::PostgreSQL(store) => Self::PostgreSQL(store), } } } diff --git a/crates/store/src/query/acl.rs b/crates/store/src/query/acl.rs index 56264fef..7fc745c2 100644 --- a/crates/store/src/query/acl.rs +++ b/crates/store/src/query/acl.rs @@ -156,11 +156,11 @@ impl Store { impl Deserialize for AclItem { fn deserialize(bytes: &[u8]) -> crate::Result { Ok(AclItem { - to_account_id: bytes.deserialize_be_u32(U32_LEN)?, + to_account_id: bytes.deserialize_be_u32(U32_LEN + 1)?, to_collection: *bytes - .get(U32_LEN * 2) + .get((U32_LEN * 2) + 1) .ok_or_else(|| Error::InternalError(format!("Corrupted acl key {bytes:?}")))?, - to_document_id: bytes.deserialize_be_u32((U32_LEN * 2) + 1)?, + to_document_id: bytes.deserialize_be_u32((U32_LEN * 2) + 2)?, permissions: 0, }) } diff --git a/crates/store/src/write/assign_id.rs b/crates/store/src/write/assign_id.rs new file mode 100644 index 00000000..22a6b923 --- /dev/null +++ b/crates/store/src/write/assign_id.rs @@ -0,0 +1,170 @@ +/* + * Copyright (c) 2023 Stalwart Labs Ltd. + * + * This file is part of the Stalwart Mail Server. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * in the LICENSE file at the top-level directory of this distribution. + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + * + * You can be released from the requirements of the AGPLv3 license by + * purchasing a commercial license. Please contact licensing@stalw.art + * for more details. +*/ + +use std::time::Instant; + +use crate::{ + backend::ID_ASSIGNMENT_EXPIRY, write::key::DeserializeBigEndian, Deserialize, IterateParams, + Serialize, Store, ValueKey, U32_LEN, +}; +use ahash::AHashMap; +use rand::Rng; +use roaring::RoaringBitmap; + +use crate::{write::now, BitmapKey}; + +use super::{BatchBuilder, ValueClass, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME}; + +impl Store { + pub async fn assign_document_id( + &self, + account_id: u32, + collection: impl Into + Sync + Send, + ) -> crate::Result { + let start = Instant::now(); + let mut retry_count = 0; + let collection = collection.into(); + + loop { + // First try to reuse an expired assigned id + let mut reserved_ids = RoaringBitmap::new(); + let mut expired_ids = AHashMap::new(); + { + let from_key = ValueKey { + account_id, + collection, + document_id: 0, + class: ValueClass::ReservedId, + }; + let to_key = ValueKey { + account_id, + collection, + document_id: u32::MAX, + class: ValueClass::ReservedId, + }; + + let expired_timestamp = now(); + self.iterate( + IterateParams::new(from_key, to_key).ascending(), + |key, value| { + let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?; + let ttl = u64::deserialize(value)?; + + if ttl <= expired_timestamp { + // Found an expired id, reuse it + expired_ids.insert(document_id, ttl); + } else { + // Keep track of all reserved ids + reserved_ids.insert(document_id); + } + + Ok(true) + }, + ) + .await?; + } + + // Prepare the patch the id + let mut batch = BatchBuilder::new(); + batch + .with_account_id(account_id) + .with_collection(collection); + + let document_id = if !expired_ids.is_empty() { + // Obtain a random id from the expired ids + let pos = if expired_ids.len() > 1 { + rand::thread_rng().gen_range(0..expired_ids.len()) + } else { + 0 + }; + let (document_id, expiry) = expired_ids.into_iter().nth(pos).unwrap(); + + batch + .update_document(document_id) + .assert_value(ValueClass::ReservedId, expiry); + document_id + } else { + // Obtain documentIds + let document_ids = if let Some(document_ids) = self + .get_bitmap(BitmapKey::document_ids(account_id, collection)) + .await? + { + if !reserved_ids.is_empty() { + document_ids | reserved_ids + } else { + document_ids + } + } else { + reserved_ids + }; + + let document_id = if retry_count == 0 { + // Find the next available id + (0..(document_ids.len() + 1) as u32) + .find(|&x| !document_ids.contains(x)) + .unwrap() + } else { + // High contention, pick a random id + const RAND_IDS: usize = 10; + let mut available_ids = Vec::with_capacity(RAND_IDS); + for id in 0..(document_ids.len() as u32 + RAND_IDS as u32) { + if !document_ids.contains(id) { + available_ids.push(id); + if available_ids.len() == RAND_IDS { + break; + } + } + } + available_ids[rand::thread_rng().gen_range(0..available_ids.len())] + }; + + batch + .update_document(document_id) + .assert_value(ValueClass::ReservedId, ()); + document_id + }; + + #[cfg(not(feature = "test_mode"))] + let expired_timestamp = now() + ID_ASSIGNMENT_EXPIRY; + #[cfg(feature = "test_mode")] + let expired_timestamp = + now() + ID_ASSIGNMENT_EXPIRY.load(std::sync::atomic::Ordering::Relaxed); + + batch.set(ValueClass::ReservedId, expired_timestamp.serialize()); + + match self.write(batch.build()).await { + Ok(_) => { + return Ok(document_id); + } + Err(crate::Error::AssertValueFailed) + if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => + { + // Retry + retry_count += 1; + continue; + } + Err(err) => return Err(err), + } + } + } +} diff --git a/crates/store/src/write/batch.rs b/crates/store/src/write/batch.rs index ecd4a11d..cfa78e43 100644 --- a/crates/store/src/write/batch.rs +++ b/crates/store/src/write/batch.rs @@ -31,7 +31,9 @@ use super::{ impl BatchBuilder { pub fn new() -> Self { - Self { ops: Vec::new() } + Self { + ops: Vec::with_capacity(16), + } } pub fn with_account_id(&mut self, account_id: u32) -> &mut Self { @@ -49,18 +51,18 @@ impl BatchBuilder { pub fn create_document(&mut self, document_id: u32) -> &mut Self { self.ops.push(Operation::DocumentId { document_id }); - // Remove reserved id - self.ops.push(Operation::Index { - field: u8::MAX, - key: vec![], - set: false, - }); - // Add document id self.ops.push(Operation::Bitmap { class: BitmapClass::DocumentIds, set: true, }); + + // Remove reserved id + self.ops.push(Operation::Value { + class: ValueClass::ReservedId, + op: ValueOp::Clear, + }); + self } diff --git a/crates/store/src/backend/foundationdb/bitmap.rs b/crates/store/src/write/bitmap.rs similarity index 67% rename from crates/store/src/backend/foundationdb/bitmap.rs rename to crates/store/src/write/bitmap.rs index 8baefab3..a292df7f 100644 --- a/crates/store/src/backend/foundationdb/bitmap.rs +++ b/crates/store/src/write/bitmap.rs @@ -24,41 +24,49 @@ use ahash::AHashSet; use roaring::RoaringBitmap; -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; -const BITS_MASK: u32 = BITS_PER_BLOCK - 1; +use crate::U64_LEN; + +pub(crate) const WORD_SIZE_BITS_L: u32 = (WORD_SIZE_L * 8) as u32; +pub(crate) const WORD_SIZE_L: usize = std::mem::size_of::(); +pub(crate) const WORDS_PER_BLOCK_L: u32 = 8; +pub(crate) const BITS_PER_BLOCK_L: u32 = WORD_SIZE_BITS_L * WORDS_PER_BLOCK_L; +pub(crate) const BITS_MASK_L: u32 = BITS_PER_BLOCK_L - 1; + +pub(crate) const WORD_SIZE_BITS_S: u32 = (WORD_SIZE_S * 8) as u32; +pub(crate) const WORD_SIZE_S: usize = U64_LEN; +pub(crate) const WORDS_PER_BLOCK_S: u32 = 16; +pub(crate) const BITS_PER_BLOCK_S: u32 = WORD_SIZE_BITS_S * WORDS_PER_BLOCK_S; +pub(crate) const BITS_MASK_S: u32 = BITS_PER_BLOCK_S - 1; pub struct DenseBitmap { - pub bitmap: [u8; WORD_SIZE * WORDS_PER_BLOCK as usize], + pub bitmap: [u8; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], } impl DenseBitmap { pub fn empty() -> Self { Self { - bitmap: [0; WORD_SIZE * WORDS_PER_BLOCK as usize], + bitmap: [0; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], } } pub fn full() -> Self { Self { - bitmap: [u8::MAX; WORD_SIZE * WORDS_PER_BLOCK as usize], + bitmap: [u8::MAX; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], } } pub fn set(&mut self, index: u32) { - let index = index & BITS_MASK; + let index = index & BITS_MASK_L; self.bitmap[(index / 8) as usize] |= 1 << (index & 7); } pub fn clear(&mut self, index: u32) { - let index = index & BITS_MASK; + let index = index & BITS_MASK_L; self.bitmap[(index / 8) as usize] &= !(1 << (index & 7)); } pub fn block_num(index: u32) -> u32 { - index / BITS_PER_BLOCK + index / BITS_PER_BLOCK_L } } @@ -83,7 +91,7 @@ pub fn next_available_index( } } - let id = (block_num * BITS_PER_BLOCK) + ((byte_pos * 8) + index) as u32; + let id = (block_num * BITS_PER_BLOCK_L) + ((byte_pos * 8) + index) as u32; if !reserved_ids.contains(&id) { return Some(id); } else if index < 7 { @@ -101,16 +109,16 @@ pub fn next_available_index( impl DeserializeBlock for RoaringBitmap { fn deserialize_block(&mut self, bytes: &[u8], block_num: u32) { - debug_assert_eq!(bytes.len(), WORD_SIZE * WORDS_PER_BLOCK as usize); + debug_assert_eq!(bytes.len(), WORD_SIZE_L * WORDS_PER_BLOCK_L as usize); - self.deserialize_word(&bytes[..WORD_SIZE], block_num, 0); - self.deserialize_word(&bytes[WORD_SIZE..WORD_SIZE * 2], block_num, 1); - self.deserialize_word(&bytes[WORD_SIZE * 2..WORD_SIZE * 3], block_num, 2); - self.deserialize_word(&bytes[WORD_SIZE * 3..WORD_SIZE * 4], block_num, 3); - self.deserialize_word(&bytes[WORD_SIZE * 4..WORD_SIZE * 5], block_num, 4); - self.deserialize_word(&bytes[WORD_SIZE * 5..WORD_SIZE * 6], block_num, 5); - self.deserialize_word(&bytes[WORD_SIZE * 6..WORD_SIZE * 7], block_num, 6); - self.deserialize_word(&bytes[WORD_SIZE * 7..], block_num, 7); + self.deserialize_word(&bytes[..WORD_SIZE_L], block_num, 0); + self.deserialize_word(&bytes[WORD_SIZE_L..WORD_SIZE_L * 2], block_num, 1); + self.deserialize_word(&bytes[WORD_SIZE_L * 2..WORD_SIZE_L * 3], block_num, 2); + self.deserialize_word(&bytes[WORD_SIZE_L * 3..WORD_SIZE_L * 4], block_num, 3); + self.deserialize_word(&bytes[WORD_SIZE_L * 4..WORD_SIZE_L * 5], block_num, 4); + self.deserialize_word(&bytes[WORD_SIZE_L * 5..WORD_SIZE_L * 6], block_num, 5); + self.deserialize_word(&bytes[WORD_SIZE_L * 6..WORD_SIZE_L * 7], block_num, 6); + self.deserialize_word(&bytes[WORD_SIZE_L * 7..], block_num, 7); } #[inline(always)] @@ -119,15 +127,16 @@ impl DeserializeBlock for RoaringBitmap { 0 => (), u128::MAX => { self.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, + block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L + ..(block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L) + + WORD_SIZE_BITS_L, ); } mut word => { while word != 0 { let trailing_zeros = word.trailing_zeros(); self.insert( - block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS + trailing_zeros, + block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L + trailing_zeros, ); word ^= 1 << trailing_zeros; } @@ -143,9 +152,7 @@ mod tests { use ahash::AHashSet; use roaring::RoaringBitmap; - use crate::backend::foundationdb::bitmap::{ - next_available_index, DenseBitmap, DeserializeBlock, BITS_PER_BLOCK, - }; + use super::*; #[test] fn serialize_bitmap_block() { @@ -155,7 +162,7 @@ mod tests { for item in range { bitmap.insert(item); blocks - .entry(item / BITS_PER_BLOCK) + .entry(item / BITS_PER_BLOCK_L) .or_insert_with(DenseBitmap::empty) .set(item); } diff --git a/crates/store/src/write/key.rs b/crates/store/src/write/key.rs index e4309ac5..b9c950a5 100644 --- a/crates/store/src/write/key.rs +++ b/crates/store/src/write/key.rs @@ -26,8 +26,8 @@ use utils::codec::leb128::Leb128_; use crate::{ BitmapKey, BlobHash, BlobKey, IndexKey, IndexKeyPrefix, Key, LogKey, ValueKey, BLOB_HASH_LEN, - SUBSPACE_ACLS, SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_VALUES, U32_LEN, - U64_LEN, + SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_INDEX_VALUES, SUBSPACE_LOGS, SUBSPACE_VALUES, + U32_LEN, U64_LEN, }; use super::{BitmapClass, BlobOp, TagValue, ValueClass}; @@ -207,10 +207,13 @@ impl Key for LogKey { impl + Sync + Send> Key for ValueKey { fn subspace(&self) -> u8 { - if !matches!(self.class.as_ref(), ValueClass::Acl(_)) { + if !matches!( + self.class.as_ref(), + ValueClass::Acl(_) | ValueClass::ReservedId + ) { SUBSPACE_VALUES } else { - SUBSPACE_ACLS + SUBSPACE_INDEX_VALUES } } @@ -226,11 +229,12 @@ impl + Sync + Send> Key for ValueKey { .write_leb128(self.document_id) .write(*field), ValueClass::Acl(grant_account_id) => if include_subspace { - KeySerializer::new(U32_LEN * 3 + 2).write(crate::SUBSPACE_ACLS) + KeySerializer::new(U32_LEN * 3 + 3).write(crate::SUBSPACE_INDEX_VALUES) } else { - KeySerializer::new(U32_LEN * 3 + 1) + KeySerializer::new(U32_LEN * 3 + 2) } .write(*grant_account_id) + .write(0u8) .write(self.account_id) .write(self.collection) .write(self.document_id), @@ -250,6 +254,15 @@ impl + Sync + Send> Key for ValueKey { .write(self.collection) .write_leb128(self.document_id) .write(u8::MAX), + ValueClass::ReservedId => if include_subspace { + KeySerializer::new(U32_LEN * 2 + 2).write(crate::SUBSPACE_INDEX_VALUES) + } else { + KeySerializer::new(U32_LEN * 2 + 1) + } + .write(self.account_id) + .write(1u8) + .write(self.collection) + .write(self.document_id), } .finalize() } diff --git a/crates/store/src/write/mod.rs b/crates/store/src/write/mod.rs index 7d21b4b1..cc6c8a2b 100644 --- a/crates/store/src/write/mod.rs +++ b/crates/store/src/write/mod.rs @@ -21,7 +21,12 @@ * for more details. */ -use std::{collections::HashSet, hash::Hash, slice::Iter, time::SystemTime}; +use std::{ + collections::HashSet, + hash::Hash, + slice::Iter, + time::{Duration, SystemTime}, +}; use nlp::tokenizers::word::WordTokenizer; use utils::codec::leb128::{Leb128Iterator, Leb128Vec}; @@ -33,17 +38,32 @@ use crate::{ use self::assert::AssertValue; pub mod assert; +pub mod assign_id; pub mod batch; +pub mod bitmap; pub mod blob; pub mod hash; pub mod key; pub mod log; +#[cfg(not(feature = "test_mode"))] +pub(crate) const ID_ASSIGNMENT_EXPIRY: u64 = 60 * 60; // seconds +#[cfg(not(feature = "test_mode"))] +pub(crate) const MAX_COMMIT_ATTEMPTS: u32 = 10; +#[cfg(not(feature = "test_mode"))] +pub(crate) const MAX_COMMIT_TIME: Duration = Duration::from_secs(10); + +#[cfg(feature = "test_mode")] +pub(crate) const MAX_COMMIT_ATTEMPTS: u32 = 1000; +#[cfg(feature = "test_mode")] +pub(crate) const MAX_COMMIT_TIME: Duration = Duration::from_secs(3600); + pub const F_VALUE: u32 = 1 << 0; pub const F_INDEX: u32 = 1 << 1; pub const F_BITMAP: u32 = 1 << 2; pub const F_CLEAR: u32 = 1 << 3; +#[derive(Debug)] pub struct Batch { pub ops: Vec, } @@ -119,6 +139,7 @@ pub enum ValueClass { Acl(u32), Named(Vec), TermIndex, + ReservedId, } #[derive(Debug, PartialEq, Eq, Hash, Default)] diff --git a/crates/utils/Cargo.toml b/crates/utils/Cargo.toml index 840e839a..fbc6e432 100644 --- a/crates/utils/Cargo.toml +++ b/crates/utils/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" resolver = "2" [dependencies] -rustls = "0.21.0" +rustls = { version = "0.21", features = ["tls12", "dangerous_configuration"]} rustls-pemfile = "1.0" tokio = { version = "1.23", features = ["net", "macros"] } tokio-rustls = { version = "0.24.0"} @@ -16,14 +16,16 @@ smtp-proto = { git = "https://github.com/stalwartlabs/smtp-proto" } mail-send = { git = "https://github.com/stalwartlabs/mail-send", default-features = false, features = ["cram-md5", "skip-ehlo"] } tracing-subscriber = { version = "0.3", features = ["env-filter"] } tracing-appender = "0.2" -tracing-opentelemetry = "0.21.0" -opentelemetry = { version = "0.20.0", features = ["rt-tokio"] } -opentelemetry-otlp = { version = "0.13.0", features = ["http-proto", "reqwest-client"] } -opentelemetry-semantic-conventions = { version = "0.12.0" } +tracing-opentelemetry = "0.22.0" +opentelemetry = { version = "0.21.0" } +opentelemetry_sdk = { version = "0.21.0", features = ["rt-tokio"] } +opentelemetry-otlp = { version = "0.14.0", features = ["http-proto", "reqwest-client"] } +opentelemetry-semantic-conventions = { version = "0.13.0" } dashmap = "5.4" ahash = { version = "0.8" } chrono = "0.4" rand = "0.8.5" +webpki-roots = { version = "0.25.2"} [target.'cfg(unix)'.dependencies] privdrop = "0.5.3" diff --git a/crates/utils/src/lib.rs b/crates/utils/src/lib.rs index 7a99997d..1ed53507 100644 --- a/crates/utils/src/lib.rs +++ b/crates/utils/src/lib.rs @@ -21,7 +21,7 @@ * for more details. */ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use config::Config; @@ -33,15 +33,17 @@ pub mod map; pub mod snowflake; pub mod suffixlist; -use opentelemetry::{ - sdk::{ - trace::{self, Sampler}, - Resource, - }, - KeyValue, -}; +use opentelemetry::KeyValue; use opentelemetry_otlp::WithExportConfig; +use opentelemetry_sdk::{ + trace::{self, Sampler}, + Resource, +}; use opentelemetry_semantic_conventions::resource::{SERVICE_NAME, SERVICE_VERSION}; +use rustls::{ + client::{ServerCertVerified, ServerCertVerifier}, + Certificate, ClientConfig, OwnedTrustAnchor, RootCertStore, ServerName, +}; use tracing_appender::non_blocking::WorkerGuard; use tracing_subscriber::{prelude::__tracing_subscriber_SubscriberExt, EnvFilter}; @@ -171,7 +173,7 @@ pub fn enable_tracing(config: &Config, message: &str) -> config::Result ClientConfig { + let config = ClientConfig::builder().with_safe_defaults(); + + if !allow_invalid_certs { + let mut root_cert_store = RootCertStore::empty(); + + root_cert_store.add_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.iter().map(|ta| { + OwnedTrustAnchor::from_subject_spki_name_constraints( + ta.subject, + ta.spki, + ta.name_constraints, + ) + })); + config + .with_root_certificates(root_cert_store) + .with_no_client_auth() + } else { + config + .with_custom_certificate_verifier(Arc::new(DummyVerifier {})) + .with_no_client_auth() + } +} + +struct DummyVerifier; + +impl ServerCertVerifier for DummyVerifier { + fn verify_server_cert( + &self, + _e: &Certificate, + _i: &[Certificate], + _sn: &ServerName, + _sc: &mut dyn Iterator, + _o: &[u8], + _n: std::time::SystemTime, + ) -> Result { + Ok(ServerCertVerified::assertion()) + } +} diff --git a/tests/src/imap/mod.rs b/tests/src/imap/mod.rs index 0060f251..dc9aef7f 100644 --- a/tests/src/imap/mod.rs +++ b/tests/src/imap/mod.rs @@ -136,8 +136,13 @@ allow-invalid-certs = true future-release = [ { if = "authenticated-as", ne = "", then = "99999999d"}, { else = false } ] -[store] -db.path = "{TMP}/sqlite.db" +[store.db] +#path = "{TMP}/sqlite.db" +host = "localhost" +port = 5432 +database = "stalwart" +user = "postgres" +password = "mysecretpassword" [store.blob] type = "local" diff --git a/tests/src/jmap/mod.rs b/tests/src/jmap/mod.rs index e0d62959..f279b906 100644 --- a/tests/src/jmap/mod.rs +++ b/tests/src/jmap/mod.rs @@ -134,8 +134,13 @@ allow-invalid-certs = true future-release = [ { if = "authenticated-as", ne = "", then = "99999999d"}, { else = false } ] -[store] -db.path = "{TMP}/sqlite.db" +[store.db] +path = "{TMP}/sqlite.db" +host = "localhost" +port = 5432 +database = "stalwart" +user = "postgres" +password = "mysecretpassword" [store.blob] type = "local" @@ -242,7 +247,7 @@ pub async fn jmap_tests() { let delete = true; let mut params = init_jmap_tests(delete).await; - //email_query::test(params.server.clone(), &mut params.client, delete).await; + email_query::test(params.server.clone(), &mut params.client, delete).await; email_get::test(params.server.clone(), &mut params.client).await; email_set::test(params.server.clone(), &mut params.client).await; email_parse::test(params.server.clone(), &mut params.client).await; diff --git a/tests/src/jmap/stress_test.rs b/tests/src/jmap/stress_test.rs index 92a9bb51..24b7c7c3 100644 --- a/tests/src/jmap/stress_test.rs +++ b/tests/src/jmap/stress_test.rs @@ -51,7 +51,10 @@ pub async fn test(server: Arc, mut client: Client) { async fn email_tests(server: Arc, client: Arc) { for pass in 0..NUM_PASSES { - println!("----------------- PASS {} -----------------", pass); + println!( + "----------------- EMAIL STRESS TEST {} -----------------", + pass + ); let mailboxes = Arc::new(vec![ client .mailbox_create("Stress 1", None::, Role::None) @@ -97,7 +100,12 @@ async fn email_tests(server: Arc, client: Arc) { .await .unwrap() .take_id(); - //println!("Inserted message {}.", message_id); + /*println!( + "Inserted message {}.", + Id::from_bytes(_message_id.as_bytes()) + .unwrap() + .document_id() + );*/ })); } @@ -110,7 +118,10 @@ async fn email_tests(server: Arc, client: Arc) { let ids = req.send_query_email().await.unwrap().take_ids(); if !ids.is_empty() { let message_id = &ids[rand::thread_rng().gen_range(0..ids.len())]; - //println!("Deleting message {}.", message_id); + /*println!( + "Deleting message {}.", + Id::from_bytes(message_id.as_bytes()).unwrap().document_id() + );*/ match client.email_destroy(message_id).await { Ok(_) => { break; @@ -121,7 +132,7 @@ async fn email_tests(server: Arc, client: Arc) { } SetErrorType::Forbidden => { // Concurrency issue, try again. - println!("Concurrent update, trying again."); + //println!("Concurrent update, trying again."); } _ => { panic!("Unexpected error: {:?}", err); @@ -169,7 +180,15 @@ async fn email_tests(server: Arc, client: Arc) { if new_mailbox_id != mailbox_id { /*println!( "Moving message {} from {} to {}.", - message_id, mailbox_id, new_mailbox_id + Id::from_bytes(message_id.as_bytes()) + .unwrap() + .document_id(), + Id::from_bytes(mailbox_id.as_bytes()) + .unwrap() + .document_id(), + Id::from_bytes(new_mailbox_id.as_bytes()) + .unwrap() + .document_id() );*/ let mut req = client.build(); req.set_email() @@ -271,6 +290,8 @@ async fn mailbox_tests(server: Arc, client: Arc) { ]); let mut futures = Vec::new(); + println!("----------------- MAILBOX STRESS TEST -----------------"); + for _ in 0..1000 { match rand::thread_rng().gen_range(0..=3) { 0 => { @@ -278,6 +299,7 @@ async fn mailbox_tests(server: Arc, client: Arc) { let client = client.clone(); let mailboxes = mailboxes.clone(); futures.push(tokio::spawn(async move { + //println!("Creating mailbox {}.", mailboxes[pos]); create_mailbox(&client, &mailboxes[pos]).await; })); } @@ -286,6 +308,7 @@ async fn mailbox_tests(server: Arc, client: Arc) { 1 => { let client = client.clone(); futures.push(tokio::spawn(async move { + //print!("Querying mailboxes..."); query_mailboxes(&client).await; })); } @@ -301,6 +324,7 @@ async fn mailbox_tests(server: Arc, client: Arc) { { let client = client.clone(); tokio::spawn(async move { + //println!("Deleting mailbox {}.", mailbox_id); delete_mailbox(&client, &mailbox_id).await; }); } @@ -318,6 +342,7 @@ async fn mailbox_tests(server: Arc, client: Arc) { if !ids.is_empty() { let id = ids.swap_remove(rand::thread_rng().gen_range(0..ids.len())); let sort_order = rand::thread_rng().gen_range(0..100); + //println!("Updating mailbox {}.", id); client.mailbox_update_sort_order(&id, sort_order).await.ok(); } })); diff --git a/tests/src/store/assign_id.rs b/tests/src/store/assign_id.rs index 0d589bcb..1e8ed0b2 100644 --- a/tests/src/store/assign_id.rs +++ b/tests/src/store/assign_id.rs @@ -31,53 +31,126 @@ use store::{write::BatchBuilder, Store}; pub async fn test(db: Store) { println!("Running Store ID assignment tests..."); - ID_ASSIGNMENT_EXPIRY.store(2, std::sync::atomic::Ordering::Relaxed); - + test_0(db.clone()).await; test_1(db.clone()).await; test_2(db.clone()).await; - test_3(db).await; + test_3(db.clone()).await; + test_4(db).await; ID_ASSIGNMENT_EXPIRY.store(60 * 60, std::sync::atomic::Ordering::Relaxed); } +async fn test_0(db: Store) { + // Test document id assignment + println!("Assigning 1000 ids concurrently..."); + ID_ASSIGNMENT_EXPIRY.store(10 * 60 * 60, std::sync::atomic::Ordering::Relaxed); + let mut handles = Vec::new(); + let mut assigned_ids = HashSet::new(); + + // Create 1000 ids concurrently + for _ in 0..1000 { + handles.push({ + let db = db.clone(); + tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await.unwrap() }) + }); + } + + for handle in handles { + let assigned_id = handle.await.unwrap(); + assert!( + assigned_ids.insert(assigned_id), + "already assigned or invalid: {assigned_id}" + ); + } + assert_eq!(assigned_ids.len(), 1000); + + db.destroy().await; +} + async fn test_1(db: Store) { // Test document id assignment + ID_ASSIGNMENT_EXPIRY.store(2, std::sync::atomic::Ordering::Relaxed); + println!("Assigning 100 ids concurrently and reassign after expiration..."); for wait_for_expiry in [true, false] { let mut handles = Vec::new(); - let mut expected_ids = HashSet::new(); + let mut assigned_ids = HashSet::new(); // Create 100 ids concurrently - for id in 0..100 { + for _ in 0..100 { handles.push({ let db = db.clone(); - tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await }) + tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await.unwrap() }) }); - expected_ids.insert(id); } for handle in handles { - let assigned_id = handle.await.unwrap().unwrap(); + let assigned_id = handle.await.unwrap(); //println!("assigned id: {assigned_id} ({wait_for_expiry})"); assert!( - expected_ids.remove(&assigned_id), + assigned_ids.insert(assigned_id), "already assigned or invalid: {assigned_id} ({wait_for_expiry})" ); } assert_eq!( - expected_ids.len(), - 0, - "{expected_ids:?} ({wait_for_expiry})" + assigned_ids.len(), + 100, + "{assigned_ids:?} ({wait_for_expiry})" ); if wait_for_expiry { tokio::time::sleep(Duration::from_secs(3)).await; } } + db.destroy().await; } async fn test_2(db: Store) { + // Test document id assignment + let mut handles = Vec::new(); + let mut assigned_ids = HashSet::new(); + + // Create 1000 ids concurrently + println!("Create 1000 documentIds concurrently..."); + ID_ASSIGNMENT_EXPIRY.store(10 * 60 * 60, std::sync::atomic::Ordering::Relaxed); + for _ in 0..1000 { + handles.push({ + let db = db.clone(); + tokio::spawn(async move { + { + let id = db.assign_document_id(0, u8::MAX).await.unwrap(); + db.write( + BatchBuilder::new() + .with_account_id(0) + .with_collection(u8::MAX) + .create_document(id) + .build_batch(), + ) + .await + .unwrap(); + id + } + }) + }); + } + + for handle in handles { + let assigned_id = handle.await.unwrap(); + assert!( + assigned_ids.insert(assigned_id), + "already assigned or invalid: {assigned_id}" + ); + } + assert_eq!(assigned_ids.len(), 1000, "{assigned_ids:?} "); + + db.destroy().await; +} + +async fn test_3(db: Store) { // Create document ids and try reassigning + println!("Assigning 100 ids concurrently and try reassigning..."); + + ID_ASSIGNMENT_EXPIRY.store(2, std::sync::atomic::Ordering::Relaxed); let mut expected_ids = AHashSet::new(); let mut batch = BatchBuilder::new(); batch.with_account_id(0).with_collection(u8::MAX); @@ -107,8 +180,10 @@ async fn test_2(db: Store) { db.destroy().await; } -async fn test_3(db: Store) { +async fn test_4(db: Store) { // Try reassigning deleted ids + println!("Create and delete 100 documentIds then try reassigning ids..."); + ID_ASSIGNMENT_EXPIRY.store(60 * 60, std::sync::atomic::Ordering::Relaxed); let mut expected_ids = AHashSet::new(); let mut batch = BatchBuilder::new(); batch.with_account_id(0).with_collection(u8::MAX); diff --git a/tests/src/store/blob.rs b/tests/src/store/blob.rs index 9bb233e6..a7dee4ba 100644 --- a/tests/src/store/blob.rs +++ b/tests/src/store/blob.rs @@ -22,7 +22,7 @@ */ use store::{ - backend::{fs::FsStore, s3::S3Store, sqlite::SqliteStore}, + backend::{fs::FsStore, postgres::PostgresStore, s3::S3Store, sqlite::SqliteStore}, write::{blob::BlobQuota, now, BatchBuilder, BlobOp, F_CLEAR}, BlobClass, BlobHash, BlobStore, Store, }; @@ -47,13 +47,19 @@ path = "{TMP}" const CONFIG_DB: &str = r#" [store.db] path = "{TMP}/db.db?mode=rwc" +host = "localhost" +post = 5432 +database = "stalwart" +user = "postgres" +password = "mysecretpassword" + "#; #[tokio::test] pub async fn blob_tests() { let temp_dir = TempDir::new("blob_tests", true); - for (store_id, store_cfg) in [("s3", CONFIG_S3), ("fs", CONFIG_LOCAL)] { + /*for (store_id, store_cfg) in [("s3", CONFIG_S3), ("fs", CONFIG_LOCAL)] { let config = Config::new(&store_cfg.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap())) .unwrap(); @@ -66,10 +72,11 @@ pub async fn blob_tests() { println!("Testing blob store {}...", store_id); test_store(blob_store_.clone()).await; - } + }*/ // Init store - let store: Store = SqliteStore::open( + //let store: Store = SqliteStore::open( + let store: Store = PostgresStore::open( //let store: Store = FdbStore::open( &Config::new(&CONFIG_DB.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap())) .unwrap(), diff --git a/tests/src/store/mod.rs b/tests/src/store/mod.rs index aca2c777..429e3434 100644 --- a/tests/src/store/mod.rs +++ b/tests/src/store/mod.rs @@ -29,28 +29,36 @@ use std::io::Read; use ::store::Store; -use store::backend::{foundationdb::FdbStore, sqlite::SqliteStore}; +use store::backend::{foundationdb::FdbStore, postgres::PostgresStore, sqlite::SqliteStore}; use utils::config::Config; pub struct TempDir { pub path: std::path::PathBuf, } +const CONFIG: &str = r#" +[store.blob] +type = "local" +local.path = "PATH" + +[store.db] +#path = "PATH/sqlite.db" +host = "localhost" +post = 5432 +database = "stalwart" +user = "postgres" +password = "mysecretpassword" + +"#; + #[tokio::test] pub async fn store_tests() { let insert = true; let temp_dir = TempDir::new("store_tests", insert); - let config_file = format!( - concat!( - "store.blob.type = \"local\"\n", - "store.blob.local.path = \"{}\"\n", - "store.db.path = \"{}/sqlite.db\"\n" - ), - temp_dir.path.display(), - temp_dir.path.display() - ); + let config_file = CONFIG.replace("PATH", &temp_dir.path.to_string_lossy()); //let db: Store = SqliteStore::open(&Config::new(&config_file).unwrap()) - let db: Store = FdbStore::open(&Config::new(&config_file).unwrap()) + //let db: Store = FdbStore::open(&Config::new(&config_file).unwrap()) + let db: Store = PostgresStore::open(&Config::new(&config_file).unwrap()) .await .unwrap() .into();