ruststream_sqlx/inbox/
batch.rs

1//! What a batch handler is handed: a payload-mode table's deliveries, or a row-mode table's rows
2//! as one slice ([`RowBatch`]), chosen by the table's lane.
3
4use std::fmt;
5use std::sync::Arc;
6use std::vec;
7
8use ruststream::CarriesBatch;
9
10#[cfg(feature = "testing")]
11use super::broker::Shared;
12use super::database::QueueDatabase;
13use super::delivery::{Hold, InboxDelivery};
14use super::engine::{Claimed, Events};
15use super::form::advisory::LockHold;
16use super::form::lease::{LeaseBook, Slot};
17use super::form::row_lock::BatchTx;
18use super::headers::HeaderCell;
19use super::queue::Queue;
20use super::subscriber::{InboxSubscriber, Taken};
21use super::{Lane, PayloadLane, PayloadRow, QueueRow, RowLane};
22use crate::inbox::pools::ServicePool;
23
24/// How a table's lane hands a batch handler what one claim took. Machinery: a subscription in the
25/// plain mode builds its batches through it, so a table's mode picks its batch type.
26#[doc(hidden)]
27pub trait BatchLane<DB: QueueDatabase, Row: Events<DB>>: Lane<Row> {
28    /// The batch a subscription to the table yields.
29    type Batch: IntoIterator<Item = InboxDelivery<DB, Row>> + Send;
30
31    /// The batch of what `claim` took.
32    fn batch(claim: BatchClaim<'_, DB, Row>) -> Self::Batch;
33}
34
35/// What one claim of a batch subscription took, and the subscription it took it for. Machinery.
36#[doc(hidden)]
37pub struct BatchClaim<'s, DB: QueueDatabase, Row: Events<DB>> {
38    subscriber: &'s mut InboxSubscriber<DB, Row>,
39    taken: Taken<DB, Row>,
40    count: usize,
41}
42
43impl<'s, DB: QueueDatabase, Row: Events<DB>> BatchClaim<'s, DB, Row> {
44    /// The `count` rows `subscriber`'s last claim took, held as `taken` says.
45    pub(crate) const fn new(
46        subscriber: &'s mut InboxSubscriber<DB, Row>,
47        taken: Taken<DB, Row>,
48        count: usize,
49    ) -> Self {
50        Self {
51            subscriber,
52            taken,
53            count,
54        }
55    }
56}
57
58impl<DB: QueueDatabase, Row: Events<DB>> fmt::Debug for BatchClaim<'_, DB, Row> {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        f.debug_struct("BatchClaim")
61            .field("subscription", &self.subscriber.queue().name)
62            .field("count", &self.count)
63            .finish_non_exhaustive()
64    }
65}
66
67/// A payload-mode batch is its deliveries, each holding its row, as a codec decodes them.
68impl<DB, Row> BatchLane<DB, Row> for PayloadLane
69where
70    DB: QueueDatabase,
71    Row: Events<DB> + PayloadRow,
72{
73    type Batch = Vec<InboxDelivery<DB, Row>>;
74
75    fn batch(claim: BatchClaim<'_, DB, Row>) -> Self::Batch {
76        let BatchClaim {
77            subscriber,
78            taken,
79            count,
80        } = claim;
81        let (queue, pool) = (subscriber.queue(), subscriber.pool());
82        #[cfg(feature = "testing")]
83        let shared = Arc::clone(&subscriber.shared);
84        let on = |delivery: InboxDelivery<DB, Row>| {
85            #[cfg(feature = "testing")]
86            let delivery = delivery.on(&shared);
87            delivery
88        };
89        match taken {
90            Taken::Locked(tx) => {
91                let batch = BatchTx::new(tx, count);
92                subscriber
93                    .take_rows()
94                    .map(|claimed| {
95                        on(InboxDelivery::batched(
96                            claimed,
97                            Arc::clone(&batch),
98                            queue,
99                            pool,
100                        ))
101                    })
102                    .collect()
103            }
104            // Each delivery of a leased batch holds its own lease and settles on its own.
105            Taken::Leased(book, lease) => subscriber
106                .take_rows()
107                .map(|claimed| {
108                    on(InboxDelivery::leased(
109                        claimed, book, lease, None, queue, pool,
110                    ))
111                })
112                .collect(),
113            // Each delivery of an advisory batch holds its own session and settles on its own.
114            Taken::Advised => subscriber
115                .take_advised()
116                .map(|(claimed, hold)| on(InboxDelivery::advised(claimed, hold, queue, pool)))
117                .collect(),
118        }
119    }
120}
121
122/// A row-mode batch lends its rows as one slice.
123impl<DB, Row> BatchLane<DB, Row> for RowLane
124where
125    DB: QueueDatabase,
126    Row: Events<DB> + QueueRow<Lane = Self>,
127{
128    type Batch = RowBatch<DB, Row>;
129
130    fn batch(claim: BatchClaim<'_, DB, Row>) -> Self::Batch {
131        RowBatch::new(claim)
132    }
133}
134
135/// The rows one claim of a table in row mode took, which a batch handler borrows as one slice,
136/// `&[Row]`, and the claims that settle them.
137///
138/// The rows lie in one vector in claim order, as the driver read them: the handler reads them
139/// where they lie, with no codec and no copy. Each row's headers column was taken into its
140/// delivery's headers when the batch was built, so the handler reads the column empty. A message
141/// assembled from a headers struct keeps its headers struct whole, and each delivery builds its
142/// header map from it on the first read. A claimed id without a row, and a row the driver could
143/// not read, follow the slice: the runtime settles them by the subscription's decode policy, its
144/// error in the log. Every delivery of the batch
145/// holds its row in the table's form, as a payload-mode batch's deliveries do: one transaction
146/// for the whole batch in the row lock form, a lease the subscription extends while the handler
147/// runs in the lease form, a session that holds the row's key in the advisory lock form.
148///
149/// A batch dropped before its deliveries settle (its handler panicked, the app shut down)
150/// returns its rows at once, as a dropped delivery does.
151///
152/// # Examples
153///
154/// ```no_run
155/// # #[cfg(feature = "postgres")]
156/// # mod demo {
157/// use ruststream_sqlx::prelude::*;
158/// use sqlx::PgPool;
159///
160/// /// A mail to send: the table has no payload column, so a handler takes the rows themselves.
161/// #[derive(Debug, Clone, Inbox, sqlx::FromRow)]
162/// #[inbox(table = "mail_jobs")]
163/// pub struct SendEmail {
164///     #[field(id)]
165///     job_id: i64,
166///     recipient: String,
167/// }
168///
169/// # async fn deliver(_: &str) -> bool { true }
170/// /// Each claim of up to 50 rows reaches the handler as one slice; each mail settles on its own.
171/// #[subscriber(InboxQueue::<SendEmail>::new("mail"))]
172/// async fn send(mails: &[SendEmail]) -> Vec<HandlerOutcome> {
173///     let mut outcomes = Vec::with_capacity(mails.len());
174///     for mail in mails {
175///         outcomes.push(if deliver(&mail.recipient).await {
176///             HandlerOutcome::ack()
177///         } else {
178///             HandlerOutcome::retry()
179///         });
180///     }
181///     outcomes
182/// }
183///
184/// pub fn app(pool: PgPool) -> RustStream {
185///     RustStream::new(AppInfo::new("mailer", "0.1.0")).with_broker(SqlxBroker::new(pool), |b| {
186///         b.include(send.batch(nonzero!(50)));
187///     })
188/// }
189/// # }
190/// # fn main() {}
191/// ```
192pub struct RowBatch<DB, Row>
193where
194    DB: QueueDatabase,
195    Row: Events<DB> + QueueRow<Lane = RowLane>,
196{
197    deliveries: RowDeliveries<DB, Row>,
198}
199
200/// The deliveries of a [`RowBatch`], built as the runtime takes them: the rows' first, in claim
201/// order, then those of the ids without a row and the rows the driver could not read. Machinery.
202///
203/// Dropped before it yielded them all, it builds and drops the rest, which returns their rows as
204/// dropped deliveries return theirs.
205#[doc(hidden)]
206pub struct RowDeliveries<DB, Row>
207where
208    DB: QueueDatabase,
209    Row: Events<DB> + QueueRow<Lane = RowLane>,
210{
211    /// The rows the handler borrows, in claim order.
212    rows: vec::IntoIter<Row>,
213    /// The headers taken from the rows, beside them; empty and unallocated while every row's are
214    /// unset (a message assembled from a headers struct builds its own on the first read), and
215    /// shorter than the rows when the last rows' are.
216    headers: vec::IntoIter<Row::Headers>,
217    /// What holds the rows, beside them.
218    claims: Claims<DB, Row>,
219    /// The deliveries with no row to lend, in claim order; empty and unallocated when every row
220    /// was read.
221    gone: vec::IntoIter<InboxDelivery<DB, Row>>,
222    queue: &'static Queue,
223    pool: &'static ServicePool<DB>,
224    /// The connection the batch was claimed on: a delivery claimed in process keeps the harness's
225    /// books.
226    #[cfg(feature = "testing")]
227    shared: Arc<Shared<DB>>,
228}
229
230/// What holds a row-mode batch's rows, by the table's form.
231enum Claims<DB: QueueDatabase, Row: Events<DB>> {
232    /// The claim's transaction, which every delivery of the batch shares.
233    Locked(Arc<BatchTx<DB>>),
234    /// The book that keeps each row's lease, and each row's place in it, entered when the batch
235    /// was built so the keeper extends the leases while the handler runs.
236    Leased(&'static LeaseBook<DB, Row>, vec::IntoIter<Slot>),
237    /// The place of each row in the book whose session holds the row's key.
238    Advised(vec::IntoIter<LockHold<DB>>),
239}
240
241impl<DB, Row> RowBatch<DB, Row>
242where
243    DB: QueueDatabase,
244    Row: Events<DB> + QueueRow<Lane = RowLane>,
245{
246    /// Moves what `claim` took into the batch: each row once, into the slice, its headers out of
247    /// it, and its claim beside it.
248    fn new(claim: BatchClaim<'_, DB, Row>) -> Self {
249        let BatchClaim {
250            subscriber,
251            taken,
252            count,
253        } = claim;
254        let (queue, pool) = (subscriber.queue(), subscriber.pool());
255        #[cfg(feature = "testing")]
256        let shared = Arc::clone(&subscriber.shared);
257        let on = |delivery: InboxDelivery<DB, Row>| {
258            #[cfg(feature = "testing")]
259            let delivery = delivery.on(&shared);
260            delivery
261        };
262        let mut rows = Vec::with_capacity(count);
263        let mut headers = Vec::new();
264        let mut gone = Vec::new();
265        let claims = match taken {
266            Taken::Locked(tx) => {
267                let batch = BatchTx::new(tx, count);
268                for claimed in subscriber.take_rows() {
269                    match claimed {
270                        Claimed::Row(row) => take(row, &mut rows, &mut headers, count),
271                        claimed => gone.push(on(InboxDelivery::batched(
272                            claimed,
273                            Arc::clone(&batch),
274                            queue,
275                            pool,
276                        ))),
277                    }
278                }
279                Claims::Locked(batch)
280            }
281            Taken::Leased(book, lease) => {
282                let mut slots = Vec::with_capacity(count);
283                for claimed in subscriber.take_rows() {
284                    match claimed {
285                        Claimed::Row(row) => {
286                            slots.push(book.enter(Row::id(&row), lease));
287                            take(row, &mut rows, &mut headers, count);
288                        }
289                        claimed => gone.push(on(InboxDelivery::leased(
290                            claimed, book, lease, None, queue, pool,
291                        ))),
292                    }
293                }
294                Claims::Leased(book, slots.into_iter())
295            }
296            Taken::Advised => {
297                let mut holds = Vec::with_capacity(count);
298                for (claimed, hold) in subscriber.take_advised() {
299                    match claimed {
300                        Claimed::Row(row) => {
301                            holds.push(hold);
302                            take(row, &mut rows, &mut headers, count);
303                        }
304                        claimed => {
305                            gone.push(on(InboxDelivery::advised(claimed, hold, queue, pool)));
306                        }
307                    }
308                }
309                Claims::Advised(holds.into_iter())
310            }
311        };
312        Self {
313            deliveries: RowDeliveries {
314                rows: rows.into_iter(),
315                headers: headers.into_iter(),
316                claims,
317                gone: gone.into_iter(),
318                queue,
319                pool,
320                #[cfg(feature = "testing")]
321                shared,
322            },
323        }
324    }
325}
326
327/// Moves `row` into `rows` and its headers into `headers`, which allocates for the first row that
328/// carries any, room for all `count` rows of the claim.
329fn take<DB: QueueDatabase, Row: Events<DB>>(
330    mut row: Row,
331    rows: &mut Vec<Row>,
332    headers: &mut Vec<Row::Headers>,
333    count: usize,
334) {
335    let taken = Row::Headers::take(&mut row);
336    if !headers.is_empty() || !taken.is_unset() {
337        if headers.is_empty() {
338            headers.reserve_exact(count);
339            headers.resize_with(rows.len(), Row::Headers::default);
340        }
341        headers.push(taken);
342    }
343    rows.push(row);
344}
345
346impl<DB, Row> CarriesBatch<Row> for RowBatch<DB, Row>
347where
348    DB: QueueDatabase,
349    Row: Events<DB> + QueueRow<Lane = RowLane>,
350{
351    fn carried(&self) -> &[Row] {
352        self.deliveries.rows.as_slice()
353    }
354}
355
356impl<DB, Row> IntoIterator for RowBatch<DB, Row>
357where
358    DB: QueueDatabase,
359    Row: Events<DB> + QueueRow<Lane = RowLane>,
360{
361    type Item = InboxDelivery<DB, Row>;
362    type IntoIter = RowDeliveries<DB, Row>;
363
364    fn into_iter(self) -> Self::IntoIter {
365        self.deliveries
366    }
367}
368
369impl<DB, Row> Iterator for RowDeliveries<DB, Row>
370where
371    DB: QueueDatabase,
372    Row: Events<DB> + QueueRow<Lane = RowLane>,
373{
374    type Item = InboxDelivery<DB, Row>;
375
376    fn next(&mut self) -> Option<Self::Item> {
377        if self.rows.len() == 0 {
378            return self.gone.next();
379        }
380        let hold = match &mut self.claims {
381            Claims::Locked(batch) => Hold::Batch(Arc::clone(batch)),
382            Claims::Leased(book, slots) => Hold::Lease {
383                book: *book,
384                slot: slots.next()?,
385                tx: None,
386            },
387            Claims::Advised(holds) => Hold::Advisory(holds.next()?),
388        };
389        // The batch built one claim per row, so a row is here whenever its claim is.
390        let row = self.rows.next()?;
391        let headers = self.headers.next().unwrap_or_default();
392        let delivery = InboxDelivery::of_batch(row, headers, hold, self.queue, self.pool);
393        #[cfg(feature = "testing")]
394        let delivery = delivery.on(&self.shared);
395        Some(delivery)
396    }
397
398    fn size_hint(&self) -> (usize, Option<usize>) {
399        let len = self.rows.len() + self.gone.len();
400        (len, Some(len))
401    }
402}
403
404impl<DB, Row> ExactSizeIterator for RowDeliveries<DB, Row>
405where
406    DB: QueueDatabase,
407    Row: Events<DB> + QueueRow<Lane = RowLane>,
408{
409}
410
411impl<DB, Row> Drop for RowDeliveries<DB, Row>
412where
413    DB: QueueDatabase,
414    Row: Events<DB> + QueueRow<Lane = RowLane>,
415{
416    fn drop(&mut self) {
417        // Each delivery not yet taken returns its row as it drops: the batch transaction's last
418        // one rolls back, a lease is released, a session closes, and the harness counts it.
419        self.for_each(drop);
420    }
421}
422
423impl<DB, Row> fmt::Debug for RowBatch<DB, Row>
424where
425    DB: QueueDatabase,
426    Row: Events<DB> + QueueRow<Lane = RowLane>,
427{
428    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
429        f.debug_struct("RowBatch")
430            .field("subscription", &self.deliveries.queue.name)
431            .field("rows", &self.deliveries.rows.len())
432            .field("gone", &self.deliveries.gone.len())
433            .finish_non_exhaustive()
434    }
435}
436
437impl<DB, Row> fmt::Debug for RowDeliveries<DB, Row>
438where
439    DB: QueueDatabase,
440    Row: Events<DB> + QueueRow<Lane = RowLane>,
441{
442    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
443        f.debug_struct("RowDeliveries")
444            .field("subscription", &self.queue.name)
445            .field("left", &self.len())
446            .finish_non_exhaustive()
447    }
448}
ruststream_sqlx/inbox/batch.rs

ruststream_sqlx/inbox/
batch.rs
