mas_storage/queue/
job.rs

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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
// Copyright 2024 New Vector Ltd.
//
// SPDX-License-Identifier: AGPL-3.0-only
// Please see LICENSE in the repository root for full details.

//! Repository to interact with jobs in the job queue

use async_trait::async_trait;
use opentelemetry::trace::TraceContextExt;
use rand_core::RngCore;
use serde::{Deserialize, Serialize};
use tracing_opentelemetry::OpenTelemetrySpanExt;
use ulid::Ulid;

use crate::{repository_impl, Clock};

/// Represents a job in the job queue
pub struct Job {
    /// The ID of the job
    pub id: Ulid,

    /// The payload of the job
    pub payload: serde_json::Value,

    /// Arbitrary metadata about the job
    pub metadata: JobMetadata,
}

/// Metadata stored alongside the job
#[derive(Serialize, Deserialize, Default)]
pub struct JobMetadata {
    #[serde(default)]
    trace_id: String,

    #[serde(default)]
    span_id: String,

    #[serde(default)]
    trace_flags: u8,
}

impl JobMetadata {
    fn new(span_context: &opentelemetry::trace::SpanContext) -> Self {
        Self {
            trace_id: span_context.trace_id().to_string(),
            span_id: span_context.span_id().to_string(),
            trace_flags: span_context.trace_flags().to_u8(),
        }
    }

    /// Get the [`opentelemetry::trace::SpanContext`] from this [`JobMetadata`]
    #[must_use]
    pub fn span_context(&self) -> opentelemetry::trace::SpanContext {
        use opentelemetry::trace::{SpanContext, SpanId, TraceFlags, TraceId, TraceState};
        SpanContext::new(
            TraceId::from_hex(&self.trace_id).unwrap_or(TraceId::INVALID),
            SpanId::from_hex(&self.span_id).unwrap_or(SpanId::INVALID),
            TraceFlags::new(self.trace_flags),
            // Trace context is remote, as it comes from another service/from the database
            true,
            TraceState::NONE,
        )
    }
}

/// A trait that represents a job which can be inserted into a queue
pub trait InsertableJob: Serialize + Send {
    /// The name of the queue this job belongs to
    const QUEUE_NAME: &'static str;
}

/// A [`QueueJobRepository`] is used to schedule jobs to be executed by a
/// worker.
#[async_trait]
pub trait QueueJobRepository: Send + Sync {
    /// The error type returned by the repository.
    type Error;

    /// Schedule a job to be executed as soon as possible by a worker.
    ///
    /// # Parameters
    ///
    /// * `rng` - The random number generator used to generate a new job ID
    /// * `clock` - The clock used to generate timestamps
    /// * `queue_name` - The name of the queue to schedule the job on
    /// * `payload` - The payload of the job
    /// * `metadata` - Arbitrary metadata about the job scheduled immediately.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying repository fails.
    async fn schedule(
        &mut self,
        rng: &mut (dyn RngCore + Send),
        clock: &dyn Clock,
        queue_name: &str,
        payload: serde_json::Value,
        metadata: serde_json::Value,
    ) -> Result<(), Self::Error>;
}

repository_impl!(QueueJobRepository:
    async fn schedule(
        &mut self,
        rng: &mut (dyn RngCore + Send),
        clock: &dyn Clock,
        queue_name: &str,
        payload: serde_json::Value,
        metadata: serde_json::Value,
    ) -> Result<(), Self::Error>;
);

/// Extension trait for [`QueueJobRepository`] to help adding a job to the queue
/// through the [`InsertableJob`] trait. This isn't in the
/// [`QueueJobRepository`] trait to keep it object safe.
#[async_trait]
pub trait QueueJobRepositoryExt: QueueJobRepository {
    /// Schedule a job to be executed as soon as possible by a worker.
    ///
    /// # Parameters
    ///
    /// * `rng` - The random number generator used to generate a new job ID
    /// * `clock` - The clock used to generate timestamps
    /// * `job` - The job to schedule
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying repository fails.
    async fn schedule_job<J: InsertableJob>(
        &mut self,
        rng: &mut (dyn RngCore + Send),
        clock: &dyn Clock,
        job: J,
    ) -> Result<(), Self::Error>;
}

#[async_trait]
impl<T> QueueJobRepositoryExt for T
where
    T: QueueJobRepository,
{
    #[tracing::instrument(
        name = "db.queue_job.schedule_job",
        fields(
            queue_job.queue_name = J::QUEUE_NAME,
        ),
        skip_all,
    )]
    async fn schedule_job<J: InsertableJob>(
        &mut self,
        rng: &mut (dyn RngCore + Send),
        clock: &dyn Clock,
        job: J,
    ) -> Result<(), Self::Error> {
        // Grab the span context from the current span
        let span = tracing::Span::current();
        let ctx = span.context();
        let span = ctx.span();
        let span_context = span.span_context();

        let metadata = JobMetadata::new(span_context);
        let metadata = serde_json::to_value(metadata).expect("Could not serialize metadata");

        let payload = serde_json::to_value(job).expect("Could not serialize job");
        self.schedule(rng, clock, J::QUEUE_NAME, payload, metadata)
            .await
    }
}