krabka_log/segment/
timestamp_scan.rs

1//! Timestamp lookups over a segment: the sparse time-index floor, the forward
2//! scan that refines it, and the restore of `max_timestamp` after a no-scan
3//! open.
4//!
5//! Kafka's `LogSegment.findOffsetByTimestamp` needs a record scan after the
6//! index lookup because the index is sparse, and all of the windowing that the
7//! scan needs lives here.
8
9use {
10    super::{Segment, header_walk::HeaderWalk as _},
11    crate::{
12        config::DEFAULT_TIMESTAMP_SCAN_WINDOW, error::LogError, record_limit::check_records_read,
13    },
14    krabka_ids::Offset,
15    krabka_protocol::records::{RecordBatch, TimestampType},
16    krabka_units::prelude::{ByteSize, ByteSizeExt, bytes},
17    krabka_verified::timestamp::{first_timestamp_record_index, timestamp_record_time},
18    std::ops::ControlFlow,
19};
20
21impl Segment {
22    /// Read the first record's timestamp once per segment for the roll check.
23    /// Empty or failed reads stay uncached so the next append can retry.
24    pub(crate) fn first_record_timestamp(&mut self) -> Option<i64> {
25        if self.first_timestamp.is_none() {
26            self.first_timestamp = self.offset_for_timestamp(i64::MIN).map(|(_, ts)| ts);
27        }
28        self.first_timestamp
29    }
30
31    /// Absolute offset and record timestamp of the first record in this
32    /// segment whose timestamp is `>= target_ts`.
33    ///
34    /// This method takes a floor position from the sparse time index, then
35    /// scans `.log` batches forward. The index is sparse, so an exact answer
36    /// needs that scan after the index lookup. This matches Kafka's
37    /// `LogSegment.findOffsetByTimestamp`. The result is `None` when no
38    /// record in this segment qualifies.
39    ///
40    /// The scan applies no `max.decompressed.message.bytes`: this is the lookup
41    /// the log's own bookkeeping uses, to read a segment's first timestamp. A
42    /// `ListOffsets` answer goes through
43    /// `Segment::offset_for_timestamp_with_window`, which takes the limit.
44    #[must_use]
45    pub fn offset_for_timestamp(&self, target_ts: i64) -> Option<(Offset, i64)> {
46        // With no limit the scan has nothing to refuse.
47        self.offset_for_timestamp_with_window(target_ts, DEFAULT_TIMESTAMP_SCAN_WINDOW, None)
48            .ok()
49            .flatten()
50    }
51
52    /// [`Segment::offset_for_timestamp`] over `scan_window`-sized reads, with
53    /// `limit` as Kafka trunk's `max.decompressed.message.bytes`.
54    ///
55    /// Kafka's `FileRecords.searchForTimestamp` decompresses a batch only when
56    /// its max timestamp reaches `target_ts`, and reads it only as far as the
57    /// record it returns, so a `limit` fails the lookup only for a record that
58    /// Kafka's iterator would have read.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`LogError::RecordTooLarge`] for such a record of a compressed
63    /// batch. A scan with no `limit` never fails; an unreadable segment
64    /// answers `None`.
65    pub(crate) fn offset_for_timestamp_with_window(
66        &self,
67        target_ts: i64,
68        scan_window: ByteSize,
69        limit: Option<ByteSize>,
70    ) -> Result<Option<(Offset, i64)>, LogError> {
71        self.offset_for_timestamp_with_window_from(target_ts, scan_window, limit, self.base_offset)
72    }
73
74    /// The bounded timestamp scan, excluding records below `minimum`, even
75    /// when the first decoded batch covers that floor.
76    ///
77    /// # Errors
78    ///
79    /// Returns [`LogError::RecordTooLarge`] when the scan reads a compressed
80    /// record above `limit`.
81    pub(crate) fn offset_for_timestamp_with_window_from(
82        &self,
83        target_ts: i64,
84        scan_window: ByteSize,
85        limit: Option<ByteSize>,
86        minimum: Offset,
87    ) -> Result<Option<(Offset, i64)>, LogError> {
88        let Some(scan_from) = self.timestamp_scan_floor(target_ts) else {
89            return Ok(None);
90        };
91        self.scan_from_floor_windowed_from(
92            scan_from.max(minimum),
93            scan_window,
94            target_ts,
95            limit,
96            minimum,
97        )
98    }
99
100    /// [`Segment::offset_of_max_timestamp_with_window`] over the default window
101    /// with no limit, for the tests that only care where the answer is.
102    #[cfg(test)]
103    pub(crate) fn offset_of_max_timestamp(&self) -> Option<(Offset, i64)> {
104        self.offset_of_max_timestamp_with_window(DEFAULT_TIMESTAMP_SCAN_WINDOW, None)
105            .expect("a lookup with no limit refuses nothing")
106    }
107
108    /// Absolute offset and timestamp of the record that carries this
109    /// segment's `max_timestamp`, read in `scan_window`-sized pieces, with
110    /// `limit` as Kafka trunk's `max.decompressed.message.bytes`: Kafka's
111    /// `RecordBatch.offsetOfMaxTimestamp` decompresses the batch that holds the
112    /// maximum, up to the record it returns.
113    ///
114    /// Ties resolve to the earliest offset, as in Kafka. The result is `None`
115    /// for an empty segment. This method starts before the first occurrence
116    /// of the maximum, then scans forward for the first record whose
117    /// timestamp equals the segment maximum.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`LogError::RecordTooLarge`] for a record of that batch that the
122    /// `limit` refuses. A scan with no `limit` never fails.
123    pub(crate) fn offset_of_max_timestamp_with_window(
124        &self,
125        scan_window: ByteSize,
126        limit: Option<ByteSize>,
127    ) -> Result<Option<(Offset, i64)>, LogError> {
128        if self.max_timestamp == i64::MIN {
129            let found = self.scan_max_timestamp_windowed(scan_window);
130            // The unknown-maximum scan reads every record, so the limit is
131            // checked on the batch that holds the maximum the way the
132            // known-maximum path checks it: by looking the maximum up again.
133            if let Some((_, timestamp)) = found
134                && limit.is_some()
135            {
136                self.scan_from_floor_windowed(self.base_offset, scan_window, timestamp, limit)?;
137            }
138            return Ok(found);
139        }
140        let Some(scan_from) = self.timestamp_scan_floor(self.max_timestamp) else {
141            return Ok(None);
142        };
143        // Equality against `max_timestamp` is safe because Kafka's batch
144        // `max_timestamp` is always a real record timestamp (the largest
145        // among the batch's records), so some record's timestamp equals
146        // the segment max exactly.
147        self.scan_from_floor_windowed(scan_from, scan_window, self.max_timestamp, limit)
148    }
149
150    fn timestamp_scan_floor(&self, target: i64) -> Option<Offset> {
151        self.base_offset
152            .0
153            .checked_add(i64::from(self.time_index.scan_start(target)))
154            .map(Offset)
155    }
156
157    /// Recover the maximum timestamp for a sealed segment opened through the
158    /// no-scan path. Those segments intentionally keep `max_timestamp` at its
159    /// unknown sentinel, so KIP-734 `MAX_TIMESTAMP` must derive the answer
160    /// from records instead of treating the segment as empty.
161    fn scan_max_timestamp_windowed(&self, window_size: ByteSize) -> Option<(Offset, i64)> {
162        let mut cursor = self.base_offset;
163        let mut window = window_size.max(bytes(1));
164        let mut best: Option<(Offset, i64)> = None;
165        loop {
166            if cursor > self.last_offset {
167                return best;
168            }
169            let batches = self.read(cursor, window).ok()?;
170            let Some(last) = batches.last() else {
171                let current = u32::try_from(window.bytes_u64()).ok()?;
172                window = bytes(krabka_verified::timestamp_scan_window(current)?);
173                continue;
174            };
175            for batch in &batches {
176                let records = Self::timestamp_records(batch)?;
177                let timestamps: Vec<_> = records.iter().map(|(_, timestamp)| *timestamp).collect();
178                if let Some(&candidate) = krabka_verified::earliest_max_timestamp_index(&timestamps)
179                    .and_then(|index| records.get(index))
180                    && best.is_none_or(|(_, best_timestamp)| candidate.1 > best_timestamp)
181                {
182                    best = Some((Offset(candidate.0), candidate.1));
183                }
184            }
185            cursor = Offset(krabka_verified::timestamp_scan_next(
186                cursor.0,
187                last.base_offset,
188                last.last_offset_delta,
189            )?);
190        }
191    }
192
193    /// Window-size-parameterized core of [`Segment::scan_from_floor`]. It is
194    /// a separate function so that tests can force multi-window scans with a
195    /// tiny window.
196    ///
197    /// Termination: each iteration does one of three things. It returns a
198    /// match. It returns `None` because `cursor > last_offset`. Or it decodes
199    /// at least one full batch and advances `cursor` strictly past that batch.
200    /// `read` caps reads at `max_bytes` and, unlike `read_raw`, gives no
201    /// anti-stall guarantee. A single batch larger than the window therefore
202    /// decodes to an empty `Vec`. This function detects that case, an empty
203    /// result while `cursor` is still within the segment, and doubles the
204    /// window before it tries again. The window is therefore bounded by the
205    /// largest batch, not by the whole tail.
206    ///
207    /// `limit` is Kafka trunk's `max.decompressed.message.bytes`. A batch whose
208    /// max timestamp reaches `target_ts` is the one Kafka's scan decompresses,
209    /// and it reads that batch up to the record it returns, or to its end when
210    /// none qualifies, so the limit is checked over exactly those records. A
211    /// scan that cannot proceed (an unreadable segment, an offset that
212    /// overflows) answers `None`, as it did before the limit existed.
213    fn scan_from_floor_windowed(
214        &self,
215        floor_offset: Offset,
216        window_size: ByteSize,
217        target_ts: i64,
218        limit: Option<ByteSize>,
219    ) -> Result<Option<(Offset, i64)>, LogError> {
220        self.scan_from_floor_windowed_from(
221            floor_offset,
222            window_size,
223            target_ts,
224            limit,
225            self.base_offset,
226        )
227    }
228
229    fn scan_from_floor_windowed_from(
230        &self,
231        floor_offset: Offset,
232        window_size: ByteSize,
233        target_ts: i64,
234        limit: Option<ByteSize>,
235        minimum: Offset,
236    ) -> Result<Option<(Offset, i64)>, LogError> {
237        let mut cursor = floor_offset;
238        let mut window = window_size.max(bytes(1));
239        loop {
240            if cursor > self.last_offset {
241                return Ok(None);
242            }
243            let Ok(batches) = self.read(cursor, window) else {
244                return Ok(None);
245            };
246            let Some(last) = batches.last() else {
247                // The batch at `cursor` is larger than the window, so it
248                // could not be fully decoded. Grow the window and retry
249                // the same cursor; bounded by the largest batch size.
250                let Some(grown) = u32::try_from(window.bytes_u64())
251                    .ok()
252                    .and_then(krabka_verified::timestamp_scan_window)
253                else {
254                    return Ok(None);
255                };
256                window = bytes(grown);
257                continue;
258            };
259            for batch in &batches {
260                let Some(records) = Self::timestamp_records(batch) else {
261                    return Ok(None);
262                };
263                let found = first_timestamp_record_index(&records, minimum.0, target_ts);
264                if batch.max_timestamp >= target_ts {
265                    let read = found.map_or(batch.records.len(), |index| index + 1);
266                    check_records_read(batch, read, limit)?;
267                }
268                if let Some(&(offset, timestamp)) = found.and_then(|index| records.get(index)) {
269                    return Ok(Some((Offset(offset), timestamp)));
270                }
271            }
272            // No match in this window; resume just past the last batch
273            // read. `read` includes the batch covering `cursor`, so
274            // `last_read` >= cursor and the cursor strictly advances.
275            let Some(next) = krabka_verified::timestamp_scan_next(
276                cursor.0,
277                last.base_offset,
278                last.last_offset_delta,
279            ) else {
280                return Ok(None);
281            };
282            cursor = Offset(next);
283        }
284    }
285
286    /// The `(offset, timestamp)` pair of every record in `batch`, as a reader
287    /// sees them.
288    ///
289    /// Under `CreateTime` a record's timestamp is the batch's `base_timestamp`
290    /// plus its own delta, which is how the v2 format stores it. Under
291    /// `LogAppendTime` every record carries the batch's `max_timestamp`
292    /// instead: the broker stamped that one field and left `base_timestamp`
293    /// and the deltas as the producer wrote them, and Kafka's
294    /// `DefaultRecord.timestamp()` substitutes `maxTimestamp` for the whole
295    /// batch whenever the timestamp-type bit is set. A scan that read the
296    /// deltas on such a batch would answer `ListOffsets` in producer time on a
297    /// topic whose whole point is that it answers in append time.
298    fn timestamp_records(batch: &RecordBatch) -> Option<Vec<(i64, i64)>> {
299        let append_time = (batch.attributes.timestamp_type() == TimestampType::LogAppendTime)
300            .then_some(batch.max_timestamp);
301        batch
302            .records
303            .iter()
304            .map(|record| {
305                let timestamp = timestamp_record_time(
306                    batch.base_timestamp,
307                    record.timestamp_delta,
308                    append_time,
309                )?;
310                krabka_verified::timestamp_record_coordinates(
311                    batch.base_offset,
312                    record.offset_delta,
313                    timestamp,
314                    0,
315                )
316            })
317            .collect()
318    }
319
320    /// Restore `max_timestamp` for a segment loaded without a scan.
321    ///
322    /// The no-scan loader leaves the field at its unknown sentinel. Retention compares
323    /// that sentinel against the age cutoff, so without this call every
324    /// reopened segment looks older than any window and the first
325    /// [`Log::tick`](crate::Log::tick) after a restart deletes all of them.
326    ///
327    /// Kafka's `LogSegment` reads `maxTimestampSoFar` back from the time
328    /// index, and this method starts there. The sparse index lags the writes,
329    /// though: its newest entry holds the running maximum as of the last
330    /// *indexed* batch, and every batch appended after it is unaccounted for.
331    /// The method therefore also walks the batch headers from the newest
332    /// offset-index entry to the end of the file. That walk is bounded by
333    /// `index_interval` plus one batch, and it reads no record body.
334    ///
335    /// # Errors
336    ///
337    /// Returns [`LogError::Io`] when the `.log` file cannot be read.
338    pub fn restore_max_timestamp(&mut self) -> Result<(), LogError> {
339        // The time index's newest entry pairs the running maximum with the
340        // last offset of the batch that set it, which is where the restored
341        // maximum's batch starts too.
342        let (mut max_timestamp, mut max_timestamp_offset) =
343            self.time_index
344                .last_entry()
345                .map_or((i64::MIN, self.base_offset - 1), |entry| {
346                    (
347                        entry.timestamp,
348                        self.base_offset + i64::from(entry.relative_offset),
349                    )
350                });
351        let scan_from = self
352            .offset_index
353            .last_entry()
354            .map_or(0, |entry| u64::from(entry.position));
355        self.walk_batch_headers(scan_from, |view| {
356            if view.max_timestamp > max_timestamp {
357                max_timestamp = view.max_timestamp;
358                max_timestamp_offset = view.last_offset;
359            }
360            ControlFlow::Continue(())
361        })?;
362        if max_timestamp > self.max_timestamp {
363            self.max_timestamp = max_timestamp;
364            self.max_timestamp_offset = max_timestamp_offset;
365        }
366        Ok(())
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use {
373        super::*,
374        crate::segment::test_support::{
375            DENSE_INDEX, FIVE_OFFSET_TIMESTAMP_BATCHES, ONE_RECORD_BATCH, TempSegment,
376            sample_batch, seeded_fixture, segment_at, test_segment,
377        },
378        assert2::check,
379        bytes::Bytes,
380        krabka_protocol::records::{Record, RecordBatch},
381        krabka_units::prelude::kibibytes,
382        tempfile::tempdir,
383    };
384
385    fn one_record_batch() -> RecordBatch {
386        sample_batch()
387            .offset(ONE_RECORD_BATCH.offset.0)
388            .records(ONE_RECORD_BATCH.records.0)
389            .timestamp(ONE_RECORD_BATCH.timestamp.0)
390            .call()
391    }
392
393    #[yare::parameterized(
394        none_create_time = {
395            krabka_compression::CompressionType::None,
396            TimestampType::CreateTime,
397        },
398        none_log_append_time = {
399            krabka_compression::CompressionType::None,
400            TimestampType::LogAppendTime,
401        },
402        lz4_create_time = {
403            krabka_compression::CompressionType::Lz4,
404            TimestampType::CreateTime,
405        },
406        lz4_log_append_time = {
407            krabka_compression::CompressionType::Lz4,
408            TimestampType::LogAppendTime,
409        },
410    )]
411    fn first_record_timestamp_survives_appends_and_reloads_after_truncation(
412        codec: krabka_compression::CompressionType,
413        timestamp_type: TimestampType,
414    ) {
415        let TempSegment {
416            dir,
417            segment: mut seg,
418        } = test_segment();
419        assert2::assert!(seg.first_record_timestamp().is_none());
420        let mut batch = sample_batch().records(2).call();
421        batch.records[0].timestamp_delta = -5;
422        batch.attributes = batch
423            .attributes
424            .with_compression(codec)
425            .with_timestamp_type(timestamp_type);
426        let expected = if timestamp_type == TimestampType::LogAppendTime {
427            batch.max_timestamp
428        } else {
429            -5
430        };
431        seg.append(&batch, DENSE_INDEX).unwrap();
432        assert2::assert!(seg.first_record_timestamp() == Some(expected));
433        seg.append(&sample_batch().offset(2).timestamp(100).call(), DENSE_INDEX)
434            .unwrap();
435        assert2::assert!(seg.first_record_timestamp() == Some(expected));
436        drop(seg);
437        let mut seg = Segment::open_active(dir.path(), Offset(0), true).unwrap();
438        assert2::assert!(seg.first_record_timestamp() == Some(expected));
439        seg.truncate_to_relative(2).unwrap();
440        assert2::assert!(seg.first_record_timestamp() == Some(expected));
441        seg.truncate_to_relative(0).unwrap();
442        assert2::assert!(seg.first_record_timestamp().is_none());
443        seg.append(&sample_batch().timestamp(200).call(), DENSE_INDEX)
444            .unwrap();
445        assert2::assert!(seg.first_record_timestamp() == Some(200));
446    }
447
448    /// The scan reports the offset of the maximum timestamp as
449    /// `batch.base_offset + record.offset_delta`, and keeps the first record
450    /// holding that timestamp when several share it.
451    #[test]
452    fn the_windowed_scan_reports_the_first_offset_at_the_maximum() {
453        let dir = tempdir().unwrap();
454        // A non-zero base offset and a non-zero delta, so a sum is telling
455        // apart from a product and from either operand alone.
456        let mut seg = Segment::create(dir.path(), Offset(10)).unwrap();
457        let mut batch = RecordBatch {
458            base_offset: 10,
459            base_timestamp: 500,
460            max_timestamp: 502,
461            last_offset_delta: 2,
462            ..RecordBatch::default()
463        };
464        // Records at offsets 10, 11, 12 with timestamps 500, 502, 502: the
465        // maximum is shared, and the first to hold it is offset 11.
466        for (delta, ts_delta) in [(0i32, 0i64), (1, 2), (2, 2)] {
467            batch.records.push(Record {
468                offset_delta: delta,
469                timestamp_delta: ts_delta,
470                key: Some(Bytes::from(format!("k{delta}"))),
471                value: Some(Bytes::from(format!("v{delta}"))),
472                ..Default::default()
473            });
474        }
475        seg.append(&batch, DENSE_INDEX).unwrap();
476
477        let found = seg.scan_max_timestamp_windowed(kibibytes(64));
478        check!(found == Some((Offset(11), 502)), "got {found:?}");
479
480        // Multiple batches: max timestamp is at last_offset
481        seg.append(
482            &sample_batch().offset(13).timestamp(600).call(),
483            DENSE_INDEX,
484        )
485        .unwrap();
486        seg.max_timestamp = i64::MIN;
487        let found_last = seg.scan_max_timestamp_windowed(bytes(1));
488        check!(found_last == Some((Offset(13), 600)));
489
490        // offset_of_max_timestamp routes through scan_max_timestamp_windowed when max_timestamp is MIN
491        let found_via_api = seg.offset_of_max_timestamp();
492        check!(found_via_api == Some((Offset(13), 600)));
493    }
494
495    #[test]
496    fn malformed_record_coordinates_fail_the_scan_closed() {
497        let offset_overflow = RecordBatch {
498            base_offset: i64::MAX,
499            base_timestamp: 0,
500            records: vec![Record {
501                offset_delta: 1,
502                ..Default::default()
503            }],
504            ..Default::default()
505        };
506        let timestamp_overflow = RecordBatch {
507            base_offset: 0,
508            base_timestamp: i64::MAX,
509            records: vec![Record {
510                timestamp_delta: 1,
511                ..Default::default()
512            }],
513            ..Default::default()
514        };
515
516        assert2::assert!(Segment::timestamp_records(&offset_overflow).is_none());
517        assert2::assert!(Segment::timestamp_records(&timestamp_overflow).is_none());
518        let mut append_time = timestamp_overflow;
519        append_time.attributes = append_time
520            .attributes
521            .with_timestamp_type(TimestampType::LogAppendTime);
522        append_time.max_timestamp = 2_400;
523        assert2::assert!(Segment::timestamp_records(&append_time) == Some(vec![(0, 2_400)]));
524        append_time.base_offset = i64::MAX;
525        append_time.records[0].offset_delta = 1;
526        assert2::assert!(Segment::timestamp_records(&append_time).is_none());
527    }
528
529    #[test]
530    fn malformed_or_stale_time_index_floor_fails_closed() {
531        let TempSegment {
532            dir: _dir,
533            segment: mut stale,
534        } = test_segment();
535        stale.append(&one_record_batch(), DENSE_INDEX).unwrap();
536        stale.time_index.append(200, u32::MAX).unwrap();
537        assert2::assert!(stale.offset_for_timestamp(200).is_none());
538
539        let TempSegment {
540            dir: _dir2,
541            segment: mut overflowing,
542        } = segment_at(Offset(i64::MAX));
543        overflowing.last_offset = Offset(i64::MAX);
544        overflowing.time_index.append(0, 1).unwrap();
545        assert2::assert!(overflowing.offset_for_timestamp(0).is_none());
546    }
547
548    #[test]
549    fn truncated_log_failure_exhausts_the_retry_window() {
550        let TempSegment {
551            dir: _dir,
552            segment: seg,
553        } = seeded_fixture().batches(&[ONE_RECORD_BATCH]).call();
554        seg.log_file.set_len(1).unwrap();
555
556        assert2::assert!(
557            seg.offset_for_timestamp_with_window(100, bytes(u32::MAX), None)
558                .unwrap()
559                .is_none()
560        );
561    }
562
563    #[yare::parameterized(
564        first_exact = { 100, Some((Offset(0), 100)) },
565        within_first_batch = { 101, Some((Offset(1), 101)) },
566        between_batches = { 150, Some((Offset(3), 200)) },
567        last_exact = { 201, Some((Offset(4), 201)) },
568        past_end = { 202, None },
569    )]
570    fn offset_for_timestamp_finds_first_ge(ts: i64, want: Option<(Offset, i64)>) {
571        // Two batches: offsets 0..=2 ts 100..=102, offsets 3..=4 ts 200..=201.
572        let TempSegment { dir, segment: seg } = seeded_fixture()
573            .batches(FIVE_OFFSET_TIMESTAMP_BATCHES)
574            .call();
575        // sample_batch sets per-record timestamp_delta = i, base_timestamp = ts_base.
576        // Batch 1 records: (off0,ts100),(off1,ts101),(off2,ts102).
577        // Batch 2 records: (off3,ts200),(off4,ts201).
578        check!(seg.offset_for_timestamp(ts) == want, "ts={ts}");
579        drop(dir);
580    }
581
582    #[yare::parameterized(
583        match_at_final_record = { 1_049, Some((Offset(49), 1_049)) },
584        no_matching_record = { 10_001, None },
585    )]
586    fn scan_from_floor_finds_match_beyond_first_window(
587        threshold: i64,
588        expected: Option<(Offset, i64)>,
589    ) {
590        let TempSegment {
591            dir,
592            segment: mut seg,
593        } = test_segment();
594        // Many single-record batches with increasing timestamps. With a
595        // tiny scan window each batch lands in its own window, so a match
596        // at the tail forces the windowed loop to advance many times.
597        let n = 50i64;
598        for off in 0..n {
599            let mut b = RecordBatch {
600                base_offset: off,
601                base_timestamp: 1_000 + off,
602                max_timestamp: 1_000 + off,
603                last_offset_delta: 0,
604                ..RecordBatch::default()
605            };
606            b.records.push(Record {
607                offset_delta: 0,
608                timestamp_delta: 0,
609                value: Some(Bytes::from(format!("v{off}"))),
610                ..Default::default()
611            });
612            seg.append(&b, DENSE_INDEX).unwrap();
613        }
614        // A window of 1 byte forces one batch per read (anti-stall rule).
615        // Target ts is the very last record's, so the loop must advance
616        // through every window before matching.
617        assert2::assert!(
618            seg.scan_from_floor_windowed(Offset(0), bytes(1), threshold, None)
619                .unwrap()
620                == expected
621        );
622        drop(dir);
623    }
624
625    #[test]
626    fn scan_returns_absolute_offset_of_matching_record() {
627        // A full-size window keeps the match in the first read so the
628        // cursor-advance path isn't involved.
629        const WINDOW: ByteSize = kibibytes(64);
630        let TempSegment {
631            dir,
632            segment: mut seg,
633        } = test_segment();
634        // A leading single-record batch at offset 0, then a 3-record batch
635        // based at offset 1 (abs offsets 1,2,3; timestamps 200,201,202). The
636        // match is the *third* record, whose absolute offset is
637        // `base_offset + offset_delta = 1 + 2 = 3` — a value that only a
638        // correct `+` reproduces (`1 - 2` or `1 * 2` both differ), so this
639        // pins the returned offset arithmetic.
640        seg.append(&one_record_batch(), DENSE_INDEX).unwrap();
641        seg.append(
642            &sample_batch().offset(1).records(3).timestamp(200).call(),
643            DENSE_INDEX,
644        )
645        .unwrap();
646        let got = seg
647            .scan_from_floor_windowed(Offset(0), WINDOW, 202, None)
648            .unwrap();
649        assert2::assert!(got == Some((Offset(3), 202)));
650        drop(dir);
651    }
652
653    #[test]
654    fn offset_of_max_timestamp_earliest_on_tie() {
655        // Batch records ts: 100,101,102 (max in batch = 102 at offset 2).
656        // Second batch: offsets 3,4 ts 200,201 — segment max becomes 201 @4.
657        let TempSegment { dir, segment: seg } = seeded_fixture()
658            .batches(FIVE_OFFSET_TIMESTAMP_BATCHES)
659            .call();
660        assert2::assert!(seg.offset_of_max_timestamp() == Some((Offset(4), 201)));
661
662        // Empty segment → None.
663        let TempSegment {
664            dir: dir2,
665            segment: empty,
666        } = test_segment();
667        assert2::assert!(empty.offset_of_max_timestamp() == None);
668        drop(dir);
669        drop(dir2);
670    }
671
672    #[test]
673    fn offset_of_max_timestamp_tie_picks_earliest() {
674        let TempSegment {
675            dir,
676            segment: mut seg,
677        } = test_segment();
678        // All three records share timestamp 500; earliest offset is 0.
679        let mut b = RecordBatch {
680            base_offset: 0,
681            base_timestamp: 500,
682            max_timestamp: 500,
683            last_offset_delta: 2,
684            ..RecordBatch::default()
685        };
686        for i in 0..3 {
687            b.records.push(Record {
688                offset_delta: i,
689                timestamp_delta: 0,
690                value: Some(Bytes::from("v")),
691                ..Default::default()
692            });
693        }
694        seg.append(&b, DENSE_INDEX).unwrap();
695        assert2::assert!(seg.offset_of_max_timestamp() == Some((Offset(0), 500)));
696        drop(dir);
697    }
698}
699
700/// Kafka trunk's `max.decompressed.message.bytes` on the by-timestamp lookups:
701/// `FileRecords.searchForTimestamp` and `RecordBatch.offsetOfMaxTimestamp`.
702#[cfg(test)]
703mod record_limit_tests {
704    use {
705        super::*,
706        crate::segment::test_support::{DENSE_INDEX, TempSegment, test_segment},
707        assert2::check,
708        bytes::Bytes,
709        krabka_compression::CompressionType,
710        krabka_protocol::records::{Attributes, Record, RecordBatch},
711        krabka_units::prelude::{ByteSize, bytes, kibibytes},
712    };
713
714    /// The body size of a record with a `1000`-byte value and nothing else:
715    /// its attributes byte, a one-byte timestamp and offset delta, a null key,
716    /// the value behind its two-byte length, and no headers.
717    const LARGE: usize = 1_007;
718    const LIMIT: ByteSize = bytes(100);
719    const WINDOW: ByteSize = kibibytes(64);
720
721    /// A record's timestamp and value length.
722    #[derive(Clone, Copy)]
723    struct TimedRecord {
724        timestamp: i64,
725        value_len: usize,
726    }
727
728    /// A batch at `base_offset` holding one record per [`TimedRecord`],
729    /// compressed with `codec`.
730    fn batch(base_offset: i64, records: &[TimedRecord], codec: CompressionType) -> RecordBatch {
731        let base_timestamp = records[0].timestamp;
732        RecordBatch {
733            base_offset,
734            base_timestamp,
735            max_timestamp: records.iter().map(|record| record.timestamp).max().unwrap(),
736            last_offset_delta: i32::try_from(records.len()).unwrap() - 1,
737            attributes: Attributes::default().with_compression(codec),
738            records: records
739                .iter()
740                .enumerate()
741                .map(|(delta, record)| Record {
742                    offset_delta: i32::try_from(delta).unwrap(),
743                    timestamp_delta: record.timestamp - base_timestamp,
744                    value: Some(Bytes::from(vec![7_u8; record.value_len])),
745                    ..Default::default()
746                })
747                .collect(),
748            ..Default::default()
749        }
750    }
751
752    fn segment(batches: &[RecordBatch]) -> TempSegment {
753        let TempSegment {
754            dir,
755            segment: mut seg,
756        } = test_segment();
757        for batch in batches {
758            seg.append(batch, DENSE_INDEX).unwrap();
759        }
760        TempSegment { dir, segment: seg }
761    }
762
763    /// What a lookup answers: the record it found, or the size of the record
764    /// it refused.
765    type Lookup = Result<Option<(Offset, i64)>, usize>;
766
767    fn answer(result: Result<Option<(Offset, i64)>, LogError>) -> Lookup {
768        match result {
769            Ok(found) => Ok(found),
770            Err(LogError::RecordTooLarge { size, limit }) => {
771                check!(limit == LIMIT.bytes_usize());
772                Err(size)
773            }
774            Err(other) => panic!("unexpected error: {other}"),
775        }
776    }
777
778    /// A batch whose max timestamp is below the target is skipped without being
779    /// decompressed, and a batch is read only up to the record the lookup
780    /// returns, so the limit fails a lookup only for a record Kafka reads. An
781    /// uncompressed batch is never held to it.
782    #[yare::parameterized(
783        before_oversized_batch = {
784            "found before the oversized batch",
785            CompressionType::Gzip,
786            Some(LIMIT),
787            1_000,
788            Ok(Some((Offset(0), 1_000))),
789        },
790        oversized_match = {
791            "the oversized record is the match",
792            CompressionType::Gzip,
793            Some(LIMIT),
794            1_001,
795            Err(LARGE),
796        },
797        gzip_without_limit = {
798            "no limit",
799            CompressionType::Gzip,
800            None,
801            1_001,
802            Ok(Some((Offset(1), 1_001))),
803        },
804        uncompressed = {
805            "an uncompressed batch is never held to it",
806            CompressionType::None,
807            Some(LIMIT),
808            1_001,
809            Ok(Some((Offset(1), 1_001))),
810        },
811    )]
812    fn a_lookup_by_timestamp_refuses_only_the_records_kafka_reads(
813        name: &str,
814        codec: CompressionType,
815        limit: Option<ByteSize>,
816        target: i64,
817        expected: Lookup,
818    ) {
819        let TempSegment {
820            dir: _dir,
821            segment: seg,
822        } = segment(&[
823            batch(
824                0,
825                &[TimedRecord {
826                    timestamp: 1_000,
827                    value_len: 10,
828                }],
829                codec,
830            ),
831            batch(
832                1,
833                &[TimedRecord {
834                    timestamp: 1_001,
835                    value_len: 1_000,
836                }],
837                codec,
838            ),
839        ]);
840        let found = seg.offset_for_timestamp_with_window(target, WINDOW, limit);
841        check!(answer(found) == expected, "{name}");
842    }
843
844    /// Kafka's iterator decodes lazily, so a record after the one it returns is
845    /// never read, and one before it always is.
846    #[yare::parameterized(
847        first_record = { 1_000, Ok(Some((Offset(0), 1_000))) },
848        oversized_match = { 1_001, Err(LARGE) },
849        oversized_before_match = { 1_002, Err(LARGE) },
850        skip_below_maximum = { 1_003, Ok(None) },
851    )]
852    fn a_lookup_stops_reading_a_batch_at_the_record_it_returns(target: i64, expected: Lookup) {
853        let TempSegment {
854            dir: _dir,
855            segment: seg,
856        } = segment(&[batch(
857            0,
858            &[
859                TimedRecord {
860                    timestamp: 1_000,
861                    value_len: 10,
862                },
863                TimedRecord {
864                    timestamp: 1_001,
865                    value_len: 1_000,
866                },
867                TimedRecord {
868                    timestamp: 1_002,
869                    value_len: 10,
870                },
871            ],
872            CompressionType::Gzip,
873        )]);
874        let found = seg.offset_for_timestamp_with_window(target, WINDOW, Some(LIMIT));
875        check!(answer(found) == expected, "target {target}");
876    }
877
878    /// `RecordBatch.offsetOfMaxTimestamp` decompresses the batch that holds the
879    /// segment's maximum, and no other.
880    #[test]
881    fn a_max_timestamp_lookup_reads_only_the_batch_holding_the_maximum() {
882        let gzip = CompressionType::Gzip;
883        let TempSegment {
884            dir: _dir,
885            segment: mut seg,
886        } = segment(&[
887            batch(
888                0,
889                &[TimedRecord {
890                    timestamp: 1_000,
891                    value_len: 1_000,
892                }],
893                gzip,
894            ),
895            batch(
896                1,
897                &[TimedRecord {
898                    timestamp: 1_001,
899                    value_len: 10,
900                }],
901                gzip,
902            ),
903        ]);
904        let lookup =
905            |seg: &Segment| answer(seg.offset_of_max_timestamp_with_window(WINDOW, Some(LIMIT)));
906        // The oversized record is in an older batch, which Kafka never reads.
907        check!(lookup(&seg) == Ok(Some((Offset(1), 1_001))));
908
909        // Once the oversized batch holds the maximum, the lookup refuses it.
910        seg.append(
911            &batch(
912                2,
913                &[TimedRecord {
914                    timestamp: 1_002,
915                    value_len: 1_000,
916                }],
917                gzip,
918            ),
919            DENSE_INDEX,
920        )
921        .unwrap();
922        check!(lookup(&seg) == Err(LARGE));
923
924        // A segment that never learned its maximum finds it by reading
925        // every record, and the limit holds there too.
926        seg.max_timestamp = i64::MIN;
927        check!(lookup(&seg) == Err(LARGE));
928        check!(
929            seg.offset_of_max_timestamp_with_window(WINDOW, None)
930                .unwrap()
931                == Some((Offset(2), 1_002))
932        );
933    }
934}
krabka_log/segment/timestamp_scan.rs

krabka_log/segment/
timestamp_scan.rs
