feat: Initial definition and plumbing of messages for separating cont…

…ract distribution from witness (#12224)

This PR introduced basic data structures and does some basic plumbing of
the data. We cover the chunk producer part of the flow and do not
implement the message handling on the chunk validator side. Also we do
not handle the handling of requesting and responding to code requests.

The high-level flow:
- Chunk producer:
- Identifies the contracts accessed (ie. called) when applying the
previous chunk, and saves them as part of state transition data.
- Sends the hashes of accessed contracts separately to chunk validators,
when sending the state witness.
- Chunk validator:
  - Upon receiving the code hashes, checks which contracts are missing.
- Requests the missing contracts (by code hash) from the chunk producer.
- Chunk producer:
- Responds to requests for retrieving missing contracts by sending raw
code bytes to chunk validators.
- Chunk validator:
- Upon receiving the missing code, includes them in the partial storage
and starts validating the witness (note that the runtime will compile
the code if missing from compiled-contract cache). Also starts comping
the code in parallel.

Changes in this PR:

Add 3 messages:
- `ChunkContractAccesses`: Contains the code-hashes of the contracts
accessed while applying the chunk. Sent from chunk producer to chunk
validators (the same set of validators as state witness).
- `ContractCodeRequest`: Contains a list of code-hashes that chunk
validators request from chunk producer. Sent from a chunk validator to
the chunk producer that sent a state witness.
- `ContractCodeResponse`: Contains a list of code bytes as response to a
`ContractCodeRequest`. Sent from a chunk producer to a chunk validator.

Introduce a new feature flag `contract_distribution` to guard some
changes that should not yet be included in the main binary (eg. new
field to database column and sending new messages).

Added some plumbing to pass a list of code-hashes from chunk application
up to the point we send state witness.
- Return the code-hashes from `TrieUpdate::finalize`. Also contain the
return values from `finalize` in a new struct called `TrieUpdateResult`
for readability.
- Pass the code-hashes from `TrieUpdateResult` to `ApplyResult` and
`ApplyChunkResult` in `validate_apply_state_update`.
- Pass the code-hashes from `ApplyChunkResult` to
`StoredChunkStateTransitionData` in
`ChainUpdate::save_state_transition_data`(currently guarded by feature
flag, since this needs a database migration).
- Read `StoredChunkStateTransitionData` in `create_state_witness` and
return a pair of state witness and list of code-hashes.

This also includes a database migration that converts
StoredChunkStateTransitionData to an enum and adds a new field to the
first version (V1).

chain/chain/src/chain_update.rs

@@ -155,6 +155,7 @@ impl<'a> ChainUpdate<'a> {
                         shard_id,
                         apply_result.proof,
                         apply_result.applied_receipts_hash,
+                        apply_result.contract_accesses,
                     );
                 }
             }
@@ -184,6 +185,7 @@ impl<'a> ChainUpdate<'a> {
                         shard_uid.shard_id(),
                         apply_result.proof,
                         apply_result.applied_receipts_hash,
+                        apply_result.contract_accesses,
                     );
                 }
             }

chain/chain/src/resharding/manager.rs

@@ -159,6 +159,9 @@ impl ReshardingManager {
                 new_shard_uid.shard_id(),
                 Some(partial_storage),
                 CryptoHash::default(),
+                // No contract code is accessed during resharding.
+                // TODO(#11099): Confirm if sending no contracts is ok here.
+                vec![],
             );
 
             // Commit `TrieChanges` directly. They are needed to serve reads of

chain/chain/src/runtime/mod.rs

@@ -466,6 +466,7 @@ impl NightshadeRuntime {
             processed_yield_timeouts: apply_result.processed_yield_timeouts,
             applied_receipts_hash: hash(&borsh::to_vec(receipts).unwrap()),
             congestion_info: apply_result.congestion_info,
+            contract_accesses: apply_result.contract_accesses,
         };
 
         Ok(result)

chain/chain/src/stateless_validation/state_transition_data.rs

@@ -115,7 +115,9 @@ mod tests {
 
     use near_primitives::block_header::{BlockHeader, BlockHeaderInnerLite, BlockHeaderV4};
     use near_primitives::hash::{hash, CryptoHash};
-    use near_primitives::stateless_validation::stored_chunk_state_transition_data::StoredChunkStateTransitionData;
+    use near_primitives::stateless_validation::stored_chunk_state_transition_data::{
+        StoredChunkStateTransitionData, StoredChunkStateTransitionDataV1,
+    };
     use near_primitives::types::{new_shard_id_tmp, BlockHeight, EpochId, ShardId};
     use near_primitives::utils::{get_block_shard_id, get_block_shard_id_rev, index_to_bytes};
     use near_store::db::STATE_TRANSITION_START_HEIGHTS;
@@ -202,10 +204,12 @@ mod tests {
             .set_ser(
                 DBCol::StateTransitionData,
                 &get_block_shard_id(&block_hash, shard_id),
-                &StoredChunkStateTransitionData {
+                &StoredChunkStateTransitionData::V1(StoredChunkStateTransitionDataV1 {
                     base_state: Default::default(),
                     receipts_hash: Default::default(),
-                },
+                    // TODO(#11099): Revisit this.
+                    contract_accesses: vec![],
+                }),
             )
             .unwrap();
         store_update

chain/chain/src/store/mod.rs

@@ -23,7 +23,10 @@ use near_primitives::sharding::{
 use near_primitives::state_sync::{
     ReceiptProofResponse, ShardStateSyncResponseHeader, StateHeaderKey, StateSyncDumpProgress,
 };
-use near_primitives::stateless_validation::stored_chunk_state_transition_data::StoredChunkStateTransitionData;
+use near_primitives::stateless_validation::contract_distribution::CodeHash;
+use near_primitives::stateless_validation::stored_chunk_state_transition_data::{
+    StoredChunkStateTransitionData, StoredChunkStateTransitionDataV1,
+};
 use near_primitives::transaction::{
     ExecutionOutcomeWithId, ExecutionOutcomeWithIdAndProof, ExecutionOutcomeWithProof,
     SignedTransaction,
@@ -2013,14 +2016,16 @@ impl<'a> ChainStoreUpdate<'a> {
         shard_id: ShardId,
         partial_storage: Option<PartialStorage>,
         applied_receipts_hash: CryptoHash,
+        contract_accesses: Vec<CodeHash>,
     ) {
         if let Some(partial_storage) = partial_storage {
             self.state_transition_data.insert(
                 (block_hash, shard_id),
-                StoredChunkStateTransitionData {
+                StoredChunkStateTransitionData::V1(StoredChunkStateTransitionDataV1 {
                     base_state: partial_storage.nodes,
                     receipts_hash: applied_receipts_hash,
-                },
+                    contract_accesses: contract_accesses,
+                }),
             );
         }
     }

chain/chain/src/test_utils/kv_runtime.rs

@@ -1343,6 +1343,8 @@ impl RuntimeAdapter for KeyValueRuntime {
             processed_yield_timeouts: vec![],
             applied_receipts_hash: hash(&borsh::to_vec(receipts).unwrap()),
             congestion_info: Self::get_congestion_info(PROTOCOL_VERSION),
+            // Since all actions are transfer actions, there is no contracts accessed.
+            contract_accesses: vec![],
         })
     }
 

chain/chain/src/types.rs

@@ -22,6 +22,7 @@ use near_primitives::receipt::{PromiseYieldTimeout, Receipt};
 use near_primitives::sandbox::state_patch::SandboxStatePatch;
 use near_primitives::shard_layout::ShardUId;
 use near_primitives::state_part::PartId;
+use near_primitives::stateless_validation::contract_distribution::CodeHash;
 use near_primitives::transaction::{ExecutionOutcomeWithId, SignedTransaction};
 use near_primitives::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
 use near_primitives::types::{
@@ -101,6 +102,8 @@ pub struct ApplyChunkResult {
     /// should be set to None for chunks before the CongestionControl protocol
     /// version and Some otherwise.
     pub congestion_info: Option<CongestionInfo>,
+    /// Hashes of the contracts accessed while applying the chunk.
+    pub contract_accesses: Vec<CodeHash>,
 }
 
 impl ApplyChunkResult {

chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs

@@ -8,12 +8,16 @@ use near_chain::Error;
 use near_chain_configs::MutableValidatorSigner;
 use near_epoch_manager::EpochManagerAdapter;
 use near_network::state_witness::{
-    ChunkStateWitnessAckMessage, PartialEncodedStateWitnessForwardMessage,
+    ChunkContractAccessesMessage, ChunkStateWitnessAckMessage, ContractCodeRequestMessage,
+    ContractCodeResponseMessage, PartialEncodedStateWitnessForwardMessage,
     PartialEncodedStateWitnessMessage,
 };
 use near_network::types::{NetworkRequests, PeerManagerAdapter, PeerManagerMessageRequest};
 use near_performance_metrics_macros::perf;
 use near_primitives::sharding::ShardChunkHeader;
+use near_primitives::stateless_validation::contract_distribution::{
+    ChunkContractAccesses, ContractCodeRequest, ContractCodeResponse,
+};
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
 use near_primitives::stateless_validation::state_witness::{
     ChunkStateWitness, ChunkStateWitnessAck, EncodedChunkStateWitness,
@@ -99,6 +103,30 @@ impl Handler<PartialEncodedStateWitnessForwardMessage> for PartialWitnessActor {
     }
 }
 
+impl Handler<ChunkContractAccessesMessage> for PartialWitnessActor {
+    fn handle(&mut self, msg: ChunkContractAccessesMessage) {
+        if let Err(err) = self.handle_chunk_contract_accesses(msg.0) {
+            tracing::error!(target: "client", ?err, "Failed to handle ChunkContractAccessesMessage");
+        }
+    }
+}
+
+impl Handler<ContractCodeRequestMessage> for PartialWitnessActor {
+    fn handle(&mut self, msg: ContractCodeRequestMessage) {
+        if let Err(err) = self.handle_contract_code_request(msg.0) {
+            tracing::error!(target: "client", ?err, "Failed to handle ContractCodeRequestMessage");
+        }
+    }
+}
+
+impl Handler<ContractCodeResponseMessage> for PartialWitnessActor {
+    fn handle(&mut self, msg: ContractCodeResponseMessage) {
+        if let Err(err) = self.handle_contract_code_response(msg.0) {
+            tracing::error!(target: "client", ?err, "Failed to handle ContractCodeResponseMessage");
+        }
+    }
+}
+
 impl PartialWitnessActor {
     pub fn new(
         clock: Clock,
@@ -319,6 +347,40 @@ impl PartialWitnessActor {
     pub fn handle_chunk_state_witness_ack(&mut self, witness_ack: ChunkStateWitnessAck) {
         self.state_witness_tracker.on_witness_ack_received(witness_ack);
     }
+
+    /// Handles contract code accesses message from chunk producer.
+    /// This is sent in parallel to a chunk state witness and contains the code-hashes
+    /// of the contracts accessed when applying the previous chunk of the witness.
+    pub fn handle_chunk_contract_accesses(
+        &mut self,
+        _accesses: ChunkContractAccesses,
+    ) -> Result<(), Error> {
+        // TODO(#11099): Implement this and remove debug message.
+        tracing::debug!(target: "client", next_chunk=?_accesses.chunk_production_key(), contracts=?_accesses.contracts(), "handle_chunk_contract_accesses");
+        unimplemented!()
+    }
+
+    /// Handles contract code requests message from chunk validators.
+    /// As response to this message, sends the contract code requested to
+    /// the requesting chunk validator for the given code hashes.
+    pub fn handle_contract_code_request(
+        &mut self,
+        _request: ContractCodeRequest,
+    ) -> Result<(), Error> {
+        // TODO(#11099): Implement this and remove debug message.
+        tracing::debug!(target: "client", next_chunk=?_request.chunk_production_key(), contracts=?_request.contracts(), "handle_contract_code_request");
+        unimplemented!()
+    }
+
+    /// Handles contract code responses message from chunk producer.
+    pub fn handle_contract_code_response(
+        &mut self,
+        _response: ContractCodeResponse,
+    ) -> Result<(), Error> {
+        // TODO(#11099): Implement this and remove debug message.
+        tracing::debug!(target: "client", "handle_contract_code_response");
+        unimplemented!()
+    }
 }
 
 fn compress_witness(witness: &ChunkStateWitness) -> Result<EncodedChunkStateWitness, Error> {

chain/client/src/stateless_validation/shadow_validate.rs

@@ -74,7 +74,7 @@ impl Client {
             ));
         };
 
-        let witness = self.create_state_witness(
+        let (witness, _contract_accesses) = self.create_state_witness(
             // Setting arbitrary chunk producer is OK for shadow validation
             "alice.near".parse().unwrap(),
             prev_block_header,