chore(prover-service): setup basic health check server

Galxe · Nov 20, 2024 · 972900d · 972900d
1 parent f7c1ca0
commit 972900d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/prover-service/Cargo.lock b/prover-service/Cargo.lock
diff --git a/prover-service/Cargo.toml b/prover-service/Cargo.toml
@@ -12,6 +12,7 @@ path = "src/bin/server.rs"
 sp1-sdk = "3.0.0"
 alloy-sol-types = "0.8.12"
 tonic = "0.10"
+tonic-health = "0.10"
 prost = "0.12"
 tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
 clap = { version = "4.5", features = ["derive"] }

diff --git a/prover-service/src/bin/server.rs b/prover-service/src/bin/server.rs
@@ -2,6 +2,7 @@ use clap::Parser;
 use prover_service::proto::prover::prover_service_server::ProverServiceServer;
 use prover_service::service::ProverServiceImpl;
 use tonic::transport::Server;
+use tonic_health::server::health_reporter;
 
 #[derive(Parser)]
 #[command(
@@ -40,10 +41,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let addr = format!("[::1]:{}", args.port).parse()?;
     let service: ProverServiceImpl =
         ProverServiceImpl::new(&args.private_key, &args.elf_path, args.timeout_secs)?;
-    log::info!("ProverService listening on {}", addr);
+
+    // Set up health checks
+    let (mut health_reporter, health_service) = health_reporter();
+    health_reporter
+        .set_serving::<ProverServiceServer<ProverServiceImpl>>()
+        .await;
 
     // Serve requests
+    log::info!("Health Server and GRPC Prover Server listening on {}", addr);
     Server::builder()
+        .add_service(health_service)
         .add_service(ProverServiceServer::new(service))
         .serve(addr)
         .await?;