grpc 1.18.0 → 1.22.0

data/src/core/ext/filters/client_channel/http_connect_handshaker.cc

@@ -31,153 +31,161 @@
 #include "src/core/ext/filters/client_channel/resolver_registry.h"
 #include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/channel/handshaker_registry.h"
-#include "src/core/lib/gpr/env.h"
 #include "src/core/lib/gpr/string.h"
+#include "src/core/lib/gprpp/sync.h"
 #include "src/core/lib/http/format_request.h"
 #include "src/core/lib/http/parser.h"
 #include "src/core/lib/slice/slice_internal.h"
 #include "src/core/lib/uri/uri_parser.h"
 
-typedef struct http_connect_handshaker {
-  // Base class.  Must be first.
-  grpc_handshaker base;
+namespace grpc_core {
 
-  gpr_refcount refcount;
-  gpr_mu mu;
+namespace {
 
-  bool shutdown;
+class HttpConnectHandshaker : public Handshaker {
+ public:
+  HttpConnectHandshaker();
+  void Shutdown(grpc_error* why) override;
+  void DoHandshake(grpc_tcp_server_acceptor* acceptor,
+                   grpc_closure* on_handshake_done,
+                   HandshakerArgs* args) override;
+  const char* name() const override { return "http_connect"; }
+
+ private:
+  virtual ~HttpConnectHandshaker();
+  void CleanupArgsForFailureLocked();
+  void HandshakeFailedLocked(grpc_error* error);
+  static void OnWriteDone(void* arg, grpc_error* error);
+  static void OnReadDone(void* arg, grpc_error* error);
+
+  gpr_mu mu_;
+
+  bool is_shutdown_ = false;
   // Endpoint and read buffer to destroy after a shutdown.
-  grpc_endpoint* endpoint_to_destroy;
-  grpc_slice_buffer* read_buffer_to_destroy;
+  grpc_endpoint* endpoint_to_destroy_ = nullptr;
+  grpc_slice_buffer* read_buffer_to_destroy_ = nullptr;
 
   // State saved while performing the handshake.
-  grpc_handshaker_args* args;
-  grpc_closure* on_handshake_done;
+  HandshakerArgs* args_ = nullptr;
+  grpc_closure* on_handshake_done_ = nullptr;
 
   // Objects for processing the HTTP CONNECT request and response.
-  grpc_slice_buffer write_buffer;
-  grpc_closure request_done_closure;
-  grpc_closure response_read_closure;
-  grpc_http_parser http_parser;
-  grpc_http_response http_response;
-} http_connect_handshaker;
+  grpc_slice_buffer write_buffer_;
+  grpc_closure request_done_closure_;
+  grpc_closure response_read_closure_;
+  grpc_http_parser http_parser_;
+  grpc_http_response http_response_;
+};
 
-// Unref and clean up handshaker.
-static void http_connect_handshaker_unref(http_connect_handshaker* handshaker) {
-  if (gpr_unref(&handshaker->refcount)) {
-    gpr_mu_destroy(&handshaker->mu);
-    if (handshaker->endpoint_to_destroy != nullptr) {
-      grpc_endpoint_destroy(handshaker->endpoint_to_destroy);
-    }
-    if (handshaker->read_buffer_to_destroy != nullptr) {
-      grpc_slice_buffer_destroy_internal(handshaker->read_buffer_to_destroy);
-      gpr_free(handshaker->read_buffer_to_destroy);
-    }
-    grpc_slice_buffer_destroy_internal(&handshaker->write_buffer);
-    grpc_http_parser_destroy(&handshaker->http_parser);
-    grpc_http_response_destroy(&handshaker->http_response);
-    gpr_free(handshaker);
+HttpConnectHandshaker::~HttpConnectHandshaker() {
+  gpr_mu_destroy(&mu_);
+  if (endpoint_to_destroy_ != nullptr) {
+    grpc_endpoint_destroy(endpoint_to_destroy_);
+  }
+  if (read_buffer_to_destroy_ != nullptr) {
+    grpc_slice_buffer_destroy_internal(read_buffer_to_destroy_);
+    gpr_free(read_buffer_to_destroy_);
   }
+  grpc_slice_buffer_destroy_internal(&write_buffer_);
+  grpc_http_parser_destroy(&http_parser_);
+  grpc_http_response_destroy(&http_response_);
 }
 
 // Set args fields to nullptr, saving the endpoint and read buffer for
 // later destruction.
-static void cleanup_args_for_failure_locked(
-    http_connect_handshaker* handshaker) {
-  handshaker->endpoint_to_destroy = handshaker->args->endpoint;
-  handshaker->args->endpoint = nullptr;
-  handshaker->read_buffer_to_destroy = handshaker->args->read_buffer;
-  handshaker->args->read_buffer = nullptr;
-  grpc_channel_args_destroy(handshaker->args->args);
-  handshaker->args->args = nullptr;
+void HttpConnectHandshaker::CleanupArgsForFailureLocked() {
+  endpoint_to_destroy_ = args_->endpoint;
+  args_->endpoint = nullptr;
+  read_buffer_to_destroy_ = args_->read_buffer;
+  args_->read_buffer = nullptr;
+  grpc_channel_args_destroy(args_->args);
+  args_->args = nullptr;
 }
 
 // If the handshake failed or we're shutting down, clean up and invoke the
 // callback with the error.
-static void handshake_failed_locked(http_connect_handshaker* handshaker,
-                                    grpc_error* error) {
+void HttpConnectHandshaker::HandshakeFailedLocked(grpc_error* error) {
   if (error == GRPC_ERROR_NONE) {
     // If we were shut down after an endpoint operation succeeded but
     // before the endpoint callback was invoked, we need to generate our
     // own error.
     error = GRPC_ERROR_CREATE_FROM_STATIC_STRING("Handshaker shutdown");
   }
-  if (!handshaker->shutdown) {
+  if (!is_shutdown_) {
     // TODO(ctiller): It is currently necessary to shutdown endpoints
     // before destroying them, even if we know that there are no
     // pending read/write callbacks.  This should be fixed, at which
     // point this can be removed.
-    grpc_endpoint_shutdown(handshaker->args->endpoint, GRPC_ERROR_REF(error));
+    grpc_endpoint_shutdown(args_->endpoint, GRPC_ERROR_REF(error));
     // Not shutting down, so the handshake failed.  Clean up before
     // invoking the callback.
-    cleanup_args_for_failure_locked(handshaker);
+    CleanupArgsForFailureLocked();
     // Set shutdown to true so that subsequent calls to
     // http_connect_handshaker_shutdown() do nothing.
-    handshaker->shutdown = true;
+    is_shutdown_ = true;
   }
   // Invoke callback.
-  GRPC_CLOSURE_SCHED(handshaker->on_handshake_done, error);
+  GRPC_CLOSURE_SCHED(on_handshake_done_, error);
 }
 
 // Callback invoked when finished writing HTTP CONNECT request.
-static void on_write_done(void* arg, grpc_error* error) {
-  http_connect_handshaker* handshaker =
-      static_cast<http_connect_handshaker*>(arg);
-  gpr_mu_lock(&handshaker->mu);
-  if (error != GRPC_ERROR_NONE || handshaker->shutdown) {
+void HttpConnectHandshaker::OnWriteDone(void* arg, grpc_error* error) {
+  auto* handshaker = static_cast<HttpConnectHandshaker*>(arg);
+  gpr_mu_lock(&handshaker->mu_);
+  if (error != GRPC_ERROR_NONE || handshaker->is_shutdown_) {
     // If the write failed or we're shutting down, clean up and invoke the
     // callback with the error.
-    handshake_failed_locked(handshaker, GRPC_ERROR_REF(error));
-    gpr_mu_unlock(&handshaker->mu);
-    http_connect_handshaker_unref(handshaker);
+    handshaker->HandshakeFailedLocked(GRPC_ERROR_REF(error));
+    gpr_mu_unlock(&handshaker->mu_);
+    handshaker->Unref();
   } else {
     // Otherwise, read the response.
     // The read callback inherits our ref to the handshaker.
-    grpc_endpoint_read(handshaker->args->endpoint,
-                       handshaker->args->read_buffer,
-                       &handshaker->response_read_closure);
-    gpr_mu_unlock(&handshaker->mu);
+    grpc_endpoint_read(handshaker->args_->endpoint,
+                       handshaker->args_->read_buffer,
+                       &handshaker->response_read_closure_, /*urgent=*/true);
+    gpr_mu_unlock(&handshaker->mu_);
   }
 }
 
 // Callback invoked for reading HTTP CONNECT response.
-static void on_read_done(void* arg, grpc_error* error) {
-  http_connect_handshaker* handshaker =
-      static_cast<http_connect_handshaker*>(arg);
-  gpr_mu_lock(&handshaker->mu);
-  if (error != GRPC_ERROR_NONE || handshaker->shutdown) {
+void HttpConnectHandshaker::OnReadDone(void* arg, grpc_error* error) {
+  auto* handshaker = static_cast<HttpConnectHandshaker*>(arg);
+
+  gpr_mu_lock(&handshaker->mu_);
+  if (error != GRPC_ERROR_NONE || handshaker->is_shutdown_) {
     // If the read failed or we're shutting down, clean up and invoke the
     // callback with the error.
-    handshake_failed_locked(handshaker, GRPC_ERROR_REF(error));
+    handshaker->HandshakeFailedLocked(GRPC_ERROR_REF(error));
     goto done;
   }
   // Add buffer to parser.
-  for (size_t i = 0; i < handshaker->args->read_buffer->count; ++i) {
-    if (GRPC_SLICE_LENGTH(handshaker->args->read_buffer->slices[i]) > 0) {
+  for (size_t i = 0; i < handshaker->args_->read_buffer->count; ++i) {
+    if (GRPC_SLICE_LENGTH(handshaker->args_->read_buffer->slices[i]) > 0) {
       size_t body_start_offset = 0;
-      error = grpc_http_parser_parse(&handshaker->http_parser,
-                                     handshaker->args->read_buffer->slices[i],
+      error = grpc_http_parser_parse(&handshaker->http_parser_,
+                                     handshaker->args_->read_buffer->slices[i],
                                      &body_start_offset);
       if (error != GRPC_ERROR_NONE) {
-        handshake_failed_locked(handshaker, error);
+        handshaker->HandshakeFailedLocked(error);
         goto done;
       }
-      if (handshaker->http_parser.state == GRPC_HTTP_BODY) {
+      if (handshaker->http_parser_.state == GRPC_HTTP_BODY) {
         // Remove the data we've already read from the read buffer,
         // leaving only the leftover bytes (if any).
         grpc_slice_buffer tmp_buffer;
         grpc_slice_buffer_init(&tmp_buffer);
         if (body_start_offset <
-            GRPC_SLICE_LENGTH(handshaker->args->read_buffer->slices[i])) {
+            GRPC_SLICE_LENGTH(handshaker->args_->read_buffer->slices[i])) {
           grpc_slice_buffer_add(
               &tmp_buffer,
-              grpc_slice_split_tail(&handshaker->args->read_buffer->slices[i],
+              grpc_slice_split_tail(&handshaker->args_->read_buffer->slices[i],
                                     body_start_offset));
         }
         grpc_slice_buffer_addn(&tmp_buffer,
-                               &handshaker->args->read_buffer->slices[i + 1],
-                               handshaker->args->read_buffer->count - i - 1);
-        grpc_slice_buffer_swap(handshaker->args->read_buffer, &tmp_buffer);
+                               &handshaker->args_->read_buffer->slices[i + 1],
+                               handshaker->args_->read_buffer->count - i - 1);
+        grpc_slice_buffer_swap(handshaker->args_->read_buffer, &tmp_buffer);
         grpc_slice_buffer_destroy_internal(&tmp_buffer);
         break;
       }
@@ -194,64 +202,53 @@ static void on_read_done(void* arg, grpc_error* error) {
   // need to fix the HTTP parser to understand when the body is
   // complete (e.g., handling chunked transfer encoding or looking
   // at the Content-Length: header).
-  if (handshaker->http_parser.state != GRPC_HTTP_BODY) {
-    grpc_slice_buffer_reset_and_unref_internal(handshaker->args->read_buffer);
-    grpc_endpoint_read(handshaker->args->endpoint,
-                       handshaker->args->read_buffer,
-                       &handshaker->response_read_closure);
-    gpr_mu_unlock(&handshaker->mu);
+  if (handshaker->http_parser_.state != GRPC_HTTP_BODY) {
+    grpc_slice_buffer_reset_and_unref_internal(handshaker->args_->read_buffer);
+    grpc_endpoint_read(handshaker->args_->endpoint,
+                       handshaker->args_->read_buffer,
+                       &handshaker->response_read_closure_, /*urgent=*/true);
+    gpr_mu_unlock(&handshaker->mu_);
     return;
   }
   // Make sure we got a 2xx response.
-  if (handshaker->http_response.status < 200 ||
-      handshaker->http_response.status >= 300) {
+  if (handshaker->http_response_.status < 200 ||
+      handshaker->http_response_.status >= 300) {
     char* msg;
     gpr_asprintf(&msg, "HTTP proxy returned response code %d",
-                 handshaker->http_response.status);
+                 handshaker->http_response_.status);
     error = GRPC_ERROR_CREATE_FROM_COPIED_STRING(msg);
     gpr_free(msg);
-    handshake_failed_locked(handshaker, error);
+    handshaker->HandshakeFailedLocked(error);
     goto done;
   }
   // Success.  Invoke handshake-done callback.
-  GRPC_CLOSURE_SCHED(handshaker->on_handshake_done, error);
+  GRPC_CLOSURE_SCHED(handshaker->on_handshake_done_, error);
 done:
   // Set shutdown to true so that subsequent calls to
   // http_connect_handshaker_shutdown() do nothing.
-  handshaker->shutdown = true;
-  gpr_mu_unlock(&handshaker->mu);
-  http_connect_handshaker_unref(handshaker);
+  handshaker->is_shutdown_ = true;
+  gpr_mu_unlock(&handshaker->mu_);
+  handshaker->Unref();
 }
 
 //
 // Public handshaker methods
 //
 
-static void http_connect_handshaker_destroy(grpc_handshaker* handshaker_in) {
-  http_connect_handshaker* handshaker =
-      reinterpret_cast<http_connect_handshaker*>(handshaker_in);
-  http_connect_handshaker_unref(handshaker);
-}
-
-static void http_connect_handshaker_shutdown(grpc_handshaker* handshaker_in,
-                                             grpc_error* why) {
-  http_connect_handshaker* handshaker =
-      reinterpret_cast<http_connect_handshaker*>(handshaker_in);
-  gpr_mu_lock(&handshaker->mu);
-  if (!handshaker->shutdown) {
-    handshaker->shutdown = true;
-    grpc_endpoint_shutdown(handshaker->args->endpoint, GRPC_ERROR_REF(why));
-    cleanup_args_for_failure_locked(handshaker);
+void HttpConnectHandshaker::Shutdown(grpc_error* why) {
+  gpr_mu_lock(&mu_);
+  if (!is_shutdown_) {
+    is_shutdown_ = true;
+    grpc_endpoint_shutdown(args_->endpoint, GRPC_ERROR_REF(why));
+    CleanupArgsForFailureLocked();
   }
-  gpr_mu_unlock(&handshaker->mu);
+  gpr_mu_unlock(&mu_);
   GRPC_ERROR_UNREF(why);
 }
 
-static void http_connect_handshaker_do_handshake(
-    grpc_handshaker* handshaker_in, grpc_tcp_server_acceptor* acceptor,
-    grpc_closure* on_handshake_done, grpc_handshaker_args* args) {
-  http_connect_handshaker* handshaker =
-      reinterpret_cast<http_connect_handshaker*>(handshaker_in);
+void HttpConnectHandshaker::DoHandshake(grpc_tcp_server_acceptor* acceptor,
+                                        grpc_closure* on_handshake_done,
+                                        HandshakerArgs* args) {
   // Check for HTTP CONNECT channel arg.
   // If not found, invoke on_handshake_done without doing anything.
   const grpc_arg* arg =
@@ -260,9 +257,9 @@ static void http_connect_handshaker_do_handshake(
   if (server_name == nullptr) {
     // Set shutdown to true so that subsequent calls to
     // http_connect_handshaker_shutdown() do nothing.
-    gpr_mu_lock(&handshaker->mu);
-    handshaker->shutdown = true;
-    gpr_mu_unlock(&handshaker->mu);
+    gpr_mu_lock(&mu_);
+    is_shutdown_ = true;
+    gpr_mu_unlock(&mu_);
     GRPC_CLOSURE_SCHED(on_handshake_done, GRPC_ERROR_NONE);
     return;
   }
@@ -280,6 +277,7 @@ static void http_connect_handshaker_do_handshake(
         gpr_malloc(sizeof(grpc_http_header) * num_header_strings));
     for (size_t i = 0; i < num_header_strings; ++i) {
       char* sep = strchr(header_strings[i], ':');
+
       if (sep == nullptr) {
         gpr_log(GPR_ERROR, "skipping unparseable HTTP CONNECT header: %s",
                 header_strings[i]);
@@ -292,9 +290,9 @@ static void http_connect_handshaker_do_handshake(
     }
   }
   // Save state in the handshaker object.
-  gpr_mu_lock(&handshaker->mu);
-  handshaker->args = args;
-  handshaker->on_handshake_done = on_handshake_done;
+  MutexLock lock(&mu_);
+  args_ = args;
+  on_handshake_done_ = on_handshake_done;
   // Log connection via proxy.
   char* proxy_name = grpc_endpoint_get_peer(args->endpoint);
   gpr_log(GPR_INFO, "Connecting to server %s via HTTP proxy %s", server_name,
@@ -302,15 +300,18 @@ static void http_connect_handshaker_do_handshake(
   gpr_free(proxy_name);
   // Construct HTTP CONNECT request.
   grpc_httpcli_request request;
-  memset(&request, 0, sizeof(request));
   request.host = server_name;
+  request.ssl_host_override = nullptr;
   request.http.method = (char*)"CONNECT";
   request.http.path = server_name;
+  request.http.version = GRPC_HTTP_HTTP10;  // Set by OnReadDone
   request.http.hdrs = headers;
   request.http.hdr_count = num_headers;
+  request.http.body_length = 0;
+  request.http.body = nullptr;
   request.handshaker = &grpc_httpcli_plaintext;
   grpc_slice request_slice = grpc_httpcli_format_connect_request(&request);
-  grpc_slice_buffer_add(&handshaker->write_buffer, request_slice);
+  grpc_slice_buffer_add(&write_buffer_, request_slice);
   // Clean up.
   gpr_free(headers);
   for (size_t i = 0; i < num_header_strings; ++i) {
@@ -318,54 +319,42 @@ static void http_connect_handshaker_do_handshake(
   }
   gpr_free(header_strings);
   // Take a new ref to be held by the write callback.
-  gpr_ref(&handshaker->refcount);
-  grpc_endpoint_write(args->endpoint, &handshaker->write_buffer,
-                      &handshaker->request_done_closure, nullptr);
-  gpr_mu_unlock(&handshaker->mu);
+  Ref().release();
+  grpc_endpoint_write(args->endpoint, &write_buffer_, &request_done_closure_,
+                      nullptr);
 }
 
-static const grpc_handshaker_vtable http_connect_handshaker_vtable = {
-    http_connect_handshaker_destroy, http_connect_handshaker_shutdown,
-    http_connect_handshaker_do_handshake, "http_connect"};
-
-static grpc_handshaker* grpc_http_connect_handshaker_create() {
-  http_connect_handshaker* handshaker =
-      static_cast<http_connect_handshaker*>(gpr_malloc(sizeof(*handshaker)));
-  memset(handshaker, 0, sizeof(*handshaker));
-  grpc_handshaker_init(&http_connect_handshaker_vtable, &handshaker->base);
-  gpr_mu_init(&handshaker->mu);
-  gpr_ref_init(&handshaker->refcount, 1);
-  grpc_slice_buffer_init(&handshaker->write_buffer);
-  GRPC_CLOSURE_INIT(&handshaker->request_done_closure, on_write_done,
-                    handshaker, grpc_schedule_on_exec_ctx);
-  GRPC_CLOSURE_INIT(&handshaker->response_read_closure, on_read_done,
-                    handshaker, grpc_schedule_on_exec_ctx);
-  grpc_http_parser_init(&handshaker->http_parser, GRPC_HTTP_RESPONSE,
-                        &handshaker->http_response);
-  return &handshaker->base;
+HttpConnectHandshaker::HttpConnectHandshaker() {
+  gpr_mu_init(&mu_);
+  grpc_slice_buffer_init(&write_buffer_);
+  GRPC_CLOSURE_INIT(&request_done_closure_, &HttpConnectHandshaker::OnWriteDone,
+                    this, grpc_schedule_on_exec_ctx);
+  GRPC_CLOSURE_INIT(&response_read_closure_, &HttpConnectHandshaker::OnReadDone,
+                    this, grpc_schedule_on_exec_ctx);
+  grpc_http_parser_init(&http_parser_, GRPC_HTTP_RESPONSE, &http_response_);
 }
 
 //
 // handshaker factory
 //
 
-static void handshaker_factory_add_handshakers(
-    grpc_handshaker_factory* factory, const grpc_channel_args* args,
-    grpc_pollset_set* interested_parties,
-    grpc_handshake_manager* handshake_mgr) {
-  grpc_handshake_manager_add(handshake_mgr,
-                             grpc_http_connect_handshaker_create());
-}
-
-static void handshaker_factory_destroy(grpc_handshaker_factory* factory) {}
+class HttpConnectHandshakerFactory : public HandshakerFactory {
+ public:
+  void AddHandshakers(const grpc_channel_args* args,
+                      grpc_pollset_set* interested_parties,
+                      HandshakeManager* handshake_mgr) override {
+    handshake_mgr->Add(MakeRefCounted<HttpConnectHandshaker>());
+  }
+  ~HttpConnectHandshakerFactory() override = default;
+};
 
-static const grpc_handshaker_factory_vtable handshaker_factory_vtable = {
-    handshaker_factory_add_handshakers, handshaker_factory_destroy};
+}  // namespace
 
-static grpc_handshaker_factory handshaker_factory = {
-    &handshaker_factory_vtable};
+}  // namespace grpc_core
 
 void grpc_http_connect_register_handshaker_factory() {
-  grpc_handshaker_factory_register(true /* at_start */, HANDSHAKER_CLIENT,
-                                   &handshaker_factory);
+  using namespace grpc_core;
+  HandshakerRegistry::RegisterHandshakerFactory(
+      true /* at_start */, HANDSHAKER_CLIENT,
+      UniquePtr<HandshakerFactory>(New<HttpConnectHandshakerFactory>()));
 }

data/src/core/ext/filters/client_channel/http_connect_handshaker.h

@@ -25,7 +25,7 @@
 
 /// Channel arg indicating HTTP CONNECT headers (string).
 /// Multiple headers are separated by newlines.  Key/value pairs are
-/// seperated by colons.
+/// separated by colons.
 #define GRPC_ARG_HTTP_CONNECT_HEADERS "grpc.http_connect_headers"
 
 /// Registers handshaker factory.

data/src/core/ext/filters/client_channel/lb_policy/grpclb/client_load_reporting_filter.cc

@@ -37,17 +37,6 @@ static void destroy_channel_elem(grpc_channel_element* elem) {}
 namespace {
 
 struct call_data {
-  call_data(const grpc_call_element_args& args) {
-    if (args.context[GRPC_GRPCLB_CLIENT_STATS].value != nullptr) {
-      // Get stats object from context and take a ref.
-      client_stats = static_cast<grpc_core::GrpcLbClientStats*>(
-                         args.context[GRPC_GRPCLB_CLIENT_STATS].value)
-                         ->Ref();
-      // Record call started.
-      client_stats->AddCallStarted();
-    }
-  }
-
   // Stats object to update.
   grpc_core::RefCountedPtr<grpc_core::GrpcLbClientStats> client_stats;
   // State for intercepting send_initial_metadata.
@@ -82,7 +71,7 @@ static void recv_initial_metadata_ready(void* arg, grpc_error* error) {
 static grpc_error* init_call_elem(grpc_call_element* elem,
                                   const grpc_call_element_args* args) {
   GPR_ASSERT(args->context != nullptr);
-  new (elem->call_data) call_data(*args);
+  new (elem->call_data) call_data();
   return GRPC_ERROR_NONE;
 }
 
@@ -96,9 +85,6 @@ static void destroy_call_elem(grpc_call_element* elem,
     calld->client_stats->AddCallFinished(
         !calld->send_initial_metadata_succeeded /* client_failed_to_send */,
         calld->recv_initial_metadata_succeeded /* known_received */);
-    // All done, so unref the stats object.
-    // TODO(roth): Eliminate this once filter stack is converted to C++.
-    calld->client_stats.reset();
   }
   calld->~call_data();
 }
@@ -107,25 +93,36 @@ static void start_transport_stream_op_batch(
     grpc_call_element* elem, grpc_transport_stream_op_batch* batch) {
   call_data* calld = static_cast<call_data*>(elem->call_data);
   GPR_TIMER_SCOPE("clr_start_transport_stream_op_batch", 0);
-  if (calld->client_stats != nullptr) {
-    // Intercept send_initial_metadata.
-    if (batch->send_initial_metadata) {
-      calld->original_on_complete_for_send = batch->on_complete;
-      GRPC_CLOSURE_INIT(&calld->on_complete_for_send, on_complete_for_send,
-                        calld, grpc_schedule_on_exec_ctx);
-      batch->on_complete = &calld->on_complete_for_send;
-    }
-    // Intercept recv_initial_metadata.
-    if (batch->recv_initial_metadata) {
-      calld->original_recv_initial_metadata_ready =
-          batch->payload->recv_initial_metadata.recv_initial_metadata_ready;
-      GRPC_CLOSURE_INIT(&calld->recv_initial_metadata_ready,
-                        recv_initial_metadata_ready, calld,
-                        grpc_schedule_on_exec_ctx);
-      batch->payload->recv_initial_metadata.recv_initial_metadata_ready =
-          &calld->recv_initial_metadata_ready;
+  // Handle send_initial_metadata.
+  if (batch->send_initial_metadata) {
+    // Grab client stats object from user_data for LB token metadata.
+    grpc_linked_mdelem* lb_token =
+        batch->payload->send_initial_metadata.send_initial_metadata->idx.named
+            .lb_token;
+    if (lb_token != nullptr) {
+      grpc_core::GrpcLbClientStats* client_stats =
+          static_cast<grpc_core::GrpcLbClientStats*>(grpc_mdelem_get_user_data(
+              lb_token->md, grpc_core::GrpcLbClientStats::Destroy));
+      if (client_stats != nullptr) {
+        calld->client_stats = client_stats->Ref();
+        // Intercept completion.
+        calld->original_on_complete_for_send = batch->on_complete;
+        GRPC_CLOSURE_INIT(&calld->on_complete_for_send, on_complete_for_send,
+                          calld, grpc_schedule_on_exec_ctx);
+        batch->on_complete = &calld->on_complete_for_send;
+      }
     }
   }
+  // Intercept completion of recv_initial_metadata.
+  if (batch->recv_initial_metadata) {
+    calld->original_recv_initial_metadata_ready =
+        batch->payload->recv_initial_metadata.recv_initial_metadata_ready;
+    GRPC_CLOSURE_INIT(&calld->recv_initial_metadata_ready,
+                      recv_initial_metadata_ready, calld,
+                      grpc_schedule_on_exec_ctx);
+    batch->payload->recv_initial_metadata.recv_initial_metadata_ready =
+        &calld->recv_initial_metadata_ready;
+  }
   // Chain to next filter.
   grpc_call_next_op(elem, batch);
 }