vortex_io/runtime/
pool.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5use std::sync::atomic::AtomicBool;
6use std::sync::atomic::Ordering;
7use std::time::Duration;
8
9use parking_lot::Mutex;
10use smol::block_on;
11use vortex_error::VortexExpect;
12use vortex_utils::parallelism::get_available_parallelism;
13
14use crate::runtime::smol::SmolExecutor;
15
16#[derive(Clone)]
17pub struct CurrentThreadWorkerPool {
18    executor: Arc<SmolExecutor>,
19    state: Arc<Mutex<PoolState>>,
20}
21
22impl CurrentThreadWorkerPool {
23    pub(super) fn new(executor: Arc<SmolExecutor>) -> Self {
24        Self {
25            executor,
26            state: Arc::new(Mutex::new(PoolState::default())),
27        }
28    }
29
30    /// Set the number of worker threads to the available system parallelism as reported by
31    /// [`get_available_parallelism()`] minus 1, to leave a slot open for the calling thread.
32    pub fn set_workers_to_available_parallelism(&self) {
33        let n = get_available_parallelism()
34            .map(|n| n.saturating_sub(1).max(1))
35            .unwrap_or(1);
36        self.set_workers(n);
37    }
38
39    /// Set the number of worker threads
40    /// - If n > current: spawns additional workers
41    /// - If n < current: signals extra workers to shut down
42    pub fn set_workers(&self, n: usize) {
43        let mut state = self.state.lock();
44        let current = state.workers.len();
45
46        if n > current {
47            // Spawn new workers
48            for _ in current..n {
49                let shutdown = Arc::new(AtomicBool::new(false));
50                let executor = Arc::clone(&self.executor);
51                let shutdown_clone = Arc::clone(&shutdown);
52
53                std::thread::Builder::new()
54                    .name("vortex-current-thread-worker".to_string())
55                    .spawn(move || {
56                        // Run the executor with a sleeping future that checks for shutdown
57                        block_on(executor.run(async move {
58                            while !shutdown_clone.load(Ordering::Relaxed) {
59                                smol::Timer::after(Duration::from_millis(100)).await;
60                            }
61                        }))
62                    })
63                    .vortex_expect("Failed to spawn current thread worker");
64
65                state.workers.push(WorkerHandle { shutdown });
66            }
67        } else if n < current {
68            // Signal extra workers to shutdown and remove them
69            while state.workers.len() > n {
70                if let Some(worker) = state.workers.pop() {
71                    worker.shutdown.store(true, Ordering::Relaxed);
72                }
73            }
74        }
75    }
76
77    /// Get the current number of worker threads
78    pub fn worker_count(&self) -> usize {
79        self.state.lock().workers.len()
80    }
81}
82
83#[derive(Default)]
84struct PoolState {
85    /// The set of worker handles for the background threads.
86    workers: Vec<WorkerHandle>,
87}
88
89struct WorkerHandle {
90    /// The shutdown flag indicating that the worker should stop.
91    shutdown: Arc<AtomicBool>,
92}
93
94impl Drop for CurrentThreadWorkerPool {
95    fn drop(&mut self) {
96        let mut state = self.state.lock();
97
98        // Signal all workers to shut down
99        for worker in state.workers.drain(..) {
100            worker.shutdown.store(true, Ordering::Relaxed);
101        }
102    }
103}
vortex_io/runtime/pool.rs

vortex_io/runtime/
pool.rs
