8356372: JVMTI heap sampling not working properly with outside TLAB allocations

Co-authored-by: Stefan Johansson <sjohanss@openjdk.org>
Reviewed-by: sjohanss, sspitsyn

Author: stefank

src/hotspot/share/gc/g1/g1YoungGCPreEvacuateTasks.cpp

@@ -59,7 +59,7 @@ class G1PreEvacuateCollectionSetBatchTask::JavaThreadRetireTLABAndFlushLogs : pu
       BarrierSet::barrier_set()->make_parsable((JavaThread*)thread);
       // Retire TLABs.
       if (UseTLAB) {
-        thread->tlab().retire(&_tlab_stats);
+        thread->retire_tlab(&_tlab_stats);
       }
       // Concatenate logs.
       G1DirtyCardQueueSet& qset = G1BarrierSet::dirty_card_queue_set();

src/hotspot/share/gc/shared/collectedHeap.cpp

@@ -517,8 +517,8 @@ void CollectedHeap::ensure_parsability(bool retire_tlabs) {
   for (JavaThreadIteratorWithHandle jtiwh; JavaThread *thread = jtiwh.next();) {
     BarrierSet::barrier_set()->make_parsable(thread);
     if (UseTLAB) {
-      if (retire_tlabs) {
-        thread->tlab().retire(&stats);
+      if (retire_tlabs || ZeroTLAB) {
+        thread->retire_tlab(&stats);
       } else {
         thread->tlab().make_parsable();
       }

src/hotspot/share/gc/shared/memAllocator.cpp

@@ -51,7 +51,6 @@ class MemAllocator::Allocation: StackObj {
   bool                _overhead_limit_exceeded;
   bool                _allocated_outside_tlab;
   size_t              _allocated_tlab_size;
-  bool                _tlab_end_reset_for_sample;
 
   bool check_out_of_memory();
   void verify_before();
@@ -74,8 +73,7 @@ class MemAllocator::Allocation: StackObj {
       _obj_ptr(obj_ptr),
       _overhead_limit_exceeded(false),
       _allocated_outside_tlab(false),
-      _allocated_tlab_size(0),
-      _tlab_end_reset_for_sample(false)
+      _allocated_tlab_size(0)
   {
     assert(Thread::current() == allocator._thread, "do not pass MemAllocator across threads");
     verify_before();
@@ -174,33 +172,31 @@ void MemAllocator::Allocation::notify_allocation_jvmti_sampler() {
     return;
   }
 
-  if (!_allocated_outside_tlab && _allocated_tlab_size == 0 && !_tlab_end_reset_for_sample) {
-    // Sample if it's a non-TLAB allocation, or a TLAB allocation that either refills the TLAB
-    // or expands it due to taking a sampler induced slow path.
-    return;
-  }
+  ThreadHeapSampler& heap_sampler = _thread->heap_sampler();
+  ThreadLocalAllocBuffer& tlab = _thread->tlab();
 
-  // If we want to be sampling, protect the allocated object with a Handle
-  // before doing the callback. The callback is done in the destructor of
-  // the JvmtiSampledObjectAllocEventCollector.
-  size_t bytes_since_last = 0;
+  // Log sample decision
+  heap_sampler.log_sample_decision(tlab.top());
 
-  {
+  if (heap_sampler.should_sample(tlab.top())) {
+    // If we want to be sampling, protect the allocated object with a Handle
+    // before doing the callback. The callback is done in the destructor of
+    // the JvmtiSampledObjectAllocEventCollector.
     PreserveObj obj_h(_thread, _obj_ptr);
     JvmtiSampledObjectAllocEventCollector collector;
-    size_t size_in_bytes = _allocator._word_size * HeapWordSize;
-    ThreadLocalAllocBuffer& tlab = _thread->tlab();
 
-    if (!_allocated_outside_tlab) {
-      bytes_since_last = tlab.bytes_since_last_sample_point();
-    }
+    // Perform the sampling
+    heap_sampler.sample(obj_h(), tlab.top());
 
-    _thread->heap_sampler().check_for_sampling(obj_h(), size_in_bytes, bytes_since_last);
+    // Note that after this point all the TLAB can have been retired, and agent
+    // code can run and allocate, don't rely on earlier calculations involving
+    // the TLAB.
   }
 
-  if (_tlab_end_reset_for_sample || _allocated_tlab_size != 0) {
-    // Tell tlab to forget bytes_since_last if we passed it to the heap sampler.
-    _thread->tlab().set_sample_end(bytes_since_last != 0);
+  // Set a new sampling point in the TLAB if it fits in the current TLAB
+  const size_t words_until_sample = heap_sampler.bytes_until_sample(tlab.top()) / HeapWordSize;
+  if (words_until_sample <= tlab.free()) {
+    tlab.set_sampling_point(tlab.top() + words_until_sample);
   }
 }
 
@@ -249,6 +245,7 @@ HeapWord* MemAllocator::mem_allocate_outside_tlab(Allocation& allocation) const
 
   size_t size_in_bytes = _word_size * HeapWordSize;
   _thread->incr_allocated_bytes(size_in_bytes);
+  _thread->heap_sampler().inc_outside_tlab_bytes(size_in_bytes);
 
   return mem;
 }
@@ -262,12 +259,16 @@ HeapWord* MemAllocator::mem_allocate_inside_tlab_slow(Allocation& allocation) co
   ThreadLocalAllocBuffer& tlab = _thread->tlab();
 
   if (JvmtiExport::should_post_sampled_object_alloc()) {
+    // When sampling we artificially set the TLAB end to the sample point.
+    // When we hit that point it looks like the TLAB is full, but it's
+    // not necessarily the case. Set the real end and retry the allocation.
+
+    // Undo previous adjustment of end.
+    // Note that notify_allocation_jvmti_sampler will set a new sample point.
     tlab.set_back_allocation_end();
-    mem = tlab.allocate(_word_size);
 
-    // We set back the allocation sample point to try to allocate this, reset it
-    // when done.
-    allocation._tlab_end_reset_for_sample = true;
+    // Retry the TLAB allocation with the proper end
+    mem = tlab.allocate(_word_size);
 
     if (mem != nullptr) {
       return mem;
@@ -282,11 +283,16 @@ HeapWord* MemAllocator::mem_allocate_inside_tlab_slow(Allocation& allocation) co
   }
 
   // Discard tlab and allocate a new one.
+
+  // Record the amount wasted
+  tlab.record_refill_waste();
+
+  // Retire the current TLAB
+  _thread->retire_tlab();
+
   // To minimize fragmentation, the last TLAB may be smaller than the rest.
   size_t new_tlab_size = tlab.compute_size(_word_size);
 
-  tlab.retire_before_allocation();
-
   if (new_tlab_size == 0) {
     return nullptr;
   }
@@ -317,7 +323,8 @@ HeapWord* MemAllocator::mem_allocate_inside_tlab_slow(Allocation& allocation) co
     Copy::fill_to_words(mem + hdr_size, allocation._allocated_tlab_size - hdr_size, badHeapWordVal);
   }
 
-  tlab.fill(mem, mem + _word_size, allocation._allocated_tlab_size);
+  _thread->fill_tlab(mem, _word_size, allocation._allocated_tlab_size);
+
   return mem;
 }
 

src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp

@@ -49,7 +49,6 @@ ThreadLocalAllocBuffer::ThreadLocalAllocBuffer() :
   _desired_size(0),
   _refill_waste_limit(0),
   _allocated_before_last_gc(0),
-  _bytes_since_last_sample_point(0),
   _number_of_refills(0),
   _refill_waste(0),
   _gc_waste(0),
@@ -125,11 +124,7 @@ void ThreadLocalAllocBuffer::insert_filler() {
 void ThreadLocalAllocBuffer::make_parsable() {
   if (end() != nullptr) {
     invariants();
-    if (ZeroTLAB) {
-      retire();
-    } else {
-      insert_filler();
-    }
+    insert_filler();
   }
 }
 
@@ -140,15 +135,13 @@ void ThreadLocalAllocBuffer::retire(ThreadLocalAllocStats* stats) {
 
   if (end() != nullptr) {
     invariants();
-    thread()->incr_allocated_bytes(used_bytes());
     insert_filler();
     initialize(nullptr, nullptr, nullptr);
   }
 }
 
-void ThreadLocalAllocBuffer::retire_before_allocation() {
+void ThreadLocalAllocBuffer::record_refill_waste() {
   _refill_waste += (unsigned int)remaining();
-  retire();
 }
 
 void ThreadLocalAllocBuffer::resize() {
@@ -312,24 +305,6 @@ void ThreadLocalAllocBuffer::print_stats(const char* tag) {
             _refill_waste * HeapWordSize);
 }
 
-void ThreadLocalAllocBuffer::set_sample_end(bool reset_byte_accumulation) {
-  size_t heap_words_remaining = pointer_delta(_end, _top);
-  size_t bytes_until_sample = thread()->heap_sampler().bytes_until_sample();
-  size_t words_until_sample = bytes_until_sample / HeapWordSize;
-
-  if (reset_byte_accumulation) {
-    _bytes_since_last_sample_point = 0;
-  }
-
-  if (heap_words_remaining > words_until_sample) {
-    HeapWord* new_end = _top + words_until_sample;
-    set_end(new_end);
-    _bytes_since_last_sample_point += bytes_until_sample;
-  } else {
-    _bytes_since_last_sample_point += heap_words_remaining * HeapWordSize;
-  }
-}
-
 Thread* ThreadLocalAllocBuffer::thread() {
   return (Thread*)(((char*)this) + in_bytes(start_offset()) - in_bytes(Thread::tlab_start_offset()));
 }
@@ -338,6 +313,14 @@ void ThreadLocalAllocBuffer::set_back_allocation_end() {
   _end = _allocation_end;
 }
 
+void ThreadLocalAllocBuffer::set_sampling_point(HeapWord* sampling_point) {
+  precond(sampling_point >= _top);
+  precond(sampling_point <= _allocation_end);
+
+  // This will trigger a slow-path, which in turn might take a sample.
+  _end = sampling_point;
+}
+
 HeapWord* ThreadLocalAllocBuffer::hard_end() {
   return _allocation_end + alignment_reserve();
 }

src/hotspot/share/gc/shared/threadLocalAllocBuffer.hpp

@@ -56,7 +56,6 @@ class ThreadLocalAllocBuffer: public CHeapObj<mtThread> {
   size_t    _desired_size;                       // desired size   (including alignment_reserve)
   size_t    _refill_waste_limit;                 // hold onto tlab if free() is larger than this
   size_t    _allocated_before_last_gc;           // total bytes allocated up until the last gc
-  size_t    _bytes_since_last_sample_point;      // bytes since last sample point.
 
   static size_t   _max_size;                          // maximum size of any TLAB
   static int      _reserve_for_allocation_prefetch;   // Reserve at the end of the TLAB
@@ -124,7 +123,6 @@ class ThreadLocalAllocBuffer: public CHeapObj<mtThread> {
   size_t free() const                            { return pointer_delta(end(), top()); }
   // Don't discard tlab if remaining space is larger than this.
   size_t refill_waste_limit() const              { return _refill_waste_limit; }
-  size_t bytes_since_last_sample_point() const   { return _bytes_since_last_sample_point; }
 
   // For external inspection.
   const HeapWord* start_relaxed() const;
@@ -158,17 +156,18 @@ class ThreadLocalAllocBuffer: public CHeapObj<mtThread> {
   // Retire an in-use tlab and optionally collect statistics.
   void retire(ThreadLocalAllocStats* stats = nullptr);
 
-  // Retire in-use tlab before allocation of a new tlab
-  void retire_before_allocation();
+  // Record refill waste before allocating (refilling) with a new TLAB.
+  void record_refill_waste();
 
   // Resize based on amount of allocation, etc.
   void resize();
 
   void fill(HeapWord* start, HeapWord* top, size_t new_size);
   void initialize();
 
+  // Support for TLAB sampling
   void set_back_allocation_end();
-  void set_sample_end(bool reset_byte_accumulation);
+  void set_sampling_point(HeapWord* sampling_point);
 
   static size_t refill_waste_limit_increment();
 

src/hotspot/share/gc/shenandoah/shenandoahHeap.cpp

@@ -1499,6 +1499,9 @@ void ShenandoahHeap::labs_make_parsable() {
   for (JavaThreadIteratorWithHandle jtiwh; JavaThread *t = jtiwh.next(); ) {
     ThreadLocalAllocBuffer& tlab = t->tlab();
     tlab.make_parsable();
+    if (ZeroTLAB) {
+      t->retire_tlab();
+    }
     cl.do_thread(t);
   }
 
@@ -1516,10 +1519,9 @@ void ShenandoahHeap::tlabs_retire(bool resize) {
   ThreadLocalAllocStats stats;
 
   for (JavaThreadIteratorWithHandle jtiwh; JavaThread *t = jtiwh.next(); ) {
-    ThreadLocalAllocBuffer& tlab = t->tlab();
-    tlab.retire(&stats);
+    t->retire_tlab(&stats);
     if (resize) {
-      tlab.resize();
+      t->tlab().resize();
     }
   }
 

src/hotspot/share/gc/shenandoah/shenandoahStackWatermark.cpp

@@ -119,7 +119,7 @@ void ShenandoahStackWatermark::retire_tlab() {
   // Retire TLAB
   if (UseTLAB) {
     _stats.reset();
-    _jt->tlab().retire(&_stats);
+    _jt->retire_tlab(&_stats);
     if (ResizeTLAB) {
       _jt->tlab().resize();
     }

src/hotspot/share/gc/z/zThreadLocalAllocBuffer.cpp

@@ -65,7 +65,7 @@ void ZThreadLocalAllocBuffer::publish_statistics() {
 void ZThreadLocalAllocBuffer::retire(JavaThread* thread, ThreadLocalAllocStats* stats) {
   if (UseTLAB) {
     stats->reset();
-    thread->tlab().retire(stats);
+    thread->retire_tlab(stats);
     if (ResizeTLAB) {
       thread->tlab().resize();
     }

src/hotspot/share/runtime/javaThread.cpp

@@ -968,7 +968,7 @@ void JavaThread::exit(bool destroy_vm, ExitType exit_type) {
   _stack_overflow_state.remove_stack_guard_pages();
 
   if (UseTLAB) {
-    tlab().retire();
+    retire_tlab();
   }
 
   if (JvmtiEnv::environments_might_exist()) {
@@ -1042,7 +1042,7 @@ void JavaThread::cleanup_failed_attach_current_thread(bool is_daemon) {
   _stack_overflow_state.remove_stack_guard_pages();
 
   if (UseTLAB) {
-    tlab().retire();
+    retire_tlab();
   }
 
   Threads::remove(this, is_daemon);

src/hotspot/share/runtime/thread.cpp

@@ -154,6 +154,25 @@ void Thread::initialize_tlab() {
   }
 }
 
+void Thread::retire_tlab(ThreadLocalAllocStats* stats) {
+  // Sampling and serviceability support
+  if (tlab().end() != nullptr) {
+    incr_allocated_bytes(tlab().used_bytes());
+    heap_sampler().retire_tlab(tlab().top());
+  }
+
+  // Retire the TLAB
+  tlab().retire(stats);
+}
+
+void Thread::fill_tlab(HeapWord* start, size_t pre_reserved, size_t new_size) {
+  // Thread allocation sampling support
+  heap_sampler().set_tlab_top_at_sample_start(start);
+
+  // Fill the TLAB
+  tlab().fill(start, start + pre_reserved, new_size);
+}
+
 void Thread::initialize_thread_current() {
   assert(_thr_current == nullptr, "Thread::current already initialized");
   _thr_current = this;