grpc 1.27.0 → 1.28.0.pre2

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

@@ -74,12 +74,12 @@ namespace {
 
 constexpr char kXds[] = "xds_experimental";
 
-class ParsedXdsConfig : public LoadBalancingPolicy::Config {
+class XdsConfig : public LoadBalancingPolicy::Config {
  public:
-  ParsedXdsConfig(RefCountedPtr<LoadBalancingPolicy::Config> child_policy,
-                  RefCountedPtr<LoadBalancingPolicy::Config> fallback_policy,
-                  std::string eds_service_name,
-                  Optional<std::string> lrs_load_reporting_server_name)
+  XdsConfig(RefCountedPtr<LoadBalancingPolicy::Config> child_policy,
+            RefCountedPtr<LoadBalancingPolicy::Config> fallback_policy,
+            std::string eds_service_name,
+            Optional<std::string> lrs_load_reporting_server_name)
       : child_policy_(std::move(child_policy)),
         fallback_policy_(std::move(fallback_policy)),
         eds_service_name_(std::move(eds_service_name)),
@@ -123,31 +123,42 @@ class XdsLb : public LoadBalancingPolicy {
  private:
   class EndpointWatcher;
 
-  // We need this wrapper for the following reasons:
-  // 1. To process per-locality load reporting.
-  // 2. Since pickers are std::unique_ptrs we use this RefCounted wrapper to
-  // control
-  //     references to it by the xds picker and the locality.
-  class EndpointPickerWrapper : public RefCounted<EndpointPickerWrapper> {
+  // A simple wrapper to convert the picker returned from a locality's child
+  // policy as a unique_ptr<> to a RefCountedPtr<>.  This allows it to be
+  // referenced by both the picker and the locality.
+  class RefCountedEndpointPicker : public RefCounted<RefCountedEndpointPicker> {
    public:
-    EndpointPickerWrapper(
-        std::unique_ptr<SubchannelPicker> picker,
-        RefCountedPtr<XdsClientStats::LocalityStats> locality_stats)
+    explicit RefCountedEndpointPicker(std::unique_ptr<SubchannelPicker> picker)
+        : picker_(std::move(picker)) {}
+    PickResult Pick(PickArgs args) { return picker_->Pick(std::move(args)); }
+
+   private:
+    std::unique_ptr<SubchannelPicker> picker_;
+  };
+
+  // A picker that wraps the RefCountedEndpointPicker and performs load
+  // reporting for the locality.
+  class LoadReportingPicker : public RefCounted<LoadReportingPicker> {
+   public:
+    LoadReportingPicker(RefCountedPtr<RefCountedEndpointPicker> picker,
+                        RefCountedPtr<XdsClusterLocalityStats> locality_stats)
         : picker_(std::move(picker)),
-          locality_stats_(std::move(locality_stats)) {
-      locality_stats_->RefByPicker();
-    }
-    ~EndpointPickerWrapper() { locality_stats_->UnrefByPicker(); }
+          locality_stats_(std::move(locality_stats)) {}
 
     PickResult Pick(PickArgs args);
 
+    RefCountedEndpointPicker* picker() const { return picker_.get(); }
+    XdsClusterLocalityStats* locality_stats() const {
+      return locality_stats_.get();
+    }
+
    private:
-    std::unique_ptr<SubchannelPicker> picker_;
-    RefCountedPtr<XdsClientStats::LocalityStats> locality_stats_;
+    RefCountedPtr<RefCountedEndpointPicker> picker_;
+    RefCountedPtr<XdsClusterLocalityStats> locality_stats_;
   };
 
-  // The picker will use a stateless weighting algorithm to pick the locality to
-  // use for each request.
+  // A picker that uses a stateless weighting algorithm to pick the locality
+  // to use for each request.
   class LocalityPicker : public SubchannelPicker {
    public:
     // Maintains a weighted list of pickers from each locality that is in ready
@@ -155,14 +166,12 @@ class XdsLb : public LoadBalancingPolicy {
     // proportional to the locality's weight. The start of the range is the
     // previous value in the vector and is 0 for the first element.
     using PickerList =
-        InlinedVector<std::pair<uint32_t, RefCountedPtr<EndpointPickerWrapper>>,
+        InlinedVector<std::pair<uint32_t, RefCountedPtr<LoadReportingPicker>>,
                       1>;
-    LocalityPicker(RefCountedPtr<XdsLb> xds_policy, PickerList pickers)
-        : xds_policy_(std::move(xds_policy)),
-          pickers_(std::move(pickers)),
-          drop_config_(xds_policy_->drop_config_) {}
-
-    ~LocalityPicker() { xds_policy_.reset(DEBUG_LOCATION, "LocalityPicker"); }
+    LocalityPicker(XdsLb* xds_policy, PickerList pickers)
+        : drop_stats_(xds_policy->drop_stats_),
+          drop_config_(xds_policy->drop_config_),
+          pickers_(std::move(pickers)) {}
 
     PickResult Pick(PickArgs args) override;
 
@@ -170,9 +179,9 @@ class XdsLb : public LoadBalancingPolicy {
     // Calls the picker of the locality that the key falls within.
     PickResult PickFromLocality(const uint32_t key, PickArgs args);
 
-    RefCountedPtr<XdsLb> xds_policy_;
+    RefCountedPtr<XdsClusterDropStats> drop_stats_;
+    RefCountedPtr<XdsApi::DropConfig> drop_config_;
     PickerList pickers_;
-    RefCountedPtr<XdsDropConfig> drop_config_;
   };
 
   class FallbackHelper : public ChannelControlHelper {
@@ -199,206 +208,168 @@ class XdsLb : public LoadBalancingPolicy {
     LoadBalancingPolicy* child_ = nullptr;
   };
 
-  // There is only one PriorityList instance, which has the same lifetime with
-  // the XdsLb instance.
-  class PriorityList {
+  // Each LocalityMap holds a ref to the XdsLb.
+  class LocalityMap : public InternallyRefCounted<LocalityMap> {
    public:
-    // Each LocalityMap holds a ref to the XdsLb.
-    class LocalityMap : public InternallyRefCounted<LocalityMap> {
+    // Each Locality holds a ref to the LocalityMap it is in.
+    class Locality : public InternallyRefCounted<Locality> {
      public:
-      // Each Locality holds a ref to the LocalityMap it is in.
-      class Locality : public InternallyRefCounted<Locality> {
-       public:
-        Locality(RefCountedPtr<LocalityMap> locality_map,
-                 RefCountedPtr<XdsLocalityName> name);
-        ~Locality();
-
-        void UpdateLocked(uint32_t locality_weight,
-                          ServerAddressList serverlist);
-        void ShutdownLocked();
-        void ResetBackoffLocked();
-        void DeactivateLocked();
-        void Orphan() override;
-
-        grpc_connectivity_state connectivity_state() const {
-          return connectivity_state_;
-        }
-        uint32_t weight() const { return weight_; }
-        RefCountedPtr<EndpointPickerWrapper> picker_wrapper() const {
-          return picker_wrapper_;
-        }
+      Locality(RefCountedPtr<LocalityMap> locality_map,
+               RefCountedPtr<XdsLocalityName> name);
+      ~Locality();
 
-        void set_locality_map(RefCountedPtr<LocalityMap> locality_map) {
-          locality_map_ = std::move(locality_map);
-        }
-
-       private:
-        class Helper : public ChannelControlHelper {
-         public:
-          explicit Helper(RefCountedPtr<Locality> locality)
-              : locality_(std::move(locality)) {}
-
-          ~Helper() { locality_.reset(DEBUG_LOCATION, "Helper"); }
-
-          RefCountedPtr<SubchannelInterface> CreateSubchannel(
-              const grpc_channel_args& args) override;
-          void UpdateState(grpc_connectivity_state state,
-                           std::unique_ptr<SubchannelPicker> picker) override;
-          // This is a no-op, because we get the addresses from the xds
-          // client, which is a watch-based API.
-          void RequestReresolution() override {}
-          void AddTraceEvent(TraceSeverity severity,
-                             StringView message) override;
-          void set_child(LoadBalancingPolicy* child) { child_ = child; }
-
-         private:
-          bool CalledByPendingChild() const;
-          bool CalledByCurrentChild() const;
-
-          RefCountedPtr<Locality> locality_;
-          LoadBalancingPolicy* child_ = nullptr;
-        };
-
-        // Methods for dealing with the child policy.
-        OrphanablePtr<LoadBalancingPolicy> CreateChildPolicyLocked(
-            const char* name, const grpc_channel_args* args);
-        grpc_channel_args* CreateChildPolicyArgsLocked(
-            const grpc_channel_args* args);
-
-        static void OnDelayedRemovalTimer(void* arg, grpc_error* error);
-        static void OnDelayedRemovalTimerLocked(void* arg, grpc_error* error);
-
-        XdsLb* xds_policy() const { return locality_map_->xds_policy(); }
-
-        // The owning locality map.
-        RefCountedPtr<LocalityMap> locality_map_;
-
-        RefCountedPtr<XdsLocalityName> name_;
-        OrphanablePtr<LoadBalancingPolicy> child_policy_;
-        OrphanablePtr<LoadBalancingPolicy> pending_child_policy_;
-        RefCountedPtr<EndpointPickerWrapper> picker_wrapper_;
-        grpc_connectivity_state connectivity_state_ = GRPC_CHANNEL_IDLE;
-        uint32_t weight_;
-
-        // States for delayed removal.
-        grpc_timer delayed_removal_timer_;
-        grpc_closure on_delayed_removal_timer_;
-        bool delayed_removal_timer_callback_pending_ = false;
-        bool shutdown_ = false;
-      };
-
-      LocalityMap(RefCountedPtr<XdsLb> xds_policy, uint32_t priority);
-
-      ~LocalityMap() { xds_policy_.reset(DEBUG_LOCATION, "LocalityMap"); }
-
-      void UpdateLocked(
-          const XdsPriorityListUpdate::LocalityMap& locality_map_update);
+      void UpdateLocked(uint32_t locality_weight, ServerAddressList serverlist,
+                        bool update_locality_stats);
+      void ShutdownLocked();
       void ResetBackoffLocked();
-      void UpdateXdsPickerLocked();
-      OrphanablePtr<Locality> ExtractLocalityLocked(
-          const RefCountedPtr<XdsLocalityName>& name);
       void DeactivateLocked();
-      // Returns true if this locality map becomes the currently used one (i.e.,
-      // its priority is selected) after reactivation.
-      bool MaybeReactivateLocked();
-      void MaybeCancelFailoverTimerLocked();
-
       void Orphan() override;
 
-      XdsLb* xds_policy() const { return xds_policy_.get(); }
-      uint32_t priority() const { return priority_; }
+      uint32_t weight() const { return weight_; }
+
       grpc_connectivity_state connectivity_state() const {
         return connectivity_state_;
       }
-      bool failover_timer_callback_pending() const {
-        return failover_timer_callback_pending_;
+
+      RefCountedPtr<LoadReportingPicker> GetLoadReportingPicker() {
+        // Recreate load reporting picker if stats object has changed.
+        if (load_reporting_picker_ == nullptr ||
+            load_reporting_picker_->picker() != picker_wrapper_.get() ||
+            load_reporting_picker_->locality_stats() != stats_.get()) {
+          load_reporting_picker_ =
+              MakeRefCounted<LoadReportingPicker>(picker_wrapper_, stats_);
+        }
+        return load_reporting_picker_;
+      }
+
+      void set_locality_map(RefCountedPtr<LocalityMap> locality_map) {
+        locality_map_ = std::move(locality_map);
       }
 
      private:
-      void OnLocalityStateUpdateLocked();
-      void UpdateConnectivityStateLocked();
+      class Helper : public ChannelControlHelper {
+       public:
+        explicit Helper(RefCountedPtr<Locality> locality)
+            : locality_(std::move(locality)) {}
+
+        ~Helper() { locality_.reset(DEBUG_LOCATION, "Helper"); }
+
+        RefCountedPtr<SubchannelInterface> CreateSubchannel(
+            const grpc_channel_args& args) override;
+        void UpdateState(grpc_connectivity_state state,
+                         std::unique_ptr<SubchannelPicker> picker) override;
+        // This is a no-op, because we get the addresses from the xds
+        // client, which is a watch-based API.
+        void RequestReresolution() override {}
+        void AddTraceEvent(TraceSeverity severity, StringView message) override;
+        void set_child(LoadBalancingPolicy* child) { child_ = child; }
+
+       private:
+        bool CalledByPendingChild() const;
+        bool CalledByCurrentChild() const;
+
+        RefCountedPtr<Locality> locality_;
+        LoadBalancingPolicy* child_ = nullptr;
+      };
+
+      // Methods for dealing with the child policy.
+      OrphanablePtr<LoadBalancingPolicy> CreateChildPolicyLocked(
+          const char* name, const grpc_channel_args* args);
+      grpc_channel_args* CreateChildPolicyArgsLocked(
+          const grpc_channel_args* args);
+
+      void UpdateLocalityStats();
+
       static void OnDelayedRemovalTimer(void* arg, grpc_error* error);
-      static void OnFailoverTimer(void* arg, grpc_error* error);
       static void OnDelayedRemovalTimerLocked(void* arg, grpc_error* error);
-      static void OnFailoverTimerLocked(void* arg, grpc_error* error);
 
-      PriorityList* priority_list() const {
-        return &xds_policy_->priority_list_;
-      }
-      const XdsPriorityListUpdate& priority_list_update() const {
-        return xds_policy_->priority_list_update_;
-      }
-      const XdsPriorityListUpdate::LocalityMap* locality_map_update() const {
-        return xds_policy_->priority_list_update_.Find(priority_);
-      }
+      XdsLb* xds_policy() const { return locality_map_->xds_policy(); }
 
-      RefCountedPtr<XdsLb> xds_policy_;
+      // The owning locality map.
+      RefCountedPtr<LocalityMap> locality_map_;
 
-      std::map<RefCountedPtr<XdsLocalityName>, OrphanablePtr<Locality>,
-               XdsLocalityName::Less>
-          localities_;
-      const uint32_t priority_;
+      RefCountedPtr<XdsLocalityName> name_;
+      RefCountedPtr<XdsClusterLocalityStats> stats_;
+      OrphanablePtr<LoadBalancingPolicy> child_policy_;
+      OrphanablePtr<LoadBalancingPolicy> pending_child_policy_;
+      RefCountedPtr<RefCountedEndpointPicker> picker_wrapper_;
+      RefCountedPtr<LoadReportingPicker> load_reporting_picker_;
       grpc_connectivity_state connectivity_state_ = GRPC_CHANNEL_IDLE;
+      uint32_t weight_;
 
       // States for delayed removal.
       grpc_timer delayed_removal_timer_;
       grpc_closure on_delayed_removal_timer_;
       bool delayed_removal_timer_callback_pending_ = false;
-
-      // States of failover.
-      grpc_timer failover_timer_;
-      grpc_closure on_failover_timer_;
-      bool failover_timer_callback_pending_ = false;
+      bool shutdown_ = false;
     };
 
-    explicit PriorityList(XdsLb* xds_policy) : xds_policy_(xds_policy) {}
+    LocalityMap(RefCountedPtr<XdsLb> xds_policy, uint32_t priority);
 
-    void UpdateLocked();
+    ~LocalityMap() { xds_policy_.reset(DEBUG_LOCATION, "LocalityMap"); }
+
+    void UpdateLocked(
+        const XdsApi::PriorityListUpdate::LocalityMap& locality_map_update,
+        bool update_locality_stats);
     void ResetBackoffLocked();
-    void ShutdownLocked();
     void UpdateXdsPickerLocked();
-
-    const XdsPriorityListUpdate& priority_list_update() const {
-      return xds_policy_->priority_list_update_;
+    OrphanablePtr<Locality> ExtractLocalityLocked(
+        const RefCountedPtr<XdsLocalityName>& name);
+    void DeactivateLocked();
+    // Returns true if this locality map becomes the currently used one (i.e.,
+    // its priority is selected) after reactivation.
+    bool MaybeReactivateLocked();
+    void MaybeCancelFailoverTimerLocked();
+
+    void Orphan() override;
+
+    XdsLb* xds_policy() const { return xds_policy_.get(); }
+    uint32_t priority() const { return priority_; }
+    grpc_connectivity_state connectivity_state() const {
+      return connectivity_state_;
+    }
+    bool failover_timer_callback_pending() const {
+      return failover_timer_callback_pending_;
     }
-    uint32_t current_priority() const { return current_priority_; }
 
    private:
-    void MaybeCreateLocalityMapLocked(uint32_t priority);
-    void FailoverOnConnectionFailureLocked();
-    void FailoverOnDisconnectionLocked(uint32_t failed_priority);
-    void SwitchToHigherPriorityLocked(uint32_t priority);
-    void DeactivatePrioritiesLowerThan(uint32_t priority);
-    OrphanablePtr<LocalityMap::Locality> ExtractLocalityLocked(
-        const RefCountedPtr<XdsLocalityName>& name, uint32_t exclude_priority);
-    // Callers should make sure the priority list is non-empty.
-    uint32_t LowestPriority() const {
-      return static_cast<uint32_t>(priorities_.size()) - 1;
+    void OnLocalityStateUpdateLocked();
+    void UpdateConnectivityStateLocked();
+    static void OnDelayedRemovalTimer(void* arg, grpc_error* error);
+    static void OnFailoverTimer(void* arg, grpc_error* error);
+    static void OnDelayedRemovalTimerLocked(void* arg, grpc_error* error);
+    static void OnFailoverTimerLocked(void* arg, grpc_error* error);
+
+    const XdsApi::PriorityListUpdate& priority_list_update() const {
+      return xds_policy_->priority_list_update_;
+    }
+    const XdsApi::PriorityListUpdate::LocalityMap* locality_map_update() const {
+      return xds_policy_->priority_list_update_.Find(priority_);
     }
-    bool Contains(uint32_t priority) { return priority < priorities_.size(); }
 
-    XdsLb* xds_policy_;
+    RefCountedPtr<XdsLb> xds_policy_;
 
-    // The list of locality maps, indexed by priority. P0 is the highest
-    // priority.
-    InlinedVector<OrphanablePtr<LocalityMap>, 2> priorities_;
-    // The priority that is being used.
-    uint32_t current_priority_ = UINT32_MAX;
+    std::map<RefCountedPtr<XdsLocalityName>, OrphanablePtr<Locality>,
+             XdsLocalityName::Less>
+        localities_;
+    const uint32_t priority_;
+    grpc_connectivity_state connectivity_state_ = GRPC_CHANNEL_IDLE;
+
+    // States for delayed removal.
+    grpc_timer delayed_removal_timer_;
+    grpc_closure on_delayed_removal_timer_;
+    bool delayed_removal_timer_callback_pending_ = false;
+
+    // States of failover.
+    grpc_timer failover_timer_;
+    grpc_closure on_failover_timer_;
+    bool failover_timer_callback_pending_ = false;
   };
 
   ~XdsLb();
 
   void ShutdownLocked() override;
 
-  // Methods for dealing with fallback state.
-  void MaybeCancelFallbackAtStartupChecks();
-  static void OnFallbackTimer(void* arg, grpc_error* error);
-  static void OnFallbackTimerLocked(void* arg, grpc_error* error);
-  void UpdateFallbackPolicyLocked();
-  OrphanablePtr<LoadBalancingPolicy> CreateFallbackPolicyLocked(
-      const char* name, const grpc_channel_args* args);
-  void MaybeExitFallbackMode();
-
   const char* eds_service_name() const {
     if (config_ != nullptr && config_->eds_service_name() != nullptr) {
       return config_->eds_service_name();
@@ -411,12 +382,36 @@ class XdsLb : public LoadBalancingPolicy {
                                                : xds_client_.get();
   }
 
+  void UpdatePrioritiesLocked(bool update_locality_stats);
+  void UpdateXdsPickerLocked();
+  void MaybeCreateLocalityMapLocked(uint32_t priority);
+  void FailoverOnConnectionFailureLocked();
+  void FailoverOnDisconnectionLocked(uint32_t failed_priority);
+  void SwitchToHigherPriorityLocked(uint32_t priority);
+  void DeactivatePrioritiesLowerThan(uint32_t priority);
+  OrphanablePtr<LocalityMap::Locality> ExtractLocalityLocked(
+      const RefCountedPtr<XdsLocalityName>& name, uint32_t exclude_priority);
+  // Callers should make sure the priority list is non-empty.
+  uint32_t LowestPriority() const {
+    return static_cast<uint32_t>(priorities_.size()) - 1;
+  }
+  bool Contains(uint32_t priority) { return priority < priorities_.size(); }
+
+  // Methods for dealing with fallback state.
+  void MaybeCancelFallbackAtStartupChecks();
+  static void OnFallbackTimer(void* arg, grpc_error* error);
+  static void OnFallbackTimerLocked(void* arg, grpc_error* error);
+  void UpdateFallbackPolicyLocked();
+  OrphanablePtr<LoadBalancingPolicy> CreateFallbackPolicyLocked(
+      const char* name, const grpc_channel_args* args);
+  void MaybeExitFallbackMode();
+
   // Server name from target URI.
   std::string server_name_;
 
   // Current channel args and config from the resolver.
   const grpc_channel_args* args_ = nullptr;
-  RefCountedPtr<ParsedXdsConfig> config_;
+  RefCountedPtr<XdsConfig> config_;
 
   // Internal state.
   bool shutting_down_ = false;
@@ -454,23 +449,26 @@ class XdsLb : public LoadBalancingPolicy {
 
   const grpc_millis locality_retention_interval_ms_;
   const grpc_millis locality_map_failover_timeout_ms_;
-  // A list of locality maps indexed by priority.
-  PriorityList priority_list_;
+  // The list of locality maps, indexed by priority. P0 is the highest
+  // priority.
+  InlinedVector<OrphanablePtr<LocalityMap>, 2> priorities_;
+  // The priority that is being used.
+  uint32_t current_priority_ = UINT32_MAX;
   // The update for priority_list_.
-  XdsPriorityListUpdate priority_list_update_;
+  XdsApi::PriorityListUpdate priority_list_update_;
 
   // The config for dropping calls.
-  RefCountedPtr<XdsDropConfig> drop_config_;
+  RefCountedPtr<XdsApi::DropConfig> drop_config_;
 
-  // The stats for client-side load reporting.
-  XdsClientStats client_stats_;
+  // Drop stats for client-side load reporting.
+  RefCountedPtr<XdsClusterDropStats> drop_stats_;
 };
 
 //
-// XdsLb::EndpointPickerWrapper
+// XdsLb::LoadReportingPicker
 //
 
-LoadBalancingPolicy::PickResult XdsLb::EndpointPickerWrapper::Pick(
+LoadBalancingPolicy::PickResult XdsLb::LoadReportingPicker::Pick(
     LoadBalancingPolicy::PickArgs args) {
   // Forward the pick to the picker returned from the child policy.
   PickResult result = picker_->Pick(args);
@@ -481,7 +479,7 @@ LoadBalancingPolicy::PickResult XdsLb::EndpointPickerWrapper::Pick(
   // Record a call started.
   locality_stats_->AddCallStarted();
   // Intercept the recv_trailing_metadata op to record call completion.
-  XdsClientStats::LocalityStats* locality_stats =
+  XdsClusterLocalityStats* locality_stats =
       locality_stats_->Ref(DEBUG_LOCATION, "LocalityStats+call").release();
   result.recv_trailing_metadata_ready =
       // Note: This callback does not run in either the control plane
@@ -503,7 +501,7 @@ XdsLb::PickResult XdsLb::LocalityPicker::Pick(PickArgs args) {
   // Handle drop.
   const std::string* drop_category;
   if (drop_config_->ShouldDrop(&drop_category)) {
-    xds_policy_->client_stats_.AddCallDropped(*drop_category);
+    if (drop_stats_ != nullptr) drop_stats_->AddCallDropped(*drop_category);
     PickResult result;
     result.type = PickResult::PICK_COMPLETE;
     return result;
@@ -620,7 +618,7 @@ class XdsLb::EndpointWatcher : public XdsClient::EndpointWatcherInterface {
 
   ~EndpointWatcher() { xds_policy_.reset(DEBUG_LOCATION, "EndpointWatcher"); }
 
-  void OnEndpointChanged(EdsUpdate update) override {
+  void OnEndpointChanged(XdsApi::EdsUpdate update) override {
     if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
       gpr_log(GPR_INFO, "[xdslb %p] Received EDS update from xds client",
               xds_policy_.get());
@@ -642,13 +640,13 @@ class XdsLb::EndpointWatcher : public XdsClient::EndpointWatcherInterface {
                 xds_policy_.get(), drop_config_changed);
       }
       if (drop_config_changed) {
-        xds_policy_->priority_list_.UpdateXdsPickerLocked();
+        xds_policy_->UpdateXdsPickerLocked();
       }
       return;
     }
     // Update the priority list.
     xds_policy_->priority_list_update_ = std::move(update.priority_list_update);
-    xds_policy_->priority_list_.UpdateLocked();
+    xds_policy_->UpdatePrioritiesLocked(false /*update_locality_stats*/);
   }
 
   void OnError(grpc_error* error) override {
@@ -701,8 +699,7 @@ XdsLb::XdsLb(Args args)
           {GRPC_XDS_DEFAULT_LOCALITY_RETENTION_INTERVAL_MS, 0, INT_MAX})),
       locality_map_failover_timeout_ms_(grpc_channel_args_find_integer(
           args.args, GRPC_ARG_XDS_FAILOVER_TIMEOUT_MS,
-          {GRPC_XDS_DEFAULT_FAILOVER_TIMEOUT_MS, 0, INT_MAX})),
-      priority_list_(this) {
+          {GRPC_XDS_DEFAULT_FAILOVER_TIMEOUT_MS, 0, INT_MAX})) {
   if (xds_client_from_channel_ != nullptr &&
       GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Using xds client %p from channel", this,
@@ -735,7 +732,8 @@ void XdsLb::ShutdownLocked() {
   }
   shutting_down_ = true;
   MaybeCancelFallbackAtStartupChecks();
-  priority_list_.ShutdownLocked();
+  priorities_.clear();
+  drop_stats_.reset();
   if (fallback_policy_ != nullptr) {
     grpc_pollset_set_del_pollset_set(fallback_policy_->interested_parties(),
                                      interested_parties());
@@ -746,19 +744,15 @@ void XdsLb::ShutdownLocked() {
   }
   fallback_policy_.reset();
   pending_fallback_policy_.reset();
-  // Cancel the endpoint watch here instead of in our dtor, because the
-  // watcher holds a ref to us.
-  xds_client()->CancelEndpointDataWatch(StringView(eds_service_name()),
-                                        endpoint_watcher_);
-  if (config_->lrs_load_reporting_server_name().has_value()) {
-    // TODO(roth): We should pass the cluster name (in addition to the
-    // eds_service_name) when adding the client stats. To do so, we need to
-    // first find a way to plumb the cluster name down into this LB policy.
-    xds_client()->RemoveClientStats(
-        StringView(config_->lrs_load_reporting_server_name().value().c_str()),
-        StringView(eds_service_name()), &client_stats_);
-  }
-  xds_client_from_channel_.reset();
+  // Cancel the endpoint watch here instead of in our dtor if we are using the
+  // XdsResolver, because the watcher holds a ref to us and we might not be
+  // destroying the Xds client leading to a situation where the Xds lb policy is
+  // never destroyed.
+  if (xds_client_from_channel_ != nullptr) {
+    xds_client()->CancelEndpointDataWatch(StringView(eds_service_name()),
+                                          endpoint_watcher_);
+    xds_client_from_channel_.reset();
+  }
   xds_client_.reset();
 }
 
@@ -771,7 +765,9 @@ void XdsLb::ResetBackoffLocked() {
   // LB policy, this is done via the resolver, so we don't need to do it
   // for xds_client_from_channel_ here.
   if (xds_client_ != nullptr) xds_client_->ResetBackoff();
-  priority_list_.ResetBackoffLocked();
+  for (size_t i = 0; i < priorities_.size(); ++i) {
+    priorities_[i]->ResetBackoffLocked();
+  }
   if (fallback_policy_ != nullptr) {
     fallback_policy_->ResetBackoffLocked();
   }
@@ -795,8 +791,6 @@ void XdsLb::UpdateLocked(UpdateArgs args) {
   grpc_channel_args_destroy(args_);
   args_ = args.args;
   args.args = nullptr;
-  // Update priority list.
-  priority_list_.UpdateLocked();
   // Update the existing fallback policy. The fallback policy config and/or the
   // fallback addresses may be new.
   if (fallback_policy_ != nullptr) UpdateFallbackPolicyLocked();
@@ -823,6 +817,31 @@ void XdsLb::UpdateLocked(UpdateArgs args) {
     fallback_at_startup_checks_pending_ = true;
     grpc_timer_init(&lb_fallback_timer_, deadline, &lb_on_fallback_);
   }
+  // Update drop stats for load reporting if needed.
+  if (is_initial_update || config_->lrs_load_reporting_server_name() !=
+                               old_config->lrs_load_reporting_server_name()) {
+    drop_stats_.reset();
+    if (config_->lrs_load_reporting_server_name().has_value()) {
+      drop_stats_ = xds_client()->AddClusterDropStats(
+          config_->lrs_load_reporting_server_name().value(),
+          // TODO(roth): We currently hard-code the assumption that
+          // cluster name and EDS service name are the same.  Fix this
+          // as part of refectoring this LB policy.
+          eds_service_name(), eds_service_name());
+    }
+  }
+  // Update priority list.
+  // Note that this comes after updating drop_stats_, since we want that
+  // to be used by any new picker we create here.
+  // No need to do this on the initial update, since there won't be any
+  // priorities to update yet.
+  if (!is_initial_update) {
+    const bool update_locality_stats =
+        config_->lrs_load_reporting_server_name() !=
+            old_config->lrs_load_reporting_server_name() ||
+        strcmp(old_eds_service_name, eds_service_name()) != 0;
+    UpdatePrioritiesLocked(update_locality_stats);
+  }
   // Update endpoint watcher if needed.
   if (is_initial_update ||
       strcmp(old_eds_service_name, eds_service_name()) != 0) {
@@ -830,40 +849,12 @@ void XdsLb::UpdateLocked(UpdateArgs args) {
       xds_client()->CancelEndpointDataWatch(StringView(old_eds_service_name),
                                             endpoint_watcher_);
     }
-    auto watcher =
-        MakeUnique<EndpointWatcher>(Ref(DEBUG_LOCATION, "EndpointWatcher"));
+    auto watcher = absl::make_unique<EndpointWatcher>(
+        Ref(DEBUG_LOCATION, "EndpointWatcher"));
     endpoint_watcher_ = watcher.get();
     xds_client()->WatchEndpointData(StringView(eds_service_name()),
                                     std::move(watcher));
   }
-  // Update load reporting if needed.
-  // TODO(roth): Ideally, we should not collect any stats if load reporting
-  // is disabled, which would require changing this code to recreate
-  // all of the pickers whenever load reporting is enabled or disabled
-  // here.
-  if (is_initial_update ||
-      (config_->lrs_load_reporting_server_name().has_value()) !=
-          (old_config->lrs_load_reporting_server_name().has_value()) ||
-      (config_->lrs_load_reporting_server_name().has_value() &&
-       old_config->lrs_load_reporting_server_name().has_value() &&
-       config_->lrs_load_reporting_server_name().value() !=
-           old_config->lrs_load_reporting_server_name().value())) {
-    if (old_config != nullptr &&
-        old_config->lrs_load_reporting_server_name().has_value()) {
-      xds_client()->RemoveClientStats(
-          StringView(
-              old_config->lrs_load_reporting_server_name().value().c_str()),
-          StringView(old_eds_service_name), &client_stats_);
-    }
-    if (config_->lrs_load_reporting_server_name().has_value()) {
-      // TODO(roth): We should pass the cluster name (in addition to the
-      // eds_service_name) when adding the client stats. To do so, we need to
-      // first find a way to plumb the cluster name down into this LB policy.
-      xds_client()->AddClientStats(
-          StringView(config_->lrs_load_reporting_server_name().value().c_str()),
-          StringView(eds_service_name()), &client_stats_);
-    }
-  }
 }
 
 //
@@ -1043,21 +1034,22 @@ void XdsLb::MaybeExitFallbackMode() {
 }
 
 //
-// XdsLb::PriorityList
+// priority list-related methods
 //
 
-void XdsLb::PriorityList::UpdateLocked() {
-  const auto& priority_list_update = xds_policy_->priority_list_update_;
+void XdsLb::UpdatePrioritiesLocked(bool update_locality_stats) {
   // 1. Remove from the priority list the priorities that are not in the update.
-  DeactivatePrioritiesLowerThan(priority_list_update.LowestPriority());
+  DeactivatePrioritiesLowerThan(priority_list_update_.LowestPriority());
   // 2. Update all the existing priorities.
   for (uint32_t priority = 0; priority < priorities_.size(); ++priority) {
     LocalityMap* locality_map = priorities_[priority].get();
-    const auto* locality_map_update = priority_list_update.Find(priority);
+    const auto* locality_map_update = priority_list_update_.Find(priority);
+    // If we have more current priorities than exist in the update, stop here.
+    if (locality_map_update == nullptr) break;
     // Propagate locality_map_update.
     // TODO(juanlishen): Find a clean way to skip duplicate update for a
     // priority.
-    locality_map->UpdateLocked(*locality_map_update);
+    locality_map->UpdateLocked(*locality_map_update, update_locality_stats);
   }
   // 3. Only create a new locality map if all the existing ones have failed.
   if (priorities_.empty() ||
@@ -1069,54 +1061,51 @@ void XdsLb::PriorityList::UpdateLocked() {
     // to be created.
     MaybeCreateLocalityMapLocked(new_priority);
   }
-}
-
-void XdsLb::PriorityList::ResetBackoffLocked() {
-  for (size_t i = 0; i < priorities_.size(); ++i) {
-    priorities_[i]->ResetBackoffLocked();
+  // 4. If we updated locality stats and we already have at least one
+  // priority, update the picker to start using the new stats object(s).
+  if (update_locality_stats && !priorities_.empty()) {
+    UpdateXdsPickerLocked();
   }
 }
 
-void XdsLb::PriorityList::ShutdownLocked() { priorities_.clear(); }
-
-void XdsLb::PriorityList::UpdateXdsPickerLocked() {
+void XdsLb::UpdateXdsPickerLocked() {
   // If we are in fallback mode, don't generate an xds picker from localities.
-  if (xds_policy_->fallback_policy_ != nullptr) return;
-  if (current_priority() == UINT32_MAX) {
+  if (fallback_policy_ != nullptr) return;
+  if (current_priority_ == UINT32_MAX) {
     grpc_error* error = grpc_error_set_int(
         GRPC_ERROR_CREATE_FROM_STATIC_STRING("no ready locality map"),
         GRPC_ERROR_INT_GRPC_STATUS, GRPC_STATUS_UNAVAILABLE);
-    xds_policy_->channel_control_helper()->UpdateState(
+    channel_control_helper()->UpdateState(
         GRPC_CHANNEL_TRANSIENT_FAILURE,
-        MakeUnique<TransientFailurePicker>(error));
+        absl::make_unique<TransientFailurePicker>(error));
     return;
   }
   priorities_[current_priority_]->UpdateXdsPickerLocked();
 }
 
-void XdsLb::PriorityList::MaybeCreateLocalityMapLocked(uint32_t priority) {
+void XdsLb::MaybeCreateLocalityMapLocked(uint32_t priority) {
   // Exhausted priorities in the update.
-  if (!priority_list_update().Contains(priority)) return;
-  auto new_locality_map = new LocalityMap(
-      xds_policy_->Ref(DEBUG_LOCATION, "LocalityMap"), priority);
+  if (!priority_list_update_.Contains(priority)) return;
+  auto new_locality_map =
+      new LocalityMap(Ref(DEBUG_LOCATION, "LocalityMap"), priority);
   priorities_.emplace_back(OrphanablePtr<LocalityMap>(new_locality_map));
-  new_locality_map->UpdateLocked(*priority_list_update().Find(priority));
+  new_locality_map->UpdateLocked(*priority_list_update_.Find(priority),
+                                 false /*update_locality_stats*/);
 }
 
-void XdsLb::PriorityList::FailoverOnConnectionFailureLocked() {
+void XdsLb::FailoverOnConnectionFailureLocked() {
   const uint32_t failed_priority = LowestPriority();
   // If we're failing over from the lowest priority, report TRANSIENT_FAILURE.
-  if (failed_priority == priority_list_update().LowestPriority()) {
+  if (failed_priority == priority_list_update_.LowestPriority()) {
     UpdateXdsPickerLocked();
   }
   MaybeCreateLocalityMapLocked(failed_priority + 1);
 }
 
-void XdsLb::PriorityList::FailoverOnDisconnectionLocked(
-    uint32_t failed_priority) {
+void XdsLb::FailoverOnDisconnectionLocked(uint32_t failed_priority) {
   current_priority_ = UINT32_MAX;
   for (uint32_t next_priority = failed_priority + 1;
-       next_priority <= priority_list_update().LowestPriority();
+       next_priority <= priority_list_update_.LowestPriority();
        ++next_priority) {
     if (!Contains(next_priority)) {
       MaybeCreateLocalityMapLocked(next_priority);
@@ -1126,17 +1115,17 @@ void XdsLb::PriorityList::FailoverOnDisconnectionLocked(
   }
 }
 
-void XdsLb::PriorityList::SwitchToHigherPriorityLocked(uint32_t priority) {
+void XdsLb::SwitchToHigherPriorityLocked(uint32_t priority) {
   current_priority_ = priority;
   DeactivatePrioritiesLowerThan(current_priority_);
   UpdateXdsPickerLocked();
 }
 
-void XdsLb::PriorityList::DeactivatePrioritiesLowerThan(uint32_t priority) {
+void XdsLb::DeactivatePrioritiesLowerThan(uint32_t priority) {
   if (priorities_.empty()) return;
   // Deactivate the locality maps from the lowest priority.
   for (uint32_t p = LowestPriority(); p > priority; --p) {
-    if (xds_policy_->locality_retention_interval_ms_ == 0) {
+    if (locality_retention_interval_ms_ == 0) {
       priorities_.pop_back();
     } else {
       priorities_[p]->DeactivateLocked();
@@ -1144,8 +1133,7 @@ void XdsLb::PriorityList::DeactivatePrioritiesLowerThan(uint32_t priority) {
   }
 }
 
-OrphanablePtr<XdsLb::PriorityList::LocalityMap::Locality>
-XdsLb::PriorityList::ExtractLocalityLocked(
+OrphanablePtr<XdsLb::LocalityMap::Locality> XdsLb::ExtractLocalityLocked(
     const RefCountedPtr<XdsLocalityName>& name, uint32_t exclude_priority) {
   for (uint32_t priority = 0; priority < priorities_.size(); ++priority) {
     if (priority == exclude_priority) continue;
@@ -1157,11 +1145,11 @@ XdsLb::PriorityList::ExtractLocalityLocked(
 }
 
 //
-// XdsLb::PriorityList::LocalityMap
+// XdsLb::LocalityMap
 //
 
-XdsLb::PriorityList::LocalityMap::LocalityMap(RefCountedPtr<XdsLb> xds_policy,
-                                              uint32_t priority)
+XdsLb::LocalityMap::LocalityMap(RefCountedPtr<XdsLb> xds_policy,
+                                uint32_t priority)
     : xds_policy_(std::move(xds_policy)), priority_(priority) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Creating priority %" PRIu32,
@@ -1179,13 +1167,15 @@ XdsLb::PriorityList::LocalityMap::LocalityMap(RefCountedPtr<XdsLb> xds_policy,
   // This is the first locality map ever created, report CONNECTING.
   if (priority_ == 0) {
     xds_policy_->channel_control_helper()->UpdateState(
-        GRPC_CHANNEL_CONNECTING, MakeUnique<QueuePicker>(xds_policy_->Ref(
-                                     DEBUG_LOCATION, "QueuePicker")));
+        GRPC_CHANNEL_CONNECTING,
+        absl::make_unique<QueuePicker>(
+            xds_policy_->Ref(DEBUG_LOCATION, "QueuePicker")));
   }
 }
 
-void XdsLb::PriorityList::LocalityMap::UpdateLocked(
-    const XdsPriorityListUpdate::LocalityMap& locality_map_update) {
+void XdsLb::LocalityMap::UpdateLocked(
+    const XdsApi::PriorityListUpdate::LocalityMap& locality_map_update,
+    bool update_locality_stats) {
   if (xds_policy_->shutting_down_) return;
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Start Updating priority %" PRIu32,
@@ -1216,7 +1206,7 @@ void XdsLb::PriorityList::LocalityMap::UpdateLocked(
     OrphanablePtr<Locality>& locality = localities_[name];
     if (locality == nullptr) {
       // Move from another locality map if possible.
-      locality = priority_list()->ExtractLocalityLocked(name, priority_);
+      locality = xds_policy_->ExtractLocalityLocked(name, priority_);
       if (locality != nullptr) {
         locality->set_locality_map(
             Ref(DEBUG_LOCATION, "LocalityMap+Locality_move"));
@@ -1228,39 +1218,38 @@ void XdsLb::PriorityList::LocalityMap::UpdateLocked(
     // Keep a copy of serverlist in the update so that we can compare it
     // with the future ones.
     locality->UpdateLocked(locality_update.lb_weight,
-                           locality_update.serverlist);
+                           locality_update.serverlist, update_locality_stats);
   }
 }
 
-void XdsLb::PriorityList::LocalityMap::ResetBackoffLocked() {
+void XdsLb::LocalityMap::ResetBackoffLocked() {
   for (auto& p : localities_) p.second->ResetBackoffLocked();
 }
 
-void XdsLb::PriorityList::LocalityMap::UpdateXdsPickerLocked() {
+void XdsLb::LocalityMap::UpdateXdsPickerLocked() {
   // Construct a new xds picker which maintains a map of all locality pickers
   // that are ready. Each locality is represented by a portion of the range
   // proportional to its weight, such that the total range is the sum of the
   // weights of all localities.
   LocalityPicker::PickerList picker_list;
   uint32_t end = 0;
-  for (const auto& p : localities_) {
+  for (auto& p : localities_) {
     const auto& locality_name = p.first;
-    const Locality* locality = p.second.get();
+    Locality* locality = p.second.get();
     // Skip the localities that are not in the latest locality map update.
     if (!locality_map_update()->Contains(locality_name)) continue;
     if (locality->connectivity_state() != GRPC_CHANNEL_READY) continue;
     end += locality->weight();
-    picker_list.push_back(std::make_pair(end, locality->picker_wrapper()));
+    picker_list.push_back(
+        std::make_pair(end, locality->GetLoadReportingPicker()));
   }
   xds_policy()->channel_control_helper()->UpdateState(
       GRPC_CHANNEL_READY,
-      grpc_core::MakeUnique<LocalityPicker>(
-          xds_policy_->Ref(DEBUG_LOCATION, "LocalityPicker"),
-          std::move(picker_list)));
+      absl::make_unique<LocalityPicker>(xds_policy(), std::move(picker_list)));
 }
 
-OrphanablePtr<XdsLb::PriorityList::LocalityMap::Locality>
-XdsLb::PriorityList::LocalityMap::ExtractLocalityLocked(
+OrphanablePtr<XdsLb::LocalityMap::Locality>
+XdsLb::LocalityMap::ExtractLocalityLocked(
     const RefCountedPtr<XdsLocalityName>& name) {
   for (auto iter = localities_.begin(); iter != localities_.end(); ++iter) {
     const auto& name_in_map = iter->first;
@@ -1273,7 +1262,7 @@ XdsLb::PriorityList::LocalityMap::ExtractLocalityLocked(
   return nullptr;
 }
 
-void XdsLb::PriorityList::LocalityMap::DeactivateLocked() {
+void XdsLb::LocalityMap::DeactivateLocked() {
   // If already deactivated, don't do it again.
   if (delayed_removal_timer_callback_pending_) return;
   MaybeCancelFailoverTimerLocked();
@@ -1294,24 +1283,24 @@ void XdsLb::PriorityList::LocalityMap::DeactivateLocked() {
   delayed_removal_timer_callback_pending_ = true;
 }
 
-bool XdsLb::PriorityList::LocalityMap::MaybeReactivateLocked() {
+bool XdsLb::LocalityMap::MaybeReactivateLocked() {
   // Don't reactivate a priority that is not higher than the current one.
-  if (priority_ >= priority_list()->current_priority()) return false;
+  if (priority_ >= xds_policy_->current_priority_) return false;
   // Reactivate this priority by cancelling deletion timer.
   if (delayed_removal_timer_callback_pending_) {
     grpc_timer_cancel(&delayed_removal_timer_);
   }
   // Switch to this higher priority if it's READY.
   if (connectivity_state_ != GRPC_CHANNEL_READY) return false;
-  priority_list()->SwitchToHigherPriorityLocked(priority_);
+  xds_policy_->SwitchToHigherPriorityLocked(priority_);
   return true;
 }
 
-void XdsLb::PriorityList::LocalityMap::MaybeCancelFailoverTimerLocked() {
+void XdsLb::LocalityMap::MaybeCancelFailoverTimerLocked() {
   if (failover_timer_callback_pending_) grpc_timer_cancel(&failover_timer_);
 }
 
-void XdsLb::PriorityList::LocalityMap::Orphan() {
+void XdsLb::LocalityMap::Orphan() {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Priority %" PRIu32 " orphaned.", xds_policy(),
             priority_);
@@ -1324,11 +1313,11 @@ void XdsLb::PriorityList::LocalityMap::Orphan() {
   Unref(DEBUG_LOCATION, "LocalityMap+Orphan");
 }
 
-void XdsLb::PriorityList::LocalityMap::OnLocalityStateUpdateLocked() {
+void XdsLb::LocalityMap::OnLocalityStateUpdateLocked() {
   UpdateConnectivityStateLocked();
   // Ignore priorities not in priority_list_update.
   if (!priority_list_update().Contains(priority_)) return;
-  const uint32_t current_priority = priority_list()->current_priority();
+  const uint32_t current_priority = xds_policy_->current_priority_;
   // Ignore lower-than-current priorities.
   if (priority_ > current_priority) return;
   // Maybe update fallback state.
@@ -1342,13 +1331,13 @@ void XdsLb::PriorityList::LocalityMap::OnLocalityStateUpdateLocked() {
     if (connectivity_state_ == GRPC_CHANNEL_READY) {
       MaybeCancelFailoverTimerLocked();
       // If a higher-than-current priority becomes READY, switch to use it.
-      priority_list()->SwitchToHigherPriorityLocked(priority_);
+      xds_policy_->SwitchToHigherPriorityLocked(priority_);
     } else if (connectivity_state_ == GRPC_CHANNEL_TRANSIENT_FAILURE) {
       // If a higher-than-current priority becomes TRANSIENT_FAILURE, only
       // handle it if it's the priority that is still in failover timeout.
       if (failover_timer_callback_pending_) {
         MaybeCancelFailoverTimerLocked();
-        priority_list()->FailoverOnConnectionFailureLocked();
+        xds_policy_->FailoverOnConnectionFailureLocked();
       }
     }
     return;
@@ -1356,7 +1345,7 @@ void XdsLb::PriorityList::LocalityMap::OnLocalityStateUpdateLocked() {
   // Update is for current priority.
   if (connectivity_state_ != GRPC_CHANNEL_READY) {
     // Fail over if it's no longer READY.
-    priority_list()->FailoverOnDisconnectionLocked(priority_);
+    xds_policy_->FailoverOnDisconnectionLocked(priority_);
   }
   // At this point, one of the following things has happened to the current
   // priority.
@@ -1364,10 +1353,10 @@ void XdsLb::PriorityList::LocalityMap::OnLocalityStateUpdateLocked() {
   // 2. It changed to a lower priority due to failover.
   // 3. It became invalid because failover didn't yield a READY priority.
   // In any case, update the xds picker.
-  priority_list()->UpdateXdsPickerLocked();
+  xds_policy_->UpdateXdsPickerLocked();
 }
 
-void XdsLb::PriorityList::LocalityMap::UpdateConnectivityStateLocked() {
+void XdsLb::LocalityMap::UpdateConnectivityStateLocked() {
   size_t num_ready = 0;
   size_t num_connecting = 0;
   size_t num_idle = 0;
@@ -1415,8 +1404,7 @@ void XdsLb::PriorityList::LocalityMap::UpdateConnectivityStateLocked() {
   }
 }
 
-void XdsLb::PriorityList::LocalityMap::OnDelayedRemovalTimer(
-    void* arg, grpc_error* error) {
+void XdsLb::LocalityMap::OnDelayedRemovalTimer(void* arg, grpc_error* error) {
   LocalityMap* self = static_cast<LocalityMap*>(arg);
   self->xds_policy_->combiner()->Run(
       GRPC_CLOSURE_INIT(&self->on_delayed_removal_timer_,
@@ -1424,14 +1412,13 @@ void XdsLb::PriorityList::LocalityMap::OnDelayedRemovalTimer(
       GRPC_ERROR_REF(error));
 }
 
-void XdsLb::PriorityList::LocalityMap::OnDelayedRemovalTimerLocked(
-    void* arg, grpc_error* error) {
+void XdsLb::LocalityMap::OnDelayedRemovalTimerLocked(void* arg,
+                                                     grpc_error* error) {
   LocalityMap* self = static_cast<LocalityMap*>(arg);
   self->delayed_removal_timer_callback_pending_ = false;
   if (error == GRPC_ERROR_NONE && !self->xds_policy_->shutting_down_) {
-    auto* priority_list = self->priority_list();
     const bool keep = self->priority_list_update().Contains(self->priority_) &&
-                      self->priority_ <= priority_list->current_priority();
+                      self->priority_ <= self->xds_policy_->current_priority_;
     if (!keep) {
       // This check is to make sure we always delete the locality maps from
       // the lowest priority even if the closures of the back-to-back timers
@@ -1440,8 +1427,8 @@ void XdsLb::PriorityList::LocalityMap::OnDelayedRemovalTimerLocked(
       // deactivated locality maps when out-of-order closures are run.
       // TODO(juanlishen): Check the timer implementation to see if this
       // defense is necessary.
-      if (self->priority_ == priority_list->LowestPriority()) {
-        priority_list->priorities_.pop_back();
+      if (self->priority_ == self->xds_policy_->LowestPriority()) {
+        self->xds_policy_->priorities_.pop_back();
       } else {
         gpr_log(GPR_ERROR,
                 "[xdslb %p] Priority %" PRIu32
@@ -1454,8 +1441,7 @@ void XdsLb::PriorityList::LocalityMap::OnDelayedRemovalTimerLocked(
   self->Unref(DEBUG_LOCATION, "LocalityMap+timer");
 }
 
-void XdsLb::PriorityList::LocalityMap::OnFailoverTimer(void* arg,
-                                                       grpc_error* error) {
+void XdsLb::LocalityMap::OnFailoverTimer(void* arg, grpc_error* error) {
   LocalityMap* self = static_cast<LocalityMap*>(arg);
   self->xds_policy_->combiner()->Run(
       GRPC_CLOSURE_INIT(&self->on_failover_timer_, OnFailoverTimerLocked, self,
@@ -1463,31 +1449,31 @@ void XdsLb::PriorityList::LocalityMap::OnFailoverTimer(void* arg,
       GRPC_ERROR_REF(error));
 }
 
-void XdsLb::PriorityList::LocalityMap::OnFailoverTimerLocked(
-    void* arg, grpc_error* error) {
+void XdsLb::LocalityMap::OnFailoverTimerLocked(void* arg, grpc_error* error) {
   LocalityMap* self = static_cast<LocalityMap*>(arg);
   self->failover_timer_callback_pending_ = false;
   if (error == GRPC_ERROR_NONE && !self->xds_policy_->shutting_down_) {
-    self->priority_list()->FailoverOnConnectionFailureLocked();
+    self->xds_policy_->FailoverOnConnectionFailureLocked();
   }
   self->Unref(DEBUG_LOCATION, "LocalityMap+OnFailoverTimerLocked");
 }
 
 //
-// XdsLb::PriorityList::LocalityMap::Locality
+// XdsLb::LocalityMap::Locality
 //
 
-XdsLb::PriorityList::LocalityMap::Locality::Locality(
-    RefCountedPtr<LocalityMap> locality_map,
-    RefCountedPtr<XdsLocalityName> name)
+XdsLb::LocalityMap::Locality::Locality(RefCountedPtr<LocalityMap> locality_map,
+                                       RefCountedPtr<XdsLocalityName> name)
     : locality_map_(std::move(locality_map)), name_(std::move(name)) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] created Locality %p for %s", xds_policy(),
             this, name_->AsHumanReadableString());
   }
+  // Initialize locality stats if load reporting is enabled.
+  UpdateLocalityStats();
 }
 
-XdsLb::PriorityList::LocalityMap::Locality::~Locality() {
+XdsLb::LocalityMap::Locality::~Locality() {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Locality %p %s: destroying locality",
             xds_policy(), this, name_->AsHumanReadableString());
@@ -1495,8 +1481,20 @@ XdsLb::PriorityList::LocalityMap::Locality::~Locality() {
   locality_map_.reset(DEBUG_LOCATION, "Locality");
 }
 
-grpc_channel_args*
-XdsLb::PriorityList::LocalityMap::Locality::CreateChildPolicyArgsLocked(
+void XdsLb::LocalityMap::Locality::UpdateLocalityStats() {
+  stats_.reset();
+  if (xds_policy()->config_->lrs_load_reporting_server_name().has_value()) {
+    stats_ = xds_policy()->xds_client()->AddClusterLocalityStats(
+        xds_policy()->config_->lrs_load_reporting_server_name().value(),
+        // TODO(roth): We currently hard-code the assumption that
+        // cluster name and EDS service name are the same.  Fix this
+        // as part of refectoring this LB policy.
+        xds_policy()->eds_service_name(), xds_policy()->eds_service_name(),
+        name_);
+  }
+}
+
+grpc_channel_args* XdsLb::LocalityMap::Locality::CreateChildPolicyArgsLocked(
     const grpc_channel_args* args_in) {
   const grpc_arg args_to_add[] = {
       // A channel arg indicating if the target is a backend inferred from a
@@ -1514,7 +1512,7 @@ XdsLb::PriorityList::LocalityMap::Locality::CreateChildPolicyArgsLocked(
 }
 
 OrphanablePtr<LoadBalancingPolicy>
-XdsLb::PriorityList::LocalityMap::Locality::CreateChildPolicyLocked(
+XdsLb::LocalityMap::Locality::CreateChildPolicyLocked(
     const char* name, const grpc_channel_args* args) {
   Helper* helper = new Helper(this->Ref(DEBUG_LOCATION, "Helper"));
   LoadBalancingPolicy::Args lb_policy_args;
@@ -1546,14 +1544,17 @@ XdsLb::PriorityList::LocalityMap::Locality::CreateChildPolicyLocked(
   return lb_policy;
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::UpdateLocked(
-    uint32_t locality_weight, ServerAddressList serverlist) {
+void XdsLb::LocalityMap::Locality::UpdateLocked(uint32_t locality_weight,
+                                                ServerAddressList serverlist,
+                                                bool update_locality_stats) {
   if (xds_policy()->shutting_down_) return;
   // Update locality weight.
   weight_ = locality_weight;
   if (delayed_removal_timer_callback_pending_) {
     grpc_timer_cancel(&delayed_removal_timer_);
   }
+  // Update locality stats.
+  if (update_locality_stats) UpdateLocalityStats();
   // Construct update args.
   UpdateArgs update_args;
   update_args.addresses = std::move(serverlist);
@@ -1656,11 +1657,12 @@ void XdsLb::PriorityList::LocalityMap::Locality::UpdateLocked(
   policy_to_update->UpdateLocked(std::move(update_args));
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::ShutdownLocked() {
+void XdsLb::LocalityMap::Locality::ShutdownLocked() {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_trace)) {
     gpr_log(GPR_INFO, "[xdslb %p] Locality %p %s: shutting down locality",
             xds_policy(), this, name_->AsHumanReadableString());
   }
+  stats_.reset();
   // Remove the child policy's interested_parties pollset_set from the
   // xDS policy.
   grpc_pollset_set_del_pollset_set(child_policy_->interested_parties(),
@@ -1674,6 +1676,7 @@ void XdsLb::PriorityList::LocalityMap::Locality::ShutdownLocked() {
   }
   // Drop our ref to the child's picker, in case it's holding a ref to
   // the child.
+  load_reporting_picker_.reset();
   picker_wrapper_.reset();
   if (delayed_removal_timer_callback_pending_) {
     grpc_timer_cancel(&delayed_removal_timer_);
@@ -1681,19 +1684,19 @@ void XdsLb::PriorityList::LocalityMap::Locality::ShutdownLocked() {
   shutdown_ = true;
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::ResetBackoffLocked() {
+void XdsLb::LocalityMap::Locality::ResetBackoffLocked() {
   child_policy_->ResetBackoffLocked();
   if (pending_child_policy_ != nullptr) {
     pending_child_policy_->ResetBackoffLocked();
   }
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::Orphan() {
+void XdsLb::LocalityMap::Locality::Orphan() {
   ShutdownLocked();
   Unref();
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::DeactivateLocked() {
+void XdsLb::LocalityMap::Locality::DeactivateLocked() {
   // If already deactivated, don't do that again.
   if (weight_ == 0) return;
   // Set the locality weight to 0 so that future xds picker won't contain this
@@ -1710,8 +1713,8 @@ void XdsLb::PriorityList::LocalityMap::Locality::DeactivateLocked() {
   delayed_removal_timer_callback_pending_ = true;
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::OnDelayedRemovalTimer(
-    void* arg, grpc_error* error) {
+void XdsLb::LocalityMap::Locality::OnDelayedRemovalTimer(void* arg,
+                                                         grpc_error* error) {
   Locality* self = static_cast<Locality*>(arg);
   self->xds_policy()->combiner()->Run(
       GRPC_CLOSURE_INIT(&self->on_delayed_removal_timer_,
@@ -1719,7 +1722,7 @@ void XdsLb::PriorityList::LocalityMap::Locality::OnDelayedRemovalTimer(
       GRPC_ERROR_REF(error));
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::OnDelayedRemovalTimerLocked(
+void XdsLb::LocalityMap::Locality::OnDelayedRemovalTimerLocked(
     void* arg, grpc_error* error) {
   Locality* self = static_cast<Locality*>(arg);
   self->delayed_removal_timer_callback_pending_ = false;
@@ -1730,23 +1733,21 @@ void XdsLb::PriorityList::LocalityMap::Locality::OnDelayedRemovalTimerLocked(
 }
 
 //
-// XdsLb::Locality::Helper
+// XdsLb::LocalityMap::Locality::Helper
 //
 
-bool XdsLb::PriorityList::LocalityMap::Locality::Helper::CalledByPendingChild()
-    const {
+bool XdsLb::LocalityMap::Locality::Helper::CalledByPendingChild() const {
   GPR_ASSERT(child_ != nullptr);
   return child_ == locality_->pending_child_policy_.get();
 }
 
-bool XdsLb::PriorityList::LocalityMap::Locality::Helper::CalledByCurrentChild()
-    const {
+bool XdsLb::LocalityMap::Locality::Helper::CalledByCurrentChild() const {
   GPR_ASSERT(child_ != nullptr);
   return child_ == locality_->child_policy_.get();
 }
 
 RefCountedPtr<SubchannelInterface>
-XdsLb::PriorityList::LocalityMap::Locality::Helper::CreateSubchannel(
+XdsLb::LocalityMap::Locality::Helper::CreateSubchannel(
     const grpc_channel_args& args) {
   if (locality_->xds_policy()->shutting_down_ ||
       (!CalledByPendingChild() && !CalledByCurrentChild())) {
@@ -1756,7 +1757,7 @@ XdsLb::PriorityList::LocalityMap::Locality::Helper::CreateSubchannel(
       args);
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::Helper::UpdateState(
+void XdsLb::LocalityMap::Locality::Helper::UpdateState(
     grpc_connectivity_state state, std::unique_ptr<SubchannelPicker> picker) {
   if (locality_->xds_policy()->shutting_down_) return;
   // If this request is from the pending child policy, ignore it until
@@ -1778,22 +1779,16 @@ void XdsLb::PriorityList::LocalityMap::Locality::Helper::UpdateState(
     // This request is from an outdated child, so ignore it.
     return;
   }
-  // Cache the picker and its state in the locality.
-  // TODO(roth): If load reporting is not configured, we should ideally
-  // pass a null LocalityStats ref to the EndpointPickerWrapper and have it
-  // not collect any stats, since they're not going to be used.  This would
-  // require recreating all of the pickers whenever we get a config update.
-  locality_->picker_wrapper_ = MakeRefCounted<EndpointPickerWrapper>(
-      std::move(picker),
-      locality_->xds_policy()->client_stats_.FindLocalityStats(
-          locality_->name_));
+  // Cache the state and picker in the locality.
   locality_->connectivity_state_ = state;
+  locality_->picker_wrapper_ =
+      MakeRefCounted<RefCountedEndpointPicker>(std::move(picker));
   // Notify the locality map.
   locality_->locality_map_->OnLocalityStateUpdateLocked();
 }
 
-void XdsLb::PriorityList::LocalityMap::Locality::Helper::AddTraceEvent(
-    TraceSeverity severity, StringView message) {
+void XdsLb::LocalityMap::Locality::Helper::AddTraceEvent(TraceSeverity severity,
+                                                         StringView message) {
   if (locality_->xds_policy()->shutting_down_ ||
       (!CalledByPendingChild() && !CalledByCurrentChild())) {
     return;
@@ -1816,9 +1811,9 @@ class XdsFactory : public LoadBalancingPolicyFactory {
   const char* name() const override { return kXds; }
 
   RefCountedPtr<LoadBalancingPolicy::Config> ParseLoadBalancingConfig(
-      const grpc_json* json, grpc_error** error) const override {
+      const Json& json, grpc_error** error) const override {
     GPR_DEBUG_ASSERT(error != nullptr && *error == GRPC_ERROR_NONE);
-    if (json == nullptr) {
+    if (json.type() == Json::Type::JSON_NULL) {
       // xds was mentioned as a policy in the deprecated loadBalancingPolicy
       // field or in the client API.
       *error = GRPC_ERROR_CREATE_FROM_STATIC_STRING(
@@ -1826,70 +1821,66 @@ class XdsFactory : public LoadBalancingPolicyFactory {
           "Please use loadBalancingConfig field of service config instead.");
       return nullptr;
     }
-    GPR_DEBUG_ASSERT(strcmp(json->key, name()) == 0);
-    InlinedVector<grpc_error*, 3> error_list;
+    std::vector<grpc_error*> error_list;
+    // Child policy.
     RefCountedPtr<LoadBalancingPolicy::Config> child_policy;
+    auto it = json.object_value().find("childPolicy");
+    if (it != json.object_value().end()) {
+      grpc_error* parse_error = GRPC_ERROR_NONE;
+      child_policy = LoadBalancingPolicyRegistry::ParseLoadBalancingConfig(
+          it->second, &parse_error);
+      if (child_policy == nullptr) {
+        GPR_DEBUG_ASSERT(parse_error != GRPC_ERROR_NONE);
+        std::vector<grpc_error*> child_errors;
+        child_errors.push_back(parse_error);
+        error_list.push_back(
+            GRPC_ERROR_CREATE_FROM_VECTOR("field:childPolicy", &child_errors));
+      }
+    }
+    // Fallback policy.
     RefCountedPtr<LoadBalancingPolicy::Config> fallback_policy;
+    it = json.object_value().find("fallbackPolicy");
+    if (it != json.object_value().end()) {
+      grpc_error* parse_error = GRPC_ERROR_NONE;
+      fallback_policy = LoadBalancingPolicyRegistry::ParseLoadBalancingConfig(
+          it->second, &parse_error);
+      if (fallback_policy == nullptr) {
+        GPR_DEBUG_ASSERT(parse_error != GRPC_ERROR_NONE);
+        std::vector<grpc_error*> child_errors;
+        child_errors.push_back(parse_error);
+        error_list.push_back(GRPC_ERROR_CREATE_FROM_VECTOR(
+            "field:fallbackPolicy", &child_errors));
+      }
+    }
+    // EDS service name.
     const char* eds_service_name = nullptr;
+    it = json.object_value().find("edsServiceName");
+    if (it != json.object_value().end()) {
+      if (it->second.type() != Json::Type::STRING) {
+        error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+            "field:edsServiceName error:type should be string"));
+      } else {
+        eds_service_name = it->second.string_value().c_str();
+      }
+    }
+    // LRS load reporting server name.
     const char* lrs_load_reporting_server_name = nullptr;
-    for (const grpc_json* field = json->child; field != nullptr;
-         field = field->next) {
-      if (field->key == nullptr) continue;
-      if (strcmp(field->key, "childPolicy") == 0) {
-        if (child_policy != nullptr) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:childPolicy error:Duplicate entry"));
-        }
-        grpc_error* parse_error = GRPC_ERROR_NONE;
-        child_policy = LoadBalancingPolicyRegistry::ParseLoadBalancingConfig(
-            field, &parse_error);
-        if (child_policy == nullptr) {
-          GPR_DEBUG_ASSERT(parse_error != GRPC_ERROR_NONE);
-          error_list.push_back(parse_error);
-        }
-      } else if (strcmp(field->key, "fallbackPolicy") == 0) {
-        if (fallback_policy != nullptr) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:fallbackPolicy error:Duplicate entry"));
-        }
-        grpc_error* parse_error = GRPC_ERROR_NONE;
-        fallback_policy = LoadBalancingPolicyRegistry::ParseLoadBalancingConfig(
-            field, &parse_error);
-        if (fallback_policy == nullptr) {
-          GPR_DEBUG_ASSERT(parse_error != GRPC_ERROR_NONE);
-          error_list.push_back(parse_error);
-        }
-      } else if (strcmp(field->key, "edsServiceName") == 0) {
-        if (eds_service_name != nullptr) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:edsServiceName error:Duplicate entry"));
-        }
-        if (field->type != GRPC_JSON_STRING) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:edsServiceName error:type should be string"));
-          continue;
-        }
-        eds_service_name = field->value;
-      } else if (strcmp(field->key, "lrsLoadReportingServerName") == 0) {
-        if (lrs_load_reporting_server_name != nullptr) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:lrsLoadReportingServerName error:Duplicate entry"));
-        }
-        if (field->type != GRPC_JSON_STRING) {
-          error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "field:lrsLoadReportingServerName error:type should be string"));
-          continue;
-        }
-        lrs_load_reporting_server_name = field->value;
+    it = json.object_value().find("lrsLoadReportingServerName");
+    if (it != json.object_value().end()) {
+      if (it->second.type() != Json::Type::STRING) {
+        error_list.push_back(GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+            "field:lrsLoadReportingServerName error:type should be string"));
+      } else {
+        lrs_load_reporting_server_name = it->second.string_value().c_str();
       }
     }
     if (error_list.empty()) {
       Optional<std::string> optional_lrs_load_reporting_server_name;
       if (lrs_load_reporting_server_name != nullptr) {
-        optional_lrs_load_reporting_server_name.set(
+        optional_lrs_load_reporting_server_name.emplace(
             std::string(lrs_load_reporting_server_name));
       }
-      return MakeRefCounted<ParsedXdsConfig>(
+      return MakeRefCounted<XdsConfig>(
           std::move(child_policy), std::move(fallback_policy),
           eds_service_name == nullptr ? "" : eds_service_name,
           std::move(optional_lrs_load_reporting_server_name));
@@ -1911,7 +1902,7 @@ class XdsFactory : public LoadBalancingPolicyFactory {
 void grpc_lb_policy_xds_init() {
   grpc_core::LoadBalancingPolicyRegistry::Builder::
       RegisterLoadBalancingPolicyFactory(
-          grpc_core::MakeUnique<grpc_core::XdsFactory>());
+          absl::make_unique<grpc_core::XdsFactory>());
 }
 
 void grpc_lb_policy_xds_shutdown() {}