rebuild.rs (3678B)
1 use std::collections::HashMap; 2 3 use crate::Args; 4 use crate::db; 5 6 #[derive(serde::Deserialize)] 7 struct NginxEntry { 8 name: String, 9 #[serde(rename = "type")] 10 entry_type: String, 11 #[serde(default)] 12 size: Option<i64>, 13 } 14 15 /// If a key has different sizes across volumes, takes the max. 16 pub fn merge_volume_scans( 17 scans: &[(String, Vec<(String, i64)>)], 18 ) -> HashMap<String, (Vec<String>, i64)> { 19 let mut index: HashMap<String, (Vec<String>, i64)> = HashMap::new(); 20 for (vol_url, keys) in scans { 21 for (key, size) in keys { 22 let entry = index 23 .entry(key.clone()) 24 .or_insert_with(|| (Vec::new(), *size)); 25 entry.0.push(vol_url.clone()); 26 if *size > entry.1 { 27 entry.1 = *size; 28 } 29 } 30 } 31 index 32 } 33 34 async fn list_volume_keys(volume_url: &str) -> Result<Vec<(String, i64)>, String> { 35 let http = reqwest::Client::new(); 36 let mut keys = Vec::new(); 37 let mut dirs = vec![String::new()]; 38 39 while let Some(prefix) = dirs.pop() { 40 let url = format!("{volume_url}/{prefix}"); 41 let resp = http 42 .get(&url) 43 .send() 44 .await 45 .map_err(|e| format!("GET {url}: {e}"))?; 46 if !resp.status().is_success() { 47 return Err(format!("GET {url}: status {}", resp.status())); 48 } 49 let entries: Vec<NginxEntry> = 50 resp.json().await.map_err(|e| format!("parse {url}: {e}"))?; 51 for entry in entries { 52 let full_path = if prefix.is_empty() { 53 entry.name.clone() 54 } else { 55 format!("{prefix}{}", entry.name) 56 }; 57 match entry.entry_type.as_str() { 58 "directory" => dirs.push(format!("{full_path}/")), 59 "file" => keys.push((full_path, entry.size.unwrap_or(0))), 60 _ => {} 61 } 62 } 63 } 64 Ok(keys) 65 } 66 67 pub async fn run(args: &Args) { 68 let db_path = &args.db_path; 69 70 if let Some(parent) = std::path::Path::new(db_path).parent() { 71 let _ = std::fs::create_dir_all(parent); 72 } 73 74 let _ = std::fs::remove_file(db_path); 75 let _ = std::fs::remove_file(format!("{db_path}-wal")); 76 let _ = std::fs::remove_file(format!("{db_path}-shm")); 77 78 let db = db::Db::new(db_path); 79 80 let mut scans = Vec::new(); 81 for vol_url in &args.volumes { 82 eprintln!("Scanning {vol_url}..."); 83 match list_volume_keys(vol_url).await { 84 Ok(keys) => { 85 eprintln!(" Found {} keys", keys.len()); 86 scans.push((vol_url.clone(), keys)); 87 } 88 Err(e) => eprintln!(" Error scanning {vol_url}: {e}"), 89 } 90 } 91 92 let index = merge_volume_scans(&scans); 93 94 let records: Vec<_> = index 95 .into_iter() 96 .map(|(k, (v, s))| (k, v, Some(s))) 97 .collect(); 98 let count = records.len(); 99 db.bulk_put(records).await.expect("bulk_put failed"); 100 eprintln!("Rebuilt index with {count} keys"); 101 } 102 103 #[cfg(test)] 104 mod tests { 105 use super::*; 106 107 #[test] 108 fn test_merge_takes_max_size() { 109 // Can happen due to incomplete writes or corruption 110 let scans = vec![ 111 ("http://vol1".to_string(), vec![("key".to_string(), 50)]), 112 ("http://vol2".to_string(), vec![("key".to_string(), 200)]), 113 ("http://vol3".to_string(), vec![("key".to_string(), 100)]), 114 ]; 115 let index = merge_volume_scans(&scans); 116 let (volumes, size) = index.get("key").unwrap(); 117 assert_eq!(volumes.len(), 3); 118 assert_eq!(*size, 200, "should take maximum size across volumes"); 119 } 120 }