From e3ee74ee38e99bee081262bad6b58d95bbd499e5 Mon Sep 17 00:00:00 2001
From: Michael Birch <michael.birch@aurora.dev>
Date: Fri, 12 Apr 2024 19:01:48 +0200
Subject: [PATCH] Fix: handle case where total stake is very small

---
 .../epoch-manager/src/validator_selection.rs  |  8 +++---
 .../src/validator_mandates/compute_price.rs   | 27 +++++++++++++++++--
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/chain/epoch-manager/src/validator_selection.rs b/chain/epoch-manager/src/validator_selection.rs
index d7f60c04a3b..6e2d23a1b7a 100644
--- a/chain/epoch-manager/src/validator_selection.rs
+++ b/chain/epoch-manager/src/validator_selection.rs
@@ -838,10 +838,10 @@ mod tests {
         // Given `epoch_info` and `proposals` above, the sample at a given height is deterministic.
         let height = 42;
         let expected_assignments = vec![
-            vec![(1, 300), (0, 300), (2, 300), (3, 60)],
-            vec![(0, 600), (2, 200), (1, 200)],
-            vec![(3, 200), (2, 300), (1, 100), (0, 400)],
-            vec![(2, 200), (4, 140), (1, 400), (0, 200)],
+            vec![(4, 56), (1, 168), (2, 300), (3, 84), (0, 364)],
+            vec![(3, 70), (1, 300), (4, 42), (2, 266), (0, 308)],
+            vec![(4, 42), (1, 238), (3, 42), (0, 450), (2, 196)],
+            vec![(2, 238), (1, 294), (3, 64), (0, 378)],
         ];
         assert_eq!(epoch_info.sample_chunk_validators(height), expected_assignments);
     }
diff --git a/core/primitives/src/validator_mandates/compute_price.rs b/core/primitives/src/validator_mandates/compute_price.rs
index 1c392ac1828..7e37ece2206 100644
--- a/core/primitives/src/validator_mandates/compute_price.rs
+++ b/core/primitives/src/validator_mandates/compute_price.rs
@@ -1,4 +1,8 @@
-use {super::ValidatorMandatesConfig, near_primitives_core::types::Balance, std::cmp::Ordering};
+use {
+    super::ValidatorMandatesConfig,
+    near_primitives_core::types::Balance,
+    std::cmp::{min, Ordering},
+};
 
 /// Given the stakes for the validators and the target number of mandates to have,
 /// this function computes the mandate price to use. It works by iterating a
@@ -18,7 +22,13 @@ where
 {
     let ValidatorMandatesConfig { target_mandates_per_shard, num_shards } = config;
     let total_stake = saturating_sum(stakes());
-    let target_mandates: u128 = num_shards.saturating_mul(target_mandates_per_shard) as u128;
+
+    // The target number of mandates cannot be larger than the total amount of stake.
+    // In production the total stake is _much_ higher than
+    // `num_shards * target_mandates_per_shard`, but in tests validators are given
+    // low staked numbers, so we need to have this condition in place.
+    let target_mandates: u128 =
+        min(num_shards.saturating_mul(target_mandates_per_shard) as u128, total_stake);
 
     let initial_price = total_stake / target_mandates;
 
@@ -96,6 +106,19 @@ mod tests {
 
     use super::*;
 
+    // Test case where the target number of mandates is larger than the total stake.
+    // This should never happen in production, but nearcore tests sometimes have
+    // low stake.
+    #[test]
+    fn test_small_total_stake() {
+        let stakes = [100_u128; 1];
+        let num_shards = 1;
+        let target_mandates_per_shard = 1000;
+        let config = ValidatorMandatesConfig::new(target_mandates_per_shard, num_shards);
+
+        assert_eq!(compute_mandate_price(config, || stakes.iter().copied()), 1);
+    }
+
     // Test cases where all stakes are equal.
     #[test]
     fn test_constant_dist() {