use std::collections::{BTreeMap, VecDeque};
use std::fmt;
use std::future::Future;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use mongreldb_consensus::engine_sink::EngineApplySink;
use mongreldb_consensus::error::ConsensusError;
use mongreldb_consensus::group::{ConsensusGroup, GroupCommitReceipt, GroupConfig};
use mongreldb_consensus::identity::{CommandKind, RaftNodeId, ReplicatedCommand};
use mongreldb_consensus::network::RaftTransport;
use mongreldb_consensus::read::{ReadConsistency, ReadConsistencyError, SessionToken};
use mongreldb_consensus::state_machine::{AppliedCommand, ApplySink, StateMachineError};
use mongreldb_log::commit_log::{DurabilityLevel, ExecutionControl, LogPosition};
use mongreldb_log::envelope::CommandEnvelope;
use mongreldb_types::hlc::{HlcClock, HlcTimestamp};
use mongreldb_types::ids::{RaftGroupId, SchemaVersion, TabletId, TransactionId};
use crate::meta::TxnStatusPartition;
use crate::tablet::TabletDescriptor;
pub const COMMAND_TYPE_DIST_TXN_COORDINATOR: u32 = 5;
pub const COMMAND_TYPE_DIST_TXN_INTENT: u32 = 6;
pub const DIST_TXN_RECORD_FORMAT_VERSION: u32 = 1;
pub const MIN_SUPPORTED_DIST_TXN_RECORD_FORMAT_VERSION: u32 = 1;
pub const DIST_TXN_CHECKPOINT_FORMAT_VERSION: u32 = 1;
pub const MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION: u32 = 1;
pub const TXN_STATUS_CHECKPOINT_FILENAME: &str = "dist-txn-status.json";
pub const INTENT_CHECKPOINT_FILENAME: &str = "dist-txn-intents.json";
pub const DIST_TXN_REJECTION_LIMIT: usize = 256;
pub const DEFAULT_PENDING_TIMEOUT: Duration = Duration::from_secs(10);
pub const DEFAULT_RESOLVED_RETENTION: Duration = Duration::from_secs(24 * 60 * 60);
pub const DEFAULT_SWEEP_LIMIT: u32 = 1024;
const TAG_BEGIN: &str = "dist-txn/begin";
const TAG_HEARTBEAT: &str = "dist-txn/heartbeat";
const TAG_MARK: &str = "dist-txn/mark-preparing";
const TAG_COMMIT: &str = "dist-txn/commit";
const TAG_ABORT: &str = "dist-txn/abort";
const TAG_PREPARE: &str = "dist-txn/intents";
const TAG_RESOLVE: &str = "dist-txn/resolve";
const TAG_SWEEP: &str = "dist-txn/sweep-resolved";
fn command_id_for(tag: &str, txn_id: &TransactionId, extra: &[u8]) -> [u8; 16] {
let mut hasher = Sha256::new();
hasher.update(tag.as_bytes());
hasher.update(txn_id.as_bytes());
hasher.update(extra);
let digest = hasher.finalize();
digest[..16].try_into().expect("sha256 digest is 32 bytes")
}
fn post_decision_after_hook(txn_id: TransactionId) -> Result<(), DistTxnError> {
mongreldb_fault::inject("txn.decision.after").map_err(|fault| DistTxnError::OutcomeAmbiguous {
txn_id,
detail: format!("txn.decision.after: {fault}"),
})
}
fn fnv1a_64(bytes: &[u8]) -> u64 {
let mut hash: u64 = 0xcbf29ce484222325;
for &byte in bytes {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
#[derive(Debug, thiserror::Error)]
pub enum DistTxnError {
#[error(transparent)]
Consensus(#[from] ConsensusError),
#[error(transparent)]
Read(#[from] ReadConsistencyError),
#[error("distributed transaction command encoding failed: {0}")]
Encode(String),
#[error("invalid distributed transaction request: {0}")]
InvalidRequest(String),
#[error("distributed transaction group I/O error: {0}")]
Io(#[from] std::io::Error),
#[error("corrupt distributed transaction checkpoint: {0}")]
CorruptCheckpoint(String),
#[error("the distributed transaction aborted: {0:?}")]
Aborted(AbortReason),
#[error("prepare refused by the participant: {0}")]
PrepareRejected(PrepareRejectionReason),
#[error("transaction {0} already exists under a different idempotency key")]
IdempotencyConflict(TransactionId),
#[error(
"transaction {txn_id} outcome is ambiguous ({detail}); \
re-run with the same transaction id and idempotency key"
)]
OutcomeAmbiguous {
txn_id: TransactionId,
detail: String,
},
#[error("no live group leader: {0}")]
Unavailable(String),
#[error("coordinator clock failure: {0}")]
Clock(String),
#[error(transparent)]
Fault(#[from] mongreldb_fault::Fault),
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum DistTxnDecodeError {
#[error("distributed transaction command decode failed: {0}")]
Malformed(String),
#[error(
"unsupported distributed transaction record version {found} (supported {min}..={max})"
)]
UnsupportedVersion {
found: u32,
min: u32,
max: u32,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AbortReason {
RolledBack,
Conflict(String),
Validation(String),
Cancelled(String),
Error(String),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum DistributedTxnState {
Pending,
Preparing,
Committed {
commit_ts: HlcTimestamp,
},
Aborted {
reason: AbortReason,
},
}
impl DistributedTxnState {
pub fn is_terminal(&self) -> bool {
matches!(
self,
DistributedTxnState::Committed { .. } | DistributedTxnState::Aborted { .. }
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TxnParticipant {
pub tablet_id: TabletId,
pub raft_group_id: RaftGroupId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TxnRecord {
pub txn_id: TransactionId,
pub state: DistributedTxnState,
pub participants: Vec<TxnParticipant>,
pub prepare_ts: BTreeMap<TabletId, HlcTimestamp>,
pub coordinator: RaftGroupId,
pub created_at: HlcTimestamp,
pub heartbeat: HlcTimestamp,
pub expiry: HlcTimestamp,
pub max_observed: HlcTimestamp,
pub idempotency_key: [u8; 16],
}
impl TxnRecord {
pub fn expired(&self, now: HlcTimestamp) -> bool {
now >= self.expiry
}
pub fn participant(&self, tablet_id: &TabletId) -> Option<&TxnParticipant> {
self.participants.iter().find(|p| p.tablet_id == *tablet_id)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WriteIntent {
pub txn_id: TransactionId,
pub key: Vec<u8>,
pub value_ref: Vec<u8>,
pub prepare_ts: HlcTimestamp,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PrepareToken {
pub txn_id: TransactionId,
pub tablet_id: TabletId,
pub raft_group_id: RaftGroupId,
pub prepare_ts: HlcTimestamp,
pub position: LogPosition,
pub command_id: [u8; 16],
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum TxnDecision {
Committed {
commit_ts: HlcTimestamp,
},
Aborted {
reason: AbortReason,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TxnOutcome {
pub txn_id: TransactionId,
pub commit_ts: HlcTimestamp,
pub participants: Vec<TxnParticipant>,
pub durability: DurabilityLevel,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ParticipantTxn {
pub txn_id: TransactionId,
pub prepare_ts: HlcTimestamp,
pub command_id: [u8; 16],
pub intents: Vec<WriteIntent>,
pub resolution: Option<TxnDecision>,
#[serde(default)]
pub applied: bool,
#[serde(default)]
pub resolved_at: Option<HlcTimestamp>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommittedWrite {
pub key: Vec<u8>,
pub value_ref: Vec<u8>,
pub commit_ts: HlcTimestamp,
pub txn_id: TransactionId,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CoordinatorSelection {
#[default]
TxnIdDerived,
FirstWriteHome,
}
pub fn txn_status_partition_index(txn_id: &TransactionId, partition_count: u32) -> u32 {
assert!(partition_count > 0, "partition count must be positive");
u32::try_from(fnv1a_64(txn_id.as_bytes()) % u64::from(partition_count))
.expect("remainder below partition count")
}
pub fn select_coordinator_group(
selection: &CoordinatorSelection,
txn_id: &TransactionId,
first_write_tablet: Option<&TabletDescriptor>,
partitions: &BTreeMap<u32, TxnStatusPartition>,
) -> Result<RaftGroupId, DistTxnError> {
match selection {
CoordinatorSelection::TxnIdDerived => {
if partitions.is_empty() {
return Err(DistTxnError::InvalidRequest(
"no transaction-status partitions are published in meta state".to_owned(),
));
}
let index = txn_status_partition_index(
txn_id,
u32::try_from(partitions.len()).unwrap_or(u32::MAX),
);
partitions
.values()
.nth(index as usize)
.map(|partition| partition.home_raft_group)
.ok_or_else(|| {
DistTxnError::InvalidRequest(format!(
"transaction-status partition index {index} is not published"
))
})
}
CoordinatorSelection::FirstWriteHome => first_write_tablet
.map(|tablet| tablet.raft_group_id)
.ok_or_else(|| {
DistTxnError::InvalidRequest(
"first-write-home coordinator selection requires the first write's tablet \
descriptor"
.to_owned(),
)
}),
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum CoordinatorCommand {
Begin {
record: TxnRecord,
},
Heartbeat {
txn_id: TransactionId,
heartbeat: HlcTimestamp,
expiry: HlcTimestamp,
},
MarkPreparing {
txn_id: TransactionId,
tablet_id: TabletId,
prepare_ts: HlcTimestamp,
observed: HlcTimestamp,
},
Commit {
txn_id: TransactionId,
commit_ts: HlcTimestamp,
},
Abort {
txn_id: TransactionId,
reason: AbortReason,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum IntentCommand {
SetTabletVersions {
schema_version: SchemaVersion,
authz_version: u64,
},
PersistIntents {
txn_id: TransactionId,
expected_schema_version: SchemaVersion,
expected_authz_version: u64,
prepare_ts: HlcTimestamp,
intents: Vec<WriteIntent>,
},
Resolve {
txn_id: TransactionId,
decision: TxnDecision,
},
SweepResolved {
older_than: HlcTimestamp,
limit: u32,
},
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CoordinatorCommandRecord {
pub format_version: u32,
pub command: CoordinatorCommand,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct IntentCommandRecord {
pub format_version: u32,
pub command: IntentCommand,
}
fn encode_record<T: Serialize>(record: &T) -> Result<Vec<u8>, DistTxnError> {
serde_json::to_vec(record).map_err(|error| DistTxnError::Encode(error.to_string()))
}
fn decode_record<T: for<'de> Deserialize<'de>>(payload: &[u8]) -> Result<T, DistTxnDecodeError> {
serde_json::from_slice(payload)
.map_err(|error| DistTxnDecodeError::Malformed(error.to_string()))
}
fn check_record_version(found: u32) -> Result<(), DistTxnDecodeError> {
if !(MIN_SUPPORTED_DIST_TXN_RECORD_FORMAT_VERSION..=DIST_TXN_RECORD_FORMAT_VERSION)
.contains(&found)
{
return Err(DistTxnDecodeError::UnsupportedVersion {
found,
min: MIN_SUPPORTED_DIST_TXN_RECORD_FORMAT_VERSION,
max: DIST_TXN_RECORD_FORMAT_VERSION,
});
}
Ok(())
}
impl CoordinatorCommandRecord {
pub fn new(command: CoordinatorCommand) -> Self {
CoordinatorCommandRecord {
format_version: DIST_TXN_RECORD_FORMAT_VERSION,
command,
}
}
pub fn encode(&self) -> Result<Vec<u8>, DistTxnError> {
encode_record(self)
}
pub fn decode(payload: &[u8]) -> Result<Self, DistTxnDecodeError> {
let record: CoordinatorCommandRecord = decode_record(payload)?;
check_record_version(record.format_version)?;
Ok(record)
}
}
impl IntentCommandRecord {
pub fn new(command: IntentCommand) -> Self {
IntentCommandRecord {
format_version: DIST_TXN_RECORD_FORMAT_VERSION,
command,
}
}
pub fn encode(&self) -> Result<Vec<u8>, DistTxnError> {
encode_record(self)
}
pub fn decode(payload: &[u8]) -> Result<Self, DistTxnDecodeError> {
let record: IntentCommandRecord = decode_record(payload)?;
check_record_version(record.format_version)?;
Ok(record)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
pub enum StatusRejectionReason {
#[error("transaction already exists under a different idempotency key")]
IdempotencyKeyConflict {
existing: [u8; 16],
},
#[error("unknown transaction {0}")]
UnknownTxn(TransactionId),
#[error("the transaction already reached a final decision: {existing:?}")]
DecisionFinal {
existing: DistributedTxnState,
},
#[error("tablet {0} is not a participant of the transaction")]
UnknownParticipant(TabletId),
#[error(
"invalid commit_ts {commit_ts}: must be strictly greater than max_observed {max_observed} \
and every prepare timestamp"
)]
InvalidCommitTs {
commit_ts: HlcTimestamp,
max_observed: HlcTimestamp,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StatusRejection {
pub position: LogPosition,
pub command_id: Option<[u8; 16]>,
pub txn_id: TransactionId,
pub reason: StatusRejectionReason,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
pub enum PrepareRejectionReason {
#[error("key conflict: another transaction holds a write intent on the key")]
KeyConflict {
key: Vec<u8>,
holder: TransactionId,
},
#[error("the transaction already prepared with a different write set")]
PayloadMismatch,
#[error("the transaction is already resolved: {decision:?}")]
AlreadyResolved {
decision: TxnDecision,
},
#[error("stale schema version {expected:?} (tablet is at {found:?})")]
StaleSchemaVersion {
expected: SchemaVersion,
found: SchemaVersion,
},
#[error("stale authorization version {expected} (tablet is at {found})")]
StaleAuthzVersion {
expected: u64,
found: u64,
},
#[error("staged write is not appliable: {detail}")]
MalformedStagedWrite {
detail: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PrepareRejection {
pub position: LogPosition,
pub command_id: Option<[u8; 16]>,
pub txn_id: TransactionId,
pub reason: PrepareRejectionReason,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TxnStatusState {
pub records: BTreeMap<TransactionId, TxnRecord>,
pub rejections: VecDeque<StatusRejection>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TxnStatusCheckpoint {
format_version: u32,
position: LogPosition,
command_id: Option<[u8; 16]>,
state: TxnStatusState,
}
pub struct TxnStatusApplySink {
state: TxnStatusState,
position: LogPosition,
command_id: Option<[u8; 16]>,
state_dir: std::path::PathBuf,
}
impl TxnStatusApplySink {
pub fn open(group_dir: &Path) -> Result<Self, DistTxnError> {
let state_dir = group_dir.join("raft").join("state");
std::fs::create_dir_all(&state_dir).map_err(DistTxnError::Io)?;
let checkpoint_path = state_dir.join(TXN_STATUS_CHECKPOINT_FILENAME);
let Some(bytes) =
crate::node::read_meta_file(&checkpoint_path).map_err(|error| match error {
crate::node::ClusterError::Io(error) => DistTxnError::Io(error),
other => DistTxnError::CorruptCheckpoint(other.to_string()),
})?
else {
return Ok(TxnStatusApplySink {
state: TxnStatusState::default(),
position: LogPosition::ZERO,
command_id: None,
state_dir,
});
};
let checkpoint: TxnStatusCheckpoint = serde_json::from_slice(&bytes)
.map_err(|error| DistTxnError::CorruptCheckpoint(format!("decode: {error}")))?;
if !(MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION..=DIST_TXN_CHECKPOINT_FORMAT_VERSION)
.contains(&checkpoint.format_version)
{
return Err(DistTxnError::CorruptCheckpoint(format!(
"unsupported format version {} (supported \
{MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION}..=\
{DIST_TXN_CHECKPOINT_FORMAT_VERSION})",
checkpoint.format_version
)));
}
Ok(TxnStatusApplySink {
state: checkpoint.state,
position: checkpoint.position,
command_id: checkpoint.command_id,
state_dir,
})
}
pub fn state(&self) -> &TxnStatusState {
&self.state
}
pub fn applied_position(&self) -> LogPosition {
self.position
}
pub fn record(&self, txn_id: &TransactionId) -> Option<&TxnRecord> {
self.state.records.get(txn_id)
}
fn checkpoint(&self) -> TxnStatusCheckpoint {
TxnStatusCheckpoint {
format_version: DIST_TXN_CHECKPOINT_FORMAT_VERSION,
position: self.position,
command_id: self.command_id,
state: self.state.clone(),
}
}
fn persist(&self) -> Result<(), StateMachineError> {
let bytes = serde_json::to_vec(&self.checkpoint()).map_err(|error| {
StateMachineError::Sink(format!("txn status checkpoint encode: {error}"))
})?;
crate::node::write_meta_atomic(&self.state_dir, TXN_STATUS_CHECKPOINT_FILENAME, &bytes)
.map_err(|error| {
StateMachineError::Sink(format!("txn status checkpoint write: {error}"))
})
}
fn journal(
&mut self,
command: &AppliedCommand,
txn_id: TransactionId,
reason: StatusRejectionReason,
) {
self.state.rejections.push_back(StatusRejection {
position: command.position,
command_id: command.command_id(),
txn_id,
reason,
});
while self.state.rejections.len() > DIST_TXN_REJECTION_LIMIT {
self.state.rejections.pop_front();
}
}
fn apply_command(
&mut self,
command: &AppliedCommand,
transition: &CoordinatorCommand,
) -> Result<(), StateMachineError> {
match transition {
CoordinatorCommand::Begin { record } => {
if record.state != DistributedTxnState::Pending {
return Err(StateMachineError::Corrupt(
"Begin record must be Pending".to_owned(),
));
}
match self.state.records.get(&record.txn_id) {
Some(existing) => {
if existing.idempotency_key == record.idempotency_key {
return Ok(());
}
let existing_key = existing.idempotency_key;
self.journal(
command,
record.txn_id,
StatusRejectionReason::IdempotencyKeyConflict {
existing: existing_key,
},
);
}
None => {
self.state.records.insert(record.txn_id, record.clone());
}
}
Ok(())
}
CoordinatorCommand::Heartbeat {
txn_id,
heartbeat,
expiry,
} => match self.state.records.get_mut(txn_id) {
None => {
self.journal(command, *txn_id, StatusRejectionReason::UnknownTxn(*txn_id));
Ok(())
}
Some(record) if record.state.is_terminal() => {
let existing = record.state.clone();
self.journal(
command,
*txn_id,
StatusRejectionReason::DecisionFinal { existing },
);
Ok(())
}
Some(record) => {
if *heartbeat > record.heartbeat {
record.heartbeat = *heartbeat;
record.expiry = *expiry;
}
Ok(())
}
},
CoordinatorCommand::MarkPreparing {
txn_id,
tablet_id,
prepare_ts,
observed,
} => match self.state.records.get_mut(txn_id) {
None => {
self.journal(command, *txn_id, StatusRejectionReason::UnknownTxn(*txn_id));
Ok(())
}
Some(record) if record.state.is_terminal() => {
let existing = record.state.clone();
self.journal(
command,
*txn_id,
StatusRejectionReason::DecisionFinal { existing },
);
Ok(())
}
Some(record) => {
if record.participant(tablet_id).is_none() {
self.journal(
command,
*txn_id,
StatusRejectionReason::UnknownParticipant(*tablet_id),
);
return Ok(());
}
record.prepare_ts.insert(*tablet_id, *prepare_ts);
record.max_observed = record.max_observed.max(*observed);
record.state = DistributedTxnState::Preparing;
Ok(())
}
},
CoordinatorCommand::Commit { txn_id, commit_ts } => {
let pre = self.state.records.get(txn_id).map(|record| {
(
record.state.is_terminal(),
record.state.clone(),
record.max_observed,
record.prepare_ts.values().any(|ts| *commit_ts <= *ts)
|| *commit_ts <= record.max_observed,
)
});
match pre {
None => {
self.journal(command, *txn_id, StatusRejectionReason::UnknownTxn(*txn_id));
Ok(())
}
Some((true, existing, _, _)) => {
self.journal(
command,
*txn_id,
StatusRejectionReason::DecisionFinal { existing },
);
Ok(())
}
Some((false, _, max_observed, true)) => {
self.journal(
command,
*txn_id,
StatusRejectionReason::InvalidCommitTs {
commit_ts: *commit_ts,
max_observed,
},
);
Ok(())
}
Some((false, _, _, false)) => {
if let Some(record) = self.state.records.get_mut(txn_id) {
record.state = DistributedTxnState::Committed {
commit_ts: *commit_ts,
};
}
Ok(())
}
}
}
CoordinatorCommand::Abort { txn_id, reason } => {
match self.state.records.get_mut(txn_id) {
None => {
self.journal(command, *txn_id, StatusRejectionReason::UnknownTxn(*txn_id));
Ok(())
}
Some(record) if record.state.is_terminal() => {
let existing = record.state.clone();
self.journal(
command,
*txn_id,
StatusRejectionReason::DecisionFinal { existing },
);
Ok(())
}
Some(record) => {
record.state = DistributedTxnState::Aborted {
reason: reason.clone(),
};
Ok(())
}
}
}
}
}
}
impl ApplySink for TxnStatusApplySink {
fn apply(&mut self, command: &AppliedCommand) -> Result<(), StateMachineError> {
if command.position.index <= self.position.index {
return Ok(());
}
match &command.command {
ReplicatedCommand::Transaction(transaction) => {
transaction.envelope.verify().map_err(|error| {
StateMachineError::Corrupt(format!("txn status envelope: {error}"))
})?;
if transaction.envelope.command_type != COMMAND_TYPE_DIST_TXN_COORDINATOR {
return Err(StateMachineError::Corrupt(format!(
"txn status command_type {} is not COMMAND_TYPE_DIST_TXN_COORDINATOR",
transaction.envelope.command_type
)));
}
let record = CoordinatorCommandRecord::decode(&transaction.envelope.payload)
.map_err(|error| StateMachineError::Corrupt(error.to_string()))?;
self.apply_command(command, &record.command)?;
}
ReplicatedCommand::Maintenance(_) | ReplicatedCommand::Noop => {}
ReplicatedCommand::Catalog(_) => {
return Err(StateMachineError::Corrupt(
"catalog command on a transaction-status group".to_owned(),
));
}
}
self.position = command.position;
if let Some(command_id) = command.command_id() {
self.command_id = Some(command_id);
}
self.persist()
}
fn snapshot(&self) -> Result<Vec<u8>, StateMachineError> {
serde_json::to_vec(&self.checkpoint()).map_err(|error| {
StateMachineError::Sink(format!("txn status snapshot encode: {error}"))
})
}
fn install(&mut self, data: &[u8]) -> Result<(), StateMachineError> {
let checkpoint: TxnStatusCheckpoint = serde_json::from_slice(data).map_err(|error| {
StateMachineError::Corrupt(format!("txn status snapshot decode: {error}"))
})?;
if !(MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION..=DIST_TXN_CHECKPOINT_FORMAT_VERSION)
.contains(&checkpoint.format_version)
{
return Err(StateMachineError::Corrupt(format!(
"unsupported txn status checkpoint format version {} (supported \
{MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION}..=\
{DIST_TXN_CHECKPOINT_FORMAT_VERSION})",
checkpoint.format_version
)));
}
self.state = checkpoint.state;
self.position = checkpoint.position;
self.command_id = checkpoint.command_id;
self.persist()
}
}
impl fmt::Debug for TxnStatusApplySink {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TxnStatusApplySink")
.field("records", &self.state.records.len())
.finish()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct IntentState {
pub schema_version: SchemaVersion,
pub authz_version: u64,
pub txns: BTreeMap<TransactionId, ParticipantTxn>,
pub committed_writes: Vec<CommittedWrite>,
pub rejections: VecDeque<PrepareRejection>,
}
impl Default for IntentState {
fn default() -> Self {
IntentState {
schema_version: SchemaVersion::ZERO,
authz_version: 0,
txns: BTreeMap::new(),
committed_writes: Vec::new(),
rejections: VecDeque::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct IntentCheckpoint {
format_version: u32,
position: LogPosition,
command_id: Option<[u8; 16]>,
state: IntentState,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CompositeSnapshot {
format_version: u32,
engine: Vec<u8>,
intent: Vec<u8>,
}
pub const COMPOSITE_SNAPSHOT_FORMAT_VERSION: u32 = 1;
pub struct IntentApplySink {
state: IntentState,
position: LogPosition,
command_id: Option<[u8; 16]>,
state_dir: std::path::PathBuf,
engine: Option<Arc<Mutex<EngineApplySink>>>,
}
fn same_write_set(a: &[WriteIntent], b: &[WriteIntent]) -> bool {
a.len() == b.len()
&& a.iter()
.zip(b.iter())
.all(|(left, right)| left.key == right.key && left.value_ref == right.value_ref)
}
impl IntentApplySink {
pub fn open(group_dir: &Path) -> Result<Self, DistTxnError> {
Self::open_with_binding(group_dir, None)
}
pub fn open_with_engine(
group_dir: &Path,
engine: Arc<Mutex<EngineApplySink>>,
) -> Result<Self, DistTxnError> {
Self::open_with_binding(group_dir, Some(engine))
}
fn open_with_binding(
group_dir: &Path,
engine: Option<Arc<Mutex<EngineApplySink>>>,
) -> Result<Self, DistTxnError> {
let state_dir = group_dir.join("raft").join("state");
std::fs::create_dir_all(&state_dir).map_err(DistTxnError::Io)?;
let checkpoint_path = state_dir.join(INTENT_CHECKPOINT_FILENAME);
let Some(bytes) =
crate::node::read_meta_file(&checkpoint_path).map_err(|error| match error {
crate::node::ClusterError::Io(error) => DistTxnError::Io(error),
other => DistTxnError::CorruptCheckpoint(other.to_string()),
})?
else {
return Ok(IntentApplySink {
state: IntentState::default(),
position: LogPosition::ZERO,
command_id: None,
state_dir,
engine,
});
};
let checkpoint: IntentCheckpoint = serde_json::from_slice(&bytes)
.map_err(|error| DistTxnError::CorruptCheckpoint(format!("decode: {error}")))?;
if !(MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION..=DIST_TXN_CHECKPOINT_FORMAT_VERSION)
.contains(&checkpoint.format_version)
{
return Err(DistTxnError::CorruptCheckpoint(format!(
"unsupported format version {} (supported \
{MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION}..=\
{DIST_TXN_CHECKPOINT_FORMAT_VERSION})",
checkpoint.format_version
)));
}
Ok(IntentApplySink {
state: checkpoint.state,
position: checkpoint.position,
command_id: checkpoint.command_id,
state_dir,
engine,
})
}
pub fn state(&self) -> &IntentState {
&self.state
}
pub fn applied_position(&self) -> LogPosition {
self.position
}
pub fn txn(&self, txn_id: &TransactionId) -> Option<&ParticipantTxn> {
self.state.txns.get(txn_id)
}
pub fn unresolved_txn_ids(&self) -> Vec<TransactionId> {
self.state
.txns
.values()
.filter(|txn| txn.resolution.is_none())
.map(|txn| txn.txn_id)
.collect()
}
pub fn engine(&self) -> Option<&Arc<Mutex<EngineApplySink>>> {
self.engine.as_ref()
}
fn checkpoint(&self) -> IntentCheckpoint {
IntentCheckpoint {
format_version: DIST_TXN_CHECKPOINT_FORMAT_VERSION,
position: self.position,
command_id: self.command_id,
state: self.state.clone(),
}
}
fn persist(&self) -> Result<(), StateMachineError> {
let bytes = serde_json::to_vec(&self.checkpoint()).map_err(|error| {
StateMachineError::Sink(format!("intent checkpoint encode: {error}"))
})?;
crate::node::write_meta_atomic(&self.state_dir, INTENT_CHECKPOINT_FILENAME, &bytes)
.map_err(|error| StateMachineError::Sink(format!("intent checkpoint write: {error}")))
}
fn journal(
&mut self,
command: &AppliedCommand,
txn_id: TransactionId,
reason: PrepareRejectionReason,
) {
self.state.rejections.push_back(PrepareRejection {
position: command.position,
command_id: command.command_id(),
txn_id,
reason,
});
while self.state.rejections.len() > DIST_TXN_REJECTION_LIMIT {
self.state.rejections.pop_front();
}
}
fn apply_committed_to_engine(
engine: &Arc<Mutex<EngineApplySink>>,
txn_id: &TransactionId,
staged: &[Vec<u8>],
commit_ts: HlcTimestamp,
) -> Result<(), StateMachineError> {
if staged.is_empty() {
return Ok(());
}
let db = engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".to_owned()))?
.database()
.ok_or_else(|| {
StateMachineError::Sink("engine sink has no open database".to_owned())
})?;
let txn_tag = fnv1a_64(txn_id.as_bytes());
db.apply_staged_txn_writes(txn_tag, staged, commit_ts)
.map_err(|error| {
StateMachineError::Sink(format!("engine apply of committed resolution: {error}"))
})?;
Ok(())
}
fn conflicting_holder(&self, txn_id: &TransactionId, key: &[u8]) -> Option<TransactionId> {
self.state
.txns
.values()
.find(|txn| {
txn.txn_id != *txn_id
&& txn.resolution.is_none()
&& txn.intents.iter().any(|intent| intent.key == key)
})
.map(|txn| txn.txn_id)
}
fn apply_persist(
&mut self,
command: &AppliedCommand,
txn_id: TransactionId,
expected_schema_version: SchemaVersion,
expected_authz_version: u64,
prepare_ts: HlcTimestamp,
intents: &[WriteIntent],
) -> Result<(), StateMachineError> {
if let Some(existing) = self.state.txns.get(&txn_id) {
if let Some(decision) = &existing.resolution {
let decision = decision.clone();
self.journal(
command,
txn_id,
PrepareRejectionReason::AlreadyResolved { decision },
);
return Ok(());
}
if same_write_set(&existing.intents, intents) {
return Ok(());
}
self.journal(command, txn_id, PrepareRejectionReason::PayloadMismatch);
return Ok(());
}
if expected_schema_version != self.state.schema_version {
let found = self.state.schema_version;
self.journal(
command,
txn_id,
PrepareRejectionReason::StaleSchemaVersion {
expected: expected_schema_version,
found,
},
);
return Ok(());
}
if expected_authz_version != self.state.authz_version {
let found = self.state.authz_version;
self.journal(
command,
txn_id,
PrepareRejectionReason::StaleAuthzVersion {
expected: expected_authz_version,
found,
},
);
return Ok(());
}
for intent in intents {
if let Some(holder) = self.conflicting_holder(&txn_id, &intent.key) {
self.journal(
command,
txn_id,
PrepareRejectionReason::KeyConflict {
key: intent.key.clone(),
holder,
},
);
return Ok(());
}
}
if let Some(engine) = &self.engine {
let staged: Vec<Vec<u8>> = intents
.iter()
.map(|intent| intent.value_ref.clone())
.collect();
let db = engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".to_owned()))?
.database()
.ok_or_else(|| {
StateMachineError::Sink("engine sink has no open database".to_owned())
})?;
if let Err(error) = db.validate_staged_txn_writes(&staged) {
self.journal(
command,
txn_id,
PrepareRejectionReason::MalformedStagedWrite {
detail: error.to_string(),
},
);
return Ok(());
}
}
self.state.txns.insert(
txn_id,
ParticipantTxn {
txn_id,
prepare_ts,
command_id: command.command_id().unwrap_or([0u8; 16]),
intents: intents.to_vec(),
resolution: None,
applied: false,
resolved_at: None,
},
);
Ok(())
}
fn apply_resolve(
&mut self,
command: &AppliedCommand,
txn_id: TransactionId,
decision: &TxnDecision,
) -> Result<(), StateMachineError> {
let engine = self.engine.clone();
let resolved_at = command.commit_ts();
match self.state.txns.get_mut(&txn_id) {
None => {
self.state.txns.insert(
txn_id,
ParticipantTxn {
txn_id,
prepare_ts: HlcTimestamp::ZERO,
command_id: [0u8; 16],
intents: Vec::new(),
resolution: Some(decision.clone()),
applied: true,
resolved_at,
},
);
Ok(())
}
Some(existing) => match (&existing.resolution, decision) {
(Some(prior), new) if *prior == *new => {
if !existing.applied {
if let (Some(engine), TxnDecision::Committed { commit_ts }) =
(&engine, prior)
{
let staged: Vec<Vec<u8>> = self
.state
.committed_writes
.iter()
.filter(|write| write.txn_id == txn_id)
.map(|write| write.value_ref.clone())
.collect();
Self::apply_committed_to_engine(engine, &txn_id, &staged, *commit_ts)?;
}
if let Some(existing) = self.state.txns.get_mut(&txn_id) {
existing.applied = true;
}
}
Ok(())
}
(Some(prior), new) => Err(StateMachineError::Corrupt(format!(
"conflicting resolve decisions for transaction {txn_id}: \
applied {prior:?}, got {new:?}"
))),
(None, TxnDecision::Committed { commit_ts }) => {
let commit_ts = *commit_ts;
if let Some(engine) = &engine {
let staged: Vec<Vec<u8>> = existing
.intents
.iter()
.map(|intent| intent.value_ref.clone())
.collect();
Self::apply_committed_to_engine(engine, &txn_id, &staged, commit_ts)?;
}
let intents = std::mem::take(&mut existing.intents);
let writes = intents.into_iter().map(|intent| CommittedWrite {
key: intent.key,
value_ref: intent.value_ref,
commit_ts,
txn_id,
});
existing.resolution = Some(decision.clone());
existing.applied = true;
existing.resolved_at = resolved_at;
self.state.committed_writes.extend(writes);
Ok(())
}
(None, TxnDecision::Aborted { .. }) => {
existing.intents.clear();
existing.resolution = Some(decision.clone());
existing.applied = true;
existing.resolved_at = resolved_at;
Ok(())
}
},
}
}
fn apply_sweep(&mut self, older_than: HlcTimestamp, limit: u32) {
let mut swept = Vec::new();
for (txn_id, txn) in &self.state.txns {
if swept.len() >= limit as usize {
break;
}
if txn.resolution.is_some() && txn.resolved_at.is_some_and(|at| at < older_than) {
swept.push(*txn_id);
}
}
for txn_id in &swept {
self.state.txns.remove(txn_id);
}
self.state
.committed_writes
.retain(|write| !swept.contains(&write.txn_id));
}
}
impl ApplySink for IntentApplySink {
fn apply(&mut self, command: &AppliedCommand) -> Result<(), StateMachineError> {
if command.position.index <= self.position.index {
return Ok(());
}
match &command.command {
ReplicatedCommand::Transaction(transaction) => {
transaction.envelope.verify().map_err(|error| {
StateMachineError::Corrupt(format!("intent envelope: {error}"))
})?;
if transaction.envelope.command_type == COMMAND_TYPE_DIST_TXN_INTENT {
let record = IntentCommandRecord::decode(&transaction.envelope.payload)
.map_err(|error| StateMachineError::Corrupt(error.to_string()))?;
match record.command {
IntentCommand::SetTabletVersions {
schema_version,
authz_version,
} => {
self.state.schema_version = schema_version;
self.state.authz_version = authz_version;
}
IntentCommand::PersistIntents {
txn_id,
expected_schema_version,
expected_authz_version,
prepare_ts,
intents,
} => {
self.apply_persist(
command,
txn_id,
expected_schema_version,
expected_authz_version,
prepare_ts,
&intents,
)?;
}
IntentCommand::Resolve { txn_id, decision } => {
self.apply_resolve(command, txn_id, &decision)?;
}
IntentCommand::SweepResolved { older_than, limit } => {
self.apply_sweep(older_than, limit);
}
}
} else if let Some(engine) = &self.engine {
engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".into()))?
.apply(command)?;
} else {
return Err(StateMachineError::Corrupt(format!(
"intent command_type {} is not COMMAND_TYPE_DIST_TXN_INTENT",
transaction.envelope.command_type
)));
}
}
ReplicatedCommand::Catalog(_) => {
if let Some(engine) = &self.engine {
engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".into()))?
.apply(command)?;
} else {
return Err(StateMachineError::Corrupt(
"catalog command on a participant intent group".to_owned(),
));
}
}
ReplicatedCommand::Maintenance(_) | ReplicatedCommand::Noop => {
if let Some(engine) = &self.engine {
engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".into()))?
.apply(command)?;
}
}
}
self.position = command.position;
if let Some(command_id) = command.command_id() {
self.command_id = Some(command_id);
}
self.persist()
}
fn snapshot(&self) -> Result<Vec<u8>, StateMachineError> {
let intent = serde_json::to_vec(&self.checkpoint())
.map_err(|error| StateMachineError::Sink(format!("intent snapshot encode: {error}")))?;
match &self.engine {
None => Ok(intent),
Some(engine) => {
let engine_image = engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".into()))?
.snapshot()?;
serde_json::to_vec(&CompositeSnapshot {
format_version: COMPOSITE_SNAPSHOT_FORMAT_VERSION,
engine: engine_image,
intent,
})
.map_err(|error| {
StateMachineError::Sink(format!("composite snapshot encode: {error}"))
})
}
}
}
fn install(&mut self, data: &[u8]) -> Result<(), StateMachineError> {
let intent_bytes;
if let Some(engine) = &self.engine {
let composite: CompositeSnapshot = serde_json::from_slice(data).map_err(|error| {
StateMachineError::Corrupt(format!("composite snapshot decode: {error}"))
})?;
if composite.format_version != COMPOSITE_SNAPSHOT_FORMAT_VERSION {
return Err(StateMachineError::Corrupt(format!(
"unsupported composite snapshot format version {} (supported \
{COMPOSITE_SNAPSHOT_FORMAT_VERSION})",
composite.format_version
)));
}
engine
.lock()
.map_err(|_| StateMachineError::Sink("engine sink lock poisoned".into()))?
.install(&composite.engine)?;
intent_bytes = composite.intent;
} else {
intent_bytes = data.to_vec();
}
let checkpoint: IntentCheckpoint =
serde_json::from_slice(&intent_bytes).map_err(|error| {
StateMachineError::Corrupt(format!("intent snapshot decode: {error}"))
})?;
if !(MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION..=DIST_TXN_CHECKPOINT_FORMAT_VERSION)
.contains(&checkpoint.format_version)
{
return Err(StateMachineError::Corrupt(format!(
"unsupported intent checkpoint format version {} (supported \
{MIN_SUPPORTED_DIST_TXN_CHECKPOINT_FORMAT_VERSION}..=\
{DIST_TXN_CHECKPOINT_FORMAT_VERSION})",
checkpoint.format_version
)));
}
self.state = checkpoint.state;
self.position = checkpoint.position;
self.command_id = checkpoint.command_id;
self.persist()
}
}
impl fmt::Debug for IntentApplySink {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IntentApplySink")
.field("txns", &self.state.txns.len())
.finish()
}
}
fn basic_node<N>(address: &str) -> Result<N, DistTxnError>
where
N: for<'de> Deserialize<'de>,
{
serde_json::from_value(serde_json::json!({ "addr": address })).map_err(|error| {
DistTxnError::InvalidRequest(format!("member address `{address}`: {error}"))
})
}
pub trait GroupMember {
fn node_id(&self) -> RaftNodeId;
fn current_leader(&self) -> Option<RaftNodeId>;
}
pub struct TxnStatusGroup<T: RaftTransport> {
group: ConsensusGroup<T>,
sink: Arc<Mutex<TxnStatusApplySink>>,
raft_group_id: RaftGroupId,
}
impl<T: RaftTransport> TxnStatusGroup<T> {
pub async fn create(
mut group_config: GroupConfig,
raft_group_id: RaftGroupId,
transport: Arc<T>,
) -> Result<Self, DistTxnError> {
group_config.cluster_name = raft_group_id.to_hex();
let sink = Arc::new(Mutex::new(TxnStatusApplySink::open(&group_config.dir)?));
let group = ConsensusGroup::create(
group_config,
transport,
sink.clone() as Arc<Mutex<dyn ApplySink>>,
)
.await?;
Ok(TxnStatusGroup {
group,
sink,
raft_group_id,
})
}
pub fn group(&self) -> &ConsensusGroup<T> {
&self.group
}
pub fn raft_group_id(&self) -> RaftGroupId {
self.raft_group_id
}
pub async fn bootstrap(&self, members: &[(RaftNodeId, String)]) -> Result<(), DistTxnError> {
let mut map = BTreeMap::new();
for (raft_id, address) in members {
map.insert(*raft_id, basic_node(address)?);
}
self.group
.bootstrap(map)
.await
.map_err(DistTxnError::Consensus)
}
pub async fn propose(
&self,
command_id: [u8; 16],
command: CoordinatorCommand,
control: &ExecutionControl,
) -> Result<(GroupCommitReceipt, Option<StatusRejectionReason>), DistTxnError> {
let payload = CoordinatorCommandRecord::new(command).encode()?;
let envelope = CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_COORDINATOR, command_id, payload);
let receipt = self
.group
.propose(CommandKind::Transaction, envelope, control)
.await?;
let rejection = {
let sink = self.sink.lock().map_err(|_| {
DistTxnError::InvalidRequest("txn status sink lock poisoned".to_owned())
})?;
sink.state()
.rejections
.iter()
.rev()
.find(|entry| entry.command_id == Some(command_id))
.map(|entry| entry.reason.clone())
};
Ok((receipt, rejection))
}
pub fn record(&self, txn_id: &TransactionId) -> Option<TxnRecord> {
self.sink
.lock()
.expect("txn status sink lock poisoned")
.record(txn_id)
.cloned()
}
pub async fn record_consistent(
&self,
txn_id: &TransactionId,
consistency: &ReadConsistency,
control: &ExecutionControl,
) -> Result<Option<TxnRecord>, DistTxnError> {
self.group
.consistent_read(consistency, control)
.await
.map_err(DistTxnError::Read)?;
Ok(self.record(txn_id))
}
pub fn state(&self) -> TxnStatusState {
self.sink
.lock()
.expect("txn status sink lock poisoned")
.state()
.clone()
}
pub fn session_token(&self, receipt: &GroupCommitReceipt) -> SessionToken {
SessionToken {
group_id: self.group.group_id().to_owned(),
commit_index: receipt.position.index,
commit_ts: receipt.commit_ts,
}
}
pub async fn shutdown(&self) -> Result<(), DistTxnError> {
self.group.shutdown().await.map_err(DistTxnError::Consensus)
}
pub async fn crash(self) {
self.group.crash().await;
}
}
impl<T: RaftTransport> GroupMember for TxnStatusGroup<T> {
fn node_id(&self) -> RaftNodeId {
self.group.node_id()
}
fn current_leader(&self) -> Option<RaftNodeId> {
self.group.metrics().current_leader
}
}
pub struct IntentGroup<T: RaftTransport> {
group: ConsensusGroup<T>,
sink: Arc<Mutex<IntentApplySink>>,
raft_group_id: RaftGroupId,
}
impl<T: RaftTransport> IntentGroup<T> {
pub async fn create(
mut group_config: GroupConfig,
raft_group_id: RaftGroupId,
transport: Arc<T>,
) -> Result<Self, DistTxnError> {
group_config.cluster_name = raft_group_id.to_hex();
let sink = Arc::new(Mutex::new(IntentApplySink::open(&group_config.dir)?));
let group = ConsensusGroup::create(
group_config,
transport,
sink.clone() as Arc<Mutex<dyn ApplySink>>,
)
.await?;
Ok(IntentGroup {
group,
sink,
raft_group_id,
})
}
pub fn group(&self) -> &ConsensusGroup<T> {
&self.group
}
pub fn raft_group_id(&self) -> RaftGroupId {
self.raft_group_id
}
pub async fn bootstrap(&self, members: &[(RaftNodeId, String)]) -> Result<(), DistTxnError> {
let mut map = BTreeMap::new();
for (raft_id, address) in members {
map.insert(*raft_id, basic_node(address)?);
}
self.group
.bootstrap(map)
.await
.map_err(DistTxnError::Consensus)
}
pub async fn propose(
&self,
command_id: [u8; 16],
command: IntentCommand,
control: &ExecutionControl,
) -> Result<(GroupCommitReceipt, Option<PrepareRejectionReason>), DistTxnError> {
let payload = IntentCommandRecord::new(command).encode()?;
let envelope = CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_INTENT, command_id, payload);
let receipt = self
.group
.propose(CommandKind::Transaction, envelope, control)
.await?;
let rejection = {
let sink = self.sink.lock().map_err(|_| {
DistTxnError::InvalidRequest("intent sink lock poisoned".to_owned())
})?;
sink.state()
.rejections
.iter()
.rev()
.find(|entry| entry.command_id == Some(command_id))
.map(|entry| entry.reason.clone())
};
Ok((receipt, rejection))
}
pub fn txn(&self, txn_id: &TransactionId) -> Option<ParticipantTxn> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.txn(txn_id)
.cloned()
}
pub fn unresolved_txn_ids(&self) -> Vec<TransactionId> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.unresolved_txn_ids()
}
pub fn committed_writes(&self) -> Vec<CommittedWrite> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.state()
.committed_writes
.clone()
}
pub fn state(&self) -> IntentState {
self.sink
.lock()
.expect("intent sink lock poisoned")
.state()
.clone()
}
pub async fn shutdown(&self) -> Result<(), DistTxnError> {
self.group.shutdown().await.map_err(DistTxnError::Consensus)
}
pub async fn crash(self) {
self.group.crash().await;
}
}
impl<T: RaftTransport> GroupMember for IntentGroup<T> {
fn node_id(&self) -> RaftNodeId {
self.group.node_id()
}
fn current_leader(&self) -> Option<RaftNodeId> {
self.group.metrics().current_leader
}
}
pub trait IntentGroupMember<T: RaftTransport>: GroupMember {
fn raft_group_id(&self) -> RaftGroupId;
fn group(&self) -> &ConsensusGroup<T>;
fn propose(
&self,
command_id: [u8; 16],
command: IntentCommand,
control: &ExecutionControl,
) -> impl Future<
Output = Result<(GroupCommitReceipt, Option<PrepareRejectionReason>), DistTxnError>,
> + Send;
fn txn(&self, txn_id: &TransactionId) -> Option<ParticipantTxn>;
fn unresolved_txn_ids(&self) -> Vec<TransactionId>;
fn state(&self) -> IntentState;
}
impl<T: RaftTransport> IntentGroupMember<T> for IntentGroup<T> {
fn raft_group_id(&self) -> RaftGroupId {
self.raft_group_id
}
fn group(&self) -> &ConsensusGroup<T> {
&self.group
}
fn propose(
&self,
command_id: [u8; 16],
command: IntentCommand,
control: &ExecutionControl,
) -> impl Future<
Output = Result<(GroupCommitReceipt, Option<PrepareRejectionReason>), DistTxnError>,
> + Send {
self.propose(command_id, command, control)
}
fn txn(&self, txn_id: &TransactionId) -> Option<ParticipantTxn> {
self.txn(txn_id)
}
fn unresolved_txn_ids(&self) -> Vec<TransactionId> {
self.unresolved_txn_ids()
}
fn state(&self) -> IntentState {
self.state()
}
}
pub struct TabletTxnGroup<T: RaftTransport> {
group: ConsensusGroup<T>,
sink: Arc<Mutex<IntentApplySink>>,
engine: Arc<Mutex<EngineApplySink>>,
raft_group_id: RaftGroupId,
}
impl<T: RaftTransport> TabletTxnGroup<T> {
pub async fn create(
mut group_config: GroupConfig,
raft_group_id: RaftGroupId,
transport: Arc<T>,
engine: Arc<Mutex<EngineApplySink>>,
) -> Result<Self, DistTxnError> {
group_config.cluster_name = raft_group_id.to_hex();
let sink = Arc::new(Mutex::new(IntentApplySink::open_with_engine(
&group_config.dir,
engine.clone(),
)?));
let group = ConsensusGroup::create(
group_config,
transport,
sink.clone() as Arc<Mutex<dyn ApplySink>>,
)
.await?;
Ok(TabletTxnGroup {
group,
sink,
engine,
raft_group_id,
})
}
pub fn group(&self) -> &ConsensusGroup<T> {
&self.group
}
pub fn raft_group_id(&self) -> RaftGroupId {
self.raft_group_id
}
pub fn engine(&self) -> &Arc<Mutex<EngineApplySink>> {
&self.engine
}
pub async fn bootstrap(&self, members: &[(RaftNodeId, String)]) -> Result<(), DistTxnError> {
let mut map = BTreeMap::new();
for (raft_id, address) in members {
map.insert(*raft_id, basic_node(address)?);
}
self.group
.bootstrap(map)
.await
.map_err(DistTxnError::Consensus)
}
pub async fn propose(
&self,
command_id: [u8; 16],
command: IntentCommand,
control: &ExecutionControl,
) -> Result<(GroupCommitReceipt, Option<PrepareRejectionReason>), DistTxnError> {
let payload = IntentCommandRecord::new(command).encode()?;
let envelope = CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_INTENT, command_id, payload);
let receipt = self
.group
.propose(CommandKind::Transaction, envelope, control)
.await?;
let rejection = {
let sink = self.sink.lock().map_err(|_| {
DistTxnError::InvalidRequest("intent sink lock poisoned".to_owned())
})?;
sink.state()
.rejections
.iter()
.rev()
.find(|entry| entry.command_id == Some(command_id))
.map(|entry| entry.reason.clone())
};
Ok((receipt, rejection))
}
pub fn txn(&self, txn_id: &TransactionId) -> Option<ParticipantTxn> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.txn(txn_id)
.cloned()
}
pub fn unresolved_txn_ids(&self) -> Vec<TransactionId> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.unresolved_txn_ids()
}
pub fn committed_writes(&self) -> Vec<CommittedWrite> {
self.sink
.lock()
.expect("intent sink lock poisoned")
.state()
.committed_writes
.clone()
}
pub fn state(&self) -> IntentState {
self.sink
.lock()
.expect("intent sink lock poisoned")
.state()
.clone()
}
pub async fn shutdown(&self) -> Result<(), DistTxnError> {
self.group.shutdown().await.map_err(DistTxnError::Consensus)
}
pub async fn crash(self) {
self.group.crash().await;
}
}
impl<T: RaftTransport> GroupMember for TabletTxnGroup<T> {
fn node_id(&self) -> RaftNodeId {
self.group.node_id()
}
fn current_leader(&self) -> Option<RaftNodeId> {
self.group.metrics().current_leader
}
}
impl<T: RaftTransport> IntentGroupMember<T> for TabletTxnGroup<T> {
fn raft_group_id(&self) -> RaftGroupId {
self.raft_group_id
}
fn group(&self) -> &ConsensusGroup<T> {
&self.group
}
fn propose(
&self,
command_id: [u8; 16],
command: IntentCommand,
control: &ExecutionControl,
) -> impl Future<
Output = Result<(GroupCommitReceipt, Option<PrepareRejectionReason>), DistTxnError>,
> + Send {
self.propose(command_id, command, control)
}
fn txn(&self, txn_id: &TransactionId) -> Option<ParticipantTxn> {
self.txn(txn_id)
}
fn unresolved_txn_ids(&self) -> Vec<TransactionId> {
self.unresolved_txn_ids()
}
fn state(&self) -> IntentState {
self.state()
}
}
#[derive(Debug, Clone)]
pub struct DistTxnConfig {
pub selection: CoordinatorSelection,
pub status_partitions: BTreeMap<u32, TxnStatusPartition>,
pub pending_timeout: Duration,
pub resolved_retention: Duration,
pub sweep_limit: u32,
pub hlc_max_skew: Duration,
pub node_tiebreaker: u32,
pub propose_retry_interval: Duration,
pub max_propose_rounds: usize,
}
impl Default for DistTxnConfig {
fn default() -> Self {
DistTxnConfig {
selection: CoordinatorSelection::TxnIdDerived,
status_partitions: BTreeMap::new(),
pending_timeout: DEFAULT_PENDING_TIMEOUT,
resolved_retention: DEFAULT_RESOLVED_RETENTION,
sweep_limit: DEFAULT_SWEEP_LIMIT,
hlc_max_skew: Duration::from_millis(500),
node_tiebreaker: 0,
propose_retry_interval: Duration::from_millis(50),
max_propose_rounds: 600,
}
}
}
#[derive(Debug, Clone)]
pub struct ParticipantWrites {
pub participant: TxnParticipant,
pub expected_schema_version: SchemaVersion,
pub expected_authz_version: u64,
pub intents: Vec<WriteIntent>,
}
#[derive(Debug, Clone)]
pub struct CommitRequest {
pub txn_id: TransactionId,
pub idempotency_key: [u8; 16],
pub writes: Vec<ParticipantWrites>,
pub observed: Vec<HlcTimestamp>,
pub first_write_tablet: Option<TabletDescriptor>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecoveryOutcome {
NotFound,
InFlight {
missing_participants: usize,
},
Decided(DistributedTxnState),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PushOutcome {
NotFound,
NotExpired,
Terminal(DistributedTxnState),
Pushed(DistributedTxnState),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OrphanOutcome {
ResolvedCommit {
commit_ts: HlcTimestamp,
},
ResolvedAbort {
reason: AbortReason,
},
PushedAbort {
reason: AbortReason,
},
InFlight,
UnknownRecord,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SweepReport {
pub resolved_committed: usize,
pub resolved_aborted: usize,
pub pushed: usize,
pub in_flight: usize,
pub unknown: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SweepResolvedReport {
pub swept: usize,
pub remaining: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ResolveBroadcast {
pub resolved: Vec<RaftGroupId>,
pub deferred: Vec<RaftGroupId>,
}
fn advance(base: HlcTimestamp, by: Duration) -> HlcTimestamp {
HlcTimestamp {
physical_micros: base
.physical_micros
.saturating_add(u64::try_from(by.as_micros()).unwrap_or(u64::MAX)),
logical: base.logical,
node_tiebreaker: base.node_tiebreaker,
}
}
fn retryable(error: &DistTxnError) -> bool {
match error {
DistTxnError::Consensus(
ConsensusError::NotLeader { .. } | ConsensusError::Closed | ConsensusError::Raft(_),
)
| DistTxnError::Read(
ReadConsistencyError::NotLeader { .. }
| ReadConsistencyError::LeaderUnknown
| ReadConsistencyError::Closed,
) => true,
DistTxnError::Consensus(ConsensusError::Cancelled | ConsensusError::DeadlineExceeded) => {
false
}
_ => false,
}
}
fn control_error(control: &ExecutionControl) -> Option<DistTxnError> {
control.check().err().map(|error| match error {
mongreldb_log::commit_log::LogError::Cancelled => {
DistTxnError::Consensus(ConsensusError::Cancelled)
}
mongreldb_log::commit_log::LogError::DeadlineExceeded => {
DistTxnError::Consensus(ConsensusError::DeadlineExceeded)
}
other => DistTxnError::Consensus(ConsensusError::Raft(other.to_string())),
})
}
async fn retry_across_members<'m, G, F, Fut, R>(
members: &'m [G],
group_label: &str,
config: &DistTxnConfig,
control: &ExecutionControl,
mut call: F,
) -> Result<(&'m G, R), DistTxnError>
where
G: GroupMember,
F: FnMut(&'m G) -> Fut,
Fut: Future<Output = Result<R, DistTxnError>>,
{
if members.is_empty() {
return Err(DistTxnError::InvalidRequest(format!(
"no members supplied for {group_label}"
)));
}
let mut rounds = 0_usize;
loop {
let mut ordered: Vec<&G> = members
.iter()
.filter(|member| member.current_leader() == Some(member.node_id()))
.collect();
ordered.extend(
members
.iter()
.filter(|member| member.current_leader() != Some(member.node_id())),
);
let mut last_error: Option<DistTxnError> = None;
for member in ordered {
match call(member).await {
Ok(answer) => return Ok((member, answer)),
Err(error) if retryable(&error) => last_error = Some(error),
Err(error) => return Err(error),
}
}
rounds += 1;
if let Some(error) = control_error(control) {
return Err(error);
}
if rounds >= config.max_propose_rounds {
return Err(DistTxnError::Unavailable(format!(
"{group_label}: no live leader after {rounds} rounds ({})",
last_error
.map(|error| error.to_string())
.unwrap_or_else(|| "no members answered".to_owned())
)));
}
tokio::time::sleep(config.propose_retry_interval).await;
}
}
pub struct DistTxnDriver {
config: DistTxnConfig,
hlc: HlcClock,
}
impl DistTxnDriver {
pub fn new(config: DistTxnConfig) -> Self {
let hlc = HlcClock::new(config.node_tiebreaker, config.hlc_max_skew);
DistTxnDriver { config, hlc }
}
pub fn config(&self) -> &DistTxnConfig {
&self.config
}
pub fn now(&self) -> Result<HlcTimestamp, DistTxnError> {
self.hlc
.now()
.map_err(|error| DistTxnError::Clock(error.to_string()))
}
fn expiry_from(&self, base: HlcTimestamp) -> HlcTimestamp {
advance(base, self.config.pending_timeout)
}
fn coordinator_for(
&self,
request: &CommitRequest,
) -> Result<(RaftGroupId, Vec<TxnParticipant>), DistTxnError> {
if request.writes.is_empty() {
return Err(DistTxnError::InvalidRequest(
"a distributed commit needs at least one participant write".to_owned(),
));
}
let mut participants = Vec::with_capacity(request.writes.len());
for writes in &request.writes {
if participants.contains(&writes.participant) {
return Err(DistTxnError::InvalidRequest(format!(
"tablet {} appears twice in the write set",
writes.participant.tablet_id
)));
}
participants.push(writes.participant);
}
let group = select_coordinator_group(
&self.config.selection,
&request.txn_id,
request.first_write_tablet.as_ref(),
&self.config.status_partitions,
)?;
Ok((group, participants))
}
fn validate_status_set<T: RaftTransport>(
members: &[TxnStatusGroup<T>],
coordinator: RaftGroupId,
) -> Result<(), DistTxnError> {
if members.is_empty() {
return Err(DistTxnError::InvalidRequest(
"no transaction-status members supplied".to_owned(),
));
}
if members
.iter()
.any(|member| member.raft_group_id() != coordinator)
{
return Err(DistTxnError::InvalidRequest(format!(
"the supplied transaction-status members are not the selected coordinator group \
{coordinator}"
)));
}
Ok(())
}
pub async fn begin<T: RaftTransport>(
&self,
status: &[TxnStatusGroup<T>],
request: &CommitRequest,
control: &ExecutionControl,
) -> Result<TxnRecord, DistTxnError> {
let (coordinator, participants) = self.coordinator_for(request)?;
Self::validate_status_set(status, coordinator)?;
let now = self.now()?;
let max_observed = request
.observed
.iter()
.copied()
.max()
.unwrap_or(HlcTimestamp::ZERO);
let record = TxnRecord {
txn_id: request.txn_id,
state: DistributedTxnState::Pending,
participants,
prepare_ts: BTreeMap::new(),
coordinator,
created_at: now,
heartbeat: now,
expiry: self.expiry_from(now),
max_observed,
idempotency_key: request.idempotency_key,
};
let command_id = command_id_for(TAG_BEGIN, &request.txn_id, &[]);
let (member, (_, rejection)) = retry_across_members(
status,
"txn-status begin",
&self.config,
control,
|member| {
let record = record.clone();
let control = control.clone();
async move {
member
.propose(command_id, CoordinatorCommand::Begin { record }, &control)
.await
}
},
)
.await?;
match rejection {
Some(StatusRejectionReason::IdempotencyKeyConflict { .. }) => {
return Err(DistTxnError::IdempotencyConflict(request.txn_id));
}
Some(reason) => {
return Err(DistTxnError::InvalidRequest(format!(
"begin refused by the coordinator: {reason}"
)));
}
None => {}
}
let applied = member.record(&request.txn_id).ok_or_else(|| {
DistTxnError::InvalidRequest("begin committed but no record is visible".to_owned())
})?;
if applied.idempotency_key != request.idempotency_key {
return Err(DistTxnError::IdempotencyConflict(request.txn_id));
}
Ok(applied)
}
pub async fn heartbeat<T: RaftTransport>(
&self,
status: &[TxnStatusGroup<T>],
coordinator: RaftGroupId,
txn_id: &TransactionId,
control: &ExecutionControl,
) -> Result<TxnRecord, DistTxnError> {
Self::validate_status_set(status, coordinator)?;
let heartbeat = self.now()?;
let expiry = self.expiry_from(heartbeat);
let mut extra = Vec::with_capacity(16);
extra.extend_from_slice(&heartbeat.physical_micros.to_le_bytes());
extra.extend_from_slice(&heartbeat.logical.to_le_bytes());
extra.extend_from_slice(&heartbeat.node_tiebreaker.to_le_bytes());
let command_id = command_id_for(TAG_HEARTBEAT, txn_id, &extra);
let (member, (_, rejection)) = retry_across_members(
status,
"txn-status heartbeat",
&self.config,
control,
|member| {
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::Heartbeat {
txn_id: *txn_id,
heartbeat,
expiry,
},
&control,
)
.await
}
},
)
.await?;
match rejection {
Some(StatusRejectionReason::UnknownTxn(_)) => {
return Err(DistTxnError::InvalidRequest(format!(
"heartbeat for unknown transaction {txn_id}"
)));
}
Some(reason) => {
return Err(DistTxnError::InvalidRequest(format!(
"heartbeat refused by the coordinator: {reason}"
)));
}
None => {}
}
member.record(txn_id).ok_or_else(|| {
DistTxnError::InvalidRequest("heartbeat committed but no record is visible".to_owned())
})
}
pub async fn prepare_participant<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
intents: &[G],
txn_id: &TransactionId,
writes: &ParticipantWrites,
control: &ExecutionControl,
) -> Result<PrepareToken, DistTxnError> {
if intents.is_empty()
|| intents
.iter()
.any(|member| member.raft_group_id() != writes.participant.raft_group_id)
{
return Err(DistTxnError::InvalidRequest(format!(
"the supplied intent members are not the participant group {}",
writes.participant.raft_group_id
)));
}
mongreldb_fault::inject("txn.prepare.before")?;
let prepare_ts = self.now()?;
let mut staged = writes.intents.clone();
for intent in &mut staged {
intent.txn_id = *txn_id;
intent.prepare_ts = prepare_ts;
}
let command_id =
command_id_for(TAG_PREPARE, txn_id, writes.participant.tablet_id.as_bytes());
let (member, (receipt, rejection)) = retry_across_members(
intents,
"participant prepare",
&self.config,
control,
|member| {
let staged = staged.clone();
let control = control.clone();
async move {
member
.propose(
command_id,
IntentCommand::PersistIntents {
txn_id: *txn_id,
expected_schema_version: writes.expected_schema_version,
expected_authz_version: writes.expected_authz_version,
prepare_ts,
intents: staged,
},
&control,
)
.await
}
},
)
.await?;
if let Some(reason) = rejection {
return Err(DistTxnError::PrepareRejected(reason));
}
let stored = member.txn(txn_id).ok_or_else(|| {
DistTxnError::InvalidRequest(
"prepare committed but no intent record is visible".to_owned(),
)
})?;
let token = PrepareToken {
txn_id: *txn_id,
tablet_id: writes.participant.tablet_id,
raft_group_id: writes.participant.raft_group_id,
prepare_ts: stored.prepare_ts,
position: receipt.position,
command_id,
};
mongreldb_fault::inject("txn.prepare.after")?;
Ok(token)
}
async fn mark_prepared<T: RaftTransport>(
&self,
status: &[TxnStatusGroup<T>],
txn_id: &TransactionId,
token: &PrepareToken,
observed: HlcTimestamp,
control: &ExecutionControl,
) -> Result<(), DistTxnError> {
let command_id = command_id_for(TAG_MARK, txn_id, token.tablet_id.as_bytes());
let (_, (_, rejection)) = retry_across_members(
status,
"txn-status mark-preparing",
&self.config,
control,
|member| {
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::MarkPreparing {
txn_id: *txn_id,
tablet_id: token.tablet_id,
prepare_ts: token.prepare_ts,
observed,
},
&control,
)
.await
}
},
)
.await?;
match rejection {
None => Ok(()),
Some(StatusRejectionReason::DecisionFinal { existing }) => Err(match existing {
DistributedTxnState::Aborted { reason } => DistTxnError::Aborted(reason),
state => DistTxnError::InvalidRequest(format!(
"mark-preparing raced a final decision: {state:?}"
)),
}),
Some(reason) => Err(DistTxnError::InvalidRequest(format!(
"mark-preparing refused by the coordinator: {reason}"
))),
}
}
pub async fn decide_commit<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
prepared: &[PrepareToken],
observed: &[HlcTimestamp],
control: &ExecutionControl,
) -> Result<TxnOutcome, DistTxnError> {
let mut minimum = HlcTimestamp::ZERO;
for token in prepared {
minimum = minimum.max(token.prepare_ts);
}
for timestamp in observed {
minimum = minimum.max(*timestamp);
}
let commit_ts = self.hlc.next_after(minimum);
mongreldb_fault::inject("txn.decision.before")?;
let command_id = command_id_for(TAG_COMMIT, txn_id, &[]);
retry_across_members(
status,
"txn-status commit",
&self.config,
control,
|member| {
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::Commit {
txn_id: *txn_id,
commit_ts,
},
&control,
)
.await
}
},
)
.await
.map_err(|error| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("decision proposal failed: {error}"),
})?;
self.finalize_decision(status, participants, txn_id, control)
.await
}
async fn finalize_decision<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
control: &ExecutionControl,
) -> Result<TxnOutcome, DistTxnError> {
let record = self
.read_record(status, txn_id, control)
.await
.map_err(|error| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("decision read-back failed: {error}"),
})?
.ok_or_else(|| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: "no coordinator record after the decision proposal".to_owned(),
})?;
match record.state.clone() {
DistributedTxnState::Committed { commit_ts } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Committed { commit_ts },
control,
)
.await;
post_decision_after_hook(*txn_id)?;
Ok(TxnOutcome {
txn_id: *txn_id,
commit_ts,
participants: record.participants.clone(),
durability: DurabilityLevel::Quorum,
})
}
DistributedTxnState::Aborted { reason } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Aborted {
reason: reason.clone(),
},
control,
)
.await;
let _ = post_decision_after_hook(*txn_id);
Err(DistTxnError::Aborted(reason))
}
state => Err(DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("decision is not durable; record state is {state:?}"),
}),
}
}
async fn finalize_abort<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
reason: AbortReason,
control: &ExecutionControl,
) -> Result<DistributedTxnState, DistTxnError> {
mongreldb_fault::inject("txn.decision.before")?;
let command_id = command_id_for(TAG_ABORT, txn_id, &[]);
retry_across_members(
status,
"txn-status abort",
&self.config,
control,
|member| {
let reason = reason.clone();
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::Abort {
txn_id: *txn_id,
reason,
},
&control,
)
.await
}
},
)
.await
.map_err(|error| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("abort proposal failed: {error}"),
})?;
let record = self
.read_record(status, txn_id, control)
.await
.map_err(|error| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("abort read-back failed: {error}"),
})?
.ok_or_else(|| DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: "no coordinator record after the abort proposal".to_owned(),
})?;
match record.state.clone() {
DistributedTxnState::Committed { commit_ts } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Committed { commit_ts },
control,
)
.await;
post_decision_after_hook(*txn_id)?;
Ok(DistributedTxnState::Committed { commit_ts })
}
DistributedTxnState::Aborted { reason } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Aborted {
reason: reason.clone(),
},
control,
)
.await;
let _ = post_decision_after_hook(*txn_id);
Ok(DistributedTxnState::Aborted { reason })
}
state => Err(DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("abort is not durable; record state is {state:?}"),
}),
}
}
pub async fn commit<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
request: CommitRequest,
control: &ExecutionControl,
) -> Result<TxnOutcome, DistTxnError> {
let record = self.begin(status, &request, control).await?;
match record.state {
DistributedTxnState::Committed { commit_ts } => {
self.broadcast_resolve(
participants,
&request.txn_id,
&record.participants,
TxnDecision::Committed { commit_ts },
control,
)
.await;
return Ok(TxnOutcome {
txn_id: request.txn_id,
commit_ts,
participants: record.participants.clone(),
durability: DurabilityLevel::Quorum,
});
}
DistributedTxnState::Aborted { reason } => {
self.broadcast_resolve(
participants,
&request.txn_id,
&record.participants,
TxnDecision::Aborted {
reason: reason.clone(),
},
control,
)
.await;
return Err(DistTxnError::Aborted(reason));
}
_ => {}
}
let general_path = request.writes.len() > 1;
let mut prepare_futures = Vec::with_capacity(request.writes.len());
for writes in &request.writes {
let intent_members = participants
.get(&writes.participant.raft_group_id)
.ok_or_else(|| {
DistTxnError::InvalidRequest(format!(
"no intent members supplied for participant group {}",
writes.participant.raft_group_id
))
})?;
prepare_futures.push(self.prepare_participant(
intent_members,
&request.txn_id,
writes,
control,
));
}
let results = run_all(prepare_futures).await;
let mut prepared = Vec::with_capacity(results.len());
for result in results {
match result {
Ok(token) => prepared.push(token),
Err(error) => {
let reason = abort_reason_of(&error);
return match self
.finalize_abort(status, participants, &request.txn_id, reason, control)
.await?
{
DistributedTxnState::Aborted { reason } => {
Err(DistTxnError::Aborted(reason))
}
DistributedTxnState::Committed { commit_ts } => Ok(TxnOutcome {
txn_id: request.txn_id,
commit_ts,
participants: record.participants.clone(),
durability: DurabilityLevel::Quorum,
}),
state => Err(DistTxnError::OutcomeAmbiguous {
txn_id: request.txn_id,
detail: format!(
"abort after a failed prepare is not durable; state is {state:?}"
),
}),
};
}
}
}
let mut observed = request.observed.clone();
if general_path {
observed.extend(prepared.iter().map(|token| token.prepare_ts));
let max_observed = observed.iter().copied().max().unwrap_or(HlcTimestamp::ZERO);
for token in &prepared {
if let Err(error) = self
.mark_prepared(status, &request.txn_id, token, max_observed, control)
.await
{
let reason = match error {
DistTxnError::Aborted(reason) => reason,
other => AbortReason::Error(other.to_string()),
};
match self
.finalize_abort(status, participants, &request.txn_id, reason, control)
.await?
{
DistributedTxnState::Aborted { reason } => {
return Err(DistTxnError::Aborted(reason));
}
DistributedTxnState::Committed { commit_ts } => {
return Ok(TxnOutcome {
txn_id: request.txn_id,
commit_ts,
participants: record.participants.clone(),
durability: DurabilityLevel::Quorum,
});
}
state => {
return Err(DistTxnError::OutcomeAmbiguous {
txn_id: request.txn_id,
detail: format!(
"abort after a failed mark is not durable; state is {state:?}"
),
});
}
}
}
}
}
self.decide_commit(
status,
participants,
&request.txn_id,
&prepared,
&observed,
control,
)
.await
}
pub async fn abort<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
reason: AbortReason,
control: &ExecutionControl,
) -> Result<DistributedTxnState, DistTxnError> {
self.finalize_abort(status, participants, txn_id, reason, control)
.await
}
pub async fn broadcast_resolve<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
txn_participants: &[TxnParticipant],
decision: TxnDecision,
control: &ExecutionControl,
) -> ResolveBroadcast {
let mut report = ResolveBroadcast::default();
for participant in txn_participants {
let command_id = command_id_for(TAG_RESOLVE, txn_id, participant.tablet_id.as_bytes());
let Some(members) = participants.get(&participant.raft_group_id) else {
report.deferred.push(participant.raft_group_id);
continue;
};
let result = retry_across_members(
members,
"participant resolve",
&self.config,
control,
|member| {
let decision = decision.clone();
let control = control.clone();
async move {
member
.propose(
command_id,
IntentCommand::Resolve {
txn_id: *txn_id,
decision,
},
&control,
)
.await
}
},
)
.await;
match result {
Ok(_) => report.resolved.push(participant.raft_group_id),
Err(_) => report.deferred.push(participant.raft_group_id),
}
}
report
}
pub async fn read_record<T: RaftTransport>(
&self,
status: &[TxnStatusGroup<T>],
txn_id: &TransactionId,
control: &ExecutionControl,
) -> Result<Option<TxnRecord>, DistTxnError> {
let (_, record) =
retry_across_members(status, "txn-status read", &self.config, control, |member| {
let control = control.clone();
async move {
member
.record_consistent(txn_id, &ReadConsistency::Linearizable, &control)
.await
}
})
.await?;
Ok(record)
}
fn probe_participant<T: RaftTransport, G: IntentGroupMember<T>>(
participants: &BTreeMap<RaftGroupId, Vec<G>>,
participant: &TxnParticipant,
txn_id: &TransactionId,
) -> Option<ParticipantTxn> {
let members = participants.get(&participant.raft_group_id)?;
members.iter().find_map(|member| member.txn(txn_id))
}
pub async fn recover<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
now: HlcTimestamp,
control: &ExecutionControl,
) -> Result<RecoveryOutcome, DistTxnError> {
let Some(record) = self.read_record(status, txn_id, control).await? else {
return Ok(RecoveryOutcome::NotFound);
};
match record.state.clone() {
DistributedTxnState::Committed { commit_ts } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Committed { commit_ts },
control,
)
.await;
Ok(RecoveryOutcome::Decided(DistributedTxnState::Committed {
commit_ts,
}))
}
DistributedTxnState::Aborted { reason } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Aborted {
reason: reason.clone(),
},
control,
)
.await;
Ok(RecoveryOutcome::Decided(DistributedTxnState::Aborted {
reason,
}))
}
DistributedTxnState::Pending | DistributedTxnState::Preparing => {
let mut prepare_ts: Vec<HlcTimestamp> = Vec::new();
let mut missing = 0_usize;
for participant in &record.participants {
match Self::probe_participant(participants, participant, txn_id) {
Some(ParticipantTxn {
resolution: None,
prepare_ts: ts,
..
}) => prepare_ts.push(ts),
_ => missing += 1,
}
}
if missing == 0 {
let mut minimum = record.max_observed;
for ts in prepare_ts {
minimum = minimum.max(ts);
}
for ts in record.prepare_ts.values() {
minimum = minimum.max(*ts);
}
let commit_ts = self.hlc.next_after(minimum);
let command_id = command_id_for(TAG_COMMIT, txn_id, &[]);
retry_across_members(
status,
"txn-status commit",
&self.config,
control,
|member| {
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::Commit {
txn_id: *txn_id,
commit_ts,
},
&control,
)
.await
}
},
)
.await?;
let outcome = self
.finalize_decision(status, participants, txn_id, control)
.await;
return match outcome {
Ok(outcome) => {
Ok(RecoveryOutcome::Decided(DistributedTxnState::Committed {
commit_ts: outcome.commit_ts,
}))
}
Err(DistTxnError::Aborted(reason)) => {
Ok(RecoveryOutcome::Decided(DistributedTxnState::Aborted {
reason,
}))
}
Err(error) => Err(error),
};
}
if record.expired(now) {
let reason = AbortReason::Cancelled(format!(
"transaction heartbeat expired at {:?} (pushed by recovery)",
record.expiry
));
let state = self
.finalize_abort(status, participants, txn_id, reason, control)
.await?;
return Ok(RecoveryOutcome::Decided(state));
}
Ok(RecoveryOutcome::InFlight {
missing_participants: missing,
})
}
}
}
pub async fn push_expired<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
participants: &BTreeMap<RaftGroupId, Vec<G>>,
txn_id: &TransactionId,
now: HlcTimestamp,
control: &ExecutionControl,
) -> Result<PushOutcome, DistTxnError> {
let Some(record) = self.read_record(status, txn_id, control).await? else {
return Ok(PushOutcome::NotFound);
};
match record.state.clone() {
DistributedTxnState::Committed { commit_ts } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Committed { commit_ts },
control,
)
.await;
Ok(PushOutcome::Terminal(DistributedTxnState::Committed {
commit_ts,
}))
}
DistributedTxnState::Aborted { reason } => {
self.broadcast_resolve(
participants,
txn_id,
&record.participants,
TxnDecision::Aborted {
reason: reason.clone(),
},
control,
)
.await;
Ok(PushOutcome::Terminal(DistributedTxnState::Aborted {
reason,
}))
}
DistributedTxnState::Pending | DistributedTxnState::Preparing => {
if !record.expired(now) {
return Ok(PushOutcome::NotExpired);
}
let reason = AbortReason::Cancelled(format!(
"transaction heartbeat expired at {:?} (pushed by a third party)",
record.expiry
));
let state = self
.finalize_abort(status, participants, txn_id, reason, control)
.await?;
Ok(PushOutcome::Pushed(state))
}
}
}
pub async fn resolve_orphan<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
intent_members: &[G],
txn_id: &TransactionId,
now: HlcTimestamp,
control: &ExecutionControl,
) -> Result<OrphanOutcome, DistTxnError> {
let Some(record) = self.read_record(status, txn_id, control).await? else {
return Ok(OrphanOutcome::UnknownRecord);
};
let (decision, pushed) = match record.state.clone() {
DistributedTxnState::Committed { commit_ts } => {
(TxnDecision::Committed { commit_ts }, false)
}
DistributedTxnState::Aborted { reason } => (TxnDecision::Aborted { reason }, false),
DistributedTxnState::Pending | DistributedTxnState::Preparing => {
if !record.expired(now) {
return Ok(OrphanOutcome::InFlight);
}
let reason = AbortReason::Cancelled(format!(
"transaction heartbeat expired at {:?} (pushed by orphan recovery)",
record.expiry
));
let decision = match self
.finalize_abort_for_one(status, intent_members, &record, reason, control)
.await?
{
DistributedTxnState::Committed { commit_ts } => {
TxnDecision::Committed { commit_ts }
}
DistributedTxnState::Aborted { reason } => TxnDecision::Aborted { reason },
state => {
return Err(DistTxnError::OutcomeAmbiguous {
txn_id: *txn_id,
detail: format!("orphan push is not durable; state is {state:?}"),
});
}
};
(decision, true)
}
};
if !pushed {
let participant = record.participants.iter().find(|participant| {
intent_members
.first()
.is_some_and(|member| member.raft_group_id() == participant.raft_group_id)
});
let tablet = participant.map_or(TabletId::ZERO, |participant| participant.tablet_id);
let command_id = command_id_for(TAG_RESOLVE, txn_id, tablet.as_bytes());
retry_across_members(
intent_members,
"participant resolve",
&self.config,
control,
|member| {
let decision = decision.clone();
let control = control.clone();
async move {
member
.propose(
command_id,
IntentCommand::Resolve {
txn_id: *txn_id,
decision,
},
&control,
)
.await
}
},
)
.await?;
}
match decision {
TxnDecision::Committed { commit_ts } => Ok(OrphanOutcome::ResolvedCommit { commit_ts }),
TxnDecision::Aborted { reason } if pushed => Ok(OrphanOutcome::PushedAbort { reason }),
TxnDecision::Aborted { reason } => Ok(OrphanOutcome::ResolvedAbort { reason }),
}
}
async fn finalize_abort_for_one<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
intent_members: &[G],
record: &TxnRecord,
reason: AbortReason,
control: &ExecutionControl,
) -> Result<DistributedTxnState, DistTxnError> {
let txn_id = record.txn_id;
let command_id = command_id_for(TAG_ABORT, &txn_id, &[]);
retry_across_members(
status,
"txn-status abort",
&self.config,
control,
|member| {
let reason = reason.clone();
let control = control.clone();
async move {
member
.propose(
command_id,
CoordinatorCommand::Abort { txn_id, reason },
&control,
)
.await
}
},
)
.await
.map_err(|error| DistTxnError::OutcomeAmbiguous {
txn_id,
detail: format!("orphan push failed: {error}"),
})?;
let final_record = self
.read_record(status, &txn_id, control)
.await?
.ok_or_else(|| DistTxnError::OutcomeAmbiguous {
txn_id,
detail: "no coordinator record after the orphan push".to_owned(),
})?;
let decision = match final_record.state.clone() {
DistributedTxnState::Committed { commit_ts } => TxnDecision::Committed { commit_ts },
DistributedTxnState::Aborted { reason } => TxnDecision::Aborted { reason },
state => {
return Err(DistTxnError::OutcomeAmbiguous {
txn_id,
detail: format!("orphan push is not durable; state is {state:?}"),
});
}
};
let participant = record.participants.iter().find(|participant| {
intent_members
.first()
.is_some_and(|member| member.raft_group_id() == participant.raft_group_id)
});
let tablet = participant.map_or(TabletId::ZERO, |participant| participant.tablet_id);
let resolve_id = command_id_for(TAG_RESOLVE, &txn_id, tablet.as_bytes());
retry_across_members(
intent_members,
"participant resolve",
&self.config,
control,
|member| {
let decision = decision.clone();
let control = control.clone();
async move {
member
.propose(
resolve_id,
IntentCommand::Resolve { txn_id, decision },
&control,
)
.await
}
},
)
.await?;
Ok(final_record.state.clone())
}
pub async fn sweep_orphans<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
status: &[TxnStatusGroup<T>],
intent_members: &[G],
now: HlcTimestamp,
control: &ExecutionControl,
) -> Result<SweepReport, DistTxnError> {
let mut report = SweepReport::default();
if intent_members.is_empty() {
return Err(DistTxnError::InvalidRequest(
"no intent members supplied for the sweep".to_owned(),
));
}
let orphans = intent_members
.iter()
.max_by_key(|member| member.group().applied_position().index)
.map(|member| member.unresolved_txn_ids())
.unwrap_or_default();
for txn_id in orphans {
match self
.resolve_orphan(status, intent_members, &txn_id, now, control)
.await?
{
OrphanOutcome::ResolvedCommit { .. } => report.resolved_committed += 1,
OrphanOutcome::ResolvedAbort { .. } => report.resolved_aborted += 1,
OrphanOutcome::PushedAbort { .. } => {
report.pushed += 1;
report.resolved_aborted += 1;
}
OrphanOutcome::InFlight => report.in_flight += 1,
OrphanOutcome::UnknownRecord => report.unknown += 1,
}
}
Ok(report)
}
pub async fn sweep_resolved<T: RaftTransport, G: IntentGroupMember<T>>(
&self,
intent_members: &[G],
now: HlcTimestamp,
control: &ExecutionControl,
) -> Result<SweepResolvedReport, DistTxnError> {
if intent_members.is_empty() {
return Err(DistTxnError::InvalidRequest(
"no intent members supplied for the resolved sweep".to_owned(),
));
}
let older_than = HlcTimestamp {
physical_micros: now.physical_micros.saturating_sub(
u64::try_from(self.config.resolved_retention.as_micros()).unwrap_or(u64::MAX),
),
logical: 0,
node_tiebreaker: 0,
};
let eligible = |state: &IntentState| {
state
.txns
.values()
.filter(|txn| {
txn.resolution.is_some() && txn.resolved_at.is_some_and(|at| at < older_than)
})
.count()
};
let before = intent_members
.iter()
.map(|member| eligible(&member.state()))
.max()
.unwrap_or(0);
let mut extra = Vec::with_capacity(24);
extra.extend_from_slice(&older_than.physical_micros.to_le_bytes());
extra.extend_from_slice(&self.config.sweep_limit.to_le_bytes());
let command_id = command_id_for(TAG_SWEEP, &TransactionId::ZERO, &extra);
let (member, _) = retry_across_members(
intent_members,
"participant resolved sweep",
&self.config,
control,
|member| {
let control = control.clone();
async move {
member
.propose(
command_id,
IntentCommand::SweepResolved {
older_than,
limit: self.config.sweep_limit,
},
&control,
)
.await
}
},
)
.await?;
let after = eligible(&member.state());
Ok(SweepResolvedReport {
swept: before.saturating_sub(after),
remaining: after,
})
}
}
async fn run_all<F, T>(futures: Vec<F>) -> Vec<T>
where
F: Future<Output = T>,
{
let mut pending: Vec<Option<std::pin::Pin<Box<F>>>> = futures
.into_iter()
.map(|future| Some(Box::pin(future)))
.collect();
let mut outputs: Vec<Option<T>> = std::iter::repeat_with(|| None)
.take(pending.len())
.collect();
let mut remaining = pending.len();
std::future::poll_fn(|cx| {
for (index, slot) in pending.iter_mut().enumerate() {
let Some(future) = slot else { continue };
match future.as_mut().poll(cx) {
std::task::Poll::Ready(output) => {
*slot = None;
outputs[index] = Some(output);
remaining -= 1;
}
std::task::Poll::Pending => {}
}
}
if remaining == 0 {
std::task::Poll::Ready(())
} else {
std::task::Poll::Pending
}
})
.await;
outputs
.into_iter()
.map(|output| output.expect("every future completed"))
.collect()
}
fn abort_reason_of(error: &DistTxnError) -> AbortReason {
match error {
DistTxnError::PrepareRejected(PrepareRejectionReason::KeyConflict { holder, .. }) => {
AbortReason::Conflict(format!("write intent conflict with transaction {holder}"))
}
DistTxnError::PrepareRejected(reason) => AbortReason::Validation(reason.to_string()),
DistTxnError::Consensus(ConsensusError::Cancelled) => {
AbortReason::Cancelled("cancelled during prepare".to_owned())
}
DistTxnError::Consensus(ConsensusError::DeadlineExceeded) => {
AbortReason::Cancelled("deadline exceeded during prepare".to_owned())
}
other => AbortReason::Error(other.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tablet::{PartitionBounds, TabletState};
use mongreldb_consensus::network::InMemoryTransport;
use mongreldb_types::ids::TableId;
use std::time::Instant;
fn ts(micros: u64) -> HlcTimestamp {
HlcTimestamp {
physical_micros: micros,
logical: 0,
node_tiebreaker: 0,
}
}
fn gid(byte: u8) -> RaftGroupId {
RaftGroupId::from_bytes([byte; 16])
}
fn tid(byte: u8) -> TabletId {
TabletId::from_bytes([byte; 16])
}
fn xid(byte: u8) -> TransactionId {
TransactionId::from_bytes([byte; 16])
}
fn later(base: HlcTimestamp, delta: Duration) -> HlcTimestamp {
HlcTimestamp {
physical_micros: base
.physical_micros
.saturating_add(u64::try_from(delta.as_micros()).unwrap_or(u64::MAX)),
logical: base.logical,
node_tiebreaker: base.node_tiebreaker,
}
}
fn one_before(ts: HlcTimestamp) -> HlcTimestamp {
HlcTimestamp {
physical_micros: ts.physical_micros.saturating_sub(1),
logical: ts.logical,
node_tiebreaker: ts.node_tiebreaker,
}
}
fn record(txn_id: TransactionId, participants: Vec<TxnParticipant>) -> TxnRecord {
TxnRecord {
txn_id,
state: DistributedTxnState::Pending,
participants,
prepare_ts: BTreeMap::new(),
coordinator: gid(90),
created_at: ts(1_000),
heartbeat: ts(1_000),
expiry: ts(2_000),
max_observed: HlcTimestamp::ZERO,
idempotency_key: [7u8; 16],
}
}
fn participant(group: u8, tablet: u8) -> TxnParticipant {
TxnParticipant {
tablet_id: tid(tablet),
raft_group_id: gid(group),
}
}
fn intent(txn: u8, key: &[u8], prepare_micros: u64) -> WriteIntent {
WriteIntent {
txn_id: xid(txn),
key: key.to_vec(),
value_ref: format!("value-{}", String::from_utf8_lossy(key)).into_bytes(),
prepare_ts: ts(prepare_micros),
}
}
fn coordinator_envelope(
id: [u8; 16],
command: CoordinatorCommand,
index_micros: u64,
) -> ReplicatedCommand {
let payload = CoordinatorCommandRecord::new(command).encode().unwrap();
ReplicatedCommand::new(
CommandKind::Transaction,
CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_COORDINATOR, id, payload),
ts(index_micros),
)
}
fn intent_envelope(
id: [u8; 16],
command: IntentCommand,
index_micros: u64,
) -> ReplicatedCommand {
let payload = IntentCommandRecord::new(command).encode().unwrap();
ReplicatedCommand::new(
CommandKind::Transaction,
CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_INTENT, id, payload),
ts(index_micros),
)
}
fn applied(index: u64, command: ReplicatedCommand) -> AppliedCommand {
AppliedCommand {
position: LogPosition { term: 1, index },
command,
}
}
#[test]
fn fnv1a_is_deterministic_and_matches_the_empty_input_constant() {
assert_eq!(fnv1a_64(&[]), 0xcbf29ce484222325);
assert_eq!(fnv1a_64(b"mongrel"), fnv1a_64(b"mongrel"));
assert_ne!(fnv1a_64(b"mongrel"), fnv1a_64(b"mongrele"));
}
#[test]
fn command_ids_are_deterministic_and_distinct_per_step() {
let txn = xid(5);
assert_eq!(
command_id_for(TAG_BEGIN, &txn, &[]),
command_id_for(TAG_BEGIN, &txn, &[])
);
let ids: Vec<[u8; 16]> = [
TAG_BEGIN,
TAG_MARK,
TAG_COMMIT,
TAG_ABORT,
TAG_PREPARE,
TAG_RESOLVE,
]
.iter()
.map(|tag| command_id_for(tag, &txn, &[]))
.collect();
for (i, left) in ids.iter().enumerate() {
for right in ids.iter().skip(i + 1) {
assert_ne!(left, right);
}
}
assert_ne!(
command_id_for(TAG_BEGIN, &xid(6), &[]),
command_id_for(TAG_BEGIN, &txn, &[])
);
}
#[test]
fn txn_id_derived_selection_is_deterministic() {
let mut partitions = BTreeMap::new();
for index in 0..4_u32 {
partitions.insert(
index,
TxnStatusPartition {
partition_id: index,
home_raft_group: gid(50 + index as u8),
},
);
}
let txn = xid(9);
let expected_index = txn_status_partition_index(&txn, 4) as usize;
let first =
select_coordinator_group(&CoordinatorSelection::TxnIdDerived, &txn, None, &partitions)
.unwrap();
let second =
select_coordinator_group(&CoordinatorSelection::TxnIdDerived, &txn, None, &partitions)
.unwrap();
assert_eq!(first, second);
assert_eq!(first, gid(50 + expected_index as u8));
assert_eq!(
CoordinatorSelection::default(),
CoordinatorSelection::TxnIdDerived
);
assert!(select_coordinator_group(
&CoordinatorSelection::TxnIdDerived,
&txn,
None,
&BTreeMap::new()
)
.is_err());
}
#[test]
fn first_write_home_selection_uses_the_tablets_group() {
let tablet = TabletDescriptor {
tablet_id: tid(7),
table_id: TableId::new(3),
database_id: mongreldb_types::ids::DatabaseId::ZERO,
raft_group_id: gid(42),
partition: PartitionBounds::unbounded(),
replicas: Vec::new(),
leader_hint: None,
generation: 1,
state: TabletState::Active,
};
let selected = select_coordinator_group(
&CoordinatorSelection::FirstWriteHome,
&xid(1),
Some(&tablet),
&BTreeMap::new(),
)
.unwrap();
assert_eq!(selected, gid(42));
assert!(select_coordinator_group(
&CoordinatorSelection::FirstWriteHome,
&xid(1),
None,
&BTreeMap::new()
)
.is_err());
}
#[test]
fn command_records_round_trip_and_fail_closed_on_bad_versions() {
let coordinator = CoordinatorCommandRecord::new(CoordinatorCommand::Begin {
record: record(xid(1), vec![participant(91, 1)]),
});
let decoded = CoordinatorCommandRecord::decode(&coordinator.encode().unwrap()).unwrap();
assert_eq!(decoded, coordinator);
let intent_record = IntentCommandRecord::new(IntentCommand::PersistIntents {
txn_id: xid(1),
expected_schema_version: SchemaVersion::new(4),
expected_authz_version: 2,
prepare_ts: ts(9),
intents: vec![intent(1, b"k", 9)],
});
let decoded = IntentCommandRecord::decode(&intent_record.encode().unwrap()).unwrap();
assert_eq!(decoded, intent_record);
let future = CoordinatorCommandRecord {
format_version: DIST_TXN_RECORD_FORMAT_VERSION + 1,
command: CoordinatorCommand::Abort {
txn_id: xid(1),
reason: AbortReason::RolledBack,
},
};
let error = CoordinatorCommandRecord::decode(&future.encode().unwrap()).unwrap_err();
assert!(matches!(
error,
DistTxnDecodeError::UnsupportedVersion { .. }
));
assert!(matches!(
CoordinatorCommandRecord::decode(b"not json"),
Err(DistTxnDecodeError::Malformed(_))
));
}
#[test]
fn protocol_types_serde_round_trip() {
let states = vec![
DistributedTxnState::Pending,
DistributedTxnState::Preparing,
DistributedTxnState::Committed { commit_ts: ts(5) },
DistributedTxnState::Aborted {
reason: AbortReason::Conflict("conflict".to_owned()),
},
];
for state in states {
let json = serde_json::to_string(&state).unwrap();
assert_eq!(
serde_json::from_str::<DistributedTxnState>(&json).unwrap(),
state
);
}
let mut full = record(xid(2), vec![participant(91, 1), participant(92, 2)]);
full.prepare_ts.insert(tid(1), ts(11));
full.prepare_ts.insert(tid(2), ts(12));
full.state = DistributedTxnState::Preparing;
full.max_observed = ts(10);
let json = serde_json::to_string(&full).unwrap();
assert_eq!(serde_json::from_str::<TxnRecord>(&json).unwrap(), full);
}
#[test]
fn status_sink_begin_is_idempotent_and_terminal_states_are_final() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = TxnStatusApplySink::open(tmp.path()).unwrap();
let txn = xid(1);
let begin = CoordinatorCommand::Begin {
record: record(txn, vec![participant(91, 1)]),
};
sink.apply(&applied(
1,
coordinator_envelope([1u8; 16], begin.clone(), 100),
))
.unwrap();
sink.apply(&applied(2, coordinator_envelope([2u8; 16], begin, 101)))
.unwrap();
assert!(sink.state().rejections.is_empty());
let mut conflicting = record(txn, vec![participant(91, 1)]);
conflicting.idempotency_key = [9u8; 16];
sink.apply(&applied(
3,
coordinator_envelope(
[3u8; 16],
CoordinatorCommand::Begin {
record: conflicting,
},
102,
),
))
.unwrap();
assert_eq!(sink.state().rejections.len(), 1);
assert!(matches!(
sink.state().rejections[0].reason,
StatusRejectionReason::IdempotencyKeyConflict { existing } if existing == [7u8; 16]
));
sink.apply(&applied(
4,
coordinator_envelope(
[4u8; 16],
CoordinatorCommand::MarkPreparing {
txn_id: txn,
tablet_id: tid(1),
prepare_ts: ts(150),
observed: ts(160),
},
150,
),
))
.unwrap();
sink.apply(&applied(
5,
coordinator_envelope(
[5u8; 16],
CoordinatorCommand::Commit {
txn_id: txn,
commit_ts: ts(200),
},
200,
),
))
.unwrap();
let stored = sink.record(&txn).unwrap();
assert_eq!(
stored.state,
DistributedTxnState::Committed { commit_ts: ts(200) }
);
assert_eq!(stored.prepare_ts[&tid(1)], ts(150));
assert_eq!(stored.max_observed, ts(160));
sink.apply(&applied(
6,
coordinator_envelope(
[6u8; 16],
CoordinatorCommand::Abort {
txn_id: txn,
reason: AbortReason::RolledBack,
},
300,
),
))
.unwrap();
assert_eq!(
sink.record(&txn).unwrap().state,
DistributedTxnState::Committed { commit_ts: ts(200) }
);
assert!(matches!(
sink.state().rejections[1].reason,
StatusRejectionReason::DecisionFinal { .. }
));
sink.apply(&applied(
7,
coordinator_envelope(
[7u8; 16],
CoordinatorCommand::Commit {
txn_id: xid(99),
commit_ts: ts(400),
},
400,
),
))
.unwrap();
assert!(matches!(
sink.state().rejections[2].reason,
StatusRejectionReason::UnknownTxn(_)
));
let mut bad = record(xid(5), vec![participant(91, 1)]);
bad.state = DistributedTxnState::Committed { commit_ts: ts(1) };
assert!(sink
.apply(&applied(
8,
coordinator_envelope([8u8; 16], CoordinatorCommand::Begin { record: bad }, 500,)
))
.is_err());
}
#[test]
fn status_sink_rejects_commit_ts_not_after_max_observed() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = TxnStatusApplySink::open(tmp.path()).unwrap();
let txn = xid(42);
sink.apply(&applied(
1,
coordinator_envelope(
[1u8; 16],
CoordinatorCommand::Begin {
record: record(txn, vec![participant(91, 1)]),
},
100,
),
))
.unwrap();
sink.apply(&applied(
2,
coordinator_envelope(
[2u8; 16],
CoordinatorCommand::MarkPreparing {
txn_id: txn,
tablet_id: tid(1),
prepare_ts: ts(150),
observed: ts(160),
},
150,
),
))
.unwrap();
sink.apply(&applied(
3,
coordinator_envelope(
[3u8; 16],
CoordinatorCommand::Commit {
txn_id: txn,
commit_ts: ts(160),
},
160,
),
))
.unwrap();
assert!(
matches!(
sink.state().rejections.back().map(|r| &r.reason),
Some(StatusRejectionReason::InvalidCommitTs { .. })
),
"expected InvalidCommitTs, got {:?}",
sink.state().rejections
);
assert!(matches!(
sink.record(&txn).unwrap().state,
DistributedTxnState::Preparing
));
sink.apply(&applied(
4,
coordinator_envelope(
[4u8; 16],
CoordinatorCommand::Commit {
txn_id: txn,
commit_ts: ts(161),
},
161,
),
))
.unwrap();
assert_eq!(
sink.record(&txn).unwrap().state,
DistributedTxnState::Committed { commit_ts: ts(161) }
);
}
#[test]
fn status_sink_heartbeat_extends_expiry_only_forward() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = TxnStatusApplySink::open(tmp.path()).unwrap();
let txn = xid(3);
sink.apply(&applied(
1,
coordinator_envelope(
[1u8; 16],
CoordinatorCommand::Begin {
record: record(txn, vec![participant(91, 1)]),
},
100,
),
))
.unwrap();
sink.apply(&applied(
2,
coordinator_envelope(
[2u8; 16],
CoordinatorCommand::Heartbeat {
txn_id: txn,
heartbeat: ts(1_500),
expiry: ts(2_500),
},
150,
),
))
.unwrap();
assert_eq!(sink.record(&txn).unwrap().heartbeat, ts(1_500));
assert_eq!(sink.record(&txn).unwrap().expiry, ts(2_500));
sink.apply(&applied(
3,
coordinator_envelope(
[3u8; 16],
CoordinatorCommand::Heartbeat {
txn_id: txn,
heartbeat: ts(1_200),
expiry: ts(2_200),
},
151,
),
))
.unwrap();
assert_eq!(sink.record(&txn).unwrap().heartbeat, ts(1_500));
}
#[test]
fn status_sink_checkpoint_survives_reopen_and_replay_is_skipped() {
let tmp = tempfile::tempdir().unwrap();
let txn = xid(4);
{
let mut sink = TxnStatusApplySink::open(tmp.path()).unwrap();
sink.apply(&applied(
1,
coordinator_envelope(
[1u8; 16],
CoordinatorCommand::Begin {
record: record(txn, vec![participant(91, 1)]),
},
100,
),
))
.unwrap();
}
let mut sink = TxnStatusApplySink::open(tmp.path()).unwrap();
assert!(sink.record(&txn).is_some());
sink.apply(&applied(
1,
coordinator_envelope(
[9u8; 16],
CoordinatorCommand::Abort {
txn_id: txn,
reason: AbortReason::RolledBack,
},
101,
),
))
.unwrap();
assert_eq!(
sink.record(&txn).unwrap().state,
DistributedTxnState::Pending
);
let foreign = ReplicatedCommand::new(
CommandKind::Catalog,
CommandEnvelope::new(COMMAND_TYPE_DIST_TXN_COORDINATOR, [5u8; 16], vec![1]),
ts(600),
);
assert!(sink.apply(&applied(2, foreign)).is_err());
let wrong_type = ReplicatedCommand::new(
CommandKind::Transaction,
CommandEnvelope::new(1, [6u8; 16], vec![1]),
ts(601),
);
assert!(sink.apply(&applied(2, wrong_type)).is_err());
}
fn persist(
txn: u8,
schema: u64,
authz: u64,
prepare_micros: u64,
keys: &[&[u8]],
) -> IntentCommand {
IntentCommand::PersistIntents {
txn_id: xid(txn),
expected_schema_version: SchemaVersion::new(schema),
expected_authz_version: authz,
prepare_ts: ts(prepare_micros),
intents: keys
.iter()
.map(|key| intent(txn, key, prepare_micros))
.collect(),
}
}
#[test]
fn intent_sink_validates_versions_and_conflicts() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = IntentApplySink::open(tmp.path()).unwrap();
sink.apply(&applied(
1,
intent_envelope(
[1u8; 16],
IntentCommand::SetTabletVersions {
schema_version: SchemaVersion::new(7),
authz_version: 3,
},
100,
),
))
.unwrap();
sink.apply(&applied(
2,
intent_envelope([2u8; 16], persist(1, 0, 3, 110, &[b"k1"]), 110),
))
.unwrap();
assert!(matches!(
sink.state().rejections[0].reason,
PrepareRejectionReason::StaleSchemaVersion { .. }
));
sink.apply(&applied(
3,
intent_envelope([3u8; 16], persist(1, 7, 0, 111, &[b"k1"]), 111),
))
.unwrap();
assert!(matches!(
sink.state().rejections[1].reason,
PrepareRejectionReason::StaleAuthzVersion { .. }
));
sink.apply(&applied(
4,
intent_envelope([4u8; 16], persist(1, 7, 3, 112, &[b"k1", b"k2"]), 112),
))
.unwrap();
assert!(sink.txn(&xid(1)).is_some());
sink.apply(&applied(
5,
intent_envelope([5u8; 16], persist(2, 7, 3, 113, &[b"k1"]), 113),
))
.unwrap();
assert!(matches!(
sink.state().rejections[2].reason,
PrepareRejectionReason::KeyConflict { holder, .. } if holder == xid(1)
));
assert!(sink.txn(&xid(2)).is_none());
sink.apply(&applied(
6,
intent_envelope([6u8; 16], persist(2, 7, 3, 114, &[b"k9"]), 114),
))
.unwrap();
assert!(sink.txn(&xid(2)).is_some());
}
#[test]
fn intent_sink_replay_keeps_the_original_prepare_timestamp() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = IntentApplySink::open(tmp.path()).unwrap();
sink.apply(&applied(
1,
intent_envelope([1u8; 16], persist(1, 0, 0, 100, &[b"k1"]), 100),
))
.unwrap();
sink.apply(&applied(
2,
intent_envelope([2u8; 16], persist(1, 0, 0, 999, &[b"k1"]), 999),
))
.unwrap();
assert_eq!(sink.txn(&xid(1)).unwrap().prepare_ts, ts(100));
assert!(sink.state().rejections.is_empty());
sink.apply(&applied(
3,
intent_envelope([3u8; 16], persist(1, 0, 0, 1000, &[b"k2"]), 1000),
))
.unwrap();
assert!(matches!(
sink.state().rejections[0].reason,
PrepareRejectionReason::PayloadMismatch
));
}
#[test]
fn intent_sink_resolution_commits_visible_or_removes_and_never_conflicts() {
let tmp = tempfile::tempdir().unwrap();
let mut sink = IntentApplySink::open(tmp.path()).unwrap();
sink.apply(&applied(
1,
intent_envelope([1u8; 16], persist(1, 0, 0, 100, &[b"k1", b"k2"]), 100),
))
.unwrap();
sink.apply(&applied(
2,
intent_envelope(
[2u8; 16],
IntentCommand::Resolve {
txn_id: xid(1),
decision: TxnDecision::Committed { commit_ts: ts(200) },
},
200,
),
))
.unwrap();
let writes = sink.state().committed_writes.clone();
assert_eq!(writes.len(), 2);
assert!(writes.iter().all(|write| write.commit_ts == ts(200)));
assert!(sink.txn(&xid(1)).unwrap().intents.is_empty());
sink.apply(&applied(
3,
intent_envelope(
[3u8; 16],
IntentCommand::Resolve {
txn_id: xid(1),
decision: TxnDecision::Committed { commit_ts: ts(200) },
},
201,
),
))
.unwrap();
assert_eq!(sink.state().committed_writes.len(), 2);
assert!(sink
.apply(&applied(
4,
intent_envelope(
[4u8; 16],
IntentCommand::Resolve {
txn_id: xid(1),
decision: TxnDecision::Aborted {
reason: AbortReason::RolledBack,
},
},
202,
),
))
.is_err());
sink.apply(&applied(
5,
intent_envelope([5u8; 16], persist(1, 0, 0, 300, &[b"k3"]), 300),
))
.unwrap();
assert!(matches!(
sink.state().rejections[0].reason,
PrepareRejectionReason::AlreadyResolved { .. }
));
sink.apply(&applied(
6,
intent_envelope([6u8; 16], persist(2, 0, 0, 400, &[b"k1"]), 400),
))
.unwrap();
sink.apply(&applied(
7,
intent_envelope(
[7u8; 16],
IntentCommand::Resolve {
txn_id: xid(2),
decision: TxnDecision::Aborted {
reason: AbortReason::Conflict("test".to_owned()),
},
},
401,
),
))
.unwrap();
assert!(sink.txn(&xid(2)).unwrap().intents.is_empty());
assert_eq!(sink.state().committed_writes.len(), 2);
sink.apply(&applied(
8,
intent_envelope([8u8; 16], persist(3, 0, 0, 500, &[b"k1"]), 500),
))
.unwrap();
assert!(sink.txn(&xid(3)).unwrap().resolution.is_none());
sink.apply(&applied(
9,
intent_envelope(
[9u8; 16],
IntentCommand::Resolve {
txn_id: xid(9),
decision: TxnDecision::Aborted {
reason: AbortReason::Cancelled("expired".to_owned()),
},
},
600,
),
))
.unwrap();
sink.apply(&applied(
10,
intent_envelope([10u8; 16], persist(9, 0, 0, 700, &[b"k7"]), 700),
))
.unwrap();
assert!(matches!(
sink.state().rejections[1].reason,
PrepareRejectionReason::AlreadyResolved { .. }
));
assert!(sink.unresolved_txn_ids().contains(&xid(3)));
assert!(!sink.unresolved_txn_ids().contains(&xid(2)));
}
#[test]
fn intent_sink_checkpoint_survives_reopen() {
let tmp = tempfile::tempdir().unwrap();
{
let mut sink = IntentApplySink::open(tmp.path()).unwrap();
sink.apply(&applied(
1,
intent_envelope([1u8; 16], persist(1, 0, 0, 100, &[b"k1"]), 100),
))
.unwrap();
}
let sink = IntentApplySink::open(tmp.path()).unwrap();
assert_eq!(sink.txn(&xid(1)).unwrap().prepare_ts, ts(100));
assert_eq!(sink.applied_position(), LogPosition { term: 1, index: 1 });
}
const LEADER_TIMEOUT: Duration = Duration::from_secs(60);
const STATUS_GROUP: u8 = 90;
const P1_GROUP: u8 = 91;
const P2_GROUP: u8 = 92;
const TEN_MINUTES: Duration = Duration::from_secs(600);
fn raft_id(group: u8, member: u8) -> RaftNodeId {
u64::from(group) * 100 + u64::from(member)
}
fn fast_group_config(dir: &Path, group: u8, member: u8) -> GroupConfig {
let mut config =
GroupConfig::new("dist-txn-test", raft_id(group, member), dir.to_path_buf());
config.heartbeat_interval = Duration::from_millis(50);
config.election_timeout_min = Duration::from_millis(150);
config.election_timeout_max = Duration::from_millis(300);
config.install_snapshot_timeout = Duration::from_millis(1_000);
config
}
struct Cell {
tmp: tempfile::TempDir,
transport: Arc<InMemoryTransport>,
status: Vec<TxnStatusGroup<InMemoryTransport>>,
participants: BTreeMap<RaftGroupId, Vec<IntentGroup<InMemoryTransport>>>,
}
impl Cell {
fn p1(&self) -> &[IntentGroup<InMemoryTransport>] {
&self.participants[&gid(P1_GROUP)]
}
fn p2(&self) -> &[IntentGroup<InMemoryTransport>] {
&self.participants[&gid(P2_GROUP)]
}
async fn shutdown(self) {
for member in &self.status {
let _ = member.shutdown().await;
}
for members in self.participants.values() {
for member in members {
let _ = member.shutdown().await;
}
}
}
}
async fn boot_status_member(
cell_dir: &Path,
member: u8,
transport: &Arc<InMemoryTransport>,
) -> TxnStatusGroup<InMemoryTransport> {
TxnStatusGroup::create(
fast_group_config(
&cell_dir.join(format!("status-{member}")),
STATUS_GROUP,
member,
),
gid(STATUS_GROUP),
transport.clone(),
)
.await
.unwrap()
}
async fn boot_intent_member(
cell_dir: &Path,
group: u8,
member: u8,
transport: &Arc<InMemoryTransport>,
) -> IntentGroup<InMemoryTransport> {
IntentGroup::create(
fast_group_config(
&cell_dir.join(format!("intent-{group}-{member}")),
group,
member,
),
gid(group),
transport.clone(),
)
.await
.unwrap()
}
fn member_addresses(group: u8) -> Vec<(RaftNodeId, String)> {
(1..=3_u8)
.map(|member| {
(
raft_id(group, member),
format!(
"127.0.0.1:{}",
9_000 + u16::from(group) * 10 + u16::from(member)
),
)
})
.collect()
}
async fn wait_leader_among<M: GroupMember>(members: &[&M]) -> RaftNodeId {
let allowed: std::collections::BTreeSet<RaftNodeId> =
members.iter().map(|member| member.node_id()).collect();
let deadline = Instant::now() + LEADER_TIMEOUT;
loop {
let mut leaders = std::collections::BTreeSet::new();
let mut seen = 0_usize;
for member in members {
if let Some(leader) = member.current_leader() {
leaders.insert(leader);
seen += 1;
}
}
if seen == members.len() && leaders.len() == 1 {
let leader = *leaders.iter().next().expect("one leader");
if allowed.contains(&leader) {
return leader;
}
}
assert!(
Instant::now() < deadline,
"no consensus leader (saw {leaders:?})"
);
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
async fn leader_index<M: GroupMember>(members: &[M]) -> usize {
let member_refs: Vec<&M> = members.iter().collect();
let leader = wait_leader_among(&member_refs).await;
members
.iter()
.position(|member| member.node_id() == leader)
.expect("a settled group has a leader member")
}
async fn wait_until<F, Fut>(what: &str, mut predicate: F)
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = bool>,
{
let deadline = Instant::now() + LEADER_TIMEOUT;
loop {
if predicate().await {
return;
}
assert!(Instant::now() < deadline, "timed out waiting for {what}");
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
async fn wait_cell_converged(cell: &Cell) {
wait_until("status convergence", || async {
let reference = cell.status[0].state();
cell.status.iter().all(|member| member.state() == reference)
})
.await;
for group in [P1_GROUP, P2_GROUP] {
let members = &cell.participants[&gid(group)];
wait_until("participant convergence", || async {
let reference = members[0].state();
members.iter().all(|member| member.state() == reference)
})
.await;
}
}
async fn boot_cell() -> Cell {
let tmp = tempfile::tempdir().unwrap();
let transport = Arc::new(InMemoryTransport::new());
let mut status = Vec::new();
for member in 1..=3_u8 {
status.push(boot_status_member(tmp.path(), member, &transport).await);
}
status[0]
.bootstrap(&member_addresses(STATUS_GROUP))
.await
.unwrap();
let mut participants: BTreeMap<RaftGroupId, Vec<IntentGroup<InMemoryTransport>>> =
BTreeMap::new();
for group in [P1_GROUP, P2_GROUP] {
let mut members = Vec::new();
for member in 1..=3_u8 {
members.push(boot_intent_member(tmp.path(), group, member, &transport).await);
}
members[0]
.bootstrap(&member_addresses(group))
.await
.unwrap();
participants.insert(gid(group), members);
}
wait_leader_among(&status.iter().collect::<Vec<_>>()).await;
for group in [P1_GROUP, P2_GROUP] {
wait_leader_among(&participants[&gid(group)].iter().collect::<Vec<_>>()).await;
}
Cell {
tmp,
transport,
status,
participants,
}
}
fn driver_config(pending_timeout: Duration) -> DistTxnConfig {
let mut status_partitions = BTreeMap::new();
status_partitions.insert(
0,
TxnStatusPartition {
partition_id: 0,
home_raft_group: gid(STATUS_GROUP),
},
);
DistTxnConfig {
selection: CoordinatorSelection::TxnIdDerived,
status_partitions,
pending_timeout,
..DistTxnConfig::default()
}
}
fn participant_writes(group: u8, tablet: u8, keys: &[&[u8]]) -> ParticipantWrites {
ParticipantWrites {
participant: participant(group, tablet),
expected_schema_version: SchemaVersion::ZERO,
expected_authz_version: 0,
intents: keys
.iter()
.map(|key| WriteIntent {
txn_id: TransactionId::ZERO,
key: key.to_vec(),
value_ref: format!("staged-{}", String::from_utf8_lossy(key)).into_bytes(),
prepare_ts: HlcTimestamp::ZERO,
})
.collect(),
}
}
fn commit_request(txn: u8, key: u8, writes: Vec<ParticipantWrites>) -> CommitRequest {
CommitRequest {
txn_id: xid(txn),
idempotency_key: [key; 16],
writes,
observed: Vec::new(),
first_write_tablet: None,
}
}
#[tokio::test]
async fn single_participant_commit_fast_path() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(1, 1, vec![participant_writes(P1_GROUP, 1, &[b"alpha"])]);
let outcome = driver
.commit(&cell.status, &cell.participants, request, &control)
.await
.unwrap();
wait_cell_converged(&cell).await;
assert_eq!(outcome.txn_id, xid(1));
assert_eq!(outcome.durability, DurabilityLevel::Quorum);
assert_eq!(outcome.participants, vec![participant(P1_GROUP, 1)]);
let stored = cell.status[0].record(&xid(1)).unwrap();
assert_eq!(
stored.state,
DistributedTxnState::Committed {
commit_ts: outcome.commit_ts
}
);
assert!(stored.prepare_ts.is_empty());
let participant_txn = cell.p1()[0].txn(&xid(1)).unwrap();
assert_eq!(
participant_txn.resolution,
Some(TxnDecision::Committed {
commit_ts: outcome.commit_ts
})
);
assert!(participant_txn.prepare_ts > HlcTimestamp::ZERO);
assert!(outcome.commit_ts > participant_txn.prepare_ts);
let writes = cell.p1()[0].committed_writes();
assert_eq!(writes.len(), 1);
assert_eq!(writes[0].key, b"alpha");
assert_eq!(writes[0].commit_ts, outcome.commit_ts);
assert!(cell.p1()[0].unresolved_txn_ids().is_empty());
cell.shutdown().await;
}
#[tokio::test]
async fn two_participant_commit_is_atomic_and_commit_ts_exceeds_all_prepare_ts() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let mut request = commit_request(
2,
2,
vec![
participant_writes(P1_GROUP, 1, &[b"a"]),
participant_writes(P2_GROUP, 2, &[b"b", b"c"]),
],
);
request.observed = vec![ts(42)];
let outcome = driver
.commit(&cell.status, &cell.participants, request, &control)
.await
.unwrap();
wait_cell_converged(&cell).await;
let stored = cell.status[0].record(&xid(2)).unwrap();
assert_eq!(stored.prepare_ts.len(), 2);
assert_eq!(
stored.state,
DistributedTxnState::Committed {
commit_ts: outcome.commit_ts
}
);
for prepare_ts in stored.prepare_ts.values() {
assert!(outcome.commit_ts > *prepare_ts);
}
assert!(outcome.commit_ts > ts(42));
assert!(stored.max_observed >= ts(42));
for (group, keys) in [
(P1_GROUP, vec![b"a".to_vec()]),
(P2_GROUP, vec![b"b".to_vec(), b"c".to_vec()]),
] {
let members = &cell.participants[&gid(group)];
let txn = members[0].txn(&xid(2)).unwrap();
assert_eq!(
txn.resolution,
Some(TxnDecision::Committed {
commit_ts: outcome.commit_ts
})
);
assert!(txn.prepare_ts < outcome.commit_ts);
assert!(members[0].unresolved_txn_ids().is_empty());
let writes = members[0].committed_writes();
let visible: Vec<&[u8]> = writes
.iter()
.filter(|write| write.txn_id == xid(2))
.map(|write| write.key.as_slice())
.collect();
assert_eq!(visible.len(), keys.len());
for key in &keys {
assert!(visible.contains(&key.as_slice()));
}
}
cell.shutdown().await;
}
#[tokio::test]
async fn prepare_conflict_aborts_everywhere_without_intent_leaks() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request_a = commit_request(10, 10, vec![participant_writes(P1_GROUP, 1, &[b"shared"])]);
driver
.begin(&cell.status, &request_a, &control)
.await
.unwrap();
driver
.prepare_participant(cell.p1(), &request_a.txn_id, &request_a.writes[0], &control)
.await
.unwrap();
let request_b = commit_request(
11,
11,
vec![
participant_writes(P2_GROUP, 2, &[b"b-key"]),
participant_writes(P1_GROUP, 1, &[b"shared"]),
],
);
let error = driver
.commit(&cell.status, &cell.participants, request_b, &control)
.await
.unwrap_err();
wait_cell_converged(&cell).await;
assert!(matches!(
error,
DistTxnError::Aborted(AbortReason::Conflict(_))
));
let record_b = cell.status[0].record(&xid(11)).unwrap();
assert!(matches!(
record_b.state,
DistributedTxnState::Aborted {
reason: AbortReason::Conflict(_)
}
));
assert!(
cell.p1()[0].txn(&xid(11)).is_none()
|| cell.p1()[0].txn(&xid(11)).unwrap().resolution.is_some()
);
assert!(cell.p2()[0].unresolved_txn_ids().is_empty());
assert!(cell.p2()[0]
.committed_writes()
.iter()
.all(|write| write.txn_id != xid(11)));
let txn_a = cell.p1()[0].txn(&xid(10)).unwrap();
assert!(txn_a.resolution.is_none());
assert!(cell.p1()[0].committed_writes().is_empty());
let final_state = driver
.abort(
&cell.status,
&cell.participants,
&xid(10),
AbortReason::RolledBack,
&control,
)
.await
.unwrap();
assert!(matches!(final_state, DistributedTxnState::Aborted { .. }));
wait_until("abort resolution", || async {
cell.p1()[0].unresolved_txn_ids().is_empty()
})
.await;
cell.shutdown().await;
}
#[tokio::test]
async fn coordinator_leader_kill_between_prepare_and_decision_recovers() {
let mut cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
20,
20,
vec![
participant_writes(P1_GROUP, 1, &[b"k1"]),
participant_writes(P2_GROUP, 2, &[b"k2"]),
],
);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let token1 = driver
.prepare_participant(cell.p1(), &request.txn_id, &request.writes[0], &control)
.await
.unwrap();
let token2 = driver
.prepare_participant(cell.p2(), &request.txn_id, &request.writes[1], &control)
.await
.unwrap();
let leader = leader_index(&cell.status).await;
let victim = cell.status.remove(leader);
victim.crash().await;
let survivor_refs: Vec<&TxnStatusGroup<InMemoryTransport>> = cell.status.iter().collect();
wait_leader_among(&survivor_refs).await;
let recovery_config = DistTxnConfig {
node_tiebreaker: 99,
..driver_config(TEN_MINUTES)
};
let recovery = DistTxnDriver::new(recovery_config);
let now = recovery.now().unwrap();
let outcome = recovery
.recover(
&cell.status,
&cell.participants,
&request.txn_id,
now,
&control,
)
.await
.unwrap();
let commit_ts = match outcome {
RecoveryOutcome::Decided(DistributedTxnState::Committed { commit_ts }) => commit_ts,
other => panic!("expected a committed recovery, got {other:?}"),
};
assert!(commit_ts > token1.prepare_ts);
assert!(commit_ts > token2.prepare_ts);
for group in [P1_GROUP, P2_GROUP] {
let members = &cell.participants[&gid(group)];
wait_until("participant resolution", || async {
members[0]
.txn(&request.txn_id)
.is_some_and(|txn| txn.resolution.is_some())
})
.await;
let txn = members[0].txn(&request.txn_id).unwrap();
assert_eq!(txn.resolution, Some(TxnDecision::Committed { commit_ts }));
assert!(members[0]
.committed_writes()
.iter()
.any(|write| write.txn_id == request.txn_id && write.commit_ts == commit_ts));
}
let healed = boot_status_member(
cell.tmp.path(),
u8::try_from(leader + 1).unwrap(),
&cell.transport,
)
.await;
wait_until("healed status member", || async {
healed.record(&request.txn_id).is_some_and(|record| {
record.state == (DistributedTxnState::Committed { commit_ts })
})
})
.await;
cell.status.push(healed);
cell.shutdown().await;
}
#[tokio::test]
async fn participant_kill_post_prepare_heals_and_resolves() {
let mut cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
21,
21,
vec![
participant_writes(P1_GROUP, 1, &[b"p1-key"]),
participant_writes(P2_GROUP, 2, &[b"p2-key"]),
],
);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let token1 = driver
.prepare_participant(cell.p1(), &request.txn_id, &request.writes[0], &control)
.await
.unwrap();
let token2 = driver
.prepare_participant(cell.p2(), &request.txn_id, &request.writes[1], &control)
.await
.unwrap();
let p1_leader = leader_index(cell.p1()).await;
let victim = cell
.participants
.get_mut(&gid(P1_GROUP))
.unwrap()
.remove(p1_leader);
victim.crash().await;
let survivor_refs: Vec<&IntentGroup<InMemoryTransport>> = cell.p1().iter().collect();
wait_leader_among(&survivor_refs).await;
let outcome = driver
.decide_commit(
&cell.status,
&cell.participants,
&request.txn_id,
&[token1, token2],
&request.observed,
&control,
)
.await
.unwrap();
wait_until("p1 survivor resolution", || async {
cell.p1().iter().all(|member| {
member
.committed_writes()
.iter()
.any(|write| write.txn_id == request.txn_id)
})
})
.await;
let healed = boot_intent_member(
cell.tmp.path(),
P1_GROUP,
u8::try_from(p1_leader + 1).unwrap(),
&cell.transport,
)
.await;
let expected = cell.p1()[0].state();
wait_until("healed p1 member", || async { healed.state() == expected }).await;
assert!(healed
.committed_writes()
.iter()
.any(|write| write.txn_id == request.txn_id && write.commit_ts == outcome.commit_ts));
cell.participants
.get_mut(&gid(P1_GROUP))
.unwrap()
.push(healed);
cell.shutdown().await;
}
#[tokio::test]
async fn expired_pending_is_pushed_to_abort_only_under_the_timeout_rule() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(30, 30, vec![participant_writes(P1_GROUP, 1, &[b"e-key"])]);
let begun = driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let pusher_config = DistTxnConfig {
node_tiebreaker: 77,
..driver_config(TEN_MINUTES)
};
let pusher = DistTxnDriver::new(pusher_config);
let renewed = driver
.heartbeat(&cell.status, begun.coordinator, &request.txn_id, &control)
.await
.unwrap();
assert!(renewed.heartbeat > begun.heartbeat);
assert!(renewed.expiry > begun.expiry);
wait_cell_converged(&cell).await;
let outcome = pusher
.push_expired(
&cell.status,
&cell.participants,
&request.txn_id,
one_before(renewed.expiry),
&control,
)
.await
.unwrap();
assert_eq!(outcome, PushOutcome::NotExpired);
assert_eq!(
cell.status[0].record(&request.txn_id).unwrap().state,
DistributedTxnState::Pending
);
let outcome = pusher
.push_expired(
&cell.status,
&cell.participants,
&request.txn_id,
renewed.expiry,
&control,
)
.await
.unwrap();
match outcome {
PushOutcome::Pushed(DistributedTxnState::Aborted {
reason: AbortReason::Cancelled(_),
}) => {}
other => panic!("expected a pushed abort, got {other:?}"),
}
wait_cell_converged(&cell).await;
let stored = cell.status[0].record(&request.txn_id).unwrap();
assert!(matches!(
stored.state,
DistributedTxnState::Aborted {
reason: AbortReason::Cancelled(_)
}
));
wait_until("abort tombstone", || async {
cell.p1()[0]
.txn(&request.txn_id)
.is_some_and(|txn| txn.resolution.is_some())
})
.await;
assert!(cell.p1()[0].committed_writes().is_empty());
assert!(cell.p1()[0].unresolved_txn_ids().is_empty());
let again = pusher
.push_expired(
&cell.status,
&cell.participants,
&request.txn_id,
renewed.expiry,
&control,
)
.await
.unwrap();
assert!(matches!(again, PushOutcome::Terminal(_)));
cell.shutdown().await;
}
#[tokio::test]
async fn duplicate_client_retry_yields_one_outcome() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
40,
40,
vec![
participant_writes(P1_GROUP, 1, &[b"r1"]),
participant_writes(P2_GROUP, 2, &[b"r2"]),
],
);
let first = driver
.commit(&cell.status, &cell.participants, request.clone(), &control)
.await
.unwrap();
let second = driver
.commit(&cell.status, &cell.participants, request, &control)
.await
.unwrap();
wait_cell_converged(&cell).await;
assert_eq!(first, second);
assert_eq!(cell.status[0].state().records.len(), 1);
for group in [P1_GROUP, P2_GROUP] {
let members = &cell.participants[&gid(group)];
assert_eq!(members[0].state().txns.len(), 1);
assert_eq!(
members[0]
.committed_writes()
.iter()
.filter(|write| write.txn_id == xid(40))
.count(),
1
);
}
let conflicting = commit_request(40, 41, vec![participant_writes(P1_GROUP, 1, &[b"r1"])]);
let error = driver
.commit(&cell.status, &cell.participants, conflicting, &control)
.await
.unwrap_err();
assert!(matches!(error, DistTxnError::IdempotencyConflict(_)));
cell.shutdown().await;
}
#[tokio::test]
async fn orphan_intent_sweep_honors_the_timeout_rules() {
let cell = boot_cell().await;
let control = ExecutionControl::default();
let live_driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let expiring_driver = DistTxnDriver::new(driver_config(Duration::from_millis(1)));
let request_c = commit_request(50, 50, vec![participant_writes(P1_GROUP, 1, &[b"c-key"])]);
live_driver
.begin(&cell.status, &request_c, &control)
.await
.unwrap();
let token_c = live_driver
.prepare_participant(cell.p1(), &request_c.txn_id, &request_c.writes[0], &control)
.await
.unwrap();
let outcome_c = live_driver
.decide_commit(
&cell.status,
&BTreeMap::<RaftGroupId, Vec<IntentGroup<InMemoryTransport>>>::new(),
&request_c.txn_id,
&[token_c],
&[],
&control,
)
.await
.unwrap();
let request_e = commit_request(51, 51, vec![participant_writes(P1_GROUP, 1, &[b"e-key"])]);
expiring_driver
.begin(&cell.status, &request_e, &control)
.await
.unwrap();
expiring_driver
.prepare_participant(cell.p1(), &request_e.txn_id, &request_e.writes[0], &control)
.await
.unwrap();
let request_l = commit_request(52, 52, vec![participant_writes(P1_GROUP, 1, &[b"l-key"])]);
live_driver
.begin(&cell.status, &request_l, &control)
.await
.unwrap();
live_driver
.prepare_participant(cell.p1(), &request_l.txn_id, &request_l.writes[0], &control)
.await
.unwrap();
let request_x = commit_request(53, 53, vec![participant_writes(P1_GROUP, 1, &[b"x-key"])]);
live_driver
.prepare_participant(cell.p1(), &request_x.txn_id, &request_x.writes[0], &control)
.await
.unwrap();
let now = later(live_driver.now().unwrap(), Duration::from_secs(300));
let report = live_driver
.sweep_orphans(&cell.status, cell.p1(), now, &control)
.await
.unwrap();
wait_cell_converged(&cell).await;
assert_eq!(
report,
SweepReport {
resolved_committed: 1,
resolved_aborted: 1,
pushed: 1,
in_flight: 1,
unknown: 1,
}
);
assert!(cell.p1()[0]
.committed_writes()
.iter()
.any(
|write| write.txn_id == request_c.txn_id && write.commit_ts == outcome_c.commit_ts
));
assert!(matches!(
cell.status[0].record(&request_e.txn_id).unwrap().state,
DistributedTxnState::Aborted { .. }
));
let txn_e = cell.p1()[0].txn(&request_e.txn_id).unwrap();
assert!(txn_e.intents.is_empty());
assert!(matches!(
txn_e.resolution,
Some(TxnDecision::Aborted { .. })
));
let txn_l = cell.p1()[0].txn(&request_l.txn_id).unwrap();
assert!(txn_l.resolution.is_none());
assert!(!txn_l.intents.is_empty());
assert_eq!(
cell.status[0].record(&request_l.txn_id).unwrap().state,
DistributedTxnState::Pending
);
let txn_x = cell.p1()[0].txn(&request_x.txn_id).unwrap();
assert!(txn_x.resolution.is_none());
assert!(!txn_x.intents.is_empty());
assert!(cell.status[0].record(&request_x.txn_id).is_none());
cell.shutdown().await;
}
#[tokio::test]
async fn record_read_api_supports_linearizable_and_session_token_reads() {
let cell = boot_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(60, 60, vec![participant_writes(P1_GROUP, 1, &[b"s-key"])]);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let record = driver
.read_record(&cell.status, &request.txn_id, &control)
.await
.unwrap()
.unwrap();
assert_eq!(record.txn_id, request.txn_id);
let leader = leader_index(&cell.status).await;
let heartbeat = driver.now().unwrap();
let expiry = later(heartbeat, TEN_MINUTES);
let (receipt, rejection) = cell.status[leader]
.propose(
[0xAB; 16],
CoordinatorCommand::Heartbeat {
txn_id: request.txn_id,
heartbeat,
expiry,
},
&control,
)
.await
.unwrap();
assert!(rejection.is_none());
let token = cell.status[leader].session_token(&receipt);
let follower = (leader + 1) % cell.status.len();
let seen = cell.status[follower]
.record_consistent(
&request.txn_id,
&ReadConsistency::ReadYourWrites { token },
&control,
)
.await
.unwrap()
.unwrap();
assert_eq!(seen.heartbeat, heartbeat);
let result = cell.status[follower]
.record_consistent(&request.txn_id, &ReadConsistency::Linearizable, &control)
.await;
match cell.status[follower].group().metrics().current_leader {
Some(leader_id) if leader_id == cell.status[follower].group().node_id() => {
assert!(result.is_ok());
}
_ => assert!(result.is_err()),
}
cell.shutdown().await;
}
use mongreldb_consensus::engine_sink::{
open_engine_sink, testing as engine_testing, EngineGroupConfig,
};
use mongreldb_types::ids::{ClusterId, DatabaseId, NodeId};
const T1_GROUP: u8 = 81;
const T2_GROUP: u8 = 82;
struct EngineCell {
tmp: tempfile::TempDir,
transport: Arc<InMemoryTransport>,
status: Vec<TxnStatusGroup<InMemoryTransport>>,
tablets: BTreeMap<RaftGroupId, Vec<TabletTxnGroup<InMemoryTransport>>>,
}
impl EngineCell {
fn t1(&self) -> &[TabletTxnGroup<InMemoryTransport>] {
&self.tablets[&gid(T1_GROUP)]
}
fn t2(&self) -> &[TabletTxnGroup<InMemoryTransport>] {
&self.tablets[&gid(T2_GROUP)]
}
async fn shutdown(self) {
for member in &self.status {
let _ = member.shutdown().await;
}
for members in self.tablets.values() {
for member in members {
let _ = member.shutdown().await;
}
}
}
}
async fn boot_tablet_member(
cell_dir: &Path,
group: u8,
member: u8,
transport: &Arc<InMemoryTransport>,
bootstrap: bool,
) -> TabletTxnGroup<InMemoryTransport> {
let node_data = cell_dir.join(format!("tablet-{group}-{member}"));
let engine_config = EngineGroupConfig::new(
node_data,
gid(group),
ClusterId::from_bytes([1; 16]),
NodeId::from_bytes([group.wrapping_add(0x40); 16]),
DatabaseId::from_bytes([group; 16]),
);
let engine = open_engine_sink(&engine_config).unwrap();
let tablet = TabletTxnGroup::create(
fast_group_config(&engine_config.group_dir(), group, member),
gid(group),
transport.clone(),
engine,
)
.await
.unwrap();
if bootstrap {
tablet
.bootstrap(&[(
raft_id(group, member),
format!(
"127.0.0.1:{}",
7_000 + u16::from(group) * 10 + u16::from(member)
),
)])
.await
.unwrap();
tablet.group().wait_leader(LEADER_TIMEOUT).await.unwrap();
}
tablet
}
async fn create_table(tablet: &TabletTxnGroup<InMemoryTransport>) {
tablet
.group()
.propose(
CommandKind::Catalog,
engine_testing::create_i64_table_envelope(1, "t", 1),
&ExecutionControl::default(),
)
.await
.unwrap();
let db = tablet.engine().lock().unwrap().database().unwrap();
assert_eq!(db.table_id("t").unwrap(), 0);
}
fn engine_writes(group: u8, tablet: u8, values: &[i64]) -> ParticipantWrites {
ParticipantWrites {
participant: participant(group, tablet),
expected_schema_version: SchemaVersion::ZERO,
expected_authz_version: 0,
intents: values
.iter()
.map(|value| WriteIntent {
txn_id: TransactionId::ZERO,
key: value.to_le_bytes().to_vec(),
value_ref: engine_testing::staged_put_i64(0, &[*value]),
prepare_ts: HlcTimestamp::ZERO,
})
.collect(),
}
}
fn visible(tablet: &TabletTxnGroup<InMemoryTransport>) -> Vec<i64> {
let db = tablet.engine().lock().unwrap().database().unwrap();
engine_testing::visible_i64s(&db, "t")
}
async fn boot_engine_cell() -> EngineCell {
let tmp = tempfile::tempdir().unwrap();
let transport = Arc::new(InMemoryTransport::new());
let mut status = Vec::new();
for member in 1..=3_u8 {
status.push(boot_status_member(tmp.path(), member, &transport).await);
}
status[0]
.bootstrap(&member_addresses(STATUS_GROUP))
.await
.unwrap();
let mut tablets: BTreeMap<RaftGroupId, Vec<TabletTxnGroup<InMemoryTransport>>> =
BTreeMap::new();
for group in [T1_GROUP, T2_GROUP] {
let member = boot_tablet_member(tmp.path(), group, 1, &transport, true).await;
create_table(&member).await;
tablets.insert(gid(group), vec![member]);
}
wait_leader_among(&status.iter().collect::<Vec<_>>()).await;
EngineCell {
tmp,
transport,
status,
tablets,
}
}
#[tokio::test]
async fn engine_backed_commit_applies_rows_on_both_tablets_after_decision() {
let cell = boot_engine_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
101,
101,
vec![
engine_writes(T1_GROUP, 1, &[10, 20]),
engine_writes(T2_GROUP, 2, &[30, 40]),
],
);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let token1 = driver
.prepare_participant(cell.t1(), &request.txn_id, &request.writes[0], &control)
.await
.unwrap();
let token2 = driver
.prepare_participant(cell.t2(), &request.txn_id, &request.writes[1], &control)
.await
.unwrap();
assert_eq!(visible(&cell.t1()[0]), Vec::<i64>::new());
assert_eq!(visible(&cell.t2()[0]), Vec::<i64>::new());
let outcome = driver
.decide_commit(
&cell.status,
&cell.tablets,
&request.txn_id,
&[token1, token2],
&request.observed,
&control,
)
.await
.unwrap();
assert_eq!(visible(&cell.t1()[0]), vec![10, 20]);
assert_eq!(visible(&cell.t2()[0]), vec![30, 40]);
for (group, keys) in [(T1_GROUP, 2_usize), (T2_GROUP, 2)] {
let member = &cell.tablets[&gid(group)][0];
let txn = member.txn(&xid(101)).unwrap();
assert_eq!(
txn.resolution,
Some(TxnDecision::Committed {
commit_ts: outcome.commit_ts
})
);
assert!(txn.applied);
assert!(txn.resolved_at.is_some());
assert!(txn.intents.is_empty());
assert_eq!(
member
.committed_writes()
.iter()
.filter(|write| write.txn_id == xid(101))
.count(),
keys
);
}
cell.shutdown().await;
}
#[tokio::test]
async fn engine_backed_abort_leaves_zero_rows() {
let cell = boot_engine_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
102,
102,
vec![
engine_writes(T1_GROUP, 1, &[10, 20]),
engine_writes(T2_GROUP, 2, &[30, 40]),
],
);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
driver
.prepare_participant(cell.t1(), &request.txn_id, &request.writes[0], &control)
.await
.unwrap();
driver
.prepare_participant(cell.t2(), &request.txn_id, &request.writes[1], &control)
.await
.unwrap();
let state = driver
.abort(
&cell.status,
&cell.tablets,
&request.txn_id,
AbortReason::RolledBack,
&control,
)
.await
.unwrap();
assert!(matches!(state, DistributedTxnState::Aborted { .. }));
assert_eq!(visible(&cell.t1()[0]), Vec::<i64>::new());
assert_eq!(visible(&cell.t2()[0]), Vec::<i64>::new());
for group in [T1_GROUP, T2_GROUP] {
let member = &cell.tablets[&gid(group)][0];
let txn = member.txn(&xid(102)).unwrap();
assert!(matches!(txn.resolution, Some(TxnDecision::Aborted { .. })));
assert!(txn.applied);
assert!(txn.intents.is_empty());
assert!(member.committed_writes().is_empty());
}
cell.shutdown().await;
}
#[tokio::test]
async fn coordinator_kill_recovers_and_applies_identically() {
let mut cell = boot_engine_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
103,
103,
vec![
engine_writes(T1_GROUP, 1, &[10, 20]),
engine_writes(T2_GROUP, 2, &[30, 40]),
],
);
driver
.begin(&cell.status, &request, &control)
.await
.unwrap();
let token1 = driver
.prepare_participant(cell.t1(), &request.txn_id, &request.writes[0], &control)
.await
.unwrap();
let token2 = driver
.prepare_participant(cell.t2(), &request.txn_id, &request.writes[1], &control)
.await
.unwrap();
let leader = leader_index(&cell.status).await;
let victim = cell.status.remove(leader);
victim.crash().await;
let survivor_refs: Vec<&TxnStatusGroup<InMemoryTransport>> = cell.status.iter().collect();
wait_leader_among(&survivor_refs).await;
let recovery = DistTxnDriver::new(DistTxnConfig {
node_tiebreaker: 99,
..driver_config(TEN_MINUTES)
});
let now = recovery.now().unwrap();
let outcome = recovery
.recover(&cell.status, &cell.tablets, &request.txn_id, now, &control)
.await
.unwrap();
let commit_ts = match outcome {
RecoveryOutcome::Decided(DistributedTxnState::Committed { commit_ts }) => commit_ts,
other => panic!("expected a committed recovery, got {other:?}"),
};
assert!(commit_ts > token1.prepare_ts);
assert!(commit_ts > token2.prepare_ts);
assert_eq!(visible(&cell.t1()[0]), vec![10, 20]);
assert_eq!(visible(&cell.t2()[0]), vec![30, 40]);
cell.shutdown().await;
}
#[tokio::test]
async fn duplicate_resolve_never_double_applies() {
let cell = boot_engine_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
104,
104,
vec![
engine_writes(T1_GROUP, 1, &[10, 20]),
engine_writes(T2_GROUP, 2, &[30, 40]),
],
);
let first = driver
.commit(&cell.status, &cell.tablets, request.clone(), &control)
.await
.unwrap();
let second = driver
.commit(&cell.status, &cell.tablets, request, &control)
.await
.unwrap();
assert_eq!(first, second);
driver
.broadcast_resolve(
&cell.tablets,
&xid(104),
&second.participants,
TxnDecision::Committed {
commit_ts: second.commit_ts,
},
&control,
)
.await;
assert_eq!(visible(&cell.t1()[0]), vec![10, 20]);
assert_eq!(visible(&cell.t2()[0]), vec![30, 40]);
for group in [T1_GROUP, T2_GROUP] {
let member = &cell.tablets[&gid(group)][0];
assert_eq!(member.state().txns.len(), 1);
assert_eq!(
member
.committed_writes()
.iter()
.filter(|write| write.txn_id == xid(104))
.count(),
2
);
}
cell.shutdown().await;
}
#[tokio::test]
async fn tablet_restart_after_resolve_replays_without_double_apply() {
let mut cell = boot_engine_cell().await;
let driver = DistTxnDriver::new(driver_config(TEN_MINUTES));
let control = ExecutionControl::default();
let request = commit_request(
105,
105,
vec![
engine_writes(T1_GROUP, 1, &[10, 20]),
engine_writes(T2_GROUP, 2, &[30, 40]),
],
);
let outcome = driver
.commit(&cell.status, &cell.tablets, request, &control)
.await
.unwrap();
assert_eq!(visible(&cell.t1()[0]), vec![10, 20]);
let victim = cell.tablets.get_mut(&gid(T1_GROUP)).unwrap().remove(0);
victim.crash().await;
let healed = boot_tablet_member(cell.tmp.path(), T1_GROUP, 1, &cell.transport, false).await;
healed.group().wait_leader(LEADER_TIMEOUT).await.unwrap();
let txn = healed.txn(&xid(105)).unwrap();
assert!(txn.applied);
assert_eq!(
txn.resolution,
Some(TxnDecision::Committed {
commit_ts: outcome.commit_ts
})
);
assert_eq!(visible(&healed), vec![10, 20]);
cell.tablets.get_mut(&gid(T1_GROUP)).unwrap().push(healed);
driver
.broadcast_resolve(
&cell.tablets,
&xid(105),
&outcome.participants,
TxnDecision::Committed {
commit_ts: outcome.commit_ts,
},
&control,
)
.await;
assert_eq!(visible(&cell.t1()[0]), vec![10, 20]);
cell.shutdown().await;
}
#[tokio::test]
async fn resolved_sweep_bounds_tombstones_and_keeps_engine_rows() {
let cell = boot_engine_cell().await;
let config = DistTxnConfig {
resolved_retention: Duration::ZERO,
sweep_limit: 2,
..driver_config(TEN_MINUTES)
};
let driver = DistTxnDriver::new(config);
let control = ExecutionControl::default();
for (index, values) in [[10_i64, 20], [30, 40], [50, 60]].iter().enumerate() {
let txn = 110 + index as u8;
let request = commit_request(txn, txn, vec![engine_writes(T1_GROUP, 1, values)]);
driver
.commit(&cell.status, &cell.tablets, request, &control)
.await
.unwrap();
}
let aborted = commit_request(113, 113, vec![engine_writes(T1_GROUP, 1, &[70])]);
driver
.begin(&cell.status, &aborted, &control)
.await
.unwrap();
driver
.prepare_participant(cell.t1(), &aborted.txn_id, &aborted.writes[0], &control)
.await
.unwrap();
driver
.abort(
&cell.status,
&cell.tablets,
&aborted.txn_id,
AbortReason::RolledBack,
&control,
)
.await
.unwrap();
assert_eq!(cell.t1()[0].state().txns.len(), 4);
assert_eq!(visible(&cell.t1()[0]), vec![10, 20, 30, 40, 50, 60]);
let now = later(driver.now().unwrap(), Duration::from_secs(3_600));
let report = driver
.sweep_resolved(cell.t1(), now, &control)
.await
.unwrap();
assert_eq!(
report,
SweepResolvedReport {
swept: 2,
remaining: 2
}
);
assert_eq!(cell.t1()[0].state().txns.len(), 2);
let later_now = later(now, Duration::from_secs(1));
let report = driver
.sweep_resolved(cell.t1(), later_now, &control)
.await
.unwrap();
assert_eq!(
report,
SweepResolvedReport {
swept: 2,
remaining: 0
}
);
assert!(cell.t1()[0].state().txns.is_empty());
assert!(cell.t1()[0].committed_writes().is_empty());
assert_eq!(visible(&cell.t1()[0]), vec![10, 20, 30, 40, 50, 60]);
cell.shutdown().await;
}
#[tokio::test]
async fn parallel_prepare_respects_the_deadline() {
let tmp = tempfile::tempdir().unwrap();
let transport = Arc::new(InMemoryTransport::new());
let mut status = Vec::new();
for member in 1..=3_u8 {
status.push(boot_status_member(tmp.path(), member, &transport).await);
}
status[0]
.bootstrap(&member_addresses(STATUS_GROUP))
.await
.unwrap();
let t1 = boot_tablet_member(tmp.path(), T1_GROUP, 1, &transport, true).await;
create_table(&t1).await;
let mut t2_members = Vec::new();
for member in 1..=3_u8 {
t2_members
.push(boot_tablet_member(tmp.path(), T2_GROUP, member, &transport, false).await);
}
t2_members[0]
.bootstrap(&member_addresses(T2_GROUP))
.await
.unwrap();
t2_members.pop().unwrap().crash().await;
t2_members.pop().unwrap().crash().await;
let mut tablets: BTreeMap<RaftGroupId, Vec<TabletTxnGroup<InMemoryTransport>>> =
BTreeMap::new();
tablets.insert(gid(T1_GROUP), vec![t1]);
tablets.insert(gid(T2_GROUP), t2_members);
wait_leader_among(&status.iter().collect::<Vec<_>>()).await;
let driver = DistTxnDriver::new(driver_config(Duration::from_millis(1)));
let control = ExecutionControl {
deadline: Some(std::time::Instant::now() + Duration::from_millis(1_000)),
cancellation: None,
};
let request = commit_request(
120,
120,
vec![
engine_writes(T2_GROUP, 2, &[30]),
engine_writes(T1_GROUP, 1, &[10, 20]),
],
);
let error = driver
.commit(&status, &tablets, request.clone(), &control)
.await
.unwrap_err();
assert!(
matches!(error, DistTxnError::OutcomeAmbiguous { .. }),
"expected an ambiguous outcome after the deadline, got {error:?}"
);
{
let t1 = &tablets[&gid(T1_GROUP)][0];
assert_eq!(t1.unresolved_txn_ids(), vec![xid(120)]);
assert_eq!(visible(t1), Vec::<i64>::new());
}
let last_t2 = tablets.get_mut(&gid(T2_GROUP)).unwrap().remove(0);
last_t2.crash().await;
let pusher = DistTxnDriver::new(driver_config(Duration::from_millis(1)));
let now = later(pusher.now().unwrap(), Duration::from_secs(3_600));
let fresh = ExecutionControl::default();
let outcome = pusher
.push_expired(&status, &tablets, &request.txn_id, now, &fresh)
.await
.unwrap();
assert!(matches!(outcome, PushOutcome::Pushed(_)));
let t1 = &tablets[&gid(T1_GROUP)][0];
wait_until("t1 abort resolution", || async {
t1.txn(&xid(120))
.is_some_and(|txn| txn.resolution.is_some())
})
.await;
let txn = t1.txn(&xid(120)).unwrap();
assert!(matches!(txn.resolution, Some(TxnDecision::Aborted { .. })));
assert!(txn.applied);
assert!(txn.intents.is_empty());
assert_eq!(visible(t1), Vec::<i64>::new());
for member in &status {
let _ = member.shutdown().await;
}
for members in tablets.values() {
for member in members {
let _ = member.shutdown().await;
}
}
}
}