8340103: Add internal set_flag function to VMATree

Reviewed-by: stuefe, azafari, gziemski

Author: jdksjolen

src/hotspot/share/nmt/nmtTreap.hpp

@@ -234,6 +234,10 @@ class Treap {
     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 @@ class Treap {
     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 @@ class VMATree {
 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,40 +172,67 @@ class VMATree {
     delta reserve;
     delta commit;
   };
+
   struct SummaryDiff {
     SingleDiff tag[mt_number_of_tags];
     SummaryDiff() {
       for (int i = 0; i < mt_number_of_tags; i++) {
         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, size size, const RegionData& metadata, bool use_tag_inplace = false) {
+    return register_mapping(from, from + size, StateType::Committed, metadata, use_tag_inplace);
   }
 
-  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);
+  // 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 uncommit_mapping(position from, position sz, const RegionData& metadata) {
-    return register_mapping(from, from + sz, StateType::Reserved, metadata, true);
+  SummaryDiff release_mapping(position from, size size) {
+    return register_mapping(from, from + size, StateType::Released, VMATree::empty_regiondata);
   }
 
-  SummaryDiff release_mapping(position from, position sz) {
-    return register_mapping(from, from + sz, StateType::Released, VMATree::empty_regiondata);
+  VMATreap& tree() {
+    return _tree;
   }
 
 public:
   template<typename F>
   void visit_in_order(F f) const {
     _tree.visit_in_order(f);
   }
+
+#ifdef ASSERT
+  void print_on(outputStream* out);
+#endif
+
 };
 
 #endif