Skip to content

feat: add ready check #868

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
duyquang6 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add ready check #868

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 54 additions & 14 deletions crates/arkflow-core/src/engine/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
*/

use crate::config::EngineConfig;
use crate::input::Input;
use crate::output::Output;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
Expand All @@ -29,10 +31,20 @@ use axum::{routing::get, Router};
use serde::Serialize;
use tokio::net::TcpListener;

/// Holds references to stream components for health checking
struct StreamComponents {
input: Arc<dyn Input>,
output: Arc<dyn Output>,
error_output: Option<Arc<dyn Output>>,
}

struct EngineApiState {
health_state: Arc<HealthState>,
components: Vec<StreamComponents>,
}

/// Health check status
struct HealthState {
/// Whether the engine has been initialized
is_ready: AtomicBool,
/// Whether the engine is currently running
is_running: AtomicBool,
}
Expand Down Expand Up @@ -82,7 +94,6 @@ impl Engine {
Self {
config,
health_state: Arc::new(HealthState {
is_ready: AtomicBool::new(false),
is_running: AtomicBool::new(false),
}),
}
Expand All @@ -99,6 +110,7 @@ impl Engine {
async fn start_health_check_server(
&self,
cancellation_token: CancellationToken,
components: Vec<StreamComponents>,
) -> Result<(), Box<dyn std::error::Error>> {
let health_check = &self.config.health_check;

Expand All @@ -113,7 +125,10 @@ impl Engine {
.route(&*health_check.health_path, get(Self::handle_health))
.route(&*health_check.readiness_path, get(Self::handle_readiness))
.route(&*health_check.liveness_path, get(Self::handle_liveness))
.with_state(health_state);
.with_state(Arc::new(EngineApiState {
health_state,
components,
}));

let addr = &health_check.address;
let addr = addr.clone();
Expand Down Expand Up @@ -149,8 +164,8 @@ impl Engine {
///
/// # Arguments
/// * `state` - The shared health state containing running status
async fn handle_health(State(state): State<Arc<HealthState>>) -> impl IntoResponse {
let is_running = state.is_running.load(Ordering::SeqCst);
async fn handle_health(State(state): State<Arc<EngineApiState>>) -> impl IntoResponse {
let is_running = state.health_state.is_running.load(Ordering::SeqCst);
let status = if is_running { "healthy" } else { "unhealthy" };

let response = HealthResponse {
Expand All @@ -174,8 +189,25 @@ impl Engine {
///
/// # Arguments
/// * `state` - The shared health state containing readiness status
async fn handle_readiness(State(state): State<Arc<HealthState>>) -> impl IntoResponse {
let is_ready = state.is_ready.load(Ordering::SeqCst);
async fn handle_readiness(State(state): State<Arc<EngineApiState>>) -> impl IntoResponse {
let mut is_ready = true;

for component in state.components.iter() {
if !component.input.check_ready().await.is_ok()
|| !component.output.check_ready().await.is_ok()
{
is_ready = false;
}
if let Some(error_output) = component.error_output.as_ref() {
if !error_output.check_ready().await.is_ok() {
is_ready = false;
}
}
if !is_ready {
break;
}
}
Comment on lines +193 to +209

@coderabbitai coderabbitai Bot Oct 12, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Readiness must wait for the engine to be running.

With the new interface-only checks, every component currently reports Ok(()), so this handler now returns “ready” even before any stream has started (or after a fatal startup failure). That breaks the readiness signal consumers rely on for gating traffic. Gate the component loop behind health_state.is_running so we only declare readiness once the engine has actually transitioned to running.

-        let mut is_ready = true;
-
-        for component in state.components.iter() {
-            if !component.input.check_ready().await.is_ok()
-                || !component.output.check_ready().await.is_ok()
-            {
-                is_ready = false;
-            }
-            if let Some(error_output) = component.error_output.as_ref() {
-                if !error_output.check_ready().await.is_ok() {
-                    is_ready = false;
-                }
-            }
-            if !is_ready {
-                break;
-            }
-        }
+        let mut is_ready = state.health_state.is_running.load(Ordering::SeqCst);
+
+        if is_ready {
+            for component in state.components.iter() {
+                if !component.input.check_ready().await.is_ok()
+                    || !component.output.check_ready().await.is_ok()
+                {
+                    is_ready = false;
+                }
+                if let Some(error_output) = component.error_output.as_ref() {
+                    if !error_output.check_ready().await.is_ok() {
+                        is_ready = false;
+                    }
+                }
+                if !is_ready {
+                    break;
+                }
+            }
+        }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
let mut is_ready = true;
for component in state.components.iter() {
if !component.input.check_ready().await.is_ok()
|| !component.output.check_ready().await.is_ok()
{
is_ready = false;
}
if let Some(error_output) = component.error_output.as_ref() {
if !error_output.check_ready().await.is_ok() {
is_ready = false;
}
}
if !is_ready {
break;
}
}
let mut is_ready = state.health_state.is_running.load(Ordering::SeqCst);
if is_ready {
for component in state.components.iter() {
if !component.input.check_ready().await.is_ok()
|| !component.output.check_ready().await.is_ok()
{
is_ready = false;
}
if let Some(error_output) = component.error_output.as_ref() {
if !error_output.check_ready().await.is_ok() {
is_ready = false;
}
}
if !is_ready {
break;
}
}
}
🤖 Prompt for AI Agents 
crates/arkflow-core/src/engine/mod.rs around lines 193 to 209: the readiness
check currently iterates components and returns ready even when the engine
hasn't started because interface-only checks all return Ok; gate the component
readiness loop behind the engine running flag by first checking
health_state.is_running (or the equivalent accessor) and only perform the
per-component check when running is true — if not running, set readiness to not
ready (or return early) so readiness is only true after health_state.is_running
is true and all component checks succeed.

@duyquang6 duyquang6 Oct 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

readiness just checking input and output is ready to serve request or not, it not related to is_running check

@coderabbitai coderabbitai Bot Oct 12, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@duyquang6, thank you for the clarification. I understand now—the readiness check is specifically focused on whether the individual components (inputs and outputs) are ready to accept and process requests, which is independent of the engine's running state.

Your design makes sense: the readiness endpoint verifies component-level readiness (e.g., connection pools, message broker connections), while is_running tracks the engine's execution state. These are indeed separate concerns, and it's valid for components to report readiness even before stream processing begins.

I appreciate you explaining the design intent. This approach aligns with common patterns where readiness indicates infrastructure availability rather than active processing state.


let status = if is_ready { "ready" } else { "not ready" };

let response = ReadinessResponse {
Expand All @@ -198,7 +230,7 @@ impl Engine {
///
/// # Arguments
/// * `_` - Unused health state parameter
async fn handle_liveness(_: State<Arc<HealthState>>) -> impl IntoResponse {
async fn handle_liveness(_: State<Arc<EngineApiState>>) -> impl IntoResponse {
// As long as the server can respond, it is considered alive
let response = LivenessResponse {
status: "alive".to_string(),
Expand All @@ -220,9 +252,6 @@ impl Engine {
pub async fn run(&self) -> Result<(), Box<dyn std::error::Error>> {
let token = CancellationToken::new();

// Start the health check server
self.start_health_check_server(token.clone()).await?;

// Create and run all flows
let mut streams = Vec::new();
let mut handles = Vec::new();
Expand All @@ -241,8 +270,19 @@ impl Engine {
}
}

// Set the readiness status
self.health_state.is_ready.store(true, Ordering::SeqCst);
let components = streams
.iter()
.map(|stream| StreamComponents {
input: stream.get_input(),
output: stream.get_output(),
error_output: stream.get_error_output(),
})
.collect();

// Start the health check server
self.start_health_check_server(token.clone(), components)
.await?;

// Set up signal handlers
let mut sigint = signal(SignalKind::interrupt()).expect("Failed to set signal handler");
let mut sigterm = signal(SignalKind::terminate()).expect("Failed to set signal handler");
Expand Down
5 changes: 5 additions & 0 deletions crates/arkflow-core/src/input/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ pub trait Input: Send + Sync {

/// Close the input source connection
async fn close(&self) -> Result<(), Error>;

/// Check if the input source is ready and healthy
async fn check_ready(&self) -> Result<(), Error> {
Ok(())
}
}

pub struct NoopAck;
Expand Down
5 changes: 5 additions & 0 deletions crates/arkflow-core/src/output/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,11 @@ pub trait Output: Send + Sync {

/// Close the output destination connection
async fn close(&self) -> Result<(), Error>;

/// Check if the output destination is ready and healthy
async fn check_ready(&self) -> Result<(), Error> {
Ok(())
}
}

/// Output configuration
Expand Down
12 changes: 12 additions & 0 deletions crates/arkflow-core/src/stream/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -425,6 +425,18 @@ impl Stream {

Ok(())
}

pub fn get_input(&self) -> Arc<dyn Input> {
self.input.clone()
}

pub fn get_output(&self) -> Arc<dyn Output> {
self.output.clone()
}

pub fn get_error_output(&self) -> Option<Arc<dyn Output>> {
self.error_output.clone()
}
}

/// Stream configuration
Expand Down
Loading