1 files changed, 110 insertions, 18 deletions

diff --git a/src/core/lib/gprpp/ref_counted.h b/src/core/lib/gprpp/ref_counted.h
index 03c293f6ed..98de1a3653 100644
--- a/src/core/lib/gprpp/ref_counted.h
+++ b/src/core/lib/gprpp/ref_counted.h
@@ -21,9 +21,12 @@
 
 #include <grpc/support/port_platform.h>
 
+#include <grpc/support/atm.h>
 #include <grpc/support/log.h>
 #include <grpc/support/sync.h>
 
+#include <atomic>
+#include <cassert>
 #include <cinttypes>
 
 #include "src/core/lib/debug/trace.h"
@@ -34,14 +37,103 @@
 
 namespace grpc_core {
 
+// PolymorphicRefCount enforces polymorphic destruction of RefCounted.
+class PolymorphicRefCount {
+ public:
+  GRPC_ABSTRACT_BASE_CLASS
+
+ protected:
+  GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
+
+  virtual ~PolymorphicRefCount() = default;
+};
+
+// NonPolymorphicRefCount does not enforce polymorphic destruction of
+// RefCounted. Please refer to grpc_core::RefCounted for more details, and
+// when in doubt use PolymorphicRefCount.
+class NonPolymorphicRefCount {
+ public:
+  GRPC_ABSTRACT_BASE_CLASS
+
+ protected:
+  GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
+
+  ~NonPolymorphicRefCount() = default;
+};
+
+// RefCount is a simple atomic ref-count.
+//
+// This is a C++ implementation of gpr_refcount, with inline functions. Due to
+// inline functions, this class is significantly more efficient than
+// gpr_refcount and should be preferred over gpr_refcount whenever possible.
+//
+// TODO(soheil): Remove gpr_refcount after submitting the GRFC and the paragraph
+//               above.
+class RefCount {
+ public:
+  using Value = intptr_t;
+
+  // `init` is the initial refcount stored in this object.
+  constexpr explicit RefCount(Value init = 1) : value_(init) {}
+
+  // Increases the ref-count by `n`.
+  void Ref(Value n = 1) {
+    GPR_ATM_INC_ADD_THEN(value_.fetch_add(n, std::memory_order_relaxed));
+  }
+
+  // Similar to Ref() with an assert on the ref-count being non-zero.
+  void RefNonZero() {
+#ifndef NDEBUG
+    const Value prior =
+        GPR_ATM_INC_ADD_THEN(value_.fetch_add(1, std::memory_order_relaxed));
+    assert(prior > 0);
+#else
+    Ref();
+#endif
+  }
+
+  // Decrements the ref-count and returns true if the ref-count reaches 0.
+  bool Unref() {
+    const Value prior =
+        GPR_ATM_INC_ADD_THEN(value_.fetch_sub(1, std::memory_order_acq_rel));
+    GPR_DEBUG_ASSERT(prior > 0);
+    return prior == 1;
+  }
+
+  Value get() const { return value_.load(std::memory_order_relaxed); }
+
+ private:
+  std::atomic<Value> value_;
+};
+
 // 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().
 //
 // This will commonly be used by CRTP (curiously-recurring template pattern)
 // e.g., class MyClass : public RefCounted<MyClass>
-template <typename Child>
-class RefCounted {
+//
+// Use PolymorphicRefCount and NonPolymorphicRefCount to select between
+// different implementations of RefCounted.
+//
+// Note that NonPolymorphicRefCount does not support polymorphic destruction.
+// So, use NonPolymorphicRefCount only when both of the following conditions
+// are guaranteed to hold:
+// (a) Child is a concrete leaf class in RefCounted<Child>, and
+// (b) you are gauranteed to call Unref only on concrete leaf classes and not
+//     their parents.
+//
+// The following example is illegal, because calling Unref() will not call
+// the dtor of Child.
+//
+//    class Parent : public RefCounted<Parent, NonPolymorphicRefCount> {}
+//    class Child : public Parent {}
+//
+//    Child* ch;
+//    ch->Unref();
+//
+template <typename Child, typename Impl = PolymorphicRefCount>
+class RefCounted : public Impl {
  public:
   RefCountedPtr<Child> Ref() GRPC_MUST_USE_RESULT {
     IncrementRefCount();
@@ -53,7 +145,7 @@ class RefCounted {
   // private, since it will only be used by RefCountedPtr<>, which is a
   // friend of this class.
   void Unref() {
-    if (gpr_unref(&refs_)) {
+    if (refs_.Unref()) {
       Delete(static_cast<Child*>(this));
     }
   }
@@ -67,18 +159,19 @@ class RefCounted {
  protected:
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
-  RefCounted() { gpr_ref_init(&refs_, 1); }
+  RefCounted() = default;
 
-  virtual ~RefCounted() {}
+  // Note: Depending on the Impl used, this dtor can be implicitly virtual.
+  ~RefCounted() = default;
 
  private:
   // Allow RefCountedPtr<> to access IncrementRefCount().
   template <typename T>
   friend class RefCountedPtr;
 
-  void IncrementRefCount() { gpr_ref(&refs_); }
+  void IncrementRefCount() { refs_.Ref(); }
 
-  gpr_refcount refs_;
+  RefCount refs_;
 };
 
 // An alternative version of the RefCounted base class that
@@ -87,8 +180,8 @@ class RefCounted {
 // pointers and legacy code that is manually calling Ref() and Unref().
 // Once all of our code is converted to idiomatic C++, we may be able to
 // eliminate this class.
-template <typename Child>
-class RefCountedWithTracing {
+template <typename Child, typename Impl = PolymorphicRefCount>
+class RefCountedWithTracing : public Impl {
  public:
   RefCountedPtr<Child> Ref() GRPC_MUST_USE_RESULT {
     IncrementRefCount();
@@ -98,7 +191,7 @@ class RefCountedWithTracing {
   RefCountedPtr<Child> Ref(const DebugLocation& location,
                            const char* reason) GRPC_MUST_USE_RESULT {
     if (location.Log() && trace_flag_ != nullptr && trace_flag_->enabled()) {
-      gpr_atm old_refs = gpr_atm_no_barrier_load(&refs_.count);
+      const RefCount::Value old_refs = refs_.get();
       gpr_log(GPR_INFO, "%s:%p %s:%d ref %" PRIdPTR " -> %" PRIdPTR " %s",
               trace_flag_->name(), this, location.file(), location.line(),
               old_refs, old_refs + 1, reason);
@@ -112,14 +205,14 @@ class RefCountedWithTracing {
   // friend of this class.
 
   void Unref() {
-    if (gpr_unref(&refs_)) {
+    if (refs_.Unref()) {
       Delete(static_cast<Child*>(this));
     }
   }
 
   void Unref(const DebugLocation& location, const char* reason) {
     if (location.Log() && trace_flag_ != nullptr && trace_flag_->enabled()) {
-      gpr_atm old_refs = gpr_atm_no_barrier_load(&refs_.count);
+      const RefCount::Value old_refs = refs_.get();
       gpr_log(GPR_INFO, "%s:%p %s:%d unref %" PRIdPTR " -> %" PRIdPTR " %s",
               trace_flag_->name(), this, location.file(), location.line(),
               old_refs, old_refs - 1, reason);
@@ -140,26 +233,25 @@ class RefCountedWithTracing {
       : RefCountedWithTracing(static_cast<TraceFlag*>(nullptr)) {}
 
   explicit RefCountedWithTracing(TraceFlag* trace_flag)
-      : trace_flag_(trace_flag) {
-    gpr_ref_init(&refs_, 1);
-  }
+      : trace_flag_(trace_flag) {}
 
 #ifdef NDEBUG
   explicit RefCountedWithTracing(DebugOnlyTraceFlag* trace_flag)
       : RefCountedWithTracing() {}
 #endif
 
-  virtual ~RefCountedWithTracing() {}
+  // Note: Depending on the Impl used, this dtor can be implicitly virtual.
+  ~RefCountedWithTracing() = default;
 
  private:
   // Allow RefCountedPtr<> to access IncrementRefCount().
   template <typename T>
   friend class RefCountedPtr;
 
-  void IncrementRefCount() { gpr_ref(&refs_); }
+  void IncrementRefCount() { refs_.Ref(); }
 
   TraceFlag* trace_flag_ = nullptr;
-  gpr_refcount refs_;
+  RefCount refs_;
 };
 
 }  // namespace grpc_core