chain: ShardManager: add ability to serve partial chunks from ShardChunk

chain/chunks/Cargo.toml

@@ -16,6 +16,7 @@ tracing = "0.1.13"
 borsh = "0.9"
 lru = "0.7.2"
 reed-solomon-erasure = "4"
+visibility = "*"
 
 near-crypto = { path = "../../core/crypto" }
 near-primitives = { path = "../../core/primitives" }
@@ -33,3 +34,4 @@ assert_matches = "1.5.0"
 [features]
 byzantine_asserts = ["near-chain/byzantine_asserts"]
 expensive_tests = []
+test_features = []

chain/chunks/src/lib.rs

@@ -101,9 +101,9 @@ use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::merkle::{merklize, verify_path, MerklePath};
 use near_primitives::receipt::Receipt;
 use near_primitives::sharding::{
-    ChunkHash, EncodedShardChunk, PartialEncodedChunk, PartialEncodedChunkPart,
-    PartialEncodedChunkV1, PartialEncodedChunkV2, ReceiptList, ReceiptProof, ReedSolomonWrapper,
-    ShardChunkHeader, ShardProof,
+    ChunkHash, EncodedShardChunk, EncodedShardChunkBody, PartialEncodedChunk,
+    PartialEncodedChunkPart, PartialEncodedChunkV1, PartialEncodedChunkV2, ReceiptList,
+    ReceiptProof, ReedSolomonWrapper, ShardChunk, ShardChunkHeader, ShardProof,
 };
 use near_primitives::time::Clock;
 use near_primitives::transaction::SignedTransaction;
@@ -995,13 +995,27 @@ impl ShardsManager {
         request: PartialEncodedChunkRequestMsg,
         route_back: CryptoHash,
         chain_store: &mut ChainStore,
+        rs: &mut ReedSolomonWrapper,
     ) {
         debug!(target: "chunks", "Received partial encoded chunk request for {:?}, part_ordinals: {:?}, shards: {:?}, I'm {:?}", request.chunk_hash.0, request.part_ords, request.tracking_shards, self.me);
 
         let response = if let Some(entry) = self.encoded_chunks.get(&request.chunk_hash) {
             Self::prepare_partial_encoded_chunk_response_from_cache(request, entry)
         } else if let Ok(partial_chunk) = chain_store.get_partial_chunk(&request.chunk_hash) {
             Self::prepare_partial_encoded_chunk_response_from_partial(request, partial_chunk)
+        } else if let Ok(chunk) = chain_store.get_chunk(&request.chunk_hash).map(|ch| ch.clone()) {
+            // Note: we need to clone the chunk because otherwise we would be
+            // holding multiple references to chain_store.  One through the
+            // chunk and another through chain_store which we need to pass down
+            // to do further fetches.
+
+            // If we are archival node we might have garbage collected the
+            // partial chunk while we still keep the chunk itself.  We can get
+            // the chunk, recalculate the parts and respond to the request.
+            //
+            // TODO(#6242): This is currently not implemented and effectively
+            // this is dead code.
+            self.prepare_partial_encoded_chunk_response_from_chunk(request, chain_store, rs, chunk)
         } else {
             None
         };
@@ -1013,11 +1027,12 @@ impl ShardsManager {
         }
     }
 
-    /// Prepares response to a partial encoded chunk request from
-    /// a corresponding encoded_chunks cache entry.  If the entry can satisfy
-    /// the requests (i.e. contains all necessary parts and shards) the method
-    /// returns a [`PartialEncodedChunkResponseMsg`] object; otherwise returns
-    /// `None`.
+    /// Prepares response to a partial encoded chunk request from an entry in
+    /// a encoded_chunks in-memory cache.
+    ///
+    /// If the entry can satisfy the requests (i.e. contains all necessary parts
+    /// and shards) the method returns a [`PartialEncodedChunkResponseMsg`]
+    /// object; otherwise returns `None`.
     fn prepare_partial_encoded_chunk_response_from_cache(
         request: PartialEncodedChunkRequestMsg,
         entry: &EncodedChunksCacheEntry,
@@ -1037,11 +1052,13 @@ impl ShardsManager {
         )
     }
 
-    /// Prepares response to a partial encoded chunk request from
-    /// a corresponding partial chunk read from the storage.  If the partial
-    /// chunk can satisfy the requests (i.e. contains all necessary parts and
-    /// shards) the method returns a [`PartialEncodedChunkResponseMsg`] object;
-    /// otherwise returns `None`.
+    /// Prepares response to a partial encoded chunk request from a partial
+    /// chunk read from the storage.
+    ///
+    /// If the partial chunk can satisfy the requests (i.e. contains all
+    /// necessary parts and shards) the method returns
+    /// a [`PartialEncodedChunkResponseMsg`] object; otherwise returns `None`.
+    #[cfg_attr(feature = "test_features", visibility::make(pub))]
     fn prepare_partial_encoded_chunk_response_from_partial(
         request: PartialEncodedChunkRequestMsg,
         partial_chunk: &PartialEncodedChunk,
@@ -1075,6 +1092,143 @@ impl ShardsManager {
         )
     }
 
+    /// Prepares response to a partial encoded chunk request from a chunk read
+    /// from the storage.
+    ///
+    /// This requires encoding the chunk and as such is computationally
+    /// expensive operation.  If possible, the request should be served from
+    /// EncodedChunksCacheEntry or PartialEncodedChunk instead.
+    #[cfg_attr(feature = "test_features", visibility::make(pub))]
+    fn prepare_partial_encoded_chunk_response_from_chunk(
+        &mut self,
+        request: PartialEncodedChunkRequestMsg,
+        chain_store: &mut ChainStore,
+        rs: &mut ReedSolomonWrapper,
+        chunk: ShardChunk,
+    ) -> Option<PartialEncodedChunkResponseMsg> {
+        // TODO(mina86): Is it correct to use rs.data_shard_count() and
+        // rs.total_shard_count() here?  Can those values change over time?
+        let total_parts = rs.total_shard_count();
+        for ord in request.part_ords.iter() {
+            let ord: usize = (*ord).try_into().unwrap();
+            if ord >= total_parts {
+                debug!(target:"chunks", "Not sending {:?}, requested part_ord {} but we only expect {} total",
+                       request.chunk_hash.0, ord, total_parts);
+                return None;
+            }
+        }
+
+        // Figure out basic parameters such as shard id, epoch id, protocol
+        // version and shard layout.  Those will be needed later when fetching
+        // other data and constructing encoded shard chunk.
+        let header = chunk.cloned_header();
+        let prev_block_hash = header.prev_block_hash_ref();
+        let shard_id = header.shard_id();
+        let shard_layout = self
+            .runtime_adapter
+            .get_shard_layout_from_prev_block(prev_block_hash)
+            .map_err(|err| {
+                debug!(target: "chunks", "Not sending {:?}, failed to get shard layout: {}", request.chunk_hash.0, err);
+            })
+            .ok()?;
+
+        // Get outgoing receipts for the chunk and construct vector of their
+        // proofs.
+        let outgoing_receipts = chain_store
+            .get_outgoing_receipts_for_shard(
+                &*self.runtime_adapter,
+                *chunk.prev_block(),
+                shard_id,
+                chunk.height_included(),
+            )
+            .map_err(|err| {
+                debug!(target: "chunks", "Not sending {:?}, failed to get outgoing receipts: {}", request.chunk_hash.0, err);
+            }).ok()?;
+        let outgoing_receipts_hashes =
+            Chain::build_receipts_hashes(&outgoing_receipts, &shard_layout);
+        let (outgoing_receipts_root, outgoing_receipts_proofs) =
+            merklize(&outgoing_receipts_hashes);
+        if header.outgoing_receipts_root() != outgoing_receipts_root {
+            error!(target: "chunks",
+                   "Not sending {:?}, expected outgoing receipts root doesn’t match calculated: {} != {}",
+                   request.chunk_hash.0, header.outgoing_receipts_root(), outgoing_receipts_root);
+            return None;
+        }
+
+        let mut receipts_by_shard =
+            Chain::group_receipts_by_shard(outgoing_receipts.clone(), &shard_layout);
+        let receipt_proofs: Vec<_> = outgoing_receipts_proofs
+            .iter()
+            .enumerate()
+            .map(|(proof_shard_id, proof)| {
+                let proof_shard_id = proof_shard_id as u64;
+                let receipts = receipts_by_shard.remove(&proof_shard_id).unwrap_or_else(Vec::new);
+                let shard_proof = ShardProof {
+                    from_shard_id: shard_id,
+                    to_shard_id: proof_shard_id,
+                    proof: proof.clone(),
+                };
+                ReceiptProof(receipts, shard_proof)
+            })
+            .collect();
+
+        // Construct EncodedShardChunk.  If we earlier determined that we will
+        // need parity parts, instruct the constructor to calculate them as
+        // well.  Otherwise we won’t bother.
+        let (parts, encoded_length) = EncodedShardChunk::encode_transaction_receipts(
+            rs,
+            chunk.transactions().to_vec(),
+            &outgoing_receipts).map_err(|err| {
+                debug!(target: "chunks", "Not sending {:?}, failed to encode transaction receipts: {}", request.chunk_hash.0, err);
+            }).ok()?;
+        if header.encoded_length() != encoded_length {
+            error!(target: "chunks",
+                   "Not sending {:?}, expected encoded length doesn’t match calculated: {} != {}",
+                   request.chunk_hash.0, header.encoded_length(), encoded_length);
+            return None;
+        }
+
+        let mut content = EncodedShardChunkBody { parts };
+        if let Err(err) = content.reconstruct(rs) {
+            error!(target: "chunks",
+                   "Not sending {:?}, failed to reconstruct RS parity parts: {}",
+                   request.chunk_hash.0, err);
+            return None;
+        }
+
+        let (encoded_merkle_root, merkle_paths) = content.get_merkle_hash_and_paths();
+        if header.encoded_merkle_root() != encoded_merkle_root {
+            error!(target: "chunks",
+                   "Not sending {:?}, expected encoded Merkle root doesn’t match calculated: {} != {}",
+                   request.chunk_hash.0, header.encoded_merkle_root(), encoded_merkle_root);
+            return None;
+        }
+
+        let parts_iter = request.part_ords.into_iter().map(|part_ord| {
+            let ord: usize = part_ord.try_into().unwrap();
+            content.parts[ord].take().map(|part| PartialEncodedChunkPart {
+                part_ord,
+                part,
+                merkle_proof: merkle_paths[ord].clone(),
+            })
+        });
+
+        // Same process for receipts as above for parts.
+        let present_receipts: HashMap<ShardId, _> =
+            receipt_proofs.iter().map(|receipt| (receipt.1.to_shard_id, receipt)).collect();
+        let receipts_iter = request
+            .tracking_shards
+            .iter()
+            .map(|shard_id| present_receipts.get(shard_id).map(|x| *x).cloned());
+
+        // Pass iterators to function, same as cache case.
+        Self::prepare_partial_encoded_chunk_response_from_iters(
+            request.chunk_hash,
+            parts_iter,
+            receipts_iter,
+        )
+    }
+
     /// Checks if `parts_iter` and `receipts_iter` contain no `None` elements.
     /// It evaluates the iterators only up to the first `None` value (if any);
     /// since iterators are lazy this saves some work if there are any `Some`

chain/client/src/client_actor.rs

@@ -560,6 +560,7 @@ impl Handler<NetworkClientMessages> for ClientActor {
                     part_request_msg,
                     route_back,
                     self.client.chain.mut_store(),
+                    &mut self.client.rs,
                 );
                 NetworkClientResponses::NoResponse
             }

chain/client/src/test_utils.rs

@@ -1410,7 +1410,7 @@ impl TestEnv {
         }
     }
 
-    fn get_partial_encoded_chunk_response(
+    pub fn get_partial_encoded_chunk_response(
         &mut self,
         id: usize,
         request: PartialEncodedChunkRequestMsg,
@@ -1420,6 +1420,7 @@ impl TestEnv {
             request,
             CryptoHash::default(),
             client.chain.mut_store(),
+            &mut client.rs,
         );
         let response = self.network_adapters[id].pop().unwrap();
         if let PeerManagerMessageRequest::NetworkRequests(

chain/client/src/tests/chunks_management.rs

@@ -6,6 +6,7 @@ use near_logger_utils::init_integration_logger;
 use near_network::types::NetworkRequests;
 use near_network_primitives::types::PartialEncodedChunkRequestMsg;
 use near_primitives::hash::CryptoHash;
+use near_primitives::sharding::ReedSolomonWrapper;
 
 #[test]
 fn test_request_chunk_restart() {
@@ -22,10 +23,14 @@ fn test_request_chunk_restart() {
         tracking_shards: HashSet::default(),
     };
     let client = &mut env.clients[0];
+    let num_total_parts = client.runtime_adapter.num_total_parts();
+    let num_data_parts = client.runtime_adapter.num_data_parts();
+    let mut rs = ReedSolomonWrapper::new(num_data_parts, num_total_parts - num_data_parts);
     client.shards_mgr.process_partial_encoded_chunk_request(
         request.clone(),
         CryptoHash::default(),
         client.chain.mut_store(),
+        &mut rs,
     );
     assert!(env.network_adapters[0].pop().is_some());
 
@@ -35,6 +40,7 @@ fn test_request_chunk_restart() {
         request,
         CryptoHash::default(),
         client.chain.mut_store(),
+        &mut rs,
     );
     let response = env.network_adapters[0].pop().unwrap().as_network_requests();
 

core/primitives/src/sharding.rs

@@ -303,6 +303,17 @@ impl ShardChunkHeader {
         }
     }
 
+    // TODO(mina86): Change return type of prev_block_hash to &CryptoHash and
+    // then this method can be deleted.
+    #[inline]
+    pub fn prev_block_hash_ref(&self) -> &CryptoHash {
+        match self {
+            Self::V1(header) => &header.inner.prev_block_hash,
+            Self::V2(header) => &header.inner.prev_block_hash,
+            Self::V3(header) => header.inner.prev_block_hash(),
+        }
+    }
+
     #[inline]
     pub fn encoded_merkle_root(&self) -> CryptoHash {
         match self {
@@ -527,11 +538,7 @@ impl PartialEncodedChunk {
     pub fn prev_block(&self) -> &CryptoHash {
         match &self {
             PartialEncodedChunk::V1(chunk) => &chunk.header.inner.prev_block_hash,
-            PartialEncodedChunk::V2(chunk) => match &chunk.header {
-                ShardChunkHeader::V1(header) => &header.inner.prev_block_hash,
-                ShardChunkHeader::V2(header) => &header.inner.prev_block_hash,
-                ShardChunkHeader::V3(header) => header.inner.prev_block_hash(),
-            },
+            PartialEncodedChunk::V2(chunk) => chunk.header.prev_block_hash_ref(),
         }
     }
 
@@ -700,6 +707,14 @@ impl ShardChunk {
         }
     }
 
+    #[inline]
+    pub fn prev_block(&self) -> &CryptoHash {
+        match &self {
+            ShardChunk::V1(chunk) => &chunk.header.inner.prev_block_hash,
+            ShardChunk::V2(chunk) => chunk.header.prev_block_hash_ref(),
+        }
+    }
+
     #[inline]
     pub fn prev_state_root(&self) -> StateRoot {
         match self {
@@ -937,7 +952,7 @@ impl EncodedShardChunk {
         TransactionReceipt::try_from_slice(&encoded_data)
     }
 
-    fn encode_transaction_receipts(
+    pub fn encode_transaction_receipts(
         rs: &mut ReedSolomonWrapper,
         transactions: Vec<SignedTransaction>,
         outgoing_receipts: &[Receipt],

integration-tests/Cargo.toml

@@ -66,7 +66,7 @@ performance_stats = [
   "near-network/performance_stats",
 ]
 expensive_tests = []
-test_features = ["nearcore/test_features"]
+test_features = ["nearcore/test_features", "near-chunks/test_features"]
 protocol_feature_alt_bn128 = [
     "near-primitives/protocol_feature_alt_bn128",
     "node-runtime/protocol_feature_alt_bn128",

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

@@ -5,3 +5,5 @@ mod runtimes;
 #[cfg(feature = "sandbox")]
 mod sandbox;
 mod sharding_upgrade;
+#[cfg(feature = "test_features")]
+mod shards_manager;