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

mas_storage_pg/queue/
worker.rs
