rushdb_storage/
parallel.rs

1//! Intra-query parallelism under one process-wide thread budget.
2//!
3//! A long read (an aggregate over a million records, an exact vector scan)
4//! can split its key range into parts and run them on scoped threads. Every
5//! query draws helper threads from one shared budget, sized to the server's
6//! worker threads (`RUSHDB_THREADS`), so concurrent queries never run more
7//! threads than the host was given: a query that finds the budget spent runs
8//! its parts on its own thread. Parts are numbered and their results are
9//! returned in part order, so a caller that merges them in order gets the
10//! same answer as a sequential run whatever the timing.
11
12use std::sync::atomic::{AtomicUsize, Ordering};
13
14/// Helper threads not currently lent to a query.
15static FREE: AtomicUsize = AtomicUsize::new(0);
16/// The configured budget (for diagnostics and tests).
17static BUDGET: AtomicUsize = AtomicUsize::new(0);
18
19/// Set the helper-thread budget: the host's worker threads less the one
20/// each query already runs on. Called once at startup; `0` = run every part
21/// on the calling thread.
22pub fn set_threads(threads: usize) {
23    let helpers = threads.saturating_sub(1);
24    let previous = BUDGET.swap(helpers, Ordering::SeqCst);
25    if helpers >= previous {
26        FREE.fetch_add(helpers - previous, Ordering::SeqCst);
27    } else {
28        let _ = FREE.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |free| {
29            Some(free.saturating_sub(previous - helpers))
30        });
31    }
32}
33
34/// The helper-thread budget [`set_threads`] configured.
35#[must_use]
36pub fn budget() -> usize {
37    BUDGET.load(Ordering::SeqCst)
38}
39
40/// Helper threads lent to one call, returned when dropped.
41struct Lease(usize);
42
43impl Lease {
44    fn take(wanted: usize) -> Self {
45        let mut taken = 0;
46        let _ = FREE.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |free| {
47            taken = free.min(wanted);
48            Some(free - taken)
49        });
50        Self(taken)
51    }
52}
53
54impl Drop for Lease {
55    fn drop(&mut self) {
56        FREE.fetch_add(self.0, Ordering::SeqCst);
57    }
58}
59
60/// Run `part(0..parts)` and return the results in part order. Parts run on
61/// the calling thread plus as many helper threads as the shared budget lends
62/// right now (at most `parts - 1`); with none free, they run one after
63/// another here. The first error (by part order) is returned.
64///
65/// # Errors
66/// Returns the error of the lowest-numbered failing part.
67///
68/// # Panics
69/// Re-raises a panic of a part; a part's slot is always filled before the
70/// threads join.
71pub fn run_parts<T, E, F>(parts: usize, part: F) -> Result<Vec<T>, E>
72where
73    T: Send,
74    E: Send,
75    F: Fn(usize) -> Result<T, E> + Sync,
76{
77    if parts <= 1 {
78        return (0..parts).map(&part).collect();
79    }
80    let lease = Lease::take(parts - 1);
81    if lease.0 == 0 {
82        return (0..parts).map(&part).collect();
83    }
84    let next = AtomicUsize::new(0);
85    let workers = lease.0 + 1;
86    let mut slots: Vec<Option<Result<T, E>>> = (0..parts).map(|_| None).collect();
87    let results = std::sync::Mutex::new(&mut slots);
88    let work = || loop {
89        let index = next.fetch_add(1, Ordering::Relaxed);
90        if index >= parts {
91            break;
92        }
93        let outcome = part(index);
94        let mut guard = results
95            .lock()
96            .unwrap_or_else(std::sync::PoisonError::into_inner);
97        guard[index] = Some(outcome);
98    };
99    // Reads a helper does count toward the caller's per-thread read
100    // accounting, so a query's tally does not depend on which thread ran
101    // which part.
102    let helper_reads = std::thread::scope(|scope| {
103        // Every helper starts before the caller takes its share.
104        let mut helpers = Vec::with_capacity(workers - 1);
105        for _ in 1..workers {
106            helpers.push(scope.spawn(|| {
107                work();
108                crate::read_accounting::take_local()
109            }));
110        }
111        work();
112        helpers
113            .into_iter()
114            .map(|helper| {
115                helper
116                    .join()
117                    .unwrap_or_else(|panic| std::panic::resume_unwind(panic))
118            })
119            .collect::<Vec<_>>()
120    });
121    for reads in helper_reads {
122        crate::read_accounting::add_local(reads);
123    }
124    drop(lease);
125    slots
126        .into_iter()
127        .map(|slot| slot.expect("every part ran"))
128        .collect()
129}
130
131/// Split `0..len` into at most `parts` contiguous ranges of near-equal size,
132/// in order (no empty ranges).
133#[must_use]
134pub fn split(len: u64, parts: usize) -> Vec<std::ops::Range<u64>> {
135    let parts = u64::try_from(parts.max(1)).unwrap_or(1).min(len.max(1));
136    let size = len.div_ceil(parts);
137    (0..parts)
138        .map(|i| (i * size)..((i + 1) * size).min(len))
139        .filter(|range| !range.is_empty())
140        .collect()
141}
142
143/// How many parts a job of `units` should split into: one per helper the
144/// budget has plus the caller, and none smaller than `min_units`.
145#[must_use]
146pub fn parts_for(units: u64, min_units: u64) -> usize {
147    let by_size = usize::try_from(units / min_units.max(1)).unwrap_or(usize::MAX);
148    by_size.clamp(1, budget() + 1)
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    /// The budget is process-wide: tests that lend from it run one at a
156    /// time.
157    static BUDGET_TESTS: std::sync::Mutex<()> = std::sync::Mutex::new(());
158
159    fn exclusive() -> std::sync::MutexGuard<'static, ()> {
160        BUDGET_TESTS
161            .lock()
162            .unwrap_or_else(std::sync::PoisonError::into_inner)
163    }
164
165    #[test]
166    fn parts_return_in_order_and_errors_by_lowest_part() {
167        let _serial = exclusive();
168        set_threads(4);
169        let out: Result<Vec<usize>, ()> = run_parts(16, |i| Ok(i * 2));
170        assert_eq!(out.unwrap(), (0..16).map(|i| i * 2).collect::<Vec<_>>());
171        let failed: Result<Vec<usize>, usize> =
172            run_parts(16, |i| if i % 5 == 3 { Err(i) } else { Ok(i) });
173        assert_eq!(failed, Err(3));
174        // Every lent thread came back.
175        assert_eq!(FREE.load(Ordering::SeqCst), budget());
176    }
177
178    #[test]
179    fn reads_on_helper_threads_count_for_the_caller() {
180        let _serial = exclusive();
181        set_threads(4);
182        crate::read_accounting::enable(true);
183        let _ = crate::read_accounting::take_local();
184        let parts: Result<Vec<()>, ()> = run_parts(8, |part| {
185            crate::read_accounting::record("get", &[0x01, 0, 0, 0, 1, u8::try_from(part).unwrap()]);
186            Ok(())
187        });
188        parts.unwrap();
189        let counted: u64 = crate::read_accounting::take_local().values().sum();
190        assert_eq!(counted, 8, "every part's read is the caller's");
191    }
192
193    #[test]
194    fn split_covers_the_range_without_gaps() {
195        let ranges = split(10, 3);
196        assert_eq!(ranges, vec![0..4, 4..8, 8..10]);
197        assert_eq!(split(2, 8), vec![0..1, 1..2]);
198        assert!(split(0, 4).is_empty());
199    }
200}
rushdb_storage/parallel.rs

rushdb_storage/
parallel.rs
