1use std::panic::{AssertUnwindSafe, catch_unwind};
6
7use rayon::{ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder};
8use serde_json::Value;
9
10use super::{ArtifactType, CapabilityExecutor, ExecutorCapability, ExecutorError, ExecutorOutput};
11
12const MIN_CAPACITY: usize = 1;
13const MAX_CAPACITY: usize = 256;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub struct ThreadPoolExecutorConfig {
18 pub capacity: usize,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum ConfigError {
25 InvalidCapacity {
27 given: usize,
29 min: usize,
31 max: usize,
33 },
34 PoolBuildFailed(String),
36}
37
38impl std::fmt::Display for ConfigError {
39 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40 match self {
41 Self::InvalidCapacity { given, min, max } => {
42 write!(
43 f,
44 "invalid thread pool capacity {given}; expected {min}..={max}"
45 )
46 }
47 Self::PoolBuildFailed(msg) => write!(f, "thread pool build failed: {msg}"),
48 }
49 }
50}
51
52impl std::error::Error for ConfigError {}
53
54impl From<ThreadPoolBuildError> for ConfigError {
55 fn from(value: ThreadPoolBuildError) -> Self {
56 Self::PoolBuildFailed(value.to_string())
57 }
58}
59
60pub struct ThreadPoolExecutor {
62 pool: ThreadPool,
63 inner: Box<dyn CapabilityExecutor>,
64}
65
66impl std::fmt::Debug for ThreadPoolExecutor {
67 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68 f.debug_struct("ThreadPoolExecutor").finish_non_exhaustive()
69 }
70}
71
72impl ThreadPoolExecutor {
73 pub fn new(
80 config: ThreadPoolExecutorConfig,
81 inner: Box<dyn CapabilityExecutor>,
82 ) -> Result<Self, ConfigError> {
83 if !(MIN_CAPACITY..=MAX_CAPACITY).contains(&config.capacity) {
84 return Err(ConfigError::InvalidCapacity {
85 given: config.capacity,
86 min: MIN_CAPACITY,
87 max: MAX_CAPACITY,
88 });
89 }
90
91 let pool = ThreadPoolBuilder::new()
92 .num_threads(config.capacity)
93 .build()?;
94
95 Ok(Self { pool, inner })
96 }
97}
98
99impl CapabilityExecutor for ThreadPoolExecutor {
100 fn execute(
101 &self,
102 capability: &ExecutorCapability,
103 input: &Value,
104 ) -> Result<ExecutorOutput, ExecutorError> {
105 if capability.artifact_type == ArtifactType::Wasm {
106 return Err(ExecutorError::UnsupportedArtifactType);
107 }
108
109 let capability = capability.clone();
110 let input = input.clone();
111 let result = self
112 .pool
113 .install(|| catch_unwind(AssertUnwindSafe(|| self.inner.execute(&capability, &input))));
114
115 match result {
116 Ok(inner_result) => inner_result,
117 Err(_) => Err(ExecutorError::ExecutionFailed(
118 "capability panicked".to_string(),
119 )),
120 }
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use std::io;
127 use std::sync::{
128 Arc,
129 atomic::{AtomicUsize, Ordering},
130 };
131 use std::thread;
132 use std::time::Duration;
133
134 use serde_json::json;
135
136 use super::*;
137 use crate::executor::NativeExecutor;
138
139 fn native_capability() -> ExecutorCapability {
140 ExecutorCapability {
141 capability_id: "test.native".to_string(),
142 artifact_type: ArtifactType::Native,
143 wasm_binary_path: None,
144 wasm_checksum: None,
145 host_abi_version: None,
146 emits: Vec::new(),
147 service_type: traverse_contracts::ServiceType::Stateless,
148 }
149 }
150
151 fn wasm_capability() -> ExecutorCapability {
152 ExecutorCapability {
153 artifact_type: ArtifactType::Wasm,
154 ..native_capability()
155 }
156 }
157
158 fn new_executor(
159 capacity: usize,
160 inner: Box<dyn CapabilityExecutor>,
161 ) -> Result<ThreadPoolExecutor, ConfigError> {
162 ThreadPoolExecutor::new(ThreadPoolExecutorConfig { capacity }, inner)
163 }
164
165 fn clone_input(input: &Value) -> Result<Value, String> {
166 if input.get("__thread_pool_test_error").is_some() {
167 return Err("requested test error".to_string());
168 }
169 Ok(input.clone())
170 }
171
172 fn result_debug<T, E: std::fmt::Debug>(result: Result<T, E>) -> Result<T, String> {
173 result.map_err(|err| format!("{err:?}"))
174 }
175
176 fn executor(capacity: usize) -> Result<ThreadPoolExecutor, String> {
177 result_debug(new_executor(
178 capacity,
179 Box::new(NativeExecutor::new(clone_input)),
180 ))
181 }
182
183 fn execute_json(executor: &ThreadPoolExecutor, input: &Value) -> Result<Value, ExecutorError> {
184 executor
185 .execute(&native_capability(), input)
186 .map(|output| output.value)
187 }
188
189 #[test]
190 fn clone_input_helper_can_return_error() {
191 let result = clone_input(&json!({ "__thread_pool_test_error": true }));
192
193 assert_eq!(result, Err("requested test error".to_string()));
194 }
195
196 #[test]
197 fn config_capacity_zero_returns_error() {
198 let result = new_executor(0, Box::new(NativeExecutor::new(clone_input)));
199
200 assert_eq!(
201 result.err(),
202 Some(ConfigError::InvalidCapacity {
203 given: 0,
204 min: 1,
205 max: 256
206 })
207 );
208 }
209
210 #[test]
211 fn config_capacity_max_valid() {
212 let result = new_executor(256, Box::new(NativeExecutor::new(clone_input)));
213
214 assert!(result.is_ok());
215 }
216
217 #[test]
218 fn config_capacity_over_max_returns_error() {
219 let result = new_executor(257, Box::new(NativeExecutor::new(clone_input)));
220
221 assert_eq!(
222 result.err(),
223 Some(ConfigError::InvalidCapacity {
224 given: 257,
225 min: 1,
226 max: 256
227 })
228 );
229 }
230
231 #[test]
232 fn config_capacity_one_valid() {
233 let result = new_executor(1, Box::new(NativeExecutor::new(clone_input)));
234
235 assert!(result.is_ok());
236 }
237
238 #[test]
239 fn config_error_display_and_from_cover_failure_shapes() {
240 let invalid = ConfigError::InvalidCapacity {
241 given: 0,
242 min: 1,
243 max: 256,
244 };
245
246 assert_eq!(
247 invalid.to_string(),
248 "invalid thread pool capacity 0; expected 1..=256"
249 );
250
251 let build_error = ThreadPoolBuilder::new()
252 .num_threads(1)
253 .spawn_handler(|_| Err(io::Error::other("spawn denied")))
254 .build()
255 .map_err(ConfigError::from);
256
257 assert!(matches!(build_error, Err(ConfigError::PoolBuildFailed(_))));
258
259 let display = build_error.map_err(|err| err.to_string()).err();
260 assert_eq!(
261 display,
262 Some("thread pool build failed: spawn denied".to_string())
263 );
264 }
265
266 #[test]
267 fn thread_pool_executor_debug_is_non_exhaustive() -> Result<(), String> {
268 let debug = executor(1).map(|executor| format!("{executor:?}"))?;
269
270 assert_eq!(debug, "ThreadPoolExecutor { .. }");
271 Ok(())
272 }
273
274 #[test]
275 fn execute_native_returns_correct_output() -> Result<(), String> {
276 let executor = executor(2)?;
277 let result = execute_json(&executor, &json!({ "value": 42 }));
278
279 assert_eq!(result, Ok(json!({ "value": 42 })));
280 Ok(())
281 }
282
283 #[test]
284 fn execute_native_error_propagates() {
285 let result = new_executor(
286 2,
287 Box::new(NativeExecutor::new(|_| Err("inner failed".to_string()))),
288 )
289 .map(|executor| executor.execute(&native_capability(), &json!({})));
290
291 assert_eq!(
292 result,
293 Ok(Err(ExecutorError::ExecutionFailed(
294 "inner failed".to_string()
295 )))
296 );
297 }
298
299 #[test]
300 fn execute_wasm_artifact_type_returns_unsupported() {
301 let result = new_executor(2, Box::new(PanickingExecutor::new(1)))
302 .map(|executor| executor.execute(&wasm_capability(), &json!({})));
303
304 assert_eq!(result, Ok(Err(ExecutorError::UnsupportedArtifactType)));
305 }
306
307 #[test]
308 fn concurrent_calls_run_in_parallel() {
309 let active_calls = Arc::new(AtomicUsize::new(0));
310 let max_active_calls = Arc::new(AtomicUsize::new(0));
311 let active_for_handler = Arc::clone(&active_calls);
312 let max_for_handler = Arc::clone(&max_active_calls);
313
314 let result = new_executor(
315 2,
316 Box::new(NativeExecutor::new(move |input| {
317 let current = active_for_handler.fetch_add(1, Ordering::SeqCst) + 1;
318 max_for_handler.fetch_max(current, Ordering::SeqCst);
319 thread::sleep(Duration::from_millis(100));
320 active_for_handler.fetch_sub(1, Ordering::SeqCst);
321 Ok(input.clone())
322 })),
323 )
324 .map(|executor| {
325 let executor = Arc::new(executor);
326 let first = {
327 let executor = Arc::clone(&executor);
328 thread::spawn(move || execute_json(&executor, &json!({ "call": 1 })))
329 };
330 let second = {
331 let executor = Arc::clone(&executor);
332 thread::spawn(move || execute_json(&executor, &json!({ "call": 2 })))
333 };
334 (result_debug(first.join()), result_debug(second.join()))
335 });
336
337 let first_result = result
338 .as_ref()
339 .map(|(first, _)| first.as_ref().map_err(String::as_str));
340 let second_result = result
341 .as_ref()
342 .map(|(_, second)| second.as_ref().map_err(String::as_str));
343
344 assert_eq!(first_result, Ok(Ok(&Ok(json!({ "call": 1 })))));
345 assert_eq!(second_result, Ok(Ok(&Ok(json!({ "call": 2 })))));
346 assert!(
347 max_active_calls.load(Ordering::SeqCst) >= 2,
348 "expected two active calls to overlap"
349 );
350 }
351
352 #[test]
353 fn pool_size_one_serialises_calls() {
354 let result = new_executor(
355 1,
356 Box::new(NativeExecutor::new(|input| {
357 thread::sleep(Duration::from_millis(50));
358 Ok(input.clone())
359 })),
360 )
361 .map(|executor| {
362 let executor = Arc::new(executor);
363 let first = {
364 let executor = Arc::clone(&executor);
365 thread::spawn(move || execute_json(&executor, &json!({ "call": 1 })))
366 };
367 let second = {
368 let executor = Arc::clone(&executor);
369 thread::spawn(move || execute_json(&executor, &json!({ "call": 2 })))
370 };
371 (result_debug(first.join()), result_debug(second.join()))
372 });
373 let first_result = result
374 .as_ref()
375 .map(|(first, _)| first.as_ref().map_err(String::as_str));
376 let second_result = result
377 .as_ref()
378 .map(|(_, second)| second.as_ref().map_err(String::as_str));
379
380 assert_eq!(first_result, Ok(Ok(&Ok(json!({ "call": 1 })))));
381 assert_eq!(second_result, Ok(Ok(&Ok(json!({ "call": 2 })))));
382 }
383
384 #[test]
385 fn independent_calls_do_not_share_state() -> Result<(), String> {
386 let executor = Arc::new(executor(4)?);
387 let mut handles = Vec::new();
388
389 for value in 0..10 {
390 let executor = Arc::clone(&executor);
391 handles.push(thread::spawn(move || {
392 execute_json(&executor, &json!({ "value": value }))
393 }));
394 }
395
396 for (value, handle) in handles.into_iter().enumerate() {
397 let result = result_debug(handle.join())?;
398 assert_eq!(result, Ok(json!({ "value": value })));
399 }
400 Ok(())
401 }
402
403 struct PanickingExecutor {
404 remaining_panics: AtomicUsize,
405 }
406
407 impl PanickingExecutor {
408 fn new(remaining_panics: usize) -> Self {
409 Self {
410 remaining_panics: AtomicUsize::new(remaining_panics),
411 }
412 }
413 }
414
415 impl CapabilityExecutor for PanickingExecutor {
416 fn execute(
417 &self,
418 _capability: &ExecutorCapability,
419 input: &Value,
420 ) -> Result<ExecutorOutput, ExecutorError> {
421 let remaining = self.remaining_panics.load(Ordering::SeqCst);
422 if remaining > 0 {
423 self.remaining_panics.fetch_sub(1, Ordering::SeqCst);
424 std::panic::resume_unwind(Box::new("boom"));
425 }
426 Ok(ExecutorOutput {
427 value: input.clone(),
428 emitted_events: Vec::new(),
429 connector_invocation_evidence: Vec::new(),
430 })
431 }
432 }
433
434 #[test]
435 fn panicking_handler_returns_execution_failed() -> Result<(), String> {
436 let executor = result_debug(new_executor(1, Box::new(PanickingExecutor::new(1))))?;
437 let result = executor.execute(&native_capability(), &json!({}));
438
439 assert_eq!(
440 result,
441 Err(ExecutorError::ExecutionFailed(
442 "capability panicked".to_string()
443 ))
444 );
445 Ok(())
446 }
447
448 #[test]
449 fn pool_usable_after_panic() -> Result<(), String> {
450 let executor = result_debug(new_executor(1, Box::new(PanickingExecutor::new(1))))?;
451 let failed = executor.execute(&native_capability(), &json!({ "first": true }));
452 let recovered = executor.execute(&native_capability(), &json!({ "second": true }));
453
454 assert_eq!(
455 failed,
456 Err(ExecutorError::ExecutionFailed(
457 "capability panicked".to_string()
458 ))
459 );
460 assert_eq!(
461 recovered,
462 Ok(ExecutorOutput {
463 value: json!({ "second": true }),
464 emitted_events: Vec::new(),
465 connector_invocation_evidence: Vec::new(),
466 })
467 );
468 Ok(())
469 }
470
471 #[test]
472 fn multiple_sequential_panics_do_not_exhaust_pool() -> Result<(), String> {
473 let executor = result_debug(new_executor(2, Box::new(PanickingExecutor::new(5))))?;
474
475 for _ in 0..5 {
476 let result = executor.execute(&native_capability(), &json!({}));
477 assert_eq!(
478 result,
479 Err(ExecutorError::ExecutionFailed(
480 "capability panicked".to_string()
481 ))
482 );
483 }
484
485 let recovered = executor.execute(&native_capability(), &json!({ "ok": true }));
486
487 assert_eq!(
488 recovered,
489 Ok(ExecutorOutput {
490 value: json!({ "ok": true }),
491 emitted_events: Vec::new(),
492 connector_invocation_evidence: Vec::new(),
493 })
494 );
495 Ok(())
496 }
497
498 #[test]
499 fn executor_is_send_sync() {
500 fn assert_send_sync<T: Send + Sync>() {}
501
502 assert_send_sync::<ThreadPoolExecutor>();
503 }
504
505 #[test]
506 fn arc_wrapped_executor_callable_from_multiple_threads() -> Result<(), String> {
507 let executor = Arc::new(executor(4)?);
508 let mut handles = Vec::new();
509
510 for value in 0..4 {
511 let executor = Arc::clone(&executor);
512 handles.push(thread::spawn(move || {
513 execute_json(&executor, &json!({ "thread": value }))
514 }));
515 }
516
517 for (value, handle) in handles.into_iter().enumerate() {
518 let result = result_debug(handle.join())?;
519 assert_eq!(result, Ok(json!({ "thread": value })));
520 }
521 Ok(())
522 }
523
524 #[test]
525 fn executor_drops_cleanly_after_use() -> Result<(), String> {
526 let executor = executor(1)?;
527 let result = execute_json(&executor, &json!({ "used": true }));
528
529 assert_eq!(result, Ok(json!({ "used": true })));
530 drop(executor);
531 Ok(())
532 }
533
534 #[test]
535 fn executor_drops_cleanly_with_no_calls() -> Result<(), String> {
536 let executor = executor(1)?;
537
538 drop(executor);
539 Ok(())
540 }
541}