mkv

server.rs (6862B)

---

 use axum::body::Bytes;
 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};
 
 #[derive(Clone)]
 pub struct AppState {
     pub db: db::Db,
     pub volumes: Arc<Vec<String>>,
     pub replicas: usize,
     pub voltimeout: Duration,
     pub http: reqwest::Client,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ProbeResult {
     Found(String),
     AllFailed,
 }
 
 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
 }
 
 pub fn shuffle_volumes(volumes: Vec<String>, seed: u64) -> Vec<String> {
     use rand::SeedableRng;
     use rand::seq::SliceRandom;
     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<AppState>,
     Path(key): Path<String>,
 ) -> Result<Response, AppError> {
     let record = state.db.get(&key).await?;
     if record.volumes.is_empty() {
         return Err(AppError::CorruptRecord { key });
     }
 
     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);
 
     let mut results = Vec::with_capacity(volumes.len());
     for vol in &volumes {
         let url = format!("{vol}/{key}");
         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());
                 false
             }
             Err(e) => {
                 tracing::warn!("volume {vol} unreachable for {key}: {e}");
                 false
             }
         };
         results.push(healthy);
         if healthy {
             break;
         }
     }
 
     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(
     State(state): State<AppState>,
     Path(key): Path<String>,
     body: Bytes,
 ) -> Result<Response, AppError> {
     let target_volumes = crate::hasher::volumes_for_key(&key, &state.volumes, state.replicas);
     if target_volumes.len() < state.replicas {
         return Err(AppError::InsufficientVolumes {
             need: state.replicas,
             have: target_volumes.len(),
         });
     }
 
     let mut handles = Vec::with_capacity(target_volumes.len());
     for vol in &target_volumes {
         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| {
                         VolumeError::Request {
                             url: url.clone(),
                             source: e,
                         }
                     })?;
                     if !resp.status().is_success() {
                         return Err(VolumeError::BadStatus {
                             url,
                             status: resp.status(),
                         });
                     }
                     Ok(())
                 }
             });
         handles.push(handle);
     }
 
     let mut failed = false;
     for handle in handles {
         match handle.await {
             Ok(Err(e)) => {
                 tracing::error!("{e}");
                 failed = true;
             }
             Err(e) => {
                 tracing::error!("volume write task failed: {e}");
                 failed = true;
             }
             Ok(Ok(())) => {}
         }
     }
 
     if failed {
         // Rollback: best-effort delete from volumes
         for vol in &target_volumes {
             let _ = state.http.delete(format!("{vol}/{key}")).send().await;
         }
         return Err(AppError::PartialWrite);
     }
 
     let size = Some(body.len() as i64);
     if let Err(e) = state
         .db
         .put(key.clone(), target_volumes.clone(), size)
         .await
     {
         for vol in &target_volumes {
             let _ = state.http.delete(format!("{vol}/{key}")).send().await;
         }
         return Err(e);
     }
     Ok(StatusCode::CREATED.into_response())
 }
 
 pub async fn delete_key(
     State(state): State<AppState>,
     Path(key): Path<String>,
 ) -> Result<Response, AppError> {
     let record = state.db.get(&key).await?;
 
     let mut handles = Vec::new();
     for vol in &record.volumes {
         let url = format!("{vol}/{key}");
         let handle = tokio::spawn({
             let client = state.http.clone();
             async move {
                 let resp = client
                     .delete(&url)
                     .send()
                     .await
                     .map_err(|e| VolumeError::Request {
                         url: url.clone(),
                         source: e,
                     })?;
                 if !resp.status().is_success() {
                     return Err(VolumeError::BadStatus {
                         url,
                         status: resp.status(),
                     });
                 }
                 Ok(())
             }
         });
         handles.push(handle);
     }
     for handle in handles {
         match handle.await {
             Ok(Err(e)) => tracing::error!("{e}"),
             Err(e) => tracing::error!("volume delete task failed: {e}"),
             Ok(Ok(())) => {}
         }
     }
 
     state.db.delete(key).await?;
     Ok(StatusCode::NO_CONTENT.into_response())
 }
 
 pub async fn head_key(
     State(state): State<AppState>,
     Path(key): Path<String>,
 ) -> Result<Response, AppError> {
     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(),
         );
     }
     Ok((StatusCode::OK, headers).into_response())
 }
 
 #[derive(serde::Deserialize)]
 pub struct ListQuery {
     #[serde(default)]
     pub prefix: String,
 }
 
 pub async fn list_keys(
     State(state): State<AppState>,
     Query(query): Query<ListQuery>,
 ) -> Result<Response, AppError> {
     let keys = state.db.list_keys(&query.prefix).await?;
     Ok((StatusCode::OK, keys.join("\n")).into_response())
 }