use std::collections::HashSet;
use aion_package::ContentHash;
use crate::EngineError;
use crate::loader::DeployedWorkerContract;
use super::api::Engine;
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub enum AdmissionReason {
RouteActive,
LiveWorkflow,
StartInFlight,
}
impl AdmissionReason {
#[must_use]
pub const fn explanation(self) -> &'static str {
match self {
Self::RouteActive => "it currently routes new starts",
Self::LiveWorkflow => "a live workflow run is pinned to it",
Self::StartInFlight => "a workflow start is in flight on it",
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RequiredContract {
pub contract: DeployedWorkerContract,
pub reason: AdmissionReason,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UnreachableContract {
pub package_version: ContentHash,
pub workflow_types: Vec<String>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct QueueAdmission {
pub required: Vec<RequiredContract>,
pub unreachable: Vec<UnreachableContract>,
}
impl QueueAdmission {
#[must_use]
pub fn reason_census(&self) -> ReasonCensus {
let mut census = ReasonCensus::default();
for required in &self.required {
match required.reason {
AdmissionReason::RouteActive => census.route_active += 1,
AdmissionReason::LiveWorkflow => census.live_workflow += 1,
AdmissionReason::StartInFlight => census.start_in_flight += 1,
}
}
census
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ReasonCensus {
pub route_active: usize,
pub live_workflow: usize,
pub start_in_flight: usize,
}
impl Engine {
pub fn worker_contracts_for_admission(
&self,
task_queue: &str,
) -> Result<QueueAdmission, EngineError> {
let catalog = self.workflow_catalog();
let contracts = catalog.worker_contracts_for_queue(task_queue)?;
if contracts.is_empty() {
return Ok(QueueAdmission::default());
}
let mut live = HashSet::new();
for handle in self.registry().list()? {
if !handle.cached_status().is_terminal() {
live.insert(handle.loaded_version().clone());
}
}
let starting = catalog
.pinned_start_versions()?
.into_iter()
.map(|(_, version)| version)
.collect::<HashSet<_>>();
let mut admission = QueueAdmission::default();
for contract in contracts {
let reason = if contract.route_active {
Some(AdmissionReason::RouteActive)
} else if live.contains(&contract.package_version) {
Some(AdmissionReason::LiveWorkflow)
} else if starting.contains(&contract.package_version) {
Some(AdmissionReason::StartInFlight)
} else {
None
};
match reason {
Some(reason) => admission
.required
.push(RequiredContract { contract, reason }),
None => admission.unreachable.push(UnreachableContract {
package_version: contract.package_version,
workflow_types: contract.workflow_types,
}),
}
}
Ok(admission)
}
}
#[cfg(test)]
#[path = "admission_tests.rs"]
mod admission_tests;