polars_core/
runtime.rs

1use std::cell::RefCell;
2use std::panic::{AssertUnwindSafe, catch_unwind};
3use std::sync::LazyLock;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::time::Duration;
6
7use polars_utils::with_drop::WithDrop;
8use rayon::{ThreadPool, ThreadPoolBuilder, Yield};
9
10pub struct RAYON;
11
12// Thread locals to allow disabling threading for specific threads.
13#[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
14thread_local! {
15    static NOOP_POOL: RefCell<ThreadPool> = RefCell::new(
16        ThreadPoolBuilder::new()
17            .use_current_thread()
18            .num_threads(1)
19            .build()
20            .expect("could not create no-op thread pool")
21    );
22}
23
24impl RAYON {
25    pub fn install<OP, R>(&self, op: OP) -> R
26    where
27        OP: FnOnce() -> R + Send,
28        R: Send,
29    {
30        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
31        {
32            op()
33        }
34
35        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
36        {
37            let mut op = Some(op);
38            let mut out = None;
39            self.install_dyn(&mut || out = Some(op.take().unwrap()()));
40            out.unwrap()
41        }
42    }
43
44    // Instantiates the rayon install machinery only once instead of for every `install` caller.
45    #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
46    #[inline(never)]
47    fn install_dyn(&self, op: &mut (dyn FnMut() + Send)) {
48        self.with(|p| p.install(op))
49    }
50
51    pub fn join<A, B, RA, RB>(&self, oper_a: A, oper_b: B) -> (RA, RB)
52    where
53        A: FnOnce() -> RA + Send,
54        B: FnOnce() -> RB + Send,
55        RA: Send,
56        RB: Send,
57    {
58        self.install(|| rayon::join(oper_a, oper_b))
59    }
60
61    pub fn scope<'scope, OP, R>(&self, op: OP) -> R
62    where
63        OP: FnOnce(&rayon::Scope<'scope>) -> R + Send,
64        R: Send,
65    {
66        self.install(|| rayon::scope(op))
67    }
68
69    pub fn spawn<OP>(&self, op: OP)
70    where
71        OP: FnOnce() + Send + 'static,
72    {
73        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
74        {
75            rayon::spawn(op)
76        }
77
78        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
79        {
80            self.with(|p| {
81                p.spawn(op);
82                if p.current_num_threads() == 1 {
83                    p.yield_now();
84                }
85            })
86        }
87    }
88
89    pub fn spawn_fifo<OP>(&self, op: OP)
90    where
91        OP: FnOnce() + Send + 'static,
92    {
93        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
94        {
95            rayon::spawn_fifo(op)
96        }
97
98        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
99        {
100            self.with(|p| {
101                p.spawn_fifo(op);
102                if p.current_num_threads() == 1 {
103                    p.yield_now();
104                }
105            })
106        }
107    }
108
109    pub fn current_thread_has_pending_tasks(&self) -> Option<bool> {
110        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
111        {
112            None
113        }
114
115        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
116        {
117            self.with(|p| p.current_thread_has_pending_tasks())
118        }
119    }
120
121    pub fn current_thread_index(&self) -> Option<usize> {
122        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
123        {
124            rayon::current_thread_index()
125        }
126
127        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
128        {
129            self.with(|p| p.current_thread_index())
130        }
131    }
132
133    pub fn current_num_threads(&self) -> usize {
134        #[cfg(not(any(target_os = "emscripten", not(target_family = "wasm"))))]
135        {
136            rayon::current_num_threads()
137        }
138
139        #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
140        {
141            self.with(|p| p.current_num_threads())
142        }
143    }
144
145    #[cfg(any(target_os = "emscripten", not(target_family = "wasm")))]
146    pub fn with<OP, R>(&self, op: OP) -> R
147    where
148        OP: FnOnce(&ThreadPool) -> R + Send,
149        R: Send,
150    {
151        if polars_async::executor::ALLOW_RAYON_THREADS.get()
152            || THREAD_POOL.current_thread_index().is_some()
153        {
154            op(&THREAD_POOL)
155        } else {
156            NOOP_POOL.with(|v| op(&v.borrow()))
157        }
158    }
159
160    /// Calls a blocking function without blocking the rayon thread pool by
161    /// moving it to a different thread.
162    ///
163    /// If this thread isn't a rayon thread this simply calls f directly.
164    pub fn block_on<R: Send, F: FnOnce() -> R + Send>(&self, f: F) -> R {
165        if THREAD_POOL.current_thread_index().is_some() {
166            // We wrap F such that the resulting closure no longer has F or R as
167            // parameters, and whose only limitation is that it isn't 'static.
168            let mut opt_f: Option<F> = Some(f);
169            let mut out: Option<std::thread::Result<R>> = None;
170            let done = AtomicBool::new(false);
171            let waiter = std::thread::current();
172            let mut wrap_f = || {
173                let f = AssertUnwindSafe(opt_f.take().unwrap());
174                out = Some(catch_unwind(f));
175                done.store(true, Ordering::Release);
176                waiter.unpark();
177            };
178
179            // SAFETY: we always await the future to completion before returning from here, meaning
180            // wrap_f stays alive for as long as it needs to. If for whatever reason we unwind we
181            // abort.
182            let abort = WithDrop::new((), |()| std::process::abort());
183            let ref_wrap_f: &mut (dyn Send + FnMut()) = &mut wrap_f;
184            let static_wrap_f: &'static mut (dyn Send + FnMut() + 'static) =
185                unsafe { core::mem::transmute(ref_wrap_f) };
186            let join_handle = ASYNC.spawn_blocking(static_wrap_f);
187
188            while !join_handle.is_finished() {
189                match rayon::yield_now() {
190                    Some(Yield::Executed) => {},
191                    // Sleep until `f` is done, but wake up now and then to run new Rayon jobs.
192                    Some(Yield::Idle) if !done.load(Ordering::Acquire) => {
193                        std::thread::park_timeout(Duration::from_millis(1))
194                    },
195                    Some(Yield::Idle) => std::thread::yield_now(),
196                    None => unreachable!(),
197                }
198            }
199
200            WithDrop::dismiss(abort);
201            match out.unwrap() {
202                Ok(v) => v,
203                Err(panic) => std::panic::resume_unwind(panic),
204            }
205        } else {
206            f()
207        }
208    }
209}
210
211// this is re-exported in utils for polars child crates
212#[cfg(not(target_family = "wasm"))] // only use this on non wasm targets
213pub static THREAD_POOL: LazyLock<ThreadPool> = LazyLock::new(|| {
214    let thread_name = std::env::var("POLARS_THREAD_NAME").unwrap_or_else(|_| "polars".to_string());
215    ThreadPoolBuilder::new()
216        .num_threads(polars_config::config().max_threads())
217        .thread_name(move |i| format!("{thread_name}-{i}"))
218        .build()
219        .expect("could not spawn threads")
220});
221
222#[cfg(all(target_os = "emscripten", target_family = "wasm"))] // Use 1 rayon thread on emscripten
223pub static THREAD_POOL: LazyLock<ThreadPool> = LazyLock::new(|| {
224    ThreadPoolBuilder::new()
225        .num_threads(1)
226        .use_current_thread()
227        .build()
228        .expect("could not create pool")
229});
230
231pub use polars_async::ASYNC;
polars_core/runtime.rs

polars_core/
runtime.rs
