dig_peer_selector/engine.rs
1//! [`PeerSelector`] — the public engine (SPEC.md §5): the closed `select → record_outcome →
2//! rebalance` decision + learning loop over the registry (§2), the measured-only quality model (§3),
3//! and the autonomous scorer (§4).
4//!
5//! The selector is pure + in-memory: `select`/`rebalance` never block on I/O, and `record_outcome`
6//! folds a measured result into the models in real time so the next decision is smarter (SPEC §5).
7//! Interior mutability ([`std::sync::Mutex`]) lets the hooks take `&self` (SPEC §5.1) while the engine
8//! mutates the registry + learned models.
9
10use std::sync::Mutex;
11
12use rand::{Rng, SeedableRng};
13use rand_pcg::Pcg64Mcg;
14
15use dig_nat::{PeerId, TraversalKind};
16
17use crate::config::SelectorConfig;
18use crate::observe::{PeerSnapshot, SelectorSnapshot};
19use crate::pool_event::{PoolEvent, PoolRemovalReason};
20use crate::quality::PeerQuality;
21use crate::registry::Registry;
22use crate::scoring::{score_peer, PeerClass, RelayModel, SaturationModel, ScoredPeer};
23use crate::types::{
24 Candidate, ContentRequest, OutcomeKind, OutcomeResult, Provenance, RangePlanDelta,
25 SelectedPeer, Selection, TransferOutcome,
26};
27
28/// The learned + live state behind the engine's interior mutability.
29struct Inner {
30 registry: Registry,
31 saturation: SaturationModel,
32 relay: RelayModel,
33 rng: Pcg64Mcg,
34 /// Records, per dispatched range, the `(peer, in_flight_at_dispatch, class)` so a later outcome
35 /// can attribute the saturation observation to the concurrency the range actually ran under
36 /// (SPEC §4.1). Keyed by `(peer_id, range_index)`.
37 dispatched: std::collections::HashMap<(PeerId, usize), DispatchRecord>,
38 /// A monotonically-increasing `select`/`rebalance` counter — the "epoch" used to round-robin
39 /// exploration coverage so no eligible peer is starved (SPEC §4.4-E: every candidate is tried,
40 /// even at parallelism 1; SPEC §4.4-D: a degraded-then-recovered peer is re-probed).
41 epoch: u64,
42 /// The last epoch each peer appeared in a selection — drives the anti-starvation re-exploration:
43 /// the most-starved eligible peer (largest `epoch - last_selected`, cold peers having `0`) is
44 /// guaranteed a periodic turn so it acquires fresh outcomes.
45 last_selected: std::collections::HashMap<PeerId, u64>,
46}
47
48impl Inner {
49 /// Prune `last_selected` + `dispatched` for every `peer_id` the registry has just reported as
50 /// evicted/removed (#179 finding 2). MUST be called after every registry-mutating operation that
51 /// can shed an entry (capacity eviction on upsert/mark-connected/set-class, or explicit removal) so
52 /// neither side map ever retains a key for a peer no longer in the registry. Cheap when nothing was
53 /// evicted (the common case): `drain_evicted` returns an empty `Vec`.
54 fn prune_evicted_side_maps(&mut self) {
55 for peer in self.registry.drain_evicted() {
56 self.last_selected.remove(&peer);
57 self.dispatched.retain(|(id, _), _| *id != peer);
58 }
59 }
60}
61
62/// What a dispatched range recorded at dispatch time, for saturation learning (SPEC §4.1).
63#[derive(Clone, Copy)]
64struct DispatchRecord {
65 in_flight_at_dispatch: u32,
66 class: PeerClass,
67}
68
69/// The self-optimizing peer selector (SPEC §5, §11). Construct with [`PeerSelector::new`], feed it
70/// registry churn + DHT candidates, ask it to [`select`](PeerSelector::select) a source subset, and
71/// stream every measured outcome back via [`record_outcome`](PeerSelector::record_outcome).
72pub struct PeerSelector {
73 config: SelectorConfig,
74 inner: Mutex<Inner>,
75}
76
77impl PeerSelector {
78 /// Construct a selector from wiring-only [`SelectorConfig`] (SPEC §5.6). No behavior knobs — every
79 /// tradeoff is learned.
80 pub fn new(config: SelectorConfig) -> Self {
81 let rng = Pcg64Mcg::seed_from_u64(config.effective_seed());
82 let registry = Registry::new(config.registry_capacity);
83 PeerSelector {
84 config,
85 inner: Mutex::new(Inner {
86 registry,
87 saturation: SaturationModel::default(),
88 relay: RelayModel::default(),
89 rng,
90 dispatched: std::collections::HashMap::new(),
91 epoch: 0,
92 last_selected: std::collections::HashMap::new(),
93 }),
94 }
95 }
96
97 /// The current time from the injected clock (SPEC §5.6).
98 fn now(&self) -> u64 {
99 self.config.clock.now()
100 }
101
102 /// **`select`** — rank a candidate set for a request (SPEC §5.1).
103 ///
104 /// Registers any fresh candidate (cold — its first pick is exploratory, SPEC §2.3, §4.4-E),
105 /// scores every eligible candidate against the learned models, and returns a ranked chosen subset
106 /// bounded by the request's parallelism. Each chosen peer carries its recommended `max_concurrency`
107 /// (its learned saturation headroom, SPEC §4.1/§4.5) and its `exploratory` flag. Anti-herd: total
108 /// dispatched concurrency is filled peer-by-peer up to each peer's headroom, spilling to the next
109 /// (SPEC §4.4-B). Dispatch bumps in-flight so the very next `select`/`rebalance` sees it.
110 pub fn select(&self, req: &ContentRequest, candidates: &[Candidate]) -> Selection {
111 let now = self.now();
112 let mut inner = self.inner.lock().expect("selector mutex poisoned");
113 // Register fresh candidates (cold) — selecting a candidate registers it (SPEC §2.3).
114 for c in candidates {
115 inner.registry.upsert_candidate(c, Provenance::Dht, now);
116 }
117 // A fresh upsert may have evicted an over-capacity entry (§2.5) — prune its side-map residue
118 // before scoring (#179 finding 2).
119 inner.prune_evicted_side_maps();
120 self.select_over(&mut inner, req, &candidate_ids(candidates), &[], now)
121 }
122
123 /// **`rebalance`** — re-query the up-to-the-moment models for still-needed ranges (SPEC §5.5).
124 ///
125 /// Mid-transfer, when a source drops or a range must be relocated, the executor calls this to get
126 /// a replacement subset for the `need`ed ranges, de-ranking the already-`active` peers (they are
127 /// counted as busy so their headroom shrinks) and reflecting every `record_outcome` received so
128 /// far. It obeys the same invariants as `select` (SPEC §4.4).
129 pub fn rebalance(
130 &self,
131 req: &ContentRequest,
132 active: &[PeerId],
133 need: &RangePlanDelta,
134 ) -> Selection {
135 let now = self.now();
136 let mut inner = self.inner.lock().expect("selector mutex poisoned");
137 // The still-needed count drives how much parallelism to fill.
138 let want = need.len().max(1);
139 let effective = ContentRequest {
140 parallelism: req.effective_parallelism().min(want).max(1),
141 ..req.clone()
142 };
143 // Candidate population = every eligible registry peer (the freshly-learned view), so a peer
144 // that recovered or a newly-fed candidate can be chosen. `active` peers are de-ranked.
145 let pool: Vec<PeerId> = inner
146 .registry
147 .iter()
148 .filter(|e| e.is_eligible())
149 .map(|e| e.peer_id)
150 .collect();
151 self.select_over(&mut inner, &effective, &pool, active, now)
152 }
153
154 /// Core selection over an explicit candidate-id `pool`, de-ranking `deranked` peers. Shared by
155 /// `select` and `rebalance` so the invariants hold identically for both (SPEC §4.4, §5.5).
156 fn select_over(
157 &self,
158 inner: &mut Inner,
159 req: &ContentRequest,
160 pool: &[PeerId],
161 deranked: &[PeerId],
162 now: u64,
163 ) -> Selection {
164 // Reclaim any dispatch whose outcome never arrived before scoring (#179 finding 1): a peer
165 // dispatched to and then gone silent must not keep counting as busy / unevictable forever.
166 inner.registry.reclaim_stale_in_flight(now);
167 if pool.is_empty() {
168 return Selection::empty();
169 }
170
171 // Score every eligible pooled peer ONCE (#179 LOW finding: a separate `proven_score_bounds`
172 // pre-pass used to re-score every non-cold peer before this loop scored everyone again,
173 // doubling score_peer calls on the hot decision path). Score with bonus=0 first — a cold
174 // peer's `effective_score` under `score_peer` is exactly the bonus and nothing else depends on
175 // it (headroom/exploratory/tie_break are bonus-independent, SPEC §4.4-E), so the proven-score
176 // bounds can be derived from this single pass and the bonus applied to cold entries afterward.
177 let mut scored: Vec<(PeerId, ScoredPeer)> = Vec::with_capacity(pool.len());
178 for id in pool {
179 let Some(entry) = inner.registry.get(id) else {
180 continue;
181 };
182 if !entry.is_eligible() {
183 continue;
184 }
185 scored.push((*id, score_peer(entry, &inner.saturation, &inner.relay, 0.0)));
186 }
187 if scored.is_empty() {
188 return Selection::empty();
189 }
190
191 // Exploration bonus: a cold peer scores just ABOVE the worst proven peer (so it gets tried)
192 // but strictly BELOW the best proven peer (so it never displaces a proven fast peer for the
193 // bulk of a transfer — SPEC §4.4-E). We place it a small fraction of the proven-score gap
194 // above the worst; when there are no proven peers (all-cold pool) the bonus is 0 and cold
195 // peers simply order among themselves. Derived from the single scored pass above (skipping
196 // cold/bad-source entries, both identifiable from the already-computed `ScoredPeer`s).
197 let (worst_proven, best_proven) = proven_score_bounds(&scored);
198 let exploration_bonus = if best_proven > worst_proven {
199 // A quarter of the way up from worst to best: above the worst proven peer, below the best.
200 worst_proven + 0.25 * (best_proven - worst_proven)
201 } else if best_proven > 0.0 {
202 // A single proven peer (worst == best): explore just below it so it stays rank 0.
203 best_proven * 0.5
204 } else {
205 0.0
206 };
207
208 // Apply the now-known bonus to every cold (exploratory) entry — this is the only field a
209 // cold peer's score depends on (SPEC §4.4-E) — then the SPEC §5.5 de-rank pass for `rebalance`.
210 for (id, s) in &mut scored {
211 if s.exploratory {
212 s.effective_score = exploration_bonus;
213 }
214 // De-rank an already-active peer: shrink its headroom to reflect that it is busy, and
215 // discount its score so a replacement is preferred (SPEC §5.5).
216 if deranked.contains(id) {
217 s.effective_score *= DERANK_FACTOR;
218 s.headroom = s.headroom.saturating_sub(1);
219 }
220 }
221
222 // Advance the epoch (this select/rebalance) for anti-starvation exploration coverage.
223 inner.epoch = inner.epoch.wrapping_add(1);
224 let epoch = inner.epoch;
225
226 // Cap how many COLD exploratory peers may appear so unmeasured peers don't crowd out proven
227 // ones for the bulk of a transfer (SPEC §4.4-E). At least one exploratory slot so an all-cold
228 // network still makes progress; otherwise a small fraction of the requested parallelism.
229 let want = req.effective_parallelism();
230 let explore_cap = explore_slots(want);
231
232 // Rank: highest effective_score first; deterministic tie-break (SPEC §4.4-F). Among equal COLD
233 // exploratory peers, order by least-recently-selected (coverage), then seeded jitter (fair yet
234 // reproducible); everything else tie-breaks by stable salt. Proven peers have strictly higher
235 // effective scores than cold peers (§4.4-A/E), so they always sort ahead.
236 let last_sel: Vec<u64> = scored
237 .iter()
238 .map(|(id, _)| inner.last_selected.get(id).copied().unwrap_or(0))
239 .collect();
240 let jitter: Vec<u64> = scored.iter().map(|_| inner.rng.gen::<u64>()).collect();
241 let mut order: Vec<usize> = (0..scored.len()).collect();
242 order.sort_by(|&a, &b| {
243 let (_, sa) = &scored[a];
244 let (_, sb) = &scored[b];
245 sb.effective_score
246 .partial_cmp(&sa.effective_score)
247 .unwrap_or(std::cmp::Ordering::Equal)
248 .then_with(|| {
249 if sa.exploratory && sb.exploratory {
250 last_sel[a]
251 .cmp(&last_sel[b])
252 .then_with(|| jitter[a].cmp(&jitter[b]))
253 } else {
254 sa.tie_break.cmp(&sb.tie_break)
255 }
256 })
257 });
258
259 // ANTI-STARVATION re-exploration (SPEC §4.4-E "every candidate is tried" + §4.4-D "a degraded
260 // peer that recovers must rise again"). At low parallelism a proven peer can monopolize the
261 // only slot forever, so a peer never re-probed can never be re-measured — a fresh newcomer, a
262 // cold candidate, or a degraded-then-recovered peer would be invisible. We therefore reserve
263 // ONE slot for the MOST-STARVED eligible peer: the one not selected for the longest (cold peers
264 // have `last_selected == 0` = maximally starved). It is force-included ONLY when its staleness
265 // exceeds a round-robin threshold (~ the eligible pool size), so proven peers still serve the
266 // bulk of every transfer; it is placed LAST so it never displaces a proven peer from rank 0
267 // when `want > 1`. Marked exploratory (it is an uncertainty probe, not a proven pick).
268 let eligible = scored.len().max(1);
269 let starve_threshold = eligible as u64; // give each peer a turn within ~one pool sweep
270 // Forcing is only meaningful when there is genuine competition for slots (`want < eligible`).
271 // When every eligible peer already fits (`want >= eligible`), natural score order stands — a
272 // proven peer keeps rank 0 and no re-probe reordering is needed (this is what keeps the
273 // measured-fast peer ahead of a measured-slow one when both are selected).
274 let forced: Option<usize> = if want < eligible {
275 scored
276 .iter()
277 .enumerate()
278 .filter(|(idx, (_, s))| {
279 // A verification-failing bad source is never force-explored (SPEC §9.4 keeps it down).
280 s.effective_score > -1.0e11 && {
281 let staleness = epoch.saturating_sub(last_sel[*idx]);
282 staleness >= starve_threshold
283 }
284 })
285 .max_by(|(ia, (_, sa)), (ib, (_, sb))| {
286 let stale_a = epoch.saturating_sub(last_sel[*ia]);
287 let stale_b = epoch.saturating_sub(last_sel[*ib]);
288 stale_a
289 .cmp(&stale_b)
290 .then_with(|| sb.tie_break.cmp(&sa.tie_break)) // deterministic
291 })
292 .map(|(idx, _)| idx)
293 } else {
294 None
295 };
296
297 // Pass 1: take the top-scored peers up to `want`, honoring the exploration cap; reserve one
298 // slot for the forced-coverage peer (cap at `want - 1`) so it always fits. At `want == 1` the
299 // forced peer takes the sole slot THIS round (a periodic re-probe); because a selected peer's
300 // staleness resets, the best peer reclaims the slot in the intervening rounds — so the stream
301 // still exploits the best peer most of the time while guaranteeing the starved peer is
302 // periodically re-measured (SPEC §4.4-D/E).
303 let cap_pass1 = if forced.is_some() {
304 want.saturating_sub(1)
305 } else {
306 want
307 };
308 // The forced-coverage pick consumes one exploration slot (it is an exploratory re-probe), so
309 // the total number of exploratory peers — pass-1 cold picks PLUS the forced pick — never
310 // exceeds `explore_cap` (SPEC §4.4-E: unmeasured/uncertain peers never crowd out proven ones).
311 let forced_is_exploratory = forced.map(|i| scored[i].1.exploratory).unwrap_or(false);
312 let explore_budget_pass1 = if forced_is_exploratory {
313 explore_cap.saturating_sub(1)
314 } else {
315 explore_cap
316 };
317 let mut chosen: Vec<usize> = Vec::with_capacity(want);
318 let mut explore_used = 0usize;
319 for &i in &order {
320 if chosen.len() >= cap_pass1 {
321 break;
322 }
323 if Some(i) == forced {
324 continue; // reserved for pass 2
325 }
326 let (_, s) = &scored[i];
327 if s.exploratory {
328 if explore_used >= explore_budget_pass1 {
329 continue;
330 }
331 explore_used += 1;
332 }
333 chosen.push(i);
334 }
335 // Pass 2: append the forced-coverage peer in its reserved slot (lowest rank), if not already in.
336 if let Some(fi) = forced {
337 if chosen.len() < want && !chosen.contains(&fi) {
338 chosen.push(fi);
339 }
340 }
341
342 // Materialize the selection, recording each chosen peer's selection epoch (anti-starvation).
343 let mut selection = Selection::empty();
344 for (rank, &i) in chosen.iter().enumerate() {
345 let (peer_id, s) = scored[i];
346 inner.last_selected.insert(peer_id, epoch);
347 // A forced-coverage pick of an otherwise-proven peer is still flagged exploratory (it is an
348 // uncertainty re-probe, not a merit selection) so the host/executor treats it accordingly.
349 let exploratory = s.exploratory || Some(i) == forced;
350 selection.peers.push(SelectedPeer {
351 peer_id,
352 rank: rank as u32,
353 max_concurrency: s.headroom.max(1),
354 exploratory,
355 });
356 }
357
358 // Account the dispatch: bump in-flight + record the concurrency each peer will run under, so a
359 // later outcome attributes the saturation observation correctly (SPEC §4.1, §5.3).
360 for sp in &selection.peers {
361 let class = inner
362 .registry
363 .get(&sp.peer_id)
364 .map(|e| PeerClass::of(e.connection_class))
365 .unwrap_or(PeerClass::Unknown);
366 let in_flight_at_dispatch = inner
367 .registry
368 .get(&sp.peer_id)
369 .map(|e| e.quality.in_flight)
370 .unwrap_or(0);
371 inner
372 .registry
373 .add_in_flight(&sp.peer_id, sp.max_concurrency, now);
374 // Record dispatch context for the range indices this peer will serve. We do not know exact
375 // indices here, so key by a rolling per-peer marker; the outcome's own range index keys
376 // the lookup and falls back to the peer's recorded context.
377 inner.dispatched.insert(
378 (sp.peer_id, usize::MAX),
379 DispatchRecord {
380 in_flight_at_dispatch,
381 class,
382 },
383 );
384 }
385
386 selection
387 }
388
389 /// **`record_outcome`** — fold a measured result back into the models in real time (SPEC §5.2).
390 ///
391 /// Derives throughput strictly from measured `bytes / duration_ms` (never a self-reported rate,
392 /// SPEC §9.3); updates throughput/RTT/reliability (SPEC §3.2, §3.4), the per-class saturation
393 /// model (SPEC §4.1), the adaptive relayed penalty (SPEC §4.2), and decrements in-flight (SPEC
394 /// §5.3). An outcome for an unknown peer upserts a cold entry first (self-healing, SPEC §5.2).
395 pub fn record_outcome(&self, outcome: &TransferOutcome) {
396 let now = self.now();
397 let mut inner = self.inner.lock().expect("selector mutex poisoned");
398
399 // Self-heal: an outcome for an unknown peer registers it cold first (SPEC §5.2).
400 if inner.registry.get(&outcome.peer_id).is_none() {
401 let cand = Candidate::new(outcome.peer_id, Vec::new());
402 inner.registry.upsert_candidate(&cand, Provenance::Nat, now);
403 // The self-heal upsert may have evicted a different over-capacity entry — prune its
404 // side-map residue (#179 finding 2).
405 inner.prune_evicted_side_maps();
406 }
407
408 // Retrieve the dispatch context for saturation attribution (SPEC §4.1) before mutating.
409 let range_index = match outcome.kind {
410 OutcomeKind::Range { index, .. } => index,
411 OutcomeKind::Request { .. } => usize::MAX,
412 };
413 let dispatch = inner
414 .dispatched
415 .remove(&(outcome.peer_id, range_index))
416 .or_else(|| inner.dispatched.remove(&(outcome.peer_id, usize::MAX)));
417 let class = dispatch.map(|d| d.class).unwrap_or_else(|| {
418 inner
419 .registry
420 .get(&outcome.peer_id)
421 .map(|e| PeerClass::of(e.connection_class))
422 .unwrap_or(PeerClass::Unknown)
423 });
424 let in_flight_at_dispatch =
425 dispatch
426 .map(|d| d.in_flight_at_dispatch)
427 .unwrap_or_else(|| {
428 inner
429 .registry
430 .get(&outcome.peer_id)
431 .map(|e| e.quality.in_flight)
432 .unwrap_or(1)
433 });
434
435 let throughput = outcome.throughput_bps();
436
437 // Fold into the learned aggregate models (saturation + relay) — measured throughput only.
438 if let (Some(bps), true) = (throughput, outcome.is_success()) {
439 inner.saturation.observe(class, in_flight_at_dispatch, bps);
440 inner.relay.observe(class.is_relayed(), bps);
441 }
442
443 // Fold into the per-peer quality model.
444 if let Some(entry) = inner.registry.get_mut(&outcome.peer_id) {
445 apply_outcome_to_quality(&mut entry.quality, outcome, throughput);
446 entry.quality.bump_samples();
447 entry.last_outcome_at = Some(outcome.at.max(now));
448 }
449 // Decrement in-flight: one range settled (SPEC §5.3, symmetric with the `select` dispatch bump).
450 inner.registry.release_in_flight(&outcome.peer_id, 1);
451 }
452
453 /// Consume a `dig-gossip` churn event to keep the registry live (SPEC §5.4, §2.3). Byte-compatible
454 /// with `dig_gossip::PoolEvent` (see [`crate::pool_event`]).
455 pub fn on_pool_event(&self, event: &PoolEvent) {
456 let now = self.now();
457 let mut inner = self.inner.lock().expect("selector mutex poisoned");
458 match event {
459 PoolEvent::PeerAdded { peer_id, .. } => {
460 inner
461 .registry
462 .mark_connected(*peer_id, Provenance::Gossip, now);
463 // Insertion may have evicted a different over-capacity entry (#179 finding 2).
464 inner.prune_evicted_side_maps();
465 }
466 PoolEvent::PeerRemoved { peer_id, reason } => {
467 let banned = matches!(reason, PoolRemovalReason::Banned);
468 inner.registry.mark_disconnected(peer_id, banned);
469 // A plain disconnect never deletes the entry (SPEC §2.3), so nothing to prune here.
470 }
471 }
472 }
473
474 /// Attach / update a live peer's connection class from `dig-nat` (SPEC §5.4, §7.3).
475 pub fn on_connection_class(&self, peer: &PeerId, class: TraversalKind) {
476 let now = self.now();
477 let mut inner = self.inner.lock().expect("selector mutex poisoned");
478 inner.registry.set_connection_class(*peer, class, now);
479 // Insertion may have evicted a different over-capacity entry (#179 finding 2).
480 inner.prune_evicted_side_maps();
481 }
482
483 /// Manually upsert a candidate (seed / bootstrap feed, SPEC §5.4). A fresh peer is cold.
484 pub fn upsert_candidate(&self, candidate: &Candidate) {
485 let now = self.now();
486 let mut inner = self.inner.lock().expect("selector mutex poisoned");
487 inner
488 .registry
489 .upsert_candidate(candidate, Provenance::Manual, now);
490 // Insertion may have evicted a different over-capacity entry (#179 finding 2).
491 inner.prune_evicted_side_maps();
492 }
493
494 /// Explicitly remove a peer (rare; churn usually drives this — SPEC §5.4). A peer with a range in
495 /// flight is retained until it settles.
496 pub fn remove_peer(&self, peer: &PeerId) {
497 let mut inner = self.inner.lock().expect("selector mutex poisoned");
498 inner.registry.remove(peer);
499 // The removal itself must prune the same peer's side-map residue (#179 finding 2).
500 inner.prune_evicted_side_maps();
501 }
502
503 /// Explicitly note that `count` ranges were dispatched to `peer` (SPEC §5.3). Optional: `select`
504 /// already accounts the dispatch it returns; the host uses this only if it dispatches outside a
505 /// `select` result.
506 pub fn on_dispatch(&self, peer: &PeerId, count: u32) {
507 let now = self.now();
508 let mut inner = self.inner.lock().expect("selector mutex poisoned");
509 inner.registry.add_in_flight(peer, count, now);
510 }
511
512 // ---- Read-only observability (SPEC §5.7) ------------------------------------------------
513
514 /// A read-only snapshot of one peer's learned model, or `None` if unknown (SPEC §5.7).
515 pub fn peer_snapshot(&self, peer: &PeerId) -> Option<PeerSnapshot> {
516 let inner = self.inner.lock().expect("selector mutex poisoned");
517 inner.registry.get(peer).map(PeerSnapshot::of)
518 }
519
520 /// A read-only snapshot of the selector's learned aggregate state (SPEC §5.7). Includes the
521 /// engine's internal side-map sizes (`last_selected_len`, `dispatched_len`) so a host can confirm
522 /// they track the live registry population rather than growing unboundedly (#179 finding 2).
523 pub fn snapshot(&self) -> SelectorSnapshot {
524 let inner = self.inner.lock().expect("selector mutex poisoned");
525 SelectorSnapshot::build(
526 inner.registry.iter(),
527 &inner.saturation,
528 &inner.relay,
529 inner.last_selected.len(),
530 inner.dispatched.len(),
531 )
532 }
533
534 /// The current registry size (SPEC §5.7).
535 pub fn registry_size(&self) -> usize {
536 let inner = self.inner.lock().expect("selector mutex poisoned");
537 inner.registry.len()
538 }
539}
540
541/// De-rank multiplier for an already-active peer during `rebalance` (SPEC §5.5): its score is scaled
542/// down so a fresh replacement is preferred for the still-needed ranges.
543const DERANK_FACTOR: f64 = 0.5;
544
545/// How many COLD exploratory peers may appear in a selection of size `want` (SPEC §4.4-E): at least
546/// one (so an all-cold network makes progress), otherwise ~a third of the requested parallelism so
547/// unmeasured peers never dominate the bulk of a transfer.
548fn explore_slots(want: usize) -> usize {
549 want.div_ceil(3).max(1)
550}
551
552/// The worst + best *proven* (non-cold, non-bad) effective score among an already-scored pool, so the
553/// exploration bonus sits just above the worst proven peer (SPEC §4.4-E). Returns `(0.0, 0.0)` when
554/// there are no proven peers (an all-cold pool) — exploration then simply orders cold peers among
555/// themselves.
556///
557/// Takes the pool's `ScoredPeer`s (scored with `exploration_bonus=0.0`) rather than re-scoring from
558/// the registry (#179 LOW finding) — a proven (non-cold, non-bad-source) peer's score does not depend
559/// on the exploration bonus at all (SPEC §4.4: the bonus only ever appears in the cold branch), so
560/// scoring once and filtering is equivalent to the old two-pass approach, at half the `score_peer`
561/// calls.
562fn proven_score_bounds(scored: &[(PeerId, ScoredPeer)]) -> (f64, f64) {
563 let mut worst = f64::INFINITY;
564 let mut best = f64::NEG_INFINITY;
565 let mut any = false;
566 for (_, s) in scored {
567 if s.exploratory {
568 continue; // cold — not a "proven" bound
569 }
570 if s.effective_score <= -1.0e11 {
571 continue; // a bad source floor — not a "proven good" bound
572 }
573 any = true;
574 worst = worst.min(s.effective_score);
575 best = best.max(s.effective_score);
576 }
577 if any {
578 (worst.max(0.0), best.max(0.0))
579 } else {
580 (0.0, 0.0)
581 }
582}
583
584/// Fold one measured outcome into a peer's quality model (SPEC §3.2, §3.4). Throughput/RTT move only
585/// on a success with a derivable rate; reliability moves on every peer-attributable result.
586fn apply_outcome_to_quality(
587 quality: &mut PeerQuality,
588 outcome: &TransferOutcome,
589 throughput: Option<f64>,
590) {
591 match outcome.result {
592 OutcomeResult::Success => {
593 if let Some(bps) = throughput {
594 quality.observe_throughput(bps);
595 }
596 if let Some(rtt) = outcome.rtt_ms {
597 quality.observe_rtt(rtt as f64);
598 }
599 quality.observe_result(true, false);
600 }
601 OutcomeResult::Failure { reason } => {
602 if reason.blames_peer() {
603 quality.observe_result(false, reason.is_hard());
604 }
605 }
606 OutcomeResult::Interrupted { .. } => {
607 // A partial transfer that dropped: a soft reliability failure (it left a range straggling).
608 quality.observe_result(false, false);
609 }
610 }
611}
612
613/// The `peer_id`s of a candidate slice (dispatch pool for `select`).
614fn candidate_ids(candidates: &[Candidate]) -> Vec<PeerId> {
615 candidates.iter().map(|c| c.peer_id).collect()
616}
617
618#[cfg(test)]
619mod tests {
620 use super::*;
621 use crate::config::SelectorConfig;
622 use crate::scoring::{SCORE_PEER_CALLS, SCORE_PEER_CALLS_LOCK};
623 use crate::types::{Candidate, OutcomeKind, OutcomeResult, TransferOutcome};
624 use dig_dht::{CandidateAddr, ContentId};
625 use std::sync::atomic::Ordering;
626
627 fn pid(b: u8) -> PeerId {
628 PeerId::from_bytes([b; 32])
629 }
630 fn candidate(b: u8) -> Candidate {
631 Candidate::new(pid(b), vec![CandidateAddr::direct("10.0.0.1", 9444)])
632 }
633 fn content() -> ContentId {
634 ContentId::store([0x77; 32])
635 }
636
637 /// #179 LOW finding (select/rebalance double-scores the pool): a single `select` call must invoke
638 /// `score_peer` AT MOST ONCE per eligible pooled peer. The old `proven_score_bounds` pre-pass
639 /// re-scored every non-cold peer before the main scoring loop scored everyone again, doubling the
640 /// work for a mixed cold/proven pool on the hot decision path (rebalance can run over up to ~4096
641 /// eligible entries on every dropped source mid-transfer).
642 #[test]
643 fn select_scores_each_pooled_peer_at_most_once() {
644 let sel = PeerSelector::new(SelectorConfig::deterministic(1000, 5));
645 // Warm several peers to "proven" (non-cold) status so the old pre-pass had work to do.
646 for b in 0..5u8 {
647 for i in 0..3u64 {
648 sel.record_outcome(&TransferOutcome {
649 peer_id: pid(b),
650 content: content(),
651 kind: OutcomeKind::Range {
652 index: i as usize,
653 offset: 0,
654 length: 1000,
655 },
656 result: OutcomeResult::Success,
657 bytes: 100_000,
658 duration_ms: 1000,
659 rtt_ms: Some(10),
660 at: 1000 + i,
661 });
662 }
663 }
664 // Mix in a couple of fresh cold candidates too (a realistic mixed pool).
665 let candidates: Vec<Candidate> = (0..7u8).map(candidate).collect();
666
667 // SCORE_PEER_CALLS is process-global (shared across cargo test's parallel threads): serialize
668 // via the dedicated lock and measure a delta across just the `select` call, so a concurrently
669 // running test's own `score_peer` calls cannot pollute this measurement.
670 let _guard = SCORE_PEER_CALLS_LOCK
671 .lock()
672 .unwrap_or_else(|e| e.into_inner());
673 let before = SCORE_PEER_CALLS.load(Ordering::SeqCst);
674 let _ = sel.select(&ContentRequest::new(content(), 3), &candidates);
675 let calls = SCORE_PEER_CALLS.load(Ordering::SeqCst) - before;
676 assert!(
677 calls <= candidates.len() as u64,
678 "score_peer must run at most once per pooled peer ({} peers), got {calls} calls \
679 (a separate proven-bounds pre-pass would double-count the non-cold ones)",
680 candidates.len()
681 );
682 }
683}