use std::collections::{BTreeMap, HashMap, VecDeque};
use std::ops::Bound;
use std::time::{Duration, Instant};
use mongreldb_types::errors::ErrorCategory;
use parking_lot::{Condvar, Mutex, MutexGuard};
use crate::{ExecutionControl, MongrelError, RowId};
const CANCELLATION_POLL_INTERVAL: Duration = Duration::from_millis(10);
const DEFAULT_PRIORITY: u64 = 0;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LockMode {
Shared,
Exclusive,
}
impl LockMode {
pub const fn compatible(self, other: Self) -> bool {
matches!((self, other), (Self::Shared, Self::Shared))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum LockKey {
Row { table_id: u64, row_id: RowId },
Key { table_id: u64, key: Vec<u8> },
Range {
table_id: u64,
low: Bound<Vec<u8>>,
high: Bound<Vec<u8>>,
},
Barrier { name: String },
}
impl LockKey {
pub fn row(table_id: u64, row_id: RowId) -> Self {
Self::Row { table_id, row_id }
}
pub fn key(table_id: u64, key: impl Into<Vec<u8>>) -> Self {
Self::Key {
table_id,
key: key.into(),
}
}
pub fn range(table_id: u64, low: Bound<Vec<u8>>, high: Bound<Vec<u8>>) -> Self {
Self::Range {
table_id,
low,
high,
}
}
pub fn barrier(name: impl Into<String>) -> Self {
Self::Barrier { name: name.into() }
}
pub fn schema_barrier() -> Self {
Self::barrier("schema")
}
pub fn sequence_barrier(sequence: &str) -> Self {
Self::barrier(format!("sequence:{sequence}"))
}
}
#[derive(Debug, Clone)]
pub struct LockRequest {
pub txn_id: u64,
pub mode: LockMode,
pub deadline: Option<Instant>,
pub control: ExecutionControl,
pub priority: Option<u64>,
}
impl LockRequest {
pub fn new(txn_id: u64, mode: LockMode, control: ExecutionControl) -> Self {
Self {
txn_id,
mode,
deadline: None,
control,
priority: None,
}
}
pub fn with_deadline(mut self, deadline: Instant) -> Self {
self.deadline = Some(deadline);
self
}
pub fn with_timeout(self, timeout: Duration) -> Self {
let now = Instant::now();
self.with_deadline(now.checked_add(timeout).unwrap_or(now))
}
pub fn with_priority(mut self, priority: u64) -> Self {
self.priority = Some(priority);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum LockError {
#[error(
"deadlock: transaction {victim} was chosen as the deadlock victim (wait-for cycle {cycle}); retry the whole transaction"
)]
Deadlock {
victim: u64,
cycle: String,
},
#[error("lock wait deadline exceeded")]
DeadlineExceeded,
#[error("lock wait cancelled")]
Cancelled,
#[error("invalid lock request: {0}")]
InvalidRequest(String),
}
impl LockError {
pub fn category(&self) -> ErrorCategory {
match self {
Self::Deadlock { .. } => ErrorCategory::Deadlock,
Self::DeadlineExceeded => ErrorCategory::DeadlineExceeded,
Self::Cancelled => ErrorCategory::Cancelled,
Self::InvalidRequest(_) => ErrorCategory::ClusterVersionMismatch,
}
}
}
impl From<LockError> for MongrelError {
fn from(error: LockError) -> Self {
match error {
LockError::Deadlock { victim, cycle } => MongrelError::Deadlock { victim, cycle },
LockError::DeadlineExceeded => MongrelError::DeadlineExceeded,
LockError::Cancelled => MongrelError::Cancelled,
LockError::InvalidRequest(message) => MongrelError::InvalidArgument(message),
}
}
}
#[derive(Debug, Clone, Copy)]
struct Holder {
txn_id: u64,
mode: LockMode,
}
#[derive(Debug)]
struct Waiter {
txn_id: u64,
deadline: Option<Instant>,
control: ExecutionControl,
mode: LockMode,
}
#[derive(Debug, Default)]
struct LockState {
holders: Vec<Holder>,
queue: VecDeque<Waiter>,
}
#[derive(Debug)]
enum Fate {
Deadlock { victim: u64, cycle: String },
DeadlineExceeded,
Cancelled,
}
impl Fate {
fn into_error(self) -> LockError {
match self {
Self::Deadlock { victim, cycle } => LockError::Deadlock { victim, cycle },
Self::DeadlineExceeded => LockError::DeadlineExceeded,
Self::Cancelled => LockError::Cancelled,
}
}
}
#[derive(Debug, Default)]
struct Inner {
locks: HashMap<LockKey, LockState>,
fates: HashMap<u64, Fate>,
priorities: HashMap<u64, u64>,
}
#[derive(Debug)]
pub struct LockManager {
inner: Mutex<Inner>,
wake: Condvar,
}
impl Default for LockManager {
fn default() -> Self {
Self::new()
}
}
impl LockManager {
pub fn new() -> Self {
Self {
inner: Mutex::new(Inner::default()),
wake: Condvar::new(),
}
}
pub fn acquire(&self, key: LockKey, request: LockRequest) -> Result<(), LockError> {
let LockRequest {
txn_id,
mode,
deadline,
control,
priority,
} = request;
let mut inner = self.inner.lock();
check_request_live(&control, deadline)?;
if let Some(priority) = priority {
inner.priorities.insert(txn_id, priority);
}
{
let state = inner.locks.entry(key.clone()).or_default();
match state.holders.iter().position(|h| h.txn_id == txn_id) {
Some(position) => {
let held = state.holders[position].mode;
if held == LockMode::Exclusive || mode == LockMode::Shared {
return Ok(());
}
if state.holders.len() == 1 && state.queue.is_empty() {
state.holders[position].mode = LockMode::Exclusive;
return Ok(());
}
}
None => {
let blocked = !state.queue.is_empty()
|| state.holders.iter().any(|h| !mode.compatible(h.mode));
if !blocked {
state.holders.push(Holder { txn_id, mode });
return Ok(());
}
}
}
if state.queue.iter().any(|w| w.txn_id == txn_id) {
return Err(LockError::InvalidRequest(format!(
"transaction {txn_id} already has a pending wait on {key:?}"
)));
}
state.queue.push_back(Waiter {
txn_id,
deadline,
control: control.clone(),
mode,
});
}
detect_deadlocks(&mut inner);
self.wake.notify_all();
let result = self.wait_for_grant(&mut inner, &key, txn_id, &control, deadline);
if result.is_err() {
process_key(&mut inner, &key);
detect_deadlocks(&mut inner);
self.wake.notify_all();
}
prune_if_idle(&mut inner, &key);
result
}
pub fn release(&self, txn_id: u64, key: &LockKey) {
let mut inner = self.inner.lock();
if inner
.locks
.get(key)
.is_some_and(|s| s.holders.iter().any(|h| h.txn_id == txn_id))
{
if let Some(state) = inner.locks.get_mut(key) {
state.holders.retain(|h| h.txn_id != txn_id);
}
process_key(&mut inner, key);
detect_deadlocks(&mut inner);
prune_if_idle(&mut inner, key);
}
drop(inner);
self.wake.notify_all();
}
pub fn release_all(&self, txn_id: u64) {
let mut inner = self.inner.lock();
inner.fates.remove(&txn_id);
inner.priorities.remove(&txn_id);
let mut affected = Vec::new();
for (key, state) in inner.locks.iter_mut() {
let touched = state.holders.iter().any(|h| h.txn_id == txn_id)
|| state.queue.iter().any(|w| w.txn_id == txn_id);
if touched {
state.holders.retain(|h| h.txn_id != txn_id);
state.queue.retain(|w| w.txn_id != txn_id);
affected.push(key.clone());
}
}
for key in &affected {
process_key(&mut inner, key);
}
detect_deadlocks(&mut inner);
inner
.locks
.retain(|_, state| !(state.holders.is_empty() && state.queue.is_empty()));
drop(inner);
self.wake.notify_all();
}
pub fn holds(&self, txn_id: u64, key: &LockKey) -> bool {
self.inner
.lock()
.locks
.get(key)
.is_some_and(|state| state.holders.iter().any(|h| h.txn_id == txn_id))
}
#[cfg(test)]
fn queued_waiters(&self, key: &LockKey) -> usize {
self.inner
.lock()
.locks
.get(key)
.map_or(0, |state| state.queue.len())
}
fn wait_for_grant(
&self,
inner: &mut MutexGuard<'_, Inner>,
key: &LockKey,
txn_id: u64,
control: &ExecutionControl,
deadline: Option<Instant>,
) -> Result<(), LockError> {
loop {
if let Some(fate) = inner.fates.remove(&txn_id) {
return Err(fate.into_error());
}
let (queued, held) = match inner.locks.get(key) {
Some(state) => (
state.queue.iter().any(|w| w.txn_id == txn_id),
state.holders.iter().any(|h| h.txn_id == txn_id),
),
None => (false, false),
};
if !queued {
if held {
return Ok(());
}
return Err(LockError::Cancelled);
}
if let Err(error) = check_request_live(control, deadline) {
remove_waiter(inner, key, txn_id);
return Err(error);
}
let wake_at = deadline
.unwrap_or_else(|| Instant::now() + CANCELLATION_POLL_INTERVAL)
.min(Instant::now() + CANCELLATION_POLL_INTERVAL);
let _ = self.wake.wait_until(inner, wake_at);
}
}
}
fn check_request_live(
control: &ExecutionControl,
deadline: Option<Instant>,
) -> Result<(), LockError> {
match control.checkpoint() {
Ok(()) => {}
Err(MongrelError::DeadlineExceeded) => return Err(LockError::DeadlineExceeded),
Err(_) => return Err(LockError::Cancelled),
}
if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
return Err(LockError::DeadlineExceeded);
}
Ok(())
}
fn remove_waiter(inner: &mut Inner, key: &LockKey, txn_id: u64) {
if let Some(state) = inner.locks.get_mut(key) {
state.queue.retain(|w| w.txn_id != txn_id);
}
}
fn prune_if_idle(inner: &mut Inner, key: &LockKey) {
if inner
.locks
.get(key)
.is_some_and(|state| state.holders.is_empty() && state.queue.is_empty())
{
inner.locks.remove(key);
}
}
fn process_key(inner: &mut Inner, key: &LockKey) {
let now = Instant::now();
let dead: Vec<u64> = match inner.locks.get(key) {
Some(state) => state
.queue
.iter()
.filter(|w| w.control.is_cancelled() || w.deadline.is_some_and(|d| now >= d))
.map(|w| w.txn_id)
.collect(),
None => return,
};
for txn_id in dead {
if let Some(state) = inner.locks.get_mut(key) {
if let Some(position) = state.queue.iter().position(|w| w.txn_id == txn_id) {
let waiter = state.queue.remove(position).expect("position found above");
let fate = match waiter.control.checkpoint() {
Err(MongrelError::DeadlineExceeded) => Fate::DeadlineExceeded,
Err(_) => Fate::Cancelled,
Ok(()) => Fate::DeadlineExceeded,
};
inner.fates.entry(txn_id).or_insert(fate);
}
}
}
loop {
let Some(state) = inner.locks.get_mut(key) else {
return;
};
let Some(front) = state.queue.front() else {
return;
};
let grantable = state
.holders
.iter()
.all(|h| h.txn_id == front.txn_id || front.mode.compatible(h.mode));
if !grantable {
return;
}
let waiter = state.queue.pop_front().expect("front checked above");
match state.holders.iter_mut().find(|h| h.txn_id == waiter.txn_id) {
Some(holder) => holder.mode = LockMode::Exclusive,
None => state.holders.push(Holder {
txn_id: waiter.txn_id,
mode: waiter.mode,
}),
}
}
}
fn build_wait_for_graph(inner: &Inner) -> BTreeMap<u64, Vec<u64>> {
let mut graph: BTreeMap<u64, Vec<u64>> = BTreeMap::new();
for state in inner.locks.values() {
for (index, waiter) in state.queue.iter().enumerate() {
let edges = graph.entry(waiter.txn_id).or_default();
for holder in &state.holders {
if holder.txn_id != waiter.txn_id && !waiter.mode.compatible(holder.mode) {
edges.push(holder.txn_id);
}
}
for ahead in state.queue.iter().take(index) {
if ahead.txn_id != waiter.txn_id && !waiter.mode.compatible(ahead.mode) {
edges.push(ahead.txn_id);
}
}
}
}
for edges in graph.values_mut() {
edges.sort_unstable();
edges.dedup();
}
graph
}
fn find_cycle(graph: &BTreeMap<u64, Vec<u64>>) -> Option<Vec<u64>> {
#[derive(Clone, Copy, PartialEq, Eq)]
enum Mark {
Visiting,
Done,
}
let mut marks: HashMap<u64, Mark> = HashMap::new();
let mut stack: Vec<u64> = Vec::new();
for &start in graph.keys() {
if marks.contains_key(&start) {
continue;
}
marks.insert(start, Mark::Visiting);
stack.push(start);
let mut frames: Vec<(u64, usize)> = vec![(start, 0)];
while let Some((node, index)) = frames.last().copied() {
let edges = graph.get(&node).map(Vec::as_slice).unwrap_or(&[]);
if index < edges.len() {
frames.last_mut().expect("last copied above").1 += 1;
let next = edges[index];
match marks.get(&next) {
None => {
marks.insert(next, Mark::Visiting);
stack.push(next);
frames.push((next, 0));
}
Some(Mark::Visiting) => {
let position = stack
.iter()
.position(|member| *member == next)
.expect("a visiting node is on the stack");
return Some(stack[position..].to_vec());
}
Some(Mark::Done) => {}
}
} else {
marks.insert(node, Mark::Done);
stack.pop();
frames.pop();
}
}
}
None
}
fn choose_victim(cycle: &[u64], priorities: &HashMap<u64, u64>) -> u64 {
*cycle
.iter()
.min_by(|a, b| {
let priority_a = priorities.get(a).copied().unwrap_or(DEFAULT_PRIORITY);
let priority_b = priorities.get(b).copied().unwrap_or(DEFAULT_PRIORITY);
priority_a.cmp(&priority_b).then_with(|| b.cmp(a))
})
.expect("a deadlock cycle is never empty")
}
fn detect_deadlocks(inner: &mut Inner) {
loop {
let graph = build_wait_for_graph(inner);
let Some(cycle) = find_cycle(&graph) else {
return;
};
let victim = choose_victim(&cycle, &inner.priorities);
let victim_key = inner.locks.iter().find_map(|(key, state)| {
state
.queue
.iter()
.any(|w| w.txn_id == victim)
.then(|| key.clone())
});
let Some(victim_key) = victim_key else {
return;
};
if let Some(state) = inner.locks.get_mut(&victim_key) {
state.queue.retain(|w| w.txn_id != victim);
}
let cycle = cycle
.iter()
.map(u64::to_string)
.collect::<Vec<_>>()
.join(" → ");
inner.fates.insert(victim, Fate::Deadlock { victim, cycle });
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::CancellationReason;
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use LockMode::{Exclusive, Shared};
fn manager() -> Arc<LockManager> {
Arc::new(LockManager::new())
}
fn request(txn_id: u64, mode: LockMode) -> LockRequest {
LockRequest::new(txn_id, mode, ExecutionControl::new(None))
}
fn timed_request(txn_id: u64, mode: LockMode, timeout: Duration) -> LockRequest {
request(txn_id, mode).with_timeout(timeout)
}
fn row_key(n: u64) -> LockKey {
LockKey::row(1, RowId(n))
}
fn spawn_acquire(
manager: &Arc<LockManager>,
key: LockKey,
request: LockRequest,
) -> mpsc::Receiver<Result<(), LockError>> {
let (sender, receiver) = mpsc::channel();
let manager = Arc::clone(manager);
thread::spawn(move || {
let result = manager.acquire(key, request);
sender.send(result).expect("receiver alive");
});
receiver
}
fn wait_for_queue(manager: &LockManager, key: &LockKey, n: usize) {
for _ in 0..200 {
if manager.queued_waiters(key) >= n {
return;
}
thread::sleep(Duration::from_millis(10));
}
panic!("queue did not reach {n} waiters on {key:?}");
}
#[test]
fn acquire_and_release_round_trip() {
let manager = LockManager::new();
let key = row_key(1);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
assert!(manager.holds(1, &key));
manager.release(1, &key);
assert!(!manager.holds(1, &key));
manager.acquire(key.clone(), request(2, Shared)).unwrap();
assert!(manager.holds(2, &key));
manager.release_all(2);
assert!(!manager.holds(2, &key));
}
#[test]
fn mode_compatibility_matrix() {
assert!(Shared.compatible(Shared));
assert!(!Shared.compatible(Exclusive));
assert!(!Exclusive.compatible(Shared));
assert!(!Exclusive.compatible(Exclusive));
}
#[test]
fn shared_locks_coexist() {
let manager = LockManager::new();
let key = row_key(7);
manager.acquire(key.clone(), request(1, Shared)).unwrap();
manager.acquire(key.clone(), request(2, Shared)).unwrap();
assert!(manager.holds(1, &key));
assert!(manager.holds(2, &key));
}
#[test]
fn exclusive_blocks_shared_and_exclusive_until_release() {
let manager = LockManager::new();
let key = row_key(8);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
let shared = manager.acquire(
key.clone(),
timed_request(2, Shared, Duration::from_millis(40)),
);
assert_eq!(shared, Err(LockError::DeadlineExceeded));
let exclusive = manager.acquire(
key.clone(),
timed_request(3, Exclusive, Duration::from_millis(40)),
);
assert_eq!(exclusive, Err(LockError::DeadlineExceeded));
manager.release(1, &key);
manager.acquire(key.clone(), request(2, Shared)).unwrap();
let still_blocked = manager.acquire(
key.clone(),
timed_request(3, Exclusive, Duration::from_millis(40)),
);
assert_eq!(still_blocked, Err(LockError::DeadlineExceeded));
}
#[test]
fn wait_respects_the_lock_deadline() {
let manager = LockManager::new();
let key = row_key(9);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
let started = Instant::now();
let result = manager.acquire(key, timed_request(2, Exclusive, Duration::from_millis(80)));
let elapsed = started.elapsed();
assert_eq!(result, Err(LockError::DeadlineExceeded));
assert_eq!(
result.unwrap_err().category(),
ErrorCategory::DeadlineExceeded
);
assert!(
elapsed >= Duration::from_millis(70),
"returned early: {elapsed:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"returned late: {elapsed:?}"
);
}
#[test]
fn expired_deadline_fails_fast_even_when_unlocked() {
let manager = LockManager::new();
let result = manager.acquire(row_key(10), timed_request(1, Exclusive, Duration::ZERO));
assert_eq!(result, Err(LockError::DeadlineExceeded));
}
#[test]
fn cancellation_wakes_the_waiter_and_clears_its_queue_entry() {
let manager = manager();
let key = row_key(11);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
let control = ExecutionControl::new(None);
let cancelled_rx = spawn_acquire(
&manager,
key.clone(),
LockRequest::new(2, Exclusive, control.clone()),
);
wait_for_queue(&manager, &key, 1);
control.cancel(CancellationReason::ClientRequest);
let result = cancelled_rx.recv_timeout(Duration::from_secs(5)).unwrap();
assert_eq!(result, Err(LockError::Cancelled));
assert!(!manager.holds(2, &key));
let next_rx = spawn_acquire(&manager, key.clone(), request(3, Exclusive));
wait_for_queue(&manager, &key, 1);
manager.release(1, &key);
assert_eq!(
next_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
Ok(())
);
}
#[test]
fn cancelled_control_fails_fast_even_when_unlocked() {
let manager = LockManager::new();
let control = ExecutionControl::new(None);
control.cancel(CancellationReason::ClientRequest);
let result = manager.acquire(row_key(12), LockRequest::new(1, Exclusive, control));
assert_eq!(result, Err(LockError::Cancelled));
}
#[test]
fn control_deadline_maps_to_deadline_exceeded() {
let manager = LockManager::new();
let key = row_key(13);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
let control = ExecutionControl::with_timeout(Duration::from_millis(50));
let result = manager.acquire(key, LockRequest::new(2, Exclusive, control));
assert_eq!(result, Err(LockError::DeadlineExceeded));
}
#[test]
fn fifo_grant_order_blocks_barging_readers() {
let manager = manager();
let key = row_key(14);
manager.acquire(key.clone(), request(1, Shared)).unwrap();
let writer_rx = spawn_acquire(&manager, key.clone(), request(2, Exclusive));
wait_for_queue(&manager, &key, 1);
let reader_rx = spawn_acquire(&manager, key.clone(), request(3, Shared));
wait_for_queue(&manager, &key, 2);
manager.release(1, &key);
assert_eq!(
writer_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
Ok(())
);
assert!(
reader_rx.recv_timeout(Duration::from_millis(150)).is_err(),
"reader barged ahead of the queued writer"
);
manager.release(2, &key);
assert_eq!(
reader_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
Ok(())
);
}
#[test]
fn release_all_grants_blocked_waiters() {
let manager = manager();
let key = row_key(24);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
let waiter_rx = spawn_acquire(&manager, key.clone(), request(2, Exclusive));
wait_for_queue(&manager, &key, 1);
manager.release_all(1);
assert_eq!(
waiter_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
Ok(())
);
assert!(manager.holds(2, &key));
}
#[test]
fn ab_ba_deadlock_kills_requester_when_it_is_youngest() {
let manager = manager();
let key_a = row_key(20);
let key_b = row_key(21);
manager
.acquire(key_a.clone(), request(1, Exclusive))
.unwrap();
manager
.acquire(key_b.clone(), request(2, Exclusive))
.unwrap();
let t1_rx = spawn_acquire(&manager, key_b.clone(), request(1, Exclusive));
wait_for_queue(&manager, &key_b, 1);
let error = manager
.acquire(key_a.clone(), request(2, Exclusive))
.unwrap_err();
match &error {
LockError::Deadlock { victim, cycle } => {
assert_eq!(*victim, 2, "youngest transaction must be the victim");
assert!(
cycle.contains('1') && cycle.contains('2'),
"cycle names both: {cycle}"
);
}
other => panic!("expected t2 to be the deadlock victim, got {other:?}"),
}
assert_eq!(error.category(), ErrorCategory::Deadlock);
manager.release_all(2);
assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
assert!(manager.holds(1, &key_b));
}
#[test]
fn ab_ba_deadlock_kills_youngest_even_when_it_is_not_the_requester() {
let manager = manager();
let key_a = row_key(22);
let key_b = row_key(23);
manager
.acquire(key_a.clone(), request(100, Exclusive))
.unwrap();
manager
.acquire(key_b.clone(), request(200, Exclusive))
.unwrap();
let t200 = {
let manager = Arc::clone(&manager);
let key_a = key_a.clone();
thread::spawn(move || {
let result = manager.acquire(key_a, request(200, Exclusive));
if result.is_err() {
manager.release_all(200);
}
result
})
};
wait_for_queue(&manager, &key_a, 1);
manager
.acquire(key_b.clone(), request(100, Exclusive))
.unwrap();
assert!(manager.holds(100, &key_b));
let victim = t200.join().unwrap();
assert!(
matches!(victim, Err(LockError::Deadlock { victim: 200, .. })),
"youngest transaction must be the victim: {victim:?}"
);
}
#[test]
fn three_transaction_cycle_kills_youngest() {
let manager = manager();
let key_a = row_key(30);
let key_b = row_key(31);
let key_c = row_key(32);
manager
.acquire(key_a.clone(), request(1, Exclusive))
.unwrap();
manager
.acquire(key_b.clone(), request(2, Exclusive))
.unwrap();
manager
.acquire(key_c.clone(), request(3, Exclusive))
.unwrap();
let t1_rx = spawn_acquire(&manager, key_b.clone(), request(1, Exclusive));
wait_for_queue(&manager, &key_b, 1);
let t2_rx = spawn_acquire(&manager, key_c.clone(), request(2, Exclusive));
wait_for_queue(&manager, &key_c, 1);
let error = manager
.acquire(key_a.clone(), request(3, Exclusive))
.unwrap_err();
assert!(
matches!(error, LockError::Deadlock { victim: 3, .. }),
"youngest of the 3-cycle must die: {error:?}"
);
manager.release_all(3);
assert_eq!(t2_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
manager.release_all(2);
assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
}
#[test]
fn lowest_priority_dies_before_age_is_considered() {
let manager = manager();
let key_a = row_key(40);
let key_b = row_key(41);
let key_c = row_key(42);
manager
.acquire(key_a.clone(), request(1, Exclusive))
.unwrap();
manager
.acquire(key_b.clone(), request(2, Exclusive))
.unwrap();
manager
.acquire(key_c.clone(), request(3, Exclusive))
.unwrap();
let t1 = {
let manager = Arc::clone(&manager);
let key_b = key_b.clone();
thread::spawn(move || {
let result = manager.acquire(key_b, request(1, Exclusive).with_priority(10));
if result.is_err() {
manager.release_all(1);
}
result
})
};
wait_for_queue(&manager, &key_b, 1);
let t2_rx = spawn_acquire(
&manager,
key_c.clone(),
request(2, Exclusive).with_priority(20),
);
wait_for_queue(&manager, &key_c, 1);
manager
.acquire(key_a.clone(), request(3, Exclusive).with_priority(30))
.unwrap();
assert!(manager.holds(3, &key_a));
let victim = t1.join().unwrap();
assert!(
matches!(victim, Err(LockError::Deadlock { victim: 1, .. })),
"lowest priority must die: {victim:?}"
);
manager.release_all(3);
assert_eq!(t2_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
}
#[test]
fn equal_priorities_fall_back_to_youngest() {
let manager = manager();
let key_a = row_key(50);
let key_b = row_key(51);
manager
.acquire(key_a.clone(), request(1, Exclusive).with_priority(5))
.unwrap();
manager
.acquire(key_b.clone(), request(2, Exclusive).with_priority(5))
.unwrap();
let t1_rx = spawn_acquire(
&manager,
key_b.clone(),
request(1, Exclusive).with_priority(5),
);
wait_for_queue(&manager, &key_b, 1);
let error = manager
.acquire(key_a.clone(), request(2, Exclusive).with_priority(5))
.unwrap_err();
assert!(
matches!(error, LockError::Deadlock { victim: 2, .. }),
"equal priorities must fall back to youngest: {error:?}"
);
manager.release_all(2);
assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
}
#[test]
fn same_transaction_requests_never_self_deadlock() {
let manager = LockManager::new();
let key = row_key(60);
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
manager.acquire(key.clone(), request(1, Exclusive)).unwrap();
manager.acquire(key.clone(), request(1, Shared)).unwrap();
assert!(manager.holds(1, &key));
let other = row_key(61);
manager.acquire(other.clone(), request(2, Shared)).unwrap();
manager
.acquire(other.clone(), request(2, Exclusive))
.unwrap();
let blocked = manager.acquire(
other,
timed_request(3, Exclusive, Duration::from_millis(40)),
);
assert_eq!(blocked, Err(LockError::DeadlineExceeded));
}
#[test]
fn detector_kills_a_self_loop_if_one_is_ever_constructed() {
let graph = BTreeMap::from([(7_u64, vec![7_u64])]);
let cycle = find_cycle(&graph).expect("a self loop is a cycle");
assert_eq!(cycle, vec![7]);
assert_eq!(choose_victim(&cycle, &HashMap::new()), 7);
}
#[test]
fn shared_to_exclusive_upgrade_deadlock_kills_youngest() {
let manager = manager();
let key = row_key(70);
manager.acquire(key.clone(), request(1, Shared)).unwrap();
manager.acquire(key.clone(), request(2, Shared)).unwrap();
let t1_rx = spawn_acquire(&manager, key.clone(), request(1, Exclusive));
wait_for_queue(&manager, &key, 1);
let error = manager
.acquire(key.clone(), request(2, Exclusive))
.unwrap_err();
assert!(
matches!(error, LockError::Deadlock { victim: 2, .. }),
"youngest upgrader must die: {error:?}"
);
manager.release_all(2);
assert_eq!(t1_rx.recv_timeout(Duration::from_secs(5)).unwrap(), Ok(()));
let blocked = manager.acquire(
key.clone(),
timed_request(3, Shared, Duration::from_millis(40)),
);
assert_eq!(
blocked,
Err(LockError::DeadlineExceeded),
"t1 must now hold Exclusive"
);
}
#[test]
fn sequence_barrier_is_exclusive_across_transactions() {
let manager = LockManager::new();
let sequence = LockKey::sequence_barrier("order_id");
manager
.acquire(sequence.clone(), request(1, Exclusive))
.unwrap();
let blocked = manager.acquire(
sequence.clone(),
timed_request(2, Exclusive, Duration::from_millis(40)),
);
assert_eq!(blocked, Err(LockError::DeadlineExceeded));
manager
.acquire(
LockKey::sequence_barrier("shipment_id"),
request(2, Exclusive),
)
.unwrap();
manager.release_all(1);
manager.acquire(sequence, request(2, Exclusive)).unwrap();
}
#[test]
fn ddl_barrier_excludes_dml_and_concurrent_ddl() {
let manager = LockManager::new();
let barrier = LockKey::schema_barrier();
manager
.acquire(barrier.clone(), request(1, Exclusive))
.unwrap();
let dml = manager.acquire(
barrier.clone(),
timed_request(2, Shared, Duration::from_millis(40)),
);
assert_eq!(dml, Err(LockError::DeadlineExceeded));
let ddl = manager.acquire(
barrier.clone(),
timed_request(3, Exclusive, Duration::from_millis(40)),
);
assert_eq!(ddl, Err(LockError::DeadlineExceeded));
manager.release_all(1);
manager
.acquire(barrier.clone(), request(2, Shared))
.unwrap();
manager.acquire(barrier, request(3, Shared)).unwrap();
}
#[test]
fn key_families_do_not_false_share() {
let manager = LockManager::new();
manager
.acquire(LockKey::row(1, RowId(5)), request(1, Exclusive))
.unwrap();
manager
.acquire(
LockKey::key(1, 5_u64.to_be_bytes().to_vec()),
request(2, Exclusive),
)
.unwrap();
manager
.acquire(
LockKey::range(
1,
Bound::Included(1_u64.to_be_bytes().to_vec()),
Bound::Excluded(6_u64.to_be_bytes().to_vec()),
),
request(3, Exclusive),
)
.unwrap();
}
#[test]
fn lock_errors_bridge_to_mongrel_error_and_the_taxonomy() {
let deadlock = LockError::Deadlock {
victim: 9,
cycle: "9 → 4 → 9".to_string(),
};
assert_eq!(deadlock.category(), ErrorCategory::Deadlock);
let mapped = MongrelError::from(deadlock);
assert!(
matches!(
&mapped,
MongrelError::Deadlock { victim: 9, cycle } if cycle == "9 → 4 → 9"
),
"deadlock maps to the dedicated variant, victim and cycle preserved: {mapped:?}"
);
assert_eq!(mapped.category(), ErrorCategory::Deadlock);
assert!(matches!(
MongrelError::from(LockError::DeadlineExceeded),
MongrelError::DeadlineExceeded
));
assert!(matches!(
MongrelError::from(LockError::Cancelled),
MongrelError::Cancelled
));
assert_eq!(LockError::Cancelled.category(), ErrorCategory::Cancelled);
assert!(matches!(
MongrelError::from(LockError::InvalidRequest("x".to_string())),
MongrelError::InvalidArgument(_)
));
}
#[test]
fn wait_for_graph_cycle_search_is_deterministic() {
let graph = BTreeMap::from([(1, vec![2]), (2, vec![3])]);
assert_eq!(find_cycle(&graph), None);
let graph = BTreeMap::from([(1, vec![2]), (2, vec![3]), (3, vec![2])]);
let cycle = find_cycle(&graph).expect("cycle");
assert_eq!(cycle, vec![2, 3]);
assert_eq!(choose_victim(&cycle, &HashMap::new()), 3);
assert_eq!(choose_victim(&cycle, &HashMap::from([(2, 5), (3, 1)])), 3);
assert_eq!(choose_victim(&cycle, &HashMap::from([(2, 1), (3, 5)])), 2);
}
}