Skip to main content

krabka_remote_storage/index/
offset.rs

1//! The Kafka `OffsetIndex` on-disk layout, its floor lookup, and the byte
2//! range a remote fetch derives from it.
3//!
4//! An entry maps a relative offset to a byte position in the segment's `.log`
5//! file, so the read path turns a requested offset into a position and then
6//! into an inclusive object-store byte range.
7
8use {
9    super::{BytePosition, RelativeOffset, corrupt_index},
10    crate::error::RemoteStorageError,
11    krabka_verified::remote_read::remote_fetch_end_position,
12    std::mem,
13    zerocopy::{BigEndian, FromBytes, Immutable, KnownLayout, Unaligned, byteorder::U32},
14};
15
16/// 8 bytes per entry: rel u32 BE, then pos u32 BE.
17///
18/// It mirrors
19/// `krabka_log::index::OffsetEntryRaw`, so the remote-tier copy of an
20/// `OffsetIndex` file decodes through the same byte layout that wrote the
21/// local index.
22#[derive(Debug, Clone, Copy, FromBytes, KnownLayout, Immutable, Unaligned)]
23#[repr(C)]
24pub struct OffsetIndexEntry {
25    /// Offset of the indexed record, relative to the segment's base offset.
26    pub relative_offset: U32<BigEndian>,
27    /// Byte position of that record's batch within the `.log` file.
28    pub position: U32<BigEndian>,
29}
30
31/// Byte length of one serialized offset-index entry.
32const OFFSET_INDEX_ENTRY_LEN: usize = mem::size_of::<OffsetIndexEntry>();
33
34const _: [(); 8] = [(); OFFSET_INDEX_ENTRY_LEN];
35
36/// Computes the inclusive `end_position` for a remote byte-range fetch.
37///
38/// It returns `None`, which means read to the end of the segment, when
39/// `start_position` plus `max_bytes` would reach or pass `segment_size`. In
40/// every other case it returns the inclusive last byte to read.
41#[must_use]
42pub fn end_position_for(
43    start_position: BytePosition,
44    segment_size: u32,
45    max_bytes: usize,
46) -> Option<BytePosition> {
47    let max_bytes_u32 = u32::try_from(max_bytes).unwrap_or(u32::MAX);
48    remote_fetch_end_position(start_position, segment_size, max_bytes_u32)
49}
50
51/// Borrows Kafka's `OffsetIndex` on-disk format as a zero-copy
52/// `&[OffsetIndexEntry]`.
53///
54/// Each entry is 8 bytes: rel u32 BE, then pos u32 BE. It ignores trailing bytes that do not complete an 8-byte entry. The result
55/// borrows from `bytes`.
56///
57/// # Errors
58///
59/// Returns [`RemoteStorageError::Io`] when the object store returned bytes that
60/// do not form an entry array.
61pub fn parse_offset_index(bytes: &[u8]) -> Result<&[OffsetIndexEntry], RemoteStorageError> {
62    let truncated_len = (bytes.len() / OFFSET_INDEX_ENTRY_LEN) * OFFSET_INDEX_ENTRY_LEN;
63    bytes
64        .get(..truncated_len)
65        .and_then(|whole| <[OffsetIndexEntry]>::ref_from_bytes(whole).ok())
66        .ok_or_else(|| corrupt_index("offset"))
67}
68
69/// Floor lookup: the byte position of the largest entry with
70/// `rel <= target_rel`.
71///
72/// It returns 0 when the index is empty, and when the
73/// target is before the first entry. It runs directly against the borrowed
74/// zero-copy slice and builds no owned `Vec`.
75#[must_use]
76pub fn position_for_relative_offset(
77    entries: &[OffsetIndexEntry],
78    target_rel: RelativeOffset,
79) -> BytePosition {
80    let floor = match entries.binary_search_by_key(&target_rel, |e| e.relative_offset.get()) {
81        Ok(i) => Some(i),
82        Err(i) => i.checked_sub(1),
83    };
84    floor
85        .and_then(|i| entries.get(i))
86        .map_or(0, |entry| entry.position.get())
87}
88
89#[cfg(test)]
90mod tests {
91    use {super::*, assert2::assert};
92
93    /// An offset index entry's decoded relative offset and byte position.
94    #[derive(Debug, PartialEq, Eq)]
95    struct OffsetSpec {
96        relative_offset: u32,
97        position: u32,
98    }
99
100    const fn offset_at(relative_offset: u32, position: u32) -> OffsetSpec {
101        OffsetSpec {
102            relative_offset,
103            position,
104        }
105    }
106
107    fn offset_entries(specs: &[OffsetSpec]) -> Vec<OffsetIndexEntry> {
108        specs
109            .iter()
110            .map(|spec| OffsetIndexEntry {
111                relative_offset: U32::new(spec.relative_offset),
112                position: U32::new(spec.position),
113            })
114            .collect()
115    }
116
117    #[yare::parameterized(
118        one_entry = { vec![(0_u32, 0_u32)] },
119        two_entries = { vec![(0_u32, 0_u32), (10, 256)] },
120        all_known_entries = { vec![(0_u32, 0_u32), (10, 256), (20, 512)] },
121    )]
122    fn parse_offset_index_round_trips_known_entries(specs: Vec<(u32, u32)>) {
123        // Mirror OffsetIndex::append: 4B rel BE, 4B pos BE.
124        let mut buf = Vec::new();
125        for (rel, pos) in &specs {
126            buf.extend_from_slice(&rel.to_be_bytes());
127            buf.extend_from_slice(&pos.to_be_bytes());
128        }
129        let entries = parse_offset_index(&buf).expect("valid offset index");
130        let decoded: Vec<OffsetSpec> = entries
131            .iter()
132            .map(|e| offset_at(e.relative_offset.get(), e.position.get()))
133            .collect();
134        assert2::assert!(
135            decoded
136                == specs
137                    .iter()
138                    .map(|(rel, pos)| offset_at(*rel, *pos))
139                    .collect::<Vec<_>>()
140        );
141    }
142
143    #[yare::parameterized(
144        exact = { 10, 256 },
145        between = { 15, 256 },
146        first_entry = { 0, 0 },
147        after_last = { 100, 1024 },
148    )]
149    fn position_for_relative_offset_returns_floor(rel: u32, want: u32) {
150        let entries = offset_entries(&[
151            offset_at(0, 0),
152            offset_at(10, 256),
153            offset_at(20, 512),
154            offset_at(30, 1024),
155        ]);
156        assert2::assert!(position_for_relative_offset(&entries, rel) == want);
157    }
158
159    #[test]
160    fn position_for_relative_offset_empty_index_returns_zero() {
161        assert!(position_for_relative_offset(&[], 50) == 0);
162    }
163
164    #[test]
165    fn position_for_relative_offset_below_first() {
166        // Synthetic: first entry isn't at rel=0. Floor below it returns 0.
167        let entries = offset_entries(&[offset_at(5, 100), offset_at(10, 200)]);
168        assert!(position_for_relative_offset(&entries, 3) == 0);
169    }
170
171    #[yare::parameterized(
172        finite_inclusive_end = { 0, 1024, 256, Some(255) },
173        cap_beyond_segment = { 512, 1024, 999_999, None },
174        zero_cap_means_no_cap = { 0, 1024, 0, None },
175        start_past_segment = { u32::MAX, 1024, 100, None },
176        wide_cap_saturates = { 0, u32::MAX, usize::MAX, None },
177    )]
178    fn end_position_for_caps_with_max_bytes(
179        start: u32,
180        segment: u32,
181        max_bytes: usize,
182        want: Option<u32>,
183    ) {
184        assert2::assert!(end_position_for(start, segment, max_bytes) == want);
185    }
186}