diff --git a/bucket_map/src/bucket_api.rs b/bucket_map/src/bucket_api.rs index 2ec495ab2..d81b4b52f 100644 --- a/bucket_map/src/bucket_api.rs +++ b/bucket_map/src/bucket_api.rs @@ -1,7 +1,7 @@ use { crate::{ bucket::Bucket, bucket_item::BucketItem, bucket_map::BucketMapError, - bucket_stats::BucketMapStats, MaxSearch, RefCount, + bucket_stats::BucketMapStats, restart::RestartableBucket, MaxSearch, RefCount, }, solana_sdk::pubkey::Pubkey, std::{ @@ -23,6 +23,11 @@ pub struct BucketApi { bucket: LockedBucket, count: Arc, + + /// keeps track of which index file this bucket is currently using + /// or at startup, which bucket file this bucket should initially use + #[allow(dead_code)] + restartable_bucket: RestartableBucket, } impl BucketApi { @@ -30,6 +35,7 @@ impl BucketApi { drives: Arc>, max_search: MaxSearch, stats: Arc, + restartable_bucket: RestartableBucket, ) -> Self { Self { drives, @@ -37,6 +43,7 @@ impl BucketApi { stats, bucket: RwLock::default(), count: Arc::default(), + restartable_bucket, } } diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs index 2adeefdcc..4f6e177c5 100644 --- a/bucket_map/src/bucket_map.rs +++ b/bucket_map/src/bucket_map.rs @@ -1,9 +1,17 @@ //! BucketMap is a mostly contention free concurrent map backed by MmapMut use { - crate::{bucket_api::BucketApi, bucket_stats::BucketMapStats, MaxSearch, RefCount}, + crate::{ + bucket_api::BucketApi, bucket_stats::BucketMapStats, restart::Restart, MaxSearch, RefCount, + }, solana_sdk::pubkey::Pubkey, - std::{convert::TryInto, fmt::Debug, fs, path::PathBuf, sync::Arc}, + std::{ + convert::TryInto, + fmt::Debug, + fs::{self}, + path::PathBuf, + sync::{Arc, Mutex}, + }, tempfile::TempDir, }; @@ -83,6 +91,11 @@ impl BucketMap { if let Some(drives) = config.drives.as_ref() { Self::erase_previous_drives(drives); } + + let stats = Arc::default(); + + let restart = Restart::new(&config); + let mut temp_dir = None; let drives = config.drives.unwrap_or_else(|| { temp_dir = Some(TempDir::new().unwrap()); @@ -90,13 +103,19 @@ impl BucketMap { }); let drives = Arc::new(drives); - let stats = Arc::default(); - let buckets = (0..config.max_buckets) - .map(|_| { + let restart = restart.map(|restart| Arc::new(Mutex::new(restart))); + + let restartable_buckets = + Restart::get_restartable_buckets(restart.as_ref(), &drives, config.max_buckets); + + let buckets = restartable_buckets + .into_iter() + .map(|restartable_bucket| { Arc::new(BucketApi::new( Arc::clone(&drives), max_search, Arc::clone(&stats), + restartable_bucket, )) }) .collect();