feat: Implement Least Request Load Balancing Policy (gRFC A48)

grpc/src/client/channel.rs

@@ -55,6 +55,7 @@ use crate::client::load_balancing::QueuingPicker;
 use crate::client::load_balancing::WorkData;
 use crate::client::load_balancing::WorkScheduler;
 use crate::client::load_balancing::graceful_switch::GracefulSwitchPolicy;
+use crate::client::load_balancing::least_request;
 use crate::client::load_balancing::pick_first;
 use crate::client::load_balancing::round_robin;
 use crate::client::load_balancing::subchannel::Subchannel;
@@ -171,6 +172,7 @@ impl Channel {
     ) -> Self {
         pick_first::reg();
         round_robin::reg();
+        least_request::reg();
         dns::reg();
         #[cfg(unix)]
         name_resolution::unix::reg();
@@ -380,11 +382,20 @@ impl Invoke for Arc<ActiveChannel> {
                     "channel has been closed",
                 ));
             };
-            let result = &state.picker.pick(&headers);
+            let result = state.picker.pick(&headers);
             match result {
-                PickResult::Pick(pr) => {
+                PickResult::Pick(mut pr) => {
                     if let Some(sc) = pr.subchannel.downcast_ref::<InternalSubchannel>() {
-                        return sc.dyn_invoke(headers, options.clone()).await;
+                        let (tx, rx) = sc.dyn_invoke(headers, options.clone()).await;
+                        let rx = if let Some(on_complete) = pr.on_complete.take() {
+                            Box::new(CallbackRecvStream {
+                                delegate: rx,
+                                on_complete: Some(on_complete),
+                            }) as Box<dyn DynRecvStream>
+                        } else {
+                            rx
+                        };
+                        return (tx, rx);
                     } else {
                         panic!(
                             "picked subchannel is not an implementation provided by the channel"
@@ -395,7 +406,7 @@ impl Invoke for Arc<ActiveChannel> {
                     // Continue and retry the RPC with the next picker.
                 }
                 PickResult::Fail(status) => {
-                    return FailingRecvStream::new_stream_pair(status.clone());
+                    return FailingRecvStream::new_stream_pair(status);
                 }
                 PickResult::Drop(status) => {
                     todo!("dropped pick: {:?}", status);
@@ -411,6 +422,39 @@ impl Drop for ActiveChannel {
     }
 }
 
+struct CallbackRecvStream {
+    delegate: Box<dyn DynRecvStream>,
+    on_complete: Option<Box<dyn Fn() + Send + Sync>>,
+}
+
+#[tonic::async_trait]
+impl DynRecvStream for CallbackRecvStream {
+    async fn dyn_recv(
+        &mut self,
+        msg: &mut dyn crate::core::RecvMessage,
+    ) -> crate::client::ResponseStreamItem {
+        let item = self.delegate.dyn_recv(msg).await;
+        if matches!(
+            item,
+            crate::client::ResponseStreamItem::Trailers(_)
+                | crate::client::ResponseStreamItem::StreamClosed
+        ) {
+            if let Some(cb) = self.on_complete.take() {
+                cb();
+            }
+        }
+        item
+    }
+}
+
+impl Drop for CallbackRecvStream {
+    fn drop(&mut self) {
+        if let Some(cb) = self.on_complete.take() {
+            cb();
+        }
+    }
+}
+
 struct ResolverWorkScheduler {
     wqtx: WorkQueueTx,
 }
@@ -462,13 +506,22 @@ impl ResolverChannelController {
 
 impl name_resolution::ChannelController for ResolverChannelController {
     fn update(&mut self, update: ResolverUpdate) -> Result<(), String> {
-        let json_config = if let Ok(Some(service_config)) = update.service_config.as_ref()
-            && service_config
+        let json_config = if let Ok(Some(service_config)) = update.service_config.as_ref() {
+            if service_config
                 .load_balancing_policy
                 .as_ref()
                 .is_some_and(|p| *p == LbPolicyType::RoundRobin)
-        {
-            json!([{round_robin::POLICY_NAME: {}}])
+            {
+                json!([{round_robin::POLICY_NAME: {}}])
+            } else if service_config
+                .load_balancing_policy
+                .as_ref()
+                .is_some_and(|p| *p == LbPolicyType::LeastRequest)
+            {
+                json!([{least_request::POLICY_NAME: {}}])
+            } else {
+                json!([{pick_first::POLICY_NAME: {"shuffleAddressList": true, "unknown_field": false}}])
+            }
         } else {
             json!([{pick_first::POLICY_NAME: {"shuffleAddressList": true, "unknown_field": false}}])
         };

grpc/src/client/load_balancing/child_manager.rs

@@ -65,10 +65,17 @@ pub(crate) struct Child<T> {
     pub identifier: T,
     pub builder: Arc<DynLbPolicyBuilder>,
     pub state: LbState,
+    pub subchannels: Vec<WeakSubchannel>,
     policy: Box<DynLbPolicy>,
     work_scheduler: Arc<ChildWorkScheduler>,
 }
 
+impl<T> Child<T> {
+    pub fn subchannels(&self) -> impl Iterator<Item = Arc<dyn Subchannel>> + '_ {
+        self.subchannels.iter().filter_map(|weak| weak.upgrade())
+    }
+}
+
 /// A collection of data sent to a child of the ChildManager.
 pub(crate) struct ChildUpdate<'a, T> {
     /// The identifier the ChildManager should use for this child.
@@ -158,6 +165,7 @@ where
         for csc in channel_controller.created_subchannels {
             self.subchannel_to_child_idx
                 .insert((&csc).into(), child_idx);
+            self.children[child_idx].subchannels.push((&csc).into());
         }
         // Update the tracked state if the child produced an update.
         if let Some(state) = channel_controller.picker_update {
@@ -220,6 +228,7 @@ where
                         policy: e.policy,
                         builder: e.builder,
                         state: e.state,
+                        subchannels: e.subchannels,
                         work_scheduler: e.work_scheduler,
                     },
                 )
@@ -237,16 +246,19 @@ where
             if let Some(old_child) = old_children.remove(&k) {
                 let old_idx = old_child.identifier;
                 let new_child_idx = self.children.len();
+                let mut subchannels = Vec::new();
                 for subchannel in mem::take(&mut old_child_subchannels[old_idx]) {
                     self.subchannel_to_child_idx
-                        .insert(subchannel, new_child_idx);
+                        .insert(subchannel.clone(), new_child_idx);
+                    subchannels.push(subchannel);
                 }
                 self.handle_to_child_idx
                     .insert(old_child.work_scheduler.handle.clone(), new_child_idx);
                 self.children.push(Child {
                     builder,
                     identifier: k.1,
                     state: old_child.state,
+                    subchannels,
                     policy: old_child.policy,
                     work_scheduler: old_child.work_scheduler,
                 });
@@ -267,6 +279,7 @@ where
                     builder,
                     identifier: k.1,
                     state: LbState::initial(),
+                    subchannels: Vec::new(),
                     policy,
                     work_scheduler,
                 });