ruststream_sqlx/outbox/
publish.rs

1//! The publish middleware: a message published under a registered name is recorded before it is
2//! sent, and carries its record's id.
3
4#[cfg(feature = "inbox")]
5use std::any::Any as Dynamic;
6use std::any::type_name;
7use std::error::Error as StdError;
8use std::fmt::{self, Display};
9use std::io::Write;
10use std::mem;
11use std::str::{self, FromStr};
12use std::sync::Arc;
13
14use ruststream::runtime::{Outgoing, PublishLayer, PublishNext, PublishPipeline};
15use ruststream::{Bytes, HeaderMap, OutgoingMessage, Publisher};
16use sqlx::Database;
17#[cfg(feature = "inbox")]
18use sqlx::any::AnyConnectOptions;
19#[cfg(feature = "inbox")]
20use sqlx::error::DatabaseError;
21#[cfg(all(feature = "inbox", feature = "sqlite"))]
22use sqlx::sqlite::SqliteConnectOptions;
23#[cfg(feature = "inbox")]
24use sqlx::{Pool, SqlStr};
25#[cfg(feature = "inbox")]
26use sqlx_core::transaction::TransactionManager;
27#[cfg(feature = "inbox")]
28use tokio::time::{error::Elapsed, timeout};
29
30use super::OUTBOX_ID_HEADER;
31use super::error::OutboxError;
32use super::events::Publish;
33use super::registry::RecordList;
34use super::switch::enabled;
35#[cfg(feature = "inbox")]
36use crate::inbox::form::advisory::database_of;
37#[cfg(feature = "inbox")]
38use crate::inbox::writers::{self, Holding};
39use crate::outbox::store::Store;
40
41/// The publish middleware of an [`Outbox`](super::Outbox), from
42/// [`publish_layer`](super::Outbox::publish_layer).
43///
44/// A message under a registered name is recorded through its record type's `Publish`, then sent
45/// with [`OUTBOX_ID_HEADER`]; a record that fails fails the publish, and the message is not sent.
46/// A message under any other name is sent as it is.
47pub struct TrackingPublishLayer<DB: Database, Records> {
48    store: Arc<Store<DB>>,
49    records: Records,
50}
51
52impl<DB: Database, Records> TrackingPublishLayer<DB, Records> {
53    pub(super) const fn new(store: Arc<Store<DB>>, records: Records) -> Self {
54        Self { store, records }
55    }
56}
57
58impl<DB: Database, Records: Copy> Clone for TrackingPublishLayer<DB, Records> {
59    fn clone(&self) -> Self {
60        Self {
61            store: Arc::clone(&self.store),
62            records: self.records,
63        }
64    }
65}
66
67impl<DB: Database, Records: fmt::Debug> fmt::Debug for TrackingPublishLayer<DB, Records> {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        f.debug_struct("TrackingPublishLayer")
70            .field("records", &self.records)
71            .finish_non_exhaustive()
72    }
73}
74
75impl<DB: Database, Records: RecordList<DB>> PublishLayer for TrackingPublishLayer<DB, Records> {
76    async fn on_publish<'a, N: PublishPipeline, P: Publisher>(
77        &'a self,
78        out: &'a mut Outgoing<'a>,
79        next: PublishNext<'a, N, P>,
80    ) -> Result<(), Box<dyn StdError + Send + Sync>> {
81        if enabled() && self.records.contains(out.name()) {
82            // The map moves out and back: the record reads it through an `OutgoingMessage`
83            // that borrows the name and the payload, so nothing is copied.
84            let headers = mem::take(out.headers_mut());
85            let msg =
86                OutgoingMessage::with_payload(out.name(), out.payload()).with_headers(headers);
87            let recorded = self.records.record(&self.store, &msg).await;
88            let (_, _, headers) = msg.into_parts();
89            *out.headers_mut() = headers;
90            if let Some(id) = recorded {
91                out.headers_mut().insert(OUTBOX_ID_HEADER, id?);
92            }
93        }
94        next.run(out).await
95    }
96}
97
98/// Records `msg` with `Record`, and returns the id header's value.
99pub(super) async fn record<DB: Database, Record: Publish<DB>>(
100    name: &'static str,
101    store: &Store<DB>,
102    msg: &OutgoingMessage<'_>,
103) -> Result<Bytes, OutboxError> {
104    let pool = store.get().ok_or(OutboxError::NoPool { name })?;
105    let failed = |source| OutboxError::Record {
106        name,
107        record: type_name::<Record>(),
108        source,
109    };
110    // One load in a process with no transactional subscription on SQLite.
111    #[cfg(feature = "inbox")]
112    if writers::any_listed()
113        && let Some(database) = store.one_writer(|| one_writer(pool))
114    {
115        match writers::holding(database) {
116            Holding::Nothing => {}
117            Holding::HoldsWriteLock => {
118                return Err(OutboxError::WriteLockHeld {
119                    name,
120                    record: type_name::<Record>(),
121                });
122            }
123            Holding::MayHoldWriteLock => {
124                return record_alone::<DB, Record>(name, store, pool, msg).await;
125            }
126        }
127    }
128    let mut conn = store.acquire(pool).await.map_err(failed)?;
129    let id = Record::publish(&mut conn, msg).await.map_err(failed)?;
130    Ok(id_value(&id))
131}
132
133/// The database `pool` reaches, as the process tells databases apart, where it takes one writer at
134/// a time: SQLite, through its own driver or `Any`.
135#[cfg(feature = "inbox")]
136fn one_writer<DB: Database>(pool: &Pool<DB>) -> Option<u64> {
137    let options = pool.connect_options();
138    let options: &dyn Dynamic = &*options;
139    #[cfg(feature = "sqlite")]
140    let sqlite = options.is::<SqliteConnectOptions>();
141    #[cfg(not(feature = "sqlite"))]
142    let sqlite = false;
143    let sqlite = sqlite
144        || options
145            .downcast_ref::<AnyConnectOptions>()
146            .is_some_and(|any| any.database_url.scheme() == "sqlite");
147    sqlite.then(|| database_of(pool))
148}
149
150/// Records `msg` with `Record` on SQLite for a handler whose own delivery transaction may hold the
151/// write lock: the record takes the lock first, for at most the pool's acquire timeout, and fails
152/// naming the lock instead of waiting for the handler that waits for it.
153///
154/// Why a bound and not a check: whether the handler's transaction wrote is the database's to know,
155/// and SQLite tells it only by the wait. Why the lock is taken on its own: a statement that waits
156/// out the bound still runs once the lock is free, and a lock taken alone writes nothing then.
157#[cfg(feature = "inbox")]
158async fn record_alone<DB: Database, Record: Publish<DB>>(
159    name: &'static str,
160    store: &Store<DB>,
161    pool: &Pool<DB>,
162    msg: &OutgoingMessage<'_>,
163) -> Result<Bytes, OutboxError> {
164    let record = type_name::<Record>();
165    let failed = |source| OutboxError::Record {
166        name,
167        record,
168        source,
169    };
170    let waited = pool.options().get_acquire_timeout();
171    let held = || OutboxError::WriteLockHeld { name, record };
172    let mut conn = store.acquire(pool).await.map_err(failed)?;
173    let begin = DB::TransactionManager::begin(&mut conn, Some(SqlStr::from_static(TAKE_LOCK)));
174    match timeout(waited, begin).await {
175        Ok(Ok(())) => {}
176        Ok(Err(error)) if locked(&error) => return Err(held()),
177        Ok(Err(error)) => {
178            conn.close_on_drop();
179            return Err(failed(error));
180        }
181        Err(Elapsed { .. }) => {
182            // The begin still waits on the connection: it closes, and its transaction with it.
183            conn.close_on_drop();
184            return Err(held());
185        }
186    }
187    let written = match Record::publish(&mut conn, msg).await {
188        Ok(id) => DB::TransactionManager::commit(&mut conn).await.map(|()| id),
189        Err(error) => Err(error),
190    };
191    match written {
192        Ok(id) => Ok(id_value(&id)),
193        Err(error) => {
194            // A transaction left open closes with its connection, which then leaves the pool.
195            if DB::TransactionManager::rollback(&mut conn).await.is_err() {
196                conn.close_on_drop();
197            }
198            Err(if locked(&error) {
199                held()
200            } else {
201                failed(error)
202            })
203        }
204    }
205}
206
207/// How the record takes SQLite's write lock before it writes.
208#[cfg(feature = "inbox")]
209const TAKE_LOCK: &str = "BEGIN IMMEDIATE";
210
211/// Whether SQLite refused `error`'s statement because another connection holds a lock it needs:
212/// `SQLITE_BUSY` or `SQLITE_LOCKED`, in any of their extended codes.
213#[cfg(feature = "inbox")]
214fn locked(error: &sqlx::Error) -> bool {
215    const SQLITE_BUSY: i32 = 5;
216    const SQLITE_LOCKED: i32 = 6;
217    error
218        .as_database_error()
219        .and_then(DatabaseError::code)
220        .and_then(|code| code.parse::<i32>().ok())
221        .is_some_and(|code| matches!(code & 0xff, SQLITE_BUSY | SQLITE_LOCKED))
222}
223
224/// The longest id text [`id_value`] formats on the stack: a UUID's 36 characters fit.
225const ID_ON_STACK: usize = 64;
226
227/// The id header's value: the id's `Display` text in a buffer of its own length.
228///
229/// One allocation. `to_string` would take two: its `String` reserves more than an integer's
230/// digits, and a `Bytes` made from a vector with spare capacity allocates a block to share it.
231pub(super) fn id_value(id: &impl Display) -> Bytes {
232    let mut text = [0_u8; ID_ON_STACK];
233    let mut rest = &mut text[..];
234    if write!(rest, "{id}").is_ok() {
235        let len = ID_ON_STACK - rest.len();
236        return Bytes::copy_from_slice(&text[..len]);
237    }
238    Bytes::copy_from_slice(id.to_string().as_bytes())
239}
240
241/// The id a delivery carries in [`OUTBOX_ID_HEADER`]: `None` without the header, the header's
242/// text when it does not parse as an `Id`.
243pub(super) fn carried_id<Id: FromStr>(headers: &HeaderMap) -> Option<Result<Id, String>> {
244    let value = headers.get(OUTBOX_ID_HEADER)?;
245    Some(
246        str::from_utf8(value)
247            .ok()
248            .and_then(|text| text.parse().ok())
249            .ok_or_else(|| String::from_utf8_lossy(value).into_owned()),
250    )
251}
252
253#[cfg(test)]
254mod tests {
255    use ruststream::HeaderMap;
256
257    use super::{ID_ON_STACK, OUTBOX_ID_HEADER, carried_id, id_value};
258
259    #[test]
260    fn an_id_survives_the_header() {
261        let mut headers = HeaderMap::new();
262        headers.insert(OUTBOX_ID_HEADER, 42_i64.to_string());
263        assert_eq!(carried_id::<i64>(&headers), Some(Ok(42)));
264        headers.insert(OUTBOX_ID_HEADER, "7f3a".to_owned());
265        assert_eq!(carried_id::<String>(&headers), Some(Ok("7f3a".to_owned())));
266    }
267
268    #[test]
269    fn an_id_value_holds_the_ids_text_whatever_its_length() {
270        assert_eq!(&id_value(&42_i64)[..], b"42");
271        assert_eq!(&id_value(&i64::MIN)[..], i64::MIN.to_string().as_bytes());
272        let long = "7".repeat(ID_ON_STACK + 1);
273        assert_eq!(&id_value(&long)[..], long.as_bytes());
274    }
275
276    #[test]
277    fn a_missing_or_foreign_header_is_told_apart() {
278        let mut headers = HeaderMap::new();
279        assert_eq!(carried_id::<i64>(&headers), None);
280        headers.insert(OUTBOX_ID_HEADER, "seven");
281        assert_eq!(carried_id::<i64>(&headers), Some(Err("seven".to_owned())));
282        headers.insert(OUTBOX_ID_HEADER, &b"\xff7"[..]);
283        assert_eq!(
284            carried_id::<i64>(&headers),
285            Some(Err("\u{fffd}7".to_owned()))
286        );
287    }
288}
ruststream_sqlx/outbox/publish.rs

ruststream_sqlx/outbox/
publish.rs
