1use super::{LocationRecord, PoolMaintenanceReport, PoolMemberId, PoolMemberState, PoolStore};
2use hashtree_core::store::StoreError;
3use hashtree_core::types::Hash;
4
5impl PoolStore {
6 pub fn maintain(&self, max_items: usize) -> Result<PoolMaintenanceReport, StoreError> {
7 let mut report = PoolMaintenanceReport::default();
8 if max_items == 0 {
9 return Ok(report);
10 }
11 for (hash, location) in self.active_moves(max_items)? {
12 if report.examined >= max_items {
13 return Ok(report);
14 }
15 report.examined += 1;
16 let LocationRecord::Moving { source, target, .. } = location else {
17 continue;
18 };
19 match self.move_blob(source, target, hash) {
20 Ok(Some(bytes)) => {
21 report.moved += 1;
22 report.bytes_moved = report.bytes_moved.saturating_add(bytes);
23 }
24 Ok(None) => {}
25 Err(error) => report.failed.push(format!("{hash:?}: {error}")),
26 }
27 }
28 let draining = self
29 .read_manifest()?
30 .members
31 .into_iter()
32 .filter(|member| member.state == PoolMemberState::Draining)
33 .map(|member| member.id)
34 .collect::<Vec<_>>();
35
36 for source in draining {
37 let hashes = self.member_hashes(source, max_items.saturating_sub(report.examined))?;
38 for hash in hashes {
39 if report.examined >= max_items {
40 return Ok(report);
41 }
42 report.examined += 1;
43 match self.move_from_draining(source, hash) {
44 Ok(Some(bytes)) => {
45 report.moved += 1;
46 report.bytes_moved = report.bytes_moved.saturating_add(bytes);
47 }
48 Ok(None) => {}
49 Err(error) => report.failed.push(format!("{hash:?}: {error}")),
50 }
51 }
52 }
53 while report.examined < max_items {
54 let Some((source, target)) = self.rebalance_pair()? else {
55 break;
56 };
57 let hashes = self.member_hashes(source, max_items - report.examined)?;
58 if hashes.is_empty() {
59 break;
60 }
61 let mut progressed = false;
62 for hash in hashes {
63 if report.examined >= max_items {
64 break;
65 }
66 report.examined += 1;
67 let Some(location) = self.read_location(&hash)? else {
68 continue;
69 };
70 if !self.move_improves_balance(source, target, location.size())? {
71 continue;
72 }
73 match self.move_blob(source, target, hash) {
74 Ok(Some(bytes)) => {
75 report.moved += 1;
76 report.bytes_moved = report.bytes_moved.saturating_add(bytes);
77 progressed = true;
78 }
79 Ok(None) => {}
80 Err(error) => report.failed.push(format!("{hash:?}: {error}")),
81 }
82 }
83 if !progressed {
84 break;
85 }
86 }
87 Ok(report)
88 }
89
90 fn move_from_draining(
91 &self,
92 source: PoolMemberId,
93 hash: Hash,
94 ) -> Result<Option<u64>, StoreError> {
95 let Some(location) = self.read_location(&hash)? else {
96 return Ok(None);
97 };
98 let (target, size, moving) = match location {
99 LocationRecord::Pending { member, size } | LocationRecord::Stored { member, size }
100 if member == source =>
101 {
102 let target = self.choose_write_member(size, Some(source))?;
103 return self.move_blob(source, target, hash);
104 }
105 LocationRecord::Moving {
106 source: moving_source,
107 target,
108 size,
109 } if moving_source == source => (target, size, true),
110 _ => return Ok(None),
111 };
112
113 self.move_blob_inner(source, target, hash, location, size, moving)
114 }
115
116 pub(super) fn move_blob(
117 &self,
118 source: PoolMemberId,
119 target: PoolMemberId,
120 hash: Hash,
121 ) -> Result<Option<u64>, StoreError> {
122 let Some(location) = self.read_location(&hash)? else {
123 return Ok(None);
124 };
125 let (actual_target, size, moving) = match location {
126 LocationRecord::Pending { member, size } | LocationRecord::Stored { member, size }
127 if member == source =>
128 {
129 (target, size, false)
130 }
131 LocationRecord::Moving {
132 source: moving_source,
133 target,
134 size,
135 } if moving_source == source => (target, size, true),
136 _ => return Ok(None),
137 };
138 self.move_blob_inner(source, actual_target, hash, location, size, moving)
139 }
140
141 fn move_blob_inner(
142 &self,
143 source: PoolMemberId,
144 target: PoolMemberId,
145 hash: Hash,
146 location: LocationRecord,
147 size: u64,
148 moving: bool,
149 ) -> Result<Option<u64>, StoreError> {
150 let target_store = self.get_member(target)?;
151 let source_store = match self.get_member(source) {
152 Ok(store) => store,
153 Err(_) if moving => {
154 target_store.verify_blob_streaming(
155 &hash,
156 size,
157 self.temperature_config.copy_chunk_bytes,
158 )?;
159 self.complete_move(hash, source, target, size)?;
160 return Ok(Some(size));
161 }
162 Err(error) => return Err(error),
163 };
164 let source_size = match source_store.blob_size_sync(&hash) {
165 Ok(size) => size,
166 Err(_) if moving => {
167 target_store.verify_blob_streaming(
168 &hash,
169 size,
170 self.temperature_config.copy_chunk_bytes,
171 )?;
172 self.complete_move(hash, source, target, size)?;
173 let _ = self.delete_member_blob(source, &source_store, &hash);
174 return Ok(Some(size));
175 }
176 Err(error) => return Err(error),
177 };
178 if source_size.is_none() && moving {
179 target_store.verify_blob_streaming(
180 &hash,
181 size,
182 self.temperature_config.copy_chunk_bytes,
183 )?;
184 self.complete_move(hash, source, target, size)?;
185 let _ = self.delete_member_blob(source, &source_store, &hash);
186 return Ok(Some(size));
187 }
188 let source_size = source_size.ok_or_else(|| {
189 StoreError::Other(format!(
190 "draining source {source} does not contain the blob"
191 ))
192 })?;
193 if source_size != size {
194 return Err(StoreError::Other(format!(
195 "pool source {source} size mismatch: catalog={size}, member={source_size}"
196 )));
197 }
198
199 let moving = LocationRecord::Moving {
200 source,
201 target,
202 size,
203 };
204 if !self.begin_move_record(hash, location, moving)? {
205 return Ok(None);
206 }
207 let source_gate = self.member_gate(source, false)?;
208 let target_gate = self.member_gate(target, true)?;
209 let _source_permit = source_gate.acquire()?;
210 let _target_permit = target_gate.acquire()?;
211 source_store.copy_blob_to_sync(
212 &target_store,
213 &hash,
214 size,
215 self.temperature_config.copy_chunk_bytes,
216 )?;
217 self.complete_move(hash, source, target, size)?;
218 let _ = self.delete_member_blob(source, &source_store, &hash);
219 Ok(Some(size))
220 }
221
222 fn complete_move(
223 &self,
224 hash: Hash,
225 source: PoolMemberId,
226 target: PoolMemberId,
227 size: u64,
228 ) -> Result<(), StoreError> {
229 let current = self.read_location(&hash)?;
230 match current {
231 Some(LocationRecord::Moving { .. }) | Some(LocationRecord::Stored { .. }) => {
232 self.finish_move_record(hash, source, target, size)
233 }
234 other => Err(StoreError::Other(format!(
235 "pool location disappeared while moving {hash:?}: {other:?}"
236 ))),
237 }
238 }
239
240 fn rebalance_pair(&self) -> Result<Option<(PoolMemberId, PoolMemberId)>, StoreError> {
241 let manifest = self.read_manifest()?;
242 let mut members = Vec::new();
243 for member in manifest
244 .members
245 .into_iter()
246 .filter(|member| member.state == PoolMemberState::Active)
247 {
248 let Ok(store) = self.get_member(member.id) else {
249 continue;
250 };
251 let stats = store.stats()?;
252 members.push((member.id, stats.total_bytes, member.config.capacity_bytes));
253 }
254 if members.len() < 2 {
255 return Ok(None);
256 }
257 let total_bytes = members
258 .iter()
259 .map(|(_, bytes, _)| *bytes)
260 .fold(0u64, u64::saturating_add);
261 let total_capacity = members
262 .iter()
263 .map(|(_, _, capacity)| *capacity)
264 .fold(0u64, u64::saturating_add);
265 if total_bytes == 0 || total_capacity == 0 {
266 return Ok(None);
267 }
268 let deviation = |bytes: u64, capacity: u64| -> i128 {
269 i128::from(bytes) * i128::from(total_capacity)
270 - i128::from(total_bytes) * i128::from(capacity)
271 };
272 let source = members
273 .iter()
274 .max_by_key(|(_, bytes, capacity)| deviation(*bytes, *capacity))
275 .copied();
276 let target = members
277 .iter()
278 .min_by_key(|(_, bytes, capacity)| deviation(*bytes, *capacity))
279 .copied();
280 match (source, target) {
281 (
282 Some((source, source_bytes, source_capacity)),
283 Some((target, target_bytes, target_capacity)),
284 ) if source != target
285 && deviation(source_bytes, source_capacity) > 0
286 && deviation(target_bytes, target_capacity) < 0 =>
287 {
288 Ok(Some((source, target)))
289 }
290 _ => Ok(None),
291 }
292 }
293
294 fn move_improves_balance(
295 &self,
296 source: PoolMemberId,
297 target: PoolMemberId,
298 blob_bytes: u64,
299 ) -> Result<bool, StoreError> {
300 let members = self.members()?;
301 let active = members
302 .iter()
303 .filter(|member| member.state == PoolMemberState::Active && member.available)
304 .collect::<Vec<_>>();
305 let Some(source_status) = active.iter().find(|member| member.id == source) else {
306 return Ok(false);
307 };
308 let Some(target_status) = active.iter().find(|member| member.id == target) else {
309 return Ok(false);
310 };
311 if blob_bytes > source_status.logical_bytes
312 || target_status.logical_bytes.saturating_add(blob_bytes) > target_status.capacity_bytes
313 {
314 return Ok(false);
315 }
316 let total_bytes = active
317 .iter()
318 .map(|member| member.logical_bytes)
319 .fold(0u64, u64::saturating_add);
320 let total_capacity = active
321 .iter()
322 .map(|member| member.capacity_bytes)
323 .fold(0u64, u64::saturating_add);
324 if total_capacity == 0 {
325 return Ok(false);
326 }
327 let deviation = |bytes: u64, capacity: u64| -> i128 {
328 i128::from(bytes) * i128::from(total_capacity)
329 - i128::from(total_bytes) * i128::from(capacity)
330 };
331 let before = deviation(source_status.logical_bytes, source_status.capacity_bytes).abs()
332 + deviation(target_status.logical_bytes, target_status.capacity_bytes).abs();
333 let after = deviation(
334 source_status.logical_bytes - blob_bytes,
335 source_status.capacity_bytes,
336 )
337 .abs()
338 + deviation(
339 target_status.logical_bytes.saturating_add(blob_bytes),
340 target_status.capacity_bytes,
341 )
342 .abs();
343 Ok(after < before)
344 }
345}