1use std::collections::{BTreeMap, HashMap, VecDeque};
93use std::ops::Bound;
94use std::time::{Duration, Instant};
95
96use mongreldb_types::errors::ErrorCategory;
97use parking_lot::{Condvar, Mutex, MutexGuard};
98
99use crate::{ExecutionControl, MongrelError, RowId};
100
101const CANCELLATION_POLL_INTERVAL: Duration = Duration::from_millis(10);
105
106const DEFAULT_PRIORITY: u64 = 0;
111
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
118pub enum LockMode {
119 Shared,
122 Exclusive,
125}
126
127impl LockMode {
128 pub const fn compatible(self, other: Self) -> bool {
130 matches!((self, other), (Self::Shared, Self::Shared))
131 }
132}
133
134#[derive(Debug, Clone, PartialEq, Eq, Hash)]
140pub enum LockKey {
141 Row { table_id: u64, row_id: RowId },
143 Key { table_id: u64, key: Vec<u8> },
147 Range {
150 table_id: u64,
151 low: Bound<Vec<u8>>,
152 high: Bound<Vec<u8>>,
153 },
154 Barrier { name: String },
157}
158
159impl LockKey {
160 pub fn row(table_id: u64, row_id: RowId) -> Self {
162 Self::Row { table_id, row_id }
163 }
164
165 pub fn key(table_id: u64, key: impl Into<Vec<u8>>) -> Self {
167 Self::Key {
168 table_id,
169 key: key.into(),
170 }
171 }
172
173 pub fn range(table_id: u64, low: Bound<Vec<u8>>, high: Bound<Vec<u8>>) -> Self {
175 Self::Range {
176 table_id,
177 low,
178 high,
179 }
180 }
181
182 pub fn barrier(name: impl Into<String>) -> Self {
184 Self::Barrier { name: name.into() }
185 }
186
187 pub fn schema_barrier() -> Self {
191 Self::barrier("schema")
192 }
193
194 pub fn sequence_barrier(sequence: &str) -> Self {
197 Self::barrier(format!("sequence:{sequence}"))
198 }
199}
200
201#[derive(Debug, Clone)]
207pub struct LockRequest {
208 pub txn_id: u64,
211 pub mode: LockMode,
213 pub deadline: Option<Instant>,
216 pub control: ExecutionControl,
221 pub priority: Option<u64>,
225}
226
227impl LockRequest {
228 pub fn new(txn_id: u64, mode: LockMode, control: ExecutionControl) -> Self {
230 Self {
231 txn_id,
232 mode,
233 deadline: None,
234 control,
235 priority: None,
236 }
237 }
238
239 pub fn with_deadline(mut self, deadline: Instant) -> Self {
241 self.deadline = Some(deadline);
242 self
243 }
244
245 pub fn with_timeout(self, timeout: Duration) -> Self {
249 let now = Instant::now();
250 self.with_deadline(now.checked_add(timeout).unwrap_or(now))
251 }
252
253 pub fn with_priority(mut self, priority: u64) -> Self {
255 self.priority = Some(priority);
256 self
257 }
258}
259
260#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
266pub enum LockError {
267 #[error(
270 "deadlock: transaction {victim} was chosen as the deadlock victim (wait-for cycle {cycle}); retry the whole transaction"
271 )]
272 Deadlock {
273 victim: u64,
275 cycle: String,
277 },
278 #[error("lock wait deadline exceeded")]
280 DeadlineExceeded,
281 #[error("lock wait cancelled")]
283 Cancelled,
284 #[error("invalid lock request: {0}")]
287 InvalidRequest(String),
288}
289
290impl LockError {
291 pub fn category(&self) -> ErrorCategory {
295 match self {
296 Self::Deadlock { .. } => ErrorCategory::Deadlock,
297 Self::DeadlineExceeded => ErrorCategory::DeadlineExceeded,
298 Self::Cancelled => ErrorCategory::Cancelled,
299 Self::InvalidRequest(_) => ErrorCategory::ClusterVersionMismatch,
301 }
302 }
303}
304
305impl From<LockError> for MongrelError {
306 fn from(error: LockError) -> Self {
311 match error {
312 LockError::Deadlock { victim, cycle } => MongrelError::Deadlock { victim, cycle },
313 LockError::DeadlineExceeded => MongrelError::DeadlineExceeded,
314 LockError::Cancelled => MongrelError::Cancelled,
315 LockError::InvalidRequest(message) => MongrelError::InvalidArgument(message),
316 }
317 }
318}
319
320#[derive(Debug, Clone, Copy)]
322struct Holder {
323 txn_id: u64,
324 mode: LockMode,
325}
326
327#[derive(Debug)]
330struct Waiter {
331 txn_id: u64,
332 deadline: Option<Instant>,
333 control: ExecutionControl,
334 mode: LockMode,
335}
336
337#[derive(Debug, Default)]
339struct LockState {
340 holders: Vec<Holder>,
341 queue: VecDeque<Waiter>,
342}
343
344#[derive(Debug)]
347enum Fate {
348 Deadlock { victim: u64, cycle: String },
349 DeadlineExceeded,
350 Cancelled,
351}
352
353impl Fate {
354 fn into_error(self) -> LockError {
355 match self {
356 Self::Deadlock { victim, cycle } => LockError::Deadlock { victim, cycle },
357 Self::DeadlineExceeded => LockError::DeadlineExceeded,
358 Self::Cancelled => LockError::Cancelled,
359 }
360 }
361}
362
363#[derive(Debug, Default)]
364struct Inner {
365 locks: HashMap<LockKey, LockState>,
366 fates: HashMap<u64, Fate>,
370 priorities: HashMap<u64, u64>,
373}
374
375#[derive(Debug)]
378pub struct LockManager {
379 inner: Mutex<Inner>,
380 wake: Condvar,
381}
382
383impl Default for LockManager {
384 fn default() -> Self {
385 Self::new()
386 }
387}
388
389impl LockManager {
390 pub fn new() -> Self {
391 Self {
392 inner: Mutex::new(Inner::default()),
393 wake: Condvar::new(),
394 }
395 }
396
397 pub fn acquire(&self, key: LockKey, request: LockRequest) -> Result<(), LockError> {
401 let LockRequest {
402 txn_id,
403 mode,
404 deadline,
405 control,
406 priority,
407 } = request;
408
409 let mut inner = self.inner.lock();
410 check_request_live(&control, deadline)?;
411 if let Some(priority) = priority {
412 inner.priorities.insert(txn_id, priority);
413 }
414
415 {
416 let state = inner.locks.entry(key.clone()).or_default();
417 match state.holders.iter().position(|h| h.txn_id == txn_id) {
418 Some(position) => {
419 let held = state.holders[position].mode;
420 if held == LockMode::Exclusive || mode == LockMode::Shared {
421 return Ok(());
425 }
426 if state.holders.len() == 1 && state.queue.is_empty() {
427 state.holders[position].mode = LockMode::Exclusive;
429 return Ok(());
430 }
431 }
435 None => {
436 let blocked = !state.queue.is_empty()
437 || state.holders.iter().any(|h| !mode.compatible(h.mode));
438 if !blocked {
439 state.holders.push(Holder { txn_id, mode });
440 return Ok(());
441 }
442 }
443 }
444 if state.queue.iter().any(|w| w.txn_id == txn_id) {
445 return Err(LockError::InvalidRequest(format!(
446 "transaction {txn_id} already has a pending wait on {key:?}"
447 )));
448 }
449 state.queue.push_back(Waiter {
450 txn_id,
451 deadline,
452 control: control.clone(),
453 mode,
454 });
455 }
456
457 detect_deadlocks(&mut inner);
461 self.wake.notify_all();
462
463 let result = self.wait_for_grant(&mut inner, &key, txn_id, &control, deadline);
464 if result.is_err() {
465 process_key(&mut inner, &key);
467 detect_deadlocks(&mut inner);
468 self.wake.notify_all();
469 }
470 prune_if_idle(&mut inner, &key);
471 result
472 }
473
474 pub fn release(&self, txn_id: u64, key: &LockKey) {
477 let mut inner = self.inner.lock();
478 if inner
479 .locks
480 .get(key)
481 .is_some_and(|s| s.holders.iter().any(|h| h.txn_id == txn_id))
482 {
483 if let Some(state) = inner.locks.get_mut(key) {
484 state.holders.retain(|h| h.txn_id != txn_id);
485 }
486 process_key(&mut inner, key);
487 detect_deadlocks(&mut inner);
489 prune_if_idle(&mut inner, key);
490 }
491 drop(inner);
492 self.wake.notify_all();
493 }
494
495 pub fn release_all(&self, txn_id: u64) {
498 let mut inner = self.inner.lock();
499 inner.fates.remove(&txn_id);
500 inner.priorities.remove(&txn_id);
501 let mut affected = Vec::new();
502 for (key, state) in inner.locks.iter_mut() {
503 let touched = state.holders.iter().any(|h| h.txn_id == txn_id)
504 || state.queue.iter().any(|w| w.txn_id == txn_id);
505 if touched {
506 state.holders.retain(|h| h.txn_id != txn_id);
507 state.queue.retain(|w| w.txn_id != txn_id);
508 affected.push(key.clone());
509 }
510 }
511 for key in &affected {
512 process_key(&mut inner, key);
513 }
514 detect_deadlocks(&mut inner);
515 inner
516 .locks
517 .retain(|_, state| !(state.holders.is_empty() && state.queue.is_empty()));
518 drop(inner);
519 self.wake.notify_all();
520 }
521
522 pub fn holds(&self, txn_id: u64, key: &LockKey) -> bool {
524 self.inner
525 .lock()
526 .locks
527 .get(key)
528 .is_some_and(|state| state.holders.iter().any(|h| h.txn_id == txn_id))
529 }
530
531 #[cfg(test)]
534 fn queued_waiters(&self, key: &LockKey) -> usize {
535 self.inner
536 .lock()
537 .locks
538 .get(key)
539 .map_or(0, |state| state.queue.len())
540 }
541
542 fn wait_for_grant(
546 &self,
547 inner: &mut MutexGuard<'_, Inner>,
548 key: &LockKey,
549 txn_id: u64,
550 control: &ExecutionControl,
551 deadline: Option<Instant>,
552 ) -> Result<(), LockError> {
553 loop {
554 if let Some(fate) = inner.fates.remove(&txn_id) {
555 return Err(fate.into_error());
556 }
557 let (queued, held) = match inner.locks.get(key) {
558 Some(state) => (
559 state.queue.iter().any(|w| w.txn_id == txn_id),
560 state.holders.iter().any(|h| h.txn_id == txn_id),
561 ),
562 None => (false, false),
563 };
564 if !queued {
565 if held {
566 return Ok(());
567 }
568 return Err(LockError::Cancelled);
572 }
573 if let Err(error) = check_request_live(control, deadline) {
574 remove_waiter(inner, key, txn_id);
575 return Err(error);
576 }
577 let wake_at = deadline
580 .unwrap_or_else(|| Instant::now() + CANCELLATION_POLL_INTERVAL)
581 .min(Instant::now() + CANCELLATION_POLL_INTERVAL);
582 let _ = self.wake.wait_until(inner, wake_at);
583 }
584 }
585}
586
587fn check_request_live(
592 control: &ExecutionControl,
593 deadline: Option<Instant>,
594) -> Result<(), LockError> {
595 match control.checkpoint() {
596 Ok(()) => {}
597 Err(MongrelError::DeadlineExceeded) => return Err(LockError::DeadlineExceeded),
598 Err(_) => return Err(LockError::Cancelled),
599 }
600 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
601 return Err(LockError::DeadlineExceeded);
602 }
603 Ok(())
604}
605
606fn remove_waiter(inner: &mut Inner, key: &LockKey, txn_id: u64) {
607 if let Some(state) = inner.locks.get_mut(key) {
608 state.queue.retain(|w| w.txn_id != txn_id);
609 }
610}
611
612fn prune_if_idle(inner: &mut Inner, key: &LockKey) {
613 if inner
614 .locks
615 .get(key)
616 .is_some_and(|state| state.holders.is_empty() && state.queue.is_empty())
617 {
618 inner.locks.remove(key);
619 }
620}
621
622fn process_key(inner: &mut Inner, key: &LockKey) {
627 let now = Instant::now();
628 let dead: Vec<u64> = match inner.locks.get(key) {
629 Some(state) => state
630 .queue
631 .iter()
632 .filter(|w| w.control.is_cancelled() || w.deadline.is_some_and(|d| now >= d))
633 .map(|w| w.txn_id)
634 .collect(),
635 None => return,
636 };
637 for txn_id in dead {
638 if let Some(state) = inner.locks.get_mut(key) {
639 if let Some(position) = state.queue.iter().position(|w| w.txn_id == txn_id) {
640 let waiter = state.queue.remove(position).expect("position found above");
641 let fate = match waiter.control.checkpoint() {
642 Err(MongrelError::DeadlineExceeded) => Fate::DeadlineExceeded,
643 Err(_) => Fate::Cancelled,
644 Ok(()) => Fate::DeadlineExceeded,
646 };
647 inner.fates.entry(txn_id).or_insert(fate);
648 }
649 }
650 }
651
652 loop {
653 let Some(state) = inner.locks.get_mut(key) else {
654 return;
655 };
656 let Some(front) = state.queue.front() else {
657 return;
658 };
659 let grantable = state
662 .holders
663 .iter()
664 .all(|h| h.txn_id == front.txn_id || front.mode.compatible(h.mode));
665 if !grantable {
666 return;
667 }
668 let waiter = state.queue.pop_front().expect("front checked above");
669 match state.holders.iter_mut().find(|h| h.txn_id == waiter.txn_id) {
670 Some(holder) => holder.mode = LockMode::Exclusive,
671 None => state.holders.push(Holder {
672 txn_id: waiter.txn_id,
673 mode: waiter.mode,
674 }),
675 }
676 }
677}
678
679fn build_wait_for_graph(inner: &Inner) -> BTreeMap<u64, Vec<u64>> {
686 let mut graph: BTreeMap<u64, Vec<u64>> = BTreeMap::new();
687 for state in inner.locks.values() {
688 for (index, waiter) in state.queue.iter().enumerate() {
689 let edges = graph.entry(waiter.txn_id).or_default();
690 for holder in &state.holders {
691 if holder.txn_id != waiter.txn_id && !waiter.mode.compatible(holder.mode) {
692 edges.push(holder.txn_id);
693 }
694 }
695 for ahead in state.queue.iter().take(index) {
696 if ahead.txn_id != waiter.txn_id && !waiter.mode.compatible(ahead.mode) {
697 edges.push(ahead.txn_id);
698 }
699 }
700 }
701 }
702 for edges in graph.values_mut() {
703 edges.sort_unstable();
704 edges.dedup();
705 }
706 graph
707}
708
709fn find_cycle(graph: &BTreeMap<u64, Vec<u64>>) -> Option<Vec<u64>> {
714 #[derive(Clone, Copy, PartialEq, Eq)]
715 enum Mark {
716 Visiting,
717 Done,
718 }
719
720 let mut marks: HashMap<u64, Mark> = HashMap::new();
721 let mut stack: Vec<u64> = Vec::new();
722 for &start in graph.keys() {
723 if marks.contains_key(&start) {
724 continue;
725 }
726 marks.insert(start, Mark::Visiting);
727 stack.push(start);
728 let mut frames: Vec<(u64, usize)> = vec![(start, 0)];
729 while let Some((node, index)) = frames.last().copied() {
730 let edges = graph.get(&node).map(Vec::as_slice).unwrap_or(&[]);
731 if index < edges.len() {
732 frames.last_mut().expect("last copied above").1 += 1;
733 let next = edges[index];
734 match marks.get(&next) {
735 None => {
736 marks.insert(next, Mark::Visiting);
737 stack.push(next);
738 frames.push((next, 0));
739 }
740 Some(Mark::Visiting) => {
741 let position = stack
742 .iter()
743 .position(|member| *member == next)
744 .expect("a visiting node is on the stack");
745 return Some(stack[position..].to_vec());
746 }
747 Some(Mark::Done) => {}
748 }
749 } else {
750 marks.insert(node, Mark::Done);
751 stack.pop();
752 frames.pop();
753 }
754 }
755 }
756 None
757}
758
759fn choose_victim(cycle: &[u64], priorities: &HashMap<u64, u64>) -> u64 {
764 *cycle
765 .iter()
766 .min_by(|a, b| {
767 let priority_a = priorities.get(a).copied().unwrap_or(DEFAULT_PRIORITY);
768 let priority_b = priorities.get(b).copied().unwrap_or(DEFAULT_PRIORITY);
769 priority_a.cmp(&priority_b).then_with(|| b.cmp(a))
770 })
771 .expect("a deadlock cycle is never empty")
772}
773
774fn detect_deadlocks(inner: &mut Inner) {
778 loop {
779 let graph = build_wait_for_graph(inner);
780 let Some(cycle) = find_cycle(&graph) else {
781 return;
782 };
783 let victim = choose_victim(&cycle, &inner.priorities);
784 let victim_key = inner.locks.iter().find_map(|(key, state)| {
787 state
788 .queue
789 .iter()
790 .any(|w| w.txn_id == victim)
791 .then(|| key.clone())
792 });
793 let Some(victim_key) = victim_key else {
794 return;
797 };
798 if let Some(state) = inner.locks.get_mut(&victim_key) {
799 state.queue.retain(|w| w.txn_id != victim);
800 }
801 let cycle = cycle
802 .iter()
803 .map(u64::to_string)
804 .collect::<Vec<_>>()
805 .join(" → ");
806 inner.fates.insert(victim, Fate::Deadlock { victim, cycle });
807 }
808}
809
810#[cfg(test)]
811mod tests {
812 use super::*;
813 use crate::CancellationReason;
814 use std::sync::mpsc;
815 use std::sync::Arc;
816 use std::thread;
817
818 use LockMode::{Exclusive, Shared};
819
820 fn manager() -> Arc<LockManager> {
821 Arc::new(LockManager::new())
822 }
823
824 fn request(txn_id: u64, mode: LockMode) -> LockRequest {
825 LockRequest::new(txn_id, mode, ExecutionControl::new(None))
826 }
827
828 fn timed_request(txn_id: u64, mode: LockMode, timeout: Duration) -> LockRequest {
829 request(txn_id, mode).with_timeout(timeout)
830 }
831
832 fn row_key(n: u64) -> LockKey {
833 LockKey::row(1, RowId(n))
834 }
835
836 fn spawn_acquire(
838 manager: &Arc<LockManager>,
839 key: LockKey,
840 request: LockRequest,
841 ) -> mpsc::Receiver<Result<(), LockError>> {
842 let (sender, receiver) = mpsc::channel();
843 let manager = Arc::clone(manager);
844 thread::spawn(move || {
845 let result = manager.acquire(key, request);
846 sender.send(result).expect("receiver alive");
847 });
848 receiver
849 }
850
851 fn wait_for_queue(manager: &LockManager, key: &LockKey, n: usize) {
855 for _ in 0..200 {
856 if manager.queued_waiters(key) >= n {
857 return;
858 }
859 thread::sleep(Duration::from_millis(10));
860 }
861 panic!("queue did not reach {n} waiters on {key:?}");
862 }
863
864 #[test]
865 fn acquire_and_release_round_trip() {
866 let manager = LockManager::new();
867 let key = row_key(1);
868 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
869 assert!(manager.holds(1, &key));
870 manager.release(1, &key);
871 assert!(!manager.holds(1, &key));
872
873 manager.acquire(key.clone(), request(2, Shared)).unwrap();
874 assert!(manager.holds(2, &key));
875 manager.release_all(2);
876 assert!(!manager.holds(2, &key));
877 }
878
879 #[test]
880 fn mode_compatibility_matrix() {
881 assert!(Shared.compatible(Shared));
882 assert!(!Shared.compatible(Exclusive));
883 assert!(!Exclusive.compatible(Shared));
884 assert!(!Exclusive.compatible(Exclusive));
885 }
886
887 #[test]
888 fn shared_locks_coexist() {
889 let manager = LockManager::new();
890 let key = row_key(7);
891 manager.acquire(key.clone(), request(1, Shared)).unwrap();
892 manager.acquire(key.clone(), request(2, Shared)).unwrap();
893 assert!(manager.holds(1, &key));
894 assert!(manager.holds(2, &key));
895 }
896
897 #[test]
898 fn exclusive_blocks_shared_and_exclusive_until_release() {
899 let manager = LockManager::new();
900 let key = row_key(8);
901 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
902
903 let shared = manager.acquire(
904 key.clone(),
905 timed_request(2, Shared, Duration::from_millis(40)),
906 );
907 assert_eq!(shared, Err(LockError::DeadlineExceeded));
908 let exclusive = manager.acquire(
909 key.clone(),
910 timed_request(3, Exclusive, Duration::from_millis(40)),
911 );
912 assert_eq!(exclusive, Err(LockError::DeadlineExceeded));
913
914 manager.release(1, &key);
915 manager.acquire(key.clone(), request(2, Shared)).unwrap();
916 let still_blocked = manager.acquire(
917 key.clone(),
918 timed_request(3, Exclusive, Duration::from_millis(40)),
919 );
920 assert_eq!(still_blocked, Err(LockError::DeadlineExceeded));
921 }
922
923 #[test]
924 fn wait_respects_the_lock_deadline() {
925 let manager = LockManager::new();
926 let key = row_key(9);
927 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
928
929 let started = Instant::now();
930 let result = manager.acquire(key, timed_request(2, Exclusive, Duration::from_millis(80)));
931 let elapsed = started.elapsed();
932 assert_eq!(result, Err(LockError::DeadlineExceeded));
933 assert_eq!(
934 result.unwrap_err().category(),
935 ErrorCategory::DeadlineExceeded
936 );
937 assert!(
938 elapsed >= Duration::from_millis(70),
939 "returned early: {elapsed:?}"
940 );
941 assert!(
942 elapsed < Duration::from_secs(5),
943 "returned late: {elapsed:?}"
944 );
945 }
946
947 #[test]
948 fn expired_deadline_fails_fast_even_when_unlocked() {
949 let manager = LockManager::new();
950 let result = manager.acquire(row_key(10), timed_request(1, Exclusive, Duration::ZERO));
951 assert_eq!(result, Err(LockError::DeadlineExceeded));
952 }
953
954 #[test]
955 fn cancellation_wakes_the_waiter_and_clears_its_queue_entry() {
956 let manager = manager();
957 let key = row_key(11);
958 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
959
960 let control = ExecutionControl::new(None);
961 let cancelled_rx = spawn_acquire(
962 &manager,
963 key.clone(),
964 LockRequest::new(2, Exclusive, control.clone()),
965 );
966 wait_for_queue(&manager, &key, 1);
967 control.cancel(CancellationReason::ClientRequest);
968 let result = cancelled_rx.recv_timeout(Duration::from_secs(5)).unwrap();
969 assert_eq!(result, Err(LockError::Cancelled));
970 assert!(!manager.holds(2, &key));
971
972 let next_rx = spawn_acquire(&manager, key.clone(), request(3, Exclusive));
975 wait_for_queue(&manager, &key, 1);
976 manager.release(1, &key);
977 assert_eq!(
978 next_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
979 Ok(())
980 );
981 }
982
983 #[test]
984 fn cancelled_control_fails_fast_even_when_unlocked() {
985 let manager = LockManager::new();
986 let control = ExecutionControl::new(None);
987 control.cancel(CancellationReason::ClientRequest);
988 let result = manager.acquire(row_key(12), LockRequest::new(1, Exclusive, control));
989 assert_eq!(result, Err(LockError::Cancelled));
990 }
991
992 #[test]
993 fn control_deadline_maps_to_deadline_exceeded() {
994 let manager = LockManager::new();
995 let key = row_key(13);
996 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
997 let control = ExecutionControl::with_timeout(Duration::from_millis(50));
998 let result = manager.acquire(key, LockRequest::new(2, Exclusive, control));
999 assert_eq!(result, Err(LockError::DeadlineExceeded));
1000 }
1001
1002 #[test]
1003 fn fifo_grant_order_blocks_barging_readers() {
1004 let manager = manager();
1005 let key = row_key(14);
1006 manager.acquire(key.clone(), request(1, Shared)).unwrap();
1007
1008 let writer_rx = spawn_acquire(&manager, key.clone(), request(2, Exclusive));
1010 wait_for_queue(&manager, &key, 1);
1011 let reader_rx = spawn_acquire(&manager, key.clone(), request(3, Shared));
1015 wait_for_queue(&manager, &key, 2);
1016
1017 manager.release(1, &key);
1018 assert_eq!(
1019 writer_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
1020 Ok(())
1021 );
1022 assert!(
1023 reader_rx.recv_timeout(Duration::from_millis(150)).is_err(),
1024 "reader barged ahead of the queued writer"
1025 );
1026 manager.release(2, &key);
1027 assert_eq!(
1028 reader_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
1029 Ok(())
1030 );
1031 }
1032
1033 #[test]
1034 fn release_all_grants_blocked_waiters() {
1035 let manager = manager();
1036 let key = row_key(24);
1037 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
1038 let waiter_rx = spawn_acquire(&manager, key.clone(), request(2, Exclusive));
1039 wait_for_queue(&manager, &key, 1);
1040 manager.release_all(1);
1041 assert_eq!(
1042 waiter_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
1043 Ok(())
1044 );
1045 assert!(manager.holds(2, &key));
1046 }
1047
1048 #[test]
1049 fn ab_ba_deadlock_kills_requester_when_it_is_youngest() {
1050 let manager = manager();
1051 let key_a = row_key(20);
1052 let key_b = row_key(21);
1053 manager
1054 .acquire(key_a.clone(), request(1, Exclusive))
1055 .unwrap();
1056 manager
1057 .acquire(key_b.clone(), request(2, Exclusive))
1058 .unwrap();
1059
1060 let t1_rx = spawn_acquire(&manager, key_b.clone(), request(1, Exclusive));
1062 wait_for_queue(&manager, &key_b, 1);
1063 let error = manager
1066 .acquire(key_a.clone(), request(2, Exclusive))
1067 .unwrap_err();
1068 match &error {
1069 LockError::Deadlock { victim, cycle } => {
1070 assert_eq!(*victim, 2, "youngest transaction must be the victim");
1071 assert!(
1072 cycle.contains('1') && cycle.contains('2'),
1073 "cycle names both: {cycle}"
1074 );
1075 }
1076 other => panic!("expected t2 to be the deadlock victim, got {other:?}"),
1077 }
1078 assert_eq!(error.category(), ErrorCategory::Deadlock);
1079
1080 manager.release_all(2);
1082 assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1083 assert!(manager.holds(1, &key_b));
1084 }
1085
1086 #[test]
1087 fn ab_ba_deadlock_kills_youngest_even_when_it_is_not_the_requester() {
1088 let manager = manager();
1091 let key_a = row_key(22);
1092 let key_b = row_key(23);
1093 manager
1094 .acquire(key_a.clone(), request(100, Exclusive))
1095 .unwrap();
1096 manager
1097 .acquire(key_b.clone(), request(200, Exclusive))
1098 .unwrap();
1099
1100 let t200 = {
1102 let manager = Arc::clone(&manager);
1103 let key_a = key_a.clone();
1104 thread::spawn(move || {
1105 let result = manager.acquire(key_a, request(200, Exclusive));
1106 if result.is_err() {
1107 manager.release_all(200);
1108 }
1109 result
1110 })
1111 };
1112 wait_for_queue(&manager, &key_a, 1);
1113
1114 manager
1117 .acquire(key_b.clone(), request(100, Exclusive))
1118 .unwrap();
1119 assert!(manager.holds(100, &key_b));
1120
1121 let victim = t200.join().unwrap();
1122 assert!(
1123 matches!(victim, Err(LockError::Deadlock { victim: 200, .. })),
1124 "youngest transaction must be the victim: {victim:?}"
1125 );
1126 }
1127
1128 #[test]
1129 fn three_transaction_cycle_kills_youngest() {
1130 let manager = manager();
1131 let key_a = row_key(30);
1132 let key_b = row_key(31);
1133 let key_c = row_key(32);
1134 manager
1135 .acquire(key_a.clone(), request(1, Exclusive))
1136 .unwrap();
1137 manager
1138 .acquire(key_b.clone(), request(2, Exclusive))
1139 .unwrap();
1140 manager
1141 .acquire(key_c.clone(), request(3, Exclusive))
1142 .unwrap();
1143
1144 let t1_rx = spawn_acquire(&manager, key_b.clone(), request(1, Exclusive));
1146 wait_for_queue(&manager, &key_b, 1);
1147 let t2_rx = spawn_acquire(&manager, key_c.clone(), request(2, Exclusive));
1148 wait_for_queue(&manager, &key_c, 1);
1149 let error = manager
1151 .acquire(key_a.clone(), request(3, Exclusive))
1152 .unwrap_err();
1153 assert!(
1154 matches!(error, LockError::Deadlock { victim: 3, .. }),
1155 "youngest of the 3-cycle must die: {error:?}"
1156 );
1157
1158 manager.release_all(3);
1160 assert_eq!(t2_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1161 manager.release_all(2);
1162 assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1163 }
1164
1165 #[test]
1166 fn lowest_priority_dies_before_age_is_considered() {
1167 let manager = manager();
1168 let key_a = row_key(40);
1169 let key_b = row_key(41);
1170 let key_c = row_key(42);
1171 manager
1172 .acquire(key_a.clone(), request(1, Exclusive))
1173 .unwrap();
1174 manager
1175 .acquire(key_b.clone(), request(2, Exclusive))
1176 .unwrap();
1177 manager
1178 .acquire(key_c.clone(), request(3, Exclusive))
1179 .unwrap();
1180
1181 let t1 = {
1184 let manager = Arc::clone(&manager);
1185 let key_b = key_b.clone();
1186 thread::spawn(move || {
1187 let result = manager.acquire(key_b, request(1, Exclusive).with_priority(10));
1188 if result.is_err() {
1189 manager.release_all(1);
1190 }
1191 result
1192 })
1193 };
1194 wait_for_queue(&manager, &key_b, 1);
1195 let t2_rx = spawn_acquire(
1196 &manager,
1197 key_c.clone(),
1198 request(2, Exclusive).with_priority(20),
1199 );
1200 wait_for_queue(&manager, &key_c, 1);
1201
1202 manager
1204 .acquire(key_a.clone(), request(3, Exclusive).with_priority(30))
1205 .unwrap();
1206 assert!(manager.holds(3, &key_a));
1207
1208 let victim = t1.join().unwrap();
1209 assert!(
1210 matches!(victim, Err(LockError::Deadlock { victim: 1, .. })),
1211 "lowest priority must die: {victim:?}"
1212 );
1213
1214 manager.release_all(3);
1215 assert_eq!(t2_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1216 }
1217
1218 #[test]
1219 fn equal_priorities_fall_back_to_youngest() {
1220 let manager = manager();
1221 let key_a = row_key(50);
1222 let key_b = row_key(51);
1223 manager
1224 .acquire(key_a.clone(), request(1, Exclusive).with_priority(5))
1225 .unwrap();
1226 manager
1227 .acquire(key_b.clone(), request(2, Exclusive).with_priority(5))
1228 .unwrap();
1229
1230 let t1_rx = spawn_acquire(
1231 &manager,
1232 key_b.clone(),
1233 request(1, Exclusive).with_priority(5),
1234 );
1235 wait_for_queue(&manager, &key_b, 1);
1236 let error = manager
1237 .acquire(key_a.clone(), request(2, Exclusive).with_priority(5))
1238 .unwrap_err();
1239 assert!(
1240 matches!(error, LockError::Deadlock { victim: 2, .. }),
1241 "equal priorities must fall back to youngest: {error:?}"
1242 );
1243 manager.release_all(2);
1244 assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1245 }
1246
1247 #[test]
1248 fn same_transaction_requests_never_self_deadlock() {
1249 let manager = LockManager::new();
1250 let key = row_key(60);
1251 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
1254 manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
1255 manager.acquire(key.clone(), request(1, Shared)).unwrap();
1256 assert!(manager.holds(1, &key));
1257
1258 let other = row_key(61);
1261 manager.acquire(other.clone(), request(2, Shared)).unwrap();
1262 manager
1263 .acquire(other.clone(), request(2, Exclusive))
1264 .unwrap();
1265 let blocked = manager.acquire(
1266 other,
1267 timed_request(3, Exclusive, Duration::from_millis(40)),
1268 );
1269 assert_eq!(blocked, Err(LockError::DeadlineExceeded));
1270 }
1271
1272 #[test]
1273 fn detector_kills_a_self_loop_if_one_is_ever_constructed() {
1274 let graph = BTreeMap::from([(7_u64, vec![7_u64])]);
1278 let cycle = find_cycle(&graph).expect("a self loop is a cycle");
1279 assert_eq!(cycle, vec![7]);
1280 assert_eq!(choose_victim(&cycle, &HashMap::new()), 7);
1281 }
1282
1283 #[test]
1284 fn shared_to_exclusive_upgrade_deadlock_kills_youngest() {
1285 let manager = manager();
1286 let key = row_key(70);
1287 manager.acquire(key.clone(), request(1, Shared)).unwrap();
1288 manager.acquire(key.clone(), request(2, Shared)).unwrap();
1289
1290 let t1_rx = spawn_acquire(&manager, key.clone(), request(1, Exclusive));
1292 wait_for_queue(&manager, &key, 1);
1293 let error = manager
1296 .acquire(key.clone(), request(2, Exclusive))
1297 .unwrap_err();
1298 assert!(
1299 matches!(error, LockError::Deadlock { victim: 2, .. }),
1300 "youngest upgrader must die: {error:?}"
1301 );
1302
1303 manager.release_all(2);
1305 assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
1306 let blocked = manager.acquire(
1307 key.clone(),
1308 timed_request(3, Shared, Duration::from_millis(40)),
1309 );
1310 assert_eq!(
1311 blocked,
1312 Err(LockError::DeadlineExceeded),
1313 "t1 must now hold Exclusive"
1314 );
1315 }
1316
1317 #[test]
1318 fn sequence_barrier_is_exclusive_across_transactions() {
1319 let manager = LockManager::new();
1320 let sequence = LockKey::sequence_barrier("order_id");
1321 manager
1322 .acquire(sequence.clone(), request(1, Exclusive))
1323 .unwrap();
1324
1325 let blocked = manager.acquire(
1327 sequence.clone(),
1328 timed_request(2, Exclusive, Duration::from_millis(40)),
1329 );
1330 assert_eq!(blocked, Err(LockError::DeadlineExceeded));
1331 manager
1333 .acquire(
1334 LockKey::sequence_barrier("shipment_id"),
1335 request(2, Exclusive),
1336 )
1337 .unwrap();
1338
1339 manager.release_all(1);
1340 manager.acquire(sequence, request(2, Exclusive)).unwrap();
1341 }
1342
1343 #[test]
1344 fn ddl_barrier_excludes_dml_and_concurrent_ddl() {
1345 let manager = LockManager::new();
1346 let barrier = LockKey::schema_barrier();
1347 manager
1349 .acquire(barrier.clone(), request(1, Exclusive))
1350 .unwrap();
1351 let dml = manager.acquire(
1353 barrier.clone(),
1354 timed_request(2, Shared, Duration::from_millis(40)),
1355 );
1356 assert_eq!(dml, Err(LockError::DeadlineExceeded));
1357 let ddl = manager.acquire(
1359 barrier.clone(),
1360 timed_request(3, Exclusive, Duration::from_millis(40)),
1361 );
1362 assert_eq!(ddl, Err(LockError::DeadlineExceeded));
1363
1364 manager.release_all(1);
1365 manager
1367 .acquire(barrier.clone(), request(2, Shared))
1368 .unwrap();
1369 manager.acquire(barrier, request(3, Shared)).unwrap();
1370 }
1371
1372 #[test]
1373 fn key_families_do_not_false_share() {
1374 let manager = LockManager::new();
1375 manager
1376 .acquire(LockKey::row(1, RowId(5)), request(1, Exclusive))
1377 .unwrap();
1378 manager
1380 .acquire(
1381 LockKey::key(1, 5_u64.to_be_bytes().to_vec()),
1382 request(2, Exclusive),
1383 )
1384 .unwrap();
1385 manager
1386 .acquire(
1387 LockKey::range(
1388 1,
1389 Bound::Included(1_u64.to_be_bytes().to_vec()),
1390 Bound::Excluded(6_u64.to_be_bytes().to_vec()),
1391 ),
1392 request(3, Exclusive),
1393 )
1394 .unwrap();
1395 }
1396
1397 #[test]
1398 fn lock_errors_bridge_to_mongrel_error_and_the_taxonomy() {
1399 let deadlock = LockError::Deadlock {
1400 victim: 9,
1401 cycle: "9 → 4 → 9".to_string(),
1402 };
1403 assert_eq!(deadlock.category(), ErrorCategory::Deadlock);
1404 let mapped = MongrelError::from(deadlock);
1405 assert!(
1406 matches!(
1407 &mapped,
1408 MongrelError::Deadlock { victim: 9, cycle } if cycle == "9 → 4 → 9"
1409 ),
1410 "deadlock maps to the dedicated variant, victim and cycle preserved: {mapped:?}"
1411 );
1412 assert_eq!(mapped.category(), ErrorCategory::Deadlock);
1413
1414 assert!(matches!(
1415 MongrelError::from(LockError::DeadlineExceeded),
1416 MongrelError::DeadlineExceeded
1417 ));
1418 assert!(matches!(
1419 MongrelError::from(LockError::Cancelled),
1420 MongrelError::Cancelled
1421 ));
1422 assert_eq!(LockError::Cancelled.category(), ErrorCategory::Cancelled);
1423 assert!(matches!(
1424 MongrelError::from(LockError::InvalidRequest("x".to_string())),
1425 MongrelError::InvalidArgument(_)
1426 ));
1427 }
1428
1429 #[test]
1430 fn wait_for_graph_cycle_search_is_deterministic() {
1431 let graph = BTreeMap::from([(1, vec![2]), (2, vec![3])]);
1433 assert_eq!(find_cycle(&graph), None);
1434
1435 let graph = BTreeMap::from([(1, vec![2]), (2, vec![3]), (3, vec![2])]);
1437 let cycle = find_cycle(&graph).expect("cycle");
1438 assert_eq!(cycle, vec![2, 3]);
1439
1440 assert_eq!(choose_victim(&cycle, &HashMap::new()), 3);
1443 assert_eq!(choose_victim(&cycle, &HashMap::from([(2, 5), (3, 1)])), 3);
1444 assert_eq!(choose_victim(&cycle, &HashMap::from([(2, 1), (3, 5)])), 2);
1445 }
1446}