605 lines
20 KiB
Rust
605 lines
20 KiB
Rust
/*
|
|
* 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,
|
|
telemetry::tracers::store::{TracingStore, build_span_document},
|
|
};
|
|
use directory::{QueryParams, Type, backend::internal::manage::ManageDirectory};
|
|
use email::message::metadata::MessageMetadata;
|
|
use groupware::{cache::GroupwareCache, calendar::CalendarEvent, contact::ContactCard};
|
|
use std::cmp::Ordering;
|
|
use store::{
|
|
IterateParams, SerializeInfallible, U32_LEN, ValueKey,
|
|
ahash::AHashMap,
|
|
roaring::RoaringBitmap,
|
|
search::{IndexDocument, SearchField, SearchFilter, SearchQuery},
|
|
write::{
|
|
BatchBuilder, BlobOp, SearchIndex, TaskQueueClass, ValueClass, key::DeserializeBigEndian,
|
|
now,
|
|
},
|
|
};
|
|
use trc::{AddContext, TaskQueueEvent};
|
|
use types::{
|
|
blob_hash::{BLOB_HASH_LEN, BlobHash},
|
|
collection::{Collection, SyncCollection},
|
|
field::EmailField,
|
|
};
|
|
|
|
pub(crate) trait SearchIndexTask: Sync + Send {
|
|
fn index(
|
|
&self,
|
|
tasks: &[Task<IndexAction>],
|
|
) -> impl Future<Output = Vec<IndexTaskResult>> + 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;
|
|
}
|
|
|
|
const NUM_INDEXES: usize = 5;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
enum TaskType {
|
|
Insert,
|
|
Delete,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
enum TaskStatus {
|
|
Success,
|
|
Failed,
|
|
Ignored,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub(crate) struct IndexTaskResult {
|
|
index: SearchIndex,
|
|
task_type: TaskType,
|
|
status: TaskStatus,
|
|
}
|
|
|
|
impl SearchIndexTask for Server {
|
|
async fn index(&self, tasks: &[Task<IndexAction>]) -> Vec<IndexTaskResult> {
|
|
let mut results: Vec<IndexTaskResult> = Vec::with_capacity(tasks.len());
|
|
let mut batch = BatchBuilder::new();
|
|
let mut document_insertions: [Vec<IndexDocument>; NUM_INDEXES] =
|
|
std::array::from_fn(|_| Vec::new());
|
|
let mut document_deletions: [AHashMap<u32, Vec<u32>>; NUM_INDEXES] =
|
|
std::array::from_fn(|_| AHashMap::new());
|
|
|
|
for task in tasks {
|
|
if task.action.is_insert {
|
|
let (idx, document) = match task.action.index {
|
|
SearchIndex::Email => (
|
|
0,
|
|
build_email_document(self, task.account_id, task.document_id).await,
|
|
),
|
|
SearchIndex::Calendar => (
|
|
1,
|
|
build_calendar_document(self, task.account_id, task.document_id).await,
|
|
),
|
|
SearchIndex::Contacts => (
|
|
2,
|
|
build_contact_document(self, task.account_id, task.document_id).await,
|
|
),
|
|
SearchIndex::File => {
|
|
// File indexing not implemented yet
|
|
continue;
|
|
}
|
|
SearchIndex::TracingSpan => (
|
|
4,
|
|
build_tracing_span_document(self, task.account_id, task.document_id).await,
|
|
),
|
|
SearchIndex::InMemory => unreachable!(),
|
|
};
|
|
|
|
let result = match document {
|
|
Ok(Some(doc)) if !doc.is_empty() => {
|
|
document_insertions[idx].push(doc);
|
|
TaskStatus::Success
|
|
}
|
|
Err(err) => {
|
|
trc::error!(
|
|
err.account_id(task.account_id)
|
|
.document_id(task.document_id)
|
|
.caused_by(trc::location!())
|
|
.ctx(trc::Key::Collection, task.action.index.name())
|
|
.details("Failed to build document for indexing")
|
|
);
|
|
TaskStatus::Failed
|
|
}
|
|
_ => {
|
|
trc::event!(
|
|
TaskQueue(TaskQueueEvent::TaskIgnored),
|
|
Collection = task.action.index.name(),
|
|
Reason = "Nothing to index",
|
|
AccountId = task.account_id,
|
|
DocumentId = task.document_id,
|
|
);
|
|
TaskStatus::Ignored
|
|
}
|
|
};
|
|
|
|
results.push(IndexTaskResult {
|
|
task_type: TaskType::Insert,
|
|
index: task.action.index,
|
|
status: result,
|
|
});
|
|
} else {
|
|
let idx = match task.action.index {
|
|
SearchIndex::Email => {
|
|
if let Err(err) = delete_email_metadata(
|
|
self,
|
|
&mut batch,
|
|
task.account_id,
|
|
task.document_id,
|
|
)
|
|
.await
|
|
{
|
|
trc::error!(
|
|
err.account_id(task.account_id)
|
|
.document_id(task.document_id)
|
|
.caused_by(trc::location!())
|
|
.details("Failed to delete email metadata from index")
|
|
);
|
|
results.push(IndexTaskResult {
|
|
task_type: TaskType::Delete,
|
|
index: task.action.index,
|
|
status: TaskStatus::Failed,
|
|
});
|
|
continue;
|
|
}
|
|
0
|
|
}
|
|
SearchIndex::Calendar => 1,
|
|
SearchIndex::Contacts => 2,
|
|
SearchIndex::File => 3,
|
|
SearchIndex::TracingSpan | SearchIndex::InMemory => unreachable!(),
|
|
};
|
|
|
|
document_deletions[idx]
|
|
.entry(task.account_id)
|
|
.or_default()
|
|
.push(task.document_id);
|
|
|
|
results.push(IndexTaskResult {
|
|
task_type: TaskType::Delete,
|
|
index: task.action.index,
|
|
status: TaskStatus::Success,
|
|
});
|
|
}
|
|
}
|
|
|
|
// Commit deletion batch to data store
|
|
if !batch.is_empty() {
|
|
if let Err(err) = self.store().write(batch.build_all()).await {
|
|
trc::error!(
|
|
err.caused_by(trc::location!())
|
|
.details("Failed to commit index deletions to data store")
|
|
);
|
|
}
|
|
for r in results.iter_mut() {
|
|
if r.task_type == TaskType::Delete
|
|
&& r.status == TaskStatus::Success
|
|
&& r.index == SearchIndex::Email
|
|
{
|
|
r.status = TaskStatus::Failed;
|
|
}
|
|
}
|
|
return results;
|
|
}
|
|
|
|
// Index documents
|
|
for (documents, index) in document_insertions.into_iter().zip([
|
|
SearchIndex::Email,
|
|
SearchIndex::Calendar,
|
|
SearchIndex::Contacts,
|
|
SearchIndex::File,
|
|
SearchIndex::TracingSpan,
|
|
]) {
|
|
if !documents.is_empty()
|
|
&& let Err(err) = self.search_store().index(index, documents).await
|
|
{
|
|
trc::error!(
|
|
err.caused_by(trc::location!())
|
|
.details("Failed to index documents")
|
|
.ctx(trc::Key::Collection, index.name())
|
|
);
|
|
for r in results.iter_mut() {
|
|
if r.task_type == TaskType::Delete && r.status == TaskStatus::Success {
|
|
r.status = TaskStatus::Failed;
|
|
}
|
|
}
|
|
return results;
|
|
}
|
|
}
|
|
|
|
// Delete documents
|
|
for (accounts, index) in document_deletions.into_iter().zip([
|
|
SearchIndex::Email,
|
|
SearchIndex::Calendar,
|
|
SearchIndex::Contacts,
|
|
]) {
|
|
let multi_account = match accounts.len().cmp(&1) {
|
|
Ordering::Greater => true,
|
|
Ordering::Equal => false,
|
|
Ordering::Less => continue,
|
|
};
|
|
|
|
let mut query = SearchQuery::new(index);
|
|
if multi_account {
|
|
query.add_filter(SearchFilter::Or);
|
|
}
|
|
|
|
for (account_id, document_ids) in accounts {
|
|
let multi_document = document_ids.len() > 1;
|
|
query
|
|
.add_filter(SearchFilter::And)
|
|
.add_filter(SearchFilter::eq(SearchField::AccountId, account_id));
|
|
|
|
if multi_document {
|
|
query.add_filter(SearchFilter::Or);
|
|
}
|
|
|
|
for document_id in document_ids {
|
|
query.add_filter(SearchFilter::eq(SearchField::DocumentId, document_id));
|
|
}
|
|
|
|
if multi_document {
|
|
query.add_filter(SearchFilter::End);
|
|
}
|
|
query.add_filter(SearchFilter::End);
|
|
}
|
|
|
|
if multi_account {
|
|
query.add_filter(SearchFilter::End);
|
|
}
|
|
|
|
if let Err(err) = self.search_store().unindex(query).await {
|
|
trc::error!(
|
|
err.caused_by(trc::location!())
|
|
.details("Failed to delete documents from index")
|
|
.ctx(trc::Key::Collection, index.name())
|
|
);
|
|
for r in results.iter_mut() {
|
|
if r.task_type == TaskType::Delete && r.status == TaskStatus::Success {
|
|
r.status = TaskStatus::Failed;
|
|
}
|
|
}
|
|
return results;
|
|
}
|
|
}
|
|
|
|
results
|
|
}
|
|
}
|
|
|
|
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(())
|
|
}
|
|
}
|
|
|
|
async fn build_email_document(
|
|
server: &Server,
|
|
account_id: u32,
|
|
document_id: u32,
|
|
) -> trc::Result<Option<IndexDocument>> {
|
|
let Some(index_fields) = server.core.jmap.index_fields.get(&SearchIndex::Email) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
match server
|
|
.archive_by_property(
|
|
account_id,
|
|
Collection::Email,
|
|
document_id,
|
|
EmailField::Metadata.into(),
|
|
)
|
|
.await?
|
|
{
|
|
Some(metadata_) => {
|
|
let metadata = metadata_
|
|
.unarchive::<MessageMetadata>()
|
|
.caused_by(trc::location!())?;
|
|
|
|
let raw_message = server
|
|
.blob_store()
|
|
.get_blob(metadata.blob_hash.0.as_slice(), 0..usize::MAX)
|
|
.await
|
|
.caused_by(trc::location!())?
|
|
.ok_or_else(|| {
|
|
trc::StoreEvent::NotFound
|
|
.into_err()
|
|
.details("Blob not found")
|
|
})?;
|
|
|
|
Ok(Some(metadata.index_document(
|
|
&raw_message,
|
|
index_fields,
|
|
server.core.jmap.index_all_headers,
|
|
)))
|
|
}
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
async fn build_calendar_document(
|
|
server: &Server,
|
|
account_id: u32,
|
|
document_id: u32,
|
|
) -> trc::Result<Option<IndexDocument>> {
|
|
let Some(index_fields) = server.core.jmap.index_fields.get(&SearchIndex::Calendar) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
match server
|
|
.archive(account_id, Collection::CalendarEvent, document_id)
|
|
.await?
|
|
{
|
|
Some(metadata_) => Ok(Some(
|
|
metadata_
|
|
.unarchive::<CalendarEvent>()
|
|
.caused_by(trc::location!())?
|
|
.index_document(index_fields),
|
|
)),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
async fn build_contact_document(
|
|
server: &Server,
|
|
account_id: u32,
|
|
document_id: u32,
|
|
) -> trc::Result<Option<IndexDocument>> {
|
|
let Some(index_fields) = server.core.jmap.index_fields.get(&SearchIndex::Contacts) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
match server
|
|
.archive(account_id, Collection::ContactCard, document_id)
|
|
.await?
|
|
{
|
|
Some(metadata_) => Ok(Some(
|
|
metadata_
|
|
.unarchive::<ContactCard>()
|
|
.caused_by(trc::location!())?
|
|
.index_document(index_fields),
|
|
)),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
async fn build_tracing_span_document(
|
|
server: &Server,
|
|
account_id: u32,
|
|
document_id: u32,
|
|
) -> trc::Result<Option<IndexDocument>> {
|
|
let Some(index_fields) = server.core.jmap.index_fields.get(&SearchIndex::TracingSpan) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
let span_id = ((account_id as u64) << 32) | document_id as u64;
|
|
let span = server.store().get_span(span_id).await?;
|
|
|
|
if !span.is_empty() {
|
|
Ok(Some(build_span_document(span_id, span, index_fields)))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
async fn delete_email_metadata(
|
|
server: &Server,
|
|
batch: &mut BatchBuilder,
|
|
account_id: u32,
|
|
document_id: u32,
|
|
) -> trc::Result<()> {
|
|
match server
|
|
.archive_by_property(
|
|
account_id,
|
|
Collection::Email,
|
|
document_id,
|
|
EmailField::Metadata.into(),
|
|
)
|
|
.await?
|
|
{
|
|
Some(metadata) => {
|
|
let tenant_id = server
|
|
.core
|
|
.storage
|
|
.directory
|
|
.query(QueryParams::id(account_id).with_return_member_of(false))
|
|
.await
|
|
.unwrap_or_default()
|
|
.and_then(|p| p.tenant());
|
|
|
|
batch
|
|
.with_account_id(account_id)
|
|
.with_collection(Collection::Email)
|
|
.with_document(document_id);
|
|
metadata
|
|
.unarchive::<MessageMetadata>()
|
|
.caused_by(trc::location!())?
|
|
.unindex(batch, account_id, tenant_id);
|
|
}
|
|
None => {
|
|
trc::event!(
|
|
TaskQueue(TaskQueueEvent::MetadataNotFound),
|
|
Details = "E-mail metadata not found",
|
|
AccountId = account_id,
|
|
DocumentId = document_id,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
impl IndexTaskResult {
|
|
pub fn is_done(&self) -> bool {
|
|
self.status != TaskStatus::Failed
|
|
}
|
|
}
|