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 FinalizeEntity(Box<crate::EntityFinalizationPlan>),
48 AddNote(Box<AddNotePlan>),
49 Update(Box<UpdatePlan>),
50 Delete(DeletePlan),
51 Link(LinkPlan),
52 Merge(MergePlan),
53 GtdTransition(GtdTransitionPlan),
54 GtdComplete(GtdCompletePlan),
55 Governance(GovernancePlan),
56}
57
58impl AtomicOpPlan {
59 fn plan_statements(&self) -> Vec<(PlanStatement, bool)> {
69 let statements = match self {
70 AtomicOpPlan::AddEntity(p) => p.statements.clone(),
71 AtomicOpPlan::FinalizeEntity(_) => Vec::new(),
73 AtomicOpPlan::AddNote(p) => p.statements.clone(),
74 AtomicOpPlan::Update(p) => p.statements.clone(),
75 AtomicOpPlan::Delete(p) => p.statements.clone(),
76 AtomicOpPlan::Link(p) => p.statements.clone(),
77 AtomicOpPlan::Merge(p) => {
78 let mut statements: Vec<PlanStatement> = p
79 .rewires
80 .iter()
81 .map(|rewire| PlanStatement {
82 statement: rewire.statement.clone(),
83 guard: None,
84 })
85 .collect();
86 statements.extend(p.lifecycle.clone());
87 statements
88 }
89 AtomicOpPlan::GtdTransition(p) => p.statements.clone(),
90 AtomicOpPlan::GtdComplete(p) => p.statements.clone(),
91 AtomicOpPlan::Governance(p) => p.statements.clone(),
92 };
93 let mut statements: Vec<_> = statements
94 .into_iter()
95 .map(|statement| (statement, false))
96 .collect();
97 if let AtomicOpPlan::Update(plan) = self {
98 statements.extend(plan.graph_effects.iter().map(|effect| match effect {
99 crate::atomic_plan::NoteUpdateStatement::Write(statement) => {
100 (statement.clone(), false)
101 }
102 crate::atomic_plan::NoteUpdateStatement::Assert(statement) => {
103 (statement.clone(), true)
104 }
105 }));
106 }
107 statements
108 }
109
110 fn post_commit_effect(&self) -> Option<PostCommitEffect> {
124 match self {
125 AtomicOpPlan::AddEntity(p) if p.post_commit != PostCommitEffect::None => {
126 Some(p.post_commit.clone())
127 }
128 AtomicOpPlan::AddNote(p) if p.post_commit != PostCommitEffect::None => {
129 Some(p.post_commit.clone())
130 }
131 AtomicOpPlan::Update(p) if p.post_commit != PostCommitEffect::None => {
132 Some(p.post_commit.clone())
133 }
134 AtomicOpPlan::Delete(p) if p.post_commit != PostCommitEffect::None => {
135 Some(p.post_commit.clone())
136 }
137 AtomicOpPlan::GtdTransition(p) if p.post_commit != PostCommitEffect::None => {
138 Some(p.post_commit.clone())
139 }
140 AtomicOpPlan::GtdComplete(p) if p.post_commit != PostCommitEffect::None => {
141 Some(p.post_commit.clone())
142 }
143 _ => None,
144 }
145 }
146}
147
148#[derive(Debug, Clone, PartialEq, Eq)]
151pub enum AtomicOpFailure {
152 NoteConflict(crate::note_write::NoteWriteConflict),
153 EntityConflict(crate::entity_write::EntityVersionConflict),
154 GuardFailed {
160 statement_label: Option<String>,
161 expected: AffectedRowGuard,
162 observed: u64,
163 },
164 SqlError {
168 statement_label: Option<String>,
169 message: String,
170 },
171}
172
173#[derive(Debug, PartialEq, Eq)]
195pub struct CommittedPostCommitEffects {
196 effects: Vec<PostCommitEffect>,
197}
198
199impl CommittedPostCommitEffects {
200 fn new(effects: Vec<PostCommitEffect>) -> Self {
201 Self { effects }
202 }
203
204 pub fn as_slice(&self) -> &[PostCommitEffect] {
207 &self.effects
208 }
209
210 pub(crate) fn into_effects(self) -> Vec<PostCommitEffect> {
211 self.effects
212 }
213}
214
215#[derive(Debug, PartialEq, Eq)]
219pub enum AtomicRunOutcome {
220 Committed {
228 post_commit: CommittedPostCommitEffects,
229 },
230 RolledBack {
238 failed_op_index: usize,
239 failure: AtomicOpFailure,
240 },
241}
242
243#[derive(Debug)]
250pub struct AtomicRunnerError(pub StorageError);
251
252impl std::fmt::Display for AtomicRunnerError {
253 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254 write!(f, "atomic_unit seam failure: {}", self.0)
255 }
256}
257
258impl std::error::Error for AtomicRunnerError {
259 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
260 Some(&self.0)
261 }
262}
263
264async fn execute_savepoint(writer: &mut dyn SqlWriter, sql: String) -> Result<(), StorageError> {
265 writer
266 .execute(SqlStatement::new(sql, vec![]))
267 .await
268 .map(|_| ())
269}
270
271async fn begin_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
272 execute_savepoint(writer, format!("SAVEPOINT {name}")).await
273}
274
275async fn release_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
276 execute_savepoint(writer, format!("RELEASE {name}")).await
277}
278
279async fn rollback_to_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
280 execute_savepoint(writer, format!("ROLLBACK TO {name}")).await
281}
282
283pub(crate) struct AppliedPlan {
290 pub(crate) effect: Option<PostCommitEffect>,
291 pub(crate) note_version: Option<i64>,
292}
293
294pub(crate) async fn apply_plan(
297 writer: &mut dyn SqlWriter,
298 plan: &AtomicOpPlan,
299 capture_note_versions: bool,
300) -> Result<AppliedPlan, AtomicOpFailure> {
301 apply_plan_preserving_storage(writer, plan, capture_note_versions)
302 .await
303 .map_err(|failure| match failure {
304 PlanExecutionFailure::Operation(failure) => failure,
305 PlanExecutionFailure::Storage {
306 statement_label,
307 source,
308 } => AtomicOpFailure::SqlError {
309 statement_label,
310 message: source.to_string(),
311 },
312 })
313}
314pub(crate) enum PlanExecutionFailure {
315 Operation(AtomicOpFailure),
316 Storage {
317 statement_label: Option<String>,
318 source: StorageError,
319 },
320}
321impl From<AtomicOpFailure> for PlanExecutionFailure {
322 fn from(failure: AtomicOpFailure) -> Self {
323 Self::Operation(failure)
324 }
325}
326
327pub(crate) async fn apply_plan_preserving_storage(
328 writer: &mut dyn SqlWriter,
329 plan: &AtomicOpPlan,
330 capture_note_versions: bool,
331) -> Result<AppliedPlan, PlanExecutionFailure> {
332 if let AtomicOpPlan::FinalizeEntity(plan) = plan {
333 return plan
334 .apply(writer)
335 .await
336 .map(|effect| AppliedPlan {
337 effect,
338 note_version: None,
339 })
340 .map_err(PlanExecutionFailure::Operation);
341 }
342 if let AtomicOpPlan::Update(plan) = plan {
343 if let Some(guard) = &plan.entity_guard {
344 if let Some(conflict) =
345 guard
346 .check(writer)
347 .await
348 .map_err(|error| PlanExecutionFailure::Storage {
349 statement_label: Some("entity-version-precondition".into()),
350 source: error,
351 })?
352 {
353 return Err(AtomicOpFailure::EntityConflict(conflict).into());
354 }
355 }
356 }
357 let note_guard = match plan {
358 AtomicOpPlan::Update(plan) => plan.note_guard.as_ref(),
359 AtomicOpPlan::AddNote(plan) => plan.note_guard.as_ref(),
360 _ => None,
361 };
362 if let Some(guard) = note_guard {
363 if let Some(conflict) =
364 guard
365 .check_fence(writer)
366 .await
367 .map_err(|error| PlanExecutionFailure::Storage {
368 statement_label: Some("note-fence".into()),
369 source: error,
370 })?
371 {
372 return Err(AtomicOpFailure::NoteConflict(conflict).into());
373 }
374 }
375 let mut post_commit = plan.post_commit_effect();
376 if let AtomicOpPlan::Update(plan) = plan {
377 if let Some(inheritance) = &plan.note_embedding_inheritance {
378 if !inheritance
379 .vectors
380 .has_rows(writer)
381 .await
382 .map_err(|error| PlanExecutionFailure::Storage {
383 statement_label: Some("note-embedding-inheritance".into()),
384 source: error,
385 })?
386 {
387 post_commit = Some(PostCommitEffect::NoteChanged {
388 note_id: plan.target_id,
389 kind: inheritance.kind.clone(),
390 });
391 }
392 }
393 }
394 for (index, (stmt, read_assertion)) in plan.plan_statements().into_iter().enumerate() {
395 let label = stmt.statement.label.clone();
396 let result = if read_assertion
399 || matches!(plan, AtomicOpPlan::GtdTransition(p) if p.idempotent_noop)
400 || matches!(plan, AtomicOpPlan::Update(p) if p.idempotent_noop)
401 {
402 writer
403 .query_all(stmt.statement)
404 .await
405 .map(|rows| rows.len() as u64)
406 } else {
407 writer.execute(stmt.statement).await
408 };
409 let affected = result.map_err(|e| PlanExecutionFailure::Storage {
410 statement_label: label.clone(),
411 source: e,
412 })?;
413 if let Some(guard) = stmt.guard {
414 if !guard.holds_for(affected) {
415 if let Some(note_guard) = note_guard.filter(|_| index == 0) {
416 if let Some(conflict) =
417 note_guard.classify_refusal(writer).await.map_err(|error| {
418 PlanExecutionFailure::Storage {
419 statement_label: label.clone(),
420 source: error,
421 }
422 })?
423 {
424 return Err(AtomicOpFailure::NoteConflict(conflict).into());
425 }
426 }
427 return Err(AtomicOpFailure::GuardFailed {
428 statement_label: label,
429 expected: guard,
430 observed: affected,
431 }
432 .into());
433 }
434 }
435 }
436 if let AtomicOpPlan::Update(plan) = plan {
437 if let Some(purge) = &plan.note_vector_purge {
438 purge
439 .apply(writer)
440 .await
441 .map_err(|error| PlanExecutionFailure::Storage {
442 statement_label: Some("note-vector-purge".into()),
443 source: error,
444 })?;
445 }
446 }
447 let note_version = if let AtomicOpPlan::AddNote(plan) = plan {
448 if capture_note_versions {
449 let current = writer
450 .query_scalar(
451 SqlStatement::new(
452 "SELECT version FROM notes WHERE id=?1",
453 vec![khive_storage::SqlValue::Text(plan.note_id.to_string())],
454 )
455 .labelled("note-write-guard"),
456 )
457 .await
458 .map_err(|error| PlanExecutionFailure::Storage {
459 statement_label: Some("note-version-receipt".into()),
460 source: error,
461 })?;
462 let Some(khive_storage::SqlValue::Integer(version)) = current else {
463 return Err(AtomicOpFailure::SqlError {
464 statement_label: Some("note-version-receipt".into()),
465 message: "missing persisted note version".into(),
466 }
467 .into());
468 };
469 Some(version)
470 } else {
471 None
472 }
473 } else {
474 None
475 };
476 Ok(AppliedPlan {
477 effect: post_commit,
478 note_version,
479 })
480}
481
482pub async fn run_atomic_unit(
505 access: &dyn SqlAccess,
506 plans: Vec<AtomicOpPlan>,
507) -> Result<AtomicRunOutcome, AtomicRunnerError> {
508 run_atomic_unit_with_note_versions(access, plans, false)
509 .await
510 .map(|(outcome, _)| outcome)
511}
512
513pub(crate) async fn run_atomic_unit_with_note_versions(
516 access: &dyn SqlAccess,
517 plans: Vec<AtomicOpPlan>,
518 capture_note_versions: bool,
519) -> Result<(AtomicRunOutcome, Vec<i64>), AtomicRunnerError> {
520 let finalizers: Vec<_> = plans
521 .iter()
522 .enumerate()
523 .filter_map(|(index, plan)| match plan {
524 AtomicOpPlan::FinalizeEntity(plan) => Some((index, plan.for_settlement())),
525 _ => None,
526 })
527 .collect();
528 let op: PreparedAtomicOp<Vec<i64>, (usize, AtomicOpFailure)> = Box::new(move |writer| {
529 Box::pin(async move {
530 let mut post_commit = Vec::new();
531 let mut note_versions = Vec::new();
532 for (op_index, plan) in plans.iter().enumerate() {
533 let savepoint = format!("adr099_atomic_op_{op_index}");
534 begin_savepoint(writer, &savepoint).await?;
535 match apply_plan(writer, plan, capture_note_versions).await {
536 Ok(applied) => {
537 release_savepoint(writer, &savepoint).await?;
538 if let Some(effect) = applied.effect {
539 post_commit.push(effect);
540 }
541 if let Some(version) = applied.note_version {
542 note_versions.push(version);
543 }
544 }
545 Err(failure) => {
546 let _ = rollback_to_savepoint(writer, &savepoint).await;
554 let _ = release_savepoint(writer, &savepoint).await;
555 return Err(PreparedAtomicError::Refused {
556 failure: (op_index, failure),
557 message: format!("ADR-099 atomic unit aborted at op {op_index}"),
558 });
559 }
560 }
561 }
562 Ok((note_versions, post_commit))
563 })
564 });
565
566 match run_prepared_atomic_unit(access, op).await {
567 Ok(PreparedAtomicOutcome::Committed { value, post_commit }) => {
568 for (_, plan) in finalizers {
569 plan.finish_commit();
570 }
571 Ok((AtomicRunOutcome::Committed { post_commit }, value))
572 }
573 Ok(PreparedAtomicOutcome::RolledBack((failed_op_index, failure))) => {
574 let mut native_source = None;
575 for (index, plan) in finalizers {
576 if index <= failed_op_index {
577 if let Some(source) = plan.finish_rollback(access).await {
578 native_source = Some(source);
579 }
580 }
581 }
582 if let Some(source) = native_source {
583 return Err(AtomicRunnerError(StorageError::WriterTaskRequestFailed {
584 request_state: khive_storage::WriterTaskRequestState::TransactionRolledBack,
585 source: Box::new(source),
586 }));
587 }
588 Ok((
589 AtomicRunOutcome::RolledBack {
590 failed_op_index,
591 failure,
592 },
593 Vec::new(),
594 ))
595 }
596 Err(error) => {
597 if matches!(
600 &error,
601 StorageError::WriterTaskRequestFailed {
602 request_state: khive_storage::WriterTaskRequestState::TransactionRolledBack,
603 ..
604 }
605 ) {
606 for (_, plan) in finalizers {
607 let _ = plan.finish_rollback(access).await;
608 }
609 }
610 Err(AtomicRunnerError(error))
611 }
612 }
613}
614
615pub(crate) enum PreparedAtomicError<E> {
616 Refused { failure: E, message: String },
617 Storage(StorageError),
618}
619
620impl<E> From<StorageError> for PreparedAtomicError<E> {
621 fn from(error: StorageError) -> Self {
622 Self::Storage(error)
623 }
624}
625
626pub(crate) enum PreparedAtomicOutcome<T, E> {
627 Committed {
628 value: T,
629 post_commit: CommittedPostCommitEffects,
630 },
631 RolledBack(E),
632}
633
634pub(crate) type PreparedAtomicOp<T, E> = Box<
635 dyn for<'w> FnOnce(
636 &'w mut dyn SqlWriter,
637 )
638 -> BoxFuture<'w, Result<(T, Vec<PostCommitEffect>), PreparedAtomicError<E>>>
639 + Send,
640>;
641
642pub(crate) async fn run_prepared_atomic_unit<T: Send + 'static, E: Send + 'static>(
645 access: &dyn SqlAccess,
646 prepared: PreparedAtomicOp<T, E>,
647) -> Result<PreparedAtomicOutcome<T, E>, StorageError> {
648 let failure_slot = Arc::new(Mutex::new(None));
649 let recorded = Arc::clone(&failure_slot);
650 let op: AtomicUnitOp = Box::new(move |writer| {
651 Box::pin(async move {
652 match prepared(writer).await {
653 Ok(value) => Ok(Box::new(value) as Box<dyn Any + Send>),
654 Err(PreparedAtomicError::Storage(error)) => Err(error),
655 Err(PreparedAtomicError::Refused { failure, message }) => {
656 *recorded
657 .lock()
658 .expect("atomic runner failure slot poisoned") = Some(failure);
659 Err(StorageError::Internal(message))
660 }
661 }
662 })
663 });
664 match access.atomic_unit(op).await {
665 Ok(boxed) => {
666 let (value, effects) = *boxed
667 .downcast::<(T, Vec<PostCommitEffect>)>()
668 .map_err(|_| StorageError::Internal("invalid prepared atomic outcome".into()))?;
669 Ok(PreparedAtomicOutcome::Committed {
670 value,
671 post_commit: CommittedPostCommitEffects::new(effects),
672 })
673 }
674 Err(storage_err) => {
675 if !atomic_unit_error_allows_recorded_refusal(&storage_err) {
677 return Err(storage_err);
678 }
679 let recorded = failure_slot
680 .lock()
681 .expect("atomic runner failure slot poisoned")
682 .take();
683 match recorded {
684 Some(failure) => Ok(PreparedAtomicOutcome::RolledBack(failure)),
685 None => Err(storage_err),
686 }
687 }
688 }
689}
690
691pub(crate) fn atomic_unit_error_allows_recorded_refusal(error: &StorageError) -> bool {
692 match error.without_sqlite_write_stage() {
693 StorageError::WriterTaskTerminated { .. } => false,
694 StorageError::WriterTaskRequestFailed { request_state, .. } => {
695 *request_state == khive_storage::WriterTaskRequestState::TransactionRolledBack
696 }
697 _ => true,
698 }
699}
700
701#[cfg(test)]
702#[path = "atomic_runner_tests.rs"]
703mod tests;