Skip to main content

diode_base/testing/
free_port.rs

1//! Free Port Utility
2//!
3//! This module provides a utility for managing free ports in tests to avoid port conflicts.
4//! The `FreePort` struct automatically reserves and releases ports using a global registry.
5
6use rand::Rng;
7use std::collections::HashSet;
8use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener};
9use std::sync::{LazyLock, Mutex};
10
11/// Global registry of allocated ports to prevent conflicts between tests
12static ALLOCATED_PORTS: LazyLock<Mutex<HashSet<u16>>> =
13    LazyLock::new(|| Mutex::new(HashSet::new()));
14
15/// A wrapper around a port number that guarantees the port is free and manages its lifecycle
16#[derive(Debug)]
17pub struct FreePort(u16);
18
19impl FreePort {
20    /// Creates a new FreePort by finding an available port
21    ///
22    /// This method will attempt up to 16 times to find a free port by:
23    /// 1. Generating a random port number in the range 8000-65000
24    /// 2. Checking that it's not already allocated in the global registry
25    /// 3. Attempting to bind to the port to verify it's actually free
26    /// 4. Adding it to the global registry to prevent other threads from using it
27    ///
28    /// # Panics
29    ///
30    /// Panics if unable to find a free port after 16 attempts
31    pub fn new() -> Self {
32        let mut rng = rand::thread_rng();
33
34        for _ in 0..16 {
35            // Generate random port in range 8000-65000
36            let port = rng.gen_range(8000..=65000);
37
38            // Check if port is already allocated
39            {
40                let allocated = ALLOCATED_PORTS.lock().unwrap();
41                if allocated.contains(&port) {
42                    continue;
43                }
44            }
45
46            // Try to bind to the port to verify it's actually free
47            if let Ok(listener) = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)) {
48                // Close the listener immediately - we just wanted to check availability
49                drop(listener);
50
51                // Add to allocated ports registry
52                {
53                    let mut allocated = ALLOCATED_PORTS.lock().unwrap();
54                    if allocated.insert(port) {
55                        // Successfully inserted (wasn't already there)
56                        return FreePort(port);
57                    }
58                    // If insert returned false, another thread beat us to it, try again
59                }
60            }
61        }
62
63        panic!("Unable to find a free port after 16 attempts");
64    }
65
66    /// Returns the port number
67    pub fn port(&self) -> u16 {
68        self.0
69    }
70
71    /// Returns the port as a formatted string for binding addresses
72    pub fn as_addr(&self) -> SocketAddr {
73        SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, self.0))
74    }
75}
76
77impl Drop for FreePort {
78    /// Automatically removes the port from the global registry when dropped
79    fn drop(&mut self) {
80        let mut allocated = ALLOCATED_PORTS.lock().unwrap();
81        allocated.remove(&self.0);
82    }
83}
84
85impl Default for FreePort {
86    fn default() -> Self {
87        FreePort::new()
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use std::thread;
95
96    #[test]
97    fn test_free_port_allocation() {
98        let port = FreePort::new();
99        assert!(port.port() >= 8000 && port.port() <= 65000);
100
101        // Verify port is in allocated registry
102        {
103            let allocated = ALLOCATED_PORTS.lock().unwrap();
104            assert!(allocated.contains(&port.port()));
105        }
106    }
107
108    #[test]
109    fn test_free_port_release() {
110        let port_num = {
111            let port = FreePort::new();
112            let port_num = port.port();
113
114            // Verify port is allocated
115            {
116                let allocated = ALLOCATED_PORTS.lock().unwrap();
117                assert!(allocated.contains(&port_num));
118            }
119
120            port_num
121        }; // port is dropped here
122
123        // Verify port is released
124        {
125            let allocated = ALLOCATED_PORTS.lock().unwrap();
126            assert!(!allocated.contains(&port_num));
127        }
128    }
129
130    #[test]
131    fn test_multiple_ports_no_conflict() {
132        let port1 = FreePort::new();
133        let port2 = FreePort::new();
134        let port3 = FreePort::new();
135
136        // All ports should be different
137        assert_ne!(port1.port(), port2.port());
138        assert_ne!(port1.port(), port3.port());
139        assert_ne!(port2.port(), port3.port());
140
141        // All should be in allocated registry
142        {
143            let allocated = ALLOCATED_PORTS.lock().unwrap();
144            assert!(allocated.contains(&port1.port()));
145            assert!(allocated.contains(&port2.port()));
146            assert!(allocated.contains(&port3.port()));
147        }
148    }
149
150    #[test]
151    fn test_concurrent_allocation() {
152        let handles: Vec<_> = (0..10)
153            .map(|_| {
154                thread::spawn(|| {
155                    let port = FreePort::new();
156                    thread::sleep(std::time::Duration::from_millis(10));
157                    port.port()
158                })
159            })
160            .collect();
161
162        let ports: Vec<u16> = handles.into_iter().map(|h| h.join().unwrap()).collect();
163
164        // All ports should be unique
165        let mut unique_ports = HashSet::new();
166        for port in &ports {
167            assert!(
168                unique_ports.insert(*port),
169                "Port {} was allocated twice",
170                port
171            );
172        }
173
174        assert_eq!(ports.len(), 10);
175    }
176
177    #[test]
178    fn test_as_addr_format() {
179        let port = FreePort::new();
180        let expected = format!("127.0.0.1:{}", port.port());
181        assert_eq!(port.as_addr().to_string(), expected);
182    }
183}