From 9d49f9e5e660a719a00d64845d3ca26de8aa181d Mon Sep 17 00:00:00 2001 From: DevITWay | Pavel Volkov Date: Sun, 5 Apr 2026 22:48:39 +0300 Subject: [PATCH] feat: add total_size_bytes to /health endpoint (#42) (#93) Add StorageBackend::total_size() trait method with efficient implementations: - LocalStorage: spawn_blocking dir walk (no async overhead) - S3Storage: list + stat iteration Health endpoint now returns storage.total_size_bytes in JSON response. Closes #42 --- nora-registry/src/health.rs | 38 +++++++++++++++++++++ nora-registry/src/storage/local.rs | 55 ++++++++++++++++++++++++++++++ nora-registry/src/storage/mod.rs | 6 ++++ nora-registry/src/storage/s3.rs | 11 ++++++ 4 files changed, 110 insertions(+) diff --git a/nora-registry/src/health.rs b/nora-registry/src/health.rs index dca4b4b..2adeba6 100644 --- a/nora-registry/src/health.rs +++ b/nora-registry/src/health.rs @@ -21,6 +21,7 @@ pub struct StorageHealth { pub backend: String, pub reachable: bool, pub endpoint: String, + pub total_size_bytes: u64, } #[derive(Serialize)] @@ -40,6 +41,7 @@ pub fn routes() -> Router> { async fn health_check(State(state): State>) -> (StatusCode, Json) { let storage_reachable = check_storage_reachable(&state).await; + let total_size = state.storage.total_size().await; let status = if storage_reachable { "healthy" @@ -60,6 +62,7 @@ async fn health_check(State(state): State>) -> (StatusCode, Json state.config.storage.s3_url.clone(), _ => state.config.storage.path.clone(), }, + total_size_bytes: total_size, }, registries: RegistriesHealth { docker: "ok".to_string(), @@ -117,6 +120,41 @@ mod tests { assert!(json.get("version").is_some()); } + #[tokio::test] + async fn test_health_json_has_storage_size() { + let ctx = create_test_context(); + + // Put some data to have non-zero size + ctx.state + .storage + .put("test/artifact", b"hello world") + .await + .unwrap(); + + let response = send(&ctx.app, Method::GET, "/health", "").await; + assert_eq!(response.status(), StatusCode::OK); + let body = body_bytes(response).await; + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + + let storage = json.get("storage").unwrap(); + let size = storage.get("total_size_bytes").unwrap().as_u64().unwrap(); + assert!( + size > 0, + "total_size_bytes should be > 0 after storing data" + ); + } + + #[tokio::test] + async fn test_health_empty_storage_size_zero() { + let ctx = create_test_context(); + let response = send(&ctx.app, Method::GET, "/health", "").await; + let body = body_bytes(response).await; + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + + let size = json["storage"]["total_size_bytes"].as_u64().unwrap(); + assert_eq!(size, 0, "empty storage should report 0 bytes"); + } + #[tokio::test] async fn test_ready_returns_200() { let ctx = create_test_context(); diff --git a/nora-registry/src/storage/local.rs b/nora-registry/src/storage/local.rs index e6c88bb..1e13819 100644 --- a/nora-registry/src/storage/local.rs +++ b/nora-registry/src/storage/local.rs @@ -138,6 +138,29 @@ impl StorageBackend for LocalStorage { fs::create_dir_all(&self.base_path).await.is_ok() } + async fn total_size(&self) -> u64 { + let base = self.base_path.clone(); + tokio::task::spawn_blocking(move || { + fn dir_size(path: &std::path::Path) -> u64 { + let mut total = 0u64; + if let Ok(entries) = std::fs::read_dir(path) { + for entry in entries.flatten() { + let path = entry.path(); + if path.is_file() { + total += entry.metadata().map(|m| m.len()).unwrap_or(0); + } else if path.is_dir() { + total += dir_size(&path); + } + } + } + total + } + dir_size(&base) + }) + .await + .unwrap_or(0) + } + fn backend_name(&self) -> &'static str { "local" } @@ -338,6 +361,38 @@ mod tests { assert_eq!(&*data, &vec![1u8; 4096]); } + #[tokio::test] + async fn test_total_size_empty() { + let temp_dir = TempDir::new().unwrap(); + let storage = LocalStorage::new(temp_dir.path().to_str().unwrap()); + assert_eq!(storage.total_size().await, 0); + } + + #[tokio::test] + async fn test_total_size_with_files() { + let temp_dir = TempDir::new().unwrap(); + let storage = LocalStorage::new(temp_dir.path().to_str().unwrap()); + + storage.put("a/file1", b"hello").await.unwrap(); // 5 bytes + storage.put("b/file2", b"world!").await.unwrap(); // 6 bytes + + let size = storage.total_size().await; + assert_eq!(size, 11); + } + + #[tokio::test] + async fn test_total_size_after_delete() { + let temp_dir = TempDir::new().unwrap(); + let storage = LocalStorage::new(temp_dir.path().to_str().unwrap()); + + storage.put("file1", b"12345").await.unwrap(); + storage.put("file2", b"67890").await.unwrap(); + assert_eq!(storage.total_size().await, 10); + + storage.delete("file1").await.unwrap(); + assert_eq!(storage.total_size().await, 5); + } + #[tokio::test(flavor = "multi_thread")] async fn test_concurrent_deletes_same_key() { let temp_dir = TempDir::new().unwrap(); diff --git a/nora-registry/src/storage/mod.rs b/nora-registry/src/storage/mod.rs index a526e80..fb30c04 100644 --- a/nora-registry/src/storage/mod.rs +++ b/nora-registry/src/storage/mod.rs @@ -46,6 +46,8 @@ pub trait StorageBackend: Send + Sync { async fn list(&self, prefix: &str) -> Vec; async fn stat(&self, key: &str) -> Option; async fn health_check(&self) -> bool; + /// Total size of all stored artifacts in bytes + async fn total_size(&self) -> u64; fn backend_name(&self) -> &'static str; } @@ -110,6 +112,10 @@ impl Storage { self.inner.health_check().await } + pub async fn total_size(&self) -> u64 { + self.inner.total_size().await + } + pub fn backend_name(&self) -> &'static str { self.inner.backend_name() } diff --git a/nora-registry/src/storage/s3.rs b/nora-registry/src/storage/s3.rs index 31434e6..47c4e27 100644 --- a/nora-registry/src/storage/s3.rs +++ b/nora-registry/src/storage/s3.rs @@ -382,6 +382,17 @@ impl StorageBackend for S3Storage { } } + async fn total_size(&self) -> u64 { + let keys = self.list("").await; + let mut total = 0u64; + for key in &keys { + if let Some(meta) = self.stat(key).await { + total += meta.size; + } + } + total + } + fn backend_name(&self) -> &'static str { "s3" }