push.rocks/smartdb
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(transactions): add single-node transaction support with session-aware reads, commits, aborts, and transaction metrics

Browse Source
This commit is contained in:
Jürgen Kunz
2026-04-29 22:14:46 +00:00
parent e79fe339aa
commit b72e8ed5e7
19 changed files with 913 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/crates/rustdb-commands/src/error.rs
+17 -1
View File
@@ -18,6 +18,12 @@ pub enum CommandError {
#[error("transaction error: {0}")]
TransactionError(String),
#[error("no such transaction: {0}")]
NoSuchTransaction(String),
#[error("write conflict: {0}")]
WriteConflict(String),
#[error("namespace not found: {0}")]
NamespaceNotFound(String),
@@ -52,6 +58,8 @@ impl CommandError {
CommandError::StorageError(_) => (1, "InternalError"),
CommandError::IndexError(_) => (27, "IndexNotFound"),
CommandError::TransactionError(_) => (112, "WriteConflict"),
CommandError::NoSuchTransaction(_) => (251, "NoSuchTransaction"),
CommandError::WriteConflict(_) => (112, "WriteConflict"),
CommandError::NamespaceNotFound(_) => (26, "NamespaceNotFound"),
CommandError::NamespaceExists(_) => (48, "NamespaceExists"),
CommandError::DuplicateKey(_) => (11000, "DuplicateKey"),
@@ -79,7 +87,15 @@ impl From<rustdb_storage::StorageError> for CommandError {
impl From<rustdb_txn::TransactionError> for CommandError {
fn from(e: rustdb_txn::TransactionError) -> Self {
CommandError::TransactionError(e.to_string())
match e {
rustdb_txn::TransactionError::NotFound(message) => {
CommandError::NoSuchTransaction(message)
}
rustdb_txn::TransactionError::WriteConflict(message) => {
CommandError::WriteConflict(message)
}
other => CommandError::TransactionError(other.to_string()),
}
}
}
rust/crates/rustdb-commands/src/handlers/admin_handler.rs
+74 -18
View File
@@ -2,8 +2,9 @@ use bson::{doc, Bson, Document};
use rustdb_index::IndexEngine;
use tracing::debug;
use crate::context::{CommandContext, CursorState};
use crate::context::{CommandContext, ConnectionState, CursorState};
use crate::error::{CommandError, CommandResult};
use crate::transactions;
/// Handle various admin / diagnostic / session / auth commands.
pub async fn handle(
@@ -11,6 +12,7 @@ pub async fn handle(
db: &str,
ctx: &CommandContext,
command_name: &str,
connection: &ConnectionState,
) -> CommandResult<Document> {
match command_name {
"ping" => Ok(doc! { "ok": 1.0 }),
@@ -24,13 +26,7 @@ pub async fn handle(
"ok": 1.0,
}),
"serverStatus" => Ok(doc! {
"host": "localhost",
"version": "7.0.0",
"process": "rustdb",
"uptime": ctx.start_time.elapsed().as_secs() as i64,
"ok": 1.0,
}),
"serverStatus" => handle_server_status(ctx),
"hostInfo" => Ok(doc! {
"system": {
@@ -90,13 +86,7 @@ pub async fn handle(
"codeName": "CommandNotFound",
}),
"connectionStatus" => Ok(doc! {
"authInfo": {
"authenticatedUsers": [],
"authenticatedUserRoles": [],
},
"ok": 1.0,
}),
"connectionStatus" => Ok(handle_connection_status(connection)),
"createUser" => handle_create_user(cmd, db, ctx).await,
@@ -156,9 +146,9 @@ pub async fn handle(
Ok(doc! { "ok": 1.0 })
}
"commitTransaction" | "abortTransaction" => Err(CommandError::IllegalOperation(
"Transaction numbers are only allowed on a replica set member or mongos".into(),
)),
"commitTransaction" => transactions::commit_transaction_command(cmd, ctx).await,
"abortTransaction" => transactions::abort_transaction_command(cmd, ctx),
// Auth stubs - accept silently.
"saslStart" => Ok(doc! {
@@ -195,6 +185,72 @@ pub async fn handle(
}
}
fn handle_server_status(ctx: &CommandContext) -> CommandResult<Document> {
let oplog_stats = ctx.oplog.stats();
Ok(doc! {
"host": "localhost",
"version": "7.0.0",
"process": "rustdb",
"uptime": ctx.start_time.elapsed().as_secs() as i64,
"connections": {
"current": 0_i32,
"available": i32::MAX,
},
"logicalSessionRecordCache": {
"activeSessionsCount": ctx.sessions.len() as i64,
},
"transactions": {
"currentActive": ctx.transactions.len() as i64,
},
"oplog": {
"currentSeq": oplog_stats.current_seq as i64,
"totalEntries": oplog_stats.total_entries as i64,
"oldestSeq": oplog_stats.oldest_seq as i64,
"entriesByOp": {
"insert": oplog_stats.inserts as i64,
"update": oplog_stats.updates as i64,
"delete": oplog_stats.deletes as i64,
},
},
"security": {
"authentication": ctx.auth.enabled(),
"users": ctx.auth.user_count() as i64,
},
"ok": 1.0,
})
}
fn handle_connection_status(connection: &ConnectionState) -> Document {
let authenticated_users: Vec<Bson> = connection
.authenticated_users
.iter()
.map(|user| {
Bson::Document(doc! {
"user": user.username.clone(),
"db": user.database.clone(),
})
})
.collect();
let authenticated_roles: Vec<Bson> = connection
.authenticated_users
.iter()
.flat_map(|user| {
user.roles
.iter()
.map(|role| Bson::Document(role_to_document(&user.database, role)))
})
.collect();
doc! {
"authInfo": {
"authenticatedUsers": authenticated_users,
"authenticatedUserRoles": authenticated_roles,
},
"ok": 1.0,
}
}
async fn handle_create_user(
cmd: &Document,
db: &str,
rust/crates/rustdb-commands/src/handlers/delete_handler.rs
+20 -1
View File
@@ -7,6 +7,7 @@ use tracing::debug;
use crate::context::CommandContext;
use crate::error::{CommandError, CommandResult};
use crate::transactions;
/// Handle the `delete` command.
pub async fn handle(
@@ -36,6 +37,7 @@ pub async fn handle(
);
let ns_key = format!("{}.{}", db, coll);
let txn_id = transactions::active_transaction_id(ctx, cmd);
let mut total_deleted: i32 = 0;
let mut write_errors: Vec<Document> = Vec::new();
@@ -69,7 +71,7 @@ pub async fn handle(
_ => 0, // default: delete all matches
};
match delete_matching(db, coll, &ns_key, &filter, limit, ctx).await {
match delete_matching(db, coll, &ns_key, &filter, limit, ctx, txn_id.as_deref()).await {
Ok(count) => {
total_deleted += count;
}
@@ -114,7 +116,24 @@ async fn delete_matching(
filter: &Document,
limit: i32,
ctx: &CommandContext,
txn_id: Option<&str>,
) -> Result<i32, CommandError> {
if let Some(txn_id) = txn_id {
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
let matched = QueryMatcher::filter(&docs, filter);
let to_delete: &[Document] = if limit == 1 && !matched.is_empty() {
&matched[..1]
} else {
&matched
};
for doc in to_delete {
transactions::record_delete(ctx, txn_id, db, coll, doc.clone()).await?;
}
return Ok(to_delete.len() as i32);
}
// Check if the collection exists; if not, nothing to delete.
match ctx.storage.collection_exists(db, coll).await {
Ok(false) => return Ok(0),
rust/crates/rustdb-commands/src/handlers/find_handler.rs
+44 -5
View File
@@ -7,6 +7,7 @@ use rustdb_query::{QueryMatcher, sort_documents, apply_projection, distinct_valu
use crate::context::{CommandContext, CursorState};
use crate::error::{CommandError, CommandResult};
use crate::transactions;
/// Atomic counter for generating unique cursor IDs.
static CURSOR_ID_COUNTER: AtomicI64 = AtomicI64::new(1);
@@ -80,9 +81,14 @@ pub async fn handle(
let limit = get_i64(cmd, "limit").unwrap_or(0).max(0) as usize;
let batch_size = get_i32(cmd, "batchSize").unwrap_or(101).max(0) as usize;
let single_batch = get_bool(cmd, "singleBatch").unwrap_or(false);
let txn_id = transactions::active_transaction_id(ctx, cmd);
// If the collection does not exist, return an empty cursor.
let exists = ctx.storage.collection_exists(db, coll).await?;
let exists = if txn_id.is_some() {
true
} else {
ctx.storage.collection_exists(db, coll).await?
};
if !exists {
return Ok(doc! {
"cursor": {
@@ -96,7 +102,9 @@ pub async fn handle(
// Try index-accelerated lookup.
let index_key = format!("{}.{}", db, coll);
let docs = if let Some(idx_ref) = ctx.indexes.get(&index_key) {
let docs = if let Some(ref txn_id) = txn_id {
transactions::load_transaction_docs(ctx, txn_id, db, coll).await?
} else if let Some(idx_ref) = ctx.indexes.get(&index_key) {
if let Some(candidate_ids) = idx_ref.find_candidate_ids(&filter) {
debug!(
ns = %ns,
@@ -298,9 +306,14 @@ pub async fn handle_count(
ctx: &CommandContext,
) -> CommandResult<Document> {
let coll = get_str(cmd, "count").unwrap_or("unknown");
let txn_id = transactions::active_transaction_id(ctx, cmd);
// Check collection existence.
let exists = ctx.storage.collection_exists(db, coll).await?;
let exists = if txn_id.is_some() {
true
} else {
ctx.storage.collection_exists(db, coll).await?
};
if !exists {
return Ok(doc! { "n": 0_i64, "ok": 1.0 });
}
@@ -309,6 +322,23 @@ pub async fn handle_count(
let skip = get_i64(cmd, "skip").unwrap_or(0).max(0) as usize;
let limit = get_i64(cmd, "limit").unwrap_or(0).max(0) as usize;
if let Some(ref txn_id) = txn_id {
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
let filtered = if query.is_empty() {
docs
} else {
QueryMatcher::filter(&docs, &query)
};
let mut n = filtered.len().saturating_sub(skip);
if limit > 0 {
n = n.min(limit);
}
return Ok(doc! {
"n": n as i64,
"ok": 1.0,
});
}
let count: u64 = if query.is_empty() && skip == 0 && limit == 0 {
// Fast path: use storage-level count.
ctx.storage.count(db, coll).await?
@@ -352,15 +382,24 @@ pub async fn handle_distinct(
let key = get_str(cmd, "key").ok_or_else(|| {
CommandError::InvalidArgument("distinct requires a 'key' field".into())
})?;
let txn_id = transactions::active_transaction_id(ctx, cmd);
// Check collection existence.
let exists = ctx.storage.collection_exists(db, coll).await?;
let exists = if txn_id.is_some() {
true
} else {
ctx.storage.collection_exists(db, coll).await?
};
if !exists {
return Ok(doc! { "values": [], "ok": 1.0 });
}
let query = get_document(cmd, "query").cloned();
let docs = ctx.storage.find_all(db, coll).await?;
let docs = if let Some(txn_id) = txn_id {
transactions::load_transaction_docs(ctx, &txn_id, db, coll).await?
} else {
ctx.storage.find_all(db, coll).await?
};
let values = distinct_values(&docs, key, query.as_ref());
Ok(doc! {
rust/crates/rustdb-commands/src/handlers/insert_handler.rs
+26 -2
View File
@@ -6,6 +6,7 @@ use tracing::debug;
use crate::context::CommandContext;
use crate::error::{CommandError, CommandResult};
use crate::transactions;
/// Handle the `insert` command.
pub async fn handle(
@@ -48,8 +49,13 @@ pub async fn handle(
"insert command"
);
// Auto-create database and collection if they don't exist.
ensure_collection_exists(db, coll, ctx).await?;
let txn_id = transactions::active_transaction_id(ctx, cmd);
// Auto-create database and collection if they don't exist. Transactional
// writes defer collection creation until commit so abort remains clean.
if txn_id.is_none() {
ensure_collection_exists(db, coll, ctx).await?;
}
let ns_key = format!("{}.{}", db, coll);
let mut inserted_count: i32 = 0;
@@ -84,6 +90,24 @@ pub async fn handle(
}
}
if let Some(ref txn_id) = txn_id {
match transactions::record_insert(ctx, txn_id, db, coll, doc.clone()).await {
Ok(_) => inserted_count += 1,
Err(e) => {
write_errors.push(doc! {
"index": idx as i32,
"code": 11000_i32,
"codeName": "DuplicateKey",
"errmsg": e.to_string(),
});
if ordered {
break;
}
}
}
continue;
}
// Attempt storage insert.
match ctx.storage.insert_one(db, coll, doc.clone()).await {
Ok(id_str) => {
rust/crates/rustdb-commands/src/handlers/update_handler.rs
+145 -6
View File
@@ -7,6 +7,7 @@ use tracing::debug;
use crate::context::CommandContext;
use crate::error::{CommandError, CommandResult};
use crate::transactions;
/// Handle `update` and `findAndModify` commands.
pub async fn handle(
@@ -47,8 +48,12 @@ async fn handle_update(
debug!(db = db, collection = coll, count = updates.len(), "update command");
// Auto-create database and collection if needed.
ensure_collection_exists(db, coll, ctx).await?;
let txn_id = transactions::active_transaction_id(ctx, cmd);
// Transactional writes defer namespace creation until commit.
if txn_id.is_none() {
ensure_collection_exists(db, coll, ctx).await?;
}
let ns_key = format!("{}.{}", db, coll);
@@ -136,7 +141,7 @@ async fn handle_update(
});
// Load all documents and filter.
let all_docs = load_filtered_docs(db, coll, &filter, &ns_key, ctx).await?;
let all_docs = load_filtered_docs(db, coll, &filter, &ns_key, ctx, txn_id.as_deref()).await?;
if all_docs.is_empty() && upsert {
// Upsert: create a new document.
@@ -166,6 +171,30 @@ async fn handle_update(
}
}
if let Some(ref txn_id) = txn_id {
match transactions::record_insert(ctx, txn_id, db, coll, updated.clone()).await {
Ok(_) => {
total_n += 1;
upserted_list.push(doc! {
"index": idx as i32,
"_id": new_id,
});
}
Err(e) => {
write_errors.push(doc! {
"index": idx as i32,
"code": 1_i32,
"codeName": "InternalError",
"errmsg": e.to_string(),
});
if ordered {
break;
}
}
}
continue;
}
// Insert the new document.
match ctx.storage.insert_one(db, coll, updated.clone()).await {
Ok(id_str) => {
@@ -258,6 +287,38 @@ async fn handle_update(
}
let id_str = extract_id_string(matched_doc);
if let Some(ref txn_id) = txn_id {
match transactions::record_update(
ctx,
txn_id,
db,
coll,
matched_doc.clone(),
updated_doc.clone(),
)
.await
{
Ok(_) => {
total_n += 1;
if matched_doc != &updated_doc {
total_n_modified += 1;
}
}
Err(e) => {
write_errors.push(doc! {
"index": idx as i32,
"code": 1_i32,
"codeName": "InternalError",
"errmsg": e.to_string(),
});
if ordered {
break;
}
}
}
continue;
}
match ctx
.storage
.update_by_id(db, coll, &id_str, updated_doc.clone())
@@ -407,8 +468,12 @@ async fn handle_find_and_modify(
.collect()
});
// Auto-create database and collection.
ensure_collection_exists(db, coll, ctx).await?;
let txn_id = transactions::active_transaction_id(ctx, cmd);
// Transactional writes defer namespace creation until commit.
if txn_id.is_none() {
ensure_collection_exists(db, coll, ctx).await?;
}
let ns_key = format!("{}.{}", db, coll);
@@ -416,7 +481,7 @@ async fn handle_find_and_modify(
drop(ctx.get_or_init_index_engine(db, coll).await);
// Load and filter documents.
let mut matched = load_filtered_docs(db, coll, &query, &ns_key, ctx).await?;
let mut matched = load_filtered_docs(db, coll, &query, &ns_key, ctx, txn_id.as_deref()).await?;
// Sort if specified.
if let Some(ref sort_spec) = sort {
@@ -430,6 +495,21 @@ async fn handle_find_and_modify(
// Remove operation.
if let Some(ref doc) = target {
let id_str = extract_id_string(doc);
if let Some(ref txn_id) = txn_id {
transactions::record_delete(ctx, txn_id, db, coll, doc.clone()).await?;
let value = apply_fields_projection(doc, &fields);
return Ok(doc! {
"value": value,
"lastErrorObject": {
"n": 1_i32,
"updatedExisting": false,
},
"ok": 1.0,
});
}
ctx.storage.delete_by_id(db, coll, &id_str).await?;
// Record in oplog.
@@ -503,6 +583,35 @@ async fn handle_find_and_modify(
}
let id_str = extract_id_string(&original_doc);
if let Some(ref txn_id) = txn_id {
transactions::record_update(
ctx,
txn_id,
db,
coll,
original_doc.clone(),
updated_doc.clone(),
)
.await?;
let return_doc = if return_new {
&updated_doc
} else {
&original_doc
};
let value = apply_fields_projection(return_doc, &fields);
return Ok(doc! {
"value": value,
"lastErrorObject": {
"n": 1_i32,
"updatedExisting": true,
},
"ok": 1.0,
});
}
ctx.storage
.update_by_id(db, coll, &id_str, updated_doc.clone())
.await?;
@@ -563,6 +672,26 @@ async fn handle_find_and_modify(
}
}
if let Some(ref txn_id) = txn_id {
transactions::record_insert(ctx, txn_id, db, coll, updated_doc.clone()).await?;
let value = if return_new {
apply_fields_projection(&updated_doc, &fields)
} else {
Bson::Null
};
return Ok(doc! {
"value": value,
"lastErrorObject": {
"n": 1_i32,
"updatedExisting": false,
"upserted": upserted_id,
},
"ok": 1.0,
});
}
let inserted_id_str = ctx.storage
.insert_one(db, coll, updated_doc.clone())
.await?;
@@ -622,7 +751,17 @@ async fn load_filtered_docs(
filter: &Document,
ns_key: &str,
ctx: &CommandContext,
txn_id: Option<&str>,
) -> CommandResult<Vec<Document>> {
if let Some(txn_id) = txn_id {
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
return if filter.is_empty() {
Ok(docs)
} else {
Ok(QueryMatcher::filter(&docs, filter))
};
}
// Try to use index to narrow candidates.
let candidate_ids: Option<HashSet<String>> = ctx
.indexes
rust/crates/rustdb-commands/src/lib.rs
+1
View File
@@ -1,6 +1,7 @@
mod context;
pub mod error;
pub mod handlers;
pub mod transactions;
mod router;
pub use context::{CommandContext, ConnectionState, CursorState};
rust/crates/rustdb-commands/src/router.rs
+8 -13
View File
@@ -8,7 +8,7 @@ use rustdb_auth::AuthAction;
use crate::context::{CommandContext, ConnectionState};
use crate::error::CommandError;
use crate::handlers;
use crate::{handlers, transactions};
/// Routes parsed wire protocol commands to the appropriate handler.
pub struct CommandRouter {
@@ -55,11 +55,12 @@ impl CommandRouter {
}
}
if transaction_command_unsupported(command_name, &cmd.command) {
return CommandError::IllegalOperation(
"Transaction numbers are only allowed on a replica set member or mongos".into(),
)
.to_error_doc();
if let Err(e) = transactions::prepare_transaction_for_command(
&self.ctx,
&cmd.command,
command_name,
) {
return e.to_error_doc();
}
// Extract session id if present, and touch the session.
@@ -136,7 +137,7 @@ impl CommandRouter {
| "grantRolesToUser" | "revokeRolesFromUser"
| "currentOp" | "killOp" | "top" | "profile"
| "compact" | "reIndex" | "fsync" | "connPoolSync" => {
handlers::admin_handler::handle(&cmd.command, db, &self.ctx, command_name).await
handlers::admin_handler::handle(&cmd.command, db, &self.ctx, command_name, connection).await
}
// -- unknown command --
@@ -207,9 +208,3 @@ fn aggregate_writes(command: &Document) -> bool {
_ => None,
}).unwrap_or(false)
}
fn transaction_command_unsupported(command_name: &str, command: &Document) -> bool {
matches!(command_name, "commitTransaction" | "abortTransaction")
|| matches!(command.get("startTransaction"), Some(Bson::Boolean(true)))
|| matches!(command.get("autocommit"), Some(Bson::Boolean(false)))
}
rust/crates/rustdb-commands/src/transactions.rs
+367
View File
@@ -0,0 +1,367 @@
use bson::{doc, Bson, Document};
use rustdb_storage::OpType;
use rustdb_txn::{TransactionState, WriteEntry, WriteOp};
use crate::context::CommandContext;
use crate::error::{CommandError, CommandResult};
pub fn command_starts_transaction(cmd: &Document) -> bool {
matches!(cmd.get("startTransaction"), Some(Bson::Boolean(true)))
}
pub fn command_uses_transaction(cmd: &Document) -> bool {
command_starts_transaction(cmd) || matches!(cmd.get("autocommit"), Some(Bson::Boolean(false)))
}
pub fn active_transaction_id(ctx: &CommandContext, cmd: &Document) -> Option<String> {
if !command_uses_transaction(cmd) {
return None;
}
let session_id = cmd
.get("lsid")
.and_then(rustdb_txn::SessionEngine::extract_session_id)?;
ctx.sessions.get_transaction_id(&session_id)
}
pub fn prepare_transaction_for_command(
ctx: &CommandContext,
cmd: &Document,
command_name: &str,
) -> CommandResult<()> {
if matches!(command_name, "commitTransaction" | "abortTransaction") {
return Ok(());
}
let starts_transaction = command_starts_transaction(cmd);
let uses_transaction = command_uses_transaction(cmd);
if !uses_transaction {
return Ok(());
}
let session_id = session_id_from_command(cmd)?;
require_txn_number(cmd)?;
ctx.sessions.get_or_create_session(&session_id);
if starts_transaction {
let txn_id = ctx.transactions.start_transaction(&session_id)?;
ctx.sessions.start_transaction(&session_id, &txn_id)?;
return Ok(());
}
if ctx.sessions.get_transaction_id(&session_id).is_none() {
return Err(CommandError::NoSuchTransaction(format!(
"session {session_id} has no active transaction"
)));
}
Ok(())
}
pub async fn load_transaction_docs(
ctx: &CommandContext,
txn_id: &str,
db: &str,
coll: &str,
) -> CommandResult<Vec<Document>> {
let ns = namespace(db, coll);
if !ctx.transactions.has_snapshot(txn_id, &ns) {
let docs = match ctx.storage.collection_exists(db, coll).await {
Ok(true) => ctx.storage.find_all(db, coll).await?,
Ok(false) => Vec::new(),
Err(_) => Vec::new(),
};
ctx.transactions.set_snapshot(txn_id, &ns, docs);
}
ctx.transactions
.get_snapshot(txn_id, &ns)
.ok_or_else(|| CommandError::NoSuchTransaction(txn_id.to_string()))
}
pub async fn record_insert(
ctx: &CommandContext,
txn_id: &str,
db: &str,
coll: &str,
doc: Document,
) -> CommandResult<String> {
let id = document_id_string(&doc)?;
let docs = load_transaction_docs(ctx, txn_id, db, coll).await?;
if docs.iter().any(|existing| document_id_string(existing).ok().as_deref() == Some(id.as_str())) {
return Err(CommandError::DuplicateKey(format!(
"duplicate _id '{}' in transaction",
id
)));
}
ctx.transactions.record_write(
txn_id,
&namespace(db, coll),
&id,
WriteOp::Insert,
Some(doc),
None,
);
Ok(id)
}
pub async fn record_update(
ctx: &CommandContext,
txn_id: &str,
db: &str,
coll: &str,
original: Document,
updated: Document,
) -> CommandResult<String> {
let id = document_id_string(&original)?;
ctx.transactions.record_write(
txn_id,
&namespace(db, coll),
&id,
WriteOp::Update,
Some(updated),
Some(original),
);
Ok(id)
}
pub async fn record_delete(
ctx: &CommandContext,
txn_id: &str,
db: &str,
coll: &str,
original: Document,
) -> CommandResult<String> {
let id = document_id_string(&original)?;
ctx.transactions.record_write(
txn_id,
&namespace(db, coll),
&id,
WriteOp::Delete,
None,
Some(original),
);
Ok(id)
}
pub async fn commit_transaction_command(
cmd: &Document,
ctx: &CommandContext,
) -> CommandResult<Document> {
let session_id = session_id_from_command(cmd)?;
let txn_id = ctx
.sessions
.get_transaction_id(&session_id)
.ok_or_else(|| CommandError::NoSuchTransaction(format!(
"session {session_id} has no active transaction"
)))?;
let state = ctx.transactions.take_transaction(&txn_id)?;
preflight_transaction(&state, ctx).await?;
apply_transaction(state, ctx).await?;
ctx.sessions.end_transaction(&session_id);
Ok(doc! { "ok": 1.0 })
}
pub fn abort_transaction_command(cmd: &Document, ctx: &CommandContext) -> CommandResult<Document> {
let session_id = session_id_from_command(cmd)?;
let txn_id = ctx
.sessions
.get_transaction_id(&session_id)
.ok_or_else(|| CommandError::NoSuchTransaction(format!(
"session {session_id} has no active transaction"
)))?;
ctx.transactions.abort_transaction(&txn_id)?;
ctx.sessions.end_transaction(&session_id);
Ok(doc! { "ok": 1.0 })
}
pub fn document_id_string(doc: &Document) -> CommandResult<String> {
match doc.get("_id") {
Some(Bson::ObjectId(oid)) => Ok(oid.to_hex()),
Some(Bson::String(s)) => Ok(s.clone()),
Some(other) => Ok(format!("{}", other)),
None => Err(CommandError::InvalidArgument("document missing _id field".into())),
}
}
fn session_id_from_command(cmd: &Document) -> CommandResult<String> {
cmd.get("lsid")
.and_then(rustdb_txn::SessionEngine::extract_session_id)
.ok_or_else(|| CommandError::InvalidArgument("transaction command requires lsid".into()))
}
fn require_txn_number(cmd: &Document) -> CommandResult<()> {
match cmd.get("txnNumber") {
Some(Bson::Int64(_)) | Some(Bson::Int32(_)) => Ok(()),
_ => Err(CommandError::InvalidArgument(
"transaction command requires txnNumber".into(),
)),
}
}
fn namespace(db: &str, coll: &str) -> String {
format!("{db}.{coll}")
}
async fn preflight_transaction(state: &TransactionState, ctx: &CommandContext) -> CommandResult<()> {
for (ns, writes) in &state.write_set {
let (db, coll) = split_namespace(ns)?;
drop(ctx.get_or_init_index_engine(db, coll).await);
for (doc_id, entry) in writes {
let current = current_doc(ctx, db, coll, doc_id).await?;
match entry.op {
WriteOp::Insert => {
if current.is_some() {
return Err(CommandError::DuplicateKey(format!(
"duplicate _id '{}' on transaction commit",
doc_id
)));
}
if let Some(ref doc) = entry.doc {
if let Some(engine) = ctx.indexes.get(ns) {
engine.check_unique_constraints(doc)?;
}
}
}
WriteOp::Update => {
assert_unchanged(doc_id, current.as_ref(), entry.original_doc.as_ref())?;
if let (Some(current_doc), Some(updated_doc)) = (current.as_ref(), entry.doc.as_ref()) {
if let Some(engine) = ctx.indexes.get(ns) {
engine.check_unique_constraints_for_update(current_doc, updated_doc)?;
}
}
}
WriteOp::Delete => {
assert_unchanged(doc_id, current.as_ref(), entry.original_doc.as_ref())?;
}
}
}
}
Ok(())
}
async fn apply_transaction(state: TransactionState, ctx: &CommandContext) -> CommandResult<()> {
let mut namespaces: Vec<_> = state.write_set.into_iter().collect();
namespaces.sort_by(|a, b| a.0.cmp(&b.0));
for (ns, writes) in namespaces {
let (db, coll) = split_namespace(&ns)?;
ensure_collection_exists(db, coll, ctx).await?;
drop(ctx.get_or_init_index_engine(db, coll).await);
let mut writes: Vec<(String, WriteEntry)> = writes.into_iter().collect();
writes.sort_by(|a, b| a.0.cmp(&b.0));
for (doc_id, entry) in writes {
match entry.op {
WriteOp::Insert => {
let Some(doc) = entry.doc else { continue; };
let inserted_id = ctx.storage.insert_one(db, coll, doc.clone()).await?;
ctx.oplog.append(OpType::Insert, db, coll, &inserted_id, Some(doc.clone()), None);
if let Some(mut engine) = ctx.indexes.get_mut(&ns) {
engine.on_insert(&doc)?;
}
}
WriteOp::Update => {
let Some(doc) = entry.doc else { continue; };
ctx.storage.update_by_id(db, coll, &doc_id, doc.clone()).await?;
ctx.oplog.append(
OpType::Update,
db,
coll,
&doc_id,
Some(doc.clone()),
entry.original_doc.clone(),
);
if let (Some(mut engine), Some(ref original)) =
(ctx.indexes.get_mut(&ns), entry.original_doc.as_ref())
{
engine.on_update(original, &doc)?;
}
}
WriteOp::Delete => {
ctx.storage.delete_by_id(db, coll, &doc_id).await?;
ctx.oplog.append(
OpType::Delete,
db,
coll,
&doc_id,
None,
entry.original_doc.clone(),
);
if let (Some(mut engine), Some(ref original)) =
(ctx.indexes.get_mut(&ns), entry.original_doc.as_ref())
{
engine.on_delete(original);
}
}
}
}
}
Ok(())
}
async fn current_doc(
ctx: &CommandContext,
db: &str,
coll: &str,
doc_id: &str,
) -> CommandResult<Option<Document>> {
match ctx.storage.collection_exists(db, coll).await {
Ok(true) => Ok(ctx.storage.find_by_id(db, coll, doc_id).await?),
Ok(false) => Ok(None),
Err(_) => Ok(None),
}
}
fn assert_unchanged(
doc_id: &str,
current: Option<&Document>,
original: Option<&Document>,
) -> CommandResult<()> {
if current == original {
return Ok(());
}
Err(CommandError::WriteConflict(format!(
"document '{}' changed during transaction",
doc_id
)))
}
async fn ensure_collection_exists(
db: &str,
coll: &str,
ctx: &CommandContext,
) -> CommandResult<()> {
if let Err(e) = ctx.storage.create_database(db).await {
let msg = e.to_string();
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
return Err(CommandError::StorageError(msg));
}
}
match ctx.storage.collection_exists(db, coll).await {
Ok(true) => Ok(()),
Ok(false) | Err(_) => {
if let Err(e) = ctx.storage.create_collection(db, coll).await {
let msg = e.to_string();
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
return Err(CommandError::StorageError(msg));
}
}
Ok(())
}
}
}
fn split_namespace(ns: &str) -> CommandResult<(&str, &str)> {
ns.split_once('.')
.ok_or_else(|| CommandError::InvalidArgument(format!("invalid namespace '{ns}'")))
}