krabka_remote_storage/index/
batch.rs

1//! Scans over a remote segment's `.log` bytes, which the offset index can only
2//! point near.
3//!
4//! Kafka offset and time indexes are sparse, so a byte range fetched from the
5//! object store usually starts before the record the caller asked for. These
6//! scans decode the batches in that range and pick the first one, or the first
7//! record, that satisfies the request.
8
9use {
10    super::{LogOffset, TimestampMs, corrupt_log},
11    crate::error::RemoteStorageError,
12    krabka_protocol::records::{Attributes, RecordBatch, TimestampType},
13    krabka_verified::timestamp::timestamp_record_time,
14};
15
16/// Decodes remote log batches and returns the earliest record at or after both
17/// `floor_offset` and `target_timestamp`.
18///
19/// Under `LogAppendTime`, every record
20/// uses the batch append stamp instead of producer base time and deltas.
21///
22/// `max_record_body` is Kafka trunk's `max.decompressed.message.bytes`, or
23/// `None` for no limit. Kafka's `RemoteLogManager.lookupTimestamp` decompresses
24/// a batch only when its max timestamp reaches `target_timestamp` and its last
25/// offset reaches `floor_offset`, and it reads that batch a record at a time up
26/// to the record it returns, so only those records are held to the limit. A
27/// batch the producer did not compress never is.
28///
29/// # Errors
30///
31/// Returns [`RemoteStorageError::Io`] when a batch does not decode, or when a
32/// record's offset or timestamp delta overflows its base, and
33/// [`RemoteStorageError::RecordTooLarge`] for a record above `max_record_body`.
34pub fn first_record_at_or_after_timestamp(
35    data: &[u8],
36    floor_offset: LogOffset,
37    target_timestamp: TimestampMs,
38    max_record_body: Option<usize>,
39) -> Result<Option<(LogOffset, TimestampMs)>, RemoteStorageError> {
40    let mut cur = data;
41    while !cur.is_empty() {
42        let batch = RecordBatch::decode(&mut cur).map_err(corrupt_log)?;
43        let limit = max_record_body.filter(|_| {
44            batch.attributes.compression() != Attributes::default().compression()
45                && batch.max_timestamp >= target_timestamp
46                && batch
47                    .base_offset
48                    .checked_add(i64::from(batch.last_offset_delta))
49                    .is_some_and(|last_offset| last_offset >= floor_offset)
50        });
51        let append_time = (batch.attributes.timestamp_type() == TimestampType::LogAppendTime)
52            .then_some(batch.max_timestamp);
53        for record in &batch.records {
54            if let Some(limit) = limit {
55                let size = record.body_len();
56                if size > limit {
57                    return Err(RemoteStorageError::RecordTooLarge { size, limit });
58                }
59            }
60            let offset = batch
61                .base_offset
62                .checked_add(i64::from(record.offset_delta))
63                .ok_or_else(|| corrupt_log("record offset overflow"))?;
64            if offset < floor_offset {
65                continue;
66            }
67            let timestamp =
68                timestamp_record_time(batch.base_timestamp, record.timestamp_delta, append_time)
69                    .ok_or_else(|| corrupt_log("record timestamp overflow"))?;
70            if timestamp >= target_timestamp {
71                return Ok(Some((offset, timestamp)));
72            }
73        }
74    }
75    Ok(None)
76}
77
78/// Decodes batches from `data` and returns the first one whose last offset is
79/// `>= floor`.
80///
81/// It skips the batches at the start of the returned byte range
82/// that the offset index pointed at but that do not cover the requested
83/// offset. Kafka offset indexes are sparse, so such batches occur.
84#[must_use]
85pub fn first_batch_at_or_after(data: &[u8], floor: LogOffset) -> Option<RecordBatch> {
86    let mut cur: &[u8] = data;
87    while !cur.is_empty() {
88        let Ok(batch) = RecordBatch::decode(&mut cur) else {
89            break;
90        };
91        let last_offset = batch
92            .base_offset
93            .checked_add(i64::from(batch.last_offset_delta))?;
94        if last_offset >= floor {
95            return Some(batch);
96        }
97    }
98    None
99}
100
101#[cfg(test)]
102mod tests {
103    use {
104        super::*,
105        assert2::assert,
106        bytes::{Bytes, BytesMut},
107        krabka_protocol::records::Record,
108    };
109
110    fn test_batch_at(base_offset: i64, record_count: i32, value_byte: u8) -> RecordBatch {
111        let mut batch = RecordBatch {
112            base_offset,
113            last_offset_delta: record_count - 1,
114            ..RecordBatch::default()
115        };
116        for offset_delta in 0..record_count {
117            batch.records.push(Record {
118                offset_delta,
119                value: Some(Bytes::from(vec![value_byte; 4])),
120                ..Default::default()
121            });
122        }
123        batch
124    }
125
126    krabka_macros::timestamped_batch!(timestamped_batch_at);
127
128    fn encoded(batches: &[RecordBatch]) -> Bytes {
129        let mut buf = BytesMut::new();
130        for batch in batches {
131            batch.encode(&mut buf).unwrap();
132        }
133        buf.freeze()
134    }
135
136    #[yare::parameterized(
137        skips_first_batch = { 10, Some(10) },
138        below_all_batches = { 0, Some(0) },
139        above_all_batches = { 1_000, None },
140    )]
141    fn first_batch_at_or_after_decodes_and_skips(floor: i64, want_base: Option<i64>) {
142        // Two adjacent batches; floor=10 should skip the first (last=9) and
143        // return the second.
144        let bytes = encoded(&[test_batch_at(0, 10, b'a'), test_batch_at(10, 10, b'b')]);
145        assert2::assert!(
146            first_batch_at_or_after(&bytes, floor).map(|b| b.base_offset) == want_base
147        );
148    }
149
150    #[test]
151    fn first_batch_at_or_after_empty_buffer_returns_none() {
152        assert!(first_batch_at_or_after(&[], 0).is_none());
153    }
154
155    #[test]
156    fn first_batch_at_or_after_rejects_floor_past_base_plus_delta() {
157        let bytes = encoded(&[test_batch_at(3, 4, b'z')]);
158
159        assert2::assert!(
160            first_batch_at_or_after(&bytes, 7).is_none(),
161            "batch 3..6 must not cover floor 7"
162        );
163    }
164
165    #[test]
166    fn first_batch_at_or_after_rejects_offset_overflow() {
167        let bytes = encoded(&[test_batch_at(i64::MAX, 2, b'z')]);
168        assert!(first_batch_at_or_after(&bytes, i64::MAX).is_none());
169    }
170
171    #[yare::parameterized(
172        both_floors_inside_first_batch = { 10, 1_600, Some((12, 1_600)) },
173        offset_floor_skips_match = { 14, 1_000, Some((14, 2_000)) },
174        no_timestamp_match = { 10, 9_999, None },
175    )]
176    fn first_record_at_or_after_timestamp_honours_both_floors(
177        floor_offset: i64,
178        target: i64,
179        want: Option<(i64, i64)>,
180    ) {
181        let bytes = encoded(&[
182            timestamped_batch_at(10, &[1_000, 1_100, 1_600, 1_700], b'a'),
183            timestamped_batch_at(14, &[2_000, 2_200, 2_400], b'b'),
184        ]);
185        assert2::assert!(
186            first_record_at_or_after_timestamp(&bytes, floor_offset, target, None).unwrap() == want
187        );
188    }
189
190    #[yare::parameterized(
191        ordinary_and_first_match = { 1_000, 0, 10, 1_500, Some((10, 2_400)) },
192        ordinary_and_second_match = { 1_000, 0, 11, 2_400, Some((11, 2_400)) },
193        ordinary_offset_past_batch = { 1_000, 0, 12, 0, None },
194        ordinary_timestamp_past_max = { 1_000, 0, 10, 2_401, None },
195        max_base_and_first_match = { i64::MAX, 1, 10, 1_500, Some((10, 2_400)) },
196        max_base_and_second_match = { i64::MAX, 1, 11, 2_400, Some((11, 2_400)) },
197        max_base_and_offset_past_batch = { i64::MAX, 1, 12, 0, None },
198        max_base_and_timestamp_past_max = { i64::MAX, 1, 10, 2_401, None },
199        min_base_and_first_match = { i64::MIN, -1, 10, 1_500, Some((10, 2_400)) },
200        min_base_and_second_match = { i64::MIN, -1, 11, 2_400, Some((11, 2_400)) },
201        min_base_and_offset_past_batch = { i64::MIN, -1, 12, 0, None },
202        min_base_and_timestamp_past_max = { i64::MIN, -1, 10, 2_401, None },
203    )]
204    fn append_time_lookup_ignores_producer_timestamps(
205        base: i64,
206        delta: i64,
207        floor: i64,
208        target: i64,
209        expected: Option<(i64, i64)>,
210    ) {
211        let mut batch = timestamped_batch_at(10, &[1_000, 1_100], b'a');
212        batch.attributes = batch
213            .attributes
214            .with_timestamp_type(TimestampType::LogAppendTime);
215        batch.max_timestamp = 2_400;
216        batch.base_timestamp = base;
217        for record in &mut batch.records {
218            record.timestamp_delta = delta;
219        }
220        let bytes = encoded(&[batch]);
221        assert2::assert!(
222            first_record_at_or_after_timestamp(&bytes, floor, target, None).unwrap() == expected
223        );
224    }
225
226    #[test]
227    fn first_record_at_or_after_timestamp_reports_corrupt_bytes() {
228        // A truncated batch header decodes to an error, never a panic.
229        let bytes = encoded(&[test_batch_at(0, 2, b'a')]);
230        let error = first_record_at_or_after_timestamp(&bytes[..12], 0, 0, None).unwrap_err();
231        assert!(matches!(error, RemoteStorageError::Io(_)));
232    }
233
234    /// A record's timestamp and the length of its value.
235    struct SizedRecord {
236        timestamp: i64,
237        value_len: usize,
238    }
239
240    const fn sized(timestamp: i64, value_len: usize) -> SizedRecord {
241        SizedRecord {
242            timestamp,
243            value_len,
244        }
245    }
246
247    /// A batch at `base_offset` with one record per [`SizedRecord`], compressed
248    /// with `codec`.
249    fn sized_batch_at(
250        base_offset: i64,
251        records: &[SizedRecord],
252        codec: krabka_compression::CompressionType,
253    ) -> RecordBatch {
254        let timestamps: Vec<i64> = records.iter().map(|record| record.timestamp).collect();
255        let mut batch = timestamped_batch_at(base_offset, &timestamps, b'z');
256        batch.attributes = batch.attributes.with_compression(codec);
257        for (record, sized) in batch.records.iter_mut().zip(records) {
258            record.value = Some(Bytes::from(vec![b'z'; sized.value_len]));
259        }
260        batch
261    }
262
263    /// The size of a record with a `1000`-byte value and nothing else: its
264    /// attributes byte, a one-byte timestamp and offset delta, a null key, the
265    /// value behind its two-byte length, and no headers.
266    const LARGE: usize = 1_007;
267
268    /// Kafka's `RemoteLogManager.lookupTimestamp` decompresses a batch only when
269    /// its max timestamp and its last offset reach the lookup's, and holds each
270    /// record it reads to the limit. A batch the producer did not compress is
271    /// never held to it.
272    #[yare::parameterized(
273        found_before_oversized_batch = {
274            krabka_compression::CompressionType::Gzip,
275            10,
276            1_000,
277            Some(100),
278            Ok(Some((10, 1_000))),
279        },
280        oversized_record_is_match = {
281            krabka_compression::CompressionType::Gzip,
282            10,
283            1_100,
284            Some(100),
285            Err(LARGE),
286        },
287        oversized_record_before_match = {
288            krabka_compression::CompressionType::Gzip,
289            10,
290            1_200,
291            Some(100),
292            Err(LARGE),
293        },
294        batch_below_offset_floor = {
295            krabka_compression::CompressionType::Gzip,
296            13,
297            1_100,
298            Some(100),
299            Ok(None),
300        },
301        no_limit = {
302            krabka_compression::CompressionType::Gzip,
303            10,
304            1_100,
305            None,
306            Ok(Some((11, 1_100))),
307        },
308        uncompressed_batch = {
309            krabka_compression::CompressionType::None,
310            10,
311            1_100,
312            Some(100),
313            Ok(Some((11, 1_100))),
314        },
315    )]
316    fn first_record_at_or_after_timestamp_refuses_only_the_records_kafka_reads(
317        codec: krabka_compression::CompressionType,
318        floor_offset: i64,
319        target: i64,
320        limit: Option<usize>,
321        want: Result<Option<(i64, i64)>, usize>,
322    ) {
323        let layout = |codec| {
324            encoded(&[
325                sized_batch_at(10, &[sized(1_000, 10)], codec),
326                sized_batch_at(11, &[sized(1_100, 1_000), sized(1_200, 10)], codec),
327            ])
328        };
329        let got =
330            match first_record_at_or_after_timestamp(&layout(codec), floor_offset, target, limit) {
331                Ok(found) => Ok(found),
332                Err(RemoteStorageError::RecordTooLarge { size, limit }) => {
333                    assert2::assert!(limit == 100);
334                    Err(size)
335                }
336                Err(other) => panic!("unexpected error: {other}"),
337            };
338        assert2::assert!(got == want);
339    }
340}
krabka_remote_storage/index/batch.rs

krabka_remote_storage/index/
batch.rs
