feat(oneshot channel): ensure msg won't be dropped on sender side when send returns ok

[wenym1]

Motivation

Related: risingwavelabs/risingwave#6617.

When we use oneshot::Sender, it is intuitive to assume that when tx.send(msg) returns with Ok(()), the msg will have been added to the buffer and either will be consumed by the rx.await, or dropped along with dropping rx, but in either way, msg should not be dropped on the tx side.

However, in the current implementation, the assumption does not hold. It might happen that, even though tx.send(msg) returns with Ok(()), the msg is still dropped on the sender side.

The current struct of Sender is

struct Sender<T> {
    inner: Option<Arc<Inner<T>>>,
}

In the current implementation, in tx.send(msg), msg will be added to inner. Before msg gets consumed by rx, the struct Inner holds the ownership of msg. Inner is protected by Arc with reference counting to avoid memory leak. When the reference counting of Arc<Inner> decreases to 0, msg gets dropped along with Inner. However, in the following execution flow, msg will be dropped on sender side when tx.send(msg) returns Ok(()).

pub fn send(mut self, t: T) -> Result<(), T> {
    let inner = self.inner.take().unwrap();

    inner.value.with_mut(|ptr| unsafe {
        *ptr = Some(t);
    });

    if !inner.complete() {
        unsafe {
            return Err(inner.consume_value().unwrap());
        }
    }

    // assume that rx not dropped yet here, ref count of `Arc<Inner>` is 2
-> rx dropped here, ref count becomes 1

    Ok(())
-> variable `inner` gets dropped,  ref count decreases to 0, `msg` gets dropped along with dropping `Inner`, but we return with `Ok(())`
}

Solution

The problem happens due to the implicit drop of variable inner. We are not aware of whether or not dropping inner will decreases the Arc ref count to 0 and further drop the Inner.

Therefore, in this PR, we can leverage the Arc::into_inner to explicitly consume the variable inner so that we can atomically decrease the ref count and check whether the Inner will be further dropped. If Inner is going to be dropped, we can then take the msg out and return with Err(msg) if the msg has not been consumed by rx yet.

[Darksonn]

Can you please add a loom test that catches the issue you have mentioned? It would need to go somewhere in this folder. Preferably the test should be as simple as possible.

[wenym1]

Can you please add a loom test that catches the issue you have mentioned? It would need to go somewhere in this folder. Preferably the test should be as simple as possible.

I have added a loom test. The test logic is as simple as followed.

let (tx, rx) = oneshot::channel();

// thread 1
if let Err(msg) = tx.send(msg) {
    std::mem::forget(msg);
}

// thread 2
drop(rx);

We can assert that msg can only be dropped in thread 2, because in thread 1, tx.send(msg) either return with Ok(()), or call std::mem::forget(msg) on Err(msg), and drop shouldn't be called when returning Ok(()).

[wenym1]

If the newly added code in oneshot.rs is removed, we will hit this assertion.

[Darksonn]

I suspect this could be solved by having impl Drop for Receiver call consume_value if the call to self.close() observes sees the VALUE_SET bit.

[wenym1]

I suspect this could be solved by having impl Drop for Receiver call consume_value if the call to self.close() observes sees the VALUE_SET bit.

It's feasible. I have updated the code. Thanks for the suggestion!. Now no dependency will be added in this PR. @Darksonn

[Darksonn]

You can just call drop(rx) on the main thread. Spawning threads is very expensive in loom tests.

[wenym1]

I had a try.

If I simply use drop(rx) in the main thread, when I revert the changes in oneshot.rs, the assertion cannot be hit, and if I do it in a spawned thread, the assertion can be hit.

I doubt that only drop(rx) in a separate thread can cover the case in loom.

[Darksonn]

Do you hit the assertion with LOOM_MAX_PREEMPTIONS=2? That's what we run with in CI.

[wenym1]

Oops, when I set it to 2, the assertion can be hit.

Should the CONTRIBUTING.md be updated? I ran the loom test following the guide in it.

tokio/CONTRIBUTING.md

Line 195 in df77063

LOOM_MAX_PREEMPTIONS=1 LOOM_MAX_BRANCHES=10000 RUSTFLAGS="--cfg loom -C debug_assertions" \

[Darksonn]

It takes a lot longer to run with the option set to 2, and most bugs can be caught with the value set to 1. It may be better to have your test use the loom test builder to explicitly set the max preemptions to 2.

[wenym1]

LGTM. Have updated the code.

[Darksonn]

Thank you.