grpc 1.51.0 → 1.52.0.pre2

data/src/core/ext/filters/client_channel/lb_policy/rls/rls.cc

@@ -34,6 +34,7 @@
 #include <random>
 #include <set>
 #include <string>
+#include <type_traits>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -55,9 +56,8 @@
 #include <grpc/byte_buffer_reader.h>
 #include <grpc/event_engine/event_engine.h>
 #include <grpc/grpc.h>
-#include <grpc/impl/codegen/connectivity_state.h>
-#include <grpc/impl/codegen/grpc_types.h>
-#include <grpc/impl/codegen/propagation_bits.h>
+#include <grpc/impl/connectivity_state.h>
+#include <grpc/impl/propagation_bits.h>
 #include <grpc/slice.h>
 #include <grpc/status.h>
 #include <grpc/support/log.h>
@@ -82,7 +82,6 @@
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/iomgr/exec_ctx.h"
 #include "src/core/lib/iomgr/pollset_set.h"
-#include "src/core/lib/iomgr/timer.h"
 #include "src/core/lib/json/json.h"
 #include "src/core/lib/json/json_args.h"
 #include "src/core/lib/json/json_object_loader.h"
@@ -110,6 +109,8 @@ TraceFlag grpc_lb_rls_trace(false, "rls_lb");
 
 namespace {
 
+using ::grpc_event_engine::experimental::EventEngine;
+
 constexpr absl::string_view kRls = "rls_experimental";
 const char kGrpc[] = "grpc";
 const char* kRlsRequestPath = "/grpc.lookup.v1.RouteLookupService/RouteLookup";
@@ -333,7 +334,7 @@ class RlsLb : public LoadBalancingPolicy {
           ServerAddress address, const ChannelArgs& args) override;
       void UpdateState(grpc_connectivity_state state,
                        const absl::Status& status,
-                       std::unique_ptr<SubchannelPicker> picker) override;
+                       RefCountedPtr<SubchannelPicker> picker) override;
       void RequestReresolution() override;
       absl::string_view GetAuthority() override;
       grpc_event_engine::experimental::EventEngine* GetEventEngine() override;
@@ -354,7 +355,7 @@ class RlsLb : public LoadBalancingPolicy {
 
     grpc_connectivity_state connectivity_state_ ABSL_GUARDED_BY(&RlsLb::mu_) =
         GRPC_CHANNEL_IDLE;
-    std::unique_ptr<LoadBalancingPolicy::SubchannelPicker> picker_
+    RefCountedPtr<LoadBalancingPolicy::SubchannelPicker> picker_
         ABSL_GUARDED_BY(&RlsLb::mu_);
   };
 
@@ -460,12 +461,11 @@ class RlsLb : public LoadBalancingPolicy {
         void Orphan() override ABSL_NO_THREAD_SAFETY_ANALYSIS;
 
        private:
-        static void OnBackoffTimer(void* args, grpc_error_handle error);
+        void OnBackoffTimerLocked();
 
         RefCountedPtr<Entry> entry_;
-        bool armed_ ABSL_GUARDED_BY(&RlsLb::mu_) = true;
-        grpc_timer backoff_timer_;
-        grpc_closure backoff_timer_callback_;
+        absl::optional<EventEngine::TaskHandle> backoff_timer_task_handle_
+            ABSL_GUARDED_BY(&RlsLb::mu_);
       };
 
       RefCountedPtr<RlsLb> lb_policy_;
@@ -521,7 +521,10 @@ class RlsLb : public LoadBalancingPolicy {
     void Shutdown() ABSL_EXCLUSIVE_LOCKS_REQUIRED(&RlsLb::mu_);
 
    private:
-    static void OnCleanupTimer(void* arg, grpc_error_handle error);
+    // Shared logic for starting the cleanup timer
+    void StartCleanupTimer() ABSL_EXCLUSIVE_LOCKS_REQUIRED(&RlsLb::mu_);
+
+    void OnCleanupTimer();
 
     // Returns the entry size for a given key.
     static size_t EntrySizeForKey(const RequestKey& key);
@@ -539,8 +542,7 @@ class RlsLb : public LoadBalancingPolicy {
     std::list<RequestKey> lru_list_ ABSL_GUARDED_BY(&RlsLb::mu_);
     std::unordered_map<RequestKey, OrphanablePtr<Entry>, absl::Hash<RequestKey>>
         map_ ABSL_GUARDED_BY(&RlsLb::mu_);
-    grpc_timer cleanup_timer_;
-    grpc_closure timer_callback_;
+    absl::optional<EventEngine::TaskHandle> cleanup_timer_handle_;
   };
 
   // Channel for communicating with the RLS server.
@@ -730,7 +732,7 @@ RlsLb::ChildPolicyWrapper::ChildPolicyWrapper(RefCountedPtr<RlsLb> lb_policy,
                                                      : nullptr),
       lb_policy_(lb_policy),
       target_(std::move(target)),
-      picker_(std::make_unique<QueuePicker>(std::move(lb_policy))) {
+      picker_(MakeRefCounted<QueuePicker>(std::move(lb_policy))) {
   lb_policy_->child_policy_map_.emplace(target_, this);
 }
 
@@ -811,7 +813,7 @@ void RlsLb::ChildPolicyWrapper::StartUpdate() {
               config.status().ToString().c_str());
     }
     pending_config_.reset();
-    picker_ = std::make_unique<TransientFailurePicker>(
+    picker_ = MakeRefCounted<TransientFailurePicker>(
         absl::UnavailableError(config.status().message()));
     child_policy_.reset();
   } else {
@@ -876,7 +878,7 @@ RlsLb::ChildPolicyWrapper::ChildPolicyHelper::CreateSubchannel(
 
 void RlsLb::ChildPolicyWrapper::ChildPolicyHelper::UpdateState(
     grpc_connectivity_state state, const absl::Status& status,
-    std::unique_ptr<SubchannelPicker> picker) {
+    RefCountedPtr<SubchannelPicker> picker) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
     gpr_log(GPR_INFO,
             "[rlslb %p] ChildPolicyWrapper=%p [%s] ChildPolicyHelper=%p: "
@@ -1116,46 +1118,50 @@ LoadBalancingPolicy::PickResult RlsLb::Picker::Pick(PickArgs args) {
 RlsLb::Cache::Entry::BackoffTimer::BackoffTimer(RefCountedPtr<Entry> entry,
                                                 Timestamp backoff_time)
     : entry_(std::move(entry)) {
-  GRPC_CLOSURE_INIT(&backoff_timer_callback_, OnBackoffTimer, this, nullptr);
-  Ref(DEBUG_LOCATION, "BackoffTimer").release();
-  grpc_timer_init(&backoff_timer_, backoff_time, &backoff_timer_callback_);
+  backoff_timer_task_handle_ =
+      entry_->lb_policy_->channel_control_helper()->GetEventEngine()->RunAfter(
+          backoff_time - Timestamp::Now(),
+          [self = Ref(DEBUG_LOCATION, "BackoffTimer")]() mutable {
+            ApplicationCallbackExecCtx callback_exec_ctx;
+            ExecCtx exec_ctx;
+            auto self_ptr = self.get();
+            self_ptr->entry_->lb_policy_->work_serializer()->Run(
+                [self = std::move(self)]() { self->OnBackoffTimerLocked(); },
+                DEBUG_LOCATION);
+          });
 }
 
 void RlsLb::Cache::Entry::BackoffTimer::Orphan() {
-  if (armed_) {
-    armed_ = false;
-    grpc_timer_cancel(&backoff_timer_);
+  if (backoff_timer_task_handle_.has_value() &&
+      entry_->lb_policy_->channel_control_helper()->GetEventEngine()->Cancel(
+          *backoff_timer_task_handle_)) {
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
+      gpr_log(GPR_INFO, "[rlslb %p] cache entry=%p %s, backoff timer canceled",
+              entry_->lb_policy_.get(), entry_.get(),
+              entry_->is_shutdown_ ? "(shut down)"
+                                   : entry_->lru_iterator_->ToString().c_str());
+    }
   }
+  backoff_timer_task_handle_.reset();
   Unref(DEBUG_LOCATION, "Orphan");
 }
 
-void RlsLb::Cache::Entry::BackoffTimer::OnBackoffTimer(
-    void* arg, grpc_error_handle /*error*/) {
-  auto* self = static_cast<BackoffTimer*>(arg);
-  self->entry_->lb_policy_->work_serializer()->Run(
-      [self]() {
-        RefCountedPtr<BackoffTimer> backoff_timer(self);
-        {
-          MutexLock lock(&self->entry_->lb_policy_->mu_);
-          if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
-            gpr_log(GPR_INFO,
-                    "[rlslb %p] cache entry=%p %s, armed_=%d: "
-                    "backoff timer fired",
-                    self->entry_->lb_policy_.get(), self->entry_.get(),
-                    self->entry_->is_shutdown_
-                        ? "(shut down)"
-                        : self->entry_->lru_iterator_->ToString().c_str(),
-                    self->armed_);
-          }
-          bool cancelled = !self->armed_;
-          self->armed_ = false;
-          if (cancelled) return;
-        }
-        // The pick was in backoff state and there could be a pick queued if
-        // wait_for_ready is true. We'll update the picker for that case.
-        self->entry_->lb_policy_->UpdatePickerLocked();
-      },
-      DEBUG_LOCATION);
+void RlsLb::Cache::Entry::BackoffTimer::OnBackoffTimerLocked() {
+  {
+    MutexLock lock(&entry_->lb_policy_->mu_);
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
+      gpr_log(GPR_INFO, "[rlslb %p] cache entry=%p %s, backoff timer fired",
+              entry_->lb_policy_.get(), entry_.get(),
+              entry_->is_shutdown_ ? "(shut down)"
+                                   : entry_->lru_iterator_->ToString().c_str());
+    }
+    // Skip the update if Orphaned
+    if (!backoff_timer_task_handle_.has_value()) return;
+    backoff_timer_task_handle_.reset();
+  }
+  // The pick was in backoff state and there could be a pick queued if
+  // wait_for_ready is true. We'll update the picker for that case.
+  entry_->lb_policy_->UpdatePickerLocked();
 }
 
 //
@@ -1361,11 +1367,7 @@ RlsLb::Cache::Entry::OnRlsResponseLocked(
 //
 
 RlsLb::Cache::Cache(RlsLb* lb_policy) : lb_policy_(lb_policy) {
-  Timestamp now = Timestamp::Now();
-  lb_policy_->Ref(DEBUG_LOCATION, "CacheCleanupTimer").release();
-  GRPC_CLOSURE_INIT(&timer_callback_, OnCleanupTimer, this, nullptr);
-  grpc_timer_init(&cleanup_timer_, now + kCacheCleanupTimerInterval,
-                  &timer_callback_);
+  StartCleanupTimer();
 }
 
 RlsLb::Cache::Entry* RlsLb::Cache::Find(const RequestKey& key) {
@@ -1419,37 +1421,49 @@ void RlsLb::Cache::ResetAllBackoff() {
 void RlsLb::Cache::Shutdown() {
   map_.clear();
   lru_list_.clear();
-  grpc_timer_cancel(&cleanup_timer_);
+  if (cleanup_timer_handle_.has_value() &&
+      lb_policy_->channel_control_helper()->GetEventEngine()->Cancel(
+          *cleanup_timer_handle_)) {
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
+      gpr_log(GPR_INFO, "[rlslb %p] cache cleanup timer canceled", lb_policy_);
+    }
+  }
+  cleanup_timer_handle_.reset();
 }
 
-void RlsLb::Cache::OnCleanupTimer(void* arg, grpc_error_handle error) {
-  Cache* cache = static_cast<Cache*>(arg);
-  cache->lb_policy_->work_serializer()->Run(
-      [cache, error]() {
-        RefCountedPtr<RlsLb> lb_policy(cache->lb_policy_);
-        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
-          gpr_log(GPR_INFO, "[rlslb %p] cache cleanup timer fired (%s)",
-                  cache->lb_policy_, StatusToString(error).c_str());
-        }
-        if (error == absl::CancelledError()) return;
-        MutexLock lock(&lb_policy->mu_);
-        if (lb_policy->is_shutdown_) return;
-        for (auto it = cache->map_.begin(); it != cache->map_.end();) {
-          if (GPR_UNLIKELY(it->second->ShouldRemove() &&
-                           it->second->CanEvict())) {
-            cache->size_ -= it->second->Size();
-            it = cache->map_.erase(it);
-          } else {
-            ++it;
-          }
-        }
-        Timestamp now = Timestamp::Now();
-        lb_policy.release();
-        grpc_timer_init(&cache->cleanup_timer_,
-                        now + kCacheCleanupTimerInterval,
-                        &cache->timer_callback_);
-      },
-      DEBUG_LOCATION);
+void RlsLb::Cache::StartCleanupTimer() {
+  cleanup_timer_handle_ =
+      lb_policy_->channel_control_helper()->GetEventEngine()->RunAfter(
+          kCacheCleanupTimerInterval,
+          [this, lb_policy = lb_policy_->Ref(DEBUG_LOCATION,
+                                             "CacheCleanupTimer")]() mutable {
+            ApplicationCallbackExecCtx callback_exec_ctx;
+            ExecCtx exec_ctx;
+            lb_policy_->work_serializer()->Run(
+                [this, lb_policy = std::move(lb_policy)]() {
+                  // The lb_policy ref is held until the callback completes
+                  OnCleanupTimer();
+                },
+                DEBUG_LOCATION);
+          });
+}
+
+void RlsLb::Cache::OnCleanupTimer() {
+  if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_rls_trace)) {
+    gpr_log(GPR_INFO, "[rlslb %p] cache cleanup timer fired", lb_policy_);
+  }
+  MutexLock lock(&lb_policy_->mu_);
+  if (!cleanup_timer_handle_.has_value()) return;
+  if (lb_policy_->is_shutdown_) return;
+  for (auto it = map_.begin(); it != map_.end();) {
+    if (GPR_UNLIKELY(it->second->ShouldRemove() && it->second->CanEvict())) {
+      size_ -= it->second->Size();
+      it = map_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  StartCleanupTimer();
 }
 
 size_t RlsLb::Cache::EntrySizeForKey(const RequestKey& key) {
@@ -2147,7 +2161,7 @@ void RlsLb::UpdatePickerLocked() {
     status = absl::UnavailableError("no children available");
   }
   channel_control_helper()->UpdateState(
-      state, status, std::make_unique<Picker>(Ref(DEBUG_LOCATION, "Picker")));
+      state, status, MakeRefCounted<Picker>(Ref(DEBUG_LOCATION, "Picker")));
 }
 
 //

data/src/core/ext/filters/client_channel/lb_policy/round_robin/round_robin.cc

@@ -31,7 +31,7 @@
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 
-#include <grpc/impl/codegen/connectivity_state.h>
+#include <grpc/impl/connectivity_state.h>
 #include <grpc/support/log.h>
 
 #include "src/core/ext/filters/client_channel/lb_policy/subchannel_list.h"
@@ -305,7 +305,7 @@ absl::Status RoundRobin::UpdateLocked(UpdateArgs args) {
                             : args.addresses.status();
     channel_control_helper()->UpdateState(
         GRPC_CHANNEL_TRANSIENT_FAILURE, status,
-        std::make_unique<TransientFailurePicker>(status));
+        MakeRefCounted<TransientFailurePicker>(status));
     return status;
   }
   // Otherwise, if this is the initial update, immediately promote it to
@@ -314,7 +314,7 @@ absl::Status RoundRobin::UpdateLocked(UpdateArgs args) {
     subchannel_list_ = std::move(latest_pending_subchannel_list_);
     channel_control_helper()->UpdateState(
         GRPC_CHANNEL_CONNECTING, absl::Status(),
-        std::make_unique<QueuePicker>(Ref(DEBUG_LOCATION, "QueuePicker")));
+        MakeRefCounted<QueuePicker>(Ref(DEBUG_LOCATION, "QueuePicker")));
   }
   return absl::OkStatus();
 }
@@ -355,12 +355,14 @@ void RoundRobin::RoundRobinSubchannelList::
   // If this is latest_pending_subchannel_list_, then swap it into
   // subchannel_list_ in the following cases:
   // - subchannel_list_ has no READY subchannels.
-  // - This list has at least one READY subchannel.
+  // - This list has at least one READY subchannel and we have seen the
+  //   initial connectivity state notification for all subchannels.
   // - All of the subchannels in this list are in TRANSIENT_FAILURE.
   //   (This may cause the channel to go from READY to TRANSIENT_FAILURE,
   //   but we're doing what the control plane told us to do.)
   if (p->latest_pending_subchannel_list_.get() == this &&
-      (p->subchannel_list_->num_ready_ == 0 || num_ready_ > 0 ||
+      (p->subchannel_list_->num_ready_ == 0 ||
+       (num_ready_ > 0 && AllSubchannelsSeenInitialState()) ||
        num_transient_failure_ == num_subchannels())) {
     if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_round_robin_trace)) {
       const std::string old_counters_string =
@@ -386,7 +388,7 @@ void RoundRobin::RoundRobinSubchannelList::
               this);
     }
     p->channel_control_helper()->UpdateState(GRPC_CHANNEL_READY, absl::Status(),
-                                             std::make_unique<Picker>(p, this));
+                                             MakeRefCounted<Picker>(p, this));
   } else if (num_connecting_ > 0) {
     if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_round_robin_trace)) {
       gpr_log(GPR_INFO, "[RR %p] reporting CONNECTING with subchannel list %p",
@@ -394,7 +396,7 @@ void RoundRobin::RoundRobinSubchannelList::
     }
     p->channel_control_helper()->UpdateState(
         GRPC_CHANNEL_CONNECTING, absl::Status(),
-        std::make_unique<QueuePicker>(p->Ref(DEBUG_LOCATION, "QueuePicker")));
+        MakeRefCounted<QueuePicker>(p->Ref(DEBUG_LOCATION, "QueuePicker")));
   } else if (num_transient_failure_ == num_subchannels()) {
     if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_round_robin_trace)) {
       gpr_log(GPR_INFO,
@@ -408,7 +410,7 @@ void RoundRobin::RoundRobinSubchannelList::
     }
     p->channel_control_helper()->UpdateState(
         GRPC_CHANNEL_TRANSIENT_FAILURE, last_failure_,
-        std::make_unique<TransientFailurePicker>(last_failure_));
+        MakeRefCounted<TransientFailurePicker>(last_failure_));
   }
 }
 

data/src/core/ext/filters/client_channel/lb_policy/subchannel_list.h

@@ -30,7 +30,7 @@
 #include "absl/status/status.h"
 #include "absl/types/optional.h"
 
-#include <grpc/impl/codegen/connectivity_state.h>
+#include <grpc/impl/connectivity_state.h>
 #include <grpc/support/log.h>
 
 #include "src/core/lib/channel/channel_args.h"
@@ -47,25 +47,25 @@
 // Code for maintaining a list of subchannels within an LB policy.
 //
 // To use this, callers must create their own subclasses, like so:
-/*
+//
 
-class MySubchannelList;  // Forward declaration.
+// class MySubchannelList;  // Forward declaration.
 
-class MySubchannelData
-    : public SubchannelData<MySubchannelList, MySubchannelData> {
- public:
-  void ProcessConnectivityChangeLocked(
-      absl::optional<grpc_connectivity_state> old_state,
-      grpc_connectivity_state new_state) override {
-    // ...code to handle connectivity changes...
-  }
-};
+// class MySubchannelData
+//   : public SubchannelData<MySubchannelList, MySubchannelData> {
+// public:
+// void ProcessConnectivityChangeLocked(
+//     absl::optional<grpc_connectivity_state> old_state,
+//     grpc_connectivity_state new_state) override {
+//   // ...code to handle connectivity changes...
+// }
+// };
 
-class MySubchannelList
-    : public SubchannelList<MySubchannelList, MySubchannelData> {
-};
+// class MySubchannelList
+//   : public SubchannelList<MySubchannelList, MySubchannelData> {
+// };
 
-*/
+//
 // All methods will be called from within the client_channel work serializer.
 
 namespace grpc_core {
@@ -200,6 +200,10 @@ class SubchannelList : public DualRefCounted<SubchannelListType> {
   // Resets connection backoff of all subchannels.
   void ResetBackoffLocked();
 
+  // Returns true if all subchannels have seen their initial
+  // connectivity state notifications.
+  bool AllSubchannelsSeenInitialState();
+
   void Orphan() override;
 
  protected:
@@ -324,6 +328,7 @@ void SubchannelData<SubchannelListType,
   pending_watcher_ =
       new Watcher(this, subchannel_list()->WeakRef(DEBUG_LOCATION, "Watcher"));
   subchannel_->WatchConnectivityState(
+      // NOLINTNEXTLINE(google-readability-casting)
       std::unique_ptr<SubchannelInterface::ConnectivityStateWatcherInterface>(
           pending_watcher_));
 }
@@ -434,6 +439,15 @@ void SubchannelList<SubchannelListType,
   }
 }
 
+template <typename SubchannelListType, typename SubchannelDataType>
+bool SubchannelList<SubchannelListType,
+                    SubchannelDataType>::AllSubchannelsSeenInitialState() {
+  for (size_t i = 0; i < num_subchannels(); ++i) {
+    if (!subchannel(i)->connectivity_state().has_value()) return false;
+  }
+  return true;
+}
+
 }  // namespace grpc_core
 
 #endif  // GRPC_CORE_EXT_FILTERS_CLIENT_CHANNEL_LB_POLICY_SUBCHANNEL_LIST_H

data/src/core/ext/filters/client_channel/lb_policy/weighted_target/weighted_target.cc

@@ -35,7 +35,7 @@
 #include "absl/types/optional.h"
 
 #include <grpc/event_engine/event_engine.h>
-#include <grpc/impl/codegen/connectivity_state.h>
+#include <grpc/impl/connectivity_state.h>
 #include <grpc/support/log.h>
 
 #include "src/core/ext/filters/client_channel/lb_policy/address_filtering.h"
@@ -45,7 +45,6 @@
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/gprpp/debug_location.h"
 #include "src/core/lib/gprpp/orphanable.h"
-#include "src/core/lib/gprpp/ref_counted.h"
 #include "src/core/lib/gprpp/ref_counted_ptr.h"
 #include "src/core/lib/gprpp/time.h"
 #include "src/core/lib/gprpp/validation_errors.h"
@@ -121,17 +120,6 @@ class WeightedTargetLb : public LoadBalancingPolicy {
   void ResetBackoffLocked() override;
 
  private:
-  // A simple wrapper for ref-counting a picker from the child policy.
-  class ChildPickerWrapper : public RefCounted<ChildPickerWrapper> {
-   public:
-    explicit ChildPickerWrapper(std::unique_ptr<SubchannelPicker> picker)
-        : picker_(std::move(picker)) {}
-    PickResult Pick(PickArgs args) { return picker_->Pick(args); }
-
-   private:
-    std::unique_ptr<SubchannelPicker> picker_;
-  };
-
   // Picks a child using stateless WRR and then delegates to that
   // child's picker.
   class WeightedPicker : public SubchannelPicker {
@@ -141,7 +129,7 @@ class WeightedTargetLb : public LoadBalancingPolicy {
     // range proportional to the child's weight. The start of the range
     // is the previous value in the vector and is 0 for the first element.
     using PickerList =
-        std::vector<std::pair<uint64_t, RefCountedPtr<ChildPickerWrapper>>>;
+        std::vector<std::pair<uint64_t, RefCountedPtr<SubchannelPicker>>>;
 
     explicit WeightedPicker(PickerList pickers)
         : pickers_(std::move(pickers)) {}
@@ -173,9 +161,7 @@ class WeightedTargetLb : public LoadBalancingPolicy {
     grpc_connectivity_state connectivity_state() const {
       return connectivity_state_;
     }
-    RefCountedPtr<ChildPickerWrapper> picker_wrapper() const {
-      return picker_wrapper_;
-    }
+    RefCountedPtr<SubchannelPicker> picker() const { return picker_; }
 
    private:
     class Helper : public ChannelControlHelper {
@@ -189,7 +175,7 @@ class WeightedTargetLb : public LoadBalancingPolicy {
           ServerAddress address, const ChannelArgs& args) override;
       void UpdateState(grpc_connectivity_state state,
                        const absl::Status& status,
-                       std::unique_ptr<SubchannelPicker> picker) override;
+                       RefCountedPtr<SubchannelPicker> picker) override;
       void RequestReresolution() override;
       absl::string_view GetAuthority() override;
       grpc_event_engine::experimental::EventEngine* GetEventEngine() override;
@@ -220,7 +206,7 @@ class WeightedTargetLb : public LoadBalancingPolicy {
 
     void OnConnectivityStateUpdateLocked(
         grpc_connectivity_state state, const absl::Status& status,
-        std::unique_ptr<SubchannelPicker> picker);
+        RefCountedPtr<SubchannelPicker> picker);
 
     // The owning LB policy.
     RefCountedPtr<WeightedTargetLb> weighted_target_policy_;
@@ -231,7 +217,7 @@ class WeightedTargetLb : public LoadBalancingPolicy {
 
     OrphanablePtr<LoadBalancingPolicy> child_policy_;
 
-    RefCountedPtr<ChildPickerWrapper> picker_wrapper_;
+    RefCountedPtr<SubchannelPicker> picker_;
     grpc_connectivity_state connectivity_state_ = GRPC_CHANNEL_CONNECTING;
 
     OrphanablePtr<DelayedRemovalTimer> delayed_removal_timer_;
@@ -364,7 +350,7 @@ absl::Status WeightedTargetLb::UpdateLocked(UpdateArgs args) {
         "no children in weighted_target policy: ", args.resolution_note));
     channel_control_helper()->UpdateState(
         GRPC_CHANNEL_TRANSIENT_FAILURE, status,
-        std::make_unique<TransientFailurePicker>(status));
+        MakeRefCounted<TransientFailurePicker>(status));
     return absl::OkStatus();
   }
   UpdateStateLocked();
@@ -409,18 +395,19 @@ void WeightedTargetLb::UpdateStateLocked() {
     if (config_->target_map().find(child_name) == config_->target_map().end()) {
       continue;
     }
+    auto child_picker = child->picker();
     if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_weighted_target_trace)) {
       gpr_log(GPR_INFO,
               "[weighted_target_lb %p]   child=%s state=%s weight=%u picker=%p",
               this, child_name.c_str(),
               ConnectivityStateName(child->connectivity_state()),
-              child->weight(), child->picker_wrapper().get());
+              child->weight(), child_picker.get());
     }
     switch (child->connectivity_state()) {
       case GRPC_CHANNEL_READY: {
         GPR_ASSERT(child->weight() > 0);
         ready_end += child->weight();
-        ready_picker_list.emplace_back(ready_end, child->picker_wrapper());
+        ready_picker_list.emplace_back(ready_end, std::move(child_picker));
         break;
       }
       case GRPC_CHANNEL_CONNECTING: {
@@ -434,7 +421,7 @@ void WeightedTargetLb::UpdateStateLocked() {
       case GRPC_CHANNEL_TRANSIENT_FAILURE: {
         GPR_ASSERT(child->weight() > 0);
         tf_end += child->weight();
-        tf_picker_list.emplace_back(tf_end, child->picker_wrapper());
+        tf_picker_list.emplace_back(tf_end, std::move(child_picker));
         break;
       }
       default:
@@ -456,19 +443,18 @@ void WeightedTargetLb::UpdateStateLocked() {
     gpr_log(GPR_INFO, "[weighted_target_lb %p] connectivity changed to %s",
             this, ConnectivityStateName(connectivity_state));
   }
-  std::unique_ptr<SubchannelPicker> picker;
+  RefCountedPtr<SubchannelPicker> picker;
   absl::Status status;
   switch (connectivity_state) {
     case GRPC_CHANNEL_READY:
-      picker = std::make_unique<WeightedPicker>(std::move(ready_picker_list));
+      picker = MakeRefCounted<WeightedPicker>(std::move(ready_picker_list));
       break;
     case GRPC_CHANNEL_CONNECTING:
     case GRPC_CHANNEL_IDLE:
-      picker =
-          std::make_unique<QueuePicker>(Ref(DEBUG_LOCATION, "QueuePicker"));
+      picker = MakeRefCounted<QueuePicker>(Ref(DEBUG_LOCATION, "QueuePicker"));
       break;
     default:
-      picker = std::make_unique<WeightedPicker>(std::move(tf_picker_list));
+      picker = MakeRefCounted<WeightedPicker>(std::move(tf_picker_list));
   }
   channel_control_helper()->UpdateState(connectivity_state, status,
                                         std::move(picker));
@@ -555,7 +541,7 @@ void WeightedTargetLb::WeightedChild::Orphan() {
   child_policy_.reset();
   // Drop our ref to the child's picker, in case it's holding a ref to
   // the child.
-  picker_wrapper_.reset();
+  picker_.reset();
   delayed_removal_timer_.reset();
   Unref();
 }
@@ -635,16 +621,16 @@ void WeightedTargetLb::WeightedChild::ResetBackoffLocked() {
 
 void WeightedTargetLb::WeightedChild::OnConnectivityStateUpdateLocked(
     grpc_connectivity_state state, const absl::Status& status,
-    std::unique_ptr<SubchannelPicker> picker) {
+    RefCountedPtr<SubchannelPicker> picker) {
   // Cache the picker in the WeightedChild.
-  picker_wrapper_ = MakeRefCounted<ChildPickerWrapper>(std::move(picker));
+  picker_ = std::move(picker);
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_weighted_target_trace)) {
     gpr_log(GPR_INFO,
             "[weighted_target_lb %p] WeightedChild %p %s: connectivity "
-            "state update: state=%s (%s) picker_wrapper=%p",
+            "state update: state=%s (%s) picker=%p",
             weighted_target_policy_.get(), this, name_.c_str(),
             ConnectivityStateName(state), status.ToString().c_str(),
-            picker_wrapper_.get());
+            picker_.get());
   }
   // If the child reports IDLE, immediately tell it to exit idle.
   if (state == GRPC_CHANNEL_IDLE) child_policy_->ExitIdleLocked();
@@ -688,7 +674,7 @@ WeightedTargetLb::WeightedChild::Helper::CreateSubchannel(
 
 void WeightedTargetLb::WeightedChild::Helper::UpdateState(
     grpc_connectivity_state state, const absl::Status& status,
-    std::unique_ptr<SubchannelPicker> picker) {
+    RefCountedPtr<SubchannelPicker> picker) {
   if (weighted_child_->weighted_target_policy_->shutting_down_) return;
   weighted_child_->OnConnectivityStateUpdateLocked(state, status,
                                                    std::move(picker));