probe_code/search/
timeout.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::Arc;
3use std::thread;
4use std::time::Duration;
5
6/// Starts a timeout thread that will terminate the process if the timeout is reached.
7/// Returns a handle to the timeout thread that can be used to stop it.
8pub fn start_timeout_thread(timeout_seconds: u64) -> Arc<AtomicBool> {
9    let should_stop = Arc::new(AtomicBool::new(false));
10    let should_stop_clone = should_stop.clone();
11
12    // For testing purposes, check if we're running in a test environment
13    let is_test = std::env::var("RUST_TEST_THREADS").is_ok();
14
15    // Use a shorter sleep interval for tests to make timeouts more reliable
16    let sleep_interval = if is_test {
17        Duration::from_millis(10) // 100ms for tests
18    } else {
19        Duration::from_secs(1) // 1 second for normal operation
20    };
21
22    thread::spawn(move || {
23        let mut elapsed_time = Duration::from_secs(0);
24        let timeout_duration = Duration::from_secs(timeout_seconds);
25
26        while elapsed_time < timeout_duration {
27            // Check if we should stop the timeout thread
28            if should_stop_clone.load(Ordering::SeqCst) {
29                return;
30            }
31
32            // Sleep for the interval
33            thread::sleep(sleep_interval);
34            elapsed_time += sleep_interval;
35        }
36
37        // Timeout reached, print a message and terminate the process
38        eprintln!("Search operation timed out after {timeout_seconds} seconds");
39        std::process::exit(1);
40    });
41
42    should_stop
43}