use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use aion_core::{ActivityId, ContentType, Payload, RunId, WorkflowId};
use aion_store::OutboxRow;
use async_trait::async_trait;
use liminal::protocol::WorkerRegistration as WireWorkerRegistration;
use liminal_sdk::{SchemaMetadata, SchemaValidate};
use liminal_server::ServerError as LiminalServerError;
use liminal_server::server::connection::{
ConnectionNotifier, ConnectionSupervisor, PushReplyAwaiter,
};
use serde::{Deserialize, Serialize};
use super::bridge::OutboxDeliveryCallback;
use super::envelope::{CompletionFences, CompletionToken};
use super::registry::{ConnectedWorkerRegistry, WorkerDelivery, WorkerRegistration};
use crate::error::ServerError;
const PUSH_REPLY_TIMEOUT: Duration = Duration::from_secs(30);
const BRIDGE_REPLY_POLL: Duration = Duration::from_secs(1);
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct DispatchRequest {
pub activity_type: String,
pub workflow_id: WorkflowId,
pub ordinal: u64,
pub run_id: Option<RunId>,
pub completion_token: String,
pub idempotency_key: String,
pub input: Vec<u8>,
#[serde(default = "first_attempt")]
pub attempt: u32,
#[serde(default)]
pub labels: std::collections::BTreeMap<String, String>,
#[serde(default)]
pub heartbeat_window_ms: u64,
}
const fn first_attempt() -> u32 {
1
}
impl SchemaValidate for DispatchRequest {
fn schema_metadata() -> SchemaMetadata {
SchemaMetadata::new(
"aion.outbox.dispatch.request",
"1",
br#"{"type":"object"}"#.as_slice(),
)
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct DispatchResponse {
pub workflow_id: WorkflowId,
pub ordinal: u64,
pub run_id: Option<RunId>,
pub completion_token: String,
pub outcome: Result<String, String>,
}
impl SchemaValidate for DispatchResponse {
fn schema_metadata() -> SchemaMetadata {
SchemaMetadata::new(
"aion.outbox.dispatch.response",
"1",
br#"{"type":"object"}"#.as_slice(),
)
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct InterventionRequest {
pub intervention: aion_core::InterventionCommand,
}
impl SchemaValidate for InterventionRequest {
fn schema_metadata() -> SchemaMetadata {
SchemaMetadata::new(
"aion.intervention.request",
"1",
br#"{"type":"object"}"#.as_slice(),
)
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct InterventionReply {
pub outcome: aion_core::InterventionOutcome,
}
impl SchemaValidate for InterventionReply {
fn schema_metadata() -> SchemaMetadata {
SchemaMetadata::new(
"aion.intervention.reply",
"1",
br#"{"type":"object"}"#.as_slice(),
)
}
}
pub const WORKER_LIVENESS_CHANNEL: &str = "aion.worker.liveness";
pub const WORKER_CAPABILITIES_CHANNEL: &str = "aion.worker.capabilities";
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerCapabilitiesAnnouncement {
pub capabilities: aion_core::InterventionCapabilities,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct WorkerLivenessBeat {
pub workflow_id: WorkflowId,
pub ordinal: u64,
}
const SEGMENT_SEPARATOR: char = '.';
const SEGMENT_ESCAPE: char = '%';
fn encode_segment(segment: &str) -> String {
if !segment.contains([SEGMENT_SEPARATOR, SEGMENT_ESCAPE]) {
return segment.to_owned();
}
let mut encoded = String::with_capacity(segment.len());
for ch in segment.chars() {
match ch {
SEGMENT_ESCAPE => encoded.push_str("%25"),
SEGMENT_SEPARATOR => encoded.push_str("%2E"),
other => encoded.push(other),
}
}
encoded
}
#[must_use]
pub fn dispatch_channel_name(namespace: &str, task_queue: &str, node: Option<&str>) -> String {
let namespace = encode_segment(namespace);
let task_queue = encode_segment(task_queue);
match node {
Some(node) => {
let node = encode_segment(node);
format!("aion.dispatch.{namespace}.{task_queue}.{node}")
}
None => format!("aion.dispatch.{namespace}.{task_queue}"),
}
}
#[must_use]
pub fn channel_for_row(row: &OutboxRow) -> String {
dispatch_channel_name(&row.namespace, &row.task_queue, row.node.as_deref())
}
fn dispatch_error(channel: &str, reason: String) -> ServerError {
ServerError::WorkerDispatch {
namespace: "liminal".to_owned(),
activity_type: channel.to_owned(),
reason,
}
}
pub struct LiminalCompletionSource {
callback: Arc<dyn OutboxDeliveryCallback>,
completion_fences: CompletionFences,
}
impl std::fmt::Debug for LiminalCompletionSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LiminalCompletionSource")
.finish_non_exhaustive()
}
}
impl LiminalCompletionSource {
#[must_use]
pub fn new(callback: Arc<dyn OutboxDeliveryCallback>) -> Self {
Self {
callback,
completion_fences: CompletionFences::default(),
}
}
#[must_use]
pub fn with_completion_fences(mut self, completion_fences: CompletionFences) -> Self {
self.completion_fences = completion_fences;
self
}
pub fn deliver(&self, response: &DispatchResponse) -> Result<bool, ServerError> {
let run_id = response.run_id.as_ref().ok_or_else(|| {
dispatch_error(
"liminal completion",
"activity result run id is missing; refusing run-gate bypass".to_owned(),
)
})?;
let activity_id = ActivityId::from_sequence_position(response.ordinal);
let completion_token = CompletionToken::from_wire(
&response.workflow_id,
&activity_id,
response.completion_token.clone(),
)?;
match self
.completion_fences
.accept(&response.workflow_id, &activity_id, &completion_token)
{
Ok(accepted) => {
tracing::debug!(
workflow_id = %response.workflow_id,
%activity_id,
attempt = accepted.attempt(),
"liminal completion consumed the execution generation"
);
}
Err(error) => {
tracing::warn!(
workflow_id = %response.workflow_id,
%activity_id,
%error,
"liminal completion fence rejected a late or stale reply"
);
return Err(error);
}
}
match &response.outcome {
Ok(result) => self.callback.deliver_completion(
&response.workflow_id,
&activity_id,
Some(run_id),
result.clone(),
),
Err(reason) => self.callback.deliver_failure(
&response.workflow_id,
&activity_id,
Some(run_id),
reason.clone(),
),
}
}
}
#[must_use]
pub fn payload_from_request(request: &DispatchRequest) -> Payload {
Payload::new(ContentType::Json, request.input.clone())
}
#[derive(Clone)]
pub struct LiminalWorkerDelivery {
supervisor: ConnectionSupervisor,
pid: u64,
}
impl std::fmt::Debug for LiminalWorkerDelivery {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LiminalWorkerDelivery")
.field("pid", &self.pid)
.finish_non_exhaustive()
}
}
impl LiminalWorkerDelivery {
#[must_use]
pub const fn new(supervisor: ConnectionSupervisor, pid: u64) -> Self {
Self { supervisor, pid }
}
#[must_use]
pub const fn pid(&self) -> u64 {
self.pid
}
pub fn dispatch_held(
&self,
request: &DispatchRequest,
keep_waiting: impl Fn() -> bool,
) -> Result<Option<DispatchResponse>, ServerError> {
let awaiter = self.push_dispatch(request)?;
receive_bridge_reply(&awaiter, keep_waiting)
}
pub(crate) fn push_dispatch(
&self,
request: &DispatchRequest,
) -> Result<PushReplyAwaiter, ServerError> {
let payload = serde_json::to_vec(request).map_err(|error| {
dispatch_error("liminal-push", format!("request serialize failed: {error}"))
})?;
self.supervisor
.push_to_connection(self.pid, payload)
.map_err(|error| classify_push_error(&error))
}
pub fn push_payload_with_deadline(
&self,
payload: Vec<u8>,
deadline: Duration,
) -> Result<PushReplyAwaiter, ServerError> {
self.supervisor
.push_to_connection_with_deadline(self.pid, payload, deadline)
.map_err(|error| classify_push_error(&error))
}
pub fn push_intervention(
&self,
request: &InterventionRequest,
) -> Result<InterventionReply, ServerError> {
let payload = serde_json::to_vec(request).map_err(|error| {
dispatch_error(
"liminal-push",
format!("intervention serialize failed: {error}"),
)
})?;
let awaiter = self
.supervisor
.push_to_connection_with_deadline(self.pid, payload, PUSH_REPLY_TIMEOUT)
.map_err(|error| {
ServerError::worker_connection_lost(
"liminal-push",
format!("push intervention to worker failed: {error}"),
)
})?;
let reply = awaiter.receive(PUSH_REPLY_TIMEOUT).map_err(|error| {
if is_connection_closed_reply_error(&error) {
ServerError::worker_connection_lost(
"liminal-push",
format!("worker connection closed before intervention ack: {error}"),
)
} else {
dispatch_error("liminal-push", format!("intervention ack failed: {error}"))
}
})?;
serde_json::from_slice(&reply).map_err(|error| {
dispatch_error(
"liminal-push",
format!("intervention ack decode failed: {error}"),
)
})
}
}
fn decode_dispatch_response(reply: &[u8]) -> Result<DispatchResponse, ServerError> {
serde_json::from_slice(reply).map_err(|error| {
dispatch_error(
"liminal-push",
format!("worker reply decode failed: {error}"),
)
})
}
pub(crate) fn receive_bridge_reply(
awaiter: &PushReplyAwaiter,
keep_waiting: impl Fn() -> bool,
) -> Result<Option<DispatchResponse>, ServerError> {
loop {
match awaiter.receive(BRIDGE_REPLY_POLL) {
Ok(reply) => return decode_dispatch_response(&reply).map(Some),
Err(LiminalServerError::PushReplyTimeout { .. }) => {
if !keep_waiting() {
return Ok(None);
}
}
Err(error) if is_connection_closed_reply_error(&error) => {
return Err(ServerError::worker_connection_lost(
"liminal-push",
format!("worker connection closed before reply: {error}"),
));
}
Err(error) => {
return Err(dispatch_error(
"liminal-push",
format!("worker reply failed: {error}"),
));
}
}
}
}
fn is_connection_closed_reply_error(error: &LiminalServerError) -> bool {
matches!(error, LiminalServerError::PushReplyDisconnected { .. })
}
fn classify_push_error(error: &LiminalServerError) -> ServerError {
if is_connection_cap_error(error) {
ServerError::worker_busy(
"liminal-push",
format!("worker connection push cap reached: {error}"),
)
} else {
ServerError::worker_connection_lost(
"liminal-push",
format!("push to worker failed: {error}"),
)
}
}
fn is_connection_cap_error(error: &LiminalServerError) -> bool {
matches!(error, LiminalServerError::ConnectionCapReached { .. })
}
pub(crate) struct AttemptOwnerGuard {
owners: super::intervention::AttemptOwnerIndex,
key: super::intervention::AttemptKey,
}
impl AttemptOwnerGuard {
pub(crate) fn bind(
owners: super::intervention::AttemptOwnerIndex,
key: super::intervention::AttemptKey,
worker: super::registry::WorkerId,
) -> Self {
owners.bind(key.clone(), worker);
Self { owners, key }
}
}
impl Drop for AttemptOwnerGuard {
fn drop(&mut self) {
self.owners.release(&self.key);
}
}
fn normalize_wire_node(node: Option<&str>) -> Option<String> {
node.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub struct LiminalConnectionNotifier {
registry: ConnectedWorkerRegistry,
contract_catalog: Option<Arc<aion::Engine>>,
supervisor: OnceLock<ConnectionSupervisor>,
guards: Mutex<HashMap<u64, WorkerRegistration>>,
intervention_capabilities: aion_core::InterventionCapabilities,
transcript: Option<TranscriptTap>,
heartbeat_tracker: Option<super::heartbeat::HeartbeatTracker>,
}
#[derive(Clone)]
struct TranscriptTap {
queue: tokio::sync::mpsc::Sender<aion_core::ActivityEvent>,
}
const TRANSCRIPT_QUEUE_CAPACITY: usize = 4096;
impl std::fmt::Debug for LiminalConnectionNotifier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LiminalConnectionNotifier")
.field("supervisor_bound", &self.supervisor.get().is_some())
.finish_non_exhaustive()
}
}
impl LiminalConnectionNotifier {
#[must_use]
pub fn new(registry: ConnectedWorkerRegistry) -> Self {
Self {
registry,
contract_catalog: None,
supervisor: OnceLock::new(),
guards: Mutex::new(HashMap::new()),
intervention_capabilities: aion_core::InterventionCapabilities::none(),
transcript: None,
heartbeat_tracker: None,
}
}
#[must_use]
pub fn with_contract_catalog(mut self, engine: Arc<aion::Engine>) -> Self {
self.contract_catalog = Some(engine);
self
}
#[must_use]
pub fn with_heartbeat_tracker(mut self, tracker: super::heartbeat::HeartbeatTracker) -> Self {
self.heartbeat_tracker = Some(tracker);
self
}
#[must_use]
pub fn with_transcript_publisher(
mut self,
publisher: crate::activity_publisher::ActivityEventPublisher,
) -> Self {
let (queue, mut events) =
tokio::sync::mpsc::channel::<aion_core::ActivityEvent>(TRANSCRIPT_QUEUE_CAPACITY);
tokio::runtime::Handle::current().spawn(async move {
let dropped = publisher.drain(&mut events, "observability_tap").await;
if dropped > 0 {
tracing::warn!(
dropped,
operation = "observability_tap",
"observability tap: transcript events were not retained"
);
}
});
self.transcript = Some(TranscriptTap { queue });
self
}
#[must_use]
pub fn with_intervention_capabilities(
mut self,
capabilities: aion_core::InterventionCapabilities,
) -> Self {
self.intervention_capabilities = capabilities;
self
}
pub fn bind_supervisor(&self, supervisor: ConnectionSupervisor) -> bool {
self.supervisor.set(supervisor).is_ok()
}
#[must_use]
pub fn liveness_targets(&self) -> Vec<super::liminal_liveness::LivenessTarget> {
let Some(supervisor) = self.supervisor.get() else {
return Vec::new();
};
let guards = match self.guards.lock() {
Ok(guards) => guards,
Err(poisoned) => poisoned.into_inner(),
};
guards
.iter()
.filter_map(|(pid, guard)| {
guard
.worker_id()
.map(|worker_id| super::liminal_liveness::LivenessTarget {
pid: *pid,
worker_id,
delivery: LiminalWorkerDelivery::new(supervisor.clone(), *pid),
})
})
.collect()
}
fn record_liveness_beat(&self, pid: u64, payload: &[u8]) {
let Some(tracker) = &self.heartbeat_tracker else {
return;
};
let beat: WorkerLivenessBeat = match serde_json::from_slice(payload) {
Ok(beat) => beat,
Err(error) => {
tracing::warn!(%error, "liveness tap: malformed WorkerLivenessBeat payload");
return;
}
};
let worker_id = match self.guards.lock() {
Ok(guards) => guards.get(&pid).and_then(WorkerRegistration::worker_id),
Err(poisoned) => poisoned
.into_inner()
.get(&pid)
.and_then(WorkerRegistration::worker_id),
};
let Some(worker_id) = worker_id else {
tracing::warn!(
connection_pid = pid,
"liveness tap: beat from a connection with no registered worker"
);
return;
};
let activity_id = ActivityId::from_sequence_position(beat.ordinal);
if let Err(error) = tracker.record_liveness(
worker_id,
&beat.workflow_id,
&activity_id,
std::time::Instant::now(),
) {
tracing::error!(
%error,
connection_pid = pid,
"liveness tap: heartbeat tracker refresh failed"
);
}
}
fn record_capabilities_announcement(&self, pid: u64, payload: &[u8]) {
let announcement: WorkerCapabilitiesAnnouncement = match serde_json::from_slice(payload) {
Ok(announcement) => announcement,
Err(error) => {
tracing::warn!(
%error,
"capabilities tap: malformed WorkerCapabilitiesAnnouncement payload"
);
return;
}
};
let worker_id = match self.guards.lock() {
Ok(guards) => guards.get(&pid).and_then(WorkerRegistration::worker_id),
Err(poisoned) => poisoned
.into_inner()
.get(&pid)
.and_then(WorkerRegistration::worker_id),
};
let Some(worker_id) = worker_id else {
tracing::warn!(
connection_pid = pid,
"capabilities tap: announcement from a connection with no registered worker"
);
return;
};
match self
.registry
.set_intervention_capabilities(worker_id, &announcement.capabilities)
{
Ok(true) => {}
Ok(false) => tracing::warn!(
connection_pid = pid,
worker_id = ?worker_id,
"capabilities tap: announcement raced the worker's deregistration"
),
Err(error) => tracing::error!(
%error,
connection_pid = pid,
"capabilities tap: registry capability update failed"
),
}
}
fn validate_registration_contract(
&self,
registration: &WireWorkerRegistration,
) -> Result<(), LiminalServerError> {
let Some(engine) = &self.contract_catalog else {
return Ok(());
};
let advertised =
registration
.activities
.iter()
.map(|activity| {
let input_schema =
serde_json::from_str(&activity.input_schema_json).map_err(|error| {
LiminalServerError::ListenerAccept {
message: format!(
"liminal worker activity `{}` input schema is invalid: {error}",
activity.name
),
}
})?;
let output_schema = serde_json::from_str(&activity.output_schema_json)
.map_err(|error| LiminalServerError::ListenerAccept {
message: format!(
"liminal worker activity `{}` output schema is invalid: {error}",
activity.name
),
})?;
Ok(aion_package::ActivityDescriptor {
name: activity.name.clone(),
input_schema,
output_schema,
})
})
.collect::<Result<Vec<_>, LiminalServerError>>()?;
let activity_types = registration
.activity_types
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
super::contracts::validate_worker_contracts(
engine,
self.registry.admission_audit(),
®istration.task_queue,
registration.node.as_deref(),
®istration.identity,
super::contracts::WorkerAdvertisement {
activity_types: &activity_types,
contracts: &advertised,
},
)
.map_err(|error| LiminalServerError::ListenerAccept {
message: error.to_string(),
})
}
fn admit_registration(
&self,
pid: u64,
registration: &WireWorkerRegistration,
) -> Result<(), LiminalServerError> {
let supervisor =
self.supervisor
.get()
.ok_or_else(|| LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} rejected: \
notifier supervisor handle not yet bound"
),
})?;
let delivery = WorkerDelivery::Liminal(LiminalWorkerDelivery::new(supervisor.clone(), pid));
let node = normalize_wire_node(registration.node.as_deref());
let guard = self
.registry
.register_delivery_with_capabilities(
registration.namespaces.iter().cloned(),
registration.task_queue.clone(),
node,
registration.activity_types.iter(),
delivery,
self.intervention_capabilities.clone(),
)
.map_err(|error| LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} rejected: {error}"
),
})?;
let worker_id = guard
.worker_id()
.ok_or_else(|| LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} has no worker id"
),
})?;
let mut guards = self.guards.lock().map_err(|_| {
LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} rejected: \
notifier guard map poisoned"
),
}
})?;
guards.insert(pid, guard);
drop(guards);
if let Some(tracker) = &self.heartbeat_tracker
&& let Err(error) = tracker.register_connection(worker_id, std::time::Instant::now())
{
let removed = match self.guards.lock() {
Ok(mut guards) => guards.remove(&pid),
Err(poisoned) => poisoned.into_inner().remove(&pid),
};
drop(removed);
return Err(LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} rejected: {error}"
),
});
}
tracing::info!(
connection_pid = pid,
identity = %registration.identity,
task_queue = %registration.task_queue,
"registered liminal worker in-band"
);
Ok(())
}
}
impl ConnectionNotifier for LiminalConnectionNotifier {
fn on_worker_registered(
&self,
pid: u64,
registration: &WireWorkerRegistration,
) -> Result<(), LiminalServerError> {
self.validate_registration_contract(registration)?;
self.admit_registration(pid, registration)
.inspect_err(|error| {
let advertised_contracts = registration
.activities
.iter()
.map(|activity| activity.name.as_str())
.collect::<Vec<_>>();
tracing::warn!(
connection_pid = pid,
identity = %registration.identity,
task_queue = %registration.task_queue,
namespaces = ?registration.namespaces,
advertised_activity_types = ?registration.activity_types,
advertised_contracts = ?advertised_contracts,
reason = %error,
"REFUSED liminal worker registration"
);
})
}
fn on_worker_unregistered(&self, pid: u64) {
let removed = match self.guards.lock() {
Ok(mut guards) => guards.remove(&pid),
Err(poisoned) => poisoned.into_inner().remove(&pid),
};
let worker_id = removed.as_ref().and_then(WorkerRegistration::worker_id);
if let (Some(tracker), Some(worker_id)) = (&self.heartbeat_tracker, worker_id)
&& let Err(error) = tracker.unregister_connection(worker_id)
{
tracing::error!(
%error,
connection_pid = pid,
worker_id = worker_id.value(),
"failed to clear liminal worker connection lease"
);
}
if removed.is_some() {
tracing::info!(
connection_pid = pid,
"deregistered liminal worker on disconnect"
);
}
}
fn on_channel_publish(&self, pid: u64, channel: &str, payload: &[u8]) -> bool {
if let Some(tracker) = &self.heartbeat_tracker {
let worker_id = match self.guards.lock() {
Ok(guards) => guards.get(&pid).and_then(WorkerRegistration::worker_id),
Err(poisoned) => poisoned
.into_inner()
.get(&pid)
.and_then(WorkerRegistration::worker_id),
};
if let Some(worker_id) = worker_id
&& let Err(error) =
tracker.record_connection_activity(worker_id, std::time::Instant::now())
{
tracing::error!(
%error,
connection_pid = pid,
worker_id = worker_id.value(),
"failed to advance liminal worker connection lease"
);
}
}
if channel == WORKER_LIVENESS_CHANNEL {
self.record_liveness_beat(pid, payload);
return true;
}
if channel == WORKER_CAPABILITIES_CHANNEL {
self.record_capabilities_announcement(pid, payload);
return true;
}
if channel != liminal_sdk::OBSERVABILITY_CHANNEL {
return false;
}
let Some(tap) = &self.transcript else {
return true;
};
let event: aion_core::ActivityEvent = match serde_json::from_slice(payload) {
Ok(event) => event,
Err(error) => {
tracing::warn!(%error, "observability tap: malformed ActivityEvent payload");
return true;
}
};
if let Err(error) = tap.queue.try_send(event) {
tracing::warn!(%error, "observability tap: transcript queue rejected event");
}
true
}
}
#[derive(Clone, Debug, Default)]
pub struct LiminalInterventionTransport;
#[async_trait]
impl super::intervention::InterventionTransport for LiminalInterventionTransport {
async fn push(
&self,
worker: &super::registry::WorkerHandle,
command: aion_core::InterventionCommand,
) -> Result<aion_core::InterventionOutcome, ServerError> {
let delivery = match worker.delivery() {
WorkerDelivery::Liminal(delivery) => delivery.clone(),
WorkerDelivery::Grpc(_) => {
return Err(ServerError::worker_connection_lost(
"liminal-push",
"owning worker is not delivered over liminal".to_owned(),
));
}
};
let request = InterventionRequest {
intervention: command,
};
let reply = tokio::task::spawn_blocking(move || delivery.push_intervention(&request))
.await
.map_err(|error| {
dispatch_error(
"liminal-push",
format!("intervention task join failed: {error}"),
)
})??;
Ok(reply.outcome)
}
}
#[cfg(test)]
#[path = "liminal_contract_tests.rs"]
mod contract_tests;
#[cfg(test)]
mod tests {
use super::{channel_for_row, dispatch_channel_name, normalize_wire_node};
use aion_core::{ActivityId, ContentType, Payload, WorkflowId};
use aion_store::{OutboxRow, OutboxStatus};
use chrono::Utc;
use uuid::Uuid;
#[tokio::test]
async fn attempt_owner_guard_releases_on_drop() -> Result<(), Box<dyn std::error::Error>> {
use super::super::intervention::{AttemptKey, AttemptOwnerIndex};
use super::super::registry::{ConnectedWorkerRegistry, WorkerDelivery};
use super::AttemptOwnerGuard;
let registry = ConnectedWorkerRegistry::default();
let (tx, _rx) = tokio::sync::mpsc::channel(1);
let types = [String::from("agent")];
let registration = registry.register_delivery_with_capabilities(
[String::from("default")],
String::from("default"),
None,
types.iter(),
WorkerDelivery::Grpc(tx),
aion_core::InterventionCapabilities::none(),
)?;
let worker = registration
.worker_id()
.ok_or("registration must assign a worker id")?;
let owners = AttemptOwnerIndex::new();
let key = AttemptKey::new(
WorkflowId::new(Uuid::nil()),
aion_core::RunId::new(Uuid::from_u128(0x11)),
ActivityId::from_sequence_position(3),
1,
);
owners.bind(key.clone(), worker);
assert_eq!(
owners.owner(&key),
Some(worker),
"owner bound before the guard"
);
{
let _guard = AttemptOwnerGuard {
owners: owners.clone(),
key: key.clone(),
};
assert_eq!(
owners.owner(&key),
Some(worker),
"still bound while in flight"
);
}
assert_eq!(
owners.owner(&key),
None,
"owner released when the dispatch returns"
);
Ok(())
}
#[test]
fn channel_format_is_pinned() {
assert_eq!(
dispatch_channel_name("remote", "gpu", None),
"aion.dispatch.remote.gpu"
);
assert_eq!(
dispatch_channel_name("local", "norn", None),
"aion.dispatch.local.norn"
);
}
#[test]
fn node_pinned_channel_appends_node_subsegment() {
assert_eq!(
dispatch_channel_name("remote", "gpu", Some("box-7")),
"aion.dispatch.remote.gpu.box-7"
);
}
#[test]
fn channel_derivation_is_stable() {
assert_eq!(
dispatch_channel_name("default", "default", None),
dispatch_channel_name("default", "default", None)
);
assert_eq!(
dispatch_channel_name("default", "default", Some("box-1")),
dispatch_channel_name("default", "default", Some("box-1"))
);
}
#[test]
fn distinct_pools_get_distinct_channels() {
assert_ne!(
dispatch_channel_name("remote", "gpu", None),
dispatch_channel_name("local", "norn", None)
);
}
#[test]
fn node_pin_separates_channels() {
let unpinned = dispatch_channel_name("remote", "gpu", None);
let box7 = dispatch_channel_name("remote", "gpu", Some("box-7"));
let box8 = dispatch_channel_name("remote", "gpu", Some("box-8"));
assert_ne!(
unpinned, box7,
"pinned dispatch must not reach unpinned pool"
);
assert_ne!(box7, box8, "distinct nodes must not collide");
}
#[test]
fn dotted_fields_do_not_collide_across_segments() {
assert_ne!(
dispatch_channel_name("a.b", "c", None),
dispatch_channel_name("a", "b.c", None),
"a '.' in a field must not bleed across the segment separator"
);
}
#[test]
fn node_subsegment_does_not_collide_with_dotted_fields() {
assert_ne!(
dispatch_channel_name("a", "b", Some("c")),
dispatch_channel_name("a", "b.c", None),
"a node sub-segment must not collide with a dotted task_queue"
);
assert_ne!(
dispatch_channel_name("a.b", "c", None),
dispatch_channel_name("a", "b", Some("c")),
"a dotted namespace must not collide with a node-pinned channel"
);
}
#[test]
fn reserved_char_shifts_stay_distinct() {
assert_ne!(
dispatch_channel_name("ns.", "tq", None),
dispatch_channel_name("ns", ".tq", None)
);
assert_ne!(
dispatch_channel_name("", "a.b", None),
dispatch_channel_name(".a", "b", None)
);
assert_ne!(
dispatch_channel_name("%2E", "x", None),
dispatch_channel_name(".", "x", None)
);
}
#[test]
fn encoding_is_injective_over_reserved_char_triples() {
let fields = ["a", "a.b", "a.", ".a", ".", "", "%", "%2E", "a%b", "%2."];
let nodes = [
None,
Some("a"),
Some("a.b"),
Some("."),
Some(""),
Some("%2E"),
];
let mut channels = std::collections::HashSet::new();
for ns in fields {
for tq in fields {
for node in nodes {
let channel = dispatch_channel_name(ns, tq, node);
assert!(
channels.insert(channel.clone()),
"collision on ({ns:?}, {tq:?}, {node:?}) -> {channel}"
);
}
}
}
}
fn row(namespace: &str, task_queue: &str) -> OutboxRow {
let workflow_id = WorkflowId::new(Uuid::new_v4());
OutboxRow {
dispatch_key: format!("{workflow_id}:0"),
workflow_id,
ordinal: 0,
run_id: Some(aion_core::RunId::new_v4()),
namespace: namespace.to_owned(),
task_queue: task_queue.to_owned(),
node: None,
activity_type: "charge-card".to_owned(),
input: Payload::new(ContentType::Json, Vec::new()),
status: OutboxStatus::Pending,
attempt: 0,
visible_after: Utc::now(),
claimed_at: None,
failure_delivered: false,
}
}
#[test]
fn channel_for_row_uses_namespace_and_task_queue_only() {
let remote_gpu = row("remote", "gpu");
let local_norn = row("local", "norn");
assert_eq!(channel_for_row(&remote_gpu), "aion.dispatch.remote.gpu");
assert_eq!(channel_for_row(&local_norn), "aion.dispatch.local.norn");
assert_ne!(channel_for_row(&remote_gpu), channel_for_row(&local_norn));
let mut other_activity = row("remote", "gpu");
other_activity.activity_type = "refund".to_owned();
assert_eq!(
channel_for_row(&remote_gpu),
channel_for_row(&other_activity),
"activity_type must not affect the channel"
);
}
#[test]
fn channel_for_row_derives_node_subchannel_when_pinned() {
let mut pinned = row("remote", "gpu");
pinned.node = Some("box-7".to_owned());
assert_eq!(channel_for_row(&pinned), "aion.dispatch.remote.gpu.box-7");
let unpinned = row("remote", "gpu");
assert_eq!(channel_for_row(&unpinned), "aion.dispatch.remote.gpu");
assert_ne!(channel_for_row(&pinned), channel_for_row(&unpinned));
}
#[test]
fn wire_node_normalizes_empty_to_none() {
assert_eq!(normalize_wire_node(None), None);
assert_eq!(normalize_wire_node(Some("")), None);
assert_eq!(normalize_wire_node(Some("box-7")), Some("box-7".to_owned()));
}
}