#![cfg(feature = "liminal-transport")]
use std::collections::BTreeMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use aion::{ActivityDispatch, ActivityDispatcher};
use aion_core::{ActivityId, WorkflowId};
use aion_server::config::{
AuthConfig, AuthoringConfig, DeployConfig, ListenConfig, MetricsConfig, NamespaceConfig,
NamespaceMode, OpsConsoleAssetSource, OpsConsoleConfig, OutboxConfig, RuntimeConfig,
WebSocketConfig, WorkerConfig as ServerWorkerConfig,
};
use aion_server::worker::{LiminalConnectionNotifier, WorkerActivityDispatcher};
use aion_server::{
NamespaceResolver, ServerState, StaticScheduleNamespaces, StaticWorkflowNamespaces,
};
use aion_worker::{
ActivityFailure, ActivityRegistry, RedialTiming, WorkerConfig, serve_with_redial,
};
use liminal_server::config::{ChannelDef, ServerConfig};
use liminal_server::server::connection::ConnectionSupervisor;
use liminal_server::server::listener::ServerListener;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
type TestError = Box<dyn std::error::Error + Send + Sync>;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
const NAMESPACE: &str = "default";
const TASK_QUEUE: &str = "default";
const PROVISION: &str = "provision";
const FLAKY: &str = "flaky-provision";
const FLAKY_REASON: &str = "provision backend briefly unavailable";
const SLOW: &str = "slow-provision";
const ATTEMPT_ECHO: &str = "attempt-echo";
const DOOMED: &str = "doomed-provision";
const DEFAULT_WINDOW: Duration = Duration::from_secs(30);
const SHORT_WINDOW: Duration = Duration::from_millis(500);
const SLOW_RUNTIME: Duration = Duration::from_secs(2);
const WATCHDOG_LIMIT: Duration = Duration::from_secs(300);
const STOP_JOIN_LIMIT: Duration = Duration::from_secs(60);
fn test_error(message: impl std::fmt::Display) -> TestError {
message.to_string().into()
}
struct TestWatchdog {
_disarm: std::sync::mpsc::Sender<()>,
}
impl TestWatchdog {
fn arm(test: &'static str) -> Result<Self, TestError> {
let (disarm, armed) = std::sync::mpsc::channel::<()>();
std::thread::Builder::new()
.name(format!("watchdog-{test}"))
.spawn(move || match armed.recv_timeout(WATCHDOG_LIMIT) {
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) | Ok(()) => {}
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
use std::io::Write as _;
let mut stderr = std::io::stderr().lock();
let _ = writeln!(
stderr,
"bridge_liminal_dispatch_e2e watchdog: test `{test}` still running \
after {WATCHDOG_LIMIT:?}; aborting the process so the runner's \
slot is freed"
);
let _ = stderr.flush();
std::process::abort();
}
})
.map_err(|error| test_error(format!("watchdog thread spawn failed: {error}")))?;
Ok(Self { _disarm: disarm })
}
}
fn join_within(
handle: std::thread::JoinHandle<Result<(), TestError>>,
what: &str,
) -> Result<(), TestError> {
let deadline = Instant::now() + STOP_JOIN_LIMIT;
while !handle.is_finished() {
if Instant::now() > deadline {
return Err(test_error(format!(
"{what} did not finish within {STOP_JOIN_LIMIT:?}; \
leaking the thread and failing the test"
)));
}
std::thread::sleep(Duration::from_millis(10));
}
handle
.join()
.map_err(|_| test_error(format!("{what} panicked")))?
}
fn run_bounded(
test: &'static str,
body: impl std::future::Future<Output = Result<(), TestError>>,
) -> Result<(), TestError> {
let _watchdog = TestWatchdog::arm(test)?;
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(4)
.enable_all()
.build()
.map_err(test_error)?;
let result = runtime.block_on(body);
runtime.shutdown_timeout(STOP_JOIN_LIMIT);
result
}
macro_rules! bounded_test {
($(#[$meta:meta])* async fn $name:ident() -> Result<(), TestError> $body:block) => {
$(#[$meta])*
#[test]
fn $name() -> Result<(), TestError> {
run_bounded(concat!(module_path!(), "::", stringify!($name)), async $body)
}
};
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ProvisionInput {
resource: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ProvisionOutput {
provisioned: bool,
resource: String,
}
struct RunningServer {
listener: Option<ServerListener>,
state: ServerState,
address: SocketAddr,
sweeper_shutdown: tokio::sync::watch::Sender<bool>,
}
impl RunningServer {
fn start(heartbeat_window: Duration) -> Result<Self, TestError> {
let resolver = NamespaceResolver::authorization_only(
NamespaceMode::SharedEngine,
StaticWorkflowNamespaces::default(),
StaticScheduleNamespaces::default(),
);
let state = ServerState::from_parts(resolver, runtime_config(heartbeat_window));
let config = ServerConfig {
listen_address: "127.0.0.1:0".parse().map_err(test_error)?,
health_listen_address: reserve_loopback_port()?,
channels: Vec::<ChannelDef>::new(),
routing_rules: Vec::new(),
persistence_path: None,
cluster: None,
auth: None,
drain_timeout_ms: 30_000,
services: liminal_server::config::ServicesConfig::default(),
limits: liminal_server::config::LimitsConfig::default(),
websocket: None,
participant: None,
};
let notifier = Arc::new(
LiminalConnectionNotifier::new(state.worker_registry().clone())
.with_heartbeat_tracker(state.heartbeat_tracker().clone()),
);
let supervisor = build_supervisor_with_notifier(&config, notifier.clone())?;
if !notifier.bind_supervisor(supervisor.clone()) {
return Err(test_error("notifier supervisor was already bound"));
}
let listener = ServerListener::bind(&config, supervisor).map_err(test_error)?;
let address = listener.local_addr();
let (sweeper_shutdown, sweeper_rx) = tokio::sync::watch::channel(false);
drop(state.spawn_heartbeat_sweeper(sweeper_rx));
Ok(Self {
listener: Some(listener),
state,
address,
sweeper_shutdown,
})
}
fn bridge_dispatcher(&self) -> WorkerActivityDispatcher {
WorkerActivityDispatcher::new(
self.state.worker_registry().clone(),
NAMESPACE,
self.state.heartbeat_tracker().clone(),
)
.with_pending(self.state.pending_activities().clone())
.with_drain_state(self.state.drain_state().clone())
.with_tokio_handle(tokio::runtime::Handle::current())
.with_attempt_owners(self.state.attempt_owners().clone())
}
fn wait_for_registered_worker(&self, activity_type: &str) -> Result<(), TestError> {
let deadline = Instant::now() + CONNECT_TIMEOUT;
while Instant::now() < deadline {
if self.worker_is_registered(activity_type)? {
return Ok(());
}
std::thread::sleep(Duration::from_millis(10));
}
Err(test_error("server never registered the in-band worker"))
}
fn worker_is_registered(&self, activity_type: &str) -> Result<bool, TestError> {
Ok(self
.state
.worker_registry()
.select_worker(NAMESPACE, TASK_QUEUE, activity_type, None)
.map_err(test_error)?
.is_some())
}
fn shutdown(mut self) -> Result<(), TestError> {
let _ = self.sweeper_shutdown.send(true);
if let Some(listener) = self.listener.take() {
listener.shutdown().map_err(test_error)?;
}
Ok(())
}
}
fn build_supervisor_with_notifier(
config: &ServerConfig,
notifier: Arc<LiminalConnectionNotifier>,
) -> Result<ConnectionSupervisor, TestError> {
use liminal_server::server::connection::LiminalConnectionServices;
let services = Arc::new(LiminalConnectionServices::from_config(config).map_err(test_error)?);
ConnectionSupervisor::with_services_and_notifier(services, notifier).map_err(test_error)
}
fn reserve_loopback_port() -> Result<SocketAddr, TestError> {
let listener = std::net::TcpListener::bind("127.0.0.1:0").map_err(test_error)?;
let address = listener.local_addr().map_err(test_error)?;
drop(listener);
Ok(address)
}
fn runtime_config(heartbeat_window: Duration) -> RuntimeConfig {
RuntimeConfig {
listen: ListenConfig {
grpc: SocketAddr::from(([127, 0, 0, 1], 0)),
http: SocketAddr::from(([127, 0, 0, 1], 0)),
},
tls: None,
auth: AuthConfig {
enabled: false,
jwks_url: None,
jwks_refresh_seconds: 300,
},
ops_console: OpsConsoleConfig {
source: OpsConsoleAssetSource::Embedded,
},
namespace: NamespaceConfig {
mode: NamespaceMode::SharedEngine,
},
worker: ServerWorkerConfig {
heartbeat_window,
..ServerWorkerConfig::default()
},
websocket: WebSocketConfig {
outbound_buffer_bound: 32,
event_broadcast_capacity: Some(64),
cluster_broadcast_capacity: Some(64),
},
workflow_packages: Vec::new(),
deploy: DeployConfig::default(),
authoring: AuthoringConfig::default(),
dev: aion_server::config::DevConfig::default(),
outbox: OutboxConfig {
enabled: true,
..OutboxConfig::default()
},
observability: aion_server::config::ObservabilityConfig::with_flush_policy(64, 0),
mcp: aion_server::config::ResolvedMcpConfig::default(),
scheduler_threads: 1,
jit_threshold: None,
query_timeout: Some(Duration::from_secs(10)),
workloop_sweep_interval: Some(std::time::Duration::from_millis(50)),
default_namespace: NAMESPACE.to_owned(),
auto_create: aion_server::config::AutoCreate::Open,
max_in_flight_activities: aion_server::config::DEFAULT_MAX_IN_FLIGHT_ACTIVITIES,
drain_timeout: Duration::from_secs(30),
metrics: MetricsConfig { enabled: false },
owned_shards: Vec::new(),
cors_allowed_origins: Vec::new(),
}
}
fn worker_config() -> Result<WorkerConfig, TestError> {
WorkerConfig::builder()
.endpoint("unused-direct-address")
.namespace(NAMESPACE)
.task_queue(TASK_QUEUE)
.node("")
.identity("bridge-liminal-worker")
.max_concurrency(1)
.reconnect_initial_backoff(Duration::from_millis(5))
.reconnect_max_backoff(Duration::from_millis(20))
.reconnect_max_attempts(3)
.build()
.map_err(test_error)
}
fn worker_activity_registry(
executions: Arc<AtomicUsize>,
) -> Result<Arc<ActivityRegistry>, TestError> {
let slow_executions = Arc::clone(&executions);
let registry = ActivityRegistry::new()
.register_activity(PROVISION, move |input: ProvisionInput, _context| {
let executions = Arc::clone(&executions);
Box::pin(async move {
executions.fetch_add(1, Ordering::SeqCst);
Ok(ProvisionOutput {
provisioned: true,
resource: input.resource,
})
})
})
.map_err(test_error)?
.register_activity(FLAKY, |_input: serde_json::Value, _context| {
Box::pin(async move {
Err::<serde_json::Value, _>(ActivityFailure::retryable(FLAKY_REASON))
})
})
.map_err(test_error)?
.register_activity(SLOW, move |input: ProvisionInput, _context| {
let executions = Arc::clone(&slow_executions);
Box::pin(async move {
tokio::time::sleep(SLOW_RUNTIME).await;
executions.fetch_add(1, Ordering::SeqCst);
Ok(ProvisionOutput {
provisioned: true,
resource: input.resource,
})
})
})
.map_err(test_error)?
.register_activity(ATTEMPT_ECHO, |_input: serde_json::Value, context| {
Box::pin(async move { Ok(serde_json::json!({ "attempt": context.attempt() })) })
})
.map_err(test_error)?;
Ok(Arc::new(registry))
}
struct ServedWorker {
stop: Arc<AtomicBool>,
handle: Option<std::thread::JoinHandle<Result<(), TestError>>>,
}
impl ServedWorker {
fn spawn(address: String, registry: Arc<ActivityRegistry>) -> Self {
let stop = Arc::new(AtomicBool::new(false));
let worker_stop = Arc::clone(&stop);
let handle = std::thread::spawn(move || -> Result<(), TestError> {
let config = worker_config()?;
serve_with_redial(
vec![address],
&config,
®istry,
RedialTiming::new(Duration::from_millis(5), Duration::from_millis(20)),
&worker_stop,
None,
|| {},
)
.map_err(test_error)
});
Self {
stop,
handle: Some(handle),
}
}
fn stop(mut self) -> Result<(), TestError> {
self.stop.store(true, Ordering::SeqCst);
if let Some(handle) = self.handle.take() {
join_within(handle, "worker serve thread")?;
}
Ok(())
}
}
fn dispatch_request(
workflow_id: &WorkflowId,
ordinal: u64,
activity_type: &str,
input: &serde_json::Value,
) -> ActivityDispatch {
ActivityDispatch {
namespace: NAMESPACE.to_owned(),
task_queue: TASK_QUEUE.to_owned(),
node: None,
workflow_id: workflow_id.clone(),
run_id: aion_core::RunId::new_v4(),
activity_id: ActivityId::from_sequence_position(ordinal),
name: activity_type.to_owned(),
input: input.to_string(),
config: "{}".to_owned(),
attempt: 1,
labels: BTreeMap::new(),
advisory: false,
}
}
async fn dispatch_via_seam(
dispatcher: &Arc<WorkerActivityDispatcher>,
request: ActivityDispatch,
) -> Result<Result<String, String>, TestError> {
let dispatcher = Arc::clone(dispatcher);
tokio::time::timeout(
Duration::from_secs(20),
tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(request))),
)
.await
.map_err(|_| test_error("bridge dispatch did not resolve within the test deadline"))?
.map_err(test_error)
}
#[path = "bridge_liminal_dispatch_e2e/bridge_liminal_dispatch_cases.rs"]
mod bridge_liminal_dispatch_cases;