1use std::future::Future;
4use std::sync::Arc;
5
6use aion_core::{ActivityError, ActivityId, Payload, RunId, WorkflowId};
7use aion_store::EventStore;
8use chrono::{DateTime, Utc};
9use tokio::runtime::Handle;
10use tokio::sync::Mutex;
11
12use crate::EngineError;
13use crate::durability::{
14 Command, DurabilityError, HistoryCursor, Recorder, ResolveOutcome, Resolver,
15};
16use crate::registry::{Registry, WorkflowHandle};
17
18#[derive(thiserror::Error, Debug)]
20pub enum NifContextError {
21 #[error("unknown workflow process pid {pid}")]
23 UnknownProcess {
24 pid: u64,
26 },
27 #[error("workflow recorder lock is poisoned")]
29 RecorderPoisoned,
30 #[error("durability error: {0}")]
32 Durability(#[from] DurabilityError),
33 #[error("term encoding error: {reason}")]
35 TermEncoding {
36 reason: String,
38 },
39}
40
41impl NifContextError {
42 pub(crate) fn error_reason(&self) -> String {
48 match self {
49 Self::UnknownProcess { pid } => format!("unknown_process:{pid}"),
50 Self::RecorderPoisoned => "recorder_poisoned".to_owned(),
51 Self::Durability(error) => format!("durability:{error}"),
52 Self::TermEncoding { reason } => format!("term_encoding:{reason}"),
53 }
54 }
55}
56
57pub struct NifContext {
59 handle: WorkflowHandle,
60 recorder: Arc<Mutex<Recorder>>,
61 tokio_handle: Handle,
62 resolver: Resolver,
63 last_recorded_at: Option<DateTime<Utc>>,
64}
65
66impl NifContext {
67 pub fn new(
77 pid: u64,
78 registry: &Registry,
79 tokio_handle: Handle,
80 birth_wait: crate::runtime::SignalDeliveryConfig,
81 ) -> Result<Self, NifContextError> {
82 Self::new_with_history_store(pid, registry, tokio_handle, None, birth_wait)
83 }
84
85 pub fn new_with_history_store(
96 pid: u64,
97 registry: &Registry,
98 tokio_handle: Handle,
99 store: Option<Arc<dyn EventStore>>,
100 birth_wait: crate::runtime::SignalDeliveryConfig,
101 ) -> Result<Self, NifContextError> {
102 let handle = resolve_handle_with_birth_wait(registry, pid, birth_wait)?;
103 let recorder = handle.recorder();
104 let workflow_id = handle.workflow_id().clone();
105 let history = match store {
106 Some(store) => tokio_handle
107 .block_on(store.read_history(&workflow_id))
108 .map_err(DurabilityError::from)?,
109 None => tokio_handle.block_on(async {
110 let recorder = recorder.lock().await;
111 recorder.read_history().await
112 })?,
113 };
114 let history = crate::durability::current_run_segment(history, handle.run_id())?;
117 let last_recorded_at = history.last().map(|event| *event.recorded_at());
118 let cursor = HistoryCursor::new(history)?;
119 let resolver = Resolver::new(workflow_id, cursor);
120
121 Ok(Self {
122 handle,
123 recorder,
124 tokio_handle,
125 resolver,
126 last_recorded_at,
127 })
128 }
129
130 #[must_use]
132 pub fn workflow_id(&self) -> &WorkflowId {
133 self.handle.workflow_id()
134 }
135
136 #[must_use]
138 pub fn run_id(&self) -> &RunId {
139 self.handle.run_id()
140 }
141
142 #[must_use]
149 pub fn next_activity_ordinal(&self) -> u64 {
150 self.handle.allocate_activity_ordinals(1)
151 }
152
153 #[must_use]
155 pub fn allocate_activity_ordinals(&self, count: u64) -> u64 {
156 self.handle.allocate_activity_ordinals(count)
157 }
158
159 #[must_use]
164 pub fn next_timer_ordinal(&self) -> u64 {
165 self.handle.allocate_timer_ordinals(1)
166 }
167
168 #[must_use]
177 pub fn next_child_ordinal(&self) -> u64 {
178 self.handle.allocate_child_ordinals(1)
179 }
180
181 #[must_use]
183 pub fn signal_receives_consumed(&self, name: &str) -> u64 {
184 self.handle.signal_receives_consumed(name)
185 }
186
187 pub fn mark_signal_receive_consumed(&self, name: &str) {
189 self.handle.mark_signal_receive_consumed(name);
190 }
191
192 #[must_use]
194 pub fn signal_sends_completed(&self, name: &str) -> u64 {
195 self.handle.signal_sends_completed(name)
196 }
197
198 pub fn mark_signal_send_completed(&self, name: &str) {
200 self.handle.mark_signal_send_completed(name);
201 }
202
203 #[must_use]
205 pub fn workflow_handle(&self) -> WorkflowHandle {
206 self.handle.clone()
207 }
208
209 #[must_use]
211 pub const fn pid(&self) -> u64 {
212 self.handle.pid()
213 }
214
215 #[must_use]
217 pub const fn last_recorded_at(&self) -> Option<DateTime<Utc>> {
218 self.last_recorded_at
219 }
220
221 #[must_use]
223 pub fn next_deterministic_sequence(&self) -> u64 {
224 self.handle.next_deterministic_nif_sequence()
225 }
226
227 #[must_use]
229 pub fn recorder(&self) -> Arc<Mutex<Recorder>> {
230 Arc::clone(&self.recorder)
231 }
232
233 pub fn block_on_recorder<T, F>(&self, f: F) -> Result<T, NifContextError>
239 where
240 F: for<'a> FnOnce(
241 &'a mut Recorder,
242 ) -> std::pin::Pin<
243 Box<dyn Future<Output = Result<T, DurabilityError>> + Send + 'a>,
244 >,
245 {
246 self.tokio_handle
247 .block_on(async {
248 let mut recorder = self.recorder.lock().await;
249 f(&mut recorder).await
250 })
251 .map_err(Into::into)
252 }
253
254 pub fn record_activity_scheduled_started(
260 &self,
261 recorded_at: chrono::DateTime<chrono::Utc>,
262 activity_id: ActivityId,
263 activity_type: String,
264 input: Payload,
265 ) -> Result<(), NifContextError> {
266 self.tokio_handle
267 .block_on(async {
268 let mut recorder = self.recorder.lock().await;
269 recorder
270 .record_activity_scheduled(
271 recorded_at,
272 activity_id.clone(),
273 activity_type,
274 input,
275 )
276 .await?;
277 recorder
278 .record_activity_started(recorded_at, activity_id)
279 .await
280 })
281 .map_err(Into::into)
282 }
283
284 pub fn record_activity_completed(
290 &self,
291 recorded_at: chrono::DateTime<chrono::Utc>,
292 activity_id: ActivityId,
293 result: Payload,
294 ) -> Result<(), NifContextError> {
295 self.tokio_handle
296 .block_on(async {
297 let mut recorder = self.recorder.lock().await;
298 recorder
299 .record_activity_completed(recorded_at, activity_id, result)
300 .await
301 })
302 .map_err(Into::into)
303 }
304
305 pub fn record_activity_failed(
311 &self,
312 recorded_at: chrono::DateTime<chrono::Utc>,
313 activity_id: ActivityId,
314 error: ActivityError,
315 attempt: u32,
316 ) -> Result<(), NifContextError> {
317 self.tokio_handle
318 .block_on(async {
319 let mut recorder = self.recorder.lock().await;
320 recorder
321 .record_activity_failed(recorded_at, activity_id, error, attempt)
322 .await
323 })
324 .map_err(Into::into)
325 }
326
327 pub fn record_activity_cancelled(
333 &self,
334 recorded_at: chrono::DateTime<chrono::Utc>,
335 activity_id: ActivityId,
336 ) -> Result<(), NifContextError> {
337 self.tokio_handle
338 .block_on(async {
339 let mut recorder = self.recorder.lock().await;
340 recorder
341 .record_activity_cancelled(recorded_at, activity_id)
342 .await
343 })
344 .map_err(Into::into)
345 }
346
347 #[must_use]
349 pub fn history(&self) -> &[aion_core::Event] {
350 self.resolver.history()
351 }
352
353 pub fn resolve_command(&mut self, command: Command) -> Result<ResolveOutcome, NifContextError> {
360 if let Some(key) = command.key() {
368 self.resolver.fast_forward_to(key);
369 } else if let Command::AwaitChild { child_workflow_id } = &command {
370 self.resolver
371 .fast_forward_to_child_terminal(child_workflow_id);
372 }
373 self.resolver.resolve(command).map_err(Into::into)
374 }
375}
376
377fn registry_error_to_context(error: &EngineError) -> NifContextError {
378 match error {
379 EngineError::RegistryPoisoned => NifContextError::RecorderPoisoned,
380 _ => NifContextError::TermEncoding {
381 reason: format!("registry lookup failed: {error}"),
382 },
383 }
384}
385
386fn resolve_handle_with_birth_wait(
407 registry: &Registry,
408 pid: u64,
409 birth_wait: crate::runtime::SignalDeliveryConfig,
410) -> Result<WorkflowHandle, NifContextError> {
411 let lookup = |registry: &Registry| -> Result<Option<WorkflowHandle>, NifContextError> {
412 Ok(registry
413 .list()
414 .map_err(|error| registry_error_to_context(&error))?
415 .into_iter()
416 .find(|handle| handle.pid() == pid))
417 };
418 if let Some(handle) = lookup(registry)? {
419 return Ok(handle);
420 }
421 let budget = birth_wait
422 .ready_timeout
423 .saturating_mul(birth_wait.max_enqueue_attempts.max(1));
424 let deadline = std::time::Instant::now() + budget;
425 let mut backoff = birth_wait.initial_backoff;
426 while std::time::Instant::now() < deadline {
427 std::thread::sleep(backoff);
428 let doubled = backoff.saturating_mul(2);
429 backoff = if doubled > birth_wait.max_backoff {
430 birth_wait.max_backoff
431 } else {
432 doubled
433 };
434 if let Some(handle) = lookup(registry)? {
435 return Ok(handle);
436 }
437 }
438 Err(NifContextError::UnknownProcess { pid })
439}
440
441#[cfg(test)]
442mod tests {
443 use std::sync::Arc;
444
445 use aion_core::{ActivityId, Event, EventEnvelope, Payload, WorkflowStatus};
446 use aion_package::ContentHash;
447 use aion_store::{EventStore, InMemoryStore, WriteToken};
448 use chrono::{TimeZone, Utc};
449 use serde_json::json;
450
451 use super::{NifContext, NifContextError};
452 use crate::durability::{Command, CorrelationKey, Recorder, Resolution, ResolveOutcome};
453 use crate::registry::{
454 CompletionNotifier, HandleResidency, Registry, WorkflowHandle, WorkflowHandleParts,
455 };
456
457 type TestResult = Result<(), Box<dyn std::error::Error>>;
458
459 fn hash() -> ContentHash {
460 ContentHash::from_bytes([7; 32])
461 }
462
463 fn birth_wait() -> crate::runtime::SignalDeliveryConfig {
465 crate::runtime::SignalDeliveryConfig::new(
466 std::time::Duration::from_millis(200),
467 1,
468 std::time::Duration::from_millis(2),
469 std::time::Duration::from_millis(8),
470 )
471 }
472
473 fn payload(label: &str) -> Result<Payload, Box<dyn std::error::Error>> {
474 Ok(Payload::from_json(&json!({ "label": label }))?)
475 }
476
477 fn envelope(
478 workflow_id: &aion_core::WorkflowId,
479 seq: u64,
480 ) -> Result<EventEnvelope, Box<dyn std::error::Error>> {
481 let recorded_at = Utc
482 .timestamp_opt(i64::try_from(seq)?, 0)
483 .single()
484 .ok_or_else(|| "invalid timestamp".to_owned())?;
485 Ok(EventEnvelope {
486 seq,
487 recorded_at,
488 workflow_id: workflow_id.clone(),
489 })
490 }
491
492 fn started_event(
493 workflow_id: &aion_core::WorkflowId,
494 run_id: &aion_core::RunId,
495 ) -> Result<Event, Box<dyn std::error::Error>> {
496 Ok(Event::WorkflowStarted {
497 envelope: envelope(workflow_id, 1)?,
498 workflow_type: "checkout".to_owned(),
499 input: payload("input")?,
500 run_id: run_id.clone(),
501 parent_run_id: None,
502 package_version: aion_core::PackageVersion::new("a".repeat(64)),
503 })
504 }
505
506 fn handle(
507 pid: u64,
508 store: Arc<dyn EventStore>,
509 workflow_id: aion_core::WorkflowId,
510 run_id: aion_core::RunId,
511 ) -> WorkflowHandle {
512 let recorder = Recorder::resume_at(workflow_id.clone(), store, 1);
513 WorkflowHandle::new(WorkflowHandleParts {
514 workflow_id,
515 run_id,
516 pid,
517 workflow_type: "checkout".to_owned(),
518 loaded_version: hash(),
519 cached_status: WorkflowStatus::Running,
520 residency: HandleResidency::Resident,
521 recorder,
522 completion: CompletionNotifier::new(),
523 })
524 }
525
526 type TestContext = (Registry, Arc<dyn EventStore>, WorkflowHandle);
527
528 fn context_with_history(
529 runtime: &tokio::runtime::Runtime,
530 pid: u64,
531 workflow_id: aion_core::WorkflowId,
532 history: &[Event],
533 ) -> Result<TestContext, Box<dyn std::error::Error>> {
534 let registry = Registry::default();
535 let run_id = aion_core::RunId::new_v4();
536 let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
537 let mut full_history = vec![started_event(&workflow_id, &run_id)?];
538 full_history.extend_from_slice(history);
539 runtime.block_on(store.append(WriteToken::recorder(), &workflow_id, &full_history, 0))?;
540 let recorder = Recorder::resume_at(
541 workflow_id.clone(),
542 Arc::clone(&store),
543 full_history.len() as u64,
544 );
545 let handle = WorkflowHandle::new(WorkflowHandleParts {
546 workflow_id: workflow_id.clone(),
547 run_id: run_id.clone(),
548 pid,
549 workflow_type: "checkout".to_owned(),
550 loaded_version: hash(),
551 cached_status: WorkflowStatus::Running,
552 residency: HandleResidency::Resident,
553 recorder,
554 completion: CompletionNotifier::new(),
555 });
556 registry.insert((workflow_id, run_id), handle.clone())?;
557 Ok((registry, store, handle))
558 }
559
560 #[test]
561 fn resolves_registered_pid_to_context() -> TestResult {
562 let runtime = tokio::runtime::Runtime::new()?;
563 let registry = Registry::default();
564 let workflow_id = aion_core::WorkflowId::new_v4();
565 let run_id = aion_core::RunId::new_v4();
566 let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
567 runtime.block_on(store.append(
568 WriteToken::recorder(),
569 &workflow_id,
570 &[started_event(&workflow_id, &run_id)?],
571 0,
572 ))?;
573 let handle = handle(44, Arc::clone(&store), workflow_id.clone(), run_id.clone());
574 registry.insert((workflow_id.clone(), run_id), handle)?;
575
576 let context = NifContext::new(44, ®istry, runtime.handle().clone(), birth_wait())?;
577
578 assert_eq!(context.workflow_id(), &workflow_id);
579 assert_eq!(context.pid(), 44);
580 Ok(())
581 }
582
583 #[test]
584 fn unknown_pid_returns_unknown_process() -> TestResult {
585 let runtime = tokio::runtime::Runtime::new()?;
586 let registry = Registry::default();
587
588 let error = NifContext::new(77, ®istry, runtime.handle().clone(), birth_wait())
589 .err()
590 .ok_or("expected unknown process error")?;
591
592 assert!(matches!(error, NifContextError::UnknownProcess { pid: 77 }));
593 Ok(())
594 }
595
596 #[test]
604 fn birth_window_registration_resolves_instead_of_failing() -> TestResult {
605 let runtime = tokio::runtime::Runtime::new()?;
606 let registry = Arc::new(Registry::default());
607 let workflow_id = aion_core::WorkflowId::new_v4();
608 let run_id = aion_core::RunId::new_v4();
609 let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
610 runtime.block_on(store.append(
611 WriteToken::recorder(),
612 &workflow_id,
613 &[started_event(&workflow_id, &run_id)?],
614 0,
615 ))?;
616 let handle = handle(91, Arc::clone(&store), workflow_id.clone(), run_id.clone());
617
618 let late_registry = Arc::clone(®istry);
621 let inserter = std::thread::spawn(move || {
622 std::thread::sleep(std::time::Duration::from_millis(30));
623 late_registry.insert((workflow_id.clone(), run_id), handle)
624 });
625
626 let context = NifContext::new(91, ®istry, runtime.handle().clone(), birth_wait())?;
627
628 assert_eq!(context.pid(), 91);
629 inserter
630 .join()
631 .map_err(|_| "registry insert thread panicked")??;
632 Ok(())
633 }
634
635 #[test]
636 fn block_on_recorder_reads_current_head_without_deadlock() -> TestResult {
637 let runtime = tokio::runtime::Runtime::new()?;
638 let registry = Registry::default();
639 let workflow_id = aion_core::WorkflowId::new_v4();
640 let run_id = aion_core::RunId::new_v4();
641 let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
642 runtime.block_on(store.append(
643 WriteToken::recorder(),
644 &workflow_id,
645 &[started_event(&workflow_id, &run_id)?],
646 0,
647 ))?;
648 let recorder = Recorder::resume_at(workflow_id.clone(), Arc::clone(&store), 5);
649 let handle = WorkflowHandle::new(WorkflowHandleParts {
650 workflow_id: workflow_id.clone(),
651 run_id: run_id.clone(),
652 pid: 55,
653 workflow_type: "checkout".to_owned(),
654 loaded_version: hash(),
655 cached_status: WorkflowStatus::Running,
656 residency: HandleResidency::Resident,
657 recorder,
658 completion: CompletionNotifier::new(),
659 });
660 registry.insert((workflow_id, run_id), handle)?;
661 let context = NifContext::new(55, ®istry, runtime.handle().clone(), birth_wait())?;
662
663 let head = context
664 .block_on_recorder(|recorder| Box::pin(async move { Ok(recorder.current_head()) }))?;
665
666 assert_eq!(head, 5);
667 Ok(())
668 }
669
670 #[test]
671 fn resolve_command_returns_recorded_activity_resolution() -> TestResult {
672 let runtime = tokio::runtime::Runtime::new()?;
673 let workflow_id = aion_core::WorkflowId::new_v4();
674 let result = payload("activity-result")?;
675 let history = vec![
676 Event::ActivityScheduled {
677 envelope: envelope(&workflow_id, 2)?,
678 activity_id: ActivityId::from_sequence_position(0),
679 activity_type: "activity".to_owned(),
680 input: payload("activity-input")?,
681 },
682 Event::ActivityCompleted {
683 envelope: envelope(&workflow_id, 3)?,
684 activity_id: ActivityId::from_sequence_position(0),
685 result: result.clone(),
686 },
687 ];
688 let (registry, store, handle) = context_with_history(&runtime, 66, workflow_id, &history)?;
689 let mut context = NifContext::new_with_history_store(
690 66,
691 ®istry,
692 runtime.handle().clone(),
693 Some(store),
694 birth_wait(),
695 )?;
696
697 assert_eq!(context.workflow_id(), handle.workflow_id());
698 assert_eq!(
699 context.resolve_command(Command::RunActivity {
700 key: CorrelationKey::Activity(0),
701 activity_type: "activity".to_owned(),
702 input: payload("activity-input")?,
703 })?,
704 ResolveOutcome::Recorded(Resolution::ActivityCompleted(result))
705 );
706 Ok(())
707 }
708
709 fn child_history(
710 workflow_id: &aion_core::WorkflowId,
711 child_workflow_id: &aion_core::WorkflowId,
712 include_terminal: bool,
713 ) -> Result<Vec<Event>, Box<dyn std::error::Error>> {
714 let timer_id = aion_core::TimerId::anonymous(0);
715 let mut history = vec![
716 Event::ActivityScheduled {
717 envelope: envelope(workflow_id, 2)?,
718 activity_id: ActivityId::from_sequence_position(0),
719 activity_type: "activity".to_owned(),
720 input: payload("activity-input")?,
721 },
722 Event::ActivityCompleted {
723 envelope: envelope(workflow_id, 3)?,
724 activity_id: ActivityId::from_sequence_position(0),
725 result: payload("activity-result")?,
726 },
727 Event::TimerStarted {
728 envelope: envelope(workflow_id, 4)?,
729 timer_id: timer_id.clone(),
730 fire_at: Utc
731 .timestamp_opt(99, 0)
732 .single()
733 .ok_or_else(|| "invalid timestamp".to_owned())?,
734 },
735 Event::TimerFired {
736 envelope: envelope(workflow_id, 5)?,
737 timer_id,
738 },
739 Event::ChildWorkflowStarted {
740 envelope: envelope(workflow_id, 6)?,
741 child_workflow_id: child_workflow_id.clone(),
742 workflow_type: "child".to_owned(),
743 input: payload("child-input")?,
744 package_version: aion_core::PackageVersion::new("a".repeat(64)),
745 },
746 ];
747 if include_terminal {
748 history.push(Event::ChildWorkflowCompleted {
749 envelope: envelope(workflow_id, 7)?,
750 child_workflow_id: child_workflow_id.clone(),
751 result: payload("child-result")?,
752 });
753 }
754 Ok(history)
755 }
756
757 #[test]
758 fn await_child_skips_consumed_commands_to_recorded_terminal() -> TestResult {
759 let runtime = tokio::runtime::Runtime::new()?;
760 let workflow_id = aion_core::WorkflowId::new_v4();
761 let child_workflow_id = aion_core::WorkflowId::new_v4();
762 let history = child_history(&workflow_id, &child_workflow_id, true)?;
767 let (registry, store, _handle) = context_with_history(&runtime, 88, workflow_id, &history)?;
768 let mut context = NifContext::new_with_history_store(
769 88,
770 ®istry,
771 runtime.handle().clone(),
772 Some(store),
773 birth_wait(),
774 )?;
775
776 assert_eq!(
777 context.resolve_command(Command::AwaitChild {
778 child_workflow_id: child_workflow_id.clone(),
779 })?,
780 ResolveOutcome::Recorded(Resolution::ChildCompleted(payload("child-result")?))
781 );
782 Ok(())
783 }
784
785 #[test]
786 fn await_child_without_recorded_terminal_resumes_live() -> TestResult {
787 let runtime = tokio::runtime::Runtime::new()?;
788 let workflow_id = aion_core::WorkflowId::new_v4();
789 let child_workflow_id = aion_core::WorkflowId::new_v4();
790 let history = child_history(&workflow_id, &child_workflow_id, false)?;
794 let (registry, store, _handle) = context_with_history(&runtime, 89, workflow_id, &history)?;
795 let mut context = NifContext::new_with_history_store(
796 89,
797 ®istry,
798 runtime.handle().clone(),
799 Some(store),
800 birth_wait(),
801 )?;
802
803 assert_eq!(
804 context.resolve_command(Command::AwaitChild { child_workflow_id })?,
805 ResolveOutcome::ResumeLive
806 );
807 Ok(())
808 }
809}