es_entity/operation/batch/
search.rs

1//! The bisect search, as a pure state machine.
2//!
3//! Index bookkeeping over the input's length: it decides *which contiguous
4//! range to probe next* and *what a probe's verdict means*, while the caller
5//! decides how a probe runs and where its transaction boundary sits. So both of
6//! these drive the same algorithm:
7//!
8//! - [`BatchIsolation::run_bisected`](super::BatchIsolation::run_bisected),
9//!   against `with_savepoint` inside a single enclosing transaction.
10//! - A caller that wants a transaction boundary *between* probes, committing
11//!   each clean range as it lands: it mints and commits its own operations
12//!   around [`next_range`](BisectSearch::next_range) and
13//!   [`report`](BisectSearch::report).
14//!
15//! Being plain bookkeeping, the probe sequence is also testable on its own.
16
17use std::{cmp::Ordering, collections::BinaryHeap, ops::Range};
18
19/// How many times a transiently-failed range may be re-probed before the search
20/// is abandoned.
21///
22/// A retry costs up to one `deadlock_timeout` wait, so this caps the time a
23/// search spends on contention.
24pub const DEFAULT_MAX_TRANSIENT_RETRIES: usize = 2;
25
26/// How many probes a bisect may spend before giving up on the ranges it has
27/// not yet resolved.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
29pub enum BisectBudget {
30    /// `2·⌈log₂N⌉ + 1` — the depth needed to isolate one failing item and
31    /// resolve every other item as clean. 9 probes at N=10, 15 at N=100.
32    #[default]
33    Auto,
34    /// An explicit cap, clamped to at least 1. `MaxProbes(1)` probes once and
35    /// leaves every item unresolved if that probe fails.
36    MaxProbes(usize),
37    /// No cap: keep splitting until every item is resolved. Worst case
38    /// `2N - 1` probes for an all-bad batch.
39    FullResolution,
40}
41
42impl BisectBudget {
43    /// The probe cap this budget implies for a batch of `n` items.
44    pub fn effective_cap(self, n: usize) -> usize {
45        match self {
46            Self::Auto => 2 * ceil_log2(n) + 1,
47            Self::MaxProbes(max) => max.max(1),
48            Self::FullResolution => usize::MAX,
49        }
50    }
51}
52
53/// `ceil(log2(n))`, defined as `0` for `n <= 1`.
54fn ceil_log2(n: usize) -> usize {
55    if n <= 1 {
56        return 0;
57    }
58    (usize::BITS - (n - 1).leading_zeros()) as usize
59}
60
61/// A range still awaiting a probe.
62///
63/// The [`Ord`] impl is the search's determinism guarantee: a
64/// [`BinaryHeap`] pops the **largest** range, and among equal lengths the one
65/// with the **earliest start**. For a given input and budget the probe sequence
66/// — and therefore the probe count — is reproducible.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68struct PendingRange {
69    start: usize,
70    end: usize,
71}
72
73impl PendingRange {
74    fn len(&self) -> usize {
75        self.end - self.start
76    }
77}
78
79impl Ord for PendingRange {
80    fn cmp(&self, other: &Self) -> Ordering {
81        self.len()
82            .cmp(&other.len())
83            .then_with(|| other.start.cmp(&self.start))
84    }
85}
86
87impl PartialOrd for PendingRange {
88    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
89        Some(self.cmp(other))
90    }
91}
92
93/// What a probe over one range turned out to be.
94#[derive(Debug)]
95pub enum ProbeVerdict<E> {
96    /// Every item in the range succeeded. They resolve here and are never
97    /// probed again.
98    Clean,
99    /// The range failed for a reason attributable to its contents. A
100    /// multi-item range splits; a single-item range resolves as that item's
101    /// failure.
102    Failed(E),
103    /// The failure describes contention: a deadlock victim, a serialization
104    /// failure, or a caller-classified conflict. The **same** range is
105    /// re-probed unsplit, and the probe is refunded to the budget.
106    Transient(E),
107}
108
109/// One input item's resolution. Positionally aligned with the input slice.
110#[derive(Debug, PartialEq, Eq)]
111pub enum ItemOutcome<E> {
112    /// Resolved by a clean probe over a range containing it.
113    Complete,
114    /// Resolved by its own single-item probe, so the error is attributable to
115    /// this item alone.
116    Failed(E),
117    /// Never resolved: the budget ran out, or the search was abandoned, while
118    /// this item was still inside an unprobed range. See
119    /// [`BisectOutcomes::last_error`] for what the search last saw.
120    Unresolved,
121}
122
123/// The result of a completed search: one outcome per input item, plus what it
124/// cost.
125#[derive(Debug)]
126pub struct BisectOutcomes<E> {
127    /// One entry per input item, in input order.
128    pub items: Vec<ItemOutcome<E>>,
129    /// Probes actually spent, net of refunded transient re-probes.
130    pub probes_used: usize,
131    /// Transient re-probes taken (refunded, so not counted in `probes_used`).
132    pub transient_retries: usize,
133    /// The most recent error from a probe spanning more than one item — what a
134    /// caller quotes when explaining an [`ItemOutcome::Unresolved`].
135    ///
136    /// `None` when every failure reached a single-item probe, since each of
137    /// those hands its error to [`ItemOutcome::Failed`].
138    pub last_error: Option<E>,
139}
140
141/// The transient allowance ran out: the same range kept failing on contention.
142///
143/// The search has nothing to attribute to any item, so the caller abandons it
144/// and retries the whole batch later.
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
146pub struct TransientLimitExceeded {
147    /// Probes spent before abandoning.
148    pub probes_used: usize,
149    /// Transient re-probes taken before abandoning.
150    pub transient_retries: usize,
151}
152
153impl std::fmt::Display for TransientLimitExceeded {
154    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
155        write!(
156            f,
157            "bisect abandoned after {} probes and {} transient retries",
158            self.probes_used, self.transient_retries
159        )
160    }
161}
162
163impl std::error::Error for TransientLimitExceeded {}
164
165/// Drives a bisect over `n` items: which range to probe next, and what each
166/// verdict means for the ranges still outstanding.
167///
168/// ```rust,ignore
169/// let mut search = BisectSearch::new(items.len(), BisectBudget::Auto);
170/// while let Some(range) = search.next_range() {
171///     // The caller owns the operation — one shared transaction, or a fresh
172///     // one committed per range.
173///     let verdict = probe(&items[range.clone()]).await;
174///     search.report(range, verdict)?;
175/// }
176/// let outcomes = search.into_outcomes();
177/// ```
178#[derive(Debug)]
179pub struct BisectSearch<E> {
180    pending: BinaryHeap<PendingRange>,
181    resolved: Vec<Option<ItemOutcome<E>>>,
182    probes_used: usize,
183    cap: usize,
184    transient_retries: usize,
185    max_transient_retries: usize,
186    last_error: Option<E>,
187}
188
189impl<E> BisectSearch<E> {
190    /// Starts a search over `n` items with the whole range as the first probe.
191    pub fn new(n: usize, budget: BisectBudget) -> Self {
192        let mut pending = BinaryHeap::new();
193        if n > 0 {
194            pending.push(PendingRange { start: 0, end: n });
195        }
196        Self {
197            pending,
198            resolved: (0..n).map(|_| None).collect(),
199            probes_used: 0,
200            cap: budget.effective_cap(n),
201            transient_retries: 0,
202            max_transient_retries: DEFAULT_MAX_TRANSIENT_RETRIES,
203            last_error: None,
204        }
205    }
206
207    /// Overrides [`DEFAULT_MAX_TRANSIENT_RETRIES`] for this search.
208    #[must_use]
209    pub fn with_max_transient_retries(mut self, max: usize) -> Self {
210        self.max_transient_retries = max;
211        self
212    }
213
214    /// The next range to probe: the largest outstanding one, earliest start
215    /// breaking ties.
216    ///
217    /// `None` when everything is resolved or the budget is spent — items still
218    /// inside unprobed ranges resolve as [`ItemOutcome::Unresolved`].
219    pub fn next_range(&mut self) -> Option<Range<usize>> {
220        if self.probes_used >= self.cap {
221            return None;
222        }
223        let range = self.pending.pop()?;
224        self.probes_used += 1;
225        Some(range.start..range.end)
226    }
227
228    /// Records what a probe found.
229    ///
230    /// Errors once the transient allowance is exhausted, at which point the
231    /// search is over — see [`TransientLimitExceeded`].
232    pub fn report(
233        &mut self,
234        range: Range<usize>,
235        verdict: ProbeVerdict<E>,
236    ) -> Result<(), TransientLimitExceeded> {
237        let range = PendingRange {
238            start: range.start,
239            end: range.end,
240        };
241
242        match verdict {
243            ProbeVerdict::Clean => {
244                for slot in &mut self.resolved[range.start..range.end] {
245                    *slot = Some(ItemOutcome::Complete);
246                }
247            }
248
249            ProbeVerdict::Transient(error) => {
250                self.last_error = Some(error);
251                // Refunded: the budget pays for bisection, and this probe
252                // produced none.
253                self.probes_used = self.probes_used.saturating_sub(1);
254                if self.transient_retries >= self.max_transient_retries {
255                    return Err(TransientLimitExceeded {
256                        probes_used: self.probes_used,
257                        transient_retries: self.transient_retries,
258                    });
259                }
260                self.transient_retries += 1;
261                self.pending.push(range);
262            }
263
264            // A single item that failed on its own probe: the error is
265            // attributable to it, so it owns it.
266            ProbeVerdict::Failed(error) if range.len() == 1 => {
267                self.resolved[range.start] = Some(ItemOutcome::Failed(error));
268            }
269
270            ProbeVerdict::Failed(error) => {
271                self.last_error = Some(error);
272                let mid = range.start + range.len() / 2;
273                self.pending.push(PendingRange {
274                    start: range.start,
275                    end: mid,
276                });
277                self.pending.push(PendingRange {
278                    start: mid,
279                    end: range.end,
280                });
281            }
282        }
283
284        Ok(())
285    }
286
287    /// Probes spent so far, net of refunds.
288    pub fn probes_used(&self) -> usize {
289        self.probes_used
290    }
291
292    /// Transient re-probes taken so far.
293    pub fn transient_retries(&self) -> usize {
294        self.transient_retries
295    }
296
297    /// The most recent error from a probe spanning more than one item.
298    pub fn last_error(&self) -> Option<&E> {
299        self.last_error.as_ref()
300    }
301
302    /// Finishes the search, resolving anything still outstanding as
303    /// [`ItemOutcome::Unresolved`] so every input gets exactly one outcome.
304    pub fn into_outcomes(self) -> BisectOutcomes<E> {
305        let items = self
306            .resolved
307            .into_iter()
308            .map(|slot| slot.unwrap_or(ItemOutcome::Unresolved))
309            .collect();
310
311        BisectOutcomes {
312            items,
313            probes_used: self.probes_used,
314            transient_retries: self.transient_retries,
315            last_error: self.last_error,
316        }
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    /// Runs a search where `culprits` are the item indices that fail, recording
325    /// the probe sequence. No database, no operation — the algorithm alone.
326    fn run(
327        n: usize,
328        budget: BisectBudget,
329        culprits: &[usize],
330    ) -> (Vec<Range<usize>>, BisectOutcomes<String>) {
331        let mut search = BisectSearch::new(n, budget);
332        let mut probes = Vec::new();
333        while let Some(range) = search.next_range() {
334            probes.push(range.clone());
335            let verdict = if culprits.iter().any(|c| range.contains(c)) {
336                ProbeVerdict::Failed(format!("bad {range:?}"))
337            } else {
338                ProbeVerdict::Clean
339            };
340            search.report(range, verdict).expect("no transients here");
341        }
342        (probes, search.into_outcomes())
343    }
344
345    #[test]
346    fn auto_budget_matches_the_documented_formula() {
347        assert_eq!(BisectBudget::Auto.effective_cap(1), 1);
348        assert_eq!(BisectBudget::Auto.effective_cap(10), 9);
349        assert_eq!(BisectBudget::Auto.effective_cap(25), 11);
350        assert_eq!(BisectBudget::Auto.effective_cap(100), 15);
351    }
352
353    #[test]
354    fn max_probes_zero_clamps_to_one() {
355        assert_eq!(BisectBudget::MaxProbes(0).effective_cap(50), 1);
356        assert_eq!(BisectBudget::MaxProbes(3).effective_cap(50), 3);
357    }
358
359    #[test]
360    fn a_clean_batch_probes_exactly_once() {
361        let (probes, outcomes) = run(5, BisectBudget::Auto, &[]);
362        assert_eq!(probes, vec![0..5]);
363        assert_eq!(outcomes.probes_used, 1);
364        assert!(
365            outcomes
366                .items
367                .iter()
368                .all(|o| matches!(o, ItemOutcome::Complete))
369        );
370    }
371
372    #[test]
373    fn ranges_are_probed_largest_first_earliest_start_on_ties() {
374        let (probes, _) = run(8, BisectBudget::FullResolution, &[0]);
375        // Whole slice, then the split halves largest-first with the earlier
376        // start winning the tie, narrowing onto index 0.
377        assert_eq!(probes[0], 0..8);
378        assert_eq!(probes[1], 0..4);
379        assert_eq!(probes[2], 4..8);
380        assert!(
381            probes.windows(2).all(|w| {
382                let (a, b) = (w[0].len(), w[1].len());
383                a > b || (a == b && w[0].start < w[1].start) || a < b
384            }),
385            "probe order was {probes:?}"
386        );
387    }
388
389    #[test]
390    fn a_single_culprit_is_isolated_and_its_siblings_are_salvaged() {
391        let (_, outcomes) = run(8, BisectBudget::Auto, &[3]);
392        for (idx, outcome) in outcomes.items.iter().enumerate() {
393            if idx == 3 {
394                assert!(
395                    matches!(outcome, ItemOutcome::Failed(_)),
396                    "index 3: {outcome:?}"
397                );
398            } else {
399                assert_eq!(outcome, &ItemOutcome::Complete, "index {idx}");
400            }
401        }
402    }
403
404    #[test]
405    fn scattered_culprits_do_not_poison_their_clean_siblings() {
406        let (_, outcomes) = run(16, BisectBudget::FullResolution, &[0, 8]);
407        let completed = outcomes
408            .items
409            .iter()
410            .filter(|o| matches!(o, ItemOutcome::Complete))
411            .count();
412        assert_eq!(completed, 14);
413        assert!(matches!(outcomes.items[0], ItemOutcome::Failed(_)));
414        assert!(matches!(outcomes.items[8], ItemOutcome::Failed(_)));
415    }
416
417    #[test]
418    fn full_resolution_resolves_every_item_of_an_all_bad_batch() {
419        let (_, outcomes) = run(6, BisectBudget::FullResolution, &[0, 1, 2, 3, 4, 5]);
420        assert!(
421            outcomes
422                .items
423                .iter()
424                .all(|o| matches!(o, ItemOutcome::Failed(_)))
425        );
426    }
427
428    #[test]
429    fn max_probes_one_is_equivalent_to_resolving_nothing() {
430        let (probes, outcomes) = run(5, BisectBudget::MaxProbes(1), &[2]);
431        assert_eq!(probes, vec![0..5]);
432        assert!(
433            outcomes
434                .items
435                .iter()
436                .all(|o| matches!(o, ItemOutcome::Unresolved))
437        );
438        assert!(
439            outcomes.last_error.is_some(),
440            "the batch-level error is kept for the caller"
441        );
442    }
443
444    #[test]
445    fn every_input_gets_exactly_one_outcome_even_under_budget_exhaustion() {
446        let (_, outcomes) = run(10, BisectBudget::MaxProbes(3), &[0]);
447        assert_eq!(outcomes.items.len(), 10);
448    }
449
450    #[test]
451    fn a_transient_probe_is_re_run_whole_and_refunded() {
452        let mut search: BisectSearch<String> = BisectSearch::new(8, BisectBudget::MaxProbes(1));
453
454        let first = search.next_range().expect("a first probe");
455        assert_eq!(first, 0..8);
456        search
457            .report(first, ProbeVerdict::Transient("40P01".into()))
458            .expect("within the allowance");
459
460        // Refunded, so the budget of 1 still admits a probe, and it covers the
461        // same range.
462        let retry = search.next_range().expect("the refunded re-probe");
463        assert_eq!(retry, 0..8);
464        assert_eq!(search.transient_retries(), 1);
465
466        search.report(retry, ProbeVerdict::Clean).expect("clean");
467        let outcomes = search.into_outcomes();
468        assert_eq!(outcomes.probes_used, 1);
469        assert_eq!(outcomes.transient_retries, 1);
470        assert!(
471            outcomes
472                .items
473                .iter()
474                .all(|o| matches!(o, ItemOutcome::Complete))
475        );
476    }
477
478    #[test]
479    fn the_transient_allowance_is_bounded() {
480        let mut search: BisectSearch<String> =
481            BisectSearch::new(4, BisectBudget::FullResolution).with_max_transient_retries(2);
482
483        for _ in 0..2 {
484            let range = search.next_range().expect("a probe");
485            search
486                .report(range, ProbeVerdict::Transient("40001".into()))
487                .expect("within the allowance");
488        }
489
490        let range = search.next_range().expect("a probe");
491        let limit = search
492            .report(range, ProbeVerdict::Transient("40001".into()))
493            .expect_err("the allowance is spent");
494        assert_eq!(limit.transient_retries, 2);
495    }
496
497    #[test]
498    fn an_empty_batch_probes_nothing() {
499        let (probes, outcomes) = run(0, BisectBudget::Auto, &[]);
500        assert!(probes.is_empty());
501        assert!(outcomes.items.is_empty());
502    }
503}
es_entity/operation/batch/search.rs

es_entity/operation/batch/
search.rs
