Test: Bind to port 0

Author: ReinierMaas

libs/opsqueue_python/tests/conftest.py

@@ -50,13 +50,9 @@ def opsqueue() -> Generator[OpsqueueProcess, None, None]:
 
 @contextmanager
 def opsqueue_service(
-    *, port: int | None = None
+    *,
+    port: int = 0,
 ) -> Generator[OpsqueueProcess, None, None]:
-    global test_opsqueue_port_offset
-
-    if port is None:
-        port = random_free_port()
-
     # This will create a SQLite database in memory.
     # We need the `cache=shared` to allow sharing this DB between all threads within the same OS process.
     temp_dbname = "file::memory:?cache=shared"
@@ -66,43 +62,47 @@ def opsqueue_service(
     # will from time to time hang for **many minutes** on initializing SQLite for some reason.
     # temp_dbname = f"/tmp/opsqueue_tests-{uuid.uuid4()}.db"
 
+    read_fd, write_fd = os.pipe()
+
     command = [
         "setpriv",
         "--pdeathsig=SIGKILL",
         str(opsqueue_bin_location()),
         "--port",
         str(port),
+        "--report-bound-port-pipe",
+        str(write_fd),
         "--database-filename",
         temp_dbname,
     ]
     env = os.environ.copy()  # We copy the env so e.g. RUST_LOG and other env vars are propagated from outside of the invocation of pytest
     if env.get("RUST_LOG") is None:
         env["RUST_LOG"] = "off"
 
-    with subprocess.Popen(command, cwd=PROJECT_ROOT, env=env) as process:
-        assert process.poll() is None, "Opsqueue process failed to start"
-        try:
-            wrapper = OpsqueueProcess(port=port, process=process)
-            yield wrapper
-            assert process.poll() is None, "Opsqueue process failed during run"
-        finally:
-            process.terminate()
-
-
-def random_free_port() -> int:
-    import random
-
-    while True:
-        port = random.randrange(10_000, 60_000)
-        if not is_port_in_use(port):
-            return port
-
-
-def is_port_in_use(port: int) -> bool:
-    import socket
-
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        return s.connect_ex(("localhost", port)) == 0
+    try:
+        with subprocess.Popen(
+            command,
+            cwd=PROJECT_ROOT,
+            env=env,
+            pass_fds=(write_fd,),
+        ) as process:
+            os.close(write_fd)
+            write_fd = -1
+
+            assert process.poll() is None, "Opsqueue process failed to start"
+            try:
+                actual_port = int.from_bytes(
+                    os.read(read_fd, 2), byteorder="big", signed=False
+                )
+                wrapper = OpsqueueProcess(port=actual_port, process=process)
+                yield wrapper
+                assert process.poll() is None, "Opsqueue process failed during run"
+            finally:
+                process.terminate()
+    finally:
+        if write_fd != -1:
+            os.close(write_fd)
+        os.close(read_fd)
 
 
 @contextmanager

opsqueue/src/config.rs

@@ -2,7 +2,13 @@
 //!
 //! We make use of the excellent `clap` crate to make customizing the configuration
 //! with command-line args easier.
-use std::num::NonZero;
+use std::{
+    fs::File,
+    io::{self, Write},
+    num::NonZero,
+    os::fd::FromRawFd,
+    sync::{Arc, Mutex},
+};
 
 use clap::Parser;
 
@@ -17,6 +23,13 @@ pub struct Config {
     #[arg(short, long, default_value_t = 3999)]
     pub port: u16,
 
+    /// Optional file descriptor where the final bound TCP port is written.
+    ///
+    /// This is useful when `--port 0` is used and a parent process wants
+    /// to receive the assigned port without filesystem IO.
+    #[arg(long, value_parser = parse_report_bound_port_fd)]
+    pub report_bound_port_pipe: ReportBoundPortPipe,
+
     /// Name of the SQLite database file used by this opsqueue.
     ///
     /// Configure this to different values when you run multiple opsqueues
@@ -83,6 +96,7 @@ impl Default for Config {
     fn default() -> Self {
         use std::str::FromStr;
         let port = 3999;
+        let report_bound_port_pipe = ReportBoundPortPipe::default();
         let database_filename = "opsqueue.db".to_string();
         let reservation_expiration =
             humantime::Duration::from_str("10 minutes").expect("valid humantime");
@@ -94,6 +108,7 @@ impl Default for Config {
         let max_submission_age = humantime::Duration::from_str("1 hour").expect("valid humantime");
         Config {
             port,
+            report_bound_port_pipe,
             database_filename,
             reservation_expiration,
             max_read_pool_size,
@@ -104,3 +119,39 @@ impl Default for Config {
         }
     }
 }
+
+#[derive(Debug, Clone, Default)]
+pub struct ReportBoundPortPipe(Arc<Mutex<Option<BoundPortPipe>>>);
+
+impl ReportBoundPortPipe {
+    pub fn take(&self) -> Option<BoundPortPipe> {
+        self.0.lock().expect("No poison").take()
+    }
+}
+
+#[derive(Debug)]
+pub struct BoundPortPipe(File);
+
+impl BoundPortPipe {
+    pub fn write_port(mut self, port: u16) -> io::Result<()> {
+        self.0.write_all(&u16::to_be_bytes(port))?;
+        self.0.flush()
+    }
+}
+
+fn parse_report_bound_port_fd(value: &str) -> Result<ReportBoundPortPipe, String> {
+    let fd = value
+        .parse::<i32>()
+        .map_err(|err| format!("invalid file descriptor {value:?}: {err}"))?;
+
+    if fd < 0 {
+        return Err(format!(
+            "invalid file descriptor {value:?}: must be non-negative"
+        ));
+    }
+
+    // SAFETY: the parent process passes ownership of this FD to us through `pass_fds`.
+    Ok(ReportBoundPortPipe(Arc::new(Mutex::new(Some(
+        BoundPortPipe(unsafe { File::from_raw_fd(fd) }),
+    )))))
+}

opsqueue/src/server.rs

@@ -45,7 +45,18 @@ pub async fn serve_producer_and_consumer(
         );
         let listener = tokio::net::TcpListener::bind(server_addr).await?;
         match listener.local_addr() {
-            Ok(addr) => tracing::info!("Server listening on {addr}"),
+            Ok(addr) => {
+                tracing::info!("Server listening on {addr}");
+                if let Some(pipe) = config.report_bound_port_pipe.take() {
+                    if let Err(err) = pipe.write_port(addr.port()) {
+                        tracing::warn!(
+                            "Failed to write bound port {} to pipe: {}",
+                            addr.port(),
+                            err
+                        );
+                    }
+                }
+            }
             Err(err) => tracing::warn!(
                 "Could not get locally bound address of the server, tried binding on {server_addr}: {err}"
             ),