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