// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    future::Future,
    marker::PhantomData,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use futures::future::FutureExt;
use pin_project::pin_project;
use tikv_util::time::Instant;
use tokio::sync::{Semaphore, SemaphorePermit};

use crate::coprocessor::metrics::*;

/// Limits the concurrency of heavy tasks by limiting the time spent on
/// executing `fut` before forcing to acquire a semaphore permit.
///
/// The future `fut` can always run for at least `time_limit_without_permit`,
/// but it needs to acquire a permit from the semaphore before it can continue.
pub fn limit_concurrency<'a, F: Future + 'a>(
    fut: F,
    semaphore: &'a Semaphore,
    time_limit_without_permit: Duration,
) -> impl Future<Output = F::Output> + 'a {
    ConcurrencyLimiter::new(
        semaphore
            .acquire()
            .map(|permit| permit.expect("the semaphore never be closed")),
        fut,
        time_limit_without_permit,
    )
}

#[pin_project]
struct ConcurrencyLimiter<'a, PF, F>
where
    PF: Future<Output = SemaphorePermit<'a>>,
    F: Future,
{
    #[pin]
    permit_fut: PF,
    #[pin]
    fut: F,
    time_limit_without_permit: Duration,
    execution_time: Duration,
    state: LimitationState<'a>,
    _phantom: PhantomData<&'a ()>,
}

enum LimitationState<'a> {
    NotLimited,
    Acquiring,
    Acquired(SemaphorePermit<'a>),
}

impl<'a, PF, F> ConcurrencyLimiter<'a, PF, F>
where
    PF: Future<Output = SemaphorePermit<'a>>,
    F: Future,
{
    fn new(permit_fut: PF, fut: F, time_limit_without_permit: Duration) -> Self {
        ConcurrencyLimiter {
            permit_fut,
            fut,
            time_limit_without_permit,
            execution_time: Duration::default(),
            state: LimitationState::NotLimited,
            _phantom: PhantomData,
        }
    }
}

impl<'a, PF, F> Future for ConcurrencyLimiter<'a, PF, F>
where
    PF: Future<Output = SemaphorePermit<'a>>,
    F: Future,
{
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        match this.state {
            LimitationState::NotLimited if this.execution_time > this.time_limit_without_permit => {
                match this.permit_fut.poll(cx) {
                    Poll::Ready(permit) => {
                        *this.state = LimitationState::Acquired(permit);
                        COPR_ACQUIRE_SEMAPHORE_TYPE.acquired.inc();
                    }
                    Poll::Pending => {
                        *this.state = LimitationState::Acquiring;
                        COPR_WAITING_FOR_SEMAPHORE.inc();
                        return Poll::Pending;
                    }
                }
            }
            LimitationState::Acquiring => match this.permit_fut.poll(cx) {
                Poll::Ready(permit) => {
                    *this.state = LimitationState::Acquired(permit);
                    COPR_WAITING_FOR_SEMAPHORE.dec();
                    COPR_ACQUIRE_SEMAPHORE_TYPE.acquired.inc();
                }
                Poll::Pending => {
                    return Poll::Pending;
                }
            },
            LimitationState::Acquired(permit) => {
                let _ = permit;
            }
            LimitationState::NotLimited => {}
        }
        let now = Instant::now();
        match this.fut.poll(cx) {
            Poll::Ready(res) => {
                if let LimitationState::NotLimited = this.state {
                    COPR_ACQUIRE_SEMAPHORE_TYPE.unacquired.inc();
                }
                Poll::Ready(res)
            }
            Poll::Pending => {
                *this.execution_time += now.saturating_elapsed();
                Poll::Pending
            }
        }
    }
}

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

    use tokio::{sync::oneshot, task::yield_now, time::timeout};

    use super::*;

    #[tokio::test(flavor = "current_thread")]
    async fn test_limit_concurrency() {
        let no_permits = Semaphore::new(0);
        let light_result =
            limit_concurrency(async { 42 }, &no_permits, Duration::from_millis(500)).await;
        assert_eq!(light_result, 42);

        let semaphore = Arc::new(Semaphore::new(1));
        let (task1_admitted_tx, task1_admitted_rx) = oneshot::channel();
        let (release_task1_tx, release_task1_rx) = oneshot::channel();
        let semaphore1 = semaphore.clone();
        let task1 = tokio::spawn(async move {
            limit_concurrency(
                async move {
                    thread::sleep(Duration::from_millis(10));
                    yield_now().await;
                    task1_admitted_tx.send(()).unwrap();
                    release_task1_rx.await.unwrap();
                },
                &semaphore1,
                Duration::ZERO,
            )
            .await;
        });
        timeout(Duration::from_secs(1), task1_admitted_rx)
            .await
            .unwrap()
            .unwrap();
        assert_eq!(semaphore.available_permits(), 0);

        let (task2_admitted_tx, mut task2_admitted_rx) = oneshot::channel();
        let semaphore2 = semaphore.clone();
        let task2 = tokio::spawn(async move {
            limit_concurrency(
                async move {
                    thread::sleep(Duration::from_millis(10));
                    yield_now().await;
                    task2_admitted_tx.send(()).unwrap();
                },
                &semaphore2,
                Duration::ZERO,
            )
            .await;
        });

        timeout(Duration::from_millis(100), &mut task2_admitted_rx)
            .await
            .expect_err("task 2 must wait while task 1 holds the only permit");
        release_task1_tx.send(()).unwrap();
        timeout(Duration::from_secs(1), task1)
            .await
            .unwrap()
            .unwrap();
        timeout(Duration::from_secs(1), task2_admitted_rx)
            .await
            .unwrap()
            .unwrap();
        timeout(Duration::from_secs(1), task2)
            .await
            .unwrap()
            .unwrap();
    }
}