refactor ClusterMap to never hand out dashmap references

glindstedt · glindstedt · commit 59810e42babd · 2025-10-21T14:13:40.000+02:00
diff --git a/benches/cluster_map.rs b/benches/cluster_map.rs
@@ -17,21 +17,19 @@ mod serde {
     fn serialize_to_protobuf(cm: &ClusterMap) -> Vec<Any> {
         let mut resources = Vec::new();
 
-        for cluster in cm.iter() {
-            resources.push(
-                Resource::Cluster(Cluster {
-                    locality: cluster.key().clone().map(From::from),
-                    endpoints: cluster
-                        .endpoints
-                        .iter()
-                        .map(TryFrom::try_from)
-                        .collect::<Result<_, _>>()
-                        .unwrap(),
-                })
-                .try_encode()
-                .unwrap(),
-            );
-        }
+        cm.iter_with(|locality, endpoint_set| {
+            Resource::Cluster(Cluster {
+                locality: locality.clone().map(From::from),
+                endpoints: endpoint_set
+                    .endpoints
+                    .iter()
+                    .map(TryFrom::try_from)
+                    .collect::<Result<_, _>>()
+                    .unwrap(),
+            })
+        })
+        .into_iter()
+        .map(|resource| resources.push(resource.try_encode().unwrap()));
 
         resources
     }
@@ -111,11 +109,11 @@ mod ops {
     use shared::{GenCluster, gen_cluster_map};
 
     fn compute_hash<const S: u64>(gc: &GenCluster) -> usize {
-        let mut total_endpoints = 0;
-
-        for kv in gc.cm.iter() {
-            total_endpoints += kv.endpoints.len();
-        }
+        let total_endpoints = gc
+            .cm
+            .iter_with(|_locality, endpoint_set| endpoint_set.len())
+            .iter()
+            .fold(0, |acc, e| acc + e);
 
         assert_eq!(total_endpoints, gc.total_endpoints);
         total_endpoints
diff --git a/benches/shared.rs b/benches/shared.rs
@@ -690,7 +690,7 @@ pub fn gen_cluster_map<const S: u64>(token_kind: TokenKind) -> GenCluster {
 
     // Now actually insert the endpoints, now that the order of keys is established,
     // annoying, but note we split out iteration versus insertion, otherwise we deadlock
-    let keys: Vec<_> = cm.iter().map(|kv| kv.key().clone()).collect();
+    let keys = cm.iter_with(|locality, _endpoint_set| locality.clone());
     let mut sets = std::collections::BTreeMap::new();
 
     let mut token_generator = match token_kind {
diff --git a/benches/token_router.rs b/benches/token_router.rs
@@ -14,13 +14,18 @@ fn token_router(b: Bencher<'_, '_>, token_kind: &str) {
     let cm = std::sync::Arc::new(gc.cm);
 
     // Calculate the amount of bytes for all the tokens
-    for eps in cm.iter() {
-        for ep in &eps.value().endpoints {
+    cm.iter_with(|_locality, endpoint_set| {
+        let mut tokens = Vec::new();
+        for ep in &endpoint_set.endpoints {
             for tok in ep.metadata.known.tokens.iter() {
                 tokens.push(tok.clone());
             }
         }
-    }
+        tokens
+    })
+    .iter()
+    .flatten()
+    .map(|tok| tokens.push(tok));
 
     let total_token_size: usize = tokens.iter().map(|t| t.len()).sum();
     let pool = std::sync::Arc::new(quilkin::collections::BufferPool::new(1, 1));
diff --git a/src/config.rs b/src/config.rs
@@ -644,40 +644,43 @@ impl Config {
                     }
                 }
                 ResourceType::Cluster => {
-                    let mut push = |key: &Option<crate::net::endpoint::Locality>,
-                                    value: &crate::net::cluster::EndpointSet|
-                     -> crate::Result<()> {
-                        let version = value.version().to_string();
-                        let key_s = key.as_ref().map(|k| k.to_string()).unwrap_or_default();
-
-                        if client_state.version_matches(&key_s, &version) {
-                            return Ok(());
-                        }
-
-                        let resource = crate::xds::Resource::Cluster(
-                            quilkin_xds::generated::quilkin::config::v1alpha1::Cluster {
-                                locality: key.clone().map(|l| l.into()),
-                                endpoints: value.endpoints.iter().map(|ep| ep.into()).collect(),
-                            },
-                        );
+                    let cluster_resource_transformer =
+                        |key: &Option<crate::net::endpoint::Locality>,
+                         value: &crate::net::cluster::EndpointSet|
+                         -> crate::Result<Option<XdsResource>> {
+                            let version = value.version().to_string();
+                            let key_s = key.as_ref().map(|k| k.to_string()).unwrap_or_default();
+
+                            if client_state.version_matches(&key_s, &version) {
+                                return Ok(None);
+                            }
 
-                        resources.push(XdsResource {
-                            name: key_s,
-                            version,
-                            resource: Some(resource.try_encode()?),
-                            ..Default::default()
-                        });
+                            let resource = crate::xds::Resource::Cluster(
+                                quilkin_xds::generated::quilkin::config::v1alpha1::Cluster {
+                                    locality: key.clone().map(|l| l.into()),
+                                    endpoints: value.endpoints.iter().map(|ep| ep.into()).collect(),
+                                },
+                            );
 
-                        Ok(())
-                    };
+                            Ok(Some(XdsResource {
+                                name: key_s,
+                                version,
+                                resource: Some(resource.try_encode()?),
+                                ..Default::default()
+                            }))
+                        };
 
                     let Some(clusters) = self.dyn_cfg.clusters() else {
                         break 'append;
                     };
 
                     if client_state.subscribed.is_empty() {
-                        for cluster in clusters.read().iter() {
-                            push(cluster.key(), cluster.value())?;
+                        for cluster_resources in
+                            clusters.read().iter_with(cluster_resource_transformer)
+                        {
+                            if let Some(clr) = cluster_resources? {
+                                resources.push(clr);
+                            }
                         }
                     } else {
                         for locality in client_state.subscribed.iter().filter_map(|name| {
@@ -687,18 +690,22 @@ impl Config {
                                 name.parse().ok().map(Some)
                             }
                         }) {
-                            if let Some(cluster) = clusters.read().get(&locality) {
-                                push(cluster.key(), cluster.value())?;
+                            if let Some(cluster_resource) =
+                                clusters.read().with_value(&locality, |entry| {
+                                    cluster_resource_transformer(entry.key(), entry.value())
+                                })
+                            {
+                                if let Some(clr) = cluster_resource? {
+                                    resources.push(clr);
+                                }
                             }
                         }
                     };
 
                     // Currently, we have exactly _one_ special case for removed resources, which
                     // is when ClusterMap::update_unlocated_endpoints is called to move the None
                     // locality endpoints to another one, so we just detect that case manually
-                    if client_state.versions.contains_key("")
-                        && clusters.read().get(&None).is_none()
-                    {
+                    if client_state.versions.contains_key("") && !clusters.read().exists(&None) {
                         removed.insert("".into());
                     }
                 }
diff --git a/src/metrics.rs b/src/metrics.rs
@@ -676,14 +676,13 @@ pub(crate) fn apply_clusters(clusters: &crate::config::Watch<crate::net::Cluster
     let clusters = clusters.read();
     crate::net::cluster::active_clusters().set(clusters.len() as i64);
 
-    for entry in clusters.iter() {
+    clusters.iter_with(|locality, endpoint_set| {
         crate::net::cluster::active_endpoints(
-            &entry
-                .key()
+            &locality
                 .clone()
                 .map(|key| key.to_string())
                 .unwrap_or_default(),
         )
-        .set(entry.value().len() as i64);
-    }
+        .set(endpoint_set.len() as i64);
+    });
 }
diff --git a/src/net/cluster.rs b/src/net/cluster.rs
@@ -263,8 +263,6 @@ pub struct ClusterMap<S = gxhash::GxBuildHasher> {
     version: AtomicU64,
 }
 
-type DashMapRef<'inner> = dashmap::mapref::one::Ref<'inner, Option<Locality>, EndpointSet>;
-
 impl ClusterMap {
     pub fn new() -> Self {
         Self::default()
@@ -367,15 +365,39 @@ where
     }
 
     #[inline]
-    pub fn get(&self, key: &Option<Locality>) -> Option<DashMapRef<'_>> {
-        self.map.get(key)
+    pub fn exists(&self, key: &Option<Locality>) -> bool {
+        self.map.get(key).is_some()
+    }
+
+    #[inline]
+    pub fn with_value<F, T>(&self, key: &Option<Locality>, func: F) -> Option<T>
+    where
+        F: FnOnce(dashmap::mapref::one::Ref<'_, Option<Locality>, EndpointSet>) -> T,
+    {
+        self.map.get(key).map(func)
     }
 
     #[inline]
     pub fn insert_default(&self, endpoints: BTreeSet<Endpoint>) {
         self.insert(None, None, endpoints);
     }
 
+    /// Iterates over the entries in the `ClusterMap` 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_with<F, T>(&self, func: F) -> Vec<T>
+    where
+        F: for<'a> Fn(&'a Option<Locality>, &'a EndpointSet) -> T,
+    {
+        let mut results: Vec<T> = Vec::new();
+        for entry in self.map.iter() {
+            results.push(func(entry.key(), entry.value()));
+        }
+        results
+    }
+
     #[inline]
     pub fn remove_endpoint(&self, needle: &Endpoint) -> bool {
         let locality = 'l: {
@@ -430,11 +452,6 @@ where
         true
     }
 
-    #[inline]
-    pub fn iter(&self) -> dashmap::iter::Iter<'_, Option<Locality>, EndpointSet, S> {
-        self.map.iter()
-    }
-
     #[inline]
     pub fn replace(
         &self,
@@ -477,7 +494,7 @@ where
     }
 
     pub fn nth_endpoint(&self, mut index: usize) -> Option<Endpoint> {
-        for set in self.iter() {
+        for set in self.map.iter() {
             let set = &set.value().endpoints;
             if index < set.len() {
                 return set.iter().nth(index).cloned();
@@ -492,7 +509,7 @@ where
     pub fn filter_endpoints(&self, f: impl Fn(&Endpoint) -> bool) -> Vec<Endpoint> {
         let mut endpoints = Vec::new();
 
-        for set in self.iter() {
+        for set in self.map.iter() {
             for endpoint in set.endpoints.iter().filter(|e| (f)(e)) {
                 endpoints.push(endpoint.clone());
             }
@@ -639,8 +656,9 @@ where
     S: Default + std::hash::BuildHasher + Clone,
 {
     fn eq(&self, rhs: &Self) -> bool {
-        for a in self.iter() {
+        for a in self.map.iter() {
             match rhs
+                .map
                 .get(a.key())
                 .filter(|b| a.value().endpoints == b.endpoints)
             {
@@ -808,16 +826,18 @@ mod tests {
         cluster1.insert(None, Some(nl1.clone()), [endpoint.clone()].into());
         cluster1.insert(None, Some(de1.clone()), [endpoint.clone()].into());
 
-        assert_eq!(cluster1.get(&Some(nl1.clone())).unwrap().len(), 1);
+        assert_eq!(cluster1.map.get(&Some(nl1.clone())).unwrap().len(), 1);
         assert!(
             cluster1
+                .map
                 .get(&Some(nl1.clone()))
                 .unwrap()
                 .contains(&endpoint)
         );
-        assert_eq!(cluster1.get(&Some(de1.clone())).unwrap().len(), 1);
+        assert_eq!(cluster1.map.get(&Some(de1.clone())).unwrap().len(), 1);
         assert!(
             cluster1
+                .map
                 .get(&Some(de1.clone()))
                 .unwrap()
                 .contains(&endpoint)
@@ -827,19 +847,20 @@ mod tests {
 
         cluster1.insert(None, Some(de1.clone()), [endpoint.clone()].into());
 
-        assert_eq!(cluster1.get(&Some(nl1.clone())).unwrap().len(), 1);
-        assert_eq!(cluster1.get(&Some(de1.clone())).unwrap().len(), 1);
+        assert_eq!(cluster1.map.get(&Some(nl1.clone())).unwrap().len(), 1);
+        assert_eq!(cluster1.map.get(&Some(de1.clone())).unwrap().len(), 1);
         assert!(
             cluster1
+                .map
                 .get(&Some(de1.clone()))
                 .unwrap()
                 .contains(&endpoint)
         );
 
         cluster1.insert(None, Some(de1.clone()), <_>::default());
 
-        assert_eq!(cluster1.get(&Some(nl1.clone())).unwrap().len(), 1);
-        assert!(cluster1.get(&Some(de1.clone())).unwrap().is_empty());
+        assert_eq!(cluster1.map.get(&Some(nl1.clone())).unwrap().len(), 1);
+        assert!(cluster1.map.get(&Some(de1.clone())).unwrap().is_empty());
     }
 
     #[test]
@@ -862,12 +883,15 @@ mod tests {
             [Endpoint::new((Ipv4Addr::new(20, 20, 20, 20), 1234).into())].into();
 
         cluster.insert(Some(nl02.into()), Some(nl1.clone()), not_expected.clone());
-        assert_eq!(cluster.get(&Some(nl1.clone())).unwrap().endpoints, expected);
+        assert_eq!(
+            cluster.map.get(&Some(nl1.clone())).unwrap().endpoints,
+            expected
+        );
 
         cluster.remove_locality(Some(nl01.into()), &Some(nl1.clone()));
 
         cluster.insert(Some(nl02.into()), Some(nl1.clone()), not_expected.clone());
-        assert_eq!(cluster.get(&Some(nl1)).unwrap().endpoints, not_expected);
+        assert_eq!(cluster.map.get(&Some(nl1)).unwrap().endpoints, not_expected);
     }
 
     #[test]
@@ -886,7 +910,7 @@ mod tests {
                 assert!(cluster.remove_endpoint(ep));
             }
 
-            assert!(cluster.get(&None).is_none());
+            assert!(cluster.map.get(&None).is_none());
         }
 
         {
@@ -895,7 +919,7 @@ mod tests {
                 assert!(cluster.remove_endpoint_if(|_ep| true));
             }
 
-            assert!(cluster.get(&None).is_none());
+            assert!(cluster.map.get(&None).is_none());
         }
     }
 }
