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 put_batch(
129 &mut self,
130 table: &str,
131 rows: Vec<Vec<(u16, Value)>>,
132 ) -> Result<Vec<Option<i64>>> {
133 let id = self.db.table_id(table)?;
134 self.reject_after_truncate(id)?;
135 let handle = self.db.table(table)?;
136 let mut t = handle.lock();
137 let mut assigned = Vec::with_capacity(rows.len());
138 for mut cells in rows {
139 let a = t.fill_auto_inc(&mut cells)?;
140 assigned.push(a);
141 self.staging.push((id, Staged::Put(cells)));
142 }
143 drop(t);
144 Ok(assigned)
145 }
146
147 pub fn delete(&mut self, table: &str, row_id: RowId) -> Result<()> {
149 let id = self.db.table_id(table)?;
150 self.reject_after_truncate(id)?;
151 self.staging.push((id, Staged::Delete(row_id)));
152 Ok(())
153 }
154
155 pub fn delete_many(&mut self, table: &str, row_ids: Vec<RowId>) -> Result<Vec<OwnedRow>> {
156 let id = self.db.table_id(table)?;
157 self.reject_after_truncate(id)?;
158 let snap = self.read;
159 let handle = self.db.table(table)?;
160 let t = handle.lock();
161 let mut pre_images = Vec::with_capacity(row_ids.len());
162 for row_id in &row_ids {
163 if let Some(row) = t.get(*row_id, snap) {
164 pre_images.push(owned_row_from_map(row.columns));
165 }
166 }
167 drop(t);
168 for row_id in row_ids {
169 self.staging.push((id, Staged::Delete(row_id)));
170 }
171 Ok(pre_images)
172 }
173
174 pub fn update_many(
175 &mut self,
176 table: &str,
177 updates: Vec<(RowId, Vec<(u16, Value)>)>,
178 ) -> Result<Vec<OwnedRow>> {
179 let id = self.db.table_id(table)?;
180 self.reject_after_truncate(id)?;
181 let snap = self.read;
182 let handle = self.db.table(table)?;
183 let t = handle.lock();
184 let mut post_images = Vec::with_capacity(updates.len());
185 let mut staged = Vec::with_capacity(updates.len() * 2);
186 for (old_id, new_cells) in updates {
187 let old_row = t
188 .get(old_id, snap)
189 .ok_or_else(|| MongrelError::NotFound(format!("row {old_id:?} not found")))?;
190 let merged = merge_cells(old_row.columns.into_iter().collect(), new_cells);
191 post_images.push(owned_row_from_cells(&merged));
192 staged.push((id, Staged::Delete(old_id)));
193 staged.push((id, Staged::Put(merged)));
194 }
195 drop(t);
196 self.staging.extend(staged);
197 Ok(post_images)
198 }
199
200 pub fn upsert(
201 &mut self,
202 table: &str,
203 mut insert_cells: Vec<(u16, Value)>,
204 action: UpsertAction,
205 ) -> Result<UpsertResult> {
206 let id = self.db.table_id(table)?;
207 self.reject_after_truncate(id)?;
208 match (self.existing_pk_row(table, &insert_cells)?, action) {
209 (None, _) => {
210 let assigned = self
211 .db
212 .table(table)?
213 .lock()
214 .fill_auto_inc(&mut insert_cells)?;
215 let row = owned_row_from_cells(&insert_cells);
216 self.staging.push((id, Staged::Put(insert_cells)));
217 Ok(UpsertResult {
218 action: UpsertActionKind::Inserted,
219 row,
220 auto_inc: assigned,
221 })
222 }
223 (Some((_old_id, old_row)), UpsertAction::DoNothing) => Ok(UpsertResult {
224 action: UpsertActionKind::Unchanged,
225 row: old_row,
226 auto_inc: None,
227 }),
228 (Some((old_id, old_row)), UpsertAction::DoUpdate(update_cells)) => {
229 let merged = merge_cells(old_row.columns.clone(), update_cells);
230 if columns_equal(&old_row.columns, &merged) {
231 return Ok(UpsertResult {
232 action: UpsertActionKind::Unchanged,
233 row: old_row,
234 auto_inc: None,
235 });
236 }
237 let row = owned_row_from_cells(&merged);
238 self.staging.push((id, Staged::Delete(old_id)));
239 self.staging.push((id, Staged::Put(merged)));
240 Ok(UpsertResult {
241 action: UpsertActionKind::Updated,
242 row,
243 auto_inc: None,
244 })
245 }
246 }
247 }
248
249 pub fn truncate(&mut self, table: &str) -> Result<()> {
250 let id = self.db.table_id(table)?;
251 for (table_id, op) in &self.staging {
252 if *table_id == id && !matches!(op, Staged::Truncate) {
253 return Err(MongrelError::InvalidArgument(
254 "truncate cannot be combined with other writes on the same table".into(),
255 ));
256 }
257 }
258 self.staging.push((id, Staged::Truncate));
259 Ok(())
260 }
261
262 fn reject_after_truncate(&self, table_id: u64) -> Result<()> {
263 if self
264 .staging
265 .iter()
266 .any(|(tid, op)| *tid == table_id && matches!(op, Staged::Truncate))
267 {
268 return Err(MongrelError::InvalidArgument(
269 "truncate cannot be combined with other writes on the same table".into(),
270 ));
271 }
272 Ok(())
273 }
274
275 fn existing_pk_row(
276 &self,
277 table: &str,
278 cells: &[(u16, Value)],
279 ) -> Result<Option<(RowId, OwnedRow)>> {
280 let handle = self.db.table(table)?;
281 let t = handle.lock();
282 let Some(pk_col) = t.schema().primary_key() else {
283 return Ok(None);
284 };
285 let Some((_, pk_value)) = cells.iter().find(|(id, _)| *id == pk_col.id) else {
286 return Ok(None);
287 };
288 if matches!(pk_value, Value::Null) {
289 return Ok(None);
290 }
291 let Some(row_id) = t.lookup_pk(&pk_value.encode_key()) else {
292 return Ok(None);
293 };
294 Ok(t.get(row_id, self.read)
295 .map(|row| (row_id, owned_row_from_map(row.columns))))
296 }
297
298 pub fn commit(self) -> Result<Epoch> {
300 self.db
301 .commit_transaction(self.txn_id, self.read.epoch, self.staging)
302 }
303
304 pub fn rollback(self) {
306 }
308}
309
310fn owned_row_from_cells(cells: &[(u16, Value)]) -> OwnedRow {
311 let mut columns = cells.to_vec();
312 columns.sort_by_key(|(id, _)| *id);
313 OwnedRow { columns }
314}
315
316fn owned_row_from_map(columns: HashMap<u16, Value>) -> OwnedRow {
317 let mut columns: Vec<(u16, Value)> = columns.into_iter().collect();
318 columns.sort_by_key(|(id, _)| *id);
319 OwnedRow { columns }
320}
321
322fn merge_cells(mut base: Vec<(u16, Value)>, updates: Vec<(u16, Value)>) -> Vec<(u16, Value)> {
323 for (id, value) in updates {
324 base.retain(|(existing, _)| *existing != id);
325 base.push((id, value));
326 }
327 base.sort_by_key(|(id, _)| *id);
328 base
329}
330
331fn columns_equal(a: &[(u16, Value)], b: &[(u16, Value)]) -> bool {
332 if a.len() != b.len() {
333 return false;
334 }
335 let mut a: Vec<_> = a.iter().collect();
336 let mut b: Vec<_> = b.iter().collect();
337 a.sort_by_key(|(id, _)| *id);
338 b.sort_by_key(|(id, _)| *id);
339 a.iter()
340 .zip(b.iter())
341 .all(|((id_a, v_a), (id_b, v_b))| id_a == id_b && v_a == v_b)
342}
343
344pub(crate) enum StagedOp {
346 Put(crate::memtable::Row),
347 Delete(RowId),
348 Truncate,
349}
350
351use std::collections::{BTreeMap, HashMap};
354use std::hash::{Hash, Hasher};
355
356#[derive(Clone, Debug)]
359pub enum WriteKey {
360 Row { table_id: u64, row_id: u64 },
362 Unique {
364 table_id: u64,
365 index_id: u16,
366 key_hash: u64,
367 },
368 Table { table_id: u64 },
370}
371
372impl Hash for WriteKey {
373 fn hash<H: Hasher>(&self, state: &mut H) {
374 match self {
375 WriteKey::Row { table_id, row_id } => {
376 0u8.hash(state);
377 table_id.hash(state);
378 row_id.hash(state);
379 }
380 WriteKey::Unique {
381 table_id,
382 index_id,
383 key_hash,
384 } => {
385 1u8.hash(state);
386 table_id.hash(state);
387 index_id.hash(state);
388 key_hash.hash(state);
389 }
390 WriteKey::Table { table_id } => {
391 2u8.hash(state);
392 table_id.hash(state);
393 }
394 }
395 }
396}
397
398impl PartialEq for WriteKey {
399 fn eq(&self, other: &Self) -> bool {
400 match (self, other) {
401 (
402 WriteKey::Row {
403 table_id: a,
404 row_id: b,
405 },
406 WriteKey::Row {
407 table_id: c,
408 row_id: d,
409 },
410 ) => a == c && b == d,
411 (
412 WriteKey::Unique {
413 table_id: a,
414 index_id: b,
415 key_hash: c,
416 },
417 WriteKey::Unique {
418 table_id: d,
419 index_id: e,
420 key_hash: f,
421 },
422 ) => a == d && b == e && c == f,
423 (WriteKey::Table { table_id: a }, WriteKey::Table { table_id: b }) => a == b,
424 _ => false,
425 }
426 }
427}
428
429impl Eq for WriteKey {}
430
431const CONFLICT_SHARDS: usize = 16;
432
433pub struct ConflictIndex {
437 shards: [parking_lot::Mutex<HashMap<WriteKey, u64>>; CONFLICT_SHARDS],
438 table_truncate_epochs: parking_lot::Mutex<HashMap<u64, u64>>,
439 table_write_epochs: parking_lot::Mutex<HashMap<u64, u64>>,
440 version: std::sync::atomic::AtomicU64,
444}
445
446impl ConflictIndex {
447 pub fn new() -> Self {
448 Self {
449 shards: std::array::from_fn(|_| parking_lot::Mutex::new(HashMap::new())),
450 table_truncate_epochs: parking_lot::Mutex::new(HashMap::new()),
451 table_write_epochs: parking_lot::Mutex::new(HashMap::new()),
452 version: std::sync::atomic::AtomicU64::new(0),
453 }
454 }
455
456 pub fn version(&self) -> u64 {
460 self.version.load(std::sync::atomic::Ordering::Acquire)
461 }
462
463 fn shard(&self, key: &WriteKey) -> &parking_lot::Mutex<HashMap<WriteKey, u64>> {
464 let mut h = std::collections::hash_map::DefaultHasher::new();
465 key.hash(&mut h);
466 let idx = (h.finish() as usize) & (CONFLICT_SHARDS - 1);
467 &self.shards[idx]
468 }
469
470 pub fn conflicts(&self, keys: &[WriteKey], read_epoch: Epoch) -> bool {
473 for k in keys {
474 let s = self.shard(k);
475 if let Some(&ce) = s.lock().get(k) {
476 if ce > read_epoch.0 {
477 return true;
478 }
479 }
480 }
481 let truncates = self.table_truncate_epochs.lock();
482 let writes = self.table_write_epochs.lock();
483 for k in keys {
484 match k {
485 WriteKey::Row { table_id, .. } | WriteKey::Unique { table_id, .. } => {
486 if truncates.get(table_id).is_some_and(|&ce| ce > read_epoch.0) {
487 return true;
488 }
489 }
490 WriteKey::Table { table_id } => {
491 if writes.get(table_id).is_some_and(|&ce| ce > read_epoch.0) {
492 return true;
493 }
494 }
495 }
496 }
497 false
498 }
499
500 pub fn record(&self, keys: &[WriteKey], commit_epoch: Epoch) {
502 for k in keys {
503 let s = self.shard(k);
504 s.lock().insert(k.clone(), commit_epoch.0);
505 }
506 let mut truncates = self.table_truncate_epochs.lock();
507 let mut writes = self.table_write_epochs.lock();
508 for k in keys {
509 match k {
510 WriteKey::Table { table_id } => {
511 truncates
512 .entry(*table_id)
513 .and_modify(|ce| *ce = (*ce).max(commit_epoch.0))
514 .or_insert(commit_epoch.0);
515 }
516 WriteKey::Row { table_id, .. } | WriteKey::Unique { table_id, .. } => {
517 writes
518 .entry(*table_id)
519 .and_modify(|ce| *ce = (*ce).max(commit_epoch.0))
520 .or_insert(commit_epoch.0);
521 }
522 }
523 }
524 self.version
525 .fetch_add(1, std::sync::atomic::Ordering::Release);
526 }
527
528 pub fn prune_below(&self, min_active: Epoch) {
531 for s in &self.shards {
532 s.lock().retain(|_, ce| *ce >= min_active.0);
533 }
534 self.table_truncate_epochs
535 .lock()
536 .retain(|_, ce| *ce >= min_active.0);
537 self.table_write_epochs
538 .lock()
539 .retain(|_, ce| *ce >= min_active.0);
540 }
541}
542
543impl Default for ConflictIndex {
544 fn default() -> Self {
545 Self::new()
546 }
547}
548
549pub struct GroupCommit {
559 inner: PlMutex<GroupState>,
560 cv: Condvar,
561}
562
563struct GroupState {
564 durable_seq: u64,
565 syncing: bool,
566 poisoned: bool,
567}
568
569impl GroupCommit {
570 pub fn new(durable_seq: u64) -> Self {
571 Self {
572 inner: PlMutex::new(GroupState {
573 durable_seq,
574 syncing: false,
575 poisoned: false,
576 }),
577 cv: Condvar::new(),
578 }
579 }
580
581 pub fn await_durable(&self, wal: &PlMutex<SharedWal>, commit_seq: u64) -> Result<()> {
587 let mut st = self.inner.lock();
588 loop {
589 if st.poisoned {
590 return Err(MongrelError::Other(
591 "database poisoned by fsync error".into(),
592 ));
593 }
594 if st.durable_seq >= commit_seq {
595 return Ok(());
596 }
597 if st.syncing {
598 self.cv.wait(&mut st);
600 continue;
601 }
602 st.syncing = true;
605 drop(st);
606 let res = wal.lock().group_sync();
607 st = self.inner.lock();
608 st.syncing = false;
609 match res {
610 Ok(durable) => {
611 if durable > st.durable_seq {
612 st.durable_seq = durable;
613 }
614 self.cv.notify_all();
615 }
618 Err(e) => {
619 st.poisoned = true;
620 self.cv.notify_all();
621 return Err(e);
622 }
623 }
624 }
625 }
626}
627
628pub struct ActiveTxns {
632 inner: parking_lot::Mutex<BTreeMap<u64, u64>>,
633}
634
635impl ActiveTxns {
636 pub fn new() -> Self {
637 Self {
638 inner: parking_lot::Mutex::new(BTreeMap::new()),
639 }
640 }
641
642 pub fn register(&self, read_epoch: Epoch) -> ActiveTxnGuard<'_> {
645 let mut g = self.inner.lock();
646 *g.entry(read_epoch.0).or_insert(0) += 1;
647 ActiveTxnGuard {
648 active: self,
649 epoch: read_epoch.0,
650 }
651 }
652
653 pub fn min_read_epoch(&self) -> u64 {
655 self.inner.lock().keys().next().copied().unwrap_or(u64::MAX)
656 }
657}
658
659impl Default for ActiveTxns {
660 fn default() -> Self {
661 Self::new()
662 }
663}
664
665pub struct ActiveTxnGuard<'a> {
667 active: &'a ActiveTxns,
668 epoch: u64,
669}
670
671impl Drop for ActiveTxnGuard<'_> {
672 fn drop(&mut self) {
673 let mut g = self.active.inner.lock();
674 if let Some(count) = g.get_mut(&self.epoch) {
675 *count -= 1;
676 if *count == 0 {
677 g.remove(&self.epoch);
678 }
679 }
680 }
681}
682
683#[cfg(test)]
684mod tests {
685 use super::*;
686
687 #[test]
688 fn conflict_index_first_committer_wins_and_prunes_safely() {
689 let ci = ConflictIndex::new();
690 let k = vec![WriteKey::Row {
691 table_id: 1,
692 row_id: 7,
693 }];
694 assert!(!ci.conflicts(&k, Epoch(5)));
695 ci.record(&k, Epoch(6));
696 assert!(ci.conflicts(&k, Epoch(5)));
697 assert!(!ci.conflicts(&k, Epoch(6)));
698 ci.prune_below(Epoch(7));
699 assert!(!ci.conflicts(&k, Epoch(5)));
700 }
701
702 #[test]
703 fn conflict_index_table_scope_conflicts_both_directions() {
704 let ci = ConflictIndex::new();
705 ci.record(&[WriteKey::Table { table_id: 1 }], Epoch(6));
706 assert!(ci.conflicts(
707 &[WriteKey::Row {
708 table_id: 1,
709 row_id: 7,
710 }],
711 Epoch(5)
712 ));
713 assert!(ci.conflicts(
714 &[WriteKey::Unique {
715 table_id: 1,
716 index_id: 0,
717 key_hash: 42,
718 }],
719 Epoch(5)
720 ));
721 assert!(!ci.conflicts(
722 &[WriteKey::Row {
723 table_id: 2,
724 row_id: 7,
725 }],
726 Epoch(5)
727 ));
728
729 let ci = ConflictIndex::new();
730 ci.record(
731 &[WriteKey::Row {
732 table_id: 1,
733 row_id: 7,
734 }],
735 Epoch(6),
736 );
737 assert!(ci.conflicts(&[WriteKey::Table { table_id: 1 }], Epoch(5)));
738 assert!(!ci.conflicts(&[WriteKey::Table { table_id: 2 }], Epoch(5)));
739 }
740
741 #[test]
742 fn writekey_eq_across_variants() {
743 let r1 = WriteKey::Row {
744 table_id: 1,
745 row_id: 2,
746 };
747 let r2 = WriteKey::Row {
748 table_id: 1,
749 row_id: 2,
750 };
751 let r3 = WriteKey::Row {
752 table_id: 1,
753 row_id: 3,
754 };
755 assert_eq!(r1, r2);
756 assert_ne!(r1, r3);
757
758 let u1 = WriteKey::Unique {
759 table_id: 1,
760 index_id: 0,
761 key_hash: 42,
762 };
763 let u2 = WriteKey::Unique {
764 table_id: 1,
765 index_id: 0,
766 key_hash: 42,
767 };
768 assert_eq!(u1, u2);
769 assert_ne!(r1, u1);
770
771 let t1 = WriteKey::Table { table_id: 5 };
772 let t2 = WriteKey::Table { table_id: 5 };
773 assert_eq!(t1, t2);
774 assert_ne!(t1, r1);
775 }
776
777 #[test]
778 fn active_txns_tracks_min_read_epoch() {
779 let at = ActiveTxns::new();
780 assert_eq!(at.min_read_epoch(), u64::MAX);
781 let g1 = at.register(Epoch(5));
782 assert_eq!(at.min_read_epoch(), 5);
783 let g2 = at.register(Epoch(3));
784 assert_eq!(at.min_read_epoch(), 3);
785 drop(g2);
786 assert_eq!(at.min_read_epoch(), 5);
787 drop(g1);
788 assert_eq!(at.min_read_epoch(), u64::MAX);
789 }
790
791 #[test]
792 fn active_txns_dedups_same_epoch() {
793 let at = ActiveTxns::new();
794 let g1 = at.register(Epoch(7));
795 let g2 = at.register(Epoch(7));
796 assert_eq!(at.min_read_epoch(), 7);
797 drop(g1);
798 assert_eq!(at.min_read_epoch(), 7);
799 drop(g2);
800 assert_eq!(at.min_read_epoch(), u64::MAX);
801 }
802}