Simplify moar

This commit is contained in:
Silas Brack 2026-03-07 13:23:38 +01:00
parent 7f3ec69cf6
commit 07490efc28
14 changed files with 261 additions and 1061 deletions

View file

@ -2,201 +2,113 @@ use axum::body::Bytes;
use axum::extract::{Path, Query, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::config::Config;
use crate::db;
use crate::error::AppError;
use crate::hasher::Ring;
use crate::volume::VolumeClient;
#[derive(Clone)]
pub struct AppState {
pub writer: db::WriterHandle,
pub reads: db::ReadPool,
pub ring: Arc<RwLock<Ring>>,
pub volume_client: VolumeClient,
pub healthy_volumes: Arc<RwLock<HashSet<String>>>,
pub config: Arc<Config>,
pub db: db::Db,
pub volumes: Arc<Vec<String>>,
pub replicas: usize,
pub http: reqwest::Client,
}
// --- Pure decision functions ---
/// Pick the first volume from the record that appears in the healthy set.
fn pick_healthy_volume<'a>(
record_volumes: &'a [String],
healthy: &HashSet<String>,
) -> Option<&'a str> {
record_volumes
.iter()
.find(|v| healthy.contains(v.as_str()))
.map(|v| v.as_str())
}
/// Select target volumes for a key, ensuring we have enough for replication.
fn select_volumes(ring: &Ring, key: &str, replication: usize) -> Result<Vec<String>, AppError> {
let target_volumes = ring.get_volumes(key, replication);
if target_volumes.len() < replication {
return Err(AppError::VolumeError(format!(
"need {replication} volumes but only {} available",
target_volumes.len()
)));
}
Ok(target_volumes)
}
/// Evaluate fan-out results. Returns the list of volumes that succeeded,
/// or the error messages if any failed.
fn evaluate_fanout(
results: Vec<Result<(), String>>,
volumes: &[String],
) -> Result<Vec<String>, Vec<String>> {
let mut succeeded = Vec::new();
let mut errors = Vec::new();
for (i, result) in results.into_iter().enumerate() {
match result {
Ok(()) => succeeded.push(volumes[i].clone()),
Err(e) => errors.push(e),
}
}
if errors.is_empty() {
Ok(succeeded)
} else {
Err(errors)
}
}
// --- Handlers ---
/// GET /:key — look up key, redirect to a healthy volume
pub async fn get_key(
State(state): State<AppState>,
Path(key): Path<String>,
) -> Result<Response, AppError> {
let record = state
.reads
.query({
let key = key.clone();
move |conn| db::get(conn, &key)
})
.await?;
let healthy = state.healthy_volumes.read().await;
let vol = pick_healthy_volume(&record.volumes, &healthy).ok_or(AppError::NoHealthyVolume)?;
let record = state.db.get(&key).await?;
let vol = &record.volumes[0];
let location = format!("{vol}/{key}");
Ok((
StatusCode::FOUND,
[(axum::http::header::LOCATION, location)],
)
.into_response())
Ok((StatusCode::FOUND, [(axum::http::header::LOCATION, location)]).into_response())
}
/// PUT /:key — store blob on volumes, record in index
pub async fn put_key(
State(state): State<AppState>,
Path(key): Path<String>,
body: Bytes,
) -> Result<Response, AppError> {
let target_volumes = {
let ring = state.ring.read().await;
select_volumes(&ring, &key, state.config.server.replication_factor)?
};
let target_volumes = crate::hasher::volumes_for_key(&key, &state.volumes, state.replicas);
if target_volumes.len() < state.replicas {
return Err(AppError::VolumeError(format!(
"need {} volumes but only {} available",
state.replicas,
target_volumes.len()
)));
}
// Fan out PUTs to all target volumes concurrently
let mut handles = Vec::with_capacity(target_volumes.len());
for vol in &target_volumes {
let client = state.volume_client.clone();
let vol = vol.clone();
let key = key.clone();
let data = body.clone();
handles.push(tokio::spawn(async move {
client.put(&vol, &key, data).await
}));
let url = format!("{vol}/{key}");
let handle = tokio::spawn({
let client = state.http.clone();
let data = body.clone();
async move {
let resp = client.put(&url).body(data).send().await.map_err(|e| format!("PUT {url}: {e}"))?;
if !resp.status().is_success() {
return Err(format!("PUT {url}: status {}", resp.status()));
}
Ok(())
}
});
handles.push(handle);
}
let mut results = Vec::with_capacity(handles.len());
let mut all_ok = true;
for handle in handles {
results.push(handle.await.unwrap());
if let Err(e) = handle.await.unwrap() {
tracing::error!("PUT to volume failed: {e}");
all_ok = false;
}
}
match evaluate_fanout(results, &target_volumes) {
Ok(succeeded_volumes) => {
let size = Some(body.len() as i64);
state.writer.put(key, succeeded_volumes, size).await?;
Ok(StatusCode::CREATED.into_response())
}
Err(errors) => {
for e in &errors {
tracing::error!("PUT to volume failed: {e}");
}
// Rollback: best-effort delete from any volumes that may have succeeded
for vol in &target_volumes {
let _ = state.volume_client.delete(vol, &key).await;
}
Err(AppError::VolumeError(
"not all volume writes succeeded".into(),
))
if !all_ok {
// Rollback: best-effort delete from volumes
for vol in &target_volumes {
let _ = state.http.delete(format!("{vol}/{key}")).send().await;
}
return Err(AppError::VolumeError("not all volume writes succeeded".into()));
}
let size = Some(body.len() as i64);
state.db.put(key, target_volumes, size).await?;
Ok(StatusCode::CREATED.into_response())
}
/// DELETE /:key — remove from volumes and index
pub async fn delete_key(
State(state): State<AppState>,
Path(key): Path<String>,
) -> Result<Response, AppError> {
let record = state
.reads
.query({
let key = key.clone();
move |conn| db::get(conn, &key)
})
.await?;
let record = state.db.get(&key).await?;
// Fan out DELETEs concurrently
let mut handles = Vec::new();
for vol in &record.volumes {
let client = state.volume_client.clone();
let vol = vol.clone();
let key = key.clone();
handles.push(tokio::spawn(
async move { client.delete(&vol, &key).await },
));
let url = format!("{vol}/{key}");
let client = state.http.clone();
handles.push(tokio::spawn(async move { client.delete(&url).send().await }));
}
for handle in handles {
if let Err(e) = handle.await.unwrap() {
tracing::error!("DELETE from volume failed: {e}");
}
}
// Remove from index regardless of volume DELETE results
state.writer.delete(key).await?;
state.db.delete(key).await?;
Ok(StatusCode::NO_CONTENT.into_response())
}
/// HEAD /:key — check if key exists, return size
pub async fn head_key(
State(state): State<AppState>,
Path(key): Path<String>,
) -> Result<Response, AppError> {
let record = state
.reads
.query(move |conn| db::get(conn, &key))
.await?;
let record = state.db.get(&key).await?;
let mut headers = HeaderMap::new();
if let Some(size) = record.size {
headers.insert(
axum::http::header::CONTENT_LENGTH,
size.to_string().parse().unwrap(),
);
headers.insert(axum::http::header::CONTENT_LENGTH, size.to_string().parse().unwrap());
}
Ok((StatusCode::OK, headers).into_response())
}
@ -206,89 +118,27 @@ pub struct ListQuery {
pub prefix: String,
}
/// GET / — list keys with optional prefix filter
pub async fn list_keys(
State(state): State<AppState>,
Query(query): Query<ListQuery>,
) -> Result<Response, AppError> {
let keys = state
.reads
.query(move |conn| db::list_keys(conn, &query.prefix))
.await?;
let body = keys.join("\n");
Ok((StatusCode::OK, body).into_response())
let keys = state.db.list_keys(&query.prefix).await?;
Ok((StatusCode::OK, keys.join("\n")).into_response())
}
// --- Tests for pure functions ---
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pick_healthy_volume_first_match() {
let volumes = vec!["http://vol1".into(), "http://vol2".into(), "http://vol3".into()];
let healthy: HashSet<String> = ["http://vol2".into(), "http://vol3".into()].into();
assert_eq!(pick_healthy_volume(&volumes, &healthy), Some("http://vol2"));
}
#[test]
fn test_pick_healthy_volume_none_healthy() {
let volumes = vec!["http://vol1".into(), "http://vol2".into()];
let healthy: HashSet<String> = HashSet::new();
assert_eq!(pick_healthy_volume(&volumes, &healthy), None);
}
#[test]
fn test_pick_healthy_volume_all_healthy() {
let volumes = vec!["http://vol1".into(), "http://vol2".into()];
let healthy: HashSet<String> = ["http://vol1".into(), "http://vol2".into()].into();
assert_eq!(pick_healthy_volume(&volumes, &healthy), Some("http://vol1"));
}
#[test]
fn test_select_volumes_sufficient() {
fn test_volumes_for_key_sufficient() {
let volumes: Vec<String> = (1..=3).map(|i| format!("http://vol{i}")).collect();
let ring = Ring::new(&volumes, 50);
let selected = select_volumes(&ring, "test-key", 2).unwrap();
let selected = crate::hasher::volumes_for_key("test-key", &volumes, 2);
assert_eq!(selected.len(), 2);
}
#[test]
fn test_select_volumes_insufficient() {
fn test_volumes_for_key_insufficient() {
let volumes: Vec<String> = vec!["http://vol1".into()];
let ring = Ring::new(&volumes, 50);
assert!(select_volumes(&ring, "test-key", 2).is_err());
}
#[test]
fn test_evaluate_fanout_all_ok() {
let volumes = vec!["http://vol1".into(), "http://vol2".into()];
let results = vec![Ok(()), Ok(())];
assert_eq!(evaluate_fanout(results, &volumes).unwrap(), volumes);
}
#[test]
fn test_evaluate_fanout_partial_failure() {
let volumes = vec!["http://vol1".into(), "http://vol2".into()];
let results = vec![Ok(()), Err("connection refused".into())];
let errors = evaluate_fanout(results, &volumes).unwrap_err();
assert_eq!(errors, vec!["connection refused"]);
}
#[test]
fn test_evaluate_fanout_all_fail() {
let volumes = vec!["http://vol1".into(), "http://vol2".into()];
let results = vec![Err("err1".into()), Err("err2".into())];
assert_eq!(evaluate_fanout(results, &volumes).unwrap_err().len(), 2);
let selected = crate::hasher::volumes_for_key("test-key", &volumes, 2);
assert_eq!(selected.len(), 1);
}
}