vmtool support limit option. #1781

arthas-vmtool/src/main/java/arthas/VmTool.java

@@ -46,12 +46,12 @@ public class VmTool implements VmToolMXBean {
     /**
      * 获取某个class在jvm中当前所有存活实例
      */
-    private static synchronized native <T> T[] getInstances0(Class<T> klass);
+    private static synchronized native <T> T[] getInstances0(Class<T> klass, int limit);
 
     /**
      * 统计某个class在jvm中当前所有存活实例的总占用内存，单位：Byte
      */
-    private static native long sumInstanceSize0(Class<?> klass);
+    private static synchronized native long sumInstanceSize0(Class<?> klass);
 
     /**
      * 获取某个实例的占用内存，单位：Byte
@@ -77,7 +77,15 @@ public class VmTool implements VmToolMXBean {
 
     @Override
     public <T> T[] getInstances(Class<T> klass) {
-        return getInstances0(klass);
+        return getInstances0(klass, -1);
+    }
+
+    @Override
+    public <T> T[] getInstances(Class<T> klass, int limit) {
+        if (limit == 0) {
+            throw new IllegalArgumentException("limit can not be 0");
+        }
+        return getInstances0(klass, limit);
     }
 
     @Override

arthas-vmtool/src/main/java/arthas/VmToolMXBean.java

@@ -20,10 +20,16 @@ public interface VmToolMXBean {
      */
     public void forceGc();
 
+    public <T> T[] getInstances(Class<T> klass);
+
     /**
      * 获取某个class在jvm中当前所有存活实例
+     * @param <T>
+     * @param klass
+     * @param limit 如果小于 0 ，则不限制
+     * @return
      */
-    public <T> T[] getInstances(Class<T> klass);
+    public <T> T[] getInstances(Class<T> klass, int limit);
 
     /**
      * 统计某个class在jvm中当前所有存活实例的总占用内存，单位：Byte

arthas-vmtool/src/main/native/src/jni-library.cpp

@@ -8,6 +8,30 @@
 static jvmtiEnv *jvmti;
 static jlong tagCounter = 0;
 
+struct LimitCounter {
+    jint currentCounter;
+    jint limitValue;
+
+    void init(jint limit) {
+        currentCounter = 0;
+        limitValue = limit;
+    }
+
+    void countDown() {
+        currentCounter++;
+    }
+
+    bool allow() {
+        if (limitValue < 0) {
+            return true;
+        }
+        return limitValue > currentCounter;
+    }
+};
+
+// 每次 IterateOverInstancesOfClass 调用前需要先 init
+static LimitCounter limitCounter = {0, 0};
+
 extern "C"
 int init_agent(JavaVM *vm, void *reserved) {
     jint rc;
@@ -61,13 +85,20 @@ jvmtiIterationControl JNICALL
 HeapObjectCallback(jlong class_tag, jlong size, jlong *tag_ptr, void *user_data) {
     jlong *data = static_cast<jlong *>(user_data);
     *tag_ptr = *data;
-    return JVMTI_ITERATION_CONTINUE;
+
+    limitCounter.countDown();
+    if (limitCounter.allow()) {
+        return JVMTI_ITERATION_CONTINUE;
+    }else {
+        return JVMTI_ITERATION_ABORT;
+    }
 }
 
 extern "C"
 JNIEXPORT jobjectArray JNICALL
-Java_arthas_VmTool_getInstances0(JNIEnv *env, jclass thisClass, jclass klass) {
+Java_arthas_VmTool_getInstances0(JNIEnv *env, jclass thisClass, jclass klass, jint limit) {
     jlong tag = getTag();
+    limitCounter.init(limit);
     jvmtiError error = jvmti->IterateOverInstancesOfClass(klass, JVMTI_HEAP_OBJECT_EITHER,
                                                HeapObjectCallback, &tag);
     if (error) {
@@ -96,6 +127,7 @@ extern "C"
 JNIEXPORT jlong JNICALL
 Java_arthas_VmTool_sumInstanceSize0(JNIEnv *env, jclass thisClass, jclass klass) {
     jlong tag = getTag();
+    limitCounter.init(-1);
     jvmtiError error = jvmti->IterateOverInstancesOfClass(klass, JVMTI_HEAP_OBJECT_EITHER,
                                                HeapObjectCallback, &tag);
     if (error) {
@@ -136,6 +168,7 @@ extern "C"
 JNIEXPORT jlong JNICALL
 Java_arthas_VmTool_countInstances0(JNIEnv *env, jclass thisClass, jclass klass) {
     jlong tag = getTag();
+    limitCounter.init(-1);
     jvmtiError error = jvmti->IterateOverInstancesOfClass(klass, JVMTI_HEAP_OBJECT_EITHER,
                                                HeapObjectCallback, &tag);
     if (error) {

arthas-vmtool/src/test/java/arthas/VmToolTest.java

@@ -2,6 +2,7 @@ package arthas;
 
 import java.io.File;
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.assertj.core.api.Assertions;
@@ -9,6 +10,8 @@ import org.junit.Test;
 
 import com.taobao.arthas.common.VmToolUtils;
 
+import arthas.VmToolTest.LimitTest;
+
 /**
  * 以下本地测试的jvm参数均为：-Xms128m -Xmx128m
  */
@@ -146,6 +149,25 @@ public class VmToolTest {
         }
     }
 
+    class LimitTest {
+    }
+
+    @Test
+    public void test_getInstances_lmiit() {
+        VmTool vmtool = initVmTool();
+
+        ArrayList<LimitTest> list = new ArrayList<LimitTest>();
+        for (int i = 0; i < 10; ++i) {
+            list.add(new LimitTest());
+        }
+        LimitTest[] instances = vmtool.getInstances(LimitTest.class, 5);
+        Assertions.assertThat(instances).hasSize(5);
+        LimitTest[] instances2 = vmtool.getInstances(LimitTest.class, -1);
+        Assertions.assertThat(instances2).hasSize(10);
+        LimitTest[] instances3 = vmtool.getInstances(LimitTest.class, 1);
+        Assertions.assertThat(instances3).hasSize(1);
+    }
+
     interface III {
     }
 

core/src/main/java/com/taobao/arthas/core/command/monitor200/VmToolCommand.java

@@ -47,6 +47,8 @@ import arthas.VmTool;
 @Description(Constants.EXAMPLE
         + "  vmtool --action getInstances --className demo.MathGame\n"
         + "  vmtool --action getInstances --className demo.MathGame --express 'instances.length'\n"
+        + "  vmtool --action getInstances --className demo.MathGame --express 'instances[0]'\n"
+        + "  vmtool --action getInstances --className java.lang.String --limit 10\n"
         + "  vmtool --action forceGc\n"
         + Constants.WIKI + Constants.WIKI_HOME + "vmtool")
 //@formatter:on
@@ -64,6 +66,11 @@ public class VmToolCommand extends AnnotatedCommand {
      */
     private int expand;
 
+    /**
+     * default value 10
+     */
+    private int limit;
+
     private static String libPath;
     private static VmTool vmTool = null;
 
@@ -117,6 +124,13 @@ public class VmToolCommand extends AnnotatedCommand {
         this.classLoaderClass = classLoaderClass;
     }
 
+    @Option(shortName = "l", longName = "limit")
+    @Description("Set the limit value of the getInstances action, default value is 10, set to -1 is unlimited")
+    @DefaultValue("10")
+    public void setLimit(int limit) {
+        this.limit = limit;
+    }
+
     @Option(longName = "express", required = false)
     @Description("The ognl expression, default valueis `instances`.")
     public void setExpress(String express) {
@@ -165,7 +179,7 @@ public class VmToolCommand extends AnnotatedCommand {
                     process.end(-1, "Found more than one class: " + matchedClasses + ".");
                     return;
                 } else {
-                    Object[] instances = vmToolInstance().getInstances(matchedClasses.get(0));
+                    Object[] instances = vmToolInstance().getInstances(matchedClasses.get(0), limit);
                     Object value = instances;
                     if (express != null) {
                         Express unpooledExpress = ExpressFactory.unpooledExpress(classLoader);

spy/src/main/java/arthas/VmTool.java

@@ -41,17 +41,17 @@ public class VmTool implements VmToolMXBean {
      */
     private static native String check0();
 
-    private static native void forceGc0();
+    private static synchronized native void forceGc0();
 
     /**
      * 获取某个class在jvm中当前所有存活实例
      */
-    private static native <T> T[] getInstances0(Class<T> klass);
+    private static synchronized native <T> T[] getInstances0(Class<T> klass, int limit);
 
     /**
      * 统计某个class在jvm中当前所有存活实例的总占用内存，单位：Byte
      */
-    private static native long sumInstanceSize0(Class<?> klass);
+    private static synchronized native long sumInstanceSize0(Class<?> klass);
 
     /**
      * 获取某个实例的占用内存，单位：Byte
@@ -61,24 +61,14 @@ public class VmTool implements VmToolMXBean {
     /**
      * 统计某个class在jvm中当前所有存活实例的总个数
      */
-    private static native long countInstances0(Class<?> klass);
+    private static synchronized native long countInstances0(Class<?> klass);
 
     /**
      * 获取所有已加载的类
+     * @param klass 这个参数必须是 Class.class
+     * @return
      */
-    private static native Class<?>[] getAllLoadedClasses0();
-
-    /**
-     * 包括小类型(如int)
-     */
-    public static Class<?>[] getAllClasses() {
-        return getInstances0(Class.class);
-    }
-
-    @Override
-    public String check() {
-        return check0();
-    }
+    private static synchronized native Class<?>[] getAllLoadedClasses0(Class<?> klass);
 
     @Override
     public void forceGc() {
@@ -87,7 +77,15 @@ public class VmTool implements VmToolMXBean {
 
     @Override
     public <T> T[] getInstances(Class<T> klass) {
-        return getInstances0(klass);
+        return getInstances0(klass, -1);
+    }
+
+    @Override
+    public <T> T[] getInstances(Class<T> klass, int limit) {
+        if (limit == 0) {
+            throw new IllegalArgumentException("limit can not be 0");
+        }
+        return getInstances0(klass, limit);
     }
 
     @Override
@@ -107,7 +105,7 @@ public class VmTool implements VmToolMXBean {
 
     @Override
     public Class<?>[] getAllLoadedClasses() {
-        return getAllLoadedClasses0();
+        return getAllLoadedClasses0(Class.class);
     }
 
 }

spy/src/main/java/arthas/VmToolMXBean.java

@@ -14,20 +14,22 @@ package arthas;
  * @author hengyunabc 2021-04-26
  */
 public interface VmToolMXBean {
-    /**
-     * 检测jni-lib是否正常，如果正常，应该输出OK
-     */
-    public String check();
 
     /**
      * https://docs.oracle.com/javase/8/docs/platform/jvmti/jvmti.html#ForceGarbageCollection
      */
     public void forceGc();
 
+    public <T> T[] getInstances(Class<T> klass);
+
     /**
      * 获取某个class在jvm中当前所有存活实例
+     * @param <T>
+     * @param klass
+     * @param limit 如果小于 0 ，则不限制
+     * @return
      */
-    public <T> T[] getInstances(Class<T> klass);
+    public <T> T[] getInstances(Class<T> klass, int limit);
 
     /**
      * 统计某个class在jvm中当前所有存活实例的总占用内存，单位：Byte