use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use bitflags::bitflags;
use smallvec::SmallVec;
use wbase::store_type::{REPLAY_TASK_ACCESS_VECTOR_BYTES, StoreType};
use crate::{
TxnState,
txn_key_entry::{LockType, TxnKeyEntries},
txn_keys_buffer::TxnKeysBuffer,
txn_lock_table::TxnLockTable,
txn_slot_verify::SlotVerifyHandle,
txn_watched_keys_container::TxnWatchedKeysContainer,
watch_version_map::WatchVersionMap,
};
static GLOBAL_TXN_VERSION_SEQ: AtomicU64 = AtomicU64::new(1);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TxnEntryType {
TxnStart,
TxnCommit,
StoredProcedure,
}
pub type SublogVirtualVectors = SmallVec<[[u8; REPLAY_TASK_ACCESS_VECTOR_BYTES]; 4]>;
#[derive(Clone, Copy, Default, Debug)]
pub struct SublogAccess<'a> {
pub physical_vector: u64,
pub virtual_vectors: &'a [[u8; REPLAY_TASK_ACCESS_VECTOR_BYTES]],
pub participant_count: usize,
}
pub trait TxnAofLog: Send + Sync {
fn size(&self) -> usize;
fn replay_task_count(&self) -> usize;
fn get_physical_sublog_idx(&self, key_hash: i64) -> usize;
fn get_replay_task_idx(&self, key_hash: i64) -> usize;
fn enqueue_txn(
&self,
op_type: TxnEntryType,
txn_version: i64,
session_id: i32,
access: &SublogAccess<'_>,
) -> waof::Result<()>;
fn enqueue_stored_proc(
&self,
op_type: TxnEntryType,
txn_version: i64,
session_id: i32,
proc_id: u8,
payload: &[u8],
access: &SublogAccess<'_>,
) -> waof::Result<()>;
}
bitflags! {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TransactionStoreTypes: u8 {
const None = 0;
const Main = 1;
const Object = 1 << 1;
const Unified = 1 << 2;
}
}
pub trait TxnProcApi {
fn get(&mut self, key: &[u8]) -> Option<Vec<u8>>;
fn set(&mut self, key: &[u8], val: &[u8]) -> bool;
fn setex(&mut self, key: &[u8], val: &[u8], expiry_ticks: i64) -> bool;
fn delete(&mut self, key: &[u8]) -> bool;
fn increment(&mut self, key: &[u8], delta: i64) -> Option<i64>;
fn sorted_set_add(&mut self, key: &[u8], score: f64, member: &[u8]) -> bool;
fn sorted_set_remove(&mut self, key: &[u8], member: &[u8]) -> bool;
}
pub trait TxnProcReadApi {
fn get(&mut self, txn: &mut TransactionManager, key: &[u8]) -> Option<Vec<u8>>;
}
pub struct TxnWatchApi<'a> {
api: &'a mut (dyn TxnProcApi + 'a),
}
impl TxnProcReadApi for TxnWatchApi<'_> {
fn get(&mut self, txn: &mut TransactionManager, key: &[u8]) -> Option<Vec<u8>> {
txn.watch(key);
self.api.get(key)
}
}
pub trait TxnProcedure {
fn id(&self) -> u8;
fn fail_fast_on_key_lock_failure(&self) -> bool {
false
}
fn key_lock_timeout(&self) -> Duration {
Duration::ZERO
}
fn prepare(
&mut self,
txn_manager: &mut TransactionManager,
api: &mut dyn TxnProcReadApi,
verifier: Option<&SlotVerifyHandle<'_>>,
) -> bool;
fn main(
&mut self,
txn_manager: &mut TransactionManager,
api: &mut dyn TxnProcApi,
output: &mut Vec<u8>,
);
fn finalize(
&mut self,
txn_manager: &mut TransactionManager,
api: &mut dyn TxnProcApi,
output: &mut Vec<u8>,
);
}
pub struct TransactionGuard<'a> {
txn_manager: Option<&'a mut TransactionManager>,
}
impl TransactionGuard<'_> {
fn null() -> Self {
Self { txn_manager: None }
}
pub fn state(&self) -> TxnState {
self
.txn_manager
.as_deref()
.map_or(TxnState::None, |txn_manager| txn_manager.state)
}
pub fn dispose(&mut self) {
if let Some(txn_manager) = self.txn_manager.take()
&& let Err(err) = txn_manager.commit(true)
{
log::error!("事务提升守卫提交失败: {err:?}");
}
}
}
impl Drop for TransactionGuard<'_> {
fn drop(&mut self) {
self.dispose();
}
}
impl Drop for TransactionManager {
fn drop(&mut self) {
self.reset();
}
}
pub struct TransactionManager {
pub state: TxnState,
pub key_entries: TxnKeyEntries,
pub watch_container: TxnWatchedKeysContainer,
pub store_types: TransactionStoreTypes,
pub txn_start_head: usize,
pub operation_cnt_txn: usize,
pub perform_writes: bool,
pub txn_version: i64,
pub aof_log: Option<Arc<dyn TxnAofLog>>,
pub session_id: i32,
pub stored_proc_mode: bool,
pub is_replaying: bool,
pub txn_keys: TxnKeysBuffer,
pub cluster_enabled: bool,
}
impl TransactionManager {
pub fn new(
lock_table: TxnLockTable,
watch_version_map: Arc<WatchVersionMap>,
aof_log: Option<Arc<dyn TxnAofLog>>,
) -> Self {
Self {
state: TxnState::None,
key_entries: TxnKeyEntries::new(16, lock_table),
watch_container: TxnWatchedKeysContainer::new(watch_version_map),
store_types: TransactionStoreTypes::None,
txn_start_head: 0,
operation_cnt_txn: 0,
perform_writes: false,
txn_version: 0,
aof_log,
session_id: 0,
stored_proc_mode: false,
is_replaying: false,
txn_keys: TxnKeysBuffer::new(),
cluster_enabled: false,
}
}
#[inline]
pub fn set_session_id(&mut self, session_id: i32) {
self.session_id = session_id;
}
pub fn aof_enabled(&self) -> bool {
self.aof_log.is_some()
}
pub fn reset(&mut self) {
self.key_entries.unlock_all_keys();
self.txn_version = 0;
self.txn_start_head = 0;
self.operation_cnt_txn = 0;
self.state = TxnState::None;
self.store_types = TransactionStoreTypes::None;
self.stored_proc_mode = false;
self.is_replaying = false;
self.perform_writes = false;
self.txn_keys.clear();
}
pub fn run(
&mut self,
internal_txn: bool,
fail_fast_on_lock: bool,
lock_timeout: Duration,
) -> bool {
if !internal_txn {
let Self {
watch_container,
key_entries,
perform_writes,
..
} = self;
for key in watch_container.save_keys_to_lock() {
Self::register_key_lock(key_entries, perform_writes, key, LockType::Shared);
}
}
self.txn_version = GLOBAL_TXN_VERSION_SEQ.fetch_add(1, Ordering::Relaxed) as i64;
let lock_success = if fail_fast_on_lock {
self.key_entries.try_lock_all_keys(lock_timeout)
} else {
self.key_entries.lock_all_keys();
true
};
if !lock_success || (!internal_txn && !self.watch_container.validate_watch_version()) {
if !lock_success {
log::error!("Transaction failed to acquire all the locks on keys to proceed.");
}
self.reset();
if !internal_txn {
self.watch_container.reset();
}
return false;
}
if self.perform_writes
&& !self.stored_proc_mode
&& let Some(log) = self.aof_log.as_deref()
&& self
.enqueue_txn_marker(log, TxnEntryType::TxnStart)
.is_err()
{
self.reset();
if !internal_txn {
self.watch_container.reset();
}
return false;
}
self.state = TxnState::Running;
true
}
pub fn commit(&mut self, internal_txn: bool) -> waof::Result<()> {
let enqueued = if self.perform_writes
&& !self.stored_proc_mode
&& let Some(log) = self.aof_log.as_deref()
{
self.enqueue_txn_marker(log, TxnEntryType::TxnCommit)
} else {
Ok(())
};
if !internal_txn {
self.watch_container.reset();
}
self.reset();
enqueued
}
pub fn abort(&mut self) {
self.state = TxnState::Aborted;
}
#[inline]
pub fn is_read_only(&self) -> bool {
self.key_entries.is_read_only()
}
#[inline]
pub fn add_txn_key(&mut self, key: &[u8]) {
if !self.cluster_enabled {
return;
}
self.txn_keys.push(key);
}
pub fn watch(&mut self, key: &[u8]) {
self.watch_container.add_watch(key);
self.add_txn_key(key);
}
#[inline]
pub fn add_transaction_store_types(&mut self, types: TransactionStoreTypes) {
self.store_types |= types;
}
pub fn add_transaction_store_type(&mut self, store_type: StoreType) {
let transaction_store_types = match store_type {
StoreType::Main => TransactionStoreTypes::Main,
StoreType::Object => TransactionStoreTypes::Object,
StoreType::All => TransactionStoreTypes::Unified,
StoreType::None => TransactionStoreTypes::None,
};
self.store_types |= transaction_store_types;
}
pub fn get_lockset(&self) -> String {
self.key_entries.get_lockset()
}
pub fn promote_to_transaction(
&mut self,
store_types: TransactionStoreTypes,
key: &[u8],
lock_type: LockType,
) -> TransactionGuard<'_> {
if self.state == TxnState::Running {
return TransactionGuard::null();
}
self.add_transaction_store_types(store_types);
self.save_key_entry_to_lock(key, lock_type);
let _ = self.run(true, false, Duration::ZERO);
TransactionGuard {
txn_manager: Some(self),
}
}
#[inline]
fn enqueue_txn_marker(&self, log: &dyn TxnAofLog, op_type: TxnEntryType) -> waof::Result<()> {
let (physical_vector, virtual_vectors, participant_count) = self.compute_sublog_access_vector();
log.enqueue_txn(
op_type,
self.txn_version,
self.session_id,
&SublogAccess {
physical_vector,
virtual_vectors: &virtual_vectors,
participant_count: participant_count as usize,
},
)
}
pub fn compute_sublog_access_vector(&self) -> (u64, SublogVirtualVectors, u32) {
let Some(log) = self.aof_log.as_deref() else {
return (0, SublogVirtualVectors::new(), 0);
};
if log.size() <= 1 && log.replay_task_count() <= 1 {
return (0, SublogVirtualVectors::new(), 0);
}
let sublog_size = log.size();
let mut virtual_vectors: SublogVirtualVectors =
smallvec::smallvec![[0u8; REPLAY_TASK_ACCESS_VECTOR_BYTES]; sublog_size];
let mut physical_vector = 0u64;
let mut participant_count = 0u32;
for key_hash in self.key_entries.key_hashes() {
let physical_idx = log.get_physical_sublog_idx(key_hash);
let replay_idx = log.get_replay_task_idx(key_hash);
if physical_idx < 64 {
physical_vector |= 1u64 << physical_idx;
}
let byte_idx = replay_idx / 8;
let bit_mask = 1u8 << (replay_idx % 8);
if let Some(sublog_vector) = virtual_vectors.get_mut(physical_idx)
&& let Some(slot) = sublog_vector.get_mut(byte_idx)
&& (*slot & bit_mask) == 0
{
*slot |= bit_mask;
participant_count += 1;
}
}
(physical_vector, virtual_vectors, participant_count)
}
fn log_proc(
&mut self,
proc: &(impl TxnProcedure + ?Sized),
proc_input: &[u8],
) -> waof::Result<()> {
debug_assert!(self.stored_proc_mode);
if self.perform_writes
&& let Some(log) = self.aof_log.as_deref()
{
let (physical_vector, virtual_vectors, participant_count) =
self.compute_sublog_access_vector();
return log.enqueue_stored_proc(
TxnEntryType::StoredProcedure,
self.txn_version,
self.session_id,
proc.id(),
proc_input,
&SublogAccess {
physical_vector,
virtual_vectors: &virtual_vectors,
participant_count: participant_count as usize,
},
);
}
Ok(())
}
pub fn run_transaction_proc(
&mut self,
proc: &mut (impl TxnProcedure + ?Sized),
proc_input: &[u8],
output: &mut Vec<u8>,
is_replaying: bool,
verifier: Option<&SlotVerifyHandle<'_>>,
api: &mut dyn TxnProcApi,
) -> bool {
self.is_replaying = is_replaying;
if let Some(v) = verifier {
v.reset_cached_slot_verification_result();
}
self.stored_proc_mode = true;
let mut watch_api = TxnWatchApi { api };
let ran = if !proc.prepare(self, &mut watch_api, verifier) {
false
} else if self.state == TxnState::Aborted {
if let Some(v) = verifier {
v.write_cached_slot_verification_message(output);
}
false
} else if !self.run(
false,
proc.fail_fast_on_key_lock_failure(),
proc.key_lock_timeout(),
) {
false
} else {
proc.main(self, api, output);
let logged = if is_replaying {
true
} else {
self.log_proc(proc, proc_input).is_ok()
};
logged && self.commit(false).is_ok()
};
if !ran {
self.reset();
}
if !is_replaying {
proc.finalize(self, api, output);
}
ran
}
}