1use super::temperature::{TemperatureCandidate, HEAT_UNIT};
2use super::{
3 LocationRecord, PoolMemberId, PoolMemberState, PoolMemberStatus, PoolStore,
4 PoolTemperatureReport,
5};
6use hashtree_core::store::StoreError;
7use std::collections::HashSet;
8
9impl PoolStore {
10 pub fn balance_temperature(&self) -> Result<PoolTemperatureReport, StoreError> {
11 self.balance_temperature_at(super::unix_timestamp_now())
12 }
13
14 pub(super) fn balance_temperature_at(
15 &self,
16 now: u64,
17 ) -> Result<PoolTemperatureReport, StoreError> {
18 if !self.temperature_config.enabled {
19 return Ok(PoolTemperatureReport::default());
20 }
21 let Ok(_cycle) = self.temperature_cycle.try_lock() else {
22 return Ok(PoolTemperatureReport::default());
23 };
24 let sampled = self.flush_sampled_accesses(now)?;
25 let mut report = PoolTemperatureReport {
26 sampled_accesses_flushed: sampled.len(),
27 ..PoolTemperatureReport::default()
28 };
29 if !self.try_acquire_temperature_lease(now)? {
30 return Ok(report);
31 }
32 report.lease_acquired = true;
33 let _lease = super::temperature_worker::TemperatureLeaseHeartbeat::start(
34 std::sync::Arc::downgrade(&self.inner),
35 self.temperature_config.lease_duration,
36 )?;
37 self.balance_temperature_owned(now, sampled, report)
38 }
39
40 fn balance_temperature_owned(
41 &self,
42 now: u64,
43 sampled: Vec<TemperatureCandidate>,
44 mut report: PoolTemperatureReport,
45 ) -> Result<PoolTemperatureReport, StoreError> {
46 let scanned = self.scan_temperature_candidates(now)?;
47 report.scanned = scanned.len();
48 {
49 let mut temperature = self
50 .temperature
51 .lock()
52 .map_err(|_| StoreError::Other("pool temperature lock poisoned".into()))?;
53 for candidate in sampled.into_iter().chain(scanned) {
54 if candidate.heat >= HEAT_UNIT {
55 temperature.hot.upsert(candidate);
56 } else {
57 temperature.cold.upsert(candidate);
58 }
59 }
60 report.candidates = temperature.hot.len().saturating_add(temperature.cold.len());
61 }
62
63 if self.temperature_foreground_busy()? {
64 report.throttled = true;
65 return Ok(report);
66 }
67
68 let mut members = self.temperature_members()?;
69 let pressured = members
70 .iter()
71 .filter(|member| {
72 member.state == PoolMemberState::Active
73 && member.available
74 && fill_percent(member.logical_bytes, member.capacity_bytes)
75 >= member.temperature_high_watermark_percent
76 })
77 .map(|member| member.id)
78 .collect::<HashSet<_>>();
79 let mut remaining_bytes = self.temperature_config.max_bytes_per_cycle;
80 let mut attempted_moves = 0usize;
81 for (hash, location) in self.active_moves(
82 self.temperature_config
83 .max_moves_per_cycle
84 .saturating_sub(report.moved),
85 )? {
86 if attempted_moves >= self.temperature_config.max_moves_per_cycle
87 || remaining_bytes < location.size()
88 {
89 break;
90 }
91 let LocationRecord::Moving { source, target, .. } = location else {
92 continue;
93 };
94 attempted_moves += 1;
95 report.attempted_moves += 1;
96 report.peak_concurrent_moves = report.peak_concurrent_moves.max(1);
97 remaining_bytes = remaining_bytes.saturating_sub(location.size());
98 match self.move_blob(source, target, hash) {
99 Ok(Some(bytes)) => {
100 report.resumed += 1;
101 report.moved += 1;
102 report.bytes_moved = report.bytes_moved.saturating_add(bytes);
103 update_member_bytes(&mut members, source, target, bytes);
104 }
105 Ok(None) => {}
106 Err(error) => report.failed.push(format!("{hash:?}: {error}")),
107 }
108 }
109 while attempted_moves < self.temperature_config.max_moves_per_cycle && remaining_bytes > 0 {
110 if self.temperature_foreground_busy()? {
111 report.throttled = true;
112 break;
113 }
114 let batch_limit = self
115 .temperature_config
116 .max_concurrent_moves
117 .max(1)
118 .min(self.temperature_config.max_moves_per_cycle - attempted_moves);
119 let mut plans = Vec::with_capacity(batch_limit);
120 while plans.len() < batch_limit && remaining_bytes > 0 {
121 let candidate = {
122 let mut temperature = self
123 .temperature
124 .lock()
125 .map_err(|_| StoreError::Other("pool temperature lock poisoned".into()))?;
126 let mut planned = None;
127 let hot_len = temperature.hot.len();
128 for _ in 0..hot_len {
129 let Some(candidate) = temperature
130 .hot
131 .remove_best_by(|candidate, current| candidate.heat > current.heat)
132 else {
133 break;
134 };
135 if let Some(target) = self.plan_hot_move(candidate, &members, now) {
136 planned = Some((candidate, target));
137 break;
138 }
139 }
140 if planned.is_none() {
141 let cold_len = temperature.cold.len();
142 for _ in 0..cold_len {
143 let Some(candidate) =
144 temperature.cold.remove_best_by(|candidate, current| {
145 candidate.heat < current.heat
146 || (candidate.heat == current.heat
147 && candidate.last_accessed_at
148 < current.last_accessed_at)
149 })
150 else {
151 break;
152 };
153 if let Some(target) =
154 self.plan_cold_move(candidate, &members, &pressured, now)
155 {
156 planned = Some((candidate, target));
157 break;
158 }
159 }
160 }
161 planned
162 };
163 let Some((candidate, target)) = candidate else {
164 break;
165 };
166 if candidate.size > remaining_bytes {
167 continue;
168 }
169 remaining_bytes = remaining_bytes.saturating_sub(candidate.size);
170 update_member_bytes(&mut members, candidate.member, target, candidate.size);
171 plans.push((candidate, target));
172 }
173 if plans.is_empty() {
174 break;
175 }
176 attempted_moves = attempted_moves.saturating_add(plans.len());
177 report.attempted_moves = report.attempted_moves.saturating_add(plans.len());
178 report.peak_concurrent_moves = report.peak_concurrent_moves.max(plans.len());
179 let outcomes = std::thread::scope(|scope| {
180 let workers = plans
181 .into_iter()
182 .map(|(candidate, target)| {
183 scope.spawn(move || {
184 (
185 candidate,
186 target,
187 self.move_blob(candidate.member, target, candidate.hash),
188 )
189 })
190 })
191 .collect::<Vec<_>>();
192 workers
193 .into_iter()
194 .map(|worker| worker.join())
195 .collect::<Vec<_>>()
196 });
197 for outcome in outcomes {
198 match outcome {
199 Ok((_, _, Ok(Some(bytes)))) => {
200 report.moved += 1;
201 report.bytes_moved = report.bytes_moved.saturating_add(bytes);
202 }
203 Ok((_, _, Ok(None))) => {}
204 Ok((candidate, _, Err(error))) => {
205 report.failed.push(format!("{:?}: {error}", candidate.hash));
206 }
207 Err(_) => report
208 .failed
209 .push("temperature move worker panicked".into()),
210 }
211 }
212 members = self.temperature_members()?;
213 }
214 Ok(report)
215 }
216
217 fn temperature_members(&self) -> Result<Vec<PoolMemberStatus>, StoreError> {
218 self.refresh_members()?;
219 let manifest = self.read_manifest()?;
220 let runtime = self
221 .runtime
222 .read()
223 .map_err(|_| StoreError::Other("pool runtime lock poisoned".into()))?;
224 manifest
225 .members
226 .into_iter()
227 .map(|member| {
228 let (logical_bytes, available, last_error) = match runtime.stores.get(&member.id) {
229 Some(store) => match store.stats() {
230 Ok(stats) => (stats.total_bytes, true, None),
231 Err(error) => (0, false, Some(error.to_string())),
232 },
233 None => (0, false, runtime.errors.get(&member.id).cloned()),
234 };
235 Ok(PoolMemberStatus {
236 id: member.id,
237 state: member.state,
238 path: member.config.path,
239 capacity_bytes: member.config.capacity_bytes,
240 map_size_bytes: member.config.map_size_bytes,
241 external_blob_dir: member.config.external_blob_dir,
242 external_blob_min_bytes: member.config.external_blob_min_bytes,
243 external_blob_sync: member.config.external_blob_sync,
244 external_pack_target_bytes: member.config.external_pack_target_bytes,
245 max_read_concurrency: member.config.max_read_concurrency,
246 max_write_concurrency: member.config.max_write_concurrency,
247 temperature_low_watermark_percent: member
248 .config
249 .temperature_low_watermark_percent,
250 temperature_high_watermark_percent: member
251 .config
252 .temperature_high_watermark_percent,
253 logical_bytes,
254 located_blobs: 0,
255 available,
256 last_error,
257 })
258 })
259 .collect()
260 }
261
262 fn plan_hot_move(
263 &self,
264 candidate: TemperatureCandidate,
265 members: &[PoolMemberStatus],
266 now: u64,
267 ) -> Option<PoolMemberId> {
268 if now.saturating_sub(candidate.placed_at)
269 < self.temperature_config.minimum_residence.as_secs()
270 {
271 return None;
272 }
273 let mut targets = members
274 .iter()
275 .filter(|member| {
276 member.id != candidate.member
277 && member.state == PoolMemberState::Active
278 && member.available
279 && below_high_watermark(member, candidate.size)
280 })
281 .map(|member| member.id)
282 .collect::<Vec<_>>();
283 let adaptive = self.adaptive.lock().ok()?;
284 adaptive.order_temperature_targets(&mut targets);
285 targets.into_iter().find(|target| {
286 adaptive.meaningfully_faster(
287 *target,
288 candidate.member,
289 self.temperature_config.promotion_hysteresis_percent,
290 )
291 })
292 }
293
294 fn plan_cold_move(
295 &self,
296 candidate: TemperatureCandidate,
297 members: &[PoolMemberStatus],
298 pressured: &HashSet<PoolMemberId>,
299 now: u64,
300 ) -> Option<PoolMemberId> {
301 if now.saturating_sub(candidate.placed_at)
302 < self.temperature_config.minimum_residence.as_secs()
303 {
304 return None;
305 }
306 let source = members.iter().find(|member| {
307 member.id == candidate.member
308 && member.state == PoolMemberState::Active
309 && member.available
310 })?;
311 if !pressured.contains(&source.id)
312 || fill_percent(source.logical_bytes, source.capacity_bytes)
313 <= source.temperature_low_watermark_percent
314 {
315 return None;
316 }
317
318 members
319 .iter()
320 .filter(|target| {
321 target.id != source.id
322 && target.state == PoolMemberState::Active
323 && target.available
324 && target.logical_bytes.saturating_add(candidate.size) <= target.capacity_bytes
325 })
326 .max_by_key(|target| {
327 (
328 target.capacity_bytes > source.capacity_bytes,
329 target.capacity_bytes.saturating_sub(target.logical_bytes),
330 target.capacity_bytes,
331 )
332 })
333 .map(|target| target.id)
334 }
335
336 fn temperature_foreground_busy(&self) -> Result<bool, StoreError> {
337 self.refresh_members()?;
338 let runtime = self
339 .runtime
340 .read()
341 .map_err(|_| StoreError::Other("pool runtime lock poisoned".into()))?;
342 let threshold = self.temperature_config.foreground_load_percent.max(1);
343 for gate in runtime
344 .read_gates
345 .values()
346 .chain(runtime.write_gates.values())
347 {
348 if gate.load_percent()? >= threshold {
349 return Ok(true);
350 }
351 }
352 Ok(false)
353 }
354}
355
356fn fill_percent(bytes: u64, capacity: u64) -> u8 {
357 if capacity == 0 {
358 return 100;
359 }
360 bytes.saturating_mul(100).saturating_div(capacity).min(100) as u8
361}
362
363fn below_high_watermark(member: &PoolMemberStatus, incoming: u64) -> bool {
364 fill_percent(
365 member.logical_bytes.saturating_add(incoming),
366 member.capacity_bytes,
367 ) <= member.temperature_high_watermark_percent
368}
369
370fn update_member_bytes(
371 members: &mut [PoolMemberStatus],
372 source: PoolMemberId,
373 target: PoolMemberId,
374 bytes: u64,
375) {
376 for member in members {
377 if member.id == source {
378 member.logical_bytes = member.logical_bytes.saturating_sub(bytes);
379 } else if member.id == target {
380 member.logical_bytes = member.logical_bytes.saturating_add(bytes);
381 }
382 }
383}