hotshot_new_protocol/

state.rs

use std::{
    collections::{BTreeMap, HashMap},
    sync::Arc,
};

use committable::{Commitment, Committable};
use hotshot_types::{
    data::{BlockNumber, EpochNumber, Leaf2, VidCommitment, ViewNumber},
    message::UpgradeLock,
    traits::{
        BlockPayload, ValidatedState,
        block_contents::{BlockHeader, BuilderFee},
        metrics::Histogram,
        node_implementation::NodeType,
    },
    utils::BuilderCommitment,
    vote::HasViewNumber,
};
use tokio::task::{AbortHandle, JoinSet};
use tracing::{debug, error, warn};

use crate::{
    coordinator::metrics::{Measurement, finish_measurement, ignore_measurement},
    helpers::proposal_commitment,
    message::Proposal,
};

/// Validated states stay this many views below the decided view so peers can serve catchup.
const STATE_GC_MARGIN: u64 = 8;

pub struct UpdateLeaf<T: NodeType> {
    pub view: ViewNumber,
    pub leaf: Leaf2<T>,
    pub state: Arc<T::ValidatedState>,
    pub delta: Option<Delta<T>>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateRequest<T: NodeType> {
    pub view: ViewNumber,
    pub parent_view: ViewNumber,
    pub epoch: EpochNumber,
    pub block: BlockNumber,
    pub proposal: Proposal<T>,
    pub parent_commitment: Commitment<Leaf2<T>>,
    pub payload_size: Option<u32>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct HeaderRequest<T: NodeType> {
    pub view: ViewNumber,
    pub epoch: EpochNumber,
    pub parent_proposal: Proposal<T>,
    pub payload_commitment: VidCommitment,
    pub builder_commitment: BuilderCommitment,
    pub metadata: <T::BlockPayload as BlockPayload<T>>::Metadata,
    pub builder_fee: BuilderFee<T>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateResponse<T: NodeType> {
    pub view: ViewNumber,
    pub commitment: Commitment<Leaf2<T>>,
    pub state: Arc<T::ValidatedState>,
    pub delta: Option<Delta<T>>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct HeaderResponse<T: NodeType> {
    pub view: ViewNumber,
    pub epoch: EpochNumber,
    pub parent_proposal: Proposal<T>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
#[allow(clippy::large_enum_variant)]
pub enum StateManagerOutput<T: NodeType> {
    State {
        response: StateResponse<T>,
        validated: bool,
    },
    Header {
        response: HeaderResponse<T>,
        header: Option<T::BlockHeader>,
    },
}

type Delta<T> = Arc<<<T as NodeType>::ValidatedState as ValidatedState<T>>::Delta>;

#[derive(Clone)]
pub struct StateEntry<T: NodeType> {
    pub state: Arc<T::ValidatedState>,
    pub delta: Option<Delta<T>>,
    pub leaf: Leaf2<T>,
}

pub struct StateManager<T: NodeType> {
    instance: Arc<T::InstanceState>,
    validated_states: BTreeMap<Commitment<Leaf2<T>>, StateEntry<T>>,
    state_requests: HashMap<Commitment<Leaf2<T>>, InFlight<T>>,
    header_requests: HashMap<(ViewNumber, Commitment<Leaf2<T>>), AbortHandle>,
    pending_requests: HashMap<Commitment<Leaf2<T>>, Vec<Pending<T>>>,
    /// Newest decided leaf whose in-flight validation made `seed_decided` skip it.
    /// Cleared by gc below the gc view.
    decided_in_flight: Option<Leaf2<T>>,
    upgrade_lock: UpgradeLock<T>,
    tasks: JoinSet<Completed<T>>,
    validate_duration_metric: Option<Arc<dyn Histogram>>,
    update_leaf_duration_metric: Option<Arc<dyn Histogram>>,
}

/// The proposal lets `gc` seed a stub for a validation it aborts.
struct InFlight<T: NodeType> {
    handle: AbortHandle,
    view: ViewNumber,
    proposal: Proposal<T>,
}

enum Pending<T: NodeType> {
    State(StateRequest<T>),
    Header(HeaderRequest<T>),
}

impl<T: NodeType> Pending<T> {
    fn view(&self) -> ViewNumber {
        match self {
            Pending::State(r) => r.view,
            Pending::Header(r) => r.view,
        }
    }
}

enum Completed<T: NodeType> {
    State {
        response: StateResponse<T>,
        leaf: Option<Leaf2<T>>,
    },
    Header {
        response: HeaderResponse<T>,
        header: Option<T::BlockHeader>,
    },
}

impl<T: NodeType> StateManager<T> {
    pub fn new(instance: Arc<T::InstanceState>, upgrade_lock: UpgradeLock<T>) -> Self {
        Self {
            instance,
            validated_states: BTreeMap::new(),
            state_requests: HashMap::new(),
            header_requests: HashMap::new(),
            pending_requests: HashMap::new(),
            decided_in_flight: None,
            upgrade_lock,
            tasks: JoinSet::new(),
            validate_duration_metric: None,
            update_leaf_duration_metric: None,
        }
    }

    pub fn with_metrics(
        mut self,
        validate: Option<Arc<dyn Histogram>>,
        update_leaf: Option<Arc<dyn Histogram>>,
    ) -> Self {
        self.validate_duration_metric = validate;
        self.update_leaf_duration_metric = update_leaf;
        self
    }

    /// Get the validated state for a given view.
    ///
    /// Lookup is by view, so below the decided view this may return a state from an undecided fork.
    pub fn get_state(&self, view: ViewNumber) -> Option<&StateEntry<T>> {
        self.validated_states
            .iter()
            .find(|(_, entry)| entry.leaf.view_number() == view)
            .map(|(_, entry)| entry)
    }

    /// Get the leaf for a given view
    pub fn get_leaf(&self, view: ViewNumber) -> Option<Leaf2<T>> {
        self.validated_states
            .iter()
            .find(|(_, entry)| entry.leaf.view_number() == view)
            .map(|(_, entry)| entry.leaf.clone())
    }

    pub fn seed_state(&mut self, view: ViewNumber, state: Arc<T::ValidatedState>, leaf: Leaf2<T>) {
        self.insert_state(view, state, None, leaf);
    }

    /// Seed a commitment-only (`from_header`) state so a child proposal can be
    /// validated against this leaf via catchup instead of being dropped.
    ///
    /// A real (validated) state for the leaf is never displaced, and any
    /// header/state requests queued on this leaf are restarted, or re-queued while
    /// its validation is in flight.
    pub(crate) fn seed_from_header(&mut self, proposal: Proposal<T>) {
        let commitment = proposal_commitment(&proposal);
        if !self.validated_states.contains_key(&commitment) {
            self.insert_empty_state(proposal);
        }
        self.start_pending(commitment);
    }

    /// Seed a stub for a decided leaf so work queued on it runs via catchup.
    /// No-op when the leaf is validated. When it is being validated, the stub is
    /// seeded only if that validation fails. An older leaf never replaces a newer one.
    pub(crate) fn seed_decided(&mut self, leaf: Leaf2<T>) {
        let commitment = leaf.commit();
        if self.validated_states.contains_key(&commitment) {
            return;
        }
        if self.state_requests.contains_key(&commitment) {
            let newest = self
                .decided_in_flight
                .as_ref()
                .is_none_or(|held| held.view_number() <= leaf.view_number());
            if newest {
                self.decided_in_flight = Some(leaf);
            }
            return;
        }
        let state = T::ValidatedState::from_header(leaf.block_header());
        self.insert_state(leaf.view_number(), Arc::new(state), None, leaf);
        self.start_pending(commitment);
    }

    pub fn request_state(&mut self, request: StateRequest<T>) {
        let commitment = proposal_commitment(&request.proposal);
        if self.state_requests.contains_key(&commitment) {
            return;
        }

        if self.state_requests.contains_key(&request.parent_commitment) {
            self.pending_requests
                .entry(request.parent_commitment)
                .or_default()
                .push(Pending::State(request));
            return;
        }

        let Some(parent_entry) = self
            .validated_states
            .get(&request.parent_commitment)
            .cloned()
        else {
            warn!(
                view = %request.view,
                parent_view = %request.parent_view,
                epoch = %request.epoch,
                block = %request.block,
                parent_commitment = %request.parent_commitment,
                "parent state unavailable; queued on parent for retry. \
                 If this persists, the parent state never arrived and the node cannot vote."
            );
            let queued = self
                .pending_requests
                .entry(request.parent_commitment)
                .or_default();
            if !queued
                .iter()
                .any(|p| matches!(p, Pending::State(r) if r.view == request.view))
            {
                queued.push(Pending::State(request));
            }
            return;
        };

        let instance = self.instance.clone();
        let header = request.proposal.block_header.clone();
        let proposal = request.proposal.clone();
        let view = request.view;
        let payload_size = request.payload_size;

        let Ok(upgrade_lock) = self.upgrade_lock.version(view) else {
            error!(%view, "unsupported version");
            return;
        };

        let duration_metric = self.validate_duration_metric.clone();
        let handle = self.tasks.spawn(async move {
            let measurement = duration_metric.map(Measurement::start);
            let result = parent_entry
                .state
                .validate_and_apply_header(
                    &instance,
                    &parent_entry.leaf,
                    &header,
                    payload_size,
                    upgrade_lock,
                    *view,
                )
                .await;
            match result {
                Ok((state, delta)) => {
                    finish_measurement(measurement);
                    Completed::State {
                        response: StateResponse {
                            view,
                            commitment,
                            state: Arc::new(state),
                            delta: Some(Arc::new(delta)),
                        },
                        leaf: Some(request.proposal.into()),
                    }
                },
                Err(err) => {
                    ignore_measurement(measurement);
                    warn!(%err, "state validation failed");
                    Completed::State {
                        response: StateResponse {
                            view,
                            commitment,
                            state: Arc::new(T::ValidatedState::from_header(&header)),
                            delta: None,
                        },
                        leaf: None,
                    }
                },
            }
        });

        self.state_requests.insert(
            commitment,
            InFlight {
                handle,
                view,
                proposal,
            },
        );
    }

    pub fn request_header(&mut self, request: HeaderRequest<T>) {
        let parent_commitment = proposal_commitment(&request.parent_proposal);
        if self
            .header_requests
            .contains_key(&(request.view, parent_commitment))
        {
            return;
        }

        if self.state_requests.contains_key(&parent_commitment) {
            self.pending_requests
                .entry(parent_commitment)
                .or_default()
                .push(Pending::Header(request));
            return;
        }

        let Some(parent_entry) = self.validated_states.get(&parent_commitment).cloned() else {
            // Parent state not available yet (e.g. its proposal is still
            // being validated).  Queue the request so it is retried once
            // the state for the parent view is inserted.
            self.pending_requests
                .entry(parent_commitment)
                .or_default()
                .push(Pending::Header(request));
            return;
        };

        let instance = self.instance.clone();
        let view = request.view;
        let epoch = request.epoch;
        let parent_proposal = request.parent_proposal;

        let Ok(version) = self.upgrade_lock.version(view) else {
            error!(%view, "unsupported version");
            return;
        };

        let handle = self.tasks.spawn(async move {
            let result = T::BlockHeader::new(
                &parent_entry.state,
                &instance,
                &parent_entry.leaf,
                request.payload_commitment,
                request.builder_commitment,
                request.metadata,
                request.builder_fee,
                version,
                *view,
            )
            .await;
            match result {
                Ok(header) => Completed::Header {
                    response: HeaderResponse {
                        view,
                        epoch,
                        parent_proposal,
                    },
                    header: Some(header),
                },
                Err(err) => {
                    warn!(%err, "header creation failed");
                    Completed::Header {
                        response: HeaderResponse {
                            view,
                            epoch,
                            parent_proposal,
                        },
                        header: None,
                    }
                },
            }
        });

        self.header_requests
            .insert((view, parent_commitment), handle);
    }

    /// Provide an externally-obtained validated state.
    ///
    /// An update without a delta only enriches an existing entry and is
    /// otherwise dropped.
    pub fn update_state(&mut self, update: UpdateLeaf<T>) {
        let UpdateLeaf {
            view,
            leaf,
            state,
            delta,
        } = update;
        let commitment = leaf.commit();
        if delta.is_none() {
            if self.validated_states.contains_key(&commitment) {
                self.insert_state(view, state, delta, leaf);
            } else {
                debug!(%view, %commitment, "dropping delta-less update: no entry for leaf");
            }
            return;
        }
        self.insert_state(view, state, delta, leaf);
        if let Some(in_flight) = self.state_requests.remove(&commitment) {
            in_flight.handle.abort();
        }
        self.start_pending(commitment);
    }

    /// Get the next output.
    pub async fn next(&mut self) -> Option<StateManagerOutput<T>> {
        loop {
            match self.tasks.join_next().await {
                Some(Ok(result)) => match result {
                    Completed::State {
                        response,
                        leaf: leaf2,
                    } => {
                        if self.state_requests.remove(&response.commitment).is_none() {
                            continue;
                        }
                        let decided = self
                            .decided_in_flight
                            .take_if(|leaf| leaf.commit() == response.commitment);
                        if let Some(leaf) = leaf2 {
                            let measurement = self
                                .update_leaf_duration_metric
                                .clone()
                                .map(Measurement::start);
                            self.insert_state(
                                response.view,
                                response.state.clone(),
                                response.delta.clone(),
                                leaf,
                            );
                            finish_measurement(measurement);
                            self.start_pending(response.commitment);
                            return Some(StateManagerOutput::State {
                                response,
                                validated: true,
                            });
                        } else {
                            if self.validated_states.contains_key(&response.commitment) {
                                self.start_pending(response.commitment);
                            } else if let Some(leaf) = decided {
                                self.seed_decided(leaf);
                            }
                            return Some(StateManagerOutput::State {
                                response,
                                validated: false,
                            });
                        }
                    },
                    Completed::Header { response, header } => {
                        let key = (
                            response.view,
                            proposal_commitment(&response.parent_proposal),
                        );
                        if self.header_requests.remove(&key).is_none() {
                            continue;
                        }
                        return Some(StateManagerOutput::Header { response, header });
                    },
                },
                Some(Err(err)) => {
                    if err.is_panic() {
                        error!(%err, "task panicked");
                    }
                },
                None => return None,
            }
        }
    }

    /// The decided view's own validation, or a header this node needs to
    /// propose, may be queued behind a validation aborted here. A stub for the
    /// aborted proposal lets them proceed via catchup instead of hanging.
    pub fn gc(&mut self, view_number: ViewNumber) {
        let margin_view = ViewNumber::new(view_number.saturating_sub(STATE_GC_MARGIN));
        self.validated_states.retain(|_, entry| {
            let view = entry.leaf.view_number();
            view >= view_number || (entry.delta.is_some() && view >= margin_view)
        });

        self.decided_in_flight
            .take_if(|leaf| leaf.view_number() < view_number);

        self.header_requests.retain(|(view, _), handle| {
            let keep = *view >= view_number;
            if !keep {
                handle.abort();
            }
            keep
        });

        self.pending_requests.retain(|_, pending| {
            pending.retain(|p| p.view() >= view_number);
            !pending.is_empty()
        });

        let stale: Vec<_> = self
            .state_requests
            .extract_if(|_, in_flight| in_flight.view < view_number)
            .collect();
        for (commitment, in_flight) in stale {
            in_flight.handle.abort();
            if self.pending_requests.contains_key(&commitment) {
                self.seed_from_header(in_flight.proposal);
            }
        }
    }

    fn start_pending(&mut self, finished_commitment: Commitment<Leaf2<T>>) {
        let Some(pending) = self.pending_requests.remove(&finished_commitment) else {
            return;
        };
        for p in pending {
            match p {
                Pending::State(r) => self.request_state(r),
                Pending::Header(r) => self.request_header(r),
            }
        }
    }

    /// Insert a state into the validated states map.
    ///
    /// States created via `from_header`
    /// have no delta. States produced by `validate_and_apply_header` carry a delta representing
    /// the state transition. This method prevents a `from_header` state from overwriting a
    /// fully validated state that already has a delta.
    fn insert_state(
        &mut self,
        view: ViewNumber,
        state: Arc<T::ValidatedState>,
        delta: Option<Delta<T>>,
        leaf: Leaf2<T>,
    ) {
        if let Some(existing) = self.validated_states.get(&leaf.commit())
            && existing.delta.is_some()
            && delta.is_none()
        {
            warn!(
                ?view,
                "Skipping state update to not override a state with a delta"
            );
            return;
        }
        self.validated_states
            .insert(leaf.commit(), StateEntry { state, delta, leaf });
    }

    fn insert_empty_state(&mut self, proposal: Proposal<T>) {
        let state = T::ValidatedState::from_header(&proposal.block_header);
        self.insert_state(
            proposal.view_number(),
            Arc::new(state),
            None,
            proposal.into(),
        );
    }

    #[cfg(test)]
    pub(crate) fn validated_contains_view(&self, v: ViewNumber) -> bool {
        self.validated_states
            .iter()
            .any(|(_, entry)| entry.leaf.view_number() == v)
    }
}