From 9642ca2dee51e2f4d04628a6fde613eb8101ef75 Mon Sep 17 00:00:00 2001
From: Philip Pronin <philipp@fb.com>
Date: Fri, 18 Sep 2015 23:38:39 -0700
Subject: [PATCH] implement ThreadCachedArena::totalSize()

Summary: Semantics matches `Arena<>::totalSize()`.

Reviewed By: @luciang, @ot

Differential Revision: D2456388
---
 folly/ThreadCachedArena.cpp          | 14 ++++++++++++--
 folly/ThreadCachedArena.h            | 21 +++++++++++++--------
 folly/test/ThreadCachedArenaTest.cpp |  4 ++++
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/folly/ThreadCachedArena.cpp b/folly/ThreadCachedArena.cpp
index 53636042..64c3eb96 100644
--- a/folly/ThreadCachedArena.cpp
+++ b/folly/ThreadCachedArena.cpp
@@ -16,6 +16,8 @@
 
 #include <folly/ThreadCachedArena.h>
 
+#include <memory>
+
 namespace folly {
 
 ThreadCachedArena::ThreadCachedArena(size_t minBlockSize, size_t maxAlign)
@@ -36,8 +38,16 @@ SysArena* ThreadCachedArena::allocateThreadLocalArena() {
 }
 
 void ThreadCachedArena::zombify(SysArena&& arena) {
-  std::lock_guard<std::mutex> lock(zombiesMutex_);
-  zombies_.merge(std::move(arena));
+  zombies_->merge(std::move(arena));
+}
+
+size_t ThreadCachedArena::totalSize() const {
+  size_t result = sizeof(ThreadCachedArena);
+  for (const auto& arena : arena_.accessAllThreads()) {
+    result += arena.totalSize();
+  }
+  result += zombies_->totalSize() - sizeof(SysArena);
+  return result;
 }
 
 }  // namespace folly
diff --git a/folly/ThreadCachedArena.h b/folly/ThreadCachedArena.h
index ff1eb8fd..e34a618d 100644
--- a/folly/ThreadCachedArena.h
+++ b/folly/ThreadCachedArena.h
@@ -17,13 +17,11 @@
 #ifndef FOLLY_THREADCACHEDARENA_H_
 #define FOLLY_THREADCACHEDARENA_H_
 
-#include <utility>
-#include <mutex>
-#include <limits>
-#include <boost/intrusive/slist.hpp>
+#include <type_traits>
 
-#include <folly/Likely.h>
 #include <folly/Arena.h>
+#include <folly/Likely.h>
+#include <folly/Synchronized.h>
 #include <folly/ThreadLocal.h>
 
 namespace folly {
@@ -58,7 +56,12 @@ class ThreadCachedArena {
     // Deallocate? Never!
   }
 
+  // Gets the total memory used by the arena
+  size_t totalSize() const;
+
  private:
+  struct ThreadLocalPtrTag {};
+
   ThreadCachedArena(const ThreadCachedArena&) = delete;
   ThreadCachedArena(ThreadCachedArena&&) = delete;
   ThreadCachedArena& operator=(const ThreadCachedArena&) = delete;
@@ -72,9 +75,11 @@ class ThreadCachedArena {
 
   const size_t minBlockSize_;
   const size_t maxAlign_;
-  SysArena zombies_;  // allocated from threads that are now dead
-  std::mutex zombiesMutex_;
-  ThreadLocalPtr<SysArena> arena_;  // per-thread arena
+
+  ThreadLocalPtr<SysArena, ThreadLocalPtrTag> arena_;  // Per-thread arena.
+
+  // Allocations from threads that are now dead.
+  Synchronized<SysArena> zombies_;
 };
 
 template <>
diff --git a/folly/test/ThreadCachedArenaTest.cpp b/folly/test/ThreadCachedArenaTest.cpp
index 9c4e273d..756f153d 100644
--- a/folly/test/ThreadCachedArenaTest.cpp
+++ b/folly/test/ThreadCachedArenaTest.cpp
@@ -119,9 +119,13 @@ TEST(ThreadCachedArena, BlockSize) {
 TEST(ThreadCachedArena, SingleThreaded) {
   static const size_t requestedBlockSize = 64;
   ThreadCachedArena arena(requestedBlockSize);
+  EXPECT_EQ(arena.totalSize(), sizeof(ThreadCachedArena));
+
   ArenaTester tester(arena);
   tester.allocate(100, 100 << 10);
   tester.verify();
+
+  EXPECT_GT(arena.totalSize(), sizeof(ThreadCachedArena));
 }
 
 TEST(ThreadCachedArena, MultiThreaded) {
-- 
2.34.1