Test: Bind to port 0 for automatic port selection and report that to the fixture

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,  # The default of 0 means "pick any free port".
 ) -> 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,82 @@ 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()
-
+    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(
+                    read_exact_fd(read_fd, 2),
+                    byteorder="big",
+                    signed=False,
+                )
+                os.close(read_fd)
+                read_fd = -1
+
+                yield OpsqueueProcess(port=actual_port, process=process)
+                assert process.poll() is None, "Opsqueue process failed during run"
+            finally:
+                # Give the process a chance to exit cleanly, but if it doesn't, kill it.
+                # `with subprocess.Popen(...) as process` will not terminate the process on its own.
+                process.terminate()
+                try:
+                    process.wait(timeout=1)
+                except subprocess.TimeoutExpired as exc:
+                    process.kill()
+                    raise AssertionError(
+                        "Opsqueue process locked up for more than 1 second on shutdown"
+                    ) from exc
 
-def random_free_port() -> int:
-    import random
+    finally:
+        if write_fd != -1:
+            os.close(write_fd)
+        if read_fd != -1:
+            os.close(read_fd)
 
-    while True:
-        port = random.randrange(10_000, 60_000)
-        if not is_port_in_use(port):
-            return port
 
+def read_exact_fd(fd: int, num_bytes: int) -> bytes:
+    """
+    Reads exactly `num_bytes` bytes from the given file descriptor `fd`.
 
-def is_port_in_use(port: int) -> bool:
-    import socket
+    `os.read` may return fewer bytes than requested, so this function will keep reading until the
+    requested number of bytes is obtained or EOF is reached.
 
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        return s.connect_ex(("localhost", port)) == 0
+    Raises EOFError if the end of the file is reached before reading the requested number of bytes.
+    """
+    data = bytearray()
+    while len(data) < num_bytes:
+        chunk = os.read(fd, num_bytes - len(data))
+        if not chunk:
+            raise EOFError(
+                f"Unexpected EOF: expected {num_bytes} bytes, got {len(data)}: {bytes(data)!r}"
+            )
+        data.extend(chunk)
+    return bytes(data)
 
 
 @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,16 @@ pub struct Config {
     #[arg(short, long, default_value_t = 3999)]
     pub port: u16,
 
+    /// 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.
+    ///
+    /// The port is written as a 16-bit (u16) big-endian integer to the given file descriptor which
+    /// is then closed. On error the file descriptor is closed without writing anything.
+    #[arg(long, value_parser = parse_report_bound_port_fd, default_value = "-1")]
+    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 +99,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 +111,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 +122,43 @@ 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 == -1 {
+        return Ok(ReportBoundPortPipe::default());
+    }
+
+    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}"
             ),
@@ -57,7 +68,11 @@ pub async fn serve_producer_and_consumer(
     })
     .retry(retry_policy())
     .notify(|e, d| tracing::error!("Error when binding server address: {e:?}, retrying in {d:?}"))
-    .await
+    .await.inspect_err(|_|{
+        // Drop the pipe after the server start retries have been exhausted. This ensures that the
+        // parent process can safely block on reading from the pipe.
+        config.report_bound_port_pipe.take();
+    })
 }
 
 #[cfg(feature = "server-logic")]