product-os-async-executor 0.0.20

use crate::SpawnError;
use alloc::sync::Arc;
use core::pin::Pin;
use core::task::{ready, Context, Poll};
use tokio::runtime::Handle;

use crate::{Executor, ExecutorPerform};
use crate::{Future, Timer};

/// The internal timer state, supporting both one-shot and interval timers.
#[cfg(feature = "exec_tokio")]
enum TimerState {
    /// A one-shot timer using `tokio::time::Sleep`.
    Once(Pin<alloc::boxed::Box<tokio::time::Sleep>>),
    /// An interval timer using `tokio::time::Interval`.
    Interval(tokio::time::Interval),
}

#[cfg(feature = "exec_tokio")]
/// Tokio-based executor implementation.
///
/// This executor wraps a Tokio runtime handle and provides timer functionality.
/// When cloned, the timer state is not preserved (only the handle is cloned).
pub struct TokioExecutor {
    handle: Option<Handle>,
    timer: Option<TimerState>,
}

#[cfg(feature = "exec_tokio")]
impl Executor<Handle> for TokioExecutor {
    async fn context() -> Result<Self, SpawnError> {
        let handle = Handle::current();
        Ok(Self {
            handle: Some(handle),
            timer: None,
        })
    }

    async fn set_context(&mut self, executor: Handle) {
        self.handle = Some(executor);
    }

    async fn enter_context(&self) {
        // Note: The EnterGuard returned by handle.enter() is dropped immediately
        // when this method returns, so the context is only entered briefly.
        // For persistent context, use handle.enter() directly and hold the guard.
        if let Some(handle) = &self.handle {
            let _ = handle.enter();
        }
    }

    async fn get_executor<'a>(&'a self) -> &'a Handle
    where
        Handle: 'a,
    {
        self.handle.as_ref().expect("No executor available")
    }

    fn context_sync() -> Result<Self, SpawnError> {
        let handle = Handle::current();
        Ok(Self {
            handle: Some(handle),
            timer: None,
        })
    }

    fn set_context_sync(&mut self, executor: Handle) {
        self.handle = Some(executor);
    }

    fn enter_context_sync(&self) {
        // Note: The EnterGuard is dropped immediately. See enter_context() docs.
        if let Some(handle) = &self.handle {
            let _ = handle.enter();
        }
    }

    fn get_executor_sync(&self) -> &Handle {
        self.handle.as_ref().expect("No executor available")
    }
}

#[cfg(feature = "exec_tokio")]
impl ExecutorPerform<Handle> for TokioExecutor {
    async fn spawn_in_context<F>(
        future: F,
    ) -> Result<Arc<dyn crate::Task<F::Output, Output = F::Output>>, SpawnError>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = tokio::task::spawn(future);
        Ok(Arc::new(Task::new(handle)))
    }

    async fn spawn_from_executor<E, F>(
        executor: &E,
        future: F,
    ) -> Result<Arc<dyn crate::Task<F::Output, Output = F::Output>>, SpawnError>
    where
        E: Executor<Handle>,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = executor.get_executor().await;
        Ok(Arc::new(Task::new(handle.spawn(future))))
    }

    async fn block_from_executor<E, F>(executor: &E, future: F) -> F::Output
    where
        E: Executor<Handle>,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = executor.get_executor().await;
        handle.block_on(future)
    }

    fn spawn_in_context_sync<F>(
        future: F,
    ) -> Result<Arc<dyn crate::Task<F::Output, Output = F::Output>>, SpawnError>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = tokio::task::spawn(future);
        Ok(Arc::new(Task::new(handle)))
    }

    fn spawn_from_executor_sync<E, F>(
        executor: &E,
        future: F,
    ) -> Result<Arc<dyn crate::Task<F::Output, Output = F::Output>>, SpawnError>
    where
        E: Executor<Handle>,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = executor.get_executor_sync();
        Ok(Arc::new(Task::new(handle.spawn(future))))
    }

    fn block_from_executor_sync<E, F>(executor: &E, future: F) -> F::Output
    where
        E: Executor<Handle>,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let handle = executor.get_executor_sync();
        handle.block_on(future)
    }
}

#[cfg(feature = "exec_tokio")]
impl Timer for TokioExecutor {
    async fn once(duration_millis: u32) -> Self {
        let duration = tokio::time::Duration::from_millis(u64::from(duration_millis));
        let sleep = tokio::time::sleep(duration);

        Self {
            handle: None,
            timer: Some(TimerState::Once(alloc::boxed::Box::pin(sleep))),
        }
    }

    async fn interval(duration_millis: u32) -> Self {
        let duration = tokio::time::Duration::from_millis(u64::from(duration_millis));
        let interval = tokio::time::interval(duration);

        Self {
            handle: None,
            timer: Some(TimerState::Interval(interval)),
        }
    }

    async fn cancel(&mut self) {
        self.timer = None;
    }

    fn once_sync(duration_millis: u32) -> Self {
        let duration = tokio::time::Duration::from_millis(u64::from(duration_millis));
        let sleep = tokio::time::sleep(duration);

        Self {
            handle: None,
            timer: Some(TimerState::Once(alloc::boxed::Box::pin(sleep))),
        }
    }

    fn interval_sync(duration_millis: u32) -> Self {
        let duration = tokio::time::Duration::from_millis(u64::from(duration_millis));
        let interval = tokio::time::interval(duration);

        Self {
            handle: None,
            timer: Some(TimerState::Interval(interval)),
        }
    }

    fn cancel_sync(&mut self) {
        self.timer = None;
    }

    async fn tick(&mut self) -> u32 {
        match &mut self.timer {
            None => panic!("Timer has been cancelled"),
            Some(TimerState::Once(sleep)) => {
                sleep.as_mut().await;
                // One-shot timer fires once, then cancel it
                self.timer = None;
                0
            }
            Some(TimerState::Interval(interval)) => {
                let instant = interval.tick().await;
                u32::try_from(instant.elapsed().as_millis()).unwrap_or_default()
            }
        }
    }
}

#[cfg(feature = "exec_tokio")]
impl Clone for TokioExecutor {
    /// Clones the executor. Note: timer state is not cloned.
    fn clone(&self) -> Self {
        Self {
            handle: self.handle.clone(),
            timer: None,
        }
    }
}

#[cfg(feature = "exec_tokio")]
pub struct Task<T> {
    detached: bool,
    handle: tokio::task::JoinHandle<T>,
}

#[cfg(feature = "exec_tokio")]
impl<T> Task<T> {
    pub const fn new(handle: tokio::task::JoinHandle<T>) -> Self {
        Self {
            detached: false,
            handle,
        }
    }
}

#[cfg(feature = "exec_tokio")]
impl<T> Future for Task<T> {
    type Output = T;

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

        match ready!(Pin::new(handle).poll(cx)) {
            Ok(task) => Poll::Ready(task),
            Err(e) => panic!("Task canceled or panicked: {}", e),
        }
    }
}

#[cfg(feature = "exec_tokio")]
impl<T: Send + 'static> crate::Task<T> for Task<T> {
    fn output(self) -> Pin<alloc::boxed::Box<dyn Future<Output = T> + Send>> {
        alloc::boxed::Box::pin(async move {
            match self.handle.await {
                Ok(out) => out,
                Err(e) => panic!("Problem obtaining output: {:?}", e),
            }
        })
    }

    fn detach(mut self) {
        self.detached = true;
    }

    fn drop(self) {
        if !self.detached {
            self.handle.abort()
        }
    }
}

#[cfg(all(feature = "hyper_executor", feature = "exec_tokio"))]
impl<F> hyper::rt::Executor<F> for TokioExecutor
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    fn execute(&self, fut: F) {
        let _ = Self::spawn_in_context_sync(fut);
    }
}