filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 2 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

115
crates/services/src/task_manager/bayes.rs
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
}
}
}
}

451
crates/services/src/task_manager/fts.rs
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(())
}
}

297
crates/services/src/task_manager/index.rs Normal file
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(())
}
}

325
crates/services/src/task_manager/lock.rs Normal file
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!())),
},
})
}
}

192
crates/services/src/task_manager/merge_threads.rs Normal file
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!()));
}
}
}
}

580
crates/services/src/task_manager/mod.rs
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!())),
},
})
}
}