fix review issues

Add weakReference for original classes map
Add try/catch on retransform call in removeAllProbes
Handle null class (newly loaded) based on the name for
re-transformation
Use AgentTaskScheduler for scheduling poll
Fix tests

Author: jpbempel

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

@@ -3,6 +3,7 @@
 import com.datadog.debugging.probes.ProbesPoller;
 import java.lang.instrument.Instrumentation;
 import java.lang.instrument.UnmodifiableClassException;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -81,23 +82,40 @@ public void accept(DebuggerProbe[] debuggerProbes) {
     }
   }
 
-  private void removeAllProbes() throws UnmodifiableClassException {
+  private void removeAllProbes() {
     if (currentTransformer == null) {
       return;
     }
     log.info("Remove probes by restoring & re-transforming original class definitions...");
-    Map<String, Class<?>> originalClasses = currentTransformer.getOriginalClasses();
+    Map<String, WeakReference<Class<?>>> originalClasses = currentTransformer.getOriginalClasses();
     instrumentation.removeTransformer(currentTransformer);
-    for (Map.Entry<String, Class<?>> entry : originalClasses.entrySet()) {
-      Class<?> clazz = entry.getValue();
+    Set<String> newlyLoadedClassNames = new HashSet<>();
+    for (Map.Entry<String, WeakReference<Class<?>>> entry : originalClasses.entrySet()) {
+      Class<?> clazz = entry.getValue().get();
       if (clazz != null) {
-        log.info("Restoring class {}", clazz.getName());
-        instrumentation.retransformClasses(clazz);
+        restoreClass(clazz);
       } else {
-        log.warn("class null for: {} cannot re-transform", entry.getKey());
+        newlyLoadedClassNames.add(entry.getKey());
+      }
+    }
+    // process nealy loaded class (no Class instance at the time of transform) based on the name
+    if (!newlyLoadedClassNames.isEmpty()) {
+      for (Class<?> clazz : instrumentation.getAllLoadedClasses()) {
+        if (newlyLoadedClassNames.contains(clazz.getName())) {
+          restoreClass(clazz);
+        }
       }
     }
     probeIds.clear();
     currentTransformer = null;
   }
+
+  private void restoreClass(Class<?> clazz) {
+    log.info("Restoring class {}", clazz.getName());
+    try {
+      instrumentation.retransformClasses(clazz);
+    } catch (UnmodifiableClassException ex) {
+      log.warn("Cannot re-transform class {}:", clazz.getName(), ex);
+    }
+  }
 }

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

@@ -7,6 +7,7 @@
 import datadog.trace.bootstrap.debugging.DebuggerContext;
 import java.io.IOException;
 import java.lang.instrument.ClassFileTransformer;
+import java.lang.ref.WeakReference;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -27,14 +28,14 @@
 public class DebuggerTransformer implements ClassFileTransformer {
   private final DebuggerProbe[] probes;
   private final Config config;
-  private final Map<String, Class<?>> originalClasses = new HashMap<>();
+  private final Map<String, WeakReference<Class<?>>> originalClasses = new HashMap<>();
 
   public DebuggerTransformer(Config config, DebuggerProbe... probes) {
     this.probes = probes;
     this.config = config;
   }
 
-  public Map<String, Class<?>> getOriginalClasses() {
+  public Map<String, WeakReference<Class<?>>> getOriginalClasses() {
     return originalClasses;
   }
 
@@ -62,7 +63,7 @@ public byte[] transform(
       String typeName = className.replace('/', '.');
       if (where.isTypeMatching(typeName)) {
         log.debug("Matched type '{}'", typeName);
-        originalClasses.put(className, classBeingRedefined);
+        originalClasses.put(className, new WeakReference<>(classBeingRedefined));
         ClassReader reader = new ClassReader(classfileBuffer);
         dumpOriginalClassFile(className, classfileBuffer);
         ClassNode classNode = new ClassNode();

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

@@ -3,9 +3,6 @@
 import com.datadog.debugging.probes.ProbesPoller;
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugging.DebuggerContext;
-import java.io.FileInputStream;
-import java.io.IOException;
-import java.io.InputStream;
 import java.lang.instrument.Instrumentation;
 import lombok.extern.slf4j.Slf4j;
 

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/probes/ProbesPoller.java

@@ -5,6 +5,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import datadog.trace.api.Config;
 import datadog.trace.api.RatelimitedLogger;
+import datadog.trace.util.AgentTaskScheduler;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -26,17 +27,16 @@
 import okhttp3.ResponseBody;
 
 @Slf4j
-public class ProbesPoller {
+public class ProbesPoller implements AgentTaskScheduler.Target<ProbesPoller> {
   private static final long NANOSECONDS_BETWEEN_ERROR_LOG = TimeUnit.MINUTES.toNanos(5);
   private static final int MAX_RUNNING_REQUESTS = 1;
 
   private final ScheduledExecutorService okHttpExecutorService;
   private final OkHttpClient client;
-  private final String url;
   private final RatelimitedLogger ratelimitedLogger;
   private final ProbeRedefinitionListener listener;
-  private ScheduledFuture<?> future;
   private final Request request;
+  private final long pollInterval;
 
   public interface ProbeRedefinitionListener {
     void accept(DebuggerProbe[] debuggerProbes);
@@ -48,12 +48,14 @@ public ProbesPoller(Config config, ProbeRedefinitionListener listener) {
 
   ProbesPoller(
       Config config, ProbeRedefinitionListener listener, RatelimitedLogger ratelimitedLogger) {
-    url = config.getDebuggingProbeUrl();
+    String url = config.getDebuggingProbeUrl();
     if (url == null || url.length() == 0) {
       throw new IllegalArgumentException("Probe url is empty");
     }
     this.listener = listener;
     this.ratelimitedLogger = ratelimitedLogger;
+    // TODO add a jitter to avoid herd issue
+    this.pollInterval = Duration.ofSeconds(config.getDebuggingPollInterval()).toMillis();
     log.debug("Started Probes Poller with target url {}", url);
     // This is the same thing OkHttp Dispatcher is doing except thread naming and daemonization
     okHttpExecutorService =
@@ -86,16 +88,17 @@ public ProbesPoller(Config config, ProbeRedefinitionListener listener) {
     request = new Request.Builder().url(url).get().build();
   }
 
-  public void start() {
-    future =
-        okHttpExecutorService.scheduleAtFixedRate(this::pollDebuggerProbes, 0, 1, TimeUnit.SECONDS);
+  @Override
+  public ProbesPoller get() {
+    return this;
   }
 
-  public void stop() {
-    future.cancel(true);
+  public void start() {
+    AgentTaskScheduler.INSTANCE.scheduleAtFixedRate(
+        this::pollDebuggerProbes, this, 0, pollInterval, TimeUnit.MILLISECONDS);
   }
 
-  void pollDebuggerProbes() {
+  void pollDebuggerProbes(AgentTaskScheduler.Target<ProbesPoller> target) {
     try {
       try (Response response = client.newCall(request).execute()) {
         ResponseBody body = response.body();

dd-java-agent/agent-debugging/src/main/java/com/datadog/debugging/uploader/SnapshotUploader.java

@@ -13,7 +13,6 @@
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-
 import lombok.extern.slf4j.Slf4j;
 import okhttp3.Call;
 import okhttp3.Callback;
@@ -37,6 +36,7 @@ public enum Kind {
     DIAGNOSTIC("diagnostic");
 
     private String apiTarget;
+
     Kind(String apiTarget) {
       this.apiTarget = apiTarget;
     }
@@ -72,7 +72,7 @@ public SnapshotUploader(Config config) {
 
   // Visible for testing
   SnapshotUploader(Config config, RatelimitedLogger ratelimitedLogger) {
-    url = config.getDebuggingSnapshotUrl();
+    String url = config.getDebuggingSnapshotUrl();
     if (url == null || url.length() == 0) {
       throw new IllegalArgumentException("Snapshot url is empty");
     }

dd-java-agent/agent-debugging/src/test/java/com/datadog/debugging/agent/DebuggerProbeRedefinitionTest.java

@@ -11,6 +11,7 @@
 
 import java.lang.instrument.Instrumentation;
 import java.lang.instrument.UnmodifiableClassException;
+import java.lang.ref.WeakReference;
 import java.util.HashMap;
 import java.util.Map;
 import org.junit.jupiter.api.Test;
@@ -64,15 +65,15 @@ private static DebuggerTransformer createTransformer(DebuggerProbe[] debuggerPro
   private static DebuggerTransformer createTransformerWithOrigClasses(
       DebuggerProbe[] debuggerProbes) {
     DebuggerTransformer transformer = mock(DebuggerTransformer.class);
-    Map<String, Class<?>> origClasses = new HashMap<>();
+    Map<String, WeakReference<Class<?>>> origClasses = new HashMap<>();
     for (DebuggerProbe probe : debuggerProbes) {
       Class<?> clazz = null;
       try {
         clazz = Class.forName(probe.getWhere().getTypeName());
       } catch (ClassNotFoundException e) {
         e.printStackTrace();
       }
-      origClasses.put(clazz.getName(), clazz);
+      origClasses.put(clazz.getName(), new WeakReference<>(clazz));
     }
     when(transformer.getOriginalClasses()).thenReturn(origClasses);
     return transformer;

dd-java-agent/agent-debugging/src/test/java/com/datadog/debugging/agent/SnapshotSinkTest.java

@@ -1,42 +1,65 @@
 package com.datadog.debugging.agent;
 
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.matches;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.datadog.debugging.uploader.SnapshotUploader;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.debugging.DiagnosticMessage;
 import datadog.trace.bootstrap.debugging.Snapshot;
+import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
 @ExtendWith(MockitoExtension.class)
 public class SnapshotSinkTest {
-  private static final String URL_PATH = "/lalala";
   private static final String PROBE_ID = "12fd-8490-c111-4374-ffde";
 
   @Mock private Config config;
   @Mock private SnapshotUploader snapshotUploader;
 
   @Test
-  public void add() {
+  public void addSnapshot() {
     when(config.getDebuggingUploadFormat()).thenReturn("json");
     SnapshotSink sink = new SnapshotSink(config, snapshotUploader);
     Snapshot snapshot = new Snapshot(Thread.currentThread(), PROBE_ID);
     sink.addSnapshot(snapshot);
     String regex =
         "\\{\"startTs\":\\d+,\"timestamp\":\\d+,\"probeId\":\"12fd-8490-c111-4374-ffde\",\"language\":\"java\",\"thread\":\\{\"id\":\\d+,\"name\":\"Test worker\"\\},\"version\":0,\"captures\":\\[\\],\"duration\":0\\}";
-    verify(snapshotUploader).upload(matches(regex), SnapshotUploader.Kind.SNAPSHOT);
+    verify(snapshotUploader).upload(matches(regex), eq(SnapshotUploader.Kind.SNAPSHOT));
   }
 
   @Test
-  public void addBinary() {
+  public void addSnapshotBinary() {
     when(config.getDebuggingUploadFormat()).thenReturn("binary");
     SnapshotSink sink = new SnapshotSink(config, snapshotUploader);
     Snapshot snapshot = new Snapshot(Thread.currentThread(), PROBE_ID);
     sink.addSnapshot(snapshot);
     // FIXME when binary format is implemented and used
   }
+
+  @Test
+  public void addDiagnostics() {
+    when(config.getDebuggingUploadFormat()).thenReturn("json");
+    SnapshotSink sink = new SnapshotSink(config, snapshotUploader);
+    sink.addDiagnostics(
+        PROBE_ID, Arrays.asList(new DiagnosticMessage(DiagnosticMessage.Kind.ERROR, "error1")));
+    String regex =
+        "\\{\"probeId\":\"12fd-8490-c111-4374-ffde\",\"messages\":\\[\\{\"timestamp\":\\d+,\"kind\":\"ERROR\",\"message\":\"error1\"\\}\\]\\}";
+    verify(snapshotUploader).upload(matches(regex), eq(SnapshotUploader.Kind.DIAGNOSTIC));
+  }
+
+  @Test
+  public void addDiagnosticsBinary() {
+    when(config.getDebuggingUploadFormat()).thenReturn("binary");
+    SnapshotSink sink = new SnapshotSink(config, snapshotUploader);
+    Snapshot snapshot = new Snapshot(Thread.currentThread(), PROBE_ID);
+    sink.addDiagnostics(
+        PROBE_ID, Arrays.asList(new DiagnosticMessage(DiagnosticMessage.Kind.ERROR, "error1")));
+    // FIXME when binary format is implemented and used
+  }
 }

dd-java-agent/agent-debugging/src/test/java/com/datadog/debugging/probes/ProbesPollerTest.java

@@ -49,15 +49,15 @@ public void receive200Response() {
         new MockResponse().setResponseCode(200).setBody(getResourceContent("/test_probe.json")));
     when(config.getDebuggingProbeUrl()).thenReturn(server.url(URL_PATH).toString());
     ProbesPoller poller = new ProbesPoller(config, ProbesPollerTest::assertGoodProbeDefinition);
-    poller.pollDebuggerProbes();
+    poller.pollDebuggerProbes(null);
   }
 
   @Test
   public void receive500Response() {
     server.enqueue(new MockResponse().setResponseCode(500));
     when(config.getDebuggingProbeUrl()).thenReturn(server.url(URL_PATH).toString());
     ProbesPoller poller = new ProbesPoller(config, ProbesPollerTest::assertGoodProbeDefinition);
-    poller.pollDebuggerProbes();
+    poller.pollDebuggerProbes(null);
   }
 
   @Test
@@ -69,7 +69,7 @@ public void receive500ResponseWithBody() {
             .setBody("{\"message\": \"error\"}"));
     when(config.getDebuggingProbeUrl()).thenReturn(server.url(URL_PATH).toString());
     ProbesPoller poller = new ProbesPoller(config, ProbesPollerTest::assertGoodProbeDefinition);
-    poller.pollDebuggerProbes();
+    poller.pollDebuggerProbes(null);
   }
 
   @Test

dd-java-agent/agent-debugging/src/test/java/com/datadog/debugging/uploader/SnapshotUploaderTest.java

@@ -120,7 +120,11 @@ void testDiagnostics() throws Exception {
 
     RecordedRequest request = server.takeRequest(5, TimeUnit.SECONDS);
     assertNotNull(request);
-    assertTrue(request.getRequestUrl().toString().endsWith("/" + SnapshotUploader.Kind.DIAGNOSTIC.getApiTarget()));
+    assertTrue(
+        request
+            .getRequestUrl()
+            .toString()
+            .endsWith("/" + SnapshotUploader.Kind.DIAGNOSTIC.getApiTarget()));
     assertTrue(request.getBody().size() > 0);
   }
 
@@ -142,7 +146,10 @@ public void testConnectionRefused() throws IOException, InterruptedException {
     // Shutting down uploader ensures all callbacks are called on http client
     uploader.shutdown();
     verify(ratelimitedLogger)
-        .warn(eq("Failed to upload snapshot to {}"), ArgumentMatchers.argThat(arg -> arg.toString().startsWith(url.toString())), any(ConnectException.class));
+        .warn(
+            eq("Failed to upload snapshot to {}"),
+            ArgumentMatchers.argThat(arg -> arg.toString().startsWith(url.toString())),
+            any(ConnectException.class));
   }
 
   @Test

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -67,6 +67,7 @@ public final class ConfigDefaults {
   static final int DEFAULT_DEBUGGING_UPLOAD_TIMEOUT = 30; // seconds
   static final String DEFAULT_DEBUGGING_UPLOAD_FORMAT = "binary";
   static final boolean DEFAULT_DEBUGGING_CLASSFILE_DUMP_ENABLED = false;
+  static final int DEFAULT_DEBUGGING_POLL_INTERVAL = 1; // seconds
 
   static final boolean DEFAULT_KAFKA_CLIENT_PROPAGATION_ENABLED = true;