use std::{path::PathBuf, sync::Arc};
use aion::{
ActivityDispatcher, EngineBuilder, RuntimeHandle, SignalRouter, signal::ConcreteSignalRouter,
};
use aion_store::{EventStore, NamespaceStore, OutboxStore};
#[cfg(feature = "libsql-backend")]
use aion_store_libsql::LibSqlStore;
use crate::dev_ui::{ActivityMockRegistry, DevMockingDispatcher};
#[cfg(feature = "auth")]
use crate::auth::JwksCache;
use crate::{
config::{RuntimeConfig, ServerConfig, StoreBackend, StoreConfig},
error::ServerError,
namespace::{NamespaceGuard, NamespaceMinter, resolver::NamespaceResolver},
observability::{
Metrics, health::HealthState, instrumented_store::InstrumentedEventStore,
metrics::MetricsError,
},
shutdown::DrainState,
worker::{
ConnectedWorkerRegistry, HeartbeatTracker, PendingActivities, WorkerActivityDispatcher,
},
};
#[derive(Clone)]
pub struct ServerState {
inner: Arc<ServerStateInner>,
}
struct ServerStateInner {
namespace_guard: NamespaceGuard,
runtime: RuntimeConfig,
worker_registry: ConnectedWorkerRegistry,
pending_activities: PendingActivities,
heartbeat_tracker: HeartbeatTracker,
drain_state: DrainState,
metrics: Option<Metrics>,
health: Option<HealthState>,
activity_mock_registry: Option<ActivityMockRegistry>,
outbox_store: Option<Arc<dyn OutboxStore>>,
namespace_store: Arc<dyn NamespaceStore>,
outbox_wake: Arc<tokio::sync::Notify>,
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
transcript_publisher: crate::activity_publisher::ActivityEventPublisher,
attempt_owners: crate::worker::AttemptOwnerIndex,
cluster_self_node: Option<String>,
#[cfg(feature = "haematite-backend")]
cluster_responder: Option<aion_store_haematite::ClusterResponder>,
#[cfg(feature = "haematite-backend")]
cluster_store: Option<Arc<aion_store_haematite::HaematiteStore>>,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec<crate::cluster::WatchedPeer>,
#[cfg(feature = "haematite-backend")]
shard_directory: Option<Arc<crate::routing::StaticShardDirectory>>,
#[cfg(feature = "haematite-backend")]
request_forwarder: Option<Arc<dyn crate::routing::RequestForwarder>>,
#[cfg(feature = "auth")]
jwks_cache: Option<JwksCache>,
}
impl ServerState {
const FALLBACK_CLUSTER_BROADCAST_CAPACITY: std::num::NonZeroUsize =
match std::num::NonZeroUsize::new(64) {
Some(value) => value,
None => std::num::NonZeroUsize::MIN,
};
pub async fn build(config: ServerConfig) -> Result<Self, ServerError> {
let (store_config, runtime) = config.into_parts();
let connected = connect_store(store_config).await?;
Self::build_with_connected_store(connected, runtime).await
}
pub async fn build_with_store<S>(store: S, runtime: RuntimeConfig) -> Result<Self, ServerError>
where
S: EventStore + NamespaceStore,
{
let leaf = Arc::new(store);
let namespace_store: Arc<dyn NamespaceStore> = leaf.clone();
Self::build_with_connected_store(
ConnectedStore::local(leaf, None, namespace_store),
runtime,
)
.await
}
async fn build_with_connected_store(
connected: ConnectedStore,
runtime: RuntimeConfig,
) -> Result<Self, ServerError> {
let outbox_store = connected.outbox_store;
let bootstrap_coordinator = connected.bootstrap_coordinator;
#[cfg(feature = "haematite-backend")]
let cluster_responder = connected.cluster_responder;
#[cfg(feature = "haematite-backend")]
let cluster_store = connected.cluster_store;
#[cfg(feature = "haematite-backend")]
let watched_peers = connected.watched_peers;
#[cfg(feature = "haematite-backend")]
let cluster_self_node = connected.self_node_id.clone();
#[cfg(not(feature = "haematite-backend"))]
let cluster_self_node: Option<String> = None;
#[cfg(feature = "haematite-backend")]
let RoutingState {
shard_directory,
request_forwarder,
} = build_routing_state(
cluster_store.as_ref(),
connected.directory_peers,
connected.self_node_id,
);
let (event_broadcast_capacity, query_timeout) = required_engine_seams(&runtime)?;
let (cluster_publisher, transcript_publisher) =
build_real_time_publishers(&runtime, connected.observability_store)?;
let metrics = Metrics::new().map_err(|error| metrics_config_error(&error))?;
let outbox_wake = Arc::new(tokio::sync::Notify::new());
let instrumented_store = Arc::new(
InstrumentedEventStore::new(
connected.event_store,
metrics.clone(),
runtime.default_namespace.clone(),
)
.with_outbox_wake(Arc::clone(&outbox_wake)),
);
let exported_metrics = runtime.metrics.enabled.then_some(metrics.clone());
let worker_registry = ConnectedWorkerRegistry::default()
.with_cluster_publisher(cluster_publisher.clone())
.with_namespace_minting(connected.namespace_store.clone(), runtime.auto_create);
let pending_activities = PendingActivities::default();
let heartbeat_tracker = HeartbeatTracker::new(runtime.worker.heartbeat_window);
let drain_state = DrainState::default();
let (dispatcher, attempt_owners) = build_bridge_dispatcher(
&runtime,
&worker_registry,
&pending_activities,
&heartbeat_tracker,
&drain_state,
);
let (activity_dispatcher, activity_mock_registry) =
decorate_activity_dispatcher(dispatcher, runtime.dev.enabled);
let engine = build_engine(EngineAssembly {
instrumented_store: &instrumented_store,
event_broadcast_capacity,
query_timeout,
activity_dispatcher,
active_registry: Arc::new(aion::Registry::default()),
bootstrap_coordinator,
runtime: &runtime,
})
.await?;
let engine = Arc::new(engine);
install_outbox_delivery(&pending_activities, &engine, runtime.outbox.enabled);
let resolver = NamespaceResolver::from_config(runtime.namespace.clone(), engine);
#[cfg(feature = "auth")]
let jwks_cache = build_jwks_cache(&runtime).await?;
Ok(Self {
inner: Arc::new(ServerStateInner {
namespace_guard: NamespaceGuard::new(resolver),
runtime,
worker_registry,
pending_activities,
heartbeat_tracker,
drain_state,
metrics: exported_metrics,
health: Some(HealthState::new(instrumented_store, true)),
activity_mock_registry,
outbox_store,
namespace_store: connected.namespace_store,
outbox_wake,
cluster_publisher,
transcript_publisher,
attempt_owners,
cluster_self_node,
#[cfg(feature = "haematite-backend")]
cluster_responder,
#[cfg(feature = "haematite-backend")]
cluster_store,
#[cfg(feature = "haematite-backend")]
watched_peers,
#[cfg(feature = "haematite-backend")]
shard_directory,
#[cfg(feature = "haematite-backend")]
request_forwarder,
#[cfg(feature = "auth")]
jwks_cache,
}),
})
}
#[must_use]
pub fn from_parts(namespace_resolver: NamespaceResolver, runtime: RuntimeConfig) -> Self {
Self::from_parts_with_namespace_store(
namespace_resolver,
runtime,
Arc::new(aion_store::InMemoryStore::default()),
)
}
#[must_use]
pub fn from_parts_with_namespace_store(
namespace_resolver: NamespaceResolver,
runtime: RuntimeConfig,
namespace_store: Arc<dyn NamespaceStore>,
) -> Self {
let heartbeat_tracker = HeartbeatTracker::new(runtime.worker.heartbeat_window);
let bounds = transcript_bounds(&runtime);
Self {
inner: Arc::new(ServerStateInner {
namespace_guard: NamespaceGuard::new(namespace_resolver),
runtime,
worker_registry: ConnectedWorkerRegistry::default(),
pending_activities: PendingActivities::default(),
heartbeat_tracker,
drain_state: DrainState::default(),
metrics: None,
health: None,
activity_mock_registry: None,
outbox_store: None,
namespace_store,
outbox_wake: Arc::new(tokio::sync::Notify::new()),
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher::new(
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
),
transcript_publisher: build_transcript_publisher(
None,
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
bounds,
),
attempt_owners: crate::worker::AttemptOwnerIndex::new(),
cluster_self_node: None,
#[cfg(feature = "haematite-backend")]
cluster_responder: None,
#[cfg(feature = "haematite-backend")]
cluster_store: None,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
shard_directory: None,
#[cfg(feature = "haematite-backend")]
request_forwarder: None,
#[cfg(feature = "auth")]
jwks_cache: None,
}),
}
}
#[cfg(feature = "auth")]
#[must_use]
pub fn from_parts_with_namespace_store_and_jwks(
namespace_resolver: NamespaceResolver,
runtime: RuntimeConfig,
namespace_store: Arc<dyn NamespaceStore>,
jwks_cache: JwksCache,
) -> Self {
let heartbeat_tracker = HeartbeatTracker::new(runtime.worker.heartbeat_window);
let bounds = transcript_bounds(&runtime);
Self {
inner: Arc::new(ServerStateInner {
namespace_guard: NamespaceGuard::new(namespace_resolver),
runtime,
worker_registry: ConnectedWorkerRegistry::default(),
pending_activities: PendingActivities::default(),
heartbeat_tracker,
drain_state: DrainState::default(),
metrics: None,
health: None,
activity_mock_registry: None,
outbox_store: None,
namespace_store,
outbox_wake: Arc::new(tokio::sync::Notify::new()),
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher::new(
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
),
transcript_publisher: build_transcript_publisher(
None,
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
bounds,
),
attempt_owners: crate::worker::AttemptOwnerIndex::new(),
cluster_self_node: None,
#[cfg(feature = "haematite-backend")]
cluster_responder: None,
#[cfg(feature = "haematite-backend")]
cluster_store: None,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
shard_directory: None,
#[cfg(feature = "haematite-backend")]
request_forwarder: None,
jwks_cache: Some(jwks_cache),
}),
}
}
#[cfg(feature = "auth")]
#[must_use]
pub fn from_parts_with_jwks(
namespace_resolver: NamespaceResolver,
runtime: RuntimeConfig,
jwks_cache: JwksCache,
) -> Self {
let heartbeat_tracker = HeartbeatTracker::new(runtime.worker.heartbeat_window);
let bounds = transcript_bounds(&runtime);
Self {
inner: Arc::new(ServerStateInner {
namespace_guard: NamespaceGuard::new(namespace_resolver),
runtime,
worker_registry: ConnectedWorkerRegistry::default(),
pending_activities: PendingActivities::default(),
heartbeat_tracker,
drain_state: DrainState::default(),
metrics: None,
health: None,
activity_mock_registry: None,
outbox_store: None,
namespace_store: Arc::new(aion_store::InMemoryStore::default()),
outbox_wake: Arc::new(tokio::sync::Notify::new()),
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher::new(
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
),
transcript_publisher: build_transcript_publisher(
None,
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
bounds,
),
attempt_owners: crate::worker::AttemptOwnerIndex::new(),
cluster_self_node: None,
#[cfg(feature = "haematite-backend")]
cluster_responder: None,
#[cfg(feature = "haematite-backend")]
cluster_store: None,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
shard_directory: None,
#[cfg(feature = "haematite-backend")]
request_forwarder: None,
jwks_cache: Some(jwks_cache),
}),
}
}
#[must_use]
pub fn from_parts_with_registry(
namespace_resolver: NamespaceResolver,
runtime: RuntimeConfig,
worker_registry: ConnectedWorkerRegistry,
) -> Self {
let heartbeat_tracker = HeartbeatTracker::new(runtime.worker.heartbeat_window);
let bounds = transcript_bounds(&runtime);
Self {
inner: Arc::new(ServerStateInner {
namespace_guard: NamespaceGuard::new(namespace_resolver),
runtime,
worker_registry,
pending_activities: PendingActivities::default(),
heartbeat_tracker,
drain_state: DrainState::default(),
metrics: None,
health: None,
activity_mock_registry: None,
outbox_store: None,
namespace_store: Arc::new(aion_store::InMemoryStore::default()),
outbox_wake: Arc::new(tokio::sync::Notify::new()),
cluster_publisher: crate::cluster_publisher::ClusterEventPublisher::new(
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
),
transcript_publisher: build_transcript_publisher(
None,
Self::FALLBACK_CLUSTER_BROADCAST_CAPACITY,
bounds,
),
attempt_owners: crate::worker::AttemptOwnerIndex::new(),
cluster_self_node: None,
#[cfg(feature = "haematite-backend")]
cluster_responder: None,
#[cfg(feature = "haematite-backend")]
cluster_store: None,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
shard_directory: None,
#[cfg(feature = "haematite-backend")]
request_forwarder: None,
#[cfg(feature = "auth")]
jwks_cache: None,
}),
}
}
#[must_use]
pub fn namespace_guard(&self) -> &NamespaceGuard {
&self.inner.namespace_guard
}
#[must_use]
pub fn deploy_guard(&self) -> crate::deploy::DeployGuard {
crate::deploy::DeployGuard::new(self.inner.namespace_guard.resolver().clone())
}
#[must_use]
pub fn runtime_config(&self) -> &RuntimeConfig {
&self.inner.runtime
}
#[must_use]
pub fn worker_registry(&self) -> &ConnectedWorkerRegistry {
&self.inner.worker_registry
}
#[must_use]
pub fn cluster_publisher(&self) -> &crate::cluster_publisher::ClusterEventPublisher {
&self.inner.cluster_publisher
}
#[must_use]
pub fn transcript_publisher(&self) -> &crate::activity_publisher::ActivityEventPublisher {
&self.inner.transcript_publisher
}
#[must_use]
pub fn attempt_owners(&self) -> &crate::worker::AttemptOwnerIndex {
&self.inner.attempt_owners
}
#[must_use]
pub fn intervention_router(&self) -> crate::worker::InterventionRouter {
let transport: std::sync::Arc<dyn crate::worker::InterventionTransport> = {
#[cfg(feature = "liminal-transport")]
{
std::sync::Arc::new(crate::worker::LiminalInterventionTransport)
}
#[cfg(not(feature = "liminal-transport"))]
{
std::sync::Arc::new(NullInterventionTransport)
}
};
crate::worker::InterventionRouter::new(
self.inner.worker_registry.clone(),
self.inner.attempt_owners.clone(),
transport,
)
.with_transcript_publisher(self.inner.transcript_publisher.clone())
}
#[must_use]
pub fn cluster_self_node(&self) -> Option<&str> {
self.inner.cluster_self_node.as_deref()
}
pub fn engine(&self) -> Result<Arc<aion::Engine>, ServerError> {
self.inner
.namespace_guard
.resolver()
.engine()
.map(Arc::clone)
}
#[must_use]
pub fn pending_activities(&self) -> &PendingActivities {
&self.inner.pending_activities
}
#[must_use]
pub fn heartbeat_tracker(&self) -> &HeartbeatTracker {
&self.inner.heartbeat_tracker
}
#[must_use]
pub fn drain_state(&self) -> &DrainState {
&self.inner.drain_state
}
#[must_use]
pub fn metrics(&self) -> Option<&Metrics> {
self.inner.metrics.as_ref()
}
#[must_use]
pub fn health(&self) -> Option<&HealthState> {
self.inner.health.as_ref()
}
#[must_use]
pub fn activity_mock_registry(&self) -> Option<&ActivityMockRegistry> {
self.inner.activity_mock_registry.as_ref()
}
#[must_use]
pub fn outbox_store(&self) -> Option<Arc<dyn OutboxStore>> {
self.inner.outbox_store.clone()
}
#[must_use]
pub fn namespace_store(&self) -> &Arc<dyn NamespaceStore> {
&self.inner.namespace_store
}
#[must_use]
pub fn namespace_minter(&self) -> NamespaceMinter {
NamespaceMinter::new(
Arc::clone(&self.inner.namespace_store),
self.inner.runtime.auto_create,
)
.with_cluster_publisher(self.inner.cluster_publisher.clone())
}
#[must_use]
pub fn outbox_wake(&self) -> Arc<tokio::sync::Notify> {
Arc::clone(&self.inner.outbox_wake)
}
#[cfg(feature = "haematite-backend")]
#[must_use]
pub fn is_clustered(&self) -> bool {
self.inner.cluster_responder.is_some()
}
#[cfg(feature = "haematite-backend")]
#[must_use]
pub fn cluster_store(&self) -> Option<&Arc<aion_store_haematite::HaematiteStore>> {
self.inner.cluster_store.as_ref()
}
#[cfg(feature = "haematite-backend")]
#[must_use]
pub fn shard_directory(&self) -> Option<&Arc<crate::routing::StaticShardDirectory>> {
self.inner.shard_directory.as_ref()
}
#[cfg(feature = "haematite-backend")]
#[must_use]
pub fn request_forwarder(&self) -> Option<&Arc<dyn crate::routing::RequestForwarder>> {
self.inner.request_forwarder.as_ref()
}
#[must_use]
pub fn spawn_heartbeat_sweeper(
&self,
shutdown: tokio::sync::watch::Receiver<bool>,
) -> tokio::task::JoinHandle<()> {
let sweeper = crate::worker::HeartbeatSweeper::new(
self.inner.heartbeat_tracker.clone(),
self.inner.worker_registry.clone(),
self.inner.pending_activities.clone(),
self.inner.drain_state.clone(),
self.inner.runtime.worker.heartbeat_window,
);
tokio::spawn(sweeper.run(shutdown))
}
#[cfg(feature = "haematite-backend")]
pub fn spawn_cluster_supervisor(
&self,
config: crate::cluster::SupervisorConfig,
shutdown: tokio::sync::watch::Receiver<bool>,
) -> Result<bool, ServerError> {
let Some(cluster_store) = self.inner.cluster_store.clone() else {
return Ok(false);
};
if self.inner.watched_peers.is_empty() {
return Ok(false);
}
let engine = Arc::clone(self.inner.namespace_guard.resolver().engine()?);
let publisher = Arc::new(self.inner.cluster_publisher.clone());
let self_node = self.inner.cluster_self_node.clone().unwrap_or_default();
let adopter = Arc::new(crate::cluster::OutboxSettlingAdopter::new(
engine,
self.inner.outbox_store.clone(),
));
let supervisor = crate::cluster::ClusterSupervisor::new(
cluster_store,
adopter,
self.inner.watched_peers.clone(),
config,
)
.with_publisher(publisher, self_node);
if !supervisor.watches_any() {
return Ok(false);
}
tokio::spawn(supervisor.run(shutdown));
Ok(true)
}
#[cfg(feature = "auth")]
#[must_use]
pub fn jwks_cache(&self) -> Option<&JwksCache> {
self.inner.jwks_cache.as_ref()
}
pub fn shutdown(&self) -> Result<(), ServerError> {
self.inner.namespace_guard.resolver().shutdown_engine()
}
}
#[cfg(feature = "auth")]
async fn build_jwks_cache(runtime: &RuntimeConfig) -> Result<Option<JwksCache>, ServerError> {
if !runtime.auth.enabled {
return Ok(None);
}
let Some(url) = runtime.auth.jwks_url.clone() else {
return Err(ServerError::Config {
message: "auth.jwks_url must not be empty when auth.enabled is true".to_owned(),
});
};
let interval = std::time::Duration::from_secs(runtime.auth.jwks_refresh_seconds);
let cache = JwksCache::new(url, interval)
.await
.map_err(|error| ServerError::Config {
message: format!("auth jwks initial fetch failed: {error}"),
})?;
Ok(Some(cache))
}
fn metrics_config_error(error: &MetricsError) -> ServerError {
ServerError::Config {
message: error.to_string(),
}
}
struct EngineAssembly<'a> {
instrumented_store: &'a Arc<InstrumentedEventStore>,
event_broadcast_capacity: std::num::NonZeroUsize,
query_timeout: std::time::Duration,
activity_dispatcher: Arc<dyn ActivityDispatcher>,
active_registry: Arc<aion::Registry>,
bootstrap_coordinator: bool,
runtime: &'a RuntimeConfig,
}
async fn build_engine(assembly: EngineAssembly<'_>) -> Result<aion::Engine, ServerError> {
let mut search_attribute_schema = aion_core::SearchAttributeSchema::new();
search_attribute_schema
.register(
crate::namespace::NAMESPACE_ATTRIBUTE,
aion_core::SearchAttributeType::String,
)
.map_err(|error| ServerError::Config {
message: format!("failed to register namespace search attribute: {error}"),
})?;
search_attribute_schema
.register(
crate::namespace::TASK_QUEUE_ATTRIBUTE,
aion_core::SearchAttributeType::String,
)
.map_err(|error| ServerError::Config {
message: format!("failed to register task_queue search attribute: {error}"),
})?;
let runtime = assembly.runtime;
let builder = EngineBuilder::new()
.store_arc(assembly.instrumented_store.clone())
.event_streaming(assembly.event_broadcast_capacity)
.in_memory_visibility()
.search_attribute_schema(search_attribute_schema)
.scheduler_threads(runtime.scheduler_threads)
.outbox_enabled(runtime.outbox.enabled)
.activity_dispatcher(assembly.activity_dispatcher)
.active_registry(assembly.active_registry)
.production_recovery_seam()
.signal_router_factory(|runtime: Arc<RuntimeHandle>, handoff| {
Arc::new(ConcreteSignalRouter::new(runtime, handoff)) as Arc<dyn SignalRouter>
})
.query_timeout(assembly.query_timeout)
.bootstrap_schedule_coordinator(assembly.bootstrap_coordinator)
.load_workflow_sources(runtime.workflow_packages.iter().map(PathBuf::as_path));
let builder = if runtime.owned_shards.is_empty() {
builder
} else {
builder.owned_shards(runtime.owned_shards.iter().copied())
};
builder.build().await.map_err(ServerError::from)
}
fn required_engine_seams(
runtime: &RuntimeConfig,
) -> Result<(std::num::NonZeroUsize, std::time::Duration), ServerError> {
let event_broadcast_capacity = runtime
.websocket
.event_broadcast_capacity
.and_then(std::num::NonZeroUsize::new)
.ok_or_else(|| ServerError::Config {
message: crate::config::EVENT_BROADCAST_CAPACITY_REQUIRED.to_owned(),
})?;
let query_timeout = runtime
.query_timeout
.filter(|timeout| !timeout.is_zero())
.ok_or_else(|| ServerError::Config {
message: crate::config::QUERY_TIMEOUT_REQUIRED.to_owned(),
})?;
Ok((event_broadcast_capacity, query_timeout))
}
fn install_outbox_delivery(
pending_activities: &PendingActivities,
engine: &Arc<aion::Engine>,
outbox_enabled: bool,
) {
if outbox_enabled {
let callback = Arc::new(crate::worker::ServerOutboxDeliveryCallback::new(
Arc::clone(engine),
));
pending_activities.set_outbox_delivery(callback);
}
}
fn build_bridge_dispatcher(
runtime: &RuntimeConfig,
worker_registry: &ConnectedWorkerRegistry,
pending_activities: &PendingActivities,
heartbeat_tracker: &HeartbeatTracker,
drain_state: &DrainState,
) -> (WorkerActivityDispatcher, crate::worker::AttemptOwnerIndex) {
let attempt_owners = crate::worker::AttemptOwnerIndex::new();
let dispatcher = WorkerActivityDispatcher::new(
worker_registry.clone(),
runtime.default_namespace.clone(),
heartbeat_tracker.clone(),
)
.with_pending(pending_activities.clone())
.with_drain_state(drain_state.clone())
.with_tokio_handle(tokio::runtime::Handle::current())
.with_attempt_owners(attempt_owners.clone());
(dispatcher, attempt_owners)
}
fn decorate_activity_dispatcher(
dispatcher: WorkerActivityDispatcher,
dev_enabled: bool,
) -> (Arc<dyn ActivityDispatcher>, Option<ActivityMockRegistry>) {
if dev_enabled {
let registry = ActivityMockRegistry::new();
let decorated = DevMockingDispatcher::new(Arc::new(dispatcher), registry.clone());
(Arc::new(decorated), Some(registry))
} else {
(Arc::new(dispatcher), None)
}
}
fn required_cluster_broadcast_capacity(
runtime: &RuntimeConfig,
) -> Result<std::num::NonZeroUsize, ServerError> {
runtime
.websocket
.cluster_broadcast_capacity
.and_then(std::num::NonZeroUsize::new)
.ok_or_else(|| ServerError::Config {
message: crate::config::CLUSTER_BROADCAST_CAPACITY_REQUIRED.to_owned(),
})
}
fn build_real_time_publishers(
runtime: &RuntimeConfig,
observability_store: Option<Arc<dyn aion_store::ObservabilityStore>>,
) -> Result<
(
crate::cluster_publisher::ClusterEventPublisher,
crate::activity_publisher::ActivityEventPublisher,
),
ServerError,
> {
let capacity = required_cluster_broadcast_capacity(runtime)?;
Ok((
crate::cluster_publisher::ClusterEventPublisher::new(capacity),
build_transcript_publisher(observability_store, capacity, transcript_bounds(runtime)),
))
}
fn transcript_bounds(runtime: &RuntimeConfig) -> crate::activity_bounds::TranscriptBounds {
crate::activity_bounds::TranscriptBounds {
max_event_bytes: runtime.observability.max_event_bytes,
max_stream_events: runtime.observability.max_stream_events,
}
}
fn build_transcript_publisher(
observability_store: Option<Arc<dyn aion_store::ObservabilityStore>>,
capacity: std::num::NonZeroUsize,
bounds: crate::activity_bounds::TranscriptBounds,
) -> crate::activity_publisher::ActivityEventPublisher {
let store = observability_store
.unwrap_or_else(|| Arc::new(aion_store::InMemoryObservabilityStore::default()));
crate::activity_publisher::ActivityEventPublisher::new(store, capacity).with_bounds(bounds)
}
#[cfg(feature = "haematite-backend")]
struct RoutingState {
shard_directory: Option<Arc<crate::routing::StaticShardDirectory>>,
request_forwarder: Option<Arc<dyn crate::routing::RequestForwarder>>,
}
#[cfg(feature = "haematite-backend")]
fn build_routing_state(
cluster_store: Option<&Arc<aion_store_haematite::HaematiteStore>>,
directory_peers: Vec<crate::routing::DirectoryPeer>,
self_node_id: Option<String>,
) -> RoutingState {
let Some(store) = cluster_store else {
return RoutingState {
shard_directory: None,
request_forwarder: None,
};
};
RoutingState {
shard_directory: Some(Arc::new(crate::routing::StaticShardDirectory::new(
Arc::clone(store),
directory_peers,
self_node_id,
))),
request_forwarder: Some(Arc::new(crate::routing::GrpcRequestForwarder::new())),
}
}
struct ConnectedStore {
event_store: Arc<dyn EventStore>,
outbox_store: Option<Arc<dyn OutboxStore>>,
namespace_store: Arc<dyn NamespaceStore>,
observability_store: Option<Arc<dyn aion_store::ObservabilityStore>>,
bootstrap_coordinator: bool,
#[cfg(feature = "haematite-backend")]
cluster_responder: Option<aion_store_haematite::ClusterResponder>,
#[cfg(feature = "haematite-backend")]
cluster_store: Option<Arc<aion_store_haematite::HaematiteStore>>,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec<crate::cluster::WatchedPeer>,
#[cfg(feature = "haematite-backend")]
directory_peers: Vec<crate::routing::DirectoryPeer>,
#[cfg(feature = "haematite-backend")]
self_node_id: Option<String>,
}
impl ConnectedStore {
fn local(
event_store: Arc<dyn EventStore>,
outbox_store: Option<Arc<dyn OutboxStore>>,
namespace_store: Arc<dyn NamespaceStore>,
) -> Self {
Self {
event_store,
outbox_store,
namespace_store,
observability_store: None,
bootstrap_coordinator: true,
#[cfg(feature = "haematite-backend")]
cluster_responder: None,
#[cfg(feature = "haematite-backend")]
cluster_store: None,
#[cfg(feature = "haematite-backend")]
watched_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
directory_peers: Vec::new(),
#[cfg(feature = "haematite-backend")]
self_node_id: None,
}
}
}
async fn connect_store(config: StoreConfig) -> Result<ConnectedStore, ServerError> {
match config.backend {
StoreBackend::Memory => {
let leaf = Arc::new(aion_store::InMemoryStore::default());
let namespace_store: Arc<dyn NamespaceStore> = leaf.clone();
Ok(ConnectedStore::local(leaf, None, namespace_store))
}
StoreBackend::LibSql => {
#[cfg(feature = "libsql-backend")]
{
connect_libsql_store(config).await
}
#[cfg(not(feature = "libsql-backend"))]
{
let _ = config;
connect_libsql_store_unavailable()
}
}
StoreBackend::Haematite => {
#[cfg(feature = "haematite-backend")]
{
connect_haematite_store(config).await
}
#[cfg(not(feature = "haematite-backend"))]
{
let _ = config;
connect_haematite_store_unavailable()
}
}
}
}
#[cfg(feature = "libsql-backend")]
async fn connect_libsql_store(config: StoreConfig) -> Result<ConnectedStore, ServerError> {
let Some(url) = config.url else {
return Err(ServerError::Config {
message: "store.url must not be empty when store.backend is libsql".to_owned(),
});
};
let store = LibSqlStore::open(url.clone())
.await
.map_err(ServerError::from)?;
store
.validate_event_compatibility()
.await
.map_err(|error| match error {
aion_store::StoreError::Serialization(_) => ServerError::Config {
message: format!(
"Database schema mismatch — delete {url} and restart, or run migrations."
),
},
other => ServerError::from(other),
})?;
let leaf = Arc::new(store);
let event_store: Arc<dyn EventStore> = leaf.clone();
let namespace_store: Arc<dyn NamespaceStore> = leaf.clone();
let outbox_store: Arc<dyn OutboxStore> = leaf;
Ok(ConnectedStore::local(
event_store,
Some(outbox_store),
namespace_store,
))
}
#[cfg(not(feature = "libsql-backend"))]
fn connect_libsql_store_unavailable() -> Result<ConnectedStore, ServerError> {
Err(ServerError::Config {
message: "store.backend = libsql requires the aion-server `libsql-backend` feature"
.to_owned(),
})
}
#[cfg(feature = "haematite-backend")]
async fn connect_haematite_store(config: StoreConfig) -> Result<ConnectedStore, ServerError> {
let Some(data_dir) = config.data_dir else {
return Err(ServerError::Config {
message: "store.data_dir must not be empty when store.backend is haematite".to_owned(),
});
};
let shard_count = config.shard_count;
let owned_shards = config.owned_shards.clone();
let cluster = config.cluster.clone();
let watched_peers: Vec<crate::cluster::WatchedPeer> = cluster
.as_ref()
.map(|cluster| {
cluster
.peers
.iter()
.map(|peer| crate::cluster::WatchedPeer {
name: peer.name.clone(),
owned_shards: peer.owned_shards.clone(),
})
.collect()
})
.unwrap_or_default();
let directory_peers: Vec<crate::routing::DirectoryPeer> = cluster
.as_ref()
.map(|cluster| {
cluster
.peers
.iter()
.map(|peer| crate::routing::DirectoryPeer {
name: peer.name.clone(),
owned_shards: peer.owned_shards.clone(),
grpc_addr: peer.grpc_address,
})
.collect()
})
.unwrap_or_default();
let self_node_id: Option<String> = cluster.as_ref().map(|cluster| cluster.node_id.clone());
let (store, responder) =
tokio::task::spawn_blocking(move || build_haematite_store(&data_dir, shard_count, cluster))
.await
.map_err(|error| ServerError::Config {
message: format!("haematite store initialization task failed: {error}"),
})??;
let bootstrap_coordinator = if owned_shards.is_empty() {
true
} else {
store.set_owned_shards(owned_shards.iter().copied());
store.owns_workflow_shard(&aion::schedule_coordinator_workflow_id())
};
let leaf = Arc::new(store);
let event_store: Arc<dyn EventStore> = leaf.clone();
let outbox_store: Arc<dyn OutboxStore> = leaf.clone();
let namespace_store: Arc<dyn NamespaceStore> = leaf.clone();
let observability_store: Arc<dyn aion_store::ObservabilityStore> = leaf.clone();
let cluster_store = responder.as_ref().map(|_| leaf);
let (watched_peers, directory_peers, self_node_id) = if cluster_store.is_some() {
(watched_peers, directory_peers, self_node_id)
} else {
(Vec::new(), Vec::new(), None)
};
Ok(ConnectedStore {
event_store,
outbox_store: Some(outbox_store),
namespace_store,
observability_store: Some(observability_store),
bootstrap_coordinator,
cluster_responder: responder,
cluster_store,
watched_peers,
directory_peers,
self_node_id,
})
}
#[cfg(feature = "haematite-backend")]
fn build_haematite_store(
data_dir: &str,
shard_count: usize,
cluster: Option<crate::config::ClusterConfig>,
) -> Result<
(
aion_store_haematite::HaematiteStore,
Option<aion_store_haematite::ClusterResponder>,
),
ServerError,
> {
use aion_store_haematite::{ClusterBootstrap, HaematiteStore};
let Some(cluster) = cluster else {
let path = std::path::Path::new(data_dir);
let store = if path.join("config.json").exists() {
HaematiteStore::open(path).map_err(ServerError::from)?
} else {
HaematiteStore::create_with_shard_count(path, shard_count).map_err(ServerError::from)?
};
return Ok((store, None));
};
let boot = ClusterBootstrap {
node_id: cluster.node_id,
bind_address: cluster.bind_address,
members: cluster.members,
peers: cluster
.peers
.into_iter()
.map(|peer| (peer.name, peer.address))
.collect(),
timeout: HAEMATITE_CLUSTER_OP_TIMEOUT,
};
let (store, responder) =
HaematiteStore::open_or_create_distributed(data_dir, shard_count, boot)
.map_err(ServerError::from)?;
Ok((store, Some(responder)))
}
#[cfg(feature = "haematite-backend")]
const HAEMATITE_CLUSTER_OP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
#[cfg(not(feature = "haematite-backend"))]
fn connect_haematite_store_unavailable() -> Result<ConnectedStore, ServerError> {
Err(ServerError::Config {
message: "store.backend = haematite requires the aion-server `haematite-backend` feature"
.to_owned(),
})
}
#[cfg(not(feature = "liminal-transport"))]
#[derive(Clone, Debug)]
struct NullInterventionTransport;
#[cfg(not(feature = "liminal-transport"))]
#[async_trait::async_trait]
impl crate::worker::InterventionTransport for NullInterventionTransport {
async fn push(
&self,
_worker: &crate::worker::WorkerHandle,
_command: aion_core::InterventionCommand,
) -> Result<aion_core::InterventionOutcome, ServerError> {
Err(ServerError::worker_connection_lost(
"intervention",
"no intervention push transport is compiled in".to_owned(),
))
}
}
#[cfg(test)]
mod tests {
use std::{net::SocketAddr, time::Duration};
use aion_store::InMemoryStore;
use super::ServerState;
use crate::config::{
AuthConfig, AuthoringConfig, DeployConfig, DevConfig, ListenConfig, MetricsConfig,
NamespaceConfig, NamespaceMode, OpsConsoleAssetSource, OpsConsoleConfig, OutboxConfig,
RuntimeConfig, WebSocketConfig, WorkerConfig,
};
fn runtime_config() -> RuntimeConfig {
RuntimeConfig {
listen: ListenConfig {
grpc: SocketAddr::from(([127, 0, 0, 1], 50051)),
http: SocketAddr::from(([127, 0, 0, 1], 8080)),
},
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: WorkerConfig {
heartbeat_window: Duration::from_millis(30_000),
},
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: DevConfig::default(),
outbox: OutboxConfig::default(),
observability: crate::config::ObservabilityConfig::default(),
scheduler_threads: 1,
query_timeout: Some(Duration::from_millis(10_000)),
default_namespace: "default".to_owned(),
auto_create: crate::config::AutoCreate::Open,
max_in_flight_activities: crate::config::DEFAULT_MAX_IN_FLIGHT_ACTIVITIES,
drain_timeout: Duration::from_secs(30),
metrics: MetricsConfig { enabled: true },
owned_shards: Vec::new(),
cors_allowed_origins: Vec::new(),
}
}
#[tokio::test]
async fn builds_state_with_in_memory_store() -> Result<(), Box<dyn std::error::Error>> {
let state =
ServerState::build_with_store(InMemoryStore::default(), runtime_config()).await?;
std::hint::black_box(state.namespace_guard());
std::hint::black_box(state.worker_registry());
Ok(())
}
#[tokio::test]
async fn namespace_store_is_reachable_and_functional_after_default_boot()
-> Result<(), Box<dyn std::error::Error>> {
use aion_store::{MintOutcome, NamespaceOrigin};
let state =
ServerState::build_with_store(InMemoryStore::default(), runtime_config()).await?;
let store = state.namespace_store();
let outcome = store
.register_namespace("orders", NamespaceOrigin::WorkerMint)
.await?;
assert_eq!(
outcome,
MintOutcome::Created,
"the first reference to a namespace mints it"
);
let again = store
.register_namespace("orders", NamespaceOrigin::WorkerMint)
.await?;
assert_eq!(
again,
MintOutcome::AlreadyExisted,
"a second reference touches the existing record rather than re-creating it"
);
let fetched = store.get_namespace("orders").await?;
let record = fetched.ok_or("registered namespace must be retrievable via get_namespace")?;
assert_eq!(record.name, "orders");
assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
let listed = store.list_namespaces().await?;
assert!(
listed.iter().any(|record| record.name == "orders"),
"list_namespaces returns the minted namespace"
);
Ok(())
}
#[cfg(feature = "haematite-backend")]
#[tokio::test(flavor = "multi_thread")]
async fn connect_store_haematite_round_trips_through_event_store()
-> Result<(), Box<dyn std::error::Error>> {
use aion_core::{ContentType, EventEnvelope, PackageVersion, Payload, RunId, WorkflowId};
use aion_store::WriteToken;
use chrono::Utc;
use crate::config::{StoreBackend, StoreConfig};
let data_dir = tempfile::tempdir()?;
let connected = super::connect_store(StoreConfig {
backend: StoreBackend::Haematite,
url: None,
owned_shards: Vec::new(),
data_dir: Some(data_dir.path().to_string_lossy().into_owned()),
shard_count: 1,
cluster: None,
})
.await?;
let event_store = connected.event_store;
assert!(
connected.outbox_store.is_some(),
"the haematite backend shares its leaf store as the dispatcher's outbox store"
);
assert!(
connected.bootstrap_coordinator,
"a single-node haematite boot owns all shards and bootstraps the coordinator"
);
assert!(
connected.cluster_responder.is_none(),
"a single-node (no [cluster]) haematite boot has no distributed responder"
);
let workflow_id = WorkflowId::new_v4();
let event = aion_core::Event::WorkflowStarted {
envelope: EventEnvelope {
seq: 1,
recorded_at: Utc::now(),
workflow_id: workflow_id.clone(),
},
workflow_type: String::from("checkout"),
input: Payload::new(ContentType::Json, b"{}".to_vec()),
run_id: RunId::new_v4(),
parent_run_id: None,
package_version: PackageVersion::new("a".repeat(64)),
};
event_store
.append(
WriteToken::recorder(),
&workflow_id,
std::slice::from_ref(&event),
0,
)
.await?;
let history = event_store.read_history(&workflow_id).await?;
assert_eq!(
history.len(),
1,
"an event appended through the server's dyn EventStore reads back"
);
Ok(())
}
#[tokio::test]
async fn connect_store_memory_backend_exposes_no_outbox_store()
-> Result<(), Box<dyn std::error::Error>> {
use crate::config::{StoreBackend, StoreConfig};
let connected = super::connect_store(StoreConfig {
backend: StoreBackend::Memory,
url: None,
owned_shards: Vec::new(),
data_dir: None,
shard_count: 1,
cluster: None,
})
.await?;
assert!(
connected.outbox_store.is_none(),
"the in-memory backend exposes no outbox store"
);
Ok(())
}
#[cfg(feature = "libsql-backend")]
#[tokio::test]
async fn connect_store_shares_outbox_store_only_for_libsql()
-> Result<(), Box<dyn std::error::Error>> {
use crate::config::{StoreBackend, StoreConfig};
let path = std::env::temp_dir().join(format!(
"aion-connect-store-{}-{}.db",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|elapsed| elapsed.as_nanos())
.unwrap_or_default()
));
let connected = super::connect_store(StoreConfig {
backend: StoreBackend::LibSql,
url: Some(path.to_string_lossy().into_owned()),
owned_shards: Vec::new(),
data_dir: None,
shard_count: 1,
cluster: None,
})
.await?;
assert!(
connected.outbox_store.is_some(),
"the libSQL backend shares its leaf store as the dispatcher's outbox store"
);
Ok(())
}
#[tokio::test]
async fn state_build_fails_without_event_broadcast_capacity()
-> Result<(), Box<dyn std::error::Error>> {
let mut runtime = runtime_config();
runtime.websocket.event_broadcast_capacity = None;
let error = ServerState::build_with_store(InMemoryStore::default(), runtime)
.await
.err()
.ok_or("state build must fail when event streaming is unsized")?;
assert!(error.is_config(), "expected a config error, got {error}");
assert!(
error
.to_string()
.contains("websocket.event_broadcast_capacity"),
"error must name the missing key: {error}"
);
Ok(())
}
#[tokio::test]
async fn state_build_fails_without_query_timeout() -> Result<(), Box<dyn std::error::Error>> {
let mut runtime = runtime_config();
runtime.query_timeout = None;
let error = ServerState::build_with_store(InMemoryStore::default(), runtime)
.await
.err()
.ok_or("state build must fail when the query reply deadline is unset")?;
assert!(error.is_config(), "expected a config error, got {error}");
assert!(
error.to_string().contains("runtime.query_timeout_ms"),
"error must name the missing key: {error}"
);
assert!(
error.to_string().contains("AION_RUNTIME_QUERY_TIMEOUT_MS"),
"error must name the environment override: {error}"
);
Ok(())
}
#[tokio::test]
async fn state_build_fails_with_zero_query_timeout() -> Result<(), Box<dyn std::error::Error>> {
let mut runtime = runtime_config();
runtime.query_timeout = Some(Duration::ZERO);
let error = ServerState::build_with_store(InMemoryStore::default(), runtime)
.await
.err()
.ok_or("state build must fail when the query reply deadline is zero")?;
assert!(error.is_config(), "expected a config error, got {error}");
assert!(
error.to_string().contains("runtime.query_timeout_ms"),
"error must name the zero-valued key: {error}"
);
Ok(())
}
}