starred/stalwart
1
0
Fork 0
mirror of https://github.com/stalwartlabs/stalwart.git synced 2026-10-29 08:12:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Database schema optimization - part 3

Browse Source
This commit is contained in:
mdecimus
2025-11-01 19:19:36 +01:00
parent d5c5f12adf
commit 8e1316895b
41 changed files with 1979 additions and 1276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/common/src/config/jmap/settings.rs
+41
View File
@@ -5,9 +5,11 @@
*/
use crate::config::groupware::GroupwareConfig;
use ahash::{AHashMap, AHashSet};
use jmap_proto::request::capability::BaseCapabilities;
use nlp::language::Language;
use std::{str::FromStr, time::Duration};
use store::{search::SearchField, write::SearchIndex};
use types::{collection::Collection, special_use::SpecialUse};
use utils::{
config::{Config, Rate, cron::SimpleCron, utils::ParseValue},
@@ -79,6 +81,10 @@ pub struct JmapConfig {
pub encrypt: bool,
pub encrypt_append: bool,
pub index_batch_size: usize,
pub index_all_headers: bool,
pub index_fields: AHashMap<SearchIndex, AHashSet<SearchField>>,
pub capabilities: BaseCapabilities,
pub account_purge_frequency: SimpleCron,
}
@@ -351,10 +357,45 @@ impl JmapConfig {
calendar_parse_max_items: config
.property("jmap.calendar.parse.max-items")
.unwrap_or(10),
index_batch_size: config.property("jmap.index.batch-size").unwrap_or(100),
index_all_headers: config
.property_or_default("jmap.index.email.all-headers", "false")
.unwrap_or(false),
index_fields: AHashMap::new(),
default_folders,
shared_folder,
};
// Parse index fields
for index in [
SearchIndex::Email,
SearchIndex::Contacts,
SearchIndex::Calendar,
] {
let mut fields = AHashSet::new();
let todo = "implement";
/*for field_str in config.values(&format!(
"jmap.index.{}.fields",
index.as_config_case()
)) {
match SearchField::try_from(field_str.1.as_str()) {
Ok(field) => {
fields.insert(field);
}
Err(_) => {
config.new_parse_error(
&format!(
"jmap.index.{}.fields",
index.as_config_case()
),
format!("Invalid search field: {}", field_str.1),
);
}
}
}*/
jmap.index_fields.insert(index, fields);
}
for collection in Bitmap::<Collection>::all() {
let key = format!("object-quota.{}", collection.as_config_case());
jmap.max_objects[collection as usize] =
crates/common/src/config/network.rs
+5
View File
@@ -44,6 +44,7 @@ pub struct ClusterRoles {
pub fts_indexing: ClusterRole,
pub bayes_training: ClusterRole,
pub imip_processing: ClusterRole,
pub merge_threads: ClusterRole,
pub calendar_alerts: ClusterRole,
pub renew_acme: ClusterRole,
pub calculate_metrics: ClusterRole,
@@ -256,6 +257,10 @@ impl Network {
&mut network.roles.calendar_alerts,
"cluster.roles.calendar-alerts",
),
(
&mut network.roles.merge_threads,
"cluster.roles.merge-threads",
),
] {
let shards = config
.properties::<NodeList>(key)
crates/common/src/telemetry/tracers/index.rs
+178
View File
@@ -0,0 +1,178 @@
fn build_index() {
let mut queue_ids = AHashSet::new();
let mut values = AHashSet::new();
for event in events.iter().chain([span, &event]) {
for (key, value) in &event.keys {
match (key, value) {
(Key::QueueId, Value::UInt(queue_id)) => {
queue_ids.insert(*queue_id);
}
(Key::From | Key::To | Key::Domain | Key::Hostname, Value::String(address)) => {
values.insert(address.clone());
}
(Key::To, Value::Array(value)) => {
for value in value {
if let Value::String(address) = value {
values.insert(address.clone());
}
}
}
(Key::RemoteIp, Value::Ipv4(ip)) => {
values.insert(ip.to_string().into());
}
(Key::RemoteIp, Value::Ipv6(ip)) => {
values.insert(ip.to_string().into());
}
_ => {}
}
}
}
// Build index
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: (span.inner.typ.code() as u16).to_be_bytes().to_vec(),
}),
vec![],
);
for queue_id in queue_ids {
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: queue_id.to_be_bytes().to_vec(),
}),
vec![],
);
}
for value in values {
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: value.as_bytes().to_vec(),
}),
vec![],
);
}
}
/*
enum SpanCollector {
Vec(Vec<u64>),
HashSet(AHashSet<u64>),
Empty,
}
impl SpanCollector {
fn new(num_params: usize) -> Self {
if num_params == 1 {
Self::Vec(Vec::new())
} else {
Self::HashSet(AHashSet::new())
}
}
fn insert(&mut self, span_id: u64) {
match self {
Self::Vec(vec) => vec.push(span_id),
Self::HashSet(set) => {
set.insert(span_id);
}
_ => unreachable!(),
}
}
fn into_vec(self) -> Vec<u64> {
match self {
Self::Vec(mut vec) => {
vec.sort_unstable_by(|a, b| b.cmp(a));
vec
}
Self::HashSet(set) => {
let mut vec: Vec<u64> = set.into_iter().collect();
vec.sort_unstable_by(|a, b| b.cmp(a));
vec
}
Self::Empty => Vec::new(),
}
}
fn intersect(&mut self, other_span: Self) -> bool {
match (self, other_span) {
(Self::HashSet(set), Self::HashSet(other_set)) => {
set.retain(|span_id| other_set.contains(span_id));
set.is_empty()
}
_ => unreachable!(),
}
}
}
let mut spans = SpanCollector::Empty;
let num_params = params.len();
let todo = "use FTS";
for (param_num, param) in params.iter().enumerate() {
let (value, exact_len) = match param {
TracingQuery::EventType(event) => (
(event.code() as u16).to_be_bytes().to_vec(),
std::mem::size_of::<u16>() + U64_LEN,
),
TracingQuery::QueueId(id) => (
id.to_be_bytes().to_vec(),
std::mem::size_of::<u64>() + U64_LEN,
),
TracingQuery::Keywords(value) => {
if let Some(value) = value.strip_prefix('"').and_then(|v| v.strip_suffix('"')) {
(value.as_bytes().to_vec(), value.len() + U64_LEN)
} else {
(value.as_bytes().to_vec(), 0)
}
}
};
let mut param_spans = SpanCollector::new(num_params);
self.iterate(
IterateParams::new(
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: 0,
value: value.clone(),
})),
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: u64::MAX,
value,
})),
)
.no_values(),
|key, _| {
if exact_len == 0 || key.len() == exact_len {
let span_id = key
.deserialize_be_u64(key.len() - U64_LEN)
.caused_by(trc::location!())?;
if (from_span_id == 0 || span_id >= from_span_id)
&& (to_span_id == 0 || span_id <= to_span_id)
{
param_spans.insert(span_id);
}
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
if param_num == 0 {
spans = param_spans;
} else if spans.intersect(param_spans) {
return Ok(Vec::new());
}
}
Ok(spans.into_vec())
*/
crates/common/src/telemetry/tracers/store.rs
+26 -230
View File
@@ -8,12 +8,12 @@
*
*/
use std::{future::Future, time::Duration};
use crate::config::telemetry::StoreTracer;
use ahash::{AHashMap, AHashSet};
use std::{future::Future, time::Duration};
use store::{
Deserialize, IterateParams, Store, U64_LEN, ValueKey,
write::{BatchBuilder, TelemetryClass, ValueClass, key::DeserializeBigEndian},
Deserialize, Store, ValueKey,
write::{BatchBuilder, SearchIndex, TaskQueueClass, TelemetryClass, ValueClass, now},
};
use trc::{
AddContext, AuthEvent, Event, EventDetails, EventType, Key, MessageIngestEvent,
@@ -23,8 +23,6 @@ use trc::{
};
use utils::snowflake::SnowflakeIdGenerator;
use crate::config::telemetry::StoreTracer;
const MAX_EVENTS: usize = 2048;
pub(crate) fn spawn_store_tracer(builder: SubscriberBuilder, settings: StoreTracer) {
@@ -35,6 +33,8 @@ pub(crate) fn spawn_store_tracer(builder: SubscriberBuilder, settings: StoreTrac
let mut batch = BatchBuilder::new();
while let Some(events) = rx.recv().await {
let now = now();
for event in events {
if let Some(span) = &event.inner.span {
let span_id = span.span_id().unwrap();
@@ -43,44 +43,16 @@ pub(crate) fn spawn_store_tracer(builder: SubscriberBuilder, settings: StoreTrac
if events.len() < MAX_EVENTS {
events.push(event);
}
} else if let Some(events) = active_spans.remove(&span_id) {
let mut queue_ids = AHashSet::new();
let mut values = AHashSet::new();
for event in events.iter().chain([span, &event]) {
for (key, value) in &event.keys {
match (key, value) {
(Key::QueueId, Value::UInt(queue_id)) => {
queue_ids.insert(*queue_id);
}
(
Key::From | Key::To | Key::Domain | Key::Hostname,
Value::String(address),
) => {
values.insert(address.clone());
}
(Key::To, Value::Array(value)) => {
for value in value {
if let Value::String(address) = value {
values.insert(address.clone());
}
}
}
(Key::RemoteIp, Value::Ipv4(ip)) => {
values.insert(ip.to_string().into());
}
(Key::RemoteIp, Value::Ipv6(ip)) => {
values.insert(ip.to_string().into());
}
_ => {}
}
}
}
if !queue_ids.is_empty() {
// Serialize events
batch.set(
} else if let Some(events) = active_spans.remove(&span_id)
&& events
.iter()
.chain([span, &event])
.flat_map(|event| event.keys.iter())
.any(|(k, v)| matches!((k, v), (Key::QueueId, Value::UInt(_))))
{
// Serialize events
batch
.set(
ValueClass::Telemetry(TelemetryClass::Span { span_id }),
serialize_events(
[span.as_ref()]
@@ -89,35 +61,17 @@ pub(crate) fn spawn_store_tracer(builder: SubscriberBuilder, settings: StoreTrac
.chain([event.as_ref()].into_iter()),
events.len() + 2,
),
);
// Build index
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: (span.inner.typ.code() as u16).to_be_bytes().to_vec(),
)
.with_account_id((span_id >> 32) as u32) // TODO: This is hacky, improve
.with_document(span_id as u32)
.set(
ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
due: now,
index: SearchIndex::TracingSpan,
is_insert: true,
}),
vec![],
);
for queue_id in queue_ids {
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: queue_id.to_be_bytes().to_vec(),
}),
vec![],
);
}
for value in values {
batch.set(
ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: value.as_bytes().to_vec(),
}),
vec![],
);
}
}
}
}
}
@@ -181,69 +135,7 @@ impl TracingStore for Store {
from_span_id: u64,
to_span_id: u64,
) -> trc::Result<Vec<u64>> {
let mut spans = SpanCollector::Empty;
let num_params = params.len();
let todo = "use FTS";
for (param_num, param) in params.iter().enumerate() {
let (value, exact_len) = match param {
TracingQuery::EventType(event) => (
(event.code() as u16).to_be_bytes().to_vec(),
std::mem::size_of::<u16>() + U64_LEN,
),
TracingQuery::QueueId(id) => (
id.to_be_bytes().to_vec(),
std::mem::size_of::<u64>() + U64_LEN,
),
TracingQuery::Keywords(value) => {
if let Some(value) = value.strip_prefix('"').and_then(|v| v.strip_suffix('"')) {
(value.as_bytes().to_vec(), value.len() + U64_LEN)
} else {
(value.as_bytes().to_vec(), 0)
}
}
};
let mut param_spans = SpanCollector::new(num_params);
self.iterate(
IterateParams::new(
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: 0,
value: value.clone(),
})),
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: u64::MAX,
value,
})),
)
.no_values(),
|key, _| {
if exact_len == 0 || key.len() == exact_len {
let span_id = key
.deserialize_be_u64(key.len() - U64_LEN)
.caused_by(trc::location!())?;
if (from_span_id == 0 || span_id >= from_span_id)
&& (to_span_id == 0 || span_id <= to_span_id)
{
param_spans.insert(span_id);
}
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
if param_num == 0 {
spans = param_spans;
} else if spans.intersect(param_spans) {
return Ok(Vec::new());
}
}
Ok(spans.into_vec())
todo!()
}
async fn purge_spans(&self, period: Duration) -> trc::Result<()> {
@@ -262,108 +154,12 @@ impl TracingStore for Store {
.await
.caused_by(trc::location!())?;
let mut delete_keys: Vec<ValueClass> = Vec::new();
self.iterate(
IterateParams::new(
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: 0,
value: vec![],
})),
ValueKey::from(ValueClass::Telemetry(TelemetryClass::Index {
span_id: u64::MAX,
value: vec![u8::MAX; 16],
})),
)
.no_values(),
|key, _| {
let span_id = key
.deserialize_be_u64(key.len() - U64_LEN)
.caused_by(trc::location!())?;
if span_id < until_span_id {
delete_keys.push(ValueClass::Telemetry(TelemetryClass::Index {
span_id,
value: key[0..key.len() - U64_LEN].to_vec(),
}));
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
if !delete_keys.is_empty() {
// Commit index
let mut batch = BatchBuilder::new();
for key in delete_keys {
if batch.is_large_batch() {
self.write(batch.build_all()).await?;
batch = BatchBuilder::new();
}
batch.clear(key);
}
if !batch.is_empty() {
self.write(batch.build_all()).await?;
}
}
let todo = "delete from index";
Ok(())
}
}
enum SpanCollector {
Vec(Vec<u64>),
HashSet(AHashSet<u64>),
Empty,
}
impl SpanCollector {
fn new(num_params: usize) -> Self {
if num_params == 1 {
Self::Vec(Vec::new())
} else {
Self::HashSet(AHashSet::new())
}
}
fn insert(&mut self, span_id: u64) {
match self {
Self::Vec(vec) => vec.push(span_id),
Self::HashSet(set) => {
set.insert(span_id);
}
_ => unreachable!(),
}
}
fn into_vec(self) -> Vec<u64> {
match self {
Self::Vec(mut vec) => {
vec.sort_unstable_by(|a, b| b.cmp(a));
vec
}
Self::HashSet(set) => {
let mut vec: Vec<u64> = set.into_iter().collect();
vec.sort_unstable_by(|a, b| b.cmp(a));
vec
}
Self::Empty => Vec::new(),
}
}
fn intersect(&mut self, other_span: Self) -> bool {
match (self, other_span) {
(Self::HashSet(set), Self::HashSet(other_set)) => {
set.retain(|span_id| other_set.contains(span_id));
set.is_empty()
}
_ => unreachable!(),
}
}
}
impl StoreTracer {
pub fn default_events() -> impl IntoIterator<Item = EventType> {
EventType::variants().into_iter().filter(|event| {
crates/email/src/message/copy.rs
+13 -4
View File
@@ -12,7 +12,7 @@ use crate::{
mailbox::UidMailbox,
message::{
index::extractors::VisitText,
ingest::{MergeThreadTask, ThreadInfo},
ingest::{MergeThreadIds, ThreadInfo},
},
};
use common::{Server, auth::ResourceToken, storage::index::ObjectIndexBuilder};
@@ -179,7 +179,7 @@ impl EmailCopy for Server {
.log_container_insert(SyncCollection::Thread);
document_id
};
let due = now();
batch
.with_collection(Collection::Email)
.with_document(document_id)
@@ -201,11 +201,20 @@ impl EmailCopy for Server {
.set(
ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
index: SearchIndex::Email,
due: now(),
due,
is_insert: true,
}),
MergeThreadTask::new(thread_result).serialize(),
vec![],
);
// Merge threads if necessary
if let Some(merge_threads) = MergeThreadIds::new(thread_result).serialize() {
batch.set(
ValueClass::TaskQueue(TaskQueueClass::MergeThreads { due }),
merge_threads,
);
}
metadata
.index(
&mut batch,
crates/email/src/message/index/search.rs
+3 -6
View File
@@ -29,11 +29,8 @@ impl ArchivedMessageMetadata {
let message_contents = &self.contents[0];
let mut document = IndexDocument::with_default_language(language);
document.index_number(
EmailSearchField::ReceivedAt,
self.received_at.to_native() as i64,
);
document.index_number(EmailSearchField::Size, self.size.to_native());
document.index_unsigned(EmailSearchField::ReceivedAt, self.received_at.to_native());
document.index_unsigned(EmailSearchField::Size, self.size.to_native());
for (part_id, part) in message_contents
.parts
@@ -95,7 +92,7 @@ impl ArchivedMessageMetadata {
}
ArchivedHeaderName::Date => {
if let Some(date) = header.value.as_datetime() {
document.index_number(
document.index_integer(
EmailSearchField::SentAt,
DateTime::from(date).to_timestamp(),
);
crates/email/src/message/ingest.rs
+136 -32
View File
@@ -27,11 +27,11 @@ use mail_parser::{
use spam_filter::{
SpamFilterInput, analysis::init::SpamFilterInit, modules::bayes::BayesClassifier,
};
use std::future::Future;
use std::{borrow::Cow, fmt::Write, time::Instant};
use std::{borrow::Cow, cmp::Ordering, fmt::Write, time::Instant};
use std::{future::Future, hash::Hasher};
use store::{
IndexKeyPrefix, IterateParams, U32_LEN, ValueKey,
ahash::AHashMap,
ahash::{AHashMap, AHashSet},
write::{
BatchBuilder, IndexPropertyClass, SearchIndex, TaskQueueClass, ValueClass,
key::DeserializeBigEndian, now,
@@ -655,9 +655,17 @@ impl EmailIngest for Server {
due,
is_insert: true,
}),
MergeThreadTask::new(thread_result).serialize(),
vec![],
);
// Merge threads if necessary
if let Some(merge_threads) = MergeThreadIds::new(thread_result).serialize() {
batch.set(
ValueClass::TaskQueue(TaskQueueClass::MergeThreads { due }),
merge_threads,
);
}
// Request spam training
if let Some(learn_spam) = train_spam {
batch.set(
@@ -748,7 +756,7 @@ impl EmailIngest for Server {
// Find thread ids
let key_len = IndexKeyPrefix::len() + result.thread_hash.len() + U32_LEN;
let document_id_pos = key_len - U32_LEN;
let mut thread_ids = AHashMap::<u32, Vec<u32>>::with_capacity(16);
let mut thread_merge = ThreadMerge::new();
self.store()
.iterate(
IterateParams::new(
@@ -788,7 +796,7 @@ impl EmailIngest for Server {
result.duplicate_ids.push(document_id);
}
thread_ids.entry(thread_id).or_default().push(document_id);
thread_merge.add(thread_id, document_id);
return Ok(true);
}
@@ -803,25 +811,18 @@ impl EmailIngest for Server {
.await
.caused_by(trc::location!())?;
match thread_ids.len() {
match thread_merge.num_thread_ids() {
0 => Ok(result),
1 => {
// Happy path, only one thread id
result.thread_id = thread_ids.into_keys().next();
result.thread_id = thread_merge.thread_ids().next().copied();
Ok(result)
}
_ => {
// Multiple thread ids that this message belongs to, merge them
let mut max_thread_id = u32::MAX;
let mut max_count = 0;
for (thread_id, ids) in thread_ids {
if ids.len() > max_count {
max_count = ids.len();
max_thread_id = thread_id;
}
result.merge_ids.extend(ids);
}
result.thread_id = Some(max_thread_id);
let thread_merge = thread_merge.merge();
result.merge_ids = thread_merge.merge_ids;
result.thread_id = Some(thread_merge.thread_id);
Ok(result)
}
}
@@ -856,47 +857,53 @@ impl IngestSource<'_> {
}
}
pub struct MergeThreadTask {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MergeThreadIds<T> {
pub thread_hash: CheekyHash,
pub duplicate_ids: Vec<u32>,
pub merge_ids: T,
}
impl MergeThreadTask {
impl MergeThreadIds<Vec<u32>> {
pub(crate) fn new(thread_result: ThreadResult) -> Self {
Self {
thread_hash: thread_result.thread_hash,
duplicate_ids: thread_result.duplicate_ids,
merge_ids: thread_result.merge_ids,
}
}
pub(crate) fn serialize(&self) -> Vec<u8> {
if !self.duplicate_ids.is_empty() {
pub(crate) fn serialize(&self) -> Option<Vec<u8>> {
if !self.merge_ids.is_empty() {
let mut buf =
Vec::with_capacity(self.thread_hash.len() + self.duplicate_ids.len() * U32_LEN);
Vec::with_capacity(self.thread_hash.len() + self.merge_ids.len() * U32_LEN);
buf.extend_from_slice(self.thread_hash.as_bytes());
for id in &self.duplicate_ids {
for id in &self.merge_ids {
buf.extend_from_slice(&id.to_be_bytes());
}
buf
Some(buf)
} else {
vec![]
None
}
}
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
impl MergeThreadIds<AHashSet<u32>> {
pub fn deserialize(document_id: u32, bytes: &[u8]) -> Option<Self> {
if !bytes.is_empty() {
let thread_hash = CheekyHash::deserialize(bytes)?;
let mut duplicate_ids = Vec::new();
let mut merge_ids =
AHashSet::with_capacity(((bytes.len() - thread_hash.len()) / U32_LEN) + 1);
let mut start_offset = thread_hash.len();
merge_ids.insert(document_id);
while let Some(id_bytes) = bytes.get(start_offset..start_offset + U32_LEN) {
duplicate_ids.push(u32::from_be_bytes(id_bytes.try_into().ok()?));
merge_ids.insert(u32::from_be_bytes(id_bytes.try_into().ok()?));
start_offset += U32_LEN;
}
Some(Self {
thread_hash,
duplicate_ids,
merge_ids,
})
} else {
None
@@ -904,6 +911,13 @@ impl MergeThreadTask {
}
}
impl std::hash::Hash for MergeThreadIds<AHashSet<u32>> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.thread_hash.hash(state);
self.merge_ids.len().hash(state);
}
}
pub(crate) struct ThreadInfo;
impl ThreadInfo {
@@ -925,3 +939,93 @@ impl ThreadInfo {
buf
}
}
pub struct ThreadMerge {
entries: AHashMap<u32, Vec<u32>>,
num_ids: usize,
}
pub struct ThreadMergeResult {
pub thread_id: u32,
pub merge_ids: Vec<u32>,
}
impl ThreadMerge {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
Self {
entries: AHashMap::with_capacity(8),
num_ids: 0,
}
}
pub fn add(&mut self, thread_id: u32, document_id: u32) {
self.entries.entry(thread_id).or_default().push(document_id);
self.num_ids += 1;
}
pub fn num_document_ids(&self) -> usize {
self.num_ids
}
pub fn num_thread_ids(&self) -> usize {
self.entries.len()
}
pub fn thread_ids(&self) -> impl Iterator<Item = &u32> {
self.entries.keys()
}
pub fn thread_groups(&self) -> impl Iterator<Item = (&u32, &Vec<u32>)> {
self.entries.iter()
}
pub fn merge_thread_id(&self) -> u32 {
let mut max_thread_id = u32::MAX;
let mut max_count = 0;
for (thread_id, ids) in &self.entries {
match ids.len().cmp(&max_count) {
Ordering::Greater => {
max_count = ids.len();
max_thread_id = *thread_id;
}
Ordering::Equal => {
if *thread_id < max_thread_id {
max_thread_id = *thread_id;
}
}
Ordering::Less => (),
}
}
max_thread_id
}
pub fn merge(self) -> ThreadMergeResult {
let mut max_thread_id = u32::MAX;
let mut max_count = 0;
let mut merge_ids = Vec::with_capacity(self.num_ids);
for (thread_id, ids) in self.entries {
match ids.len().cmp(&max_count) {
Ordering::Greater => {
max_count = ids.len();
max_thread_id = thread_id;
}
Ordering::Equal => {
if thread_id < max_thread_id {
max_thread_id = thread_id;
}
}
Ordering::Less => (),
}
merge_ids.extend(ids);
}
ThreadMergeResult {
thread_id: max_thread_id,
merge_ids,
}
}
}
crates/groupware/src/cache/calcard.rs
Vendored
-2
View File
@@ -26,7 +26,6 @@ use trc::AddContext;
use types::{
acl::AclGrant,
collection::{Collection, SyncCollection},
field::CalendarNotificationField,
};
use utils::map::bitmap::Bitmap;
@@ -128,7 +127,6 @@ pub(super) async fn build_calcard_resources(
}
}
let todo = "fix fdb range scan to support chunked reads";
let parent_range = cache.resources.len();
server
.archives(account_id, item_collection, &(), |document_id, archive| {
crates/groupware/src/calendar/index.rs
+1 -1
View File
@@ -339,7 +339,7 @@ impl ArchivedCalendarEvent {
pub fn index_document(&self) -> IndexDocument {
let mut document = IndexDocument::with_default_language(Language::Unknown);
document.index_number(CalendarSearchField::Start, self.data.event_range_start());
document.index_integer(CalendarSearchField::Start, self.data.event_range_start());
for component in self
.data
crates/groupware/src/contact/index.rs
+1 -1
View File
@@ -245,7 +245,7 @@ impl ArchivedContactCard {
pub fn index_document(&self) -> IndexDocument {
let mut document = IndexDocument::with_default_language(Language::Unknown);
document.index_number(ContactSearchField::Created, self.created.to_native());
document.index_integer(ContactSearchField::Created, self.created.to_native());
for entry in self.card.entries.iter() {
let field = match entry.name {
crates/http/src/management/principal.rs
+33 -6
View File
@@ -20,6 +20,7 @@ use hyper::{Method, header};
use serde_json::json;
use std::future::Future;
use std::sync::Arc;
use store::{search::SearchQuery, write::SearchIndex};
use trc::AddContext;
use utils::url_params::UrlParams;
@@ -399,11 +400,23 @@ impl PrincipalManager for Server {
}
if matches!(typ, Type::Individual | Type::Group) {
// Remove FTS index
if let Err(err) =
server.core.storage.fts.remove_all(principal.id()).await
{
trc::error!(err.details("Failed to delete FTS index"));
// Remove search index
for index in [
SearchIndex::Email,
SearchIndex::Contacts,
SearchIndex::Calendar,
] {
if let Err(err) = server
.core
.storage
.fts
.unindex(
SearchQuery::new(index).with_account_id(principal.id()),
)
.await
{
trc::error!(err.details("Failed to delete FTS index"));
}
}
// Delete bayes model
@@ -516,7 +529,21 @@ impl PrincipalManager for Server {
if matches!(typ, Type::Individual | Type::Group) {
// Remove FTS index
self.core.storage.fts.remove_all(account_id).await?;
for index in [
SearchIndex::Email,
SearchIndex::Contacts,
SearchIndex::Calendar,
] {
if let Err(err) = self
.core
.storage
.fts
.unindex(SearchQuery::new(index).with_account_id(account_id))
.await
{
trc::error!(err.details("Failed to delete FTS index"));
}
}
// Delete bayes model
if self
crates/http/src/management/stores.rs
+7 -4
View File
@@ -23,11 +23,11 @@ use email::{
use http_proto::{request::decode_path_element, *};
use hyper::Method;
use serde_json::json;
use services::task_manager::fts::FtsIndexTask;
use services::task_manager::index::ReindexIndexTask;
use std::future::Future;
use store::{
Serialize, rand,
write::{Archiver, BatchBuilder, ValueClass},
write::{Archiver, BatchBuilder, SearchIndex, ValueClass},
};
use trc::AddContext;
use types::{
@@ -221,7 +221,7 @@ impl ManageStore for Server {
}))
.await
}
(Some("reindex"), id, None, &Method::GET) => {
(Some("reindex"), Some(index), id, &Method::GET) => {
// Validate the access token
access_token.assert_has_permission(Permission::FtsReindex)?;
@@ -237,10 +237,13 @@ impl ManageStore for Server {
None
};
let tenant_id = access_token.tenant.map(|t| t.id);
let index = SearchIndex::try_from_str(index).ok_or_else(|| {
trc::ResourceEvent::BadParameters.reason("Invalid search index specified")
})?;
let jmap = self.clone();
tokio::spawn(async move {
if let Err(err) = jmap.fts_reindex(account_id, tenant_id).await {
if let Err(err) = jmap.reindex(index, account_id, tenant_id).await {
trc::error!(err.details("Failed to reindex FTS"));
}
});
crates/imap/src/op/search.rs
+8 -13
View File
@@ -26,16 +26,12 @@ use std::{borrow::Cow, str::FromStr, sync::Arc, time::Instant};
use store::{
query::log::Query,
roaring::RoaringBitmap,
search::{EmailSearchField, SearchComparator, SearchFilter},
write::now,
search::{EmailSearchField, SearchComparator, SearchFilter, SearchQuery},
write::{SearchIndex, now},
};
use tokio::sync::watch;
use trc::AddContext;
use types::{
collection::{Collection, SyncCollection},
id::Id,
keyword::Keyword,
};
use types::{collection::SyncCollection, id::Id, keyword::Keyword};
impl<T: SessionStream> Session<T> {
pub async fn handle_search(
@@ -226,8 +222,6 @@ impl<T: SessionStream> SessionData<T> {
.map(|m| m.document_id),
);
filters.push(SearchFilter::is_in_set(message_ids.clone()));
// Convert query
let mut include_highest_modseq = false;
for filter in imap_filter {
@@ -594,10 +588,11 @@ impl<T: SessionStream> SessionData<T> {
self.server
.search_store()
.query(
mailbox.id.account_id,
Collection::Email,
filters,
comparators,
SearchQuery::new(SearchIndex::Email)
.with_filters(filters)
.with_comparators(comparators)
.with_account_id(mailbox.id.account_id)
.with_mask(message_ids),
)
.await
.map(|res| (res, include_highest_modseq))
crates/jmap/src/calendar_event/query.rs
+14 -17
View File
@@ -18,7 +18,8 @@ use nlp::language::Language;
use std::{cmp::Ordering, sync::Arc};
use store::{
roaring::RoaringBitmap,
search::{CalendarSearchField, SearchComparator, SearchFilter},
search::{CalendarSearchField, SearchComparator, SearchFilter, SearchQuery},
write::SearchIndex,
};
use trc::AddContext;
use types::{
@@ -208,9 +209,20 @@ impl CalendarEventQuery for Server {
} else {
vec![]
};
let results = self
.search_store()
.query(account_id, Collection::CalendarEvent, filters, comparators)
.query(
SearchQuery::new(SearchIndex::Calendar)
.with_filters(filters)
.with_comparators(comparators)
.with_account_id(account_id)
.with_mask(if access_token.is_shared(account_id) {
cache.shared_items(access_token, [Acl::ReadItems], true)
} else {
cache.document_ids(false).collect()
}),
)
.await?;
let mut response = QueryResponseBuilder::new(
@@ -227,8 +239,6 @@ impl CalendarEventQuery for Server {
.as_deref()
.filter(|s| !s.is_empty())
.unwrap_or_default();
let filter_mask = (access_token.is_shared(account_id))
.then(|| cache.shared_items(access_token, [Acl::ReadItems], true));
if expand_recurrences {
let Some(time_range) = filter.filter(|f| f.start != i64::MIN && f.end != i64::MAX)
@@ -244,13 +254,6 @@ impl CalendarEventQuery for Server {
.any(|c| matches!(c.property, CalendarEventComparator::Uid));
for document_id in results {
if filter_mask
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
}
let Some(_calendar_event) = self
.archive(account_id, Collection::CalendarEvent, document_id)
.await?
@@ -327,12 +330,6 @@ impl CalendarEventQuery for Server {
}
} else {
for document_id in results {
if filter_mask
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
}
if !response.add(0, document_id) {
break;
}
crates/jmap/src/calendar_event_notification/query.rs
+11 -13
View File
@@ -19,8 +19,8 @@ use store::{
IterateParams, U32_LEN, U64_LEN, ValueKey,
ahash::AHashSet,
roaring::RoaringBitmap,
search::SearchFilter,
write::{IndexPropertyClass, ValueClass, key::DeserializeBigEndian},
search::{SearchFilter, SearchQuery},
write::{IndexPropertyClass, SearchIndex, ValueClass, key::DeserializeBigEndian},
};
use trc::AddContext;
use types::{
@@ -58,6 +58,7 @@ impl CalendarEventNotificationQuery for Server {
)
.await?;
let mut notifications = Vec::with_capacity(16);
let mut document_ids = RoaringBitmap::new();
self.store()
.iterate(
@@ -83,11 +84,13 @@ impl CalendarEventNotificationQuery for Server {
)
.ascending(),
|key, value| {
let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?;
notifications.push(Notification {
document_id: key.deserialize_be_u32(key.len() - U32_LEN)?,
document_id,
created: key.deserialize_be_u64(key.len() - U32_LEN - U64_LEN)?,
event_id: value.deserialize_be_u32(0)?,
});
document_ids.insert(document_id);
Ok(true)
},
@@ -164,18 +167,13 @@ impl CalendarEventNotificationQuery for Server {
notifications.reverse();
}
let results = self
.search_store()
.query(
account_id,
Collection::CalendarEventNotification,
filters,
vec![],
)
.await?;
let results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_mask(document_ids)
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
cache.get_state(false),
&request,
crates/jmap/src/contact/query.rs
+14 -14
View File
@@ -14,12 +14,10 @@ use jmap_proto::{
};
use store::{
roaring::RoaringBitmap,
search::{ContactSearchField, SearchComparator, SearchFilter},
};
use types::{
acl::Acl,
collection::{Collection, SyncCollection},
search::{ContactSearchField, SearchComparator, SearchFilter, SearchQuery},
write::SearchIndex,
};
use types::{acl::Acl, collection::SyncCollection};
use utils::sanitize_email;
pub trait ContactCardQuery: Sync + Send {
@@ -41,8 +39,6 @@ impl ContactCardQuery for Server {
let cache = self
.fetch_dav_resources(access_token, account_id, SyncCollection::AddressBook)
.await?;
let filter_mask = (access_token.is_shared(account_id))
.then(|| cache.shared_items(access_token, [Acl::ReadItems], true));
for cond in std::mem::take(&mut request.filter) {
match cond {
@@ -212,7 +208,17 @@ impl ContactCardQuery for Server {
let results = self
.search_store()
.query(account_id, Collection::ContactCard, filters, comparators)
.query(
SearchQuery::new(SearchIndex::Contacts)
.with_filters(filters)
.with_comparators(comparators)
.with_account_id(account_id)
.with_mask(if access_token.is_shared(account_id) {
cache.shared_items(access_token, [Acl::ReadItems], true)
} else {
cache.document_ids(false).collect()
}),
)
.await?;
let mut response = QueryResponseBuilder::new(
@@ -223,12 +229,6 @@ impl ContactCardQuery for Server {
);
for document_id in results {
if filter_mask
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
}
if !response.add(0, document_id) {
break;
}
crates/jmap/src/email/query.rs
+19 -16
View File
@@ -17,10 +17,11 @@ use std::{borrow::Cow, future::Future};
use store::{
ahash::{AHashMap, AHashSet},
roaring::RoaringBitmap,
search::{EmailSearchField, SearchComparator, SearchFilter},
search::{EmailSearchField, SearchComparator, SearchFilter, SearchQuery},
write::SearchIndex,
};
use trc::AddContext;
use types::{acl::Acl, collection::Collection, keyword::Keyword};
use types::{acl::Acl, keyword::Keyword};
pub trait EmailQuery: Sync + Send {
fn email_query(
@@ -331,7 +332,22 @@ impl EmailQuery for Server {
let results = self
.search_store()
.query(account_id, Collection::Email, filters, comparators)
.query(
SearchQuery::new(SearchIndex::Email)
.with_filters(filters)
.with_comparators(comparators)
.with_account_id(account_id)
.with_mask(if access_token.is_shared(account_id) {
cached_messages.shared_messages(access_token, Acl::ReadItems)
} else {
cached_messages
.emails
.items
.iter()
.map(|item| item.document_id)
.collect()
}),
)
.await?;
let mut response = QueryResponseBuilder::new(
@@ -342,23 +358,10 @@ impl EmailQuery for Server {
);
if !results.is_empty() {
let filter_ids = if access_token.is_shared(account_id) {
cached_messages
.shared_messages(access_token, Acl::ReadItems)
.into()
} else {
None
};
let collapse_threads = request.arguments.collapse_threads.unwrap_or(false);
let mut seen_thread_ids = AHashSet::new();
for document_id in results {
if filter_ids
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
}
let Some(thread_id) = cached_messages
.email_by_id(&document_id)
.map(|email| email.thread_id)
crates/jmap/src/file/query.rs
+14 -17
View File
@@ -12,11 +12,12 @@ use jmap_proto::{
object::file_node::{FileNode, FileNodeFilter},
request::MaybeInvalid,
};
use store::{roaring::RoaringBitmap, search::SearchFilter};
use types::{
acl::Acl,
collection::{Collection, SyncCollection},
use store::{
roaring::RoaringBitmap,
search::{SearchFilter, SearchQuery},
write::SearchIndex,
};
use types::{acl::Acl, collection::SyncCollection};
pub trait FileNodeQuery: Sync + Send {
fn file_node_query(
@@ -37,8 +38,6 @@ impl FileNodeQuery for Server {
let cache = self
.fetch_dav_resources(access_token, account_id, SyncCollection::FileNode)
.await?;
let filter_mask = (access_token.is_shared(account_id))
.then(|| cache.shared_containers(access_token, [Acl::ReadItems], true));
for cond in std::mem::take(&mut request.filter) {
match cond {
@@ -143,25 +142,23 @@ impl FileNodeQuery for Server {
.details("Sorting is not supported on FileNode"));
}
let results = self
.search_store()
.query(account_id, Collection::FileNode, filters, vec![])
.await?;
let results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_mask(if access_token.is_shared(account_id) {
cache.shared_containers(access_token, [Acl::ReadItems], true)
} else {
cache.document_ids(false).collect()
})
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
cache.get_state(false),
&request,
);
for document_id in results {
if filter_mask
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
}
if !response.add(0, document_id) {
break;
}
crates/jmap/src/mailbox/query.rs
+23 -23
View File
@@ -17,9 +17,10 @@ use std::{
};
use store::{
roaring::RoaringBitmap,
search::{SearchComparator, SearchFilter},
search::{SearchComparator, SearchFilter, SearchQuery},
write::SearchIndex,
};
use types::{acl::Acl, collection::Collection, special_use::SpecialUse};
use types::{acl::Acl, special_use::SpecialUse};
pub trait MailboxQuery: Sync + Send {
fn mailbox_query(
@@ -225,43 +226,47 @@ impl MailboxQuery for Server {
});
}
let results = self
.search_store()
.query(account_id, Collection::Mailbox, filters, comparators)
.await?;
let results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_comparators(comparators)
.with_mask(if access_token.is_shared(account_id) {
mailboxes.shared_mailboxes(access_token, Acl::Read)
} else {
mailboxes
.mailboxes
.items
.iter()
.map(|m| m.document_id)
.collect()
})
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
mailboxes.get_state(true),
&request,
);
if !results.is_empty() {
let filter_ids = if access_token.is_shared(account_id) {
mailboxes.shared_mailboxes(access_token, Acl::Read).into()
} else {
None
};
// Filter as tree
if filter_as_tree {
let mut total_filtered = 0;
let mut is_page_full = false;
for document_id in &results {
let mut check_id = *document_id;
let mut check_id = document_id;
for _ in 0..self.core.jmap.mailbox_max_depth {
if let Some(mailbox) = mailboxes.mailbox_by_id(&check_id) {
if let Some(parent_id) = mailbox.parent_id() {
if results.contains(&parent_id) {
if results.contains(parent_id) {
check_id = parent_id;
} else {
break;
}
} else {
total_filtered += 1;
if !is_page_full && !response.add(0, *document_id) {
if !is_page_full && !response.add(0, document_id) {
is_page_full = true;
}
}
@@ -270,16 +275,11 @@ impl MailboxQuery for Server {
}
if total_filtered != results.len() {
response.response.total = Some(total_filtered);
response.response.total = Some(total_filtered as usize);
}
} else {
for document_id in results {
if filter_ids
.as_ref()
.is_some_and(|filter_ids| !filter_ids.contains(document_id))
{
continue;
} else if !response.add(0, document_id) {
if !response.add(0, document_id) {
break;
}
}
crates/jmap/src/principal/query.rs
+12 -10
View File
@@ -4,6 +4,7 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::api::query::QueryResponseBuilder;
use common::{Server, auth::AccessToken};
use directory::{Permission, QueryParams, Type, backend::internal::manage::ManageDirectory};
use http_proto::HttpSessionData;
@@ -13,11 +14,12 @@ use jmap_proto::{
types::state::State,
};
use std::future::Future;
use store::{roaring::RoaringBitmap, search::SearchFilter};
use store::{
roaring::RoaringBitmap,
search::{SearchFilter, SearchQuery},
write::SearchIndex,
};
use trc::AddContext;
use types::collection::Collection;
use crate::api::query::QueryResponseBuilder;
pub trait PrincipalQuery: Sync + Send {
fn principal_query(
@@ -171,20 +173,20 @@ impl PrincipalQuery for Server {
}
}
let results = self
.search_store()
.query(u32::MAX, Collection::Principal, filters, vec![])
.await?;
let results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_mask(principal_ids)
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
State::Initial,
&request,
);
for document_id in results {
if principal_ids.contains(document_id) && !response.add(0, document_id) {
if !response.add(0, document_id) {
break;
}
}
crates/jmap/src/sieve/query.rs
+11 -11
View File
@@ -13,8 +13,10 @@ use jmap_proto::{
};
use std::future::Future;
use store::{
IndexKeyPrefix, IterateParams, U32_LEN, ahash::AHashSet, roaring::RoaringBitmap,
search::SearchFilter, write::key::DeserializeBigEndian,
IndexKeyPrefix, IterateParams, U32_LEN,
roaring::RoaringBitmap,
search::{SearchFilter, SearchQuery},
write::{SearchIndex, key::DeserializeBigEndian},
};
use trc::AddContext;
use types::{
@@ -153,15 +155,13 @@ impl SieveScriptQuery for Server {
};
}
let mut results = self
.search_store()
.query(account_id, Collection::SieveScript, filters, vec![])
.await?
.into_iter()
.collect::<AHashSet<_>>();
let mut results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_mask(document_ids)
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
self.get_state(account_id, SyncCollection::SieveScript)
.await?,
@@ -170,7 +170,7 @@ impl SieveScriptQuery for Server {
if !results.is_empty() {
if matches!(sort_by_active, Some(true))
&& results.remove(&active_script_id.unwrap_or_default())
&& results.remove(active_script_id.unwrap_or_default())
&& !response.add(0, active_script_id.unwrap())
{
return response.build();
@@ -178,7 +178,7 @@ impl SieveScriptQuery for Server {
let mut last_id = None;
for (document_id, _) in names {
if results.contains(&document_id) {
if results.contains(document_id) {
if sort_by_active.is_some() && Some(document_id) == active_script_id {
last_id = Some(document_id);
} else if !response.add(0, document_id) {
crates/jmap/src/submission/query.rs
+14 -11
View File
@@ -17,8 +17,8 @@ use store::{
IterateParams, U32_LEN, U64_LEN, ValueKey,
ahash::AHashSet,
roaring::RoaringBitmap,
search::SearchFilter,
write::{IndexPropertyClass, ValueClass, key::DeserializeBigEndian, now},
search::{SearchFilter, SearchQuery},
write::{IndexPropertyClass, SearchIndex, ValueClass, key::DeserializeBigEndian, now},
};
use trc::AddContext;
use types::{
@@ -50,6 +50,7 @@ impl EmailSubmissionQuery for Server {
let account_id = request.account_id.document_id();
let mut submissions = Vec::with_capacity(16);
let mut document_ids = RoaringBitmap::new();
self.store()
.iterate(
@@ -75,8 +76,10 @@ impl EmailSubmissionQuery for Server {
)
.ascending(),
|key, value| {
let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?;
submissions.push(Submission {
document_id: key.deserialize_be_u32(key.len() - U32_LEN)?,
document_id,
send_at: key.deserialize_be_u64(key.len() - U32_LEN - U64_LEN)?,
email_id: value.deserialize_be_u32(0)?,
thread_id: value.deserialize_be_u32(U32_LEN)?,
@@ -84,6 +87,8 @@ impl EmailSubmissionQuery for Server {
undo_status: value.last().copied().unwrap(),
});
document_ids.insert(document_id);
Ok(true)
},
)
@@ -188,15 +193,13 @@ impl EmailSubmissionQuery for Server {
}
}
let results = self
.search_store()
.query(account_id, Collection::ContactCard, filters, vec![])
.await?
.into_iter()
.collect::<AHashSet<_>>();
let results = SearchQuery::new(SearchIndex::InMemory)
.with_filters(filters)
.with_mask(document_ids)
.execute();
let mut response = QueryResponseBuilder::new(
results.len(),
results.len() as usize,
self.core.jmap.query_max_results,
self.get_state(account_id, SyncCollection::EmailSubmission)
.await?,
@@ -232,7 +235,7 @@ impl EmailSubmissionQuery for Server {
}
for submission in submissions {
if results.contains(&submission.document_id)
if results.contains(submission.document_id)
&& !response.add(0, submission.document_id)
{
break;
crates/services/src/task_manager/bayes.rs
+75 -40
View File
@@ -5,67 +5,102 @@
*/
use common::Server;
use email::message::metadata::MessageMetadata;
use mail_parser::MessageParser;
use spam_filter::{
SpamFilterInput, analysis::init::SpamFilterInit, modules::bayes::BayesClassifier,
};
use std::time::Instant;
use trc::{SpamEvent, TaskQueueEvent};
use types::{blob_hash::BlobHash, collection::Collection};
use types::{collection::Collection, field::EmailField};
pub trait BayesTrainTask: Sync + Send {
fn bayes_train(
&self,
account_id: u32,
document_id: u32,
hash: &BlobHash,
learn_spam: bool,
) -> impl Future<Output = bool> + Send;
}
impl BayesTrainTask for Server {
async fn bayes_train(
&self,
account_id: u32,
document_id: u32,
hash: &BlobHash,
learn_spam: bool,
) -> bool {
async fn bayes_train(&self, account_id: u32, document_id: u32, learn_spam: bool) -> bool {
let op_start = Instant::now();
// Obtain raw message
if let Ok(Some(raw_message)) = self
.blob_store()
.get_blob(hash.as_slice(), 0..usize::MAX)
// Obtain metadata
let metadata_ = match self
.archive_by_property(
account_id,
Collection::Email,
document_id,
EmailField::Metadata.into(),
)
.await
{
// Train bayes classifier for account
self.bayes_train_if_balanced(
&self.spam_filter_init(SpamFilterInput::from_account_message(
&MessageParser::new().parse(&raw_message).unwrap_or_default(),
account_id,
0,
)),
learn_spam,
)
.await;
Ok(Some(metadata)) => metadata,
Ok(None) => {
trc::event!(
TaskQueue(TaskQueueEvent::MetadataNotFound),
AccountId = account_id,
Collection = Collection::Email,
DocumentId = document_id,
);
return false;
}
Err(err) => {
trc::error!(err.caused_by(trc::location!()));
return false;
}
};
trc::event!(
Spam(SpamEvent::TrainAccount),
AccountId = account_id,
Collection = Collection::Email,
DocumentId = document_id,
Details = if learn_spam { "spam" } else { "ham" },
Elapsed = op_start.elapsed(),
);
true
} else {
trc::event!(
TaskQueue(TaskQueueEvent::BlobNotFound),
AccountId = account_id,
DocumentId = document_id,
BlobId = hash.as_slice(),
);
false
let metadata = match metadata_.unarchive::<MessageMetadata>() {
Ok(metadata) => metadata,
Err(err) => {
trc::error!(err.caused_by(trc::location!()));
return false;
}
};
// Obtain raw message
match self
.blob_store()
.get_blob(metadata.blob_hash.0.as_slice(), 0..usize::MAX)
.await
{
Ok(Some(raw_message)) => {
// Train bayes classifier for account
self.bayes_train_if_balanced(
&self.spam_filter_init(SpamFilterInput::from_account_message(
&MessageParser::new().parse(&raw_message).unwrap_or_default(),
account_id,
0,
)),
learn_spam,
)
.await;
trc::event!(
Spam(SpamEvent::TrainAccount),
AccountId = account_id,
Collection = Collection::Email,
DocumentId = document_id,
Details = if learn_spam { "spam" } else { "ham" },
Elapsed = op_start.elapsed(),
);
true
}
Ok(None) => {
trc::event!(
TaskQueue(TaskQueueEvent::BlobNotFound),
AccountId = account_id,
DocumentId = document_id,
BlobId = metadata.blob_hash.0.as_slice(),
);
false
}
Err(err) => {
trc::error!(err.caused_by(trc::location!()));
false
}
}
}
}
crates/services/src/task_manager/fts.rs
-451
View File
@@ -1,451 +0,0 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use common::Server;
use directory::{Type, backend::internal::manage::ManageDirectory};
use email::message::metadata::MessageMetadata;
use std::time::Instant;
use store::{
IterateParams, SerializeInfallible, U32_LEN, ValueKey,
ahash::AHashMap,
roaring::RoaringBitmap,
write::{
BatchBuilder, BlobOp, SearchIndex, TaskQueueClass, ValueClass, key::DeserializeBigEndian,
now,
},
};
use trc::{AddContext, MessageIngestEvent, TaskQueueEvent};
use types::{
blob_hash::{BLOB_HASH_LEN, BlobHash},
collection::Collection,
field::EmailField,
};
pub trait FtsIndexTask: Sync + Send {
fn fts_index(
&self,
account_id: u32,
document_id: u32,
hash: &BlobHash,
) -> impl Future<Output = bool> + Send;
fn fts_reindex(
&self,
account_id: Option<u32>,
tenant_id: Option<u32>,
) -> impl Future<Output = trc::Result<()>> + Send;
}
impl FtsIndexTask for Server {
async fn fts_index(&self, account_id: u32, document_id: u32, hash: &BlobHash) -> bool {
let todo = "merge threads";
let todo = "combine task with bayes train if needed";
let todo = "delete Threading field on delete";
/*loop {
// Find messages with a matching subject
let mut subj_results = RoaringBitmap::new();
self.store()
.iterate(
IterateParams::new(
IndexKey {
account_id,
collection: Collection::Email.into(),
document_id: 0,
field: EmailField::Subject.into(),
key: thread_name.clone(),
},
IndexKey {
account_id,
collection: Collection::Email.into(),
document_id: u32::MAX,
field: EmailField::Subject.into(),
key: thread_name.clone(),
},
)
.no_values()
.ascending(),
|key, _| {
let id_pos = key.len() - U32_LEN;
let value = key.get(IndexKeyPrefix::len()..id_pos).ok_or_else(|| {
trc::Error::corrupted_key(key, None, trc::location!())
})?;
if value == thread_name {
subj_results.insert(key.deserialize_be_u32(id_pos)?);
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
// No matching subjects were found, skip early
if subj_results.is_empty() {
return Ok(ThreadResult::Id(None));
}
// Find messages with matching references
let mut results = RoaringBitmap::new();
let mut found_message_id = Vec::new();
self.store()
.iterate(
IterateParams::new(
IndexKey {
account_id,
collection: Collection::Email.into(),
document_id: 0,
field: EmailField::References.into(),
key: references.first().unwrap().to_vec(),
},
IndexKey {
account_id,
collection: Collection::Email.into(),
document_id: u32::MAX,
field: EmailField::References.into(),
key: references.last().unwrap().to_vec(),
},
)
.no_values()
.ascending(),
|key, _| {
let id_pos = key.len() - U32_LEN;
let mut value =
key.get(IndexKeyPrefix::len()..id_pos).ok_or_else(|| {
trc::Error::corrupted_key(key, None, trc::location!())
})?;
let document_id = key.deserialize_be_u32(id_pos)?;
if let Some(message_id) = value.strip_suffix(&[0]) {
value = message_id;
if skip_duplicate.is_some_and(|(message_id, _)| message_id == value) {
found_message_id.push(document_id);
}
}
if subj_results.contains(document_id)
&& references.binary_search(&value).is_ok()
{
results.insert(document_id);
if subj_results.len() == results.len() {
return Ok(false);
}
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
// No matching messages
if results.is_empty() {
return Ok(ThreadResult::Id(None));
}
// Fetch cached messages
let cache = self
.get_cached_messages(account_id)
.await
.caused_by(trc::location!())?;
// Skip duplicate messages
if !found_message_id.is_empty()
&& cache
.in_mailbox(skip_duplicate.unwrap().1)
.any(|m| found_message_id.contains(&m.document_id))
{
return Ok(ThreadResult::Skip);
}
// Find the most common threadId
let mut thread_counts = AHashMap::<u32, u32>::with_capacity(16);
let mut thread_id = u32::MAX;
let mut thread_count = 0;
for item in &cache.emails.items {
if results.contains(item.document_id) {
let tc = thread_counts.entry(item.thread_id).or_default();
*tc += 1;
if *tc > thread_count {
thread_count = *tc;
thread_id = item.thread_id;
}
}
}
if thread_id == u32::MAX {
return Ok(ThreadResult::Id(None));
} else if thread_counts.len() == 1 {
return Ok(ThreadResult::Id(Some(thread_id)));
}
// Delete all but the most common threadId
let mut batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Thread);
for &delete_thread_id in thread_counts.keys() {
if delete_thread_id != thread_id {
batch
.with_document(delete_thread_id)
.log_container_delete(SyncCollection::Thread);
}
}
// Move messages to the new threadId
batch.with_collection(Collection::Email);
for item in &cache.emails.items {
if thread_id == item.thread_id || !thread_counts.contains_key(&item.thread_id) {
continue;
}
if let Some(data_) = self
.archive(account_id, Collection::Email, item.document_id)
.await
.caused_by(trc::location!())?
{
let data = data_
.to_unarchived::<MessageData>()
.caused_by(trc::location!())?;
if data.inner.thread_id != item.thread_id {
continue;
}
let mut new_data = data.deserialize().caused_by(trc::location!())?;
new_data.thread_id = thread_id;
batch
.with_document(item.document_id)
.custom(
ObjectIndexBuilder::new()
.with_current(data)
.with_changes(new_data),
)
.caused_by(trc::location!())?;
}
}
match self.commit_batch(batch).await {
Ok(_) => return Ok(ThreadResult::Id(Some(thread_id))),
Err(err) if err.is_assertion_failure() && try_count < MAX_RETRIES => {
let backoff = store::rand::rng().random_range(50..=300);
tokio::time::sleep(Duration::from_millis(backoff)).await;
try_count += 1;
}
Err(err) => {
return Err(err.caused_by(trc::location!()));
}
}
}*/
// Obtain raw message
let op_start = Instant::now();
let raw_message = if let Ok(Some(raw_message)) = self
.blob_store()
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
{
raw_message
} else {
trc::event!(
TaskQueue(TaskQueueEvent::BlobNotFound),
AccountId = account_id,
DocumentId = document_id,
BlobId = hash.as_slice(),
);
return false;
};
match self
.archive_by_property(
account_id,
Collection::Email,
document_id,
EmailField::Metadata.into(),
)
.await
{
Ok(Some(metadata_)) => {
match metadata_.unarchive::<MessageMetadata>() {
Ok(metadata) if metadata.blob_hash.0.as_slice() == hash.as_slice() => {
// Index message
/*let document =
FtsDocument::with_default_language(self.core.jmap.default_language)
.with_account_id(account_id)
.with_collection(Collection::Email)
.with_document_id(document_id)
.index_message(metadata, &raw_message);
if let Err(err) = self.core.storage.fts.index(document).await {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.details("Failed to index email in FTS index")
);
return false;
}*/
trc::event!(
MessageIngest(MessageIngestEvent::FtsIndex),
AccountId = account_id,
Collection = Collection::Email,
DocumentId = document_id,
Elapsed = op_start.elapsed(),
);
}
Err(err) => {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.details("Failed to unarchive email metadata")
);
}
_ => {
// The message was probably deleted or overwritten
trc::event!(
TaskQueue(TaskQueueEvent::MetadataNotFound),
Details = "E-mail blob hash mismatch",
AccountId = account_id,
DocumentId = document_id,
);
}
}
true
}
Err(err) => {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.caused_by(trc::location!())
.details("Failed to retrieve email metadata")
);
false
}
_ => {
// The message was probably deleted or overwritten
trc::event!(
TaskQueue(TaskQueueEvent::MetadataNotFound),
Details = "E-mail metadata not found",
AccountId = account_id,
DocumentId = document_id,
);
true
}
}
}
async fn fts_reindex(
&self,
account_id: Option<u32>,
tenant_id: Option<u32>,
) -> trc::Result<()> {
let accounts = if let Some(account_id) = account_id {
RoaringBitmap::from_sorted_iter([account_id]).unwrap()
} else {
let mut accounts = RoaringBitmap::new();
for principal in self
.core
.storage
.data
.list_principals(
None,
tenant_id,
&[Type::Individual, Type::Group],
false,
0,
0,
)
.await
.caused_by(trc::location!())?
.items
{
accounts.insert(principal.id());
}
accounts
};
// Validate linked blobs
let from_key = ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Blob(BlobOp::Link {
hash: BlobHash::default(),
}),
};
let to_key = ValueKey {
account_id: u32::MAX,
collection: u8::MAX,
document_id: u32::MAX,
class: ValueClass::Blob(BlobOp::Link {
hash: BlobHash::new_max(),
}),
};
let mut document_ids: AHashMap<u32, Vec<u32>> = AHashMap::new();
self.core
.storage
.data
.iterate(
IterateParams::new(from_key, to_key).ascending().no_values(),
|key, _| {
let account_id = key.deserialize_be_u32(BLOB_HASH_LEN)?;
let collection = *key
.get(BLOB_HASH_LEN + U32_LEN)
.ok_or_else(|| trc::Error::corrupted_key(key, None, trc::location!()))?;
if accounts.contains(account_id) && collection == Collection::Email as u8 {
document_ids
.entry(account_id)
.or_default()
.push(key.deserialize_be_u32(key.len() - U32_LEN)?);
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
let due = now();
for (account_id, document_ids) in document_ids {
let mut batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Email);
for document_id in document_ids {
batch.with_document(document_id).set(
ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
due,
index: SearchIndex::Email,
is_insert: true,
}),
0u64.serialize(),
);
if batch.len() >= 2000 {
self.core.storage.data.write(batch.build_all()).await?;
batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Email);
}
}
if !batch.is_empty() {
self.core.storage.data.write(batch.build_all()).await?;
}
}
// Request indexing
self.notify_task_queue();
Ok(())
}
}
crates/services/src/task_manager/index.rs
+297
View File
@@ -0,0 +1,297 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::task_manager::{IndexAction, Task};
use common::Server;
use directory::{Type, backend::internal::manage::ManageDirectory};
use groupware::cache::GroupwareCache;
use store::{
IterateParams, SerializeInfallible, U32_LEN, ValueKey,
ahash::AHashMap,
roaring::RoaringBitmap,
write::{
BatchBuilder, BlobOp, SearchIndex, TaskQueueClass, ValueClass, key::DeserializeBigEndian,
now,
},
};
use trc::AddContext;
use types::{
blob_hash::{BLOB_HASH_LEN, BlobHash},
collection::{Collection, SyncCollection},
};
pub(crate) trait SearchIndexTask: Sync + Send {
fn index(&self, tasks: &[Task<IndexAction>]) -> impl Future<Output = bool> + Send;
}
pub trait ReindexIndexTask: Sync + Send {
fn reindex(
&self,
index: SearchIndex,
account_id: Option<u32>,
tenant_id: Option<u32>,
) -> impl Future<Output = trc::Result<()>> + Send;
}
impl SearchIndexTask for Server {
async fn index(&self, tasks: &[Task<IndexAction>]) -> bool {
todo!()
// Obtain raw message
/*let op_start = Instant::now();
let raw_message = if let Ok(Some(raw_message)) = self
.blob_store()
.get_blob(hash.as_slice(), 0..usize::MAX)
.await
{
raw_message
} else {
trc::event!(
TaskQueue(TaskQueueEvent::BlobNotFound),
AccountId = account_id,
DocumentId = document_id,
BlobId = hash.as_slice(),
);
return false;
};
match self
.archive_by_property(
account_id,
Collection::Email,
document_id,
EmailField::Metadata.into(),
)
.await
{
Ok(Some(metadata_)) => {
match metadata_.unarchive::<MessageMetadata>() {
Ok(metadata) if metadata.blob_hash.0.as_slice() == hash.as_slice() => {
// Index message
/*let document =
FtsDocument::with_default_language(self.core.jmap.default_language)
.with_account_id(account_id)
.with_collection(Collection::Email)
.with_document_id(document_id)
.index_message(metadata, &raw_message);
if let Err(err) = self.core.storage.fts.index(document).await {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.details("Failed to index email in FTS index")
);
return false;
}*/
trc::event!(
MessageIngest(MessageIngestEvent::FtsIndex),
AccountId = account_id,
Collection = Collection::Email,
DocumentId = document_id,
Elapsed = op_start.elapsed(),
);
}
Err(err) => {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.details("Failed to unarchive email metadata")
);
}
_ => {
// The message was probably deleted or overwritten
trc::event!(
TaskQueue(TaskQueueEvent::MetadataNotFound),
Details = "E-mail blob hash mismatch",
AccountId = account_id,
DocumentId = document_id,
);
}
}
true
}
Err(err) => {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.caused_by(trc::location!())
.details("Failed to retrieve email metadata")
);
false
}
_ => {
// The message was probably deleted or overwritten
trc::event!(
TaskQueue(TaskQueueEvent::MetadataNotFound),
Details = "E-mail metadata not found",
AccountId = account_id,
DocumentId = document_id,
);
true
}
}*/
}
}
impl ReindexIndexTask for Server {
async fn reindex(
&self,
index: SearchIndex,
account_id: Option<u32>,
tenant_id: Option<u32>,
) -> trc::Result<()> {
let accounts = if let Some(account_id) = account_id {
RoaringBitmap::from_sorted_iter([account_id]).unwrap()
} else {
let mut accounts = RoaringBitmap::new();
for principal in self
.core
.storage
.data
.list_principals(
None,
tenant_id,
&[Type::Individual, Type::Group],
false,
0,
0,
)
.await
.caused_by(trc::location!())?
.items
{
accounts.insert(principal.id());
}
accounts
};
let due = now();
match index {
SearchIndex::Email => {
// Validate linked blobs
let from_key = ValueKey {
account_id: 0,
collection: 0,
document_id: 0,
class: ValueClass::Blob(BlobOp::Link {
hash: BlobHash::default(),
}),
};
let to_key = ValueKey {
account_id: u32::MAX,
collection: u8::MAX,
document_id: u32::MAX,
class: ValueClass::Blob(BlobOp::Link {
hash: BlobHash::new_max(),
}),
};
let mut document_ids: AHashMap<u32, Vec<u32>> = AHashMap::new();
self.core
.storage
.data
.iterate(
IterateParams::new(from_key, to_key).ascending().no_values(),
|key, _| {
let account_id = key.deserialize_be_u32(BLOB_HASH_LEN)?;
let collection =
*key.get(BLOB_HASH_LEN + U32_LEN).ok_or_else(|| {
trc::Error::corrupted_key(key, None, trc::location!())
})?;
if accounts.contains(account_id)
&& collection == Collection::Email as u8
{
document_ids
.entry(account_id)
.or_default()
.push(key.deserialize_be_u32(key.len() - U32_LEN)?);
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
for (account_id, document_ids) in document_ids {
let mut batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Email);
for document_id in document_ids {
batch.with_document(document_id).set(
ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
due,
index: SearchIndex::Email,
is_insert: true,
}),
0u64.serialize(),
);
if batch.len() >= 2000 {
self.core.storage.data.write(batch.build_all()).await?;
batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Email);
}
}
if !batch.is_empty() {
self.core.storage.data.write(batch.build_all()).await?;
}
}
}
SearchIndex::Calendar | SearchIndex::Contacts => {
for account_id in accounts {
let Some(cache) = self.cached_dav_resources(
account_id,
if index == SearchIndex::Calendar {
SyncCollection::Calendar
} else {
SyncCollection::AddressBook
},
) else {
continue;
};
let mut batch = BatchBuilder::new();
batch.with_account_id(account_id);
for document_id in cache.document_ids(false) {
batch.with_document(document_id).set(
ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
due,
index,
is_insert: true,
}),
0u64.serialize(),
);
if batch.len() >= 2000 {
self.core.storage.data.write(batch.build_all()).await?;
batch = BatchBuilder::new();
batch.with_account_id(account_id);
}
}
if !batch.is_empty() {
self.core.storage.data.write(batch.build_all()).await?;
}
}
}
SearchIndex::File | SearchIndex::TracingSpan | SearchIndex::InMemory => (),
}
// Request indexing
self.notify_task_queue();
Ok(())
}
}
crates/services/src/task_manager/lock.rs
+325
View File
@@ -0,0 +1,325 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::task_manager::*;
pub(crate) trait TaskLockManager: Sync + Send {
fn try_lock_task(
&self,
account_id: u32,
document_id: u32,
lock_key: Vec<u8>,
lock_expiry: u64,
) -> impl Future<Output = bool> + Send;
fn remove_index_lock(&self, lock_key: Vec<u8>) -> impl Future<Output = ()> + Send;
}
impl TaskLockManager for Server {
async fn try_lock_task(
&self,
account_id: u32,
document_id: u32,
lock_key: Vec<u8>,
lock_expiry: u64,
) -> bool {
match self
.in_memory_store()
.try_lock(KV_LOCK_TASK, &lock_key, lock_expiry)
.await
{
Ok(result) => {
if !result {
trc::event!(
TaskQueue(TaskQueueEvent::TaskLocked),
AccountId = account_id,
DocumentId = document_id,
Expires = trc::Value::Timestamp(now() + lock_expiry),
);
}
result
}
Err(err) => {
trc::error!(
err.account_id(account_id)
.document_id(document_id)
.details("Failed to lock task")
);
false
}
}
}
async fn remove_index_lock(&self, lock_key: Vec<u8>) {
if let Err(err) = self
.in_memory_store()
.remove_lock(KV_LOCK_TASK, &lock_key)
.await
{
trc::error!(
err.details("Failed to unlock task")
.ctx(trc::Key::Key, lock_key)
.caused_by(trc::location!())
);
}
}
}
pub(crate) trait TaskLock {
fn account_id(&self) -> u32;
fn document_id(&self) -> u32;
fn remove_lock(&self) -> bool;
fn lock_key(&self) -> Vec<u8>;
fn lock_expiry(&self) -> u64;
fn value_classes(&self) -> impl Iterator<Item = ValueClass>;
}
impl TaskLock for Task<IndexAction> {
fn account_id(&self) -> u32 {
self.account_id
}
fn document_id(&self) -> u32 {
self.document_id
}
fn remove_lock(&self) -> bool {
true
}
fn lock_key(&self) -> Vec<u8> {
KeySerializer::new((U32_LEN * 2) + U64_LEN + 2)
.write(0u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.write(self.action.index.to_u8())
.finalize()
}
fn lock_expiry(&self) -> u64 {
INDEX_EXPIRY
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
std::iter::once(ValueClass::TaskQueue(TaskQueueClass::UpdateIndex {
due: self.due,
index: self.action.index,
is_insert: self.action.is_insert,
}))
}
}
impl TaskLock for Task<bool> {
fn account_id(&self) -> u32 {
self.account_id
}
fn document_id(&self) -> u32 {
self.document_id
}
fn remove_lock(&self) -> bool {
false
}
fn lock_key(&self) -> Vec<u8> {
KeySerializer::new((U32_LEN * 2) + 1)
.write(1u8)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize()
}
fn lock_expiry(&self) -> u64 {
BAYES_LOCK_EXPIRY
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
std::iter::once(ValueClass::TaskQueue(TaskQueueClass::BayesTrain {
due: self.due,
learn_spam: self.action,
}))
}
}
impl TaskLock for Task<CalendarAlarm> {
fn account_id(&self) -> u32 {
self.account_id
}
fn document_id(&self) -> u32 {
self.document_id
}
fn remove_lock(&self) -> bool {
true
}
fn lock_key(&self) -> Vec<u8> {
KeySerializer::new((U32_LEN * 2) + U64_LEN + 1)
.write(2u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize()
}
fn lock_expiry(&self) -> u64 {
ALARM_EXPIRY
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
std::iter::once(ValueClass::TaskQueue(TaskQueueClass::SendAlarm {
event_id: self.action.event_id,
alarm_id: self.action.alarm_id,
due: self.due,
is_email_alert: matches!(self.action.typ, CalendarAlarmType::Email { .. }),
}))
}
}
impl TaskLock for Task<ImipAction> {
fn account_id(&self) -> u32 {
self.account_id
}
fn document_id(&self) -> u32 {
self.document_id
}
fn remove_lock(&self) -> bool {
true
}
fn lock_key(&self) -> Vec<u8> {
KeySerializer::new((U32_LEN * 2) + U64_LEN + 1)
.write(3u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize()
}
fn lock_expiry(&self) -> u64 {
ALARM_EXPIRY
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
[
Some(ValueClass::TaskQueue(TaskQueueClass::SendImip {
due: self.due,
is_payload: false,
})),
Some(ValueClass::TaskQueue(TaskQueueClass::SendImip {
due: self.due,
is_payload: true,
})),
]
.into_iter()
.flatten()
}
}
impl TaskLock for Task<MergeThreadIds<AHashSet<u32>>> {
fn account_id(&self) -> u32 {
self.account_id
}
fn document_id(&self) -> u32 {
self.document_id
}
fn remove_lock(&self) -> bool {
true
}
fn lock_key(&self) -> Vec<u8> {
KeySerializer::new((U32_LEN * 2) + U64_LEN + 1)
.write(4u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize()
}
fn lock_expiry(&self) -> u64 {
ALARM_EXPIRY
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
std::iter::once(ValueClass::TaskQueue(TaskQueueClass::MergeThreads {
due: self.due,
}))
}
}
impl Task<TaskAction> {
pub(crate) fn lock_expiry(&self) -> u64 {
match &self.action {
TaskAction::UpdateIndex(_) => INDEX_EXPIRY,
TaskAction::BayesTrain(_) => BAYES_LOCK_EXPIRY,
TaskAction::SendAlarm(_) => ALARM_EXPIRY,
_ => ALARM_EXPIRY,
}
}
pub(crate) fn deserialize(key: &[u8], value: &[u8]) -> trc::Result<Self> {
let document_id = key.deserialize_be_u32(U64_LEN + U32_LEN + 1)?;
Ok(Task {
due: key.deserialize_be_u64(0)?,
account_id: key.deserialize_be_u32(U64_LEN)?,
document_id,
action: match key.get(U64_LEN + U32_LEN) {
Some(v @ (7 | 8)) => TaskAction::UpdateIndex(IndexAction {
index: key
.last()
.copied()
.and_then(SearchIndex::try_from_u8)
.ok_or_else(|| trc::Error::corrupted_key(key, None, trc::location!()))?,
is_insert: *v == 7,
}),
Some(v @ (1 | 2)) => TaskAction::BayesTrain(*v == 1),
Some(3) => TaskAction::SendAlarm(CalendarAlarm {
event_id: key.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + 1)?,
alarm_id: key.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + U16_LEN + 1)?,
alarm_time: 0,
typ: CalendarAlarmType::Email {
event_start: value.deserialize_be_u64(0)? as i64,
event_end: value.deserialize_be_u64(U64_LEN)? as i64,
event_start_tz: value.deserialize_be_u16(U64_LEN * 2)?,
event_end_tz: value.deserialize_be_u16((U64_LEN * 2) + U16_LEN)?,
},
}),
Some(6) => {
let recurrence_id = value.deserialize_be_u64(0)? as i64;
TaskAction::SendAlarm(CalendarAlarm {
event_id: key.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + 1)?,
alarm_id: key
.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + U16_LEN + 1)?,
alarm_time: 0,
typ: CalendarAlarmType::Display {
recurrence_id: if recurrence_id != 0 {
Some(recurrence_id)
} else {
None
},
},
})
}
Some(4) => TaskAction::SendImip,
Some(9) => TaskAction::MergeThreads(
MergeThreadIds::deserialize(document_id, value).ok_or_else(|| {
trc::Error::corrupted_key(key, value.into(), trc::location!())
})?,
),
_ => return Err(trc::Error::corrupted_key(key, None, trc::location!())),
},
})
}
}
crates/services/src/task_manager/merge_threads.rs
+192
View File
@@ -0,0 +1,192 @@
/*
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use common::{Server, storage::index::ObjectIndexBuilder};
use email::message::{
ingest::{MergeThreadIds, ThreadMerge},
metadata::MessageData,
};
use std::time::Duration;
use store::{
IndexKeyPrefix, IterateParams, U32_LEN, ValueKey,
ahash::{AHashMap, AHashSet},
rand::Rng,
write::{BatchBuilder, IndexPropertyClass, ValueClass, key::DeserializeBigEndian},
};
use trc::AddContext;
use types::{
collection::{Collection, SyncCollection},
field::EmailField,
};
const MAX_RETRIES: usize = 5;
pub trait MergeThreadsTask: Sync + Send {
fn merge_threads(
&self,
account_id: u32,
threads: &MergeThreadIds<AHashSet<u32>>,
) -> impl Future<Output = bool> + Send;
}
impl MergeThreadsTask for Server {
async fn merge_threads(
&self,
account_id: u32,
threads: &MergeThreadIds<AHashSet<u32>>,
) -> bool {
match merge_threads(self, account_id, threads).await {
Ok(_) => true,
Err(err) => {
trc::error!(
err.account_id(account_id)
.details("Failed to merge threads")
);
false
}
}
}
}
async fn merge_threads(
server: &Server,
account_id: u32,
merge_threads: &MergeThreadIds<AHashSet<u32>>,
) -> trc::Result<()> {
let key_len = IndexKeyPrefix::len() + merge_threads.thread_hash.len() + U32_LEN;
let document_id_pos = key_len - U32_LEN;
let mut thread_merge = ThreadMerge::new();
let mut thread_index = AHashMap::new();
let mut try_count = 0;
'retry: loop {
// Find thread ids
server
.store()
.iterate(
IterateParams::new(
ValueKey {
account_id,
collection: Collection::Email.into(),
document_id: 0,
class: ValueClass::IndexProperty(IndexPropertyClass::Hash {
property: EmailField::Threading.into(),
hash: merge_threads.thread_hash,
}),
},
ValueKey {
account_id,
collection: Collection::Email.into(),
document_id: u32::MAX,
class: ValueClass::IndexProperty(IndexPropertyClass::Hash {
property: EmailField::Threading.into(),
hash: merge_threads.thread_hash,
}),
},
)
.ascending(),
|key, value| {
if key.len() == key_len {
let document_id = key.deserialize_be_u32(document_id_pos)?;
if merge_threads.merge_ids.contains(&document_id) {
let thread_id = value.deserialize_be_u32(0)?;
thread_merge.add(thread_id, document_id);
thread_index.insert(document_id, value.to_vec());
return Ok(
thread_merge.num_document_ids() != merge_threads.merge_ids.len()
);
}
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
if thread_merge.num_thread_ids() < 2 {
// Another process merged the threads already?
return Ok(());
}
let thread_id = thread_merge.merge_thread_id();
// Delete all but the most common threadId
let mut batch = BatchBuilder::new();
batch
.with_account_id(account_id)
.with_collection(Collection::Thread);
for &delete_thread_id in thread_merge.thread_ids() {
if delete_thread_id != thread_id {
batch
.with_document(delete_thread_id)
.log_container_delete(SyncCollection::Thread);
}
}
// Move messages to the new threadId
batch.with_collection(Collection::Email);
for (&group_thread_id, document_ids) in thread_merge.thread_groups() {
if thread_id != group_thread_id {
for &document_id in document_ids {
if let Some(data_) = server
.archive(account_id, Collection::Email, document_id)
.await
.caused_by(trc::location!())?
{
let data = data_
.to_unarchived::<MessageData>()
.caused_by(trc::location!())?;
if data.inner.thread_id != group_thread_id {
try_count += 1;
continue 'retry;
}
// Update thread id
let mut new_data = data
.deserialize::<MessageData>()
.caused_by(trc::location!())?;
new_data.thread_id = thread_id;
batch
.with_document(document_id)
.custom(
ObjectIndexBuilder::new()
.with_current(data)
.with_changes(new_data),
)
.caused_by(trc::location!())?;
// Update thread index property
let mut thread_index = thread_index.remove(&document_id).unwrap();
thread_index[0..U32_LEN].copy_from_slice(&thread_id.to_be_bytes());
batch.set(
ValueClass::IndexProperty(IndexPropertyClass::Hash {
property: EmailField::Threading.into(),
hash: merge_threads.thread_hash,
}),
thread_index,
);
}
}
}
}
match server.commit_batch(batch).await {
Ok(_) => return Ok(()),
Err(err) if err.is_assertion_failure() && try_count < MAX_RETRIES => {
let backoff = store::rand::rng().random_range(50..=300);
tokio::time::sleep(Duration::from_millis(backoff)).await;
try_count += 1;
}
Err(err) => {
return Err(err.caused_by(trc::location!()));
}
}
}
}
crates/services/src/task_manager/mod.rs
+312 -268
View File
@@ -4,18 +4,24 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use crate::task_manager::bayes::BayesTrainTask;
use crate::task_manager::imip::SendImipTask;
use crate::task_manager::index::SearchIndexTask;
use crate::task_manager::lock::{TaskLock, TaskLockManager};
use crate::task_manager::merge_threads::MergeThreadsTask;
use alarm::SendAlarmTask;
use common::IPC_CHANNEL_BUFFER;
use common::config::server::ServerProtocol;
use common::listener::limiter::ConcurrencyLimiter;
use common::listener::{ServerInstance, TcpAcceptor};
use common::{Inner, KV_LOCK_TASK, Server, core::BuildServer};
use email::message::ingest::MergeThreadIds;
use groupware::calendar::alarm::{CalendarAlarm, CalendarAlarmType};
use std::collections::hash_map::Entry;
use std::future::Future;
use std::time::Duration;
use std::{sync::Arc, time::Instant};
use store::ahash::AHashSet;
use store::rand;
use store::rand::seq::SliceRandom;
use store::write::SearchIndex;
@@ -30,40 +36,52 @@ use store::{
};
use tokio::sync::{mpsc, watch};
use trc::TaskQueueEvent;
use types::blob_hash::{BLOB_HASH_LEN, BlobHash};
use utils::snowflake::SnowflakeIdGenerator;
pub mod alarm;
pub mod bayes;
pub mod fts;
pub mod imip;
pub mod index;
pub mod lock;
pub mod merge_threads;
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Task {
pub struct Task<T> {
pub account_id: u32,
pub document_id: u32,
pub due: u64,
pub action: TaskAction,
pub action: T,
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum TaskAction {
UpdateIndex { index: SearchIndex, is_insert: bool },
BayesTrain { learn_spam: bool },
SendAlarm { alarm: CalendarAlarm },
pub(crate) enum TaskAction {
UpdateIndex(IndexAction),
BayesTrain(bool),
SendAlarm(CalendarAlarm),
SendImip,
MergeThreads(MergeThreadIds<AHashSet<u32>>),
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub(crate) struct IndexAction {
pub index: SearchIndex,
pub is_insert: bool,
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub(crate) struct ImipAction;
const INDEX_EXPIRY: u64 = 60 * 5; // 5 minutes
const BAYES_LOCK_EXPIRY: u64 = 60 * 30; // 30 minutes
const ALARM_EXPIRY: u64 = 60 * 2; // 2 minutes
const QUEUE_REFRESH_INTERVAL: u64 = 60 * 5; // 5 minutes
pub(crate) struct TaskManagerIpc {
tx_fts: mpsc::Sender<Task>,
tx_bayes: mpsc::Sender<Task>,
tx_alarm: mpsc::Sender<Task>,
tx_imip: mpsc::Sender<Task>,
tx_fts: mpsc::Sender<Task<IndexAction>>,
tx_bayes: mpsc::Sender<Task<bool>>,
tx_alarm: mpsc::Sender<Task<CalendarAlarm>>,
tx_imip: mpsc::Sender<Task<ImipAction>>,
tx_threads: mpsc::Sender<Task<MergeThreadIds<AHashSet<u32>>>>,
locked: AHashMap<Vec<u8>, Locked>,
revision: u64,
}
@@ -75,10 +93,12 @@ struct Locked {
pub fn spawn_task_manager(inner: Arc<Inner>) {
// Create three mpsc channels for the different task types
let (tx_index_1, rx_index_1) = mpsc::channel::<Task>(IPC_CHANNEL_BUFFER);
let (tx_index_2, rx_index_2) = mpsc::channel::<Task>(IPC_CHANNEL_BUFFER);
let (tx_index_3, rx_index_3) = mpsc::channel::<Task>(IPC_CHANNEL_BUFFER);
let (tx_index_4, rx_index_4) = mpsc::channel::<Task>(IPC_CHANNEL_BUFFER);
let (tx_index_1, mut rx_index_1) = mpsc::channel::<Task<IndexAction>>(IPC_CHANNEL_BUFFER);
let (tx_index_2, mut rx_index_2) = mpsc::channel::<Task<bool>>(IPC_CHANNEL_BUFFER);
let (tx_index_3, mut rx_index_3) = mpsc::channel::<Task<CalendarAlarm>>(IPC_CHANNEL_BUFFER);
let (tx_index_4, mut rx_index_4) = mpsc::channel::<Task<ImipAction>>(IPC_CHANNEL_BUFFER);
let (tx_index_5, mut rx_index_5) =
mpsc::channel::<Task<MergeThreadIds<AHashSet<u32>>>>(IPC_CHANNEL_BUFFER);
// Create dummy server instance for alarms
let server_instance = Arc::new(ServerInstance {
@@ -91,83 +111,182 @@ pub fn spawn_task_manager(inner: Arc<Inner>) {
span_id_gen: Arc::new(SnowflakeIdGenerator::new()),
});
for mut rx_index in [rx_index_1, rx_index_2, rx_index_3, rx_index_4] {
// Indexing worker
{
let inner = inner.clone();
let server_instance = server_instance.clone();
tokio::spawn(async move {
while let Some(task) = rx_index.recv().await {
while let Some(task) = rx_index_1.recv().await {
let server = inner.build_server();
let batch_size = server.core.jmap.index_batch_size;
let mut batch = Vec::with_capacity(batch_size);
batch.push(task);
while batch.len() < batch_size {
match rx_index_1.try_recv() {
Ok(task) => batch.push(task),
Err(_) => break,
}
}
if batch.len() > 1 {
batch.shuffle(&mut rand::rng());
}
// Lock tasks
let mut locked_batch = Vec::with_capacity(batch.len());
for task in batch {
if server
.try_lock_task(
task.account_id,
task.document_id,
task.lock_key(),
task.lock_expiry(),
)
.await
{
locked_batch.push(task);
}
}
// Dispatch
if !locked_batch.is_empty() {
let success = server.index(&locked_batch).await;
// Remove entries from queue
if success {
delete_tasks(&server, &locked_batch).await;
}
}
}
});
}
// Bayes training worker
{
let inner = inner.clone();
tokio::spawn(async move {
while let Some(task) = rx_index_2.recv().await {
let server = inner.build_server();
// Lock task
if server.try_lock_task(&task).await {
let success = match &task.action {
TaskAction::UpdateIndex { index, is_insert } => {
let todo = "implement";
/*server
.fts_index(task.account_id, task.document_id, hash)
.await*/
true
}
TaskAction::BayesTrain { learn_spam } => {
let todo = "implement";
/*server
.bayes_train(task.account_id, task.document_id, hash, *learn_spam)
.await*/
true
}
TaskAction::SendAlarm { alarm } => {
if server.core.groupware.alarms_enabled {
server
.send_alarm(
task.account_id,
task.document_id,
alarm,
server_instance.clone(),
)
.await
} else {
true
}
}
TaskAction::SendImip => {
if server.core.groupware.itip_enabled {
server
.send_imip(
task.account_id,
task.document_id,
task.due,
server_instance.clone(),
)
.await
} else {
true
}
}
};
if server
.try_lock_task(
task.account_id,
task.document_id,
task.lock_key(),
task.lock_expiry(),
)
.await
{
let success = server
.bayes_train(task.account_id, task.document_id, task.action)
.await;
// Remove entry from queue
if success {
let mut batch = BatchBuilder::new();
batch
.with_account_id(task.account_id)
.with_document(task.document_id);
delete_tasks(&server, &[task]).await;
}
}
}
});
}
for value in task.value_classes() {
batch.clear(value);
}
// Send alarm worker
{
let inner = inner.clone();
let server_instance = server_instance.clone();
tokio::spawn(async move {
while let Some(task) = rx_index_3.recv().await {
let server = inner.build_server();
if let Err(err) = server.core.storage.data.write(batch.build_all()).await {
trc::error!(
err.account_id(task.account_id)
.document_id(task.document_id)
.details("Failed to remove task from queue.")
);
}
// Lock task
if server.core.groupware.alarms_enabled
&& server
.try_lock_task(
task.account_id,
task.document_id,
task.lock_key(),
task.lock_expiry(),
)
.await
{
let success = server
.send_alarm(
task.account_id,
task.document_id,
&task.action,
server_instance.clone(),
)
.await;
if task.remove_lock() {
server.remove_index_lock(&task).await;
}
// Remove entry from queue
if success {
delete_tasks(&server, &[task]).await;
}
}
}
});
}
// Send iMIP worker
{
let inner = inner.clone();
let server_instance = server_instance.clone();
tokio::spawn(async move {
while let Some(task) = rx_index_4.recv().await {
let server = inner.build_server();
// Lock task
if server.core.groupware.itip_enabled
&& server
.try_lock_task(
task.account_id,
task.document_id,
task.lock_key(),
task.lock_expiry(),
)
.await
{
let success = server
.send_imip(
task.account_id,
task.document_id,
task.due,
server_instance.clone(),
)
.await;
// Remove entry from queue
if success {
delete_tasks(&server, &[task]).await;
}
}
}
});
}
// Merge threads worker
{
let inner = inner.clone();
tokio::spawn(async move {
while let Some(task) = rx_index_5.recv().await {
let server = inner.build_server();
// Lock task
if server
.try_lock_task(
task.account_id,
task.document_id,
task.lock_key(),
task.lock_expiry(),
)
.await
{
let success = server.merge_threads(task.account_id, &task.action).await;
// Remove entry from queue
if success {
delete_tasks(&server, &[task]).await;
}
}
}
@@ -180,6 +299,7 @@ pub fn spawn_task_manager(inner: Arc<Inner>) {
tx_bayes: tx_index_2,
tx_alarm: tx_index_3,
tx_imip: tx_index_4,
tx_threads: tx_index_5,
locked: Default::default(),
revision: 0,
};
@@ -196,8 +316,6 @@ pub fn spawn_task_manager(inner: Arc<Inner>) {
pub(crate) trait TaskQueueManager: Sync + Send {
fn process_tasks(&self, ipc: &mut TaskManagerIpc) -> impl Future<Output = Duration> + Send;
fn try_lock_task(&self, event: &Task) -> impl Future<Output = bool> + Send;
fn remove_index_lock(&self, event: &Task) -> impl Future<Output = ()> + Send;
}
impl TaskQueueManager for Server {
@@ -287,24 +405,109 @@ impl TaskQueueManager for Server {
// Dispatch tasks
let roles = &self.core.network.roles;
for event in tasks {
let tx = match &event.action {
TaskAction::UpdateIndex { .. }
match event.action {
TaskAction::UpdateIndex(index)
if roles.fts_indexing.is_enabled_for_hash(&event) =>
{
&ipc.tx_fts
if ipc
.tx_fts
.send(Task {
account_id: event.account_id,
document_id: event.document_id,
due: event.due,
action: index,
})
.await
.is_err()
{
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
TaskAction::BayesTrain { .. }
TaskAction::BayesTrain(learn_spam)
if roles.bayes_training.is_enabled_for_hash(&event) =>
{
&ipc.tx_bayes
if ipc
.tx_bayes
.send(Task {
account_id: event.account_id,
document_id: event.document_id,
due: event.due,
action: learn_spam,
})
.await
.is_err()
{
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
TaskAction::SendAlarm { .. }
TaskAction::SendAlarm(alarm)
if roles.calendar_alerts.is_enabled_for_hash(&event) =>
{
&ipc.tx_alarm
if ipc
.tx_alarm
.send(Task {
account_id: event.account_id,
document_id: event.document_id,
due: event.due,
action: alarm,
})
.await
.is_err()
{
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
TaskAction::SendImip if roles.imip_processing.is_enabled_for_hash(&event) => {
&ipc.tx_imip
if ipc
.tx_imip
.send(Task {
account_id: event.account_id,
document_id: event.document_id,
due: event.due,
action: ImipAction,
})
.await
.is_err()
{
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
TaskAction::MergeThreads(info)
if roles.merge_threads.is_enabled_for_hash(&event) =>
{
if ipc
.tx_threads
.send(Task {
account_id: event.account_id,
document_id: event.document_id,
due: event.due,
action: info,
})
.await
.is_err()
{
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
_ => {
trc::event!(
@@ -315,13 +518,6 @@ impl TaskQueueManager for Server {
continue;
}
};
if tx.send(event).await.is_err() {
trc::event!(
Server(trc::ServerEvent::ThreadError),
Details = "Error sending task.",
CausedBy = trc::location!()
);
}
}
@@ -333,180 +529,28 @@ impl TaskQueueManager for Server {
timestamp.saturating_sub(store::write::now())
}))
}
}
async fn try_lock_task(&self, event: &Task) -> bool {
match self
.in_memory_store()
.try_lock(KV_LOCK_TASK, &event.lock_key(), event.lock_expiry())
.await
{
Ok(result) => {
if !result {
trc::event!(
TaskQueue(TaskQueueEvent::TaskLocked),
AccountId = event.account_id,
DocumentId = event.document_id,
Expires = trc::Value::Timestamp(now() + event.lock_expiry()),
);
}
result
}
Err(err) => {
trc::error!(
err.account_id(event.account_id)
.document_id(event.document_id)
.details("Failed to lock task")
);
async fn delete_tasks<T: TaskLock>(server: &Server, tasks: &[T]) {
let mut batch = BatchBuilder::new();
false
}
for task in tasks {
batch
.with_account_id(task.account_id())
.with_document(task.document_id());
for value in task.value_classes() {
batch.clear(value);
}
}
async fn remove_index_lock(&self, event: &Task) {
let key = event.lock_key();
if let Err(err) = self.in_memory_store().remove_lock(KV_LOCK_TASK, &key).await {
trc::error!(
err.details("Failed to unlock task")
.ctx(trc::Key::Key, key)
.caused_by(trc::location!())
);
if let Err(err) = server.store().write(batch.build_all()).await {
trc::error!(err.details("Failed to remove task(s) from queue."));
}
for task in tasks {
if task.remove_lock() {
server.remove_index_lock(task.lock_key()).await;
}
}
}
impl Task {
fn remove_lock(&self) -> bool {
// Bayes locks are not removed to avoid constant retraining
!matches!(self.action, TaskAction::BayesTrain { .. })
}
fn lock_key(&self) -> Vec<u8> {
match &self.action {
TaskAction::UpdateIndex { index, .. } => {
KeySerializer::new((U32_LEN * 2) + U64_LEN + 2)
.write(0u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.write(index.to_u8())
.finalize()
}
TaskAction::BayesTrain { .. } => KeySerializer::new((U32_LEN * 2) + 1)
.write(1u8)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize(),
TaskAction::SendAlarm { .. } => KeySerializer::new((U32_LEN * 2) + U64_LEN + 1)
.write(2u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize(),
TaskAction::SendImip => KeySerializer::new((U32_LEN * 2) + U64_LEN + 1)
.write(3u8)
.write(self.due)
.write_leb128(self.account_id)
.write_leb128(self.document_id)
.finalize(),
}
}
fn lock_expiry(&self) -> u64 {
match self.action {
TaskAction::UpdateIndex { .. } => INDEX_EXPIRY,
TaskAction::BayesTrain { .. } => BAYES_LOCK_EXPIRY,
TaskAction::SendAlarm { .. } | TaskAction::SendImip => ALARM_EXPIRY,
}
}
fn value_classes(&self) -> impl Iterator<Item = ValueClass> {
[
Some(ValueClass::TaskQueue(match &self.action {
TaskAction::UpdateIndex { index, is_insert } => TaskQueueClass::UpdateIndex {
due: self.due,
index: *index,
is_insert: *is_insert,
},
TaskAction::BayesTrain { learn_spam } => TaskQueueClass::BayesTrain {
due: self.due,
learn_spam: *learn_spam,
},
TaskAction::SendAlarm { alarm } => TaskQueueClass::SendAlarm {
event_id: alarm.event_id,
alarm_id: alarm.alarm_id,
due: self.due,
is_email_alert: matches!(alarm.typ, CalendarAlarmType::Email { .. }),
},
TaskAction::SendImip => TaskQueueClass::SendImip {
due: self.due,
is_payload: false,
},
})),
(matches!(self.action, TaskAction::SendImip)).then_some(ValueClass::TaskQueue(
TaskQueueClass::SendImip {
due: self.due,
is_payload: true,
},
)),
]
.into_iter()
.flatten()
}
fn deserialize(key: &[u8], value: &[u8]) -> trc::Result<Self> {
Ok(Task {
due: key.deserialize_be_u64(0)?,
account_id: key.deserialize_be_u32(U64_LEN)?,
document_id: key.deserialize_be_u32(U64_LEN + U32_LEN + 1)?,
action: match key.get(U64_LEN + U32_LEN) {
Some(v @ (7 | 8)) => TaskAction::UpdateIndex {
index: key
.last()
.copied()
.and_then(SearchIndex::try_from_u8)
.ok_or_else(|| trc::Error::corrupted_key(key, None, trc::location!()))?,
is_insert: *v == 7,
},
Some(v @ (1 | 2)) => TaskAction::BayesTrain {
learn_spam: *v == 1,
},
Some(3) => TaskAction::SendAlarm {
alarm: CalendarAlarm {
event_id: key.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + 1)?,
alarm_id: key
.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + U16_LEN + 1)?,
alarm_time: 0,
typ: CalendarAlarmType::Email {
event_start: value.deserialize_be_u64(0)? as i64,
event_end: value.deserialize_be_u64(U64_LEN)? as i64,
event_start_tz: value.deserialize_be_u16(U64_LEN * 2)?,
event_end_tz: value.deserialize_be_u16((U64_LEN * 2) + U16_LEN)?,
},
},
},
Some(6) => {
let recurrence_id = value.deserialize_be_u64(0)? as i64;
TaskAction::SendAlarm {
alarm: CalendarAlarm {
event_id: key.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + 1)?,
alarm_id: key
.deserialize_be_u16(U64_LEN + U32_LEN + U32_LEN + U16_LEN + 1)?,
alarm_time: 0,
typ: CalendarAlarmType::Display {
recurrence_id: if recurrence_id != 0 {
Some(recurrence_id)
} else {
None
},
},
},
}
}
Some(4) => TaskAction::SendImip,
_ => return Err(trc::Error::corrupted_key(key, None, trc::location!())),
},
})
}
}
crates/store/src/backend/foundationdb/read.rs
+2
View File
@@ -47,6 +47,8 @@ impl FdbStore {
let begin = params.begin.serialize(WITH_SUBSPACE);
let end = params.end.serialize(WITH_SUBSPACE);
let todo = "fix fdb range scan to support chunked reads";
if !params.first {
let mut last_key = vec![];
crates/store/src/backend/mysql/main.rs
-1
View File
@@ -104,7 +104,6 @@ impl MysqlStore {
SUBSPACE_LOGS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_INDEX,
] {
let table = char::from(table);
conn.query_drop(format!(
crates/store/src/backend/postgres/main.rs
-1
View File
@@ -92,7 +92,6 @@ impl PostgresStore {
SUBSPACE_BLOBS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_INDEX,
] {
let table = char::from(table);
conn.execute(
crates/store/src/backend/rocksdb/main.rs
-1
View File
@@ -74,7 +74,6 @@ impl RocksDbStore {
SUBSPACE_BLOBS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_INDEX,
] {
let cf_opts = Options::default();
cfs.push(ColumnFamilyDescriptor::new(
crates/store/src/backend/sqlite/main.rs
-1
View File
@@ -105,7 +105,6 @@ impl SqliteStore {
SUBSPACE_BLOBS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_INDEX,
] {
let table = char::from(table);
conn.execute(
crates/store/src/dispatch/search.rs
+28 -44
View File
@@ -7,51 +7,44 @@
use super::DocumentSet;
use crate::{
SearchStore,
search::{IndexDocument, SearchComparator, SearchFilter},
backend::elastic::query,
search::{IndexDocument, SearchComparator, SearchDocumentId, SearchFilter, SearchQuery},
};
use trc::AddContext;
use types::collection::Collection;
impl SearchStore {
pub async fn index(&self, document: IndexDocument) -> trc::Result<()> {
match self {
pub async fn query<R: SearchDocumentId>(&self, query: SearchQuery) -> trc::Result<Vec<R>> {
todo!()
/*match self {
SearchStore::Store(store) => {
store
.index_query(account_id, collection, filters, comparators)
.await
}
#[cfg(feature = "elastic")]
SearchStore::ElasticSearch(store) => {
store
.index_query(account_id, collection, filters, comparators)
.await
}
}
.caused_by(trc::location!())*/
}
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
todo!()
/*match self {
SearchStore::Store(store) => store.index_insert(document).await,
#[cfg(feature = "elastic")]
SearchStore::ElasticSearch(store) => store.index_insert(document).await,
}
.caused_by(trc::location!())
.caused_by(trc::location!())*/
}
pub async fn query(
&self,
account_id: u32,
collection: Collection,
filters: Vec<SearchFilter>,
comparators: Vec<SearchComparator>,
) -> trc::Result<Vec<u32>> {
match self {
SearchStore::Store(store) => {
store
.index_query(account_id, collection, filters, comparators)
.await
}
#[cfg(feature = "elastic")]
SearchStore::ElasticSearch(store) => {
store
.index_query(account_id, collection, filters, comparators)
.await
}
}
.caused_by(trc::location!())
}
pub async fn remove(
&self,
account_id: u32,
collection: Collection,
document_ids: &impl DocumentSet,
) -> trc::Result<()> {
match self {
pub async fn unindex(&self, query: SearchQuery) -> trc::Result<()> {
todo!()
/*match self {
SearchStore::Store(store) => {
store
.index_remove(account_id, collection, document_ids)
@@ -64,15 +57,6 @@ impl SearchStore {
.await
}
}
.caused_by(trc::location!())
}
pub async fn remove_all(&self, account_id: u32) -> trc::Result<()> {
match self {
SearchStore::Store(store) => store.index_remove_all(account_id).await,
#[cfg(feature = "elastic")]
SearchStore::ElasticSearch(store) => store.index_remove_all(account_id).await,
}
.caused_by(trc::location!())
.caused_by(trc::location!())*/
}
}
crates/store/src/dispatch/store.rs
-2
View File
@@ -488,7 +488,6 @@ impl Store {
SUBSPACE_REPORT_IN,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_INDEX,
] {
self.delete_range(
AnyKey {
@@ -667,7 +666,6 @@ impl Store {
(SUBSPACE_INDEXES, false),
(SUBSPACE_TELEMETRY_SPAN, true),
(SUBSPACE_TELEMETRY_METRIC, true),
(SUBSPACE_TELEMETRY_INDEX, true),
] {
let from_key = crate::write::AnyKey {
subspace,
crates/store/src/lib.rs
+1 -2
View File
@@ -102,7 +102,6 @@ pub const SUBSPACE_QUOTA: u8 = b'u';
pub const SUBSPACE_REPORT_OUT: u8 = b'h';
pub const SUBSPACE_REPORT_IN: u8 = b'r';
pub const SUBSPACE_TELEMETRY_SPAN: u8 = b'o';
pub const SUBSPACE_TELEMETRY_INDEX: u8 = b'w';
pub const SUBSPACE_TELEMETRY_METRIC: u8 = b'x';
pub const SUBSPACE_RESERVED_2: u8 = b'z';
@@ -111,7 +110,7 @@ pub const SUBSPACE_BITMAP_ID: u8 = b'b';
pub const SUBSPACE_BITMAP_TAG: u8 = b'c';
pub const SUBSPACE_BITMAP_TEXT: u8 = b'v';
pub const SUBSPACE_FTS_INDEX: u8 = b'g';
pub const SUBSPACE_TELEMETRY_INDEX: u8 = b'w';
*/
#[derive(Clone)]
crates/store/src/search/index.rs
+3 -3
View File
@@ -29,7 +29,7 @@ use crate::{
pub const TERM_INDEX_VERSION: u8 = 1;
impl Store {
pub async fn index_insert(&self, document: IndexDocument) -> trc::Result<()> {
pub(crate) async fn index_insert(&self, document: IndexDocument) -> trc::Result<()> {
/*let mut detect = LanguageDetector::new();
let mut tokens: AHashMap<BitmapHash, Postings> = AHashMap::new();
let mut parts = Vec::new();
@@ -142,7 +142,7 @@ impl Store {
Ok(())
}
pub async fn index_remove(
pub(crate) async fn index_remove(
&self,
account_id: u32,
collection: Collection,
@@ -244,7 +244,7 @@ impl Store {
Ok(())
}
pub async fn index_remove_all(&self, _: u32) -> trc::Result<()> {
pub(crate) async fn index_remove_all(&self, _: u32) -> trc::Result<()> {
// No-op
// Term indexes are stored in the same key range as the document
crates/store/src/search/mod.rs
+130 -14
View File
@@ -13,6 +13,8 @@ use nlp::language::Language;
use roaring::RoaringBitmap;
use std::{borrow::Cow, collections::hash_map::Entry};
use crate::write::SearchIndex;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchOperator {
LowerThan,
@@ -26,6 +28,9 @@ pub enum SearchOperator {
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum SearchField {
AccountId,
DocumentId,
Id,
Email(EmailSearchField),
Calendar(CalendarSearchField),
Contact(ContactSearchField),
@@ -84,10 +89,25 @@ pub enum FileSearchField {
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SearchValue {
Text { value: String, language: Language },
Number(i64),
Int(i64),
Uint(u64),
Boolean(bool),
}
pub trait SearchDocumentId: Sized {
fn from_u32(id: u32) -> Self;
fn from_u64(id: u64) -> Self;
fn field(&self) -> SearchField;
}
#[derive(Debug)]
pub struct SearchQuery {
index: SearchIndex,
filters: Vec<SearchFilter>,
comparators: Vec<SearchComparator>,
mask: RoaringBitmap,
}
#[derive(Debug)]
pub enum SearchFilter {
Operator {
@@ -111,8 +131,6 @@ pub enum SearchComparator {
#[derive(Debug)]
pub struct IndexDocument {
pub(crate) account_id: u32,
pub(crate) document_id: u32,
pub(crate) fields: AHashMap<SearchField, SearchValue>,
pub(crate) default_language: Language,
}
@@ -275,18 +293,25 @@ impl IndexDocument {
Self {
fields: Default::default(),
default_language,
account_id: 0,
document_id: 0,
}
}
pub fn with_account_id(mut self, account_id: u32) -> Self {
self.account_id = account_id;
self.fields
.insert(SearchField::AccountId, SearchValue::Uint(account_id as u64));
self
}
pub fn with_document_id(mut self, document_id: u32) -> Self {
self.document_id = document_id;
self.fields.insert(
SearchField::DocumentId,
SearchValue::Uint(document_id as u64),
);
self
}
pub fn with_id(mut self, id: u64) -> Self {
self.fields.insert(SearchField::Id, SearchValue::Uint(id));
self
}
@@ -316,9 +341,14 @@ impl IndexDocument {
.insert(field.into(), SearchValue::Boolean(value));
}
pub fn index_number<N: Into<i64>>(&mut self, field: impl Into<SearchField>, value: N) {
pub fn index_integer<N: Into<i64>>(&mut self, field: impl Into<SearchField>, value: N) {
self.fields
.insert(field.into(), SearchValue::Number(value.into()));
.insert(field.into(), SearchValue::Int(value.into()));
}
pub fn index_unsigned<N: Into<u64>>(&mut self, field: impl Into<SearchField>, value: N) {
self.fields
.insert(field.into(), SearchValue::Uint(value.into()));
}
pub fn has_field(&self, field: &SearchField) -> bool {
@@ -326,6 +356,64 @@ impl IndexDocument {
}
}
impl SearchQuery {
pub fn new(index: SearchIndex) -> Self {
Self {
index,
filters: Vec::new(),
comparators: Vec::new(),
mask: RoaringBitmap::new(),
}
}
pub fn with_filters(mut self, filters: Vec<SearchFilter>) -> Self {
if self.filters.is_empty() {
self.filters = filters;
} else {
self.filters.extend(filters);
}
self
}
pub fn with_comparators(mut self, comparators: Vec<SearchComparator>) -> Self {
if self.comparators.is_empty() {
self.comparators = comparators;
} else {
self.comparators.extend(comparators);
}
self
}
pub fn with_filter(mut self, filter: SearchFilter) -> Self {
self.filters.push(filter);
self
}
pub fn with_comparator(mut self, comparator: SearchComparator) -> Self {
self.comparators.push(comparator);
self
}
pub fn with_mask(mut self, mask: RoaringBitmap) -> Self {
self.mask = mask;
self
}
pub fn with_account_id(mut self, account_id: u32) -> Self {
self.filters.push(SearchFilter::cond(
SearchField::AccountId,
SearchOperator::Equal,
SearchValue::Uint(account_id as u64),
));
self
}
pub fn execute(&self) -> RoaringBitmap {
let todo = "implement search execution logic";
todo!()
}
}
impl From<EmailSearchField> for SearchField {
fn from(field: EmailSearchField) -> Self {
SearchField::Email(field)
@@ -352,31 +440,31 @@ impl From<FileSearchField> for SearchField {
impl From<u64> for SearchValue {
fn from(value: u64) -> Self {
SearchValue::Number(value as i64)
SearchValue::Uint(value)
}
}
impl From<i64> for SearchValue {
fn from(value: i64) -> Self {
SearchValue::Number(value)
SearchValue::Int(value)
}
}
impl From<u32> for SearchValue {
fn from(value: u32) -> Self {
SearchValue::Number(value as i64)
SearchValue::Uint(value as u64)
}
}
impl From<i32> for SearchValue {
fn from(value: i32) -> Self {
SearchValue::Number(value as i64)
SearchValue::Int(value as i64)
}
}
impl From<usize> for SearchValue {
fn from(value: usize) -> Self {
SearchValue::Number(value as i64)
SearchValue::Uint(value as u64)
}
}
@@ -394,3 +482,31 @@ impl From<String> for SearchValue {
}
}
}
impl SearchDocumentId for u32 {
fn from_u32(id: u32) -> Self {
id
}
fn from_u64(id: u64) -> Self {
id as u32
}
fn field(&self) -> SearchField {
SearchField::DocumentId
}
}
impl SearchDocumentId for u64 {
fn from_u32(id: u32) -> Self {
id as u64
}
fn from_u64(id: u64) -> Self {
id
}
fn field(&self) -> SearchField {
SearchField::Id
}
}
crates/store/src/search/query.rs
+1 -1
View File
@@ -44,7 +44,7 @@ enum FtsTokenized {
}*/
impl Store {
pub async fn index_query(
pub(crate) async fn index_query(
&self,
account_id: u32,
collection: Collection,
crates/store/src/write/key.rs
+24 -11
View File
@@ -13,9 +13,8 @@ use crate::{
SUBSPACE_BLOB_RESERVE, SUBSPACE_COUNTER, SUBSPACE_DIRECTORY, SUBSPACE_IN_MEMORY_COUNTER,
SUBSPACE_IN_MEMORY_VALUE, SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_PROPERTY,
SUBSPACE_QUEUE_EVENT, SUBSPACE_QUEUE_MESSAGE, SUBSPACE_QUOTA, SUBSPACE_REPORT_IN,
SUBSPACE_REPORT_OUT, SUBSPACE_SETTINGS, SUBSPACE_TASK_QUEUE, SUBSPACE_TELEMETRY_INDEX,
SUBSPACE_TELEMETRY_METRIC, SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN, U64_LEN, ValueKey,
WITH_SUBSPACE,
SUBSPACE_REPORT_OUT, SUBSPACE_SETTINGS, SUBSPACE_TASK_QUEUE, SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN, U64_LEN, ValueKey, WITH_SUBSPACE,
write::{IndexPropertyClass, SearchIndex},
};
use std::convert::TryInto;
@@ -322,6 +321,11 @@ impl ValueClass {
.write(*due)
}
}
TaskQueueClass::MergeThreads { due } => serializer
.write(*due)
.write(account_id)
.write(9u8)
.write(document_id),
},
ValueClass::Blob(op) => match op {
BlobOp::Reserve { hash, until } => serializer
@@ -418,9 +422,6 @@ impl ValueClass {
},
ValueClass::Telemetry(telemetry) => match telemetry {
TelemetryClass::Span { span_id } => serializer.write(*span_id),
TelemetryClass::Index { span_id, value } => {
serializer.write(value.as_slice()).write(*span_id)
}
TelemetryClass::Metric {
timestamp,
metric_id,
@@ -512,7 +513,9 @@ impl ValueClass {
ValueClass::TaskQueue(e) => match e {
TaskQueueClass::UpdateIndex { .. } => (U64_LEN * 2) + 2,
TaskQueueClass::BayesTrain { .. } => (U64_LEN * 2) + 1,
TaskQueueClass::SendAlarm { .. } => U64_LEN + (U32_LEN * 3) + 1,
TaskQueueClass::SendAlarm { .. } | TaskQueueClass::MergeThreads { .. } => {
U64_LEN + (U32_LEN * 3) + 1
}
TaskQueueClass::SendImip { is_payload, .. } => {
if *is_payload {
(U64_LEN * 2) + (U32_LEN * 2) + 1
@@ -535,7 +538,6 @@ impl ValueClass {
ValueClass::Report(_) => U64_LEN * 2 + 1,
ValueClass::Telemetry(telemetry) => match telemetry {
TelemetryClass::Span { .. } => U64_LEN + 1,
TelemetryClass::Index { value, .. } => U64_LEN + value.len() + 1,
TelemetryClass::Metric { .. } => U64_LEN * 2 + 1,
},
ValueClass::DocumentId => U32_LEN + 1,
@@ -584,7 +586,6 @@ impl ValueClass {
ValueClass::Report(_) => SUBSPACE_REPORT_IN,
ValueClass::Telemetry(telemetry) => match telemetry {
TelemetryClass::Span { .. } => SUBSPACE_TELEMETRY_SPAN,
TelemetryClass::Index { .. } => SUBSPACE_TELEMETRY_INDEX,
TelemetryClass::Metric { .. } => SUBSPACE_TELEMETRY_METRIC,
},
ValueClass::DocumentId | ValueClass::ChangeId => SUBSPACE_COUNTER,
@@ -666,7 +667,8 @@ impl SearchIndex {
SearchIndex::Calendar => 1,
SearchIndex::Contacts => 2,
SearchIndex::File => 3,
SearchIndex::DeliveryHistory => 4,
SearchIndex::TracingSpan => 4,
SearchIndex::InMemory => unreachable!(),
}
}
@@ -676,7 +678,18 @@ impl SearchIndex {
1 => Some(SearchIndex::Calendar),
2 => Some(SearchIndex::Contacts),
3 => Some(SearchIndex::File),
4 => Some(SearchIndex::DeliveryHistory),
4 => Some(SearchIndex::TracingSpan),
_ => None,
}
}
pub fn try_from_str(value: &str) -> Option<Self> {
match value {
"email" => Some(SearchIndex::Email),
"calendar" => Some(SearchIndex::Calendar),
"contacts" => Some(SearchIndex::Contacts),
"file" => Some(SearchIndex::File),
"tracing" => Some(SearchIndex::TracingSpan),
_ => None,
}
}
crates/store/src/write/mod.rs
+5 -5
View File
@@ -209,6 +209,9 @@ pub enum TaskQueueClass {
due: u64,
is_payload: bool,
},
MergeThreads {
due: u64,
},
}
#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)]
@@ -217,7 +220,8 @@ pub enum SearchIndex {
Calendar,
Contacts,
File,
DeliveryHistory,
TracingSpan,
InMemory,
}
#[derive(Debug, PartialEq, Clone, Eq, Hash)]
@@ -272,10 +276,6 @@ pub enum TelemetryClass {
metric_id: u64,
node_id: u64,
},
Index {
span_id: u64,
value: Vec<u8>,
},
}
#[derive(Debug, PartialEq, Clone, Eq, Hash)]