feat: Readiness endpoint

objectstore-server/src/cli.rs

@@ -1,3 +1,4 @@
+use std::fs::File;
 use std::path::PathBuf;
 use std::thread;
 use std::time::Duration;
@@ -6,6 +7,7 @@ use anyhow::Result;
 use argh::FromArgs;
 
 use crate::config::Config;
+use crate::endpoints::health::SHUTDOWN_MARKER_PATH;
 use crate::{healthcheck, observability, web};
 
 /// Objectstore API webserver.
@@ -26,6 +28,7 @@ enum Command {
     Healthcheck(HealthcheckCommand),
     Version(VersionCommand),
     Sleep(SleepCommand),
+    Down(DownCommand),
 }
 
 /// run the objectstore web server
@@ -54,6 +57,15 @@ struct SleepCommand {
     seconds: u64,
 }
 
+/// mark the server as not ready
+#[derive(Debug, FromArgs)]
+#[argh(subcommand, name = "down")]
+struct DownCommand {
+    /// seconds to sleep after creating the marker file
+    #[argh(positional)]
+    sleep_seconds: Option<u64>,
+}
+
 /// Bootstrap the runtime and execute the CLI command.
 pub fn execute() -> Result<()> {
     let args: Args = argh::from_env();
@@ -69,6 +81,14 @@ pub fn execute() -> Result<()> {
         return Ok(());
     }
 
+    if let Command::Down(DownCommand { sleep_seconds }) = args.command {
+        File::create(SHUTDOWN_MARKER_PATH)?;
+        if let Some(seconds) = sleep_seconds {
+            thread::sleep(Duration::from_secs(seconds));
+        }
+        return Ok(());
+    }
+
     let config = Config::load(args.config.as_deref())?;
 
     // Ensure a rustls crypto provider is installed, required on distroless.
@@ -95,7 +115,9 @@ pub fn execute() -> Result<()> {
         match args.command {
             Command::Run(RunCommand {}) => web::server(config).await,
             Command::Healthcheck(HealthcheckCommand {}) => healthcheck::healthcheck(config).await,
-            Command::Version(VersionCommand {}) | Command::Sleep(SleepCommand { .. }) => {
+            Command::Version(VersionCommand {})
+            | Command::Sleep(SleepCommand { .. })
+            | Command::Down(DownCommand { .. }) => {
                 unreachable!()
             }
         }

objectstore-server/src/endpoints/health.rs

@@ -1,12 +1,29 @@
+use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::{Router, routing};
 
 use crate::state::ServiceState;
 
+pub const SHUTDOWN_MARKER_PATH: &str = "/tmp/objectstore.down";
+
 pub fn router() -> Router<ServiceState> {
-    Router::new().route("/health", routing::get(health))
+    Router::new()
+        .route("/health", routing::get(health))
+        .route("/ready", routing::get(ready))
 }
 
 async fn health() -> impl IntoResponse {
     "OK"
 }
+
+async fn ready() -> impl IntoResponse {
+    if tokio::fs::try_exists(SHUTDOWN_MARKER_PATH)
+        .await
+        .unwrap_or(false)
+    {
+        tracing::debug!("Shutdown marker exists, failing readiness");
+        (StatusCode::SERVICE_UNAVAILABLE, "Shutting down")
+    } else {
+        (StatusCode::OK, "OK")
+    }
+}

objectstore-server/src/endpoints/mod.rs

@@ -8,7 +8,7 @@ use crate::state::ServiceState;
 
 mod batch;
 pub mod common;
-mod health;
+pub mod health;
 mod objects;
 
 pub fn routes() -> Router<ServiceState> {