diff --git a/slatedb/src/batch.rs b/slatedb/src/batch.rs index c54586c66..8dca5a976 100644 --- a/slatedb/src/batch.rs +++ b/slatedb/src/batch.rs @@ -53,6 +53,18 @@ pub struct WriteBatch { pub(crate) ops: BTreeMap>, pub(crate) op_count: usize, pub(crate) has_merge_ops: bool, + /// Keys whose surviving operations are exclusively commutative merges. + /// + /// This is transaction conflict metadata only. It is never encoded into a + /// [`RowEntry`] or persisted by the WAL, memtable, or compaction paths. The + /// set is boxed and allocated lazily so ordinary write batches retain their + /// existing inline size and pay no allocation for this opt-in metadata. + commutative_merge_keys: Option>, +} + +#[derive(Clone, Debug, Default)] +struct CommutativeMergeKeys { + keys: HashSet, } impl Default for WriteBatch { @@ -146,6 +158,7 @@ impl WriteBatch { ops: BTreeMap::new(), op_count: 0, has_merge_ops: false, + commutative_merge_keys: None, } } @@ -224,6 +237,7 @@ impl WriteBatch { /// - if the value size is larger than u32::MAX pub fn put_bytes_with_options(&mut self, key: Bytes, value: Bytes, options: &PutOptions) { self.assert_kv(&key, &value); + self.remove_commutative_merge_key(&key); // put will overwrite the existing key so we can safely // remove all previous entries. @@ -250,16 +264,57 @@ impl WriteBatch { where K: AsRef<[u8]>, V: AsRef<[u8]>, + { + self.merge_with_conflict_kind(key, value, options, false); + } + + /// Merge a key-value pair whose operands are commutative. + /// + /// This remains crate-private because compatibility affects transaction + /// conflict detection, while public non-transactional batches do not run + /// write/write conflict checks. + pub(crate) fn merge_commutative(&mut self, key: K, value: V) + where + K: AsRef<[u8]>, + V: AsRef<[u8]>, + { + self.merge_with_conflict_kind(key, value, &MergeOptions::default(), true); + } + + fn merge_with_conflict_kind( + &mut self, + key: K, + value: V, + options: &MergeOptions, + commutative: bool, + ) where + K: AsRef<[u8]>, + V: AsRef<[u8]>, { self.assert_kv(&key, &value); - let key = key.as_ref(); + let key = Bytes::copy_from_slice(key.as_ref()); let value = value.as_ref(); let op = WriteOp::Merge(Bytes::copy_from_slice(value), options.clone()); - if let Some(ops) = self.ops.get_mut(key) { + let existing_operations_are_commutative = self.ops.contains_key(&key) + && self + .commutative_merge_keys + .as_ref() + .is_some_and(|keys| keys.keys.contains(&key)); + let is_first_operation = !self.ops.contains_key(&key); + if commutative && (is_first_operation || existing_operations_are_commutative) { + self.commutative_merge_keys + .get_or_insert_default() + .keys + .insert(key.clone()); + } else { + self.remove_commutative_merge_key(&key); + } + + if let Some(ops) = self.ops.get_mut(&key) { ops.push(op); } else { - self.ops.insert(Bytes::copy_from_slice(key), smallvec![op]); + self.ops.insert(key, smallvec![op]); } self.has_merge_ops = true; @@ -271,6 +326,7 @@ impl WriteBatch { self.assert_kv(&key, &[]); let key = Bytes::copy_from_slice(key.as_ref()); + self.remove_commutative_merge_key(&key); // delete will overwrite the existing key so we can safely // remove all previous entries. @@ -306,6 +362,23 @@ impl WriteBatch { self.ops.keys().cloned().collect() } + /// Returns whether every surviving operation for `key` is a commutative merge. + pub(crate) fn is_commutative_merge_key(&self, key: &[u8]) -> bool { + self.commutative_merge_keys + .as_ref() + .is_some_and(|keys| keys.keys.contains(key)) + } + + fn remove_commutative_merge_key(&mut self, key: &[u8]) { + let Some(keys) = self.commutative_merge_keys.as_mut() else { + return; + }; + keys.keys.remove(key); + if keys.keys.is_empty() { + self.commutative_merge_keys = None; + } + } + /// Converts a WriteBatch into a vector of RowEntry objects with /// seq and timestamp set, applying the merge operator to any /// mergeable entries. @@ -864,6 +937,49 @@ mod tests { } } + #[test] + fn commutative_merge_classification_requires_only_commutative_operations() { + let mut only_commutative = WriteBatch::new(); + only_commutative.merge_commutative(b"key", b"first"); + only_commutative.merge_commutative(b"key", b"second"); + assert!(only_commutative.is_commutative_merge_key(b"key")); + assert!(only_commutative.commutative_merge_keys.is_some()); + + let mut commutative_then_ordinary = WriteBatch::new(); + commutative_then_ordinary.merge_commutative(b"key", b"first"); + commutative_then_ordinary.merge(b"key", b"second"); + assert!(!commutative_then_ordinary.is_commutative_merge_key(b"key")); + assert!(commutative_then_ordinary.commutative_merge_keys.is_none()); + + let mut ordinary_then_commutative = WriteBatch::new(); + ordinary_then_commutative.merge(b"key", b"first"); + ordinary_then_commutative.merge_commutative(b"key", b"second"); + assert!(!ordinary_then_commutative.is_commutative_merge_key(b"key")); + } + + #[test] + fn put_and_delete_make_commutative_merge_keys_exclusive() { + let mut put_after_merge = WriteBatch::new(); + put_after_merge.merge_commutative(b"key", b"merge"); + put_after_merge.put(b"key", b"put"); + assert!(!put_after_merge.is_commutative_merge_key(b"key")); + + let mut merge_after_put = WriteBatch::new(); + merge_after_put.put(b"key", b"put"); + merge_after_put.merge_commutative(b"key", b"merge"); + assert!(!merge_after_put.is_commutative_merge_key(b"key")); + + let mut delete_after_merge = WriteBatch::new(); + delete_after_merge.merge_commutative(b"key", b"merge"); + delete_after_merge.delete(b"key"); + assert!(!delete_after_merge.is_commutative_merge_key(b"key")); + + let mut merge_after_delete = WriteBatch::new(); + merge_after_delete.delete(b"key"); + merge_after_delete.merge_commutative(b"key", b"merge"); + assert!(!merge_after_delete.is_commutative_merge_key(b"key")); + } + #[test] fn should_create_merge_operation_with_custom_options() { // Given: an empty WriteBatch and custom merge options diff --git a/slatedb/src/db_transaction.rs b/slatedb/src/db_transaction.rs index 305cd003e..580e06970 100644 --- a/slatedb/src/db_transaction.rs +++ b/slatedb/src/db_transaction.rs @@ -1,6 +1,6 @@ use bytes::Bytes; use parking_lot::{Mutex, RwLock}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; @@ -13,7 +13,7 @@ use crate::db_iter::{DbIterator, DbIteratorRangeTracker}; use crate::error::SlateDBError; use crate::iter::IterationOrder; use crate::reader::ScanContext; -use crate::transaction_manager::{IsolationLevel, TransactionManager}; +use crate::transaction_manager::{IsolationLevel, TransactionManager, TransactionWriteKind}; use crate::types::KeyValue; use crate::{DbReadOps, DbTransactionOps}; @@ -634,6 +634,165 @@ impl DbTransaction { self.merge_with_options(key, value, &MergeOptions::default()) } + /// Buffers a merge operand whose effect is independent of its ordering + /// relative to other commutative merge operands for the same key. + /// + /// This is an opt-in write/write conflict optimization for transactions + /// performing blind, merge-only updates. Unlike [`Self::merge`], concurrent + /// transactions using this method for the same key may both commit. Both + /// operands are retained and applied by the configured [`crate::MergeOperator`]. + /// + /// # Why this exists + /// + /// Ordinary merge operations participate in the same write/write conflict + /// checks as puts and deletes. That is conservative for merge operators such + /// as counters and set unions, where retaining and applying both operands is + /// equivalent to serially executing the transactions in either order. + /// + /// This method allows callers to state that stronger algebraic property + /// explicitly, avoiding unnecessary transaction retries on hot keys. + /// + /// # Required algebraic contract + /// + /// Let `apply(base, operand)` represent the configured merge operator. For + /// every base value and every pair of operands that may concurrently target + /// the key, the caller must ensure: + /// + /// ```text + /// apply(apply(base, a), b) == apply(apply(base, b), a) + /// ``` + /// + /// The existing associativity requirements of [`crate::MergeOperator`] and + /// [`crate::MergeOperator::merge_batch`] also continue to apply. SlateDB + /// cannot inspect an application-defined merge operator or prove this + /// property. + /// + /// Idempotence is not required: additive counter deltas are valid even + /// though applying the same delta twice changes the result. Set insertion + /// and bitmap union are both commutative and idempotent. + /// + /// Do not use this method for order-sensitive appends, replacements, + /// compare-and-set operations, removals that race with additions, or any + /// operand whose meaning depends on which transaction commits first. + /// + /// # Conflict and isolation guarantees + /// + /// This remains a tracked write. It does not have the broad semantics of + /// [`Self::unmark_write`]. + /// + /// - A write/write conflict is omitted only when both transactions classify + /// their final operations for the key exclusively as commutative merges. + /// - A put, delete, ordinary merge, or mixture of operation kinds makes the + /// key exclusive and restores normal write/write conflict detection. + /// - Serializable point-read and range-read dependencies continue to see + /// this key as written and can still cause the transaction to abort. + /// - Reading the key before calling this method does not suppress that read + /// dependency. + /// - Conflicts involving other keys in the transaction are unaffected. + /// - Calling [`Self::unmark_write`] separately still removes the key from + /// conflict tracking according to that method's broader contract. + /// + /// Consequently, this method does not guarantee that a transaction will + /// commit; it only makes compatible same-key merges cease being a + /// write/write conflict. + /// + /// # Atomicity and persistence guarantees + /// + /// The operand remains part of the transaction's ordinary atomic write + /// batch. Commit ordering, durability, snapshot visibility, merge-operator + /// execution, and error propagation are unchanged. The commutative marker + /// is transaction conflict metadata only and is not stored on disk. + /// + /// This is a safe Rust API and cannot cause undefined behavior. Violating + /// the algebraic contract can, however, admit executions that violate the + /// caller's application-level ordering or consistency invariants. + /// + /// Custom [`MergeOptions`] are intentionally unsupported because TTL or + /// other order-sensitive options would require a separate compatibility + /// contract. + /// + /// # Errors + /// + /// Returns [`crate::Error`] when the database has no configured + /// [`crate::MergeOperator`]. + /// + /// # Panics + /// + /// Panics under the same input constraints as other write operations: + /// + /// - the key is empty; + /// - the key exceeds `u16::MAX` bytes; + /// - the operand exceeds `u32::MAX` bytes. + /// + /// # Example + /// + /// ``` + /// use std::sync::Arc; + /// + /// use bytes::Bytes; + /// use slatedb::object_store::{memory::InMemory, ObjectStore}; + /// use slatedb::{ + /// Db, Error, IsolationLevel, MergeOperator, MergeOperatorError, + /// }; + /// + /// struct Counter; + /// + /// impl MergeOperator for Counter { + /// fn merge( + /// &self, + /// _key: &Bytes, + /// existing: Option, + /// operand: Bytes, + /// ) -> Result { + /// let current = existing + /// .map(|value| { + /// u64::from_le_bytes(value.as_ref().try_into().unwrap()) + /// }) + /// .unwrap_or(0); + /// let delta = + /// u64::from_le_bytes(operand.as_ref().try_into().unwrap()); + /// Ok(Bytes::copy_from_slice(&(current + delta).to_le_bytes())) + /// } + /// } + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Error> { + /// let store: Arc = Arc::new(InMemory::new()); + /// let db = Db::builder("commutative-merge-example", store) + /// .with_merge_operator(Arc::new(Counter)) + /// .build() + /// .await?; + /// + /// let first = db.begin(IsolationLevel::SerializableSnapshot).await?; + /// let second = db.begin(IsolationLevel::SerializableSnapshot).await?; + /// + /// first.merge_commutative(b"counter", 1_u64.to_le_bytes())?; + /// second.merge_commutative(b"counter", 1_u64.to_le_bytes())?; + /// + /// first.commit().await?; + /// second.commit().await?; + /// + /// let value = db.get(b"counter").await?.unwrap(); + /// assert_eq!( + /// u64::from_le_bytes(value.as_ref().try_into().unwrap()), + /// 2, + /// ); + /// Ok(()) + /// } + /// ``` + pub fn merge_commutative(&self, key: K, operand: V) -> Result<(), crate::Error> + where + K: AsRef<[u8]>, + V: AsRef<[u8]>, + { + if self.db_inner.flush_merge_operator.is_none() { + return Err(SlateDBError::MergeOperatorMissing.into()); + } + + self.write_batch.write().merge_commutative(key, operand); + Ok(()) + } + /// Merge a key-value pair into the transaction with custom options. /// /// ## Errors @@ -738,17 +897,25 @@ impl DbTransaction { return Ok(None); } - // Track only write keys that were not explicitly unmarked. - let tracked_write_keys = { + // Track only writes that were not explicitly unmarked. Compatibility + // is derived from the final surviving operations for each key. + let tracked_writes: HashMap = { let untracked_write_keys = self.untracked_write_keys.read(); write_batch .keys() .into_iter() .filter(|key| !untracked_write_keys.contains(key)) + .map(|key| { + let kind = if write_batch.is_commutative_merge_key(&key) { + TransactionWriteKind::CommutativeMerge + } else { + TransactionWriteKind::Exclusive + }; + (key, kind) + }) .collect() }; - self.txn_manager - .track_write_keys(&self.txn_id, &tracked_write_keys); + self.txn_manager.track_writes(&self.txn_id, &tracked_writes); // Submit the WriteBatch to the database for processing. The batch is sent to a // dedicated background task (in batch_write.rs) that processes all WriteBatches @@ -877,6 +1044,14 @@ impl DbTransactionOps for DbTransaction { DbTransaction::merge_with_options(self, key, value, options) } + fn merge_commutative(&self, key: K, operand: V) -> Result<(), crate::Error> + where + K: AsRef<[u8]>, + V: AsRef<[u8]>, + { + DbTransaction::merge_commutative(self, key, operand) + } + fn mark_read(&self, keys: I) -> Result<(), crate::Error> where K: AsRef<[u8]>, @@ -2322,6 +2497,146 @@ mod tests { assert_eq!(total, EXPECTED); } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_commutative_merge_counter_aggregates_under_high_concurrency() { + const CONCURRENT_TXNS: usize = 32; + const ROUNDS: usize = 8; + const MERGE_INCREMENT: [u8; 8] = 1u64.to_le_bytes(); + const EXPECTED: u64 = (CONCURRENT_TXNS * ROUNDS) as u64; + + let object_store: Arc = Arc::new(InMemory::new()); + let db = crate::Db::builder("test_commutative_merge_counter", object_store) + .with_merge_operator(Arc::new(CounterMergeOperator)) + .build() + .await + .unwrap(); + + for _ in 0..ROUNDS { + let barrier = Arc::new(tokio::sync::Barrier::new(CONCURRENT_TXNS)); + let mut handles = Vec::with_capacity(CONCURRENT_TXNS); + + for _ in 0..CONCURRENT_TXNS { + let db = db.clone(); + let barrier = barrier.clone(); + handles.push(tokio::spawn(async move { + let txn = db + .begin(IsolationLevel::SerializableSnapshot) + .await + .unwrap(); + txn.merge_commutative(b"counter", MERGE_INCREMENT).unwrap(); + barrier.wait().await; + txn.commit().await.unwrap(); + })); + } + + for handle in handles { + handle.await.unwrap(); + } + } + + let value = db.get(b"counter").await.unwrap().unwrap(); + let total = u64::from_le_bytes(value.as_ref().try_into().unwrap()); + assert_eq!(total, EXPECTED); + } + + #[tokio::test] + async fn test_commutative_merge_conflicts_with_ordinary_merge_put_and_delete() { + const INITIAL: [u8; 8] = 0u64.to_le_bytes(); + const INCREMENT: [u8; 8] = 1u64.to_le_bytes(); + + let object_store: Arc = Arc::new(InMemory::new()); + let db = crate::Db::builder("test_commutative_merge_exclusive_conflicts", object_store) + .with_merge_operator(Arc::new(CounterMergeOperator)) + .build() + .await + .unwrap(); + + for key in [b"merge".as_slice(), b"put".as_slice(), b"delete".as_slice()] { + db.put(key, INITIAL).await.unwrap(); + let commutative = db.begin(IsolationLevel::Snapshot).await.unwrap(); + let exclusive = db.begin(IsolationLevel::Snapshot).await.unwrap(); + commutative.merge_commutative(key, INCREMENT).unwrap(); + + match key { + b"merge" => exclusive.merge(key, INCREMENT).unwrap(), + b"put" => exclusive.put(key, INITIAL).unwrap(), + b"delete" => exclusive.delete(key).unwrap(), + _ => unreachable!(), + } + + commutative.commit().await.unwrap(); + assert!(exclusive.commit().await.is_err()); + } + } + + #[tokio::test] + async fn test_mixed_same_key_operations_restore_exclusive_conflicts() { + const INITIAL: [u8; 8] = 0u64.to_le_bytes(); + const INCREMENT: [u8; 8] = 1u64.to_le_bytes(); + + let object_store: Arc = Arc::new(InMemory::new()); + let db = crate::Db::builder("test_commutative_merge_mixed_operations", object_store) + .with_merge_operator(Arc::new(CounterMergeOperator)) + .build() + .await + .unwrap(); + db.put(b"counter", INITIAL).await.unwrap(); + + let mixed = db.begin(IsolationLevel::Snapshot).await.unwrap(); + mixed.merge_commutative(b"counter", INCREMENT).unwrap(); + mixed.put(b"counter", INITIAL).unwrap(); + + let commutative = db.begin(IsolationLevel::Snapshot).await.unwrap(); + commutative + .merge_commutative(b"counter", INCREMENT) + .unwrap(); + commutative.commit().await.unwrap(); + + assert!(mixed.commit().await.is_err()); + } + + #[tokio::test] + async fn test_commutative_merge_remains_visible_to_ssi_point_and_range_reads() { + const INCREMENT: [u8; 8] = 1u64.to_le_bytes(); + + let object_store: Arc = Arc::new(InMemory::new()); + let db = crate::Db::builder("test_commutative_merge_ssi_reads", object_store) + .with_merge_operator(Arc::new(CounterMergeOperator)) + .build() + .await + .unwrap(); + + let point_reader = db + .begin(IsolationLevel::SerializableSnapshot) + .await + .unwrap(); + assert_eq!(point_reader.get(b"point").await.unwrap(), None); + let point_writer = db.begin(IsolationLevel::Snapshot).await.unwrap(); + point_writer.merge_commutative(b"point", INCREMENT).unwrap(); + point_writer.commit().await.unwrap(); + point_reader.put(b"point-reader-write", b"value").unwrap(); + assert!(point_reader.commit().await.is_err()); + + let range_reader = db + .begin(IsolationLevel::SerializableSnapshot) + .await + .unwrap(); + let mut range = range_reader + .scan(&b"range-a"[..]..=&b"range-z"[..]) + .await + .unwrap(); + while range.next().await.unwrap().is_some() {} + drop(range); + + let range_writer = db.begin(IsolationLevel::Snapshot).await.unwrap(); + range_writer + .merge_commutative(b"range-m", INCREMENT) + .unwrap(); + range_writer.commit().await.unwrap(); + range_reader.put(b"range-reader-write", b"value").unwrap(); + assert!(range_reader.commit().await.is_err()); + } + #[tokio::test] async fn test_txn_merge_requires_merge_operator() { let object_store: Arc = Arc::new(InMemory::new()); @@ -2339,6 +2654,20 @@ mod tests { assert_eq!(db.get(b"counter").await.unwrap(), None); } + #[tokio::test] + async fn test_txn_commutative_merge_requires_merge_operator() { + let object_store: Arc = Arc::new(InMemory::new()); + let db = crate::Db::open("test_txn_commutative_merge_requires_operator", object_store) + .await + .unwrap(); + + let txn = db.begin(IsolationLevel::Snapshot).await.unwrap(); + let err = txn + .merge_commutative(b"counter", 1u64.to_le_bytes()) + .unwrap_err(); + assert_eq!(err.kind(), crate::ErrorKind::Invalid); + } + #[tokio::test] async fn test_txn_commit_rejects_same_key_merge_different_ttls() { let object_store: Arc = Arc::new(InMemory::new()); diff --git a/slatedb/src/ops.rs b/slatedb/src/ops.rs index 2a7c83031..f01cd115e 100644 --- a/slatedb/src/ops.rs +++ b/slatedb/src/ops.rs @@ -490,6 +490,25 @@ pub trait DbTransactionOps: DbReadOps { self.merge_with_options(key, value, &MergeOptions::default()) } + /// Buffers a merge operand that is compatible with concurrent + /// commutative merges for the same key. + /// + /// See [`DbTransaction::merge_commutative`](crate::DbTransaction::merge_commutative) + /// for the algebraic caller contract and the precise conflict, isolation, + /// atomicity, and persistence guarantees. + /// + /// The default implementation conservatively delegates to [`Self::merge`] + /// so third-party transaction implementations remain source-compatible and + /// retain ordinary write/write conflicts unless they explicitly override + /// this method. + fn merge_commutative(&self, key: K, operand: V) -> Result<(), crate::Error> + where + K: AsRef<[u8]>, + V: AsRef<[u8]>, + { + self.merge(key, operand) + } + /// Merge a key-value pair into the transaction with custom `MergeOptions`. /// /// ## Errors diff --git a/slatedb/src/transaction_manager.rs b/slatedb/src/transaction_manager.rs index a9cba127e..b50aa81c9 100644 --- a/slatedb/src/transaction_manager.rs +++ b/slatedb/src/transaction_manager.rs @@ -21,6 +21,16 @@ pub enum IsolationLevel { SerializableSnapshot, } +/// Conflict behavior for a transaction's final operations on one key. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum TransactionWriteKind { + /// The key conflicts with every concurrent write to the same key. + Exclusive, + /// The key is merge-only and may coexist with another merge-only, + /// commutative write to the same key. + CommutativeMerge, +} + #[derive(Debug)] pub(crate) struct TransactionState { /// The sequence number when the transaction started. This is used to establish @@ -30,9 +40,9 @@ pub(crate) struct TransactionState { /// a transaction is committed and is used to check conflicts with recent committed /// transactions. committed_seq: Option, - /// The write keys of the transaction for write-write conflict detection. + /// The transaction's final per-key write behavior for conflict detection. /// Used in both Snapshot Isolation and Serializable Snapshot Isolation. - write_keys: HashSet, + writes: HashMap, /// The read keys of the transaction for read-write conflict detection. /// Only used in Serializable Snapshot Isolation mode. read_keys: HashSet, @@ -42,9 +52,29 @@ pub(crate) struct TransactionState { } impl TransactionState { - /// Add write keys to this transaction's write set for conflict detection. + /// Add writes to this transaction's write set for conflict detection. + /// + /// Exclusive behavior dominates when a key is tracked more than once so a + /// mixed operation sequence cannot accidentally become compatible. + fn track_writes(&mut self, writes: impl IntoIterator) { + for (key, kind) in writes { + self.writes + .entry(key) + .and_modify(|existing| { + if *existing != kind { + *existing = TransactionWriteKind::Exclusive; + } + }) + .or_insert(kind); + } + } + + #[cfg(test)] fn track_write_keys(&mut self, keys: impl IntoIterator) { - self.write_keys.extend(keys); + self.track_writes( + keys.into_iter() + .map(|key| (key, TransactionWriteKind::Exclusive)), + ); } /// Add read keys to this transaction's read set for SSI conflict detection. @@ -120,7 +150,7 @@ impl TransactionManager { TransactionState { started_seq: seq, committed_seq: None, - write_keys: HashSet::new(), + writes: HashMap::new(), read_keys: HashSet::new(), read_ranges: Vec::new(), }, @@ -134,7 +164,7 @@ impl TransactionManager { let txn_state = TransactionState { started_seq: seq, committed_seq: None, - write_keys: HashSet::new(), + writes: HashMap::new(), read_keys: HashSet::new(), read_ranges: Vec::new(), }; @@ -158,15 +188,29 @@ impl TransactionManager { inner.recycle_recent_committed_txns(); } - /// Track write keys for a transaction. This is used for conflict detection. - /// Keys should be tracked before calling commit-related methods. - pub(crate) fn track_write_keys(&self, txn_id: &Uuid, write_keys: &HashSet) { + /// Track writes for a transaction. This is used for conflict detection. + /// Writes should be tracked before calling commit-related methods. + pub(crate) fn track_writes( + &self, + txn_id: &Uuid, + writes: &HashMap, + ) { let mut inner = self.inner.write(); if let Some(txn_state) = inner.active_txns.get_mut(txn_id) { - txn_state.track_write_keys(write_keys.iter().cloned()); + txn_state.track_writes(writes.iter().map(|(key, kind)| (key.clone(), *kind))); } } + #[cfg(test)] + fn track_write_keys(&self, txn_id: &Uuid, write_keys: &HashSet) { + let writes = write_keys + .iter() + .cloned() + .map(|key| (key, TransactionWriteKind::Exclusive)) + .collect(); + self.track_writes(txn_id, &writes); + } + /// Track a key read operation (for SSI) pub(crate) fn track_read_keys( &self, @@ -200,8 +244,7 @@ impl TransactionManager { }; // both SI and SSI need to check write-write conflicts - let ww_conflict = - inner.has_write_write_conflict(&txn_state.write_keys, txn_state.started_seq); + let ww_conflict = inner.has_write_write_conflict(&txn_state.writes, txn_state.started_seq); if ww_conflict { return true; } @@ -249,7 +292,11 @@ impl TransactionManager { inner.track_recent_committed_state(TransactionState { started_seq: committed_seq, committed_seq: Some(committed_seq), - write_keys: keys.clone(), + writes: keys + .iter() + .cloned() + .map(|key| (key, TransactionWriteKind::Exclusive)) + .collect(), read_keys: HashSet::new(), read_ranges: Vec::new(), }); @@ -319,9 +366,13 @@ impl TransactionManagerInner { } } - fn has_write_write_conflict(&self, write_keys: &HashSet, started_seq: u64) -> bool { + fn has_write_write_conflict( + &self, + writes: &HashMap, + started_seq: u64, + ) -> bool { // If the current transaction has no write operations, there's no write-write conflict - if write_keys.is_empty() { + if writes.is_empty() { return false; } @@ -331,12 +382,21 @@ impl TransactionManagerInner { "all txns in recent_committed_txns should be committed with committed_seq set", ); - // if another transaction committed after the current transaction started, - // and they have overlapping write keys, then there's a conflict. - if other_committed_seq > started_seq - && !write_keys.is_disjoint(&committed_txn.write_keys) - { - return true; + if other_committed_seq > started_seq { + for (key, kind) in writes { + let Some(other_kind) = committed_txn.writes.get(key) else { + continue; + }; + if !matches!( + (kind, other_kind), + ( + TransactionWriteKind::CommutativeMerge, + TransactionWriteKind::CommutativeMerge + ) + ) { + return true; + } + } } } @@ -368,7 +428,10 @@ impl TransactionManagerInner { if other_committed_seq > started_seq { // Check if any of the current transaction's read keys were written by // the committed transaction. - if !read_keys.is_disjoint(&committed_txn.write_keys) { + if read_keys + .iter() + .any(|read_key| committed_txn.writes.contains_key(read_key)) + { return true; } @@ -376,8 +439,8 @@ impl TransactionManagerInner { // committed transaction write keys. for read_range in &read_ranges { if committed_txn - .write_keys - .iter() + .writes + .keys() .any(|write_key| read_range.contains(write_key)) { return true; @@ -400,6 +463,22 @@ mod tests { use slatedb_common::DbRand; use std::collections::HashSet; + fn exclusive_writes( + keys: impl IntoIterator, + ) -> HashMap { + keys.into_iter() + .map(|key| (Bytes::from(key), TransactionWriteKind::Exclusive)) + .collect() + } + + fn commutative_writes( + keys: impl IntoIterator, + ) -> HashMap { + keys.into_iter() + .map(|key| (Bytes::from(key), TransactionWriteKind::CommutativeMerge)) + .collect() + } + struct CheckConflictTestCase { name: &'static str, recent_committed_txns: Vec, @@ -495,7 +574,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: 50, committed_seq: Some(80), - write_keys: ["key1", "key2"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1", "key2"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -508,7 +587,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: 50, committed_seq: Some(150), - write_keys: ["key1"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -522,14 +601,14 @@ mod tests { TransactionState { started_seq: 30, committed_seq: Some(50), - write_keys: ["key1"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }, TransactionState { started_seq: 80, committed_seq: Some(150), - write_keys: ["key2"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key2"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }, @@ -543,7 +622,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: 30, committed_seq: Some(50), - write_keys: ["key1"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -556,7 +635,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: 100, committed_seq: Some(100), - write_keys: ["key1"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -569,10 +648,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: 80, committed_seq: Some(150), - write_keys: ["key1", "key2", "key3"] - .into_iter() - .map(Bytes::from) - .collect(), + writes: exclusive_writes(["key1", "key2", "key3"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -585,7 +661,7 @@ mod tests { recent_committed_txns: vec![TransactionState { started_seq: u64::MAX - 1, committed_seq: Some(u64::MAX), - write_keys: ["key1"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }], @@ -603,15 +679,12 @@ mod tests { } // Convert current transaction write keys - let conflict_keys: HashSet = case - .current_write_keys - .into_iter() - .map(Bytes::from) - .collect(); + let conflict_writes = exclusive_writes(case.current_write_keys); // Call the method under test let inner = txn_manager.inner.read(); - let has_conflict = inner.has_write_write_conflict(&conflict_keys, case.current_started_seq); + let has_conflict = + inner.has_write_write_conflict(&conflict_writes, case.current_started_seq); // Verify result assert_eq!( @@ -621,6 +694,131 @@ mod tests { ); } + #[rstest] + #[case::commutative_with_commutative( + TransactionWriteKind::CommutativeMerge, + TransactionWriteKind::CommutativeMerge, + false + )] + #[case::commutative_with_exclusive( + TransactionWriteKind::CommutativeMerge, + TransactionWriteKind::Exclusive, + true + )] + #[case::exclusive_with_commutative( + TransactionWriteKind::Exclusive, + TransactionWriteKind::CommutativeMerge, + true + )] + #[case::exclusive_with_exclusive( + TransactionWriteKind::Exclusive, + TransactionWriteKind::Exclusive, + true + )] + fn test_write_write_conflict_respects_write_kind( + #[case] current_kind: TransactionWriteKind, + #[case] committed_kind: TransactionWriteKind, + #[case] expected_conflict: bool, + ) { + let txn_manager = create_transaction_manager(); + txn_manager + .inner + .write() + .recent_committed_txns + .push_back(TransactionState { + started_seq: 50, + committed_seq: Some(150), + writes: [(Bytes::from_static(b"key"), committed_kind)] + .into_iter() + .collect(), + read_keys: HashSet::new(), + read_ranges: Vec::new(), + }); + let writes = [(Bytes::from_static(b"key"), current_kind)] + .into_iter() + .collect(); + + assert_eq!( + txn_manager + .inner + .read() + .has_write_write_conflict(&writes, 100), + expected_conflict + ); + } + + #[test] + fn test_mixed_write_kinds_are_exclusive_in_either_order() { + let mut repeated_commutative = TransactionState { + started_seq: 0, + committed_seq: None, + writes: HashMap::new(), + read_keys: HashSet::new(), + read_ranges: Vec::new(), + }; + repeated_commutative.track_writes(commutative_writes(["key"])); + repeated_commutative.track_writes(commutative_writes(["key"])); + assert_eq!( + repeated_commutative.writes.get(b"key".as_slice()), + Some(&TransactionWriteKind::CommutativeMerge) + ); + + let mut commutative_then_exclusive = TransactionState { + started_seq: 0, + committed_seq: None, + writes: HashMap::new(), + read_keys: HashSet::new(), + read_ranges: Vec::new(), + }; + commutative_then_exclusive.track_writes(commutative_writes(["key"])); + commutative_then_exclusive.track_writes(exclusive_writes(["key"])); + assert_eq!( + commutative_then_exclusive.writes.get(b"key".as_slice()), + Some(&TransactionWriteKind::Exclusive) + ); + + let mut exclusive_then_commutative = TransactionState { + started_seq: 0, + committed_seq: None, + writes: HashMap::new(), + read_keys: HashSet::new(), + read_ranges: Vec::new(), + }; + exclusive_then_commutative.track_writes(exclusive_writes(["key"])); + exclusive_then_commutative.track_writes(commutative_writes(["key"])); + assert_eq!( + exclusive_then_commutative.writes.get(b"key".as_slice()), + Some(&TransactionWriteKind::Exclusive) + ); + } + + #[test] + fn test_commutative_writes_remain_visible_to_point_and_range_read_conflicts() { + let txn_manager = create_transaction_manager(); + txn_manager + .inner + .write() + .recent_committed_txns + .push_back(TransactionState { + started_seq: 50, + committed_seq: Some(150), + writes: commutative_writes(["foo5"]), + read_keys: HashSet::new(), + read_ranges: Vec::new(), + }); + let read_keys = [Bytes::from_static(b"foo5")].into_iter().collect(); + + let inner = txn_manager.inner.read(); + assert!(inner.has_read_write_conflict(&read_keys, Vec::new(), 100)); + assert!(inner.has_read_write_conflict( + &HashSet::new(), + vec![BytesRange::from( + Bytes::from_static(b"foo0")..=Bytes::from_static(b"foo9") + )], + 100 + )); + } + #[derive(Debug)] struct MinActiveSeqTestCase { name: &'static str, @@ -686,7 +884,7 @@ mod tests { expected_recent_committed_txn: Some(TransactionState { started_seq: 100, committed_seq: Some(150), - write_keys: ["key1", "key2"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1", "key2"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }), @@ -720,7 +918,7 @@ mod tests { expected_recent_committed_txn: Some(TransactionState { started_seq: 100, committed_seq: Some(100), - write_keys: ["key1", "key2"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["key1", "key2"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }), @@ -751,7 +949,7 @@ mod tests { expected_recent_committed_txn: Some(TransactionState { started_seq: 100, committed_seq: Some(150), - write_keys: ["existing_key", "key1", "key2"].into_iter().map(Bytes::from).collect(), + writes: exclusive_writes(["existing_key", "key1", "key2"]), read_keys: HashSet::new(), read_ranges: Vec::new(), }), @@ -806,9 +1004,9 @@ mod tests { test_case.name, expected_txn.committed_seq, actual_txn.committed_seq ); assert_eq!( - actual_txn.write_keys, expected_txn.write_keys, - "Test case '{}' failed: expected write_keys {:?}, got {:?}", - test_case.name, expected_txn.write_keys, actual_txn.write_keys + actual_txn.writes, expected_txn.writes, + "Test case '{}' failed: expected writes {:?}, got {:?}", + test_case.name, expected_txn.writes, actual_txn.writes ); } } @@ -919,7 +1117,7 @@ mod tests { inner.recent_committed_txns.push_back(TransactionState { started_seq: 50, committed_seq: None, // This should not happen in practice but let's test - write_keys: HashSet::new(), + writes: HashMap::new(), read_keys: HashSet::new(), read_ranges: Vec::new(), }); @@ -1403,13 +1601,16 @@ mod tests { } // Direct read-write conflict on keys - let direct_conflict = !txn.read_keys.is_disjoint(&committed.write_keys); + let direct_conflict = txn + .read_keys + .iter() + .any(|read_key| committed.writes.contains_key(read_key)); // Phantom conflict via range containment let mut phantom_conflict = false; - if !txn.read_ranges.is_empty() && !committed.write_keys.is_empty() { + if !txn.read_ranges.is_empty() && !committed.writes.is_empty() { 'outer: for range in txn.read_ranges.iter() { - for w in committed.write_keys.iter() { + for w in committed.writes.keys() { if range.contains(w) { phantom_conflict = true; break 'outer;