use std::collections::BTreeMap;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use aion::{ActivityDispatch, ActivityDispatcher, LeaseSignal};
use aion_core::{ActivityId, RunId, WorkflowId};
use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
use dashmap::DashMap;
use dashmap::mapref::entry::Entry;
use super::dispatch::{ActivityCompletion, ActivityCompletionSink};
use super::envelope::{CompletionFences, CompletionToken, idempotency_key};
use super::heartbeat::{HeartbeatTracker, InFlightActivity};
use super::queue_service::{
DeliveryRefusal, ExpiredClock, PARK_POLL_INTERVAL, PoolCensus, QueueDeclarationSource,
QueueServiceConfig, QueueServiceReason, QueueServiceState, SelectionRefusal, ServiceAddress,
ServiceWait, WorkerUnavailable, deliver_within_schedule_to_start, select_worker_or_refuse,
};
use super::registry::{
ConnectedWorkerRegistry, DispatchReservation, WorkerArrival, WorkerDelivery, WorkerHandle,
WorkerId, WorkerMessage,
};
use crate::error::ServerError;
use crate::shutdown::DrainState;
mod completion_outcome;
use tracing::info_span;
type SyncSender = std::sync::mpsc::SyncSender<Result<String, String>>;
struct PendingWaiter {
sender: SyncSender,
attempt: u32,
}
fn superseded_site_reason(new_attempt: u32) -> String {
format!(
"retryable: superseded by attempt {new_attempt} at this execution site — the prior \
attempt reached its bound, its completion token is fence-refused, and its worker \
has been asked to stop"
)
}
type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
type PendingActivityKey = (WorkflowId, ActivityId);
pub trait OutboxDeliveryCallback: Send + Sync {
fn deliver_completion(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
result: String,
) -> Result<bool, ServerError>;
fn deliver_failure(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
reason: String,
) -> Result<bool, ServerError>;
}
#[derive(Clone)]
pub struct PendingActivities {
pending: Arc<DashMap<PendingActivityKey, PendingWaiter>>,
completion_fences: CompletionFences,
outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
lease_recorder: super::lease_record::LeaseRecorderSeam,
transport_losses: super::transport_loss::TransportLossLedger,
}
impl std::fmt::Debug for PendingActivities {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("PendingActivities")
.field("pending", &self.pending.len())
.field("completion_fences", &self.completion_fences)
.field(
"outbox_delivery_installed",
&self.outbox_delivery.get().is_some(),
)
.field("lease_recorder", &self.lease_recorder)
.field("transport_losses", &self.transport_losses)
.finish()
}
}
impl PendingActivities {
fn insert(
&self,
workflow_id: WorkflowId,
run_id: &RunId,
activity_id: ActivityId,
attempt: u32,
) -> Result<(CompletionToken, SyncReceiver, Option<u32>), ServerError> {
let completion_token =
self.completion_fences
.issue(&workflow_id, run_id, &activity_id, attempt)?;
let (tx, rx) = std::sync::mpsc::sync_channel(1);
let key = (workflow_id, activity_id);
match self.pending.entry(key) {
Entry::Vacant(entry) => {
entry.insert(PendingWaiter {
sender: tx,
attempt,
});
Ok((completion_token, rx, None))
}
Entry::Occupied(mut entry) => {
let held_attempt = entry.get().attempt;
if attempt > held_attempt {
let superseded = entry.insert(PendingWaiter {
sender: tx,
attempt,
});
let _ = superseded.sender.send(Err(superseded_site_reason(attempt)));
Ok((completion_token, rx, Some(held_attempt)))
} else {
let (workflow_id, activity_id) = entry.key();
Err(ServerError::PendingActivityCollision {
workflow_id: workflow_id.clone(),
activity_id: activity_id.clone(),
held_attempt,
incoming_attempt: attempt,
})
}
}
}
}
#[must_use]
pub fn completion_fences(&self) -> CompletionFences {
self.completion_fences.clone()
}
#[cfg(test)]
pub(crate) fn insert_for_test(
&self,
workflow_id: WorkflowId,
run_id: &RunId,
activity_id: ActivityId,
attempt: u32,
) -> Result<(CompletionToken, SyncReceiver, Option<u32>), ServerError> {
self.insert(workflow_id, run_id, activity_id, attempt)
}
pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
if self.outbox_delivery.set(callback).is_err() {
tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
}
}
fn complete(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
result: Result<String, String>,
) -> bool {
let matched = self
.pending
.remove(&(workflow_id.clone(), activity_id.clone()));
if let Some((_, waiter)) = matched {
return waiter.sender.send(result).is_ok();
}
let Some(callback) = self.outbox_delivery.get() else {
return false;
};
let outcome = match result {
Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
};
match outcome {
Ok(true) => true,
Ok(false) => {
tracing::debug!(
workflow_id = %workflow_id,
activity_id = %activity_id,
"unmatched outbox completion for a workflow that is not currently live; \
recovery will re-arm it"
);
false
}
Err(error) => {
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"failed to deliver unmatched outbox completion to the live workflow"
);
false
}
}
}
fn complete_fenced_after_accept(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
completion_token: &CompletionToken,
result: Result<String, String>,
after_accept: impl FnOnce() -> Result<bool, ServerError>,
) -> Result<bool, ServerError> {
let accepted = self
.completion_fences
.accept(workflow_id, activity_id, completion_token)
.inspect_err(|error| {
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"activity completion rejected by execution-generation fence"
);
})?;
let should_publish = match after_accept() {
Ok(should_publish) => should_publish,
Err(error) => {
if let Err(restore_error) =
self.completion_fences
.restore_if_absent(workflow_id, activity_id, &accepted)
{
tracing::error!(%restore_error, "failed to restore completion generation after settlement failure");
}
return Err(error);
}
};
if !should_publish {
if let Err(restore_error) =
self.completion_fences
.restore_if_absent(workflow_id, activity_id, &accepted)
{
tracing::error!(%restore_error, "failed to restore completion generation after a non-publishing settlement");
}
return Ok(false);
}
let transport_domain = result
.as_ref()
.err()
.is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
if !transport_domain
&& let Err(error) = self.transport_losses.clear(workflow_id, activity_id)
{
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"failed to retire the transport-loss budget for a resolved activity"
);
}
Ok(self.complete(workflow_id, activity_id, run_id, result))
}
#[must_use]
pub fn new(heartbeat_window: std::time::Duration) -> Self {
Self {
pending: Arc::default(),
completion_fences: CompletionFences::default(),
outbox_delivery: Arc::default(),
lease_recorder: super::lease_record::LeaseRecorderSeam::default(),
transport_losses: super::transport_loss::TransportLossLedger::new(heartbeat_window),
}
}
#[must_use]
pub fn lease_recorder(&self) -> super::lease_record::LeaseRecorderSeam {
self.lease_recorder.clone()
}
pub fn set_lease_recorder(
&self,
recorder: Arc<dyn super::lease_record::ActivityLeaseRecorder>,
metrics: Option<crate::observability::Metrics>,
) {
self.lease_recorder.install(recorder, metrics);
}
}
impl ActivityCompletionSink for PendingActivities {
fn park_activity(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
) -> Result<(), ServerError> {
self.completion_fences
.revoke_current(workflow_id, activity_id)?;
let matched = self
.pending
.remove(&(workflow_id.clone(), activity_id.clone()));
if let Some((_, waiter)) = matched {
let _ = waiter
.sender
.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
}
Ok(())
}
fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
self.complete_activity_after_accept(completion, || Ok(()))
}
}
pub struct WorkerActivityDispatcher {
registry: ConnectedWorkerRegistry,
namespace: String,
pending: PendingActivities,
heartbeat_tracker: HeartbeatTracker,
drain_state: DrainState,
tokio_handle: Option<tokio::runtime::Handle>,
attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
queue_service: QueueServiceConfig,
queue_declarations: QueueDeclarationSource,
queue_state: QueueServiceState,
cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
}
impl std::fmt::Debug for WorkerActivityDispatcher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WorkerActivityDispatcher")
.field("namespace", &self.namespace)
.finish_non_exhaustive()
}
}
impl WorkerActivityDispatcher {
#[must_use]
pub fn new(
registry: ConnectedWorkerRegistry,
namespace: impl Into<String>,
heartbeat_tracker: HeartbeatTracker,
) -> Self {
Self {
registry,
namespace: namespace.into(),
pending: PendingActivities::new(heartbeat_tracker.heartbeat_window()),
heartbeat_tracker,
drain_state: DrainState::default(),
tokio_handle: None,
attempt_owners: None,
queue_service: QueueServiceConfig::default(),
queue_declarations: QueueDeclarationSource::default(),
queue_state: QueueServiceState::default(),
cluster_publisher: None,
}
}
#[must_use]
pub fn with_cluster_publisher(
mut self,
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
) -> Self {
self.cluster_publisher = Some(cluster_publisher);
self
}
#[must_use]
pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
self.queue_service = queue_service;
self
}
#[must_use]
pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
self.queue_declarations = queue_declarations;
self
}
#[must_use]
pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
self.queue_state = queue_state;
self
}
#[must_use]
pub fn with_attempt_owners(
mut self,
attempt_owners: super::intervention::AttemptOwnerIndex,
) -> Self {
self.attempt_owners = Some(attempt_owners);
self
}
#[must_use]
pub const fn pending(&self) -> &PendingActivities {
&self.pending
}
#[must_use]
pub fn with_pending(mut self, pending: PendingActivities) -> Self {
self.pending = pending;
self
}
#[must_use]
pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
self.drain_state = drain_state;
self
}
#[must_use]
pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
self.tokio_handle = Some(tokio_handle);
self
}
}
impl WorkerActivityDispatcher {
fn ensure_accepting(
&self,
namespace: &str,
activity_type: &str,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
worker_id: Option<WorkerId>,
) -> Result<(), String> {
self.drain_state
.ensure_accepting(namespace, activity_type)
.map_err(|error| {
log_worker_error(
"WorkerDispatch",
namespace,
activity_type,
workflow_id,
activity_id,
worker_id,
&error.to_string(),
);
aion::PARKED_ACTIVITY_REASON.to_owned()
})
}
fn select_worker_or_wait(
&self,
address: &ServiceAddress,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
) -> Result<(WorkerHandle, DispatchReservation), String> {
let wait = ServiceWait {
registry: &self.registry,
declarations: &self.queue_declarations,
config: &self.queue_service,
state: &self.queue_state,
address,
workflow_id,
activity_id,
publisher: self.cluster_publisher.as_ref(),
};
let mut accepting = || {
self.ensure_accepting(
&address.namespace,
&address.activity_type,
workflow_id,
activity_id,
None,
)
};
let mut park = |arrival: WorkerArrival, budget: Option<Duration>| {
self.park_for_worker(arrival, budget);
};
select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
let reason = refusal.reason_string();
if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
let error_type = match refusal {
SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
_ => "WorkerRegistry",
};
log_worker_error(
error_type,
&address.namespace,
&address.activity_type,
workflow_id,
activity_id,
None,
&reason,
);
}
reason
})
}
fn park_for_worker(&self, arrival: WorkerArrival, budget: Option<Duration>) {
let handle = self
.tokio_handle
.clone()
.or_else(|| tokio::runtime::Handle::try_current().ok());
let Some(handle) = handle else {
std::thread::sleep(
budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
);
return;
};
handle.block_on(async {
let raced = async {
tokio::select! {
() = arrival => {}
() = self.drain_state.wait_for_drain() => {}
}
};
match budget {
None => raced.await,
Some(budget) => {
drop(tokio::time::timeout(budget, raced).await);
}
}
});
}
fn track_worker_task(
&self,
worker_id: WorkerId,
activity_type: &str,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
attempt: u32,
lease: (CompletionToken, DispatchReservation),
) -> Result<(), String> {
let (completion_token, reservation) = lease;
self.heartbeat_tracker
.track_task(
worker_id,
InFlightActivity {
workflow_id: workflow_id.clone(),
activity_id: activity_id.clone(),
attempt,
completion_token,
},
Instant::now(),
&self.registry,
Some(reservation),
)
.map_err(|error| {
let reason = error.to_string();
log_worker_error(
"WorkerHeartbeatTracker",
&self.namespace,
activity_type,
workflow_id,
activity_id,
Some(worker_id),
&reason,
);
reason
})
}
fn cleanup_activity(
&self,
worker_id: WorkerId,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
attempt: u32,
completion_token: &CompletionToken,
) {
self.pending
.pending
.remove_if(&(workflow_id.clone(), activity_id.clone()), |_, waiter| {
waiter.attempt == attempt
});
if let Err(error) =
self.pending
.completion_fences
.revoke(workflow_id, activity_id, completion_token)
{
tracing::error!(
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"failed to revoke undelivered activity generation"
);
}
let _ = self.heartbeat_tracker.complete_task(
worker_id,
workflow_id,
activity_id,
&self.registry,
);
self.drain_state.notify_activity_drained();
}
fn send_activity_task(
&self,
worker: &WorkerHandle,
task: ProtoActivityTask,
address: &ServiceAddress,
handoff: &super::lease_record::LeaseHandoff,
) -> Result<(), String> {
let workflow_id = &handoff.key().workflow_id;
let activity_id = &handoff.key().activity_id;
let completion_token = handoff.token();
match worker.delivery() {
WorkerDelivery::Grpc(sender) => {
let worker_id = worker.id();
let mut accepting = || {
self.ensure_accepting(
&address.namespace,
&address.activity_type,
workflow_id,
activity_id,
Some(worker_id),
)
};
let handed_over = deliver_within_schedule_to_start(
sender,
WorkerMessage::ActivityTask(Box::new(task)),
self.queue_service.schedule_to_start_timeout,
&mut accepting,
);
let Err(refusal) = handed_over else {
handoff.accepted_blocking(self.tokio_handle.as_ref());
return Ok(());
};
self.cleanup_activity(
worker_id,
workflow_id,
activity_id,
handoff.key().attempt,
completion_token,
);
let (error_type, reason) = self.hand_off_failure(&refusal, address);
if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
log_worker_error(
error_type,
&address.namespace,
&address.activity_type,
workflow_id,
activity_id,
Some(worker_id),
&reason,
);
}
Err(reason)
}
#[cfg(feature = "liminal-transport")]
WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
worker.id(),
delivery,
task,
&address.activity_type,
handoff,
),
}
}
fn hand_off_failure(
&self,
refusal: &DeliveryRefusal,
address: &ServiceAddress,
) -> (&'static str, String) {
match refusal {
DeliveryRefusal::Saturated { waited } => {
let census = self
.registry
.pool_census(
&address.namespace,
&address.task_queue,
&address.activity_type,
address.node.as_deref(),
)
.unwrap_or_else(|error| {
tracing::error!(
namespace = %address.namespace,
task_queue = %address.task_queue,
activity_type = %address.activity_type,
%error,
"poller census failed while reporting a saturated queue; \
the refusal carries an empty census"
);
PoolCensus::default()
});
let unavailable = WorkerUnavailable {
reason: QueueServiceReason::Saturated,
clock: Some(ExpiredClock::ScheduleToStart),
waited: *waited,
address: address.clone(),
census,
};
("WorkerUnavailable", unavailable.reason_string())
}
DeliveryRefusal::Full => (
"WorkerChannelClosed",
"worker task channel full or closed: no available capacity".to_owned(),
),
DeliveryRefusal::Closed => (
"WorkerChannelClosed",
"worker task channel full or closed: channel closed".to_owned(),
),
DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
}
}
#[cfg(feature = "liminal-transport")]
fn send_liminal_activity_task(
&self,
worker_id: WorkerId,
delivery: &super::liminal_transport::LiminalWorkerDelivery,
task: ProtoActivityTask,
activity_type: &str,
handoff: &super::lease_record::LeaseHandoff,
) -> Result<(), String> {
let workflow_id = &handoff.key().workflow_id;
let activity_id = &handoff.key().activity_id;
let completion_token =
CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
.map_err(|error| error.to_string())?;
let heartbeat_window_ms =
u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
.unwrap_or(u64::MAX);
let attempt = task.attempt;
let run_id = task
.run_id
.map(RunId::try_from)
.transpose()
.map_err(|error| error.to_string())?
.ok_or_else(|| {
"activity task run id is missing; refusing to dispatch an unidentified run"
.to_owned()
})?;
let request = super::liminal_transport::DispatchRequest {
activity_type: activity_type.to_owned(),
workflow_id: workflow_id.clone(),
ordinal: activity_id.sequence_position(),
run_id: Some(run_id.clone()),
attempt,
completion_token: task.completion_token,
idempotency_key: task.idempotency_key,
labels: task.labels.into_iter().collect(),
heartbeat_window_ms,
input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
};
let awaiter = match delivery.push_dispatch(&request) {
Ok(awaiter) => awaiter,
Err(error) => {
let reason = format!("worker liminal push failed: {error}");
self.cleanup_activity(
worker_id,
workflow_id,
activity_id,
attempt,
&completion_token,
);
log_worker_error(
"WorkerChannelClosed",
&self.namespace,
activity_type,
workflow_id,
activity_id,
Some(worker_id),
&reason,
);
return Err(reason);
}
};
handoff.accepted_blocking(self.tokio_handle.as_ref());
let owner_binding = self.attempt_owners.as_ref().map(|owners| {
super::liminal_transport::AttemptOwnerGuard::bind(
owners.clone(),
super::intervention::AttemptKey::new(
workflow_id.clone(),
run_id.clone(),
activity_id.clone(),
attempt,
),
worker_id,
)
});
self.spawn_liminal_reply_router(
worker_id,
awaiter,
workflow_id,
activity_id,
&completion_token,
owner_binding,
);
Ok(())
}
#[cfg(feature = "liminal-transport")]
fn spawn_liminal_reply_router(
&self,
worker_id: WorkerId,
awaiter: liminal_server::server::connection::PushReplyAwaiter,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
completion_token: &CompletionToken,
owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
) {
let pending = self.pending.clone();
let heartbeat_tracker = self.heartbeat_tracker.clone();
let registry = self.registry.clone();
let drain_state = self.drain_state.clone();
let workflow_id = workflow_id.clone();
let activity_id = activity_id.clone();
let completion_token = completion_token.clone();
std::thread::spawn(move || {
let _owner_binding = owner_binding;
route_liminal_reply(
&pending,
&heartbeat_tracker,
®istry,
&drain_state,
&awaiter,
(worker_id, &workflow_id, &activity_id, &completion_token),
);
});
}
fn await_activity_result(
&self,
context: &ActivityDispatchContext<'_>,
rx: &SyncReceiver,
) -> Result<String, String> {
match self.registry.is_registered(context.worker_id) {
Ok(true) => {}
Ok(false) => {
if let Ok(result) = rx.try_recv() {
return self.deliver_result(context, result);
}
self.cleanup_activity(
context.worker_id,
context.workflow_id,
context.activity_id,
context.attempt,
&context.completion_token,
);
let reason = self.pending.classify_worker_loss(
context.workflow_id,
context.activity_id,
context.worker_id,
);
log_worker_error(
"WorkerLost",
&self.namespace,
context.activity_type,
context.workflow_id,
context.activity_id,
Some(context.worker_id),
&reason,
);
return Err(reason);
}
Err(error) => {
self.cleanup_activity(
context.worker_id,
context.workflow_id,
context.activity_id,
context.attempt,
&context.completion_token,
);
let reason = format!("worker registry inspection failed: {error}");
log_worker_error(
"WorkerRegistry",
&self.namespace,
context.activity_type,
context.workflow_id,
context.activity_id,
Some(context.worker_id),
&reason,
);
return Err(reason);
}
}
if let Ok(result) = rx.recv() {
return self.deliver_result(context, result);
}
self.cleanup_activity(
context.worker_id,
context.workflow_id,
context.activity_id,
context.attempt,
&context.completion_token,
);
let reason = "activity response channel dropped".to_owned();
log_worker_error(
"WorkerChannelClosed",
&self.namespace,
context.activity_type,
context.workflow_id,
context.activity_id,
Some(context.worker_id),
&reason,
);
Err(reason)
}
fn deliver_result(
&self,
context: &ActivityDispatchContext<'_>,
result: Result<String, String>,
) -> Result<String, String> {
self.pending.pending.remove_if(
&(context.workflow_id.clone(), context.activity_id.clone()),
|_, waiter| waiter.attempt == context.attempt,
);
if let Err(reason) = &result
&& aion::is_parked_reason(reason)
{
tracing::info!(
operation = "activity_dispatch",
namespace = %self.namespace,
workflow_id = %context.workflow_id,
activity_id = %context.activity_id,
activity_type = context.activity_type,
worker_id = ?context.worker_id,
"activity parked for restart recovery"
);
return result;
}
log_activity_completion(context, result.is_ok());
result.inspect_err(|reason| {
log_worker_error(
"ActivityFailed",
&self.namespace,
context.activity_type,
context.workflow_id,
context.activity_id,
Some(context.worker_id),
reason,
);
})
}
}
impl ActivityDispatcher for WorkerActivityDispatcher {
fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
match tokio::runtime::Handle::try_current() {
Ok(handle) => match handle.runtime_flavor() {
tokio::runtime::RuntimeFlavor::MultiThread => {
tokio::task::block_in_place(|| {
self.dispatch_blocking(request, &LeaseSignal::none())
})
}
flavor => Err(format!(
"activity dispatch blocks the calling thread until the worker responds; \
a {flavor:?} tokio runtime cannot host that wait because the worker \
stream forwarder shares its only executor thread and the task could \
never be delivered — run the engine on a multi-thread tokio runtime"
)),
},
Err(_) => self.dispatch_blocking(request, &LeaseSignal::none()),
}
}
fn dispatch_async(
self: std::sync::Arc<Self>,
request: ActivityDispatch,
lease: LeaseSignal,
) -> futures::future::BoxFuture<'static, Result<String, String>> {
Box::pin(async move {
let blocking =
tokio::task::spawn_blocking(move || self.dispatch_blocking(request, &lease));
match blocking.await {
Ok(result) => result,
Err(join_error) => Err(format!("activity dispatch task failed: {join_error}")),
}
})
}
}
impl WorkerActivityDispatcher {
fn dispatch_blocking(
&self,
request: ActivityDispatch,
lease: &LeaseSignal,
) -> Result<String, String> {
let ActivityDispatch {
namespace,
task_queue,
node,
workflow_id,
run_id,
activity_id,
name,
input,
config: _,
attempt,
labels,
advisory: _,
} = request;
let started_at = Instant::now();
self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
let address = ServiceAddress {
namespace: namespace.clone(),
task_queue: task_queue.clone(),
activity_type: name.clone(),
node: node.clone(),
};
let (worker, reservation) =
self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
let worker_id = worker.id();
let span = info_span!(
"activity_dispatch",
operation = "activity_dispatch",
namespace = %namespace,
task_queue = %task_queue,
node = node.as_deref(),
workflow_id = %workflow_id,
activity_id = %activity_id,
activity_type = %name,
worker_id = ?worker_id,
);
let _span_guard = span.enter();
self.ensure_accepting(
&namespace,
&name,
&workflow_id,
&activity_id,
Some(worker_id),
)?;
let (completion_token, rx, superseded_attempt) = self
.pending
.insert(workflow_id.clone(), &run_id, activity_id.clone(), attempt)
.map_err(|error| error.to_string())?;
if let Some(held_attempt) = superseded_attempt {
self.release_superseded_attempt(&workflow_id, &activity_id, held_attempt, attempt);
}
let handoff = super::lease_record::LeaseHandoff::arm(
self.pending.lease_recorder(),
super::lease_record::LeaseKey {
workflow_id: workflow_id.clone(),
run_id: run_id.clone(),
activity_id: activity_id.clone(),
attempt,
},
super::lease_record::attribution_for(&worker),
self.pending.completion_fences(),
completion_token.clone(),
lease.clone(),
)
.map_err(|error| error.to_string())?;
let task = activity_task(
&name,
&input,
(&workflow_id, &run_id, &activity_id),
attempt,
labels,
&completion_token,
);
self.track_then_send(
(worker_id, &worker),
task,
&address,
(&name, &workflow_id, &activity_id, attempt),
(&completion_token, &handoff),
reservation,
)?;
let context = ActivityDispatchContext {
namespace: &namespace,
activity_type: &name,
worker_id,
workflow_id: &workflow_id,
activity_id: &activity_id,
attempt,
completion_token,
started_at,
};
self.await_activity_result(&context, &rx)
}
fn track_then_send(
&self,
worker: (WorkerId, &WorkerHandle),
task: ProtoActivityTask,
address: &ServiceAddress,
site: (&str, &WorkflowId, &ActivityId, u32),
lease: (&CompletionToken, &super::lease_record::LeaseHandoff),
reservation: DispatchReservation,
) -> Result<(), String> {
let (worker_id, handle) = worker;
let (name, workflow_id, activity_id, attempt) = site;
let (completion_token, handoff) = lease;
if let Err(error) = self.track_worker_task(
worker_id,
name,
workflow_id,
activity_id,
attempt,
(completion_token.clone(), reservation),
) {
self.cleanup_activity(
worker_id,
workflow_id,
activity_id,
attempt,
completion_token,
);
return Err(error);
}
self.send_activity_task(handle, task, address, handoff)
}
fn release_superseded_attempt(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
held_attempt: u32,
new_attempt: u32,
) {
let in_flight = match self.heartbeat_tracker.in_flight_for_workflow(workflow_id) {
Ok(in_flight) => in_flight,
Err(error) => {
tracing::error!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
new_attempt,
%error,
"superseded-attempt release could not read the in-flight tracker; \
the corpse's worker was not asked to stop"
);
return;
}
};
let holders = in_flight
.into_iter()
.filter(|task| &task.activity_id == activity_id && task.attempt == held_attempt);
let mut asked_anyone = false;
for task in holders {
asked_anyone = true;
let delivery = match super::activity_cancel::ask_worker_to_stop(
&self.registry,
task.worker_id,
workflow_id,
activity_id,
) {
Ok(delivery) => Some(delivery),
Err(error) => {
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
worker_id = task.worker_id.value(),
%error,
"superseded-attempt release could not ask the holding worker to stop; \
untracking the corpse anyway so its capacity slot is not leaked"
);
None
}
};
if let Err(error) = self.heartbeat_tracker.complete_task(
task.worker_id,
workflow_id,
activity_id,
&self.registry,
) {
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
worker_id = task.worker_id.value(),
%error,
"superseded-attempt release could not untrack the corpse's heartbeat task"
);
}
let Some(delivery) = delivery else {
continue;
};
if delivery.was_requested() {
tracing::info!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
new_attempt,
worker_id = task.worker_id.value(),
"attempt takeover: asked the worker still executing the superseded \
attempt to stop"
);
} else {
tracing::warn!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
new_attempt,
worker_id = task.worker_id.value(),
outcome = ?delivery,
"attempt takeover: the superseded attempt's worker could NOT be asked \
to stop — its work runs on unwitnessed until it next reports, and its \
completion will be refused as a stale generation"
);
}
}
if !asked_anyone {
tracing::info!(
workflow_id = %workflow_id,
activity_id = %activity_id,
held_attempt,
new_attempt,
"attempt takeover: no tracked worker holds the superseded attempt \
(already untracked, or its transport records no heartbeat task); \
its completion token is fence-refused regardless"
);
}
}
}
#[cfg(feature = "liminal-transport")]
fn synthesized_failure(
current_token: &CompletionToken,
reason: String,
) -> (Option<RunId>, CompletionToken, Result<String, String>, bool) {
(None, current_token.clone(), Err(reason), true)
}
#[cfg(feature = "liminal-transport")]
fn route_liminal_reply(
pending: &PendingActivities,
heartbeat_tracker: &HeartbeatTracker,
registry: &crate::worker::ConnectedWorkerRegistry,
drain_state: &DrainState,
awaiter: &liminal_server::server::connection::PushReplyAwaiter,
execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
) {
let (worker_id, workflow_id, activity_id, current_token) = execution;
let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
heartbeat_tracker
.is_tracked(worker_id, workflow_id, activity_id)
.unwrap_or(false)
});
let (run_id, submitted_token, outcome, synthesized) = match waited {
Ok(Some(response)) => {
let submitted_token = match CompletionToken::from_wire(
workflow_id,
activity_id,
response.completion_token,
) {
Ok(token) => token,
Err(error) => {
tracing::warn!(
worker_id = ?worker_id,
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"liminal activity completion omitted its generation proof"
);
return;
}
};
(response.run_id, submitted_token, response.outcome, false)
}
Ok(None) => {
tracing::debug!(
worker_id = ?worker_id,
workflow_id = %workflow_id,
activity_id = %activity_id,
"liminal dispatch resolved by another path; abandoning reply wait"
);
return;
}
Err(error) if error.is_worker_dispatch_unservable() => {
synthesized_failure(current_token, format!("terminal:{error}"))
}
Err(error) if error.is_worker_connection_lost() => synthesized_failure(
current_token,
pending.classify_worker_loss(workflow_id, activity_id, worker_id),
),
Err(error) => synthesized_failure(
current_token,
format!("retryable:worker liminal reply failed: {error}"),
),
};
let after_accept = || {
if synthesized {
let was_tracked = clear_completed_task_tracking(
heartbeat_tracker,
registry,
worker_id,
workflow_id,
activity_id,
);
if !was_tracked {
tracing::debug!(
worker_id = ?worker_id,
workflow_id = %workflow_id,
activity_id = %activity_id,
"liminal dispatch already resolved; dropping synthesized lost-worker failure"
);
return Ok(false);
}
} else {
let _ = clear_completed_task_tracking(
heartbeat_tracker,
registry,
worker_id,
workflow_id,
activity_id,
);
}
drain_state.notify_activity_drained();
Ok(true)
};
if let Err(error) = pending.complete_fenced_after_accept(
workflow_id,
activity_id,
run_id.as_ref(),
&submitted_token,
outcome,
after_accept,
) {
tracing::warn!(
worker_id = ?worker_id,
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"liminal activity completion handoff rejected"
);
}
}
pub(crate) fn clear_completed_task_tracking(
heartbeat_tracker: &HeartbeatTracker,
registry: &crate::worker::ConnectedWorkerRegistry,
worker_id: WorkerId,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
) -> bool {
heartbeat_tracker
.complete_task(worker_id, workflow_id, activity_id, registry)
.unwrap_or_else(|error| {
tracing::error!(
worker_id = ?worker_id,
workflow_id = %workflow_id,
activity_id = %activity_id,
%error,
"failed to clear in-flight tracking for a completed activity; delivering anyway — a poisoned tracker must not withhold a result"
);
true
})
}
struct ActivityDispatchContext<'a> {
namespace: &'a str,
activity_type: &'a str,
worker_id: WorkerId,
workflow_id: &'a WorkflowId,
activity_id: &'a ActivityId,
attempt: u32,
completion_token: CompletionToken,
started_at: Instant,
}
fn activity_task(
activity_type: &str,
input: &str,
execution: (&WorkflowId, &RunId, &ActivityId),
attempt: u32,
labels: BTreeMap<String, String>,
completion_token: &CompletionToken,
) -> ProtoActivityTask {
let (workflow_id, run_id, activity_id) = execution;
ProtoActivityTask {
workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
activity_id: Some(ProtoActivityId::from(activity_id.clone())),
activity_type: activity_type.to_owned(),
input: Some(ProtoPayload {
content_type: String::from("application/json"),
bytes: input.as_bytes().to_vec(),
}),
attempt,
labels: labels.into_iter().collect(),
run_id: Some(run_id.clone().into()),
completion_token: completion_token.as_str().to_owned(),
idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
}
}
fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
let duration_ms = duration_ms(context.started_at.elapsed());
tracing::info!(
operation = "activity_complete",
namespace = context.namespace,
workflow_id = %context.workflow_id,
activity_id = %context.activity_id,
activity_type = context.activity_type,
worker_id = ?context.worker_id,
duration_ms,
outcome = if succeeded { "succeeded" } else { "failed" },
"activity completed"
);
}
fn duration_ms(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn log_worker_error(
error_type: &'static str,
namespace: &str,
activity_type: &str,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
worker_id: Option<super::registry::WorkerId>,
reason: &str,
) {
tracing::error!(
operation = "activity_dispatch",
namespace,
workflow_id = %workflow_id,
activity_id = %activity_id,
activity_type,
worker_id = ?worker_id,
error_type,
reason,
"worker interaction failed"
);
}
#[cfg(test)]
mod tests {
use std::sync::{
Mutex,
atomic::{AtomicBool, Ordering},
};
use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
use crate::worker::dispatch::ActivityCompletionOutcome;
use super::*;
fn activity_id(pos: u64) -> ActivityId {
ActivityId::from_sequence_position(pos)
}
fn insert_attempt_one(
pending: &PendingActivities,
workflow_id: &WorkflowId,
id: &ActivityId,
) -> Result<(CompletionToken, SyncReceiver), ServerError> {
let (token, rx, superseded) =
pending.insert(workflow_id.clone(), &RunId::new_v4(), id.clone(), 1)?;
assert!(
superseded.is_none(),
"a first insert must never report a superseded holder"
);
Ok((token, rx))
}
#[test]
fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(1);
let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
assert_eq!(
rx.recv_timeout(Duration::from_millis(50)),
Ok(Ok("done".to_owned()))
);
Ok(())
}
#[test]
fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(2);
let run_id = RunId::new_v4();
let _first = pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
let error = pending
.insert(workflow_id, &run_id, id, 1)
.err()
.ok_or_else(|| {
ServerError::worker_dispatch(
"test",
"same-worker-collision",
"a live responder collision was silently replaced",
)
})?;
assert!(
error.to_string().contains("pending activity collision"),
"{error}"
);
Ok(())
}
#[test]
fn a_higher_attempt_takes_over_a_stale_execution_site() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(3);
let run_id = RunId::new_v4();
let (_stale_token, stale_rx, none) =
pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
assert!(none.is_none());
let (fresh_token, fresh_rx, superseded) =
pending.insert(workflow_id.clone(), &run_id, id.clone(), 2)?;
assert_eq!(
superseded,
Some(1),
"the takeover must name the attempt it displaced"
);
let stale_resolution = stale_rx
.recv_timeout(Duration::from_millis(50))
.map_err(|_| {
ServerError::worker_dispatch(
"test",
"takeover",
"the superseded waiter was never woken",
)
})?;
let reason = match stale_resolution {
Ok(delivered) => {
return Err(ServerError::worker_dispatch(
"test",
"takeover",
format!("the superseded waiter resolved Ok({delivered}) instead of an error"),
));
}
Err(reason) => reason,
};
assert!(
reason.starts_with("retryable:") && reason.contains("superseded by attempt 2"),
"{reason}"
);
pending.complete_activity(ActivityCompletion {
workflow_id: workflow_id.clone(),
activity_id: id.clone(),
run_id: None,
completion_token: fresh_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""second attempt wins""#.to_vec(),
)),
})?;
assert_eq!(
fresh_rx.recv_timeout(Duration::from_millis(50)),
Ok(Ok(r#""second attempt wins""#.to_owned()))
);
Ok(())
}
#[test]
fn a_lower_attempt_is_still_refused_by_name() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(4);
let run_id = RunId::new_v4();
let _held = pending.insert(workflow_id.clone(), &run_id, id.clone(), 3)?;
let error = pending
.insert(workflow_id, &run_id, id, 2)
.err()
.ok_or_else(|| {
ServerError::worker_dispatch(
"test",
"stale-redispatch",
"a lower attempt silently replaced a live responder",
)
})?;
let text = error.to_string();
assert!(
text.contains("attempt 2 arrived while attempt 3 still holds"),
"{text}"
);
Ok(())
}
#[test]
fn pending_complete_unknown_returns_false() {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
assert!(!pending.complete(
&WorkflowId::new_v4(),
&activity_id(99),
None,
Ok("orphan".to_owned())
));
}
#[derive(Default)]
struct RecordingOutboxCallback {
completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
live: bool,
}
impl OutboxDeliveryCallback for RecordingOutboxCallback {
fn deliver_completion(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
result: String,
) -> Result<bool, ServerError> {
let _ = run_id;
self.completions
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
.push((workflow_id.clone(), activity_id.clone(), result));
Ok(self.live)
}
fn deliver_failure(
&self,
workflow_id: &WorkflowId,
activity_id: &ActivityId,
run_id: Option<&RunId>,
reason: String,
) -> Result<bool, ServerError> {
let _ = run_id;
self.failures
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
.push((workflow_id.clone(), activity_id.clone(), reason));
Ok(self.live)
}
}
#[test]
fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let callback = Arc::new(RecordingOutboxCallback {
live: true,
..RecordingOutboxCallback::default()
});
pending.clone().set_outbox_delivery(callback.clone());
let workflow_id = WorkflowId::new_v4();
let id = activity_id(7);
assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
let completions = callback
.completions
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
assert_eq!(completions.len(), 1);
assert_eq!(completions[0].0, workflow_id);
assert_eq!(completions[0].1, id);
assert_eq!(completions[0].2, "done");
Ok(())
}
#[test]
fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
-> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let callback = Arc::new(RecordingOutboxCallback::default());
pending.set_outbox_delivery(callback.clone());
let workflow_id = WorkflowId::new_v4();
let id = activity_id(8);
assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
let failures = callback
.failures
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
assert_eq!(failures.len(), 1);
assert_eq!(failures[0].2, "retryable:boom");
Ok(())
}
#[test]
fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
assert!(!pending.complete(
&WorkflowId::new_v4(),
&activity_id(9),
None,
Ok("x".to_owned())
));
}
#[test]
fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let callback = Arc::new(RecordingOutboxCallback {
live: true,
..RecordingOutboxCallback::default()
});
pending.set_outbox_delivery(callback.clone());
let workflow_id = WorkflowId::new_v4();
let id = activity_id(10);
let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
assert_eq!(
rx.recv_timeout(Duration::from_millis(50)),
Ok(Ok("matched".to_owned()))
);
assert!(
callback
.completions
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
.is_empty(),
"a matched completion must deliver to its waiter, not the outbox callback"
);
Ok(())
}
#[test]
fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(11);
let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
pending.park_activity(&workflow_id, &id)?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
Ok(())
}
#[test]
fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let callback = Arc::new(RecordingOutboxCallback {
live: true,
..RecordingOutboxCallback::default()
});
pending.set_outbox_delivery(callback.clone());
pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
assert!(
callback
.failures
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
.is_empty(),
"a park must never be delivered as an outbox failure"
);
assert!(
callback
.completions
.lock()
.map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
.is_empty(),
"a park must never be delivered as an outbox completion"
);
Ok(())
}
#[test]
fn completion_sink_routes_success() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(2);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
pending.complete_activity(ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(payload),
})?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
Ok(())
}
#[test]
fn accepted_hook_runs_before_the_result_is_observable() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(13);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let accepted = AtomicBool::new(false);
pending.complete_activity_after_accept(
ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""ordered""#.to_vec(),
)),
},
|| {
accepted.store(true, Ordering::SeqCst);
Ok(())
},
)?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
assert_eq!(result, Ok(r#""ordered""#.to_owned()));
assert!(
accepted.load(Ordering::SeqCst),
"the accepted-path cleanup hook must finish before the waiter can observe the result"
);
Ok(())
}
#[test]
fn a_non_publishing_settlement_restores_the_generation_for_the_true_resolver()
-> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(16);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let settled = pending.complete_fenced_after_accept(
&workflow_id,
&id,
None,
&completion_token,
Err("retryable:worker lost".to_owned()),
|| Ok(false),
)?;
assert!(!settled, "a non-publishing settlement must report false");
assert!(
rx.try_recv().is_err(),
"a non-publishing settlement must not publish"
);
let published = pending.complete_fenced_after_accept(
&workflow_id,
&id,
None,
&completion_token,
Err("retryable:worker lost".to_owned()),
|| Ok(true),
)?;
assert!(
published,
"the true resolver must be accepted after a non-publishing settlement"
);
assert_eq!(
rx.recv_timeout(Duration::from_millis(50))
.map_err(|error| ServerError::worker_dispatch(
"",
"",
format!("channel: {error}")
))?,
Err("retryable:worker lost".to_owned())
);
Ok(())
}
#[test]
fn a_poisoned_tracker_fails_open_and_never_withholds_a_completion() {
let tracker = HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW);
tracker.poison_for_tests();
let workflow_id = WorkflowId::new_v4();
let id = activity_id(17);
let registry = ConnectedWorkerRegistry::default();
let was_tracked = clear_completed_task_tracking(
&tracker,
®istry,
WorkerId::from_value(1),
&workflow_id,
&id,
);
assert!(
was_tracked,
"a poisoned tracker must fail open (deliver), never withhold"
);
}
#[test]
fn accepted_hook_failure_blocks_publication_and_restores_generation() -> Result<(), ServerError>
{
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(15);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let rejected = pending.complete_activity_after_accept(
ActivityCompletion {
workflow_id: workflow_id.clone(),
activity_id: id.clone(),
run_id: None,
completion_token: completion_token.clone(),
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""ordered""#.to_vec(),
)),
},
|| Err(ServerError::lock_poisoned("accepted-path settlement")),
);
assert!(matches!(rejected, Err(ServerError::LockPoisoned { .. })));
assert!(
rx.try_recv().is_err(),
"failed accepted-path settlement must not publish the completion"
);
pending.complete_activity_after_accept(
ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""ordered""#.to_vec(),
)),
},
|| Ok(()),
)?;
assert_eq!(
rx.recv_timeout(Duration::from_millis(50))
.map_err(|error| ServerError::worker_dispatch(
"",
"",
format!("channel: {error}")
))?,
Ok(r#""ordered""#.to_owned())
);
Ok(())
}
#[test]
fn refused_generation_never_runs_the_accepted_hook() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(14);
let (_completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let accepted = AtomicBool::new(false);
let rejected = pending.complete_activity_after_accept(
ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token: CompletionToken::for_test(),
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""stale""#.to_vec(),
)),
},
|| {
accepted.store(true, Ordering::SeqCst);
Ok(())
},
);
assert!(matches!(
rejected,
Err(ServerError::ActivityCompletionRejected { .. })
));
assert!(
!accepted.load(Ordering::SeqCst),
"a refused generation must preserve liveness by skipping accepted-path cleanup"
);
assert!(
rx.try_recv().is_err(),
"a refused generation must leave the current waiter unresolved"
);
Ok(())
}
#[test]
fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(12);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
let malformed = pending.complete_activity(ActivityCompletion {
workflow_id: workflow_id.clone(),
activity_id: id.clone(),
run_id: None,
completion_token: completion_token.clone(),
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
vec![0xff],
)),
});
assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
assert!(
rx.try_recv().is_err(),
"an invalid result must leave the waiter unresolved"
);
pending.complete_activity(ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""valid""#.to_vec(),
)),
})?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
assert_eq!(result, Ok(r#""valid""#.to_owned()));
Ok(())
}
#[test]
fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(3);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
pending.complete_activity(ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Failed(ActivityError {
kind: ActivityErrorKind::Retryable,
message: "temporary".to_owned(),
details: None,
}),
})?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
assert_eq!(result, Err("retryable:temporary".to_owned()));
Ok(())
}
#[test]
fn completion_sink_routes_policy_refusal_with_its_own_prefix() -> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let workflow_id = WorkflowId::new_v4();
let id = activity_id(4);
let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
pending.complete_activity(ActivityCompletion {
workflow_id,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Failed(ActivityError {
kind: ActivityErrorKind::PolicyRefused,
message: "provider safety policy".to_owned(),
details: None,
}),
})?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
assert_eq!(
result,
Err("policy_refused:provider safety policy".to_owned())
);
Ok(())
}
#[test]
fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
-> Result<(), ServerError> {
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let post_restart_workflow = WorkflowId::new_v4();
let pre_restart_workflow = WorkflowId::new_v4();
let id = activity_id(1);
let (completion_token, rx) = insert_attempt_one(&pending, &post_restart_workflow, &id)?;
let rejected = pending.complete_activity(ActivityCompletion {
workflow_id: pre_restart_workflow,
activity_id: id.clone(),
run_id: None,
completion_token: CompletionToken::for_test(),
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""stale""#.to_vec(),
)),
});
assert!(matches!(
rejected,
Err(ServerError::ActivityCompletionRejected { .. })
));
assert!(
rx.try_recv().is_err(),
"stale result for a different workflow must not complete this dispatch"
);
pending.complete_activity(ActivityCompletion {
workflow_id: post_restart_workflow,
activity_id: id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""fresh""#.to_vec(),
)),
})?;
let result = rx
.recv_timeout(Duration::from_millis(50))
.map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
assert_eq!(result, Ok(r#""fresh""#.to_owned()));
Ok(())
}
const TEST_HEARTBEAT_WINDOW: Duration = Duration::from_secs(5);
#[test]
fn a_sink_derives_its_transport_budget_from_the_operator_window() {
let window = Duration::from_secs(7);
let sink = PendingActivities::new(window);
assert_eq!(
sink.transport_losses().budget(),
window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
"the transport-loss budget must be the operator's heartbeat window times \
TRANSPORT_LOSS_BUDGET_WINDOWS; a sink carrying any other budget is one no \
operator declared"
);
}
#[test]
fn a_dispatcher_carries_the_transport_budget_of_the_tracker_it_was_given() {
let window = Duration::from_secs(11);
let registry = ConnectedWorkerRegistry::default();
let dispatcher =
WorkerActivityDispatcher::new(registry, "default", HeartbeatTracker::new(window));
assert_eq!(
dispatcher.pending().transport_losses().budget(),
window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
"the dispatcher's transport budget must be derived from the heartbeat window it \
was constructed with; a zero or mismatched budget here means a worker loss is \
declared transport-exhausted after ONE re-dispatch instead of after the \
operator's window"
);
assert_ne!(
dispatcher.pending().transport_losses().budget(),
Duration::ZERO,
"a ZERO budget is the specific regression this pin exists for: it grants exactly \
one re-dispatchable loss and then declares the infrastructure flapping"
);
}
fn test_tracker() -> HeartbeatTracker {
HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW)
}
fn greet_request() -> ActivityDispatch {
ActivityDispatch {
namespace: "default".to_owned(),
task_queue: "default".to_owned(),
node: None,
workflow_id: WorkflowId::new_v4(),
run_id: RunId::new_v4(),
activity_id: ActivityId::from_sequence_position(0),
name: "greet".to_owned(),
input: "{}".to_owned(),
config: "{}".to_owned(),
attempt: 1,
labels: std::collections::BTreeMap::new(),
advisory: false,
}
}
#[test]
fn dispatcher_fails_immediately_when_draining_without_workers() {
let registry = ConnectedWorkerRegistry::default();
let drain = DrainState::default();
let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
.with_drain_state(drain.clone());
let _ = drain.begin();
let result = dispatcher.dispatch(greet_request());
assert!(result.is_err());
let err = result.err().unwrap_or_default();
assert!(
err.contains("drain"),
"expected drain rejection, got: {err}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_retry_takes_over_a_timed_out_attempt_and_asks_its_worker_to_stop()
-> Result<(), Box<dyn std::error::Error>> {
let registry = ConnectedWorkerRegistry::default();
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
let activity_types = [String::from("greet")];
let registration = registry.register(
"default",
activity_types.iter(),
worker_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let worker_id = registration
.worker_id()
.ok_or("test worker registration missing id")?;
let tracker = test_tracker();
let dispatcher = Arc::new(
WorkerActivityDispatcher::new(registry, "default", tracker.clone())
.with_pending(pending.clone()),
);
let request = greet_request();
let workflow_id = request.workflow_id.clone();
let activity_id = request.activity_id.clone();
let first_dispatcher = Arc::clone(&dispatcher);
let first_request = request.clone();
let first_attempt = tokio::spawn(futures::future::lazy(move |_| {
first_dispatcher.dispatch(first_request)
}));
let Some(WorkerMessage::ActivityTask(first_task)) = worker_rx.recv().await else {
return Err("expected attempt 1's task on the worker channel".into());
};
assert_eq!(first_task.attempt, 1);
let mut retry_request = request.clone();
retry_request.attempt = 2;
let retry_dispatcher = Arc::clone(&dispatcher);
let retry_attempt = tokio::spawn(futures::future::lazy(move |_| {
retry_dispatcher.dispatch(retry_request)
}));
let Some(WorkerMessage::CancelActivity(cancel)) = worker_rx.recv().await else {
return Err(
"expected the superseded attempt's cancel frame before attempt 2's task".into(),
);
};
assert_eq!(
cancel
.activity_id
.map(|id| id.sequence_position)
.unwrap_or_default(),
activity_id.sequence_position()
);
let Some(WorkerMessage::ActivityTask(retry_task)) = worker_rx.recv().await else {
return Err("expected attempt 2's task on the worker channel".into());
};
assert_eq!(retry_task.attempt, 2);
let in_flight = tracker.in_flight_for_workflow(&workflow_id)?;
assert_eq!(
in_flight
.iter()
.map(|task| task.attempt)
.collect::<Vec<_>>(),
vec![2],
"exactly attempt 2 must remain tracked for worker {worker_id:?}"
);
let retry_token = CompletionToken::from_wire(
&workflow_id,
&activity_id,
retry_task.completion_token.clone(),
)
.map_err(|error| error.to_string())?;
pending.complete_activity(ActivityCompletion {
workflow_id: workflow_id.clone(),
activity_id: activity_id.clone(),
run_id: None,
completion_token: retry_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#""attempt 2 delivered""#.to_vec(),
)),
})?;
let retry_result = retry_attempt.await.map_err(|error| error.to_string())?;
assert_eq!(retry_result, Ok(r#""attempt 2 delivered""#.to_owned()));
let first_result = first_attempt.await.map_err(|error| error.to_string())?;
let first_reason = match first_result {
Ok(delivered) => {
return Err(
format!("attempt 1 resolved Ok({delivered}) instead of superseded").into(),
);
}
Err(reason) => reason,
};
assert!(
first_reason.starts_with("retryable:")
&& first_reason.contains("superseded by attempt 2"),
"{first_reason}"
);
registration.deregister()?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
-> Result<(), Box<dyn std::error::Error>> {
let registry = ConnectedWorkerRegistry::default();
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
let activity_types = [String::from("greet")];
let registration = registry.register(
"default",
activity_types.iter(),
worker_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let sink = pending.clone();
let echo_worker = tokio::spawn(async move {
let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
return Err("expected an activity task on the worker channel".to_owned());
};
let workflow_id = task
.workflow_id
.ok_or("task missing workflow id")
.and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
let activity_id = task
.activity_id
.map(ActivityId::from)
.ok_or("task missing activity id")?;
let completion_token =
CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
.map_err(|error| error.to_string())?;
sink.complete_activity(ActivityCompletion {
workflow_id,
activity_id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#"{"greeting":"hello"}"#.to_vec(),
)),
})
.map_err(|error| error.to_string())
});
let dispatcher = Arc::new(
WorkerActivityDispatcher::new(registry, "default", test_tracker())
.with_pending(pending),
);
let started = Instant::now();
let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
dispatcher.dispatch(greet_request())
}));
let result = dispatch_task.await.map_err(|error| error.to_string())?;
let elapsed = started.elapsed();
assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
assert!(
elapsed < Duration::from_secs(5),
"dispatch round trip took {elapsed:?}; task delivery must not \
depend on the blocked dispatch thread"
);
echo_worker.await.map_err(|error| error.to_string())??;
registration.deregister()?;
Ok(())
}
#[tokio::test]
async fn dispatch_on_current_thread_runtime_fails_fast()
-> Result<(), Box<dyn std::error::Error>> {
let registry = ConnectedWorkerRegistry::default();
let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
let activity_types = [String::from("greet")];
let registration = registry.register(
"default",
activity_types.iter(),
worker_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
let started = Instant::now();
let result = dispatcher.dispatch(greet_request());
let elapsed = started.elapsed();
let err = result.err().ok_or("expected dispatch to fail")?;
assert!(
err.contains("multi-thread tokio runtime"),
"unexpected error: {err}"
);
assert!(
elapsed < Duration::from_secs(5),
"fail-fast path took {elapsed:?}"
);
registration.deregister()?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn dispatch_pinned_to_node_reaches_only_that_node()
-> Result<(), Box<dyn std::error::Error>> {
let registry = ConnectedWorkerRegistry::default();
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let activity_types = [String::from("greet")];
let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
let on_n2 = registry.register_namespaces(
[String::from("default")],
"default",
Some(String::from("n2")),
activity_types.iter(),
n2_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let on_n1 = registry.register_namespaces(
[String::from("default")],
"default",
Some(String::from("n1")),
activity_types.iter(),
n1_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let sink = pending.clone();
let echo_n1 = tokio::spawn(async move {
let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
return Err("expected an activity task on the n1 worker channel".to_owned());
};
let workflow_id = task
.workflow_id
.ok_or("task missing workflow id")
.and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
let activity_id = task
.activity_id
.map(ActivityId::from)
.ok_or("task missing activity id")?;
let completion_token =
CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
.map_err(|error| error.to_string())?;
sink.complete_activity(ActivityCompletion {
workflow_id,
activity_id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#"{"greeting":"hello"}"#.to_vec(),
)),
})
.map_err(|error| error.to_string())
});
let dispatcher = Arc::new(
WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
.with_pending(pending),
);
let pinned = ActivityDispatch {
node: Some(String::from("n1")),
..greet_request()
};
let started = Instant::now();
let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
.await
.map_err(|error| error.to_string())?;
let elapsed = started.elapsed();
assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
assert!(
elapsed < Duration::from_secs(5),
"pinned dispatch round trip took {elapsed:?}; the task must route to n1"
);
echo_n1.await.map_err(|error| error.to_string())??;
assert!(
n2_rx.try_recv().is_err(),
"node=Some(\"n1\") dispatch must not reach the n2 worker"
);
on_n1.deregister()?;
on_n2.deregister()?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
-> Result<(), Box<dyn std::error::Error>> {
let registry = ConnectedWorkerRegistry::default();
let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
let activity_types = [String::from("greet")];
let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
let on_n1 = registry.register_namespaces(
[String::from("default")],
"default",
Some(String::from("n1")),
activity_types.iter(),
n1_tx,
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)?;
let sink = pending.clone();
let echo = tokio::spawn(async move {
let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
return Err("expected an activity task on the worker channel".to_owned());
};
let workflow_id = task
.workflow_id
.ok_or("task missing workflow id")
.and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
let activity_id = task
.activity_id
.map(ActivityId::from)
.ok_or("task missing activity id")?;
let completion_token =
CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
.map_err(|error| error.to_string())?;
sink.complete_activity(ActivityCompletion {
workflow_id,
activity_id,
run_id: None,
completion_token,
outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
ContentType::Json,
br#"{"greeting":"hello"}"#.to_vec(),
)),
})
.map_err(|error| error.to_string())
});
let dispatcher = Arc::new(
WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
.with_pending(pending),
);
let result = tokio::spawn(futures::future::lazy(move |_| {
dispatcher.dispatch(greet_request())
}))
.await
.map_err(|error| error.to_string())?;
assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
echo.await.map_err(|error| error.to_string())??;
on_n1.deregister()?;
Ok(())
}
}