Enhancing Blockstore Functionality

Cargo.toml

@@ -28,6 +28,7 @@ chrono = "0.4.24"
 rand = "0.8.5"
 crossbeam-channel = "0.5.13"
 tempfile = "3.3.0"
+assert_matches = "1.5.0"
 
 solana-bpf-loader-program = { version = "2.0.6" }
 solana-compute-budget = { version = "2.0.6" }

storage/Cargo.toml

@@ -14,6 +14,7 @@ log = { workspace = true }
 tempfile = { workspace = true }
 crossbeam-channel = { workspace = true }
 rand = { workspace = true }
+assert_matches = { workspace = true }
 
 solana-ledger = { workspace = true }
 solana-accounts-db = { workspace = true }

storage/src/accounts/tests.rs

@@ -33,6 +33,7 @@ async fn init_from_snapshot_works() -> Result<()> {
         .storage
         .snapshot_config
         .full_snapshot_archive_interval_slots = 1; // set snapshot interval to 1
+    config.storage.wait_snapshot_complete = true;
     let mut store = RollupStorage::new(config)?;
     store.init()?;
     let keypairs = store.config.keypairs.clone();
@@ -95,8 +96,6 @@ async fn init_from_snapshot_works() -> Result<()> {
 
     // 3. save and close
     store.force_save().await?;
-    // TODO: sleep is needed here, improve later
-    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
     store.close().await?;
 
     // 4. open with snapshot

storage/src/background.rs

@@ -1,4 +1,7 @@
-use std::sync::{atomic::AtomicBool, Arc, RwLock};
+use std::{
+    sync::{atomic::AtomicBool, Arc, RwLock},
+    time::Duration,
+};
 
 use crossbeam_channel::unbounded;
 use solana_core::{
@@ -12,10 +15,13 @@ use solana_runtime::{
         SnapshotRequestHandler,
     },
     bank_forks::BankForks,
+    snapshot_config::SnapshotConfig,
     snapshot_hash::StartingSnapshotHashes,
+    snapshot_utils,
 };
+use tokio::time::sleep;
 
-use crate::{config::StorageConfig, sig_hub::SignalHub, Error, Result};
+use crate::{config::StorageConfig, sig_hub::SignalHub, Error, Result, RollupStorage};
 
 #[allow(dead_code)]
 pub struct StorageBackground {
@@ -25,6 +31,68 @@ pub struct StorageBackground {
     pub(crate) snapshot_packager_service: SnapshotPackagerService,
 }
 
+impl RollupStorage {
+    pub async fn try_wait_snapshot_complete(&self) {
+        let storage_cfg = &self.config.storage;
+        if !storage_cfg.snapshot_config.should_generate_snapshots()
+            || !storage_cfg.wait_snapshot_complete
+        {
+            return;
+        }
+
+        let full_latest = self.wait_full_snapshot(&storage_cfg.snapshot_config).await;
+        if let Some(full_latest) = full_latest {
+            self.wait_incremental_snapshot(&storage_cfg.snapshot_config, full_latest)
+                .await;
+        }
+    }
+
+    async fn wait_full_snapshot(&self, snapshot_config: &SnapshotConfig) -> Option<u64> {
+        let expect_slot = self.get_latest_snapshot(
+            self.current_height(),
+            snapshot_config.full_snapshot_archive_interval_slots,
+        );
+        if let Some(slot) = expect_slot {
+            while snapshot_utils::get_highest_full_snapshot_archive_slot(
+                &snapshot_config.full_snapshot_archives_dir,
+            ) != Some(slot)
+            {
+                debug!("Waiting for full snapshot {slot}");
+                sleep(Duration::from_millis(20)).await;
+            }
+        }
+
+        expect_slot
+    }
+
+    async fn wait_incremental_snapshot(&self, snapshot_config: &SnapshotConfig, full_latest: u64) {
+        let expect_slot = self.get_latest_snapshot(
+            self.current_height(),
+            snapshot_config.incremental_snapshot_archive_interval_slots,
+        );
+
+        if let Some(slot) = expect_slot {
+            while snapshot_utils::get_highest_incremental_snapshot_archive_slot(
+                &snapshot_config.incremental_snapshot_archives_dir,
+                full_latest,
+            ) != Some(slot)
+            {
+                debug!("Waiting for incremental snapshot {slot}");
+                sleep(Duration::from_millis(20)).await;
+            }
+        }
+    }
+
+    fn get_latest_snapshot(&self, highest: u64, inteval: u64) -> Option<u64> {
+        let reminder = highest % inteval;
+        if reminder >= highest {
+            None
+        } else {
+            Some(highest - reminder)
+        }
+    }
+}
+
 impl StorageBackground {
     pub fn new(
         bank_forks: Arc<RwLock<BankForks>>,

storage/src/config.rs

@@ -2,6 +2,7 @@ use std::{
     collections::HashSet,
     path::{Path, PathBuf},
     sync::Arc,
+    time::Duration,
 };
 
 use rollups_interface::l2::executor::Config;
@@ -86,6 +87,8 @@ pub struct StorageConfig {
     pub expected_genesis_hash: Option<Hash>,
     pub account_paths: Vec<PathBuf>,
     pub snapshot_config: SnapshotConfig,
+    pub wait_snapshot_complete: bool,
+    pub wait_timeout: Duration,
     pub enforce_ulimit_nofile: bool,
     pub fixed_leader_schedule: Option<FixedSchedule>,
     pub new_hard_forks: Option<Vec<Slot>>,
@@ -112,6 +115,8 @@ impl Default for StorageConfig {
             expected_genesis_hash: None,
             account_paths: Vec::new(),
             snapshot_config: SnapshotConfig::new_load_only(),
+            wait_snapshot_complete: false,
+            wait_timeout: Duration::from_secs(10),
             enforce_ulimit_nofile: true,
             fixed_leader_schedule: None,
             new_hard_forks: None,

storage/src/impls.rs

@@ -193,6 +193,7 @@ impl StorageOperations for RollupStorage {
     }
 
     async fn close(self) -> Result<(), Self::Error> {
+        self.try_wait_snapshot_complete().await;
         self.exit.store(true, std::sync::atomic::Ordering::Relaxed);
         self.blockstore.drop_signal();
         self.join();

storage/src/ledger/mod.rs

@@ -0,0 +1,116 @@
+use std::{
+    collections::HashSet,
+    sync::{Arc, RwLock},
+};
+
+use rollups_interface::l2::bank;
+use solana_runtime::{
+    bank::Bank, bank_forks::BankForks, installed_scheduler_pool::BankWithScheduler,
+};
+use solana_sdk::{account::ReadableAccount, pubkey::Pubkey};
+
+use crate::{error::BankError, Result, RollupStorage};
+
+#[cfg(test)]
+mod tests;
+
+impl RollupStorage {
+    pub fn get_mixed_heights(&self) -> Result<(u64, Option<u64>)> {
+        let bank_height = self.bank_forks.read().unwrap().highest_slot();
+        let store_height = self.blockstore.highest_slot()?;
+        Ok((bank_height, store_height))
+    }
+
+    pub fn current_height(&self) -> u64 {
+        self.bank.slot()
+    }
+
+    pub fn balance(&self, pubkey: &Pubkey) -> u64 {
+        self.bank
+            .get_account(pubkey)
+            .map(|d| d.lamports())
+            .unwrap_or(0)
+    }
+
+    pub fn reorg(&mut self, slot: u64, finalized: Option<u64>) -> Result<Vec<BankWithScheduler>> {
+        let removed = self.set_root(slot, finalized)?;
+        let ancestors = find_ancestors(slot, finalized, self.bank_forks.clone(), &removed);
+
+        removed
+            .iter()
+            .filter(|b| !ancestors.contains(&b.slot()))
+            .for_each(|bank| {
+                if let Err(e) = self.blockstore.set_dead_slot(bank.slot()) {
+                    error!("set dead slot failed: {}", e.to_string());
+                }
+            });
+
+        let bank_forks = self.bank_forks.read().unwrap();
+        self.bank = bank_forks.working_bank();
+
+        Ok(removed)
+    }
+
+    pub fn set_root(
+        &mut self,
+        slot: u64,
+        finalized: Option<u64>,
+    ) -> Result<Vec<BankWithScheduler>> {
+        self.blockstore
+            .set_roots(std::iter::once(&slot))
+            .map_err(|e| BankError::SetRootFailed(e.to_string()))?;
+        let removed_banks = self
+            .bank_forks
+            .write()
+            .unwrap()
+            .set_root(
+                slot,
+                &self.background_service.accounts_background_request_sender,
+                finalized,
+            )
+            .map_err(|e| BankError::SetRootFailed(e.to_string()))?;
+        Ok(removed_banks)
+    }
+}
+
+pub(crate) fn find_ancestors(
+    mut root_slot: u64,
+    finalized: Option<u64>,
+    bank_forks: Arc<RwLock<BankForks>>,
+    removed: &[BankWithScheduler],
+) -> HashSet<u64> {
+    let forks = bank_forks.read().unwrap();
+    let stop_at = finalized.unwrap_or(root_slot);
+    while let Some(b) = forks.get(root_slot) {
+        if b.parent_slot() >= stop_at {
+            root_slot = b.parent_slot();
+        } else {
+            break;
+        }
+    }
+    let mut root_parent = if let Some(b) = forks.get(root_slot) {
+        b.parent_slot()
+    } else {
+        return Default::default();
+    };
+    drop(forks);
+
+    let mut ancestors = HashSet::new();
+    loop {
+        let mut found = false;
+        for b in removed {
+            let slot = b.slot();
+            if !ancestors.contains(&slot) && slot == root_parent {
+                ancestors.insert(slot);
+                root_parent = b.parent_slot();
+                found = true;
+                break;
+            }
+        }
+
+        if !found {
+            break;
+        }
+    }
+    ancestors
+}