krabka_log/leader_epoch_checkpoint/
mutation.rs

1//! The mutations a live checkpoint accepts -- `append`, `assign`,
2//! `truncate_from_end`, `truncate_from_start` and `clear` -- each of which
3//! persists the file only when it actually changed the entry list. They wrap the pure cores in [`super`]
4//! and mirror Kafka's `LeaderEpochFileCache`.
5
6use {
7    super::{EpochEntry, LeaderEpochCheckpoint, append_to, is_strict_successor, truncate_to},
8    crate::error::LogError,
9    krabka_ids::{LeaderEpoch, Offset},
10    tracing::instrument,
11};
12
13impl LeaderEpochCheckpoint {
14    /// Append `(epoch, start_offset)`. This method is idempotent. A second
15    /// append of an entry with the same epoch does nothing and keeps the
16    /// earliest recorded `start_offset`. The method rewrites the file
17    /// atomically.
18    #[instrument(level = "debug", skip(self), fields(epoch = epoch.0, start_offset = start_offset.0), err)]
19    /// # Errors
20    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
21    pub fn append(&mut self, epoch: LeaderEpoch, start_offset: Offset) -> Result<(), LogError> {
22        if self.entries.iter().any(|entry| entry.epoch == epoch) {
23            return Ok(());
24        }
25        let entry = EpochEntry {
26            epoch,
27            start_offset,
28        };
29        if self
30            .entries
31            .last()
32            .is_some_and(|previous| !is_strict_successor(previous, &entry))
33        {
34            return Err(LogError::InvalidArgument(format!(
35                "leader epoch checkpoint entry ({}, {}) is not strictly after the previous entry",
36                epoch.0, start_offset.0
37            )));
38        }
39        if append_to(&mut self.entries, epoch, start_offset) {
40            self.flush()?;
41        }
42        Ok(())
43    }
44
45    /// Kafka's `LeaderEpochFileCache.assign(epoch, startOffset)`, which
46    /// `UnifiedLog.assignEpochStartOffset` calls when `Partition.makeLeader`
47    /// records a new leader epoch at the log end before anything is written.
48    ///
49    /// It is a no-op when `epoch` is already the latest recorded epoch and
50    /// `start_offset` is at or after that entry's start
51    /// (`isUpdateNeeded`). Otherwise it first drops every trailing entry whose
52    /// epoch is `>= epoch` or whose start is `>= start_offset`
53    /// (`maybeTruncateNonMonotonicEntries`) and then records the new entry, so
54    /// the history stays strictly increasing in both columns. The file is
55    /// rewritten only when the entry list changed, and a failed rewrite
56    /// leaves the entries as they were.
57    ///
58    /// # Errors
59    /// Returns [`LogError::InvalidArgument`] for a negative epoch or start
60    /// offset, as Kafka throws `IllegalArgumentException`, and an I/O error
61    /// when the checkpoint cannot be persisted.
62    #[instrument(level = "debug", skip(self), fields(epoch = epoch.0, start_offset = start_offset.0), err)]
63    pub fn assign(&mut self, epoch: LeaderEpoch, start_offset: Offset) -> Result<(), LogError> {
64        if epoch.0 < 0 || start_offset.0 < 0 {
65            return Err(LogError::InvalidArgument(format!(
66                "invalid leader epoch checkpoint entry ({}, {})",
67                epoch.0, start_offset.0
68            )));
69        }
70        // `isUpdateNeeded` first, so the per-batch call on every append costs
71        // one comparison and no copy of the history.
72        if !assign_needed(&self.entries, epoch, start_offset) {
73            return Ok(());
74        }
75        let previous = self.entries.clone();
76        assign_to(&mut self.entries, epoch, start_offset);
77        self.flush_or_restore(previous)
78    }
79
80    /// Remove epoch entries that begin at or after `end_offset`. This mirrors
81    /// Kafka's LeaderEpochFileCache.truncateFromEnd. The method persists the
82    /// file if anything changed.
83    #[instrument(level = "debug", skip(self), fields(end_offset = end_offset.0), err)]
84    /// # Errors
85    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
86    pub fn truncate_from_end(&mut self, end_offset: Offset) -> Result<(), LogError> {
87        let before = self.entries.len();
88        truncate_to(&mut self.entries, end_offset);
89        if self.entries.len() != before {
90            self.flush()?;
91        }
92        Ok(())
93    }
94
95    /// Kafka's `LeaderEpochFileCache.truncateFromStart(startOffset)`, which
96    /// `UnifiedLog.maybeIncrementLogStartOffset` calls when the log start
97    /// moves up and `LogLoader.load` calls with the checkpointed log start.
98    ///
99    /// Every entry that starts at or below `start_offset` is removed, and the
100    /// newest of them comes back with its start raised to `start_offset`, so
101    /// the epoch that covers the new log start still answers for it and no
102    /// entry points below the log. The call is exclusive: an entry that
103    /// starts exactly at `start_offset` keeps its place. The file is
104    /// rewritten only when an entry was removed, and a failed rewrite leaves
105    /// the entries as they were.
106    ///
107    /// # Errors
108    /// Returns an I/O error when the checkpoint cannot be persisted.
109    #[instrument(level = "debug", skip(self), fields(start_offset = start_offset.0), err)]
110    pub fn truncate_from_start(&mut self, start_offset: Offset) -> Result<(), LogError> {
111        let previous = self.entries.clone();
112        if truncate_start_of(&mut self.entries, start_offset) {
113            self.flush_or_restore(previous)?;
114        }
115        Ok(())
116    }
117
118    fn flush_or_restore(&mut self, previous: Vec<EpochEntry>) -> Result<(), LogError> {
119        self.flush().inspect_err(|_| {
120            // The file still holds `previous`; so must memory.
121            self.entries = previous;
122        })
123    }
124
125    /// Drop every recorded epoch. This mirrors Kafka's
126    /// `LeaderEpochFileCache.clearAndFlush`, which
127    /// `LocalLog.truncateFullyAndStartAt` invokes.
128    ///
129    /// [`crate::Log::reset_to`] uses this method. Once the log is empty, no
130    /// offset has a backing record, so the broker may advertise no epoch. The
131    /// method persists the now-empty file only when it removed something.
132    #[instrument(level = "debug", skip(self), fields(cleared = self.entries.len()), err)]
133    /// # Errors
134    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
135    pub fn clear(&mut self) -> Result<(), LogError> {
136        if self.entries.is_empty() {
137            return Ok(());
138        }
139        self.entries.clear();
140        self.flush()
141    }
142}
143
144/// Kafka's `LeaderEpochFileCache.isUpdateNeeded`: whether `assign` would
145/// change `entries`. Only the latest epoch at or after its own start is a
146/// no-op.
147fn assign_needed(entries: &[EpochEntry], epoch: LeaderEpoch, start_offset: Offset) -> bool {
148    entries
149        .last()
150        .is_none_or(|latest| latest.epoch != epoch || start_offset < latest.start_offset)
151}
152
153/// Pure core of [`LeaderEpochCheckpoint::assign`], for an entry that
154/// [`assign_needed`] admitted: Kafka's `maybeTruncateNonMonotonicEntries`
155/// followed by the put.
156fn assign_to(entries: &mut Vec<EpochEntry>, epoch: LeaderEpoch, start_offset: Offset) {
157    while entries
158        .last()
159        .is_some_and(|last| last.epoch >= epoch || last.start_offset >= start_offset)
160    {
161        entries.pop();
162    }
163    entries.push(EpochEntry {
164        epoch,
165        start_offset,
166    });
167}
168
169/// Pure core of [`LeaderEpochCheckpoint::truncate_from_start`]. Returns
170/// `true` when it removed an entry, so the caller knows that it must flush.
171fn truncate_start_of(entries: &mut Vec<EpochEntry>, start_offset: Offset) -> bool {
172    let removed = entries
173        .iter()
174        .take_while(|entry| entry.start_offset <= start_offset)
175        .count();
176    let Some(newest_removed) = removed
177        .checked_sub(1)
178        .and_then(|index| entries.get(index))
179        .map(|entry| entry.epoch)
180    else {
181        return false;
182    };
183    if removed == 1
184        && entries
185            .first()
186            .is_some_and(|entry| entry.start_offset == start_offset)
187    {
188        // The one entry at the new start is put back unchanged.
189        return false;
190    }
191    entries.splice(..removed, [EpochEntry {
192        epoch: newest_removed,
193        start_offset,
194    }]);
195    true
196}
197
198#[cfg(test)]
199mod tests {
200    use {
201        super::*,
202        crate::{
203            leader_epoch_checkpoint::{
204                EpochEntry,
205                test_support::{TempCheckpoint, check_persisted, checkpoint, entry, fresh},
206            },
207            test_support::TempPath,
208        },
209        std::fs,
210    };
211
212    #[test]
213    fn append_preserves_existing_rows() {
214        let TempPath { dir: _d, path } = fresh();
215        {
216            let mut c = LeaderEpochCheckpoint::open(path.clone()).unwrap();
217            c.append(LeaderEpoch(0), Offset(0)).unwrap();
218        }
219        let mut c2 = LeaderEpochCheckpoint::open(path).unwrap();
220        c2.append(LeaderEpoch(1), Offset(50)).unwrap();
221        assert2::assert!(c2.entries() == &[entry(0, 0), entry(1, 50)]);
222    }
223
224    #[test]
225    fn append_idempotent_for_same_epoch() {
226        let TempCheckpoint {
227            dir: _d,
228            checkpoint: mut c,
229        } = checkpoint(&[entry(0, 0)]);
230        c.append(LeaderEpoch(0), Offset(999)).unwrap(); // ignored; epoch 0 already recorded
231        assert2::assert!(c.entries() == &[entry(0, 0)]);
232    }
233
234    #[yare::parameterized(
235        epoch_out_of_order = { 2, 150 },
236        offset_not_increasing = { 6, 100 },
237    )]
238    fn append_rejects_novel_out_of_order_entries_without_mutating(epoch: i32, offset: i64) {
239        let TempCheckpoint {
240            dir: _d,
241            checkpoint: mut c,
242        } = checkpoint(&[entry(0, 0), entry(5, 100)]);
243        let before = fs::read(&c.path).unwrap();
244
245        let error = c.append(LeaderEpoch(epoch), Offset(offset)).unwrap_err();
246        assert2::assert!(matches!(error, LogError::InvalidArgument(_)));
247
248        assert2::assert!(fs::read(&c.path).unwrap() == before);
249        assert2::assert!(c.entries() == &[entry(0, 0), entry(5, 100)]);
250    }
251
252    #[test]
253    fn append_after_truncation_restores_a_strict_history() {
254        let TempCheckpoint {
255            dir: _d,
256            checkpoint: mut c,
257        } = checkpoint(&[entry(0, 0), entry(5, 100)]);
258        c.truncate_from_end(Offset(100)).unwrap();
259        c.append(LeaderEpoch(2), Offset(50)).unwrap();
260        c.append(LeaderEpoch(4), Offset(80)).unwrap();
261
262        assert2::assert!(
263            c.epoch_and_offset_for(LeaderEpoch(3), Offset(120)) == (LeaderEpoch(2), Offset(80))
264        );
265    }
266
267    /// One row per `LeaderEpochFileCache.assign` outcome, each checked in
268    /// memory and after a reopen of the persisted file.
269    #[yare::parameterized(
270        first = { "first entry", &[][..], 2, 6, &[entry(2, 6)][..] },
271        appends = {
272            "new epoch after the log end appends",
273            &[entry(0, 0), entry(2, 6)][..],
274            3,
275            9,
276            &[entry(0, 0), entry(2, 6), entry(3, 9)][..],
277        },
278        later_same_epoch = {
279            "latest epoch at a later offset is a no-op",
280            &[entry(2, 6)][..],
281            2,
282            7,
283            &[entry(2, 6)][..],
284        },
285        same_offset = {
286            "latest epoch at its own offset is a no-op",
287            &[entry(2, 6)][..],
288            2,
289            6,
290            &[entry(2, 6)][..],
291        },
292        earlier_offset = {
293            "latest epoch at an earlier offset moves its start",
294            &[entry(1, 0), entry(2, 6)][..],
295            2,
296            5,
297            &[entry(1, 0), entry(2, 5)][..],
298        },
299        same_start_new_epoch = {
300            "newer epoch at the same start replaces an unwritten epoch",
301            &[entry(1, 0), entry(2, 6)][..],
302            3,
303            6,
304            &[entry(1, 0), entry(3, 6)][..],
305        },
306        older_epoch = {
307            "older epoch drops every entry it does not strictly follow",
308            &[entry(1, 0), entry(2, 6), entry(3, 8)][..],
309            2,
310            7,
311            &[entry(1, 0), entry(2, 7)][..],
312        },
313    )]
314    fn assign_follows_kafka_leader_epoch_file_cache_assign(
315        name: &str,
316        recorded: &[EpochEntry],
317        epoch: i32,
318        start: i64,
319        expected: &[EpochEntry],
320    ) {
321        let TempCheckpoint {
322            dir: _d,
323            checkpoint: mut c,
324        } = checkpoint(recorded);
325        c.assign(LeaderEpoch(epoch), Offset(start)).unwrap();
326        check_persisted(&c, expected, name);
327    }
328
329    #[yare::parameterized(
330        negative_epoch = { -1, 5 },
331        negative_start = { 2, -1 },
332    )]
333    fn assign_rejects_negative_entries_without_mutating(epoch: i32, start: i64) {
334        let TempCheckpoint {
335            dir: _d,
336            checkpoint: mut c,
337        } = checkpoint(&[entry(1, 0)]);
338        let error = c.assign(LeaderEpoch(epoch), Offset(start)).unwrap_err();
339        assert2::assert!(matches!(error, LogError::InvalidArgument(_)));
340        assert2::assert!(c.entries() == &[entry(1, 0)]);
341    }
342
343    #[test]
344    fn truncate_from_end_removes_entries_at_or_after_end_offset() {
345        let TempCheckpoint {
346            dir: _d,
347            checkpoint: mut c,
348        } = checkpoint(&[entry(1, 0), entry(7, 4)]);
349        c.truncate_from_end(Offset(4)).unwrap();
350        assert2::assert!(c.latest_epoch() == Some(LeaderEpoch(1)));
351        assert2::assert!(c.end_offset_for_epoch(LeaderEpoch(7), Offset(4)) == Offset(-1));
352        assert2::assert!(c.end_offset_for_epoch(LeaderEpoch(1), Offset(4)) == Offset(4));
353        // Persisted: a reopen sees only epoch 1.
354        let reopened = LeaderEpochCheckpoint::open(c.path.clone()).unwrap();
355        assert2::assert!(reopened.latest_epoch() == Some(LeaderEpoch(1)));
356    }
357
358    /// One row per `LeaderEpochFileCache.truncateFromStart` outcome, over the
359    /// history `(0, 5) (1, 10) (2, 15)`, each checked in memory and after a
360    /// reopen of the persisted file.
361    #[yare::parameterized(
362        empty_cache = { "an empty cache stays empty", &[][..], 7, &[][..] },
363        below = {
364            "a start below every entry changes nothing",
365            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
366            4,
367            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
368        },
369        on_first = {
370            "a start on the first entry keeps it",
371            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
372            5,
373            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
374        },
375        inside_first = {
376            "a start inside the first epoch raises its start",
377            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
378            8,
379            &[entry(0, 8), entry(1, 10), entry(2, 15)][..],
380        },
381        on_later = {
382            "a start on a later entry drops every earlier one",
383            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
384            10,
385            &[entry(1, 10), entry(2, 15)][..],
386        },
387        past_all = {
388            "a start past every entry keeps the latest epoch at the start",
389            &[entry(0, 5), entry(1, 10), entry(2, 15)][..],
390            17,
391            &[entry(2, 17)][..],
392        },
393    )]
394    fn truncate_from_start_follows_kafka_leader_epoch_file_cache(
395        name: &str,
396        recorded: &[EpochEntry],
397        start: i64,
398        expected: &[EpochEntry],
399    ) {
400        let TempCheckpoint {
401            dir: _d,
402            checkpoint: mut c,
403        } = checkpoint(recorded);
404        c.truncate_from_start(Offset(start)).unwrap();
405        check_persisted(&c, expected, name);
406    }
407
408    #[test]
409    fn clear_removes_all_entries_and_persists_empty() {
410        let TempCheckpoint {
411            dir: _d,
412            checkpoint: mut c,
413        } = checkpoint(&[entry(1, 0), entry(2, 50)]);
414        c.clear().unwrap();
415        assert2::assert!(c.entries() == &[][..]);
416        assert2::assert!(c.latest_epoch() == None);
417        // Persisted: a reopen sees no entries.
418        let reopened = LeaderEpochCheckpoint::open(c.path.clone()).unwrap();
419        assert2::assert!(reopened.entries().is_empty());
420    }
421
422    #[test]
423    fn clear_on_empty_cache_skips_flush_and_writes_no_file() {
424        let TempPath { dir: _d, path } = fresh();
425        let mut c = LeaderEpochCheckpoint::open(path.clone()).unwrap();
426        c.clear().unwrap();
427        assert2::assert!(c.entries().is_empty());
428        // The early-return skips the flush for an already-empty cache, so no
429        // checkpoint file is written. A forced-`false` empty-guard would flush
430        // an empty file here instead.
431        assert2::assert!(!path.exists());
432    }
433}
krabka_log/leader_epoch_checkpoint/mutation.rs

krabka_log/leader_epoch_checkpoint/
mutation.rs
