happy maps (#682)

* Use iddqd::BiHashMap for scope_id/interface map

The unnumbered manager needs bidirectional lookup between scope_id and
interface name. The prod code used a single HashMap with O(n) reverse
scans; the mock maintained two manually synchronized HashMaps. Replace
both with a ScopeMap newtype wrapping iddqd::BiHashMap.

* Use iddqd IdOrdMap for Router.sessions

Updates Router.sessions to be a new SessionMap which wraps an IdOrdMap,
allowing us to use the SessionRunner's PeerId as a HashMap key. This
ensures the map keys are borrowed from the values, so it's impossible to
insert a SessionRunner into the map using a key that isn't derived from
itself, i.e. you can't insert a SessionRunner behind the wrong PeerId.

* Remove redundant peer_to_session map

The peer_to_session map was very similar to the SessionMap in its uses.
It mapped a PeerId to a SessionEndpoint, which was just a lightweight
wrapper for a SessionRunner's event_tx (FSM event sender) and config
(SessionInfo) used by the dispatcher during lookups for incoming TCP
connections. Since both event_tx and SessionInfo are info coming from
SessionRunner, this can all be accessed using the existing SessionMap
without needing to maintain two separate maps. This gets rid of the
peer_to_session map in lieu of using the SessionMap.

* Review Feedback

Signed-off-by: Trey Aspelund <trey@oxidecomputer.com>

---------

Signed-off-by: Trey Aspelund <trey@oxidecomputer.com>

Author: taspelund

Cargo.lock

@@ -135,6 +135,7 @@ serial_test = "3.3"
 internet-checksum = "0.2.1"
 network-interface = { git = "https://github.com/oxidecomputer/network-interface", branch = "illumos" }
 natord = "1.0"
+iddqd = "0.3"
 
 
 [workspace.dependencies.opte-ioctl]

Cargo.toml

@@ -23,6 +23,7 @@ itertools.workspace = true
 oxnet.workspace = true
 uuid.workspace = true
 rand.workspace = true
+iddqd.workspace = true
 clap = { workspace = true, optional = true }
 
 [target.'cfg(target_os = "illumos")'.dependencies]

bgp/Cargo.toml

@@ -6,12 +6,12 @@ use crate::{
     clock::ConnectionClock,
     error::Error,
     messages::Message,
-    session::{FsmEvent, PeerId, SessionEndpoint, SessionInfo},
+    router::SessionMap,
+    session::{FsmEvent, SessionInfo},
     unnumbered::UnnumberedManager,
 };
 use slog::Logger;
 use std::{
-    collections::BTreeMap,
     net::{SocketAddr, ToSocketAddrs},
     sync::{Arc, Mutex, mpsc::Sender},
     thread::JoinHandle,
@@ -49,11 +49,11 @@ pub trait BgpListener<Cnx: BgpConnection> {
     /// Accept a connection. This Listener is non-blocking, so the timeout
     /// is used as a sleep between accept attempts. This function may be called
     /// multiple times, returning a new connection each time. Policy application
-    /// is handled by the Dispatcher after the peer_to_session lookup.
+    /// is handled by the Dispatcher after the session lookup.
     fn accept(
         &self,
         log: Logger,
-        peer_to_session: Arc<Mutex<BTreeMap<PeerId, SessionEndpoint<Cnx>>>>,
+        sessions: Arc<Mutex<SessionMap<Cnx>>>,
         timeout: Duration,
     ) -> Result<Cnx, Error>;
 

bgp/src/connection.rs

@@ -17,15 +17,14 @@ use crate::{
     error::Error,
     log::{connection_log, connection_log_lite},
     messages::Message,
-    session::{
-        ConnectionEvent, FsmEvent, PeerId, SessionEndpoint, SessionInfo,
-    },
+    router::SessionMap,
+    session::{ConnectionEvent, FsmEvent, PeerId, SessionInfo},
     unnumbered::UnnumberedManager,
 };
 use mg_common::lock;
 use slog::{Logger, info};
 use std::{
-    collections::{BTreeMap, HashMap},
+    collections::HashMap,
     net::{SocketAddr, ToSocketAddrs},
     sync::{
         Arc, Mutex,
@@ -244,38 +243,32 @@ impl BgpListener<BgpConnectionChannel> for BgpListenerChannel {
     fn accept(
         &self,
         log: Logger,
-        peer_to_session: Arc<
-            Mutex<BTreeMap<PeerId, SessionEndpoint<BgpConnectionChannel>>>,
-        >,
+        sessions: Arc<Mutex<SessionMap<BgpConnectionChannel>>>,
         timeout: Duration,
     ) -> Result<BgpConnectionChannel, Error> {
         let (peer, endpoint) = self.listener.accept(timeout)?;
 
-        // For channel-based test connections, we use the bind address as the local
-        // address. In a real network scenario (like TCP), we would need to get the
-        // actual connection's local address to handle dual-stack correctly, but for
-        // testing purposes with channels, the bind address is the connection address.
         let local = self.bind_addr;
 
         // Resolve peer address to appropriate PeerId (IP or Interface)
         let key = self.resolve_session_key(peer);
 
-        match lock!(peer_to_session).get(&key) {
-            Some(session_endpoint) => {
-                let config = lock!(session_endpoint.config);
-                Ok(BgpConnectionChannel::with_conn(
-                    local,
-                    peer,
-                    endpoint,
-                    session_endpoint.event_tx.clone(),
-                    IO_TIMEOUT,
-                    log,
-                    ConnectionDirection::Inbound,
-                    &config,
-                ))
-            }
-            None => Err(Error::UnknownPeer(peer.ip())),
-        }
+        let runner = lock!(sessions)
+            .get(&key)
+            .cloned()
+            .ok_or(Error::UnknownPeer(peer.ip()))?;
+
+        let config = lock!(runner.session);
+        Ok(BgpConnectionChannel::with_conn(
+            local,
+            peer,
+            endpoint,
+            runner.event_tx.clone(),
+            IO_TIMEOUT,
+            log,
+            ConnectionDirection::Inbound,
+            &config,
+        ))
     }
 
     fn apply_policy(

bgp/src/connection_channel.rs

@@ -20,16 +20,13 @@ use crate::{
         RouteRefreshParseErrorReason, notification_message_from_wire,
         open_message_from_wire, route_refresh_message_from_wire,
     },
-    session::{
-        ConnectionEvent, FsmEvent, PeerId, SessionEndpoint, SessionEvent,
-        SessionInfo,
-    },
+    router::SessionMap,
+    session::{ConnectionEvent, FsmEvent, PeerId, SessionEvent, SessionInfo},
     unnumbered::UnnumberedManager,
 };
 use mg_common::lock;
 use slog::{Logger, info};
 use std::{
-    collections::BTreeMap,
     io::Read,
     io::Write,
     net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs},
@@ -138,9 +135,7 @@ impl BgpListener<BgpConnectionTcp> for BgpListenerTcp {
     fn accept(
         &self,
         log: Logger,
-        peer_to_session: Arc<
-            Mutex<BTreeMap<PeerId, SessionEndpoint<BgpConnectionTcp>>>,
-        >,
+        sessions: Arc<Mutex<SessionMap<BgpConnectionTcp>>>,
         timeout: Duration,
     ) -> Result<BgpConnectionTcp, Error> {
         let start = Instant::now();
@@ -172,23 +167,24 @@ impl BgpListener<BgpConnectionTcp> for BgpListenerTcp {
                     // Resolve peer address to appropriate PeerId (IP or Interface)
                     let key = self.resolve_session_key(peer);
 
-                    // Check if we have a session for this peer
-                    match lock!(peer_to_session).get(&key) {
-                        Some(session_endpoint) => {
-                            let config = lock!(session_endpoint.config);
-                            return BgpConnectionTcp::with_conn(
-                                local,
-                                peer,
-                                conn,
-                                IO_TIMEOUT,
-                                session_endpoint.event_tx.clone(),
-                                log,
-                                ConnectionDirection::Inbound,
-                                &config,
-                            );
-                        }
-                        None => return Err(Error::UnknownPeer(ip)),
-                    }
+                    // Look up the session runner, clone the Arc, then release
+                    // the sessions lock before accessing session config.
+                    let runner = lock!(sessions)
+                        .get(&key)
+                        .cloned()
+                        .ok_or(Error::UnknownPeer(ip))?;
+
+                    let config = lock!(runner.session);
+                    return BgpConnectionTcp::with_conn(
+                        local,
+                        peer,
+                        conn,
+                        IO_TIMEOUT,
+                        runner.event_tx.clone(),
+                        log,
+                        ConnectionDirection::Inbound,
+                        &config,
+                    );
                 }
                 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                     // Check if we've exceeded the timeout

bgp/src/connection_tcp.rs

@@ -5,13 +5,13 @@
 use crate::{
     IO_TIMEOUT,
     connection::{BgpConnection, BgpListener},
-    session::{FsmEvent, PeerId, SessionEndpoint, SessionEvent},
+    router::SessionMap,
+    session::{FsmEvent, PeerId, SessionEvent},
     unnumbered::UnnumberedManager,
 };
 use mg_common::lock;
 use slog::{Logger, debug, error, info, warn};
 use std::{
-    collections::BTreeMap,
     net::SocketAddr,
     sync::atomic::{AtomicBool, Ordering},
     sync::{Arc, Mutex},
@@ -21,10 +21,9 @@ use std::{
 
 const UNIT_DISPATCHER: &str = "dispatcher";
 
-pub struct Dispatcher<Cnx: BgpConnection> {
-    /// Session endpoint map indexed by PeerId (IP or interface name)
-    /// This unified map supports both numbered and unnumbered BGP sessions
-    pub peer_to_session: Arc<Mutex<BTreeMap<PeerId, SessionEndpoint<Cnx>>>>,
+pub struct Dispatcher<Cnx: BgpConnection + 'static> {
+    /// Session map shared with all Routers, indexed by PeerId.
+    pub sessions: Arc<Mutex<SessionMap<Cnx>>>,
 
     /// Optional unnumbered neighbor manager for link-local connection routing.
     /// When present, enables routing of IPv6 link-local connections to
@@ -38,7 +37,7 @@ pub struct Dispatcher<Cnx: BgpConnection> {
 
 impl<Cnx: BgpConnection + 'static> Dispatcher<Cnx> {
     pub fn new(
-        peer_to_session: Arc<Mutex<BTreeMap<PeerId, SessionEndpoint<Cnx>>>>,
+        sessions: Arc<Mutex<SessionMap<Cnx>>>,
         listen: String,
         log: Logger,
         unnumbered_manager: Option<Arc<dyn UnnumberedManager>>,
@@ -50,7 +49,7 @@ impl<Cnx: BgpConnection + 'static> Dispatcher<Cnx> {
         ));
 
         Self {
-            peer_to_session,
+            sessions,
             unnumbered_manager,
             listen,
             log: Mutex::new(log),
@@ -145,7 +144,7 @@ impl<Cnx: BgpConnection + 'static> Dispatcher<Cnx> {
 
                 let accepted = match listener.accept(
                     log.clone(),
-                    self.peer_to_session.clone(),
+                    self.sessions.clone(),
                     IO_TIMEOUT,
                 ) {
                     Ok(c) => {
@@ -171,41 +170,40 @@ impl<Cnx: BgpConnection + 'static> Dispatcher<Cnx> {
                     "session_key" => format!("{key:?}"),
                 ));
 
-                match lock!(self.peer_to_session).get(&key).cloned() {
-                    Some(session_endpoint) => {
-                        // Apply connection policy from the session configuration
-                        let min_ttl = lock!(session_endpoint.config).min_ttl;
-                        let md5_key =
-                            lock!(session_endpoint.config).md5_auth_key.clone();
-
-                        if let Err(e) =
-                            Listener::apply_policy(&accepted, min_ttl, md5_key)
-                        {
-                            warn!(session_log,
-                                "failed to apply policy for connection";
-                                "error" => format!("{e}")
-                            );
-                        }
-
-                        if let Err(e) =
-                            session_endpoint.event_tx.send(FsmEvent::Session(
-                                SessionEvent::TcpConnectionAcked(accepted),
-                            ))
-                        {
-                            error!(session_log,
-                                "failed to send connected event to session";
-                                "error" => format!("{e}")
-                            );
-                            continue 'listener;
-                        }
-                    }
-                    None => {
+                let (runner, min_ttl, md5_key) = {
+                    let sessions = lock!(self.sessions);
+                    let Some(runner) = sessions.get(&key).cloned() else {
                         debug!(
                             session_log,
                             "no session found for peer, dropping connection"
                         );
                         continue 'accept;
-                    }
+                    };
+                    let config = lock!(runner.session);
+                    (
+                        runner.clone(),
+                        config.min_ttl,
+                        config.md5_auth_key.clone(),
+                    )
+                };
+
+                if let Err(e) =
+                    Listener::apply_policy(&accepted, min_ttl, md5_key)
+                {
+                    warn!(session_log,
+                        "failed to apply policy for connection";
+                        "error" => format!("{e}")
+                    );
+                }
+
+                if let Err(e) = runner.event_tx.send(FsmEvent::Session(
+                    SessionEvent::TcpConnectionAcked(accepted),
+                )) {
+                    error!(session_log,
+                        "failed to send connected event to session";
+                        "error" => format!("{e}")
+                    );
+                    continue 'listener;
                 }
             }
         }
@@ -225,7 +223,7 @@ impl<Cnx: BgpConnection + 'static> Dispatcher<Cnx> {
     }
 }
 
-impl<Cnx: BgpConnection> Drop for Dispatcher<Cnx> {
+impl<Cnx: BgpConnection + 'static> Drop for Dispatcher<Cnx> {
     fn drop(&mut self) {
         debug!(lock!(self.log), "dropping dispatcher");
     }