grpc 1.58.0 → 1.60.0

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

@@ -18,6 +18,7 @@
 
 #include "src/core/ext/filters/client_channel/lb_policy/xds/xds_override_host.h"
 
+#include <inttypes.h>
 #include <stddef.h>
 
 #include <algorithm>
@@ -28,6 +29,7 @@
 #include <set>
 #include <string>
 #include <tuple>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -35,8 +37,11 @@
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
+#include "absl/types/span.h"
 #include "absl/types/variant.h"
 
 #include <grpc/impl/connectivity_state.h>
@@ -50,10 +55,12 @@
 #include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/config/core_configuration.h"
 #include "src/core/lib/debug/trace.h"
+#include "src/core/lib/experiments/experiments.h"
 #include "src/core/lib/gprpp/debug_location.h"
 #include "src/core/lib/gprpp/match.h"
 #include "src/core/lib/gprpp/orphanable.h"
 #include "src/core/lib/gprpp/ref_counted_ptr.h"
+#include "src/core/lib/gprpp/ref_counted_string.h"
 #include "src/core/lib/gprpp/sync.h"
 #include "src/core/lib/gprpp/validation_errors.h"
 #include "src/core/lib/gprpp/work_serializer.h"
@@ -62,6 +69,7 @@
 #include "src/core/lib/iomgr/exec_ctx.h"
 #include "src/core/lib/iomgr/iomgr_fwd.h"
 #include "src/core/lib/iomgr/pollset_set.h"
+#include "src/core/lib/iomgr/resolved_address.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"
@@ -70,7 +78,7 @@
 #include "src/core/lib/load_balancing/lb_policy_factory.h"
 #include "src/core/lib/load_balancing/lb_policy_registry.h"
 #include "src/core/lib/load_balancing/subchannel_interface.h"
-#include "src/core/lib/resolver/server_address.h"
+#include "src/core/lib/resolver/endpoint_addresses.h"
 #include "src/core/lib/transport/connectivity_state.h"
 
 namespace grpc_core {
@@ -94,9 +102,9 @@ struct PtrLessThan {
   }
 };
 
-XdsHealthStatus GetAddressHealthStatus(const ServerAddress& address) {
+XdsHealthStatus GetEndpointHealthStatus(const EndpointAddresses& endpoint) {
   return XdsHealthStatus(static_cast<XdsHealthStatus::HealthStatus>(
-      address.args()
+      endpoint.args()
           .GetInt(GRPC_ARG_XDS_HEALTH_STATUS)
           .value_or(XdsHealthStatus::HealthStatus::kUnknown)));
 }
@@ -120,8 +128,7 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
   class SubchannelWrapper : public DelegatingSubchannel {
    public:
     SubchannelWrapper(RefCountedPtr<SubchannelInterface> subchannel,
-                      RefCountedPtr<XdsOverrideHostLb> policy,
-                      absl::string_view key);
+                      RefCountedPtr<XdsOverrideHostLb> policy);
 
     ~SubchannelWrapper() override;
 
@@ -137,6 +144,9 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
 
     XdsOverrideHostLb* policy() { return policy_.get(); }
 
+    void set_key(absl::string_view key) { key_ = std::string(key); }
+    const absl::optional<std::string>& key() const { return key_; }
+
    private:
     class ConnectivityStateWatcher : public ConnectivityStateWatcherInterface {
      public:
@@ -168,6 +178,60 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
         GRPC_CHANNEL_IDLE};
   };
 
+  class SubchannelEntry {
+   public:
+    explicit SubchannelEntry(XdsHealthStatus eds_health_status)
+        : eds_health_status_(eds_health_status) {}
+
+    void SetSubchannel(SubchannelWrapper* subchannel) {
+      if (eds_health_status_.status() == XdsHealthStatus::kDraining) {
+        subchannel_ = subchannel->Ref();
+      } else {
+        subchannel_ = subchannel->WeakRef();
+      }
+    }
+
+    void UnsetSubchannel() {
+      subchannel_ = WeakRefCountedPtr<SubchannelWrapper>(nullptr);
+    }
+
+    SubchannelWrapper* GetSubchannel() const {
+      return Match(
+          subchannel_,
+          [](WeakRefCountedPtr<XdsOverrideHostLb::SubchannelWrapper>
+                 subchannel) { return subchannel.get(); },
+          [](RefCountedPtr<XdsOverrideHostLb::SubchannelWrapper> subchannel) {
+            return subchannel.get();
+          });
+    }
+
+    void SetEdsHealthStatus(XdsHealthStatus eds_health_status) {
+      eds_health_status_ = eds_health_status;
+      auto subchannel = GetSubchannel();
+      if (subchannel == nullptr) return;
+      if (eds_health_status_.status() == XdsHealthStatus::kDraining) {
+        subchannel_ = subchannel->Ref();
+      } else {
+        subchannel_ = subchannel->WeakRef();
+      }
+    }
+
+    XdsHealthStatus eds_health_status() const { return eds_health_status_; }
+
+    void set_address_list(RefCountedStringValue address_list) {
+      address_list_ = std::move(address_list);
+    }
+
+    RefCountedStringValue address_list() const { return address_list_; }
+
+   private:
+    absl::variant<WeakRefCountedPtr<SubchannelWrapper>,
+                  RefCountedPtr<SubchannelWrapper>>
+        subchannel_;
+    XdsHealthStatus eds_health_status_;
+    RefCountedStringValue address_list_;
+  };
+
   // A picker that wraps the picker from the child for cases when cookie is
   // present.
   class Picker : public SubchannelPicker {
@@ -206,7 +270,7 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
     };
 
     absl::optional<LoadBalancingPolicy::PickResult> PickOverridenHost(
-        absl::string_view override_host);
+        XdsOverrideHostAttribute* override_host_attr) const;
 
     RefCountedPtr<XdsOverrideHostLb> policy_;
     RefCountedPtr<SubchannelPicker> picker_;
@@ -221,58 +285,12 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
               std::move(xds_override_host_policy)) {}
 
     RefCountedPtr<SubchannelInterface> CreateSubchannel(
-        ServerAddress address, const ChannelArgs& args) override;
+        const grpc_resolved_address& address,
+        const ChannelArgs& per_address_args, const ChannelArgs& args) override;
     void UpdateState(grpc_connectivity_state state, const absl::Status& status,
                      RefCountedPtr<SubchannelPicker> picker) override;
   };
 
-  class SubchannelEntry {
-   public:
-    explicit SubchannelEntry(XdsHealthStatus eds_health_status)
-        : eds_health_status_(eds_health_status) {}
-
-    void SetSubchannel(SubchannelWrapper* subchannel) {
-      if (eds_health_status_.status() == XdsHealthStatus::kDraining) {
-        subchannel_ = subchannel->Ref();
-      } else {
-        subchannel_ = subchannel;
-      }
-    }
-
-    void UnsetSubchannel() { subchannel_ = nullptr; }
-
-    SubchannelWrapper* GetSubchannel() const {
-      return Match(
-          subchannel_,
-          [](XdsOverrideHostLb::SubchannelWrapper* subchannel) {
-            return subchannel;
-          },
-          [](RefCountedPtr<XdsOverrideHostLb::SubchannelWrapper> subchannel) {
-            return subchannel.get();
-          });
-    }
-
-    void SetEdsHealthStatus(XdsHealthStatus eds_health_status) {
-      eds_health_status_ = eds_health_status;
-      auto subchannel = GetSubchannel();
-      if (subchannel == nullptr) {
-        return;
-      }
-      if (eds_health_status_.status() == XdsHealthStatus::kDraining) {
-        subchannel_ = subchannel->Ref();
-      } else {
-        subchannel_ = subchannel;
-      }
-    }
-
-    XdsHealthStatus eds_health_status() const { return eds_health_status_; }
-
-   private:
-    absl::variant<SubchannelWrapper*, RefCountedPtr<SubchannelWrapper>>
-        subchannel_;
-    XdsHealthStatus eds_health_status_;
-  };
-
   ~XdsOverrideHostLb() override;
 
   void ShutdownLocked() override;
@@ -282,20 +300,18 @@ class XdsOverrideHostLb : public LoadBalancingPolicy {
 
   void MaybeUpdatePickerLocked();
 
-  absl::StatusOr<ServerAddressList> UpdateAddressMap(
-      absl::StatusOr<ServerAddressList> addresses);
+  absl::StatusOr<EndpointAddressesList> UpdateAddressMap(
+      absl::StatusOr<EndpointAddressesList> endpoints);
 
   RefCountedPtr<SubchannelWrapper> AdoptSubchannel(
-      ServerAddress address, RefCountedPtr<SubchannelInterface> subchannel);
+      const grpc_resolved_address& address,
+      RefCountedPtr<SubchannelInterface> subchannel);
 
   void UnsetSubchannel(absl::string_view key, SubchannelWrapper* subchannel);
 
-  RefCountedPtr<SubchannelWrapper> GetSubchannelByAddress(
-      absl::string_view address, XdsHealthStatusSet overriden_health_statuses);
-
   void OnSubchannelConnectivityStateChange(absl::string_view subchannel_key)
-      ABSL_NO_THREAD_SAFETY_ANALYSIS;  // Called from within the worker
-                                       // serializer and does not require
+      ABSL_NO_THREAD_SAFETY_ANALYSIS;  // Called from within the
+                                       // WorkSerializer and does not require
                                        // additional synchronization
 
   // Current config from the resolver.
@@ -333,39 +349,89 @@ XdsOverrideHostLb::Picker::Picker(
 }
 
 absl::optional<LoadBalancingPolicy::PickResult>
-XdsOverrideHostLb::Picker::PickOverridenHost(absl::string_view override_host) {
-  if (override_host.length() == 0) {
-    return absl::nullopt;
-  }
-  auto subchannel = policy_->GetSubchannelByAddress(
-      override_host, override_host_health_status_set_);
-  if (subchannel == nullptr) {
-    return absl::nullopt;
-  }
-  auto connectivity_state = subchannel->connectivity_state();
-  if (connectivity_state == GRPC_CHANNEL_READY) {
-    return PickResult::Complete(subchannel->wrapped_subchannel());
-  } else if (connectivity_state == GRPC_CHANNEL_CONNECTING) {
+XdsOverrideHostLb::Picker::PickOverridenHost(
+    XdsOverrideHostAttribute* override_host_attr) const {
+  GPR_ASSERT(override_host_attr != nullptr);
+  auto cookie_address_list = override_host_attr->cookie_address_list();
+  if (cookie_address_list.empty()) return absl::nullopt;
+  // The cookie has an address list, so look through the addresses in order.
+  RefCountedPtr<SubchannelWrapper> idle_subchannel;
+  bool found_connecting = false;
+  {
+    MutexLock lock(&policy_->subchannel_map_mu_);
+    for (absl::string_view address : absl::StrSplit(cookie_address_list, ',')) {
+      RefCountedPtr<SubchannelWrapper> subchannel;
+      auto it = policy_->subchannel_map_.find(address);
+      if (it != policy_->subchannel_map_.end()) {
+        subchannel = it->second.GetSubchannel()->RefIfNonZero();
+      }
+      if (subchannel == nullptr) {
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO, "Subchannel %s was not found",
+                  std::string(address).c_str());
+        }
+        continue;
+      }
+      if (!override_host_health_status_set_.Contains(
+              it->second.eds_health_status())) {
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO,
+                  "Subchannel %s health status is not overridden (%s)",
+                  std::string(address).c_str(),
+                  it->second.eds_health_status().ToString());
+        }
+        continue;
+      }
+      auto connectivity_state = subchannel->connectivity_state();
+      if (connectivity_state == GRPC_CHANNEL_READY) {
+        // Found a READY subchannel.  Pass back the actual address list
+        // and return the subchannel.
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO, "Picker override found READY subchannel %s",
+                  std::string(address).c_str());
+        }
+        override_host_attr->set_actual_address_list(it->second.address_list());
+        return PickResult::Complete(subchannel->wrapped_subchannel());
+      } else if (connectivity_state == GRPC_CHANNEL_IDLE) {
+        if (idle_subchannel == nullptr) idle_subchannel = std::move(subchannel);
+      } else if (connectivity_state == GRPC_CHANNEL_CONNECTING) {
+        found_connecting = true;
+      }
+    }
+  }
+  // No READY subchannel found.  If we found an IDLE subchannel, trigger
+  // a connection attempt and queue the pick until that attempt completes.
+  if (idle_subchannel != nullptr) {
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+      gpr_log(GPR_INFO, "Picker override found IDLE subchannel");
+    }
+    // Deletes itself after the connection is requested.
+    new SubchannelConnectionRequester(std::move(idle_subchannel));
     return PickResult::Queue();
-  } else if (connectivity_state == GRPC_CHANNEL_IDLE) {
-    // Deleted after the connection is requested
-    new SubchannelConnectionRequester(std::move(subchannel));
+  }
+  // No READY or IDLE subchannels.  If we found a CONNECTING subchannel,
+  // queue the pick and wait for the connection attempt to complete.
+  if (found_connecting) {
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+      gpr_log(GPR_INFO, "Picker override found CONNECTING subchannel");
+    }
     return PickResult::Queue();
   }
+  // No READY, IDLE, or CONNECTING subchannels found.
   return absl::nullopt;
 }
 
-LoadBalancingPolicy::PickResult XdsOverrideHostLb::Picker::Pick(
-    LoadBalancingPolicy::PickArgs args) {
+LoadBalancingPolicy::PickResult XdsOverrideHostLb::Picker::Pick(PickArgs args) {
   auto* call_state = static_cast<ClientChannelLbCallState*>(args.call_state);
-  auto* override_host = static_cast<XdsOverrideHostAttribute*>(
+  auto* override_host_attr = static_cast<XdsOverrideHostAttribute*>(
       call_state->GetCallAttribute(XdsOverrideHostAttribute::TypeName()));
-  auto overridden_host_pick =
-      PickOverridenHost(override_host != nullptr ? override_host->host_name()
-                                                 : absl::string_view());
-  if (overridden_host_pick.has_value()) {
-    return std::move(*overridden_host_pick);
+  if (override_host_attr != nullptr) {
+    auto overridden_host_pick = PickOverridenHost(override_host_attr);
+    if (overridden_host_pick.has_value()) {
+      return std::move(*overridden_host_pick);
+    }
   }
+  // No usable override.  Delegate to child picker.
   if (picker_ == nullptr) {  // Should never happen.
     return PickResult::Fail(absl::InternalError(
         "xds_override_host picker not given any child picker"));
@@ -373,9 +439,23 @@ LoadBalancingPolicy::PickResult XdsOverrideHostLb::Picker::Pick(
   auto result = picker_->Pick(args);
   auto complete_pick = absl::get_if<PickResult::Complete>(&result.result);
   if (complete_pick != nullptr) {
-    complete_pick->subchannel =
-        static_cast<SubchannelWrapper*>(complete_pick->subchannel.get())
-            ->wrapped_subchannel();
+    auto* wrapper =
+        static_cast<SubchannelWrapper*>(complete_pick->subchannel.get());
+    // Populate the address list in the override host attribute so that
+    // the StatefulSession filter can set the cookie.
+    if (override_host_attr != nullptr) {
+      auto& key = wrapper->key();
+      if (key.has_value()) {
+        MutexLock lock(&policy_->subchannel_map_mu_);
+        auto it = policy_->subchannel_map_.find(*key);
+        if (it != policy_->subchannel_map_.end()) {  // Should always be true.
+          override_host_attr->set_actual_address_list(
+              it->second.address_list());
+        }
+      }
+    }
+    // Unwrap the subchannel.
+    complete_pick->subchannel = wrapper->wrapped_subchannel();
   }
   return result;
 }
@@ -430,7 +510,10 @@ void XdsOverrideHostLb::ResetBackoffLocked() {
 
 absl::Status XdsOverrideHostLb::UpdateLocked(UpdateArgs args) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
-    gpr_log(GPR_INFO, "[xds_override_host_lb %p] Received update", this);
+    gpr_log(GPR_INFO,
+            "[xds_override_host_lb %p] Received update with %" PRIuPTR
+            " addresses",
+            this, args.addresses.ok() ? args.addresses->size() : 0);
   }
   auto old_config = std::move(config_);
   // Update config.
@@ -495,62 +578,133 @@ OrphanablePtr<LoadBalancingPolicy> XdsOverrideHostLb::CreateChildPolicyLocked(
   return lb_policy;
 }
 
-absl::StatusOr<ServerAddressList> XdsOverrideHostLb::UpdateAddressMap(
-    absl::StatusOr<ServerAddressList> addresses) {
-  if (!addresses.ok()) {
-    return addresses;
-  }
-  ServerAddressList return_value;
-  std::map<const std::string, XdsHealthStatus> addresses_for_map;
-  for (const auto& address : *addresses) {
-    XdsHealthStatus status = GetAddressHealthStatus(address);
+absl::StatusOr<EndpointAddressesList> XdsOverrideHostLb::UpdateAddressMap(
+    absl::StatusOr<EndpointAddressesList> endpoints) {
+  if (!endpoints.ok()) {
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+      gpr_log(GPR_INFO, "[xds_override_host_lb %p] address error: %s", this,
+              endpoints.status().ToString().c_str());
+    }
+    return endpoints;
+  }
+  // Construct the list of addresses to pass to the child policy and a
+  // map of address info from which to update subchannel_map_.
+  EndpointAddressesList child_addresses;
+  struct AddressInfo {
+    XdsHealthStatus eds_health_status;
+    RefCountedStringValue address_list;
+    AddressInfo(XdsHealthStatus status, RefCountedStringValue addresses)
+        : eds_health_status(status), address_list(std::move(addresses)) {}
+  };
+  std::map<const std::string, AddressInfo> addresses_for_map;
+  for (const auto& endpoint : *endpoints) {
+    XdsHealthStatus status = GetEndpointHealthStatus(endpoint);
     if (status.status() != XdsHealthStatus::kDraining) {
-      return_value.push_back(address);
+      if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+        gpr_log(GPR_INFO,
+                "[xds_override_host_lb %p] endpoint %s: not draining, "
+                "passing to child",
+                this, endpoint.ToString().c_str());
+      }
+      child_addresses.push_back(endpoint);
     } else if (!config_->override_host_status_set().Contains(status)) {
       // Skip draining hosts if not in the override status set.
+      if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+        gpr_log(GPR_INFO,
+                "[xds_override_host_lb %p] endpoint %s: draining but not in "
+                "override_host_status set -- ignoring",
+                this, endpoint.ToString().c_str());
+      }
       continue;
     }
-    auto key = grpc_sockaddr_to_uri(&address.address());
-    if (key.ok()) {
-      addresses_for_map.emplace(std::move(*key), status);
+    std::vector<std::string> addresses;
+    addresses.reserve(endpoint.addresses().size());
+    for (const auto& address : endpoint.addresses()) {
+      auto key = grpc_sockaddr_to_string(&address, /*normalize=*/false);
+      if (key.ok()) {
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO,
+                  "[xds_override_host_lb %p] endpoint %s: adding map key %s",
+                  this, endpoint.ToString().c_str(), key->c_str());
+        }
+        addresses.push_back(*std::move(key));
+      }
+    }
+    absl::Span<const std::string> addresses_span = addresses;
+    for (size_t i = 0; i < addresses.size(); ++i) {
+      std::string start = absl::StrJoin(addresses_span.subspan(0, i), ",");
+      std::string end = absl::StrJoin(addresses_span.subspan(i + 1), ",");
+      RefCountedStringValue address_list(
+          absl::StrCat(addresses[i], (start.empty() ? "" : ","), start,
+                       (end.empty() ? "" : ","), end));
+      addresses_for_map.emplace(
+          std::piecewise_construct, std::forward_as_tuple(addresses[i]),
+          std::forward_as_tuple(status, std::move(address_list)));
     }
   }
+  // Now grab the lock and update subchannel_map_ from addresses_for_map.
   {
     MutexLock lock(&subchannel_map_mu_);
     for (auto it = subchannel_map_.begin(); it != subchannel_map_.end();) {
       if (addresses_for_map.find(it->first) == addresses_for_map.end()) {
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO, "[xds_override_host_lb %p] removing map key %s",
+                  this, it->first.c_str());
+        }
         it = subchannel_map_.erase(it);
       } else {
         ++it;
       }
     }
-    for (const auto& key_status : addresses_for_map) {
-      auto it = subchannel_map_.find(key_status.first);
+    for (auto& p : addresses_for_map) {
+      const auto& address = p.first;
+      auto& address_info = p.second;
+      auto it = subchannel_map_.find(address);
       if (it == subchannel_map_.end()) {
-        subchannel_map_.emplace(std::piecewise_construct,
-                                std::forward_as_tuple(key_status.first),
-                                std::forward_as_tuple(key_status.second));
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO, "[xds_override_host_lb %p] adding map key %s", this,
+                  address.c_str());
+        }
+        it = subchannel_map_
+                 .emplace(std::piecewise_construct,
+                          std::forward_as_tuple(address),
+                          std::forward_as_tuple(address_info.eds_health_status))
+                 .first;
       } else {
-        it->second.SetEdsHealthStatus(key_status.second);
+        if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+          gpr_log(GPR_INFO,
+                  "[xds_override_host_lb %p] setting EDS health status for "
+                  "%s to %s",
+                  this, address.c_str(),
+                  address_info.eds_health_status.ToString());
+        }
+        it->second.SetEdsHealthStatus(address_info.eds_health_status);
+      }
+      if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
+        gpr_log(GPR_INFO,
+                "[xds_override_host_lb %p] setting address list for %s to %s",
+                this, address.c_str(), address_info.address_list.c_str());
       }
+      it->second.set_address_list(std::move(address_info.address_list));
     }
   }
-  return return_value;
+  return child_addresses;
 }
 
 RefCountedPtr<XdsOverrideHostLb::SubchannelWrapper>
 XdsOverrideHostLb::AdoptSubchannel(
-    ServerAddress address, RefCountedPtr<SubchannelInterface> subchannel) {
-  auto key = grpc_sockaddr_to_uri(&address.address());
-  if (!key.ok()) {
-    return subchannel;
-  }
+    const grpc_resolved_address& address,
+    RefCountedPtr<SubchannelInterface> subchannel) {
+  auto key = grpc_sockaddr_to_string(&address, /*normalize=*/false);
   auto wrapper =
-      MakeRefCounted<SubchannelWrapper>(std::move(subchannel), Ref(), *key);
-  MutexLock lock(&subchannel_map_mu_);
-  auto it = subchannel_map_.find(*key);
-  if (it != subchannel_map_.end()) {
-    it->second.SetSubchannel(wrapper.get());
+      MakeRefCounted<SubchannelWrapper>(std::move(subchannel), Ref());
+  if (key.ok()) {
+    MutexLock lock(&subchannel_map_mu_);
+    auto it = subchannel_map_.find(*key);
+    if (it != subchannel_map_.end()) {
+      wrapper->set_key(*key);
+      it->second.SetSubchannel(wrapper.get());
+    }
   }
   return wrapper;
 }
@@ -566,29 +720,6 @@ void XdsOverrideHostLb::UnsetSubchannel(absl::string_view key,
   }
 }
 
-RefCountedPtr<XdsOverrideHostLb::SubchannelWrapper>
-XdsOverrideHostLb::GetSubchannelByAddress(
-    absl::string_view address, XdsHealthStatusSet overriden_health_statuses) {
-  MutexLock lock(&subchannel_map_mu_);
-  auto it = subchannel_map_.find(address);
-  if (it == subchannel_map_.end() || it->second.GetSubchannel() == nullptr) {
-    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
-      gpr_log(GPR_INFO, "Subchannel %s was not found",
-              std::string(address).c_str());
-    }
-    return nullptr;
-  }
-  if (!overriden_health_statuses.Contains(it->second.eds_health_status())) {
-    if (GRPC_TRACE_FLAG_ENABLED(grpc_lb_xds_override_host_trace)) {
-      gpr_log(GPR_INFO, "Subchannel %s health status is not overridden (%s)",
-              std::string(address).c_str(),
-              it->second.eds_health_status().ToString());
-    }
-    return nullptr;
-  }
-  return it->second.GetSubchannel()->Ref();
-}
-
 void XdsOverrideHostLb::OnSubchannelConnectivityStateChange(
     absl::string_view subchannel_key) {
   auto it = subchannel_map_.find(subchannel_key);
@@ -605,10 +736,11 @@ void XdsOverrideHostLb::OnSubchannelConnectivityStateChange(
 //
 
 RefCountedPtr<SubchannelInterface> XdsOverrideHostLb::Helper::CreateSubchannel(
-    ServerAddress address, const ChannelArgs& args) {
-  auto subchannel =
-      parent()->channel_control_helper()->CreateSubchannel(address, args);
-  return parent()->AdoptSubchannel(address, subchannel);
+    const grpc_resolved_address& address, const ChannelArgs& per_address_args,
+    const ChannelArgs& args) {
+  auto subchannel = parent()->channel_control_helper()->CreateSubchannel(
+      address, per_address_args, args);
+  return parent()->AdoptSubchannel(address, std::move(subchannel));
 }
 
 void XdsOverrideHostLb::Helper::UpdateState(
@@ -629,10 +761,8 @@ void XdsOverrideHostLb::Helper::UpdateState(
 
 XdsOverrideHostLb::SubchannelWrapper::SubchannelWrapper(
     RefCountedPtr<SubchannelInterface> subchannel,
-    RefCountedPtr<XdsOverrideHostLb> policy, absl::string_view key)
-    : DelegatingSubchannel(std::move(subchannel)),
-      key_(key),
-      policy_(std::move(policy)) {
+    RefCountedPtr<XdsOverrideHostLb> policy)
+    : DelegatingSubchannel(std::move(subchannel)), policy_(std::move(policy)) {
   auto watcher = std::make_unique<ConnectivityStateWatcher>(WeakRef());
   watcher_ = watcher.get();
   wrapped_subchannel()->WatchConnectivityState(std::move(watcher));
@@ -663,7 +793,8 @@ void XdsOverrideHostLb::SubchannelWrapper::UpdateConnectivityState(
   // Sending connectivity state notifications to the watchers may cause the set
   // of watchers to change, so we can't be iterating over the set of watchers
   // while we send the notifications
-  std::vector<ConnectivityStateWatcherInterface*> watchers(watchers_.size());
+  std::vector<ConnectivityStateWatcherInterface*> watchers;
+  watchers.reserve(watchers_.size());
   for (const auto& watcher : watchers_) {
     watchers.push_back(watcher.get());
   }
@@ -678,8 +809,19 @@ void XdsOverrideHostLb::SubchannelWrapper::UpdateConnectivityState(
 }
 
 void XdsOverrideHostLb::SubchannelWrapper::Orphan() {
-  key_.reset();
-  wrapped_subchannel()->CancelConnectivityStateWatch(watcher_);
+  if (!IsWorkSerializerDispatchEnabled()) {
+    key_.reset();
+    wrapped_subchannel()->CancelConnectivityStateWatch(watcher_);
+    return;
+  }
+  WeakRefCountedPtr<SubchannelWrapper> self = WeakRef();
+  policy_->work_serializer()->Run(
+      [self = std::move(self)]() {
+        self->key_.reset();
+        self->wrapped_subchannel()->CancelConnectivityStateWatch(
+            self->watcher_);
+      },
+      DEBUG_LOCATION);
 }
 
 grpc_pollset_set* XdsOverrideHostLb::SubchannelWrapper::

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

@@ -51,7 +51,7 @@
 #include "src/core/lib/load_balancing/lb_policy.h"
 #include "src/core/lib/load_balancing/lb_policy_factory.h"
 #include "src/core/lib/load_balancing/lb_policy_registry.h"
-#include "src/core/lib/resolver/server_address.h"
+#include "src/core/lib/resolver/endpoint_addresses.h"
 
 namespace grpc_core {