Restructure TrieStorage and implement correct semantics for state proofs

chain/chain/src/flat_storage_creator.rs

@@ -110,7 +110,7 @@ impl FlatStorageShardCreator {
             trie_iter.visit_nodes_interval(&path_begin, &path_end).unwrap()
         {
             if let Some(key) = key {
-                let value = trie.storage.retrieve_raw_bytes(&hash).unwrap();
+                let value = trie.retrieve_value(&hash).unwrap();
                 store_helper::set_flat_state_value(
                     &mut store_update,
                     shard_uid,

core/primitives/src/types.rs

@@ -960,7 +960,7 @@ pub trait EpochInfoProvider {
 }
 
 /// Mode of the trie cache.
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum TrieCacheMode {
     /// In this mode we put each visited node to LRU cache to optimize performance.
     /// Presence of any exact node is not guaranteed.

core/store/src/test_utils.rs

@@ -1,15 +1,17 @@
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
+use near_primitives::state::{FlatStateValue, ValueRef};
 use near_primitives::trie_key::TrieKey;
 use rand::seq::SliceRandom;
 use rand::Rng;
 
 use crate::db::TestDB;
+use crate::flat::{store_helper, BlockInfo, FlatStorageReadyStatus};
 use crate::metadata::{DbKind, DbVersion, DB_VERSION};
 use crate::{get, get_delayed_receipt_indices, DBCol, NodeStorage, ShardTries, Store};
 use near_primitives::account::id::AccountId;
-use near_primitives::hash::CryptoHash;
+use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::receipt::{DataReceipt, Receipt, ReceiptEnum};
 use near_primitives::shard_layout::{ShardUId, ShardVersion};
 use near_primitives::types::{NumShards, StateRoot};
@@ -87,6 +89,34 @@ pub fn test_populate_trie(
     root
 }
 
+pub fn test_populate_flat_storage(
+    tries: &ShardTries,
+    shard_uid: ShardUId,
+    block_hash: &CryptoHash,
+    prev_block_hash: &CryptoHash,
+    changes: &Vec<(Vec<u8>, Option<Vec<u8>>)>,
+) {
+    let mut store_update = tries.store_update();
+    store_helper::set_flat_storage_status(
+        &mut store_update,
+        shard_uid,
+        crate::flat::FlatStorageStatus::Ready(FlatStorageReadyStatus {
+            flat_head: BlockInfo { hash: *block_hash, prev_hash: *prev_block_hash, height: 1 },
+        }),
+    );
+    for (key, value) in changes {
+        store_helper::set_flat_state_value(
+            &mut store_update,
+            shard_uid,
+            key.clone(),
+            value.as_ref().map(|value| {
+                FlatStateValue::Ref(ValueRef { hash: hash(value), length: value.len() as u32 })
+            }),
+        );
+    }
+    store_update.commit().unwrap();
+}
+
 /// Insert values to non-reference-counted columns in the store.
 pub fn test_populate_store(store: &Store, data: &[(DBCol, Vec<u8>, Vec<u8>)]) {
     let mut update = store.store_update();

core/store/src/trie/chunk_cache.rs

@@ -0,0 +1,117 @@
+use near_o11y::metrics::prometheus;
+use near_o11y::metrics::prometheus::core::{GenericCounter, GenericGauge};
+use near_primitives::errors::StorageError;
+use near_primitives::hash::CryptoHash;
+use near_primitives::shard_layout::ShardUId;
+use near_vm_runner::logic::TrieNodesCount;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use crate::{metrics, TrieStorage};
+
+/// Deterministic cache to store trie nodes that have been accessed so far
+/// during the processing of a single chunk.
+///
+/// This cache's correctness is critical as it contributes to the gas
+/// accounting of storage operations during contract execution. For that
+/// reason, a new TrieChunkCache must be created at the beginning of a chunk's
+/// execution, and the db_read_nodes and mem_read_nodes must be taken into
+/// account whenever a storage operation is performed to calculate what kind
+/// of operation it was.
+///
+/// Note that we don't have a size limit for values in the chunk cache. There
+/// are two reasons:
+///   - for nodes, value size is an implementation detail. If we change
+///     internal representation of a node (e.g. change `memory_usage` field
+///     from `RawTrieNodeWithSize`), this would have to be a protocol upgrade.
+///   - total size of all values is limited by the runtime fees. More
+///     thoroughly:
+///       - number of nodes is limited by receipt gas limit / touching trie
+///         node fee ~= 500 Tgas / 16 Ggas = 31_250;
+///       - size of trie keys and values is limited by receipt gas limit /
+///         lowest per byte fee (`storage_read_value_byte`) ~=
+///         (500 * 10**12 / 5611005) / 2**20 ~= 85 MB.
+/// All values are given as of 16/03/2022. We may consider more precise limit
+/// for the chunk cache as well.
+///
+/// Note that in general, it is NOT true that all storage access is either a
+/// db read or mem read. It can also be a flat storage read, which is not
+/// tracked via TrieChunkCache.
+pub struct TrieChunkCache {
+    /// Whether the cache is enabled. By default it is not, but it can be
+    /// turned on or off on the fly.
+    enable: bool,
+    /// Cache of trie node hash -> trie node body, or a leaf value hash ->
+    /// leaf value.
+    cache: HashMap<CryptoHash, Arc<[u8]>>,
+    /// The number of times a key was accessed by reading from the underlying
+    /// storage. (This does not necessarily mean it was accessed from *disk*,
+    /// as the underlying storage layer may have a best-effort cache.)
+    db_read_nodes: u64,
+    /// The number of times a key was accessed when it was deterministically
+    /// already cached during the processing of this chunk.
+    mem_read_nodes: u64,
+    /// Prometheus metrics. It's optional - in testing it can be None.
+    metrics: Option<TrieChunkCacheMetrics>,
+}
+
+struct TrieChunkCacheMetrics {
+    chunk_cache_hits: GenericCounter<prometheus::core::AtomicU64>,
+    chunk_cache_misses: GenericCounter<prometheus::core::AtomicU64>,
+    chunk_cache_size: GenericGauge<prometheus::core::AtomicI64>,
+}
+
+impl TrieChunkCache {
+    /// Constructs a new cache. By default it is not enabled.
+    /// The optional parameter is passed in if prometheus metrics are desired.
+    pub fn new(shard_uid_and_is_view: Option<(ShardUId, bool)>) -> Self {
+        let metrics = shard_uid_and_is_view.map(|(shard_uid, is_view)| {
+            let mut buffer = itoa::Buffer::new();
+            let shard_id = buffer.format(shard_uid.shard_id);
+
+            let metrics_labels: [&str; 2] = [&shard_id, if is_view { "1" } else { "0" }];
+            TrieChunkCacheMetrics {
+                chunk_cache_hits: metrics::CHUNK_CACHE_HITS.with_label_values(&metrics_labels),
+                chunk_cache_misses: metrics::CHUNK_CACHE_MISSES.with_label_values(&metrics_labels),
+                chunk_cache_size: metrics::CHUNK_CACHE_SIZE.with_label_values(&metrics_labels),
+            }
+        });
+        Self { enable: false, cache: HashMap::new(), db_read_nodes: 0, mem_read_nodes: 0, metrics }
+    }
+
+    pub fn set_enabled(&mut self, enabled: bool) {
+        self.enable = enabled;
+    }
+
+    pub fn retrieve_trie_node_with_cache(
+        &mut self,
+        hash: &CryptoHash,
+        storage: &dyn TrieStorage,
+    ) -> Result<Arc<[u8]>, StorageError> {
+        if let Some(node) = self.cache.get(hash) {
+            self.mem_read_nodes += 1;
+            if let Some(metrics) = &self.metrics {
+                metrics.chunk_cache_hits.inc();
+            }
+            Ok(node.clone())
+        } else {
+            self.db_read_nodes += 1;
+            if let Some(metrics) = &self.metrics {
+                metrics.chunk_cache_misses.inc();
+            }
+            let node = storage.retrieve_raw_bytes(hash)?;
+
+            if self.enable {
+                self.cache.insert(*hash, node.clone());
+                if let Some(metrics) = &self.metrics {
+                    metrics.chunk_cache_size.set(self.cache.len() as i64);
+                }
+            }
+            Ok(node)
+        }
+    }
+
+    pub fn get_trie_nodes_count(&self) -> TrieNodesCount {
+        TrieNodesCount { db_reads: self.db_read_nodes, mem_reads: self.mem_read_nodes }
+    }
+}

core/store/src/trie/iterator.rs

@@ -372,7 +372,7 @@ impl<'a> TrieIterator<'a> {
                     if self.key_nibbles[prefix..] >= path_end[prefix..] {
                         break;
                     }
-                    self.trie.storage.retrieve_raw_bytes(&hash)?;
+                    self.trie.retrieve_value(&hash)?;
                     nodes_list.push(TrieTraversalItem {
                         hash,
                         key: self.has_value().then(|| self.key()),
@@ -417,10 +417,7 @@ impl<'a> Iterator for TrieIterator<'a> {
                 },
                 (IterStep::Value(hash), true) => {
                     return Some(
-                        self.trie
-                            .storage
-                            .retrieve_raw_bytes(&hash)
-                            .map(|value| (self.key(), value.to_vec())),
+                        self.trie.retrieve_value(&hash).map(|value| (self.key(), value.to_vec())),
                     )
                 }
             }