fix to calculate state size

Author: kobayurii

database/src/base/rpc_server.rs

@@ -93,6 +93,14 @@ pub trait ReaderDbManager {
         method_name: &str,
     ) -> anyhow::Result<readnode_primitives::QueryData<Vec<u8>>>;
 
+    // Returns the contract code size at the given block height
+    async fn get_contract_code_size(
+        &self,
+        account_id: &near_primitives::types::AccountId,
+        request_block_height: near_primitives::types::BlockHeight,
+        method_name: &str,
+    ) -> anyhow::Result<u64>;
+
     /// Returns the near_primitives::account::AccessKey at the given block height
     async fn get_access_key(
         &self,

database/src/postgres/rpc_server.rs

@@ -357,6 +357,37 @@ impl crate::ReaderDbManager for crate::PostgresDBManager {
         })
     }
 
+    async fn get_contract_code_size(
+        &self,
+        account_id: &near_primitives::types::AccountId,
+        request_block_height: near_primitives::types::BlockHeight,
+        method_name: &str,
+    ) -> anyhow::Result<u64> {
+        let shard_id_pool = self.get_shard_connection(account_id).await?;
+        crate::metrics::SHARD_DATABASE_READ_QUERIES
+            .with_label_values(&[
+                &shard_id_pool.shard_id.to_string(),
+                method_name,
+                "state_changes_contract",
+            ])
+            .inc();
+        let (code_size,): (bigdecimal::BigDecimal,) = sqlx::query_as(
+            "
+                SELECT pg_column_size(data_value)
+                FROM state_changes_contract
+                WHERE account_id = $1 
+                    AND block_height <= $2
+                ORDER BY block_height DESC
+                LIMIT 1;
+                ",
+        )
+        .bind(account_id.to_string())
+        .bind(bigdecimal::BigDecimal::from(request_block_height))
+        .fetch_one(shard_id_pool.pool)
+        .await?;
+        Ok(code_size.to_u64().map(|size| size).unwrap_or_default())
+    }
+
     async fn get_access_key(
         &self,
         account_id: &near_primitives::types::AccountId,

rpc-server/src/config.rs

@@ -153,7 +153,11 @@ impl ServerContext {
 
         let compiled_contract_code_cache =
             std::sync::Arc::new(CompiledCodeCache::new(contract_code_cache_size_in_bytes));
-
+        
+        crate::metrics::CARGO_PKG_VERSION
+            .with_label_values(&[NEARD_VERSION])
+            .inc();
+        
         Ok(Self {
             s3_client,
             db_manager: std::sync::Arc::new(Box::new(db_manager)),

rpc-server/src/metrics.rs

@@ -1,5 +1,5 @@
 use actix_web::{get, Responder};
-use prometheus::{Encoder, IntCounterVec, IntGauge, IntGaugeVec, Opts};
+use prometheus::{CounterVec, Encoder, IntCounterVec, IntGauge, IntGaugeVec, Opts};
 
 type Result<T, E> = std::result::Result<T, E>;
 
@@ -113,6 +113,14 @@ lazy_static! {
         "Optimistic updating status. 0: working, 1: not working",
     ).unwrap();
 
+    pub(crate) static ref CARGO_PKG_VERSION: CounterVec = {
+        let opts = Opts::new("cargo_pkg_version", "Cargo package version. This is used to track the version of the running server.")
+            .variable_label("version");
+        let counter_vec = CounterVec::new(opts, &["version"]).expect("metric can be created");
+        prometheus::register(Box::new(counter_vec.clone())).unwrap();
+        counter_vec
+    };
+
     pub(crate) static ref LEGACY_DATABASE_TX_DETAILS: IntCounterVec = register_int_counter_vec(
         "legacy_database_tx_details",
         "Total number of calls to the legacy database for transaction details",

rpc-server/src/modules/blocks/mod.rs

@@ -335,10 +335,6 @@ impl BlocksInfoByFinality {
     // Update final block info in the cache.
     // Executes every second.
     pub async fn update_final_block(&self, block_info: BlockInfo) {
-        tracing::debug!(
-            "Update final block info: {:?}",
-            block_info.block_cache.block_height
-        );
         let mut final_block_lock = self.final_block.write().await;
         final_block_lock.block_cache = block_info.block_cache;
         final_block_lock.block_view = block_info.block_view;
@@ -348,11 +344,6 @@ impl BlocksInfoByFinality {
     // Update optimistic block changes and optimistic block info in the cache.
     // Executes every second.
     pub async fn update_optimistic_block(&self, block_info: BlockInfo) {
-        tracing::debug!(
-            "Update optimistic block info: {:?}",
-            block_info.block_cache.block_height
-        );
-
         let mut optimistic_changes_lock = self.optimistic_changes.write().await;
         optimistic_changes_lock.account_changes = block_info.changes_in_block_account_map().await;
 

rpc-server/src/modules/queries/contract_runner/mod.rs

@@ -1,9 +1,8 @@
 use std::collections::HashMap;
 
-use near_vm_runner::ContractRuntimeCache;
-
 use crate::modules::blocks::BlocksInfoByFinality;
 use code_storage::CodeStorage;
+use near_vm_runner::ContractRuntimeCache;
 
 mod code_storage;
 
@@ -143,14 +142,30 @@ pub async fn run_contract(
         }
     };
 
+    // We need to calculate the state size of the contract to determine if we should prefetch the state or not.
+    // The state size is the storage usage minus the code size.
+    // If the state size is less than the prefetch_state_size_limit, we prefetch the state.
+    let code_size = if let Some(contract_code) = &contract_code.contract_code {
+        size_of_val(&contract_code.code()) as u64
+    } else {
+        if let Some(code) = contract_code_cache.get(&code_hash).await {
+            size_of_val(&code) as u64
+        } else {
+            db_manager
+                .get_contract_code_size(account_id, block.block_height, "query_call_function")
+                .await
+                .unwrap_or(0)
+        }
+    };
+    let state_size = contract.data.storage_usage() - code_size;
     // Init an external database interface for the Runtime logic
     let code_storage = CodeStorage::init(
         db_manager.clone(),
         account_id.clone(),
         block.block_height,
         validators,
         optimistic_data,
-        contract.data.storage_usage() <= prefetch_state_size_limit,
+        state_size <= prefetch_state_size_limit,
     )
     .await;
 

rpc-server/src/modules/queries/methods.rs

@@ -443,7 +443,17 @@ async fn view_state(
                 block_hash: block.block_hash,
             },
         )?;
-    if prefix.is_empty() && account.data.storage_usage() > data.prefetch_state_size_limit {
+
+    // Calculate the state size excluding the contract code size to check if it's too large to fetch.
+    // The state size is the storage usage minus the code size.
+    let contract_code_size = data
+        .db_manager
+        .get_contract_code_size(account_id, block.block_height, "query_call_function")
+        .await
+        .unwrap_or(0);
+    let state_size = account.data.storage_usage() - contract_code_size;
+    // If the prefix is empty and the state size is larger than the limit, return an error.
+    if prefix.is_empty() && state_size > data.prefetch_state_size_limit {
         return Err(
             near_jsonrpc::primitives::types::query::RpcQueryError::TooLargeContractState {
                 contract_account_id: account_id.clone(),