//! A oneshot broadcast channel.

//!

//! The motivation for this channel type was to allow multiple

//! receivers to either wait for something to finish,

//! or to have an inexpensive method of checking if it has finished.

//!

//! See [`channel()`].

use std::future::{Future, IntoFuture};

use std::ops::Drop;

use std::pin::Pin;

use std::sync::{Arc, Mutex, OnceLock, Weak};

use std::task::{Context, Poll, Waker, ready};

use slotmap_careful::DenseSlotMap;

slotmap_careful::new_key_type! { struct WakerKey; }

/// A [oneshot broadcast][crate::oneshot_broadcast] sender.

#[derive(Debug)]

pub struct Sender<T> {

    /// State shared with all [`Receiver`]s.

    shared: Weak<Shared<T>>,

}

/// A [oneshot broadcast][crate::oneshot_broadcast] receiver.

///

/// The `Receiver` offers two methods for receiving the message:

///

/// 1. [`Receiver::into_future`]

///     ```rust

///     # use tor_async_utils::oneshot_broadcast::{channel, SenderDropped};

///     # async fn x() -> Result<(), SenderDropped> {

///     let (tx, rx) = channel();

///     tx.send(0);

///     let message: u32 = rx.await.unwrap();

///     # Ok(())

///     # }

///     ```

///

/// 2. [`Receiver::borrowed`]

///     ```rust

///     # use tor_async_utils::oneshot_broadcast::{channel, SenderDropped};

///     # async fn x() -> Result<(), SenderDropped> {

///     let (tx, rx) = channel();

///     tx.send(0);

///     let message: &u32 = rx.borrowed().await.unwrap();

///     # Ok(())

///     # }

///     ```

#[derive(Clone, Debug)]

pub struct Receiver<T> {

    /// State shared with the sender and all other receivers.

    shared: Arc<Shared<T>>,

}

/// State shared between the sender and receivers.

/// Correctness:

///

/// Sending a message:

///  - set the message OnceLock (A)

///  - acquire the wakers Mutex

///  - take all wakers (B)

///  - release the wakers Mutex (C)

///  - wake all wakers

///

/// Polling:

///  - if message was set, return it (fast path)

///  - acquire the wakers Mutex (D)

///  - if message was set, return it (E)

///  - add waker (F)

///  - release the wakers Mutex

///

/// When the wakers Mutex is released at (C), a release-store operation is performed by the Mutex,

/// which means that the message set at (A) will be seen by all future acquire-load operations by

/// that same Mutex. More specifically, after (C) has occurred and when the same mutex is acquired at

/// (D), the message set at (A) is guaranteed to be visible at (E). This means that after the wakers

/// are taken at (B), no future wakers will be added at (F) and no waker will be "lost".

#[derive(Debug)]

struct Shared<T> {

    /// The message sent from the [`Sender`] to the [`Receiver`]s.

    msg: OnceLock<Result<T, SenderDropped>>,

    /// The wakers waiting for a value to be sent.

    /// Will be set to `Err` after the wakers have been woken.

    // the `Result` isn't technically needed here,

    // but we use it to help detect bugs;

    // see `WakersAlreadyWoken` for details

    wakers: Mutex<Result<DenseSlotMap<WakerKey, Waker>, WakersAlreadyWoken>>,

}

/// The future from [`Receiver::borrowed`].

///

/// Will be ready, yielding `&'a T`,

/// when the sender sends a message or is dropped.

#[derive(Debug)]

pub struct BorrowedReceiverFuture<'a, T> {

    /// State shared with the sender and all other receivers.

    shared: &'a Shared<T>,

    /// The key for any waker that we've added to [`Shared::wakers`].

    waker_key: Option<WakerKey>,

}

/// The future from [`Receiver::into_future`].

///

/// Will be ready, yielding a clone of `T`,

/// when the sender sends a message or is dropped.

// Both `ReceiverFuture` and `BorrowedReceiverFuture` have similar fields

// but there's no nice way to deduplicated them.

// It would have been nice if we could store a `BorrowedReceiverFuture`

// holding a reference to our `Arc<Shared>`,

// but that would be a self-referential struct,

// so we need to duplicate the fields here instead.

#[derive(Debug)]

pub struct ReceiverFuture<T> {

    /// State shared with the sender and all other receivers.

    shared: Arc<Shared<T>>,

    /// The key for any waker that we've added to [`Shared::wakers`].

    waker_key: Option<WakerKey>,

}

/// The wakers have already been woken.

///

/// This is used to help detect if we're trying to access the wakers after they've already been

/// woken, which likely indicates a bug. For example, it is a bug if a receiver attempts to add a

/// waker after the sender has already sent its message and woken the wakers, since the new waker

/// would never be woken.

#[derive(Copy, Clone, Debug)]

struct WakersAlreadyWoken;

/// The message has already been set, and we can't set it again.

#[derive(Copy, Clone, Debug, thiserror::Error)]

#[error("the message was already set")]

struct MessageAlreadySet;

/// The sender was dropped, so the channel is closed.

#[derive(Copy, Clone, Debug, PartialEq, Eq, thiserror::Error)]

#[error("the sender was dropped")]

#[allow(clippy::exhaustive_structs)]

pub struct SenderDropped;

/// All the receivers were dropped, so the channel is closed.

#[derive(Copy, Clone, Debug, PartialEq, Eq, thiserror::Error)]

#[error("all the receivers were dropped")]

#[allow(clippy::exhaustive_structs)]

pub struct AllReceiversDropped;

/// Create a new oneshot broadcast channel.

///

/// ```rust

/// # use tor_async_utils::oneshot_broadcast::{channel, SenderDropped};

/// # async fn x() -> Result<(), SenderDropped> {

/// let (tx, rx) = channel();

/// let rx_clone = rx.clone();

/// tx.send(0_u8);

/// assert_eq!(rx.await, Ok(0));

/// assert_eq!(rx_clone.await, Ok(0));

/// # Ok(())

/// # }

/// ```

impl<T> Sender<T> {

    /// Send the message to the [`Receiver`]s.

    ///

    /// The message may be lost if all receivers have been dropped.

        // set the message and inform the wakers

            // this 'send()` method takes an owned self,

            // and we don't send a message outside of here and the drop handler,

            // so this shouldn't be possible

    /// Send the message, and wake and clear all wakers.

    ///

    /// If all receivers have been dropped, then always returns `Ok`.

    ///

    /// If the message was unable to be set, returns `Err(MessageAlreadySet)`.

        // Even if the `Weak` upgrade is successful,

        // it's possible that the last receiver

        // will be dropped during this `send_and_wake` method,

        // in which case we will be holding the last `Arc`.

            // all receivers have dropped; nothing to do

        };

        // set the message

            // Take the wakers and drop the mutex guard, releasing the lock.

            //

            // We could just drain the wakers map in-place here, but instead we replace the map with

            // an explicit `WakersAlreadyWoken` state to help catch bugs if something tries adding a

            // new waker later after we've already woken the wakers.

            //

            // The above `msg.set()` will only ever succeed once,

            // which means that we should only end up here once.

        };

        // Once we drop the mutex guard, which does a release-store on its own atomic, any other

        // code which later acquires the wakers mutex is guaranteed to see the msg as "set".

        // See comments on `Shared`.

        // Wake while not holding the lock.

        // Since the lock is used in `ReceiverFuture::poll` and `ReceiverFuture::drop` and

        // should not block for long periods of time,

        // we'd prefer not to run third-party waker code here while holding the mutex,

        // even if `wake` should typically be fast.

    /// Returns `true` if all [`Receiver`]s (and all futures created from the receivers) have been

    /// dropped.

    ///

    /// This can be useful to skip doing extra work to generate the message if the message will be

    /// discarded anyways.

    // This is for external use.

    // It is not always valid to call this internally.

    // For example when we've done a `Weak::upgrade` internally, like in `send_and_wake`,

    // this won't return the correct value.

    /// Subscribe to this channel, creating a new receiver.

    ///

    /// Returns an error if all the existing [`Receiver`]s (and all futures created

    /// from the receivers) have already been dropped.

        Ok(Receiver {

        })

}

impl<T> Drop for Sender<T> {

        // set an error message to indicate that the sender was dropped and inform the wakers;

        // it's fine if setting the message fails since it might have been set previously during a

        // `send()`

}

impl<T> Receiver<T> {

    /// Receive a borrowed message from the [`Sender`].

    ///

    /// This may be more efficient than [`Receiver::into_future`]

    /// and doesn't require `T: Clone`.

    ///

    /// This is cancellation-safe.

    /// The receiver is ready.

    ///

    /// If `true`, the [`Sender`] has either sent its message or been dropped.

}

impl<T: Clone> IntoFuture for Receiver<T> {

    type Output = Result<T, SenderDropped>;

    type IntoFuture = ReceiverFuture<T>;

    /// This future is cancellation-safe.

}

impl<'a, T> Future for BorrowedReceiverFuture<'a, T> {

    type Output = Result<&'a T, SenderDropped>;

}

impl<T> Drop for BorrowedReceiverFuture<'_, T> {

}

impl<T: Clone> Future for ReceiverFuture<T> {

    type Output = Result<T, SenderDropped>;

}

impl<T> Drop for ReceiverFuture<T> {

}

/// The shared poll implementation for receiver futures.

    // if the message was already set, return it

    // check again now that we've acquired the mutex

    // we have acquired the wakers mutex and checked that the message wasn't set,

    // so we know that wakers have not yet been woken

    // and it's okay to add our waker to the wakers map

        // we have added a waker previously

        // we have never added a waker

    }

/// The shared drop implementation for receiver futures.

            // this is the only place that removes the waker from the map,

            // so the waker should never be missing

#[cfg(test)]

mod test {

    #![allow(clippy::unwrap_used)]

    use super::*;

    use futures::future::FutureExt;

    use tor_rtcompat::SpawnExt;

    impl<T> Shared<T> {

        /// Count the number of wakers.

        fn count_wakers(&self) -> usize {

            self.wakers

                .lock()

                .expect("poisoned")

                .as_ref()

                .map(|x| x.len())

                .unwrap_or(0)

        }

    }

    #[test]

    fn standard_usage() {

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

            let (tx, rx) = channel();

            tx.send(0_u8);

            assert_eq!(rx.borrowed().await, Ok(&0));

            let (tx, rx) = channel();

            tx.send(0_u8);

            assert_eq!(rx.await, Ok(0));

        });

    }

    #[test]

    fn immediate_drop() {

        let _ = channel::<()>();

        let (tx, rx) = channel::<()>();

        drop(tx);

        drop(rx);

        let (tx, rx) = channel::<()>();

        drop(rx);

        drop(tx);

    }

    #[test]

    fn drop_sender() {

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

            let (tx, rx_1) = channel::<u8>();

            let rx_2 = rx_1.clone();

            drop(tx);

            let rx_3 = rx_1.clone();

            assert_eq!(rx_1.borrowed().await, Err(SenderDropped));

            assert_eq!(rx_2.borrowed().await, Err(SenderDropped));

            assert_eq!(rx_3.borrowed().await, Err(SenderDropped));

        });

    }

    #[test]

    fn clone_before_send() {

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

            let (tx, rx_1) = channel();

            let rx_2 = rx_1.clone();

            tx.send(0_u8);

            assert_eq!(rx_1.borrowed().await, Ok(&0));

            assert_eq!(rx_2.borrowed().await, Ok(&0));

        });

    }

    #[test]

    fn clone_after_send() {

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

            let (tx, rx_1) = channel();

            tx.send(0_u8);

            let rx_2 = rx_1.clone();

            assert_eq!(rx_1.borrowed().await, Ok(&0));

            assert_eq!(rx_2.borrowed().await, Ok(&0));

        });

    }

    #[test]

    fn clone_after_borrowed() {

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

            let (tx, rx_1) = channel();

            tx.send(0_u8);

            assert_eq!(rx_1.borrowed().await, Ok(&0));

            let rx_2 = rx_1.clone();

            assert_eq!(rx_2.borrowed().await, Ok(&0));

        });

    }

    #[test]

    fn drop_one_receiver() {

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

            let (tx, rx_1) = channel();

            let rx_2 = rx_1.clone();

            drop(rx_1);

            tx.send(0_u8);

            assert_eq!(rx_2.borrowed().await, Ok(&0));

        });

    }

    #[test]

    fn drop_all_receivers() {

        let (tx, rx_1) = channel();

        let rx_2 = rx_1.clone();

        drop(rx_1);

        drop(rx_2);

        tx.send(0_u8);

    }

    #[test]

    fn drop_fut() {

        let (_tx, rx) = channel::<u8>();

        let fut = rx.borrowed();

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after sending

        let (tx, rx) = channel();

        tx.send(0_u8);

        let fut = rx.borrowed();

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after polling once

        let (_tx, rx) = channel::<u8>();

        let mut fut = Box::pin(rx.borrowed());

        assert_eq!(rx.shared.count_wakers(), 0);

        assert_eq!(fut.as_mut().now_or_never(), None);

        assert_eq!(rx.shared.count_wakers(), 1);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after polling once and send

        let (tx, rx) = channel();

        let mut fut = Box::pin(rx.borrowed());

        assert_eq!(rx.shared.count_wakers(), 0);

        assert_eq!(fut.as_mut().now_or_never(), None);

        assert_eq!(rx.shared.count_wakers(), 1);

        tx.send(0_u8);

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

    }

    #[test]

    fn drop_owned_fut() {

        let (_tx, rx) = channel::<u8>();

        let fut = rx.clone().into_future();

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after sending

        let (tx, rx) = channel();

        tx.send(0_u8);

        let fut = rx.clone().into_future();

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after polling once

        let (_tx, rx) = channel::<u8>();

        let mut fut = Box::pin(rx.clone().into_future());

        assert_eq!(rx.shared.count_wakers(), 0);

        assert_eq!(fut.as_mut().now_or_never(), None);

        assert_eq!(rx.shared.count_wakers(), 1);

        drop(fut);

        assert_eq!(rx.shared.count_wakers(), 0);

        // drop after polling once and send

        let (tx, rx) = channel();

        let mut fut = Box::pin(rx.clone().into_future());

        assert_eq!(rx.shared.count_wakers(), 0);

        assert_eq!(fut.as_mut().now_or_never(), None);

        assert_eq!(rx.shared.count_wakers(), 1);

        tx.send(0_u8);

        assert_eq!(rx.shared.count_wakers(), 0);

        drop(fut);

    }

    #[test]

    fn is_ready_after_send() {

        let (tx, rx_1) = channel();

        assert!(!rx_1.is_ready());

        let rx_2 = rx_1.clone();

        assert!(!rx_2.is_ready());

        tx.send(0_u8);

        assert!(rx_1.is_ready());

        assert!(rx_2.is_ready());

        let rx_3 = rx_1.clone();

        assert!(rx_3.is_ready());

    }

    #[test]

    fn is_ready_after_drop() {

        let (tx, rx_1) = channel::<u8>();

        assert!(!rx_1.is_ready());

        let rx_2 = rx_1.clone();

        assert!(!rx_2.is_ready());

        drop(tx);

        assert!(rx_1.is_ready());

        assert!(rx_2.is_ready());

        let rx_3 = rx_1.clone();

        assert!(rx_3.is_ready());

    }

    #[test]

    fn is_cancelled() {

        let (tx, rx) = channel::<u8>();

        assert!(!tx.is_cancelled());

        drop(rx);

        assert!(tx.is_cancelled());

        let (tx, rx_1) = channel::<u8>();

        assert!(!tx.is_cancelled());

        let rx_2 = rx_1.clone();

        drop(rx_1);

        assert!(!tx.is_cancelled());

        drop(rx_2);

        assert!(tx.is_cancelled());

    }

    #[test]

    fn recv_in_task() {

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

            let (tx, rx) = channel();

            let join = rt

                .spawn_with_handle(async move {

                    assert_eq!(rx.borrowed().await, Ok(&0));

                    assert_eq!(rx.await, Ok(0));

                })

                .unwrap();

            tx.send(0_u8);

            join.await;

        });

    }

    #[test]

    fn recv_multiple_in_task() {

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

            let (tx, rx) = channel();

            let rx_1 = rx.clone();

            let rx_2 = rx.clone();

            let join_1 = rt

                .spawn_with_handle(async move {

                    assert_eq!(rx_1.borrowed().await, Ok(&0));

                })

                .unwrap();

            let join_2 = rt

                .spawn_with_handle(async move {

                    assert_eq!(rx_2.await, Ok(0));

                })

                .unwrap();

            tx.send(0_u8);

            join_1.await;

            join_2.await;

            assert_eq!(rx.borrowed().await, Ok(&0));

        });

    }

    #[test]

    fn recv_multiple_times() {

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

            let (tx, rx) = channel();

            let rx_subscribed = tx.subscribe().unwrap();

            tx.send(0_u8);

            assert_eq!(rx.borrowed().await, Ok(&0));

            assert_eq!(rx.borrowed().await, Ok(&0));

            assert_eq!(rx.clone().await, Ok(0));

            assert_eq!(rx.await, Ok(0));

            assert_eq!(rx_subscribed.await, Ok(0));

        });

    }

    #[test]

    fn stress() {

        // In general we don't have control over the runtime and where/when tasks are scheduled,

        // so we try as best as possible to send the message while simultaneously creating new

        // receivers and waiting on them.

        // It's possible this might be entirely ineffective since we don't enforce any specific

        // scheduler behaviour here,

        // but in the worst case it's still a test with multiple receivers on different tasks,

        // so is useful to have.

        //

        // The `test_with_various` helper uses `MockExecutor` with two different deterministic

        // scheduling policies.

        // At least at the time of writing,

        // when this test uses `MockExecutor` with its "queue" scheduling policy

        // the "send" occurs after 20 of the tasks have begun waiting.

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

            let (tx, rx) = channel();

            rt.spawn(async move {

                // this tries to delay the send a little bit

                // to give time for some of the receiver tasks to start

                for _ in 0..20 {

                    tor_rtcompat::task::yield_now().await;

                }

                tx.send(0_u8);

            })

            .unwrap();

            let mut joins = vec![];

            for _ in 0..100 {

                let rx_clone = rx.clone();

                let join = rt

                    .spawn_with_handle(async move { rx_clone.borrowed().await.cloned() })

                    .unwrap();

                joins.push(join);

                // allows the send task to make progress if single-threaded

                tor_rtcompat::task::yield_now().await;

            }

            for join in joins {

                assert!(matches!(join.await, Ok(0)));

            }

        });

    }

}