parallel_processor/buckets/
mod.rs

1use arc_swap::ArcSwap;
2use bincode::{Decode, Encode};
3use parking_lot::Mutex;
4use std::num::NonZeroU64;
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7
8use crate::memory_fs::file::reader::FileRangeReference;
9use crate::DEFAULT_BINCODE_CONFIG;
10
11pub mod bucket_writer;
12pub mod concurrent;
13pub mod readers;
14pub mod single;
15pub mod writers;
16
17/// This enum serves as a way to specifying the behavior of the
18/// bucket portions created after setting the checkpoint data.
19/// If set on passtrough there is the option to directly read binary data and copy it somewhere else
20#[derive(Encode, Decode, Copy, Clone, Debug, PartialEq, Eq)]
21pub enum CheckpointStrategy {
22    Decompress,
23    Passtrough,
24}
25
26#[derive(Encode, Decode, Clone, Debug, PartialEq, Eq)]
27pub(crate) struct CheckpointData {
28    offset: u64,
29    data: Option<Vec<u8>>,
30}
31
32impl PartialOrd for CheckpointData {
33    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
34        self.offset.partial_cmp(&other.offset)
35    }
36}
37
38impl Ord for CheckpointData {
39    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
40        self.offset.cmp(&other.offset)
41    }
42}
43
44pub trait LockFreeBucket: Sized {
45    type InitData: Clone;
46
47    fn new_serialized_data_format(
48        path: &Path,
49        data: &Self::InitData,
50        index: usize,
51        data_format: &[u8],
52    ) -> Self;
53
54    fn new<T: Encode>(path: &Path, data: &Self::InitData, index: usize, data_format: &T) -> Self {
55        Self::new_serialized_data_format(
56            path,
57            data,
58            index,
59            &bincode::encode_to_vec(data_format, DEFAULT_BINCODE_CONFIG).unwrap(),
60        )
61    }
62
63    fn set_checkpoint_data<T: Encode>(
64        &self,
65        data: Option<&T>,
66        passtrough_range: Option<FileRangeReference>,
67    );
68
69    fn get_bucket_size(&self) -> u64;
70    fn write_data(&self, bytes: &[u8]) -> u64;
71    fn get_path(&self) -> PathBuf;
72    fn finalize(self);
73}
74
75#[derive(Encode, Decode, Debug, Clone)]
76pub struct MultiChunkBucket {
77    pub index: usize,
78    pub chunks: Vec<PathBuf>,
79    pub extra_bucket_data: Option<ExtraBucketData>,
80}
81
82impl MultiChunkBucket {
83    pub fn into_single(mut self) -> SingleBucket {
84        assert!(self.chunks.len() == 1);
85        SingleBucket {
86            index: self.index,
87            path: self.chunks.pop().unwrap(),
88            extra_bucket_data: self.extra_bucket_data,
89        }
90    }
91}
92
93#[derive(Encode, Decode)]
94pub struct SingleBucket {
95    pub index: usize,
96    pub path: PathBuf,
97    pub extra_bucket_data: Option<ExtraBucketData>,
98}
99
100impl SingleBucket {
101    pub fn to_multi_chunk(self) -> MultiChunkBucket {
102        MultiChunkBucket {
103            index: self.index,
104            chunks: vec![self.path],
105            extra_bucket_data: self.extra_bucket_data,
106        }
107    }
108}
109
110#[derive(Debug, Clone)]
111pub enum ChunkingStatus {
112    SameChunk,
113    NewChunk,
114}
115
116#[derive(Encode, Decode, Copy, Clone, Debug, PartialEq, Eq)]
117pub struct BucketsCount {
118    pub normal_buckets_count: usize,
119    pub normal_buckets_count_log: usize,
120    pub total_buckets_count: usize,
121    pub extra_buckets_count: ExtraBuckets,
122}
123
124impl BucketsCount {
125    pub const ONE: Self = Self::new(0, ExtraBuckets::None);
126
127    pub fn from_power_of_two(size: usize, extra_buckets: ExtraBuckets) -> Self {
128        assert_eq!(size, size.next_power_of_two());
129        Self::new(size.ilog2() as usize, extra_buckets)
130    }
131
132    pub const fn new(size_log: usize, extra_buckets: ExtraBuckets) -> Self {
133        let normal_buckets_count = 1 << size_log;
134        let extra_buckets_count = match extra_buckets {
135            ExtraBuckets::None => 0,
136            ExtraBuckets::Extra { count, .. } => count,
137        };
138
139        Self {
140            normal_buckets_count,
141            normal_buckets_count_log: size_log,
142            total_buckets_count: normal_buckets_count + extra_buckets_count,
143            extra_buckets_count: extra_buckets,
144        }
145    }
146
147    pub fn get_extra_buckets_count(&self) -> usize {
148        match self.extra_buckets_count {
149            ExtraBuckets::None => 0,
150            ExtraBuckets::Extra { count, .. } => count,
151        }
152    }
153}
154
155pub struct MultiThreadBuckets<B: LockFreeBucket> {
156    active_buckets: Vec<ArcSwap<B>>,
157    stored_buckets: Vec<Mutex<MultiChunkBucket>>,
158    chunking_size_threshold: Option<NonZeroU64>,
159    bucket_count_lock: Mutex<usize>,
160    base_path: Option<PathBuf>,
161    init_data: Option<B::InitData>,
162    serialized_format_info: Vec<u8>,
163    size: BucketsCount,
164}
165
166#[derive(Encode, Decode, Copy, Clone, Debug, PartialEq, Eq)]
167pub struct ExtraBucketData(pub usize);
168
169#[derive(Encode, Decode, Copy, Clone, Debug, PartialEq, Eq)]
170pub enum ExtraBuckets {
171    None,
172    Extra { count: usize, data: ExtraBucketData },
173}
174
175impl<B: LockFreeBucket> MultiThreadBuckets<B> {
176    pub const EMPTY: Self = Self {
177        active_buckets: vec![],
178        stored_buckets: vec![],
179        chunking_size_threshold: None,
180        bucket_count_lock: Mutex::new(0),
181        base_path: None,
182        init_data: None,
183        serialized_format_info: vec![],
184        size: BucketsCount {
185            normal_buckets_count: 0,
186            normal_buckets_count_log: 0,
187            total_buckets_count: 0,
188            extra_buckets_count: ExtraBuckets::None,
189        },
190    };
191
192    pub fn get_buckets_count(&self) -> &BucketsCount {
193        &self.size
194    }
195
196    pub fn create_matching_multichunks(&self) -> Vec<Mutex<MultiChunkBucket>> {
197        self.stored_buckets
198            .iter()
199            .map(|s| {
200                let s = s.lock();
201                Mutex::new(MultiChunkBucket {
202                    index: s.index,
203                    chunks: vec![],
204                    extra_bucket_data: s.extra_bucket_data.clone(),
205                })
206            })
207            .collect()
208    }
209
210    pub fn new(
211        size: BucketsCount,
212        path: PathBuf,
213        chunking_size_threshold: Option<u64>,
214        init_data: &B::InitData,
215        format_info: &impl Encode,
216    ) -> MultiThreadBuckets<B> {
217        let mut buckets = Vec::with_capacity(size.total_buckets_count);
218
219        for i in 0..size.total_buckets_count {
220            buckets.push(ArcSwap::from_pointee(B::new(
221                &path,
222                init_data,
223                i,
224                format_info,
225            )));
226        }
227        MultiThreadBuckets {
228            active_buckets: buckets,
229            stored_buckets: (0..size.total_buckets_count)
230                .map(|index| {
231                    Mutex::new(MultiChunkBucket {
232                        index,
233                        chunks: vec![],
234                        extra_bucket_data: match size.extra_buckets_count {
235                            ExtraBuckets::None => None,
236                            ExtraBuckets::Extra { data, .. } => {
237                                if index >= size.normal_buckets_count {
238                                    Some(data)
239                                } else {
240                                    None
241                                }
242                            }
243                        },
244                    })
245                })
246                .collect(),
247            chunking_size_threshold: chunking_size_threshold.map(NonZeroU64::new).flatten(),
248            bucket_count_lock: Mutex::new(size.total_buckets_count),
249            base_path: Some(path),
250            init_data: Some(init_data.clone()),
251            serialized_format_info: bincode::encode_to_vec(format_info, DEFAULT_BINCODE_CONFIG)
252                .unwrap(),
253            size,
254        }
255    }
256
257    pub fn get_stored_buckets(&self) -> &Vec<Mutex<MultiChunkBucket>> {
258        &self.stored_buckets
259    }
260
261    pub fn into_buckets(mut self) -> impl Iterator<Item = B> {
262        assert!(
263            self.stored_buckets
264                .iter_mut()
265                .all(|bucket| bucket.get_mut().chunks.is_empty())
266                && self.chunking_size_threshold.is_none()
267        );
268        let buckets = std::mem::take(&mut self.active_buckets);
269        buckets
270            .into_iter()
271            .map(|bucket| Arc::into_inner(bucket.into_inner()).unwrap())
272    }
273
274    pub fn get_path(&self, bucket: u16) -> PathBuf {
275        self.active_buckets[bucket as usize].load().get_path()
276    }
277
278    fn take_active_bucket(&self, index: u16, next_progressive: usize) -> PathBuf {
279        let mut taken_bucket =
280            self.active_buckets[index as usize].swap(Arc::new(B::new_serialized_data_format(
281                &self.base_path.as_deref().unwrap(),
282                &self.init_data.as_ref().unwrap(),
283                next_progressive,
284                &self.serialized_format_info,
285            )));
286
287        let taken_bucket = loop {
288            // Wait for the bucket to end all the pending writes before finalizing it
289            match Arc::try_unwrap(taken_bucket) {
290                Ok(bucket) => break bucket,
291                Err(waiting_arc) => {
292                    taken_bucket = waiting_arc;
293                    std::hint::spin_loop();
294                }
295            }
296        };
297        let path = taken_bucket.get_path();
298        taken_bucket.finalize();
299        path
300    }
301
302    pub fn add_data(&self, index: u16, data: &[u8]) -> ChunkingStatus {
303        let bucket_guard = self.active_buckets[index as usize].load();
304        let last_bucket_size = bucket_guard.write_data(data);
305
306        drop(bucket_guard);
307
308        // If the disk usage limit is set, check if the disk usage is not exceeded
309        if let Some(chunk_threshold) = self.chunking_size_threshold {
310            if last_bucket_size >= chunk_threshold.get() {
311                let mut buckets_count = self.bucket_count_lock.lock();
312
313                if self.active_buckets[index as usize].load().get_bucket_size()
314                    < chunk_threshold.get()
315                {
316                    // Do not add a new chunk if a previous writer already swapped it
317                    return ChunkingStatus::SameChunk;
318                }
319
320                // Take the largest bucket and add it to the stored buckets
321
322                let bucket_path = self.take_active_bucket(index, *buckets_count);
323
324                // Add the bucket to the stored buckets and clear its active usage
325                self.stored_buckets[index as usize]
326                    .lock()
327                    .chunks
328                    .push(bucket_path);
329
330                *buckets_count += 1;
331                ChunkingStatus::NewChunk
332            } else {
333                ChunkingStatus::SameChunk
334            }
335        } else {
336            ChunkingStatus::SameChunk
337        }
338    }
339
340    pub fn take_bucket(self: &Arc<Self>, index: u16) -> MultiChunkBucket {
341        let mut files = vec![];
342
343        let mut buckets_count = self.bucket_count_lock.lock();
344
345        let active_bucket = self.take_active_bucket(index, *buckets_count);
346        *buckets_count += 1;
347        drop(buckets_count);
348
349        files.push(active_bucket);
350        let mut stored_buckets = self.stored_buckets[index as usize].lock();
351        files.extend(stored_buckets.chunks.drain(..));
352        MultiChunkBucket {
353            index: index as usize,
354            chunks: files,
355            extra_bucket_data: stored_buckets.extra_bucket_data,
356        }
357    }
358
359    pub fn finalize_single(self: Arc<Self>) -> Vec<SingleBucket> {
360        assert!(self.chunking_size_threshold.is_none());
361        let buckets = self.finalize();
362        buckets
363            .into_iter()
364            .map(|mut bucket| {
365                assert!(bucket.chunks.len() == 1);
366                SingleBucket {
367                    index: bucket.index,
368                    path: bucket.chunks.pop().unwrap(),
369                    extra_bucket_data: bucket.extra_bucket_data,
370                }
371            })
372            .collect()
373    }
374
375    pub fn finalize(self: Arc<Self>) -> Vec<MultiChunkBucket> {
376        let mut self_ = Arc::try_unwrap(self)
377            .unwrap_or_else(|_| panic!("Cannot take full ownership of multi thread buckets!"));
378
379        self_
380            .active_buckets
381            .drain(..)
382            .zip(self_.stored_buckets.drain(..))
383            .map(|(bucket, stored)| {
384                let mut stored = stored.into_inner();
385                let bucket = Arc::into_inner(bucket.into_inner()).unwrap();
386                stored.chunks.push(bucket.get_path());
387                bucket.finalize();
388                stored
389            })
390            .collect()
391    }
392}
393
394impl<B: LockFreeBucket> Drop for MultiThreadBuckets<B> {
395    fn drop(&mut self) {
396        self.active_buckets.drain(..).for_each(|bucket| {
397            let bucket = Arc::into_inner(bucket.into_inner()).unwrap();
398            bucket.finalize();
399        });
400    }
401}
402
403unsafe impl<B: LockFreeBucket> Send for MultiThreadBuckets<B> {}
404
405unsafe impl<B: LockFreeBucket> Sync for MultiThreadBuckets<B> {}
parallel_processor/buckets/mod.rs

parallel_processor/buckets/
mod.rs
