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