From a9657b560dec8e65643f9eda671abab35856f151 Mon Sep 17 00:00:00 2001
From: "Mark D. Roth" <roth@google.com>
Date: Wed, 21 Jul 2021 16:09:38 -0700
Subject: [PATCH] make on_call_destruction_complete always invoked via
 LoadBalancedCall

---
 .../filters/client_channel/client_channel.cc  | 16 +++++++------
 .../filters/client_channel/client_channel.h   | 24 +++++++++----------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/src/core/ext/filters/client_channel/client_channel.cc b/src/core/ext/filters/client_channel/client_channel.cc
index d985a9f0387..73b5ab5dd71 100644
--- a/src/core/ext/filters/client_channel/client_channel.cc
+++ b/src/core/ext/filters/client_channel/client_channel.cc
@@ -319,13 +319,11 @@ class DynamicTerminationFilter::CallData {
                       const grpc_call_final_info* /*final_info*/,
                       grpc_closure* then_schedule_closure) {
     auto* calld = static_cast<CallData*>(elem->call_data);
-    RefCountedPtr<SubchannelCall> subchannel_call;
-    if (GPR_LIKELY(calld->lb_call_ != nullptr)) {
-      subchannel_call = calld->lb_call_->subchannel_call();
-    }
+    RefCountedPtr<ClientChannel::LoadBalancedCall> lb_call =
+        std::move(calld->lb_call_);
     calld->~CallData();
-    if (GPR_LIKELY(subchannel_call != nullptr)) {
-      subchannel_call->SetAfterCallStackDestroy(then_schedule_closure);
+    if (GPR_LIKELY(lb_call != nullptr)) {
+      lb_call->set_on_call_destruction_complete(then_schedule_closure);
     } else {
       // TODO(yashkt) : This can potentially be a Closure::Run
       ExecCtx::Run(DEBUG_LOCATION, then_schedule_closure, GRPC_ERROR_NONE);
@@ -2601,7 +2599,11 @@ ClientChannel::LoadBalancedCall::~LoadBalancedCall() {
   for (size_t i = 0; i < GPR_ARRAY_SIZE(pending_batches_); ++i) {
     GPR_ASSERT(pending_batches_[i] == nullptr);
   }
-  if (on_call_destruction_complete_ != nullptr) {
+  // Arrange to invoke on_call_destruction_complete_ once we know that
+  // the call stack has been destroyed.
+  if (GPR_LIKELY(subchannel_call_ != nullptr)) {
+    subchannel_call_->SetAfterCallStackDestroy(on_call_destruction_complete_);
+  } else {
     ExecCtx::Run(DEBUG_LOCATION, on_call_destruction_complete_,
                  GRPC_ERROR_NONE);
   }
diff --git a/src/core/ext/filters/client_channel/client_channel.h b/src/core/ext/filters/client_channel/client_channel.h
index 0e66c2b0a22..f267f3bb087 100644
--- a/src/core/ext/filters/client_channel/client_channel.h
+++ b/src/core/ext/filters/client_channel/client_channel.h
@@ -372,19 +372,16 @@ class ClientChannel {
 // ClientChannel::LoadBalancedCall
 //
 
-// This object is ref-counted, but it cannot inherit from RefCounted<>,
-// because it is allocated on the arena and can't free its memory when
-// its refcount goes to zero.  So instead, it manually implements the
-// same API as RefCounted<>, so that it can be used with RefCountedPtr<>.
+// TODO(roth): Consider whether this actually needs to be RefCounted<>.
+// Ideally, it should be single-owner, or at least InternallyRefCounted<>.
 class ClientChannel::LoadBalancedCall
     : public RefCounted<LoadBalancedCall, PolymorphicRefCount, kUnrefCallDtor> {
  public:
   // If on_call_destruction_complete is non-null, then it will be
   // invoked once the LoadBalancedCall is completely destroyed.
-  // If it is null, then the caller is responsible for checking whether
-  // the LB call has a subchannel call and ensuring that the
-  // on_call_destruction_complete closure passed down from the surface
-  // is not invoked until after the subchannel call stack is destroyed.
+  // If it is null, then the caller is responsible for calling
+  // set_on_call_destruction_complete() before the LoadBalancedCall is
+  // destroyed.
   LoadBalancedCall(
       ClientChannel* chand, const grpc_call_element_args& args,
       grpc_polling_entity* pollent, grpc_closure* on_call_destruction_complete,
@@ -392,6 +389,13 @@ class ClientChannel::LoadBalancedCall
       bool is_transparent_retry);
   ~LoadBalancedCall() override;
 
+  // Callers must call this before unreffing if they did not set the
+  // closure via the ctor.
+  void set_on_call_destruction_complete(
+      grpc_closure* on_call_destruction_complete) {
+    on_call_destruction_complete_ = on_call_destruction_complete;
+  }
+
   void StartTransportStreamOpBatch(grpc_transport_stream_op_batch* batch);
 
   // Invoked by channel for queued LB picks when the picker is updated.
@@ -405,10 +409,6 @@ class ClientChannel::LoadBalancedCall
   // will not run until after this method returns.
   void AsyncPickDone(grpc_error_handle error);
 
-  RefCountedPtr<SubchannelCall> subchannel_call() const {
-    return subchannel_call_;
-  }
-
  private:
   class LbQueuedCallCanceller;
   class Metadata;