use std::time::Duration;
use enum_dispatch::enum_dispatch;
use parking_lot::Mutex;
use wtxn::{LockType, SlotVerifyHandle, TransactionManager, TxnProcedure};
#[enum_dispatch]
pub trait CustomTransactionProcedure: TxnProcedure {
fn bind_args(&mut self, _args: &[Vec<u8>]) {}
fn add_key(
&self,
txn_manager: &mut TransactionManager,
verifier: Option<&SlotVerifyHandle<'_>>,
key: &[u8],
lock_type: LockType,
) {
txn_manager.save_key_entry_to_lock(key, lock_type);
if !txn_manager.is_replaying
&& let Some(v) = verifier
&& !v.network_iterative_slot_verify(key, lock_type == LockType::Shared)
{
txn_manager.abort();
}
}
}
#[derive(Debug, Clone, Default)]
pub struct DefaultTxnProc {
pub id: u8,
}
impl TxnProcedure for DefaultTxnProc {
fn id(&self) -> u8 {
self.id
}
fn prepare(
&mut self,
_txn_manager: &mut TransactionManager,
_verifier: Option<&SlotVerifyHandle<'_>>,
) -> bool {
true
}
fn main(&mut self, _txn_manager: &mut TransactionManager, _output: &mut Vec<u8>) {}
fn finalize(&mut self, _txn_manager: &mut TransactionManager, _output: &mut Vec<u8>) {}
}
impl CustomTransactionProcedure for DefaultTxnProc {}
pub static LAST_SET_KV: Mutex<Option<(Vec<u8>, Vec<u8>)>> = Mutex::new(None);
#[derive(Debug, Clone, Default)]
pub struct SetTxnProc {
pub id: u8,
pub args: Vec<Vec<u8>>,
}
impl TxnProcedure for SetTxnProc {
fn id(&self) -> u8 {
self.id
}
fn prepare(
&mut self,
txn_manager: &mut TransactionManager,
verifier: Option<&SlotVerifyHandle<'_>>,
) -> bool {
for key in self.args.iter().step_by(2) {
self.add_key(txn_manager, verifier, key, LockType::Exclusive);
}
!self.args.is_empty()
}
fn main(&mut self, _txn_manager: &mut TransactionManager, _output: &mut Vec<u8>) {
if self.args.len() >= 2 {
let mut guard = LAST_SET_KV.lock();
*guard = Some((self.args[0].clone(), self.args[1].clone()));
}
}
fn finalize(&mut self, _txn_manager: &mut TransactionManager, _output: &mut Vec<u8>) {}
}
impl CustomTransactionProcedure for SetTxnProc {
fn bind_args(&mut self, args: &[Vec<u8>]) {
self.args = args.to_vec();
}
}
#[enum_dispatch(CustomTransactionProcedure)]
#[derive(Debug, Clone)]
pub enum CustomTxnProc {
Default(DefaultTxnProc),
Set(SetTxnProc),
}
impl TxnProcedure for CustomTxnProc {
fn id(&self) -> u8 {
match self {
Self::Default(p) => p.id(),
Self::Set(p) => p.id(),
}
}
fn fail_fast_on_key_lock_failure(&self) -> bool {
match self {
Self::Default(p) => p.fail_fast_on_key_lock_failure(),
Self::Set(p) => p.fail_fast_on_key_lock_failure(),
}
}
fn key_lock_timeout(&self) -> Duration {
match self {
Self::Default(p) => p.key_lock_timeout(),
Self::Set(p) => p.key_lock_timeout(),
}
}
fn prepare(
&mut self,
txn_manager: &mut TransactionManager,
verifier: Option<&SlotVerifyHandle<'_>>,
) -> bool {
match self {
Self::Default(p) => p.prepare(txn_manager, verifier),
Self::Set(p) => p.prepare(txn_manager, verifier),
}
}
fn main(&mut self, txn_manager: &mut TransactionManager, output: &mut Vec<u8>) {
match self {
Self::Default(p) => p.main(txn_manager, output),
Self::Set(p) => p.main(txn_manager, output),
}
}
fn finalize(&mut self, txn_manager: &mut TransactionManager, output: &mut Vec<u8>) {
match self {
Self::Default(p) => p.finalize(txn_manager, output),
Self::Set(p) => p.finalize(txn_manager, output),
}
}
}