--- old/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/libHeapMonitor.c	2017-10-06 14:50:24.838594041 -0700
+++ new/test/hotspot/jtreg/serviceability/jvmti/HeapMonitor/libHeapMonitor.c	2017-10-06 14:50:24.466595293 -0700
@@ -22,6 +22,7 @@
  */
 
 #include <stdio.h>
+#include <stdlib.h>
 #include <string.h>
 #include "jvmti.h"
 
@@ -141,7 +142,7 @@
   memset(&caps, 0, sizeof(caps));
   // Get line numbers, sample heap, and filename for the test.
   caps.can_get_line_numbers = 1;
-  caps.can_sample_heap= 1;
+  caps.can_sample_heap = 1;
   caps.can_get_source_file_name = 1;
   if (check_error((*jvmti)->AddCapabilities(jvmti, &caps),
                   "Add capabilities\n")){
@@ -474,11 +475,15 @@
               "Start Heap Sampling");
 }
 
+static void enable_sampling_with_max_traces(int max_traces) {
+  check_error((*jvmti)->StartHeapSampling(jvmti, 1 << 19, max_traces),
+              "Start Heap Sampling");
+}
+
 static void disable_sampling() {
   check_error((*jvmti)->StopHeapSampling(jvmti), "Stop Heap Sampling");
 }
 
-// HeapMonitorTest JNI.
 JNIEXPORT jint JNICALL
 Java_MyPackage_HeapMonitorTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
   // We want the frames in each part.
@@ -493,7 +498,6 @@
   enable_sampling();
 }
 
-// HeapMonitorOnOffTest JNI.
 JNIEXPORT jint JNICALL
 Java_MyPackage_HeapMonitorOnOffTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
   // We want the frames in each part.
@@ -522,7 +526,6 @@
   disable_sampling();
 }
 
-// HeapMonitorRecentTest JNI.
 JNIEXPORT jint JNICALL
 Java_MyPackage_HeapMonitorRecentTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
   // We want the frames in each part.
@@ -553,7 +556,6 @@
   enable_sampling();
 }
 
-// HeapMonitorFrequentTest JNI.
 JNIEXPORT jint JNICALL
 Java_MyPackage_HeapMonitorFrequentTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
   // We want the frames in each part.
@@ -692,11 +694,139 @@
   jvmtiHeapSamplingStats stats;
   check_error((*jvmti)->GetHeapSamplingStats(jvmti, &stats),
               "Heap Sampling Statistics");
-    (int) stats.sample_rate_accumulation, (int) stats.sample_rate_count);
-
   return ((double) stats.sample_rate_accumulation) / stats.sample_rate_count;
 }
 
+JNIEXPORT void JNICALL
+Java_MyPackage_HeapMonitorStackDepthTest_enableSampling(JNIEnv *env, jclass cls) {
+  enable_sampling();
+}
+
+JNIEXPORT void JNICALL
+Java_MyPackage_HeapMonitorStackDepthTest_disableSampling(JNIEnv *env, jclass cls) {
+  disable_sampling();
+}
+
+JNIEXPORT jdouble JNICALL
+Java_MyPackage_HeapMonitorStackDepthTest_getAverageStackDepth(JNIEnv *env, jclass cls) {
+  jvmtiStackTraces traces;
+  jvmtiError error = (*jvmti)->GetLiveTraces(jvmti, &traces);;
+
+  if (error != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+
+  int trace_count = traces.trace_count;
+
+  if (trace_count == 0) {
+    return 0;
+  }
+
+  int i;
+  jvmtiStackTrace* stack_traces = traces.stack_traces;
+  double sum = 0;
+  for (i = 0; i < trace_count; i++) {
+    jvmtiStackTrace *stack_trace = stack_traces + i;
+    sum += stack_trace->frame_count;
+  }
+
+  return sum / i;
+}
+
+JNIEXPORT void JNICALL
+Java_MyPackage_HeapMonitorThreadTest_enableSampling(JNIEnv *env, jclass cls) {
+  // Remember a lot of the garbage collected samples to ensure test correctness.
+  enable_sampling_with_max_traces(100000);
+}
+
+typedef struct sThreadsFound {
+  jint *threads;
+  int num_threads;
+}ThreadsFound;
+
+static void find_threads_in_traces(jvmtiStackTraces* traces, ThreadsFound* thread_data) {
+  int i;
+  jvmtiStackTrace* stack_traces = traces->stack_traces;
+  int trace_count = traces->trace_count;
+
+  jint *threads = thread_data->threads;
+  int num_threads = thread_data->num_threads;
+
+  // We are looking for at last expected_num_threads different traces.
+  for (i = 0; i < trace_count; i++) {
+    jvmtiStackTrace *stack_trace = stack_traces + i;
+    jlong thread_id = stack_trace->thread_id;
+
+    // Check it is the right frame: only accept helper top framed traces.
+    jmethodID methodid = stack_trace->frames[0].method;
+    char *name = NULL, *signature = NULL, *file_name = NULL;
+    (*jvmti)->GetMethodName(jvmti, methodid, &name, &signature, 0);
+
+    if (strcmp(name, "helper")) {
+      continue;
+    }
+
+    // Really not efficient look-up but it's for a test...
+    int found = 0;
+    int j;
+    for (j = 0; j < num_threads; j++) {
+      if (thread_id == threads[j]) {
+        found = 1;
+        break;
+      }
+    }
+
+    if (!found) {
+      threads[num_threads] = thread_id;
+      num_threads++;
+    }
+  }
+  thread_data->num_threads = num_threads;
+}
+
+JNIEXPORT jboolean JNICALL
+Java_MyPackage_HeapMonitorThreadTest_checkSamples(JNIEnv* env, jclass cls, jintArray threads) {
+  jvmtiStackTraces traces;
+  ThreadsFound thread_data;
+  thread_data.threads = (*env)->GetIntArrayElements(env, threads, 0);
+  thread_data.num_threads = 0;
+
+  // Get live and garbage traces to ensure we capture all the threads that have
+  // been sampled.
+  if ((*jvmti)->GetLiveTraces(jvmti, &traces) != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+
+  find_threads_in_traces(&traces, &thread_data);
+
+  if ((*jvmti)->ReleaseTraces(jvmti, &traces) != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+
+  if ((*jvmti)->GetGarbageTraces(jvmti, &traces) != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+
+  find_threads_in_traces(&traces, &thread_data);
+
+  if ((*jvmti)->ReleaseTraces(jvmti, &traces) != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+
+  (*env)->ReleaseIntArrayElements(env, threads, thread_data.threads, 0);
+  return 1;
+}
+
+JNIEXPORT void JNICALL
+Java_MyPackage_Switch_enableSampling(JNIEnv *env, jclass cls) {
+  enable_sampling();
+}
+
+JNIEXPORT void JNICALL
+Java_MyPackage_Switch_disableSampling(JNIEnv *env, jclass cls) {
+  disable_sampling();
+}
+
 #ifdef __cplusplus
 }
 #endif