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::{MetricsRecorder, 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_recorder: Option<Arc<dyn MetricsRecorder>>,
55    ) -> Self {
56        let name = name.into();
57        let metrics_recorder = metrics_recorder.unwrap_or_else(|| Arc::new(NoopMetricsRecorder));
58
59        let batching_policy = batching_policy.normalise(limits);
60
61        let (handle, worker_guard, item_tx) = Worker::spawn(
62            name.clone(),
63            processor,
64            limits,
65            batching_policy,
66            metrics_recorder,
67        );
68
69        Self {
70            name,
71            worker: Arc::new(handle),
72            worker_guard: Arc::new(worker_guard),
73            item_tx,
74        }
75    }
76
77    /// Add an item to be batched and processed, and await the result.
78    pub async fn add(&self, key: P::Key, input: P::Input) -> BatchResult<P::Output, P::Error> {
79        // Record the span ID so we can link the shared processing span.
80        let requesting_span = Span::current().clone();
81
82        let (tx, rx) = oneshot::channel();
83        self.item_tx
84            .send(BatchItem {
85                key,
86                input,
87                submitted_at: tokio::time::Instant::now(),
88                tx,
89                requesting_span,
90            })
91            .await?;
92
93        let (output, batch_span) = rx.await?;
94
95        {
96            let link_back_span = span!(Level::INFO, "batch finished");
97            if let Some(span) = batch_span {
98                // WARNING: It's very important that we don't drop the span until _after_
99                // follows_from().
100                //
101                // If we did e.g. `.follows_from(span)` then the span would get converted into an ID
102                // and dropped. Any attempt to look up the span by ID _inside_ follows_from() would
103                // then panic, because the span will have been closed and no longer exist.
104                //
105                // Don't ask me how long this took me to debug.
106                link_back_span.follows_from(&span);
107                link_back_span.in_scope(|| {
108                    // Do nothing. This span is just here to work around a Honeycomb limitation:
109                    //
110                    // If the batch span is linked to a parent span like so:
111                    //
112                    // parent_span_1 <-link- batch_span
113                    //
114                    // then in Honeycomb, the link is only shown on the batch span. It it not possible
115                    // to click through to the batch span from the parent.
116                    //
117                    // So, here we link back to the batch to make this easier.
118                });
119            }
120        }
121        output
122    }
123
124    /// Get a handle to the worker.
125    pub fn worker_handle(&self) -> Arc<WorkerHandle> {
126        Arc::clone(&self.worker)
127    }
128}
129
130impl<P: Processor> Clone for Batcher<P> {
131    fn clone(&self) -> Self {
132        Self {
133            name: self.name.clone(),
134            worker: self.worker.clone(),
135            worker_guard: self.worker_guard.clone(),
136            item_tx: self.item_tx.clone(),
137        }
138    }
139}
batch_aint_one/batcher.rs

batch_aint_one/
batcher.rs
