From 26463b93defff760a7341618cd1845038ebe30cd Mon Sep 17 00:00:00 2001
From: Vaibhav Gogte <vgogte@google.com>
Date: Fri, 18 Oct 2024 14:24:22 -0700
Subject: [PATCH] Atomically shrink to usage limit when accounting for reused
 slab.

Currently, we report impending new slab bytes before we shrink to usage limit. Also, the two operations aren't atomic. So, we might end up shrinking incorrect bytes due to interleaving allocations on other threads.

PiperOrigin-RevId: 687424762
Change-Id: I3beda652c49ce814bff44a367f1b02005ca48c92
---
 tcmalloc/cpu_cache.h       | 11 +++--------
 tcmalloc/cpu_cache_test.cc | 11 ++++++-----
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/tcmalloc/cpu_cache.h b/tcmalloc/cpu_cache.h
index 79d23ae96..931b36ae0 100644
--- a/tcmalloc/cpu_cache.h
+++ b/tcmalloc/cpu_cache.h
@@ -116,15 +116,12 @@ class StaticForwarder {
       ABSL_LOCKS_EXCLUDED(pageheap_lock) {
     TC_ASSERT(tc_globals.IsInited());
     PageHeapSpinLockHolder l;
+    if (allocated > 0) {
+      tc_globals.page_allocator().ShrinkToUsageLimit(Length(allocated));
+    }
     tc_globals.arena().UpdateAllocatedAndNonresident(allocated, nonresident);
   }
 
-  static void ShrinkToUsageLimit() ABSL_LOCKS_EXCLUDED(pageheap_lock) {
-    TC_ASSERT(tc_globals.IsInited());
-    PageHeapSpinLockHolder l;
-    tc_globals.page_allocator().ShrinkToUsageLimit(Length(0));
-  }
-
   static bool per_cpu_caches_dynamic_slab_enabled() {
     return Parameters::per_cpu_caches_dynamic_slab_enabled();
   }
@@ -1551,7 +1548,6 @@ void CpuCache<Forwarder>::ResizeSizeClassMaxCapacities()
   // Account for impending allocation/reusing of new slab so that we can avoid
   // going over memory limit.
   forwarder_.ArenaUpdateAllocatedAndNonresident(new_slabs_size, 0);
-  forwarder_.ShrinkToUsageLimit();
 
   int64_t reused_bytes;
   ResizeSlabsInfo info;
@@ -2294,7 +2290,6 @@ void CpuCache<Forwarder>::ResizeSlabIfNeeded() ABSL_NO_THREAD_SAFETY_ANALYSIS {
   // Account for impending allocation/reusing of new slab so that we can avoid
   // going over memory limit.
   forwarder_.ArenaUpdateAllocatedAndNonresident(new_slabs_size, 0);
-  forwarder_.ShrinkToUsageLimit();
 
   for (int cpu = 0; cpu < num_cpus; ++cpu) resize_[cpu].lock.Lock();
   ResizeSlabsInfo info;
diff --git a/tcmalloc/cpu_cache_test.cc b/tcmalloc/cpu_cache_test.cc
index 15aea2eb4..1ee06dc68 100644
--- a/tcmalloc/cpu_cache_test.cc
+++ b/tcmalloc/cpu_cache_test.cc
@@ -46,6 +46,7 @@
 #include "tcmalloc/internal/percpu_tcmalloc.h"
 #include "tcmalloc/internal/sysinfo.h"
 #include "tcmalloc/mock_transfer_cache.h"
+#include "tcmalloc/pages.h"
 #include "tcmalloc/parameters.h"
 #include "tcmalloc/sizemap.h"
 #include "tcmalloc/static_vars.h"
@@ -133,6 +134,11 @@ class TestStaticForwarder {
 
   void ArenaUpdateAllocatedAndNonresident(int64_t allocated,
                                           int64_t nonresident) {
+    if (allocated > 0) {
+      EXPECT_EQ(arena_reported_impending_bytes_, 0);
+      ++shrink_to_usage_limit_calls_;
+    }
+
     if (nonresident == 0) {
       arena_reported_impending_bytes_ += allocated;
     } else {
@@ -141,11 +147,6 @@ class TestStaticForwarder {
     arena_reported_nonresident_bytes_ += nonresident;
   }
 
-  void ShrinkToUsageLimit() {
-    EXPECT_GT(arena_reported_impending_bytes_, 0);
-    ++shrink_to_usage_limit_calls_;
-  }
-
   bool per_cpu_caches_dynamic_slab_enabled() { return dynamic_slab_enabled_; }
 
   double per_cpu_caches_dynamic_slab_grow_threshold() {