move everything questionable to cold (#8942)

Addressing some of the TODOs by Waclaw.
Making PartialChunks, ChunkHashesByHeight, and BlockPerHeight cold instead of implementing logic to recompute them.
This adds ~30Gb to cold DB.
Should be included in 1.34.

core/store/src/cold_storage.rs

@@ -274,6 +274,7 @@ fn get_keys_from_store(
         key_type_to_keys.insert(
             key_type,
             match key_type {
+                DBKeyType::BlockHeight => vec![height_key.to_vec()],
                 DBKeyType::BlockHash => vec![block_hash_key.clone()],
                 DBKeyType::PreviousBlockHash => {
                     vec![block.header().prev_hash().as_bytes().to_vec()]

core/store/src/columns.rs

@@ -398,12 +398,18 @@ impl DBCol {
             DBCol::Block
             | DBCol::BlockExtra
             | DBCol::BlockInfo
+            // TODO can be reconstruction from BlockHeight instead of saving to cold storage.
+            | DBCol::BlockPerHeight
             | DBCol::ChunkExtra
+            // TODO can be changed to reconstruction from Block instead of saving to cold storage.
+            | DBCol::ChunkHashesByHeight
             | DBCol::Chunks
             | DBCol::IncomingReceipts
             | DBCol::NextBlockHashes
             | DBCol::OutcomeIds
             | DBCol::OutgoingReceipts
+            // TODO can be changed to reconstruction on request instead of saving in cold storage.
+            | DBCol::PartialChunks
             | DBCol::ReceiptIdToShardId
             | DBCol::Receipts
             | DBCol::State
@@ -420,27 +426,12 @@ impl DBCol {
             DBCol::BlocksToCatchup => false,
             // BlockRefCount is only needed when handling forks and it is not immutable.
             DBCol::BlockRefCount => false,
-            // PartialChunks can be recomputed and take a lot of space.
-            // TODO - but could it be used by the view client? If so then this
-            // column needs to be added to cold or the relevant code needs to be
-            // able to recompute it on demand.
-            DBCol::PartialChunks => false,
             // InvalidChunks is only needed at head when accepting new chunks.
             DBCol::InvalidChunks => false,
-            // TODO go/cold-store says: "this can be handled by reading Block
-            // and iterating over hashes there". Needs to be implemented or
-            // confirmed that view client doesn't use that column.
-            DBCol::ChunkHashesByHeight => false,
             // StateParts is only needed while syncing.
             DBCol::StateParts => false,
             // TrieChanges is only needed for GC.
             DBCol::TrieChanges => false,
-            // BlockPerHeight becomes equivalent to BlockHeight since in cold
-            // storage there are only final blocks.
-            // TODO but could it be used by the view client? If so then this
-            // column needs to be added to cold or the relevant code needs to be
-            // able to recompute it on demand.
-            DBCol::BlockPerHeight => false,
             // StateDlInfos is only needed when syncing and it is not immutable.
             DBCol::StateDlInfos => false,
             // TODO

integration-tests/src/tests/client/cold_storage.rs

@@ -7,7 +7,7 @@ use near_client::test_utils::TestEnv;
 use near_crypto::{InMemorySigner, KeyType};
 use near_o11y::testonly::init_test_logger;
 use near_primitives::block::Tip;
-use near_primitives::sharding::ShardChunk;
+use near_primitives::sharding::{PartialEncodedChunk, ShardChunk};
 use near_primitives::transaction::{
     Action, DeployContractAction, FunctionCallAction, SignedTransaction,
 };
@@ -30,20 +30,24 @@ fn check_key(first_store: &Store, second_store: &Store, col: DBCol, key: &[u8])
     let first_res = first_store.get(col, key);
     let second_res = second_store.get(col, key);
 
+    if col == DBCol::PartialChunks {
+        tracing::debug!("{:?}", first_store.get_ser::<PartialEncodedChunk>(col, key));
+    }
+
     assert_eq!(first_res.unwrap(), second_res.unwrap(), "col: {:?} key: {:?}", col, pretty_key);
 }
 
 fn check_iter(
     first_store: &Store,
     second_store: &Store,
     col: DBCol,
-    no_check_rules: &Vec<Box<dyn Fn(DBCol, &Box<[u8]>) -> bool>>,
+    no_check_rules: &Vec<Box<dyn Fn(DBCol, &Box<[u8]>, &Box<[u8]>) -> bool>>,
 ) -> u64 {
     let mut num_checks = 0;
     for (key, value) in first_store.iter(col).map(Result::unwrap) {
         let mut check = true;
         for no_check in no_check_rules {
-            if no_check(col, &value) {
+            if no_check(col, &key, &value) {
                 check = false;
             }
         }
@@ -163,8 +167,8 @@ fn test_storage_after_commit_of_cold_update() {
     assert_eq!(state_changes_reads, test_get_store_reads(DBCol::StateChanges));
 
     // We still need to filter out one chunk
-    let mut no_check_rules: Vec<Box<dyn Fn(DBCol, &Box<[u8]>) -> bool>> = vec![];
-    no_check_rules.push(Box::new(move |col, value| -> bool {
+    let mut no_check_rules: Vec<Box<dyn Fn(DBCol, &Box<[u8]>, &Box<[u8]>) -> bool>> = vec![];
+    no_check_rules.push(Box::new(move |col, _key, value| -> bool {
         if col == DBCol::Chunks {
             let chunk = ShardChunk::try_from_slice(&*value).unwrap();
             if *chunk.prev_block() == last_hash {
@@ -173,6 +177,24 @@ fn test_storage_after_commit_of_cold_update() {
         }
         false
     }));
+    no_check_rules.push(Box::new(move |col, _key, value| -> bool {
+        if col == DBCol::PartialChunks {
+            let chunk = PartialEncodedChunk::try_from_slice(&*value).unwrap();
+            if *chunk.prev_block() == last_hash {
+                return true;
+            }
+        }
+        false
+    }));
+    no_check_rules.push(Box::new(move |col, key, _value| -> bool {
+        if col == DBCol::ChunkHashesByHeight {
+            let height = u64::from_le_bytes(key[0..8].try_into().unwrap());
+            if height == max_height {
+                return true;
+            }
+        }
+        false
+    }));
 
     for col in DBCol::iter() {
         if col.is_cold() {
@@ -313,8 +335,8 @@ fn test_cold_db_copy_with_height_skips() {
     }
 
     // We still need to filter out one chunk
-    let mut no_check_rules: Vec<Box<dyn Fn(DBCol, &Box<[u8]>) -> bool>> = vec![];
-    no_check_rules.push(Box::new(move |col, value| -> bool {
+    let mut no_check_rules: Vec<Box<dyn Fn(DBCol, &Box<[u8]>, &Box<[u8]>) -> bool>> = vec![];
+    no_check_rules.push(Box::new(move |col, _key, value| -> bool {
         if col == DBCol::Chunks {
             let chunk = ShardChunk::try_from_slice(&*value).unwrap();
             if *chunk.prev_block() == last_hash {
@@ -323,9 +345,18 @@ fn test_cold_db_copy_with_height_skips() {
         }
         false
     }));
+    no_check_rules.push(Box::new(move |col, _key, value| -> bool {
+        if col == DBCol::PartialChunks {
+            let chunk = PartialEncodedChunk::try_from_slice(&*value).unwrap();
+            if *chunk.prev_block() == last_hash {
+                return true;
+            }
+        }
+        false
+    }));
 
     for col in DBCol::iter() {
-        if col.is_cold() {
+        if col.is_cold() && col != DBCol::ChunkHashesByHeight {
             let num_checks = check_iter(
                 &env.clients[0].runtime_adapter.store(),
                 &storage.get_cold_store().unwrap(),