diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs index 4a9933f264..49d5e44ffa 100644 --- a/bucket_map/src/bucket.rs +++ b/bucket_map/src/bucket.rs @@ -161,6 +161,18 @@ impl Bucket { Err(BucketMapError::IndexNoSpace(index.capacity_pow2)) } + pub fn addref(&mut self, key: &Pubkey) -> Option { + let (elem, _) = self.find_entry_mut(key)?; + elem.ref_count += 1; + Some(elem.ref_count) + } + + pub fn unref(&mut self, key: &Pubkey) -> Option { + let (elem, _) = self.find_entry_mut(key)?; + elem.ref_count -= 1; + Some(elem.ref_count) + } + fn create_key(&self, key: &Pubkey, ref_count: u64) -> Result { Self::bucket_create_key( &self.index, diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs index dacd5bf1bb..12a3da78a1 100644 --- a/bucket_map/src/bucket_map.rs +++ b/bucket_map/src/bucket_map.rs @@ -196,6 +196,20 @@ impl BucketMap { 0 } } + + /// Increment the refcount for Pubkey `key` + pub fn addref(&self, key: &Pubkey) -> Option { + let ix = self.bucket_ix(key); + let mut bucket = self.buckets[ix].write().unwrap(); + bucket.as_mut()?.addref(key) + } + + /// Decrement the refcount for Pubkey `key` + pub fn unref(&self, key: &Pubkey) -> Option { + let ix = self.bucket_ix(key); + let mut bucket = self.buckets[ix].write().unwrap(); + bucket.as_mut()?.unref(key) + } } /// Look at the first 8 bytes of the input and reinterpret them as a u64