From 317b1968ecd8dd15b4abfa0d64b5c674b8f04071 Mon Sep 17 00:00:00 2001
From: Tao Zhu <82401714+taozhu-chicago@users.noreply.github.com>
Date: Wed, 7 Jul 2021 23:44:51 -0500
Subject: [PATCH] update ledger tool to restore cost table from blockstore
 (#18489)

* update ledger tool to restore cost model from blockstore when compute-slot-cost

* Move initialize_cost_table into cost_model, so the function can be tested and shared between validator and ledger-tool

* refactor and simplify a test
---
 core/src/cost_model.rs  | 42 +++++++++++++++++++++++++++++++++++++++++
 core/src/validator.rs   | 35 ++++------------------------------
 ledger-tool/src/main.rs |  9 +++++----
 3 files changed, 51 insertions(+), 35 deletions(-)

diff --git a/core/src/cost_model.rs b/core/src/cost_model.rs
index d12a6fc4804261..0a54f5e085a7f9 100644
--- a/core/src/cost_model.rs
+++ b/core/src/cost_model.rs
@@ -91,6 +91,29 @@ impl CostModel {
         self.block_cost_limit
     }
 
+    pub fn initialize_cost_table(&mut self, cost_table: &[(Pubkey, u64)]) {
+        for (program_id, cost) in cost_table {
+            match self.upsert_instruction_cost(program_id, cost) {
+                Ok(c) => {
+                    debug!(
+                        "initiating cost table, instruction {:?} has cost {}",
+                        program_id, c
+                    );
+                }
+                Err(err) => {
+                    debug!(
+                        "initiating cost table, failed for instruction {:?}, err: {}",
+                        program_id, err
+                    );
+                }
+            }
+        }
+        debug!(
+            "restored cost model instruction cost table from blockstore, current values: {:?}",
+            self.get_instruction_cost_table()
+        );
+    }
+
     pub fn calculate_cost(&mut self, transaction: &Transaction) -> &TransactionCost {
         self.transaction_cost.reset();
 
@@ -450,4 +473,23 @@ mod tests {
             th.join().unwrap();
         }
     }
+
+    #[test]
+    fn test_cost_model_init_cost_table() {
+        // build cost table
+        let cost_table = vec![
+            (Pubkey::new_unique(), 10),
+            (Pubkey::new_unique(), 20),
+            (Pubkey::new_unique(), 30),
+        ];
+
+        // init cost model
+        let mut cost_model = CostModel::default();
+        cost_model.initialize_cost_table(&cost_table);
+
+        // verify
+        for (id, cost) in cost_table.iter() {
+            assert_eq!(*cost, cost_model.find_instruction_cost(id));
+        }
+    }
 }
diff --git a/core/src/validator.rs b/core/src/validator.rs
index f449bb732e3174..ccc74a9d7ff419 100644
--- a/core/src/validator.rs
+++ b/core/src/validator.rs
@@ -7,7 +7,7 @@ use {
         cluster_info_vote_listener::VoteTracker,
         completed_data_sets_service::CompletedDataSetsService,
         consensus::{reconcile_blockstore_roots_with_tower, Tower},
-        cost_model::{CostModel, ACCOUNT_MAX_COST, BLOCK_MAX_COST},
+        cost_model::CostModel,
         rewards_recorder_service::{RewardsRecorderSender, RewardsRecorderService},
         sample_performance_service::SamplePerformanceService,
         serve_repair::ServeRepair,
@@ -682,11 +682,9 @@ impl Validator {
             bank_forks.read().unwrap().root_bank().deref(),
         ));
 
-        let cost_model = Arc::new(RwLock::new(CostModel::new(
-            ACCOUNT_MAX_COST,
-            BLOCK_MAX_COST,
-        )));
-        Self::initiate_cost_model(&cost_model, &blockstore.read_program_costs().unwrap());
+        let mut cost_model = CostModel::default();
+        cost_model.initialize_cost_table(&blockstore.read_program_costs().unwrap());
+        let cost_model = Arc::new(RwLock::new(cost_model));
 
         let (retransmit_slots_sender, retransmit_slots_receiver) = unbounded();
         let (verified_vote_sender, verified_vote_receiver) = unbounded();
@@ -919,31 +917,6 @@ impl Validator {
             ip_echo_server.shutdown_background();
         }
     }
-
-    fn initiate_cost_model(cost_model: &RwLock<CostModel>, cost_table: &[(Pubkey, u64)]) {
-        let mut cost_model_mutable = cost_model.write().unwrap();
-        for (program_id, cost) in cost_table {
-            match cost_model_mutable.upsert_instruction_cost(program_id, cost) {
-                Ok(c) => {
-                    debug!(
-                        "initiating cost table, instruction {:?} has cost {}",
-                        program_id, c
-                    );
-                }
-                Err(err) => {
-                    debug!(
-                        "initiating cost table, failed for instruction {:?}, err: {}",
-                        program_id, err
-                    );
-                }
-            }
-        }
-        drop(cost_model_mutable);
-        debug!(
-            "restored cost model instruction cost table from blockstore, current values: {:?}",
-            cost_model.read().unwrap().get_instruction_cost_table()
-        );
-    }
 }
 
 fn active_vote_account_exists_in_bank(bank: &Arc<Bank>, vote_account: &Pubkey) -> bool {
diff --git a/ledger-tool/src/main.rs b/ledger-tool/src/main.rs
index 9c594eb4e9b048..6d406b096ebce1 100644
--- a/ledger-tool/src/main.rs
+++ b/ledger-tool/src/main.rs
@@ -15,6 +15,8 @@ use solana_clap_utils::{
         is_parsable, is_pubkey, is_pubkey_or_keypair, is_slot, is_valid_percentage,
     },
 };
+use solana_core::cost_model::CostModel;
+use solana_core::cost_tracker::CostTracker;
 use solana_ledger::entry::Entry;
 use solana_ledger::{
     ancestor_iterator::AncestorIterator,
@@ -739,10 +741,9 @@ fn compute_slot_cost(blockstore: &Blockstore, slot: Slot) -> Result<(), String>
     let mut transactions = 0;
     let mut programs = 0;
     let mut program_ids = HashMap::new();
-    let cost_model = Arc::new(RwLock::new(CostModel::new(
-        ACCOUNT_MAX_COST,
-        BLOCK_MAX_COST,
-    )));
+    let mut cost_model = CostModel::default();
+    cost_model.initialize_cost_table(&blockstore.read_program_costs().unwrap());
+    let cost_model = Arc::new(RwLock::new(cost_model));
     let mut cost_tracker = CostTracker::new(cost_model.clone());
 
     for entry in &entries {