Pull probe definitions from backend

Add a ProbesPoller to poll regularly (1s) from the backend
we store classes instrumented to be able to re-tranform classes to
remove probes
we store probeIds to be able to make a diff of new definitions

Author: jpbempel

dd-java-agent/agent-debugging/agent-debugging.gradle

@@ -18,6 +18,7 @@ excludedClassesCoverage += [
   'com.datadog.debugging.parser.*Exception',
   // for generics type information trick see com.fasterxml.jackson.core.type.TypeReference javadoc
   'com.datadog.debugging.agent.DebuggingAgent.1',
+  'com.datadog.debugging.agent.ProbesPoller.1',
   'com.datadog.debugging.agent.DebuggerTransformer.SafeClassWriter'
 ]
 

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/agent/DebuggerProbe.java

@@ -72,14 +72,19 @@ public static SourceLine fromString(String lineDef) {
   public static final class Where {
     private String typeName;
     private String methodName;
-    private String source;
+    private String sourceFile;
     private String signature;
     private SourceLine[] lines;
 
-    Where(String typeName, String methodName, String signature, SourceLine[] lines, String source) {
+    Where(
+        String typeName,
+        String methodName,
+        String signature,
+        SourceLine[] lines,
+        String sourceFile) {
       this.typeName = typeName;
       this.methodName = methodName;
-      this.source = source;
+      this.sourceFile = sourceFile;
       this.signature = signature;
       this.lines = lines;
     }
@@ -90,8 +95,8 @@ public static final class Where {
         @JsonProperty("methodName") String methodName,
         @JsonProperty("signature") String signature,
         @JsonProperty("lines") String[] lines,
-        @JsonProperty("source") String source) {
-      this(typeName, methodName, signature, sourceLines(lines), source);
+        @JsonProperty("sourceFile") String sourceFile) {
+      this(typeName, methodName, signature, sourceLines(lines), sourceFile);
     }
 
     public Where() {}
@@ -116,8 +121,8 @@ public Where lines(String... lines) {
       return this;
     }
 
-    public Where source(String source) {
-      this.source = source;
+    public Where sourceFile(String sourceFile) {
+      this.sourceFile = sourceFile;
       return this;
     }
 
@@ -140,8 +145,8 @@ public String getMethodName() {
       return methodName;
     }
 
-    public String getSource() {
-      return source;
+    public String getSourceFile() {
+      return sourceFile;
     }
 
     public String getSignature() {
@@ -282,6 +287,8 @@ public static Tag fromString(String def) {
   private final int version;
   private final String language;
   private final String probeId;
+  private final long orgId;
+  private final String appId;
   private Where where = null;
   private When when = null;
   private Tag[] tags = null;
@@ -292,12 +299,16 @@ public DebuggerProbe(
       @JsonProperty("version") int version,
       @JsonProperty("language") String language,
       @JsonProperty("probeId") String probeId,
+      @JsonProperty("orgId") long orgId,
+      @JsonProperty("appId") String appId,
       @JsonProperty("tags") String[] tags,
       @JsonProperty("where") Where where,
       @JsonProperty("when") When when) {
     this.version = version;
     this.language = language;
     this.probeId = probeId;
+    this.orgId = orgId;
+    this.appId = appId;
     this.where = where;
     this.when = when;
     if (tags != null) {
@@ -310,7 +321,7 @@ public DebuggerProbe(
   }
 
   public DebuggerProbe() {
-    this(VERSION, LANGUAGE, UUID.randomUUID().toString(), null, null, null);
+    this(VERSION, LANGUAGE, UUID.randomUUID().toString(), -1, null, null, null, null);
   }
 
   public Where getWhere() {

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/agent/DebuggerProbeRedefinition.java

@@ -0,0 +1,103 @@
+package com.datadog.debugging.agent;
+
+import com.datadog.debugging.probes.ProbesPoller;
+import java.lang.instrument.Instrumentation;
+import java.lang.instrument.UnmodifiableClassException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class DebuggerProbeRedefinition implements ProbesPoller.ProbeRedefinitionListener {
+  private final Instrumentation instrumentation;
+  private final Function<DebuggerProbe[], DebuggerTransformer> transformerSupplier;
+  private DebuggerTransformer currentTransformer;
+  private Set<String> probeIds = new HashSet<>();
+
+  public DebuggerProbeRedefinition(
+      Instrumentation instrumentation,
+      Function<DebuggerProbe[], DebuggerTransformer> transformerSupplier) {
+    this.instrumentation = instrumentation;
+    this.transformerSupplier = transformerSupplier;
+  }
+
+  // Should be called by only one thread
+  public void accept(DebuggerProbe[] debuggerProbes) {
+    try {
+      if (debuggerProbes == null || debuggerProbes.length == 0) {
+        log.debug("No debugger probes defined.");
+        removeAllProbes();
+        return;
+      }
+      // check already defined probes
+      Map<String, DebuggerProbe> newDefinitions = new HashMap<>();
+      for (DebuggerProbe definition : debuggerProbes) {
+        if (!probeIds.contains(definition.getProbeId())) {
+          newDefinitions.put(definition.getWhere().getTypeName(), definition);
+        }
+      }
+      if (newDefinitions.isEmpty()) {
+        return;
+      }
+      removeAllProbes();
+      log.info("Applying new probe definitions by re-transforming classes...");
+      // install new probe definitions
+      currentTransformer = transformerSupplier.apply(debuggerProbes);
+      instrumentation.addTransformer(currentTransformer, true);
+      log.info("re-transforming {} classes", instrumentation.getAllLoadedClasses().length);
+      // Build list of classes needed to be re-transformed, matching probe definitions
+      List<Class<?>> classesToBeTransformed = new ArrayList<>();
+      for (Class<?> clazz : instrumentation.getAllLoadedClasses()) {
+        String typeName = clazz.getName().replace('/', '.');
+        DebuggerProbe debuggerProbe = newDefinitions.get(typeName);
+        if (debuggerProbe != null) {
+          classesToBeTransformed.add(clazz);
+        }
+      }
+      // re-transforms matching classes
+      for (Class<?> clazz : classesToBeTransformed) {
+        try {
+          log.info("re-transformed {}", clazz.getCanonicalName());
+          instrumentation.retransformClasses(clazz);
+        } catch (Exception ex) {
+          log.warn("re-transform error: ", ex);
+        } catch (Throwable ex) {
+          log.warn("throwable: ", ex);
+        }
+      }
+      probeIds =
+          newDefinitions.values().stream()
+              .map(DebuggerProbe::getProbeId)
+              .collect(Collectors.toSet());
+      log.info("Re-transformation done.");
+    } catch (Exception e) {
+      log.warn("Error during redefinition of the probes: ", e);
+    }
+  }
+
+  private void removeAllProbes() throws UnmodifiableClassException {
+    if (currentTransformer == null) {
+      return;
+    }
+    log.info("Remove probes by restoring & re-transforming original class definitions...");
+    Map<String, Class<?>> originalClasses = currentTransformer.getOriginalClasses();
+    instrumentation.removeTransformer(currentTransformer);
+    for (Map.Entry<String, Class<?>> entry : originalClasses.entrySet()) {
+      Class<?> clazz = entry.getValue();
+      if (clazz != null) {
+        log.info("Restoring class {}", clazz.getName());
+        instrumentation.retransformClasses(clazz);
+      } else {
+        log.warn("class null for: {} cannot re-transform", entry.getKey());
+      }
+    }
+    probeIds.clear();
+    currentTransformer = null;
+  }
+}

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/agent/DebuggerTransformer.java

@@ -12,6 +12,8 @@
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
 import java.security.ProtectionDomain;
+import java.util.HashMap;
+import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
@@ -22,23 +24,27 @@
 import org.objectweb.asm.tree.MethodNode;
 
 @Slf4j
-public final class DebuggerTransformer implements ClassFileTransformer {
+public class DebuggerTransformer implements ClassFileTransformer {
   private final DebuggerProbe[] probes;
   private final Config config;
+  private final Map<String, Class<?>> originalClasses = new HashMap<>();
 
   public DebuggerTransformer(Config config, DebuggerProbe... probes) {
     this.probes = probes;
     this.config = config;
   }
 
+  public Map<String, Class<?>> getOriginalClasses() {
+    return originalClasses;
+  }
+
   @Override
   public byte[] transform(
       ClassLoader loader,
       String className,
       Class<?> classBeingRedefined,
       ProtectionDomain protectionDomain,
       byte[] classfileBuffer) {
-    log.debug("Try transform: {}", className);
     if (probes == null || probes.length == 0) {
       // no transformation
       log.warn("No debugger probes defined.");
@@ -56,6 +62,7 @@ public byte[] transform(
       String typeName = className.replace('/', '.');
       if (where.isTypeMatching(typeName)) {
         log.debug("Matched type '{}'", typeName);
+        originalClasses.put(className, classBeingRedefined);
         ClassReader reader = new ClassReader(classfileBuffer);
         dumpOriginalClassFile(className, classfileBuffer);
         ClassNode classNode = new ClassNode();
@@ -81,6 +88,8 @@ public byte[] transform(
           byte[] data = writer.toByteArray();
           dumpInstrumentedClassFile(className, data);
           return data;
+        } else {
+          log.warn("type matched but no transformation");
         }
       }
     } catch (Exception ex) {

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/agent/DebuggingAgent.java

@@ -1,7 +1,6 @@
 package com.datadog.debugging.agent;
 
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
+import com.datadog.debugging.probes.ProbesPoller;
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugging.DebuggerContext;
 import java.io.FileInputStream;
@@ -13,25 +12,23 @@
 @Slf4j
 public class DebuggingAgent {
 
+  private static ProbesPoller probesPoller;
+
   public static synchronized void run(Instrumentation instrumentation) {
     if (Config.get().isDebuggingEnabled()) {
       log.info("starting debugging agent");
       SnapshotSink sink = new SnapshotSink(Config.get());
       DebuggerContext.init(sink);
-      String configFileName = System.getProperty("dd.debugger.config-file", "debugger-config.json");
-      try {
-        DebuggerProbe[] probes = null;
-        try (InputStream is = new FileInputStream(configFileName)) {
-          ObjectMapper mapper = new ObjectMapper();
-          probes = mapper.readValue(is, new TypeReference<DebuggerProbe[]>() {});
-        }
-        instrumentation.addTransformer(new DebuggerTransformer(Config.get(), probes), true);
-      } catch (IOException ex) {
-        log.warn(
-            "Unable to load config file: " + configFileName + ". Debugger agent is disabled.", ex);
-      }
+      DebuggerProbeRedefinition debuggerProbeRedefinition =
+          new DebuggerProbeRedefinition(instrumentation, DebuggingAgent::createTransformer);
+      probesPoller = new ProbesPoller(Config.get(), debuggerProbeRedefinition);
+      probesPoller.start();
     } else {
       log.info("Debugging agent disabled");
     }
   }
+
+  private static DebuggerTransformer createTransformer(DebuggerProbe[] debuggerProbes) {
+    return new DebuggerTransformer(Config.get(), debuggerProbes);
+  }
 }