Fix @Tag returnedObject bug #5036 (#5045)

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/util/CustomizeExpression.java

@@ -49,13 +49,29 @@ public class CustomizeExpression {
 
     public static Map<String, Object> evaluationReturnContext(Object ret)  {
         Map<String, Object> context = new HashMap<>();
-        Field[] fields = ret.getClass().getDeclaredFields();
-        for (Field field : fields) {
-            field.setAccessible(true);
-            try {
-                context.put(field.getName(), field.get(ret));
-            } catch (Exception e) {
-                logger.debug("evaluationReturnContext error, ret is {}, exception is {}", ret, e.getMessage());
+        context.put("returnedObj", ret.toString());
+        if (ret instanceof List) {
+            List retList = (List) ret;
+            int retLength = retList.size();
+            for (int i = 0; i < retLength; i++) {
+                context.put(String.valueOf(i), retList.get(i));
+            }
+        } else if (ret.getClass().isArray()) {
+            int length = Array.getLength(ret);
+            for (int i = 0; i < length; i++) {
+                context.put(String.valueOf(i), Array.get(ret, i));
+            }
+        } else if (ret instanceof Map) {
+            context.putAll((Map) ret);
+        } else {
+            Field[] fields = ret.getClass().getDeclaredFields();
+            for (Field field : fields) {
+                field.setAccessible(true);
+                try {
+                    context.put(field.getName(), field.get(ret));
+                } catch (Exception e) {
+                    logger.debug("evaluationReturnContext error, ret is {}, exception is {}", ret, e.getMessage());
+                }
             }
         }
         return context;
@@ -75,6 +91,9 @@ public class CustomizeExpression {
     public static String parseReturnExpression(String expression, Map<String, Object> context) {
         try {
             String[] es = expression.split("\\.");
+            if (es.length == 1) {
+                return String.valueOf(context.get(es[0]));
+            }
             Object o = context.get(es[1]);
             return o == null ? "null" : String.valueOf(parse(es, o, 1));
         } catch (Exception e) {

apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/util/TagUtil.java

@@ -38,6 +38,6 @@ public class TagUtil {
 
     public static Boolean isReturnTag(String expression) {
         String[] es = expression.split("\\.");
-        return es.length == 2 && "returnedObj".equals(es[0]);
+        return "returnedObj".equals(es[0]);
     }
 }

apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/test/java/org/apache/skywalking/apm/toolkit/activation/trace/TagAnnotationTest.java

@@ -19,7 +19,11 @@
 package org.apache.skywalking.apm.toolkit.activation.trace;
 
 import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
 import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
@@ -41,6 +45,8 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.modules.junit4.PowerMockRunnerDelegate;
+
+import static org.hamcrest.CoreMatchers.anything;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 
@@ -138,6 +144,142 @@ public class TagAnnotationTest {
 
     }
 
+    @Test
+    public void testTraceWithReturnList() throws Throwable {
+        Method testMethodWithReturnList = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnList", String.class, Integer.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnList, new Object[]{"wangwu", 18}, null, null);
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnList, null, null, Arrays.asList(new User("wangwu", 18)));
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 2);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("username"));
+        assertThat(tags.get(0).getValue(), is("wangwu"));
+        assertThat(tags.get(1).getKey().key(), is("info"));
+        assertThat(tags.get(1).getValue(), is("username=wangwu,age=18"));
+
+    }
+
+    @Test
+    public void testTraceWithReturnArray() throws Throwable {
+        Method testMethodWithReturnArray = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnArray", String.class, Integer.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnArray, new Object[]{"wangwu", 18}, null, null);
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnArray, null, null, new User[]{new User("wangwu", 18)});
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 2);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("username"));
+        assertThat(tags.get(0).getValue(), is("wangwu"));
+        assertThat(tags.get(1).getKey().key(), is("info"));
+        assertThat(tags.get(1).getValue(), is("username=wangwu,age=18"));
+    }
+
+    @Test
+    public void testTraceWithReturnMap() throws Throwable {
+        Method testMethodWithReturnMap = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnMap", String.class, Integer.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnMap, new Object[]{"wangwu", 18}, null, null);
+
+        Map<String, User> userMap = new HashMap<>();
+        userMap.put("user", new User("wangwu", 18));
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnMap, null, null, userMap);
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 2);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("username"));
+        assertThat(tags.get(0).getValue(), is("wangwu"));
+        assertThat(tags.get(1).getKey().key(), is("info"));
+        assertThat(tags.get(1).getValue(), is("username=wangwu,age=18"));
+    }
+
+    @Test
+    public void testTraceWithReturnString() throws Throwable {
+        Method testMethodWithReturnString = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnString", String.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnString, new Object[]{"wangwu"}, null, null);
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnString, null, null, "wangwu");
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 1);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("result"));
+        assertThat(tags.get(0).getValue(), is("wangwu"));
+    }
+
+    @Test
+    public void testTraceWithReturnInteger() throws Throwable {
+        Method testMethodWithReturnInteger = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnInteger", Integer.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnInteger, new Object[]{18}, null, null);
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnInteger, null, null, 18);
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 1);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("result"));
+        assertThat(Integer.valueOf(tags.get(0).getValue()), is(18));
+    }
+
+    @Test
+    public void testTraceWithReturnObject() throws Throwable {
+        Method testMethodWithReturnObject = TestAnnotationMethodClass.class.getDeclaredMethod("testMethodWithReturnObject", String.class, Integer.class);
+        methodInterceptor.beforeMethod(enhancedInstance, testMethodWithReturnObject, new Object[]{"wangwu", 18}, null, null);
+        methodInterceptor.afterMethod(enhancedInstance, testMethodWithReturnObject, null, null, new User("wangwu", 18));
+
+        ContextManager.stopSpan();
+        assertThat(storage.getTraceSegments().size(), is(1));
+        TraceSegment traceSegment = storage.getTraceSegments().get(0);
+        List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
+        assertThat(spans.size(), is(1));
+        AbstractTracingSpan tracingSpan = spans.get(0);
+        assertThat(tracingSpan.getOperationName(), is("testMethod"));
+        SpanAssert.assertLogSize(tracingSpan, 0);
+        SpanAssert.assertTagSize(tracingSpan, 1);
+        List<TagValuePair> tags = SpanHelper.getTags(tracingSpan);
+
+        assertThat(tags.get(0).getKey().key(), is("result"));
+        assertThat(tags.get(0).getValue(), anything());
+    }
+
     private class TestAnnotationMethodClass {
 
         @Tag(key = "username", value = "arg[0]")
@@ -153,6 +295,38 @@ public class TagAnnotationTest {
         public User testMethodWithTags(String username, Integer age) {
             return new User(username, age);
         }
+
+        @Tags({@Tag(key = "username", value = "arg[0]"), @Tag(key = "info", value = "returnedObj.0.info")})
+        public List<User> testMethodWithReturnList(String username, Integer age) {
+            return Arrays.asList(new User(username, age));
+        }
+
+        @Tags({@Tag(key = "username", value = "arg[0]"), @Tag(key = "info", value = "returnedObj.0.info")})
+        public User[] testMethodWithReturnArray(String username, Integer age) {
+            return new User[]{new User(username, age)};
+        }
+
+        @Tags({@Tag(key = "username", value = "arg[0]"), @Tag(key = "info", value = "returnedObj.user.info")})
+        public Map<String, User> testMethodWithReturnMap(String username, Integer age) {
+            Map<String, User> userMap = new HashMap<>();
+            userMap.put("user", new User(username, age));
+            return userMap;
+        }
+
+        @Tag(key = "result", value = "returnedObj")
+        public String testMethodWithReturnString(String username) {
+            return username;
+        }
+
+        @Tag(key = "result", value = "returnedObj")
+        public Integer testMethodWithReturnInteger(Integer age) {
+            return age;
+        }
+
+        @Tag(key = "result", value = "returnedObj")
+        public User testMethodWithReturnObject(String username, Integer age) {
+            return new User(username, age);
+        }
     }
 
     private class User {