From c6960bfba8b993d0c3bdbd8789360b1aea07fbd2 Mon Sep 17 00:00:00 2001 From: mdecimus Date: Tue, 23 Apr 2024 20:03:30 +0200 Subject: [PATCH] documentId generation within the database transaction --- .../directory/src/backend/internal/manage.rs | 191 ++++---- crates/directory/src/lib.rs | 8 + crates/imap/src/op/create.rs | 34 +- crates/imap/src/op/rename.rs | 72 +-- crates/imap/src/op/thread.rs | 5 +- crates/jmap-proto/src/object/index.rs | 6 +- crates/jmap-proto/src/types/keyword.rs | 4 +- crates/jmap/src/api/http.rs | 13 +- crates/jmap/src/email/copy.rs | 91 ++-- crates/jmap/src/email/get.rs | 2 +- crates/jmap/src/email/ingest.rs | 116 ++--- crates/jmap/src/identity/get.rs | 10 +- crates/jmap/src/identity/set.rs | 7 +- crates/jmap/src/lib.rs | 44 +- crates/jmap/src/mailbox/mod.rs | 8 +- crates/jmap/src/mailbox/query.rs | 5 +- crates/jmap/src/mailbox/set.rs | 87 ++-- crates/jmap/src/push/set.rs | 7 +- crates/jmap/src/services/index.rs | 4 +- crates/jmap/src/sieve/get.rs | 20 +- crates/jmap/src/sieve/set.rs | 26 +- crates/jmap/src/submission/set.rs | 7 +- crates/jmap/src/vacation/set.rs | 73 +-- crates/main/Cargo.toml | 3 +- crates/managesieve/src/op/putscript.rs | 18 +- crates/store/src/backend/foundationdb/read.rs | 70 ++- .../store/src/backend/foundationdb/write.rs | 328 ++++---------- crates/store/src/backend/mod.rs | 6 - crates/store/src/backend/mysql/read.rs | 6 +- crates/store/src/backend/mysql/write.rs | 196 ++++---- crates/store/src/backend/postgres/read.rs | 6 +- .../store/src/backend/postgres/read_dense.rs | 155 ------- .../src/backend/postgres/read_roaring.rs | 151 ------- crates/store/src/backend/postgres/write.rs | 203 +++++---- .../store/src/backend/postgres/write_dense.rs | 386 ---------------- .../src/backend/postgres/write_roaring.rs | 417 ------------------ crates/store/src/backend/rocksdb/main.rs | 34 +- crates/store/src/backend/rocksdb/read.rs | 39 +- crates/store/src/backend/rocksdb/write.rs | 410 ++++++----------- crates/store/src/backend/sqlite/read.rs | 6 +- crates/store/src/backend/sqlite/write.rs | 129 +++--- crates/store/src/dispatch/lookup.rs | 14 +- crates/store/src/dispatch/store.rs | 32 +- crates/store/src/fts/query.rs | 10 +- crates/store/src/lib.rs | 7 +- crates/store/src/query/filter.rs | 2 +- crates/store/src/query/mod.rs | 14 +- crates/store/src/query/sort.rs | 4 +- crates/store/src/write/assign_id.rs | 170 ------- crates/store/src/write/batch.rs | 68 ++- crates/store/src/write/bitmap.rs | 64 ++- crates/store/src/write/blob.rs | 53 +-- crates/store/src/write/hash.rs | 2 +- crates/store/src/write/key.rs | 189 ++++---- crates/store/src/write/log.rs | 86 +++- crates/store/src/write/mod.rs | 269 +++++++++-- tests/Cargo.toml | 3 +- tests/src/directory/internal.rs | 269 ++++++----- tests/src/directory/ldap.rs | 2 +- tests/src/directory/mod.rs | 14 +- tests/src/imap/mod.rs | 15 +- tests/src/jmap/auth_limits.rs | 5 +- tests/src/jmap/email_query.rs | 64 ++- tests/src/jmap/email_query_changes.rs | 11 +- tests/src/jmap/email_set.rs | 19 +- tests/src/jmap/mailbox.rs | 76 ++-- tests/src/jmap/mod.rs | 19 +- tests/src/jmap/push_subscription.rs | 26 +- tests/src/jmap/quota.rs | 2 +- tests/src/jmap/thread_merge.rs | 18 +- tests/src/store/assign_id.rs | 188 +++----- tests/src/store/mod.rs | 2 +- tests/src/store/ops.rs | 107 ++++- tests/src/store/query.rs | 16 +- 74 files changed, 2132 insertions(+), 3111 deletions(-) delete mode 100644 crates/store/src/backend/postgres/read_dense.rs delete mode 100644 crates/store/src/backend/postgres/read_roaring.rs delete mode 100644 crates/store/src/backend/postgres/write_dense.rs delete mode 100644 crates/store/src/backend/postgres/write_roaring.rs delete mode 100644 crates/store/src/write/assign_id.rs diff --git a/crates/directory/src/backend/internal/manage.rs b/crates/directory/src/backend/internal/manage.rs index 1b67d546..3cda4ffd 100644 --- a/crates/directory/src/backend/internal/manage.rs +++ b/crates/directory/src/backend/internal/manage.rs @@ -24,7 +24,8 @@ use jmap_proto::types::collection::Collection; use store::{ write::{ - assert::HashedValue, key::DeserializeBigEndian, BatchBuilder, DirectoryClass, ValueClass, + assert::HashedValue, key::DeserializeBigEndian, AssignedIds, BatchBuilder, DirectoryClass, + MaybeDynamicId, MaybeDynamicValue, SerializeWithId, ValueClass, }, Deserialize, IterateParams, Serialize, Store, ValueKey, U32_LEN, }; @@ -117,11 +118,6 @@ impl ManageDirectory for Store { return Ok(account_id); } - // Assign new ID - let account_id = self - .assign_document_id(u32::MAX, Collection::Principal) - .await?; - // Write account ID let name_key = ValueClass::Directory(DirectoryClass::NameToId(name.as_bytes().to_vec())); @@ -129,25 +125,24 @@ impl ManageDirectory for Store { batch .with_account_id(u32::MAX) .with_collection(Collection::Principal) - .create_document(account_id) .assert_value(name_key.clone(), ()) + .create_document() + .set(name_key, DynamicPrincipalIdType(Type::Individual)) .set( - name_key, - PrincipalIdType::new(account_id, Type::Individual).serialize(), - ) - .set( - ValueClass::Directory(DirectoryClass::Principal(account_id)), + ValueClass::Directory(DirectoryClass::Principal(MaybeDynamicId::Dynamic(0))), Principal { - id: account_id, typ: Type::Individual, name: name.to_string(), ..Default::default() - } - .serialize(), + }, ); - match self.write(batch.build()).await { - Ok(_) => { + match self + .write(batch.build()) + .await + .and_then(|r| r.last_document_id()) + { + Ok(account_id) => { return Ok(account_id); } Err(store::Error::AssertValueFailed) if try_count < 3 => { @@ -208,18 +203,13 @@ impl ManageDirectory for Store { } } - // Assign accountId - principal.id = self - .assign_document_id(u32::MAX, Collection::Principal) - .await?; - // Write principal let mut batch = BatchBuilder::new(); - let ptype = PrincipalIdType::new(principal.id, principal.typ.into_base_type()).serialize(); + let ptype = DynamicPrincipalIdType(principal.typ.into_base_type()); batch .with_account_id(u32::MAX) .with_collection(Collection::Principal) - .create_document(principal.id) + .create_document() .assert_value( ValueClass::Directory(DirectoryClass::NameToId( principal.name.clone().into_bytes(), @@ -227,19 +217,19 @@ impl ManageDirectory for Store { (), ) .set( - ValueClass::Directory(DirectoryClass::Principal(principal.id)), - (&principal).serialize(), + ValueClass::Directory(DirectoryClass::Principal(MaybeDynamicId::Dynamic(0))), + principal.clone(), ) .set( ValueClass::Directory(DirectoryClass::NameToId(principal.name.into_bytes())), - ptype.clone(), + ptype, ); // Write email to id mapping for email in principal.emails { batch.set( ValueClass::Directory(DirectoryClass::EmailToId(email.into_bytes())), - ptype.clone(), + ptype, ); } @@ -247,15 +237,15 @@ impl ManageDirectory for Store { for member_of in principal.member_of { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: principal.id, - member_of, + principal_id: MaybeDynamicId::Dynamic(0), + member_of: MaybeDynamicId::Static(member_of), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: member_of, - has_member: principal.id, + principal_id: MaybeDynamicId::Static(member_of), + has_member: MaybeDynamicId::Dynamic(0), }), vec![], ); @@ -263,23 +253,24 @@ impl ManageDirectory for Store { for member_id in members { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: member_id, - member_of: principal.id, + principal_id: MaybeDynamicId::Static(member_id), + member_of: MaybeDynamicId::Dynamic(0), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: principal.id, - has_member: member_id, + principal_id: MaybeDynamicId::Dynamic(0), + has_member: MaybeDynamicId::Static(member_id), }), vec![], ); } - self.write(batch.build()).await?; - - Ok(principal.id) + self.write(batch.build()) + .await + .and_then(|r| r.last_document_id()) + .map_err(Into::into) } async fn delete_account(&self, by: QueryBy<'_>) -> crate::Result<()> { @@ -314,7 +305,9 @@ impl ManageDirectory for Store { batch .with_account_id(account_id) .clear(DirectoryClass::NameToId(principal.name.into_bytes())) - .clear(DirectoryClass::Principal(account_id)) + .clear(DirectoryClass::Principal(MaybeDynamicId::Static( + account_id, + ))) .clear(DirectoryClass::UsedQuota(account_id)); for email in principal.emails { @@ -323,23 +316,23 @@ impl ManageDirectory for Store { for member_id in self.get_member_of(account_id).await? { batch.clear(DirectoryClass::MemberOf { - principal_id: account_id, - member_of: member_id, + principal_id: MaybeDynamicId::Static(account_id), + member_of: MaybeDynamicId::Static(member_id), }); batch.clear(DirectoryClass::Members { - principal_id: member_id, - has_member: account_id, + principal_id: MaybeDynamicId::Static(member_id), + has_member: MaybeDynamicId::Static(account_id), }); } for member_id in self.get_members(account_id).await? { batch.clear(DirectoryClass::MemberOf { - principal_id: member_id, - member_of: account_id, + principal_id: MaybeDynamicId::Static(member_id), + member_of: MaybeDynamicId::Static(account_id), }); batch.clear(DirectoryClass::Members { - principal_id: account_id, - has_member: member_id, + principal_id: MaybeDynamicId::Static(account_id), + has_member: MaybeDynamicId::Static(member_id), }); } @@ -386,7 +379,9 @@ impl ManageDirectory for Store { if update_principal { batch.assert_value( - ValueClass::Directory(DirectoryClass::Principal(account_id)), + ValueClass::Directory(DirectoryClass::Principal(MaybeDynamicId::Static( + account_id, + ))), &principal, ); } @@ -556,15 +551,15 @@ impl ManageDirectory for Store { if !member_of.contains(&member_id) { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: account_id, - member_of: member_id, + principal_id: MaybeDynamicId::Static(account_id), + member_of: MaybeDynamicId::Static(member_id), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: member_id, - has_member: account_id, + principal_id: MaybeDynamicId::Static(member_id), + has_member: MaybeDynamicId::Static(account_id), }), vec![], ); @@ -576,12 +571,12 @@ impl ManageDirectory for Store { for member_id in &member_of { if !new_member_of.contains(member_id) { batch.clear(ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: account_id, - member_of: *member_id, + principal_id: MaybeDynamicId::Static(account_id), + member_of: MaybeDynamicId::Static(*member_id), })); batch.clear(ValueClass::Directory(DirectoryClass::Members { - principal_id: *member_id, - has_member: account_id, + principal_id: MaybeDynamicId::Static(*member_id), + has_member: MaybeDynamicId::Static(account_id), })); } } @@ -599,15 +594,15 @@ impl ManageDirectory for Store { if !member_of.contains(&member_id) { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: account_id, - member_of: member_id, + principal_id: MaybeDynamicId::Static(account_id), + member_of: MaybeDynamicId::Static(member_id), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: member_id, - has_member: account_id, + principal_id: MaybeDynamicId::Static(member_id), + has_member: MaybeDynamicId::Static(account_id), }), vec![], ); @@ -622,12 +617,12 @@ impl ManageDirectory for Store { if let Some(member_id) = self.get_account_id(&member).await? { if let Some(pos) = member_of.iter().position(|v| *v == member_id) { batch.clear(ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: account_id, - member_of: member_id, + principal_id: MaybeDynamicId::Static(account_id), + member_of: MaybeDynamicId::Static(member_id), })); batch.clear(ValueClass::Directory(DirectoryClass::Members { - principal_id: member_id, - has_member: account_id, + principal_id: MaybeDynamicId::Static(member_id), + has_member: MaybeDynamicId::Static(account_id), })); member_of.remove(pos); } @@ -647,15 +642,15 @@ impl ManageDirectory for Store { if !members.contains(&member_id) { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: member_id, - member_of: account_id, + principal_id: MaybeDynamicId::Static(member_id), + member_of: MaybeDynamicId::Static(account_id), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: account_id, - has_member: member_id, + principal_id: MaybeDynamicId::Static(account_id), + has_member: MaybeDynamicId::Static(member_id), }), vec![], ); @@ -667,12 +662,12 @@ impl ManageDirectory for Store { for member_id in &members { if !new_members.contains(member_id) { batch.clear(ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: *member_id, - member_of: account_id, + principal_id: MaybeDynamicId::Static(*member_id), + member_of: MaybeDynamicId::Static(account_id), })); batch.clear(ValueClass::Directory(DirectoryClass::Members { - principal_id: account_id, - has_member: *member_id, + principal_id: MaybeDynamicId::Static(account_id), + has_member: MaybeDynamicId::Static(*member_id), })); } } @@ -690,15 +685,15 @@ impl ManageDirectory for Store { if !members.contains(&member_id) { batch.set( ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: member_id, - member_of: account_id, + principal_id: MaybeDynamicId::Static(member_id), + member_of: MaybeDynamicId::Static(account_id), }), vec![], ); batch.set( ValueClass::Directory(DirectoryClass::Members { - principal_id: account_id, - has_member: member_id, + principal_id: MaybeDynamicId::Static(account_id), + has_member: MaybeDynamicId::Static(member_id), }), vec![], ); @@ -713,12 +708,12 @@ impl ManageDirectory for Store { if let Some(member_id) = self.get_account_id(&member).await? { if let Some(pos) = members.iter().position(|v| *v == member_id) { batch.clear(ValueClass::Directory(DirectoryClass::MemberOf { - principal_id: member_id, - member_of: account_id, + principal_id: MaybeDynamicId::Static(member_id), + member_of: MaybeDynamicId::Static(account_id), })); batch.clear(ValueClass::Directory(DirectoryClass::Members { - principal_id: account_id, - has_member: member_id, + principal_id: MaybeDynamicId::Static(account_id), + has_member: MaybeDynamicId::Static(member_id), })); members.remove(pos); } @@ -733,7 +728,9 @@ impl ManageDirectory for Store { if update_principal { batch.set( - ValueClass::Directory(DirectoryClass::Principal(account_id)), + ValueClass::Directory(DirectoryClass::Principal(MaybeDynamicId::Static( + account_id, + ))), principal.inner.serialize(), ); } @@ -971,6 +968,36 @@ impl ManageDirectory for Store { } } +impl SerializeWithId for Principal { + fn serialize_with_id(&self, ids: &AssignedIds) -> store::Result> { + let mut principal = self.clone(); + principal.id = ids.last_document_id()?; + Ok(principal.serialize()) + } +} + +impl From> for MaybeDynamicValue { + fn from(principal: Principal) -> Self { + MaybeDynamicValue::Dynamic(Box::new(principal)) + } +} + +#[derive(Clone, Copy)] +struct DynamicPrincipalIdType(Type); + +impl SerializeWithId for DynamicPrincipalIdType { + fn serialize_with_id(&self, ids: &AssignedIds) -> store::Result> { + ids.last_document_id() + .map(|account_id| PrincipalIdType::new(account_id, self.0).serialize()) + } +} + +impl From for MaybeDynamicValue { + fn from(value: DynamicPrincipalIdType) -> Self { + MaybeDynamicValue::Dynamic(Box::new(value)) + } +} + impl From> for Principal { fn from(principal: Principal) -> Self { Principal { diff --git a/crates/directory/src/lib.rs b/crates/directory/src/lib.rs index b82fe222..abcda5be 100644 --- a/crates/directory/src/lib.rs +++ b/crates/directory/src/lib.rs @@ -214,6 +214,14 @@ impl Principal { } } +impl Principal { + pub fn into_sorted(mut self) -> Self { + self.member_of.sort_unstable(); + self.emails.sort_unstable(); + self + } +} + impl From for DirectoryError { fn from(error: LdapError) -> Self { tracing::warn!( diff --git a/crates/imap/src/op/create.rs b/crates/imap/src/op/create.rs index 651a262c..9727f82e 100644 --- a/crates/imap/src/op/create.rs +++ b/crates/imap/src/op/create.rs @@ -90,23 +90,10 @@ impl SessionData { return StatusResponse::database_failure().with_tag(arguments.tag); } }; - let mut batch = BatchBuilder::new(); - batch - .with_account_id(params.account_id) - .with_collection(Collection::Mailbox); + let mut parent_id = params.parent_mailbox_id.map(|id| id + 1).unwrap_or(0); let mut create_ids = Vec::with_capacity(params.path.len()); for (pos, &path_item) in params.path.iter().enumerate() { - let mailbox_id = match self - .jmap - .assign_document_id(params.account_id, Collection::Mailbox) - .await - { - Ok(mailbox_id) => mailbox_id, - Err(_) => { - return StatusResponse::database_failure().with_tag(arguments.tag); - } - }; let mut mailbox = Object::with_capacity(4) .with_property(Property::Name, path_item) .with_property(Property::ParentId, Value::Id(Id::from(parent_id))) @@ -119,15 +106,30 @@ impl SessionData { mailbox.set(Property::Role, mailbox_role); } } + let mut batch = BatchBuilder::new(); batch - .create_document(mailbox_id) + .with_account_id(params.account_id) + .with_collection(Collection::Mailbox) + .create_document() .custom(ObjectIndexBuilder::new(SCHEMA).with_changes(mailbox)); + let mailbox_id = match self.jmap.write_batch_expect_id(batch).await { + Ok(mailbox_id) => mailbox_id, + Err(_) => { + return StatusResponse::database_failure().with_tag(arguments.tag); + } + }; changes.log_insert(Collection::Mailbox, mailbox_id); parent_id = mailbox_id + 1; create_ids.push(mailbox_id); } + + // Write changes let change_id = changes.change_id; - batch.custom(changes); + let mut batch = BatchBuilder::new(); + batch + .with_account_id(params.account_id) + .with_collection(Collection::Mailbox) + .custom(changes); if self.jmap.write_batch(batch).await.is_err() { return StatusResponse::database_failure().with_tag(arguments.tag); } diff --git a/crates/imap/src/op/rename.rs b/crates/imap/src/op/rename.rs index 299c5f01..01fb45b4 100644 --- a/crates/imap/src/op/rename.rs +++ b/crates/imap/src/op/rename.rs @@ -145,51 +145,57 @@ impl SessionData { return StatusResponse::database_failure().with_tag(arguments.tag); } }; - let mut batch = BatchBuilder::new(); - batch - .with_account_id(params.account_id) - .with_collection(Collection::Mailbox); + let mut parent_id = params.parent_mailbox_id.map(|id| id + 1).unwrap_or(0); let mut create_ids = Vec::with_capacity(params.path.len()); for &path_item in params.path.iter() { - let mailbox_id = match self - .jmap - .assign_document_id(params.account_id, Collection::Mailbox) - .await - { + let mut batch = BatchBuilder::new(); + batch + .with_account_id(params.account_id) + .with_collection(Collection::Mailbox) + .create_document() + .custom( + ObjectIndexBuilder::new(SCHEMA).with_changes( + Object::with_capacity(3) + .with_property(Property::Name, path_item) + .with_property(Property::ParentId, Value::Id(Id::from(parent_id))) + .with_property( + Property::Cid, + Value::UnsignedInt(rand::random::() as u64), + ), + ), + ); + + let mailbox_id = match self.jmap.write_batch_expect_id(batch).await { Ok(mailbox_id) => mailbox_id, Err(_) => { return StatusResponse::database_failure().with_tag(arguments.tag); } }; - batch.create_document(mailbox_id).custom( - ObjectIndexBuilder::new(SCHEMA).with_changes( - Object::with_capacity(3) - .with_property(Property::Name, path_item) - .with_property(Property::ParentId, Value::Id(Id::from(parent_id))) - .with_property( - Property::Cid, - Value::UnsignedInt(rand::random::() as u64), - ), - ), - ); + changes.log_insert(Collection::Mailbox, mailbox_id); parent_id = mailbox_id + 1; create_ids.push(mailbox_id); } - batch.update_document(mailbox_id).custom( - ObjectIndexBuilder::new(SCHEMA) - .with_current(mailbox) - .with_changes( - Object::with_capacity(3) - .with_property(Property::Name, new_mailbox_name) - .with_property(Property::ParentId, Value::Id(Id::from(parent_id))) - .with_property( - Property::Cid, - Value::UnsignedInt(rand::random::() as u64), - ), - ), - ); + + let mut batch = BatchBuilder::new(); + batch + .with_account_id(params.account_id) + .with_collection(Collection::Mailbox) + .update_document(mailbox_id) + .custom( + ObjectIndexBuilder::new(SCHEMA) + .with_current(mailbox) + .with_changes( + Object::with_capacity(3) + .with_property(Property::Name, new_mailbox_name) + .with_property(Property::ParentId, Value::Id(Id::from(parent_id))) + .with_property( + Property::Cid, + Value::UnsignedInt(rand::random::() as u64), + ), + ), + ); changes.log_update(Collection::Mailbox, mailbox_id); let change_id = changes.change_id; diff --git a/crates/imap/src/op/thread.rs b/crates/imap/src/op/thread.rs index 0f5d0ee6..dcbfb9ed 100644 --- a/crates/imap/src/op/thread.rs +++ b/crates/imap/src/op/thread.rs @@ -108,7 +108,10 @@ impl SessionData { let mut threads = threads .into_iter() - .map(|(_, messages)| messages) + .map(|(_, mut messages)| { + messages.sort_unstable(); + messages + }) .collect::>(); threads.sort_unstable(); diff --git a/crates/jmap-proto/src/object/index.rs b/crates/jmap-proto/src/object/index.rs index 7ba65eea..0faccec6 100644 --- a/crates/jmap-proto/src/object/index.rs +++ b/crates/jmap-proto/src/object/index.rs @@ -440,7 +440,7 @@ fn merge_batch( if has_changes { batch.ops.push(Operation::Value { class: Property::Value.into(), - op: ValueOp::Set(current.serialize()), + op: ValueOp::Set(current.serialize().into()), }); } } @@ -629,13 +629,13 @@ impl IntoIndex for &Id { } } -impl From for ValueClass { +impl From for ValueClass { fn from(value: Property) -> Self { ValueClass::Property(value.into()) } } -impl From for BitmapClass { +impl From for BitmapClass { fn from(value: Property) -> Self { BitmapClass::Tag { field: value.into(), diff --git a/crates/jmap-proto/src/types/keyword.rs b/crates/jmap-proto/src/types/keyword.rs index 1d5f75d0..8bf9cac0 100644 --- a/crates/jmap-proto/src/types/keyword.rs +++ b/crates/jmap-proto/src/types/keyword.rs @@ -285,7 +285,7 @@ impl DeserializeFrom for Keyword { } } -impl From for TagValue { +impl From for TagValue { fn from(value: Keyword) -> Self { match value { Keyword::Seen => TagValue::Static(SEEN as u8), @@ -305,7 +305,7 @@ impl From for TagValue { } } -impl From<&Keyword> for TagValue { +impl From<&Keyword> for TagValue { fn from(value: &Keyword) -> Self { match value { Keyword::Seen => TagValue::Static(SEEN as u8), diff --git a/crates/jmap/src/api/http.rs b/crates/jmap/src/api/http.rs index 0f9102a7..5346fad4 100644 --- a/crates/jmap/src/api/http.rs +++ b/crates/jmap/src/api/http.rs @@ -98,6 +98,8 @@ impl JMAP { .await .ok_or_else(|| RequestError::limit(RequestLimitError::SizeRequest)) .and_then(|bytes| { + //let c = println!("<- {}", String::from_utf8_lossy(&bytes)); + Request::parse( &bytes, self.core.jmap.request_max_calls, @@ -105,13 +107,18 @@ impl JMAP { ) }) { Ok(request) => { - //let _ = println!("<- {}", String::from_utf8_lossy(&bytes)); - match self .handle_request(request, access_token, &session.instance) .await { - Ok(response) => response.into_http_response(), + Ok(response) => { + /*let c = println!( + "-> {}", + serde_json::to_string_pretty(&response).unwrap() + );*/ + + response.into_http_response() + } Err(err) => err.into_http_response(), } } diff --git a/crates/jmap/src/email/copy.rs b/crates/jmap/src/email/copy.rs index 7b9f7a88..145eaefa 100644 --- a/crates/jmap/src/email/copy.rs +++ b/crates/jmap/src/email/copy.rs @@ -48,7 +48,10 @@ use jmap_proto::{ }; use mail_parser::{parsers::fields::thread::thread_name, HeaderName, HeaderValue}; use store::{ - write::{BatchBuilder, Bincode, ValueClass, F_BITMAP, F_VALUE}, + write::{ + log::{Changes, LogInsert}, + BatchBuilder, Bincode, MaybeDynamicId, TagValue, ValueClass, F_BITMAP, F_VALUE, + }, BlobClass, }; use utils::map::vec_map::VecMap; @@ -57,7 +60,7 @@ use crate::{auth::AccessToken, mailbox::UidMailbox, services::housekeeper::Event use super::{ index::{EmailIndexBuilder, TrimTextValue, VisitValues, MAX_ID_LENGTH, MAX_SORT_FIELD_LENGTH}, - ingest::IngestedEmail, + ingest::{IngestedEmail, LogEmailInsert}, metadata::MessageMetadata, }; @@ -359,21 +362,11 @@ impl JMAP { }; // Assign id - let message_id = self - .assign_document_id(account_id, Collection::Email) - .await?; let mut email = IngestedEmail { - blob_id: BlobId::new( - metadata.blob_hash.clone(), - BlobClass::Linked { - account_id, - collection: Collection::Email.into(), - document_id: message_id, - }, - ), size: metadata.size, ..Default::default() }; + let blob_hash = metadata.blob_hash.clone(); // Assign IMAP UIDs let mut mailbox_ids = Vec::with_capacity(mailboxes.len()); @@ -395,47 +388,42 @@ impl JMAP { } // Prepare batch + let change_id = self.assign_change_id(account_id).await?; let mut batch = BatchBuilder::new(); - batch.with_account_id(account_id); - - // Build change log - let mut changes = self.begin_changes(account_id).await?; - let thread_id = if let Some(thread_id) = thread_id { - changes.log_child_update(Collection::Thread, thread_id); - thread_id + batch + .with_account_id(account_id) + .with_change_id(change_id) + .with_collection(Collection::Thread); + if let Some(thread_id) = thread_id { + batch.log(Changes::update([thread_id])); } else { - let thread_id = self - .assign_document_id(account_id, Collection::Thread) - .await?; - batch - .with_collection(Collection::Thread) - .create_document(thread_id); - changes.log_insert(Collection::Thread, thread_id); - thread_id + batch.create_document().log(LogInsert()); }; - email.id = Id::from_parts(thread_id, message_id); - email.change_id = changes.change_id; - changes.log_insert(Collection::Email, email.id); - for mailbox_id in &mailboxes { - changes.log_child_update(Collection::Mailbox, *mailbox_id); - } // Build batch + let maybe_thread_id = thread_id + .map(MaybeDynamicId::Static) + .unwrap_or(MaybeDynamicId::Dynamic(0)); batch + .with_collection(Collection::Mailbox) + .log(Changes::child_update(mailboxes.iter().copied())) .with_collection(Collection::Email) - .create_document(message_id) - .value(Property::ThreadId, thread_id, F_VALUE | F_BITMAP) + .create_document() + .log(LogEmailInsert::new(thread_id)) + .set(Property::ThreadId, maybe_thread_id) + .tag(Property::ThreadId, TagValue::Id(maybe_thread_id), 0) .value(Property::MailboxIds, mailbox_ids, F_VALUE | F_BITMAP) .value(Property::Keywords, keywords, F_VALUE | F_BITMAP) - .value(Property::Cid, changes.change_id, F_VALUE) + .value(Property::Cid, change_id, F_VALUE) .set( ValueClass::IndexEmail(self.generate_snowflake_id()?), - metadata.blob_hash.clone(), + metadata.blob_hash.as_ref(), ) - .custom(EmailIndexBuilder::set(metadata)) - .custom(changes); + .custom(EmailIndexBuilder::set(metadata)); - self.core + // Insert and obtain ids + let ids = self + .core .storage .data .write(batch.build()) @@ -448,10 +436,31 @@ impl JMAP { "Failed to write message to database."); MethodError::ServerPartialFail })?; + let thread_id = match thread_id { + Some(thread_id) => thread_id, + None => ids + .first_document_id() + .map_err(|_| MethodError::ServerPartialFail)?, + }; + let document_id = ids + .last_document_id() + .map_err(|_| MethodError::ServerPartialFail)?; // Request FTS index let _ = self.inner.housekeeper_tx.send(Event::IndexStart).await; + // Update response + email.id = Id::from_parts(thread_id, document_id); + email.change_id = change_id; + email.blob_id = BlobId::new( + blob_hash, + BlobClass::Linked { + account_id, + collection: Collection::Email.into(), + document_id, + }, + ); + Ok(Ok(email)) } } diff --git a/crates/jmap/src/email/get.rs b/crates/jmap/src/email/get.rs index d43e2cb5..7d4e6d6e 100644 --- a/crates/jmap/src/email/get.rs +++ b/crates/jmap/src/email/get.rs @@ -121,7 +121,7 @@ impl JMAP { MethodError::ServerPartialFail })? .into_iter() - .filter_map(|(thread_id, document_id)| { + .filter_map(|(document_id, thread_id)| { Id::from_parts(thread_id, document_id).into() }) .collect() diff --git a/crates/jmap/src/email/ingest.rs b/crates/jmap/src/email/ingest.rs index 611b3706..71583eef 100644 --- a/crates/jmap/src/email/ingest.rs +++ b/crates/jmap/src/email/ingest.rs @@ -39,10 +39,11 @@ use store::{ ahash::AHashSet, query::Filter, write::{ - log::ChangeLogBuilder, now, BatchBuilder, BitmapClass, TagValue, ValueClass, F_BITMAP, - F_CLEAR, F_VALUE, + log::{ChangeLogBuilder, Changes, LogInsert}, + now, AssignedIds, BatchBuilder, BitmapClass, MaybeDynamicId, MaybeDynamicValue, + SerializeWithId, TagValue, ValueClass, F_BITMAP, F_CLEAR, F_VALUE, }, - BitmapKey, BlobClass, + BitmapKey, BlobClass, Serialize, }; use utils::map::vec_map::VecMap; @@ -262,20 +263,6 @@ impl JMAP { } // Obtain a documentId and changeId - let document_id = self - .core - .storage - .data - .assign_document_id(params.account_id, Collection::Email) - .await - .map_err(|err| { - tracing::error!( - event = "error", - context = "email_ingest", - error = ?err, - "Failed to assign documentId."); - IngestError::Temporary - })?; let change_id = self .assign_change_id(params.account_id) .await @@ -322,44 +309,26 @@ impl JMAP { // Prepare batch let mut batch = BatchBuilder::new(); - batch.with_account_id(params.account_id); - - // Build change log - let mut changes = ChangeLogBuilder::with_change_id(change_id); - let thread_id = if let Some(thread_id) = thread_id { - changes.log_child_update(Collection::Thread, thread_id); - thread_id + batch + .with_change_id(change_id) + .with_account_id(params.account_id) + .with_collection(Collection::Thread); + if let Some(thread_id) = thread_id { + batch.log(Changes::update([thread_id])); } else { - let thread_id = self - .core - .storage - .data - .assign_document_id(params.account_id, Collection::Thread) - .await - .map_err(|err| { - tracing::error!( - event = "error", - context = "email_ingest", - error = ?err, - "Failed to assign documentId for new thread."); - IngestError::Temporary - })?; - batch - .with_collection(Collection::Thread) - .create_document(thread_id); - changes.log_insert(Collection::Thread, thread_id); - thread_id - }; - let id = Id::from_parts(thread_id, document_id); - changes.log_insert(Collection::Email, id); - for mailbox_id in ¶ms.mailbox_ids { - changes.log_child_update(Collection::Mailbox, *mailbox_id); + batch.create_document().log(LogInsert()); } // Build write batch + let maybe_thread_id = thread_id + .map(MaybeDynamicId::Static) + .unwrap_or(MaybeDynamicId::Dynamic(0)); batch + .with_collection(Collection::Mailbox) + .log(Changes::child_update(params.mailbox_ids.iter().copied())) .with_collection(Collection::Email) - .create_document(document_id) + .create_document() + .log(LogEmailInsert(thread_id)) .index_message( message, blob_id.hash.clone(), @@ -368,16 +337,19 @@ impl JMAP { params.received_at.unwrap_or_else(now), ) .value(Property::Cid, change_id, F_VALUE) - .value(Property::ThreadId, thread_id, F_VALUE | F_BITMAP) - .custom(changes) + .set(Property::ThreadId, maybe_thread_id) + .tag(Property::ThreadId, TagValue::Id(maybe_thread_id), 0) .set( ValueClass::IndexEmail( self.generate_snowflake_id() .map_err(|_| IngestError::Temporary)?, ), - blob_id.hash.clone(), + blob_id.hash.as_ref(), ); - self.core + + // Insert and obtain ids + let ids = self + .core .storage .data .write(batch.build()) @@ -390,6 +362,14 @@ impl JMAP { "Failed to write message to database."); IngestError::Temporary })?; + let thread_id = match thread_id { + Some(thread_id) => thread_id, + None => ids + .first_document_id() + .map_err(|_| IngestError::Temporary)?, + }; + let document_id = ids.last_document_id().map_err(|_| IngestError::Temporary)?; + let id = Id::from_parts(thread_id, document_id); // Request FTS index let _ = self.inner.housekeeper_tx.send(Event::IndexStart).await; @@ -545,7 +525,7 @@ impl JMAP { field: Property::ThreadId.into(), value: TagValue::Id(old_thread_id), }, - block_num: 0, + document_id: 0, }) .await .map_err(|err| { @@ -605,7 +585,33 @@ impl JMAP { .data .write(batch.build()) .await - .map(|v| v.expect("UID next") as u32) + .and_then(|v| v.last_counter_id().map(|id| id as u32)) + } +} + +pub struct LogEmailInsert(Option); + +impl LogEmailInsert { + pub fn new(thread_id: Option) -> Self { + Self(thread_id) + } +} + +impl SerializeWithId for LogEmailInsert { + fn serialize_with_id(&self, ids: &AssignedIds) -> store::Result> { + let thread_id = match self.0 { + Some(thread_id) => thread_id, + None => ids.first_document_id()?, + }; + let document_id = ids.last_document_id()?; + + Ok(Changes::insert([Id::from_parts(thread_id, document_id)]).serialize()) + } +} + +impl From for MaybeDynamicValue { + fn from(log: LogEmailInsert) -> Self { + MaybeDynamicValue::Dynamic(Box::new(log)) } } diff --git a/crates/jmap/src/identity/get.rs b/crates/jmap/src/identity/get.rs index 2181b0d1..1be070e6 100644 --- a/crates/jmap/src/identity/get.rs +++ b/crates/jmap/src/identity/get.rs @@ -159,14 +159,12 @@ impl JMAP { .trim() .to_string(); let has_many = principal.emails.len() > 1; - for email in principal.emails { + for (idx, email) in principal.emails.into_iter().enumerate() { + let document_id = idx as u32; let email = sanitize_email(&email).unwrap_or_default(); if email.is_empty() { continue; } - let identity_id = self - .assign_document_id(account_id, Collection::Identity) - .await?; let name = if name.is_empty() { email.clone() } else if has_many { @@ -174,14 +172,14 @@ impl JMAP { } else { name.clone() }; - batch.create_document(identity_id).value( + batch.create_document_with_id(document_id).value( Property::Value, Object::with_capacity(4) .with_property(Property::Name, name) .with_property(Property::Email, email), F_VALUE, ); - identity_ids.insert(identity_id); + identity_ids.insert(document_id); } self.core .storage diff --git a/crates/jmap/src/identity/set.rs b/crates/jmap/src/identity/set.rs index 8f8de881..686beea9 100644 --- a/crates/jmap/src/identity/set.rs +++ b/crates/jmap/src/identity/set.rs @@ -106,16 +106,13 @@ impl JMAP { // Insert record let mut batch = BatchBuilder::new(); - let document_id = self - .assign_document_id(account_id, Collection::Identity) - .await?; batch .with_account_id(account_id) .with_collection(Collection::Identity) - .create_document(document_id) + .create_document() .value(Property::Value, identity, F_VALUE); + let document_id = self.write_batch_expect_id(batch).await?; identity_ids.insert(document_id); - self.write_batch(batch).await?; changes.log_insert(Collection::Identity, document_id); response.created(id, document_id); } diff --git a/crates/jmap/src/lib.rs b/crates/jmap/src/lib.rs index f8f25775..93602f77 100644 --- a/crates/jmap/src/lib.rs +++ b/crates/jmap/src/lib.rs @@ -48,7 +48,8 @@ use store::{ query::{sort::Pagination, Comparator, Filter, ResultSet, SortedResultSet}, roaring::RoaringBitmap, write::{ - key::DeserializeBigEndian, BatchBuilder, BitmapClass, DirectoryClass, TagValue, ValueClass, + key::DeserializeBigEndian, AssignedIds, BatchBuilder, BitmapClass, DirectoryClass, + TagValue, ValueClass, }, BitmapKey, Deserialize, IterateParams, ValueKey, U32_LEN, }; @@ -174,26 +175,6 @@ impl JMAP { jmap_instance } - pub async fn assign_document_id( - &self, - account_id: u32, - collection: Collection, - ) -> Result { - self.core - .storage - .data - .assign_document_id(account_id, collection) - .await - .map_err(|err| { - tracing::error!( - event = "error", - context = "assign_document_id", - error = ?err, - "Failed to assign documentId."); - MethodError::ServerPartialFail - }) - } - pub async fn get_property( &self, account_id: u32, @@ -322,7 +303,7 @@ impl JMAP { account_id: u32, collection: Collection, property: impl AsRef, - value: impl Into, + value: impl Into>, ) -> Result, MethodError> { let property = property.as_ref(); match self @@ -336,7 +317,7 @@ impl JMAP { field: property.into(), value: value.into(), }, - block_num: 0, + document_id: 0, }) .await { @@ -544,13 +525,12 @@ impl JMAP { Ok(response) } - pub async fn write_batch(&self, batch: BatchBuilder) -> Result<(), MethodError> { + pub async fn write_batch(&self, batch: BatchBuilder) -> Result { self.core .storage .data .write(batch.build()) .await - .map(|_| ()) .map_err(|err| { match err { store::Error::InternalError(err) => { @@ -573,6 +553,20 @@ impl JMAP { } }) } + + pub async fn write_batch_expect_id(&self, batch: BatchBuilder) -> Result { + self.write_batch(batch).await.and_then(|ids| { + ids.last_document_id().map_err(|err| { + tracing::error!( + event = "error", + context = "write_batch_expect_id", + error = ?err, + "Failed to obtain last document id." + ); + MethodError::ServerPartialFail + }) + }) + } } impl From for JMAP { diff --git a/crates/jmap/src/mailbox/mod.rs b/crates/jmap/src/mailbox/mod.rs index b4a0ac1c..b4bfa6b4 100644 --- a/crates/jmap/src/mailbox/mod.rs +++ b/crates/jmap/src/mailbox/mod.rs @@ -24,7 +24,9 @@ use std::slice::Iter; use store::{ - write::{BitmapClass, DeserializeFrom, Operation, SerializeInto, TagValue, ToBitmaps}, + write::{ + BitmapClass, DeserializeFrom, MaybeDynamicId, Operation, SerializeInto, TagValue, ToBitmaps, + }, Serialize, U32_LEN, }; use utils::codec::leb128::{Leb128Iterator, Leb128Vec}; @@ -36,6 +38,8 @@ pub mod set; pub const INBOX_ID: u32 = 0; pub const TRASH_ID: u32 = 1; pub const JUNK_ID: u32 = 2; +pub const DRAFTS_ID: u32 = 3; +pub const SENT_ID: u32 = 4; #[derive(Debug, Clone, Copy)] pub struct UidMailbox { @@ -56,7 +60,7 @@ impl ToBitmaps for UidMailbox { ops.push(Operation::Bitmap { class: BitmapClass::Tag { field, - value: TagValue::Id(self.mailbox_id), + value: TagValue::Id(MaybeDynamicId::Static(self.mailbox_id)), }, set, }); diff --git a/crates/jmap/src/mailbox/query.rs b/crates/jmap/src/mailbox/query.rs index 5f5d9277..5799d459 100644 --- a/crates/jmap/src/mailbox/query.rs +++ b/crates/jmap/src/mailbox/query.rs @@ -223,8 +223,9 @@ impl JMAP { }; while let Some(&id) = it.next() { - jmap_id = id.document_id() + 1; - if children.remove(&jmap_id) { + let next_id = id.document_id() + 1; + if children.remove(&next_id) { + jmap_id = next_id; if !paginate.add(0, id.document_id()) { break 'outer; } else { diff --git a/crates/jmap/src/mailbox/set.rs b/crates/jmap/src/mailbox/set.rs index 50373c59..c83aba1f 100644 --- a/crates/jmap/src/mailbox/set.rs +++ b/crates/jmap/src/mailbox/set.rs @@ -60,6 +60,7 @@ use crate::{ #[allow(unused_imports)] use super::{UidMailbox, INBOX_ID, JUNK_ID, TRASH_ID}; +use super::{DRAFTS_ID, SENT_ID}; struct SetContext<'x> { account_id: u32, @@ -115,9 +116,6 @@ impl JMAP { match self.mailbox_set_item(object, None, &ctx).await? { Ok(builder) => { let mut batch = BatchBuilder::new(); - let document_id = self - .assign_document_id(account_id, Collection::Mailbox) - .await?; batch .with_account_id(account_id) .with_collection(Collection::Mailbox); @@ -133,11 +131,19 @@ impl JMAP { } } - batch.create_document(document_id).custom(builder); - changes.log_insert(Collection::Mailbox, document_id); - ctx.mailbox_ids.insert(document_id); - match self.core.storage.data.write(batch.build()).await { - Ok(_) => { + batch.create_document().custom(builder); + + match self + .core + .storage + .data + .write(batch.build()) + .await + .and_then(|ids| ids.last_document_id()) + { + Ok(document_id) => { + changes.log_insert(Collection::Mailbox, document_id); + ctx.mailbox_ids.insert(document_id); ctx.response.created(id, document_id); } Err(store::Error::AssertValueFailed) => { @@ -833,17 +839,14 @@ impl JMAP { .with_collection(Collection::Mailbox); // Create mailboxes - for (name, role) in [ - ("Inbox", "inbox"), - ("Deleted Items", "trash"), - ("Junk Mail", "junk"), - ("Drafts", "drafts"), - ("Sent Items", "sent"), + for (name, role, document_id) in [ + ("Inbox", "inbox", INBOX_ID), + ("Deleted Items", "trash", TRASH_ID), + ("Junk Mail", "junk", JUNK_ID), + ("Drafts", "drafts", DRAFTS_ID), + ("Sent Items", "sent", SENT_ID), ] { - let mailbox_id = self - .assign_document_id(account_id, Collection::Mailbox) - .await?; - batch.create_document(mailbox_id).custom( + batch.create_document_with_id(document_id).custom( ObjectIndexBuilder::new(SCHEMA).with_changes( Object::with_capacity(4) .with_property(Property::Name, name) @@ -855,7 +858,7 @@ impl JMAP { ), ), ); - mailbox_ids.insert(mailbox_id); + mailbox_ids.insert(document_id); } self.core .storage @@ -901,36 +904,42 @@ impl JMAP { // Create missing folders if path.peek().is_some() { - let mut batch = BatchBuilder::new(); let mut changes = self.begin_changes(account_id).await?; - batch - .with_account_id(account_id) - .with_collection(Collection::Mailbox); for name in path { if name.len() > self.core.jmap.mailbox_name_max_len { return Ok(None); } - - let document_id = self - .assign_document_id(account_id, Collection::Mailbox) - .await?; - batch.create_document(document_id).custom( - ObjectIndexBuilder::new(SCHEMA).with_changes( - Object::with_capacity(3) - .with_property(Property::Name, name) - .with_property(Property::ParentId, Value::Id(Id::from(next_parent_id))) - .with_property( - Property::Cid, - Value::UnsignedInt(rand::random::() as u64), - ), - ), - ); + let mut batch = BatchBuilder::new(); + batch + .with_account_id(account_id) + .with_collection(Collection::Mailbox) + .create_document() + .custom( + ObjectIndexBuilder::new(SCHEMA).with_changes( + Object::with_capacity(3) + .with_property(Property::Name, name) + .with_property( + Property::ParentId, + Value::Id(Id::from(next_parent_id)), + ) + .with_property( + Property::Cid, + Value::UnsignedInt(rand::random::() as u64), + ), + ), + ); + let document_id = self.write_batch_expect_id(batch).await?; changes.log_insert(Collection::Mailbox, document_id); next_parent_id = document_id + 1; } let change_id = changes.change_id; - batch.custom(changes); + let mut batch = BatchBuilder::new(); + + batch + .with_account_id(account_id) + .with_collection(Collection::Mailbox) + .custom(changes); self.write_batch(batch).await?; Ok(Some((next_parent_id - 1, Some(change_id)))) diff --git a/crates/jmap/src/push/set.rs b/crates/jmap/src/push/set.rs index 42d92206..bb8f5e65 100644 --- a/crates/jmap/src/push/set.rs +++ b/crates/jmap/src/push/set.rs @@ -121,16 +121,13 @@ impl JMAP { // Insert record let mut batch = BatchBuilder::new(); - let document_id = self - .assign_document_id(account_id, Collection::PushSubscription) - .await?; batch .with_account_id(account_id) .with_collection(Collection::PushSubscription) - .create_document(document_id) + .create_document() .value(Property::Value, push, F_VALUE); + let document_id = self.write_batch_expect_id(batch).await?; push_ids.insert(document_id); - self.write_batch(batch).await?; response.created.insert( id, Object::with_capacity(1) diff --git a/crates/jmap/src/services/index.rs b/crates/jmap/src/services/index.rs index 350c2c49..eef9014c 100644 --- a/crates/jmap/src/services/index.rs +++ b/crates/jmap/src/services/index.rs @@ -44,13 +44,13 @@ struct IndexEmail { impl JMAP { pub async fn fts_index_queued(&self) { - let from_key = ValueKey:: { + let from_key = ValueKey::> { account_id: 0, collection: 0, document_id: 0, class: ValueClass::IndexEmail(0), }; - let to_key = ValueKey:: { + let to_key = ValueKey::> { account_id: u32::MAX, collection: u8::MAX, document_id: u32::MAX, diff --git a/crates/jmap/src/sieve/get.rs b/crates/jmap/src/sieve/get.rs index 8a2a1a7b..c17d0d79 100644 --- a/crates/jmap/src/sieve/get.rs +++ b/crates/jmap/src/sieve/get.rs @@ -33,7 +33,7 @@ use sieve::Sieve; use store::{ query::Filter, write::{assert::HashedValue, BatchBuilder, Bincode, BlobOp}, - Deserialize, Serialize, + BlobClass, Deserialize, Serialize, }; use crate::{sieve::SeenIds, JMAP}; @@ -99,9 +99,25 @@ impl JMAP { Property::Id => { result.append(Property::Id, Value::Id(id)); } - Property::Name | Property::BlobId | Property::IsActive => { + Property::Name | Property::IsActive => { result.append(property.clone(), push.remove(property)); } + Property::BlobId => { + result.append( + Property::BlobId, + match push.remove(&Property::BlobId) { + Value::BlobId(mut blob_id) => { + blob_id.class = BlobClass::Linked { + account_id, + collection: Collection::SieveScript.into(), + document_id, + }; + Value::BlobId(blob_id) + } + other => other, + }, + ); + } property => { result.append(property.clone(), Value::Null); } diff --git a/crates/jmap/src/sieve/set.rs b/crates/jmap/src/sieve/set.rs index c2f254d6..946f58c1 100644 --- a/crates/jmap/src/sieve/set.rs +++ b/crates/jmap/src/sieve/set.rs @@ -100,28 +100,18 @@ impl JMAP { if sieve_ids.len() as usize <= self.core.jmap.sieve_max_scripts { match self.sieve_set_item(object, None, &ctx).await? { Ok((mut builder, Some(blob))) => { - // Obtain document id - let document_id = self - .assign_document_id(account_id, Collection::SieveScript) - .await?; - // Store blob let blob_id = builder.changes_mut().unwrap().blob_id_mut().unwrap(); blob_id.hash = self.put_blob(account_id, &blob, false).await?.hash; - blob_id.class = BlobClass::Linked { - account_id, - collection: Collection::SieveScript.into(), - document_id, - }; let script_size = blob_id.section.as_ref().unwrap().size; - let blob_id = blob_id.clone(); + let mut blob_id = blob_id.clone(); // Write record let mut batch = BatchBuilder::new(); batch .with_account_id(account_id) .with_collection(Collection::SieveScript) - .create_document(document_id) + .create_document() .add(DirectoryClass::UsedQuota(account_id), script_size as i64) .set( BlobOp::Link { @@ -130,11 +120,17 @@ impl JMAP { Vec::new(), ) .custom(builder); + + let document_id = self.write_batch_expect_id(batch).await?; sieve_ids.insert(document_id); - self.write_batch(batch).await?; changes.log_insert(Collection::SieveScript, document_id); // Add result with updated blobId + blob_id.class = BlobClass::Linked { + account_id, + collection: Collection::SieveScript.into(), + document_id, + }; ctx.response.created.insert( id, Object::with_capacity(1) @@ -206,11 +202,11 @@ impl JMAP { // Store blob let blob_id = builder.changes_mut().unwrap().blob_id_mut().unwrap(); blob_id.hash = self.put_blob(account_id, &blob, false).await?.hash; - blob_id.class = BlobClass::Linked { + /*blob_id.class = BlobClass::Linked { account_id, collection: Collection::SieveScript.into(), document_id, - }; + };*/ let script_size = blob_id.section.as_ref().unwrap().size as i64; let prev_script_size = prev_blob_id.section.as_ref().unwrap().size as i64; diff --git a/crates/jmap/src/submission/set.rs b/crates/jmap/src/submission/set.rs index 68ad947d..7cb9adbc 100644 --- a/crates/jmap/src/submission/set.rs +++ b/crates/jmap/src/submission/set.rs @@ -95,15 +95,12 @@ impl JMAP { // Insert record let mut batch = BatchBuilder::new(); - let document_id = self - .assign_document_id(account_id, Collection::EmailSubmission) - .await?; batch .with_account_id(account_id) .with_collection(Collection::EmailSubmission) - .create_document(document_id) + .create_document() .custom(ObjectIndexBuilder::new(SCHEMA).with_changes(submission)); - self.write_batch(batch).await?; + let document_id = self.write_batch_expect_id(batch).await?; changes.log_insert(Collection::EmailSubmission, document_id); response.created(id, document_id); } diff --git a/crates/jmap/src/vacation/set.rs b/crates/jmap/src/vacation/set.rs index a9730da5..186100bf 100644 --- a/crates/jmap/src/vacation/set.rs +++ b/crates/jmap/src/vacation/set.rs @@ -41,12 +41,10 @@ use jmap_proto::{ }; use mail_builder::MessageBuilder; use mail_parser::decoders::html::html_to_text; -use store::{ - write::{ - assert::HashedValue, log::ChangeLogBuilder, BatchBuilder, BlobOp, DirectoryClass, F_CLEAR, - F_VALUE, - }, - BlobClass, +use store::write::{ + assert::HashedValue, + log::{Changes, LogInsert}, + BatchBuilder, BlobOp, DirectoryClass, F_CLEAR, F_VALUE, }; use crate::{ @@ -121,8 +119,15 @@ impl JMAP { } } + // Prepare write batch + let mut batch = BatchBuilder::new(); + let change_id = self.assign_change_id(account_id).await?; + batch + .with_change_id(change_id) + .with_account_id(account_id) + .with_collection(Collection::SieveScript); + // Process changes - let mut change_log = ChangeLogBuilder::new(); if let Some(changes_) = changes { // Parse properties let mut changes = Object::with_capacity(changes_.properties.len()); @@ -199,12 +204,6 @@ impl JMAP { } } - // Prepare write batch - let mut batch = BatchBuilder::new(); - batch - .with_account_id(account_id) - .with_collection(Collection::SieveScript); - // Obtain current script let document_id = self.get_vacation_sieve_script_id(account_id).await?; let mut was_active = false; @@ -231,20 +230,14 @@ impl JMAP { .with_changes(changes); // Update id - let document_id = if let Some(document_id) = document_id { + if let Some(document_id) = document_id { batch .update_document(document_id) - .value(Property::EmailIds, (), F_VALUE | F_CLEAR); - change_log.log_insert(Collection::SieveScript, document_id); - document_id + .value(Property::EmailIds, (), F_VALUE | F_CLEAR) + .log(Changes::update([document_id])); } else { - let document_id = self - .assign_document_id(account_id, Collection::SieveScript) - .await?; - batch.create_document(document_id); - change_log.log_update(Collection::SieveScript, document_id); - document_id - }; + batch.create_document().log(LogInsert()); + } // Create sieve script only if there are changes if build_script { @@ -255,11 +248,11 @@ impl JMAP { .hash; let blob_id = obj.changes_mut().unwrap().blob_id_mut().unwrap(); blob_id.hash = hash; - blob_id.class = BlobClass::Linked { + /*blob_id.class = BlobClass::Linked { account_id, collection: Collection::SieveScript.into(), - document_id, - }; + document_id: u32::MAX, + };*/ // Link blob batch.set( @@ -305,9 +298,18 @@ impl JMAP { // Write changes batch.custom(obj); - if !batch.is_empty() { - self.write_batch(batch).await?; - } + let document_id = if !batch.is_empty() { + let ids = self.write_batch(batch).await?; + response.new_state = Some(change_id.into()); + match document_id { + Some(document_id) => document_id, + None => ids + .last_document_id() + .map_err(|_| MethodError::ServerPartialFail)?, + } + } else { + document_id.unwrap_or(u32::MAX) + }; // Deactivate other sieve scripts if !was_active && is_active { @@ -331,7 +333,7 @@ impl JMAP { { self.sieve_script_delete(account_id, document_id, false) .await?; - change_log.log_delete(Collection::SieveScript, document_id); + batch.log(Changes::delete([document_id])); response.destroyed.push(id); continue; } @@ -339,11 +341,12 @@ impl JMAP { response.not_destroyed.append(id, SetError::not_found()); } - } - // Write changes - if !change_log.is_empty() { - response.new_state = Some(self.commit_changes(account_id, change_log).await?.into()); + // Write changes + if !batch.is_empty() { + self.write_batch(batch).await?; + response.new_state = Some(change_id.into()); + } } Ok(response) diff --git a/crates/main/Cargo.toml b/crates/main/Cargo.toml index b34847bf..a37385b3 100644 --- a/crates/main/Cargo.toml +++ b/crates/main/Cargo.toml @@ -32,7 +32,8 @@ tracing = "0.1" jemallocator = "0.5.0" [features] -default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis"] +#default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis"] +default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis", "foundationdb"] sqlite = ["store/sqlite"] foundationdb = ["store/foundation"] postgres = ["store/postgres"] diff --git a/crates/managesieve/src/op/putscript.rs b/crates/managesieve/src/op/putscript.rs index 452ade76..e41bc809 100644 --- a/crates/managesieve/src/op/putscript.rs +++ b/crates/managesieve/src/op/putscript.rs @@ -30,7 +30,7 @@ use jmap_proto::{ use sieve::compiler::ErrorType; use store::{ query::Filter, - write::{assert::HashedValue, BatchBuilder, BlobOp, DirectoryClass}, + write::{assert::HashedValue, log::LogInsert, BatchBuilder, BlobOp, DirectoryClass}, BlobClass, }; use tokio::io::{AsyncRead, AsyncWrite}; @@ -165,12 +165,6 @@ impl Session { ); self.jmap.write_batch(batch).await?; } else { - // Obtain document id - let document_id = self - .jmap - .assign_document_id(account_id, Collection::SieveScript) - .await?; - // Write script blob let blob_id = BlobId::new( self.jmap @@ -180,19 +174,18 @@ impl Session { BlobClass::Linked { account_id, collection: Collection::SieveScript.into(), - document_id, + document_id: 0, }, ) .with_section_size(script_size as usize); // Write record - let mut changelog = self.jmap.begin_changes(account_id).await?; - changelog.log_insert(Collection::SieveScript, document_id); let mut batch = BatchBuilder::new(); batch .with_account_id(account_id) .with_collection(Collection::SieveScript) - .create_document(document_id) + .create_document() + .log(LogInsert()) .add(DirectoryClass::UsedQuota(account_id), script_size) .set( BlobOp::Link { @@ -207,8 +200,7 @@ impl Session { .with_property(Property::IsActive, Value::Bool(false)) .with_property(Property::BlobId, Value::BlobId(blob_id)), ), - ) - .custom(changelog); + ); self.jmap.write_batch(batch).await?; } diff --git a/crates/store/src/backend/foundationdb/read.rs b/crates/store/src/backend/foundationdb/read.rs index 0b72aaa9..88fd779d 100644 --- a/crates/store/src/backend/foundationdb/read.rs +++ b/crates/store/src/backend/foundationdb/read.rs @@ -32,7 +32,6 @@ use roaring::RoaringBitmap; use crate::{ backend::deserialize_i64_le, write::{ - bitmap::DeserializeBlock, key::{DeserializeBigEndian, KeySerializer}, BitmapClass, ValueClass, }, @@ -41,13 +40,6 @@ use crate::{ use super::{FdbStore, MAX_VALUE_SIZE}; -#[cfg(feature = "fdb-chunked-bm")] -pub(crate) enum ChunkedBitmap { - Single(RoaringBitmap), - Chunked { n_chunks: u8, bitmap: RoaringBitmap }, - None, -} - #[allow(dead_code)] pub(crate) enum ChunkedValue { Single(FdbSlice), @@ -72,47 +64,35 @@ impl FdbStore { pub(crate) async fn get_bitmap( &self, - mut key: BitmapKey, + mut key: BitmapKey>, ) -> crate::Result> { - #[cfg(feature = "fdb-chunked-bm")] - { - read_chunked_bitmap(&key.serialize(WITH_SUBSPACE), &self.db.create_trx()?, true) - .await - .map(Into::into) - } + let mut bm = RoaringBitmap::new(); + let begin = key.serialize(WITH_SUBSPACE); + key.document_id = u32::MAX; + let end = key.serialize(WITH_SUBSPACE); + let key_len = begin.len(); + let trx = self.db.create_trx()?; + let mut values = trx.get_ranges( + RangeOption { + begin: KeySelector::first_greater_or_equal(begin), + end: KeySelector::first_greater_or_equal(end), + mode: StreamingMode::WantAll, + reverse: false, + ..RangeOption::default() + }, + true, + ); - #[cfg(not(feature = "fdb-chunked-bm"))] - { - let mut bm = RoaringBitmap::new(); - let begin = key.serialize(WITH_SUBSPACE); - key.block_num = u32::MAX; - let end = key.serialize(WITH_SUBSPACE); - let key_len = begin.len(); - let trx = self.db.create_trx()?; - let mut values = trx.get_ranges( - RangeOption { - begin: KeySelector::first_greater_or_equal(begin), - end: KeySelector::first_greater_or_equal(end), - mode: StreamingMode::WantAll, - reverse: false, - ..RangeOption::default() - }, - true, - ); - - while let Some(values) = values.next().await { - for value in values? { - let key = value.key(); - if key.len() == key_len { - bm.deserialize_block( - value.value(), - key.deserialize_be_u32(key.len() - U32_LEN)?, - ); - } + while let Some(values) = values.next().await { + for value in values? { + let key = value.key(); + if key.len() == key_len { + bm.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); } } - Ok(if !bm.is_empty() { Some(bm) } else { None }) } + + Ok(if !bm.is_empty() { Some(bm) } else { None }) } pub(crate) async fn iterate( @@ -155,7 +135,7 @@ impl FdbStore { pub(crate) async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { let key = key.into().serialize(WITH_SUBSPACE); if let Some(bytes) = self.db.create_trx()?.get(&key, true).await? { diff --git a/crates/store/src/backend/foundationdb/write.rs b/crates/store/src/backend/foundationdb/write.rs index 9d211d79..82a00f72 100644 --- a/crates/store/src/backend/foundationdb/write.rs +++ b/crates/store/src/backend/foundationdb/write.rs @@ -26,23 +26,22 @@ use std::{ time::{Duration, Instant}, }; -use ahash::AHashMap; use foundationdb::{ - options::{self, MutationType}, + options::{self, MutationType, StreamingMode}, FdbError, KeySelector, RangeOption, }; use futures::StreamExt; use rand::Rng; +use roaring::RoaringBitmap; use crate::{ backend::deserialize_i64_le, write::{ - bitmap::{block_contains, DenseBitmap}, - key::KeySerializer, - Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, + key::{DeserializeBigEndian, KeySerializer}, + AssignedIds, Batch, BitmapClass, Operation, RandomAvailableId, ValueOp, + MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, }, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_BITMAPS, SUBSPACE_COUNTERS, - SUBSPACE_VALUES, WITH_SUBSPACE, + BitmapKey, IndexKey, Key, LogKey, SUBSPACE_COUNTERS, SUBSPACE_VALUES, U32_LEN, WITH_SUBSPACE, }; use super::{ @@ -50,41 +49,16 @@ use super::{ FdbStore, MAX_VALUE_SIZE, }; -#[cfg(feature = "fdb-chunked-bm")] -use super::read::{read_chunked_bitmap, ChunkedBitmap}; - -#[cfg(feature = "fdb-chunked-bm")] -use roaring::RoaringBitmap; - -#[cfg(feature = "fdb-chunked-bm")] -struct BitmapOp { - document_id: u32, - set: bool, -} - -#[cfg(feature = "fdb-chunked-bm")] -impl BitmapOp { - fn new(document_id: u32, set: bool) -> Self { - Self { document_id, set } - } -} - impl FdbStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result { let start = Instant::now(); let mut retry_count = 0; - #[cfg(not(feature = "fdb-chunked-bm"))] - let mut set_bitmaps = AHashMap::new(); - #[cfg(not(feature = "fdb-chunked-bm"))] - let mut clear_bitmaps = AHashMap::new(); - #[cfg(feature = "fdb-chunked-bm")] - let mut bitmaps = AHashMap::new(); loop { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; - let mut result = None; + let mut result = AssignedIds::default(); let trx = self.db.create_trx()?; @@ -106,17 +80,18 @@ impl FdbStore { document_id = *document_id_; } Operation::Value { class, op } => { - let mut key = ValueKey { + let mut key = class.serialize( account_id, collection, document_id, - class, - } - .serialize(WITH_SUBSPACE); + WITH_SUBSPACE, + (&result).into(), + ); let do_chunk = key[0] == SUBSPACE_VALUES; match op { ValueOp::Set(value) => { + let value = value.resolve(&result)?; if !value.is_empty() && do_chunk { for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() { match pos.cmp(&1) { @@ -138,29 +113,7 @@ impl FdbStore { trx.set(&key, chunk); } } else { - trx.set(&key, value); - } - - if matches!(class, ValueClass::ReservedId) { - let block_num = DenseBitmap::block_num(document_id); - if let Ok(Some(bytes)) = trx - .get( - &BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num, - } - .serialize(WITH_SUBSPACE), - true, - ) - .await - { - if block_contains(&bytes, block_num, document_id) { - trx.cancel(); - return Err(crate::Error::AssertValueFailed); - } - } + trx.set(&key, value.as_ref()); } } ValueOp::AtomicAdd(by) => { @@ -173,7 +126,7 @@ impl FdbStore { *by }; trx.set(&key, &num.to_le_bytes()[..]); - result = Some(num); + result.push_counter_id(num); } ValueOp::Clear => { if do_chunk { @@ -207,64 +160,100 @@ impl FdbStore { } } Operation::Bitmap { class, set } => { - if retry_count == 0 { - #[cfg(not(feature = "fdb-chunked-bm"))] - if *set { - &mut set_bitmaps - } else { - &mut clear_bitmaps + // Find the next available document id + let assign_id = *set + && matches!(class, BitmapClass::DocumentIds) + && document_id == u32::MAX; + if assign_id { + let begin = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: 0, } - .entry( - BitmapKey { - account_id, - collection, - class, - block_num: DenseBitmap::block_num(document_id), + .serialize(WITH_SUBSPACE); + let end = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: u32::MAX, + } + .serialize(WITH_SUBSPACE); + let key_len = begin.len(); + let mut values = trx.get_ranges( + RangeOption { + begin: KeySelector::first_greater_or_equal(begin), + end: KeySelector::first_greater_or_equal(end), + mode: StreamingMode::WantAll, + reverse: false, + ..RangeOption::default() + }, + true, + ); + let mut found_ids = RoaringBitmap::new(); + while let Some(values) = values.next().await { + for value in values? { + let key = value.key(); + if key.len() == key_len { + found_ids + .insert(key.deserialize_be_u32(key_len - U32_LEN)?); + } else { + break; + } } - .serialize(WITH_SUBSPACE), - ) - .or_insert_with(DenseBitmap::empty) - .set(document_id); + } + document_id = found_ids.random_available_id(); + result.push_document_id(document_id); + } - #[cfg(feature = "fdb-chunked-bm")] - bitmaps - .entry( - BitmapKey { + let key = class.serialize( + account_id, + collection, + document_id, + WITH_SUBSPACE, + (&result).into(), + ); + + if *set { + if assign_id { + trx.add_conflict_range( + &key, + &class.serialize( account_id, collection, - class, - block_num: 0, - } - .serialize(WITH_SUBSPACE), - ) - .or_insert(Vec::new()) - .push(BitmapOp::new(document_id, *set)); + document_id + 1, + WITH_SUBSPACE, + (&result).into(), + ), + options::ConflictRangeType::Read, + )?; + } + + trx.set(&key, &[]); + } else { + trx.clear(&key); } } - Operation::Log { - collection, - change_id, - set, - } => { + Operation::Log { set } => { let key = LogKey { account_id, - collection: *collection, - change_id: *change_id, + collection, + change_id: batch.change_id, } .serialize(WITH_SUBSPACE); - trx.set(&key, set); + trx.set(&key, set.resolve(&result)?.as_ref()); } Operation::AssertValue { class, assert_value, } => { - let key = ValueKey { + let key = class.serialize( account_id, collection, document_id, - class, - } - .serialize(WITH_SUBSPACE); + WITH_SUBSPACE, + (&result).into(), + ); let matches = match read_chunked_value(&key, &trx, false).await { Ok(ChunkedValue::Single(bytes)) => assert_value.matches(bytes.as_ref()), @@ -283,100 +272,6 @@ impl FdbStore { } } - #[cfg(not(feature = "fdb-chunked-bm"))] - { - for (key, bitmap) in &set_bitmaps { - trx.atomic_op(key, &bitmap.bitmap, MutationType::BitOr); - } - - for (key, bitmap) in &clear_bitmaps { - trx.atomic_op(key, &bitmap.bitmap, MutationType::BitXor); - } - } - - // Write bitmaps - #[cfg(feature = "fdb-chunked-bm")] - for (key, bitmap_ops) in &bitmaps { - let (mut bitmap, exists, n_chunks) = - match read_chunked_bitmap(key, &trx, false).await? { - ChunkedBitmap::Single(bitmap) => (bitmap, true, 0u8), - ChunkedBitmap::Chunked { n_chunks, bitmap } => (bitmap, true, n_chunks), - ChunkedBitmap::None => (RoaringBitmap::new(), false, 0u8), - }; - - for bitmap_op in bitmap_ops { - if bitmap_op.set { - bitmap.insert(bitmap_op.document_id); - } else { - bitmap.remove(bitmap_op.document_id); - } - } - - if !bitmap.is_empty() { - let mut bytes = Vec::with_capacity(bitmap.serialized_size()); - bitmap.serialize_into(&mut bytes).map_err(|_| { - crate::Error::InternalError("Failed to serialize bitmap".into()) - })?; - let mut key = KeySerializer::new(key.len() + 1) - .write(key.as_slice()) - .finalize(); - let mut chunk_diff = n_chunks; - - for (pos, chunk) in bytes.chunks(MAX_VALUE_SIZE).enumerate() { - match pos.cmp(&1) { - Ordering::Less => {} - Ordering::Equal => { - key.push(0); - if n_chunks > 0 { - chunk_diff -= 1; - } - } - Ordering::Greater => { - if pos < u8::MAX as usize { - *key.last_mut().unwrap() += 1; - if n_chunks > 0 { - chunk_diff -= 1; - } - } else { - trx.cancel(); - return Err(crate::Error::InternalError( - "Bitmap value too large".into(), - )); - } - } - } - trx.set(&key, chunk); - } - - // Delete any additional chunks - if chunk_diff > 0 { - let mut key = KeySerializer::new(key.len() + 1) - .write(key.as_slice()) - .write(0u8) - .finalize(); - for chunk in (0..n_chunks).rev().take(chunk_diff as usize) { - *key.last_mut().unwrap() = chunk; - trx.clear(&key); - } - } - } else if exists { - // Delete main key - trx.clear(key); - - // Delete additional chunked keys - if n_chunks > 0 { - let mut key = KeySerializer::new(key.len() + 1) - .write(key.as_slice()) - .write(0u8) - .finalize(); - for chunk in 0..n_chunks { - *key.last_mut().unwrap() = chunk; - trx.clear(&key); - } - } - } - } - match trx.commit().await { Ok(_) => { return Ok(result); @@ -396,38 +291,8 @@ impl FdbStore { } pub(crate) async fn purge_store(&self) -> crate::Result<()> { - // Obtain all empty bitmaps - let trx = self.db.create_trx()?; - let mut iter = trx.get_ranges( - RangeOption { - begin: KeySelector::first_greater_or_equal(&[SUBSPACE_BITMAPS, 0u8][..]), - end: KeySelector::first_greater_or_equal( - &[ - SUBSPACE_BITMAPS, - u8::MAX, - u8::MAX, - u8::MAX, - u8::MAX, - u8::MAX, - ][..], - ), - mode: options::StreamingMode::WantAll, - reverse: false, - ..Default::default() - }, - true, - ); - let mut delete_keys = Vec::new(); - - while let Some(values) = iter.next().await { - for value in values? { - if value.value().iter().all(|byte| *byte == 0) { - delete_keys.push(value.key().to_vec()); - } - } - } - // Obtain all zero counters + let mut delete_keys = Vec::new(); let trx = self.db.create_trx()?; let mut iter = trx.get_ranges( RangeOption { @@ -462,22 +327,13 @@ impl FdbStore { } // Delete keys - let bitmap = DenseBitmap::empty(); let integer = 0i64.to_le_bytes(); for chunk in delete_keys.chunks(1024) { let mut retry_count = 0; loop { let trx = self.db.create_trx()?; for key in chunk { - trx.atomic_op( - key, - if key[0] == SUBSPACE_BITMAPS { - &bitmap.bitmap - } else { - &integer - }, - MutationType::CompareAndClear, - ); + trx.atomic_op(key, &integer, MutationType::CompareAndClear); } match trx.commit().await { Ok(_) => { diff --git a/crates/store/src/backend/mod.rs b/crates/store/src/backend/mod.rs index eb4fa00e..cf0b5cfc 100644 --- a/crates/store/src/backend/mod.rs +++ b/crates/store/src/backend/mod.rs @@ -43,12 +43,6 @@ pub mod sqlite; pub const MAX_TOKEN_LENGTH: usize = (u8::MAX >> 1) as usize; pub const MAX_TOKEN_MASK: usize = MAX_TOKEN_LENGTH - 1; -#[cfg(feature = "test_mode")] -pub static ID_ASSIGNMENT_EXPIRY: std::sync::atomic::AtomicU64 = - std::sync::atomic::AtomicU64::new(60 * 60); // seconds -#[cfg(not(feature = "test_mode"))] -pub const ID_ASSIGNMENT_EXPIRY: u64 = 60 * 60; // seconds - impl From for crate::Error { fn from(err: std::io::Error) -> Self { Self::InternalError(format!("IO error: {}", err)) diff --git a/crates/store/src/backend/mysql/read.rs b/crates/store/src/backend/mysql/read.rs index 0e6e2637..01bbc715 100644 --- a/crates/store/src/backend/mysql/read.rs +++ b/crates/store/src/backend/mysql/read.rs @@ -59,10 +59,10 @@ impl MysqlStore { pub(crate) async fn get_bitmap( &self, - mut key: BitmapKey, + mut key: BitmapKey>, ) -> crate::Result> { let begin = key.serialize(0); - key.block_num = u32::MAX; + key.document_id = u32::MAX; let key_len = begin.len(); let end = key.serialize(0); let mut conn = self.conn_pool.get_conn().await?; @@ -142,7 +142,7 @@ impl MysqlStore { pub(crate) async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { let key = key.into().serialize(0); let mut conn = self.conn_pool.get_conn().await?; diff --git a/crates/store/src/backend/mysql/write.rs b/crates/store/src/backend/mysql/write.rs index d2046eac..67bbbd38 100644 --- a/crates/store/src/backend/mysql/write.rs +++ b/crates/store/src/backend/mysql/write.rs @@ -24,20 +24,30 @@ use std::time::{Duration, Instant}; use ahash::AHashMap; -use mysql_async::{params, prelude::Queryable, Conn, Error, IsolationLevel, Row, TxOpts}; +use futures::TryStreamExt; +use mysql_async::{params, prelude::Queryable, Conn, Error, IsolationLevel, TxOpts}; use rand::Rng; +use roaring::RoaringBitmap; use crate::{ write::{ - Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, + key::DeserializeBigEndian, AssignedIds, Batch, BitmapClass, Operation, RandomAvailableId, + ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, }, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, + BitmapKey, IndexKey, Key, LogKey, SUBSPACE_COUNTERS, U32_LEN, }; use super::MysqlStore; +#[derive(Debug)] +enum CommitError { + Mysql(mysql_async::Error), + Internal(crate::Error), + Retry, +} + impl MysqlStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result { let start = Instant::now(); let mut retry_count = 0; let mut conn = self.conn_pool.get_conn().await?; @@ -45,29 +55,32 @@ impl MysqlStore { loop { match self.write_trx(&mut conn, &batch).await { Ok(result) => { - return result; + return Ok(result); } - Err(Error::Server(err)) + Err(CommitError::Mysql(Error::Server(err))) if [1062, 1213].contains(&err.code) && retry_count < MAX_COMMIT_ATTEMPTS - && start.elapsed() < MAX_COMMIT_TIME => - { - let backoff = rand::thread_rng().gen_range(50..=300); - tokio::time::sleep(Duration::from_millis(backoff)).await; - retry_count += 1; + && start.elapsed() < MAX_COMMIT_TIME => {} + Err(CommitError::Retry) => { + if retry_count > MAX_COMMIT_ATTEMPTS || start.elapsed() > MAX_COMMIT_TIME { + return Err(crate::Error::AssertValueFailed); + } } - Err(err) => { + Err(CommitError::Mysql(err)) => { return Err(err.into()); } + Err(CommitError::Internal(err)) => { + return Err(err); + } } + + let backoff = rand::thread_rng().gen_range(50..=300); + tokio::time::sleep(Duration::from_millis(backoff)).await; + retry_count += 1; } } - async fn write_trx( - &self, - conn: &mut Conn, - batch: &Batch, - ) -> Result>, mysql_async::Error> { + async fn write_trx(&self, conn: &mut Conn, batch: &Batch) -> Result { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; @@ -77,7 +90,7 @@ impl MysqlStore { .with_consistent_snapshot(false) .with_isolation_level(IsolationLevel::ReadCommitted); let mut trx = conn.start_transaction(tx_opts).await?; - let mut result = None; + let mut result = AssignedIds::default(); for op in &batch.ops { match op { @@ -97,14 +110,9 @@ impl MysqlStore { document_id = *document_id_; } Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + let table = char::from(class.subspace(collection)); match op { ValueOp::Set(value) => { @@ -128,32 +136,22 @@ impl MysqlStore { .await? }; - match trx.exec_drop(&s, params! {"k" => key, "v" => value}).await { + match trx + .exec_drop( + &s, + params! {"k" => key, "v" => value.resolve(&result)?.as_ref()}, + ) + .await + { Ok(_) => { if exists.is_some() && trx.affected_rows() == 0 { trx.rollback().await?; - return Ok(Err(crate::Error::AssertValueFailed)); + return Err(crate::Error::AssertValueFailed.into()); } } Err(err) => { trx.rollback().await?; - return Err(err); - } - } - - if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let s = trx.prep("SELECT 1 FROM b WHERE k = ?").await?; - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: document_id, - } - .serialize(0); - if trx.exec_first::(&s, (key,)).await?.is_some() { - trx.rollback().await?; - return Ok(Err(crate::Error::AssertValueFailed)); + return Err(err.into()); } } } @@ -180,18 +178,16 @@ impl MysqlStore { .await?; trx.exec_drop(&s, params! {"k" => key, "v" => by}).await?; let s = trx.prep("SELECT LAST_INSERT_ID()").await?; - result = trx - .exec_first::(&s, ()) - .await? - .ok_or_else(|| { + result.push_counter_id( + trx.exec_first::(&s, ()).await?.ok_or_else(|| { mysql_async::Error::Io(mysql_async::IoError::Io( std::io::Error::new( std::io::ErrorKind::Other, "LAST_INSERT_ID() did not return a value", ), )) - }) - .map(Some)?; + })?, + ); } ValueOp::Clear => { let s = trx @@ -219,16 +215,45 @@ impl MysqlStore { trx.exec_drop(&s, (key,)).await?; } Operation::Bitmap { class, set } => { - let key = BitmapKey { - account_id, - collection, - class, - block_num: document_id, + // Find the next available document id + let is_document_id = matches!(class, BitmapClass::DocumentIds); + if *set && is_document_id && document_id == u32::MAX { + let begin = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: 0, + } + .serialize(0); + let end = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: u32::MAX, + } + .serialize(0); + let key_len = begin.len(); + + let s = trx.prep("SELECT k FROM b WHERE k >= ? AND k <= ?").await?; + let mut rows = trx.exec_stream::, _, _>(&s, (begin, end)).await?; + let mut found_ids = RoaringBitmap::new(); + + while let Some(key) = rows.try_next().await? { + if key.len() == key_len { + found_ids.insert( + key.as_slice().deserialize_be_u32(key.len() - U32_LEN)?, + ); + } + } + + document_id = found_ids.random_available_id(); + result.push_document_id(document_id); } - .serialize(0); + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); let s = if *set { - if matches!(class, BitmapClass::DocumentIds) { + if is_document_id { trx.prep("INSERT INTO b (k) VALUES (?)").await? } else { trx.prep("INSERT IGNORE INTO b (k) VALUES (?)").await? @@ -236,37 +261,42 @@ impl MysqlStore { } else { trx.prep("DELETE FROM b WHERE k = ?").await? }; - trx.exec_drop(&s, (key,)).await?; + + if let Err(err) = trx.exec_drop(&s, (key,)).await { + return Err( + if is_document_id + && matches!(&err, Error::Server(err) if [1062, 1213].contains(&err.code)) + { + trx.rollback().await?; + CommitError::Retry + } else { + CommitError::Mysql(err) + }, + ); + } } - Operation::Log { - collection, - change_id, - set, - } => { + Operation::Log { set } => { let key = LogKey { account_id, - collection: *collection, - change_id: *change_id, + collection, + change_id: batch.change_id, } .serialize(0); let s = trx .prep("INSERT INTO l (k, v) VALUES (?, ?) ON DUPLICATE KEY UPDATE v = VALUES(v)") .await?; - trx.exec_drop(&s, (key, set)).await?; + + trx.exec_drop(&s, (key, set.resolve(&result)?.as_ref())) + .await?; } Operation::AssertValue { class, assert_value, } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + let table = char::from(class.subspace(collection)); let s = trx .prep(&format!("SELECT v FROM {} WHERE k = ? FOR UPDATE", table)) @@ -278,14 +308,14 @@ impl MysqlStore { .unwrap_or_else(|| (false, assert_value.is_none())); if !matches { trx.rollback().await?; - return Ok(Err(crate::Error::AssertValueFailed)); + return Err(crate::Error::AssertValueFailed.into()); } asserted_values.insert(key, exists); } } } - trx.commit().await.map(|_| Ok(result)) + trx.commit().await.map(|_| result).map_err(Into::into) } pub(crate) async fn purge_store(&self) -> crate::Result<()> { @@ -314,3 +344,15 @@ impl MysqlStore { .map_err(Into::into) } } + +impl From for CommitError { + fn from(err: crate::Error) -> Self { + CommitError::Internal(err) + } +} + +impl From for CommitError { + fn from(err: mysql_async::Error) -> Self { + CommitError::Mysql(err) + } +} diff --git a/crates/store/src/backend/postgres/read.rs b/crates/store/src/backend/postgres/read.rs index 1b67db3a..d253b137 100644 --- a/crates/store/src/backend/postgres/read.rs +++ b/crates/store/src/backend/postgres/read.rs @@ -58,10 +58,10 @@ impl PostgresStore { pub(crate) async fn get_bitmap( &self, - mut key: BitmapKey, + mut key: BitmapKey>, ) -> crate::Result> { let begin = key.serialize(0); - key.block_num = u32::MAX; + key.document_id = u32::MAX; let key_len = begin.len(); let end = key.serialize(0); let conn = self.conn_pool.get().await?; @@ -140,7 +140,7 @@ impl PostgresStore { pub(crate) async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { let key = key.into().serialize(0); let conn = self.conn_pool.get().await?; diff --git a/crates/store/src/backend/postgres/read_dense.rs b/crates/store/src/backend/postgres/read_dense.rs deleted file mode 100644 index 5d4da589..00000000 --- a/crates/store/src/backend/postgres/read_dense.rs +++ /dev/null @@ -1,155 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use futures::{pin_mut, TryStreamExt}; -use roaring::RoaringBitmap; - -use crate::{ - write::{bitmap::DeserializeBlock, key::DeserializeBigEndian, BitmapClass, ValueClass}, - BitmapKey, Deserialize, IterateParams, Key, ValueKey, U32_LEN, -}; - -use super::PostgresStore; - -impl PostgresStore { - pub(crate) async fn get_value(&self, key: impl Key) -> crate::Result> - where - U: Deserialize + 'static, - { - let conn = self.conn_pool.get().await?; - let s = conn - .prepare_cached(&format!( - "SELECT v FROM {} WHERE k = $1", - char::from(key.subspace()) - )) - .await?; - let key = key.serialize(0); - conn.query_opt(&s, &[&key]) - .await - .map_err(Into::into) - .and_then(|r| { - if let Some(r) = r { - Ok(Some(U::deserialize(r.get(0))?)) - } else { - Ok(None) - } - }) - } - - pub(crate) async fn get_bitmap( - &self, - mut key: BitmapKey, - ) -> crate::Result> { - let begin = key.serialize(0); - key.block_num = u32::MAX; - let key_len = begin.len(); - let end = key.serialize(0); - let conn = self.conn_pool.get().await?; - - let mut bm = RoaringBitmap::new(); - let s = conn - .prepare_cached("SELECT k, v FROM b WHERE k >= $1 AND k <= $2") - .await?; - let rows = conn.query_raw(&s, &[&begin, &end]).await?; - - pin_mut!(rows); - - while let Some(row) = rows.try_next().await? { - let key: &[u8] = row.try_get(0)?; - if key.len() == key_len { - let value: &[u8] = row.try_get(0)?; - bm.deserialize_block(value, key.deserialize_be_u32(key.len() - U32_LEN)?); - } - } - Ok(if !bm.is_empty() { Some(bm) } else { None }) - } - - pub(crate) async fn iterate( - &self, - params: IterateParams, - mut cb: impl for<'x> FnMut(&'x [u8], &'x [u8]) -> crate::Result + Sync + Send, - ) -> crate::Result<()> { - let conn = self.conn_pool.get().await?; - let table = char::from(params.begin.subspace()); - let begin = params.begin.serialize(0); - let end = params.end.serialize(0); - let keys = if params.values { "k, v" } else { "k" }; - - let s = conn - .prepare_cached(&match (params.first, params.ascending) { - (true, true) => { - format!( - "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1" - ) - } - (true, false) => { - format!( - "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC LIMIT 1" - ) - } - (false, true) => { - format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC") - } - (false, false) => { - format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC") - } - }) - .await?; - let rows = conn.query_raw(&s, &[&begin, &end]).await?; - - pin_mut!(rows); - - if params.values { - while let Some(row) = rows.try_next().await? { - let key = row.try_get::<_, &[u8]>(0)?; - let value = row.try_get::<_, &[u8]>(1)?; - - if !cb(key, value)? { - break; - } - } - } else { - while let Some(row) = rows.try_next().await? { - if !cb(row.try_get::<_, &[u8]>(0)?, b"")? { - break; - } - } - } - - Ok(()) - } - - pub(crate) async fn get_counter( - &self, - key: impl Into> + Sync + Send, - ) -> crate::Result { - let key = key.into().serialize(0); - let conn = self.conn_pool.get().await?; - let s = conn.prepare_cached("SELECT v FROM c WHERE k = $1").await?; - match conn.query_opt(&s, &[&key]).await { - Ok(Some(row)) => row.try_get(0).map_err(Into::into), - Ok(None) => Ok(0), - Err(e) => Err(e.into()), - } - } -} diff --git a/crates/store/src/backend/postgres/read_roaring.rs b/crates/store/src/backend/postgres/read_roaring.rs deleted file mode 100644 index c7b946c3..00000000 --- a/crates/store/src/backend/postgres/read_roaring.rs +++ /dev/null @@ -1,151 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use futures::{pin_mut, TryStreamExt}; -use roaring::RoaringBitmap; - -use crate::{ - write::{BitmapClass, ValueClass}, - BitmapKey, Deserialize, IterateParams, Key, ValueKey, WITHOUT_BLOCK_NUM, -}; - -use super::{deserialize_bitmap, PostgresStore}; - -impl PostgresStore { - pub(crate) async fn get_value(&self, key: impl Key) -> crate::Result> - where - U: Deserialize + 'static, - { - let conn = self.conn_pool.get().await?; - let s = conn - .prepare_cached(&format!( - "SELECT v FROM {} WHERE k = $1", - char::from(key.subspace()) - )) - .await?; - let key = key.serialize(0); - conn.query_opt(&s, &[&key]) - .await - .map_err(Into::into) - .and_then(|r| { - if let Some(r) = r { - Ok(Some(U::deserialize(r.get(0))?)) - } else { - Ok(None) - } - }) - } - - pub(crate) async fn get_bitmap( - &self, - key: BitmapKey, - ) -> crate::Result> { - let conn = self.conn_pool.get().await?; - let s = conn.prepare_cached("SELECT v FROM b WHERE k = $1").await?; - let key = key.serialize(WITHOUT_BLOCK_NUM); - conn.query_opt(&s, &[&key]) - .await - .map_err(Into::into) - .and_then(|r| { - if let Some(r) = r { - let bm = deserialize_bitmap(r.get(0))?; - if !bm.is_empty() { - Ok(Some(bm)) - } else { - Ok(None) - } - } else { - Ok(None) - } - }) - } - - pub(crate) async fn iterate( - &self, - params: IterateParams, - mut cb: impl for<'x> FnMut(&'x [u8], &'x [u8]) -> crate::Result + Sync + Send, - ) -> crate::Result<()> { - let conn = self.conn_pool.get().await?; - let table = char::from(params.begin.subspace()); - let begin = params.begin.serialize(0); - let end = params.end.serialize(0); - let keys = if params.values { "k, v" } else { "k" }; - - let s = conn - .prepare_cached(&match (params.first, params.ascending) { - (true, true) => { - format!( - "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1" - ) - } - (true, false) => { - format!( - "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC LIMIT 1" - ) - } - (false, true) => { - format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC") - } - (false, false) => { - format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC") - } - }) - .await?; - let rows = conn.query_raw(&s, &[&begin, &end]).await?; - - pin_mut!(rows); - - if params.values { - while let Some(row) = rows.try_next().await? { - let key = row.try_get::<_, &[u8]>(0)?; - let value = row.try_get::<_, &[u8]>(1)?; - - if !cb(key, value)? { - break; - } - } - } else { - while let Some(row) = rows.try_next().await? { - if !cb(row.try_get::<_, &[u8]>(0)?, b"")? { - break; - } - } - } - - Ok(()) - } - - pub(crate) async fn get_counter( - &self, - key: impl Into> + Sync + Send, - ) -> crate::Result { - let key = key.into().serialize(0); - let conn = self.conn_pool.get().await?; - let s = conn.prepare_cached("SELECT v FROM c WHERE k = $1").await?; - match conn.query_opt(&s, &[&key]).await { - Ok(Some(row)) => row.try_get(0).map_err(Into::into), - Ok(None) => Ok(0), - Err(e) => Err(e.into()), - } - } -} diff --git a/crates/store/src/backend/postgres/write.rs b/crates/store/src/backend/postgres/write.rs index 07e284ae..21832f9b 100644 --- a/crates/store/src/backend/postgres/write.rs +++ b/crates/store/src/backend/postgres/write.rs @@ -25,20 +25,30 @@ use std::time::{Duration, Instant}; use ahash::AHashMap; use deadpool_postgres::Object; +use futures::{pin_mut, TryStreamExt}; use rand::Rng; +use roaring::RoaringBitmap; use tokio_postgres::{error::SqlState, IsolationLevel}; use crate::{ write::{ - Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, + key::DeserializeBigEndian, AssignedIds, Batch, BitmapClass, Operation, RandomAvailableId, + ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, }, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, + BitmapKey, IndexKey, Key, LogKey, SUBSPACE_COUNTERS, U32_LEN, }; use super::PostgresStore; +#[derive(Debug)] +enum CommitError { + Postgres(tokio_postgres::Error), + Internal(crate::Error), + Retry, +} + impl PostgresStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result { let mut conn = self.conn_pool.get().await?; let start = Instant::now(); let mut retry_count = 0; @@ -46,21 +56,35 @@ impl PostgresStore { loop { match self.write_trx(&mut conn, &batch).await { Ok(result) => { - return result; + return Ok(result); } - Err(err) => match err.code() { - Some( - &SqlState::T_R_SERIALIZATION_FAILURE | &SqlState::T_R_DEADLOCK_DETECTED, - ) if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => { - let backoff = rand::thread_rng().gen_range(50..=300); - tokio::time::sleep(Duration::from_millis(backoff)).await; - retry_count += 1; + Err(err) => { + match err { + CommitError::Postgres(err) => match err.code() { + Some( + &SqlState::T_R_SERIALIZATION_FAILURE + | &SqlState::T_R_DEADLOCK_DETECTED, + ) if retry_count < MAX_COMMIT_ATTEMPTS + && start.elapsed() < MAX_COMMIT_TIME => {} + Some(&SqlState::UNIQUE_VIOLATION) => { + return Err(crate::Error::AssertValueFailed); + } + _ => return Err(err.into()), + }, + CommitError::Internal(err) => return Err(err), + CommitError::Retry => { + if retry_count > MAX_COMMIT_ATTEMPTS + || start.elapsed() > MAX_COMMIT_TIME + { + return Err(crate::Error::AssertValueFailed); + } + } } - Some(&SqlState::UNIQUE_VIOLATION) => { - return Err(crate::Error::AssertValueFailed); - } - _ => return Err(err.into()), - }, + + let backoff = rand::thread_rng().gen_range(50..=300); + tokio::time::sleep(Duration::from_millis(backoff)).await; + retry_count += 1; + } } } } @@ -69,7 +93,7 @@ impl PostgresStore { &self, conn: &mut Object, batch: &Batch, - ) -> Result>, tokio_postgres::Error> { + ) -> Result { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; @@ -79,7 +103,7 @@ impl PostgresStore { .isolation_level(IsolationLevel::ReadCommitted) .start() .await?; - let mut result = None; + let mut result = AssignedIds::default(); for op in &batch.ops { match op { @@ -99,14 +123,9 @@ impl PostgresStore { document_id = *document_id_; } Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + let table = char::from(class.subspace(collection)); match op { ValueOp::Set(value) => { @@ -135,23 +154,12 @@ impl PostgresStore { .await? }; - if trx.execute(&s, &[&key, value]).await? == 0 { - return Ok(Err(crate::Error::AssertValueFailed)); - } - - if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let s = trx.prepare_cached("SELECT 1 FROM b WHERE k = $1").await?; - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: document_id, - } - .serialize(0); - if trx.query_opt(&s, &[&key]).await?.is_some() { - return Ok(Err(crate::Error::AssertValueFailed)); - } + if trx + .execute(&s, &[&key, &value.resolve(&result)?.as_ref()]) + .await? + == 0 + { + return Err(crate::Error::AssertValueFailed.into()); } } ValueOp::AtomicAdd(by) => { @@ -177,11 +185,11 @@ impl PostgresStore { "ON CONFLICT(k) DO UPDATE SET v = c.v + EXCLUDED.v RETURNING v" )) .await?; - result = trx - .query_one(&s, &[&key, &by]) - .await - .and_then(|row| row.try_get::<_, i64>(0))? - .into(); + result.push_counter_id( + trx.query_one(&s, &[&key, &by]) + .await + .and_then(|row| row.try_get::<_, i64>(0))?, + ); } ValueOp::Clear => { let s = trx @@ -212,16 +220,50 @@ impl PostgresStore { trx.execute(&s, &[&key]).await?; } Operation::Bitmap { class, set } => { - let key = BitmapKey { - account_id, - collection, - class, - block_num: document_id, + // Find the next available document id + let is_document_id = matches!(class, BitmapClass::DocumentIds); + if *set && is_document_id && document_id == u32::MAX { + let begin = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: 0, + } + .serialize(0); + let end = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: u32::MAX, + } + .serialize(0); + let key_len = begin.len(); + + let s = trx + .prepare_cached("SELECT k FROM b WHERE k >= $1 AND k <= $2") + .await?; + let rows = trx.query_raw(&s, &[&begin, &end]).await?; + + pin_mut!(rows); + + let mut found_ids = RoaringBitmap::new(); + + while let Some(row) = rows.try_next().await? { + let key: &[u8] = row.try_get(0)?; + if key.len() == key_len { + found_ids.insert(key.deserialize_be_u32(key_len - U32_LEN)?); + } + } + + document_id = found_ids.random_available_id(); + result.push_document_id(document_id); } - .serialize(0); + + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); let s = if *set { - if matches!(class, BitmapClass::DocumentIds) { + if is_document_id { trx.prepare_cached("INSERT INTO b (k) VALUES ($1)").await? } else { trx.prepare_cached( @@ -232,17 +274,21 @@ impl PostgresStore { } else { trx.prepare_cached("DELETE FROM b WHERE k = $1").await? }; - trx.execute(&s, &[&key]).await?; + + trx.execute(&s, &[&key]).await.map_err(|err| { + if is_document_id && matches!(err.code(), Some(&SqlState::UNIQUE_VIOLATION)) + { + CommitError::Retry + } else { + CommitError::Postgres(err) + } + })?; } - Operation::Log { - collection, - change_id, - set, - } => { + Operation::Log { set } => { let key = LogKey { account_id, - collection: *collection, - change_id: *change_id, + collection, + change_id: batch.change_id, } .serialize(0); @@ -252,20 +298,17 @@ impl PostgresStore { "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" )) .await?; - trx.execute(&s, &[&key, set]).await?; + + trx.execute(&s, &[&key, &set.resolve(&result)?.as_ref()]) + .await?; } Operation::AssertValue { class, assert_value, } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + let table = char::from(class.subspace(collection)); let s = trx .prepare_cached(&format!("SELECT v FROM {} WHERE k = $1 FOR UPDATE", table)) @@ -279,14 +322,14 @@ impl PostgresStore { }) .unwrap_or_else(|| (false, assert_value.is_none())); if !matches { - return Ok(Err(crate::Error::AssertValueFailed)); + return Err(crate::Error::AssertValueFailed.into()); } asserted_values.insert(key, exists); } } } - trx.commit().await.map(|_| Ok(result)) + trx.commit().await.map(|_| result).map_err(Into::into) } pub(crate) async fn purge_store(&self) -> crate::Result<()> { @@ -316,3 +359,15 @@ impl PostgresStore { .map_err(Into::into) } } + +impl From for CommitError { + fn from(err: crate::Error) -> Self { + CommitError::Internal(err) + } +} + +impl From for CommitError { + fn from(err: tokio_postgres::Error) -> Self { + CommitError::Postgres(err) + } +} diff --git a/crates/store/src/backend/postgres/write_dense.rs b/crates/store/src/backend/postgres/write_dense.rs deleted file mode 100644 index 67d1126e..00000000 --- a/crates/store/src/backend/postgres/write_dense.rs +++ /dev/null @@ -1,386 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use std::{ - collections::BTreeMap, - time::{Duration, Instant}, -}; - -use ahash::AHashMap; -use deadpool_postgres::Object; -use rand::Rng; -use tokio_postgres::{error::SqlState, IsolationLevel}; - -use crate::{ - write::{ - bitmap::{block_contains, DenseBitmap}, - Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, - }, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, -}; - -use super::PostgresStore; - -impl PostgresStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { - let mut conn = self.conn_pool.get().await?; - let start = Instant::now(); - let mut retry_count = 0; - - loop { - match self.write_trx(&mut conn, &batch).await { - Ok(success) => { - return if success { - Ok(()) - } else { - Err(crate::Error::AssertValueFailed) - }; - } - Err(err) => match err.code() { - Some( - &SqlState::T_R_SERIALIZATION_FAILURE - | &SqlState::T_R_DEADLOCK_DETECTED - | &SqlState::UNIQUE_VIOLATION, - ) if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => { - let backoff = rand::thread_rng().gen_range(50..=300); - tokio::time::sleep(Duration::from_millis(backoff)).await; - retry_count += 1; - } - Some(&SqlState::UNIQUE_VIOLATION) => { - return Err(crate::Error::AssertValueFailed); - } - _ => return Err(err.into()), - }, - } - } - } - - async fn write_trx( - &self, - conn: &mut Object, - batch: &Batch, - ) -> Result { - let mut account_id = u32::MAX; - let mut collection = u8::MAX; - let mut document_id = u32::MAX; - let mut asserted_values = AHashMap::new(); - let trx = conn - .build_transaction() - .isolation_level(IsolationLevel::ReadCommitted) - .start() - .await?; - - // Sort the operations by key to avoid deadlocks - let mut assert_values = BTreeMap::new(); - let mut bitmap_updates = BTreeMap::new(); - for op in &batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::AssertValue { - class, - assert_value, - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - assert_values.insert(key.serialize(0), (table, assert_value)); - } - Operation::Bitmap { class, set } => { - bitmap_updates - .entry( - BitmapKey { - account_id, - collection, - class, - block_num: DenseBitmap::block_num(document_id), - } - .serialize(0), - ) - .or_insert_with(Vec::new) - .push((*set, document_id)); - } - _ => {} - } - } - - // Assert values - for (key, (table, assert_value)) in assert_values { - let s = trx - .prepare_cached(&format!("SELECT v FROM {} WHERE k = $1 FOR UPDATE", table)) - .await?; - let (exists, matches) = trx - .query_opt(&s, &[&key]) - .await? - .map(|row| { - row.try_get::<_, &[u8]>(0) - .map_or((true, false), |v| (true, assert_value.matches(v))) - }) - .unwrap_or_else(|| (false, assert_value.is_none())); - if !matches { - return Ok(false); - } - asserted_values.insert(key, exists); - } - - // Update bitmaps - for (key, changes) in bitmap_updates { - // Try updating the bitmap first - let mut update_query = String::from("v"); - let mut has_inserts = false; - for (set, document_id) in &changes { - update_query = format!( - "set_bit({update_query},{},{})", - DenseBitmap::block_index(*document_id), - *set as i8 - ); - has_inserts = has_inserts || *set; - } - - let s = trx - .prepare(&format!("UPDATE b SET v = {update_query} WHERE k = $1")) - .await?; - if trx.execute(&s, &[&key]).await? == 0 && has_inserts { - // The bitmap does not exist, create it - let mut dense_bm = DenseBitmap::empty(); - for (set, document_id) in changes { - if set { - dense_bm.set(document_id); - } - } - let s = trx - .prepare(&format!( - "INSERT INTO b (k, v) VALUES ($1, $2) ON CONFLICT(k) DO UPDATE SET v = {}", - update_query.replace("(v,", "(b.v,") - )) - .await?; - trx.execute(&s, &[&key, &&dense_bm.bitmap[..]]).await?; - } - } - - // Apply the operations - account_id = u32::MAX; - collection = u8::MAX; - document_id = u32::MAX; - for op in &batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::Value { - class, - op: ValueOp::AtomicAdd(by), - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(0); - - if *by >= 0 { - let s = trx - .prepare_cached(concat!( - "INSERT INTO c (k, v) VALUES ($1, $2) ", - "ON CONFLICT(k) DO UPDATE SET v = c.v + EXCLUDED.v" - )) - .await?; - trx.execute(&s, &[&key, &by]).await?; - } else { - let s = trx - .prepare_cached("UPDATE c SET v = v + $1 WHERE k = $2") - .await?; - trx.execute(&s, &[&by, &key]).await?; - } - } - Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); - - if let ValueOp::Set(value) = op { - let s = if let Some(exists) = asserted_values.get(&key) { - if *exists { - trx.prepare_cached(&format!( - "UPDATE {} SET v = $2 WHERE k = $1", - table - )) - .await? - } else { - trx.prepare_cached(&format!( - "INSERT INTO {} (k, v) VALUES ($1, $2)", - table - )) - .await? - } - } else { - trx.prepare_cached(&format!( - concat!( - "INSERT INTO {} (k, v) VALUES ($1, $2) ", - "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" - ), - table - )) - .await? - }; - - if trx.execute(&s, &[&key, value]).await? == 0 { - return Ok(false); - } - - if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let block_num = DenseBitmap::block_num(document_id); - let s = trx.prepare_cached("SELECT v FROM b WHERE k = $1").await?; - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num, - } - .serialize(0); - - if let Some(row) = trx.query_opt(&s, &[&key]).await? { - if block_contains(row.get(0), block_num, document_id) { - return Ok(false); - } - } - } - } else { - let s = trx - .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) - .await?; - trx.execute(&s, &[&key]).await?; - } - } - Operation::Index { field, key, set } => { - let key = IndexKey { - account_id, - collection, - document_id, - field: *field, - key, - } - .serialize(0); - - let s = if *set { - trx.prepare_cached( - "INSERT INTO i (k) VALUES ($1) ON CONFLICT (k) DO NOTHING", - ) - .await? - } else { - trx.prepare_cached("DELETE FROM i WHERE k = $1").await? - }; - trx.execute(&s, &[&key]).await?; - } - - Operation::Log { - collection, - change_id, - set, - } => { - let key = LogKey { - account_id, - collection: *collection, - change_id: *change_id, - } - .serialize(0); - - let s = trx - .prepare_cached(concat!( - "INSERT INTO l (k, v) VALUES ($1, $2) ", - "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" - )) - .await?; - trx.execute(&s, &[&key, set]).await?; - } - Operation::Bitmap { .. } | Operation::AssertValue { .. } => {} - } - } - - trx.commit().await.map(|_| true) - } - - pub(crate) async fn purge_store(&self) -> crate::Result<()> { - let todo = "delete bitmaps"; - let conn = self.conn_pool.get().await?; - - let s = conn - .prepare_cached(&format!( - "DELETE FROM {} WHERE v = 0", - char::from(SUBSPACE_COUNTERS), - )) - .await?; - conn.execute(&s, &[]).await.map(|_| ()).map_err(Into::into) - } - - pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> crate::Result<()> { - let conn = self.conn_pool.get().await?; - - let s = conn - .prepare_cached(&format!( - "DELETE FROM {} WHERE k >= $1 AND k < $2", - char::from(from.subspace()), - )) - .await?; - conn.execute(&s, &[&from.serialize(0), &to.serialize(0)]) - .await - .map(|_| ()) - .map_err(Into::into) - } -} diff --git a/crates/store/src/backend/postgres/write_roaring.rs b/crates/store/src/backend/postgres/write_roaring.rs deleted file mode 100644 index ba550d64..00000000 --- a/crates/store/src/backend/postgres/write_roaring.rs +++ /dev/null @@ -1,417 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use std::{ - collections::{BTreeMap, BTreeSet}, - time::{Duration, Instant}, -}; - -use ahash::AHashMap; -use deadpool_postgres::Object; -use rand::Rng; -use roaring::RoaringBitmap; -use tokio_postgres::{error::SqlState, IsolationLevel}; - -use crate::{ - write::{Batch, Operation, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME}, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, WITHOUT_BLOCK_NUM, -}; - -use super::{deserialize_bitmap, PostgresStore}; - -impl PostgresStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result<()> { - let mut conn = self.conn_pool.get().await?; - let start = Instant::now(); - let mut retry_count = 0; - - loop { - match self.write_trx(&mut conn, &batch).await { - Ok(success) => { - return if success { - Ok(()) - } else { - Err(crate::Error::AssertValueFailed) - }; - } - Err(err) => match err.code() { - Some( - &SqlState::T_R_SERIALIZATION_FAILURE - | &SqlState::T_R_DEADLOCK_DETECTED - | &SqlState::UNIQUE_VIOLATION, - ) if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => { - let backoff = rand::thread_rng().gen_range(50..=300); - tokio::time::sleep(Duration::from_millis(backoff)).await; - retry_count += 1; - } - Some(&SqlState::UNIQUE_VIOLATION) => { - return Err(crate::Error::AssertValueFailed); - } - _ => return Err(err.into()), - }, - } - } - } - - async fn write_trx( - &self, - conn: &mut Object, - batch: &Batch, - ) -> Result { - let mut account_id = u32::MAX; - let mut collection = u8::MAX; - let mut document_id = u32::MAX; - let mut asserted_values = AHashMap::new(); - let trx = conn - .build_transaction() - .isolation_level(IsolationLevel::ReadCommitted) - .start() - .await?; - - // Sort the operations by key to avoid deadlocks - let mut assert_values = BTreeMap::new(); - let mut bitmap_updates = BTreeMap::new(); - let mut advisory_locks = BTreeSet::new(); - for op in &batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::AssertValue { - class, - assert_value, - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - assert_values.insert(key.serialize(0), (table, assert_value)); - if account_id != u32::MAX { - advisory_locks.insert((account_id as u64) << 32 | collection as u64); - } - } - Operation::Bitmap { class, set } => { - bitmap_updates - .entry( - BitmapKey { - account_id, - collection, - class, - block_num: 0, - } - .serialize(WITHOUT_BLOCK_NUM), - ) - .or_insert_with(Vec::new) - .push((*set, document_id)); - if account_id != u32::MAX { - advisory_locks.insert((account_id as u64) << 32 | collection as u64); - } - } - _ => {} - } - } - - // Acquire advisory locks - for lock in advisory_locks { - trx.execute("SELECT pg_advisory_xact_lock($1)", &[&(lock as i64)]) - .await?; - } - - // Assert values - for (key, (table, assert_value)) in assert_values { - let s = trx - .prepare_cached(&format!("SELECT v FROM {} WHERE k = $1 FOR UPDATE", table)) - .await?; - let (exists, matches) = trx - .query_opt(&s, &[&key]) - .await? - .map(|row| { - row.try_get::<_, &[u8]>(0) - .map_or((true, false), |v| (true, assert_value.matches(v))) - }) - .unwrap_or_else(|| (false, assert_value.is_none())); - if !matches { - return Ok(false); - } - asserted_values.insert(key, exists); - } - - // Update bitmaps - for (key, changes) in bitmap_updates { - let s = trx - .prepare_cached("SELECT v FROM b WHERE k = $1 FOR UPDATE") - .await?; - let (value_exists, mut bm) = match trx - .query_opt(&s, &[&key]) - .await? - .map(|r| deserialize_bitmap(r.get(0))) - { - Some(Ok(bm)) => (true, bm), - None => (false, RoaringBitmap::new()), - Some(Err(e)) => { - tracing::error!("Failed to deserialize bitmap: {:?}", e); - return Ok(false); - } - }; - - let mut has_changes = false; - for (set, document_id) in changes { - if set { - if bm.insert(document_id) { - has_changes = true; - } - } else if bm.remove(document_id) { - has_changes = true; - } - } - - if has_changes { - if !bm.is_empty() { - let mut bytes = Vec::with_capacity(bm.serialized_size() + 1); - let _ = bm.serialize_into(&mut bytes); - let s = if value_exists { - trx.prepare_cached("UPDATE b SET v = $2 WHERE k = $1") - .await? - } else { - trx.prepare_cached("INSERT INTO b (k, V) VALUES ($1, $2)") - .await? - }; - trx.execute(&s, &[&key, &bytes]).await?; - } else if value_exists { - let s = trx.prepare_cached("DELETE FROM b WHERE k = $1").await?; - trx.execute(&s, &[&key]).await?; - } - } - } - - // Apply the operations - account_id = u32::MAX; - collection = u8::MAX; - document_id = u32::MAX; - for op in &batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::Value { - class, - op: ValueOp::AtomicAdd(by), - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(0); - - if *by >= 0 { - let s = trx - .prepare_cached(concat!( - "INSERT INTO c (k, v) VALUES ($1, $2) ", - "ON CONFLICT(k) DO UPDATE SET v = c.v + EXCLUDED.v" - )) - .await?; - trx.execute(&s, &[&key, &by]).await?; - } else { - let s = trx - .prepare_cached("UPDATE c SET v = v + $1 WHERE k = $2") - .await?; - trx.execute(&s, &[&by, &key]).await?; - } - } - Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); - - if let ValueOp::Set(value) = op { - let s = if let Some(exists) = asserted_values.get(&key) { - if *exists { - trx.prepare_cached(&format!( - "UPDATE {} SET v = $2 WHERE k = $1", - table - )) - .await? - } else { - trx.prepare_cached(&format!( - "INSERT INTO {} (k, v) VALUES ($1, $2)", - table - )) - .await? - } - } else { - trx.prepare_cached(&format!( - concat!( - "INSERT INTO {} (k, v) VALUES ($1, $2) ", - "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" - ), - table - )) - .await? - }; - - if trx.execute(&s, &[&key, value]).await? == 0 { - return Ok(false); - } - - /*if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let s = trx.prepare_cached("SELECT v FROM b WHERE k = $1").await?; - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: document_id, - } - .serialize(WITHOUT_BLOCK_NUM); - - match trx - .query_opt(&s, &[&key]) - .await? - .map(|r| deserialize_bitmap(r.get(0))) - { - Some(Ok(bm)) if bm.contains(document_id) => { - return Ok(false); - } - Some(Err(e)) => { - tracing::error!("Failed to deserialize bitmap: {:?}", e); - return Ok(false); - } - _ => {} - } - }*/ - } else { - let s = trx - .prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table)) - .await?; - trx.execute(&s, &[&key]).await?; - } - } - Operation::Index { field, key, set } => { - let key = IndexKey { - account_id, - collection, - document_id, - field: *field, - key, - } - .serialize(0); - - let s = if *set { - trx.prepare_cached( - "INSERT INTO i (k) VALUES ($1) ON CONFLICT (k) DO NOTHING", - ) - .await? - } else { - trx.prepare_cached("DELETE FROM i WHERE k = $1").await? - }; - trx.execute(&s, &[&key]).await?; - } - - Operation::Log { - collection, - change_id, - set, - } => { - let key = LogKey { - account_id, - collection: *collection, - change_id: *change_id, - } - .serialize(0); - - let s = trx - .prepare_cached(concat!( - "INSERT INTO l (k, v) VALUES ($1, $2) ", - "ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v" - )) - .await?; - trx.execute(&s, &[&key, set]).await?; - } - Operation::Bitmap { .. } | Operation::AssertValue { .. } => {} - } - } - - trx.commit().await.map(|_| true) - } - - pub(crate) async fn purge_store(&self) -> crate::Result<()> { - let conn = self.conn_pool.get().await?; - - let s = conn - .prepare_cached(&format!( - "DELETE FROM {} WHERE v = 0", - char::from(SUBSPACE_COUNTERS), - )) - .await?; - conn.execute(&s, &[]).await.map(|_| ()).map_err(Into::into) - } - - pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> crate::Result<()> { - let conn = self.conn_pool.get().await?; - - let s = conn - .prepare_cached(&format!( - "DELETE FROM {} WHERE k >= $1 AND k < $2", - char::from(from.subspace()), - )) - .await?; - conn.execute(&s, &[&from.serialize(0), &to.serialize(0)]) - .await - .map(|_| ()) - .map_err(Into::into) - } -} diff --git a/crates/store/src/backend/rocksdb/main.rs b/crates/store/src/backend/rocksdb/main.rs index 6900e69d..d51ca5a8 100644 --- a/crates/store/src/backend/rocksdb/main.rs +++ b/crates/store/src/backend/rocksdb/main.rs @@ -23,17 +23,11 @@ use std::path::PathBuf; -use roaring::RoaringBitmap; -use rocksdb::{ - compaction_filter::Decision, ColumnFamilyDescriptor, MergeOperands, OptimisticTransactionDB, - Options, -}; +use rocksdb::{ColumnFamilyDescriptor, MergeOperands, OptimisticTransactionDB, Options}; use tokio::sync::oneshot; use utils::config::{utils::AsKey, Config}; -use crate::Deserialize; - use super::{RocksDbStore, CF_BITMAPS, CF_BLOBS, CF_COUNTERS, CF_INDEXES, CF_LOGS, CF_VALUES}; impl RocksDbStore { @@ -59,8 +53,6 @@ impl RocksDbStore { // Bitmaps let mut cf_opts = Options::default(); cf_opts.set_max_write_buffer_number(16); - cf_opts.set_merge_operator("merge", bitmap_merge, bitmap_partial_merge); - cf_opts.set_compaction_filter("compact", bitmap_compact); cfs.push(ColumnFamilyDescriptor::new(CF_BITMAPS, cf_opts)); // Counters @@ -170,27 +162,3 @@ pub fn numeric_value_merge( bytes.extend_from_slice(&value.to_le_bytes()); Some(bytes) } - -pub fn bitmap_merge( - _new_key: &[u8], - existing_val: Option<&[u8]>, - operands: &MergeOperands, -) -> Option> { - super::bitmap::bitmap_merge(existing_val, operands.len(), operands) -} - -pub fn bitmap_partial_merge( - _new_key: &[u8], - _existing_val: Option<&[u8]>, - _operands: &MergeOperands, -) -> Option> { - // Force a full merge - None -} - -pub fn bitmap_compact(_level: u32, _key: &[u8], value: &[u8]) -> Decision { - match RoaringBitmap::deserialize(value) { - Ok(bm) if bm.is_empty() => Decision::Remove, - _ => Decision::Keep, - } -} diff --git a/crates/store/src/backend/rocksdb/read.rs b/crates/store/src/backend/rocksdb/read.rs index 287d3e2f..504ee141 100644 --- a/crates/store/src/backend/rocksdb/read.rs +++ b/crates/store/src/backend/rocksdb/read.rs @@ -25,8 +25,8 @@ use roaring::RoaringBitmap; use rocksdb::{Direction, IteratorMode}; use crate::{ - write::{BitmapClass, ValueClass}, - BitmapKey, Deserialize, IterateParams, Key, ValueKey, WITHOUT_BLOCK_NUM, + write::{key::DeserializeBigEndian, BitmapClass, ValueClass}, + BitmapKey, Deserialize, IterateParams, Key, ValueKey, U32_LEN, }; use super::{RocksDbStore, CF_BITMAPS, CF_COUNTERS}; @@ -57,28 +57,29 @@ impl RocksDbStore { pub(crate) async fn get_bitmap( &self, - key: BitmapKey, + mut key: BitmapKey>, ) -> crate::Result> { let db = self.db.clone(); self.spawn_worker(move || { - db.get_pinned_cf( + let mut bm = RoaringBitmap::new(); + let begin = key.serialize(0); + key.document_id = u32::MAX; + let end = key.serialize(0); + let key_len = begin.len(); + for row in db.iterator_cf( &db.cf_handle(CF_BITMAPS).unwrap(), - &key.serialize(WITHOUT_BLOCK_NUM), - ) - .map_err(Into::into) - .and_then(|value| { - if let Some(value) = value { - RoaringBitmap::deserialize(&value).map(|rb| { - if !rb.is_empty() { - Some(rb) - } else { - None - } - }) + IteratorMode::From(&begin, Direction::Forward), + ) { + let (key, _) = row?; + let key = key.as_ref(); + if key.len() == key_len && key >= begin.as_slice() && key <= end.as_slice() { + bm.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); } else { - Ok(None) + break; } - }) + } + + Ok(if !bm.is_empty() { Some(bm) } else { None }) }) .await } @@ -120,7 +121,7 @@ impl RocksDbStore { pub(crate) async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { let key = key.into().serialize(0); let db = self.db.clone(); diff --git a/crates/store/src/backend/rocksdb/write.rs b/crates/store/src/backend/rocksdb/write.rs index 777cc686..3e670126 100644 --- a/crates/store/src/backend/rocksdb/write.rs +++ b/crates/store/src/backend/rocksdb/write.rs @@ -34,20 +34,18 @@ use rocksdb::{ OptimisticTransactionOptions, WriteOptions, }; -use super::{ - bitmap::{clear_bit, set_bit}, - RocksDbStore, CF_BITMAPS, CF_COUNTERS, CF_INDEXES, CF_LOGS, CF_VALUES, -}; +use super::{RocksDbStore, CF_BITMAPS, CF_COUNTERS, CF_INDEXES, CF_LOGS, CF_VALUES}; use crate::{ backend::deserialize_i64_le, write::{ - Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, + key::DeserializeBigEndian, AssignedIds, Batch, BitmapClass, Operation, RandomAvailableId, + ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, }, - BitmapKey, Deserialize, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, WITHOUT_BLOCK_NUM, + BitmapKey, Deserialize, IndexKey, Key, LogKey, SUBSPACE_COUNTERS, U32_LEN, }; impl RocksDbStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result { let db = self.db.clone(); self.spawn_worker(move || { @@ -175,288 +173,168 @@ enum CommitError { } impl<'x> RocksDBTransaction<'x> { - fn commit(&self) -> Result, CommitError> { + fn commit(&self) -> Result { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; - let mut result = None; + let mut result = AssignedIds::default(); let txn = self .db .transaction_opt(&WriteOptions::default(), &self.txn_opts); - if !self.batch.is_atomic() { - for op in &self.batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; + for op in &self.batch.ops { + match op { + Operation::AccountId { + account_id: account_id_, + } => { + account_id = *account_id_; + } + Operation::Collection { + collection: collection_, + } => { + collection = *collection_; + } + Operation::DocumentId { + document_id: document_id_, + } => { + document_id = *document_id_; + } + Operation::Value { class, op } => { + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + let is_counter = class.is_counter(collection); - let is_counter = key.is_counter(); - let key = key.serialize(0); - - match op { - ValueOp::Set(value) => { - txn.put_cf(&self.cf_values, &key, value)?; - - if matches!(class, ValueClass::ReservedId) { - if let Some(bitmap) = txn - .get_pinned_for_update_cf( - &self.cf_bitmaps, - &BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: 0, - } - .serialize(WITHOUT_BLOCK_NUM), - true, - ) - .map_err(CommitError::from) - .and_then(|bytes| { - if let Some(bytes) = bytes { - RoaringBitmap::deserialize(&bytes) - .map(Some) - .map_err(CommitError::from) - } else { - Ok(None) - } - })? - { - if bitmap.contains(document_id) { - txn.rollback()?; - return Err(CommitError::Internal( - crate::Error::AssertValueFailed, - )); - } + match op { + ValueOp::Set(value) => { + txn.put_cf(&self.cf_values, &key, value.resolve(&result)?.as_ref())?; + } + ValueOp::AtomicAdd(by) => { + txn.merge_cf(&self.cf_counters, &key, &by.to_le_bytes()[..])?; + } + ValueOp::AddAndGet(by) => { + let num = txn + .get_pinned_for_update_cf(&self.cf_counters, &key, true) + .map_err(CommitError::from) + .and_then(|bytes| { + if let Some(bytes) = bytes { + deserialize_i64_le(&bytes) + .map(|v| v + *by) + .map_err(CommitError::from) + } else { + Ok(*by) } - } - } - ValueOp::AtomicAdd(by) => { - txn.merge_cf(&self.cf_counters, &key, &by.to_le_bytes()[..])?; - } - ValueOp::AddAndGet(by) => { - let num = txn - .get_pinned_for_update_cf(&self.cf_counters, &key, true) - .map_err(CommitError::from) - .and_then(|bytes| { - if let Some(bytes) = bytes { - deserialize_i64_le(&bytes) - .map(|v| v + *by) - .map_err(CommitError::from) - } else { - Ok(*by) - } - })?; - txn.put_cf(&self.cf_counters, &key, &num.to_le_bytes()[..])?; - result = Some(num); - } - ValueOp::Clear => { - txn.delete_cf( - if is_counter { - &self.cf_counters - } else { - &self.cf_values - }, - &key, - )?; - } + })?; + txn.put_cf(&self.cf_counters, &key, &num.to_le_bytes()[..])?; + result.push_counter_id(num); } - } - Operation::Index { field, key, set } => { - let key = IndexKey { - account_id, - collection, - document_id, - field: *field, - key, - } - .serialize(0); - - if *set { - txn.put_cf(&self.cf_indexes, &key, [])?; - } else { - txn.delete_cf(&self.cf_indexes, &key)?; - } - } - Operation::Bitmap { class, set } => { - let key = BitmapKey { - account_id, - collection, - class, - block_num: 0, - } - .serialize(WITHOUT_BLOCK_NUM); - - let value = if *set { - set_bit(document_id) - } else { - clear_bit(document_id) - }; - - txn.merge_cf(&self.cf_bitmaps, key, value)?; - } - Operation::Log { - collection, - change_id, - set, - } => { - let key = LogKey { - account_id, - collection: *collection, - change_id: *change_id, - } - .serialize(0); - - txn.put_cf(&self.cf_logs, &key, set)?; - } - Operation::AssertValue { - class, - assert_value, - } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - } - .serialize(0); - let matches = txn - .get_pinned_for_update_cf(&self.cf_values, &key, true)? - .map(|value| assert_value.matches(&value)) - .unwrap_or_else(|| assert_value.is_none()); - - if !matches { - txn.rollback()?; - return Err(CommitError::Internal(crate::Error::AssertValueFailed)); + ValueOp::Clear => { + txn.delete_cf( + if is_counter { + &self.cf_counters + } else { + &self.cf_values + }, + &key, + )?; } } } - } - - txn.commit().map(|_| result).map_err(Into::into) - } else { - let mut wb = txn.get_writebatch(); - for op in &self.batch.ops { - match op { - Operation::AccountId { - account_id: account_id_, - } => { - account_id = *account_id_; - } - Operation::Collection { - collection: collection_, - } => { - collection = *collection_; - } - Operation::DocumentId { - document_id: document_id_, - } => { - document_id = *document_id_; - } - Operation::Value { class, op } => { - let key = ValueKey { - account_id, - collection, - document_id, - class, - }; - - let is_counter = key.is_counter(); - let key = key.serialize(0); - - match op { - ValueOp::Set(value) => { - wb.put_cf(&self.cf_values, &key, value); - } - ValueOp::AtomicAdd(by) => { - wb.merge_cf(&self.cf_counters, &key, &by.to_le_bytes()[..]); - } - ValueOp::Clear => { - wb.delete_cf( - if is_counter { - &self.cf_counters - } else { - &self.cf_values - }, - &key, - ); - } - ValueOp::AddAndGet(_) => unreachable!(), - } - } - Operation::Index { field, key, set } => { - let key = IndexKey { - account_id, - collection, - document_id, - field: *field, - key, - } - .serialize(0); - - if *set { - wb.put_cf(&self.cf_indexes, &key, []); - } else { - wb.delete_cf(&self.cf_indexes, &key); - } - } - Operation::Bitmap { class, set } => { - let key = BitmapKey { - account_id, - collection, - class, - block_num: 0, - } - .serialize(WITHOUT_BLOCK_NUM); - - let value = if *set { - set_bit(document_id) - } else { - clear_bit(document_id) - }; - - wb.merge_cf(&self.cf_bitmaps, key, value); - } - Operation::Log { + Operation::Index { field, key, set } => { + let key = IndexKey { + account_id, collection, - change_id, - set, - } => { - let key = LogKey { + document_id, + field: *field, + key, + } + .serialize(0); + + if *set { + txn.put_cf(&self.cf_indexes, &key, [])?; + } else { + txn.delete_cf(&self.cf_indexes, &key)?; + } + } + Operation::Bitmap { class, set } => { + let is_document_id = matches!(class, BitmapClass::DocumentIds); + if *set && is_document_id && document_id == u32::MAX { + let begin = BitmapKey { account_id, - collection: *collection, - change_id: *change_id, + collection, + class: BitmapClass::DocumentIds, + document_id: 0, } .serialize(0); + let end = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: u32::MAX, + } + .serialize(0); + let key_len = begin.len(); + let mut found_ids = RoaringBitmap::new(); - wb.put_cf(&self.cf_logs, &key, set); + for row in txn.iterator_cf( + &self.cf_bitmaps, + IteratorMode::From(&begin, Direction::Forward), + ) { + let (key, _) = row?; + let key = key.as_ref(); + if key.len() == key_len + && key >= begin.as_slice() + && key <= end.as_slice() + { + found_ids.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); + } else { + break; + } + } + + document_id = found_ids.random_available_id(); + result.push_document_id(document_id); + } + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + + if *set { + txn.put_cf(&self.cf_bitmaps, &key, [])?; + } else { + txn.delete_cf(&self.cf_bitmaps, &key)?; + } + } + Operation::Log { set } => { + let key = LogKey { + account_id, + collection, + change_id: self.batch.change_id, + } + .serialize(0); + + txn.put_cf(&self.cf_logs, &key, set.resolve(&result)?.as_ref())?; + } + Operation::AssertValue { + class, + assert_value, + } => { + let key = + class.serialize(account_id, collection, document_id, 0, (&result).into()); + + let matches = txn + .get_pinned_for_update_cf(&self.cf_values, &key, true)? + .map(|value| assert_value.matches(&value)) + .unwrap_or_else(|| assert_value.is_none()); + + if !matches { + txn.rollback()?; + return Err(CommitError::Internal(crate::Error::AssertValueFailed)); } - Operation::AssertValue { .. } => unreachable!(), } } - - self.db.write(wb).map(|_| result).map_err(Into::into) } + + txn.commit().map(|_| result).map_err(Into::into) } } diff --git a/crates/store/src/backend/sqlite/read.rs b/crates/store/src/backend/sqlite/read.rs index f1fe3391..0870a3f2 100644 --- a/crates/store/src/backend/sqlite/read.rs +++ b/crates/store/src/backend/sqlite/read.rs @@ -56,10 +56,10 @@ impl SqliteStore { pub(crate) async fn get_bitmap( &self, - mut key: BitmapKey, + mut key: BitmapKey>, ) -> crate::Result> { let begin = key.serialize(0); - key.block_num = u32::MAX; + key.document_id = u32::MAX; let key_len = begin.len(); let end = key.serialize(0); let conn = self.conn_pool.get()?; @@ -137,7 +137,7 @@ impl SqliteStore { pub(crate) async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { let key = key.into().serialize(0); let conn = self.conn_pool.get()?; diff --git a/crates/store/src/backend/sqlite/write.rs b/crates/store/src/backend/sqlite/write.rs index 27863606..c2797456 100644 --- a/crates/store/src/backend/sqlite/write.rs +++ b/crates/store/src/backend/sqlite/write.rs @@ -21,24 +21,28 @@ * for more details. */ +use roaring::RoaringBitmap; use rusqlite::{params, OptionalExtension, TransactionBehavior}; use crate::{ - write::{Batch, BitmapClass, Operation, ValueClass, ValueOp}, - BitmapKey, IndexKey, Key, LogKey, ValueKey, SUBSPACE_COUNTERS, + write::{ + key::DeserializeBigEndian, AssignedIds, Batch, BitmapClass, Operation, RandomAvailableId, + ValueOp, + }, + BitmapKey, IndexKey, Key, LogKey, SUBSPACE_COUNTERS, U32_LEN, }; use super::SqliteStore; impl SqliteStore { - pub(crate) async fn write(&self, batch: Batch) -> crate::Result> { + pub(crate) async fn write(&self, batch: Batch) -> crate::Result { let mut conn = self.conn_pool.get()?; self.spawn_worker(move || { let mut account_id = u32::MAX; let mut collection = u8::MAX; let mut document_id = u32::MAX; let trx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; - let mut result = None; + let mut result = AssignedIds::default(); for op in &batch.ops { match op { @@ -58,14 +62,14 @@ impl SqliteStore { document_id = *document_id_; } Operation::Value { class, op } => { - let key = ValueKey { + let key = class.serialize( account_id, collection, document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + 0, + (&result).into(), + ); + let table = char::from(class.subspace(collection)); match op { ValueOp::Set(value) => { @@ -73,27 +77,7 @@ impl SqliteStore { "INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)", table ))? - .execute([&key, value])?; - - if matches!(class, ValueClass::ReservedId) { - // Make sure the reserved id is not already in use - let key = BitmapKey { - account_id, - collection, - class: BitmapClass::DocumentIds, - block_num: document_id, - } - .serialize(0); - if trx - .prepare_cached("SELECT 1 FROM b WHERE k = ?")? - .query_row([&key], |_| Ok(true)) - .optional()? - .unwrap_or(false) - { - trx.rollback()?; - return Err(crate::Error::AssertValueFailed); - } - } + .execute([&key, value.resolve(&result)?.as_ref()])?; } ValueOp::AtomicAdd(by) => { if *by >= 0 { @@ -108,13 +92,14 @@ impl SqliteStore { } } ValueOp::AddAndGet(by) => { - result = trx - .prepare_cached(concat!( + result.push_counter_id( + trx.prepare_cached(concat!( "INSERT INTO c (k, v) VALUES (?, ?) ", - "ON CONFLICT(k) DO UPDATE SET v = v + excluded.v RETURNING v" + "ON CONFLICT(k) DO UPDATE SET v = v + ", + "excluded.v RETURNING v" ))? - .query_row(params![&key, &by], |row| row.get::<_, i64>(0))? - .into(); + .query_row(params![&key, &by], |row| row.get::<_, i64>(0))?, + ); } ValueOp::Clear => { trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table))? @@ -141,49 +126,83 @@ impl SqliteStore { } } Operation::Bitmap { class, set } => { - let key = BitmapKey { + // Find the next available document id + let is_document_id = matches!(class, BitmapClass::DocumentIds); + if *set && is_document_id && document_id == u32::MAX { + let begin = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: 0, + } + .serialize(0); + let end = BitmapKey { + account_id, + collection, + class: BitmapClass::DocumentIds, + document_id: u32::MAX, + } + .serialize(0); + let key_len = begin.len(); + + let mut query = + trx.prepare_cached("SELECT k FROM b WHERE k >= ? AND k <= ?")?; + let mut rows = query.query([&begin, &end])?; + let mut found_ids = RoaringBitmap::new(); + while let Some(row) = rows.next()? { + let key = row.get_ref(0)?.as_bytes()?; + if key.len() == key_len { + found_ids.insert(key.deserialize_be_u32(key.len() - U32_LEN)?); + } + } + + document_id = found_ids.random_available_id(); + result.push_document_id(document_id); + } + let key = class.serialize( account_id, collection, - class, - block_num: document_id, - } - .serialize(0); + document_id, + 0, + (&result).into(), + ); if *set { - trx.prepare_cached("INSERT OR IGNORE INTO b (k) VALUES (?)")? - .execute(params![&key])?; + if is_document_id { + trx.prepare_cached("INSERT INTO b (k) VALUES (?)")? + .execute(params![&key])?; + } else { + trx.prepare_cached("INSERT OR IGNORE INTO b (k) VALUES (?)")? + .execute(params![&key])?; + } } else { trx.prepare_cached("DELETE FROM b WHERE k = ?")? .execute(params![&key])?; }; } - Operation::Log { - collection, - change_id, - set, - } => { + Operation::Log { set } => { let key = LogKey { account_id, - collection: *collection, - change_id: *change_id, + collection, + change_id: batch.change_id, } .serialize(0); trx.prepare_cached("INSERT OR REPLACE INTO l (k, v) VALUES (?, ?)")? - .execute([&key, set])?; + .execute([&key, set.resolve(&result)?.as_ref()])?; } Operation::AssertValue { class, assert_value, } => { - let key = ValueKey { + let key = class.serialize( account_id, collection, document_id, - class, - }; - let table = char::from(key.subspace()); - let key = key.serialize(0); + 0, + (&result).into(), + ); + let table = char::from(class.subspace(collection)); let matches = trx .prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table))? diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index 21c4d989..74cb185f 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -73,7 +73,8 @@ impl LookupStore { KeySerializer::new(value.len() + U64_LEN) .write(expires.map_or(u64::MAX, |expires| now() + expires)) .write(value.as_slice()) - .finalize(), + .finalize() + .into(), ), }); store.write(batch.build()).await.map(|_| ()) @@ -111,7 +112,8 @@ impl LookupStore { op: ValueOp::Set( KeySerializer::new(U64_LEN) .write(now() + expires) - .finalize(), + .finalize() + .into(), ), }); } @@ -125,7 +127,13 @@ impl LookupStore { }, }); - store.write(batch.build()).await.map(|r| r.unwrap_or(0)) + store.write(batch.build()).await.and_then(|r| { + if return_value { + r.last_counter_id() + } else { + Ok(0) + } + }) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_incr(key, value, expires).await, diff --git a/crates/store/src/dispatch/store.rs b/crates/store/src/dispatch/store.rs index 0667e890..602abfbe 100644 --- a/crates/store/src/dispatch/store.rs +++ b/crates/store/src/dispatch/store.rs @@ -26,7 +26,9 @@ use std::ops::{BitAndAssign, Range}; use roaring::RoaringBitmap; use crate::{ - write::{key::KeySerializer, now, AnyKey, Batch, BitmapClass, ReportClass, ValueClass}, + write::{ + key::KeySerializer, now, AnyKey, AssignedIds, Batch, BitmapClass, ReportClass, ValueClass, + }, BitmapKey, Deserialize, IterateParams, Key, Store, ValueKey, SUBSPACE_BITMAPS, SUBSPACE_INDEXES, SUBSPACE_LOGS, U32_LEN, }; @@ -59,7 +61,7 @@ impl Store { pub async fn get_bitmap( &self, - key: BitmapKey, + key: BitmapKey>, ) -> crate::Result> { match self { #[cfg(feature = "sqlite")] @@ -78,7 +80,7 @@ impl Store { pub async fn get_bitmaps_intersection( &self, - keys: Vec>, + keys: Vec>>, ) -> crate::Result> { let mut result: Option = None; for key in keys { @@ -120,7 +122,7 @@ impl Store { pub async fn get_counter( &self, - key: impl Into> + Sync + Send, + key: impl Into>> + Sync + Send, ) -> crate::Result { match self { #[cfg(feature = "sqlite")] @@ -137,7 +139,7 @@ impl Store { } } - pub async fn write(&self, batch: Batch) -> crate::Result> { + pub async fn write(&self, batch: Batch) -> crate::Result { #[cfg(feature = "test_mode")] if std::env::var("PARANOID_WRITE").map_or(false, |v| v == "1") { use crate::write::Operation; @@ -146,6 +148,7 @@ impl Store { let mut document_id = u32::MAX; let mut bitmaps = Vec::new(); + let mut result = AssignedIds::default(); for op in &batch.ops { match op { @@ -165,13 +168,19 @@ impl Store { document_id = *document_id_; } Operation::Bitmap { class, set } => { - let key = BitmapKey { + if *set && matches!(class, BitmapClass::DocumentIds) { + let id = result.document_ids.len() as u32; + result.document_ids.push(id); + } + + let key = class.serialize( account_id, collection, - block_num: 0, - class, - } - .serialize(0); + document_id, + 0, + (&result).into(), + ); + bitmaps.push((key, class.clone(), document_id, *set)); } _ => {} @@ -216,7 +225,7 @@ impl Store { } } - return Ok(None); + return Ok(AssignedIds::default()); } match self { @@ -310,7 +319,6 @@ impl Store { for (from_class, to_class) in [ (ValueClass::Acl(account_id), ValueClass::Acl(account_id + 1)), - (ValueClass::ReservedId, ValueClass::ReservedId), (ValueClass::Property(0), ValueClass::Property(0)), (ValueClass::TermIndex, ValueClass::TermIndex), ] { diff --git a/crates/store/src/fts/query.rs b/crates/store/src/fts/query.rs index f57153f9..7fe30645 100644 --- a/crates/store/src/fts/query.rs +++ b/crates/store/src/fts/query.rs @@ -81,7 +81,7 @@ impl Store { account_id, collection, class: BitmapClass::word(token.word.as_ref(), field), - block_num: 0, + document_id: 0, }); if !last_token.is_empty() { @@ -141,7 +141,7 @@ impl Store { account_id, collection, class: BitmapClass::word(token.word.as_ref(), field), - block_num: 0, + document_id: 0, }; let token2 = BitmapKey { account_id, @@ -155,7 +155,7 @@ impl Store { .as_ref(), field, ), - block_num: 0, + document_id: 0, }; match self.get_bitmaps_union(vec![token1, token2]).await? { @@ -184,7 +184,7 @@ impl Store { account_id, collection, class: BitmapClass::word(text, field), - block_num: 0, + document_id: 0, }) .await? } @@ -264,7 +264,7 @@ impl Store { async fn get_bitmaps_union( &self, - keys: Vec>, + keys: Vec>>, ) -> crate::Result> { let mut bm = RoaringBitmap::new(); diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index eeb20259..5d95e4f3 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -73,7 +73,6 @@ pub trait Serialize { // Key serialization flags pub(crate) const WITH_SUBSPACE: u32 = 1; -pub(crate) const WITHOUT_BLOCK_NUM: u32 = 1 << 1; pub trait Key: Sync + Send { fn serialize(&self, flags: u32) -> Vec; @@ -81,11 +80,11 @@ pub trait Key: Sync + Send { } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct BitmapKey> { +pub struct BitmapKey>> { pub account_id: u32, pub collection: u8, pub class: T, - pub block_num: u32, + pub document_id: u32, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -105,7 +104,7 @@ pub struct IndexKeyPrefix { } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct ValueKey> { +pub struct ValueKey>> { pub account_id: u32, pub collection: u8, pub document_id: u32, diff --git a/crates/store/src/query/filter.rs b/crates/store/src/query/filter.rs index 135004b0..a260dafb 100644 --- a/crates/store/src/query/filter.rs +++ b/crates/store/src/query/filter.rs @@ -98,7 +98,7 @@ impl Store { account_id, collection, class, - block_num: 0, + document_id: 0, }) .await? } diff --git a/crates/store/src/query/mod.rs b/crates/store/src/query/mod.rs index ffc673a0..b62faeed 100644 --- a/crates/store/src/query/mod.rs +++ b/crates/store/src/query/mod.rs @@ -54,7 +54,7 @@ pub enum Filter { text: String, tokenize: bool, }, - InBitmap(BitmapClass), + InBitmap(BitmapClass), DocumentSet(RoaringBitmap), And, Or, @@ -160,7 +160,7 @@ impl Filter { } } - pub fn is_in_bitmap(field: impl Into, value: impl Into) -> Self { + pub fn is_in_bitmap(field: impl Into, value: impl Into>) -> Self { Self::InBitmap(BitmapClass::Tag { field: field.into(), value: value.into(), @@ -199,13 +199,13 @@ impl Comparator { } } -impl BitmapKey { +impl BitmapKey> { pub fn document_ids(account_id: u32, collection: impl Into) -> Self { BitmapKey { account_id, collection: collection.into(), class: BitmapClass::DocumentIds, - block_num: 0, + document_id: 0, } } @@ -222,7 +222,7 @@ impl BitmapKey { field: field.into(), token: BitmapHash::new(token), }, - block_num: 0, + document_id: 0, } } @@ -230,7 +230,7 @@ impl BitmapKey { account_id: u32, collection: impl Into, field: impl Into, - value: impl Into, + value: impl Into>, ) -> Self { BitmapKey { account_id, @@ -239,7 +239,7 @@ impl BitmapKey { field: field.into(), value: value.into(), }, - block_num: 0, + document_id: 0, } } } diff --git a/crates/store/src/query/sort.rs b/crates/store/src/query/sort.rs index d609ab70..ed21e042 100644 --- a/crates/store/src/query/sort.rs +++ b/crates/store/src/query/sort.rs @@ -41,7 +41,7 @@ pub struct Pagination { has_anchor: bool, anchor_found: bool, pub ids: Vec, - prefix_key: Option>, + prefix_key: Option>>, prefix_unique: bool, } @@ -303,7 +303,7 @@ impl Pagination { } } - pub fn with_prefix_key(mut self, prefix_key: ValueKey) -> Self { + pub fn with_prefix_key(mut self, prefix_key: ValueKey>) -> Self { self.prefix_key = Some(prefix_key); self } diff --git a/crates/store/src/write/assign_id.rs b/crates/store/src/write/assign_id.rs deleted file mode 100644 index 22a6b923..00000000 --- a/crates/store/src/write/assign_id.rs +++ /dev/null @@ -1,170 +0,0 @@ -/* - * Copyright (c) 2023 Stalwart Labs Ltd. - * - * This file is part of the Stalwart Mail Server. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as - * published by the Free Software Foundation, either version 3 of - * the License, or (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * in the LICENSE file at the top-level directory of this distribution. - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - * - * You can be released from the requirements of the AGPLv3 license by - * purchasing a commercial license. Please contact licensing@stalw.art - * for more details. -*/ - -use std::time::Instant; - -use crate::{ - backend::ID_ASSIGNMENT_EXPIRY, write::key::DeserializeBigEndian, Deserialize, IterateParams, - Serialize, Store, ValueKey, U32_LEN, -}; -use ahash::AHashMap; -use rand::Rng; -use roaring::RoaringBitmap; - -use crate::{write::now, BitmapKey}; - -use super::{BatchBuilder, ValueClass, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME}; - -impl Store { - pub async fn assign_document_id( - &self, - account_id: u32, - collection: impl Into + Sync + Send, - ) -> crate::Result { - let start = Instant::now(); - let mut retry_count = 0; - let collection = collection.into(); - - loop { - // First try to reuse an expired assigned id - let mut reserved_ids = RoaringBitmap::new(); - let mut expired_ids = AHashMap::new(); - { - let from_key = ValueKey { - account_id, - collection, - document_id: 0, - class: ValueClass::ReservedId, - }; - let to_key = ValueKey { - account_id, - collection, - document_id: u32::MAX, - class: ValueClass::ReservedId, - }; - - let expired_timestamp = now(); - self.iterate( - IterateParams::new(from_key, to_key).ascending(), - |key, value| { - let document_id = key.deserialize_be_u32(key.len() - U32_LEN)?; - let ttl = u64::deserialize(value)?; - - if ttl <= expired_timestamp { - // Found an expired id, reuse it - expired_ids.insert(document_id, ttl); - } else { - // Keep track of all reserved ids - reserved_ids.insert(document_id); - } - - Ok(true) - }, - ) - .await?; - } - - // Prepare the patch the id - let mut batch = BatchBuilder::new(); - batch - .with_account_id(account_id) - .with_collection(collection); - - let document_id = if !expired_ids.is_empty() { - // Obtain a random id from the expired ids - let pos = if expired_ids.len() > 1 { - rand::thread_rng().gen_range(0..expired_ids.len()) - } else { - 0 - }; - let (document_id, expiry) = expired_ids.into_iter().nth(pos).unwrap(); - - batch - .update_document(document_id) - .assert_value(ValueClass::ReservedId, expiry); - document_id - } else { - // Obtain documentIds - let document_ids = if let Some(document_ids) = self - .get_bitmap(BitmapKey::document_ids(account_id, collection)) - .await? - { - if !reserved_ids.is_empty() { - document_ids | reserved_ids - } else { - document_ids - } - } else { - reserved_ids - }; - - let document_id = if retry_count == 0 { - // Find the next available id - (0..(document_ids.len() + 1) as u32) - .find(|&x| !document_ids.contains(x)) - .unwrap() - } else { - // High contention, pick a random id - const RAND_IDS: usize = 10; - let mut available_ids = Vec::with_capacity(RAND_IDS); - for id in 0..(document_ids.len() as u32 + RAND_IDS as u32) { - if !document_ids.contains(id) { - available_ids.push(id); - if available_ids.len() == RAND_IDS { - break; - } - } - } - available_ids[rand::thread_rng().gen_range(0..available_ids.len())] - }; - - batch - .update_document(document_id) - .assert_value(ValueClass::ReservedId, ()); - document_id - }; - - #[cfg(not(feature = "test_mode"))] - let expired_timestamp = now() + ID_ASSIGNMENT_EXPIRY; - #[cfg(feature = "test_mode")] - let expired_timestamp = - now() + ID_ASSIGNMENT_EXPIRY.load(std::sync::atomic::Ordering::Relaxed); - - batch.set(ValueClass::ReservedId, expired_timestamp.serialize()); - - match self.write(batch.build()).await { - Ok(_) => { - return Ok(document_id); - } - Err(crate::Error::AssertValueFailed) - if retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME => - { - // Retry - retry_count += 1; - continue; - } - Err(err) => return Err(err), - } - } - } -} diff --git a/crates/store/src/write/batch.rs b/crates/store/src/write/batch.rs index 954facff..3f8a430b 100644 --- a/crates/store/src/write/batch.rs +++ b/crates/store/src/write/batch.rs @@ -22,17 +22,24 @@ */ use super::{ - assert::ToAssertValue, Batch, BatchBuilder, BitmapClass, HasFlag, IntoOperations, Operation, - Serialize, TagValue, ToBitmaps, ValueClass, ValueOp, F_BITMAP, F_CLEAR, F_INDEX, F_VALUE, + assert::ToAssertValue, Batch, BatchBuilder, BitmapClass, HasFlag, IntoOperations, + MaybeDynamicId, MaybeDynamicValue, Operation, Serialize, TagValue, ToBitmaps, ValueClass, + ValueOp, F_BITMAP, F_CLEAR, F_INDEX, F_VALUE, }; impl BatchBuilder { pub fn new() -> Self { Self { ops: Vec::with_capacity(16), + change_id: u64::MAX, } } + pub fn with_change_id(&mut self, change_id: u64) -> &mut Self { + self.change_id = change_id; + self + } + pub fn with_account_id(&mut self, account_id: u32) -> &mut Self { self.ops.push(Operation::AccountId { account_id }); self @@ -45,8 +52,10 @@ impl BatchBuilder { self } - pub fn create_document(&mut self, document_id: u32) -> &mut Self { - self.ops.push(Operation::DocumentId { document_id }); + pub fn create_document(&mut self) -> &mut Self { + self.ops.push(Operation::DocumentId { + document_id: u32::MAX, + }); // Add document id self.ops.push(Operation::Bitmap { @@ -54,10 +63,16 @@ impl BatchBuilder { set: true, }); - // Remove reserved id - self.ops.push(Operation::Value { - class: ValueClass::ReservedId, - op: ValueOp::Clear, + self + } + + pub fn create_document_with_id(&mut self, document_id: u32) -> &mut Self { + self.ops.push(Operation::DocumentId { document_id }); + + // Add document id + self.ops.push(Operation::Bitmap { + class: BitmapClass::DocumentIds, + set: true, }); self @@ -79,7 +94,7 @@ impl BatchBuilder { pub fn assert_value( &mut self, - class: impl Into, + class: impl Into>, value: impl ToAssertValue, ) -> &mut Self { self.ops.push(Operation::AssertValue { @@ -116,7 +131,7 @@ impl BatchBuilder { self.ops.push(Operation::Value { class: ValueClass::Property(field), op: if is_set { - ValueOp::Set(value) + ValueOp::Set(value.into()) } else { ValueOp::Clear }, @@ -129,7 +144,7 @@ impl BatchBuilder { pub fn tag( &mut self, field: impl Into, - value: impl Into, + value: impl Into>, options: u32, ) -> &mut Self { self.ops.push(Operation::Bitmap { @@ -142,7 +157,7 @@ impl BatchBuilder { self } - pub fn add(&mut self, class: impl Into, value: i64) -> &mut Self { + pub fn add(&mut self, class: impl Into>, value: i64) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::AtomicAdd(value), @@ -150,7 +165,11 @@ impl BatchBuilder { self } - pub fn add_and_get(&mut self, class: impl Into, value: i64) -> &mut Self { + pub fn add_and_get( + &mut self, + class: impl Into>, + value: i64, + ) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::AddAndGet(value), @@ -158,7 +177,11 @@ impl BatchBuilder { self } - pub fn set(&mut self, class: impl Into, value: impl Into>) -> &mut Self { + pub fn set( + &mut self, + class: impl Into>, + value: impl Into, + ) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::Set(value.into()), @@ -166,7 +189,7 @@ impl BatchBuilder { self } - pub fn clear(&mut self, class: impl Into) -> &mut Self { + pub fn clear(&mut self, class: impl Into>) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::Clear, @@ -174,18 +197,27 @@ impl BatchBuilder { self } + pub fn log(&mut self, value: impl Into) -> &mut Self { + self.ops.push(Operation::Log { set: value.into() }); + self + } + pub fn custom(&mut self, value: impl IntoOperations) -> &mut Self { value.build(self); self } pub fn build(self) -> Batch { - Batch { ops: self.ops } + Batch { + ops: self.ops, + change_id: self.change_id, + } } pub fn build_batch(&mut self) -> Batch { Batch { ops: std::mem::take(&mut self.ops), + change_id: self.change_id, } } @@ -216,10 +248,6 @@ impl Batch { matches!( op, Operation::AssertValue { .. } - | Operation::Value { - class: ValueClass::ReservedId, - op: ValueOp::Set(_) - } | Operation::Value { op: ValueOp::AddAndGet(_), .. diff --git a/crates/store/src/write/bitmap.rs b/crates/store/src/write/bitmap.rs index 7e3fccbf..a4b88214 100644 --- a/crates/store/src/write/bitmap.rs +++ b/crates/store/src/write/bitmap.rs @@ -24,55 +24,48 @@ use ahash::AHashSet; use roaring::RoaringBitmap; -use crate::U64_LEN; - -pub const WORD_SIZE_BITS_L: u32 = (WORD_SIZE_L * 8) as u32; -pub const WORD_SIZE_L: usize = std::mem::size_of::(); -pub const WORDS_PER_BLOCK_L: u32 = 8; -pub const BITS_PER_BLOCK_L: u32 = WORD_SIZE_BITS_L * WORDS_PER_BLOCK_L; -pub const BITS_MASK_L: u32 = BITS_PER_BLOCK_L - 1; - -pub const WORD_SIZE_BITS_S: u32 = (WORD_SIZE_S * 8) as u32; -pub const WORD_SIZE_S: usize = U64_LEN; -pub const WORDS_PER_BLOCK_S: u32 = 16; -pub const BITS_PER_BLOCK_S: u32 = WORD_SIZE_BITS_S * WORDS_PER_BLOCK_S; -pub const BITS_MASK_S: u32 = BITS_PER_BLOCK_S - 1; +pub const WORD_SIZE_BITS: u32 = (WORD_SIZE * 8) as u32; +pub const WORD_SIZE: usize = std::mem::size_of::(); +pub const WORDS_PER_BLOCK: u32 = 8; +pub const BITS_PER_BLOCK: u32 = WORD_SIZE_BITS * WORDS_PER_BLOCK; +pub const BITS_MASK: u32 = BITS_PER_BLOCK - 1; +pub const BITMAP_SIZE: usize = WORD_SIZE * WORDS_PER_BLOCK as usize; pub struct DenseBitmap { - pub bitmap: [u8; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], + pub bitmap: [u8; BITMAP_SIZE], } impl DenseBitmap { pub fn empty() -> Self { Self { - bitmap: [0; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], + bitmap: [0; BITMAP_SIZE], } } pub fn full() -> Self { Self { - bitmap: [u8::MAX; WORD_SIZE_L * WORDS_PER_BLOCK_L as usize], + bitmap: [u8::MAX; BITMAP_SIZE], } } pub fn set(&mut self, index: u32) { - let index = index & BITS_MASK_L; + let index = index & BITS_MASK; self.bitmap[(index / 8) as usize] |= 1 << (index & 7); } pub fn clear(&mut self, index: u32) { - let index = index & BITS_MASK_L; + let index = index & BITS_MASK; self.bitmap[(index / 8) as usize] &= !(1 << (index & 7)); } #[inline(always)] pub fn block_num(index: u32) -> u32 { - index / BITS_PER_BLOCK_L + index / BITS_PER_BLOCK } #[inline(always)] pub fn block_index(index: u32) -> u32 { - index & BITS_MASK_L + index & BITS_MASK } } @@ -97,7 +90,7 @@ pub fn next_available_index( } } - let id = (block_num * BITS_PER_BLOCK_L) + ((byte_pos * 8) + index) as u32; + let id = (block_num * BITS_PER_BLOCK) + ((byte_pos * 8) + index) as u32; if !reserved_ids.contains(&id) { return Some(id); } else if index < 7 { @@ -125,7 +118,7 @@ pub fn block_contains(bytes: &[u8], block_num: u32, document_id: u32) -> bool { } } - let id = (block_num * BITS_PER_BLOCK_L) + ((byte_pos * 8) + index) as u32; + let id = (block_num * BITS_PER_BLOCK) + ((byte_pos * 8) + index) as u32; if id == document_id { return true; } else if index < 7 { @@ -143,16 +136,16 @@ pub fn block_contains(bytes: &[u8], block_num: u32, document_id: u32) -> bool { impl DeserializeBlock for RoaringBitmap { fn deserialize_block(&mut self, bytes: &[u8], block_num: u32) { - debug_assert_eq!(bytes.len(), WORD_SIZE_L * WORDS_PER_BLOCK_L as usize); + debug_assert_eq!(bytes.len(), BITMAP_SIZE); - self.deserialize_word(&bytes[..WORD_SIZE_L], block_num, 0); - self.deserialize_word(&bytes[WORD_SIZE_L..WORD_SIZE_L * 2], block_num, 1); - self.deserialize_word(&bytes[WORD_SIZE_L * 2..WORD_SIZE_L * 3], block_num, 2); - self.deserialize_word(&bytes[WORD_SIZE_L * 3..WORD_SIZE_L * 4], block_num, 3); - self.deserialize_word(&bytes[WORD_SIZE_L * 4..WORD_SIZE_L * 5], block_num, 4); - self.deserialize_word(&bytes[WORD_SIZE_L * 5..WORD_SIZE_L * 6], block_num, 5); - self.deserialize_word(&bytes[WORD_SIZE_L * 6..WORD_SIZE_L * 7], block_num, 6); - self.deserialize_word(&bytes[WORD_SIZE_L * 7..], block_num, 7); + self.deserialize_word(&bytes[..WORD_SIZE], block_num, 0); + self.deserialize_word(&bytes[WORD_SIZE..WORD_SIZE * 2], block_num, 1); + self.deserialize_word(&bytes[WORD_SIZE * 2..WORD_SIZE * 3], block_num, 2); + self.deserialize_word(&bytes[WORD_SIZE * 3..WORD_SIZE * 4], block_num, 3); + self.deserialize_word(&bytes[WORD_SIZE * 4..WORD_SIZE * 5], block_num, 4); + self.deserialize_word(&bytes[WORD_SIZE * 5..WORD_SIZE * 6], block_num, 5); + self.deserialize_word(&bytes[WORD_SIZE * 6..WORD_SIZE * 7], block_num, 6); + self.deserialize_word(&bytes[WORD_SIZE * 7..], block_num, 7); } #[inline(always)] @@ -161,16 +154,15 @@ impl DeserializeBlock for RoaringBitmap { 0 => (), u128::MAX => { self.insert_range( - block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L - ..(block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L) - + WORD_SIZE_BITS_L, + block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS + ..(block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS) + WORD_SIZE_BITS, ); } mut word => { while word != 0 { let trailing_zeros = word.trailing_zeros(); self.insert( - block_num * BITS_PER_BLOCK_L + word_num * WORD_SIZE_BITS_L + trailing_zeros, + block_num * BITS_PER_BLOCK + word_num * WORD_SIZE_BITS + trailing_zeros, ); word ^= 1 << trailing_zeros; } @@ -196,7 +188,7 @@ mod tests { for item in range { bitmap.insert(item); blocks - .entry(item / BITS_PER_BLOCK_L) + .entry(item / BITS_PER_BLOCK) .or_insert_with(DenseBitmap::empty) .set(item); } diff --git a/crates/store/src/write/blob.rs b/crates/store/src/write/blob.rs index 8f78218e..f4fd8e08 100644 --- a/crates/store/src/write/blob.rs +++ b/crates/store/src/write/blob.rs @@ -169,12 +169,7 @@ impl Store { .unwrap(); let until = key.deserialize_be_u64(key.len() - U64_LEN)?; if until <= now { - delete_keys.push(ValueKey { - account_id: key.deserialize_be_u32(1)?, - collection: 0, - document_id: 0, - class: ValueClass::Blob(BlobOp::Reserve { until, hash }), - }); + delete_keys.push((key.deserialize_be_u32(1)?, BlobOp::Reserve { until, hash })); } else { active_hashes.insert(hash); } @@ -220,12 +215,7 @@ impl Store { } } else if last_hash != hash && !active_hashes.contains(&hash) { // Unlinked or expired blob, delete. - delete_keys.push(ValueKey { - account_id: 0, - collection: 0, - document_id: 0, - class: ValueClass::Blob(BlobOp::Commit { hash }), - }); + delete_keys.push((0, BlobOp::Commit { hash })); } Ok(true) @@ -234,8 +224,8 @@ impl Store { .await?; // Delete expired or unlinked blobs - for key in &delete_keys { - if let ValueClass::Blob(BlobOp::Commit { hash }) = &key.class { + for (_, op) in &delete_keys { + if let BlobOp::Commit { hash } = op { blob_store.delete_blob(hash.as_ref()).await?; } } @@ -243,20 +233,18 @@ impl Store { // Delete hashes let mut batch = BatchBuilder::new(); let mut last_account_id = u32::MAX; - for key in delete_keys.into_iter() { + for (account_id, op) in delete_keys.into_iter() { if batch.ops.len() >= 1000 { last_account_id = u32::MAX; self.write(batch.build()).await?; batch = BatchBuilder::new(); } - if matches!(key.class, ValueClass::Blob(BlobOp::Reserve { .. })) - && key.account_id != last_account_id - { - batch.with_account_id(key.account_id); - last_account_id = key.account_id; + if matches!(op, BlobOp::Reserve { .. }) && account_id != last_account_id { + batch.with_account_id(account_id); + last_account_id = account_id; } batch.ops.push(Operation::Value { - class: key.class, + class: ValueClass::Blob(op), op: ValueOp::Clear, }) } @@ -294,11 +282,10 @@ impl Store { if document_id != u32::MAX && key.deserialize_be_u32(1 + BLOB_HASH_LEN)? == account_id { - delete_keys.push(ValueKey { - account_id, - collection: key[1 + BLOB_HASH_LEN + U32_LEN], + delete_keys.push(( + key[1 + BLOB_HASH_LEN + U32_LEN], document_id, - class: ValueClass::Blob(BlobOp::Link { + BlobOp::Link { hash: BlobHash::try_from_hash_slice( key.get(1..1 + BLOB_HASH_LEN).ok_or_else(|| { crate::Error::InternalError(format!( @@ -307,8 +294,8 @@ impl Store { })?, ) .unwrap(), - }), - }); + }, + )); } Ok(true) @@ -320,20 +307,20 @@ impl Store { let mut batch = BatchBuilder::new(); batch.with_account_id(account_id); let mut last_collection = u8::MAX; - for key in delete_keys.into_iter() { + for (collection, document_id, op) in delete_keys.into_iter() { if batch.ops.len() >= 1000 { self.write(batch.build()).await?; batch = BatchBuilder::new(); batch.with_account_id(account_id); last_collection = u8::MAX; } - if key.collection != last_collection { - batch.with_collection(key.collection); - last_collection = key.collection; + if collection != last_collection { + batch.with_collection(collection); + last_collection = collection; } - batch.update_document(key.document_id); + batch.update_document(document_id); batch.ops.push(Operation::Value { - class: key.class, + class: ValueClass::Blob(op), op: ValueOp::Clear, }); } diff --git a/crates/store/src/write/hash.rs b/crates/store/src/write/hash.rs index f0c882f8..aed8a7a5 100644 --- a/crates/store/src/write/hash.rs +++ b/crates/store/src/write/hash.rs @@ -25,7 +25,7 @@ use crate::backend::MAX_TOKEN_LENGTH; use super::{BitmapClass, BitmapHash}; -impl BitmapClass { +impl BitmapClass { pub fn word(token: impl AsRef<[u8]>, field: impl Into) -> Self { BitmapClass::Text { field: field.into(), diff --git a/crates/store/src/write/key.rs b/crates/store/src/write/key.rs index 4c877b92..43bce81a 100644 --- a/crates/store/src/write/key.rs +++ b/crates/store/src/write/key.rs @@ -27,12 +27,12 @@ use utils::{codec::leb128::Leb128_, BLOB_HASH_LEN}; use crate::{ BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, Key, LogKey, ValueKey, SUBSPACE_BITMAPS, SUBSPACE_COUNTERS, SUBSPACE_INDEXES, SUBSPACE_LOGS, SUBSPACE_VALUES, U32_LEN, U64_LEN, - WITHOUT_BLOCK_NUM, WITH_SUBSPACE, + WITH_SUBSPACE, }; use super::{ - AnyKey, BitmapClass, BlobOp, DirectoryClass, LookupClass, QueueClass, ReportClass, ReportEvent, - TagValue, ValueClass, + AnyKey, AssignedIds, BitmapClass, BlobOp, DirectoryClass, LookupClass, QueueClass, ReportClass, + ReportEvent, ResolveId, TagValue, ValueClass, }; pub struct KeySerializer { @@ -148,13 +148,13 @@ impl DeserializeBigEndian for &[u8] { } } -impl> ValueKey { +impl>> ValueKey { pub fn property( account_id: u32, collection: impl Into, document_id: u32, field: impl Into, - ) -> ValueKey { + ) -> ValueKey> { ValueKey { account_id, collection: collection.into(), @@ -171,13 +171,7 @@ impl> ValueKey { } pub fn is_counter(&self) -> bool { - match self.class.as_ref() { - ValueClass::Directory(DirectoryClass::UsedQuota(_)) - | ValueClass::Lookup(LookupClass::Counter(_)) - | ValueClass::Queue(QueueClass::QuotaCount(_) | QueueClass::QuotaSize(_)) => true, - ValueClass::Property(84) if self.collection == 1 => true, // TODO: Find a more elegant way to do this - _ => false, - } + self.class.as_ref().is_counter(self.collection) } } @@ -228,54 +222,64 @@ impl Key for LogKey { } } -impl + Sync + Send> Key for ValueKey { +impl> + Sync + Send> Key for ValueKey { fn subspace(&self) -> u8 { - if self.is_counter() { - SUBSPACE_COUNTERS - } else { - SUBSPACE_VALUES - } + self.class.as_ref().subspace(self.collection) } fn serialize(&self, flags: u32) -> Vec { + self.class.as_ref().serialize( + self.account_id, + self.collection, + self.document_id, + flags, + None, + ) + } +} + +impl ValueClass { + pub fn serialize( + &self, + account_id: u32, + collection: u8, + document_id: u32, + flags: u32, + assigned_ids: Option<&AssignedIds>, + ) -> Vec { let serializer = if (flags & WITH_SUBSPACE) != 0 { - KeySerializer::new(self.class.as_ref().serialized_size() + 2).write(self.subspace()) + KeySerializer::new(self.serialized_size() + 2).write(self.subspace(collection)) } else { - KeySerializer::new(self.class.as_ref().serialized_size() + 1) + KeySerializer::new(self.serialized_size() + 1) }; - match self.class.as_ref() { + match self { ValueClass::Property(field) => serializer .write(0u8) - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(*field) - .write(self.document_id), + .write(document_id), ValueClass::TermIndex => serializer .write(1u8) - .write(self.account_id) - .write(self.collection) - .write_leb128(self.document_id), + .write(account_id) + .write(collection) + .write_leb128(document_id), ValueClass::Acl(grant_account_id) => serializer .write(2u8) .write(*grant_account_id) - .write(self.account_id) - .write(self.collection) - .write(self.document_id), - ValueClass::ReservedId => serializer - .write(3u8) - .write(self.account_id) - .write(self.collection) - .write(self.document_id), + .write(account_id) + .write(collection) + .write(document_id), ValueClass::IndexEmail(seq) => serializer .write(5u8) .write(*seq) - .write(self.account_id) - .write(self.document_id), + .write(account_id) + .write(document_id), ValueClass::Blob(op) => match op { BlobOp::Reserve { hash, until } => serializer .write(6u8) - .write(self.account_id) + .write(account_id) .write::<&[u8]>(hash.as_ref()) .write(*until), BlobOp::Commit { hash } => serializer @@ -287,9 +291,9 @@ impl + Sync + Send> Key for ValueKey { BlobOp::Link { hash } => serializer .write(7u8) .write::<&[u8]>(hash.as_ref()) - .write(self.account_id) - .write(self.collection) - .write(self.document_id), + .write(account_id) + .write(collection) + .write(document_id), }, ValueClass::Config(key) => serializer.write(8u8).write(key.as_slice()), ValueClass::Lookup(lookup) => match lookup { @@ -300,7 +304,9 @@ impl + Sync + Send> Key for ValueKey { ValueClass::Directory(directory) => match directory { DirectoryClass::NameToId(name) => serializer.write(20u8).write(name.as_slice()), DirectoryClass::EmailToId(email) => serializer.write(21u8).write(email.as_slice()), - DirectoryClass::Principal(uid) => serializer.write(22u8).write_leb128(*uid), + DirectoryClass::Principal(uid) => serializer + .write(22u8) + .write_leb128(uid.resolve_id(assigned_ids)), DirectoryClass::Domain(name) => serializer.write(23u8).write(name.as_slice()), DirectoryClass::UsedQuota(uid) => serializer.write(24u8).write_leb128(*uid), DirectoryClass::MemberOf { @@ -308,15 +314,15 @@ impl + Sync + Send> Key for ValueKey { member_of, } => serializer .write(25u8) - .write(*principal_id) - .write(*member_of), + .write(principal_id.resolve_id(assigned_ids)) + .write(member_of.resolve_id(assigned_ids)), DirectoryClass::Members { principal_id, has_member, } => serializer .write(26u8) - .write(*principal_id) - .write(*has_member), + .write(principal_id.resolve_id(assigned_ids)) + .write(has_member.resolve_id(assigned_ids)), }, ValueClass::Queue(queue) => match queue { QueueClass::Message(queue_id) => serializer.write(50u8).write(*queue_id), @@ -393,12 +399,31 @@ impl + Sync + Send> Key for IndexKey { } } -impl + Sync + Send> Key for BitmapKey { +impl> + Sync + Send> Key for BitmapKey { fn subspace(&self) -> u8 { SUBSPACE_BITMAPS } fn serialize(&self, flags: u32) -> Vec { + self.class.as_ref().serialize( + self.account_id, + self.collection, + self.document_id, + flags, + None, + ) + } +} + +impl BitmapClass { + pub fn serialize( + &self, + account_id: u32, + collection: u8, + document_id: u32, + flags: u32, + assigned_ids: Option<&AssignedIds>, + ) -> Vec { const BM_DOCUMENT_IDS: u8 = 0; const BM_TAG: u8 = 1 << 6; const BM_TEXT: u8 = 1 << 7; @@ -407,14 +432,14 @@ impl + Sync + Send> Key for BitmapKey { const TAG_TEXT: u8 = 1 << 0; const TAG_STATIC: u8 = 1 << 1; - let serializer = match self.class.as_ref() { + match self { BitmapClass::DocumentIds => if (flags & WITH_SUBSPACE) != 0 { KeySerializer::new(U32_LEN + 3).write(SUBSPACE_BITMAPS) } else { KeySerializer::new(U32_LEN + 2) } - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(BM_DOCUMENT_IDS), BitmapClass::Tag { field, value } => match value { TagValue::Id(id) => if (flags & WITH_SUBSPACE) != 0 { @@ -422,18 +447,18 @@ impl + Sync + Send> Key for BitmapKey { } else { KeySerializer::new((U32_LEN * 2) + 3) } - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(BM_TAG | TAG_ID) .write(*field) - .write_leb128(*id), + .write_leb128(id.resolve_id(assigned_ids)), TagValue::Text(text) => if (flags & WITH_SUBSPACE) != 0 { KeySerializer::new(U32_LEN + 4 + text.len()).write(SUBSPACE_BITMAPS) } else { KeySerializer::new(U32_LEN + 3 + text.len()) } - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(BM_TAG | TAG_TEXT) .write(*field) .write(text.as_slice()), @@ -442,8 +467,8 @@ impl + Sync + Send> Key for BitmapKey { } else { KeySerializer::new(U32_LEN + 4) } - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(BM_TAG | TAG_STATIC) .write(*field) .write(*id), @@ -453,18 +478,14 @@ impl + Sync + Send> Key for BitmapKey { } else { KeySerializer::new(U32_LEN + 16 + 3) } - .write(self.account_id) - .write(self.collection) + .write(account_id) + .write(collection) .write(BM_TEXT | token.len) .write(*field) .write(token.hash.as_slice()), - }; - - if (flags & WITHOUT_BLOCK_NUM) != 0 { - serializer.finalize() - } else { - serializer.write(self.block_num).finalize() } + .write(document_id) + .finalize() } } @@ -485,12 +506,10 @@ impl + Sync + Send> Key for AnyKey { } } -impl ValueClass { +impl ValueClass { pub fn serialized_size(&self) -> usize { match self { - ValueClass::Property(_) | ValueClass::TermIndex | ValueClass::ReservedId => { - U32_LEN * 2 + 3 - } + ValueClass::Property(_) | ValueClass::TermIndex => U32_LEN * 2 + 3, ValueClass::Acl(_) => U32_LEN * 3 + 2, ValueClass::Lookup( LookupClass::Counter(v) | LookupClass::CounterExpiry(v) | LookupClass::Key(v), @@ -522,10 +541,28 @@ impl ValueClass { ValueClass::Report(_) => U64_LEN * 2 + 1, } } + + pub fn subspace(&self, collection: u8) -> u8 { + if self.is_counter(collection) { + SUBSPACE_COUNTERS + } else { + SUBSPACE_VALUES + } + } + + pub fn is_counter(&self, collection: u8) -> bool { + match self { + ValueClass::Directory(DirectoryClass::UsedQuota(_)) + | ValueClass::Lookup(LookupClass::Counter(_)) + | ValueClass::Queue(QueueClass::QuotaCount(_) | QueueClass::QuotaSize(_)) => true, + ValueClass::Property(84) if collection == 1 => true, // TODO: Find a more elegant way to do this + _ => false, + } + } } -impl From for ValueKey { - fn from(class: ValueClass) -> Self { +impl From> for ValueKey> { + fn from(class: ValueClass) -> Self { ValueKey { account_id: 0, collection: 0, @@ -535,8 +572,8 @@ impl From for ValueKey { } } -impl From for ValueKey { - fn from(value: DirectoryClass) -> Self { +impl From> for ValueKey> { + fn from(value: DirectoryClass) -> Self { ValueKey { account_id: 0, collection: 0, @@ -546,13 +583,13 @@ impl From for ValueKey { } } -impl From for ValueClass { - fn from(value: DirectoryClass) -> Self { +impl From> for ValueClass { + fn from(value: DirectoryClass) -> Self { ValueClass::Directory(value) } } -impl From for ValueClass { +impl From for ValueClass { fn from(value: BlobOp) -> Self { ValueClass::Blob(value) } diff --git a/crates/store/src/write/log.rs b/crates/store/src/write/log.rs index 18bdd40c..923d0754 100644 --- a/crates/store/src/write/log.rs +++ b/crates/store/src/write/log.rs @@ -26,7 +26,7 @@ use utils::{codec::leb128::Leb128Vec, map::vec_map::VecMap}; use crate::Serialize; -use super::{IntoOperations, Operation}; +use super::{IntoOperations, MaybeDynamicValue, Operation, SerializeWithId}; #[derive(Default)] pub struct ChangeLogBuilder { @@ -144,17 +144,63 @@ impl ChangeLogBuilder { impl IntoOperations for ChangeLogBuilder { fn build(self, batch: &mut super::BatchBuilder) { + batch.change_id = self.change_id; for (collection, changes) in self.changes { + batch.ops.push(Operation::Collection { collection }); batch.ops.push(Operation::Log { - change_id: self.change_id, - collection, - set: changes.serialize(), + set: changes.serialize().into(), }); } } } -impl Serialize for Changes { +impl Changes { + pub fn insert(id: T) -> Self + where + T: IntoIterator, + I: Into, + { + Changes { + inserts: id.into_iter().map(Into::into).collect(), + ..Default::default() + } + } + + pub fn update(id: T) -> Self + where + T: IntoIterator, + I: Into, + { + Changes { + updates: id.into_iter().map(Into::into).collect(), + ..Default::default() + } + } + + pub fn child_update(id: T) -> Self + where + T: IntoIterator, + I: Into, + { + Changes { + child_updates: id.into_iter().map(Into::into).collect(), + ..Default::default() + } + } + + pub fn delete(id: T) -> Self + where + T: IntoIterator, + I: Into, + { + Changes { + deletes: id.into_iter().map(Into::into).collect(), + ..Default::default() + } + } +} + +impl Serialize for &Changes { fn serialize(self) -> Vec { let mut buf = Vec::with_capacity( 1 + (self.inserts.len() @@ -170,11 +216,37 @@ impl Serialize for Changes { buf.push_leb128(self.child_updates.len()); buf.push_leb128(self.deletes.len()); - for list in [self.inserts, self.updates, self.child_updates, self.deletes] { + for list in [ + &self.inserts, + &self.updates, + &self.child_updates, + &self.deletes, + ] { for id in list { - buf.push_leb128(id); + buf.push_leb128(*id); } } buf } } + +impl From for MaybeDynamicValue { + fn from(changes: Changes) -> Self { + MaybeDynamicValue::Static(changes.serialize()) + } +} + +pub struct LogInsert(); + +impl SerializeWithId for LogInsert { + fn serialize_with_id(&self, ids: &super::AssignedIds) -> crate::Result> { + ids.last_document_id() + .map(|id| Changes::insert([id]).serialize()) + } +} + +impl From for MaybeDynamicValue { + fn from(value: LogInsert) -> Self { + MaybeDynamicValue::Dynamic(Box::new(value)) + } +} diff --git a/crates/store/src/write/mod.rs b/crates/store/src/write/mod.rs index cf6225c4..5b2f854d 100644 --- a/crates/store/src/write/mod.rs +++ b/crates/store/src/write/mod.rs @@ -22,13 +22,17 @@ */ use std::{ + borrow::Cow, collections::HashSet, + fmt::{self, Formatter}, hash::Hash, slice::Iter, time::{Duration, SystemTime}, }; use nlp::tokenizers::word::WordTokenizer; +use rand::Rng; +use roaring::RoaringBitmap; use utils::{ codec::leb128::{Leb128Iterator, Leb128Vec}, BlobHash, @@ -39,7 +43,6 @@ use crate::{backend::MAX_TOKEN_LENGTH, BlobClass, Deserialize, Serialize, Value} use self::assert::AssertValue; pub mod assert; -pub mod assign_id; pub mod batch; pub mod bitmap; pub mod blob; @@ -48,6 +51,34 @@ pub mod key; pub mod log; pub mod purge; +pub trait SerializeWithId: Send + Sync { + fn serialize_with_id(&self, ids: &AssignedIds) -> crate::Result>; +} + +pub trait ResolveId { + fn resolve_id(&self, ids: Option<&AssignedIds>) -> u32; +} + +pub enum MaybeDynamicValue { + Static(Vec), + Dynamic(Box), +} + +#[derive(Debug, PartialEq, Clone, Copy, Eq, Hash)] +pub enum MaybeDynamicId { + Static(u32), + Dynamic(usize), +} + +#[derive(Debug, PartialEq, Clone, Eq, Hash)] +pub struct DynamicDocumentId(pub usize); + +#[derive(Debug, Default)] +pub struct AssignedIds { + pub document_ids: Vec, + pub counter_ids: Vec, +} + #[cfg(not(feature = "test_mode"))] pub(crate) const MAX_COMMIT_ATTEMPTS: u32 = 10; #[cfg(not(feature = "test_mode"))] @@ -66,11 +97,13 @@ pub const F_CLEAR: u32 = 1 << 3; #[derive(Debug)] pub struct Batch { pub ops: Vec, + pub change_id: u64, } #[derive(Debug)] pub struct BatchBuilder { pub ops: Vec, + pub change_id: u64, } #[derive(Debug, PartialEq, Eq, Hash)] @@ -85,11 +118,11 @@ pub enum Operation { document_id: u32, }, AssertValue { - class: ValueClass, + class: ValueClass, assert_value: AssertValue, }, Value { - class: ValueClass, + class: ValueClass, op: ValueOp, }, Index { @@ -98,20 +131,18 @@ pub enum Operation { set: bool, }, Bitmap { - class: BitmapClass, + class: BitmapClass, set: bool, }, Log { - change_id: u64, - collection: u8, - set: Vec, + set: MaybeDynamicValue, }, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum BitmapClass { +pub enum BitmapClass { DocumentIds, - Tag { field: u8, value: TagValue }, + Tag { field: u8, value: TagValue }, Text { field: u8, token: BitmapHash }, } @@ -122,20 +153,19 @@ pub struct BitmapHash { } #[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum TagValue { - Id(u32), +pub enum TagValue { + Id(T), Text(Vec), Static(u8), } #[derive(Debug, PartialEq, Clone, Eq, Hash)] -pub enum ValueClass { +pub enum ValueClass { Property(u8), Acl(u32), Lookup(LookupClass), TermIndex, - ReservedId, - Directory(DirectoryClass), + Directory(DirectoryClass), Blob(BlobOp), IndexEmail(u64), Config(Vec), @@ -151,13 +181,13 @@ pub enum LookupClass { } #[derive(Debug, PartialEq, Clone, Eq, Hash)] -pub enum DirectoryClass { +pub enum DirectoryClass { NameToId(Vec), EmailToId(Vec), - MemberOf { principal_id: u32, member_of: u32 }, - Members { principal_id: u32, has_member: u32 }, + MemberOf { principal_id: T, member_of: T }, + Members { principal_id: T, has_member: T }, Domain(Vec), - Principal(u32), + Principal(T), UsedQuota(u32), } @@ -196,7 +226,7 @@ pub struct ReportEvent { #[derive(Debug, PartialEq, Eq, Hash, Default)] pub enum ValueOp { - Set(Vec), + Set(MaybeDynamicValue), AtomicAdd(i64), AddAndGet(i64), #[default] @@ -216,31 +246,37 @@ pub struct AnyKey> { pub key: T, } -impl From for TagValue { +impl From for TagValue { + fn from(value: u32) -> Self { + TagValue::Id(MaybeDynamicId::Static(value)) + } +} + +impl From for TagValue { fn from(value: u32) -> Self { TagValue::Id(value) } } -impl From> for TagValue { +impl From> for TagValue { fn from(value: Vec) -> Self { TagValue::Text(value) } } -impl From for TagValue { +impl From for TagValue { fn from(value: String) -> Self { TagValue::Text(value.into_bytes()) } } -impl From for TagValue { +impl From for TagValue { fn from(value: u8) -> Self { TagValue::Static(value) } } -impl From<()> for TagValue { +impl From<()> for TagValue { fn from(_: ()) -> Self { TagValue::Text(vec![]) } @@ -493,7 +529,7 @@ impl ToBitmaps for u32 { ops.push(Operation::Bitmap { class: BitmapClass::Tag { field, - value: TagValue::Id(*self), + value: TagValue::Id(MaybeDynamicId::Static(*self)), }, set, }); @@ -505,7 +541,7 @@ impl ToBitmaps for u64 { ops.push(Operation::Bitmap { class: BitmapClass::Tag { field, - value: TagValue::Id(*self as u32), + value: TagValue::Id(MaybeDynamicId::Static(*self as u32)), }, set, }); @@ -552,7 +588,9 @@ impl Operation { pub fn acl(grant_account_id: u32, set: Option>) -> Self { Operation::Value { class: ValueClass::Acl(grant_account_id), - op: set.map(ValueOp::Set).unwrap_or(ValueOp::Clear), + op: set + .map(|op| ValueOp::Set(op.into())) + .unwrap_or(ValueOp::Clear), } } } @@ -564,23 +602,26 @@ pub fn now() -> u64 { .map_or(0, |d| d.as_secs()) } -impl AsRef for ValueClass { - fn as_ref(&self) -> &ValueClass { +impl AsRef> for ValueClass { + fn as_ref(&self) -> &ValueClass { self } } -impl AsRef for BitmapClass { - fn as_ref(&self) -> &BitmapClass { +impl AsRef> for BitmapClass { + fn as_ref(&self) -> &BitmapClass { self } } -impl BitmapClass { - pub fn tag_id(property: impl Into, id: u32) -> Self { +impl BitmapClass { + pub fn tag_id(property: impl Into, id: u32) -> Self + where + TagValue: From, + { BitmapClass::Tag { field: property.into(), - value: TagValue::Id(id), + value: id.into(), } } } @@ -668,3 +709,163 @@ impl ToBitmaps for &Bincode crate::Result { + self.document_ids + .get(idx) + .copied() + .ok_or_else(|| crate::Error::InternalError("No document ids were created".to_string())) + } + + pub fn first_document_id(&self) -> crate::Result { + self.get_document_id(0) + } + + pub fn last_document_id(&self) -> crate::Result { + self.document_ids + .last() + .copied() + .ok_or_else(|| crate::Error::InternalError("No document ids were created".to_string())) + } + + pub fn last_counter_id(&self) -> crate::Result { + self.counter_ids + .last() + .copied() + .ok_or_else(|| crate::Error::InternalError("No counter ids were created".to_string())) + } +} + +impl From for MaybeDynamicValue { + fn from(value: String) -> Self { + MaybeDynamicValue::Static(value.into_bytes()) + } +} + +impl From<&[u8]> for MaybeDynamicValue { + fn from(value: &[u8]) -> Self { + MaybeDynamicValue::Static(value.to_vec()) + } +} + +impl From> for MaybeDynamicValue { + fn from(value: Vec) -> Self { + MaybeDynamicValue::Static(value) + } +} + +impl MaybeDynamicValue { + pub fn resolve(&self, ids: &AssignedIds) -> crate::Result> { + match self { + MaybeDynamicValue::Static(value) => Ok(Cow::Borrowed(value.as_slice())), + MaybeDynamicValue::Dynamic(value) => value.serialize_with_id(ids).map(Cow::Owned), + } + } +} + +impl MaybeDynamicId { + pub fn resolve(&self, ids: &AssignedIds) -> crate::Result { + match self { + MaybeDynamicId::Static(id) => Ok(*id), + MaybeDynamicId::Dynamic(idx) => ids.get_document_id(*idx), + } + } +} + +impl ResolveId for u32 { + fn resolve_id(&self, _: Option<&AssignedIds>) -> u32 { + *self + } +} + +impl ResolveId for MaybeDynamicId { + fn resolve_id(&self, ids: Option<&AssignedIds>) -> u32 { + match self { + MaybeDynamicId::Static(id) => *id, + MaybeDynamicId::Dynamic(idx) => ids + .and_then(|ids| ids.document_ids.get(*idx)) + .copied() + .unwrap_or(u32::MAX), + } + } +} + +impl std::fmt::Debug for MaybeDynamicValue { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + MaybeDynamicValue::Static(value) => write!(f, "{:?}", value), + MaybeDynamicValue::Dynamic(_) => write!(f, "Dynamic"), + } + } +} + +impl PartialEq for MaybeDynamicValue { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (MaybeDynamicValue::Static(a), MaybeDynamicValue::Static(b)) => a == b, + (MaybeDynamicValue::Dynamic(_), MaybeDynamicValue::Dynamic(_)) => true, + _ => false, + } + } +} + +impl Eq for MaybeDynamicValue {} + +impl Hash for MaybeDynamicValue { + fn hash(&self, state: &mut H) { + match self { + MaybeDynamicValue::Static(value) => value.hash(state), + MaybeDynamicValue::Dynamic(_) => 0.hash(state), + } + } +} + +impl From for MaybeDynamicValue { + fn from(value: MaybeDynamicId) -> Self { + match value { + MaybeDynamicId::Static(id) => MaybeDynamicValue::Static(id.serialize()), + MaybeDynamicId::Dynamic(idx) => { + MaybeDynamicValue::Dynamic(Box::new(DynamicDocumentId(idx))) + } + } + } +} + +impl SerializeWithId for DynamicDocumentId { + fn serialize_with_id(&self, ids: &AssignedIds) -> crate::Result> { + ids.get_document_id(self.0).map(|id| id.serialize()) + } +} + +pub(crate) trait RandomAvailableId { + fn random_available_id(&self) -> u32; +} + +impl RandomAvailableId for RoaringBitmap { + fn random_available_id(&self) -> u32 { + let mut last_id = 0; + let mut available_ids = Vec::with_capacity(100); + for id in self.iter() { + for i in last_id..id { + available_ids.push(i); + } + last_id = id + 1; + } + + while available_ids.len() < 100 { + available_ids.push(last_id); + last_id += 1; + } + + available_ids[rand::thread_rng().gen_range(0..available_ids.len())] + } +} diff --git a/tests/Cargo.toml b/tests/Cargo.toml index 127a5f57..b3b609e6 100644 --- a/tests/Cargo.toml +++ b/tests/Cargo.toml @@ -5,7 +5,8 @@ edition = "2021" resolver = "2" [features] -default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis"] +#default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis"] +default = ["sqlite", "postgres", "mysql", "rocks", "elastic", "s3", "redis", "foundationdb"] sqlite = ["store/sqlite"] foundationdb = ["store/foundation"] postgres = ["store/postgres"] diff --git a/tests/src/directory/internal.rs b/tests/src/directory/internal.rs index f04c1caa..fb669841 100644 --- a/tests/src/directory/internal.rs +++ b/tests/src/directory/internal.rs @@ -21,6 +21,7 @@ * for more details. */ +use ahash::AHashSet; use directory::{ backend::internal::{ lookup::DirectoryStore, manage::ManageDirectory, PrincipalField, PrincipalUpdate, @@ -55,20 +56,18 @@ async fn internal_directory() { ); // Basic account creation - assert_eq!( - store - .create_account( - Principal { - name: "john".to_string(), - description: Some("John Doe".to_string()), - secrets: vec!["secret".to_string(), "secret2".to_string()], - ..Default::default() - }, - vec![] - ) - .await, - Ok(0) - ); + let john_id = store + .create_account( + Principal { + name: "john".to_string(), + description: Some("John Doe".to_string()), + secrets: vec!["secret".to_string(), "secret2".to_string()], + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); // Two accounts with the same name should fail assert_eq!( @@ -125,7 +124,7 @@ async fn internal_directory() { assert!(store.rcpt("john@example.org").await.unwrap()); assert_eq!( store.email_to_ids("john@example.org").await.unwrap(), - vec![0] + vec![john_id] ); // Using non-existent domain should fail @@ -145,27 +144,26 @@ async fn internal_directory() { ); // Create an account with an email address - assert_eq!( - store - .create_account( - Principal { - name: "jane".to_string(), - description: Some("Jane Doe".to_string()), - secrets: vec!["my_secret".to_string(), "my_secret2".to_string()], - emails: vec!["jane@example.org".to_string()], - quota: 123, - ..Default::default() - }, - vec![] - ) - .await, - Ok(1) - ); + let jane_id = store + .create_account( + Principal { + name: "jane".to_string(), + description: Some("Jane Doe".to_string()), + secrets: vec!["my_secret".to_string(), "my_secret2".to_string()], + emails: vec!["jane@example.org".to_string()], + quota: 123, + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); + assert!(store.rcpt("jane@example.org").await.unwrap()); assert!(!store.rcpt("jane@otherdomain.org").await.unwrap()); assert_eq!( store.email_to_ids("jane@example.org").await.unwrap(), - vec![1] + vec![jane_id] ); assert_eq!(store.vrfy("jane").await.unwrap(), vec!["jane@example.org"]); assert_eq!( @@ -180,7 +178,7 @@ async fn internal_directory() { .await .unwrap(), Some(Principal { - id: 1, + id: jane_id, name: "jane".to_string(), description: Some("Jane Doe".to_string()), emails: vec!["jane@example.org".to_string()], @@ -223,20 +221,18 @@ async fn internal_directory() { ); // Create a mailing list - assert_eq!( - store - .create_account( - Principal { - name: "list".to_string(), - typ: Type::List, - emails: vec!["list@example.org".to_string()], - ..Default::default() - }, - vec![] - ) - .await, - Ok(2) - ); + let list_id = store + .create_account( + Principal { + name: "list".to_string(), + typ: Type::List, + emails: vec!["list@example.org".to_string()], + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); assert_eq!( store .update_account( @@ -251,8 +247,13 @@ async fn internal_directory() { ); assert!(store.rcpt("list@example.org").await.unwrap()); assert_eq!( - store.email_to_ids("list@example.org").await.unwrap(), - vec![0, 1] + store + .email_to_ids("list@example.org") + .await + .unwrap() + .into_iter() + .collect::>(), + [john_id, jane_id].into_iter().collect::>(), ); assert_eq!( store @@ -262,46 +263,50 @@ async fn internal_directory() { .unwrap(), Principal { name: "list".to_string(), - id: 2, + id: list_id, typ: Type::List, emails: vec!["list@example.org".to_string()], ..Default::default() } ); assert_eq!( - store.expn("list@example.org").await.unwrap(), - vec!["john@example.org", "jane@example.org"] + store + .expn("list@example.org") + .await + .unwrap() + .into_iter() + .collect::>(), + ["john@example.org", "jane@example.org"] + .into_iter() + .map(|s| s.to_string()) + .collect::>() ); // Create groups - assert_eq!( - store - .create_account( - Principal { - name: "sales".to_string(), - description: Some("Sales Team".to_string()), - typ: Type::Group, - ..Default::default() - }, - vec![] - ) - .await, - Ok(3) - ); - assert_eq!( - store - .create_account( - Principal { - name: "support".to_string(), - description: Some("Support Team".to_string()), - typ: Type::Group, - ..Default::default() - }, - vec![] - ) - .await, - Ok(4) - ); + store + .create_account( + Principal { + name: "sales".to_string(), + description: Some("Sales Team".to_string()), + typ: Type::Group, + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); + store + .create_account( + Principal { + name: "support".to_string(), + description: Some("Support Team".to_string()), + typ: Type::Group, + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); // Add John to the Sales and Support groups assert_eq!( @@ -332,8 +337,10 @@ async fn internal_directory() { .unwrap() ) .await - .unwrap(), + .unwrap() + .into_sorted(), Principal { + id: john_id, name: "john".to_string(), description: Some("John Doe".to_string()), secrets: vec!["secret".to_string(), "secret2".to_string()], @@ -386,8 +393,10 @@ async fn internal_directory() { .unwrap() ) .await - .unwrap(), + .unwrap() + .into_sorted(), Principal { + id: john_id, name: "john".to_string(), description: Some("John Doe".to_string()), secrets: vec!["secret".to_string(), "secret2".to_string()], @@ -443,8 +452,10 @@ async fn internal_directory() { .unwrap() ) .await - .unwrap(), + .unwrap() + .into_sorted(), Principal { + id: john_id, name: "john.doe".to_string(), description: Some("Johnny Doe".to_string()), secrets: vec!["12345".to_string()], @@ -452,7 +463,6 @@ async fn internal_directory() { quota: 1024, typ: Type::Superuser, member_of: vec!["list".to_string(), "sales".to_string()], - ..Default::default() } ); assert_eq!(store.get_account_id("john").await.unwrap(), None); @@ -474,7 +484,7 @@ async fn internal_directory() { ); assert_eq!( store.email_to_ids("list@example.org").await.unwrap(), - vec![1] + vec![jane_id] ); assert_eq!( store @@ -489,8 +499,13 @@ async fn internal_directory() { Ok(()) ); assert_eq!( - store.email_to_ids("list@example.org").await.unwrap(), - vec![0, 1] + store + .email_to_ids("list@example.org") + .await + .unwrap() + .into_iter() + .collect::>(), + [john_id, jane_id].into_iter().collect::>() ); // Field validation @@ -527,8 +542,16 @@ async fn internal_directory() { // List accounts assert_eq!( - store.list_accounts(None, None).await.unwrap(), - vec!["jane", "john.doe", "list", "sales", "support"] + store + .list_accounts(None, None) + .await + .unwrap() + .into_iter() + .collect::>(), + ["jane", "john.doe", "list", "sales", "support"] + .into_iter() + .map(|s| s.to_string()) + .collect::>() ); assert_eq!( store.list_accounts("john".into(), None).await.unwrap(), @@ -538,12 +561,25 @@ async fn internal_directory() { store .list_accounts(None, Type::Individual.into()) .await - .unwrap(), - vec!["jane", "john.doe"] + .unwrap() + .into_iter() + .collect::>(), + ["jane", "john.doe"] + .into_iter() + .map(|s| s.to_string()) + .collect::>() ); assert_eq!( - store.list_accounts(None, Type::Group.into()).await.unwrap(), - vec!["sales", "support"] + store + .list_accounts(None, Type::Group.into()) + .await + .unwrap() + .into_iter() + .collect::>(), + ["sales", "support"] + .into_iter() + .map(|s| s.to_string()) + .collect::>() ); assert_eq!( store.list_accounts(None, Type::List.into()).await.unwrap(), @@ -551,21 +587,20 @@ async fn internal_directory() { ); // Write records on John's and Jane's accounts - for account_id in [0, 1] { - let document_id = store - .assign_document_id(account_id, Collection::Email) - .await - .unwrap(); - store + let mut document_id = u32::MAX; + for account_id in [john_id, jane_id] { + document_id = store .write( BatchBuilder::new() .with_account_id(account_id) .with_collection(Collection::Email) - .create_document(document_id) + .create_document() .set(ValueClass::Property(0), "hello".as_bytes()) .build_batch(), ) .await + .unwrap() + .last_document_id() .unwrap(); assert_eq!( store @@ -582,7 +617,7 @@ async fn internal_directory() { } // Delete John's account and make sure his records are gone - store.delete_account(QueryBy::Id(0)).await.unwrap(); + store.delete_account(QueryBy::Id(john_id)).await.unwrap(); assert_eq!(store.get_account_id("john.doe").await.unwrap(), None); assert_eq!( store.email_to_ids("john.doe@example.org").await.unwrap(), @@ -590,16 +625,24 @@ async fn internal_directory() { ); assert!(!store.rcpt("john.doe@example.org").await.unwrap()); assert_eq!( - store.list_accounts(None, None).await.unwrap(), - vec!["jane", "list", "sales", "support"] + store + .list_accounts(None, None) + .await + .unwrap() + .into_iter() + .collect::>(), + ["jane", "list", "sales", "support"] + .into_iter() + .map(|s| s.to_string()) + .collect::>() ); assert_eq!( store .get_bitmap(BitmapKey { - account_id: 0, + account_id: john_id, collection: Collection::Email.into(), class: BitmapClass::DocumentIds, - block_num: 0 + document_id: 0 }) .await .unwrap(), @@ -608,7 +651,7 @@ async fn internal_directory() { assert_eq!( store .get_value::(ValueKey { - account_id: 0, + account_id: john_id, collection: Collection::Email.into(), document_id: 0, class: ValueClass::Property(0) @@ -619,30 +662,30 @@ async fn internal_directory() { ); // Make sure Jane's records are still there - assert_eq!(store.get_account_id("jane").await.unwrap(), Some(1)); + assert_eq!(store.get_account_id("jane").await.unwrap(), Some(jane_id)); assert_eq!( store.email_to_ids("jane@example.org").await.unwrap(), - vec![1] + vec![jane_id] ); assert!(store.rcpt("jane@example.org").await.unwrap()); assert_eq!( store .get_bitmap(BitmapKey { - account_id: 1, + account_id: jane_id, collection: Collection::Email.into(), class: BitmapClass::DocumentIds, - block_num: 0 + document_id: 0 }) .await .unwrap(), - Some(RoaringBitmap::from_sorted_iter([0]).unwrap()) + Some(RoaringBitmap::from_sorted_iter([document_id]).unwrap()) ); assert_eq!( store .get_value::(ValueKey { - account_id: 1, + account_id: jane_id, collection: Collection::Email.into(), - document_id: 0, + document_id, class: ValueClass::Property(0) }) .await diff --git a/tests/src/directory/ldap.rs b/tests/src/directory/ldap.rs index 929225c5..770c0797 100644 --- a/tests/src/directory/ldap.rs +++ b/tests/src/directory/ldap.rs @@ -26,7 +26,7 @@ use std::fmt::Debug; use directory::{backend::internal::manage::ManageDirectory, Principal, QueryBy, Type}; use mail_send::Credentials; -use crate::directory::{map_account_ids, DirectoryTest, IntoSortedPrincipal}; +use crate::directory::{map_account_ids, DirectoryTest}; #[tokio::test] async fn ldap_directory() { diff --git a/tests/src/directory/mod.rs b/tests/src/directory/mod.rs index 614fae62..4f9a4d23 100644 --- a/tests/src/directory/mod.rs +++ b/tests/src/directory/mod.rs @@ -28,7 +28,7 @@ pub mod smtp; pub mod sql; use common::{config::smtp::session::AddressMapping, Core}; -use directory::{backend::internal::manage::ManageDirectory, Directories, Principal}; +use directory::{backend::internal::manage::ManageDirectory, Directories}; use mail_send::Credentials; use rustls::ServerConfig; use rustls_pemfile::{certs, pkcs8_private_keys}; @@ -669,15 +669,3 @@ async fn map_account_ids(store: &Store, names: Vec>) -> Vec } ids } - -trait IntoSortedPrincipal: Sized { - fn into_sorted(self) -> Self; -} - -impl IntoSortedPrincipal for Principal { - fn into_sorted(mut self) -> Self { - self.member_of.sort_unstable(); - self.emails.sort_unstable(); - self - } -} diff --git a/tests/src/imap/mod.rs b/tests/src/imap/mod.rs index 3f01018f..799a0e43 100644 --- a/tests/src/imap/mod.rs +++ b/tests/src/imap/mod.rs @@ -35,7 +35,11 @@ pub mod search; pub mod store; pub mod thread; -use std::{path::PathBuf, sync::Arc, time::Duration}; +use std::{ + path::PathBuf, + sync::Arc, + time::{Duration, Instant}, +}; use ::managesieve::core::ManageSieveSessionManager; use common::{ @@ -394,6 +398,7 @@ pub async fn imap_tests() { } // Prepare settings + let start_time = Instant::now(); let delete = true; let handle = init_imap_tests( &std::env::var("STORE") @@ -448,6 +453,14 @@ pub async fn imap_tests() { // Run ManageSieve tests managesieve::test().await; + // Print elapsed time + let elapsed = start_time.elapsed(); + println!( + "Elapsed: {}.{:03}s", + elapsed.as_secs(), + elapsed.subsec_millis() + ); + // Remove test data if delete { handle.temp_dir.delete(); diff --git a/tests/src/jmap/auth_limits.rs b/tests/src/jmap/auth_limits.rs index 9245122a..70d2b1f9 100644 --- a/tests/src/jmap/auth_limits.rs +++ b/tests/src/jmap/auth_limits.rs @@ -243,13 +243,12 @@ pub async fn test(params: &mut JMAPTest) { for _ in 0..8 { let client_ = client.clone(); tokio::spawn(async move { - client_ + let _ = client_ .mailbox_query( mailbox::query::Filter::name("__sleep").into(), [mailbox::query::Comparator::name()].into(), ) - .await - .unwrap(); + .await; }); } tokio::time::sleep(Duration::from_millis(500)).await; diff --git a/tests/src/jmap/email_query.rs b/tests/src/jmap/email_query.rs index 9889e0c6..08b31046 100644 --- a/tests/src/jmap/email_query.rs +++ b/tests/src/jmap/email_query.rs @@ -33,11 +33,11 @@ use jmap_client::{ email, }; use jmap_proto::types::{collection::Collection, id::Id, property::Property}; -use mail_parser::HeaderName; +use mail_parser::{DateTime, HeaderName}; use store::{ ahash::AHashMap, - write::{BatchBuilder, ValueClass}, + write::{now, BatchBuilder, ValueClass}, }; use super::JMAPTest; @@ -61,7 +61,7 @@ pub async fn test(params: &mut JMAPTest, insert: bool) { .with_account_id(account_id) .with_collection(Collection::Mailbox); for mailbox_id in 1545..3010 { - batch.create_document(mailbox_id); + batch.create_document_with_id(mailbox_id); } server.core.storage.data.write(batch.build()).await.unwrap(); @@ -81,25 +81,14 @@ pub async fn test(params: &mut JMAPTest, insert: bool) { } server.core.storage.data.write(batch.build()).await.unwrap(); - for thread_id in 0..MAX_THREADS { - assert!( - client - .thread_get(&Id::new(thread_id as u64).to_string()) - .await - .unwrap() - .is_some(), - "thread {} not found", - thread_id - ); - } - - assert!( - client - .thread_get(&Id::new(MAX_THREADS as u64).to_string()) + assert_eq!( + params + .server + .get_document_ids(account_id, Collection::Thread) .await .unwrap() - .is_none(), - "thread {} found", + .unwrap() + .len() as usize, MAX_THREADS ); @@ -136,7 +125,10 @@ pub async fn query(client: &mut Client) { (email::query::Filter::after(1850)), (email::query::Filter::from("george")), ]), - vec![email::query::Comparator::subject()], + vec![ + email::query::Comparator::subject(), + email::query::Comparator::sent_at(), + ], vec![ "N01389", "T10115", "N00618", "N03500", "T01587", "T00397", "N01561", "N05250", "N03973", "N04973", "N04057", "N01940", "N01539", "N01612", "N04484", "N01954", @@ -229,7 +221,10 @@ pub async fn query(client: &mut Client) { (email::query::Filter::none_in_thread_have_keyword("N")), (email::query::Filter::after(1995)), ]), - vec![email::query::Comparator::from()], + vec![ + email::query::Comparator::from(), + email::query::Comparator::sent_at(), + ], vec![ "AR00163", "AR00164", "AR00472", "P11481", "AR00066", "AR00178", "P77895", "P77896", "P77897", @@ -252,6 +247,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::all_in_thread_have_keyword("N"), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "N01496", "N05916", "N01046", "N00675", "N01320", "N01321", "N00273", "N01453", @@ -266,6 +262,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::all_in_thread_have_keyword("N").descending(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "T09417", "T01882", "T08820", "N04689", "T08891", "T00986", "N00316", "N03544", @@ -283,6 +280,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::some_in_thread_have_keyword("Bronze"), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "N04326", "N01610", "N02920", "N01587", "T00167", "T00168", "N01554", "N01535", @@ -298,6 +296,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::some_in_thread_have_keyword("Bronze").descending(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "N01559", "N02123", "N01940", "N03594", "N01494", "N04271", "N04326", "N01610", @@ -314,6 +313,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::has_keyword("attributed to"), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "T09455", "T09334", "T10965", "T08626", "T09417", "T08951", "T01851", "T01852", @@ -333,6 +333,7 @@ pub async fn query(client: &mut Client) { vec![ email::query::Comparator::has_keyword("attributed to").descending(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], vec![ "T09417", "T08951", "T01851", "T01852", "T08761", "T08123", "T08756", "T10561", @@ -398,6 +399,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: None, @@ -419,6 +421,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 10, anchor: None, @@ -440,6 +443,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: -10, anchor: None, @@ -461,6 +465,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: -20, anchor: None, @@ -482,6 +487,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: -100000, anchor: None, @@ -497,6 +503,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: -1, anchor: None, @@ -512,6 +519,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "N01205").await, @@ -533,6 +541,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "N01205").await, @@ -554,6 +563,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "N01205").await, @@ -575,6 +585,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "N01496").await, @@ -590,6 +601,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "AR00164").await, @@ -605,6 +617,7 @@ pub async fn query_options(client: &mut Client) { sort: vec![ email::query::Comparator::subject(), email::query::Comparator::from(), + email::query::Comparator::sent_at(), ], position: 0, anchor: get_anchor(client, "AR00164").await, @@ -674,6 +687,7 @@ pub async fn query_options(client: &mut Client) { } pub async fn create(client: &mut Client) { + let sent_at = now(); let now = Instant::now(); let mut fields = AHashMap::default(); for (field_num, field) in FIELDS.iter().enumerate() { @@ -685,10 +699,11 @@ pub async fn create(client: &mut Client) { let mut thread_count = AHashMap::default(); let mut artist_count = AHashMap::default(); - 'outer: for record in csv::ReaderBuilder::new() + 'outer: for (idx, record) in csv::ReaderBuilder::new() .has_headers(true) .from_reader(&deflate_test_resource("artwork_data.csv.gz")[..]) .records() + .enumerate() { let record = record.unwrap(); let mut values_str = AHashMap::default(); @@ -752,10 +767,11 @@ pub async fn create(client: &mut Client) { .email_import( format!( concat!( - "From: \"{}\" \nCc: \"{}\" \nMessage-ID: <{}>\n", + "Date: {}\nFrom: \"{}\" \nCc: \"{}\" \nMessage-ID: <{}>\n", "References: <{}>\nComments: {}\nSubject: [{}]", " Year {}\n\n{}\n{}\n" ), + DateTime::from_timestamp(sent_at as i64 + idx as i64).to_rfc822(), values_str["artist"], values_str["medium"], values_str["accession_number"], diff --git a/tests/src/jmap/email_query_changes.rs b/tests/src/jmap/email_query_changes.rs index c07b6165..e8c18c5e 100644 --- a/tests/src/jmap/email_query_changes.rs +++ b/tests/src/jmap/email_query_changes.rs @@ -30,7 +30,7 @@ use jmap_proto::types::{collection::Collection, id::Id, property::Property, stat use store::{ ahash::{AHashMap, AHashSet}, - write::{log::ChangeLogBuilder, BatchBuilder, F_BITMAP, F_CLEAR, F_VALUE}, + write::{log::ChangeLogBuilder, BatchBuilder, MaybeDynamicId, TagValue, F_BITMAP, F_CLEAR}, }; use crate::jmap::{ @@ -156,11 +156,16 @@ pub async fn test(params: &mut JMAPTest) { BatchBuilder::new() .with_account_id(1) .with_collection(Collection::Thread) - .create_document(thread_id) + .create_document() .with_collection(Collection::Email) .update_document(id.document_id()) .value(Property::ThreadId, id.prefix_id(), F_BITMAP | F_CLEAR) - .value(Property::ThreadId, thread_id, F_VALUE | F_BITMAP) + .set(Property::ThreadId, MaybeDynamicId::Dynamic(0)) + .tag( + Property::ThreadId, + TagValue::Id(MaybeDynamicId::Dynamic(0)), + 0, + ) .custom(server.begin_changes(1).await.unwrap().with_log_move( Collection::Email, id, diff --git a/tests/src/jmap/email_set.rs b/tests/src/jmap/email_set.rs index 024d1413..ac0ffa06 100644 --- a/tests/src/jmap/email_set.rs +++ b/tests/src/jmap/email_set.rs @@ -24,6 +24,7 @@ use std::{fs, path::PathBuf}; use crate::jmap::{assert_is_empty, mailbox::destroy_all_mailboxes}; +use ahash::AHashSet; use jmap::mailbox::INBOX_ID; use jmap_client::{ client::Client, @@ -326,11 +327,17 @@ pub async fn assert_email_properties( .unwrap() .unwrap(); - let mut mailbox_ids_ = result.mailbox_ids().to_vec(); - let mut keywords_ = result.keywords().to_vec(); - mailbox_ids_.sort_unstable(); - keywords_.sort_unstable(); + assert_eq!( + mailbox_ids.iter().copied().collect::>(), + result + .mailbox_ids() + .iter() + .copied() + .collect::>() + ); - assert_eq!(mailbox_ids_, mailbox_ids); - assert_eq!(keywords_, keywords); + assert_eq!( + keywords.iter().copied().collect::>(), + result.keywords().iter().copied().collect::>() + ); } diff --git a/tests/src/jmap/mailbox.rs b/tests/src/jmap/mailbox.rs index b92606c0..49705f18 100644 --- a/tests/src/jmap/mailbox.rs +++ b/tests/src/jmap/mailbox.rs @@ -64,9 +64,9 @@ pub async fn test(params: &mut JMAPTest) { "drafts", "spam2", "inbox", - "1", - "2", - "3", + "l.1", + "l.2", + "l.3", "sent", "spam", "1.1", @@ -99,15 +99,15 @@ pub async fn test(params: &mut JMAPTest) { [ "drafts", "inbox", - "1", + "l.1", "1.1", "1.1.1", "1.1.1.1", "1.1.1.1.1", "1.2", "1.2.1", - "2", - "3", + "l.2", + "l.3", "sent", "spam", "spam1", @@ -134,15 +134,15 @@ pub async fn test(params: &mut JMAPTest) { .map(|id| id_map.get(id).unwrap()) .collect::>(), [ - "1", + "l.1", "1.1", "1.1.1", "1.1.1.1", "1.1.1.1.1", "1.2", "1.2.1", - "2", - "3" + "l.2", + "l.3" ] ); @@ -231,31 +231,35 @@ pub async fn test(params: &mut JMAPTest) { // Duplicate name let mut request = client.build(); - request.set_mailbox().update(&id_map["2"]).name("Level 3"); - assert!(matches!( - request - .send_set_mailbox() - .await - .unwrap() - .updated(&id_map["2"]), - Err(Error::Set(SetError { - type_: SetErrorType::InvalidProperties, - .. - })) - )); + request.set_mailbox().update(&id_map["l.2"]).name("Level 3"); + let result = request + .send_set_mailbox() + .await + .unwrap() + .updated(&id_map["l.2"]); + assert!( + matches!( + result, + Err(Error::Set(SetError { + type_: SetErrorType::InvalidProperties, + .. + })) + ), + "{result:?}", + ); // Circular relationship let mut request = client.build(); request .set_mailbox() - .update(&id_map["1"]) + .update(&id_map["l.1"]) .parent_id((&id_map["1.1.1.1.1"]).into()); assert!(matches!( request .send_set_mailbox() .await .unwrap() - .updated(&id_map["1"]), + .updated(&id_map["l.1"]), Err(Error::Set(SetError { type_: SetErrorType::InvalidProperties, .. @@ -265,14 +269,14 @@ pub async fn test(params: &mut JMAPTest) { let mut request = client.build(); request .set_mailbox() - .update(&id_map["1"]) - .parent_id((&id_map["1"]).into()); + .update(&id_map["l.1"]) + .parent_id((&id_map["l.1"]).into()); assert!(matches!( request .send_set_mailbox() .await .unwrap() - .updated(&id_map["1"]), + .updated(&id_map["l.1"]), Err(Error::Set(SetError { type_: SetErrorType::InvalidProperties, .. @@ -283,14 +287,14 @@ pub async fn test(params: &mut JMAPTest) { let mut request = client.build(); request .set_mailbox() - .update(&id_map["1"]) + .update(&id_map["l.1"]) .parent_id(Id::new(u64::MAX).to_string().into()); assert!(matches!( request .send_set_mailbox() .await .unwrap() - .updated(&id_map["1"]), + .updated(&id_map["l.1"]), Err(Error::Set(SetError { type_: SetErrorType::InvalidProperties, .. @@ -311,7 +315,7 @@ pub async fn test(params: &mut JMAPTest) { .set_mailbox() .update(&id_map["1.1.1.1.1"]) .name("Renamed and moved") - .parent_id((&id_map["2"]).into()); + .parent_id((&id_map["l.2"]).into()); assert!(request .send_set_mailbox() .await @@ -472,13 +476,13 @@ pub async fn test(params: &mut JMAPTest) { // Deleting folders with children is not allowed let mut request = client.build(); - request.set_mailbox().destroy([&id_map["1"]]); + request.set_mailbox().destroy([&id_map["l.1"]]); assert!(matches!( request .send_set_mailbox() .await .unwrap() - .destroyed(&id_map["1"]), + .destroyed(&id_map["l.1"]), Err(Error::Set(SetError { type_: SetErrorType::MailboxHasChild, .. @@ -533,7 +537,7 @@ pub async fn test(params: &mut JMAPTest) { .update(&id_map["drafts"]) .name("Borradores") .sort_order(100) - .parent_id((&id_map["2"]).into()) + .parent_id((&id_map["l.2"]).into()) .role(Role::None); assert!(request .send_set_mailbox() @@ -546,7 +550,7 @@ pub async fn test(params: &mut JMAPTest) { .mailbox_query( Filter::and([ mailbox::query::Filter::name("Borradores").into(), - mailbox::query::Filter::parent_id((&id_map["2"]).into()).into(), + mailbox::query::Filter::parent_id((&id_map["l.2"]).into()).into(), Filter::not([mailbox::query::Filter::has_any_role(true)]) ]) .into(), @@ -691,7 +695,7 @@ const TEST_MAILBOXES: &[u8] = br#" "children": [ { "name": "Level 1", - "id": "1", + "id": "l.1", "order": 4, "children": [ { @@ -737,12 +741,12 @@ const TEST_MAILBOXES: &[u8] = br#" }, { "name": "Level 2", - "id": "2", + "id": "l.2", "order": 8 }, { "name": "Level 3", - "id": "3", + "id": "l.3", "order": 9 } ] diff --git a/tests/src/jmap/mod.rs b/tests/src/jmap/mod.rs index 9d1b3b2e..a5f8e9b7 100644 --- a/tests/src/jmap/mod.rs +++ b/tests/src/jmap/mod.rs @@ -21,7 +21,7 @@ * for more details. */ -use std::{sync::Arc, time::Duration}; +use std::{path::PathBuf, sync::Arc, time::Duration}; use base64::{ engine::general_purpose::{self, STANDARD}, @@ -29,6 +29,7 @@ use base64::{ }; use common::{ config::server::{ServerProtocol, Servers}, + manager::config::{ConfigManager, Patterns}, Core, }; use hyper::{header::AUTHORIZATION, Method}; @@ -80,7 +81,7 @@ pub mod websocket; const SERVER: &str = r#" [server] hostname = "'jmap.example.org'" -url = "'https://127.0.0.1:8899'" +http.url = "'https://127.0.0.1:8899'" [server.listener.jmap] bind = ["127.0.0.1:8899"] @@ -340,7 +341,7 @@ pub async fn jmap_stress_tests() { .with_env_filter( tracing_subscriber::EnvFilter::builder() .parse( - format!("smtp={level},imap={level},jmap={level},store={level},utils={level},directory={level}"), + format!("smtp={level},imap={level},jmap={level},store={level},utils={level},directory={level},common={level}"), ) .unwrap(), ) @@ -419,7 +420,17 @@ async fn init_jmap_tests(store_id: &str, delete_if_exists: bool) -> JMAPTest { let stores = Stores::parse_all(&mut config).await; // Parse core - let core = Core::parse(&mut config, stores, Default::default()).await; + let config_manager = ConfigManager { + cfg_local: Default::default(), + cfg_local_path: PathBuf::new(), + cfg_local_patterns: Patterns::parse(&mut config).into(), + cfg_store: config + .value("storage.data") + .and_then(|id| stores.stores.get(id)) + .cloned() + .unwrap_or_default(), + }; + let core = Core::parse(&mut config, stores, config_manager).await; let store = core.storage.data.clone(); let shared_core = core.into_shared(); diff --git a/tests/src/jmap/push_subscription.rs b/tests/src/jmap/push_subscription.rs index de2c3f08..8936528f 100644 --- a/tests/src/jmap/push_subscription.rs +++ b/tests/src/jmap/push_subscription.rs @@ -61,23 +61,20 @@ use super::JMAPTest; const SERVER: &str = r#" [server] hostname = "'jmap-push.example.org'" -url = "'https://127.0.0.1:9000'" +http.url = "'https://127.0.0.1:9000'" [server.listener.jmap] bind = ['127.0.0.1:9000'] protocol = 'http' +tls.implicit = true [server.socket] reuse-addr = true -[server.tls] -enable = true -implicit = false -certificate = 'default' - [certificate.default] cert = '%{file:{CERT}}%' private-key = '%{file:{PK}}%' +default = true "#; pub async fn test(params: &mut JMAPTest) { @@ -119,7 +116,11 @@ pub async fn test(params: &mut JMAPTest) { // Start mock push server let mut settings = Config::new(add_test_certs(SERVER)).unwrap(); settings.resolve_macros().await; - let mock_core = Core::default().into_shared(); + let mock_core = Core::parse(&mut settings, Default::default(), Default::default()) + .await + .into_shared(); + settings.errors.clear(); + settings.warnings.clear(); let mut servers = Servers::parse(&mut settings); servers.parse_tcp_acceptors(&mut settings, mock_core.clone()); @@ -304,16 +305,7 @@ impl common::listener::SessionManager for SessionManager { let _ = http1::Builder::new() .keep_alive(false) .serve_connection( - TokioIo::new( - session - .instance - .acceptor - .accept(session.stream, None) - .await - .unwrap_tls() - .await - .unwrap(), - ), + TokioIo::new(session.stream), service_fn(|mut req: hyper::Request| { let push = push.clone(); diff --git a/tests/src/jmap/quota.rs b/tests/src/jmap/quota.rs index 661e3b33..8b6f62ed 100644 --- a/tests/src/jmap/quota.rs +++ b/tests/src/jmap/quota.rs @@ -318,7 +318,7 @@ pub async fn test(params: &mut JMAPTest) { "jane@example.com", "robert@example.com", &format!("Ingest test {i}"), - 100, + 513, )) .unwrap(), ) diff --git a/tests/src/jmap/thread_merge.rs b/tests/src/jmap/thread_merge.rs index 94589afe..798cae31 100644 --- a/tests/src/jmap/thread_merge.rs +++ b/tests/src/jmap/thread_merge.rs @@ -27,7 +27,7 @@ use crate::{ jmap::{assert_is_empty, mailbox::destroy_all_mailboxes}, store::deflate_test_resource, }; -use jmap::{email::ingest::IngestEmail, mailbox::INBOX_ID, IngestError}; +use jmap::{email::ingest::IngestEmail, IngestError}; use jmap_client::{email, mailbox::Role}; use jmap_proto::types::{collection::Collection, id::Id}; use mail_parser::{mailbox::mbox::MessageIterator, MessageParser}; @@ -229,14 +229,20 @@ async fn test_multi_thread(params: &mut JMAPTest) { println!("Running Email Merge Threads tests (multi-threaded)..."); //let semaphore = sync::Arc::Arc::new(tokio::sync::Semaphore::new(100)); let mut handles = vec![]; - for num in 0..3 { + + let mailbox_id = Id::from_bytes( params .client .set_default_account_id(Id::new(0u64).to_string()) - .mailbox_create(format!("Mailbox {num}"), None::, Role::None) + .mailbox_create("Inbox", None::, Role::None) .await - .unwrap(); - } + .unwrap() + .id() + .unwrap() + .as_bytes(), + ) + .unwrap() + .document_id(); for message in MessageIterator::new(Cursor::new(deflate_test_resource("mailbox.gz"))) .collect::>() @@ -255,7 +261,7 @@ async fn test_multi_thread(params: &mut JMAPTest) { message: MessageParser::new().parse(message.contents()), account_id: 0, account_quota: 0, - mailbox_ids: vec![INBOX_ID], + mailbox_ids: vec![mailbox_id], keywords: vec![], received_at: None, skip_duplicates: true, diff --git a/tests/src/store/assign_id.rs b/tests/src/store/assign_id.rs index 1e8ed0b2..3c551c42 100644 --- a/tests/src/store/assign_id.rs +++ b/tests/src/store/assign_id.rs @@ -21,29 +21,19 @@ * for more details. */ -use std::{collections::HashSet, time::Duration}; +use std::collections::HashSet; -use store::ahash::AHashSet; - -use store::backend::ID_ASSIGNMENT_EXPIRY; use store::{write::BatchBuilder, Store}; pub async fn test(db: Store) { println!("Running Store ID assignment tests..."); - test_0(db.clone()).await; - test_1(db.clone()).await; - test_2(db.clone()).await; - test_3(db.clone()).await; - test_4(db).await; - - ID_ASSIGNMENT_EXPIRY.store(60 * 60, std::sync::atomic::Ordering::Relaxed); + test_0(db).await; } async fn test_0(db: Store) { // Test document id assignment - println!("Assigning 1000 ids concurrently..."); - ID_ASSIGNMENT_EXPIRY.store(10 * 60 * 60, std::sync::atomic::Ordering::Relaxed); + println!("Creating 1000 documentIds concurrently..."); let mut handles = Vec::new(); let mut assigned_ids = HashSet::new(); @@ -51,7 +41,19 @@ async fn test_0(db: Store) { for _ in 0..1000 { handles.push({ let db = db.clone(); - tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await.unwrap() }) + tokio::spawn(async move { + db.write( + BatchBuilder::new() + .with_account_id(0) + .with_collection(u8::MAX) + .create_document() + .build_batch(), + ) + .await + .unwrap() + .last_document_id() + .unwrap() + }) }); } @@ -64,72 +66,48 @@ async fn test_0(db: Store) { } assert_eq!(assigned_ids.len(), 1000); - db.destroy().await; -} - -async fn test_1(db: Store) { - // Test document id assignment - ID_ASSIGNMENT_EXPIRY.store(2, std::sync::atomic::Ordering::Relaxed); - println!("Assigning 100 ids concurrently and reassign after expiration..."); - for wait_for_expiry in [true, false] { - let mut handles = Vec::new(); - let mut assigned_ids = HashSet::new(); - - // Create 100 ids concurrently - for _ in 0..100 { - handles.push({ - let db = db.clone(); - tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await.unwrap() }) - }); - } - - for handle in handles { - let assigned_id = handle.await.unwrap(); - //println!("assigned id: {assigned_id} ({wait_for_expiry})"); - assert!( - assigned_ids.insert(assigned_id), - "already assigned or invalid: {assigned_id} ({wait_for_expiry})" - ); - } - assert_eq!( - assigned_ids.len(), - 100, - "{assigned_ids:?} ({wait_for_expiry})" - ); - - if wait_for_expiry { - tokio::time::sleep(Duration::from_secs(3)).await; - } + // Create 1000 ids concurrently + println!("Deleting 1000 documentIds concurrently..."); + let mut handles = Vec::new(); + for document_id in assigned_ids { + let db = db.clone(); + handles.push({ + tokio::spawn(async move { + db.write( + BatchBuilder::new() + .with_account_id(0) + .with_collection(u8::MAX) + .delete_document(document_id) + .build_batch(), + ) + .await + .unwrap(); + }) + }); + } + for handle in handles { + handle.await.unwrap(); } - db.destroy().await; -} - -async fn test_2(db: Store) { - // Test document id assignment + // Reuse 1000 ids concurrently + println!("Reusing 1000 freed documentIds concurrently..."); let mut handles = Vec::new(); let mut assigned_ids = HashSet::new(); - - // Create 1000 ids concurrently - println!("Create 1000 documentIds concurrently..."); - ID_ASSIGNMENT_EXPIRY.store(10 * 60 * 60, std::sync::atomic::Ordering::Relaxed); for _ in 0..1000 { handles.push({ let db = db.clone(); tokio::spawn(async move { - { - let id = db.assign_document_id(0, u8::MAX).await.unwrap(); - db.write( - BatchBuilder::new() - .with_account_id(0) - .with_collection(u8::MAX) - .create_document(id) - .build_batch(), - ) - .await - .unwrap(); - id - } + db.write( + BatchBuilder::new() + .with_account_id(0) + .with_collection(u8::MAX) + .create_document() + .build_batch(), + ) + .await + .unwrap() + .last_document_id() + .unwrap() }) }); } @@ -138,72 +116,10 @@ async fn test_2(db: Store) { let assigned_id = handle.await.unwrap(); assert!( assigned_ids.insert(assigned_id), - "already assigned or invalid: {assigned_id}" + "freed id already assigned or invalid: {assigned_id}" ); } - assert_eq!(assigned_ids.len(), 1000, "{assigned_ids:?} "); - - db.destroy().await; -} - -async fn test_3(db: Store) { - // Create document ids and try reassigning - println!("Assigning 100 ids concurrently and try reassigning..."); - - ID_ASSIGNMENT_EXPIRY.store(2, std::sync::atomic::Ordering::Relaxed); - let mut expected_ids = AHashSet::new(); - let mut batch = BatchBuilder::new(); - batch.with_account_id(0).with_collection(u8::MAX); - for pos in 0..100 { - let id = db.assign_document_id(0, u8::MAX).await.unwrap(); - if pos % 2 == 0 { - batch.create_document(id); - } else { - expected_ids.insert(id); - } - } - db.write(batch.build()).await.unwrap(); - - // Wait for ids to expire - tokio::time::sleep(Duration::from_secs(3)).await; - - for _ in 0..expected_ids.len() { - let id = db.assign_document_id(0, u8::MAX).await.unwrap(); - assert!( - expected_ids.remove(&id), - "already assigned or invalid: {id}" - ); - } - assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 100); - assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 101); - - db.destroy().await; -} - -async fn test_4(db: Store) { - // Try reassigning deleted ids - println!("Create and delete 100 documentIds then try reassigning ids..."); - ID_ASSIGNMENT_EXPIRY.store(60 * 60, std::sync::atomic::Ordering::Relaxed); - let mut expected_ids = AHashSet::new(); - let mut batch = BatchBuilder::new(); - batch.with_account_id(0).with_collection(u8::MAX); - for id in 0..100 { - if id % 2 == 0 { - batch.create_document(id); - } else { - expected_ids.insert(id); - } - } - db.write(batch.build()).await.unwrap(); - for _ in 0..expected_ids.len() { - let id = db.assign_document_id(0, u8::MAX).await.unwrap(); - assert!( - expected_ids.remove(&id), - "already assigned or invalid: {id}" - ); - } - assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 100); - assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 101); + assert_eq!(assigned_ids.len(), 1000); db.destroy().await; } diff --git a/tests/src/store/mod.rs b/tests/src/store/mod.rs index 91ba18e4..28035d63 100644 --- a/tests/src/store/mod.rs +++ b/tests/src/store/mod.rs @@ -109,9 +109,9 @@ pub async fn store_tests() { } import_export::test(store.clone()).await; + assign_id::test(store.clone()).await; ops::test(store.clone()).await; query::test(store.clone(), FtsStore::Store(store.clone()), insert).await; - assign_id::test(store).await; if insert { temp_dir.delete(); diff --git a/tests/src/store/ops.rs b/tests/src/store/ops.rs index ac32e10d..279bea81 100644 --- a/tests/src/store/ops.rs +++ b/tests/src/store/ops.rs @@ -23,15 +23,108 @@ use std::collections::HashSet; +use jmap_proto::types::{collection::Collection, property::Property}; use store::{ - write::{BatchBuilder, DirectoryClass, ValueClass}, - Store, ValueKey, + write::{ + BatchBuilder, BitmapClass, DirectoryClass, MaybeDynamicId, TagValue, ValueClass, F_CLEAR, + }, + BitmapKey, Store, ValueKey, }; // FDB max value const MAX_VALUE_SIZE: usize = 100000; pub async fn test(db: Store) { + // Testing ID assignment + println!("Running dynamic ID assignment tests..."); + let mut builder = BatchBuilder::new(); + builder + .with_account_id(0) + .with_collection(Collection::Thread) + .create_document() + .with_collection(Collection::Email) + .create_document() + .tag( + Property::ThreadId, + TagValue::Id(MaybeDynamicId::Dynamic(0)), + 0, + ) + .set(Property::ThreadId, MaybeDynamicId::Dynamic(0)); + + let assigned_ids = db.write(builder.build_batch()).await.unwrap(); + assert_eq!(assigned_ids.document_ids.len(), 2); + let thread_id = assigned_ids.first_document_id().unwrap(); + let email_id = assigned_ids.last_document_id().unwrap(); + + let email_ids = db + .get_bitmap(BitmapKey { + account_id: 0, + collection: Collection::Email.into(), + class: BitmapClass::DocumentIds, + document_id: 0, + }) + .await + .unwrap() + .unwrap(); + assert_eq!(email_ids.len(), 1); + assert!(email_ids.contains(email_id)); + + let thread_ids = db + .get_bitmap(BitmapKey { + account_id: 0, + collection: Collection::Thread.into(), + class: BitmapClass::DocumentIds, + document_id: 0, + }) + .await + .unwrap() + .unwrap(); + assert_eq!(thread_ids.len(), 1); + assert!(thread_ids.contains(thread_id)); + + let tagged_ids = db + .get_bitmap(BitmapKey { + account_id: 0, + collection: Collection::Email.into(), + class: BitmapClass::Tag { + field: Property::ThreadId.into(), + value: TagValue::Id(thread_id), + }, + document_id: 0, + }) + .await + .unwrap() + .unwrap(); + assert_eq!(tagged_ids.len(), 1); + assert!(tagged_ids.contains(email_id)); + + let stored_thread_id = db + .get_value::(ValueKey { + account_id: 0, + collection: Collection::Email.into(), + document_id: email_id, + class: ValueClass::Property(Property::ThreadId.into()), + }) + .await + .unwrap() + .unwrap(); + assert_eq!(stored_thread_id, thread_id); + + let mut builder = BatchBuilder::new(); + builder + .with_account_id(0) + .with_collection(Collection::Thread) + .delete_document(thread_id) + .with_collection(Collection::Email) + .delete_document(email_id) + .tag( + Property::ThreadId, + TagValue::Id(MaybeDynamicId::Static(thread_id)), + F_CLEAR, + ) + .clear(Property::ThreadId); + db.write(builder.build_batch()).await.unwrap(); + // Increment a counter 1000 times concurrently let mut handles = Vec::new(); let mut assigned_ids = HashSet::new(); @@ -46,7 +139,11 @@ pub async fn test(db: Store) { .with_collection(0) .update_document(0) .add_and_get(ValueClass::Directory(DirectoryClass::UsedQuota(0)), 1); - db.write(builder.build_batch()).await.unwrap().unwrap() + db.write(builder.build_batch()) + .await + .unwrap() + .last_counter_id() + .unwrap() }) }); } @@ -100,8 +197,8 @@ pub async fn test(db: Store) { .with_collection(0) .update_document(0) .set(ValueClass::Property(1), value.as_slice()) - .set(ValueClass::Property(0), "check1") - .set(ValueClass::Property(2), "check2") + .set(ValueClass::Property(0), "check1".as_bytes()) + .set(ValueClass::Property(2), "check2".as_bytes()) .build_batch(), ) .await diff --git a/tests/src/store/query.rs b/tests/src/store/query.rs index 9b3b024a..c65f98df 100644 --- a/tests/src/store/query.rs +++ b/tests/src/store/query.rs @@ -23,6 +23,7 @@ use std::{ fmt::Display, + io::Write, sync::{Arc, Mutex}, time::Instant, }; @@ -155,7 +156,7 @@ pub async fn test(db: Store, fts_store: FtsStore, do_insert: bool) { builder .with_account_id(0) .with_collection(COLLECTION_ID) - .create_document(document_id as u32); + .create_document_with_id(document_id as u32); for (pos, field) in record.iter().enumerate() { let field_id = pos as u8; match FIELDS_OPTIONS[pos] { @@ -218,6 +219,7 @@ pub async fn test(db: Store, fts_store: FtsStore, do_insert: bool) { let mut chunk = Vec::new(); let mut fts_chunk = Vec::new(); + print!("Inserting... ",); for (batch, fts_batch) in batches { let chunk_instance = Instant::now(); chunk.push({ @@ -235,10 +237,8 @@ pub async fn test(db: Store, fts_store: FtsStore, do_insert: bool) { for handle in fts_chunk { handle.await.unwrap().unwrap(); } - println!( - "Store insert took {} ms.", - chunk_instance.elapsed().as_millis() - ); + print!(" [{} ms]", chunk_instance.elapsed().as_millis()); + std::io::stdout().flush().unwrap(); chunk = Vec::new(); fts_chunk = Vec::new(); } @@ -250,14 +250,18 @@ pub async fn test(db: Store, fts_store: FtsStore, do_insert: bool) { } } - println!("Insert took {} ms.", now.elapsed().as_millis()); + println!("\nInsert took {} ms.", now.elapsed().as_millis()); } println!("Running filter tests..."); + let now = Instant::now(); test_filter(db.clone(), fts_store).await; + println!("Filtering took {} ms.", now.elapsed().as_millis()); println!("Running sort tests..."); + let now = Instant::now(); test_sort(db).await; + println!("Sorting took {} ms.", now.elapsed().as_millis()); } pub async fn test_filter(db: Store, fts: FtsStore) {