Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

disk index: refactoring get/get_mut fns #31025

Merged
merged 1 commit into from
Apr 3, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions bucket_map/src/bucket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use {
bucket_item::BucketItem,
bucket_map::BucketMapError,
bucket_stats::BucketMapStats,
bucket_storage::{BucketOccupied, BucketStorage, DEFAULT_CAPACITY_POW2},
bucket_storage::{BucketOccupied, BucketStorage, IncludeHeader, DEFAULT_CAPACITY_POW2},
index_entry::{
DataBucket, IndexBucket, IndexEntry, IndexEntryPlaceInBucket, MultipleSlots,
OccupiedEnum,
Expand Down Expand Up @@ -328,7 +328,11 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
if best_fit_bucket == bucket_ix as u64 {
// in place update in same data file
assert!(!current_bucket.is_free(elem_loc));
let slice: &mut [T] = current_bucket.get_mut_cell_slice(elem_loc, data_len as u64);
let slice: &mut [T] = current_bucket.get_mut_cell_slice(
elem_loc,
data_len as u64,
IncludeHeader::NoHeader,
);
multiple_slots.set_num_slots(num_slots);

slice.iter_mut().zip(data).for_each(|(dest, src)| {
Expand Down Expand Up @@ -375,7 +379,8 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
// copy slotlist into the data bucket
let best_bucket = &mut self.data[best_fit_bucket as usize];
best_bucket.occupy(ix, false).unwrap();
let slice = best_bucket.get_mut_cell_slice(ix, num_slots);
let slice =
best_bucket.get_mut_cell_slice(ix, num_slots, IncludeHeader::NoHeader);
slice.iter_mut().zip(data).for_each(|(dest, src)| {
*dest = *src;
});
Expand Down
37 changes: 27 additions & 10 deletions bucket_map/src/bucket_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,14 @@ impl<O: BucketOccupied> Drop for BucketStorage<O> {
}
}

#[allow(dead_code)]
pub(crate) enum IncludeHeader {
/// caller wants header
Header,
/// caller wants header skipped
NoHeader,
}

impl<O: BucketOccupied> BucketStorage<O> {
pub fn new_with_capacity(
drives: Arc<Vec<PathBuf>>,
Expand Down Expand Up @@ -175,18 +183,22 @@ impl<O: BucketOccupied> BucketStorage<O> {
(self.cell_size * ix) as usize
}

fn get_start_offset_no_header(&self, ix: u64) -> usize {
self.get_start_offset_with_header(ix) + O::offset_to_first_data()
fn get_start_offset(&self, ix: u64, header: IncludeHeader) -> usize {
self.get_start_offset_with_header(ix)
+ match header {
IncludeHeader::Header => 0,
IncludeHeader::NoHeader => O::offset_to_first_data(),
}
}

pub fn get<T>(&self, ix: u64) -> &T {
let slice = self.get_cell_slice::<T>(ix, 1);
pub(crate) fn get<T>(&self, ix: u64) -> &T {
let slice = self.get_cell_slice::<T>(ix, 1, IncludeHeader::NoHeader);
// SAFETY: `get_cell_slice` ensures there's at least one element in the slice
unsafe { slice.get_unchecked(0) }
}

pub fn get_mut<T>(&mut self, ix: u64) -> &mut T {
let slice = self.get_mut_cell_slice::<T>(ix, 1);
pub(crate) fn get_mut<T>(&mut self, ix: u64) -> &mut T {
let slice = self.get_mut_cell_slice::<T>(ix, 1, IncludeHeader::NoHeader);
// SAFETY: `get_mut_cell_slice` ensures there's at least one element in the slice
unsafe { slice.get_unchecked_mut(0) }
}
Expand All @@ -203,8 +215,8 @@ impl<O: BucketOccupied> BucketStorage<O> {
unsafe { &*item }
}

pub fn get_cell_slice<T>(&self, ix: u64, len: u64) -> &[T] {
let start = self.get_start_offset_no_header(ix);
pub(crate) fn get_cell_slice<T>(&self, ix: u64, len: u64, header: IncludeHeader) -> &[T] {
let start = self.get_start_offset(ix, header);
let slice = {
let size = std::mem::size_of::<T>() * len as usize;
let slice = &self.mmap[start..];
Expand All @@ -219,8 +231,13 @@ impl<O: BucketOccupied> BucketStorage<O> {
unsafe { std::slice::from_raw_parts(ptr, len as usize) }
}

pub fn get_mut_cell_slice<T>(&mut self, ix: u64, len: u64) -> &mut [T] {
let start = self.get_start_offset_no_header(ix);
pub(crate) fn get_mut_cell_slice<T>(
&mut self,
ix: u64,
len: u64,
header: IncludeHeader,
) -> &mut [T] {
let start = self.get_start_offset(ix, header);
let slice = {
let size = std::mem::size_of::<T>() * len as usize;
let slice = &mut self.mmap[start..];
Expand Down
8 changes: 6 additions & 2 deletions bucket_map/src/index_entry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

use {
crate::{
bucket_storage::{BucketOccupied, BucketStorage},
bucket_storage::{BucketOccupied, BucketStorage, IncludeHeader},
RefCount,
},
bv::BitVec,
Expand Down Expand Up @@ -308,7 +308,11 @@ impl<T: Copy> IndexEntryPlaceInBucket<T> {
let data_bucket = &data_buckets[data_bucket_ix as usize];
let loc = multiple_slots.data_loc(data_bucket);
assert!(!data_bucket.is_free(loc));
data_bucket.get_cell_slice::<T>(loc, multiple_slots.num_slots)
data_bucket.get_cell_slice::<T>(
loc,
multiple_slots.num_slots,
IncludeHeader::NoHeader,
)
}
_ => {
panic!("trying to read data from a free entry");
Expand Down