krabka_log/segment/
lifecycle.rs

1//! State transitions of a segment after it holds data: sealing, flushing, and
2//! truncation.
3//!
4//! Each one rewrites the segment's own view of what it holds -- the sealed
5//! flag, the byte length, the last offset, the maximum timestamp, and the
6//! sparse indexes -- so they belong together rather than beside a read path.
7
8#[cfg(all(test, not(target_os = "wasi")))]
9use std::cell::Cell;
10#[cfg(not(target_os = "wasi"))]
11use std::fs::File;
12use {
13    super::{
14        ReserveRequest, Segment,
15        io::{SourceIo as _, seek_to_log_size},
16    },
17    crate::{error::LogError, io::LogIo},
18    krabka_ids::Offset,
19    krabka_protocol::records::RecordBatch,
20    krabka_units::prelude::{ByteSize, ByteSizeExt as _},
21    std::{
22        io::{self, ErrorKind},
23        mem,
24        sync::{Arc, atomic::Ordering},
25    },
26    tracing::{field, instrument},
27};
28
29/// Log a block reservation, a release, or a switch to `O_DIRECT` that the
30/// filesystem refused.
31///
32/// A filesystem without `fallocate` or `O_DIRECT` refuses every one, so that
33/// is not news worth more than a debug line; anything else, a full disk above
34/// all, is.
35/// The least a reservation step covers: a quarter of the segment, or this
36/// when the segment is small. A 1 GiB segment reserves 256 MiB at a time.
37const RESERVE_STEP_MIN: u64 = 1 << 20;
38
39pub(super) fn log_refused(operation: &'static str, base_offset: Offset, error: &io::Error) {
40    if error.kind() == ErrorKind::Unsupported {
41        tracing::debug!(operation, base_offset = base_offset.0, %error, "segment preallocation unsupported");
42    } else {
43        tracing::warn!(operation, base_offset = base_offset.0, %error, "segment preallocation refused");
44    }
45}
46
47/// What [`Segment::write_snapshot`] saves for [`Segment::rollback_failed_write`].
48#[derive(Debug, Clone, Copy)]
49pub(super) struct WriteSnapshot {
50    last_offset: Offset,
51    max_timestamp: i64,
52    max_timestamp_offset: Offset,
53}
54
55#[cfg(all(test, not(target_os = "wasi")))]
56std::thread_local! {
57    static FAIL_FLUSH_HANDLES: Cell<bool> = const { Cell::new(false) };
58}
59
60impl Segment {
61    #[cfg(all(test, not(target_os = "wasi")))]
62    pub(crate) fn test_fail_flush_handles(fail: bool) { FAIL_FLUSH_HANDLES.set(fail); }
63
64    #[cfg(not(target_os = "wasi"))]
65    pub(crate) fn flush_handles(&self) -> io::Result<[Arc<File>; 3]> {
66        #[cfg(test)]
67        if FAIL_FLUSH_HANDLES.get() {
68            return Err(io::Error::other("injected flush handle clone failure"));
69        }
70        Ok([
71            Arc::clone(&self.log_file),
72            Arc::new(self.offset_index.flush_handle()?),
73            Arc::new(self.time_index.flush_handle()?),
74        ])
75    }
76
77    /// The `.log` file, for a background sync of its page-cache data.
78    #[cfg(not(target_os = "wasi"))]
79    pub(crate) fn log_file_handle(&self) -> Arc<File> { Arc::clone(&self.log_file) }
80
81    /// The `.log` file's shared handle, which the page-cache ledger records
82    /// its synced ranges against.
83    #[cfg(not(target_os = "wasi"))]
84    pub(crate) fn log_file(&self) -> &Arc<File> { &self.log_file }
85
86    /// Mark this segment as sealed. No more appends.
87    ///
88    /// Sealing first writes the segment's final time-index entry, Kafka's
89    /// `LogSegment.onBecomeInactiveSegment`: the sparse index lags the writes,
90    /// so without it a reopened segment would not learn the timestamp of the
91    /// batches after its last index point from the index. A segment that
92    /// wrote through `O_DIRECT` then goes back to the page cache, which cuts
93    /// the padding past its last batch, and what is left of its block
94    /// reservation is given back.
95    ///
96    /// # Errors
97    /// Returns an error when the time-index entry cannot be written or its
98    /// offset overflows the index range, or when the padding cannot be cut.
99    /// The segment stays open then.
100    pub fn seal(&mut self) -> Result<(), LogError> {
101        self.append_running_max_time_entry()?;
102        self.write_buffered()?;
103        self.sealed = true;
104        self.release_reservation();
105        Ok(())
106    }
107
108    /// Where the blocks reserved for a segment of `size` bytes should reach
109    /// now: the whole segment while a quarter of it is under
110    /// [`RESERVE_STEP_MIN`], otherwise the next quarter boundary at least a
111    /// sixteenth of the segment past the bytes it holds, so the reservation
112    /// runs ahead of the appends and a broker with many partitions commits a
113    /// quarter of a segment per partition at a time rather than a whole one.
114    fn reservation_target(&self, size: u64) -> u64 {
115        let step = (size / 4).max(RESERVE_STEP_MIN);
116        let needed = self.log_size.saturating_add(step / 4);
117        if needed >= size {
118            return size;
119        }
120        needed.div_ceil(step).saturating_mul(step).min(size)
121    }
122
123    /// Whether the next append wants more blocks reserved than the segment
124    /// holds or has reserved. See [`Self::reserve`].
125    pub(crate) fn needs_reservation(&self, size: ByteSize) -> bool {
126        self.reserve_requested != ReserveRequest::Refused
127            && self.reservation_target(size.bytes_u64()) > self.log_size.max(self.reserved_end)
128    }
129
130    /// Reserve disk blocks for this segment to grow towards `size` without
131    /// allocating as it goes: Kafka's `preallocate`, with the file's length
132    /// left at the bytes it holds. See
133    /// [`SegmentAllocation::Preallocate`](crate::SegmentAllocation::Preallocate).
134    ///
135    /// The log calls this before every append. It reserves a step at a time,
136    /// see [`Self::reservation_target`], so most calls do nothing. A refused
137    /// reservation leaves a segment that grows as it is written, which is
138    /// the segment it would have been without one, so it is logged and not
139    /// returned, and not asked for again.
140    pub(crate) fn reserve(&mut self, size: ByteSize) {
141        if !self.needs_reservation(size) {
142            return;
143        }
144        self.reserve_requested = ReserveRequest::Size(size.bytes_u64());
145        let target = self.reservation_target(size.bytes_u64());
146        let start = self.log_size.max(self.reserved_end);
147        match self.io.reserve(&self.log_file, start, target - start) {
148            Ok(()) => self.reserved_end = target,
149            Err(error) => {
150                self.reserve_requested = ReserveRequest::Refused;
151                log_refused("reserve", self.base_offset, &error);
152            }
153        }
154    }
155
156    /// Take the reservation again after a truncate.
157    ///
158    /// A truncate frees every block past the new end, the reservation's with
159    /// them -- it is how [`Self::release_reservation`] gives them back -- so
160    /// the blocks counted before it are gone, and a reservation this process
161    /// asked for has to be asked for again.
162    pub(super) fn renew_reservation(&mut self) {
163        self.reserved_end = 0;
164        if let ReserveRequest::Size(size) = self.reserve_requested {
165            self.reserve(ByteSize::from_bytes(size));
166        }
167    }
168
169    /// Give back the blocks reserved past the bytes this segment holds.
170    ///
171    /// A segment sealed before it filled -- by `segment.ms`, a full index, or
172    /// a batch that did not fit -- would otherwise keep its reservation for as
173    /// long as retention keeps the segment. A refused release keeps them only
174    /// until then, so it is logged and not returned.
175    fn release_reservation(&mut self) {
176        let reserved_end = mem::take(&mut self.reserved_end);
177        if reserved_end <= self.log_size {
178            return;
179        }
180        if let Err(error) = self.io.release(&self.log_file, self.log_size) {
181            log_refused("release", self.base_offset, &error);
182        }
183    }
184
185    /// Kafka's `timeIndex().maybeAppend(maxTimestampSoFar(),
186    /// shallowOffsetOfMaxTimestampSoFar())`: the running maximum timestamp
187    /// with the last offset of the batch that set it. The index keeps only an
188    /// entry whose timestamp is newer than its last, so it stays strictly
189    /// increasing, and a segment without a timestamp (`i64::MIN`, or
190    /// `NO_TIMESTAMP`) adds nothing.
191    pub(super) fn append_running_max_time_entry(&mut self) -> Result<(), LogError> {
192        if self.max_timestamp < 0 {
193            return Ok(());
194        }
195        let relative = krabka_verified::truncation_relative_offset(
196            self.base_offset.0,
197            self.max_timestamp_offset.0,
198        )
199        .ok_or_else(|| LogError::BadSegmentName("offset overflow in segment".into()))?;
200        self.time_index.maybe_append(self.max_timestamp, relative)
201    }
202
203    /// Seal a segment loaded without a scan and
204    /// set its `last_offset` to `last`.
205    ///
206    /// Callers pass `next_segment.base_offset - 1`, the highest offset this
207    /// sealed segment can hold. The no-scan loader leaves
208    /// `last_offset = base_offset - 1` because it does not scan the `.log`.
209    /// Without this fix, a sealed segment recovered on
210    /// [`Log::open`](crate::Log) reports that stale `last_offset`.
211    /// `Log::read_raw` skips any segment whose
212    /// `last_offset() < fetch_offset`, so it would skip the first sealed
213    /// segment after a restart and serve a later segment's base offset. That
214    /// creates an offset gap, and a follower that fetches at 0 then loops on
215    /// the resulting append mismatch.
216    pub fn seal_at(&mut self, last: Offset) {
217        self.sealed = true;
218        self.last_offset = last;
219    }
220
221    /// Force-sync everything to disk.
222    #[instrument(
223        level = "debug",
224        skip_all,
225        fields(base_offset = self.base_offset.0, log_size = self.log_size),
226        err,
227    )]
228    /// # Errors
229    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
230    pub fn flush(&mut self) -> Result<(), LogError> {
231        self.drain_direct()?;
232        self.io.sync_data(&self.log_file)?;
233        self.offset_index.flush()?;
234        self.time_index.flush()?;
235        Ok(())
236    }
237
238    /// The part of the segment's state an append changes, saved before the
239    /// write so a failed one can put it back.
240    pub(super) fn write_snapshot(&self) -> WriteSnapshot {
241        WriteSnapshot {
242            last_offset: self.last_offset,
243            max_timestamp: self.max_timestamp,
244            max_timestamp_offset: self.max_timestamp_offset,
245        }
246    }
247
248    pub(super) fn rollback_failed_write(
249        &mut self,
250        position: u64,
251        snapshot: WriteSnapshot,
252    ) -> Result<(), LogError> {
253        let WriteSnapshot {
254            last_offset,
255            max_timestamp,
256            max_timestamp_offset,
257        } = snapshot;
258        self.settle_direct();
259        self.log_file.set_len(position)?;
260        self.first_timestamp = None;
261        seek_to_log_size(&self.log_file, position)?;
262        self.log_size = position;
263        self.renew_reservation();
264        self.resync_direct()?;
265        self.last_offset = last_offset;
266        self.max_timestamp = max_timestamp;
267        self.max_timestamp_offset = max_timestamp_offset;
268        let position = u32::try_from(position)
269            .map_err(|_| LogError::BadSegmentName("position overflow".into()))?;
270        self.offset_index.truncate_by_position(position)?;
271        let next = last_offset
272            .0
273            .checked_add(1)
274            .and_then(|next| next.checked_sub(self.base_offset.0))
275            .ok_or_else(|| LogError::BadSegmentName("offset overflow".into()))?;
276        let next_relative =
277            u32::try_from(next).map_err(|_| LogError::BadSegmentName("offset overflow".into()))?;
278        self.time_index.truncate_by_relative_offset(next_relative)?;
279        Ok(())
280    }
281
282    pub(crate) fn set_io(&mut self, io: Arc<dyn LogIo>) {
283        self.offset_index.set_io(Arc::clone(&io));
284        self.time_index.set_io(Arc::clone(&io));
285        if let Some(mode) = self.direct.as_ref() {
286            mode.writer.set_io(Arc::clone(&io));
287        }
288        self.io = io;
289        self.cache_probe_unsupported.store(false, Ordering::Relaxed);
290    }
291
292    /// Truncate the `.log` file and the indexes so that no batch at
293    /// `relative_offset` `>= rel` remains. `Log::truncate_to` uses this
294    /// method. The segment stays unsealed.
295    #[instrument(
296        level = "info",
297        skip(self),
298        fields(base_offset = self.base_offset.0, new_last_offset = field::Empty),
299        err,
300    )]
301    /// # Errors
302    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
303    pub fn truncate_to_relative(&mut self, rel: u32) -> Result<(), LogError> {
304        // Read only as far as the cut can be: every kept batch lives below
305        // the first index entry at or after `rel`. When `rel` is past the
306        // last index entry, fall back to the whole file. This avoids
307        // slurping the discarded tail on each truncate.
308        let read_limit = self
309            .offset_index
310            .position_at_or_after(rel)
311            .map_or(self.log_size, u64::from);
312        let mut buf = Vec::new();
313        let to_read = usize::try_from(read_limit).unwrap_or(usize::MAX);
314        self.read_log_range(0, &mut buf, to_read)?;
315
316        let target_abs = self
317            .base_offset
318            .0
319            .checked_add(i64::from(rel))
320            .map(Offset)
321            .ok_or_else(|| LogError::BadSegmentName("offset overflow".into()))?;
322        let mut cur: &[u8] = &buf;
323        let mut pos: u64 = 0;
324        let mut last_kept_offset = self.base_offset - 1;
325        let mut last_kept_ts = i64::MIN;
326        let mut last_kept_ts_offset = last_kept_offset;
327        while !cur.is_empty() {
328            let before = cur.len();
329            let Ok(batch) = RecordBatch::decode(&mut cur) else {
330                break;
331            };
332            let batch_last_offset = batch
333                .base_offset
334                .checked_add(i64::from(batch.last_offset_delta))
335                .map(Offset)
336                .ok_or_else(|| LogError::Corrupt("batch last offset overflow".into()))?;
337            if !krabka_verified::truncation_batch_retained(batch_last_offset.0, target_abs.0) {
338                break;
339            }
340            pos += (before - cur.len()) as u64;
341            last_kept_offset = batch_last_offset;
342            if batch.max_timestamp > last_kept_ts {
343                last_kept_ts = batch.max_timestamp;
344                last_kept_ts_offset = batch_last_offset;
345            }
346        }
347
348        self.settle_direct();
349        self.log_file.set_len(pos)?;
350        self.first_timestamp = None;
351        seek_to_log_size(&self.log_file, pos)?;
352        self.log_size = pos;
353        self.renew_reservation();
354        self.resync_direct()?;
355        self.last_offset = last_kept_offset;
356        self.max_timestamp = last_kept_ts;
357        self.max_timestamp_offset = last_kept_ts_offset;
358
359        let pos_u32 =
360            u32::try_from(pos).map_err(|_| LogError::BadSegmentName("position overflow".into()))?;
361        self.offset_index.truncate_by_position(pos_u32)?;
362        self.time_index.truncate_by_relative_offset(rel)?;
363        self.sealed = false;
364        tracing::Span::current().record("new_last_offset", self.last_offset.0);
365        Ok(())
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use {
372        super::*,
373        crate::segment::test_support::{
374            DENSE_INDEX, NO_LIMIT, SECOND_THREE_RECORD_BATCH, THREE_RECORD_BATCH, TempSegment,
375            sample_batch, seeded_fixture, segment_at, test_segment, two_record_fixture,
376        },
377        assert2::check,
378        krabka_units::prelude::{ByteSize, ByteSizeExt, kibibytes},
379        tempfile::tempdir,
380    };
381
382    /// Truncating to a relative offset keeps every batch that ends before it,
383    /// and leaves the segment describing exactly what it kept.
384    ///
385    /// Three things are rewritten from the walk and each is read back later:
386    /// the byte length, which decides where the next append lands; the last
387    /// offset, which decides what the next batch is numbered; and the maximum
388    /// timestamp, which time-based retention and `MAX_TIMESTAMP` both read.
389    /// Truncating everything away is the case that pins the last offset down --
390    /// it has to fall back to one before the base.
391    #[test]
392    fn truncating_a_segment_leaves_it_describing_what_it_kept() {
393        let dir = tempdir().unwrap();
394        let mut seg = Segment::create(dir.path(), Offset(100)).unwrap();
395        // Three batches: offsets 100..=101, 102..=103, 104..=105, with
396        // timestamps 500, 600, 700.
397        for i in 0..3i64 {
398            seg.append(
399                &sample_batch()
400                    .offset(100 + i * 2)
401                    .records(2)
402                    .timestamp(500 + i * 100)
403                    .call(),
404                DENSE_INDEX,
405            )
406            .unwrap();
407        }
408        let full_size = seg.size();
409        check!(seg.last_offset() == Offset(105));
410        check!(
411            seg.max_timestamp() == 701,
412            "batch 3 carries the newest record"
413        );
414
415        // Keep only the first batch: relative offset 2 is the start of the
416        // second, and the bound is exclusive of the batch containing it.
417        seg.truncate_to_relative(2).unwrap();
418        check!(
419            seg.last_offset() == Offset(101),
420            "last kept batch ends at 101"
421        );
422        check!(seg.max_timestamp() == 501, "the newest surviving record");
423        let kept = seg.size();
424        check!(
425            kept > ByteSize::ZERO && kept < full_size,
426            "shorter, not empty"
427        );
428
429        // Truncating everything away: nothing is kept, so the segment reports
430        // one before its base and no timestamp at all.
431        seg.truncate_to_relative(0).unwrap();
432        check!(seg.size() == ByteSize::ZERO, "no bytes survive");
433        check!(
434            seg.last_offset() == Offset(99),
435            "one before the base, got {:?}",
436            seg.last_offset()
437        );
438        check!(seg.max_timestamp() == i64::MIN, "no records, no timestamp");
439    }
440
441    /// Sealing is what `is_sealed` reports.
442    #[test]
443    fn is_sealed_follows_seal() {
444        let TempSegment {
445            dir: _dir,
446            segment: mut seg,
447        } = two_record_fixture();
448        check!(!seg.is_sealed(), "a fresh segment is open");
449        seg.seal().unwrap();
450        check!(seg.is_sealed(), "a sealed segment reports it");
451    }
452
453    /// Sealing writes Kafka's final time-index entry, the running maximum
454    /// timestamp with the last offset of the batch that set it, when the
455    /// sparse index has not already recorded it. A reopened segment restores
456    /// its `max_timestamp` from that entry, and the entries stay strictly
457    /// increasing. Each case is `(label, batches as (base offset, timestamp,
458    /// index interval), the entry count and last entry after sealing)`.
459    #[test]
460    fn seal_appends_the_final_time_index_entry() {
461        let sparse = kibibytes(4);
462        let cases = [
463            (
464                "no index point yet, and the newest timestamp came first",
465                vec![(0, 100, sparse), (1, 300, sparse), (2, 200, sparse)],
466                (1, Some((300, 1))),
467            ),
468            (
469                "the index already holds the newest timestamp",
470                vec![(0, 100, DENSE_INDEX), (1, 200, DENSE_INDEX)],
471                (1, Some((200, 1))),
472            ),
473            (
474                "an index point lags the newest batch",
475                vec![
476                    (0, 100, DENSE_INDEX),
477                    (1, 200, DENSE_INDEX),
478                    (2, 300, sparse),
479                ],
480                (2, Some((300, 2))),
481            ),
482        ];
483        for (label, batches, expected) in cases {
484            let TempSegment {
485                dir: _dir,
486                segment: mut seg,
487            } = test_segment();
488            for (base, timestamp, interval) in batches {
489                seg.append(
490                    &sample_batch().offset(base).timestamp(timestamp).call(),
491                    interval,
492                )
493                .unwrap();
494            }
495            seg.seal().unwrap();
496            check!(
497                (
498                    seg.time_index.entry_count(),
499                    seg.time_index
500                        .last_entry()
501                        .map(|entry| (entry.timestamp, entry.relative_offset))
502                ) == expected,
503                "{label}"
504            );
505        }
506    }
507
508    /// `truncate_to_relative` decides which batches to drop by each batch's
509    /// last offset, `batch.base_offset + last_offset_delta`, compared against
510    /// `target_abs`. MULTI-record batches make the `+` load-bearing. Batch A
511    /// spans 0..=2 and batch B spans 3..=5, so a truncate to rel 3, where
512    /// `target_abs = 3`, must keep A, whose last offset 2 is < 3, and drop B,
513    /// whose last offset 5 is >= 3. A mutation of `+` to `-` computes A's last
514    /// offset as -2 and B's as 1, so it wrongly keeps B and the read still
515    /// returns batch B.
516    #[test]
517    fn truncate_to_relative_uses_batch_last_offset() {
518        let TempSegment {
519            dir: _dir,
520            segment: mut seg,
521        } = seeded_fixture()
522            .batches(&[
523                THREE_RECORD_BATCH,        // offsets 0..=2
524                SECOND_THREE_RECORD_BATCH, // offsets 3..=5
525            ])
526            .call();
527        assert2::assert!(seg.last_offset() == 5);
528
529        // target_abs = base(0) + rel(3) = 3. Drop batches with last >= 3.
530        seg.truncate_to_relative(3).unwrap();
531        let read = seg.read(Offset(0), NO_LIMIT).unwrap();
532        assert2::assert!(seg.last_offset() == Offset(2));
533        assert2::assert!(
534            read == vec![
535                sample_batch()
536                    .offset(THREE_RECORD_BATCH.offset.0)
537                    .records(THREE_RECORD_BATCH.records.0)
538                    .timestamp(THREE_RECORD_BATCH.timestamp.0)
539                    .call()
540            ]
541        );
542        let expected_size: usize = read.iter().map(RecordBatch::encoded_len).sum();
543        assert2::assert!(seg.size().bytes_usize() == expected_size);
544    }
545
546    #[test]
547    fn truncate_to_relative_rejects_an_absolute_offset_overflow() {
548        let TempSegment {
549            dir: _dir,
550            segment: mut seg,
551        } = segment_at(Offset(i64::MAX));
552
553        let error = seg
554            .truncate_to_relative(1)
555            .expect_err("base plus relative cut must be checked");
556        assert2::assert!(error.to_string().contains("offset overflow"));
557    }
558
559    #[test]
560    fn flush_succeeds() {
561        let TempSegment {
562            dir: _dir,
563            segment: mut seg,
564        } = test_segment();
565        seg.append(&sample_batch().timestamp(42).call(), kibibytes(4))
566            .unwrap();
567        seg.flush().unwrap();
568    }
569}
krabka_log/segment/lifecycle.rs

krabka_log/segment/
lifecycle.rs
