From f19656486a6caee2147272908b9cce525bf56eba Mon Sep 17 00:00:00 2001
From: Silas Brack <silasbrack@gmail.com>
Date: Sun, 8 Mar 2026 17:41:58 +0100
Subject: [PATCH] Add timeout

---
 src/lib.rs           | 3 +++
 src/main.rs          | 7 +++++++
 src/server.rs        | 4 +++-
 tests/integration.rs | 1 +
 4 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index ba576f6..498c78a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,6 +6,7 @@ pub mod rebuild;
 pub mod server;
 
 use std::sync::Arc;
+use std::time::Duration;
 
 const DEFAULT_BODY_LIMIT: usize = 256 * 1024 * 1024; // 256 MB
 
@@ -13,6 +14,7 @@ pub struct Args {
     pub db_path: String,
     pub volumes: Vec<String>,
     pub replicas: usize,
+    pub voltimeout: Duration,
 }
 
 pub fn build_app(args: &Args) -> axum::Router {
@@ -36,6 +38,7 @@ pub fn build_app(args: &Args) -> axum::Router {
         db: db::Db::new(&args.db_path),
         volumes: Arc::new(args.volumes.clone()),
         replicas: args.replicas,
+        voltimeout: args.voltimeout,
         http: reqwest::Client::new(),
     };
 
diff --git a/src/main.rs b/src/main.rs
index f838428..7ec4727 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use clap::{Parser, Subcommand};
 
 #[derive(Parser)]
@@ -18,6 +20,10 @@ struct Cli {
     #[arg(short, long, env = "MKV_REPLICAS", default_value_t = 2)]
     replicas: usize,
 
+    /// Timeout for volume health checks (in milliseconds)
+    #[arg(long, env = "MKV_VOLTIMEOUT", default_value_t = 1000)]
+    voltimeout: u64,
+
     #[command(subcommand)]
     command: Commands,
 }
@@ -65,6 +71,7 @@ async fn main() {
         db_path: cli.db,
         volumes: cli.volumes,
         replicas: cli.replicas,
+        voltimeout: Duration::from_millis(cli.voltimeout),
     };
 
     match cli.command {
diff --git a/src/server.rs b/src/server.rs
index c251b1b..cf391fc 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -3,6 +3,7 @@ use axum::extract::{Path, Query, State};
 use axum::http::{HeaderMap, StatusCode};
 use axum::response::{IntoResponse, Response};
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::db;
 use crate::error::{AppError, VolumeError};
@@ -12,6 +13,7 @@ pub struct AppState {
     pub db: db::Db,
     pub volumes: Arc<Vec<String>>,
     pub replicas: usize,
+    pub voltimeout: Duration,
     pub http: reqwest::Client,
 }
 
@@ -57,7 +59,7 @@ pub async fn get_key(
     let mut results = Vec::with_capacity(volumes.len());
     for vol in &volumes {
         let url = format!("{vol}/{key}");
-        let healthy = match state.http.head(&url).send().await {
+        let healthy = match state.http.head(&url).timeout(state.voltimeout).send().await {
             Ok(resp) if resp.status().is_success() => true,
             Ok(resp) => {
                 tracing::warn!("volume {vol} returned {} for {key}", resp.status());
diff --git a/tests/integration.rs b/tests/integration.rs
index 890b23a..cfd46ad 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -19,6 +19,7 @@ async fn start_server() -> String {
             "http://localhost:3103".into(),
         ],
         replicas: 2,
+        voltimeout: std::time::Duration::from_secs(1),
     };
 
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();