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