1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
//! Worker and worker implementations for Ora.

#![warn(clippy::pedantic, missing_docs)]
#![allow(clippy::module_name_repetitions)]

use async_trait::async_trait;
use ora_common::task::{TaskDataFormat, TaskDefinition, WorkerSelector};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

pub mod worker;
pub mod store;

/// A context that is passed to each worker task execution.
#[derive(Debug, Clone)]
pub struct TaskContext {
    task_id: Uuid,
    cancellation: CancellationToken,
}

impl TaskContext {
    /// Return the task's ID.
    #[must_use]
    #[inline]
    pub const fn task_id(&self) -> Uuid {
        self.task_id
    }

    /// Return whether the task was cancelled.
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.cancellation.is_cancelled()
    }

    /// Wait for task cancellation.
    pub async fn cancelled(&self) {
        self.cancellation.cancelled().await;
    }
}

/// A handler that works with raw input and output
/// without any task type information attached.
#[async_trait]
pub trait RawHandler {
    /// Return the selector that should be used to
    /// match tasks to this handler.
    fn selector(&self) -> &WorkerSelector;

    /// The data format of the task output.
    fn output_format(&self) -> TaskDataFormat;

    /// Execute a task.
    async fn run(&self, context: TaskContext, task: TaskDefinition) -> eyre::Result<Vec<u8>>;
}

// Not public API, do not use!
#[doc(hidden)]
pub mod _private {
    use tokio_util::sync::CancellationToken;
    use uuid::Uuid;

    use crate::TaskContext;

    #[must_use]
    pub fn new_context(task_id: Uuid, cancellation: CancellationToken) -> TaskContext {
        TaskContext {
            task_id,
            cancellation,
        }
    }
}