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

ic_query/nns/neuron/report/
collection.rs
