use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard};
use crate::error::KgError;
use crate::graph::dir_graph::DirGraph;
impl DirGraph {
pub(crate) fn fork_transaction(&self) -> Self {
let mut child = self.clone();
child.graph.adopt_transaction_lineage(&self.graph);
child
}
}
pub struct Session {
graph: Mutex<Arc<DirGraph>>,
}
pub struct SessionWriteGuard<'a> {
guard: MutexGuard<'a, Arc<DirGraph>>,
}
impl Deref for SessionWriteGuard<'_> {
type Target = DirGraph;
fn deref(&self) -> &Self::Target {
&self.guard
}
}
impl DerefMut for SessionWriteGuard<'_> {
fn deref_mut(&mut self) -> &mut Self::Target {
if Arc::get_mut(&mut self.guard).is_none() {
let child = self.guard.fork_transaction();
*self.guard = Arc::new(child);
}
Arc::get_mut(&mut self.guard).expect("Session write guard owns the active graph")
}
}
impl Session {
pub fn new(graph: DirGraph) -> Self {
Self {
graph: Mutex::new(Arc::new(graph)),
}
}
pub fn from_arc(graph: Arc<DirGraph>) -> Self {
Self {
graph: Mutex::new(graph),
}
}
pub fn snapshot(&self) -> Arc<DirGraph> {
Arc::clone(&self.graph.lock().unwrap_or_else(|p| p.into_inner()))
}
pub fn version(&self) -> u64 {
self.snapshot().version()
}
pub fn write(&self) -> SessionWriteGuard<'_> {
SessionWriteGuard {
guard: self.graph.lock().unwrap_or_else(|p| p.into_inner()),
}
}
pub fn transact<T, E>(
&self,
operation: impl FnOnce(&mut DirGraph) -> Result<T, E>,
) -> Result<T, E> {
let mut guard = self.graph.lock().unwrap_or_else(|p| p.into_inner());
let current_version = guard.version();
let mut working = guard.fork_transaction();
let value = operation(&mut working)?;
working.set_version(current_version + 1);
*guard = Arc::new(working);
Ok(value)
}
pub fn begin(&self) -> Transaction {
let snapshot = self.snapshot();
let base_version = snapshot.version();
Transaction {
snapshot: Some(snapshot),
working: None,
base_version,
read_only: false,
}
}
pub fn begin_read(&self) -> Transaction {
let mut tx = self.begin();
tx.read_only = true;
tx
}
pub fn commit(&self, tx: Transaction, check_occ: bool) -> CommitOutcome {
let (working_opt, base_version) = tx.take_working();
let Some(mut working) = working_opt else {
return CommitOutcome::NoWritesNoOp;
};
let mut guard = self.graph.lock().unwrap_or_else(|p| p.into_inner());
let current_version = guard.version();
if check_occ && current_version != base_version {
return CommitOutcome::ConflictDetected {
current_version,
base_version,
};
}
let new_version = current_version + 1;
working.set_version(new_version);
*guard = Arc::new(working);
CommitOutcome::Committed { new_version }
}
pub fn rollback(&self, _tx: Transaction) {
}
}
pub struct Transaction {
pub(super) snapshot: Option<Arc<DirGraph>>,
pub(super) working: Option<DirGraph>,
pub(super) base_version: u64,
pub(super) read_only: bool,
}
impl Transaction {
pub fn is_read_only(&self) -> bool {
self.read_only
}
pub fn base_version(&self) -> u64 {
self.base_version
}
pub fn has_writes(&self) -> bool {
self.working.is_some()
}
pub fn current(&self) -> Option<&DirGraph> {
self.working.as_ref().or(self.snapshot.as_deref())
}
pub fn working_mut(&mut self) -> Result<&mut DirGraph, KgError> {
if self.read_only {
return Err(KgError::Argument(
"read-only transaction does not support mutations \
(CREATE/SET/DELETE/REMOVE/MERGE) — open a read-write tx \
via Session::begin"
.to_string(),
));
}
if self.working.is_none() {
let snap = self.snapshot.take().ok_or_else(|| {
KgError::Argument("transaction already committed or rolled back".to_string())
})?;
let working = Arc::try_unwrap(snap).unwrap_or_else(|arc| arc.fork_transaction());
self.working = Some(working);
}
Ok(self
.working
.as_mut()
.expect("invariant: just materialized above"))
}
pub fn take_working(self) -> (Option<DirGraph>, u64) {
(self.working, self.base_version)
}
}
#[derive(Debug)]
pub enum CommitOutcome {
NoWritesNoOp,
Committed { new_version: u64 },
ConflictDetected {
current_version: u64,
base_version: u64,
},
}
#[cfg(test)]
mod tests {
use super::*;
fn empty_graph() -> DirGraph {
DirGraph::new()
}
#[test]
fn new_session_version_is_zero() {
let s = Session::new(empty_graph());
assert_eq!(s.version(), 0);
}
#[test]
fn snapshot_is_cheap_arc_clone() {
let s = Session::new(empty_graph());
let snap1 = s.snapshot();
let snap2 = s.snapshot();
assert!(Arc::ptr_eq(&snap1, &snap2));
}
#[test]
fn serialized_write_uses_unique_arc_in_place() {
use crate::graph::storage::mode::{new_dir_graph_in_mode, StorageMode};
for mode in [StorageMode::Memory, StorageMode::Mapped] {
let s = Session::new(new_dir_graph_in_mode(mode, None).unwrap());
let before = {
let guard = s.graph.lock().unwrap();
Arc::as_ptr(&guard)
};
{
let mut graph = s.write();
graph.bump_version();
}
let after = {
let guard = s.graph.lock().unwrap();
Arc::as_ptr(&guard)
};
assert_eq!(
before, after,
"unique Session write must not clone the graph"
);
assert_eq!(s.version(), 1);
}
}
#[test]
fn serialized_execute_skips_checkpoint_for_proven_single_node_create() {
use crate::graph::session::execute::{execute_mut, ExecuteOptions};
use crate::graph::storage::backend::{backend_clone_count, reset_backend_clone_count};
let params = std::collections::HashMap::new();
let mut opts = ExecuteOptions::eager(¶ms);
opts.deadline = Some(std::time::Instant::now() + std::time::Duration::from_secs(60));
let unique = Session::new(empty_graph());
reset_backend_clone_count();
execute_mut(&mut unique.write(), "CREATE (:N {id: 1})", &opts).unwrap();
assert_eq!(
backend_clone_count(),
0,
"proven single-node CREATE must not clone a uniquely-owned graph"
);
let shared = Session::new(empty_graph());
let _reader = shared.snapshot();
reset_backend_clone_count();
execute_mut(&mut shared.write(), "CREATE (:N {id: 1})", &opts).unwrap();
assert_eq!(
backend_clone_count(),
1,
"held reader needs only the working fork, not a second checkpoint"
);
let checkpointed = Session::new(empty_graph());
reset_backend_clone_count();
execute_mut(
&mut checkpointed.write(),
"CREATE (:N {id: 1}), (:N {id: 2})",
&opts,
)
.unwrap();
assert_eq!(
backend_clone_count(),
1,
"multi-element CREATE must retain its atomic rollback checkpoint"
);
reset_backend_clone_count();
execute_mut(
&mut checkpointed.write(),
"MATCH (n:N {id: 1}) DELETE n",
&opts,
)
.unwrap();
assert_eq!(
backend_clone_count(),
0,
"terminal preflighted DELETE must not clone the graph"
);
}
#[test]
fn serialized_write_forks_when_reader_snapshot_is_held() {
let s = Session::new(empty_graph());
let old = s.snapshot();
{
let mut graph = s.write();
graph.bump_version();
}
let current = s.snapshot();
assert!(!Arc::ptr_eq(&old, ¤t));
assert_eq!(old.version(), 0);
assert_eq!(current.version(), 1);
}
#[test]
fn serialized_transaction_swaps_once_or_discards_on_error() {
let s = Session::new(empty_graph());
let old = s.snapshot();
let value = s
.transact(|working| {
working.bump_version();
working.bump_version();
Ok::<_, &'static str>(42)
})
.unwrap();
assert_eq!(value, 42);
assert_eq!(s.version(), 1, "one transaction is one committed version");
assert_eq!(old.version(), 0);
let failed = s.transact(|working| {
working.bump_version();
Err::<(), _>("cancelled")
});
assert_eq!(failed, Err("cancelled"));
assert_eq!(
s.version(),
1,
"failed transaction must not reach the live Arc"
);
}
#[test]
fn begin_then_commit_no_writes_is_noop() {
let s = Session::new(empty_graph());
let tx = s.begin();
let outcome = s.commit(tx, true);
assert!(matches!(outcome, CommitOutcome::NoWritesNoOp));
assert_eq!(s.version(), 0);
}
#[test]
fn begin_then_rollback_is_noop() {
let s = Session::new(empty_graph());
let tx = s.begin();
s.rollback(tx);
assert_eq!(s.version(), 0);
}
#[test]
fn working_mut_materializes_only_on_first_call() {
let s = Session::new(empty_graph());
let mut tx = s.begin();
assert!(!tx.has_writes());
assert!(tx.current().is_some());
let _ = tx.working_mut().unwrap();
assert!(tx.has_writes());
assert!(tx.snapshot.is_none());
assert!(tx.working.is_some());
}
#[test]
fn current_routes_through_working_after_materialize() {
let s = Session::new(empty_graph());
let mut tx = s.begin();
let _ = tx.working_mut().unwrap();
let _: &DirGraph = tx.current().unwrap();
}
#[test]
fn commit_with_writes_bumps_version() {
let s = Session::new(empty_graph());
let mut tx = s.begin();
let _ = tx.working_mut().unwrap();
let outcome = s.commit(tx, true);
match outcome {
CommitOutcome::Committed { new_version } => assert_eq!(new_version, 1),
other => panic!("expected Committed, got {other:?}"),
}
assert_eq!(s.version(), 1);
}
#[test]
fn read_only_tx_rejects_working_mut() {
let s = Session::new(empty_graph());
let mut tx = s.begin_read();
assert!(tx.is_read_only());
match tx.working_mut() {
Err(KgError::Argument(msg)) => assert!(msg.contains("read-only")),
Err(other) => panic!("expected Argument, got different error: {other}"),
Ok(_) => panic!("expected read-only rejection but got Ok"),
}
}
#[test]
fn read_only_tx_commit_is_noop() {
let s = Session::new(empty_graph());
let tx = s.begin_read();
let outcome = s.commit(tx, true);
assert!(matches!(outcome, CommitOutcome::NoWritesNoOp));
assert_eq!(s.version(), 0);
}
#[test]
fn occ_conflict_detected_when_other_writer_commits() {
let s = Arc::new(Session::new(empty_graph()));
let mut tx_a = s.begin();
let _ = tx_a.working_mut().unwrap();
let mut tx_b = s.begin();
let _ = tx_b.working_mut().unwrap();
let outcome_b = s.commit(tx_b, true);
assert!(matches!(
outcome_b,
CommitOutcome::Committed { new_version: 1 }
));
let outcome_a = s.commit(tx_a, true);
match outcome_a {
CommitOutcome::ConflictDetected {
current_version,
base_version,
} => {
assert_eq!(current_version, 1);
assert_eq!(base_version, 0);
}
other => panic!("expected ConflictDetected, got {other:?}"),
}
assert_eq!(s.version(), 1);
}
#[test]
fn occ_disabled_means_last_writer_wins() {
let s = Arc::new(Session::new(empty_graph()));
let mut tx_a = s.begin();
let _ = tx_a.working_mut().unwrap();
let mut tx_b = s.begin();
let _ = tx_b.working_mut().unwrap();
let outcome_a = s.commit(tx_a, false);
let outcome_b = s.commit(tx_b, false);
assert!(matches!(outcome_a, CommitOutcome::Committed { .. }));
assert!(matches!(outcome_b, CommitOutcome::Committed { .. }));
assert_eq!(s.version(), 2);
}
#[test]
fn snapshot_after_commit_sees_new_graph() {
let s = Session::new(empty_graph());
let pre = s.snapshot();
assert_eq!(pre.version(), 0);
let mut tx = s.begin();
let _ = tx.working_mut().unwrap();
let _ = s.commit(tx, true);
let post = s.snapshot();
assert_eq!(post.version(), 1);
assert!(!Arc::ptr_eq(&pre, &post));
}
#[test]
fn double_commit_via_take_working_drops_state() {
let s = Session::new(empty_graph());
let tx = s.begin();
let _ = s.commit(tx, true);
}
#[test]
fn concurrent_writers_compose_with_occ_retry() {
const N: u64 = 8;
const PER: u64 = 200;
let session = Arc::new(Session::new(empty_graph()));
let handles: Vec<_> = (0..N)
.map(|_| {
let s = Arc::clone(&session);
std::thread::spawn(move || {
for _ in 0..PER {
loop {
let mut tx = s.begin();
tx.working_mut().expect("rw tx");
match s.commit(tx, true) {
CommitOutcome::Committed { .. } => break,
CommitOutcome::ConflictDetected { .. } => continue,
CommitOutcome::NoWritesNoOp => {
panic!("materialised tx must commit as a write")
}
}
}
}
})
})
.collect();
for h in handles {
h.join().expect("worker thread panicked");
}
assert_eq!(
session.version(),
N * PER,
"OCC-retried commits must compose with no lost or doubled updates"
);
}
#[test]
fn occ_detects_conflict_between_overlapping_txs() {
let s = Session::new(empty_graph());
let mut tx1 = s.begin();
tx1.working_mut().unwrap();
let mut tx2 = s.begin();
tx2.working_mut().unwrap();
assert!(matches!(
s.commit(tx1, true),
CommitOutcome::Committed { .. }
));
assert!(matches!(
s.commit(tx2, true),
CommitOutcome::ConflictDetected { .. }
));
assert_eq!(s.version(), 1, "only the winning commit bumped the version");
}
#[test]
fn concurrent_snapshots_consistent_under_commits() {
const WRITERS: u64 = 4;
const PER: u64 = 250;
const READERS: usize = 4;
let total = WRITERS * PER;
let session = Arc::new(Session::new(empty_graph()));
let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
let readers: Vec<_> = (0..READERS)
.map(|_| {
let s = Arc::clone(&session);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
let mut last = 0u64;
while !stop.load(std::sync::atomic::Ordering::Relaxed) {
let snap = s.snapshot();
let v = snap.version();
assert!(v <= total, "snapshot version {v} exceeds total {total}");
assert!(v >= last, "version went backwards: {v} < {last}");
last = v;
}
})
})
.collect();
let writers: Vec<_> = (0..WRITERS)
.map(|_| {
let s = Arc::clone(&session);
std::thread::spawn(move || {
for _ in 0..PER {
loop {
let mut tx = s.begin();
tx.working_mut().unwrap();
if matches!(s.commit(tx, true), CommitOutcome::Committed { .. }) {
break;
}
}
}
})
})
.collect();
for h in writers {
h.join().unwrap();
}
stop.store(true, std::sync::atomic::Ordering::Relaxed);
for h in readers {
h.join().unwrap();
}
assert_eq!(session.version(), total);
}
}