mas_storage_pg/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//! A module containing the PostgreSQL implementation of the
8//! [`QueueWorkerRepository`].
9
10use async_trait::async_trait;
11use chrono::Duration;
12use mas_data_model::{Clock, UlidExt as _};
13use mas_storage::queue::{QueueWorkerRepository, Worker};
14use rand::Rng;
15use sqlx::PgConnection;
16use ulid::Ulid;
17use uuid::Uuid;
18
19use crate::{DatabaseError, ExecuteExt};
20
21/// An implementation of [`QueueWorkerRepository`] for a PostgreSQL connection.
22pub struct PgQueueWorkerRepository<'c> {
23    conn: &'c mut PgConnection,
24}
25
26impl<'c> PgQueueWorkerRepository<'c> {
27    /// Create a new [`PgQueueWorkerRepository`] from an active PostgreSQL
28    /// connection.
29    #[must_use]
30    pub fn new(conn: &'c mut PgConnection) -> Self {
31        Self { conn }
32    }
33}
34
35#[async_trait]
36impl QueueWorkerRepository for PgQueueWorkerRepository<'_> {
37    type Error = DatabaseError;
38
39    #[tracing::instrument(
40        name = "db.queue_worker.register",
41        skip_all,
42        fields(
43            worker.id,
44            db.query.text,
45        ),
46        err,
47    )]
48    async fn register(
49        &mut self,
50        rng: &mut (dyn Rng + Send),
51        clock: &dyn Clock,
52    ) -> Result<Worker, Self::Error> {
53        let now = clock.now();
54        let worker_id = Ulid::from_datetime_with_rng(now, rng);
55        tracing::Span::current().record("worker.id", tracing::field::display(worker_id));
56
57        sqlx::query!(
58            r#"
59                INSERT INTO queue_workers (queue_worker_id, registered_at, last_seen_at)
60                VALUES ($1, $2, $2)
61            "#,
62            Uuid::from(worker_id),
63            now,
64        )
65        .traced()
66        .execute(&mut *self.conn)
67        .await?;
68
69        Ok(Worker { id: worker_id })
70    }
71
72    #[tracing::instrument(
73        name = "db.queue_worker.heartbeat",
74        skip_all,
75        fields(
76            %worker.id,
77            db.query.text,
78        ),
79        err,
80    )]
81    async fn heartbeat(&mut self, clock: &dyn Clock, worker: &Worker) -> Result<(), Self::Error> {
82        let now = clock.now();
83        let res = sqlx::query!(
84            r#"
85                UPDATE queue_workers
86                SET last_seen_at = $2
87                WHERE queue_worker_id = $1 AND shutdown_at IS NULL
88            "#,
89            Uuid::from(worker.id),
90            now,
91        )
92        .traced()
93        .execute(&mut *self.conn)
94        .await?;
95
96        // If no row was updated, the worker was shutdown so we return an error
97        DatabaseError::ensure_affected_rows(&res, 1)?;
98
99        Ok(())
100    }
101
102    #[tracing::instrument(
103        name = "db.queue_worker.shutdown",
104        skip_all,
105        fields(
106            %worker.id,
107            db.query.text,
108        ),
109        err,
110    )]
111    async fn shutdown(&mut self, clock: &dyn Clock, worker: &Worker) -> Result<(), Self::Error> {
112        let now = clock.now();
113        let res = sqlx::query!(
114            r#"
115                UPDATE queue_workers
116                SET shutdown_at = $2
117                WHERE queue_worker_id = $1
118            "#,
119            Uuid::from(worker.id),
120            now,
121        )
122        .traced()
123        .execute(&mut *self.conn)
124        .await?;
125
126        DatabaseError::ensure_affected_rows(&res, 1)?;
127
128        // Remove the leader lease if we were holding it
129        let res = sqlx::query!(
130            r#"
131                DELETE FROM queue_leader
132                WHERE queue_worker_id = $1
133            "#,
134            Uuid::from(worker.id),
135        )
136        .traced()
137        .execute(&mut *self.conn)
138        .await?;
139
140        // If we were holding the leader lease, notify workers
141        if res.rows_affected() > 0 {
142            sqlx::query!(
143                r#"
144                    NOTIFY queue_leader_stepdown
145                "#,
146            )
147            .traced()
148            .execute(&mut *self.conn)
149            .await?;
150        }
151
152        Ok(())
153    }
154
155    #[tracing::instrument(
156        name = "db.queue_worker.shutdown_dead_workers",
157        skip_all,
158        fields(
159            db.query.text,
160        ),
161        err,
162    )]
163    async fn shutdown_dead_workers(
164        &mut self,
165        clock: &dyn Clock,
166        threshold: Duration,
167    ) -> Result<(), Self::Error> {
168        // Here the threshold is usually set to a few minutes, so we don't need to use
169        // the database time, as we can assume worker clocks have less than a minute
170        // skew between each other, else other things would break
171        let now = clock.now();
172        sqlx::query!(
173            r#"
174                UPDATE queue_workers
175                SET shutdown_at = $1
176                WHERE shutdown_at IS NULL
177                  AND last_seen_at < $2
178            "#,
179            now,
180            now - threshold,
181        )
182        .traced()
183        .execute(&mut *self.conn)
184        .await?;
185
186        Ok(())
187    }
188
189    #[tracing::instrument(
190        name = "db.queue_worker.remove_leader_lease_if_expired",
191        skip_all,
192        fields(
193            db.query.text,
194        ),
195        err,
196    )]
197    async fn remove_leader_lease_if_expired(
198        &mut self,
199        _clock: &dyn Clock,
200    ) -> Result<(), Self::Error> {
201        // `expires_at` is a rare exception where we use the database time, as this
202        // would be very sensitive to clock skew between workers
203        sqlx::query!(
204            r#"
205                DELETE FROM queue_leader
206                WHERE expires_at < NOW()
207            "#,
208        )
209        .traced()
210        .execute(&mut *self.conn)
211        .await?;
212
213        Ok(())
214    }
215
216    #[tracing::instrument(
217        name = "db.queue_worker.try_get_leader_lease",
218        skip_all,
219        fields(
220            %worker.id,
221            db.query.text,
222        ),
223        err,
224    )]
225    async fn try_get_leader_lease(
226        &mut self,
227        clock: &dyn Clock,
228        worker: &Worker,
229    ) -> Result<bool, Self::Error> {
230        let now = clock.now();
231        // The queue_leader table is meant to only have a single row, which conflicts on
232        // the `active` column
233
234        // If there is a conflict, we update the `expires_at` column ONLY IF the current
235        // leader is ourselves.
236
237        // `expires_at` is a rare exception where we use the database time, as this
238        // would be very sensitive to clock skew between workers
239        let res = sqlx::query!(
240            r#"
241                INSERT INTO queue_leader (elected_at, expires_at, queue_worker_id)
242                VALUES ($1, NOW() + INTERVAL '5 seconds', $2)
243                ON CONFLICT (active)
244                DO UPDATE SET expires_at = EXCLUDED.expires_at
245                WHERE queue_leader.queue_worker_id = $2
246            "#,
247            now,
248            Uuid::from(worker.id)
249        )
250        .traced()
251        .execute(&mut *self.conn)
252        .await?;
253
254        // We can then detect whether we are the leader or not by checking how many rows
255        // were affected by the upsert
256        let am_i_the_leader = res.rows_affected() == 1;
257
258        Ok(am_i_the_leader)
259    }
260}
mas_storage_pg/queue/worker.rs

mas_storage_pg/queue/
worker.rs
