Skip to main content

krabka_log/leader_epoch_checkpoint/
file.rs

1//! Reading and writing the checkpoint file itself: `open` and its strict
2//! `parse`, and the atomic `flush` that every mutation ends with. The parser
3//! and the writer live together because they are the two halves of one
4//! byte-for-byte Kafka format, and a change to either has to move the other.
5
6use {
7    super::{EpochEntry, LeaderEpochCheckpoint, is_strict_successor},
8    crate::{
9        error::LogError,
10        io::{IoTarget, file_io, write_atomic},
11    },
12    krabka_ids::{LeaderEpoch, Offset},
13    std::{
14        fmt::Write as _,
15        fs,
16        io::ErrorKind,
17        path::{Path, PathBuf},
18        str::{FromStr, SplitWhitespace},
19    },
20};
21
22/// The `leader-epoch-checkpoint` format version this build writes and reads:
23/// the first line of the file, as in Kafka's `CheckpointFile`.
24///
25/// It is part of the 1.x on-disk contract, and it is Kafka's number too:
26/// `LeaderEpochCheckpointFile.CURRENT_VERSION` is 0.
27pub const LEADER_EPOCH_CHECKPOINT_VERSION: i32 = 0;
28
29/// The artifact name the version errors carry, the file's own name.
30const ARTIFACT: &str = "leader-epoch-checkpoint";
31
32impl LeaderEpochCheckpoint {
33    #[krabka_macros::index_open]
34    /// Open or recover the checkpoint at `path`. A missing file gives an
35    /// empty checkpoint.
36    /// # Errors
37    /// Returns [`LogError::UnsupportedFormatVersion`] when the header line is
38    /// a version other than `LEADER_EPOCH_CHECKPOINT_VERSION`,
39    /// [`LogError::MissingFormatVersion`] when it is not a version at all, and
40    /// an error when log I/O fails or a row is corrupt.
41    pub fn open(path: PathBuf) -> Result<Self, LogError> {
42        let entries = match fs::read_to_string(&path) {
43            Ok(s) => Self::parse(&path, &s)?,
44            Err(e) if e.kind() == ErrorKind::NotFound => Vec::new(),
45            Err(e) => return Err(LogError::Io(e)),
46        };
47        tracing::Span::current().record("entries", entries.len());
48        Ok(Self {
49            path,
50            io: file_io(),
51            entries,
52        })
53    }
54
55    /// Parse the file's text. As Kafka's `CheckpointFile.read`, an empty
56    /// file has no entries, and a header other than the current version is
57    /// refused: Kafka throws "Unrecognized version of the ... file".
58    fn parse(path: &Path, s: &str) -> Result<Vec<EpochEntry>, LogError> {
59        let mut lines = s.lines();
60        let Some(header) = lines.next() else {
61            return Ok(Vec::new());
62        };
63        // Kafka's `toInt` takes the whole line, so neither side may carry
64        // padding. A first line that is no integer is a file with no header.
65        let version: i64 = header.parse().map_err(|_| LogError::MissingFormatVersion {
66            artifact: ARTIFACT,
67            path: path.to_path_buf(),
68        })?;
69        if version != i64::from(LEADER_EPOCH_CHECKPOINT_VERSION) {
70            return Err(LogError::UnsupportedFormatVersion {
71                artifact: ARTIFACT,
72                path: path.to_path_buf(),
73                found: version,
74            });
75        }
76        let count: usize = lines
77            .next()
78            .and_then(|l| l.trim().parse().ok())
79            .unwrap_or(0);
80        // Do NOT pre-size from the untrusted `count`: a corrupt or hostile
81        // checkpoint (local dir, or bytes restored from tiered storage) could
82        // declare a huge count and trigger a multi-GB allocation before the
83        // bounded `lines.take(count)` loop ever runs. `count` is used only to
84        // bound the number of rows read; the Vec grows as entries are parsed.
85        // Matches Kafka's CheckpointFile, which reads entries line-by-line.
86        let mut out: Vec<EpochEntry> = Vec::new();
87        for line in lines.take(count) {
88            let mut parts = line.split_whitespace();
89            let epoch = parse_column(&mut parts, line)?;
90            let start_offset = parse_column(&mut parts, line)?;
91            let entry = EpochEntry {
92                epoch: LeaderEpoch(epoch),
93                start_offset: Offset(start_offset),
94            };
95            if out
96                .last()
97                .is_some_and(|previous| !is_strict_successor(previous, &entry))
98            {
99                return Err(LogError::Corrupt(format!(
100                    "leader epoch checkpoint rows are not strictly increasing: {line:?}"
101                )));
102            }
103            out.push(entry);
104        }
105        Ok(out)
106    }
107
108    /// Rewrite the whole file atomically. `mutation` calls this after every
109    /// change that altered the entry list.
110    pub(super) fn flush(&self) -> Result<(), LogError> {
111        let mut s = String::new();
112        let _ = writeln!(s, "{LEADER_EPOCH_CHECKPOINT_VERSION}");
113        let _ = writeln!(s, "{}", self.entries.len());
114        for e in &self.entries {
115            let _ = writeln!(s, "{} {}", e.epoch.0, e.start_offset.0);
116        }
117        let tmp = self.path.with_extension("tmp");
118        write_atomic(
119            &*self.io,
120            IoTarget::LeaderEpochCheckpoint,
121            &tmp,
122            &self.path,
123            s.as_bytes(),
124            false,
125        )
126        .map_err(LogError::Io)?;
127        Ok(())
128    }
129}
130
131fn parse_column<T: FromStr>(parts: &mut SplitWhitespace<'_>, line: &str) -> Result<T, LogError> {
132    parts
133        .next()
134        .and_then(|token| token.parse().ok())
135        .ok_or_else(|| LogError::Corrupt(format!("bad checkpoint row: {line:?}")))
136}
137
138#[cfg(test)]
139mod tests {
140    use {
141        super::*,
142        crate::{
143            leader_epoch_checkpoint::test_support::{TempCheckpoint, checkpoint, entry, fresh},
144            test_support::TempPath,
145        },
146        std::fs,
147    };
148
149    /// The golden file: Kafka's text layout, version line `0` first.
150    const GOLDEN: &str = "0\n3\n0 0\n1 50\n2 100\n";
151
152    fn golden_entries() -> Vec<EpochEntry> {
153        [(0, 0), (1, 50), (2, 100)]
154            .into_iter()
155            .map(|(epoch, start_offset)| EpochEntry {
156                epoch: LeaderEpoch(epoch),
157                start_offset: Offset(start_offset),
158            })
159            .collect()
160    }
161
162    #[test]
163    fn round_trip_byte_compat_format() {
164        let TempCheckpoint {
165            dir: _d,
166            checkpoint: c,
167        } = checkpoint(&[entry(0, 0), entry(1, 50), entry(2, 100)]);
168        let s = fs::read_to_string(&c.path).unwrap();
169        assert2::assert!(s == GOLDEN);
170    }
171
172    #[test]
173    fn the_golden_file_decodes_to_its_entries() {
174        let TempPath { dir: _d, path } = fresh();
175        fs::write(&path, GOLDEN).unwrap();
176        assert2::assert!(LeaderEpochCheckpoint::open(path).unwrap().entries() == golden_entries());
177    }
178
179    /// An empty file and a header with no count line both hold no entries,
180    /// as Kafka's reader returns an empty list when a line runs out.
181    #[yare::parameterized(
182        empty_file = { "" },
183        version_without_count = { "0\n" },
184    )]
185    fn a_file_that_ends_early_has_no_entries(contents: &str) {
186        let TempPath { dir: _d, path } = fresh();
187        fs::write(&path, contents).unwrap();
188        let c = LeaderEpochCheckpoint::open(path).unwrap();
189        assert2::assert!(c.entries() == &[][..], "{contents:?}");
190    }
191
192    /// A header that is an integer other than 0 is an unrecognized version,
193    /// as in Kafka. A first line that is no integer is a file without a
194    /// header.
195    #[yare::parameterized(
196        version_1 = { "1\n3\n0 0\n1 50\n2 100\n", Some(1) },
197        version_negative = { "-1\n0\n", Some(-1) },
198        version_overflow = { "99999999999\n0\n", Some(99_999_999_999) },
199        version_3 = { "3\n0 0\n1 50\n2 100\n", Some(3) },
200        missing_count = { "0 0\n1 50\n", None },
201        leading_space = { " 0\n0\n", None },
202        nonnumeric_version = { "v0\n0\n", None },
203    )]
204    fn open_refuses_an_unknown_or_missing_version(contents: &str, found: Option<i64>) {
205        let TempPath { dir: _d, path } = fresh();
206        fs::write(&path, contents).unwrap();
207        let error = LeaderEpochCheckpoint::open(path.clone()).unwrap_err();
208        let got = match error {
209            LogError::UnsupportedFormatVersion {
210                artifact: ARTIFACT,
211                path: named,
212                found,
213            } if named == path => Some(Some(found)),
214            LogError::MissingFormatVersion {
215                artifact: ARTIFACT,
216                path: named,
217            } if named == path => Some(None),
218            _ => None,
219        };
220        assert2::assert!(got == Some(found), "{contents:?}");
221    }
222
223    #[test]
224    fn missing_file_yields_empty() {
225        let TempPath { dir: _d, path } = fresh();
226        let c = LeaderEpochCheckpoint::open(path).unwrap();
227        assert2::assert!(c.entries() == &[][..]);
228        assert2::assert!(c.latest_epoch() == None);
229    }
230
231    #[test]
232    fn open_rejects_out_of_order_checkpoint_rows() {
233        let TempPath { dir: _d, path } = fresh();
234        fs::write(&path, "0\n4\n0 0\n5 100\n2 50\n4 80\n").unwrap();
235
236        let error = LeaderEpochCheckpoint::open(path).unwrap_err();
237        assert2::assert!(matches!(error, LogError::Corrupt(_)));
238    }
239
240    #[test]
241    fn absurd_declared_count_does_not_over_allocate() {
242        // Hostile/corrupt checkpoint: header declares billions of rows but only
243        // one actual entry line follows. Parsing must not pre-size a giant Vec;
244        // it should grow to fit the real rows and return just those.
245        let s = "0\n9999999999999\n3 42\n";
246        let entries = LeaderEpochCheckpoint::parse(Path::new("checkpoint"), s).unwrap();
247        assert2::assert!(
248            entries
249                == [EpochEntry {
250                    epoch: LeaderEpoch(3),
251                    start_offset: Offset(42),
252                }]
253        );
254        // `lines.take(count)` bounds reads to the available lines, so capacity
255        // stays at the grown size, not the untrusted billions.
256        assert2::assert!(entries.capacity() < 9_999_999_999_999);
257    }
258}