1mod map;
17mod policy;
18mod pool;
19
20pub(crate) use pool::{run_ordered_batches, run_ordered_pool};
21
22pub use map::{
23 par_map, parallel_worthwhile, try_par_map, try_par_map_force, try_par_map_with_local,
24};
25pub use policy::{ParallelOp, ParallelOptions, ThreadPolicy};
26
27use std::num::NonZeroUsize;
28
29#[derive(Debug, Clone, PartialEq, Eq)]
34pub enum ParallelHashError<E> {
35 Task(E),
37}
38
39impl<E: std::fmt::Display> std::fmt::Display for ParallelHashError<E> {
40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 match self {
42 Self::Task(e) => write!(f, "{e}"),
43 }
44 }
45}
46
47impl<E: std::fmt::Debug + std::fmt::Display> std::error::Error for ParallelHashError<E> {}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct Parallelism {
55 threads: NonZeroUsize,
56}
57
58impl Parallelism {
59 #[must_use]
61 pub fn new(threads: NonZeroUsize) -> Self {
62 Self { threads }
63 }
64
65 #[must_use]
67 pub fn resolve(requested: Option<usize>) -> Self {
68 let avail = std::thread::available_parallelism()
69 .map(NonZeroUsize::get)
70 .unwrap_or(1);
71 let n = match requested {
72 None | Some(0) => avail,
73 Some(n) => n.min(avail),
74 };
75 Self {
76 threads: NonZeroUsize::new(n.max(1)).unwrap_or(NonZeroUsize::MIN),
77 }
78 }
79
80 #[must_use]
82 pub fn threads(self) -> NonZeroUsize {
83 self.threads
84 }
85}
86
87#[must_use]
89pub fn index_parallelism_from_config(config: &crate::config::ConfigSet) -> Parallelism {
90 ThreadPolicy::default()
91 .resolve(ParallelOp::IndexRefresh, config)
92 .parallelism()
93}
94
95pub fn index_pack_threads_for_ingest(
100 policy: ThreadPolicy,
101 config: &crate::config::ConfigSet,
102 explicit: Option<usize>,
103) -> Result<Option<usize>, crate::error::ConfigError> {
104 if let Some(n) = explicit {
105 return Ok(Some(n));
106 }
107 if policy.threads_override.is_some() {
108 return Ok(Some(
109 policy
110 .try_resolve(ParallelOp::IndexPack, config)?
111 .threads
112 .get(),
113 ));
114 }
115 Ok(None)
116}
117
118#[must_use]
120pub fn index_load_lazy_thread_count(
121 policy: &ThreadPolicy,
122 config: &crate::config::ConfigSet,
123 cache_nr: usize,
124 try_threaded: bool,
125) -> usize {
126 if !try_threaded {
127 return 0;
128 }
129 let opts = policy.resolve(ParallelOp::IndexLoad, config);
130 if opts.threads.get() <= 1 {
131 return 0;
132 }
133 if cache_nr < 2 * crate::index_name_hash_lazy::LAZY_THREAD_COST {
134 return 0;
135 }
136 let mut nr = opts.threads.get();
137 if cache_nr < nr * crate::index_name_hash_lazy::LAZY_THREAD_COST {
138 nr = cache_nr / crate::index_name_hash_lazy::LAZY_THREAD_COST;
139 }
140 nr
141}
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146 use std::sync::atomic::{AtomicUsize, Ordering};
147 use std::sync::Arc;
148
149 fn threads(n: usize) -> NonZeroUsize {
150 NonZeroUsize::new(n).expect("non-zero threads")
151 }
152
153 #[test]
154 fn try_par_map_preserves_order() {
155 let items: Vec<usize> = (0..256).collect();
156 let opts = ParallelOptions {
157 threads: threads(4),
158 min_items: 0,
159 min_total_bytes: None,
160 signed_entry_threshold: None,
161 };
162 let got = par_map(&items, opts, items.len(), |&x| x * 2);
163 let expect: Vec<_> = items.iter().map(|x| x * 2).collect();
164 assert_eq!(got, expect);
165 }
166
167 #[test]
168 fn try_par_map_surfaces_task_error() {
169 let items: Vec<i32> = (0..128).collect();
170 let opts = ParallelOptions {
171 threads: threads(4),
172 min_items: 0,
173 min_total_bytes: None,
174 signed_entry_threshold: None,
175 };
176 let err = try_par_map(&items, opts, items.len(), |&x| {
177 if x == 64 {
178 Err("boom")
179 } else {
180 Ok(x)
181 }
182 })
183 .unwrap_err();
184 assert_eq!(err, ParallelHashError::Task("boom"));
185 }
186
187 #[test]
188 fn ordered_batches_consume_in_item_order_with_worker_state() {
189 let items: Vec<usize> = (0..1000).collect();
190 let inits = AtomicUsize::new(0);
191 let mut seen = Vec::new();
192 let mut batch_lens = Vec::new();
193 run_ordered_batches(
194 &items,
195 7,
196 4,
197 || {
198 inits.fetch_add(1, Ordering::Relaxed);
199 0usize
200 },
201 |mapped: &mut usize, &x| -> Result<usize, ()> {
202 *mapped += 1;
203 Ok(x * 3)
204 },
205 |batch, results| {
206 assert_eq!(batch.len(), results.len());
207 batch_lens.push(batch.len());
208 seen.extend(results);
209 Ok(())
210 },
211 )
212 .expect("batches");
213 assert_eq!(seen, items.iter().map(|x| x * 3).collect::<Vec<_>>());
214 assert_eq!(batch_lens.iter().sum::<usize>(), items.len());
215 assert!(batch_lens[..batch_lens.len() - 1].iter().all(|&n| n == 7));
216 assert_eq!(
217 inits.load(Ordering::Relaxed),
218 4,
219 "init runs once per worker"
220 );
221 }
222
223 #[test]
224 fn ordered_batches_surface_map_and_consume_errors() {
225 let items: Vec<i32> = (0..500).collect();
226 let err = run_ordered_batches(
227 &items,
228 10,
229 3,
230 || (),
231 |_, &x| if x == 250 { Err("map") } else { Ok(x) },
232 |_, _| Ok(()),
233 )
234 .unwrap_err();
235 assert_eq!(err, "map");
236
237 let mut consumed = 0;
238 let err = run_ordered_batches(
239 &items,
240 10,
241 3,
242 || (),
243 |_, &x| Ok::<_, &str>(x),
244 |_, _| {
245 consumed += 1;
246 if consumed == 2 {
247 Err("consume")
248 } else {
249 Ok(())
250 }
251 },
252 )
253 .unwrap_err();
254 assert_eq!(err, "consume");
255 assert_eq!(consumed, 2, "no batch is consumed after an error");
256
257 run_ordered_batches(
258 &[] as &[i32],
259 10,
260 3,
261 || (),
262 |_, &x| Ok::<_, ()>(x),
263 |_, _| panic!("nothing to consume"),
264 )
265 .expect("empty input");
266 }
267
268 #[test]
269 fn ordered_batches_resume_worker_panics() {
270 let items: Vec<i32> = (0..200).collect();
271 let outcome = std::panic::catch_unwind(|| {
272 run_ordered_batches(
273 &items,
274 5,
275 4,
276 || (),
277 |_, &x| -> Result<i32, ()> {
278 assert!(x != 120, "worker boom");
279 Ok(x)
280 },
281 |_, _| Ok(()),
282 )
283 });
284 assert!(outcome.is_err());
285 }
286
287 #[test]
288 fn index_pack_threads_for_ingest_honors_override_only() {
289 let cfg = config_from("[pack]\n\tthreads = 2\n");
290 assert_eq!(
291 index_pack_threads_for_ingest(ThreadPolicy::default(), &cfg, None).expect("ingest"),
292 None
293 );
294 assert_eq!(
295 index_pack_threads_for_ingest(ThreadPolicy::default(), &cfg, Some(3)).expect("ingest"),
296 Some(3)
297 );
298 let policy = ThreadPolicy::with_threads_override(Some(6));
299 assert_eq!(
300 index_pack_threads_for_ingest(policy, &cfg, None).expect("ingest"),
301 Some(policy.resolve(ParallelOp::IndexPack, &cfg).threads.get())
302 );
303 }
304
305 fn config_from(text: &str) -> crate::config::ConfigSet {
306 use crate::config::{ConfigFile, ConfigScope, ConfigSet};
307 use std::path::Path;
308 let file = ConfigFile::parse(Path::new("test"), text, ConfigScope::Local).expect("parse");
309 let mut set = ConfigSet::new();
310 set.merge(&file);
311 set
312 }
313
314 #[test]
315 fn try_par_map_with_local_reuses_per_worker_state() {
316 let items: Vec<i32> = (0..256).collect();
317 let opts = ParallelOptions {
318 threads: threads(4),
319 min_items: 0,
320 min_total_bytes: None,
321 signed_entry_threshold: None,
322 };
323 let init_count = Arc::new(AtomicUsize::new(0));
324 let init_count_for_init = Arc::clone(&init_count);
325 let got = try_par_map_with_local(
326 &items,
327 opts,
328 items.len(),
329 move || {
330 init_count_for_init.fetch_add(1, Ordering::SeqCst);
331 0_i32
332 },
333 |state, &x| -> Result<i32, ()> {
334 *state += 1;
335 Ok(x + *state)
336 },
337 )
338 .expect("map");
339 assert_eq!(got.len(), items.len());
340 let inits = init_count.load(Ordering::SeqCst);
341 assert!(
342 inits >= 2,
343 "expected multiple workers to init state, got {inits}"
344 );
345 assert!(
346 inits <= 4,
347 "init count should not exceed worker count, got {inits}"
348 );
349 }
350}