fix(epoch manager): Fix fishermen unstake (#2212)

* Fix fishermen unstake

* fix kickout set

* add test

chain/epoch_manager/src/lib.rs

@@ -2437,6 +2437,46 @@ mod tests {
         assert_eq!(em.get_epoch_info(&EpochId(CryptoHash::default())).unwrap(), &epoch_info)
     }
 
+    #[test]
+    fn test_fishermen_unstake() {
+        let stake_amount = 1_000;
+        let fishermen_threshold = 100;
+        let validators = vec![
+            ("test1", stake_amount),
+            ("test2", fishermen_threshold),
+            ("test3", fishermen_threshold),
+        ];
+        let mut em = setup_epoch_manager(
+            validators,
+            2,
+            1,
+            1,
+            0,
+            90,
+            70,
+            fishermen_threshold,
+            default_reward_calculator(),
+        );
+        let h = hash_range(5);
+        record_block(&mut em, CryptoHash::default(), h[0], 0, vec![]);
+        // fishermen unstake
+        record_block(&mut em, h[0], h[1], 1, vec![stake("test2", 0)]);
+        record_block(&mut em, h[1], h[2], 2, vec![stake("test3", 1)]);
+        assert_eq!(
+            em.get_epoch_info(&EpochId(h[2])).unwrap(),
+            &epoch_info(
+                vec![("test1", stake_amount)],
+                vec![0],
+                vec![vec![0]],
+                vec![],
+                vec![],
+                change_stake(vec![("test1", stake_amount), ("test2", 0), ("test3", 0)]),
+                reward(vec![("test1", 0), ("near", 0)]),
+                0
+            )
+        );
+    }
+
     #[test]
     fn test_validator_consistency() {
         let stake_amount = 1_000;

chain/epoch_manager/src/proposals.rs

@@ -75,7 +75,9 @@ pub fn proposals_to_epoch_info(
     }
 
     for r in epoch_info.fishermen.iter() {
-        if !ordered_proposals.contains_key(&r.account_id) {
+        if !ordered_proposals.contains_key(&r.account_id)
+            && !validator_kickout.contains(&r.account_id)
+        {
             // safe to do this here because fishermen from previous epoch is guaranteed to have no
             // duplicates.
             fishermen_to_index.insert(r.account_id.clone(), fishermen.len() as ValidatorId);
@@ -113,7 +115,9 @@ pub fn proposals_to_epoch_info(
                     }
                 })
                 .or_insert((0, p.stake));
-            if epoch_info.validator_to_index.contains_key(&account_id) {
+            if epoch_info.validator_to_index.contains_key(&account_id)
+                || epoch_info.fishermen_to_index.contains_key(&account_id)
+            {
                 validator_kickout.insert(account_id);
             }
         }
@@ -166,7 +170,9 @@ pub fn proposals_to_epoch_info(
             fishermen.push(p);
         } else {
             stake_change.insert(p.account_id.clone(), (0, p.stake));
-            if epoch_info.validator_to_index.contains_key(&p.account_id) {
+            if epoch_info.validator_to_index.contains_key(&p.account_id)
+                || epoch_info.fishermen_to_index.contains_key(&p.account_id)
+            {
                 validator_kickout.insert(p.account_id);
             }
         }
@@ -298,4 +304,62 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn test_fishermen_allocation() {
+        // 4 proposals of stake 10, fishermen threshold 10 --> 1 validator and 3 fishermen
+        assert_eq!(
+            proposals_to_epoch_info(
+                &epoch_config(2, 2, 1, 0, 90, 60, 10),
+                [0; 32],
+                &EpochInfo::default(),
+                vec![
+                    stake("test1", 10),
+                    stake("test2", 10),
+                    stake("test3", 10),
+                    stake("test4", 10)
+                ],
+                HashSet::default(),
+                HashMap::default(),
+                0
+            )
+            .unwrap(),
+            epoch_info(
+                vec![("test1", 10)],
+                vec![0],
+                vec![vec![0], vec![0]],
+                vec![],
+                vec![("test2", 10), ("test3", 10), ("test4", 10)],
+                change_stake(vec![("test1", 10), ("test2", 10), ("test3", 10), ("test4", 10)]),
+                HashMap::default(),
+                0
+            )
+        );
+
+        // 4 proposals of stake 9, fishermen threshold 10 --> 1 validator and 0 fishermen
+        let mut epoch_info = epoch_info(
+            vec![("test1", 9)],
+            vec![0],
+            vec![vec![0], vec![0]],
+            vec![],
+            vec![],
+            change_stake(vec![("test1", 9), ("test2", 0), ("test3", 0), ("test4", 0)]),
+            HashMap::default(),
+            0,
+        );
+        epoch_info.validator_kickout = HashSet::default();
+        assert_eq!(
+            proposals_to_epoch_info(
+                &epoch_config(2, 2, 1, 0, 90, 60, 10),
+                [0; 32],
+                &EpochInfo::default(),
+                vec![stake("test1", 9), stake("test2", 9), stake("test3", 9), stake("test4", 9)],
+                HashSet::default(),
+                HashMap::default(),
+                0
+            )
+            .unwrap(),
+            epoch_info
+        );
+    }
 }

near/src/runtime.rs

@@ -2337,6 +2337,60 @@ mod test {
         assert!(response.current_fishermen.is_empty());
     }
 
+    /// Test that when fishermen unstake they get their tokens back.
+    #[test]
+    fn test_fishermen_unstake() {
+        init_test_logger();
+        let num_nodes = 2;
+        let validators = (0..num_nodes).map(|i| format!("test{}", i + 1)).collect::<Vec<_>>();
+        let mut env = TestEnv::new(
+            "test_validator_stake_change_multiple_times",
+            vec![validators.clone()],
+            2,
+            vec![],
+            vec![],
+            false,
+        );
+        let block_producers: Vec<_> = validators
+            .iter()
+            .map(|id| InMemoryValidatorSigner::from_seed(id, KeyType::ED25519, id))
+            .collect();
+        let signers: Vec<_> = validators
+            .iter()
+            .map(|id| InMemorySigner::from_seed(id, KeyType::ED25519, id))
+            .collect();
+
+        let staking_transaction = stake(1, &signers[0], &block_producers[0], FISHERMEN_THRESHOLD);
+        env.step_default(vec![staking_transaction]);
+        for _ in 2..9 {
+            env.step_default(vec![]);
+        }
+
+        let account0 = env.view_account(&block_producers[0].validator_id());
+        assert_eq!(account0.locked, FISHERMEN_THRESHOLD);
+        assert_eq!(account0.amount, TESTING_INIT_BALANCE - FISHERMEN_THRESHOLD);
+        let response = env.runtime.get_validator_info(&env.head.last_block_hash).unwrap();
+        assert_eq!(
+            response
+                .current_fishermen
+                .into_iter()
+                .map(|fishermen| fishermen.account_id)
+                .collect::<Vec<_>>(),
+            vec!["test1"]
+        );
+        let staking_transaction = stake(2, &signers[0], &block_producers[0], 0);
+        env.step_default(vec![staking_transaction]);
+        for _ in 10..17 {
+            env.step_default(vec![]);
+        }
+
+        let account0 = env.view_account(&block_producers[0].validator_id());
+        assert_eq!(account0.locked, 0);
+        assert_eq!(account0.amount, TESTING_INIT_BALANCE);
+        let response = env.runtime.get_validator_info(&env.head.last_block_hash).unwrap();
+        assert!(response.current_fishermen.is_empty());
+    }
+
     /// Enable reward and make sure that validators get reward proportional to their stake.
     #[test]
     fn test_validator_reward() {