Lib.rs

Wipe-on-fork 为 Rust 提供的 OnceCell, LazyCell, Once, OnceLock, LazyLock

关于金字塔的创造者是谁，埃及人还是外星人，一直有阴谋论。同样，几千年后，我们可以预期未来的人会问：互联网是人类还是外星人发明的？

当时的史学家可以引用 HTTP 状态码 418 I'm a teapot，这个状态码最初是一个愚人节玩笑，用来证明 HTTP 是由喝茶、制作茶壶并具有幽默感的人类发明的。创建 save418.com 网站的 Shane Brunswick，这个网站对于保留 418 状态码至关重要，他说：“这是提醒我们，计算机背后的过程仍然是由人类制造的。”

类似的事情也发生在计算机科学的其他领域。其中之一就是 fork()。它是一个进程创建另一个进程的方式。在 HotOS 2019 上，四位享有盛誉的计算机系统研究人员——Andrew Baumann、Jonathan Appavoo、Orran Krieger 和 Timothy Roscoe——在他们题为 A fork() in the road 的论文中，讨论了为什么人们应该避免使用 fork()，尽管它在 POSIX 和操作系统设计中是核心，并且被广泛使用。平行宇宙可能没有 fork()。

由于 fork() 是一种低级操作系统原语，它以使应用程序几乎透明的这种方式进行更改，因此它始终存在兼容性问题。对于 Rust 来说，这一直是个头疼的问题（见 https://github.com/rust-lang/rust/issues/6930、https://github.com/rust-lang/rust/issues/9373、https://github.com/rust-lang/rust/issues/9568、https://github.com/rust-lang/rust/issues/16799）。

这就需要替换上述的 lazy_static!，以便子进程可以运行自己的，而不是从父进程中继承。这与 Linux 中的“在 fork 期间擦除”（在 Linux madvise 函数中称为“wipe-on-fork”的概念密切相关），该函数允许程序在 fork 时通知操作系统擦除页面。

MADV_WIPEONFORK（自 Linux 4.14 以来）

在 fork(2) 之后，向子进程提供该范围内已填充零的内存。这在为服务器进行 fork 以确保敏感的进程数据（例如，PRNG 种子、加密机密等）不会传递给子进程时很有用。

因此，我们采用了这种命名约定，并创建了许多数据结构。

此存储库重新实现了（复制粘贴，但进行了修改）以下来自 Rust 的 std 库的结构

大部分代码，包括文档测试，都是从rust-lang/rust中的 Rust std 库中复制粘贴的。我们之所以这样做，而不是以黑盒方式使用现有的原语——这始终是首选选择——是因为（1）其中一些仍在等待稳定，（2）一些必要的类型或函数仅可在 std 包中使用，（3）我们的一些更改需要更多底层操作。

使用 OnceCell、LazyCell、Once、OnceLock、LazyLock 的擦除-on-fork 版本类似于它们的 keep-on-fork 对应版本。需要注意的是，这些擦除-on-fork 版本并不是“更好”或“更通用”的实现。某些应用程序会 特别 需要擦除-on-fork，而其他应用程序会 特别 需要 keep-on-fork。这就是为什么我们在前缀 WipeOnFork* 中包含，以提醒它们虽然相关但基本不同，因为它们都涉及 fork()。

请注意，fork() 并不是创建子进程的唯一解决方案。实际上，一个更受欢迎的解决方案（尽管不那么方便）是使用 posix_spawn() 创建新进程。这已经在 Dask 中使用，但在 Ray 中没有使用（见 https://github.com/ray-project/ray/issues/13568 中的讨论）。使用擦除-on-fork 原语是为了在可组合性上提供兼容性，因为任何进程的任何线程都可以调用 fork()，而更少破坏性的解决方案就像 线程安全 一样，编写具有 fork 安全性 的代码。

fork 检测

检测后台分叉有两大方法。

• 检查代码是否在另一个进程ID (PID)下运行，该进程ID可以从以下代码获取：std::process::id()
• 通过pthread_atfork()注册分叉处理程序

我们最终没有采用PID方法，因为它有一个固有的限制。在Unix中，无法保证PID不会重复。事实上，假设整个操作系统已经拥有pid_max - 3个进程（注意：由于64位系统中的pid_max是2^22，这非常不可能）

• 父亲的PID为a
• 父亲分叉并创建了一个子进程，其PID为b
• 子进程分叉并创建了一个孙子进程，其PID为c
• 父亲去世，留下a可供操作系统重用
• 孙子进程分叉，可以期待为曾孙进程获得PID a
• 如果父亲创建并初始化了一个Once，并且这个Once被儿子和孙子保持不变，当曾孙进程第一次使用它时，它无法区分这个Once是否应该被清除。

尽管这非常罕见，因为大多数消费级内存不太可能达到pid_max进程的数量，但我们选择采用更稳健的方法。

我们引入了“代”的概念。当一个分叉时清除对象首次初始化时，它将是第一代（即，具有代ID 0）。这个代ID存储在一个全局变量中。

pub struct GenerationCounter {
    pub(crate) gen: Mutex<Option<u64>>,
}

// implementations of `GenerationCounter`

pub(crate) static GENERATION: GenerationCounter = GenerationCounter::new();

我们使用pthread_atfork()注册分叉处理程序。每次分叉时，我们都要求增加这个计数器。因为未来的代将继承这个分叉处理程序。

unsafe extern "C" fn update_generations() {
    let mut lock = GENERATION.gen.lock().unwrap();
    if lock.is_some() {
        *lock = Some(lock.unwrap() + 1);
    } else {
        panic!("The generation counter is expected to have started.");
    }
}

impl GenerationCounter {
    pub fn get(&self) -> u64 {
        // code before pthread_atfork
        unsafe {
            libc::pthread_atfork(None, None, Some(update_generations));
        }
        // code after pthread_atfork
    }
}

这解决了问题，因为这里的曾孙进程保证具有代ID为3。

实现细节

现在我们突出显示每个是如何实现的。

OnceCell

回忆一下，std::cell::OnceCell的实现如下。

pub struct OnceCell<T> {
    inner: UnsafeCell<Option<T>>,
}

我们的实现如下。

pub struct WipeOnForkOnceCell<T> {
    generation_id: Cell<Option<u64>>,
    inner: UnsafeCell<Option<T>>,
    _not_send_sync: PhantomData<*const ()>,
}

我们使用Cell来托管Option<u64>，这样我们就可以在只有对WipeOnForkOnceCell<T>的非可变引用的情况下修改代ID。使用_not_send_sync是一种黑客手段，用于实现不支持在标准库之外的负特征!Sync。有关此黑客手段的讨论，请参阅此处。

#[inline]
fn wipe_if_should_wipe(&self) {
    if self.check_if_should_wipe() {
        self.generation_id.set(None);
        unsafe {
            *self.inner.get() = None;
        }
    }
}

enum State<T, F> {
    Uninit(F),
    Init(T),
    Poisoned,
}

pub struct LazyCell<T, F = fn() -> T> {
    state: UnsafeCell<State<T, F>>,
}

enum State<T, F> {
    Uninit(F),
    Init(T, F),
    Poisoned,
}

pub struct WipeOnForkLazyCell<T, F = fn() -> T> {
    generation_id: Cell<Option<u64>>,
    state: UnsafeCell<State<T, F>>,
    _not_send_sync: PhantomData<*const ()>,
}

#[inline]
fn wipe_if_should_wipe(&self) {
    if self.check_if_should_wipe() {
        self.generation_id.set(None);
        
        let is_state_init = unsafe {
            match *self.state.get() {
                State::Init(_, _) => true,
                _ => false,
            }
        };
        
        if is_state_init {
            let state = unsafe { &mut *self.state.get() };
            let State::Init(_, f) = core::mem::replace(state, State::Poisoned) else {
                unreachable!()
            };
            
            unsafe { self.state.get().write(State::Uninit(f)) };
        }
    }
}

pub struct Once {
    inner: sys::Once,
}

pub struct Once {
    state: AtomicU32,
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum State {
    Incomplete,
    Poisoned,
    Running,
    Complete,
}

pub struct WipeOnForkOnce {
    generation_id: Mutex<Option<u64>>,
    state: Mutex<State>,
}

pub struct Once {
    state: Cell<State>,
}

#[inline]
fn wipe_if_should_wipe(&self) {
    let mut lock = self.generation_id.lock().unwrap();

    let res = match *lock {
        None => false,
        Some(generation_id) => generation_id != crate::utils::GENERATION.get(),
    };

    if res {
        *lock = None;
        *self.state.lock().unwrap() = State::Incomplete;
    }
}

pub struct OnceLock<T> {
    once: Once,
    value: UnsafeCell<MaybeUninit<T>>,
    _marker: PhantomData<T>,
}

pub struct WipeOnForkOnceLock<T> {
    once: WipeOnForkOnce,
    value: UnsafeCell<Option<T>>,
    _marker: PhantomData<T>,
}

unsafe impl<#[may_dangle] T> Drop for OnceLock<T>  {
    // the implementation 
}

union Data<T, F> {
    value: ManuallyDrop<T>,
    f: ManuallyDrop<F>,
}

pub struct LazyLock<T, F = fn() -> T> {
    once: Once,
    data: UnsafeCell<Data<T, F>>,
}

pub struct WipeOnForkLazyLock<T, F = fn() -> T> {
    once: WipeOnForkOnce,
    func: UnsafeCell<ManuallyDrop<F>>,
    data: UnsafeCell<ManuallyDrop<Option<T>>>,
}