rpc: add rollback to getLatestBlockhash (agave#2023)

rpcpool · Oct 7, 2024 · 5b25172 · 5b25172
1 parent 9e7b00b
commit 5b25172
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 6 deletions.
diff --git a/rpc-client-api/src/config.rs b/rpc-client-api/src/config.rs
@@ -358,3 +358,12 @@ pub struct RpcContextConfig {
 pub struct RpcRecentPrioritizationFeesConfig {
     pub percentile: Option<u16>,
 }
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct RpcLatestBlockhashConfig {
+    #[serde(flatten)]
+    pub context: RpcContextConfig,
+    #[serde(default)]
+    pub rollback: usize,
+}
diff --git a/rpc-client/src/nonblocking/rpc_client.rs b/rpc-client/src/nonblocking/rpc_client.rs
@@ -5290,6 +5290,26 @@ impl RpcClient {
         Ok(blockhash)
     }
 
+    pub async fn get_latest_blockhash_with_config(
+        &self,
+        config: RpcLatestBlockhashConfig,
+    ) -> ClientResult<(Hash, u64)> {
+        let RpcBlockhash {
+            blockhash,
+            last_valid_block_height,
+        } = self
+            .send::<Response<RpcBlockhash>>(RpcRequest::GetLatestBlockhash, json!([config]))
+            .await?
+            .value;
+        let blockhash = blockhash.parse().map_err(|_| {
+            ClientError::new_with_request(
+                RpcError::ParseError("Hash".to_string()).into(),
+                RpcRequest::GetLatestBlockhash,
+            )
+        })?;
+        Ok((blockhash, last_valid_block_height))
+    }
+
     #[allow(deprecated)]
     pub async fn get_latest_blockhash_with_commitment(
         &self,

diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
@@ -2149,8 +2149,28 @@ impl JsonRpcRequestProcessor {
         }
     }
 
-    fn get_latest_blockhash(&self, config: RpcContextConfig) -> Result<RpcResponse<RpcBlockhash>> {
-        let bank = self.get_bank_with_config(config)?;
+    fn get_latest_blockhash(
+        &self,
+        config: RpcLatestBlockhashConfig,
+    ) -> Result<RpcResponse<RpcBlockhash>> {
+        let mut bank = self.get_bank_with_config(config.context)?;
+        if config.rollback > MAX_RECENT_BLOCKHASHES {
+            return Err(Error::invalid_params("rollback exceeds 300"));
+        }
+        if config.rollback > 0 {
+            let r_bank_forks = self.bank_forks.read().unwrap();
+            for _ in 0..config.rollback {
+                bank = match r_bank_forks.get(bank.parent_slot()).or_else(|| {
+                    r_bank_forks
+                        .banks_frozen
+                        .get(&bank.parent_slot())
+                        .map(Clone::clone)
+                }) {
+                    Some(bank) => bank,
+                    None => return Err(Error::invalid_params("failed to rollback block")),
+                };
+            }
+        }
         let blockhash = bank.last_blockhash();
         let last_valid_block_height = bank
             .get_blockhash_last_valid_block_height(&blockhash)
@@ -3397,7 +3417,7 @@ pub mod rpc_full {
         fn get_latest_blockhash(
             &self,
             meta: Self::Metadata,
-            config: Option<RpcContextConfig>,
+            config: Option<RpcLatestBlockhashConfig>,
         ) -> Result<RpcResponse<RpcBlockhash>>;
 
         #[rpc(meta, name = "isBlockhashValid")]
@@ -4059,7 +4079,7 @@ pub mod rpc_full {
         fn get_latest_blockhash(
             &self,
             meta: Self::Metadata,
-            config: Option<RpcContextConfig>,
+            config: Option<RpcLatestBlockhashConfig>,
         ) -> Result<RpcResponse<RpcBlockhash>> {
             debug!("get_latest_blockhash rpc request received");
             meta.get_latest_blockhash(config.unwrap_or_default())

diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs
@@ -13,12 +13,12 @@ use {
     solana_measure::measure::Measure,
     solana_program_runtime::loaded_programs::{BlockRelation, ForkGraph},
     solana_sdk::{
-        clock::{Epoch, Slot},
+        clock::{Epoch, Slot, MAX_RECENT_BLOCKHASHES},
         hash::Hash,
         timing,
     },
     std::{
-        collections::{hash_map::Entry, HashMap, HashSet},
+        collections::{hash_map::Entry, BTreeMap, HashMap, HashSet},
         ops::Index,
         sync::{
             atomic::{AtomicBool, AtomicU64, Ordering},
@@ -64,6 +64,7 @@ struct SetRootTimings {
 
 pub struct BankForks {
     banks: HashMap<Slot, BankWithScheduler>,
+    pub banks_frozen: BTreeMap<Slot, Arc<Bank>>,
     descendants: HashMap<Slot, HashSet<Slot>>,
     root: Arc<AtomicSlot>,
 
@@ -117,6 +118,7 @@ impl BankForks {
         let bank_forks = Arc::new(RwLock::new(Self {
             root: Arc::new(AtomicSlot::new(root_slot)),
             banks,
+            banks_frozen: Default::default(),
             descendants,
             snapshot_config: None,
             accounts_hash_interval_slots: std::u64::MAX,
@@ -249,6 +251,13 @@ impl BankForks {
 
     pub fn remove(&mut self, slot: Slot) -> Option<Arc<Bank>> {
         let bank = self.banks.remove(&slot)?;
+        if bank.is_frozen() {
+            self.banks_frozen
+                .insert(bank.slot(), bank.clone_without_scheduler());
+            while self.banks_frozen.len() > MAX_RECENT_BLOCKHASHES {
+                self.banks_frozen.pop_first();
+            }
+        }
         for parent in bank.proper_ancestors() {
             let Entry::Occupied(mut entry) = self.descendants.entry(parent) else {
                 panic!("this should not happen!");