rskit_worker/
task.rs

1use tokio::sync::{broadcast, oneshot};
2use tokio_util::sync::CancellationToken;
3use uuid::Uuid;
4
5use rskit_errors::AppResult;
6
7use crate::event::Event;
8
9/// Handle returned to a caller after submitting a task to the pool.
10///
11/// - `result()` awaits the final task output.
12/// - `events()` returns a broadcast receiver for intermediate [`Event`]s.
13/// - `cancel()` requests cooperative cancellation.
14pub struct TaskHandle<O: Clone + Send + 'static> {
15    /// Unique identifier assigned to this task by the pool.
16    pub id: Uuid,
17    events_rx: broadcast::Receiver<Event<O>>,
18    result_rx: oneshot::Receiver<AppResult<O>>,
19    cancel: CancellationToken,
20}
21
22impl<O: Clone + Send + 'static> TaskHandle<O> {
23    pub(crate) fn new(
24        id: Uuid,
25        events_rx: broadcast::Receiver<Event<O>>,
26        result_rx: oneshot::Receiver<AppResult<O>>,
27        cancel: CancellationToken,
28    ) -> Self {
29        Self {
30            id,
31            events_rx,
32            result_rx,
33            cancel,
34        }
35    }
36
37    /// Await the final result of the task.
38    pub async fn result(self) -> AppResult<O> {
39        match self.result_rx.await {
40            Ok(r) => r,
41            Err(_) => Err(rskit_errors::AppError::new(
42                rskit_errors::ErrorCode::Internal,
43                "worker task dropped before completing",
44            )),
45        }
46    }
47
48    /// Get a new broadcast receiver for intermediate events.
49    pub fn events(&self) -> broadcast::Receiver<Event<O>> {
50        self.events_rx.resubscribe()
51    }
52
53    /// Signal the task to cancel.
54    pub fn cancel(&self) {
55        self.cancel.cancel();
56    }
57
58    /// Clone the cancellation token
59    /// so it can be stored separately (e.g., for cancelling after the handle is consumed by `result()`).
60    pub fn cancel_token(&self) -> CancellationToken {
61        self.cancel.clone()
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use tokio::sync::{broadcast, oneshot};
68
69    use super::*;
70
71    #[tokio::test]
72    async fn dropped_result_sender_maps_to_internal_error() {
73        let (events_tx, events_rx) = broadcast::channel(1);
74        let (result_tx, result_rx) = oneshot::channel();
75        drop(events_tx);
76        drop(result_tx);
77        let handle = TaskHandle::<u32>::new(
78            Uuid::new_v4(),
79            events_rx,
80            result_rx,
81            CancellationToken::new(),
82        );
83
84        let error = handle.result().await.unwrap_err();
85
86        assert_eq!(error.code(), rskit_errors::ErrorCode::Internal);
87    }
88}
rskit_worker/task.rs

rskit_worker/
task.rs
