Skip to content

Commit

Permalink
fix to calculate state size
Browse files Browse the repository at this point in the history
  • Loading branch information
kobayurii committed Nov 4, 2024
1 parent 17db12a commit 7a2080d
Show file tree
Hide file tree
Showing 7 changed files with 79 additions and 14 deletions.
8 changes: 8 additions & 0 deletions database/src/base/rpc_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
31 changes: 31 additions & 0 deletions database/src/postgres/rpc_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions rpc-server/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,10 @@ 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)),
Expand Down
10 changes: 9 additions & 1 deletion rpc-server/src/metrics.rs
Original file line number Diff line number Diff line change
@@ -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>;

Expand Down Expand Up @@ -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",
Expand Down
9 changes: 0 additions & 9 deletions rpc-server/src/modules/blocks/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand Down
19 changes: 16 additions & 3 deletions rpc-server/src/modules/queries/contract_runner/mod.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -143,14 +142,28 @@ 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;

Expand Down
12 changes: 11 additions & 1 deletion rpc-server/src/modules/queries/methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down

0 comments on commit 7a2080d

Please sign in to comment.