--- old/src/share/vm/gc/shared/threadLocalAllocBuffer.cpp	2017-06-21 11:43:11.167543290 -0700
+++ new/src/share/vm/gc/shared/threadLocalAllocBuffer.cpp	2017-06-21 11:43:11.075543632 -0700
@@ -29,6 +29,7 @@
 #include "memory/resourceArea.hpp"
 #include "memory/universe.inline.hpp"
 #include "oops/oop.inline.hpp"
+#include "runtime/heapMonitoring.hpp"
 #include "runtime/thread.inline.hpp"
 #include "utilities/copy.hpp"
 
@@ -121,6 +122,7 @@
       set_top(NULL);
       set_pf_top(NULL);
       set_end(NULL);
+      set_actual_end(NULL);
     }
   }
   assert(!(retire || ZeroTLAB)  ||
@@ -183,7 +185,9 @@
   set_top(top);
   set_pf_top(top);
   set_end(end);
+  set_actual_end(end);
   invariants();
+  _bytes_until_sample = 0;
 }
 
 void ThreadLocalAllocBuffer::initialize() {
@@ -306,12 +310,56 @@
   guarantee(p == top(), "end of last object must match end of space");
 }
 
+void ThreadLocalAllocBuffer::pick_next_sample() {
+  if (!HeapMonitoring::enabled()) {
+    return;
+  }
+
+  if (bytes_until_sample() == 0) {
+    HeapMonitoring::pick_next_sample(bytes_until_sample_addr());
+  }
+
+  // Finally, fix up the sampling bytes left and _end.
+  size_t heap_words_remaining = _end - _top;
+  size_t bytes_left = bytes_until_sample();
+  size_t words_until_sample = bytes_left / HeapWordSize;
+
+  if (heap_words_remaining > words_until_sample) {
+    set_end(_top + words_until_sample);
+    set_bytes_until_sample(0);
+  } else {
+    bytes_left -= heap_words_remaining * HeapWordSize;
+    set_bytes_until_sample(bytes_left);
+  }
+}
+
 Thread* ThreadLocalAllocBuffer::myThread() {
   return (Thread*)(((char *)this) +
                    in_bytes(start_offset()) -
                    in_bytes(Thread::tlab_start_offset()));
 }
 
+void ThreadLocalAllocBuffer::handle_sample(Thread* thread, HeapWord* result,
+                                           size_t size) {
+  if (!HeapMonitoring::enabled()) {
+    return;
+  }
+
+  set_bytes_until_sample(bytes_until_sample() - size);
+
+  // Should we sample now?
+  if (should_sample()) {
+    set_back_actual_end();
+    HeapMonitoring::object_alloc_do_sample(thread,
+                                           reinterpret_cast<oopDesc*>(result),
+                                           size);
+    pick_next_sample();
+  } else {
+    // Update the size and end fields.
+    set_back_actual_end();
+    pick_next_sample();
+  }
+}
 
 GlobalTLABStats::GlobalTLABStats() :
   _allocating_threads_avg(TLABAllocationWeight) {