Avoid duplicate calls by caching WTN proposals per NNS proposal (#33)

CodeGov-org · Nov 15, 2024 · e848e14 · e848e14
1 parent f854ba4
commit e848e14
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 25 deletions.
diff --git a/canister/src/jobs/process_votes.rs b/canister/src/jobs/process_votes.rs
@@ -32,39 +32,70 @@ async fn process_vote(vote: VoteToProcess) {
 
     match vote {
         VoteToProcess::NnsVote(pair_id, nns_vote) => {
-            let canister_id = state::read(|s| s.wtn_protocol_canister_id());
-            let response = get_wtn_proposal_id(canister_id, nns_vote.proposal_id).await;
-            let vote_to_process = match response {
-                Ok(Ok(wtn_proposal_id)) => Some(VoteToProcess::PendingWtnVote(
+            let vote_to_process = match state::read(|s| {
+                s.get_cached_wtn_proposal_for_nns_proposal(nns_vote.proposal_id)
+            }) {
+                Some(Some(wtn_proposal_id)) => Some(VoteToProcess::PendingWtnVote(
                     pair_id,
                     WtnVote {
                         nns_proposal_id: nns_vote.proposal_id,
-                        wtn_proposal_id: wtn_proposal_id.id,
+                        wtn_proposal_id,
                         adopt: nns_vote.adopt,
                     },
                 )),
-                Ok(Err(latest_processed_nns_proposal_id)) => {
-                    if latest_processed_nns_proposal_id.id >= nns_vote.proposal_id {
-                        log(format!(
-                            "No WTN proposal found for NNS proposal {}",
-                            nns_vote.proposal_id
-                        ));
-                        None
-                    } else {
-                        // The WTN canister hasn't processed this NNS proposal yet, so put the NNS
-                        // proposal back in the queue for it to be attempted again shortly
-                        log(format!(
-                            "WTN canister has not processed NNS proposal yet. ProposalId: {}. Latest processed: {}",
-                            nns_vote.proposal_id,
-                            latest_processed_nns_proposal_id.id
-                        ));
-                        Some(VoteToProcess::NnsVote(pair_id, nns_vote))
+                Some(None) => None,
+                None => {
+                    // Didn't find the WTN proposal in the cache, so call
+                    // into WTN canister to retrieve it
+                    let canister_id = state::read(|s| s.wtn_protocol_canister_id());
+                    let response = get_wtn_proposal_id(canister_id, nns_vote.proposal_id).await;
+                    match response {
+                        Ok(Ok(wtn_proposal_id)) => {
+                            state::mutate(|s| {
+                                s.record_wtn_proposal_for_nns_proposal(
+                                    nns_vote.proposal_id,
+                                    Some(wtn_proposal_id.id),
+                                )
+                            });
+                            Some(VoteToProcess::PendingWtnVote(
+                                pair_id,
+                                WtnVote {
+                                    nns_proposal_id: nns_vote.proposal_id,
+                                    wtn_proposal_id: wtn_proposal_id.id,
+                                    adopt: nns_vote.adopt,
+                                },
+                            ))
+                        }
+                        Ok(Err(latest_processed_nns_proposal_id)) => {
+                            if latest_processed_nns_proposal_id.id >= nns_vote.proposal_id {
+                                state::mutate(|s| {
+                                    s.record_wtn_proposal_for_nns_proposal(
+                                        nns_vote.proposal_id,
+                                        None,
+                                    )
+                                });
+                                log(format!(
+                                    "No WTN proposal found for NNS proposal {}",
+                                    nns_vote.proposal_id
+                                ));
+                                None
+                            } else {
+                                // The WTN canister hasn't processed this NNS proposal yet, so put the NNS
+                                // proposal back in the queue for it to be attempted again shortly
+                                log(format!(
+                                    "WTN canister has not processed NNS proposal yet. ProposalId: {}. Latest processed: {}",
+                                    nns_vote.proposal_id,
+                                    latest_processed_nns_proposal_id.id
+                                ));
+                                Some(VoteToProcess::NnsVote(pair_id, nns_vote))
+                            }
+                        }
+                        Err(error) => {
+                            ic_cdk::eprintln!("Error calling `get_wtn_proposal_id`: {error:?}");
+                            Some(VoteToProcess::NnsVote(pair_id, nns_vote))
+                        }
                     }
                 }
-                Err(error) => {
-                    ic_cdk::eprintln!("Error calling `get_wtn_proposal_id`: {error:?}");
-                    Some(VoteToProcess::NnsVote(pair_id, nns_vote))
-                }
             };
             if let Some(vote) = vote_to_process {
                 state::mutate(|s| s.push_vote_to_process(vote));

diff --git a/canister/src/state.rs b/canister/src/state.rs
@@ -25,6 +25,8 @@ pub struct State {
     wtn_protocol_canister_id: Principal,
     neuron_pairs: BTreeMap<u64, NeuronPair>,
     votes_to_process: VecDeque<VoteToProcess>,
+    #[serde(default)]
+    cached_wtn_proposals_per_nns_proposal: BTreeMap<u64, Option<u64>>,
 }
 
 const STATE_ALREADY_INITIALIZED: &str = "State has already been initialized";
@@ -66,6 +68,7 @@ impl State {
                 .unwrap_or(DEFAULT_WTN_PROTOCOL_CANISTER_ID),
             neuron_pairs: BTreeMap::new(),
             votes_to_process: VecDeque::new(),
+            cached_wtn_proposals_per_nns_proposal: BTreeMap::new(),
         }
     }
 
@@ -153,6 +156,28 @@ impl State {
     pub fn votes_to_process_count(&self) -> usize {
         self.votes_to_process.len()
     }
+
+    pub fn record_wtn_proposal_for_nns_proposal(
+        &mut self,
+        nns_proposal_id: u64,
+        wtn_proposal_id: Option<u64>,
+    ) {
+        self.cached_wtn_proposals_per_nns_proposal
+            .insert(nns_proposal_id, wtn_proposal_id);
+
+        while self.cached_wtn_proposals_per_nns_proposal.len() > 500 {
+            self.cached_wtn_proposals_per_nns_proposal.pop_first();
+        }
+    }
+
+    pub fn get_cached_wtn_proposal_for_nns_proposal(
+        &self,
+        nns_proposal_id: u64,
+    ) -> Option<Option<u64>> {
+        self.cached_wtn_proposals_per_nns_proposal
+            .get(&nns_proposal_id)
+            .copied()
+    }
 }
 
 #[cfg(test)]