De-facto linear types and `Arc`. Need for an “unwrap_or_drop” API?

On an unrelated train of thought about what the difference between Pin<Rc<T>> and Rc<T> is, I learned about the try_unwrap(this: Rc<T>) -> Result<T, Rc<T>> function on Rc and similarly on Arc (which could arguably be named try_into_inner, but that ship has sailed).

When you want to have linear types (in the sense of types that cannot be explicitly or implicitly dropped) in Rust, for example because proper dropping needs extra arguments, can fail, or is async, etc., a function like try_unwrap is super useful to get something back out of an Rc without dropping that thing. When you have a “linear” type today, that e.g. uses a drop bomb, you can do

fn get_rid_of(x: Rc<MyLinearType>) {
    match Rc::try_unwrap(x) {
        Ok(inner) => inner.safely_discard_with("extra argument"),
        Err(x) => drop(x) // won't call destructor of `MyLinearType`
                          // because there’s another `Rc` left
}

This approach does however not work anymore when you have an Arc, since an interleaved execution of two get_rid_of functions, implemented analogously for Arc, could both first determine there's more than one copy of the Arc left, and then both drop their Arc. I feel like we’re missing some API here like

fn unwrap_or_drop<T>(this: Arc<T>) -> Option<T> { /* ... */ }

whose implementation is essentially an atomic version of

{
    Arc::try_unwrap(this).ok()
}

For convenience one might also want to have a

fn drop_with<T>(this: Arc<T>, f: impl FnOnce(T)) { /* ... */ }

or, for supporting unsized types that only panic on drop in certain states but can be mutated back into a safe state

fn drop_with_mut<T: ?Sized>(this: Arc<T>, f: impl FnOnce(&mut T)) { /* ... */ }

As an alternative for these drop_with* functions, one could of course also just use an Arc<Wrapper<T>> where Wrapper makes the extra function call on drop. Of course drop_with_mut has the advantage of accepting closures with shorter life-time.


But I digress. My main question here is just:

  • Do you think, unwrap_or_drop is a reasonable extension of the Arc API?
    • Is there any way I’ve overlooked how the same can be archieved with the existing API or any technical reason I’ve overlooked why an atomic unwrap_or_drop doesn’t work?
  • Should it be added to Rc, too, for consistency?
  • If we want this, how is the process on such API additions? Can a PR with a non-stabilized implementation be created and merged directly or would this need an RFC first?
  • And if you care thinking about those as-well, what do you think of the drop_with* functions?
    • (I’m open for suggestions for different names, too.)
3 Likes

My immediate gut reaction is "how would you do this at all." But going into it, I think it is implementable, fairly simply:

    // comments removed for space
    fn drop(&mut self) {
        if self.inner().strong.fetch_sub(1, Release) != 1 { return; }
        acquire!(self.inner().strong);
        unsafe { self.drop_slow(); }
    }

    unsafe fn drop_slow(&mut self) {
        unsafe { ptr::drop_in_place(Self::get_mut_unchecked(self)) };
        drop(Weak { ptr: self.ptr });
    }

    pub fn try_unwrap(this: Self) -> Result<T, Self> {
        if this.inner().strong.compare_exchange(1, 0, Relaxed, Relaxed).is_err() {
            return Err(this);
        }
        acquire!(this.inner().strong);
        unsafe {
            let elem = ptr::read(&this.ptr.as_ref().data);
            let _weak = Weak { ptr: this.ptr };
            mem::forget(this);
            Ok(elem)
        }
    }

    pub fn unwrap_or_drop(this: Self) -> Option<T> {
        let this = ManuallyDrop::new(this);
        if this.inner().strong.fetch_sub(1, Release) != 1 { return None; }
        acquire!(this.inner().strong);
        unsafe {
            let elem = ptr::read(&this.ptr.as_ref().data);
            let _weak = Weak { ptr: this.ptr };
            Ok(elem)
        }
    }

I'm not 100% sure on the memory ordering of unwrap_or_drop, but it should work as intended with the correct one.

1 Like

Bike shedding: I'd use into in the name since you want the Arc consumed regardless -- probably into_inner. That's a method name with a lot of precedence, though it's not usually an Option return, but that's the different nature of reference counting.

1 Like

I think you need to add a mem::forget(this) before return None in unwrap_or_drop.

Whoops, you're right. Fixed.

If we call this into_inner, then try_unwrap should be deprecate-renamed to try_into_inner. I personally think that keeping these two operations named the same thing is more important than the choice of into_inner vs unwrap. (Though perhaps using unwrap suggests there should be a panicking form of Arc::unwrap? The current way to get that behavior is arc.try_unwrap().unwrap(), which is silly.)

1 Like

you mean try_unwrap

Obviously having just started graduate school has done a number on my basic attention :sweat:

Fixed.

1 Like

I guess, we could combine all those implementations, something like this:

// private helper function can take `self`
impl<T: ?Sized> Arc<T> {
    // #inline[always] // <- uncomment if we want to make this explicit
    unsafe fn drop_with_ptr<F>(&mut self, f: F)
    where
        F: FnOnce(*mut T),
    {
        if self.inner().strong.fetch_sub(1, Release) != 1 {
            return;
        }
        acquire!(self.inner().strong);
        unsafe { self.drop_with_mut_slow() };
    }

    #[inline(never)]
    unsafe fn drop_with_ptr_slow<F>(&mut self, f: F)
    where
        F: FnOnce(*mut T),
    {
        unsafe { f(Self::get_mut_unchecked(self)) };
        drop(Weak { ptr: self.ptr });
    }
}

unsafe impl<#[may_dangle] T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        unsafe { self.drop_with_ptr(|p| ptr::drop_in_place(p)) };
    }
}

impl<T> Arc<T> {
    pub fn unwrap_or_drop(this: Self) -> Option<T> {
        let this = ManuallyDrop::new(this);
        let mut result = None;
        unsafe { this.drop_with_ptr(|p| result = ptr::read(p)) }
        result
    }

    pub fn drop_with<F>(this: Self, f: F)
    where
        F: FnOnce(T),
    {
        let this = ManuallyDrop::new(this);
        unsafe { this.drop_with_ptr(|p| f(ptr::read(p))) }
    }
}

impl<T: ?Sized> Arc<T> {
    pub fn drop_with_mut<F>(this: Self, f: F)
    where
        F: FnOnce(&mut T),
    {
        let this = ManuallyDrop::new(this);
        unsafe {
            this.drop_with_ptr(|p| {
                // like
                // defer!(ptr::drop_in_place(p));
                // but I suspect we don’t have scopeguard in std
                struct Guard<T: ?Sized>(*mut T);
                impl<T: ?Sized> Drop for Guard<T> {
                    fn drop(&mut self) {
                        unsafe { ptr::drop_in_place(self.0) }
                    }
                }
                let guard = Guard(p);
                
                f(&mut *p);
            })
        }
    }
}

I’m not really sure about the importance of the non-inlined drop_slow function and how useful it is to have the same thing in the unwrap_or_drop implementation, etc.

try_unwrap is of course still implemented separately.

I turned a similar implementation as the one @CAD97 posted into a PR. first contribution, let’s go! I didn’t get much feedback here on my questions regarding Rc, drop_with*, or general procedures around adding new things to the standard library. But I mostly care about the unwrap_or_drop for Arc, so I’m trying to get that one moving, learning the best procedures on the way, and leaving matters like possibly renaming try_unwrap or whether Rc needs the same API additions as side questions that need not be addressed if nobody cares too much.

I forgot (i.e. wasn’t aware of) another use-case: Types where drop can cause stack overflows!

See this comment on my PR:

I’ve found another use-case while reading: this book about linked lists in rust.

The author presents code like this

impl<T> Drop for List<T> {
    fn drop(&mut self) {
        let mut head = self.head.take();
        while let Some(node) = head {
            if let Ok(mut node) = Rc::try_unwrap(node) {
                head = node.next.take();
            } else {
                break;
            }
        }
    }
}

to avoid stack overflows on drop. This code does use try_unwrap and drop the Rc in the Err case. On the next page they state “In order to get thread safety, we have to use Arc. [...] All we need to do [...] is replace every reference to Rc with std::sync::Arc. That's it. We're thread safe. Done!”

At least they acknowledge that in Rust, “[...] it's impossible to cause data races, (not to be mistaken with the more general issue of race conditions).” They do not seem to notice that by replacing Rc with Arc, they are introducing a race condition themselves that can cause a stack overflow.

This example can of course be solved by unwrap_or_drop. Since it comes from starting out with Rc and only migrating the code to Arc afterwards, it convinces me that Rc should get a version of unwrap_or_drop, too.