preemptive_threads/
preemption.rs

1#[cfg(target_os = "linux")]
2pub struct Preemption {
3    enabled: bool,
4}
5
6#[cfg(target_os = "linux")]
7impl Default for Preemption {
8    fn default() -> Self {
9        Self::new()
10    }
11}
12
13#[cfg(target_os = "linux")]
14impl Preemption {
15    pub const fn new() -> Self {
16        Preemption { enabled: false }
17    }
18
19    /// # Safety
20    /// This function sets up signal handlers and timer interrupts which can affect
21    /// the entire process. Only one instance should manage preemption at a time.
22    /// The caller must ensure thread safety when accessing the global scheduler.
23    pub unsafe fn enable(&mut self, interval_us: u64) {
24        if self.enabled {
25            return;
26        }
27
28        extern "C" {
29            fn signal(sig: i32, handler: extern "C" fn(i32)) -> i32;
30            fn setitimer(which: i32, new_value: *const itimerval, old_value: *mut itimerval)
31                -> i32;
32        }
33
34        const SIGALRM: i32 = 14;
35        const ITIMER_REAL: i32 = 0;
36
37        #[repr(C)]
38        struct timeval {
39            tv_sec: i64,
40            tv_usec: i64,
41        }
42
43        #[repr(C)]
44        struct itimerval {
45            it_interval: timeval,
46            it_value: timeval,
47        }
48
49        if interval_us == 0 {
50            return;
51        }
52
53        let result = signal(SIGALRM, timer_handler);
54        if result == -1 {
55            return;
56        }
57
58        let timer = itimerval {
59            it_interval: timeval {
60                tv_sec: (interval_us / 1_000_000) as i64,
61                tv_usec: (interval_us % 1_000_000) as i64,
62            },
63            it_value: timeval {
64                tv_sec: (interval_us / 1_000_000) as i64,
65                tv_usec: (interval_us % 1_000_000) as i64,
66            },
67        };
68
69        let result = setitimer(ITIMER_REAL, &timer, core::ptr::null_mut());
70        if result == 0 {
71            self.enabled = true;
72        }
73    }
74
75    /// # Safety
76    /// This function modifies process-wide timer settings. The caller must ensure
77    /// no other code depends on SIGALRM or ITIMER_REAL functionality.
78    pub unsafe fn disable(&mut self) {
79        if !self.enabled {
80            return;
81        }
82
83        extern "C" {
84            fn setitimer(which: i32, new_value: *const itimerval, old_value: *mut itimerval)
85                -> i32;
86        }
87
88        const ITIMER_REAL: i32 = 0;
89
90        #[repr(C)]
91        struct timeval {
92            tv_sec: i64,
93            tv_usec: i64,
94        }
95
96        #[repr(C)]
97        struct itimerval {
98            it_interval: timeval,
99            it_value: timeval,
100        }
101
102        let timer = itimerval {
103            it_interval: timeval {
104                tv_sec: 0,
105                tv_usec: 0,
106            },
107            it_value: timeval {
108                tv_sec: 0,
109                tv_usec: 0,
110            },
111        };
112
113        setitimer(ITIMER_REAL, &timer, core::ptr::null_mut());
114        self.enabled = false;
115    }
116}
117
118#[cfg(target_os = "linux")]
119extern "C" fn timer_handler(_sig: i32) {
120    use core::sync::atomic::{AtomicBool, Ordering};
121
122    static IN_HANDLER: AtomicBool = AtomicBool::new(false);
123
124    if IN_HANDLER.swap(true, Ordering::SeqCst) {
125        return;
126    }
127
128    unsafe {
129        let scheduler = crate::scheduler::SCHEDULER.get();
130
131        if let Some(current_id) = scheduler.get_current_thread() {
132            if let Some(next_id) = scheduler.schedule() {
133                if current_id != next_id {
134                    scheduler.set_current_thread(Some(next_id));
135                    let _ = scheduler.switch_context(current_id, next_id);
136                }
137            }
138        }
139    }
140
141    IN_HANDLER.store(false, Ordering::SeqCst);
142}
143
144#[cfg(not(target_os = "linux"))]
145pub struct Preemption;
146
147#[cfg(not(target_os = "linux"))]
148impl Default for Preemption {
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154#[cfg(not(target_os = "linux"))]
155impl Preemption {
156    pub const fn new() -> Self {
157        Preemption
158    }
159
160    /// # Safety
161    /// Enables timer-based preemption. May affect signal handlers.
162    pub unsafe fn enable(&mut self, _interval_us: u64) {}
163
164    /// # Safety
165    /// Disables timer-based preemption. May affect signal handlers.
166    pub unsafe fn disable(&mut self) {}
167}