use std::{fmt, sync::Arc, thread, time::Duration};
use coarsetime::Instant;
use itoa::Buffer;
use smallvec::SmallVec;
use windex::HashIndex;
use super::{txn_key_entry_comparison::TxnKeyEntryComparison, txn_lock_table::TxnLockTable};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(u8)]
pub enum LockType {
None = 0,
Exclusive = 1,
Shared = 2,
}
#[derive(Debug, Clone, Copy)]
pub struct TxnKeyEntry {
pub key_hash: i64,
pub lock_type: LockType,
}
impl TxnKeyEntry {
pub fn new(key_hash: i64, lock_type: LockType) -> Self {
Self {
key_hash,
lock_type,
}
}
}
impl fmt::Display for TxnKeyEntry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let key_hash_sign = if self.key_hash < 0 { "-" } else { "" };
let lock_str = match self.lock_type {
LockType::None => "-",
LockType::Shared => "s",
LockType::Exclusive => "x",
};
let mut buf = Buffer::new();
let num_str = buf.format(self.key_hash.unsigned_abs());
write!(f, "{key_hash_sign}{num_str}:{lock_str}")
}
}
#[derive(Debug, Clone, Copy)]
struct LockPlanSlot {
bucket: usize,
exclusive: bool,
}
pub struct TxnKeyEntries {
lock_table: TxnLockTable,
keys: SmallVec<[TxnKeyEntry; 8]>,
unified_store_key_locked: bool,
pub phase: i32,
held: SmallVec<[LockPlanSlot; 4]>,
latch: Option<Arc<HashIndex>>,
}
impl TxnKeyEntries {
pub fn new(initial_count: usize, lock_table: TxnLockTable) -> Self {
Self {
lock_table,
keys: if initial_count <= 8 {
SmallVec::new()
} else {
SmallVec::with_capacity(initial_count)
},
unified_store_key_locked: false,
phase: 0,
held: SmallVec::new(),
latch: None,
}
}
#[inline]
pub fn lock_table(&self) -> &TxnLockTable {
&self.lock_table
}
pub fn is_read_only(&self) -> bool {
!self.keys.iter().any(|k| k.lock_type == LockType::Exclusive)
}
pub fn count(&self) -> usize {
self.keys.len()
}
#[inline]
pub fn key_hashes(&self) -> impl Iterator<Item = i64> + '_ {
self.keys.iter().map(|k| k.key_hash)
}
pub fn add_key(&mut self, key_hash: i64, lock_type: LockType) {
self.keys.push(TxnKeyEntry {
key_hash,
lock_type,
});
}
fn lock_plan(&mut self, index: &HashIndex) -> SmallVec<[LockPlanSlot; 4]> {
if self.keys.is_empty() {
return SmallVec::new();
}
self
.keys
.sort_unstable_by(|a, b| TxnKeyEntryComparison::compare(index, a, b));
let mut plan: SmallVec<[LockPlanSlot; 4]> = SmallVec::with_capacity(self.keys.len());
for entry in &self.keys {
let bucket = index.bucket_index_for_hash(entry.key_hash as u64);
let exclusive = entry.lock_type == LockType::Exclusive;
if let Some(last) = plan.last_mut()
&& last.bucket == bucket
{
last.exclusive |= exclusive;
continue;
}
plan.push(LockPlanSlot { bucket, exclusive });
}
plan
}
fn release_held(&mut self) {
let Some(index) = self.latch.clone() else {
self.held.clear();
return;
};
for slot in self.held.drain(..).rev() {
if slot.exclusive {
index.bucket(slot.bucket).unlock_exclusive();
} else {
index.bucket(slot.bucket).unlock_shared();
}
}
}
fn acquire_plan(&mut self, plan: &[LockPlanSlot]) -> bool {
self.held.clear();
let index = self
.latch
.clone()
.expect("取锁前必已钉定本笔事务的索引版本");
for slot in plan {
let bucket = index.bucket(slot.bucket);
let taken = if slot.exclusive {
bucket.try_lock_exclusive()
} else {
bucket.try_lock_shared()
};
if taken {
self.held.push(*slot);
continue;
}
self.release_held();
return false;
}
true
}
pub fn lock_all_keys(&mut self) {
self.phase = 1;
let index = self.lock_table.pin();
let plan = self.lock_plan(&index);
if !plan.is_empty() {
self.latch = Some(index);
while !self.acquire_plan(&plan) {
thread::yield_now();
}
self.unified_store_key_locked = true;
}
self.phase = 0;
}
pub fn try_lock_all_keys(&mut self, lock_timeout: Duration) -> bool {
self.phase = 1;
let index = self.lock_table.pin();
let plan = self.lock_plan(&index);
if !plan.is_empty() {
self.latch = Some(index);
let start = Instant::now();
loop {
if self.acquire_plan(&plan) {
self.unified_store_key_locked = true;
self.phase = 0;
return true;
}
if lock_timeout.is_zero() || Duration::from(start.elapsed()) >= lock_timeout {
self.unified_store_key_locked = false;
self.latch = None;
self.phase = 0;
return false;
}
thread::yield_now();
}
}
self.phase = 0;
true
}
pub fn unlock_all_keys(&mut self) {
self.phase = 2;
if self.unified_store_key_locked {
self.release_held();
}
self.held.clear();
self.keys.clear();
self.latch = None;
self.unified_store_key_locked = false;
self.phase = 0;
}
pub fn get_lockset(&self) -> String {
use std::fmt::Write as _;
if self.keys.is_empty() {
return String::new();
}
let mut sb = String::with_capacity(self.keys.len() * 16 + 24);
for entry in &self.keys {
let _ = write!(sb, "{entry}");
}
let phase_str = match self.phase {
0 => "none",
1 => "lock",
_ => "unlock",
};
let _ = write!(sb, " (phase: {phase_str}))");
sb
}
}
impl Drop for TxnKeyEntries {
#[inline]
fn drop(&mut self) {
self.unlock_all_keys();
}
}