🪞 9979 - Instrumentation test error visibility

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

@@ -24,7 +24,10 @@ private InstrumentationContext() {}
   public static <K, C> ContextStore<K, C> get(
       final Class<K> keyClass, final Class<C> contextClass) {
     throw new RuntimeException(
-        "Calls to this method will be rewritten by Instrumentation Context Provider (e.g. FieldBackedProvider)");
+        "Calls to this method will be rewritten by Instrumentation Context Provider (e.g. FieldBackedProvider)."
+            + " If you get this exception, this method has not been rewritten."
+            + " Ensure instrumentation class has a contextStore method and the call to InstrumentationContext.get happens directly in an instrumentation Advice class."
+            + " See how_instrumentations_work.md for details.");
   }
 
   /**
@@ -35,6 +38,9 @@ public static <K, C> ContextStore<K, C> get(
    */
   public static <K, C> ContextStore<K, C> get(final String keyClass, final String contextClass) {
     throw new RuntimeException(
-        "Calls to this method will be rewritten by Instrumentation Context Provider (e.g. FieldBackedProvider)");
+        "Calls to this method will be rewritten by Instrumentation Context Provider (e.g. FieldBackedProvider)."
+            + " If you get this exception, this method has not been rewritten."
+            + " Ensure instrumentation class has a contextStore method and the call to InstrumentationContext.get happens directly in an instrumentation Advice class."
+            + " See how_instrumentations_work.md for details.");
   }
 }

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

@@ -1,21 +1,36 @@
 package datadog.trace.bootstrap;
 
-import java.util.concurrent.atomic.AtomicLong;
+import java.io.PrintWriter;
+import java.io.StringWriter;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 public class InstrumentationErrors {
-  private static final AtomicLong COUNTER = new AtomicLong();
+  private static final List<String> ERRORS = new CopyOnWriteArrayList<>();
+  private static volatile boolean recordErrors = false;
 
-  public static long getErrorCount() {
-    return COUNTER.get();
+  /**
+   * Record an error from {@link datadog.trace.agent.tooling.bytebuddy.ExceptionHandlers} for test
+   * visibility.
+   */
+  @SuppressWarnings("unused")
+  public static void recordError(final Throwable throwable) {
+    if (recordErrors) {
+      StringWriter stackTrace = new StringWriter();
+      throwable.printStackTrace(new PrintWriter(stackTrace));
+      ERRORS.add(stackTrace.toString());
+    }
   }
 
-  @SuppressWarnings("unused")
-  public static void incrementErrorCount() {
-    COUNTER.incrementAndGet();
+  // Visible for testing
+  public static void enableRecordingAndReset() {
+    recordErrors = true;
+    ERRORS.clear();
   }
 
   // Visible for testing
-  public static void resetErrorCount() {
-    COUNTER.set(0);
+  public static List<String> getErrors() {
+    return Collections.unmodifiableList(ERRORS);
   }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/ExceptionHandlers.java

@@ -39,24 +39,14 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
               final boolean exitOnFailure = InstrumenterConfig.get().isInternalExitOnFailure();
               final String logMethod = exitOnFailure ? "error" : "debug";
 
-              // Writes the following bytecode if exitOnFailure is false:
+              // Writes the following bytecode (note that two statements are conditionally written):
               //
-              // BlockingExceptionHandler.rethrowIfBlockingException(t);
+              // BlockingExceptionHandler.rethrowIfBlockingException(t); // when appSecEnabled=true
               // try {
-              //   InstrumentationErrors.incrementErrorCount();
-              //   org.slf4j.LoggerFactory.getLogger((Class)ExceptionLogger.class)
-              //     .debug("Failed to handle exception in instrumentation for ...", t);
-              // } catch (Throwable t2) {
-              // }
-              //
-              // And the following bytecode if exitOnFailure is true:
-              //
-              // BlockingExceptionHandler.rethrowIfBlockingException(t);
-              // try {
-              //   InstrumentationErrors.incrementErrorCount();
+              //   InstrumentationErrors.recordError(t);
               //   org.slf4j.LoggerFactory.getLogger((Class)ExceptionLogger.class)
               //     .error("Failed to handle exception in instrumentation for ...", t);
-              //   System.exit(1);
+              //   System.exit(1); // when exitOnFailure=true
               // } catch (Throwable t2) {
               // }
               //
@@ -81,12 +71,14 @@ public Size apply(final MethodVisitor mv, final Implementation.Context context)
 
               mv.visitTryCatchBlock(logStart, logEnd, eatException, "java/lang/Throwable");
               mv.visitLabel(logStart);
-              // invoke incrementAndGet on our exception counter
+              mv.visitInsn(Opcodes.DUP); // stack: (top) throwable,throwable
+              // invoke recordError on our exception tracker
               mv.visitMethodInsn(
                   Opcodes.INVOKESTATIC,
                   "datadog/trace/bootstrap/InstrumentationErrors",
-                  "incrementErrorCount",
-                  "()V");
+                  "recordError",
+                  "(Ljava/lang/Throwable;)V");
+
               // stack: (top) throwable
               mv.visitLdcInsn(Type.getType("L" + HANDLER_NAME + ";"));
               mv.visitMethodInsn(

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleCheck.java

@@ -3,14 +3,19 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.InstrumenterState;
 import datadog.trace.agent.tooling.Utils;
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class MuzzleCheck implements ElementMatcher<ClassLoader> {
   private static final Logger log = LoggerFactory.getLogger(MuzzleCheck.class);
 
+  private static final List<String> ERRORS = new CopyOnWriteArrayList<>();
+  private static volatile boolean recordErrors = false;
+
   private final int instrumentationId;
   private final String instrumentationClass;
   private final ReferenceProvider runtimeMuzzleReferences;
@@ -33,9 +38,13 @@ public boolean matches(ClassLoader classLoader) {
       InstrumenterState.applyInstrumentation(classLoader, instrumentationId);
     } else {
       InstrumenterState.blockInstrumentation(classLoader, instrumentationId);
-      if (log.isDebugEnabled()) {
+      if (recordErrors || log.isDebugEnabled()) {
         final List<Reference.Mismatch> mismatches =
             muzzle.getMismatchedReferenceSources(classLoader);
+        if (recordErrors) {
+          recordMuzzleMismatch(
+              InstrumenterState.describe(instrumentationId), classLoader, mismatches);
+        }
         log.debug(
             "Muzzled - {} instrumentation.target.classloader={}",
             InstrumenterState.describe(instrumentationId),
@@ -61,4 +70,38 @@ private ReferenceMatcher muzzle() {
     }
     return muzzle;
   }
+
+  /**
+   * Record a muzzle mismatch error for test visibility.
+   *
+   * @param instrumentationDescription the description of the instrumentation that was blocked
+   * @param classLoader the classloader where the instrumentation was blocked
+   * @param mismatches the list of mismatch details
+   */
+  private static void recordMuzzleMismatch(
+      String instrumentationDescription,
+      ClassLoader classLoader,
+      List<Reference.Mismatch> mismatches) {
+    StringBuilder sb = new StringBuilder();
+    sb.append("Muzzled - ")
+        .append(instrumentationDescription)
+        .append(" instrumentation.target.classloader=")
+        .append(classLoader)
+        .append("\n");
+    for (Reference.Mismatch mismatch : mismatches) {
+      sb.append("  Mismatch: ").append(mismatch).append("\n");
+    }
+    ERRORS.add(sb.toString());
+  }
+
+  // Visible for testing
+  public static void enableRecordingAndReset() {
+    recordErrors = true;
+    ERRORS.clear();
+  }
+
+  // Visible for testing
+  public static List<String> getErrors() {
+    return Collections.unmodifiableList(ERRORS);
+  }
 }

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -25,6 +25,7 @@ import datadog.trace.agent.tooling.InstrumenterModule
 import datadog.trace.agent.tooling.TracerInstaller
 import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers
 import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores
+import datadog.trace.agent.tooling.muzzle.MuzzleCheck
 import datadog.trace.api.Config
 import datadog.trace.api.IdGenerationStrategy
 import datadog.trace.api.ProcessTags
@@ -73,9 +74,9 @@ import java.lang.instrument.ClassFileTransformer
 import java.lang.instrument.Instrumentation
 import java.nio.ByteBuffer
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.TimeoutException
-import java.util.concurrent.atomic.AtomicInteger
 
 import static datadog.communication.http.OkHttpUtils.buildHttpClient
 import static datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_HOST
@@ -165,7 +166,7 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
 
   @SuppressWarnings('PropertyName')
   @Shared
-  AtomicInteger INSTRUMENTATION_ERROR_COUNT = new AtomicInteger(0)
+  List<String> INSTRUMENTATION_ERRORS = new CopyOnWriteArrayList<String>()
 
   // don't use mocks because it will break too many exhaustive interaction-verifying tests
   @SuppressWarnings('PropertyName')
@@ -463,7 +464,8 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
     if (forceAppSecActive) {
       ActiveSubsystems.APPSEC_ACTIVE = true
     }
-    InstrumentationErrors.resetErrorCount()
+    InstrumentationErrors.enableRecordingAndReset()
+    MuzzleCheck.enableRecordingAndReset()
     ProcessTags.reset()
   }
 
@@ -505,7 +507,10 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
       spanFinishLocations.clear()
       originalToTrackingSpan.clear()
     }
-    assert InstrumentationErrors.errorCount == 0
+    def instrumentationErrorCount = InstrumentationErrors.getErrors().size()
+    assert instrumentationErrorCount == 0, "${instrumentationErrorCount} instrumentation errors were seen:\n${InstrumentationErrors.getErrors().join("\n---\n")}"
+    def muzzleErrorCount = MuzzleCheck.getErrors().size()
+    assert muzzleErrorCount == 0, "${muzzleErrorCount} muzzle errors were seen:\n${MuzzleCheck.getErrors().join("\n---\n")}"
   }
 
   private void doCheckRepeatedFinish() {
@@ -547,7 +552,8 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
     cleanupAfterAgent()
 
     // All cleanup should happen before these assertion.  If not, a failing assertion may prevent cleanup
-    assert INSTRUMENTATION_ERROR_COUNT.get() == 0: INSTRUMENTATION_ERROR_COUNT.get() + " Instrumentation errors during test"
+    def instrumentationErrors = INSTRUMENTATION_ERRORS.size()
+    assert instrumentationErrors == 0: "${instrumentationErrors} Instrumentation errors during test:\n${INSTRUMENTATION_ERRORS.join("\n---\n")}"
 
     assert TRANSFORMED_CLASSES_TYPES.findAll {
       GlobalIgnores.isAdditionallyIgnored(it.getActualName())
@@ -641,7 +647,9 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
 
     println "Unexpected instrumentation error when instrumenting ${typeName} on ${classLoader}"
     throwable.printStackTrace()
-    INSTRUMENTATION_ERROR_COUNT.incrementAndGet()
+    def stackTrace = new StringWriter()
+    throwable.printStackTrace(new PrintWriter(stackTrace))
+    INSTRUMENTATION_ERRORS.add(stackTrace.toString())
   }
 
   @Override

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/ListWriterAssert.groovy

@@ -2,6 +2,8 @@ package datadog.trace.agent.test.asserts
 
 import static TraceAssert.assertTrace
 
+import datadog.trace.agent.tooling.muzzle.MuzzleCheck
+import datadog.trace.bootstrap.InstrumentationErrors
 import datadog.trace.common.writer.ListWriter
 import datadog.trace.core.DDSpan
 import groovy.transform.stc.ClosureParams
@@ -47,7 +49,15 @@ class ListWriterAssert {
     @ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.ListWriterAssert'])
     @DelegatesTo(value = ListWriterAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     try {
-      writer.waitForTraces(expectedSize)
+      // Fail fast if we see any muzzle errors or instrumentation errors
+      for (int timeoutRemaining = ListWriter.WAIT_FOR_TRACES_TIMEOUT_SECONDS; timeoutRemaining > 0; timeoutRemaining--) {
+        if (writer.waitForTracesMax(expectedSize, 1)) {
+          break
+        }
+        assertNoErrors()
+      }
+      assertNoErrors() // even if we break immediately above, make sure we check for errors at least once
+      writer.waitForTraces(expectedSize, 0)
       def array = writer.toArray()
       assert array.length == expectedSize
       def traces = (Arrays.asList(array) as List<List<DDSpan>>)
@@ -91,6 +101,12 @@ class ListWriterAssert {
     }
   }
 
+  static void assertNoErrors() {
+    assert InstrumentationErrors.getErrors().size() == 0 : "Instrumentation errors were seen, failing fast"
+    assert MuzzleCheck.getErrors().size() == 0 : "Muzzle errors were seen, failing fast"
+    // InstrumentationSpecification.{cleanup,cleanupAfterAgent} will provide details later
+  }
+
   void sortSpansByStart() {
     traces = traces.collect {
       return new ArrayList<DDSpan>(it).sort { a, b ->

dd-trace-core/src/main/java/datadog/trace/common/writer/ListWriter.java

@@ -18,6 +18,8 @@
 
 /** List writer used by tests mostly */
 public class ListWriter extends CopyOnWriteArrayList<List<DDSpan>> implements Writer {
+  public static final int WAIT_FOR_TRACES_TIMEOUT_SECONDS = 20;
+
   private static final Logger log = LoggerFactory.getLogger(ListWriter.class);
   private static final Filter ACCEPT_ALL = trace -> true;
 
@@ -85,7 +87,12 @@ public boolean waitForTracesMax(final int number, int seconds) throws Interrupte
   }
 
   public void waitForTraces(final int number) throws InterruptedException, TimeoutException {
-    if (!waitForTracesMax(number, 20)) {
+    waitForTraces(number, WAIT_FOR_TRACES_TIMEOUT_SECONDS);
+  }
+
+  public void waitForTraces(final int number, final int seconds)
+      throws InterruptedException, TimeoutException {
+    if (!waitForTracesMax(number, seconds)) {
       String msg =
           "Timeout waiting for "
               + number