Skip to main content

krabka_log/segment/
read.rs

1//! The decoding read path: batches come back as [`RecordBatch`] values.
2//!
3//! This is the read that callers use when they need the records themselves.
4//! The verbatim, decode-free counterpart lives in the sibling `read_raw`
5//! module.
6
7use {
8    super::{Segment, io::SourceIo as _},
9    crate::{config::DEFAULT_READ_BUFFER_CAP, error::LogError},
10    krabka_ids::Offset,
11    krabka_protocol::records::RecordBatch,
12    krabka_units::prelude::{ByteSize, ByteSizeExt},
13    tracing::{field, instrument},
14};
15
16impl Segment {
17    /// Read batches from `offset` or just before it, up to about `max_bytes`
18    /// of `.log` data. The result is an empty `Vec` when `offset` is past
19    /// `last_offset`.
20    #[instrument(
21        level = "debug",
22        skip(self),
23        fields(base_offset = self.base_offset.0, batches = field::Empty),
24        err,
25    )]
26    /// # Errors
27    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
28    pub fn read(&self, offset: Offset, max_size: ByteSize) -> Result<Vec<RecordBatch>, LogError> {
29        self.read_with_buffer_cap(offset, max_size, DEFAULT_READ_BUFFER_CAP)
30    }
31
32    pub(crate) fn read_with_buffer_cap(
33        &self,
34        offset: Offset,
35        max_size: ByteSize,
36        read_buffer_cap: ByteSize,
37    ) -> Result<Vec<RecordBatch>, LogError> {
38        if offset > self.last_offset {
39            return Ok(vec![]);
40        }
41        let target_rel = u32::try_from((offset.0 - self.base_offset.0).max(0))
42            .map_err(|_| LogError::BadSegmentName("target offset out of range".into()))?;
43        let start_pos = self.read_start_position(target_rel)?;
44
45        // Below this line the budget is a buffer length and a file-read
46        // count, so it crosses back to `usize` once, here.
47        let max_bytes = max_size.bytes_usize();
48        let initial_cap = max_size.min(read_buffer_cap).bytes_usize();
49        let mut buf: Vec<u8> = Vec::with_capacity(initial_cap);
50        self.read_log_range(start_pos, &mut buf, max_bytes)?;
51
52        let mut out: Vec<RecordBatch> = Vec::new();
53        let mut total: usize = 0;
54        let mut cursor: &[u8] = &buf;
55        while !cursor.is_empty() {
56            let before = cursor.len();
57            let Ok(batch) = RecordBatch::decode(&mut cursor) else {
58                break; // partial trailing batch — stop.
59            };
60            let consumed = before - cursor.len();
61            let batch_last = Offset(batch.base_offset + i64::from(batch.last_offset_delta));
62            if batch_last >= offset {
63                out.push(batch);
64                total += consumed;
65                if total >= max_bytes {
66                    break;
67                }
68            }
69        }
70        tracing::Span::current().record("batches", out.len());
71        Ok(out)
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use {
78        super::*,
79        crate::segment::test_support::{
80            FIRST_LARGE_TIMESTAMP_BATCH, INDEXED_READ_BATCHES, NO_LIMIT, ONE_RECORD_BATCH,
81            SECOND_LARGE_TIMESTAMP_BATCH, SECOND_SINGLE_RECORD_BATCH, THIRD_SINGLE_RECORD_BATCH,
82            TempSegment, sample_batch, seeded_fixture, test_segment, three_single_record_batches,
83        },
84        krabka_units::prelude::kibibytes,
85    };
86
87    // `Segment::read` maps the absolute fetch offset to a relative index key
88    // via `offset - base_offset`. With a dense index and base_offset 100,
89    // reading from offset 103 must start at the batch containing 103 and
90    // return the batches at 103 and 105. Mutating `-`→`+` computes
91    // `103 + 100 = 203`, whose index lookup lands at (or past) the last
92    // batch, skipping the offset-103 batch.
93    #[test]
94    fn read_uses_relative_offset_for_index_lookup() {
95        let TempSegment {
96            dir: _dir,
97            segment: seg,
98        } = super::super::test_support::indexed_segment();
99        let read = seg.read(Offset(103), NO_LIMIT).unwrap();
100        assert2::assert!(seg.last_offset() == Offset(105));
101        assert2::assert!(
102            read == vec![
103                sample_batch()
104                    .offset(INDEXED_READ_BATCHES[1].offset.0)
105                    .records(INDEXED_READ_BATCHES[1].records.0)
106                    .timestamp(INDEXED_READ_BATCHES[1].timestamp.0)
107                    .call(),
108                sample_batch()
109                    .offset(INDEXED_READ_BATCHES[2].offset.0)
110                    .records(INDEXED_READ_BATCHES[2].records.0)
111                    .timestamp(INDEXED_READ_BATCHES[2].timestamp.0)
112                    .call()
113            ]
114        );
115    }
116
117    /// `Segment::read` accumulates consumed bytes as `before - cursor.len()`,
118    /// the exact number of bytes each batch decode advanced, to enforce the
119    /// `max_bytes` budget. With `max_bytes` set to the segment's full size,
120    /// all three batches fit and the read returns them. A mutation of `-` to
121    /// `+` inflates `consumed` on the first batch past `max_bytes`, so the
122    /// read stops after one batch.
123    #[test]
124    fn read_consumed_bytes_gates_max_bytes_budget() {
125        let TempSegment {
126            dir: _dir,
127            segment: seg,
128        } = seeded_fixture()
129            .batches(&[
130                ONE_RECORD_BATCH,
131                SECOND_SINGLE_RECORD_BATCH,
132                THIRD_SINGLE_RECORD_BATCH,
133            ])
134            .call();
135        // Exactly the whole segment: correct consumed accounting fits all three
136        // batches; inflated accounting overshoots after the first.
137        let max_size = seg.size();
138
139        let read = seg.read(Offset(0), max_size).unwrap();
140        assert2::assert!(read == three_single_record_batches());
141    }
142
143    #[test]
144    fn read_at_higher_offset_skips_earlier_batches() {
145        let TempSegment {
146            dir: _dir,
147            segment: mut seg,
148        } = test_segment();
149        seg.append(
150            &sample_batch()
151                .offset(FIRST_LARGE_TIMESTAMP_BATCH.offset.0)
152                .records(FIRST_LARGE_TIMESTAMP_BATCH.records.0)
153                .timestamp(FIRST_LARGE_TIMESTAMP_BATCH.timestamp.0)
154                .call(),
155            kibibytes(4),
156        )
157        .unwrap();
158        seg.append(
159            &sample_batch()
160                .offset(SECOND_LARGE_TIMESTAMP_BATCH.offset.0)
161                .records(SECOND_LARGE_TIMESTAMP_BATCH.records.0)
162                .timestamp(SECOND_LARGE_TIMESTAMP_BATCH.timestamp.0)
163                .call(),
164            kibibytes(4),
165        )
166        .unwrap();
167        let read = seg.read(Offset(4), NO_LIMIT).unwrap();
168        // Offset 4 falls inside the second batch (offsets 3..=4).
169        assert2::assert!(
170            read == vec![
171                sample_batch()
172                    .offset(SECOND_LARGE_TIMESTAMP_BATCH.offset.0)
173                    .records(SECOND_LARGE_TIMESTAMP_BATCH.records.0)
174                    .timestamp(SECOND_LARGE_TIMESTAMP_BATCH.timestamp.0)
175                    .call()
176            ]
177        );
178    }
179
180    #[test]
181    fn read_past_last_offset_returns_empty() {
182        let TempSegment {
183            dir: _dir,
184            segment: mut seg,
185        } = test_segment();
186        seg.append(
187            &sample_batch().records(2).timestamp(1_000).call(),
188            kibibytes(4),
189        )
190        .unwrap();
191        let read = seg.read(Offset(100), NO_LIMIT).unwrap();
192        assert2::assert!(read.is_empty());
193    }
194}