[model-gateway] improve workflow engine code quality (#16977)

This commit is contained in:
Simo Lin
2026-01-12 12:14:46 -08:00
committed by GitHub
parent ed729d22b3
commit 6e158e55b4
4 changed files with 81 additions and 14 deletions

View File

@@ -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()));
}
}

View File

@@ -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))
}
}
}

View File

@@ -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};

View File

@@ -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]