Remove pageheap_lock from TransferCache operations.

This allows us to simplify the mock for now, which uncovered an overflow
triggered by inconsistent results in FakeCpuLayout.

PiperOrigin-RevId: 708412071
Change-Id: I4485f674520b4cf92f793cbfce300e677ec021c2

tcmalloc/mock_transfer_cache.h

@@ -66,14 +66,12 @@ class FakeTransferCacheManager {
   std::vector<std::unique_ptr<AlignedPtr>> memory_;
 };
 
-// A transfer cache manager which allocates memory from a fixed size arena. This
-// is necessary to prevent running into deadlocks in some cases, e.g. when the
-// `ShardedTransferCacheManager` calls `Alloc()` which in turns tries to
-// allocate memory using the normal malloc machinery, it leads to a deadlock
-// on the pageheap_lock.
+// A transfer cache manager which wraps malloc.
+//
+// TODO(b/175334169): Remove this once the locks are no longer used.
 class ArenaBasedFakeTransferCacheManager {
  public:
-  ArenaBasedFakeTransferCacheManager() { bytes_.resize(kTotalSize); }
+  ArenaBasedFakeTransferCacheManager() = default;
   constexpr static size_t class_to_size(int size_class) {
     // Chosen >= min size for the sharded transfer cache to kick in.
     if (size_class == kSizeClass) return 4096;
@@ -84,17 +82,14 @@ class ArenaBasedFakeTransferCacheManager {
     return 0;
   }
   void* Alloc(size_t size, std::align_val_t alignment = kAlignment) {
-    size_t space = kTotalSize - used_;
-    if (space < size) return nullptr;
-    void* head = &bytes_[used_];
-    void* aligned =
-        std::align(static_cast<size_t>(alignment), size, head, space);
-    if (aligned != nullptr) {
-      // Increase by the allocated size plus the alignment offset.
-      used_ += size + (kTotalSize - space);
-      TC_CHECK_LE(used_, bytes_.capacity());
+    {
+      // Bounce pageheap_lock to verify we can take it.
+      //
+      // TODO(b/175334169): Remove this.
+      PageHeapSpinLockHolder l;
     }
-    return aligned;
+    used_ += size;
+    return ::operator new(size, alignment);
   }
   size_t used() const { return used_; }
 
@@ -103,10 +98,6 @@ class ArenaBasedFakeTransferCacheManager {
   }
 
  private:
-  static constexpr size_t kTotalSize = 10000000;
-  // We're not changing the size of this vector during the life of this object,
-  // to avoid running into deadlocks.
-  std::vector<char> bytes_;
   size_t used_ = 0;
   static bool partial_legacy_transfer_cache_;
 };
@@ -303,7 +294,9 @@ class FakeCpuLayout {
 
   unsigned NumShards() { return num_shards_; }
   int CurrentCpu() { return current_cpu_; }
-  unsigned CpuShard(int cpu) { return cpu / kCpusPerShard; }
+  unsigned CpuShard(int cpu) {
+    return std::min(cpu / kCpusPerShard, num_shards_ - 1);
+  }
 
  private:
   int current_cpu_ = 0;
@@ -313,9 +306,7 @@ class FakeCpuLayout {
 // Defines transfer cache manager for testing legacy transfer cache.
 class FakeMultiClassTransferCacheManager : public TransferCacheManager {
  public:
-  void Init() ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) {
-    InitCaches();
-  }
+  void Init() { InitCaches(); }
 };
 
 // Wires up a largely functional TransferCache + TransferCacheManager +

tcmalloc/transfer_cache.h

@@ -299,12 +299,14 @@ class ShardedTransferCacheManagerBase {
   int tc_length(int cpu, int size_class) const {
     if (shards_ == nullptr) return 0;
     const uint8_t shard = cpu_layout_->CpuShard(cpu);
+    TC_ASSERT_LT(shard, num_shards_);
     if (!shard_initialized(shard)) return 0;
     return shards_[shard].transfer_caches[size_class].tc_length();
   }
 
   bool shard_initialized(int shard) const {
     if (shards_ == nullptr) return false;
+    TC_ASSERT_LT(shard, num_shards_);
     return shards_[shard].initialized.load(std::memory_order_acquire);
   }
 
@@ -413,9 +415,7 @@ class TransferCacheManager : public StaticForwarder {
   TransferCacheManager(const TransferCacheManager &) = delete;
   TransferCacheManager &operator=(const TransferCacheManager &) = delete;
 
-  void Init() ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) {
-    InitCaches();
-  }
+  void Init() { InitCaches(); }
 
   void InsertRange(int size_class, absl::Span<void *> batch) {
     cache_[size_class].tc.InsertRange(size_class, batch);
@@ -464,7 +464,7 @@ class TransferCacheManager : public StaticForwarder {
     }
   }
 
-  void InitCaches() ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) {
+  void InitCaches() {
     for (int i = 0; i < kNumClasses; ++i) {
       new (&cache_[i].tc) TransferCache(this, i);
     }
@@ -545,7 +545,7 @@ class TransferCacheManager {
   TransferCacheManager(const TransferCacheManager&) = delete;
   TransferCacheManager& operator=(const TransferCacheManager&) = delete;
 
-  void Init() ABSL_EXCLUSIVE_LOCKS_REQUIRED(pageheap_lock) {
+  void Init() {
     for (int i = 0; i < kNumClasses; ++i) {
       freelist_[i].Init(i);
     }