minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] make StateStore trait async for external persistence (#16979)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 12:55:59 -08:00
committed by GitHub
parent 6e158e55b4
commit 6620548fd8
4 changed files with 167 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/workflow/engine.rs
+103 -68
View File
@@ -229,7 +229,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
/// to ensure all workflows are stopped. Note that this cancels workflows /// to ensure all workflows are stopped. Note that this cancels workflows
/// at the state level; running steps may still complete. /// at the state level; running steps may still complete.
pub async fn force_cancel_all(&self) -> usize { pub async fn force_cancel_all(&self) -> usize {
let active_states = match self.state_store.list_active() { let active_states = match self.state_store.list_active().await {
Ok(states) => states, Ok(states) => states,
Err(e) => { Err(e) => {
tracing::error!(error = ?e, "Failed to list active workflows for force cancel"); tracing::error!(error = ?e, "Failed to list active workflows for force cancel");
@@ -304,7 +304,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
loop { loop {
tokio::select! { tokio::select! {
_ = ticker.tick() => { _ = ticker.tick() => {
state_store.cleanup_old_workflows(ttl); state_store.cleanup_old_workflows(ttl).await;
} }
_ = shutdown_rx.changed() => { _ = shutdown_rx.changed() => {
tracing::info!("Cleanup task stopping due to shutdown"); tracing::info!("Cleanup task stopping due to shutdown");
@@ -373,7 +373,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.insert(step.id.clone(), StepState::default()); .insert(step.id.clone(), StepState::default());
} }
self.state_store.save(state)?; self.state_store.save(state).await?;
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowStarted { .publish(WorkflowEvent::WorkflowStarted {
@@ -421,7 +421,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.collect(); .collect();
loop { loop {
if self.state_store.is_cancelled(instance_id)? { if self.state_store.is_cancelled(instance_id).await? {
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowCancelled { instance_id }) .publish(WorkflowEvent::WorkflowCancelled { instance_id })
.await; .await;
@@ -460,9 +460,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
"Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string() "Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
}; };
self.state_store.update(instance_id, |s| { self.state_store
s.status = WorkflowStatus::Failed; .update(instance_id, |s| {
})?; s.status = WorkflowStatus::Failed;
})
.await?;
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowFailed { .publish(WorkflowEvent::WorkflowFailed {
instance_id, instance_id,
@@ -498,37 +500,31 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
Err(_) => StepResult::Failure, Err(_) => StepResult::Failure,
}; };
{ // Track whether we need to update state to Skipped after releasing lock
let needs_skip_update = {
let mut t = tracker.write(); let mut t = tracker.write();
t.running.remove(&step_id); t.running.remove(&step_id);
match result { let needs_update = match result {
Ok(StepResult::Success) => { Ok(StepResult::Success) => {
t.completed.insert(step_id.clone()); t.completed.insert(step_id.clone());
false
} }
Ok(StepResult::Skip) => { Ok(StepResult::Skip) => {
t.skipped.insert(step_id.clone()); t.skipped.insert(step_id.clone());
false
} }
Ok(StepResult::Failure) | Err(_) => match step.on_failure { Ok(StepResult::Failure) | Err(_) => match step.on_failure {
FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => { FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => {
t.failed.insert(step_id.clone()); t.failed.insert(step_id.clone());
false
} }
FailureAction::ContinueNextStep => { FailureAction::ContinueNextStep => {
if let Err(e) = engine.state_store.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step_id) {
step_state.status = StepStatus::Skipped;
}
}) {
tracing::warn!(
step_id = %step_id,
error = ?e,
"Failed to update step state to Skipped"
);
}
t.skipped.insert(step_id.clone()); t.skipped.insert(step_id.clone());
true // Need to update state store after releasing lock
} }
}, },
} };
if let Err(e) = tx.try_send((step_id.clone(), signal)) { if let Err(e) = tx.try_send((step_id.clone(), signal)) {
use mpsc::error::TrySendError; use mpsc::error::TrySendError;
@@ -547,6 +543,27 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
} }
} }
} }
needs_update
};
// Perform async state update after releasing the tracker lock
if needs_skip_update {
if let Err(e) = engine
.state_store
.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step_id) {
step_state.status = StepStatus::Skipped;
}
})
.await
{
tracing::warn!(
step_id = %step_id,
error = ?e,
"Failed to update step state to Skipped"
);
}
} }
}); });
} }
@@ -577,9 +594,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}; };
if let Some(ref step) = failed_step { if let Some(ref step) = failed_step {
self.state_store.update(instance_id, |s| { self.state_store
s.status = WorkflowStatus::Failed; .update(instance_id, |s| {
})?; s.status = WorkflowStatus::Failed;
})
.await?;
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowFailed { .publish(WorkflowEvent::WorkflowFailed {
instance_id, instance_id,
@@ -588,9 +607,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}) })
.await; .await;
} else { } else {
self.state_store.update(instance_id, |s| { self.state_store
s.status = WorkflowStatus::Completed; .update(instance_id, |s| {
})?; s.status = WorkflowStatus::Completed;
})
.await?;
let duration = start_time.elapsed(); let duration = start_time.elapsed();
self.event_bus self.event_bus
@@ -624,23 +645,25 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let mut backoff = Self::create_backoff(&retry_policy.backoff); let mut backoff = Self::create_backoff(&retry_policy.backoff);
loop { loop {
if self.state_store.is_cancelled(instance_id)? { if self.state_store.is_cancelled(instance_id).await? {
return Err(WorkflowError::Cancelled(instance_id)); return Err(WorkflowError::Cancelled(instance_id));
} }
// Update step state // Update step state
self.state_store.update(instance_id, |s| { self.state_store
s.current_step = Some(step.id.clone()); .update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step.id) { s.current_step = Some(step.id.clone());
step_state.status = if attempt == 1 { if let Some(step_state) = s.step_states.get_mut(&step.id) {
StepStatus::Running step_state.status = if attempt == 1 {
} else { StepStatus::Running
StepStatus::Retrying } else {
}; StepStatus::Retrying
step_state.attempt = attempt; };
step_state.started_at = Some(Utc::now()); step_state.attempt = attempt;
} step_state.started_at = Some(Utc::now());
})?; }
})
.await?;
// Emit step started event // Emit step started event
self.event_bus self.event_bus
@@ -651,7 +674,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}) })
.await; .await;
let mut context = self.state_store.get_context(instance_id)?; let mut context = self.state_store.get_context(instance_id).await?;
// Execute step with timeout // Execute step with timeout
let step_start = std::time::Instant::now(); let step_start = std::time::Instant::now();
@@ -659,19 +682,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let step_duration = step_start.elapsed(); let step_duration = step_start.elapsed();
self.state_store.update(instance_id, |s| { self.state_store
s.context = context.clone(); .update(instance_id, |s| {
})?; s.context = context.clone();
})
.await?;
match result { match result {
Ok(Ok(StepResult::Success)) => { Ok(Ok(StepResult::Success)) => {
// Step succeeded // Step succeeded
self.state_store.update(instance_id, |s| { self.state_store
if let Some(step_state) = s.step_states.get_mut(&step.id) { .update(instance_id, |s| {
step_state.status = StepStatus::Succeeded; if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.completed_at = Some(Utc::now()); step_state.status = StepStatus::Succeeded;
} step_state.completed_at = Some(Utc::now());
})?; }
})
.await?;
self.event_bus self.event_bus
.publish(WorkflowEvent::StepSucceeded { .publish(WorkflowEvent::StepSucceeded {
@@ -708,19 +735,21 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let will_retry = should_retry && attempt < max_attempts; let will_retry = should_retry && attempt < max_attempts;
// Update step state // Update step state
self.state_store.update(instance_id, |s| { self.state_store
if let Some(step_state) = s.step_states.get_mut(&step.id) { .update(instance_id, |s| {
step_state.status = if will_retry { if let Some(step_state) = s.step_states.get_mut(&step.id) {
StepStatus::Retrying step_state.status = if will_retry {
} else { StepStatus::Retrying
StepStatus::Failed } else {
}; StepStatus::Failed
step_state.last_error = Some(error_msg.clone()); };
if !will_retry { step_state.last_error = Some(error_msg.clone());
step_state.completed_at = Some(Utc::now()); if !will_retry {
step_state.completed_at = Some(Utc::now());
}
} }
} })
})?; .await?;
// Emit step failed event // Emit step failed event
self.event_bus self.event_bus
@@ -787,9 +816,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
/// Cancel a running workflow /// Cancel a running workflow
pub async fn cancel_workflow(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> { pub async fn cancel_workflow(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.state_store.update(instance_id, |s| { self.state_store
s.status = WorkflowStatus::Cancelled; .update(instance_id, |s| {
})?; s.status = WorkflowStatus::Cancelled;
})
.await?;
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowCancelled { instance_id }) .publish(WorkflowEvent::WorkflowCancelled { instance_id })
@@ -799,8 +830,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
} }
/// Get workflow status /// Get workflow status
pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> { pub async fn get_status(
self.state_store.load(instance_id) &self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowState<D>> {
self.state_store.load(instance_id).await
} }
/// Wait for a workflow to complete with adaptive polling /// Wait for a workflow to complete with adaptive polling
@@ -829,6 +863,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let state = self let state = self
.get_status(instance_id) .get_status(instance_id)
.await
.map_err(|e| format!("Failed to get workflow status: {:?}", e))?; .map_err(|e| format!("Failed to get workflow status: {:?}", e))?;
let result = match state.status { let result = match state.status {
@@ -857,7 +892,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
} }
}; };
self.state_store.cleanup_if_terminal(instance_id); self.state_store.cleanup_if_terminal(instance_id).await;
return result; return result;
} }
} }
sgl-model-gateway/src/workflow/state.rs
+52 -43
View File
@@ -2,6 +2,7 @@
use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration}; use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration};
use async_trait::async_trait;
use parking_lot::RwLock; use parking_lot::RwLock;
use super::types::{ use super::types::{
@@ -13,39 +14,43 @@ use super::types::{
/// ///
/// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis). /// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis).
/// The default implementation is `InMemoryStore` which keeps state in memory. /// The default implementation is `InMemoryStore` which keeps state in memory.
#[async_trait]
pub trait StateStore<D: WorkflowData>: Send + Sync + Clone { pub trait StateStore<D: WorkflowData>: Send + Sync + Clone {
/// Save workflow state /// Save workflow state
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>; async fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>;
/// Load workflow state by instance ID /// Load workflow state by instance ID
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>; async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>;
/// Update workflow state using a closure /// Update workflow state using a closure
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> async fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where where
F: FnOnce(&mut WorkflowState<D>); F: FnOnce(&mut WorkflowState<D>) + Send;
/// Delete workflow state /// Delete workflow state
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>; async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>;
/// List all active workflows (Running or Pending) /// List all active workflows (Running or Pending)
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>; async fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// List all workflows /// List all workflows
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>; async fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// Check if workflow is cancelled without loading full state /// Check if workflow is cancelled without loading full state
fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>; async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>;
/// Clean up old completed/failed/cancelled workflows beyond a time threshold /// Clean up old completed/failed/cancelled workflows beyond a time threshold
fn cleanup_old_workflows(&self, ttl: Duration) -> usize; async fn cleanup_old_workflows(&self, ttl: Duration) -> usize;
/// Get just the workflow context without cloning the entire state /// Get just the workflow context without cloning the entire state
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>>; async fn get_context(
&self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowContext<D>>;
/// Clean up a specific workflow immediately if it's in a terminal state /// Clean up a specific workflow immediately if it's in a terminal state
/// Returns true if the workflow was removed, false otherwise /// Returns true if the workflow was removed, false otherwise
fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool; async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool;
} }
/// In-memory state storage for workflow instances /// In-memory state storage for workflow instances
@@ -84,13 +89,14 @@ impl<D: WorkflowData> Default for InMemoryStore<D> {
} }
} }
#[async_trait]
impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> { impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> { async fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> {
self.states.write().insert(state.instance_id, state); self.states.write().insert(state.instance_id, state);
Ok(()) Ok(())
} }
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> { async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
self.states self.states
.read() .read()
.get(&instance_id) .get(&instance_id)
@@ -98,28 +104,9 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
.ok_or(WorkflowError::NotFound(instance_id)) .ok_or(WorkflowError::NotFound(instance_id))
} }
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> { async fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states.write().remove(&instance_id);
Ok(())
}
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where where
F: FnOnce(&mut WorkflowState<D>), F: FnOnce(&mut WorkflowState<D>) + Send,
{ {
let mut states = self.states.write(); let mut states = self.states.write();
let state = states let state = states
@@ -130,15 +117,26 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
Ok(()) Ok(())
} }
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>> { async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states self.states.write().remove(&instance_id);
.read() Ok(())
.get(&instance_id)
.map(|s| s.context.clone())
.ok_or(WorkflowError::NotFound(instance_id))
} }
fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> { async fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
async fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> {
self.states self.states
.read() .read()
.get(&instance_id) .get(&instance_id)
@@ -146,7 +144,7 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
.ok_or(WorkflowError::NotFound(instance_id)) .ok_or(WorkflowError::NotFound(instance_id))
} }
fn cleanup_old_workflows(&self, ttl: Duration) -> usize { async fn cleanup_old_workflows(&self, ttl: Duration) -> usize {
let now = chrono::Utc::now(); let now = chrono::Utc::now();
let mut states = self.states.write(); let mut states = self.states.write();
let initial_count = states.len(); let initial_count = states.len();
@@ -179,7 +177,18 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
removed_count removed_count
} }
fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool { async fn get_context(
&self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowContext<D>> {
self.states
.read()
.get(&instance_id)
.map(|s| s.context.clone())
.ok_or(WorkflowError::NotFound(instance_id))
}
async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
let mut states = self.states.write(); let mut states = self.states.write();
if let Some(state) = states.get(&instance_id) { if let Some(state) = states.get(&instance_id) {
if matches!( if matches!(
sgl-model-gateway/tests/wasm_test.rs
+1
View File
@@ -710,6 +710,7 @@ async fn test_wasm_module_execution() {
let state = engines let state = engines
.wasm_registration .wasm_registration
.get_status(instance_id) .get_status(instance_id)
.await
.expect("Failed to get workflow status"); .expect("Failed to get workflow status");
match state.status { match state.status {
sgl-model-gateway/tests/workflow_test.rs
+11 -11
View File
@@ -104,7 +104,7 @@ async fn test_simple_workflow_execution() {
sleep(Duration::from_millis(100)).await; sleep(Duration::from_millis(100)).await;
// Check status // Check status
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
assert_eq!(state.step_states.len(), 2); assert_eq!(state.step_states.len(), 2);
} }
@@ -149,7 +149,7 @@ async fn test_workflow_with_retry() {
sleep(Duration::from_millis(500)).await; sleep(Duration::from_millis(500)).await;
// Check that step was retried and eventually succeeded // Check that step was retried and eventually succeeded
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
let step_state = state.step_states.get(&StepId::new("retry_step")).unwrap(); let step_state = state.step_states.get(&StepId::new("retry_step")).unwrap();
@@ -200,7 +200,7 @@ async fn test_workflow_failure_after_max_retries() {
sleep(Duration::from_millis(500)).await; sleep(Duration::from_millis(500)).await;
// Check that workflow failed // Check that workflow failed
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Failed); assert_eq!(state.status, WorkflowStatus::Failed);
let step_state = state.step_states.get(&StepId::new("failing_step")).unwrap(); let step_state = state.step_states.get(&StepId::new("failing_step")).unwrap();
@@ -257,7 +257,7 @@ async fn test_workflow_continue_on_failure() {
sleep(Duration::from_millis(500)).await; sleep(Duration::from_millis(500)).await;
// Workflow should complete despite first step failing // Workflow should complete despite first step failing
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
// First step should be skipped // First step should be skipped
@@ -341,7 +341,7 @@ async fn test_workflow_context_sharing() {
sleep(Duration::from_millis(100)).await; sleep(Duration::from_millis(100)).await;
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
} }
@@ -434,7 +434,7 @@ async fn test_parallel_execution_no_dependencies() {
// Wait for completion - give enough time for async scheduling // Wait for completion - give enough time for async scheduling
for _ in 0..50 { for _ in 0..50 {
sleep(Duration::from_millis(50)).await; sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running { if state.status != WorkflowStatus::Running {
break; break;
} }
@@ -442,7 +442,7 @@ async fn test_parallel_execution_no_dependencies() {
let overall_duration = overall_start.elapsed(); let overall_duration = overall_start.elapsed();
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
// Check that all steps completed // Check that all steps completed
@@ -528,13 +528,13 @@ async fn test_dag_with_dependencies() {
// Poll until workflow completes (or timeout) // Poll until workflow completes (or timeout)
for _ in 0..50 { for _ in 0..50 {
sleep(Duration::from_millis(50)).await; sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running { if state.status != WorkflowStatus::Running {
break; break;
} }
} }
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed); assert_eq!(state.status, WorkflowStatus::Completed);
// Verify step C started after both A and B finished // Verify step C started after both A and B finished
@@ -617,13 +617,13 @@ async fn test_dag_dependency_failure_blocks_dependents() {
// Poll until workflow completes (or timeout) // Poll until workflow completes (or timeout)
for _ in 0..50 { for _ in 0..50 {
sleep(Duration::from_millis(50)).await; sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running { if state.status != WorkflowStatus::Running {
break; break;
} }
} }
let state = engine.get_status(instance_id).unwrap(); let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Failed); assert_eq!(state.status, WorkflowStatus::Failed);
// Step B should not have executed // Step B should not have executed