Value sent through `mpsc::Permit` should be immediately dropped when receiver is already dropped

[tabokie]

Version

Latest (1.48.0)

Platform

Doesn't matter.

Description

Even though not documented explicitly, but the current behavior is, all values stored in the channel are dropped when the sole receiver is dropped:

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

Lines 490 to 526 in 5a709e3

	fn drop(&mut self) {
	use super::block::Read::Value;
	self.close();
	self.inner.rx_fields.with_mut(|rx_fields_ptr| {
	let rx_fields = unsafe { &mut *rx_fields_ptr };
	struct Guard<'a, T, S: Semaphore> {
	list: &'a mut list::Rx<T>,
	tx: &'a list::Tx<T>,
	sem: &'a S,
	}
	impl<'a, T, S: Semaphore> Guard<'a, T, S> {
	fn drain(&mut self) {
	// call T's destructor.
	while let Some(Value(_)) = self.list.pop(self.tx) {
	self.sem.add_permit();
	}
	}
	}
	impl<'a, T, S: Semaphore> Drop for Guard<'a, T, S> {
	fn drop(&mut self) {
	self.drain();
	}
	}
	let mut guard = Guard {
	list: &mut rx_fields.list,
	tx: &self.inner.tx,
	sem: &self.inner.semaphore,
	};
	guard.drain();
	});
	}

It is natural for user to consider "receiver dropped -> values dropped" an invariant.

But it doesn't hold in this sequence of event:

• Sender acquires a permit
• Receiver is dropped
• Sender publishes a value using the permit

[Darksonn]

Interesting. I would be happy to see a PR fixing this, but I think it might be tricky to implement.

[seakerOner]

Reproduction using tracing

I made a small visual representation of the behavior so it's easier to follow along.
Behavior reproduced using tokio 1.48.0.

Here's the minimal repro:

use tokio::sync::mpsc;
use tracing::{info, instrument};
use tracing_subscriber::fmt::time::Uptime;

struct Msg(&'static str);

impl Drop for Msg {
    fn drop(&mut self) {
        info!("--> dropped message: {}", self.0);
    }
}

#[tokio::main]
#[instrument]
async fn main() {
    tracing_subscriber::fmt()
        .with_timer(Uptime::default())
        .with_target(false)
        .init();

    info!("Starting Tokio MPSC drop test... ");

    let (tx, rx) = mpsc::channel::<Msg>(1);
    let message = Msg("Banana");

    info!("Getting the permit for sending message");
    let permit = tx.reserve().await.unwrap();

    info!("Dropping Receiver before sending the message with Sender's Permit");
    drop(rx);
    info!("Dropped Receiver");

    permit.send(message);
    info!("Sent message after receiver dropped");

    drop(tx);
    info!("Sender dropped, program ending");
}

Output:

0.000109641s  INFO Starting Tokio MPSC drop test...
0.000131681s  INFO Getting the permit for sending message
0.000140271s  INFO Dropping Receiver before sending the message with Sender's Permit
0.000148221s  INFO Dropped Receiver
0.000153841s  INFO Sent message after receiver dropped
0.000159661s  INFO --> dropped message: Banana
0.000165741s  INFO Sender dropped, program ending

As shown, the value is only dropped after permit.send(), confirming that it’s not immediately destroyed when the receiver is gone.

I would like to tackle this but I'm not sure how I would do so.
My first thought was to validate inside .send() if the receiver is still alive, and if it's not, simply drop the value instead of sending it.

Would an approach likes this work? Or is there a better place in the channel internals to handle this? Adding drop logic inside a function called "send" feels weird :D

[martin-g]

the value is only dropped after permit.send()

It is really late/early here, so I may not think well but the value/message is not in the channel yet when the receiver is dropped. Why do you expect that it will be dropped earlier?
I'd expect it to be dropped only if it is sent to the channel but not yet delivered.

[seakerOner]

Maybe I expressed myself unclearly.

I don't expect the message to be dropped before calling Permit.send()
After dropping the receiver, I would expect that if I used Permit.send(), it would either fail (and drop the value) or at least drop it without attempting to send to a channel that has already been closed and drained.

Another option would be to not send the value to the channel and, instead of dropping it, return the value so the user decides what to do with it.

Right now the sender silently succeeds, which I believe goes against the invariant "receiver dropped -> values dropped".
It's not about dropping before .send(), but about handling the closed-receiver case inside .send().

I'd also need to look more closely at the code to see how that would work, the Permit assumes the receiver will never be dropped.

Does this sound more in line with what you'd expect?

[tabokie]

Here's another one:

use tokio::sync::{mpsc, oneshot};

#[tokio::main]
async fn main() {
    let (tx, rx) = mpsc::channel(1);
    
    let permit = tx.reserve().await.unwrap();

    let (inner_tx, mut inner_rx) = oneshot::channel::<()>();
    drop(rx);
    permit.send(inner_tx);
    
    // The item in channel should be dropped now.
    assert_eq!(
        inner_rx.try_recv().unwrap_err(),
        oneshot::error::TryRecvError::Closed,
        "The item in channel should be dropped now"
    );
}

My first thought was to validate inside .send() if the receiver is still alive

My thoughts too. It adds one atomic read per send(), would love see some performance results first.

[Darksonn]

I don't think the sender has any way to tell the difference between Receiver::close() and Receiver::drop(), so this cannot be easily implemented. The behavior is correct for the Receiver::close() case and that part cannot be changed.

[seakerOner]

I've investigating this issue and put together some tests and benchmarks to confirm the behavior and effect of a possible fix

Current behavior:

Using the current mpsc implementation (before any patch), I wrote two tests (one of them is the same I showed earlier in this issue) to observe the current behavior when sending via a Permit after the receiver is dropped or closed.

Test before patch

async fn test_msg_drop_timing_after_rx_dropped() {
    info!("Starting Tokio MPSC close test... ");

    let (tx, mut rx) = mpsc::channel::<Msg>(1);
    let message = Msg("Banana");

    info!("Getting the permit for sending message");
    let permit = tx.reserve().await.unwrap();

    info!("Dropping Receiver before sending the message with Sender's Permit");
    rx.close();
    drop(rx);
    info!("Dropped Receiver");

    permit.send(message);
    info!("Sent message after receiver dropped");

    drop(tx);
    info!("Sender dropped, program ending");
}

Output:

 0.000109641s  INFO Starting Tokio MPSC drop test...
 0.000131681s  INFO Getting the permit for sending message
 0.000140271s  INFO Dropping Receiver before sending the message with Sender's Permit
 0.000148221s  INFO Dropped Receiver
 0.000153841s  INFO Sent message after receiver dropped
 0.000159661s  INFO --> dropped message: Banana
 0.000165741s  INFO Sender dropped, program ending

Test before patch for the .close() scenario:

async fn test_msg_drop_timing_after_rx_closed() {
    info!("Starting Tokio MPSC drop test... ");

    let (tx, mut rx) = mpsc::channel::<Msg>(1);
    let message = Msg("Banana");

    info!("Getting the permit for sending message");
    let permit = tx.reserve().await.unwrap();

    info!("Close Receiver before sending the message with Sender's Permit");
    rx.close();
    info!("Receiver Closed");

    permit.send(message);
    info!("Sent message after receiver closed");

    drop(tx);
    info!("Sender dropped, program ending");
}

Results:

0.000128931s  INFO Starting Tokio MPSC drop test...
0.000138141s  INFO Getting the permit for sending message
0.000160361s  INFO Close Receiver before sending the message with Sender's Permit
0.000170381s  INFO Receiver Closed
0.000177801s  INFO Sent message after receiver closed
0.000185382s  INFO Sender dropped, program ending
0.000194752s  INFO --> dropped message: Banana

Patched behavior:

After applying the patch, the behavior is corrected. Messages are dropped immediately when the receiver is dropped, and the behavior remains unchanged on the .close() scenario

Tests after patch

#[maybe_tokio_test]
async fn test_permit_send_after_rx_drop_drops_value() {
    use std::sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    };

    struct FlaggedDrop(Arc<AtomicBool>);
    impl Drop for FlaggedDrop {
        fn drop(&mut self) {
            self.0.store(true, Ordering::Release);
        }
    }

    let (sender, receiver) = mpsc::channel::<FlaggedDrop>(1);

    let permit = sender.reserve().await.unwrap();

    let foo = Arc::new(AtomicBool::new(false));

    drop(receiver);

    permit.send(FlaggedDrop(foo.clone()));

    assert_eq!(foo.load(Ordering::Acquire), true); // Message is dropped immediately
}

#[maybe_tokio_test]
async fn test_permit_send_after_rx_close_still_sends_value() {
    use std::sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    };

    struct FlaggedDrop(Arc<AtomicBool>);
    impl Drop for FlaggedDrop {
        fn drop(&mut self) {
            self.0.store(true, Ordering::Release);
        }
    }

    let (sender, mut receiver) = mpsc::channel::<FlaggedDrop>(1);

    let permit = sender.reserve().await.unwrap();

    let foo = Arc::new(AtomicBool::new(false));

    receiver.close();

    permit.send(FlaggedDrop(foo.clone()));

    assert_eq!(foo.load(Ordering::Acquire), false); // Message is still sent normally

    drop(receiver);
    assert_eq!(foo.load(Ordering::Acquire), true);
}

Benchmarks

To ensure the fix didn't introduce regressions, I added two benchmarks:

fn permit_send_alive(g: &mut BenchmarkGroup<WallTime>) {
    let rt = rt();

    g.bench_function("bounded_permit_send_alive", |b| {
        b.iter(|| {
            rt.block_on(async move {
                let (tx, mut rx) = mpsc::channel::<usize>(1024);

                tokio::spawn(async move { while rx.recv().await.is_some() {} });

                for i in 0..10_000 {
                    let permit = tx.reserve().await.unwrap();
                    permit.send(i);
                }
            })
        })
    });
}

fn permit_send_dropped(g: &mut BenchmarkGroup<WallTime>) {
    let rt = rt();

    g.bench_function("bounded_permit_send_dropped", |b| {
        b.iter(|| {
            rt.block_on(async move {
                let (tx, rx) = mpsc::channel::<usize>(1024);

                let permit = tx.reserve().await.unwrap();
                drop(rx);

                permit.send(50);
            })
        })
    });
}

And ran them with Criterion:

# Baseline (master)
cargo bench --bench sync_mpsc permit_send -- --save-baseline master

# Patched version
cargo bench --bench sync_mpsc permit_send -- --baseline master

Results

    Finished `bench` profile [optimized] target(s) in 0.08s
     Running sync_mpsc.rs (/home/seaker/Documents/open source/tokio/target/release/deps/sync_mpsc-09b4a052af7e16dd)

permit_send/bounded_permit_send_alive
                        time:   [724.81 µs 736.20 µs 747.64 µs]
                        change: [-8.1294% -5.5881% -3.2245%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 100 measurements (1.00%)
  1 (1.00%) high mild

permit_send/bounded_permit_send_dropped
                        time:   [250.33 ns 250.93 ns 251.59 ns]
                        change: [-7.6936% -7.4763% -7.2593%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  5 (5.00%) high mild
  1 (1.00%) high severe

No measurable regressions observed. (I don't understand why it says the performance has improved, if someone could explain what I am not seeing I would appreciate it).

Criterion report

You can see created HTML below in PDF Format

permit_send Summary - Criterion.rs.pdf
permit_send_bounded_permit_send_dropped - Criterion.rs.pdf
permit_send_bounded_permit_send_alive - Criterion.rs.pdf

TLDR:

• The patch fixes the drop semantics when calling Permit::send() after the receiver is dropped.
• .close() behavior remains the same.
• No performance regressions observed in benchmarks.

The fix introduces an AtomicBool flag in the channel state that tracks whether the receiver has been dropped. When Receiver is dropped, the flag is set to true. Then, in Permit::send(), this flag is checked before proceeding — if it’s true, the value is immediately dropped instead of being sent into a dropped receiver.

If everything looks good I can create a PR so everyone can review the changes in detail.

[happycoder0011]

Does this solves the problem
Key Changes (demonstrated via Parseltongue temporal editing):

// BEFORE (current code)

  fn send(&self, value: T) {
      self.tx.push(value);          // Race condition possible here
      self.rx_waker.wake();
  }

// AFTER (fixed code)

fn send(&self, value: T) {
     // FIXED: Check receiver liveness after acquiring permit but before sending
     if !self.inner.rx_alive.load(Acquire) {
         // Receiver is gone, drop the value immediately to prevent leak
         drop(value);
         return;
     }
     self.tx.push(value);
     self.rx_waker.wake();
 }

Additional Required Changes:

1. Add rx_alive: AtomicBool to Chan<T,S> struct
2. Initialize to true in channel()
3. Set to false in Rx::drop() before draining

[seakerOner]

Hey @happycoder0011 , I am a bit confused by your proposed solution.

The Permit.send() function doesn't have that structure; It consumes 'self' and calls self.chan.send(value) followed by mem::forget(self).
I also believe you can't directly access the inner fields from there; In my solution I was required to add a helper method.

I already have an implementation that should handle this correctly, with tests and benchmarks to confirm the behavior.
Probably best to wait for one of the maintainers to weigh in before moving forward.
@Darksonn any thoughts?

[Darksonn]

@happycoder0011 That solution doesn't look right. If the receiver is dropped in parallel with that, I believe the value may not get dropped.