Add unit test for failover
This commit is contained in:
parent
1fc59674f5
commit
1d3b9dddf5
1 changed files with 125 additions and 15 deletions
140
src/server.rs
140
src/server.rs
|
|
@ -15,42 +15,83 @@ pub struct AppState {
|
||||||
pub http: reqwest::Client,
|
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<String>, seed: u64) -> Vec<String> {
|
||||||
|
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(
|
pub async fn get_key(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
Path(key): Path<String>,
|
Path(key): Path<String>,
|
||||||
) -> Result<Response, AppError> {
|
) -> Result<Response, AppError> {
|
||||||
use rand::seq::SliceRandom;
|
|
||||||
|
|
||||||
let record = state.db.get(&key).await?;
|
let record = state.db.get(&key).await?;
|
||||||
if record.volumes.is_empty() {
|
if record.volumes.is_empty() {
|
||||||
return Err(AppError::CorruptRecord { key });
|
return Err(AppError::CorruptRecord { key });
|
||||||
}
|
}
|
||||||
|
|
||||||
// Shuffle volumes for load balancing
|
// Shuffle for load balancing (random seed in production)
|
||||||
let mut volumes = record.volumes.clone();
|
let seed = std::time::SystemTime::now()
|
||||||
volumes.shuffle(&mut rand::thread_rng());
|
.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 {
|
for vol in &volumes {
|
||||||
let url = format!("{vol}/{key}");
|
let url = format!("{vol}/{key}");
|
||||||
match state.http.head(&url).send().await {
|
let healthy = match state.http.head(&url).send().await {
|
||||||
Ok(resp) if resp.status().is_success() => {
|
Ok(resp) if resp.status().is_success() => true,
|
||||||
return Ok((
|
|
||||||
StatusCode::FOUND,
|
|
||||||
[(axum::http::header::LOCATION, url)],
|
|
||||||
)
|
|
||||||
.into_response());
|
|
||||||
}
|
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
tracing::warn!("volume {vol} returned {} for {key}", resp.status());
|
tracing::warn!("volume {vol} returned {} for {key}", resp.status());
|
||||||
|
false
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!("volume {vol} unreachable for {key}: {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(
|
pub async fn put_key(
|
||||||
|
|
@ -204,6 +245,8 @@ pub async fn list_keys(
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_volumes_for_key_sufficient() {
|
fn test_volumes_for_key_sufficient() {
|
||||||
let volumes: Vec<String> = (1..=3).map(|i| format!("http://vol{i}")).collect();
|
let volumes: Vec<String> = (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);
|
let selected = crate::hasher::volumes_for_key("test-key", &volumes, 2);
|
||||||
assert_eq!(selected.len(), 1);
|
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<String> = (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<String> = (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<String> = (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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue