diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/util/CustomizeExpression.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/util/CustomizeExpression.java index c03a12101..9a301d762 100644 --- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/util/CustomizeExpression.java +++ b/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 evaluationReturnContext(Object ret) { Map 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 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) { diff --git a/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/util/TagUtil.java b/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/util/TagUtil.java index f5b5153b7..3a20d1a03 100644 --- a/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/util/TagUtil.java +++ b/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]); } } diff --git a/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/test/java/org/apache/skywalking/apm/toolkit/activation/trace/TagAnnotationTest.java b/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/test/java/org/apache/skywalking/apm/toolkit/activation/trace/TagAnnotationTest.java index 02ddf68a3..1a3b093c0 100644 --- a/apm-sniffer/apm-toolkit-activation/apm-toolkit-trace-activation/src/test/java/org/apache/skywalking/apm/toolkit/activation/trace/TagAnnotationTest.java +++ b/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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 testMethodWithReturnMap(String username, Integer age) { + Map 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 {