fix calculate state size limit

near · Nov 6, 2024 · 31d837e · 31d837e
1 parent 8f1f4f7
commit 31d837e
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 49 deletions.
diff --git a/database/src/base/rpc_server.rs b/database/src/base/rpc_server.rs
@@ -93,14 +93,6 @@ 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,

diff --git a/database/src/postgres/rpc_server.rs b/database/src/postgres/rpc_server.rs
@@ -357,37 +357,6 @@ 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,): (i32,) = sqlx::query_as(
-            "
-                SELECT pg_column_size(data_value) as code_size
-                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().unwrap_or_default())
-    }
-
     async fn get_access_key(
         &self,
         account_id: &near_primitives::types::AccountId,

diff --git a/rpc-server/src/modules/queries/contract_runner/mod.rs b/rpc-server/src/modules/queries/contract_runner/mod.rs
@@ -136,6 +136,7 @@ pub async fn run_contract(
                             block_hash: block.block_hash,
                         }
                     })?;
+                println!("Contract code len {}", code.data.len());
                 contract_code_cache.put(code_hash, code.data.clone()).await;
                 Contract::new(Some(code.data), code_hash)
             }
@@ -145,17 +146,21 @@ 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
+    let code_len = if let Some(contract_code) = &contract_code.contract_code {
+        contract_code.code().len()
     } else if let Some(code) = contract_code_cache.get(&code_hash).await {
-        size_of_val(&code) as u64
+        code.len()
     } else {
         db_manager
-            .get_contract_code_size(account_id, block.block_height, "query_call_function")
+            .get_contract_code(account_id, block.block_height, "query_call_function")
             .await
-            .unwrap_or(0)
+            .map(|code| code.data.len())
+            .unwrap_or_default()
     };
-    let state_size = contract.data.storage_usage() - code_size;
+    let state_size = contract
+        .data
+        .storage_usage()
+        .saturating_sub(code_len as u64);
     // Init an external database interface for the Runtime logic
     let code_storage = CodeStorage::init(
         db_manager.clone(),

diff --git a/rpc-server/src/modules/queries/methods.rs b/rpc-server/src/modules/queries/methods.rs
@@ -446,12 +446,14 @@ async fn view_state(
 
     // 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
+    // more details: nearcore/runtime/runtime/src/state_viewer/mod.rs:150
+    let code_len = data
         .db_manager
-        .get_contract_code_size(account_id, block.block_height, "query_call_function")
+        .get_contract_code(account_id, block.block_height, "query_view_state")
         .await
-        .unwrap_or(0);
-    let state_size = account.data.storage_usage() - contract_code_size;
+        .map(|code| code.data.len() as u64)
+        .unwrap_or_default();
+    let state_size = account.data.storage_usage().saturating_sub(code_len);
     // 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(