feat(coprocessor): implement health-checking logic for zkproof-worker (#318)

* feat(coprocessor): add healthz server as a common util

* chore(coprocessor): implement HealthCheckService for zkproof worker

* chore(coprocessor): refactor the health-checking logic to be aligned with the spec

- Add dependencies in Json response
- Use ok/fail
- Add is_alive func

* feat(coprocessor): implement HealthCheckService::is_alive for zkproof worker

* chore(coprocessor): add version endpoint support in healthz

* chore(coprocessor): add simple impl of HealthCheckService::get_version

Author: goshawk-3

coprocessor/fhevm-engine/Cargo.lock

@@ -52,6 +52,7 @@ tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["fmt", "json"] }
 humantime = "2.2.0"
 bytesize = "2.0.1"
+http = "1.3.1"
 
 [profile.dev.package.tfhe]
 overflow-checks = false

coprocessor/fhevm-engine/Cargo.toml

@@ -23,6 +23,10 @@ tonic  = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 bytesize = { workspace = true}
+tokio-util = { workspace = true}
+axum = { workspace = true}
+serde_json = { workspace = true}
+http = {workspace = true}
 
 # crates.io dependencies
 paste = "1.0.15"
@@ -34,6 +38,8 @@ opentelemetry-otlp = { workspace = true }
 opentelemetry_sdk = { workspace = true }
 opentelemetry-semantic-conventions = { workspace = true }
 
+
+
 [features]
 nightly-avx512 = ["tfhe/nightly-avx512"]
 gpu = ["tfhe/gpu"]

coprocessor/fhevm-engine/fhevm-engine-common/Cargo.toml

@@ -0,0 +1,176 @@
+use axum::{
+    extract::State,
+    http::StatusCode,
+    response::{IntoResponse, Json},
+    routing::{get, Router},
+};
+use serde::Serialize;
+use sqlx::PgPool;
+use std::{collections::HashMap, net::SocketAddr, sync::Arc};
+use tokio::net::TcpListener;
+use tokio_util::sync::CancellationToken;
+use tracing::{error, info};
+
+#[derive(Serialize)]
+struct HealthResponse {
+    status_code: String,
+    status: String,
+    dependencies: HashMap<&'static str, &'static str>,
+    details: String,
+}
+
+impl From<HealthStatus> for HealthResponse {
+    fn from(status: HealthStatus) -> Self {
+        let details = status.error_details();
+        let dependencies: HashMap<&'static str, &'static str> = status
+            .checks
+            .iter()
+            .map(|(&key, &value)| (key, if value { "ok" } else { "fail" }))
+            .collect();
+
+        Self {
+            status_code: if status.is_healthy() { "200" } else { "503" }.to_string(),
+            status: if status.is_healthy() {
+                "healthy".to_string()
+            } else {
+                "unhealthy".to_string()
+            },
+            dependencies,
+            details,
+        }
+    }
+}
+
+#[derive(Serialize)]
+pub struct Version {
+    pub name: &'static str,
+    pub version: &'static str,
+    pub build: &'static str,
+}
+
+pub trait HealthCheckService: Send + Sync {
+    fn health_check(&self) -> impl std::future::Future<Output = HealthStatus> + Send;
+    fn is_alive(&self) -> impl std::future::Future<Output = bool> + Send;
+    fn get_version(&self) -> Version;
+}
+
+pub struct HttpServer<S: HealthCheckService + Send + Sync + 'static> {
+    service: Arc<S>,
+    port: u16,
+    cancel_token: CancellationToken,
+}
+
+impl<S: HealthCheckService + Send + Sync + 'static> HttpServer<S> {
+    pub fn new(service: Arc<S>, port: u16, cancel_token: CancellationToken) -> Self {
+        Self {
+            service,
+            port,
+            cancel_token,
+        }
+    }
+    pub async fn start(&self) -> anyhow::Result<()> {
+        let app = Router::new()
+            .route("/healthz", get(Self::health_handler))
+            .route("/liveness", get(Self::liveness_handler))
+            .route("/version", get(Self::version_handler))
+            .with_state(self.service.clone());
+
+        let addr = SocketAddr::from(([0, 0, 0, 0], self.port));
+        info!("Starting HTTP server on {}", addr);
+
+        let shutdown = {
+            let cancel_token = self.cancel_token.clone();
+            async move {
+                cancel_token.cancelled().await;
+            }
+        };
+
+        let listener = TcpListener::bind(addr).await?;
+        let server =
+            axum::serve(listener, app.into_make_service()).with_graceful_shutdown(shutdown);
+
+        if let Err(err) = server.await {
+            error!("HTTP server error: {}", err);
+            return Err(anyhow::anyhow!("HTTP server error: {}", err));
+        }
+
+        Ok(())
+    }
+
+    async fn health_handler(State(service): State<Arc<S>>) -> impl IntoResponse {
+        let status = service.health_check().await;
+        let http_status = if status.is_healthy() {
+            StatusCode::OK
+        } else {
+            StatusCode::SERVICE_UNAVAILABLE
+        };
+
+        (http_status, Json(HealthResponse::from(status)))
+    }
+
+    async fn liveness_handler(State(service): State<Arc<S>>) -> impl IntoResponse {
+        if service.is_alive().await {
+            (
+                StatusCode::OK,
+                Json(serde_json::json!({
+                    "status_code": "200",
+                    "status": "alive"
+                })),
+            )
+        } else {
+            (
+                StatusCode::SERVICE_UNAVAILABLE,
+                Json(serde_json::json!({
+                    "status_code": "503",
+                    "status": "not_responding"
+                })),
+            )
+        }
+    }
+
+    async fn version_handler(State(service): State<Arc<S>>) -> impl IntoResponse {
+        let version = service.get_version();
+        (
+            StatusCode::OK,
+            Json(serde_json::json!(version)),
+        )
+    }
+}
+
+#[derive(Clone, Default)]
+pub struct HealthStatus {
+    pub checks: HashMap<&'static str, bool>,
+    pub error_details: Vec<String>,
+}
+
+impl HealthStatus {
+    /// Checks DB availability by reusing the service internal DB connection pool
+    ///
+    /// query has its internal timeout
+    pub async fn set_db_connected(&mut self, pool: &PgPool) {
+        let mut is_connected = false;
+        match sqlx::query("SELECT 1").execute(pool).await {
+            Ok(_) => {
+                is_connected = true;
+            }
+            Err(e) => {
+                self.error_details
+                    .push(format!("Database query error: {}", e));
+            }
+        }
+        self.checks.insert("database", is_connected);
+    }
+
+    pub fn is_healthy(&self) -> bool {
+        self.checks.iter().all(|(_, s)| *s)
+    }
+
+    pub fn error_details(&self) -> String {
+        self.error_details
+            .iter()
+            .filter(|s| !s.is_empty())
+            .cloned()
+            .collect::<Vec<_>>()
+            .join("; ")
+    }
+}

coprocessor/fhevm-engine/fhevm-engine-common/src/healthz_server.rs

@@ -1,3 +1,4 @@
+pub mod healthz_server;
 pub mod keys;
 pub mod telemetry;
 pub mod tenant_keys;

coprocessor/fhevm-engine/fhevm-engine-common/src/lib.rs

@@ -1,6 +1,11 @@
 use clap::{command, Parser};
+use fhevm_engine_common::healthz_server::HttpServer;
 use fhevm_engine_common::telemetry;
+use std::sync::Arc;
+use tokio::{join, task};
+use tokio_util::sync::CancellationToken;
 use tracing::{error, info, Level};
+use zkproof_worker::verifier::ZkProofService;
 
 #[derive(Parser, Debug, Clone)]
 #[command(version, about, long_about = None)]
@@ -40,6 +45,10 @@ pub struct Args {
         value_parser = clap::value_parser!(Level),
         default_value_t = Level::INFO)]
     pub log_level: Level,
+
+    /// HTTP server port for health checks
+    #[arg(long, default_value_t = 8080)]
+    health_check_port: u16,
 }
 
 pub fn parse_args() -> Args {
@@ -74,8 +83,33 @@ async fn main() {
         std::process::exit(1);
     }
 
-    info!("Starting zkProof worker...");
-    if let Err(err) = zkproof_worker::verifier::execute_verify_proofs_loop(&conf).await {
-        error!("Worker failed: {:?}", err);
-    }
+    let cancel_token = CancellationToken::new();
+    let service = ZkProofService::create(conf, cancel_token.child_token()).await;
+    let service = Arc::new(service);
+
+    let http_server = HttpServer::new(
+        service.clone(),
+        args.health_check_port,
+        cancel_token.child_token(),
+    );
+
+    let http_task = task::spawn(async move {
+        if let Err(err) = http_server.start().await {
+            error!(
+                task = "health_check",
+                "Error while running server: {:?}", err
+            );
+        }
+        anyhow::Ok(())
+    });
+
+    let service_task = async {
+        info!("Starting worker...");
+        if let Err(err) = service.run().await {
+            error!("Worker failed: {:?}", err);
+        }
+        Ok::<_, anyhow::Error>(())
+    };
+
+    let (_http_result, _service_result) = join!(http_task, service_task);
 }

coprocessor/fhevm-engine/zkproof-worker/src/bin/zkproof_worker.rs

@@ -1,5 +1,6 @@
-use std::time::Duration;
-
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
+use tokio::sync::RwLock;
 use fhevm_engine_common::{tenant_keys, utils::safe_serialize};
 use test_harness::instance::DBInstance;
 use tokio::time::sleep;
@@ -31,8 +32,10 @@ pub async fn setup() -> anyhow::Result<(sqlx::PgPool, DBInstance)> {
         .await
         .unwrap();
 
+    let last_active_at = Arc::new(RwLock::new(SystemTime::now()));
+    let db_pool = pool.clone();
     tokio::spawn(async move {
-        crate::verifier::execute_verify_proofs_loop(&conf)
+        crate::verifier::execute_verify_proofs_loop(db_pool, conf.clone(), last_active_at.clone() )
             .await
             .unwrap();
     });