fix(node/bft): ensure gateways are shut down during prod tests

Author: kaimast

node/bft/src/bft.rs

@@ -899,7 +899,10 @@ mod tests {
     use aleo_std::StorageMode;
     use anyhow::Result;
     use indexmap::{IndexMap, IndexSet};
-    use std::sync::Arc;
+    use std::{
+        net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+        sync::Arc,
+    };
 
     type CurrentNetwork = snarkvm::console::network::MainnetV0;
 
@@ -934,13 +937,17 @@ mod tests {
     ) -> anyhow::Result<BFT<CurrentNetwork>> {
         // Create the block synchronization logic.
         let block_sync = Arc::new(BlockSync::new(ledger.clone()));
+
+        // Pick a random port so we can run tests concurrently.
+        let any_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0));
+
         // Initialize the BFT.
         BFT::new(
             account.clone(),
             storage.clone(),
             ledger.clone(),
             block_sync,
-            None,
+            Some(any_addr),
             &[],
             StorageMode::new_test(None),
             None,

node/bft/src/gateway.rs

@@ -20,29 +20,31 @@ use crate::{
     MAX_BATCH_DELAY_IN_MS,
     MEMORY_POOL_PORT,
     Worker,
-    events::{BatchPropose, BatchSignature, EventCodec, PrimaryPing},
+    events::{
+        BatchPropose,
+        BatchSignature,
+        BlockRequest,
+        BlockResponse,
+        CertificateRequest,
+        CertificateResponse,
+        ChallengeRequest,
+        ChallengeResponse,
+        DataBlocks,
+        DisconnectReason,
+        Event,
+        EventCodec,
+        EventTrait,
+        PrimaryPing,
+        TransmissionRequest,
+        TransmissionResponse,
+        ValidatorsRequest,
+        ValidatorsResponse,
+    },
     helpers::{Cache, CallbackHandle, Resolver, Storage, WorkerSender, assign_to_worker},
-    spawn_blocking,
+    ledger_service::LedgerService,
 };
-use aleo_std::StorageMode;
+
 use snarkos_account::Account;
-use snarkos_node_bft_events::{
-    BlockRequest,
-    BlockResponse,
-    CertificateRequest,
-    CertificateResponse,
-    ChallengeRequest,
-    ChallengeResponse,
-    DataBlocks,
-    DisconnectReason,
-    Event,
-    EventTrait,
-    TransmissionRequest,
-    TransmissionResponse,
-    ValidatorsRequest,
-    ValidatorsResponse,
-};
-use snarkos_node_bft_ledger_service::LedgerService;
 use snarkos_node_router::{NodeType, Peer, PeerPoolHandling};
 use snarkos_node_sync::{MAX_BLOCKS_BEHIND, communication_service::CommunicationService, locators::BlockLocators};
 use snarkos_node_tcp::{
@@ -53,6 +55,8 @@ use snarkos_node_tcp::{
     Tcp,
     protocols::{Disconnect, Handshake, OnConnect, Reading, Writing},
 };
+
+use aleo_std::StorageMode;
 use snarkvm::{
     console::prelude::*,
     ledger::{
@@ -61,6 +65,7 @@ use snarkvm::{
         narwhal::{BatchCertificate, BatchHeader, Data},
     },
     prelude::{Address, Field},
+    utilities::task::{self, JoinHandle},
 };
 
 use colored::Colorize;
@@ -82,7 +87,6 @@ use std::{
 use tokio::{
     net::TcpStream,
     sync::{OnceCell, oneshot},
-    task::{self, JoinHandle},
 };
 use tokio_stream::StreamExt;
 use tokio_util::codec::Framed;
@@ -215,8 +219,14 @@ impl<N: Network> Gateway<N> {
             (None, None) => SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, MEMORY_POOL_PORT)),
             (Some(ip), _) => ip,
         };
+
+        // Allow at most as many connections as the maximum committe size.
+        // and fail if the chosen port is not available.
+        let mut tcp_config = Config::new(ip, Committee::<N>::max_committee_size()?);
+        tcp_config.allow_random_port = false;
+
         // Initialize the TCP stack.
-        let tcp = Tcp::new(Config::new(ip, Committee::<N>::max_committee_size()?));
+        let tcp = Tcp::new(tcp_config);
 
         // Prepare the collection of the initial peers.
         let mut initial_peers = HashMap::new();
@@ -609,19 +619,14 @@ impl<N: Network> Gateway<N> {
                 }
 
                 let self_ = self.clone();
-                let blocks = match task::spawn_blocking(move || {
+                let blocks = task::spawn_blocking(move || {
                     // Retrieve the blocks within the requested range.
                     match self_.ledger.get_blocks(start_height..end_height) {
                         Ok(blocks) => Ok(Data::Object(DataBlocks(blocks))),
                         Err(error) => bail!("Missing blocks {start_height} to {end_height} from ledger - {error}"),
                     }
                 })
-                .await
-                {
-                    Ok(Ok(blocks)) => blocks,
-                    Ok(Err(error)) => return Err(error),
-                    Err(error) => return Err(anyhow!("[BlockRequest] {error}")),
-                };
+                .await?;
 
                 let self_ = self.clone();
                 tokio::spawn(async move {
@@ -870,7 +875,7 @@ impl<N: Network> Gateway<N> {
     /// Spawns a task with the given future; it should only be used for long-running tasks.
     #[allow(dead_code)]
     fn spawn<T: Future<Output = ()> + Send + 'static>(&self, future: T) {
-        self.handles.lock().push(tokio::spawn(future));
+        self.handles.lock().push(task::spawn(future));
     }
 
     /// Shuts down the gateway.
@@ -1510,7 +1515,7 @@ impl<N: Network> Gateway<N> {
             return Some(DisconnectReason::InvalidChallengeResponse);
         }
         // Perform the deferred non-blocking deserialization of the signature.
-        let Ok(signature) = spawn_blocking!(signature.deserialize_blocking()) else {
+        let Ok(signature) = task::spawn_blocking(|| signature.deserialize_blocking()).await else {
             warn!("{CONTEXT} Gateway handshake with '{peer_addr}' failed (cannot deserialize the signature)");
             return Some(DisconnectReason::InvalidChallengeResponse);
         };
@@ -1621,17 +1626,11 @@ mod prop_tests {
 
     impl GatewayAddress {
         fn ip(&self) -> Option<SocketAddr> {
-            if let GatewayAddress::Prod(ip) = self {
-                return *ip;
-            }
-            None
+            if let GatewayAddress::Prod(ip) = self { *ip } else { None }
         }
 
         fn port(&self) -> Option<u16> {
-            if let GatewayAddress::Dev(port) = self {
-                return Some(*port as u16);
-            }
-            None
+            if let GatewayAddress::Dev(port) = self { Some(*port as u16) } else { None }
         }
     }
 
@@ -1689,8 +1688,8 @@ mod prop_tests {
             .boxed()
     }
 
-    #[proptest]
-    fn gateway_dev_initialization(#[strategy(any_valid_dev_gateway())] input: GatewayInput) {
+    #[proptest(async = "tokio")]
+    async fn gateway_dev_initialization(#[strategy(any_valid_dev_gateway())] input: GatewayInput) {
         let (storage, _, private_key, dev) = input;
         let account = Account::try_from(private_key).unwrap();
 
@@ -1711,10 +1710,13 @@ mod prop_tests {
         let tcp_config = gateway.tcp().config();
         assert_eq!(tcp_config.max_connections, Committee::<CurrentNetwork>::max_committee_size().unwrap());
         assert_eq!(gateway.account().address(), account.address());
+
+        // Ensure the gateway shuts down and unbinds the TCP port.
+        gateway.shut_down().await;
     }
 
-    #[proptest]
-    fn gateway_prod_initialization(#[strategy(any_valid_prod_gateway())] input: GatewayInput) {
+    #[proptest(async = "tokio")]
+    async fn gateway_prod_initialization(#[strategy(any_valid_prod_gateway())] input: GatewayInput) {
         let (storage, _, private_key, dev) = input;
         let account = Account::try_from(private_key).unwrap();
 
@@ -1740,6 +1742,9 @@ mod prop_tests {
         let tcp_config = gateway.tcp().config();
         assert_eq!(tcp_config.max_connections, Committee::<CurrentNetwork>::max_committee_size().unwrap());
         assert_eq!(gateway.account().address(), account.address());
+
+        // Ensure the gateway shuts down and unbinds the TCP port.
+        gateway.shut_down().await;
     }
 
     #[proptest(async = "tokio")]
@@ -1793,6 +1798,9 @@ mod prop_tests {
             SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), MEMORY_POOL_PORT + dev.port().unwrap())
         );
         assert_eq!(gateway.num_workers(), workers.len() as u8);
+
+        // Ensure the gateway shuts down and unbinds the TCP port.
+        gateway.shut_down().await;
     }
 
     #[proptest]

node/bft/src/lib.rs

@@ -69,14 +69,3 @@ pub const MAX_WORKERS: u8 = 1; // worker(s)
 pub const PRIMARY_PING_IN_MS: u64 = 2 * MAX_BATCH_DELAY_IN_MS; // ms
 /// The interval at which each worker broadcasts a ping to every other node.
 pub const WORKER_PING_IN_MS: u64 = 4 * MAX_BATCH_DELAY_IN_MS; // ms
-
-/// A helper macro to spawn a blocking task.
-#[macro_export]
-macro_rules! spawn_blocking {
-    ($expr:expr) => {
-        match tokio::task::spawn_blocking(move || $expr).await {
-            Ok(value) => value,
-            Err(error) => Err(anyhow::anyhow!("[tokio::spawn_blocking] {error}")),
-        }
-    };
-}