krabka_verified/
leader_epoch.rs

1//! KIP-101/320 leader-epoch divergence lookup kernel: Kafka's
2//! `LeaderEpochFileCache.endOffsetFor`.
3//!
4//! The log crate keeps the checkpoint entries in strictly increasing epoch and
5//! start-offset order. This kernel answers the `(found_epoch, end_offset)`
6//! pair that an `OffsetForLeaderEpoch` row carries and that a `Fetch` row's
7//! `diverging_epoch` is built from.
8
9use {
10    creusot_std::prelude::*,
11    krabka_ids::{LeaderEpoch, Offset},
12};
13
14/// Kafka `LeaderEpochFileCache.UNDEFINED_EPOCH`.
15const UNDEFINED_EPOCH: LeaderEpoch = LeaderEpoch(-1);
16/// Kafka `LeaderEpochFileCache.UNDEFINED_EPOCH_OFFSET`.
17const UNDEFINED_EPOCH_OFFSET: Offset = Offset(-1);
18
19/// One leader epoch and the offset where it begins.
20#[cfg_attr(not(creusot), derive(Debug, Clone, Copy, PartialEq, Eq, Hash))]
21pub struct EpochEntry {
22    pub epoch: LeaderEpoch,
23    pub start_offset: Offset,
24}
25
26// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
27#[cfg(creusot)]
28#[cfg_attr(test, mutants::skip)]
29#[logic(open)]
30/// `h` is the position of Kafka's `epochs.higherEntry(requested)`: the least
31/// recorded epoch strictly above `requested`. Over entries in increasing
32/// epoch order that is the first entry above `requested`, and the entry just
33/// before it, if any, is `epochs.floorEntry(requested)`.
34pub fn higher_entry_at(entries: Seq<EpochEntry>, requested: Int, h: Int) -> bool {
35    pearlite! {
36        0 <= h && h < entries.len()
37            && entries[h].epoch.0@ > requested
38            && forall<j: Int> 0 <= j && j < h ==> entries[j].epoch.0@ <= requested
39    }
40}
41
42// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
43#[cfg(creusot)]
44#[cfg_attr(test, mutants::skip)]
45#[logic(open)]
46/// Kafka's `epochs.higherEntry(requested)` is `null`: no recorded epoch is
47/// above `requested`. An empty cache has none.
48pub fn no_higher_entry(entries: Seq<EpochEntry>, requested: Int) -> bool {
49    pearlite! {
50        forall<i: Int> 0 <= i && i < entries.len() ==> entries[i].epoch.0@ <= requested
51    }
52}
53
54// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
55#[cfg(creusot)]
56#[cfg_attr(test, mutants::skip)]
57#[logic(open)]
58/// Kafka's `LeaderEpochFileCache.endOffsetFor(requestedEpoch, logEndOffset)`
59/// case table, in Kafka's order, as the relation between the inputs and the
60/// answer `(found, end)`:
61///
62/// 1. `requested == UNDEFINED_EPOCH` → `(UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)`.
63/// 2. `requested` is the latest recorded epoch → `(requested, log_end)`.
64/// 3. no `higherEntry` (this includes an empty cache) →
65///    `(UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)`.
66/// 4. a `higherEntry` but no `floorEntry` (`requested` is below every
67///    recorded epoch) → `(requested, higher.start_offset)`.
68/// 5. both (an exact older epoch, or a gap epoch) →
69///    `(floor.epoch, higher.start_offset)`.
70///
71/// The relation is functional: every input has exactly one answer.
72pub fn kafka_end_offset_for(
73    entries: Seq<EpochEntry>,
74    requested: Int,
75    log_end: Int,
76    found: Int,
77    end: Int,
78) -> bool {
79    pearlite! {
80        if requested == -1 {
81            found == -1 && end == -1
82        } else if entries.len() > 0 && entries[entries.len() - 1].epoch.0@ == requested {
83            found == requested && end == log_end
84        } else if no_higher_entry(entries, requested) {
85            found == -1 && end == -1
86        } else {
87            exists<h: Int> higher_entry_at(entries, requested, h)
88                && end == entries[h].start_offset.0@
89                && found == if h == 0 { requested } else { entries[h - 1].epoch.0@ }
90        }
91    }
92}
93
94/// Resolve a requested leader epoch to `(found_epoch, end_offset)`, exactly as
95/// Kafka's `LeaderEpochFileCache.endOffsetFor` does (see
96/// `kafka_end_offset_for` for the case table).
97///
98/// `(UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)` = `(-1, -1)` is the "cannot
99/// place this epoch" answer: `UnifiedLog.endOffsetForEpoch` turns it into
100/// `None`, which `Partition.lastOffsetForLeaderEpoch` answers with the
101/// schema defaults `(-1, -1)` and `Partition.readRecords` answers with
102/// `OFFSET_OUT_OF_RANGE`. Those two uses are the caller's.
103///
104/// Entries must be strictly increasing in both fields.
105#[requires(forall<i: Int, j: Int> 0 <= i && i < j && j < [email protected]()
106    ==> entries@[i].epoch.0@ < entries@[j].epoch.0@
107        && entries@[i].start_offset.0@ < entries@[j].start_offset.0@)]
108#[ensures(kafka_end_offset_for(
109    entries@, requested_epoch.0@, log_end_offset.0@, result.0.0@, result.1.0@))]
110#[must_use]
111pub fn epoch_and_offset_for_entries(
112    entries: &[EpochEntry],
113    requested_epoch: LeaderEpoch,
114    log_end_offset: Offset,
115) -> (LeaderEpoch, Offset) {
116    if requested_epoch.0 == UNDEFINED_EPOCH.0 {
117        return (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
118    }
119
120    // `i` becomes the position of Kafka's `higherEntry`, or `len` if none.
121    let mut i = 0usize;
122    #[invariant(i@ <= [email protected]())]
123    #[invariant(forall<j: Int> 0 <= j && j < i@
124        ==> entries@[j].epoch.0@ <= requested_epoch.0@)]
125    #[variant([email protected]() - i@)]
126    while i < entries.len()
127        && matches!(entries.get(i), Some(entry) if entry.epoch.0 <= requested_epoch.0)
128    {
129        i += 1;
130    }
131
132    if i == entries.len() {
133        if i > 0 && matches!(entries.get(i - 1), Some(last) if last.epoch.0 == requested_epoch.0) {
134            return (requested_epoch, log_end_offset);
135        }
136        return (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
137    }
138    proof_assert!(higher_entry_at(entries@, requested_epoch.0@, i@));
139    proof_assert!(entries@[[email protected]() - 1].epoch.0@ >= entries@[i@].epoch.0@);
140    let Some(higher) = entries.get(i) else {
141        return (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
142    };
143    if i == 0 {
144        return (requested_epoch, higher.start_offset);
145    }
146    let Some(previous) = entries.get(i - 1) else {
147        return (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
148    };
149    (previous.epoch, higher.start_offset)
150}
151
152#[cfg(test)]
153mod tests {
154    use {
155        super::*,
156        proptest::{collection, prelude::*},
157        std::collections::BTreeMap,
158    };
159
160    /// Kafka's own `endOffsetFor`, written against a `TreeMap`-shaped index
161    /// (`higherEntry` / `floorEntry`) rather than the kernel's linear scan.
162    fn tree_map_oracle(
163        entries: &[EpochEntry],
164        requested_epoch: LeaderEpoch,
165        log_end_offset: Offset,
166    ) -> (LeaderEpoch, Offset) {
167        let epochs: BTreeMap<i32, EpochEntry> = entries.iter().map(|e| (e.epoch.0, *e)).collect();
168        if requested_epoch == UNDEFINED_EPOCH {
169            return (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
170        }
171        if epochs.last_key_value().map(|(epoch, _)| *epoch) == Some(requested_epoch.0) {
172            return (requested_epoch, log_end_offset);
173        }
174        let higher = epochs
175            .range(requested_epoch.0.saturating_add(1)..)
176            .next()
177            .map(|(_, entry)| *entry);
178        let floor = epochs
179            .range(..=requested_epoch.0)
180            .next_back()
181            .map(|(_, entry)| *entry);
182        match (higher, floor) {
183            (None, _) => (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
184            (Some(higher), None) => (requested_epoch, higher.start_offset),
185            (Some(higher), Some(floor)) => (floor.epoch, higher.start_offset),
186        }
187    }
188
189    const fn entry(epoch: i32, start_offset: i64) -> EpochEntry {
190        EpochEntry {
191            epoch: LeaderEpoch(epoch),
192            start_offset: Offset(start_offset),
193        }
194    }
195
196    /// One row per `endOffsetFor` branch, plus the gap-epoch variants.
197    #[yare::parameterized(
198        empty_cache = { "empty cache", &[][..], 0, 9, (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET) },
199        undefined_requested_empty = {
200            "undefined requested epoch, empty cache",
201            &[][..],
202            -1,
203            9,
204            (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
205        },
206        undefined_requested = {
207            "undefined requested epoch",
208            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
209            -1,
210            200,
211            (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
212        },
213        latest_epoch = {
214            "latest epoch answers the log end offset",
215            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
216            5,
217            200,
218            (LeaderEpoch(5), Offset(200)),
219        },
220        above_latest = {
221            "requested epoch above the latest",
222            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
223            7,
224            200,
225            (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
226        },
227        below_first = {
228            "below the first epoch: requested epoch, first start",
229            &[entry(3, 30), entry(4, 40)][..],
230            1,
231            200,
232            (LeaderEpoch(1), Offset(30)),
233        },
234        exact_older = {
235            "exact older epoch: that epoch, next start",
236            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
237            2,
238            200,
239            (LeaderEpoch(2), Offset(100)),
240        },
241        gap_epoch = {
242            "gap epoch: floor epoch, next start",
243            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
244            3,
245            200,
246            (LeaderEpoch(2), Offset(100)),
247        },
248        gap_below_latest = {
249            "gap epoch below the latest: floor epoch 2, start of 5",
250            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
251            4,
252            200,
253            (LeaderEpoch(2), Offset(100)),
254        },
255        gap_above_first = {
256            "gap epoch just above the first: floor epoch 0, start of 2",
257            &[entry(0, 0), entry(2, 50), entry(5, 100)][..],
258            1,
259            200,
260            (LeaderEpoch(0), Offset(50)),
261        },
262    )]
263    fn lookup_matches_kafka_end_offset_for_rows(
264        name: &str,
265        entries: &[EpochEntry],
266        requested: i32,
267        log_end: i64,
268        expected: (LeaderEpoch, Offset),
269    ) {
270        assert2::check!(
271            epoch_and_offset_for_entries(entries, LeaderEpoch(requested), Offset(log_end))
272                == expected,
273            "case {name}"
274        );
275    }
276
277    proptest! {
278        #[test]
279        fn lookup_matches_tree_map_oracle(
280            epochs in collection::btree_set(0i32..100, 0..32),
281            requested in -3i32..110,
282            log_end in 0i64..10_000,
283        ) {
284            let entries = epochs
285                .into_iter()
286                .enumerate()
287                .map(|(i, epoch)| {
288                    entry(epoch, i64::try_from(i).expect("epoch set length is bounded to 32") * 10)
289                })
290                .collect::<Vec<_>>();
291            prop_assert_eq!(
292                epoch_and_offset_for_entries(
293                    &entries,
294                    LeaderEpoch(requested),
295                    Offset(log_end),
296                ),
297                tree_map_oracle(&entries, LeaderEpoch(requested), Offset(log_end))
298            );
299        }
300    }
301}
krabka_verified/leader_epoch.rs

krabka_verified/
leader_epoch.rs
