Skip to main content

snarkos_node_bft/
worker.rs

1// Copyright (c) 2019-2026 Provable Inc.
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
16#[cfg(not(test))]
17use crate::Gateway;
18use crate::{
19    MAX_FETCH_TIMEOUT,
20    MAX_WORKERS,
21    ProposedBatch,
22    ProposedBatchState,
23    Transport,
24    events::{Event, TransmissionRequest, TransmissionResponse},
25    helpers::{Pending, Ready, Storage, WorkerReceiver, fmt_id, max_redundant_requests},
26    spawn_blocking,
27};
28use snarkos_node_bft_ledger_service::{LedgerService, deserialize_transaction_strict};
29use snarkvm::{
30    console::prelude::*,
31    ledger::{
32        Transaction,
33        narwhal::{BatchHeader, Data, Transmission, TransmissionID},
34        puzzle::{Solution, SolutionID},
35    },
36};
37
38use anyhow::Context;
39use colored::{ColoredString, Colorize};
40use indexmap::{IndexMap, IndexSet};
41#[cfg(feature = "locktick")]
42use locktick::parking_lot::{Mutex, RwLock};
43#[cfg(not(feature = "locktick"))]
44use parking_lot::{Mutex, RwLock};
45use rand::seq::IteratorRandom;
46
47use std::{future::Future, net::SocketAddr, sync::Arc};
48use tokio::{sync::oneshot, task::JoinHandle, time::timeout};
49
50/// A worker's main role is maintaining a queue of verified ("ready") transmissions,
51/// which will eventually be fetched by the primary when the primary generates a new batch.
52#[derive(Clone)]
53pub struct Worker<N: Network> {
54    /// The worker ID.
55    id: u8,
56    /// The gateway.
57    #[cfg(not(test))]
58    gateway: Arc<Gateway<N>>,
59    #[cfg(test)]
60    gateway: Arc<dyn Transport<N>>,
61    /// The storage.
62    storage: Storage<N>,
63    /// The ledger service.
64    ledger: Arc<dyn LedgerService<N>>,
65    /// The proposed batch.
66    proposed_batch: Arc<ProposedBatch<N>>,
67    /// The ready queue.
68    ready: Arc<RwLock<Ready<N>>>,
69    /// The pending transmissions queue.
70    pending: Arc<Pending<TransmissionID<N>, Transmission<N>>>,
71    /// The spawned handles.
72    handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
73}
74
75impl<N: Network> Worker<N> {
76    /// Initializes a new worker instance.
77    pub fn new(
78        id: u8,
79        #[cfg(not(test))] gateway: Arc<Gateway<N>>,
80        #[cfg(test)] gateway: Arc<dyn Transport<N>>,
81        storage: Storage<N>,
82        ledger: Arc<dyn LedgerService<N>>,
83        proposed_batch: Arc<ProposedBatch<N>>,
84    ) -> Result<Self> {
85        // Ensure the worker ID is valid.
86        ensure!(id < MAX_WORKERS, "Invalid worker ID '{id}'");
87        // Return the worker.
88        Ok(Self {
89            id,
90            gateway,
91            storage,
92            ledger,
93            proposed_batch,
94            ready: Default::default(),
95            pending: Default::default(),
96            handles: Default::default(),
97        })
98    }
99
100    /// Run the worker instance.
101    pub fn run(&self, receiver: WorkerReceiver<N>) {
102        info!("Starting worker instance {} of the memory pool...", self.id);
103        // Start the worker handlers.
104        self.start_handlers(receiver);
105    }
106
107    /// Returns the worker ID.
108    pub const fn id(&self) -> u8 {
109        self.id
110    }
111
112    /// Returns a reference to the pending transmissions queue.
113    pub fn pending(&self) -> &Arc<Pending<TransmissionID<N>, Transmission<N>>> {
114        &self.pending
115    }
116}
117
118impl<N: Network> Worker<N> {
119    /// The maximum number of transmissions allowed in a worker.
120    pub const MAX_TRANSMISSIONS_PER_WORKER: usize =
121        BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH / MAX_WORKERS as usize;
122    /// The maximum number of transmissions allowed in a worker ping.
123    pub const MAX_TRANSMISSIONS_PER_WORKER_PING: usize = BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH / 10;
124
125    /// Returns the number of transmissions in the ready queue.
126    pub fn num_transmissions(&self) -> usize {
127        self.ready.read().num_transmissions()
128    }
129
130    /// Returns the number of ratifications in the ready queue.
131    pub fn num_ratifications(&self) -> usize {
132        self.ready.read().num_ratifications()
133    }
134
135    /// Returns the number of solutions in the ready queue.
136    pub fn num_solutions(&self) -> usize {
137        self.ready.read().num_solutions()
138    }
139
140    /// Returns the number of transactions in the ready queue.
141    pub fn num_transactions(&self) -> usize {
142        self.ready.read().num_transactions()
143    }
144}
145
146impl<N: Network> Worker<N> {
147    /// Returns the transmission IDs in the ready queue.
148    pub fn transmission_ids(&self) -> IndexSet<TransmissionID<N>> {
149        self.ready.read().transmission_ids()
150    }
151
152    /// Returns the transmissions in the ready queue.
153    pub fn transmissions(&self) -> IndexMap<TransmissionID<N>, Transmission<N>> {
154        self.ready.read().transmissions()
155    }
156
157    /// Returns the solutions in the ready queue.
158    pub fn solutions(&self) -> impl '_ + Iterator<Item = (SolutionID<N>, Data<Solution<N>>)> {
159        self.ready.read().solutions().into_iter()
160    }
161
162    /// Returns the transactions in the ready queue.
163    pub fn transactions(&self) -> impl '_ + Iterator<Item = (N::TransactionID, Data<Transaction<N>>)> {
164        self.ready.read().transactions().into_iter()
165    }
166}
167
168impl<N: Network> Worker<N> {
169    /// Clears the solutions from the ready queue.
170    pub(super) fn clear_solutions(&self) {
171        self.ready.write().clear_solutions()
172    }
173}
174
175impl<N: Network> Worker<N> {
176    // Helper to print the transmission ID and checksum (if any).
177    fn format_transmission_id(&self, transmission_id: TransmissionID<N>) -> ColoredString {
178        if let Some(checksum) = transmission_id.checksum() {
179            // fmt_id will suffix with `..`, so we should not use `.`  as a separator.
180            format!("{}:{}", fmt_id(transmission_id), fmt_id(checksum))
181        } else {
182            fmt_id(transmission_id)
183        }
184        .dimmed()
185    }
186
187    /// Returns `true` if the transmission ID exists in the ready queue, proposed batch, storage, or ledger.
188    pub fn contains_transmission(&self, transmission_id: impl Into<TransmissionID<N>>) -> bool {
189        let transmission_id = transmission_id.into();
190        // Check if the transmission ID exists in the ready queue, proposed batch, storage, or ledger.
191        self.ready.read().contains(transmission_id)
192            || matches!(&*self.proposed_batch.read(), ProposedBatchState::Certifying(p) if p.contains_transmission(transmission_id))
193            || self.storage.contains_transmission(transmission_id)
194            || self.ledger.contains_transmission(&transmission_id).unwrap_or(false)
195    }
196
197    /// Returns the transmission if it exists in the ready queue, proposed batch, storage.
198    ///
199    /// Note: We explicitly forbid retrieving a transmission from the ledger, as transmissions
200    /// in the ledger are not guaranteed to be invalid for the current batch.
201    pub fn get_transmission(&self, transmission_id: TransmissionID<N>) -> Option<Transmission<N>> {
202        // Check if the transmission ID exists in the ready queue.
203        if let Some(transmission) = self.ready.read().get(transmission_id) {
204            return Some(transmission);
205        }
206        // Check if the transmission ID exists in storage.
207        if let Some(transmission) = self.storage.get_transmission(transmission_id) {
208            return Some(transmission);
209        }
210        // Check if the transmission ID exists in the proposed batch.
211        if let Some(transmission) = match &*self.proposed_batch.read() {
212            ProposedBatchState::Certifying(p) => p.get_transmission(transmission_id),
213            _ => None,
214        } {
215            return Some(transmission.clone());
216        }
217        None
218    }
219
220    /// Returns the transmissions if it exists in the worker, or requests it from the specified peer.
221    pub async fn get_or_fetch_transmission(
222        &self,
223        peer_ip: SocketAddr,
224        transmission_id: TransmissionID<N>,
225    ) -> Result<(TransmissionID<N>, Transmission<N>)> {
226        // Attempt to get the transmission from the worker.
227        if let Some(transmission) = self.get_transmission(transmission_id) {
228            return Ok((transmission_id, transmission));
229        }
230        // Send a transmission request to the peer.
231        let (candidate_id, transmission) = self.send_transmission_request(peer_ip, transmission_id).await?;
232        // Ensure the transmission ID matches.
233        ensure!(candidate_id == transmission_id, "Invalid transmission ID");
234        // Return the transmission.
235        Ok((transmission_id, transmission))
236    }
237
238    /// Inserts the transmission at the front of the ready queue.
239    pub(crate) fn insert_front(&self, key: TransmissionID<N>, value: Transmission<N>) {
240        self.ready.write().insert_front(key, value);
241    }
242
243    /// Removes and returns the transmission at the front of the ready queue.
244    pub(crate) fn remove_front(&self) -> Option<(TransmissionID<N>, Transmission<N>)> {
245        self.ready.write().remove_front()
246    }
247
248    /// Reinserts the specified transmission into the ready queue.
249    pub(crate) fn reinsert(&self, transmission_id: TransmissionID<N>, transmission: Transmission<N>) -> bool {
250        // Check if the transmission ID exists.
251        if !self.contains_transmission(transmission_id) {
252            // Insert the transmission into the ready queue.
253            return self.ready.write().insert(transmission_id, transmission);
254        }
255        false
256    }
257
258    /// Broadcasts a worker ping event.
259    pub(crate) fn broadcast_ping(&self) {
260        // Retrieve the transmission IDs.
261        let transmission_ids = self
262            .ready
263            .read()
264            .transmission_ids()
265            .into_iter()
266            .sample(&mut rand::rng(), Self::MAX_TRANSMISSIONS_PER_WORKER_PING)
267            .into_iter()
268            .collect::<IndexSet<_>>();
269
270        // Broadcast the ping event.
271        if !transmission_ids.is_empty() {
272            self.gateway.broadcast(Event::WorkerPing(transmission_ids.into()));
273        }
274    }
275}
276
277impl<N: Network> Worker<N> {
278    /// Handles the incoming transmission ID from a worker ping event.
279    fn process_transmission_id_from_ping(&self, peer_ip: SocketAddr, transmission_id: TransmissionID<N>) {
280        // Check if the transmission ID exists.
281        if self.contains_transmission(transmission_id) {
282            return;
283        }
284        // If the ready queue is full, then skip this transmission.
285        // Note: We must prioritize the unconfirmed solutions and unconfirmed transactions, not transmissions.
286        if self.ready.read().num_transmissions() > Self::MAX_TRANSMISSIONS_PER_WORKER {
287            return;
288        }
289        // Attempt to fetch the transmission from the peer.
290        let self_ = self.clone();
291        tokio::spawn(async move {
292            // Send a transmission request to the peer.
293            match self_.send_transmission_request(peer_ip, transmission_id).await {
294                // If the transmission was fetched, then process it.
295                Ok((candidate_id, transmission)) => {
296                    // Ensure the transmission ID matches.
297                    if candidate_id == transmission_id {
298                        // Insert the transmission into the ready queue.
299                        // Note: This method checks `contains_transmission` again, because by the time the transmission is fetched,
300                        // it could have already been inserted into the ready queue.
301                        self_.process_transmission_from_ping(peer_ip, transmission_id, transmission);
302                    }
303                }
304                // If the transmission was not fetched, then attempt to fetch it again.
305                Err(e) => {
306                    warn!(
307                        "Worker {} - Failed to fetch transmission '{}' from '{peer_ip}' (ping) - {e}",
308                        self_.id,
309                        self_.format_transmission_id(transmission_id),
310                    );
311                }
312            }
313        });
314    }
315
316    /// Handles a transmission fetched in response to a worker ping event.
317    ///
318    /// Unlike [`Worker::process_transmission_from_peer`], this enforces the ready-queue capacity
319    /// limit (atomically with the insertion), dropping the transmission if the queue is full.
320    ///
321    /// This is required to close a time-of-check-to-time-of-use race: the capacity check in
322    /// `process_transmission_id_from_ping` runs before the transmission is fetched asynchronously,
323    /// so without a re-check at insertion time, many concurrent fetches could each observe a
324    /// non-full ready queue and then collectively exceed `MAX_TRANSMISSIONS_PER_WORKER`.
325    ///
326    /// Note: We must prioritize the unconfirmed solutions and unconfirmed transactions, not
327    /// transmissions fetched from peer pings, which is why the limit is enforced on this path.
328    fn process_transmission_from_ping(
329        &self,
330        peer_ip: SocketAddr,
331        transmission_id: TransmissionID<N>,
332        transmission: Transmission<N>,
333    ) {
334        self.insert_transmission_from_peer(peer_ip, transmission_id, transmission, true);
335    }
336
337    /// Handles the incoming transmission from a peer.
338    ///
339    /// Note: This does *not* enforce the ready-queue capacity limit, because this path also
340    /// materializes the transmissions of a peer's batch proposal prior to signing it (see
341    /// `Primary::insert_missing_transmissions_into_workers`); dropping a transmission there would
342    /// cause the node to sign a batch whose transmissions it does not hold. The capacity limit is
343    /// enforced only on the worker-ping path, via [`Worker::process_transmission_from_ping`].
344    pub(crate) fn process_transmission_from_peer(
345        &self,
346        peer_ip: SocketAddr,
347        transmission_id: TransmissionID<N>,
348        transmission: Transmission<N>,
349    ) {
350        self.insert_transmission_from_peer(peer_ip, transmission_id, transmission, false);
351    }
352
353    /// Inserts a peer-provided transmission into the ready queue.
354    ///
355    /// If `enforce_capacity_limit` is `true`, the transmission is dropped when the ready queue is
356    /// already full. This must only be set for the worker-ping fetch path; callers that must retain
357    /// the transmission (e.g. materializing a peer's batch proposal before signing it) must leave it
358    /// `false`. See [`Worker::process_transmission_from_ping`] and
359    /// [`Worker::process_transmission_from_peer`].
360    fn insert_transmission_from_peer(
361        &self,
362        peer_ip: SocketAddr,
363        transmission_id: TransmissionID<N>,
364        transmission: Transmission<N>,
365        enforce_capacity_limit: bool,
366    ) {
367        // If the transmission ID already exists, then do not store it.
368        if self.contains_transmission(transmission_id) {
369            return;
370        }
371        // Ensure the transmission ID and transmission type matches.
372        let is_well_formed = match (&transmission_id, &transmission) {
373            (TransmissionID::Solution(_, _), Transmission::Solution(_)) => true,
374            (TransmissionID::Transaction(_, _), Transmission::Transaction(_)) => true,
375            // Note: We explicitly forbid inserting ratifications into the ready queue,
376            // as the protocol currently does not support ratifications.
377            (TransmissionID::Ratification, Transmission::Ratification) => false,
378            // All other combinations are clearly invalid.
379            _ => false,
380        };
381        // If the transmission type does not match the transmission ID, then do not store it.
382        if !is_well_formed {
383            return;
384        }
385        // If the transmission is a deserialized execution, capture it so it can be verified
386        // immediately below (only if it is actually inserted). This takes heavy transaction
387        // verification out of the hot path during block generation.
388        let execution_to_verify = match (transmission_id, &transmission) {
389            (TransmissionID::Transaction(tx_id, _), Transmission::Transaction(Data::Object(tx))) if tx.is_execute() => {
390                Some((tx_id, tx.clone()))
391            }
392            _ => None,
393        };
394        // Insert the transmission into the ready queue. When the capacity limit is enforced, the
395        // check is performed atomically with the insertion (i.e. under a single write lock), so that
396        // concurrent inserts cannot collectively exceed `MAX_TRANSMISSIONS_PER_WORKER`.
397        let inserted = {
398            let mut ready = self.ready.write();
399            // If the ready queue is full, then skip this transmission.
400            if enforce_capacity_limit && ready.num_transmissions() > Self::MAX_TRANSMISSIONS_PER_WORKER {
401                return;
402            }
403            ready.insert(transmission_id, transmission)
404        };
405        // If the transmission was newly inserted, then process it further.
406        if inserted {
407            // Eagerly verify a newly-inserted execution to warm the verification cache.
408            if let Some((tx_id, tx)) = execution_to_verify {
409                let self_ = self.clone();
410                tokio::spawn(async move {
411                    let _ = self_.ledger.check_transaction_basic(tx_id, tx).await;
412                });
413            }
414            trace!(
415                "Worker {} - Added transmission '{}' from '{peer_ip}'",
416                self.id,
417                self.format_transmission_id(transmission_id),
418            );
419        }
420    }
421
422    /// Handles the incoming unconfirmed solution.
423    /// Note: This method assumes the incoming solution is valid and does not exist in the ledger.
424    ///
425    /// # Returns
426    /// - `Ok(true)` if the solution was added to the ready queue.
427    /// - `Ok(false)` if the solution was valid but already exists in the ready queue.
428    /// - `Err(anyhow::Error)` if the solution is invalid.
429    pub(crate) async fn process_unconfirmed_solution(
430        &self,
431        solution_id: SolutionID<N>,
432        solution: Data<Solution<N>>,
433    ) -> Result<bool> {
434        // Construct the transmission.
435        let transmission = Transmission::Solution(solution.clone());
436        // Compute the checksum.
437        let checksum = solution.to_checksum::<N>()?;
438        // Construct the transmission ID.
439        let transmission_id = TransmissionID::Solution(solution_id, checksum);
440        // Remove the solution ID from the pending queue.
441        self.pending.remove(transmission_id, Some(transmission.clone()));
442        // Check if the solution exists.
443        if self.contains_transmission(transmission_id) {
444            return Ok(false);
445        }
446        // Check that the solution is well-formed and unique.
447        self.ledger.check_solution_basic(solution_id, solution).await?;
448        // Adds the solution to the ready queue.
449        if self.ready.write().insert(transmission_id, transmission) {
450            trace!(
451                "Worker {} - Added unconfirmed solution '{}'",
452                self.id,
453                self.format_transmission_id(transmission_id),
454            );
455        }
456        Ok(true)
457    }
458
459    /// Handles the incoming unconfirmed transaction.
460    ///
461    /// # Returns
462    /// - `Ok(true)` if the transaction was added to the ready queue.
463    /// - `Ok(false)` if the transaction was valid but already exists in the ready queue.
464    /// - `Err(anyhow::Error)` if the transaction was invalid.
465    pub(crate) async fn process_unconfirmed_transaction(
466        &self,
467        transaction_id: N::TransactionID,
468        transaction: Data<Transaction<N>>,
469    ) -> Result<bool> {
470        // Construct the transmission.
471        let transmission = Transmission::Transaction(transaction.clone());
472        // Compute the checksum.
473        let checksum = transaction.to_checksum::<N>()?;
474        // Construct the transmission ID.
475        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
476        // Remove the transaction from the pending queue.
477        self.pending.remove(transmission_id, Some(transmission.clone()));
478        // Check if the transaction ID exists.
479        if self.contains_transmission(transmission_id) {
480            return Ok(false);
481        }
482        // Deserialize the transaction. If the transaction exceeds the maximum size, then return an error.
483        let transaction = spawn_blocking!(deserialize_transaction_strict(transaction))?;
484
485        // Check that the transaction is well-formed and unique.
486        self.ledger.check_transaction_basic(transaction_id, transaction).await?;
487        // Adds the transaction to the ready queue.
488        if self.ready.write().insert(transmission_id, transmission) {
489            trace!(
490                "Worker {} - Added unconfirmed transaction '{}'",
491                self.id,
492                self.format_transmission_id(transmission_id),
493            );
494        }
495        Ok(true)
496    }
497}
498
499impl<N: Network> Worker<N> {
500    /// Starts the worker handlers.
501    fn start_handlers(&self, receiver: WorkerReceiver<N>) {
502        let WorkerReceiver { mut rx_worker_ping, mut rx_transmission_request, mut rx_transmission_response } = receiver;
503
504        // Start the pending queue expiration loop.
505        let self_ = self.clone();
506        self.spawn(async move {
507            loop {
508                // Sleep briefly.
509                tokio::time::sleep(MAX_FETCH_TIMEOUT).await;
510
511                // Remove the expired pending certificate requests.
512                let self__ = self_.clone();
513                let _ = spawn_blocking!({
514                    self__.pending.clear_expired_callbacks();
515                    Ok(())
516                });
517            }
518        });
519
520        // Process the ping events.
521        let self_ = self.clone();
522        self.spawn(async move {
523            while let Some((peer_ip, transmission_id)) = rx_worker_ping.recv().await {
524                self_.process_transmission_id_from_ping(peer_ip, transmission_id);
525            }
526        });
527
528        // Process the transmission requests.
529        let self_ = self.clone();
530        self.spawn(async move {
531            while let Some((peer_ip, transmission_request)) = rx_transmission_request.recv().await {
532                self_.send_transmission_response(peer_ip, transmission_request);
533            }
534        });
535
536        // Process the transmission responses.
537        let self_ = self.clone();
538        self.spawn(async move {
539            while let Some((peer_ip, transmission_response)) = rx_transmission_response.recv().await {
540                // Process the transmission response.
541                let self__ = self_.clone();
542                let _ = spawn_blocking!({
543                    self__.finish_transmission_request(peer_ip, transmission_response);
544                    Ok(())
545                });
546            }
547        });
548    }
549
550    /// Sends a transmission request to the specified peer.
551    async fn send_transmission_request(
552        &self,
553        peer_ip: SocketAddr,
554        transmission_id: TransmissionID<N>,
555    ) -> Result<(TransmissionID<N>, Transmission<N>)> {
556        // Initialize a oneshot channel.
557        let (callback_sender, callback_receiver) = oneshot::channel();
558        // Determine how many sent requests are pending.
559        let num_sent_requests = self.pending.num_sent_requests(transmission_id);
560        // Determine if we've already sent a request to the peer.
561        let contains_peer_with_sent_request = self.pending.contains_peer_with_sent_request(transmission_id, peer_ip);
562        // Determine the maximum number of redundant requests.
563        let num_redundant_requests = max_redundant_requests(self.ledger.clone(), self.storage.current_round())?;
564        // Establish whether the peers who already got the request collectively hold sufficient stake.
565        #[cfg(test)]
566        let stake_redundancy_reached = || Ok::<_, anyhow::Error>(true);
567        #[cfg(not(test))]
568        let stake_redundancy_reached = || self.pending.request_stake_redundancy_reached(&self.gateway, transmission_id);
569        // Determine if we should send a transmission request to the peer.
570        // Each peer can only receive one request at a time.
571        // We send at most `num_redundant_requests` requests, unless the stake redundancy factor hasn't been reached.
572        let should_send_request = !contains_peer_with_sent_request
573            && (num_sent_requests < num_redundant_requests || !stake_redundancy_reached()?);
574
575        // Insert the transmission ID into the pending queue.
576        self.pending.insert(transmission_id, peer_ip, Some((callback_sender, should_send_request)));
577
578        // If the number of requests is less than or equal to the the redundancy factor, send the transmission request to the peer.
579        if should_send_request {
580            trace!("Requesting transmission {} from peer '{peer_ip}'", self.format_transmission_id(transmission_id));
581            // Send the transmission request to the peer.
582            if self.gateway.send(peer_ip, Event::TransmissionRequest(transmission_id.into())).await.is_none() {
583                bail!(
584                    "Unable to fetch transmission {} - failed to send request",
585                    self.format_transmission_id(transmission_id)
586                )
587            }
588        } else {
589            debug!(
590                "Skipped sending request for transmission {} to '{peer_ip}' ({num_sent_requests} redundant requests)",
591                self.format_transmission_id(transmission_id)
592            );
593        }
594
595        // Wait for the transmission to be fetched.
596        let transmission = timeout(MAX_FETCH_TIMEOUT, callback_receiver)
597            .await
598            .with_context(|| {
599                format!("Unable to fetch transmission {} (timeout)", self.format_transmission_id(transmission_id))
600            })?
601            .with_context(|| {
602                format!("Unable to fetch transmission {}", self.format_transmission_id(transmission_id))
603            })?;
604
605        Ok((transmission_id, transmission))
606    }
607
608    /// Handles the incoming transmission response.
609    /// This method ensures the transmission response is well-formed and matches the transmission ID.
610    fn finish_transmission_request(&self, peer_ip: SocketAddr, response: TransmissionResponse<N>) {
611        let TransmissionResponse { transmission_id, mut transmission } = response;
612        // Check if the peer IP exists in the pending queue for the given transmission ID.
613        let exists = self.pending.get_peers(transmission_id).unwrap_or_default().contains(&peer_ip);
614        // If the peer IP exists, finish the pending request.
615        if exists {
616            // Ensure the transmission is not a fee and matches the transmission ID.
617            match self.ledger.ensure_transmission_is_well_formed(transmission_id, &mut transmission) {
618                Ok(()) => {
619                    trace!(
620                        "Received valid transmission response from peer '{peer_ip}' for transmission '{}'",
621                        self.format_transmission_id(transmission_id)
622                    );
623                    // Remove the transmission ID from the pending queue.
624                    self.pending.remove(transmission_id, Some(transmission));
625                }
626                Err(err) => warn!("Failed to finish transmission response from peer '{peer_ip}': {err}"),
627            };
628        }
629    }
630
631    /// Sends the requested transmission to the specified peer.
632    fn send_transmission_response(&self, peer_ip: SocketAddr, request: TransmissionRequest<N>) {
633        let TransmissionRequest { transmission_id } = request;
634        // Attempt to retrieve the transmission.
635        if let Some(transmission) = self.get_transmission(transmission_id) {
636            // Send the transmission response to the peer.
637            let self_ = self.clone();
638            tokio::spawn(async move {
639                self_.gateway.send(peer_ip, Event::TransmissionResponse((transmission_id, transmission).into())).await;
640            });
641        }
642    }
643
644    /// Spawns a task with the given future; it should only be used for long-running tasks.
645    fn spawn<T: Future<Output = ()> + Send + 'static>(&self, future: T) {
646        self.handles.lock().push(tokio::spawn(future));
647    }
648
649    /// Shuts down the worker.
650    pub(crate) fn shut_down(&self) {
651        trace!("Shutting down worker {}...", self.id);
652        // Abort the tasks.
653        self.handles.lock().iter().for_each(|handle| handle.abort());
654    }
655}
656
657#[cfg(test)]
658mod tests {
659    use super::*;
660    use crate::helpers::CALLBACK_EXPIRATION_IN_SECS;
661    use snarkos_node_bft_ledger_service::{BeginLedgerUpdateError, LedgerService, LedgerUpdateService};
662    use snarkos_node_bft_storage_service::BFTMemoryService;
663    use snarkvm::{
664        console::{network::Network, types::Field},
665        ledger::{
666            Block,
667            CheckBlockError,
668            PendingBlock,
669            committee::Committee,
670            narwhal::{BatchCertificate, Transmission, TransmissionID},
671            test_helpers::sample_execution_transaction_with_fee,
672        },
673        prelude::Address,
674    };
675
676    use bytes::Bytes;
677    use mockall::mock;
678    use rand::RngExt;
679    use std::{io, ops::Range, time::Duration};
680
681    type CurrentNetwork = snarkvm::prelude::MainnetV0;
682
683    const ITERATIONS: usize = 100;
684
685    mock! {
686        Gateway<N: Network> {}
687        #[async_trait]
688        impl<N:Network> Transport<N> for Gateway<N> {
689            fn broadcast(&self, event: Event<N>);
690            async fn send(&self, peer_ip: SocketAddr, event: Event<N>) -> Option<oneshot::Receiver<io::Result<()>>>;
691        }
692    }
693
694    mock! {
695        #[derive(Debug)]
696        Ledger<N: Network> {}
697        #[async_trait]
698        impl<N: Network> LedgerService<N> for Ledger<N> {
699            fn latest_round(&self) -> u64;
700            fn latest_block_height(&self) -> u32;
701            fn latest_block(&self) -> Block<N>;
702            fn latest_restrictions_id(&self) -> Field<N>;
703            fn latest_leader(&self) -> Option<(u64, Address<N>)>;
704            fn update_latest_leader(&self, round: u64, leader: Address<N>);
705            fn contains_block_height(&self, height: u32) -> bool;
706            fn get_block_height(&self, hash: &N::BlockHash) -> Result<u32>;
707            fn get_block_hash(&self, height: u32) -> Result<N::BlockHash>;
708            fn get_block_round(&self, height: u32) -> Result<u64>;
709            fn get_block(&self, height: u32) -> Result<Block<N>>;
710            fn get_blocks(&self, heights: Range<u32>) -> Result<Vec<Block<N>>>;
711            fn get_solution(&self, solution_id: &SolutionID<N>) -> Result<Option<Solution<N>>>;
712            fn get_unconfirmed_transaction(&self, transaction_id: N::TransactionID) -> Result<Option<Transaction<N>>>;
713            fn get_batch_certificate(&self, certificate_id: &Field<N>) -> Result<BatchCertificate<N>>;
714            fn current_committee(&self) -> Result<Committee<N>>;
715            fn get_committee_for_round(&self, round: u64) -> Result<Committee<N>>;
716            fn get_committee_lookback_for_round(&self, round: u64) -> Result<Committee<N>>;
717            fn contains_certificate(&self, certificate_id: &Field<N>) -> Result<bool>;
718            fn contains_transmission(&self, transmission_id: &TransmissionID<N>) -> Result<bool>;
719            fn ensure_transmission_is_well_formed(
720                &self,
721                transmission_id: TransmissionID<N>,
722                transmission: &mut Transmission<N>,
723            ) -> Result<()>;
724            async fn check_solution_basic(
725                &self,
726                solution_id: SolutionID<N>,
727                solution: Data<Solution<N>>,
728            ) -> Result<()>;
729            async fn check_transaction_basic(
730                &self,
731                transaction_id: N::TransactionID,
732                transaction: Transaction<N>,
733            ) -> Result<()>;
734            fn check_block_subdag(&self, _block: Block<N>, _prefix: &[PendingBlock<N>]) -> Result<PendingBlock<N>, CheckBlockError<N>>;
735            fn begin_ledger_update<'a>(&'a self) -> Result<Box<dyn LedgerUpdateService<N> + 'a>, BeginLedgerUpdateError>;
736            fn transaction_spend_in_microcredits(&self, transaction: &Transaction<N>, consensus_version: ConsensusVersion) -> Result<u64>;
737            fn is_stopped(&self) -> bool;
738        }
739    }
740
741    #[tokio::test]
742    async fn test_max_redundant_requests() {
743        let num_nodes: u16 = CurrentNetwork::MAX_CERTIFICATES.first().unwrap().1;
744
745        let rng = &mut TestRng::default();
746        // Sample a committee.
747        let committee =
748            snarkvm::ledger::committee::test_helpers::sample_committee_for_round_and_size(0, num_nodes, rng);
749        let committee_clone = committee.clone();
750        // Setup the mock ledger.
751        let mut mock_ledger = MockLedger::default();
752        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
753        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
754        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
755        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
756        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
757
758        // Ensure the maximum number of redundant requests is correct and consistent across iterations.
759        assert_eq!(max_redundant_requests(ledger, 0).unwrap(), 6, "Update me if the formula changes");
760    }
761
762    #[tokio::test]
763    async fn test_process_transmission() {
764        let rng = &mut TestRng::default();
765        // Sample a committee.
766        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
767        let committee_clone = committee.clone();
768        // Setup the mock gateway and ledger.
769        let gateway = MockGateway::default();
770        let mut mock_ledger = MockLedger::default();
771        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
772        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
773        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
774        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
775        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
776        // Initialize the storage.
777        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
778
779        // Create the Worker.
780        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
781        let data =
782            |rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
783        let transmission_id = TransmissionID::Solution(
784            rng.random::<u64>().into(),
785            rng.random::<<CurrentNetwork as Network>::TransmissionChecksum>(),
786        );
787        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
788        let transmission = Transmission::Solution(data(rng));
789
790        // Process the transmission.
791        worker.process_transmission_from_peer(peer_ip, transmission_id, transmission.clone());
792        assert!(worker.contains_transmission(transmission_id));
793        assert!(worker.ready.read().contains(transmission_id));
794        assert_eq!(worker.get_transmission(transmission_id), Some(transmission));
795        // Take the transmission from the ready set.
796        assert!(worker.ready.write().remove_front().is_some());
797        assert!(!worker.ready.read().contains(transmission_id));
798    }
799
800    #[tokio::test]
801    async fn test_process_transmission_from_ping_respects_capacity() {
802        let rng = &mut TestRng::default();
803        // Sample a committee.
804        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
805        let committee_clone = committee.clone();
806        // Setup the mock gateway and ledger.
807        let gateway = MockGateway::default();
808        let mut mock_ledger = MockLedger::default();
809        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
810        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
811        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
812        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
813        // Initialize the storage.
814        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
815
816        // Create the Worker.
817        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
818        let data =
819            |rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
820        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
821
822        // Emulate the asynchronous worker-ping fetch path inserting far more transmissions than the
823        // per-worker limit. Each fetched transmission is inserted via `process_transmission_from_ping`,
824        // exactly as it would be when a `send_transmission_request` future resolves. Without an insertion
825        // time capacity check, the TOCTOU in `process_transmission_id_from_ping` lets these accumulate
826        // without bound.
827        let num_to_insert = Worker::<CurrentNetwork>::MAX_TRANSMISSIONS_PER_WORKER * 4;
828        for _ in 0..num_to_insert {
829            let transmission_id = TransmissionID::Solution(
830                rng.random::<u64>().into(),
831                rng.random::<<CurrentNetwork as Network>::TransmissionChecksum>(),
832            );
833            let transmission = Transmission::Solution(data(rng));
834            worker.process_transmission_from_ping(peer_ip, transmission_id, transmission);
835        }
836
837        // The worker-ping path must not grow the ready queue beyond the per-worker limit.
838        // At most one extra transmission is tolerated (the insertion that crosses the boundary).
839        assert!(
840            worker.num_transmissions() <= Worker::<CurrentNetwork>::MAX_TRANSMISSIONS_PER_WORKER + 1,
841            "ready queue exceeded the per-worker limit: {} > {}",
842            worker.num_transmissions(),
843            Worker::<CurrentNetwork>::MAX_TRANSMISSIONS_PER_WORKER + 1,
844        );
845    }
846
847    #[tokio::test]
848    async fn test_process_transmission_from_peer_ignores_capacity() {
849        let rng = &mut TestRng::default();
850        // Sample a committee.
851        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
852        let committee_clone = committee.clone();
853        // Setup the mock gateway and ledger.
854        let gateway = MockGateway::default();
855        let mut mock_ledger = MockLedger::default();
856        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
857        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
858        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
859        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
860        // Initialize the storage.
861        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
862
863        // Create the Worker.
864        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
865        let data =
866            |rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
867        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
868
869        // The batch-materialization path (`process_transmission_from_peer`) must retain every
870        // transmission, even beyond the per-worker limit, so the node never signs a peer's batch
871        // whose transmissions it does not hold.
872        let num_to_insert = Worker::<CurrentNetwork>::MAX_TRANSMISSIONS_PER_WORKER * 4;
873        for _ in 0..num_to_insert {
874            let transmission_id = TransmissionID::Solution(
875                rng.random::<u64>().into(),
876                rng.random::<<CurrentNetwork as Network>::TransmissionChecksum>(),
877            );
878            let transmission = Transmission::Solution(data(rng));
879            worker.process_transmission_from_peer(peer_ip, transmission_id, transmission);
880        }
881
882        // Every transmission was retained.
883        assert_eq!(worker.num_transmissions(), num_to_insert);
884    }
885
886    #[tokio::test]
887    async fn test_send_transmission() {
888        let rng = &mut TestRng::default();
889        // Sample a committee.
890        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
891        let committee_clone = committee.clone();
892        // Setup the mock gateway and ledger.
893        let mut gateway = MockGateway::default();
894        gateway.expect_send().returning(|_, _| {
895            let (_tx, rx) = oneshot::channel();
896            Some(rx)
897        });
898        let mut mock_ledger = MockLedger::default();
899        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
900        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
901        mock_ledger.expect_ensure_transmission_is_well_formed().returning(|_, _| Ok(()));
902        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
903        // Initialize the storage.
904        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
905
906        // Create the Worker.
907        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
908        let transmission_id = TransmissionID::Solution(
909            rng.random::<u64>().into(),
910            rng.random::<<CurrentNetwork as Network>::TransmissionChecksum>(),
911        );
912        let worker_ = worker.clone();
913        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
914        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
915        assert!(worker.pending.contains(transmission_id));
916        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
917        // Fake the transmission response.
918        worker.finish_transmission_request(peer_ip, TransmissionResponse {
919            transmission_id,
920            transmission: Transmission::Solution(Data::Buffer(Bytes::from(vec![0; 512]))),
921        });
922        // Check the transmission was removed from the pending set.
923        assert!(!worker.pending.contains(transmission_id));
924    }
925
926    #[tokio::test]
927    async fn test_process_solution_ok() {
928        let rng = &mut TestRng::default();
929        // Sample a committee.
930        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
931        let committee_clone = committee.clone();
932        // Setup the mock gateway and ledger.
933        let mut gateway = MockGateway::default();
934        gateway.expect_send().returning(|_, _| {
935            let (_tx, rx) = oneshot::channel();
936            Some(rx)
937        });
938        let mut mock_ledger = MockLedger::default();
939        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
940        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
941        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
942        mock_ledger.expect_check_solution_basic().returning(|_, _| Ok(()));
943        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
944        // Initialize the storage.
945        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
946
947        // Create the Worker.
948        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
949        let solution = Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
950        let solution_id = rng.random::<u64>().into();
951        let solution_checksum = solution.to_checksum::<CurrentNetwork>().unwrap();
952        let transmission_id = TransmissionID::Solution(solution_id, solution_checksum);
953        let worker_ = worker.clone();
954        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
955        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
956        assert!(worker.pending.contains(transmission_id));
957        let result = worker.process_unconfirmed_solution(solution_id, solution).await;
958        assert!(result.is_ok());
959        assert!(!worker.pending.contains(transmission_id));
960        assert!(worker.ready.read().contains(transmission_id));
961    }
962
963    #[tokio::test]
964    async fn test_process_solution_nok() {
965        let rng = &mut TestRng::default();
966        // Sample a committee.
967        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
968        let committee_clone = committee.clone();
969        // Setup the mock gateway and ledger.
970        let mut gateway = MockGateway::default();
971        gateway.expect_send().returning(|_, _| {
972            let (_tx, rx) = oneshot::channel();
973            Some(rx)
974        });
975        let mut mock_ledger = MockLedger::default();
976        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
977        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
978        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
979        mock_ledger.expect_check_solution_basic().returning(|_, _| Err(anyhow!("")));
980        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
981        // Initialize the storage.
982        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
983
984        // Create the Worker.
985        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
986        let solution_id = rng.random::<u64>().into();
987        let solution = Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
988        let checksum = solution.to_checksum::<CurrentNetwork>().unwrap();
989        let transmission_id = TransmissionID::Solution(solution_id, checksum);
990        let worker_ = worker.clone();
991        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
992        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
993        assert!(worker.pending.contains(transmission_id));
994        let result = worker.process_unconfirmed_solution(solution_id, solution).await;
995        assert!(result.is_err());
996        assert!(!worker.pending.contains(transmission_id));
997        assert!(!worker.ready.read().contains(transmission_id));
998    }
999
1000    #[tokio::test]
1001    async fn test_process_transaction_ok() {
1002        let rng = &mut TestRng::default();
1003        // Sample a committee.
1004        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
1005        let committee_clone = committee.clone();
1006        // Setup the mock gateway and ledger.
1007        let mut gateway = MockGateway::default();
1008        gateway.expect_send().returning(|_, _| {
1009            let (_tx, rx) = oneshot::channel();
1010            Some(rx)
1011        });
1012        let mut mock_ledger = MockLedger::default();
1013        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
1014        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
1015        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
1016        mock_ledger.expect_check_transaction_basic().returning(|_, _| Ok(()));
1017        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
1018        // Initialize the storage.
1019        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
1020
1021        // Create the Worker.
1022        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
1023        let transaction = sample_execution_transaction_with_fee(false, rng, 0);
1024        let transaction_id = transaction.id();
1025        let transaction_data = Data::Object(transaction);
1026        let checksum = transaction_data.to_checksum::<CurrentNetwork>().unwrap();
1027        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
1028        let worker_ = worker.clone();
1029        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
1030        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
1031        assert!(worker.pending.contains(transmission_id));
1032        let result = worker.process_unconfirmed_transaction(transaction_id, transaction_data).await;
1033        assert!(result.is_ok());
1034        assert!(!worker.pending.contains(transmission_id));
1035        assert!(worker.ready.read().contains(transmission_id));
1036    }
1037
1038    #[tokio::test]
1039    async fn test_process_transaction_nok() {
1040        let mut rng = &mut TestRng::default();
1041        // Sample a committee.
1042        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
1043        let committee_clone = committee.clone();
1044        // Setup the mock gateway and ledger.
1045        let mut gateway = MockGateway::default();
1046        gateway.expect_send().returning(|_, _| {
1047            let (_tx, rx) = oneshot::channel();
1048            Some(rx)
1049        });
1050        let mut mock_ledger = MockLedger::default();
1051        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
1052        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
1053        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
1054        mock_ledger.expect_check_transaction_basic().returning(|_, _| Err(anyhow!("")));
1055        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
1056        // Initialize the storage.
1057        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
1058
1059        // Create the Worker.
1060        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
1061        let transaction_id: <CurrentNetwork as Network>::TransactionID = Field::<CurrentNetwork>::rand(&mut rng).into();
1062        let transaction = Data::Buffer(Bytes::from((0..512).map(|_| rng.random::<u8>()).collect::<Vec<_>>()));
1063        let checksum = transaction.to_checksum::<CurrentNetwork>().unwrap();
1064        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
1065        let worker_ = worker.clone();
1066        let peer_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
1067        let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
1068        assert!(worker.pending.contains(transmission_id));
1069        let result = worker.process_unconfirmed_transaction(transaction_id, transaction).await;
1070        assert!(result.is_err());
1071        assert!(!worker.pending.contains(transmission_id));
1072        assert!(!worker.ready.read().contains(transmission_id));
1073    }
1074
1075    #[tokio::test]
1076    async fn test_flood_transmission_requests() {
1077        let rng = &mut TestRng::default();
1078        // Sample a committee.
1079        let committee = snarkvm::ledger::committee::test_helpers::sample_committee(rng);
1080        let committee_clone = committee.clone();
1081        // Setup the mock gateway and ledger.
1082        let mut gateway = MockGateway::default();
1083        gateway.expect_send().returning(|_, _| {
1084            let (_tx, rx) = oneshot::channel();
1085            Some(rx)
1086        });
1087        let mut mock_ledger = MockLedger::default();
1088        mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
1089        mock_ledger.expect_get_committee_lookback_for_round().returning(move |_| Ok(committee_clone.clone()));
1090        mock_ledger.expect_contains_transmission().returning(|_| Ok(false));
1091        mock_ledger.expect_check_transaction_basic().returning(|_, _| Ok(()));
1092        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
1093        // Initialize the storage.
1094        let storage = Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), 1).unwrap();
1095
1096        // Create the Worker.
1097        let worker = Worker::new(0, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
1098        let transaction = sample_execution_transaction_with_fee(false, rng, 0);
1099        let transaction_id = transaction.id();
1100        let transaction_data = Data::Object(transaction);
1101        let checksum = transaction_data.to_checksum::<CurrentNetwork>().unwrap();
1102        let transmission_id = TransmissionID::Transaction(transaction_id, checksum);
1103
1104        // Determine the number of redundant requests are sent.
1105        let num_redundant_requests =
1106            max_redundant_requests(worker.ledger.clone(), worker.storage.current_round()).unwrap();
1107        let num_flood_requests = num_redundant_requests * 10;
1108        let mut peer_ips =
1109            (0..num_flood_requests).map(|i| SocketAddr::from(([127, 0, 0, 1], 1234 + i as u16))).collect_vec();
1110        let first_peer_ip = peer_ips[0];
1111
1112        // Flood the pending queue with transmission requests.
1113        for i in 1..=num_flood_requests {
1114            let worker_ = worker.clone();
1115            let peer_ip = peer_ips.pop().unwrap();
1116            tokio::spawn(async move {
1117                let _ = worker_.send_transmission_request(peer_ip, transmission_id).await;
1118            });
1119            tokio::time::sleep(Duration::from_millis(10)).await;
1120            // Check that the number of sent requests does not exceed the maximum number of redundant requests.
1121            assert!(worker.pending.num_sent_requests(transmission_id) <= num_redundant_requests);
1122            assert_eq!(worker.pending.num_callbacks(transmission_id), i);
1123        }
1124        // Check that the number of sent requests does not exceed the maximum number of redundant requests.
1125        assert_eq!(worker.pending.num_sent_requests(transmission_id), num_redundant_requests);
1126        assert_eq!(worker.pending.num_callbacks(transmission_id), num_flood_requests);
1127
1128        // Let all the requests expire.
1129        tokio::time::sleep(Duration::from_secs(CALLBACK_EXPIRATION_IN_SECS as u64 + 1)).await;
1130        assert_eq!(worker.pending.num_sent_requests(transmission_id), 0);
1131        assert_eq!(worker.pending.num_callbacks(transmission_id), 0);
1132
1133        // Flood the pending queue with transmission requests again, this time to a single peer
1134        for i in 1..=num_flood_requests {
1135            let worker_ = worker.clone();
1136            tokio::spawn(async move {
1137                let _ = worker_.send_transmission_request(first_peer_ip, transmission_id).await;
1138            });
1139            tokio::time::sleep(Duration::from_millis(10)).await;
1140            assert!(worker.pending.num_sent_requests(transmission_id) <= num_redundant_requests);
1141            assert_eq!(worker.pending.num_callbacks(transmission_id), i);
1142        }
1143        // Check that the number of sent requests does not exceed the maximum number of redundant requests.
1144        assert_eq!(worker.pending.num_sent_requests(transmission_id), 1);
1145        assert_eq!(worker.pending.num_callbacks(transmission_id), num_flood_requests);
1146
1147        // Check that fulfilling a transmission request clears the pending queue.
1148        let result = worker.process_unconfirmed_transaction(transaction_id, transaction_data).await;
1149        assert!(result.is_ok());
1150        assert_eq!(worker.pending.num_sent_requests(transmission_id), 0);
1151        assert_eq!(worker.pending.num_callbacks(transmission_id), 0);
1152        assert!(!worker.pending.contains(transmission_id));
1153        assert!(worker.ready.read().contains(transmission_id));
1154    }
1155
1156    #[tokio::test]
1157    async fn test_storage_gc_on_initialization() {
1158        let rng = &mut TestRng::default();
1159
1160        for _ in 0..ITERATIONS {
1161            // Mock the ledger round.
1162            let max_gc_rounds = rng.random_range(50..=100);
1163            let latest_ledger_round = rng.random_range((max_gc_rounds + 1)..1000);
1164            let expected_gc_round = latest_ledger_round - max_gc_rounds;
1165
1166            // Sample a committee.
1167            let committee =
1168                snarkvm::ledger::committee::test_helpers::sample_committee_for_round(latest_ledger_round, rng);
1169
1170            // Setup the mock gateway and ledger.
1171            let mut mock_ledger = MockLedger::default();
1172            mock_ledger.expect_current_committee().returning(move || Ok(committee.clone()));
1173
1174            let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(mock_ledger);
1175            // Initialize the storage.
1176            let storage =
1177                Storage::<CurrentNetwork>::new(ledger.clone(), Arc::new(BFTMemoryService::new()), max_gc_rounds)
1178                    .unwrap();
1179
1180            // Ensure that the storage GC round is correct.
1181            assert_eq!(storage.gc_round(), expected_gc_round);
1182        }
1183    }
1184}
1185
1186#[cfg(test)]
1187mod prop_tests {
1188    use super::*;
1189    use crate::Gateway;
1190    use snarkos_node_bft_ledger_service::MockLedgerService;
1191    use snarkvm::{
1192        console::account::Address,
1193        ledger::committee::{Committee, MIN_VALIDATOR_STAKE},
1194    };
1195
1196    use rand::RngExt;
1197    use test_strategy::proptest;
1198
1199    type CurrentNetwork = snarkvm::prelude::MainnetV0;
1200
1201    // Initializes a new test committee.
1202    fn new_test_committee(n: u16) -> Committee<CurrentNetwork> {
1203        let mut members = IndexMap::with_capacity(n as usize);
1204        for i in 0..n {
1205            // Sample the address.
1206            let rng = &mut TestRng::fixed(i as u64);
1207            let address = Address::new(rng.random());
1208            info!("Validator {i}: {address}");
1209            members.insert(address, (MIN_VALIDATOR_STAKE, false, rng.random_range(0..100)));
1210        }
1211        // Initialize the committee.
1212        Committee::<CurrentNetwork>::new(1u64, members).unwrap()
1213    }
1214
1215    #[proptest]
1216    fn worker_initialization(
1217        #[strategy(0..MAX_WORKERS)] id: u8,
1218        gateway: Gateway<CurrentNetwork>,
1219        storage: Storage<CurrentNetwork>,
1220    ) {
1221        let committee = new_test_committee(4);
1222        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(MockLedgerService::new(committee));
1223        let worker = Worker::new(id, Arc::new(gateway), storage, ledger, Default::default()).unwrap();
1224        assert_eq!(worker.id(), id);
1225    }
1226
1227    #[proptest]
1228    fn invalid_worker_id(
1229        #[strategy(MAX_WORKERS..)] id: u8,
1230        gateway: Gateway<CurrentNetwork>,
1231        storage: Storage<CurrentNetwork>,
1232    ) {
1233        let committee = new_test_committee(4);
1234        let ledger: Arc<dyn LedgerService<CurrentNetwork>> = Arc::new(MockLedgerService::new(committee));
1235        let worker = Worker::new(id, Arc::new(gateway), storage, ledger, Default::default());
1236        // TODO once Worker implements Debug, simplify this with `unwrap_err`
1237        if let Err(error) = worker {
1238            assert_eq!(error.to_string(), format!("Invalid worker ID '{id}'"));
1239        }
1240    }
1241}