Merge branch 'gracefulswitch' into child_manager_work_scheduler

Author: dfawley

grpc/src/client/load_balancing/child_manager.rs

@@ -47,7 +47,7 @@ use super::{Subchannel, SubchannelState};
 // An LbPolicy implementation that manages multiple children.
 #[derive(Debug)]
 pub(crate) struct ChildManager<T: Debug, S: ResolverUpdateSharder<T>> {
-    subchannel_child_map: HashMap<WeakSubchannel, usize>,
+    subchannel_to_child_idx: HashMap<WeakSubchannel, usize>,
     children: Vec<Child<T>>,
     update_sharder: S,
     pending_work: Arc<Mutex<HashSet<usize>>>,
@@ -60,10 +60,9 @@ pub(crate) struct ChildManager<T: Debug, S: ResolverUpdateSharder<T>> {
 #[derive(Debug)]
 pub(crate) struct Child<T> {
     pub identifier: T,
-    pub policy: Box<dyn LbPolicy>,
     pub builder: Arc<dyn LbPolicyBuilder>,
     pub state: LbState,
-    pub updated: bool, // Set when the child updates its picker; cleared in child_states is called.
+    policy: Box<dyn LbPolicy>,
     work_scheduler: Arc<ChildWorkScheduler>,
 }
 
@@ -107,7 +106,7 @@ where
     ) -> Self {
         Self {
             update_sharder,
-            subchannel_child_map: Default::default(),
+            subchannel_to_child_idx: Default::default(),
             children: Default::default(),
             pending_work: Default::default(),
             runtime,
@@ -170,12 +169,11 @@ where
     ) {
         // Add all created subchannels into the subchannel_child_map.
         for csc in channel_controller.created_subchannels {
-            self.subchannel_child_map.insert(csc.into(), child_idx);
+            self.subchannel_to_child_idx.insert(csc.into(), child_idx);
         }
         // Update the tracked state if the child produced an update.
         if let Some(state) = channel_controller.picker_update {
             self.children[child_idx].state = state;
-            self.children[child_idx].updated = true;
             self.updated = true;
         };
     }
@@ -219,26 +217,35 @@ where
         let old_children = mem::take(&mut self.children);
 
         // Replace the subchannel map with an empty map.
-        let old_subchannel_child_map = mem::take(&mut self.subchannel_child_map);
+        let old_subchannel_child_map = mem::take(&mut self.subchannel_to_child_idx);
 
-        // Reverse the old subchannel map.
-        let mut old_child_subchannels_map: HashMap<usize, Vec<WeakSubchannel>> = HashMap::new();
+        // Reverse the old subchannel map into a vector indexed by the old child ID.
+        let mut old_child_subchannels: Vec<Vec<WeakSubchannel>> = Vec::new();
+        old_child_subchannels.resize_with(old_children.len(), Vec::new);
 
-        for (subchannel, child_idx) in old_subchannel_child_map {
-            old_child_subchannels_map
-                .entry(child_idx)
-                .or_default()
-                .push(subchannel);
+        for (subchannel, old_idx) in old_subchannel_child_map {
+            old_child_subchannels[old_idx].push(subchannel);
         }
 
         // Build a map of the old children from their IDs for efficient lookups.
-        let old_children = old_children.into_iter().enumerate().map(|(old_idx, e)| {
-            (
-                (e.builder.name(), e.identifier),
-                (e.policy, e.state, old_idx, e.work_scheduler, e.updated),
-            )
-        });
-        let mut old_children: HashMap<(&'static str, T), _> = old_children.collect();
+        // This leverages a Child<usize> to hold all the entries where the
+        // identifier becomes the index within the old self.children vector.
+        let mut old_children: HashMap<(&'static str, T), _> = old_children
+            .into_iter()
+            .enumerate()
+            .map(|(old_idx, e)| {
+                (
+                    (e.builder.name(), e.identifier),
+                    Child {
+                        identifier: old_idx,
+                        policy: e.policy,
+                        builder: e.builder,
+                        state: e.state,
+                        work_scheduler: e.work_scheduler,
+                    },
+                )
+            })
+            .collect();
 
         // Split the child updates into the IDs and builders, and the
         // ResolverUpdates/LbConfigs.
@@ -251,26 +258,21 @@ where
         // subchannel map.
         for (new_idx, (identifier, builder)) in ids_builders.into_iter().enumerate() {
             let k = (builder.name(), identifier);
-            if let Some((policy, state, old_idx, work_scheduler, updated)) = old_children.remove(&k)
-            {
-                for subchannel in old_child_subchannels_map
-                    .remove(&old_idx)
-                    .into_iter()
-                    .flatten()
-                {
-                    self.subchannel_child_map.insert(subchannel, new_idx);
+            if let Some(old_child) = old_children.remove(&k) {
+                let old_idx = old_child.identifier;
+                for subchannel in mem::take(&mut old_child_subchannels[old_idx]) {
+                    self.subchannel_to_child_idx.insert(subchannel, new_idx);
                 }
                 if old_pending_work.contains(&old_idx) {
                     pending_work.insert(new_idx);
                 }
-                *work_scheduler.idx.lock().unwrap() = Some(new_idx);
+                *old_child.work_scheduler.idx.lock().unwrap() = Some(new_idx);
                 self.children.push(Child {
                     builder,
                     identifier: k.1,
-                    state,
-                    policy,
-                    work_scheduler,
-                    updated,
+                    state: old_child.state,
+                    policy: old_child.policy,
+                    work_scheduler: old_child.work_scheduler,
                 });
             } else {
                 let work_scheduler = Arc::new(ChildWorkScheduler {
@@ -288,14 +290,13 @@ where
                     state: LbState::initial(),
                     policy,
                     work_scheduler,
-                    updated: false,
                 });
             };
         }
 
         // Invalidate all deleted children's work_schedulers.
-        for (_, (_, _, _, work_scheduler, _)) in old_children {
-            *work_scheduler.idx.lock().unwrap() = None;
+        for (_, old_child) in old_children {
+            *old_child.work_scheduler.idx.lock().unwrap() = None;
         }
 
         // Release the pending_work mutex before calling into the children to
@@ -333,7 +334,7 @@ where
     ) {
         // Determine which child created this subchannel.
         let child_idx = *self
-            .subchannel_child_map
+            .subchannel_to_child_idx
             .get(&WeakSubchannel::new(&subchannel))
             .unwrap();
         let policy = &mut self.children[child_idx].policy;

grpc/src/client/load_balancing/graceful_switch.rs

@@ -2,8 +2,8 @@ use crate::client::load_balancing::child_manager::{
     self, ChildManager, ChildUpdate, ResolverUpdateSharder,
 };
 use crate::client::load_balancing::{
-    ChannelController, LbConfig, LbPolicy, LbPolicyBuilder, ParsedJsonLbConfig, Subchannel,
-    SubchannelState, WorkScheduler, GLOBAL_LB_REGISTRY,
+    ChannelController, LbConfig, LbPolicy, LbPolicyBuilder, LbState, ParsedJsonLbConfig,
+    Subchannel, SubchannelState, WorkScheduler, GLOBAL_LB_REGISTRY,
 };
 use crate::client::name_resolution::ResolverUpdate;
 use crate::client::ConnectivityState;
@@ -103,6 +103,7 @@ impl UpdateSharder {
 #[derive(Debug)]
 pub(crate) struct GracefulSwitchPolicy {
     child_manager: ChildManager<(), UpdateSharder>, // Child ID is the name of the child policy.
+    last_update: LbState, // Saves the last output LbState to determine if an update is needed.
 }
 
 impl LbPolicy for GracefulSwitchPolicy {
@@ -115,7 +116,7 @@ impl LbPolicy for GracefulSwitchPolicy {
         let res = self
             .child_manager
             .resolver_update(update, config, channel_controller)?;
-        self.maybe_swap(channel_controller);
+        self.update_picker(channel_controller);
         Ok(())
     }
 
@@ -127,17 +128,17 @@ impl LbPolicy for GracefulSwitchPolicy {
     ) {
         self.child_manager
             .subchannel_update(subchannel, state, channel_controller);
-        self.maybe_swap(channel_controller);
+        self.update_picker(channel_controller);
     }
 
     fn work(&mut self, channel_controller: &mut dyn ChannelController) {
         self.child_manager.work(channel_controller);
-        self.maybe_swap(channel_controller);
+        self.update_picker(channel_controller);
     }
 
     fn exit_idle(&mut self, channel_controller: &mut dyn ChannelController) {
         self.child_manager.exit_idle(channel_controller);
-        self.maybe_swap(channel_controller);
+        self.update_picker(channel_controller);
     }
 }
 
@@ -152,6 +153,7 @@ impl GracefulSwitchPolicy {
     pub fn new(runtime: Arc<dyn Runtime>, work_scheduler: Arc<dyn WorkScheduler>) -> Self {
         GracefulSwitchPolicy {
             child_manager: ChildManager::new(UpdateSharder::new(), runtime, work_scheduler),
+            last_update: LbState::initial(),
         }
     }
 
@@ -191,10 +193,27 @@ impl GracefulSwitchPolicy {
         Err("no supported policies found in config".into())
     }
 
-    fn maybe_swap(&mut self, channel_controller: &mut dyn ChannelController) {
-        if !self.child_manager.child_updated() {
+    fn update_picker(&mut self, channel_controller: &mut dyn ChannelController) {
+        let Some(update) = self.maybe_swap(channel_controller) else {
+            return;
+        };
+        if self.last_update.connectivity_state == update.connectivity_state
+            && std::ptr::addr_eq(
+                Arc::as_ptr(&self.last_update.picker),
+                Arc::as_ptr(&update.picker),
+            )
+        {
             return;
         }
+        channel_controller.update_picker(update.clone());
+        self.last_update = update;
+    }
+
+    // Determines the appropriate state to output
+    fn maybe_swap(&mut self, channel_controller: &mut dyn ChannelController) -> Option<LbState> {
+        if !self.child_manager.child_updated() {
+            return None;
+        }
 
         let active_name = self
             .child_manager
@@ -203,36 +222,33 @@ impl GracefulSwitchPolicy {
             .as_ref()
             .unwrap()
             .name();
-        let mut could_switch = false;
-        let mut active_state_if_updated = None;
-        let mut pending_builder = None;
-        let mut pending_state = None;
+
+        let mut active_child = None;
+        let mut pending_child = None;
         for child in self.child_manager.children() {
             if child.builder.name() == active_name {
-                could_switch |= child.state.connectivity_state != ConnectivityState::Ready;
-                if child.updated {
-                    active_state_if_updated = Some(child.state.clone());
-                }
+                active_child = Some(child);
             } else {
-                pending_builder = Some(child.builder.clone());
-                pending_state = Some(child.state.clone());
-                could_switch |= child.state.connectivity_state != ConnectivityState::Connecting;
+                pending_child = Some(child);
             }
         }
-        if could_switch && pending_builder.is_some() {
-            self.child_manager
-                .resolver_update(
-                    ResolverUpdate::default(),
-                    Some(&LbConfig::new(GracefulSwitchLbConfig::Swap(
-                        pending_builder.unwrap(),
-                    ))),
-                    channel_controller,
-                )
-                .expect("resolver_update with an empty update should not fail");
-            channel_controller.update_picker(pending_state.unwrap().clone());
-        } else if active_state_if_updated.is_some() {
-            channel_controller.update_picker(active_state_if_updated.unwrap());
+        let active_child = active_child.expect("There should always be an active child policy");
+        let Some(pending_child) = pending_child else {
+            return Some(active_child.state.clone());
+        };
+
+        if active_child.state.connectivity_state == ConnectivityState::Ready
+            && pending_child.state.connectivity_state == ConnectivityState::Connecting
+        {
+            return Some(active_child.state.clone());
         }
+
+        let config = &LbConfig::new(GracefulSwitchLbConfig::Swap(pending_child.builder.clone()));
+        let state = pending_child.state.clone();
+        self.child_manager
+            .resolver_update(ResolverUpdate::default(), Some(config), channel_controller)
+            .expect("resolver_update with an empty update should not fail");
+        return Some(state);
     }
 }
 
@@ -809,11 +825,8 @@ mod test {
             tcc.as_mut(),
             ConnectivityState::Connecting,
         );
-        verify_correct_picker_from_policy(
-            &mut rx_events,
-            "stub-gracefulswitch_current_leaving_ready-one",
-        )
-        .await;
+        // This should not produce an update.
+        assert_channel_empty(&mut rx_events).await;
         move_subchannel_to_state(
             &mut *graceful_switch,
             current_subchannel,

grpc/src/client/load_balancing/registry.rs

@@ -37,4 +37,5 @@ impl Default for LbPolicyRegistry {
 
 /// The registry used if a local registry is not provided to a channel or if it
 /// does not exist in the local registry.
-pub(crate) static GLOBAL_LB_REGISTRY: LazyLock<LbPolicyRegistry> = LazyLock::new(LbPolicyRegistry::new);
+pub(crate) static GLOBAL_LB_REGISTRY: LazyLock<LbPolicyRegistry> =
+    LazyLock::new(LbPolicyRegistry::new);

grpc/src/client/mod.rs

@@ -45,7 +45,7 @@ pub(crate) mod transport;
 ///
 /// Channels may re-enter the Idle state if they are unused for longer than
 /// their configured idleness timeout.
-#[derive(Copy, Clone, PartialEq, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum ConnectivityState {
     Idle,
     Connecting,