Add a method to drain items out of the transfer cache if it's not used.

This will be used in the future to clean up transfer caches of L3 nodes that are
no longer used.

PiperOrigin-RevId: 392657424
Change-Id: Iab91fc986b515a4f7fc0f2b10d85b10cf5f694de

Author: Philipp Brüschweiler

tcmalloc/background.cc

@@ -23,6 +23,7 @@
 #include "tcmalloc/internal_malloc_extension.h"
 #include "tcmalloc/malloc_extension.h"
 #include "tcmalloc/parameters.h"
+#include "tcmalloc/static_vars.h"
 
 GOOGLE_MALLOC_SECTION_BEGIN
 namespace tcmalloc {
@@ -133,6 +134,7 @@ void MallocExtension_Internal_ProcessBackgroundActions() {
       }
     }
 
+    tcmalloc::tcmalloc_internal::Static().sharded_transfer_cache().Plunder();
     prev_time = now;
     absl::SleepFor(kSleepTime);
   }

tcmalloc/transfer_cache.h

@@ -95,6 +95,18 @@ class TransferCacheManager : public StaticForwarder {
     }
   }
 
+  // All caches which have not been modified since the last time this method has
+  // been called will return all objects to the freelist.
+  void Plunder() {
+    for (int i = 0; i < kNumClasses; ++i) {
+      if (use_ringbuffer_) {
+        cache_[i].rbtc.TryPlunder(i);
+      } else {
+        cache_[i].tc.TryPlunder(i);
+      }
+    }
+  }
+
   // This is not const because the underlying ring-buffer transfer cache
   // function requires acquiring a lock.
   size_t tc_length(int size_class) {
@@ -197,6 +209,7 @@ struct ShardedTransferCacheManager {
   static constexpr void *Pop(int cl) { return nullptr; }
   static constexpr void Push(int cl, void *ptr) {}
   static constexpr size_t TotalBytes() { return 0; }
+  static constexpr void Plunder() {}
 };
 
 }  // namespace tcmalloc_internal

tcmalloc/transfer_cache_internals.h

@@ -85,6 +85,7 @@ class TransferCache {
         lock_(absl::kConstInit, absl::base_internal::SCHEDULE_KERNEL_ONLY),
         max_capacity_(capacity.max_capacity),
         slot_info_(SizeInfo({0, capacity.capacity})),
+        touched_(false),
         slots_(nullptr),
         freelist_do_not_access_directly_() {
     freelist().Init(cl);
@@ -150,6 +151,7 @@ class TransferCache {
     const int B = Manager::num_objects_to_move(size_class);
     ASSERT(0 < N && N <= B);
     auto info = slot_info_.load(std::memory_order_relaxed);
+    touched_.store(true, std::memory_order_release);
     if (N == B) {
       if (info.used + N <= max_capacity_) {
         absl::base_internal::SpinLockHolder h(&lock_);
@@ -183,6 +185,7 @@ class TransferCache {
     ASSERT(N > 0);
     const int B = Manager::num_objects_to_move(size_class);
     auto info = slot_info_.load(std::memory_order_relaxed);
+    touched_.store(true, std::memory_order_release);
     if (N == B) {
       if (info.used >= N) {
         absl::base_internal::SpinLockHolder h(&lock_);
@@ -208,6 +211,35 @@ class TransferCache {
     return freelist().RemoveRange(batch, N);
   }
 
+  // If this object has not been touched since the last attempt, then
+  // return all objects to 'freelist()'.
+  void TryPlunder(int size_class) ABSL_LOCKS_EXCLUDED(lock_) {
+    if (max_capacity_ == 0) return;
+    if (touched_.load(std::memory_order_acquire)) {
+      touched_.store(false, std::memory_order_release);
+      return;
+    }
+    // If the lock is being held, someone is modifying the cache.
+    if (!lock_.TryLock()) return;
+    const int B = Manager::num_objects_to_move(size_class);
+    SizeInfo info = GetSlotInfo();
+    while (info.used > 0) {
+      const size_t num_to_move = std::min(B, info.used);
+      void *buf[kMaxObjectsToMove];
+      void **const entry = GetSlot(info.used - B);
+      memcpy(buf, entry, sizeof(void *) * B);
+      info.used -= num_to_move;
+      SetSlotInfo(info);
+      lock_.Unlock();
+      freelist().InsertRange({buf, num_to_move});
+      // If someone is starting to use the cache, stop doing this.
+      if (touched_.load(std::memory_order_acquire)) return;
+      if (!lock_.TryLock()) return;
+      info = GetSlotInfo();
+    }
+    lock_.Unlock();
+  }
+
   // Returns the number of free objects in the transfer cache.
   size_t tc_length() const {
     return static_cast<size_t>(slot_info_.load(std::memory_order_relaxed).used);
@@ -371,6 +403,9 @@ class TransferCache {
   // INVARIANT: [0 <= slot_info_.used <= slot_info.capacity <= max_cache_slots_]
   std::atomic<SizeInfo> slot_info_;
 
+  // Has this transfer cache been changed since the last call to TryPlunder()?
+  std::atomic<bool> touched_;
+
   // Pointer to array of free objects.  Use GetSlot() to get pointers to
   // entries.
   void **slots_ ABSL_GUARDED_BY(lock_);
@@ -448,6 +483,7 @@ class RingBufferTransferCache {
 
     {
       absl::base_internal::SpinLockHolder h(&lock_);
+      touched_ = true;
       RingBufferSizeInfo info = GetSlotInfo();
       if (info.used + N <= max_capacity_) {
         const bool cache_grown = MakeCacheSpace(size_class, N);
@@ -517,6 +553,7 @@ class RingBufferTransferCache {
 
     {
       absl::base_internal::SpinLockHolder h(&lock_);
+      touched_ = true;
       RingBufferSizeInfo info = GetSlotInfo();
       if (info.used > 0) {
         // Return up to however much we have in our local cache.
@@ -534,6 +571,34 @@ class RingBufferTransferCache {
     return freelist().RemoveRange(batch, N);
   }
 
+  // If this object has not been touched since the last attempt, then
+  // return all objects to 'freelist()'.
+  void TryPlunder(int size_class) ABSL_LOCKS_EXCLUDED(lock_) {
+    if (max_capacity_ == 0) return;
+    // If the lock is being held, someone is modifying the cache.
+    if (!lock_.TryLock()) return;
+    if (touched_) {
+      touched_ = false;
+      lock_.Unlock();
+      return;
+    }
+    const int B = Manager::num_objects_to_move(size_class);
+    while (slot_info_.used > 0) {
+      const size_t num_to_move(std::min(B, slot_info_.used));
+      void *buf[kMaxObjectsToMove];
+      CopyOutOfEnd(buf, num_to_move, slot_info_);
+      lock_.Unlock();
+      freelist().InsertRange({buf, num_to_move});
+      // If someone is starting to use the cache, stop doing this.
+      if (!lock_.TryLock()) return;
+      if (touched_) {
+        lock_.Unlock();
+        return;
+      }
+    }
+    lock_.Unlock();
+  }
+
   // Returns the number of free objects in the central cache.
   size_t central_length() const { return freelist().length(); }
 
@@ -778,6 +843,9 @@ class RingBufferTransferCache {
   // INVARIANT: [0 <= slot_info_.used <= slot_info.capacity <= max_cache_slots_]
   RingBufferSizeInfo slot_info_ ABSL_GUARDED_BY(lock_);
 
+  // Has this transfer cache been changed since the last call to TryPlunder()?
+  bool touched_ ABSL_GUARDED_BY(lock_) = false;
+
   // Maximum size of the cache.
   const int32_t max_capacity_;
   // This is a bitmask used instead of a modulus in the ringbuffer index

tcmalloc/transfer_cache_test.cc

@@ -311,6 +311,33 @@ TYPED_TEST_P(TransferCacheTest, WrappingFlex) {
   }
 }
 
+TYPED_TEST_P(TransferCacheTest, Plunder) {
+  TypeParam env;
+
+  // Fill in some elements so the test runs without problems.
+  env.Insert(TypeParam::kBatchSize);
+  env.Insert(TypeParam::kBatchSize);
+
+  EXPECT_CALL(env.central_freelist(), RemoveRange).Times(0);
+  EXPECT_CALL(env.central_freelist(), InsertRange).Times(2);
+
+  void* buf[TypeParam::kBatchSize];
+
+  // Make changes and make sure plundering does not do anything.
+  (void)env.transfer_cache().RemoveRange(kSizeClass, buf,
+                                         TypeParam::kBatchSize);
+  env.transfer_cache().TryPlunder(kSizeClass);
+  ASSERT_GT(env.transfer_cache().tc_length(), 0);
+
+  env.transfer_cache().InsertRange(kSizeClass, {buf, TypeParam::kBatchSize});
+  env.transfer_cache().TryPlunder(kSizeClass);
+  ASSERT_GT(env.transfer_cache().tc_length(), 0);
+
+  // And now plunder again and confirm this time we emptied the cache.
+  env.transfer_cache().TryPlunder(kSizeClass);
+  ASSERT_EQ(env.transfer_cache().tc_length(), 0);
+}
+
 // PickCoprimeBatchSize picks a batch size in [2, max_batch_size) that is
 // coprime with 2^32.  We choose the largest possible batch size within that
 // constraint to minimize the number of iterations of insert/remove required.
@@ -414,7 +441,7 @@ REGISTER_TYPED_TEST_SUITE_P(TransferCacheTest, IsolatedSmoke, ReadStats,
                             FetchesFromFreelist, PartialFetchFromFreelist,
                             EvictsOtherCaches, PushesToFreelist, WrappingWorks,
                             SingleItemSmoke, EvictsOtherCachesFlex,
-                            FullCacheFlex, WrappingFlex);
+                            FullCacheFlex, WrappingFlex, Plunder);
 template <typename Env>
 using FuzzTest = ::testing::Test;
 TYPED_TEST_SUITE_P(FuzzTest);