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}