JMAP Registry API implementation - part 2

This commit is contained in:
mdecimus
2026-02-24 17:19:44 +01:00
parent ec148c5550
commit 2fea9a46b0
61 changed files with 1746 additions and 1992 deletions

View File

@@ -10,13 +10,16 @@ use super::{
};
use crate::{
SerializeInfallible, U32_LEN,
write::{LogCollection, MergeFnc, MergeOperation, Params, SetFnc, SetOperation},
write::{
LogCollection, MergeFnc, MergeOperation, Params, SetFnc, SetOperation, TaskQueueClass,
},
};
use registry::{pickle::Pickle, schema::structs::Task, types::EnumImpl};
use types::{
collection::{Collection, SyncCollection, VanishedCollection},
field::FieldType,
};
use utils::map::vec_map::VecMap;
use utils::{map::vec_map::VecMap, snowflake::SnowflakeIdGenerator};
impl BatchBuilder {
pub fn new() -> Self {
@@ -390,6 +393,13 @@ impl BatchBuilder {
}
pub fn any_op(&mut self, op: Operation) -> &mut Self {
if let Operation::Value { class, op } = &op {
self.batch_size += class.serialized_size();
if let ValueOp::Set(value) = op {
self.batch_size += value.len();
}
}
self.ops.push(op);
self.batch_ops += 1;
self
@@ -459,6 +469,20 @@ impl BatchBuilder {
pub fn is_empty(&self) -> bool {
self.batch_ops == 0
}
pub fn schedule_task(&mut self, task: Task) -> &mut Self {
let due = task.due_timestamp();
let class = task.object_type().to_id();
let task = task.to_pickled_vec();
let id = SnowflakeIdGenerator::from_sequence_id(xxhash_rust::xxh3::xxh3_64(&task))
.unwrap_or_default();
self.set(ValueClass::TaskQueue(TaskQueueClass::Task { id }), task)
.set(
ValueClass::TaskQueue(TaskQueueClass::Due { id, due }),
class.serialize(),
)
}
}
pub struct CommitPointIterator {

View File

@@ -4,13 +4,17 @@
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*/
use std::time::Instant;
use super::{BlobOp, Operation, ValueClass, ValueOp, key::DeserializeBigEndian, now};
use crate::{
BlobStore, IterateParams, Store, U32_LEN, U64_LEN, ValueKey,
write::{BatchBuilder, BlobLink},
BlobStore, Deserialize, IterateParams, SerializeInfallible, Store, U16_LEN, U32_LEN, U64_LEN,
ValueKey,
write::{BatchBuilder, BlobLink, RegistryClass},
};
use registry::{
schema::prelude::Property,
types::{EnumImpl, id::ObjectId},
};
use std::time::Instant;
use trc::{AddContext, PurgeEvent};
use types::{
blob::BlobClass,
@@ -38,49 +42,6 @@ impl Store {
.caused_by(trc::location!())
}
pub async fn blob_quota(&self, account_id: u32) -> trc::Result<BlobQuota> {
let from_key = ValueKey {
account_id,
collection: 0,
document_id: 0,
class: ValueClass::Blob(BlobOp::Quota {
hash: BlobHash::default(),
until: 0,
}),
};
let to_key = ValueKey {
account_id: account_id + 1,
collection: 0,
document_id: 0,
class: ValueClass::Blob(BlobOp::Quota {
hash: BlobHash::default(),
until: u64::MAX,
}),
};
let now = now();
let mut quota = BlobQuota { bytes: 0, count: 0 };
self.iterate(
IterateParams::new(from_key, to_key).ascending(),
|key, value| {
let until = key.deserialize_be_u64(key.len() - U64_LEN)?;
if until > now {
let bytes = value.deserialize_be_u32(0)?;
if bytes > 0 {
quota.bytes += bytes as usize;
quota.count += 1;
}
}
Ok(true)
},
)
.await
.caused_by(trc::location!())?;
Ok(quota)
}
pub async fn blob_has_access(
&self,
hash: impl AsRef<BlobHash> + Sync + Send,
@@ -195,6 +156,31 @@ impl Store {
op: ValueOp::Clear,
});
}
for (account_id, object_id) in state.delete_registry {
if batch.is_large_batch() {
self.write(batch.build_all())
.await
.caused_by(trc::location!())?;
batch = BatchBuilder::new();
}
let item_id = object_id.id().id();
let object_id = object_id.object().to_id();
if let Some(account_id) = account_id {
batch.clear(ValueClass::Registry(RegistryClass::Index {
index_id: Property::AccountId.to_id(),
object_id,
item_id,
key: account_id.serialize(),
}));
}
batch.clear(ValueClass::Registry(RegistryClass::Id {
object_id,
item_id,
}));
}
if !batch.is_empty() {
self.write(batch.build_all())
.await
@@ -220,7 +206,7 @@ struct BlobPurgeState {
last_hash: BlobHash,
last_hash_is_linked: bool,
delete_keys: Vec<(Option<u32>, BlobOp)>,
spam_train_samples: Vec<(u32, u64)>,
delete_registry: Vec<(Option<u32>, ObjectId)>,
now: u64,
total_deleted: u64,
total_active: u64,
@@ -232,7 +218,7 @@ impl BlobPurgeState {
last_hash: BlobHash::default(),
last_hash_is_linked: true, // Avoid deleting non-existing last_hash on first iteration
delete_keys: Vec::new(),
spam_train_samples: Vec::new(),
delete_registry: Vec::new(),
now: now(),
total_deleted: 0,
total_active: 0,
@@ -257,42 +243,6 @@ impl BlobPurgeState {
));
} else {
self.total_active += 1;
if !self.spam_train_samples.is_empty() {
if self.spam_train_samples.len() > 1 {
// Sort by account_id ascending, then until descending
self.spam_train_samples
.sort_unstable_by(|(a_id, a_until), (b_id, b_until)| {
a_id.cmp(b_id).then_with(|| b_until.cmp(a_until))
});
let mut samples = self.spam_train_samples.iter().peekable();
while let Some((account_id, _)) = samples.next() {
// Keep only the latest sample per account
while let Some((next_account_id, next_until)) = samples.peek() {
if next_account_id == account_id {
self.delete_keys.push((
Some(*account_id),
BlobOp::SpamSample {
hash: self.last_hash.clone(),
until: *next_until,
},
));
self.delete_keys.push((
Some(*account_id),
BlobOp::Link {
hash: self.last_hash.clone(),
to: BlobLink::Temporary { until: *next_until },
},
));
samples.next();
} else {
break;
}
}
}
}
self.spam_train_samples.clear();
}
self.last_hash = new_hash;
}
}
@@ -319,44 +269,12 @@ impl BlobPurgeState {
to: BlobLink::Temporary { until },
},
));
match value.first().copied() {
Some(BlobLink::QUOTA_LINK) => {
self.delete_keys.push((
Some(account_id),
BlobOp::Quota {
hash: self.last_hash.clone(),
until,
},
));
}
Some(BlobLink::UNDELETE_LINK) => {
self.delete_keys.push((
Some(account_id),
BlobOp::Undelete {
hash: self.last_hash.clone(),
until,
},
));
}
Some(BlobLink::SPAM_SAMPLE_LINK) => {
self.delete_keys.push((
Some(account_id),
BlobOp::SpamSample {
hash: self.last_hash.clone(),
until,
},
));
}
_ => {}
if value.len() == U16_LEN + U64_LEN {
self.delete_registry.push((
(account_id != u32::MAX).then_some(account_id),
ObjectId::deserialize(value)?,
));
}
} else {
// Delete attempts to train the same message multiple times
if matches!(value.first(), Some(&BlobLink::SPAM_SAMPLE_LINK)) {
let account_id = key.deserialize_be_u32(BLOB_HASH_LEN)?;
self.spam_train_samples.push((account_id, until));
}
self.last_hash_is_linked = true;
}
Ok(())
}

View File

@@ -9,16 +9,18 @@ use super::{
TelemetryClass, ValueClass,
};
use crate::{
Deserialize, IndexKey, IndexKeyPrefix, Key, LogKey, SUBSPACE_ACL, SUBSPACE_BLOB_EXTRA,
SUBSPACE_BLOB_LINK, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_IN_MEMORY_VALUE,
Deserialize, IndexKey, IndexKeyPrefix, Key, LogKey, SUBSPACE_ACL, SUBSPACE_BLOB_LINK,
SUBSPACE_COUNTER, SUBSPACE_DELETED_ITEMS, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_IN_MEMORY_VALUE,
SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_PROPERTY, SUBSPACE_QUEUE_EVENT,
SUBSPACE_QUEUE_MESSAGE, SUBSPACE_QUOTA, SUBSPACE_REGISTRY, SUBSPACE_REPORT_IN,
SUBSPACE_REPORT_OUT, SUBSPACE_SEARCH_INDEX, SUBSPACE_TASK_QUEUE, SUBSPACE_TELEMETRY_METRIC,
SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN, U64_LEN, ValueKey, WITH_SUBSPACE,
SUBSPACE_QUEUE_MESSAGE, SUBSPACE_QUOTA, SUBSPACE_REGISTRY, SUBSPACE_REGISTRY_DIRECTORY,
SUBSPACE_REPORT_IN, SUBSPACE_REPORT_OUT, SUBSPACE_SEARCH_INDEX, SUBSPACE_SPAM_SAMPLES,
SUBSPACE_TASK_QUEUE, SUBSPACE_TELEMETRY_METRIC, SUBSPACE_TELEMETRY_SPAN, U16_LEN, U32_LEN,
U64_LEN, ValueKey, WITH_SUBSPACE,
write::{
BlobLink, IndexPropertyClass, RegistryClass, SearchIndex, SearchIndexId, SearchIndexType,
},
};
use registry::schema::prelude::ObjectType;
use std::convert::TryInto;
use types::{
blob_hash::BLOB_HASH_LEN,
@@ -295,49 +297,8 @@ impl ValueClass {
.write(collection)
.write(document_id),
ValueClass::TaskQueue(task) => match task {
TaskQueueClass::UpdateIndex {
index,
is_insert,
due,
} => serializer
.write(due.inner())
.write(account_id)
.write(if *is_insert { 7u8 } else { 8u8 })
.write(document_id)
.write(index.to_u8()),
TaskQueueClass::SendAlarm {
due,
event_id,
alarm_id,
is_email_alert,
} => serializer
.write(due.inner())
.write(account_id)
.write(if *is_email_alert { 3u8 } else { 6u8 })
.write(document_id)
.write(*event_id)
.write(*alarm_id),
TaskQueueClass::SendImip { due, is_payload } => {
if !*is_payload {
serializer
.write(due.inner())
.write(account_id)
.write(4u8)
.write(document_id)
} else {
serializer
.write(u64::MAX)
.write(account_id)
.write(5u8)
.write(document_id)
.write(due.inner())
}
}
TaskQueueClass::MergeThreads { due } => serializer
.write(due.inner())
.write(account_id)
.write(9u8)
.write(document_id),
TaskQueueClass::Task { id } => serializer.write(*id),
TaskQueueClass::Due { id, due } => serializer.write(*due).write(*id),
},
ValueClass::Blob(op) => match op {
BlobOp::Commit { hash } => serializer.write::<&[u8]>(hash.as_ref()),
@@ -353,21 +314,6 @@ impl ValueClass {
.write(account_id)
.write(*until),
},
BlobOp::Quota { hash, until } => serializer
.write(BlobLink::QUOTA_LINK)
.write(account_id)
.write::<&[u8]>(hash.as_ref())
.write(*until),
BlobOp::Undelete { hash, until } => serializer
.write(BlobLink::UNDELETE_LINK)
.write(account_id)
.write::<&[u8]>(hash.as_ref())
.write(*until),
BlobOp::SpamSample { hash, until } => serializer
.write(BlobLink::SPAM_SAMPLE_LINK)
.write(*until)
.write(account_id)
.write::<&[u8]>(hash.as_ref()),
},
ValueClass::InMemory(lookup) => match lookup {
InMemoryClass::Key(key) => serializer.write(key.as_slice()),
@@ -545,12 +491,6 @@ impl ValueClass {
}
}
impl BlobLink {
pub const QUOTA_LINK: u8 = 0;
pub const UNDELETE_LINK: u8 = 1;
pub const SPAM_SAMPLE_LINK: u8 = 2;
}
impl<T: AsRef<[u8]> + Sync + Send + Clone> Key for IndexKey<T> {
fn subspace(&self) -> u8 {
SUBSPACE_INDEXES
@@ -592,6 +532,19 @@ impl<T: AsRef<[u8]> + Sync + Send + Clone> Key for AnyKey<T> {
}
}
const MAILBOX_COLLECTION: u8 = Collection::Mailbox as u8;
const MAILBOX_COUNTER_FIELD: u8 = MailboxField::UidCounter as u8;
const REG_DELETED_ITEM: u16 = ObjectType::DeletedItem as u16;
const REG_SPAM_SAMPLE: u16 = ObjectType::SpamTrainingSample as u16;
const REG_ACCOUNT: u16 = ObjectType::Account as u16;
const REG_DOMAIN: u16 = ObjectType::Domain as u16;
const REG_TENANT: u16 = ObjectType::Tenant as u16;
const REG_ROLE: u16 = ObjectType::Role as u16;
const REG_OAUTH_CLIENT: u16 = ObjectType::OAuthClient as u16;
const REG_MAILING_LIST: u16 = ObjectType::MailingList as u16;
const REG_MASKED_EMAIL: u16 = ObjectType::MaskedEmail as u16;
const REG_PUBLIC_KEY: u16 = ObjectType::PublicKey as u16;
impl ValueClass {
pub fn serialized_size(&self) -> usize {
match self {
@@ -621,23 +574,10 @@ impl ValueClass {
BlobLink::Temporary { .. } => U32_LEN + U64_LEN,
}
}
BlobOp::Quota { .. } | BlobOp::Undelete { .. } => {
BLOB_HASH_LEN + U32_LEN + U64_LEN + 1
}
BlobOp::SpamSample { .. } => BLOB_HASH_LEN + U32_LEN + 2,
},
ValueClass::TaskQueue(e) => match e {
TaskQueueClass::UpdateIndex { .. } => (U64_LEN * 2) + 2,
TaskQueueClass::SendAlarm { .. } | TaskQueueClass::MergeThreads { .. } => {
U64_LEN + (U32_LEN * 3) + 1
}
TaskQueueClass::SendImip { is_payload, .. } => {
if *is_payload {
(U64_LEN * 2) + (U32_LEN * 2) + 1
} else {
U64_LEN + (U32_LEN * 2) + 1
}
}
TaskQueueClass::Task { .. } => U64_LEN + 1,
TaskQueueClass::Due { .. } => (U64_LEN * 2) + 1,
},
ValueClass::Queue(q) => match q {
QueueClass::Message(_) => U64_LEN,
@@ -671,9 +611,6 @@ impl ValueClass {
}
pub fn subspace(&self, collection: u8) -> u8 {
const MAILBOX_COLLECTION: u8 = Collection::Mailbox as u8;
const MAILBOX_COUNTER_FIELD: u8 = MailboxField::UidCounter as u8;
match self {
ValueClass::Property(field) => {
if collection == MAILBOX_COLLECTION && *field == MAILBOX_COUNTER_FIELD {
@@ -687,17 +624,26 @@ impl ValueClass {
ValueClass::TaskQueue { .. } => SUBSPACE_TASK_QUEUE,
ValueClass::Blob(op) => match op {
BlobOp::Commit { .. } | BlobOp::Link { .. } => SUBSPACE_BLOB_LINK,
BlobOp::Quota { .. } | BlobOp::Undelete { .. } | BlobOp::SpamSample { .. } => {
SUBSPACE_BLOB_EXTRA
}
},
ValueClass::Registry(registry) => {
if matches!(registry, RegistryClass::IdCounter { .. }) {
SUBSPACE_COUNTER
} else {
SUBSPACE_REGISTRY
}
}
ValueClass::Registry(registry) => match registry {
RegistryClass::Item { object_id, .. }
| RegistryClass::Id { object_id, .. }
| RegistryClass::Index { object_id, .. }
| RegistryClass::Reference {
to_object_id: object_id,
..
} => match *object_id {
REG_ACCOUNT | REG_DOMAIN | REG_TENANT | REG_ROLE | REG_OAUTH_CLIENT
| REG_MAILING_LIST | REG_MASKED_EMAIL | REG_PUBLIC_KEY => {
SUBSPACE_REGISTRY_DIRECTORY
}
REG_DELETED_ITEM => SUBSPACE_DELETED_ITEMS,
REG_SPAM_SAMPLE => SUBSPACE_SPAM_SAMPLES,
_ => SUBSPACE_REGISTRY,
},
RegistryClass::IndexGlobal { .. } => SUBSPACE_REGISTRY,
RegistryClass::IdCounter { .. } => SUBSPACE_COUNTER,
},
ValueClass::InMemory(lookup) => match lookup {
InMemoryClass::Key(_) => SUBSPACE_IN_MEMORY_VALUE,
InMemoryClass::Counter(_) => SUBSPACE_IN_MEMORY_COUNTER,

View File

@@ -222,30 +222,10 @@ pub enum SearchIndexId {
#[derive(Debug, PartialEq, Clone, Eq, Hash)]
pub enum TaskQueueClass {
UpdateIndex {
due: TaskEpoch,
index: SearchIndex,
is_insert: bool,
},
SendAlarm {
due: TaskEpoch,
event_id: u16,
alarm_id: u16,
is_email_alert: bool,
},
SendImip {
due: TaskEpoch,
is_payload: bool,
},
MergeThreads {
due: TaskEpoch,
},
Task { id: u64 },
Due { id: u64, due: u64 },
}
#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)]
#[repr(transparent)]
pub struct TaskEpoch(pub(crate) u64);
#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)]
pub enum SearchIndex {
Email,
@@ -396,9 +376,6 @@ pub struct SetOperation {
pub enum BlobOp {
Commit { hash: BlobHash },
Link { hash: BlobHash, to: BlobLink },
Quota { hash: BlobHash, until: u64 },
Undelete { hash: BlobHash, until: u64 },
SpamSample { hash: BlobHash, until: u64 },
}
#[derive(Debug, PartialEq, Clone, Eq, Hash)]
@@ -743,56 +720,3 @@ impl AsRef<[Param]> for Params {
&self.0
}
}
impl TaskEpoch {
/*
Structure of the 64-bit epoch:
4 bytes: seconds since custom epoch (1632280000)
2 bytes: attempt number
2 bytes: sequence id
*/
const EPOCH_OFFSET: u64 = 1632280000;
pub fn now() -> Self {
Self::new(now())
}
pub fn new(timestamp: u64) -> Self {
Self(timestamp.saturating_sub(Self::EPOCH_OFFSET) << 32)
}
pub fn with_attempt(mut self, attempt: u16) -> Self {
self.0 |= (attempt as u64) << 16;
self
}
pub fn with_sequence_id(mut self, sequence_id: u16) -> Self {
self.0 |= sequence_id as u64;
self
}
pub fn with_random_sequence_id(self) -> Self {
self.with_sequence_id(rand::random())
}
pub fn due(&self) -> u64 {
(self.0 >> 32) + Self::EPOCH_OFFSET
}
pub fn attempt(&self) -> u16 {
(self.0 >> 16) as u16
}
pub fn sequence_id(&self) -> u16 {
self.0 as u16
}
pub fn inner(&self) -> u64 {
self.0
}
pub fn from_inner(inner: u64) -> Self {
Self(inner)
}
}