//! A facility for an MPSC channel that counts the number of outstanding entries on the channel.

//

// (Tokio makes this possible by default, but we don't require tokio.  Crossbeam channels also allow

// this, but they aren't async, and they're MPMC. If a future version of the

// `futures` crate adds this functionality, we can use that instead. )

use std::{

    pin::Pin,

    sync::{

        Arc,

        atomic::{AtomicUsize, Ordering},

    },

    task::ready,

    task::{Context, Poll},

};

use futures::{Stream, sink::Sink, stream::FusedStream};

use pin_project::pin_project;

/// A wrapper around an arbitrary [`Sink`], to count the items inserted.

#[derive(Clone, Debug)]

#[pin_project]

pub struct CountingSink<S> {

    /// The inner sink whose items we're counting.

    #[pin]

    inner: S,

    /// A shared counter for items inserted into the channel

    ///

    /// We add 1 every time we enqueue an item.

    count: Arc<AtomicUsize>,

}

/// A wrapper around an arbitrary [`Stream`], to count the items inserted.

#[derive(Clone, Debug)]

#[pin_project]

pub struct CountingStream<S> {

    /// The inner stream whose items we're counting.

    #[pin]

    inner: S,

    /// A shared counter for items inserted into the channel.

    ///

    /// We remove 1 every time we dequeue an item.

    count: Arc<AtomicUsize>,

}

impl<T, S: Sink<T>> Sink<T> for CountingSink<S> {

    type Error = S::Error;

}

impl<S: Stream> Stream for CountingStream<S> {

    type Item = S::Item;

}

impl<S: FusedStream> FusedStream for CountingStream<S> {

}

impl<S> CountingStream<S> {

    /// Return an approximate count of the number of items currently on this channel.

    ///

    /// This count is necessarily approximate because the count can be changed by any of this

    /// channel's Senders or Receivers between when the caller

    /// gets the count and when the caller uses the count.

    /// Return a reference to the inner stream.

    ///

    /// If the stream has interior mutability, the caller must take care

    /// not to do anything with the stream that would invalidate the current counter.

    /// Return a mutable reference to the inner stream.

    ///

    /// If the stream has interior mutability, the caller must take care

    /// not to do anything with the stream that would invalidate the current counter.

}

impl<S> CountingSink<S> {

    /// Return an approximate count of the number of items currently on this channel.

    ///

    /// This count is necessarily approximate because the count can be changed by any of this

    /// channel's Senders or Receivers between when the caller

    /// gets the count and when the caller uses the count.

    /// Return a reference to the inner sink.

    ///

    /// If the sink has interior mutability, the caller must take care

    /// not to do anything with the sink that would invalidate the current counter.

    /// Return a mutable reference to the inner sink.

    ///

    /// If the sink has interior mutability, the caller must take care

    /// not to do anything with the sink that would invalidate the current counter.

}

/// Wrap a [`Sink`]/[`Stream`] pair into a [`CountingSink`] and [`CountingStream`] pair.

///

/// # Correctness

///

/// The sink and the stream should match and form a channel:

/// items sent on the sink should be received from the stream.

///

/// There should be no other handles in use for adding or removing items from the channel.

///

/// If these requirements aren't met, then the counts returned by the sink and stream

/// will not be accurate.

#[cfg(test)]

mod test {

    // @@ begin test lint list maintained by maint/add_warning @@

    #![allow(clippy::bool_assert_comparison)]

    #![allow(clippy::clone_on_copy)]

    #![allow(clippy::dbg_macro)]

    #![allow(clippy::mixed_attributes_style)]

    #![allow(clippy::print_stderr)]

    #![allow(clippy::print_stdout)]

    #![allow(clippy::single_char_pattern)]

    #![allow(clippy::unwrap_used)]

    #![allow(clippy::unchecked_time_subtraction)]

    #![allow(clippy::useless_vec)]

    #![allow(clippy::needless_pass_by_value)]

    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->

    use futures::{SinkExt as _, StreamExt as _};

    #[test]

    fn send_only_onetask() {

        tor_rtmock::MockRuntime::test_with_various(|_rt| async move {

            let (tx, rx) = futures::channel::mpsc::unbounded::<usize>();

            let (mut tx, rx) = super::channel(tx, rx);

            for n in 1..10 {

                tx.send(n).await.unwrap();

                assert_eq!(tx.approx_count(), n);

                assert_eq!(rx.approx_count(), n);

            }

        });

    }

    #[test]

    fn send_only_twotasks() {

        tor_rtmock::MockRuntime::test_with_various(|rt| async move {

            let (tx, rx) = futures::channel::mpsc::unbounded::<usize>();

            let (mut tx, rx) = super::channel(tx, rx);

            let mut tx2 = tx.clone();

            let j1 = rt.spawn_join("thread1", async move {

                for n in 1..=10 {

                    tx.send(n).await.unwrap();

                    assert!(tx.approx_count() >= n);

                }

            });

            let j2 = rt.spawn_join("thread2", async move {

                for n in 1..=10 {

                    tx2.send(n).await.unwrap();

                    assert!(tx2.approx_count() >= n);

                }

            });

            j1.await;

            j2.await;

            assert_eq!(rx.approx_count(), 20);

        });

    }

    #[test]

    fn send_and_receive() {

        tor_rtmock::MockRuntime::test_with_various(|rt| async move {

            let (tx, rx) = futures::channel::mpsc::unbounded::<usize>();

            let (mut tx, mut rx) = super::channel(tx, rx);

            const MAX: usize = 10000;

            let mut tx2 = tx.clone();

            let j1 = rt.spawn_join("thread1", async move {

                for n in 1..=MAX {

                    tx.send(n).await.unwrap();

                }

            });

            let j2 = rt.spawn_join("thread2", async move {

                for n in 1..=MAX {

                    tx2.send(n).await.unwrap();

                }

            });

            let j3 = rt.spawn_join("receiver", async move {

                let mut total = 0;

                while let Some(x) = rx.next().await {

                    total += x; // spot check

                    let count = rx.approx_count();

                    assert!(count <= MAX * 2);

                }

                assert_eq!(total, MAX * (MAX + 1)); // two senders, so no "/2".

                rx

            });

            j1.await;

            j2.await;

            let rx = j3.await;

            assert_eq!(rx.approx_count(), 0);

        });

    }

}