refactor: liveness & readiness

argumentcomputer · Sep 13, 2024 · 9d381dd · 9d381dd
1 parent 69f4ead
commit 9d381dd
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 9 deletions.
diff --git a/aptos/proof-server/src/bin/proof_server.rs b/aptos/proof-server/src/bin/proof_server.rs
@@ -109,6 +109,7 @@ async fn main() -> Result<()> {
 
     let app = Router::new()
         .route("/health", get(health_check))
+        .route("/ready", get(ready_check))
         .route("/inclusion/proof", post(inclusion_proof))
         .route("/epoch/proof", post(epoch_proof))
         .route("/epoch/verify", post(epoch_verify))
@@ -128,6 +129,19 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+async fn health_check(State(state): State<ServerState>) -> impl IntoResponse {
+    StatusCode::OK
+}
+
+async fn ready_check(State(state): State<ServerState>) -> impl IntoResponse {
+    let active_requests = state.active_requests.load(Ordering::SeqCst);
+    if active_requests > 0 {
+        StatusCode::CONFLICT
+    } else {
+        StatusCode::OK
+    }
+}
+
 async fn inclusion_proof(
     State(state): State<ServerState>,
     request: axum::extract::Request,
@@ -403,15 +417,6 @@ async fn forward_request(
     Ok(res_bytes.to_vec())
 }
 
-async fn health_check(State(state): State<ServerState>) -> impl IntoResponse {
-    let active_requests = state.active_requests.load(Ordering::SeqCst);
-    if active_requests > 0 {
-        StatusCode::CONFLICT
-    } else {
-        StatusCode::OK
-    }
-}
-
 async fn count_requests_middleware(
     State(state): State<ServerState>,
     req: axum::http::Request<Body>,

diff --git a/docker/proof-server-chart/templates/proof-server-deployment.yaml b/docker/proof-server-chart/templates/proof-server-deployment.yaml
@@ -43,6 +43,12 @@ spec:
             port: 8080
           initialDelaySeconds: 10
           periodSeconds: 5
+        readinessProbe:
+          httpGet:
+            path: /ready
+            port: 8080
+          initialDelaySeconds: 15
+          periodSeconds: 10
         name: proof-server
         ports:
         - containerPort: 8080

diff --git a/ethereum/light-client/src/bin/proof_server.rs b/ethereum/light-client/src/bin/proof_server.rs
@@ -81,6 +81,7 @@ async fn main() -> Result<()> {
 
     let app = Router::new()
         .route("/health", get(health_check))
+        .route("/ready", get(ready_check))
         .route("/inclusion/proof", post(inclusion_proof))
         .route("/committee/proof", post(committee_proof))
         .route("/committee/verify", post(committee_verify))
@@ -100,6 +101,10 @@ async fn main() -> Result<()> {
 }
 
 async fn health_check(State(state): State<ServerState>) -> impl IntoResponse {
+        StatusCode::OK
+}
+
+async fn ready_check(State(state): State<ServerState>) -> impl IntoResponse {
     let active_requests = state.active_requests.load(Ordering::SeqCst);
     if active_requests > 0 {
         StatusCode::CONFLICT

diff --git a/kadena/light-client/src/bin/proof_server.rs b/kadena/light-client/src/bin/proof_server.rs
@@ -65,6 +65,7 @@ async fn main() -> Result<()> {
 
     let app = Router::new()
         .route("/health", get(health_check))
+        .route("/ready", get(ready_check))
         .layer(axum::middleware::from_fn_with_state(
             state.clone(),
             count_requests_middleware,
@@ -80,6 +81,10 @@ async fn main() -> Result<()> {
 }
 
 async fn health_check(State(state): State<ServerState>) -> impl IntoResponse {
+    StatusCode::OK
+}
+
+async fn ready_check(State(state): State<ServerState>) -> impl IntoResponse {
     let active_requests = state.active_requests.load(Ordering::SeqCst);
     if active_requests > 0 {
         StatusCode::CONFLICT
@@ -88,6 +93,7 @@ async fn health_check(State(state): State<ServerState>) -> impl IntoResponse {
     }
 }
 
+
 #[allow(dead_code)]
 async fn forward_request(request_bytes: &[u8], snd_addr: &str) -> Result<Vec<u8>, StatusCode> {
     info!("Connecting to the secondary server");