Skip to main content

c/
pthread_atfork.rs

1// pthread_atfork() handler registry and fork hook, musl-compatible.
2// Include in libc/src/lib.rs; fork() must call __fork_handler(-1/0/1).
3
4const MAX_ATFORK: usize = 64;
5
6#[derive(Copy, Clone)]
7struct atfork_funcs {
8    prepare: Option<unsafe extern "C" fn()>,
9    parent: Option<unsafe extern "C" fn()>,
10    child: Option<unsafe extern "C" fn()>,
11}
12
13const EMPTY: atfork_funcs = atfork_funcs { prepare: None, parent: None, child: None };
14
15// ponytail: static array, no malloc; upgrade to linked list + malloc if >64 handlers needed
16static mut ATFORK_FUNCS: [atfork_funcs; MAX_ATFORK] = [EMPTY; MAX_ATFORK];
17static ATFORK_NFUNCS: AtomicUsize = AtomicUsize::new(0);
18static ATFORK_LOCK: AtomicI32 = AtomicI32::new(0);
19
20#[inline]
21unsafe fn atfork_lock() {
22    while ATFORK_LOCK
23        .compare_exchange_weak(0, 1, Ordering::Acquire, Ordering::Relaxed)
24        .is_err()
25    {
26        core::hint::spin_loop();
27    }
28}
29
30#[inline]
31unsafe fn atfork_unlock() {
32    ATFORK_LOCK.store(0, Ordering::Release);
33}
34
35// who < 0: prepare (reverse order, acquires lock)
36// who == 0: parent (forward order, releases lock)
37// who != 0: child (forward order, releases lock)
38#[no_mangle]
39pub unsafe extern "C" fn __fork_handler(who: c_int) {
40    let n = ATFORK_NFUNCS.load(Ordering::Relaxed);
41    if n == 0 {
42        return;
43    }
44
45    let base = core::ptr::addr_of!(ATFORK_FUNCS).cast::<atfork_funcs>();
46    if who < 0 {
47        atfork_lock();
48        let n = ATFORK_NFUNCS.load(Ordering::Relaxed);
49        let mut i = n;
50        while i > 0 {
51            i -= 1;
52            let entry = &*base.add(i);
53            if let Some(f) = entry.prepare {
54                f();
55            }
56        }
57    } else {
58        let n = ATFORK_NFUNCS.load(Ordering::Relaxed);
59        for i in 0..n {
60            let entry = &*base.add(i);
61            if who == 0 {
62                if let Some(f) = entry.parent {
63                    f();
64                }
65            } else if let Some(f) = entry.child {
66                f();
67            }
68        }
69        atfork_unlock();
70    }
71}
72
73#[no_mangle]
74pub unsafe extern "C" fn pthread_atfork(
75    prepare: Option<unsafe extern "C" fn()>,
76    parent: Option<unsafe extern "C" fn()>,
77    child: Option<unsafe extern "C" fn()>,
78) -> c_int {
79    atfork_lock();
80    let n = ATFORK_NFUNCS.load(Ordering::Relaxed);
81    if n >= MAX_ATFORK {
82        atfork_unlock();
83        return ENOMEM;
84    }
85    let dst = core::ptr::addr_of_mut!(ATFORK_FUNCS).cast::<atfork_funcs>().add(n);
86    *dst = atfork_funcs { prepare, parent, child };
87    ATFORK_NFUNCS.store(n + 1, Ordering::Relaxed);
88    atfork_unlock();
89    0
90}