hexz_core/api/file.rs
1//! High-level snapshot file API and logical stream types.
2
3use crate::algo::compression::{Compressor, create_compressor};
4use crate::algo::encryption::Encryptor;
5use crate::cache::lru::{BlockCache, ShardedPageCache};
6use crate::cache::prefetch::Prefetcher;
7use crate::format::header::Header;
8use crate::format::index::{BlockInfo, IndexPage, MasterIndex, PageEntry};
9use crate::format::magic::{HEADER_SIZE, MAGIC_BYTES};
10use crate::format::version::{VersionCompatibility, check_version, compatibility_message};
11use crate::store::StorageBackend;
12use crate::store::local::file::FileBackend;
13use bytes::Bytes;
14use crc32fast::hash as crc32_hash;
15use std::mem::MaybeUninit;
16use std::path::Path;
17use std::ptr;
18use std::sync::{Arc, Mutex};
19
20use hexz_common::constants::{BLOCK_OFFSET_PARENT, DEFAULT_BLOCK_SIZE};
21use hexz_common::{Error, Result};
22use rayon::prelude::*;
23
24/// Shared zero block for the default block size to avoid allocating when returning zero blocks.
25static ZEROS_64K: [u8; DEFAULT_BLOCK_SIZE as usize] = [0u8; DEFAULT_BLOCK_SIZE as usize];
26
27/// Work item for block decompression: (block_idx, info, buf_offset, offset_in_block, to_copy)
28type WorkItem = (u64, BlockInfo, usize, usize, usize);
29
30/// Result of fetching a block from cache or storage.
31///
32/// Eliminates TOCTOU races by tracking data state at fetch time rather than
33/// re-checking the cache later (which can give a different answer if a
34/// background prefetch thread modifies the cache between check and use).
35enum FetchResult {
36 /// Data is already decompressed (came from L1 cache or is a zero block).
37 Decompressed(Bytes),
38 /// Data is raw compressed bytes from storage (needs decompression).
39 Compressed(Bytes),
40}
41
42/// Logical stream identifier for dual-stream snapshots.
43///
44/// Hexz snapshots can store two independent data streams:
45/// - **Disk**: Persistent storage (disk image, filesystem data)
46/// - **Memory**: Volatile state (RAM contents, process memory)
47///
48/// # Example
49///
50/// ```no_run
51/// use hexz_core::{File, SnapshotStream};
52/// # use std::sync::Arc;
53/// # fn example(snapshot: Arc<File>) -> Result<(), Box<dyn std::error::Error>> {
54/// // Read 4KB from disk stream
55/// let disk_data = snapshot.read_at(SnapshotStream::Disk, 0, 4096)?;
56///
57/// // Read 4KB from memory stream (if present)
58/// let mem_data = snapshot.read_at(SnapshotStream::Memory, 0, 4096)?;
59/// # Ok(())
60/// # }
61/// ```
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63#[repr(u8)]
64pub enum SnapshotStream {
65 /// Persistent disk/storage stream
66 Disk = 0,
67 /// Volatile memory stream
68 Memory = 1,
69}
70
71/// Read-only interface for accessing Hexz snapshot data.
72///
73/// `File` is the primary API for reading compressed, block-indexed snapshots.
74/// It handles:
75/// - Block-level decompression with LRU caching
76/// - Optional AES-256-GCM decryption
77/// - Thin snapshot parent chaining
78/// - Dual-stream access (disk and memory)
79/// - Random access with minimal I/O
80///
81/// # Thread Safety
82///
83/// `File` is `Send + Sync` and can be safely shared across threads via `Arc`.
84/// Internal caches use `Mutex` for synchronization.
85///
86/// # Performance
87///
88/// - **Cache hit latency**: ~80μs (warm cache)
89/// - **Cache miss latency**: ~1ms (cold cache, local storage)
90/// - **Sequential throughput**: ~2-3 GB/s (NVMe + LZ4)
91/// - **Memory overhead**: ~150MB typical (configurable)
92///
93/// # Examples
94///
95/// ## Basic Usage
96///
97/// ```no_run
98/// use hexz_core::{File, SnapshotStream};
99/// use hexz_core::store::local::FileBackend;
100/// use hexz_core::algo::compression::lz4::Lz4Compressor;
101/// use std::sync::Arc;
102///
103/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
104/// let backend = Arc::new(FileBackend::new("snapshot.hxz".as_ref())?);
105/// let compressor = Box::new(Lz4Compressor::new());
106/// let snapshot = File::new(backend, compressor, None)?;
107///
108/// // Read 4KB at offset 1MB
109/// let data = snapshot.read_at(SnapshotStream::Disk, 1024 * 1024, 4096)?;
110/// assert_eq!(data.len(), 4096);
111/// # Ok(())
112/// # }
113/// ```
114///
115/// ## Thin Snapshots (with parent)
116///
117/// ```no_run
118/// use hexz_core::File;
119/// use hexz_core::store::local::FileBackend;
120/// use hexz_core::algo::compression::lz4::Lz4Compressor;
121/// use std::sync::Arc;
122///
123/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
124/// // Open base snapshot
125/// let base_backend = Arc::new(FileBackend::new("base.hxz".as_ref())?);
126/// let base = File::new(
127/// base_backend,
128/// Box::new(Lz4Compressor::new()),
129/// None
130/// )?;
131///
132/// // The thin snapshot will automatically load its parent based on
133/// // the parent_path field in the header
134/// let thin_backend = Arc::new(FileBackend::new("incremental.hxz".as_ref())?);
135/// let thin = File::new(
136/// thin_backend,
137/// Box::new(Lz4Compressor::new()),
138/// None
139/// )?;
140///
141/// // Reads automatically fall back to base for unchanged blocks
142/// let data = thin.read_at(hexz_core::SnapshotStream::Disk, 0, 4096)?;
143/// # Ok(())
144/// # }
145/// ```
146pub struct File {
147 /// Snapshot metadata (sizes, compression, encryption settings)
148 pub header: Header,
149
150 /// Master index containing top-level page entries
151 master: MasterIndex,
152
153 /// Storage backend for reading raw snapshot data
154 backend: Arc<dyn StorageBackend>,
155
156 /// Compression algorithm (LZ4 or Zstandard)
157 compressor: Box<dyn Compressor>,
158
159 /// Optional encryption (AES-256-GCM)
160 encryptor: Option<Box<dyn Encryptor>>,
161
162 /// Optional parent snapshot for thin (incremental) snapshots.
163 /// When a block's offset is BLOCK_OFFSET_PARENT, data is fetched from parent.
164 parent: Option<Arc<File>>,
165
166 /// LRU cache for decompressed blocks (per-stream, per-block-index)
167 cache_l1: BlockCache,
168
169 /// Sharded LRU cache for deserialized index pages
170 page_cache: ShardedPageCache,
171
172 /// Optional prefetcher for background data loading
173 prefetcher: Option<Prefetcher>,
174}
175
176impl File {
177 /// Opens a Hexz snapshot with default cache settings.
178 ///
179 /// This is the primary constructor for `File`. It:
180 /// 1. Reads and validates the snapshot header (magic bytes, version)
181 /// 2. Deserializes the master index
182 /// 3. Recursively loads parent snapshots (for thin snapshots)
183 /// 4. Initializes block and page caches
184 ///
185 /// # Parameters
186 ///
187 /// - `backend`: Storage backend (local file, HTTP, S3, etc.)
188 /// - `compressor`: Compression algorithm matching the snapshot format
189 /// - `encryptor`: Optional decryption handler (pass `None` for unencrypted snapshots)
190 ///
191 /// # Returns
192 ///
193 /// - `Ok(File)` on success
194 /// - `Err(Error::Format)` if magic bytes or version are invalid
195 /// - `Err(Error::Io)` if storage backend fails
196 ///
197 /// # Examples
198 ///
199 /// ```no_run
200 /// use hexz_core::{File, SnapshotStream};
201 /// use hexz_core::store::local::FileBackend;
202 /// use hexz_core::algo::compression::lz4::Lz4Compressor;
203 /// use std::sync::Arc;
204 ///
205 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
206 /// let backend = Arc::new(FileBackend::new("snapshot.hxz".as_ref())?);
207 /// let compressor = Box::new(Lz4Compressor::new());
208 /// let snapshot = File::new(backend, compressor, None)?;
209 ///
210 /// println!("Disk size: {} bytes", snapshot.size(SnapshotStream::Disk));
211 /// # Ok(())
212 /// # }
213 /// ```
214 /// Opens a snapshot, auto-detecting compression and dictionary from the header.
215 ///
216 /// This eliminates the 3-step boilerplate of: read header, load dict, create
217 /// compressor. Equivalent to `File::new(backend, auto_compressor, encryptor)`.
218 pub fn open(
219 backend: Arc<dyn StorageBackend>,
220 encryptor: Option<Box<dyn Encryptor>>,
221 ) -> Result<Arc<Self>> {
222 Self::open_with_cache(backend, encryptor, None, None)
223 }
224
225 /// Like [`open`](Self::open) but with custom cache and prefetch settings.
226 pub fn open_with_cache(
227 backend: Arc<dyn StorageBackend>,
228 encryptor: Option<Box<dyn Encryptor>>,
229 cache_capacity_bytes: Option<usize>,
230 prefetch_window_size: Option<u32>,
231 ) -> Result<Arc<Self>> {
232 let header = Header::read_from_backend(backend.as_ref())?;
233 let dictionary = header.load_dictionary(backend.as_ref())?;
234 let compressor = create_compressor(header.compression, None, dictionary);
235 Self::with_cache(
236 backend,
237 compressor,
238 encryptor,
239 cache_capacity_bytes,
240 prefetch_window_size,
241 )
242 }
243
244 pub fn new(
245 backend: Arc<dyn StorageBackend>,
246 compressor: Box<dyn Compressor>,
247 encryptor: Option<Box<dyn Encryptor>>,
248 ) -> Result<Arc<Self>> {
249 Self::with_cache(backend, compressor, encryptor, None, None)
250 }
251
252 /// Opens a Hexz snapshot with custom cache capacity and prefetching.
253 ///
254 /// Identical to [`new`](Self::new) but allows specifying cache size and prefetch window.
255 ///
256 /// # Parameters
257 ///
258 /// - `backend`: Storage backend
259 /// - `compressor`: Compression algorithm
260 /// - `encryptor`: Optional decryption handler
261 /// - `cache_capacity_bytes`: Block cache size in bytes (default: ~400MB for 4KB blocks)
262 /// - `prefetch_window_size`: Number of blocks to prefetch ahead (default: disabled)
263 ///
264 /// # Cache Sizing
265 ///
266 /// The cache stores decompressed blocks. Given a block size of 4KB:
267 /// - `Some(100_000_000)` → ~24,000 blocks (~96MB effective)
268 /// - `None` → 1000 blocks (~4MB effective)
269 ///
270 /// Larger caches reduce repeated decompression but increase memory usage.
271 ///
272 /// # Prefetching
273 ///
274 /// When `prefetch_window_size` is set, the system will automatically fetch the next N blocks
275 /// in the background after each read, optimizing sequential access patterns:
276 /// - `Some(4)` → Prefetch 4 blocks ahead
277 /// - `None` or `Some(0)` → Disable prefetching
278 ///
279 /// # Examples
280 ///
281 /// ```no_run
282 /// use hexz_core::File;
283 /// use hexz_core::store::local::FileBackend;
284 /// use hexz_core::algo::compression::lz4::Lz4Compressor;
285 /// use std::sync::Arc;
286 ///
287 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
288 /// let backend = Arc::new(FileBackend::new("snapshot.hxz".as_ref())?);
289 /// let compressor = Box::new(Lz4Compressor::new());
290 ///
291 /// // Allocate 256MB for cache, prefetch 4 blocks ahead
292 /// let snapshot = File::with_cache(
293 /// backend,
294 /// compressor,
295 /// None,
296 /// Some(256 * 1024 * 1024),
297 /// Some(4)
298 /// )?;
299 /// # Ok(())
300 /// # }
301 /// ```
302 pub fn with_cache(
303 backend: Arc<dyn StorageBackend>,
304 compressor: Box<dyn Compressor>,
305 encryptor: Option<Box<dyn Encryptor>>,
306 cache_capacity_bytes: Option<usize>,
307 prefetch_window_size: Option<u32>,
308 ) -> Result<Arc<Self>> {
309 let header_bytes = backend.read_exact(0, HEADER_SIZE)?;
310 let header: Header = bincode::deserialize(&header_bytes)?;
311
312 if &header.magic != MAGIC_BYTES {
313 return Err(Error::Format("Invalid magic bytes".into()));
314 }
315
316 // Check version compatibility
317 let compatibility = check_version(header.version);
318 match compatibility {
319 VersionCompatibility::Full => {
320 // Perfect match, proceed silently
321 }
322 VersionCompatibility::Degraded => {
323 // Newer version, issue warning but allow
324 tracing::warn!("{}", compatibility_message(header.version));
325 }
326 VersionCompatibility::Incompatible => {
327 // Too old or too new, reject
328 return Err(Error::Format(compatibility_message(header.version)));
329 }
330 }
331
332 let index_bytes = backend.read_exact(
333 header.index_offset,
334 (backend.len() - header.index_offset) as usize,
335 )?;
336
337 let master: MasterIndex = bincode::deserialize(&index_bytes)?;
338
339 // Recursively load parent if present
340 let parent = if let Some(parent_path) = &header.parent_path {
341 tracing::info!("Loading parent snapshot: {}", parent_path);
342 let p_backend = Arc::new(FileBackend::new(Path::new(parent_path))?);
343 Some(File::open(p_backend, None)?)
344 } else {
345 None
346 };
347
348 let block_size = header.block_size as usize;
349 let l1_capacity = if let Some(bytes) = cache_capacity_bytes {
350 (bytes / block_size).max(1)
351 } else {
352 1000
353 };
354
355 // Initialize prefetcher if window size is specified and > 0
356 let prefetcher = prefetch_window_size.filter(|&w| w > 0).map(Prefetcher::new);
357
358 Ok(Arc::new(Self {
359 header,
360 master,
361 backend,
362 compressor,
363 encryptor,
364 parent,
365 cache_l1: BlockCache::with_capacity(l1_capacity),
366 page_cache: ShardedPageCache::default(),
367 prefetcher,
368 }))
369 }
370
371 /// Returns the logical size of a stream in bytes.
372 ///
373 /// # Parameters
374 ///
375 /// - `stream`: The stream to query (Disk or Memory)
376 ///
377 /// # Returns
378 ///
379 /// The uncompressed, logical size of the stream. This is the size you would
380 /// get if you decompressed all blocks and concatenated them.
381 ///
382 /// # Examples
383 ///
384 /// ```no_run
385 /// use hexz_core::{File, SnapshotStream};
386 /// # use std::sync::Arc;
387 /// # fn example(snapshot: Arc<File>) {
388 /// let disk_bytes = snapshot.size(SnapshotStream::Disk);
389 /// let mem_bytes = snapshot.size(SnapshotStream::Memory);
390 ///
391 /// println!("Disk: {} GB", disk_bytes / (1024 * 1024 * 1024));
392 /// println!("Memory: {} MB", mem_bytes / (1024 * 1024));
393 /// # }
394 /// ```
395 /// Returns the total number of prefetch operations spawned since this file was opened.
396 /// Returns 0 if prefetching is disabled.
397 pub fn prefetch_spawn_count(&self) -> u64 {
398 self.prefetcher.as_ref().map_or(0, |p| p.spawn_count())
399 }
400
401 pub fn size(&self, stream: SnapshotStream) -> u64 {
402 match stream {
403 SnapshotStream::Disk => self.master.disk_size,
404 SnapshotStream::Memory => self.master.memory_size,
405 }
406 }
407
408 /// Reads data from a snapshot stream at a given offset.
409 ///
410 /// This is the primary read method for random access. It:
411 /// 1. Identifies which blocks overlap the requested range
412 /// 2. Fetches blocks from cache or decompresses from storage
413 /// 3. Handles thin snapshot fallback to parent
414 /// 4. Assembles the final buffer from block slices
415 ///
416 /// # Parameters
417 ///
418 /// - `stream`: Which stream to read from (Disk or Memory)
419 /// - `offset`: Starting byte offset (0-indexed)
420 /// - `len`: Number of bytes to read
421 ///
422 /// # Returns
423 ///
424 /// A `Vec<u8>` containing up to `len` bytes. The returned vector may be shorter
425 /// if:
426 /// - `offset` is beyond the stream size (returns empty vector)
427 /// - `offset + len` exceeds stream size (returns partial data)
428 ///
429 /// Missing data (sparse regions) is zero-filled.
430 ///
431 /// # Errors
432 ///
433 /// - `Error::Io` if backend read fails (e.g. truncated file)
434 /// - `Error::Corruption(block_idx)` if block checksum does not match
435 /// - `Error::Decompression` if block decompression fails
436 /// - `Error::Decryption` if block decryption fails
437 ///
438 /// # Performance
439 ///
440 /// - **Cache hit**: ~80μs latency, no I/O
441 /// - **Cache miss**: ~1ms latency (local storage), includes decompression
442 /// - **Remote storage**: Latency depends on network (HTTP: ~50ms, S3: ~100ms)
443 ///
444 /// Aligned reads (offset % block_size == 0) are most efficient.
445 ///
446 /// # Examples
447 ///
448 /// ```no_run
449 /// use hexz_core::{File, SnapshotStream};
450 /// # use std::sync::Arc;
451 /// # fn example(snapshot: Arc<File>) -> Result<(), Box<dyn std::error::Error>> {
452 /// // Read first 512 bytes of disk stream
453 /// let boot_sector = snapshot.read_at(SnapshotStream::Disk, 0, 512)?;
454 ///
455 /// // Read from arbitrary offset
456 /// let chunk = snapshot.read_at(SnapshotStream::Disk, 1024 * 1024, 4096)?;
457 ///
458 /// // Reading beyond stream size returns empty vector
459 /// let empty = snapshot.read_at(SnapshotStream::Disk, u64::MAX, 100)?;
460 /// assert!(empty.is_empty());
461 /// # Ok(())
462 /// # }
463 /// ```
464 /// Reads a byte range. Uses parallel block decompression when the range spans multiple blocks.
465 pub fn read_at(
466 self: &Arc<Self>,
467 stream: SnapshotStream,
468 offset: u64,
469 len: usize,
470 ) -> Result<Vec<u8>> {
471 let stream_size = self.size(stream);
472 if offset >= stream_size {
473 return Ok(Vec::new());
474 }
475 let actual_len = std::cmp::min(len as u64, stream_size - offset) as usize;
476 if actual_len == 0 {
477 return Ok(Vec::new());
478 }
479
480 let pages = match stream {
481 SnapshotStream::Disk => &self.master.disk_pages,
482 SnapshotStream::Memory => &self.master.memory_pages,
483 };
484
485 if pages.is_empty() {
486 if let Some(parent) = &self.parent {
487 return parent.read_at(stream, offset, actual_len);
488 }
489 return Ok(vec![0u8; actual_len]);
490 }
491
492 let mut buf: Vec<MaybeUninit<u8>> = Vec::new();
493 buf.resize_with(actual_len, MaybeUninit::uninit);
494 self.read_at_into_uninit(stream, offset, &mut buf)?;
495 let ptr = buf.as_mut_ptr().cast::<u8>();
496 let len = buf.len();
497 let cap = buf.capacity();
498 std::mem::forget(buf);
499 // SAFETY: `buf` was a Vec<MaybeUninit<u8>> that we fully initialized via
500 // `read_at_into_uninit` (which writes every byte). We `forget` the original
501 // Vec to avoid a double-free and reconstruct it with the same ptr/len/cap.
502 // MaybeUninit<u8> has the same layout as u8.
503 Ok(unsafe { Vec::from_raw_parts(ptr, len, cap) })
504 }
505
506 /// Reads into a provided buffer. Unused suffix is zero-filled. Uses parallel decompression when spanning multiple blocks.
507 pub fn read_at_into(
508 self: &Arc<Self>,
509 stream: SnapshotStream,
510 offset: u64,
511 buffer: &mut [u8],
512 ) -> Result<()> {
513 let len = buffer.len();
514 if len == 0 {
515 return Ok(());
516 }
517 let stream_size = self.size(stream);
518 if offset >= stream_size {
519 buffer.fill(0);
520 return Ok(());
521 }
522 let actual_len = std::cmp::min(len as u64, stream_size - offset) as usize;
523 if actual_len < len {
524 buffer[actual_len..].fill(0);
525 }
526 self.read_at_into_uninit_bytes(stream, offset, &mut buffer[0..actual_len])
527 }
528
529 /// Minimum number of local blocks to use the parallel decompression path.
530 const PARALLEL_MIN_BLOCKS: usize = 2;
531
532 /// Collects work items for blocks that need decompression.
533 ///
534 /// This method iterates through index pages and blocks, handling:
535 /// - Parent blocks: delegate to parent snapshot or zero-fill
536 /// - Zero blocks: zero-fill directly
537 /// - Regular blocks: add to work queue for later decompression
538 ///
539 /// Returns the work items to process and updates the tracking variables.
540 fn collect_work_items(
541 &self,
542 stream: SnapshotStream,
543 pages: &[PageEntry],
544 page_idx: usize,
545 target: &mut [MaybeUninit<u8>],
546 offset: u64,
547 actual_len: usize,
548 ) -> Result<(Vec<WorkItem>, usize)> {
549 let mut local_work: Vec<WorkItem> = Vec::new();
550 let mut buf_offset = 0;
551 let mut current_pos = offset;
552 let mut remaining = actual_len;
553
554 for page_entry in pages.iter().skip(page_idx) {
555 if remaining == 0 {
556 break;
557 }
558 if page_entry.start_logical > current_pos + remaining as u64 {
559 break;
560 }
561
562 let page = self.get_page(page_entry)?;
563 let mut block_logical_start = page_entry.start_logical;
564
565 for (block_idx_in_page, block) in page.blocks.iter().enumerate() {
566 let block_end = block_logical_start + block.logical_len as u64;
567
568 if block_end > current_pos {
569 let global_block_idx = page_entry.start_block + block_idx_in_page as u64;
570 let offset_in_block = (current_pos - block_logical_start) as usize;
571 let to_copy = std::cmp::min(
572 remaining,
573 (block.logical_len as usize).saturating_sub(offset_in_block),
574 );
575
576 if block.offset == BLOCK_OFFSET_PARENT {
577 // Parent block: delegate or zero-fill
578 if let Some(parent) = &self.parent {
579 let dest = &mut target[buf_offset..buf_offset + to_copy];
580 parent.read_at_into_uninit(stream, current_pos, dest)?;
581 } else {
582 Self::zero_fill_uninit(&mut target[buf_offset..buf_offset + to_copy]);
583 }
584 current_pos += to_copy as u64;
585 buf_offset += to_copy;
586 remaining -= to_copy;
587 } else if block.length == 0 {
588 // Zero block: fill with zeros
589 Self::zero_fill_uninit(&mut target[buf_offset..buf_offset + to_copy]);
590 current_pos += to_copy as u64;
591 buf_offset += to_copy;
592 remaining -= to_copy;
593 } else {
594 // Regular block: add to work queue
595 if to_copy > 0 {
596 local_work.push((
597 global_block_idx,
598 *block,
599 buf_offset,
600 offset_in_block,
601 to_copy,
602 ));
603 buf_offset += to_copy;
604 current_pos += to_copy as u64;
605 remaining -= to_copy;
606 }
607 }
608
609 if remaining == 0 {
610 break;
611 }
612 }
613 block_logical_start += block.logical_len as u64;
614 }
615 }
616
617 Ok((local_work, buf_offset))
618 }
619
620 /// Executes parallel decompression for multiple blocks.
621 ///
622 /// Uses a two-phase approach:
623 /// 1. Parallel I/O: Fetch all raw blocks concurrently
624 /// 2. Parallel CPU: Decompress and copy to target buffer
625 fn execute_parallel_decompression(
626 self: &Arc<Self>,
627 stream: SnapshotStream,
628 work_items: &[WorkItem],
629 target: &mut [MaybeUninit<u8>],
630 actual_len: usize,
631 ) -> Result<()> {
632 let snap = Arc::clone(self);
633 let target_addr = target.as_mut_ptr() as usize;
634
635 // Phase 1: Parallel fetch all raw blocks. Each result tracks whether
636 // the data is already decompressed (cache hit / zero block) or still
637 // compressed (storage read), eliminating a TOCTOU race where a background
638 // prefetch thread could modify the cache between fetch and decompression.
639 let raw_blocks: Vec<Result<FetchResult>> = work_items
640 .par_iter()
641 .map(|(block_idx, info, _, _, _)| snap.fetch_raw_block(stream, *block_idx, info))
642 .collect();
643
644 // Phase 2: Parallel decompress and copy
645 let err: Mutex<Option<Error>> = Mutex::new(None);
646 work_items
647 .par_iter()
648 .zip(raw_blocks)
649 .for_each(|(work_item, fetch_result)| {
650 if err.lock().map_or(true, |e| e.is_some()) {
651 return;
652 }
653
654 let (block_idx, info, buf_offset, offset_in_block, to_copy) = work_item;
655
656 // Handle fetch errors
657 let fetched = match fetch_result {
658 Ok(r) => r,
659 Err(e) => {
660 if let Ok(mut guard) = err.lock() {
661 let _ = guard.replace(e);
662 }
663 return;
664 }
665 };
666
667 // Use the FetchResult to determine if decompression is needed,
668 // rather than re-checking the cache (which could give a stale answer).
669 let data = match fetched {
670 FetchResult::Decompressed(data) => data,
671 FetchResult::Compressed(raw) => {
672 match snap.decompress_and_verify(raw, *block_idx, info) {
673 Ok(d) => {
674 // Cache the result
675 snap.cache_l1.insert(stream, *block_idx, d.clone());
676 d
677 }
678 Err(e) => {
679 if let Ok(mut guard) = err.lock() {
680 let _ = guard.replace(e);
681 }
682 return;
683 }
684 }
685 }
686 };
687
688 // Copy to target buffer
689 let src = data.as_ref();
690 let start = *offset_in_block;
691 let len = *to_copy;
692 if start < src.len() && len <= src.len() - start {
693 // Defensive assertion: ensure destination write is within bounds
694 debug_assert!(
695 buf_offset + len <= actual_len,
696 "Buffer overflow: attempting to write {} bytes at offset {} into buffer of length {}",
697 len,
698 buf_offset,
699 actual_len
700 );
701 let dest = (target_addr + buf_offset) as *mut u8;
702 // SAFETY: `src[start..start+len]` is in-bounds (checked above).
703 // `dest` points into the `target` MaybeUninit buffer at a unique
704 // non-overlapping offset (each work item has a distinct `buf_offset`),
705 // and the rayon par_iter ensures each item writes to a disjoint region.
706 // The debug_assert above validates buf_offset + len <= actual_len.
707 unsafe { ptr::copy_nonoverlapping(src[start..].as_ptr(), dest, len) };
708 }
709 });
710
711 if let Some(e) = err.lock().ok().and_then(|mut guard| guard.take()) {
712 return Err(e);
713 }
714
715 Ok(())
716 }
717
718 /// Executes serial decompression for a small number of blocks.
719 fn execute_serial_decompression(
720 &self,
721 stream: SnapshotStream,
722 work_items: &[WorkItem],
723 target: &mut [MaybeUninit<u8>],
724 actual_len: usize,
725 ) -> Result<()> {
726 for (block_idx, info, buf_offset, offset_in_block, to_copy) in work_items {
727 let data = self.resolve_block_data(stream, *block_idx, info)?;
728 let src = data.as_ref();
729 let start = *offset_in_block;
730 if start < src.len() && *to_copy <= src.len() - start {
731 // Defensive assertion: ensure destination write is within bounds
732 debug_assert!(
733 *buf_offset + *to_copy <= actual_len,
734 "Buffer overflow: attempting to write {} bytes at offset {} into buffer of length {}",
735 to_copy,
736 buf_offset,
737 actual_len
738 );
739 // SAFETY: `src[start..start+to_copy]` is in-bounds (checked above).
740 // `target[buf_offset..]` has sufficient room because `buf_offset + to_copy`
741 // never exceeds `actual_len` (tracked during work-item collection).
742 // The debug_assert above validates this invariant.
743 // MaybeUninit<u8> has the same layout as u8.
744 unsafe {
745 ptr::copy_nonoverlapping(
746 src[start..].as_ptr(),
747 target[*buf_offset..].as_mut_ptr() as *mut u8,
748 *to_copy,
749 );
750 }
751 }
752 }
753 Ok(())
754 }
755
756 /// Zero-fills a slice of uninitialized memory.
757 ///
758 /// This helper centralizes all unsafe zero-filling operations to improve
759 /// safety auditing and reduce code duplication.
760 ///
761 /// # Safety
762 ///
763 /// This function writes zeros to the provided buffer, making it fully initialized.
764 /// The caller must ensure the buffer is valid for writes.
765 #[inline]
766 fn zero_fill_uninit(buffer: &mut [MaybeUninit<u8>]) {
767 if !buffer.is_empty() {
768 // SAFETY: buffer is a valid &mut [MaybeUninit<u8>] slice, so writing
769 // buffer.len() zero bytes through its pointer is in-bounds.
770 unsafe { ptr::write_bytes(buffer.as_mut_ptr(), 0, buffer.len()) };
771 }
772 }
773
774 /// Writes into uninitialized memory. Unused suffix is zero-filled. Uses parallel decompression when spanning multiple blocks.
775 ///
776 /// **On error:** The buffer contents are undefined (possibly partially written).
777 pub fn read_at_into_uninit(
778 self: &Arc<Self>,
779 stream: SnapshotStream,
780 offset: u64,
781 buffer: &mut [MaybeUninit<u8>],
782 ) -> Result<()> {
783 self.read_at_uninit_inner(stream, offset, buffer, false)
784 }
785
786 /// Inner implementation of [`read_at_into_uninit`](Self::read_at_into_uninit).
787 ///
788 /// The `is_prefetch` flag prevents recursive prefetch thread spawning:
789 /// when `true`, the prefetch block is skipped to avoid unbounded thread creation.
790 fn read_at_uninit_inner(
791 self: &Arc<Self>,
792 stream: SnapshotStream,
793 offset: u64,
794 buffer: &mut [MaybeUninit<u8>],
795 is_prefetch: bool,
796 ) -> Result<()> {
797 // Early validation
798 let len = buffer.len();
799 if len == 0 {
800 return Ok(());
801 }
802
803 let stream_size = self.size(stream);
804 if offset >= stream_size {
805 Self::zero_fill_uninit(buffer);
806 return Ok(());
807 }
808
809 // Calculate actual read length and zero-fill suffix if needed
810 let actual_len = std::cmp::min(len as u64, stream_size - offset) as usize;
811 if actual_len < len {
812 Self::zero_fill_uninit(&mut buffer[actual_len..]);
813 }
814
815 let target = &mut buffer[0..actual_len];
816
817 // Get page list for stream
818 let pages = match stream {
819 SnapshotStream::Disk => &self.master.disk_pages,
820 SnapshotStream::Memory => &self.master.memory_pages,
821 };
822
823 // Delegate to parent if no index pages
824 if pages.is_empty() {
825 if let Some(parent) = &self.parent {
826 return parent.read_at_into_uninit(stream, offset, target);
827 }
828 Self::zero_fill_uninit(target);
829 return Ok(());
830 }
831
832 // Find starting page index
833 let page_idx = match pages.binary_search_by(|p| p.start_logical.cmp(&offset)) {
834 Ok(idx) => idx,
835 Err(idx) => idx.saturating_sub(1),
836 };
837
838 // Collect work items (handles parent blocks, zero blocks, and queues regular blocks)
839 let (work_items, buf_offset) =
840 self.collect_work_items(stream, pages, page_idx, target, offset, actual_len)?;
841
842 // Choose parallel or serial decompression based on work item count
843 if work_items.len() >= Self::PARALLEL_MIN_BLOCKS {
844 self.execute_parallel_decompression(stream, &work_items, target, actual_len)?;
845 } else {
846 self.execute_serial_decompression(stream, &work_items, target, actual_len)?;
847 }
848
849 // Handle any remaining unprocessed data
850 let remaining = actual_len - buf_offset;
851 if remaining > 0 {
852 if let Some(parent) = &self.parent {
853 let current_pos = offset + buf_offset as u64;
854 parent.read_at_into_uninit(stream, current_pos, &mut target[buf_offset..])?;
855 } else {
856 Self::zero_fill_uninit(&mut target[buf_offset..]);
857 }
858 }
859
860 // Trigger prefetch for next sequential blocks if enabled.
861 // Guards:
862 // 1. `is_prefetch` prevents recursive spawning (prefetch thread spawning another)
863 // 2. `try_start()` limits to one in-flight prefetch at a time, preventing
864 // unbounded thread creation under rapid sequential reads
865 if let Some(prefetcher) = &self.prefetcher {
866 if !is_prefetch && !work_items.is_empty() && prefetcher.try_start() {
867 let next_offset = offset + actual_len as u64;
868 let prefetch_len = (self.header.block_size * 4) as usize;
869 let snap = Arc::clone(self);
870 let stream_copy = stream;
871 std::thread::spawn(move || {
872 let mut buf = vec![MaybeUninit::uninit(); prefetch_len];
873 let _ = snap.read_at_uninit_inner(stream_copy, next_offset, &mut buf, true);
874 // Release the in-flight guard so the next read can prefetch
875 if let Some(pf) = &snap.prefetcher {
876 pf.clear_in_flight();
877 }
878 });
879 }
880 }
881
882 Ok(())
883 }
884
885 /// Like [`read_at_into_uninit`](Self::read_at_into_uninit) but accepts `&mut [u8]`. Use from FFI (e.g. Python).
886 #[inline]
887 pub fn read_at_into_uninit_bytes(
888 self: &Arc<Self>,
889 stream: SnapshotStream,
890 offset: u64,
891 buf: &mut [u8],
892 ) -> Result<()> {
893 if buf.is_empty() {
894 return Ok(());
895 }
896 // SAFETY: &mut [u8] and &mut [MaybeUninit<u8>] have identical layout (both
897 // are slices of single-byte types). Initialized u8 values are valid MaybeUninit<u8>.
898 // The borrow is derived from `buf` so no aliasing occurs.
899 let uninit = unsafe { &mut *(buf as *mut [u8] as *mut [MaybeUninit<u8>]) };
900 self.read_at_into_uninit(stream, offset, uninit)
901 }
902
903 /// Fetches an index page from cache or storage.
904 ///
905 /// Index pages map logical offsets to physical block locations. This method
906 /// maintains an LRU cache to avoid repeated deserialization.
907 ///
908 /// # Parameters
909 ///
910 /// - `entry`: Page metadata from master index
911 ///
912 /// # Returns
913 ///
914 /// A shared reference to the deserialized index page.
915 ///
916 /// # Thread Safety
917 ///
918 /// This method acquires a lock on the page cache only for cache lookup and insertion.
919 /// I/O and deserialization are performed without holding the lock to avoid blocking
920 /// other threads during cache misses.
921 fn get_page(&self, entry: &PageEntry) -> Result<Arc<IndexPage>> {
922 // Fast path: check sharded cache
923 if let Some(p) = self.page_cache.get(entry.offset) {
924 return Ok(p);
925 }
926
927 // Slow path: I/O and deserialization without holding any lock
928 let bytes = self
929 .backend
930 .read_exact(entry.offset, entry.length as usize)?;
931 let page: IndexPage = bincode::deserialize(&bytes)?;
932 let arc = Arc::new(page);
933
934 // Check again in case another thread inserted while we were doing I/O
935 if let Some(p) = self.page_cache.get(entry.offset) {
936 return Ok(p);
937 }
938 self.page_cache.insert(entry.offset, arc.clone());
939
940 Ok(arc)
941 }
942
943 /// Fetches raw compressed block data from cache or storage.
944 ///
945 /// This is the I/O portion of block resolution, separated to enable parallel I/O.
946 /// It:
947 /// 1. Checks the block cache
948 /// 2. Handles zero-length blocks
949 /// 3. Reads raw compressed data from backend
950 ///
951 /// # Parameters
952 ///
953 /// - `stream`: Stream identifier (for cache key)
954 /// - `block_idx`: Global block index
955 /// - `info`: Block metadata (offset, length)
956 ///
957 /// # Returns
958 ///
959 /// Raw block data (potentially compressed/encrypted) or cached decompressed data.
960 fn fetch_raw_block(
961 &self,
962 stream: SnapshotStream,
963 block_idx: u64,
964 info: &BlockInfo,
965 ) -> Result<FetchResult> {
966 // Check cache first - return decompressed data if available
967 if let Some(data) = self.cache_l1.get(stream, block_idx) {
968 return Ok(FetchResult::Decompressed(data));
969 }
970
971 // Handle zero blocks
972 if info.length == 0 {
973 let len = info.logical_len as usize;
974 if len == 0 {
975 return Ok(FetchResult::Decompressed(Bytes::new()));
976 }
977 if len == ZEROS_64K.len() {
978 return Ok(FetchResult::Decompressed(Bytes::from_static(&ZEROS_64K)));
979 }
980 return Ok(FetchResult::Decompressed(Bytes::from(vec![0u8; len])));
981 }
982
983 // Fetch raw compressed data (THIS IS THE PARALLEL PART)
984 self.backend
985 .read_exact(info.offset, info.length as usize)
986 .map(FetchResult::Compressed)
987 }
988
989 /// Decompresses and verifies a raw block.
990 ///
991 /// This is the CPU portion of block resolution, separated to enable parallel decompression.
992 /// It:
993 /// 1. Verifies CRC32 checksum
994 /// 2. Decrypts (if encrypted)
995 /// 3. Decompresses
996 ///
997 /// # Parameters
998 ///
999 /// - `raw`: Raw block data (potentially compressed/encrypted)
1000 /// - `block_idx`: Global block index (for error reporting and decryption)
1001 /// - `info`: Block metadata (checksum)
1002 ///
1003 /// # Returns
1004 ///
1005 /// Decompressed block data as `Bytes`.
1006 ///
1007 /// # Performance
1008 ///
1009 /// Decompression throughput:
1010 /// - LZ4: ~2 GB/s per core
1011 /// - Zstd: ~500 MB/s per core
1012 fn decompress_and_verify(&self, raw: Bytes, block_idx: u64, info: &BlockInfo) -> Result<Bytes> {
1013 // Verify stored checksum (CRC32 of compressed/encrypted data) before decrypt/decompress
1014 if info.checksum != 0 {
1015 let computed = crc32_hash(&raw);
1016 if computed != info.checksum {
1017 return Err(Error::Corruption(block_idx));
1018 }
1019 }
1020
1021 // Pre-allocate exact output buffer to avoid over-allocation inside decompressor.
1022 // We use decompress_into() instead of decompress() to eliminate the allocation
1023 // and potential reallocation overhead inside the compression library.
1024 //
1025 // Performance impact: Avoids zero-initialization overhead (~16% improvement for
1026 // high-thread-count workloads based on benchmarks).
1027 let out_len = info.logical_len as usize;
1028 let mut out = Vec::with_capacity(out_len);
1029
1030 // SAFETY: This unsafe block is required to create an uninitialized buffer for
1031 // decompress_into() to write into. This is safe because:
1032 //
1033 // 1. Contract guarantee: Both LZ4 and Zstd decompress_into() implementations
1034 // promise to either:
1035 // a) Write exactly `out.len()` bytes (the full decompressed size), OR
1036 // b) Return an Err() if decompression fails (buffer underrun, corruption, etc.)
1037 //
1038 // 2. Size accuracy: We set out.len() to info.logical_len, which is the exact
1039 // decompressed size recorded in the block metadata during compression.
1040 // The decompressor will write exactly this many bytes or fail.
1041 //
1042 // 3. Error propagation: If decompress_into() returns Err(), we propagate it
1043 // immediately via the ? operator. The uninitialized buffer is dropped
1044 // without ever being read.
1045 //
1046 // 4. No partial writes: The decompressor APIs do not support partial writes.
1047 // They either fully succeed or fully fail. We never access a partially
1048 // initialized buffer.
1049 //
1050 // 5. Memory safety: We never read from `out` before decompress_into() succeeds.
1051 // The only subsequent access is Bytes::from(out), which transfers ownership
1052 // of the now-fully-initialized buffer.
1053 //
1054 // This is a well-established pattern for zero-copy decompression. The clippy
1055 // lint is conservative and warns about ANY use of set_len() after with_capacity(),
1056 // but in this case we have explicit API guarantees from the decompressor.
1057 #[allow(clippy::uninit_vec)]
1058 unsafe {
1059 out.set_len(out_len);
1060 }
1061
1062 if let Some(enc) = &self.encryptor {
1063 let compressed = enc.decrypt(&raw, block_idx)?;
1064 self.compressor.decompress_into(&compressed, &mut out)?;
1065 } else {
1066 self.compressor.decompress_into(raw.as_ref(), &mut out)?;
1067 }
1068
1069 Ok(Bytes::from(out))
1070 }
1071
1072 /// Resolves raw block data by fetching from cache or decompressing from storage.
1073 ///
1074 /// This is the core decompression path. It:
1075 /// 1. Checks the block cache
1076 /// 2. Reads compressed block from backend
1077 /// 3. Verifies CRC32 checksum (if stored) and returns `Corruption(block_idx)` on mismatch
1078 /// 4. Decrypts (if encrypted)
1079 /// 5. Decompresses
1080 /// 6. Caches the result
1081 ///
1082 /// # Parameters
1083 ///
1084 /// - `stream`: Stream identifier (for cache key)
1085 /// - `block_idx`: Global block index
1086 /// - `info`: Block metadata (offset, length, compression)
1087 ///
1088 /// # Returns
1089 ///
1090 /// Decompressed block data as `Bytes` (zero-copy on cache hit).
1091 ///
1092 /// # Performance
1093 ///
1094 /// This method is hot path for cache misses. Decompression throughput:
1095 /// - LZ4: ~2 GB/s per core
1096 /// - Zstd: ~500 MB/s per core
1097 fn resolve_block_data(
1098 &self,
1099 stream: SnapshotStream,
1100 block_idx: u64,
1101 info: &BlockInfo,
1102 ) -> Result<Bytes> {
1103 // Fetch block (from cache or I/O). The FetchResult tracks whether
1104 // data is already decompressed, avoiding a TOCTOU race where a
1105 // background prefetch thread could modify the cache between fetch
1106 // and the decompression decision.
1107 match self.fetch_raw_block(stream, block_idx, info)? {
1108 FetchResult::Decompressed(data) => Ok(data),
1109 FetchResult::Compressed(raw) => {
1110 let data = self.decompress_and_verify(raw, block_idx, info)?;
1111 self.cache_l1.insert(stream, block_idx, data.clone());
1112 Ok(data)
1113 }
1114 }
1115 }
1116}