Add CI Visibility instrumentation for JMH benchmarks

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/decorator/TestDecorator.java

@@ -4,6 +4,7 @@
 
 public interface TestDecorator {
   String TEST_TYPE = "test";
+  String TEST_TYPE_BENCHMARK = "benchmark";
 
   AgentSpan afterStart(final AgentSpan span);
 

dd-java-agent/instrumentation/jmh/jmh-1.0/build.gradle

@@ -0,0 +1,23 @@
+plugins {
+  id 'dd-trace-java.instrumentation.testing-framework-tests'
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+muzzle {
+  pass {
+    group = 'org.openjdk.jmh'
+    module = 'jmh-core'
+    versions = '[1.0,)'
+  }
+}
+
+dependencies {
+  compileOnly group: 'org.openjdk.jmh', name: 'jmh-core', version: '1.0'
+
+  testImplementation group: 'org.openjdk.jmh', name: 'jmh-core', version: '1.37'
+  testImplementation project(':dd-java-agent:agent-ci-visibility:civisibility-instrumentation-test-fixtures')
+
+  // JMH annotation processor generates META-INF/BenchmarkList from @Benchmark annotations
+  testAnnotationProcessor group: 'org.openjdk.jmh', name: 'jmh-generator-annprocess', version: '1.37'
+}

dd-java-agent/instrumentation/jmh/jmh-1.0/src/main/java/datadog/trace/instrumentation/jmh/DDOutputFormat.java

@@ -0,0 +1,214 @@
+package datadog.trace.instrumentation.jmh;
+
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_ERROR;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_FORKS;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_ITERATIONS;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_MAX;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_MIN;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_MODE;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_P50;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_P90;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_P95;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_P99;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_SAMPLE_COUNT;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_THREADS;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_TIME_UNIT;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_UNIT;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_VALUE;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.BENCHMARK_WARMUP_ITERATIONS;
+
+import datadog.trace.api.civisibility.InstrumentationBridge;
+import datadog.trace.api.civisibility.config.TestSourceData;
+import datadog.trace.api.civisibility.events.TestEventsHandler;
+import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import java.util.Collections;
+import org.openjdk.jmh.infra.BenchmarkParams;
+import org.openjdk.jmh.results.BenchmarkResult;
+import org.openjdk.jmh.results.Result;
+import org.openjdk.jmh.results.RunResult;
+import org.openjdk.jmh.runner.format.OutputFormat;
+import org.openjdk.jmh.util.Statistics;
+
+/**
+ * Wraps a JMH {@link OutputFormat} to emit CI Visibility spans for each benchmark method.
+ *
+ * <p>Hooks only fire once per benchmark method (after all forks and iterations complete), so there
+ * is zero overhead on the benchmark hot path.
+ */
+public class DDOutputFormat implements OutputFormat {
+
+  private final OutputFormat delegate;
+  private final TestEventsHandler<String, String> handler;
+  private final String frameworkVersion;
+
+  // Keys used as suite/test descriptors in the handler — just the full benchmark name strings.
+  // We keep suite and test keys separate so the handler can manage their lifetimes independently.
+  private volatile String currentSuiteKey;
+  private volatile String currentTestKey;
+
+  public DDOutputFormat(OutputFormat delegate, String frameworkVersion) {
+    this.delegate = delegate;
+    this.frameworkVersion = frameworkVersion;
+    this.handler =
+        InstrumentationBridge.createTestEventsHandler(
+            JmhUtils.FRAMEWORK_NAME, null, null, Collections.emptyList());
+  }
+
+  @Override
+  public void startBenchmark(BenchmarkParams benchParams) {
+    delegate.startBenchmark(benchParams);
+
+    String fullName = benchParams.getBenchmark();
+    String[] parts = JmhUtils.splitBenchmarkName(fullName);
+    String suiteName = parts[0];
+    String testName = parts[1];
+    String testParameters = JmhUtils.testParameters(fullName);
+
+    currentSuiteKey = suiteName + "#" + fullName;
+    currentTestKey = fullName;
+
+    handler.onTestSuiteStart(
+        currentSuiteKey,
+        suiteName,
+        JmhUtils.FRAMEWORK_NAME,
+        frameworkVersion,
+        null,
+        Collections.emptyList(),
+        false,
+        TestFrameworkInstrumentation.JMH,
+        null);
+
+    handler.onTestStart(
+        currentSuiteKey,
+        currentTestKey,
+        testName,
+        JmhUtils.FRAMEWORK_NAME,
+        frameworkVersion,
+        testParameters,
+        Collections.emptyList(),
+        TestSourceData.UNKNOWN,
+        null,
+        null);
+  }
+
+  @Override
+  public void endBenchmark(BenchmarkResult result) {
+    String suiteKey = currentSuiteKey;
+    String testKey = currentTestKey;
+    if (suiteKey == null || testKey == null) {
+      delegate.endBenchmark(result);
+      return;
+    }
+
+    try {
+      tagBenchmarkMetrics(result);
+      handler.onTestFinish(testKey, null, null);
+      handler.onTestSuiteFinish(suiteKey, null);
+    } finally {
+      delegate.endBenchmark(result);
+    }
+  }
+
+  private void tagBenchmarkMetrics(BenchmarkResult result) {
+    AgentSpan span = AgentTracer.activeSpan();
+    if (span == null) {
+      return;
+    }
+
+    BenchmarkParams params = result.getParams();
+    span.setTag(BENCHMARK_MODE, params.getMode().shortLabel());
+    span.setTag(BENCHMARK_ITERATIONS, params.getMeasurement().getCount());
+    span.setTag(BENCHMARK_WARMUP_ITERATIONS, params.getWarmup().getCount());
+    span.setTag(BENCHMARK_FORKS, params.getForks());
+    span.setTag(BENCHMARK_THREADS, params.getThreads());
+    span.setTag(BENCHMARK_TIME_UNIT, params.getTimeUnit().name());
+
+    Result<?> primary = result.getPrimaryResult();
+    span.setMetric(BENCHMARK_VALUE, primary.getScore());
+    span.setTag(BENCHMARK_UNIT, primary.getScoreUnit());
+
+    double error = primary.getScoreError();
+    if (!Double.isNaN(error)) {
+      span.setMetric(BENCHMARK_ERROR, error);
+    }
+
+    Statistics stats = primary.getStatistics();
+    if (stats.getN() > 1) {
+      span.setMetric(BENCHMARK_P50, stats.getPercentile(50));
+      span.setMetric(BENCHMARK_P90, stats.getPercentile(90));
+      span.setMetric(BENCHMARK_P95, stats.getPercentile(95));
+      span.setMetric(BENCHMARK_P99, stats.getPercentile(99));
+      span.setMetric(BENCHMARK_MIN, stats.getMin());
+      span.setMetric(BENCHMARK_MAX, stats.getMax());
+      span.setMetric(BENCHMARK_SAMPLE_COUNT, stats.getN());
+    }
+  }
+
+  // ---- Delegation-only methods ----
+
+  @Override
+  public void iteration(
+      BenchmarkParams benchParams, org.openjdk.jmh.infra.IterationParams params, int iteration) {
+    delegate.iteration(benchParams, params, iteration);
+  }
+
+  @Override
+  public void iterationResult(
+      BenchmarkParams benchParams,
+      org.openjdk.jmh.infra.IterationParams params,
+      int iteration,
+      org.openjdk.jmh.results.IterationResult data) {
+    delegate.iterationResult(benchParams, params, iteration, data);
+  }
+
+  @Override
+  public void startRun() {
+    delegate.startRun();
+  }
+
+  @Override
+  public void endRun(java.util.Collection<RunResult> result) {
+    try {
+      handler.close();
+    } finally {
+      delegate.endRun(result);
+    }
+  }
+
+  @Override
+  public void print(String s) {
+    delegate.print(s);
+  }
+
+  @Override
+  public void println(String s) {
+    delegate.println(s);
+  }
+
+  @Override
+  public void verbosePrintln(String s) {
+    delegate.verbosePrintln(s);
+  }
+
+  @Override
+  public void write(int b) {
+    delegate.write(b);
+  }
+
+  @Override
+  public void write(byte[] b) throws java.io.IOException {
+    delegate.write(b);
+  }
+
+  @Override
+  public void flush() {
+    delegate.flush();
+  }
+
+  @Override
+  public void close() {
+    delegate.close();
+  }
+}

dd-java-agent/instrumentation/jmh/jmh-1.0/src/main/java/datadog/trace/instrumentation/jmh/JmhInstrumentation.java

@@ -0,0 +1,60 @@
+package datadog.trace.instrumentation.jmh;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import net.bytebuddy.asm.Advice;
+import org.openjdk.jmh.runner.format.OutputFormat;
+
+@AutoService(InstrumenterModule.class)
+public class JmhInstrumentation extends InstrumenterModule.CiVisibility
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public JmhInstrumentation() {
+    super("ci-visibility", "jmh");
+  }
+
+  @Override
+  public String instrumentedType() {
+    // Instrument BaseRunner (where the 'out' field is declared) so that the final-field write
+    // in RunnerConstructorAdvice is legal: JDK 17+ rejects writing a final field declared in a
+    // superclass from advice injected into the subclass (Runner).
+    return "org.openjdk.jmh.runner.BaseRunner";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".JmhUtils", packageName + ".DDOutputFormat",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor()
+            .and(takesArgument(0, named("org.openjdk.jmh.runner.options.Options")))
+            .and(takesArgument(1, named("org.openjdk.jmh.runner.format.OutputFormat"))),
+        JmhInstrumentation.class.getName() + "$RunnerConstructorAdvice");
+  }
+
+  public static class RunnerConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.FieldValue(value = "out", readOnly = false) OutputFormat out) {
+      if (out instanceof DDOutputFormat) {
+        return;
+      }
+      String version = null;
+      try {
+        version = org.openjdk.jmh.Main.class.getPackage().getImplementationVersion();
+      } catch (Throwable ignored) {
+      }
+      out = new DDOutputFormat(out, version);
+    }
+  }
+}

dd-java-agent/instrumentation/jmh/jmh-1.0/src/main/java/datadog/trace/instrumentation/jmh/JmhUtils.java

@@ -0,0 +1,52 @@
+package datadog.trace.instrumentation.jmh;
+
+import javax.annotation.Nullable;
+
+public final class JmhUtils {
+
+  static final String FRAMEWORK_NAME = "jmh";
+
+  /**
+   * Splits a JMH benchmark name into suite (class) and method parts.
+   *
+   * <p>JMH names have the form {@code "com.example.MyBenchmark.myMethod"} or, when {@code @Param}
+   * combinations are present, {@code "com.example.MyBenchmark.myMethod:size=1000,threads=4"}.
+   */
+  public static String[] splitBenchmarkName(String fullName) {
+    // Strip any @Param suffix before splitting on the class/method boundary
+    int colonIdx = fullName.indexOf(':');
+    String baseName = colonIdx >= 0 ? fullName.substring(0, colonIdx) : fullName;
+
+    int lastDot = baseName.lastIndexOf('.');
+    if (lastDot < 0) {
+      return new String[] {"", fullName};
+    }
+    return new String[] {baseName.substring(0, lastDot), baseName.substring(lastDot + 1)};
+  }
+
+  /**
+   * Returns the {@code test.parameters} JSON string for a parameterized benchmark, or {@code null}
+   * for an unparameterized one.
+   *
+   * <p>Follows the same convention as JUnit 5 parameterized tests: {@code
+   * {"metadata":{"test_name":"<displayName>"}}}.
+   */
+  @Nullable
+  public static String testParameters(String fullName) {
+    int colonIdx = fullName.indexOf(':');
+    if (colonIdx < 0) {
+      return null;
+    }
+    // fullName after last dot includes the param suffix, e.g. "myMethod:size=1000"
+    int lastDot = fullName.lastIndexOf('.', colonIdx);
+    String displayName = lastDot >= 0 ? fullName.substring(lastDot + 1) : fullName;
+    return "{\"metadata\":{\"test_name\":\"" + escapeJson(displayName) + "\"}}";
+  }
+
+  /** Minimal JSON string escaping for benchmark names (no unicode escaping needed). */
+  private static String escapeJson(String s) {
+    return s.replace("\\", "\\\\").replace("\"", "\\\"");
+  }
+
+  private JmhUtils() {}
+}