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}