Skip to main content

mas_storage/queue/
worker.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3//
4// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
5// Please see LICENSE files in the repository root for full details.
6
7//! Repository to interact with workers in the job queue
8
9use async_trait::async_trait;
10use chrono::Duration;
11use mas_data_model::Clock;
12use rand_core::Rng;
13use ulid::Ulid;
14
15use crate::repository_impl;
16
17/// A worker is an entity which can execute jobs.
18pub struct Worker {
19    /// The ID of the worker.
20    pub id: Ulid,
21}
22
23/// A [`QueueWorkerRepository`] is used to schedule jobs to be executed by a
24/// worker.
25#[async_trait]
26pub trait QueueWorkerRepository: Send + Sync {
27    /// The error type returned by the repository.
28    type Error;
29
30    /// Register a new worker.
31    ///
32    /// Returns a reference to the worker.
33    ///
34    /// # Errors
35    ///
36    /// Returns an error if the underlying repository fails.
37    async fn register(
38        &mut self,
39        rng: &mut (dyn Rng + Send),
40        clock: &dyn Clock,
41    ) -> Result<Worker, Self::Error>;
42
43    /// Send a heartbeat for the given worker.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error if the underlying repository fails or if the worker was
48    /// shutdown.
49    async fn heartbeat(&mut self, clock: &dyn Clock, worker: &Worker) -> Result<(), Self::Error>;
50
51    /// Mark the given worker as shutdown.
52    ///
53    /// # Errors
54    ///
55    /// Returns an error if the underlying repository fails.
56    async fn shutdown(&mut self, clock: &dyn Clock, worker: &Worker) -> Result<(), Self::Error>;
57
58    /// Find dead workers and shut them down.
59    ///
60    /// # Errors
61    ///
62    /// Returns an error if the underlying repository fails.
63    async fn shutdown_dead_workers(
64        &mut self,
65        clock: &dyn Clock,
66        threshold: Duration,
67    ) -> Result<(), Self::Error>;
68
69    /// Remove the leader lease if it is expired, sending a notification to
70    /// trigger a new leader election.
71    ///
72    /// # Errors
73    ///
74    /// Returns an error if the underlying repository fails.
75    async fn remove_leader_lease_if_expired(
76        &mut self,
77        clock: &dyn Clock,
78    ) -> Result<(), Self::Error>;
79
80    /// Try to get the leader lease, renewing it if we already have it
81    ///
82    /// Returns `true` if we got the leader lease, `false` if we didn't
83    ///
84    /// # Errors
85    ///
86    /// Returns an error if the underlying repository fails.
87    async fn try_get_leader_lease(
88        &mut self,
89        clock: &dyn Clock,
90        worker: &Worker,
91    ) -> Result<bool, Self::Error>;
92}
93
94repository_impl!(QueueWorkerRepository:
95    async fn register(
96        &mut self,
97        rng: &mut (dyn Rng + Send),
98        clock: &dyn Clock,
99    ) -> Result<Worker, Self::Error>;
100
101    async fn heartbeat(
102        &mut self,
103        clock: &dyn Clock,
104        worker: &Worker,
105    ) -> Result<(), Self::Error>;
106
107    async fn shutdown(
108        &mut self,
109        clock: &dyn Clock,
110        worker: &Worker,
111    ) -> Result<(), Self::Error>;
112
113    async fn shutdown_dead_workers(
114        &mut self,
115        clock: &dyn Clock,
116        threshold: Duration,
117    ) -> Result<(), Self::Error>;
118
119    async fn remove_leader_lease_if_expired(
120        &mut self,
121        clock: &dyn Clock,
122    ) -> Result<(), Self::Error>;
123
124    async fn try_get_leader_lease(
125        &mut self,
126        clock: &dyn Clock,
127        worker: &Worker,
128    ) -> Result<bool, Self::Error>;
129);