Node family management (#1670)

* Family management messages

* Add family queries

* Add queries to client

* Layer assignment message

* Paged family queries, annotate mixnodes with family

* Add layer assignments to epoch operations

* Remove family layer peristence

* Add NotImplemented error for kick

Co-authored-by: Jędrzej Stuczyński <jedrzej.stuczynski@gmail.com>

validator-api/src/contract_cache/mod.rs

@@ -5,6 +5,7 @@ use crate::nymd_client::Client;
 use crate::storage::ValidatorApiStorage;
 use ::time::OffsetDateTime;
 use anyhow::Result;
+use mixnet_contract_common::families::FamilyHead;
 use mixnet_contract_common::mixnode::MixNodeDetails;
 use mixnet_contract_common::reward_params::{Performance, RewardingParams};
 use mixnet_contract_common::{
@@ -141,7 +142,11 @@ impl<C> ValidatorCacheRefresher<C> {
         interval_reward_params: RewardingParams,
         current_interval: Interval,
         rewarded_set: &HashMap<MixId, RewardedSetNodeStatus>,
+        mix_to_family: Vec<(IdentityKey, FamilyHead)>,
     ) -> Vec<MixNodeBondAnnotated> {
+        let mix_to_family = mix_to_family
+            .into_iter()
+            .collect::<HashMap<IdentityKey, FamilyHead>>();
         let mut annotated = Vec::new();
         for mixnode in mixnodes {
             let stake_saturation = mixnode
@@ -174,6 +179,10 @@ impl<C> ValidatorCacheRefresher<C> {
                     current_interval,
                 );
 
+            let family = mix_to_family
+                .get(&mixnode.bond_information.identity().to_string())
+                .cloned();
+
             annotated.push(MixNodeBondAnnotated {
                 mixnode_details: mixnode,
                 stake_saturation,
@@ -181,6 +190,7 @@ impl<C> ValidatorCacheRefresher<C> {
                 performance,
                 estimated_operator_apy,
                 estimated_delegators_apy,
+                family,
             });
         }
         annotated
@@ -226,10 +236,18 @@ impl<C> ValidatorCacheRefresher<C> {
         let mixnodes = self.nymd_client.get_mixnodes().await?;
         let gateways = self.nymd_client.get_gateways().await?;
 
+        let mix_to_family = self.nymd_client.get_all_family_members().await?;
+
         let rewarded_set = self.get_rewarded_set_map().await;
 
         let mixnodes = self
-            .annotate_node_with_details(mixnodes, rewarding_params, current_interval, &rewarded_set)
+            .annotate_node_with_details(
+                mixnodes,
+                rewarding_params,
+                current_interval,
+                &rewarded_set,
+                mix_to_family,
+            )
             .await;
 
         let (rewarded_set, active_set) =
@@ -322,6 +340,7 @@ impl ValidatorCache {
         })
     }
 
+    #[allow(clippy::too_many_arguments)]
     async fn update_cache(
         &self,
         mixnodes: Vec<MixNodeBondAnnotated>,

validator-api/src/epoch_operations/error.rs

@@ -25,6 +25,11 @@ pub enum RewardingError {
         #[from]
         source: std::num::TryFromIntError,
     },
+    #[error("{source}")]
+    WeightedError {
+        #[from]
+        source: rand::distributions::WeightedError,
+    },
 }
 
 impl From<NymdError> for RewardingError {

validator-api/src/epoch_operations/mod.rs

@@ -19,11 +19,14 @@ use crate::storage::ValidatorApiStorage;
 use mixnet_contract_common::{
     reward_params::Performance, CurrentIntervalResponse, ExecuteMsg, Interval, MixId,
 };
+use mixnet_contract_common::{Layer, LayerAssignment};
 use rand::prelude::SliceRandom;
 use rand::rngs::OsRng;
+use std::collections::HashMap;
 use std::collections::HashSet;
 use std::time::Duration;
 use tokio::time::sleep;
+use validator_api_requests::models::MixNodeBondAnnotated;
 use validator_client::nymd::SigningNymdClient;
 
 pub(crate) mod error;
@@ -32,7 +35,6 @@ mod helpers;
 use crate::epoch_operations::helpers::stake_to_f64;
 use crate::node_status_api::ONE_DAY;
 use error::RewardingError;
-use mixnet_contract_common::mixnode::MixNodeDetails;
 use task::ShutdownListener;
 
 #[derive(Debug, Clone, Copy)]
@@ -57,6 +59,16 @@ pub struct RewardedSetUpdater {
     storage: ValidatorApiStorage,
 }
 
+// Weight of a layer being chose is reciprocal to current count in layer
+fn layer_weight(l: &Layer, layer_assignments: &HashMap<Layer, f32>) -> f32 {
+    let total = layer_assignments.values().fold(0., |acc, i| acc + i);
+    if total == 0. {
+        1.
+    } else {
+        1. - (layer_assignments.get(l).unwrap_or(&0.) / total)
+    }
+}
+
 impl RewardedSetUpdater {
     pub(crate) async fn current_interval_details(
         &self,
@@ -76,23 +88,62 @@ impl RewardedSetUpdater {
         })
     }
 
+    async fn determine_layers(
+        &self,
+        rewarded_set: &[MixNodeBondAnnotated],
+    ) -> Result<(Vec<LayerAssignment>, HashMap<String, Layer>), RewardingError> {
+        let mut families_in_layer: HashMap<String, Layer> = HashMap::new();
+        let mut assignments = vec![];
+        let mut layer_assignments: HashMap<Layer, f32> = HashMap::new();
+        let mut rng = OsRng;
+        let layers = vec![Layer::One, Layer::Two, Layer::Three];
+
+        for mix in rewarded_set {
+            // Get layer already assigned to nodes family, if any
+            let family_layer = mix
+                .family
+                .as_ref()
+                .and_then(|h| families_in_layer.get(h.identity()));
+
+            // Same node families are always assigned to the same layer, otherwise layer selected by a random weighted choice
+            let layer = if let Some(layer) = family_layer {
+                layer.to_owned()
+            } else {
+                layers
+                    .choose_weighted(&mut rng, |l| layer_weight(l, &layer_assignments))?
+                    .to_owned()
+            };
+
+            assignments.push(LayerAssignment::new(mix.mix_id(), layer));
+
+            // layer accounting
+            let layer_entry = layer_assignments.entry(layer).or_insert(0.);
+            *layer_entry += 1.;
+            if let Some(ref family) = mix.family {
+                families_in_layer.insert(family.identity().to_string(), layer);
+            }
+        }
+
+        Ok((assignments, families_in_layer))
+    }
+
     fn determine_rewarded_set(
         &self,
-        mixnodes: Vec<MixNodeDetails>,
+        mixnodes: &[MixNodeBondAnnotated],
         nodes_to_select: u32,
-    ) -> Vec<MixId> {
+    ) -> Result<Vec<MixNodeBondAnnotated>, RewardingError> {
         if mixnodes.is_empty() {
-            return Vec::new();
+            return Ok(Vec::new());
         }
 
         let mut rng = OsRng;
 
         // generate list of mixnodes and their relatively weight (by total stake)
         let choices = mixnodes
-            .into_iter()
+            .iter()
             .map(|mix| {
-                let total_stake = stake_to_f64(mix.total_stake());
-                (mix.mix_id(), total_stake)
+                let total_stake = stake_to_f64(mix.mixnode_details.total_stake());
+                (mix.to_owned(), total_stake)
             })
             .collect::<Vec<_>>();
 
@@ -101,11 +152,10 @@ impl RewardedSetUpdater {
         // - we have invalid weights, i.e. less than zero or NaNs - it shouldn't happen in our case as we safely cast down from u128
         // - all weights are zero - it's impossible in our case as the list of nodes is not empty and weight is proportional to stake. You must have non-zero stake in order to bond
         // - we have more than u32::MAX values (which is incredibly unrealistic to have 4B mixnodes bonded... literally every other person on the planet would need one)
-        choices
-            .choose_multiple_weighted(&mut rng, nodes_to_select as usize, |item| item.1)
-            .unwrap()
-            .map(|(mix_id, _weight)| *mix_id)
-            .collect()
+        Ok(choices
+            .choose_multiple_weighted(&mut rng, nodes_to_select as usize, |item| item.1)?
+            .map(|(mix, _weight)| mix.to_owned())
+            .collect())
     }
 
     async fn reward_current_rewarded_set(
@@ -186,15 +236,19 @@ impl RewardedSetUpdater {
 
     async fn update_rewarded_set_and_advance_epoch(
         &self,
-        all_mixnodes: Vec<MixNodeDetails>,
+        all_mixnodes: &[MixNodeBondAnnotated],
     ) -> Result<(), RewardingError> {
         // we grab rewarding parameters here as they might have gotten updated when performing epoch actions
         let rewarding_parameters = self.nymd_client.get_current_rewarding_parameters().await?;
 
         let new_rewarded_set =
-            self.determine_rewarded_set(all_mixnodes, rewarding_parameters.rewarded_set_size);
+            self.determine_rewarded_set(all_mixnodes, rewarding_parameters.rewarded_set_size)?;
+
+        let (layer_assignments, _families_in_layer) =
+            self.determine_layers(&new_rewarded_set).await?;
+
         self.nymd_client
-            .advance_current_epoch(new_rewarded_set, rewarding_parameters.active_set_size)
+            .advance_current_epoch(layer_assignments, rewarding_parameters.active_set_size)
             .await?;
 
         Ok(())
@@ -219,15 +273,11 @@ impl RewardedSetUpdater {
 
         let epoch_end = interval.current_epoch_end();
 
-        let all_nodes = self.validator_cache.mixnodes().await;
-        if all_nodes.is_empty() {
+        let all_mixnodes = self.validator_cache.mixnodes_detailed().await;
+        if all_mixnodes.is_empty() {
             log::warn!("there don't seem to be any mixnodes on the network!")
         }
 
-        // get list of all mixnodes BEFORE rewarding happens as to now be biased by rewards
-        // that might be given to them
-        let all_mixnodes = self.validator_cache.mixnodes().await;
-
         // Reward all the nodes in the still current, soon to be previous rewarded set
         log::info!("Rewarding the current rewarded set...");
         if let Err(err) = self.reward_current_rewarded_set(interval).await {
@@ -257,7 +307,7 @@ impl RewardedSetUpdater {
 
         log::info!("Advancing epoch and updating the rewarded set...");
         if let Err(err) = self
-            .update_rewarded_set_and_advance_epoch(all_mixnodes)
+            .update_rewarded_set_and_advance_epoch(&all_mixnodes)
             .await
         {
             log::error!("FAILED to advance the current epoch... - {}", err);

validator-api/src/nymd_client.rs

@@ -4,10 +4,12 @@
 use crate::config::Config;
 use crate::epoch_operations::MixnodeToReward;
 use config::defaults::{NymNetworkDetails, DEFAULT_VALIDATOR_API_PORT};
+use mixnet_contract_common::families::{Family, FamilyHead};
 use mixnet_contract_common::mixnode::MixNodeDetails;
 use mixnet_contract_common::reward_params::RewardingParams;
 use mixnet_contract_common::{
-    CurrentIntervalResponse, ExecuteMsg, GatewayBond, MixId, RewardedSetNodeStatus,
+    CurrentIntervalResponse, ExecuteMsg, GatewayBond, IdentityKey, LayerAssignment, MixId,
+    RewardedSetNodeStatus,
 };
 use std::sync::Arc;
 use tokio::sync::RwLock;
@@ -191,6 +193,23 @@ impl<C> Client<C> {
             .await
     }
 
+    #[allow(dead_code)]
+    pub(crate) async fn get_all_node_families(&self) -> Result<Vec<Family>, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync + Send,
+    {
+        self.0.read().await.get_all_node_families().await
+    }
+
+    pub(crate) async fn get_all_family_members(
+        &self,
+    ) -> Result<Vec<(IdentityKey, FamilyHead)>, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync + Send,
+    {
+        self.0.read().await.get_all_family_members().await
+    }
+
     pub(crate) async fn send_rewarding_messages(
         &self,
         nodes: &[MixnodeToReward],
@@ -237,7 +256,7 @@ impl<C> Client<C> {
 
     pub(crate) async fn advance_current_epoch(
         &self,
-        new_rewarded_set: Vec<MixId>,
+        new_rewarded_set: Vec<LayerAssignment>,
         expected_active_set_size: u32,
     ) -> Result<(), ValidatorClientError>
     where