krabka_log/segment/
read_raw.rs

1//! The verbatim, decode-free read path.
2//!
3//! A fetch serves the producer's own batch bytes, so this walk reads only the
4//! fixed v2 batch headers to find run boundaries and never touches a record
5//! body. The zero-copy descriptor form of the same walk is the sibling
6//! `read_raw_desc` module.
7
8use {
9    super::{
10        RawSegmentRead, Segment,
11        header_walk::HeaderWalk,
12        io::SourceIo,
13        read_source::{ReadSource, relative_target},
14    },
15    crate::{
16        config::{DEFAULT_READ_AHEAD_MAX, DEFAULT_READ_BUFFER_CAP},
17        error::LogError,
18    },
19    bytes::Bytes,
20    krabka_ids::Offset,
21    krabka_protocol::records::{HEADER_LEN, RecordBatchHeader},
22    krabka_units::prelude::{ByteSize, ByteSizeExt},
23    std::ops::Range,
24    tracing::{field, instrument},
25    zerocopy::FromBytes,
26};
27
28/// Header fields used to choose a verbatim range without decoding records.
29pub(super) struct RawBatch {
30    start_offset: Offset,
31    last_offset: Offset,
32    pub(super) len: usize,
33}
34
35impl RawBatch {
36    pub(super) fn decode(bytes: &[u8]) -> Result<Self, LogError> {
37        let header = RecordBatchHeader::ref_from_bytes(bytes)
38            .map_err(|_| LogError::Corrupt("record batch header".into()))?;
39        let base = header.base_offset.get();
40        Ok(Self {
41            start_offset: Offset(base),
42            last_offset: Offset(base + i64::from(header.last_offset_delta.get())),
43            len: 12 + usize::try_from(header.batch_length.get().max(0)).unwrap_or(0),
44        })
45    }
46}
47
48/// A selected byte range relative to the first read position.
49pub(super) struct RawRange {
50    pub(super) positions: Range<usize>,
51    pub(super) start_offset: Offset,
52    pub(super) last_offset: Offset,
53}
54
55/// Select the same complete batches for buffered and file-region fetches.
56///
57/// Only the first eligible batch may extend past the initial window. The
58/// caller checks that this anti-stall batch is complete in its source.
59pub(super) fn select_raw_range(
60    fetch_offset: Offset,
61    limit_offset: Offset,
62    max_bytes: usize,
63    window: usize,
64    mut read_header: impl FnMut(usize) -> Result<Option<RawBatch>, LogError>,
65) -> Result<Option<RawRange>, LogError> {
66    let mut pos = 0;
67    let mut range: Option<RawRange> = None;
68    while pos + HEADER_LEN <= window {
69        let Some(batch) = read_header(pos)? else {
70            break;
71        };
72        let end = pos + batch.len;
73        if batch.last_offset < fetch_offset {
74            pos = end;
75            continue;
76        }
77        if batch.start_offset >= limit_offset || (end > window && range.is_some()) {
78            break;
79        }
80        let selected = range.get_or_insert(RawRange {
81            positions: pos..end,
82            start_offset: batch.start_offset,
83            last_offset: batch.last_offset,
84        });
85        selected.positions.end = end;
86        selected.last_offset = batch.last_offset;
87        pos = end;
88        if end > window || selected.positions.len() >= max_bytes {
89            break;
90        }
91    }
92    Ok(range)
93}
94
95impl Segment {
96    /// Read a contiguous run of **complete, verbatim** record-batch bytes.
97    ///
98    /// The run starts at the batch that contains `fetch_offset`. It includes
99    /// only batches whose `base_offset < limit_offset`, up to about
100    /// `max_bytes`, and always at least one batch. That last rule is Kafka's
101    /// anti-stall rule. This method decodes no records. It reads only the
102    /// fixed batch headers to find the boundaries.
103    #[instrument(
104        level = "debug",
105        skip(self),
106        fields(base_offset = self.base_offset.0, bytes = field::Empty),
107        err,
108    )]
109    /// # Errors
110    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
111    pub fn read_raw(
112        &self,
113        fetch_offset: Offset,
114        limit_offset: Offset,
115        max_size: ByteSize,
116    ) -> Result<RawSegmentRead, LogError> {
117        self.read_raw_with_policy(
118            fetch_offset,
119            limit_offset,
120            max_size,
121            DEFAULT_READ_BUFFER_CAP,
122            DEFAULT_READ_AHEAD_MAX,
123        )
124    }
125
126    #[cfg(test)]
127    pub(super) fn raw_start_position(
128        &self,
129        fetch_offset: Offset,
130        limit_offset: Offset,
131        operation: &str,
132    ) -> Result<Option<u64>, LogError> {
133        self.raw_start_position_with(fetch_offset, limit_offset, operation, |pos, buf, len| {
134            self.read_log_range(pos, buf, len)
135        })
136    }
137}
138
139/// The verbatim byte-copy read of a [`ReadSource`]. Provided for every
140/// source, so it runs the same over a live segment under the log lock and
141/// over a view of one after the lock is released.
142pub trait RawReadSource: ReadSource {
143    /// [`Segment::raw_start_position`] over the positional reads of
144    /// `read_range`.
145    fn raw_start_position_with(
146        &self,
147        fetch_offset: Offset,
148        limit_offset: Offset,
149        operation: &str,
150        read_range: impl FnMut(u64, &mut Vec<u8>, usize) -> Result<(), LogError>,
151    ) -> Result<Option<u64>, LogError> {
152        if fetch_offset > self.last() || fetch_offset >= limit_offset {
153            return Ok(None);
154        }
155        let target_rel = relative_target(self.base(), fetch_offset, operation)?;
156        self.read_start_position_with(target_rel, read_range)
157            .map(Some)
158    }
159
160    /// [`Segment::read_raw`] under the log's configured read policy: the
161    /// initial allocation is capped at `read_buffer_cap`, and the readahead
162    /// hint reaches at most `read_ahead_max` past the read's own window.
163    ///
164    /// # Errors
165    /// Returns an error when log I/O fails, a record or index is corrupt, or
166    /// the requested offset violates the segment state.
167    fn read_raw_with_policy(
168        &self,
169        fetch_offset: Offset,
170        limit_offset: Offset,
171        max_size: ByteSize,
172        read_buffer_cap: ByteSize,
173        read_ahead_max: ByteSize,
174    ) -> Result<RawSegmentRead, LogError> {
175        let Some(start_pos) = self.raw_start_position_with(
176            fetch_offset,
177            limit_offset,
178            "read_raw",
179            |pos, buf, len| self.read_log_range(pos, buf, len),
180        )?
181        else {
182            return Ok(RawSegmentRead::empty());
183        };
184
185        // Below this line the budget indexes into a byte buffer, so it
186        // crosses back to `usize` once, here.
187        let max_bytes = max_size.bytes_usize();
188        let first_read = max_bytes.max(HEADER_LEN);
189        // A segment that writes directly holds its newest batches in memory.
190        // Serve them as a slice of what it holds, without the copy below:
191        // the fetch runs under the log lock, which the producer's appends
192        // wait for.
193        if let Some(cached) = self.cached_bytes(start_pos, first_read) {
194            let range =
195                select_raw_range(fetch_offset, limit_offset, max_bytes, cached.len(), |pos| {
196                    cached
197                        .get(pos..pos + HEADER_LEN)
198                        .map(RawBatch::decode)
199                        .transpose()
200                })?;
201            if let Some(range) = range.filter(|range| range.positions.end <= cached.len()) {
202                let bytes = cached.slice(range.positions);
203                tracing::Span::current().record("bytes", bytes.len());
204                return Ok(RawSegmentRead {
205                    start_offset: range.start_offset,
206                    last_offset: range.last_offset,
207                    bytes,
208                    file_range: None,
209                });
210            }
211        }
212        self.advise_read(start_pos, first_read as u64, read_ahead_max);
213        // A tailing consumer can request megabytes while only one small
214        // batch remains. The returned Bytes retains this allocation, so
215        // reserve at most what read_log_range can actually read.
216        let available =
217            usize::try_from(self.log_size().saturating_sub(start_pos)).unwrap_or(usize::MAX);
218        let mut buf: Vec<u8> =
219            Vec::with_capacity(first_read.min(read_buffer_cap.bytes_usize()).min(available));
220        self.read_log_range(start_pos, &mut buf, first_read)?;
221
222        let Some(range) =
223            select_raw_range(fetch_offset, limit_offset, max_bytes, buf.len(), |pos| {
224                buf.get(pos..pos + HEADER_LEN)
225                    .map(RawBatch::decode)
226                    .transpose()
227            })?
228        else {
229            return Ok(RawSegmentRead::empty());
230        };
231        let file_range =
232            start_pos + range.positions.start as u64..start_pos + range.positions.end as u64;
233        let bytes = if range.positions.end > buf.len() {
234            let len = range.positions.len();
235            let mut one = Vec::with_capacity(len);
236            self.read_log_range(file_range.start, &mut one, len)?;
237            if one.len() < len {
238                return Ok(RawSegmentRead::empty());
239            }
240            Bytes::from(one)
241        } else {
242            Bytes::from(buf).slice(range.positions)
243        };
244        tracing::Span::current().record("bytes", bytes.len());
245        Ok(RawSegmentRead {
246            start_offset: range.start_offset,
247            last_offset: range.last_offset,
248            bytes,
249            file_range: Some(file_range),
250        })
251    }
252}
253
254impl<S: ReadSource + ?Sized> RawReadSource for S {}
255
256#[cfg(test)]
257mod tests {
258    use {
259        super::*,
260        assert2::check,
261        krabka_units::prelude::{bytes, mebibytes},
262    };
263
264    fn assert_read_offsets(read: &RawSegmentRead, start: i64, last: i64) {
265        check!(
266            read.start_offset == Offset(start),
267            "start {:?}",
268            read.start_offset
269        );
270        check!(
271            read.last_offset == Offset(last),
272            "last {:?}",
273            read.last_offset
274        );
275    }
276    use crate::segment::test_support::{
277        DENSE_INDEX, NO_LIMIT, SECOND_TWO_RECORD_BATCH, THIRD_TWO_RECORD_BATCH, THREE_RECORD_BATCH,
278        TWO_RECORD_BATCH, TempSegment, sample_batch, seeded_fixture, test_batch_at, test_segment,
279        two_record_fixture,
280    };
281
282    /// A fetch reads nothing when it starts past the segment, and nothing when
283    /// it starts at or past the limit. Either condition alone is enough --
284    /// joined with `&&` a fetch would have to be both before it read nothing.
285    #[derive(Clone, Copy)]
286    struct BatchCount(i64);
287
288    fn append_two_record_batches(seg: &mut Segment, count: BatchCount) {
289        for i in 0..count.0 {
290            seg.append(
291                &sample_batch()
292                    .offset(i * 2)
293                    .records(2)
294                    .timestamp(100 + i)
295                    .call(),
296                DENSE_INDEX,
297            )
298            .unwrap();
299        }
300    }
301
302    #[test]
303    fn a_fetch_past_the_segment_or_at_the_limit_reads_nothing() {
304        let TempSegment {
305            dir: _dir,
306            segment: seg,
307        } = seeded_fixture()
308            .batches(&[
309                THREE_RECORD_BATCH, // offsets 0..=2
310            ])
311            .call();
312
313        let past = seg.read_raw(Offset(3), Offset(99), NO_LIMIT).unwrap();
314        check!(past.is_empty(), "a fetch past the last offset");
315
316        let at_limit = seg.read_raw(Offset(0), Offset(0), NO_LIMIT).unwrap();
317        check!(at_limit.is_empty(), "a fetch at the limit");
318
319        let inside = seg.read_raw(Offset(0), Offset(3), NO_LIMIT).unwrap();
320        check!(
321            !inside.is_empty(),
322            "a fetch inside the segment and below the limit"
323        );
324    }
325
326    #[test]
327    fn a_small_tail_does_not_retain_the_consumers_large_read_budget() {
328        let TempSegment {
329            dir: _dir,
330            segment: mut seg,
331        } = test_segment();
332        append_two_record_batches(&mut seg, BatchCount(4));
333        let read = seg.read_raw(Offset(6), Offset(99), mebibytes(10)).unwrap();
334        assert_read_offsets(&read, 6, 7);
335        let length = read.bytes.len();
336        let allocation = read.bytes.try_into_mut().expect("unique read buffer");
337        check!(length > HEADER_LEN);
338        check!(allocation.capacity() == length);
339    }
340
341    /// Batches before the fetch offset are stepped over by their own length.
342    /// Advancing by anything else lands mid-batch and the walk reads a header
343    /// out of the middle of a record.
344    #[test]
345    fn a_fetch_mid_segment_steps_over_the_batches_before_it() {
346        let TempSegment {
347            dir: _dir,
348            segment: mut seg,
349        } = test_segment();
350        // A sparse index, so only the first batch gets an entry and the lookup
351        // lands at the segment head. The walk then has to step over the two
352        // batches before the one asked for -- with a dense index it would jump
353        // straight there and never step at all.
354        let sparse = mebibytes(1);
355        seg.append(
356            &sample_batch()
357                .offset(TWO_RECORD_BATCH.offset.0)
358                .records(TWO_RECORD_BATCH.records.0)
359                .timestamp(TWO_RECORD_BATCH.timestamp.0)
360                .call(),
361            sparse,
362        )
363        .unwrap(); // 0..=1
364        seg.append(
365            &sample_batch()
366                .offset(SECOND_TWO_RECORD_BATCH.offset.0)
367                .records(SECOND_TWO_RECORD_BATCH.records.0)
368                .timestamp(SECOND_TWO_RECORD_BATCH.timestamp.0)
369                .call(),
370            sparse,
371        )
372        .unwrap(); // 2..=3
373        seg.append(
374            &sample_batch()
375                .offset(THIRD_TWO_RECORD_BATCH.offset.0)
376                .records(THIRD_TWO_RECORD_BATCH.records.0)
377                .timestamp(THIRD_TWO_RECORD_BATCH.timestamp.0)
378                .call(),
379            sparse,
380        )
381        .unwrap(); // 4..=5
382
383        let read = seg
384            .read_raw_with_policy(
385                Offset(4),
386                Offset(99),
387                NO_LIMIT,
388                mebibytes(1),
389                DEFAULT_READ_AHEAD_MAX,
390            )
391            .unwrap();
392        assert_read_offsets(&read, 4, 5);
393    }
394
395    /// A batch that will not fit the first read is fetched on its own, from
396    /// its own position in the file.
397    ///
398    /// That position is `start_pos + pos`, and with a dense index and a fetch
399    /// past the first batch `start_pos` is well away from the file head -- so
400    /// reading from anywhere else returns a different batch, or nothing.
401    #[test]
402    fn a_batch_too_large_for_the_first_read_is_fetched_from_its_own_position() {
403        let TempSegment {
404            dir: _dir,
405            segment: mut seg,
406        } = test_segment();
407        append_two_record_batches(&mut seg, BatchCount(4));
408
409        // A one-byte budget makes the first read header-sized, so the batch
410        // cannot fit it and takes the read-one-batch path.
411        let read = seg.read_raw(Offset(4), Offset(99), bytes(1)).unwrap();
412        assert_read_offsets(&read, 4, 5);
413        check!(!read.is_empty());
414    }
415
416    #[test]
417    fn batch_too_large_with_nonzero_start_pos_and_pos() {
418        let TempSegment {
419            dir: _dir,
420            segment: mut seg,
421        } = two_record_fixture();
422        let p1 = seg.log_size;
423        seg.append(
424            &sample_batch()
425                .offset(SECOND_TWO_RECORD_BATCH.offset.0)
426                .records(SECOND_TWO_RECORD_BATCH.records.0)
427                .timestamp(SECOND_TWO_RECORD_BATCH.timestamp.0)
428                .call(),
429            DENSE_INDEX,
430        )
431        .unwrap();
432        let b1_len = seg.log_size - p1;
433        let no_index = mebibytes(1);
434        seg.append(
435            &sample_batch()
436                .offset(THIRD_TWO_RECORD_BATCH.offset.0)
437                .records(THIRD_TWO_RECORD_BATCH.records.0)
438                .timestamp(THIRD_TWO_RECORD_BATCH.timestamp.0)
439                .call(),
440            no_index,
441        )
442        .unwrap();
443        seg.append(
444            &sample_batch().offset(6).records(2).timestamp(400).call(),
445            no_index,
446        )
447        .unwrap();
448
449        let budget = bytes(u32::try_from(b1_len).unwrap() + 61);
450        let read = seg.read_raw(Offset(4), Offset(99), budget).unwrap();
451        check!(read.start_offset == Offset(4));
452        check!(read.last_offset == Offset(5));
453        check!(!read.is_empty());
454    }
455
456    /// The byte budget stops the walk once the selected range reaches it, so a
457    /// small budget returns fewer batches than an unlimited one.
458    #[test]
459    fn the_byte_budget_bounds_how_much_a_fetch_returns() {
460        let TempSegment {
461            dir: _dir,
462            segment: mut seg,
463        } = test_segment();
464        append_two_record_batches(&mut seg, BatchCount(6));
465
466        let everything = seg.read_raw(Offset(0), Offset(99), NO_LIMIT).unwrap();
467        let clipped = seg.read_raw(Offset(0), Offset(99), bytes(1)).unwrap();
468        check!(
469            !clipped.is_empty(),
470            "a budget of one byte still returns a batch"
471        );
472        check!(
473            clipped.last_offset < everything.last_offset,
474            "clipped to {:?}, unlimited reached {:?}",
475            clipped.last_offset,
476            everything.last_offset
477        );
478
479        let batch_size = u32::try_from(
480            sample_batch()
481                .offset(TWO_RECORD_BATCH.offset.0)
482                .records(TWO_RECORD_BATCH.records.0)
483                .timestamp(TWO_RECORD_BATCH.timestamp.0)
484                .call()
485                .encoded_len(),
486        )
487        .unwrap();
488        let two_batches = seg
489            .read_raw(Offset(0), Offset(99), bytes(batch_size * 2))
490            .unwrap();
491        check!(two_batches.start_offset == Offset(0));
492        check!(two_batches.last_offset == Offset(3));
493    }
494
495    // `Segment::read_raw` maps the fetch offset to the relative index key the
496    // same way. base_offset 100, dense index, `read_raw(103)` must begin at
497    // the offset-103 batch (`start_offset == 103`). Mutating `-`→`+` computes
498    // `203`, whose lookup skips past the offset-103 batch → `start_offset`
499    // becomes 105.
500    #[test]
501    fn read_raw_uses_relative_offset_for_index_lookup() {
502        let TempSegment {
503            dir: _dir,
504            segment: seg,
505        } = super::super::test_support::indexed_segment();
506
507        let r = seg.read_raw(Offset(103), Offset(1000), NO_LIMIT).unwrap();
508        assert2::assert!(!r.is_empty());
509        assert2::assert!(r.start_offset == Offset(103));
510    }
511
512    #[test]
513    fn read_raw_is_byte_exact_and_multi_batch() {
514        let TempSegment {
515            dir,
516            segment: mut seg,
517        } = test_segment();
518        let mut wire = bytes::BytesMut::new();
519        for off in 0..3i64 {
520            let b = test_batch_at(Offset(off));
521            seg.append(&b, DENSE_INDEX).unwrap();
522            b.encode(&mut wire).unwrap();
523        }
524        let wire = wire.freeze();
525        let r = seg.read_raw(Offset(0), Offset(3), mebibytes(10)).unwrap();
526        assert2::assert!(r.start_offset == Offset(0));
527        assert2::assert!(r.last_offset == Offset(2));
528        assert2::assert!(&r.bytes[..] == &wire[..]);
529        drop(dir);
530    }
531
532    #[test]
533    fn read_raw_clamps_at_limit_offset() {
534        let TempSegment {
535            dir,
536            segment: mut seg,
537        } = test_segment();
538        let mut expected = bytes::BytesMut::new();
539        for off in 0..3i64 {
540            let batch = test_batch_at(Offset(off));
541            seg.append(&batch, DENSE_INDEX).unwrap();
542            if off < 2 {
543                batch.encode(&mut expected).unwrap();
544            }
545        }
546        let r = seg.read_raw(Offset(0), Offset(2), mebibytes(10)).unwrap();
547        assert2::assert!(r.start_offset == Offset(0));
548        assert2::assert!(r.last_offset == Offset(1));
549        assert2::assert!(&r.bytes[..] == &expected[..]);
550        drop(dir);
551    }
552
553    #[test]
554    fn read_raw_returns_at_least_one_batch_over_budget() {
555        let TempSegment {
556            dir,
557            segment: mut seg,
558        } = test_segment();
559        let batch = test_batch_at(Offset(0));
560        let mut expected = bytes::BytesMut::new();
561        batch.encode(&mut expected).unwrap();
562        seg.append(&batch, DENSE_INDEX).unwrap();
563        let r = seg.read_raw(Offset(0), Offset(1), bytes(1)).unwrap();
564        assert2::assert!(r.start_offset == Offset(0));
565        assert2::assert!(r.last_offset == Offset(0));
566        assert2::assert!(&r.bytes[..] == &expected[..]);
567        drop(dir);
568    }
569}
krabka_log/segment/read_raw.rs

krabka_log/segment/
read_raw.rs
