feat(nns): Add pagination to list_neurons API

rs/nervous_system/integration_tests/src/pocket_ic_helpers.rs

@@ -1,12 +1,9 @@
 use candid::{Decode, Encode, Nat, Principal};
 use canister_test::Wasm;
-use futures::stream;
-use futures::StreamExt;
+use futures::{stream, StreamExt};
 use ic_base_types::{CanisterId, PrincipalId, SubnetId};
 use ic_ledger_core::Tokens;
-use ic_nervous_system_agent::pocketic_impl::PocketIcCallError;
-use ic_nervous_system_agent::sns::Sns;
-use ic_nervous_system_agent::CallCanisters;
+use ic_nervous_system_agent::{pocketic_impl::PocketIcCallError, sns::Sns, CallCanisters};
 use ic_nervous_system_common::{E8, ONE_DAY_SECONDS};
 use ic_nervous_system_common_test_keys::{TEST_NEURON_1_ID, TEST_NEURON_1_OWNER_PRINCIPAL};
 use ic_nns_common::pb::v1::{NeuronId, ProposalId};
@@ -60,17 +57,15 @@ use icrc_ledger_types::icrc1::{
     account::Account,
     transfer::{TransferArg, TransferError},
 };
-use itertools::EitherOrBoth;
-use itertools::Itertools;
+use itertools::{EitherOrBoth, Itertools};
 use maplit::btreemap;
 use pocket_ic::{
     management_canister::CanisterSettings, nonblocking::PocketIc, ErrorCode, PocketIcBuilder,
     UserError, WasmResult,
 };
 use prost::Message;
 use rust_decimal::prelude::ToPrimitive;
-use std::ops::Range;
-use std::{collections::BTreeMap, fmt::Write, time::Duration};
+use std::{collections::BTreeMap, fmt::Write, ops::Range, time::Duration};
 
 pub const STARTING_CYCLES_PER_CANISTER: u128 = 2_000_000_000_000_000;
 
@@ -806,6 +801,7 @@ pub mod nns {
                         include_neurons_readable_by_caller: true,
                         include_empty_neurons_readable_by_caller: None,
                         include_public_neurons_in_full_neurons: None,
+                        start_from_neuron_id: None
                     })
                     .unwrap(),
                 )
@@ -1402,8 +1398,9 @@ pub mod sns {
         use assert_matches::assert_matches;
         use ic_crypto_sha2::Sha256;
         use ic_nervous_system_agent::sns::governance::GovernanceCanister;
-        use ic_sns_governance::governance::UPGRADE_STEPS_INTERVAL_REFRESH_BACKOFF_SECONDS;
-        use ic_sns_governance::pb::v1::get_neuron_response;
+        use ic_sns_governance::{
+            governance::UPGRADE_STEPS_INTERVAL_REFRESH_BACKOFF_SECONDS, pb::v1::get_neuron_response,
+        };
         use pocket_ic::ErrorCode;
 
         pub const EXPECTED_UPGRADE_DURATION_MAX_SECONDS: u64 = 1000;

rs/nns/governance/api/src/ic_nns_governance.pb.v1.rs

@@ -3350,6 +3350,12 @@ pub struct ListNeurons {
     /// existing (unmigrated) callers.
     #[prost(bool, optional, tag = "4")]
     pub include_public_neurons_in_full_neurons: Option<bool>,
+    /// If this is set, it skips all neuron IDs until this value is reached, then returns
+    /// the remaining neurons up until the internal limit is reached.  This is only needed
+    /// if the caller is paginating through the list of neurons, and did not get the full
+    /// list in the first request.
+    #[prost(uint64, optional, tag = "5")]
+    pub start_from_neuron_id: Option<u64>,
 }
 /// A response to a `ListNeurons` request.
 ///
@@ -3368,6 +3374,10 @@ pub struct ListNeuronsResponse {
     /// `ManageNeuron` topic).
     #[prost(message, repeated, tag = "2")]
     pub full_neurons: Vec<Neuron>,
+    /// This is returned if there are more neurons to list.  The caller can repeat the original
+    /// request with `start_from_neuron_id` set to the value of this field to get the next batch.
+    #[prost(uint64, optional, tag = "3")]
+    pub next_start_from_neuron_id: Option<u64>,
 }
 /// A response to "ListKnownNeurons"
 #[derive(candid::CandidType, candid::Deserialize, serde::Serialize, comparable::Comparable)]

rs/nns/governance/canister/governance.did

@@ -439,11 +439,13 @@ type ListNeurons = record {
   neuron_ids : vec nat64;
   include_empty_neurons_readable_by_caller : opt bool;
   include_neurons_readable_by_caller : bool;
+  start_from_neuron_id: opt nat64;
 };
 
 type ListNeuronsResponse = record {
   neuron_infos : vec record { nat64; NeuronInfo };
   full_neurons : vec Neuron;
+  next_start_from_neuron_id: opt nat64;
 };
 
 type ListNodeProviderRewardsRequest = record {

rs/nns/governance/canister/governance_test.did

@@ -438,11 +438,13 @@ type ListNeurons = record {
   neuron_ids : vec nat64;
   include_empty_neurons_readable_by_caller : opt bool;
   include_neurons_readable_by_caller : bool;
+  start_from_neuron_id: opt nat64;
 };
 
 type ListNeuronsResponse = record {
   neuron_infos : vec record { nat64; NeuronInfo };
   full_neurons : vec Neuron;
+  next_start_from_neuron_id: opt nat64;
 };
 
 type ListNodeProviderRewardsRequest = record {

rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto

@@ -2659,6 +2659,11 @@ message ListNeurons {
   // since this feature was added later, it is opt in to avoid confusing
   // existing (unmigrated) callers.
   optional bool include_public_neurons_in_full_neurons = 4;
+  // If this is set, it skips all neuron IDs until this value is reached, then returns
+  // the remaining neurons up until the internal limit is reached.  This is only needed
+  // if the caller is paginating through the list of neurons, and did not get the full
+  // list in the first request.
+  optional fixed64 start_from_neuron_id = 5;
 }
 
 // A response to a `ListNeurons` request.
@@ -2673,6 +2678,9 @@ message ListNeuronsResponse {
   // hot key, or controller or hot key of some followee on the
   // `ManageNeuron` topic).
   repeated Neuron full_neurons = 2;
+  // This is returned if there are more neurons to list.  The caller can repeat the original
+  // request with `start_from_neuron_id` set to the value of this field to get the next batch.
+  optional fixed64 next_start_from_neuron_id = 3;
 }
 
 // A response to "ListKnownNeurons"

rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs

@@ -3968,6 +3968,12 @@ pub struct ListNeurons {
     /// existing (unmigrated) callers.
     #[prost(bool, optional, tag = "4")]
     pub include_public_neurons_in_full_neurons: ::core::option::Option<bool>,
+    /// If this is set, it skips all neuron IDs until this value is reached, then returns
+    /// the remaining neurons up until the internal limit is reached.  This is only needed
+    /// if the caller is paginating through the list of neurons, and did not get the full
+    /// list in the first request.
+    #[prost(fixed64, optional, tag = "5")]
+    pub start_from_neuron_id: ::core::option::Option<u64>,
 }
 /// A response to a `ListNeurons` request.
 ///
@@ -3992,6 +3998,10 @@ pub struct ListNeuronsResponse {
     /// `ManageNeuron` topic).
     #[prost(message, repeated, tag = "2")]
     pub full_neurons: ::prost::alloc::vec::Vec<Neuron>,
+    /// This is returned if there are more neurons to list.  The caller can repeat the original
+    /// request with `start_from_neuron_id` set to the value of this field to get the next batch.
+    #[prost(fixed64, optional, tag = "3")]
+    pub next_start_from_neuron_id: ::core::option::Option<u64>,
 }
 /// A response to "ListKnownNeurons"
 #[derive(

rs/nns/governance/src/governance.rs

@@ -117,7 +117,7 @@ use rust_decimal_macros::dec;
 use std::{
     borrow::Cow,
     cmp::{max, Ordering},
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     convert::{TryFrom, TryInto},
     fmt,
     future::Future,
@@ -143,8 +143,6 @@ pub mod tla;
 
 use crate::storage::with_voting_state_machines_mut;
 #[cfg(feature = "tla")]
-use std::collections::BTreeSet;
-#[cfg(feature = "tla")]
 pub use tla::{
     tla_update_method, InstrumentationState, ToTla, CLAIM_NEURON_DESC, MERGE_NEURONS_DESC,
     SPAWN_NEURONS_DESC, SPAWN_NEURON_DESC, SPLIT_NEURON_DESC, TLA_INSTRUMENTATION_STATE,
@@ -2419,6 +2417,7 @@ impl Governance {
             include_neurons_readable_by_caller,
             include_empty_neurons_readable_by_caller,
             include_public_neurons_in_full_neurons,
+            start_from_neuron_id,
         } = list_neurons;
 
         let include_empty_neurons_readable_by_caller = include_empty_neurons_readable_by_caller
@@ -2455,12 +2454,26 @@ impl Governance {
             neuron_ids.iter().map(|id| NeuronId { id: *id }).collect();
         requested_neuron_ids.append(&mut implicitly_requested_neuron_ids);
 
+        if let Some(start_from_neuron_id) = start_from_neuron_id {
+            requested_neuron_ids = requested_neuron_ids
+                .into_iter()
+                .skip_while(|id| id.id != *start_from_neuron_id)
+                .collect();
+        }
+
         // These will be assembled into the final result.
         let mut neuron_infos = hashmap![];
         let mut full_neurons = vec![];
 
+        const MAX_LIST_NEURONS_COUNT: usize = 500;
+
+        let mut next_start_from_neuron_id = None;
         // Populate the above two neuron collections.
-        for neuron_id in requested_neuron_ids {
+        for (count, neuron_id) in &mut requested_neuron_ids.into_iter().enumerate() {
+            if count >= MAX_LIST_NEURONS_COUNT {
+                next_start_from_neuron_id = Some(neuron_id.id);
+                break;
+            }
             // Ignore when a neuron is not found. It is not guaranteed that a
             // neuron will be found, because some of the elements in
             // requested_neuron_ids are supplied by the caller.
@@ -2494,6 +2507,7 @@ impl Governance {
         ListNeuronsResponse {
             neuron_infos,
             full_neurons,
+            next_start_from_neuron_id,
         }
     }
 
@@ -3869,7 +3883,7 @@ impl Governance {
         match proposal_data {
             None => None,
             Some(pd) => {
-                let caller_neurons: HashSet<NeuronId> =
+                let caller_neurons: BTreeSet<NeuronId> =
                     self.neuron_store.get_neuron_ids_readable_by_caller(*caller);
                 let now = self.env.now();
                 Some(self.proposal_data_to_info(pd, &caller_neurons, now, false))
@@ -3955,7 +3969,7 @@ impl Governance {
     ///   retrieve dropped payloads by calling `get_proposal_info` for
     ///   each proposal of interest.
     pub fn get_pending_proposals(&self, caller: &PrincipalId) -> Vec<ProposalInfo> {
-        let caller_neurons: HashSet<NeuronId> =
+        let caller_neurons: BTreeSet<NeuronId> =
             self.neuron_store.get_neuron_ids_readable_by_caller(*caller);
         let now = self.env.now();
         self.get_pending_proposals_data()
@@ -3993,7 +4007,7 @@ impl Governance {
     fn proposal_data_to_info(
         &self,
         data: &ProposalData,
-        caller_neurons: &HashSet<NeuronId>,
+        caller_neurons: &BTreeSet<NeuronId>,
         now_seconds: u64,
         multi_query: bool,
     ) -> ProposalInfo {
@@ -4025,7 +4039,7 @@ impl Governance {
         /// in `except_from`.
         fn remove_ballots_not_cast_by(
             all_ballots: &HashMap<u64, Ballot>,
-            except_from: &HashSet<NeuronId>,
+            except_from: &BTreeSet<NeuronId>,
         ) -> HashMap<u64, Ballot> {
             let mut ballots = HashMap::new();
             for neuron_id in except_from.iter() {
@@ -4065,7 +4079,7 @@ impl Governance {
     fn proposal_is_visible_to_neurons(
         &self,
         info: &ProposalData,
-        caller_neurons: &HashSet<NeuronId>,
+        caller_neurons: &BTreeSet<NeuronId>,
     ) -> bool {
         // Is 'info' a manage neuron proposal?
         if let Some(ref managed_id) = info.proposal.as_ref().and_then(|x| x.managed_neuron()) {
@@ -4135,7 +4149,7 @@ impl Governance {
         caller: &PrincipalId,
         req: &ListProposalInfo,
     ) -> ListProposalInfoResponse {
-        let caller_neurons: HashSet<NeuronId> =
+        let caller_neurons: BTreeSet<NeuronId> =
             self.neuron_store.get_neuron_ids_readable_by_caller(*caller);
         let exclude_topic: HashSet<i32> = req.exclude_topic.iter().cloned().collect();
         let include_reward_status: HashSet<i32> =

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

@@ -557,6 +557,7 @@ fn list_neurons_benchmark() -> BenchResult {
         include_neurons_readable_by_caller: true,
         include_empty_neurons_readable_by_caller: Some(false),
         include_public_neurons_in_full_neurons: None,
+        start_from_neuron_id: None,
     };
 
     bench_fn(|| {

rs/nns/governance/src/governance/tests/list_neurons.rs

@@ -0,0 +1,87 @@
+use crate::{
+    governance::Governance,
+    pb::v1::{neuron::DissolveState, ListNeurons, NetworkEconomics, Neuron},
+    test_utils::{MockEnvironment, StubCMC, StubIcpLedger},
+};
+use ic_base_types::PrincipalId;
+use ic_nns_common::pb::v1::NeuronId;
+
+#[test]
+fn test_list_neurons_with_paging() {
+    let user_id = PrincipalId::new_user_test_id(100);
+
+    let neurons = (1..1000u64)
+        .map(|id| {
+            let dissolve_state = DissolveState::DissolveDelaySeconds(100);
+            let account = crate::test_utils::test_subaccount_for_neuron_id(id);
+            (
+                id,
+                Neuron {
+                    id: Some(NeuronId::from_u64(id)),
+                    controller: Some(user_id),
+                    account,
+                    dissolve_state: Some(dissolve_state),
+                    // Fill in the rest as needed (stake, maturity, etc.)
+                    ..Default::default()
+                },
+            )
+        })
+        .collect();
+
+    let governance = Governance::new(
+        crate::pb::v1::Governance {
+            neurons,
+            economics: Some(NetworkEconomics {
+                voting_power_economics: Some(Default::default()),
+                ..Default::default()
+            }),
+            ..crate::pb::v1::Governance::default()
+        },
+        Box::new(MockEnvironment::new(Default::default(), 0)),
+        Box::new(StubIcpLedger {}),
+        Box::new(StubCMC {}),
+    );
+
+    let response = governance.list_neurons(
+        &ListNeurons {
+            neuron_ids: vec![],
+            include_neurons_readable_by_caller: true,
+            include_empty_neurons_readable_by_caller: None,
+            include_public_neurons_in_full_neurons: None,
+            start_from_neuron_id: None,
+        },
+        user_id,
+    );
+
+    assert_eq!(response.full_neurons.len(), 500);
+    assert_eq!(response.next_start_from_neuron_id, Some(501));
+
+    let response = governance.list_neurons(
+        &ListNeurons {
+            neuron_ids: vec![],
+            include_neurons_readable_by_caller: true,
+            include_empty_neurons_readable_by_caller: None,
+            include_public_neurons_in_full_neurons: None,
+            start_from_neuron_id: Some(501),
+        },
+        user_id,
+    );
+
+    assert_eq!(response.full_neurons.len(), 499);
+    assert_eq!(response.next_start_from_neuron_id, None);
+
+    // Edge case, just barely fit all the neurons in second response
+    let response = governance.list_neurons(
+        &ListNeurons {
+            neuron_ids: vec![],
+            include_neurons_readable_by_caller: true,
+            include_empty_neurons_readable_by_caller: None,
+            include_public_neurons_in_full_neurons: None,
+            start_from_neuron_id: Some(500),
+        },
+        user_id,
+    );
+
+    assert_eq!(response.full_neurons.len(), 500);
+    assert_eq!(response.next_start_from_neuron_id, None);
+}

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

@@ -24,6 +24,7 @@ use lazy_static::lazy_static;
 use maplit::{btreemap, hashmap};
 use std::convert::TryFrom;
 
+mod list_neurons;
 mod neurons_fund;
 mod node_provider_rewards;
 mod stake_maturity;

rs/nns/governance/src/neuron_store.rs

@@ -31,7 +31,7 @@ use ic_nns_common::pb::v1::{NeuronId, ProposalId};
 use icp_ledger::{AccountIdentifier, Subaccount};
 use std::{
     borrow::Cow,
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashMap},
     fmt::{Debug, Display, Formatter},
     ops::{Bound, Deref, RangeBounds},
 };
@@ -1241,7 +1241,7 @@ impl NeuronStore {
     pub fn get_neuron_ids_readable_by_caller(
         &self,
         principal_id: PrincipalId,
-    ) -> HashSet<NeuronId> {
+    ) -> BTreeSet<NeuronId> {
         with_stable_neuron_indexes(|indexes| {
             indexes
                 .principal()