Merge branch 'master' into alejandro.gonzalez/APPSEC-61873-3

Author: jandro996

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

@@ -20,7 +20,7 @@
 import datadog.trace.core.DDSpan;
 import datadog.trace.core.PendingTrace;
 import datadog.trace.core.TraceCollector;
-import de.thetaphi.forbiddenapis.SuppressForbidden;
+import datadog.trace.junit.utils.context.AllowContextTestingExtension;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
 import java.util.List;
@@ -31,7 +31,6 @@
 import java.util.function.Predicate;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.opentest4j.AssertionFailedError;
@@ -42,7 +41,7 @@
  * current implementation is inspired and kept close to it Groovy / Spock counterpart, the {@code
  * InstrumentationSpecification}.
  */
-@ExtendWith(TestClassShadowingExtension.class)
+@ExtendWith({TestClassShadowingExtension.class, AllowContextTestingExtension.class})
 public abstract class AbstractInstrumentationTest {
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
 
@@ -55,19 +54,6 @@ public abstract class AbstractInstrumentationTest {
   protected ClassFileTransformer activeTransformer;
   protected ClassFileTransformerListener transformerLister;
 
-  @SuppressForbidden // Class.forName() used to dynamically configure context if present
-  @BeforeAll
-  static void allowContextTesting() {
-    // Allow re-registration of context managers so each test can use a fresh tracer.
-    // This mirrors DDSpecification.allowContextTesting() for the Spock test framework.
-    try {
-      Class.forName("datadog.context.ContextManager").getMethod("allowTesting").invoke(null);
-      Class.forName("datadog.context.ContextBinder").getMethod("allowTesting").invoke(null);
-    } catch (Throwable ignore) {
-      // don't block testing if context types aren't available
-    }
-  }
-
   @BeforeEach
   public void init() {
     // If this fails, it's likely the result of another test loading Config before it can be

dd-smoke-tests/src/main/groovy/datadog/smoketest/ProcessManager.groovy

@@ -50,7 +50,7 @@ abstract class ProcessManager extends Specification {
 
   // Here for backwards compatibility with single process case
   @Shared
-  def logFilePath = logFilePaths[0]
+  def logFilePath = logFilePaths.length > 0 ? logFilePaths[0] : null
 
   def setup() {
     testedProcesses.each {

dd-smoke-tests/websphere-jmx/build.gradle

@@ -0,0 +1,15 @@
+apply from: "$rootDir/gradle/java.gradle"
+
+dependencies {
+  testImplementation project(':dd-smoke-tests')
+  testImplementation libs.testcontainers
+}
+
+testJvmConstraints {
+  // there is no need to run it multiple times since it runs on a container
+  maxJavaVersion = JavaVersion.VERSION_1_8
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+}

dd-smoke-tests/websphere-jmx/src/test/groovy/datadog/smoketest/WebSphereJmxSmokeTest.groovy

@@ -0,0 +1,118 @@
+package datadog.smoketest
+
+
+import java.time.Duration
+import java.util.concurrent.ArrayBlockingQueue
+import java.util.concurrent.BlockingQueue
+import java.util.concurrent.TimeUnit
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+import org.testcontainers.containers.GenericContainer
+import org.testcontainers.containers.output.Slf4jLogConsumer
+import org.testcontainers.containers.wait.strategy.Wait
+import org.testcontainers.utility.MountableFile
+import spock.lang.Shared
+
+/**
+ * Smoke test for the websphere-jmx instrumentation.
+ *
+ * Builds and starts a WebSphere traditional (tWAS) container with the dd-java-agent baked in
+ * via jvm-config.props, then verifies that jmxfetch reports WebSphere thread pool metrics
+ * via statsd UDP.
+ *
+ * Note that the websphere related metrics will only arrive if our instrumentation is applied.
+ */
+class WebSphereJmxSmokeTest extends AbstractSmokeTest {
+
+  private static final Logger LOG = LoggerFactory.getLogger(WebSphereJmxSmokeTest)
+
+  @Override
+  protected int numberOfProcesses() {
+    return 0
+  }
+
+  @Shared
+  DatagramSocket statsdSocket
+
+  @Shared
+  int statsdPort
+
+  @Shared
+  BlockingQueue<String> statsdMessages = new ArrayBlockingQueue<>(256)
+
+  @Shared
+  Thread listenerThread
+
+  @Shared
+  GenericContainer websphere
+
+  def setupSpec() {
+    statsdSocket = new DatagramSocket()
+    statsdPort = statsdSocket.getLocalPort()
+    LOG.info("StatsDServer listening on UDP port {}", statsdPort)
+
+    listenerThread = Thread.start {
+      byte[] buf = new byte[2048]
+      while (!Thread.currentThread().interrupted()) {
+        try {
+          DatagramPacket packet = new DatagramPacket(buf, buf.length)
+          statsdSocket.receive(packet)
+          String msg = new String(packet.getData(), 0, packet.getLength())
+          LOG.debug("Received statsd: {}", msg)
+          statsdMessages.offer(msg, 1, TimeUnit.SECONDS)
+        } catch (Exception ignored) {
+          break
+        }
+      }
+    }
+
+    websphere = new GenericContainer("icr.io/appcafe/websphere-traditional:latest")
+      // inject wished jvm props for the server we are running
+      .withCopyFileToContainer(MountableFile.forClasspathResource("jvm-config.props"), "/work/config/")
+      // copy the agent jar
+      .withCopyFileToContainer(MountableFile.forHostPath(shadowJarPath), "/opt/dd-java-agent.jar")
+      // let it run on a macos for dev
+      .withCreateContainerCmdModifier { it.withPlatform('linux/amd64') }
+      // this is required to send back udp datagrams to us
+      .withExtraHost('host.docker.internal', 'host-gateway')
+      // set jmxfetch props
+      .withEnv('DD_JMXFETCH_STATSD_HOST', 'host.docker.internal')
+      .withEnv('DD_JMXFETCH_STATSD_PORT', String.valueOf(statsdPort))
+      .withLogConsumer(new Slf4jLogConsumer(LOG).withPrefix('websphere'))
+      // the server will restart 2 times. First to update the jvm props
+      .waitingFor(Wait.forLogMessage('.*open for e-business.*', 2))
+      // it can be long
+      .withStartupTimeout(Duration.ofMinutes(8))
+      // override the command (by default it's /work/start_server.sh)
+      .withCommand("bash", "-c", "/work/configure.sh && /work/start_server.sh")
+
+    websphere.start()
+    LOG.info("WebSphere container started")
+  }
+
+  def cleanupSpec() {
+    websphere?.stop()
+    statsdSocket?.close()
+    listenerThread?.join()
+  }
+
+  def "jmxfetch reports WebSphere thread pool metrics via statsd"() {
+    when: "waiting for websphere.thread_pool metrics"
+    String metric = waitForMetric('websphere.thread_pool', Duration.ofMinutes(3))
+
+    then: "at least one thread pool metric arrives"
+    metric != null
+    metric.contains('websphere.thread_pool')
+  }
+
+  def waitForMetric(String prefix, Duration timeout) {
+    long deadline = System.currentTimeMillis() + timeout.toMillis()
+    while (System.currentTimeMillis() < deadline) {
+      String msg = statsdMessages.poll(5, TimeUnit.SECONDS)
+      if (msg?.contains(prefix)) {
+        return msg
+      }
+    }
+    return null
+  }
+}

dd-smoke-tests/websphere-jmx/src/test/resources/jvm-config.props

@@ -0,0 +1,6 @@
+ResourceType=JavaVirtualMachine
+ImplementingResourceType=Server
+ResourceId=Cell=!{cellName}:Node=!{nodeName}:Server=!{serverName}:JavaProcessDef=:JavaVirtualMachine=
+AttributeInfo=jvmEntries
+
+genericJvmArguments=-javaagent:/opt/dd-java-agent.jar -Ddd.trace.debug=true -Ddd.jmxfetch.websphere.enabled=true

dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java

@@ -1,6 +1,6 @@
 package datadog.trace.api;
 
-import datadog.trace.test.util.TableTestTypeConverters;
+import datadog.trace.junit.utils.tabletest.TableTestTypeConverters;
 import org.tabletest.junit.TypeConverter;
 
 /** TableTest converters shared by dd-trace-api test classes for unparsable constants. */

dd-trace-core/src/jmh/java/datadog/trace/core/TimeSourceBenchmark.java

@@ -0,0 +1,99 @@
+package datadog.trace.core;
+
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
+
+import datadog.trace.api.DDTraceId;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+@State(Scope.Thread)
+@Warmup(iterations = 3)
+@Measurement(iterations = 5)
+@BenchmarkMode(Mode.Throughput)
+@Threads(8)
+@OutputTimeUnit(NANOSECONDS)
+@Fork(value = 1)
+public class TimeSourceBenchmark {
+
+  static final CoreTracer TRACER = CoreTracer.builder().build();
+
+  private PendingTrace pendingTrace;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    TraceCollector collector = TRACER.createTraceCollector(DDTraceId.ONE);
+    pendingTrace = (PendingTrace) collector;
+  }
+
+  @TearDown(Level.Trial)
+  public void teardown() {
+    pendingTrace = null;
+  }
+
+  @Benchmark
+  public long getCurrentTimeNano() {
+    return pendingTrace.getCurrentTimeNano();
+  }
+
+  @Benchmark
+  public long systemNanoTime() {
+    return System.nanoTime();
+  }
+
+  @Benchmark
+  public long systemCurrentTimeMillis() {
+    return System.currentTimeMillis();
+  }
+
+  @Benchmark
+  public long traceGetTimeWithNanoTicks() {
+    return TRACER.getTimeWithNanoTicks(System.nanoTime());
+  }
+
+  /**
+   * Measures a full span start + finish cycle, exercising both the {@code rootSpan} CAS guard in
+   * {@link PendingTrace#registerSpan} and the {@code lazySet} of {@code lastReferenced} in {@link
+   * PendingTrace#getCurrentTimeNano}.
+   */
+  @Benchmark
+  public void startAndFinishSpan() {
+    TRACER.startSpan("benchmark", "op").finish();
+  }
+
+  @State(Scope.Benchmark)
+  public static class SharedState {
+    PendingTrace sharedTrace;
+
+    @Setup(Level.Trial)
+    public void setup() {
+      TraceCollector collector = TRACER.createTraceCollector(DDTraceId.ONE);
+      sharedTrace = (PendingTrace) collector;
+    }
+
+    @TearDown(Level.Trial)
+    public void teardown() {
+      sharedTrace = null;
+    }
+  }
+
+  /**
+   * Measures {@link PendingTrace#getCurrentTimeNano()} under cross-thread contention on a single
+   * shared {@code PendingTrace}. All threads write to the same {@code lastReferenced} field,
+   * demonstrating the benefit of {@code lazySet} over a volatile store under contention.
+   */
+  @Benchmark
+  public long getCurrentTimeNano_contended(SharedState shared) {
+    return shared.sharedTrace.getCurrentTimeNano();
+  }
+}

dd-trace-core/src/main/java/datadog/trace/core/PendingTrace.java

@@ -131,10 +131,14 @@ public PendingTrace create(@Nonnull DDTraceId traceId, ConfigSnapshot traceConfi
 
   /**
    * Updated with the latest nanoTicks each time getCurrentTimeNano is called (at the start and
-   * finish of each span).
+   * finish of each span). Uses lazySet for writes (release-store semantics) since the value is only
+   * read for approximate timeout detection by the PendingTraceBuffer background thread.
    */
   private volatile long lastReferenced = 0;
 
+  private static final AtomicLongFieldUpdater<PendingTrace> LAST_REFERENCED =
+      AtomicLongFieldUpdater.newUpdater(PendingTrace.class, "lastReferenced");
+
   private PendingTrace(
       @Nonnull CoreTracer tracer,
       @Nonnull DDTraceId traceId,
@@ -163,25 +167,27 @@ private PendingTrace(
   @Override
   public long getCurrentTimeNano() {
     long nanoTicks = timeSource.getNanoTicks();
-    lastReferenced = nanoTicks;
+    LAST_REFERENCED.lazySet(this, nanoTicks);
     return tracer.getTimeWithNanoTicks(nanoTicks);
   }
 
   @Override
   void touch() {
-    lastReferenced = timeSource.getNanoTicks();
+    LAST_REFERENCED.lazySet(this, timeSource.getNanoTicks());
   }
 
   @Override
   public boolean lastReferencedNanosAgo(long nanos) {
     long currentNanoTicks = timeSource.getNanoTicks();
-    long age = currentNanoTicks - lastReferenced;
+    long age = currentNanoTicks - LAST_REFERENCED.get(this);
     return nanos < age;
   }
 
   @Override
   void registerSpan(final DDSpan span) {
-    ROOT_SPAN.compareAndSet(this, null, span);
+    if (rootSpan == null) {
+      ROOT_SPAN.compareAndSet(this, null, span);
+    }
     PENDING_REFERENCE_COUNT.incrementAndGet(this);
     healthMetrics.onCreateSpan();
     if (pendingTraceBuffer.longRunningSpansEnabled()) {

dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java

@@ -32,11 +32,6 @@
 class W3CHttpCodec {
   private static final Logger log = LoggerFactory.getLogger(W3CHttpCodec.class);
 
-  static final String TRACE_PARENT_KEY = "traceparent";
-  static final String TRACE_STATE_KEY = "tracestate";
-  static final String OT_BAGGAGE_PREFIX = "ot-baggage-";
-  private static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME;
-
   private static final int TRACE_PARENT_TID_START = 2 + 1;
   private static final int TRACE_PARENT_TID_END = TRACE_PARENT_TID_START + 32;
   private static final int TRACE_PARENT_SID_START = TRACE_PARENT_TID_END + 1;
@@ -45,6 +40,12 @@ class W3CHttpCodec {
   private static final int TRACE_PARENT_FLAGS_SAMPLED = 1;
   private static final int TRACE_PARENT_LENGTH = TRACE_PARENT_FLAGS_START + 2;
 
+  // Package-protected for testing
+  static final String TRACE_PARENT_KEY = "traceparent";
+  static final String TRACE_STATE_KEY = "tracestate";
+  static final String OT_BAGGAGE_PREFIX = "ot-baggage-";
+  static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME;
+
   private W3CHttpCodec() {
     // This class should not be created. This also makes code coverage checks happy.
   }