Cross-shard transaction tests and fixes

- New tests to cover cross-shard txs and receipts
- Fixing receipts storage after applying transactions: they need to be stored indexed by the source shard, not the destination shard

chain/chain/src/chain.rs

@@ -777,11 +777,10 @@ impl<'a> ChainUpdate<'a> {
                     self.chain_store_update.save_post_state_root(&chunk_hash, &state_root);
                     // Save resulting receipts.
                     for (receipt_shard_id, receipts) in new_receipts.drain() {
-                        self.chain_store_update.save_receipt(
-                            &block.hash(),
-                            receipt_shard_id,
-                            receipts,
-                        );
+                        // The receipts in store are indexed by the SOURCE shard_id, not destination,
+                        //    since they are later retrieved by the chunk producer of the source
+                        //    shard to be distributed to the recipients.
+                        self.chain_store_update.save_receipt(&block.hash(), shard_id, receipts);
                     }
                     // Save receipt and transaction results.
                     for (i, tx_result) in tx_results.drain(..).enumerate() {

chain/chain/src/test_utils.rs

@@ -1,15 +1,16 @@
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use chrono::Utc;
 
+use log::info;
 use near_primitives::crypto::signature::{PublicKey, Signature};
 use near_primitives::crypto::signer::InMemorySigner;
 use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::rpc::{AccountViewCallResult, QueryResponse};
 use near_primitives::test_utils::get_public_key_from_seed;
 use near_primitives::transaction::{
-    ReceiptTransaction, SignedTransaction, TransactionResult, TransactionStatus,
+    AsyncCall, ReceiptTransaction, SignedTransaction, TransactionResult, TransactionStatus,
 };
 use near_primitives::types::{AccountId, BlockIndex, MerkleHash, ShardId};
 use near_store::test_utils::create_test_store;
@@ -18,7 +19,10 @@ use near_store::{Store, StoreUpdate, Trie, TrieChanges, WrappedTrieChanges};
 use crate::error::{Error, ErrorKind};
 use crate::types::{BlockHeader, ReceiptResult, RuntimeAdapter, Weight};
 use crate::{Chain, ValidTransaction};
+use near_primitives::merkle::merklize;
 use near_primitives::sharding::ShardChunkHeader;
+use near_primitives::transaction::ReceiptBody::NewCall;
+use near_primitives::transaction::TransactionBody::SendMoney;
 
 /// Simple key value runtime for tests.
 pub struct KeyValueRuntime {
@@ -27,6 +31,13 @@ pub struct KeyValueRuntime {
     root: MerkleHash,
     validators: Vec<(AccountId, PublicKey)>,
     validators_per_shard: u64,
+
+    // A mapping state_root => {account id => amounts}, for transactions and receipts
+    amounts: RwLock<HashMap<MerkleHash, HashMap<AccountId, u128>>>,
+}
+
+pub fn account_id_to_shard_id(account_id: &AccountId, num_shards: ShardId) -> ShardId {
+    ((hash(&account_id.clone().into_bytes()).0).0[0] as u64) % num_shards
 }
 
 impl KeyValueRuntime {
@@ -40,6 +51,13 @@ impl KeyValueRuntime {
         validators_per_shard: u64,
     ) -> Self {
         let trie = Arc::new(Trie::new(store.clone()));
+        let mut initial_amounts = HashMap::new();
+        for (i, validator) in validators.iter().enumerate() {
+            initial_amounts.insert(validator.clone(), (1000 + 100 * i) as u128);
+        }
+
+        let mut amounts = HashMap::new();
+        amounts.insert(MerkleHash::default(), initial_amounts);
         KeyValueRuntime {
             store,
             trie,
@@ -49,6 +67,7 @@ impl KeyValueRuntime {
                 .map(|account_id| (account_id.clone(), get_public_key_from_seed(account_id)))
                 .collect(),
             validators_per_shard,
+            amounts: RwLock::new(amounts),
         }
     }
 
@@ -98,7 +117,7 @@ impl RuntimeAdapter for KeyValueRuntime {
         height: BlockIndex,
         shard_id: ShardId,
     ) -> Result<AccountId, Box<dyn std::error::Error>> {
-        assert!((self.validators.len() as u64) % self.num_shards() == 0);
+        assert_eq!((self.validators.len() as u64) % self.num_shards(), 0);
         let validators_per_shard = self.validators_per_shard;
         let offset = (shard_id / validators_per_shard * validators_per_shard) as usize;
         // The +1 is so that if all validators validate all shards in a test, the chunk producer
@@ -128,7 +147,7 @@ impl RuntimeAdapter for KeyValueRuntime {
     }
 
     fn account_id_to_shard_id(&self, account_id: &AccountId) -> ShardId {
-        ((hash(&account_id.clone().into_bytes()).0).0[0] as u64) % self.num_shards()
+        account_id_to_shard_id(account_id, self.num_shards())
     }
 
     fn get_part_owner(
@@ -145,7 +164,7 @@ impl RuntimeAdapter for KeyValueRuntime {
         _height: BlockIndex,
         shard_id: ShardId,
     ) -> bool {
-        assert!((self.validators.len() as u64) % self.num_shards() == 0);
+        assert_eq!((self.validators.len() as u64) % self.num_shards(), 0);
         let validators_per_shard = self.validators_per_shard;
         let offset = (shard_id / validators_per_shard * validators_per_shard) as usize;
         for validator in self.validators[offset..offset + (validators_per_shard as usize)].iter() {
@@ -167,44 +186,151 @@ impl RuntimeAdapter for KeyValueRuntime {
 
     fn apply_transactions(
         &self,
-        _shard_id: ShardId,
+        shard_id: ShardId,
         state_root: &MerkleHash,
         _block_index: BlockIndex,
         _prev_block_hash: &CryptoHash,
-        _receipts: &Vec<ReceiptTransaction>,
+        receipts: &Vec<ReceiptTransaction>,
         transactions: &Vec<SignedTransaction>,
     ) -> Result<
         (WrappedTrieChanges, MerkleHash, Vec<TransactionResult>, ReceiptResult),
         Box<dyn std::error::Error>,
     > {
         let mut tx_results = vec![];
-        for _ in transactions {
-            tx_results.push(TransactionResult {
-                status: TransactionStatus::Completed,
-                logs: vec![],
-                receipts: vec![],
-                result: None,
-            });
+
+        let mut accounts_mapping =
+            self.amounts.read().unwrap().get(state_root).cloned().unwrap_or_else(|| HashMap::new());
+
+        let mut balance_transfers = vec![];
+
+        for receipt in receipts.iter() {
+            if let NewCall(call) = &receipt.body {
+                info!(
+                    "MOO applying receipt in shard {} from {} to {} amt {}",
+                    shard_id, receipt.originator, receipt.receiver, call.amount
+                );
+                assert_eq!(self.account_id_to_shard_id(&receipt.receiver), shard_id);
+                balance_transfers.push((
+                    receipt.originator.clone(),
+                    receipt.receiver.clone(),
+                    call.amount,
+                ));
+            } else {
+                unreachable!();
+            }
         }
+
+        for transaction in transactions {
+            if let SendMoney(send_money_tx) = &transaction.body {
+                info!(
+                    "MOO applying transaction from {} to {} amt {}",
+                    send_money_tx.originator, send_money_tx.receiver, send_money_tx.amount
+                );
+                assert_eq!(self.account_id_to_shard_id(&send_money_tx.originator), shard_id);
+                balance_transfers.push((
+                    send_money_tx.originator.clone(),
+                    send_money_tx.receiver.clone(),
+                    send_money_tx.amount,
+                ));
+            } else {
+                unreachable!();
+            }
+        }
+
+        let mut new_receipts = HashMap::new();
+
+        for (from, to, amount) in balance_transfers {
+            let mut good_to_go = false;
+
+            if self.account_id_to_shard_id(&from) != shard_id {
+                // This is a receipt, was already debited
+                good_to_go = true;
+            } else if let Some(balance) = accounts_mapping.get(&from) {
+                if *balance >= amount {
+                    let new_balance = balance - amount;
+                    accounts_mapping.insert(from.clone(), new_balance);
+                    good_to_go = true;
+                }
+            }
+
+            if good_to_go {
+                let new_receipt_hashes = if self.account_id_to_shard_id(&to) == shard_id {
+                    accounts_mapping
+                        .insert(to.clone(), accounts_mapping.get(&to).unwrap_or(&0) + amount);
+                    vec![]
+                } else {
+                    info!(
+                        "MOO creating receipt from shard {} from {} to {} amt {}",
+                        shard_id, from, to, amount
+                    );
+                    let receipt = ReceiptTransaction::new(
+                        from.clone(),
+                        to,
+                        CryptoHash::default(),
+                        NewCall(AsyncCall::new(vec![], vec![], amount, from)),
+                    );
+                    let receipt_hash = receipt.get_hash();
+                    new_receipts.entry(receipt.shard_id()).or_insert_with(|| vec![]).push(receipt);
+                    vec![receipt_hash]
+                };
+
+                tx_results.push(TransactionResult {
+                    status: TransactionStatus::Completed,
+                    logs: vec![],
+                    receipts: new_receipt_hashes,
+                    result: None,
+                });
+            }
+        }
+
+        let mut new_balances = vec![];
+        for (name, _) in self.validators.iter() {
+            let mut seen = false;
+            for (key, value) in accounts_mapping.iter() {
+                if key == name {
+                    assert!(!seen);
+                    seen = true;
+                    new_balances.push(*value);
+                }
+            }
+            if !seen {
+                new_balances.push(0);
+            }
+        }
+        let (new_state_root, _) = merklize(&new_balances);
+        self.amounts.write().unwrap().insert(new_state_root, accounts_mapping);
+
+        info!(
+            "MOO Applied transactions in shard {}, new state root {}, balances {:?}",
+            shard_id, new_state_root, new_balances
+        );
+
         Ok((
             WrappedTrieChanges::new(self.trie.clone(), TrieChanges::empty(state_root.clone())),
-            *state_root,
+            new_state_root,
             tx_results,
-            HashMap::default(),
+            new_receipts,
         ))
     }
 
     fn query(
         &self,
-        _state_root: MerkleHash,
+        state_root: MerkleHash,
         _height: BlockIndex,
         path: Vec<&str>,
         _data: &[u8],
     ) -> Result<QueryResponse, Box<dyn std::error::Error>> {
+        let account_id = path[1].to_string();
+        let account_id2 = account_id.clone();
         Ok(QueryResponse::ViewAccount(AccountViewCallResult {
-            account_id: path[1].to_string(),
+            account_id,
             nonce: 0,
-            amount: 1000,
+            amount: self
+                .amounts
+                .read()
+                .unwrap()
+                .get(&state_root)
+                .map_or_else(|| 0, |mapping| *mapping.get(&account_id2).unwrap_or(&0)),
             stake: 0,
             public_keys: vec![],
             code_hash: CryptoHash::default(),

chain/client/Cargo.toml

@@ -21,3 +21,6 @@ near-chain = { path = "../chain" }
 near-network = { path = "../network" }
 near-pool = { path = "../pool" }
 near-chunks = { path = "../chunks" }
+
+[features]
+expensive_tests = []

chain/client/src/client.rs

@@ -520,7 +520,7 @@ impl ClientActor {
 
         let transactions =
             self.shards_mgr.prepare_transactions(shard_id, self.config.block_expected_weight)?;
-        debug!("Creating a chunk with {} transactions for shard {}", transactions.len(), shard_id);
+        info!("Creating a chunk with {} transactions for shard {}", transactions.len(), shard_id);
 
         let mut receipts = vec![];
         let mut receipts_block_hash = prev_block_hash.clone();