refactor to move voting into Governance context

rs/nns/governance/src/governance.rs

@@ -135,7 +135,7 @@ mod benches;
 pub mod tla_macros;
 #[cfg(feature = "tla")]
 pub mod tla;
-use crate::voting::cast_vote_and_cascade_follow;
+
 #[cfg(feature = "tla")]
 pub use tla::{
     claim_neuron_desc, split_neuron_desc, tla_update_method, InstrumentationState, ToTla,
@@ -5073,19 +5073,6 @@ impl Governance {
             .expect("NetworkEconomics not present")
     }
 
-    /// Inserts a proposals that has already been validated in the state.
-    ///
-    /// This is a low-level function that makes no verification whatsoever.
-    fn insert_proposal(&mut self, pid: u64, data: ProposalData) {
-        let voting_period_seconds = self.voting_period_seconds()(data.topic());
-        self.closest_proposal_deadline_timestamp_seconds = std::cmp::min(
-            data.proposal_timestamp_seconds + voting_period_seconds,
-            self.closest_proposal_deadline_timestamp_seconds,
-        );
-        self.heap_data.proposals.insert(pid, data);
-        self.process_proposal(pid);
-    }
-
     /// The proposal id of the next proposal.
     fn next_proposal_id(&self) -> u64 {
         // Correctness is based on the following observations:
@@ -5597,17 +5584,17 @@ impl Governance {
         })
         .expect("Proposer not found.");
 
-        // Cast self-vote, including following.
-        cast_vote_and_cascade_follow(
-            &proposal_id,
-            &mut proposal_data.ballots,
-            proposer_id,
-            Vote::Yes,
-            topic,
-            &mut self.neuron_store,
-        );
         // Finally, add this proposal as an open proposal.
-        self.insert_proposal(proposal_num, proposal_data);
+        let voting_period_seconds = self.voting_period_seconds()(proposal_data.topic());
+        self.closest_proposal_deadline_timestamp_seconds = std::cmp::min(
+            proposal_data.proposal_timestamp_seconds + voting_period_seconds,
+            self.closest_proposal_deadline_timestamp_seconds,
+        );
+        self.heap_data.proposals.insert(proposal_num, proposal_data);
+
+        self.cast_vote_and_cascade_follow(proposal_id, *proposer_id, Vote::Yes, topic);
+
+        self.process_proposal(proposal_num);
 
         self.refresh_voting_power(proposer_id);
 
@@ -5823,14 +5810,12 @@ impl Governance {
             ));
         }
 
-        cast_vote_and_cascade_follow(
+        self.cast_vote_and_cascade_follow(
             // Actually update the ballot, including following.
-            proposal_id,
-            &mut proposal.ballots,
-            neuron_id,
+            *proposal_id,
+            *neuron_id,
             vote,
             topic,
-            &mut self.neuron_store,
         );
 
         self.process_proposal(proposal_id.id);

rs/nns/governance/src/governance/benches.rs

@@ -4,14 +4,13 @@ use crate::{
     neuron_store::NeuronStore,
     pb::v1::{
         neuron::Followees, proposal::Action, Ballot, BallotInfo, Governance as GovernanceProto,
-        KnownNeuron, Neuron as NeuronProto, Topic, Vote,
+        KnownNeuron, Neuron as NeuronProto, ProposalData, Topic, Vote,
     },
     temporarily_disable_active_neurons_in_stable_memory,
     temporarily_disable_stable_memory_following_index,
     temporarily_enable_active_neurons_in_stable_memory,
     temporarily_enable_stable_memory_following_index,
     test_utils::{MockEnvironment, StubCMC, StubIcpLedger},
-    voting::cast_vote_and_cascade_follow,
 };
 use canbench_rs::{bench, bench_fn, BenchResult};
 use ic_base_types::PrincipalId;
@@ -47,10 +46,19 @@ enum SetUpStrategy {
 fn set_up<R: Rng>(
     strategy: SetUpStrategy,
     rng: &mut R,
-    neuron_store: &mut NeuronStore,
-    ballots: &mut HashMap<u64, Ballot>,
+    governance: &mut Governance,
     topic: Topic,
 ) -> NeuronId {
+    let neuron_store = &mut governance.neuron_store;
+    governance.heap_data.proposals.insert(
+        1,
+        ProposalData {
+            id: Some(ProposalId { id: 1 }),
+            ..Default::default()
+        },
+    );
+    let ballots = &mut governance.heap_data.proposals.get_mut(&1).unwrap().ballots;
+
     match strategy {
         SetUpStrategy::Centralized { num_neurons } => {
             set_up_centralized(num_neurons, rng, neuron_store, ballots, topic)
@@ -301,21 +309,23 @@ fn set_up_chain<R: Rng>(
 }
 
 fn cast_vote_cascade_helper(strategy: SetUpStrategy, topic: Topic) -> BenchResult {
-    let mut ballots = HashMap::new();
     let mut rng = ChaCha20Rng::seed_from_u64(0);
-    let mut neuron_store = NeuronStore::new(btreemap! {});
-    let neuron_id = set_up(strategy, &mut rng, &mut neuron_store, &mut ballots, topic);
+
+    let governance_proto = GovernanceProto {
+        ..Default::default()
+    };
+    let mut governance = Governance::new(
+        governance_proto,
+        Box::new(MockEnvironment::new(Default::default(), 0)),
+        Box::new(StubIcpLedger {}),
+        Box::new(StubCMC {}),
+    );
+
+    let neuron_id = set_up(strategy, &mut rng, &mut governance, topic);
 
     let proposal_id = ProposalId { id: 1 };
     bench_fn(|| {
-        cast_vote_and_cascade_follow(
-            &proposal_id,
-            &mut ballots,
-            &neuron_id.into(),
-            Vote::Yes,
-            topic,
-            &mut neuron_store,
-        );
+        governance.cast_vote_and_cascade_follow(proposal_id, neuron_id.into(), Vote::Yes, topic);
         // let yes_votes = ballots
         //     .iter()
         //     .filter(|(id, ballot)| ballot.vote == Vote::Yes as i32)

rs/nns/governance/src/governance/tests/mod.rs

@@ -1136,16 +1136,15 @@ mod neuron_archiving_tests {
 
 mod cast_vote_and_cascade_follow {
     use crate::{
-        governance::MIN_DISSOLVE_DELAY_FOR_VOTE_ELIGIBILITY_SECONDS,
+        governance::{Governance, MIN_DISSOLVE_DELAY_FOR_VOTE_ELIGIBILITY_SECONDS},
         neuron::{DissolveStateAndAge, Neuron, NeuronBuilder},
-        neuron_store::NeuronStore,
-        pb::v1::{neuron::Followees, Ballot, Topic, Vote},
-        voting::cast_vote_and_cascade_follow,
+        pb::v1::{neuron::Followees, Ballot, ProposalData, Topic, Vote},
+        test_utils::{MockEnvironment, StubCMC, StubIcpLedger},
     };
     use ic_base_types::PrincipalId;
     use ic_nns_common::pb::v1::{NeuronId, ProposalId};
     use icp_ledger::Subaccount;
-    use maplit::hashmap;
+    use maplit::{btreemap, hashmap};
     use std::collections::{BTreeMap, HashMap};
 
     fn make_ballot(voting_power: u64, vote: Vote) -> Ballot {
@@ -1231,24 +1230,42 @@ mod cast_vote_and_cascade_follow {
         // Add a neuron without a ballot for neuron 6 to follow.
         add_neuron_without_ballot(&mut heap_neurons, 7, vec![1]);
 
-        let mut neuron_store = NeuronStore::new(heap_neurons);
+        let governance_proto = crate::pb::v1::Governance {
+            neurons: heap_neurons
+                .into_iter()
+                .map(|(id, neuron)| (id, neuron.into()))
+                .collect(),
+            proposals: btreemap! {
+                1 => ProposalData {
+                    id: Some(ProposalId {id: 1}),
+                    ballots,
+                    ..Default::default()
+                }
+            },
+            ..Default::default()
+        };
+        let mut governance = Governance::new(
+            governance_proto,
+            Box::new(MockEnvironment::new(Default::default(), 0)),
+            Box::new(StubIcpLedger {}),
+            Box::new(StubCMC {}),
+        );
 
-        cast_vote_and_cascade_follow(
-            &ProposalId { id: 1 },
-            &mut ballots,
-            &NeuronId { id: 1 },
+        governance.cast_vote_and_cascade_follow(
+            ProposalId { id: 1 },
+            NeuronId { id: 1 },
             Vote::Yes,
             topic,
-            &mut neuron_store,
         );
 
         let deciding_voting_power = |neuron_id| {
-            neuron_store
+            governance
+                .neuron_store
                 .with_neuron(&neuron_id, |n| n.deciding_voting_power(now))
                 .unwrap()
         };
         assert_eq!(
-            ballots,
+            governance.heap_data.proposals.get(&1).unwrap().ballots,
             hashmap! {
                 1 => make_ballot(deciding_voting_power(NeuronId { id: 1}), Vote::Yes),
                 2 => make_ballot(deciding_voting_power(NeuronId { id: 2}), Vote::Unspecified),
@@ -1300,24 +1317,42 @@ mod cast_vote_and_cascade_follow {
         // Add a neuron without a ballot for neuron 6 to follow.
         add_neuron_without_ballot(&mut neurons, 7, vec![1]);
 
-        let mut neuron_store = NeuronStore::new(neurons);
+        let governance_proto = crate::pb::v1::Governance {
+            neurons: neurons
+                .into_iter()
+                .map(|(id, neuron)| (id, neuron.into()))
+                .collect(),
+            proposals: btreemap! {
+                1 => ProposalData {
+                    id: Some(ProposalId {id: 1}),
+                    ballots,
+                    ..Default::default()
+                }
+            },
+            ..Default::default()
+        };
+        let mut governance = Governance::new(
+            governance_proto,
+            Box::new(MockEnvironment::new(Default::default(), 0)),
+            Box::new(StubIcpLedger {}),
+            Box::new(StubCMC {}),
+        );
 
-        cast_vote_and_cascade_follow(
-            &ProposalId { id: 1 },
-            &mut ballots,
-            &NeuronId { id: 1 },
+        governance.cast_vote_and_cascade_follow(
+            ProposalId { id: 1 },
+            NeuronId { id: 1 },
             Vote::Yes,
             topic,
-            &mut neuron_store,
         );
 
         let deciding_voting_power = |neuron_id| {
-            neuron_store
+            governance
+                .neuron_store
                 .with_neuron(&neuron_id, |n| n.deciding_voting_power(now))
                 .unwrap()
         };
         assert_eq!(
-            ballots,
+            governance.heap_data.proposals.get(&1).unwrap().ballots,
             hashmap! {
                 1 => make_ballot(deciding_voting_power(NeuronId { id: 1 }), Vote::Yes),
                 2 => make_ballot(deciding_voting_power(NeuronId { id: 2 }), Vote::Yes),

rs/nns/governance/src/voting.rs

@@ -21,24 +21,31 @@ impl Governance {
         voting_neuron_id: NeuronId,
         vote_of_neuron: Vote,
         topic: Topic,
-    ) -> Result<(), GovernanceError> {
-        let mut state_machine =
-            VotingStateMachine::try_new(proposal_id, topic, voting_neuron_id, vote_of_neuron)
-                .map_err(|e| {
-                    GovernanceError::new_with_message(ErrorType::InvalidCommand, format!("{:?}", e))
-                })?;
-
-        let mut ballots;
-        while !state_machine.is_done() {
-            ballots = &mut self
-                .heap_data
-                .proposals
-                .get_mut(&proposal_id.id)
-                .unwrap()
-                .ballots;
-            state_machine.continue_processing(&mut self.neuron_store, ballots);
-        }
-        Ok(())
+    ) {
+        let neuron_store = &mut self.neuron_store;
+        let ballots = &mut self
+            .heap_data
+            .proposals
+            .get_mut(&proposal_id.id)
+            .unwrap()
+            .ballots;
+        // Use of thread local storage to store the state machines prevents
+        // more than one state machine per proposal, which limits the overall
+        // memory usage for voting, which will be relevant when this can be used
+        // across multiple messages, which would cause the memory usage to accumulate.
+        VOTING_STATE_MACHINES.with(|vsm| {
+            let mut voting_state_machines = vsm.borrow_mut();
+            let proposal_voting_machine =
+                voting_state_machines.get_or_create_machine(*proposal_id, topic);
+
+            proposal_voting_machine.cast_vote(ballots, *voting_neuron_id, vote_of_neuron);
+
+            while !proposal_voting_machine.is_done() {
+                proposal_voting_machine.continue_processing(neuron_store, ballots);
+            }
+
+            voting_state_machines.remove_if_done(proposal_id);
+        });
     }
 }