use { crate::{ bucket::Bucket, bucket_item::BucketItem, bucket_map::BucketMapError, bucket_stats::BucketMapStats, MaxSearch, RefCount, }, solana_sdk::pubkey::Pubkey, std::{ ops::RangeBounds, path::PathBuf, sync::{ atomic::{AtomicU64, Ordering}, Arc, RwLock, RwLockWriteGuard, }, }, }; type LockedBucket = RwLock>>; pub struct BucketApi { drives: Arc>, max_search: MaxSearch, pub stats: Arc, bucket: LockedBucket, count: Arc, } impl BucketApi { pub fn new( drives: Arc>, max_search: MaxSearch, stats: Arc, count: Arc, ) -> Self { Self { drives, max_search, stats, bucket: RwLock::default(), count, } } /// Get the items for bucket pub fn items_in_range(&self, range: &Option<&R>) -> Vec> where R: RangeBounds, { self.bucket .read() .unwrap() .as_ref() .map(|bucket| bucket.items_in_range(range)) .unwrap_or_default() } /// Get the Pubkeys pub fn keys(&self) -> Vec { self.bucket .read() .unwrap() .as_ref() .map_or_else(Vec::default, |bucket| bucket.keys()) } /// Get the values for Pubkey `key` pub fn read_value(&self, key: &Pubkey) -> Option<(Vec, RefCount)> { self.bucket.read().unwrap().as_ref().and_then(|bucket| { bucket .read_value(key) .map(|(value, ref_count)| (value.to_vec(), ref_count)) }) } pub fn bucket_len(&self) -> u64 { self.bucket .read() .unwrap() .as_ref() .map(|bucket| bucket.bucket_len()) .unwrap_or_default() } pub fn delete_key(&self, key: &Pubkey) { let mut bucket = self.get_write_bucket(); if let Some(bucket) = bucket.as_mut() { bucket.delete_key(key) } } fn get_write_bucket(&self) -> RwLockWriteGuard>> { 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), )); } else { let write = bucket.as_mut().unwrap(); write.handle_delayed_grows(); self.count.store(write.bucket_len(), Ordering::Relaxed); } bucket } pub fn addref(&self, key: &Pubkey) -> Option { self.get_write_bucket() .as_mut() .and_then(|bucket| bucket.addref(key)) } pub fn unref(&self, key: &Pubkey) -> Option { self.get_write_bucket() .as_mut() .and_then(|bucket| bucket.unref(key)) } pub fn insert(&self, pubkey: &Pubkey, value: (&[T], RefCount)) { let mut bucket = self.get_write_bucket(); bucket.as_mut().unwrap().insert(pubkey, value) } pub fn grow(&self, err: BucketMapError) { // grows are special - they get a read lock and modify 'reallocated' // the grown changes are applied the next time there is a write lock taken if let Some(bucket) = self.bucket.read().unwrap().as_ref() { bucket.grow(err) } } pub fn update(&self, key: &Pubkey, updatefn: F) where F: FnMut(Option<(&[T], RefCount)>) -> Option<(Vec, RefCount)>, { let mut bucket = self.get_write_bucket(); bucket.as_mut().unwrap().update(key, updatefn) } pub fn try_write( &self, pubkey: &Pubkey, value: (&[T], RefCount), ) -> Result<(), BucketMapError> { let mut bucket = self.get_write_bucket(); bucket.as_mut().unwrap().try_write(pubkey, value.0, value.1) } }