Holding an OwnedPermit for mpsc deadlocks other senders, even when there are more free slots than permits

[Pr0methean]

Version

│   │   │   │   ├── tokio v1.48.0
│   │   │   │   │   └── tokio-macros v2.6.0 (proc-macro)
│   │   │   │   ├── tokio-util v0.7.16
│   │   │   │   │   └── tokio v1.48.0 (*)
│   │   │   ├── tokio v1.48.0 (*)
│   │   │   ├── tokio v1.48.0 (*)
│   │   ├── tokio v1.48.0 (*)
│   │   │   ├── tokio v1.48.0 (*)
│   ├── tokio v1.48.0 (*)
│   │   ├── tokio v1.48.0 (*)
│   │   ├── tokio-native-tls v0.3.1
│   │   │   └── tokio v1.48.0 (*)
│   ├── tokio v1.48.0 (*)
│   ├── tokio-native-tls v0.3.1 (*)
├── tokio v1.48.0 (*)

Platform
Linux [private hostname] 6.14.0-33-generic #33~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Fri Sep 19 17:02:30 UTC 2 x86_64 x86_64 x86_64 GNU/Linux

Description
When a receiver holds a send permit, e.g. to send the next received item back to the channel in case of a retryable error, other senders are deadlocked, even when the channel has multiple free slots and only one outstanding permit.

I tried this code:

use tokio::sync::mpsc::{OwnedPermit, Receiver, Sender};

pub struct PushbackReceiver<T> {
    sender: Sender<T>,
    receiver: Receiver<T>,
    permit: Option<OwnedPermit<T>>,
}

impl<T> PushbackReceiver<T> {
    pub fn new(receiver: Receiver<T>, sender: &Sender<T>) -> Self {
        let sender = sender.clone();
        let permit = sender.clone().try_reserve_owned().ok();
        PushbackReceiver {
            receiver,
            sender,
            permit,
        }
    }

    pub async fn recv(&mut self) -> T {
        let result = self.receiver.recv().await.unwrap();
        if self.permit.is_none() {
            self.permit = self.sender.clone().try_reserve_owned().ok();
        }
        result
    }

    pub fn try_recv(&mut self) -> Option<T> {
        let result = self.receiver.try_recv().ok()?;
        if self.permit.is_none() {
            self.permit = self.sender.clone().try_reserve_owned().ok();
        }
        Some(result)
    }

    pub fn try_send(&mut self, value: T) -> bool {
        if let Some(permit) = self.permit.take() {
            permit.send(value);
            self.permit = self.sender.clone().try_reserve_owned().ok();
            true
        } else {
            let result = self.sender.try_send(value).is_ok();
            if result {
                self.permit = self.sender.clone().try_reserve_owned().ok();
            }
            result
        }
    }
}

I expected to see this happen: While the receiving thread holds this permit, send() called from another thread does not block unless the number of outstanding permits equals the number of free slots in the channel; and reserve_many() does not block unless it would lead the total number of outstanding permits to exceed the number of free slots. Messages arrive in the order sent, even when it's not the order in which the permits are reserved. For example, the following sequence:

[thread 1] let a = sender.reserve();
[thread 2] let b = sender.reserve_many(2);
[thread 3] let c = sender.reserve();
[thread 4] sender.send("E");
[thread 3] c.send("C");
[thread 1] a.send("A");
[thread 2] b.next().unwrap().send("B");
[thread 2] b.next().unwrap().send("D");

leads to the messages being received in the order "E", "C", "A", "B", "D" with none of the above calls blocking if the channel capacity is 5 or more, and "C", "A", "B", "D", "E" with only thread 4's send blocking if it is exactly 4.

When a slot previously assigned to a permit is filled this way without consuming the permit, the nth outstanding permit is reassigned to the nth free slot for all n.

Instead, this happened: The sending thread became blocked, leading to deadlock. (NB: The sending thread doesn't hold a copy of the PushbackReceiver; instead, it holds the Sender from which its self.sender was cloned.)

[Pr0methean]

A workaround seems to be to create a second channel with the same capacity as the main one, but this is probably inefficient because that second channel pays the cost of thread-safety unnecessarily, plus it doubles the memory utilization. It causes try_send() to fail more often and send() to block more often using the original sender, but also makes returns infallible. Fortunately my use cases involve 16..=32-byte messages, 128..=1024-message channel capacities, and mostly modern Intel CPUs with a 256KiB L2 cache per socket, so it's likely that the only non-negligible impact will be on the L1 cache hit rate.

use tokio::select;
use tokio::sync::mpsc::{OwnedPermit, Receiver, Sender, channel};

pub struct PushbackReceiver<T> {
    receiver: Receiver<T>,
    return_sender: Sender<T>,
    return_receiver: Receiver<T>,
}

impl<T> PushbackReceiver<T> {
    pub fn new(receiver: Receiver<T>, sender: &Sender<T>) -> Self {
        let (return_sender, return_receiver) = channel(sender.max_capacity());
        PushbackReceiver {
            receiver,
            return_sender,
            return_receiver,
        }
    }

    pub async fn recv(&mut self) -> (T, OwnedPermit<T>) {
        let return_permit = self.return_sender.clone().reserve_owned().await.unwrap();
        select! {
            biased;
            result = self.receiver.recv() => {
                (result.unwrap(), return_permit)
            },
            result = self.return_receiver.recv() => {
                (result.unwrap(), return_permit)
            }
        }
    }

    pub fn try_recv(&mut self) -> Option<(T, OwnedPermit<T>)> {
        match self.return_sender.clone().try_reserve_owned().ok() {
            Some(return_permit) => {
                if let Ok(received) = self.receiver.try_recv() {
                    Some((received, return_permit))
                } else if let Ok(received_return) = self.return_receiver.try_recv() {
                    Some((received_return, return_permit))
                } else {
                    None
                }
            }
            None => {
                let received = self.return_receiver.try_recv().expect(
                    "Failed to receive a returned item when no return permits are available",
                );
                let return_permit = self
                    .return_sender
                    .clone()
                    .try_reserve_owned()
                    .expect("Failed to obtain a return permit after receiving a returned item");
                Some((received, return_permit))
            }
        }
    }
}

[ADD-SP]

It would be better to have a reproducible example that leads to the deadlock by using the PushbackReceiver.

[Pr0methean]

I'll see what I can do about a minimal-reproducing example, but it will be difficult since the only program where I've encountered this issue is a remote client.

[Darksonn]

All calls to reserve or reserve_many() go into a queue and permits are given out in the order they appear in the queue. So if reserve_many(4) is at the front of the queue, then the next 4 permits are given to that reserve_many() call before any other reserve() calls can be given any permits. Does this explain what you are seeing?

[Pr0methean]

It may in part, but I believe I also observed send() calls without permits being blocked as though also queued that way, and the queue blocking while the oldest permit was waiting to be spent.