snarkos_node_bft/
worker.rs

1// Copyright 2024 Aleo Network Foundation
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use crate::{
17    MAX_FETCH_TIMEOUT_IN_MS,
18    MAX_WORKERS,
19    ProposedBatch,
20    Transport,
21    events::{Event, TransmissionRequest, TransmissionResponse},
22    helpers::{Pending, Ready, Storage, WorkerReceiver, fmt_id, max_redundant_requests},
23    spawn_blocking,
24};
25use snarkos_node_bft_ledger_service::LedgerService;
26use snarkvm::{
27    console::prelude::*,
28    ledger::{
29        block::Transaction,
30        narwhal::{BatchHeader, Data, Transmission, TransmissionID},
31        puzzle::{Solution, SolutionID},
32    },
33};
34
35use colored::Colorize;
36use indexmap::{IndexMap, IndexSet};
37use parking_lot::Mutex;
38use rand::seq::IteratorRandom;
39use std::{future::Future, net::SocketAddr, sync::Arc, time::Duration};
40use tokio::{sync::oneshot, task::JoinHandle, time::timeout};
41
42#[derive(Clone)]
43pub struct Worker<N: Network> {
44    /// The worker ID.
45    id: u8,
46    /// The gateway.
47    gateway: Arc<dyn Transport<N>>,
48    /// The storage.
49    storage: Storage<N>,
50    /// The ledger service.
51    ledger: Arc<dyn LedgerService<N>>,
52    /// The proposed batch.
53    proposed_batch: Arc<ProposedBatch<N>>,
54    /// The ready queue.
55    ready: Ready<N>,
56    /// The pending transmissions queue.
57    pending: Arc<Pending<TransmissionID<N>, Transmission<N>>>,
58    /// The spawned handles.
59    handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
60}
61
62impl<N: Network> Worker<N> {
63    /// Initializes a new worker instance.
64    pub fn new(
65        id: u8,
66        gateway: Arc<dyn Transport<N>>,
67        storage: Storage<N>,
68        ledger: Arc<dyn LedgerService<N>>,
69        proposed_batch: Arc<ProposedBatch<N>>,
70    ) -> Result<Self> {
71        // Ensure the worker ID is valid.
72        ensure!(id < MAX_WORKERS, "Invalid worker ID '{id}'");
73        // Return the worker.
74        Ok(Self {
75            id,
76            gateway,
77            storage,
78            ledger,
79            proposed_batch,
80            ready: Default::default(),
81            pending: Default::default(),
82            handles: Default::default(),
83        })
84    }
85
86    /// Run the worker instance.
87    pub fn run(&self, receiver: WorkerReceiver<N>) {
88        info!("Starting worker instance {} of the memory pool...", self.id);
89        // Start the worker handlers.
90        self.start_handlers(receiver);
91    }
92
93    /// Returns the worker ID.
94    pub const fn id(&self) -> u8 {
95        self.id
96    }
97
98    /// Returns a reference to the pending transmissions queue.
99    pub fn pending(&self) -> &Arc<Pending<TransmissionID<N>, Transmission<N>>> {
100        &self.pending
101    }
102}
103
104impl<N: Network> Worker<N> {
105    /// The maximum number of transmissions allowed in a worker.
106    pub const MAX_TRANSMISSIONS_PER_WORKER: usize =
107        BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH / MAX_WORKERS as usize;
108    /// The maximum number of transmissions allowed in a worker ping.
109    pub const MAX_TRANSMISSIONS_PER_WORKER_PING: usize = BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH / 10;
110
111    // transmissions
112
113    /// Returns the number of transmissions in the ready queue.
114    pub fn num_transmissions(&self) -> usize {
115        self.ready.num_transmissions()
116    }
117
118    /// Returns the number of ratifications in the ready queue.
119    pub fn num_ratifications(&self) -> usize {
120        self.ready.num_ratifications()
121    }
122
123    /// Returns the number of solutions in the ready queue.
124    pub fn num_solutions(&self) -> usize {
125        self.ready.num_solutions()
126    }
127
128    /// Returns the number of transactions in the ready queue.
129    pub fn num_transactions(&self) -> usize {
130        self.ready.num_transactions()
131    }
132}
133
134impl<N: Network> Worker<N> {
135    /// Returns the transmission IDs in the ready queue.
136    pub fn transmission_ids(&self) -> IndexSet<TransmissionID<N>> {
137        self.ready.transmission_ids()
138    }
139
140    /// Returns the transmissions in the ready queue.
141    pub fn transmissions(&self) -> IndexMap<TransmissionID<N>, Transmission<N>> {
142        self.ready.transmissions()
143    }
144
145    /// Returns the solutions in the ready queue.
146    pub fn solutions(&self) -> impl '_ + Iterator<Item = (SolutionID<N>, Data<Solution<N>>)> {
147        self.ready.solutions()
148    }
149
150    /// Returns the transactions in the ready queue.
151    pub fn transactions(&self) -> impl '_ + Iterator<Item = (N::TransactionID, Data<Transaction<N>>)> {
152        self.ready.transactions()
153    }
154}
155
156impl<N: Network> Worker<N> {
157    /// Clears the solutions from the ready queue.
158    pub(super) fn clear_solutions(&self) {
159        self.ready.clear_solutions()
160    }
161}
162
163impl<N: Network> Worker<N> {
164    /// Returns `true` if the transmission ID exists in the ready queue, proposed batch, storage, or ledger.
165    pub fn contains_transmission(&self, transmission_id: impl Into<TransmissionID<N>>) -> bool {
166        let transmission_id = transmission_id.into();
167        // Check if the transmission ID exists in the ready queue, proposed batch, storage, or ledger.
168        self.ready.contains(transmission_id)
169            || self.proposed_batch.read().as_ref().map_or(false, |p| p.contains_transmission(transmission_id))
170            || self.storage.contains_transmission(transmission_id)
171            || self.ledger.contains_transmission(&transmission_id).unwrap_or(false)
172    }
173
174    /// Returns the transmission if it exists in the ready queue, proposed batch, storage.
175    ///
176    /// Note: We explicitly forbid retrieving a transmission from the ledger, as transmissions
177    /// in the ledger are not guaranteed to be invalid for the current batch.
178    pub fn get_transmission(&self, transmission_id: TransmissionID<N>) -> Option<Transmission<N>> {
179        // Check if the transmission ID exists in the ready queue.
180        if let Some(transmission) = self.ready.get(transmission_id) {
181            return Some(transmission);
182        }
183        // Check if the transmission ID exists in storage.
184        if let Some(transmission) = self.storage.get_transmission(transmission_id) {
185            return Some(transmission);
186        }
187        // Check if the transmission ID exists in the proposed batch.
188        if let Some(transmission) =
189            self.proposed_batch.read().as_ref().and_then(|p| p.get_transmission(transmission_id))
190        {
191            return Some(transmission.clone());
192        }
193        None
194    }
195
196    /// Returns the transmissions if it exists in the worker, or requests it from the specified peer.
197    pub async fn get_or_fetch_transmission(
198        &self,
199        peer_ip: SocketAddr,
200        transmission_id: TransmissionID<N>,
201    ) -> Result<(TransmissionID<N>, Transmission<N>)> {
202        // Attempt to get the transmission from the worker.
203        if let Some(transmission) = self.get_transmission(transmission_id) {
204            return Ok((transmission_id, transmission));
205        }
206        // Send a transmission request to the peer.
207        let (candidate_id, transmission) = self.send_transmission_request(peer_ip, transmission_id).await?;
208        // Ensure the transmission ID matches.
209        ensure!(candidate_id == transmission_id, "Invalid transmission ID");
210        // Return the transmission.
211        Ok((transmission_id, transmission))
212    }
213
214    /// Removes up to the specified number of transmissions from the ready queue, and returns them.
215    pub(crate) fn drain(&self, num_transmissions: usize) -> impl Iterator<Item = (TransmissionID<N>, Transmission<N>)> {
216        self.ready.drain(num_transmissions).into_iter()
217    }
218
219    /// Reinserts the specified transmission into the ready queue.
220    pub(crate) fn reinsert(&self, transmission_id: TransmissionID<N>, transmission: Transmission<N>) -> bool {
221        // Check if the transmission ID exists.
222        if !self.contains_transmission(transmission_id) {
223            // Insert the transmission into the ready queue.
224            return self.ready.insert(transmission_id, transmission);
225        }
226        false
227    }
228
229    /// Broadcasts a worker ping event.
230    pub(crate) fn broadcast_ping(&self) {
231        // Retrieve the transmission IDs.
232        let transmission_ids = self
233            .ready
234            .transmission_ids()
235            .into_iter()
236            .choose_multiple(&mut rand::thread_rng(), Self::MAX_TRANSMISSIONS_PER_WORKER_PING)
237            .into_iter()
238            .collect::<IndexSet<_>>();
239
240        // Broadcast the ping event.
241        if !transmission_ids.is_empty() {
242            self.gateway.broadcast(Event::WorkerPing(transmission_ids.into()));
243        }
244    }
245}
246
247impl<N: Network> Worker<N> {
248    /// Handles the incoming transmission ID from a worker ping event.
249    fn process_transmission_id_from_ping(&self, peer_ip: SocketAddr, transmission_id: TransmissionID<N>) {
250        // Check if the transmission ID exists.
251        if self.contains_transmission(transmission_id) {
252            return;
253        }
254        // If the ready queue is full, then skip this transmission.
255        // Note: We must prioritize the unconfirmed solutions and unconfirmed transactions, not transmissions.
256        if self.ready.num_transmissions() > Self::MAX_TRANSMISSIONS_PER_WORKER {
257            return;
258        }
259        // Attempt to fetch the transmission from the peer.
260        let self_ = self.clone();
261        tokio::spawn(async move {
262            // Send a transmission request to the peer.
263            match self_.send_transmission_request(peer_ip, transmission_id).await {
264                // If the transmission was fetched, then process it.
265                Ok((candidate_id, transmission)) => {
266                    // Ensure the transmission ID matches.
267                    if candidate_id == transmission_id {
268                        // Insert the transmission into the ready queue.
269                        // Note: This method checks `contains_transmission` again, because by the time the transmission is fetched,
270                        // it could have already been inserted into the ready queue.
271                        self_.process_transmission_from_peer(peer_ip, transmission_id, transmission);
272                    }
273                }
274                // If the transmission was not fetched, then attempt to fetch it again.
275                Err(e) => {
276                    warn!(
277                        "Worker {} - Failed to fetch transmission '{}.{}' from '{peer_ip}' (ping) - {e}",
278                        self_.id,
279                        fmt_id(transmission_id),
280                        fmt_id(transmission_id.checksum().unwrap_or_default()).dimmed()
281                    );
282                }
283            }
284        });
285    }
286
287    /// Handles the incoming transmission from a peer.
288    pub(crate) fn process_transmission_from_peer(
289        &self,
290        peer_ip: SocketAddr,
291        transmission_id: TransmissionID<N>,
292        transmission: Transmission<N>,
293    ) {
294        // If the transmission ID already exists, then do not store it.
295        if self.contains_transmission(transmission_id) {
296            return;
297        }
298        // Ensure the transmission ID and transmission type matches.
299        let is_well_formed = match (&transmission_id, &transmission) {
300            (TransmissionID::Solution(_, _), Transmission::Solution(_)) => true,
301            (TransmissionID::Transaction(_, _), Transmission::Transaction(_)) => true,
302            // Note: We explicitly forbid inserting ratifications into the ready queue,
303            // as the protocol currently does not support ratifications.
304            (TransmissionID::Ratification, Transmission::Ratification) => false,
305            // All other combinations are clearly invalid.
306            _ => false,
307        };
308        // If the transmission ID and transmission type matches, then insert the transmission into the ready queue.
309        if is_well_formed && self.ready.insert(transmission_id, transmission) {
310            trace!(
311                "Worker {} - Added transmission '{}.{}' from '{peer_ip}'",
312                self.id,
313                fmt_id(transmission_id),
314                fmt_id(transmission_id.checksum().unwrap_or_default()).dimmed()
315            );
316        }
317    }
318
319    /// Handles the incoming unconfirmed solution.
320    /// Note: This method assumes the incoming solution is valid and does not exist in the ledger.
321    pub(crate) async fn process_unconfirmed_solution(
322        &self,
323        solution_id: SolutionID<N>,
324        solution: Data<Solution<N>>,
325    ) -> Result<()> {
326        // Construct the transmission.
327        let transmission = Transmission::Solution(solution.clone());
328        // Compute the checksum.
329        let checksum = solution.to_checksum::<N>()?;
330        // Construct the transmission ID.
331        let transmission_id = TransmissionID::Solution(solution_id, checksum);
332        // Remove the solution ID from the pending queue.
333        self.pending.remove(transmission_id, Some(transmission.clone()));
334        // Check if the solution exists.
335        if self.contains_transmission(transmission_id) {
336            bail!("Solution '{}.{}' already exists.", fmt_id(solution_id), fmt_id(checksum).dimmed());
337        }
338        // Check that the solution is well-formed and unique.
339        self.ledger.check_solution_basic(solution_id, solution).await?;
340        // Adds the solution to the ready queue.
341        if self.ready.insert(transmission_id, transmission) {
342            trace!(
343                "Worker {} - Added unconfirmed solution '{}.{}'",
344                self.id,
345                fmt_id(solution_id),
346                fmt_id(checksum).dimmed()
347            );
348        }
349        Ok(())
350    }
351
352    /// Handles the incoming unconfirmed transaction.
353    pub(crate) async fn process_unconfirmed_transaction(
354        &self,
355        transaction_id: N::TransactionID,
356        transaction: Data<Transaction<N>>,
357    ) -> Result<()> {
358        // Construct the transmission.
359        let transmission = Transmission::Transaction(transaction.clone());
360        // Compute the checksum.
361        let checksum = transaction.to_checksum::<N>()?;
362        // Construct the transmission ID.
363        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
364        // Remove the transaction from the pending queue.
365        self.pending.remove(transmission_id, Some(transmission.clone()));
366        // Check if the transaction ID exists.
367        if self.contains_transmission(transmission_id) {
368            bail!("Transaction '{}.{}' already exists.", fmt_id(transaction_id), fmt_id(checksum).dimmed());
369        }
370        // Check that the transaction is well-formed and unique.
371        self.ledger.check_transaction_basic(transaction_id, transaction).await?;
372        // Adds the transaction to the ready queue.
373        if self.ready.insert(transmission_id, transmission) {
374            trace!(
375                "Worker {}.{} - Added unconfirmed transaction '{}'",
376                self.id,
377                fmt_id(transaction_id),
378                fmt_id(checksum).dimmed()
379            );
380        }
381        Ok(())
382    }
383}
384
385impl<N: Network> Worker<N> {
386    /// Starts the worker handlers.
387    fn start_handlers(&self, receiver: WorkerReceiver<N>) {
388        let WorkerReceiver { mut rx_worker_ping, mut rx_transmission_request, mut rx_transmission_response } = receiver;
389
390        // Start the pending queue expiration loop.
391        let self_ = self.clone();
392        self.spawn(async move {
393            loop {
394                // Sleep briefly.
395                tokio::time::sleep(Duration::from_millis(MAX_FETCH_TIMEOUT_IN_MS)).await;
396
397                // Remove the expired pending certificate requests.
398                let self__ = self_.clone();
399                let _ = spawn_blocking!({
400                    self__.pending.clear_expired_callbacks();
401                    Ok(())
402                });
403            }
404        });
405
406        // Process the ping events.
407        let self_ = self.clone();
408        self.spawn(async move {
409            while let Some((peer_ip, transmission_id)) = rx_worker_ping.recv().await {
410                self_.process_transmission_id_from_ping(peer_ip, transmission_id);
411            }
412        });
413
414        // Process the transmission requests.
415        let self_ = self.clone();
416        self.spawn(async move {
417            while let Some((peer_ip, transmission_request)) = rx_transmission_request.recv().await {
418                self_.send_transmission_response(peer_ip, transmission_request);
419            }
420        });
421
422        // Process the transmission responses.
423        let self_ = self.clone();
424        self.spawn(async move {
425            while let Some((peer_ip, transmission_response)) = rx_transmission_response.recv().await {
426                // Process the transmission response.
427                let self__ = self_.clone();
428                let _ = spawn_blocking!({
429                    self__.finish_transmission_request(peer_ip, transmission_response);
430                    Ok(())
431                });
432            }
433        });
434    }
435
436    /// Sends a transmission request to the specified peer.
437    async fn send_transmission_request(
438        &self,
439        peer_ip: SocketAddr,
440        transmission_id: TransmissionID<N>,
441    ) -> Result<(TransmissionID<N>, Transmission<N>)> {
442        // Initialize a oneshot channel.
443        let (callback_sender, callback_receiver) = oneshot::channel();
444        // Determine how many sent requests are pending.
445        let num_sent_requests = self.pending.num_sent_requests(transmission_id);
446        // Determine if we've already sent a request to the peer.
447        let contains_peer_with_sent_request = self.pending.contains_peer_with_sent_request(transmission_id, peer_ip);
448        // Determine the maximum number of redundant requests.
449        let num_redundant_requests = max_redundant_requests(self.ledger.clone(), self.storage.current_round());
450        // Determine if we should send a transmission request to the peer.
451        // We send at most `num_redundant_requests` requests and each peer can only receive one request at a time.
452        let should_send_request = num_sent_requests < num_redundant_requests && !contains_peer_with_sent_request;
453
454        // Insert the transmission ID into the pending queue.
455        self.pending.insert(transmission_id, peer_ip, Some((callback_sender, should_send_request)));
456
457        // If the number of requests is less than or equal to the the redundancy factor, send the transmission request to the peer.
458        if should_send_request {
459            // Send the transmission request to the peer.
460            if self.gateway.send(peer_ip, Event::TransmissionRequest(transmission_id.into())).await.is_none() {
461                bail!("Unable to fetch transmission - failed to send request")
462            }
463        } else {
464            debug!(
465                "Skipped sending request for transmission {}.{} to '{peer_ip}' ({num_sent_requests} redundant requests)",
466                fmt_id(transmission_id),
467                fmt_id(transmission_id.checksum().unwrap_or_default()).dimmed()
468            );
469        }
470        // Wait for the transmission to be fetched.
471        match timeout(Duration::from_millis(MAX_FETCH_TIMEOUT_IN_MS), callback_receiver).await {
472            // If the transmission was fetched, return it.
473            Ok(result) => Ok((transmission_id, result?)),
474            // If the transmission was not fetched, return an error.
475            Err(e) => bail!("Unable to fetch transmission - (timeout) {e}"),
476        }
477    }
478
479    /// Handles the incoming transmission response.
480    /// This method ensures the transmission response is well-formed and matches the transmission ID.
481    fn finish_transmission_request(&self, peer_ip: SocketAddr, response: TransmissionResponse<N>) {
482        let TransmissionResponse { transmission_id, mut transmission } = response;
483        // Check if the peer IP exists in the pending queue for the given transmission ID.
484        let exists = self.pending.get_peers(transmission_id).unwrap_or_default().contains(&peer_ip);
485        // If the peer IP exists, finish the pending request.
486        if exists {
487            // Ensure the transmission is not a fee and matches the transmission ID.
488            match self.ledger.ensure_transmission_is_well_formed(transmission_id, &mut transmission) {
489                Ok(()) => {
490                    // Remove the transmission ID from the pending queue.
491                    self.pending.remove(transmission_id, Some(transmission));
492                }
493                Err(err) => warn!("Failed to finish transmission response from peer '{peer_ip}': {err}"),
494            };
495        }
496    }
497
498    /// Sends the requested transmission to the specified peer.
499    fn send_transmission_response(&self, peer_ip: SocketAddr, request: TransmissionRequest<N>) {
500        let TransmissionRequest { transmission_id } = request;
501        // Attempt to retrieve the transmission.
502        if let Some(transmission) = self.get_transmission(transmission_id) {
503            // Send the transmission response to the peer.
504            let self_ = self.clone();
505            tokio::spawn(async move {
506                self_.gateway.send(peer_ip, Event::TransmissionResponse((transmission_id, transmission).into())).await;
507            });
508        }
509    }
510
511    /// Spawns a task with the given future; it should only be used for long-running tasks.
512    fn spawn<T: Future<Output = ()> + Send + 'static>(&self, future: T) {
513        self.handles.lock().push(tokio::spawn(future));
514    }
515
516    /// Shuts down the worker.
517    pub(crate) fn shut_down(&self) {
518        trace!("Shutting down worker {}...", self.id);
519        // Abort the tasks.
520        self.handles.lock().iter().for_each(|handle| handle.abort());
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use super::*;
527    use crate::helpers::CALLBACK_EXPIRATION_IN_SECS;
528    use snarkos_node_bft_ledger_service::LedgerService;
529    use snarkos_node_bft_storage_service::BFTMemoryService;
530    use snarkvm::{
531        console::{network::Network, types::Field},
532        ledger::{
533            block::Block,
534            committee::Committee,
535            narwhal::{BatchCertificate, Subdag, Transmission, TransmissionID},
536        },
537        prelude::Address,
538    };
539
540    use bytes::Bytes;
541    use indexmap::IndexMap;
542    use mockall::mock;
543    use std::{io, ops::Range};
544
545    type CurrentNetwork = snarkvm::prelude::MainnetV0;
546
547    const ITERATIONS: usize = 100;
548
549    mock! {
550        Gateway<N: Network> {}
551        #[async_trait]
552        impl<N:Network> Transport<N> for Gateway<N> {
553            fn broadcast(&self, event: Event<N>);
554            async fn send(&self, peer_ip: SocketAddr, event: Event<N>) -> Option<oneshot::Receiver<io::Result<()>>>;
555        }
556    }
557
558    mock! {
559        #[derive(Debug)]
560        Ledger<N: Network> {}
561        #[async_trait]
562        impl<N: Network> LedgerService<N> for Ledger<N> {
563            fn latest_round(&self) -> u64;
564            fn latest_block_height(&self) -> u32;
565            fn latest_block(&self) -> Block<N>;
566            fn latest_restrictions_id(&self) -> Field<N>;
567            fn latest_leader(&self) -> Option<(u64, Address<N>)>;
568            fn update_latest_leader(&self, round: u64, leader: Address<N>);
569            fn contains_block_height(&self, height: u32) -> bool;
570            fn get_block_height(&self, hash: &N::BlockHash) -> Result<u32>;
571            fn get_block_hash(&self, height: u32) -> Result<N::BlockHash>;
572            fn get_block_round(&self, height: u32) -> Result<u64>;
573            fn get_block(&self, height: u32) -> Result<Block<N>>;
574            fn get_blocks(&self, heights: Range<u32>) -> Result<Vec<Block<N>>>;
575            fn get_solution(&self, solution_id: &SolutionID<N>) -> Result<Solution<N>>;
576            fn get_unconfirmed_transaction(&self, transaction_id: N::TransactionID) -> Result<Transaction<N>>;
577            fn get_batch_certificate(&self, certificate_id: &Field<N>) -> Result<BatchCertificate<N>>;
578            fn current_committee(&self) -> Result<Committee<N>>;
579            fn get_committee_for_round(&self, round: u64) -> Result<Committee<N>>;
580            fn get_committee_lookback_for_round(&self, round: u64) -> Result<Committee<N>>;
581            fn contains_certificate(&self, certificate_id: &Field<N>) -> Result<bool>;
582            fn contains_transmission(&self, transmission_id: &TransmissionID<N>) -> Result<bool>;
583            fn ensure_transmission_is_well_formed(
584                &self,
585                transmission_id: TransmissionID<N>,
586                transmission: &mut Transmission<N>,
587            ) -> Result<()>;
588            async fn check_solution_basic(
589                &self,
590                solution_id: SolutionID<N>,
591                solution: Data<Solution<N>>,
592            ) -> Result<()>;
593            async fn check_transaction_basic(
594                &self,
595                transaction_id: N::TransactionID,
596                transaction: Data<Transaction<N>>,
597            ) -> Result<()>;
598            fn check_next_block(&self, block: &Block<N>) -> Result<()>;
599            fn prepare_advance_to_next_quorum_block(
600                &self,
601                subdag: Subdag<N>,
602                transmissions: IndexMap<TransmissionID<N>, Transmission<N>>,
603            ) -> Result<Block<N>>;
604            fn advance_to_next_block(&self, block: &Block<N>) -> Result<()>;
605        }
606    }
607
608    #[tokio::test]
609    async fn test_max_redundant_requests() {
610        const NUM_NODES: u16 = Committee::<CurrentNetwork>::MAX_COMMITTEE_SIZE;
611
612        let rng = &mut TestRng::default();
613        // Sample a committee.
614        let committee =
615            snarkvm::ledger::committee::test_helpers::sample_committee_for_round_and_size(0, NUM_NODES, rng);
616        let committee_clone = committee.clone();
617        // Setup the mock ledger.
618        let mut mock_ledger = MockLedger::default();
619        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
620        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
621        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
622        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
623        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
624
625        // Ensure the maximum number of redundant requests is correct and consistent across iterations.
626        assert_eq!(max_redundant_requests(ledger, 0), 34, "Update me if the formula changes");
627    }
628
629    #[tokio::test]
630    async fn test_process_transmission() {
631        let rng = &mut TestRng::default();
632        // Sample a committee.
633        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
634        let committee_clone = committee.clone();
635        // Setup the mock gateway and ledger.
636        let gateway = MockGateway::default();
637        let mut mock_ledger = MockLedger::default();
638        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
639        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
640        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
641        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
642        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
643        // Initialize the storage.
644        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
645
646        // Create the Worker.
647        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
648        let data = |rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
649        let transmission_id = TransmissionID::Solution(
650            rng.gen::<u64>().into(),
651            rng.gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
652        );
653        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
654        let transmission = Transmission::Solution(data(rng));
655
656        // Process the transmission.
657        worker.process_transmission_from_peer(peer_ip, transmission_id, transmission.clone());
658        assert!(worker.contains_transmission(transmission_id));
659        assert!(worker.ready.contains(transmission_id));
660        assert_eq!(worker.get_transmission(transmission_id), Some(transmission));
661        // Take the transmission from the ready set.
662        let transmission: Vec<_> = worker.drain(1).collect();
663        assert_eq!(transmission.len(), 1);
664        assert!(!worker.ready.contains(transmission_id));
665    }
666
667    #[tokio::test]
668    async fn test_send_transmission() {
669        let rng = &mut TestRng::default();
670        // Sample a committee.
671        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
672        let committee_clone = committee.clone();
673        // Setup the mock gateway and ledger.
674        let mut gateway = MockGateway::default();
675        gateway.expect_send().returning(|_, _| {
676            let (_tx, rx) = oneshot::channel();
677            Some(rx)
678        });
679        let mut mock_ledger = MockLedger::default();
680        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
681        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
682        mock_ledger.expect_ensure_transmission_is_well_formed().returning(|_, _| Ok(()));
683        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
684        // Initialize the storage.
685        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
686
687        // Create the Worker.
688        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
689        let transmission_id = TransmissionID::Solution(
690            rng.gen::<u64>().into(),
691            rng.gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
692        );
693        let worker_ = worker.clone();
694        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
695        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
696        assert!(worker.pending.contains(transmission_id));
697        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
698        // Fake the transmission response.
699        worker.finish_transmission_request(peer_ip, TransmissionResponse {
700            transmission_id,
701            transmission: Transmission::Solution(Data::Buffer(Bytes::from(vec![0; 512]))),
702        });
703        // Check the transmission was removed from the pending set.
704        assert!(!worker.pending.contains(transmission_id));
705    }
706
707    #[tokio::test]
708    async fn test_process_solution_ok() {
709        let rng = &mut TestRng::default();
710        // Sample a committee.
711        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
712        let committee_clone = committee.clone();
713        // Setup the mock gateway and ledger.
714        let mut gateway = MockGateway::default();
715        gateway.expect_send().returning(|_, _| {
716            let (_tx, rx) = oneshot::channel();
717            Some(rx)
718        });
719        let mut mock_ledger = MockLedger::default();
720        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
721        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
722        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
723        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
724        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
725        // Initialize the storage.
726        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
727
728        // Create the Worker.
729        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
730        let solution = Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
731        let solution_id = rng.gen::<u64>().into();
732        let solution_checksum = solution.to_checksum::<CurrentNetwork>().unwrap();
733        let transmission_id = TransmissionID::Solution(solution_id, solution_checksum);
734        let worker_ = worker.clone();
735        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
736        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
737        assert!(worker.pending.contains(transmission_id));
738        let result = worker.process_unconfirmed_solution(solution_id, solution).await;
739        assert!(result.is_ok());
740        assert!(!worker.pending.contains(transmission_id));
741        assert!(worker.ready.contains(transmission_id));
742    }
743
744    #[tokio::test]
745    async fn test_process_solution_nok() {
746        let rng = &mut TestRng::default();
747        // Sample a committee.
748        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
749        let committee_clone = committee.clone();
750        // Setup the mock gateway and ledger.
751        let mut gateway = MockGateway::default();
752        gateway.expect_send().returning(|_, _| {
753            let (_tx, rx) = oneshot::channel();
754            Some(rx)
755        });
756        let mut mock_ledger = MockLedger::default();
757        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
758        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
759        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
760        mock_ledger.expect_check_solution_basic().returning(|_, _| Err(anyhow!("")));
761        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
762        // Initialize the storage.
763        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
764
765        // Create the Worker.
766        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
767        let solution_id = rng.gen::<u64>().into();
768        let solution = Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
769        let checksum = solution.to_checksum::<CurrentNetwork>().unwrap();
770        let transmission_id = TransmissionID::Solution(solution_id, checksum);
771        let worker_ = worker.clone();
772        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
773        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
774        assert!(worker.pending.contains(transmission_id));
775        let result = worker.process_unconfirmed_solution(solution_id, solution).await;
776        assert!(result.is_err());
777        assert!(!worker.pending.contains(transmission_id));
778        assert!(!worker.ready.contains(transmission_id));
779    }
780
781    #[tokio::test]
782    async fn test_process_transaction_ok() {
783        let mut rng = &mut TestRng::default();
784        // Sample a committee.
785        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
786        let committee_clone = committee.clone();
787        // Setup the mock gateway and ledger.
788        let mut gateway = MockGateway::default();
789        gateway.expect_send().returning(|_, _| {
790            let (_tx, rx) = oneshot::channel();
791            Some(rx)
792        });
793        let mut mock_ledger = MockLedger::default();
794        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
795        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
796        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
797        mock_ledger.expect_check_transaction_basic().returning(|_, _| Ok(()));
798        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
799        // Initialize the storage.
800        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
801
802        // Create the Worker.
803        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
804        let transaction_id: <CurrentNetwork as Network>::TransactionID = Field::<CurrentNetwork>::rand(&mut rng).into();
805        let transaction = Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
806        let checksum = transaction.to_checksum::<CurrentNetwork>().unwrap();
807        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
808        let worker_ = worker.clone();
809        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
810        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
811        assert!(worker.pending.contains(transmission_id));
812        let result = worker.process_unconfirmed_transaction(transaction_id, transaction).await;
813        assert!(result.is_ok());
814        assert!(!worker.pending.contains(transmission_id));
815        assert!(worker.ready.contains(transmission_id));
816    }
817
818    #[tokio::test]
819    async fn test_process_transaction_nok() {
820        let mut rng = &mut TestRng::default();
821        // Sample a committee.
822        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
823        let committee_clone = committee.clone();
824        // Setup the mock gateway and ledger.
825        let mut gateway = MockGateway::default();
826        gateway.expect_send().returning(|_, _| {
827            let (_tx, rx) = oneshot::channel();
828            Some(rx)
829        });
830        let mut mock_ledger = MockLedger::default();
831        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
832        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
833        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
834        mock_ledger.expect_check_transaction_basic().returning(|_, _| Err(anyhow!("")));
835        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
836        // Initialize the storage.
837        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
838
839        // Create the Worker.
840        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
841        let transaction_id: <CurrentNetwork as Network>::TransactionID = Field::<CurrentNetwork>::rand(&mut rng).into();
842        let transaction = Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
843        let checksum = transaction.to_checksum::<CurrentNetwork>().unwrap();
844        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
845        let worker_ = worker.clone();
846        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
847        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
848        assert!(worker.pending.contains(transmission_id));
849        let result = worker.process_unconfirmed_transaction(transaction_id, transaction).await;
850        assert!(result.is_err());
851        assert!(!worker.pending.contains(transmission_id));
852        assert!(!worker.ready.contains(transmission_id));
853    }
854
855    #[tokio::test]
856    async fn test_flood_transmission_requests() {
857        let mut rng = &mut TestRng::default();
858        // Sample a committee.
859        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
860        let committee_clone = committee.clone();
861        // Setup the mock gateway and ledger.
862        let mut gateway = MockGateway::default();
863        gateway.expect_send().returning(|_, _| {
864            let (_tx, rx) = oneshot::channel();
865            Some(rx)
866        });
867        let mut mock_ledger = MockLedger::default();
868        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
869        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
870        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
871        mock_ledger.expect_check_transaction_basic().returning(|_, _| Ok(()));
872        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
873        // Initialize the storage.
874        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1);
875
876        // Create the Worker.
877        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
878        let transaction_id: <CurrentNetwork as Network>::TransactionID = Field::<CurrentNetwork>::rand(&mut rng).into();
879        let transaction = Data::Buffer(Bytes::from((0..512).map(|_| rng.gen::<u8>()).collect::<Vec<_>>()));
880        let checksum = transaction.to_checksum::<CurrentNetwork>().unwrap();
881        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
882
883        // Determine the number of redundant requests are sent.
884        let num_redundant_requests = max_redundant_requests(worker.ledger.clone(), worker.storage.current_round());
885        let num_flood_requests = num_redundant_requests * 10;
886        let mut peer_ips =
887            (0..num_flood_requests).map(|i| SocketAddr::from(([127, 0, 0, 1], 1234 + i as u16))).collect_vec();
888        let first_peer_ip = peer_ips[0];
889
890        // Flood the pending queue with transmission requests.
891        for i in 1..=num_flood_requests {
892            let worker_ = worker.clone();
893            let peer_ip = peer_ips.pop().unwrap();
894            tokio::spawn(async move {
895                let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
896            });
897            tokio::time::sleep(Duration::from_millis(10)).await;
898            // Check that the number of sent requests does not exceed the maximum number of redundant requests.
899            assert!(worker.pending.num_sent_requests(transmission_id) <= num_redundant_requests);
900            assert_eq!(worker.pending.num_callbacks(transmission_id), i);
901        }
902        // Check that the number of sent requests does not exceed the maximum number of redundant requests.
903        assert_eq!(worker.pending.num_sent_requests(transmission_id), num_redundant_requests);
904        assert_eq!(worker.pending.num_callbacks(transmission_id), num_flood_requests);
905
906        // Let all the requests expire.
907        tokio::time::sleep(Duration::from_secs(CALLBACK_EXPIRATION_IN_SECS as u64 + 1)).await;
908        assert_eq!(worker.pending.num_sent_requests(transmission_id), 0);
909        assert_eq!(worker.pending.num_callbacks(transmission_id), 0);
910
911        // Flood the pending queue with transmission requests again, this time to a single peer
912        for i in 1..=num_flood_requests {
913            let worker_ = worker.clone();
914            tokio::spawn(async move {
915                let _ = worker_.send_transmission_request(first_peer_ip, transmission_id).await;
916            });
917            tokio::time::sleep(Duration::from_millis(10)).await;
918            assert!(worker.pending.num_sent_requests(transmission_id) <= num_redundant_requests);
919            assert_eq!(worker.pending.num_callbacks(transmission_id), i);
920        }
921        // Check that the number of sent requests does not exceed the maximum number of redundant requests.
922        assert_eq!(worker.pending.num_sent_requests(transmission_id), 1);
923        assert_eq!(worker.pending.num_callbacks(transmission_id), num_flood_requests);
924
925        // Check that fulfilling a transmission request clears the pending queue.
926        let result = worker.process_unconfirmed_transaction(transaction_id, transaction).await;
927        assert!(result.is_ok());
928        assert_eq!(worker.pending.num_sent_requests(transmission_id), 0);
929        assert_eq!(worker.pending.num_callbacks(transmission_id), 0);
930        assert!(!worker.pending.contains(transmission_id));
931        assert!(worker.ready.contains(transmission_id));
932    }
933
934    #[tokio::test]
935    async fn test_storage_gc_on_initialization() {
936        let rng = &mut TestRng::default();
937
938        for _ in 0..ITERATIONS {
939            // Mock the ledger round.
940            let max_gc_rounds = rng.gen_range(50..=100);
941            let latest_ledger_round = rng.gen_range((max_gc_rounds + 1)..1000);
942            let expected_gc_round = latest_ledger_round - max_gc_rounds;
943
944            // Sample a committee.
945            let committee =
946                snarkvm::ledger::committee::test_helpers::sample_committee_for_round(latest_ledger_round, rng);
947
948            // Setup the mock gateway and ledger.
949            let mut mock_ledger = MockLedger::default();
950            mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
951
952            let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
953            // Initialize the storage.
954            let storage =
955                Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), max_gc_rounds);
956
957            // Ensure that the storage GC round is correct.
958            assert_eq!(storage.gc_round(), expected_gc_round);
959        }
960    }
961}
962
963#[cfg(test)]
964mod prop_tests {
965    use super::*;
966    use crate::Gateway;
967    use snarkos_node_bft_ledger_service::MockLedgerService;
968    use snarkvm::{
969        console::account::Address,
970        ledger::committee::{Committee, MIN_VALIDATOR_STAKE},
971    };
972
973    use test_strategy::proptest;
974
975    type CurrentNetwork = snarkvm::prelude::MainnetV0;
976
977    // Initializes a new test committee.
978    fn new_test_committee(n: u16) -> Committee<CurrentNetwork> {
979        let mut members = IndexMap::with_capacity(n as usize);
980        for i in 0..n {
981            // Sample the address.
982            let rng = &mut TestRng::fixed(i as u64);
983            let address = Address::new(rng.gen());
984            info!("Validator {i}: {address}");
985            members.insert(address, (MIN_VALIDATOR_STAKE, false, rng.gen_range(0..100)));
986        }
987        // Initialize the committee.
988        Committee::<CurrentNetwork>::new(1u64, members).unwrap()
989    }
990
991    #[proptest]
992    fn worker_initialization(
993        #[strategy(0..MAX_WORKERS)] id: u8,
994        gateway: Gateway<CurrentNetwork>,
995        storage: Storage<CurrentNetwork>,
996    ) {
997        let committee = new_test_committee(4);
998        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(MockLedgerService::new(committee));
999        let worker = Worker::new(id, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
1000        assert_eq!(worker.id(), id);
1001    }
1002
1003    #[proptest]
1004    fn invalid_worker_id(
1005        #[strategy(MAX_WORKERS..)] id: u8,
1006        gateway: Gateway<CurrentNetwork>,
1007        storage: Storage<CurrentNetwork>,
1008    ) {
1009        let committee = new_test_committee(4);
1010        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(MockLedgerService::new(committee));
1011        let worker = Worker::new(id, Arc::new(gateway), storage, ledger, Default::default());
1012        // TODO once Worker implements Debug, simplify this with `unwrap_err`
1013        if let Err(error) = worker {
1014            assert_eq!(error.to_string(), format!("Invalid worker ID '{}'", id));
1015        }
1016    }
1017}