Add health endpoint

planetary-social · Aug 15, 2024 · 2c29d53 · 2c29d53
1 parent bd80699
commit 2c29d53
Show file tree

Hide file tree

Showing 7 changed files with 145 additions and 5 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.86"
+axum = "0.7.5"
 cached = { version = "0.53.1", features = ["async"] }
 chrono = { version = "0.4.38", features = ["serde"] }
 config = "0.14.0"
@@ -21,5 +22,6 @@ thiserror = "1.0.63"
 time = "0.3.36"
 tokio = { version = "1.39.2", features = ["full"] }
 tokio-util = { version = "0.7.11", features = ["rt"] }
+tower-http = { version = "0.5.2", features = ["timeout", "trace"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
diff --git a/compose.yaml b/compose.yaml
@@ -20,6 +20,8 @@ services:
     depends_on:
       db:
         condition: service_healthy
+    ports:
+      - "3000:3000"
     restart: always
     attach: true
 

diff --git a/src/http_server.rs b/src/http_server.rs
@@ -0,0 +1,51 @@
+mod router;
+use anyhow::{Context, Result};
+use axum::Router;
+use router::create_router;
+use std::net::SocketAddr;
+use std::time::Duration;
+use tokio::time::timeout;
+use tokio_util::sync::CancellationToken;
+use tracing::info;
+
+pub struct HttpServer;
+impl HttpServer {
+    pub async fn run(cancellation_token: CancellationToken) -> Result<()> {
+        let router = create_router()?;
+
+        start_http_server(router, cancellation_token).await
+    }
+}
+
+async fn start_http_server(router: Router, cancellation_token: CancellationToken) -> Result<()> {
+    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
+    let listener = tokio::net::TcpListener::bind(addr).await?;
+    let token_clone = cancellation_token.clone();
+    let server_future = tokio::spawn(async {
+        axum::serve(listener, router)
+            .with_graceful_shutdown(shutdown_hook(token_clone))
+            .await
+            .context("Failed to start HTTP server")
+    });
+
+    await_shutdown(cancellation_token, server_future).await;
+
+    Ok(())
+}
+
+async fn await_shutdown(
+    cancellation_token: CancellationToken,
+    server_future: tokio::task::JoinHandle<Result<()>>,
+) {
+    cancellation_token.cancelled().await;
+    info!("Shutdown signal received.");
+    match timeout(Duration::from_secs(5), server_future).await {
+        Ok(_) => info!("HTTP service exited successfully."),
+        Err(e) => info!("HTTP service exited after timeout: {}", e),
+    }
+}
+
+async fn shutdown_hook(cancellation_token: CancellationToken) {
+    cancellation_token.cancelled().await;
+    info!("Exiting the process");
+}
diff --git a/src/http_server/router.rs b/src/http_server/router.rs
@@ -0,0 +1,38 @@
+use anyhow::Result;
+use axum::{http::HeaderMap, response::Html};
+use axum::{response::IntoResponse, routing::get, Router};
+use std::time::Duration;
+use tower_http::trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer};
+use tower_http::LatencyUnit;
+use tower_http::{timeout::TimeoutLayer, trace::DefaultOnFailure};
+use tracing::Level;
+
+pub fn create_router() -> Result<Router> {
+    let tracing_layer = TraceLayer::new_for_http()
+        .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
+        .on_response(
+            DefaultOnResponse::new()
+                .level(Level::INFO)
+                .latency_unit(LatencyUnit::Millis),
+        )
+        .on_failure(DefaultOnFailure::new().level(Level::ERROR));
+
+    Ok(Router::new()
+        .route("/", get(serve_root_page))
+        .layer(tracing_layer)
+        .layer(TimeoutLayer::new(Duration::from_secs(1))))
+}
+
+async fn serve_root_page(_headers: HeaderMap) -> impl IntoResponse {
+    let body = "
+    <html>
+        <head>
+            <title>Nos</title>
+        </head>
+        <body>
+            <h1>Healthy</h1>
+        </body>
+        ";
+
+    Html(body)
+}
diff --git a/src/main.rs b/src/main.rs
@@ -4,6 +4,7 @@ mod fetch_friendly_id;
 mod follow_change_handler;
 mod follows_differ;
 mod google_publisher;
+mod http_server;
 mod migrations;
 mod relay_subscriber;
 mod repo;
@@ -14,6 +15,7 @@ use crate::config::Config;
 use crate::domain::follow_change::FollowChange;
 use follow_change_handler::FollowChangeHandler;
 use follows_differ::FollowsDiffer;
+use http_server::HttpServer;
 use migrations::apply_migrations;
 use neo4rs::Graph;
 use nostr_sdk::prelude::*;
@@ -89,14 +91,21 @@ async fn main() -> Result<()> {
         .since(five_minutes_ago)
         .kind(Kind::ContactList)];
 
-    start_nostr_subscription(
+    let nostr_sub = start_nostr_subscription(
         shared_nostr_client,
-        &[relay],
+        [relay].into(),
         filters,
         event_sender,
         cancellation_token.clone(),
-    )
-    .await?;
+    );
+
+    let http_server = HttpServer::run(cancellation_token.clone());
+
+    tokio::select! {
+        _ = nostr_sub => info!("Nostr subscription ended"),
+        _ = http_server => info!("HTTP server ended"),
+        _ = cancellation_token.cancelled() => info!("Cancellation token cancelled"),
+    }
 
     info!("Finished Nostr subscription");
 

diff --git a/src/relay_subscriber.rs b/src/relay_subscriber.rs
@@ -21,7 +21,7 @@ pub fn create_client() -> Client {
 
 pub async fn start_nostr_subscription(
     nostr_client: Client,
-    relays: &[String],
+    relays: Vec<String>,
     filters: Vec<Filter>,
     event_tx: Sender<Box<Event>>,
     cancellation_token: CancellationToken,