use aion_core::{
ActivityError, ActivityId, Event, EventEnvelope, Payload, RunId, status::run_segment,
};
use aion_store::OutboxRow;
use chrono::{DateTime, Utc};
use super::Recorder;
use crate::durability::DurabilityError;
#[derive(Clone, Debug)]
pub enum FanOutOutcome {
Completed {
result: Payload,
attempt: u32,
},
Failed {
error: ActivityError,
attempt: u32,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FanOutCompletionResult {
Recorded,
Dropped,
}
#[derive(Clone, Debug)]
pub struct FanOutItem {
pub ordinal: u64,
pub namespace: String,
pub task_queue: String,
pub node: Option<String>,
pub activity_type: String,
pub input: Payload,
pub attempt: u32,
}
impl Recorder {
pub async fn record_fan_out_dispatch(
&mut self,
recorded_at: DateTime<Utc>,
items: &[FanOutItem],
) -> Result<(), DurabilityError> {
if items.is_empty() {
return Ok(());
}
let mut events = Vec::with_capacity(items.len() * 2);
let mut outbox_rows = Vec::with_capacity(items.len());
let mut previous: Option<EventEnvelope> = None;
for item in items {
let scheduled_envelope = match &previous {
Some(previous) => self.envelope_after(previous, recorded_at)?,
None => self.next_envelope(recorded_at)?,
};
let started_envelope = self.envelope_after(&scheduled_envelope, recorded_at)?;
previous = Some(started_envelope.clone());
let activity_id = ActivityId::from_sequence_position(item.ordinal);
events.push(Event::ActivityScheduled {
envelope: scheduled_envelope,
activity_id: activity_id.clone(),
activity_type: item.activity_type.clone(),
input: item.input.clone(),
task_queue: item.task_queue.clone(),
node: item.node.clone(),
});
events.push(Event::ActivityStarted {
envelope: started_envelope,
activity_id,
attempt: item.attempt,
});
outbox_rows.push(
OutboxRow::pending(
self.workflow_id.clone(),
item.ordinal,
item.activity_type.clone(),
item.input.clone(),
recorded_at,
)
.with_run_id(self.run_id.clone())
.with_namespace(item.namespace.clone())
.with_task_queue(item.task_queue.clone())
.with_node(item.node.clone()),
);
}
let expected_seq = self.sequence.current();
self.store
.append_with_outbox(
self.write_token,
&self.workflow_id,
&events,
expected_seq,
&outbox_rows,
)
.await?;
self.sequence.mark_append_success(events.len())
}
pub async fn rearm_outbox_pending(
&self,
recorded_at: DateTime<Utc>,
items: &[FanOutItem],
) -> Result<(), DurabilityError> {
if items.is_empty() {
return Ok(());
}
let rows: Vec<OutboxRow> = items
.iter()
.map(|item| {
OutboxRow::pending(
self.workflow_id.clone(),
item.ordinal,
item.activity_type.clone(),
item.input.clone(),
recorded_at,
)
.with_run_id(self.run_id.clone())
.with_namespace(item.namespace.clone())
.with_task_queue(item.task_queue.clone())
.with_node(item.node.clone())
})
.collect();
self.store.rearm_outbox_pending(&rows).await?;
Ok(())
}
pub async fn record_fan_out_completion(
&mut self,
recorded_at: DateTime<Utc>,
ordinal: u64,
run_id: Option<RunId>,
outcome: FanOutOutcome,
) -> Result<FanOutCompletionResult, DurabilityError> {
let activity_id = ActivityId::from_sequence_position(ordinal);
let history = self.store.read_history(&self.workflow_id).await?;
if run_id_mismatches_recorder(run_id.as_ref(), self.run_id.as_ref()) {
return Ok(FanOutCompletionResult::Dropped);
}
if ordinal_is_resolved(completion_history(&history, run_id.as_ref()), &activity_id) {
return Ok(FanOutCompletionResult::Dropped);
}
match outcome {
FanOutOutcome::Completed { result, attempt } => {
self.append_with(recorded_at, |envelope| Event::ActivityCompleted {
envelope,
activity_id,
result,
attempt,
})
.await?;
}
FanOutOutcome::Failed { error, attempt } => {
self.append_with(recorded_at, |envelope| Event::ActivityFailed {
envelope,
activity_id,
error,
attempt,
})
.await?;
}
}
Ok(FanOutCompletionResult::Recorded)
}
}
fn run_id_mismatches_recorder(completion: Option<&RunId>, current: Option<&RunId>) -> bool {
matches!((completion, current), (Some(completion), Some(current)) if completion != current)
}
fn completion_history<'a>(history: &'a [Event], run_id: Option<&RunId>) -> &'a [Event] {
match run_id {
Some(run_id) => run_segment(history, run_id),
None => history,
}
}
fn ordinal_is_resolved(history: &[Event], activity_id: &ActivityId) -> bool {
history.iter().any(|event| match event {
Event::ActivityCompleted {
activity_id: recorded,
..
}
| Event::ActivityFailed {
activity_id: recorded,
..
}
| Event::ActivityCancelled {
activity_id: recorded,
..
} => recorded == activity_id,
_ => false,
})
}