diff --git a/grpc/src/client/load_balancing/hierarchy.rs b/grpc/src/client/load_balancing/hierarchy.rs new file mode 100644 index 000000000..21a997237 --- /dev/null +++ b/grpc/src/client/load_balancing/hierarchy.rs @@ -0,0 +1,281 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use std::collections::HashMap; +use std::fmt; +use std::fmt::Debug; +use std::ops::Deref; +use std::sync::Arc; + +use crate::client::name_resolution::Endpoint; + +#[derive(Clone)] +struct ArcSlice { + data: Arc>, + start: usize, + end: usize, +} + +impl ArcSlice { + fn new(vec: Vec) -> Self { + let len = vec.len(); + Self { + data: Arc::new(vec), + start: 0, + end: len, + } + } + + fn pop_front(&self) -> Self { + assert!(self.end - self.start > 0); + Self { + data: Arc::clone(&self.data), + start: self.start + 1, + end: self.end, + } + } +} + +impl Deref for ArcSlice { + type Target = [String]; + fn deref(&self) -> &[String] { + &self.data[self.start..self.end] + } +} + +impl Debug for ArcSlice { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Debug::fmt(&**self, f) + } +} + +impl PartialEq for ArcSlice { + fn eq(&self, other: &Self) -> bool { + // Fast path: if they point to the exact same Arc allocation and have + // identical bounds. + if Arc::ptr_eq(&self.data, &other.data) + && self.start == other.start + && self.end == other.end + { + return true; + } + + // Compare the actual slices element-by-element by value. + **self == **other + } +} + +impl Eq for ArcSlice {} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct HierarchicalPath { + parts: ArcSlice, +} + +impl HierarchicalPath { + pub(crate) fn new(parts: Vec) -> Self { + Self { + parts: ArcSlice::new(parts), + } + } + + fn new_with_arc_slice(parts: ArcSlice) -> Self { + Self { parts } + } +} + +/// Returns the hierarchical path of endpoint. +fn from_endpoint(endpoint: &Endpoint) -> Option<&HierarchicalPath> { + endpoint.attributes.get::() +} + +fn set_in_endpoint_path(mut endpoint: Endpoint, path: HierarchicalPath) -> Endpoint { + endpoint.attributes = endpoint.attributes.add(path); + endpoint +} + +/// Overrides the hierarchical path in endpoint with path. +pub fn set_in_endpoint(endpoint: Endpoint, path: Vec) -> Endpoint { + set_in_endpoint_path(endpoint, HierarchicalPath::new(path)) +} + +/// Group splits a slice of endpoints into groups based on the first hierarchy +/// path. The first hierarchy path will be removed from the result. +/// +/// If hierarchical path is not set, or has no path in it, the endpoint is +/// dropped. +pub fn group(endpoints: Vec) -> HashMap> { + let mut ret = HashMap::new(); + for endpoint in endpoints { + if let Some(path) = from_endpoint(&endpoint) { + if path.parts.is_empty() { + continue; + } + let cur_path = path.parts[0].clone(); + let new_path = HierarchicalPath::new_with_arc_slice(path.parts.pop_front()); + let new_endpoint = set_in_endpoint_path(endpoint, new_path); + ret.entry(cur_path) + .or_insert_with(Vec::new) + .push(new_endpoint); + } + } + ret +} + +#[cfg(test)] +mod test { + use super::*; + use crate::attributes::Attributes; + use crate::byte_str::ByteStr; + use crate::core::Address; + + fn test_endpoint(addr: &str) -> Endpoint { + Endpoint { + addresses: vec![Address { + network_type: "tcp", + address: ByteStr::from(addr.to_string()), + attributes: Attributes::new(), + }], + attributes: Attributes::new(), + } + } + + #[test] + fn test_from_endpoint() { + let ep_not_set = test_endpoint("a"); + assert_eq!(from_endpoint(&ep_not_set), None); + + let ep_set = set_in_endpoint(test_endpoint("a"), vec!["a".to_string(), "b".to_string()]); + assert_eq!( + from_endpoint(&ep_set), + Some(&HierarchicalPath::new(vec![ + "a".to_string(), + "b".to_string() + ])) + ); + } + + #[test] + fn test_set_in_endpoint() { + // before is not set + let ep = test_endpoint("a"); + let path = vec!["a".to_string(), "b".to_string()]; + let new_ep = set_in_endpoint(ep, path.clone()); + assert_eq!( + from_endpoint(&new_ep), + Some(&HierarchicalPath::new(path.clone())) + ); + + // before is set + let ep = set_in_endpoint( + test_endpoint("a"), + vec!["before".to_string(), "a".to_string(), "b".to_string()], + ); + let path = vec!["a".to_string(), "b".to_string()]; + let new_ep = set_in_endpoint(ep, path.clone()); + assert_eq!(from_endpoint(&new_ep), Some(&HierarchicalPath::new(path))); + } + + #[test] + fn test_group_with_hierarchy() { + let eps = vec![ + set_in_endpoint( + test_endpoint("a0"), + vec!["a".to_string(), "b".to_string(), "c".to_string()], + ), + set_in_endpoint(test_endpoint("a1"), vec!["a".to_string(), "b".to_string()]), + set_in_endpoint( + test_endpoint("b0"), + vec!["b".to_string(), "c".to_string(), "d".to_string()], + ), + set_in_endpoint(test_endpoint("b1"), vec!["b".to_string(), "c".to_string()]), + ]; + + // 1st group call: splits by first level ("a" and "b") + let mut g1 = group(eps); + assert_eq!(g1.len(), 2); + + let eps_a = g1.remove("a").unwrap(); + let eps_b = g1.remove("b").unwrap(); + + assert_eq!( + eps_a, + vec![ + set_in_endpoint(test_endpoint("a0"), vec!["b".to_string(), "c".to_string()],), + set_in_endpoint(test_endpoint("a1"), vec!["b".to_string()],), + ] + ); + assert_eq!( + eps_b, + vec![ + set_in_endpoint(test_endpoint("b0"), vec!["c".to_string(), "d".to_string()],), + set_in_endpoint(test_endpoint("b1"), vec!["c".to_string()],), + ] + ); + + // 2nd group call: splits by second level ("b" for group "a", "c" for + // group "b") + let mut g2_a = group(eps_a); + let mut g2_b = group(eps_b); + assert_eq!(g2_a.len(), 1); + assert_eq!(g2_b.len(), 1); + + let eps_ab = g2_a.remove("b").unwrap(); + let eps_bc = g2_b.remove("c").unwrap(); + + assert_eq!( + eps_ab, + vec![ + set_in_endpoint(test_endpoint("a0"), vec!["c".to_string()],), + set_in_endpoint(test_endpoint("a1"), vec![]), + ] + ); + assert_eq!( + eps_bc, + vec![ + set_in_endpoint(test_endpoint("b0"), vec!["d".to_string()],), + set_in_endpoint(test_endpoint("b1"), vec![]), + ] + ); + + // 3rd group call: 2-level endpoints (a1, b1) have empty hierarchy and + // are dropped. 3-level endpoints (a0, b0) have their hierarchy become + // empty after 3 calls. + let mut g3_a = group(eps_ab); + let mut g3_b = group(eps_bc); + assert_eq!(g3_a.len(), 1); + assert_eq!(g3_b.len(), 1); + + let eps_abc = g3_a.remove("c").unwrap(); + let eps_bcd = g3_b.remove("d").unwrap(); + + assert_eq!(eps_abc, vec![set_in_endpoint(test_endpoint("a0"), vec![],)]); + assert_eq!(eps_bcd, vec![set_in_endpoint(test_endpoint("b0"), vec![],)]); + + // Calling group on endpoints that now have empty hierarchy returns an + // empty map. + assert!(group(eps_abc).is_empty()); + assert!(group(eps_bcd).is_empty()); + } +} diff --git a/grpc/src/client/load_balancing/mod.rs b/grpc/src/client/load_balancing/mod.rs index e517fbdcc..c25b7bdfa 100644 --- a/grpc/src/client/load_balancing/mod.rs +++ b/grpc/src/client/load_balancing/mod.rs @@ -40,6 +40,7 @@ use crate::core::Address; use crate::metadata::MetadataMap; use crate::rt::GrpcRuntime; +pub(crate) mod hierarchy; pub(crate) mod subchannel_sharing; pub mod child_manager;