c/
pthread_atfork.rs

1// Copyright (c) 2026 Meng Zhuo
2// SPDX-License-Identifier: BSD-3-Clause
3
4// pthread_atfork() handler registry and fork hook, musl-compatible.
5// Included from libc/src/lib.rs; fork() must call __fork_handler(-1/0/1).
6
7#[derive(Copy, Clone)]
8struct atfork_funcs {
9    prepare: Option<unsafe extern "C" fn()>,
10    parent: Option<unsafe extern "C" fn()>,
11    child: Option<unsafe extern "C" fn()>,
12}
13
14// The registry grows with realloc; musl also allocates one node per handler.
15static mut ATFORK_FUNCS: *mut atfork_funcs = core::ptr::null_mut();
16static mut ATFORK_CAP: usize = 0;
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    if who < 0 {
41        atfork_lock();
42        let n = ATFORK_NFUNCS.load(Ordering::Acquire);
43        let base = ATFORK_FUNCS;
44        if n == 0 || base.is_null() {
45            return;
46        }
47        let mut i = n;
48        while i > 0 {
49            i -= 1;
50            let entry = &*base.add(i);
51            if let Some(f) = entry.prepare {
52                f();
53            }
54        }
55    } else {
56        let n = ATFORK_NFUNCS.load(Ordering::Acquire);
57        let base = ATFORK_FUNCS;
58        if n != 0 && !base.is_null() {
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        }
70        atfork_unlock();
71    }
72}
73
74#[no_mangle]
75pub unsafe extern "C" fn pthread_atfork(
76    prepare: Option<unsafe extern "C" fn()>,
77    parent: Option<unsafe extern "C" fn()>,
78    child: Option<unsafe extern "C" fn()>,
79) -> c_int {
80    atfork_lock();
81    let n = ATFORK_NFUNCS.load(Ordering::Relaxed);
82    if n == ATFORK_CAP {
83        let new_cap = if ATFORK_CAP == 0 { 8 } else { ATFORK_CAP * 2 };
84        let new = crate::realloc(
85            ATFORK_FUNCS as *mut c_void,
86            new_cap * core::mem::size_of::<atfork_funcs>(),
87        ) as *mut atfork_funcs;
88        if new.is_null() {
89            atfork_unlock();
90            return ENOMEM;
91        }
92        ATFORK_FUNCS = new;
93        ATFORK_CAP = new_cap;
94    }
95    *ATFORK_FUNCS.add(n) = atfork_funcs { prepare, parent, child };
96    ATFORK_NFUNCS.store(n + 1, Ordering::Release);
97    atfork_unlock();
98    0
99}
c/pthread_atfork.rs

c/
pthread_atfork.rs
