diode_base/testing/
free_port.rs1use rand::Rng;
7use std::collections::HashSet;
8use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener};
9use std::sync::{LazyLock, Mutex};
10
11static ALLOCATED_PORTS: LazyLock<Mutex<HashSet<u16>>> =
13 LazyLock::new(|| Mutex::new(HashSet::new()));
14
15#[derive(Debug)]
17pub struct FreePort(u16);
18
19impl FreePort {
20 pub fn new() -> Self {
32 let mut rng = rand::thread_rng();
33
34 for _ in 0..16 {
35 let port = rng.gen_range(8000..=65000);
37
38 {
40 let allocated = ALLOCATED_PORTS.lock().unwrap();
41 if allocated.contains(&port) {
42 continue;
43 }
44 }
45
46 if let Ok(listener) = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)) {
48 drop(listener);
50
51 {
53 let mut allocated = ALLOCATED_PORTS.lock().unwrap();
54 if allocated.insert(port) {
55 return FreePort(port);
57 }
58 }
60 }
61 }
62
63 panic!("Unable to find a free port after 16 attempts");
64 }
65
66 pub fn port(&self) -> u16 {
68 self.0
69 }
70
71 pub fn as_addr(&self) -> SocketAddr {
73 SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, self.0))
74 }
75}
76
77impl Drop for FreePort {
78 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 {
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 {
116 let allocated = ALLOCATED_PORTS.lock().unwrap();
117 assert!(allocated.contains(&port_num));
118 }
119
120 port_num
121 }; {
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 assert_ne!(port1.port(), port2.port());
138 assert_ne!(port1.port(), port3.port());
139 assert_ne!(port2.port(), port3.port());
140
141 {
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 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}