vector/test_util/
addr.rs

1//! Test utilities for allocating unique network addresses.
2//!
3//! This module provides thread-safe port allocation for tests using a guard pattern
4//! to prevent port reuse race conditions. The design eliminates intra-process races by:
5//! 1. Binding to get a port while holding a TCP listener
6//! 2. Registering the port atomically (while still holding the listener and registry lock)
7//! 3. Only then releasing the listener (port now protected by registry entry)
8//!
9//! This ensures no race window between port allocation and registration.
10
11#[cfg(windows)]
12use std::net::UdpSocket;
13use std::{
14    collections::HashSet,
15    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener},
16    sync::{LazyLock, Mutex},
17};
18
19/// Maximum number of attempts to allocate a unique port before panicking.
20/// This should be far more than needed since port collisions are rare,
21/// but provides a safety net against infinite loops.
22const MAX_PORT_ALLOCATION_ATTEMPTS: usize = 100;
23
24/// A guard that reserves a port in the registry, preventing port reuse until dropped.
25/// The guard does NOT hold the actual listener - it just marks the port as reserved
26/// so that concurrent calls to next_addr() won't return the same port.
27pub struct PortGuard {
28    addr: SocketAddr,
29}
30
31impl PortGuard {
32    /// Get the socket address that this guard is holding.
33    pub const fn addr(&self) -> SocketAddr {
34        self.addr
35    }
36}
37
38impl Drop for PortGuard {
39    fn drop(&mut self) {
40        // Remove from the reserved ports set when dropped
41        RESERVED_PORTS
42            .lock()
43            .expect("poisoned lock potentially due to test panicking")
44            .remove(&self.addr.port());
45    }
46}
47
48/// Global set of reserved ports for collision detection. When a test allocates a port, we check this set to ensure the
49/// OS didn't recycle a port that's still in use by another test.
50/// Ports are tracked by number only (u16). This means IPv4 and IPv6 may block each other from using the same port.
51/// This simplification is acceptable for our tests.
52static RESERVED_PORTS: LazyLock<Mutex<HashSet<u16>>> = LazyLock::new(|| Mutex::new(HashSet::new()));
53
54/// Allocates a unique port and returns a guard that keeps it reserved.
55///
56/// The returned `PortGuard` must be kept alive for as long as you need the port reserved.
57/// When the guard is dropped, the port is automatically released.
58///
59/// If the OS assigns a port that's already reserved by another test, this function will
60/// automatically retry with a new port, ensuring each test gets a unique port.
61///
62/// # Example
63/// ```ignore
64/// let (_guard, addr) = next_addr_for_ip(IpAddr::V4(Ipv4Addr::LOCALHOST));
65/// // Use addr for your test
66/// // Port is automatically released when _guard goes out of scope
67/// ```
68pub fn next_addr_for_ip(ip: IpAddr) -> (PortGuard, SocketAddr) {
69    for _ in 0..MAX_PORT_ALLOCATION_ATTEMPTS {
70        let listener = StdTcpListener::bind((ip, 0)).expect("Failed to bind to OS-assigned port");
71        let addr = listener.local_addr().expect("Failed to get local address");
72        let port = addr.port();
73
74        // Check if this port is already reserved by another test WHILE still holding the listener
75        let mut reserved = RESERVED_PORTS
76            .lock()
77            .expect("poisoned lock potentially due to test panicking");
78        if reserved.contains(&port) {
79            // OS recycled a port that's still reserved by another test.
80            // Lock and listener will be dropped implicitly after continuing
81            continue;
82        }
83
84        // On Windows, certain ports are in OS-excluded ranges (e.g. set by Hyper-V/WSL).
85        // TCP bind(0) may return such a port, but UDP bind to the same port will fail with
86        // WSAEACCES (10013). Probe with a UDP socket and retry if it is excluded.
87        #[cfg(windows)]
88        if UdpSocket::bind(addr).is_err() {
89            continue;
90        }
91
92        // Port is unique, mark it as reserved BEFORE dropping the listener
93        // This ensures no race window between dropping listener and registering the port
94        reserved.insert(port);
95        drop(reserved);
96
97        // Now it's safe to drop the listener - the registry protects the port
98        drop(listener);
99
100        let guard = PortGuard { addr };
101        return (guard, addr);
102    }
103
104    panic!("Failed to allocate a unique port after {MAX_PORT_ALLOCATION_ATTEMPTS} attempts");
105}
106
107pub fn next_addr() -> (PortGuard, SocketAddr) {
108    next_addr_for_ip(IpAddr::V4(Ipv4Addr::LOCALHOST))
109}
110
111pub fn next_addr_any() -> (PortGuard, SocketAddr) {
112    next_addr_for_ip(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
113}
114
115pub fn next_addr_v6() -> (PortGuard, SocketAddr) {
116    next_addr_for_ip(IpAddr::V6(Ipv6Addr::LOCALHOST))
117}
vector/test_util/addr.rs

vector/test_util/
addr.rs
