allow RefCounted<> to be used for objects allocated on the arena (#25848)

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

@@ -2510,9 +2510,9 @@ class ClientChannel::LoadBalancedCall::LbCallState
 ClientChannel::LoadBalancedCall::LoadBalancedCall(
     ClientChannel* chand, const grpc_call_element_args& args,
     grpc_polling_entity* pollent)
-    : refs_(1, GRPC_TRACE_FLAG_ENABLED(grpc_client_channel_routing_trace)
-                   ? "LoadBalancedCall"
-                   : nullptr),
+    : RefCounted(GRPC_TRACE_FLAG_ENABLED(grpc_client_channel_routing_trace)
+                     ? "LoadBalancedCall"
+                     : nullptr),
       chand_(chand),
       path_(grpc_slice_ref_internal(args.path)),
       call_start_time_(args.start_time),
@@ -2536,39 +2536,6 @@ ClientChannel::LoadBalancedCall::~LoadBalancedCall() {
   }
 }
 
-RefCountedPtr<ClientChannel::LoadBalancedCall>
-ClientChannel::LoadBalancedCall::Ref() {
-  IncrementRefCount();
-  return RefCountedPtr<LoadBalancedCall>(this);
-}
-
-RefCountedPtr<ClientChannel::LoadBalancedCall>
-ClientChannel::LoadBalancedCall::Ref(const DebugLocation& location,
-                                     const char* reason) {
-  IncrementRefCount(location, reason);
-  return RefCountedPtr<LoadBalancedCall>(this);
-}
-
-void ClientChannel::LoadBalancedCall::Unref() {
-  if (GPR_UNLIKELY(refs_.Unref())) {
-    this->~LoadBalancedCall();
-  }
-}
-
-void ClientChannel::LoadBalancedCall::Unref(const DebugLocation& location,
-                                            const char* reason) {
-  if (GPR_UNLIKELY(refs_.Unref(location, reason))) {
-    this->~LoadBalancedCall();
-  }
-}
-
-void ClientChannel::LoadBalancedCall::IncrementRefCount() { refs_.Ref(); }
-
-void ClientChannel::LoadBalancedCall::IncrementRefCount(
-    const DebugLocation& location, const char* reason) {
-  refs_.Ref(location, reason);
-}
-
 void* ClientChannel::LoadBalancedCall::GetParentData() {
   return reinterpret_cast<char*>(this) +
          GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(LoadBalancedCall));

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

@@ -355,20 +355,12 @@ class ClientChannel {
 // 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<>.
-class ClientChannel::LoadBalancedCall {
+class ClientChannel::LoadBalancedCall
+    : public RefCounted<LoadBalancedCall, PolymorphicRefCount, kUnrefCallDtor> {
  public:
   LoadBalancedCall(ClientChannel* chand, const grpc_call_element_args& args,
                    grpc_polling_entity* pollent);
-  ~LoadBalancedCall();
-
-  // Interface of RefCounted<>.
-  RefCountedPtr<LoadBalancedCall> Ref() GRPC_MUST_USE_RESULT;
-  RefCountedPtr<LoadBalancedCall> Ref(const DebugLocation& location,
-                                      const char* reason) GRPC_MUST_USE_RESULT;
-  // When refcount drops to 0, destroys itself and the associated call stack,
-  // but does NOT free the memory because it's in the call arena.
-  void Unref();
-  void Unref(const DebugLocation& location, const char* reason);
+  ~LoadBalancedCall() override;
 
   void* GetParentData();
 
@@ -390,18 +382,10 @@ class ClientChannel::LoadBalancedCall {
   }
 
  private:
-  // Allow RefCountedPtr<> to access IncrementRefCount().
-  template <typename T>
-  friend class ::grpc_core::RefCountedPtr;
-
   class LbQueuedCallCanceller;
   class Metadata;
   class LbCallState;
 
-  // Interface of RefCounted<>.
-  void IncrementRefCount();
-  void IncrementRefCount(const DebugLocation& location, const char* reason);
-
   // Returns the index into pending_batches_ to be used for batch.
   static size_t GetBatchIndex(grpc_transport_stream_op_batch* batch);
   void PendingBatchesAdd(grpc_transport_stream_op_batch* batch);
@@ -444,8 +428,6 @@ class ClientChannel::LoadBalancedCall {
   void MaybeAddCallToLbQueuedCallsLocked()
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(&ClientChannel::data_plane_mu_);
 
-  RefCount refs_;
-
   ClientChannel* chand_;
 
   // TODO(roth): Instead of duplicating these fields in every filter

src/core/ext/filters/client_channel/resolver/xds/xds_resolver.cc

@@ -122,7 +122,7 @@ class XdsResolver : public Resolver {
   };
 
   class ClusterState
-      : public RefCounted<ClusterState, PolymorphicRefCount, false> {
+      : public RefCounted<ClusterState, PolymorphicRefCount, kUnrefNoDelete> {
    public:
     using ClusterStateMap =
         std::map<std::string, std::unique_ptr<ClusterState>>;

src/core/lib/gprpp/ref_counted.h

@@ -215,31 +215,45 @@ class NonPolymorphicRefCount {
   ~NonPolymorphicRefCount() = default;
 };
 
+// Behavior of RefCounted<> upon ref count reaching 0.
+enum UnrefBehavior {
+  // Default behavior: Delete the object.
+  kUnrefDelete,
+  // Do not delete the object upon unref.  This is useful in cases where all
+  // existing objects must be tracked in a registry but the object's entry in
+  // the registry cannot be removed from the object's dtor due to
+  // synchronization issues.  In this case, the registry can be cleaned up
+  // later by identifying entries for which RefIfNonZero() returns null.
+  kUnrefNoDelete,
+  // Call the object's dtor but do not delete it.  This is useful for cases
+  // where the object is stored in memory allocated elsewhere (e.g., the call
+  // arena).
+  kUnrefCallDtor,
+};
+
 namespace internal {
-template <typename T, bool DoDelete>
+template <typename T, UnrefBehavior UnrefBehaviorArg>
 class Delete;
 template <typename T>
-class Delete<T, true> {
+class Delete<T, kUnrefDelete> {
  public:
   explicit Delete(T* t) { delete t; }
 };
 template <typename T>
-class Delete<T, false> {
+class Delete<T, kUnrefNoDelete> {
  public:
   explicit Delete(T* /*t*/) {}
 };
+template <typename T>
+class Delete<T, kUnrefCallDtor> {
+ public:
+  explicit Delete(T* t) { t->~T(); }
+};
 }  // namespace internal
 
 // A base class for reference-counted objects.
 // New objects should be created via new and start with a refcount of 1.
-// When the refcount reaches 0, the object will be deleted via delete.
-//
-// If DeleteUponUnref is false, deletion will not occur when the ref
-// count reaches 0.  This is useful in cases where all existing objects
-// must be tracked in a registry but the object's entry in the registry
-// cannot be removed from the object's dtor due to synchronization issues.
-// In this case, the registry can be cleaned up later by identifying
-// entries for which RefIfNonZero() returns null.
+// When the refcount reaches 0, executes the specified UnrefBehavior.
 //
 // This will commonly be used by CRTP (curiously-recurring template pattern)
 // e.g., class MyClass : public RefCounted<MyClass>
@@ -264,7 +278,7 @@ class Delete<T, false> {
 //    ch->Unref();
 //
 template <typename Child, typename Impl = PolymorphicRefCount,
-          bool DeleteUponUnref = true>
+          UnrefBehavior UnrefBehaviorArg = kUnrefDelete>
 class RefCounted : public Impl {
  public:
   // Note: Depending on the Impl used, this dtor can be implicitly virtual.
@@ -287,12 +301,12 @@ class RefCounted : public Impl {
   // friend of this class.
   void Unref() {
     if (GPR_UNLIKELY(refs_.Unref())) {
-      internal::Delete<Child, DeleteUponUnref>(static_cast<Child*>(this));
+      internal::Delete<Child, UnrefBehaviorArg>(static_cast<Child*>(this));
     }
   }
   void Unref(const DebugLocation& location, const char* reason) {
     if (GPR_UNLIKELY(refs_.Unref(location, reason))) {
-      internal::Delete<Child, DeleteUponUnref>(static_cast<Child*>(this));
+      internal::Delete<Child, UnrefBehaviorArg>(static_cast<Child*>(this));
     }
   }
 

test/core/gprpp/ref_counted_test.cc

@@ -51,7 +51,7 @@ TEST(RefCounted, ExtraRef) {
   foo->Unref();
 }
 
-class Value : public RefCounted<Value, PolymorphicRefCount, false> {
+class Value : public RefCounted<Value, PolymorphicRefCount, kUnrefNoDelete> {
  public:
   Value(int value, std::set<std::unique_ptr<Value>>* registry) : value_(value) {
     registry->emplace(this);
@@ -104,6 +104,26 @@ TEST(RefCounted, NoDeleteUponUnref) {
   EXPECT_THAT(registry, ::testing::UnorderedElementsAre());
 }
 
+class ValueInExternalAllocation
+    : public RefCounted<ValueInExternalAllocation, PolymorphicRefCount,
+                        kUnrefCallDtor> {
+ public:
+  explicit ValueInExternalAllocation(int value) : value_(value) {}
+
+  int value() const { return value_; }
+
+ private:
+  int value_;
+};
+
+TEST(RefCounted, CallDtorUponUnref) {
+  std::aligned_storage<sizeof(ValueInExternalAllocation),
+                       alignof(ValueInExternalAllocation)>::type storage;
+  RefCountedPtr<ValueInExternalAllocation> value(
+      new (&storage) ValueInExternalAllocation(5));
+  EXPECT_EQ(value->value(), 5);
+}
+
 class FooNonPolymorphic
     : public RefCounted<FooNonPolymorphic, NonPolymorphicRefCount> {
  public: