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);