[model-gateway] improve workflow engine code quality (#16977)
This commit is contained in:
@@ -2,6 +2,7 @@
|
||||
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
fmt,
|
||||
sync::Arc,
|
||||
time::Duration,
|
||||
};
|
||||
@@ -11,6 +12,18 @@ use super::{
|
||||
types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId},
|
||||
};
|
||||
|
||||
/// Errors that can occur during workflow validation
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum ValidationError {
|
||||
/// A step depends on another step that doesn't exist
|
||||
#[error("Step '{step}' depends on non-existent step '{dependency}'")]
|
||||
MissingDependency { step: StepId, dependency: StepId },
|
||||
|
||||
/// A cycle was detected in the workflow DAG
|
||||
#[error("Cycle detected involving step '{0}'")]
|
||||
CycleDetected(StepId),
|
||||
}
|
||||
|
||||
/// Definition of a single step within a workflow
|
||||
pub struct StepDefinition<D: WorkflowData> {
|
||||
pub id: StepId,
|
||||
@@ -22,6 +35,19 @@ pub struct StepDefinition<D: WorkflowData> {
|
||||
pub depends_on: Vec<StepId>,
|
||||
}
|
||||
|
||||
impl<D: WorkflowData> fmt::Debug for StepDefinition<D> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("StepDefinition")
|
||||
.field("id", &self.id)
|
||||
.field("name", &self.name)
|
||||
.field("retry_policy", &self.retry_policy)
|
||||
.field("timeout", &self.timeout)
|
||||
.field("on_failure", &self.on_failure)
|
||||
.field("depends_on", &self.depends_on)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl<D: WorkflowData> StepDefinition<D> {
|
||||
pub fn new(
|
||||
id: impl Into<String>,
|
||||
@@ -76,6 +102,18 @@ pub struct WorkflowDefinition<D: WorkflowData> {
|
||||
initial_step_indices: Vec<usize>,
|
||||
}
|
||||
|
||||
impl<D: WorkflowData> fmt::Debug for WorkflowDefinition<D> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("WorkflowDefinition")
|
||||
.field("id", &self.id)
|
||||
.field("name", &self.name)
|
||||
.field("steps", &self.steps)
|
||||
.field("default_retry_policy", &self.default_retry_policy)
|
||||
.field("default_timeout", &self.default_timeout)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl<D: WorkflowData> WorkflowDefinition<D> {
|
||||
pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
@@ -122,7 +160,8 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
|
||||
/// - There's a cycle in the dependencies
|
||||
///
|
||||
/// On success, pre-computes reverse dependencies for O(1) dependent lookup.
|
||||
pub fn validate(&mut self) -> Result<(), String> {
|
||||
#[must_use = "validation result should be checked"]
|
||||
pub fn validate(&mut self) -> Result<(), ValidationError> {
|
||||
// Build HashMap for O(1) lookup instead of O(n) linear search
|
||||
let steps_map: HashMap<&StepId, &StepDefinition<D>> =
|
||||
self.steps.iter().map(|s| (&s.id, s)).collect();
|
||||
@@ -131,10 +170,10 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
|
||||
for step in &self.steps {
|
||||
for dep in &step.depends_on {
|
||||
if !steps_map.contains_key(dep) {
|
||||
return Err(format!(
|
||||
"Step '{}' depends on non-existent step '{}'",
|
||||
step.id, dep
|
||||
));
|
||||
return Err(ValidationError::MissingDependency {
|
||||
step: step.id.clone(),
|
||||
dependency: dep.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -147,7 +186,7 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
|
||||
if !visited.contains(&step.id)
|
||||
&& Self::has_cycle(&step.id, &steps_map, &mut visited, &mut rec_stack)
|
||||
{
|
||||
return Err(format!("Cycle detected involving step '{}'", step.id));
|
||||
return Err(ValidationError::CycleDetected(step.id.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -100,6 +100,23 @@ impl Backoff for LinearBackoff {
|
||||
}
|
||||
}
|
||||
|
||||
/// Enum-based backoff implementation to avoid heap allocation
|
||||
enum BackoffImpl {
|
||||
Fixed(FixedBackoff),
|
||||
Exponential(backoff::ExponentialBackoff),
|
||||
Linear(LinearBackoff),
|
||||
}
|
||||
|
||||
impl BackoffImpl {
|
||||
fn next_backoff(&mut self) -> Option<Duration> {
|
||||
match self {
|
||||
BackoffImpl::Fixed(b) => b.next_backoff(),
|
||||
BackoffImpl::Exponential(b) => b.next_backoff(),
|
||||
BackoffImpl::Linear(b) => b.next_backoff(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Main workflow execution engine
|
||||
///
|
||||
/// # Type Parameters
|
||||
@@ -299,7 +316,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
|
||||
}
|
||||
|
||||
/// Register a workflow definition
|
||||
pub fn register_workflow(&self, mut definition: WorkflowDefinition<D>) -> Result<(), String> {
|
||||
#[must_use = "registration result should be checked"]
|
||||
pub fn register_workflow(
|
||||
&self,
|
||||
mut definition: WorkflowDefinition<D>,
|
||||
) -> Result<(), super::definition::ValidationError> {
|
||||
// Validate DAG and build dependency graph once at registration
|
||||
definition.validate()?;
|
||||
|
||||
@@ -321,6 +342,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
|
||||
/// Start a new workflow instance
|
||||
///
|
||||
/// Returns `Err(WorkflowError::ShuttingDown)` if the engine is shutting down.
|
||||
#[must_use = "workflow instance ID should be stored or awaited"]
|
||||
pub async fn start_workflow(
|
||||
&self,
|
||||
definition_id: WorkflowId,
|
||||
@@ -746,19 +768,19 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
|
||||
}
|
||||
}
|
||||
|
||||
fn create_backoff(strategy: &BackoffStrategy) -> Box<dyn Backoff + Send> {
|
||||
fn create_backoff(strategy: &BackoffStrategy) -> BackoffImpl {
|
||||
match strategy {
|
||||
BackoffStrategy::Fixed(duration) => Box::new(FixedBackoff(*duration)),
|
||||
BackoffStrategy::Fixed(duration) => BackoffImpl::Fixed(FixedBackoff(*duration)),
|
||||
BackoffStrategy::Exponential { base, max } => {
|
||||
let backoff = ExponentialBackoffBuilder::new()
|
||||
.with_initial_interval(*base)
|
||||
.with_max_interval(*max)
|
||||
.with_max_elapsed_time(None)
|
||||
.build();
|
||||
Box::new(backoff)
|
||||
BackoffImpl::Exponential(backoff)
|
||||
}
|
||||
BackoffStrategy::Linear { increment, max } => {
|
||||
Box::new(LinearBackoff::new(*increment, *max))
|
||||
BackoffImpl::Linear(LinearBackoff::new(*increment, *max))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ mod executor;
|
||||
mod state;
|
||||
pub mod types;
|
||||
|
||||
pub use definition::{StepDefinition, WorkflowDefinition};
|
||||
pub use definition::{StepDefinition, ValidationError, WorkflowDefinition};
|
||||
pub use engine::WorkflowEngine;
|
||||
pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent};
|
||||
pub use executor::{FunctionStep, StepExecutor};
|
||||
|
||||
@@ -649,7 +649,10 @@ fn test_dag_validation_cycle_detection() {
|
||||
|
||||
let result = workflow.validate();
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("Cycle detected"));
|
||||
assert!(matches!(
|
||||
result.unwrap_err(),
|
||||
ValidationError::CycleDetected(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -668,7 +671,10 @@ fn test_dag_validation_missing_dependency() {
|
||||
|
||||
let result = workflow.validate();
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("non-existent step"));
|
||||
assert!(matches!(
|
||||
result.unwrap_err(),
|
||||
ValidationError::MissingDependency { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user