Use byte buddy for instead of byteman.
RENAME
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/TraceAnnotationsManagerTest.java
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java
-package com.datadoghq.agent;
+package com.datadoghq.agent.instrumentation.annotation;
import static org.assertj.core.api.Assertions.assertThat;
@@ -14,7 +14,7 @@ import java.lang.reflect.Field;
import org.junit.Before;
import org.junit.Test;
-public class TraceAnnotationsManagerTest {
+public class TraceAnnotationsTest {
private final ListWriter writer = new ListWriter();
private final DDTracer tracer = new DDTracer(writer);
@@ -39,7 +39,8 @@ public class TraceAnnotationsManagerTest {
SayTracedHello.sayHello();
assertThat(writer.firstTrace().size()).isEqualTo(1);
- assertThat(writer.firstTrace().get(0).getOperationName()).isEqualTo("SAY_HELLO");
+ assertThat(writer.firstTrace().get(0).getOperationName())
+ .isEqualTo("com.datadoghq.agent.test.SayTracedHello.sayHello");
assertThat(writer.firstTrace().get(0).getServiceName()).isEqualTo("test");
}
@@ -56,7 +57,8 @@ public class TraceAnnotationsManagerTest {
assertThat(writer.firstTrace().get(0).context().getParentId()).isEqualTo(0);
assertThat(writer.firstTrace().get(0).getServiceName()).isEqualTo("test2");
- assertThat(writer.firstTrace().get(1).getOperationName()).isEqualTo("SAY_HELLO");
+ assertThat(writer.firstTrace().get(1).getOperationName())
+ .isEqualTo("com.datadoghq.agent.test.SayTracedHello.sayHello");
assertThat(writer.firstTrace().get(1).getServiceName()).isEqualTo("test");
assertThat(writer.firstTrace().get(1).getParentId()).isEqualTo(parentId);
@@ -82,7 +84,7 @@ public class TraceAnnotationsManagerTest {
final DDBaseSpan<?> span = writer.firstTrace().get(0);
assertThat(span.getOperationName()).isEqualTo("ERROR");
- assertThat(span.getTags().get("error")).isEqualTo("true");
+ assertThat(span.getTags().get("error")).isEqualTo(true);
assertThat(span.getTags().get("error.msg")).isEqualTo(error.getMessage());
assertThat(span.getTags().get("error.type")).isEqualTo(error.getClass().getName());
assertThat(span.getTags().get("error.stack")).isEqualTo(errorString.toString());

MODIFY
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/test/SayTracedHello.java
dd-java-agent-ittests/src/test/java/com/datadoghq/agent/test/SayTracedHello.java
@@ -7,7 +7,7 @@ import io.opentracing.util.GlobalTracer;
public class SayTracedHello {
- @Trace(operationName = "SAY_HELLO")
+ @Trace
public static String sayHello() {
new StringTag(DDTags.SERVICE_NAME).set(GlobalTracer.get().activeSpan(), "test");
return "hello!";

MODIFY
dd-java-agent/dd-java-agent.gradle
dd-java-agent/dd-java-agent.gradle
@@ -23,7 +23,6 @@ dependencies {
compile group: 'net.bytebuddy', name: 'byte-buddy', version: '1.7.6'
compile group: 'org.jboss.byteman', name: 'byteman', version: '3.0.10'
- compile group: 'org.reflections', name: 'reflections', version: '0.9.11'
compile group: 'com.google.auto.service', name: 'auto-service', version: '1.0-rc3'
compile deps.slf4j
@@ -87,9 +86,7 @@ shadowJar {
relocate 'org.msgpack', 'dd.deps.org.msgpack'
relocate 'com.fasterxml', 'dd.deps.com.fasterxml'
- relocate 'javassist', 'dd.deps.javassist'
relocate 'net.bytebuddy', 'dd.deps.net.bytebuddy'
- relocate 'org.reflections', 'dd.deps.org.reflections'
relocate('org.jboss.byteman', 'dd.deps.org.jboss.byteman') {
// Renaming these causes a verify error in the tests.
exclude 'org.jboss.byteman.rule.*'

MODIFY
dd-java-agent/src/main/java/com/datadoghq/agent/AgentRulesManager.java
dd-java-agent/src/main/java/com/datadoghq/agent/AgentRulesManager.java
@@ -32,12 +32,10 @@ public class AgentRulesManager {
protected final Retransformer transformer;
protected final TracingAgentConfig agentTracerConfig;
protected final InstrumentationRulesManager instrumentationRulesManager;
- protected final TraceAnnotationsManager traceAnnotationsManager;
public AgentRulesManager(final Retransformer trans, final TracingAgentConfig config) {
transformer = trans;
agentTracerConfig = config;
- traceAnnotationsManager = new TraceAnnotationsManager(trans, config);
instrumentationRulesManager = new InstrumentationRulesManager(trans, config, this);
}
@@ -62,7 +60,6 @@ public class AgentRulesManager {
INSTANCE = manager;
manager.loadRules(INITIALIZER_RULES, ClassLoader.getSystemClassLoader());
- manager.traceAnnotationsManager.initialize();
}
/**

DELETE
dd-java-agent/src/main/java/com/datadoghq/agent/TraceAnnotationsManager.java
-package com.datadoghq.agent;
-
-import com.datadoghq.trace.Trace;
-import io.opentracing.ActiveSpan;
-import io.opentracing.tag.Tags;
-import java.io.PrintWriter;
-import java.io.StringWriter;
-import java.lang.reflect.Method;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.Set;
-import javassist.ClassPool;
-import javassist.CtClass;
-import javassist.CtMethod;
-import javassist.bytecode.Descriptor;
-import lombok.extern.slf4j.Slf4j;
-import org.jboss.byteman.agent.Location;
-import org.jboss.byteman.agent.LocationType;
-import org.jboss.byteman.agent.Retransformer;
-import org.jboss.byteman.agent.RuleScript;
-import org.reflections.Reflections;
-import org.reflections.scanners.MethodAnnotationsScanner;
-import org.reflections.util.ConfigurationBuilder;
-import org.reflections.util.FilterBuilder;
-
-/**
- * This manager is loaded at pre-main.
- *
- * <p>This class adds rules to instrument all the methods annotated with the @Trace.
- */
-@Slf4j
-public class TraceAnnotationsManager {
- private static final String CURRENT_SPAN_EXISTS = "IF TRUE\n";
- private static final String BUILD_SPAN = "DO\n" + "getTracer().buildSpan(\"";
- private static final String CLOSE_PARENTHESIS = "\")";
- private static final String START = ".startActive();";
- private static final String EXIT_RULE =
- "IF getTracer().activeSpan() != null\n" + "DO\n" + "getTracer().activeSpan().deactivate();\n";
- private static final String EXCEPTION_EXIT_RULE =
- "BIND span:"
- + ActiveSpan.class.getName()
- + " = getTracer().activeSpan()\n"
- + "IF span != null\n"
- + "DO\n"
- + "span.setTag("
- + Tags.class.getName()
- + ".ERROR.getKey(),\"true\");\n"
- + "span.log("
- + Collections.class.getName()
- + ".singletonMap(\"error.object\",$^));\n"
- + "span.deactivate();\n";
-
- private final Retransformer transformer;
- private final TracingAgentConfig agentTracerConfig;
-
- public TraceAnnotationsManager(final Retransformer trans, final TracingAgentConfig config) {
- transformer = trans;
- agentTracerConfig = config;
- }
-
- /** This method initializes the manager. */
- public void initialize() {
- log.debug("Initializing {}", TraceAnnotationsManager.class.getSimpleName());
-
- // Check if annotations are enabled
- if (agentTracerConfig != null
- && agentTracerConfig.getEnableCustomAnnotationTracingOver() != null
- && agentTracerConfig.getEnableCustomAnnotationTracingOver().length > 0) {
- loadAnnotationsRules(agentTracerConfig.getEnableCustomAnnotationTracingOver());
- }
- }
-
- /** Find all the methods annoted with @Trace and inject rules */
- // FIXME: this needs to be done for each classloader.
- private void loadAnnotationsRules(final String... scannedPackages) {
- log.info(
- "Looking for annotations over the following packages: {}", Arrays.asList(scannedPackages));
- final long cur = System.currentTimeMillis();
-
- final Reflections reflections =
- new Reflections(
- new ConfigurationBuilder()
- .forPackages(scannedPackages)
- .filterInputsBy(
- new FilterBuilder().includePackage(scannedPackages).include(".*\\.class"))
- .setScanners(new MethodAnnotationsScanner()));
-
- final Set<Method> methods = reflections.getMethodsAnnotatedWith(Trace.class);
-
- final StringBuilder generatedScripts = new StringBuilder();
- for (final Method method : methods) {
- try {
- final ClassPool pool = ClassPool.getDefault();
- final CtClass cc = pool.get(method.getDeclaringClass().getCanonicalName());
- final CtMethod javassistMethod = cc.getDeclaredMethod(method.getName());
-
- // AT ENTRY: child of current case
- final String ruleText = CURRENT_SPAN_EXISTS + buildSpan(javassistMethod) + START;
- RuleScript script =
- createRuleScript(
- "Start Active Span ",
- cc,
- javassistMethod,
- Location.create(LocationType.ENTRY, ""),
- ruleText);
- generatedScripts.append(script).append("\n");
-
- // AT EXIT
- script =
- createRuleScript(
- "Close span ",
- cc,
- javassistMethod,
- Location.create(LocationType.EXIT, ""),
- EXIT_RULE);
- generatedScripts.append(script).append("\n");
-
- // AT EXCEPTION EXIT
- script =
- createRuleScript(
- "Close span in error ",
- cc,
- javassistMethod,
- Location.create(LocationType.EXCEPTION_EXIT, ""),
- EXCEPTION_EXIT_RULE);
- generatedScripts.append(script).append("\n");
-
- } catch (final Exception e) {
- log.warn(
- "Could not create rule for method " + method + ". Proceed to next annoted method.", e);
- }
- log.trace("Instrumenting annotated method: {}", method.getName());
- }
- try {
- final StringWriter sw = new StringWriter();
- try (PrintWriter pr = new PrintWriter(sw)) {
- transformer.installScript(
- Arrays.asList(generatedScripts.toString()), Arrays.asList("@Trace annotations"), pr);
- }
- log.debug("Install new rules: \n{}", sw.toString());
- } catch (final Exception e) {
- log.warn("Could not install annotation scripts.", e);
- }
-
- log.info(
- "Finished annotation scanning in "
- + (System.currentTimeMillis() - cur)
- + " ms. You can accelerate this process by refining the packages you want to scan with `scannedPackages` in the dd-trace.yaml configuration file.");
- }
-
- private RuleScript createRuleScript(
- final String ruleNamePrefix,
- final CtClass cc,
- final CtMethod javassistMethod,
- final Location loc,
- final String ruleText) {
- final int lineNumber = javassistMethod.getMethodInfo().getLineNumber(0);
- final String[] imports = new String[0];
-
- final RuleScript ruleScript =
- new RuleScript(
- ruleNamePrefix + loc + " " + javassistMethod.getLongName(),
- cc.getName(),
- false,
- false,
- javassistMethod.getName() + Descriptor.toString(javassistMethod.getSignature()),
- OpenTracingHelper.class.getName(),
- imports,
- loc,
- ruleText,
- lineNumber,
- "",
- false);
- return ruleScript;
- }
-
- private String buildSpan(final CtMethod javassistMethod) {
- try {
- final Trace trace = (Trace) javassistMethod.getAnnotation(Trace.class);
- if (trace.operationName() != null & !trace.operationName().isEmpty()) {
- return BUILD_SPAN + trace.operationName() + CLOSE_PARENTHESIS;
- }
- } catch (final Exception e) {
- log.warn(
- "Error when building injection rule on method "
- + javassistMethod
- + ". Fallback on default value.",
- e);
- }
- return BUILD_SPAN + javassistMethod.getName() + CLOSE_PARENTHESIS;
- }
-}

ADD
dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationInstrumentation.java
+package com.datadoghq.agent.instrumentation.annotation;
+
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
+import com.datadoghq.trace.Trace;
+import com.google.auto.service.AutoService;
+import io.opentracing.ActiveSpan;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Method;
+import java.sql.PreparedStatement;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class TraceAnnotationInstrumentation implements Instrumenter {
+ public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
+
+ @Override
+ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+ return agentBuilder
+ .type(hasSuperType(declaresMethod(isAnnotatedWith(Trace.class))))
+ .transform(
+ new AgentBuilder.Transformer.ForAdvice()
+ .advice(isAnnotatedWith(Trace.class), TraceAdvice.class.getName()));
+ }
+
+ public static class TraceAdvice {
+
+ @Advice.OnMethodEnter
+ public static ActiveSpan startSpan(@Advice.Origin final Method method) {
+ final Trace trace = method.getAnnotation(Trace.class);
+ String operationName = trace.operationName();
+ if (operationName == null || operationName.isEmpty()) {
+ operationName = method.getDeclaringClass().getName() + "." + method.getName();
+ }
+
+ return GlobalTracer.get().buildSpan(operationName).startActive();
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class)
+ public static void stopSpan(
+ @Advice.Enter final ActiveSpan activeSpan, @Advice.Thrown final Throwable throwable) {
+ if (throwable != null) {
+ Tags.ERROR.set(activeSpan, true);
+ activeSpan.log(Collections.singletonMap("error.object", throwable));
+ }
+ activeSpan.deactivate();
+ }
+ }
+}

Use byte buddy for @Trace instead of byteman.