add bucketmap 'insert' to eliminate copies (#20020)

bucket_map/src/bucket.rs

@@ -366,19 +366,10 @@ impl<T: Clone + Copy> Bucket<T> {
         //debug!(            "INDEX_IX: {:?} uid:{} loc: {} cap:{}",            key,            uid,            location,            index.num_cells()        );
     }
 
-    pub fn update<F>(&mut self, key: &Pubkey, updatefn: F)
-    where
-        F: Fn(Option<(&[T], RefCount)>) -> Option<(Vec<T>, 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<T: Clone + Copy> Bucket<T> {
             }
         }
     }
+
+    pub fn update<F>(&mut self, key: &Pubkey, updatefn: F)
+    where
+        F: Fn(Option<(&[T], RefCount)>) -> Option<(Vec<T>, 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));
+    }
 }

bucket_map/src/bucket_map.rs

@@ -166,6 +166,21 @@ impl<T: Clone + Copy + Debug> BucketMap<T> {
         }
     }
 
+    /// 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<F>(&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));