1use std::collections::{HashMap, HashSet};
49use std::sync::Arc;
50use std::sync::atomic::{AtomicU64, Ordering};
51
52use dashmap::DashMap;
53use parking_lot::RwLock;
54use smallvec::SmallVec;
55
56use crate::durable_storage::InlineKey;
57
58pub type TxnId = u64;
60
61pub type Timestamp = u64;
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum SsiTxnStatus {
67 Active,
69 Committed(Timestamp),
71 Aborted,
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub enum ConflictType {
78 WriteWrite,
80 ReadWriteAnti,
82 DangerousStructure,
84}
85
86#[derive(Debug, Clone)]
88pub struct SsiConflictError {
89 pub victim_txn: TxnId,
91 pub winner_txn: Option<TxnId>,
93 pub conflict_type: ConflictType,
95 pub message: String,
97}
98
99impl std::fmt::Display for SsiConflictError {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 write!(
102 f,
103 "SSI conflict ({:?}): {}",
104 self.conflict_type, self.message
105 )
106 }
107}
108
109impl std::error::Error for SsiConflictError {}
110
111#[derive(Debug, Clone)]
113pub struct SsiVersionInfo {
114 pub xmin: TxnId,
116 pub xmax: TxnId,
118 pub begin_ts: Timestamp,
120 pub end_ts: Timestamp,
122 pub commit_ts: Option<Timestamp>,
124}
125
126impl SsiVersionInfo {
127 pub fn new(xmin: TxnId, begin_ts: Timestamp) -> Self {
129 Self {
130 xmin,
131 xmax: 0,
132 begin_ts,
133 end_ts: Timestamp::MAX,
134 commit_ts: None,
135 }
136 }
137
138 pub fn is_visible(&self, snapshot_ts: Timestamp, txn_states: &SsiTxnStates) -> bool {
144 match txn_states.get_status(self.xmin) {
146 Some(SsiTxnStatus::Committed(commit_ts)) => {
147 if commit_ts > snapshot_ts {
148 return false; }
150 }
151 Some(SsiTxnStatus::Active) | Some(SsiTxnStatus::Aborted) | None => {
152 return false; }
154 }
155
156 if self.xmax == 0 {
158 return true; }
160
161 match txn_states.get_status(self.xmax) {
162 Some(SsiTxnStatus::Committed(commit_ts)) => {
163 commit_ts > snapshot_ts }
165 Some(SsiTxnStatus::Active) | Some(SsiTxnStatus::Aborted) | None => {
166 true }
168 }
169 }
170}
171
172#[derive(Debug, Clone, PartialEq, Eq, Hash)]
174pub struct RwDependency {
175 pub reader: TxnId,
177 pub writer: TxnId,
179 pub key: InlineKey,
181}
182
183#[derive(Debug)]
185pub struct SsiTransaction {
186 pub txn_id: TxnId,
188 pub start_ts: Timestamp,
190 pub status: SsiTxnStatus,
192 pub commit_ts: Option<Timestamp>,
194 pub read_set: HashSet<InlineKey>,
197 pub write_set: HashSet<InlineKey>,
199 read_bloom: [u64; 4],
202 pub in_rw_deps: HashSet<TxnId>,
204 pub out_rw_deps: HashSet<TxnId>,
206 pub has_committed_in_rw: bool,
208 pub has_committed_out_rw: bool,
210}
211
212impl SsiTransaction {
213 pub fn new(txn_id: TxnId, start_ts: Timestamp) -> Self {
215 Self {
216 txn_id,
217 start_ts,
218 status: SsiTxnStatus::Active,
219 commit_ts: None,
220 read_set: HashSet::new(),
221 write_set: HashSet::new(),
222 read_bloom: [0u64; 4],
223 in_rw_deps: HashSet::new(),
224 out_rw_deps: HashSet::new(),
225 has_committed_in_rw: false,
226 has_committed_out_rw: false,
227 }
228 }
229
230 pub fn record_read(&mut self, key: &[u8]) {
232 let ik = SmallVec::from_slice(key);
233 self.read_set.insert(ik);
234 let h = twox_hash::xxh3::hash64(key);
236 let h1 = (h & 0xFF) as usize; let h2 = ((h >> 8) & 0xFF) as usize; self.read_bloom[h1 / 64] |= 1u64 << (h1 % 64);
239 self.read_bloom[h2 / 64] |= 1u64 << (h2 % 64);
240 }
241
242 pub fn record_write(&mut self, key: &[u8]) {
244 self.write_set.insert(SmallVec::from_slice(key));
245 }
246
247 pub fn maybe_read(&self, key: &[u8]) -> bool {
250 let h = twox_hash::xxh3::hash64(key);
251 let h1 = (h & 0xFF) as usize;
252 let h2 = ((h >> 8) & 0xFF) as usize;
253 (self.read_bloom[h1 / 64] & (1u64 << (h1 % 64))) != 0
254 && (self.read_bloom[h2 / 64] & (1u64 << (h2 % 64))) != 0
255 }
256
257 pub fn is_dangerous(&self) -> bool {
263 self.has_committed_in_rw && self.has_committed_out_rw
264 }
265}
266
267pub struct SsiTxnStates {
269 states: RwLock<HashMap<TxnId, SsiTxnStatus>>,
271}
272
273impl SsiTxnStates {
274 pub fn new() -> Self {
275 Self {
276 states: RwLock::new(HashMap::new()),
277 }
278 }
279
280 pub fn get_status(&self, txn_id: TxnId) -> Option<SsiTxnStatus> {
281 self.states.read().get(&txn_id).copied()
282 }
283
284 pub fn set_status(&self, txn_id: TxnId, status: SsiTxnStatus) {
285 self.states.write().insert(txn_id, status);
286 }
287}
288
289impl Default for SsiTxnStates {
290 fn default() -> Self {
291 Self::new()
292 }
293}
294
295pub struct SsiManager {
317 next_txn_id: AtomicU64,
319 timestamp: AtomicU64,
321 transactions: RwLock<HashMap<TxnId, SsiTransaction>>,
323 txn_states: Arc<SsiTxnStates>,
325 key_writers: DashMap<Vec<u8>, (TxnId, Timestamp)>,
329 key_readers: DashMap<Vec<u8>, HashSet<TxnId>>,
333}
334
335impl SsiManager {
336 pub fn new() -> Self {
338 Self {
339 next_txn_id: AtomicU64::new(1),
340 timestamp: AtomicU64::new(1),
341 transactions: RwLock::new(HashMap::new()),
342 txn_states: Arc::new(SsiTxnStates::new()),
343 key_writers: DashMap::new(),
344 key_readers: DashMap::new(),
345 }
346 }
347
348 pub fn begin(&self) -> Result<(TxnId, Timestamp), SsiConflictError> {
350 let txn_id = self.next_txn_id.fetch_add(1, Ordering::SeqCst);
351 let start_ts = self.begin_with_id(txn_id)?;
352 Ok((txn_id, start_ts))
353 }
354
355 pub fn begin_with_id(&self, txn_id: TxnId) -> Result<Timestamp, SsiConflictError> {
362 let start_ts = self.timestamp.fetch_add(1, Ordering::SeqCst);
363
364 let txn = SsiTransaction::new(txn_id, start_ts);
365 self.transactions.write().insert(txn_id, txn);
366 self.txn_states.set_status(txn_id, SsiTxnStatus::Active);
367
368 Ok(start_ts)
369 }
370
371 pub fn record_read(&self, txn_id: TxnId, key: &[u8]) -> Result<(), SsiConflictError> {
376 let snapshot_ts = {
378 let mut txns = self.transactions.write();
379 let txn = txns.get_mut(&txn_id).ok_or_else(|| SsiConflictError {
380 victim_txn: txn_id,
381 winner_txn: None,
382 conflict_type: ConflictType::ReadWriteAnti,
383 message: "Transaction not found".into(),
384 })?;
385 let ts = txn.start_ts;
386 txn.record_read(key);
387 ts
388 };
389
390 self.key_readers
392 .entry(key.to_vec())
393 .or_default()
394 .insert(txn_id);
395
396 let writer_info = self.key_writers.get(key).map(|r| *r);
398 if let Some((writer_txn, write_ts)) = writer_info
399 && write_ts > snapshot_ts
400 && writer_txn != txn_id
401 {
402 let writer_committed = matches!(
405 self.txn_states.get_status(writer_txn),
406 Some(SsiTxnStatus::Committed(_))
407 );
408
409 let mut txns = self.transactions.write();
410
411 if let Some(reader_txn) = txns.get_mut(&txn_id) {
413 reader_txn.out_rw_deps.insert(writer_txn);
414 if writer_committed {
415 reader_txn.has_committed_out_rw = true;
416 if reader_txn.is_dangerous() {
418 return Err(SsiConflictError {
419 victim_txn: txn_id,
420 winner_txn: Some(writer_txn),
421 conflict_type: ConflictType::DangerousStructure,
422 message: format!(
423 "Transaction {} would create serialization anomaly with {}",
424 txn_id, writer_txn
425 ),
426 });
427 }
428 }
429 }
430
431 if let Some(writer_txn_entry) = txns.get_mut(&writer_txn) {
433 writer_txn_entry.in_rw_deps.insert(txn_id);
434 }
435 }
436
437 Ok(())
438 }
439
440 pub fn record_write(&self, txn_id: TxnId, key: &[u8]) -> Result<(), SsiConflictError> {
445 let mut txns = self.transactions.write();
446 let txn = txns.get_mut(&txn_id).ok_or_else(|| SsiConflictError {
447 victim_txn: txn_id,
448 winner_txn: None,
449 conflict_type: ConflictType::WriteWrite,
450 message: "Transaction not found".into(),
451 })?;
452
453 let snapshot_ts = txn.start_ts;
454
455 if let Some(entry) = self.key_writers.get(key) {
457 let (prev_writer, write_ts) = *entry;
458 if write_ts > snapshot_ts && prev_writer != txn_id {
459 return Err(SsiConflictError {
460 victim_txn: txn_id,
461 winner_txn: Some(prev_writer),
462 conflict_type: ConflictType::WriteWrite,
463 message: format!(
464 "Write-write conflict: transaction {} already wrote to key, ts {}",
465 prev_writer, write_ts
466 ),
467 });
468 }
469 }
470
471 txn.record_write(key);
473
474 let write_ts = self.timestamp.fetch_add(1, Ordering::SeqCst);
476 drop(txns);
477
478 self.key_writers.insert(key.to_vec(), (txn_id, write_ts));
479
480 if let Some(readers) = self.key_readers.get(key) {
482 let mut txns = self.transactions.write();
483 for reader_id in readers.value() {
484 if *reader_id != txn_id
485 && let Some(reader_txn) = txns.get(reader_id)
486 && reader_txn.start_ts < write_ts
487 {
488 if let Some(writer_txn) = txns.get_mut(&txn_id) {
490 writer_txn.in_rw_deps.insert(*reader_id);
491
492 if let Some(SsiTxnStatus::Committed(_)) =
494 self.txn_states.get_status(*reader_id)
495 {
496 writer_txn.has_committed_in_rw = true;
497 }
498 }
499 if let Some(reader_txn) = txns.get_mut(reader_id) {
500 reader_txn.out_rw_deps.insert(txn_id);
501 }
502 }
503 }
504 }
505
506 Ok(())
507 }
508
509 pub fn commit(&self, txn_id: TxnId) -> Result<Timestamp, SsiConflictError> {
513 let commit_ts = self.timestamp.fetch_add(1, Ordering::SeqCst);
514
515 let (is_dangerous, out_deps, in_deps) = {
517 let txns = self.transactions.read();
518 let txn = txns.get(&txn_id).ok_or_else(|| SsiConflictError {
519 victim_txn: txn_id,
520 winner_txn: None,
521 conflict_type: ConflictType::DangerousStructure,
522 message: "Transaction not found".into(),
523 })?;
524 (
525 txn.is_dangerous(),
526 txn.out_rw_deps.clone(),
527 txn.in_rw_deps.clone(),
528 )
529 };
530
531 if is_dangerous {
532 let mut txns = self.transactions.write();
533 if let Some(txn) = txns.get_mut(&txn_id) {
534 txn.status = SsiTxnStatus::Aborted;
535 }
536 self.txn_states.set_status(txn_id, SsiTxnStatus::Aborted);
537 return Err(SsiConflictError {
538 victim_txn: txn_id,
539 winner_txn: None,
540 conflict_type: ConflictType::DangerousStructure,
541 message: "Transaction would create serialization anomaly (dangerous structure)"
542 .into(),
543 });
544 }
545
546 {
548 let mut txns = self.transactions.write();
549
550 if let Some(txn) = txns.get_mut(&txn_id) {
552 txn.status = SsiTxnStatus::Committed(commit_ts);
553 txn.commit_ts = Some(commit_ts);
554 }
555
556 for out_dep in &out_deps {
558 if let Some(other_txn) = txns.get_mut(out_dep) {
559 other_txn.has_committed_in_rw = true;
560 }
561 }
562
563 for in_dep in &in_deps {
565 if let Some(other_txn) = txns.get_mut(in_dep) {
566 other_txn.has_committed_out_rw = true;
567 }
568 }
569 }
570
571 self.txn_states
572 .set_status(txn_id, SsiTxnStatus::Committed(commit_ts));
573 Ok(commit_ts)
574 }
575
576 pub fn abort(&self, txn_id: TxnId) {
578 let mut txns = self.transactions.write();
579 if let Some(txn) = txns.get_mut(&txn_id) {
580 txn.status = SsiTxnStatus::Aborted;
581 self.txn_states.set_status(txn_id, SsiTxnStatus::Aborted);
582 }
583
584 self.key_writers.retain(|_, (writer, _)| *writer != txn_id);
586
587 for mut entry in self.key_readers.iter_mut() {
589 entry.value_mut().remove(&txn_id);
590 }
591 }
592
593 pub fn get_status(&self, txn_id: TxnId) -> Option<SsiTxnStatus> {
595 self.txn_states.get_status(txn_id)
596 }
597
598 pub fn get_snapshot_ts(&self, txn_id: TxnId) -> Option<Timestamp> {
600 self.transactions.read().get(&txn_id).map(|t| t.start_ts)
601 }
602
603 pub fn is_visible(&self, txn_id: TxnId, version: &SsiVersionInfo) -> bool {
605 if let Some(snapshot_ts) = self.get_snapshot_ts(txn_id) {
606 version.is_visible(snapshot_ts, &self.txn_states)
607 } else {
608 false
609 }
610 }
611
612 pub fn gc(&self, watermark: Timestamp) -> usize {
614 let mut removed = 0;
615
616 self.transactions.write().retain(|_, txn| {
618 if let Some(commit_ts) = txn.commit_ts
619 && commit_ts < watermark
620 {
621 removed += 1;
622 return false;
623 }
624 true
625 });
626
627 removed
628 }
629}
630
631impl Default for SsiManager {
632 fn default() -> Self {
633 Self::new()
634 }
635}
636
637pub struct HybridLogicalClock {
645 timestamp: AtomicU64,
647}
648
649impl HybridLogicalClock {
650 const LOGICAL_BITS: u32 = 20;
651 const LOGICAL_MASK: u64 = (1 << Self::LOGICAL_BITS) - 1;
652
653 pub fn new() -> Self {
655 let now_ms = Self::physical_time_ms();
656 Self {
657 timestamp: AtomicU64::new(now_ms << Self::LOGICAL_BITS),
658 }
659 }
660
661 fn physical_time_ms() -> u64 {
663 use std::time::{SystemTime, UNIX_EPOCH};
664 SystemTime::now()
665 .duration_since(UNIX_EPOCH)
666 .unwrap()
667 .as_millis() as u64
668 }
669
670 pub fn get_physical(ts: u64) -> u64 {
672 ts >> Self::LOGICAL_BITS
673 }
674
675 pub fn get_logical(ts: u64) -> u64 {
677 ts & Self::LOGICAL_MASK
678 }
679
680 pub fn next(&self) -> u64 {
687 loop {
688 let current = self.timestamp.load(Ordering::Acquire);
689 let current_physical = Self::get_physical(current);
690 let current_logical = Self::get_logical(current);
691
692 let now_physical = Self::physical_time_ms();
693
694 let (new_physical, new_logical) = if now_physical > current_physical {
695 (now_physical, 0)
697 } else {
698 if current_logical >= Self::LOGICAL_MASK {
700 std::thread::yield_now();
702 continue;
703 }
704 (current_physical, current_logical + 1)
705 };
706
707 let new_ts = (new_physical << Self::LOGICAL_BITS) | new_logical;
708
709 if self
710 .timestamp
711 .compare_exchange(current, new_ts, Ordering::AcqRel, Ordering::Acquire)
712 .is_ok()
713 {
714 return new_ts;
715 }
716 }
717 }
718
719 pub fn update(&self, msg_ts: u64) {
723 loop {
724 let current = self.timestamp.load(Ordering::Acquire);
725 let now_physical = Self::physical_time_ms();
726
727 let new_ts = if msg_ts > current {
728 let msg_physical = Self::get_physical(msg_ts);
730 let msg_logical = Self::get_logical(msg_ts);
731
732 if msg_physical > now_physical {
733 let bounded_physical = now_physical.max(msg_physical.saturating_sub(1000));
735 (bounded_physical << Self::LOGICAL_BITS) | (msg_logical + 1)
736 } else {
737 (now_physical << Self::LOGICAL_BITS) | (msg_logical + 1)
738 }
739 } else {
740 return;
742 };
743
744 if self
745 .timestamp
746 .compare_exchange(current, new_ts, Ordering::AcqRel, Ordering::Acquire)
747 .is_ok()
748 {
749 return;
750 }
751 }
752 }
753
754 pub fn now(&self) -> u64 {
756 self.timestamp.load(Ordering::Acquire)
757 }
758}
759
760impl Default for HybridLogicalClock {
761 fn default() -> Self {
762 Self::new()
763 }
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769
770 #[test]
771 fn test_ssi_basic_commit() {
772 let ssi = SsiManager::new();
773
774 let (txn1, _) = ssi.begin().unwrap();
775 ssi.record_read(txn1, b"key1").unwrap();
776 ssi.record_write(txn1, b"key1").unwrap();
777 let commit_ts = ssi.commit(txn1).unwrap();
778
779 assert!(commit_ts > 0);
780 assert!(matches!(
781 ssi.get_status(txn1),
782 Some(SsiTxnStatus::Committed(_))
783 ));
784 }
785
786 #[test]
787 fn test_ssi_write_write_conflict() {
788 let ssi = SsiManager::new();
789
790 let (txn1, _) = ssi.begin().unwrap();
792
793 let (txn2, _) = ssi.begin().unwrap();
795 ssi.record_write(txn2, b"key1").unwrap();
796 ssi.commit(txn2).unwrap();
797
798 let result = ssi.record_write(txn1, b"key1");
800 assert!(result.is_err());
801 assert!(matches!(
802 result.unwrap_err().conflict_type,
803 ConflictType::WriteWrite
804 ));
805 }
806
807 #[test]
808 fn test_ssi_rw_antidependency() {
809 let ssi = SsiManager::new();
810
811 let (txn1, _) = ssi.begin().unwrap();
813 ssi.record_read(txn1, b"key1").unwrap();
814
815 let (txn2, _) = ssi.begin().unwrap();
817 ssi.record_write(txn2, b"key1").unwrap();
818 ssi.commit(txn2).unwrap();
819
820 let result = ssi.commit(txn1);
823 assert!(result.is_ok());
824 }
825
826 #[test]
827 fn test_ssi_dangerous_structure() {
828 let ssi = SsiManager::new();
829
830 let (txn1, _) = ssi.begin().unwrap();
834 ssi.record_read(txn1, b"key1").unwrap();
835
836 let (txn2, _) = ssi.begin().unwrap();
837 ssi.record_write(txn2, b"key1").unwrap();
838 ssi.commit(txn2).unwrap(); ssi.record_write(txn1, b"key2").unwrap();
841
842 let (txn3, _) = ssi.begin().unwrap();
844 ssi.record_read(txn3, b"key2").unwrap();
845 ssi.record_write(txn3, b"key1").unwrap();
848 ssi.commit(txn3).unwrap(); let _result = ssi.commit(txn1);
852 }
856
857 #[test]
858 fn test_hlc_monotonic() {
859 let hlc = HybridLogicalClock::new();
860
861 let mut prev = hlc.next();
862 for _ in 0..1000 {
863 let curr = hlc.next();
864 assert!(curr > prev, "HLC must be monotonic");
865 prev = curr;
866 }
867 }
868
869 #[test]
870 fn test_hlc_physical_extraction() {
871 let hlc = HybridLogicalClock::new();
872 let ts = hlc.next();
873
874 let physical = HybridLogicalClock::get_physical(ts);
875 let logical = HybridLogicalClock::get_logical(ts);
876
877 assert!(physical > 1577836800000); assert!(logical < 1000);
882 }
883
884 #[test]
885 fn test_version_visibility() {
886 let states = SsiTxnStates::new();
887
888 states.set_status(1, SsiTxnStatus::Committed(100));
890
891 let version = SsiVersionInfo::new(1, 100);
893
894 assert!(version.is_visible(150, &states));
896
897 assert!(!version.is_visible(50, &states));
899 }
900
901 #[test]
902 fn test_ssi_abort_cleanup() {
903 let ssi = SsiManager::new();
904
905 let (txn1, _) = ssi.begin().unwrap();
906 ssi.record_write(txn1, b"key1").unwrap();
907 ssi.abort(txn1);
908
909 let (txn2, _) = ssi.begin().unwrap();
911 let result = ssi.record_write(txn2, b"key1");
912 assert!(result.is_ok());
913 }
914
915 #[test]
916 fn test_ssi_bloom_filter_negative() {
917 let mut txn = SsiTransaction::new(1, 100);
919 txn.record_read(b"alpha");
920 txn.record_read(b"beta");
921
922 assert!(txn.maybe_read(b"alpha"));
924 assert!(txn.maybe_read(b"beta"));
925
926 let _ = txn.maybe_read(b"gamma");
930 }
931
932 #[test]
933 fn test_ssi_inline_key_read_write_sets() {
934 let mut txn = SsiTransaction::new(1, 100);
936
937 txn.record_read(b"short");
939 txn.record_write(b"short_w");
940
941 let k32 = [0xABu8; 32];
943 txn.record_read(&k32);
944 txn.record_write(&k32);
945
946 let k64 = [0xCDu8; 64];
948 txn.record_read(&k64);
949 txn.record_write(&k64);
950
951 assert_eq!(txn.read_set.len(), 3);
952 assert_eq!(txn.write_set.len(), 3);
953 }
954
955 #[test]
956 fn test_ssi_dashmap_concurrent_disjoint_keys() {
957 let ssi = SsiManager::new();
959 let mut txns = Vec::new();
960 for i in 0..10 {
961 let (tid, _) = ssi.begin().unwrap();
962 let key = format!("key_{}", i);
963 ssi.record_write(tid, key.as_bytes()).unwrap();
964 txns.push(tid);
965 }
966 for tid in txns {
968 assert!(ssi.commit(tid).is_ok());
969 }
970 }
971
972 #[test]
973 fn test_ssi_dashmap_abort_cleans_all_shards() {
974 let ssi = SsiManager::new();
975
976 let (txn1, _) = ssi.begin().unwrap();
977 for i in 0..20 {
979 let key = format!("shard_test_{}", i);
980 ssi.record_write(txn1, key.as_bytes()).unwrap();
981 }
982
983 ssi.abort(txn1);
984
985 for i in 0..20 {
987 let key = format!("shard_test_{}", i);
988 let has_entry = ssi.key_writers.get(key.as_bytes()).is_some();
989 assert!(!has_entry, "key_writers should be cleaned for {}", key);
990 }
991
992 let (txn2, _) = ssi.begin().unwrap();
994 ssi.record_write(txn2, b"shard_test_5").unwrap();
995 assert!(ssi.commit(txn2).is_ok());
996 }
997}