Enforce continuation diagnostics on smoke tests

dd-java-agent/instrumentation-testing/build.gradle

@@ -1,7 +1,50 @@
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
+
 plugins {
   id 'dd-trace-java.module.testing-support'
 }
 
+sourceSets {
+  smokeAgent {
+    compileClasspath += sourceSets.main.output + sourceSets.main.compileClasspath
+  }
+}
+
+// Reuse the production agent's logging API from the tracer's isolated loader.
+def scopeDiagnosticsPayloadJar = tasks.register('scopeDiagnosticsPayloadJar', ShadowJar) {
+  archiveClassifier = 'scope-diagnostics-payload'
+  configurations = []
+  from(sourceSets.main.output) {
+    include 'datadog/trace/agent/test/scopediag/**'
+    exclude '**/ScopeDiagnosticsExtension*', '**/ScopeDiagnosticsSpock*'
+  }
+  from(sourceSets.smokeAgent.output) {
+    include 'datadog/trace/agent/test/scopediag/**'
+  }
+  from(project(':dd-java-agent:testing').sourceSets.main.output) {
+    include 'datadog/trace/agent/test/scopediag/ScopeDiagnosticsCheck.class'
+  }
+  relocate 'org.slf4j', 'datadog.slf4j'
+}
+
+tasks.register('scopeDiagnosticsAgentJar', Jar) {
+  archiveClassifier = 'scope-diagnostics-agent'
+  manifest {
+    attributes 'Premain-Class': 'datadog.trace.agent.test.smoke.ScopeDiagnosticsAgent',
+    'Can-Retransform-Classes': 'true', 'Can-Redefine-Classes': 'true'
+  }
+  from(sourceSets.smokeAgent.output) {
+    include 'datadog/trace/agent/test/smoke/**'
+  }
+  into('payload') {
+    from(zipTree(scopeDiagnosticsPayloadJar.flatMap { it.archiveFile })) {
+      include '**/*.class'
+    }
+    rename '(.*)\\.class', '$1.classdata'
+  }
+  dependsOn scopeDiagnosticsPayloadJar
+}
+
 dependencies {
   api libs.bytebuddy
   api libs.bytebuddyagent

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/scopediag/ScopeContinuationTransformer.java

@@ -1,16 +1,25 @@
 package datadog.trace.agent.test.scopediag;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.is;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import java.lang.instrument.Instrumentation;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * Installs test-only advice with a separate {@link AgentBuilder} because the tracer ignores its own
@@ -55,115 +64,164 @@ final class ScopeContinuationTransformer {
   private ScopeContinuationTransformer() {}
 
   static synchronized void install() {
+    if (transformer == null) {
+      install(
+          ByteBuddyAgent.getInstrumentation(),
+          "datadog.trace.core",
+          ClassFileLocator.ForClassLoader.of(ScopeContinuationTransformer.class.getClassLoader()));
+    }
+  }
+
+  static synchronized void install(
+      Instrumentation instrumentation, String corePackage, ClassFileLocator locator) {
     if (transformer != null) {
       return;
     }
+    ClassLoader loader = ScopeContinuationTransformer.class.getClassLoader();
     try {
       // Related core types can otherwise load these targets reentrantly while they are transformed.
-      ClassLoader loader = ScopeContinuationTransformer.class.getClassLoader();
-      Class.forName("datadog.trace.core.scopemanager.ScopeContinuation", false, loader);
-      Class.forName("datadog.trace.core.scopemanager.ScopeStack", false, loader);
+      for (String target :
+          new String[] {
+            ".scopemanager.ScopeContinuation",
+            ".scopemanager.ScopeStack",
+            ".scopemanager.ContinuableScope",
+            ".scopemanager.ContinuableScopeManager",
+            ".PendingTrace"
+          }) {
+        Class.forName(corePackage + target, false, loader);
+      }
     } catch (ClassNotFoundException missingCoreTracer) {
       throw new IllegalStateException(
           "Scope continuation diagnostics require dd-trace-core", missingCoreTracer);
     }
-    Instrumentation instrumentation = ByteBuddyAgent.getInstrumentation();
+    AtomicReference<Throwable> failure = new AtomicReference<>();
+    Set<String> transformed = Collections.synchronizedSet(new HashSet<String>());
     transformer =
         new AgentBuilder.Default()
+            .with(
+                new AgentBuilder.Listener.Adapter() {
+                  @Override
+                  public void onTransformation(
+                      TypeDescription type,
+                      ClassLoader loader,
+                      JavaModule module,
+                      boolean loaded,
+                      DynamicType dynamicType) {
+                    transformed.add(type.getName());
+                  }
+
+                  @Override
+                  public void onError(
+                      String typeName,
+                      ClassLoader loader,
+                      JavaModule module,
+                      boolean loaded,
+                      Throwable error) {
+                    failure.compareAndSet(
+                        null, new IllegalStateException("Cannot instrument " + typeName, error));
+                  }
+                })
             .disableClassFormatChanges()
             .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
             .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
-            .type(named("datadog.trace.core.scopemanager.ScopeContinuation"))
+            .type(named(corePackage + ".scopemanager.ScopeContinuation"), is(loader))
             .transform(
                 (builder, type, classLoader, module, pd) ->
                     builder
                         .visit(
-                            Advice.to(ContinuationAdvice.Register.class)
+                            Advice.to(ContinuationAdvice.Register.class, locator)
                                 .on(
                                     isMethod()
                                         .and(named("register"))
                                         .and(takesArguments(0))
                                         .and(
                                             returns(
                                                 named(
-                                                    "datadog.trace.core.scopemanager.ScopeContinuation")))))
+                                                    corePackage
+                                                        + ".scopemanager.ScopeContinuation")))))
                         .visit(
-                            Advice.to(ContinuationAdvice.Activate.class)
+                            Advice.to(ContinuationAdvice.Activate.class, locator)
                                 .on(
                                     isMethod()
                                         .and(named("resume"))
                                         .and(takesArguments(0))
                                         .and(returns(named("datadog.context.ContextScope")))))
                         .visit(
-                            Advice.to(ContinuationAdvice.Cancel.class)
+                            Advice.to(ContinuationAdvice.Cancel.class, locator)
                                 .on(
                                     isMethod()
                                         .and(
                                             named("release")
                                                 .or(named("cancelFromContinuedScopeClose")))
                                         .and(takesArguments(0))
                                         .and(returns(void.class)))))
-            .type(named("datadog.trace.core.PendingTrace"))
+            .type(named(corePackage + ".PendingTrace"), is(loader))
             .transform(
                 (builder, type, classLoader, module, pd) ->
                     builder.visit(
-                        Advice.to(PendingTraceAdvice.Write.class)
+                        Advice.to(PendingTraceAdvice.Write.class, locator)
                             .on(
                                 isMethod()
                                     .and(named("write"))
                                     .and(takesArguments(boolean.class))
                                     .and(returns(int.class)))))
-            .type(named("datadog.trace.core.scopemanager.ContinuableScope"))
+            .type(named(corePackage + ".scopemanager.ContinuableScope"), is(loader))
             .transform(
                 (builder, type, classLoader, module, pd) ->
                     builder
                         .visit(
-                            Advice.to(ContinuableScopeAdvice.OnProperClose.class)
+                            Advice.to(ContinuableScopeAdvice.OnProperClose.class, locator)
                                 .on(
                                     isMethod()
                                         .and(named("onProperClose"))
                                         .and(takesArguments(0))
                                         .and(returns(void.class))))
                         .visit(
-                            Advice.to(ContinuableScopeAdvice.Close.class)
+                            Advice.to(ContinuableScopeAdvice.Close.class, locator)
                                 .on(
                                     isMethod()
                                         .and(named("close"))
                                         .and(takesArguments(0))
                                         .and(returns(void.class)))))
-            .type(named("datadog.trace.core.scopemanager.ScopeStack"))
+            .type(named(corePackage + ".scopemanager.ScopeStack"), is(loader))
             .transform(
                 (builder, type, classLoader, module, pd) ->
                     builder.visit(
-                        Advice.to(ScopeStackAdvice.Push.class)
+                        Advice.to(ScopeStackAdvice.Push.class, locator)
                             .on(
                                 isMethod()
                                     .and(named("push"))
                                     .and(takesArguments(1))
                                     .and(
                                         takesArgument(
                                             0,
-                                            named(
-                                                "datadog.trace.core.scopemanager.ContinuableScope")))
+                                            named(corePackage + ".scopemanager.ContinuableScope")))
                                     .and(returns(void.class)))))
-            .type(named("datadog.trace.core.scopemanager.ContinuableScopeManager"))
+            .type(named(corePackage + ".scopemanager.ContinuableScopeManager"), is(loader))
             .transform(
                 (builder, type, classLoader, module, pd) ->
                     builder.visit(
-                        Advice.to(ContinuableScopeManagerAdvice.ScheduleRootIterationCleanup.class)
+                        Advice.to(
+                                ContinuableScopeManagerAdvice.ScheduleRootIterationCleanup.class,
+                                locator)
                             .on(
                                 isMethod()
                                     .and(named("scheduleRootIterationScopeCleanup"))
                                     .and(takesArguments(2))
                                     .and(
                                         takesArgument(
-                                            0, named("datadog.trace.core.scopemanager.ScopeStack")))
+                                            0, named(corePackage + ".scopemanager.ScopeStack")))
                                     .and(
                                         takesArgument(
                                             1,
-                                            named(
-                                                "datadog.trace.core.scopemanager.ContinuableScope")))
+                                            named(corePackage + ".scopemanager.ContinuableScope")))
                                     .and(returns(void.class)))))
             .installOn(instrumentation);
+    if (failure.get() != null || transformed.size() != 5) {
+      transformer.reset(instrumentation, AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);
+      transformer = null;
+      throw new IllegalStateException(
+          "Scope diagnostics installation failed; transformed " + transformed, failure.get());
+    }
   }
 }

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/scopediag/StackFilter.java

@@ -8,6 +8,7 @@ final class StackFilter {
   private static final String[] DROP_PREFIXES = {
     "datadog.trace.agent.test.scopediag.",
     "datadog.trace.core.",
+    "datadog.trace.agent.core.",
     "datadog.trace.bootstrap.instrumentation.java.concurrent.",
     "datadog.trace.bootstrap.instrumentation.api.",
     "java.util.concurrent.Executors$",

dd-java-agent/instrumentation-testing/src/smokeAgent/java/datadog/trace/agent/test/scopediag/SmokeDiagnosticsBridge.java

@@ -0,0 +1,101 @@
+package datadog.trace.agent.test.scopediag;
+
+import java.lang.instrument.Instrumentation;
+import java.util.EnumSet;
+import java.util.Map;
+import java.util.Properties;
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+/** Runs in the tracer's loader so all recorded objects retain their actual type identity. */
+public final class SmokeDiagnosticsBridge {
+  private static Properties completed;
+  private static boolean recording;
+  private static String exclusionReason = "";
+  private static EnumSet<ScopeDiagnosticsCheck> enabledChecks =
+      EnumSet.allOf(ScopeDiagnosticsCheck.class);
+
+  private SmokeDiagnosticsBridge() {}
+
+  public static void install(
+      Instrumentation instrumentation,
+      Map<String, byte[]> classes,
+      String disabledChecks,
+      String reason) {
+    EnumSet<ScopeDiagnosticsCheck> checks = EnumSet.allOf(ScopeDiagnosticsCheck.class);
+    if (!disabledChecks.isEmpty()) {
+      if (reason.trim().isEmpty()) {
+        throw new IllegalArgumentException("Disabling scope diagnostic checks requires a reason");
+      }
+      int start = 0;
+      while (true) {
+        int comma = disabledChecks.indexOf(',', start);
+        String check =
+            comma < 0 ? disabledChecks.substring(start) : disabledChecks.substring(start, comma);
+        checks.remove(ScopeDiagnosticsCheck.valueOf(check.trim()));
+        if (comma < 0) {
+          break;
+        }
+        start = comma + 1;
+      }
+    }
+    enabledChecks = checks;
+    exclusionReason = reason;
+    String corePackage = "datadog.trace.agent.core";
+    try {
+      Class.forName(
+          corePackage + ".scopemanager.ScopeContinuation",
+          false,
+          SmokeDiagnosticsBridge.class.getClassLoader());
+    } catch (ClassNotFoundException unrelocatedAgent) {
+      corePackage = "datadog.trace.core";
+    }
+    ScopeContinuationTransformer.install(
+        instrumentation, corePackage, new ClassFileLocator.Simple(classes));
+  }
+
+  public static synchronized void start() {
+    if (recording) {
+      throw new IllegalStateException("Scope diagnostics are already recording");
+    }
+    completed = null;
+    ScopeDiagnostics.startRecording();
+    recording = true;
+  }
+
+  public static synchronized Properties finish() {
+    if (completed != null) {
+      return completed;
+    }
+    if (!recording) {
+      throw new IllegalStateException("Scope diagnostics have not started recording");
+    }
+    ScopeDiagnostics.awaitQuiescence();
+    ScopeDiagnostics.stop();
+    recording = false;
+    ScopeDiagnosticsReport report = ScopeDiagnostics.report();
+    Properties result = new Properties();
+    result.setProperty("status", report.hasViolations(enabledChecks) ? "problems" : "ok");
+    result.setProperty("detail", report.renderSummary(enabledChecks));
+    result.setProperty("timeline", report.renderTimeline());
+    result.setProperty("enabledChecks", enabledChecks.toString());
+    result.setProperty("reason", exclusionReason);
+    long events = 0;
+    for (ContinuationRecord record : report.records()) {
+      events +=
+          (record.capture() == null ? 0 : 1)
+              + record.resumes().size()
+              + record.failedActivations().size()
+              + (record.terminal() == null ? 0 : 1)
+              + record.extraTerminals().size();
+    }
+    for (ScopeRecord record : report.scopeRecords()) {
+      events +=
+          (record.open() == null ? 0 : 1)
+              + (record.close() == null ? 0 : 1)
+              + record.wrongThreadCloses().size();
+    }
+    result.setProperty("eventCount", Long.toString(events));
+    completed = result;
+    return result;
+  }
+}