8340103: Add internal set_flag function to VMATree

Reviewed-by: stuefe, azafari, gziemski

src/hotspot/share/nmt/nmtTreap.hpp

@@ -234,6 +234,10 @@ public:
     this->remove_all();
   }
 
+  int size() {
+    return _node_count;
+  }
+
   void upsert(const K& k, const V& v) {
     TreapNode* found = find(_root, k);
     if (found != nullptr) {
@@ -304,6 +308,38 @@ public:
     return candidate;
   }
 
+  TreapNode* closest_gt(const K& key) {
+    TreapNode* candidate = nullptr;
+    TreapNode* pos = _root;
+    while (pos != nullptr) {
+      int cmp_r = COMPARATOR::cmp(pos->key(), key);
+      if (cmp_r > 0) {
+        // Found a match, try to find a better one.
+        candidate = pos;
+        pos = pos->_left;
+      } else if (cmp_r <= 0) {
+        pos = pos->_right;
+      }
+    }
+    return candidate;
+  }
+
+  struct Range {
+    TreapNode* start;
+    TreapNode* end;
+    Range(TreapNode* start, TreapNode* end)
+    : start(start), end(end) {}
+  };
+
+  // Return the range [start, end)
+  // where start->key() <= addr < end->key().
+  // Failure to find the range leads to start and/or end being null.
+  Range find_enclosing_range(K addr) {
+    TreapNode* start = closest_leq(addr);
+    TreapNode* end = closest_gt(addr);
+    return Range(start, end);
+  }
+
   // Visit all TreapNodes in ascending key order.
   template<typename F>
   void visit_in_order(F f) const {

src/hotspot/share/nmt/vmatree.cpp

@@ -26,6 +26,7 @@
 #include "precompiled.hpp"
 #include "logging/log.hpp"
 #include "nmt/vmatree.hpp"
+#include "utilities/globalDefinitions.hpp"
 #include "utilities/growableArray.hpp"
 
 const VMATree::RegionData VMATree::empty_regiondata{NativeCallStackStorage::StackIndex{}, mtNone};
@@ -82,12 +83,12 @@ VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType
 
     // 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;
+    stB.out = out_state(leqA_n);
 
     // Direct address match.
     if (leqA_n->key() == A) {
       // Take over in state from old address.
-      stA.in = leqA_n->val().in;
+      stA.in = in_state(leqA_n);
 
       // We may now be able to merge two regions:
       // If the node's old state matches the new, it becomes a noop. That happens, for example,
@@ -113,7 +114,7 @@ VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType
       // We add a new node, but only if there would be a state change. If there would not be a
       // state change, we just omit the node.
       // That happens, for example, when reserving within an already reserved region with identical metadata.
-      stA.in = leqA_n->val().out; // .. and the region's prior state is the incoming state
+      stA.in = out_state(leqA_n); // .. and the region's prior state is the incoming state
       if (stA.is_noop()) {
         // Nothing to do.
       } else {
@@ -134,7 +135,7 @@ VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType
   // outgoing state.
   _tree.visit_range_in_order(A + 1, B + 1, [&](TreapNode* head) {
     int cmp_B = PositionComparator::cmp(head->key(), B);
-    stB.out = head->val().out;
+    stB.out = out_state(head);
     if (cmp_B < 0) {
       // Record all nodes preceding B.
       to_be_deleted_inbetween_a_b.push({head->key(), head->val()});
@@ -215,3 +216,99 @@ VMATree::SummaryDiff VMATree::register_mapping(position A, position B, StateType
   }
   return diff;
 }
+
+#ifdef ASSERT
+void VMATree::print_on(outputStream* out) {
+  visit_in_order([&](TreapNode* current) {
+    out->print(SIZE_FORMAT " (%s) - %s - ", current->key(), NMTUtil::tag_to_name(out_state(current).mem_tag()),
+               statetype_to_string(out_state(current).type()));
+  });
+  out->cr();
+}
+#endif
+
+VMATree::SummaryDiff VMATree::set_tag(const position start, const size size, const MemTag tag) {
+  auto pos = [](TreapNode* n) { return n->key(); };
+  position from = start;
+  position end  = from+size;
+  size_t remsize = size;
+  VMATreap::Range range(nullptr, nullptr);
+
+  // Find the next range to adjust and set range, remsize and from
+  // appropriately. If it returns false, there is no valid next range.
+  auto find_next_range = [&]() -> bool {
+    range = _tree.find_enclosing_range(from);
+    if ((range.start == nullptr && range.end == nullptr) ||
+        (range.start != nullptr && range.end == nullptr)) {
+      // There is no range containing the starting address
+      assert(range.start->val().out.type() == StateType::Released, "must be");
+      return false;
+    } else if (range.start == nullptr && range.end != nullptr) {
+      position found_end = pos(range.end);
+      if (found_end >= end) {
+        // The found address is outside of our range, we can end now.
+        return false;
+      }
+      // There is at least one range [found_end, ?) which starts within [start, end)
+      // Use this as the range instead.
+      range = _tree.find_enclosing_range(found_end);
+      remsize = end - found_end;
+      from = found_end;
+    }
+    return true;
+  };
+
+  bool success = find_next_range();
+  if (!success) return SummaryDiff();
+  assert(range.start != nullptr && range.end != nullptr, "must be");
+
+  end = MIN2(from + remsize, pos(range.end));
+  IntervalState& out = out_state(range.start);
+  StateType type = out.type();
+
+  SummaryDiff diff;
+  // Ignore any released ranges, these must be mtNone and have no stack
+  if (type != StateType::Released) {
+    RegionData new_data = RegionData(out.stack(), tag);
+    SummaryDiff result = register_mapping(from, end, type, new_data);
+    diff.add(result);
+  }
+
+  remsize = remsize - (end - from);
+  from = end;
+
+  // If end < from + sz then there are multiple ranges for which to set the flag.
+  while (end < from + remsize) {
+    // Using register_mapping may invalidate the already found range, so we must
+    // use find_next_range repeatedly
+    bool success = find_next_range();
+    if (!success) return diff;
+    assert(range.start != nullptr && range.end != nullptr, "must be");
+
+    end = MIN2(from + remsize, pos(range.end));
+    IntervalState& out = out_state(range.start);
+    StateType type = out.type();
+
+    if (type != StateType::Released) {
+      RegionData new_data = RegionData(out.stack(), tag);
+      SummaryDiff result = register_mapping(from, end, type, new_data);
+      diff.add(result);
+    }
+    remsize = remsize - (end - from);
+    from = end;
+  }
+
+  return diff;
+}
+
+#ifdef ASSERT
+void VMATree::SummaryDiff::print_on(outputStream* out) {
+  for (int i = 0; i < mt_number_of_tags; i++) {
+    if (tag[i].reserve == 0 && tag[i].commit == 0) {
+      continue;
+    }
+    out->print_cr("Tag %s R: " INT64_FORMAT " C: " INT64_FORMAT, NMTUtil::tag_to_enum_name((MemTag)i), tag[i].reserve,
+                  tag[i].commit);
+  }
+}
+#endif

src/hotspot/share/nmt/vmatree.hpp

@@ -26,10 +26,11 @@
 #ifndef SHARE_NMT_VMATREE_HPP
 #define SHARE_NMT_VMATREE_HPP
 
+#include "nmt/memTag.hpp"
 #include "nmt/nmtNativeCallStackStorage.hpp"
 #include "nmt/nmtTreap.hpp"
-#include "runtime/os.hpp"
 #include "utilities/globalDefinitions.hpp"
+#include "utilities/ostream.hpp"
 #include <cstdint>
 
 // A VMATree stores a sequence of points on the natural number line.
@@ -42,6 +43,7 @@ class VMATree {
   // A position in memory.
 public:
   using position = size_t;
+  using size = size_t;
 
   class PositionComparator {
   public:
@@ -140,6 +142,14 @@ public:
 private:
   VMATreap _tree;
 
+  static IntervalState& in_state(TreapNode* node) {
+    return node->val().in;
+  }
+
+  static IntervalState& out_state(TreapNode* node) {
+    return node->val().out;
+  }
+
   // AddressState saves the necessary information for performing online summary accounting.
   struct AddressState {
     position address;
@@ -162,6 +172,7 @@ public:
     delta reserve;
     delta commit;
   };
+
   struct SummaryDiff {
     SingleDiff tag[mt_number_of_tags];
     SummaryDiff() {
@@ -169,26 +180,47 @@ public:
         tag[i] = SingleDiff{0, 0};
       }
     }
+
+    void add(SummaryDiff& other) {
+      for (int i = 0; i < mt_number_of_tags; i++) {
+        tag[i].reserve += other.tag[i].reserve;
+        tag[i].commit += other.tag[i].commit;
+      }
+    }
+
+#ifdef ASSERT
+    void print_on(outputStream* out);
+#endif
   };
 
  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, false);
+  SummaryDiff reserve_mapping(position from, size size, const RegionData& metadata) {
+    return register_mapping(from, from + size, 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, size size, const RegionData& metadata, bool use_tag_inplace = false) {
+    return register_mapping(from, from + size, StateType::Committed, metadata, use_tag_inplace);
   }
 
-  SummaryDiff uncommit_mapping(position from, position sz, const RegionData& metadata) {
-    return register_mapping(from, from + sz, StateType::Reserved, metadata, true);
+  // Given an interval and a tag, find all reserved and committed ranges at least
+  // partially contained within that interval and set their tag to the one provided.
+  // This may cause merging and splitting of ranges.
+  // Released regions are ignored.
+  SummaryDiff set_tag(position from, size size, MemTag tag);
+
+  SummaryDiff uncommit_mapping(position from, size size, const RegionData& metadata) {
+    return register_mapping(from, from + size, StateType::Reserved, metadata, true);
   }
 
-  SummaryDiff release_mapping(position from, position sz) {
-    return register_mapping(from, from + sz, StateType::Released, VMATree::empty_regiondata);
+  SummaryDiff release_mapping(position from, size size) {
+    return register_mapping(from, from + size, StateType::Released, VMATree::empty_regiondata);
+  }
+
+  VMATreap& tree() {
+    return _tree;
   }
 
 public:
@@ -196,6 +228,11 @@ public:
   void visit_in_order(F f) const {
     _tree.visit_in_order(f);
   }
+
+#ifdef ASSERT
+  void print_on(outputStream* out);
+#endif
+
 };
 
 #endif

test/hotspot/gtest/nmt/test_vmatree.cpp

@@ -284,6 +284,188 @@ TEST_VM_F(NMTVMATreeTest, LowLevel) {
   }
 }
 
+TEST_VM_F(NMTVMATreeTest, SetTag) {
+  using State = VMATree::StateType;
+  struct testrange {
+    VMATree::position from;
+    VMATree::position to;
+    MemTag tag;
+    NCS::StackIndex stack;
+    State state;
+  };
+
+  // Take a sorted list of testranges and check that those and only those are found in the tree.
+  auto expect_equivalent_form = [&](auto& expected, VMATree& tree) {
+    // With auto& our arrays do not deteriorate to pointers but are kept as testrange[N]
+    // so this actually works!
+    int len = sizeof(expected) / sizeof(testrange);
+    VMATree::position previous_to = 0;
+    for (int i = 0; i < len; i++) {
+      testrange expect = expected[i];
+      assert(previous_to == 0 || previous_to <= expect.from, "the expected list must be sorted");
+      previous_to = expect.to;
+
+      VMATree::VMATreap::Range found = tree.tree().find_enclosing_range(expect.from);
+      ASSERT_NE(nullptr, found.start);
+      ASSERT_NE(nullptr, found.end);
+      // Same region
+      EXPECT_EQ(expect.from, found.start->key());
+      EXPECT_EQ(expect.to, found.end->key());
+      // Same tag
+      EXPECT_EQ(expect.tag, found.start->val().out.mem_tag());
+      EXPECT_EQ(expect.tag, found.end->val().in.mem_tag());
+      // Same stack
+      EXPECT_EQ(expect.stack, found.start->val().out.stack());
+      EXPECT_EQ(expect.stack, found.end->val().in.stack());
+      // Same state
+      EXPECT_EQ(expect.state, found.start->val().out.type());
+      EXPECT_EQ(expect.state, found.end->val().in.type());
+    }
+    // expected must cover all nodes
+    EXPECT_EQ(len+1, tree.tree().size());
+  };
+  NCS::StackIndex si = NCS::StackIndex();
+  Tree::RegionData rd(si, mtNone);
+
+  { // The gc/cds case with only reserved data
+    testrange expected[2]{
+        {  0, 500,          mtGC, si, State::Reserved},
+        {500, 600, mtClassShared, si, State::Reserved}
+    };
+    VMATree tree;
+
+    tree.reserve_mapping(0, 600, rd);
+
+    tree.set_tag(0, 500, mtGC);
+    tree.set_tag(500, 100, mtClassShared);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Now let's add in some committed data
+    testrange expected[]{
+        {  0, 100,          mtGC, si, State::Reserved},
+        {100, 225,          mtGC, si, State::Committed},
+        {225, 500,          mtGC, si, State::Reserved},
+        {500, 550, mtClassShared, si, State::Reserved},
+        {550, 560, mtClassShared, si, State::Committed},
+        {560, 565, mtClassShared, si, State::Reserved},
+        {565, 575, mtClassShared, si, State::Committed},
+        {575, 600, mtClassShared, si, State::Reserved}
+    };
+    VMATree tree;
+
+    tree.reserve_mapping(0, 600, rd);
+    // The committed areas
+    tree.commit_mapping(100, 125, rd);
+    tree.commit_mapping(550, 10, rd);
+    tree.commit_mapping(565, 10, rd);
+    // OK, set tag
+    tree.set_tag(0, 500, mtGC);
+    tree.set_tag(500, 100, mtClassShared);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Setting the tag for adjacent regions with same stacks should merge the regions
+    testrange expected[]{
+        {0, 200, mtGC, si, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData gc(si, mtGC);
+    Tree::RegionData compiler(si, mtCompiler);
+    tree.reserve_mapping(0, 100, gc);
+    tree.reserve_mapping(100, 100, compiler);
+    tree.set_tag(0, 200, mtGC);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Setting the tag for adjacent regions with different stacks should NOT merge the regions
+    NCS::StackIndex si1 = 1;
+    NCS::StackIndex si2 = 2;
+    testrange expected[]{
+        {  0, 100, mtGC, si1, State::Reserved},
+        {100, 200, mtGC, si2, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData gc(si1, mtGC);
+    Tree::RegionData compiler(si2, mtCompiler);
+    tree.reserve_mapping(0, 100, gc);
+    tree.reserve_mapping(100, 100, compiler);
+    tree.set_tag(0, 200, mtGC);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Setting the tag in the middle of a range causes a split
+    testrange expected[]{
+        {  0, 100, mtCompiler, si, State::Reserved},
+        {100, 150,       mtGC, si, State::Reserved},
+        {150, 200, mtCompiler, si, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData compiler(si, mtCompiler);
+    tree.reserve_mapping(0, 200, compiler);
+    tree.set_tag(100, 50, mtGC);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Setting the tag in between two ranges causes a split
+    testrange expected[]{
+        {  0,  75,       mtGC, si, State::Reserved},
+        { 75, 125,    mtClass, si, State::Reserved},
+        {125, 200, mtCompiler, si, State::Reserved},
+    };
+    VMATree tree;
+    Tree::RegionData gc(si, mtGC);
+    Tree::RegionData compiler(si, mtCompiler);
+    tree.reserve_mapping(0, 100, gc);
+    tree.reserve_mapping(100, 100, compiler);
+    tree.set_tag(75, 50, mtClass);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Holes in the address range are acceptable and untouched
+    testrange expected[]{
+        { 0,  50,          mtGC, si, State::Reserved},
+        {50,  75,        mtNone, si, State::Released},
+        {75,  80,          mtGC, si, State::Reserved},
+        {80, 100, mtClassShared, si, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData class_shared(si, mtClassShared);
+    tree.reserve_mapping(0, 50, class_shared);
+    tree.reserve_mapping(75, 25, class_shared);
+    tree.set_tag(0, 80, mtGC);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Check that setting tag with 'hole' not consisting of any regions work
+    testrange expected[]{
+        {10, 20, mtCompiler, si, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData class_shared(si, mtClassShared);
+    tree.reserve_mapping(10, 10, class_shared);
+    tree.set_tag(0, 100, mtCompiler);
+    expect_equivalent_form(expected, tree);
+  }
+
+  { // Check that multiple holes still work
+    testrange expected[]{
+        { 0,   1,   mtGC, si, State::Reserved},
+        { 1,  50, mtNone, si, State::Released},
+        {50,  75,   mtGC, si, State::Reserved},
+        {75,  99, mtNone, si, State::Released},
+        {99, 100,   mtGC, si, State::Reserved}
+    };
+    VMATree tree;
+    Tree::RegionData class_shared(si, mtClassShared);
+    tree.reserve_mapping(0, 100, class_shared);
+    tree.release_mapping(1, 49);
+    tree.release_mapping(75, 24);
+    tree.set_tag(0, 100, mtGC);
+    expect_equivalent_form(expected, tree);
+  }
+}
+
 // Tests for summary accounting
 TEST_VM_F(NMTVMATreeTest, SummaryAccounting) {
   { // Fully enclosed re-reserving works correctly.
@@ -330,7 +512,7 @@ TEST_VM_F(NMTVMATreeTest, SummaryAccounting) {
     diff = all_diff.tag[NMTUtil::tag_to_index(mtTest)];
     EXPECT_EQ(100, diff.reserve);
   }
-  { // Adjacent reserved mappings with different flags
+  { // Adjacent reserved mappings with different tags
   Tree::RegionData rd(NCS::StackIndex(), mtTest);
     Tree::RegionData rd2(NCS::StackIndex(), mtNMT);
     Tree tree;