Add support for obtaining cache size of the current thread and softer idling

7 files changed, 57 insertions, 1 deletions

diff --git a/src/gperftools/malloc_extension.h b/src/gperftools/malloc_extension.h
index f0a3895..689b5f1 100644
--- a/src/gperftools/malloc_extension.h
+++ b/src/gperftools/malloc_extension.h
@@ -395,6 +395,15 @@ class PERFTOOLS_DLL_DECL MallocExtension {
   // Like ReadStackTraces(), but returns stack traces that caused growth
   // in the address space size.
   virtual void** ReadHeapGrowthStackTraces();
+
+  // Returns the size in bytes of the calling threads cache.
+  virtual size_t GetThreadCacheSize();
+
+  // Like MarkThreadIdle, but does not destroy the internal data
+  // structures of the thread cache. When the thread resumes, it wil
+  // have an empty cache but will not need to pay to reconstruct the
+  // cache data structures.
+  virtual void MarkThreadTemporarilyIdle();
 };
 
 namespace base {
diff --git a/src/gperftools/malloc_extension_c.h b/src/gperftools/malloc_extension_c.h
index baa013d..70ff686 100644
--- a/src/gperftools/malloc_extension_c.h
+++ b/src/gperftools/malloc_extension_c.h
@@ -79,6 +79,8 @@ PERFTOOLS_DLL_DECL void MallocExtension_ReleaseToSystem(size_t num_bytes);
 PERFTOOLS_DLL_DECL void MallocExtension_ReleaseFreeMemory(void);
 PERFTOOLS_DLL_DECL size_t MallocExtension_GetEstimatedAllocatedSize(size_t size);
 PERFTOOLS_DLL_DECL size_t MallocExtension_GetAllocatedSize(const void* p);
+PERFTOOLS_DLL_DECL size_t MallocExtension_GetThreadCacheSize(void);
+PERFTOOLS_DLL_DECL void MallocExtension_MarkThreadTemporarilyIdle(void);
 
 /*
  * NOTE: These enum values MUST be kept in sync with the version in
diff --git a/src/malloc_extension.cc b/src/malloc_extension.cc
index 4ff719c..6e69552 100644
--- a/src/malloc_extension.cc
+++ b/src/malloc_extension.cc
@@ -193,6 +193,14 @@ void MallocExtension::GetFreeListSizes(
   v->clear();
 }
 
+size_t MallocExtension::GetThreadCacheSize() {
+  return 0;
+}
+
+void MallocExtension::MarkThreadTemporarilyIdle() {
+  // Default implementation does nothing
+}
+
 // The current malloc extension object.
 
 static MallocExtension* current_instance;
@@ -369,6 +377,8 @@ C_SHIM(ReleaseFreeMemory, void, (void), ());
 C_SHIM(ReleaseToSystem, void, (size_t num_bytes), (num_bytes));
 C_SHIM(GetEstimatedAllocatedSize, size_t, (size_t size), (size));
 C_SHIM(GetAllocatedSize, size_t, (const void* p), (p));
+C_SHIM(GetThreadCacheSize, size_t, (void), ());
+C_SHIM(MarkThreadTemporarilyIdle, void, (void), ());
 
 // Can't use the shim here because of the need to translate the enums.
 extern "C"
diff --git a/src/tcmalloc.cc b/src/tcmalloc.cc
index 0941bff..2a4e7a6 100644
--- a/src/tcmalloc.cc
+++ b/src/tcmalloc.cc
@@ -629,6 +629,17 @@ class TCMallocImplementation : public MallocExtension {
     return DumpHeapGrowthStackTraces();
   }
 
+  virtual size_t GetThreadCacheSize() {
+    ThreadCache* tc = ThreadCache::GetCacheIfPresent();
+    if (!tc)
+      return 0;
+    return tc->Size();
+  }
+
+  virtual void MarkThreadTemporarilyIdle() {
+    ThreadCache::BecomeTemporarilyIdle();
+  }
+
   virtual void Ranges(void* arg, RangeFunction func) {
     IterateOverRanges(arg, func);
   }
diff --git a/src/tests/markidle_unittest.cc b/src/tests/markidle_unittest.cc
index 827609f..7e59e51 100644
--- a/src/tests/markidle_unittest.cc
+++ b/src/tests/markidle_unittest.cc
@@ -98,11 +98,28 @@ static void TestIdleUsage() {
   VLOG(0, "Post idle: %" PRIuS "\n", post_idle);
 }
 
+static void TestTemporarilyIdleUsage() {
+  const size_t original = MallocExtension::instance()->GetThreadCacheSize();
+
+  TestAllocation();
+  const size_t post_allocation = MallocExtension::instance()->GetThreadCacheSize();
+  CHECK_GT(post_allocation, original);
+
+  MallocExtension::instance()->MarkThreadIdle();
+  const size_t post_idle = MallocExtension::instance()->GetThreadCacheSize();
+  CHECK_EQ(post_idle, 0);
+
+  // Log after testing because logging can allocate heap memory.
+  VLOG(0, "Original usage: %" PRIuS "\n", original);
+  VLOG(0, "Post allocation: %" PRIuS "\n", post_allocation);
+  VLOG(0, "Post idle: %" PRIuS "\n", post_idle);
+}
+
 int main(int argc, char** argv) {
   RunThread(&TestIdleUsage);
   RunThread(&TestAllocation);
   RunThread(&MultipleIdleCalls);
-  RunThread(&MultipleIdleNonIdlePhases);
+  RunThread(&TestTemporarilyIdleUsage);
 
   printf("PASS\n");
   return 0;
diff --git a/src/thread_cache.cc b/src/thread_cache.cc
index 6e0c2bc..ef1f435 100644
--- a/src/thread_cache.cc
+++ b/src/thread_cache.cc
@@ -396,6 +396,12 @@ void ThreadCache::BecomeIdle() {
   DeleteCache(heap);
 }
 
+void ThreadCache::BecomeTemporarilyIdle() {
+  ThreadCache* heap = GetCacheIfPresent();
+  if (heap)
+    heap->Cleanup();
+}
+
 void ThreadCache::DestroyThreadCache(void* ptr) {
   // Note that "ptr" cannot be NULL since pthread promises not
   // to invoke the destructor on NULL values, but for safety,
diff --git a/src/thread_cache.h b/src/thread_cache.h
index ed02c33..e3f4d67 100644
--- a/src/thread_cache.h
+++ b/src/thread_cache.h
@@ -108,6 +108,7 @@ class ThreadCache {
   static ThreadCache* GetCacheWhichMustBePresent();
   static ThreadCache* CreateCacheIfNecessary();
   static void         BecomeIdle();
+  static void         BecomeTemporarilyIdle();
   static size_t       MinSizeForSlowPath();
   static void         SetMinSizeForSlowPath(size_t size);