1use std::any::Any;
29use std::sync::{Arc, Mutex};
30
31use khive_storage::{AtomicUnitOp, BoxFuture, SqlAccess, SqlStatement, SqlWriter, StorageError};
32
33use crate::atomic_plan::{
34 AddEntityPlan, AddNotePlan, AffectedRowGuard, DeletePlan, GovernancePlan, GtdCompletePlan,
35 GtdTransitionPlan, LinkPlan, MergePlan, PlanStatement, PostCommitEffect, UpdatePlan,
36};
37
38#[derive(Debug, Clone)]
45pub enum AtomicOpPlan {
46 AddEntity(AddEntityPlan),
47 AddNote(Box<AddNotePlan>),
48 Update(Box<UpdatePlan>),
49 Delete(DeletePlan),
50 Link(LinkPlan),
51 Merge(MergePlan),
52 GtdTransition(GtdTransitionPlan),
53 GtdComplete(GtdCompletePlan),
54 Governance(GovernancePlan),
55}
56
57impl AtomicOpPlan {
58 fn plan_statements(&self) -> Vec<(PlanStatement, bool)> {
68 let statements = match self {
69 AtomicOpPlan::AddEntity(p) => p.statements.clone(),
70 AtomicOpPlan::AddNote(p) => p.statements.clone(),
71 AtomicOpPlan::Update(p) => p.statements.clone(),
72 AtomicOpPlan::Delete(p) => p.statements.clone(),
73 AtomicOpPlan::Link(p) => p.statements.clone(),
74 AtomicOpPlan::Merge(p) => {
75 let mut statements: Vec<PlanStatement> = p
76 .rewires
77 .iter()
78 .map(|rewire| PlanStatement {
79 statement: rewire.statement.clone(),
80 guard: None,
81 })
82 .collect();
83 statements.extend(p.lifecycle.clone());
84 statements
85 }
86 AtomicOpPlan::GtdTransition(p) => p.statements.clone(),
87 AtomicOpPlan::GtdComplete(p) => p.statements.clone(),
88 AtomicOpPlan::Governance(p) => p.statements.clone(),
89 };
90 let mut statements: Vec<_> = statements
91 .into_iter()
92 .map(|statement| (statement, false))
93 .collect();
94 if let AtomicOpPlan::Update(plan) = self {
95 statements.extend(plan.graph_effects.iter().map(|effect| match effect {
96 crate::atomic_plan::NoteUpdateStatement::Write(statement) => {
97 (statement.clone(), false)
98 }
99 crate::atomic_plan::NoteUpdateStatement::Assert(statement) => {
100 (statement.clone(), true)
101 }
102 }));
103 }
104 statements
105 }
106
107 fn post_commit_effect(&self) -> Option<PostCommitEffect> {
121 match self {
122 AtomicOpPlan::AddEntity(p) if p.post_commit != PostCommitEffect::None => {
123 Some(p.post_commit.clone())
124 }
125 AtomicOpPlan::AddNote(p) if p.post_commit != PostCommitEffect::None => {
126 Some(p.post_commit.clone())
127 }
128 AtomicOpPlan::Update(p) if p.post_commit != PostCommitEffect::None => {
129 Some(p.post_commit.clone())
130 }
131 AtomicOpPlan::Delete(p) if p.post_commit != PostCommitEffect::None => {
132 Some(p.post_commit.clone())
133 }
134 AtomicOpPlan::GtdTransition(p) if p.post_commit != PostCommitEffect::None => {
135 Some(p.post_commit.clone())
136 }
137 AtomicOpPlan::GtdComplete(p) if p.post_commit != PostCommitEffect::None => {
138 Some(p.post_commit.clone())
139 }
140 _ => None,
141 }
142 }
143}
144
145#[derive(Debug, Clone, PartialEq, Eq)]
148pub enum AtomicOpFailure {
149 NoteConflict(crate::note_write::NoteWriteConflict),
150 EntityConflict(crate::entity_write::EntityVersionConflict),
151 GuardFailed {
157 statement_label: Option<String>,
158 expected: AffectedRowGuard,
159 observed: u64,
160 },
161 SqlError {
165 statement_label: Option<String>,
166 message: String,
167 },
168}
169
170#[derive(Debug, PartialEq, Eq)]
192pub struct CommittedPostCommitEffects {
193 effects: Vec<PostCommitEffect>,
194}
195
196impl CommittedPostCommitEffects {
197 fn new(effects: Vec<PostCommitEffect>) -> Self {
198 Self { effects }
199 }
200
201 pub fn as_slice(&self) -> &[PostCommitEffect] {
204 &self.effects
205 }
206
207 pub(crate) fn into_effects(self) -> Vec<PostCommitEffect> {
208 self.effects
209 }
210}
211
212#[derive(Debug, PartialEq, Eq)]
216pub enum AtomicRunOutcome {
217 Committed {
225 post_commit: CommittedPostCommitEffects,
226 },
227 RolledBack {
235 failed_op_index: usize,
236 failure: AtomicOpFailure,
237 },
238}
239
240#[derive(Debug)]
247pub struct AtomicRunnerError(pub StorageError);
248
249impl std::fmt::Display for AtomicRunnerError {
250 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251 write!(f, "atomic_unit seam failure: {}", self.0)
252 }
253}
254
255impl std::error::Error for AtomicRunnerError {
256 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
257 Some(&self.0)
258 }
259}
260
261fn raw_stmt(sql: String) -> SqlStatement {
262 SqlStatement {
263 sql,
264 params: vec![],
265 label: None,
266 }
267}
268
269async fn begin_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
270 writer
271 .execute(raw_stmt(format!("SAVEPOINT {name}")))
272 .await
273 .map(|_| ())
274}
275
276async fn release_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
277 writer
278 .execute(raw_stmt(format!("RELEASE {name}")))
279 .await
280 .map(|_| ())
281}
282
283async fn rollback_to_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
284 writer
285 .execute(raw_stmt(format!("ROLLBACK TO {name}")))
286 .await
287 .map(|_| ())
288}
289
290pub(crate) struct AppliedPlan {
297 pub(crate) effect: Option<PostCommitEffect>,
298 pub(crate) note_version: Option<i64>,
299}
300
301pub(crate) async fn apply_plan(
302 writer: &mut dyn SqlWriter,
303 plan: &AtomicOpPlan,
304 capture_note_versions: bool,
305) -> Result<AppliedPlan, AtomicOpFailure> {
306 if let AtomicOpPlan::Update(plan) = plan {
307 if let Some(guard) = &plan.entity_guard {
308 if let Some(conflict) =
309 guard
310 .check(writer)
311 .await
312 .map_err(|error| AtomicOpFailure::SqlError {
313 statement_label: Some("entity-version-precondition".into()),
314 message: error.to_string(),
315 })?
316 {
317 return Err(AtomicOpFailure::EntityConflict(conflict));
318 }
319 }
320 }
321 let note_guard = match plan {
322 AtomicOpPlan::Update(plan) => plan.note_guard.as_ref(),
323 AtomicOpPlan::AddNote(plan) => plan.note_guard.as_ref(),
324 _ => None,
325 };
326 if let Some(guard) = note_guard {
327 if let Some(conflict) =
328 guard
329 .check_fence(writer)
330 .await
331 .map_err(|error| AtomicOpFailure::SqlError {
332 statement_label: Some("note-fence".into()),
333 message: error.to_string(),
334 })?
335 {
336 return Err(AtomicOpFailure::NoteConflict(conflict));
337 }
338 }
339 let mut post_commit = plan.post_commit_effect();
340 if let AtomicOpPlan::Update(plan) = plan {
341 if let Some(inheritance) = &plan.note_embedding_inheritance {
342 if !inheritance
343 .vectors
344 .has_rows(writer)
345 .await
346 .map_err(|error| AtomicOpFailure::SqlError {
347 statement_label: Some("note-embedding-inheritance".into()),
348 message: error.to_string(),
349 })?
350 {
351 post_commit = Some(PostCommitEffect::NoteChanged {
352 note_id: plan.target_id,
353 kind: inheritance.kind.clone(),
354 });
355 }
356 }
357 }
358 for (index, (stmt, read_assertion)) in plan.plan_statements().into_iter().enumerate() {
359 let label = stmt.statement.label.clone();
360 let result = if read_assertion
363 || matches!(plan, AtomicOpPlan::GtdTransition(p) if p.idempotent_noop)
364 || matches!(plan, AtomicOpPlan::Update(p) if p.idempotent_noop)
365 {
366 writer
367 .query_all(stmt.statement)
368 .await
369 .map(|rows| rows.len() as u64)
370 } else {
371 writer.execute(stmt.statement).await
372 };
373 let affected = result.map_err(|e| AtomicOpFailure::SqlError {
374 statement_label: label.clone(),
375 message: e.to_string(),
376 })?;
377 if let Some(guard) = stmt.guard {
378 if !guard.holds_for(affected) {
379 if let Some(note_guard) = note_guard.filter(|_| index == 0) {
380 if let Some(conflict) =
381 note_guard.classify_refusal(writer).await.map_err(|error| {
382 AtomicOpFailure::SqlError {
383 statement_label: label.clone(),
384 message: error.to_string(),
385 }
386 })?
387 {
388 return Err(AtomicOpFailure::NoteConflict(conflict));
389 }
390 }
391 return Err(AtomicOpFailure::GuardFailed {
392 statement_label: label,
393 expected: guard,
394 observed: affected,
395 });
396 }
397 }
398 }
399 if let AtomicOpPlan::Update(plan) = plan {
400 if let Some(purge) = &plan.note_vector_purge {
401 purge
402 .apply(writer)
403 .await
404 .map_err(|error| AtomicOpFailure::SqlError {
405 statement_label: Some("note-vector-purge".into()),
406 message: error.to_string(),
407 })?;
408 }
409 }
410 let note_version = if let AtomicOpPlan::AddNote(plan) = plan {
411 if capture_note_versions {
412 let current = writer
413 .query_scalar(crate::note_write::statement(
414 "SELECT version FROM notes WHERE id=?1",
415 vec![khive_storage::SqlValue::Text(plan.note_id.to_string())],
416 ))
417 .await
418 .map_err(|error| AtomicOpFailure::SqlError {
419 statement_label: Some("note-version-receipt".into()),
420 message: error.to_string(),
421 })?;
422 let Some(khive_storage::SqlValue::Integer(version)) = current else {
423 return Err(AtomicOpFailure::SqlError {
424 statement_label: Some("note-version-receipt".into()),
425 message: "missing persisted note version".into(),
426 });
427 };
428 Some(version)
429 } else {
430 None
431 }
432 } else {
433 None
434 };
435 Ok(AppliedPlan {
436 effect: post_commit,
437 note_version,
438 })
439}
440
441pub async fn run_atomic_unit(
464 access: &dyn SqlAccess,
465 plans: Vec<AtomicOpPlan>,
466) -> Result<AtomicRunOutcome, AtomicRunnerError> {
467 run_atomic_unit_with_note_versions(access, plans, false)
468 .await
469 .map(|(outcome, _)| outcome)
470}
471
472pub(crate) async fn run_atomic_unit_with_note_versions(
475 access: &dyn SqlAccess,
476 plans: Vec<AtomicOpPlan>,
477 capture_note_versions: bool,
478) -> Result<(AtomicRunOutcome, Vec<i64>), AtomicRunnerError> {
479 let op: PreparedAtomicOp<Vec<i64>, (usize, AtomicOpFailure)> = Box::new(move |writer| {
480 Box::pin(async move {
481 let mut post_commit = Vec::new();
482 let mut note_versions = Vec::new();
483 for (op_index, plan) in plans.iter().enumerate() {
484 let savepoint = format!("adr099_atomic_op_{op_index}");
485 begin_savepoint(writer, &savepoint).await?;
486 match apply_plan(writer, plan, capture_note_versions).await {
487 Ok(applied) => {
488 release_savepoint(writer, &savepoint).await?;
489 if let Some(effect) = applied.effect {
490 post_commit.push(effect);
491 }
492 if let Some(version) = applied.note_version {
493 note_versions.push(version);
494 }
495 }
496 Err(failure) => {
497 let _ = rollback_to_savepoint(writer, &savepoint).await;
505 let _ = release_savepoint(writer, &savepoint).await;
506 return Err(PreparedAtomicError::Refused {
507 failure: (op_index, failure),
508 message: format!("ADR-099 atomic unit aborted at op {op_index}"),
509 });
510 }
511 }
512 }
513 Ok((note_versions, post_commit))
514 })
515 });
516
517 match run_prepared_atomic_unit(access, op)
518 .await
519 .map_err(AtomicRunnerError)?
520 {
521 PreparedAtomicOutcome::Committed { value, post_commit } => {
522 Ok((AtomicRunOutcome::Committed { post_commit }, value))
523 }
524 PreparedAtomicOutcome::RolledBack((failed_op_index, failure)) => Ok((
525 AtomicRunOutcome::RolledBack {
526 failed_op_index,
527 failure,
528 },
529 Vec::new(),
530 )),
531 }
532}
533
534pub(crate) enum PreparedAtomicError<E> {
535 Refused { failure: E, message: String },
536 Storage(StorageError),
537}
538
539impl<E> From<StorageError> for PreparedAtomicError<E> {
540 fn from(error: StorageError) -> Self {
541 Self::Storage(error)
542 }
543}
544
545pub(crate) enum PreparedAtomicOutcome<T, E> {
546 Committed {
547 value: T,
548 post_commit: CommittedPostCommitEffects,
549 },
550 RolledBack(E),
551}
552
553pub(crate) type PreparedAtomicOp<T, E> = Box<
554 dyn for<'w> FnOnce(
555 &'w mut dyn SqlWriter,
556 )
557 -> BoxFuture<'w, Result<(T, Vec<PostCommitEffect>), PreparedAtomicError<E>>>
558 + Send,
559>;
560
561pub(crate) async fn run_prepared_atomic_unit<T: Send + 'static, E: Send + 'static>(
564 access: &dyn SqlAccess,
565 prepared: PreparedAtomicOp<T, E>,
566) -> Result<PreparedAtomicOutcome<T, E>, StorageError> {
567 let failure_slot = Arc::new(Mutex::new(None));
568 let recorded = Arc::clone(&failure_slot);
569 let op: AtomicUnitOp = Box::new(move |writer| {
570 Box::pin(async move {
571 match prepared(writer).await {
572 Ok(value) => Ok(Box::new(value) as Box<dyn Any + Send>),
573 Err(PreparedAtomicError::Storage(error)) => Err(error),
574 Err(PreparedAtomicError::Refused { failure, message }) => {
575 *recorded
576 .lock()
577 .expect("atomic runner failure slot poisoned") = Some(failure);
578 Err(StorageError::Internal(message))
579 }
580 }
581 })
582 });
583 match access.atomic_unit(op).await {
584 Ok(boxed) => {
585 let (value, effects) = *boxed
586 .downcast::<(T, Vec<PostCommitEffect>)>()
587 .map_err(|_| StorageError::Internal("invalid prepared atomic outcome".into()))?;
588 Ok(PreparedAtomicOutcome::Committed {
589 value,
590 post_commit: CommittedPostCommitEffects::new(effects),
591 })
592 }
593 Err(storage_err) => {
594 if !atomic_unit_error_allows_recorded_refusal(&storage_err) {
596 return Err(storage_err);
597 }
598 let recorded = failure_slot
599 .lock()
600 .expect("atomic runner failure slot poisoned")
601 .take();
602 match recorded {
603 Some(failure) => Ok(PreparedAtomicOutcome::RolledBack(failure)),
604 None => Err(storage_err),
605 }
606 }
607 }
608}
609
610pub(crate) fn atomic_unit_error_allows_recorded_refusal(error: &StorageError) -> bool {
611 match error {
612 StorageError::WriterTaskTerminated { .. } => false,
613 StorageError::WriterTaskRequestFailed { request_state, .. } => {
614 *request_state == khive_storage::WriterTaskRequestState::TransactionRolledBack
615 }
616 _ => true,
617 }
618}
619
620#[cfg(test)]
621mod tests {
622 use super::*;
623
624 use std::sync::Arc as StdArc;
625
626 use khive_db::{ConnectionPool, PoolConfig, SqlBridge};
627 use khive_storage::types::{SqlValue, StorageResult as StorageResultAlias};
628 use uuid::Uuid;
629
630 struct TestPool {
632 pool: StdArc<ConnectionPool>,
633 _temp_dir: tempfile::TempDir,
634 }
635
636 impl std::ops::Deref for TestPool {
637 type Target = StdArc<ConnectionPool>;
638
639 fn deref(&self) -> &Self::Target {
640 &self.pool
641 }
642 }
643
644 fn scratch_pool(name: &str) -> TestPool {
651 let dir = tempfile::tempdir().expect("tempdir");
652 let path = dir.path().join(format!("{name}.db"));
653 let pool = StdArc::new(
654 ConnectionPool::new(PoolConfig {
655 path: Some(path),
656 write_queue_enabled: Some(true),
657 ..PoolConfig::for_test()
658 })
659 .expect("pool open"),
660 );
661 TestPool {
662 pool,
663 _temp_dir: dir,
664 }
665 }
666
667 fn seed_schema(pool: &ConnectionPool) {
675 let writer = pool.try_writer().expect("writer");
676 writer
677 .conn()
678 .execute_batch(
679 "CREATE TABLE entities (
680 id TEXT PRIMARY KEY,
681 namespace TEXT NOT NULL,
682 kind TEXT NOT NULL,
683 entity_type TEXT,
684 name TEXT NOT NULL,
685 description TEXT,
686 properties TEXT,
687 tags TEXT NOT NULL DEFAULT '[]',
688 created_at INTEGER NOT NULL,
689 updated_at INTEGER NOT NULL,
690 version INTEGER NOT NULL DEFAULT 1,
691 deleted_at INTEGER,
692 merged_into TEXT,
693 merge_event_id TEXT
694 );
695 CREATE TABLE graph_edges (
696 namespace TEXT NOT NULL,
697 id TEXT NOT NULL,
698 source_id TEXT NOT NULL,
699 target_id TEXT NOT NULL,
700 relation TEXT NOT NULL,
701 weight REAL NOT NULL DEFAULT 1.0,
702 created_at INTEGER NOT NULL,
703 updated_at INTEGER NOT NULL,
704 deleted_at INTEGER,
705 metadata TEXT,
706 target_backend TEXT,
707 PRIMARY KEY (namespace, id)
708 );",
709 )
710 .expect("seed schema");
711 }
712
713 fn insert_entity(pool: &ConnectionPool, id: Uuid, name: &str) {
714 let writer = pool.try_writer().expect("writer");
715 writer
716 .conn()
717 .execute(
718 "INSERT INTO entities \
719 (id, namespace, kind, name, created_at, updated_at) \
720 VALUES (?1, 'local', 'concept', ?2, 0, 0)",
721 rusqlite::params![id.to_string(), name],
722 )
723 .expect("insert entity");
724 }
725
726 fn entity_exists(pool: &ConnectionPool, id: Uuid) -> bool {
727 let writer = pool.try_writer().expect("writer");
728 let count: i64 = writer
729 .conn()
730 .query_row(
731 "SELECT COUNT(*) FROM entities WHERE id = ?1 AND deleted_at IS NULL",
732 rusqlite::params![id.to_string()],
733 |row| row.get(0),
734 )
735 .expect("query entity");
736 count > 0
737 }
738
739 fn edge_count(pool: &ConnectionPool, source: Uuid, target: Uuid) -> i64 {
740 let writer = pool.try_writer().expect("writer");
741 writer
742 .conn()
743 .query_row(
744 "SELECT COUNT(*) FROM graph_edges \
745 WHERE source_id = ?1 AND target_id = ?2 AND deleted_at IS NULL",
746 rusqlite::params![source.to_string(), target.to_string()],
747 |row| row.get(0),
748 )
749 .expect("query edge")
750 }
751
752 fn entities_snapshot(pool: &ConnectionPool) -> Vec<String> {
753 let writer = pool.try_writer().expect("writer");
754 let mut stmt = writer
755 .conn()
756 .prepare("SELECT id, name, deleted_at FROM entities ORDER BY id")
757 .expect("prepare snapshot");
758 let rows = stmt
759 .query_map([], |row| {
760 let id: String = row.get(0)?;
761 let name: String = row.get(1)?;
762 let deleted_at: Option<i64> = row.get(2)?;
763 Ok(format!("{id}:{name}:{deleted_at:?}"))
764 })
765 .expect("query snapshot");
766 rows.collect::<Result<Vec<_>, _>>()
767 .expect("collect snapshot")
768 }
769
770 fn delete_plan(id: Uuid, label: &str) -> AtomicOpPlan {
771 AtomicOpPlan::Delete(DeletePlan {
772 target_id: id,
773 statements: vec![PlanStatement {
774 statement: SqlStatement {
775 sql: "DELETE FROM entities WHERE id = ?1 AND deleted_at IS NULL".to_string(),
776 params: vec![SqlValue::Text(id.to_string())],
777 label: Some(label.to_string()),
778 },
779 guard: Some(AffectedRowGuard::exactly(1)),
780 }],
781 post_commit: PostCommitEffect::None,
782 })
783 }
784
785 fn rename_plan(id: Uuid, new_name: &str, label: &str) -> AtomicOpPlan {
786 AtomicOpPlan::Update(Box::new(UpdatePlan {
787 graph_effects: Vec::new(),
788 note_vector_purge: None,
789 note_embedding_inheritance: None,
790 entity_guard: None,
791 note_guard: None,
792 target_id: id,
793 statements: vec![PlanStatement {
794 statement: SqlStatement {
795 sql: "UPDATE entities SET version = version + 1, name = ?1, updated_at = 1 \
796 WHERE id = ?2 AND deleted_at IS NULL"
797 .to_string(),
798 params: vec![
799 SqlValue::Text(new_name.to_string()),
800 SqlValue::Text(id.to_string()),
801 ],
802 label: Some(label.to_string()),
803 },
804 guard: Some(AffectedRowGuard::exactly(1)),
805 }],
806 post_commit: PostCommitEffect::None,
807 edge_natural_key: None,
808 idempotent_noop: false,
809 }))
810 }
811
812 fn link_plan(edge_id: Uuid, source: Uuid, target: Uuid) -> AtomicOpPlan {
817 AtomicOpPlan::Link(LinkPlan {
818 source_id: source,
819 target_id: target,
820 statements: vec![PlanStatement {
821 statement: SqlStatement {
822 sql: "INSERT INTO graph_edges \
823 (namespace, id, source_id, target_id, relation, created_at, updated_at) \
824 SELECT 'local', ?1, ?2, ?3, 'annotates', 0, 0 \
825 WHERE EXISTS (SELECT 1 FROM entities WHERE id = ?2 AND deleted_at IS NULL) \
826 AND EXISTS (SELECT 1 FROM entities WHERE id = ?3 AND deleted_at IS NULL)"
827 .to_string(),
828 params: vec![
829 SqlValue::Text(edge_id.to_string()),
830 SqlValue::Text(source.to_string()),
831 SqlValue::Text(target.to_string()),
832 ],
833 label: Some("insert-edge-where-exists".to_string()),
834 },
835 guard: Some(AffectedRowGuard::exactly(1)),
836 }],
837 disposition: khive_storage::EdgeUpsertDisposition::Created,
838 })
839 }
840
841 fn merge_plan(into_id: Uuid, from_id: Uuid) -> AtomicOpPlan {
842 AtomicOpPlan::Merge(MergePlan {
843 into_id,
844 from_id,
845 rewires: vec![crate::atomic_plan::PlanPredicate {
846 description: "source_id = :from".to_string(),
847 statement: SqlStatement {
848 sql: "UPDATE graph_edges SET source_id = ?1, updated_at = 1 \
849 WHERE source_id = ?2"
850 .to_string(),
851 params: vec![
852 SqlValue::Text(into_id.to_string()),
853 SqlValue::Text(from_id.to_string()),
854 ],
855 label: Some("merge-rewire".to_string()),
856 },
857 }],
858 lifecycle: vec![PlanStatement {
859 statement: SqlStatement {
860 sql: "UPDATE entities SET version = version + 1, deleted_at = 1, merged_into = ?1 \
861 WHERE id = ?2 AND deleted_at IS NULL"
862 .to_string(),
863 params: vec![
864 SqlValue::Text(into_id.to_string()),
865 SqlValue::Text(from_id.to_string()),
866 ],
867 label: Some("tombstone-from-entity".to_string()),
868 },
869 guard: Some(AffectedRowGuard::exactly(1)),
870 }],
871 })
872 }
873
874 #[tokio::test]
879 async fn rollback_end_to_end_leaves_zero_partial_state() {
880 let pool = scratch_pool("rollback_end_to_end");
881 seed_schema(&pool);
882 let a = Uuid::new_v4();
883 let b = Uuid::new_v4();
884 insert_entity(&pool, a, "alpha");
885 insert_entity(&pool, b, "bravo");
886 let before = entities_snapshot(&pool);
887
888 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
889 let plans = vec![
890 rename_plan(a, "alpha-renamed", "rename-a"),
891 delete_plan(Uuid::new_v4(), "delete-nonexistent"),
894 rename_plan(b, "bravo-renamed", "rename-b"),
895 ];
896
897 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
898 match outcome {
899 AtomicRunOutcome::RolledBack {
900 failed_op_index,
901 failure,
902 } => {
903 assert_eq!(failed_op_index, 1);
904 assert!(matches!(failure, AtomicOpFailure::GuardFailed { .. }));
905 }
906 other => panic!("expected RolledBack, got {other:?}"),
907 }
908
909 let after = entities_snapshot(&pool);
910 assert_eq!(
911 before, after,
912 "a mid-unit failure must leave the database byte-for-byte unchanged"
913 );
914 }
915
916 #[tokio::test]
923 async fn commit_pass_resolves_on_first_poll_and_commits() {
924 let pool = scratch_pool("suspend_trap_happy_path");
925 seed_schema(&pool);
926 let a = Uuid::new_v4();
927 insert_entity(&pool, a, "alpha");
928
929 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
930 let plans = vec![rename_plan(a, "alpha-v2", "rename-a")];
931
932 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
933 assert!(
934 matches!(outcome, AtomicRunOutcome::Committed { .. }),
935 "the real commit-pass closure (SAVEPOINT + guarded DML only) must \
936 resolve on block_on_sync's first poll and commit: {outcome:?}"
937 );
938
939 let writer = pool.try_writer().expect("writer");
940 let name: String = writer
941 .conn()
942 .query_row(
943 "SELECT name FROM entities WHERE id = ?1",
944 rusqlite::params![a.to_string()],
945 |row| row.get(0),
946 )
947 .expect("query renamed entity");
948 assert_eq!(name, "alpha-v2");
949 }
950
951 #[tokio::test]
952 async fn hand_built_suspending_closure_fails_loudly_through_the_same_seam() {
953 let pool = scratch_pool("suspend_trap_misuse");
963 seed_schema(&pool);
964
965 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
966 let suspending_op: AtomicUnitOp = Box::new(|_writer| {
967 Box::pin(async move {
968 tokio::task::yield_now().await;
969 Ok(Box::new(()) as Box<dyn Any + Send>)
970 })
971 });
972
973 let result: StorageResultAlias<Box<dyn Any + Send>> =
974 khive_storage::SqlAccess::atomic_unit(&bridge, suspending_op).await;
975 assert!(
976 result.is_err(),
977 "a closure that suspends on first poll must fail loudly through \
978 `atomic_unit`, never silently succeed or wedge; got {result:?}"
979 );
980 }
981
982 #[tokio::test]
987 async fn concurrent_mutation_between_prepare_and_apply_trips_guard_and_rolls_back() {
988 let pool = scratch_pool("staleness_tripped");
989 seed_schema(&pool);
990 let a = Uuid::new_v4();
991 insert_entity(&pool, a, "alpha");
992
993 let plan = rename_plan(a, "alpha-v2", "rename-a");
995
996 {
999 let writer = pool.try_writer().expect("writer");
1000 writer
1001 .conn()
1002 .execute(
1003 "DELETE FROM entities WHERE id = ?1",
1004 rusqlite::params![a.to_string()],
1005 )
1006 .expect("concurrent delete");
1007 }
1008
1009 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1010 let outcome = run_atomic_unit(&bridge, vec![plan])
1011 .await
1012 .expect("seam call ok");
1013 match outcome {
1014 AtomicRunOutcome::RolledBack {
1015 failed_op_index,
1016 failure,
1017 } => {
1018 assert_eq!(failed_op_index, 0);
1019 assert!(matches!(
1020 failure,
1021 AtomicOpFailure::GuardFailed { observed: 0, .. }
1022 ));
1023 }
1024 other => panic!("expected RolledBack from the stale guard, got {other:?}"),
1025 }
1026 }
1027
1028 #[tokio::test]
1029 async fn no_mutation_twin_commits() {
1030 let pool = scratch_pool("staleness_untripped");
1031 seed_schema(&pool);
1032 let a = Uuid::new_v4();
1033 insert_entity(&pool, a, "alpha");
1034
1035 let plan = rename_plan(a, "alpha-v2", "rename-a");
1036 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1037 let outcome = run_atomic_unit(&bridge, vec![plan])
1038 .await
1039 .expect("seam call ok");
1040 assert!(matches!(outcome, AtomicRunOutcome::Committed { .. }));
1041 assert!(entity_exists(&pool, a));
1042 }
1043
1044 #[tokio::test]
1049 async fn dangling_edge_from_delete_then_link_rolls_back_whole_unit() {
1050 let pool = scratch_pool("dangling_edge");
1051 seed_schema(&pool);
1052 let x = Uuid::new_v4();
1053 let a = Uuid::new_v4();
1054 insert_entity(&pool, x, "x");
1055 insert_entity(&pool, a, "a");
1056 let before = entities_snapshot(&pool);
1057
1058 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1059 let edge_id = Uuid::new_v4();
1060 let plans = vec![delete_plan(x, "delete-x"), link_plan(edge_id, a, x)];
1061
1062 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1063 match outcome {
1064 AtomicRunOutcome::RolledBack {
1065 failed_op_index,
1066 failure,
1067 } => {
1068 assert_eq!(
1069 failed_op_index, 1,
1070 "delete(x) succeeds; link(a, x) is the op whose \
1071 endpoint-existence guard fails once x is gone"
1072 );
1073 assert!(matches!(
1074 failure,
1075 AtomicOpFailure::GuardFailed { observed: 0, .. }
1076 ));
1077 }
1078 other => panic!("expected RolledBack, got {other:?}"),
1079 }
1080
1081 assert_eq!(
1082 edge_count(&pool, a, x),
1083 0,
1084 "no dangling edge may be committed"
1085 );
1086 assert_eq!(
1087 entities_snapshot(&pool),
1088 before,
1089 "x must still exist — the whole unit rolled back, including the delete"
1090 );
1091 }
1092
1093 #[tokio::test]
1098 async fn merge_rewire_sees_earlier_in_file_edge_write() {
1099 let pool = scratch_pool("merge_rewire");
1100 seed_schema(&pool);
1101 let z = Uuid::new_v4();
1102 let from = Uuid::new_v4();
1103 let into = Uuid::new_v4();
1104 insert_entity(&pool, z, "z");
1105 insert_entity(&pool, from, "from-entity");
1106 insert_entity(&pool, into, "into-entity");
1107
1108 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1109 let edge_id = Uuid::new_v4();
1110 let plans = vec![link_plan(edge_id, from, z), merge_plan(into, from)];
1115
1116 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1117 assert!(
1118 matches!(outcome, AtomicRunOutcome::Committed { .. }),
1119 "expected the unit to commit: {outcome:?}"
1120 );
1121
1122 assert_eq!(
1123 edge_count(&pool, from, z),
1124 0,
1125 "no live edge may remain sourced from the merged-away entity"
1126 );
1127 assert_eq!(
1128 edge_count(&pool, into, z),
1129 1,
1130 "the edge must be rewired onto `into`"
1131 );
1132 assert!(
1133 !entity_exists(&pool, from),
1134 "the `from` entity must be tombstoned (soft-deleted)"
1135 );
1136 }
1137
1138 #[tokio::test]
1143 async fn zero_row_apply_fails_the_whole_unit() {
1144 let pool = scratch_pool("zero_row");
1145 seed_schema(&pool);
1146 let x = Uuid::new_v4();
1147 insert_entity(&pool, x, "x");
1148 let before = entities_snapshot(&pool);
1149
1150 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1151 let plans = vec![
1154 delete_plan(x, "delete-x"),
1155 rename_plan(x, "x-v2", "update-x"),
1156 ];
1157
1158 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1159 match outcome {
1160 AtomicRunOutcome::RolledBack {
1161 failed_op_index,
1162 failure,
1163 } => {
1164 assert_eq!(failed_op_index, 1);
1165 assert_eq!(
1166 failure,
1167 AtomicOpFailure::GuardFailed {
1168 statement_label: Some("update-x".to_string()),
1169 expected: AffectedRowGuard::exactly(1),
1170 observed: 0,
1171 }
1172 );
1173 }
1174 other => panic!("expected RolledBack, got {other:?}"),
1175 }
1176
1177 assert_eq!(
1178 entities_snapshot(&pool),
1179 before,
1180 "x must be restored — the whole unit rolled back"
1181 );
1182 }
1183
1184 #[tokio::test]
1205 async fn atomic_unit_and_concurrent_normal_write_share_the_same_writer_queue() {
1206 let pool = scratch_pool("daemon_coexistence");
1207 seed_schema(&pool);
1208 let a = Uuid::new_v4();
1209 let b = Uuid::new_v4();
1210 insert_entity(&pool, a, "alpha");
1211
1212 let writer_task = pool
1213 .writer_task_handle()
1214 .expect("writer task lookup")
1215 .expect("writer task must be spawned with the flag on for a file-backed pool");
1216
1217 let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>();
1218 let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
1219 let occupier = {
1220 let writer_task = writer_task.clone();
1221 tokio::spawn(async move {
1222 writer_task
1223 .send(move |_conn| {
1224 let _ = started_tx.send(());
1225 let _ = release_rx.blocking_recv();
1226 Ok::<(), StorageError>(())
1227 })
1228 .await
1229 })
1230 };
1231
1232 started_rx
1233 .await
1234 .expect("occupier must signal it has started running inside the writer task");
1235 assert_eq!(
1236 writer_task.queue_depth(),
1237 0,
1238 "channel must start empty once the occupier has been dequeued and is running"
1239 );
1240
1241 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1243 let atomic_plans = vec![rename_plan(a, "alpha-atomic", "rename-a-atomic")];
1244 let atomic_task = tokio::spawn(async move { run_atomic_unit(&bridge, atomic_plans).await });
1245
1246 let normal_write_task = {
1255 let writer_task = writer_task.clone();
1256 let b_str = b.to_string();
1257 tokio::spawn(async move {
1258 writer_task
1259 .send(move |conn| {
1260 conn.execute(
1261 "INSERT INTO entities \
1262 (id, namespace, kind, name, created_at, updated_at) \
1263 VALUES (?1, 'local', 'concept', ?2, 0, 0)",
1264 rusqlite::params![b_str, "bravo"],
1265 )
1266 .map_err(|e| StorageError::Internal(e.to_string()))
1267 })
1268 .await
1269 })
1270 };
1271
1272 let mut saw_both_enqueued = false;
1277 for _ in 0..200 {
1278 if writer_task.queue_depth() >= 2 {
1279 saw_both_enqueued = true;
1280 break;
1281 }
1282 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1283 }
1284 assert!(
1285 saw_both_enqueued,
1286 "expected BOTH the atomic unit's atomic_unit() call and the \
1287 concurrent ordinary write to appear in the writer task's queue \
1288 while the occupier held the single drain slot (queue_depth \
1289 should have reached 2) — observed {}; run_atomic_unit is not \
1290 sharing the same queue as an ordinary write",
1291 writer_task.queue_depth()
1292 );
1293
1294 release_tx.send(()).expect("release occupier");
1295 occupier
1296 .await
1297 .expect("occupier task join")
1298 .expect("occupier op ok");
1299
1300 let atomic_outcome = atomic_task
1301 .await
1302 .expect("atomic task join")
1303 .expect("seam call ok");
1304 assert!(
1305 matches!(atomic_outcome, AtomicRunOutcome::Committed { .. }),
1306 "atomic unit must commit once the occupier releases: {atomic_outcome:?}"
1307 );
1308 normal_write_task
1309 .await
1310 .expect("normal write task join")
1311 .expect("normal write op ok");
1312
1313 assert!(
1314 entity_exists(&pool, a),
1315 "a must exist (renamed) after commit"
1316 );
1317 {
1318 let writer = pool.try_writer().expect("writer");
1325 let renamed: String = writer
1326 .conn()
1327 .query_row(
1328 "SELECT name FROM entities WHERE id = ?1",
1329 rusqlite::params![a.to_string()],
1330 |row| row.get(0),
1331 )
1332 .expect("query renamed entity");
1333 assert_eq!(renamed, "alpha-atomic");
1334 }
1335 assert!(
1336 entity_exists(&pool, b),
1337 "the concurrent ordinary write must also have landed"
1338 );
1339 }
1340
1341 #[tokio::test]
1342 async fn post_commit_effects_are_collected_in_op_order_and_only_for_update() {
1343 let pool = scratch_pool("post_commit_collection");
1344 seed_schema(&pool);
1345 let a = Uuid::new_v4();
1346 let b = Uuid::new_v4();
1347 insert_entity(&pool, a, "alpha");
1348 insert_entity(&pool, b, "bravo");
1349
1350 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1351 let mut update_a = match rename_plan(a, "alpha-v2", "rename-a") {
1352 AtomicOpPlan::Update(p) => p,
1353 _ => unreachable!(),
1354 };
1355 update_a.post_commit = PostCommitEffect::ReindexEntity { entity_id: a };
1356 let plans = vec![
1357 AtomicOpPlan::Update(update_a),
1358 delete_plan(b, "delete-b"), ];
1360
1361 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1362 match outcome {
1363 AtomicRunOutcome::Committed { post_commit } => {
1364 fn require_commit_provenance(_: &CommittedPostCommitEffects) {}
1365 require_commit_provenance(&post_commit);
1366 assert_eq!(
1367 post_commit.as_slice(),
1368 &[PostCommitEffect::ReindexEntity { entity_id: a }]
1369 );
1370 }
1371 other => panic!("expected Committed, got {other:?}"),
1372 }
1373 }
1374
1375 mod acknowledgement {
1376 use super::*;
1377 use khive_storage::{SqlReader, WriterTaskRequestState};
1378
1379 #[derive(Clone, Copy, Debug)]
1380 enum Fault {
1381 Request(WriterTaskRequestState),
1382 Terminated(WriterTaskRequestState),
1383 }
1384
1385 impl Fault {
1386 fn error(self) -> StorageError {
1387 match self {
1388 Self::Request(request_state) => StorageError::WriterTaskRequestFailed {
1389 request_state,
1390 source: Box::new(StorageError::Internal(
1391 "lost atomic acknowledgement".into(),
1392 )),
1393 },
1394 Self::Terminated(request_state) => {
1395 StorageError::WriterTaskTerminated { request_state }
1396 }
1397 }
1398 }
1399
1400 fn assert_preserved(self, error: StorageError) {
1401 match (self, error) {
1402 (
1403 Self::Request(expected),
1404 StorageError::WriterTaskRequestFailed {
1405 request_state,
1406 source,
1407 },
1408 ) => {
1409 assert_eq!(request_state, expected);
1410 assert!(matches!(*source, StorageError::Internal(ref message)
1411 if message == "lost atomic acknowledgement"));
1412 }
1413 (
1414 Self::Terminated(expected),
1415 StorageError::WriterTaskTerminated { request_state },
1416 ) => assert_eq!(request_state, expected),
1417 (expected, actual) => panic!("expected {expected:?}, got {actual:?}"),
1418 }
1419 }
1420 }
1421
1422 struct FaultAccess {
1423 bridge: SqlBridge,
1424 fault: Fault,
1425 invoke: bool,
1426 callback_result: StdArc<Mutex<Option<bool>>>,
1427 }
1428
1429 #[async_trait::async_trait]
1430 impl SqlAccess for FaultAccess {
1431 async fn reader(&self) -> StorageResultAlias<Box<dyn SqlReader>> {
1432 self.bridge.reader().await
1433 }
1434
1435 async fn writer(&self) -> StorageResultAlias<Box<dyn SqlWriter>> {
1436 self.bridge.writer().await
1437 }
1438
1439 async fn atomic_unit(
1440 &self,
1441 op: AtomicUnitOp,
1442 ) -> StorageResultAlias<Box<dyn Any + Send>> {
1443 if self.invoke {
1444 let callback_result = StdArc::clone(&self.callback_result);
1445 let result = self
1446 .bridge
1447 .atomic_unit(Box::new(move |writer| {
1448 Box::pin(async move {
1449 let result = op(writer).await;
1450 *callback_result.lock().expect("callback result") =
1451 Some(result.is_ok());
1452 result
1453 })
1454 }))
1455 .await;
1456 assert_eq!(
1457 Some(result.is_ok()),
1458 *self.callback_result.lock().expect("callback result"),
1459 "the real transaction must finish before replacing its acknowledgement"
1460 );
1461 }
1462 Err(self.fault.error())
1463 }
1464 }
1465
1466 fn fault_access(pool: &TestPool, fault: Fault, invoke: bool) -> FaultAccess {
1467 FaultAccess {
1468 bridge: SqlBridge::new(StdArc::clone(pool), true),
1469 fault,
1470 invoke,
1471 callback_result: StdArc::new(Mutex::new(None)),
1472 }
1473 }
1474
1475 fn effect_plan(id: Uuid) -> AtomicOpPlan {
1476 let AtomicOpPlan::Update(mut plan) = rename_plan(id, "changed", "rename") else {
1477 unreachable!()
1478 };
1479 plan.post_commit = PostCommitEffect::ReindexEntity { entity_id: id };
1480 AtomicOpPlan::Update(plan)
1481 }
1482
1483 #[tokio::test]
1484 async fn recorded_failure_recovers_only_after_confirmed_rollback() {
1485 let pool = scratch_pool("recorded_rollback");
1486 seed_schema(&pool);
1487 let id = Uuid::new_v4();
1488 insert_entity(&pool, id, "original");
1489 let before = entities_snapshot(&pool);
1490 let access = fault_access(
1491 &pool,
1492 Fault::Request(WriterTaskRequestState::TransactionRolledBack),
1493 true,
1494 );
1495 let result = run_atomic_unit(
1496 &access,
1497 vec![
1498 effect_plan(id),
1499 rename_plan(Uuid::new_v4(), "missing", "missing-target"),
1500 ],
1501 )
1502 .await
1503 .expect("confirmed rollback recovers the recorded refusal");
1504 assert_eq!(
1505 result,
1506 AtomicRunOutcome::RolledBack {
1507 failed_op_index: 1,
1508 failure: AtomicOpFailure::GuardFailed {
1509 statement_label: Some("missing-target".into()),
1510 expected: AffectedRowGuard::exactly(1),
1511 observed: 0,
1512 },
1513 },
1514 );
1515 assert_eq!(*access.callback_result.lock().unwrap(), Some(false));
1516 assert_eq!(entities_snapshot(&pool), before);
1517 }
1518
1519 #[tokio::test]
1520 async fn recorded_failure_does_not_mask_unknown_or_terminal_outcomes() {
1521 for fault in [
1522 Fault::Request(WriterTaskRequestState::SideEffectsUnknown),
1523 Fault::Terminated(WriterTaskRequestState::SideEffectsUnknown),
1524 Fault::Terminated(WriterTaskRequestState::TransactionRolledBack),
1525 ] {
1526 let pool = scratch_pool("recorded_uncertain");
1527 seed_schema(&pool);
1528 let id = Uuid::new_v4();
1529 insert_entity(&pool, id, "original");
1530 let before = entities_snapshot(&pool);
1531 let access = fault_access(&pool, fault, true);
1532 let error = run_atomic_unit(
1533 &access,
1534 vec![
1535 effect_plan(id),
1536 rename_plan(Uuid::new_v4(), "missing", "missing-target"),
1537 ],
1538 )
1539 .await
1540 .expect_err("a recorded refusal cannot replace the outer outcome");
1541 fault.assert_preserved(error.0);
1542 assert_eq!(*access.callback_result.lock().unwrap(), Some(false));
1543 assert_eq!(entities_snapshot(&pool), before);
1544 }
1545 }
1546
1547 #[tokio::test]
1548 async fn failure_before_callback_cannot_manufacture_recorded_refusal() {
1549 for fault in [
1550 Fault::Request(WriterTaskRequestState::NotStarted),
1551 Fault::Request(WriterTaskRequestState::TransactionRolledBack),
1552 Fault::Terminated(WriterTaskRequestState::NotStarted),
1553 ] {
1554 let pool = scratch_pool("empty_failure_slot");
1555 seed_schema(&pool);
1556 let id = Uuid::new_v4();
1557 insert_entity(&pool, id, "original");
1558 let before = entities_snapshot(&pool);
1559 let access = fault_access(&pool, fault, false);
1560 let error = run_atomic_unit(&access, vec![effect_plan(id)])
1561 .await
1562 .expect_err("an empty slot must preserve the outer failure");
1563 fault.assert_preserved(error.0);
1564 assert_eq!(*access.callback_result.lock().unwrap(), None);
1565 assert_eq!(entities_snapshot(&pool), before);
1566 }
1567 }
1568
1569 #[tokio::test]
1570 async fn successful_callback_with_lost_acknowledgement_returns_no_commit_token() {
1571 for fault in [
1572 Fault::Request(WriterTaskRequestState::SideEffectsUnknown),
1573 Fault::Terminated(WriterTaskRequestState::SideEffectsUnknown),
1574 ] {
1575 let pool = scratch_pool("committed_ack_lost");
1576 seed_schema(&pool);
1577 let id = Uuid::new_v4();
1578 insert_entity(&pool, id, "original");
1579 let access = fault_access(&pool, fault, true);
1580 let error = run_atomic_unit(&access, vec![effect_plan(id)])
1581 .await
1582 .expect_err("callback success is not a committed-effects token");
1583 fault.assert_preserved(error.0);
1584 assert_eq!(*access.callback_result.lock().unwrap(), Some(true));
1585 let writer = pool.try_writer().expect("writer");
1586 let name: String = writer
1587 .conn()
1588 .query_row(
1589 "SELECT name FROM entities WHERE id=?1",
1590 rusqlite::params![id.to_string()],
1591 |row| row.get(0),
1592 )
1593 .expect("persisted write");
1594 assert_eq!(
1595 name, "changed",
1596 "the lost acknowledgement does not imply rollback"
1597 );
1598 }
1599 }
1600 }
1601}