diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs index 94f4a7e25a..48b73c10cb 100644 --- a/bucket_map/src/bucket.rs +++ b/bucket_map/src/bucket.rs @@ -366,19 +366,10 @@ impl Bucket { //debug!( "INDEX_IX: {:?} uid:{} loc: {} cap:{}", key, uid, location, index.num_cells() ); } - pub fn update(&mut self, key: &Pubkey, updatefn: F) - where - F: Fn(Option<(&[T], RefCount)>) -> Option<(Vec, RefCount)>, - { - let current = self.read_value(key); - let new = updatefn(current); - if new.is_none() { - self.delete_key(key); - return; - } - let (new, refct) = new.unwrap(); + pub fn insert(&mut self, key: &Pubkey, value: (&[T], RefCount)) { + let (new, refct) = value; loop { - let rv = self.try_write(key, &new, refct); + let rv = self.try_write(key, new, refct); match rv { Err(BucketMapError::DataNoSpace(sz)) => { //debug!("GROWING SPACE {:?}", sz); @@ -394,4 +385,18 @@ impl Bucket { } } } + + pub fn update(&mut self, key: &Pubkey, updatefn: F) + where + F: Fn(Option<(&[T], RefCount)>) -> Option<(Vec, RefCount)>, + { + let current = self.read_value(key); + let new = updatefn(current); + if new.is_none() { + self.delete_key(key); + return; + } + let (new, refct) = new.unwrap(); + self.insert(key, (&new, refct)); + } } diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs index 90edd09ffe..7699f2faf8 100644 --- a/bucket_map/src/bucket_map.rs +++ b/bucket_map/src/bucket_map.rs @@ -166,6 +166,21 @@ impl BucketMap { } } + /// Update Pubkey `key`'s value with 'value' + pub fn insert(&self, key: &Pubkey, value: (&[T], RefCount)) { + let ix = self.bucket_ix(key); + let mut bucket = self.buckets[ix].write().unwrap(); + if bucket.is_none() { + *bucket = Some(Bucket::new( + Arc::clone(&self.drives), + self.max_search, + Arc::clone(&self.stats), + )); + } + let bucket = bucket.as_mut().unwrap(); + bucket.insert(key, value) + } + /// Update Pubkey `key`'s value with function `updatefn` pub fn update(&self, key: &Pubkey, updatefn: F) where @@ -231,6 +246,26 @@ mod tests { assert_eq!(index.read_value(&key), Some((vec![0], 0))); } + #[test] + fn bucket_map_test_insert2() { + let key = Pubkey::new_unique(); + let config = BucketMapConfig::new(1 << 1); + let index = BucketMap::new(config); + index.insert(&key, (&[0], 0)); + assert_eq!(index.read_value(&key), Some((vec![0], 0))); + } + + #[test] + fn bucket_map_test_update2() { + let key = Pubkey::new_unique(); + let config = BucketMapConfig::new(1 << 1); + let index = BucketMap::new(config); + index.insert(&key, (&[0], 0)); + assert_eq!(index.read_value(&key), Some((vec![0], 0))); + index.insert(&key, (&[1], 0)); + assert_eq!(index.read_value(&key), Some((vec![1], 0))); + } + #[test] fn bucket_map_test_update() { let key = Pubkey::new_unique(); @@ -439,11 +474,16 @@ mod tests { let k = solana_sdk::pubkey::new_rand(); let v = gen_rand_value(); hash_map.write().unwrap().insert(k, v.clone()); + let insert = thread_rng().gen_range(0, 2) == 0; maps.iter().for_each(|map| { - map.update(&k, |current| { - assert!(current.is_none()); - Some(v.clone()) - }) + if insert { + map.insert(&k, (&v.0, v.1)) + } else { + map.update(&k, |current| { + assert!(current.is_none()); + Some(v.clone()) + }) + } }); return_key(k); } @@ -453,11 +493,16 @@ mod tests { let hm = hash_map.read().unwrap(); let (v, rc) = gen_rand_value(); let v_old = hm.get(&k); + let insert = thread_rng().gen_range(0, 2) == 0; maps.iter().for_each(|map| { - map.update(&k, |current| { - assert_eq!(current, v_old.map(|(v, rc)| (&v[..], *rc)), "{}", k); - Some((v.clone(), rc)) - }) + if insert { + map.insert(&k, (&v, rc)) + } else { + map.update(&k, |current| { + assert_eq!(current, v_old.map(|(v, rc)| (&v[..], *rc)), "{}", k); + Some((v.clone(), rc)) + }) + } }); drop(hm); hash_map.write().unwrap().insert(k, (v, rc));