parallel_processor/buckets/readers/
binary_reader.rs

1use crate::buckets::readers::compressed_decoder::CompressedStreamDecoder;
2use crate::buckets::readers::lock_free_decoder::LockFreeStreamDecoder;
3use crate::buckets::writers::{BucketCheckpoints, BucketHeader};
4use crate::memory_fs::file::reader::{FileRangeReference, FileReader};
5use crate::memory_fs::{MemoryFs, RemoveFileMode};
6use crate::DEFAULT_BINCODE_CONFIG;
7use bincode::Decode;
8use desse::Desse;
9use desse::DesseSized;
10use parking_lot::Mutex;
11use std::collections::VecDeque;
12use std::fmt::Debug;
13use std::io::{Read, Seek, SeekFrom};
14use std::path::{Path, PathBuf};
15use std::sync::Arc;
16
17pub trait ChunkDecoder {
18    const MAGIC_HEADER: &'static [u8; 16];
19    type ReaderType: Read;
20    fn decode_stream(reader: FileReader, size: u64) -> Self::ReaderType;
21    fn dispose_stream(stream: Self::ReaderType) -> FileReader;
22}
23
24pub(crate) struct RemoveFileGuard {
25    path: PathBuf,
26    remove_mode: RemoveFileMode,
27}
28
29impl Drop for RemoveFileGuard {
30    fn drop(&mut self) {
31        MemoryFs::remove_file(&self.path, self.remove_mode).unwrap();
32    }
33}
34
35#[derive(Clone)]
36pub struct BinaryReaderChunk {
37    pub(crate) reader: FileReader,
38    pub(crate) length: u64,
39    pub(crate) remove_guard: Arc<RemoveFileGuard>,
40    pub(crate) extra_data: Option<Vec<u8>>,
41    pub(crate) decoder_type: DecoderType,
42}
43
44impl Debug for BinaryReaderChunk {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("BinaryReaderChunk")
47            .field("length", &self.length)
48            .field("extra_data", &self.extra_data)
49            .field("decoder_type", &self.decoder_type)
50            .field("file_name", &self.reader.get_file_path().display())
51            .finish()
52    }
53}
54
55impl BinaryReaderChunk {
56    pub fn get_extra_data<E: Decode<()>>(&self) -> Option<E> {
57        self.extra_data.as_ref().map(|data| {
58            bincode::decode_from_slice(&data, DEFAULT_BINCODE_CONFIG)
59                .unwrap()
60                .0
61        })
62    }
63
64    pub fn get_decoder_type(&self) -> DecoderType {
65        self.decoder_type
66    }
67
68    pub fn get_unique_file_id(&self) -> usize {
69        self.reader.get_unique_file_id()
70    }
71
72    pub fn get_length(&self) -> u64 {
73        self.length
74    }
75}
76
77#[derive(Clone, Copy, PartialEq, Eq, Debug)]
78pub enum DecoderType {
79    LockFree,
80    Compressed,
81}
82
83pub struct ChunkedBinaryReaderIndex {
84    file_path: PathBuf,
85    chunks: Vec<BinaryReaderChunk>,
86    format_data_info: Vec<u8>,
87    file_size: u64,
88}
89
90pub enum DecodeItemsStatus<T> {
91    Decompressed,
92    Passtrough {
93        file_range: FileRangeReference,
94        data: Option<T>,
95    },
96}
97
98impl ChunkedBinaryReaderIndex {
99    fn get_chunk_size(
100        checkpoints: &BucketCheckpoints,
101        last_byte_position: u64,
102        index: usize,
103    ) -> u64 {
104        if checkpoints.index.len() > (index + 1) as usize {
105            checkpoints.index[(index + 1) as usize].offset
106                - checkpoints.index[index as usize].offset
107        } else {
108            last_byte_position - checkpoints.index[index as usize].offset
109        }
110    }
111
112    pub fn get_file_size(&self) -> u64 {
113        self.file_size
114    }
115
116    pub fn from_file(
117        name: impl AsRef<Path>,
118        remove_file: RemoveFileMode,
119    ) -> Self {
120        let mut file = FileReader::open(&name)
121            .unwrap_or_else(|| panic!("Cannot open file {}", name.as_ref().display()));
122
123        let mut header_buffer = [0; BucketHeader::SIZE];
124        file.read_exact(&mut header_buffer)
125            .unwrap_or_else(|_| panic!("File {} is corrupted", name.as_ref().display()));
126
127        let header: BucketHeader = BucketHeader::deserialize_from(&header_buffer);
128
129        file.seek(SeekFrom::Start(header.index_offset)).unwrap();
130        let index: BucketCheckpoints =
131            bincode::decode_from_std_read(&mut file, DEFAULT_BINCODE_CONFIG).unwrap();
132
133        let remove_guard = Arc::new(RemoveFileGuard {
134            path: name.as_ref().to_path_buf(),
135            remove_mode: remove_file,
136        });
137
138        let mut chunks = Vec::with_capacity(index.index.len());
139
140        let decoder_type = match &header.magic {
141            LockFreeStreamDecoder::MAGIC_HEADER => DecoderType::LockFree,
142            CompressedStreamDecoder::MAGIC_HEADER => DecoderType::Compressed,
143            _ => panic!(
144                "Invalid decode bucket header magic: {:?}. This is a bug",
145                header.magic
146            ),
147        };
148
149        if index.index.len() > 0 {
150            file.seek(SeekFrom::Start(index.index[0].offset)).unwrap();
151
152            for (chunk_idx, chunk) in index.index.iter().enumerate() {
153                let length = Self::get_chunk_size(&index, header.index_offset, chunk_idx);
154
155                chunks.push(BinaryReaderChunk {
156                    reader: file.clone(),
157                    length,
158                    remove_guard: remove_guard.clone(),
159                    extra_data: chunk.data.clone(),
160                    decoder_type,
161                });
162
163                file.seek(SeekFrom::Current(length as i64)).unwrap();
164            }
165        }
166
167        Self {
168            file_path: name.as_ref().to_path_buf(),
169            chunks,
170            format_data_info: header.data_format_info.to_vec(),
171            file_size: MemoryFs::get_file_size(name).unwrap() as u64,
172        }
173    }
174
175    pub fn get_data_format_info<T: Decode<()>>(&self) -> T {
176        bincode::decode_from_slice(&self.format_data_info, DEFAULT_BINCODE_CONFIG)
177            .unwrap()
178            .0
179    }
180
181    pub fn get_path(&self) -> &Path {
182        &self.file_path
183    }
184
185    pub fn into_chunks(self) -> Vec<BinaryReaderChunk> {
186        self.chunks
187    }
188
189    pub fn into_parallel_chunks(self) -> Mutex<VecDeque<BinaryReaderChunk>> {
190        Mutex::new(self.chunks.into())
191    }
192}
193
194pub struct BinaryChunkReader<D: ChunkDecoder> {
195    reader: D::ReaderType,
196    _remove_guard: Arc<RemoveFileGuard>,
197}
198
199impl<D: ChunkDecoder> BinaryChunkReader<D> {
200    pub fn new(chunk: BinaryReaderChunk) -> Self {
201        Self {
202            reader: D::decode_stream(chunk.reader, chunk.length),
203            _remove_guard: chunk.remove_guard,
204        }
205    }
206}
207
208impl<D: ChunkDecoder> Read for BinaryChunkReader<D> {
209    #[inline(always)]
210    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
211        self.reader.read(buf)
212    }
213}
parallel_processor/buckets/readers/binary_reader.rs

parallel_processor/buckets/readers/
binary_reader.rs
