krabka_remote_storage_topic/
snapshot.rs

1//! On-disk RLMM snapshot: a versioned envelope wrapping a
2//! [`RlmmCacheDump`] plus the per-metadata-partition committed
3//! offsets.
4//!
5//! The offsets let a restarting broker resume the metadata consumer from
6//! `committed + 1` instead of replaying `__remote_log_metadata` from
7//! offset 0.
8//!
9//! The per-segment and per-partition-delete encoding reuses the
10//! [`MetadataEvent`] codec. The envelope adds a format version, the
11//! committed-offsets vector, and length-prefixed entries. Writes are atomic,
12//! through a temp file and a rename, so a crash mid-write never yields a torn
13//! snapshot. A corrupt or truncated file decodes to an error and never to a
14//! panic, and the caller then falls back to a full replay.
15
16use {
17    crate::{
18        error::{CodecError, SnapshotError},
19        serde::{MetadataEvent, Reader, read_uvarint, write_uvarint},
20    },
21    bytes::{BufMut, BytesMut},
22    krabka_remote_storage::{PartitionDump, RlmmCacheDump},
23    krabka_units::prelude::{ByteSize, ByteSizeExt as _, bytes as byte_size},
24    std::{
25        collections::BTreeMap,
26        fs::{self, File},
27        io::ErrorKind,
28        path::Path,
29    },
30    tracing::instrument,
31};
32
33/// Format version at the head of every snapshot file.
34///
35/// Part of the 1.x on-disk contract: a 1.x broker reads every snapshot that
36/// any earlier 1.x broker wrote, so a change to the layout bumps this version
37/// and keeps a decoder arm for every earlier 1.x version. See
38/// `docs/persisted_formats.md`.
39pub const SNAPSHOT_FORMAT_VERSION: u16 = 0;
40
41/// Default snapshot file name under the snapshot directory.
42pub const SNAPSHOT_FILE_NAME: &str = "snapshot";
43
44/// Initial encode-buffer reservation. It covers the envelope and a few
45/// entries, and the buffer grows for anything larger.
46const ENCODE_BUFFER_HINT: ByteSize = byte_size(256);
47
48/// A decoded snapshot: the per-metadata-partition committed offsets and
49/// the cache dump to seed an [`InmemoryRemoteLogMetadataManager`].
50///
51/// [`InmemoryRemoteLogMetadataManager`]: krabka_remote_storage::InmemoryRemoteLogMetadataManager
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct Snapshot {
54    /// Highest offset applied into the cache per metadata partition, indexed
55    /// by metadata-partition. `committed[p] == -1` means the snapshot covers
56    /// no events for `p`.
57    pub committed_offsets: Vec<i64>,
58    /// The cache contents at the moment the snapshot was taken.
59    pub dump: RlmmCacheDump,
60}
61
62/// A partition's identity in a snapshot being decoded.
63#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
64struct PartitionKey {
65    topic_id: uuid::Uuid,
66    partition: i32,
67}
68
69/// A partition's dump and the order its first event appeared in.
70struct OrderedDump {
71    ordinal: usize,
72    dump: PartitionDump,
73}
74
75impl Snapshot {
76    /// Encode this snapshot into freshly-allocated bytes.
77    ///
78    /// # Errors
79    ///
80    /// Returns [`SnapshotError::Malformed`] if `committed_offsets` contains
81    /// more than `i32::MAX + 1` offsets, because its partition indices must
82    /// fit in the encoded `i32` field, or if an entry fails to encode.
83    pub fn encode(&self) -> Result<Vec<u8>, SnapshotError> {
84        if self.committed_offsets.len() > i32::MAX as usize + 1 {
85            return Err(SnapshotError::Malformed(CodecError::LengthOverflow(
86                self.committed_offsets.len() as u64,
87            )));
88        }
89        let mut buf = BytesMut::with_capacity(ENCODE_BUFFER_HINT.bytes_usize());
90        buf.put_u16(SNAPSHOT_FORMAT_VERSION);
91        // Committed offsets.
92        write_uvarint(self.committed_offsets.len() as u64, &mut buf);
93        // The check above leaves an `i32` index for every offset.
94        for (&off, p) in self.committed_offsets.iter().zip(0..=i32::MAX) {
95            buf.put_i32(p);
96            buf.put_i64(off);
97        }
98        // Entries: one MetadataEvent per segment, then one per
99        // partition-delete (so import sees segments before delete state).
100        let mut entries: Vec<bytes::Bytes> = Vec::new();
101        for p in &self.dump.partitions {
102            for seg in &p.segments {
103                entries.push(MetadataEvent::AddSegment(seg.clone()).encode()?);
104            }
105            if let Some(state) = p.delete_state {
106                entries.push(
107                    MetadataEvent::PartitionDelete(
108                        krabka_remote_storage::RemotePartitionDeleteMetadata {
109                            topic_id_partition: p.topic_id_partition.clone(),
110                            state,
111                            event_timestamp_ms: 0,
112                            broker_id: 0,
113                        },
114                    )
115                    .encode()?,
116                );
117            }
118        }
119        write_uvarint(entries.len() as u64, &mut buf);
120        for entry in entries {
121            write_uvarint(entry.len() as u64, &mut buf);
122            buf.put_slice(&entry);
123        }
124        Ok(buf.to_vec())
125    }
126
127    /// Decode a snapshot from `bytes`.
128    ///
129    /// # Errors
130    ///
131    /// Returns [`SnapshotError`] for any malformed input: a bad version, a
132    /// short or truncated buffer, a contained event that fails to decode, or
133    /// trailing bytes after the declared entries.
134    pub fn decode(bytes: &[u8]) -> Result<Self, SnapshotError> {
135        let mut r = Reader::new(bytes);
136        let version = read_u16(&mut r)?;
137        if version != SNAPSHOT_FORMAT_VERSION {
138            return Err(SnapshotError::UnsupportedVersion(version));
139        }
140        let n_offsets = usize::try_from(read_uvarint(&mut r)?)
141            .map_err(|_| CodecError::LengthOverflow(u64::MAX))?;
142        // Every offset entry takes an i32 partition and an i64 offset. Check
143        // the input before allocating from its untrusted count.
144        if n_offsets > r.remaining() / 12 {
145            return Err(CodecError::UnexpectedEof(bytes.len()).into());
146        }
147        let mut committed_offsets = vec![-1i64; n_offsets];
148        for _ in 0..n_offsets {
149            let p = r.read_i32()?;
150            let off = r.read_i64()?;
151            let idx = usize::try_from(p).map_err(|_| CodecError::LengthOverflow(u64::MAX))?;
152            let Some(slot) = committed_offsets.get_mut(idx) else {
153                return Err(SnapshotError::Malformed(CodecError::LengthOverflow(
154                    idx as u64,
155                )));
156            };
157            *slot = off;
158        }
159        let n_entries = usize::try_from(read_uvarint(&mut r)?)
160            .map_err(|_| CodecError::LengthOverflow(u64::MAX))?;
161
162        // Accumulate per-partition while preserving first-seen order.
163        let mut by_tp: BTreeMap<PartitionKey, OrderedDump> = BTreeMap::new();
164        for _ in 0..n_entries {
165            let len = usize::try_from(read_uvarint(&mut r)?)
166                .map_err(|_| CodecError::LengthOverflow(u64::MAX))?;
167            let raw = r.read_n(len)?;
168            match MetadataEvent::decode(raw)? {
169                MetadataEvent::AddSegment(md) => {
170                    let tp = md.remote_log_segment_id().topic_id_partition.clone();
171                    let key = PartitionKey {
172                        topic_id: tp.topic_id,
173                        partition: tp.partition,
174                    };
175                    let ordinal = by_tp.len();
176                    by_tp
177                        .entry(key)
178                        .or_insert_with(|| OrderedDump {
179                            ordinal,
180                            dump: PartitionDump {
181                                topic_id_partition: tp,
182                                segments: Vec::new(),
183                                delete_state: None,
184                            },
185                        })
186                        .dump
187                        .segments
188                        .push(md);
189                }
190                MetadataEvent::PartitionDelete(d) => {
191                    let key = PartitionKey {
192                        topic_id: d.topic_id_partition.topic_id,
193                        partition: d.topic_id_partition.partition,
194                    };
195                    let ordinal = by_tp.len();
196                    by_tp
197                        .entry(key)
198                        .or_insert_with(|| OrderedDump {
199                            ordinal,
200                            dump: PartitionDump {
201                                topic_id_partition: d.topic_id_partition,
202                                segments: Vec::new(),
203                                delete_state: None,
204                            },
205                        })
206                        .dump
207                        .delete_state = Some(d.state);
208                }
209                MetadataEvent::UpdateSegment(_) => {
210                    // Snapshots only ever encode Add + PartitionDelete;
211                    // segment updates are folded into the dumped cache, not
212                    // stored as standalone events.
213                    return Err(SnapshotError::Malformed(CodecError::Domain(
214                        "snapshot must not contain an UpdateSegment event".to_owned(),
215                    )));
216                }
217            }
218        }
219        if r.remaining() != 0 {
220            return Err(SnapshotError::TrailingBytes(r.remaining()));
221        }
222        let mut ordered: Vec<_> = by_tp.into_values().collect();
223        ordered.sort_unstable_by_key(|entry| entry.ordinal);
224        let partitions = ordered.into_iter().map(|entry| entry.dump).collect();
225        Ok(Self {
226            committed_offsets,
227            dump: RlmmCacheDump { partitions },
228        })
229    }
230
231    /// Atomically write this snapshot to `path`.
232    ///
233    /// This method writes to a sibling temp file, fsyncs the file, renames it
234    /// over `path`, then fsyncs the parent directory so the rename itself is
235    /// durable. A crash at any point leaves either the old snapshot or none,
236    /// and never a torn file. This method removes the temp file if any step
237    /// before the rename fails, so a persistent error does not leave a stale
238    /// `.tmp` behind.
239    ///
240    /// The directory fsync is best-effort. This method ignores filesystems
241    /// that do not support it, and platforms where an open of a directory
242    /// fails, rather than fails the write.
243    ///
244    /// # Errors
245    ///
246    /// Returns [`SnapshotError::Io`] on any filesystem failure up to and
247    /// including the rename.
248    #[instrument(
249        skip_all,
250        fields(path = %path.display(), partitions = self.dump.partitions.len()),
251        err
252    )]
253    pub fn write_atomic(&self, path: &Path) -> Result<(), SnapshotError> {
254        use std::io::Write;
255        let bytes = self.encode()?;
256        let parent = path.parent().unwrap_or_else(|| Path::new("."));
257        fs::create_dir_all(parent)?;
258        let tmp = path.with_extension("tmp");
259        let write_tmp = || -> Result<(), SnapshotError> {
260            let mut f = File::create(&tmp)?;
261            f.write_all(&bytes)?;
262            f.sync_all()?;
263            fs::rename(&tmp, path)?;
264            Ok(())
265        };
266        if let Err(e) = write_tmp() {
267            // Don't leave a stale temp file behind on a failed write.
268            let _ = fs::remove_file(&tmp);
269            return Err(e);
270        }
271        // Best-effort durable rename: fsync the parent directory so the
272        // rename survives a crash. Ignore unsupported/not-a-dir cases.
273        if let Ok(dir) = File::open(parent) {
274            let _ = dir.sync_all();
275        }
276        Ok(())
277    }
278
279    /// Load a snapshot from `path`. Returns `Ok(None)` when the file does
280    /// not exist, which is the first boot. Returns `Err` when the file exists
281    /// but is corrupt. Returns `Ok(Some)` on success.
282    ///
283    /// # Errors
284    ///
285    /// Returns [`SnapshotError::Io`] for read failures other than
286    /// not-found, or a decode error for a present-but-malformed file.
287    #[instrument(skip_all, fields(path = %path.display()), err)]
288    pub fn load(path: &Path) -> Result<Option<Self>, SnapshotError> {
289        let bytes = match fs::read(path) {
290            Ok(b) => b,
291            Err(e) if e.kind() == ErrorKind::NotFound => return Ok(None),
292            Err(e) => return Err(SnapshotError::Io(e)),
293        };
294        Ok(Some(Self::decode(&bytes)?))
295    }
296}
297
298fn read_u16(r: &mut Reader<'_>) -> Result<u16, CodecError> {
299    let hi = u16::from(r.read_u8()?);
300    let lo = u16::from(r.read_u8()?);
301    Ok((hi << 8) | lo)
302}
303
304#[cfg(test)]
305mod tests {
306    use {
307        super::*,
308        crate::manager::test_support::{started, tp},
309        assert2::assert,
310        krabka_remote_storage::RemotePartitionDeleteState,
311        std::{env, fs, process},
312    };
313
314    fn sample_snapshot() -> Snapshot {
315        let dump = RlmmCacheDump {
316            partitions: vec![PartitionDump {
317                topic_id_partition: tp(),
318                segments: vec![started(10, 0, 99), started(11, 100, 199)],
319                delete_state: Some(RemotePartitionDeleteState::DeletePartitionMarked),
320            }],
321        };
322        Snapshot {
323            committed_offsets: vec![5, -1, 2],
324            dump,
325        }
326    }
327
328    #[test]
329    fn encode_decode_round_trip() {
330        let snap = sample_snapshot();
331        let bytes = snap.encode().unwrap();
332        let back = Snapshot::decode(&bytes).expect("decodes");
333        assert!(back == snap);
334    }
335
336    #[test]
337    fn decode_preserves_first_seen_partition_order() {
338        let mut snap = sample_snapshot();
339        snap.dump.partitions = [2, 0, 1]
340            .into_iter()
341            .map(|partition| {
342                let mut topic_id_partition = tp();
343                topic_id_partition.partition = partition;
344                PartitionDump {
345                    topic_id_partition,
346                    segments: Vec::new(),
347                    delete_state: Some(RemotePartitionDeleteState::DeletePartitionMarked),
348                }
349            })
350            .collect();
351        assert!(Snapshot::decode(&snap.encode().unwrap()).unwrap() == snap);
352    }
353
354    #[test]
355    fn truncated_file_is_error_not_panic() {
356        let bytes = sample_snapshot().encode().unwrap();
357        let err = Snapshot::decode(&bytes[..bytes.len() - 3]).unwrap_err();
358        assert!(matches!(
359            err,
360            SnapshotError::Malformed(_) | SnapshotError::TrailingBytes(_)
361        ));
362    }
363
364    #[yare::parameterized(
365        one = { 1 },
366        two = { 2 },
367        maximum = { u64::MAX },
368    )]
369    fn offset_counts_are_bounded_by_the_encoded_entries(count: u64) {
370        let mut bytes = BytesMut::new();
371        bytes.put_u16(SNAPSHOT_FORMAT_VERSION);
372        write_uvarint(count, &mut bytes);
373        assert2::assert!(matches!(
374            Snapshot::decode(&bytes),
375            Err(SnapshotError::Malformed(
376                CodecError::UnexpectedEof(_) | CodecError::LengthOverflow(_)
377            ))
378        ));
379    }
380
381    #[test]
382    fn valid_snapshot_round_trips_after_malformed_count_cases() {
383        let snapshot = Snapshot {
384            committed_offsets: vec![-1],
385            dump: RlmmCacheDump::default(),
386        };
387        assert!(Snapshot::decode(&snapshot.encode().unwrap()).unwrap() == snapshot);
388    }
389
390    #[test]
391    fn garbage_bytes_are_error_not_panic() {
392        let err = Snapshot::decode(&[0xFF, 0xFF, 0xFF, 0xFF]).unwrap_err();
393        assert!(matches!(
394            err,
395            SnapshotError::UnsupportedVersion(_) | SnapshotError::Malformed(_)
396        ));
397    }
398
399    #[test]
400    fn empty_buffer_is_error_not_panic() {
401        let err = Snapshot::decode(&[]).unwrap_err();
402        assert!(matches!(err, SnapshotError::Malformed(_)));
403    }
404
405    #[test]
406    fn write_then_load_round_trips_through_a_file() {
407        let dir = env::temp_dir().join(format!("krabka-snap-{}", process::id()));
408        fs::create_dir_all(&dir).unwrap();
409        let path = dir.join("snapshot");
410        let snap = sample_snapshot();
411        snap.write_atomic(&path).expect("write");
412        let loaded = Snapshot::load(&path).expect("load").expect("present");
413        assert!(loaded == snap);
414        // No temp file left behind.
415        assert!(
416            fs::read_dir(&dir)
417                .unwrap()
418                .filter_map(Result::ok)
419                .all(|e| e.file_name() == "snapshot")
420        );
421        let _ = fs::remove_dir_all(&dir);
422    }
423
424    #[test]
425    fn load_absent_file_is_ok_none() {
426        let path = env::temp_dir().join("krabka-snap-does-not-exist-xyz");
427        let _ = fs::remove_file(&path);
428        assert!(Snapshot::load(&path).unwrap() == None);
429    }
430
431    #[test]
432    fn load_corrupt_file_is_err() {
433        let dir = env::temp_dir().join(format!("krabka-snap-corrupt-{}", process::id()));
434        fs::create_dir_all(&dir).unwrap();
435        let path = dir.join("snapshot");
436        fs::write(&path, [0xFF, 0xFF, 0x00, 0x01]).unwrap();
437        assert!(Snapshot::load(&path).is_err());
438        let _ = fs::remove_dir_all(&dir);
439    }
440}
krabka_remote_storage_topic/snapshot.rs

krabka_remote_storage_topic/
snapshot.rs
