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(
414 SqlStatement::new(
415 "SELECT version FROM notes WHERE id=?1",
416 vec![khive_storage::SqlValue::Text(plan.note_id.to_string())],
417 )
418 .labelled("note-write-guard"),
419 )
420 .await
421 .map_err(|error| AtomicOpFailure::SqlError {
422 statement_label: Some("note-version-receipt".into()),
423 message: error.to_string(),
424 })?;
425 let Some(khive_storage::SqlValue::Integer(version)) = current else {
426 return Err(AtomicOpFailure::SqlError {
427 statement_label: Some("note-version-receipt".into()),
428 message: "missing persisted note version".into(),
429 });
430 };
431 Some(version)
432 } else {
433 None
434 }
435 } else {
436 None
437 };
438 Ok(AppliedPlan {
439 effect: post_commit,
440 note_version,
441 })
442}
443
444pub async fn run_atomic_unit(
467 access: &dyn SqlAccess,
468 plans: Vec<AtomicOpPlan>,
469) -> Result<AtomicRunOutcome, AtomicRunnerError> {
470 run_atomic_unit_with_note_versions(access, plans, false)
471 .await
472 .map(|(outcome, _)| outcome)
473}
474
475pub(crate) async fn run_atomic_unit_with_note_versions(
478 access: &dyn SqlAccess,
479 plans: Vec<AtomicOpPlan>,
480 capture_note_versions: bool,
481) -> Result<(AtomicRunOutcome, Vec<i64>), AtomicRunnerError> {
482 let op: PreparedAtomicOp<Vec<i64>, (usize, AtomicOpFailure)> = Box::new(move |writer| {
483 Box::pin(async move {
484 let mut post_commit = Vec::new();
485 let mut note_versions = Vec::new();
486 for (op_index, plan) in plans.iter().enumerate() {
487 let savepoint = format!("adr099_atomic_op_{op_index}");
488 begin_savepoint(writer, &savepoint).await?;
489 match apply_plan(writer, plan, capture_note_versions).await {
490 Ok(applied) => {
491 release_savepoint(writer, &savepoint).await?;
492 if let Some(effect) = applied.effect {
493 post_commit.push(effect);
494 }
495 if let Some(version) = applied.note_version {
496 note_versions.push(version);
497 }
498 }
499 Err(failure) => {
500 let _ = rollback_to_savepoint(writer, &savepoint).await;
508 let _ = release_savepoint(writer, &savepoint).await;
509 return Err(PreparedAtomicError::Refused {
510 failure: (op_index, failure),
511 message: format!("ADR-099 atomic unit aborted at op {op_index}"),
512 });
513 }
514 }
515 }
516 Ok((note_versions, post_commit))
517 })
518 });
519
520 match run_prepared_atomic_unit(access, op)
521 .await
522 .map_err(AtomicRunnerError)?
523 {
524 PreparedAtomicOutcome::Committed { value, post_commit } => {
525 Ok((AtomicRunOutcome::Committed { post_commit }, value))
526 }
527 PreparedAtomicOutcome::RolledBack((failed_op_index, failure)) => Ok((
528 AtomicRunOutcome::RolledBack {
529 failed_op_index,
530 failure,
531 },
532 Vec::new(),
533 )),
534 }
535}
536
537pub(crate) enum PreparedAtomicError<E> {
538 Refused { failure: E, message: String },
539 Storage(StorageError),
540}
541
542impl<E> From<StorageError> for PreparedAtomicError<E> {
543 fn from(error: StorageError) -> Self {
544 Self::Storage(error)
545 }
546}
547
548pub(crate) enum PreparedAtomicOutcome<T, E> {
549 Committed {
550 value: T,
551 post_commit: CommittedPostCommitEffects,
552 },
553 RolledBack(E),
554}
555
556pub(crate) type PreparedAtomicOp<T, E> = Box<
557 dyn for<'w> FnOnce(
558 &'w mut dyn SqlWriter,
559 )
560 -> BoxFuture<'w, Result<(T, Vec<PostCommitEffect>), PreparedAtomicError<E>>>
561 + Send,
562>;
563
564pub(crate) async fn run_prepared_atomic_unit<T: Send + 'static, E: Send + 'static>(
567 access: &dyn SqlAccess,
568 prepared: PreparedAtomicOp<T, E>,
569) -> Result<PreparedAtomicOutcome<T, E>, StorageError> {
570 let failure_slot = Arc::new(Mutex::new(None));
571 let recorded = Arc::clone(&failure_slot);
572 let op: AtomicUnitOp = Box::new(move |writer| {
573 Box::pin(async move {
574 match prepared(writer).await {
575 Ok(value) => Ok(Box::new(value) as Box<dyn Any + Send>),
576 Err(PreparedAtomicError::Storage(error)) => Err(error),
577 Err(PreparedAtomicError::Refused { failure, message }) => {
578 *recorded
579 .lock()
580 .expect("atomic runner failure slot poisoned") = Some(failure);
581 Err(StorageError::Internal(message))
582 }
583 }
584 })
585 });
586 match access.atomic_unit(op).await {
587 Ok(boxed) => {
588 let (value, effects) = *boxed
589 .downcast::<(T, Vec<PostCommitEffect>)>()
590 .map_err(|_| StorageError::Internal("invalid prepared atomic outcome".into()))?;
591 Ok(PreparedAtomicOutcome::Committed {
592 value,
593 post_commit: CommittedPostCommitEffects::new(effects),
594 })
595 }
596 Err(storage_err) => {
597 if !atomic_unit_error_allows_recorded_refusal(&storage_err) {
599 return Err(storage_err);
600 }
601 let recorded = failure_slot
602 .lock()
603 .expect("atomic runner failure slot poisoned")
604 .take();
605 match recorded {
606 Some(failure) => Ok(PreparedAtomicOutcome::RolledBack(failure)),
607 None => Err(storage_err),
608 }
609 }
610 }
611}
612
613pub(crate) fn atomic_unit_error_allows_recorded_refusal(error: &StorageError) -> bool {
614 match error {
615 StorageError::WriterTaskTerminated { .. } => false,
616 StorageError::WriterTaskRequestFailed { request_state, .. } => {
617 *request_state == khive_storage::WriterTaskRequestState::TransactionRolledBack
618 }
619 _ => true,
620 }
621}
622
623#[cfg(test)]
624mod tests {
625 use super::*;
626
627 use std::sync::Arc as StdArc;
628
629 use khive_db::{ConnectionPool, PoolConfig, SqlBridge};
630 use khive_storage::types::{SqlValue, StorageResult as StorageResultAlias};
631 use uuid::Uuid;
632
633 struct TestPool {
635 pool: StdArc<ConnectionPool>,
636 _temp_dir: tempfile::TempDir,
637 }
638
639 impl std::ops::Deref for TestPool {
640 type Target = StdArc<ConnectionPool>;
641
642 fn deref(&self) -> &Self::Target {
643 &self.pool
644 }
645 }
646
647 fn scratch_pool(name: &str) -> TestPool {
654 let dir = tempfile::tempdir().expect("tempdir");
655 let path = dir.path().join(format!("{name}.db"));
656 let pool = StdArc::new(
657 ConnectionPool::new(PoolConfig {
658 path: Some(path),
659 write_queue_enabled: Some(true),
660 ..PoolConfig::for_test()
661 })
662 .expect("pool open"),
663 );
664 TestPool {
665 pool,
666 _temp_dir: dir,
667 }
668 }
669
670 fn seed_schema(pool: &ConnectionPool) {
678 let writer = pool.try_writer().expect("writer");
679 writer
680 .conn()
681 .execute_batch(
682 "CREATE TABLE entities (
683 id TEXT PRIMARY KEY,
684 namespace TEXT NOT NULL,
685 kind TEXT NOT NULL,
686 entity_type TEXT,
687 name TEXT NOT NULL,
688 description TEXT,
689 properties TEXT,
690 tags TEXT NOT NULL DEFAULT '[]',
691 created_at INTEGER NOT NULL,
692 updated_at INTEGER NOT NULL,
693 version INTEGER NOT NULL DEFAULT 1,
694 deleted_at INTEGER,
695 merged_into TEXT,
696 merge_event_id TEXT
697 );
698 CREATE TABLE graph_edges (
699 namespace TEXT NOT NULL,
700 id TEXT NOT NULL,
701 source_id TEXT NOT NULL,
702 target_id TEXT NOT NULL,
703 relation TEXT NOT NULL,
704 weight REAL NOT NULL DEFAULT 1.0,
705 created_at INTEGER NOT NULL,
706 updated_at INTEGER NOT NULL,
707 deleted_at INTEGER,
708 metadata TEXT,
709 target_backend TEXT,
710 PRIMARY KEY (namespace, id)
711 );",
712 )
713 .expect("seed schema");
714 }
715
716 fn insert_entity(pool: &ConnectionPool, id: Uuid, name: &str) {
717 let writer = pool.try_writer().expect("writer");
718 writer
719 .conn()
720 .execute(
721 "INSERT INTO entities \
722 (id, namespace, kind, name, created_at, updated_at) \
723 VALUES (?1, 'local', 'concept', ?2, 0, 0)",
724 rusqlite::params![id.to_string(), name],
725 )
726 .expect("insert entity");
727 }
728
729 fn entity_exists(pool: &ConnectionPool, id: Uuid) -> bool {
730 let writer = pool.try_writer().expect("writer");
731 let count: i64 = writer
732 .conn()
733 .query_row(
734 "SELECT COUNT(*) FROM entities WHERE id = ?1 AND deleted_at IS NULL",
735 rusqlite::params![id.to_string()],
736 |row| row.get(0),
737 )
738 .expect("query entity");
739 count > 0
740 }
741
742 fn edge_count(pool: &ConnectionPool, source: Uuid, target: Uuid) -> i64 {
743 let writer = pool.try_writer().expect("writer");
744 writer
745 .conn()
746 .query_row(
747 "SELECT COUNT(*) FROM graph_edges \
748 WHERE source_id = ?1 AND target_id = ?2 AND deleted_at IS NULL",
749 rusqlite::params![source.to_string(), target.to_string()],
750 |row| row.get(0),
751 )
752 .expect("query edge")
753 }
754
755 fn entities_snapshot(pool: &ConnectionPool) -> Vec<String> {
756 let writer = pool.try_writer().expect("writer");
757 let mut stmt = writer
758 .conn()
759 .prepare("SELECT id, name, deleted_at FROM entities ORDER BY id")
760 .expect("prepare snapshot");
761 let rows = stmt
762 .query_map([], |row| {
763 let id: String = row.get(0)?;
764 let name: String = row.get(1)?;
765 let deleted_at: Option<i64> = row.get(2)?;
766 Ok(format!("{id}:{name}:{deleted_at:?}"))
767 })
768 .expect("query snapshot");
769 rows.collect::<Result<Vec<_>, _>>()
770 .expect("collect snapshot")
771 }
772
773 fn delete_plan(id: Uuid, label: &str) -> AtomicOpPlan {
774 AtomicOpPlan::Delete(DeletePlan {
775 target_id: id,
776 statements: vec![PlanStatement {
777 statement: SqlStatement {
778 sql: "DELETE FROM entities WHERE id = ?1 AND deleted_at IS NULL".to_string(),
779 params: vec![SqlValue::Text(id.to_string())],
780 label: Some(label.to_string()),
781 },
782 guard: Some(AffectedRowGuard::exactly(1)),
783 }],
784 post_commit: PostCommitEffect::None,
785 })
786 }
787
788 fn rename_plan(id: Uuid, new_name: &str, label: &str) -> AtomicOpPlan {
789 AtomicOpPlan::Update(Box::new(UpdatePlan {
790 graph_effects: Vec::new(),
791 note_vector_purge: None,
792 note_embedding_inheritance: None,
793 entity_guard: None,
794 note_guard: None,
795 target_id: id,
796 statements: vec![PlanStatement {
797 statement: SqlStatement {
798 sql: "UPDATE entities SET version = version + 1, name = ?1, updated_at = 1 \
799 WHERE id = ?2 AND deleted_at IS NULL"
800 .to_string(),
801 params: vec![
802 SqlValue::Text(new_name.to_string()),
803 SqlValue::Text(id.to_string()),
804 ],
805 label: Some(label.to_string()),
806 },
807 guard: Some(AffectedRowGuard::exactly(1)),
808 }],
809 post_commit: PostCommitEffect::None,
810 edge_natural_key: None,
811 idempotent_noop: false,
812 }))
813 }
814
815 fn link_plan(edge_id: Uuid, source: Uuid, target: Uuid) -> AtomicOpPlan {
820 AtomicOpPlan::Link(LinkPlan {
821 source_id: source,
822 target_id: target,
823 statements: vec![PlanStatement {
824 statement: SqlStatement {
825 sql: "INSERT INTO graph_edges \
826 (namespace, id, source_id, target_id, relation, created_at, updated_at) \
827 SELECT 'local', ?1, ?2, ?3, 'annotates', 0, 0 \
828 WHERE EXISTS (SELECT 1 FROM entities WHERE id = ?2 AND deleted_at IS NULL) \
829 AND EXISTS (SELECT 1 FROM entities WHERE id = ?3 AND deleted_at IS NULL)"
830 .to_string(),
831 params: vec![
832 SqlValue::Text(edge_id.to_string()),
833 SqlValue::Text(source.to_string()),
834 SqlValue::Text(target.to_string()),
835 ],
836 label: Some("insert-edge-where-exists".to_string()),
837 },
838 guard: Some(AffectedRowGuard::exactly(1)),
839 }],
840 disposition: khive_storage::EdgeUpsertDisposition::Created,
841 })
842 }
843
844 fn merge_plan(into_id: Uuid, from_id: Uuid) -> AtomicOpPlan {
845 AtomicOpPlan::Merge(MergePlan {
846 into_id,
847 from_id,
848 rewires: vec![crate::atomic_plan::PlanPredicate {
849 description: "source_id = :from".to_string(),
850 statement: SqlStatement {
851 sql: "UPDATE graph_edges SET source_id = ?1, updated_at = 1 \
852 WHERE source_id = ?2"
853 .to_string(),
854 params: vec![
855 SqlValue::Text(into_id.to_string()),
856 SqlValue::Text(from_id.to_string()),
857 ],
858 label: Some("merge-rewire".to_string()),
859 },
860 }],
861 lifecycle: vec![PlanStatement {
862 statement: SqlStatement {
863 sql: "UPDATE entities SET version = version + 1, deleted_at = 1, merged_into = ?1 \
864 WHERE id = ?2 AND deleted_at IS NULL"
865 .to_string(),
866 params: vec![
867 SqlValue::Text(into_id.to_string()),
868 SqlValue::Text(from_id.to_string()),
869 ],
870 label: Some("tombstone-from-entity".to_string()),
871 },
872 guard: Some(AffectedRowGuard::exactly(1)),
873 }],
874 })
875 }
876
877 #[tokio::test]
882 async fn rollback_end_to_end_leaves_zero_partial_state() {
883 let pool = scratch_pool("rollback_end_to_end");
884 seed_schema(&pool);
885 let a = Uuid::new_v4();
886 let b = Uuid::new_v4();
887 insert_entity(&pool, a, "alpha");
888 insert_entity(&pool, b, "bravo");
889 let before = entities_snapshot(&pool);
890
891 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
892 let plans = vec![
893 rename_plan(a, "alpha-renamed", "rename-a"),
894 delete_plan(Uuid::new_v4(), "delete-nonexistent"),
897 rename_plan(b, "bravo-renamed", "rename-b"),
898 ];
899
900 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
901 match outcome {
902 AtomicRunOutcome::RolledBack {
903 failed_op_index,
904 failure,
905 } => {
906 assert_eq!(failed_op_index, 1);
907 assert!(matches!(failure, AtomicOpFailure::GuardFailed { .. }));
908 }
909 other => panic!("expected RolledBack, got {other:?}"),
910 }
911
912 let after = entities_snapshot(&pool);
913 assert_eq!(
914 before, after,
915 "a mid-unit failure must leave the database byte-for-byte unchanged"
916 );
917 }
918
919 #[tokio::test]
926 async fn commit_pass_resolves_on_first_poll_and_commits() {
927 let pool = scratch_pool("suspend_trap_happy_path");
928 seed_schema(&pool);
929 let a = Uuid::new_v4();
930 insert_entity(&pool, a, "alpha");
931
932 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
933 let plans = vec![rename_plan(a, "alpha-v2", "rename-a")];
934
935 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
936 assert!(
937 matches!(outcome, AtomicRunOutcome::Committed { .. }),
938 "the real commit-pass closure (SAVEPOINT + guarded DML only) must \
939 resolve on block_on_sync's first poll and commit: {outcome:?}"
940 );
941
942 let writer = pool.try_writer().expect("writer");
943 let name: String = writer
944 .conn()
945 .query_row(
946 "SELECT name FROM entities WHERE id = ?1",
947 rusqlite::params![a.to_string()],
948 |row| row.get(0),
949 )
950 .expect("query renamed entity");
951 assert_eq!(name, "alpha-v2");
952 }
953
954 #[tokio::test]
955 async fn hand_built_suspending_closure_fails_loudly_through_the_same_seam() {
956 let pool = scratch_pool("suspend_trap_misuse");
966 seed_schema(&pool);
967
968 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
969 let suspending_op: AtomicUnitOp = Box::new(|_writer| {
970 Box::pin(async move {
971 tokio::task::yield_now().await;
972 Ok(Box::new(()) as Box<dyn Any + Send>)
973 })
974 });
975
976 let result: StorageResultAlias<Box<dyn Any + Send>> =
977 khive_storage::SqlAccess::atomic_unit(&bridge, suspending_op).await;
978 assert!(
979 result.is_err(),
980 "a closure that suspends on first poll must fail loudly through \
981 `atomic_unit`, never silently succeed or wedge; got {result:?}"
982 );
983 }
984
985 #[tokio::test]
990 async fn concurrent_mutation_between_prepare_and_apply_trips_guard_and_rolls_back() {
991 let pool = scratch_pool("staleness_tripped");
992 seed_schema(&pool);
993 let a = Uuid::new_v4();
994 insert_entity(&pool, a, "alpha");
995
996 let plan = rename_plan(a, "alpha-v2", "rename-a");
998
999 {
1002 let writer = pool.try_writer().expect("writer");
1003 writer
1004 .conn()
1005 .execute(
1006 "DELETE FROM entities WHERE id = ?1",
1007 rusqlite::params![a.to_string()],
1008 )
1009 .expect("concurrent delete");
1010 }
1011
1012 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1013 let outcome = run_atomic_unit(&bridge, vec![plan])
1014 .await
1015 .expect("seam call ok");
1016 match outcome {
1017 AtomicRunOutcome::RolledBack {
1018 failed_op_index,
1019 failure,
1020 } => {
1021 assert_eq!(failed_op_index, 0);
1022 assert!(matches!(
1023 failure,
1024 AtomicOpFailure::GuardFailed { observed: 0, .. }
1025 ));
1026 }
1027 other => panic!("expected RolledBack from the stale guard, got {other:?}"),
1028 }
1029 }
1030
1031 #[tokio::test]
1032 async fn no_mutation_twin_commits() {
1033 let pool = scratch_pool("staleness_untripped");
1034 seed_schema(&pool);
1035 let a = Uuid::new_v4();
1036 insert_entity(&pool, a, "alpha");
1037
1038 let plan = rename_plan(a, "alpha-v2", "rename-a");
1039 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1040 let outcome = run_atomic_unit(&bridge, vec![plan])
1041 .await
1042 .expect("seam call ok");
1043 assert!(matches!(outcome, AtomicRunOutcome::Committed { .. }));
1044 assert!(entity_exists(&pool, a));
1045 }
1046
1047 #[tokio::test]
1052 async fn dangling_edge_from_delete_then_link_rolls_back_whole_unit() {
1053 let pool = scratch_pool("dangling_edge");
1054 seed_schema(&pool);
1055 let x = Uuid::new_v4();
1056 let a = Uuid::new_v4();
1057 insert_entity(&pool, x, "x");
1058 insert_entity(&pool, a, "a");
1059 let before = entities_snapshot(&pool);
1060
1061 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1062 let edge_id = Uuid::new_v4();
1063 let plans = vec![delete_plan(x, "delete-x"), link_plan(edge_id, a, x)];
1064
1065 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1066 match outcome {
1067 AtomicRunOutcome::RolledBack {
1068 failed_op_index,
1069 failure,
1070 } => {
1071 assert_eq!(
1072 failed_op_index, 1,
1073 "delete(x) succeeds; link(a, x) is the op whose \
1074 endpoint-existence guard fails once x is gone"
1075 );
1076 assert!(matches!(
1077 failure,
1078 AtomicOpFailure::GuardFailed { observed: 0, .. }
1079 ));
1080 }
1081 other => panic!("expected RolledBack, got {other:?}"),
1082 }
1083
1084 assert_eq!(
1085 edge_count(&pool, a, x),
1086 0,
1087 "no dangling edge may be committed"
1088 );
1089 assert_eq!(
1090 entities_snapshot(&pool),
1091 before,
1092 "x must still exist — the whole unit rolled back, including the delete"
1093 );
1094 }
1095
1096 #[tokio::test]
1101 async fn merge_rewire_sees_earlier_in_file_edge_write() {
1102 let pool = scratch_pool("merge_rewire");
1103 seed_schema(&pool);
1104 let z = Uuid::new_v4();
1105 let from = Uuid::new_v4();
1106 let into = Uuid::new_v4();
1107 insert_entity(&pool, z, "z");
1108 insert_entity(&pool, from, "from-entity");
1109 insert_entity(&pool, into, "into-entity");
1110
1111 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1112 let edge_id = Uuid::new_v4();
1113 let plans = vec![link_plan(edge_id, from, z), merge_plan(into, from)];
1118
1119 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1120 assert!(
1121 matches!(outcome, AtomicRunOutcome::Committed { .. }),
1122 "expected the unit to commit: {outcome:?}"
1123 );
1124
1125 assert_eq!(
1126 edge_count(&pool, from, z),
1127 0,
1128 "no live edge may remain sourced from the merged-away entity"
1129 );
1130 assert_eq!(
1131 edge_count(&pool, into, z),
1132 1,
1133 "the edge must be rewired onto `into`"
1134 );
1135 assert!(
1136 !entity_exists(&pool, from),
1137 "the `from` entity must be tombstoned (soft-deleted)"
1138 );
1139 }
1140
1141 #[tokio::test]
1146 async fn zero_row_apply_fails_the_whole_unit() {
1147 let pool = scratch_pool("zero_row");
1148 seed_schema(&pool);
1149 let x = Uuid::new_v4();
1150 insert_entity(&pool, x, "x");
1151 let before = entities_snapshot(&pool);
1152
1153 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1154 let plans = vec![
1157 delete_plan(x, "delete-x"),
1158 rename_plan(x, "x-v2", "update-x"),
1159 ];
1160
1161 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1162 match outcome {
1163 AtomicRunOutcome::RolledBack {
1164 failed_op_index,
1165 failure,
1166 } => {
1167 assert_eq!(failed_op_index, 1);
1168 assert_eq!(
1169 failure,
1170 AtomicOpFailure::GuardFailed {
1171 statement_label: Some("update-x".to_string()),
1172 expected: AffectedRowGuard::exactly(1),
1173 observed: 0,
1174 }
1175 );
1176 }
1177 other => panic!("expected RolledBack, got {other:?}"),
1178 }
1179
1180 assert_eq!(
1181 entities_snapshot(&pool),
1182 before,
1183 "x must be restored — the whole unit rolled back"
1184 );
1185 }
1186
1187 #[tokio::test]
1208 async fn atomic_unit_and_concurrent_normal_write_share_the_same_writer_queue() {
1209 let pool = scratch_pool("daemon_coexistence");
1210 seed_schema(&pool);
1211 let a = Uuid::new_v4();
1212 let b = Uuid::new_v4();
1213 insert_entity(&pool, a, "alpha");
1214
1215 let writer_task = pool
1216 .writer_task_handle()
1217 .expect("writer task lookup")
1218 .expect("writer task must be spawned with the flag on for a file-backed pool");
1219
1220 let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>();
1221 let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
1222 let occupier = {
1223 let writer_task = writer_task.clone();
1224 tokio::spawn(async move {
1225 writer_task
1226 .send(move |_conn| {
1227 let _ = started_tx.send(());
1228 let _ = release_rx.blocking_recv();
1229 Ok::<(), StorageError>(())
1230 })
1231 .await
1232 })
1233 };
1234
1235 started_rx
1236 .await
1237 .expect("occupier must signal it has started running inside the writer task");
1238 assert_eq!(
1239 writer_task.queue_depth(),
1240 0,
1241 "channel must start empty once the occupier has been dequeued and is running"
1242 );
1243
1244 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1246 let atomic_plans = vec![rename_plan(a, "alpha-atomic", "rename-a-atomic")];
1247 let atomic_task = tokio::spawn(async move { run_atomic_unit(&bridge, atomic_plans).await });
1248
1249 let normal_write_task = {
1258 let writer_task = writer_task.clone();
1259 let b_str = b.to_string();
1260 tokio::spawn(async move {
1261 writer_task
1262 .send(move |conn| {
1263 conn.execute(
1264 "INSERT INTO entities \
1265 (id, namespace, kind, name, created_at, updated_at) \
1266 VALUES (?1, 'local', 'concept', ?2, 0, 0)",
1267 rusqlite::params![b_str, "bravo"],
1268 )
1269 .map_err(|e| StorageError::Internal(e.to_string()))
1270 })
1271 .await
1272 })
1273 };
1274
1275 let mut saw_both_enqueued = false;
1280 for _ in 0..200 {
1281 if writer_task.queue_depth() >= 2 {
1282 saw_both_enqueued = true;
1283 break;
1284 }
1285 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1286 }
1287 assert!(
1288 saw_both_enqueued,
1289 "expected BOTH the atomic unit's atomic_unit() call and the \
1290 concurrent ordinary write to appear in the writer task's queue \
1291 while the occupier held the single drain slot (queue_depth \
1292 should have reached 2) — observed {}; run_atomic_unit is not \
1293 sharing the same queue as an ordinary write",
1294 writer_task.queue_depth()
1295 );
1296
1297 release_tx.send(()).expect("release occupier");
1298 occupier
1299 .await
1300 .expect("occupier task join")
1301 .expect("occupier op ok");
1302
1303 let atomic_outcome = atomic_task
1304 .await
1305 .expect("atomic task join")
1306 .expect("seam call ok");
1307 assert!(
1308 matches!(atomic_outcome, AtomicRunOutcome::Committed { .. }),
1309 "atomic unit must commit once the occupier releases: {atomic_outcome:?}"
1310 );
1311 normal_write_task
1312 .await
1313 .expect("normal write task join")
1314 .expect("normal write op ok");
1315
1316 assert!(
1317 entity_exists(&pool, a),
1318 "a must exist (renamed) after commit"
1319 );
1320 {
1321 let writer = pool.try_writer().expect("writer");
1328 let renamed: String = writer
1329 .conn()
1330 .query_row(
1331 "SELECT name FROM entities WHERE id = ?1",
1332 rusqlite::params![a.to_string()],
1333 |row| row.get(0),
1334 )
1335 .expect("query renamed entity");
1336 assert_eq!(renamed, "alpha-atomic");
1337 }
1338 assert!(
1339 entity_exists(&pool, b),
1340 "the concurrent ordinary write must also have landed"
1341 );
1342 }
1343
1344 #[tokio::test]
1345 async fn post_commit_effects_are_collected_in_op_order_and_only_for_update() {
1346 let pool = scratch_pool("post_commit_collection");
1347 seed_schema(&pool);
1348 let a = Uuid::new_v4();
1349 let b = Uuid::new_v4();
1350 insert_entity(&pool, a, "alpha");
1351 insert_entity(&pool, b, "bravo");
1352
1353 let bridge = SqlBridge::new(StdArc::clone(&pool), true);
1354 let mut update_a = match rename_plan(a, "alpha-v2", "rename-a") {
1355 AtomicOpPlan::Update(p) => p,
1356 _ => unreachable!(),
1357 };
1358 update_a.post_commit = PostCommitEffect::ReindexEntity { entity_id: a };
1359 let plans = vec![
1360 AtomicOpPlan::Update(update_a),
1361 delete_plan(b, "delete-b"), ];
1363
1364 let outcome = run_atomic_unit(&bridge, plans).await.expect("seam call ok");
1365 match outcome {
1366 AtomicRunOutcome::Committed { post_commit } => {
1367 fn require_commit_provenance(_: &CommittedPostCommitEffects) {}
1368 require_commit_provenance(&post_commit);
1369 assert_eq!(
1370 post_commit.as_slice(),
1371 &[PostCommitEffect::ReindexEntity { entity_id: a }]
1372 );
1373 }
1374 other => panic!("expected Committed, got {other:?}"),
1375 }
1376 }
1377
1378 mod acknowledgement {
1379 use super::*;
1380 use khive_storage::{SqlReader, WriterTaskRequestState};
1381
1382 #[derive(Clone, Copy, Debug)]
1383 enum Fault {
1384 Request(WriterTaskRequestState),
1385 Terminated(WriterTaskRequestState),
1386 }
1387
1388 impl Fault {
1389 fn error(self) -> StorageError {
1390 match self {
1391 Self::Request(request_state) => StorageError::WriterTaskRequestFailed {
1392 request_state,
1393 source: Box::new(StorageError::Internal(
1394 "lost atomic acknowledgement".into(),
1395 )),
1396 },
1397 Self::Terminated(request_state) => {
1398 StorageError::writer_task_terminated(request_state)
1399 }
1400 }
1401 }
1402
1403 fn assert_preserved(self, error: StorageError) {
1404 match (self, error) {
1405 (
1406 Self::Request(expected),
1407 StorageError::WriterTaskRequestFailed {
1408 request_state,
1409 source,
1410 },
1411 ) => {
1412 assert_eq!(request_state, expected);
1413 assert!(matches!(*source, StorageError::Internal(ref message)
1414 if message == "lost atomic acknowledgement"));
1415 }
1416 (
1417 Self::Terminated(expected),
1418 StorageError::WriterTaskTerminated { request_state, .. },
1419 ) => assert_eq!(request_state, expected),
1420 (expected, actual) => panic!("expected {expected:?}, got {actual:?}"),
1421 }
1422 }
1423 }
1424
1425 struct FaultAccess {
1426 bridge: SqlBridge,
1427 fault: Fault,
1428 invoke: bool,
1429 callback_result: StdArc<Mutex<Option<bool>>>,
1430 }
1431
1432 #[async_trait::async_trait]
1433 impl SqlAccess for FaultAccess {
1434 async fn reader(&self) -> StorageResultAlias<Box<dyn SqlReader>> {
1435 self.bridge.reader().await
1436 }
1437
1438 async fn writer(&self) -> StorageResultAlias<Box<dyn SqlWriter>> {
1439 self.bridge.writer().await
1440 }
1441
1442 async fn atomic_unit(
1443 &self,
1444 op: AtomicUnitOp,
1445 ) -> StorageResultAlias<Box<dyn Any + Send>> {
1446 if self.invoke {
1447 let callback_result = StdArc::clone(&self.callback_result);
1448 let result = self
1449 .bridge
1450 .atomic_unit(Box::new(move |writer| {
1451 Box::pin(async move {
1452 let result = op(writer).await;
1453 *callback_result.lock().expect("callback result") =
1454 Some(result.is_ok());
1455 result
1456 })
1457 }))
1458 .await;
1459 assert_eq!(
1460 Some(result.is_ok()),
1461 *self.callback_result.lock().expect("callback result"),
1462 "the real transaction must finish before replacing its acknowledgement"
1463 );
1464 }
1465 Err(self.fault.error())
1466 }
1467 }
1468
1469 fn fault_access(pool: &TestPool, fault: Fault, invoke: bool) -> FaultAccess {
1470 FaultAccess {
1471 bridge: SqlBridge::new(StdArc::clone(pool), true),
1472 fault,
1473 invoke,
1474 callback_result: StdArc::new(Mutex::new(None)),
1475 }
1476 }
1477
1478 fn effect_plan(id: Uuid) -> AtomicOpPlan {
1479 let AtomicOpPlan::Update(mut plan) = rename_plan(id, "changed", "rename") else {
1480 unreachable!()
1481 };
1482 plan.post_commit = PostCommitEffect::ReindexEntity { entity_id: id };
1483 AtomicOpPlan::Update(plan)
1484 }
1485
1486 #[tokio::test]
1487 async fn recorded_failure_recovers_only_after_confirmed_rollback() {
1488 let pool = scratch_pool("recorded_rollback");
1489 seed_schema(&pool);
1490 let id = Uuid::new_v4();
1491 insert_entity(&pool, id, "original");
1492 let before = entities_snapshot(&pool);
1493 let access = fault_access(
1494 &pool,
1495 Fault::Request(WriterTaskRequestState::TransactionRolledBack),
1496 true,
1497 );
1498 let result = run_atomic_unit(
1499 &access,
1500 vec![
1501 effect_plan(id),
1502 rename_plan(Uuid::new_v4(), "missing", "missing-target"),
1503 ],
1504 )
1505 .await
1506 .expect("confirmed rollback recovers the recorded refusal");
1507 assert_eq!(
1508 result,
1509 AtomicRunOutcome::RolledBack {
1510 failed_op_index: 1,
1511 failure: AtomicOpFailure::GuardFailed {
1512 statement_label: Some("missing-target".into()),
1513 expected: AffectedRowGuard::exactly(1),
1514 observed: 0,
1515 },
1516 },
1517 );
1518 assert_eq!(*access.callback_result.lock().unwrap(), Some(false));
1519 assert_eq!(entities_snapshot(&pool), before);
1520 }
1521
1522 #[tokio::test]
1523 async fn recorded_failure_does_not_mask_unknown_or_terminal_outcomes() {
1524 for fault in [
1525 Fault::Request(WriterTaskRequestState::SideEffectsUnknown),
1526 Fault::Terminated(WriterTaskRequestState::SideEffectsUnknown),
1527 Fault::Terminated(WriterTaskRequestState::TransactionRolledBack),
1528 ] {
1529 let pool = scratch_pool("recorded_uncertain");
1530 seed_schema(&pool);
1531 let id = Uuid::new_v4();
1532 insert_entity(&pool, id, "original");
1533 let before = entities_snapshot(&pool);
1534 let access = fault_access(&pool, fault, true);
1535 let error = run_atomic_unit(
1536 &access,
1537 vec![
1538 effect_plan(id),
1539 rename_plan(Uuid::new_v4(), "missing", "missing-target"),
1540 ],
1541 )
1542 .await
1543 .expect_err("a recorded refusal cannot replace the outer outcome");
1544 fault.assert_preserved(error.0);
1545 assert_eq!(*access.callback_result.lock().unwrap(), Some(false));
1546 assert_eq!(entities_snapshot(&pool), before);
1547 }
1548 }
1549
1550 #[tokio::test]
1551 async fn failure_before_callback_cannot_manufacture_recorded_refusal() {
1552 for fault in [
1553 Fault::Request(WriterTaskRequestState::NotStarted),
1554 Fault::Request(WriterTaskRequestState::TransactionRolledBack),
1555 Fault::Terminated(WriterTaskRequestState::NotStarted),
1556 ] {
1557 let pool = scratch_pool("empty_failure_slot");
1558 seed_schema(&pool);
1559 let id = Uuid::new_v4();
1560 insert_entity(&pool, id, "original");
1561 let before = entities_snapshot(&pool);
1562 let access = fault_access(&pool, fault, false);
1563 let error = run_atomic_unit(&access, vec![effect_plan(id)])
1564 .await
1565 .expect_err("an empty slot must preserve the outer failure");
1566 fault.assert_preserved(error.0);
1567 assert_eq!(*access.callback_result.lock().unwrap(), None);
1568 assert_eq!(entities_snapshot(&pool), before);
1569 }
1570 }
1571
1572 #[tokio::test]
1573 async fn successful_callback_with_lost_acknowledgement_returns_no_commit_token() {
1574 for fault in [
1575 Fault::Request(WriterTaskRequestState::SideEffectsUnknown),
1576 Fault::Terminated(WriterTaskRequestState::SideEffectsUnknown),
1577 ] {
1578 let pool = scratch_pool("committed_ack_lost");
1579 seed_schema(&pool);
1580 let id = Uuid::new_v4();
1581 insert_entity(&pool, id, "original");
1582 let access = fault_access(&pool, fault, true);
1583 let error = run_atomic_unit(&access, vec![effect_plan(id)])
1584 .await
1585 .expect_err("callback success is not a committed-effects token");
1586 fault.assert_preserved(error.0);
1587 assert_eq!(*access.callback_result.lock().unwrap(), Some(true));
1588 let writer = pool.try_writer().expect("writer");
1589 let name: String = writer
1590 .conn()
1591 .query_row(
1592 "SELECT name FROM entities WHERE id=?1",
1593 rusqlite::params![id.to_string()],
1594 |row| row.get(0),
1595 )
1596 .expect("persisted write");
1597 assert_eq!(
1598 name, "changed",
1599 "the lost acknowledgement does not imply rollback"
1600 );
1601 }
1602 }
1603 }
1604}