krabka_log/segment/
open.rs

1//! Constructors for a segment and the tail recovery that an active segment
2//! needs when it is reopened.
3//!
4//! Fresh segments are created, sealed segments are validated for reading,
5//! and active segments recover their tail. Log recovery also uses a private
6//! no-scan loader before it supplies the sealed segment bounds.
7
8#[cfg(unix)]
9use std::os::unix::fs::MetadataExt;
10use {
11    super::{
12        ReserveRequest, Segment,
13        io::{SourceIo as _, seek_to_log_size},
14    },
15    crate::{
16        config::DEFAULT_INDEX_INTERVAL,
17        error::LogError,
18        index::{OffsetIndex, TimeIndex},
19        io::{FileIo, LogIo},
20        name,
21    },
22    krabka_ids::Offset,
23    krabka_protocol::records::RecordBatch,
24    krabka_units::prelude::{ByteSize, ByteSizeExt},
25    std::{
26        fs::{File, Metadata, OpenOptions},
27        path::Path,
28        sync::{Arc, atomic::AtomicBool},
29    },
30    tracing::{field, instrument},
31};
32
33/// One sparse-index point the tail scan found, as Kafka's `LogSegment.recover`
34/// writes it: the offset index gets the indexed batch's last offset and its
35/// byte position, and the time index gets the running maximum timestamp with
36/// the last offset of the batch that set it.
37#[derive(Debug, Clone, PartialEq, Eq)]
38struct IndexPoint {
39    /// Last offset of the indexed batch, relative to the segment base.
40    last_offset: u32,
41    /// Byte position where the indexed batch starts.
42    position: u32,
43    /// Maximum timestamp of every batch up to and including the indexed one.
44    max_timestamp: i64,
45    /// Last offset of the batch that set `max_timestamp`, relative to the
46    /// segment base.
47    max_timestamp_offset: u32,
48}
49
50/// What a walk over a segment's batches derives: the end of their valid
51/// prefix, the segment's frontiers, and the sparse-index points Kafka's
52/// `LogSegment.recover` writes.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct TailScan {
55    valid_end: u64,
56    last_offset: Offset,
57    max_timestamp: i64,
58    max_timestamp_offset: Offset,
59    index_points: Vec<IndexPoint>,
60}
61
62/// Builds a [`TailScan`] one batch at a time, in file order.
63///
64/// The tail scan feeds it the batches it decodes from a file, and a
65/// compaction rewrite the batches it writes, so the two derive the same
66/// frontiers and index points by the same rule.
67pub struct FrontierScan {
68    base_offset: Offset,
69    file_end: u64,
70    index_interval: ByteSize,
71    next_offset: i64,
72    scan: TailScan,
73}
74
75impl FrontierScan {
76    /// A walk over a segment at `base_offset` whose file holds `file_end`
77    /// bytes, indexing every `index_interval` bytes.
78    pub(crate) fn new(
79        base_offset: Offset,
80        file_end: u64,
81        index_interval: ByteSize,
82    ) -> Result<Self, LogError> {
83        let last_offset = base_offset
84            .0
85            .checked_sub(1)
86            .map(Offset)
87            .ok_or_else(|| LogError::Corrupt("recovery offset underflow".into()))?;
88        Ok(Self {
89            base_offset,
90            file_end,
91            index_interval,
92            next_offset: base_offset.0,
93            scan: TailScan {
94                valid_end: 0,
95                last_offset,
96                max_timestamp: i64::MIN,
97                max_timestamp_offset: last_offset,
98                index_points: Vec::new(),
99            },
100        })
101    }
102
103    /// Take `batch`, `encoded_len` bytes long, which starts where the valid
104    /// prefix ends. Returns `false`, and takes nothing, when the batch does not
105    /// extend the valid prefix: it runs past the file, or its offsets do not
106    /// follow the prefix's.
107    pub(crate) fn push(&mut self, batch: &RecordBatch, encoded_len: u64) -> Result<bool, LogError> {
108        let batch_position = self.scan.valid_end;
109        let Some(step) = krabka_verified::local_recovery_batch_step(
110            batch_position,
111            self.file_end,
112            self.next_offset,
113            batch.base_offset,
114            batch.last_offset_delta,
115            encoded_len,
116        ) else {
117            return Ok(false);
118        };
119        let scan = &mut self.scan;
120        scan.valid_end = step.valid_end;
121        scan.last_offset = Offset(step.last_offset);
122        self.next_offset = step.next_offset;
123        if batch.max_timestamp > scan.max_timestamp {
124            scan.max_timestamp = batch.max_timestamp;
125            scan.max_timestamp_offset = scan.last_offset;
126        }
127        // Kafka's `LogSegment.recover`: `validBytes - lastIndexEntry >
128        // indexIntervalBytes`, the last entry starting at zero.
129        let last_indexed = scan
130            .index_points
131            .last()
132            .map_or(0, |previous| u64::from(previous.position));
133        if batch_position.saturating_sub(last_indexed) > self.index_interval.bytes_u64() {
134            let base_offset = self.base_offset;
135            let relative = |offset: Offset| {
136                krabka_verified::truncation_relative_offset(base_offset.0, offset.0).ok_or_else(
137                    || LogError::Corrupt("recovered batch offset exceeds index range".into()),
138                )
139            };
140            let position = u32::try_from(batch_position).map_err(|_| {
141                LogError::Corrupt("recovered batch position exceeds index range".into())
142            })?;
143            scan.index_points.push(IndexPoint {
144                last_offset: relative(scan.last_offset)?,
145                position,
146                max_timestamp: scan.max_timestamp,
147                max_timestamp_offset: relative(scan.max_timestamp_offset)?,
148            });
149        }
150        Ok(true)
151    }
152
153    /// The walk's result.
154    pub(crate) fn finish(self) -> TailScan { self.scan }
155}
156
157impl TailScan {
158    /// The end of the valid prefix: the length of a file that holds exactly
159    /// what was scanned.
160    pub(crate) fn valid_end(&self) -> u64 { self.valid_end }
161
162    /// Check that the scanned prefix fits a segment: its end in a `u32`
163    /// position and its offsets in the 32-bit relative index.
164    fn check_index_range(&self, base_offset: Offset) -> Result<(), LogError> {
165        u32::try_from(self.valid_end)
166            .map_err(|_| LogError::Corrupt("recovered segment position exceeds u32".into()))?;
167        krabka_verified::local_recovery_index_frontier(base_offset.0, self.last_offset.0)
168            .ok_or_else(|| {
169                LogError::Corrupt("recovered segment offset exceeds index range".into())
170            })?;
171        Ok(())
172    }
173
174    /// Append the scan's index points to `offset_index` and `time_index`,
175    /// then the running maximum timestamp, as a segment that closed normally
176    /// always had its newest timestamp in the index (Kafka's
177    /// `LogSegment.recover` ends the same way).
178    fn write_index_points(
179        &self,
180        base_offset: Offset,
181        offset_index: &mut OffsetIndex,
182        time_index: &mut TimeIndex,
183    ) -> Result<(), LogError> {
184        for point in &self.index_points {
185            offset_index.append(point.last_offset, point.position)?;
186            time_index.maybe_append(point.max_timestamp, point.max_timestamp_offset)?;
187        }
188        if self.max_timestamp >= 0 {
189            let relative = krabka_verified::truncation_relative_offset(
190                base_offset.0,
191                self.max_timestamp_offset.0,
192            )
193            .ok_or_else(|| LogError::BadSegmentName("offset overflow in segment".into()))?;
194            time_index.maybe_append(self.max_timestamp, relative)?;
195        }
196        Ok(())
197    }
198
199    /// Write the scan's sparse indexes for a segment at `base_offset` into
200    /// the index files at `index_path` and `timeindex_path`, which a
201    /// compaction rewrite creates empty beside its output, through `io`.
202    ///
203    /// The index points are then dropped: the files hold them, and a
204    /// compaction pass keeps every output's scan until it installs them all.
205    pub(crate) fn write_indexes(
206        &mut self,
207        io: &Arc<dyn LogIo>,
208        base_offset: Offset,
209        index_path: &Path,
210        timeindex_path: &Path,
211    ) -> Result<(), LogError> {
212        self.check_index_range(base_offset)?;
213        let mut offset_index = OffsetIndex::open(index_path)?;
214        offset_index.set_io(Arc::clone(io));
215        let mut time_index = TimeIndex::open(timeindex_path)?;
216        time_index.set_io(Arc::clone(io));
217        self.write_index_points(base_offset, &mut offset_index, &mut time_index)?;
218        self.index_points = Vec::new();
219        Ok(())
220    }
221}
222
223impl Segment {
224    /// Attach the sparse indexes and initial frontiers to an already-open log file.
225    fn from_log_file(
226        dir: &Path,
227        base_offset: Offset,
228        log_file: File,
229        log_size: u64,
230    ) -> Result<Self, LogError> {
231        let offset_index = OffsetIndex::open(&name::index_path(dir, base_offset.0))?;
232        let time_index = TimeIndex::open(&name::timeindex_path(dir, base_offset.0))?;
233        Ok(Self {
234            dir: dir.to_path_buf(),
235            base_offset,
236            log_file: Arc::new(log_file),
237            io: Arc::new(FileIo),
238            log_size,
239            offset_index,
240            time_index,
241            sealed: false,
242            max_timestamp: i64::MIN,
243            max_timestamp_offset: base_offset - 1,
244            first_timestamp: None,
245            last_offset: base_offset - 1,
246            reserved_end: 0,
247            reserve_requested: ReserveRequest::None,
248            direct: None,
249            direct_asked: false,
250            wrote_direct: false,
251            cache_probe_unsupported: Arc::new(AtomicBool::new(false)),
252        })
253    }
254
255    /// Create a fresh active segment at the given base offset. This fails if
256    /// the `.log` file already exists.
257    #[instrument(
258        level = "debug",
259        skip_all,
260        fields(dir = %dir.display(), base_offset = base_offset.0),
261        err,
262    )]
263    /// # Errors
264    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
265    pub fn create(dir: &Path, base_offset: Offset) -> Result<Self, LogError> {
266        let path = name::log_path(dir, base_offset.0);
267        let mut opts = OpenOptions::new();
268        opts.read(true).write(true).create_new(true);
269        let log_file = opts.open(&path)?;
270        Self::from_log_file(dir, base_offset, log_file, 0)
271    }
272
273    /// Open as the active segment.
274    ///
275    /// When `validate` is true, this method scans from byte zero to EOF. It
276    /// truncates a partial or invalid trailing batch, rebuilds both sparse
277    /// indexes, and updates `last_offset` and `max_timestamp` from the same
278    /// maximal valid prefix.
279    #[instrument(
280        level = "info",
281        skip_all,
282        fields(dir = %dir.display(), base_offset = base_offset.0, validate),
283        err,
284    )]
285    /// # Errors
286    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
287    pub fn open_active(dir: &Path, base_offset: Offset, validate: bool) -> Result<Self, LogError> {
288        Self::open_active_with_index_interval(dir, base_offset, validate, DEFAULT_INDEX_INTERVAL)
289    }
290
291    pub(crate) fn open_active_with_index_interval(
292        dir: &Path,
293        base_offset: Offset,
294        validate: bool,
295        index_interval: ByteSize,
296    ) -> Result<Self, LogError> {
297        let mut seg = Self::open_unscanned(dir, base_offset)?;
298        if validate {
299            seg.recover_active_tail(index_interval)?;
300        }
301        Ok(seg)
302    }
303
304    #[instrument(
305        level = "debug",
306        skip_all,
307        fields(
308            base_offset = self.base_offset.0,
309            log_size = self.log_size,
310            recovered_last_offset = field::Empty,
311        ),
312        err,
313    )]
314    fn recover_active_tail(&mut self, index_interval: ByteSize) -> Result<(), LogError> {
315        let recovered = self.scan_valid_tail(index_interval)?;
316        recovered.check_index_range(self.base_offset)?;
317        if recovered.valid_end < self.log_size {
318            self.log_file.set_len(recovered.valid_end)?;
319            self.renew_reservation();
320        }
321        self.log_size = recovered.valid_end;
322        seek_to_log_size(&self.log_file, self.log_size)?;
323        self.offset_index.truncate_by_position(0)?;
324        self.time_index.truncate_by_relative_offset(0)?;
325        recovered.write_index_points(
326            self.base_offset,
327            &mut self.offset_index,
328            &mut self.time_index,
329        )?;
330        self.last_offset = recovered.last_offset;
331        self.max_timestamp = recovered.max_timestamp;
332        self.max_timestamp_offset = recovered.max_timestamp_offset;
333        tracing::Span::current().record("recovered_last_offset", self.last_offset.0);
334        Ok(())
335    }
336
337    fn scan_valid_tail(&self, index_interval: ByteSize) -> Result<TailScan, LogError> {
338        let mut buf = Vec::new();
339        let to_read = usize::try_from(self.log_size).unwrap_or(usize::MAX);
340        self.read_log_range(0, &mut buf, to_read)?;
341
342        let mut scan = FrontierScan::new(self.base_offset, self.log_size, index_interval)?;
343        let mut cur: &[u8] = &buf;
344        while !cur.is_empty() {
345            let before = cur.len();
346            let Ok(batch) = RecordBatch::decode(&mut cur) else {
347                break;
348            };
349            let encoded_len = u64::try_from(before - cur.len())
350                .map_err(|_| LogError::Corrupt("decoded batch length exceeds u64".into()))?;
351            if !scan.push(&batch, encoded_len)? {
352                break;
353            }
354        }
355        Ok(scan.finish())
356    }
357
358    /// Open an existing segment for reading and refuse further appends.
359    ///
360    /// Scans the batches to restore the last offset and maximum timestamp. A corrupt or incomplete batch returns an error without truncating the file. Use [`Segment::open_active`] to recover a writable segment.
361    ///
362    /// # Errors
363    ///
364    /// Returns an error when the segment or its indexes cannot be opened, its batches cannot be read, or its offset chain or batch bytes are corrupt.
365    #[instrument(
366        level = "debug",
367        skip_all,
368        fields(dir = %dir.display(), base_offset = base_offset.0),
369        err,
370    )]
371    pub fn open(dir: &Path, base_offset: Offset) -> Result<Self, LogError> {
372        let mut segment = Self::open_unscanned(dir, base_offset)?;
373        let scanned = segment.scan_valid_tail(DEFAULT_INDEX_INTERVAL)?;
374        if scanned.valid_end != segment.log_size {
375            return Err(LogError::Corrupt(
376                "sealed segment contains an invalid batch".into(),
377            ));
378        }
379        segment.last_offset = scanned.last_offset;
380        segment.max_timestamp = scanned.max_timestamp;
381        segment.max_timestamp_offset = scanned.max_timestamp_offset;
382        segment.sealed = true;
383        Ok(segment)
384    }
385
386    /// Open the segment a compaction rewrite wrote at `base_offset`, from the
387    /// [`TailScan`] the rewrite derived while it wrote it and the indexes it
388    /// wrote with [`TailScan::write_indexes`], without decoding the file
389    /// again. The file must hold exactly the bytes the scan covers.
390    pub(crate) fn open_rewritten(
391        dir: &Path,
392        base_offset: Offset,
393        scan: &TailScan,
394    ) -> Result<Self, LogError> {
395        let mut segment = Self::open_unscanned(dir, base_offset)?;
396        if segment.log_size != scan.valid_end {
397            return Err(LogError::Corrupt(format!(
398                "rewritten segment {} holds {} bytes, the rewrite wrote {}",
399                base_offset.0, segment.log_size, scan.valid_end
400            )));
401        }
402        scan.check_index_range(base_offset)?;
403        segment.last_offset = scan.last_offset;
404        segment.max_timestamp = scan.max_timestamp;
405        segment.max_timestamp_offset = scan.max_timestamp_offset;
406        Ok(segment)
407    }
408
409    // Log recovery supplies sealed bounds from the following segment and
410    // restores timestamps separately, without decoding every record body.
411    pub(crate) fn open_unscanned(dir: &Path, base_offset: Offset) -> Result<Self, LogError> {
412        let log_path = name::log_path(dir, base_offset.0);
413        let log_file = OpenOptions::new().read(true).write(true).open(&log_path)?;
414        let metadata = log_file.metadata()?;
415        let log_size = metadata.len();
416        seek_to_log_size(&log_file, log_size)?;
417        let mut segment = Self::from_log_file(dir, base_offset, log_file, log_size)?;
418        // A reservation outlives the process that made it, and nothing on
419        // disk records it but the blocks themselves. Count them, so sealing
420        // this segment later gives back what an earlier run reserved.
421        segment.reserved_end = allocated_bytes(&metadata);
422        Ok(segment)
423    }
424}
425
426/// Bytes of disk the file behind `metadata` occupies, which is past its
427/// length when blocks are reserved beyond its end.
428#[cfg(unix)]
429fn allocated_bytes(metadata: &Metadata) -> u64 { MetadataExt::blocks(metadata).saturating_mul(512) }
430
431/// The targets without block counts cannot reserve blocks either.
432#[cfg(not(unix))]
433fn allocated_bytes(metadata: &Metadata) -> u64 {
434    let _ = metadata;
435    0
436}
437
438#[cfg(test)]
439mod tests {
440    use {
441        super::*,
442        crate::segment::test_support::{
443            DENSE_INDEX, FIVE_OFFSET_TIMESTAMP_BATCHES, NO_LIMIT, ONE_RECORD_BATCH,
444            SECOND_SINGLE_RECORD_BATCH, THIRD_SINGLE_RECORD_BATCH, sample_batch, seeded_segment,
445            three_single_record_batches,
446        },
447        krabka_units::prelude::kibibytes,
448        std::{fs, io::Write},
449        tempfile::tempdir,
450    };
451
452    #[test]
453    fn append_after_open_active_writes_at_eof() {
454        let dir = tempdir().unwrap();
455        drop(
456            seeded_segment()
457                .dir(dir.path())
458                .batches(&[ONE_RECORD_BATCH, SECOND_SINGLE_RECORD_BATCH])
459                .call(),
460        );
461
462        let mut seg = Segment::open_active(dir.path(), Offset(0), true).unwrap();
463        let position = seg
464            .append(
465                &sample_batch()
466                    .offset(THIRD_SINGLE_RECORD_BATCH.offset.0)
467                    .records(THIRD_SINGLE_RECORD_BATCH.records.0)
468                    .timestamp(THIRD_SINGLE_RECORD_BATCH.timestamp.0)
469                    .call(),
470                DENSE_INDEX,
471            )
472            .unwrap();
473
474        let read = seg.read(Offset(0), NO_LIMIT).unwrap();
475        assert2::assert!(position > 0);
476        assert2::assert!(seg.last_offset() == Offset(2));
477        assert2::assert!(read == three_single_record_batches());
478    }
479
480    /// Tail recovery must physically truncate a partial or garbage tail and
481    /// rebuild both indexes from exactly the batches in the retained prefix.
482    #[test]
483    fn recover_active_tail_truncates_trailing_garbage() {
484        let dir = tempdir().unwrap();
485        let valid_size = {
486            let mut seg = seeded_segment()
487                .dir(dir.path())
488                .batches(FIVE_OFFSET_TIMESTAMP_BATCHES)
489                .call();
490            seg.flush().unwrap();
491            let valid_size = seg.log_size;
492            let stale_position = u32::try_from(valid_size).unwrap();
493            seg.offset_index.append(999, stale_position).unwrap();
494            seg.time_index.append(i64::MAX, 999).unwrap();
495            valid_size
496        };
497
498        // Append 16 bytes of garbage (an undecodable partial batch tail).
499        let log_path = name::log_path(dir.path(), 0);
500        let mut f = OpenOptions::new().append(true).open(&log_path).unwrap();
501        f.write_all(&[0xCD; 16]).unwrap();
502        f.sync_data().unwrap();
503        drop(f);
504
505        // Reopen with validation: the tail scan must clip the garbage.
506        let seg =
507            Segment::open_active_with_index_interval(dir.path(), Offset(0), true, DENSE_INDEX)
508                .unwrap();
509        assert2::assert!(seg.last_offset() == Offset(4));
510        assert2::assert!(seg.log_size == valid_size);
511        assert2::assert!(fs::metadata(&log_path).unwrap().len() == valid_size);
512        // The first batch takes no entry (Kafka's `LogSegment.recover` indexes
513        // a batch once more than the interval lies between it and the last
514        // entry, which starts at zero), so the second batch's is the one.
515        assert2::assert!(seg.offset_index.entry_count() == 1);
516        assert2::assert!(seg.time_index.entry_count() == 1);
517        assert2::assert!(u64::from(seg.offset_index.lookup(999)) < valid_size);
518        assert2::assert!(seg.time_index.last_entry().unwrap().relative_offset < 999);
519        drop(seg);
520
521        // Recovery is idempotent once the file and both indexes share the
522        // proved frontiers.
523        let retry =
524            Segment::open_active_with_index_interval(dir.path(), Offset(0), true, DENSE_INDEX)
525                .unwrap();
526        assert2::assert!(retry.last_offset() == Offset(4));
527        assert2::assert!(retry.log_size == valid_size);
528        assert2::assert!(fs::metadata(&log_path).unwrap().len() == valid_size);
529        assert2::assert!(retry.offset_index.entry_count() == 1);
530        assert2::assert!(retry.time_index.entry_count() == 1);
531    }
532
533    /// Recovery ends with Kafka's final time-index entry, `(maxTimestampSoFar,
534    /// shallowOffsetOfMaxTimestampSoFar)`: the sparse index holds no point for
535    /// these small batches, so without it the index would not know any of
536    /// their timestamps.
537    #[test]
538    fn recover_active_tail_appends_the_final_time_index_entry() {
539        let dir = tempdir().unwrap();
540        {
541            let mut seg = Segment::create(dir.path(), Offset(0)).unwrap();
542            for (base, timestamp) in [(0, 100), (1, 300), (2, 200)] {
543                seg.append(
544                    &sample_batch().offset(base).timestamp(timestamp).call(),
545                    kibibytes(4),
546                )
547                .unwrap();
548            }
549        }
550
551        let seg = Segment::open_active(dir.path(), Offset(0), true).unwrap();
552
553        assert2::assert!(seg.time_index.entry_count() == 1);
554        assert2::assert!(
555            seg.time_index
556                .last_entry()
557                .map(|entry| (entry.timestamp, entry.relative_offset))
558                == Some((300, 1))
559        );
560    }
561
562    #[test]
563    fn unvalidated_open_reports_last_offset_as_base_minus_one() {
564        let dir = tempdir().unwrap();
565        Segment::create(dir.path(), Offset(10)).unwrap();
566        let seg = Segment::open_unscanned(dir.path(), Offset(10)).unwrap();
567        assert2::assert!(seg.last_offset() == Offset(9));
568    }
569
570    /// Writing a scan's indexes keeps its frontiers and drops its index
571    /// points, which the files then hold.
572    #[test]
573    fn writing_the_indexes_drops_the_points_and_keeps_the_frontiers() {
574        let dir = tempdir().unwrap();
575        let seg = seeded_segment()
576            .dir(dir.path())
577            .batches(FIVE_OFFSET_TIMESTAMP_BATCHES)
578            .call();
579        let mut scan = seg.scan_valid_tail(DENSE_INDEX).unwrap();
580        let scanned = scan.clone();
581        let out = tempdir().unwrap();
582        let index = name::index_path(out.path(), 0);
583        let timeindex = name::timeindex_path(out.path(), 0);
584        let io: Arc<dyn LogIo> = Arc::new(FileIo);
585
586        scan.write_indexes(&io, Offset(0), &index, &timeindex)
587            .unwrap();
588
589        assert2::check!(!scanned.index_points.is_empty());
590        assert2::check!(
591            scan == TailScan {
592                index_points: Vec::new(),
593                ..scanned
594            }
595        );
596        assert2::check!(OffsetIndex::open(&index).unwrap().entry_count() > 0);
597    }
598}
krabka_log/segment/open.rs

krabka_log/segment/
open.rs
