Add grpc_core::RefCount and use it for RefCountedPtr

1 files changed, 57 insertions, 16 deletions

diff --git a/src/core/lib/gprpp/ref_counted.h b/src/core/lib/gprpp/ref_counted.h
index 81772f3403..e366445bff 100644
--- a/src/core/lib/gprpp/ref_counted.h
+++ b/src/core/lib/gprpp/ref_counted.h
@@ -24,6 +24,8 @@
 #include <grpc/support/log.h>
 #include <grpc/support/sync.h>
 
+#include <atomic>
+#include <cassert>
 #include <cinttypes>
 
 #include "src/core/lib/debug/trace.h"
@@ -42,7 +44,7 @@ class PolymorphicRefCount {
  protected:
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
-  virtual ~PolymorphicRefCount() {}
+  virtual ~PolymorphicRefCount() = default;
 };
 
 // NonPolymorphicRefCount does not enforce polymorphic destruction of
@@ -55,7 +57,48 @@ class NonPolymorphicRefCount {
  protected:
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
-  ~NonPolymorphicRefCount() {}
+  ~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) { 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 = 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 = 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.
@@ -97,7 +140,7 @@ class RefCounted : public Impl {
   // 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));
     }
   }
@@ -111,19 +154,19 @@ class RefCounted : public Impl {
  protected:
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
-  RefCounted() { gpr_ref_init(&refs_, 1); }
+  RefCounted() = default;
 
   // Note: Depending on the Impl used, this dtor can be implicitly virtual.
-  ~RefCounted() {}
+  ~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
@@ -143,7 +186,7 @@ class RefCountedWithTracing : public Impl {
   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);
@@ -157,14 +200,14 @@ class RefCountedWithTracing : public Impl {
   // 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);
@@ -185,9 +228,7 @@ class RefCountedWithTracing : public Impl {
       : 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)
@@ -195,17 +236,17 @@ class RefCountedWithTracing : public Impl {
 #endif
 
   // Note: Depending on the Impl used, this dtor can be implicitly virtual.
-  ~RefCountedWithTracing() {}
+  ~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