From f80fe8b873bb818b2e8645fb0068fdfb57c30526 Mon Sep 17 00:00:00 2001 From: John Coffey Date: Thu, 8 Oct 2026 10:54:04 -0700 Subject: [PATCH] Read message archives in batches Store::get_values reads a set of keys in one round trip on PostgreSQL (k = ANY($1)) and falls back to one get_value per key elsewhere. A replicated store serves the batch from a replica and re-reads from the primary the keys the replica has not received yet. IMAP FETCH and STORE, JMAP Email/get and the message and mailbox cache refreshers that every node runs after a write now read their archives READ_BATCH_SIZE (128) at a time instead of one round trip per message. FETCH also prefetches the message data of the messages it is about to mark \Seen, saving the second round trip per message. --- crates/email/src/cache/email.rs | 43 +- crates/email/src/cache/mailbox.rs | 42 +- crates/imap/src/op/fetch.rs | 503 ++++++++++--------- crates/imap/src/op/store.rs | 269 ++++++----- crates/jmap/src/email/get.rs | 563 +++++++++++----------- crates/store/src/backend/postgres/read.rs | 39 ++ crates/store/src/dispatch/store.rs | 75 +++ crates/store/src/lib.rs | 4 + tests/src/store/ops.rs | 63 +++ 9 files changed, 933 insertions(+), 668 deletions(-) diff --git a/crates/email/src/cache/email.rs b/crates/email/src/cache/email.rs index fd7f539e..6607ebf3 100644 --- a/crates/email/src/cache/email.rs +++ b/crates/email/src/cache/email.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::message::metadata::{ArchivedMessageData, MessageData}; @@ -10,7 +12,7 @@ use common::{ sharing::EffectiveAcl, }; use store::write::{AlignedBytes, Archive}; -use store::{ValueKey, ahash::AHashMap, roaring::RoaringBitmap}; +use store::{READ_BATCH_SIZE, ValueKey, ahash::AHashMap, roaring::RoaringBitmap}; use trc::AddContext; use types::{ acl::Acl, @@ -41,23 +43,32 @@ pub(crate) async fn update_email_cache( keywords: store_cache.emails.keywords.to_vec(), }; - for (document_id, is_update) in changed_ids { - if *is_update - && let Some(archive) = server + // inbuxa: one round trip per batch of changed messages, not one per + // message; every node runs this after any write it hears about + let updated = changed_ids + .iter() + .filter(|(_, is_update)| **is_update) + .map(|(document_id, _)| *document_id) + .collect::>(); + for document_ids in updated.chunks(READ_BATCH_SIZE) { + let keys = document_ids + .iter() + .map(|document_id| ValueKey::archive(account_id, Collection::Email, *document_id)) + .collect::>(); + for (document_id, archive) in document_ids.iter().zip( + server .store() - .get_value::>(ValueKey::archive( - account_id, - Collection::Email, - *document_id, - )) + .get_values::>(&keys) .await - .caused_by(trc::location!())? - { - insert_item( - &mut new_cache, - *document_id, - archive.to_unarchived::()?, - ); + .caused_by(trc::location!())?, + ) { + if let Some(archive) = archive { + insert_item( + &mut new_cache, + *document_id, + archive.to_unarchived::()?, + ); + } } } diff --git a/crates/email/src/cache/mailbox.rs b/crates/email/src/cache/mailbox.rs index a315a6db..9b6e97c0 100644 --- a/crates/email/src/cache/mailbox.rs +++ b/crates/email/src/cache/mailbox.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::mailbox::{ArchivedMailbox, Mailbox, manage::MailboxFnc}; @@ -9,11 +11,11 @@ use common::{ MailboxCache, MailboxesCache, MessageStoreCache, Server, auth::AccessToken, sharing::EffectiveAcl, }; +use store::{READ_BATCH_SIZE, ahash::AHashMap, roaring::RoaringBitmap}; use store::{ ValueKey, write::{AlignedBytes, Archive}, }; -use store::{ahash::AHashMap, roaring::RoaringBitmap}; use trc::AddContext; use types::{ acl::{Acl, AclGrant}, @@ -42,23 +44,31 @@ pub(crate) async fn update_mailbox_cache( change_id: 0, }; - for (document_id, is_update) in changed_ids { - if *is_update - && let Some(archive) = server + // inbuxa: one round trip per batch of changed mailboxes (see email.rs) + let updated = changed_ids + .iter() + .filter(|(_, is_update)| **is_update) + .map(|(document_id, _)| *document_id) + .collect::>(); + for document_ids in updated.chunks(READ_BATCH_SIZE) { + let keys = document_ids + .iter() + .map(|document_id| ValueKey::archive(account_id, Collection::Mailbox, *document_id)) + .collect::>(); + for (document_id, archive) in document_ids.iter().zip( + server .store() - .get_value::>(ValueKey::archive( - account_id, - Collection::Mailbox, - *document_id, - )) + .get_values::>(&keys) .await - .caused_by(trc::location!())? - { - insert_item( - &mut new_cache, - *document_id, - archive.unarchive::()?, - ); + .caused_by(trc::location!())?, + ) { + if let Some(archive) = archive { + insert_item( + &mut new_cache, + *document_id, + archive.unarchive::()?, + ); + } } } diff --git a/crates/imap/src/op/fetch.rs b/crates/imap/src/op/fetch.rs index 7953d677..771d72a0 100644 --- a/crates/imap/src/op/fetch.rs +++ b/crates/imap/src/op/fetch.rs @@ -37,7 +37,7 @@ use imap_proto::{ use registry::schema::enums::Permission; use std::{borrow::Cow, sync::Arc, time::Instant}; use store::{ - ValueKey, + READ_BATCH_SIZE, ValueKey, write::{AlignedBytes, Archive}, }; use store::{ @@ -359,261 +359,294 @@ impl SessionData { .await .imap_ctx(&arguments.tag, trc::location!())?; - for (seqnum, uid, id) in ids { - // Obtain attributes and keywords - let (metadata_, data) = if let (Some(email), Some(data)) = ( - self.server - .store() - .get_value::>(ValueKey::property( - account_id, - Collection::Email, - id, - EmailField::Metadata, - )) - .await - .imap_ctx(&arguments.tag, trc::location!())?, - message_cache.email_by_id(&id), - ) { - (email, data) - } else { - trc::event!( - Store(trc::StoreEvent::NotFound), - AccountId = account_id, - DocumentId = id, - Collection = Collection::Email, - Details = "Message metadata not found.", - CausedBy = trc::location!(), - ); - continue; - }; - let metadata = metadata_ - .unarchive::() + for chunk in ids.chunks(READ_BATCH_SIZE) { + // inbuxa: the metadata of a batch of messages in one round trip, + // and the message data of those about to be marked \Seen in a + // second, instead of one or two round trips per message + let keys = chunk + .iter() + .map(|(_, _, id)| { + ValueKey::property(account_id, Collection::Email, *id, EmailField::Metadata) + }) + .collect::>(); + let metadatas = self + .server + .store() + .get_values::>(&keys) + .await .imap_ctx(&arguments.tag, trc::location!())?; - let raw_body; + let mut seen_archives = chunk + .iter() + .map(|_| None::>) + .collect::>(); + if set_seen_flags { + let unseen = chunk + .iter() + .enumerate() + .filter(|(_, (_, _, id))| { + message_cache + .email_by_id(id) + .is_some_and(|data| !message_cache.has_keyword(data, &Keyword::Seen)) + }) + .map(|(pos, (_, _, id))| (pos, *id)) + .collect::>(); + if !unseen.is_empty() { + let keys = unseen + .iter() + .map(|(_, id)| ValueKey::archive(account_id, Collection::Email, *id)) + .collect::>(); + for ((pos, _), archive) in unseen.into_iter().zip( + self.server + .store() + .get_values::>(&keys) + .await + .imap_ctx(&arguments.tag, trc::location!())?, + ) { + seen_archives[pos] = archive; + } + } + } - // Fetch and parse blob - let mut raw_message = ChainedBytes::new(metadata.raw_headers.as_ref()); - if needs_blobs { - // Retrieve raw message if needed - raw_body = self - .server - .blob_store() - .get_blob(metadata.blob_hash.0.as_slice(), 0..usize::MAX) - .await - .imap_ctx(&arguments.tag, trc::location!())?; - - if let Some(raw_body) = &raw_body { - raw_message.append( - raw_body - .get(metadata.blob_body_offset.to_native() as usize..) - .unwrap_or_default(), - ); + for (&(seqnum, uid, id), (metadata_, seen_archive)) in + chunk.iter().zip(metadatas.into_iter().zip(seen_archives)) + { + // Obtain attributes and keywords + let (metadata_, data) = if let (Some(email), Some(data)) = + (metadata_, message_cache.email_by_id(&id)) + { + (email, data) } else { trc::event!( Store(trc::StoreEvent::NotFound), AccountId = account_id, DocumentId = id, Collection = Collection::Email, - BlobId = metadata.blob_hash.0.as_slice(), - Details = "Blob not found.", + Details = "Message metadata not found.", CausedBy = trc::location!(), ); - continue; - } - } + }; + let metadata = metadata_ + .unarchive::() + .imap_ctx(&arguments.tag, trc::location!())?; + let raw_body; - let message = &metadata.contents[0]; - let decoded = metadata.decode_contents(raw_message.clone()); + // Fetch and parse blob + let mut raw_message = ChainedBytes::new(metadata.raw_headers.as_ref()); + if needs_blobs { + // Retrieve raw message if needed + raw_body = self + .server + .blob_store() + .get_blob(metadata.blob_hash.0.as_slice(), 0..usize::MAX) + .await + .imap_ctx(&arguments.tag, trc::location!())?; - // Build response - let mut items = Vec::with_capacity(arguments.attributes.len()); - let set_seen_flag = set_seen_flags && !message_cache.has_keyword(data, &Keyword::Seen); - - for attribute in &arguments.attributes { - match attribute { - Attribute::Envelope => { - items.push(DataItem::Envelope { - envelope: message.envelope(), - }); - } - Attribute::Flags => { - let mut flags = message_cache - .expand_keywords(data) - .map(Flag::from) - .collect::>(); - if set_seen_flag { - flags.push(Flag::Seen); - } - items.push(DataItem::Flags { flags }); - } - Attribute::InternalDate => { - items.push(DataItem::InternalDate { - date: (metadata.rcvd_attach.to_native() & MESSAGE_RECEIVED_MASK) as i64, - }); - } - Attribute::Preview { .. } => { - items.push(DataItem::Preview { - contents: if !metadata.preview.is_empty() { - Some(metadata.preview.as_bytes().into()) - } else { - None - }, - }); - } - Attribute::Rfc822Size => { - items.push(DataItem::Rfc822Size { - size: data.size as usize, - }); - } - Attribute::Uid => { - items.push(DataItem::Uid { uid }); - } - Attribute::Rfc822 => { - items.push(DataItem::Rfc822 { - contents: raw_message.get_full_range(), - }); - } - Attribute::Rfc822Header => { - let contents = raw_message.get_slice_range( - 0..u32::from(metadata.root_part().offset_body) as usize, + if let Some(raw_body) = &raw_body { + raw_message.append( + raw_body + .get(metadata.blob_body_offset.to_native() as usize..) + .unwrap_or_default(), + ); + } else { + trc::event!( + Store(trc::StoreEvent::NotFound), + AccountId = account_id, + DocumentId = id, + Collection = Collection::Email, + BlobId = metadata.blob_hash.0.as_slice(), + Details = "Blob not found.", + CausedBy = trc::location!(), ); - if contents != SliceRange::None { - items.push(DataItem::Rfc822Header { contents }); - } - } - Attribute::Rfc822Text => { - items.push(DataItem::Rfc822Text { - contents: raw_message.get_full_range(), - }); - } - Attribute::Body => { - items.push(DataItem::Body { - part: metadata.body_structure(&decoded, false), - }); - } - Attribute::BodyStructure => { - items.push(DataItem::BodyStructure { - part: metadata.body_structure(&decoded, true), - }); - } - Attribute::BodySection { - sections, partial, .. - } => { - if let Some(contents) = metadata.body_section(&decoded, sections, *partial) - { - items.push(DataItem::BodySection { - sections: sections.to_vec(), - origin_octet: partial.map(|(start, _)| start), - contents, - }); - } - } - - Attribute::Binary { - sections, partial, .. - } => match metadata.binary(&decoded, sections, *partial) { - Ok(Some(contents)) => { - items.push(DataItem::Binary { - sections: sections.to_vec(), - offset: partial.map(|(start, _)| start), - contents, - }); - } - Err(_) => { - self.write_error( - trc::ImapEvent::Error - .into_err() - .details(format!( - "Failed to decode part {} of message {}.", - sections - .iter() - .map(|s| s.to_string()) - .collect::>() - .join("."), - if is_uid { uid } else { seqnum } - )) - .code(ResponseCode::UnknownCte), - ) - .await?; - continue; - } - _ => (), - }, - Attribute::BinarySize { sections } => { - if let Some(size) = metadata.binary_size(&decoded, sections) { - items.push(DataItem::BinarySize { - sections: sections.to_vec(), - size, - }); - } - } - Attribute::ModSeq => { - items.push(DataItem::ModSeq { - modseq: data.change_id + 1, - }); - } - Attribute::ObjectId => { - items.push(DataItem::ObjectId(ObjectId { - email_id: Some(Id::from_parts(data.thread_id, id)), - thread_id: Some(Id::from(data.thread_id)), - ..Default::default() - })); + continue; } } - } - // Add flags to the response if the message was unseen - if set_seen_flag && !arguments.attributes.contains(&Attribute::Flags) { - let mut flags = message_cache - .expand_keywords(data) - .map(Flag::from) - .collect::>(); - flags.push(Flag::Seen); - items.push(DataItem::Flags { flags }); - } + let message = &metadata.contents[0]; + let decoded = metadata.decode_contents(raw_message.clone()); - // Serialize fetch item - let mut buf = Vec::with_capacity(128); - FetchItem { - id: if is_uidonly { uid } else { seqnum }, - is_uidonly, - items, - } - .serialize(&mut buf, is_utf8); - self.write_bytes(buf).await?; + // Build response + let mut items = Vec::with_capacity(arguments.attributes.len()); + let set_seen_flag = + set_seen_flags && !message_cache.has_keyword(data, &Keyword::Seen); - // Add to set flags - if set_seen_flag - && let Some(data_) = self - .server - .store() - .get_value::>(ValueKey::archive( - account_id, - Collection::Email, - id, - )) - .await - .imap_ctx(&arguments.tag, trc::location!())? - { - let data = data_ - .to_unarchived::() - .imap_ctx(&arguments.tag, trc::location!())?; - let mut new_data = data.inner.to_builder(); - new_data.keywords.push(Keyword::Seen); + for attribute in &arguments.attributes { + match attribute { + Attribute::Envelope => { + items.push(DataItem::Envelope { + envelope: message.envelope(), + }); + } + Attribute::Flags => { + let mut flags = message_cache + .expand_keywords(data) + .map(Flag::from) + .collect::>(); + if set_seen_flag { + flags.push(Flag::Seen); + } + items.push(DataItem::Flags { flags }); + } + Attribute::InternalDate => { + items.push(DataItem::InternalDate { + date: (metadata.rcvd_attach.to_native() & MESSAGE_RECEIVED_MASK) + as i64, + }); + } + Attribute::Preview { .. } => { + items.push(DataItem::Preview { + contents: if !metadata.preview.is_empty() { + Some(metadata.preview.as_bytes().into()) + } else { + None + }, + }); + } + Attribute::Rfc822Size => { + items.push(DataItem::Rfc822Size { + size: data.size as usize, + }); + } + Attribute::Uid => { + items.push(DataItem::Uid { uid }); + } + Attribute::Rfc822 => { + items.push(DataItem::Rfc822 { + contents: raw_message.get_full_range(), + }); + } + Attribute::Rfc822Header => { + let contents = raw_message.get_slice_range( + 0..u32::from(metadata.root_part().offset_body) as usize, + ); - batch - .with_account_id(account_id) - .with_collection(Collection::Email) - .with_document(id) - .custom( - ObjectIndexBuilder::new() - .with_current(data) - .with_changes(new_data.seal()), - ) - .imap_ctx(&arguments.tag, trc::location!())? - .commit_point(); + if contents != SliceRange::None { + items.push(DataItem::Rfc822Header { contents }); + } + } + Attribute::Rfc822Text => { + items.push(DataItem::Rfc822Text { + contents: raw_message.get_full_range(), + }); + } + Attribute::Body => { + items.push(DataItem::Body { + part: metadata.body_structure(&decoded, false), + }); + } + Attribute::BodyStructure => { + items.push(DataItem::BodyStructure { + part: metadata.body_structure(&decoded, true), + }); + } + Attribute::BodySection { + sections, partial, .. + } => { + if let Some(contents) = + metadata.body_section(&decoded, sections, *partial) + { + items.push(DataItem::BodySection { + sections: sections.to_vec(), + origin_octet: partial.map(|(start, _)| start), + contents, + }); + } + } + + Attribute::Binary { + sections, partial, .. + } => match metadata.binary(&decoded, sections, *partial) { + Ok(Some(contents)) => { + items.push(DataItem::Binary { + sections: sections.to_vec(), + offset: partial.map(|(start, _)| start), + contents, + }); + } + Err(_) => { + self.write_error( + trc::ImapEvent::Error + .into_err() + .details(format!( + "Failed to decode part {} of message {}.", + sections + .iter() + .map(|s| s.to_string()) + .collect::>() + .join("."), + if is_uid { uid } else { seqnum } + )) + .code(ResponseCode::UnknownCte), + ) + .await?; + continue; + } + _ => (), + }, + Attribute::BinarySize { sections } => { + if let Some(size) = metadata.binary_size(&decoded, sections) { + items.push(DataItem::BinarySize { + sections: sections.to_vec(), + size, + }); + } + } + Attribute::ModSeq => { + items.push(DataItem::ModSeq { + modseq: data.change_id + 1, + }); + } + Attribute::ObjectId => { + items.push(DataItem::ObjectId(ObjectId { + email_id: Some(Id::from_parts(data.thread_id, id)), + thread_id: Some(Id::from(data.thread_id)), + ..Default::default() + })); + } + } + } + + // Add flags to the response if the message was unseen + if set_seen_flag && !arguments.attributes.contains(&Attribute::Flags) { + let mut flags = message_cache + .expand_keywords(data) + .map(Flag::from) + .collect::>(); + flags.push(Flag::Seen); + items.push(DataItem::Flags { flags }); + } + + // Serialize fetch item + let mut buf = Vec::with_capacity(128); + FetchItem { + id: if is_uidonly { uid } else { seqnum }, + is_uidonly, + items, + } + .serialize(&mut buf, is_utf8); + self.write_bytes(buf).await?; + + // Add to set flags + if set_seen_flag && let Some(data_) = seen_archive { + let data = data_ + .to_unarchived::() + .imap_ctx(&arguments.tag, trc::location!())?; + let mut new_data = data.inner.to_builder(); + new_data.keywords.push(Keyword::Seen); + + batch + .with_account_id(account_id) + .with_collection(Collection::Email) + .with_document(id) + .custom( + ObjectIndexBuilder::new() + .with_current(data) + .with_changes(new_data.seal()), + ) + .imap_ctx(&arguments.tag, trc::location!())? + .commit_point(); + } } } diff --git a/crates/imap/src/op/store.rs b/crates/imap/src/op/store.rs index 8a33843e..a97e509a 100644 --- a/crates/imap/src/op/store.rs +++ b/crates/imap/src/op/store.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use super::{FromModSeq, ImapContext}; @@ -27,7 +29,7 @@ use imap_proto::{ use registry::schema::enums::Permission; use std::{sync::Arc, time::Instant}; use store::{ - ValueKey, + READ_BATCH_SIZE, ValueKey, query::log::{Change, Query}, write::{AlignedBytes, Archive, BatchBuilder}, }; @@ -252,158 +254,161 @@ impl SessionData { let mut changed_mailboxes = AHashSet::new(); let mut batch = BatchBuilder::new(); - for (id, imap_id) in &ids { - // Obtain message data - let data_ = if let Some(data) = self + let ids = ids.into_iter().collect::>(); + for chunk in ids.chunks(READ_BATCH_SIZE) { + // inbuxa: one round trip per batch of messages instead of one each + let keys = chunk + .iter() + .map(|(id, _)| ValueKey::archive(account_id, Collection::Email, *id)) + .collect::>(); + let archives = self .server .store() - .get_value::>(ValueKey::archive( - account_id, - Collection::Email, - *id, - )) + .get_values::>(&keys) .await - .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())? - { - data - } else { - continue; - }; - - // Deserialize - let data = data_ - .to_unarchived::() .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())?; - let mut new_data = data.inner.to_builder(); - // Apply changes - let mut seen_changed = false; - match arguments.operation { - Operation::Set => { - seen_changed = set_keywords.contains(&Keyword::Seen) - != new_data.has_keyword(&Keyword::Seen); - new_data.set_keywords(set_keywords.clone()); - } - Operation::Add => { - for keyword in &set_keywords { - if new_data.add_keyword(keyword.clone()) && keyword == &Keyword::Seen { - seen_changed = true; + for ((id, imap_id), data_) in chunk.iter().zip(archives) { + // Obtain message data + let Some(data_) = data_ else { + continue; + }; + + // Deserialize + let data = data_ + .to_unarchived::() + .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())?; + let mut new_data = data.inner.to_builder(); + + // Apply changes + let mut seen_changed = false; + match arguments.operation { + Operation::Set => { + seen_changed = set_keywords.contains(&Keyword::Seen) + != new_data.has_keyword(&Keyword::Seen); + new_data.set_keywords(set_keywords.clone()); + } + Operation::Add => { + for keyword in &set_keywords { + if new_data.add_keyword(keyword.clone()) && keyword == &Keyword::Seen { + seen_changed = true; + } + } + } + Operation::Clear => { + for keyword in &set_keywords { + if new_data.remove_keyword(keyword) && keyword == &Keyword::Seen { + seen_changed = true; + } } } } - Operation::Clear => { - for keyword in &set_keywords { - if new_data.remove_keyword(keyword) && keyword == &Keyword::Seen { - seen_changed = true; - } - } - } - } - if !new_data.has_keyword_changes(data.inner) { - continue; - } - - // Train spam filter - let mut train_spam = None; - for keyword in new_data.added_keywords(data.inner) { - if keyword == &Keyword::Junk { - train_spam = Some(true); - break; - } else if keyword == &Keyword::NotJunk && !data.inner.has_mailbox_id(TRASH_ID) { - // Only train as ham if not in Trash (Apple likes to add NotJunk to trashed items, which would be spammy) - train_spam = Some(false); - break; + if !new_data.has_keyword_changes(data.inner) { + continue; } - } - if train_spam.is_none() { - for keyword in new_data.removed_keywords(data.inner) { + + // Train spam filter + let mut train_spam = None; + for keyword in new_data.added_keywords(data.inner) { if keyword == &Keyword::Junk { - if !data.inner.has_mailbox_id(TRASH_ID) { - train_spam = Some(false); - } + train_spam = Some(true); + break; + } else if keyword == &Keyword::NotJunk && !data.inner.has_mailbox_id(TRASH_ID) { + // Only train as ham if not in Trash (Apple likes to add NotJunk to trashed items, which would be spammy) + train_spam = Some(false); break; } } - } - - // Convert keywords to flags - let flags = if !arguments.is_silent { - new_data - .keywords - .iter() - .cloned() - .map(Flag::from) - .collect::>() - } else { - vec![] - }; - - // Set all current mailboxes as changed if the Seen tag changed - if seen_changed { - for mailbox_id in new_data.mailboxes.iter() { - changed_mailboxes.insert(mailbox_id.mailbox_id); + if train_spam.is_none() { + for keyword in new_data.removed_keywords(data.inner) { + if keyword == &Keyword::Junk { + if !data.inner.has_mailbox_id(TRASH_ID) { + train_spam = Some(false); + } + break; + } + } } - } - // Write changes - batch - .with_account_id(account_id) - .with_collection(Collection::Email) - .with_document(*id) - .custom( - ObjectIndexBuilder::new() - .with_current(data) - .with_changes(new_data.seal()), - ) - .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())?; + // Convert keywords to flags + let flags = if !arguments.is_silent { + new_data + .keywords + .iter() + .cloned() + .map(Flag::from) + .collect::>() + } else { + vec![] + }; - // Add spam train task - if let Some(learn_spam) = train_spam { - self.server - .add_account_spam_sample( - &mut batch, - account_id, - *id, - learn_spam, - self.session_id, + // Set all current mailboxes as changed if the Seen tag changed + if seen_changed { + for mailbox_id in new_data.mailboxes.iter() { + changed_mailboxes.insert(mailbox_id.mailbox_id); + } + } + + // Write changes + batch + .with_account_id(account_id) + .with_collection(Collection::Email) + .with_document(*id) + .custom( + ObjectIndexBuilder::new() + .with_current(data) + .with_changes(new_data.seal()), ) - .await .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())?; - } - // Set commit point - batch.commit_point(); + // Add spam train task + if let Some(learn_spam) = train_spam { + self.server + .add_account_spam_sample( + &mut batch, + account_id, + *id, + learn_spam, + self.session_id, + ) + .await + .imap_ctx(response.tag.as_ref().unwrap(), trc::location!())?; + } - // Add item to response - if !arguments.is_silent { - let mut data_items = vec![DataItem::Flags { flags }]; - if is_uid { - data_items.push(DataItem::Uid { uid: imap_id.uid }); + // Set commit point + batch.commit_point(); + + // Add item to response + if !arguments.is_silent { + let mut data_items = vec![DataItem::Flags { flags }]; + if is_uid { + data_items.push(DataItem::Uid { uid: imap_id.uid }); + } + items.items.push(FetchItem { + id: if is_uidonly { + imap_id.uid + } else { + imap_id.seqnum + }, + is_uidonly, + items: data_items, + }); + } else if is_condstore { + items.items.push(FetchItem { + id: if is_uidonly { + imap_id.uid + } else { + imap_id.seqnum + }, + is_uidonly, + items: if is_uid { + vec![DataItem::Uid { uid: imap_id.uid }] + } else { + vec![] + }, + }); } - items.items.push(FetchItem { - id: if is_uidonly { - imap_id.uid - } else { - imap_id.seqnum - }, - is_uidonly, - items: data_items, - }); - } else if is_condstore { - items.items.push(FetchItem { - id: if is_uidonly { - imap_id.uid - } else { - imap_id.seqnum - }, - is_uidonly, - items: if is_uid { - vec![DataItem::Uid { uid: imap_id.uid }] - } else { - vec![] - }, - }); } } @@ -449,7 +454,7 @@ impl SessionData { MailboxId = mailbox.id.mailbox_id, DocumentId = ids .iter() - .map(|id| trc::Value::from(*id.0)) + .map(|id| trc::Value::from(id.0)) .collect::>(), Type = format!("{:?}", arguments.operation), Details = arguments diff --git a/crates/jmap/src/email/get.rs b/crates/jmap/src/email/get.rs index cf4804da..ea8934e3 100644 --- a/crates/jmap/src/email/get.rs +++ b/crates/jmap/src/email/get.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::changes::state::JmapCacheState; @@ -27,7 +29,7 @@ use jmap_tools::{Key, Map, Value}; use mail_parser::HeaderValue; use std::future::Future; use store::{ - ValueKey, + READ_BATCH_SIZE, ValueKey, write::{AlignedBytes, Archive}, }; use trc::{AddContext, StoreEvent}; @@ -150,298 +152,321 @@ impl EmailGet for Server { } } - for id in ids { - // Obtain the email object - if !message_ids.contains(id.document_id()) { - response.push_not_found(id); - continue; - } - let metadata_ = match self + for ids in ids.chunks(READ_BATCH_SIZE) { + // inbuxa: the metadata of a batch of messages in one round trip + let keys = ids + .iter() + .filter(|id| message_ids.contains(id.document_id())) + .map(|id| { + ValueKey::property( + account_id, + Collection::Email, + id.document_id(), + EmailField::Metadata, + ) + }) + .collect::>(); + let mut metadatas = self .store() - .get_value::>(ValueKey::property( - account_id, - Collection::Email, - id.document_id(), - EmailField::Metadata, - )) + .get_values::>(&keys) .await? - { - Some(metadata) => metadata, - None => { + .into_iter(); + + for &id in ids { + // Obtain the email object + if !message_ids.contains(id.document_id()) { response.push_not_found(id); continue; } - }; - let metadata = metadata_ - .unarchive::() - .caused_by(trc::location!())?; - - // Obtain message data - let data = match cache.email_by_id(&id.document_id()) { - Some(data) => data, - None => { - response.push_not_found(id); - continue; - } - }; - - // Retrieve raw message if needed - let blob_hash = BlobHash::from(&metadata.blob_hash); - let raw_body; - let mut raw_message = ChainedBytes::new(metadata.raw_headers.as_ref()); - if needs_body { - raw_body = self - .blob_store() - .get_blob(blob_hash.as_slice(), 0..usize::MAX) - .await?; - - if let Some(raw_body) = &raw_body { - raw_message.append( - raw_body - .get(metadata.blob_body_offset.to_native() as usize..) - .unwrap_or_default(), - ); - } else { - trc::event!( - Store(StoreEvent::NotFound), - AccountId = account_id, - DocumentId = id.document_id(), - Collection = Collection::Email, - BlobId = blob_hash.to_hex(), - Details = "Blob not found.", - CausedBy = trc::location!(), - ); - - response.push_not_found(id); - continue; - } - } - let blob_id = BlobId { - hash: blob_hash, - class: BlobClass::Linked { - account_id, - collection: Collection::Email.into(), - document_id: id.document_id(), - }, - section: None, - }; - - // Prepare response - let mut email: Map<'_, EmailProperty, EmailValue> = - Map::with_capacity(properties.len()); - let contents = &metadata.contents[0]; - let root_part = &contents.parts[0]; - let blob_body_offset = metadata.blob_body_offset.to_native() as isize - - root_part.offset_body.to_native() as isize; - for property in &properties { - match property { - EmailProperty::Id => { - email.insert_unchecked(EmailProperty::Id, Id::from(*id)); + let metadata_ = match metadatas.next().flatten() { + Some(metadata) => metadata, + None => { + response.push_not_found(id); + continue; } - EmailProperty::ThreadId => { - email.insert_unchecked(EmailProperty::ThreadId, Id::from(id.prefix_id())); - } - EmailProperty::BlobId => { - email.insert_unchecked(EmailProperty::BlobId, blob_id.clone()); - } - EmailProperty::MailboxIds => { - let mut obj = Map::with_capacity(data.mailboxes.len()); - for id in data.mailboxes.iter() { - debug_assert!(id.uid != 0); - obj.insert_unchecked( - EmailProperty::IdValue(Id::from(id.mailbox_id)), - true, - ); - } + }; + let metadata = metadata_ + .unarchive::() + .caused_by(trc::location!())?; - email.insert_unchecked(property.clone(), Value::Object(obj)); + // Obtain message data + let data = match cache.email_by_id(&id.document_id()) { + Some(data) => data, + None => { + response.push_not_found(id); + continue; } - EmailProperty::Keywords => { - let mut obj = Map::with_capacity(2); - for keyword in cache.expand_keywords(data) { - obj.insert_unchecked(EmailProperty::Keyword(keyword), true); - } - email.insert_unchecked(property.clone(), Value::Object(obj)); - } - EmailProperty::Size => { - email.insert_unchecked(EmailProperty::Size, data.size); - } - EmailProperty::ReceivedAt => { - email.insert_unchecked( - EmailProperty::ReceivedAt, - EmailValue::Date(UTCDate::from_timestamp( - (metadata.rcvd_attach.to_native() & MESSAGE_RECEIVED_MASK) as i64, - )), + }; + + // Retrieve raw message if needed + let blob_hash = BlobHash::from(&metadata.blob_hash); + let raw_body; + let mut raw_message = ChainedBytes::new(metadata.raw_headers.as_ref()); + if needs_body { + raw_body = self + .blob_store() + .get_blob(blob_hash.as_slice(), 0..usize::MAX) + .await?; + + if let Some(raw_body) = &raw_body { + raw_message.append( + raw_body + .get(metadata.blob_body_offset.to_native() as usize..) + .unwrap_or_default(), ); + } else { + trc::event!( + Store(StoreEvent::NotFound), + AccountId = account_id, + DocumentId = id.document_id(), + Collection = Collection::Email, + BlobId = blob_hash.to_hex(), + Details = "Blob not found.", + CausedBy = trc::location!(), + ); + + response.push_not_found(id); + continue; } - EmailProperty::Preview => { - if !metadata.preview.is_empty() { + } + let blob_id = BlobId { + hash: blob_hash, + class: BlobClass::Linked { + account_id, + collection: Collection::Email.into(), + document_id: id.document_id(), + }, + section: None, + }; + + // Prepare response + let mut email: Map<'_, EmailProperty, EmailValue> = + Map::with_capacity(properties.len()); + let contents = &metadata.contents[0]; + let root_part = &contents.parts[0]; + let blob_body_offset = metadata.blob_body_offset.to_native() as isize + - root_part.offset_body.to_native() as isize; + for property in &properties { + match property { + EmailProperty::Id => { + email.insert_unchecked(EmailProperty::Id, Id::from(*id)); + } + EmailProperty::ThreadId => { email.insert_unchecked( - EmailProperty::Preview, - metadata.preview.to_string(), + EmailProperty::ThreadId, + Id::from(id.prefix_id()), ); } - } - EmailProperty::HasAttachment => { - email.insert_unchecked( - EmailProperty::HasAttachment, - (metadata.rcvd_attach.to_native() & MESSAGE_HAS_ATTACHMENT) != 0, - ); - } - EmailProperty::Subject => { - email.insert_unchecked( - EmailProperty::Subject, - root_part - .header_value(&MetadataHeaderName::Subject) - .map(|value| HeaderValue::from(value).into_form(&HeaderForm::Text)) - .unwrap_or_default(), - ); - } - EmailProperty::SentAt => { - email.insert_unchecked( - EmailProperty::SentAt, - root_part - .header_value(&MetadataHeaderName::Date) - .map(|value| HeaderValue::from(value).into_form(&HeaderForm::Date)) - .unwrap_or_default(), - ); - } - EmailProperty::MessageId - | EmailProperty::InReplyTo - | EmailProperty::References => { - email.insert_unchecked( - property.clone(), - root_part - .header_value(&match property { - EmailProperty::MessageId => MetadataHeaderName::MessageId, - EmailProperty::InReplyTo => MetadataHeaderName::InReplyTo, - EmailProperty::References => MetadataHeaderName::References, - _ => unreachable!(), - }) - .map(|value| { - HeaderValue::from(value).into_form(&HeaderForm::MessageIds) - }) - .unwrap_or_default(), - ); - } - - EmailProperty::Sender - | EmailProperty::From - | EmailProperty::To - | EmailProperty::Cc - | EmailProperty::Bcc - | EmailProperty::ReplyTo => { - email.insert_unchecked( - property.clone(), - root_part - .header_value(&match property { - EmailProperty::Sender => MetadataHeaderName::Sender, - EmailProperty::From => MetadataHeaderName::From, - EmailProperty::To => MetadataHeaderName::To, - EmailProperty::Cc => MetadataHeaderName::Cc, - EmailProperty::Bcc => MetadataHeaderName::Bcc, - EmailProperty::ReplyTo => MetadataHeaderName::ReplyTo, - _ => unreachable!(), - }) - .map(|value| { - HeaderValue::from(value).into_form(&HeaderForm::Addresses) - }) - .unwrap_or_default(), - ); - } - EmailProperty::Header(_) => { - email.insert_unchecked( - property.clone(), - root_part.header_to_value(property, &raw_message), - ); - } - EmailProperty::Headers => { - email.insert_unchecked( - EmailProperty::Headers, - root_part.headers_to_value(&raw_message), - ); - } - EmailProperty::TextBody - | EmailProperty::HtmlBody - | EmailProperty::Attachments => { - let list = match property { - EmailProperty::TextBody => &contents.text_body, - EmailProperty::HtmlBody => &contents.html_body, - EmailProperty::Attachments => &contents.attachments, - _ => unreachable!(), + EmailProperty::BlobId => { + email.insert_unchecked(EmailProperty::BlobId, blob_id.clone()); } - .iter(); - email.insert_unchecked( - property.clone(), - list.map(|part_id| { + EmailProperty::MailboxIds => { + let mut obj = Map::with_capacity(data.mailboxes.len()); + for id in data.mailboxes.iter() { + debug_assert!(id.uid != 0); + obj.insert_unchecked( + EmailProperty::IdValue(Id::from(id.mailbox_id)), + true, + ); + } + + email.insert_unchecked(property.clone(), Value::Object(obj)); + } + EmailProperty::Keywords => { + let mut obj = Map::with_capacity(2); + for keyword in cache.expand_keywords(data) { + obj.insert_unchecked(EmailProperty::Keyword(keyword), true); + } + email.insert_unchecked(property.clone(), Value::Object(obj)); + } + EmailProperty::Size => { + email.insert_unchecked(EmailProperty::Size, data.size); + } + EmailProperty::ReceivedAt => { + email.insert_unchecked( + EmailProperty::ReceivedAt, + EmailValue::Date(UTCDate::from_timestamp( + (metadata.rcvd_attach.to_native() & MESSAGE_RECEIVED_MASK) + as i64, + )), + ); + } + EmailProperty::Preview => { + if !metadata.preview.is_empty() { + email.insert_unchecked( + EmailProperty::Preview, + metadata.preview.to_string(), + ); + } + } + EmailProperty::HasAttachment => { + email.insert_unchecked( + EmailProperty::HasAttachment, + (metadata.rcvd_attach.to_native() & MESSAGE_HAS_ATTACHMENT) != 0, + ); + } + EmailProperty::Subject => { + email.insert_unchecked( + EmailProperty::Subject, + root_part + .header_value(&MetadataHeaderName::Subject) + .map(|value| { + HeaderValue::from(value).into_form(&HeaderForm::Text) + }) + .unwrap_or_default(), + ); + } + EmailProperty::SentAt => { + email.insert_unchecked( + EmailProperty::SentAt, + root_part + .header_value(&MetadataHeaderName::Date) + .map(|value| { + HeaderValue::from(value).into_form(&HeaderForm::Date) + }) + .unwrap_or_default(), + ); + } + EmailProperty::MessageId + | EmailProperty::InReplyTo + | EmailProperty::References => { + email.insert_unchecked( + property.clone(), + root_part + .header_value(&match property { + EmailProperty::MessageId => MetadataHeaderName::MessageId, + EmailProperty::InReplyTo => MetadataHeaderName::InReplyTo, + EmailProperty::References => MetadataHeaderName::References, + _ => unreachable!(), + }) + .map(|value| { + HeaderValue::from(value).into_form(&HeaderForm::MessageIds) + }) + .unwrap_or_default(), + ); + } + + EmailProperty::Sender + | EmailProperty::From + | EmailProperty::To + | EmailProperty::Cc + | EmailProperty::Bcc + | EmailProperty::ReplyTo => { + email.insert_unchecked( + property.clone(), + root_part + .header_value(&match property { + EmailProperty::Sender => MetadataHeaderName::Sender, + EmailProperty::From => MetadataHeaderName::From, + EmailProperty::To => MetadataHeaderName::To, + EmailProperty::Cc => MetadataHeaderName::Cc, + EmailProperty::Bcc => MetadataHeaderName::Bcc, + EmailProperty::ReplyTo => MetadataHeaderName::ReplyTo, + _ => unreachable!(), + }) + .map(|value| { + HeaderValue::from(value).into_form(&HeaderForm::Addresses) + }) + .unwrap_or_default(), + ); + } + EmailProperty::Header(_) => { + email.insert_unchecked( + property.clone(), + root_part.header_to_value(property, &raw_message), + ); + } + EmailProperty::Headers => { + email.insert_unchecked( + EmailProperty::Headers, + root_part.headers_to_value(&raw_message), + ); + } + EmailProperty::TextBody + | EmailProperty::HtmlBody + | EmailProperty::Attachments => { + let list = match property { + EmailProperty::TextBody => &contents.text_body, + EmailProperty::HtmlBody => &contents.html_body, + EmailProperty::Attachments => &contents.attachments, + _ => unreachable!(), + } + .iter(); + email.insert_unchecked( + property.clone(), + list.map(|part_id| { + contents.to_body_part( + u16::from(part_id) as u32, + &body_properties, + &raw_message, + &blob_id, + blob_body_offset, + ) + }) + .collect::>(), + ); + } + EmailProperty::BodyStructure => { + email.insert_unchecked( + EmailProperty::BodyStructure, contents.to_body_part( - u16::from(part_id) as u32, + 0, &body_properties, &raw_message, &blob_id, blob_body_offset, - ) - }) - .collect::>(), - ); - } - EmailProperty::BodyStructure => { - email.insert_unchecked( - EmailProperty::BodyStructure, - contents.to_body_part( - 0, - &body_properties, - &raw_message, - &blob_id, - blob_body_offset, - ), - ); - } - EmailProperty::BodyValues => { - let mut body_values = Map::with_capacity(contents.parts.len()); - for (part_id, part) in contents.parts.iter().enumerate() { - if part.is_text_mime_type() - && (fetch_all_body_values - || (fetch_html_body_values - && contents.is_html_part(part_id as u16)) - || (fetch_text_body_values - && contents.is_text_part(part_id as u16))) - { - let contents = part.decode_contents(&raw_message); - - let (is_truncated, value) = match &part.body { - ArchivedMetadataPartType::Text => { - truncate_plain(contents.as_str(), max_body_value_bytes) - } - ArchivedMetadataPartType::Html => { - truncate_html(contents.as_str(), max_body_value_bytes) - } - _ => unreachable!(), - }; - - body_values.insert_unchecked( - Key::Owned(part_id.to_string()), - Map::with_capacity(3) - .with_key_value( - EmailProperty::IsEncodingProblem, - (part.flags & PART_ENCODING_PROBLEM) != 0, - ) - .with_key_value(EmailProperty::IsTruncated, is_truncated) - .with_key_value(EmailProperty::Value, value), - ); - } + ), + ); } - email.insert_unchecked(EmailProperty::BodyValues, body_values); - } + EmailProperty::BodyValues => { + let mut body_values = Map::with_capacity(contents.parts.len()); + for (part_id, part) in contents.parts.iter().enumerate() { + if part.is_text_mime_type() + && (fetch_all_body_values + || (fetch_html_body_values + && contents.is_html_part(part_id as u16)) + || (fetch_text_body_values + && contents.is_text_part(part_id as u16))) + { + let contents = part.decode_contents(&raw_message); - _ => { - return Err(trc::JmapEvent::InvalidArguments - .into_err() - .details(format!("Invalid property {property:?}"))); + let (is_truncated, value) = match &part.body { + ArchivedMetadataPartType::Text => { + truncate_plain(contents.as_str(), max_body_value_bytes) + } + ArchivedMetadataPartType::Html => { + truncate_html(contents.as_str(), max_body_value_bytes) + } + _ => unreachable!(), + }; + + body_values.insert_unchecked( + Key::Owned(part_id.to_string()), + Map::with_capacity(3) + .with_key_value( + EmailProperty::IsEncodingProblem, + (part.flags & PART_ENCODING_PROBLEM) != 0, + ) + .with_key_value( + EmailProperty::IsTruncated, + is_truncated, + ) + .with_key_value(EmailProperty::Value, value), + ); + } + } + email.insert_unchecked(EmailProperty::BodyValues, body_values); + } + + _ => { + return Err(trc::JmapEvent::InvalidArguments + .into_err() + .details(format!("Invalid property {property:?}"))); + } } } + response.list.push(email.into()); } - response.list.push(email.into()); } Ok(response) diff --git a/crates/store/src/backend/postgres/read.rs b/crates/store/src/backend/postgres/read.rs index 681c3d46..ca7a6cb3 100644 --- a/crates/store/src/backend/postgres/read.rs +++ b/crates/store/src/backend/postgres/read.rs @@ -11,6 +11,7 @@ use crate::{ Deserialize, IterateParams, Key, ValueKey, backend::postgres::into_pool_error, write::ValueClass, }; +use ahash::AHashMap; use futures::{TryStreamExt, pin_mut}; impl PostgresStore { @@ -44,6 +45,44 @@ impl PostgresStore { bounded(conn, result, limit) } + /// inbuxa: reads every key in one round trip (`k = ANY($1)`) and + /// returns the values in the keys' order, `None` for a missing key. + /// All keys must live in the same subspace; the dispatcher checks. + pub(crate) async fn get_values(&self, keys: &[impl Key]) -> trc::Result>> + where + U: Deserialize + 'static, + { + let Some(subspace) = keys.first().map(|key| key.subspace()) else { + return Ok(Vec::new()); + }; + let conn = self.conn_pool.get().await.map_err(into_pool_error)?; + let limit = self.timeouts.query; + let result = tokio::time::timeout(limit, async { + let s = conn + .prepare_cached(&format!( + "SELECT k, v FROM {} WHERE k = ANY($1)", + char::from(subspace) + )) + .await + .map_err(into_error)?; + let keys = keys.iter().map(|key| key.serialize(0)).collect::>(); + let mut found = AHashMap::with_capacity(keys.len()); + for row in conn.query(&s, &[&keys]).await.map_err(into_error)? { + found.insert(row.get::<_, Vec>(0), row.get::<_, Vec>(1)); + } + keys.iter() + .map(|key| { + found + .get(key) + .map(|value| U::deserialize_with_key(key, value)) + .transpose() + }) + .collect() + }) + .await; + bounded(conn, result, limit) + } + pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result { let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let limit = self.timeouts.query; diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index d65b2a98..4ddaae17 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -59,6 +59,81 @@ impl Store { .caused_by(trc::location!()) } + /// inbuxa: reads a set of keys in one round trip where the backend can + /// (PostgreSQL: `k = ANY($1)`), one `get_value` per key elsewhere. + /// Values come back in the keys' order, `None` for a missing key. + /// Callers bound a batch at [`crate::READ_BATCH_SIZE`] keys so a + /// FETCH of a whole mailbox never holds every archive at once. + pub async fn get_values(&self, keys: &[impl Key]) -> trc::Result>> + where + U: Deserialize + 'static, + { + async fn read(store: &Store, keys: &[impl Key]) -> trc::Result>> + where + U: Deserialize + 'static, + { + match store { + #[cfg(feature = "postgres")] + Store::PostgreSQL(db) => db.get_values(keys).await, + _ => { + let mut values = Vec::with_capacity(keys.len()); + for key in keys { + values.push(store.get_value(key.clone()).await?); + } + Ok(values) + } + } + } + + match self { + // ST-6 to ST-8, ST-12: one replica serves the batch when every + // key is in a replicated subspace; keys it has not received yet + // are read again from the primary + Self::Replicated(store) => { + let subspace = keys + .first() + .map(|key| key.subspace()) + .filter(|subspace| keys.iter().all(|key| key.subspace() == *subspace)); + match subspace { + Some(subspace) => match store.read_target(subspace).await { + Some(index) => match read(&store.replicas[index].store, keys).await { + Ok(mut values) => { + store.served(index); + let missing = values + .iter() + .enumerate() + .filter(|(_, value)| value.is_none()) + .map(|(pos, _)| pos) + .collect::>(); + if !missing.is_empty() { + let retry = missing + .iter() + .map(|pos| keys[*pos].clone()) + .collect::>(); + for (pos, value) in missing + .into_iter() + .zip(read::(&store.primary, &retry).await?) + { + values[pos] = value; + } + } + Ok(values) + } + Err(err) => { + store.failed(index, err); + read(&store.primary, keys).await + } + }, + None => read(&store.primary, keys).await, + }, + None => read(self, keys).await, + } + } + _ => read(self, keys).await, + } + .caused_by(trc::location!()) + } + pub async fn key_exists(&self, key: impl Key) -> trc::Result { match self { #[cfg(feature = "sqlite")] diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 4a88fc5f..382d847f 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -71,6 +71,10 @@ pub trait Key: Sync + Send + Clone { fn subspace(&self) -> u8; } +/// inbuxa: how many keys a caller hands `Store::get_values` at once. Bounds +/// the archives held in memory while a long FETCH or Email/get streams out. +pub const READ_BATCH_SIZE: usize = 128; + #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct IndexKey> { pub account_id: u32, diff --git a/tests/src/store/ops.rs b/tests/src/store/ops.rs index 01439f74..90967db1 100644 --- a/tests/src/store/ops.rs +++ b/tests/src/store/ops.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::utils::{cleanup::store_assert_is_empty, server::TestServer}; @@ -30,6 +32,8 @@ fn value_gen(chunks: impl IntoIterator) -> Vec { pub async fn test(test: &TestServer) { let db = test.server.store().clone(); + batched_reads(&db).await; + #[cfg(feature = "foundationdb")] if matches!(db, store::Store::FoundationDb(_)) { use store::write::RegistryClass; @@ -799,3 +803,62 @@ pub async fn test(test: &TestServer) { store_assert_is_empty(&db, db.clone().into(), false).await; } } + +// inbuxa: `Store::get_values` answers in the keys' order, `None` for a +// missing key, and agrees with `get_value` key by key +async fn batched_reads(db: &store::Store) { + println!("Running batched reads test..."); + const FIELD: u8 = 7; + let key = |document_id: u32| ValueKey::property(0, Collection::Email, document_id, FIELD); + let present = |document_id: u32| document_id % 7 != 3; + + let mut batch = BatchBuilder::new(); + batch.with_account_id(0).with_collection(Collection::Email); + for document_id in 0..300u32 { + if present(document_id) { + batch.with_document(document_id).set( + ValueClass::Property(FIELD), + (document_id as u64).to_be_bytes().to_vec(), + ); + } + } + db.write(batch.build_all()).await.unwrap(); + + // A batch wider than READ_BATCH_SIZE, with gaps, a repeat and keys that + // were never written + let mut ids = (0..300u32).collect::>(); + ids.push(42); + ids.push(5000); + ids.push(42); + let keys = ids.iter().map(|id| key(*id)).collect::>(); + let values = db.get_values::(&keys).await.unwrap(); + assert_eq!(values.len(), ids.len()); + for (id, value) in ids.iter().zip(&values) { + let expected = db.get_value::(key(*id)).await.unwrap(); + assert_eq!(value, &expected, "document {id}"); + assert_eq!(value.is_some(), *id < 300 && present(*id), "document {id}"); + } + + assert!( + db.get_values::(&[] as &[ValueKey]) + .await + .unwrap() + .is_empty() + ); + + let mut batch = BatchBuilder::new(); + batch.with_account_id(0).with_collection(Collection::Email); + for document_id in 0..300u32 { + batch + .with_document(document_id) + .clear(ValueClass::Property(FIELD)); + } + db.write(batch.build_all()).await.unwrap(); + assert!( + db.get_values::(&keys) + .await + .unwrap() + .iter() + .all(Option::is_none) + ); +} -- 2.54.0