s2_sdk/
batching.rs

1//! Utilities for batching [AppendRecord]s.
2
3use std::{
4    pin::Pin,
5    task::{Context, Poll},
6    time::Duration,
7};
8
9use futures_core::Stream;
10use futures_util::{StreamExt, stream};
11use s2_common::{
12    caps::RECORD_BATCH_MAX,
13    read_extent::CountOrBytes,
14    record::{Metered, MeteredSize},
15};
16use tokio::time::Instant;
17
18use crate::types::{AppendInput, AppendRecord, AppendRecordBatch, FencingToken, ValidationError};
19
20const RECORD_BATCH_MIN: CountOrBytes = CountOrBytes { count: 1, bytes: 8 };
21
22#[derive(Debug, Clone)]
23/// Limits for batching [`AppendRecord`]s.
24pub struct BatchLimits {
25    max_batch_bytes: usize,
26    max_batch_records: usize,
27}
28
29impl Default for BatchLimits {
30    fn default() -> Self {
31        Self {
32            max_batch_bytes: RECORD_BATCH_MAX.bytes,
33            max_batch_records: RECORD_BATCH_MAX.count,
34        }
35    }
36}
37
38impl BatchLimits {
39    /// Create new [`BatchLimits`] with default settings.
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    /// Set the maximum metered bytes per batch.
45    ///
46    /// **Note:** It must be at least `8B` and must not exceed `1MiB`.
47    ///
48    /// Defaults to `1MiB`.
49    pub fn with_max_batch_bytes(self, max_batch_bytes: usize) -> Result<Self, ValidationError> {
50        if max_batch_bytes < RECORD_BATCH_MIN.bytes {
51            return Err(ValidationError(format!(
52                "max_batch_bytes ({max_batch_bytes}) must be at least {}",
53                RECORD_BATCH_MIN.bytes
54            )));
55        }
56        if max_batch_bytes > RECORD_BATCH_MAX.bytes {
57            return Err(ValidationError(format!(
58                "max_batch_bytes ({max_batch_bytes}) must not exceed {}",
59                RECORD_BATCH_MAX.bytes
60            )));
61        }
62        Ok(Self {
63            max_batch_bytes,
64            ..self
65        })
66    }
67
68    /// Set the maximum number of records per batch.
69    ///
70    /// **Note:** It must be at least `1` and must not exceed `1000`.
71    ///
72    /// Defaults to `1000`.
73    pub fn with_max_batch_records(self, max_batch_records: usize) -> Result<Self, ValidationError> {
74        if max_batch_records < RECORD_BATCH_MIN.count {
75            return Err(ValidationError(format!(
76                "max_batch_records ({max_batch_records}) must be at least {}",
77                RECORD_BATCH_MIN.count
78            )));
79        }
80        if max_batch_records > RECORD_BATCH_MAX.count {
81            return Err(ValidationError(format!(
82                "max_batch_records ({max_batch_records}) must not exceed {}",
83                RECORD_BATCH_MAX.count
84            )));
85        }
86        Ok(Self {
87            max_batch_records,
88            ..self
89        })
90    }
91}
92
93#[derive(Debug, Clone)]
94/// Configuration for batching [`AppendRecord`]s.
95pub struct BatchingConfig {
96    linger: Duration,
97    limits: BatchLimits,
98}
99
100impl Default for BatchingConfig {
101    fn default() -> Self {
102        Self {
103            linger: Duration::from_millis(5),
104            limits: BatchLimits::default(),
105        }
106    }
107}
108
109impl BatchingConfig {
110    /// Create a new [`BatchingConfig`] with default settings.
111    pub fn new() -> Self {
112        Self::default()
113    }
114
115    /// Set the duration for how long to wait for more records before flushing a batch.
116    ///
117    /// Defaults to `5ms`.
118    pub fn with_linger(self, linger: Duration) -> Self {
119        Self { linger, ..self }
120    }
121
122    /// Set the batch limits.
123    pub fn with_limits(self, limits: BatchLimits) -> Self {
124        Self { limits, ..self }
125    }
126}
127
128/// A [`Stream`] that batches [`AppendRecord`]s into [`AppendInput`]s.
129pub struct AppendInputs {
130    pub(crate) batches: AppendRecordBatches,
131    pub(crate) fencing_token: Option<FencingToken>,
132    pub(crate) match_seq_num: Option<u64>,
133}
134
135impl AppendInputs {
136    /// Create a new [`AppendInputs`] from pre-batched records.
137    pub fn new(batches: AppendRecordBatches) -> Self {
138        Self {
139            batches,
140            fencing_token: None,
141            match_seq_num: None,
142        }
143    }
144
145    /// Set the fencing token for all [`AppendInput`]s.
146    pub fn with_fencing_token(self, fencing_token: FencingToken) -> Self {
147        Self {
148            fencing_token: Some(fencing_token),
149            ..self
150        }
151    }
152
153    /// Set the match sequence number for the initial [`AppendInput`]. It will be auto-incremented
154    /// for the subsequent ones.
155    pub fn with_match_seq_num(self, seq_num: u64) -> Self {
156        Self {
157            match_seq_num: Some(seq_num),
158            ..self
159        }
160    }
161}
162
163impl Stream for AppendInputs {
164    type Item = Result<AppendInput, ValidationError>;
165
166    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
167        match self.batches.poll_next_unpin(cx) {
168            Poll::Ready(Some(Ok(batch))) => {
169                let match_seq_num = self.match_seq_num;
170                if let Some(seq_num) = self.match_seq_num.as_mut() {
171                    *seq_num += batch.len() as u64;
172                }
173                Poll::Ready(Some(Ok(AppendInput {
174                    records: batch,
175                    match_seq_num,
176                    fencing_token: self.fencing_token.clone(),
177                })))
178            }
179            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
180            Poll::Ready(None) => Poll::Ready(None),
181            Poll::Pending => Poll::Pending,
182        }
183    }
184}
185
186/// A [`Stream`] that batches [`AppendRecord`]s into [`AppendRecordBatch`]es.
187pub struct AppendRecordBatches {
188    inner: Pin<Box<dyn Stream<Item = Result<AppendRecordBatch, ValidationError>> + Send>>,
189}
190
191impl AppendRecordBatches {
192    /// Create a new [`AppendRecordBatches`] from a record stream and config.
193    pub fn from_stream(
194        records: impl Stream<Item = impl Into<AppendRecord> + Send> + Send + Unpin + 'static,
195        config: BatchingConfig,
196    ) -> Self {
197        Self {
198            inner: Box::pin(append_record_batches(records, config)),
199        }
200    }
201
202    /// Eagerly batch in-memory records with the given limits.
203    ///
204    /// The returned value is a stream of validated batches and can be collected
205    /// into a `Vec<AppendRecordBatch>` when eager results are needed.
206    ///
207    /// ```rust
208    /// # use s2_sdk::batching::{AppendRecordBatches, BatchLimits};
209    /// # use s2_sdk::types::AppendRecord;
210    /// let records = [AppendRecord::new("one")?, AppendRecord::new("two")?];
211    /// let batches = AppendRecordBatches::from_iter(records, BatchLimits::new())?;
212    /// # let _ = batches;
213    /// # Ok::<(), s2_sdk::types::ValidationError>(())
214    /// ```
215    pub fn from_iter(
216        records: impl IntoIterator<Item = impl Into<AppendRecord>>,
217        limits: BatchLimits,
218    ) -> Result<Self, ValidationError> {
219        let mut batches = Vec::new();
220        let mut batch = Metered::with_capacity(limits.max_batch_records);
221
222        for item in records {
223            let record = Metered::from(item.into());
224            if record.metered_size() > limits.max_batch_bytes {
225                return Err(ValidationError(format!(
226                    "record size in metered bytes ({}) exceeds max_batch_bytes ({})",
227                    record.metered_size(),
228                    limits.max_batch_bytes
229                )));
230            }
231
232            if !batch.is_empty() && would_overflow_batch(&limits, &batch, &record) {
233                batches.push(AppendRecordBatch::from(std::mem::replace(
234                    &mut batch,
235                    Metered::with_capacity(limits.max_batch_records),
236                )));
237            }
238
239            batch.push(record);
240            if is_batch_full(&limits, &batch) {
241                batches.push(AppendRecordBatch::from(std::mem::replace(
242                    &mut batch,
243                    Metered::with_capacity(limits.max_batch_records),
244                )));
245            }
246        }
247
248        if !batch.is_empty() {
249            batches.push(batch.into());
250        }
251
252        Ok(Self {
253            inner: Box::pin(stream::iter(batches.into_iter().map(Ok))),
254        })
255    }
256}
257
258impl Stream for AppendRecordBatches {
259    type Item = Result<AppendRecordBatch, ValidationError>;
260
261    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
262        self.inner.as_mut().poll_next(cx)
263    }
264}
265
266fn is_batch_full(limits: &BatchLimits, batch: &Metered<Vec<AppendRecord>>) -> bool {
267    batch.len() >= limits.max_batch_records || batch.metered_size() >= limits.max_batch_bytes
268}
269
270fn would_overflow_batch(
271    limits: &BatchLimits,
272    batch: &Metered<Vec<AppendRecord>>,
273    record: &Metered<AppendRecord>,
274) -> bool {
275    batch.len() + 1 > limits.max_batch_records
276        || batch.metered_size() + record.metered_size() > limits.max_batch_bytes
277}
278
279fn append_record_batches(
280    mut records: impl Stream<Item = impl Into<AppendRecord> + Send> + Send + Unpin + 'static,
281    config: BatchingConfig,
282) -> impl Stream<Item = Result<AppendRecordBatch, ValidationError>> + Send + 'static {
283    async_stream::try_stream! {
284        let mut batch = Metered::with_capacity(config.limits.max_batch_records);
285        let mut overflowed_record: Option<Metered<AppendRecord>> = None;
286
287        let linger_deadline = tokio::time::sleep(config.linger);
288        tokio::pin!(linger_deadline);
289
290        'outer: loop {
291            let first_record = match overflowed_record.take() {
292                Some(pair) => pair,
293                None => match records.next().await {
294                    Some(item) => Metered::from(item.into()),
295                    None => break,
296                },
297            };
298
299            if first_record.metered_size() > config.limits.max_batch_bytes {
300                Err(ValidationError(format!(
301                    "record size in metered bytes ({}) exceeds max_batch_bytes ({})",
302                    first_record.metered_size(),
303                    config.limits.max_batch_bytes
304                )))?;
305            }
306            batch.push(first_record);
307
308            while !is_batch_full(&config.limits, &batch) && overflowed_record.is_none() {
309                if batch.len() == 1 {
310                    linger_deadline
311                        .as_mut()
312                        .reset(Instant::now() + config.linger);
313                }
314
315                tokio::select! {
316                    next_record = records.next() => {
317                        match next_record {
318                            Some(record) => {
319                                let record = Metered::from(record.into());
320                                if would_overflow_batch(&config.limits, &batch, &record) {
321                                    overflowed_record = Some(record);
322                                } else {
323                                    batch.push(record);
324                                }
325                            }
326                            None => {
327                                yield AppendRecordBatch::from(std::mem::replace(
328                                    &mut batch,
329                                    Metered::with_capacity(config.limits.max_batch_records),
330                                ));
331                                break 'outer;
332                            }
333                        }
334                    },
335                    _ = &mut linger_deadline, if !batch.is_empty() => {
336                        break;
337                    }
338                };
339            }
340
341            yield AppendRecordBatch::from(std::mem::replace(
342                &mut batch,
343                Metered::with_capacity(config.limits.max_batch_records),
344            ));
345        }
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use assert_matches::assert_matches;
352    use futures_util::TryStreamExt;
353
354    use super::*;
355    use crate::types::MeteredBytes as _;
356
357    #[tokio::test]
358    async fn batches_should_be_empty_when_record_stream_is_empty() {
359        let batches: Vec<_> = AppendRecordBatches::from_stream(
360            futures_util::stream::iter::<Vec<AppendRecord>>(vec![]),
361            BatchingConfig::default(),
362        )
363        .collect()
364        .await;
365        assert_eq!(batches.len(), 0);
366    }
367
368    #[tokio::test]
369    async fn batches_respect_count_limit() -> Result<(), ValidationError> {
370        let records: Vec<_> = (0..10)
371            .map(|i| AppendRecord::new(format!("record{i}")))
372            .collect::<Result<_, _>>()?;
373        let config = BatchingConfig::default()
374            .with_limits(BatchLimits::default().with_max_batch_records(3)?);
375        let batches: Vec<_> =
376            AppendRecordBatches::from_stream(futures_util::stream::iter(records), config)
377                .try_collect()
378                .await?;
379
380        assert_eq!(batches.len(), 4);
381        assert_eq!(batches[0].len(), 3);
382        assert_eq!(batches[1].len(), 3);
383        assert_eq!(batches[2].len(), 3);
384        assert_eq!(batches[3].len(), 1);
385
386        Ok(())
387    }
388
389    #[tokio::test]
390    async fn batches_respect_bytes_limit() -> Result<(), ValidationError> {
391        let records: Vec<_> = (0..10)
392            .map(|i| AppendRecord::new(format!("record{i}")))
393            .collect::<Result<_, _>>()?;
394        let single_record_bytes = records[0].metered_bytes();
395        let max_batch_bytes = single_record_bytes * 3;
396
397        let config = BatchingConfig::default()
398            .with_limits(BatchLimits::default().with_max_batch_bytes(max_batch_bytes)?);
399        let batches: Vec<_> =
400            AppendRecordBatches::from_stream(futures_util::stream::iter(records), config)
401                .try_collect()
402                .await?;
403
404        assert_eq!(batches.len(), 4);
405        assert_eq!(batches[0].metered_bytes(), max_batch_bytes);
406        assert_eq!(batches[1].metered_bytes(), max_batch_bytes);
407        assert_eq!(batches[2].metered_bytes(), max_batch_bytes);
408        assert_eq!(batches[3].metered_bytes(), single_record_bytes);
409
410        Ok(())
411    }
412
413    #[tokio::test(start_paused = true)]
414    async fn batches_flush_after_linger_when_stream_remains_open() -> Result<(), ValidationError> {
415        let records: Vec<_> = (0..2)
416            .map(|i| AppendRecord::new(format!("record{i}")))
417            .collect::<Result<_, _>>()?;
418        let records = futures_util::stream::iter(records).chain(futures_util::stream::pending());
419        let config = BatchingConfig::default().with_linger(Duration::from_millis(5));
420        let mut batches = AppendRecordBatches::from_stream(records, config);
421        let next_batch = tokio::spawn(async move { batches.next().await });
422
423        tokio::task::yield_now().await;
424        tokio::time::advance(Duration::from_millis(6)).await;
425
426        let batch = next_batch.await.unwrap().unwrap()?;
427        assert_eq!(batch.len(), 2);
428        Ok(())
429    }
430
431    #[tokio::test]
432    async fn batching_should_error_when_it_sees_oversized_record() -> Result<(), ValidationError> {
433        let record = AppendRecord::new("hello-world")?;
434        let record_bytes = record.metered_bytes();
435        let max_batch_bytes = 10;
436
437        let config = BatchingConfig::default()
438            .with_limits(BatchLimits::default().with_max_batch_bytes(max_batch_bytes)?);
439        let results: Vec<_> =
440            AppendRecordBatches::from_stream(futures_util::stream::iter(vec![record]), config)
441                .collect()
442                .await;
443
444        assert_eq!(results.len(), 1);
445        assert_matches!(&results[0], Err(err) => {
446            assert_eq!(
447                err.to_string(),
448                format!("record size in metered bytes ({record_bytes}) exceeds max_batch_bytes ({max_batch_bytes})")
449            );
450        });
451
452        Ok(())
453    }
454
455    #[tokio::test]
456    async fn eager_batches_should_be_empty_when_record_iter_is_empty() {
457        let batches =
458            AppendRecordBatches::from_iter(std::iter::empty::<AppendRecord>(), BatchLimits::new())
459                .unwrap()
460                .try_collect::<Vec<_>>()
461                .await
462                .unwrap();
463        assert!(batches.is_empty());
464    }
465
466    #[tokio::test]
467    async fn eager_batches_respect_count_limit() -> Result<(), ValidationError> {
468        let records: Vec<_> = (0..10)
469            .map(|i| AppendRecord::new(format!("record{i}")))
470            .collect::<Result<_, _>>()?;
471        let limits = BatchLimits::new().with_max_batch_records(3)?;
472        let batches = AppendRecordBatches::from_iter(records, limits)?
473            .try_collect::<Vec<_>>()
474            .await?;
475
476        assert_eq!(batches.len(), 4);
477        assert_eq!(batches[0].len(), 3);
478        assert_eq!(batches[1].len(), 3);
479        assert_eq!(batches[2].len(), 3);
480        assert_eq!(batches[3].len(), 1);
481        Ok(())
482    }
483
484    #[tokio::test]
485    async fn eager_batches_respect_bytes_limit() -> Result<(), ValidationError> {
486        let records: Vec<_> = (0..10)
487            .map(|i| AppendRecord::new(format!("record{i}")))
488            .collect::<Result<_, _>>()?;
489        let single_record_bytes = records[0].metered_bytes();
490        let max_batch_bytes = single_record_bytes * 3;
491        let limits = BatchLimits::new().with_max_batch_bytes(max_batch_bytes)?;
492        let batches = AppendRecordBatches::from_iter(records, limits)?
493            .try_collect::<Vec<_>>()
494            .await?;
495
496        assert_eq!(batches.len(), 4);
497        assert_eq!(batches[0].metered_bytes(), max_batch_bytes);
498        assert_eq!(batches[1].metered_bytes(), max_batch_bytes);
499        assert_eq!(batches[2].metered_bytes(), max_batch_bytes);
500        assert_eq!(batches[3].metered_bytes(), single_record_bytes);
501        Ok(())
502    }
503
504    #[tokio::test]
505    async fn eager_batching_should_error_when_record_is_oversized() -> Result<(), ValidationError> {
506        let record = AppendRecord::new("hello-world")?;
507        let record_bytes = record.metered_bytes();
508        let max_batch_bytes = 10;
509        let limits = BatchLimits::new().with_max_batch_bytes(max_batch_bytes)?;
510        let err = AppendRecordBatches::from_iter([record], limits)
511            .err()
512            .unwrap();
513
514        assert_eq!(
515            err.to_string(),
516            format!(
517                "record size in metered bytes ({record_bytes}) exceeds max_batch_bytes ({max_batch_bytes})"
518            )
519        );
520        Ok(())
521    }
522}
s2_sdk/batching.rs

s2_sdk/
batching.rs
