portalis_transpiler/wasi_threading/
pool.rs

1//! Thread Pool Implementation
2//!
3//! Provides fixed and dynamic thread pools with work stealing and task execution.
4
5use anyhow::{Result, Context, anyhow};
6use std::sync::Arc;
7use std::time::Duration;
8use super::{WasiQueue, ThreadingError, WasiThread, ThreadHandle};
9
10#[cfg(not(target_arch = "wasm32"))]
11use rayon;
12
13/// Thread pool configuration
14#[derive(Debug, Clone)]
15pub struct ThreadPoolConfig {
16    /// Number of worker threads (None = number of CPUs)
17    pub num_threads: Option<usize>,
18    /// Thread name prefix
19    pub thread_name_prefix: Option<String>,
20    /// Stack size per thread
21    pub stack_size: Option<usize>,
22    /// Maximum pending tasks (None = unbounded)
23    pub max_pending_tasks: Option<usize>,
24    /// Enable work stealing (rayon-based)
25    pub enable_work_stealing: bool,
26}
27
28impl ThreadPoolConfig {
29    /// Create a new thread pool configuration
30    pub fn new() -> Self {
31        Self {
32            num_threads: None,
33            thread_name_prefix: Some("worker".to_string()),
34            stack_size: None,
35            max_pending_tasks: None,
36            enable_work_stealing: true,
37        }
38    }
39
40    /// Set the number of worker threads
41    pub fn num_threads(mut self, n: usize) -> Self {
42        self.num_threads = Some(n);
43        self
44    }
45
46    /// Set the thread name prefix
47    pub fn thread_name_prefix(mut self, prefix: impl Into<String>) -> Self {
48        self.thread_name_prefix = Some(prefix.into());
49        self
50    }
51
52    /// Set the stack size per thread
53    pub fn stack_size(mut self, size: usize) -> Self {
54        self.stack_size = Some(size);
55        self
56    }
57
58    /// Set the maximum pending tasks
59    pub fn max_pending_tasks(mut self, max: usize) -> Self {
60        self.max_pending_tasks = Some(max);
61        self
62    }
63
64    /// Enable or disable work stealing
65    pub fn enable_work_stealing(mut self, enabled: bool) -> Self {
66        self.enable_work_stealing = enabled;
67        self
68    }
69}
70
71impl Default for ThreadPoolConfig {
72    fn default() -> Self {
73        Self::new()
74    }
75}
76
77/// Thread pool builder
78pub struct ThreadPoolBuilder {
79    config: ThreadPoolConfig,
80}
81
82impl ThreadPoolBuilder {
83    /// Create a new thread pool builder
84    pub fn new() -> Self {
85        Self {
86            config: ThreadPoolConfig::default(),
87        }
88    }
89
90    /// Set the number of worker threads
91    pub fn num_threads(mut self, n: usize) -> Self {
92        self.config.num_threads = Some(n);
93        self
94    }
95
96    /// Set the thread name prefix
97    pub fn thread_name_prefix(mut self, prefix: impl Into<String>) -> Self {
98        self.config.thread_name_prefix = Some(prefix.into());
99        self
100    }
101
102    /// Set the stack size per thread
103    pub fn stack_size(mut self, size: usize) -> Self {
104        self.config.stack_size = Some(size);
105        self
106    }
107
108    /// Set the maximum pending tasks
109    pub fn max_pending_tasks(mut self, max: usize) -> Self {
110        self.config.max_pending_tasks = Some(max);
111        self
112    }
113
114    /// Enable or disable work stealing
115    pub fn enable_work_stealing(mut self, enabled: bool) -> Self {
116        self.config.enable_work_stealing = enabled;
117        self
118    }
119
120    /// Build the thread pool
121    pub fn build(self) -> Result<ThreadPool> {
122        ThreadPool::with_config(self.config)
123    }
124}
125
126impl Default for ThreadPoolBuilder {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132/// Work item type
133type WorkItem = Box<dyn FnOnce() + Send + 'static>;
134
135/// Result of work execution
136pub struct WorkResult<T> {
137    handle: ThreadHandle<T>,
138}
139
140impl<T> WorkResult<T> {
141    /// Wait for the work to complete and get the result
142    pub fn wait(self) -> Result<T> {
143        self.handle.join()
144    }
145
146    /// Wait for the work with a timeout
147    #[cfg(not(target_arch = "wasm32"))]
148    pub fn wait_timeout(self, timeout: Duration) -> Result<T>
149    where
150        T: Send + 'static,
151    {
152        self.handle.join_timeout(timeout)
153    }
154
155    #[cfg(target_arch = "wasm32")]
156    pub fn wait_timeout(self, _timeout: Duration) -> Result<T> {
157        self.wait()
158    }
159}
160
161/// Thread pool for executing tasks
162pub struct ThreadPool {
163    #[cfg(not(target_arch = "wasm32"))]
164    task_queue: Option<Arc<WasiQueue<WorkItem>>>,
165
166    #[cfg(not(target_arch = "wasm32"))]
167    workers: Vec<ThreadHandle<()>>,
168
169    #[cfg(not(target_arch = "wasm32"))]
170    rayon_pool: Option<Arc<rayon::ThreadPool>>,
171
172    #[cfg(target_arch = "wasm32")]
173    task_queue: Arc<WasiQueue<WorkItem>>,
174
175    #[allow(dead_code)]
176    config: ThreadPoolConfig,
177    is_shutdown: Arc<parking_lot::Mutex<bool>>,
178}
179
180impl ThreadPool {
181    /// Create a new thread pool with default configuration
182    pub fn new(num_threads: usize) -> Result<Self> {
183        let config = ThreadPoolConfig::new().num_threads(num_threads);
184        Self::with_config(config)
185    }
186
187    /// Create a thread pool with custom configuration
188    pub fn with_config(config: ThreadPoolConfig) -> Result<Self> {
189        let num_threads = config.num_threads.unwrap_or_else(|| {
190            #[cfg(not(target_arch = "wasm32"))]
191            {
192                WasiThread::available_parallelism()
193            }
194            #[cfg(target_arch = "wasm32")]
195            {
196                1
197            }
198        });
199
200        #[cfg(not(target_arch = "wasm32"))]
201        {
202            // Use rayon for work stealing if enabled
203            if config.enable_work_stealing {
204                let thread_name_prefix = config.thread_name_prefix.clone().unwrap_or_else(|| "worker".to_string());
205
206                let pool = rayon::ThreadPoolBuilder::new()
207                    .num_threads(num_threads)
208                    .thread_name(move |i| {
209                        format!("{}-{}", thread_name_prefix, i)
210                    })
211                    .build()
212                    .context("Failed to create rayon thread pool")?;
213
214                Ok(Self {
215                    task_queue: None,
216                    workers: Vec::new(),
217                    rayon_pool: Some(Arc::new(pool)),
218                    config,
219                    is_shutdown: Arc::new(parking_lot::Mutex::new(false)),
220                })
221            } else {
222                // Use manual thread pool
223                let task_queue: Arc<WasiQueue<WorkItem>> = if let Some(max) = config.max_pending_tasks {
224                    Arc::new(WasiQueue::with_capacity(max))
225                } else {
226                    Arc::new(WasiQueue::new())
227                };
228
229                let mut workers = Vec::with_capacity(num_threads);
230                let is_shutdown = Arc::new(parking_lot::Mutex::new(false));
231
232                for i in 0..num_threads {
233                    let queue_clone = task_queue.clone();
234                    let shutdown_clone = is_shutdown.clone();
235                    let name = format!("{}-{}", config.thread_name_prefix.as_deref().unwrap_or("worker"), i);
236
237                    let thread_config = super::ThreadConfig::new().with_name(name);
238
239                    let handle = WasiThread::spawn_with_config(move || {
240                        loop {
241                            // Check for shutdown
242                            if *shutdown_clone.lock() {
243                                break;
244                            }
245
246                            // Try to get a task
247                            if let Some(task) = queue_clone.try_pop() {
248                                task();
249                            } else {
250                                // No task available, sleep briefly
251                                super::thread_sleep(Duration::from_millis(10));
252                            }
253                        }
254                    }, thread_config)?;
255
256                    workers.push(handle);
257                }
258
259                Ok(Self {
260                    task_queue: Some(task_queue),
261                    workers,
262                    rayon_pool: None,
263                    config,
264                    is_shutdown,
265                })
266            }
267        }
268
269        #[cfg(target_arch = "wasm32")]
270        {
271            // WASM: Simple queue-based pool (no true parallelism)
272            let task_queue = if let Some(max) = config.max_pending_tasks {
273                Arc::new(WasiQueue::with_capacity(max))
274            } else {
275                Arc::new(WasiQueue::new())
276            };
277
278            Ok(Self {
279                task_queue,
280                config,
281                is_shutdown: Arc::new(parking_lot::Mutex::new(false)),
282            })
283        }
284    }
285
286    /// Execute a task in the thread pool (fire-and-forget)
287    pub fn execute<F>(&self, f: F) -> Result<()>
288    where
289        F: FnOnce() + Send + 'static,
290    {
291        if *self.is_shutdown.lock() {
292            return Err(anyhow!(ThreadingError::InvalidOperation("Thread pool is shut down".to_string())));
293        }
294
295        #[cfg(not(target_arch = "wasm32"))]
296        {
297            if let Some(ref rayon_pool) = self.rayon_pool {
298                rayon_pool.spawn(f);
299                Ok(())
300            } else if let Some(ref queue) = self.task_queue {
301                queue.push(Box::new(f))
302            } else {
303                Err(anyhow!(ThreadingError::InvalidOperation("Thread pool not properly initialized".to_string())))
304            }
305        }
306
307        #[cfg(target_arch = "wasm32")]
308        {
309            self.task_queue.push(Box::new(f))
310        }
311    }
312
313    /// Submit a task that returns a value
314    pub fn submit<F, T>(&self, f: F) -> Result<WorkResult<T>>
315    where
316        F: FnOnce() -> T + Send + 'static,
317        T: Send + 'static,
318    {
319        if *self.is_shutdown.lock() {
320            return Err(anyhow!(ThreadingError::InvalidOperation("Thread pool is shut down".to_string())));
321        }
322
323        // Spawn in a separate thread for now
324        // In a production implementation, this would use the pool's threads
325        let handle = WasiThread::spawn(f)?;
326        Ok(WorkResult { handle })
327    }
328
329    /// Execute work in parallel using rayon (native only)
330    #[cfg(not(target_arch = "wasm32"))]
331    pub fn parallel_for_each<T, F>(&self, items: Vec<T>, f: F) -> Result<()>
332    where
333        T: Send,
334        F: Fn(T) + Send + Sync,
335    {
336        if let Some(ref pool) = self.rayon_pool {
337            pool.install(|| {
338                rayon::scope(|s| {
339                    for item in items {
340                        s.spawn(|_| f(item));
341                    }
342                });
343            });
344            Ok(())
345        } else {
346            Err(anyhow!(ThreadingError::InvalidOperation("Parallel operations require work stealing enabled".to_string())))
347        }
348    }
349
350    /// Map operation in parallel (native only)
351    #[cfg(not(target_arch = "wasm32"))]
352    pub fn parallel_map<T, R, F>(&self, items: Vec<T>, f: F) -> Result<Vec<R>>
353    where
354        T: Send,
355        R: Send,
356        F: Fn(T) -> R + Send + Sync,
357    {
358        if let Some(ref pool) = self.rayon_pool {
359            Ok(pool.install(|| {
360                use rayon::prelude::*;
361                items.into_par_iter().map(f).collect()
362            }))
363        } else {
364            Err(anyhow!(ThreadingError::InvalidOperation("Parallel operations require work stealing enabled".to_string())))
365        }
366    }
367
368    /// Get the number of worker threads
369    pub fn num_threads(&self) -> usize {
370        #[cfg(not(target_arch = "wasm32"))]
371        {
372            if let Some(ref pool) = self.rayon_pool {
373                pool.current_num_threads()
374            } else {
375                self.workers.len()
376            }
377        }
378
379        #[cfg(target_arch = "wasm32")]
380        {
381            self.config.num_threads.unwrap_or(1)
382        }
383    }
384
385    /// Get pending task count
386    pub fn pending_tasks(&self) -> usize {
387        #[cfg(not(target_arch = "wasm32"))]
388        {
389            if let Some(ref queue) = self.task_queue {
390                queue.len()
391            } else {
392                0
393            }
394        }
395
396        #[cfg(target_arch = "wasm32")]
397        {
398            self.task_queue.len()
399        }
400    }
401
402    /// Shutdown the thread pool gracefully
403    pub fn shutdown(self) -> Result<()> {
404        *self.is_shutdown.lock() = true;
405
406        #[cfg(not(target_arch = "wasm32"))]
407        {
408            // Wait for all workers to finish
409            for worker in self.workers {
410                worker.join()?;
411            }
412        }
413
414        Ok(())
415    }
416
417    /// Shutdown and wait with timeout
418    #[cfg(not(target_arch = "wasm32"))]
419    pub fn shutdown_timeout(self, timeout: Duration) -> Result<()> {
420        *self.is_shutdown.lock() = true;
421
422        let start = std::time::Instant::now();
423        for worker in self.workers {
424            let remaining = timeout.saturating_sub(start.elapsed());
425            if remaining.is_zero() {
426                return Err(anyhow!(ThreadingError::Timeout("Thread pool shutdown timed out".to_string())));
427            }
428            worker.join_timeout(remaining)?;
429        }
430
431        Ok(())
432    }
433}
434
435/// Global thread pool for convenience
436static GLOBAL_POOL: once_cell::sync::Lazy<Result<ThreadPool>> = once_cell::sync::Lazy::new(|| {
437    ThreadPool::new(WasiThread::available_parallelism())
438});
439
440/// Execute a task on the global thread pool
441pub fn spawn<F>(f: F) -> Result<()>
442where
443    F: FnOnce() + Send + 'static,
444{
445    GLOBAL_POOL.as_ref()
446        .map_err(|e| anyhow!("Global thread pool not available: {}", e))?
447        .execute(f)
448}
449
450/// Submit a task on the global thread pool
451pub fn submit<F, T>(f: F) -> Result<WorkResult<T>>
452where
453    F: FnOnce() -> T + Send + 'static,
454    T: Send + 'static,
455{
456    GLOBAL_POOL.as_ref()
457        .map_err(|e| anyhow!("Global thread pool not available: {}", e))?
458        .submit(f)
459}
460
461#[cfg(test)]
462mod tests {
463    use super::*;
464    use std::sync::atomic::{AtomicUsize, Ordering};
465
466    #[test]
467    #[cfg(not(target_arch = "wasm32"))]
468    fn test_thread_pool_basic() {
469        let pool = ThreadPool::new(4).unwrap();
470        let counter = Arc::new(AtomicUsize::new(0));
471
472        for _ in 0..10 {
473            let counter_clone = counter.clone();
474            pool.execute(move || {
475                counter_clone.fetch_add(1, Ordering::SeqCst);
476            }).unwrap();
477        }
478
479        // Give threads time to execute
480        std::thread::sleep(Duration::from_millis(100));
481
482        assert_eq!(counter.load(Ordering::SeqCst), 10);
483    }
484
485    #[test]
486    #[cfg(not(target_arch = "wasm32"))]
487    fn test_thread_pool_submit() {
488        let pool = ThreadPool::new(2).unwrap();
489
490        let result = pool.submit(|| {
491            42
492        }).unwrap();
493
494        assert_eq!(result.wait().unwrap(), 42);
495    }
496
497    #[test]
498    #[cfg(not(target_arch = "wasm32"))]
499    fn test_thread_pool_rayon() {
500        let pool = ThreadPoolBuilder::new()
501            .num_threads(4)
502            .enable_work_stealing(true)
503            .build()
504            .unwrap();
505
506        let items = vec![1, 2, 3, 4, 5];
507        let results = pool.parallel_map(items, |x| x * 2).unwrap();
508
509        assert_eq!(results, vec![2, 4, 6, 8, 10]);
510    }
511
512    #[test]
513    fn test_thread_pool_config() {
514        let config = ThreadPoolConfig::new()
515            .num_threads(8)
516            .thread_name_prefix("test")
517            .max_pending_tasks(100);
518
519        assert_eq!(config.num_threads, Some(8));
520        assert_eq!(config.thread_name_prefix.as_deref(), Some("test"));
521        assert_eq!(config.max_pending_tasks, Some(100));
522    }
523}
portalis_transpiler/wasi_threading/pool.rs

portalis_transpiler/wasi_threading/
pool.rs
