Simplify moar
This commit is contained in:
parent
7f3ec69cf6
commit
07490efc28
14 changed files with 261 additions and 1061 deletions
262
src/server.rs
262
src/server.rs
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue