use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tracing::Instrument;
use crate::error::Error;
use crate::instance::Executor;
use crate::queue::DEFAULT_POLLING_INTERVAL;
use crate::registry::WorkflowKey;
use crate::sysdb;
use crate::sysdb::INTERNAL_QUEUE;
use crate::sysdb::types::{
Applications, NewWorkflow, QueueRecord, ResolvedLimits, Submission, WorkflowFilter,
WorkflowRecord, WorkflowStatus,
};
use crate::workflow::{MAX_RECOVERY_ATTEMPTS, spawn_execution, spawn_tracked};
const SUPERVISOR_INTERVAL: Duration = Duration::from_secs(1);
const MAX_POLLING_INTERVAL: Duration = Duration::from_secs(120);
const BACKOFF_FACTOR: f64 = 2.0;
const SCALEBACK_FACTOR: f64 = 0.9;
const JITTER: (f64, f64) = (0.95, 1.05);
#[derive(Default)]
pub(crate) struct Running(Mutex<HashMap<Key, i64>>);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum Key {
Queue(String),
Partition(String, String),
}
impl Running {
fn claim(self: &Arc<Self>, queue: &str, partition: Option<&str>) -> Slot {
let mut keys = vec![Key::Queue(queue.to_owned())];
if let Some(partition) = partition {
keys.push(Key::Partition(queue.to_owned(), partition.to_owned()));
}
let mut counts = self.0.lock().expect("the running tally is poisoned");
for key in &keys {
*counts.entry(key.clone()).or_default() += 1;
}
drop(counts);
Slot {
running: Arc::clone(self),
keys,
}
}
fn count(&self, queue: &str) -> i64 {
self.get(&Key::Queue(queue.to_owned()))
}
fn count_for_partition(&self, queue: &str, partition: &str) -> i64 {
self.get(&Key::Partition(queue.to_owned(), partition.to_owned()))
}
fn get(&self, key: &Key) -> i64 {
self.0
.lock()
.expect("the running tally is poisoned")
.get(key)
.copied()
.unwrap_or(0)
}
}
pub(crate) struct Slot {
running: Arc<Running>,
keys: Vec<Key>,
}
impl Drop for Slot {
fn drop(&mut self) {
let mut counts = match self.running.0.lock() {
Ok(counts) => counts,
Err(_) => return,
};
for key in &self.keys {
if let Some(count) = counts.get_mut(key) {
*count -= 1;
if *count <= 0 {
counts.remove(key);
}
}
}
}
}
type Queues = Arc<Mutex<HashMap<String, QueueRecord>>>;
pub(crate) fn spawn(executor: Arc<Executor>) {
spawn_tracked(
&executor,
{
let executor = Arc::clone(&executor);
supervise(executor)
}
.instrument(tracing::info_span!("queues")),
);
}
async fn supervise(executor: Arc<Executor>) {
let queues: Queues = Arc::default();
let running: Arc<Running> = Arc::default();
let mut workers: HashMap<String, tokio::task::JoinHandle<()>> = HashMap::new();
let mut warned_internal = false;
loop {
match executor.sysdb().transition_delayed_workflows().await {
Ok(0) => {}
Ok(moved) => tracing::debug!(workflows = moved, "delayed workflows are now enqueued"),
Err(error) => {
tracing::warn!(error = %error, "could not transition delayed workflows");
}
}
let names = refresh_queue_set(&executor, &queues, &mut warned_internal).await;
workers.retain(|_, handle| !handle.is_finished());
for name in names {
if workers.contains_key(&name) {
continue;
}
tracing::debug!(queue = name, "starting a worker for the queue");
let handle = spawn_tracked(
&executor,
poll_queue(
Arc::clone(&executor),
name.clone(),
Arc::clone(&queues),
Arc::clone(&running),
)
.instrument(tracing::info_span!("queue", queue = %name)),
);
workers.insert(name, handle);
}
tokio::time::sleep(SUPERVISOR_INTERVAL).await;
}
}
async fn refresh_queue_set(
executor: &Arc<Executor>,
queues: &Queues,
warned_internal: &mut bool,
) -> Vec<String> {
let listed = match executor.sysdb().list_queues(&Applications::Unset).await {
Ok(listed) => listed,
Err(error) => {
tracing::warn!(error = %error, "could not list queues; keeping the current set");
return queues
.lock()
.expect("the queue set is poisoned")
.keys()
.cloned()
.collect();
}
};
let mut current: HashMap<String, QueueRecord> =
HashMap::from([(INTERNAL_QUEUE.to_owned(), internal_queue())]);
let listened = executor.listen_queues();
for queue in listed {
if queue.name == INTERNAL_QUEUE {
if !*warned_internal {
*warned_internal = true;
tracing::warn!(
queue = INTERNAL_QUEUE,
"the queues table holds a row for the engine's internal queue; its stored \
limits are ignored. Delete the row: it can only throttle `resume` and `fork`"
);
}
continue;
}
if let Some(listened) = listened
&& !listened.contains(&queue.name)
{
continue;
}
current.insert(queue.name.clone(), queue);
}
let names = current.keys().cloned().collect();
*queues.lock().expect("the queue set is poisoned") = current;
names
}
fn internal_queue() -> QueueRecord {
QueueRecord {
name: INTERNAL_QUEUE.to_owned(),
concurrency: None,
worker_concurrency: None,
rate_limit: None,
priority_enabled: false,
partition_queue: false,
partition_concurrency: None,
partition_worker_concurrency: None,
partition_rate_limit: None,
polling_interval: DEFAULT_POLLING_INTERVAL,
application_name: None,
}
}
async fn poll_queue(executor: Arc<Executor>, name: String, queues: Queues, running: Arc<Running>) {
let mut interval = DEFAULT_POLLING_INTERVAL;
loop {
let Some(queue) = queues
.lock()
.expect("the queue set is poisoned")
.get(&name)
.cloned()
else {
tracing::info!("the queue is no longer registered; stopping its worker");
return;
};
let floor = queue.polling_interval;
let ceiling = floor.max(MAX_POLLING_INTERVAL);
interval = interval.clamp(floor, ceiling);
let contended = poll_once(&executor, &queue, &running).await;
interval = if contended {
interval.mul_f64(BACKOFF_FACTOR).min(ceiling)
} else {
interval.mul_f64(SCALEBACK_FACTOR).max(floor)
};
tokio::time::sleep(jitter(interval)).await;
}
}
fn worker_budget(limits: &ResolvedLimits, running: i64) -> Option<i64> {
if limits.partition_worker_concurrency.is_some_and(|w| w <= 0) {
return Some(0);
}
limits
.worker_concurrency
.map(|worker| (i64::from(worker) - running).max(0))
}
async fn poll_once(executor: &Arc<Executor>, queue: &QueueRecord, running: &Arc<Running>) -> bool {
let limits = queue.resolved_limits();
if !limits.is_partitioned() {
return match executor
.sysdb()
.start_queued_workflows(
queue,
executor.executor_id(),
executor.app_version(),
None,
running.count(&queue.name),
0,
)
.await
{
Ok(claimed) => {
dispatch_claimed(executor, queue, None, running, claimed).await;
false
}
Err(error) => report_dequeue_error(&error),
};
}
let already_running = running.count(&queue.name);
let budget = worker_budget(&limits, already_running);
if budget == Some(0) {
return false;
}
if limits.partition_concurrency == Some(1)
&& limits.concurrency.is_none()
&& limits.rate_limit.is_none()
&& limits.partition_rate_limit.is_none()
{
return match executor
.sysdb()
.start_queued_partitioned_workflows(
queue,
executor.executor_id(),
executor.app_version(),
budget,
)
.await
{
Ok(claimed) => {
dispatch_claimed(executor, queue, None, running, claimed).await;
false
}
Err(error) => report_dequeue_error(&error),
};
}
let partitions = match executor.sysdb().get_queue_partitions(&queue.name).await {
Ok(partitions) => shuffled(partitions),
Err(error) => return report_dequeue_error(&error),
};
let mut claimed_here = 0i64;
for partition in partitions {
if worker_budget(&limits, already_running + claimed_here) == Some(0) {
break;
}
let claimed = match executor
.sysdb()
.start_queued_workflows(
queue,
executor.executor_id(),
executor.app_version(),
Some(&partition),
already_running + claimed_here,
running.count_for_partition(&queue.name, &partition),
)
.await
{
Ok(claimed) => claimed,
Err(error) if is_contention(&error) => {
tracing::debug!(
partition,
"a peer is mid-dequeue on this partition; skipping it"
);
continue;
}
Err(error) => {
tracing::warn!(error = %error, partition, "could not dequeue from the partition");
continue;
}
};
claimed_here += i64::try_from(claimed.len()).unwrap_or(i64::MAX);
dispatch_claimed(executor, queue, Some(&partition), running, claimed).await;
}
false
}
fn report_dequeue_error(error: &sysdb::Error) -> bool {
if is_contention(error) {
tracing::debug!(error = %error, "a peer is mid-dequeue; backing off");
return true;
}
tracing::warn!(error = %error, "could not dequeue from the queue");
false
}
fn shuffled(mut keys: Vec<String>) -> Vec<String> {
let mut state = (uuid::Uuid::new_v4().as_u128() as u64) | 1;
for i in (1..keys.len()).rev() {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
keys.swap(i, (state % (i as u64 + 1)) as usize);
}
keys
}
async fn dispatch_claimed(
executor: &Arc<Executor>,
queue: &QueueRecord,
partition: Option<&str>,
running: &Arc<Running>,
claimed: Vec<String>,
) {
if claimed.is_empty() {
return;
}
let ids: Vec<&str> = claimed.iter().map(String::as_str).collect();
let rows = match executor
.sysdb()
.list_workflows(
&WorkflowFilter {
workflow_ids: ids,
load_output: false,
..WorkflowFilter::default()
},
None,
)
.await
{
Ok(rows) => rows,
Err(error) => {
tracing::warn!(
error = %error,
claimed = claimed.len(),
"could not read the claimed workflows; they stay PENDING for recovery"
);
return;
}
};
if rows.len() != claimed.len() {
tracing::warn!(
claimed = claimed.len(),
found = rows.len(),
"some claimed workflows have no row"
);
}
tracing::debug!(workflows = rows.len(), "dequeued workflows");
let mut by_id: HashMap<String, _> = rows
.into_iter()
.map(|row| (row.workflow_id.clone(), row))
.collect();
for id in &claimed {
let Some(row) = by_id.remove(id) else {
continue;
};
let workflow_id = row.workflow_id.clone();
let partition = partition.or(row.queue_partition_key.as_deref());
let slot = running.claim(&queue.name, partition);
if let Err(error) = dispatch(executor, row, Submission::Dequeue, Some(slot)).await {
tracing::warn!(
workflow_id,
error = %error,
"could not start the dequeued workflow; it stays PENDING for recovery"
);
}
}
}
async fn dispatch(
executor: &Arc<Executor>,
row: WorkflowRecord,
submission: Submission,
slot: Option<Slot>,
) -> crate::Result<()> {
let workflow_id = row.workflow_id;
let Some(name) = row.name else {
tracing::warn!(workflow_id, "the row names no workflow; skipped");
return Ok(());
};
let key = WorkflowKey::from_row(name, row.class_name.as_deref(), row.config_name.as_deref());
if !executor.workflows().contains_key(&key) {
tracing::warn!(
workflow_id,
workflow = %key,
"no workflow is registered under the row's name; it stays where it is"
);
return Ok(());
}
let initialized = match executor
.sysdb()
.init_workflow(
&NewWorkflow {
name: Some(&key.name),
class_name: key.class_name.as_deref(),
config_name: key.config_name.as_deref(),
input: row.input.as_deref(),
serialization: row.serialization.as_deref(),
queue_name: row.queue_name.as_deref(),
executor_id: Some(executor.executor_id()),
application_name: Some(executor.app_name()),
application_version: Some(executor.app_version()),
..NewWorkflow::new(&workflow_id)
},
Some(MAX_RECOVERY_ATTEMPTS),
submission,
None,
)
.await
{
Ok(initialized) => initialized,
Err(error @ sysdb::Error::MaxRecoveryAttemptsExceeded { .. }) => {
tracing::warn!(workflow_id, error = %error, "the workflow is parked");
return Ok(());
}
Err(error) => return Err(Error::SystemDatabase(error)),
};
if initialized.status != WorkflowStatus::Pending {
tracing::warn!(
workflow_id,
status = %initialized.status,
"the dequeued workflow is no longer pending; it is not started"
);
return Ok(());
}
tracing::debug!(
workflow_id,
workflow = %key,
attempt = initialized.recovery_attempts,
submission = ?submission,
"running the workflow"
);
spawn_execution(
executor,
key,
workflow_id,
row.input,
initialized.deadline,
slot,
);
Ok(())
}
fn is_contention(error: &sysdb::Error) -> bool {
matches!(
error,
sysdb::Error::Backend(backend) if backend.sqlstate.as_deref() == Some("55P03")
)
}
fn jitter(interval: Duration) -> Duration {
let (low, high) = JITTER;
let bits = uuid::Uuid::new_v4().as_u128() as u32;
let factor = low + (high - low) * (f64::from(bits) / f64::from(u32::MAX).next_up());
interval.mul_f64(factor)
}