|
| 1 | +use std::sync::OnceLock; |
| 2 | + |
1 | 3 | /// Initialize Rayon global thread pool with specified number of threads. |
2 | 4 | /// |
3 | 5 | /// If `--threads` option is not used, or `--threads 0` is given, |
4 | 6 | /// default to the number of available CPU cores. |
5 | 7 | /// |
| 8 | +/// Repeated calls reuse the pool and thread count from the first call. |
| 9 | +/// |
6 | 10 | /// # Panics |
7 | | -/// Panics if the global thread pool has already been initialized. |
| 11 | +/// Panics if another consumer has already initialized the global thread pool. |
8 | 12 | #[expect(clippy::print_stderr)] |
9 | 13 | pub fn init_rayon(threads: Option<usize>) { |
10 | | - // Always initialize thread pool, even if using default thread count, |
11 | | - // to ensure thread pool's thread count is locked after this point. |
12 | | - // `rayon::current_num_threads()` will always return the same number after this point. |
13 | | - // |
14 | | - // If you don't initialize the global thread pool explicitly, or don't specify `num_threads`, |
15 | | - // Rayon will initialize the thread pool when it's first used, with a thread count of |
16 | | - // `std::thread::available_parallelism()`, and that thread count won't change thereafter. |
17 | | - // So we don't *need* to initialize the thread pool here if we just want the default thread count. |
18 | | - // |
19 | | - // However, Rayon's docs state that: |
20 | | - // > In the future, the default behavior may change to dynamically add or remove threads as needed. |
21 | | - // https://docs.rs/rayon/1.11.0/rayon/struct.ThreadPoolBuilder.html#method.num_threads |
22 | | - // |
23 | | - // To ensure we continue to have a "locked" thread count, even after future Rayon upgrades, |
24 | | - // we always initialize the thread pool and explicitly specify thread count here. |
| 14 | + static RAYON_INIT: OnceLock<()> = OnceLock::new(); |
| 15 | + RAYON_INIT.get_or_init(|| { |
| 16 | + // Always initialize thread pool, even if using default thread count, |
| 17 | + // to ensure thread pool's thread count is locked after this point. |
| 18 | + // `rayon::current_num_threads()` will always return the same number after this point. |
| 19 | + // |
| 20 | + // If you don't initialize the global thread pool explicitly, or don't specify `num_threads`, |
| 21 | + // Rayon will initialize the thread pool when it's first used, with a thread count of |
| 22 | + // `std::thread::available_parallelism()`, and that thread count won't change thereafter. |
| 23 | + // So we don't *need* to initialize the thread pool here if we just want the default thread count. |
| 24 | + // |
| 25 | + // However, Rayon's docs state that: |
| 26 | + // > In the future, the default behavior may change to dynamically add or remove threads as needed. |
| 27 | + // https://docs.rs/rayon/1.11.0/rayon/struct.ThreadPoolBuilder.html#method.num_threads |
| 28 | + // |
| 29 | + // To ensure we continue to have a "locked" thread count, even after future Rayon upgrades, |
| 30 | + // we always initialize the thread pool and explicitly specify thread count here. |
25 | 31 |
|
26 | | - let thread_count = if let Some(thread_count) = threads |
27 | | - && thread_count > 0 |
28 | | - { |
29 | | - thread_count |
30 | | - } else if let Ok(thread_count) = std::thread::available_parallelism() { |
31 | | - thread_count.get() |
32 | | - } else { |
33 | | - eprintln!( |
34 | | - "Unable to determine available thread count. Defaulting to 1.\nConsider specifying the number of threads explicitly with `--threads` option." |
35 | | - ); |
36 | | - 1 |
37 | | - }; |
| 32 | + let thread_count = if let Some(thread_count) = threads |
| 33 | + && thread_count > 0 |
| 34 | + { |
| 35 | + thread_count |
| 36 | + } else if let Ok(thread_count) = std::thread::available_parallelism() { |
| 37 | + thread_count.get() |
| 38 | + } else { |
| 39 | + eprintln!( |
| 40 | + "Unable to determine available thread count. Defaulting to 1.\nConsider specifying the number of threads explicitly with `--threads` option." |
| 41 | + ); |
| 42 | + 1 |
| 43 | + }; |
38 | 44 |
|
39 | | - rayon::ThreadPoolBuilder::new().num_threads(thread_count).build_global().unwrap(); |
| 45 | + rayon::ThreadPoolBuilder::new().num_threads(thread_count).build_global().unwrap(); |
| 46 | + }); |
40 | 47 | } |
0 commit comments