--- a/src/hotspot/share/gc/shared/memAllocator.cpp	Tue Apr 30 12:53:32 2019 -0700
+++ b/src/hotspot/share/gc/shared/memAllocator.cpp	Tue Apr 30 15:39:42 2019 -0700
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -198,20 +198,27 @@
     return;
   }
 
-  if (JvmtiExport::should_post_sampled_object_alloc()) {
-    // 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.
+  // 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;
+
+  {
     PreserveObj obj_h(_thread, _obj_ptr);
     JvmtiSampledObjectAllocEventCollector collector;
     size_t size_in_bytes = _allocator._word_size * HeapWordSize;
     ThreadLocalAllocBuffer& tlab = _thread->tlab();
-    size_t bytes_since_last = _allocated_outside_tlab ? 0 : tlab.bytes_since_last_sample_point();
+
+    if (!_allocated_outside_tlab) {
+      bytes_since_last = tlab.bytes_since_last_sample_point();
+    }
+
     _thread->heap_sampler().check_for_sampling(obj_h(), size_in_bytes, bytes_since_last);
   }
 
   if (_tlab_end_reset_for_sample || _allocated_tlab_size != 0) {
-    _thread->tlab().set_sample_end();
+    // Tell tlab to forget bytes_since_last if we passed it to the heap sampler.
+    _thread->tlab().set_sample_end(bytes_since_last != 0);
   }
 }
 
@@ -283,12 +290,14 @@
   ThreadLocalAllocBuffer& tlab = _thread->tlab();
 
   if (JvmtiExport::should_post_sampled_object_alloc()) {
-    // Try to allocate the sampled object from TLAB, it is possible a sample
-    // point was put and the TLAB still has space.
     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;
+
     if (mem != NULL) {
-      allocation._tlab_end_reset_for_sample = true;
       return mem;
     }
   }

--- a/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp	Tue Apr 30 12:53:32 2019 -0700
+++ b/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp	Tue Apr 30 15:39:42 2019 -0700
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 1999, 2018, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 1999, 2019, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -281,17 +281,21 @@
             _fast_refill_waste * HeapWordSize);
 }
 
-void ThreadLocalAllocBuffer::set_sample_end() {
+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;
+    _bytes_since_last_sample_point += bytes_until_sample;
   } else {
-    _bytes_since_last_sample_point = heap_words_remaining * HeapWordSize;
+    _bytes_since_last_sample_point += heap_words_remaining * HeapWordSize;
   }
 }
 

--- a/src/hotspot/share/gc/shared/threadLocalAllocBuffer.hpp	Tue Apr 30 12:53:32 2019 -0700
+++ b/src/hotspot/share/gc/shared/threadLocalAllocBuffer.hpp	Tue Apr 30 15:39:42 2019 -0700
@@ -171,7 +171,7 @@
   void initialize();
 
   void set_back_allocation_end();
-  void set_sample_end();
+  void set_sample_end(bool reset_byte_accumulation);
 
   static size_t refill_waste_limit_increment()   { return TLABWasteIncrement; }
 

--- a/src/hotspot/share/runtime/threadHeapSampler.cpp	Tue Apr 30 12:53:32 2019 -0700
+++ b/src/hotspot/share/runtime/threadHeapSampler.cpp	Tue Apr 30 15:39:42 2019 -0700
@@ -121,11 +121,6 @@
   }
 
   pick_next_geometric_sample();
-
-  // Try to correct sample size by removing extra space from last allocation.
-  if (overflowed_bytes > 0 && _bytes_until_sample > overflowed_bytes) {
-    _bytes_until_sample -= overflowed_bytes;
-  }
 }
 
 void ThreadHeapSampler::check_for_sampling(oop obj, size_t allocation_size, size_t bytes_since_allocation) {

--- a/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitor.java	Tue Apr 30 12:53:32 2019 -0700
+++ b/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitor.java	Tue Apr 30 15:39:42 2019 -0700
@@ -1,6 +1,6 @@
 /*
- * Copyright (c) 2018, Oracle and/or its affiliates. All rights reserved.
- * Copyright (c) 2018, Google and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Google and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -263,8 +263,13 @@
   }
 
   public static boolean statsHaveExpectedNumberSamples(int expected, int acceptedErrorPercentage) {
-    double actual = getEventStorageElementCount();
-    double diffPercentage = Math.abs(actual - expected) / expected;
+    double actual = sampledEvents();
+    double diffPercentage = 100 * Math.abs(actual - expected) / expected;
+
+    if (diffPercentage >= acceptedErrorPercentage) {
+      System.err.println("Unexpected high difference percentage: " + diffPercentage
+          + " due to the count being " + actual + " instead of " + expected);
+    }
     return diffPercentage < acceptedErrorPercentage;
   }
 

--- a/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitorArrayAllSampledTest.java	Tue Apr 30 12:53:32 2019 -0700
+++ b/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitorArrayAllSampledTest.java	Tue Apr 30 15:39:42 2019 -0700
@@ -1,6 +1,6 @@
 /*
- * Copyright (c) 2018, Oracle and/or its affiliates. All rights reserved.
- * Copyright (c) 2018, Google and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Google and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -47,8 +47,7 @@
   public static void main(String[] args) {
     int sizes[] = {1000, 10000, 100000, 1000000};
 
-    HeapMonitor.setSamplingInterval(0);
-    HeapMonitor.enableSamplingEvents();
+    HeapMonitor.sampleEverything();
 
     for (int currentSize : sizes) {
       System.out.println("Testing size " + currentSize);

--- a/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitorStatArrayCorrectnessTest.java	Tue Apr 30 12:53:32 2019 -0700
+++ b/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/MyPackage/HeapMonitorStatArrayCorrectnessTest.java	Tue Apr 30 15:39:42 2019 -0700
@@ -1,6 +1,6 @@
 /*
- * Copyright (c) 2018, Oracle and/or its affiliates. All rights reserved.
- * Copyright (c) 2018, Google and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2019, Google and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -55,12 +55,12 @@
         throw new RuntimeException("Should not have any events stored yet.");
       }
 
-      HeapMonitor.enableSamplingEvents();
-
       // 111 is as good a number as any.
       final int samplingMultiplier = 111;
       HeapMonitor.setSamplingInterval(samplingMultiplier * currentSize);
 
+      HeapMonitor.enableSamplingEvents();
+
       allocate(currentSize);
 
       HeapMonitor.disableSamplingEvents();
@@ -84,7 +84,9 @@
       // statistical geometric variable around the sampling interval. This means that the test could be
       // unlucky and not achieve the mean average fast enough for the test case.
       if (!HeapMonitor.statsHaveExpectedNumberSamples((int) expected, 10)) {
-        throw new RuntimeException("Statistics should show about " + expected + " samples.");
+        throw new RuntimeException("Statistics should show about " + expected + " samples; "
+            + " but have " + HeapMonitor.sampledEvents() + " instead for the size "
+            + currentSize);
       }
     }
   }