use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, Ordering};
use aion_core::{RunId, WorkflowId};
use dashmap::DashMap;
use dashmap::mapref::entry::Entry;
use tokio::task::JoinHandle;
use crate::EngineError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct CompletionRetryKey {
pub(crate) workflow_id: WorkflowId,
pub(crate) run_id: RunId,
pub(crate) monitor_pid: super::Pid,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ArmOutcome {
Armed,
AlreadyArmed,
EpochClosed,
}
pub(crate) struct EngineTaskRuntime {
runtime: Mutex<Option<tokio::runtime::Runtime>>,
watches: DashMap<(u64, WorkflowId), JoinHandle<()>>,
spawn_retries: DashMap<WorkflowId, JoinHandle<()>>,
completion_retries: DashMap<CompletionRetryKey, JoinHandle<()>>,
shutting_down: AtomicBool,
}
impl EngineTaskRuntime {
pub(crate) fn new() -> Result<Self, EngineError> {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.thread_name("aion-engine-tasks")
.enable_all()
.build()
.map_err(|error| EngineError::Runtime {
reason: format!("failed to start the engine-task runtime: {error}"),
})?;
Ok(Self {
runtime: Mutex::new(Some(runtime)),
watches: DashMap::new(),
spawn_retries: DashMap::new(),
completion_retries: DashMap::new(),
shutting_down: AtomicBool::new(false),
})
}
pub(crate) fn arm_watch<F>(&self, parent_pid: u64, child_id: WorkflowId, task: F) -> ArmOutcome
where
F: Future<Output = ()> + Send + 'static,
{
Self::arm(
&self.shutting_down,
&self.runtime,
&self.watches,
(parent_pid, child_id),
task,
)
}
pub(crate) fn arm_spawn_retry<F>(&self, child_id: WorkflowId, task: F) -> ArmOutcome
where
F: Future<Output = ()> + Send + 'static,
{
Self::arm(
&self.shutting_down,
&self.runtime,
&self.spawn_retries,
child_id,
task,
)
}
pub(crate) fn arm_completion_retry<F>(&self, lease: CompletionRetryKey, task: F) -> ArmOutcome
where
F: Future<Output = ()> + Send + 'static,
{
Self::arm(
&self.shutting_down,
&self.runtime,
&self.completion_retries,
lease,
task,
)
}
fn arm<K, F>(
shutting_down: &AtomicBool,
runtime: &Mutex<Option<tokio::runtime::Runtime>>,
registry: &DashMap<K, JoinHandle<()>>,
key: K,
task: F,
) -> ArmOutcome
where
K: std::hash::Hash + Eq + Clone,
F: Future<Output = ()> + Send + 'static,
{
if shutting_down.load(Ordering::Acquire) {
return ArmOutcome::EpochClosed;
}
let handle = {
let guard = match runtime.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
let Some(owned) = guard.as_ref() else {
return ArmOutcome::EpochClosed;
};
owned.handle().clone()
};
let undo_key = key.clone();
let spawned_id;
match registry.entry(key) {
Entry::Occupied(slot) => {
if slot.get().is_finished() {
let spawned = handle.spawn(task);
spawned_id = spawned.id();
let (key, _finished) = slot.replace_entry(spawned);
let _ = key;
} else {
return ArmOutcome::AlreadyArmed;
}
}
Entry::Vacant(slot) => {
let spawned = handle.spawn(task);
spawned_id = spawned.id();
slot.insert(spawned);
}
}
if shutting_down.load(Ordering::Acquire) {
if let Some((_, handle)) = registry.remove_if(&undo_key, |_, h| h.id() == spawned_id) {
handle.abort();
}
return ArmOutcome::EpochClosed;
}
ArmOutcome::Armed
}
pub(crate) fn remove_watch(&self, parent_pid: u64, child_id: &WorkflowId) {
self.watches.remove(&(parent_pid, child_id.clone()));
}
pub(crate) fn remove_spawn_retry(&self, child_id: &WorkflowId) {
self.spawn_retries.remove(child_id);
}
pub(crate) fn remove_completion_retry(
&self,
lease: &CompletionRetryKey,
task: tokio::task::Id,
) {
self.completion_retries
.remove_if(lease, |_, handle| handle.id() == task);
}
pub(crate) fn abort_watch(&self, parent_pid: u64, child_id: &WorkflowId) {
if let Some((_, handle)) = self.watches.remove(&(parent_pid, child_id.clone())) {
handle.abort();
}
}
pub(crate) fn abort_watches_for_parent(&self, parent_pid: u64) {
self.watches.retain(|(pid, _), handle| {
if *pid == parent_pid {
handle.abort();
false
} else {
true
}
});
}
#[cfg(test)]
pub(crate) fn armed_watch_count(&self) -> usize {
self.watches.len()
}
#[cfg(test)]
pub(crate) fn armed_spawn_retry_count(&self) -> usize {
self.spawn_retries.len()
}
#[cfg(test)]
pub(crate) fn armed_completion_retry_count(&self) -> usize {
self.completion_retries.len()
}
#[cfg(test)]
pub(crate) fn owns_runtime(&self) -> bool {
match self.runtime.lock() {
Ok(guard) => guard.is_some(),
Err(poisoned) => poisoned.into_inner().is_some(),
}
}
pub(crate) fn shutdown(&self) {
self.gate_and_abort();
let runtime = {
let mut guard = match self.runtime.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
guard.take()
};
let Some(runtime) = runtime else {
return;
};
shutdown_runtime_and_join(runtime);
}
pub(crate) fn is_epoch_open(&self) -> bool {
!self.shutting_down.load(Ordering::Acquire)
}
fn gate_and_abort(&self) {
self.shutting_down.store(true, Ordering::Release);
self.watches.retain(|_, handle| {
handle.abort();
false
});
self.spawn_retries.retain(|_, handle| {
handle.abort();
false
});
self.completion_retries.retain(|_, handle| {
handle.abort();
false
});
}
pub(crate) fn begin_close(&self) {
self.gate_and_abort();
}
}
pub(crate) async fn sleep_backoff(current: &mut std::time::Duration, ceiling: std::time::Duration) {
tokio::time::sleep(*current).await;
let doubled = current.saturating_mul(2);
*current = if doubled > ceiling { ceiling } else { doubled };
}
fn shutdown_runtime_and_join(runtime: tokio::runtime::Runtime) {
match std::thread::Builder::new()
.name("aion-engine-tasks-shutdown".to_owned())
.spawn(move || drop(runtime))
{
Ok(joiner) => {
if joiner.join().is_err() {
tracing::error!("engine-task runtime shutdown thread panicked");
}
}
Err(error) => {
tracing::error!(
error = %error,
"could not spawn the engine-task runtime shutdown thread; the \
runtime was dropped on the calling thread instead"
);
}
}
}
impl Drop for EngineTaskRuntime {
fn drop(&mut self) {
self.gate_and_abort();
let runtime = {
let mut guard = match self.runtime.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
guard.take()
};
if let Some(runtime) = runtime {
runtime.shutdown_background();
}
}
}
#[cfg(test)]
mod tests;