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

vector/test_util/
addr.rs
