[Wisp] Support threadMxbean in AllThreadAsWisp

Summary: Support ThreadMxBean to inspect coroutine
infomation, including getThreadInfo, dumpAllThreads,
getThreadCount, getAllThreadIds.

Test Plan: com/alibaba/wisp2/WispThreadMXBeanTest.java

Reviewed-by: zhengxiaolinX, yuleil, sanhongli

Issue: dragonwell-project/dragonwell8#132

src/linux/classes/com/alibaba/wisp/engine/WispTask.java

@@ -603,7 +603,7 @@ void countExecutionTime(long beginTime) {
     }
 
     StackTraceElement[] getStackTrace() {
-        return this.ctx.getCoroutineStack();
+        return this.threadWrapper.getStackTrace();
     }
 
     private static final AtomicReferenceFieldUpdater<WispTask, WispCarrier> CARRIER_UPDATER;

src/share/classes/java/dyn/Coroutine.java

@@ -152,8 +152,4 @@ protected void run() {
     private static native void registerNatives();
 
     private static native void setWispTask(long coroutine, int id, Object task, Object engine);
-
-    public StackTraceElement[] getCoroutineStack() {
-        return CoroutineSupport.getCoroutineStack(this.nativeCoroutine);
-    }
 }

src/share/classes/java/lang/Thread.java

@@ -1670,26 +1670,12 @@ public void setContextClassLoader(ClassLoader cl) {
      * @since 1.5
      */
     public StackTraceElement[] getStackTrace() {
-        boolean slowPath;
-        if (WEA != null) {
-            WispTask task = this.wispTask;
-            if (task == null) {
-                // When we create a thread, coroutine will not be created immediately.
-                // So if we take the stack trace at once, it will NCE.
-                // See: NullStackTrace.java and 6571589, Thread return an array which size is 0.
-                return new StackTraceElement[0];
-            }
-            if (!WEA.runningAsCoroutine(this)) {
-                slowPath = this != Thread.currentThread();
-            } else {
-                boolean isCurrentTask = WEA.getCurrentTask() == task;
-                slowPath = !isCurrentTask;
-                if (!WEA.isThreadTask(task) && !isCurrentTask) {
-                    return WEA.getStackTrace(task);
-                }
-            }
-        } else {
-            slowPath = this != Thread.currentThread();
+        boolean slowPath = this != Thread.currentThread();
+        if (WEA != null && this.wispTask == null) {
+            // When we create a thread, coroutine will not be created immediately.
+            // So if we take the stack trace at once, it will NCE.
+            // See: NullStackTrace.java and 6571589, Thread return an array which size is 0.
+            return new StackTraceElement[0];
         }
 
         if (slowPath) {

src/share/classes/sun/management/ThreadImpl.java

@@ -25,6 +25,9 @@
 
 package sun.management;
 
+import sun.misc.SharedSecrets;
+import sun.misc.WispEngineAccess;
+
 import java.lang.management.ManagementFactory;
 
 import java.lang.management.ThreadInfo;
@@ -40,6 +43,8 @@
  */
 class ThreadImpl implements com.sun.management.ThreadMXBean {
 
+    private static WispEngineAccess WEA = SharedSecrets.getWispEngineAccess();
+
     private final VMManagement jvm;
 
     // default for thread contention monitoring is disabled.
@@ -57,7 +62,7 @@ class ThreadImpl implements com.sun.management.ThreadMXBean {
     }
 
     public int getThreadCount() {
-        return jvm.getLiveThreadCount();
+        return WEA != null && WEA.isAllThreadAsWisp()? getAllThreadIds().length : jvm.getLiveThreadCount();
     }
 
     public int getPeakThreadCount() {

test/com/alibaba/wisp2/WispThreadMXBeanTest.java

@@ -0,0 +1,122 @@
+/*
+ * @test
+ * @library /lib/testlibrary
+ * @summary
+ * @requires os.family == "linux"
+ * @run main/othervm  -XX:+UnlockExperimentalVMOptions -XX:+UseWisp2 WispThreadMXBeanTest
+ */
+// * @run main/othervm -XX:+UnlockExperimentalVMOptions -XX:+UseWisp2 WispThreadMXBeanTest
+
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.lang.management.ThreadMXBean;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Phaser;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+import java.util.stream.Stream;
+
+import static jdk.testlibrary.Asserts.assertTrue;
+
+
+public class WispThreadMXBeanTest {
+    private static final ThreadMXBean mbean = ManagementFactory.getThreadMXBean();
+    private static final int THREAD_SIZE = 25;
+    private static final Phaser startupCheck = new Phaser(THREAD_SIZE + 1);
+
+    public static void main(String[] args) throws Exception {
+        getAllThreadIds();
+        dumpAllThreads();
+        getThreadInfo();
+    }
+
+    // also getThreadCount
+    private static void getAllThreadIds() throws Exception {
+        long[] startIds = mbean.getAllThreadIds();
+        int startCnt = mbean.getThreadCount();
+        assertTrue(startIds.length == startCnt);
+
+        long[] newIds = new long[THREAD_SIZE];
+        for (int i = 0; i < THREAD_SIZE; i++) {
+            Thread t = new Thread(() -> {
+                startupCheck.arrive();
+                while (true) {
+                    sleep();
+                }
+            });
+            newIds[i] = t.getId();
+            t.start();
+        }
+        startupCheck.arriveAndAwaitAdvance();
+
+        long[] endIds = mbean.getAllThreadIds();
+        long endCnt = mbean.getThreadCount();
+
+        assertTrue(endIds.length == endCnt);
+        assertTrue(endCnt - startCnt == THREAD_SIZE);
+        for (int i = 0; i < THREAD_SIZE; i++) {
+            final int idx = i;
+            assertTrue(Arrays.stream(endIds).anyMatch(e -> e == newIds[idx]));
+        }
+    }
+
+    private static void dumpAllThreads() {
+        long[] ids = mbean.getAllThreadIds();
+        long[] newThreadIds = new long[THREAD_SIZE];
+
+        for (int i = 0; i < THREAD_SIZE; i++) {
+            Thread t = new Thread(() -> {
+                startupCheck.arrive();
+                while (true) {
+                    sleep();
+                }
+            });
+            t.setName("dumpAll");
+            newThreadIds[i] = t.getId();
+            t.start();
+        }
+        startupCheck.arriveAndAwaitAdvance();
+
+        boolean[] lockedMonitor = new boolean[] {false, true};
+        boolean[] lockedSyncronizer = new boolean[] {false, true};
+        for (boolean lm : lockedMonitor) {
+            for (boolean ls : lockedSyncronizer) {
+                ThreadInfo[] threadInfosNew = mbean.dumpAllThreads(lm, ls);
+                assertTrue(threadInfosNew.length - ids.length >= THREAD_SIZE, " " + threadInfosNew.length + "," + ids.length);
+                assertTrue(Arrays.stream(threadInfosNew).allMatch(info -> info.getStackTrace() != null));
+                assertTrue(Arrays.stream(threadInfosNew).anyMatch(info -> info.getThreadName().equals("dumpAll")));
+                assertTrue(Arrays.stream(threadInfosNew).noneMatch(Objects::isNull));
+                for (int i = 0; i < newThreadIds.length; i++) {
+                    final int idx = i;
+                    assertTrue(Arrays.stream(threadInfosNew).anyMatch(info -> info.getThreadId() == newThreadIds[idx] && Objects.equals(info.getStackTrace()[0].getFileName(), "CoroutineSupport.java")));
+                }
+            }
+        }
+    }
+
+    private static void getThreadInfo() throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        Thread thread = new Thread(() -> {
+            latch.countDown();
+            while (true) { sleep(); }
+        });
+        thread.start();
+
+        latch.await();
+        ThreadInfo[] infos = mbean.getThreadInfo(new long[]{thread.getId()}, Integer.MAX_VALUE);
+        StackTraceElement[] stack1 =  thread.getStackTrace();
+        StackTraceElement[] stack2 =  infos[0].getStackTrace();
+        assertTrue(Arrays.equals(stack1, stack2));
+    }
+
+
+
+    private static void sleep() {
+        try {
+            Thread.sleep(20);
+        } catch (InterruptedException e) {}
+    }
+}