Move rewards to runtime (#1497)

* Move validator rewards to runtime
* Merge branch 'staging' of github.com:nearprotocol/nearcore into move-rewards-to-runtime
* Merge branch 'staging' of github.com:nearprotocol/nearcore into move-rewards-to-runtime
* Update runtime/runtime/src/lib.rs

Co-Authored-By: Maksym Zavershynskyi <35039879+nearmax@users.noreply.github.com>
* Merge branch 'staging' of github.com:nearprotocol/nearcore into move-rewards-to-runtime
* Merge branch 'move-rewards-to-runtime' of github.com:nearprotocol/nearcore into move-rewards-to-runtime
* Merge branch 'staging' into move-rewards-to-runtime
* Merge refs/heads/staging into move-rewards-to-runtime

Author: Evgeny Kuzyakov

core/primitives/src/errors.rs

@@ -170,6 +170,7 @@ impl Display for InvalidAccessKeyError {
 #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq, Eq)]
 pub struct BalanceMismatchError {
     // Input balances
+    pub incoming_validator_rewards: Balance,
     pub initial_accounts_balance: Balance,
     pub incoming_receipts_balance: Balance,
     pub initial_postponed_receipts_balance: Balance,
@@ -186,7 +187,8 @@ impl Display for BalanceMismatchError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
         // Using saturating add to avoid overflow in display
         let initial_balance = self
-            .initial_accounts_balance
+            .incoming_validator_rewards
+            .saturating_add(self.initial_accounts_balance)
             .saturating_add(self.incoming_receipts_balance)
             .saturating_add(self.initial_postponed_receipts_balance);
         let final_balance = self
@@ -200,6 +202,7 @@ impl Display for BalanceMismatchError {
             f,
             "Balance Mismatch Error. The input balance {} doesn't match output balance {}\n\
              Inputs:\n\
+             \tIncoming validator rewards sum: {}\n\
              \tInitial accounts balance sum: {}\n\
              \tIncoming receipts balance sum: {}\n\
              \tInitial postponed receipts balance sum: {}\n\
@@ -212,6 +215,7 @@ impl Display for BalanceMismatchError {
              \tTotal balance burnt: {}",
             initial_balance,
             final_balance,
+            self.incoming_validator_rewards,
             self.initial_accounts_balance,
             self.incoming_receipts_balance,
             self.initial_postponed_receipts_balance,

core/store/src/trie/update.rs

@@ -12,9 +12,7 @@ use crate::trie::TrieChanges;
 use super::{Trie, TrieIterator};
 use crate::StorageError;
 
-// TODO(#1481): Remove clone after Runtime bug is fixed.
 /// Provides a way to access Storage and record changes with future commit.
-#[derive(Clone)]
 pub struct TrieUpdate {
     pub trie: Arc<Trie>,
     root: MerkleHash,
@@ -276,18 +274,6 @@ mod tests {
 
     use super::*;
 
-    // TODO(#1481): Remove clone after Runtime bug is fixed.
-    #[test]
-    fn trie_clone() {
-        let trie = create_trie();
-        let root = MerkleHash::default();
-        let mut trie_update = TrieUpdate::new(trie.clone(), root);
-        trie_update.set(b"dog".to_vec(), DBValue::from_slice(b"puppy"));
-        let trie_cloned = trie_update.clone();
-        trie_update.set(b"dog".to_vec(), DBValue::from_slice(b"batman"));
-        assert_eq!(trie_cloned.get(b"dog"), Ok(Some(DBValue::from_slice(b"puppy"))));
-    }
-
     #[test]
     fn trie() {
         let trie = create_trie();

near/src/runtime.rs

@@ -1,4 +1,3 @@
-use std::cmp::max;
 use std::collections::{HashMap, HashSet};
 use std::convert::{TryFrom, TryInto};
 use std::fs::File;
@@ -30,12 +29,11 @@ use near_primitives::views::{
     AccessKeyInfoView, CallResult, QueryError, QueryResponse, ViewStateResult,
 };
 use near_store::{
-    get_access_key_raw, get_account, set_account, StorageError, Store, StoreUpdate, Trie,
-    TrieUpdate, WrappedTrieChanges, COL_STATE,
+    get_access_key_raw, Store, StoreUpdate, Trie, TrieUpdate, WrappedTrieChanges, COL_STATE,
 };
 use node_runtime::adapter::ViewRuntimeAdapter;
 use node_runtime::state_viewer::TrieViewer;
-use node_runtime::{ApplyState, Runtime, StateRecord};
+use node_runtime::{ApplyState, Runtime, StateRecord, ValidatorAccountsUpdate};
 
 use crate::config::GenesisConfig;
 use crate::shard_tracker::{account_id_to_shard_id, ShardTracker};
@@ -126,68 +124,6 @@ impl NightshadeRuntime {
         }
     }
 
-    /// Iterates over validator accounts in the given shard and updates their accounts to return stake
-    /// and allocate rewards.
-    fn update_validator_accounts(
-        &self,
-        shard_id: ShardId,
-        block_hash: &CryptoHash,
-        last_validator_proposals: &[ValidatorStake],
-        state_update: &mut TrieUpdate,
-    ) -> Result<(), Box<dyn std::error::Error>> {
-        let mut epoch_manager = self.epoch_manager.write().expect(POISONED_LOCK_ERR);
-        let (stake_info, validator_reward) = epoch_manager.compute_stake_return_info(block_hash)?;
-        let account_to_stake =
-            last_validator_proposals.iter().fold(HashMap::new(), |mut acc, v| {
-                acc.insert(v.account_id.clone(), v.amount);
-                acc
-            });
-
-        for (account_id, max_of_stakes) in stake_info {
-            if self.account_id_to_shard_id(&account_id) == shard_id {
-                let account: Option<Account> = get_account(state_update, &account_id)?;
-                if let Some(mut account) = account {
-                    if let Some(reward) = validator_reward.get(&account_id) {
-                        debug!(target: "runtime", "account {} adding reward {} to stake {}", account_id, reward, account.locked);
-                        account.locked += *reward;
-                    }
-
-                    debug!(target: "runtime",
-                           "account {} stake {} max_of_stakes: {}",
-                           account_id, account.locked, max_of_stakes
-                    );
-                    assert!(
-                        account.locked >= max_of_stakes,
-                        "FATAL: staking invariant does not hold. Account stake {} is less than maximum of stakes {} in the past three epochs",
-                        account.locked,
-                        max_of_stakes
-                    );
-                    let last_stake = *account_to_stake.get(&account_id).unwrap_or(&0);
-                    let return_stake = account.locked - max(max_of_stakes, last_stake);
-                    debug!(target: "runtime", "account {} return stake {}", account_id, return_stake);
-                    account.locked -= return_stake;
-                    account.amount += return_stake;
-
-                    set_account(state_update, &account_id, &account);
-                }
-            }
-        }
-        if self.account_id_to_shard_id(&self.genesis_config.protocol_treasury_account) == shard_id {
-            let mut protocol_treasury_account =
-                get_account(state_update, &self.genesis_config.protocol_treasury_account)?.unwrap();
-            protocol_treasury_account.amount +=
-                *validator_reward.get(&self.genesis_config.protocol_treasury_account).unwrap();
-            set_account(
-                state_update,
-                &self.genesis_config.protocol_treasury_account,
-                &protocol_treasury_account,
-            );
-        }
-        state_update.commit();
-
-        Ok(())
-    }
-
     fn genesis_state_from_dump(&self) -> (StoreUpdate, Vec<StateRoot>) {
         let store_update = self.store.store_update();
         let mut state_file = self.home_dir.clone();
@@ -603,33 +539,45 @@ impl RuntimeAdapter for NightshadeRuntime {
         } else {
             self.trie.clone()
         };
-        let mut state_update = TrieUpdate::new(trie.clone(), state_root.hash);
-        let should_update_account = {
+        let validator_accounts_update = {
             let mut epoch_manager = self.epoch_manager.write().expect(POISONED_LOCK_ERR);
             debug!(target: "runtime",
                 "block index: {}, is next_block_epoch_start {}",
                 block_index,
                 epoch_manager.is_next_block_epoch_start(prev_block_hash).unwrap()
             );
-            epoch_manager.is_next_block_epoch_start(prev_block_hash)?
+            if epoch_manager.is_next_block_epoch_start(prev_block_hash)? {
+                let (stake_info, validator_reward) =
+                    epoch_manager.compute_stake_return_info(prev_block_hash)?;
+                let stake_info = stake_info
+                    .into_iter()
+                    .filter(|(account_id, _)| self.account_id_to_shard_id(account_id) == shard_id)
+                    .collect();
+                let validator_rewards = validator_reward
+                    .into_iter()
+                    .filter(|(account_id, _)| self.account_id_to_shard_id(account_id) == shard_id)
+                    .collect();
+                let last_proposals = last_validator_proposals
+                    .iter()
+                    .filter(|v| self.account_id_to_shard_id(&v.account_id) == shard_id)
+                    .fold(HashMap::new(), |mut acc, v| {
+                        acc.insert(v.account_id.clone(), v.amount);
+                        acc
+                    });
+                Some(ValidatorAccountsUpdate {
+                    stake_info,
+                    validator_rewards,
+                    last_proposals,
+                    protocol_treasury_account_id: Some(
+                        self.genesis_config.protocol_treasury_account.clone(),
+                    )
+                    .filter(|account_id| self.account_id_to_shard_id(account_id) == shard_id),
+                })
+            } else {
+                None
+            }
         };
 
-        // If we are starting to apply 1st block in the new epoch.
-        if should_update_account {
-            self.update_validator_accounts(
-                shard_id,
-                prev_block_hash,
-                last_validator_proposals,
-                &mut state_update,
-            )
-            .map_err(|e| {
-                if let Some(e) = e.downcast_ref::<StorageError>() {
-                    panic!(e.to_string())
-                }
-                Error::from(ErrorKind::ValidatorError(e.to_string()))
-            })?;
-        }
-
         let apply_state = ApplyState {
             block_index,
             epoch_length: self.genesis_config.epoch_length,
@@ -639,7 +587,14 @@ impl RuntimeAdapter for NightshadeRuntime {
 
         let apply_result = self
             .runtime
-            .apply(state_update, &apply_state, &receipts, &transactions)
+            .apply(
+                trie.clone(),
+                state_root.hash,
+                &validator_accounts_update,
+                &apply_state,
+                &receipts,
+                &transactions,
+            )
             .map_err(|e| match e {
                 RuntimeError::InvalidTxError(_) => ErrorKind::InvalidTransactions,
                 RuntimeError::BalanceMismatch(e) => panic!("{}", e),

runtime/runtime-params-estimator/src/testbed.rs

@@ -3,7 +3,7 @@ use near::get_store_path;
 use near_primitives::receipt::Receipt;
 use near_primitives::transaction::{ExecutionStatus, SignedTransaction};
 use near_primitives::types::MerkleHash;
-use near_store::{create_store, Trie, TrieUpdate, COL_STATE};
+use near_store::{create_store, Trie, COL_STATE};
 use node_runtime::config::RuntimeConfig;
 use node_runtime::{ApplyState, Runtime};
 use std::fs::File;
@@ -64,10 +64,16 @@ impl RuntimeTestbed {
     }
 
     pub fn process_block(&mut self, transactions: &[SignedTransaction], allow_failures: bool) {
-        let state_update = TrieUpdate::new(self.trie.clone(), self.root);
         let apply_result = self
             .runtime
-            .apply(state_update, &self.apply_state, &self.prev_receipts, transactions)
+            .apply(
+                self.trie.clone(),
+                self.root,
+                &None,
+                &self.apply_state,
+                &self.prev_receipts,
+                transactions,
+            )
             .unwrap();
 
         let (store_update, root) = apply_result.trie_changes.into(self.trie.clone()).unwrap();