Ndev 3385 Caching ProgramData accounts (#569)

* Hash , rpc + db_work

* fix clippy errors

* test fixing

* apply iface fixes for rocksdb

* account check

* add a unit test and mock rpc

* fixed RPC mock and add unit test

* applied mr comments

* change rpc API && revert clickhouse job

evm_loader/lib/Cargo.toml

@@ -52,6 +52,7 @@ clap = "2.34.0"
 lazy_static = "1.5.0"
 elsa = "1.10.0"
 arrayref = "0.3.8"
+futures = "0.3.30"
 
 [dev-dependencies]
 hex-literal = "0.4.1"

evm_loader/lib/src/account_storage_tests.rs

@@ -3,6 +3,7 @@ use crate::rpc;
 use crate::tracing::AccountOverride;
 use evm_loader::types::vector::VectorVecExt;
 use hex_literal::hex;
+use solana_account_decoder::UiDataSliceConfig;
 use std::collections::HashMap;
 use std::str::FromStr;
 
@@ -19,6 +20,8 @@ mod mock_rpc_client {
     use solana_sdk::pubkey::Pubkey;
     use std::collections::HashMap;
 
+    use solana_account_decoder::UiDataSliceConfig;
+
     pub struct MockRpcClient {
         accounts: HashMap<Pubkey, Account>,
     }
@@ -33,7 +36,31 @@ mod mock_rpc_client {
 
     #[async_trait(?Send)]
     impl Rpc for MockRpcClient {
-        async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
+        async fn get_account_slice(
+            &self,
+            key: &Pubkey,
+            slice: Option<UiDataSliceConfig>,
+        ) -> ClientResult<Option<Account>> {
+            if let Some(data_slice) = slice {
+                if let Some(orig_acc) = self.accounts.get(key) {
+                    let cut_to =
+                        usize::min(data_slice.offset + data_slice.length, orig_acc.data.len());
+                    let sliced_data = if data_slice.offset < orig_acc.data.len() {
+                        orig_acc.data[data_slice.offset..cut_to].to_vec()
+                    } else {
+                        vec![]
+                    };
+
+                    return Ok(Some(Account {
+                        lamports: orig_acc.lamports,
+                        data: sliced_data,
+                        owner: orig_acc.owner,
+                        executable: orig_acc.executable,
+                        rent_epoch: orig_acc.rent_epoch,
+                    }));
+                }
+            }
+
             let result = self.accounts.get(key).cloned();
             Ok(result)
         }
@@ -1790,3 +1817,35 @@ async fn test_storage_new_from_other_and_override() {
         expected_balance
     );
 }
+
+#[tokio::test]
+async fn test_storage_get_account_slice() {
+    let slice_from = 2;
+    let slice_size = 20;
+    let test_key = Pubkey::new_unique();
+    let acc = Account::new(10, 1 * 1024 * 1024, &solana_sdk::sysvar::rent::id());
+
+    let account_tuple = (test_key, acc);
+    let accounts_for_rpc = vec![
+        (solana_sdk::sysvar::rent::id(), account_tuple.1.clone()),
+        account_tuple.clone(),
+    ];
+    let rpc_client = mock_rpc_client::MockRpcClient::new(&accounts_for_rpc);
+    let acc_no_slice = rpc_client
+        .get_account(&test_key)
+        .await
+        .expect("Failed to get account slice");
+
+    let slice_cfg = UiDataSliceConfig {
+        offset: slice_from,
+        length: slice_size,
+    };
+    let sliced_acc = rpc_client
+        .get_account_slice(&test_key, Some(slice_cfg))
+        .await
+        .expect("Failed to get account slice");
+    assert!(acc_no_slice.is_some());
+    assert!(sliced_acc.is_some());
+    assert!(acc_no_slice.unwrap().data.len() > 2000);
+    assert_eq!(sliced_acc.unwrap().data.len(), slice_size);
+}

evm_loader/lib/src/commands/get_config.rs

@@ -85,7 +85,6 @@ impl BuildConfigSimulator for CallDbClient {
     async fn build_config_simulator(&self, program_id: Pubkey) -> NeonResult<ConfigSimulator> {
         let mut simulator = SolanaSimulator::new_without_sync(self).await?;
         simulator.sync_accounts(self, &[program_id]).await?;
-
         Ok(ConfigSimulator::ProgramTestContext {
             program_id,
             simulator,

evm_loader/lib/src/rpc/db_call_client.rs

@@ -1,4 +1,4 @@
-use super::{e, Rpc};
+use super::{e, Rpc, SliceConfig};
 use crate::types::{TracerDb, TracerDbTrait};
 use crate::NeonError;
 use crate::NeonError::RocksDb;
@@ -42,18 +42,26 @@ impl CallDbClient {
         })
     }
 
-    async fn get_account_at(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
+    async fn get_account_at(
+        &self,
+        key: &Pubkey,
+        slice: Option<SliceConfig>,
+    ) -> ClientResult<Option<Account>> {
         self.tracer_db
-            .get_account_at(key, self.slot, self.tx_index_in_block)
+            .get_account_at(key, self.slot, self.tx_index_in_block, slice)
             .await
             .map_err(|e| e!("load account error", key, e))
     }
 }
 
 #[async_trait(?Send)]
 impl Rpc for CallDbClient {
-    async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
-        self.get_account_at(key).await
+    async fn get_account_slice(
+        &self,
+        key: &Pubkey,
+        slice: Option<SliceConfig>,
+    ) -> ClientResult<Option<Account>> {
+        self.get_account_at(key, slice).await
     }
 
     async fn get_multiple_accounts(
@@ -62,7 +70,7 @@ impl Rpc for CallDbClient {
     ) -> ClientResult<Vec<Option<Account>>> {
         let mut result = Vec::new();
         for key in pubkeys {
-            result.push(self.get_account_at(key).await?);
+            result.push(self.get_account_at(key, None).await?);
         }
         debug!("get_multiple_accounts: pubkeys={pubkeys:?} result={result:?}");
         Ok(result)

evm_loader/lib/src/rpc/emulator_client.rs

@@ -10,21 +10,35 @@ use solana_sdk::{
 
 use crate::account_storage::{fake_operator, EmulatorAccountStorage};
 
-use super::Rpc;
+use super::{Rpc, SliceConfig};
 
 #[async_trait(?Send)]
 impl<'rpc, T: Rpc> Rpc for EmulatorAccountStorage<'rpc, T> {
-    async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
-        if *key == self.operator() {
-            return Ok(Some(fake_operator()));
-        }
-
-        if let Some(account_data) = self.accounts_get(key) {
-            return Ok(Some(Account::from(&*account_data)));
+    async fn get_account_slice(
+        &self,
+        key: &Pubkey,
+        slice: Option<SliceConfig>,
+    ) -> ClientResult<Option<Account>> {
+        let answer_account = if *key == self.operator() {
+            Some(fake_operator())
+        } else if let Some(account_data) = self.accounts_get(key) {
+            Some(Account::from(&*account_data))
+        } else {
+            self._get_account_from_rpc(*key).await?.cloned()
+        };
+
+        if let Some(data_slice) = slice {
+            // if only slice is necessary - cut data
+            if let Some(mut account) = answer_account {
+                if data_slice.offset != 0 {
+                    account.data.drain(0..data_slice.offset);
+                }
+                account.data.truncate(data_slice.length);
+                return Ok(Some(account));
+            }
         }
 
-        let account = self._get_account_from_rpc(*key).await?.cloned();
-        Ok(account)
+        Ok(answer_account)
     }
 
     async fn get_multiple_accounts(

evm_loader/lib/src/rpc/mod.rs

@@ -9,20 +9,29 @@ use crate::commands::get_config::{BuildConfigSimulator, ConfigSimulator};
 use crate::{NeonError, NeonResult};
 use async_trait::async_trait;
 use enum_dispatch::enum_dispatch;
+pub use solana_account_decoder::UiDataSliceConfig as SliceConfig;
 use solana_cli::cli::CliError;
 use solana_client::client_error::Result as ClientResult;
-use solana_sdk::message::Message;
-use solana_sdk::native_token::lamports_to_sol;
 use solana_sdk::{
     account::Account,
     clock::{Slot, UnixTimestamp},
+    message::Message,
+    native_token::lamports_to_sol,
     pubkey::Pubkey,
 };
 
 #[async_trait(?Send)]
 #[enum_dispatch]
 pub trait Rpc {
-    async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>>;
+    async fn get_account_slice(
+        &self,
+        key: &Pubkey,
+        slice: Option<SliceConfig>,
+    ) -> ClientResult<Option<Account>>;
+    async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
+        self.get_account_slice(key, None).await
+    }
+
     async fn get_multiple_accounts(&self, pubkeys: &[Pubkey])
         -> ClientResult<Vec<Option<Account>>>;
     async fn get_block_time(&self, slot: Slot) -> ClientResult<UnixTimestamp>;

evm_loader/lib/src/rpc/validator_client.rs

@@ -1,6 +1,6 @@
 use crate::{config::APIOptions, Config};
 
-use super::Rpc;
+use super::{Rpc, SliceConfig};
 use async_trait::async_trait;
 use solana_account_decoder::{UiAccount, UiAccountEncoding};
 use solana_client::{
@@ -15,6 +15,7 @@ use solana_sdk::{
     clock::{Slot, UnixTimestamp},
     pubkey::Pubkey,
 };
+
 use std::{error::Error, ops::Deref, time::Duration};
 use std::{future::Future, sync::Arc};
 use tracing::debug;
@@ -113,12 +114,16 @@ impl Deref for CloneRpcClient {
 
 #[async_trait(?Send)]
 impl Rpc for CloneRpcClient {
-    async fn get_account(&self, key: &Pubkey) -> ClientResult<Option<Account>> {
+    async fn get_account_slice(
+        &self,
+        key: &Pubkey,
+        slice: Option<SliceConfig>,
+    ) -> ClientResult<Option<Account>> {
         let request = || {
             let config = RpcAccountInfoConfig {
                 encoding: Some(UiAccountEncoding::Base64Zstd),
                 commitment: Some(self.commitment()),
-                data_slice: None,
+                data_slice: slice,
                 min_context_slot: None,
             };
             let params = serde_json::json!([key.to_string(), config]);

evm_loader/lib/src/solana_simulator/mod.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 pub use error::Error;
 use evm_loader::solana_program::bpf_loader_upgradeable::UpgradeableLoaderState;
 use evm_loader::solana_program::clock::Slot;
+
 use evm_loader::solana_program::loader_v4;
 use evm_loader::solana_program::loader_v4::{LoaderV4State, LoaderV4Status};
 use evm_loader::solana_program::message::SanitizedMessage;
@@ -51,6 +52,7 @@ use solana_sdk::{
 pub use utils::SyncState;
 
 use crate::rpc::Rpc;
+use crate::types::programs_cache::programdata_cache_get_values_by_keys;
 
 mod error;
 mod utils;
@@ -113,7 +115,6 @@ impl SolanaSimulator {
             let Some(account) = account else {
                 continue;
             };
-
             if account.executable && bpf_loader_upgradeable::check_id(&account.owner) {
                 let programdata_address = utils::program_data_address(account)?;
                 debug!(
@@ -129,7 +130,9 @@ impl SolanaSimulator {
             storable_accounts.push((key, account));
         }
 
-        let mut programdata_accounts = rpc.get_multiple_accounts(&programdata_keys).await?;
+        let mut programdata_accounts =
+            programdata_cache_get_values_by_keys(&programdata_keys, rpc).await?;
+
         for (key, account) in programdata_keys.iter().zip(&mut programdata_accounts) {
             let Some(account) = account else {
                 continue;

evm_loader/lib/src/types/mod.rs

@@ -1,5 +1,6 @@
 pub mod tracer_ch_common;
 
+pub mod programs_cache;
 pub(crate) mod tracer_ch_db;
 pub mod tracer_rocks_db;
 
@@ -26,6 +27,8 @@ use evm_loader::{
 use serde::{Deserialize, Serialize};
 use serde_with::skip_serializing_none;
 use serde_with::{hex::Hex, serde_as, DisplayFromStr, OneOrMany};
+
+use crate::rpc::SliceConfig;
 use solana_sdk::signature::Signature;
 use solana_sdk::{account::Account, pubkey::Pubkey};
 use std::collections::HashMap;
@@ -79,6 +82,7 @@ pub trait TracerDbTrait {
         pubkey: &Pubkey,
         slot: u64,
         tx_index_in_block: Option<u64>,
+        data_slice: Option<SliceConfig>,
     ) -> DbResult<Option<Account>>;
 
     async fn get_transaction_index(&self, signature: Signature) -> DbResult<u64>;