use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use noxu::xa::{PrepareResult, XaEnvironment, XaFlags, XaResource, Xid};
use noxu::{Database, DatabaseConfig, Environment, EnvironmentConfig};
use crate::datastore::noxu::{
delete_object_in, get_object_in, map_txn_store_error, NoxuDatastoreError,
};
use crate::txn::{TxnBatch, TxnOp, TxnOutcome, TxnStoreError};
pub(crate) const DYNIAK_XA_FORMAT_ID: i32 = 0x6479_6e6b;
const XA_DB_NAME: &str = "riak.objects";
pub struct XaParticipant {
xa: XaEnvironment,
db: Database,
name: Vec<u8>,
}
impl XaParticipant {
pub fn open(path: &Path, name: Vec<u8>) -> Result<Self, NoxuDatastoreError> {
let env_config = EnvironmentConfig::new(path.to_path_buf())
.with_allow_create(true)
.with_transactional(true);
let env = Environment::open(env_config)?;
let xa = XaEnvironment::new(env).with_prepared_log()?;
let db_config = DatabaseConfig::new()
.with_allow_create(true)
.with_transactional(true);
let db = xa.inner().open_database(None, XA_DB_NAME, &db_config)?;
Ok(Self { xa, db, name })
}
#[must_use]
pub fn name(&self) -> &[u8] {
&self.name
}
#[must_use]
pub fn xa(&self) -> &XaEnvironment {
&self.xa
}
pub fn apply_op(&self, xid: &Xid, op: &TxnOp) -> Result<(), NoxuDatastoreError> {
self.apply(xid, op)
}
pub fn get_object(
&self,
bucket: &[u8],
key: &[u8],
) -> Result<Option<Vec<u8>>, NoxuDatastoreError> {
get_object_in(&self.db, None, bucket, key)
}
fn apply(&self, xid: &Xid, op: &TxnOp) -> Result<(), NoxuDatastoreError> {
let txn = self
.xa
.get_transaction(xid)
.map_err(|e| NoxuDatastoreError::Xa(e.to_string()))?;
let txn: &noxu::Transaction = &txn;
match op {
TxnOp::Put {
bucket,
key,
value,
indexes,
} => crate::datastore::noxu::put_object_in(
&self.db,
Some(txn),
bucket,
key,
value,
indexes,
),
TxnOp::Delete { bucket, key } => {
delete_object_in(&self.db, Some(txn), bucket, key).map(|_| ())
}
}
}
}
pub struct XaCoordinator {
participants: Vec<XaParticipant>,
next_gtid: AtomicU64,
}
impl XaCoordinator {
#[must_use]
pub fn new(participants: Vec<XaParticipant>) -> Self {
Self {
participants,
next_gtid: AtomicU64::new(1),
}
}
#[must_use]
pub fn len(&self) -> usize {
self.participants.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.participants.is_empty()
}
#[must_use]
pub fn participant(&self, index: usize) -> Option<&XaParticipant> {
self.participants.get(index)
}
pub fn execute<R>(&self, batch: &TxnBatch, route: R) -> Result<TxnOutcome, TxnStoreError>
where
R: Fn(&TxnOp) -> usize,
{
if batch.ops.is_empty() {
return Err(TxnStoreError::EmptyBatch);
}
let mut per_branch: Vec<Vec<usize>> = vec![Vec::new(); self.participants.len()];
for (op_idx, op) in batch.ops.iter().enumerate() {
let branch = route(op);
if branch >= self.participants.len() {
return Err(TxnStoreError::Backend(format!(
"routing returned participant index {branch} but only {} participants exist",
self.participants.len()
)));
}
per_branch[branch].push(op_idx);
}
let gtid = self.next_gtid.fetch_add(1, Ordering::Relaxed);
let gtid_bytes = gtid.to_be_bytes();
let mut active: Vec<(usize, Xid)> = Vec::new();
for (branch, ops) in per_branch.iter().enumerate() {
if ops.is_empty() {
continue;
}
let xid = Xid::new(
DYNIAK_XA_FORMAT_ID,
>id_bytes,
self.participants[branch].name(),
)
.map_err(|e| TxnStoreError::Backend(format!("xid: {e}")))?;
active.push((branch, xid));
}
for (started, (branch, xid)) in active.iter().enumerate() {
let participant = &self.participants[*branch];
let phase1 = Self::run_branch_work(participant, xid, &per_branch[*branch], &batch.ops);
if let Err(e) = phase1 {
let _ = participant.xa.xa_rollback(xid, XaFlags::NOFLAGS);
self.rollback_active(&active[..started]);
return Err(e);
}
}
let mut prepared: Vec<(usize, Xid)> = Vec::new();
for (branch, xid) in &active {
let participant = &self.participants[*branch];
match participant.xa.xa_prepare(xid, XaFlags::NOFLAGS) {
Ok(PrepareResult::Ok) => prepared.push((*branch, xid.clone())),
Ok(PrepareResult::ReadOnly) => {}
Err(e) => {
let _ = participant.xa.xa_rollback(xid, XaFlags::NOFLAGS);
self.rollback_active(&prepared);
self.rollback_remaining(&active, *branch, &prepared);
return Err(map_xa_error(&e));
}
}
}
if batch.force_abort {
self.rollback_active(&prepared);
return Ok(TxnOutcome::Aborted {
reason: "client requested abort".to_string(),
});
}
for (branch, xid) in &prepared {
self.participants[*branch]
.xa
.xa_commit(xid, XaFlags::NOFLAGS)
.map_err(|e| map_xa_error(&e))?;
}
Ok(TxnOutcome::Committed {
operations: batch.ops.len(),
})
}
fn run_branch_work(
participant: &XaParticipant,
xid: &Xid,
op_indices: &[usize],
ops: &[TxnOp],
) -> Result<(), TxnStoreError> {
participant
.xa
.xa_start(xid, XaFlags::NOFLAGS)
.map_err(|e| map_xa_error(&e))?;
for &op_idx in op_indices {
participant
.apply(xid, &ops[op_idx])
.map_err(|e| map_txn_store_error(&e))?;
}
participant
.xa
.mark_write(xid)
.map_err(|e| map_xa_error(&e))?;
participant
.xa
.xa_end(xid, XaFlags::TMSUCCESS)
.map_err(|e| map_xa_error(&e))?;
Ok(())
}
fn rollback_active(&self, branches: &[(usize, Xid)]) {
for (branch, xid) in branches {
let _ = self.participants[*branch]
.xa
.xa_rollback(xid, XaFlags::NOFLAGS);
}
}
fn rollback_remaining(
&self,
active: &[(usize, Xid)],
failed: usize,
prepared: &[(usize, Xid)],
) {
for (branch, xid) in active {
if *branch == failed {
continue;
}
if prepared.iter().any(|(b, _)| b == branch) {
continue;
}
let _ = self.participants[*branch]
.xa
.xa_rollback(xid, XaFlags::NOFLAGS);
}
}
}
fn map_xa_error(err: &noxu::xa::XaError) -> TxnStoreError {
let text = err.to_string();
let lowered = text.to_ascii_lowercase();
if lowered.contains("deadlock")
|| lowered.contains("lock timeout")
|| lowered.contains("conflict")
|| lowered.contains("would block")
{
TxnStoreError::Conflict(text)
} else {
TxnStoreError::Backend(text)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn route_by_key_parity(op: &TxnOp) -> usize {
usize::from(op.key().first().copied().unwrap_or(0) & 1)
}
fn two_branch_coordinator() -> (XaCoordinator, TempDir, TempDir) {
let d0 = TempDir::new().expect("tempdir0");
let d1 = TempDir::new().expect("tempdir1");
let p0 = XaParticipant::open(d0.path(), b"east".to_vec()).expect("open east");
let p1 = XaParticipant::open(d1.path(), b"west".to_vec()).expect("open west");
(XaCoordinator::new(vec![p0, p1]), d0, d1)
}
fn put(key: &[u8], value: &[u8]) -> TxnOp {
TxnOp::Put {
bucket: b"u".to_vec(),
key: key.to_vec(),
value: value.to_vec(),
indexes: vec![],
}
}
#[test]
fn commit_spans_two_branches() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch {
ops: vec![put(b"alice", b"a"), put(b"bob", b"b")],
force_abort: false,
};
let outcome = coord.execute(&batch, route_by_key_parity).expect("commit");
assert_eq!(outcome, TxnOutcome::Committed { operations: 2 });
let east = coord.participant(0).expect("east");
let west = coord.participant(1).expect("west");
assert_eq!(
east.get_object(b"u", b"bob").unwrap().as_deref(),
Some(&b"b"[..])
);
assert!(east.get_object(b"u", b"alice").unwrap().is_none());
assert_eq!(
west.get_object(b"u", b"alice").unwrap().as_deref(),
Some(&b"a"[..])
);
assert!(west.get_object(b"u", b"bob").unwrap().is_none());
}
#[test]
fn force_abort_rolls_back_every_branch() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch {
ops: vec![put(b"alice", b"a"), put(b"bob", b"b")],
force_abort: true,
};
let outcome = coord.execute(&batch, route_by_key_parity).expect("abort");
assert!(matches!(outcome, TxnOutcome::Aborted { .. }));
let east = coord.participant(0).expect("east");
let west = coord.participant(1).expect("west");
assert!(east.get_object(b"u", b"bob").unwrap().is_none());
assert!(west.get_object(b"u", b"alice").unwrap().is_none());
}
#[test]
fn empty_batch_is_rejected() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch::default();
assert!(matches!(
coord.execute(&batch, route_by_key_parity),
Err(TxnStoreError::EmptyBatch)
));
}
#[test]
fn coordinator_accessors_reflect_participants() {
let (coord, _d0, _d1) = two_branch_coordinator();
assert_eq!(coord.len(), 2);
assert!(!coord.is_empty());
assert!(coord.participant(0).is_some());
assert!(coord.participant(2).is_none());
let empty = XaCoordinator::new(vec![]);
assert!(empty.is_empty());
assert_eq!(empty.len(), 0);
}
#[test]
fn delete_op_commits_through_branch_work() {
let (coord, _d0, _d1) = two_branch_coordinator();
coord
.execute(
&TxnBatch {
ops: vec![put(b"alice", b"a")],
force_abort: false,
},
route_by_key_parity,
)
.expect("put");
let del = TxnBatch {
ops: vec![TxnOp::Delete {
bucket: b"u".to_vec(),
key: b"alice".to_vec(),
}],
force_abort: false,
};
let outcome = coord.execute(&del, route_by_key_parity).expect("delete");
assert_eq!(outcome, TxnOutcome::Committed { operations: 1 });
let west = coord.participant(1).expect("west");
assert!(west.get_object(b"u", b"alice").unwrap().is_none());
}
#[test]
fn map_xa_error_classifies_conflict_vs_backend() {
use noxu::xa::XaError;
for needle in [
"deadlock detected",
"lock timeout",
"write conflict",
"would block",
] {
let err = XaError::Protocol(needle.to_string());
assert!(
matches!(map_xa_error(&err), TxnStoreError::Conflict(_)),
"{needle} should map to Conflict"
);
}
let other = XaError::Protocol("branch in bad state".to_string());
assert!(matches!(map_xa_error(&other), TxnStoreError::Backend(_)));
}
#[test]
fn phase_one_apply_failure_rolls_back_started_branches() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch {
ops: vec![
put(b"bob", b"b"), TxnOp::Put {
bucket: b"u\0bad".to_vec(), key: b"alice".to_vec(), value: b"a".to_vec(),
indexes: vec![],
},
],
force_abort: false,
};
let err = coord
.execute(&batch, route_by_key_parity)
.expect_err("phase-1 apply failure");
assert!(matches!(err, TxnStoreError::Backend(_)));
assert!(coord
.participant(0)
.unwrap()
.get_object(b"u", b"bob")
.unwrap()
.is_none());
}
#[test]
fn out_of_range_route_is_a_backend_error() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch {
ops: vec![put(b"alice", b"a")],
force_abort: false,
};
let err = coord.execute(&batch, |_| 99).expect_err("out of range");
assert!(matches!(err, TxnStoreError::Backend(_)));
assert!(coord
.participant(0)
.unwrap()
.get_object(b"u", b"alice")
.unwrap()
.is_none());
}
#[test]
fn commit_then_read_after_reopen_is_durable() {
let (coord, _d0, _d1) = two_branch_coordinator();
let batch = TxnBatch {
ops: vec![put(b"carol", b"c")], force_abort: false,
};
coord.execute(&batch, route_by_key_parity).expect("commit");
let west = coord.participant(1).expect("west");
assert_eq!(
west.get_object(b"u", b"carol").unwrap().as_deref(),
Some(&b"c"[..])
);
}
}