Skip to main content

qubit_tokio_executor/
tokio_task_handle.rs

1/*******************************************************************************
2 *
3 *    Copyright (c) 2025 - 2026 Haixing Hu.
4 *
5 *    SPDX-License-Identifier: Apache-2.0
6 *
7 *    Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10use std::{
11    future::Future,
12    pin::Pin,
13    task::{
14        Context,
15        Poll,
16    },
17};
18
19use tokio::task::{
20    JoinError,
21    JoinHandle,
22};
23
24use qubit_executor::{
25    CancelResult,
26    TaskExecutionError,
27    TaskResult,
28};
29
30/// Async handle returned by Tokio-backed executor services.
31///
32/// Awaiting this handle reports the accepted task's final result, including
33/// task failure, panic, or cancellation.
34///
35/// # Type Parameters
36///
37/// * `R` - The task success value.
38/// * `E` - The task error value.
39///
40pub struct TokioTaskHandle<R, E> {
41    /// Tokio task whose output is the accepted task's final result.
42    handle: JoinHandle<TaskResult<R, E>>,
43}
44
45impl<R, E> TokioTaskHandle<R, E> {
46    /// Creates a handle from a Tokio join handle.
47    ///
48    /// # Parameters
49    ///
50    /// * `handle` - The Tokio join handle that resolves to a task result.
51    ///
52    /// # Returns
53    ///
54    /// A task handle that can be awaited.
55    #[inline]
56    pub(crate) fn new(handle: JoinHandle<TaskResult<R, E>>) -> Self {
57        Self { handle }
58    }
59
60    /// Sends a best-effort abort request to the underlying Tokio task.
61    ///
62    /// `CancelResult::Cancelled` means this handle requested Tokio abort for a
63    /// task that was not observed as finished at the instant of the check.
64    /// Completion may still win the race with cancellation, so the final task
65    /// outcome is always the value produced by awaiting this handle.
66    ///
67    /// # Returns
68    ///
69    /// [`CancelResult::Cancelled`] when an abort request was sent, or
70    /// [`CancelResult::AlreadyFinished`] if the Tokio task had already
71    /// completed.
72    #[must_use]
73    #[inline]
74    pub fn cancel(&self) -> CancelResult {
75        if self.handle.is_finished() {
76            return CancelResult::AlreadyFinished;
77        }
78        self.handle.abort();
79        CancelResult::Cancelled
80    }
81
82    /// Returns whether the underlying Tokio task has finished.
83    ///
84    /// # Returns
85    ///
86    /// `true` if the Tokio task is complete.
87    #[inline]
88    pub fn is_done(&self) -> bool {
89        self.handle.is_finished()
90    }
91}
92
93impl<R, E> Future for TokioTaskHandle<R, E> {
94    type Output = TaskResult<R, E>;
95
96    /// Polls the underlying Tokio task.
97    ///
98    /// # Parameters
99    ///
100    /// * `cx` - Async task context used to register the current waker.
101    ///
102    /// # Returns
103    ///
104    /// `Poll::Ready` with the task result when the Tokio task completes, or
105    /// `Poll::Pending` while it is still running. Tokio cancellation and panic
106    /// join errors are converted to [`TaskExecutionError`] values.
107    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
108        let this = self.get_mut();
109        match Pin::new(&mut this.handle).poll(cx) {
110            Poll::Ready(Ok(result)) => Poll::Ready(result),
111            Poll::Ready(Err(error)) => Poll::Ready(Err(join_error_to_task_error(error))),
112            Poll::Pending => Poll::Pending,
113        }
114    }
115}
116
117/// Converts a Tokio join error into a task execution error.
118///
119/// # Parameters
120///
121/// * `error` - Join error returned by Tokio for an aborted or panicked task.
122///
123/// # Returns
124///
125/// [`TaskExecutionError::Cancelled`] for aborted tasks, otherwise
126/// [`TaskExecutionError::Panicked`].
127fn join_error_to_task_error<E>(error: JoinError) -> TaskExecutionError<E> {
128    if error.is_cancelled() {
129        TaskExecutionError::Cancelled
130    } else {
131        TaskExecutionError::Panicked
132    }
133}