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

ruststream_sqlx/inbox/
batch.rs
