rebalance.rs (6218B)
1 use crate::Args; 2 use crate::db; 3 4 pub struct KeyMove { 5 pub key: String, 6 pub size: Option<i64>, 7 pub current_volumes: Vec<String>, 8 pub desired_volumes: Vec<String>, 9 pub to_add: Vec<String>, 10 pub to_remove: Vec<String>, 11 } 12 13 pub fn plan_rebalance( 14 records: &[db::Record], 15 volumes: &[String], 16 replication: usize, 17 ) -> Vec<KeyMove> { 18 let mut moves = Vec::new(); 19 for record in records { 20 let desired = crate::hasher::volumes_for_key(&record.key, volumes, replication); 21 let to_add: Vec<String> = desired 22 .iter() 23 .filter(|v| !record.volumes.contains(v)) 24 .cloned() 25 .collect(); 26 let to_remove: Vec<String> = record 27 .volumes 28 .iter() 29 .filter(|v| !desired.contains(v)) 30 .cloned() 31 .collect(); 32 33 if !to_add.is_empty() || !to_remove.is_empty() { 34 moves.push(KeyMove { 35 key: record.key.clone(), 36 size: record.size, 37 current_volumes: record.volumes.clone(), 38 desired_volumes: desired, 39 to_add, 40 to_remove, 41 }); 42 } 43 } 44 moves 45 } 46 47 pub async fn run(args: &Args, dry_run: bool) { 48 let db = db::Db::new(&args.db_path); 49 let records = db.all_records_sync().expect("failed to read records"); 50 let moves = plan_rebalance(&records, &args.volumes, args.replicas); 51 52 if moves.is_empty() { 53 eprintln!("Nothing to rebalance — all keys are already correctly placed."); 54 return; 55 } 56 57 let total_bytes: i64 = moves.iter().filter_map(|m| m.size).sum(); 58 eprintln!("{} keys to move ({} bytes)", moves.len(), total_bytes); 59 60 if dry_run { 61 for m in &moves { 62 eprintln!(" {} : add {:?}, remove {:?}", m.key, m.to_add, m.to_remove); 63 } 64 return; 65 } 66 67 let client = reqwest::Client::new(); 68 let mut moved = 0; 69 let mut errors = 0; 70 71 for m in &moves { 72 let Some(src) = m.current_volumes.first() else { 73 eprintln!(" SKIP {} : no source volume", m.key); 74 errors += 1; 75 continue; 76 }; 77 let mut copy_ok = true; 78 79 for dst in &m.to_add { 80 let src_url = format!("{src}/{}", m.key); 81 match client.get(&src_url).send().await { 82 Ok(resp) if resp.status().is_success() => { 83 let data = match resp.bytes().await { 84 Ok(b) => b, 85 Err(e) => { 86 eprintln!(" ERROR read body {} from {}: {}", m.key, src, e); 87 copy_ok = false; 88 errors += 1; 89 break; 90 } 91 }; 92 let dst_url = format!("{dst}/{}", m.key); 93 match client.put(&dst_url).body(data).send().await { 94 Ok(resp) if !resp.status().is_success() => { 95 eprintln!( 96 " ERROR copy {} to {}: status {}", 97 m.key, 98 dst, 99 resp.status() 100 ); 101 copy_ok = false; 102 errors += 1; 103 } 104 Err(e) => { 105 eprintln!(" ERROR copy {} to {}: {}", m.key, dst, e); 106 copy_ok = false; 107 errors += 1; 108 } 109 Ok(_) => {} 110 } 111 } 112 Ok(resp) => { 113 eprintln!( 114 " ERROR read {} from {}: status {}", 115 m.key, 116 src, 117 resp.status() 118 ); 119 copy_ok = false; 120 errors += 1; 121 } 122 Err(e) => { 123 eprintln!(" ERROR read {} from {}: {}", m.key, src, e); 124 copy_ok = false; 125 errors += 1; 126 } 127 } 128 } 129 130 if !copy_ok { 131 continue; 132 } 133 134 db.put(m.key.clone(), m.desired_volumes.clone(), m.size) 135 .await 136 .expect("failed to update index"); 137 138 for old in &m.to_remove { 139 let url = format!("{old}/{}", m.key); 140 if let Err(e) = client.delete(&url).send().await { 141 eprintln!(" WARN delete {} from {}: {}", m.key, old, e); 142 } 143 } 144 moved += 1; 145 } 146 147 eprintln!("Rebalanced {moved}/{} keys ({errors} errors)", moves.len()); 148 } 149 150 #[cfg(test)] 151 mod tests { 152 use super::*; 153 154 #[test] 155 fn test_plan_rebalance_no_change() { 156 let volumes: Vec<String> = (1..=3).map(|i| format!("http://vol{i}")).collect(); 157 let records: Vec<db::Record> = (0..100) 158 .map(|i| { 159 let key = format!("key-{i}"); 160 let vols = crate::hasher::volumes_for_key(&key, &volumes, 2); 161 db::Record { 162 key, 163 volumes: vols, 164 size: Some(100), 165 } 166 }) 167 .collect(); 168 169 let moves = plan_rebalance(&records, &volumes, 2); 170 assert!(moves.is_empty()); 171 } 172 173 #[test] 174 fn test_plan_rebalance_new_volume() { 175 let volumes3: Vec<String> = (1..=3).map(|i| format!("http://vol{i}")).collect(); 176 let records: Vec<db::Record> = (0..1000) 177 .map(|i| { 178 let key = format!("key-{i}"); 179 let vols = crate::hasher::volumes_for_key(&key, &volumes3, 2); 180 db::Record { 181 key, 182 volumes: vols, 183 size: Some(100), 184 } 185 }) 186 .collect(); 187 188 let volumes4: Vec<String> = (1..=4).map(|i| format!("http://vol{i}")).collect(); 189 let moves = plan_rebalance(&records, &volumes4, 2); 190 191 assert!(!moves.is_empty()); 192 assert!(moves.len() < 800, "too many moves: {}", moves.len()); 193 } 194 }