From 4ea66291df5490e5dc90655901c36253732dc72b Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Wed, 17 Jan 2024 16:07:52 -0800
Subject: [PATCH] feature-gate

---
 cost-model/src/cost_model.rs | 37 +++++++++++++++++++++++++++---------
 sdk/src/feature_set.rs       |  5 +++++
 2 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/cost-model/src/cost_model.rs b/cost-model/src/cost_model.rs
index 3efc5d94148293..1e15735426737f 100644
--- a/cost-model/src/cost_model.rs
+++ b/cost-model/src/cost_model.rs
@@ -18,7 +18,7 @@ use {
     solana_sdk::{
         borsh1::try_from_slice_unchecked,
         compute_budget::{self, ComputeBudgetInstruction},
-        feature_set::{include_loaded_accounts_data_size_in_fee_calculation, FeatureSet},
+        feature_set::{self, include_loaded_accounts_data_size_in_fee_calculation, FeatureSet},
         fee::FeeStructure,
         instruction::CompiledInstruction,
         program_utils::limited_deserialize,
@@ -44,7 +44,7 @@ impl CostModel {
             let mut tx_cost = UsageCostDetails::new_with_default_capacity();
 
             tx_cost.signature_cost = Self::get_signature_cost(transaction);
-            Self::get_write_lock_cost(&mut tx_cost, transaction);
+            Self::get_write_lock_cost(&mut tx_cost, transaction, feature_set);
             Self::get_transaction_cost(&mut tx_cost, transaction, feature_set);
             tx_cost.account_data_size = Self::calculate_account_data_size(transaction);
 
@@ -73,10 +73,19 @@ impl CostModel {
             .collect()
     }
 
-    fn get_write_lock_cost(tx_cost: &mut UsageCostDetails, transaction: &SanitizedTransaction) {
+    fn get_write_lock_cost(
+        tx_cost: &mut UsageCostDetails,
+        transaction: &SanitizedTransaction,
+        feature_set: &FeatureSet,
+    ) {
         tx_cost.writable_accounts = Self::get_writable_accounts(transaction);
-        tx_cost.write_lock_cost =
-            WRITE_LOCK_UNITS.saturating_mul(transaction.message().num_write_locks());
+        let num_write_locks =
+            if feature_set.is_active(&feature_set::cost_model_requested_write_lock_cost::id()) {
+                transaction.message().num_write_locks()
+            } else {
+                tx_cost.writable_accounts.len() as u64
+            };
+        tx_cost.write_lock_cost = WRITE_LOCK_UNITS.saturating_mul(num_write_locks);
     }
 
     fn get_transaction_cost(
@@ -339,10 +348,20 @@ mod tests {
             system_transaction::transfer(&mint_keypair, &system_program::id(), 2, start_hash),
         );
 
-        // Cost should be fore 2 write-locks, but only 1 is actually writable.
-        let tx_cost = CostModel::calculate_cost(&simple_transaction, &FeatureSet::all_enabled());
-        assert_eq!(2 * WRITE_LOCK_UNITS, tx_cost.write_lock_cost());
-        assert_eq!(1, tx_cost.writable_accounts().len());
+        // Feature not enabled - write lock is demoted and does not count towards cost
+        {
+            let tx_cost = CostModel::calculate_cost(&simple_transaction, &FeatureSet::default());
+            assert_eq!(WRITE_LOCK_UNITS, tx_cost.write_lock_cost());
+            assert_eq!(1, tx_cost.writable_accounts().len());
+        }
+
+        // Feature enabled - write lock is demoted but still counts towards cost
+        {
+            let tx_cost =
+                CostModel::calculate_cost(&simple_transaction, &FeatureSet::all_enabled());
+            assert_eq!(2 * WRITE_LOCK_UNITS, tx_cost.write_lock_cost());
+            assert_eq!(1, tx_cost.writable_accounts().len());
+        }
     }
 
     #[test]
diff --git a/sdk/src/feature_set.rs b/sdk/src/feature_set.rs
index f2e9c63ff1b2c9..2941c94ae81cb3 100644
--- a/sdk/src/feature_set.rs
+++ b/sdk/src/feature_set.rs
@@ -768,6 +768,10 @@ pub mod enable_zk_proof_from_account {
     solana_sdk::declare_id!("zkiTNuzBKxrCLMKehzuQeKZyLtX2yvFcEKMML8nExU8");
 }
 
+pub mod cost_model_requested_write_lock_cost {
+    solana_sdk::declare_id!("wLckV1a64ngtcKPRGU4S4grVTestXjmNjxBjaKZrAcn");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -955,6 +959,7 @@ lazy_static! {
         (deprecate_executable_meta_update_in_bpf_loader::id(), "deprecate executable meta flag update in bpf loader #34194"),
         (enable_zk_proof_from_account::id(), "Enable zk token proof program to read proof from accounts instead of instruction data #34750"),
         (curve25519_restrict_msm_length::id(), "restrict curve25519 multiscalar multiplication vector lengths #34763"),
+        (cost_model_requested_write_lock_cost::id(), "cost model uses number of requested write locks #34819"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()