Merge branch 'staging' into T1048

near · Nov 26, 2019 · 0f6e533 · 0f6e533
2 parents 7f08cb7 + 8b0ab7b
commit 0f6e533
Show file tree

Hide file tree

Showing 20 changed files with 381 additions and 146 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/chain/chain/Cargo.toml b/chain/chain/Cargo.toml
@@ -13,7 +13,7 @@ kvdb = "0.1"
 rand = "0.7"
 serde = "1.0"
 serde_derive = "1.0"
-cached = "0.9.0"
+cached = "0.11.0"
 lazy_static = "1.4"
 rust-lzma = "0.5"
 

diff --git a/chain/chunks/Cargo.toml b/chain/chunks/Cargo.toml
@@ -15,7 +15,7 @@ borsh = "0.2.9"
 serde = "1.0"
 serde_derive = "1.0"
 reed-solomon-erasure = "3.1.1"
-cached = "0.9.0"
+cached = "0.11.0"
 
 near-crypto = { path = "../../core/crypto" }
 near-primitives = { path = "../../core/primitives" }

diff --git a/chain/chunks/src/chunk_cache.rs b/chain/chunks/src/chunk_cache.rs
@@ -1,12 +1,16 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
+use cached::{Cached, SizedCache};
+use near_primitives::hash::CryptoHash;
 use near_primitives::sharding::{
     ChunkHash, PartialEncodedChunk, PartialEncodedChunkPart, ReceiptProof, ShardChunkHeader,
 };
 use near_primitives::types::{BlockIndex, ShardId};
 
 const HEIGHT_HORIZON: u64 = 1024;
 const MAX_HEIGHTS_AHEAD: u64 = 5;
+const CHUNK_HEADER_HEIGHT_HORIZON: u64 = 10;
+const NUM_BLOCK_HASH_TO_CHUNK_HEADER: usize = 30;
 
 pub struct EncodedChunksCacheEntry {
     pub header: ShardChunkHeader,
@@ -18,7 +22,8 @@ pub struct EncodedChunksCache {
     largest_seen_height: BlockIndex,
 
     encoded_chunks: HashMap<ChunkHash, EncodedChunksCacheEntry>,
-    height_map: HashMap<BlockIndex, Vec<ChunkHash>>,
+    height_map: HashMap<BlockIndex, HashSet<ChunkHash>>,
+    block_hash_to_chunk_headers: SizedCache<CryptoHash, Vec<(ShardId, ShardChunkHeader)>>,
 }
 
 impl EncodedChunksCacheEntry {
@@ -49,6 +54,7 @@ impl EncodedChunksCache {
             largest_seen_height: 0,
             encoded_chunks: HashMap::new(),
             height_map: HashMap::new(),
+            block_hash_to_chunk_headers: SizedCache::with_size(NUM_BLOCK_HASH_TO_CHUNK_HEADER),
         }
     }
 
@@ -61,6 +67,10 @@ impl EncodedChunksCache {
     }
 
     pub fn insert(&mut self, chunk_hash: ChunkHash, entry: EncodedChunksCacheEntry) {
+        self.height_map
+            .entry(entry.header.inner.height_created)
+            .or_insert_with(|| HashSet::default())
+            .insert(chunk_hash.clone());
         self.encoded_chunks.insert(chunk_hash, entry);
     }
 
@@ -91,8 +101,13 @@ impl EncodedChunksCache {
     ) -> bool {
         let chunk_hash = partial_encoded_chunk.chunk_hash.clone();
         if self.encoded_chunks.contains_key(&chunk_hash) || partial_encoded_chunk.header.is_some() {
-            self.get_or_insert_from_header(chunk_hash, partial_encoded_chunk.header.as_ref())
-                .merge_in_partial_encoded_chunk(&partial_encoded_chunk);
+            let entry = self.get_or_insert_from_header(
+                chunk_hash.clone(),
+                partial_encoded_chunk.header.as_ref(),
+            );
+            let height = entry.header.inner.height_created;
+            entry.merge_in_partial_encoded_chunk(&partial_encoded_chunk);
+            self.height_map.entry(height).or_insert_with(|| HashSet::default()).insert(chunk_hash);
             return true;
         } else {
             return false;
@@ -127,4 +142,78 @@ impl EncodedChunksCache {
             }
         }
     }
+
+    pub fn insert_chunk_header(&mut self, shard_id: ShardId, header: ShardChunkHeader) {
+        let height = header.inner.height_created;
+        if height >= self.largest_seen_height.saturating_sub(CHUNK_HEADER_HEIGHT_HORIZON)
+            && height <= self.largest_seen_height + MAX_HEIGHTS_AHEAD
+        {
+            let mut block_hash_to_chunk_headers = self
+                .block_hash_to_chunk_headers
+                .cache_remove(&header.inner.prev_block_hash)
+                .unwrap_or_else(|| vec![]);
+            let prev_block_hash = header.inner.prev_block_hash;
+            block_hash_to_chunk_headers.push((shard_id, header));
+            self.block_hash_to_chunk_headers
+                .cache_set(prev_block_hash, block_hash_to_chunk_headers);
+        }
+    }
+
+    pub fn get_chunk_headers_for_block(
+        &mut self,
+        prev_block_hash: &CryptoHash,
+    ) -> Vec<(ShardId, ShardChunkHeader)> {
+        self.block_hash_to_chunk_headers.cache_remove(prev_block_hash).unwrap_or_else(|| vec![])
+    }
+
+    pub fn num_chunks_for_block(&mut self, prev_block_hash: &CryptoHash) -> ShardId {
+        self.block_hash_to_chunk_headers
+            .cache_get(prev_block_hash)
+            .map(|x| x.len() as ShardId)
+            .unwrap_or_else(|| 0)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::chunk_cache::EncodedChunksCache;
+    use crate::ChunkRequestInfo;
+    use near_crypto::{InMemorySigner, KeyType};
+    use near_primitives::hash::CryptoHash;
+    use near_primitives::sharding::{PartialEncodedChunk, ShardChunkHeader};
+    use std::collections::HashMap;
+
+    #[test]
+    fn test_cache_removal() {
+        let mut cache = EncodedChunksCache::new();
+        let signer = InMemorySigner::from_random("test".to_string(), KeyType::ED25519);
+        let partial_encoded_chunk = PartialEncodedChunk {
+            shard_id: 0,
+            chunk_hash: Default::default(),
+            header: Some(ShardChunkHeader::new(
+                CryptoHash::default(),
+                CryptoHash::default(),
+                CryptoHash::default(),
+                CryptoHash::default(),
+                1,
+                1,
+                0,
+                0,
+                0,
+                0,
+                0,
+                0,
+                CryptoHash::default(),
+                CryptoHash::default(),
+                vec![],
+                &signer,
+            )),
+            parts: vec![],
+            receipts: vec![],
+        };
+        assert!(cache.merge_in_partial_encoded_chunk(&partial_encoded_chunk));
+        cache.update_largest_seen_height::<ChunkRequestInfo>(2000, &HashMap::default());
+        assert!(cache.encoded_chunks.is_empty());
+        assert!(cache.height_map.is_empty());
+    }
 }