Sanity and correctnes of my toy `OnceLock` implementation?

code review
1

Hi! I got a bit bored and decided to speed run my own implementation of OnceLock<T>

I wrote a small test in the main function, and it works, although I'm not sure if it's 100% correct.

The code:

#![allow(unused)]

use std::{
    cell::UnsafeCell,
    mem::MaybeUninit,
    sync::atomic::{AtomicU32, Ordering::*},
    thread,
    time::Duration,
};

const UNINIT: u32 = 0;
const INITIALIZING: u32 = 1;
const INITIALIZED: u32 = 2;

struct Once<T> {
    x: UnsafeCell<MaybeUninit<T>>,
    flag: AtomicU32,
}

unsafe impl<T: Sync> Sync for Once<T> {}
unsafe impl<T: Sync> Send for Once<T> {}

impl<T> Once<T> {
    pub const fn new() -> Self {
        Self {
            x: UnsafeCell::new(MaybeUninit::uninit()),
            flag: AtomicU32::new(UNINIT),
        }
    }

    pub fn get_or_init<F>(&self, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        if let Ok(_) = self
            .flag
            .compare_exchange(UNINIT, INITIALIZING, Acquire, Relaxed)
        {
            let x = f();
            unsafe {
                (&mut *self.x.get()).write(x);
            }
            self.flag.store(INITIALIZED, Release);
            atomic_wait::wake_all(&self.flag);
        }

        while self.flag.load(Relaxed) == INITIALIZING {
            atomic_wait::wait(&self.flag, INITIALIZING);
        }

        /// Safety:
        /// If we get to this point, it means that the value has been initialized.
        /// Therefore we can safely read it and return an immutable reference to it.
        unsafe { (&*self.x.get()).assume_init_ref() }
    }

    pub fn get(&self) -> Option<&T> {
        if self.flag.load(Relaxed) != INITIALIZED {
            return None;
        }

        /// Safety: the value is initialized, so we can safely read it.
        Some(unsafe { (&*self.x.get()).assume_init_ref() })
    }
}

static INT: Once<i32> = Once::new();

fn main() {
    let handle = thread::spawn(|| {
        let _ = INT.get_or_init(|| {
            thread::sleep(Duration::from_secs(3));
            42
        });
    });

    thread::sleep(Duration::from_millis(100));

    assert_eq!(INT.get(), None);

    let x = INT.get_or_init(|| 0);
    assert_eq!(*x, 42);

    handle.join();

    assert_eq!(*INT.get().unwrap(), 42);
    assert_eq!(*INT.get_or_init(|| 43), 42);
}

Thanks in advance!

1 Like
2

When checking that the flag is set before getting the value reference, the ordering needs to be Acquire in order to synchronize with the Release store and be allowed to read from the UnsafeCell. Without that synchronization edge, you have a data race. The futex might imply sufficient synchronization in the get_or_init case, but I wouldn't rely on that, and the get case absolutely needs an Acquire.

Miri contains simple data race / synchronization detection, so try running your test cases under miri. It's more invasive, but supporting testing under loom also helps build confidence that your orderings are sufficient.

4 Likes
3

I'm definitely not qualified to audit unsafe code, but here are a few observations (which may be inaccurate or entirely incorrect [1]).

Your Once type leaks any initialized value. This is safe, but not friendly. It should have impl Drop unless the leaking is intentional for some specific reason (which is not documented).

This seems to have the wrong constraint:

By requiring Sync, this allows sending a type that is not Send but is Sync to another thread, violating its invariants. One such type is MutexGuard. It has this invariant for POSIX hosts, according to this legendary post: Example of a type that is not `Send`? - #3 by Yandros

Which means it can (in theory) be exploited to send a MutexGuard to another thread to unlock a mutex it didn't originally lock, in violation of the pthreads contract:

fn main()
    static M: Mutex<()> = Mutex::new(());
    let o = Once::new();
    o.get_or_init(|| M.lock().unwrap());
    thread::spawn(move || drop(o)).join().unwrap();
    M.lock().unwrap();
}

We create a static mutex and then wrap the guard in a Once and drop it in another thread.

This is currently sound as-written -- because Once does not impl Drop! It leaks the guard and the mutex is never unlocked. The last line deadlocks.

Changing the constraint to impl<T: Send> Send [2] will allow a sound Drop impl.

  1. Standard disclaimer. ↩︎

  2. Constraining the Sync impl with T: Sync is perfectly fine, AFAIK. ↩︎

3 Likes
4

Thanks a lot, that's very informative!

5

That's something I really overlooked, thanks.

Initially, I required only T: Sync for Send for Once<T>, because I thought that the value won't be moved (hence there is no way to get a mutable access, or drop it (I was only thinking of static variables :'/)). But your example proves me wrong, thanks.

Edit: Hmm, now I'm not sure that the bound is wrong. See, as there is no Drop implementation, nothing happens (so that's, the value is never accessed (mutably and immutably) in a different thread. However, if we add the Drop implementation, then that bound becomes unsound.

@parasyte