1use crate::database::Database;
12use crate::epoch::{Epoch, Snapshot};
13use crate::error::{MongrelError, Result};
14use crate::memtable::Value;
15use crate::rowid::RowId;
16use crate::wal::SharedWal;
17use parking_lot::{Condvar, Mutex as PlMutex};
18
19pub(crate) enum Staged {
21 Put(Vec<(u16, Value)>),
22 Delete(RowId),
23 Truncate,
24}
25
26#[derive(Debug, Clone)]
27pub struct OwnedRow {
28 pub columns: Vec<(u16, Value)>,
29}
30
31#[derive(Debug, Clone)]
32pub struct PutResult {
33 pub auto_inc: Option<i64>,
34 pub row: OwnedRow,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum UpsertActionKind {
39 Inserted,
40 Updated,
41 Unchanged,
42}
43
44#[derive(Debug, Clone)]
45pub enum UpsertAction {
46 DoNothing,
47 DoUpdate(Vec<(u16, Value)>),
48}
49
50#[derive(Debug, Clone)]
51pub struct UpsertResult {
52 pub action: UpsertActionKind,
53 pub row: OwnedRow,
54 pub auto_inc: Option<i64>,
55}
56
57pub struct Transaction<'db> {
60 db: &'db Database,
61 txn_id: u64,
62 read: Snapshot,
63 staging: Vec<(u64 , Staged)>,
64 _active: Option<ActiveTxnGuard<'db>>,
65}
66
67impl<'db> Transaction<'db> {
68 pub(crate) fn new(db: &'db Database, txn_id: u64, read: Snapshot) -> Self {
69 let guard = db.register_active(read.epoch);
70 Self {
71 db,
72 txn_id,
73 read,
74 staging: Vec::new(),
75 _active: Some(guard),
76 }
77 }
78
79 pub fn read_snapshot(&self) -> Snapshot {
80 self.read
81 }
82
83 pub fn txn_id(&self) -> u64 {
86 self.txn_id
87 }
88
89 pub fn put(&mut self, table: &str, mut cells: Vec<(u16, Value)>) -> Result<Option<i64>> {
96 let id = self.db.table_id(table)?;
97 self.reject_after_truncate(id)?;
98 let handle = self.db.table(table)?;
99 let mut t = handle.lock();
100 let assigned = t.fill_auto_inc(&mut cells)?;
101 drop(t);
102 self.staging.push((id, Staged::Put(cells)));
103 Ok(assigned)
104 }
105
106 pub fn put_returning(
107 &mut self,
108 table: &str,
109 mut cells: Vec<(u16, Value)>,
110 ) -> Result<PutResult> {
111 let id = self.db.table_id(table)?;
112 self.reject_after_truncate(id)?;
113 let handle = self.db.table(table)?;
114 let assigned = handle.lock().fill_auto_inc(&mut cells)?;
115 let row = owned_row_from_cells(&cells);
116 self.staging.push((id, Staged::Put(cells)));
117 Ok(PutResult {
118 auto_inc: assigned,
119 row,
120 })
121 }
122
123 pub fn delete(&mut self, table: &str, row_id: RowId) -> Result<()> {
125 let id = self.db.table_id(table)?;
126 self.reject_after_truncate(id)?;
127 self.staging.push((id, Staged::Delete(row_id)));
128 Ok(())
129 }
130
131 pub fn delete_many(&mut self, table: &str, row_ids: Vec<RowId>) -> Result<Vec<OwnedRow>> {
132 let id = self.db.table_id(table)?;
133 self.reject_after_truncate(id)?;
134 let snap = self.read;
135 let handle = self.db.table(table)?;
136 let t = handle.lock();
137 let mut pre_images = Vec::with_capacity(row_ids.len());
138 for row_id in &row_ids {
139 if let Some(row) = t.get(*row_id, snap) {
140 pre_images.push(owned_row_from_map(row.columns));
141 }
142 }
143 drop(t);
144 for row_id in row_ids {
145 self.staging.push((id, Staged::Delete(row_id)));
146 }
147 Ok(pre_images)
148 }
149
150 pub fn update_many(
151 &mut self,
152 table: &str,
153 updates: Vec<(RowId, Vec<(u16, Value)>)>,
154 ) -> Result<Vec<OwnedRow>> {
155 let id = self.db.table_id(table)?;
156 self.reject_after_truncate(id)?;
157 let snap = self.read;
158 let handle = self.db.table(table)?;
159 let t = handle.lock();
160 let mut post_images = Vec::with_capacity(updates.len());
161 let mut staged = Vec::with_capacity(updates.len() * 2);
162 for (old_id, new_cells) in updates {
163 let old_row = t
164 .get(old_id, snap)
165 .ok_or_else(|| MongrelError::NotFound(format!("row {old_id:?} not found")))?;
166 let merged = merge_cells(old_row.columns.into_iter().collect(), new_cells);
167 post_images.push(owned_row_from_cells(&merged));
168 staged.push((id, Staged::Delete(old_id)));
169 staged.push((id, Staged::Put(merged)));
170 }
171 drop(t);
172 self.staging.extend(staged);
173 Ok(post_images)
174 }
175
176 pub fn upsert(
177 &mut self,
178 table: &str,
179 mut insert_cells: Vec<(u16, Value)>,
180 action: UpsertAction,
181 ) -> Result<UpsertResult> {
182 let id = self.db.table_id(table)?;
183 self.reject_after_truncate(id)?;
184 match (self.existing_pk_row(table, &insert_cells)?, action) {
185 (None, _) => {
186 let assigned = self
187 .db
188 .table(table)?
189 .lock()
190 .fill_auto_inc(&mut insert_cells)?;
191 let row = owned_row_from_cells(&insert_cells);
192 self.staging.push((id, Staged::Put(insert_cells)));
193 Ok(UpsertResult {
194 action: UpsertActionKind::Inserted,
195 row,
196 auto_inc: assigned,
197 })
198 }
199 (Some((_old_id, old_row)), UpsertAction::DoNothing) => Ok(UpsertResult {
200 action: UpsertActionKind::Unchanged,
201 row: old_row,
202 auto_inc: None,
203 }),
204 (Some((old_id, old_row)), UpsertAction::DoUpdate(update_cells)) => {
205 let merged = merge_cells(old_row.columns.clone(), update_cells);
206 if columns_equal(&old_row.columns, &merged) {
207 return Ok(UpsertResult {
208 action: UpsertActionKind::Unchanged,
209 row: old_row,
210 auto_inc: None,
211 });
212 }
213 let row = owned_row_from_cells(&merged);
214 self.staging.push((id, Staged::Delete(old_id)));
215 self.staging.push((id, Staged::Put(merged)));
216 Ok(UpsertResult {
217 action: UpsertActionKind::Updated,
218 row,
219 auto_inc: None,
220 })
221 }
222 }
223 }
224
225 pub fn truncate(&mut self, table: &str) -> Result<()> {
226 let id = self.db.table_id(table)?;
227 for (table_id, op) in &self.staging {
228 if *table_id == id && !matches!(op, Staged::Truncate) {
229 return Err(MongrelError::InvalidArgument(
230 "truncate cannot be combined with other writes on the same table".into(),
231 ));
232 }
233 }
234 self.staging.push((id, Staged::Truncate));
235 Ok(())
236 }
237
238 fn reject_after_truncate(&self, table_id: u64) -> Result<()> {
239 if self
240 .staging
241 .iter()
242 .any(|(tid, op)| *tid == table_id && matches!(op, Staged::Truncate))
243 {
244 return Err(MongrelError::InvalidArgument(
245 "truncate cannot be combined with other writes on the same table".into(),
246 ));
247 }
248 Ok(())
249 }
250
251 fn existing_pk_row(
252 &self,
253 table: &str,
254 cells: &[(u16, Value)],
255 ) -> Result<Option<(RowId, OwnedRow)>> {
256 let handle = self.db.table(table)?;
257 let t = handle.lock();
258 let Some(pk_col) = t.schema().primary_key() else {
259 return Ok(None);
260 };
261 let Some((_, pk_value)) = cells.iter().find(|(id, _)| *id == pk_col.id) else {
262 return Ok(None);
263 };
264 if matches!(pk_value, Value::Null) {
265 return Ok(None);
266 }
267 let Some(row_id) = t.lookup_pk(&pk_value.encode_key()) else {
268 return Ok(None);
269 };
270 Ok(t.get(row_id, self.read)
271 .map(|row| (row_id, owned_row_from_map(row.columns))))
272 }
273
274 pub fn commit(self) -> Result<Epoch> {
276 self.db
277 .commit_transaction(self.txn_id, self.read.epoch, self.staging)
278 }
279
280 pub fn rollback(self) {
282 }
284}
285
286fn owned_row_from_cells(cells: &[(u16, Value)]) -> OwnedRow {
287 let mut columns = cells.to_vec();
288 columns.sort_by_key(|(id, _)| *id);
289 OwnedRow { columns }
290}
291
292fn owned_row_from_map(columns: HashMap<u16, Value>) -> OwnedRow {
293 let mut columns: Vec<(u16, Value)> = columns.into_iter().collect();
294 columns.sort_by_key(|(id, _)| *id);
295 OwnedRow { columns }
296}
297
298fn merge_cells(mut base: Vec<(u16, Value)>, updates: Vec<(u16, Value)>) -> Vec<(u16, Value)> {
299 for (id, value) in updates {
300 base.retain(|(existing, _)| *existing != id);
301 base.push((id, value));
302 }
303 base.sort_by_key(|(id, _)| *id);
304 base
305}
306
307fn columns_equal(a: &[(u16, Value)], b: &[(u16, Value)]) -> bool {
308 if a.len() != b.len() {
309 return false;
310 }
311 let mut a: Vec<_> = a.iter().collect();
312 let mut b: Vec<_> = b.iter().collect();
313 a.sort_by_key(|(id, _)| *id);
314 b.sort_by_key(|(id, _)| *id);
315 a.iter()
316 .zip(b.iter())
317 .all(|((id_a, v_a), (id_b, v_b))| id_a == id_b && v_a == v_b)
318}
319
320pub(crate) enum StagedOp {
322 Put(crate::memtable::Row),
323 Delete(RowId),
324 Truncate,
325}
326
327use std::collections::{BTreeMap, HashMap};
330use std::hash::{Hash, Hasher};
331
332#[derive(Clone, Debug)]
335pub enum WriteKey {
336 Row { table_id: u64, row_id: u64 },
338 Unique {
340 table_id: u64,
341 index_id: u16,
342 key_hash: u64,
343 },
344 Table { table_id: u64 },
346}
347
348impl Hash for WriteKey {
349 fn hash<H: Hasher>(&self, state: &mut H) {
350 match self {
351 WriteKey::Row { table_id, row_id } => {
352 0u8.hash(state);
353 table_id.hash(state);
354 row_id.hash(state);
355 }
356 WriteKey::Unique {
357 table_id,
358 index_id,
359 key_hash,
360 } => {
361 1u8.hash(state);
362 table_id.hash(state);
363 index_id.hash(state);
364 key_hash.hash(state);
365 }
366 WriteKey::Table { table_id } => {
367 2u8.hash(state);
368 table_id.hash(state);
369 }
370 }
371 }
372}
373
374impl PartialEq for WriteKey {
375 fn eq(&self, other: &Self) -> bool {
376 match (self, other) {
377 (
378 WriteKey::Row {
379 table_id: a,
380 row_id: b,
381 },
382 WriteKey::Row {
383 table_id: c,
384 row_id: d,
385 },
386 ) => a == c && b == d,
387 (
388 WriteKey::Unique {
389 table_id: a,
390 index_id: b,
391 key_hash: c,
392 },
393 WriteKey::Unique {
394 table_id: d,
395 index_id: e,
396 key_hash: f,
397 },
398 ) => a == d && b == e && c == f,
399 (WriteKey::Table { table_id: a }, WriteKey::Table { table_id: b }) => a == b,
400 _ => false,
401 }
402 }
403}
404
405impl Eq for WriteKey {}
406
407const CONFLICT_SHARDS: usize = 16;
408
409pub struct ConflictIndex {
413 shards: [parking_lot::Mutex<HashMap<WriteKey, u64>>; CONFLICT_SHARDS],
414 table_truncate_epochs: parking_lot::Mutex<HashMap<u64, u64>>,
415 table_write_epochs: parking_lot::Mutex<HashMap<u64, u64>>,
416 version: std::sync::atomic::AtomicU64,
420}
421
422impl ConflictIndex {
423 pub fn new() -> Self {
424 Self {
425 shards: std::array::from_fn(|_| parking_lot::Mutex::new(HashMap::new())),
426 table_truncate_epochs: parking_lot::Mutex::new(HashMap::new()),
427 table_write_epochs: parking_lot::Mutex::new(HashMap::new()),
428 version: std::sync::atomic::AtomicU64::new(0),
429 }
430 }
431
432 pub fn version(&self) -> u64 {
436 self.version.load(std::sync::atomic::Ordering::Acquire)
437 }
438
439 fn shard(&self, key: &WriteKey) -> &parking_lot::Mutex<HashMap<WriteKey, u64>> {
440 let mut h = std::collections::hash_map::DefaultHasher::new();
441 key.hash(&mut h);
442 let idx = (h.finish() as usize) & (CONFLICT_SHARDS - 1);
443 &self.shards[idx]
444 }
445
446 pub fn conflicts(&self, keys: &[WriteKey], read_epoch: Epoch) -> bool {
449 for k in keys {
450 let s = self.shard(k);
451 if let Some(&ce) = s.lock().get(k) {
452 if ce > read_epoch.0 {
453 return true;
454 }
455 }
456 }
457 let truncates = self.table_truncate_epochs.lock();
458 let writes = self.table_write_epochs.lock();
459 for k in keys {
460 match k {
461 WriteKey::Row { table_id, .. } | WriteKey::Unique { table_id, .. } => {
462 if truncates.get(table_id).is_some_and(|&ce| ce > read_epoch.0) {
463 return true;
464 }
465 }
466 WriteKey::Table { table_id } => {
467 if writes.get(table_id).is_some_and(|&ce| ce > read_epoch.0) {
468 return true;
469 }
470 }
471 }
472 }
473 false
474 }
475
476 pub fn record(&self, keys: &[WriteKey], commit_epoch: Epoch) {
478 for k in keys {
479 let s = self.shard(k);
480 s.lock().insert(k.clone(), commit_epoch.0);
481 }
482 let mut truncates = self.table_truncate_epochs.lock();
483 let mut writes = self.table_write_epochs.lock();
484 for k in keys {
485 match k {
486 WriteKey::Table { table_id } => {
487 truncates
488 .entry(*table_id)
489 .and_modify(|ce| *ce = (*ce).max(commit_epoch.0))
490 .or_insert(commit_epoch.0);
491 }
492 WriteKey::Row { table_id, .. } | WriteKey::Unique { table_id, .. } => {
493 writes
494 .entry(*table_id)
495 .and_modify(|ce| *ce = (*ce).max(commit_epoch.0))
496 .or_insert(commit_epoch.0);
497 }
498 }
499 }
500 self.version
501 .fetch_add(1, std::sync::atomic::Ordering::Release);
502 }
503
504 pub fn prune_below(&self, min_active: Epoch) {
507 for s in &self.shards {
508 s.lock().retain(|_, ce| *ce >= min_active.0);
509 }
510 self.table_truncate_epochs
511 .lock()
512 .retain(|_, ce| *ce >= min_active.0);
513 self.table_write_epochs
514 .lock()
515 .retain(|_, ce| *ce >= min_active.0);
516 }
517}
518
519impl Default for ConflictIndex {
520 fn default() -> Self {
521 Self::new()
522 }
523}
524
525pub struct GroupCommit {
535 inner: PlMutex<GroupState>,
536 cv: Condvar,
537}
538
539struct GroupState {
540 durable_seq: u64,
541 syncing: bool,
542 poisoned: bool,
543}
544
545impl GroupCommit {
546 pub fn new(durable_seq: u64) -> Self {
547 Self {
548 inner: PlMutex::new(GroupState {
549 durable_seq,
550 syncing: false,
551 poisoned: false,
552 }),
553 cv: Condvar::new(),
554 }
555 }
556
557 pub fn await_durable(&self, wal: &PlMutex<SharedWal>, commit_seq: u64) -> Result<()> {
563 let mut st = self.inner.lock();
564 loop {
565 if st.poisoned {
566 return Err(MongrelError::Other(
567 "database poisoned by fsync error".into(),
568 ));
569 }
570 if st.durable_seq >= commit_seq {
571 return Ok(());
572 }
573 if st.syncing {
574 self.cv.wait(&mut st);
576 continue;
577 }
578 st.syncing = true;
581 drop(st);
582 let res = wal.lock().group_sync();
583 st = self.inner.lock();
584 st.syncing = false;
585 match res {
586 Ok(durable) => {
587 if durable > st.durable_seq {
588 st.durable_seq = durable;
589 }
590 self.cv.notify_all();
591 }
594 Err(e) => {
595 st.poisoned = true;
596 self.cv.notify_all();
597 return Err(e);
598 }
599 }
600 }
601 }
602}
603
604pub struct ActiveTxns {
608 inner: parking_lot::Mutex<BTreeMap<u64, u64>>,
609}
610
611impl ActiveTxns {
612 pub fn new() -> Self {
613 Self {
614 inner: parking_lot::Mutex::new(BTreeMap::new()),
615 }
616 }
617
618 pub fn register(&self, read_epoch: Epoch) -> ActiveTxnGuard<'_> {
621 let mut g = self.inner.lock();
622 *g.entry(read_epoch.0).or_insert(0) += 1;
623 ActiveTxnGuard {
624 active: self,
625 epoch: read_epoch.0,
626 }
627 }
628
629 pub fn min_read_epoch(&self) -> u64 {
631 self.inner.lock().keys().next().copied().unwrap_or(u64::MAX)
632 }
633}
634
635impl Default for ActiveTxns {
636 fn default() -> Self {
637 Self::new()
638 }
639}
640
641pub struct ActiveTxnGuard<'a> {
643 active: &'a ActiveTxns,
644 epoch: u64,
645}
646
647impl Drop for ActiveTxnGuard<'_> {
648 fn drop(&mut self) {
649 let mut g = self.active.inner.lock();
650 if let Some(count) = g.get_mut(&self.epoch) {
651 *count -= 1;
652 if *count == 0 {
653 g.remove(&self.epoch);
654 }
655 }
656 }
657}
658
659#[cfg(test)]
660mod tests {
661 use super::*;
662
663 #[test]
664 fn conflict_index_first_committer_wins_and_prunes_safely() {
665 let ci = ConflictIndex::new();
666 let k = vec![WriteKey::Row {
667 table_id: 1,
668 row_id: 7,
669 }];
670 assert!(!ci.conflicts(&k, Epoch(5)));
671 ci.record(&k, Epoch(6));
672 assert!(ci.conflicts(&k, Epoch(5)));
673 assert!(!ci.conflicts(&k, Epoch(6)));
674 ci.prune_below(Epoch(7));
675 assert!(!ci.conflicts(&k, Epoch(5)));
676 }
677
678 #[test]
679 fn conflict_index_table_scope_conflicts_both_directions() {
680 let ci = ConflictIndex::new();
681 ci.record(&[WriteKey::Table { table_id: 1 }], Epoch(6));
682 assert!(ci.conflicts(
683 &[WriteKey::Row {
684 table_id: 1,
685 row_id: 7,
686 }],
687 Epoch(5)
688 ));
689 assert!(ci.conflicts(
690 &[WriteKey::Unique {
691 table_id: 1,
692 index_id: 0,
693 key_hash: 42,
694 }],
695 Epoch(5)
696 ));
697 assert!(!ci.conflicts(
698 &[WriteKey::Row {
699 table_id: 2,
700 row_id: 7,
701 }],
702 Epoch(5)
703 ));
704
705 let ci = ConflictIndex::new();
706 ci.record(
707 &[WriteKey::Row {
708 table_id: 1,
709 row_id: 7,
710 }],
711 Epoch(6),
712 );
713 assert!(ci.conflicts(&[WriteKey::Table { table_id: 1 }], Epoch(5)));
714 assert!(!ci.conflicts(&[WriteKey::Table { table_id: 2 }], Epoch(5)));
715 }
716
717 #[test]
718 fn writekey_eq_across_variants() {
719 let r1 = WriteKey::Row {
720 table_id: 1,
721 row_id: 2,
722 };
723 let r2 = WriteKey::Row {
724 table_id: 1,
725 row_id: 2,
726 };
727 let r3 = WriteKey::Row {
728 table_id: 1,
729 row_id: 3,
730 };
731 assert_eq!(r1, r2);
732 assert_ne!(r1, r3);
733
734 let u1 = WriteKey::Unique {
735 table_id: 1,
736 index_id: 0,
737 key_hash: 42,
738 };
739 let u2 = WriteKey::Unique {
740 table_id: 1,
741 index_id: 0,
742 key_hash: 42,
743 };
744 assert_eq!(u1, u2);
745 assert_ne!(r1, u1);
746
747 let t1 = WriteKey::Table { table_id: 5 };
748 let t2 = WriteKey::Table { table_id: 5 };
749 assert_eq!(t1, t2);
750 assert_ne!(t1, r1);
751 }
752
753 #[test]
754 fn active_txns_tracks_min_read_epoch() {
755 let at = ActiveTxns::new();
756 assert_eq!(at.min_read_epoch(), u64::MAX);
757 let g1 = at.register(Epoch(5));
758 assert_eq!(at.min_read_epoch(), 5);
759 let g2 = at.register(Epoch(3));
760 assert_eq!(at.min_read_epoch(), 3);
761 drop(g2);
762 assert_eq!(at.min_read_epoch(), 5);
763 drop(g1);
764 assert_eq!(at.min_read_epoch(), u64::MAX);
765 }
766
767 #[test]
768 fn active_txns_dedups_same_epoch() {
769 let at = ActiveTxns::new();
770 let g1 = at.register(Epoch(7));
771 let g2 = at.register(Epoch(7));
772 assert_eq!(at.min_read_epoch(), 7);
773 drop(g1);
774 assert_eq!(at.min_read_epoch(), 7);
775 drop(g2);
776 assert_eq!(at.min_read_epoch(), u64::MAX);
777 }
778}