krabka_log/segment/
append.rs

1//! The two append paths: the encoding one and the verbatim passthrough.
2//!
3//! Both write to the active `.log` file and update the same segment state and
4//! sparse indexes, so the shared side effects stay in one file. The verbatim
5//! path additionally has to leave the producer's CRC-covered bytes untouched.
6
7use {
8    super::{
9        Segment,
10        io::{write_all, write_all_vectored},
11        lifecycle::WriteSnapshot,
12    },
13    crate::{VerbatimBatch, error::LogError, index::OffsetIndexEntry},
14    bytes::Bytes,
15    krabka_ids::{LeaderEpoch, Offset},
16    krabka_protocol::records::{HEADER_LEN, RecordBatch, patch_base_offset_and_leader_epoch},
17    krabka_units::prelude::{ByteSize, ByteSizeExt},
18    std::{io, io::IoSlice},
19    tracing::{field, instrument},
20};
21
22/// One batch's encoded bytes, as an append path has them.
23#[derive(Clone, Copy)]
24enum BatchBytes<'a> {
25    /// The encoding path's single buffer.
26    Whole(&'a Bytes),
27    /// The verbatim path's patched header and the producer's untouched body.
28    Split(&'a [u8], &'a [u8]),
29}
30
31/// The slices of one batch handed to the Direct I/O writer, and the joined
32/// bytes that go into the tail cache.
33struct DirectParts<'a> {
34    parts: &'a [&'a [u8]],
35    cached: Bytes,
36}
37
38impl BatchBytes<'_> {
39    fn len(self) -> usize {
40        match self {
41            Self::Whole(bytes) => bytes.len(),
42            Self::Split(header, body) => header.len() + body.len(),
43        }
44    }
45}
46
47/// A verbatim batch together with the base offset it is written at.
48#[derive(Clone, Copy)]
49pub struct PlacedBatch<'a> {
50    pub(crate) batch: &'a VerbatimBatch,
51    pub(crate) base: Offset,
52}
53
54/// The last offset a batch ends at, and the segment's byte length after the
55/// batch is written.
56struct AppendCoordinates {
57    last_offset: i64,
58    size: u64,
59}
60
61impl Segment {
62    fn append_coordinates(
63        &self,
64        base: i64,
65        delta: i32,
66        len: usize,
67    ) -> Result<AppendCoordinates, LogError> {
68        let (last, _) = krabka_verified::local_append_coordinates(base, base, delta)
69            .ok_or_else(|| LogError::InvalidArgument("invalid batch offset interval".into()))?;
70        let len = u64::try_from(len)
71            .map_err(|_| LogError::InvalidArgument("encoded batch length overflow".into()))?;
72        let size = self
73            .log_size
74            .checked_add(len)
75            .ok_or_else(|| LogError::InvalidArgument("segment byte length overflow".into()))?;
76        Ok(AppendCoordinates {
77            last_offset: last,
78            size,
79        })
80    }
81
82    /// Append a record batch and return the byte position where the batch
83    /// starts.
84    ///
85    /// Side effects:
86    /// - Updates `log_size`, `max_timestamp`, and `last_offset`.
87    /// - Adds sparse index entries when the byte count since the last entry
88    ///   exceeds `index_interval`, which Kafka's `LogSegment.append` does not
89    ///   do for the first batch either.
90    #[instrument(
91        level = "debug",
92        skip(self, batch),
93        fields(
94            base_offset = self.base_offset.0,
95            batch_base = batch.base_offset,
96            bytes = batch.encoded_len(),
97            position = field::Empty,
98        ),
99        err,
100    )]
101    /// # Errors
102    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
103    pub fn append(
104        &mut self,
105        batch: &RecordBatch,
106        index_interval: ByteSize,
107    ) -> Result<u64, LogError> {
108        self.ensure_writable()?;
109
110        let mut buf = bytes::BytesMut::with_capacity(batch.encoded_len());
111        batch.encode(&mut buf)?;
112        let bytes = buf.freeze();
113
114        self.append_encoded(batch, index_interval, &bytes)
115    }
116
117    /// Write an owned batch encoded at the caller's chosen compression level.
118    pub(crate) fn append_encoded(
119        &mut self,
120        batch: &RecordBatch,
121        index_interval: ByteSize,
122        bytes: &Bytes,
123    ) -> Result<u64, LogError> {
124        self.ensure_writable()?;
125        // Owned appends retain the single-buffer write path and its I/O seam.
126        self.write_batch(
127            batch.base_offset,
128            batch.last_offset_delta,
129            batch.max_timestamp,
130            index_interval,
131            BatchBytes::Whole(bytes),
132        )
133    }
134
135    /// Append a batch **verbatim** and write the producer's exact wire bytes.
136    ///
137    /// This method does not decode, re-encode, recompress, or recompute the
138    /// CRC. `bytes` is the producer's verbatim v2 batch, which the caller has
139    /// already CRC-validated through the borrowed header-only path. This
140    /// method patches only `base_offset`, bytes 0..8, and
141    /// `partition_leader_epoch`, bytes 12..16, into a writable copy, then
142    /// writes those bytes. Both fields sit outside the CRC-covered region. The
143    /// stored CRC stays byte-identical to the producer's CRC, because no
144    /// CRC-covered byte changes.
145    ///
146    /// `base_offset`, `last_offset_delta`, `max_timestamp`, and
147    /// `leader_epoch` come from the caller's borrowed header read. This method
148    /// updates the segment side effects, `log_size`, `last_offset`,
149    /// `max_timestamp`, and the sparse index, exactly as [`Segment::append`]
150    /// does.
151    ///
152    /// It returns the byte position where the batch starts.
153    #[instrument(
154        level = "debug",
155        skip(self, bytes),
156        fields(
157            seg_base_offset = self.base_offset.0,
158            base_offset = base_offset.0,
159            bytes_len = bytes.len(),
160            position = field::Empty,
161        ),
162        err,
163    )]
164    /// # Errors
165    /// Returns an error when log I/O fails, a record or index is corrupt, or the requested offset violates the segment state.
166    pub fn append_verbatim(
167        &mut self,
168        bytes: &[u8],
169        base_offset: Offset,
170        last_offset_delta: i32,
171        max_timestamp: i64,
172        leader_epoch: LeaderEpoch,
173        index_interval: ByteSize,
174    ) -> Result<u64, LogError> {
175        self.ensure_writable()?;
176        let Some((&header, body)) = bytes.split_first_chunk::<HEADER_LEN>() else {
177            return Err(LogError::Corrupt(
178                "verbatim batch shorter than v2 header".into(),
179            ));
180        };
181
182        // Patch base_offset + partition_leader_epoch in a copy of *just* the
183        // fixed-size header — both fields live below byte 16, well under the
184        // CRC-covered region (byte 21), so the producer's CRC stays valid (no
185        // recompute). The batch BODY is written straight from the input slice
186        // with no copy: the previous `bytes.to_vec()` was a full-payload memcpy
187        // on the produce hot path (100 KiB+ per batch for large messages), the
188        // dominant remaining produce-side cost.
189        let mut header = header;
190        // The protocol patcher writes the raw KIP-320 wire `int32`; unwrap here.
191        patch_base_offset_and_leader_epoch(&mut header, base_offset.0, leader_epoch.0);
192
193        self.write_batch(
194            base_offset.0,
195            last_offset_delta,
196            max_timestamp,
197            index_interval,
198            BatchBytes::Split(&header, body),
199        )
200    }
201
202    /// Write one batch's bytes at the end of the `.log` file.
203    ///
204    /// Through the page cache, the buffered handle's cursor is kept at
205    /// `log_size` by open, recovery and truncate, so the hot append path
206    /// needs no `lseek` before a write, and a split batch goes out in one
207    /// `writev` without a copy. Through `O_DIRECT`, the writer copies the
208    /// batch behind the partial last block and writes whole blocks, and the
209    /// batch joins the bytes kept in memory for reads.
210    fn write_bytes(&mut self, batch: BatchBytes<'_>, chunk_ends: &[usize]) -> io::Result<()> {
211        let Some(mode) = self.direct.as_mut() else {
212            return match batch {
213                BatchBytes::Whole(bytes) => write_all(&*self.io, &self.log_file, bytes),
214                BatchBytes::Split(header, body) => {
215                    write_all_vectored(&*self.io, &self.log_file, &mut [
216                        IoSlice::new(header),
217                        IoSlice::new(body),
218                    ])
219                }
220            };
221        };
222        let DirectParts { parts, cached } = match batch {
223            BatchBytes::Whole(bytes) => DirectParts {
224                parts: &[bytes],
225                cached: bytes.clone(),
226            },
227            BatchBytes::Split(header, body) => DirectParts {
228                parts: &[header, body],
229                cached: Bytes::from([header, body].concat()),
230            },
231        };
232        mode.writer.append(parts)?;
233        mode.cache.set_written(mode.writer.written());
234        if chunk_ends.is_empty() {
235            mode.cache.push(cached);
236        } else {
237            mode.cache.push_group(&cached, chunk_ends);
238        }
239        Ok(())
240    }
241
242    fn ensure_writable(&self) -> Result<(), LogError> {
243        if self.sealed {
244            return Err(LogError::Io(io::Error::other("segment is sealed")));
245        }
246        Ok(())
247    }
248
249    fn write_batch(
250        &mut self,
251        base: i64,
252        delta: i32,
253        max_timestamp: i64,
254        index_interval: ByteSize,
255        batch: BatchBytes<'_>,
256    ) -> Result<u64, LogError> {
257        let AppendCoordinates {
258            last_offset,
259            size: new_log_size,
260        } = self.append_coordinates(base, delta, batch.len())?;
261        let position = self.log_size;
262        let previous = self.write_snapshot();
263        let index_entry = self.index_entry_for(position, last_offset, index_interval)?;
264        if let Err(error) = self.write_bytes(batch, &[]) {
265            self.rollback_failed_write(position, previous)?;
266            return Err(error.into());
267        }
268        self.log_size = new_log_size;
269        self.finish_append(position, previous, index_entry, last_offset, max_timestamp)
270    }
271
272    /// Write a validated group once, then build each batch's sparse indexes.
273    /// Any failure restores the segment to the start of the group.
274    pub(crate) fn append_verbatim_group(
275        &mut self,
276        batches: &[PlacedBatch<'_>],
277        index_interval: ByteSize,
278    ) -> Result<(), LogError> {
279        self.ensure_writable()?;
280        let size = batches.iter().map(|placed| placed.batch.bytes.len()).sum();
281        let mut bytes = bytes::BytesMut::with_capacity(size);
282        let mut chunk_ends = Vec::with_capacity(batches.len());
283        for &PlacedBatch { batch, base } in batches {
284            let Some((&header, body)) = batch.bytes.split_first_chunk::<HEADER_LEN>() else {
285                return Err(LogError::Corrupt(
286                    "verbatim batch shorter than v2 header".into(),
287                ));
288            };
289            let mut header = header;
290            patch_base_offset_and_leader_epoch(&mut header, base.0, batch.leader_epoch.0);
291            bytes.extend_from_slice(&header);
292            bytes.extend_from_slice(body);
293            chunk_ends.push(bytes.len());
294        }
295        let bytes = bytes.freeze();
296        let start = self.log_size;
297        let snapshot = self.write_snapshot();
298        let result = (|| {
299            self.write_bytes(BatchBytes::Whole(&bytes), &chunk_ends)?;
300            for &PlacedBatch { batch, base } in batches {
301                let position = self.log_size;
302                let AppendCoordinates {
303                    last_offset: last,
304                    size,
305                } = self.append_coordinates(base.0, batch.last_offset_delta, batch.bytes.len())?;
306                let entry = self.index_entry_for(position, last, index_interval)?;
307                self.log_size = size;
308                self.record_batch(last, batch.max_timestamp);
309                if let Some(entry) = entry {
310                    self.append_index_entries(entry)?;
311                }
312            }
313            Ok(())
314        })();
315        if result.is_err() {
316            self.rollback_failed_write(start, snapshot)?;
317        }
318        result
319    }
320
321    /// Publish the completed write and roll back its bytes if an index write fails.
322    fn finish_append(
323        &mut self,
324        position: u64,
325        previous: WriteSnapshot,
326        index_entry: Option<OffsetIndexEntry>,
327        last_offset: i64,
328        max_timestamp: i64,
329    ) -> Result<u64, LogError> {
330        self.record_batch(last_offset, max_timestamp);
331        if let Some(entry) = index_entry
332            && let Err(error) = self.append_index_entries(entry)
333        {
334            self.rollback_failed_write(position, previous)?;
335            return Err(error);
336        }
337        tracing::Span::current().record("position", position);
338        Ok(position)
339    }
340
341    /// Whether the batch about to be written at `position` gets a sparse-index
342    /// point, and if so its offset-index entry: the batch's **last** offset,
343    /// relative to the segment base, with the position where the batch starts.
344    /// Kafka indexes a batch by its last offset (`LogSegment.append` writes
345    /// `offsetIndex().append(batchLastOffset, physicalPosition)`), which is
346    /// what `kafka-dump-log` checks against the batch at that position.
347    ///
348    /// The conversions can fail, so they run before anything is written.
349    fn index_entry_for(
350        &self,
351        position: u64,
352        last_offset: i64,
353        index_interval: ByteSize,
354    ) -> Result<Option<OffsetIndexEntry>, LogError> {
355        // Kafka's `bytesSinceLastIndexEntry > indexIntervalBytes`: the bytes
356        // written since the last entry, counted from that entry's batch (or
357        // from the segment start while there is none), must exceed the
358        // interval. So the first batch of a segment is never indexed, and a
359        // batch exactly one interval past the last entry is not either.
360        let last_indexed = self
361            .offset_index
362            .last_entry()
363            .map_or(0, |entry| u64::from(entry.position));
364        let should_index = position.saturating_sub(last_indexed) > index_interval.bytes_u64();
365        if !should_index {
366            return Ok(None);
367        }
368        let rel = u32::try_from(last_offset - self.base_offset.0)
369            .map_err(|_| LogError::BadSegmentName("offset overflow in segment".into()))?;
370        let pos = u32::try_from(position)
371            .map_err(|_| LogError::BadSegmentName("position overflow in segment".into()))?;
372        Ok(Some(OffsetIndexEntry {
373            relative_offset: rel,
374            position: pos,
375        }))
376    }
377
378    /// Account for a batch just written: it moves the last offset, and it
379    /// becomes the running maximum timestamp's batch when its timestamp is
380    /// strictly higher, Kafka's `maxTimestampAndOffsetSoFar`.
381    fn record_batch(&mut self, last_offset: i64, max_timestamp: i64) {
382        self.last_offset = Offset(last_offset);
383        if max_timestamp > self.max_timestamp {
384            self.max_timestamp = max_timestamp;
385            self.max_timestamp_offset = Offset(last_offset);
386        }
387    }
388
389    /// Write one sparse-index point. The time index gets the running maximum
390    /// timestamp with the last offset of the batch that set it, and only when
391    /// that timestamp is newer than its last entry, as Kafka's
392    /// `timeIndex().maybeAppend(maxTimestampSoFar(), shallowOffsetOfMaxTimestampSoFar())`.
393    fn append_index_entries(&mut self, entry: OffsetIndexEntry) -> Result<(), LogError> {
394        self.offset_index
395            .append(entry.relative_offset, entry.position)?;
396        self.append_running_max_time_entry()
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use {
403        super::*,
404        crate::{
405            name::log_path,
406            segment::{
407                io::SourceIo as _,
408                test_support::{
409                    DENSE_INDEX, FIRST_LARGE_TIMESTAMP_BATCH, NO_LIMIT, ONE_RECORD_BATCH,
410                    SECOND_LARGE_TIMESTAMP_BATCH, SECOND_SINGLE_RECORD_BATCH, TempSegment,
411                    sample_batch, seeded_fixture, segment_at, test_batch_at, test_segment,
412                },
413            },
414        },
415        krabka_protocol::records::RecordBatch,
416        krabka_units::prelude::kibibytes,
417        std::fs,
418    };
419
420    #[test]
421    fn append_then_read_back() {
422        let TempSegment {
423            dir: _dir,
424            segment: mut seg,
425        } = test_segment();
426        let b1 = sample_batch()
427            .offset(FIRST_LARGE_TIMESTAMP_BATCH.offset.0)
428            .records(FIRST_LARGE_TIMESTAMP_BATCH.records.0)
429            .timestamp(FIRST_LARGE_TIMESTAMP_BATCH.timestamp.0)
430            .call();
431        let b2 = sample_batch()
432            .offset(SECOND_LARGE_TIMESTAMP_BATCH.offset.0)
433            .records(SECOND_LARGE_TIMESTAMP_BATCH.records.0)
434            .timestamp(SECOND_LARGE_TIMESTAMP_BATCH.timestamp.0)
435            .call();
436        seg.append(&b1, kibibytes(4)).unwrap();
437        seg.append(&b2, kibibytes(4)).unwrap();
438        let read = seg.read(Offset(0), NO_LIMIT).unwrap();
439        assert2::assert!(seg.last_offset() == Offset(4));
440        assert2::assert!(read == vec![b1, b2]);
441    }
442
443    #[test]
444    fn append_rejects_offset_successor_overflow_before_writing() {
445        let TempSegment {
446            dir,
447            segment: mut seg,
448        } = segment_at(Offset(i64::MAX));
449        let batch = sample_batch().offset(i64::MAX).timestamp(100).call();
450
451        let error = seg.append(&batch, DENSE_INDEX).unwrap_err();
452
453        assert2::assert!(matches!(error, LogError::InvalidArgument(_)));
454        assert2::assert!(seg.size().bytes_u64() == 0);
455        assert2::assert!(fs::metadata(log_path(dir.path(), i64::MAX)).unwrap().len() == 0);
456    }
457
458    #[test]
459    fn append_after_truncate_writes_at_new_eof() {
460        let TempSegment {
461            dir: _dir,
462            segment: mut seg,
463        } = seeded_fixture().batches(&[ONE_RECORD_BATCH]).call();
464        let expected_position = seg.size().bytes_u64();
465        seg.append(
466            &sample_batch()
467                .offset(SECOND_SINGLE_RECORD_BATCH.offset.0)
468                .records(SECOND_SINGLE_RECORD_BATCH.records.0)
469                .timestamp(SECOND_SINGLE_RECORD_BATCH.timestamp.0)
470                .call(),
471            DENSE_INDEX,
472        )
473        .unwrap();
474
475        seg.truncate_to_relative(1).unwrap();
476        let position = seg
477            .append(&sample_batch().offset(1).timestamp(300).call(), DENSE_INDEX)
478            .unwrap();
479
480        let read = seg.read(Offset(0), NO_LIMIT).unwrap();
481        assert2::assert!(position == expected_position);
482        assert2::assert!(seg.last_offset() == Offset(1));
483        assert2::assert!(
484            read == vec![
485                sample_batch()
486                    .offset(ONE_RECORD_BATCH.offset.0)
487                    .records(ONE_RECORD_BATCH.records.0)
488                    .timestamp(ONE_RECORD_BATCH.timestamp.0)
489                    .call(),
490                sample_batch().offset(1).timestamp(300).call()
491            ]
492        );
493    }
494
495    #[test]
496    fn append_to_sealed_segment_errors() {
497        let TempSegment {
498            dir: _dir,
499            segment: mut seg,
500        } = test_segment();
501        seg.seal().unwrap();
502        assert2::assert!(seg.is_sealed());
503        let err = seg
504            .append(&sample_batch().call(), kibibytes(4))
505            .unwrap_err();
506        assert2::assert!(matches!(err, LogError::Io(_)));
507    }
508
509    #[test]
510    fn append_verbatim_is_byte_exact_except_offset_and_epoch() {
511        let TempSegment {
512            dir,
513            segment: mut seg,
514        } = test_segment();
515        // Build a batch as a "producer" would, with its own base_offset and
516        // leader epoch, then encode to verbatim wire bytes.
517        let mut producer = test_batch_at(Offset(0));
518        producer.base_offset = 999; // producer-supplied (to be overwritten)
519        producer.partition_leader_epoch = -1; // producer-supplied
520        producer.last_offset_delta = 0;
521        producer.max_timestamp = 1_000;
522        let mut wire = bytes::BytesMut::new();
523        producer.encode(&mut wire).unwrap();
524        let wire = wire.freeze();
525
526        // Append verbatim with an assigned base_offset and a stamped epoch.
527        let assigned_base = Offset(0);
528        let stamped_epoch = 7i32;
529        seg.append_verbatim(
530            &wire,
531            assigned_base,
532            0,
533            1_000,
534            LeaderEpoch(stamped_epoch),
535            DENSE_INDEX,
536        )
537        .unwrap();
538        assert2::assert!(seg.last_offset() == 0);
539
540        // Read back the raw .log bytes.
541        let mut on_disk = Vec::new();
542        seg.read_log_range(0, &mut on_disk, usize::MAX).unwrap();
543        let mut expected_wire = wire.to_vec();
544        expected_wire[0..8].copy_from_slice(&assigned_base.0.to_be_bytes());
545        expected_wire[12..16].copy_from_slice(&stamped_epoch.to_be_bytes());
546        assert2::assert!(on_disk == expected_wire);
547
548        // And it decodes (CRC still valid).
549        let mut cur: &[u8] = &on_disk;
550        let decoded = RecordBatch::decode(&mut cur).unwrap();
551        assert2::assert!(decoded.base_offset == assigned_base.0);
552        assert2::assert!(decoded.partition_leader_epoch == stamped_epoch);
553        drop(dir);
554    }
555
556    #[test]
557    fn append_verbatim_updates_index_and_last_offset() {
558        let TempSegment {
559            dir,
560            segment: mut seg,
561        } = test_segment();
562        let mut producer = test_batch_at(Offset(0));
563        producer.last_offset_delta = 2; // spans 3 offsets
564        producer.max_timestamp = 5_000;
565        let mut wire = bytes::BytesMut::new();
566        producer.encode(&mut wire).unwrap();
567        let wire = wire.freeze();
568
569        seg.append_verbatim(&wire, Offset(0), 2, 5_000, LeaderEpoch(0), DENSE_INDEX)
570            .unwrap();
571        // Reading at offset 2 (inside the batch) returns the batch.
572        let read = seg.read(Offset(2), NO_LIMIT).unwrap();
573        assert2::assert!(seg.last_offset() == Offset(2));
574        assert2::assert!(seg.max_timestamp() == 5_000);
575        assert2::assert!(read == vec![producer]);
576        drop(dir);
577    }
578
579    #[test]
580    fn append_verbatim_to_sealed_segment_errors() {
581        let TempSegment {
582            dir,
583            segment: mut seg,
584        } = test_segment();
585        seg.seal().unwrap();
586        let mut wire = bytes::BytesMut::new();
587        test_batch_at(Offset(0)).encode(&mut wire).unwrap();
588        let err = seg
589            .append_verbatim(&wire.freeze(), Offset(0), 0, 0, LeaderEpoch(0), DENSE_INDEX)
590            .unwrap_err();
591        assert2::assert!(matches!(err, LogError::Io(_)));
592        drop(dir);
593    }
594}
krabka_log/segment/append.rs

krabka_log/segment/
append.rs
