solana_net_utils/
sockets.rs

1use {
2    crate::{bind_common_in_range_with_config, bind_common_with_config, PortRange, SocketConfig},
3    std::{
4        net::{IpAddr, SocketAddr, TcpListener, UdpSocket},
5        sync::atomic::{AtomicU16, Ordering},
6    },
7};
8// base port for deconflicted allocations
9const BASE_PORT: u16 = 5000;
10// how much to allocate per individual process.
11// we expect to have at most 64 concurrent tests in CI at any moment on a given host.
12const SLICE_PER_PROCESS: u16 = (u16::MAX - BASE_PORT) / 64;
13/// Retrieve a free 20-port slice for unit tests
14///
15/// When running under nextest, this will try to provide
16/// a unique slice of port numbers (assuming no other nextest processes
17/// are running on the same host) based on NEXTEST_TEST_GLOBAL_SLOT variable
18/// The port ranges will be reused following nextest logic.
19///
20/// When running without nextest, this will only bump an atomic and eventually
21/// panic when it runs out of port numbers to assign.
22#[allow(clippy::arithmetic_side_effects)]
23pub fn localhost_port_range_for_tests() -> (u16, u16) {
24    static SLICE: AtomicU16 = AtomicU16::new(0);
25    let offset = SLICE.fetch_add(20, Ordering::Relaxed);
26    let start = offset
27        + match std::env::var("NEXTEST_TEST_GLOBAL_SLOT") {
28            Ok(slot) => {
29                let slot: u16 = slot.parse().unwrap();
30                assert!(
31                    offset < SLICE_PER_PROCESS,
32                    "Overrunning into the port range of another test! Consider using fewer ports per test."
33                );
34                BASE_PORT + slot * SLICE_PER_PROCESS
35            }
36            Err(_) => BASE_PORT,
37        };
38    assert!(start < u16::MAX - 20, "ran out of port numbers!");
39    (start, start + 20)
40}
41
42pub fn bind_gossip_port_in_range(
43    gossip_addr: &SocketAddr,
44    port_range: PortRange,
45    bind_ip_addr: IpAddr,
46) -> (u16, (UdpSocket, TcpListener)) {
47    let config = SocketConfig::default();
48    if gossip_addr.port() != 0 {
49        (
50            gossip_addr.port(),
51            bind_common_with_config(bind_ip_addr, gossip_addr.port(), config).unwrap_or_else(|e| {
52                panic!("gossip_addr bind_to port {}: {}", gossip_addr.port(), e)
53            }),
54        )
55    } else {
56        bind_common_in_range_with_config(bind_ip_addr, port_range, config).expect("Failed to bind")
57    }
58}