fix recursive trace

sunwu51 · Apr 22, 2024 · 07578e4 · 07578e4
1 parent 03a7792
commit 07578e4
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 19 deletions.
diff --git a/src/main/java/w/core/model/TraceTransformer.java b/src/main/java/w/core/model/TraceTransformer.java
@@ -17,6 +17,7 @@
 @Data
 public class TraceTransformer extends BaseClassTransformer {
     public static ThreadLocal<Map<String, int[]>> traceContent = ThreadLocal.withInitial(LinkedHashMap::new);
+    public static ThreadLocal<Integer> stackDeep = ThreadLocal.withInitial(()->1);
 
     transient TraceMessage message;
 
@@ -61,6 +62,10 @@ public void edit(MethodCall m) throws CannotCompileException {
                     return;
                 }
                 String code = "{" +
+                        "   if (\"" + m.getMethodName() + "\".equals(\"" + method + "\")) {" +
+                        "       int deep = ((Integer)w.core.model.TraceTransformer.stackDeep.get()).intValue();" +
+                        "       w.core.model.TraceTransformer.stackDeep.set(new Integer(++deep));" +
+                        "   }" +
                         "   long start = System.currentTimeMillis();\n" +
                         "   $_ = $proceed($$);\n" +
                         "   long duration = System.currentTimeMillis() - start;\n" +
@@ -69,33 +74,37 @@ public void edit(MethodCall m) throws CannotCompileException {
                         "   if (!map.containsKey(sig)) {map.put(sig, new int[]{0,0});}" +
                         "   int[] arr = (int[])map.get(sig); \n" +
                         "   arr[0] += duration; arr[1] += 1;" +
+
                         "}";
                 m.replace(code);
             }
         });
-        String str ="";
-
         ctMethod.addLocalVariable("s", CtClass.longType);
         ctMethod.addLocalVariable("cost", CtClass.longType);
         ctMethod.insertBefore("s = System.currentTimeMillis();");
         ctMethod.insertAfter("{" +
-                "cost = System.currentTimeMillis() - s;\n" +
-                "if (cost >= " + minCost +") {" +
-                "  w.Global.checkCountAndUnload(\"" + uuid + "\");\n"+
-                "  w.util.RequestUtils.fillCurThread(\"" + message.getId() + "\");\n" +
-                "  String str = \"" + className + "#" + method + ", total cost:\"+cost+\"ms\\\n\";\n" +
-                "  LinkedHashMap map = (LinkedHashMap)w.core.model.TraceTransformer.traceContent.get(); \n" +
-                "  Iterator it = map.entrySet().iterator(); \n" +
-                "  while (it.hasNext()) {\n" +
-                "      java.util.Map.Entry e = (java.util.Map.Entry)it.next();\n" +
-                "      String k = e.getKey().toString();\n" +
-                "      int[] v = (int[])e.getValue();\n" +
-                "      if (v[0] == 0 && " + ignoreZero + ") \n {} else {" +
-                "       str += \">>\" + k + \" hit:\" + v[1] + \"times, total cost:\" + v[0] + \"ms\\\n\";}\n" +
-                "  }" +
-                "  w.core.model.TraceTransformer.traceContent.remove();\n" +
-                "  w.Global.info(str);\n" +
-                "  w.util.RequestUtils.clearRequestCtx();\n" +
+                "int deep = ((Integer)w.core.model.TraceTransformer.stackDeep.get()).intValue();\n" +
+                "w.core.model.TraceTransformer.stackDeep.set(new Integer(--deep));" +
+                "if (deep <= 0) {" +
+                "   w.core.model.TraceTransformer.stackDeep.remove();\n" +
+                "   cost = System.currentTimeMillis() - s;\n" +
+                "   if (cost >= " + minCost +") {" +
+                "     w.Global.checkCountAndUnload(\"" + uuid + "\");\n"+
+                "     w.util.RequestUtils.fillCurThread(\"" + message.getId() + "\");\n" +
+                "     String str = \"" + className + "#" + method + ", total cost:\"+cost+\"ms\\\n\";\n" +
+                "     LinkedHashMap map = (LinkedHashMap)w.core.model.TraceTransformer.traceContent.get(); \n" +
+                "     Iterator it = map.entrySet().iterator(); \n" +
+                "     while (it.hasNext()) {\n" +
+                "         java.util.Map.Entry e = (java.util.Map.Entry)it.next();\n" +
+                "         String k = e.getKey().toString();\n" +
+                "         int[] v = (int[])e.getValue();\n" +
+                "         if (v[0] == 0 && " + ignoreZero + ") \n {} else {" +
+                "          str += \">>\" + k + \" hit:\" + v[1] + \"times, total cost:\" + v[0] + \"ms\\\n\";}\n" +
+                "     }" +
+                "     w.core.model.TraceTransformer.traceContent.remove();\n" +
+                "     w.Global.info(str);\n" +
+                "     w.util.RequestUtils.clearRequestCtx();\n" +
+                "   }" +
                 "}" +
                 "}");
     }

diff --git a/src/test/java/w/core/SwapperTest.java b/src/test/java/w/core/SwapperTest.java
@@ -111,4 +111,13 @@ public void replaceClassTest() throws Exception {
         Assertions.assertTrue(swapper.swap(message));
         Assertions.assertEquals("hi frank", t.hello("frank"));
     }
+
+    @Test
+    public void traceRecursiveTest() {
+        TraceMessage message = new TraceMessage();
+        message.setSignature("w.core.TestClass#recursive");
+        message.setIgnoreZero(false);
+        Assertions.assertTrue(swapper.swap(message));
+        t.recursive(3);
+    }
 }
diff --git a/src/test/java/w/core/TestClass.java b/src/test/java/w/core/TestClass.java
@@ -1,5 +1,7 @@
 package w.core;
 
+import java.util.UUID;
+
 /**
  * @author Frank
  * @date 2024/4/21 11:02
@@ -14,4 +16,10 @@ public String wrapperHello(String name) {
         return hello(name);
     }
 
+
+    public int recursive(int n) {
+        if (n <= 1) return 1;
+        UUID.randomUUID();
+        return recursive(n - 1) + recursive(n - 2);
+    }
 }