Skip to main content

neptune_job_queue/
job_handle.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::Context;
4use std::task::Poll;
5
6use super::channels::JobCancelSender;
7use super::channels::JobResultReceiver;
8use super::errors::JobHandleError;
9use super::job_completion::JobCompletion;
10use super::job_id::JobId;
11
12/// A job-handle enables cancelling a job and awaiting results
13///
14/// A JobHandle can be awaited directly.  It returns a
15/// `Result<JobCompletion, JobHandleError>`
16///
17/// See [JobCompletion] and [JobHandleError] for details.
18///
19/// When the `JobHandle` is dropped a cancellation message is sent to the job
20/// task.
21#[derive(Debug)]
22pub struct JobHandle {
23    job_id: JobId,
24    result_rx: JobResultReceiver,
25    cancel_tx: JobCancelSender,
26}
27impl JobHandle {
28    // private instantiation fn.  only for use by JobQueue
29    pub(super) fn new(
30        job_id: JobId,
31        result_rx: JobResultReceiver,
32        cancel_tx: JobCancelSender,
33    ) -> Self {
34        Self {
35            job_id,
36            result_rx,
37            cancel_tx,
38        }
39    }
40
41    // indicates if job has finished processing or not.
42    //
43    // returns true if job completed normally or was cancelled or panicked.
44    //
45    // returns false if job is still waiting in the queue or is presently
46    // processing.
47    pub fn is_finished(&self) -> bool {
48        self.cancel_tx.is_closed()
49    }
50
51    /// sends cancel message to job and returns immediately.
52    ///
53    /// note: await the JobHandle after calling `cancel()` to ensure the job has
54    /// ended and obtain a [JobCompletion]
55    ///
56    /// Basic example:
57    ///
58    /// ```
59    /// use neptune_job_queue::JobQueue;
60    /// use neptune_job_queue::JobCompletion;
61    /// use neptune_job_queue::traits::Job;
62    /// use neptune_job_queue::errors::JobHandleError;
63    ///
64    /// async fn add_and_cancel_job(job_queue: &mut JobQueue<u8>, job: Box<dyn Job>) -> Result<JobCompletion, JobHandleError> {
65    ///
66    ///     let job_priority: u8 = 10;
67    ///     let job_handle = job_queue.add_job(job, job_priority).unwrap();
68    ///
69    ///     // some time later...
70    ///     tokio::time::sleep(std::time::Duration::from_secs(5)).await;
71    ///
72    ///     job_handle.cancel()?;
73    ///
74    ///     let job_completion_result = job_handle.await;
75    ///     assert!(matches!(job_completion_result, Ok(JobCompletion::Cancelled)));
76    ///
77    ///     job_completion_result
78    /// }
79    /// ```
80    ///
81    /// Sometimes it is necessary to listen for an application message that
82    /// the job needs to cancel.  This can be achieved with tokio::select!{}
83    ///
84    /// Example:
85    ///
86    /// ```
87    /// use neptune_job_queue::JobQueue;
88    /// use neptune_job_queue::JobCompletion;
89    /// use neptune_job_queue::traits::Job;
90    /// use neptune_job_queue::errors::JobHandleError;
91    ///
92    /// async fn do_some_work(
93    ///     job_queue: &mut JobQueue<u8>,
94    ///     job: Box<dyn Job>,
95    ///     cancel_work_rx: tokio::sync::oneshot::Receiver<()>,
96    /// ) -> Result<JobCompletion, JobHandleError> {
97    ///
98    ///     // add the job to queue
99    ///     let job_priority: u8 = 10;
100    ///     let job_handle = job_queue.add_job(job, job_priority).unwrap();
101    ///
102    ///     // pin job_handle, so borrow checker knows the address can't change
103    ///     // and it is safe to use in both select branches
104    ///     tokio::pin!(job_handle);
105    ///
106    ///     // execute job and simultaneously listen for cancel msg from elsewhere
107    ///     let job_completion_result = tokio::select! {
108    ///         // case: job completion.
109    ///         completion = &mut job_handle => completion,
110    ///
111    ///         // case: sender cancelled, or sender dropped.
112    ///         _ = cancel_work_rx => {
113    ///             job_handle.cancel()?;
114    ///             job_handle.await
115    ///         }
116    ///     };
117    ///     job_completion_result
118    /// }
119    /// ```
120    pub fn cancel(&self) -> Result<(), JobHandleError> {
121        let result = self.cancel_tx.send(()).map_err(JobHandleError::from);
122
123        match result {
124            Ok(_) => tracing::debug!("Sent job-cancel msg to job: {}", self.job_id),
125            Err(ref e) => tracing::error!("{}", e),
126        };
127
128        result
129    }
130
131    /// obtain the job identifier
132    pub fn job_id(&self) -> JobId {
133        self.job_id
134    }
135}
136
137// we implement Future for JobHandle so that a JobHandle can be
138// directly awaited (like a tokio JoinHandle).
139impl Future for JobHandle {
140    type Output = Result<JobCompletion, JobHandleError>;
141
142    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
143        // Directly poll the underlying result_rx
144        let result_rx = &mut self.get_mut().result_rx;
145        Pin::new(result_rx).poll(cx).map_err(|e| e.into())
146    }
147}
148
149impl Drop for JobHandle {
150    fn drop(&mut self) {
151        tracing::debug!("JobHandle dropping for job: {}", self.job_id);
152        if !self.cancel_tx.is_closed() {
153            let _ = self.cancel();
154        }
155    }
156}