hotshot_new_protocol/vid/
disperse.rs

1use std::{
2    collections::BTreeMap,
3    mem,
4    ops::Range,
5    sync::{Arc, LazyLock},
6};
7
8use hotshot::traits::BlockPayload;
9use hotshot_types::{
10    data::{
11        EpochNumber, VidCommitment2, VidDisperse2, ViewNumber,
12        vid_disperse::{AvidmGf2DisperseShareFragment, AvidmGf2NamespacePiece},
13    },
14    epoch_membership::EpochMembershipCoordinator,
15    message::Proposal as SignedProposal,
16    traits::{metrics::Histogram, node_implementation::NodeType, signature_key::SignatureKey},
17    vid::avidm_gf2::AvidmGf2Scheme,
18};
19use hotshot_utils::anytrace::{self, Wrap};
20use rayon::prelude::*;
21use tokio::task::{AbortHandle, JoinSet};
22use tracing::{error, warn};
23
24use crate::{
25    coordinator::metrics::{Measurement, finish_measurement},
26    message::{ConsensusMessage, Message, MessageType},
27    network::{NetworkError, Sender},
28};
29
30static NUM_THREADS: LazyLock<usize> = LazyLock::new(|| rayon::current_num_threads().max(1));
31
32#[derive(Clone, Eq, PartialEq, Debug)]
33pub struct VidDisperseRequest<T: NodeType> {
34    pub view: ViewNumber,
35    pub epoch: EpochNumber,
36    pub block: T::BlockPayload,
37    pub metadata: <T::BlockPayload as BlockPayload<T>>::Metadata,
38    pub payload_commitment: VidCommitment2,
39}
40
41pub struct VidDisperseOutput {
42    pub view: ViewNumber,
43    pub payload_commitment: VidCommitment2,
44}
45
46pub struct VidDisperser<T: NodeType> {
47    calculations: BTreeMap<(ViewNumber, VidCommitment2), AbortHandle>,
48    epoch_membership_coordinator: EpochMembershipCoordinator<T>,
49    network: Sender<T>,
50    public_key: T::SignatureKey,
51    private_key: <T::SignatureKey as SignatureKey>::PrivateKey,
52    tasks: JoinSet<Result<VidDisperseOutput, VidDisperseError>>,
53    duration_metric: Option<Arc<dyn Histogram>>,
54}
55
56impl<T: NodeType> VidDisperser<T> {
57    pub fn new(
58        epoch_membership_coordinator: EpochMembershipCoordinator<T>,
59        network: Sender<T>,
60        public_key: T::SignatureKey,
61        private_key: <T::SignatureKey as SignatureKey>::PrivateKey,
62    ) -> Self {
63        Self {
64            calculations: BTreeMap::new(),
65            epoch_membership_coordinator,
66            network,
67            public_key,
68            private_key,
69            tasks: JoinSet::new(),
70            duration_metric: None,
71        }
72    }
73
74    pub fn with_metrics(mut self, hist: Option<Arc<dyn Histogram>>) -> Self {
75        self.duration_metric = hist;
76        self
77    }
78
79    pub fn request_vid_disperse(&mut self, vid_disperse_request: VidDisperseRequest<T>) {
80        let key = (
81            vid_disperse_request.view,
82            vid_disperse_request.payload_commitment,
83        );
84        if self.calculations.contains_key(&key) {
85            return;
86        }
87        let membership = self.epoch_membership_coordinator.clone();
88        let network = self.network.clone();
89        let public_key = self.public_key.clone();
90        let private_key = self.private_key.clone();
91        let duration_metric = self.duration_metric.clone();
92        let handle = self.tasks.spawn_blocking(move || {
93            let measurement = duration_metric.map(Measurement::start);
94            let result = handle_vid_disperse_request(
95                membership,
96                network,
97                public_key,
98                private_key,
99                vid_disperse_request,
100            );
101            finish_measurement(measurement);
102            result
103        });
104        self.calculations.insert(key, handle);
105    }
106
107    pub async fn next(&mut self) -> Option<Result<VidDisperseOutput, VidDisperseError>> {
108        loop {
109            match self.tasks.join_next().await {
110                Some(Ok(result)) => return Some(result),
111                Some(Err(err)) => {
112                    if err.is_panic() {
113                        error!(%err, "vid disperse task panic");
114                    }
115                    continue;
116                },
117                None => return None,
118            }
119        }
120    }
121
122    pub fn gc(&mut self, view_number: ViewNumber) {
123        let keep = self
124            .calculations
125            .split_off(&(view_number, VidCommitment2::default()));
126        for handle in self.calculations.values_mut() {
127            handle.abort();
128        }
129        self.calculations = keep;
130    }
131}
132
133fn handle_vid_disperse_request<T: NodeType>(
134    epoch_membership_coordinator: EpochMembershipCoordinator<T>,
135    network: Sender<T>,
136    public_key: T::SignatureKey,
137    private_key: <T::SignatureKey as SignatureKey>::PrivateKey,
138    vid_disperse_request: VidDisperseRequest<T>,
139) -> Result<VidDisperseOutput, VidDisperseError> {
140    let view = vid_disperse_request.view;
141    let epoch = vid_disperse_request.epoch;
142    let payload_commitment = vid_disperse_request.payload_commitment;
143
144    let params = VidDisperse2::<T>::disperse_params(
145        &vid_disperse_request.block,
146        &epoch_membership_coordinator,
147        Some(epoch),
148        &vid_disperse_request.metadata,
149    )
150    .map_err(VidDisperseError::Vid)?;
151
152    let signature = T::SignatureKey::sign(&private_key, payload_commitment.as_ref())
153        .map_err(|err| VidDisperseError::Sign(err.into()))?;
154
155    let num_namespaces = params.ns_table.len();
156
157    // Coalesce namespaces into size-balanced buckets, roughly one per core,
158    // but never below the minimum size, so a block of many tiny namespaces
159    // goes out as a few balanced messages per recipient instead of one tiny
160    // message each. Each bucket is one parallel unit and one message per
161    // recipient.
162    let buckets: Vec<Vec<usize>> = {
163        // We want buckets to be at least 256 KiB, otherwise the per-message
164        // overhead is too large. Larger payloads are divided by the number
165        // of available threads to fully utilise rayon.
166        let threshold = params.payload.len().div_ceil(*NUM_THREADS).max(256 * 1024);
167        bucketize(&params.ns_table, threshold)
168    };
169
170    buckets.par_iter().try_for_each_with(
171        network,
172        |network, bucket| -> Result<(), VidDisperseError> {
173            let mut pieces = vec![Vec::new(); params.recipients.len()];
174
175            for &ns_index in bucket {
176                let dispersal = AvidmGf2Scheme::ns_disperse_one(
177                    &params.param,
178                    &params.weights,
179                    &params.payload[params.ns_table[ns_index].clone()],
180                    ns_index,
181                )
182                .wrap()
183                .map_err(VidDisperseError::Vid)?;
184
185                let ns_payload_byte_len = dispersal.payload_byte_len;
186                let ns_commit = dispersal.commit;
187                for (pieces, ns_share) in pieces.iter_mut().zip(dispersal.shares) {
188                    pieces.push(AvidmGf2NamespacePiece {
189                        ns_index: dispersal.ns_index,
190                        ns_payload_byte_len,
191                        ns_commit,
192                        ns_share,
193                    });
194                }
195            }
196
197            for (recipient, pieces) in params.recipients.iter().zip(pieces) {
198                let fragment = AvidmGf2DisperseShareFragment {
199                    view_number: view,
200                    epoch: Some(epoch),
201                    target_epoch: Some(epoch),
202                    payload_commitment,
203                    recipient_key: recipient.clone(),
204                    param: params.param.clone(),
205                    num_namespaces,
206                    namespaces: pieces,
207                };
208                let message = Message {
209                    sender: public_key.clone(),
210                    message_type: MessageType::Consensus(ConsensusMessage::VidShareFragment(
211                        SignedProposal::new(fragment, signature.clone()),
212                    )),
213                };
214                if let Err(err) = network.unicast(view, recipient, &message) {
215                    if err.is_critical() {
216                        return Err(err.into());
217                    }
218                    warn!(%err, "network error while sending vid share fragment");
219                }
220            }
221            Ok(())
222        },
223    )?;
224
225    Ok(VidDisperseOutput {
226        view,
227        payload_commitment,
228    })
229}
230
231/// Group namespace indices into buckets whose payload sizes are each at least
232/// `threshold` bytes (except possibly the last), coalescing small namespaces so
233/// the disperser sends one message per bucket rather than one per namespace.
234///
235/// A `threshold` of 0 puts every namespace in its own bucket (no coalescing); a
236/// threshold larger than the whole payload yields a single bucket. A namespace
237/// at least `threshold` bytes on its own seals its bucket immediately, so large
238/// namespaces stay separate while small ones accumulate.
239fn bucketize(ns_table: &[Range<usize>], threshold: usize) -> Vec<Vec<usize>> {
240    let mut buckets = Vec::new();
241    let mut current = Vec::new();
242    let mut current_bytes = 0usize;
243    for (ns_index, range) in ns_table.iter().enumerate() {
244        current.push(ns_index);
245        current_bytes += range.len();
246        if current_bytes >= threshold {
247            buckets.push(mem::take(&mut current));
248            current_bytes = 0;
249        }
250    }
251    if !current.is_empty() {
252        buckets.push(current);
253    }
254    buckets
255}
256
257#[derive(Debug, thiserror::Error)]
258pub enum VidDisperseError {
259    #[error("network error: {0}")]
260    Net(#[from] NetworkError),
261
262    #[error("vid error: {0}")]
263    Vid(#[source] anytrace::Error),
264
265    #[error("sign error: {0}")]
266    Sign(#[source] Box<dyn std::error::Error + Send + Sync>),
267}
268
269impl VidDisperseError {
270    pub fn is_critical(&self) -> bool {
271        match self {
272            Self::Net(e) => e.is_critical(),
273            Self::Vid(_) | Self::Sign(_) => false,
274        }
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use std::ops::Range;
281
282    use quickcheck::{TestResult, quickcheck};
283
284    use super::bucketize;
285
286    /// Build a contiguous namespace table from per-namespace byte lengths.
287    /// Only the lengths matter to `bucketize`; the offsets are incidental.
288    fn ns_ranges(lens: &[u8]) -> Vec<Range<usize>> {
289        let mut start = 0;
290        lens.iter()
291            .map(|&len| {
292                let range = start..start + usize::from(len);
293                start += usize::from(len);
294                range
295            })
296            .collect()
297    }
298
299    /// Total payload bytes of the namespaces at `indices`.
300    fn bytes(indices: &[usize], ns_table: &[Range<usize>]) -> usize {
301        indices.iter().map(|&i| ns_table[i].len()).sum()
302    }
303
304    quickcheck! {
305        /// The buckets partition the namespace indices: every index appears
306        /// exactly once, contiguously, and in namespace-table order. This pins
307        /// completeness, disjointness, and ordering in a single property.
308        fn prop_buckets_partition(lens: Vec<u8>, threshold: u16) -> bool {
309            let ns_table = ns_ranges(&lens);
310            bucketize(&ns_table, threshold.into())
311                .into_iter()
312                .flatten()
313                .eq(0..ns_table.len())
314        }
315
316        /// No bucket is ever empty.
317        fn prop_buckets_non_empty(lens: Vec<u8>, threshold: u16) -> bool {
318            let ns_table = ns_ranges(&lens);
319            bucketize(&ns_table, threshold.into())
320                .iter()
321                .all(|bucket| !bucket.is_empty())
322        }
323
324        /// Every bucket except the last reaches the threshold: small namespaces
325        /// are coalesced until the bucket is large enough.
326        fn prop_non_last_buckets_meet_threshold(lens: Vec<u8>, threshold: u16) -> bool {
327            let ns_table = ns_ranges(&lens);
328            let buckets = bucketize(&ns_table, threshold.into());
329            let non_last = buckets.len().saturating_sub(1);
330            buckets
331                .iter()
332                .take(non_last)
333                .all(|bucket| bytes(bucket, &ns_table) >= usize::from(threshold))
334        }
335
336        /// Buckets are minimal: dropping a non-last bucket's final namespace
337        /// leaves it below the threshold, so the disperser seals as soon as it
338        /// reaches the cutoff rather than over-coalescing. (Vacuous at
339        /// `threshold == 0`, where every namespace is its own bucket.)
340        fn prop_non_last_buckets_are_minimal(lens: Vec<u8>, threshold: u16) -> TestResult {
341            if threshold == 0 {
342                return TestResult::discard();
343            }
344            let ns_table = ns_ranges(&lens);
345            let buckets = bucketize(&ns_table, threshold.into());
346            let non_last = buckets.len().saturating_sub(1);
347            TestResult::from_bool(buckets.iter().take(non_last).all(|bucket| {
348                bytes(&bucket[..bucket.len() - 1], &ns_table) < usize::from(threshold)
349            }))
350        }
351    }
352}
hotshot_new_protocol/vid/disperse.rs

hotshot_new_protocol/vid/
disperse.rs
