use std::collections::HashMap;
use std::future::Future;
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;
use crate::namespace::{CallerIdentity, NamespaceGuard};
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,
#[serde(with = "aion_core::payload_bytes")]
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,
pub max_concurrency: u32,
}
#[derive(Debug, Deserialize)]
struct AnnouncementProbe {
capabilities: aion_core::InterventionCapabilities,
max_concurrency: Option<u32>,
}
#[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
}
#[must_use]
pub fn is_connected(&self) -> bool {
self.supervisor.is_tracked(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) => return Err(classify_reply_error(&error)),
}
}
}
fn classify_reply_error(error: &LiminalServerError) -> ServerError {
if let Some(unservable) = unservable_frame_refusal(error) {
return unservable;
}
if is_connection_closed_reply_error(error) {
return ServerError::worker_connection_lost(
"liminal-push",
format!("worker connection closed before reply: {error}"),
);
}
dispatch_error("liminal-push", format!("worker reply failed: {error}"))
}
fn is_connection_closed_reply_error(error: &LiminalServerError) -> bool {
matches!(error, LiminalServerError::PushReplyDisconnected { .. })
}
fn unservable_frame_refusal(error: &LiminalServerError) -> Option<ServerError> {
let LiminalServerError::PushFrameExceedsOutboundCapacity {
needed, capacity, ..
} = error
else {
return None;
};
Some(ServerError::worker_dispatch_unservable(
"liminal-push",
format!(
"the dispatch frame is {needed} bytes and this worker connection's whole outbound \
buffer is {capacity} bytes (`outbox.{OUTBOUND_BOUND_KEY}`); no retry can carry \
it: {error}"
),
))
}
fn classify_push_error(error: &LiminalServerError) -> ServerError {
if let Some(unservable) = unservable_frame_refusal(error) {
return unservable;
}
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}"),
)
}
}
pub(crate) const OUTBOUND_BOUND_KEY: &str = "liminal_max_connection_outbound_bytes";
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>,
admission: Option<RegistrationAdmission>,
}
#[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,
admission: None,
}
}
#[must_use]
pub fn with_admission(
mut self,
guard: NamespaceGuard,
auth_enabled: bool,
handle: tokio::runtime::Handle,
) -> Self {
self.admission = Some(RegistrationAdmission {
guard,
auth_enabled,
handle,
});
self
}
#[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: AnnouncementProbe = match serde_json::from_slice(payload) {
Ok(announcement) => announcement,
Err(error) => {
tracing::warn!(
%error,
connection_pid = pid,
"capabilities tap: malformed WorkerCapabilitiesAnnouncement payload"
);
return;
}
};
let Some(max_concurrency) = announcement.max_concurrency else {
tracing::error!(
connection_pid = pid,
"capabilities tap: this worker announced no capacity and will never be selected \
for a dispatch. The liminal registration frame cannot carry a capacity, so the \
server holds a worker at capacity UNKNOWN until it announces one; a worker that \
never does is registered, idle and unselectable. It is an older build that \
predates the capacity announcement, or its announcement was truncated"
);
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_advertised_capacity(worker_id, max_concurrency)
{
Ok(true) => {}
Ok(false) => tracing::warn!(
connection_pid = pid,
worker_id = ?worker_id,
"capabilities tap: capacity announcement raced the worker's deregistration"
),
Err(error) => tracing::error!(
%error,
connection_pid = pid,
worker_id = ?worker_id,
announced_max_concurrency = max_concurrency,
"capabilities tap: worker announced a capacity the registry refused; it stays \
unknown-capacity and no dispatch will be selected for it"
),
}
if announcement.capabilities.supported.is_empty() {
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 Some(admission) = &self.admission else {
return Err(LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} rejected: this listener \
has no namespace admission bound, so it cannot judge which namespaces the \
worker may serve; a listener is commissioned with one before it accepts"
),
});
};
if admission.auth_enabled {
return Err(LiminalServerError::ListenerAccept {
message: format!(
"liminal worker registration for connection {pid} (identity `{}`) rejected: \
this server authenticates callers (auth.enabled = true) and the liminal \
registration frame carries no credential to authenticate, so the worker's \
namespaces cannot be scoped to a caller; register this worker over gRPC \
with a bearer token, or run the server with auth disabled",
registration.identity
),
});
}
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 caller = liminal_caller_identity(®istration.identity);
let proto = proto_registration_from_wire(registration);
let guard = bridge_admission(
&admission.handle,
self.registry.admit_delivery(
&admission.guard,
&caller,
&proto,
delivery,
self.intervention_capabilities.clone(),
),
)
.map_err(|error| LiminalServerError::ListenerAccept {
message: format!("liminal worker registration for connection {pid} rejected: {error}"),
})?
.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(())
}
}
struct RegistrationAdmission {
guard: NamespaceGuard,
auth_enabled: bool,
handle: tokio::runtime::Handle,
}
fn bridge_admission<F: Future>(
handle: &tokio::runtime::Handle,
admission: F,
) -> Result<F::Output, ServerError> {
match tokio::runtime::Handle::try_current() {
Err(_) => Ok(handle.block_on(admission)),
Ok(current) => match current.runtime_flavor() {
tokio::runtime::RuntimeFlavor::MultiThread => {
Ok(tokio::task::block_in_place(|| handle.block_on(admission)))
}
_ => Err(ServerError::Config {
message: "the liminal registration callback ran on a current-thread Tokio \
runtime worker, which cannot be blocked while the namespace \
admission runs; drive the listener from a multi-thread runtime \
or from a plain thread, as liminal's connection process does"
.to_owned(),
}),
},
}
}
fn liminal_caller_identity(identity: &str) -> CallerIdentity {
let subject = identity.trim();
CallerIdentity::operator(if subject.is_empty() {
"operator"
} else {
subject
})
}
fn wire_activity_to_proto(
activity: &liminal::protocol::WorkerActivityDescriptor,
) -> aion_proto::ProtoActivityDescriptor {
aion_proto::ProtoActivityDescriptor {
name: activity.name.clone(),
input_schema_json: activity.input_schema_json.clone(),
output_schema_json: activity.output_schema_json.clone(),
}
}
fn proto_registration_from_wire(
registration: &WireWorkerRegistration,
) -> aion_proto::ProtoRegisterWorker {
aion_proto::ProtoRegisterWorker {
namespaces: registration.namespaces.clone(),
activity_types: registration.activity_types.clone(),
task_queue: registration.task_queue.clone(),
node: normalize_wire_node(registration.node.as_deref()).unwrap_or_default(),
activities: registration
.activities
.iter()
.map(wire_activity_to_proto)
.collect(),
identity: registration.identity.clone(),
instance: None,
max_concurrency: None,
}
}
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::{
AnnouncementProbe, DispatchRequest, LiminalServerError, OUTBOUND_BOUND_KEY,
channel_for_row, classify_push_error, classify_reply_error, dispatch_channel_name,
normalize_wire_node,
};
#[test]
fn an_oversize_frame_on_the_reply_awaiter_is_unservable_and_names_the_bound() {
let error = LiminalServerError::PushFrameExceedsOutboundCapacity {
correlation_id: 11,
needed: 5_089_012,
capacity: 1_048_576,
queued: 0,
};
let classified = classify_reply_error(&error);
assert!(classified.is_worker_dispatch_unservable(), "{classified}");
assert!(!classified.is_worker_connection_lost(), "{classified}");
let said = classified.to_string();
for needle in ["5089012", "1048576", OUTBOUND_BOUND_KEY] {
assert!(said.contains(needle), "refusal must name {needle}: {said}");
}
}
#[test]
fn a_closed_connection_on_the_reply_awaiter_is_still_a_lost_connection() {
let classified =
classify_reply_error(&LiminalServerError::PushReplyDisconnected { correlation_id: 3 });
assert!(classified.is_worker_connection_lost(), "{classified}");
assert!(!classified.is_worker_dispatch_unservable(), "{classified}");
}
#[test]
fn any_other_reply_fault_stays_the_retryable_dispatch_class() {
let classified =
classify_reply_error(&LiminalServerError::PushReplyExpired { correlation_id: 5 });
assert!(
matches!(classified, crate::error::ServerError::WorkerDispatch { .. }),
"{classified}"
);
assert!(!classified.is_worker_dispatch_unservable(), "{classified}");
assert!(!classified.is_worker_connection_lost(), "{classified}");
}
#[test]
fn an_oversize_frame_is_unservable_and_names_the_bound_and_the_key() {
let error = LiminalServerError::PushFrameExceedsOutboundCapacity {
correlation_id: 7,
needed: 6_214_149,
capacity: 4_194_304,
queued: 0,
};
let classified = classify_push_error(&error);
assert!(classified.is_worker_dispatch_unservable(), "{classified}");
assert!(!classified.is_worker_connection_lost(), "{classified}");
assert!(!classified.is_worker_busy(), "{classified}");
let said = classified.to_string();
for needle in ["6214149", "4194304", OUTBOUND_BOUND_KEY] {
assert!(said.contains(needle), "refusal must name {needle}: {said}");
}
}
#[test]
fn the_cap_and_lost_classes_are_unchanged() {
let cap = LiminalServerError::ConnectionCapReached {
operation: "push".to_owned(),
cap: "max_pending_pushes_per_connection",
limit: 32,
};
assert!(classify_push_error(&cap).is_worker_busy());
let other = LiminalServerError::PushReplyDisconnected { correlation_id: 9 };
assert!(classify_push_error(&other).is_worker_connection_lost());
}
#[test]
fn the_refusal_names_the_real_config_key() -> Result<(), Box<dyn std::error::Error>> {
let parsed: crate::config::OutboxConfig =
toml::from_str(&format!("{OUTBOUND_BOUND_KEY} = 1"))?;
assert_eq!(parsed.liminal_max_connection_outbound_bytes, Some(1));
Ok(())
}
#[test]
fn the_dispatch_request_twins_agree_field_for_field() -> Result<(), Box<dyn std::error::Error>>
{
use base64::Engine as _;
use std::collections::BTreeSet;
let ours = DispatchRequest {
activity_type: "charge-card".to_owned(),
workflow_id: WorkflowId::new(Uuid::new_v4()),
ordinal: 7,
run_id: Some(RunId::new(Uuid::new_v4())),
completion_token: "generation-3".to_owned(),
idempotency_key: "effect-key".to_owned(),
input: br#"{"amount":42}"#.to_vec(),
attempt: 3,
labels: std::collections::BTreeMap::from([("region".to_owned(), "apac".to_owned())]),
heartbeat_window_ms: 30_000,
};
let wire = serde_json::to_value(&ours)?;
let theirs: aion_worker::runtime::liminal::DispatchRequest =
serde_json::from_value(wire.clone())?;
assert_eq!(
serde_json::to_value(&theirs)?,
wire,
"the worker twin must re-encode this frame identically"
);
let keys = |value: &serde_json::Value| -> BTreeSet<String> {
value
.as_object()
.map(|object| object.keys().cloned().collect())
.unwrap_or_default()
};
assert_eq!(keys(&wire), keys(&serde_json::to_value(&theirs)?));
assert_eq!(theirs.input, ours.input);
assert_eq!(theirs.completion_token, ours.completion_token);
assert_eq!(theirs.attempt, ours.attempt);
assert_eq!(theirs.heartbeat_window_ms, ours.heartbeat_window_ms);
let mut as_base64 = wire;
as_base64["input"] = serde_json::Value::String(
base64::engine::general_purpose::STANDARD.encode(&ours.input),
);
let ours_from_base64: DispatchRequest = serde_json::from_value(as_base64.clone())?;
let theirs_from_base64: aion_worker::runtime::liminal::DispatchRequest =
serde_json::from_value(as_base64)?;
assert_eq!(ours_from_base64, ours);
assert_eq!(theirs_from_base64.input, ours.input);
Ok(())
}
use crate::worker::registry::RegistrationOptions;
use aion_core::{ActivityId, ContentType, Payload, RunId, WorkflowId};
use aion_store::{OutboxRow, OutboxStatus};
use chrono::Utc;
use uuid::Uuid;
#[test]
fn an_announcement_without_capacity_decodes_as_absent_not_as_malformed()
-> Result<(), Box<dyn std::error::Error>> {
let pre_capacity = br#"{"capabilities":{"supported":[]}}"#;
let decoded: AnnouncementProbe = serde_json::from_slice(pre_capacity)?;
assert!(
decoded.max_concurrency.is_none(),
"the absence must be readable as absence, so the tap can name the condition and its \
consequence rather than reporting a parse failure"
);
let current = br#"{"capabilities":{"supported":[]},"max_concurrency":3}"#;
let decoded: AnnouncementProbe = serde_json::from_slice(current)?;
assert_eq!(decoded.max_concurrency, Some(3));
assert!(
serde_json::from_slice::<AnnouncementProbe>(br#"{"capabilities":7}"#).is_err(),
"a malformed payload must still fail to decode"
);
Ok(())
}
#[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(
[String::from("default")],
String::from("default"),
None,
types.iter(),
WorkerDelivery::Grpc(tx),
RegistrationOptions::identified(
String::from("liminal-transport-worker"),
crate::worker::UNBOUNDED_SENDER_WORKER_CONCURRENCY,
)
.with_intervention_capabilities(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,
started_attempt: 1,
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()));
}
}