hotshot_new_protocol/
state.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    sync::Arc,
4};
5
6use committable::{Commitment, Committable};
7use hotshot::traits::{BlockPayload, ValidatedState};
8use hotshot_types::{
9    data::{BlockNumber, EpochNumber, Leaf2, VidCommitment, ViewNumber},
10    message::UpgradeLock,
11    traits::{
12        block_contents::{BlockHeader, BuilderFee},
13        metrics::Histogram,
14        node_implementation::NodeType,
15    },
16    utils::BuilderCommitment,
17    vote::HasViewNumber,
18};
19use tokio::task::{AbortHandle, JoinSet};
20use tracing::{error, warn};
21
22use crate::{
23    coordinator::metrics::{Measurement, finish_measurement, ignore_measurement},
24    helpers::proposal_commitment,
25    message::Proposal,
26};
27
28pub struct UpdateLeaf<T: NodeType> {
29    pub view: ViewNumber,
30    pub leaf: Leaf2<T>,
31    pub state: Arc<T::ValidatedState>,
32    pub delta: Option<Delta<T>>,
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub struct StateRequest<T: NodeType> {
37    pub view: ViewNumber,
38    pub parent_view: ViewNumber,
39    pub epoch: EpochNumber,
40    pub block: BlockNumber,
41    pub proposal: Proposal<T>,
42    pub parent_commitment: Commitment<Leaf2<T>>,
43    pub payload_size: u32,
44}
45
46#[derive(Clone, Debug, Eq, PartialEq)]
47pub struct HeaderRequest<T: NodeType> {
48    pub view: ViewNumber,
49    pub epoch: EpochNumber,
50    pub parent_proposal: Proposal<T>,
51    pub payload_commitment: VidCommitment,
52    pub builder_commitment: BuilderCommitment,
53    pub metadata: <T::BlockPayload as BlockPayload<T>>::Metadata,
54    pub builder_fee: BuilderFee<T>,
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct StateResponse<T: NodeType> {
59    pub view: ViewNumber,
60    pub commitment: Commitment<Leaf2<T>>,
61    pub state: Arc<T::ValidatedState>,
62    pub delta: Option<Delta<T>>,
63}
64
65#[derive(Clone, Debug, Eq, PartialEq)]
66pub struct HeaderResponse<T: NodeType> {
67    pub view: ViewNumber,
68    pub epoch: EpochNumber,
69    pub parent_proposal: Proposal<T>,
70}
71
72#[derive(Clone, Debug, Eq, PartialEq)]
73#[allow(clippy::large_enum_variant)]
74pub enum StateManagerOutput<T: NodeType> {
75    State {
76        response: StateResponse<T>,
77        validated: bool,
78    },
79    Header {
80        response: HeaderResponse<T>,
81        header: Option<T::BlockHeader>,
82    },
83}
84
85type Delta<T> = Arc<<<T as NodeType>::ValidatedState as ValidatedState<T>>::Delta>;
86
87#[derive(Clone)]
88pub struct StateEntry<T: NodeType> {
89    pub state: Arc<T::ValidatedState>,
90    pub delta: Option<Delta<T>>,
91    pub leaf: Leaf2<T>,
92}
93
94pub struct StateManager<T: NodeType> {
95    instance: Arc<T::InstanceState>,
96    validated_states: BTreeMap<Commitment<Leaf2<T>>, StateEntry<T>>,
97    state_requests: HashMap<Commitment<Leaf2<T>>, (AbortHandle, ViewNumber)>,
98    header_requests: HashMap<(ViewNumber, Commitment<Leaf2<T>>), AbortHandle>,
99    pending_requests: HashMap<Commitment<Leaf2<T>>, Vec<Pending<T>>>,
100    upgrade_lock: UpgradeLock<T>,
101    tasks: JoinSet<Completed<T>>,
102    validate_duration_metric: Option<Arc<dyn Histogram>>,
103    update_leaf_duration_metric: Option<Arc<dyn Histogram>>,
104}
105
106enum Pending<T: NodeType> {
107    State(StateRequest<T>),
108    Header(HeaderRequest<T>),
109}
110
111impl<T: NodeType> Pending<T> {
112    fn view(&self) -> ViewNumber {
113        match self {
114            Pending::State(r) => r.view,
115            Pending::Header(r) => r.view,
116        }
117    }
118}
119
120enum Completed<T: NodeType> {
121    State {
122        response: StateResponse<T>,
123        leaf: Option<Leaf2<T>>,
124    },
125    Header {
126        response: HeaderResponse<T>,
127        header: Option<T::BlockHeader>,
128    },
129}
130
131impl<T: NodeType> StateManager<T> {
132    pub fn new(instance: Arc<T::InstanceState>, upgrade_lock: UpgradeLock<T>) -> Self {
133        Self {
134            instance,
135            validated_states: BTreeMap::new(),
136            state_requests: HashMap::new(),
137            header_requests: HashMap::new(),
138            pending_requests: HashMap::new(),
139            upgrade_lock,
140            tasks: JoinSet::new(),
141            validate_duration_metric: None,
142            update_leaf_duration_metric: None,
143        }
144    }
145
146    pub fn with_metrics(
147        mut self,
148        validate: Option<Arc<dyn Histogram>>,
149        update_leaf: Option<Arc<dyn Histogram>>,
150    ) -> Self {
151        self.validate_duration_metric = validate;
152        self.update_leaf_duration_metric = update_leaf;
153        self
154    }
155
156    /// Get the validated state for a given view.
157    pub fn get_state(&self, view: ViewNumber) -> Option<&StateEntry<T>> {
158        self.validated_states
159            .iter()
160            .find(|(_, entry)| entry.leaf.view_number() == view)
161            .map(|(_, entry)| entry)
162    }
163
164    /// Get the leaf for a given view
165    pub fn get_leaf(&self, view: ViewNumber) -> Option<Leaf2<T>> {
166        self.validated_states
167            .iter()
168            .find(|(_, entry)| entry.leaf.view_number() == view)
169            .map(|(_, entry)| entry.leaf.clone())
170    }
171
172    pub fn seed_state(&mut self, view: ViewNumber, state: Arc<T::ValidatedState>, leaf: Leaf2<T>) {
173        self.insert_state(view, state, None, leaf);
174    }
175
176    /// Seed a commitment-only (`from_header`) state so a child proposal can be
177    /// validated against this leaf via catchup instead of being dropped.
178    pub(crate) fn seed_from_header(&mut self, proposal: Proposal<T>) {
179        self.insert_empty_state(proposal);
180    }
181
182    pub fn request_state(&mut self, request: StateRequest<T>) {
183        let commitment = proposal_commitment(&request.proposal);
184        if self.state_requests.contains_key(&commitment) {
185            return;
186        }
187
188        if self.state_requests.contains_key(&request.parent_commitment) {
189            self.pending_requests
190                .entry(request.parent_commitment)
191                .or_default()
192                .push(Pending::State(request));
193            return;
194        }
195
196        let Some(parent_entry) = self
197            .validated_states
198            .get(&request.parent_commitment)
199            .cloned()
200        else {
201            warn!(
202                view = %request.view,
203                parent_view = %request.parent_view,
204                epoch = %request.epoch,
205                block = %request.block,
206                parent_commitment = %request.parent_commitment,
207                "parent state unavailable; deferring state validation (from_header stub inserted). \
208                 If this persists, the node cannot vote until the parent state is recovered."
209            );
210            self.insert_empty_state(request.proposal);
211            self.start_pending(commitment);
212            return;
213        };
214
215        let instance = self.instance.clone();
216        let header = request.proposal.block_header.clone();
217        let view = request.view;
218        let payload_size = request.payload_size;
219
220        let Ok(upgrade_lock) = self.upgrade_lock.version(view) else {
221            error!(%view, "unsupported version");
222            return;
223        };
224
225        let duration_metric = self.validate_duration_metric.clone();
226        let handle = self.tasks.spawn(async move {
227            let measurement = duration_metric.map(Measurement::start);
228            let result = parent_entry
229                .state
230                .validate_and_apply_header(
231                    &instance,
232                    &parent_entry.leaf,
233                    &header,
234                    payload_size,
235                    upgrade_lock,
236                    *view,
237                )
238                .await;
239            match result {
240                Ok((state, delta)) => {
241                    finish_measurement(measurement);
242                    Completed::State {
243                        response: StateResponse {
244                            view,
245                            commitment,
246                            state: Arc::new(state),
247                            delta: Some(Arc::new(delta)),
248                        },
249                        leaf: Some(request.proposal.into()),
250                    }
251                },
252                Err(err) => {
253                    ignore_measurement(measurement);
254                    warn!(%err, "state validation failed");
255                    Completed::State {
256                        response: StateResponse {
257                            view,
258                            commitment,
259                            state: Arc::new(T::ValidatedState::from_header(&header)),
260                            delta: None,
261                        },
262                        leaf: None,
263                    }
264                },
265            }
266        });
267
268        self.state_requests.insert(commitment, (handle, view));
269    }
270
271    pub fn request_header(&mut self, request: HeaderRequest<T>) {
272        let parent_commitment = proposal_commitment(&request.parent_proposal);
273        if self
274            .header_requests
275            .contains_key(&(request.view, parent_commitment))
276        {
277            return;
278        }
279
280        if self.state_requests.contains_key(&parent_commitment) {
281            self.pending_requests
282                .entry(parent_commitment)
283                .or_default()
284                .push(Pending::Header(request));
285            return;
286        }
287
288        let Some(parent_entry) = self.validated_states.get(&parent_commitment).cloned() else {
289            // Parent state not available yet (e.g. its proposal is still
290            // being validated).  Queue the request so it is retried once
291            // the state for the parent view is inserted.
292            self.pending_requests
293                .entry(parent_commitment)
294                .or_default()
295                .push(Pending::Header(request));
296            return;
297        };
298
299        let instance = self.instance.clone();
300        let view = request.view;
301        let epoch = request.epoch;
302        let parent_proposal = request.parent_proposal;
303
304        let Ok(version) = self.upgrade_lock.version(view) else {
305            error!(%view, "unsupported version");
306            return;
307        };
308
309        let handle = self.tasks.spawn(async move {
310            let result = T::BlockHeader::new(
311                &parent_entry.state,
312                &instance,
313                &parent_entry.leaf,
314                request.payload_commitment,
315                request.builder_commitment,
316                request.metadata,
317                request.builder_fee,
318                version,
319                *view,
320            )
321            .await;
322            match result {
323                Ok(header) => Completed::Header {
324                    response: HeaderResponse {
325                        view,
326                        epoch,
327                        parent_proposal,
328                    },
329                    header: Some(header),
330                },
331                Err(err) => {
332                    warn!(%err, "header creation failed");
333                    Completed::Header {
334                        response: HeaderResponse {
335                            view,
336                            epoch,
337                            parent_proposal,
338                        },
339                        header: None,
340                    }
341                },
342            }
343        });
344
345        self.header_requests
346            .insert((view, parent_commitment), handle);
347    }
348
349    /// Provide an externally-obtained validated state.
350    pub fn update_state(&mut self, update: UpdateLeaf<T>) {
351        let UpdateLeaf {
352            view,
353            leaf,
354            state,
355            delta,
356        } = update;
357        let commitment = leaf.commit();
358        self.insert_state(view, state, delta, leaf);
359        if let Some((task, _)) = self.state_requests.remove(&commitment) {
360            task.abort();
361        }
362        self.start_pending(commitment);
363    }
364
365    /// Get the next output.
366    pub async fn next(&mut self) -> Option<StateManagerOutput<T>> {
367        loop {
368            match self.tasks.join_next().await {
369                Some(Ok(result)) => match result {
370                    Completed::State {
371                        response,
372                        leaf: leaf2,
373                    } => {
374                        if self.state_requests.remove(&response.commitment).is_none() {
375                            continue;
376                        }
377                        if let Some(leaf) = leaf2 {
378                            let measurement = self
379                                .update_leaf_duration_metric
380                                .clone()
381                                .map(Measurement::start);
382                            self.insert_state(
383                                response.view,
384                                response.state.clone(),
385                                response.delta.clone(),
386                                leaf,
387                            );
388                            finish_measurement(measurement);
389                            self.start_pending(response.commitment);
390                            return Some(StateManagerOutput::State {
391                                response,
392                                validated: true,
393                            });
394                        } else {
395                            self.pending_requests.remove(&response.commitment);
396                            return Some(StateManagerOutput::State {
397                                response,
398                                validated: false,
399                            });
400                        }
401                    },
402                    Completed::Header { response, header } => {
403                        let key = (
404                            response.view,
405                            proposal_commitment(&response.parent_proposal),
406                        );
407                        if self.header_requests.remove(&key).is_none() {
408                            continue;
409                        }
410                        return Some(StateManagerOutput::Header { response, header });
411                    },
412                },
413                Some(Err(err)) => {
414                    if err.is_panic() {
415                        error!(%err, "task panicked");
416                    }
417                },
418                None => return None,
419            }
420        }
421    }
422
423    pub fn gc(&mut self, view_number: ViewNumber) {
424        self.validated_states
425            .retain(|_, entry| entry.leaf.view_number() >= view_number);
426
427        for (task, view) in self.state_requests.values() {
428            if *view < view_number {
429                task.abort();
430            }
431        }
432
433        self.state_requests
434            .retain(|_, (_, view)| *view >= view_number);
435
436        self.header_requests.retain(|(view, _), handle| {
437            let keep = *view >= view_number;
438            if !keep {
439                handle.abort();
440            }
441            keep
442        });
443
444        self.pending_requests.retain(|_, pending| {
445            pending.retain(|p| p.view() >= view_number);
446            !pending.is_empty()
447        });
448    }
449
450    fn start_pending(&mut self, finished_commitment: Commitment<Leaf2<T>>) {
451        let Some(pending) = self.pending_requests.remove(&finished_commitment) else {
452            return;
453        };
454        for p in pending {
455            match p {
456                Pending::State(r) => self.request_state(r),
457                Pending::Header(r) => self.request_header(r),
458            }
459        }
460    }
461
462    /// Insert a state into the validated states map.
463    ///
464    /// States created via `from_header`
465    /// have no delta. States produced by `validate_and_apply_header` carry a delta representing
466    /// the state transition. This method prevents a `from_header` state from overwriting a
467    /// fully validated state that already has a delta.
468    fn insert_state(
469        &mut self,
470        view: ViewNumber,
471        state: Arc<T::ValidatedState>,
472        delta: Option<Delta<T>>,
473        leaf: Leaf2<T>,
474    ) {
475        if let Some(existing) = self.validated_states.get(&leaf.commit())
476            && existing.delta.is_some()
477            && delta.is_none()
478        {
479            warn!(
480                ?view,
481                "Skipping state update to not override a state with a delta"
482            );
483            return;
484        }
485        self.validated_states
486            .insert(leaf.commit(), StateEntry { state, delta, leaf });
487    }
488
489    fn insert_empty_state(&mut self, proposal: Proposal<T>) {
490        let state = T::ValidatedState::from_header(&proposal.block_header);
491        self.insert_state(
492            proposal.view_number(),
493            Arc::new(state),
494            None,
495            proposal.into(),
496        );
497    }
498
499    #[cfg(test)]
500    pub(crate) fn validated_contains_view(&self, v: ViewNumber) -> bool {
501        self.validated_states
502            .iter()
503            .any(|(_, entry)| entry.leaf.view_number() == v)
504    }
505
506    #[cfg(test)]
507    pub(crate) fn pending_contains_commitment(&self, c: &Commitment<Leaf2<T>>) -> bool {
508        self.pending_requests.contains_key(c)
509    }
510}
hotshot_new_protocol/state.rs

hotshot_new_protocol/
state.rs
