use std::sync::Arc;
use aion::ActivityDispatch;
use aion_core::{ActivityError, ActivityErrorKind, ActivityId, ContentType, Payload, RunId};
use aion_package::ActionBodyContract;
use aion_store::OutboxRow;
use async_trait::async_trait;
use super::DeclaredCommandExecutor;
use super::declared_body::{DeclaredBodyLookup, DeclaredBodySource, DispatchingRun};
use super::declared_body_ambiguity::ambiguous_body_refusal;
use super::dispatch::{ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink};
use super::envelope::{CompletionFences, CompletionToken};
use super::outbox_dispatcher::OutboxRowDispatch;
use crate::error::ServerError;
const ROW_DISPATCH_CONFIG: &str = "{}";
pub struct DeclaredBodyOutboxDispatch {
inner: Arc<dyn OutboxRowDispatch>,
bodies: DeclaredBodySource,
executor: Arc<DeclaredCommandExecutor>,
fences: CompletionFences,
sink: Arc<dyn ActivityCompletionSink + Send + Sync>,
handle: tokio::runtime::Handle,
}
impl std::fmt::Debug for DeclaredBodyOutboxDispatch {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("DeclaredBodyOutboxDispatch")
.field("bodies", &self.bodies)
.finish_non_exhaustive()
}
}
impl DeclaredBodyOutboxDispatch {
#[must_use]
pub fn new(
inner: Arc<dyn OutboxRowDispatch>,
bodies: DeclaredBodySource,
executor: Arc<DeclaredCommandExecutor>,
fences: CompletionFences,
sink: Arc<dyn ActivityCompletionSink + Send + Sync>,
handle: tokio::runtime::Handle,
) -> Self {
Self {
inner,
bodies,
executor,
fences,
sink,
handle,
}
}
fn run_id(row: &OutboxRow) -> Result<RunId, ServerError> {
row.run_id.clone().ok_or_else(|| {
ServerError::worker_dispatch(
&row.namespace,
&row.activity_type,
format!(
"outbox row {} carries no run id; a declared body cannot be completed \
against a run it does not name",
row.dispatch_key
),
)
})
}
fn request_for(row: &OutboxRow, run_id: RunId) -> Result<ActivityDispatch, String> {
let input = String::from_utf8(row.input.bytes().to_vec()).map_err(|_| {
"terminal:declared body input on the outbox row is not valid UTF-8".to_owned()
})?;
Ok(ActivityDispatch {
namespace: row.namespace.clone(),
task_queue: row.task_queue.clone(),
node: row.node.clone(),
workflow_id: row.workflow_id.clone(),
run_id,
activity_id: ActivityId::from_sequence_position(row.ordinal),
name: row.activity_type.clone(),
input,
config: ROW_DISPATCH_CONFIG.to_owned(),
attempt: row.started_attempt,
labels: std::collections::BTreeMap::new(),
advisory: false,
})
}
fn mint_token(&self, row: &OutboxRow, run_id: &RunId) -> Result<CompletionToken, ServerError> {
self.fences.issue(
&row.workflow_id,
run_id,
&ActivityId::from_sequence_position(row.ordinal),
row.started_attempt,
)
}
fn complete(
sink: &(dyn ActivityCompletionSink + Send + Sync),
row: &OutboxRow,
run_id: RunId,
token: CompletionToken,
ended: Result<String, String>,
) -> Result<(), ServerError> {
let outcome = match ended {
Ok(encoded) => ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
encoded.into_bytes(),
)),
Err(reason) => ActivityCompletionOutcome::Failed(classify_failure(&reason)),
};
sink.complete_activity(ActivityCompletion {
workflow_id: row.workflow_id.clone(),
activity_id: ActivityId::from_sequence_position(row.ordinal),
run_id: Some(run_id),
completion_token: token,
outcome,
})
}
fn place(&self, row: &OutboxRow, contract: ActionBodyContract) -> Result<(), ServerError> {
let run_id = Self::run_id(row)?;
let token = self.mint_token(row, &run_id)?;
let request = match Self::request_for(row, run_id.clone()) {
Ok(request) => request,
Err(reason) => {
return Self::complete(self.sink.as_ref(), row, run_id, token, Err(reason));
}
};
tracing::info!(
operation = "declared_command_outbox_dispatch",
workflow_id = %row.workflow_id,
activity_id = %request.activity_id,
activity_name = %row.activity_type,
task_queue = %row.task_queue,
attempt = row.started_attempt,
"executing a declared action body at the server for an outbox (fan-out) row"
);
let executor = Arc::clone(&self.executor);
let sink = Arc::clone(&self.sink);
let fences = self.fences.clone();
let row = row.clone();
self.handle.spawn(async move {
let contract_for_run = contract;
let request_for_run = request;
let executor_for_run = Arc::clone(&executor);
let ended = tokio::task::spawn_blocking(move || {
executor_for_run.execute(&request_for_run, &contract_for_run)
})
.await
.unwrap_or_else(|join_error| {
Err(format!(
"terminal:the declared body's execution task ended abnormally: {join_error}"
))
});
if ended
.as_ref()
.err()
.is_some_and(|reason| reason == aion::PARKED_ACTIVITY_REASON)
{
if let Err(error) = fences.revoke(
&row.workflow_id,
&ActivityId::from_sequence_position(row.ordinal),
&token,
) {
tracing::warn!(
workflow_id = %row.workflow_id,
ordinal = row.ordinal,
%error,
"could not withdraw the completion token of a parked declared body"
);
}
tracing::info!(
operation = "declared_command_outbox_dispatch",
workflow_id = %row.workflow_id,
ordinal = row.ordinal,
"declared body parked: this server is draining and starts no new work"
);
return;
}
if let Err(error) = Self::complete(sink.as_ref(), &row, run_id, token, ended) {
tracing::error!(
operation = "declared_command_outbox_dispatch",
workflow_id = %row.workflow_id,
ordinal = row.ordinal,
%error,
"a declared body ran for an outbox row but its completion was refused"
);
}
});
Ok(())
}
}
fn classify_failure(reason: &str) -> ActivityError {
let (kind, message) = if let Some(rest) = reason.strip_prefix("terminal:") {
(ActivityErrorKind::Terminal, rest.to_owned())
} else if let Some(rest) = reason.strip_prefix("policy_refused:") {
(ActivityErrorKind::PolicyRefused, rest.to_owned())
} else if let Some(rest) = reason.strip_prefix("retryable:") {
(ActivityErrorKind::Retryable, rest.to_owned())
} else {
(ActivityErrorKind::Retryable, reason.to_owned())
};
ActivityError {
kind,
message,
details: None,
}
}
#[async_trait]
impl OutboxRowDispatch for DeclaredBodyOutboxDispatch {
async fn dispatch(&self, row: &OutboxRow) -> Result<(), ServerError> {
let Some(run_id) = row.run_id.as_ref() else {
return self.inner.dispatch(row).await;
};
let run = DispatchingRun {
workflow_id: &row.workflow_id,
run_id,
};
match self
.bodies
.body_for(&row.task_queue, &row.activity_type, run)
{
DeclaredBodyLookup::None => self.inner.dispatch(row).await,
DeclaredBodyLookup::Unreadable(reason) => {
tracing::error!(
operation = "declared_command_outbox_dispatch",
workflow_id = %row.workflow_id,
activity_name = %row.activity_type,
task_queue = %row.task_queue,
%reason,
"declared-body catalog read failed; delegating the outbox row to the worker path"
);
self.inner.dispatch(row).await
}
DeclaredBodyLookup::Ambiguous { declaring } => {
let run_id = Self::run_id(row)?;
let token = self.mint_token(row, &run_id)?;
let refusal =
ambiguous_body_refusal(&row.activity_type, &row.task_queue, &declaring);
Self::complete(self.sink.as_ref(), row, run_id, token, Err(refusal))
}
DeclaredBodyLookup::Declared(contract) => self.place(row, contract),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aion_core::{ActivityId, Payload, RunId, WorkflowId};
use aion_package::ActionBodyContract;
use aion_store::{OutboxRow, OutboxStatus};
use async_trait::async_trait;
use super::super::declared_body::{DeclaredBodies, DeclaredBodyLookup, DispatchingRun};
use super::super::declared_body_cancel::DeclaredCommandAttempts;
use super::super::workspace_root::WorkspaceRoot;
use super::*;
type TestResult = Result<(), Box<dyn std::error::Error + Send + Sync>>;
struct RecordingInner {
rows: Arc<Mutex<Vec<String>>>,
}
#[async_trait]
impl OutboxRowDispatch for RecordingInner {
async fn dispatch(&self, row: &OutboxRow) -> Result<(), ServerError> {
self.rows
.lock()
.map_err(|_| ServerError::worker_dispatch("", "", "recording inner poisoned"))?
.push(row.dispatch_key.clone());
Err(ServerError::worker_dispatch(
&row.namespace,
&row.activity_type,
"no worker can currently serve this queue (NO_LIVE_POLLERS)",
))
}
}
struct FixedBodies {
lookup: DeclaredBodyLookup,
}
impl DeclaredBodies for FixedBodies {
fn body_for(
&self,
_task_queue: &str,
_action: &str,
_run: DispatchingRun<'_>,
) -> DeclaredBodyLookup {
self.lookup.clone()
}
}
#[derive(Default)]
struct RecordingSink {
completions: Arc<Mutex<Vec<ActivityCompletion>>>,
}
impl ActivityCompletionSink for RecordingSink {
fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
self.completions
.lock()
.map_err(|_| ServerError::worker_dispatch("", "", "recording sink poisoned"))?
.push(completion);
Ok(())
}
fn park_activity(
&self,
_workflow_id: &WorkflowId,
_activity_id: &ActivityId,
) -> Result<(), ServerError> {
Ok(())
}
}
const TRANSCRIPT_CAPACITY: std::num::NonZeroUsize = match std::num::NonZeroUsize::new(64) {
Some(capacity) => capacity,
None => unreachable!(),
};
fn executor() -> Arc<DeclaredCommandExecutor> {
let store: Arc<dyn aion_store::ObservabilityStore> =
Arc::new(aion_store::InMemoryObservabilityStore::default());
let transcript = crate::activity_publisher::ActivityEventPublisher::new(
store,
TRANSCRIPT_CAPACITY,
crate::activity_publisher::TranscriptBatchPolicy::UNBATCHED,
);
Arc::new(DeclaredCommandExecutor::new(
DeclaredCommandAttempts::new(crate::shutdown::DrainState::default()),
tokio::runtime::Handle::current(),
WorkspaceRoot::from_resolution(Ok(std::env::temp_dir())),
transcript,
))
}
fn row(ordinal: u64, activity_type: &str, input: &str) -> OutboxRow {
let workflow_id = WorkflowId::new(uuid::Uuid::new_v4());
OutboxRow {
dispatch_key: format!("{workflow_id}:{ordinal}"),
workflow_id,
ordinal,
run_id: Some(RunId::new(uuid::Uuid::new_v4())),
namespace: "default".to_owned(),
task_queue: "json_box".to_owned(),
node: None,
activity_type: activity_type.to_owned(),
input: Payload::new(ContentType::Json, input.as_bytes().to_vec()),
status: OutboxStatus::Claimed,
attempt: 0,
started_attempt: 1,
visible_after: chrono::Utc::now(),
claimed_at: Some(chrono::Utc::now()),
failure_delivered: false,
}
}
struct Rig {
dispatch: DeclaredBodyOutboxDispatch,
inner_rows: Arc<Mutex<Vec<String>>>,
completions: Arc<Mutex<Vec<ActivityCompletion>>>,
fences: CompletionFences,
}
fn rig(lookup: DeclaredBodyLookup) -> Rig {
let inner_rows = Arc::new(Mutex::new(Vec::new()));
let bodies = DeclaredBodySource::default();
bodies.install(Arc::new(FixedBodies { lookup }));
let sink = RecordingSink::default();
let completions = Arc::clone(&sink.completions);
let fences = CompletionFences::default();
let dispatch = DeclaredBodyOutboxDispatch::new(
Arc::new(RecordingInner {
rows: Arc::clone(&inner_rows),
}),
bodies,
executor(),
fences.clone(),
Arc::new(sink),
tokio::runtime::Handle::current(),
);
Rig {
dispatch,
inner_rows,
completions,
fences,
}
}
async fn await_completions(
completions: &Arc<Mutex<Vec<ActivityCompletion>>>,
count: usize,
) -> Result<Vec<ActivityCompletion>, Box<dyn std::error::Error + Send + Sync>> {
let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
loop {
let seen = completions
.lock()
.map_err(|_| "completions poisoned")?
.iter()
.map(|completion| ActivityCompletion {
workflow_id: completion.workflow_id.clone(),
activity_id: completion.activity_id.clone(),
run_id: completion.run_id.clone(),
completion_token: completion.completion_token.clone(),
outcome: completion.outcome.clone(),
})
.collect::<Vec<_>>();
if seen.len() >= count {
return Ok(seen);
}
if tokio::time::Instant::now() >= deadline {
return Err(format!("expected {count} completion(s), saw {}", seen.len()).into());
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_declared_body_row_executes_at_the_server_and_completes_through_the_sink()
-> TestResult {
let rig = rig(DeclaredBodyLookup::Declared(ActionBodyContract::Run {
command: "printf 'sha=%s' {{sha}}".to_owned(),
}));
let row = row(3, "size_of", r#"{"sha":"abc123"}"#);
rig.dispatch.dispatch(&row).await?;
assert!(
rig.inner_rows
.lock()
.map_err(|_| "inner poisoned")?
.is_empty(),
"a declared body must never reach the worker path"
);
let completions = await_completions(&rig.completions, 1).await?;
let completion = &completions[0];
assert_eq!(completion.workflow_id, row.workflow_id);
assert_eq!(
completion.activity_id,
ActivityId::from_sequence_position(3)
);
assert_eq!(completion.run_id, row.run_id);
let ActivityCompletionOutcome::Succeeded(payload) = &completion.outcome else {
return Err(format!("expected a success, got {:?}", completion.outcome).into());
};
let encoded = std::str::from_utf8(payload.bytes())?;
let outcome: serde_json::Value = serde_json::from_str(encoded)?;
assert_eq!(
outcome["stdout"], "sha=abc123",
"the body ran with the row's input bound"
);
assert_eq!(outcome["exit_code"], 0);
rig.fences.accept(
&row.workflow_id,
&completion.activity_id,
&completion.completion_token,
)?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_row_without_a_body_goes_to_the_worker_path_untouched() -> TestResult {
let rig = rig(DeclaredBodyLookup::None);
let row = row(1, "size_of", r#"{"sha":"abc123"}"#);
let refused = rig.dispatch.dispatch(&row).await;
assert!(
refused.is_err(),
"the inner transport's refusal must surface as the row's"
);
assert_eq!(
*rig.inner_rows.lock().map_err(|_| "inner poisoned")?,
vec![row.dispatch_key.clone()]
);
assert!(
rig.completions
.lock()
.map_err(|_| "completions poisoned")?
.is_empty()
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn an_unreadable_catalog_delegates_to_the_worker_path() -> TestResult {
let rig = rig(DeclaredBodyLookup::Unreadable("catalog offline".to_owned()));
let row = row(1, "size_of", r#"{"sha":"abc123"}"#);
let _ = rig.dispatch.dispatch(&row).await;
assert_eq!(
*rig.inner_rows.lock().map_err(|_| "inner poisoned")?,
vec![row.dispatch_key.clone()]
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn an_ambiguous_body_is_a_terminal_failure_delivered_to_the_workflow() -> TestResult {
let rig = rig(DeclaredBodyLookup::Ambiguous {
declaring: Vec::new(),
});
let row = row(2, "size_of", r#"{"sha":"abc123"}"#);
rig.dispatch.dispatch(&row).await?;
assert!(
rig.inner_rows
.lock()
.map_err(|_| "inner poisoned")?
.is_empty()
);
let completions = await_completions(&rig.completions, 1).await?;
let ActivityCompletionOutcome::Failed(error) = &completions[0].outcome else {
return Err("expected a failure".into());
};
assert_eq!(error.kind, ActivityErrorKind::Terminal);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_failing_body_completes_as_the_executors_retryable_failure() -> TestResult {
let rig = rig(DeclaredBodyLookup::Declared(ActionBodyContract::Run {
command: "sh -c 'exit 7'".to_owned(),
}));
let row = row(4, "size_of", "{}");
rig.dispatch.dispatch(&row).await?;
let completions = await_completions(&rig.completions, 1).await?;
let ActivityCompletionOutcome::Failed(error) = &completions[0].outcome else {
return Err(format!("expected a failure, got {:?}", completions[0].outcome).into());
};
assert_eq!(error.kind, ActivityErrorKind::Retryable);
assert!(
error.message.contains("exited 7"),
"the executor's own sentence must travel: {}",
error.message
);
Ok(())
}
#[test]
fn failure_classification_follows_the_executor_prefixes() {
assert_eq!(
classify_failure("terminal:x").kind,
ActivityErrorKind::Terminal
);
assert_eq!(classify_failure("terminal:x").message, "x");
assert_eq!(
classify_failure("policy_refused:y").kind,
ActivityErrorKind::PolicyRefused
);
assert_eq!(
classify_failure("retryable:z").kind,
ActivityErrorKind::Retryable
);
let timeout = classify_failure("timeout:attempt outlived its bound");
assert_eq!(timeout.kind, ActivityErrorKind::Retryable);
assert_eq!(timeout.message, "timeout:attempt outlived its bound");
}
}