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

@@ -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;

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() {