use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError, TrySendError}; cfg_time! { use crate::sync::mpsc::error::SendTimeoutError; use crate::time::Duration; } use std::fmt; #[cfg(any(feature = "signal", feature = "process", feature = "stream"))] use std::task::{Context, Poll}; /// Send values to the associated `Receiver`. /// /// Instances are created by the [`channel`](channel) function. pub struct Sender { chan: chan::Tx, } /// Permit to send one value into the channel. /// /// `Permit` values are returned by [`Sender::reserve()`] and are used to /// guarantee channel capacity before generating a message to send. /// /// [`Sender::reserve()`]: Sender::reserve pub struct Permit<'a, T> { chan: &'a chan::Tx, } /// Receive values from the associated `Sender`. /// /// Instances are created by the [`channel`](channel) function. pub struct Receiver { /// The channel receiver chan: chan::Rx, } /// Creates a bounded mpsc channel for communicating between asynchronous tasks /// with backpressure. /// /// The channel will buffer up to the provided number of messages. Once the /// buffer is full, attempts to `send` new messages will wait until a message is /// received from the channel. The provided buffer capacity must be at least 1. /// /// All data sent on `Sender` will become available on `Receiver` in the same /// order as it was sent. /// /// The `Sender` can be cloned to `send` to the same channel from multiple code /// locations. Only one `Receiver` is supported. /// /// If the `Receiver` is disconnected while trying to `send`, the `send` method /// will return a `SendError`. Similarly, if `Sender` is disconnected while /// trying to `recv`, the `recv` method will return a `RecvError`. /// /// # Panics /// /// Panics if the buffer capacity is 0. /// /// # Examples /// /// ```rust /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(_) = tx.send(i).await { /// println!("receiver dropped"); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// } /// } /// ``` pub fn channel(buffer: usize) -> (Sender, Receiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let semaphore = (semaphore::Semaphore::new(buffer), buffer); let (tx, rx) = chan::channel(semaphore); let tx = Sender::new(tx); let rx = Receiver::new(rx); (tx, rx) } /// Channel semaphore is a tuple of the semaphore implementation and a `usize` /// representing the channel bound. type Semaphore = (semaphore::Semaphore, usize); impl Receiver { pub(crate) fn new(chan: chan::Rx) -> Receiver { Receiver { chan } } /// Receives the next value for this receiver. /// /// `None` is returned when all `Sender` halves have dropped, indicating /// that no further values can be sent on the channel. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tokio::spawn(async move { /// tx.send("hello").await.unwrap(); /// }); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(None, rx.recv().await); /// } /// ``` /// /// Values are buffered: /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(100); /// /// tx.send("hello").await.unwrap(); /// tx.send("world").await.unwrap(); /// /// assert_eq!(Some("hello"), rx.recv().await); /// assert_eq!(Some("world"), rx.recv().await); /// } /// ``` pub async fn recv(&mut self) -> Option { use crate::future::poll_fn; poll_fn(|cx| self.chan.recv(cx)).await } #[cfg(any(feature = "signal", feature = "process"))] pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { self.chan.recv(cx) } /// Blocking receive to call outside of asynchronous contexts. /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::runtime::Runtime; /// use tokio::sync::mpsc; /// /// fn main() { /// let (tx, mut rx) = mpsc::channel::(10); /// /// let sync_code = thread::spawn(move || { /// assert_eq!(Some(10), rx.blocking_recv()); /// }); /// /// Runtime::new() /// .unwrap() /// .block_on(async move { /// let _ = tx.send(10).await; /// }); /// sync_code.join().unwrap() /// } /// ``` #[cfg(feature = "sync")] pub fn blocking_recv(&mut self) -> Option { crate::future::block_on(self.recv()) } /// Attempts to return a pending value on this receiver without blocking. /// /// This method will never block the caller in order to wait for data to /// become available. Instead, this will always return immediately with /// a possible option of pending data on the channel. /// /// This is useful for a flavor of "optimistic check" before deciding to /// block on a receiver. /// /// Compared with recv, this function has two failure cases instead of /// one (one for disconnection, one for an empty buffer). pub fn try_recv(&mut self) -> Result { self.chan.try_recv() } /// Closes the receiving half of a channel, without dropping it. /// /// This prevents any further messages from being sent on the channel while /// still enabling the receiver to drain messages that are buffered. Any /// outstanding [`Permit`] values will still be able to send messages. /// /// In order to guarantee no messages are dropped, after calling `close()`, /// `recv()` must be called until `None` is returned. /// /// [`Permit`]: Permit /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(20); /// /// tokio::spawn(async move { /// let mut i = 0; /// while let Ok(permit) = tx.reserve().await { /// permit.send(i); /// i += 1; /// } /// }); /// /// rx.close(); /// /// while let Some(msg) = rx.recv().await { /// println!("got {}", msg); /// } /// /// // Channel closed and no messages are lost. /// } /// ``` pub fn close(&mut self) { self.chan.close(); } } impl fmt::Debug for Receiver { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Receiver") .field("chan", &self.chan) .finish() } } impl Unpin for Receiver {} cfg_stream! { impl crate::stream::Stream for Receiver { type Item = T; fn poll_next(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.chan.recv(cx) } } } impl Sender { pub(crate) fn new(chan: chan::Tx) -> Sender { Sender { chan } } /// Sends a value, waiting until there is capacity. /// /// A successful send occurs when it is determined that the other end of the /// channel has not hung up already. An unsuccessful send would be one where /// the corresponding receiver has already been closed. Note that a return /// value of `Err` means that the data will never be received, but a return /// value of `Ok` does not mean that the data will be received. It is /// possible for the corresponding receiver to hang up immediately after /// this function returns `Ok`. /// /// # Errors /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] handle dropping, the function returns /// an error. The error includes the value passed to `send`. /// /// [`close`]: Receiver::close /// [`Receiver`]: Receiver /// /// # Examples /// /// In the following example, each call to `send` will block until the /// previously sent value was received. /// /// ```rust /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(_) = tx.send(i).await { /// println!("receiver dropped"); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// } /// } /// ``` pub async fn send(&self, value: T) -> Result<(), SendError> { match self.reserve().await { Ok(permit) => { permit.send(value); Ok(()) } Err(_) => Err(SendError(value)), } } /// Completes when the receiver has dropped. /// /// This allows the producers to get notified when interest in the produced /// values is canceled and immediately stop doing work. /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx1, rx) = mpsc::channel::<()>(1); /// let tx2 = tx1.clone(); /// let tx3 = tx1.clone(); /// let tx4 = tx1.clone(); /// let tx5 = tx1.clone(); /// tokio::spawn(async move { /// drop(rx); /// }); /// /// futures::join!( /// tx1.closed(), /// tx2.closed(), /// tx3.closed(), /// tx4.closed(), /// tx5.closed() /// ); //// println!("Receiver dropped"); /// } /// ``` pub async fn closed(&self) { self.chan.closed().await } /// Attempts to immediately send a message on this `Sender` /// /// This method differs from [`send`] by returning immediately if the channel's /// buffer is full or no receiver is waiting to acquire some data. Compared /// with [`send`], this function has two failure cases instead of one (one for /// disconnection, one for a full buffer). /// /// # Errors /// /// If the channel capacity has been reached, i.e., the channel has `n` /// buffered values where `n` is the argument passed to [`channel`], then an /// error is returned. /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] handle dropping, the function returns /// an error. The error includes the value passed to `send`. /// /// [`send`]: Sender::send /// [`channel`]: channel /// [`close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// // Create a channel with buffer size 1 /// let (tx1, mut rx) = mpsc::channel(1); /// let tx2 = tx1.clone(); /// /// tokio::spawn(async move { /// tx1.send(1).await.unwrap(); /// tx1.send(2).await.unwrap(); /// // task waits until the receiver receives a value. /// }); /// /// tokio::spawn(async move { /// // This will return an error and send /// // no message if the buffer is full /// let _ = tx2.try_send(3); /// }); /// /// let mut msg; /// msg = rx.recv().await.unwrap(); /// println!("message {} received", msg); /// /// msg = rx.recv().await.unwrap(); /// println!("message {} received", msg); /// /// // Third message may have never been sent /// match rx.recv().await { /// Some(msg) => println!("message {} received", msg), /// None => println!("the third message was never sent"), /// } /// } /// ``` pub fn try_send(&self, message: T) -> Result<(), TrySendError> { match self.chan.semaphore().0.try_acquire(1) { Ok(_) => {} Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(message)), Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(message)), } // Send the message self.chan.send(message); Ok(()) } /// Sends a value, waiting until there is capacity, but only for a limited time. /// /// Shares the same success and error conditions as [`send`], adding one more /// condition for an unsuccessful send, which is when the provided timeout has /// elapsed, and there is no capacity available. /// /// [`send`]: Sender::send /// /// # Errors /// /// If the receive half of the channel is closed, either due to [`close`] /// being called or the [`Receiver`] having been dropped, /// the function returns an error. The error includes the value passed to `send`. /// /// [`close`]: Receiver::close /// [`Receiver`]: Receiver /// /// # Examples /// /// In the following example, each call to `send_timeout` will block until the /// previously sent value was received, unless the timeout has elapsed. /// /// ```rust /// use tokio::sync::mpsc; /// use tokio::time::{sleep, Duration}; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// tokio::spawn(async move { /// for i in 0..10 { /// if let Err(e) = tx.send_timeout(i, Duration::from_millis(100)).await { /// println!("send error: #{:?}", e); /// return; /// } /// } /// }); /// /// while let Some(i) = rx.recv().await { /// println!("got = {}", i); /// sleep(Duration::from_millis(200)).await; /// } /// } /// ``` #[cfg(feature = "time")] #[cfg_attr(docsrs, doc(cfg(feature = "time")))] pub async fn send_timeout( &self, value: T, timeout: Duration, ) -> Result<(), SendTimeoutError> { let permit = match crate::time::timeout(timeout, self.reserve()).await { Err(_) => { return Err(SendTimeoutError::Timeout(value)); } Ok(Err(_)) => { return Err(SendTimeoutError::Closed(value)); } Ok(Ok(permit)) => permit, }; permit.send(value); Ok(()) } /// Blocking send to call outside of asynchronous contexts. /// /// # Panics /// /// This function panics if called within an asynchronous execution /// context. /// /// # Examples /// /// ``` /// use std::thread; /// use tokio::runtime::Runtime; /// use tokio::sync::mpsc; /// /// fn main() { /// let (tx, mut rx) = mpsc::channel::(1); /// /// let sync_code = thread::spawn(move || { /// tx.blocking_send(10).unwrap(); /// }); /// /// Runtime::new().unwrap().block_on(async move { /// assert_eq!(Some(10), rx.recv().await); /// }); /// sync_code.join().unwrap() /// } /// ``` #[cfg(feature = "sync")] pub fn blocking_send(&self, value: T) -> Result<(), SendError> { crate::future::block_on(self.send(value)) } /// Checks if the channel has been closed. This happens when the /// [`Receiver`] is dropped, or when the [`Receiver::close`] method is /// called. /// /// [`Receiver`]: crate::sync::mpsc::Receiver /// [`Receiver::close`]: crate::sync::mpsc::Receiver::close /// /// ``` /// let (tx, rx) = tokio::sync::mpsc::channel::<()>(42); /// assert!(!tx.is_closed()); /// /// let tx2 = tx.clone(); /// assert!(!tx2.is_closed()); /// /// drop(rx); /// assert!(tx.is_closed()); /// assert!(tx2.is_closed()); /// ``` pub fn is_closed(&self) -> bool { self.chan.is_closed() } /// Wait for channel capacity. Once capacity to send one message is /// available, it is reserved for the caller. /// /// If the channel is full, the function waits for the number of unreceived /// messages to become less than the channel capacity. Capacity to send one /// message is reserved for the caller. A [`Permit`] is returned to track /// the reserved capacity. The [`send`] function on [`Permit`] consumes the /// reserved capacity. /// /// Dropping [`Permit`] without sending a message releases the capacity back /// to the channel. /// /// [`Permit`]: Permit /// [`send`]: Permit::send /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.reserve().await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Sending on the permit succeeds /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// } /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { match self.chan.semaphore().0.acquire(1).await { Ok(_) => {} Err(_) => return Err(SendError(())), } Ok(Permit { chan: &self.chan }) } } impl Clone for Sender { fn clone(&self) -> Self { Sender { chan: self.chan.clone(), } } } impl fmt::Debug for Sender { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Sender") .field("chan", &self.chan) .finish() } } // ===== impl Permit ===== impl Permit<'_, T> { /// Sends a value using the reserved capacity. /// /// Capacity for the message has already been reserved. The message is sent /// to the receiver and the permit is consumed. The operation will succeed /// even if the receiver half has been closed. See [`Receiver::close`] for /// more details on performing a clean shutdown. /// /// [`Receiver::close`]: Receiver::close /// /// # Examples /// /// ``` /// use tokio::sync::mpsc; /// /// #[tokio::main] /// async fn main() { /// let (tx, mut rx) = mpsc::channel(1); /// /// // Reserve capacity /// let permit = tx.reserve().await.unwrap(); /// /// // Trying to send directly on the `tx` will fail due to no /// // available capacity. /// assert!(tx.try_send(123).is_err()); /// /// // Send a message on the permit /// permit.send(456); /// /// // The value sent on the permit is received /// assert_eq!(rx.recv().await.unwrap(), 456); /// } /// ``` pub fn send(self, value: T) { use std::mem; self.chan.send(value); // Avoid the drop logic mem::forget(self); } } impl Drop for Permit<'_, T> { fn drop(&mut self) { use chan::Semaphore; let semaphore = self.chan.semaphore(); // Add the permit back to the semaphore semaphore.add_permit(); if semaphore.is_closed() && semaphore.is_idle() { self.chan.wake_rx(); } } } impl fmt::Debug for Permit<'_, T> { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Permit") .field("chan", &self.chan) .finish() } }