hide logic in obtain

core/store/src/flat/chunk_view.rs

@@ -51,4 +51,8 @@ impl FlatStorageChunkView {
     ) -> impl Iterator<Item = (Vec<u8>, Box<[u8]>)> + 'a {
         store_helper::iter_flat_state_entries(&self.store, from, to)
     }
+
+    pub fn get_head_hash(&self) -> CryptoHash {
+        self.flat_storage.get_head_hash()
+    }
 }

core/store/src/flat/manager.rs

@@ -91,47 +91,28 @@ impl FlatStorageManager {
         assert!(original_value.is_none());
     }
 
-    /// Creates `FlatStorageChunkView` to access state for `shard_id` and block `block_hash`. Note that
-    /// the state includes changes by the block `block_hash`.
-    /// `block_hash`: only create FlatStorageChunkView if it is not None. This is a hack we have temporarily
-    ///               to not introduce too many changes in the trie interface.
-    /// `is_view`: whether this flat state is used for view client. We use a separate set of caches
-    ///            for flat state for client vs view client. For now, we don't support flat state
-    ///            for view client, so we simply return None if `is_view` is True.
-    /// TODO (#7327): take block_hash as CryptoHash instead of Option<CryptoHash>
-    /// TODO (#7327): implement support for view_client
+    /// Creates `FlatStorageChunkView` to access state for `shard_id` and block `block_hash`.
+    /// Note that:
+    /// * the state includes changes by the block `block_hash`;
+    /// * request to get value locks shared `FlatStorage` struct which may cause write slowdowns.
     pub fn chunk_view(
         &self,
         shard_uid: ShardUId,
-        block_hash: Option<CryptoHash>,
-        is_view: bool,
+        block_hash: CryptoHash,
     ) -> Option<FlatStorageChunkView> {
-        let block_hash = match block_hash {
-            Some(block_hash) => block_hash,
-            None => {
-                return None;
+        let flat_storage = {
+            let flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR);
+            // It is possible that flat storage state does not exist yet because it is being created in
+            // background.
+            match flat_storages.get(&shard_uid) {
+                Some(flat_storage) => flat_storage.clone(),
+                None => {
+                    debug!(target: "chain", "FlatStorage is not ready");
+                    return None;
+                }
             }
         };
-
-        if is_view {
-            // TODO (#7327): Technically, like TrieCache, we should have a separate set of caches for Client and
-            // ViewClient. Right now, we can get by by not enabling flat state for view trie
-            None
-        } else {
-            let flat_storage = {
-                let flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR);
-                // It is possible that flat storage state does not exist yet because it is being created in
-                // background.
-                match flat_storages.get(&shard_uid) {
-                    Some(flat_storage) => flat_storage.clone(),
-                    None => {
-                        debug!(target: "chain", "FlatStorage is not ready");
-                        return None;
-                    }
-                }
-            };
-            Some(FlatStorageChunkView::new(self.0.store.clone(), block_hash, flat_storage))
-        }
+        Some(FlatStorageChunkView::new(self.0.store.clone(), block_hash, flat_storage))
     }
 
     // TODO (#7327): consider returning Result<FlatStorage, Error> when we expect flat storage to exist

core/store/src/flat/storage.rs

@@ -588,8 +588,7 @@ mod tests {
             let block_hash = chain.get_block_hash(i);
             let blocks = flat_storage.get_blocks_to_head(&block_hash).unwrap();
             assert_eq!(blocks.len(), i as usize);
-            let chunk_view =
-                flat_storage_manager.chunk_view(shard_uid, Some(block_hash), false).unwrap();
+            let chunk_view = flat_storage_manager.chunk_view(shard_uid, block_hash).unwrap();
             assert_eq!(
                 chunk_view.get_value(&[1]).unwrap(),
                 Some(FlatStateValue::value_ref(&[i as u8]))
@@ -614,12 +613,10 @@ mod tests {
         //    Verify that they return the correct values
         let blocks = flat_storage.get_blocks_to_head(&chain.get_block_hash(10)).unwrap();
         assert_eq!(blocks.len(), 10);
-        let chunk_view0 = flat_storage_manager
-            .chunk_view(shard_uid, Some(chain.get_block_hash(10)), false)
-            .unwrap();
-        let chunk_view1 = flat_storage_manager
-            .chunk_view(shard_uid, Some(chain.get_block_hash(4)), false)
-            .unwrap();
+        let chunk_view0 =
+            flat_storage_manager.chunk_view(shard_uid, chain.get_block_hash(10)).unwrap();
+        let chunk_view1 =
+            flat_storage_manager.chunk_view(shard_uid, chain.get_block_hash(4)).unwrap();
         assert_eq!(chunk_view0.get_value(&[1]).unwrap(), None);
         assert_eq!(chunk_view0.get_value(&[2]).unwrap(), Some(FlatStateValue::value_ref(&[1])));
         assert_eq!(chunk_view1.get_value(&[1]).unwrap(), Some(FlatStateValue::value_ref(&[4])));

core/store/src/flat/store_helper.rs

@@ -140,15 +140,13 @@ pub fn iter_flat_state_entries<'a>(
     from: &'a [u8],
     to: &'a [u8],
 ) -> impl Iterator<Item = (Vec<u8>, Box<[u8]>)> + 'a {
-    store.iter_range(FlatStateColumn::State.to_db_col(), Some(from), Some(to)).filter_map(
-        move |result| {
-            if let Ok((key, value)) = result {
-                let (_, trie_key) = decode_flat_state_db_key(&key).unwrap();
-                return Some((trie_key, value));
-            }
-            return None;
-        },
-    )
+    store.iter_range(DBCol::FlatState, Some(from), Some(to)).filter_map(move |result| {
+        if let Ok((key, value)) = result {
+            let (_, trie_key) = decode_flat_state_db_key(&key).unwrap();
+            return Some((trie_key, value));
+        }
+        return None;
+    })
 }
 
 /// Currently all the data in flat storage is 'together' - so we have to parse the key,

core/store/src/trie/mod.rs

@@ -857,16 +857,13 @@ impl Trie {
     {
         let mut memory = NodesStorage::new();
         let mut root_node = self.move_node_to_mutable(&mut memory, &self.root)?;
-        let mut changes_num = 0;
         for (key, value) in changes {
-            changes_num += 1;
             let key = NibbleSlice::new(&key);
             root_node = match value {
                 Some(arr) => self.insert(&mut memory, root_node, key, arr),
                 None => self.delete(&mut memory, root_node, key),
             }?;
         }
-        eprintln!("FS KV pairs = {}", changes_num);
 
         #[cfg(test)]
         {

core/store/src/trie/shard_tries.rs

@@ -146,8 +146,9 @@ impl ShardTries {
             is_view,
             prefetch_api,
         ));
-        let flat_storage_chunk_view =
-            self.0.flat_storage_manager.chunk_view(shard_uid, block_hash, is_view);
+        let flat_storage_chunk_view = block_hash
+            .map(|block_hash| self.0.flat_storage_manager.chunk_view(shard_uid, block_hash))
+            .flatten();
 
         Trie::new(storage, state_root, flat_storage_chunk_view)
     }
@@ -161,8 +162,9 @@ impl ShardTries {
         shard_uid: ShardUId,
         state_root: StateRoot,
         block_hash: &CryptoHash,
+        is_view: bool,
     ) -> Trie {
-        self.get_trie_for_shard_internal(shard_uid, state_root, false, Some(*block_hash))
+        self.get_trie_for_shard_internal(shard_uid, state_root, is_view, Some(*block_hash))
     }
 
     pub fn get_view_trie_for_shard(&self, shard_uid: ShardUId, state_root: StateRoot) -> Trie {