batch_aint_one/
batcher.rs

1use std::{fmt::Debug, sync::Arc};
2
3use bon::bon;
4use tokio::sync::{mpsc, oneshot};
5use tracing::{Level, Span, span};
6
7use crate::{
8    batch::BatchItem,
9    error::BatchResult,
10    limits::Limits,
11    metrics::{MetricsRecorderFactory, NoopMetricsRecorder},
12    policies::BatchingPolicy,
13    processor::Processor,
14    worker::{Worker, WorkerDropGuard, WorkerHandle},
15};
16
17/// Groups items to be processed in batches.
18///
19/// Takes inputs one at a time and sends them to a background worker task which groups them into
20/// batches according to the specified [`BatchingPolicy`] and [`Limits`], and processes them using
21/// the provided [`Processor`].
22///
23/// Cheap to clone. Cloned instances share the same background worker task.
24///
25/// ## Drop
26///
27/// When the last instance of a `Batcher` is dropped, the worker task will be aborted (ungracefully
28/// shut down). Any callers still awaiting [`Batcher::add()`] will receive a
29/// [`BatchError::Rx`](crate::BatchError::Rx) error.
30///
31/// If you want to shut down the worker gracefully, call [`WorkerHandle::shut_down()`].
32#[derive(Debug)]
33pub struct Batcher<P: Processor> {
34    name: String,
35    worker: Arc<WorkerHandle>,
36    worker_guard: Arc<WorkerDropGuard>,
37    item_tx: mpsc::Sender<BatchItem<P>>,
38}
39
40#[bon]
41impl<P: Processor> Batcher<P> {
42    /// Create a new batcher.
43    ///
44    /// # Notes
45    ///
46    /// If `batching_policy` is `Balanced { min_size_hint }` where `min_size_hint` is greater than
47    /// `limits.max_batch_size`, the `min_size_hint` will be clamped to `max_batch_size`.
48    #[builder]
49    pub fn new(
50        name: impl Into<String>,
51        processor: P,
52        limits: Limits,
53        batching_policy: BatchingPolicy,
54        metrics: Option<Box<dyn MetricsRecorderFactory>>,
55    ) -> Self {
56        let name = name.into();
57        let metrics_recorder = metrics
58            .map(|f| f.create_recorder(&name))
59            .unwrap_or_else(|| Box::new(NoopMetricsRecorder));
60
61        let batching_policy = batching_policy.normalise(limits);
62
63        let (handle, worker_guard, item_tx) = Worker::spawn(
64            name.clone(),
65            processor,
66            limits,
67            batching_policy,
68            metrics_recorder,
69        );
70
71        Self {
72            name,
73            worker: Arc::new(handle),
74            worker_guard: Arc::new(worker_guard),
75            item_tx,
76        }
77    }
78
79    /// Add an item to be batched and processed, and await the result.
80    pub async fn add(&self, key: P::Key, input: P::Input) -> BatchResult<P::Output, P::Error> {
81        // Record the span ID so we can link the shared processing span.
82        let requesting_span = Span::current().clone();
83
84        let (tx, rx) = oneshot::channel();
85        self.item_tx
86            .send(BatchItem {
87                key,
88                input,
89                submitted_at: tokio::time::Instant::now(),
90                received_at: None,
91                tx,
92                requesting_span,
93            })
94            .await?;
95
96        let (output, batch_span) = rx.await?;
97
98        {
99            let link_back_span = span!(Level::INFO, "batch finished");
100            if let Some(span) = batch_span {
101                // WARNING: It's very important that we don't drop the span until _after_
102                // follows_from().
103                //
104                // If we did e.g. `.follows_from(span)` then the span would get converted into an ID
105                // and dropped. Any attempt to look up the span by ID _inside_ follows_from() would
106                // then panic, because the span will have been closed and no longer exist.
107                //
108                // Don't ask me how long this took me to debug.
109                link_back_span.follows_from(&span);
110                link_back_span.in_scope(|| {
111                    // Do nothing. This span is just here to work around a Honeycomb limitation:
112                    //
113                    // If the batch span is linked to a parent span like so:
114                    //
115                    // parent_span_1 <-link- batch_span
116                    //
117                    // then in Honeycomb, the link is only shown on the batch span. It it not possible
118                    // to click through to the batch span from the parent.
119                    //
120                    // So, here we link back to the batch to make this easier.
121                });
122            }
123        }
124        output
125    }
126
127    /// Get a handle to the worker.
128    pub fn worker_handle(&self) -> Arc<WorkerHandle> {
129        Arc::clone(&self.worker)
130    }
131}
132
133impl<P: Processor> Clone for Batcher<P> {
134    fn clone(&self) -> Self {
135        Self {
136            name: self.name.clone(),
137            worker: self.worker.clone(),
138            worker_guard: self.worker_guard.clone(),
139            item_tx: self.item_tx.clone(),
140        }
141    }
142}
batch_aint_one/batcher.rs

batch_aint_one/
batcher.rs
