ruststream_sqlx/inbox/database/
notify.rs

1//! `LISTEN/NOTIFY`: the databases whose sessions notify each other, the broker's listening
2//! connection, and the notification a publish sends.
3//!
4//! A broker with [`listen_notify`](crate::SqlxBroker::listen_notify) holds one connection of the
5//! pool for its life, driven by one task: each subscription's open listens on its table's channel
6//! there, and each notification wakes the subscriptions of the table and group it names. A publish
7//! announces its row on the same connection right after the write, so a listener never claims
8//! before the row is visible.
9
10use std::future::Future;
11use std::io;
12use std::pin::Pin;
13use std::sync::{Arc, Mutex, PoisonError};
14
15use sqlx::{Database, Pool};
16use tokio::sync::{mpsc, oneshot};
17use tokio::task::JoinHandle;
18use tokio_util::sync::CancellationToken;
19
20use super::QueueDatabase;
21use crate::inbox::broker::Shared;
22use crate::inbox::error::SqlxBrokerError;
23use crate::inbox::publish::TableWake;
24#[cfg(feature = "testing")]
25use crate::inbox::testing::{cancelled, off_clock};
26
27/// The longest channel name Postgres keeps: an identifier of `NAMEDATALEN - 1` bytes.
28const CHANNEL_LIMIT: usize = 63;
29
30/// A database whose sessions notify each other: a broker on it may
31/// [`listen_notify`](crate::SqlxBroker::listen_notify).
32///
33/// Postgres alone implements it, with `LISTEN` and `pg_notify`.
34///
35/// # Examples
36///
37/// ```
38/// # #[cfg(feature = "postgres")]
39/// # mod demo {
40/// use ruststream_sqlx::prelude::*;
41/// use ruststream_sqlx::{BuiltInDialect, Notifies};
42/// use serde::Deserialize;
43/// use sqlx::{PgPool, Pool};
44///
45/// #[derive(Inbox, sqlx::FromRow)]
46/// #[inbox(table = "email_jobs")]
47/// pub struct SendEmail {
48///     #[field(id, generated)]
49///     job_id: i64,
50///     #[field(group)]
51///     name: String,
52///     #[field(payload)]
53///     payload: Vec<u8>,
54/// }
55///
56/// #[derive(Deserialize)]
57/// struct Email {
58///     to: String,
59/// }
60///
61/// #[subscriber(InboxQueue::<SendEmail>::new("emails"))]
62/// async fn send(email: &Email) -> HandlerOutcome {
63///     tracing::info!(to = %email.to, "sending");
64///     HandlerOutcome::ack()
65/// }
66///
67/// // The service's broker wakes on notifications, on whichever database sends them.
68/// pub fn notified<DB: Notifies + BuiltInDialect>(pool: Pool<DB>) -> SqlxBroker<DB> {
69///     SqlxBroker::new(pool).listen_notify()
70/// }
71///
72/// pub fn app(pool: PgPool) -> RustStream {
73///     RustStream::new(AppInfo::new("mailer", "0.1.0")).with_broker(notified(pool), |b| {
74///         b.include(send);
75///     })
76/// }
77/// # }
78/// # fn main() {}
79/// ```
80#[diagnostic::on_unimplemented(
81    message = "`LISTEN/NOTIFY` is Postgres's: `{Self}` has no notifications",
82    label = "this database cannot listen",
83    note = "drop `.listen_notify()`: the subscriptions wake on the poll interval and on publishes \
84            from this process"
85)]
86pub trait Notifies: QueueDatabase + sealed::Listens {}
87
88#[cfg(feature = "postgres")]
89impl Notifies for sqlx::Postgres {}
90
91/// What starts a broker's listening connection on its pool, once, at `connect`.
92pub(crate) type StartListening<DB> = fn(Pool<DB>) -> Starting;
93
94/// The start of a listening connection: boxed, once per broker.
95pub(crate) type Starting = Pin<Box<dyn Future<Output = Result<Listening, sqlx::Error>> + Send>>;
96
97/// The start of the listening connection of a database that notifies.
98pub(crate) fn start_listening<DB: Notifies>() -> StartListening<DB> {
99    <DB as sealed::Listens>::listen
100}
101
102mod sealed {
103    use sqlx::{Database, Pool};
104
105    use super::Starting;
106
107    /// The machinery of [`Notifies`](super::Notifies): a database outside this crate cannot
108    /// implement it.
109    pub trait Listens: Database {
110        /// Opens the listening connection on `pool` and starts the task that drives it.
111        #[expect(
112            private_interfaces,
113            reason = "the trait is sealed: nothing outside the crate names or implements it"
114        )]
115        fn listen(pool: Pool<Self>) -> Starting;
116    }
117}
118
119/// A broker's listening connection: the task that drives it, and the way to it.
120pub(crate) struct Listening {
121    /// The statement a publish announces its row with, binding the channel and the group.
122    notify: &'static str,
123    commands: mpsc::UnboundedSender<Listen>,
124    stop: CancellationToken,
125    task: Mutex<Option<JoinHandle<()>>>,
126}
127
128/// A subscription's request to listen on its table's channel, answered once it is listened.
129#[cfg_attr(
130    not(feature = "postgres"),
131    expect(dead_code, reason = "a listening task runs on Postgres alone")
132)]
133struct Listen {
134    table: &'static TableWake,
135    listened: oneshot::Sender<Result<(), sqlx::Error>>,
136}
137
138impl Listening {
139    /// The statement a publish announces its row with.
140    pub(crate) const fn notify(&self) -> &'static str {
141        self.notify
142    }
143
144    /// Listens on the channel of `table`, and returns once the connection listens there.
145    async fn listen(&self, table: &'static TableWake) -> Result<(), sqlx::Error> {
146        let (listened, answer) = oneshot::channel();
147        // Why a run-time check: the task stops when the broker shuts down, and a subscription may
148        // open while it does.
149        if self.commands.send(Listen { table, listened }).is_err() {
150            return Err(stopped());
151        }
152        answer.await.unwrap_or_else(|_| Err(stopped()))
153    }
154
155    /// Stops the task, which stops listening on every channel before it ends.
156    async fn close(&self) {
157        self.stop.cancel();
158        let task = self
159            .task
160            .lock()
161            .unwrap_or_else(PoisonError::into_inner)
162            .take();
163        if let Some(task) = task {
164            let _ = task.await;
165        }
166    }
167}
168
169/// The error of a listen the listening task did not answer: it stopped.
170fn stopped() -> sqlx::Error {
171    sqlx::Error::Io(io::Error::other("the listening connection is closed"))
172}
173
174/// Listens on the channel of `table` for the subscription `subscription` of `row`, when the
175/// broker listens, off a paused clock where the connection runs in process.
176pub(crate) async fn listen<DB: Database>(
177    shared: &Arc<Shared<DB>>,
178    table: &'static TableWake,
179    subscription: &str,
180    row: &'static str,
181) -> Result<(), SqlxBrokerError> {
182    let Some(listening) = &shared.listening else {
183        return Ok(());
184    };
185    // Why a startup refusal: the limit is Postgres's, and `pg_notify` refuses a longer channel
186    // while `LISTEN` cuts it, so the subscription would never hear its table.
187    if table.table().len() > CHANNEL_LIMIT {
188        return Err(SqlxBrokerError::Declaration {
189            subscription: subscription.to_owned(),
190            table: table.table().to_owned(),
191            row,
192            reason: format!(
193                "the table's qualified name is {} bytes, and a Postgres notification channel \
194                 holds at most {CHANNEL_LIMIT}: shorten the name or drop `.listen_notify()`",
195                table.table().len()
196            ),
197        });
198    }
199    #[cfg(feature = "testing")]
200    let listened = if shared.harness.in_process() {
201        let shared = Arc::clone(shared);
202        off_clock(async move {
203            match &shared.listening {
204                Some(listening) => listening.listen(table).await,
205                None => Ok(()),
206            }
207        })
208        .await
209        .unwrap_or_else(|| Err(cancelled()))
210    } else {
211        listening.listen(table).await
212    };
213    #[cfg(not(feature = "testing"))]
214    let listened = listening.listen(table).await;
215    listened.map_err(|source| SqlxBrokerError::Listen {
216        subscription: subscription.to_owned(),
217        table: table.table().to_owned(),
218        row,
219        source: Box::new(source),
220    })
221}
222
223/// Closes the broker's listening connection, when it listens, off a paused clock where the
224/// connection runs in process.
225pub(crate) async fn close<DB: Database>(shared: &Arc<Shared<DB>>) {
226    let Some(listening) = &shared.listening else {
227        return;
228    };
229    #[cfg(feature = "testing")]
230    if shared.harness.in_process() {
231        let shared = Arc::clone(shared);
232        let _ = off_clock(async move {
233            if let Some(listening) = &shared.listening {
234                listening.close().await;
235            }
236        })
237        .await;
238        return;
239    }
240    listening.close().await;
241}
242
243/// Announces a row written into `table`'s table for the group `name` on `conn`, the connection
244/// that wrote it, with the statement `notify`.
245///
246/// The row is written whatever this answers, and the poll interval still reaches it, so a failed
247/// notification is a warning, not a failed publish.
248pub(crate) async fn announce<DB: QueueDatabase>(
249    notify: &'static str,
250    conn: &mut DB::Connection,
251    table: &TableWake,
252    name: &str,
253) {
254    let mut arguments = DB::Arguments::default();
255    let announced = match DB::bind_str(&mut arguments, table.table())
256        .and_then(|()| DB::bind_str(&mut arguments, name))
257    {
258        Ok(()) => DB::execute(conn, notify, arguments).await.map(drop),
259        Err(error) => Err(error),
260    };
261    if let Err(error) = announced {
262        tracing::warn!(
263            target: "ruststream_sqlx",
264            %error,
265            name,
266            table = table.table(),
267            "the row is written, and its notification failed: subscriptions in other processes \
268             reach it on their poll interval"
269        );
270    }
271}
272
273#[cfg(feature = "postgres")]
274mod postgres {
275    use std::ptr;
276    use std::sync::Mutex;
277    use std::time::Duration;
278
279    use sqlx::postgres::{PgListener, PgNotification};
280    use sqlx::{PgPool, Postgres};
281    use tokio::sync::mpsc;
282    use tokio_util::sync::CancellationToken;
283
284    use super::sealed::Listens;
285    use super::{Listen, Listening, Starting};
286    use crate::inbox::publish::TableWake;
287
288    /// How long the task waits before it opens the listening connection again after an attempt
289    /// failed, so a database out of reach cannot spin it.
290    const RETRY: Duration = Duration::from_secs(1);
291
292    impl Listens for Postgres {
293        #[expect(
294            private_interfaces,
295            reason = "the trait is sealed: nothing outside the crate names or implements it"
296        )]
297        fn listen(pool: PgPool) -> Starting {
298            Box::pin(async move {
299                let listener = open(&pool, &[]).await?;
300                let (commands, received) = mpsc::unbounded_channel();
301                let stop = CancellationToken::new();
302                let task = tokio::spawn(run(pool, listener, received, stop.clone()));
303                Ok(Listening {
304                    notify: "SELECT pg_notify($1, $2)",
305                    commands,
306                    stop,
307                    task: Mutex::new(Some(task)),
308                })
309            })
310        }
311    }
312
313    /// A listener on a connection of `pool`, listening on the channel of each of `tables`.
314    async fn open(pool: &PgPool, tables: &[&'static TableWake]) -> Result<PgListener, sqlx::Error> {
315        let mut listener = PgListener::connect_with(pool).await?;
316        // A lost connection then comes back from `try_recv` as `Ok(None)` at once, and this task
317        // opens a listener again itself, outside the `select!` that may drop a `try_recv`: a
318        // reconnect dropped halfway would leave a pooled connection listening.
319        listener.eager_reconnect(false);
320        if !tables.is_empty() {
321            listener
322                .listen_all(tables.iter().map(|table| table.table()))
323                .await?;
324        }
325        Ok(listener)
326    }
327
328    /// Drives `listener` until `stop`: listens for each subscription that asks, wakes the
329    /// subscriptions each notification names, and listens again after a lost connection.
330    async fn run(
331        pool: PgPool,
332        mut listener: PgListener,
333        mut commands: mpsc::UnboundedReceiver<Listen>,
334        stop: CancellationToken,
335    ) {
336        let mut tables: Vec<&'static TableWake> = Vec::new();
337        loop {
338            // `try_recv` is dropped only while it waits for the server's next message, which
339            // sqlx reads cancel-safely; with `eager_reconnect` off it awaits nothing else.
340            let lost = tokio::select! {
341                biased;
342                () = stop.cancelled() => break,
343                command = commands.recv() => {
344                    // Every handle of the broker dropped without a shutdown: nobody listens.
345                    let Some(Listen { table, listened }) = command else { break };
346                    let _ = listened.send(listen(&mut listener, &mut tables, table).await);
347                    continue;
348                }
349                received = listener.try_recv() => match received {
350                    Ok(Some(notification)) => {
351                        wake(&tables, &notification);
352                        continue;
353                    }
354                    Ok(None) => None,
355                    Err(error) => Some(error),
356                },
357            };
358            // FIXME(sqlx-postgres 0.9.0): `try_recv` answers `Ok(None)` for a lost connection only
359            // when the socket reports an abort, an end of file, a timeout or a broken pipe; a
360            // reset connection comes back as an error, with the dead connection kept. Every error
361            // is taken as a lost connection here; once sqlx reports every loss as `Ok(None)`, an
362            // error goes to the log alone.
363            tracing::warn!(
364                target: "ruststream_sqlx",
365                error = lost.as_ref().map(tracing::field::display),
366                tables = ?tables.iter().map(|table| table.table()).collect::<Vec<_>>(),
367                "the listening connection was lost; listening again"
368            );
369            let Some(again) = reopen(&pool, &tables, &stop).await else {
370                break;
371            };
372            listener = again;
373            // A notification sent while no session listened is gone, so every subscription
374            // claims once.
375            for table in &tables {
376                table.wake_all();
377            }
378        }
379        // FIXME(sqlx-postgres 0.9.0): a `PgListener` hands its connection back only from its
380        // `Drop`, which spawns a task that runs `UNLISTEN *` and returns the connection to the
381        // pool, unawaited. The listener stops listening here first, awaited, so no channel is
382        // listened once the task ends; once sqlx offers an awaited close, it replaces this.
383        if let Err(error) = listener.unlisten_all().await {
384            tracing::debug!(
385                target: "ruststream_sqlx",
386                %error,
387                "the listening connection did not stop listening; it is closed"
388            );
389        }
390        drop(listener);
391    }
392
393    /// Listens on the channel of `table`, unless the connection listens there already.
394    async fn listen(
395        listener: &mut PgListener,
396        tables: &mut Vec<&'static TableWake>,
397        table: &'static TableWake,
398    ) -> Result<(), sqlx::Error> {
399        if tables.iter().any(|listened| ptr::eq(*listened, table)) {
400            return Ok(());
401        }
402        listener.listen(table.table()).await?;
403        tables.push(table);
404        Ok(())
405    }
406
407    /// Wakes the subscriptions `notification` names: its table's subscription of the group its
408    /// payload names, or every subscription of the table for an empty payload.
409    fn wake(tables: &[&'static TableWake], notification: &PgNotification) {
410        let Some(table) = tables
411            .iter()
412            .find(|table| table.table() == notification.channel())
413        else {
414            return;
415        };
416        match notification.payload() {
417            "" => table.wake_all(),
418            group => table.wake(group),
419        }
420    }
421
422    /// A listener on the channels of `tables` again, or `None` once `stop` is cancelled.
423    async fn reopen(
424        pool: &PgPool,
425        tables: &[&'static TableWake],
426        stop: &CancellationToken,
427    ) -> Option<PgListener> {
428        loop {
429            if stop.is_cancelled() {
430                return None;
431            }
432            // Not raced against `stop`: an open dropped halfway would leave its connection
433            // listening; the pool's acquire timeout bounds it.
434            match open(pool, tables).await {
435                Ok(listener) => return Some(listener),
436                Err(error) => {
437                    tracing::warn!(
438                        target: "ruststream_sqlx",
439                        %error,
440                        tables = ?tables.iter().map(|table| table.table()).collect::<Vec<_>>(),
441                        "the listening connection did not open; trying again"
442                    );
443                    tokio::select! {
444                        () = stop.cancelled() => return None,
445                        () = tokio::time::sleep(RETRY) => {}
446                    }
447                }
448            }
449        }
450    }
451}
ruststream_sqlx/inbox/database/notify.rs

ruststream_sqlx/inbox/database/
notify.rs
