bucket map stores uid in use field in memory bitvec

bucket_map/Cargo.toml

@@ -11,6 +11,7 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+bv = { workspace = true, features = ["serde"] }
 log = { workspace = true }
 memmap2 = { workspace = true }
 modular-bitfield = { workspace = true }

bucket_map/src/bucket.rs

@@ -82,6 +82,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         max_search: MaxSearch,
         stats: Arc<BucketMapStats>,
         count: Arc<AtomicU64>,
+        use_bit_field: bool,
     ) -> Self {
         let index = BucketStorage::new(
             Arc::clone(&drives),
@@ -90,6 +91,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
             max_search,
             Arc::clone(&stats.index),
             count,
+            use_bit_field,
         );
         Self {
             random: thread_rng().gen(),
@@ -142,10 +144,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         Self::bucket_find_entry(&self.index, key, self.random)
     }
 
-    fn find_entry_mut<'a>(
-        &'a self,
-        key: &Pubkey,
-    ) -> Result<(bool, &'a mut IndexEntry, u64), BucketMapError> {
+    fn find_entry_mut<'a>(&mut self, key: &Pubkey) -> Result<(bool, u64), BucketMapError> {
         let ix = Self::bucket_index_ix(&self.index, key, self.random);
         let mut first_free = None;
         let mut m = Measure::start("bucket_find_entry_mut");
@@ -164,7 +163,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                     .index
                     .find_entry_mut_us
                     .fetch_add(m.as_us(), Ordering::Relaxed);
-                return Ok((true, elem, ii));
+                return Ok((true, ii));
             }
         }
         m.stop();
@@ -175,7 +174,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         match first_free {
             Some(ii) => {
                 let elem: &mut IndexEntry = self.index.get_mut(ii);
-                Ok((false, elem, ii))
+                Ok((false, ii))
             }
             None => Err(self.index_no_space()),
         }
@@ -189,6 +188,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         let ix = Self::bucket_index_ix(index, key, random);
         for i in ix..ix + index.max_search() {
             let ii = i % index.capacity();
+            assert!(ii < index.capacity(), "{}, {}", ii, index.capacity());
             if index.is_free(ii) {
                 continue;
             }
@@ -236,8 +236,9 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
     }
 
     pub fn addref(&mut self, key: &Pubkey) -> Option<RefCount> {
-        if let Ok((found, elem, _)) = self.find_entry_mut(key) {
+        if let Ok((found, elem_ix)) = self.find_entry_mut(key) {
             if found {
+                let elem: &mut IndexEntry = self.index.get_mut(elem_ix);
                 elem.ref_count += 1;
                 return Some(elem.ref_count);
             }
@@ -246,8 +247,9 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
     }
 
     pub fn unref(&mut self, key: &Pubkey) -> Option<RefCount> {
-        if let Ok((found, elem, _)) = self.find_entry_mut(key) {
+        if let Ok((found, elem_ix)) = self.find_entry_mut(key) {
             if found {
+                let elem: &mut IndexEntry = self.index.get_mut(elem_ix);
                 elem.ref_count -= 1;
                 return Some(elem.ref_count);
             }
@@ -277,14 +279,18 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
             // fail early if the data bucket we need doesn't exist - we don't want the index entry partially allocated
             return Err(BucketMapError::DataNoSpace((best_fit_bucket, 0)));
         }
-        let (found, elem, elem_ix) = self.find_entry_mut(key)?;
+        let (found, elem_ix) = self.find_entry_mut(key)?;
+        let elem: &mut IndexEntry;
         if !found {
             let is_resizing = false;
-            let elem_uid = IndexEntry::key_uid(key);
+            let elem_uid = IndexEntry::key_uid2(key);
             self.index.allocate(elem_ix, elem_uid, is_resizing).unwrap();
+            elem = self.index.get_mut(elem_ix);
             // These fields will be overwritten after allocation by callers.
             // Since this part of the mmapped file could have previously been used by someone else, there can be garbage here.
             elem.init(key);
+        } else {
+            elem = self.index.get_mut(elem_ix);
         }
         elem.ref_count = ref_count;
         let elem_uid = self.index.uid_unchecked(elem_ix);
@@ -378,6 +384,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                     self.index.max_search,
                     Arc::clone(&self.stats.index),
                     Arc::clone(&self.index.count),
+                    self.index.bit_field.is_some(),
                 );
                 let random = thread_rng().gen();
                 let mut valid = true;
@@ -440,6 +447,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                     self.index.max_search,
                     Arc::clone(&self.stats.data),
                     Arc::default(),
+                    false,
                 ))
             }
             self.data.push(bucket);
@@ -450,7 +458,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
 
     /// grow a data bucket
     /// The application of the new bucket is deferred until the next write lock.
-    pub fn grow_data(&self, data_index: u64, current_capacity_pow2: u8) {
+    pub fn grow_data(&self, data_index: u64, current_capacity_pow2: u8, use_bit_field: bool) {
         let new_bucket = BucketStorage::new_resized(
             &self.drives,
             self.index.max_search,
@@ -459,14 +467,15 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
             1 << data_index,
             Self::elem_size(),
             &self.stats.data,
+            use_bit_field,
         );
         self.reallocated.add_reallocation();
         let mut items = self.reallocated.items.lock().unwrap();
         items.data = Some((data_index, new_bucket));
     }
 
     fn bucket_index_ix(index: &BucketStorage, key: &Pubkey, random: u64) -> u64 {
-        let uid = IndexEntry::key_uid(key);
+        let uid = IndexEntry::full_key_uid(key);
         let mut s = DefaultHasher::new();
         uid.hash(&mut s);
         //the locally generated random will make it hard for an attacker
@@ -484,7 +493,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         match err {
             BucketMapError::DataNoSpace((data_index, current_capacity_pow2)) => {
                 //debug!("GROWING SPACE {:?}", (data_index, current_capacity_pow2));
-                self.grow_data(data_index, current_capacity_pow2);
+                self.grow_data(data_index, current_capacity_pow2, false);
             }
             BucketMapError::IndexNoSpace(current_capacity_pow2) => {
                 //debug!("GROWING INDEX {}", sz);

bucket_map/src/bucket_api.rs

@@ -83,13 +83,15 @@ impl<T: Clone + Copy> BucketApi<T> {
     }
 
     fn get_write_bucket(&self) -> RwLockWriteGuard<Option<Bucket<T>>> {
+        let use_bit_field = true;
         let mut bucket = self.bucket.write().unwrap();
         if bucket.is_none() {
             *bucket = Some(Bucket::new(
                 Arc::clone(&self.drives),
                 self.max_search,
                 Arc::clone(&self.stats),
                 Arc::clone(&self.count),
+                use_bit_field,
             ));
         } else {
             let write = bucket.as_mut().unwrap();