[Factory Experiment] connect: use deterministic BTreeSet::first for peer/IP selection

src/connect.rs

@@ -8,6 +8,7 @@ use futures::{future::BoxFuture, Stream};
 use hyper_util::rt::TokioIo;
 use rustls::pki_types::PrivateKeyDer;
 use rustls::{ClientConfig, ServerConfig};
+use std::collections::{BTreeMap, BTreeSet};
 use std::io;
 use std::net::IpAddr;
 use std::pin::Pin;
@@ -97,33 +98,24 @@ impl IntermeshClient {
     }
 
     /// Connect to a mesh name.
+    ///
+    /// If the name resolves to multiple IMIDs, the IMID with the minimum
+    /// value (per `Ord`) is selected deterministically.
     pub(crate) async fn connect_name(
         &mut self,
         name: &Name,
         port: u16,
     ) -> Result<ClientTlsStream<TcpStream>> {
-        let imid = self
-            .te
-            .derivation()
-            .name_to_imid
-            .get(name)
-            .and_then(|s| s.iter().next())
-            .cloned()
-            .ok_or_else(|| anyhow!("failed to resolve name to IMID: {name}"))?;
-
+        let imid = resolve_name_to_imid(name, &self.te.derivation().name_to_imid)?;
         self.connect_imid(&imid, port).await
     }
 
     /// Connect to an IMID.
+    ///
+    /// If the IMID resolves to multiple IPs, the IP with the minimum value
+    /// (per `Ord`) is selected deterministically.
     async fn connect_imid(&mut self, imid: &Imid, port: u16) -> Result<ClientTlsStream<TcpStream>> {
-        // Trust engine is authoritative; bootstrap hint is fallback for bootstrap.
-        let te_ip = self
-            .te
-            .derivation()
-            .imid_to_ip
-            .get(imid)
-            .and_then(|s| s.iter().next())
-            .copied();
+        let te_ip = resolve_imid_to_ip(imid, &self.te.derivation().imid_to_ip).ok();
         let hint_ip = self
             .bootstrap_hint
             .take_if(|(hint_imid, _)| hint_imid == imid)
@@ -176,6 +168,32 @@ impl Service<Uri> for IntermeshClient {
     }
 }
 
+/// Resolve a mesh name to an IMID, deterministically selecting the minimum
+/// IMID (per `Ord`) when the name maps to multiple candidates.
+fn resolve_name_to_imid(
+    name: &Name,
+    name_to_imid: &BTreeMap<Name, BTreeSet<Imid>>,
+) -> Result<Imid> {
+    name_to_imid
+        .get(name)
+        .and_then(|s| s.first())
+        .cloned()
+        .ok_or_else(|| anyhow!("failed to resolve name to IMID: {name}"))
+}
+
+/// Resolve an IMID to an IP address, deterministically selecting the minimum
+/// IP (per `Ord`) when the IMID maps to multiple candidates.
+fn resolve_imid_to_ip(
+    imid: &Imid,
+    imid_to_ip: &BTreeMap<Imid, BTreeSet<IpAddr>>,
+) -> Result<IpAddr> {
+    imid_to_ip
+        .get(imid)
+        .and_then(|s| s.first())
+        .copied()
+        .ok_or_else(|| anyhow!("failed to find IP for IMID: {imid}"))
+}
+
 /// Creates a stream of TLS-wrapped TCP connections for serving gRPC requests.
 ///
 /// This function creates a `TlsStream` that uses the `IntermeshVerifier` to
@@ -294,7 +312,7 @@ mod tests {
     use crate::proto::intermesh::GossipUpdate;
     use crate::test_utils::TestFixture;
     use futures::TryStreamExt;
-    use std::collections::BTreeMap;
+    use std::collections::{BTreeMap, BTreeSet};
     use tokio::io::{AsyncReadExt, AsyncWriteExt};
     use tokio::sync::mpsc;
     use tokio::time::{timeout, Duration};
@@ -474,6 +492,67 @@ mod tests {
         server_task.await.assert();
     }
 
+    #[test]
+    fn test_resolve_name_to_imid_first() {
+        let a = ImidKeypair::test_keypair("a").to_imid();
+        let b = ImidKeypair::test_keypair("b").to_imid();
+        let (min, max) = if a < b {
+            (a.clone(), b.clone())
+        } else {
+            (b.clone(), a.clone())
+        };
+
+        let map = BTreeMap::from([(
+            "test.local".parse().assert(),
+            BTreeSet::from([min.clone(), max]),
+        )]);
+
+        let result = resolve_name_to_imid(&"test.local".parse().assert(), &map).assert();
+        assert_eq!(
+            result, min,
+            "should select the minimum IMID per Ord (BTreeSet::first)"
+        );
+    }
+
+    #[test]
+    fn test_resolve_name_to_imid_not_found() {
+        let a = ImidKeypair::test_keypair("a").to_imid();
+        let map = BTreeMap::from([("test.local".parse().assert(), BTreeSet::from([a]))]);
+
+        let err = resolve_name_to_imid(&"other.local".parse().assert(), &map).unwrap_err();
+        assert!(
+            err.to_string().contains("failed to resolve name to IMID"),
+            "unexpected error: {err}"
+        );
+    }
+
+    #[test]
+    fn test_resolve_imid_to_ip_first() {
+        let imid = ImidKeypair::test_keypair("test").to_imid();
+        let ip: IpAddr = "127.0.0.1".parse().assert();
+
+        let map = BTreeMap::from([(imid.clone(), BTreeSet::from([ip]))]);
+
+        let result = resolve_imid_to_ip(&imid, &map).assert();
+        assert_eq!(result, ip);
+    }
+
+    #[test]
+    fn test_resolve_imid_to_ip_not_found() {
+        let imid = ImidKeypair::test_keypair("test").to_imid();
+        let other = ImidKeypair::test_keypair("other").to_imid();
+        let map = BTreeMap::from([(
+            imid,
+            BTreeSet::from(["127.0.0.1".parse::<IpAddr>().assert()]),
+        )]);
+
+        let err = resolve_imid_to_ip(&other, &map).unwrap_err();
+        assert!(
+            err.to_string().contains("failed to find IP for IMID"),
+            "unexpected error: {err}"
+        );
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_tonic_integration() {
         timeout(Duration::from_secs(30), async {