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
264fn raw_stmt(sql: String) -> SqlStatement {
265 SqlStatement {
266 sql,
267 params: vec![],
268 label: None,
269 }
270}
271
272async fn begin_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
273 writer
274 .execute(raw_stmt(format!("SAVEPOINT {name}")))
275 .await
276 .map(|_| ())
277}
278
279async fn release_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
280 writer
281 .execute(raw_stmt(format!("RELEASE {name}")))
282 .await
283 .map(|_| ())
284}
285
286async fn rollback_to_savepoint(writer: &mut dyn SqlWriter, name: &str) -> Result<(), StorageError> {
287 writer
288 .execute(raw_stmt(format!("ROLLBACK TO {name}")))
289 .await
290 .map(|_| ())
291}
292
293pub(crate) struct AppliedPlan {
300 pub(crate) effect: Option<PostCommitEffect>,
301 pub(crate) note_version: Option<i64>,
302}
303
304pub(crate) async fn apply_plan(
305 writer: &mut dyn SqlWriter,
306 plan: &AtomicOpPlan,
307 capture_note_versions: bool,
308) -> Result<AppliedPlan, AtomicOpFailure> {
309 if let AtomicOpPlan::FinalizeEntity(plan) = plan {
310 return plan.apply(writer).await.map(|effect| AppliedPlan {
311 effect,
312 note_version: None,
313 });
314 }
315 if let AtomicOpPlan::Update(plan) = plan {
316 if let Some(guard) = &plan.entity_guard {
317 if let Some(conflict) =
318 guard
319 .check(writer)
320 .await
321 .map_err(|error| AtomicOpFailure::SqlError {
322 statement_label: Some("entity-version-precondition".into()),
323 message: error.to_string(),
324 })?
325 {
326 return Err(AtomicOpFailure::EntityConflict(conflict));
327 }
328 }
329 }
330 let note_guard = match plan {
331 AtomicOpPlan::Update(plan) => plan.note_guard.as_ref(),
332 AtomicOpPlan::AddNote(plan) => plan.note_guard.as_ref(),
333 _ => None,
334 };
335 if let Some(guard) = note_guard {
336 if let Some(conflict) =
337 guard
338 .check_fence(writer)
339 .await
340 .map_err(|error| AtomicOpFailure::SqlError {
341 statement_label: Some("note-fence".into()),
342 message: error.to_string(),
343 })?
344 {
345 return Err(AtomicOpFailure::NoteConflict(conflict));
346 }
347 }
348 let mut post_commit = plan.post_commit_effect();
349 if let AtomicOpPlan::Update(plan) = plan {
350 if let Some(inheritance) = &plan.note_embedding_inheritance {
351 if !inheritance
352 .vectors
353 .has_rows(writer)
354 .await
355 .map_err(|error| AtomicOpFailure::SqlError {
356 statement_label: Some("note-embedding-inheritance".into()),
357 message: error.to_string(),
358 })?
359 {
360 post_commit = Some(PostCommitEffect::NoteChanged {
361 note_id: plan.target_id,
362 kind: inheritance.kind.clone(),
363 });
364 }
365 }
366 }
367 for (index, (stmt, read_assertion)) in plan.plan_statements().into_iter().enumerate() {
368 let label = stmt.statement.label.clone();
369 let result = if read_assertion
372 || matches!(plan, AtomicOpPlan::GtdTransition(p) if p.idempotent_noop)
373 || matches!(plan, AtomicOpPlan::Update(p) if p.idempotent_noop)
374 {
375 writer
376 .query_all(stmt.statement)
377 .await
378 .map(|rows| rows.len() as u64)
379 } else {
380 writer.execute(stmt.statement).await
381 };
382 let affected = result.map_err(|e| AtomicOpFailure::SqlError {
383 statement_label: label.clone(),
384 message: e.to_string(),
385 })?;
386 if let Some(guard) = stmt.guard {
387 if !guard.holds_for(affected) {
388 if let Some(note_guard) = note_guard.filter(|_| index == 0) {
389 if let Some(conflict) =
390 note_guard.classify_refusal(writer).await.map_err(|error| {
391 AtomicOpFailure::SqlError {
392 statement_label: label.clone(),
393 message: error.to_string(),
394 }
395 })?
396 {
397 return Err(AtomicOpFailure::NoteConflict(conflict));
398 }
399 }
400 return Err(AtomicOpFailure::GuardFailed {
401 statement_label: label,
402 expected: guard,
403 observed: affected,
404 });
405 }
406 }
407 }
408 if let AtomicOpPlan::Update(plan) = plan {
409 if let Some(purge) = &plan.note_vector_purge {
410 purge
411 .apply(writer)
412 .await
413 .map_err(|error| AtomicOpFailure::SqlError {
414 statement_label: Some("note-vector-purge".into()),
415 message: error.to_string(),
416 })?;
417 }
418 }
419 let note_version = if let AtomicOpPlan::AddNote(plan) = plan {
420 if capture_note_versions {
421 let current = writer
422 .query_scalar(
423 SqlStatement::new(
424 "SELECT version FROM notes WHERE id=?1",
425 vec![khive_storage::SqlValue::Text(plan.note_id.to_string())],
426 )
427 .labelled("note-write-guard"),
428 )
429 .await
430 .map_err(|error| AtomicOpFailure::SqlError {
431 statement_label: Some("note-version-receipt".into()),
432 message: error.to_string(),
433 })?;
434 let Some(khive_storage::SqlValue::Integer(version)) = current else {
435 return Err(AtomicOpFailure::SqlError {
436 statement_label: Some("note-version-receipt".into()),
437 message: "missing persisted note version".into(),
438 });
439 };
440 Some(version)
441 } else {
442 None
443 }
444 } else {
445 None
446 };
447 Ok(AppliedPlan {
448 effect: post_commit,
449 note_version,
450 })
451}
452
453pub async fn run_atomic_unit(
476 access: &dyn SqlAccess,
477 plans: Vec<AtomicOpPlan>,
478) -> Result<AtomicRunOutcome, AtomicRunnerError> {
479 run_atomic_unit_with_note_versions(access, plans, false)
480 .await
481 .map(|(outcome, _)| outcome)
482}
483
484pub(crate) async fn run_atomic_unit_with_note_versions(
487 access: &dyn SqlAccess,
488 plans: Vec<AtomicOpPlan>,
489 capture_note_versions: bool,
490) -> Result<(AtomicRunOutcome, Vec<i64>), AtomicRunnerError> {
491 let finalizers: Vec<_> = plans
492 .iter()
493 .enumerate()
494 .filter_map(|(index, plan)| match plan {
495 AtomicOpPlan::FinalizeEntity(plan) => Some((index, plan.for_settlement())),
496 _ => None,
497 })
498 .collect();
499 let op: PreparedAtomicOp<Vec<i64>, (usize, AtomicOpFailure)> = Box::new(move |writer| {
500 Box::pin(async move {
501 let mut post_commit = Vec::new();
502 let mut note_versions = Vec::new();
503 for (op_index, plan) in plans.iter().enumerate() {
504 let savepoint = format!("adr099_atomic_op_{op_index}");
505 begin_savepoint(writer, &savepoint).await?;
506 match apply_plan(writer, plan, capture_note_versions).await {
507 Ok(applied) => {
508 release_savepoint(writer, &savepoint).await?;
509 if let Some(effect) = applied.effect {
510 post_commit.push(effect);
511 }
512 if let Some(version) = applied.note_version {
513 note_versions.push(version);
514 }
515 }
516 Err(failure) => {
517 let _ = rollback_to_savepoint(writer, &savepoint).await;
525 let _ = release_savepoint(writer, &savepoint).await;
526 return Err(PreparedAtomicError::Refused {
527 failure: (op_index, failure),
528 message: format!("ADR-099 atomic unit aborted at op {op_index}"),
529 });
530 }
531 }
532 }
533 Ok((note_versions, post_commit))
534 })
535 });
536
537 match run_prepared_atomic_unit(access, op).await {
538 Ok(PreparedAtomicOutcome::Committed { value, post_commit }) => {
539 for (_, plan) in finalizers {
540 plan.finish_commit();
541 }
542 Ok((AtomicRunOutcome::Committed { post_commit }, value))
543 }
544 Ok(PreparedAtomicOutcome::RolledBack((failed_op_index, failure))) => {
545 let mut native_source = None;
546 for (index, plan) in finalizers {
547 if index <= failed_op_index {
548 if let Some(source) = plan.finish_rollback(access).await {
549 native_source = Some(source);
550 }
551 }
552 }
553 if let Some(source) = native_source {
554 return Err(AtomicRunnerError(StorageError::WriterTaskRequestFailed {
555 request_state: khive_storage::WriterTaskRequestState::TransactionRolledBack,
556 source: Box::new(source),
557 }));
558 }
559 Ok((
560 AtomicRunOutcome::RolledBack {
561 failed_op_index,
562 failure,
563 },
564 Vec::new(),
565 ))
566 }
567 Err(error) => {
568 if matches!(
571 &error,
572 StorageError::WriterTaskRequestFailed {
573 request_state: khive_storage::WriterTaskRequestState::TransactionRolledBack,
574 ..
575 }
576 ) {
577 for (_, plan) in finalizers {
578 let _ = plan.finish_rollback(access).await;
579 }
580 }
581 Err(AtomicRunnerError(error))
582 }
583 }
584}
585
586pub(crate) enum PreparedAtomicError<E> {
587 Refused { failure: E, message: String },
588 Storage(StorageError),
589}
590
591impl<E> From<StorageError> for PreparedAtomicError<E> {
592 fn from(error: StorageError) -> Self {
593 Self::Storage(error)
594 }
595}
596
597pub(crate) enum PreparedAtomicOutcome<T, E> {
598 Committed {
599 value: T,
600 post_commit: CommittedPostCommitEffects,
601 },
602 RolledBack(E),
603}
604
605pub(crate) type PreparedAtomicOp<T, E> = Box<
606 dyn for<'w> FnOnce(
607 &'w mut dyn SqlWriter,
608 )
609 -> BoxFuture<'w, Result<(T, Vec<PostCommitEffect>), PreparedAtomicError<E>>>
610 + Send,
611>;
612
613pub(crate) async fn run_prepared_atomic_unit<T: Send + 'static, E: Send + 'static>(
616 access: &dyn SqlAccess,
617 prepared: PreparedAtomicOp<T, E>,
618) -> Result<PreparedAtomicOutcome<T, E>, StorageError> {
619 let failure_slot = Arc::new(Mutex::new(None));
620 let recorded = Arc::clone(&failure_slot);
621 let op: AtomicUnitOp = Box::new(move |writer| {
622 Box::pin(async move {
623 match prepared(writer).await {
624 Ok(value) => Ok(Box::new(value) as Box<dyn Any + Send>),
625 Err(PreparedAtomicError::Storage(error)) => Err(error),
626 Err(PreparedAtomicError::Refused { failure, message }) => {
627 *recorded
628 .lock()
629 .expect("atomic runner failure slot poisoned") = Some(failure);
630 Err(StorageError::Internal(message))
631 }
632 }
633 })
634 });
635 match access.atomic_unit(op).await {
636 Ok(boxed) => {
637 let (value, effects) = *boxed
638 .downcast::<(T, Vec<PostCommitEffect>)>()
639 .map_err(|_| StorageError::Internal("invalid prepared atomic outcome".into()))?;
640 Ok(PreparedAtomicOutcome::Committed {
641 value,
642 post_commit: CommittedPostCommitEffects::new(effects),
643 })
644 }
645 Err(storage_err) => {
646 if !atomic_unit_error_allows_recorded_refusal(&storage_err) {
648 return Err(storage_err);
649 }
650 let recorded = failure_slot
651 .lock()
652 .expect("atomic runner failure slot poisoned")
653 .take();
654 match recorded {
655 Some(failure) => Ok(PreparedAtomicOutcome::RolledBack(failure)),
656 None => Err(storage_err),
657 }
658 }
659 }
660}
661
662pub(crate) fn atomic_unit_error_allows_recorded_refusal(error: &StorageError) -> bool {
663 match error.without_sqlite_write_stage() {
664 StorageError::WriterTaskTerminated { .. } => false,
665 StorageError::WriterTaskRequestFailed { request_state, .. } => {
666 *request_state == khive_storage::WriterTaskRequestState::TransactionRolledBack
667 }
668 _ => true,
669 }
670}
671
672#[cfg(test)]
673#[path = "atomic_runner_tests.rs"]
674mod tests;