feat: health check in worker

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Changes
 
+- Added health check endpoints to the prover service (#1006).
 - Implemented serialization for `AccountHeader` (#996).
 - Updated Pingora crates to 0.4 and added polling time to the configuration file (#997).
 - Added support for `miden-tx-prover` proxy to update workers on a running proxy (#989).

bin/tx-prover/Cargo.toml

@@ -54,6 +54,7 @@ serde_qs = { version = "0.13" }
 tokio = { version = "1.38", optional = true, features = ["full"] }
 tokio-stream = { version = "0.1", optional = true, features = [ "net" ]}
 toml = { version = "0.8" }
+tonic-health = { version = "0.12" }
 tonic-web = { version = "0.12", optional = true }
 tracing = { version = "0.1", optional = true }
 tracing-subscriber = { version = "0.3", features = ["fmt",  "json",  "env-filter"], optional = true }

bin/tx-prover/README.md

@@ -59,6 +59,8 @@ max_queue_items = 10
 max_retries_per_request = 1
 # Maximum amount of requests that a given IP address can make per second
 max_req_per_sec = 5
+# Interval to check the health of the workers
+health_check_interval_secs = 1
 
 [[workers]]
 host = "0.0.0.0"
@@ -102,6 +104,12 @@ This changes will be persisted to the configuration file.
 
 Note that, in order to update the workers, the proxy must be running in the same computer as the command is being executed because it will check if the client address is localhost to avoid any security issues.
 
+### Health check
+
+The worker service implements the [gRPC Health Check](https://grpc.io/docs/guides/health-checking/) standard, and includes the methods described in this [official proto file](https://github.com/grpc/grpc-proto/blob/master/grpc/health/v1/health.proto).
+
+The proxy service uses this health check to determine if a worker is available to receive requests. If a worker is not available, it will be removed from the set of workers that the proxy can use to send requests, and will persist this change in the configuration file.
+
 ## Logging
 
 Both the worker and the proxy will use the `info` log level by default, but it can be changed by setting the `RUST_LOG` environment variable.

bin/tx-prover/src/commands/mod.rs

@@ -9,6 +9,7 @@ use init::Init;
 use miden_tx_prover::PROVER_SERVICE_CONFIG_FILE_NAME;
 use proxy::StartProxy;
 use serde::{Deserialize, Serialize};
+use tracing::debug;
 use update_workers::{AddWorkers, RemoveWorkers, UpdateWorkers};
 use worker::StartWorker;
 
@@ -41,6 +42,8 @@ pub struct ProxyConfig {
     pub max_req_per_sec: isize,
     /// Time in milliseconds to poll available workers.
     pub available_workers_polling_time_ms: u64,
+    /// Health check interval in seconds.
+    pub health_check_interval_secs: u64,
 }
 
 impl Default for ProxyConfig {
@@ -55,6 +58,7 @@ impl Default for ProxyConfig {
             max_retries_per_request: 1,
             max_req_per_sec: 5,
             available_workers_polling_time_ms: 20,
+            health_check_interval_secs: 1,
         }
     }
 }
@@ -96,10 +100,19 @@ impl ProxyConfig {
             .write(config_as_toml_string.as_bytes())
             .map_err(|err| format!("error writing to file: {err}"))?;
 
-        println!("Config updated successfully");
+        debug!("Config updated successfully");
 
         Ok(())
     }
+
+    /// Updates the workers in the configuration with the new list.
+    pub(crate) fn set_workers(workers: Vec<WorkerConfig>) -> Result<(), String> {
+        let mut proxy_config = Self::load_config_from_file()?;
+
+        proxy_config.workers = workers;
+
+        proxy_config.save_to_config_file()
+    }
 }
 
 /// Configuration for a worker
@@ -137,7 +150,7 @@ pub enum Command {
     /// values. The file will be named as defined in the
     /// [miden_tx_prover::PROVER_SERVICE_CONFIG_FILE_NAME] constant.
     Init(Init),
-    /// Starts the workers defined in the config file.
+    /// Starts the workers with the configuration defined in the command.
     StartWorker(StartWorker),
     /// Starts the proxy defined in the config file.
     StartProxy(StartProxy),
@@ -155,26 +168,22 @@ pub enum Command {
 
 /// CLI entry point
 impl Cli {
-    pub fn execute(&self) -> Result<(), String> {
+    pub async fn execute(&self) -> Result<(), String> {
         match &self.action {
             // For the `StartWorker` command, we need to create a new runtime and run the worker
-            Command::StartWorker(worker_init) => {
-                let rt = tokio::runtime::Runtime::new()
-                    .map_err(|e| format!("Failed to create runtime: {:?}", e))?;
-                rt.block_on(worker_init.execute())
-            },
-            Command::StartProxy(proxy_init) => proxy_init.execute(),
+            Command::StartWorker(worker_init) => worker_init.execute().await,
+            Command::StartProxy(proxy_init) => proxy_init.execute().await,
             Command::Init(init) => {
                 // Init does not require async, so run directly
                 init.execute()
             },
             Command::AddWorkers(update_workers) => {
                 let update_workers: UpdateWorkers = update_workers.clone().into();
-                update_workers.execute()
+                update_workers.execute().await
             },
             Command::RemoveWorkers(update_workers) => {
                 let update_workers: UpdateWorkers = update_workers.clone().into();
-                update_workers.execute()
+                update_workers.execute().await
             },
         }
     }

bin/tx-prover/src/commands/proxy.rs

@@ -1,8 +1,13 @@
 use clap::Parser;
-use pingora::{apps::HttpServerOptions, lb::Backend, prelude::Opt, server::Server};
+use pingora::{
+    apps::HttpServerOptions,
+    lb::Backend,
+    prelude::{background_service, Opt},
+    server::Server,
+};
 use pingora_proxy::http_proxy_service;
 
-use crate::proxy::LoadBalancer;
+use crate::proxy::{LoadBalancer, LoadBalancerState};
 
 /// Starts the proxy defined in the config file.
 #[derive(Debug, Parser)]
@@ -13,8 +18,8 @@ impl StartProxy {
     ///
     /// This method will first read the config file to get the list of workers to start. It will
     /// then start a proxy with each worker as a backend.
-    pub fn execute(&self) -> Result<(), String> {
-        let mut server = Server::new(Some(Opt::default())).expect("Failed to create server");
+    pub async fn execute(&self) -> Result<(), String> {
+        let mut server = Server::new(Some(Opt::default())).map_err(|err| err.to_string())?;
         server.bootstrap();
 
         let proxy_config = super::ProxyConfig::load_config_from_file()?;
@@ -23,25 +28,33 @@ impl StartProxy {
             .workers
             .iter()
             .map(|worker| format!("{}:{}", worker.host, worker.port))
-            .map(|worker| Backend::new(&worker).expect("Failed to create backend"))
-            .collect::<Vec<Backend>>();
+            .map(|worker| Backend::new(&worker).map_err(|err| err.to_string()))
+            .collect::<Result<Vec<Backend>, String>>()?;
 
-        let worker_lb = LoadBalancer::new(workers, &proxy_config);
+        let worker_lb = LoadBalancerState::new(workers, &proxy_config).await?;
+
+        let health_check_service = background_service("health_check", worker_lb);
+        let worker_lb = health_check_service.task();
 
         // Set up the load balancer
-        let mut lb = http_proxy_service(&server.configuration, worker_lb);
+        let mut lb = http_proxy_service(&server.configuration, LoadBalancer(worker_lb));
 
         let proxy_host = proxy_config.host;
         let proxy_port = proxy_config.port.to_string();
         lb.add_tcp(format!("{}:{}", proxy_host, proxy_port).as_str());
-        let logic = lb.app_logic_mut().expect("No app logic found");
+        let logic = lb.app_logic_mut().ok_or("Failed to get app logic")?;
         let mut http_server_options = HttpServerOptions::default();
 
         // Enable HTTP/2 for plaintext
         http_server_options.h2c = true;
         logic.server_options = Some(http_server_options);
 
+        server.add_service(health_check_service);
         server.add_service(lb);
-        server.run_forever();
+        tokio::task::spawn_blocking(|| server.run_forever())
+            .await
+            .map_err(|err| err.to_string())?;
+
+        Ok(())
     }
 }

bin/tx-prover/src/commands/update_workers.rs

@@ -57,10 +57,8 @@ impl UpdateWorkers {
     /// - If the request fails.
     /// - If the status code is not successful.
     /// - If the X-Worker-Count header is missing.
-    pub fn execute(&self) -> Result<(), String> {
+    pub async fn execute(&self) -> Result<(), String> {
         // Define a runtime
-        let rt = tokio::runtime::Runtime::new()
-            .map_err(|e| format!("Failed to create runtime: {:?}", e))?;
 
         let query_params = serde_qs::to_string(&self).map_err(|err| err.to_string())?;
 
@@ -79,7 +77,7 @@ impl UpdateWorkers {
             .map_err(|err| err.to_string())?;
 
         // Make the request
-        let response = rt.block_on(client.get(url).send()).map_err(|err| err.to_string())?;
+        let response = client.get(url).send().await.map_err(|err| err.to_string())?;
 
         // Check status code
         if !response.status().is_success() {

bin/tx-prover/src/commands/worker.rs

@@ -1,6 +1,8 @@
 use clap::Parser;
+use miden_tx_prover::generated::api_server::ApiServer;
 use tokio::net::TcpListener;
 use tokio_stream::wrappers::TcpListenerStream;
+use tonic_health::server::health_reporter;
 use tracing::info;
 
 use crate::api::RpcListener;
@@ -22,6 +24,10 @@ impl StartWorker {
     /// This method receives the host and port from the CLI and starts a worker on that address.
     /// In case that one of the parameters is not provided, it will default to `0.0.0.0` for the
     /// host and `50051` for the port.
+    ///
+    /// The worker includes a health reporter that will mark the service as serving, following the
+    /// [gRPC health checking protocol](
+    /// https://github.com/grpc/grpc-proto/blob/master/grpc/health/v1/health.proto).
     pub async fn execute(&self) -> Result<(), String> {
         let worker_addr = format!("{}:{}", self.host, self.port);
         let rpc =
@@ -32,9 +38,16 @@ impl StartWorker {
             rpc.listener.local_addr().map_err(|err| err.to_string())?
         );
 
+        // Create a health reporter
+        let (mut health_reporter, health_service) = health_reporter();
+
+        // Mark the service as serving
+        health_reporter.set_serving::<ApiServer<RpcListener>>().await;
+
         tonic::transport::Server::builder()
             .accept_http1(true)
             .add_service(tonic_web::enable(rpc.api_service))
+            .add_service(health_service)
             .serve_with_incoming(TcpListenerStream::new(rpc.listener))
             .await
             .map_err(|err| err.to_string())?;

bin/tx-prover/src/main.rs

@@ -5,7 +5,8 @@ mod utils;
 use commands::Cli;
 use utils::setup_tracing;
 
-fn main() -> Result<(), String> {
+#[tokio::main]
+async fn main() -> Result<(), String> {
     use clap::Parser;
 
     setup_tracing();
@@ -14,7 +15,7 @@ fn main() -> Result<(), String> {
     let cli = Cli::parse();
 
     // execute cli action
-    cli.execute()
+    cli.execute().await
 }
 
 // TESTS