Allow RefCounted<> objects to defer deletion from Unref().

src/core/lib/gprpp/ref_counted.h

@@ -37,20 +37,6 @@
 
 namespace grpc_core {
 
-// PolymorphicRefCount enforces polymorphic destruction of RefCounted.
-class PolymorphicRefCount {
- public:
-  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:
-  ~NonPolymorphicRefCount() = default;
-};
-
 // RefCount is a simple atomic ref-count.
 //
 // This is a C++ implementation of gpr_refcount, with inline functions. Due to
@@ -218,10 +204,46 @@ class RefCount {
   Atomic<Value> value_;
 };
 
+// PolymorphicRefCount enforces polymorphic destruction of RefCounted.
+class PolymorphicRefCount {
+ public:
+  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:
+  ~NonPolymorphicRefCount() = default;
+};
+
+namespace internal {
+template <typename T, bool DoDelete>
+class Delete;
+template <typename T>
+class Delete<T, true> {
+ public:
+  Delete(T* t) { delete t; }
+};
+template <typename T>
+class Delete<T, false> {
+ public:
+  Delete(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.
+// 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 false.
+//
 // This will commonly be used by CRTP (curiously-recurring template pattern)
 // e.g., class MyClass : public RefCounted<MyClass>
 //
@@ -244,7 +266,8 @@ class RefCount {
 //    Child* ch;
 //    ch->Unref();
 //
-template <typename Child, typename Impl = PolymorphicRefCount>
+template <typename Child, typename Impl = PolymorphicRefCount,
+          bool DeleteUponUnref = true>
 class RefCounted : public Impl {
  public:
   // Note: Depending on the Impl used, this dtor can be implicitly virtual.
@@ -267,12 +290,12 @@ class RefCounted : public Impl {
   // friend of this class.
   void Unref() {
     if (GPR_UNLIKELY(refs_.Unref())) {
-      delete static_cast<Child*>(this);
+      internal::Delete<Child, DeleteUponUnref>(static_cast<Child*>(this));
     }
   }
   void Unref(const DebugLocation& location, const char* reason) {
     if (GPR_UNLIKELY(refs_.Unref(location, reason))) {
-      delete static_cast<Child*>(this);
+      internal::Delete<Child, DeleteUponUnref>(static_cast<Child*>(this));
     }
   }
 

test/core/gprpp/ref_counted_test.cc

@@ -18,6 +18,9 @@
 
 #include "src/core/lib/gprpp/ref_counted.h"
 
+#include <set>
+
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include "src/core/lib/gprpp/memory.h"
@@ -48,6 +51,59 @@ TEST(RefCounted, ExtraRef) {
   foo->Unref();
 }
 
+class Value : public RefCounted<Value, PolymorphicRefCount, false> {
+ public:
+  Value(int value, std::set<Value*>* registry) : value_(value) {
+    registry->insert(this);
+  }
+
+  int value() const { return value_; }
+
+ private:
+  int value_;
+};
+
+void GarbageCollectRegistry(std::set<Value*>* registry) {
+  for (auto it = registry->begin(); it != registry->end();) {
+    Value* v = *it;
+    // Check if the object has any refs remaining.
+    if (v->RefIfNonZero()) {
+      // It has refs remaining, so we do not delete it.
+      v->Unref();  // Remove the ref we just added.
+      ++it;
+    } else {
+      // No refs remaining, so delete it and remove from registry.
+      delete v;
+      it = registry->erase(it);
+    }
+  }
+}
+
+TEST(RefCounted, NoDeleteUponUnref) {
+  std::set<Value*> registry;
+  // Add two objects to the registry.
+  auto v1 = MakeRefCounted<Value>(1, &registry);
+  auto v2 = MakeRefCounted<Value>(2, &registry);
+  EXPECT_THAT(registry, ::testing::UnorderedElementsAre(
+                            ::testing::Property(&Value::value, 1),
+                            ::testing::Property(&Value::value, 2)));
+  // Running garbage collection should not delete anything, since both
+  // entries still have refs.
+  GarbageCollectRegistry(&registry);
+  EXPECT_THAT(registry, ::testing::UnorderedElementsAre(
+                            ::testing::Property(&Value::value, 1),
+                            ::testing::Property(&Value::value, 2)));
+  // Unref v2 and run GC to remove it.
+  v2.reset();
+  GarbageCollectRegistry(&registry);
+  EXPECT_THAT(registry, ::testing::UnorderedElementsAre(
+                            ::testing::Property(&Value::value, 1)));
+  // Now unref v1 and run GC again.
+  v1.reset();
+  GarbageCollectRegistry(&registry);
+  EXPECT_THAT(registry, ::testing::UnorderedElementsAre());
+}
+
 class FooNonPolymorphic
     : public RefCounted<FooNonPolymorphic, NonPolymorphicRefCount> {
  public: