refactor DatacenterMap to never hand out dashmap references

Author: glindstedt

crates/test/tests/mesh.rs

@@ -189,12 +189,17 @@ trace_test!(datacenter_discovery, {
     ) -> bool {
         let dcs = config.dyn_cfg.datacenters().unwrap().read();
 
-        for i in dcs.iter() {
-            dbg!(which, i.key(), i.value());
-        }
+        drop(dcs.iter_with(|key, value| {
+            dbg!(which, key, value);
+        }));
+
+        let entries = dcs
+            .iter_with(|ip_addr, dc| (ip_addr.clone(), dc.clone()))
+            .into_iter()
+            .collect::<std::collections::HashMap<std::net::IpAddr, quilkin::config::Datacenter>>();
 
-        let ipv4_dc = dcs.get(&std::net::Ipv4Addr::LOCALHOST.into());
-        let ipv6_dc = dcs.get(&std::net::Ipv6Addr::LOCALHOST.into());
+        let ipv4_dc = entries.get(&std::net::Ipv4Addr::LOCALHOST.into());
+        let ipv6_dc = entries.get(&std::net::Ipv6Addr::LOCALHOST.into());
 
         if counter > 0 {
             match (ipv4_dc, ipv6_dc) {

src/config.rs

@@ -594,32 +594,40 @@ impl Config {
                     }
 
                     if let Some(datacenters) = self.dyn_cfg.datacenters() {
-                        for entry in datacenters.read().iter() {
-                            let host = entry.key().to_string();
-                            let qcmp_port = entry.qcmp_port;
-                            let version =
-                                resource_version(entry.icao_code.to_string().as_str(), qcmp_port);
-
-                            if client_state.version_matches(&host, &version) {
-                                continue;
-                            }
+                        let dc_resource_transformer =
+                            |ip_addr: &std::net::IpAddr,
+                             dc: &Datacenter|
+                             -> eyre::Result<Option<XdsResource>> {
+                                let host = ip_addr.to_string();
+                                let qcmp_port = dc.qcmp_port;
+                                let version =
+                                    resource_version(dc.icao_code.to_string().as_str(), qcmp_port);
+
+                                if client_state.version_matches(&host, &version) {
+                                    return Ok(None);
+                                }
 
-                            let resource = crate::xds::Resource::Datacenter(
-                                crate::net::cluster::proto::Datacenter {
-                                    qcmp_port: qcmp_port as _,
-                                    icao_code: entry.icao_code.to_string(),
-                                    host: host.clone(),
-                                },
-                            );
+                                let resource = crate::xds::Resource::Datacenter(
+                                    crate::net::cluster::proto::Datacenter {
+                                        qcmp_port: qcmp_port as _,
+                                        icao_code: dc.icao_code.to_string(),
+                                        host: host.clone(),
+                                    },
+                                );
+
+                                Ok(Some(XdsResource {
+                                    name: host,
+                                    version,
+                                    resource: Some(resource.try_encode()?),
+                                    aliases: Vec::new(),
+                                    ttl: None,
+                                    cache_control: None,
+                                }))
+                            };
 
-                            resources.push(XdsResource {
-                                name: host,
-                                version,
-                                resource: Some(resource.try_encode()?),
-                                aliases: Vec::new(),
-                                ttl: None,
-                                cache_control: None,
-                            });
+                        for resource in datacenters.read().iter_with(dc_resource_transformer) {
+                            let Some(resource) = resource? else { continue };
+                            resources.push(resource);
                         }
 
                         {
@@ -628,7 +636,7 @@ impl Config {
                                 let Ok(addr) = key.parse() else {
                                     continue;
                                 };
-                                if dc.get(&addr).is_none() {
+                                if !dc.exists(&addr) {
                                     removed.insert(key.clone());
                                 }
                             }

src/config/datacenter.rs

@@ -37,13 +37,24 @@ impl DatacenterMap {
     }
 
     #[inline]
-    pub fn get(&self, key: &IpAddr) -> Option<dashmap::mapref::one::Ref<'_, IpAddr, Datacenter>> {
-        self.map.get(key)
+    pub fn exists(&self, key: &IpAddr) -> bool {
+        self.map.get(key).is_some()
     }
 
+    /// Iterates over the entries in the `DatacenterMap` with the given func
+    ///
+    /// This ensures that the dashmap entry references are never held across await boundaries as
+    /// the func cannot be async.
     #[inline]
-    pub fn iter(&self) -> dashmap::iter::Iter<'_, IpAddr, Datacenter> {
-        self.map.iter()
+    pub fn iter_with<F, T>(&self, func: F) -> Vec<T>
+    where
+        F: for<'a> Fn(&'a IpAddr, &'a Datacenter) -> T,
+    {
+        let mut results: Vec<T> = Vec::new();
+        for entry in self.map.iter() {
+            results.push(func(entry.key(), entry.value()));
+        }
+        results
     }
 
     #[inline]
@@ -113,8 +124,8 @@ impl PartialEq for DatacenterMap {
             return false;
         }
 
-        for a in self.iter() {
-            match rhs.get(a.key()).filter(|b| *a.value() == **b) {
+        for a in self.map.iter() {
+            match rhs.map.get(a.key()).filter(|b| *a.value() == **b) {
                 Some(_) => {}
                 None => return false,
             }

src/net/phoenix.rs

@@ -388,10 +388,10 @@ impl<M> Phoenix<M> {
             self.nodes.remove(&removed);
         }
 
-        for entry in dcs.iter() {
-            let addr = (*entry.key(), entry.value().qcmp_port).into();
-            self.add_node_if_not_exists(addr, entry.value().icao_code);
-        }
+        drop(dcs.iter_with(|ip_addr, dc| {
+            let socket_addr = (*ip_addr, dc.qcmp_port).into();
+            self.add_node_if_not_exists(socket_addr, dc.icao_code);
+        }));
     }
 }