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