From 1d3b9dddf562e24ec90a4958761a56909f6c8927 Mon Sep 17 00:00:00 2001 From: Silas Brack Date: Sun, 8 Mar 2026 13:09:55 +0100 Subject: [PATCH] Add unit test for failover --- src/server.rs | 140 ++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 125 insertions(+), 15 deletions(-) diff --git a/src/server.rs b/src/server.rs index ea27108..65e742c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -15,42 +15,83 @@ pub struct AppState { pub http: reqwest::Client, } +/// Result of probing volumes for a key. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProbeResult { + /// Found a healthy volume at this URL + Found(String), + /// All volumes were probed, none healthy + AllFailed, +} + +/// Pure function: given volumes and their probe results (true = healthy), +/// returns the first healthy volume URL for the key, or None. +pub fn first_healthy_volume(key: &str, volumes: &[String], results: &[bool]) -> ProbeResult { + for (vol, &healthy) in volumes.iter().zip(results) { + if healthy { + return ProbeResult::Found(format!("{vol}/{key}")); + } + } + ProbeResult::AllFailed +} + +/// Pure function: shuffle volumes for load balancing. +/// Takes a seed for deterministic testing. +pub fn shuffle_volumes(volumes: Vec, seed: u64) -> Vec { + use rand::seq::SliceRandom; + use rand::SeedableRng; + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut vols = volumes; + vols.shuffle(&mut rng); + vols +} + pub async fn get_key( State(state): State, Path(key): Path, ) -> Result { - use rand::seq::SliceRandom; - let record = state.db.get(&key).await?; if record.volumes.is_empty() { return Err(AppError::CorruptRecord { key }); } - // Shuffle volumes for load balancing - let mut volumes = record.volumes.clone(); - volumes.shuffle(&mut rand::thread_rng()); + // Shuffle for load balancing (random seed in production) + let seed = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0); + let volumes = shuffle_volumes(record.volumes, seed); - // Probe each volume until we find one that's reachable + // Probe volumes and collect results + let mut results = Vec::with_capacity(volumes.len()); for vol in &volumes { let url = format!("{vol}/{key}"); - match state.http.head(&url).send().await { - Ok(resp) if resp.status().is_success() => { - return Ok(( - StatusCode::FOUND, - [(axum::http::header::LOCATION, url)], - ) - .into_response()); - } + let healthy = match state.http.head(&url).send().await { + Ok(resp) if resp.status().is_success() => true, Ok(resp) => { tracing::warn!("volume {vol} returned {} for {key}", resp.status()); + false } Err(e) => { tracing::warn!("volume {vol} unreachable for {key}: {e}"); + false } + }; + results.push(healthy); + // Early exit on first healthy volume + if healthy { + break; } } - Err(AppError::AllVolumesUnreachable) + match first_healthy_volume(&key, &volumes, &results) { + ProbeResult::Found(url) => Ok(( + StatusCode::FOUND, + [(axum::http::header::LOCATION, url)], + ) + .into_response()), + ProbeResult::AllFailed => Err(AppError::AllVolumesUnreachable), + } } pub async fn put_key( @@ -204,6 +245,8 @@ pub async fn list_keys( #[cfg(test)] mod tests { + use super::*; + #[test] fn test_volumes_for_key_sufficient() { let volumes: Vec = (1..=3).map(|i| format!("http://vol{i}")).collect(); @@ -217,4 +260,71 @@ mod tests { let selected = crate::hasher::volumes_for_key("test-key", &volumes, 2); assert_eq!(selected.len(), 1); } + + #[test] + fn test_first_healthy_volume_finds_first() { + let volumes = vec!["http://vol1".into(), "http://vol2".into(), "http://vol3".into()]; + let results = vec![true, true, true]; + assert_eq!( + first_healthy_volume("key", &volumes, &results), + ProbeResult::Found("http://vol1/key".into()) + ); + } + + #[test] + fn test_first_healthy_volume_skips_unhealthy() { + let volumes = vec!["http://vol1".into(), "http://vol2".into(), "http://vol3".into()]; + let results = vec![false, false, true]; + assert_eq!( + first_healthy_volume("key", &volumes, &results), + ProbeResult::Found("http://vol3/key".into()) + ); + } + + #[test] + fn test_first_healthy_volume_all_failed() { + let volumes = vec!["http://vol1".into(), "http://vol2".into()]; + let results = vec![false, false]; + assert_eq!( + first_healthy_volume("key", &volumes, &results), + ProbeResult::AllFailed + ); + } + + #[test] + fn test_first_healthy_volume_early_exit() { + // Simulates early exit: only first two volumes were probed + let volumes = vec!["http://vol1".into(), "http://vol2".into(), "http://vol3".into()]; + let results = vec![false, true]; // Only 2 results because we stopped early + assert_eq!( + first_healthy_volume("key", &volumes, &results), + ProbeResult::Found("http://vol2/key".into()) + ); + } + + #[test] + fn test_shuffle_volumes_deterministic_with_seed() { + let volumes: Vec = (1..=5).map(|i| format!("http://vol{i}")).collect(); + let a = shuffle_volumes(volumes.clone(), 42); + let b = shuffle_volumes(volumes.clone(), 42); + assert_eq!(a, b, "same seed should produce same order"); + } + + #[test] + fn test_shuffle_volumes_different_seeds() { + let volumes: Vec = (1..=10).map(|i| format!("http://vol{i}")).collect(); + let a = shuffle_volumes(volumes.clone(), 1); + let b = shuffle_volumes(volumes.clone(), 2); + assert_ne!(a, b, "different seeds should produce different orders"); + } + + #[test] + fn test_shuffle_volumes_preserves_elements() { + let volumes: Vec = (1..=5).map(|i| format!("http://vol{i}")).collect(); + let mut shuffled = shuffle_volumes(volumes.clone(), 123); + shuffled.sort(); + let mut original = volumes; + original.sort(); + assert_eq!(shuffled, original); + } }