From c1dd3ab18408c97e52856f11d135dcb3237f29b4 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Tue, 9 Aug 2016 19:42:57 -0700
Subject: [PATCH] ReadMostlyMainPtrDeleter

Summary:
Deleter helper object, which may release multiple ReadMostlyMainPtrs at once.
This allows underlying ref count implementation to perform expensive synchronization operations (e.g. asymmetric heavy barrier only once).

Reviewed By: yfeldblum

Differential Revision: D3691524

fbshipit-source-id: 3ac593b0d813345daba3a81ff4e2eb71b4db292e
---
 folly/experimental/RCURefCount.h              | 22 +++++--
 folly/experimental/ReadMostlySharedPtr.h      | 33 ++++++++++
 folly/experimental/TLRefCount.h               | 33 ++++++----
 .../test/ReadMostlySharedPtrTest.cpp          | 60 +++++++++++++++++++
 4 files changed, 131 insertions(+), 17 deletions(-)

diff --git a/folly/experimental/RCURefCount.h b/folly/experimental/RCURefCount.h
index 5acda32c..059394c8 100644
--- a/folly/experimental/RCURefCount.h
+++ b/folly/experimental/RCURefCount.h
@@ -97,17 +97,27 @@ class RCURefCount {
   }
 
   void useGlobal() noexcept {
-    state_ = State::GLOBAL_TRANSITION;
+    std::array<RCURefCount*, 1> ptrs{{this}};
+    useGlobal(ptrs);
+  }
+
+  template <typename Container>
+  static void useGlobal(const Container& refCountPtrs) {
+    for (auto refCountPtr : refCountPtrs) {
+      refCountPtr->state_ = State::GLOBAL_TRANSITION;
+    }
 
     synchronize_rcu();
     // At this point everyone is using the global count
 
-    auto accessor = localCount_.accessAllThreads();
-    for (auto& count : accessor) {
-      count.collect();
-    }
+    for (auto refCountPtr : refCountPtrs) {
+      auto accessor = refCountPtr->localCount_.accessAllThreads();
+      for (auto& count : accessor) {
+        count.collect();
+      }
 
-    state_ = State::GLOBAL;
+      refCountPtr->state_ = State::GLOBAL;
+    }
 
     synchronize_rcu();
     // After this ++ or -- can return 0.
diff --git a/folly/experimental/ReadMostlySharedPtr.h b/folly/experimental/ReadMostlySharedPtr.h
index 1839b435..f81ce856 100644
--- a/folly/experimental/ReadMostlySharedPtr.h
+++ b/folly/experimental/ReadMostlySharedPtr.h
@@ -28,6 +28,8 @@ template <typename T, typename RefCount>
 class ReadMostlyWeakPtr;
 template <typename T, typename RefCount>
 class ReadMostlySharedPtr;
+template <typename RefCount>
+class ReadMostlyMainPtrDeleter;
 
 using DefaultRefCount = TLRefCount;
 
@@ -79,6 +81,7 @@ class ReadMostlySharedPtrCore {
 
  private:
   friend class ReadMostlyMainPtr<T, RefCount>;
+  friend class ReadMostlyMainPtrDeleter<RefCount>;
 
   explicit ReadMostlySharedPtrCore(std::shared_ptr<T> ptr) :
       ptrRaw_(ptr.get()),
@@ -183,6 +186,7 @@ class ReadMostlyMainPtr {
  private:
   friend class ReadMostlyWeakPtr<T, RefCount>;
   friend class ReadMostlySharedPtr<T, RefCount>;
+  friend class ReadMostlyMainPtrDeleter<RefCount>;
 
   detail::ReadMostlySharedPtrCore<T, RefCount>* impl_{nullptr};
 };
@@ -362,4 +366,33 @@ class ReadMostlySharedPtr {
   detail::ReadMostlySharedPtrCore<T, RefCount>* impl_{nullptr};
 };
 
+/**
+ * This can be used to destroy multiple ReadMostlyMainPtrs at once.
+ */
+template <typename RefCount = DefaultRefCount>
+class ReadMostlyMainPtrDeleter {
+ public:
+  ~ReadMostlyMainPtrDeleter() noexcept {
+    RefCount::useGlobal(refCounts_);
+    for (auto& decref : decrefs_) {
+      decref();
+    }
+  }
+
+  template <typename T>
+  void add(ReadMostlyMainPtr<T, RefCount> ptr) noexcept {
+    if (!ptr.impl_) {
+      return;
+    }
+
+    refCounts_.push_back(&ptr.impl_->count_);
+    refCounts_.push_back(&ptr.impl_->weakCount_);
+    decrefs_.push_back([impl = ptr.impl_] { impl->decref(); });
+    ptr.impl_ = nullptr;
+  }
+
+ private:
+  std::vector<RefCount*> refCounts_;
+  std::vector<folly::Function<void()>> decrefs_;
+};
 }
diff --git a/folly/experimental/TLRefCount.h b/folly/experimental/TLRefCount.h
index 78054edd..c72086c9 100644
--- a/folly/experimental/TLRefCount.h
+++ b/folly/experimental/TLRefCount.h
@@ -81,25 +81,36 @@ class TLRefCount {
   }
 
   void useGlobal() noexcept {
-    std::lock_guard<std::mutex> lg(globalMutex_);
+    std::array<TLRefCount*, 1> ptrs{{this}};
+    useGlobal(ptrs);
+  }
+
+  template <typename Container>
+  static void useGlobal(const Container& refCountPtrs) {
+    std::vector<std::unique_lock<std::mutex>> lgs_;
+    for (auto refCountPtr : refCountPtrs) {
+      lgs_.emplace_back(refCountPtr->globalMutex_);
 
-    state_ = State::GLOBAL_TRANSITION;
+      refCountPtr->state_ = State::GLOBAL_TRANSITION;
+    }
 
     asymmetricHeavyBarrier();
 
-    std::weak_ptr<void> collectGuardWeak = collectGuard_;
+    for (auto refCountPtr : refCountPtrs) {
+      std::weak_ptr<void> collectGuardWeak = refCountPtr->collectGuard_;
 
-    // Make sure we can't create new LocalRefCounts
-    collectGuard_.reset();
+      // Make sure we can't create new LocalRefCounts
+      refCountPtr->collectGuard_.reset();
 
-    while (!collectGuardWeak.expired()) {
-      auto accessor = localCount_.accessAllThreads();
-      for (auto& count : accessor) {
-        count.collect();
+      while (!collectGuardWeak.expired()) {
+        auto accessor = refCountPtr->localCount_.accessAllThreads();
+        for (auto& count : accessor) {
+          count.collect();
+        }
       }
-    }
 
-    state_ = State::GLOBAL;
+      refCountPtr->state_ = State::GLOBAL;
+    }
   }
 
  private:
diff --git a/folly/experimental/test/ReadMostlySharedPtrTest.cpp b/folly/experimental/test/ReadMostlySharedPtrTest.cpp
index 05681170..4f683586 100644
--- a/folly/experimental/test/ReadMostlySharedPtrTest.cpp
+++ b/folly/experimental/test/ReadMostlySharedPtrTest.cpp
@@ -29,6 +29,7 @@
 using folly::ReadMostlyMainPtr;
 using folly::ReadMostlyWeakPtr;
 using folly::ReadMostlySharedPtr;
+using folly::ReadMostlyMainPtrDeleter;
 
 // send SIGALRM to test process after this many seconds
 const unsigned int TEST_TIMEOUT = 10;
@@ -237,3 +238,62 @@ TEST_F(ReadMostlySharedPtrTest, ClearingCache) {
   c.completed();
   t.join();
 }
+
+size_t useGlobalCalls = 0;
+
+class TestRefCount {
+ public:
+  ~TestRefCount() noexcept {
+    DCHECK_EQ(count_.load(), 0);
+  }
+
+  int64_t operator++() noexcept {
+    auto ret = ++count_;
+    DCHECK_GT(ret, 0);
+    return ret;
+  }
+
+  int64_t operator--() noexcept {
+    auto ret = --count_;
+    DCHECK_GE(ret, 0);
+    return ret;
+  }
+
+  int64_t operator*() noexcept {
+    return count_.load();
+  }
+
+  void useGlobal() {
+    ++useGlobalCalls;
+  }
+
+  template <typename Container>
+  static void useGlobal(const Container&) {
+    ++useGlobalCalls;
+  }
+
+ private:
+  std::atomic<int64_t> count_{1};
+};
+
+TEST_F(ReadMostlySharedPtrTest, ReadMostlyMainPtrDeleter) {
+  EXPECT_EQ(0, useGlobalCalls);
+  {
+    ReadMostlyMainPtr<int, TestRefCount> ptr1(std::make_shared<int>(42));
+    ReadMostlyMainPtr<int, TestRefCount> ptr2(std::make_shared<int>(42));
+  }
+
+  EXPECT_EQ(4, useGlobalCalls);
+
+  useGlobalCalls = 0;
+  {
+    ReadMostlyMainPtr<int, TestRefCount> ptr1(std::make_shared<int>(42));
+    ReadMostlyMainPtr<int, TestRefCount> ptr2(std::make_shared<int>(42));
+
+    ReadMostlyMainPtrDeleter<TestRefCount> deleter;
+    deleter.add(std::move(ptr1));
+    deleter.add(std::move(ptr2));
+  }
+
+  EXPECT_EQ(1, useGlobalCalls);
+}
-- 
2.34.1