--- old/src/hotspot/share/runtime/threadHeapSampler.cpp	2018-03-21 10:34:11.139179633 -0700
+++ new/src/hotspot/share/runtime/threadHeapSampler.cpp	2018-03-21 10:34:10.847180634 -0700
@@ -27,7 +27,8 @@
 
 // Cheap random number generator
 uint64_t ThreadHeapSampler::_rnd;
-int ThreadHeapSampler::_sampling_rate;
+// Default is 512kb.
+int ThreadHeapSampler::_sampling_rate = 512 * 1024;
 int ThreadHeapSampler::_enabled;
 
 // Statics for the fast log
@@ -68,7 +69,7 @@
 // -log_e(q)/m = x
 // log_2(q) * (-log_e(2) * 1/m) = x
 // In the code, q is actually in the range 1 to 2**26, hence the -26 below
-void ThreadHeapSampler::pick_next_sample(size_t overflowed_bytes) {
+void ThreadHeapSampler::pick_next_geometric_sample() {
   _rnd = next_random(_rnd);
   // Take the top 26 bits as the random number
   // (This plus a 1<<58 sampling bound gives a max possible step of
@@ -89,6 +90,15 @@
   size_t rate = static_cast<size_t>(
       (0.0 < log_val ? 0.0 : log_val) * (-log(2.0) * (_sampling_rate)) + 1);
   _bytes_until_sample = rate;
+}
+
+void ThreadHeapSampler::pick_next_sample(size_t overflowed_bytes) {
+  if (_sampling_rate == 1) {
+    _bytes_until_sample = 1;
+    return;
+  }
+
+  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) {
@@ -128,11 +138,6 @@
 }
 
 void ThreadHeapSampler::set_tlab_heap_sampling(int sampling_rate) {
-  if (sampling_rate == 0) {
-    disable();
-    _sampling_rate = sampling_rate;
-  } else {
-    _sampling_rate = sampling_rate;
-    enable();
-  }
+  MutexLocker mu(ThreadHeapSampler_lock);
+  _sampling_rate = sampling_rate;
 }