rustpython_host_env/
multiprocessing.rs

1#![allow(
2    clippy::not_unsafe_ptr_arg_deref,
3    reason = "Semaphore helpers intentionally mirror OS handle and pointer APIs."
4)]
5#![allow(
6    clippy::result_unit_err,
7    reason = "These helpers preserve the existing host-facing error surface."
8)]
9
10#[cfg(unix)]
11use alloc::ffi::CString;
12#[cfg(windows)]
13use std::io;
14
15#[cfg(unix)]
16use libc::sem_t;
17#[cfg(unix)]
18use nix::errno::Errno;
19
20#[cfg(unix)]
21#[repr(transparent)]
22#[derive(Debug)]
23pub struct SemHandle {
24    raw: *mut sem_t,
25}
26
27// POSIX named semaphores are safe to post/wait from any thread that holds
28// a `sem_t *` to the same kernel object.
29#[cfg(unix)]
30unsafe impl Send for SemHandle {}
31#[cfg(unix)]
32unsafe impl Sync for SemHandle {}
33
34#[cfg(unix)]
35#[derive(Copy, Clone, Debug, Eq, PartialEq)]
36pub enum SemError {
37    WouldBlock,
38    TimedOut,
39    Interrupted,
40    AlreadyExists,
41    NotFound,
42    InvalidInput,
43    InteriorNul,
44    Other(i32),
45}
46
47#[cfg(unix)]
48impl SemError {
49    const fn from_errno(err: Errno) -> Self {
50        match err {
51            Errno::EAGAIN => Self::WouldBlock,
52            Errno::ETIMEDOUT => Self::TimedOut,
53            Errno::EINTR => Self::Interrupted,
54            Errno::EEXIST => Self::AlreadyExists,
55            Errno::ENOENT => Self::NotFound,
56            Errno::EINVAL => Self::InvalidInput,
57            other => Self::Other(other as i32),
58        }
59    }
60
61    pub const fn raw_os_error(self) -> i32 {
62        match self {
63            Self::WouldBlock => Errno::EAGAIN as i32,
64            Self::TimedOut => Errno::ETIMEDOUT as i32,
65            Self::Interrupted => Errno::EINTR as i32,
66            Self::AlreadyExists => Errno::EEXIST as i32,
67            Self::NotFound => Errno::ENOENT as i32,
68            Self::InvalidInput | Self::InteriorNul => Errno::EINVAL as i32,
69            Self::Other(code) => code,
70        }
71    }
72
73    pub fn description(self) -> String {
74        Errno::from_raw(self.raw_os_error()).desc().to_owned()
75    }
76}
77
78#[cfg(unix)]
79#[derive(Copy, Clone, Debug, Eq, PartialEq)]
80pub enum TryAcquireStatus {
81    Acquired,
82    WouldBlock,
83    Interrupted,
84    Error(SemError),
85}
86
87#[cfg(unix)]
88#[derive(Copy, Clone, Debug, Eq, PartialEq)]
89pub enum WaitStatus {
90    Acquired,
91    TimedOut,
92    Interrupted,
93    Error(SemError),
94}
95
96#[cfg(unix)]
97#[derive(Copy, Clone, Debug)]
98pub struct Deadline {
99    spec: libc::timespec,
100}
101
102#[cfg(windows)]
103use windows_sys::Win32::{
104    Foundation::{
105        CloseHandle, ERROR_TOO_MANY_POSTS, GetLastError, HANDLE, INVALID_HANDLE_VALUE, WAIT_FAILED,
106        WAIT_OBJECT_0, WAIT_TIMEOUT,
107    },
108    Networking::WinSock::{SOCKET, WSAGetLastError, closesocket, recv, send},
109    System::Threading::{
110        CreateSemaphoreW, GetCurrentThreadId, INFINITE, ReleaseSemaphore, WaitForSingleObjectEx,
111    },
112};
113
114#[cfg(windows)]
115pub type RawHandle = HANDLE;
116#[cfg(windows)]
117pub type RawSocket = SOCKET;
118#[cfg(windows)]
119pub const INFINITE_TIMEOUT: u32 = INFINITE;
120
121#[cfg(windows)]
122#[repr(transparent)]
123#[derive(Debug)]
124pub struct SemHandle {
125    raw: HANDLE,
126}
127
128#[cfg(windows)]
129unsafe impl Send for SemHandle {}
130#[cfg(windows)]
131unsafe impl Sync for SemHandle {}
132
133#[cfg(unix)]
134impl SemHandle {
135    pub fn create(
136        name: &str,
137        value: u32,
138        unlink: bool,
139    ) -> Result<(Self, Option<String>), SemError> {
140        let cname = semaphore_name(name)?;
141        let raw = unsafe {
142            libc::sem_open(
143                cname.as_ptr(),
144                crate::os::O_CREAT | crate::os::O_EXCL,
145                0o600,
146                value,
147            )
148        };
149        if raw == libc::SEM_FAILED {
150            return Err(SemError::from_errno(Errno::last()));
151        }
152        if unlink {
153            if unsafe { libc::sem_unlink(cname.as_ptr()) } != 0 {
154                let err = SemError::from_errno(Errno::last());
155                unsafe {
156                    libc::sem_close(raw);
157                }
158                Err(err)
159            } else {
160                Ok((Self { raw }, None))
161            }
162        } else {
163            Ok((Self { raw }, Some(name.to_owned())))
164        }
165    }
166
167    pub fn open_existing(name: &str) -> Result<Self, SemError> {
168        let cname = semaphore_name(name)?;
169        let raw = unsafe { libc::sem_open(cname.as_ptr(), 0) };
170        if raw == libc::SEM_FAILED {
171            Err(SemError::from_errno(Errno::last()))
172        } else {
173            Ok(Self { raw })
174        }
175    }
176
177    #[inline]
178    pub fn as_handle_int(&self) -> isize {
179        self.raw as isize
180    }
181
182    /// Rebuild a handle from the integer a `SemLock` stored.
183    ///
184    /// # Safety
185    ///
186    /// `raw` must be a live `sem_open` handle. `Drop` closes it, so a
187    /// caller that still owns that close must forget this value.
188    #[inline]
189    pub const unsafe fn from_raw(raw: *mut sem_t) -> Self {
190        Self { raw }
191    }
192
193    #[inline]
194    pub fn trywait(&self) -> TryAcquireStatus {
195        sem_trywait_status(self.raw)
196    }
197
198    #[inline]
199    pub fn post(&self) -> Result<(), SemError> {
200        sem_post(self.raw)
201    }
202
203    #[inline]
204    pub fn wait(&self, deadline: Option<&Deadline>) -> WaitStatus {
205        sem_wait_status(self.raw, deadline.map(|deadline| &deadline.spec))
206    }
207
208    #[cfg(not(target_vendor = "apple"))]
209    #[inline]
210    pub fn value(&self) -> Result<i32, SemError> {
211        // Safety: `self.raw` is a live `sem_open` handle owned by this object.
212        unsafe { get_semaphore_value(self.raw) }
213    }
214
215    #[cfg(target_vendor = "apple")]
216    #[inline]
217    pub fn poll_wait_step(
218        &self,
219        deadline: &Deadline,
220        delay: u64,
221    ) -> Result<PollWaitStep, SemError> {
222        sem_timedwait_poll_step(self.raw, &deadline.spec, delay)
223    }
224}
225
226#[cfg(windows)]
227impl SemHandle {
228    pub fn create(value: i32, maxvalue: i32) -> io::Result<Self> {
229        use crate::windows::CheckWin32Handle;
230        let handle =
231            unsafe { CreateSemaphoreW(core::ptr::null(), value, maxvalue, core::ptr::null()) }
232                .check_nonnull()?;
233        Ok(Self { raw: handle })
234    }
235
236    #[inline]
237    pub fn from_raw(raw: HANDLE) -> Self {
238        Self { raw }
239    }
240
241    #[inline]
242    pub fn as_handle_int(&self) -> isize {
243        self.raw as isize
244    }
245
246    #[inline]
247    pub fn as_raw(&self) -> HANDLE {
248        self.raw
249    }
250
251    #[inline]
252    pub fn wait(&self, timeout_ms: u32) -> u32 {
253        wait_for_single_object(self.raw, timeout_ms)
254    }
255
256    #[inline]
257    pub fn release(&self) -> Result<(), u32> {
258        release_semaphore(self.raw)
259    }
260
261    #[inline]
262    pub fn value(&self) -> Result<i32, ()> {
263        get_semaphore_value(self.raw)
264    }
265}
266
267#[cfg(unix)]
268impl Drop for SemHandle {
269    fn drop(&mut self) {
270        if !self.raw.is_null() {
271            unsafe {
272                libc::sem_close(self.raw);
273            }
274        }
275    }
276}
277
278#[cfg(windows)]
279impl Drop for SemHandle {
280    fn drop(&mut self) {
281        if self.raw != 0 as HANDLE && self.raw != INVALID_HANDLE_VALUE {
282            unsafe {
283                CloseHandle(self.raw);
284            }
285        }
286    }
287}
288
289#[cfg(unix)]
290#[inline]
291pub fn current_thread_id() -> u64 {
292    unsafe { libc::pthread_self() as u64 }
293}
294
295#[cfg(windows)]
296#[inline]
297pub fn current_thread_id() -> u32 {
298    unsafe { GetCurrentThreadId() }
299}
300
301#[cfg(windows)]
302#[inline]
303pub fn wait_for_single_object(handle: HANDLE, timeout_ms: u32) -> u32 {
304    unsafe { WaitForSingleObjectEx(handle, timeout_ms, 0) }
305}
306
307#[cfg(windows)]
308#[inline]
309pub fn wait_object_0() -> u32 {
310    WAIT_OBJECT_0
311}
312
313#[cfg(windows)]
314#[inline]
315pub fn wait_timeout() -> u32 {
316    WAIT_TIMEOUT
317}
318
319#[cfg(windows)]
320#[inline]
321pub fn close_socket(socket: SOCKET) -> io::Result<()> {
322    let res = unsafe { closesocket(socket) };
323    if res != 0 {
324        Err(io::Error::from_raw_os_error(unsafe { WSAGetLastError() }))
325    } else {
326        Ok(())
327    }
328}
329
330#[cfg(windows)]
331pub fn recv_socket(socket: SOCKET, size: usize) -> io::Result<Vec<u8>> {
332    let len = i32::try_from(size).map_err(|_| {
333        io::Error::new(io::ErrorKind::InvalidInput, "socket receive size too large")
334    })?;
335    let mut buf = vec![0u8; size];
336    let n_read = unsafe { recv(socket, buf.as_mut_ptr() as *mut _, len, 0) };
337    if n_read < 0 {
338        Err(io::Error::from_raw_os_error(unsafe { WSAGetLastError() }))
339    } else {
340        buf.truncate(n_read as usize);
341        Ok(buf)
342    }
343}
344
345#[cfg(windows)]
346pub fn send_socket(socket: SOCKET, buf: &[u8]) -> io::Result<i32> {
347    let len = i32::try_from(buf.len())
348        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "socket send buffer too large"))?;
349    let ret = unsafe { send(socket, buf.as_ptr() as *const _, len, 0) };
350    if ret < 0 {
351        Err(io::Error::from_raw_os_error(unsafe { WSAGetLastError() }))
352    } else {
353        Ok(ret)
354    }
355}
356
357#[cfg(windows)]
358#[inline]
359pub fn wait_failed() -> u32 {
360    WAIT_FAILED
361}
362
363#[cfg(windows)]
364pub fn release_semaphore(handle: HANDLE) -> Result<(), u32> {
365    if unsafe { ReleaseSemaphore(handle, 1, core::ptr::null_mut()) } == 0 {
366        Err(unsafe { GetLastError() })
367    } else {
368        Ok(())
369    }
370}
371
372#[cfg(windows)]
373pub fn get_semaphore_value(handle: HANDLE) -> Result<i32, ()> {
374    match wait_for_single_object(handle, 0) {
375        WAIT_OBJECT_0 => {
376            let mut previous: i32 = 0;
377            if unsafe { ReleaseSemaphore(handle, 1, &mut previous) } == 0 {
378                Err(())
379            } else {
380                Ok(previous + 1)
381            }
382        }
383        WAIT_TIMEOUT => Ok(0),
384        _ => Err(()),
385    }
386}
387
388#[cfg(windows)]
389#[inline]
390pub fn is_too_many_posts(err: u32) -> bool {
391    err == ERROR_TOO_MANY_POSTS
392}
393
394#[cfg(unix)]
395pub fn semaphore_name(name: &str) -> Result<CString, SemError> {
396    let mut full = String::with_capacity(name.len() + 2);
397    if !name.starts_with('/') {
398        full.push('/');
399    }
400    full.push_str(name);
401    CString::new(full).map_err(|_| SemError::InteriorNul)
402}
403
404#[cfg(unix)]
405pub fn sem_unlink(name: &str) -> Result<(), SemError> {
406    let cname = semaphore_name(name)?;
407    let res = unsafe { libc::sem_unlink(cname.as_ptr()) };
408    if res < 0 {
409        Err(SemError::from_errno(Errno::last()))
410    } else {
411        Ok(())
412    }
413}
414
415#[cfg(all(unix, not(target_vendor = "apple")))]
416/// # Safety
417///
418/// `handle` must point to a valid `sem_t` that remains alive for the duration
419/// of this call and is valid to pass to `sem_getvalue`.
420unsafe fn get_semaphore_value(handle: *mut sem_t) -> Result<i32, SemError> {
421    let mut sval: libc::c_int = 0;
422    let res = unsafe { libc::sem_getvalue(handle, &mut sval) };
423    if res < 0 {
424        Err(SemError::from_errno(Errno::last()))
425    } else {
426        Ok(if sval < 0 { 0 } else { sval })
427    }
428}
429
430#[cfg(unix)]
431#[allow(clippy::not_unsafe_ptr_arg_deref)]
432fn sem_trywait_status(handle: *mut sem_t) -> TryAcquireStatus {
433    if unsafe { libc::sem_trywait(handle) } == 0 {
434        TryAcquireStatus::Acquired
435    } else {
436        match Errno::last() {
437            Errno::EAGAIN => TryAcquireStatus::WouldBlock,
438            Errno::EINTR => TryAcquireStatus::Interrupted,
439            err => TryAcquireStatus::Error(SemError::from_errno(err)),
440        }
441    }
442}
443
444#[cfg(unix)]
445#[allow(clippy::not_unsafe_ptr_arg_deref)]
446fn sem_post(handle: *mut sem_t) -> Result<(), SemError> {
447    if unsafe { libc::sem_post(handle) } < 0 {
448        Err(SemError::from_errno(Errno::last()))
449    } else {
450        Ok(())
451    }
452}
453
454#[cfg(unix)]
455pub fn sem_value_max() -> i32 {
456    let val = unsafe { libc::sysconf(libc::_SC_SEM_VALUE_MAX) };
457    if val < 0 || val > i32::MAX as libc::c_long {
458        i32::MAX
459    } else {
460        val as i32
461    }
462}
463
464#[cfg(unix)]
465fn gettimeofday() -> Result<libc::timeval, SemError> {
466    let mut tv = libc::timeval {
467        tv_sec: 0,
468        tv_usec: 0,
469    };
470    if unsafe { libc::gettimeofday(&mut tv, core::ptr::null_mut()) } < 0 {
471        Err(SemError::from_errno(Errno::last()))
472    } else {
473        Ok(tv)
474    }
475}
476
477#[cfg(unix)]
478pub fn deadline_from_timeout(timeout: f64) -> Result<Deadline, SemError> {
479    let timeout = if timeout < 0.0 { 0.0 } else { timeout };
480    if !timeout.is_finite() {
481        return Err(SemError::InvalidInput);
482    }
483    let tv = gettimeofday()?;
484    let sec_f64 = timeout.floor();
485    if sec_f64 > libc::time_t::MAX as f64 {
486        return Err(SemError::InvalidInput);
487    }
488    let sec = sec_f64 as libc::time_t;
489    let nsec = (1e9 * (timeout - sec as f64) + 0.5) as libc::c_long;
490    let tv_nsec = (tv.tv_usec as libc::c_long)
491        .checked_mul(1000)
492        .and_then(|base| base.checked_add(nsec))
493        .ok_or(SemError::InvalidInput)?;
494    let mut deadline = libc::timespec {
495        tv_sec: tv.tv_sec.checked_add(sec).ok_or(SemError::InvalidInput)?,
496        tv_nsec: tv_nsec as _,
497    };
498    deadline.tv_sec = deadline
499        .tv_sec
500        .checked_add((deadline.tv_nsec / 1_000_000_000) as libc::time_t)
501        .ok_or(SemError::InvalidInput)?;
502    deadline.tv_nsec %= 1_000_000_000;
503    Ok(Deadline { spec: deadline })
504}
505
506#[cfg(unix)]
507#[allow(clippy::not_unsafe_ptr_arg_deref)]
508fn sem_wait_status(handle: *mut sem_t, deadline: Option<&libc::timespec>) -> WaitStatus {
509    #[cfg(not(target_vendor = "apple"))]
510    if let Some(deadline) = deadline {
511        if unsafe { libc::sem_timedwait(handle, deadline) } == 0 {
512            WaitStatus::Acquired
513        } else {
514            match Errno::last() {
515                Errno::ETIMEDOUT => WaitStatus::TimedOut,
516                Errno::EINTR => WaitStatus::Interrupted,
517                err => WaitStatus::Error(SemError::from_errno(err)),
518            }
519        }
520    } else {
521        if unsafe { libc::sem_wait(handle) } == 0 {
522            WaitStatus::Acquired
523        } else {
524            match Errno::last() {
525                Errno::EINTR => WaitStatus::Interrupted,
526                err => WaitStatus::Error(SemError::from_errno(err)),
527            }
528        }
529    }
530
531    #[cfg(target_vendor = "apple")]
532    {
533        debug_assert!(deadline.is_none());
534        if unsafe { libc::sem_wait(handle) } == 0 {
535            WaitStatus::Acquired
536        } else {
537            match Errno::last() {
538                Errno::EINTR => WaitStatus::Interrupted,
539                err => WaitStatus::Error(SemError::from_errno(err)),
540            }
541        }
542    }
543}
544
545#[cfg(target_vendor = "apple")]
546pub enum PollWaitStep {
547    Acquired,
548    Timeout,
549    Continue(u64),
550}
551
552#[cfg(target_vendor = "apple")]
553#[allow(clippy::not_unsafe_ptr_arg_deref)]
554fn sem_timedwait_poll_step(
555    handle: *mut sem_t,
556    deadline: &libc::timespec,
557    delay: u64,
558) -> Result<PollWaitStep, SemError> {
559    if unsafe { libc::sem_trywait(handle) } == 0 {
560        return Ok(PollWaitStep::Acquired);
561    }
562    let err = Errno::last();
563    if err != Errno::EAGAIN {
564        return Err(SemError::from_errno(err));
565    }
566
567    let now = gettimeofday()?;
568    let deadline_usec = deadline.tv_sec * 1_000_000 + deadline.tv_nsec / 1000;
569    #[allow(clippy::unnecessary_cast)]
570    let now_usec = now.tv_sec as i64 * 1_000_000 + now.tv_usec as i64;
571    if now_usec >= deadline_usec {
572        return Ok(PollWaitStep::Timeout);
573    }
574
575    let difference = (deadline_usec - now_usec) as u64;
576    let mut delay = delay + 1000;
577    if delay > 20000 {
578        delay = 20000;
579    }
580    if delay > difference {
581        delay = difference;
582    }
583
584    let mut tv_delay = libc::timeval {
585        tv_sec: (delay / 1_000_000) as _,
586        tv_usec: (delay % 1_000_000) as _,
587    };
588    unsafe {
589        libc::select(
590            0,
591            core::ptr::null_mut(),
592            core::ptr::null_mut(),
593            core::ptr::null_mut(),
594            &mut tv_delay,
595        );
596    }
597    Ok(PollWaitStep::Continue(delay))
598}
rustpython_host_env/multiprocessing.rs

rustpython_host_env/
multiprocessing.rs
