use std::mem;
use std::sync::Arc;
use std::path::Path;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, Ordering};
use parking_lot::Mutex;
use futures::{future::{FutureExt, BoxFuture}, stream::{StreamExt, BoxStream}};
use async_stream::stream;
use pi_atom::Atom;
use pi_guid::Guid;
use pi_hash::XHashMap;
use pi_ordmap::{ordmap::OrdMap, asbtree::Tree};
use pi_async_rt::lock::spin_lock::SpinLock;
use pi_async_transaction::{AsyncTransaction,
Transaction2Pc,
Transaction2PcAllConflicts,
UnitTransaction,
SequenceTransaction,
TransactionTree,
TransactionError,
AsyncCommitLog,
ErrorLevel,
manager_2pc::Transaction2PcStatus};
use pi_ordmap::ordmap::ImOrdMap;
use crate::{Binary,
KVAction,
TableTrQos,
KVActionLog,
KVDBCommitConfirm,
KVTableTrError,
TableKeyConflict,
db::{KVDBTransaction, KVDBChildTrList},
key_version::{KeyVersions,
PrepareMode,
PreparedActions,
PreparedCommitError,
TableVersionContext,
Version,
VersionConflictKind,
VersionReceipt,
binary_state_equal,
has_prepared_conflict,
has_prepared_transaction,
take_prepared_for_commit},
tables::{KVTable, ordmap_snapshot::OrdMapSnapshot}};
#[derive(Clone)]
pub struct MemoryOrderedTable<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
>(Arc<InnerMemoryOrderedTable<C, Log>>);
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Send for MemoryOrderedTable<C, Log> {}
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Sync for MemoryOrderedTable<C, Log> {}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> KVTable for MemoryOrderedTable<C, Log> {
type Name = Atom;
type Tr = MemOrdTabTr<C, Log>;
type Error = KVTableTrError;
fn name(&self) -> <Self as KVTable>::Name {
self.0.name.clone()
}
fn path(&self) -> Option<&Path> {
None
}
#[inline]
fn is_persistent(&self) -> bool {
self.0.persistence
}
fn is_ordered(&self) -> bool {
true
}
fn len(&self) -> usize {
self.0.root.lock().size()
}
fn size(&self) -> u64 {
let root_copy = self.0.root.lock().clone();
root_copy.full_bytes_size()
}
fn transaction(&self,
source: Atom,
is_writable: bool,
is_persistent: bool,
prepare_timeout: u64,
commit_timeout: u64) -> Self::Tr {
MemOrdTabTr::new(source,
is_writable,
is_persistent,
prepare_timeout,
commit_timeout,
self.clone())
}
fn ready_collect(&self) -> BoxFuture<Result<(), Self::Error>> {
async move {
Ok(())
}.boxed()
}
fn collect(&self) -> BoxFuture<Result<(), Self::Error>> {
async move {
Ok(())
}.boxed()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> MemoryOrderedTable<C, Log> {
pub(crate) fn query_committed(&self, key: &Binary) -> Option<Binary> {
self.0.root.lock().get(key).cloned()
}
pub fn new(name: Atom,
is_persistence: bool) -> Self {
let root = Mutex::new(OrdMap::new(None));
let prepare = Mutex::new(XHashMap::default());
let inner = InnerMemoryOrderedTable {
name,
persistence: is_persistence,
root,
prepare,
marker: PhantomData,
};
MemoryOrderedTable(Arc::new(inner))
}
}
struct InnerMemoryOrderedTable<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> {
name: Atom, persistence: bool, root: Mutex<OrdMap<Tree<Binary, Binary>>>, prepare: Mutex<XHashMap<Guid, PreparedActions>>, marker: PhantomData<(C, Log)>, }
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Send for InnerMemoryOrderedTable<C, Log> {}
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Sync for InnerMemoryOrderedTable<C, Log> {}
#[derive(Clone)]
pub struct MemOrdTabTr<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
>(Arc<InnerMemOrdTabTr<C, Log>>);
#[derive(Clone, Copy)]
enum PrepareConflictKind {
Common,
First,
All,
}
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Send for MemOrdTabTr<C, Log> {}
unsafe impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Sync for MemOrdTabTr<C, Log> {}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> AsyncTransaction for MemOrdTabTr<C, Log> {
type Output = ();
type Error = KVTableTrError;
fn is_writable(&self) -> bool {
self.0.writable
}
fn is_concurrent_commit(&self) -> bool {
false
}
fn is_concurrent_rollback(&self) -> bool {
false
}
fn get_source(&self) -> Atom {
self.0.source.clone()
}
fn init(&self)
-> BoxFuture<Result<<Self as AsyncTransaction>::Output, <Self as AsyncTransaction>::Error>> {
async move {
Ok(())
}.boxed()
}
fn rollback(&self)
-> BoxFuture<Result<<Self as AsyncTransaction>::Output, <Self as AsyncTransaction>::Error>> {
let tr = self.clone();
async move {
let transaction_uid = tr.get_transaction_uid().unwrap();
let _ = tr.0.table.0.prepare.lock().remove(&transaction_uid);
if let Some(context) = tr.0.version_context.as_ref() {
context.release_snapshot();
}
Ok(())
}.boxed()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Transaction2Pc for MemOrdTabTr<C, Log> {
type Tid = Guid;
type Pid = Guid;
type Cid = Guid;
type PrepareOutput = Vec<u8>;
type PrepareError = KVTableTrError;
type ConfirmOutput = ();
type ConfirmError = KVTableTrError;
type CommitConfirm = KVDBCommitConfirm<C, Log>;
fn is_require_persistence(&self) -> bool {
self.0.persistence.load(Ordering::Relaxed)
}
fn require_persistence(&self) {
self.0.persistence.store(true, Ordering::Relaxed);
}
fn is_concurrent_prepare(&self) -> bool {
false
}
fn is_enable_inherit_uid(&self) -> bool {
true
}
fn get_transaction_uid(&self) -> Option<<Self as Transaction2Pc>::Tid> {
self.0.tid.lock().clone()
}
fn set_transaction_uid(&self, uid: <Self as Transaction2Pc>::Tid) {
*self.0.tid.lock() = Some(uid);
}
fn get_prepare_uid(&self) -> Option<<Self as Transaction2Pc>::Pid> {
None
}
fn set_prepare_uid(&self, _uid: <Self as Transaction2Pc>::Pid) {
}
fn get_commit_uid(&self) -> Option<<Self as Transaction2Pc>::Cid> {
self.0.cid.lock().clone()
}
fn set_commit_uid(&self, uid: <Self as Transaction2Pc>::Cid) {
*self.0.cid.lock() = Some(uid);
}
fn get_prepare_timeout(&self) -> u64 {
self.0.prepare_timeout
}
fn get_commit_timeout(&self) -> u64 {
self.0.commit_timeout
}
fn prepare(&self)
-> BoxFuture<Result<Option<<Self as Transaction2Pc>::PrepareOutput>, <Self as Transaction2Pc>::PrepareError>> {
let tr = self.clone();
async move {
tr.prepare_registered(PrepareConflictKind::Common).await
}.boxed()
}
fn prepare_conflicts(&self) -> BoxFuture<Result<Option<<Self as Transaction2Pc>::PrepareOutput>, <Self as Transaction2Pc>::PrepareError>> {
let tr = self.clone();
async move {
tr.prepare_registered(PrepareConflictKind::First).await
}.boxed()
}
fn commit(&self, confirm: <Self as Transaction2Pc>::CommitConfirm)
-> BoxFuture<Result<<Self as AsyncTransaction>::Output, <Self as AsyncTransaction>::Error>> {
let tr = self.clone();
async move {
let transaction_uid = tr.get_transaction_uid().unwrap();
let publication = match tr.0.version_context.as_ref() {
Some(context) => Some(context.versions().publication().write().await),
None => None,
};
let expected_mode = tr
.0
.version_context
.as_ref()
.map(TableVersionContext::mode)
.unwrap_or(PrepareMode::Ordinary);
let prepared = {
let mut prepare = tr.0.table.0.prepare.lock();
take_prepared_for_commit(&mut prepare,
&transaction_uid,
expected_mode,
tr.is_writable())
};
let prepared = match prepared {
Ok(prepared) => prepared,
Err(PreparedCommitError::ModeMismatch(prepared_mode)) => {
drop(publication);
if let Some(context) = tr.0.version_context.as_ref() {
context.release_snapshot();
}
return Err(KVTableTrError::new_transaction_error(
ErrorLevel::Fatal,
format!("Commit memory ordered table failed, table: {:?}, source: {:?}, transaction_uid: {:?}, expected_mode: {:?}, prepared_mode: {:?}, reason: prepared action protocol mismatch after entering non-rollbackable commit",
tr.0.table.name().as_str(),
tr.0.source,
transaction_uid,
expected_mode,
prepared_mode)));
},
Err(PreparedCommitError::Missing) => {
drop(publication);
if let Some(context) = tr.0.version_context.as_ref() {
context.release_snapshot();
}
return Err(KVTableTrError::new_transaction_error(
ErrorLevel::Fatal,
format!("Commit memory ordered table failed, table: {:?}, source: {:?}, transaction_uid: {:?}, expected_mode: {:?}, reason: prepared actions missing after entering non-rollbackable commit",
tr.0.table.name().as_str(),
tr.0.source,
transaction_uid,
expected_mode)));
},
};
let mut committed_versions = Vec::new();
if let Some(prepared) = prepared {
let has_writes = prepared.actions.values().any(|action| {
matches!(action, KVActionLog::Write(_) | KVActionLog::DirtyWrite(_))
});
if has_writes {
let revision = match tr.0.version_context.as_ref() {
Some(context) => {
match context.versions().checked_next_revision() {
Some(revision) => Some(revision),
None => {
drop(publication);
context.release_snapshot();
return Err(KVTableTrError::new_transaction_error(
ErrorLevel::Fatal,
format!("Commit memory ordered table failed, table: {:?}, source: {:?}, transaction_uid: {:?}, reason: key version revision exhausted",
tr.0.table.name().as_str(),
tr.0.source,
transaction_uid)));
},
}
},
None => None,
};
let mut root = tr.0.table.0.root.lock();
if root.ptr_eq(&tr.0.root_ref) {
*root = tr.0.root_mut.lock().clone();
} else {
for (key, action) in &prepared.actions {
match action {
KVActionLog::Write(None) | KVActionLog::DirtyWrite(None) => {
root.delete(key, false);
},
KVActionLog::Write(Some(value)) | KVActionLog::DirtyWrite(Some(value)) => {
root.upsert(key.clone(), value.clone(), false);
},
KVActionLog::Read => (),
}
}
}
if let (Some(context), Some(revision)) =
(tr.0.version_context.as_ref(), revision) {
for (key, action) in &prepared.actions {
let value = match action {
KVActionLog::Write(value) | KVActionLog::DirtyWrite(value) => value,
KVActionLog::Read => continue,
};
committed_versions.push(context.versions().publish(
tr.0.table.name(),
key.clone(),
value.as_ref(),
transaction_uid.clone(),
revision));
}
context.versions().complete_revision(revision);
if let Some(receipt) = context.receipt() {
receipt.append(committed_versions);
}
}
}
}
drop(publication);
if let Some(context) = tr.0.version_context.as_ref() {
context.release_snapshot();
}
if tr.is_require_persistence() {
let commit_uid = tr.get_commit_uid().unwrap();
if let Err(e) = confirm(transaction_uid.clone(), commit_uid, Ok(())) {
return Err(e);
}
}
Ok(())
}.boxed()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> Transaction2PcAllConflicts for MemOrdTabTr<C, Log> {
fn precheck_all_conflicts(&self)
-> BoxFuture<'_, Result<(), <Self as Transaction2Pc>::PrepareError>> {
let tr = self.clone();
async move {
tr.precheck_versions().await
}.boxed()
}
fn prepare_all_conflicts(&self)
-> BoxFuture<'_, Result<Option<<Self as Transaction2Pc>::PrepareOutput>, <Self as Transaction2Pc>::PrepareError>> {
let tr = self.clone();
async move {
tr.prepare_registered(PrepareConflictKind::All).await
}.boxed()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> UnitTransaction for MemOrdTabTr<C, Log> {
type Status = Transaction2PcStatus;
type Qos = TableTrQos;
fn is_unit(&self) -> bool {
true
}
fn get_status(&self) -> <Self as UnitTransaction>::Status {
self.0.status.lock().clone()
}
fn set_status(&self, status: <Self as UnitTransaction>::Status) {
*self.0.status.lock() = status;
}
fn qos(&self) -> <Self as UnitTransaction>::Qos {
if self.is_require_persistence() {
TableTrQos::Safe
} else {
TableTrQos::ThreadSafe
}
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> SequenceTransaction for MemOrdTabTr<C, Log> {
type Item = Self;
fn is_sequence(&self) -> bool {
false
}
fn prev_item(&self) -> Option<<Self as SequenceTransaction>::Item> {
None
}
fn next_item(&self) -> Option<<Self as SequenceTransaction>::Item> {
None
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> TransactionTree for MemOrdTabTr<C, Log> {
type Node = KVDBTransaction<C, Log>;
type NodeInterator = KVDBChildTrList<C, Log>;
fn is_tree(&self) -> bool {
false
}
fn children_len(&self) -> usize {
0
}
fn to_children(&self) -> Self::NodeInterator {
KVDBChildTrList::new()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> KVAction for MemOrdTabTr<C, Log> {
type Key = Binary;
type Value = Binary;
type Error = KVTableTrError;
fn dirty_query(&self, key: <Self as KVAction>::Key)
-> BoxFuture<Option<<Self as KVAction>::Value>> {
let tr = self.clone();
async move {
if let Some(value) = tr.0.root_mut.lock().get(&key) {
return Some(value.clone());
}
None
}.boxed()
}
fn query(&self, key: <Self as KVAction>::Key)
-> BoxFuture<Option<<Self as KVAction>::Value>> {
let tr = self.clone();
async move {
let mut actions_locked = tr.0.actions.lock();
if let None = actions_locked.get(&key) {
let _ = actions_locked.insert(key.clone(), KVActionLog::Read);
}
if let Some(value) = tr.0.root_mut.lock().get(&key) {
return Some(value.clone());
}
None
}.boxed()
}
fn dirty_upsert(&self,
key: <Self as KVAction>::Key,
value: <Self as KVAction>::Value)
-> BoxFuture<Result<(), <Self as KVAction>::Error>> {
let tr = self.clone();
async move {
let _ = tr.0.actions.lock().insert(key.clone(), KVActionLog::DirtyWrite(Some(value.clone())));
let _ = tr.0.root_mut.lock().upsert(key, value, false);
Ok(())
}.boxed()
}
fn upsert(&self,
key: <Self as KVAction>::Key,
value: <Self as KVAction>::Value)
-> BoxFuture<Result<(), <Self as KVAction>::Error>> {
let tr = self.clone();
async move {
let _ = tr.0.actions.lock().insert(key.clone(), KVActionLog::Write(Some(value.clone())));
let _ = tr.0.root_mut.lock().upsert(key, value, false);
Ok(())
}.boxed()
}
fn dirty_delete(&self, key: <Self as KVAction>::Key)
-> BoxFuture<Result<Option<<Self as KVAction>::Value>, <Self as KVAction>::Error>> {
let tr = self.clone();
async move {
let _ = tr.0.actions.lock().insert(key.clone(), KVActionLog::DirtyWrite(None));
if let Some(Some(value)) = tr.0.root_mut.lock().delete(&key, false) {
return Ok(Some(value));
}
Ok(None)
}.boxed()
}
fn delete(&self, key: <Self as KVAction>::Key)
-> BoxFuture<Result<Option<<Self as KVAction>::Value>, <Self as KVAction>::Error>> {
let tr = self.clone();
async move {
let _ = tr.0.actions.lock().insert(key.clone(), KVActionLog::Write(None));
if let Some(Some(value)) = tr.0.root_mut.lock().delete(&key, false) {
return Ok(Some(value));
}
Ok(None)
}.boxed()
}
fn keys<'a>(&self,
key: Option<<Self as KVAction>::Key>,
descending: bool)
-> BoxStream<'a, <Self as KVAction>::Key> {
let root = self.0.root_mut.lock().clone();
let mut iterator = OrdMapSnapshot::new(root, key.as_ref(), descending);
let stream = stream! {
while let Some(key) = iterator.next_key() {
yield key;
}
};
stream.boxed()
}
fn values<'a>(&self,
key: Option<<Self as KVAction>::Key>,
descending: bool)
-> BoxStream<'a, (<Self as KVAction>::Key, <Self as KVAction>::Value)> {
let root = self.0.root_mut.lock().clone();
let mut iterator = OrdMapSnapshot::new(root, key.as_ref(), descending);
let stream = stream! {
while let Some((key, value)) = iterator.next_entry() {
yield (key, value);
}
};
stream.boxed()
}
fn lock_key(&self, _key: <Self as KVAction>::Key)
-> BoxFuture<Result<(), <Self as KVAction>::Error>> {
async move {
Ok(())
}.boxed()
}
fn unlock_key(&self, _key: <Self as KVAction>::Key)
-> BoxFuture<Result<(), <Self as KVAction>::Error>> {
async move {
Ok(())
}.boxed()
}
}
impl<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> MemOrdTabTr<C, Log> {
#[inline]
fn new(source: Atom,
is_writable: bool,
is_persistence: bool,
prepare_timeout: u64,
commit_timeout: u64,
table: MemoryOrderedTable<C, Log>) -> Self {
let root_ref = table.0.root.lock().clone();
let inner = InnerMemOrdTabTr {
source,
tid: SpinLock::new(None),
cid: SpinLock::new(None),
status: SpinLock::new(Transaction2PcStatus::default()),
writable: is_writable,
persistence: AtomicBool::new(is_persistence),
prepare_timeout,
commit_timeout,
root_mut: SpinLock::new(root_ref.clone()),
root_ref,
table,
actions: SpinLock::new(XHashMap::default()),
version_context: None,
};
MemOrdTabTr(Arc::new(inner))
}
pub(crate) fn new_managed(source: Atom,
is_writable: bool,
is_persistence: bool,
prepare_timeout: u64,
commit_timeout: u64,
table: MemoryOrderedTable<C, Log>,
versions: KeyVersions,
mode: PrepareMode,
expected: XHashMap<Binary, Version>,
receipt: Option<VersionReceipt>,
actions: XHashMap<Binary, KVActionLog>) -> Self {
let root_locked = table.0.root.lock();
let root_ref = root_locked.clone();
let snapshot = versions.lease_current();
drop(root_locked);
let mut root_mut = root_ref.clone();
for (key, action) in &actions {
match action {
KVActionLog::Write(Some(value)) | KVActionLog::DirtyWrite(Some(value)) => {
let _ = root_mut.upsert(key.clone(), value.clone(), false);
},
KVActionLog::Write(None) | KVActionLog::DirtyWrite(None) => {
let _ = root_mut.delete(key, false);
},
KVActionLog::Read => (),
}
}
let version_context = TableVersionContext::new(versions,
snapshot,
mode,
expected,
receipt);
let inner = InnerMemOrdTabTr {
source,
tid: SpinLock::new(None),
cid: SpinLock::new(None),
status: SpinLock::new(Transaction2PcStatus::default()),
writable: is_writable,
persistence: AtomicBool::new(is_persistence),
prepare_timeout,
commit_timeout,
root_mut: SpinLock::new(root_mut),
root_ref,
table,
actions: SpinLock::new(actions),
version_context: Some(version_context),
};
MemOrdTabTr(Arc::new(inner))
}
async fn precheck_versions(&self) -> Result<(), KVTableTrError> {
let Some(context) = self.0.version_context.as_ref() else {
return Ok(());
};
if context.mode() != PrepareMode::Versioned {
return Ok(());
}
let _publication = context.versions().publication().read().await;
let mut conflicts = Vec::new();
for (key, expected) in context.expected() {
if context.versions().current_version(key).as_ref() != Some(expected) {
conflicts.push(TableKeyConflict {
table: self.0.table.name(),
key: key.clone(),
kind: VersionConflictKind::ReadSetVersionMismatch,
});
}
}
if conflicts.is_empty() {
Ok(())
} else {
Err(KVTableTrError::new_all_conflicts_error(conflicts))
}
}
async fn prepare_registered(&self,
conflict_kind: PrepareConflictKind)
-> Result<Option<Vec<u8>>, KVTableTrError> {
if !self.is_writable() {
return Ok(None);
}
let _publication = match self.0.version_context.as_ref() {
Some(context) => Some(context.versions().publication().read().await),
None => None,
};
let actions = self.0.actions.lock().clone();
let mode = self
.0
.version_context
.as_ref()
.map(TableVersionContext::mode)
.unwrap_or(PrepareMode::Ordinary);
let mut conflict_keys = Vec::new();
if let Some(context) = self.0.version_context.as_ref() {
if context.mode() == PrepareMode::Versioned {
for (key, expected) in context.expected() {
if context.versions().current_version(key).as_ref() != Some(expected) {
conflict_keys.push((key.clone(),
VersionConflictKind::ReadSetVersionMismatch));
}
}
}
}
let current_root = self.0.table.0.root.lock().clone();
for (key, action) in &actions {
let require_state_check = !self.is_require_persistence() || !action.is_dirty_writed();
if !require_state_check {
continue;
}
if let Some(context) = self.0.version_context.as_ref() {
if context
.versions()
.has_committed_after(key, context.snapshot_revision()) {
conflict_keys.push((key.clone(),
VersionConflictKind::TransactionConflict));
continue;
}
}
if !binary_state_equal(self.0.root_ref.get(key), current_root.get(key)) {
conflict_keys.push((key.clone(),
VersionConflictKind::TransactionConflict));
}
}
let write_buf = self.prepare_output(&actions);
let mut prepare = self.0.table.0.prepare.lock();
let transaction_uid = self.get_transaction_uid().unwrap();
if has_prepared_transaction(&prepare, &transaction_uid) {
return Err(KVTableTrError::new_transaction_error(
ErrorLevel::Normal,
format!("Prepare memory ordered table failed, table: {:?}, source: {:?}, transaction_uid: {:?}, reason: duplicate prepared transaction uid",
self.0.table.name().as_str(),
self.0.source,
transaction_uid)));
}
for (key, action) in &actions {
if has_prepared_conflict(&prepare, key, mode, action) {
conflict_keys.push((key.clone(),
VersionConflictKind::TransactionConflict));
}
}
if !conflict_keys.is_empty() {
return Err(self.prepare_conflict_error(conflict_kind, conflict_keys));
}
let _ = mem::replace(&mut *self.0.actions.lock(), XHashMap::default());
prepare.insert(transaction_uid, PreparedActions {
mode,
actions,
});
Ok(write_buf)
}
fn prepare_output(&self,
actions: &XHashMap<Binary, KVActionLog>) -> Option<Vec<u8>> {
if !self.is_require_persistence() {
return None;
}
let writed_count = actions
.values()
.filter(|action| matches!(action,
KVActionLog::Write(_) | KVActionLog::DirtyWrite(_)))
.count() as u64;
if writed_count == 0 {
return None;
}
let mut buf = Vec::new();
self.0.table.init_table_prepare_output(&mut buf, writed_count);
for (key, action) in actions {
match action {
KVActionLog::Write(None) | KVActionLog::DirtyWrite(None) => {
self.0.table.append_key_value_to_table_prepare_output(&mut buf, key, None);
},
KVActionLog::Write(Some(value)) | KVActionLog::DirtyWrite(Some(value)) => {
self.0.table.append_key_value_to_table_prepare_output(&mut buf,
key,
Some(value));
},
KVActionLog::Read => (),
}
}
Some(buf)
}
fn prepare_conflict_error(&self,
conflict_kind: PrepareConflictKind,
keys: Vec<(Binary, VersionConflictKind)>) -> KVTableTrError {
let key = keys[0].0.clone();
match conflict_kind {
PrepareConflictKind::Common => {
KVTableTrError::new_transaction_error(
ErrorLevel::Normal,
format!("Prepare memory ordered table conflicted, table: {:?}, key: {:?}, source: {:?}, transaction_uid: {:?}, reason: committed state or prepared reservation changed",
self.0.table.name().as_str(),
key,
self.0.source,
self.get_transaction_uid()))
},
PrepareConflictKind::First => {
KVTableTrError::new_conflicts_error(self.0.table.name(), key)
},
PrepareConflictKind::All => {
KVTableTrError::new_all_conflicts_error(keys
.into_iter()
.map(|(key, kind)| TableKeyConflict {
table: self.0.table.name(),
key,
kind,
})
.collect())
},
}
}
pub(crate) fn prepare_repair(&self, transaction_uid: Guid) {
let actions = mem::replace(&mut *self.0.actions.lock(), XHashMap::default());
for (key, action) in &actions {
match action {
KVActionLog::Write(Some(value)) | KVActionLog::DirtyWrite(Some(value)) => {
self
.0
.table
.0
.root
.lock()
.upsert(key.clone(), value.clone(), false);
},
KVActionLog::Write(None) | KVActionLog::DirtyWrite(None) => {
self.0.table.0.root.lock().delete(key, false);
},
KVActionLog::Read => (), }
}
self.0.table.0.prepare.lock().insert(transaction_uid, PreparedActions {
mode: PrepareMode::Ordinary,
actions,
});
}
}
struct InnerMemOrdTabTr<
C: Clone + Send + 'static,
Log: AsyncCommitLog<C = C, Cid = Guid>,
> {
source: Atom, tid: SpinLock<Option<Guid>>, cid: SpinLock<Option<Guid>>, status: SpinLock<Transaction2PcStatus>, writable: bool, persistence: AtomicBool, prepare_timeout: u64, commit_timeout: u64, root_mut: SpinLock<OrdMap<Tree<Binary, Binary>>>, root_ref: OrdMap<Tree<Binary, Binary>>, table: MemoryOrderedTable<C, Log>, actions: SpinLock<XHashMap<Binary, KVActionLog>>, version_context: Option<TableVersionContext>, }
#[cfg(test)]
mod memory_local_contract_tests {
use std::sync::Arc;
use futures::executor::block_on;
use pi_bon::{Encode, WriteBuffer};
use pi_store::commit_logger::CommitLogger;
use super::*;
type TestTable = MemoryOrderedTable<usize, CommitLogger>;
fn bon_usize(value: usize) -> Binary {
let mut buffer = WriteBuffer::new();
value.encode(&mut buffer);
Binary::new(buffer.bytes)
}
fn assert_binary(actual: Option<Binary>, expected: Option<&Binary>, label: &str) {
match (actual, expected) {
(Some(actual), Some(expected)) => {
assert_eq!(actual.as_ref(), expected.as_ref(), "{label}: value mismatch");
},
(None, None) => (),
(actual, expected) => {
panic!("{label}: presence mismatch, actual: {}, expected: {}",
actual.is_some(),
expected.is_some());
},
}
}
#[test]
fn test_memory_metadata_leaf_identity_and_qos_contract() {
let table = TestTable::new(Atom::from("memory_local_identity"), true);
assert_eq!(table.name().as_str(), "memory_local_identity");
assert!(table.path().is_none());
assert!(table.is_persistent());
assert!(table.is_ordered());
assert_eq!(table.len(), 0);
assert_eq!(table.size(), 0);
let transaction = table.transaction(Atom::from("memory local identity source"),
true,
false,
1_234,
5_678);
assert!(transaction.is_writable());
assert!(!transaction.is_concurrent_prepare());
assert!(!transaction.is_concurrent_commit());
assert!(!transaction.is_concurrent_rollback());
assert!(transaction.is_enable_inherit_uid());
assert_eq!(transaction.get_source().as_str(), "memory local identity source");
assert_eq!(transaction.get_prepare_timeout(), 1_234);
assert_eq!(transaction.get_commit_timeout(), 5_678);
assert_eq!(transaction.get_status(), Transaction2PcStatus::Start);
assert!(transaction.is_unit());
assert!(!transaction.is_sequence());
assert!(!transaction.is_tree());
assert!(transaction.prev_item().is_none());
assert!(transaction.next_item().is_none());
assert_eq!(transaction.children_len(), 0);
assert_eq!(transaction.to_children().count(), 0);
assert_eq!(transaction.qos(), TableTrQos::ThreadSafe);
assert!(block_on(transaction.init()).is_ok());
let tid = Guid(101);
let cid = Guid(102);
let unused_prepare_uid = Guid(103);
assert!(transaction.get_transaction_uid().is_none());
assert!(transaction.get_commit_uid().is_none());
assert!(transaction.get_prepare_uid().is_none());
transaction.set_transaction_uid(tid.clone());
transaction.set_commit_uid(cid.clone());
transaction.set_prepare_uid(unused_prepare_uid);
assert_eq!(transaction.get_transaction_uid(), Some(tid));
assert_eq!(transaction.get_commit_uid(), Some(cid));
assert!(transaction.get_prepare_uid().is_none());
transaction.set_status(Transaction2PcStatus::Actioning);
assert_eq!(transaction.get_status(), Transaction2PcStatus::Actioning);
transaction.require_persistence();
transaction.require_persistence();
assert!(transaction.is_require_persistence());
assert_eq!(transaction.qos(), TableTrQos::Safe);
let read_only = table.transaction(Atom::from("memory local read only"),
false,
true,
7,
9);
assert!(matches!(block_on(read_only.prepare()), Ok(None)));
assert!(table.0.prepare.lock().is_empty());
}
#[test]
fn test_memory_private_cow_and_final_action_contract() {
let table = TestTable::new(Atom::from("memory_local_actions"), false);
let retained_key = bon_usize(1);
let deleted_key = bon_usize(2);
let committed_value = bon_usize(10);
let first_value = bon_usize(11);
let final_value = bon_usize(12);
table.0.root.lock().upsert(deleted_key.clone(), committed_value.clone(), false);
let transaction = table.transaction(Atom::from("memory local actions source"),
true,
false,
100,
200);
block_on(transaction.upsert(retained_key.clone(), first_value))
.expect("first private upsert must succeed");
assert_binary(block_on(transaction.query(retained_key.clone())),
Some(&bon_usize(11)),
"transaction must observe first private upsert");
block_on(transaction.upsert(retained_key.clone(), final_value.clone()))
.expect("final private upsert must succeed");
let removed = block_on(transaction.delete(deleted_key.clone()))
.expect("private delete must succeed");
assert!(removed.is_none(), "Memory delete must not expose the old value");
assert_binary(transaction.0.root_mut.lock().get(&retained_key).cloned(),
Some(&final_value),
"private root must contain final upsert");
assert_binary(transaction.0.root_mut.lock().get(&deleted_key).cloned(),
None,
"private root must contain final delete");
assert_binary(table.query_committed(&retained_key),
None,
"uncommitted upsert must not reach shared root");
assert_binary(table.query_committed(&deleted_key),
Some(&committed_value),
"uncommitted delete must not reach shared root");
let actions = transaction.0.actions.lock();
assert_eq!(actions.len(), 2);
assert!(matches!(actions.get(&retained_key),
Some(KVActionLog::Write(Some(value)))
if value.as_ref() == final_value.as_ref()));
assert!(matches!(actions.get(&deleted_key), Some(KVActionLog::Write(None))));
}
#[test]
fn test_memory_prepare_wal_ownership_and_rollback_contract() {
let table = TestTable::new(Atom::from("memory_local_prepare"), true);
let upsert_key = bon_usize(21);
let delete_key = bon_usize(22);
let read_key = bon_usize(23);
let old_value = bon_usize(210);
let new_value = bon_usize(211);
table.0.root.lock().upsert(delete_key.clone(), old_value.clone(), false);
let transaction = table.transaction(Atom::from("memory local prepare source"),
true,
true,
300,
400);
let tid = Guid(201);
transaction.set_transaction_uid(tid.clone());
block_on(transaction.upsert(upsert_key.clone(), new_value.clone()))
.expect("private upsert before prepare must succeed");
block_on(transaction.delete(delete_key.clone()))
.expect("private delete before prepare must succeed");
assert!(block_on(transaction.query(read_key.clone())).is_none());
let output = block_on(transaction.prepare_conflicts())
.expect("prepare must succeed")
.expect("persistent writes must produce a Memory WAL fragment");
let (table_name, write_count, offset) =
<TestTable as KVTable>::get_init_table_prepare_output(&output, 0);
let (writes, end) =
<TestTable as KVTable>::get_all_key_value_from_table_prepare_output(
&output,
&table_name,
write_count,
offset);
assert_eq!(table_name.as_str(), "memory_local_prepare");
assert_eq!(write_count, 2, "Read must not be encoded into the WAL fragment");
assert_eq!(writes.len(), 2);
assert_eq!(end, output.len(), "decoder must consume the complete fragment");
assert!(writes.iter().any(|entry| {
entry.key.as_ref() == upsert_key.as_ref()
&& entry.value.as_ref().map(Binary::as_ref) == Some(new_value.as_ref())
}));
assert!(writes.iter().any(|entry| {
entry.key.as_ref() == delete_key.as_ref() && entry.value.is_none()
}));
assert!(transaction.0.actions.lock().is_empty(),
"successful prepare must transfer transaction-local actions");
{
let prepared = table.0.prepare.lock();
let item = prepared.get(&tid).expect("prepare map must reserve the root TID");
assert_eq!(item.mode, PrepareMode::Ordinary);
assert_eq!(item.actions.len(), 3, "prepared state must retain the Read dependency");
assert!(matches!(item.actions.get(&read_key), Some(KVActionLog::Read)));
assert!(matches!(item.actions.get(&delete_key), Some(KVActionLog::Write(None))));
}
assert_binary(table.query_committed(&upsert_key),
None,
"prepare must not publish the private upsert");
assert_binary(table.query_committed(&delete_key),
Some(&old_value),
"prepare must not publish the private delete");
block_on(transaction.rollback()).expect("rollback must release prepared state");
assert!(table.0.prepare.lock().is_empty());
assert_binary(table.query_committed(&upsert_key),
None,
"rollback must not publish private upsert");
assert_binary(table.query_committed(&delete_key),
Some(&old_value),
"rollback must retain committed value");
}
#[test]
fn test_memory_dirty_prepare_current_persistence_branch() {
let table = TestTable::new(Atom::from("memory_local_dirty"), true);
let key = bon_usize(31);
let initial_value = bon_usize(310);
let concurrent_value = bon_usize(311);
let nonpersistent_value = bon_usize(312);
let persistent_value = bon_usize(313);
let later_value = bon_usize(314);
table.0.root.lock().upsert(key.clone(), initial_value, false);
let nonpersistent = table.transaction(Atom::from("memory local nonpersistent dirty"),
true,
false,
500,
600);
nonpersistent.set_transaction_uid(Guid(301));
block_on(nonpersistent.dirty_upsert(key.clone(), nonpersistent_value))
.expect("nonpersistent dirty upsert must succeed locally");
table.0.root.lock().upsert(key.clone(), concurrent_value, false);
let error = block_on(nonpersistent.prepare_conflicts())
.expect_err("nonpersistent dirty write currently checks committed value state");
assert!(error.is_conflicts());
assert!(table.0.prepare.lock().is_empty());
let persistent = table.transaction(Atom::from("memory local persistent dirty"),
true,
true,
700,
800);
let persistent_tid = Guid(302);
persistent.set_transaction_uid(persistent_tid.clone());
block_on(persistent.dirty_upsert(key.clone(), persistent_value))
.expect("persistent dirty upsert must succeed locally");
table.0.root.lock().upsert(key.clone(), later_value.clone(), false);
assert!(block_on(persistent.prepare_conflicts())
.expect("persistent dirty write currently skips committed value-state comparison")
.is_some());
assert!(table.0.prepare.lock().contains_key(&persistent_tid));
assert_binary(table.query_committed(&key),
Some(&later_value),
"prepare must not publish persistent dirty value");
block_on(persistent.rollback()).expect("persistent dirty rollback must succeed");
assert!(table.0.prepare.lock().is_empty());
}
#[test]
fn test_memory_repair_final_state_and_local_idempotence() {
let table = TestTable::new(Atom::from("memory_local_repair"), true);
let upsert_key = bon_usize(41);
let delete_key = bon_usize(42);
let old_value = bon_usize(410);
let repaired_value = bon_usize(411);
table.0.root.lock().upsert(delete_key.clone(), old_value, false);
let first = table.transaction(Atom::from("memory local repair first"),
true,
true,
900,
1_000);
block_on(first.upsert(upsert_key.clone(), repaired_value.clone()))
.expect("first repair upsert action must be staged");
block_on(first.delete(delete_key.clone()))
.expect("first repair delete action must be staged");
let first_tid = Guid(401);
first.prepare_repair(first_tid.clone());
assert_binary(table.query_committed(&upsert_key),
Some(&repaired_value),
"repair must apply upsert directly");
assert_binary(table.query_committed(&delete_key),
None,
"repair must apply delete directly");
assert!(first.0.actions.lock().is_empty());
{
let prepared = table.0.prepare.lock();
let item = prepared.get(&first_tid).expect("repair must register first TID");
assert_eq!(item.mode, PrepareMode::Ordinary);
assert_eq!(item.actions.len(), 2);
}
let second = table.transaction(Atom::from("memory local repair second"),
true,
true,
1_100,
1_200);
block_on(second.upsert(upsert_key.clone(), repaired_value.clone()))
.expect("second repair upsert action must be staged");
block_on(second.delete(delete_key.clone()))
.expect("second repair delete action must be staged");
let second_tid = Guid(402);
second.prepare_repair(second_tid.clone());
assert_eq!(table.len(), 1);
assert_binary(table.query_committed(&upsert_key),
Some(&repaired_value),
"repeated repair must retain final upsert");
assert_binary(table.query_committed(&delete_key),
None,
"repeated repair must retain final delete");
let mut prepared = table.0.prepare.lock();
assert_eq!(prepared.len(), 2);
assert_eq!(prepared.remove(&first_tid).map(|item| item.mode),
Some(PrepareMode::Ordinary));
assert_eq!(prepared.remove(&second_tid).map(|item| item.mode),
Some(PrepareMode::Ordinary));
assert!(prepared.is_empty());
}
#[test]
fn test_memory_transaction_table_arc_lifecycle_has_no_cycle() {
let table = TestTable::new(Atom::from("memory_local_arc_lifecycle"), false);
let weak = Arc::downgrade(&table.0);
let transaction = table.transaction(Atom::from("memory local arc lifecycle source"),
true,
false,
1_300,
1_400);
assert_eq!(Arc::strong_count(&table.0), 2);
drop(table);
assert!(weak.upgrade().is_some(), "transaction must keep its table owner alive");
drop(transaction);
assert!(weak.upgrade().is_none(), "dropping the last transaction must release the table");
}
}