/* * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL */ use super::{ Batch, BatchBuilder, ChangedCollection, IntoOperations, Operation, ValueClass, ValueOp, assert::ToAssertValue, log::VanishedItem, }; use crate::{ SerializeInfallible, U32_LEN, write::{LogCollection, MergeFnc, MergeOperation, Params, SetFnc, SetOperation}, }; use types::{ collection::{Collection, SyncCollection, VanishedCollection}, field::FieldType, }; use utils::map::vec_map::VecMap; impl BatchBuilder { pub fn new() -> Self { Self { ops: Vec::with_capacity(32), current_account_id: None, current_collection: None, current_document_id: None, changes: Default::default(), changed_collections: Default::default(), batch_size: 0, batch_ops: 0, has_assertions: false, commit_points: Vec::new(), } } pub fn with_account_id(&mut self, account_id: u32) -> &mut Self { if self .current_account_id .is_none_or(|current_account_id| current_account_id != account_id) { self.current_account_id = account_id.into(); self.ops.push(Operation::AccountId { account_id }); } self } pub fn with_collection(&mut self, collection: Collection) -> &mut Self { let collection_ = Some(collection); if collection_ != self.current_collection { self.current_collection = collection_; self.ops.push(Operation::Collection { collection }); } self } pub fn with_document(&mut self, document_id: u32) -> &mut Self { self.ops.push(Operation::DocumentId { document_id }); self.current_document_id = Some(document_id); self.has_assertions = false; self } pub fn assert_value( &mut self, class: impl Into, value: impl ToAssertValue, ) -> &mut Self { self.ops.push(Operation::AssertValue { class: class.into(), assert_value: value.to_assert_value(), }); self.batch_ops += 1; self.has_assertions = true; self } pub fn index(&mut self, field: impl FieldType, value: impl Into>) -> &mut Self { let field = field.into(); let value = value.into(); let value_len = value.len(); self.ops.push(Operation::Index { field, key: value, set: true, }); self.batch_size += (U32_LEN * 3) + value_len; self.batch_ops += 1; self } pub fn unindex(&mut self, field: impl FieldType, value: impl Into>) -> &mut Self { let field = field.into(); let value = value.into(); let value_len = value.len(); self.ops.push(Operation::Index { field, key: value, set: false, }); self.batch_size += (U32_LEN * 3) + value_len; self.batch_ops += 1; self } #[inline(always)] pub fn tag(&mut self, field: impl FieldType) -> &mut Self { self.index(field, vec![]) } #[inline(always)] pub fn untag(&mut self, field: impl FieldType) -> &mut Self { self.unindex(field, vec![]) } pub fn add(&mut self, class: impl Into, value: i64) -> &mut Self { let class = class.into(); self.batch_size += class.serialized_size() + std::mem::size_of::(); self.ops.push(Operation::Value { class, op: ValueOp::AtomicAdd(value), }); self.batch_ops += 1; self } pub fn add_and_get(&mut self, class: impl Into, value: i64) -> &mut Self { let class = class.into(); self.batch_size += class.serialized_size() + (std::mem::size_of::() * 2); self.ops.push(Operation::Value { class, op: ValueOp::AddAndGet(value), }); self.batch_ops += 1; self } pub fn set(&mut self, class: impl Into, value: impl Into>) -> &mut Self { let class = class.into(); let value = value.into(); self.batch_size += class.serialized_size() + value.len(); self.ops.push(Operation::Value { class, op: ValueOp::Set(value), }); self.batch_ops += 1; self } pub fn set_fnc( &mut self, class: impl Into, params: Params, fnc: SetFnc, ) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::SetFnc(SetOperation { fnc, params }), }); self } pub fn merge_fnc( &mut self, class: impl Into, params: Params, fnc: MergeFnc, ) -> &mut Self { self.ops.push(Operation::Value { class: class.into(), op: ValueOp::MergeFnc(MergeOperation { fnc, params }), }); self } pub fn clear(&mut self, class: impl Into) -> &mut Self { let class = class.into(); self.batch_size += class.serialized_size(); self.ops.push(Operation::Value { class, op: ValueOp::Clear, }); self.batch_ops += 1; self } pub fn acl_grant(&mut self, grant_account_id: u32, op: Vec) -> &mut Self { self.batch_size += (U32_LEN * 3) + op.len(); self.ops.push(Operation::Value { class: ValueClass::Acl(grant_account_id), op: ValueOp::Set(op), }); self.batch_ops += 1; self } pub fn acl_revoke(&mut self, grant_account_id: u32) -> &mut Self { self.batch_size += U32_LEN * 3; self.ops.push(Operation::Value { class: ValueClass::Acl(grant_account_id), op: ValueOp::Clear, }); self.batch_ops += 1; self } pub fn log_item_insert( &mut self, collection: SyncCollection, prefix: Option, ) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes.get_mut_or_insert(account_id).log_item_insert( collection, prefix, document_id, ); } self } pub fn log_item_update( &mut self, collection: SyncCollection, prefix: Option, ) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes.get_mut_or_insert(account_id).log_item_update( collection, prefix, document_id, ); } self } pub fn log_item_delete( &mut self, collection: SyncCollection, prefix: Option, ) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes.get_mut_or_insert(account_id).log_item_delete( collection, prefix, document_id, ); } self } pub fn log_container_insert(&mut self, collection: SyncCollection) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes .get_mut_or_insert(account_id) .log_container_insert(collection, document_id); } self } pub fn log_container_update(&mut self, collection: SyncCollection) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes .get_mut_or_insert(account_id) .log_container_update(collection, document_id); } self } pub fn log_container_delete(&mut self, collection: SyncCollection) -> &mut Self { if let (Some(account_id), Some(document_id)) = (self.current_account_id, self.current_document_id) { self.changes .get_mut_or_insert(account_id) .log_container_delete(collection, document_id); } self } pub fn log_container_property_change( &mut self, collection: SyncCollection, document_id: u32, ) -> &mut Self { if let Some(account_id) = self.current_account_id { self.changes .get_mut_or_insert(account_id) .log_container_property_update(collection, document_id); } self } pub fn log_vanished_item( &mut self, collection: VanishedCollection, item: impl Into, ) -> &mut Self { if let Some(account_id) = self.current_account_id { let item = item.into(); self.batch_size += item.serialized_size(); self.changes .get_mut_or_insert(account_id) .log_vanished_item(collection, item); } self } pub fn log_share_notification( &mut self, notification_id: u64, notify_account_id: u32, value: impl SerializeInfallible, ) -> &mut Self { self.changed_collections .get_mut_or_insert(notify_account_id) .share_notification_id = Some(notification_id); self.set( ValueClass::ShareNotification { notification_id, notify_account_id, }, value.serialize(), ) } fn serialize_changes(&mut self) { if !self.changes.is_empty() { for (account_id, changelog) in std::mem::take(&mut self.changes) { self.with_account_id(account_id); // Serialize changes for (collection, changes) in changelog.changes.into_iter() { let cc = self.changed_collections.get_mut_or_insert(account_id); if changes.has_container_changes() { cc.changed_containers.insert(collection); } if changes.has_item_changes() { cc.changed_items.insert(collection); } self.ops.push(Operation::Log { collection: LogCollection::Sync(collection), set: changes.serialize(), }); } // Serialize vanished items for (collection, vanished) in changelog.vanished.into_iter() { self.ops.push(Operation::Log { collection: LogCollection::Vanished(collection), set: vanished.serialize(), }); } } } } pub fn commit_point(&mut self) -> &mut Self { if self.is_large_batch() { self.serialize_changes(); self.commit_points.push(self.ops.len()); self.batch_ops = 0; self.batch_size = 0; if let Some(account_id) = self.current_account_id { self.ops.push(Operation::AccountId { account_id }); } if let Some(collection) = self.current_collection { self.ops.push(Operation::Collection { collection }); } } self } #[inline] pub fn is_large_batch(&self) -> bool { self.batch_size > 5_000_000 || self.batch_ops > 1000 } pub fn any_op(&mut self, op: Operation) -> &mut Self { self.ops.push(op); self.batch_ops += 1; self } pub fn custom(&mut self, value: impl IntoOperations) -> trc::Result<&mut Self> { value.build(self)?; Ok(self) } pub fn last_account_id(&self) -> Option { self.current_account_id } pub fn last_collection(&self) -> Option { self.current_collection } pub fn last_document_id(&self) -> Option { self.current_document_id } pub fn commit_points(&mut self) -> CommitPointIterator { self.serialize_changes(); CommitPointIterator { commit_points: std::mem::take(&mut self.commit_points), commit_point_last: self.ops.len(), offset_start: 0, } } pub fn build_one(&mut self, commit_point: CommitPoint) -> Batch<'_> { Batch { changes: &self.changed_collections, ops: &mut self.ops[commit_point.offset_start..commit_point.offset_end], } } pub fn build_all(&mut self) -> Batch<'_> { self.serialize_changes(); Batch { changes: &self.changed_collections, ops: self.ops.as_mut_slice(), } } pub fn changes(self) -> Option> { if self.has_changes() { Some(self.changed_collections) } else { None } } pub fn has_changes(&self) -> bool { !self.changed_collections.is_empty() } pub fn ops(&self) -> &[Operation] { self.ops.as_slice() } pub fn len(&self) -> usize { self.batch_size } pub fn is_empty(&self) -> bool { self.batch_ops == 0 } } pub struct CommitPointIterator { commit_points: Vec, commit_point_last: usize, offset_start: usize, } pub struct CommitPoint { pub offset_start: usize, pub offset_end: usize, } impl CommitPointIterator { pub fn iter(&mut self) -> impl Iterator { self.commit_points .iter() .copied() .chain([self.commit_point_last]) .map(|offset_end| { let point = CommitPoint { offset_start: self.offset_start, offset_end, }; self.offset_start = offset_end; point }) } } impl Batch<'_> { pub fn is_atomic(&self) -> bool { !self.ops.iter().any(|op| { matches!( op, Operation::AssertValue { .. } | Operation::Value { op: ValueOp::AddAndGet(_), .. } ) }) } pub fn first_account_id(&self) -> Option { self.ops.iter().find_map(|op| match op { Operation::AccountId { account_id } => Some(*account_id), _ => None, }) } } impl Default for BatchBuilder { fn default() -> Self { Self::new() } }