Skip to main content

qubit_rayon_executor/
rayon_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    sync::Arc,
14    task::{
15        Context,
16        Poll,
17    },
18};
19
20use qubit_executor::{
21    TaskHandle,
22    TaskResult,
23};
24
25use crate::{
26    pending_cancel::PendingCancel,
27    rayon_executor_service_state::RayonExecutorServiceState,
28};
29
30/// Handle returned by [`crate::RayonExecutorService`] for accepted tasks.
31///
32/// This handle supports blocking [`Self::get`], asynchronous `.await`, and
33/// best-effort cancellation before a Rayon worker starts the task.
34pub struct RayonTaskHandle<R, E> {
35    /// Shared task result observed through blocking and async APIs.
36    inner: TaskHandle<R, E>,
37    /// Stable identifier assigned by the owning executor service.
38    task_id: usize,
39    /// Shared service state used to keep cancellation counters consistent.
40    state: Arc<RayonExecutorServiceState>,
41    /// Cancellation hook that wins only before task start.
42    cancel: PendingCancel,
43}
44
45impl<R, E> RayonTaskHandle<R, E> {
46    /// Creates a Rayon task handle from a shared task handle and cancel hook.
47    ///
48    /// # Parameters
49    ///
50    /// * `inner` - Shared task handle used for result observation.
51    /// * `task_id` - Stable identifier assigned to the accepted task.
52    /// * `state` - Shared service state that owns lifecycle counters.
53    /// * `cancel` - Cancellation hook that may cancel the task before start.
54    ///
55    /// # Returns
56    ///
57    /// A handle for the accepted Rayon task.
58    pub(crate) fn new(
59        inner: TaskHandle<R, E>,
60        task_id: usize,
61        state: Arc<RayonExecutorServiceState>,
62        cancel: PendingCancel,
63    ) -> Self {
64        Self {
65            inner,
66            task_id,
67            state,
68            cancel,
69        }
70    }
71
72    /// Waits for the task to finish and returns its final result.
73    ///
74    /// # Returns
75    ///
76    /// The final task result reported through the underlying [`TaskHandle`].
77    #[inline]
78    pub fn get(self) -> TaskResult<R, E> {
79        self.inner.get()
80    }
81
82    /// Attempts to cancel the task before any Rayon worker starts it.
83    ///
84    /// # Returns
85    ///
86    /// `true` if the task was cancelled before start, or `false` if it had
87    /// already started or completed.
88    #[inline]
89    pub fn cancel(&self) -> bool {
90        self.state.cancel_pending_task(self.task_id, &self.cancel)
91    }
92
93    /// Returns whether the task has reported completion.
94    ///
95    /// # Returns
96    ///
97    /// `true` after the task has finished or has been cancelled.
98    #[inline]
99    pub fn is_done(&self) -> bool {
100        self.inner.is_done()
101    }
102}
103
104impl<R, E> Future for RayonTaskHandle<R, E> {
105    type Output = TaskResult<R, E>;
106
107    /// Polls the accepted Rayon task for completion.
108    ///
109    /// # Parameters
110    ///
111    /// * `cx` - Async task context used to register the current waker.
112    ///
113    /// # Returns
114    ///
115    /// `Poll::Ready` with the task result after completion, or
116    /// `Poll::Pending` while the task is still running or queued.
117    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
118        let this = self.get_mut();
119        Pin::new(&mut this.inner).poll(cx)
120    }
121}