cratestack_sqlx/backplane/
pg_notify.rs

1//! A backplane over Postgres `LISTEN`/`NOTIFY`: every replica sharing the
2//! database receives every message, its own included
3//! (`docs/design/subscriptions-websocket.md` §4.2, §4.3).
4//!
5//! One listener connection per process, opened by the first subscription and
6//! kept for the life of the backplane (dropping the backplane closes it, idle
7//! or not); local subscribers fan out from it. Model
8//! events are notified by the writing transaction itself
9//! (`enqueue_event_outbox`), so Postgres delivers them only if it commits, to
10//! every listener, in commit order. When the listener's connection drops,
11//! notifications sent meanwhile are lost: every local subscription gets
12//! `Resync` and ends (§4.5).
13
14use std::sync::{Arc, Weak};
15use std::time::Duration;
16
17use cratestack_core::backplane::LocalHub;
18use cratestack_core::{BackplaneFuture, BackplaneReceiver, CratestackError, SubscriptionBackplane};
19use tokio::sync::{OnceCell, watch};
20
21use super::notify_payload::{NOTIFY_CHANNEL, Notification, encode};
22use crate::error::cratestack_error_from_sqlx;
23use crate::sqlx;
24use crate::sqlx::postgres::PgListener;
25
26pub struct PgNotifyBackplane {
27    pool: sqlx::PgPool,
28    hub: Arc<LocalHub>,
29    listening: OnceCell<()>,
30    /// Dropped with the backplane, which ends the listener task: `try_recv`
31    /// on an idle connection would otherwise wait forever.
32    alive: watch::Sender<()>,
33}
34
35impl PgNotifyBackplane {
36    /// The listener takes one connection from `pool`'s settings (not from the
37    /// pool itself). It needs a session that keeps `LISTEN` state, so not a
38    /// PgBouncer in transaction mode.
39    pub fn new(pool: sqlx::PgPool) -> Self {
40        Self {
41            pool,
42            hub: Arc::default(),
43            listening: OnceCell::new(),
44            alive: watch::channel(()).0,
45        }
46    }
47
48    async fn ensure_listening(&self) -> Result<(), CratestackError> {
49        self.listening
50            .get_or_try_init(|| async {
51                let mut listener = PgListener::connect_with(&self.pool)
52                    .await
53                    .map_err(cratestack_error_from_sqlx)?;
54                listener
55                    .listen(NOTIFY_CHANNEL)
56                    .await
57                    .map_err(cratestack_error_from_sqlx)?;
58                tokio::spawn(listen(
59                    listener,
60                    Arc::downgrade(&self.hub),
61                    self.pool.clone(),
62                    self.alive.subscribe(),
63                ));
64                Ok(())
65            })
66            .await
67            .map(|_| ())
68    }
69}
70
71impl SubscriptionBackplane for PgNotifyBackplane {
72    fn publish<'a>(
73        &'a self,
74        channel: &'a str,
75        payload: &'a str,
76    ) -> BackplaneFuture<'a, Result<(), CratestackError>> {
77        Box::pin(async move {
78            let Some(notification) = encode(channel, payload, None) else {
79                return Err(CratestackError::BadRequest(
80                    "a backplane payload must stay under Postgres' 8000-byte NOTIFY limit"
81                        .to_owned(),
82                ));
83            };
84            sqlx::query("SELECT pg_notify($1, $2)")
85                .bind(NOTIFY_CHANNEL)
86                .bind(notification)
87                .execute(&self.pool)
88                .await
89                .map_err(cratestack_error_from_sqlx)?;
90            Ok(())
91        })
92    }
93
94    fn subscribe<'a>(
95        &'a self,
96        channel: &'a str,
97    ) -> BackplaneFuture<'a, Result<BackplaneReceiver, CratestackError>> {
98        Box::pin(async move {
99            // Attach first, then make sure LISTEN is live: nothing delivered
100            // after this returns can be missed.
101            let receiver = self.hub.subscribe(channel);
102            self.ensure_listening().await?;
103            Ok(receiver)
104        })
105    }
106
107    fn publishes_model_events_in_tx(&self) -> bool {
108        true
109    }
110}
111
112/// Runs until the backplane is dropped, then closes its connection.
113/// `try_recv` reports a lost connection as `Ok(None)` and reconnects
114/// (re-`LISTEN`ing) on the next call.
115async fn listen(
116    mut listener: PgListener,
117    hub: Weak<LocalHub>,
118    pool: sqlx::PgPool,
119    mut alive: watch::Receiver<()>,
120) {
121    loop {
122        let received = tokio::select! {
123            received = listener.try_recv() => received,
124            // `Err` once the backplane (the only sender) is dropped.
125            _ = alive.changed() => return,
126        };
127        let Some(hub) = hub.upgrade() else { return };
128        match received {
129            Ok(Some(notification)) => deliver(&hub, &pool, notification.payload()).await,
130            Ok(None) => {
131                tracing::warn!(target: "cratestack", "backplane listener reconnecting; subscriptions resync");
132                hub.resync_all();
133            }
134            Err(error) => {
135                tracing::warn!(target: "cratestack", %error, "backplane listener failed; retrying");
136                hub.resync_all();
137                tokio::time::sleep(Duration::from_secs(1)).await;
138            }
139        }
140    }
141}
142
143async fn deliver(hub: &LocalHub, pool: &sqlx::PgPool, raw: &str) {
144    let Ok(notification) = serde_json::from_str::<Notification<'_>>(raw) else {
145        tracing::warn!(target: "cratestack", "unreadable backplane notification ignored");
146        return;
147    };
148    let payload = match (notification.payload, notification.event_id) {
149        (Some(payload), _) => Arc::from(payload.as_ref()),
150        (None, Some(event_id)) => {
151            match crate::descriptor::outbox_envelope_json(pool, event_id).await {
152                Ok(payload) => Arc::from(payload),
153                Err(error) => {
154                    // The event is lost for this replica: say so rather than skip it.
155                    tracing::warn!(target: "cratestack", %error, "large model event unreadable; its channel resyncs");
156                    hub.resync(&notification.channel);
157                    return;
158                }
159            }
160        }
161        (None, None) => return,
162    };
163    hub.deliver(&notification.channel, payload);
164}
cratestack_sqlx/backplane/pg_notify.rs

cratestack_sqlx/backplane/
pg_notify.rs
