From ff279b3a39c0192ef6d9e728b9b2a852399fc1e4 Mon Sep 17 00:00:00 2001 From: mdecimus Date: Mon, 4 Mar 2024 18:51:24 +0100 Subject: [PATCH] Global threadId caching --- crates/directory/src/core/config.rs | 2 - crates/imap/src/core/mod.rs | 47 ++-------------- crates/imap/src/lib.rs | 36 +++--------- crates/imap/src/op/select.rs | 4 +- crates/imap/src/op/thread.rs | 67 +++-------------------- crates/jmap/src/api/config.rs | 5 +- crates/jmap/src/auth/oauth/mod.rs | 2 - crates/jmap/src/email/cache.rs | 85 +++++++++++++++++++++++++++++ crates/jmap/src/email/get.rs | 27 +++++---- crates/jmap/src/email/ingest.rs | 15 +---- crates/jmap/src/email/mod.rs | 1 + crates/jmap/src/lib.rs | 22 +++++--- crates/jmap/src/mailbox/get.rs | 27 +++++---- crates/smtp/src/config/resolver.rs | 5 +- crates/smtp/src/config/scripts.rs | 6 +- crates/utils/src/lib.rs | 41 +++++++++++++- 16 files changed, 208 insertions(+), 184 deletions(-) create mode 100644 crates/jmap/src/email/cache.rs diff --git a/crates/directory/src/core/config.rs b/crates/directory/src/core/config.rs index 06a81035..1b8744b3 100644 --- a/crates/directory/src/core/config.rs +++ b/crates/directory/src/core/config.rs @@ -69,8 +69,6 @@ impl ConfigDirectory for Config { stores.get_lookup_store(self, "storage.lookup")?, )); - let todo = "store quota as u64"; - for id in self.sub_keys("directory", ".type") { if id.ends_with(".columns") || id.ends_with(".attributes") || id.contains(".principals") { diff --git a/crates/imap/src/core/mod.rs b/crates/imap/src/core/mod.rs index 11053df4..ce8c8c3b 100644 --- a/crates/imap/src/core/mod.rs +++ b/crates/imap/src/core/mod.rs @@ -24,14 +24,11 @@ use std::{ collections::BTreeMap, net::IpAddr, - sync::{ - atomic::{AtomicU32, AtomicU64}, - Arc, - }, + sync::{atomic::AtomicU32, Arc}, time::Duration, }; -use ahash::{AHashMap, AHashSet}; +use ahash::AHashMap; use dashmap::DashMap; use imap_proto::{ protocol::{list::Attribute, ProtocolVersion}, @@ -42,7 +39,7 @@ use jmap::{ auth::{rate_limit::ConcurrencyLimiters, AccessToken}, JMAP, }; -use store::{roaring::RoaringBitmap, write::now}; +use store::roaring::RoaringBitmap; use tokio::{ io::{ReadHalf, WriteHalf}, sync::watch, @@ -50,6 +47,7 @@ use tokio::{ use utils::{ config::Rate, listener::{limiter::InFlight, ServerInstance, SessionStream}, + CachedItem, }; pub mod client; @@ -88,17 +86,8 @@ pub struct IMAP { pub rate_concurrent: u64, pub cache_account: DashMap>, - pub cache_account_expiry: u64, pub cache_mailbox: DashMap>, - pub cache_mailbox_expiry: u64, - pub cache_threads: DashMap>, - pub cache_threads_expiry: u64, -} - -#[derive(Clone)] -pub struct CachedItem { - last_access: Arc, - item: Arc>, + pub cache_expiry: u64, } pub struct Session { @@ -198,12 +187,6 @@ pub struct MailboxSync { pub deleted: Vec, } -#[derive(Debug, Default)] -pub struct Threads { - pub threads: AHashMap, - pub modseq: Option, -} - pub enum SavedSearch { InFlight { rx: watch::Receiver>>, @@ -284,23 +267,3 @@ impl SessionData { } } } - -impl CachedItem { - pub fn new(item: T) -> Self { - Self { - last_access: Arc::new(AtomicU64::new(now())), - item: Arc::new(tokio::sync::Mutex::new(item)), - } - } - - pub async fn get(&self) -> tokio::sync::MutexGuard<'_, T> { - let lock = self.item.lock().await; - self.last_access - .store(now(), std::sync::atomic::Ordering::Relaxed); - lock - } - - pub fn last_access(&self) -> u64 { - self.last_access.load(std::sync::atomic::Ordering::Relaxed) - } -} diff --git a/crates/imap/src/lib.rs b/crates/imap/src/lib.rs index bb826f2f..dcbdf14b 100644 --- a/crates/imap/src/lib.rs +++ b/crates/imap/src/lib.rs @@ -21,7 +21,6 @@ * for more details. */ -use core::mailbox; use std::{collections::hash_map::RandomState, sync::Arc, time::Duration}; use crate::core::IMAP; @@ -47,8 +46,6 @@ impl IMAP { .unwrap_or(32) .next_power_of_two() as usize; - let todo = "document imap.cache.rate-limit.size and imap.cache.mailbox.size"; - Ok(Arc::new(IMAP { max_request_size: config.property_or_static("imap.request.max-size", "52428800")?, max_auth_failures: config.property_or_static("imap.auth.max-failures", "3")?, @@ -70,9 +67,7 @@ impl IMAP { }) .into_bytes(), rate_limiter: DashMap::with_capacity_and_hasher_and_shard_amount( - config - .property("imap.cache.rate-limit.size")? - .unwrap_or(2048), + config.property("cache.rate-limit.size")?.unwrap_or(2048), RandomState::default(), shard_amount, ), @@ -81,42 +76,27 @@ impl IMAP { allow_plain_auth: config.property_or_static("imap.auth.allow-plain-text", "false")?, enable_uidplus: config.property_or_static("imap.protocol.uidplus", "false")?, cache_account: DashMap::with_capacity_and_hasher_and_shard_amount( - config.property("imap.cache.account.size")?.unwrap_or(2048), + config.property("cache.messages.size")?.unwrap_or(2048), RandomState::default(), shard_amount, ), cache_mailbox: DashMap::with_capacity_and_hasher_and_shard_amount( - config.property("imap.cache.mailbox.size")?.unwrap_or(2048), + config.property("cache.messages.size")?.unwrap_or(2048), RandomState::default(), shard_amount, ), - cache_threads: DashMap::with_capacity_and_hasher_and_shard_amount( - config.property("imap.cache.thread.size")?.unwrap_or(2048), - RandomState::default(), - shard_amount, - ), - cache_account_expiry: config - .property_or_static::("imap.cache.account.expiry", "1h")? - .as_secs(), - cache_mailbox_expiry: config - .property_or_static::("imap.cache.mailbox.expiry", "1h")? - .as_secs(), - cache_threads_expiry: config - .property_or_static::("imap.cache.thread.expiry", "1h")? + cache_expiry: config + .property_or_static::("cache.messages.ttl", "1h")? .as_secs(), })) } pub fn purge(&self) { - let account_expiry = now() - self.cache_account_expiry; - let mailbox_expiry = now() - self.cache_mailbox_expiry; - let thread_expiry = now() - self.cache_threads_expiry; + let expiry = now() - self.cache_expiry; self.cache_account - .retain(|_, item| item.last_access() > account_expiry); + .retain(|_, item| item.last_access() > expiry); self.cache_mailbox - .retain(|_, item| item.last_access() > mailbox_expiry); - self.cache_threads - .retain(|_, item| item.last_access() > thread_expiry); + .retain(|_, item| item.last_access() > expiry); } } diff --git a/crates/imap/src/op/select.rs b/crates/imap/src/op/select.rs index d13a8ca7..e622b03d 100644 --- a/crates/imap/src/op/select.rs +++ b/crates/imap/src/op/select.rs @@ -35,9 +35,9 @@ use imap_proto::{ }; use jmap_proto::types::id::Id; -use utils::listener::SessionStream; +use utils::{listener::SessionStream, CachedItem}; -use crate::core::{CachedItem, MailboxState, SavedSearch, SelectedMailbox, Session, State}; +use crate::core::{MailboxState, SavedSearch, SelectedMailbox, Session, State}; use super::ToModSeq; diff --git a/crates/imap/src/op/thread.rs b/crates/imap/src/op/thread.rs index 1dab546d..2b406eec 100644 --- a/crates/imap/src/op/thread.rs +++ b/crates/imap/src/op/thread.rs @@ -33,11 +33,9 @@ use imap_proto::{ Command, StatusResponse, }; -use jmap_proto::types::{collection::Collection, property::Property}; -use store::{write::ValueClass, ValueKey}; use utils::listener::SessionStream; -use crate::core::{CachedItem, SelectedMailbox, Session, SessionData, Threads}; +use crate::core::{SelectedMailbox, Session, SessionData}; impl Session { pub async fn handle_thread( @@ -87,67 +85,20 @@ impl SessionData { }); } - // Obtain current state - let modseq = self + // Lock the cache + let thread_ids = self .jmap - .store - .get_last_change_id(mailbox.id.account_id, Collection::Thread) + .get_cached_thread_ids(mailbox.id.account_id, result_set.results.iter()) .await .map_err(|err| { - tracing::error!(event = "error", - context = "store", - account_id = mailbox.id.account_id, - collection = ?Collection::Thread, - error = ?err, - "Failed to obtain state"); + tracing::error!( + event = "error", + context = "thread_query", + error = ?err, + "Failed to obtain threadId."); StatusResponse::database_failure() })?; - // Lock the cache - let thread_cache_ = self - .imap - .cache_threads - .entry(mailbox.id.account_id) - .or_insert_with(|| CachedItem::new(Threads::default())); - let mut thread_cache = thread_cache_.get().await; - - // Invalidate cache if the modseq has changed - if thread_cache.modseq != modseq { - thread_cache.threads.clear(); - } - - // Obtain threadIds for matching messages - let mut thread_ids = Vec::with_capacity(result_set.results.len() as usize); - for document_id in &result_set.results { - if let Some(thread_id) = thread_cache.threads.get(&document_id) { - thread_ids.push((*thread_id).into()); - } else if let Some(thread_id) = self - .jmap - .store - .get_value::(ValueKey { - account_id: mailbox.id.account_id, - collection: Collection::Email.into(), - document_id, - class: ValueClass::Property(Property::ThreadId.into()), - }) - .await - .map_err(|err| { - tracing::error!( - event = "error", - context = "thread_query", - error = ?err, - "Failed to obtain threadId."); - StatusResponse::database_failure() - })? - { - thread_ids.push(thread_id.into()); - thread_cache.threads.insert(document_id, thread_id); - } else { - thread_ids.push(None); - } - } - thread_cache.modseq = modseq; - // Group messages by thread let mut threads: AHashMap> = AHashMap::new(); let state = mailbox.state.lock(); diff --git a/crates/jmap/src/api/config.rs b/crates/jmap/src/api/config.rs index 59c79e68..083dd4a7 100644 --- a/crates/jmap/src/api/config.rs +++ b/crates/jmap/src/api/config.rs @@ -97,7 +97,7 @@ impl crate::Config { .unwrap_or(256), capabilities: BaseCapabilities::default(), session_cache_ttl: settings - .property("jmap.session.cache.ttl")? + .property("cache.session.ttl")? .unwrap_or(Duration::from_secs(3600)), rate_authenticated: settings .property_or_static("jmap.rate-limit.account", "1000/1m")?, @@ -177,6 +177,9 @@ impl crate::Config { } }) .collect::, String>>()?, + cache_expiry: settings + .property_or_static::("cache.messages.ttl", "1h")? + .as_secs(), }; config.add_capabilites(settings); Ok(config) diff --git a/crates/jmap/src/auth/oauth/mod.rs b/crates/jmap/src/auth/oauth/mod.rs index 02a55501..f0632c79 100644 --- a/crates/jmap/src/auth/oauth/mod.rs +++ b/crates/jmap/src/auth/oauth/mod.rs @@ -81,8 +81,6 @@ pub struct OAuthCode { pub redirect_uri: Option, } -struct TodoStoreOAuthCodeInDb {} - #[derive(Debug, Serialize, Deserialize)] pub struct DeviceAuthGet { code: Option, diff --git a/crates/jmap/src/email/cache.rs b/crates/jmap/src/email/cache.rs new file mode 100644 index 00000000..42e1e630 --- /dev/null +++ b/crates/jmap/src/email/cache.rs @@ -0,0 +1,85 @@ +/* + * 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 jmap_proto::types::{collection::Collection, property::Property}; +use store::{ahash::AHashMap, write::ValueClass, ValueKey}; +use utils::CachedItem; + +use crate::JMAP; + +#[derive(Debug, Default)] +pub struct Threads { + pub threads: AHashMap, + pub modseq: Option, +} + +impl JMAP { + pub async fn get_cached_thread_ids( + &self, + account_id: u32, + message_ids: impl Iterator, + ) -> store::Result>> { + // Obtain current state + let modseq = self + .store + .get_last_change_id(account_id, Collection::Thread) + .await?; + + // Lock the cache + let thread_cache_ = self + .cache_threads + .entry(account_id) + .or_insert_with(|| CachedItem::new(Threads::default())); + let mut thread_cache = thread_cache_.get().await; + + // Invalidate cache if the modseq has changed + if thread_cache.modseq != modseq { + thread_cache.threads.clear(); + } + + // Obtain threadIds for matching messages + let mut thread_ids = Vec::with_capacity(message_ids.size_hint().0); + for document_id in message_ids { + if let Some(thread_id) = thread_cache.threads.get(&document_id) { + thread_ids.push((*thread_id).into()); + } else if let Some(thread_id) = self + .store + .get_value::(ValueKey { + account_id, + collection: Collection::Email.into(), + document_id, + class: ValueClass::Property(Property::ThreadId.into()), + }) + .await? + { + thread_ids.push(thread_id.into()); + thread_cache.threads.insert(document_id, thread_id); + } else { + thread_ids.push(None); + } + } + thread_cache.modseq = modseq; + + Ok(thread_ids) + } +} diff --git a/crates/jmap/src/email/get.rs b/crates/jmap/src/email/get.rs index b2a688bd..9b6bd239 100644 --- a/crates/jmap/src/email/get.rs +++ b/crates/jmap/src/email/get.rs @@ -110,17 +110,22 @@ impl JMAP { .iter() .take(self.config.get_max_objects) .collect::>(); - self.get_properties::( - account_id, - Collection::Email, - document_ids.iter().copied(), - Property::ThreadId, - ) - .await? - .into_iter() - .zip(document_ids) - .filter_map(|(thread_id, document_id)| Id::from_parts(thread_id?, document_id).into()) - .collect() + self.get_cached_thread_ids(account_id, document_ids.iter().copied()) + .await + .map_err(|err| { + tracing::error!(event = "error", + context = "store", + account_id = account_id, + error = ?err, + "Failed to retrieve thread Ids"); + MethodError::ServerPartialFail + })? + .into_iter() + .zip(document_ids) + .filter_map(|(thread_id, document_id)| { + Id::from_parts(thread_id?, document_id).into() + }) + .collect() }; let mut response = GetResponse { account_id: request.account_id.into(), diff --git a/crates/jmap/src/email/ingest.rs b/crates/jmap/src/email/ingest.rs index bb83e39e..0ce44f33 100644 --- a/crates/jmap/src/email/ingest.rs +++ b/crates/jmap/src/email/ingest.rs @@ -42,7 +42,7 @@ use store::{ log::ChangeLogBuilder, now, BatchBuilder, BitmapClass, TagValue, ValueClass, F_BITMAP, F_CLEAR, F_VALUE, }, - BitmapKey, BlobClass, ValueKey, + BitmapKey, BlobClass, }; use utils::map::vec_map::VecMap; @@ -437,18 +437,7 @@ impl JMAP { // Obtain threadIds for matching messages let thread_ids = self - .store - .get_values::( - results - .iter() - .map(|document_id| ValueKey { - account_id, - collection: Collection::Email.into(), - document_id, - class: ValueClass::Property(Property::ThreadId.into()), - }) - .collect(), - ) + .get_cached_thread_ids(account_id, results.iter()) .await .map_err(|err| { tracing::error!( diff --git a/crates/jmap/src/email/mod.rs b/crates/jmap/src/email/mod.rs index 4c534a67..e2a7d3c1 100644 --- a/crates/jmap/src/email/mod.rs +++ b/crates/jmap/src/email/mod.rs @@ -22,6 +22,7 @@ */ pub mod body; +pub mod cache; pub mod copy; pub mod crypto; pub mod get; diff --git a/crates/jmap/src/lib.rs b/crates/jmap/src/lib.rs index fc235497..1e595339 100644 --- a/crates/jmap/src/lib.rs +++ b/crates/jmap/src/lib.rs @@ -28,6 +28,7 @@ use api::session::BaseCapabilities; use auth::{oauth::OAuthCode, rate_limit::ConcurrencyLimiters, AccessToken}; use dashmap::DashMap; use directory::{Directories, Directory, QueryBy}; +use email::cache::Threads; use jmap_proto::{ error::method::MethodError, method::{ @@ -57,7 +58,7 @@ use utils::{ ipc::DeliveryEvent, map::ttl_dashmap::{TtlDashMap, TtlMap}, snowflake::SnowflakeIdGenerator, - UnwrapFailure, + CachedItem, UnwrapFailure, }; pub mod api; @@ -98,6 +99,8 @@ pub struct JMAP { pub housekeeper_tx: mpsc::Sender, pub smtp: Arc, + pub cache_threads: DashMap>, + pub sieve_compiler: Compiler, pub sieve_runtime: Runtime<()>, } @@ -152,6 +155,8 @@ pub struct Config { pub oauth_expiry_refresh_token_renew: u64, pub oauth_max_auth_attempts: u32, + pub cache_expiry: u64, + pub spam_header: Option<(HeaderName<'static>, String)>, pub http_headers: Vec<(hyper::header::HeaderName, hyper::header::HeaderValue)>, @@ -207,22 +212,25 @@ impl JMAP { lookup_store: stores.get_lookup_store(config, "storage.lookup")?, config: Config::new(config).failed("Invalid configuration file"), sessions: TtlDashMap::with_capacity( - config.property("jmap.session.cache.size")?.unwrap_or(100), + config.property("cache.session.size")?.unwrap_or(100), shard_amount, ), access_tokens: TtlDashMap::with_capacity( - config.property("jmap.session.cache.size")?.unwrap_or(100), + config.property("cache.session.size")?.unwrap_or(100), shard_amount, ), concurrency_limiter: DashMap::with_capacity_and_hasher_and_shard_amount( - config - .property("jmap.rate-limit.cache.size")? - .unwrap_or(1024), + config.property("cache.rate-limit.size")?.unwrap_or(1024), RandomState::default(), shard_amount, ), oauth_codes: TtlDashMap::with_capacity( - config.property("oauth.cache.size")?.unwrap_or(128), + config.property("cache.oauth.size")?.unwrap_or(128), + shard_amount, + ), + cache_threads: DashMap::with_capacity_and_hasher_and_shard_amount( + config.property("cache.messages.size")?.unwrap_or(2048), + RandomState::default(), shard_amount, ), state_tx, diff --git a/crates/jmap/src/mailbox/get.rs b/crates/jmap/src/mailbox/get.rs index 5133e3ed..75cad373 100644 --- a/crates/jmap/src/mailbox/get.rs +++ b/crates/jmap/src/mailbox/get.rs @@ -261,18 +261,21 @@ impl JMAP { ) -> Result { if let Some(document_ids) = document_ids { let mut thread_ids = AHashSet::default(); - self.get_properties::( - account_id, - Collection::Email, - document_ids.into_iter(), - Property::ThreadId, - ) - .await? - .into_iter() - .flatten() - .for_each(|thread_id| { - thread_ids.insert(thread_id); - }); + self.get_cached_thread_ids(account_id, document_ids.into_iter()) + .await + .map_err(|err| { + tracing::error!(event = "error", + context = "store", + account_id = account_id, + error = ?err, + "Failed to retrieve thread Ids"); + MethodError::ServerPartialFail + })? + .into_iter() + .flatten() + .for_each(|thread_id| { + thread_ids.insert(thread_id); + }); Ok(thread_ids.len()) } else { Ok(0) diff --git a/crates/smtp/src/config/resolver.rs b/crates/smtp/src/config/resolver.rs index 37105f08..e2547771 100644 --- a/crates/smtp/src/config/resolver.rs +++ b/crates/smtp/src/config/resolver.rs @@ -96,10 +96,11 @@ impl ConfigResolver for Config { .map_err(|err| format!("Failed to build DNSSEC resolver: {err}"))?, cache: crate::core::DnsCache { tlsa: LruCache::with_capacity( - self.property("resolver.cache.tlsa")?.unwrap_or(1024), + self.property("cache.resolver.tlsa.size")?.unwrap_or(1024), ), mta_sts: LruCache::with_capacity( - self.property("resolver.cache.mta-sts")?.unwrap_or(1024), + self.property("cache.resolver.mta-sts.size")? + .unwrap_or(1024), ), }, }) diff --git a/crates/smtp/src/config/scripts.rs b/crates/smtp/src/config/scripts.rs index 06a3bac9..e1caa805 100644 --- a/crates/smtp/src/config/scripts.rs +++ b/crates/smtp/src/config/scripts.rs @@ -77,9 +77,9 @@ impl ConfigSieve for Config { let sieve_ctx = SieveContext { psl: self.parse_public_suffix()?, bayes_cache: BayesTokenCache::new( - self.property_or_static("bayes.cache.capacity", "8192")?, - self.property_or_static("bayes.cache.ttl.positive", "1h")?, - self.property_or_static("bayes.cache.ttl.negative", "1h")?, + self.property_or_static("cache.bayes.capacity", "8192")?, + self.property_or_static("cache.bayes.ttl.positive", "1h")?, + self.property_or_static("cache.bayes.ttl.negative", "1h")?, ), remote_lists: Default::default(), }; diff --git a/crates/utils/src/lib.rs b/crates/utils/src/lib.rs index 64e15067..2d24a4fc 100644 --- a/crates/utils/src/lib.rs +++ b/crates/utils/src/lib.rs @@ -21,7 +21,11 @@ * for more details. */ -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{atomic::AtomicU64, Arc}, + time::SystemTime, +}; use config::Config; @@ -110,6 +114,41 @@ impl AsMut<[u8]> for BlobHash { self.0.as_mut() } } + +#[derive(Clone)] +pub struct CachedItem { + last_access: Arc, + item: Arc>, +} + +impl CachedItem { + pub fn new(item: T) -> Self { + Self { + last_access: Arc::new(AtomicU64::new( + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map_or(0, |d| d.as_secs()), + )), + item: Arc::new(tokio::sync::Mutex::new(item)), + } + } + + pub async fn get(&self) -> tokio::sync::MutexGuard<'_, T> { + let lock = self.item.lock().await; + self.last_access.store( + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map_or(0, |d| d.as_secs()), + std::sync::atomic::Ordering::Relaxed, + ); + lock + } + + pub fn last_access(&self) -> u64 { + self.last_access.load(std::sync::atomic::Ordering::Relaxed) + } +} + pub trait UnwrapFailure { fn failed(self, action: &str) -> T; }