krabka_log/segment/
read.rs1use {
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 #[instrument(
21 level = "debug",
22 skip(self),
23 fields(base_offset = self.base_offset.0, batches = field::Empty),
24 err,
25 )]
26 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 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; };
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 #[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 #[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 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 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}