krabka_remote_storage/
local.rs

1//! [`LocalTieredStorage`] is a filesystem-backed reference
2//! [`RemoteStorageManager`]. It mirrors Kafka's test fixture of the same
3//! name. It uses the same partition directories and segment filenames as
4//! Kafka 4.0, so a JVM `LocalTieredStorage` can read files copied by Krabka.
5//! It is useful for tests and single-node setups. Production deployments use
6//! an object-store-backed implementation behind the same trait.
7
8use {
9    crate::{
10        error::RemoteStorageError,
11        metadata::{CustomMetadata, RemoteLogSegmentMetadata},
12        storage_manager::{
13            ArtifactBody, IndexType, LogSegmentData, RemoteStorageManager, SegmentArtifact,
14            partition_dir_name, segment_file_name, segment_suffixes,
15        },
16    },
17    std::{
18        fs::{self, File},
19        io::{ErrorKind, Read, Seek, SeekFrom},
20        path::PathBuf,
21    },
22};
23
24/// A [`RemoteStorageManager`] that keeps offloaded segments on a local
25/// filesystem under `root`.
26///
27/// On-disk layout, per segment:
28///
29/// ```text
30/// <root>/<topic>-<partition>-<topic_id_base64>/
31///     <base_offset>-<segment_id_base64>.log
32///     <base_offset>-<segment_id_base64>.index
33///     <base_offset>-<segment_id_base64>.timeindex
34///     <base_offset>-<segment_id_base64>.snapshot
35///     <base_offset>-<segment_id_base64>.leader_epoch_checkpoint
36///     <base_offset>-<segment_id_base64>.txnindex  (when present)
37/// ```
38#[derive(Debug, Clone)]
39pub struct LocalTieredStorage {
40    root: PathBuf,
41}
42
43impl LocalTieredStorage {
44    /// Constructs a store rooted at `root`. The store creates the directory
45    /// on the first copy.
46    ///
47    /// Kafka's JVM implementation appends `kafka-tiered-storage` to its
48    /// configured parent directory. To share a tier with it, pass
49    /// `<remote.log.storage.local.dir>/kafka-tiered-storage` here.
50    #[must_use]
51    pub fn new(root: impl Into<PathBuf>) -> Self { Self { root: root.into() } }
52
53    /// The directory that holds every remote segment for one partition.
54    fn partition_dir(&self, metadata: &RemoteLogSegmentMetadata) -> PathBuf {
55        self.root.join(partition_dir_name(metadata))
56    }
57
58    fn segment_path(&self, metadata: &RemoteLogSegmentMetadata, suffix: &str) -> PathBuf {
59        self.partition_dir(metadata)
60            .join(segment_file_name(metadata, suffix))
61    }
62
63    fn log_path(&self, metadata: &RemoteLogSegmentMetadata) -> PathBuf {
64        self.segment_path(metadata, ".log")
65    }
66
67    fn index_path(&self, metadata: &RemoteLogSegmentMetadata, index_type: IndexType) -> PathBuf {
68        self.segment_path(metadata, index_type.suffix())
69    }
70}
71
72fn read_segment_range(
73    reader: &mut (impl Read + Seek),
74    segment_len: u64,
75    start_position: u32,
76    end_position: Option<u32>,
77) -> Result<Vec<u8>, RemoteStorageError> {
78    let start = u64::from(start_position);
79    if start > segment_len {
80        return Err(RemoteStorageError::InvalidArgument(format!(
81            "start_position {start} exceeds segment length {segment_len}"
82        )));
83    }
84    let end_exclusive = match end_position {
85        Some(end) => {
86            let end = u64::from(end);
87            if end < start {
88                return Err(RemoteStorageError::InvalidArgument(format!(
89                    "end_position {end} < start_position {start}"
90                )));
91            }
92            // `end` is inclusive; clamp to the segment length.
93            end.saturating_add(1).min(segment_len)
94        }
95        None => segment_len,
96    };
97    let requested_len = end_exclusive - start;
98    let capacity = usize::try_from(requested_len).map_err(|_| {
99        RemoteStorageError::InvalidArgument(format!(
100            "requested range length {requested_len} exceeds platform capacity"
101        ))
102    })?;
103
104    reader.seek(SeekFrom::Start(start))?;
105    let mut bytes = Vec::with_capacity(capacity);
106    reader.take(requested_len).read_to_end(&mut bytes)?;
107    Ok(bytes)
108}
109
110impl RemoteStorageManager for LocalTieredStorage {
111    #[krabka_macros::remote_operation]
112    fn copy_log_segment_data(
113        &self,
114        metadata: &RemoteLogSegmentMetadata,
115        data: &LogSegmentData,
116    ) -> Result<Option<CustomMetadata>, RemoteStorageError> {
117        let dir = self.partition_dir(metadata);
118        fs::create_dir_all(&dir)?;
119
120        for SegmentArtifact { suffix, body } in data.artifacts() {
121            let path = self.segment_path(metadata, suffix);
122            match body {
123                ArtifactBody::File(source) => {
124                    fs::copy(source, path)?;
125                }
126                ArtifactBody::Memory(bytes) => fs::write(path, bytes)?,
127            }
128        }
129        // A local store needs no opaque key echoed back.
130        Ok(None)
131    }
132
133    #[krabka_macros::remote_operation]
134    fn fetch_log_segment(
135        &self,
136        metadata: &RemoteLogSegmentMetadata,
137        start_position: u32,
138        end_position: Option<u32>,
139    ) -> Result<Vec<u8>, RemoteStorageError> {
140        let path = self.log_path(metadata);
141        if !path.exists() {
142            return Err(RemoteStorageError::SegmentNotFound(
143                metadata.remote_log_segment_id().clone(),
144            ));
145        }
146        let mut file = File::open(path)?;
147        let len = file.metadata()?.len();
148        read_segment_range(&mut file, len, start_position, end_position)
149    }
150
151    #[krabka_macros::remote_operation]
152    fn fetch_index(
153        &self,
154        metadata: &RemoteLogSegmentMetadata,
155        index_type: IndexType,
156    ) -> Result<Vec<u8>, RemoteStorageError> {
157        let path = self.index_path(metadata, index_type);
158        if !path.exists() {
159            return Err(RemoteStorageError::SegmentNotFound(
160                metadata.remote_log_segment_id().clone(),
161            ));
162        }
163        Ok(fs::read(&path)?)
164    }
165
166    #[krabka_macros::remote_operation]
167    fn delete_log_segment_data(
168        &self,
169        metadata: &RemoteLogSegmentMetadata,
170    ) -> Result<(), RemoteStorageError> {
171        for path in segment_suffixes().map(|suffix| self.segment_path(metadata, suffix)) {
172            match fs::remove_file(path) {
173                Ok(()) => {}
174                Err(error) if error.kind() == ErrorKind::NotFound => {}
175                Err(error) => return Err(RemoteStorageError::Io(error)),
176            }
177        }
178        Ok(())
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use {
185        super::*,
186        crate::test_support::{
187            TransactionIndex, check_sample_indexes, sample_data, sample_metadata as metadata,
188        },
189        assert2::{assert, check},
190        std::io,
191    };
192
193    /// A local tiered store, the directories around it, and one segment's metadata.
194    struct LocalFixture {
195        remote: tempfile::TempDir,
196        source: tempfile::TempDir,
197        store: LocalTieredStorage,
198        metadata: RemoteLogSegmentMetadata,
199    }
200
201    fn local_store() -> LocalFixture {
202        let remote = tempfile::tempdir().unwrap();
203        let source = tempfile::tempdir().unwrap();
204        let store = LocalTieredStorage::new(remote.path());
205        LocalFixture {
206            remote,
207            source,
208            store,
209            metadata: metadata(uuid::Uuid::from_u128(10)),
210        }
211    }
212
213    fn copied_store(with_transaction_index: TransactionIndex) -> LocalFixture {
214        let fixture = local_store();
215        fixture
216            .store
217            .copy_log_segment_data(
218                &fixture.metadata,
219                &sample_data(fixture.source.path(), with_transaction_index),
220            )
221            .unwrap();
222        fixture
223    }
224
225    #[test]
226    fn copy_then_fetch_full_segment() {
227        let LocalFixture {
228            remote: _remote,
229            source: src,
230            store: rsm,
231            metadata: md,
232        } = local_store();
233        assert!(
234            rsm.copy_log_segment_data(&md, &sample_data(src.path(), TransactionIndex::Present))
235                .unwrap()
236                .is_none()
237        );
238        let full = rsm.fetch_log_segment(&md, 0, None).unwrap();
239        assert!(full == b"0123456789");
240    }
241
242    #[yare::parameterized(
243        inclusive_range = { 2, Some(5), b"2345".as_ref() },
244        open_ended_range = { 7, None, b"789".as_ref() },
245        end_clamps_at_eof = { 8, Some(99), b"89".as_ref() },
246        start_at_eof = { 10, None, b"".as_ref() },
247    )]
248    fn fetch_partial_byte_ranges(start: u32, end: Option<u32>, want: &[u8]) {
249        let LocalFixture {
250            remote: _remote,
251            source: _src,
252            store: rsm,
253            metadata: md,
254        } = copied_store(TransactionIndex::Omitted);
255        check!(
256            rsm.fetch_log_segment(&md, start, end).unwrap() == want,
257            "range [{start}, {end:?}]"
258        );
259    }
260
261    #[test]
262    fn fetch_single_byte_range_start_equals_end() {
263        let LocalFixture {
264            remote: _remote,
265            source: _src,
266            store: rsm,
267            metadata: md,
268        } = copied_store(TransactionIndex::Omitted);
269        // Inclusive [3, 3] is a valid single-byte range -> "3". (The guard is
270        // `end < start`, not `<=`/`==`, so an equal start/end must succeed.)
271        assert!(rsm.fetch_log_segment(&md, 3, Some(3)).unwrap() == b"3");
272    }
273
274    #[yare::parameterized(
275        start_past_eof = { 11, None, "invalid argument: start_position 11 exceeds segment length 10" },
276        inverted = { 6, Some(5), "invalid argument: end_position 5 < start_position 6" },
277    )]
278    fn invalid_fetch_ranges_retain_error_semantics(start: u32, end: Option<u32>, expected: &str) {
279        let LocalFixture {
280            remote: _remote,
281            source: _src,
282            store: rsm,
283            metadata: md,
284        } = copied_store(TransactionIndex::Omitted);
285
286        let error = rsm.fetch_log_segment(&md, start, end).unwrap_err();
287        check!(matches!(error, RemoteStorageError::InvalidArgument(_)));
288        check!(error.to_string() == expected);
289    }
290
291    struct CountingSegment {
292        len: u64,
293        position: u64,
294        bytes_read: u64,
295    }
296
297    impl Read for CountingSegment {
298        fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
299            let remaining = self.len.saturating_sub(self.position);
300            let read_len =
301                usize::try_from(remaining.min(u64::try_from(buffer.len()).unwrap())).unwrap();
302            buffer[..read_len].fill(b'x');
303            self.position += u64::try_from(read_len).unwrap();
304            self.bytes_read += u64::try_from(read_len).unwrap();
305            Ok(read_len)
306        }
307    }
308
309    impl Seek for CountingSegment {
310        fn seek(&mut self, position: SeekFrom) -> io::Result<u64> {
311            let SeekFrom::Start(position) = position else {
312                return Err(io::Error::new(
313                    io::ErrorKind::Unsupported,
314                    "test reader supports absolute seeks only",
315                ));
316            };
317            self.position = position;
318            Ok(position)
319        }
320    }
321
322    #[yare::parameterized(
323        one_byte_of_one_mib = { 1_u64 << 20, 1_u32 },
324        sixty_four_kib_of_one_mib = { 1_u64 << 20, 1_u32 << 16 },
325        one_byte_of_sixty_four_mib = { 1_u64 << 26, 1_u32 },
326        sixty_four_kib_of_sixty_four_mib = { 1_u64 << 26, 1_u32 << 16 },
327    )]
328    fn bounded_fetch_reads_only_requested_bytes(segment_len: u64, requested_len: u32) {
329        let start = 17;
330        let end = start + requested_len - 1;
331        let mut segment = CountingSegment {
332            len: segment_len,
333            position: 0,
334            bytes_read: 0,
335        };
336
337        let bytes = read_segment_range(&mut segment, segment_len, start, Some(end)).unwrap();
338
339        check!(bytes == vec![b'x'; usize::try_from(requested_len).unwrap()]);
340        check!(segment.bytes_read == u64::from(requested_len));
341        check!(segment.position == u64::from(end) + 1);
342    }
343
344    #[test]
345    fn fetch_each_index_type() {
346        let LocalFixture {
347            remote: _remote,
348            source: _src,
349            store: rsm,
350            metadata: md,
351        } = copied_store(TransactionIndex::Present);
352        check_sample_indexes(&rsm, &md);
353    }
354
355    #[yare::parameterized(
356        log = { ".log" },
357        offset_index = { ".index" },
358        time_index = { ".timeindex" },
359        snapshot = { ".snapshot" },
360        leader_epoch_checkpoint = { ".leader_epoch_checkpoint" },
361        transaction_index = { ".txnindex" },
362    )]
363    fn copied_files_use_kafka_local_tiered_storage_layout(suffix: &str) {
364        let LocalFixture { remote, .. } = copied_store(TransactionIndex::Present);
365
366        let partition = remote.path().join("orders-0-AAAAAAAAAAAAAAAAAAAAAQ");
367        check!(
368            partition
369                .join(format!(
370                    "00000000000000000000-AAAAAAAAAAAAAAAAAAAACg{suffix}"
371                ))
372                .is_file(),
373            "missing Kafka layout artifact {suffix}"
374        );
375    }
376
377    #[test]
378    fn missing_optional_txn_index_is_not_found() {
379        let LocalFixture {
380            remote: _remote,
381            source: _src,
382            store: rsm,
383            metadata: md,
384        } = copied_store(TransactionIndex::Omitted);
385        let err = rsm.fetch_index(&md, IndexType::Transaction).unwrap_err();
386        assert!(matches!(err, RemoteStorageError::SegmentNotFound(_)));
387    }
388
389    #[test]
390    fn fetch_before_copy_is_not_found() {
391        let remote = tempfile::tempdir().unwrap();
392        let rsm = LocalTieredStorage::new(remote.path());
393        let md = metadata(uuid::Uuid::from_u128(404));
394        let err = rsm.fetch_log_segment(&md, 0, None).unwrap_err();
395        assert!(matches!(err, RemoteStorageError::SegmentNotFound(_)));
396    }
397
398    #[test]
399    fn delete_is_idempotent_and_removes_data() {
400        let LocalFixture {
401            remote: _remote,
402            source: _src,
403            store: rsm,
404            metadata: md,
405        } = copied_store(TransactionIndex::Present);
406        rsm.delete_log_segment_data(&md).unwrap();
407        // Second delete is a no-op.
408        rsm.delete_log_segment_data(&md).unwrap();
409        assert!(matches!(
410            rsm.fetch_log_segment(&md, 0, None).unwrap_err(),
411            RemoteStorageError::SegmentNotFound(_)
412        ));
413    }
414
415    #[test]
416    fn segments_are_isolated_by_id() {
417        let remote = tempfile::tempdir().unwrap();
418        let src = tempfile::tempdir().unwrap();
419        let rsm = LocalTieredStorage::new(remote.path());
420        let a = metadata(uuid::Uuid::from_u128(10));
421        let b = metadata(uuid::Uuid::from_u128(11));
422        rsm.copy_log_segment_data(&a, &sample_data(src.path(), TransactionIndex::Omitted))
423            .unwrap();
424        rsm.copy_log_segment_data(&b, &sample_data(src.path(), TransactionIndex::Omitted))
425            .unwrap();
426        rsm.delete_log_segment_data(&a).unwrap();
427        // Deleting `a` leaves `b` intact.
428        assert!(rsm.fetch_log_segment(&b, 0, None).unwrap() == b"0123456789");
429    }
430}
krabka_remote_storage/local.rs

krabka_remote_storage/
local.rs
