8315362: NMT: summary diff reports threads count incorrectly

src/hotspot/share/services/mallocTracker.hpp

@@ -176,11 +176,6 @@ class MallocMemorySnapshot : public ResourceObj {
   // Total malloc'd memory used by arenas
   size_t total_arena() const;
 
-  inline size_t thread_count() const {
-    MallocMemorySnapshot* s = const_cast<MallocMemorySnapshot*>(this);
-    return s->by_type(mtThreadStack)->malloc_count();
-  }
-
   void copy_to(MallocMemorySnapshot* s) {
      // Need to make sure that mtChunks don't get deallocated while the
      // copy is going on, because their size is adjusted using this

src/hotspot/share/services/memBaseline.cpp

@@ -190,6 +190,7 @@ void MemBaseline::baseline(bool summaryOnly) {
 
   _instance_class_count = ClassLoaderDataGraph::num_instance_classes();
   _array_class_count = ClassLoaderDataGraph::num_array_classes();
+  _thread_count = ThreadStackTracker::thread_count();
   baseline_summary();
 
   _baseline_type = Summary_baselined;

src/hotspot/share/services/memBaseline.hpp

@@ -64,6 +64,7 @@ class MemBaseline {
 
   size_t                 _instance_class_count;
   size_t                 _array_class_count;
+  size_t                 _thread_count;
 
   // Allocation sites information
   // Malloc allocation sites
@@ -84,7 +85,7 @@ class MemBaseline {
  public:
   // create a memory baseline
   MemBaseline():
-    _instance_class_count(0), _array_class_count(0),
+    _instance_class_count(0), _array_class_count(0), _thread_count(0),
     _baseline_type(Not_baselined) {
   }
 
@@ -171,7 +172,7 @@ class MemBaseline {
 
   size_t thread_count() const {
     assert(baseline_type() != Not_baselined, "Not yet baselined");
-    return _malloc_memory_snapshot.thread_count();
+    return _thread_count;
   }
 
   // reset the baseline for reuse
@@ -180,6 +181,7 @@ class MemBaseline {
     // _malloc_memory_snapshot and _virtual_memory_snapshot are copied over.
     _instance_class_count  = 0;
     _array_class_count = 0;
+    _thread_count = 0;
 
     _malloc_sites.clear();
     _virtual_memory_sites.clear();

src/hotspot/share/services/threadStackTracker.cpp

@@ -49,40 +49,38 @@ int ThreadStackTracker::compare_thread_stack_base(const SimpleThreadStackSite& s
 void ThreadStackTracker::new_thread_stack(void* base, size_t size, const NativeCallStack& stack) {
   assert(MemTracker::tracking_level() >= NMT_summary, "Must be");
   assert(base != nullptr, "Should have been filtered");
+  ThreadCritical tc;
   if (track_as_vm()) {
-    ThreadCritical tc;
     VirtualMemoryTracker::add_reserved_region((address)base, size, stack, mtThreadStack);
-    _thread_count ++;
   } else {
     // Use a slot in mallocMemorySummary for thread stack bookkeeping
     MallocMemorySummary::record_malloc(size, mtThreadStack);
     if (MemTracker::tracking_level() == NMT_detail) {
-      ThreadCritical tc;
       assert(_simple_thread_stacks != nullptr, "Must be initialized");
       SimpleThreadStackSite site((address)base, size, stack);
       _simple_thread_stacks->add(site);
     }
   }
+  _thread_count++;
 }
 
 void ThreadStackTracker::delete_thread_stack(void* base, size_t size) {
   assert(MemTracker::tracking_level() >= NMT_summary, "Must be");
   assert(base != nullptr, "Should have been filtered");
+  ThreadCritical tc;
   if(track_as_vm()) {
-    ThreadCritical tc;
     VirtualMemoryTracker::remove_released_region((address)base, size);
-    _thread_count--;
   } else {
     // Use a slot in mallocMemorySummary for thread stack bookkeeping
     MallocMemorySummary::record_free(size, mtThreadStack);
     if (MemTracker::tracking_level() == NMT_detail) {
-      ThreadCritical tc;
       assert(_simple_thread_stacks != nullptr, "Must be initialized");
       SimpleThreadStackSite site((address)base, size, NativeCallStack::empty_stack()); // Fake object just to serve as compare target for delete
       bool removed = _simple_thread_stacks->remove(site);
       assert(removed, "Must exist");
     }
   }
+  _thread_count--;
 }
 
 bool ThreadStackTracker::walk_simple_thread_stack_site(MallocSiteWalker* walker) {

test/hotspot/jtreg/runtime/NMT/SummaryDiffThreadCount.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2023, Azul Systems, Inc. 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
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/*
+ * @test
+ * @summary run NMT baseline, create threads and verify output from summary.diff
+ * @author Evgeny Ignatenko
+ * @library /test/lib
+ * @modules java.base/jdk.internal.misc
+ *          java.management
+ * @run main/othervm -XX:+UnlockDiagnosticVMOptions -XX:NativeMemoryTracking=summary SummaryDiffThreadCount
+ */
+
+import jdk.test.lib.process.ProcessTools;
+import jdk.test.lib.process.OutputAnalyzer;
+import jdk.test.lib.JDKToolFinder;
+
+public class SummaryDiffThreadCount {
+    public static void main(String args[]) throws Exception {
+        ProcessBuilder pb = new ProcessBuilder();
+        OutputAnalyzer output;
+        // Grab my own PID.
+        String pid = Long.toString(ProcessTools.getProcessId());
+
+        pb.command(new String[] { JDKToolFinder.getJDKTool("jcmd"), pid, "VM.native_memory", "baseline=true"});
+        pb.start().waitFor();
+
+        output = new OutputAnalyzer(pb.start());
+        output.shouldContain("Baseline taken");
+
+        // Creating 10 threads.
+        for (int i = 0; i < 10; i++) {
+            new Thread(()-> {
+                while (true) { continue; }
+            }).start();
+        }
+
+        // Running "jcmd <pid> VM.native_memory summary.diff" and checking for five new threads reported.
+        pb.command(new String[] { JDKToolFinder.getJDKTool("jcmd"), pid, "VM.native_memory", "summary.diff"});
+        output = new OutputAnalyzer(pb.start());
+
+        // Trailing '+' is needed to check that NMT now reports that now we have more threads than it
+        // was during the baseline.
+        output.shouldMatch("threads #\\d+ \\+");
+    }
+}