/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Futures with specialized dropping behaviour

use std::future::Future;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use futures::Stream;
use pin_project::pin_project;

/// A future with a value attached that will be dropped together when the main future is dropped.
#[pin_project]
pub struct DropTogether<F, D> {
    #[pin]
    fut: F,
    // we don't really care about this. Anything is droppable.
    _drop_with: D,
}

impl<F, D> DropTogether<F, D> {
    pub fn new(fut: F, t: D) -> Self {
        Self { fut, _drop_with: t }
    }
}

impl<F, D> Future for DropTogether<F, D>
where
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        this.fut.poll(cx)
    }
}

impl<S, D> Stream for DropTogether<S, D>
where
    S: Stream,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();
        this.fut.poll_next(cx)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use dupe::Dupe;
    use futures::future;
    use futures::stream::FuturesOrdered;
    use futures::StreamExt;

    use crate::drop::DropTogether;

    #[test]
    fn drops_together() {
        let ref_counted = Arc::new(());

        let fut = async {};
        let drop_together = DropTogether::new(fut, ref_counted.dupe());

        assert_eq!(Arc::strong_count(&ref_counted), 2);

        drop(drop_together);
        assert_eq!(Arc::strong_count(&ref_counted), 1);
    }

    #[test]
    fn pollable() {
        let fut = async { 1 };
        let drop_together = DropTogether::new(fut, ());

        assert_eq!(futures::executor::block_on(drop_together), 1);
    }

    #[test]
    fn streamable() {
        let fut = {
            let mut fut = FuturesOrdered::new();
            fut.push_back(future::ready(1));
            fut.push_back(future::ready(2));
            fut
        };
        let mut drop_together = DropTogether::new(fut, ());

        futures::executor::block_on(async move {
            assert_eq!(Some(1), drop_together.next().await);
            assert_eq!(Some(2), drop_together.next().await);
        });
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use allocative::Allocative;
use buck2_data::SharedTaskEnd;
use buck2_data::SharedTaskStart;
use buck2_events::dispatch::get_dispatcher;
use buck2_events::dispatch::Span;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use futures::future::Future;
use futures::future::Shared;
use futures::FutureExt;
use pin_project::pin_project;

/// A custom future type that wraps a Shared future that is possibly polled in different threads with
/// their own distinct dispatchers. The poll behaviour sends a SharedTask span to the thread
/// sharing the future.
#[pin_project]
#[derive(Allocative)]
pub struct SharedEventsFuture<F>
where
    F: Future,
    F::Output: Clone,
{
    #[allocative(skip)] // TODO(nga): implement.
    #[pin]
    inner: Shared<F>,
    /// Because poll is called repeatedly, we don't want to send duplicate start events.
    /// `span` is set on first-entry of a cross-thread poll and used on completion for
    /// sending a SpanEnd event.
    span: SharedSpanState,
    /// Used to compare against task_local dispatcher, to determine if future is being
    /// polled in the context of another dispatcher.
    owner_trace_id: TraceId,
}

#[derive(Allocative)]
enum SharedSpanState {
    Undefined,
    Started(Box<Span>),
    Ended,
}

impl<F> SharedEventsFuture<F>
where
    F: Future,
    F::Output: Clone,
{
    pub fn new(fut: F) -> SharedEventsFuture<F> {
        SharedEventsFuture {
            inner: fut.shared(),
            span: SharedSpanState::Undefined,
            owner_trace_id: get_dispatcher().trace_id().dupe(),
        }
    }

    pub fn inner(&self) -> &Shared<F> {
        &self.inner
    }
}

impl<F> Clone for SharedEventsFuture<F>
where
    F: Future,
    F::Output: Clone,
{
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            // span is None because it's used for tracking the state of a specific future.
            span: SharedSpanState::Undefined,
            owner_trace_id: self.owner_trace_id.dupe(),
        }
    }
}

impl<F> Future for SharedEventsFuture<F>
where
    F: Future,
    F::Output: Clone,
{
    type Output = F::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.as_mut().project();

        // Have not yet sent span start
        if matches!(this.span, SharedSpanState::Undefined) {
            let dispatcher = get_dispatcher();

            if *this.owner_trace_id != *dispatcher.trace_id() {
                // Send start event
                *this.span = SharedSpanState::Started(Box::new(Span::start(
                    dispatcher,
                    SharedTaskStart {
                        owner_trace_id: this.owner_trace_id.to_string(),
                    },
                )))
            } else {
                *this.span = SharedSpanState::Ended;
            }
        }

        match this.inner.poll(cx) {
            Poll::Ready(t) => {
                // Poll is in the context of another dispatcher
                match std::mem::replace(this.span, SharedSpanState::Ended) {
                    SharedSpanState::Started(span) => {
                        span.end(SharedTaskEnd {});
                    }
                    SharedSpanState::Undefined | SharedSpanState::Ended => {}
                }
                Poll::Ready(t)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

pub trait SharedEvents: Future {
    fn instrumented_shared(self) -> SharedEventsFuture<Self>
    where
        Self: Sized,
        Self::Output: Clone,
    {
        SharedEventsFuture::new(self)
    }
}

impl<T> SharedEvents for dyn futures::Future<Output = T> {}

impl<T> SharedEvents for Pin<Box<dyn futures::Future<Output = T> + Send>> {}

#[cfg(test)]
mod tests {
    use std::assert_matches::assert_matches;
    use std::sync::Arc;

    use buck2_events::create_source_sink_pair;
    use buck2_events::dispatch::span;
    use buck2_events::dispatch::with_dispatcher_async;
    use buck2_events::dispatch::EventDispatcher;
    use buck2_events::BuckEvent;
    use buck2_events::EventSource;
    use buck2_wrapper_common::invocation_id::TraceId;
    use futures::future::BoxFuture;
    use futures::FutureExt;
    use tokio::sync::Barrier;
    use tokio::sync::Mutex;

    use super::*;

    fn create_dispatcher() -> (impl EventSource, EventDispatcher, TraceId) {
        let (events, sink) = create_source_sink_pair();
        let trace_id = TraceId::new();
        let dispatcher = EventDispatcher::new(trace_id.dupe(), sink);
        (events, dispatcher, trace_id)
    }

    async fn next_event<E>(source: &mut E) -> BuckEvent
    where
        E: EventSource,
    {
        let event = tokio::time::timeout(std::time::Duration::from_secs(3), async {
            source.receive()
        })
        .await
        .unwrap();
        event.unwrap().unpack_buck().unwrap().clone()
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn poll_shared_future() {
        let mut handles = vec![];
        let n_thread = 2;
        let (mut events1, dispatcher1, _trace_id1) = create_dispatcher();
        let (mut events2, dispatcher2, _trace_id2) = create_dispatcher();

        // Spawned future that is created by task1, shared with task2 - each with their own
        // dispatchers.
        let fut: Arc<Mutex<Option<SharedEventsFuture<BoxFuture<'static, &str>>>>> =
            Arc::new(Mutex::new(None));

        let barrier = Arc::new(Barrier::new(n_thread));

        let fut1 = fut.dupe();
        let b1 = barrier.dupe();
        // task1 - creates shared future
        handles.push(tokio::spawn(async move {
            let shared_future = with_dispatcher_async(dispatcher1, async {
                span(
                    buck2_data::FakeStart {
                        caramba: "".to_owned(),
                    },
                    || ("Hello world!", buck2_data::FakeEnd {}),
                )
            })
            .boxed()
            .instrumented_shared();
            let mut fut = fut1.lock().await;
            *fut = Some(shared_future.clone());
            // Release mutex before waiting on barrier.
            drop(fut);
            // Raise barrier to let task1 and task2 poll shared future concurrently.
            b1.wait().await;
            shared_future.await;
        }));

        let fut2 = fut.dupe();
        let b2 = barrier.dupe();
        // task2 - shares shared future created by task1. Should receive SharedTask span events.
        handles.push(tokio::spawn(with_dispatcher_async(
            dispatcher2,
            async move {
                b2.wait().await;
                // Ensures that the shared future has been created.
                let mut guard = fut2.lock().await;
                let shared_future = guard.as_mut().unwrap();
                shared_future.await;
            },
        )));

        futures::future::join_all(handles).await;

        // The "owning" dispatcher should receive the fake start/end events
        let owned_start_event = next_event(&mut events1).await;
        assert_matches!(
            owned_start_event.data(),
            buck2_data::buck_event::Data::SpanStart(buck2_data::SpanStartEvent {
                data: Some(buck2_data::span_start_event::Data::Fake(_)),
            }),
            "expected FakeStart, got {:?}",
            &owned_start_event,
        );
        let owned_end_event = next_event(&mut events1).await;
        assert_matches!(
            owned_end_event.data(),
            buck2_data::buck_event::Data::SpanEnd(buck2_data::SpanEndEvent {
                data: Some(buck2_data::span_end_event::Data::Fake(_)),
                ..
            }),
            "expected FakeEnd, got {:?}",
            &owned_end_event,
        );

        let shared_start_event = next_event(&mut events2).await;
        assert_matches!(
            shared_start_event.data(),
            buck2_data::buck_event::Data::SpanStart(buck2_data::SpanStartEvent {
                data: Some(buck2_data::span_start_event::Data::SharedTask(_)),
            }),
            "expected SharedTaskStart, got {:?}",
            &shared_start_event,
        );

        let shared_end_event = next_event(&mut events2).await;
        assert_matches!(
            shared_end_event.data(),
            buck2_data::buck_event::Data::SpanEnd(buck2_data::SpanEndEvent {
                data: Some(buck2_data::span_end_event::Data::SharedTask(
                    SharedTaskEnd {}
                )),
                ..
            }),
            "expected SharedTaskEnd, got {:?}",
            shared_end_event
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![feature(assert_matches)]
#![feature(pin_deref_mut)]
// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

pub mod cancellable_future;
pub mod cancellation;
pub mod drop;
pub mod instrumented_shared;
mod maybe_future;
pub mod owning_future;
pub mod spawn;
pub mod spawner;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::future::Future;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use pin_project::pin_project;

/// Future that might or might not be there
#[pin_project(project = MaybeFutureProj)]
pub(crate) enum MaybeFuture<F> {
    Fut(#[pin] F),
    None,
}

impl<F: Future> MaybeFuture<F> {
    pub(crate) fn take(mut self: Pin<&mut Self>) {
        self.set(MaybeFuture::None);
    }

    pub(crate) fn set_fut(mut self: Pin<&mut Self>, f: F) {
        self.set(MaybeFuture::Fut(f));
    }
}

impl<F> Future for MaybeFuture<F>
where
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.project() {
            MaybeFutureProj::Fut(f) => f.poll(cx),
            MaybeFutureProj::None => {
                panic!("polled again after completion")
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//!
//! A future that can be canceled via an explicit `CancellationHandle`.
//! This future is intended to be spawned on tokio-runtime directly, and for its results to be
//! accessed via the joinhandle.
//! It is not intended to be polled directly.
//!

use std::future::Future;
use std::marker::PhantomPinned;
use std::mem;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use futures::future::BoxFuture;
use pin_project::pin_project;

use crate::maybe_future::MaybeFuture;

#[pin_project]
pub struct OwningFuture<T, D> {
    // NOTE: The order of these two fields is important. The `OwningFuture` holds a
    // reference to the `context`, and so absolutely must be dropped first.
    #[pin]
    fut: MaybeFuture<BoxFuture<'static, T>>, // not actually static but the lifetime of this struct
    data: D,

    _p: PhantomPinned,
}

impl<T, D> OwningFuture<T, D> {
    pub fn new<F>(data: D, f: F) -> Pin<Box<Self>>
    where
        F: for<'a> FnOnce(&'a D) -> BoxFuture<'a, T> + Send,
    {
        let this = OwningFuture {
            data,
            fut: MaybeFuture::None,
            _p: Default::default(),
        };

        let mut this = Box::pin(this);
        let fut = unsafe {
            // SAFETY: The `OwningFuture` is always immediately pinned upon
            // construction, so the self-reference will remain valid at least until this value
            // is dropped. Furthermore, the inner future is dropped before the `context` is dropped,
            // and so the reference will remain valid even during the drop of the `Future`
            mem::transmute::<BoxFuture<'_, T>, BoxFuture<'static, T>>(f(&this.data))
        };

        {
            let this = this.as_mut().project();
            this.fut.set_fut(fut);
        }

        this
    }
}

impl<T, D> Future for OwningFuture<T, D> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().fut.poll(cx)
    }
}

#[cfg(test)]
mod tests {
    use futures::future::BoxFuture;
    use futures::FutureExt;

    use crate::owning_future::OwningFuture;

    #[tokio::test]
    async fn test_future() {
        struct SomeData(usize);

        fn some_fn(d: &SomeData) -> BoxFuture<usize> {
            async move { d.0 }.boxed()
        }

        let data = SomeData(2);
        let fut = OwningFuture::new(data, some_fn);

        // we can spawn the future such that it owns the data
        let data = tokio::spawn(fut).await.unwrap();
        assert_eq!(data, 2);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! The future that is spawned, but has various more strict cancellation behaviour than
//! tokio's JoinHandle
//!

use std::any::Any;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use allocative::Allocative;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::FutureExt;
use pin_project::pin_project;
use pin_project::pinned_drop;
use thiserror::Error;
use tracing::Instrument;
use tracing::Span;

use crate::cancellable_future::CancellableFuture;
use crate::cancellable_future::StrongRefCount;
use crate::cancellable_future::WeakRefCount;
use crate::cancellation::future::make_cancellable_future;
use crate::cancellation::future::CancellationHandle;
use crate::cancellation::ExplicitCancellationContext;
use crate::instrumented_shared::SharedEvents;
use crate::instrumented_shared::SharedEventsFuture;
use crate::spawner::Spawner;

#[derive(Debug, Error, Copy, Clone, PartialEq)]
pub enum WeakFutureError {
    #[error("Join Error")]
    JoinError,

    #[error("Cancelled")]
    Cancelled,
}

/// A unit of computation within Dice. Futures to the result of this computation should be obtained
/// via this task struct
#[derive(Allocative)]
pub struct WeakJoinHandle<T: Clone> {
    #[allocative(skip)] // TODO(nga): `Shared` requires `Clone`.
    join_handle: SharedEventsFuture<BoxFuture<'static, T>>,
    #[allocative(skip)]
    guard: WeakRefCount,
}

impl<T: Send + Sync + Clone + 'static> WeakJoinHandle<T> {
    /// Return `None` if the task has been canceled.
    pub fn pollable(&self) -> Option<StrongJoinHandle<SharedEventsFuture<BoxFuture<'static, T>>>> {
        self.guard.upgrade().map(|inner| StrongJoinHandle {
            guard: inner,
            fut: self.join_handle.clone(),
        })
    }
}

impl<T: Send + Sync + Clone + 'static> WeakJoinHandle<Result<T, WeakFutureError>> {
    pub fn into_completion_observer(self) -> CompletionObserver<T> {
        CompletionObserver { inner: self }
    }
}

/// The actual pollable future that returns the result of the task. This keeps the future alive.
#[pin_project]
pub struct StrongJoinHandle<F> {
    guard: StrongRefCount,
    #[pin]
    fut: F,
}

impl<F> StrongJoinHandle<F> {
    fn map<F2>(self, map: impl FnOnce(F) -> F2) -> StrongJoinHandle<F2> {
        StrongJoinHandle {
            guard: self.guard,
            fut: map(self.fut),
        }
    }
}

impl<T> StrongJoinHandle<SharedEventsFuture<BoxFuture<'static, T>>>
where
    T: Clone,
{
    fn weak_handle(&self) -> WeakJoinHandle<T> {
        WeakJoinHandle {
            join_handle: self.fut.clone(),
            guard: self.guard.downgrade(),
        }
    }
}

impl<F: Future> StrongJoinHandle<F> {
    pub fn inner(&self) -> &F {
        &self.fut
    }
}

impl<F, T> Future for StrongJoinHandle<F>
where
    F: Future<Output = Result<T, WeakFutureError>>,
{
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // When we have a StrongJoinHandle, we expect the future to not have been cancelled.
        let this = self.project();
        this.fut.poll(cx).map(|r| r.unwrap())
    }
}

#[pin_project]
pub struct CompletionObserver<T: Clone> {
    inner: WeakJoinHandle<Result<T, WeakFutureError>>,
}

impl<T: Clone> Future for CompletionObserver<T> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        this.inner.join_handle.poll_unpin(cx).map(|_res| ())
    }
}

/// Spawn a cancellable future. The preamble is a non-cancellable portion that can come before.
pub fn spawn_dropcancel_with_preamble<T, S, P>(
    future: T,
    preamble: P,
    spawner: &dyn Spawner<S>,
    ctx: &S,
    span: Span,
) -> (
    WeakJoinHandle<Result<T::Output, WeakFutureError>>,
    StrongJoinHandle<SharedEventsFuture<BoxFuture<'static, Result<T::Output, WeakFutureError>>>>,
)
where
    T: Future + Send + 'static,
    T::Output: Any + Clone + Send + 'static,
    P: Future<Output = ()> + Send + 'static,
{
    let strong = spawn_inner(future, preamble, spawner, ctx, span).map(|f| f.instrumented_shared());
    (strong.weak_handle(), strong)
}

/// Spawn a cancellable future. The preamble is a non-cancellable portion that can come before.
fn spawn_inner<T, S, P>(
    future: T,
    preamble: P,
    spawner: &dyn Spawner<S>,
    ctx: &S,
    span: Span,
) -> StrongJoinHandle<BoxFuture<'static, Result<T::Output, WeakFutureError>>>
where
    T: Future + Send + 'static,
    T::Output: Any + Send + 'static,
    P: Future<Output = ()> + Send + 'static,
{
    // For Ready<()> and BoxFuture<()> futures we get these sizes:
    // future alone: 196/320 bits
    // future + no-op preamble via async block: 448/704 bits
    // future + no-op preamble via FuturesExt::then: 256/384 bits
    // future + no-op preamble + instrument: 512/640 bits

    // As the spawner is going to take a boxed future and erase its concrete type,
    // we can have different future types for different scenarios in order to
    // minimize the size of them.
    //
    // While we could feasibly distinguish the no-op preamble case, one extra pointer
    // is an okay cost for the simpler api (for now).
    let (future, guard) = CancellableFuture::new_refcounted(future);
    let future = future.map(|v| Box::new(v) as _);
    let future = preamble.then(|_| future);
    let future = if span.is_disabled() {
        future.boxed()
    } else {
        future.instrument(span).boxed()
    };

    let task = spawner.spawn(ctx, future.boxed());
    let task = task
        .map(|v| {
            v.map_err(|_e: tokio::task::JoinError| WeakFutureError::JoinError)?
                .downcast::<Option<T::Output>>()
                .expect("Spawned task returned the wrong type")
                .ok_or(WeakFutureError::Cancelled)
        })
        .boxed();

    StrongJoinHandle { guard, fut: task }
}

/// Spawn a cancellable future.
pub fn spawn_dropcancel<T, S>(
    future: T,
    spawner: &dyn Spawner<S>,
    ctx: &S,
    span: Span,
) -> StrongJoinHandle<BoxFuture<'static, Result<T::Output, WeakFutureError>>>
where
    T: Future + Send + 'static,
    T::Output: Any + Send + 'static,
{
    spawn_inner(future, futures::future::ready(()), spawner, ctx, span)
}

pub struct FutureAndCancellationHandle<T> {
    pub future: CancellableJoinHandle<T>,
    pub cancellation_handle: CancellationHandle,
}

impl<T> FutureAndCancellationHandle<T> {
    pub fn into_drop_cancel(self) -> DropCancelFuture<T> {
        self.future.into_drop_cancel(self.cancellation_handle)
    }
}

/// Spawn a future that's cancellable via an CancellationHandle. Dropping the future or the handle
/// does not cancel the future
pub fn spawn_cancellable<F, T, S>(
    f: F,
    spawner: &dyn Spawner<S>,
    ctx: &S,
) -> FutureAndCancellationHandle<T>
where
    for<'a> F: FnOnce(&'a ExplicitCancellationContext) -> BoxFuture<'a, T> + Send,
    T: Any + Send + 'static,
{
    let (future, cancellation_handle) = make_cancellable_future(f);

    // For Ready<()> and BoxFuture<()> futures we get these sizes:
    // future alone: 196/320 bits
    // future via async block: 448/704 bits

    // As the spawner is going to take a boxed future and erase its concrete type,
    // we can have different future types for different scenarios in order to
    // minimize the size of them.
    //
    // While we could feasibly distinguish the no-op preamble case, one extra pointer
    // is an okay cost for the simpler api (for now).
    let future = future.map(|v| Box::new(v) as _);

    let task = spawner.spawn(ctx, future.boxed());
    let task = task
        .map(|v| {
            v.map_err(|_e: tokio::task::JoinError| WeakFutureError::JoinError)?
                .downcast::<Option<T>>()
                .expect("Spawned task returned the wrong type")
                .ok_or(WeakFutureError::Cancelled)
        })
        .boxed();

    FutureAndCancellationHandle {
        future: CancellableJoinHandle(task),
        cancellation_handle,
    }
}

#[pin_project]
pub struct CancellableJoinHandle<T>(#[pin] BoxFuture<'static, Result<T, WeakFutureError>>);

impl<T> Future for CancellableJoinHandle<T> {
    type Output = Result<T, WeakFutureError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().0.poll(cx)
    }
}

#[pin_project(PinnedDrop)]
pub struct DropCancelFuture<T> {
    #[pin]
    fut: BoxFuture<'static, Result<T, WeakFutureError>>,
    cancellation_handle: Option<CancellationHandle>,
}

impl<T> Future for DropCancelFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().fut.poll(cx).map(|r| {
            // since the lifetime of this future is responsible for cancellation, this future can't
            // possibly have been canceled.
            // We can only have join errors, which is when the executor shuts down. To be consistent
            // with existing spawn behaviour, we can ignore that.
            r.unwrap()
        })
    }
}

#[pinned_drop]
impl<T> PinnedDrop for DropCancelFuture<T> {
    fn drop(mut self: Pin<&mut Self>) {
        // ignore the termination future of when we actually shutdown. The creator of this
        // DropCancelFuture has the termination future as well that it can use to observe termination
        // if it cares
        let _cancel = self
            .cancellation_handle
            .take()
            .expect("dropped twice")
            .cancel();
    }
}

impl<T> CancellableJoinHandle<T> {
    fn into_drop_cancel(self, cancellation_handle: CancellationHandle) -> DropCancelFuture<T> {
        DropCancelFuture {
            fut: self.0,
            cancellation_handle: Some(cancellation_handle),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use tokio::sync::oneshot;

    use super::*;
    use crate::spawner::TokioSpawner;

    #[derive(Default)]
    struct MockCtx;

    #[tokio::test]
    async fn test_cancellation() {
        let (release_task, recv_release_task) = oneshot::channel();
        let (notify_success, recv_success) = oneshot::channel();

        let sp = Arc::new(TokioSpawner);

        let (_task, poll) = spawn_dropcancel_with_preamble(
            async move {
                recv_release_task.await.unwrap();
                notify_success.send(()).unwrap();
            },
            futures::future::ready(()),
            sp.as_ref(),
            &MockCtx,
            tracing::debug_span!("test"),
        );

        // Throw away the strong handle.
        drop(poll);

        // Now, release the task. In all likelihood it will have already exited, but
        let _ignored = release_task.send(());

        // The task should never get to sending in notify_success since all its referenced had been
        // dropped at that point, but it *should* drop the channel itself.
        recv_success.await.unwrap_err();
    }

    #[tokio::test]
    async fn test_spawn() {
        let sp = Arc::new(TokioSpawner);
        let fut = async { "Hello world!" };

        let (_task, poll) = spawn_dropcancel_with_preamble(
            fut,
            futures::future::ready(()),
            sp.as_ref(),
            &MockCtx,
            tracing::debug_span!("test"),
        );

        let res = poll.await;
        assert_eq!(res, "Hello world!");
    }

    #[tokio::test]
    async fn test_cancellation_of_cancellable() {
        let (release_task, recv_release_task) = oneshot::channel();
        let (notify_success, recv_success) = oneshot::channel();

        let sp = Arc::new(TokioSpawner);

        let FutureAndCancellationHandle {
            cancellation_handle,
            ..
        } = spawn_cancellable(
            move |_| {
                async move {
                    recv_release_task.await.unwrap();
                    notify_success.send(()).unwrap();
                }
                .boxed()
            },
            sp.as_ref(),
            &MockCtx,
        );

        // Trigger cancellation
        cancellation_handle.cancel();

        // Now, release the task. In all likelihood it will have already exited, but
        let _ignored = release_task.send(());

        // The task should never get to sending in notify_success since cancellation was trigger,
        // but it *should* drop the channel itself.
        recv_success.await.unwrap_err();
    }

    #[tokio::test]
    async fn test_cancellation_of_cancellable_convert_dropcancel() {
        let (release_task, recv_release_task) = oneshot::channel();
        let (notify_success, recv_success) = oneshot::channel();

        let sp = Arc::new(TokioSpawner);

        let FutureAndCancellationHandle {
            future: task,
            cancellation_handle,
        } = spawn_cancellable(
            move |_| {
                async move {
                    recv_release_task.await.unwrap();
                    notify_success.send(()).unwrap();
                }
                .boxed()
            },
            sp.as_ref(),
            &MockCtx,
        );

        let future = task.into_drop_cancel(cancellation_handle);

        drop(future);

        // Now, release the task. In all likelihood it will have already exited, but
        let _ignored = release_task.send(());

        // The task should never get to sending in notify_success since all its referenced had been
        // dropped at that point, but it *should* drop the channel itself.
        recv_success.await.unwrap_err();
    }

    #[tokio::test]
    async fn test_spawn_cancellable() {
        let sp = Arc::new(TokioSpawner);
        let fut = async { "Hello world!" }.boxed();

        let FutureAndCancellationHandle { future: task, .. } =
            spawn_cancellable(|_| fut, sp.as_ref(), &MockCtx);

        let res = task.await;
        assert_eq!(res, Ok("Hello world!"));
    }

    #[tokio::test]
    async fn test_spawn_cancellable_convert_to_dropcancel() {
        let sp = Arc::new(TokioSpawner);
        let fut = async { "Hello world!" }.boxed();

        let FutureAndCancellationHandle {
            future: task,
            cancellation_handle,
        } = spawn_cancellable(|_| fut, sp.as_ref(), &MockCtx);

        let future = task.into_drop_cancel(cancellation_handle);

        let res = future.await;
        assert_eq!(res, "Hello world!");
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::any::Any;

use futures::future::BoxFuture;
use tokio::task::JoinHandle;

/// A spawner that utilizes a generic context to decorate/wrap a future and spawn it.
/// If a future needs to return a value, use a oneshot channel instead.
pub trait Spawner<T>: Send + Sync {
    fn spawn(
        &self,
        ctx: &T,
        fut: BoxFuture<'static, Box<dyn Any + Send + 'static>>,
    ) -> JoinHandle<Box<dyn Any + Send + 'static>>;
}

#[derive(Default)]
pub struct TokioSpawner;

impl<T> Spawner<T> for TokioSpawner {
    fn spawn(
        &self,
        _ctx: &T,
        fut: BoxFuture<'static, Box<dyn Any + Send + 'static>>,
    ) -> JoinHandle<Box<dyn Any + Send + 'static>> {
        tokio::spawn(fut)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//!
//! A future that can be canceled via an explicit `CancellationHandle`.
//! This future is intended to be spawned on tokio-runtime directly, and for its results to be
//! accessed via the joinhandle.
//! It is not intended to be polled directly.
//!

use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use std::task::Waker;

use dupe::Clone_;
use dupe::Dupe;
use dupe::Dupe_;
use futures::future::BoxFuture;
use futures::future::Shared;
use futures::FutureExt;
use parking_lot::Mutex;
use pin_project::pin_project;
use tokio::sync::oneshot;

use crate::cancellable_future::CancellationObserver;
use crate::cancellation::ExplicitCancellationContext;
use crate::maybe_future::MaybeFuture;
use crate::owning_future::OwningFuture;

pub(crate) fn make_cancellable_future<F, T>(
    f: F,
) -> (ExplicitlyCancellableFuture<T>, CancellationHandle)
where
    F: for<'a> FnOnce(&'a ExplicitCancellationContext) -> BoxFuture<'a, T> + Send,
{
    let context = ExecutionContext::new();

    let fut = {
        let context = context.dupe();
        let cancel = ExplicitCancellationContext { inner: context };

        OwningFuture::new(cancel, f)
    };

    let state = SharedState::new();

    let fut = ExplicitlyCancellableFuture::new(fut, state.dupe(), context);
    let handle = CancellationHandle::new(state);

    (fut, handle)
}

/// Defines a future that operates with the 'CancellationContext' to provide explicit cancellation.
///
/// NOTE: this future is intended only to be polled in a consistent tokio runtime, and never moved
/// from one executor to another.
/// The general safe way of using this future is to spawn it directly via `tokio::spawn`.
#[pin_project]
pub struct ExplicitlyCancellableFuture<T> {
    #[pin]
    fut: MaybeFuture<ExplicitlyCancellableFutureInner<T>>,
}

struct ExplicitlyCancellableFutureInner<T> {
    shared: SharedState,

    execution: ExecutionContext,

    /// NOTE: this is duplicative of the `SharedState`, but unlike that state this is not behind a
    /// lock. This avoids us needing to grab the lock to check if we're Pending every time we poll.
    started: bool,

    future: Pin<Box<OwningFuture<T, ExplicitCancellationContext>>>,
}

impl<T> ExplicitlyCancellableFuture<T> {
    fn new(
        future: Pin<Box<OwningFuture<T, ExplicitCancellationContext>>>,
        shared: SharedState,
        execution: ExecutionContext,
    ) -> Self {
        ExplicitlyCancellableFuture {
            fut: MaybeFuture::Fut(ExplicitlyCancellableFutureInner {
                shared,
                execution,
                started: false,
                future,
            }),
        }
    }
}

impl<T> Future for ExplicitlyCancellableFuture<T> {
    type Output = Option<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();

        match this.fut.as_mut().poll(cx) {
            Poll::Ready(res) => {
                this.fut.take();
                Poll::Ready(res)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<T> ExplicitlyCancellableFutureInner<T> {
    fn poll_inner(self: &mut Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
        let is_cancelled = self.shared.inner.cancelled.load(Ordering::SeqCst);

        if is_cancelled {
            let mut execution = self.execution.shared.lock();
            if execution.can_exit() {
                return Poll::Ready(None);
            }
            execution.notify_cancelled();
        }

        let res = Pin::new(&mut self.future).poll(cx).map(Some);

        // If we were using structured cancellation but just exited the critical section, then we
        // should exit now.
        if is_cancelled && self.execution.shared.lock().can_exit() {
            return Poll::Ready(None);
        }

        res
    }
}

impl<T> Future for ExplicitlyCancellableFutureInner<T> {
    type Output = Option<T>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // Update the state before we check for cancellation so that the cancellation logic can
        // observe whether this future has entered `poll` or not. This lets cancellation set the
        // termination observer correctly so that the state is picked up.
        // Once we start, the `poll_inner` will check whether we are actually canceled and return
        // the proper poll value.
        if !self.started {
            // we only update the Waker once at the beginning of the poll. For the same tokio
            // runtime, this is always safe and behaves correctly, as such, this future is
            // restricted to be ran on the same tokio executor and never moved from one runtime to
            // another
            take_mut::take(
                &mut *self.shared.inner.state.lock(),
                |future| match future {
                    State::Pending => State::Polled {
                        waker: cx.waker().clone(),
                    },
                    other => other,
                },
            );

            self.started = true;
        }

        let poll = self.poll_inner(cx);

        // When we exit, release our waker to ensure we don't keep create a reference cycle for
        // this task.
        if poll.is_ready() {
            let inner = self.shared.inner.dupe();
            let mut locked_state = inner.state.lock();
            let state = mem::replace(&mut *locked_state, State::Exited);

            match state {
                State::Cancelled => {
                    if self.execution.shared.lock().can_exit() {
                        return Poll::Ready(None);
                    }
                }
                _ => {}
            }
        } else if self.execution.shared.lock().should_exit() {
            // the future itself indicated that we should cancel

            return Poll::Ready(None);
        }

        poll
    }
}

pub struct CancellationHandle {
    shared_state: SharedState,
}

impl CancellationHandle {
    fn new(shared_state: SharedState) -> Self {
        CancellationHandle { shared_state }
    }

    /// Attempts to cancel the future this handle is associated with as soon as possible, returning
    /// a future that completes when the future is canceled.
    pub fn cancel(self) {
        // Store to the boolean first before we write to state.
        // This is because on `poll`, the future will update the state first then check the boolean.
        // This ordering ensures that either the `poll` has read our cancellation, and hence will
        // later notify the termination observer via the channel we store in `State::Cancelled`,
        // or that we will observe the terminated state of the future and directly notify the
        // `TerminationObserver` ourselves.
        self.shared_state
            .inner
            .cancelled
            .store(true, Ordering::SeqCst);

        match &mut *self.shared_state.inner.state.lock() {
            State::Cancelled { .. } => {
                unreachable!("We consume the CancellationHandle on cancel, so this isn't possible")
            }
            State::Exited => {
                // Nothing to do, that future is done.
            }
            state @ State::Pending => {
                // we wait for the future to `poll` once even if it has yet to do so.
                // Since we always should be spawning the `ExplicitlyCancellableFuture` on tokio,
                // it should be polled once.
                let _old = std::mem::replace(state, State::Cancelled);
            }
            state @ State::Polled { .. } => {
                let old = std::mem::replace(state, State::Cancelled);
                match old {
                    State::Polled { waker } => waker.wake(),
                    _ => {
                        unreachable!()
                    }
                }
            }
        };
    }
}

#[derive(Clone_, Dupe_)]
struct SharedState {
    inner: Arc<SharedStateData>,
}

impl SharedState {
    fn new() -> Self {
        Self {
            inner: Arc::new(SharedStateData {
                state: Mutex::new(State::Pending),
                cancelled: AtomicBool::new(false),
            }),
        }
    }
}

struct SharedStateData {
    state: Mutex<State>,

    /// When set, this future has been cancelled and should attempt to exit as soon as possible.
    cancelled: AtomicBool,
}

enum State {
    /// This future has been constructed, but not polled yet.
    Pending,

    /// This future has been polled. A waker is available.
    Polled { waker: Waker },

    /// This future has already been cancelled.
    Cancelled,

    /// This future has already finished executing.
    Exited,
}

/// Context relating to execution of the `poll` of the future. This will contain the information
/// required for the `CancellationContext` that the future holds to enter critical sections and
/// structured cancellations.
#[derive(Clone, Dupe)]
pub(crate) struct ExecutionContext {
    shared: Arc<Mutex<ExecutionContextData>>,
}

impl ExecutionContext {
    fn new() -> Self {
        Self {
            shared: Arc::new(Mutex::new(ExecutionContextData {
                cancellation_notification: {
                    let (tx, rx) = oneshot::channel();
                    CancellationNotificationData {
                        tx: CancellationNotification::Pending(tx),
                        rx: rx.shared(),
                    }
                },
                prevent_cancellation: 0,
                should_exit: false,
            })),
        }
    }

    pub fn testing() -> Self {
        Self::new()
    }

    pub(crate) fn enter_structured_cancellation(
        &self,
    ) -> (CancellationObserver, CriticalSectionGuard) {
        let mut shared = self.shared.lock();

        let observer = shared.enter_structured_cancellation();

        (observer, CriticalSectionGuard::new(&self.shared))
    }
}

pub(crate) struct CriticalSectionGuard<'a> {
    shared: Option<&'a Mutex<ExecutionContextData>>,
}

impl<'a> CriticalSectionGuard<'a> {
    fn new(shared: &'a Mutex<ExecutionContextData>) -> Self {
        Self {
            shared: Some(shared),
        }
    }

    pub(crate) fn exit_prevent_cancellation(mut self) -> bool {
        self.shared
            .take()
            .expect("should be set")
            .lock()
            .exit_prevent_cancellation()
    }

    pub(crate) fn try_to_disable_cancellation(mut self) -> bool {
        let mut shared = self.shared.take().expect("should be set").lock();
        if shared.try_to_disable_cancellation() {
            true
        } else {
            // couldn't prevent cancellation, so release our hold onto the counter
            shared.exit_prevent_cancellation();
            false
        }
    }
}

impl<'a> Drop for CriticalSectionGuard<'a> {
    fn drop(&mut self) {
        if let Some(shared) = self.shared.take() {
            // never actually exited during normal poll, but dropping this means we'll never poll
            // again, so just release the `prevent_cancellation`

            shared.lock().exit_prevent_cancellation();
        }
    }
}

struct ExecutionContextData {
    cancellation_notification: CancellationNotificationData,

    /// How many observers are preventing immediate cancellation.
    prevent_cancellation: usize,

    should_exit: bool,
}

impl ExecutionContextData {
    /// Does this future not currently prevent its cancellation?
    fn can_exit(&self) -> bool {
        self.prevent_cancellation == 0
    }

    fn should_exit(&self) -> bool {
        self.should_exit
    }

    fn enter_structured_cancellation(&mut self) -> CancellationObserver {
        self.prevent_cancellation += 1;

        CancellationObserver {
            rx: Some(self.cancellation_notification.rx.clone()),
        }
    }

    fn notify_cancelled(&mut self) {
        match &self.cancellation_notification.tx {
            CancellationNotification::Pending(_) => {
                let old = mem::replace(
                    &mut self.cancellation_notification.tx,
                    CancellationNotification::Notified,
                );
                match old {
                    CancellationNotification::Pending(tx) => {
                        let _ignored = tx.send(());
                    }
                    _ => unreachable!(),
                }
            }
            CancellationNotification::Notified => {}
            CancellationNotification::Disabled(..) => {}
        }
    }

    fn exit_prevent_cancellation(&mut self) -> bool {
        self.prevent_cancellation -= 1;

        self.prevent_cancellation == 0
    }

    fn try_to_disable_cancellation(&mut self) -> bool {
        match &self.cancellation_notification.tx {
            CancellationNotification::Pending(_) => {
                // since we know we'll never be canceled, delete any cancellation data we hold since
                // we should never notify cancelled.
                take_mut::take(
                    &mut self.cancellation_notification.tx,
                    |state| match state {
                        CancellationNotification::Pending(tx) => {
                            CancellationNotification::Disabled(tx)
                        }
                        x => x,
                    },
                );
                true
            }
            CancellationNotification::Notified => {
                // we've already sent our cancelled notification, so we can't record this future
                // as never cancelled
                false
            }
            CancellationNotification::Disabled(..) => {
                // already never cancelled
                true
            }
        }
    }
}

enum CancellationNotification {
    Pending(oneshot::Sender<()>),
    Notified,
    Disabled(oneshot::Sender<()>), // this just holds the sender alive so receives don't receive the "drop"
}

struct CancellationNotificationData {
    tx: CancellationNotification,
    rx: Shared<oneshot::Receiver<()>>,
}

#[cfg(test)]
mod tests {
    use std::future::Future;
    use std::pin::Pin;
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering;
    use std::sync::Arc;
    use std::task::Context;
    use std::task::Poll;
    use std::time::Duration;

    use assert_matches::assert_matches;
    use dupe::Dupe;
    use futures::FutureExt;
    use parking_lot::Mutex;
    use pin_project::pin_project;
    use pin_project::pinned_drop;

    use crate::cancellation::future::make_cancellable_future;
    use crate::cancellation::future::CancellationHandle;

    struct MaybePanicOnDrop {
        panic: bool,
    }

    impl Drop for MaybePanicOnDrop {
        fn drop(&mut self) {
            if self.panic {
                panic!()
            }
        }
    }

    #[tokio::test]
    async fn test_ready() {
        let (fut, _handle) = make_cancellable_future(|_| futures::future::ready(()).boxed());
        futures::pin_mut!(fut);
        assert_matches!(futures::poll!(fut), Poll::Ready(Some(())));
    }

    #[tokio::test]
    async fn test_cancel() {
        let (fut, handle) = make_cancellable_future(|_| futures::future::pending::<()>().boxed());

        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();

        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_cancel_never_polled() {
        let (fut, handle) = make_cancellable_future(|_| futures::future::pending::<()>().boxed());

        futures::pin_mut!(fut);

        handle.cancel();

        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_cancel_already_finished() {
        let (fut, handle) = make_cancellable_future(|_| futures::future::ready::<()>(()).boxed());

        futures::pin_mut!(fut);
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(Some(())));

        handle.cancel();
        // this is okay
    }

    #[tokio::test]
    async fn test_wakeup() {
        let (fut, handle) = make_cancellable_future(|_| futures::future::pending::<()>().boxed());

        let task = tokio::task::spawn(fut);
        futures::pin_mut!(task);

        assert_matches!(
            tokio::time::timeout(Duration::from_millis(100), &mut task).await,
            Err(..)
        );

        handle.cancel();

        assert_matches!(
            tokio::time::timeout(Duration::from_millis(100), &mut task).await,
            Ok(Ok(None))
        );
    }

    #[tokio::test]
    async fn test_is_dropped() {
        let dropped = Arc::new(Mutex::new(false));

        struct SetOnDrop {
            dropped: Arc<Mutex<bool>>,
        }

        impl Drop for SetOnDrop {
            fn drop(&mut self) {
                *self.dropped.lock() = true;
            }
        }

        impl Future for SetOnDrop {
            type Output = ();

            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
                Poll::Ready(())
            }
        }

        let (fut, _handle) = make_cancellable_future({
            let dropped = dropped.dupe();
            |_| SetOnDrop { dropped }.boxed()
        });

        let task = tokio::task::spawn(fut);

        task.await.unwrap();
        assert!(*dropped.lock());
    }

    #[tokio::test]
    async fn test_critical_section() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                {
                    cancellations.critical_section(tokio::task::yield_now).await;
                }
                futures::future::pending::<()>().await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // We reach the first yield. At this point there is one guard held by the critical section.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Cancel, then poll again. Cancellation is checked, *then* the guard in the future
        // is dropped and then immediately check for cancellation and yield.
        handle.cancel();

        // Poll again, this time we don't enter the future's poll because it is cancelled.
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_critical_section_noop_drop_is_allowed() {
        let (fut, _handle) = make_cancellable_future(|cancellations| {
            async {
                let section = cancellations.critical_section(futures::future::pending::<()>);
                drop(section); // Drop it within an ExecutionContext
            }
            .boxed()
        });

        fut.await;
    }

    #[tokio::test]
    async fn test_nested_critical_section() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                {
                    cancellations
                        .critical_section(|| async move { tokio::task::yield_now().await })
                        .await;
                }
                futures::future::pending::<()>().await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // We reach the first yield.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        let res = fut.await;

        assert_eq!(res, None);
    }

    #[tokio::test]
    async fn test_critical_section_cancelled_during_poll() {
        let handle_slot = Arc::new(Mutex::new(None::<CancellationHandle>));

        let (fut, handle) = make_cancellable_future({
            let handle_slot = handle_slot.dupe();

            move |cancellations| {
                async move {
                    {
                        let _cancel = handle_slot
                            .lock()
                            .take()
                            .expect("Expected the guard to be here by now")
                            .cancel();

                        cancellations
                            .critical_section(|| async {
                                let mut panic = MaybePanicOnDrop { panic: true };
                                tokio::task::yield_now().await;
                                panic.panic = false;
                            })
                            .await;
                    }
                    futures::future::pending::<()>().await
                }
                .boxed()
            }
        });
        futures::pin_mut!(fut);

        *handle_slot.lock() = Some(handle);

        // Run the future. It'll drop the guard (and cancel itself) after entering the critical
        // section while it's being polled, but it'll proceed to the end.
        fut.await;
    }

    // Cases to test:
    // - Basic
    // - Reentrant
    // - Cancel when exiting critical section (with no further wakeups)

    #[tokio::test]
    async fn test_structured_cancellation_notifies() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                cancellations
                    .with_structured_cancellation(|observer| observer)
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // Proceed all the way to awaiting the observer
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. At this point we'll cancel, and notify the observer.
        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_is_blocking() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                cancellations
                    .with_structured_cancellation(|_observer| async move {
                        let mut panic = MaybePanicOnDrop { panic: true };
                        tokio::task::yield_now().await;
                        panic.panic = false;
                    })
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // Proceed all the way to the first pending.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. We should resume and disarm the guard.
        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_cancels_on_exit() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                cancellations
                    .with_structured_cancellation(|observer| observer)
                    .await;
                futures::future::pending::<()>().await
            }
            .boxed()
        });

        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    // This is a bit of an implementation detail.
    #[tokio::test]
    async fn test_structured_cancellation_returns_to_executor() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                cancellations
                    .with_structured_cancellation(|observer| observer)
                    .await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_structured_cancellation_is_reentrant() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            {
                async move {
                    cancellations
                        .with_structured_cancellation(|o1| async move {
                            cancellations
                                .with_structured_cancellation(|o2| async move {
                                    o2.await;
                                    o1.await;
                                })
                                .await;
                        })
                        .await;
                }
                .boxed()
            }
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_with_critical_section() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                cancellations
                    .critical_section(|| async move {
                        cancellations
                            .with_structured_cancellation(|observer| async move {
                                let mut panic = MaybePanicOnDrop { panic: true };
                                tokio::task::yield_now().await;
                                panic.panic = false;

                                // we should get the cancel notification
                                observer.await;
                            })
                            .await;
                    })
                    .await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // Proceed all the way to the first pending.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. We should resume and disarm the guard.
        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_structured_cancellation_can_be_reentered() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                cancellations
                    .with_structured_cancellation(|_o1| async move {})
                    .await;
                cancellations
                    .with_structured_cancellation(|o2| async move {
                        o2.await;
                    })
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }

    #[tokio::test]
    async fn test_structured_cancellation_works_after_cancel() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                cancellations
                    .with_structured_cancellation(|_o1| async move {
                        tokio::task::yield_now().await;
                        // At this point we'll get cancelled.
                        cancellations
                            .with_structured_cancellation(|o2| async move {
                                o2.await;
                            })
                            .await;
                    })
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_disable_cancellation() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                assert!(cancellations.try_to_keep_going_on_cancellation().is_some());
                tokio::task::yield_now().await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(Some(())));
    }

    #[tokio::test]
    async fn test_disable_cancellation_already_canceled() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                assert!(cancellations.try_to_keep_going_on_cancellation().is_none());
                tokio::task::yield_now().await;
                panic!("already canceled")
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_disable_cancellation_synced_with_structured_cancellation_already_cancelled() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                cancellations
                    .with_structured_cancellation(|obs| async move {
                        tokio::task::yield_now().await;
                        futures::pin_mut!(obs);
                        assert_matches!(futures::poll!(&mut obs), Poll::Ready(()));

                        assert!(cancellations.try_to_keep_going_on_cancellation().is_none());
                    })
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_disable_cancellation_synced_with_structured_cancellation_not_cancelled() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                assert!(cancellations.try_to_keep_going_on_cancellation().is_some());

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

                cancellations
                    .with_structured_cancellation(|obs| async move {
                        futures::pin_mut!(obs);
                        assert_matches!(futures::poll!(&mut obs), Poll::Pending);

                        assert!(cancellations.try_to_keep_going_on_cancellation().is_some());
                    })
                    .await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();

        assert_matches!(futures::poll!(&mut fut), Poll::Ready(Some(())));
    }

    #[tokio::test]
    async fn test_finished_future_dropped_when_ready() {
        #[pin_project(PinnedDrop)]
        struct DropFuture(Arc<AtomicBool>);

        impl Future for DropFuture {
            type Output = ();

            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
                Poll::Ready(())
            }
        }

        #[pinned_drop]
        impl PinnedDrop for DropFuture {
            fn drop(self: Pin<&mut Self>) {
                self.0.store(true, Ordering::SeqCst);
            }
        }

        let is_dropped = Arc::new(AtomicBool::new(false));
        let fut = DropFuture(is_dropped.dupe());

        let (fut, _handle) = make_cancellable_future(|_cancellations| fut.boxed());
        futures::pin_mut!(fut);

        assert_matches!(futures::poll!(&mut fut), Poll::Ready(Some(())));

        assert!(is_dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_finished_future_dropped_when_cancelled() {
        struct DropFuture(Arc<AtomicBool>);

        impl Future for DropFuture {
            type Output = ();

            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
                Poll::Pending
            }
        }

        impl Drop for DropFuture {
            fn drop(&mut self) {
                self.0.store(true, Ordering::SeqCst);
            }
        }

        let is_dropped = Arc::new(AtomicBool::new(false));
        let fut = DropFuture(is_dropped.dupe());

        let (fut, handle) = make_cancellable_future(|_cancellations| fut.boxed());

        futures::pin_mut!(fut);
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();

        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
        assert!(is_dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_lambda_is_ran_without_poll() {
        let mut panic = MaybePanicOnDrop { panic: true };
        tokio::task::yield_now().await;
        panic.panic = false;

        let (fut, handle) = make_cancellable_future(move |_cancellations| {
            panic.panic = false;

            async move {
                panic!("polled");
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // cancel before any polls
        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_critical_section_via_prevent_cancellation() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                {
                    let prevent_cancellation = cancellations.begin_ignore_cancellation();
                    tokio::task::yield_now().await;

                    prevent_cancellation.allow_cancellations_again().await;
                }
                futures::future::pending::<()>().await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // We reach the first yield. At this point there is one guard held by the critical section.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Cancel, then poll again. Cancellation is checked, *then* the guard in the future
        // is dropped and then immediately check for cancellation and yield.
        handle.cancel();

        // Poll again, this time we don't enter the future's poll because it is cancelled.
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(None));
    }

    #[tokio::test]
    async fn test_prevent_cancellation_drop_is_allowed() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                let prevent_cancellation = cancellations.begin_ignore_cancellation();
                drop(prevent_cancellation);

                futures::future::pending::<()>().await
            }
            .boxed()
        });

        futures::pin_mut!(fut);
        // We reach the first yield.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();

        fut.await;
    }

    #[tokio::test]
    async fn test_prevent_cancellation_is_reentrant() {
        let mut panic = MaybePanicOnDrop { panic: true };
        tokio::task::yield_now().await;
        panic.panic = false;

        let (fut, handle) = make_cancellable_future(|cancellations| {
            async move {
                {
                    let prevent1 = cancellations.begin_ignore_cancellation();
                    let prevent2 = cancellations.begin_ignore_cancellation();

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

                    prevent1.allow_cancellations_again().await;

                    panic.panic = false;

                    prevent2.allow_cancellations_again().await;
                }
                futures::future::pending::<()>().await
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // We reach the first yield.
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        handle.cancel();
        let res = fut.await;

        assert_eq!(res, None);
    }

    #[tokio::test]
    async fn test_prevent_cancellation_cancellation_observer_notifies() {
        let (fut, handle) = make_cancellable_future(|cancellations| {
            async {
                let prevent_cancellation = cancellations.begin_ignore_cancellation();
                prevent_cancellation.cancellation_observer().await;
            }
            .boxed()
        });
        futures::pin_mut!(fut);

        // Proceed all the way to awaiting the observer
        assert_matches!(futures::poll!(&mut fut), Poll::Pending);

        // Drop our guard. At this point we'll cancel, and notify the observer.
        handle.cancel();
        assert_matches!(futures::poll!(&mut fut), Poll::Ready(..));
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

/// Run 3B instructions.
fn main() {
    let mut x = 1_000_000_000u64;

    #[cfg(target_arch = "x86_64")]
    unsafe {
        std::arch::asm!(
            "2:",
            "sub {0:r}, 1",
            "cmp {0:r}, 0",
            "jne 2b",
            inout(reg) x,
        );
    }

    #[cfg(target_arch = "aarch64")]
    unsafe {
        std::arch::asm!(
            "2:",
            "sub {0:x}, {0:x}, 1",
            "cmp {0:x}, 0",
            "bne 2b",
            inout(reg) x,
        );
    }

    assert_eq!(x, 0);
}