sync: add Sender<T>::closed future

[evanrittenhouse]

Motivation

Adds a closed Future to broadcast::Sender<T>, similar to the oneshot or mpsc variants, which completes when all subscribed receivers have been dropped.

Solution

This is a simple poll_fn which wraps a check around shared.tail.rx_cnt, returning Ready if the remaining count is 0.

Closes: #6649

[tglane]

I might be wrong here butI think this future will only get polled once because you are not waking up the task again.

Consider this example:

#[tokio::main]
async fn main() {
    let (tx, mut rx1) = tokio::sync::broadcast::channel(16);

    tokio::spawn(async move {
        let r = rx1.recv().await.unwrap();
        println!("Rx1 received: {}", r);
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
    });

    tx.send(20).unwrap();

    println!("Waiting for subscribers to drop");
    tx.closed().await;
    println!("Subscribers dropped");
}

Here the call to closed will block because it will only be polled once right when you call it. But the subscriber is not dropped at that point, so the poll_fn will return Pending. But because the future is never be polled again, this will not resolve after that.

[Darksonn]

Indeed, this is not correct. When returning Pending, you must register for wakeups so that you are polled again when progress can be made.

Compare with the mpsc channel:

tokio/tokio/src/sync/mpsc/chan.rs

Lines 58 to 59 in 4825c44

	/// Notifies all tasks listening for the receiver being dropped.
	notify_rx_closed: Notify,

[evanrittenhouse]

True, and I guess poll_fn is flagged as unstable in CI anyway. I'll find another way around this - thanks both.

[Darksonn]

Tokio has a poll_fn implementation you can use if you find that you need it.

[Darksonn]

What's the status on this?

[evanrittenhouse]

Work has been nuts - haven't had much time for anything else. This is on my backlog though

[Darksonn]

No rush from my side. Just wanted to check in.

[evanrittenhouse]

Pushed up what should be closer to the proper implementation. I'd like to add a task.is_woken() check after each rx is dropped in the integration test, but for some reason, dropping rx2 doesn't wake the task (though it does make it Ready).

I'll dig more into that when I have a bit more time, as well as fix the various CI issues and clean up the code (for example the tail lock will drop when it goes out of scope).

[Darksonn]

The poll fn needs to interact with the notified() future for this to work. Check out this for inspiration:

tokio/tokio-util/src/sync/cancellation_token.rs

Lines 303 to 319 in 5b9a290

	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
	let mut this = self.project();
	loop {
	if this.cancellation_token.is_cancelled() {
	return Poll::Ready(());
	}
	// No wakeups can be lost here because there is always a call to
	// `is_cancelled` between the creation of the future and the call to
	// `poll`, and the code that sets the cancelled flag does so before
	// waking the `Notified`.
	if this.future.as_mut().poll(cx).is_pending() {
	return Poll::Pending;
	}
	this.future.set(this.cancellation_token.inner.notified());
	}

You should be able to do something similar without creating a whole future struct if you use tokio::pin! on the notified() future.

[evanrittenhouse]

If we now notify only once, when the last receiver drops, can we just await the Notified directly? I originally had this block because I notified on every Rx drop, so had to include an additional check to ensure that we had just dropped the last Rx.

I'm also trying to get a better understanding of the internals here. For my own understanding - this doesn't work because, after the first call, the internal Future state machine is progressed to the point where we're already in the poll_fn and have a Pending from a non-terminal Rx dropping. That means that we'd never await the notified() again. Is that correct?

[Darksonn]

Yes, Notified doesn't have spurious wakeups in the scenario we are using it, so it should work. Well, except for the fact that we probably want it to return immediately if it's already closed when we call it. We can do it like this:

loop {
    let notified = self.shared.notify_rx_drop.notified();
    if self.is_closed() { return; }
    notified.await;
}