diff --git a/Cargo.lock b/Cargo.lock index 16ef4193..ad084f5d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6521,6 +6521,7 @@ dependencies = [ "chrono", "dashmap", "futures", + "lru-cache", "mail-auth", "mail-send", "opentelemetry", diff --git a/crates/directory/src/backend/internal/manage.rs b/crates/directory/src/backend/internal/manage.rs index fa2b6580..88376550 100644 --- a/crates/directory/src/backend/internal/manage.rs +++ b/crates/directory/src/backend/internal/manage.rs @@ -758,7 +758,10 @@ impl ManageDirectory for Store { ValueClass::Directory(DirectoryClass::Domain(domain.to_lowercase().into_bytes())), vec![], ); - self.write(batch.build()).await.map_err(Into::into) + self.write(batch.build()) + .await + .map_err(Into::into) + .map(|_| ()) } async fn delete_domain(&self, domain: &str) -> crate::Result<()> { @@ -771,7 +774,10 @@ impl ManageDirectory for Store { batch.clear(ValueClass::Directory(DirectoryClass::Domain( domain.to_lowercase().into_bytes(), ))); - self.write(batch.build()).await.map_err(Into::into) + self.write(batch.build()) + .await + .map_err(Into::into) + .map(|_| ()) } async fn map_group_ids(&self, principal: Principal) -> crate::Result> { diff --git a/crates/imap/src/core/client.rs b/crates/imap/src/core/client.rs index 76ede758..a8d02ab8 100644 --- a/crates/imap/src/core/client.rs +++ b/crates/imap/src/core/client.rs @@ -389,15 +389,7 @@ impl State { } pub fn close_mailbox(&self) -> bool { - match self { - State::Selected { mailbox, data } => { - if mailbox.is_select { - data.clear_recent(&mailbox.id); - } - true - } - _ => false, - } + matches!(self, State::Selected { .. }) } } diff --git a/crates/imap/src/core/mailbox.rs b/crates/imap/src/core/mailbox.rs index 6a924d0c..c6f846a6 100644 --- a/crates/imap/src/core/mailbox.rs +++ b/crates/imap/src/core/mailbox.rs @@ -1,4 +1,7 @@ -use std::{collections::BTreeMap, sync::atomic::Ordering}; +use std::{ + collections::BTreeMap, + sync::{atomic::Ordering, Arc}, +}; use ahash::AHashMap; use directory::QueryBy; @@ -13,12 +16,13 @@ use jmap_proto::{ }; use parking_lot::Mutex; use store::query::log::{Change, Query}; -use utils::listener::{limiter::InFlight, SessionStream}; - -use super::{ - Account, AccountId, CachedItem, Mailbox, MailboxId, MailboxSync, Session, SessionData, +use utils::{ + listener::{limiter::InFlight, SessionStream}, + lru_cache::LruCached, }; +use super::{Account, AccountId, Mailbox, MailboxId, MailboxSync, Session, SessionData}; + impl SessionData { pub async fn new( session: &Session, @@ -96,25 +100,25 @@ impl SessionData { .get_last_change_id(account_id, Collection::Email) .await .map_err(|_| {})?; - let cached_account_ = self - .imap - .cache_account - .entry(AccountId { - account_id, - primary_id: access_token.primary_id(), - }) - .or_insert_with(|| { - CachedItem::new(Account { - account_id: u32::MAX, - ..Default::default() + let cached_account_id = AccountId { + account_id, + primary_id: access_token.primary_id(), + }; + if let Some(cached_account) = + self.imap + .cache_account + .get(&cached_account_id) + .and_then(|cached_account| { + if cached_account.state_mailbox == state_mailbox + && cached_account.state_email == state_email + { + Some(cached_account) + } else { + None + } }) - }); - let mut cached_account = cached_account_.get().await; - if cached_account.state_mailbox == state_mailbox - && cached_account.state_email == state_email - && cached_account.account_id != u32::MAX { - return Ok((*cached_account).clone()); + return Ok(cached_account.as_ref().clone()); } let mailbox_ids = if access_token.is_primary_id(account_id) @@ -271,7 +275,9 @@ impl SessionData { } // Update cache - *cached_account = account.clone(); + self.imap + .cache_account + .insert(cached_account_id, Arc::new(account.clone())); Ok(account) } @@ -426,14 +432,16 @@ impl SessionData { } // Update cache - if let Some(cached_account) = self.imap.cache_account.get(&AccountId { - account_id, - primary_id: access_token.primary_id(), - }) { - let mut cached_account = cached_account.get().await; - if cached_account.state_mailbox != state_mailbox - || cached_account.state_email != state_email + if let Some(cached_account_) = + self.imap.cache_account.lock().get_mut(&AccountId { + account_id, + primary_id: access_token.primary_id(), + }) + { + if cached_account_.state_mailbox != state_mailbox + || cached_account_.state_email != state_email { + let mut cached_account = cached_account_.as_ref().clone(); cached_account.mailbox_state.values_mut().for_each(|v| { v.total_deleted = None; v.total_unseen = None; @@ -443,6 +451,7 @@ impl SessionData { }); cached_account.state_mailbox = state_mailbox; cached_account.state_email = state_email; + *cached_account_ = Arc::new(cached_account); } } } else { diff --git a/crates/imap/src/core/message.rs b/crates/imap/src/core/message.rs index 8fde6249..40ec2b11 100644 --- a/crates/imap/src/core/message.rs +++ b/crates/imap/src/core/message.rs @@ -21,7 +21,7 @@ * for more details. */ -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::Arc}; use ahash::AHashMap; use imap_proto::{ @@ -33,17 +33,12 @@ use jmap_proto::{ object::Object, types::{collection::Collection, property::Property, value::Value}, }; -use store::{ - roaring::RoaringBitmap, - write::{assert::HashedValue, BatchBuilder, F_VALUE}, -}; -use utils::listener::SessionStream; +use store::write::assert::HashedValue; +use utils::{listener::SessionStream, lru_cache::LruCached}; use crate::core::ImapId; -use super::{ - CachedItem, Mailbox, MailboxId, MailboxState, NextMailboxState, SelectedMailbox, SessionData, -}; +use super::{ImapUidToId, MailboxId, MailboxState, NextMailboxState, SelectedMailbox, SessionData}; pub(crate) const MAX_RETRIES: usize = 10; @@ -61,26 +56,8 @@ impl SessionData { .await? .unwrap_or_default(); - // Obtain mailbox data - let uid_validity = self - .jmap - .get_property::>( - mailbox.account_id, - Collection::Mailbox, - mailbox.mailbox_id, - &Property::Value, - ) - .await? - .and_then(|obj| obj.get(&Property::Cid).as_uint()) - .ok_or_else(|| { - tracing::debug!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Mailbox, - mailbox_id = mailbox.mailbox_id, - "Failed to obtain uid validity"); - StatusResponse::no("Mailbox unavailable.") - })? as u32; + // Obtain UID validity + let uid_validity = self.get_uid_validity(mailbox).await?; // Obtain current state let modseq = self @@ -98,11 +75,8 @@ impl SessionData { StatusResponse::database_failure() })?; - // Retrieve message ids - let mut assigned = BTreeMap::new(); - let mut unassigned = Vec::new(); - // Obtain all message ids + let mut uid_map = BTreeMap::new(); for (message_id, uid_mailbox) in self .jmap .get_properties::>, _, _>( @@ -120,190 +94,28 @@ impl SessionData { .iter() .find(|item| item.mailbox_id == mailbox.mailbox_id) { - if item.uid > 0 { - if assigned.insert(item.uid, message_id).is_some() { - tracing::warn!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Mailbox, - mailbox_id = mailbox.mailbox_id, - message_id = message_id, - "Duplicate UID"); - } - } else { - unassigned.push((message_id, uid_mailbox)); + debug_assert!(item.uid != 0, "UID is zero for message {item:?}"); + if uid_map.insert(item.uid, message_id).is_some() { + tracing::warn!(event = "error", + context = "store", + account_id = mailbox.account_id, + collection = ?Collection::Mailbox, + mailbox_id = mailbox.mailbox_id, + message_id = message_id, + "Duplicate UID"); } } } // Obtain UID next and assign UIDs - let mut try_count = 0; - let mut uid_next = 1; - let mut uid_other = 0; - let mut recent_messages = RoaringBitmap::new(); + let mut uid_max = 0; + let mut id_to_imap = AHashMap::with_capacity(uid_map.len()); + let mut uid_to_id = AHashMap::with_capacity(uid_map.len()); - // Shuffle unassigned - /*if unassigned.len() > 1 { - let mut rng = rand::thread_rng(); - unassigned.shuffle(&mut rng); - }*/ - - loop { - let last_uid = self - .jmap - .get_property::( - mailbox.account_id, - Collection::Mailbox, - mailbox.mailbox_id, - Property::EmailIds, - ) - .await?; - - if !unassigned.is_empty() { - // Increment UID next - let mut batch = BatchBuilder::new(); - batch - .with_account_id(mailbox.account_id) - .with_collection(Collection::Mailbox) - .update_document(mailbox.mailbox_id); - - if let Some(last_uid) = last_uid { - batch.assert_value(Property::EmailIds, last_uid).value( - Property::EmailIds, - last_uid + unassigned.len() as u32, - F_VALUE, - ); - uid_next = last_uid + 1; - } else { - batch.assert_value(Property::EmailIds, ()).value( - Property::EmailIds, - unassigned.len() as u32, - F_VALUE, - ); - } - - match self.jmap.store.write(batch.build()).await { - Ok(_) => (), - Err(store::Error::AssertValueFailed) if try_count < MAX_RETRIES => { - try_count += 1; - continue; - } - Err(err) => { - tracing::error!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Mailbox, - mailbox_id = mailbox.mailbox_id, - error = ?err, - "Failed to update UID next"); - return Err(StatusResponse::database_failure()); - } - } - - // Assign UIDs - for (message_id, mut uid_mailbox) in unassigned { - let uid = uid_next; - uid_next += 1; - try_count = 0; - - loop { - if let Some(item) = uid_mailbox - .inner - .iter_mut() - .find(|item| item.mailbox_id == mailbox.mailbox_id) - { - if item.uid == 0 { - item.uid = uid; - - // Increment UID next - let mut batch = BatchBuilder::new(); - batch - .with_account_id(mailbox.account_id) - .with_collection(Collection::Email) - .update_document(message_id) - .assert_value(Property::MailboxIds, &uid_mailbox) - .value(Property::MailboxIds, uid_mailbox.inner, F_VALUE); - - match self.jmap.store.write(batch.build()).await { - Ok(_) => { - if assigned.insert(uid, message_id).is_some() { - tracing::warn!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Mailbox, - mailbox_id = mailbox.mailbox_id, - message_id = message_id, - "Duplicate UID"); - } - recent_messages.insert(message_id); - } - Err(store::Error::AssertValueFailed) - if try_count < MAX_RETRIES => - { - // Another process modified the mailbox ids - if let Some(modified_uid_mailbox) = self - .jmap - .get_property::>>( - mailbox.account_id, - Collection::Email, - message_id, - Property::MailboxIds, - ) - .await? - { - uid_mailbox = modified_uid_mailbox; - try_count += 1; - continue; - } - } - Err(err) => { - tracing::error!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Email, - mailbox_id = message_id, - error = ?err, - "Failed to store UID"); - return Err(StatusResponse::database_failure()); - } - } - } else { - // Another thread has already assigned a UID - if item.uid > uid_other { - // Keep track of highest UID assigned by another thread - uid_other = item.uid; - } - - if assigned.insert(item.uid, message_id).is_some() { - tracing::warn!(event = "error", - context = "store", - account_id = mailbox.account_id, - collection = ?Collection::Mailbox, - mailbox_id = mailbox.mailbox_id, - message_id = message_id, - "Duplicate UID assigned by another thread"); - } - } - } - - break; - } - } - } else { - uid_next = last_uid.unwrap_or(0) + 1; + for (seqnum, (uid, message_id)) in uid_map.into_iter().enumerate() { + if uid > uid_max { + uid_max = uid; } - break; - } - - // Other processes might have assigned a higher UID - if uid_next <= uid_other { - uid_next = uid_other + 1; - } - - let mut id_to_imap = AHashMap::with_capacity(assigned.len()); - let mut uid_to_id = AHashMap::with_capacity(assigned.len()); - - for (seqnum, (uid, message_id)) in assigned.into_iter().enumerate() { id_to_imap.insert( message_id, ImapId { @@ -314,28 +126,13 @@ impl SessionData { uid_to_id.insert(uid, message_id); } - // Update recent flags - for account in self.mailboxes.lock().iter_mut() { - if account.account_id == mailbox.account_id { - let mailbox = account - .mailbox_state - .entry(mailbox.mailbox_id) - .or_insert_with(Mailbox::default); - mailbox.recent_messages &= &message_ids; - if !recent_messages.is_empty() { - mailbox.recent_messages |= &recent_messages; - } - break; - } - } - Ok(MailboxState { - uid_next, + uid_next: uid_max + 1, uid_validity, total_messages: id_to_imap.len(), id_to_imap, uid_to_id, - uid_max: uid_next.saturating_sub(1), + uid_max, modseq, next_state: None, }) @@ -375,7 +172,7 @@ impl SessionData { // Update cache self.imap .cache_mailbox - .insert(mailbox.id, CachedItem::new(new_state.clone())); + .insert(mailbox.id, Arc::new(new_state.clone())); // Update state current_state.modseq = new_state.modseq; @@ -450,36 +247,26 @@ impl SessionData { } } - pub fn get_recent(&self, mailbox: &MailboxId) -> RoaringBitmap { - for account in self.mailboxes.lock().iter() { - if account.account_id == mailbox.account_id { - if let Some(mailbox) = account.mailbox_state.get(&mailbox.mailbox_id) { - return mailbox.recent_messages.clone(); - } - } - } - RoaringBitmap::new() - } - - pub fn get_recent_count(&self, mailbox: &MailboxId) -> usize { - for account in self.mailboxes.lock().iter() { - if account.account_id == mailbox.account_id { - if let Some(mailbox) = account.mailbox_state.get(&mailbox.mailbox_id) { - return mailbox.recent_messages.len() as usize; - } - } - } - 0 - } - - pub fn clear_recent(&self, mailbox: &MailboxId) { - for account in self.mailboxes.lock().iter_mut() { - if account.account_id == mailbox.account_id { - if let Some(mailbox) = account.mailbox_state.get_mut(&mailbox.mailbox_id) { - mailbox.recent_messages.clear(); - } - } - } + pub async fn get_uid_validity(&self, mailbox: &MailboxId) -> crate::op::Result { + self.jmap + .get_property::>( + mailbox.account_id, + Collection::Mailbox, + mailbox.mailbox_id, + &Property::Value, + ) + .await? + .and_then(|obj| obj.get(&Property::Cid).as_uint()) + .ok_or_else(|| { + tracing::debug!(event = "error", + context = "store", + account_id = mailbox.account_id, + collection = ?Collection::Mailbox, + mailbox_id = mailbox.mailbox_id, + "Failed to obtain uid validity"); + StatusResponse::no("Mailbox unavailable.") + }) + .map(|v| v as u32) } } @@ -557,4 +344,27 @@ impl SelectedMailbox { deleted_ids.sort_unstable(); deleted_ids } + + pub fn append_messages(&self, ids: Vec, modseq: Option) -> u32 { + let mut mailbox = self.state.lock(); + if modseq.unwrap_or(0) > mailbox.modseq.unwrap_or(0) { + let mut uid_max = 0; + for id in ids { + mailbox.total_messages += 1; + let seqnum = mailbox.total_messages as u32; + mailbox.uid_to_id.insert(id.uid, id.uid); + mailbox.id_to_imap.insert( + id.id, + ImapId { + uid: id.uid, + seqnum, + }, + ); + uid_max = id.uid; + } + mailbox.uid_max = uid_max; + mailbox.uid_next = uid_max + 1; + } + mailbox.uid_validity + } } diff --git a/crates/imap/src/core/mod.rs b/crates/imap/src/core/mod.rs index ce8c8c3b..7b988239 100644 --- a/crates/imap/src/core/mod.rs +++ b/crates/imap/src/core/mod.rs @@ -39,7 +39,6 @@ use jmap::{ auth::{rate_limit::ConcurrencyLimiters, AccessToken}, JMAP, }; -use store::roaring::RoaringBitmap; use tokio::{ io::{ReadHalf, WriteHalf}, sync::watch, @@ -47,7 +46,7 @@ use tokio::{ use utils::{ config::Rate, listener::{limiter::InFlight, ServerInstance, SessionStream}, - CachedItem, + lru_cache::LruCache, }; pub mod client; @@ -72,7 +71,6 @@ pub struct IMAP { pub max_auth_failures: u32, pub name_shared: String, pub allow_plain_auth: bool, - pub enable_uidplus: bool, pub timeout_auth: Duration, pub timeout_unauth: Duration, @@ -85,9 +83,8 @@ pub struct IMAP { pub rate_requests: Rate, pub rate_concurrent: u64, - pub cache_account: DashMap>, - pub cache_mailbox: DashMap>, - pub cache_expiry: u64, + pub cache_account: LruCache>, + pub cache_mailbox: LruCache>, } pub struct Session { @@ -129,7 +126,6 @@ pub struct Mailbox { pub uid_validity: Option, pub uid_next: Option, pub size: Option, - pub recent_messages: RoaringBitmap, } #[derive(Debug, Clone, Default)] @@ -203,6 +199,12 @@ pub struct ImapId { pub seqnum: u32, } +#[derive(Debug, Clone, Copy, Default)] +pub struct ImapUidToId { + pub uid: u32, + pub id: u32, +} + pub enum State { NotAuthenticated { auth_failures: u32, diff --git a/crates/imap/src/lib.rs b/crates/imap/src/lib.rs index dcbdf14b..a573a2ea 100644 --- a/crates/imap/src/lib.rs +++ b/crates/imap/src/lib.rs @@ -21,14 +21,16 @@ * for more details. */ -use std::{collections::hash_map::RandomState, sync::Arc, time::Duration}; +use std::{collections::hash_map::RandomState, sync::Arc}; use crate::core::IMAP; use dashmap::DashMap; use imap_proto::{protocol::capability::Capability, ResponseCode, StatusResponse}; -use store::write::now; -use utils::config::Config; +use utils::{ + config::Config, + lru_cache::{LruCache, LruCached}, +}; pub mod core; pub mod op; @@ -74,30 +76,14 @@ impl IMAP { rate_requests: config.property_or_static("imap.rate-limit.requests", "2000/1m")?, rate_concurrent: config.property("imap.rate-limit.concurrent")?.unwrap_or(4), 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( + cache_account: LruCache::with_capacity( config.property("cache.messages.size")?.unwrap_or(2048), - RandomState::default(), - shard_amount, ), - cache_mailbox: DashMap::with_capacity_and_hasher_and_shard_amount( + cache_mailbox: LruCache::with_capacity( config.property("cache.messages.size")?.unwrap_or(2048), - RandomState::default(), - shard_amount, ), - cache_expiry: config - .property_or_static::("cache.messages.ttl", "1h")? - .as_secs(), })) } - - pub fn purge(&self) { - let expiry = now() - self.cache_expiry; - self.cache_account - .retain(|_, item| item.last_access() > expiry); - self.cache_mailbox - .retain(|_, item| item.last_access() > expiry); - } } pub struct ImapError; diff --git a/crates/imap/src/op/append.rs b/crates/imap/src/op/append.rs index 8bf9ff42..af55b5d1 100644 --- a/crates/imap/src/op/append.rs +++ b/crates/imap/src/op/append.rs @@ -34,7 +34,7 @@ use jmap_proto::types::{acl::Acl, keyword::Keyword, state::StateChange, type_sta use mail_parser::MessageParser; use utils::listener::SessionStream; -use crate::core::{MailboxId, SelectedMailbox, Session, SessionData}; +use crate::core::{ImapUidToId, MailboxId, SelectedMailbox, Session, SessionData}; use super::ToModSeq; @@ -138,7 +138,10 @@ impl SessionData { .await { Ok(email) => { - created_ids.push(email.id.document_id()); + created_ids.push(ImapUidToId { + uid: email.imap_uids[0], + id: email.id.document_id(), + }); last_change_id = Some(email.change_id); } Err(err) => { @@ -172,49 +175,26 @@ impl SessionData { } if !created_ids.is_empty() { - let (uids, uid_validity) = match selected_mailbox { + let uids = created_ids.iter().map(|id| id.uid).collect(); + let uid_validity = match selected_mailbox { Some(selected_mailbox) if selected_mailbox.id == mailbox => { - let modseq = self - .write_mailbox_changes(&selected_mailbox, is_qresync) - .await - .map_err(|r| r.with_tag(&arguments.tag))?; - // Write updated modseq if is_qresync { - self.write_bytes(HighestModSeq::new(modseq.to_modseq()).into_bytes()) - .await; + self.write_bytes( + HighestModSeq::new(last_change_id.to_modseq()).into_bytes(), + ) + .await; } - let mailbox = selected_mailbox.state.lock(); - ( - created_ids - .into_iter() - .filter_map(|id| mailbox.id_to_imap.get(&id)) - .map(|id| id.uid) - .collect::>(), - mailbox.uid_validity, - ) + selected_mailbox.append_messages(created_ids, last_change_id) } - - _ if self.imap.enable_uidplus => { - let mailbox = self - .fetch_messages(&mailbox) - .await - .map_err(|r| r.with_tag(&arguments.tag))?; - ( - created_ids - .into_iter() - .filter_map(|id| mailbox.id_to_imap.get(&id)) - .map(|id| id.uid) - .collect(), - mailbox.uid_validity, - ) - } - _ => (vec![], 0), + _ => self + .get_uid_validity(&mailbox) + .await + .map_err(|r| r.with_tag(&arguments.tag))?, }; - if !uids.is_empty() { - response = response.with_code(ResponseCode::AppendUid { uid_validity, uids }); - } + + response = response.with_code(ResponseCode::AppendUid { uid_validity, uids }); } Ok(response.with_tag(arguments.tag)) diff --git a/crates/imap/src/op/copy_move.rs b/crates/imap/src/op/copy_move.rs index 31c9bf03..31a6138e 100644 --- a/crates/imap/src/op/copy_move.rs +++ b/crates/imap/src/op/copy_move.rs @@ -184,7 +184,7 @@ impl SessionData { if src_mailbox.id.account_id == dest_mailbox.account_id { // Mailboxes are in the same account let account_id = src_mailbox.id.account_id; - let dest_mailbox_id = UidMailbox::from(dest_mailbox_id); + let dest_mailbox_id = UidMailbox::new_unassigned(dest_mailbox_id); for (id, imap_id) in ids { // Obtain mailbox tags let (mut mailboxes, thread_id) = if let Some(result) = self @@ -196,10 +196,11 @@ impl SessionData { } else { continue; }; + // Make sure the message still belongs to this mailbox if !mailboxes .current() - .contains(&UidMailbox::from(src_mailbox.id.mailbox_id)) + .contains(&UidMailbox::new_unassigned(src_mailbox.id.mailbox_id)) || mailboxes.current().contains(&dest_mailbox_id) { tracing::debug!( @@ -213,7 +214,30 @@ impl SessionData { // Add destination folder mailboxes.update(dest_mailbox_id, true); if is_move { - mailboxes.update(UidMailbox::from(src_mailbox.id.mailbox_id), false); + mailboxes.update(UidMailbox::new_unassigned(src_mailbox.id.mailbox_id), false); + } + + // Assign IMAP UIDs + for uid_mailbox in mailboxes.inner_tags_mut() { + if uid_mailbox.uid == 0 { + uid_mailbox.uid = match self + .jmap + .assign_imap_uid(account_id, uid_mailbox.mailbox_id) + .await + { + Ok(assigned_uid) => { + debug_assert!(assigned_uid > 0); + copied_ids.push((imap_id.uid, assigned_uid)); + + assigned_uid + } + Err(_) => { + return Err( + StatusResponse::database_failure().with_tag(&arguments.tag) + ); + } + }; + } } // Write changes @@ -239,7 +263,6 @@ impl SessionData { .log_child_update(Collection::Mailbox, src_mailbox.id.mailbox_id); did_move = true; } - copied_ids.push((imap_id, id)); } Err(MethodError::ServerUnavailable) => { response.rtype = ResponseType::No; @@ -280,7 +303,10 @@ impl SessionData { { Ok(Ok(email)) => { dest_change_id = email.change_id.into(); - copied_ids.push((imap_id, email.id.document_id())); + if let Some(assigned_uid) = email.imap_uids.first() { + debug_assert!(*assigned_uid > 0); + copied_ids.push((imap_id.uid, *assigned_uid)); + } } Ok(Err(err)) => { if err.type_ != SetErrorType::NotFound { @@ -310,7 +336,7 @@ impl SessionData { }; // Make sure the message is still in the mailbox - let src_mailbox_id = UidMailbox::from(src_mailbox.id.mailbox_id); + let src_mailbox_id = UidMailbox::new_unassigned(src_mailbox.id.mailbox_id); if !mailboxes.current().contains(&src_mailbox_id) { continue; } else if mailboxes.current().len() == 1 { @@ -410,22 +436,16 @@ impl SessionData { .with_tag(arguments.tag)); } - let dest_mailbox = self - .fetch_messages(&dest_mailbox) + // Prepare response + let uid_validity = self + .get_uid_validity(&dest_mailbox) .await .map_err(|r| r.with_tag(&arguments.tag))?; - - // Prepare response - let uid_validity = dest_mailbox.uid_validity; let mut src_uids = Vec::with_capacity(copied_ids.len()); let mut dest_uids = Vec::with_capacity(copied_ids.len()); - for (src_id, dest_id) in copied_ids { - if let Some(dest_uid) = dest_mailbox.id_to_imap.get(&dest_id) { - src_uids.push(src_id.uid); - dest_uids.push(dest_uid.uid); - } else { - tracing::debug!("Could not map JMAP ID {} to IMAP UID", dest_id); - } + for (src_uid, dest_uid) in copied_ids { + src_uids.push(src_uid); + dest_uids.push(dest_uid); } src_uids.sort_unstable(); dest_uids.sort_unstable(); diff --git a/crates/imap/src/op/create.rs b/crates/imap/src/op/create.rs index b09dc38e..6b8ad4fd 100644 --- a/crates/imap/src/op/create.rs +++ b/crates/imap/src/op/create.rs @@ -34,7 +34,7 @@ use jmap_proto::{ type_state::DataType, value::Value, }, }; -use store::{query::Filter, roaring::RoaringBitmap, write::BatchBuilder}; +use store::{query::Filter, write::BatchBuilder}; use utils::listener::SessionStream; use crate::core::{Account, Mailbox, Session, SessionData}; @@ -217,7 +217,6 @@ impl SessionData { } else { None }, - recent_messages: RoaringBitmap::new(), }, ); } diff --git a/crates/imap/src/op/expunge.rs b/crates/imap/src/op/expunge.rs index fc832bc9..c7b6c910 100644 --- a/crates/imap/src/op/expunge.rs +++ b/crates/imap/src/op/expunge.rs @@ -177,7 +177,7 @@ impl SessionData { } else { continue; }; - let mailbox_id = UidMailbox::from(mailbox_id); + let mailbox_id = UidMailbox::new_unassigned(mailbox_id); if !mailboxes.current().contains(&mailbox_id) { continue; } else if mailboxes.current().len() > 1 { diff --git a/crates/imap/src/op/fetch.rs b/crates/imap/src/op/fetch.rs index 4ea3c7af..231a2433 100644 --- a/crates/imap/src/op/fetch.rs +++ b/crates/imap/src/op/fetch.rs @@ -106,7 +106,7 @@ impl SessionData { mailbox: Arc, is_uid: bool, is_qresync: bool, - is_rev2: bool, + _is_rev2: bool, enabled_condstore: bool, ) -> StatusResponse { // Validate VANISHED parameter @@ -126,11 +126,6 @@ impl SessionData { Ok(modseq) => modseq, Err(response) => return response.with_tag(arguments.tag), }; - let recent_messages = if !is_rev2 { - self.get_recent(&mailbox.id).into() - } else { - None - }; // Convert IMAP ids to JMAP ids. let mut ids = match mailbox @@ -371,12 +366,6 @@ impl SessionData { if set_seen_flag { flags.push(Flag::Seen); } - if recent_messages - .as_ref() - .map_or(false, |recent| recent.contains(id)) - { - flags.push(Flag::Recent); - } items.push(DataItem::Flags { flags }); } Attribute::InternalDate => { diff --git a/crates/imap/src/op/search.rs b/crates/imap/src/op/search.rs index 39b55274..47a67582 100644 --- a/crates/imap/src/op/search.rs +++ b/crates/imap/src/op/search.rs @@ -570,10 +570,10 @@ impl SessionData { filters.push(query::Filter::End); } search::Filter::Recent => { - filters.push(query::Filter::is_in_set(self.get_recent(&mailbox.id))); + //filters.push(query::Filter::is_in_set(self.get_recent(&mailbox.id))); } search::Filter::New => { - filters.push(query::Filter::And); + /*filters.push(query::Filter::And); filters.push(query::Filter::is_in_set(self.get_recent(&mailbox.id))); filters.push(query::Filter::Not); filters.push(query::Filter::is_in_bitmap( @@ -581,12 +581,12 @@ impl SessionData { Keyword::Seen, )); filters.push(query::Filter::End); - filters.push(query::Filter::End); + filters.push(query::Filter::End);*/ } search::Filter::Old => { - filters.push(query::Filter::Not); + /*filters.push(query::Filter::Not); filters.push(query::Filter::is_in_set(self.get_recent(&mailbox.id))); - filters.push(query::Filter::End); + filters.push(query::Filter::End);*/ } search::Filter::Older(secs) => { filters.push(query::Filter::le( diff --git a/crates/imap/src/op/select.rs b/crates/imap/src/op/select.rs index 960403a8..3ac36be0 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, CachedItem}; +use utils::{listener::SessionStream, lru_cache::LruCached}; -use crate::core::{MailboxState, SavedSearch, SelectedMailbox, Session, State}; +use crate::core::{SavedSearch, SelectedMailbox, Session, State}; use super::ToModSeq; @@ -59,33 +59,33 @@ impl Session { if let Some(mailbox) = data.get_mailbox_by_name(&arguments.mailbox_name) { // Try obtaining the mailbox from the cache let state = { - let cached_state_ = self - .imap - .cache_mailbox - .entry(mailbox) - .or_insert_with(|| CachedItem::new(MailboxState::default())); - let mut cached_state = cached_state_.get().await; - let is_cache_miss = - cached_state.uid_validity == 0 && cached_state.uid_max == 0; - let mut modseq = None; - - if !is_cache_miss { - match data.get_modseq(mailbox.account_id).await { - Ok(modseq_) => { - modseq = modseq_; - } - Err(mut response) => { - response.tag = arguments.tag.into(); - return self.write_bytes(response.into_bytes()).await; - } + let modseq = match data.get_modseq(mailbox.account_id).await { + Ok(modseq) => modseq, + Err(mut response) => { + response.tag = arguments.tag.into(); + return self.write_bytes(response.into_bytes()).await; } - } + }; - // Refresh the mailbox if the modseq has changed or if it's a cache miss - if is_cache_miss || cached_state.modseq.unwrap_or(0) < modseq.unwrap_or(0) { + if let Some(cached_state) = + self.imap + .cache_mailbox + .get(&mailbox) + .and_then(|cached_state| { + if cached_state.modseq.unwrap_or(0) >= modseq.unwrap_or(0) { + Some(cached_state) + } else { + None + } + }) + { + cached_state.as_ref().clone() + } else { match data.fetch_messages(&mailbox).await { Ok(new_state) => { - *cached_state = new_state; + let new_state = Arc::new(new_state); + self.imap.cache_mailbox.insert(mailbox, new_state.clone()); + new_state.as_ref().clone() } Err(mut response) => { response.tag = arguments.tag.into(); @@ -93,8 +93,6 @@ impl Session { } } } - - (*cached_state).clone() }; // Synchronize messages @@ -160,7 +158,7 @@ impl Session { let response = Response { mailbox: ListItem::new(arguments.mailbox_name), total_messages, - recent_messages: data.get_recent_count(&mailbox.id), + recent_messages: 0, unseen_seq: 0, uid_validity, uid_next, diff --git a/crates/imap/src/op/status.rs b/crates/imap/src/op/status.rs index 942cf65b..839ca986 100644 --- a/crates/imap/src/op/status.rs +++ b/crates/imap/src/op/status.rs @@ -34,7 +34,9 @@ use jmap_proto::{ types::{collection::Collection, id::Id, keyword::Keyword, property::Property, value::Value}, }; use store::{ - roaring::RoaringBitmap, write::key::DeserializeBigEndian, IndexKeyPrefix, IterateParams, + roaring::RoaringBitmap, + write::{key::DeserializeBigEndian, ValueClass}, + IndexKeyPrefix, IterateParams, ValueKey, }; use store::{Deserialize, U32_LEN}; use utils::listener::SessionStream; @@ -201,10 +203,7 @@ impl SessionData { } Status::Recent => { if !update_recent { - items_response.push(( - *item, - StatusItemType::Number(mailbox_state.recent_messages.len()), - )); + items_response.push((*item, StatusItemType::Number(0))); } else { items_update.push_unique(*item); } @@ -239,14 +238,24 @@ impl SessionData { Status::UidNext => { (self .jmap - .get_property::( - mailbox.account_id, - Collection::Mailbox, - mailbox.mailbox_id, - Property::EmailIds, - ) - .await? - .unwrap_or(0) + .store + .get_counter(ValueKey { + account_id: mailbox.account_id, + collection: Collection::Mailbox.into(), + document_id: mailbox.mailbox_id, + class: ValueClass::Property(Property::EmailIds.into()), + }) + .await + .map_err(|err| { + tracing::debug!(event = "error", + context = "store", + account_id = mailbox.account_id, + collection = ?Collection::Mailbox, + mailbox_id = mailbox.mailbox_id, + reason = ?err, + "Failed to obtain uid next"); + StatusResponse::no("Mailbox unavailable.") + })? + 1) as u64 } Status::UidValidity => self @@ -352,8 +361,7 @@ impl SessionData { .iter_mut() .find(|(i, _)| *i == Status::Recent) .unwrap() - .1 = - StatusItemType::Number(mailbox_state.recent_messages.len()); + .1 = StatusItemType::Number(0); } Status::HighestModSeq | Status::MailboxId => { unreachable!() diff --git a/crates/jmap/src/api/config.rs b/crates/jmap/src/api/config.rs index 083dd4a7..2b9f3a3b 100644 --- a/crates/jmap/src/api/config.rs +++ b/crates/jmap/src/api/config.rs @@ -177,9 +177,6 @@ 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/blob/get.rs b/crates/jmap/src/blob/get.rs index 2ec97e0c..7e0eaa94 100644 --- a/crates/jmap/src/blob/get.rs +++ b/crates/jmap/src/blob/get.rs @@ -254,7 +254,10 @@ impl JMAP { DataType::Mailbox, mailboxes .into_iter() - .map(|m| Id::from(m.mailbox_id)) + .map(|m| { + debug_assert!(m.uid != 0); + Id::from(m.mailbox_id) + }) .collect::>(), ); } diff --git a/crates/jmap/src/email/cache.rs b/crates/jmap/src/email/cache.rs index 2df0eb26..89e60e51 100644 --- a/crates/jmap/src/email/cache.rs +++ b/crates/jmap/src/email/cache.rs @@ -21,14 +21,14 @@ * for more details. */ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use futures_util::TryFutureExt; use jmap_proto::{ error::method::MethodError, types::{collection::Collection, property::Property}, }; -use utils::CachedItem; +use utils::lru_cache::LruCached; use crate::JMAP; @@ -59,21 +59,32 @@ impl JMAP { .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.unwrap_or(0) < modseq.unwrap_or(0) { - thread_cache.threads = self - .get_properties::(account_id, Collection::Email, &(), Property::ThreadId) - .await? - .into_iter() - .collect(); - thread_cache.modseq = modseq; - } + let thread_cache = if let Some(thread_cache) = + self.cache_threads.get(&account_id).and_then(|t| { + if t.modseq.unwrap_or(0) >= modseq.unwrap_or(0) { + Some(t) + } else { + None + } + }) { + thread_cache + } else { + let thread_cache = Arc::new(Threads { + threads: self + .get_properties::( + account_id, + Collection::Email, + &(), + Property::ThreadId, + ) + .await? + .into_iter() + .collect(), + modseq, + }); + self.cache_threads.insert(account_id, thread_cache.clone()); + thread_cache + }; // Obtain threadIds for matching messages let mut thread_ids = Vec::with_capacity(message_ids.size_hint().0); diff --git a/crates/jmap/src/email/copy.rs b/crates/jmap/src/email/copy.rs index 7c4cb5a6..ae9e7d46 100644 --- a/crates/jmap/src/email/copy.rs +++ b/crates/jmap/src/email/copy.rs @@ -375,6 +375,25 @@ impl JMAP { ..Default::default() }; + // Assign IMAP UIDs + let mut mailbox_ids = Vec::with_capacity(mailboxes.len()); + email.imap_uids = Vec::with_capacity(mailboxes.len()); + for mailbox_id in &mailboxes { + let uid = self + .assign_imap_uid(account_id, *mailbox_id) + .await + .map_err(|err| { + tracing::error!( + event = "error", + context = "email_copy", + error = ?err, + "Failed to assign IMAP UID."); + MethodError::ServerPartialFail + })?; + mailbox_ids.push(UidMailbox::new(*mailbox_id, uid)); + email.imap_uids.push(uid); + } + // Prepare batch let mut batch = BatchBuilder::new(); batch.with_account_id(account_id); @@ -406,14 +425,7 @@ impl JMAP { .with_collection(Collection::Email) .create_document(message_id) .value(Property::ThreadId, thread_id, F_VALUE | F_BITMAP) - .value( - Property::MailboxIds, - mailboxes - .into_iter() - .map(UidMailbox::from) - .collect::>(), - F_VALUE | F_BITMAP, - ) + .value(Property::MailboxIds, mailbox_ids, F_VALUE | F_BITMAP) .value(Property::Keywords, keywords, F_VALUE | F_BITMAP) .value(Property::Cid, changes.change_id, F_VALUE) .set( diff --git a/crates/jmap/src/email/get.rs b/crates/jmap/src/email/get.rs index ddd403e1..ad32e585 100644 --- a/crates/jmap/src/email/get.rs +++ b/crates/jmap/src/email/get.rs @@ -224,6 +224,7 @@ impl JMAP { .map(|ids| { let mut obj = Object::with_capacity(ids.len()); for id in ids { + debug_assert!(id.uid != 0); obj.append( Property::_T(Id::from(id.mailbox_id).to_string()), true, diff --git a/crates/jmap/src/email/ingest.rs b/crates/jmap/src/email/ingest.rs index b5caeea5..6c0b5449 100644 --- a/crates/jmap/src/email/ingest.rs +++ b/crates/jmap/src/email/ingest.rs @@ -64,6 +64,7 @@ pub struct IngestedEmail { pub change_id: u64, pub blob_id: BlobId, pub size: usize, + pub imap_uids: Vec, } pub struct IngestEmail<'x> { @@ -193,6 +194,7 @@ impl JMAP { id: Id::default(), change_id: u64::MAX, blob_id: BlobId::default(), + imap_uids: Vec::new(), size: 0, }); } @@ -295,6 +297,25 @@ impl JMAP { IngestError::Temporary })?; + // Assign IMAP UIDs + let mut mailbox_ids = Vec::with_capacity(params.mailbox_ids.len()); + let mut imap_uids = Vec::with_capacity(params.mailbox_ids.len()); + for mailbox_id in ¶ms.mailbox_ids { + let uid = self + .assign_imap_uid(params.account_id, *mailbox_id) + .await + .map_err(|err| { + tracing::error!( + event = "error", + context = "email_ingest", + error = ?err, + "Failed to assign IMAP UID."); + IngestError::Temporary + })?; + mailbox_ids.push(UidMailbox::new(*mailbox_id, uid)); + imap_uids.push(uid); + } + // Prepare batch let mut batch = BatchBuilder::new(); batch.with_account_id(params.account_id); @@ -337,11 +358,7 @@ impl JMAP { message, blob_id.hash.clone(), params.keywords, - params - .mailbox_ids - .iter() - .map(|id| UidMailbox::from(*id)) - .collect(), + mailbox_ids, params.received_at.unwrap_or_else(now), ) .value(Property::Cid, change_id, F_VALUE) @@ -390,6 +407,7 @@ impl JMAP { section: blob_id.section, }, size: raw_message_len as usize, + imap_uids, }) } @@ -558,6 +576,20 @@ impl JMAP { } } } + + pub async fn assign_imap_uid(&self, account_id: u32, mailbox_id: u32) -> store::Result { + // Increment UID next + let mut batch = BatchBuilder::new(); + batch + .with_account_id(account_id) + .with_collection(Collection::Mailbox) + .update_document(mailbox_id) + .add_and_get(Property::EmailIds, 1); + self.store + .write(batch.build()) + .await + .map(|v| v.expect("UID next") as u32) + } } impl From for Object { diff --git a/crates/jmap/src/email/set.rs b/crates/jmap/src/email/set.rs index 779243e8..b84cf7bb 100644 --- a/crates/jmap/src/email/set.rs +++ b/crates/jmap/src/email/set.rs @@ -21,7 +21,7 @@ * for more details. */ -use std::{borrow::Cow, collections::HashMap}; +use std::{borrow::Cow, collections::HashMap, slice::IterMut}; use jmap_proto::{ error::{ @@ -811,7 +811,8 @@ impl JMAP { mailboxes.set( ids.into_iter() .filter_map(|id| { - UidMailbox::from(id.try_unwrap_id()?.document_id()).into() + UidMailbox::new_unassigned(id.try_unwrap_id()?.document_id()) + .into() }) .collect(), ); @@ -820,7 +821,7 @@ impl JMAP { let mut patch = patch.into_iter(); if let Some(id) = patch.next().unwrap().try_unwrap_id() { mailboxes.update( - UidMailbox::from(id.document_id()), + UidMailbox::new_unassigned(id.document_id()), patch.next().unwrap().try_unwrap_bool().unwrap_or_default(), ); } @@ -957,6 +958,23 @@ impl JMAP { } } + // Obtain IMAP UIDs for added mailboxes + for uid_mailbox in mailboxes.inner_tags_mut() { + if uid_mailbox.uid == 0 { + uid_mailbox.uid = self + .assign_imap_uid(account_id, uid_mailbox.mailbox_id) + .await + .map_err(|err| { + tracing::error!( + event = "error", + context = "email_copy", + error = ?err, + "Failed to assign IMAP UID."); + MethodError::ServerPartialFail + })?; + } + } + // Update mailboxIds property mailboxes.update_batch(&mut batch, Property::MailboxIds); } @@ -1101,6 +1119,7 @@ impl JMAP { return Ok(Err(SetError::not_found())); }; for mailbox_id in &mailboxes.inner { + debug_assert!(mailbox_id.uid != 0); changes.log_child_update(Collection::Mailbox, mailbox_id.mailbox_id); } batch.assert_value(Property::MailboxIds, &mailboxes).value( @@ -1326,6 +1345,10 @@ impl< self.added.iter().chain(self.removed.iter()) } + pub fn inner_tags_mut(&mut self) -> IterMut<'_, T> { + self.current.inner.iter_mut() + } + pub fn has_tags(&self) -> bool { !self.current.inner.is_empty() } diff --git a/crates/jmap/src/lib.rs b/crates/jmap/src/lib.rs index 45ab15ab..966dd48f 100644 --- a/crates/jmap/src/lib.rs +++ b/crates/jmap/src/lib.rs @@ -59,9 +59,10 @@ use tokio::sync::mpsc; use utils::{ config::{Rate, Servers}, ipc::DeliveryEvent, + lru_cache::{LruCache, LruCached}, map::ttl_dashmap::{TtlDashMap, TtlMap}, snowflake::SnowflakeIdGenerator, - CachedItem, UnwrapFailure, + UnwrapFailure, }; pub mod api; @@ -102,7 +103,7 @@ pub struct JMAP { pub housekeeper_tx: mpsc::Sender, pub smtp: Arc, - pub cache_threads: DashMap>, + pub cache_threads: LruCache>, pub sieve_compiler: Compiler, pub sieve_runtime: Runtime<()>, @@ -158,8 +159,6 @@ 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)>, @@ -231,10 +230,8 @@ impl JMAP { config.property("cache.oauth.size")?.unwrap_or(128), shard_amount, ), - cache_threads: DashMap::with_capacity_and_hasher_and_shard_amount( + cache_threads: LruCache::with_capacity( config.property("cache.messages.size")?.unwrap_or(2048), - RandomState::default(), - shard_amount, ), state_tx, housekeeper_tx, @@ -747,27 +744,31 @@ impl JMAP { } pub async fn write_batch(&self, batch: BatchBuilder) -> Result<(), MethodError> { - self.store.write(batch.build()).await.map_err(|err| { - match err { - store::Error::InternalError(err) => { - tracing::error!( + self.store + .write(batch.build()) + .await + .map(|_| ()) + .map_err(|err| { + match err { + store::Error::InternalError(err) => { + tracing::error!( event = "error", context = "write_batch", error = ?err, "Failed to write batch."); - MethodError::ServerPartialFail + MethodError::ServerPartialFail + } + store::Error::AssertValueFailed => { + // This should not occur, as we are not using assertions. + tracing::debug!( + event = "assert_failed", + context = "write_batch", + "Failed to assert value." + ); + MethodError::ServerUnavailable + } } - store::Error::AssertValueFailed => { - // This should not occur, as we are not using assertions. - tracing::debug!( - event = "assert_failed", - context = "write_batch", - "Failed to assert value." - ); - MethodError::ServerUnavailable - } - } - }) + }) } } diff --git a/crates/jmap/src/mailbox/mod.rs b/crates/jmap/src/mailbox/mod.rs index bbc1ba91..b4a0ac1c 100644 --- a/crates/jmap/src/mailbox/mod.rs +++ b/crates/jmap/src/mailbox/mod.rs @@ -87,8 +87,12 @@ impl Serialize for UidMailbox { } } -impl From for UidMailbox { - fn from(mailbox_id: u32) -> Self { +impl UidMailbox { + pub fn new(mailbox_id: u32, uid: u32) -> Self { + UidMailbox { mailbox_id, uid } + } + + pub fn new_unassigned(mailbox_id: u32) -> Self { UidMailbox { mailbox_id, uid: 0 } } } diff --git a/crates/jmap/src/mailbox/set.rs b/crates/jmap/src/mailbox/set.rs index 26d4c916..0a3ed960 100644 --- a/crates/jmap/src/mailbox/set.rs +++ b/crates/jmap/src/mailbox/set.rs @@ -375,10 +375,10 @@ impl JMAP { ) .await? .and_then(|mut ids| { - let idx = ids - .inner - .iter() - .position(|&id| id.mailbox_id == document_id)?; + let idx = ids.inner.iter().position(|&id| { + debug_assert!(id.uid != 0); + id.mailbox_id == document_id + })?; ids.inner.swap_remove(idx); Some(ids) }) diff --git a/crates/jmap/src/sieve/ingest.rs b/crates/jmap/src/sieve/ingest.rs index f8a75e4d..e7d4276a 100644 --- a/crates/jmap/src/sieve/ingest.rs +++ b/crates/jmap/src/sieve/ingest.rs @@ -117,6 +117,7 @@ impl JMAP { change_id: u64::MAX, blob_id: Default::default(), size: raw_message.len(), + imap_uids: Vec::new(), }; while let Some(event) = instance.run(input) { diff --git a/crates/store/src/backend/foundationdb/mod.rs b/crates/store/src/backend/foundationdb/mod.rs index 71323eed..7ae21a62 100644 --- a/crates/store/src/backend/foundationdb/mod.rs +++ b/crates/store/src/backend/foundationdb/mod.rs @@ -43,3 +43,9 @@ impl From for Error { Self::InternalError(format!("FoundationDB error: {}", error.message())) } } + +fn deserialize_i64_le(bytes: &[u8]) -> crate::Result { + Ok(i64::from_le_bytes(bytes[..].try_into().map_err(|_| { + crate::Error::InternalError("Invalid counter value.".to_string()) + })?)) +} diff --git a/crates/store/src/backend/foundationdb/read.rs b/crates/store/src/backend/foundationdb/read.rs index 65875763..254e63ab 100644 --- a/crates/store/src/backend/foundationdb/read.rs +++ b/crates/store/src/backend/foundationdb/read.rs @@ -38,7 +38,7 @@ use crate::{ BitmapKey, Deserialize, IterateParams, Key, ValueKey, U32_LEN, WITH_SUBSPACE, }; -use super::{FdbStore, MAX_VALUE_SIZE}; +use super::{deserialize_i64_le, FdbStore, MAX_VALUE_SIZE}; #[cfg(feature = "fdb-chunked-bm")] pub(crate) enum ChunkedBitmap { @@ -158,9 +158,7 @@ impl FdbStore { ) -> crate::Result { let key = key.into().serialize(WITH_SUBSPACE); if let Some(bytes) = self.db.create_trx()?.get(&key, true).await? { - Ok(i64::from_le_bytes(bytes[..].try_into().map_err(|_| { - crate::Error::InternalError("Invalid counter value.".to_string()) - })?)) + deserialize_i64_le(&bytes) } else { Ok(0) } diff --git a/crates/store/src/backend/foundationdb/write.rs b/crates/store/src/backend/foundationdb/write.rs index 9fe091a2..d6992cff 100644 --- a/crates/store/src/backend/foundationdb/write.rs +++ b/crates/store/src/backend/foundationdb/write.rs @@ -45,6 +45,7 @@ use crate::{ }; use super::{ + deserialize_i64_le, read::{read_chunked_value, ChunkedValue}, FdbStore, MAX_VALUE_SIZE, }; @@ -69,7 +70,7 @@ impl BitmapOp { } impl FdbStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { let start = Instant::now(); let mut retry_count = 0; #[cfg(not(feature = "fdb-chunked-bm"))] @@ -83,6 +84,7 @@ impl FdbStore { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; + let mut result = None; let trx = self.db.create_trx()?; @@ -103,20 +105,6 @@ impl FdbStore { } => { document_id = *document_id_; } - Operation::Value { - class, - op: ValueOp::Add(by), - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(WITH_SUBSPACE); - - trx.atomic_op(&key, &by.to_le_bytes()[..], MutationType::Add); - } Operation::Value { class, op } => { let mut key = ValueKey { account_id, @@ -127,62 +115,79 @@ impl FdbStore { .serialize(WITH_SUBSPACE); let do_chunk = key[0] == SUBSPACE_VALUES; - if let ValueOp::Set(value) = op { - if !value.is_empty() && do_chunk { - for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() { - match pos.cmp(&1) { - Ordering::Less => {} - Ordering::Equal => { - key.push(0); - } - Ordering::Greater => { - if pos < u8::MAX as usize { - *key.last_mut().unwrap() += 1; - } else { - trx.cancel(); - return Err(crate::Error::InternalError( - "Value too large".into(), - )); + match op { + ValueOp::Set(value) => { + if !value.is_empty() && do_chunk { + for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() { + match pos.cmp(&1) { + Ordering::Less => {} + Ordering::Equal => { + key.push(0); + } + Ordering::Greater => { + if pos < u8::MAX as usize { + *key.last_mut().unwrap() += 1; + } else { + trx.cancel(); + return Err(crate::Error::InternalError( + "Value too large".into(), + )); + } } } + trx.set(&key, chunk); } - trx.set(&key, chunk); + } else { + trx.set(&key, value); } - } else { - trx.set(&key, value); - } - if matches!(class, ValueClass::ReservedId) { - let block_num = DenseBitmap::block_num(document_id); - if let Ok(Some(bytes)) = trx - .get( - &BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num, + if matches!(class, ValueClass::ReservedId) { + let block_num = DenseBitmap::block_num(document_id); + if let Ok(Some(bytes)) = trx + .get( + &BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + block_num, + } + .serialize(WITH_SUBSPACE), + true, + ) + .await + { + if block_contains(&bytes, block_num, document_id) { + trx.cancel(); + return Err(crate::Error::AssertValueFailed); } - .serialize(WITH_SUBSPACE), - true, - ) - .await - { - if block_contains(&bytes, block_num, document_id) { - trx.cancel(); - return Err(crate::Error::AssertValueFailed); } } } - } else if do_chunk { - trx.clear_range( - &key, - &KeySerializer::new(key.len() + 1) - .write(key.as_slice()) - .write(u8::MAX) - .finalize(), - ); - } else { - trx.clear(&key); + ValueOp::AtomicAdd(by) => { + trx.atomic_op(&key, &by.to_le_bytes()[..], MutationType::Add); + } + ValueOp::AddAndGet(by) => { + let num = if let Some(bytes) = trx.get(&key, false).await? { + deserialize_i64_le(&bytes)? + *by + } else { + *by + }; + trx.set(&key, &num.to_le_bytes()[..]); + result = Some(num); + } + ValueOp::Clear => { + if do_chunk { + trx.clear_range( + &key, + &KeySerializer::new(key.len() + 1) + .write(key.as_slice()) + .write(u8::MAX) + .finalize(), + ); + } else { + trx.clear(&key); + } + } } } Operation::Index { field, key, set } => { @@ -374,7 +379,7 @@ impl FdbStore { match trx.commit().await { Ok(_) => { - return Ok(()); + return Ok(result); } Err(err) => { if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME { diff --git a/crates/store/src/backend/mysql/write.rs b/crates/store/src/backend/mysql/write.rs index bd0c1a1c..2c4324d7 100644 --- a/crates/store/src/backend/mysql/write.rs +++ b/crates/store/src/backend/mysql/write.rs @@ -97,7 +97,7 @@ impl MysqlStore { } Operation::Value { class, - op: ValueOp::Add(by), + op: ValueOp::AtomicAdd(by), } => { let key = ValueKey { account_id, diff --git a/crates/store/src/backend/postgres/write.rs b/crates/store/src/backend/postgres/write.rs index fa86df72..07e284ae 100644 --- a/crates/store/src/backend/postgres/write.rs +++ b/crates/store/src/backend/postgres/write.rs @@ -38,19 +38,15 @@ use crate::{ use super::PostgresStore; impl PostgresStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { + 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) - }; + Ok(result) => { + return result; } Err(err) => match err.code() { Some( @@ -73,7 +69,7 @@ impl PostgresStore { &self, conn: &mut Object, batch: &Batch, - ) -> Result { + ) -> Result>, tokio_postgres::Error> { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; @@ -83,6 +79,7 @@ impl PostgresStore { .isolation_level(IsolationLevel::ReadCommitted) .start() .await?; + let mut result = None; for op in &batch.ops { match op { @@ -101,33 +98,6 @@ impl PostgresStore { } => { document_id = *document_id_; } - Operation::Value { - class, - op: ValueOp::Add(by), - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(0); - - 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, @@ -138,55 +108,87 @@ impl PostgresStore { let table = char::from(key.subspace()); let key = key.serialize(0); - 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? + match op { + ValueOp::Set(value) => { + 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)", + concat!( + "INSERT INTO {} (k, v) VALUES ($1, $2) ", + "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" + ), table )) .await? - } - } else { - trx.prepare_cached(&format!( - concat!( - "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, + if trx.execute(&s, &[&key, value]).await? == 0 { + return Ok(Err(crate::Error::AssertValueFailed)); } - .serialize(0); - if trx.query_opt(&s, &[&key]).await?.is_some() { - 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(0); + if trx.query_opt(&s, &[&key]).await?.is_some() { + return Ok(Err(crate::Error::AssertValueFailed)); + } } } - } else { - let s = trx - .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) - .await?; - trx.execute(&s, &[&key]).await?; + ValueOp::AtomicAdd(by) => { + 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?; + } + } + ValueOp::AddAndGet(by) => { + 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 RETURNING v" + )) + .await?; + result = trx + .query_one(&s, &[&key, &by]) + .await + .and_then(|row| row.try_get::<_, i64>(0))? + .into(); + } + ValueOp::Clear => { + let s = trx + .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) + .await?; + trx.execute(&s, &[&key]).await?; + } } } Operation::Index { field, key, set } => { @@ -277,14 +279,14 @@ impl PostgresStore { }) .unwrap_or_else(|| (false, assert_value.is_none())); if !matches { - return Ok(false); + return Ok(Err(crate::Error::AssertValueFailed)); } asserted_values.insert(key, exists); } } } - trx.commit().await.map(|_| true) + trx.commit().await.map(|_| Ok(result)) } pub(crate) async fn purge_store(&self) -> crate::Result<()> { diff --git a/crates/store/src/backend/postgres/write_dense.rs b/crates/store/src/backend/postgres/write_dense.rs index 91969f7b..67d1126e 100644 --- a/crates/store/src/backend/postgres/write_dense.rs +++ b/crates/store/src/backend/postgres/write_dense.rs @@ -218,7 +218,7 @@ impl PostgresStore { } Operation::Value { class, - op: ValueOp::Add(by), + op: ValueOp::AtomicAdd(by), } => { let key = ValueKey { account_id, diff --git a/crates/store/src/backend/postgres/write_roaring.rs b/crates/store/src/backend/postgres/write_roaring.rs index 9702d58f..ba550d64 100644 --- a/crates/store/src/backend/postgres/write_roaring.rs +++ b/crates/store/src/backend/postgres/write_roaring.rs @@ -242,7 +242,7 @@ impl PostgresStore { } Operation::Value { class, - op: ValueOp::Add(by), + op: ValueOp::AtomicAdd(by), } => { let key = ValueKey { account_id, diff --git a/crates/store/src/backend/rocksdb/write.rs b/crates/store/src/backend/rocksdb/write.rs index 0770a650..72c420da 100644 --- a/crates/store/src/backend/rocksdb/write.rs +++ b/crates/store/src/backend/rocksdb/write.rs @@ -204,7 +204,7 @@ impl<'x> RocksDBTransaction<'x> { } Operation::Value { class, - op: ValueOp::Add(by), + op: ValueOp::AtomicAdd(by), } => { let key = ValueKey { account_id, @@ -364,7 +364,7 @@ impl<'x> RocksDBTransaction<'x> { } Operation::Value { class, - op: ValueOp::Add(by), + op: ValueOp::AtomicAdd(by), } => { let key = ValueKey { account_id, diff --git a/crates/store/src/backend/sqlite/write.rs b/crates/store/src/backend/sqlite/write.rs index 83fb7622..27863606 100644 --- a/crates/store/src/backend/sqlite/write.rs +++ b/crates/store/src/backend/sqlite/write.rs @@ -31,13 +31,14 @@ use crate::{ use super::SqliteStore; impl SqliteStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { let mut conn = self.conn_pool.get()?; self.spawn_worker(move || { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; let trx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let mut result = None; for op in &batch.ops { match op { @@ -56,29 +57,6 @@ impl SqliteStore { } => { document_id = *document_id_; } - Operation::Value { - class, - op: ValueOp::Add(by), - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(0); - - if *by >= 0 { - trx.prepare_cached(concat!( - "INSERT INTO c (k, v) VALUES (?, ?) ", - "ON CONFLICT(k) DO UPDATE SET v = v + excluded.v" - ))? - .execute(params![&key, *by])?; - } else { - trx.prepare_cached("UPDATE c SET v = v + ? WHERE k = ?")? - .execute(params![*by, &key])?; - } - } Operation::Value { class, op } => { let key = ValueKey { account_id, @@ -89,35 +67,59 @@ impl SqliteStore { let table = char::from(key.subspace()); let key = key.serialize(0); - if let ValueOp::Set(value) = op { - trx.prepare_cached(&format!( - "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", - table - ))? - .execute([&key, value])?; + match op { + ValueOp::Set(value) => { + trx.prepare_cached(&format!( + "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", + table + ))? + .execute([&key, value])?; - if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: document_id, - } - .serialize(0); - if trx - .prepare_cached("SELECT 1 FROM b WHERE k = ?")? - .query_row([&key], |_| Ok(true)) - .optional()? - .unwrap_or(false) - { - trx.rollback()?; - return Err(crate::Error::AssertValueFailed); + if matches!(class, ValueClass::ReservedId) { + // Make sure the reserved id is not already in use + let key = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + block_num: document_id, + } + .serialize(0); + if trx + .prepare_cached("SELECT 1 FROM b WHERE k = ?")? + .query_row([&key], |_| Ok(true)) + .optional()? + .unwrap_or(false) + { + trx.rollback()?; + return Err(crate::Error::AssertValueFailed); + } } } - } else { - trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table))? - .execute([&key])?; + ValueOp::AtomicAdd(by) => { + if *by >= 0 { + trx.prepare_cached(concat!( + "INSERT INTO c (k, v) VALUES (?, ?) ", + "ON CONFLICT(k) DO UPDATE SET v = v + excluded.v" + ))? + .execute(params![&key, *by])?; + } else { + trx.prepare_cached("UPDATE c SET v = v + ? WHERE k = ?")? + .execute(params![*by, &key])?; + } + } + ValueOp::AddAndGet(by) => { + result = trx + .prepare_cached(concat!( + "INSERT INTO c (k, v) VALUES (?, ?) ", + "ON CONFLICT(k) DO UPDATE SET v = v + excluded.v RETURNING v" + ))? + .query_row(params![&key, &by], |row| row.get::<_, i64>(0))? + .into(); + } + ValueOp::Clear => { + trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table))? + .execute([&key])?; + } } } Operation::Index { field, key, set } => { @@ -198,7 +200,7 @@ impl SqliteStore { } } - trx.commit().map_err(Into::into) + trx.commit().map(|_| result).map_err(Into::into) }) .await } diff --git a/crates/store/src/dispatch/config.rs b/crates/store/src/dispatch/config.rs index 0d610474..774f2bdb 100644 --- a/crates/store/src/dispatch/config.rs +++ b/crates/store/src/dispatch/config.rs @@ -64,13 +64,13 @@ impl Store { for key in keys { batch.set(ValueClass::Config(key.key.into_bytes()), key.value); } - self.write(batch.build()).await + self.write(batch.build()).await.map(|_| ()) } pub async fn config_clear(&self, key: impl Into) -> crate::Result<()> { let mut batch = BatchBuilder::new(); batch.clear(ValueClass::Config(key.into().into_bytes())); - self.write(batch.build()).await + self.write(batch.build()).await.map(|_| ()) } pub async fn config_clear_prefix(&self, key: impl AsRef) -> crate::Result<()> { diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index 6f23e586..364a7089 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -76,7 +76,7 @@ impl LookupStore { .finalize(), ), }); - store.write(batch.build()).await + store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_set(key, value, expires).await, @@ -103,17 +103,6 @@ impl LookupStore { ) -> crate::Result { match self { LookupStore::Store(store) => { - let result = if return_value { - store - .get_counter(ValueKey::from(ValueClass::Lookup(LookupClass::Counter( - key.clone(), - )))) - .await? - + 1 - } else { - 0 - }; - let mut batch = BatchBuilder::new(); if let Some(expires) = expires { @@ -129,12 +118,14 @@ impl LookupStore { batch.ops.push(Operation::Value { class: ValueClass::Lookup(LookupClass::Counter(key)), - op: ValueOp::Add(value), + op: if return_value { + ValueOp::AddAndGet(value) + } else { + ValueOp::AtomicAdd(value) + }, }); - store.write(batch.build()).await?; - - Ok(result) + store.write(batch.build()).await.map(|r| r.unwrap_or(0)) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_incr(key, value, expires).await, @@ -152,7 +143,7 @@ impl LookupStore { class: ValueClass::Lookup(LookupClass::Key(key)), op: ValueOp::Clear, }); - store.write(batch.build()).await + store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_delete(key).await, @@ -170,7 +161,7 @@ impl LookupStore { class: ValueClass::Lookup(LookupClass::Counter(key)), op: ValueOp::Clear, }); - store.write(batch.build()).await + store.write(batch.build()).await.map(|_| ()) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_delete(key).await, diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index 9ef51a27..e5f2eec4 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -133,7 +133,7 @@ impl Store { } } - pub async fn write(&self, batch: Batch) -> crate::Result<()> { + pub async fn write(&self, batch: Batch) -> crate::Result> { #[cfg(feature = "test_mode")] if std::env::var("PARANOID_WRITE").map_or(false, |v| v == "1") { use crate::write::Operation; @@ -211,7 +211,7 @@ impl Store { } } - return Ok(()); + return Ok(None); } match self { diff --git a/crates/store/src/write/batch.rs b/crates/store/src/write/batch.rs index e191de6c..a82634f0 100644 --- a/crates/store/src/write/batch.rs +++ b/crates/store/src/write/batch.rs @@ -145,7 +145,15 @@ impl BatchBuilder { pub fn add(&mut self, class: impl Into, value: i64) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), - op: ValueOp::Add(value), + op: ValueOp::AtomicAdd(value), + }); + self + } + + pub fn add_and_get(&mut self, class: impl Into, value: i64) -> &mut Self { + self.ops.push(Operation::Value { + class: class.into(), + op: ValueOp::AddAndGet(value), }); self } diff --git a/crates/store/src/write/key.rs b/crates/store/src/write/key.rs index bb24dd5b..7b8c3bdc 100644 --- a/crates/store/src/write/key.rs +++ b/crates/store/src/write/key.rs @@ -220,15 +220,14 @@ impl Key for LogKey { impl + Sync + Send> Key for ValueKey { fn subspace(&self) -> u8 { - if !matches!( - self.class.as_ref(), + match self.class.as_ref() { ValueClass::Directory(DirectoryClass::UsedQuota(_)) - | ValueClass::Lookup(LookupClass::Counter(_)) - | ValueClass::Queue(QueueClass::QuotaCount(_) | QueueClass::QuotaSize(_)) - ) { - SUBSPACE_VALUES - } else { - SUBSPACE_COUNTERS + | ValueClass::Lookup(LookupClass::Counter(_)) + | ValueClass::Queue(QueueClass::QuotaCount(_) | QueueClass::QuotaSize(_)) => { + SUBSPACE_COUNTERS + } + ValueClass::Property(84) if self.collection == 1 => SUBSPACE_COUNTERS, // TODO: Find a more elegant way to do this + _ => SUBSPACE_VALUES, } } diff --git a/crates/store/src/write/mod.rs b/crates/store/src/write/mod.rs index 7db6e2dd..53c3d42d 100644 --- a/crates/store/src/write/mod.rs +++ b/crates/store/src/write/mod.rs @@ -197,7 +197,8 @@ pub struct ReportEvent { #[derive(Debug, PartialEq, Eq, Hash, Default)] pub enum ValueOp { Set(Vec), - Add(i64), + AtomicAdd(i64), + AddAndGet(i64), #[default] Clear, } @@ -257,6 +258,12 @@ impl Serialize for u64 { } } +impl Serialize for i64 { + fn serialize(self) -> Vec { + self.to_be_bytes().to_vec() + } +} + impl Serialize for u16 { fn serialize(self) -> Vec { self.to_be_bytes().to_vec() diff --git a/crates/utils/Cargo.toml b/crates/utils/Cargo.toml index c670cf79..76a1f2b2 100644 --- a/crates/utils/Cargo.toml +++ b/crates/utils/Cargo.toml @@ -40,6 +40,7 @@ futures = "0.3" proxy-header = { version = "0.1.0", features = ["tokio"] } regex = "1.7.0" blake3 = "1.3.3" +lru-cache = "0.1.2" [target.'cfg(unix)'.dependencies] privdrop = "0.5.3" diff --git a/crates/utils/src/lib.rs b/crates/utils/src/lib.rs index 2d24a4fc..d0842f40 100644 --- a/crates/utils/src/lib.rs +++ b/crates/utils/src/lib.rs @@ -21,11 +21,7 @@ * for more details. */ -use std::{ - collections::HashMap, - sync::{atomic::AtomicU64, Arc}, - time::SystemTime, -}; +use std::{collections::HashMap, sync::Arc}; use config::Config; @@ -35,6 +31,7 @@ pub mod config; pub mod expr; pub mod ipc; pub mod listener; +pub mod lru_cache; pub mod map; pub mod snowflake; pub mod suffixlist; @@ -115,40 +112,6 @@ impl AsMut<[u8]> for BlobHash { } } -#[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; } diff --git a/crates/utils/src/lru_cache.rs b/crates/utils/src/lru_cache.rs new file mode 100644 index 00000000..af08f28d --- /dev/null +++ b/crates/utils/src/lru_cache.rs @@ -0,0 +1,58 @@ +/* + * 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 std::{borrow::Borrow, hash::Hash}; + +use parking_lot::Mutex; + +pub type LruCache = Mutex>; + +pub trait LruCached: Sized { + fn with_capacity(capacity: usize) -> Self; + fn get(&self, name: &Q) -> Option + where + K: Borrow, + Q: Hash + Eq; + fn insert(&self, name: K, value: V) -> Option; +} + +impl LruCached for LruCache { + fn with_capacity(capacity: usize) -> Self { + Mutex::new(lru_cache::LruCache::with_hasher( + capacity, + ahash::RandomState::new(), + )) + } + + fn get(&self, name: &Q) -> Option + where + K: Borrow, + Q: Hash + Eq, + { + self.lock().get_mut(name).map(|entry| entry.clone()) + } + + fn insert(&self, name: K, item: V) -> Option { + self.lock().insert(name, item) + } +} diff --git a/tests/src/imap/acl.rs b/tests/src/imap/acl.rs index f5744d07..3ab547f7 100644 --- a/tests/src/imap/acl.rs +++ b/tests/src/imap/acl.rs @@ -29,6 +29,7 @@ use super::{append::assert_append_message, AssertResult, ImapConnection, Type}; pub async fn test(mut imap_john: &mut ImapConnection, _imap_check: &mut ImapConnection) { // Delivery to support account + println!("Running ACL tests..."); let mut lmtp = SmtpConnection::connect_port(11201).await; lmtp.ingest( "bill@example.com", diff --git a/tests/src/imap/append.rs b/tests/src/imap/append.rs index f503d358..358b353f 100644 --- a/tests/src/imap/append.rs +++ b/tests/src/imap/append.rs @@ -30,6 +30,8 @@ use crate::jmap::wait_for_index; use super::{resources_dir, AssertResult, IMAPTest, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, _imap_check: &mut ImapConnection, handle: &IMAPTest) { + println!("Running APPEND tests..."); + // Invalid APPEND commands imap.send("APPEND \"Does not exist\" {1+}\r\na").await; imap.assert_read(Type::Tagged, ResponseType::No) diff --git a/tests/src/imap/basic.rs b/tests/src/imap/basic.rs index e52c86be..1bbfec19 100644 --- a/tests/src/imap/basic.rs +++ b/tests/src/imap/basic.rs @@ -29,6 +29,8 @@ use mail_send::Credentials; use super::{AssertResult, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, _imap_check: &mut ImapConnection) { + println!("Running basic tests..."); + // Test CAPABILITY imap.send("CAPABILITY").await; imap.assert_read(Type::Tagged, ResponseType::Ok).await; diff --git a/tests/src/imap/body_structure.rs b/tests/src/imap/body_structure.rs index e26b3efb..2ce32b36 100644 --- a/tests/src/imap/body_structure.rs +++ b/tests/src/imap/body_structure.rs @@ -34,6 +34,8 @@ use super::resources_dir; #[test] fn body_structure() { + println!("Running BODYSTRUCTURE..."); + for file_name in fs::read_dir(resources_dir()).unwrap() { let mut file_name = file_name.as_ref().unwrap().path(); if file_name.extension().map_or(true, |e| e != "txt") { diff --git a/tests/src/imap/condstore.rs b/tests/src/imap/condstore.rs index 1e5a4f9f..d8146022 100644 --- a/tests/src/imap/condstore.rs +++ b/tests/src/imap/condstore.rs @@ -31,6 +31,8 @@ use crate::imap::{ use super::{ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, imap_check: &mut ImapConnection) { + println!("Running CONDSTORE..."); + // Test CONDSTORE parameter imap.send("SELECT INBOX (CONDSTORE)").await; let hms = imap diff --git a/tests/src/imap/copy_move.rs b/tests/src/imap/copy_move.rs index fac79573..9b3ccbf0 100644 --- a/tests/src/imap/copy_move.rs +++ b/tests/src/imap/copy_move.rs @@ -26,6 +26,8 @@ use imap_proto::ResponseType; use super::{AssertResult, ImapConnection, Type}; pub async fn test(_imap: &mut ImapConnection, imap_check: &mut ImapConnection) { + println!("Running COPY/MOVE tests..."); + // Check status imap_check .send("LIST \"\" % RETURN (STATUS (UIDNEXT MESSAGES UNSEEN SIZE RECENT))") @@ -77,13 +79,13 @@ pub async fn test(_imap: &mut ImapConnection, imap_check: &mut ImapConnection) { .assert_read(Type::Tagged, ResponseType::Ok) .await .assert_contains("MESSAGES 4") - .assert_contains("RECENT 4") + //.assert_contains("RECENT 4") .assert_contains("UNSEEN 4") .assert_contains("UIDNEXT 5") .assert_contains("SIZE 5851"); // Check \Recent flag - imap_check.send("SELECT \"Scamorza Affumicata\"").await; + /*imap_check.send("SELECT \"Scamorza Affumicata\"").await; imap_check .assert_read(Type::Tagged, ResponseType::Ok) .await @@ -115,9 +117,11 @@ pub async fn test(_imap: &mut ImapConnection, imap_check: &mut ImapConnection) { imap_check .assert_read(Type::Tagged, ResponseType::Ok) .await - .assert_count("\\Recent", 0); + .assert_count("\\Recent", 0);*/ // Move all messages to Burrata + imap_check.send("SELECT \"Scamorza Affumicata\"").await; + imap_check.assert_read(Type::Tagged, ResponseType::Ok).await; imap_check.send("MOVE 1:* \"Burrata al Tartufo\"").await; imap_check .assert_read(Type::Tagged, ResponseType::Ok) diff --git a/tests/src/imap/fetch.rs b/tests/src/imap/fetch.rs index 1df018e1..f5c6fe48 100644 --- a/tests/src/imap/fetch.rs +++ b/tests/src/imap/fetch.rs @@ -26,6 +26,8 @@ use imap_proto::ResponseType; use super::{AssertResult, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, _imap_check: &mut ImapConnection) { + println!("Running FETCH tests..."); + // Examine INBOX imap.send("EXAMINE INBOX").await; imap.assert_read(Type::Tagged, ResponseType::Ok) diff --git a/tests/src/imap/idle.rs b/tests/src/imap/idle.rs index fa31563c..21b73007 100644 --- a/tests/src/imap/idle.rs +++ b/tests/src/imap/idle.rs @@ -28,6 +28,8 @@ use crate::jmap::delivery::SmtpConnection; use super::{AssertResult, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, imap_check: &mut ImapConnection) { + println!("Running IDLE tests..."); + // Switch connection to IDLE mode imap_check.send("CREATE Parmeggiano").await; imap_check.assert_read(Type::Tagged, ResponseType::Ok).await; diff --git a/tests/src/imap/mailbox.rs b/tests/src/imap/mailbox.rs index d902a110..84517774 100644 --- a/tests/src/imap/mailbox.rs +++ b/tests/src/imap/mailbox.rs @@ -27,6 +27,8 @@ use imap_proto::ResponseType; use super::{AssertResult, ImapConnection, Type}; pub async fn test(mut imap: &mut ImapConnection, mut imap_check: &mut ImapConnection) { + println!("Running mailbox tests..."); + // Create third connection for testing let mut other_conn = ImapConnection::connect(b"_z ").await; other_conn diff --git a/tests/src/imap/managesieve.rs b/tests/src/imap/managesieve.rs index 5ccf7749..e8f7e161 100644 --- a/tests/src/imap/managesieve.rs +++ b/tests/src/imap/managesieve.rs @@ -35,6 +35,8 @@ use tokio_rustls::client::TlsStream; use super::AssertResult; pub async fn test() { + println!("Running ManageSieve tests..."); + // Connect to ManageSieve let mut sieve = SieveConnection::connect().await; sieve diff --git a/tests/src/imap/search.rs b/tests/src/imap/search.rs index 40770883..b1bfd12a 100644 --- a/tests/src/imap/search.rs +++ b/tests/src/imap/search.rs @@ -26,6 +26,8 @@ use imap_proto::ResponseType; use super::{AssertResult, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, imap_check: &mut ImapConnection) { + println!("Running SEARCH tests..."); + // Searches without selecting a mailbox should fail. imap.send("SEARCH RETURN (MIN MAX COUNT ALL) ALL").await; imap.assert_read(Type::Tagged, ResponseType::Bad).await; diff --git a/tests/src/imap/store.rs b/tests/src/imap/store.rs index c4c4993f..bc0fd6e6 100644 --- a/tests/src/imap/store.rs +++ b/tests/src/imap/store.rs @@ -28,6 +28,8 @@ use crate::jmap::wait_for_index; use super::{AssertResult, IMAPTest, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, _imap_check: &mut ImapConnection, handle: &IMAPTest) { + println!("Running STORE tests..."); + // Select INBOX imap.send("SELECT INBOX").await; imap.assert_read(Type::Tagged, ResponseType::Ok) diff --git a/tests/src/imap/thread.rs b/tests/src/imap/thread.rs index ca6ff2d0..cf043655 100644 --- a/tests/src/imap/thread.rs +++ b/tests/src/imap/thread.rs @@ -28,6 +28,8 @@ use crate::imap::{expand_uid_list, AssertResult}; use super::{append::build_messages, ImapConnection, Type}; pub async fn test(imap: &mut ImapConnection, _imap_check: &mut ImapConnection) { + println!("Running THREAD tests..."); + // Create test messages let messages = build_messages(); diff --git a/tests/src/jmap/email_query.rs b/tests/src/jmap/email_query.rs index 3dbe1e9b..4cc42fdc 100644 --- a/tests/src/jmap/email_query.rs +++ b/tests/src/jmap/email_query.rs @@ -32,10 +32,13 @@ use jmap_client::{ core::query::{Comparator, Filter}, email, }; -use jmap_proto::types::{collection::Collection, id::Id}; +use jmap_proto::types::{collection::Collection, id::Id, property::Property}; use mail_parser::HeaderName; -use store::{ahash::AHashMap, write::BatchBuilder}; +use store::{ + ahash::AHashMap, + write::{BatchBuilder, ValueClass}, +}; use super::JMAPTest; @@ -57,7 +60,7 @@ pub async fn test(params: &mut JMAPTest, insert: bool) { batch .with_account_id(account_id) .with_collection(Collection::Mailbox); - for mailbox_id in 0..99999 { + for mailbox_id in 1545..3010 { batch.create_document(mailbox_id); } server.store.write(batch.build()).await.unwrap(); @@ -71,8 +74,10 @@ pub async fn test(params: &mut JMAPTest, insert: bool) { batch .with_account_id(account_id) .with_collection(Collection::Mailbox); - for mailbox_id in 0..99999 { - batch.delete_document(mailbox_id); + for mailbox_id in 1545..3010 { + batch + .delete_document(mailbox_id) + .clear(ValueClass::Property(Property::EmailIds.into())); } server.store.write(batch.build()).await.unwrap(); @@ -747,10 +752,10 @@ pub async fn create(client: &mut Client) { .email_import( format!( concat!( - "From: \"{}\" \nCc: \"{}\" \nMessage-ID: <{}>\n", - "References: <{}>\nComments: {}\nSubject: [{}]", - " Year {}\n\n{}\n{}\n" - ), + "From: \"{}\" \nCc: \"{}\" \nMessage-ID: <{}>\n", + "References: <{}>\nComments: {}\nSubject: [{}]", + " Year {}\n\n{}\n{}\n" + ), values_str["artist"], values_str["medium"], values_str["accession_number"], diff --git a/tests/src/jmap/stress_test.rs b/tests/src/jmap/stress_test.rs index e4e62399..5bacacab 100644 --- a/tests/src/jmap/stress_test.rs +++ b/tests/src/jmap/stress_test.rs @@ -256,6 +256,7 @@ async fn email_tests(server: Arc, client: Arc) { ); } let mailbox_tag = mailbox_tags[0]; + assert!(mailbox_tag.uid != 0); if mailbox_tag.mailbox_id != mailbox_id { panic!( concat!( diff --git a/tests/src/store/ops.rs b/tests/src/store/ops.rs index 5fa70593..ac32e10d 100644 --- a/tests/src/store/ops.rs +++ b/tests/src/store/ops.rs @@ -21,8 +21,10 @@ * for more details. */ +use std::collections::HashSet; + use store::{ - write::{BatchBuilder, ValueClass}, + write::{BatchBuilder, DirectoryClass, ValueClass}, Store, ValueKey, }; @@ -30,6 +32,46 @@ use store::{ const MAX_VALUE_SIZE: usize = 100000; pub async fn test(db: Store) { + // Increment a counter 1000 times concurrently + let mut handles = Vec::new(); + let mut assigned_ids = HashSet::new(); + println!("Incrementing counter 1000 times concurrently..."); + for _ in 0..1000 { + handles.push({ + let db = db.clone(); + tokio::spawn(async move { + let mut builder = BatchBuilder::new(); + builder + .with_account_id(0) + .with_collection(0) + .update_document(0) + .add_and_get(ValueClass::Directory(DirectoryClass::UsedQuota(0)), 1); + db.write(builder.build_batch()).await.unwrap().unwrap() + }) + }); + } + + for handle in handles { + let assigned_id = handle.await.unwrap(); + assert!( + assigned_ids.insert(assigned_id), + "counter assigned {assigned_id} twice or more times." + ); + } + assert_eq!(assigned_ids.len(), 1000); + assert_eq!( + db.get_counter(ValueKey { + account_id: 0, + collection: 0, + document_id: 0, + class: ValueClass::Directory(DirectoryClass::UsedQuota(0)), + }) + .await + .unwrap(), + 1000 + ); + + println!("Running chunking tests..."); for (test_num, value) in [ vec![b'A'; 0], vec![b'A'; 1], @@ -132,6 +174,7 @@ pub async fn test(db: Store) { .update_document(0) .clear(ValueClass::Property(0)) .clear(ValueClass::Property(2)) + .clear(ValueClass::Directory(DirectoryClass::UsedQuota(0))) .build_batch(), ) .await