headless_lms_credit_registration/use_cases/
batch_flow.rs

1//! The loop every "claim rows, send one batch, write one answer per row" flow shares: the claim's
2//! transaction, splitting a batch the study registry refused as malformed, the halves held back and
3//! released when a shutdown or an error stops the split, the refusals, one row's failed write not
4//! costing the rest theirs, the moved-on rows and the batch summary. What goes on the wire, request
5//! ids, the limiter and the breakers are the registry adapter's.
6
7use chrono::TimeDelta;
8use headless_lms_utils::prelude::Utc;
9use sqlx::{Connection, PgConnection, PgPool};
10use std::time::Instant;
11use tokio_util::sync::CancellationToken;
12use uuid::Uuid;
13
14use crate::error::CreditRegistrationResult;
15use crate::error_reports::ErrorReporter;
16use crate::phase::CreditRegistrationPhase;
17use crate::registry::{
18    BatchEntry, BatchOptions, BatchReply, BatchRequest, ExchangeAudit, StudyRegistry,
19};
20use crate::workflow::{Applied, Claimed, Counts, RefusalPolicy, write_decision};
21use headless_lms_models::credit_registrations::RegistrationScope;
22
23/// What a [`RegistryBatchFlow`] may touch.
24pub(crate) struct BatchFlowContext<'a> {
25    pub pool: &'a PgPool,
26    pub scope: &'a RegistrationScope,
27    pub phase: CreditRegistrationPhase,
28    pub errors: ErrorReporter<'a>,
29    /// The worker's SIGTERM; `None` for a run no signal can stop, such as an on-demand one.
30    pub shutdown: Option<&'a CancellationToken>,
31    /// The longest the iteration's study registry calls may take together, which a lease keeping
32    /// its rows from a concurrent iteration has to outlast.
33    pub study_registry_wait: TimeDelta,
34}
35
36/// A flow whose iteration is "claim rows, decide in one transaction what may be asked, send one
37/// batch, write one answer per row": `import`, `resolve-enrolments`, and each of `verify`'s two
38/// flows. [`run_registry_batch_flow`] is the loop they share. Implementations hold no state.
39pub(super) trait RegistryBatchFlow {
40    /// What the claim read for each row alongside it, which its answer is applied with.
41    type Extra;
42    /// What each row asks, which also names the operation it is asked under.
43    type Request: BatchRequest;
44
45    /// The iteration's error when every item came back unavailable.
46    const ALL_UNAVAILABLE_ERROR: &'static str;
47    /// What the rows get when the study registry refused the whole request, which also decides
48    /// whether a batch refused as malformed is split.
49    const REFUSAL: RefusalPolicy<Self::Extra>;
50
51    /// Claims at most `limit` rows and decides what may be asked about them. Whatever has to be
52    /// true before the request leaves is written here, in the caller's transaction.
53    async fn claim(
54        ctx: &BatchFlowContext<'_>,
55        conn: &mut PgConnection,
56        limit: usize,
57    ) -> CreditRegistrationResult<Prepared<Self::Extra, Self::Request>>;
58
59    /// Decides what one answer, or the absence of one, does to its row and writes it, with the
60    /// exchange behind it. May read and write in transactions of its own on the way.
61    async fn apply_answer(
62        conn: &mut PgConnection,
63        row: &Claimed<Self::Extra>,
64        answer: Option<&<Self::Request as BatchRequest>::Answer>,
65        audit: &ExchangeAudit,
66    ) -> CreditRegistrationResult<Applied>;
67
68    /// Called before each send after a split, with every row the split still holds, so the ones
69    /// waiting their turn are not taken for a worker that died mid-call.
70    async fn keep_in_flight(
71        _conn: &mut PgConnection,
72        _rows: &[&Claimed<Self::Extra>],
73    ) -> CreditRegistrationResult<()> {
74        Ok(())
75    }
76
77    /// Called with every row a split still holds unsent when a shutdown or an error stops the
78    /// iteration, each of which would otherwise be condemned as a lost submission.
79    async fn release_unsent(
80        _conn: &mut PgConnection,
81        _rows: &[&Claimed<Self::Extra>],
82    ) -> CreditRegistrationResult<()> {
83        Ok(())
84    }
85}
86
87/// What a batch flow's claim settled before anything is sent: the rows to send, each with what it
88/// asks, and the rows it already wrote a decision for. Private, so the batch and the preflight's
89/// counts only change together through [`Self::send`] and the `record_` methods.
90pub(super) struct Prepared<Extra, Request> {
91    sendable: Vec<BatchEntry<Claimed<Extra>, Request>>,
92    decided: Counts,
93}
94
95impl<Extra, Request> Prepared<Extra, Request> {
96    pub(super) fn new() -> Self {
97        Self {
98            sendable: Vec::new(),
99            decided: Counts::default(),
100        }
101    }
102
103    /// Puts `row` in the batch, asking `request`.
104    pub(super) fn send(&mut self, row: Claimed<Extra>, request: Request) {
105        self.sendable.push(BatchEntry { row, request });
106    }
107
108    /// A row the claim failed without asking.
109    pub(super) fn record_failed(&mut self) {
110        self.decided.record_decided(true);
111    }
112
113    /// A row the claim wrote a guarded decision for, which may have found it moved on.
114    pub(super) fn record_applied(&mut self, registration_id: Uuid, applied: Applied) {
115        self.decided.record_applied(registration_id, applied);
116    }
117
118    pub(super) fn sendable(&self) -> &[BatchEntry<Claimed<Extra>, Request>] {
119        &self.sendable
120    }
121
122    /// The batch, and the counts the claim's own decisions start the iteration with.
123    fn into_parts(self) -> (Vec<BatchEntry<Claimed<Extra>, Request>>, Counts) {
124        (self.sendable, self.decided)
125    }
126}
127
128/// One request's rows, still to be sent.
129struct PendingBatch<Extra, R> {
130    entries: Vec<BatchEntry<Claimed<Extra>, R>>,
131    /// Split off a batch refused as malformed.
132    is_resent_half: bool,
133}
134
135/// What the iteration put on the wire, for the batch summary.
136#[derive(Default)]
137struct SentTally {
138    requests: i32,
139    items: i32,
140}
141
142/// Runs one [`RegistryBatchFlow`] through the iteration's registry, and logs the batch summary.
143///
144/// `processed` counts the rows this flow wrote a decision for: the claim's included, the ones
145/// another writer had moved on before the answer could be applied excluded.
146pub(super) async fn run_registry_batch_flow<F: RegistryBatchFlow, R: StudyRegistry>(
147    ctx: &BatchFlowContext<'_>,
148    registry: &mut R,
149) -> CreditRegistrationResult<Counts> {
150    let operation = <F::Request as BatchRequest>::OPERATION;
151    let limit = registry.allowance(operation);
152    if limit == 0 {
153        return Ok(Counts::default());
154    }
155    let started_at = Instant::now();
156    let mut conn = ctx.pool.acquire().await?;
157    let mut tx = conn.begin().await?;
158    let prepared = F::claim(ctx, &mut tx, limit).await?;
159    tx.commit().await?;
160    // Held only for the claim; the study registry call below can pin it for the whole request
161    // timeout.
162    drop(conn);
163
164    let (entries, mut counts) = prepared.into_parts();
165    let claimed = counts.processed_count()
166        + counts.moved_on_count()
167        + i32::try_from(entries.len()).unwrap_or(i32::MAX);
168    let mut sent = SentTally::default();
169    // The halves a split holds back wait in whatever state the claim left them, which for import is
170    // `submitting`: no phase claims that, so none of them can be sent twice meanwhile. Each answered
171    // half is written before the next one is sent.
172    let mut batches = Vec::new();
173    if !entries.is_empty() {
174        batches.push(PendingBatch {
175            entries,
176            is_resent_half: false,
177        });
178    }
179    let drained = send_batches::<F, R>(ctx, registry, &mut batches, &mut counts, &mut sent).await;
180    // Whatever is left was never sent: a shutdown or an error stopped the split.
181    if !batches.is_empty() {
182        let released = release_unsent::<F>(ctx, &batches).await;
183        match (&drained, released) {
184            (Ok(()), released) => released?,
185            (Err(_), Ok(())) => {}
186            (Err(_), Err(error)) => error!(
187                phase = ctx.phase.as_str(),
188                error = %error,
189                "Could not release the unsent halves of a split batch"
190            ),
191        }
192    }
193    drained?;
194    if claimed > 0 {
195        let written = counts.processed_count();
196        let waiting = counts.waiting_count();
197        let failed = counts.failed_count();
198        let moved_on = counts.moved_on_count();
199        let SentTally {
200            requests: requests_sent,
201            items: items_sent,
202        } = sent;
203        let duration_ms = started_at.elapsed().as_millis() as u64;
204        info!(
205            phase = ctx.phase.as_str(),
206            ?operation,
207            claimed,
208            requests_sent,
209            items_sent,
210            written,
211            waiting,
212            failed,
213            moved_on,
214            duration_ms,
215            "claimed {claimed}, sent {items_sent} in {requests_sent} requests, wrote {written} \
216             ({waiting} waiting, {failed} failed), moved on {moved_on}, took {duration_ms}ms"
217        );
218    }
219    Ok(counts)
220}
221
222/// Sends `batches` one after another, splitting the ones refused as malformed, and writes each
223/// row's answer. Stops early on shutdown or on an error, leaving what is still unsent in `batches`.
224async fn send_batches<F: RegistryBatchFlow, R: StudyRegistry>(
225    ctx: &BatchFlowContext<'_>,
226    registry: &mut R,
227    batches: &mut Vec<PendingBatch<F::Extra, F::Request>>,
228    counts: &mut Counts,
229    sent: &mut SentTally,
230) -> CreditRegistrationResult<()> {
231    let operation = <F::Request as BatchRequest>::OPERATION;
232    while let Some(batch) = batches.pop() {
233        if batch.is_resent_half {
234            let may_send = may_resend::<F>(ctx, &batch, batches).await;
235            if !matches!(may_send, Ok(true)) {
236                // Back on the stack, so it is released with the halves still waiting.
237                batches.push(batch);
238                return may_send.map(|_| ());
239            }
240        }
241        sent.requests += 1;
242        sent.items += i32::try_from(batch.entries.len()).unwrap_or(i32::MAX);
243        let options = BatchOptions {
244            may_split: F::REFUSAL.may_split(),
245            is_resent_half: batch.is_resent_half,
246            all_unavailable_error: F::ALL_UNAVAILABLE_ERROR,
247            registration_ids: batch
248                .entries
249                .iter()
250                .map(|entry| entry.row.claim.id())
251                .collect(),
252        };
253        match <F::Request as BatchRequest>::send(registry, batch.entries, options).await {
254            BatchReply::RefusedAsMalformed {
255                entries: mut first,
256                error,
257            } => {
258                warn!(
259                    batch_size = first.len(),
260                    error = error.message.as_str(),
261                    "The study registry refused a batch as a whole; splitting it to find the rows it refuses"
262                );
263                let second = first.split_off(first.len() / 2);
264                // The stack pops the first half first.
265                for entries in [second, first] {
266                    batches.push(PendingBatch {
267                        entries,
268                        is_resent_half: true,
269                    });
270                }
271            }
272            BatchReply::Refused {
273                rows,
274                error,
275                refused_for,
276            } => {
277                let mut conn = ctx.pool.acquire().await?;
278                for refused in &rows {
279                    let decision = F::REFUSAL.decision(
280                        &refused.row,
281                        operation,
282                        &error,
283                        refused_for,
284                        Utc::now(),
285                    );
286                    let written =
287                        write_decision(&mut conn, &refused.row.claim, decision, &refused.audit)
288                            .await;
289                    record_row_write(ctx, counts, refused.row.claim.id(), written).await?;
290                }
291            }
292            BatchReply::Answered(rows) => {
293                let mut conn = ctx.pool.acquire().await?;
294                for answered in &rows {
295                    let row = &answered.row;
296                    let written =
297                        F::apply_answer(&mut conn, row, answered.answer.as_ref(), &answered.audit)
298                            .await;
299                    record_row_write(ctx, counts, row.claim.id(), written).await?;
300                }
301            }
302        }
303    }
304    Ok(())
305}
306
307/// Before a resent half: `false` once a shutdown is under way, since each half can take the whole
308/// request timeout and sending the rest would outlast the termination grace period. Otherwise keeps
309/// `next` and every row still `waiting` in flight.
310async fn may_resend<F: RegistryBatchFlow>(
311    ctx: &BatchFlowContext<'_>,
312    next: &PendingBatch<F::Extra, F::Request>,
313    waiting: &[PendingBatch<F::Extra, F::Request>],
314) -> CreditRegistrationResult<bool> {
315    if ctx.shutdown.is_some_and(CancellationToken::is_cancelled) {
316        return Ok(false);
317    }
318    let rows: Vec<_> = std::iter::once(next)
319        .chain(waiting)
320        .flat_map(|pending| &pending.entries)
321        .map(|entry| &entry.row)
322        .collect();
323    let mut conn = ctx.pool.acquire().await?;
324    F::keep_in_flight(&mut conn, &rows).await?;
325    Ok(true)
326}
327
328async fn release_unsent<F: RegistryBatchFlow>(
329    ctx: &BatchFlowContext<'_>,
330    batches: &[PendingBatch<F::Extra, F::Request>],
331) -> CreditRegistrationResult<()> {
332    let rows: Vec<_> = batches
333        .iter()
334        .flat_map(|pending| &pending.entries)
335        .map(|entry| &entry.row)
336        .collect();
337    let mut conn = ctx.pool.acquire().await?;
338    F::release_unsent(&mut conn, &rows).await
339}
340
341/// Counts one row's write. A row whose write failed is reported and counted as failed while the
342/// batch's other answers are still written: dropping them would leave rows Sisu already answered for
343/// to be recovered as uncertain. Only a lost database connection stops the batch.
344async fn record_row_write(
345    ctx: &BatchFlowContext<'_>,
346    counts: &mut Counts,
347    registration_id: Uuid,
348    written: CreditRegistrationResult<Applied>,
349) -> CreditRegistrationResult<()> {
350    match written {
351        Ok(applied) => counts.record_applied(registration_id, applied),
352        Err(error) if error.is_db_disconnect() => return Err(error),
353        Err(error) => {
354            error!(
355                phase = ctx.phase.as_str(),
356                credit_registration_id = %registration_id,
357                error = %error,
358                "Could not write the study registry's answer for a credit registration; leaving it where it stands"
359            );
360            ctx.errors
361                .report(
362                    &error.cause_chain(),
363                    Some(format!("{error:?}")),
364                    serde_json::json!({ "credit_registration_id": registration_id }),
365                )
366                .await;
367            counts.record_decided(true);
368        }
369    }
370    Ok(())
371}
372
373#[cfg(test)]
374mod tests {
375    use headless_lms_models::credit_registrations::CreditRegistrationState as State;
376
377    use super::*;
378    use crate::test_fixtures::registration;
379    use crate::workflow::ClaimedRegistration;
380
381    fn claimed() -> Claimed<()> {
382        Claimed {
383            claim: ClaimedRegistration::left_in_place(registration(State::AwaitingVerification)),
384            extra: (),
385        }
386    }
387
388    #[test]
389    fn a_claim_keeps_its_own_decisions_apart_from_the_batch_it_sends() {
390        let (first, second) = (claimed(), claimed());
391        let ids = [first.claim.id(), second.claim.id()];
392        let mut prepared = Prepared::new();
393        prepared.send(first, "first");
394        prepared.record_failed();
395        prepared.send(second, "second");
396        prepared.record_applied(
397            Uuid::new_v4(),
398            Applied::Written {
399                is_failure: false,
400                is_waiting: false,
401            },
402        );
403        assert_eq!(prepared.sendable().len(), 2);
404        let (sendable, decided) = prepared.into_parts();
405        let sent: Vec<_> = sendable
406            .iter()
407            .map(|entry| (entry.row.claim.id(), entry.request))
408            .collect();
409        assert_eq!(sent, [(ids[0], "first"), (ids[1], "second")]);
410        assert_eq!(decided.processed_count(), 2);
411        assert_eq!(decided.failed_count(), 1);
412        assert_eq!(decided.moved_on_count(), 0);
413    }
414}
headless_lms_credit_registration/use_cases/batch_flow.rs

headless_lms_credit_registration/use_cases/
batch_flow.rs
