8325890: NMT: Arena Chunk value correction causes multiple errors in reporting

src/hotspot/share/memory/arena.cpp

@@ -87,6 +87,7 @@ class ChunkPool {
   // Our four static pools
   static constexpr int _num_pools = 4;
   static ChunkPool _pools[_num_pools];
+  volatile static bool _suspend_cleaning;
 
   Chunk*       _first;
   const size_t _size;         // (inner payload) size of the chunks this pool serves
@@ -110,7 +111,6 @@ class ChunkPool {
   // Clear this pool of all contained chunks
   void prune() {
     // Free all chunks with ChunkPoolLocker lock
-    // so NMT adjustment is stable.
     ChunkPoolLocker lock;
     Chunk* cur = _first;
     Chunk* next = nullptr;
@@ -136,6 +136,7 @@ class ChunkPool {
   ChunkPool(size_t size) : _first(nullptr), _size(size) {}
 
   static void clean() {
+    if (_suspend_cleaning) return;
     NativeHeapTrimmer::SuspendMark sm("chunk pool cleaner");
     for (int i = 0; i < _num_pools; i++) {
       _pools[i].prune();
@@ -145,6 +146,9 @@ class ChunkPool {
   // Returns an initialized and null-terminated Chunk of requested size
   static Chunk* allocate_chunk(Arena* arena, size_t length, AllocFailType alloc_failmode);
   static void deallocate_chunk(Chunk* p);
+  static void set_suspend_cleaning(bool suspend) {
+    _suspend_cleaning = suspend;
+  }
 };
 
 static bool on_compiler_thread() {
@@ -194,6 +198,7 @@ Chunk* ChunkPool::allocate_chunk(Arena* arena, size_t length, AllocFailType allo
     }
     chunk = (Chunk*)p;
   }
+  MemTracker::chunk_assigned_to_arena(chunk, arena->get_mem_tag());
   ::new(chunk) Chunk(length);
   // We rely on arena alignment <= malloc alignment.
   assert(is_aligned(chunk, ARENA_AMALLOC_ALIGNMENT), "Chunk start address misaligned.");
@@ -218,18 +223,19 @@ void ChunkPool::deallocate_chunk(Chunk* c) {
     c->set_stamp(0);
   }
 
+  MemTracker::add_chunk_to_pool(c);
+
   // If this is a standard-sized chunk, return it to its pool; otherwise free it.
   ChunkPool* pool = ChunkPool::get_pool_for_size(c->length());
   if (pool != nullptr) {
     pool->return_to_pool(c);
   } else {
-    // Free chunks under a lock so that NMT adjustment is stable.
-    ChunkPoolLocker lock;
     os::free(c);
   }
 }
 
 ChunkPool ChunkPool::_pools[] = { Chunk::size, Chunk::medium_size, Chunk::init_size, Chunk::tiny_size };
+volatile bool ChunkPool::_suspend_cleaning = false;
 
 class ChunkPoolCleaner : public PeriodicTask {
   static const int cleaning_interval = 5000; // cleaning interval in ms
@@ -241,6 +247,10 @@ class ChunkPoolCleaner : public PeriodicTask {
    }
 };
 
+void Arena::suspend_chunk_pool_cleaning(bool suspend) {
+  ChunkPool::set_suspend_cleaning(suspend);
+}
+
 void Arena::start_chunk_pool_cleaner_task() {
 #ifdef ASSERT
   static bool task_created = false;

src/hotspot/share/memory/arena.hpp

@@ -159,6 +159,7 @@ class Arena : public CHeapObjBase {
 
  public:
   static void initialize_chunk_pool();
+  static void suspend_chunk_pool_cleaning(bool suspend);
 
   // Start the chunk_pool cleaner task
   static void start_chunk_pool_cleaner_task();

src/hotspot/share/nmt/mallocHeader.hpp

@@ -99,8 +99,8 @@ class MallocHeader {
   NONCOPYABLE(MallocHeader);
   NOT_LP64(uint32_t _alt_canary);
   const size_t _size;
-  const uint32_t _mst_marker;
-  const MemTag _mem_tag;
+  uint32_t _mst_marker;
+  MemTag _mem_tag;
   const uint8_t _unused;
   uint16_t _canary;
 
@@ -140,7 +140,9 @@ class MallocHeader {
   inline static void revive_block(void* memblock);
   inline size_t size()  const { return _size; }
   inline MemTag mem_tag() const { return _mem_tag; }
+  inline void set_mem_tag(MemTag new_tag) { _mem_tag = new_tag; }
   inline uint32_t mst_marker() const { return _mst_marker; }
+  inline void set_mst_marker(uint32_t new_marker) { _mst_marker = new_marker; }
 
   // Return the necessary data to deaccount the block with NMT.
   FreeInfo free_info() {

src/hotspot/share/nmt/mallocSiteTable.cpp

@@ -117,7 +117,7 @@ bool MallocSiteTable::walk(MallocSiteWalker* walker) {
  */
 MallocSite* MallocSiteTable::lookup_or_add(const NativeCallStack& key, uint32_t* marker, MemTag mem_tag) {
   assert(mem_tag != mtNone, "Should have a real memory tag");
-  const unsigned int hash = key.calculate_hash();
+  const unsigned int hash = malloc_site_hash(key, mem_tag);
   const unsigned int index = hash_to_index(hash);
   *marker = 0;
 

src/hotspot/share/nmt/mallocSiteTable.hpp

@@ -54,6 +54,10 @@ class MallocSite : public AllocationSite {
   const MemoryCounter* counter() const { return &_c; }
 };
 
+inline unsigned int malloc_site_hash(const NativeCallStack& stack, MemTag tag) {
+  return stack.calculate_hash() + NMTUtil::tag_to_index(tag) * 73;
+}
+
 // Malloc site hashtable entry
 class MallocSiteHashtableEntry : public CHeapObj<mtNMT> {
  private:
@@ -64,7 +68,7 @@ class MallocSiteHashtableEntry : public CHeapObj<mtNMT> {
  public:
 
   MallocSiteHashtableEntry(NativeCallStack stack, MemTag mem_tag):
-    _malloc_site(stack, mem_tag), _hash(stack.calculate_hash()), _next(nullptr) {
+    _malloc_site(stack, mem_tag), _hash(malloc_site_hash(stack, mem_tag)), _next(nullptr) {
     assert(mem_tag != mtNone, "Expect a real memory tag");
   }
 

src/hotspot/share/nmt/mallocTracker.cpp

@@ -62,23 +62,16 @@ void MemoryCounter::update_peak(size_t size, size_t cnt) {
 }
 
 void MallocMemorySnapshot::copy_to(MallocMemorySnapshot* s) {
-  // Use lock to make sure that mtChunks don't get deallocated while the
-  // copy is going on, because their size is adjusted using this
-  // buffer in make_adjustment().
-  ChunkPoolLocker::LockStrategy ls = ChunkPoolLocker::LockStrategy::Lock;
-  if (VMError::is_error_reported() && VMError::is_error_reported_in_current_thread()) {
-    ls = ChunkPoolLocker::LockStrategy::Try;
-  }
-  ChunkPoolLocker cpl(ls);
   s->_all_mallocs = _all_mallocs;
   size_t total_size = 0;
   size_t total_count = 0;
   for (int index = 0; index < mt_number_of_tags; index ++) {
     s->_malloc[index] = _malloc[index];
-    total_size += s->_malloc[index].malloc_size();
+    total_size += s->_malloc[index].malloc_size() + s->_malloc[index].arena_size();
     total_count += s->_malloc[index].malloc_count();
   }
   // malloc counters may be updated concurrently
+  // Update total here for consistency with saved values of individual sizes.
   s->_all_mallocs.set_size_and_count(total_size, total_count);
 }
 
@@ -91,15 +84,6 @@ size_t MallocMemorySnapshot::total_arena() const {
   return amount;
 }
 
-// Make adjustment by subtracting chunks used by arenas
-// from total chunks to get total free chunk size
-void MallocMemorySnapshot::make_adjustment() {
-  size_t arena_size = total_arena();
-  int chunk_idx = NMTUtil::tag_to_index(mtChunk);
-  _malloc[chunk_idx].record_free(arena_size);
-  _all_mallocs.deallocate(arena_size);
-}
-
 void MallocMemorySummary::initialize() {
   _snapshot.initialize();
   MallocLimitHandler::initialize(MallocLimit);
@@ -203,6 +187,61 @@ void* MallocTracker::record_malloc(void* malloc_base, size_t size, MemTag mem_ta
   return memblock;
 }
 
+void MallocTracker::chunk_assigned_to_arena(void* memblock, MemTag new_tag) {
+  MallocHeader* header = (MallocHeader*)memblock - 1;
+
+  // Only decrement mtChunk, leave the total malloc amounts unchanged.
+  assert(header->mem_tag() == mtChunk, "Should only be operating on heap chunks");
+  MallocMemorySummary::as_snapshot()->by_tag(mtChunk)->record_free(header->size());
+
+  uint32_t new_mst_marker = 0;
+  if (MemTracker::tracking_level() == NMT_detail) {
+    // retrieve the old stack from MST
+    NativeCallStack old_stack;
+    if (!MallocSiteTable::access_stack(old_stack, *header)) {
+      fatal("NMT is now out of sync.");
+    }
+    MallocSiteTable::deallocation_at(header->size(), header->mst_marker());
+    // update MST with new tag
+    if (!MallocSiteTable::allocation_at(old_stack, header->size(), &new_mst_marker, new_tag)) {
+      fatal("NMT is now out of sync.");
+    }
+  }
+
+  // Arena only accounts for payload. Account for the header size and count. Attribute it to the arena's tag.
+  MallocMemorySummary::as_snapshot()->by_tag(new_tag)->record_malloc(Chunk::aligned_overhead_size());
+
+  header->set_mem_tag(new_tag);
+  header->set_mst_marker(new_mst_marker);
+}
+
+void MallocTracker::add_chunk_to_pool(void* memblock) {
+  MallocHeader* header = (MallocHeader*)memblock - 1;
+  assert(header->mem_tag() != mtChunk, "Should only be operating on arena chunks");
+
+  // Only increment mtChunk, leave the total malloc amounts unchanged.
+  MallocMemorySummary::as_snapshot()->by_tag(mtChunk)->record_malloc(header->size());
+
+  uint32_t new_mst_marker = 0;
+  if (MemTracker::tracking_level() == NMT_detail) {
+    NativeCallStack old_stack;
+    // retrieve the old stack from MST
+    if (!MallocSiteTable::access_stack(old_stack, *header)) {
+      fatal("NMT is now out of sync.");
+    }
+    MallocSiteTable::deallocation_at(header->size(), header->mst_marker());
+    // update MST with new tag
+    if (!MallocSiteTable::allocation_at(old_stack, header->size(), &new_mst_marker, mtChunk)) {
+      fatal("NMT is now out of sync.");
+    }
+  }
+
+  MallocMemorySummary::as_snapshot()->by_tag(header->mem_tag())->record_free(Chunk::aligned_overhead_size());
+
+  header->set_mem_tag(mtChunk);
+  header->set_mst_marker(new_mst_marker);
+}
+
 void* MallocTracker::record_free_block(void* memblock) {
   assert(MemTracker::enabled(), "Sanity");
   assert(memblock != nullptr, "precondition");

src/hotspot/share/nmt/mallocTracker.hpp

@@ -176,7 +176,7 @@ class MallocMemorySnapshot {
 
   // Total malloc'd memory amount
   size_t total() const {
-    return _all_mallocs.size() + malloc_overhead() + total_arena();
+    return _all_mallocs.size() + malloc_overhead();
   }
 
   // Total peak malloc
@@ -193,10 +193,6 @@ class MallocMemorySnapshot {
   size_t total_arena() const;
 
   void copy_to(MallocMemorySnapshot* s);
-
-  // Make adjustment by subtracting chunks used by arenas
-  // from total chunks to get total free chunk size
-  void make_adjustment();
 };
 
 /*
@@ -243,7 +239,6 @@ class MallocMemorySummary : AllStatic {
 
    static void snapshot(MallocMemorySnapshot* s) {
      as_snapshot()->copy_to(s);
-     s->make_adjustment();
    }
 
    // The memory used by malloc tracking headers
@@ -283,6 +278,9 @@ class MallocTracker : AllStatic {
   static void* record_malloc(void* malloc_base, size_t size, MemTag mem_tag,
     const NativeCallStack& stack);
 
+  static void chunk_assigned_to_arena(void* memblock, MemTag new_tag);
+  static void add_chunk_to_pool(void* memblock);
+
   // Given a block returned by os::malloc() or os::realloc():
   // deaccount block from NMT, mark its header as dead and return pointer to header.
   static void* record_free_block(void* memblock);

src/hotspot/share/nmt/memReporter.cpp

@@ -88,7 +88,6 @@ void MemReporterBase::print_malloc(const MemoryCounter* c, MemTag mem_tag) const
       amount_in_current_scale(amount), scale);
   }
 
-  // blends out mtChunk count number
   if (count > 0) {
     out->print(" #%zu", count);
   }
@@ -705,7 +704,7 @@ void MemSummaryDiffReporter::diff_summary_of_tag(MemTag mem_tag,
     if (amount_in_current_scale(current_malloc_amount) > 0 ||
         diff_in_current_scale(current_malloc_amount, early_malloc_amount) != 0) {
       out->print("(");
-      print_malloc_diff(current_malloc_amount, (mem_tag == mtChunk) ? 0 : current_malloc->malloc_count(),
+      print_malloc_diff(current_malloc_amount, current_malloc->malloc_count(),
         early_malloc_amount, early_malloc->malloc_count(), mtNone);
       out->print_cr(")");
     }

src/hotspot/share/nmt/memTracker.hpp

@@ -90,6 +90,22 @@ class MemTracker : AllStatic {
     return mem_base;
   }
 
+  static inline void chunk_assigned_to_arena(void* memblock, MemTag new_tag) {
+    assert(memblock != nullptr, "caller should handle null");
+    if (!enabled()) {
+      return;
+    }
+    MallocTracker::chunk_assigned_to_arena(memblock, new_tag);
+  }
+
+  static inline void add_chunk_to_pool(void* memblock) {
+    assert(memblock != nullptr, "caller should handle null");
+    if (!enabled()) {
+      return;
+    }
+    MallocTracker::add_chunk_to_pool(memblock);
+  }
+
   // Record malloc free and return malloc base address
   static inline void* record_free(void* memblock) {
     // Never turned on

src/hotspot/share/nmt/nmtUsage.cpp

@@ -22,7 +22,6 @@
  *
  */
 
-#include "memory/arena.hpp"
 #include "nmt/mallocTracker.hpp"
 #include "nmt/memoryFileTracker.hpp"
 #include "nmt/memTracker.hpp"
@@ -55,33 +54,17 @@ void NMTUsage::walk_thread_stacks() {
 }
 
 void NMTUsage::update_malloc_usage() {
-  MallocMemorySnapshot* ms;
-  // Lock needed to keep values in sync, total area size
-  // is deducted from mtChunk in the end to give correct values.
-  {
-    ChunkPoolLocker::LockStrategy ls = ChunkPoolLocker::LockStrategy::Lock;
-    if (VMError::is_error_reported() && VMError::is_error_reported_in_current_thread()) {
-      ls = ChunkPoolLocker::LockStrategy::Try;
-    }
-    ChunkPoolLocker cpl(ls);
-    ms = MallocMemorySummary::as_snapshot();
-  }
+  MallocMemorySnapshot* ms = MallocMemorySummary::as_snapshot();
 
-  size_t total_arena_size = 0;
   for (int i = 0; i < mt_number_of_tags; i++) {
     MemTag mem_tag = NMTUtil::index_to_tag(i);
     const MallocMemory* mm = ms->by_tag(mem_tag);
     _malloc_by_type[i] = mm->malloc_size() + mm->arena_size();
-    total_arena_size +=  mm->arena_size();
   }
 
   // Total malloc size.
   _malloc_total = ms->total();
 
-  // Adjustment due to mtChunk double counting.
-  _malloc_by_type[NMTUtil::tag_to_index(mtChunk)] -= total_arena_size;
-  _malloc_total -= total_arena_size;
-
   // Adjust mtNMT to include malloc overhead.
   _malloc_by_type[NMTUtil::tag_to_index(mtNMT)] += ms->malloc_overhead();
 }