      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    300
    301
    302
    303
    304
    305
    306
    307
    308
    309
    310

    use crate::loom::sync::atomic::AtomicUsize;
    use crate::sync::mpsc::chan;
    use crate::sync::mpsc::error::SendError;

    use std::fmt;
    use std::task::{Context, Poll};

    /// Send values to the associated `UnboundedReceiver`.
    ///
    /// Instances are created by the
    /// [`unbounded_channel`](unbounded_channel) function.
    pub struct UnboundedSender<T> {
        chan: chan::Tx<T, Semaphore>,
    }

    impl<T> Clone for UnboundedSender<T> {
        fn clone(&self) -> Self {
            UnboundedSender {
                chan: self.chan.clone(),
            }
        }
    }

    impl<T> fmt::Debug for UnboundedSender<T> {
        fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
            fmt.debug_struct("UnboundedSender")
                .field("chan", &self.chan)
                .finish()
        }
    }

    /// Receive values from the associated `UnboundedSender`.
    ///
    /// Instances are created by the
    /// [`unbounded_channel`](unbounded_channel) function.
    ///
    /// This receiver can be turned into a `Stream` using [`UnboundedReceiverStream`].
    ///
    /// [`UnboundedReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.UnboundedReceiverStream.html
    pub struct UnboundedReceiver<T> {
        /// The channel receiver
        chan: chan::Rx<T, Semaphore>,
    }

    impl<T> fmt::Debug for UnboundedReceiver<T> {
        fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
            fmt.debug_struct("UnboundedReceiver")
                .field("chan", &self.chan)
                .finish()
        }
    }

    /// Creates an unbounded mpsc channel for communicating between asynchronous
    /// tasks without backpressure.
    ///
    /// A `send` on this channel will always succeed as long as the receive half has
    /// not been closed. If the receiver falls behind, messages will be arbitrarily
    /// buffered.
    ///
    /// **Note** that the amount of available system memory is an implicit bound to
    /// the channel. Using an `unbounded` channel has the ability of causing the
    /// process to run out of memory. In this case, the process will be aborted.
    pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
        let (tx, rx) = chan::channel(AtomicUsize::new(0));

        let tx = UnboundedSender::new(tx);
        let rx = UnboundedReceiver::new(rx);

        (tx, rx)
    }

    /// No capacity
    type Semaphore = AtomicUsize;

    impl<T> UnboundedReceiver<T> {
        pub(crate) fn new(chan: chan::Rx<T, Semaphore>) -> UnboundedReceiver<T> {
            UnboundedReceiver { 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::unbounded_channel();
        ///
        ///     tokio::spawn(async move {
        ///         tx.send("hello").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::unbounded_channel();
        ///
        ///     tx.send("hello").unwrap();
        ///     tx.send("world").unwrap();
        ///
        ///     assert_eq!(Some("hello"), rx.recv().await);
        ///     assert_eq!(Some("world"), rx.recv().await);
        /// }
        /// ```
        pub async fn recv(&mut self) -> Option<T> {
            use crate::future::poll_fn;

            poll_fn(|cx| self.poll_recv(cx)).await
        }

        /// 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::sync::mpsc;
        ///
        /// #[tokio::main]
        /// async fn main() {
        ///     let (tx, mut rx) = mpsc::unbounded_channel::<u8>();
        ///
        ///     let sync_code = thread::spawn(move || {
        ///         assert_eq!(Some(10), rx.blocking_recv());
        ///     });
        ///
        ///     let _ = tx.send(10);
        ///     sync_code.join().unwrap();
        /// }
        /// ```
        #[cfg(feature = "sync")]
        pub fn blocking_recv(&mut self) -> Option<T> {
            crate::future::block_on(self.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.
        pub fn close(&mut self) {
            self.chan.close();
        }

        /// Polls to receive the next message on this channel.
        ///
        /// This method returns:
        ///
        ///  * `Poll::Pending` if no messages are available but the channel is not
        ///    closed.
        ///  * `Poll::Ready(Some(message))` if a message is available.
        ///  * `Poll::Ready(None)` if the channel has been closed and all messages
        ///    sent before it was closed have been received.
        ///
        /// When the method returns `Poll::Pending`, the `Waker` in the provided
        /// `Context` is scheduled to receive a wakeup when a message is sent on any
        /// receiver, or when the channel is closed.  Note that on multiple calls to
        /// `poll_recv`, only the `Waker` from the `Context` passed to the most
        /// recent call is scheduled to receive a wakeup.
        pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
            self.chan.recv(cx)
        }
    }

    impl<T> UnboundedSender<T> {
        pub(crate) fn new(chan: chan::Tx<T, Semaphore>) -> UnboundedSender<T> {
            UnboundedSender { chan }
        }

        /// Attempts to send a message on this `UnboundedSender` without blocking.
        ///
        /// This method is not marked async because sending a message to an unbounded channel
        /// never requires any form of waiting. Because of this, the `send` method can be
        /// used in both synchronous and asynchronous code without problems.
        ///
        /// If the receive half of the channel is closed, either due to [`close`]
        /// being called or the [`UnboundedReceiver`] having been dropped, this
        /// function returns an error. The error includes the value passed to `send`.
        ///
        /// [`close`]: UnboundedReceiver::close
        /// [`UnboundedReceiver`]: UnboundedReceiver
        pub fn send(&self, message: T) -> Result<(), SendError<T>> {
            if !self.inc_num_messages() {
                return Err(SendError(message));
            }

            self.chan.send(message);
            Ok(())
        }

        fn inc_num_messages(&self) -> bool {
            use std::process;
            use std::sync::atomic::Ordering::{AcqRel, Acquire};

            let mut curr = self.chan.semaphore().load(Acquire);

            loop {
                if curr & 1 == 1 {
                    return false;
                }

                if curr == usize::MAX ^ 1 {
                    // Overflowed the ref count. There is no safe way to recover, so
                    // abort the process. In practice, this should never happen.
                    process::abort()
                }

                match self
                    .chan
                    .semaphore()
                    .compare_exchange(curr, curr + 2, AcqRel, Acquire)
                {
                    Ok(_) => return true,
                    Err(actual) => {
                        curr = actual;
                    }
                }
            }
        }

        /// 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::unbounded_channel::<()>();
        ///     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
        }
        /// Checks if the channel has been closed. This happens when the
        /// [`UnboundedReceiver`] is dropped, or when the
        /// [`UnboundedReceiver::close`] method is called.
        ///
        /// [`UnboundedReceiver`]: crate::sync::mpsc::UnboundedReceiver
        /// [`UnboundedReceiver::close`]: crate::sync::mpsc::UnboundedReceiver::close
        ///
        /// ```
        /// let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>();
        /// 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()
        }

        /// Returns `true` if senders belong to the same channel.
        ///
        /// # Examples
        ///
        /// ```
        /// let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>();
        /// let  tx2 = tx.clone();
        /// assert!(tx.same_channel(&tx2));
        ///
        /// let (tx3, rx3) = tokio::sync::mpsc::unbounded_channel::<()>();
        /// assert!(!tx3.same_channel(&tx2));
        /// ```
        pub fn same_channel(&self, other: &Self) -> bool {
            self.chan.same_channel(&other.chan)
        }
    }