solana/bucket_map/src/bucket.rs

use crate::bucket_item::BucketItem;
use crate::bucket_map::BucketMapError;
use crate::bucket_stats::BucketMapStats;
use crate::bucket_storage::{BucketStorage, UID_UNLOCKED};
use crate::index_entry::IndexEntry;
use crate::{MaxSearch, RefCount};
use rand::thread_rng;
use rand::Rng;
use solana_measure::measure::Measure;
use solana_sdk::pubkey::Pubkey;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::ops::RangeBounds;
use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::Arc;

// >= 2 instances of BucketStorage per 'bucket' in the bucket map. 1 for index, >= 1 for data
pub struct Bucket<T> {
    drives: Arc<Vec<PathBuf>>,
    //index
    index: BucketStorage,
    //random offset for the index
    random: u64,
    //storage buckets to store SlotSlice up to a power of 2 in len
    pub data: Vec<BucketStorage>,
    _phantom: PhantomData<T>,
    stats: Arc<BucketMapStats>,
}

impl<T: Clone + Copy> Bucket<T> {
    pub fn new(
        drives: Arc<Vec<PathBuf>>,
        max_search: MaxSearch,
        stats: Arc<BucketMapStats>,
    ) -> Self {
        let index = BucketStorage::new(
            Arc::clone(&drives),
            1,
            std::mem::size_of::<IndexEntry>() as u64,
            max_search,
            Arc::clone(&stats.index),
        );
        Self {
            random: thread_rng().gen(),
            drives,
            index,
            data: vec![],
            _phantom: PhantomData::default(),
            stats,
        }
    }

    pub fn bucket_len(&self) -> u64 {
        self.index.used.load(Ordering::Relaxed)
    }

    pub fn keys(&self) -> Vec<Pubkey> {
        let mut rv = vec![];
        for i in 0..self.index.num_cells() {
            if self.index.uid(i) == UID_UNLOCKED {
                continue;
            }
            let ix: &IndexEntry = self.index.get(i);
            rv.push(ix.key);
        }
        rv
    }

    pub fn items_in_range<R>(&self, range: Option<&R>) -> Vec<BucketItem<T>>
    where
        R: RangeBounds<Pubkey>,
    {
        let mut result = Vec::with_capacity(self.index.used.load(Ordering::Relaxed) as usize);
        for i in 0..self.index.num_cells() {
            let ii = i % self.index.num_cells();
            if self.index.uid(ii) == UID_UNLOCKED {
                continue;
            }
            let ix: &IndexEntry = self.index.get(ii);
            let key = ix.key;
            if range.map(|r| r.contains(&key)).unwrap_or(true) {
                let val = ix.read_value(self);
                result.push(BucketItem {
                    pubkey: key,
                    ref_count: ix.ref_count(),
                    slot_list: val.map(|(v, _ref_count)| v.to_vec()).unwrap_or_default(),
                });
            }
        }
        result
    }

    pub fn find_entry(&self, key: &Pubkey) -> Option<(&IndexEntry, u64)> {
        Self::bucket_find_entry(&self.index, key, self.random)
    }

    fn find_entry_mut(&self, key: &Pubkey) -> Option<(&mut IndexEntry, u64)> {
        Self::bucket_find_entry_mut(&self.index, key, self.random)
    }

    fn bucket_find_entry_mut<'a>(
        index: &'a BucketStorage,
        key: &Pubkey,
        random: u64,
    ) -> Option<(&'a mut IndexEntry, u64)> {
        let ix = Self::bucket_index_ix(index, key, random);
        for i in ix..ix + index.max_search() {
            let ii = i % index.num_cells();
            if index.uid(ii) == UID_UNLOCKED {
                continue;
            }
            let elem: &mut IndexEntry = index.get_mut(ii);
            if elem.key == *key {
                return Some((elem, ii));
            }
        }
        None
    }

    fn bucket_find_entry<'a>(
        index: &'a BucketStorage,
        key: &Pubkey,
        random: u64,
    ) -> Option<(&'a IndexEntry, u64)> {
        let ix = Self::bucket_index_ix(index, key, random);
        for i in ix..ix + index.max_search() {
            let ii = i % index.num_cells();
            if index.uid(ii) == UID_UNLOCKED {
                continue;
            }
            let elem: &IndexEntry = index.get(ii);
            if elem.key == *key {
                return Some((elem, ii));
            }
        }
        None
    }

    fn bucket_create_key(
        index: &BucketStorage,
        key: &Pubkey,
        elem_uid: u64,
        random: u64,
        ref_count: u64,
    ) -> Result<u64, BucketMapError> {
        let ix = Self::bucket_index_ix(index, key, random);
        for i in ix..ix + index.max_search() {
            let ii = i as u64 % index.num_cells();
            if index.uid(ii) != UID_UNLOCKED {
                continue;
            }
            index.allocate(ii, elem_uid).unwrap();
            let mut elem: &mut IndexEntry = index.get_mut(ii);
            elem.key = *key;
            elem.ref_count = ref_count;
            elem.storage_offset = 0;
            elem.storage_capacity_when_created_pow2 = 0;
            elem.num_slots = 0;
            //debug!(                "INDEX ALLOC {:?} {} {} {}",                key, ii, index.capacity, elem_uid            );
            return Ok(ii);
        }
        Err(BucketMapError::IndexNoSpace(index.capacity_pow2))
    }

    pub fn addref(&mut self, key: &Pubkey) -> Option<RefCount> {
        let (elem, _) = self.find_entry_mut(key)?;
        elem.ref_count += 1;
        Some(elem.ref_count)
    }

    pub fn unref(&mut self, key: &Pubkey) -> Option<RefCount> {
        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<u64, BucketMapError> {
        Self::bucket_create_key(
            &self.index,
            key,
            IndexEntry::key_uid(key),
            self.random,
            ref_count,
        )
    }

    pub fn read_value(&self, key: &Pubkey) -> Option<(&[T], RefCount)> {
        //debug!("READ_VALUE: {:?}", key);
        let (elem, _) = self.find_entry(key)?;
        elem.read_value(self)
    }

    pub fn try_write(
        &mut self,
        key: &Pubkey,
        data: &[T],
        ref_count: u64,
    ) -> Result<(), BucketMapError> {
        let index_entry = self.find_entry_mut(key);
        let (elem, elem_ix) = match index_entry {
            None => {
                let ii = self.create_key(key, ref_count)?;
                let elem = self.index.get_mut(ii); // get_mut here is right?
                (elem, ii)
            }
            Some(res) => {
                if ref_count != res.0.ref_count {
                    res.0.ref_count = ref_count;
                }
                res
            }
        };
        let elem_uid = self.index.uid(elem_ix);
        let best_fit_bucket = IndexEntry::data_bucket_from_num_slots(data.len() as u64);
        if self.data.get(best_fit_bucket as usize).is_none() {
            //error!("resizing because missing bucket");
            return Err(BucketMapError::DataNoSpace((best_fit_bucket, 0)));
        }
        let bucket_ix = elem.data_bucket_ix();
        let current_bucket = &self.data[bucket_ix as usize];
        if best_fit_bucket == bucket_ix && elem.num_slots > 0 {
            //in place update
            let elem_loc = elem.data_loc(current_bucket);
            let slice: &mut [T] = current_bucket.get_mut_cell_slice(elem_loc, data.len() as u64);
            //let elem: &mut IndexEntry = self.index.get_mut(elem_ix);
            assert!(current_bucket.uid(elem_loc) == elem_uid);
            elem.num_slots = data.len() as u64;
            slice.clone_from_slice(data);
            Ok(())
        } else {
            //need to move the allocation to a best fit spot
            let best_bucket = &self.data[best_fit_bucket as usize];
            let cap_power = best_bucket.capacity_pow2;
            let cap = best_bucket.num_cells();
            let pos = thread_rng().gen_range(0, cap);
            for i in pos..pos + self.index.max_search() {
                let ix = i % cap;
                if best_bucket.uid(ix) == UID_UNLOCKED {
                    let elem_loc = elem.data_loc(current_bucket);
                    if elem.num_slots > 0 {
                        current_bucket.free(elem_loc, elem_uid);
                    }
                    // elem: &mut IndexEntry = self.index.get_mut(elem_ix);
                    elem.storage_offset = ix;
                    elem.storage_capacity_when_created_pow2 = best_bucket.capacity_pow2;
                    elem.num_slots = data.len() as u64;
                    //debug!(                        "DATA ALLOC {:?} {} {} {}",                        key, elem.data_location, best_bucket.capacity, elem_uid                    );
                    if elem.num_slots > 0 {
                        best_bucket.allocate(ix, elem_uid).unwrap();
                        let slice = best_bucket.get_mut_cell_slice(ix, data.len() as u64);
                        slice.copy_from_slice(data);
                    }
                    return Ok(());
                }
            }
            Err(BucketMapError::DataNoSpace((best_fit_bucket, cap_power)))
        }
    }

    pub fn delete_key(&mut self, key: &Pubkey) {
        if let Some((elem, elem_ix)) = self.find_entry(key) {
            let elem_uid = self.index.uid(elem_ix);
            if elem.num_slots > 0 {
                let data_bucket = &self.data[elem.data_bucket_ix() as usize];
                let loc = elem.data_loc(data_bucket);
                //debug!(                    "DATA FREE {:?} {} {} {}",                    key, elem.data_location, data_bucket.capacity, elem_uid                );
                data_bucket.free(loc, elem_uid);
            }
            //debug!("INDEX FREE {:?} {}", key, elem_uid);
            self.index.free(elem_ix, elem_uid);
        }
    }

    pub fn grow_index(&mut self, sz: u8) {
        if self.index.capacity_pow2 == sz {
            let mut m = Measure::start("");
            //debug!("GROW_INDEX: {}", sz);
            let increment = 1;
            for i in increment.. {
                //increasing the capacity by ^4 reduces the
                //likelyhood of a re-index collision of 2^(max_search)^2
                //1 in 2^32
                let index = BucketStorage::new_with_capacity(
                    Arc::clone(&self.drives),
                    1,
                    std::mem::size_of::<IndexEntry>() as u64,
                    self.index.capacity_pow2 + i, // * 2,
                    self.index.max_search,
                    Arc::clone(&self.stats.index),
                );
                let random = thread_rng().gen();
                let mut valid = true;
                for ix in 0..self.index.num_cells() {
                    let uid = self.index.uid(ix);
                    if UID_UNLOCKED != uid {
                        let elem: &IndexEntry = self.index.get(ix);
                        let ref_count = 0; // ??? TODO
                        let new_ix =
                            Self::bucket_create_key(&index, &elem.key, uid, random, ref_count);
                        if new_ix.is_err() {
                            valid = false;
                            break;
                        }
                        let new_ix = new_ix.unwrap();
                        let new_elem: &mut IndexEntry = index.get_mut(new_ix);
                        *new_elem = *elem;
                        /*
                        let dbg_elem: IndexEntry = *new_elem;
                        assert_eq!(
                            Self::bucket_find_entry(&index, &elem.key, random).unwrap(),
                            (&dbg_elem, new_ix)
                        );
                        */
                    }
                }
                if valid {
                    self.index = index;
                    self.random = random;
                    break;
                }
            }
            m.stop();
            let sz = 1 << self.index.capacity_pow2;
            {
                let mut max = self.stats.index.max_size.lock().unwrap();
                *max = std::cmp::max(*max, sz);
            }
            self.stats.index.resizes.fetch_add(1, Ordering::Relaxed);
            self.stats
                .index
                .resize_us
                .fetch_add(m.as_us(), Ordering::Relaxed);
        }
    }

    pub fn grow_data(&mut self, sz: (u64, u8)) {
        if self.data.get(sz.0 as usize).is_none() {
            for i in self.data.len() as u64..(sz.0 + 1) {
                self.data.push(BucketStorage::new(
                    Arc::clone(&self.drives),
                    1 << i,
                    std::mem::size_of::<T>() as u64,
                    self.index.max_search,
                    Arc::clone(&self.stats.data),
                ))
            }
        }
        if self.data[sz.0 as usize].capacity_pow2 == sz.1 {
            //debug!("GROW_DATA: {} {}", sz.0, sz.1);
            self.data[sz.0 as usize].grow();
        }
    }

    fn bucket_index_ix(index: &BucketStorage, key: &Pubkey, random: u64) -> u64 {
        let uid = IndexEntry::key_uid(key);
        let mut s = DefaultHasher::new();
        uid.hash(&mut s);
        //the locally generated random will make it hard for an attacker
        //to deterministically cause all the pubkeys to land in the same
        //location in any bucket on all validators
        random.hash(&mut s);
        let ix = s.finish();
        ix % index.num_cells()
        //debug!(            "INDEX_IX: {:?} uid:{} loc: {} cap:{}",            key,            uid,            location,            index.num_cells()        );
    }

    pub fn insert(&mut self, key: &Pubkey, value: (&[T], RefCount)) {
        let (new, refct) = value;
        loop {
            let rv = self.try_write(key, new, refct);
            match rv {
                Err(BucketMapError::DataNoSpace(sz)) => {
                    //debug!("GROWING SPACE {:?}", sz);
                    self.grow_data(sz);
                    continue;
                }
                Err(BucketMapError::IndexNoSpace(sz)) => {
                    //debug!("GROWING INDEX {}", sz);
                    self.grow_index(sz);
                    continue;
                }
                Ok(()) => return,
            }
        }
    }

    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));
    }
}