feat: support partial mandates in chunk validator assignment (#10241)

Adds support for partial mandates in chunk validator assignment.

Currently only full mandates are considered. Assume `stake_per_mandate =
5`, then:

- Valdiator `V1` with stake 12 holds 2 full mandates corresponding to a
stake of 10. The remaining stake of 2 does not participate in
validation.
- Validator `V2` with stake 4 holds 0 full mandates and does not
participate in validation at all.

With support for partial mandates `V1` gets a partial mandate with
weight 2 and `V2` gets a partial mandate of 4. So partial mandates allow
to increase the number of validators and allows them to commit their
entire stake to validation.

### Probability of shard corruption

In [these
simulations](https://near.zulipchat.com/#narrow/stream/407237-pagoda.2Fcore.2Fstateless-validation/topic/chunk.20validator.20assignment.20simulation.20with.20partial.20seats),
the probability of shard corruption decreases as partial seats are
included. In general, the effect of including partial seats depends on
the distribution of (malicious) stake across validators and the stake
required per mandate.

### `ValidatorMandatesAssignment`

This type is introduced to more clearly document the result of sampling
chunk validators and to avoid passing around a raw
`Vec<HashMap<ValidatorId, _>>`.

### Shuffling of shard ids

[This
thread](mooori/sim-validator-assignment#12 (comment))
points out a bias towards small shard ids and brings up shard id
shuffling as a way to fix it. For now shard ids are not shuffled in this
PR, since I think it would make the diff harder to review. I would
suggest to introduce shard id shuffling in a separate PR. This should
not pose a risk as chunk validator assignment is not enabled in
production.

chain/epoch-manager/src/validator_selection.rs

@@ -391,6 +391,8 @@ mod tests {
     use near_primitives::epoch_manager::ValidatorSelectionConfig;
     use near_primitives::shard_layout::ShardLayout;
     use near_primitives::types::validator_stake::ValidatorStake;
+    #[cfg(feature = "protocol_feature_chunk_validation")]
+    use near_primitives::validator_mandates::{AssignmentWeight, ValidatorMandatesAssignment};
     use near_primitives::version::PROTOCOL_VERSION;
     use num_rational::Ratio;
 
@@ -658,13 +660,27 @@ mod tests {
             ValidatorSelectionConfig {
                 num_chunk_only_producer_seats: 0,
                 minimum_validators_per_shard: 1,
-                minimum_stake_ratio: Ratio::new(160, 1_000_000),
+                // for example purposes, we choose a higher ratio than in production
+                minimum_stake_ratio: Ratio::new(1, 10),
             },
         );
         let prev_epoch_height = 7;
         let prev_epoch_info = create_prev_epoch_info(prev_epoch_height, &["test1", "test2"], &[]);
-        let proposals =
-            create_proposals(&[("test1", 15), ("test2", 9), ("test3", 5), ("test4", 3)]);
+
+        // Choosing proposals s.t. the `threshold` (i.e. seat price) calculated in
+        // `proposals_to_epoch_info` below will be 100. For now, this `threshold` is used as the
+        // stake required for a chunk validator mandate.
+        //
+        // Note that `proposals_to_epoch_info` will not include `test6` in the set of validators,
+        // hence it will not hold a (partial) mandate
+        let proposals = create_proposals(&[
+            ("test1", 1500),
+            ("test2", 1000),
+            ("test3", 1000),
+            ("test4", 260),
+            ("test5", 140),
+            ("test6", 50),
+        ]);
 
         let epoch_info = proposals_to_epoch_info(
             &epoch_config,
@@ -681,11 +697,30 @@ mod tests {
 
         // Given `epoch_info` and `proposals` above, the sample at a given height is deterministic.
         let height = 42;
-        let expected_assignments: Vec<HashMap<ValidatorId, u16>> = vec![
-            HashMap::from([(0, 1), (1, 5), (2, 1), (3, 1)]),
-            HashMap::from([(0, 6), (1, 1), (2, 1)]),
-            HashMap::from([(0, 5), (1, 2), (2, 1)]),
-            HashMap::from([(0, 3), (1, 1), (2, 2), (3, 2)]),
+        let expected_assignments: ValidatorMandatesAssignment = vec![
+            HashMap::from([
+                (0, AssignmentWeight::new(2, 0)),
+                (1, AssignmentWeight::new(4, 0)),
+                (2, AssignmentWeight::new(2, 0)),
+                (3, AssignmentWeight::new(1, 60)),
+                (4, AssignmentWeight::new(1, 0)),
+            ]),
+            HashMap::from([
+                (0, AssignmentWeight::new(4, 0)),
+                (1, AssignmentWeight::new(2, 0)),
+                (2, AssignmentWeight::new(4, 0)),
+                (4, AssignmentWeight::new(0, 40)),
+            ]),
+            HashMap::from([
+                (0, AssignmentWeight::new(7, 0)),
+                (1, AssignmentWeight::new(1, 0)),
+                (3, AssignmentWeight::new(1, 0)),
+            ]),
+            HashMap::from([
+                (0, AssignmentWeight::new(2, 0)),
+                (1, AssignmentWeight::new(3, 0)),
+                (2, AssignmentWeight::new(4, 0)),
+            ]),
         ];
         assert_eq!(epoch_info.sample_chunk_validators(height), expected_assignments);
     }

core/primitives/src/epoch_manager.rs

@@ -522,7 +522,7 @@ pub mod epoch_info {
     use crate::epoch_manager::ValidatorWeight;
     use crate::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
     use crate::types::{BlockChunkValidatorStats, ValidatorKickoutReason};
-    use crate::validator_mandates::ValidatorMandates;
+    use crate::validator_mandates::{ValidatorMandates, ValidatorMandatesAssignment};
     use crate::version::PROTOCOL_VERSION;
     use borsh::{BorshDeserialize, BorshSerialize};
     use near_primitives_core::hash::CryptoHash;
@@ -1101,10 +1101,7 @@ pub mod epoch_info {
             }
         }
 
-        pub fn sample_chunk_validators(
-            &self,
-            height: BlockHeight,
-        ) -> Vec<HashMap<ValidatorId, u16>> {
+        pub fn sample_chunk_validators(&self, height: BlockHeight) -> ValidatorMandatesAssignment {
             // Chunk validator assignment was introduced with `V4`.
             match &self {
                 Self::V1(_) => Default::default(),

core/primitives/src/types.rs

@@ -687,6 +687,23 @@ pub mod validator_stake {
             u16::try_from(self.stake() / stake_per_mandate)
                 .expect("number of mandats should fit u16")
         }
+
+        /// Returns the weight attributed to the validator's partial mandate.
+        ///
+        /// A validator has a partial mandate if its stake cannot be divided evenly by
+        /// `stake_per_mandate`. The remainder of that division is the weight of the partial
+        /// mandate.
+        ///
+        /// Due to this definintion a validator has exactly one partial mandate with `0 <= weight <
+        /// stake_per_mandate`.
+        ///
+        /// # Example
+        ///
+        /// Let `V` be a validator with stake of 12. If `stake_per_mandate` equals 5 then the weight
+        /// of `V`'s partial mandate is `12 % 5 = 2`.
+        pub fn partial_mandate_weight(&self, stake_per_mandate: Balance) -> Balance {
+            self.stake() % stake_per_mandate
+        }
     }
 }
 
@@ -1029,3 +1046,33 @@ pub struct StateChangesForShard {
     pub shard_id: ShardId,
     pub state_changes: Vec<RawStateChangesWithTrieKey>,
 }
+
+#[cfg(test)]
+mod tests {
+    use near_crypto::{KeyType, PublicKey};
+    use near_primitives_core::types::Balance;
+
+    use super::validator_stake::ValidatorStake;
+
+    fn new_validator_stake(stake: Balance) -> ValidatorStake {
+        ValidatorStake::new(
+            "test_account".parse().unwrap(),
+            PublicKey::empty(KeyType::ED25519),
+            stake,
+        )
+    }
+
+    #[test]
+    fn test_validator_stake_num_mandates() {
+        assert_eq!(new_validator_stake(0).num_mandates(5), 0);
+        assert_eq!(new_validator_stake(10).num_mandates(5), 2);
+        assert_eq!(new_validator_stake(12).num_mandates(5), 2);
+    }
+
+    #[test]
+    fn test_validator_partial_mandate_weight() {
+        assert_eq!(new_validator_stake(0).partial_mandate_weight(5), 0);
+        assert_eq!(new_validator_stake(10).partial_mandate_weight(5), 0);
+        assert_eq!(new_validator_stake(12).partial_mandate_weight(5), 2);
+    }
+}