grpc 1.60.0 → 1.61.0.pre2

data/src/core/ext/xds/xds_route_config.cc

@@ -52,7 +52,7 @@
 #include "google/protobuf/wrappers.upb.h"
 #include "re2/re2.h"
 #include "upb/base/string_view.h"
-#include "upb/collections/map.h"
+#include "upb/message/map.h"
 #include "upb/text/encode.h"
 
 #include <grpc/status.h>
@@ -288,33 +288,44 @@ std::string XdsRouteConfigResource::Route::ToString() const {
   return absl::StrJoin(contents, "\n");
 }
 
+//
+// XdsRouteConfigResource::Route
+//
+
+std::string XdsRouteConfigResource::VirtualHost::ToString() const {
+  std::vector<std::string> parts;
+  parts.push_back(
+      absl::StrCat("vhost={\n"
+                   "  domains=[",
+                   absl::StrJoin(domains, ", "),
+                   "]\n"
+                   "  routes=[\n"));
+  for (const XdsRouteConfigResource::Route& route : routes) {
+    parts.push_back("    {\n");
+    parts.push_back(route.ToString());
+    parts.push_back("\n    }\n");
+  }
+  parts.push_back("  ]\n");
+  parts.push_back("  typed_per_filter_config={\n");
+  for (const auto& p : typed_per_filter_config) {
+    const std::string& name = p.first;
+    const auto& config = p.second;
+    parts.push_back(absl::StrCat("    ", name, "=", config.ToString(), "\n"));
+  }
+  parts.push_back("  }\n");
+  parts.push_back("}\n");
+  return absl::StrJoin(parts, "");
+}
+
 //
 // XdsRouteConfigResource
 //
 
 std::string XdsRouteConfigResource::ToString() const {
   std::vector<std::string> parts;
+  parts.reserve(virtual_hosts.size());
   for (const VirtualHost& vhost : virtual_hosts) {
-    parts.push_back(
-        absl::StrCat("vhost={\n"
-                     "  domains=[",
-                     absl::StrJoin(vhost.domains, ", "),
-                     "]\n"
-                     "  routes=[\n"));
-    for (const XdsRouteConfigResource::Route& route : vhost.routes) {
-      parts.push_back("    {\n");
-      parts.push_back(route.ToString());
-      parts.push_back("\n    }\n");
-    }
-    parts.push_back("  ]\n");
-    parts.push_back("  typed_per_filter_config={\n");
-    for (const auto& p : vhost.typed_per_filter_config) {
-      const std::string& name = p.first;
-      const auto& config = p.second;
-      parts.push_back(absl::StrCat("    ", name, "=", config.ToString(), "\n"));
-    }
-    parts.push_back("  }\n");
-    parts.push_back("]\n");
+    parts.push_back(vhost.ToString());
   }
   parts.push_back("cluster_specifier_plugins={\n");
   for (const auto& it : cluster_specifier_plugin_map) {
@@ -1137,7 +1148,8 @@ void MaybeLogRouteConfiguration(
     const upb_MessageDef* msg_type =
         envoy_config_route_v3_RouteConfiguration_getmsgdef(context.symtab);
     char buf[10240];
-    upb_TextEncode(route_config, msg_type, nullptr, 0, buf, sizeof(buf));
+    upb_TextEncode(reinterpret_cast<const upb_Message*>(route_config), msg_type,
+                   nullptr, 0, buf, sizeof(buf));
     gpr_log(GPR_DEBUG, "[xds_client %p] RouteConfiguration: %s", context.client,
             buf);
   }

data/src/core/ext/xds/xds_route_config.h

@@ -209,6 +209,7 @@ struct XdsRouteConfigResource : public XdsResourceType::ResourceData {
       return domains == other.domains && routes == other.routes &&
              typed_per_filter_config == other.typed_per_filter_config;
     }
+    std::string ToString() const;
   };
 
   std::vector<VirtualHost> virtual_hosts;

data/src/core/ext/xds/xds_server_config_fetcher.cc

@@ -52,6 +52,7 @@
 #include "src/core/ext/xds/xds_bootstrap_grpc.h"
 #include "src/core/ext/xds/xds_certificate_provider.h"
 #include "src/core/ext/xds/xds_channel_stack_modifier.h"
+#include "src/core/ext/xds/xds_client.h"
 #include "src/core/ext/xds/xds_client_grpc.h"
 #include "src/core/ext/xds/xds_common_types.h"
 #include "src/core/ext/xds/xds_http_filters.h"
@@ -91,6 +92,8 @@
 namespace grpc_core {
 namespace {
 
+using ReadDelayHandle = XdsClient::ReadDelayHandle;
+
 TraceFlag grpc_xds_server_config_fetcher_trace(false,
                                                "xds_server_config_fetcher");
 
@@ -151,11 +154,14 @@ class XdsServerConfigFetcher::ListenerWatcher
   }
 
   void OnResourceChanged(
-      std::shared_ptr<const XdsListenerResource> listener) override;
+      std::shared_ptr<const XdsListenerResource> listener,
+      RefCountedPtr<ReadDelayHandle> read_delay_handle) override;
 
-  void OnError(absl::Status status) override;
+  void OnError(absl::Status status,
+               RefCountedPtr<ReadDelayHandle> read_delay_handle) override;
 
-  void OnResourceDoesNotExist() override;
+  void OnResourceDoesNotExist(
+      RefCountedPtr<ReadDelayHandle> read_delay_handle) override;
 
   const std::string& listening_address() const { return listening_address_; }
 
@@ -225,15 +231,6 @@ class XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager
   }
 
  private:
-  struct CertificateProviders {
-    // We need to save our own refs to the root and instance certificate
-    // providers since the xds certificate provider just stores a ref to their
-    // distributors.
-    RefCountedPtr<grpc_tls_certificate_provider> root;
-    RefCountedPtr<grpc_tls_certificate_provider> instance;
-    RefCountedPtr<XdsCertificateProvider> xds;
-  };
-
   class RouteConfigWatcher;
   struct RdsUpdateState {
     RouteConfigWatcher* watcher;
@@ -271,7 +268,8 @@ class XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager
   size_t rds_resources_yet_to_fetch_ ABSL_GUARDED_BY(mu_) = 0;
   std::map<std::string /* resource_name */, RdsUpdateState> rds_map_
       ABSL_GUARDED_BY(mu_);
-  std::map<const XdsListenerResource::FilterChainData*, CertificateProviders>
+  std::map<const XdsListenerResource::FilterChainData*,
+           RefCountedPtr<XdsCertificateProvider>>
       certificate_providers_map_ ABSL_GUARDED_BY(mu_);
 };
 
@@ -292,16 +290,20 @@ class XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
         filter_chain_match_manager_(std::move(filter_chain_match_manager)) {}
 
   void OnResourceChanged(
-      std::shared_ptr<const XdsRouteConfigResource> route_config) override {
+      std::shared_ptr<const XdsRouteConfigResource> route_config,
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
     filter_chain_match_manager_->OnRouteConfigChanged(resource_name_,
                                                       std::move(route_config));
   }
 
-  void OnError(absl::Status status) override {
+  void OnError(
+      absl::Status status,
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
     filter_chain_match_manager_->OnError(resource_name_, status);
   }
 
-  void OnResourceDoesNotExist() override {
+  void OnResourceDoesNotExist(
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
     filter_chain_match_manager_->OnResourceDoesNotExist(resource_name_);
   }
 
@@ -488,13 +490,21 @@ class XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
       : parent_(std::move(parent)) {}
 
   void OnResourceChanged(
-      std::shared_ptr<const XdsRouteConfigResource> route_config) override {
+      std::shared_ptr<const XdsRouteConfigResource> route_config,
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
     parent_->OnRouteConfigChanged(std::move(route_config));
   }
 
-  void OnError(absl::Status status) override { parent_->OnError(status); }
+  void OnError(
+      absl::Status status,
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
+    parent_->OnError(status);
+  }
 
-  void OnResourceDoesNotExist() override { parent_->OnResourceDoesNotExist(); }
+  void OnResourceDoesNotExist(
+      RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) override {
+    parent_->OnResourceDoesNotExist();
+  }
 
  private:
   WeakRefCountedPtr<DynamicXdsServerConfigSelectorProvider> parent_;
@@ -527,7 +537,7 @@ void XdsServerConfigFetcher::StartWatch(
     std::unique_ptr<grpc_server_config_fetcher::WatcherInterface> watcher) {
   grpc_server_config_fetcher::WatcherInterface* watcher_ptr = watcher.get();
   auto listener_watcher = MakeRefCounted<ListenerWatcher>(
-      xds_client_->Ref(DEBUG_LOCATION, "ListenerWatcher"), std::move(watcher),
+      xds_client_.Ref(DEBUG_LOCATION, "ListenerWatcher"), std::move(watcher),
       serving_status_notifier_, listening_address);
   auto* listener_watcher_ptr = listener_watcher.get();
   XdsListenerResourceType::StartWatch(
@@ -574,7 +584,8 @@ XdsServerConfigFetcher::ListenerWatcher::ListenerWatcher(
       listening_address_(std::move(listening_address)) {}
 
 void XdsServerConfigFetcher::ListenerWatcher::OnResourceChanged(
-    std::shared_ptr<const XdsListenerResource> listener) {
+    std::shared_ptr<const XdsListenerResource> listener,
+    RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) {
   if (GRPC_TRACE_FLAG_ENABLED(grpc_xds_server_config_fetcher_trace)) {
     gpr_log(GPR_INFO,
             "[ListenerWatcher %p] Received LDS update from xds client %p: %s",
@@ -595,7 +606,7 @@ void XdsServerConfigFetcher::ListenerWatcher::OnResourceChanged(
     return;
   }
   auto new_filter_chain_match_manager = MakeRefCounted<FilterChainMatchManager>(
-      xds_client_->Ref(DEBUG_LOCATION, "FilterChainMatchManager"),
+      xds_client_.Ref(DEBUG_LOCATION, "FilterChainMatchManager"),
       tcp_listener->filter_chain_map, tcp_listener->default_filter_chain);
   MutexLock lock(&mu_);
   if (filter_chain_match_manager_ == nullptr ||
@@ -605,11 +616,14 @@ void XdsServerConfigFetcher::ListenerWatcher::OnResourceChanged(
             filter_chain_match_manager_->default_filter_chain())) {
     pending_filter_chain_match_manager_ =
         std::move(new_filter_chain_match_manager);
-    pending_filter_chain_match_manager_->StartRdsWatch(Ref());
+    pending_filter_chain_match_manager_->StartRdsWatch(
+        RefAsSubclass<ListenerWatcher>());
   }
 }
 
-void XdsServerConfigFetcher::ListenerWatcher::OnError(absl::Status status) {
+void XdsServerConfigFetcher::ListenerWatcher::OnError(
+    absl::Status status,
+    RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) {
   MutexLock lock(&mu_);
   if (filter_chain_match_manager_ != nullptr ||
       pending_filter_chain_match_manager_ != nullptr) {
@@ -652,7 +666,8 @@ void XdsServerConfigFetcher::ListenerWatcher::OnFatalError(
   }
 }
 
-void XdsServerConfigFetcher::ListenerWatcher::OnResourceDoesNotExist() {
+void XdsServerConfigFetcher::ListenerWatcher::OnResourceDoesNotExist(
+    RefCountedPtr<ReadDelayHandle> /* read_delay_handle */) {
   MutexLock lock(&mu_);
   OnFatalError(absl::NotFoundError("Requested listener does not exist"));
 }
@@ -743,8 +758,8 @@ void XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
     MutexLock lock(&mu_);
     for (const auto& resource_name : resource_names) {
       ++rds_resources_yet_to_fetch_;
-      auto route_config_watcher =
-          MakeRefCounted<RouteConfigWatcher>(resource_name, WeakRef());
+      auto route_config_watcher = MakeRefCounted<RouteConfigWatcher>(
+          resource_name, WeakRefAsSubclass<FilterChainMatchManager>());
       rds_map_.emplace(resource_name, RdsUpdateState{route_config_watcher.get(),
                                                      absl::nullopt});
       watchers_to_start.push_back(
@@ -786,10 +801,7 @@ XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
         const XdsListenerResource::FilterChainData* filter_chain) {
   MutexLock lock(&mu_);
   auto it = certificate_providers_map_.find(filter_chain);
-  if (it != certificate_providers_map_.end()) {
-    return it->second.xds;
-  }
-  CertificateProviders certificate_providers;
+  if (it != certificate_providers_map_.end()) return it->second;
   // Configure root cert.
   absl::string_view root_provider_instance_name =
       filter_chain->downstream_tls_context.common_tls_context
@@ -799,11 +811,12 @@ XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
       filter_chain->downstream_tls_context.common_tls_context
           .certificate_validation_context.ca_certificate_provider_instance
           .certificate_name;
+  RefCountedPtr<grpc_tls_certificate_provider> root_cert_provider;
   if (!root_provider_instance_name.empty()) {
-    certificate_providers.root =
+    root_cert_provider =
         xds_client_->certificate_provider_store()
             .CreateOrGetCertificateProvider(root_provider_instance_name);
-    if (certificate_providers.root == nullptr) {
+    if (root_cert_provider == nullptr) {
       return absl::NotFoundError(
           absl::StrCat("Certificate provider instance name: \"",
                        root_provider_instance_name, "\" not recognized."));
@@ -816,33 +829,23 @@ XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
   absl::string_view identity_provider_cert_name =
       filter_chain->downstream_tls_context.common_tls_context
           .tls_certificate_provider_instance.certificate_name;
+  RefCountedPtr<grpc_tls_certificate_provider> identity_cert_provider;
   if (!identity_provider_instance_name.empty()) {
-    certificate_providers.instance =
+    identity_cert_provider =
         xds_client_->certificate_provider_store()
             .CreateOrGetCertificateProvider(identity_provider_instance_name);
-    if (certificate_providers.instance == nullptr) {
+    if (identity_cert_provider == nullptr) {
       return absl::NotFoundError(
           absl::StrCat("Certificate provider instance name: \"",
                        identity_provider_instance_name, "\" not recognized."));
     }
   }
-  certificate_providers.xds = MakeRefCounted<XdsCertificateProvider>();
-  certificate_providers.xds->UpdateRootCertNameAndDistributor(
-      "", root_provider_cert_name,
-      certificate_providers.root == nullptr
-          ? nullptr
-          : certificate_providers.root->distributor());
-  certificate_providers.xds->UpdateIdentityCertNameAndDistributor(
-      "", identity_provider_cert_name,
-      certificate_providers.instance == nullptr
-          ? nullptr
-          : certificate_providers.instance->distributor());
-  certificate_providers.xds->UpdateRequireClientCertificate(
-      "", filter_chain->downstream_tls_context.require_client_certificate);
-  auto xds_certificate_provider = certificate_providers.xds;
-  certificate_providers_map_.emplace(filter_chain,
-                                     std::move(certificate_providers));
-  return xds_certificate_provider;
+  auto xds_cert_provider = MakeRefCounted<XdsCertificateProvider>(
+      std::move(root_cert_provider), root_provider_cert_name,
+      std::move(identity_cert_provider), identity_provider_cert_name,
+      filter_chain->downstream_tls_context.require_client_certificate);
+  certificate_providers_map_.emplace(filter_chain, xds_cert_provider);
+  return xds_cert_provider;
 }
 
 void XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
@@ -1122,8 +1125,8 @@ absl::StatusOr<ChannelArgs> XdsServerConfigFetcher::ListenerWatcher::
         }
         server_config_selector_provider =
             MakeRefCounted<DynamicXdsServerConfigSelectorProvider>(
-                xds_client_->Ref(DEBUG_LOCATION,
-                                 "DynamicXdsServerConfigSelectorProvider"),
+                xds_client_.Ref(DEBUG_LOCATION,
+                                "DynamicXdsServerConfigSelectorProvider"),
                 rds_name, std::move(initial_resource),
                 filter_chain->http_connection_manager.http_filters);
       },
@@ -1131,8 +1134,8 @@ absl::StatusOr<ChannelArgs> XdsServerConfigFetcher::ListenerWatcher::
       [&](const std::shared_ptr<const XdsRouteConfigResource>& route_config) {
         server_config_selector_provider =
             MakeRefCounted<StaticXdsServerConfigSelectorProvider>(
-                xds_client_->Ref(DEBUG_LOCATION,
-                                 "StaticXdsServerConfigSelectorProvider"),
+                xds_client_.Ref(DEBUG_LOCATION,
+                                "StaticXdsServerConfigSelectorProvider"),
                 route_config,
                 filter_chain->http_connection_manager.http_filters);
       });
@@ -1270,7 +1273,8 @@ XdsServerConfigFetcher::ListenerWatcher::FilterChainMatchManager::
   // RouteConfigWatcher is being created here instead of in Watch() to avoid
   // deadlocks from invoking XdsRouteConfigResourceType::StartWatch whilst in a
   // critical region.
-  auto route_config_watcher = MakeRefCounted<RouteConfigWatcher>(WeakRef());
+  auto route_config_watcher = MakeRefCounted<RouteConfigWatcher>(
+      WeakRefAsSubclass<DynamicXdsServerConfigSelectorProvider>());
   route_config_watcher_ = route_config_watcher.get();
   XdsRouteConfigResourceType::StartWatch(xds_client_.get(), resource_name_,
                                          std::move(route_config_watcher));

data/src/core/ext/xds/xds_transport.h

@@ -58,6 +58,9 @@ class XdsTransportFactory : public InternallyRefCounted<XdsTransportFactory> {
       // Only one message will be in flight at a time; subsequent
       // messages will not be sent until this one is done.
       virtual void SendMessage(std::string payload) = 0;
+
+      // Starts a recv_message operation on the stream.
+      virtual void StartRecvMessage() = 0;
     };
 
     // Create a streaming call on this transport for the specified method.

data/src/core/ext/xds/xds_transport_grpc.cc

@@ -86,39 +86,31 @@ GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::GrpcStreamingCall(
   GRPC_CLOSURE_INIT(&on_request_sent_, OnRequestSent, this, nullptr);
   // Start ops on the call.
   grpc_call_error call_error;
-  grpc_op ops[3];
+  grpc_op ops[2];
   memset(ops, 0, sizeof(ops));
-  // Send initial metadata.  No callback for this, since we don't really
-  // care when it finishes.
+  // Send initial metadata.
   grpc_op* op = ops;
   op->op = GRPC_OP_SEND_INITIAL_METADATA;
   op->data.send_initial_metadata.count = 0;
   op->flags = GRPC_INITIAL_METADATA_WAIT_FOR_READY |
               GRPC_INITIAL_METADATA_WAIT_FOR_READY_EXPLICITLY_SET;
   op->reserved = nullptr;
-  op++;
-  call_error = grpc_call_start_batch_and_execute(
-      call_, ops, static_cast<size_t>(op - ops), nullptr);
-  GPR_ASSERT(GRPC_CALL_OK == call_error);
-  // Start a batch with recv_initial_metadata and recv_message.
-  op = ops;
+  ++op;
   op->op = GRPC_OP_RECV_INITIAL_METADATA;
   op->data.recv_initial_metadata.recv_initial_metadata =
       &initial_metadata_recv_;
   op->flags = 0;
   op->reserved = nullptr;
-  op++;
-  op->op = GRPC_OP_RECV_MESSAGE;
-  op->data.recv_message.recv_message = &recv_message_payload_;
-  op->flags = 0;
-  op->reserved = nullptr;
-  op++;
-  Ref(DEBUG_LOCATION, "OnResponseReceived").release();
-  GRPC_CLOSURE_INIT(&on_response_received_, OnResponseReceived, this, nullptr);
+  ++op;
+  // Ref will be released in the callback
+  GRPC_CLOSURE_INIT(
+      &on_recv_initial_metadata_, OnRecvInitialMetadata,
+      this->Ref(DEBUG_LOCATION, "OnRecvInitialMetadata").release(), nullptr);
   call_error = grpc_call_start_batch_and_execute(
-      call_, ops, static_cast<size_t>(op - ops), &on_response_received_);
+      call_, ops, static_cast<size_t>(op - ops), &on_recv_initial_metadata_);
   GPR_ASSERT(GRPC_CALL_OK == call_error);
   // Start a batch for recv_trailing_metadata.
+  memset(ops, 0, sizeof(ops));
   op = ops;
   op->op = GRPC_OP_RECV_STATUS_ON_CLIENT;
   op->data.recv_status_on_client.trailing_metadata = &trailing_metadata_recv_;
@@ -126,7 +118,7 @@ GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::GrpcStreamingCall(
   op->data.recv_status_on_client.status_details = &status_details_;
   op->flags = 0;
   op->reserved = nullptr;
-  op++;
+  ++op;
   // This callback signals the end of the call, so it relies on the initial
   // ref instead of a new ref. When it's invoked, it's the initial ref that is
   // unreffed.
@@ -134,11 +126,11 @@ GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::GrpcStreamingCall(
   call_error = grpc_call_start_batch_and_execute(
       call_, ops, static_cast<size_t>(op - ops), &on_status_received_);
   GPR_ASSERT(GRPC_CALL_OK == call_error);
+  GRPC_CLOSURE_INIT(&on_response_received_, OnResponseReceived, this, nullptr);
 }
 
 GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
     ~GrpcStreamingCall() {
-  grpc_metadata_array_destroy(&initial_metadata_recv_);
   grpc_metadata_array_destroy(&trailing_metadata_recv_);
   grpc_byte_buffer_destroy(send_message_payload_);
   grpc_byte_buffer_destroy(recv_message_payload_);
@@ -175,55 +167,59 @@ void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::SendMessage(
   GPR_ASSERT(GRPC_CALL_OK == call_error);
 }
 
+void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
+    StartRecvMessage() {
+  Ref(DEBUG_LOCATION, "StartRecvMessage").release();
+  grpc_op op;
+  memset(&op, 0, sizeof(op));
+  op.op = GRPC_OP_RECV_MESSAGE;
+  op.data.recv_message.recv_message = &recv_message_payload_;
+  GPR_ASSERT(call_ != nullptr);
+  const grpc_call_error call_error =
+      grpc_call_start_batch_and_execute(call_, &op, 1, &on_response_received_);
+  GPR_ASSERT(GRPC_CALL_OK == call_error);
+}
+
+void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
+    OnRecvInitialMetadata(void* arg, grpc_error_handle /*error*/) {
+  RefCountedPtr<GrpcStreamingCall> self(static_cast<GrpcStreamingCall*>(arg));
+  grpc_metadata_array_destroy(&self->initial_metadata_recv_);
+}
+
 void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
     OnRequestSent(void* arg, grpc_error_handle error) {
-  auto* self = static_cast<GrpcStreamingCall*>(arg);
+  RefCountedPtr<GrpcStreamingCall> self(static_cast<GrpcStreamingCall*>(arg));
   // Clean up the sent message.
   grpc_byte_buffer_destroy(self->send_message_payload_);
   self->send_message_payload_ = nullptr;
   // Invoke request handler.
   self->event_handler_->OnRequestSent(error.ok());
-  // Drop the ref.
-  self->Unref(DEBUG_LOCATION, "OnRequestSent");
 }
 
 void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
     OnResponseReceived(void* arg, grpc_error_handle /*error*/) {
-  auto* self = static_cast<GrpcStreamingCall*>(arg);
+  RefCountedPtr<GrpcStreamingCall> self(static_cast<GrpcStreamingCall*>(arg));
   // If there was no payload, then we received status before we received
   // another message, so we stop reading.
-  if (self->recv_message_payload_ == nullptr) {
-    self->Unref(DEBUG_LOCATION, "OnResponseReceived");
-    return;
+  if (self->recv_message_payload_ != nullptr) {
+    // Process the response.
+    grpc_byte_buffer_reader bbr;
+    grpc_byte_buffer_reader_init(&bbr, self->recv_message_payload_);
+    grpc_slice response_slice = grpc_byte_buffer_reader_readall(&bbr);
+    grpc_byte_buffer_reader_destroy(&bbr);
+    grpc_byte_buffer_destroy(self->recv_message_payload_);
+    self->recv_message_payload_ = nullptr;
+    self->event_handler_->OnRecvMessage(StringViewFromSlice(response_slice));
+    CSliceUnref(response_slice);
   }
-  // Process the response.
-  grpc_byte_buffer_reader bbr;
-  grpc_byte_buffer_reader_init(&bbr, self->recv_message_payload_);
-  grpc_slice response_slice = grpc_byte_buffer_reader_readall(&bbr);
-  grpc_byte_buffer_reader_destroy(&bbr);
-  grpc_byte_buffer_destroy(self->recv_message_payload_);
-  self->recv_message_payload_ = nullptr;
-  self->event_handler_->OnRecvMessage(StringViewFromSlice(response_slice));
-  CSliceUnref(response_slice);
-  // Keep reading.
-  grpc_op op;
-  memset(&op, 0, sizeof(op));
-  op.op = GRPC_OP_RECV_MESSAGE;
-  op.data.recv_message.recv_message = &self->recv_message_payload_;
-  GPR_ASSERT(self->call_ != nullptr);
-  // Reuses the "OnResponseReceived" ref taken in ctor.
-  const grpc_call_error call_error = grpc_call_start_batch_and_execute(
-      self->call_, &op, 1, &self->on_response_received_);
-  GPR_ASSERT(GRPC_CALL_OK == call_error);
 }
 
 void GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall::
     OnStatusReceived(void* arg, grpc_error_handle /*error*/) {
-  auto* self = static_cast<GrpcStreamingCall*>(arg);
+  RefCountedPtr<GrpcStreamingCall> self(static_cast<GrpcStreamingCall*>(arg));
   self->event_handler_->OnStatusReceived(
       absl::Status(static_cast<absl::StatusCode>(self->status_code_),
                    StringViewFromSlice(self->status_details_)));
-  self->Unref(DEBUG_LOCATION, "OnStatusReceived");
 }
 
 //
@@ -321,8 +317,9 @@ GrpcXdsTransportFactory::GrpcXdsTransport::CreateStreamingCall(
     const char* method,
     std::unique_ptr<StreamingCall::EventHandler> event_handler) {
   return MakeOrphanable<GrpcStreamingCall>(
-      factory_->Ref(DEBUG_LOCATION, "StreamingCall"), channel_, method,
-      std::move(event_handler));
+      factory_->RefAsSubclass<GrpcXdsTransportFactory>(DEBUG_LOCATION,
+                                                       "StreamingCall"),
+      channel_, method, std::move(event_handler));
 }
 
 void GrpcXdsTransportFactory::GrpcXdsTransport::ResetBackoff() {

data/src/core/ext/xds/xds_transport_grpc.h

@@ -100,7 +100,10 @@ class GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall
 
   void SendMessage(std::string payload) override;
 
+  void StartRecvMessage() override;
+
  private:
+  static void OnRecvInitialMetadata(void* arg, grpc_error_handle /*error*/);
   static void OnRequestSent(void* arg, grpc_error_handle error);
   static void OnResponseReceived(void* arg, grpc_error_handle /*error*/);
   static void OnStatusReceived(void* arg, grpc_error_handle /*error*/);
@@ -114,6 +117,7 @@ class GrpcXdsTransportFactory::GrpcXdsTransport::GrpcStreamingCall
 
   // recv_initial_metadata
   grpc_metadata_array initial_metadata_recv_;
+  grpc_closure on_recv_initial_metadata_;
 
   // send_message
   grpc_byte_buffer* send_message_payload_ = nullptr;

data/src/core/lib/channel/call_tracer.cc

@@ -60,6 +60,11 @@ void ServerCallTracerFactory::RegisterGlobal(ServerCallTracerFactory* factory) {
   g_server_call_tracer_factory_ = factory;
 }
 
+void ServerCallTracerFactory::TestOnlyReset() {
+  delete g_server_call_tracer_factory_;
+  g_server_call_tracer_factory_ = nullptr;
+}
+
 absl::string_view ServerCallTracerFactory::ChannelArgName() {
   return kServerCallTracerFactoryChannelArgName;
 }
@@ -146,6 +151,13 @@ class DelegatingClientCallTracer : public ClientCallTracer {
     std::shared_ptr<TcpTracerInterface> StartNewTcpTrace() override {
       return nullptr;
     }
+    void AddOptionalLabels(
+        OptionalLabelComponent component,
+        std::shared_ptr<std::map<std::string, std::string>> labels) override {
+      for (auto* tracer : tracers_) {
+        tracer->AddOptionalLabels(component, labels);
+      }
+    }
     std::string TraceId() override { return tracers_[0]->TraceId(); }
     std::string SpanId() override { return tracers_[0]->SpanId(); }
     bool IsSampled() override { return tracers_[0]->IsSampled(); }

data/src/core/lib/channel/call_tracer.h

@@ -44,9 +44,9 @@ namespace grpc_core {
 
 // The interface hierarchy is as follows -
 //                 CallTracerAnnotationInterface
-//                      /          \
+//                    |                  |
 //        ClientCallTracer       CallTracerInterface
-//                                /             \
+//                                |              |
 //                      CallAttemptTracer    ServerCallTracer
 
 // The base class for all tracer implementations.
@@ -128,6 +128,11 @@ class ClientCallTracer : public CallTracerAnnotationInterface {
   // as transparent retry attempts.)
   class CallAttemptTracer : public CallTracerInterface {
    public:
+    enum class OptionalLabelComponent : std::uint8_t {
+      kXdsServiceLabels = 0,
+      kSize = 1,  // keep last
+    };
+
     ~CallAttemptTracer() override {}
     // TODO(yashykt): The following two methods `RecordReceivedTrailingMetadata`
     // and `RecordEnd` should be moved into CallTracerInterface.
@@ -140,6 +145,11 @@ class ClientCallTracer : public CallTracerAnnotationInterface {
     // Should be the last API call to the object. Once invoked, the tracer
     // library is free to destroy the object.
     virtual void RecordEnd(const gpr_timespec& latency) = 0;
+
+    // Adds optional labels to be reported by the underlying tracer in a call.
+    virtual void AddOptionalLabels(
+        OptionalLabelComponent component,
+        std::shared_ptr<std::map<std::string, std::string>> labels) = 0;
   };
 
   ~ClientCallTracer() override {}
@@ -175,7 +185,8 @@ class ServerCallTracerFactory {
 
   virtual ~ServerCallTracerFactory() {}
 
-  virtual ServerCallTracer* CreateNewServerCallTracer(Arena* arena) = 0;
+  virtual ServerCallTracer* CreateNewServerCallTracer(
+      Arena* arena, const ChannelArgs& channel_args) = 0;
 
   // Returns true if a server is to be traced, false otherwise.
   virtual bool IsServerTraced(const ChannelArgs& /*args*/) { return true; }
@@ -190,6 +201,9 @@ class ServerCallTracerFactory {
   // this for the lifetime of the process.
   static void RegisterGlobal(ServerCallTracerFactory* factory);
 
+  // Deletes any previous registered ServerCallTracerFactory.
+  static void TestOnlyReset();
+
   static absl::string_view ChannelArgName();
 };