280 lines
8.9 KiB
Rust
280 lines
8.9 KiB
Rust
use std::collections::{HashMap, HashSet};
|
|
|
|
use bson::Document;
|
|
use dashmap::DashMap;
|
|
use tracing::{debug, warn};
|
|
use uuid::Uuid;
|
|
|
|
use rustdb_storage::StorageAdapter;
|
|
|
|
use crate::error::{TransactionError, TransactionResult};
|
|
|
|
/// The status of a transaction.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum TransactionStatus {
|
|
Active,
|
|
Committed,
|
|
Aborted,
|
|
}
|
|
|
|
/// Describes a write operation within a transaction.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum WriteOp {
|
|
Insert,
|
|
Update,
|
|
Delete,
|
|
}
|
|
|
|
/// A single write entry recorded within a transaction.
|
|
#[derive(Debug, Clone)]
|
|
pub struct WriteEntry {
|
|
pub op: WriteOp,
|
|
pub doc: Option<Document>,
|
|
pub original_doc: Option<Document>,
|
|
}
|
|
|
|
/// Full state of an in-flight transaction.
|
|
#[derive(Debug)]
|
|
pub struct TransactionState {
|
|
pub id: String,
|
|
pub session_id: String,
|
|
pub status: TransactionStatus,
|
|
/// Tracks which documents were read: namespace -> set of doc ids.
|
|
pub read_set: HashMap<String, HashSet<String>>,
|
|
/// Tracks writes: namespace -> (doc_id -> WriteEntry).
|
|
pub write_set: HashMap<String, HashMap<String, WriteEntry>>,
|
|
/// Snapshot of collections at transaction start: namespace -> documents.
|
|
pub snapshots: HashMap<String, Vec<Document>>,
|
|
}
|
|
|
|
/// Engine that manages transaction lifecycle and conflict detection.
|
|
pub struct TransactionEngine {
|
|
transactions: DashMap<String, TransactionState>,
|
|
}
|
|
|
|
impl TransactionEngine {
|
|
/// Create a new transaction engine.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
transactions: DashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Start a new transaction for the given session.
|
|
/// Returns a unique transaction id (UUID v4).
|
|
pub fn start_transaction(&self, session_id: &str) -> TransactionResult<String> {
|
|
let txn_id = Uuid::new_v4().to_string();
|
|
debug!(txn_id = %txn_id, session_id = %session_id, "starting transaction");
|
|
|
|
let state = TransactionState {
|
|
id: txn_id.clone(),
|
|
session_id: session_id.to_string(),
|
|
status: TransactionStatus::Active,
|
|
read_set: HashMap::new(),
|
|
write_set: HashMap::new(),
|
|
snapshots: HashMap::new(),
|
|
};
|
|
|
|
self.transactions.insert(txn_id.clone(), state);
|
|
Ok(txn_id)
|
|
}
|
|
|
|
/// Commit a transaction: check for conflicts, then apply buffered writes
|
|
/// to the underlying storage adapter.
|
|
pub async fn commit_transaction(
|
|
&self,
|
|
txn_id: &str,
|
|
storage: &dyn StorageAdapter,
|
|
) -> TransactionResult<()> {
|
|
// Remove the transaction so we own it exclusively.
|
|
let mut state = self
|
|
.transactions
|
|
.remove(txn_id)
|
|
.map(|(_, s)| s)
|
|
.ok_or_else(|| TransactionError::NotFound(txn_id.to_string()))?;
|
|
|
|
if state.status != TransactionStatus::Active {
|
|
return Err(TransactionError::InvalidState(format!(
|
|
"transaction {} is {:?}, cannot commit",
|
|
txn_id, state.status
|
|
)));
|
|
}
|
|
|
|
// Conflict detection: check if any documents in the read set have
|
|
// been modified since the snapshot was taken.
|
|
// (Simplified: we skip real snapshot timestamps for now.)
|
|
|
|
// Apply buffered writes to storage.
|
|
for (ns, writes) in &state.write_set {
|
|
let parts: Vec<&str> = ns.splitn(2, '.').collect();
|
|
if parts.len() != 2 {
|
|
warn!(namespace = %ns, "invalid namespace format, skipping");
|
|
continue;
|
|
}
|
|
let (db, coll) = (parts[0], parts[1]);
|
|
|
|
for (doc_id, entry) in writes {
|
|
match entry.op {
|
|
WriteOp::Insert => {
|
|
if let Some(ref doc) = entry.doc {
|
|
let _ = storage.insert_one(db, coll, doc.clone()).await;
|
|
}
|
|
}
|
|
WriteOp::Update => {
|
|
if let Some(ref doc) = entry.doc {
|
|
let _ = storage.update_by_id(db, coll, doc_id, doc.clone()).await;
|
|
}
|
|
}
|
|
WriteOp::Delete => {
|
|
let _ = storage.delete_by_id(db, coll, doc_id).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
state.status = TransactionStatus::Committed;
|
|
debug!(txn_id = %txn_id, "transaction committed");
|
|
Ok(())
|
|
}
|
|
|
|
/// Abort a transaction, discarding all buffered writes.
|
|
pub fn abort_transaction(&self, txn_id: &str) -> TransactionResult<()> {
|
|
let mut state = self
|
|
.transactions
|
|
.get_mut(txn_id)
|
|
.ok_or_else(|| TransactionError::NotFound(txn_id.to_string()))?;
|
|
|
|
if state.status != TransactionStatus::Active {
|
|
return Err(TransactionError::InvalidState(format!(
|
|
"transaction {} is {:?}, cannot abort",
|
|
txn_id, state.status
|
|
)));
|
|
}
|
|
|
|
state.status = TransactionStatus::Aborted;
|
|
debug!(txn_id = %txn_id, "transaction aborted");
|
|
|
|
// Drop the mutable ref before removing.
|
|
drop(state);
|
|
self.transactions.remove(txn_id);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Check whether a transaction is currently active.
|
|
pub fn is_active(&self, txn_id: &str) -> bool {
|
|
self.transactions
|
|
.get(txn_id)
|
|
.map(|s| s.status == TransactionStatus::Active)
|
|
.unwrap_or(false)
|
|
}
|
|
|
|
/// Record a document read within a transaction (for conflict detection).
|
|
pub fn record_read(&self, txn_id: &str, ns: &str, doc_id: &str) {
|
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
|
state
|
|
.read_set
|
|
.entry(ns.to_string())
|
|
.or_default()
|
|
.insert(doc_id.to_string());
|
|
}
|
|
}
|
|
|
|
/// Record a document write within a transaction (buffered until commit).
|
|
pub fn record_write(
|
|
&self,
|
|
txn_id: &str,
|
|
ns: &str,
|
|
doc_id: &str,
|
|
op: WriteOp,
|
|
doc: Option<Document>,
|
|
original: Option<Document>,
|
|
) {
|
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
|
let entry = WriteEntry {
|
|
op,
|
|
doc,
|
|
original_doc: original,
|
|
};
|
|
state
|
|
.write_set
|
|
.entry(ns.to_string())
|
|
.or_default()
|
|
.insert(doc_id.to_string(), entry);
|
|
}
|
|
}
|
|
|
|
/// Get a snapshot of documents for a namespace within a transaction,
|
|
/// applying the write overlay (inserts, updates, deletes) on top.
|
|
pub fn get_snapshot(&self, txn_id: &str, ns: &str) -> Option<Vec<Document>> {
|
|
let state = self.transactions.get(txn_id)?;
|
|
|
|
// Start with the base snapshot.
|
|
let mut docs: Vec<Document> = state
|
|
.snapshots
|
|
.get(ns)
|
|
.cloned()
|
|
.unwrap_or_default();
|
|
|
|
// Apply write overlay.
|
|
if let Some(writes) = state.write_set.get(ns) {
|
|
// Collect ids to delete.
|
|
let delete_ids: HashSet<&String> = writes
|
|
.iter()
|
|
.filter(|(_, e)| e.op == WriteOp::Delete)
|
|
.map(|(id, _)| id)
|
|
.collect();
|
|
|
|
// Remove deleted docs.
|
|
docs.retain(|d| {
|
|
if let Some(id) = d.get_object_id("_id").ok().map(|oid| oid.to_hex()) {
|
|
!delete_ids.contains(&id)
|
|
} else {
|
|
true
|
|
}
|
|
});
|
|
|
|
// Apply updates.
|
|
for (doc_id, entry) in writes {
|
|
if entry.op == WriteOp::Update {
|
|
if let Some(ref new_doc) = entry.doc {
|
|
// Replace existing doc with updated version.
|
|
let hex_id = doc_id.clone();
|
|
if let Some(pos) = docs.iter().position(|d| {
|
|
d.get_object_id("_id")
|
|
.ok()
|
|
.map(|oid| oid.to_hex()) == Some(hex_id.clone())
|
|
}) {
|
|
docs[pos] = new_doc.clone();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Apply inserts.
|
|
for (_doc_id, entry) in writes {
|
|
if entry.op == WriteOp::Insert {
|
|
if let Some(ref doc) = entry.doc {
|
|
docs.push(doc.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Some(docs)
|
|
}
|
|
|
|
/// Store a base snapshot for a namespace within a transaction.
|
|
pub fn set_snapshot(&self, txn_id: &str, ns: &str, docs: Vec<Document>) {
|
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
|
state.snapshots.insert(ns.to_string(), docs);
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for TransactionEngine {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|