Skip to main content

ic_query/nns/neuron/report/
collection.rs

1//! Module: nns::neuron::report::collection
2//!
3//! Responsibility: advance caller-persisted complete public-neuron walks one page at a time.
4//! Does not own: stable memory, filesystem caches, scheduling, retries, or publication.
5//! Boundary: validates resumable input and derives the next state from one admitted source page.
6
7#[cfg(feature = "nns-host")]
8use super::NnsNeuronHostError;
9use super::{
10    NnsNeuronError,
11    model::{NnsNeuronListReport, NnsNeuronListRequest},
12    source::{NnsNeuronSource, build_nns_neuron_list_report_with_source, validate_page_size},
13};
14#[cfg(feature = "nns-host")]
15use crate::nns::LiveNnsSource;
16use crate::nns::governance::{
17    GovernanceCollectionState, NnsGovernanceCollectionStatus, NnsGovernanceRequest,
18    NnsGovernanceSourceProvenance, NnsGovernanceSourceSelection, validate_governance_continuation,
19    validate_governance_request,
20};
21use serde::{Deserialize, Serialize};
22
23/// Version of the persistable resumable NNS neuron collection state.
24pub const NNS_NEURON_COLLECTION_STATE_SCHEMA_VERSION: u32 = 1;
25
26///
27/// NnsNeuronCollectionState
28///
29/// Serializable continuation state for an explicitly bounded public-neuron walk.
30///
31
32#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
33pub struct NnsNeuronCollectionState {
34    #[serde(flatten)]
35    collection: GovernanceCollectionState,
36    neurons_fetched: u64,
37    next_start_neuron_id: Option<u64>,
38}
39
40impl NnsNeuronCollectionState {
41    /// Start an empty neuron walk with explicit per-page and cumulative call ceilings.
42    pub fn new(
43        request: &NnsGovernanceRequest,
44        page_size: u32,
45        max_pages: u32,
46    ) -> Result<Self, NnsNeuronError> {
47        validate_governance_request(request)?;
48        validate_page_size(page_size)?;
49        if max_pages == 0 {
50            return Err(NnsNeuronError::InvalidCollectionMaxPages);
51        }
52        Ok(Self {
53            collection: GovernanceCollectionState::new(request, page_size, max_pages),
54            neurons_fetched: 0,
55            next_start_neuron_id: None,
56        })
57    }
58
59    #[cfg(feature = "nns-host")]
60    pub(in crate::nns) const fn from_complete_collection(
61        collection: GovernanceCollectionState,
62        rows: u64,
63    ) -> Self {
64        Self {
65            collection,
66            neurons_fetched: rows,
67            next_start_neuron_id: None,
68        }
69    }
70
71    #[cfg(feature = "nns-host")]
72    pub(in crate::nns) fn collection_evidence(
73        &self,
74    ) -> crate::nns::governance::collection::NnsGovernanceCollectionEvidence {
75        crate::nns::governance::collection::NnsGovernanceCollectionEvidence::from_state(
76            &self.collection,
77        )
78    }
79
80    /// Return the state schema version.
81    #[must_use]
82    pub const fn schema_version(&self) -> u32 {
83        self.collection.schema_version
84    }
85
86    /// Return the fixed network identity.
87    #[must_use]
88    pub fn network(&self) -> &str {
89        &self.collection.network
90    }
91
92    /// Return the fixed Governance canister identity.
93    #[must_use]
94    pub fn governance_canister_id(&self) -> &str {
95        &self.collection.governance_canister_id
96    }
97
98    /// Return the source selection fixed when collection started.
99    #[must_use]
100    pub const fn requested_source(&self) -> &NnsGovernanceSourceSelection {
101        &self.collection.requested_source
102    }
103
104    /// Return the concrete source provenance after the first admitted page.
105    #[must_use]
106    pub const fn source(&self) -> Option<&NnsGovernanceSourceProvenance> {
107        self.collection.source.as_ref()
108    }
109
110    /// Return the maximum rows requested from each Governance call.
111    #[must_use]
112    pub const fn page_size(&self) -> u32 {
113        self.collection.page_size
114    }
115
116    /// Return the cumulative source-call ceiling.
117    #[must_use]
118    pub const fn max_pages(&self) -> u32 {
119        self.collection.max_pages
120    }
121
122    /// Return the number of successfully admitted pages.
123    #[must_use]
124    pub const fn pages_fetched(&self) -> u32 {
125        self.collection.pages_fetched
126    }
127
128    /// Return the number of successfully admitted neuron rows.
129    #[must_use]
130    pub const fn neurons_fetched(&self) -> u64 {
131        self.neurons_fetched
132    }
133
134    /// Return the exclusive lower neuron-id bound for the next page.
135    #[must_use]
136    pub const fn next_start_neuron_id(&self) -> Option<u64> {
137        self.next_start_neuron_id
138    }
139
140    /// Return the caller-supplied time at which the collection state was created.
141    #[must_use]
142    pub fn started_at(&self) -> &str {
143        &self.collection.started_at
144    }
145
146    /// Return the caller-supplied time attached to the latest admitted page.
147    #[must_use]
148    pub fn updated_at(&self) -> &str {
149        &self.collection.updated_at
150    }
151
152    /// Return the collection lifecycle status.
153    #[must_use]
154    pub const fn status(&self) -> NnsGovernanceCollectionStatus {
155        self.collection.status
156    }
157
158    /// Return whether Governance API exhaustion was observed.
159    #[must_use]
160    pub const fn is_complete(&self) -> bool {
161        matches!(
162            self.collection.status,
163            NnsGovernanceCollectionStatus::Complete
164        )
165    }
166}
167
168///
169/// NnsNeuronCollectionStep
170///
171/// One admitted bounded page and the continuation state that follows it.
172///
173
174#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
175pub struct NnsNeuronCollectionStep {
176    /// Page returned by the shared public-neuron report builder.
177    pub page: NnsNeuronListReport,
178    /// Validated state to persist only after retaining the page.
179    pub state: NnsNeuronCollectionState,
180}
181
182/// Advance a resumable public-neuron walk through the native replica adapter.
183#[cfg(feature = "nns-host")]
184pub async fn advance_nns_neuron_collection(
185    request: &NnsGovernanceRequest,
186    state: &NnsNeuronCollectionState,
187) -> Result<NnsNeuronCollectionStep, NnsNeuronHostError> {
188    Ok(
189        advance_nns_neuron_collection_with_source(request, state, &LiveNnsSource::default())
190            .await?,
191    )
192}
193
194/// Advance a resumable public-neuron walk by exactly one caller-runtime source call.
195pub async fn advance_nns_neuron_collection_with_source(
196    request: &NnsGovernanceRequest,
197    state: &NnsNeuronCollectionState,
198    source: &dyn NnsNeuronSource,
199) -> Result<NnsNeuronCollectionStep, NnsNeuronError> {
200    validate_governance_request(request)?;
201    validate_collection_state(state)?;
202    validate_governance_continuation!(request, state, NnsNeuronError)?;
203
204    let mut page_request = NnsNeuronListRequest::new(request.clone(), state.collection.page_size);
205    page_request.exclusive_start_neuron_id = state.next_start_neuron_id;
206    let page = build_nns_neuron_list_report_with_source(&page_request, source).await?;
207    if let Some(expected) = &state.collection.source
208        && *expected != page.context.source
209    {
210        return Err(NnsNeuronError::CollectionSourceChanged {
211            expected: expected.clone(),
212            actual: page.context.source,
213        });
214    }
215
216    let (pages_fetched, neurons_fetched) = crate::nns::governance::advance_collection_counts(
217        state.collection.pages_fetched,
218        state.neurons_fetched,
219        u64::try_from(page.returned_neuron_count)
220            .map_err(|_| NnsNeuronError::CollectionAccountingOverflow)?,
221    )
222    .ok_or(NnsNeuronError::CollectionAccountingOverflow)?;
223    let next_start_neuron_id = page.next_start_neuron_id;
224    let next_state = NnsNeuronCollectionState {
225        collection: state.collection.after_page(
226            request,
227            page.context.source.clone(),
228            pages_fetched,
229            next_start_neuron_id.is_some(),
230        ),
231        neurons_fetched,
232        next_start_neuron_id,
233    };
234    Ok(NnsNeuronCollectionStep {
235        page,
236        state: next_state,
237    })
238}
239
240pub(super) fn validate_collection_state(
241    state: &NnsNeuronCollectionState,
242) -> Result<(), NnsNeuronError> {
243    let invalid = |reason| NnsNeuronError::InvalidCollectionState { reason };
244    state
245        .collection
246        .validate(NNS_NEURON_COLLECTION_STATE_SCHEMA_VERSION, invalid)?;
247    validate_page_size(state.collection.page_size)?;
248
249    let neurons_fetched = state.neurons_fetched;
250    let maximum_rows = crate::nns::governance::validate_governance_collection_row_bounds(
251        "neurons_fetched",
252        neurons_fetched,
253        state.collection.pages_fetched,
254        state.collection.page_size,
255        invalid,
256    )?;
257
258    let valid_lifecycle = crate::nns::governance::collection_lifecycle_matches(
259        state.collection.status,
260        state.collection.pages_fetched,
261        state.collection.max_pages,
262        state.neurons_fetched,
263        maximum_rows,
264        state.next_start_neuron_id.is_some(),
265        state.collection.source.is_some(),
266    ) && (state.collection.status != NnsGovernanceCollectionStatus::Complete
267        || neurons_fetched < maximum_rows);
268    if !valid_lifecycle {
269        return Err(invalid(format!(
270            "status {} disagrees with cursor, provenance, or counters",
271            state.collection.status
272        )));
273    }
274    Ok(())
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use crate::nns::governance::{NnsGovernanceError, validate_governance_time_interval};
281
282    struct UnexpectedSource;
283    impl NnsNeuronSource for UnexpectedSource {
284        fn fetch_neuron_page<'a>(
285            &'a self,
286            _: &'a NnsGovernanceRequest,
287            _: Option<u64>,
288            _: u32,
289        ) -> crate::nns::neuron::NnsNeuronSourceFuture<'a, crate::nns::neuron::NnsNeuronPage>
290        {
291            panic!("unexpected source call")
292        }
293        fn fetch_neuron<'a>(
294            &'a self,
295            _: &'a NnsGovernanceRequest,
296            _: u64,
297        ) -> crate::nns::neuron::NnsNeuronSourceFuture<'a, crate::nns::neuron::NnsNeuronRow>
298        {
299            panic!("unexpected source call")
300        }
301    }
302
303    #[tokio::test(flavor = "current_thread")]
304    async fn collection_rejects_backward_continuations_and_restored_intervals() {
305        let start = NnsGovernanceRequest::replicated_inter_canister_call_from_unix_secs("ic", 100);
306        let state = NnsNeuronCollectionState::new(&start, 2, 3).unwrap();
307        let before = NnsGovernanceRequest::replicated_inter_canister_call_from_unix_secs("ic", 99);
308        let mut future = std::pin::pin!(advance_nns_neuron_collection_with_source(
309            &before,
310            &state,
311            &UnexpectedSource
312        ));
313        let mut context = std::task::Context::from_waker(std::task::Waker::noop());
314        let result = std::future::Future::poll(future.as_mut(), &mut context);
315        assert!(matches!(
316            result,
317            std::task::Poll::Ready(Err(NnsNeuronError::Governance(
318                NnsGovernanceError::InvalidTimestampOrder { .. }
319            )))
320        ));
321        let mut value = serde_json::to_value(&state).unwrap();
322        value["updated_at"] = serde_json::json!(before.fetched_at);
323        let restored = serde_json::from_value(value).unwrap();
324        assert!(matches!(
325            validate_collection_state(&restored),
326            Err(NnsNeuronError::Governance(
327                NnsGovernanceError::InvalidTimestampOrder { .. }
328            ))
329        ));
330        assert!(
331            validate_governance_time_interval(
332                "updated_at",
333                &state.collection.updated_at,
334                "fetched_at",
335                &start.fetched_at
336            )
337            .is_ok()
338        );
339    }
340
341    #[test]
342    fn collection_validates_request_and_restored_timestamps() {
343        let mut request = NnsGovernanceRequest::replicated_inter_canister_call_from_unix_secs(
344            "ic",
345            1_700_000_000,
346        );
347        let valid_state = NnsNeuronCollectionState::new(&request, 2, 3).expect("valid state");
348        for (field, value) in [
349            ("governance_canister_id", serde_json::json!("aaaaa-aa")),
350            ("max_pages", serde_json::json!(0)),
351            ("pages_fetched", serde_json::json!(4)),
352        ] {
353            let mut restored = serde_json::to_value(&valid_state).unwrap();
354            restored[field] = value;
355            let restored = serde_json::from_value(restored).unwrap();
356            assert!(
357                matches!(
358                    validate_collection_state(&restored),
359                    Err(NnsNeuronError::InvalidCollectionState { .. })
360                ),
361                "invalid restored {field}"
362            );
363        }
364        for timestamp in ["not a UTC timestamp", "+2023-11-14T22:13:20Z"] {
365            request.fetched_at = timestamp.to_string();
366            assert!(matches!(
367                NnsNeuronCollectionState::new(&request, 2, 3),
368                Err(NnsNeuronError::Governance(
369                    NnsGovernanceError::InvalidTimestamp {
370                        field: "fetched_at",
371                        ..
372                    }
373                ))
374            ));
375            for field in ["started_at", "updated_at"] {
376                let mut value = serde_json::to_value(&valid_state).expect("serialize state");
377                value[field] = serde_json::json!(timestamp);
378                let state = serde_json::from_value(value).expect("restore caller state");
379                assert!(matches!(
380                    validate_collection_state(&state),
381                    Err(NnsNeuronError::Governance(
382                        NnsGovernanceError::InvalidTimestamp { value, .. }
383                    )) if value == timestamp
384                ));
385            }
386        }
387    }
388}