Fix DebuggerAgent tests (#9951)

Inject a config instance to DebuggerAgent and stop rely on
Config.get() to be able to provide clean config at any time.
add some reset method for reseting static fields
Fix readFromFile test that wasn't at all executed

Author: jpbempel

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -1396,8 +1396,8 @@ private static synchronized void startDebuggerAgent(
       final Class<?> debuggerAgentClass =
           AGENT_CLASSLOADER.loadClass("com.datadog.debugger.agent.DebuggerAgent");
       final Method debuggerInstallerMethod =
-          debuggerAgentClass.getMethod("run", Instrumentation.class, scoClass);
-      debuggerInstallerMethod.invoke(null, inst, sco);
+          debuggerAgentClass.getMethod("run", Config.class, Instrumentation.class, scoClass);
+      debuggerInstallerMethod.invoke(null, Config.get(), inst, sco);
     } catch (final Throwable ex) {
       log.error("Throwable thrown while starting debugger agent", ex);
     } finally {

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java

@@ -72,23 +72,23 @@ public class DebuggerAgent {
   static final AtomicBoolean distributedDebuggerEnabled = new AtomicBoolean();
   private static ClassesToRetransformFinder classesToRetransformFinder;
 
-  public static synchronized void run(Instrumentation inst, SharedCommunicationObjects sco) {
+  public static synchronized void run(
+      Config config, Instrumentation inst, SharedCommunicationObjects sco) {
     instrumentation = inst;
     sharedCommunicationObjects = sco;
-    Config config = Config.get();
     classesToRetransformFinder = new ClassesToRetransformFinder();
     setupSourceFileTracking(instrumentation, classesToRetransformFinder);
     // set config updater after setup is done, as some deferred updates might be immediately called
-    DebuggerConfigBridge.setUpdater(new DefaultDebuggerConfigUpdater());
+    DebuggerConfigBridge.setUpdater(new DefaultDebuggerConfigUpdater(config));
     if (config.isDebuggerCodeOriginEnabled()) {
-      startCodeOriginForSpans();
+      startCodeOriginForSpans(config);
     }
     if (config.isDebuggerExceptionEnabled()) {
-      startExceptionReplay();
+      startExceptionReplay(config);
     }
     if (config.isDynamicInstrumentationEnabled()) {
-      startDynamicInstrumentation();
-      startCodeOriginForSpans();
+      startDynamicInstrumentation(config);
+      startCodeOriginForSpans(config);
       if (config.getDynamicInstrumentationInstrumentTheWorld() != null) {
         setupInstrumentTheWorldTransformer(config, instrumentation, sink);
       }
@@ -145,12 +145,11 @@ private static void initClassNameFilter() {
     }
   }
 
-  public static void startDynamicInstrumentation() {
+  public static void startDynamicInstrumentation(Config config) {
     if (!dynamicInstrumentationEnabled.compareAndSet(false, true)) {
       return;
     }
     LOGGER.info("Starting Dynamic Instrumentation");
-    Config config = Config.get();
     commonInit(config);
     String probeFileLocation = config.getDynamicInstrumentationProbeFile();
     if (probeFileLocation != null) {
@@ -206,12 +205,11 @@ public static void stopDynamicInstrumentation() {
     }
   }
 
-  public static void startExceptionReplay() {
+  public static void startExceptionReplay(Config config) {
     if (!exceptionReplayEnabled.compareAndSet(false, true)) {
       return;
     }
     LOGGER.info("Starting Exception Replay");
-    Config config = Config.get();
     commonInit(config);
     initClassNameFilter();
     if (config.isCiVisibilityEnabled()) {
@@ -238,12 +236,11 @@ public static void stopExceptionReplay() {
     DebuggerContext.initExceptionDebugger(null);
   }
 
-  public static void startCodeOriginForSpans() {
+  public static void startCodeOriginForSpans(Config config) {
     if (!codeOriginEnabled.compareAndSet(false, true)) {
       return;
     }
     LOGGER.debug("Starting Code Origin for spans");
-    Config config = Config.get();
     commonInit(config);
     initClassNameFilter();
     DebuggerContext.initClassNameFilter(classNameFilter);
@@ -263,7 +260,7 @@ public static void stopCodeOriginForSpans() {
     DebuggerContext.initCodeOrigin(null);
   }
 
-  public static void startDistributedDebugger() {
+  public static void startDistributedDebugger(Config config) {
     if (!distributedDebuggerEnabled.compareAndSet(false, true)) {
       return;
     }
@@ -441,6 +438,26 @@ public static JsonSnapshotSerializer getSnapshotSerializer() {
     return DebuggerAgent.snapshotSerializer;
   }
 
+  // only used for tests
+  static void reset() {
+    instrumentation = null;
+    sharedCommunicationObjects = null;
+    configurationPoller = null;
+    sink = null;
+    agentVersion = null;
+    snapshotSerializer = null;
+    classNameFilter = null;
+    symDBEnablement = null;
+    configurationUpdater = null;
+    exceptionDebugger = null;
+    commonInitDone.set(false);
+    dynamicInstrumentationEnabled.set(false);
+    exceptionReplayEnabled.set(false);
+    codeOriginEnabled.set(false);
+    distributedDebuggerEnabled.set(false);
+    classesToRetransformFinder = null;
+  }
+
   private static class ShutdownHook extends Thread {
 
     private final WeakReference<ConfigurationPoller> pollerRef;

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

@@ -697,7 +697,7 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(
       // Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
       // and therefore need to be instrumented once
       // note: exception probes are log probes and are handled the same way
-      if (!Config.get().isDistributedDebuggerEnabled() && definition instanceof TriggerProbe) {
+      if (!config.isDistributedDebuggerEnabled() && definition instanceof TriggerProbe) {
         LOGGER.debug(
             "The distributed debugger feature is disabled. Trigger probes will not be installed.");
       } else if (isCapturedContextProbe(definition)) {

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DefaultDebuggerConfigUpdater.java

@@ -5,31 +5,42 @@
 import datadog.trace.api.config.TraceInstrumentationConfig;
 import datadog.trace.api.debugger.DebuggerConfigUpdate;
 import datadog.trace.api.debugger.DebuggerConfigUpdater;
+import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 class DefaultDebuggerConfigUpdater implements DebuggerConfigUpdater {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DefaultDebuggerConfigUpdater.class);
 
+  private final Config config;
+
+  public DefaultDebuggerConfigUpdater(Config config) {
+    this.config = config;
+  }
+
   @Override
   public void updateConfig(DebuggerConfigUpdate update) {
     startOrStopFeature(
+        config,
         DebuggerConfig.DYNAMIC_INSTRUMENTATION_ENABLED,
         update.getDynamicInstrumentationEnabled(),
         DebuggerAgent::startDynamicInstrumentation,
         DebuggerAgent::stopDynamicInstrumentation);
     startOrStopFeature(
+        config,
         DebuggerConfig.EXCEPTION_REPLAY_ENABLED,
         update.getExceptionReplayEnabled(),
         DebuggerAgent::startExceptionReplay,
         DebuggerAgent::stopExceptionReplay);
     startOrStopFeature(
+        config,
         TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED,
         update.getCodeOriginEnabled(),
         DebuggerAgent::startCodeOriginForSpans,
         DebuggerAgent::stopCodeOriginForSpans);
     startOrStopFeature(
+        config,
         DebuggerConfig.DISTRIBUTED_DEBUGGER_ENABLED,
         update.getDistributedDebuggerEnabled(),
         DebuggerAgent::startDistributedDebugger,
@@ -62,14 +73,18 @@ private static boolean isExplicitlyDisabled(String booleanKey) {
   }
 
   private static void startOrStopFeature(
-      String booleanKey, Boolean currentStatus, Runnable start, Runnable stop) {
+      Config config,
+      String booleanKey,
+      Boolean currentStatus,
+      Consumer<Config> start,
+      Runnable stop) {
     if (isExplicitlyDisabled(booleanKey)) {
       LOGGER.debug("Feature {} is explicitly disabled", booleanKey);
       return;
     }
     if (currentStatus != null) {
       if (currentStatus) {
-        start.run();
+        start.accept(config);
       } else {
         stop.run();
       }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerAgentTest.java

@@ -2,7 +2,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assumptions.assumeTrue;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.condition.JRE.JAVA_11;
 import static org.junit.jupiter.api.condition.JRE.JAVA_8;
 import static org.mockito.ArgumentMatchers.any;
@@ -12,7 +12,6 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
-import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.util.RemoteConfigHelper;
 import datadog.common.container.ContainerInfo;
@@ -22,7 +21,6 @@
 import datadog.trace.api.git.GitInfoProvider;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.test.util.ControllableEnvironmentVariables;
-import datadog.trace.test.util.Flaky;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
@@ -51,10 +49,10 @@
 public class DebuggerAgentTest {
 
   public static final String URL_PATH = "/foo";
+  @Mock Config config;
   @Mock Instrumentation inst;
   final MockWebServer server = new MockWebServer();
   HttpUrl url;
-
   private ControllableEnvironmentVariables env = ControllableEnvironmentVariables.setup();
 
   private static void setFieldInContainerInfo(
@@ -76,6 +74,7 @@ public void setUp() {
 
   @AfterEach
   public void tearDown() throws IOException {
+    DebuggerAgent.reset();
     try {
       server.shutdown();
     } catch (final IOException e) {
@@ -86,64 +85,66 @@ public void tearDown() throws IOException {
   @Test
   @EnabledOnJre({JAVA_8, JAVA_11})
   public void runDisabled() {
-    setFieldInConfig(Config.get(), "dynamicInstrumentationEnabled", false);
-    DebuggerAgent.run(inst, new SharedCommunicationObjects());
+    when(config.isDynamicInstrumentationEnabled()).thenReturn(false);
+    DebuggerAgent.run(config, inst, new SharedCommunicationObjects());
     verify(inst, never()).addTransformer(any(), eq(true));
   }
 
-  @Flaky
   @Test
   @EnabledOnJre({JAVA_8, JAVA_11})
   public void runEnabledWithDatadogAgent() throws InterruptedException, IOException {
-    MockWebServer datadogAgentServer = new MockWebServer();
-    HttpUrl datadogAgentUrl = datadogAgentServer.url(URL_PATH);
-    setFieldInConfig(Config.get(), "dynamicInstrumentationEnabled", true);
-    setFieldInConfig(Config.get(), "remoteConfigEnabled", true);
-    setFieldInConfig(Config.get(), "dynamicInstrumentationSnapshotUrl", datadogAgentUrl.toString());
-    setFieldInConfig(Config.get(), "agentUrl", datadogAgentUrl.toString());
-    setFieldInConfig(Config.get(), "agentHost", "localhost");
-    setFieldInConfig(Config.get(), "agentPort", datadogAgentServer.getPort());
-    setFieldInConfig(Config.get(), "dynamicInstrumentationMaxPayloadSize", 4096L);
+    when(config.isDynamicInstrumentationEnabled()).thenReturn(true);
+    when(config.isRemoteConfigEnabled()).thenReturn(true);
+    when(config.getAgentUrl()).thenReturn(url.toString());
+    when(config.getDynamicInstrumentationUploadBatchSize()).thenReturn(100);
+    when(config.getRemoteConfigTargetsKeyId())
+        .thenReturn(Config.get().getRemoteConfigTargetsKeyId());
+    when(config.getRemoteConfigTargetsKey()).thenReturn(Config.get().getRemoteConfigTargetsKey());
+    when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
+    when(config.getRemoteConfigMaxPayloadSizeBytes())
+        .thenReturn(Config.get().getRemoteConfigMaxPayloadSizeBytes());
+    assertTrue(config.isDynamicInstrumentationEnabled());
     setFieldInContainerInfo(ContainerInfo.get(), "containerId", "");
     String infoContent =
         "{\"endpoints\": [\"v0.4/traces\", \"debugger/v1/input\", \"debugger/v1/diagnostics\", \"v0.7/config\"] }";
-    datadogAgentServer.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
-    datadogAgentServer.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
+    server.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
+    server.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
     try (BufferedReader reader =
         new BufferedReader(
             new InputStreamReader(
                 DebuggerAgentTest.class.getResourceAsStream("/test_probe.json")))) {
       String content = reader.lines().collect(Collectors.joining("\n"));
       String rcContent = RemoteConfigHelper.encode(content, "petclinic");
-      datadogAgentServer.enqueue(new MockResponse().setResponseCode(200).setBody(rcContent));
+      server.enqueue(new MockResponse().setResponseCode(200).setBody(rcContent));
     } catch (IOException e) {
       e.printStackTrace();
     }
     SharedCommunicationObjects sharedCommunicationObjects = new SharedCommunicationObjects();
-    DebuggerAgent.run(inst, sharedCommunicationObjects);
+    DebuggerAgent.run(config, inst, sharedCommunicationObjects);
     ConfigurationPoller configurationPoller =
-        (ConfigurationPoller) sharedCommunicationObjects.configurationPoller(Config.get());
+        sharedCommunicationObjects.configurationPoller(config);
     configurationPoller.start();
     RecordedRequest request;
     do {
-      request = datadogAgentServer.takeRequest(5, TimeUnit.SECONDS);
+      request = server.takeRequest(5, TimeUnit.SECONDS);
       assertNotNull(request);
     } while ("/info".equals(request.getPath()));
     assertEquals("/v0.7/config", request.getPath());
     DebuggerAgent.stop();
-    datadogAgentServer.shutdown();
   }
 
   @Test
   @EnabledOnJre({JAVA_8, JAVA_11})
   public void runEnabledWithUnsupportedDatadogAgent() throws InterruptedException {
-    setFieldInConfig(Config.get(), "dynamicInstrumentationEnabled", true);
-    setFieldInConfig(Config.get(), "dynamicInstrumentationSnapshotUrl", url.toString());
-    setFieldInConfig(Config.get(), "agentUrl", url.toString());
-    setFieldInConfig(Config.get(), "dynamicInstrumentationMaxPayloadSize", 1024L);
+    when(config.isDynamicInstrumentationEnabled()).thenReturn(true);
+    when(config.getAgentUrl()).thenReturn(url.toString());
+    when(config.getFinalDebuggerSnapshotUrl())
+        .thenReturn("http://localhost:8126/debugger/v1/input");
+    when(config.getDynamicInstrumentationUploadBatchSize()).thenReturn(100);
+    when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
     String infoContent = "{\"endpoints\": [\"v0.4/traces\"]}";
     server.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
-    DebuggerAgent.run(inst, new SharedCommunicationObjects());
+    DebuggerAgent.run(config, inst, new SharedCommunicationObjects());
     verify(inst, never()).addTransformer(any(), eq(true));
   }
 
@@ -152,18 +153,18 @@ public void runEnabledWithUnsupportedDatadogAgent() throws InterruptedException
   public void readFromFile() throws URISyntaxException {
     URL res = getClass().getClassLoader().getResource("test_probe_file.json");
     String probeDefinitionPath = Paths.get(res.toURI()).toFile().getAbsolutePath();
-    setFieldInConfig(Config.get(), "serviceName", "petclinic");
-    setFieldInConfig(Config.get(), "dynamicInstrumentationEnabled", true);
-    setFieldInConfig(Config.get(), "dynamicInstrumentationSnapshotUrl", url.toString());
-    setFieldInConfig(Config.get(), "agentUrl", url.toString());
-    setFieldInConfig(Config.get(), "dynamicInstrumentationMaxPayloadSize", 4096L);
-    setFieldInConfig(Config.get(), "dynamicInstrumentationProbeFile", probeDefinitionPath);
+    when(config.getServiceName()).thenReturn("petclinic");
+    when(config.isDynamicInstrumentationEnabled()).thenReturn(true);
+    when(config.getAgentUrl()).thenReturn(url.toString());
+    when(config.getDynamicInstrumentationMaxPayloadSize()).thenReturn(4096L);
+    when(config.getDynamicInstrumentationProbeFile()).thenReturn(probeDefinitionPath);
+    when(config.getDynamicInstrumentationUploadBatchSize()).thenReturn(100);
+    when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
     String infoContent =
-        "{\"endpoints\": [\"v0.4/traces\", \"debugger/v1/input\", \"v0.7/config\"] }";
+        "{\"endpoints\": [\"v0.4/traces\", \"debugger/v1/input\", \"debugger/v1/diagnostics\", \"v0.7/config\"] }";
     server.enqueue(new MockResponse().setResponseCode(200).setBody(infoContent));
-    // sometimes this test fails because getAllLoadedClasses returns null
-    assumeTrue(inst.getAllLoadedClasses() != null);
-    DebuggerAgent.run(inst, new SharedCommunicationObjects());
+    when(inst.getAllLoadedClasses()).thenReturn(new Class[0]);
+    DebuggerAgent.run(config, inst, new SharedCommunicationObjects());
     verify(inst, atLeastOnce()).addTransformer(any(), eq(true));
   }
 

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DefaultDebuggerConfigUpdaterTest.java

@@ -7,6 +7,7 @@
 
 import datadog.communication.ddagent.DDAgentFeaturesDiscovery;
 import datadog.communication.ddagent.SharedCommunicationObjects;
+import datadog.trace.api.Config;
 import datadog.trace.api.debugger.DebuggerConfigUpdate;
 import java.lang.instrument.Instrumentation;
 import org.junit.jupiter.api.Test;
@@ -17,8 +18,9 @@ class DefaultDebuggerConfigUpdaterTest {
   public void enableDisable() {
     SharedCommunicationObjects sco = mock(SharedCommunicationObjects.class);
     when(sco.featuresDiscovery(any())).thenReturn(mock(DDAgentFeaturesDiscovery.class));
-    DebuggerAgent.run(mock(Instrumentation.class), sco);
-    DefaultDebuggerConfigUpdater productConfigUpdater = new DefaultDebuggerConfigUpdater();
+    DebuggerAgent.run(Config.get(), mock(Instrumentation.class), sco);
+    DefaultDebuggerConfigUpdater productConfigUpdater =
+        new DefaultDebuggerConfigUpdater(Config.get());
     productConfigUpdater.updateConfig(new DebuggerConfigUpdate());
     productConfigUpdater.updateConfig(new DebuggerConfigUpdate(true, true, true, true));
     assertTrue(productConfigUpdater.isDynamicInstrumentationEnabled());