8335091: NMT: VMATree reserve_mapping and commit_mapping APIs need MEMFLAGS while un/-committing API has no MEMFLAGS arg

Reviewed-by: jsjolen, gziemski

Author: Afshin Zafari

src/hotspot/share/nmt/vmatree.cpp

@@ -24,6 +24,7 @@
  */
 
 #include "precompiled.hpp"
+#include "logging/log.hpp"
 #include "nmt/vmatree.hpp"
 #include "utilities/growableArray.hpp"
 
@@ -34,7 +35,9 @@ const char* VMATree::statetype_strings[3] = {
 };
 
 VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType state,
-                                               const RegionData& metadata) {
+                                               const RegionData& metadata, bool use_tag_inplace) {
+  assert(!use_tag_inplace || metadata.mem_tag == mtNone,
+         "If using use_tag_inplace, then the supplied tag should be mtNone, was instead: %s", NMTUtil::tag_to_name(metadata.mem_tag));
   if (A == B) {
     // A 0-sized mapping isn't worth recording.
     return SummaryDiff();
@@ -55,13 +58,28 @@ VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType
   AddressState LEQ_A;
   TreapNode* leqA_n = _tree.closest_leq(A);
   if (leqA_n == nullptr) {
+    assert(!use_tag_inplace, "Cannot use the tag inplace if no pre-existing tag exists. From: " PTR_FORMAT " To: " PTR_FORMAT, A, B);
+    if (use_tag_inplace) {
+      log_debug(nmt)("Cannot use the tag inplace if no pre-existing tag exists. From: " PTR_FORMAT " To: " PTR_FORMAT, A, B);
+    }
     // No match. We add the A node directly, unless it would have no effect.
     if (!stA.is_noop()) {
       _tree.upsert(A, stA);
     }
   } else {
     LEQ_A_found = true;
     LEQ_A = AddressState{leqA_n->key(), leqA_n->val()};
+    StateType leqA_state = leqA_n->val().out.type();
+    StateType new_state = stA.out.type();
+    // If we specify use_tag_inplace then the new region takes over the current tag instead of the tag in metadata.
+    // This is important because the VirtualMemoryTracker API doesn't require supplying the tag for some operations.
+    if (use_tag_inplace) {
+      assert(leqA_n->val().out.type() != StateType::Released, "Should not use inplace the tag of a released region");
+      MemTag tag = leqA_n->val().out.mem_tag();
+      stA.out.set_tag(tag);
+      stB.in.set_tag(tag);
+    }
+
     // Unless we know better, let B's outgoing state be the outgoing state of the node at or preceding A.
     // Consider the case where the found node is the start of a region enclosing [A,B)
     stB.out = leqA_n->val().out;

src/hotspot/share/nmt/vmatree.hpp

@@ -66,7 +66,7 @@ class VMATree {
     return statetype_strings[static_cast<uint8_t>(type)];
   }
 
-  // Each point has some stack and a flag associated with it.
+  // Each point has some stack and a tag associated with it.
   struct RegionData {
     const NativeCallStackStorage::StackIndex stack_idx;
     const MemTag mem_tag;
@@ -88,30 +88,34 @@ class VMATree {
   struct IntervalState {
   private:
     // Store the type and mem_tag as two bytes
-    uint8_t type_flag[2];
+    uint8_t type_tag[2];
     NativeCallStackStorage::StackIndex sidx;
 
   public:
-    IntervalState() : type_flag{0,0}, sidx() {}
+    IntervalState() : type_tag{0,0}, sidx() {}
     IntervalState(const StateType type, const RegionData data) {
       assert(!(type == StateType::Released) || data.mem_tag == mtNone, "Released type must have memory tag mtNone");
-      type_flag[0] = static_cast<uint8_t>(type);
-      type_flag[1] = static_cast<uint8_t>(data.mem_tag);
+      type_tag[0] = static_cast<uint8_t>(type);
+      type_tag[1] = static_cast<uint8_t>(data.mem_tag);
       sidx = data.stack_idx;
     }
 
     StateType type() const {
-      return static_cast<StateType>(type_flag[0]);
+      return static_cast<StateType>(type_tag[0]);
     }
 
     MemTag mem_tag() const {
-      return static_cast<MemTag>(type_flag[1]);
+      return static_cast<MemTag>(type_tag[1]);
     }
 
     RegionData regiondata() const {
       return RegionData{sidx, mem_tag()};
     }
 
+    void set_tag(MemTag tag) {
+      type_tag[1] = static_cast<uint8_t>(tag);
+    }
+
     NativeCallStackStorage::StackIndex stack() const {
      return sidx;
     }
@@ -167,14 +171,20 @@ class VMATree {
     }
   };
 
-  SummaryDiff register_mapping(position A, position B, StateType state, const RegionData& metadata);
+ private:
+  SummaryDiff register_mapping(position A, position B, StateType state, const RegionData& metadata, bool use_tag_inplace = false);
 
+ public:
   SummaryDiff reserve_mapping(position from, position sz, const RegionData& metadata) {
-    return register_mapping(from, from + sz, StateType::Reserved, metadata);
+    return register_mapping(from, from + sz, StateType::Reserved, metadata, false);
+  }
+
+  SummaryDiff commit_mapping(position from, position sz, const RegionData& metadata, bool use_tag_inplace = false) {
+    return register_mapping(from, from + sz, StateType::Committed, metadata, use_tag_inplace);
   }
 
-  SummaryDiff commit_mapping(position from, position sz, const RegionData& metadata) {
-    return register_mapping(from, from + sz, StateType::Committed, metadata);
+  SummaryDiff uncommit_mapping(position from, position sz, const RegionData& metadata) {
+    return register_mapping(from, from + sz, StateType::Reserved, metadata, true);
   }
 
   SummaryDiff release_mapping(position from, position sz) {

test/hotspot/gtest/nmt/test_vmatree.cpp

@@ -171,7 +171,6 @@ class NMTVMATreeTest : public testing::Test {
 };
 
 
-
 TEST_VM_F(NMTVMATreeTest, OverlappingReservationsResultInTwoNodes) {
   VMATree::RegionData rd{si[0], mtTest};
   Tree tree;
@@ -181,6 +180,23 @@ TEST_VM_F(NMTVMATreeTest, OverlappingReservationsResultInTwoNodes) {
   EXPECT_EQ(2, count_nodes(tree));
 }
 
+TEST_VM_F(NMTVMATreeTest, UseFlagInplace) {
+  Tree tree;
+  VMATree::RegionData rd1(si[0], mtTest);
+  VMATree::RegionData rd2(si[1], mtNone);
+  tree.reserve_mapping(0, 100, rd1);
+  tree.commit_mapping(20, 50, rd2, true);
+  tree.uncommit_mapping(30, 10, rd2);
+  tree.visit_in_order([&](Node* node) {
+    if (node->key() != 100) {
+      EXPECT_EQ(mtTest, node->val().out.mem_tag()) << "failed at: " << node->key();
+      if (node->key() != 20 && node->key() != 40) {
+        EXPECT_EQ(VMATree::StateType::Reserved, node->val().out.type());
+      }
+    }
+  });
+}
+
 // Low-level tests inspecting the state of the tree.
 TEST_VM_F(NMTVMATreeTest, LowLevel) {
   adjacent_2_nodes(VMATree::empty_regiondata);