puerto_sqlx/events/
outbox.rs

1use async_trait::async_trait;
2use puerto_core::EventEnvelope;
3use puerto_core::events::transport::{Outbox, TransportError};
4use sqlx::{PgConnection, PgPool};
5use tokio::sync::Mutex;
6use uuid::Uuid;
7
8use super::{ENVELOPE_COLUMNS, envelope_from_pg, unavailable};
9
10/// The advisory lock the relay holds: only one worker replica relays, so the outbox order is
11/// kept. The bytes of `"puerto"`.
12pub const RELAY_LOCK_KEY: i64 = 0x7075_6572_746f;
13
14/// The outbox, as the relay reads it.
15pub struct PgOutbox {
16    pool: PgPool,
17    lock_key: i64,
18    /// The connection holding the advisory lock, once this relay leads. A session lock lives as
19    /// long as its connection, so it is kept out of the pool.
20    leader: Mutex<Option<PgConnection>>,
21}
22
23impl PgOutbox {
24    pub fn new(pool: PgPool) -> Self {
25        Self {
26            pool,
27            lock_key: RELAY_LOCK_KEY,
28            leader: Mutex::new(None),
29        }
30    }
31
32    /// Another lock key: two projects relaying from one database each need their own.
33    pub fn with_lock_key(mut self, key: i64) -> Self {
34        self.lock_key = key;
35        self
36    }
37}
38
39#[async_trait]
40impl Outbox for PgOutbox {
41    async fn try_lead(&self) -> Result<bool, TransportError> {
42        let mut leader = self.leader.lock().await;
43        if let Some(conn) = leader.as_mut() {
44            if sqlx::query("SELECT 1").execute(&mut *conn).await.is_ok() {
45                return Ok(true);
46            }
47            // The connection died, and the lock with it: compete for it again.
48            tracing::warn!("outbox: lost the relay lock's connection");
49            *leader = None;
50        }
51        let mut conn = self
52            .pool
53            .acquire()
54            .await
55            .map_err(unavailable("outbox: try_lead"))?
56            .detach();
57        let leads: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
58            .bind(self.lock_key)
59            .fetch_one(&mut conn)
60            .await
61            .map_err(unavailable("outbox: try_lead"))?;
62        if leads {
63            tracing::info!("outbox: this worker relays");
64            *leader = Some(conn);
65        }
66        Ok(leads)
67    }
68
69    async fn fetch_unpublished(&self, limit: u32) -> Result<Vec<EventEnvelope>, TransportError> {
70        let rows = sqlx::query(&format!(
71            "SELECT {ENVELOPE_COLUMNS} FROM outbox o \
72             WHERE o.published_at IS NULL ORDER BY o.seq LIMIT $1"
73        ))
74        .bind(i64::from(limit))
75        .fetch_all(&self.pool)
76        .await
77        .map_err(unavailable("outbox: fetch_unpublished"))?;
78        rows.iter()
79            .map(|row| envelope_from_pg(row, "outbox: fetch_unpublished"))
80            .collect()
81    }
82
83    async fn mark_published(&self, ids: &[Uuid]) -> Result<(), TransportError> {
84        sqlx::query(
85            "UPDATE outbox SET published_at = now() \
86             WHERE id = ANY($1) AND published_at IS NULL",
87        )
88        .bind(ids)
89        .execute(&self.pool)
90        .await
91        .map_err(unavailable("outbox: mark_published"))?;
92        Ok(())
93    }
94}
puerto_sqlx/events/outbox.rs

puerto_sqlx/events/
outbox.rs
