8297718: Make NMT free:ing protocol more granular

Reviewed-by: stuefe, gziemski

Author: jdksjolen

src/hotspot/share/runtime/os.cpp

@@ -62,6 +62,7 @@
 #include "runtime/vm_version.hpp"
 #include "services/attachListener.hpp"
 #include "services/mallocTracker.hpp"
+#include "services/mallocHeader.inline.hpp"
 #include "services/memTracker.hpp"
 #include "services/nmtPreInit.hpp"
 #include "services/nmtCommon.hpp"
@@ -714,30 +715,31 @@ void* os::realloc(void *memblock, size_t size, MEMFLAGS memflags, const NativeCa
       return NULL;
     }
 
-    const size_t old_size = MallocTracker::malloc_header(memblock)->size();
-
-    // De-account the old block from NMT *before* calling the real realloc(3) since it
-    // may invalidate old block including its header. This will also perform integrity checks
-    // on the old block (e.g. overwriters) and mark the old header as dead.
-    void* const old_outer_ptr = MemTracker::record_free(memblock);
+    // Perform integrity checks on and mark the old block as dead *before* calling the real realloc(3) since it
+    // may invalidate the old block, including its header.
+    MallocHeader* header = MallocTracker::malloc_header(memblock);
+    header->assert_block_integrity(); // Assert block hasn't been tampered with.
+    const MallocHeader::FreeInfo free_info = header->free_info();
+    header->mark_block_as_dead();
 
     // the real realloc
-    ALLOW_C_FUNCTION(::realloc, void* const new_outer_ptr = ::realloc(old_outer_ptr, new_outer_size);)
+    ALLOW_C_FUNCTION(::realloc, void* const new_outer_ptr = ::realloc(header, new_outer_size);)
 
     if (new_outer_ptr == NULL) {
-      // If realloc(3) failed, the old block still exists. We must re-instantiate the old
-      // NMT header then, since we marked it dead already. Otherwise subsequent os::realloc()
-      // or os::free() calls would trigger block integrity asserts.
-      void* p = MemTracker::record_malloc(old_outer_ptr, old_size, memflags, stack);
-      assert(p == memblock, "sanity");
-      return NULL;
+      // realloc(3) failed and the block still exists.
+      // We have however marked it as dead, revert this change.
+      header->revive();
+      return nullptr;
     }
+    // realloc(3) succeeded, variable header now points to invalid memory and we need to deaccount the old block.
+    MemTracker::deaccount(free_info);
 
-    // After a successful realloc(3), we re-account the resized block with its new size
-    // to NMT. This re-instantiates the NMT header.
+    // After a successful realloc(3), we account the resized block with its new size
+    // to NMT.
     void* const new_inner_ptr = MemTracker::record_malloc(new_outer_ptr, size, memflags, stack);
 
 #ifdef ASSERT
+    size_t old_size = free_info.size;
     if (old_size < size) {
       // We also zap the newly extended region.
       ::memset((char*)new_inner_ptr + old_size, uninitBlockPad, size - old_size);
@@ -774,7 +776,7 @@ void  os::free(void *memblock) {
 
   DEBUG_ONLY(break_if_ptr_caught(memblock);)
 
-  // If NMT is enabled, this checks for heap overwrites, then de-accounts the old block.
+  // When NMT is enabled this checks for heap overwrites, then deaccounts the old block.
   void* const old_outer_ptr = MemTracker::record_free(memblock);
 
   ALLOW_C_FUNCTION(::free, ::free(old_outer_ptr);)

src/hotspot/share/services/mallocHeader.hpp

@@ -88,7 +88,7 @@ class outputStream;
  */
 
 class MallocHeader {
-
+  NONCOPYABLE(MallocHeader);
   NOT_LP64(uint32_t _alt_canary);
   const size_t _size;
   const uint32_t _mst_marker;
@@ -115,6 +115,12 @@ class MallocHeader {
   void set_footer(uint16_t v)       { footer_address()[0] = v >> 8; footer_address()[1] = (uint8_t)v; }
 
  public:
+  // Contains all of the necessary data to to deaccount block with NMT.
+  struct FreeInfo {
+    const size_t size;
+    const MEMFLAGS flags;
+    const uint32_t mst_marker;
+  };
 
   inline MallocHeader(size_t size, MEMFLAGS flags, uint32_t mst_marker);
 
@@ -123,7 +129,12 @@ class MallocHeader {
   inline uint32_t mst_marker() const { return _mst_marker; }
   bool get_stack(NativeCallStack& stack) const;
 
+  // Return the necessary data to deaccount the block with NMT.
+  FreeInfo free_info() {
+    return FreeInfo{this->size(), this->flags(), this->mst_marker()};
+  }
   inline void mark_block_as_dead();
+  inline void revive();
 
   // If block is broken, fill in a short descriptive text in out,
   // an option pointer to the corruption in p_corruption, and return false.

src/hotspot/share/services/mallocHeader.inline.hpp

@@ -45,6 +45,16 @@ inline MallocHeader::MallocHeader(size_t size, MEMFLAGS flags, uint32_t mst_mark
   set_footer(_footer_canary_life_mark); // set after initializing _size
 }
 
+inline void MallocHeader::revive() {
+  assert(_canary == _header_canary_dead_mark, "must be dead");
+  assert(get_footer() == _footer_canary_dead_mark, "must be dead");
+  NOT_LP64(assert(_alt_canary == _header_alt_canary_dead_mark, "must be dead"));
+  _canary = _header_canary_life_mark;
+  NOT_LP64(_alt_canary = _header_alt_canary_life_mark);
+  set_footer(_footer_canary_life_mark);
+}
+
+// The effects of this method must be reversible with MallocHeader::revive()
 inline void MallocHeader::mark_block_as_dead() {
   _canary = _header_canary_dead_mark;
   NOT_LP64(_alt_canary = _header_alt_canary_dead_mark);

src/hotspot/share/services/mallocTracker.cpp

@@ -180,23 +180,27 @@ void* MallocTracker::record_malloc(void* malloc_base, size_t size, MEMFLAGS flag
   return memblock;
 }
 
-void* MallocTracker::record_free(void* memblock) {
+void* MallocTracker::record_free_block(void* memblock) {
   assert(MemTracker::enabled(), "Sanity");
   assert(memblock != NULL, "precondition");
 
   MallocHeader* const header = malloc_header(memblock);
   header->assert_block_integrity();
 
-  MallocMemorySummary::record_free(header->size(), header->flags());
-  if (MemTracker::tracking_level() == NMT_detail) {
-    MallocSiteTable::deallocation_at(header->size(), header->mst_marker());
-  }
+  deaccount(header->free_info());
 
   header->mark_block_as_dead();
 
   return (void*)header;
 }
 
+void MallocTracker::deaccount(MallocHeader::FreeInfo free_info) {
+  MallocMemorySummary::record_free(free_info.size, free_info.flags);
+  if (MemTracker::tracking_level() == NMT_detail) {
+    MallocSiteTable::deallocation_at(free_info.size, free_info.mst_marker);
+  }
+}
+
 // Given a pointer, if it seems to point to the start of a valid malloced block,
 // print the block. Note that since there is very low risk of memory looking
 // accidentally like a valid malloc block header (canaries and all) this is not

src/hotspot/share/services/mallocTracker.hpp

@@ -303,8 +303,11 @@ class MallocTracker : AllStatic {
   static void* record_malloc(void* malloc_base, size_t size, MEMFLAGS flags,
     const NativeCallStack& stack);
 
-  // Record free on specified memory block
-  static void* record_free(void* memblock);
+  // Given a block returned by os::malloc() or os::realloc():
+  // deaccount block from NMT, mark its header as dead and return pointer to header.
+  static void* record_free_block(void* memblock);
+  // Given the free info from a block, de-account block from NMT.
+  static void deaccount(MallocHeader::FreeInfo free_info);
 
   static inline void record_new_arena(MEMFLAGS flags) {
     MallocMemorySummary::record_new_arena(flags);

src/hotspot/share/services/memTracker.hpp

@@ -109,7 +109,11 @@ class MemTracker : AllStatic {
     if (!enabled()) {
       return memblock;
     }
-    return MallocTracker::record_free(memblock);
+    return MallocTracker::record_free_block(memblock);
+  }
+  static inline void deaccount(MallocHeader::FreeInfo free_info) {
+    assert(enabled(), "NMT must be enabled");
+    MallocTracker::deaccount(free_info);
   }
 
   // Record creation of an arena

test/hotspot/gtest/nmt/test_nmt_cornercases.cpp

@@ -141,3 +141,16 @@ TEST_VM(NMT, random_reallocs) {
 
   os::free(p);
 }
+
+TEST_VM(NMT, HeaderKeepsIntegrityAfterRevival) {
+  if (!MemTracker::enabled()) {
+    return;
+  }
+  size_t some_size = 16;
+  void* p = os::malloc(some_size, mtTest);
+  ASSERT_NOT_NULL(p) << "Failed to malloc()";
+  MallocHeader* hdr = MallocTracker::malloc_header(p);
+  hdr->mark_block_as_dead();
+  hdr->revive();
+  check_expected_malloc_header(p, mtTest, some_size);
+}