1use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2use std::sync::Arc;
3
4use parking_lot::Mutex;
5
6use linkrs_core::types::EdgeIdentifier;
7use linkrs_core::types::Timestamp;
8use linkrs_core::{StorageError, StorageResult};
9use linkrs_transaction::undo_log::UndoLogManager;
10use linkrs_transaction::{
11 MutationEntityKey, MutationResult, TransactionError, UndoLogEntry, VertexId,
12};
13
14use super::GraphStorageContext;
15use crate::StorageOperationContext;
16
17#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
19pub struct WriteGateStats {
20 pub acquisitions: u64,
22 pub wait_nanos: u64,
24}
25
26impl WriteGateStats {
27 pub fn share_since(
32 &self,
33 before: &WriteGateStats,
34 wall: std::time::Duration,
35 threads: usize,
36 ) -> f64 {
37 let waited = self.wait_nanos.saturating_sub(before.wait_nanos) as f64;
38 let budget = wall.as_nanos() as f64 * threads.max(1) as f64;
39 if budget <= 0.0 {
40 0.0
41 } else {
42 waited / budget
43 }
44 }
45
46 pub fn prefer_batch_commits(share: f64) -> bool {
56 share >= 0.20
57 }
58
59 pub fn contention_advice(share: f64) -> &'static str {
66 if Self::prefer_batch_commits(share) {
67 "gate-wait share above 20%: batch writes via batch_insert_edges or begin_auto_commit_group over per-statement auto-commit; gate sharding review only on true-machine (>= 8 cores) data, finer locks below the gate are not justified"
68 } else {
69 "gate-wait share below 20%: per-statement auto-commit is fine; keep explicit caller chunking"
70 }
71 }
72}
73
74pub(crate) struct AutoCommitWriteGate {
76 state: Mutex<Option<(std::thread::ThreadId, u64)>>,
84 condvar: parking_lot::Condvar,
85 acquisitions: AtomicU64,
87 wait_nanos: AtomicU64,
88}
89
90impl AutoCommitWriteGate {
91 pub(crate) fn new() -> Arc<Self> {
92 Arc::new(Self {
93 state: Mutex::new(None),
94 condvar: parking_lot::Condvar::new(),
95 acquisitions: AtomicU64::new(0),
96 wait_nanos: AtomicU64::new(0),
97 })
98 }
99
100 pub(crate) fn acquire(self: &Arc<Self>) -> Arc<AutoCommitWriteLease> {
101 let start = std::time::Instant::now();
102 let current = std::thread::current().id();
103 let mut guard = self.state.lock();
104 if let Some((holder, depth)) = guard.as_mut() {
106 if *holder == current {
107 *depth += 1;
108 drop(guard);
109 self.acquisitions.fetch_add(1, Ordering::Relaxed);
110 return Arc::new(AutoCommitWriteLease {
111 gate: self.clone(),
112 held: AtomicBool::new(true),
113 });
114 }
115 }
116 while guard.is_some() {
117 self.condvar.wait(&mut guard);
118 }
119 *guard = Some((current, 1));
120 drop(guard);
121 self.acquisitions.fetch_add(1, Ordering::Relaxed);
122 self.wait_nanos
123 .fetch_add(start.elapsed().as_nanos() as u64, Ordering::Relaxed);
124 Arc::new(AutoCommitWriteLease {
125 gate: self.clone(),
126 held: AtomicBool::new(true),
127 })
128 }
129
130 pub(crate) fn stats(&self) -> WriteGateStats {
131 WriteGateStats {
132 acquisitions: self.acquisitions.load(Ordering::Relaxed),
133 wait_nanos: self.wait_nanos.load(Ordering::Relaxed),
134 }
135 }
136
137 pub(crate) fn release(&self) {
138 let mut guard = self.state.lock();
139 match guard.as_mut() {
140 Some((_, depth)) if *depth > 1 => *depth -= 1,
141 Some(_) => {
142 *guard = None;
143 drop(guard);
144 self.condvar.notify_one();
145 }
146 None => {}
147 }
148 }
149}
150
151pub(crate) struct AutoCommitWriteLease {
152 gate: Arc<AutoCommitWriteGate>,
153 held: AtomicBool,
154}
155
156impl AutoCommitWriteLease {
157 pub(crate) fn release(&self) {
158 if self.held.swap(false, Ordering::AcqRel) {
159 self.gate.release();
160 }
161 }
162}
163
164impl Drop for AutoCommitWriteLease {
165 fn drop(&mut self) {
166 if self.held.swap(false, Ordering::AcqRel) {
167 self.gate.release();
168 }
169 }
170}
171
172pub(crate) struct AutoCommitMutationRecorder {
175 pub(crate) undo: Arc<Mutex<UndoLogManager>>,
176 pub(crate) write_set: Arc<Mutex<linkrs_transaction::types::WriteSet>>,
178}
179
180impl std::fmt::Debug for AutoCommitMutationRecorder {
181 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182 f.debug_struct("AutoCommitMutationRecorder").finish()
183 }
184}
185
186impl linkrs_transaction::TransactionMutationRecorder for AutoCommitMutationRecorder {
187 fn record_mutation(&self, mutation: MutationResult) -> Result<(), TransactionError> {
188 for entity_key in mutation.entity_keys {
189 match entity_key {
190 MutationEntityKey::Vertex(vertex_id) => self.record_vertex_write(vertex_id),
191 MutationEntityKey::Edge(edge) => self.record_edge_write(edge),
192 }
193 }
194 if let Some(entry) = mutation.undo_entry {
195 self.add_undo_log(entry)
196 } else {
197 Ok(())
198 }
199 }
200
201 fn record_vertex_write(&self, vertex_id: VertexId) {
202 self.write_set.lock().record_vertex(vertex_id);
203 }
204
205 fn record_vertex_delete(&self, vertex_id: VertexId) {
206 self.write_set.lock().record_vertex_delete(vertex_id);
207 }
208
209 fn record_edge_write(&self, edge: EdgeIdentifier) {
210 self.write_set.lock().record_edge(edge);
211 }
212
213 fn add_undo_log(&self, entry: UndoLogEntry) -> Result<(), TransactionError> {
214 self.undo
215 .lock()
216 .add(entry)
217 .map_err(|e| TransactionError::internal(e.to_string()))
218 }
219
220 fn record_table_modification(&self, _table_name: &str) {}
221}
222
223pub struct AutoCommitBatchWindow {
225 pub(crate) base_ctx: Arc<super::GraphStorageContext>,
226 pub(crate) gate_lease: Arc<AutoCommitWriteLease>,
227 pub(crate) first_ts: Mutex<Option<Timestamp>>,
228 pub(crate) statement_count: AtomicU64,
229 pub(crate) snapshot_rounds: AtomicU64,
230 pub(crate) group: AtomicBool,
233 pub(crate) group_undo: Option<Arc<Mutex<UndoLogManager>>>,
235 pub(crate) group_transaction_id: Mutex<Option<linkrs_core::types::TransactionId>>,
239 pub(crate) group_write_sets: Mutex<Vec<linkrs_transaction::types::WriteSet>>,
244}
245
246impl AutoCommitBatchWindow {
247 pub(crate) fn bind_statement(self: &Arc<Self>) -> StorageResult<super::GraphStorageContext> {
248 let base = &self.base_ctx;
249 let is_group = self.group.load(Ordering::Acquire);
250 let ts = {
251 let mut first = self.first_ts.lock();
252 match *first {
253 Some(ts) if is_group => ts,
254 Some(_) => base
255 .persistent
256 .version_manager
257 .try_next_write_timestamp()
258 .map_err(|error| StorageError::db_error(error.to_string()))?,
259 None => {
260 let ts = base
261 .persistent
262 .version_manager
263 .try_next_write_timestamp()
264 .map_err(|error| StorageError::db_error(error.to_string()))?;
265 *first = Some(ts);
266 self.snapshot_rounds.fetch_add(1, Ordering::SeqCst);
267 ts
268 }
269 }
270 };
271
272 let transaction_id = if is_group {
273 let mut shared = self.group_transaction_id.lock();
274 if shared.is_none() {
275 *shared = Some(linkrs_core::types::TransactionId::new(
276 base.persistent
277 .next_auto_transaction_id
278 .fetch_add(1, Ordering::SeqCst),
279 ));
280 }
281 *shared.as_ref().expect("group transaction id is set above")
282 } else {
283 linkrs_core::types::TransactionId::new(
284 base.persistent
285 .next_auto_transaction_id
286 .fetch_add(1, Ordering::SeqCst),
287 )
288 };
289 let (staging_start, wal_start) = if is_group {
293 (
294 Some(base.peek_txn_staging_mark(transaction_id).unwrap_or((0, 0))),
295 base.staged_wal_len_for(transaction_id),
296 )
297 } else {
298 (None, 0)
299 };
300 let undo_log = if is_group {
301 Arc::clone(
302 self.group_undo
303 .as_ref()
304 .expect("group mode requires group_undo to be initialized"),
305 )
306 } else {
307 Arc::new(Mutex::new(UndoLogManager::new()))
308 };
309 let group_undo_start = if is_group {
310 Some(undo_log.lock().len())
311 } else {
312 None
313 };
314 let write_set = Arc::new(parking_lot::Mutex::new(
315 linkrs_transaction::types::WriteSet::new(),
316 ));
317 let context = StorageOperationContext {
318 transaction_id: Some(transaction_id),
319 read_timestamp: ts,
320 write_timestamp: Some(ts),
321 read_only: false,
322 auto_commit: true,
323 mutation_recorder: Some(Arc::new(AutoCommitMutationRecorder {
324 undo: undo_log.clone(),
325 write_set: write_set.clone(),
326 })),
327 auto_commit_group_start: group_undo_start,
328 auto_commit_staging_start: staging_start,
329 auto_commit_wal_start: wal_start,
330 };
331
332 self.statement_count.fetch_add(1, Ordering::SeqCst);
333 let mut bound = (**base).clone();
334 bound.operation_context = Some(Arc::new(context));
335 bound.write_timestamp_lease = if is_group {
340 None
341 } else {
342 Some(Arc::new(super::WriteTimestampLease {
343 version_manager: base.persistent.version_manager.clone(),
344 timestamp: ts,
345 finalized: AtomicBool::new(false),
346 }))
347 };
348 bound.write_gate_lease = None;
349 bound.auto_commit_undo = Some(undo_log);
350 bound.auto_commit_write_set = Some(write_set);
351 bound.auto_commit_window = Some(self.clone());
352 Ok(bound)
353 }
354
355 pub fn finalize(&self) -> StorageResult<()> {
356 self.unregister_snapshots();
357 self.gate_lease.release();
358 Ok(())
359 }
360
361 pub fn statement_count(&self) -> u64 {
362 self.statement_count.load(Ordering::SeqCst)
363 }
364
365 pub fn snapshot_rounds(&self) -> u64 {
366 self.snapshot_rounds.load(Ordering::SeqCst)
367 }
368
369 pub fn is_grouped(&self) -> bool {
370 self.group.load(Ordering::Acquire)
371 }
372
373 pub fn finalize_group(&self) -> StorageResult<()> {
380 let group_txid = *self.group_transaction_id.lock();
384 if let Some(txid) = group_txid {
385 if let Err(error) = self.base_ctx.apply_txn_staging(txid) {
386 let _ = self.rollback_group();
387 return Err(error);
388 }
389 if let Err(error) = self.base_ctx.commit_staged_writes_grouped(txid, &[]) {
390 let _ = self.rollback_group();
391 return Err(error);
392 }
393 }
394 if let Some(persistence) = self.base_ctx.persistent.persistence.as_ref() {
396 if let Some(wal) = persistence.read().wal_manager() {
397 wal.read().sync()?;
398 let durable = wal.read().durable_lsn();
399 self.base_ctx
400 .persistent
401 .index_data_manager
402 .read()
403 .advance_barriers(linkrs_core::types::CommitLsn::new(durable.as_u64()));
404 }
405 }
406 if let Some(ts) = *self.first_ts.lock() {
409 let commit_ts = self.base_ctx.commit_write_timestamp_ordered(ts)?;
410 let sets = std::mem::take(&mut *self.group_write_sets.lock());
411 if !sets.is_empty() {
412 let mut merged = linkrs_transaction::types::WriteSet::new();
413 for set in sets {
414 merged.vertices.extend(set.vertices);
415 merged.edges.extend(set.edges);
416 merged.edge_endpoints.extend(set.edge_endpoints);
417 merged.deleted_vertices.extend(set.deleted_vertices);
418 merged.schema_resources.extend(set.schema_resources);
419 merged.index_resources.extend(set.index_resources);
420 merged.read_ranges.extend(set.read_ranges);
421 }
422 if !merged.is_empty() {
423 self.base_ctx.publish_committed_write_set(commit_ts, merged);
424 }
425 }
426 }
427 self.unregister_snapshots();
429 self.gate_lease.release();
430 Ok(())
431 }
432
433 pub fn rollback_group(&self) -> StorageResult<()> {
440 if let Some(undo) = &self.group_undo {
441 let mut manager = undo.lock();
442 let start_ts = self.first_ts.lock().unwrap_or(0);
443 manager
444 .execute_undo(&*self.base_ctx, start_ts)
445 .map_err(|e| StorageError::db_error(e.to_string()))?;
446 drop(manager);
447 if let Some(ts) = *self.first_ts.lock() {
448 self.base_ctx.abort_write_timestamp(ts);
449 }
450 }
451 if let Some(txid) = *self.group_transaction_id.lock() {
454 self.base_ctx.drop_staged_wal_for(txid);
455 self.base_ctx.discard_txn_staging_for(txid);
456 }
457 self.group_write_sets.lock().clear();
458 self.unregister_snapshots();
459 self.gate_lease.release();
460 Ok(())
461 }
462
463 fn unregister_snapshots(&self) {}
464}
465
466impl Drop for AutoCommitBatchWindow {
467 fn drop(&mut self) {
468 self.unregister_snapshots();
469 }
470}
471
472impl GraphStorageContext {
473 pub(crate) fn begin_auto_commit_batch(
474 &self,
475 ) -> StorageResult<Arc<super::AutoCommitBatchWindow>> {
476 let write_gate_lease = self.persistent.auto_commit_write_gate.acquire();
477 let mut clean = self.clone();
478 clean.operation_context = None;
479 clean.write_timestamp_lease = None;
480 clean.write_gate_lease = None;
481 clean.auto_commit_undo = None;
482 clean.auto_commit_write_set = None;
483 clean.auto_commit_window = None;
484 Ok(Arc::new(super::AutoCommitBatchWindow {
485 base_ctx: Arc::new(clean),
486 gate_lease: write_gate_lease,
487 first_ts: parking_lot::Mutex::new(None),
488 statement_count: std::sync::atomic::AtomicU64::new(0),
489 snapshot_rounds: std::sync::atomic::AtomicU64::new(0),
490 group: std::sync::atomic::AtomicBool::new(false),
491 group_undo: None,
492 group_write_sets: parking_lot::Mutex::new(Vec::new()),
493 group_transaction_id: parking_lot::Mutex::new(None),
494 }))
495 }
496
497 pub(crate) fn begin_auto_commit_group(
498 &self,
499 ) -> StorageResult<Arc<super::AutoCommitBatchWindow>> {
500 let write_gate_lease = self.persistent.auto_commit_write_gate.acquire();
501 let mut clean = self.clone();
502 clean.operation_context = None;
503 clean.write_timestamp_lease = None;
504 clean.write_gate_lease = None;
505 clean.auto_commit_undo = None;
506 clean.auto_commit_write_set = None;
507 clean.auto_commit_window = None;
508 Ok(Arc::new(super::AutoCommitBatchWindow {
509 base_ctx: Arc::new(clean),
510 gate_lease: write_gate_lease,
511 first_ts: parking_lot::Mutex::new(None),
512 statement_count: std::sync::atomic::AtomicU64::new(0),
513 snapshot_rounds: std::sync::atomic::AtomicU64::new(0),
514 group: std::sync::atomic::AtomicBool::new(true),
515 group_undo: Some(Arc::new(parking_lot::Mutex::new(
516 linkrs_transaction::UndoLogManager::new(),
517 ))),
518 group_write_sets: parking_lot::Mutex::new(Vec::new()),
519 group_transaction_id: parking_lot::Mutex::new(None),
520 }))
521 }
522
523 pub(crate) fn restore_auto_transaction_id(&self, max_transaction_id: u64) {
524 self.persistent
525 .next_auto_transaction_id
526 .fetch_max(max_transaction_id.saturating_add(1), Ordering::SeqCst);
527 }
528
529 pub(crate) fn abort_write_timestamp(&self, timestamp: Timestamp) {
530 if let Some(lease) = &self.write_timestamp_lease {
531 lease.abort();
532 } else if self.operation_context.is_none() {
533 self.persistent
534 .version_manager
535 .abort_write_timestamp(timestamp);
536 }
537 }
538
539 pub(crate) fn commit_write_timestamp_ordered(
554 &self,
555 start: Timestamp,
556 ) -> StorageResult<Timestamp> {
557 let version_manager = &self.persistent.version_manager;
558 let commit_ts = match version_manager.commit_ordered(start) {
559 Ok(commit_ts) => commit_ts,
560 Err(linkrs_transaction::VersionManagerError::InvalidTimestamp(_)) => start,
561 Err(error) => return Err(StorageError::db_error(error.to_string())),
562 };
563 if let Some(lease) = &self.write_timestamp_lease {
564 lease.finalized.store(true, Ordering::SeqCst);
565 }
566 Ok(commit_ts)
567 }
568
569 pub(crate) fn finalize_operation(&self, committed: bool) -> StorageResult<()> {
570 let Some(operation) = &self.operation_context else {
571 return Ok(());
572 };
573 if !operation.auto_commit {
574 return Ok(());
575 }
576
577 if self.auto_commit_window.is_none() {
578 self.unregister_statement_snapshots(operation);
579 }
580
581 if operation.read_only {
582 return Ok(());
583 }
584
585 if let Some(window) = &self.auto_commit_window {
592 if window.is_grouped() {
593 let timestamp = operation.write_timestamp.ok_or_else(|| {
594 StorageError::db_error("Group operation has no write timestamp")
595 })?;
596 let rewind = |ctx: &Self| {
597 if let Some(txid) = operation.transaction_id {
598 let staging_start = operation.auto_commit_staging_start.unwrap_or((0, 0));
599 ctx.rewind_grouped_statement(
600 txid,
601 staging_start,
602 operation.auto_commit_wal_start,
603 );
604 }
605 };
606 if committed {
607 if let Some(conflict) = self.auto_commit_conflict(operation) {
608 if let Some(undo) = &self.auto_commit_undo {
609 let mut log = undo.lock();
610 let start = operation.auto_commit_group_start.unwrap_or(0);
611 if let Err(error) = log.execute_undo_from_index(self, timestamp, start)
612 {
613 log::error!("Group statement rollback failed: {}", error);
614 }
615 }
616 rewind(self);
617 self.maybe_run_index_gc();
618 return Err(conflict);
619 }
620 if let Some(write_set) = self.auto_commit_write_set.as_ref() {
621 let set = write_set.lock().clone();
622 if !set.is_empty() {
623 window.group_write_sets.lock().push(set);
624 }
625 }
626 } else {
627 if let Some(undo) = &self.auto_commit_undo {
628 let mut log = undo.lock();
629 let start = operation.auto_commit_group_start.unwrap_or(0);
630 if let Err(error) = log.execute_undo_from_index(self, timestamp, start) {
631 log::error!("Group statement rollback failed: {}", error);
632 }
633 }
634 rewind(self);
635 }
636 self.maybe_run_index_gc();
637 return Ok(());
638 }
639 }
640
641 let timestamp = operation.write_timestamp.ok_or_else(|| {
642 StorageError::db_error("Auto-commit operation has no write timestamp")
643 })?;
644 let transaction_id = operation.transaction_id;
645
646 if committed {
647 if let Some(conflict) = self.auto_commit_conflict(operation) {
648 if let Some(undo) = &self.auto_commit_undo {
649 let mut log = undo.lock();
650 if let Err(error) = log.execute_undo(self, timestamp) {
651 log::error!("Auto-commit rollback failed: {}", error);
652 }
653 }
654 self.abort_write_timestamp(timestamp);
655 if let Some(lease) = &self.write_gate_lease {
656 lease.release();
657 }
658 if let Some(transaction_id) = transaction_id {
659 self.abort_staged_writes(transaction_id);
660 }
661 self.maybe_run_index_gc();
662 return Err(conflict);
663 }
664 if let Some(transaction_id) = transaction_id {
669 if let Err(error) = self.apply_txn_staging(transaction_id) {
670 if let Some(undo) = &self.auto_commit_undo {
671 let mut log = undo.lock();
672 if let Err(undo_error) = log.execute_undo(self, timestamp) {
673 log::error!("Auto-commit rollback failed: {}", undo_error);
674 }
675 }
676 self.abort_write_timestamp(timestamp);
677 if let Some(lease) = &self.write_gate_lease {
678 lease.release();
679 }
680 self.abort_staged_writes(transaction_id);
681 self.maybe_run_index_gc();
682 return Err(error);
683 }
684 }
685 let commit_ts = self.commit_write_timestamp_ordered(timestamp)?;
689 self.publish_auto_commit_write_set(commit_ts);
690 } else {
691 if let Some(undo) = &self.auto_commit_undo {
692 let mut log = undo.lock();
693 if let Err(error) = log.execute_undo(self, timestamp) {
694 log::error!("Auto-commit rollback failed: {}", error);
695 }
696 }
697 self.abort_write_timestamp(timestamp);
698 }
699 if let Some(lease) = &self.write_gate_lease {
700 lease.release();
701 }
702 if let Some(transaction_id) = transaction_id {
703 self.persistent.staged_wal.remove(&transaction_id);
704 self.discard_txn_staging_for(transaction_id);
707 }
708 self.maybe_run_index_gc();
709 Ok(())
710 }
711
712 fn auto_commit_conflict(&self, operation: &StorageOperationContext) -> Option<StorageError> {
716 let write_set = self.auto_commit_write_set.as_ref()?.lock().clone();
717 if write_set.is_empty() {
718 return None;
719 }
720 if self
721 .persistent
722 .committed_write_sets
723 .has_conflict(&write_set, operation.read_timestamp)
724 {
725 return Some(StorageError::write_write_conflict(format!(
726 "auto-commit statement at ts={} overlaps a recently committed write",
727 operation.read_timestamp,
728 )));
729 }
730 None
731 }
732
733 #[cfg(test)]
736 pub(crate) fn committed_write_conflict_probe(
737 &self,
738 write_set: &linkrs_transaction::types::WriteSet,
739 read_ts: Timestamp,
740 ) -> bool {
741 self.persistent
742 .committed_write_sets
743 .has_conflict(write_set, read_ts)
744 }
745
746 fn publish_auto_commit_write_set(&self, commit_ts: Timestamp) {
749 let Some(write_set) = self.auto_commit_write_set.as_ref() else {
750 return;
751 };
752 let write_set = write_set.lock().clone();
753 self.publish_committed_write_set(commit_ts, write_set);
754 }
755
756 pub(crate) fn publish_committed_write_set(
761 &self,
762 commit_ts: Timestamp,
763 write_set: linkrs_transaction::types::WriteSet,
764 ) {
765 let horizon = self
766 .persistent
767 .version_manager
768 .snapshot_tracker()
769 .cleanup_threshold();
770 self.persistent
771 .committed_write_sets
772 .publish(commit_ts, write_set, horizon);
773 }
774
775 fn unregister_statement_snapshots(&self, _operation: &StorageOperationContext) {}
776}