Relay: reserve a bridge slot at the match, give it back on every exit

max_bridges compared against bridges already running, so joins waiting
(up to attach_wait) for their host's attach all passed the check and a
burst could exceed RELAY_MAX_BRIDGES.

- BridgeSlot, taken under the lookup lock when a guest is matched with
  its host, counts toward max_bridges from then on; dropping it gives it
  back, so every way out of a join (no attach, a refused or forged
  attach, an error inside bridge(), a panic) releases it.
- counts().bridges now includes matched joins awaiting attach.
- tests/bridge_slots.rs: a waiting join holds its slot (a second guest
  gets Full); the count returns to 0 after a timeout, after a forged
  attach, and after a bridge closes; refusals never take one.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01SJ3stH5q6PgMAVUSmJqKDe

crates/horizon-relay/src/lib.rs

@@ -116,6 +116,8 @@ impl Config {
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct Counts {
     pub hosts: usize,
+    /// Bridges open, and joins matched with a host and waiting for its attach (each holds one of
+    /// `max_bridges` from the match on).
     pub bridges: usize,
     pub registrations: u64,
     pub joins: u64,
@@ -485,8 +487,28 @@ fn serve_host(
     Ok(())
 }
 
-enum Lookup {
-    Found(Sender<Out>),
+/// One of `max_bridges`, reserved when a guest is matched with its host and given back when it drops:
+/// on every way out of the join (no attach, a refusal, an error, a panic) and at the bridge's end. A
+/// reservation counts from the match, not from the bridge, so joins waiting for their host's attach
+/// can't add up past the cap.
+struct BridgeSlot<'a>(&'a AtomicUsize);
+impl<'a> BridgeSlot<'a> {
+    fn take(count: &'a AtomicUsize, max: usize) -> Option<Self> {
+        if count.fetch_add(1, Ordering::AcqRel) >= max {
+            count.fetch_sub(1, Ordering::AcqRel);
+            return None;
+        }
+        Some(Self(count))
+    }
+}
+impl Drop for BridgeSlot<'_> {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::AcqRel);
+    }
+}
+
+enum Lookup<'a> {
+    Found(Sender<Out>, BridgeSlot<'a>),
     Unknown,
     Ended,
     Version(String),
@@ -521,18 +543,20 @@ fn join(
             .ended
             .get(&nameplate)
             .is_some_and(|t| now.saturating_duration_since(*t) < state.config.ended_kept);
-        let busy = state.bridges.load(Ordering::Relaxed) >= state.config.max_bridges;
         let lookup = match inner.hosts.get_mut(&nameplate) {
             None if ended => Lookup::Ended,
             None => Lookup::Unknown,
             Some(h) if h.protocol != protocol => Lookup::Version(h.game.clone()),
-            Some(h) if busy || h.pending >= state.config.max_pending_per_host => Lookup::Full,
-            Some(h) => {
-                h.pending += 1;
-                Lookup::Found(h.to_host.clone())
-            }
+            Some(h) if h.pending >= state.config.max_pending_per_host => Lookup::Full,
+            Some(h) => match BridgeSlot::take(&state.bridges, state.config.max_bridges) {
+                Some(slot) => {
+                    h.pending += 1;
+                    Lookup::Found(h.to_host.clone(), slot)
+                }
+                None => Lookup::Full,
+            },
         };
-        if matches!(lookup, Lookup::Found(_)) {
+        if matches!(lookup, Lookup::Found(..)) {
             inner.pending.insert(
                 session,
                 Pending {
@@ -543,8 +567,8 @@ fn join(
         }
         lookup
     };
-    let to_host = match lookup {
-        Lookup::Found(tx) => tx,
+    let (to_host, slot) = match lookup {
+        Lookup::Found(tx, slot) => (tx, slot),
         Lookup::Unknown => {
             state.unknown.fetch_add(1, Ordering::Relaxed);
             locked(&state.limiter).unknown(peer.ip(), now);
@@ -573,8 +597,11 @@ fn join(
         }
     }
     match attached {
-        Ok((host, host_peer)) => bridge(state, stream, peer, host, host_peer, session),
-        Err(_) => refuse(&mut stream, ErrorCode::HostUnreachable),
+        Ok((host, host_peer)) => bridge(state, stream, peer, host, host_peer, session, slot),
+        Err(_) => {
+            drop(slot);
+            refuse(&mut stream, ErrorCode::HostUnreachable)
+        }
     }
 }
 
@@ -610,7 +637,8 @@ fn attach(
     }
 }
 
-/// Tell both ends they are bridged, then copy each way until either ends.
+/// Tell both ends they are bridged, then copy each way until either ends. `_slot` (the guest's
+/// reservation of one of `max_bridges`) is given back when this returns, however it returns.
 fn bridge(
     state: &Arc<State>,
     mut guest: TcpStream,
@@ -618,9 +646,9 @@ fn bridge(
     mut host: TcpStream,
     host_peer: SocketAddr,
     session: u64,
+    _slot: BridgeSlot<'_>,
 ) -> io::Result<()> {
-    state.bridges.fetch_add(1, Ordering::Relaxed);
-    let result = (|| -> io::Result<()> {
+    (|| -> io::Result<()> {
         wire::send(
             &mut guest,
             &FromRelay::Bridged {
@@ -646,9 +674,7 @@ fn bridge(
         copy(state, host, guest);
         let _ = up.join();
         Ok(())
-    })();
-    state.bridges.fetch_sub(1, Ordering::Relaxed);
-    result
+    })()
 }
 
 fn copy(state: &State, mut from: TcpStream, mut to: TcpStream) {

crates/horizon-relay/tests/bridge_slots.rs

@@ -0,0 +1,217 @@
+//! `max_bridges` counts from the moment a guest is matched with its host, and every way out of a
+//! join gives the slot back, spoken to the relay over raw sockets.
+use horizon_friends::wire::{self, ErrorCode, FromRelay, ToRelay};
+use horizon_relay::{Config, Relay};
+use std::io::BufReader;
+use std::net::{Shutdown, SocketAddr, TcpStream};
+use std::time::{Duration, Instant};
+
+const PROTOCOL: i32 = 777;
+
+fn relay(max_bridges: usize) -> (Relay, SocketAddr) {
+    let (relay, address, _) = Relay::spawn_local(Config {
+        max_bridges,
+        per_address: false,
+        attach_wait: Duration::from_millis(400),
+        ..Config::default()
+    })
+    .unwrap();
+    (relay, address)
+}
+
+fn connect(address: SocketAddr) -> TcpStream {
+    let s = TcpStream::connect(address).unwrap();
+    s.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
+    s
+}
+
+/// A registered host's control connection (reader, writer), its nameplate and token.
+fn host(address: SocketAddr) -> (BufReader<TcpStream>, TcpStream, String, String) {
+    let mut s = connect(address);
+    wire::send(
+        &mut s,
+        &ToRelay::Host {
+            v: wire::VERSION,
+            protocol: PROTOCOL,
+            game: "26.3".into(),
+        },
+    )
+    .unwrap();
+    let mut reader = BufReader::new(s.try_clone().unwrap());
+    match wire::recv_buffered::<_, FromRelay>(&mut reader).unwrap() {
+        FromRelay::Registered {
+            nameplate, token, ..
+        } => (reader, s, nameplate, token),
+        other => panic!("{other:?}"),
+    }
+}
+
+/// A guest's join, sent; the stream to read the relay's answer from.
+fn join(address: SocketAddr, nameplate: &str) -> TcpStream {
+    let mut s = connect(address);
+    wire::send(
+        &mut s,
+        &ToRelay::Join {
+            v: wire::VERSION,
+            nameplate: nameplate.into(),
+            protocol: PROTOCOL,
+            game: "26.3".into(),
+        },
+    )
+    .unwrap();
+    s
+}
+
+/// The next message that isn't a ping.
+fn next(reader: &mut BufReader<TcpStream>) -> FromRelay {
+    loop {
+        match wire::recv_buffered::<_, FromRelay>(reader).unwrap() {
+            FromRelay::Ping => continue,
+            m => return m,
+        }
+    }
+}
+
+fn answer(s: &mut TcpStream) -> FromRelay {
+    wire::recv(s).unwrap()
+}
+
+fn settles_to_zero(relay: &Relay) {
+    let deadline = Instant::now() + Duration::from_secs(5);
+    while relay.counts().bridges != 0 {
+        assert!(
+            Instant::now() < deadline,
+            "bridges stuck at {}",
+            relay.counts().bridges
+        );
+        std::thread::sleep(Duration::from_millis(10));
+    }
+}
+
+#[test]
+fn a_waiting_join_holds_its_slot_and_a_timeout_gives_it_back() {
+    let (relay, address) = relay(1);
+    let (mut control, _keep, nameplate, _) = host(address);
+    // The first guest is matched: the host is asked to attach, and doesn't.
+    let mut first = join(address, &nameplate);
+    assert!(matches!(next(&mut control), FromRelay::Join { .. }));
+    assert_eq!(relay.counts().bridges, 1, "reserved at the match");
+    // A second guest meanwhile finds the relay full.
+    let mut second = join(address, &nameplate);
+    assert!(matches!(
+        answer(&mut second),
+        FromRelay::Error {
+            code: ErrorCode::Full,
+            ..
+        }
+    ));
+    // The first one's wait ends: refused, and the slot is back.
+    assert!(matches!(
+        answer(&mut first),
+        FromRelay::Error {
+            code: ErrorCode::HostUnreachable,
+            ..
+        }
+    ));
+    settles_to_zero(&relay);
+    // So the next guest is matched again.
+    let mut third = join(address, &nameplate);
+    assert!(matches!(next(&mut control), FromRelay::Join { .. }));
+    assert!(matches!(
+        answer(&mut third),
+        FromRelay::Error {
+            code: ErrorCode::HostUnreachable,
+            ..
+        }
+    ));
+    settles_to_zero(&relay);
+}
+
+#[test]
+fn a_bridge_gives_its_slot_back_when_either_end_closes() {
+    let (relay, address) = relay(1);
+    let (mut control, _keep, nameplate, token) = host(address);
+    let mut guest = join(address, &nameplate);
+    let FromRelay::Join { session } = next(&mut control) else {
+        panic!()
+    };
+    let mut attach = connect(address);
+    wire::send(
+        &mut attach,
+        &ToRelay::Attach {
+            v: wire::VERSION,
+            session,
+            token,
+        },
+    )
+    .unwrap();
+    assert!(matches!(answer(&mut guest), FromRelay::Bridged { .. }));
+    assert!(matches!(answer(&mut attach), FromRelay::Bridged { .. }));
+    assert_eq!(relay.counts().bridges, 1);
+    guest.shutdown(Shutdown::Both).unwrap();
+    settles_to_zero(&relay);
+}
+
+#[test]
+fn a_wrong_attach_or_a_refusal_holds_no_slot() {
+    let (relay, address) = relay(1);
+    let (mut control, _keep, nameplate, _) = host(address);
+    // Refused before any match: an unknown code, another protocol.
+    let mut unknown = join(address, "ZZZZ");
+    assert!(matches!(
+        answer(&mut unknown),
+        FromRelay::Error {
+            code: ErrorCode::UnknownCode,
+            ..
+        }
+    ));
+    let mut other = connect(address);
+    wire::send(
+        &mut other,
+        &ToRelay::Join {
+            v: wire::VERSION,
+            nameplate: nameplate.clone(),
+            protocol: PROTOCOL + 1,
+            game: "x".into(),
+        },
+    )
+    .unwrap();
+    assert!(matches!(
+        answer(&mut other),
+        FromRelay::Error {
+            code: ErrorCode::VersionMismatch,
+            ..
+        }
+    ));
+    assert_eq!(relay.counts().bridges, 0);
+    // Matched, then an attach with the wrong token: refused, and the join's wait runs out.
+    let mut guest = join(address, &nameplate);
+    let FromRelay::Join { session } = next(&mut control) else {
+        panic!()
+    };
+    let mut forged = connect(address);
+    wire::send(
+        &mut forged,
+        &ToRelay::Attach {
+            v: wire::VERSION,
+            session,
+            token: "0".repeat(32),
+        },
+    )
+    .unwrap();
+    assert!(matches!(
+        answer(&mut forged),
+        FromRelay::Error {
+            code: ErrorCode::BadRequest,
+            ..
+        }
+    ));
+    assert!(matches!(
+        answer(&mut guest),
+        FromRelay::Error {
+            code: ErrorCode::HostUnreachable,
+            ..
+        }
+    ));
+    settles_to_zero(&relay);
+}