From ecf2f7b0e9d449a85423eec98772fc23f66cf58c Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Fri, 5 Jul 2024 23:59:33 -0500
Subject: [PATCH] rpc: add rollback to getLatestBlockhash (agave#2023)

---
 rpc-client-api/src/config.rs             |  9 ++++++++
 rpc-client/src/nonblocking/rpc_client.rs | 21 ++++++++++++++++++
 rpc/src/rpc.rs                           | 28 ++++++++++++++++++++----
 runtime/src/bank_forks.rs                | 13 +++++++++--
 4 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/rpc-client-api/src/config.rs b/rpc-client-api/src/config.rs
index 50176d309c..5d49f75bb9 100644
--- a/rpc-client-api/src/config.rs
+++ b/rpc-client-api/src/config.rs
@@ -360,3 +360,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
index 7543b4f926..9ca28337fe 100644
--- a/rpc-client/src/nonblocking/rpc_client.rs
+++ b/rpc-client/src/nonblocking/rpc_client.rs
@@ -4578,6 +4578,27 @@ 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,
         commitment: CommitmentConfig,
diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
index dbf69149ea..88053e3ed9 100644
--- a/rpc/src/rpc.rs
+++ b/rpc/src/rpc.rs
@@ -2190,8 +2190,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)
@@ -3428,7 +3448,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")]
@@ -4104,7 +4124,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
index 3dd82c1fe8..a6af2151ca 100644
--- a/runtime/src/bank_forks.rs
+++ b/runtime/src/bank_forks.rs
@@ -14,12 +14,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},
@@ -72,6 +72,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>,
 
@@ -125,6 +126,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: u64::MAX,
@@ -255,6 +257,13 @@ impl BankForks {
 
     pub fn remove(&mut self, slot: Slot) -> Option<BankWithScheduler> {
         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!");