apollo_infra_utils: lease test port ranges below the ephemeral range

crates/apollo_infra_utils/src/test_utils.rs

@@ -1,4 +1,9 @@
+use std::env;
+use std::fs::{create_dir_all, File, OpenOptions, TryLockError};
+use std::io::ErrorKind;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::path::Path;
+use std::sync::Mutex;
 
 use assert_json_diff::{assert_json_matches_no_panic, CompareMode, Config};
 use num_enum::IntoPrimitive;
@@ -12,32 +17,97 @@ use tracing::instrument;
 mod test_utils_test;
 
 const PORTS_PER_INSTANCE: u16 = 80;
-pub const MAX_NUMBER_OF_INSTANCES_PER_TEST: u16 = 26;
-#[allow(clippy::as_conversions)]
-const MAX_NUMBER_OF_TESTS: u16 = TestIdentifier::COUNT as u16;
 const BASE_PORT: u16 = 11000;
 
-// A port between `is_port_in_use` returning false and the child process binding it can be handed
-// to an outbound connection by the kernel, surfacing as `Os { code: 98, kind: AddrInUse }`. The
-// allocated range (11000..63000) overlaps the ephemeral range that governs this
-// (`/proc/sys/net/ipv4/ip_local_port_range`, 32768..60999 by default), and one integration test
-// run allocated 40 of its 203 ports inside it.
-//
-// Moving the range below 32768 leaves 21768 usable ports, and the budget needs
-// MAX_NUMBER_OF_TESTS * MAX_NUMBER_OF_INSTANCES_PER_TEST * PORTS_PER_INSTANCE = 52000. Shrinking
-// PORTS_PER_INSTANCE is not the way out either: `integration_test_manager` draws
-// 5 * <services in the node> ports from a single instance, which is 60 for a distributed node and
-// 35 for a hybrid one. Fixing this needs a different allocation scheme, not a smaller constant.
+/// Lowest port the kernel hands out for outbound connections, per the default
+/// `/proc/sys/net/ipv4/ip_local_port_range` of 32768..60999.
+///
+/// Every port handed out stays below this. A port inside the ephemeral range can be claimed by an
+/// outbound connection between `is_port_in_use` returning false and the child process binding it,
+/// which surfaces as `Os { code: 98, kind: AddrInUse }` and takes the whole run down.
+const LOWEST_EPHEMERAL_PORT: u16 = 32768;
+
+/// Number of port ranges that fit below the ephemeral range.
+const NUM_PORT_SLOTS: u16 = (LOWEST_EPHEMERAL_PORT - BASE_PORT) / PORTS_PER_INSTANCE;
+
+const PORT_SLOT_LEASE_DIR_NAME: &str = "apollo_test_port_slots";
 
-// Ensure available ports don't exceed u16::MAX.
 const _: () = {
+    assert!(NUM_PORT_SLOTS > 0, "No port slots fit below the ephemeral port range.");
     assert!(
-        BASE_PORT + MAX_NUMBER_OF_TESTS * MAX_NUMBER_OF_INSTANCES_PER_TEST * PORTS_PER_INSTANCE
-            < u16::MAX,
-        "Port numbers potentially exceeding u16::MAX"
+        BASE_PORT + NUM_PORT_SLOTS * PORTS_PER_INSTANCE <= LOWEST_EPHEMERAL_PORT,
+        "Port slots reach into the ephemeral port range."
     );
 };
 
+/// Slots leased by this process, held until it exits.
+///
+/// Leases are deliberately not released when an `AvailablePorts` is dropped. Call sites take the
+/// ports they need and let the `AvailablePorts` go while those ports stay in use for the rest of
+/// the test, as `create_hybrid_component_configs` does, so releasing on drop would hand a live
+/// range to another test. The OS releases these locks when the process exits, so a crashed test
+/// cannot leak a slot either.
+static LEASED_PORT_SLOTS: Mutex<Vec<File>> = Mutex::new(Vec::new());
+
+/// Leases a port range that no other process holds, and returns its first port.
+///
+/// The lease is an exclusive lock on a file per slot, taken in a directory shared machine-wide, so
+/// concurrent test processes, including ones run from different checkouts, cannot be handed the
+/// same range. `label` identifies the holder in the log line.
+fn lease_port_slot(label: &str) -> u16 {
+    let lease_dir = env::temp_dir().join(PORT_SLOT_LEASE_DIR_NAME);
+    create_dir_all(&lease_dir)
+        .unwrap_or_else(|error| panic!("Failed to create {lease_dir:?}: {error}"));
+
+    for slot_index in 0..NUM_PORT_SLOTS {
+        let lease_path = lease_dir.join(format!("slot_{slot_index}"));
+        let lease_file = open_lease_file(&lease_path);
+
+        match lease_file.try_lock() {
+            Ok(()) => {
+                let start_port = BASE_PORT + slot_index * PORTS_PER_INSTANCE;
+                println!(
+                    "Leased port slot {slot_index} [{start_port},{}) for {label}",
+                    start_port + PORTS_PER_INSTANCE
+                );
+                LEASED_PORT_SLOTS
+                    .lock()
+                    .expect("Port slot lease registry was poisoned.")
+                    .push(lease_file);
+                return start_port;
+            }
+            Err(TryLockError::WouldBlock) => continue,
+            Err(TryLockError::Error(error)) => {
+                panic!("Failed to lock {lease_path:?}: {error}")
+            }
+        }
+    }
+
+    panic!(
+        "All {NUM_PORT_SLOTS} port slots below the ephemeral range are leased. Either too many \
+         test processes are running at once, or a slot is held by a process that outlived its \
+         test."
+    );
+}
+
+/// Opens a slot's lease file for locking, preferring a read-only handle.
+///
+/// `flock` does not care whether the handle is writable, and the lease directory is shared by every
+/// user on the machine: a file another user created is typically mode 0644, so asking for write
+/// access would fail with a permission error on a slot that is perfectly lockable.
+fn open_lease_file(lease_path: &Path) -> File {
+    match File::open(lease_path) {
+        Ok(lease_file) => lease_file,
+        Err(error) if error.kind() == ErrorKind::NotFound => OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(false)
+            .open(lease_path)
+            .unwrap_or_else(|error| panic!("Failed to create {lease_path:?}: {error}")),
+        Err(error) => panic!("Failed to open {lease_path:?}: {error}"),
+    }
+}
+
 #[repr(u16)]
 #[derive(Debug, Copy, Clone, IntoPrimitive, EnumCount)]
 // TODO(Nadin): Come up with a better name for this enum.
@@ -77,23 +147,18 @@ pub struct AvailablePorts {
 }
 
 impl AvailablePorts {
+    /// Leases a port range for this instance. `test_unique_index` and `instance_index` identify the
+    /// holder in the lease log line; the range itself comes from whichever slot is free, so tests
+    /// no longer need a statically partitioned budget large enough for every test at once.
     pub fn new(test_unique_index: u16, instance_index: u16) -> Self {
-        assert!(
-            test_unique_index < MAX_NUMBER_OF_TESTS,
-            "Test unique index {test_unique_index:?} exceeded bound {MAX_NUMBER_OF_TESTS:?}"
-        );
-        assert!(
-            instance_index < MAX_NUMBER_OF_INSTANCES_PER_TEST,
-            "Instance index {instance_index:?} exceeded bound {MAX_NUMBER_OF_INSTANCES_PER_TEST:?}",
-        );
-
-        let test_offset: u16 =
-            test_unique_index * MAX_NUMBER_OF_INSTANCES_PER_TEST * PORTS_PER_INSTANCE;
-        let instance_in_test_offset: u16 = instance_index * PORTS_PER_INSTANCE;
-        let current_port = BASE_PORT + test_offset + instance_in_test_offset;
-        let max_port: u16 = current_port + PORTS_PER_INSTANCE;
-
-        AvailablePorts { start_port: current_port, current_port, max_port }
+        let start_port =
+            lease_port_slot(&format!("test {test_unique_index} instance {instance_index}"));
+
+        AvailablePorts {
+            start_port,
+            current_port: start_port,
+            max_port: start_port + PORTS_PER_INSTANCE,
+        }
     }
 
     #[instrument]

crates/apollo_infra_utils/src/test_utils_test.rs

@@ -1,7 +1,7 @@
 use std::io::ErrorKind;
 use std::net::{Ipv4Addr, SocketAddr, TcpListener};
 
-use super::is_port_in_use;
+use super::{is_port_in_use, AvailablePorts, LOWEST_EPHEMERAL_PORT, PORTS_PER_INSTANCE};
 
 /// A free port, taken from the ephemeral range so it cannot collide with the ranges
 /// `AvailablePorts` hands out.
@@ -49,3 +49,50 @@ fn port_held_on_another_interface_is_reported_in_use() {
     assert!(is_port_in_use(port));
     drop(holder);
 }
+
+/// Ranges leased at the same time must not overlap, which is the property that lets concurrent test
+/// processes bind their own ports without coordinating.
+#[test]
+fn leases_held_at_once_do_not_overlap() {
+    let leases: Vec<AvailablePorts> =
+        (0..8).map(|instance_index| AvailablePorts::new(0, instance_index)).collect();
+
+    let mut ranges: Vec<(u16, u16)> =
+        leases.iter().map(|lease| (lease.start_port, lease.max_port)).collect();
+    ranges.sort_unstable();
+
+    for adjacent_ranges in ranges.windows(2) {
+        let (_, earlier_end) = adjacent_ranges[0];
+        let (later_start, _) = adjacent_ranges[1];
+        assert!(
+            earlier_end <= later_start,
+            "Leased ranges overlap: {:?} and {:?}",
+            adjacent_ranges[0],
+            adjacent_ranges[1]
+        );
+    }
+}
+
+/// The kernel can hand a port at or above `LOWEST_EPHEMERAL_PORT` to an outbound connection while a
+/// test is still about to bind it, so no allocated port may fall there.
+#[test]
+fn allocated_ports_are_below_the_ephemeral_range() {
+    let mut available_ports = AvailablePorts::new(0, 0);
+
+    for _ in 0..8 {
+        let port = available_ports.get_next_port();
+        assert!(
+            port < LOWEST_EPHEMERAL_PORT,
+            "Allocated port {port} is inside the ephemeral range starting at \
+             {LOWEST_EPHEMERAL_PORT}"
+        );
+    }
+}
+
+#[test]
+#[should_panic(expected = "No available ports found in range")]
+fn exhausting_a_leased_range_panics() {
+    let mut available_ports = AvailablePorts::new(0, 0);
+
+    available_ports.get_next_ports(usize::from(PORTS_PER_INSTANCE) + 1);
+}