Safe abstraction for storing a sync::Weak

I am writing a small abstraction for storing a sync::Weak. I would probably normally use a Mutex<Weak<T>>, but I want to support no_std without other dependencies. I would appreciate a second set of eyes to make sure I haven't made any mistakes.

struct FlagPole {
    ptr: AtomicPtr<AtomicUsize>,
}

impl Drop for FlagPole {
    fn drop(&mut self) {
        let flag_ptr: *mut AtomicUsize = *self.ptr.get_mut();
        // drop one weak reference
        unsafe {
            Weak::from_raw(flag_ptr);
        }
    }
}

impl FlagPole {
    fn new() -> Self {
        Self {
            ptr: AtomicPtr::new(Weak::new().into_raw()),
        }
    }

    fn set(&self, value: Weak<AtomicUsize>) {
        let flag_ptr = value.into_raw();
        let previous = self.ptr.swap(flag_ptr, Ordering::AcqRel);
        // drop one weak reference for previous value
        unsafe {
            Weak::from_raw(previous);
        }
    }

    fn get(&self) -> Weak<AtomicUsize> {
        let flag_ptr = self.ptr.load(Ordering::Acquire);
        let current = unsafe { Weak::from_raw(flag_ptr) };
        // increment one weak ref before returning, so the pointer
        // stored in the atomic remains valid
        mem::forget(Weak::clone(&current));
        current
    }
}

I think there can be a race condition if:

  1. (thread 1) get loads a pointer.
  2. (thread 2) set replaces the pointer and drops the previous Weak
  3. (thread 1) get increments the weak reference count for the Weak that has already been dropped.
3 Likes

Great catch. My actual use case only needs to call set once, so perhaps I can use a compare_exchange to only replace the value if the original is a default (that doesn't need to be dropped).

That might work but if that is all you need then maybe the once_cell crate is a better alternative, it seems to have some no_std support.

once_cell does not, I believe, support the sync version of it's titular type on no_std, and related media seems to indicate, if it did, it would likely do so with a spin-lock, which is not ideal. It is also an extra dependency, even if it is an easy one to justify.

Spin locks are ideal solution for short operations. Even general purpose locks do some spin lock for the first few tries as an optimization. Storing/cloning out Weak is definitely a short operation.

I'd use spin::RwLock for it. It's even possible that it outperform the std one for this specific case, but only benchmark can prove it.

After reading once_cell's docs in more detail it does seem to require boxing the Weak reference which isn't ideal. It has a list of related crates however and the lazycell crate seems to do all that is required. It is small, quite popular and has support for no_std. Its AtomicLazyCell::fill method can return Err(T) even though filled returns false if it is currently being initlized from another thread, but just busy waiting a few times (in a loop) should solve that.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.