Replace ThreadCritical with a specific NMT lock

Author: roberttoyonaga

src/hotspot/share/nmt/mallocTracker.cpp

@@ -62,10 +62,10 @@ void MemoryCounter::update_peak(size_t size, size_t cnt) {
 }
 
 void MallocMemorySnapshot::copy_to(MallocMemorySnapshot* s) {
-  // Use ThreadCritical to make sure that mtChunks don't get deallocated while the
+  // Use NmtGuard 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().
-  ThreadCritical tc;
+  NmtGuard guard;
   s->_all_mallocs = _all_mallocs;
   size_t total_size = 0;
   size_t total_count = 0;

src/hotspot/share/nmt/memTracker.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2013, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -31,7 +31,6 @@
 #include "nmt/threadStackTracker.hpp"
 #include "nmt/virtualMemoryTracker.hpp"
 #include "runtime/mutexLocker.hpp"
-#include "runtime/threadCritical.hpp"
 #include "utilities/debug.hpp"
 #include "utilities/nativeCallStack.hpp"
 
@@ -125,7 +124,7 @@ class MemTracker : AllStatic {
     assert_post_init();
     if (!enabled()) return;
     if (addr != nullptr) {
-      ThreadCritical tc;
+      NmtGuard guard;
       VirtualMemoryTracker::add_reserved_region((address)addr, size, stack, flag);
     }
   }
@@ -151,7 +150,7 @@ class MemTracker : AllStatic {
     assert_post_init();
     if (!enabled()) return;
     if (addr != nullptr) {
-      ThreadCritical tc;
+      NmtGuard guard;
       VirtualMemoryTracker::add_reserved_region((address)addr, size, stack, flag);
       VirtualMemoryTracker::add_committed_region((address)addr, size, stack);
     }
@@ -162,7 +161,7 @@ class MemTracker : AllStatic {
     assert_post_init();
     if (!enabled()) return;
     if (addr != nullptr) {
-      ThreadCritical tc;
+      NmtGuard guard;
       VirtualMemoryTracker::add_committed_region((address)addr, size, stack);
     }
   }
@@ -210,7 +209,7 @@ class MemTracker : AllStatic {
     assert_post_init();
     if (!enabled()) return;
     if (addr != nullptr) {
-      ThreadCritical tc;
+      NmtGuard guard;
       VirtualMemoryTracker::split_reserved_region((address)addr, size, split, flag, split_flag);
     }
   }
@@ -219,7 +218,7 @@ class MemTracker : AllStatic {
     assert_post_init();
     if (!enabled()) return;
     if (addr != nullptr) {
-      ThreadCritical tc;
+      NmtGuard guard;
       VirtualMemoryTracker::set_reserved_region_type((address)addr, flag);
     }
   }

src/hotspot/share/nmt/nmtCommon.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2013, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -36,6 +36,9 @@ NMTUtil::S NMTUtil::_strings[] = {
   MEMORY_TYPES_DO(MEMORY_TYPE_DECLARE_NAME)
 };
 
+Semaphore NmtGuard::_nmt_semaphore(1);
+intx NmtGuard::owner_id(-1);
+
 const char* NMTUtil::scale_name(size_t scale) {
   switch(scale) {
     case 1: return "";

src/hotspot/share/nmt/nmtCommon.hpp

@@ -29,6 +29,8 @@
 
 #include "memory/allStatic.hpp"
 #include "nmt/memflags.hpp"
+#include "runtime/os.hpp"
+#include "runtime/semaphore.hpp"
 #include "utilities/align.hpp"
 #include "utilities/globalDefinitions.hpp"
 
@@ -137,5 +139,24 @@ class NMTUtil : AllStatic {
   static S _strings[mt_number_of_types];
 };
 
+class NmtGuard : public StackObj {
+private:
+    static Semaphore _nmt_semaphore;
+    static intx owner_id;
+public:
+    NmtGuard() {
+      intx tid = os::current_thread_id();
+      assert(tid != owner_id, "We already own the NMT lock. Not reentrant.");
+      _nmt_semaphore.wait();
+      owner_id = tid;
+    }
+    ~NmtGuard() {
+      owner_id = -1;
+      _nmt_semaphore.signal();
+    }
+    static void assert_locked(){
+      assert(!_nmt_semaphore.trywait(), "NMT lock should be acquired in this section.");
+    }
+};
 
 #endif // SHARE_NMT_NMTCOMMON_HPP

src/hotspot/share/nmt/nmtUsage.cpp

@@ -28,7 +28,6 @@
 #include "nmt/nmtUsage.hpp"
 #include "nmt/threadStackTracker.hpp"
 #include "nmt/virtualMemoryTracker.hpp"
-#include "runtime/threadCritical.hpp"
 
 // Enabled all options for snapshot.
 const NMTUsageOptions NMTUsage::OptionsAll = { true, true, true };
@@ -53,7 +52,7 @@ void NMTUsage::walk_thread_stacks() {
 void NMTUsage::update_malloc_usage() {
   // Thread critical needed keep values in sync, total area size
   // is deducted from mtChunk in the end to give correct values.
-  ThreadCritical tc;
+  NmtGuard guard;
   const MallocMemorySnapshot* ms = MallocMemorySummary::as_snapshot();
 
   size_t total_arena_size = 0;

src/hotspot/share/nmt/threadStackTracker.cpp

@@ -29,7 +29,6 @@
 #include "nmt/threadStackTracker.hpp"
 #include "nmt/virtualMemoryTracker.hpp"
 #include "runtime/os.hpp"
-#include "runtime/threadCritical.hpp"
 #include "utilities/align.hpp"
 #include "utilities/debug.hpp"
 #include "utilities/globalDefinitions.hpp"
@@ -53,7 +52,7 @@ void ThreadStackTracker::new_thread_stack(void* base, size_t size, const NativeC
   assert(base != nullptr, "Should have been filtered");
   align_thread_stack_boundaries_inward(base, size);
 
-  ThreadCritical tc;
+  NmtGuard guard;
   VirtualMemoryTracker::add_reserved_region((address)base, size, stack, mtThreadStack);
   _thread_count++;
 }
@@ -63,7 +62,7 @@ void ThreadStackTracker::delete_thread_stack(void* base, size_t size) {
   assert(base != nullptr, "Should have been filtered");
   align_thread_stack_boundaries_inward(base, size);
 
-  ThreadCritical tc;
+  NmtGuard guard;
   VirtualMemoryTracker::remove_released_region((address)base, size);
   _thread_count--;
 }

src/hotspot/share/nmt/virtualMemoryTracker.cpp

@@ -30,7 +30,6 @@
 #include "nmt/threadStackTracker.hpp"
 #include "nmt/virtualMemoryTracker.hpp"
 #include "runtime/os.hpp"
-#include "runtime/threadCritical.hpp"
 #include "utilities/ostream.hpp"
 
 VirtualMemorySnapshot VirtualMemorySummary::_snapshot;
@@ -621,6 +620,7 @@ class SnapshotThreadStackWalker : public VirtualMemoryWalker {
   SnapshotThreadStackWalker() {}
 
   bool do_allocation_site(const ReservedMemoryRegion* rgn) {
+    NmtGuard::assert_locked();
     if (rgn->flag() == mtThreadStack) {
       address stack_bottom = rgn->thread_stack_uncommitted_bottom();
       address committed_start;
@@ -661,7 +661,7 @@ void VirtualMemoryTracker::snapshot_thread_stacks() {
 
 bool VirtualMemoryTracker::walk_virtual_memory(VirtualMemoryWalker* walker) {
   assert(_reserved_regions != nullptr, "Sanity check");
-  ThreadCritical tc;
+  NmtGuard guard;
   // Check that the _reserved_regions haven't been deleted.
   if (_reserved_regions != nullptr) {
     LinkedListNode<ReservedMemoryRegion>* head = _reserved_regions->head();