feat(gateway): Adding broker and readyz endpoint (#6)

* feat(publish): Add core publish service with testing in
frastructure

* adding newline

* Fix linting

* cleaning up broker logic

* feat(gateway): Adding broker and readyz endpoint

Author: IanWoodard

Cargo.lock

@@ -4,9 +4,11 @@ version = "0.1.0"
 edition.workspace = true
 
 [dependencies]
+broker = { path = "../../shared/broker" }
 anyhow = { workspace = true }
 axum = { workspace = true }
 dotenvy = { workspace = true }
+redis = { workspace = true }
 serde = { workspace = true }
 tokio = { workspace = true }
 tower-http = { workspace = true }

services/gateway/Cargo.toml

@@ -1,15 +1,30 @@
+use std::{collections::HashMap, time::Duration};
+
 use axum::{Json, extract::State, http::StatusCode};
+use broker::RedisOperations;
 use serde::Serialize;
+use tokio::time::timeout;
+
+use crate::state::{AppState, HealthState};
 
-use crate::state::AppState;
+const REDIS_PING_TIMEOUT: u64 = 1;
+
+#[derive(Clone, Copy, Debug, PartialEq, Serialize)]
+#[serde(rename_all = "lowercase")]
+enum Status {
+    Ok,
+    Fail,
+}
 
 #[derive(Serialize)]
 pub struct HealthResponse {
     status: String,
     uptime: u64,
 }
 
-pub async fn healthz_handler(State(state): State<AppState>) -> (StatusCode, Json<HealthResponse>) {
+pub async fn healthz_handler(
+    State(state): State<HealthState>,
+) -> (StatusCode, Json<HealthResponse>) {
     let uptime = state.start_time.elapsed().as_secs();
     (
         StatusCode::OK,
@@ -20,24 +35,91 @@ pub async fn healthz_handler(State(state): State<AppState>) -> (StatusCode, Json
     )
 }
 
+#[derive(Serialize)]
+pub struct ReadyResponse {
+    status: Status,
+    dependencies: HashMap<&'static str, Status>,
+}
+
+async fn check_readiness<R: RedisOperations>(redis: &R) -> ReadyResponse {
+    let mut deps = HashMap::new();
+    let mut status = Status::Ok;
+
+    match timeout(Duration::from_secs(REDIS_PING_TIMEOUT), redis.ping()).await {
+        Ok(Ok(_)) => {
+            deps.insert("redis", Status::Ok);
+        }
+        _ => {
+            deps.insert("redis", Status::Fail);
+            status = Status::Fail;
+        }
+    }
+
+    ReadyResponse {
+        status,
+        dependencies: deps,
+    }
+}
+
+pub async fn readyz_handler(State(state): State<AppState>) -> (StatusCode, Json<ReadyResponse>) {
+    let response = check_readiness(&state.redis).await;
+    let status_code = match response.status {
+        Status::Ok => StatusCode::OK,
+        Status::Fail => StatusCode::SERVICE_UNAVAILABLE,
+    };
+
+    (status_code, Json(response))
+}
+
 #[cfg(test)]
 mod tests {
+    use crate::state::HealthState;
+
     use super::*;
     use axum::http::StatusCode;
+    use broker::{Error as BrokerError, MockRedisOperations};
     use tokio::time::Instant;
 
-    fn mock_app_state() -> AppState {
-        AppState {
-            start_time: Instant::now(),
-        }
-    }
-
     #[tokio::test]
     async fn test_healthz_returns_ok() {
-        let state = mock_app_state();
+        let state = HealthState {
+            start_time: Instant::now(),
+        };
         let (status, response) = healthz_handler(State(state)).await;
 
         assert_eq!(status, StatusCode::OK);
         assert_eq!(response.status, "ok");
     }
+
+    #[tokio::test]
+    async fn test_readyz_redis_healthy() {
+        let mut mock_redis = MockRedisOperations::new();
+
+        mock_redis
+            .expect_ping()
+            .times(1)
+            .returning(|| Ok("PONG".to_string()));
+
+        let response = check_readiness(&mock_redis).await;
+
+        assert_eq!(response.status, Status::Ok);
+        assert_eq!(response.dependencies.get("redis"), Some(&Status::Ok));
+    }
+
+    #[tokio::test]
+    async fn test_readyz_redis_error() {
+        let mut mock_redis = MockRedisOperations::new();
+
+        mock_redis.expect_ping().times(1).returning(|| {
+            Err(BrokerError::Redis(redis::RedisError::from((
+                redis::ErrorKind::IoError,
+                "Connection refused",
+            ))))
+        });
+
+        let response = check_readiness(&mock_redis).await;
+
+        assert_eq!(response.status, Status::Fail);
+        assert_eq!(response.dependencies.get("redis"), Some(&Status::Fail));
+    }
 }

services/gateway/src/handlers/health.rs

@@ -1,8 +1,12 @@
 use std::net::SocketAddr;
 
 use axum::{Router, routing::get};
+use broker::RedisClient;
 use dotenvy::dotenv;
-use gateway::{handlers::health::healthz_handler, state::AppState};
+use gateway::{
+    handlers::health::{healthz_handler, readyz_handler},
+    state::AppState,
+};
 use tokio::{net::TcpListener, runtime::Runtime, time::Instant};
 
 fn main() -> anyhow::Result<()> {
@@ -13,13 +17,18 @@ fn main() -> anyhow::Result<()> {
 
 async fn async_main() -> anyhow::Result<(), anyhow::Error> {
     let port = std::env::var("GATEWAY_PORT")?;
+    let redis_url = std::env::var("REDIS_URL")?;
+
+    let broker = RedisClient::new(redis_url.as_str()).await?;
 
     let state = AppState {
+        redis: broker,
         start_time: Instant::now(),
     };
 
     let mut app = Router::new()
         .route("/healthz", get(healthz_handler))
+        .route("/readyz", get(readyz_handler))
         .with_state(state);
 
     #[cfg(debug_assertions)]

services/gateway/src/main.rs

@@ -1,6 +1,21 @@
+use broker::RedisClient;
 use tokio::time::Instant;
 
 #[derive(Clone)]
 pub struct AppState {
+    pub redis: RedisClient,
     pub start_time: Instant,
 }
+
+#[derive(Clone)]
+pub struct HealthState {
+    pub start_time: Instant,
+}
+
+impl axum::extract::FromRef<AppState> for HealthState {
+    fn from_ref(app: &AppState) -> Self {
+        Self {
+            start_time: app.start_time,
+        }
+    }
+}