feat(consensus):Utilize priority function in test framework CON-275

Author: ninegua

rs/consensus/tests/framework/driver.rs

@@ -3,7 +3,7 @@ use ic_artifact_pool::{
     certification_pool::CertificationPoolImpl, consensus_pool::ConsensusPoolImpl, dkg_pool,
 };
 use ic_config::artifact_pool::ArtifactPoolConfig;
-use ic_consensus::consensus::ConsensusImpl;
+use ic_consensus::consensus::{ConsensusGossipImpl, ConsensusImpl};
 use ic_interfaces::{
     artifact_pool::{ChangeSetProducer, MutablePool},
     certification,
@@ -27,6 +27,7 @@ impl<'a> ConsensusDriver<'a> {
         node_id: NodeId,
         pool_config: ArtifactPoolConfig,
         consensus: ConsensusImpl,
+        consensus_gossip: ConsensusGossipImpl,
         dkg: ic_consensus::dkg::DkgImpl,
         certifier: Box<
             dyn ChangeSetProducer<CertificationPoolImpl, ChangeSet = certification::ChangeSet> + 'a,
@@ -42,15 +43,19 @@ impl<'a> ConsensusDriver<'a> {
             logger.clone(),
             metrics_registry,
         )));
+        let consensus_priority =
+            PriorityFnState::new(&consensus_gossip, &*consensus_pool.read().unwrap());
         ConsensusDriver {
             consensus,
+            consensus_gossip,
             dkg,
             certifier,
             logger,
             consensus_pool,
             certification_pool,
             ingress_pool,
             dkg_pool,
+            consensus_priority,
         }
     }
 

rs/consensus/tests/framework/execution.rs

@@ -5,13 +5,14 @@ use ic_interfaces::{
 };
 use ic_logger::{trace, ReplicaLogger};
 use ic_test_utilities::types::ids::node_test_id;
-use ic_types::time::Time;
+use ic_types::{artifact::Priority, time::Time};
 use rand::seq::SliceRandom;
 use std::time::Duration;
 
 fn execute_instance<'a, 'b>(
     instance: &'b ConsensusInstance<'a>,
     time_source: &dyn TimeSource,
+    use_priority_fn: bool,
     logger: &ReplicaLogger,
 ) -> Option<Time> {
     let mut in_queue = instance.in_queue.borrow_mut();
@@ -33,6 +34,24 @@ fn execute_instance<'a, 'b>(
         match inp {
             Input::Message(x) => match x.message {
                 InputMessage::Consensus(msg) => {
+                    if use_priority_fn {
+                        match instance
+                            .driver
+                            .consensus_priority
+                            .borrow()
+                            .get_priority(&msg)
+                        {
+                            Priority::Drop => return Some(timestamp),
+                            Priority::Stash | Priority::Later => {
+                                instance
+                                    .buffered
+                                    .borrow_mut()
+                                    .push(InputMessage::Consensus(msg));
+                                return Some(timestamp);
+                            }
+                            Priority::Fetch | Priority::FetchNow => (),
+                        };
+                    }
                     let mut pool = instance.driver.consensus_pool.write().unwrap();
                     pool.insert(UnvalidatedArtifact {
                         message: msg,
@@ -58,9 +77,25 @@ fn execute_instance<'a, 'b>(
                 }
             },
             // Repeat the polling
-            Input::TimerExpired(x) => in_queue.push(Input::TimerExpired(
-                x + Duration::from_millis(POLLING_INTERVAL),
-            )),
+            Input::TimerExpired(x) => {
+                if use_priority_fn {
+                    let mut priority = instance.driver.consensus_priority.borrow_mut();
+                    if priority.last_updated + PRIORITY_FN_REFRESH_INTERVAL < timestamp {
+                        priority.refresh(
+                            &instance.driver.consensus_gossip,
+                            &*instance.driver.consensus_pool.read().unwrap(),
+                            timestamp,
+                        );
+                    }
+                    let mut buffered = instance.buffered.borrow_mut();
+                    for message in buffered.drain(..) {
+                        in_queue.push(Input::Message(Message { message, timestamp }));
+                    }
+                }
+                in_queue.push(Input::TimerExpired(
+                    x + Duration::from_millis(POLLING_INTERVAL),
+                ));
+            }
         }
         // Move new messages into out_queue.
         for message in instance.driver.step(time_source) {
@@ -76,11 +111,13 @@ fn execute_instance<'a, 'b>(
 /// timestamp(min(i)) value globally. This ensures that all input messages are
 /// always executed in order, for all nodes.
 #[derive(Debug)]
-pub struct GlobalMessage;
+pub struct GlobalMessage {
+    use_priority_fn: bool,
+}
 
 impl GlobalMessage {
-    pub fn new() -> Box<GlobalMessage> {
-        Box::new(GlobalMessage)
+    pub fn new(use_priority_fn: bool) -> Box<GlobalMessage> {
+        Box::new(GlobalMessage { use_priority_fn })
     }
 }
 
@@ -95,16 +132,20 @@ impl ExecutionStrategy for GlobalMessage {
                 let t_j = j.in_queue.borrow().peek().map(|x| x.timestamp());
                 compare_timestamp(t_i, t_j)
             })
-            .and_then(|instance| execute_instance(instance, runner.time_source(), logger))
+            .and_then(|instance| {
+                execute_instance(instance, runner.time_source(), self.use_priority_fn, logger)
+            })
     }
 }
 
 #[derive(Debug)]
-pub struct RandomExecute;
+pub struct RandomExecute {
+    use_priority_fn: bool,
+}
 
 impl RandomExecute {
-    pub fn new() -> Box<RandomExecute> {
-        Box::new(RandomExecute)
+    pub fn new(use_priority_fn: bool) -> Box<RandomExecute> {
+        Box::new(RandomExecute { use_priority_fn })
     }
 }
 
@@ -115,7 +156,8 @@ impl ExecutionStrategy for RandomExecute {
         let mut rng = runner.rng();
         instances.shuffle(&mut *rng);
         while let Some(instance) = instances.pop() {
-            let result = execute_instance(instance, runner.time_source(), logger);
+            let result =
+                execute_instance(instance, runner.time_source(), self.use_priority_fn, logger);
             if result.is_some() {
                 return result;
             }
@@ -125,11 +167,13 @@ impl ExecutionStrategy for RandomExecute {
 }
 
 #[derive(Debug)]
-pub struct GlobalClock;
+pub struct GlobalClock {
+    use_priority_fn: bool,
+}
 
 impl GlobalClock {
-    pub fn new() -> Box<GlobalClock> {
-        Box::new(GlobalClock)
+    pub fn new(use_priority_fn: bool) -> Box<GlobalClock> {
+        Box::new(GlobalClock { use_priority_fn })
     }
 }
 
@@ -144,6 +188,8 @@ impl ExecutionStrategy for GlobalClock {
                 let t_j = j.in_queue.borrow().peek().map(|_| *j.clock.borrow());
                 compare_timestamp(t_i, t_j)
             })
-            .and_then(|instance| execute_instance(instance, runner.time_source(), logger))
+            .and_then(|instance| {
+                execute_instance(instance, runner.time_source(), self.use_priority_fn, logger)
+            })
     }
 }

rs/consensus/tests/framework/runner.rs

@@ -5,7 +5,6 @@ use ic_config::artifact_pool::ArtifactPoolConfig;
 use ic_consensus::consensus::dkg_key_manager::DkgKeyManager;
 use ic_consensus::{
     certification::{CertificationCrypto, CertifierImpl},
-    consensus::ConsensusImpl,
     dkg,
 };
 use ic_consensus_utils::crypto::ConsensusCrypto;
@@ -141,7 +140,7 @@ impl<'a> ConsensusRunner<'a> {
         let fake_local_store_certified_time_reader =
             Arc::new(FakeLocalStoreCertifiedTimeReader::new(self.time.clone()));
 
-        let consensus = ConsensusImpl::new(
+        let (consensus, consensus_gossip) = ic_consensus::consensus::setup(
             deps.replica_config.clone(),
             Default::default(),
             Arc::clone(&deps.registry_client),
@@ -157,11 +156,11 @@ impl<'a> ConsensusRunner<'a> {
             deps.message_routing.clone(),
             deps.state_manager.clone(),
             Arc::clone(&self.time) as Arc<_>,
-            Duration::from_secs(0),
             MaliciousFlags::default(),
             deps.metrics_registry.clone(),
             replica_logger.clone(),
             fake_local_store_certified_time_reader,
+            0,
         );
         let dkg = dkg::DkgImpl::new(
             deps.replica_config.node_id,
@@ -188,12 +187,14 @@ impl<'a> ConsensusRunner<'a> {
             node_id,
             deps,
             in_queue,
+            buffered: Default::default(),
             out_queue: Default::default(),
 
             driver: ConsensusDriver::new(
                 node_id,
                 pool_config,
                 consensus,
+                consensus_gossip,
                 dkg,
                 Box::new(certifier),
                 deps.consensus_pool.clone(),
@@ -271,7 +272,8 @@ impl Default for ConsensusRunnerConfig {
             num_nodes: 10,
             num_rounds: 20,
             degree: 9,
-            execution: GlobalMessage::new(),
+            use_priority_fn: false,
+            execution: GlobalMessage::new(false),
             delivery: Sequential::new(),
         }
     }
@@ -325,6 +327,7 @@ impl ConsensusRunnerConfig {
         config.num_rounds = rng.gen_range(10..101);
         config.degree = rng
             .gen_range(std::cmp::min(5, config.num_nodes / 2)..std::cmp::min(config.num_nodes, 20));
+        config.use_priority_fn = rng.gen_bool(0.5);
         config.reset_strategies();
         config
     }
@@ -340,9 +343,9 @@ impl ConsensusRunnerConfig {
     fn strategies<R: Rng>(&self, rng: &mut R) -> Strategies {
         (
             vec![
-                GlobalMessage::new(),
-                GlobalClock::new(),
-                RandomExecute::new(),
+                GlobalMessage::new(self.use_priority_fn),
+                GlobalClock::new(self.use_priority_fn),
+                RandomExecute::new(self.use_priority_fn),
             ],
             vec![
                 Sequential::new(),
@@ -353,8 +356,8 @@ impl ConsensusRunnerConfig {
     }
 
     /// Parse and update configuration from environment: NUM_NODES,
-    /// NUM_ROUNDS, MAX_DELTA, DEGREE, EXECUTION and DELIVERY (except
-    /// RANDOM_SEED, which should be used when first creating the config).
+    /// NUM_ROUNDS, MAX_DELTA, DEGREE, USE_PRIORITY_FN, EXECUTION and DELIVERY
+    /// (except RANDOM_SEED, which should be used when first creating the config).
     /// Return the updated config if parsing is successful, or an error message
     /// in string otherwise.
     pub fn parse_extra_config(mut self) -> Result<Self, String> {
@@ -376,6 +379,11 @@ impl ConsensusRunnerConfig {
                         .parse()
                         .map_err(|_| "DEGREE must be an unsigned integer")?
                 }
+                "use_priority_fn" => {
+                    self.use_priority_fn = value
+                        .parse()
+                        .map_err(|_| "USE_PRIORITY_FN must be either true or false")?
+                }
                 _ => (),
             }
         }