Service naming: split by jee deployment (#8064)

* Service naming: split by jee deployment

* Do not use invokedynamic

* fix tests

* collapse enter and local

* Adjust service name for non legacy tracing

* Add wildfly ejb test

* limit max java verson to 11 for wildfly 15

* renaming

* refactor according to the review

* remove helper

* avoid casting

* review

Author: amarziali

dd-java-agent/instrumentation/jboss-modules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleInstrumentation.java

@@ -1,13 +1,15 @@
 package datadog.trace.instrumentation.jbossmodules;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.naming.ClassloaderServiceNames;
 import datadog.trace.bootstrap.AgentClassLoading;
 import java.io.IOException;
 import java.io.InputStream;
@@ -33,7 +35,8 @@ public String[] helperClassNames() {
     return new String[] {
       "org.jboss.modules.ModuleLinkageHelper",
       "org.jboss.modules.ModuleLinkageHelper$1",
-      "org.jboss.modules.ModuleLinkageHelper$2"
+      "org.jboss.modules.ModuleLinkageHelper$2",
+      packageName + ".ModuleNameHelper",
     };
   }
 
@@ -60,6 +63,7 @@ public void methodAdvice(MethodTransformer transformer) {
                             .and(takesArgument(0, String.class))
                             .and(takesArgument(1, boolean.class)))),
         ModuleInstrumentation.class.getName() + "$WidenLoadClassAdvice");
+    transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureModuleNameAdvice");
   }
 
   /**
@@ -154,4 +158,14 @@ public static void onExit(
       }
     }
   }
+
+  public static class CaptureModuleNameAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void afterConstruct(@Advice.This final Module module) {
+      final String name = ModuleNameHelper.extractDeploymentName(module.getClassLoader());
+      if (name != null && !name.isEmpty()) {
+        ClassloaderServiceNames.addServiceName(module.getClassLoader(), name);
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/jboss-modules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleNameHelper.java

@@ -0,0 +1,22 @@
+package datadog.trace.instrumentation.jbossmodules;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nonnull;
+import org.jboss.modules.ModuleClassLoader;
+
+public class ModuleNameHelper {
+  private ModuleNameHelper() {}
+
+  private static final Pattern SUBDEPLOYMENT_MATCH =
+      Pattern.compile("deployment(?>.+\\.ear)?\\.(.+)\\.[j|w]ar");
+
+  public static String extractDeploymentName(@Nonnull final ModuleClassLoader classLoader) {
+    final Matcher matcher =
+        SUBDEPLOYMENT_MATCH.matcher(classLoader.getModule().getIdentifier().getName());
+    if (matcher.matches()) {
+      return matcher.group(1);
+    }
+    return null;
+  }
+}

dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/BundleNameHelper.java

@@ -0,0 +1,21 @@
+package datadog.trace.instrumentation.liberty20;
+
+import com.ibm.ws.classloading.internal.ThreadContextClassLoader;
+
+public class BundleNameHelper {
+  private BundleNameHelper() {}
+
+  public static String extractDeploymentName(final ThreadContextClassLoader classLoader) {
+    final String id = classLoader.getKey();
+    // id is something like <type>:name#somethingelse
+    final int head = id.indexOf(':');
+    if (head < 0) {
+      return null;
+    }
+    final int tail = id.lastIndexOf('#');
+    if (tail <= head) {
+      return null;
+    }
+    return id.substring(head + 1, tail);
+  }
+}

dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java

@@ -12,11 +12,15 @@
 import com.google.auto.service.AutoService;
 import com.ibm.ws.webcontainer.srt.SRTServletRequest;
 import com.ibm.ws.webcontainer.srt.SRTServletResponse;
+import com.ibm.ws.webcontainer.webapp.WebApp;
+import com.ibm.wsspi.webcontainer.webapp.IWebAppDispatcherContext;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.naming.ClassloaderServiceNames;
 import datadog.trace.bootstrap.ActiveSubsystems;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
@@ -104,7 +108,18 @@ public static class HandleRequestAdvice {
       request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
       final AgentSpan span = DECORATE.startSpan(request, extractedContext);
       scope = activateSpan(span, true);
-
+      if (Config.get().isJeeSplitByDeployment()) {
+        final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext();
+        if (dispatcherContext != null) {
+          final WebApp webapp = dispatcherContext.getWebApp();
+          if (webapp != null) {
+            final ClassLoader cl = webapp.getClassLoader();
+            if (cl != null) {
+              ClassloaderServiceNames.maybeSetToSpan(span, cl);
+            }
+          }
+        }
+      }
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request, request, extractedContext);
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);

dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ThreadContextClassloaderInstrumentation.java

@@ -0,0 +1,47 @@
+package datadog.trace.instrumentation.liberty20;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import com.ibm.ws.classloading.internal.ThreadContextClassLoader;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.naming.ClassloaderServiceNames;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public class ThreadContextClassloaderInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType {
+
+  public ThreadContextClassloaderInstrumentation() {
+    super("liberty", "liberty-classloading");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.ibm.ws.classloading.internal.ThreadContextClassLoader";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".BundleNameHelper",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor(), getClass().getName() + "$ThreadContextClassloaderAdvice");
+  }
+
+  public static class ThreadContextClassloaderAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void afterConstruct(@Advice.This ThreadContextClassLoader self) {
+      final String name = BundleNameHelper.extractDeploymentName(self);
+      if (name != null && !name.isEmpty()) {
+        ClassloaderServiceNames.addServiceName(self, name);
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/liberty-20/src/test/groovy/datadog/trace/instrumentation/liberty20/Liberty20Test.groovy

@@ -213,6 +213,20 @@ class Liberty20AsyncForkedTest extends Liberty20Test implements TestingGenericHt
   }
 }
 
+@IgnoreIf({
+  // failing because org.apache.xalan.transformer.TransformerImpl is
+  // instrumented while on the the global ignores list
+  System.getProperty('java.vm.name') == 'IBM J9 VM' &&
+  System.getProperty('java.specification.version') == '1.8' })
+class LibertyServletClassloaderNamingForkedTest extends Liberty20V0ForkedTest {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    // will not set the service name according to the servlet context value
+    injectSysConfig("trace.experimental.jee.split-by-deployment", "true")
+  }
+}
+
 @IgnoreIf({
   // failing because org.apache.xalan.transformer.TransformerImpl is
   // instrumented while on the the global ignores list

dd-java-agent/instrumentation/liberty-23/build.gradle

@@ -34,6 +34,7 @@ dependencies {
   testImplementation testFixtures(project(':dd-java-agent:appsec'))
   testRuntimeOnly project(':dd-java-agent:instrumentation:osgi-4.3')
   testRuntimeOnly files({ tasks.filterLogbackClassic.filteredLogbackDir })
+  testRuntimeOnly project(':dd-java-agent:instrumentation:liberty-20')
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')
   testRuntimeOnly files({ tasks.shadowJar.archiveFile })
 

dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java

@@ -12,11 +12,15 @@
 import com.google.auto.service.AutoService;
 import com.ibm.ws.webcontainer.srt.SRTServletRequest;
 import com.ibm.ws.webcontainer.srt.SRTServletResponse;
+import com.ibm.ws.webcontainer.webapp.WebApp;
+import com.ibm.wsspi.webcontainer.webapp.IWebAppDispatcherContext;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.naming.ClassloaderServiceNames;
 import datadog.trace.bootstrap.ActiveSubsystems;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
@@ -106,7 +110,18 @@ public static class HandleRequestAdvice {
       request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
       final AgentSpan span = DECORATE.startSpan(request, extractedContext);
       scope = activateSpan(span, true);
-
+      if (Config.get().isJeeSplitByDeployment()) {
+        final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext();
+        if (dispatcherContext != null) {
+          final WebApp webapp = dispatcherContext.getWebApp();
+          if (webapp != null) {
+            final ClassLoader cl = webapp.getClassLoader();
+            if (cl != null) {
+              ClassloaderServiceNames.maybeSetToSpan(span, cl);
+            }
+          }
+        }
+      }
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request, request, extractedContext);
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);

dd-java-agent/instrumentation/liberty-23/src/test/groovy/datadog/trace/instrumentation/liberty23/Liberty23Test.groovy

@@ -170,3 +170,17 @@ class Liberty23V0ForkedTest extends Liberty23Test implements TestingGenericHttpN
   System.getProperty('java.specification.version') == '1.8' })
 class Liberty23V1ForkedTest extends Liberty23Test implements TestingGenericHttpNamingConventions.ServerV1 {
 }
+
+@IgnoreIf({
+  // failing because org.apache.xalan.transformer.TransformerImpl is
+  // instrumented while on the the global ignores list
+  System.getProperty('java.vm.name') == 'IBM J9 VM' &&
+  System.getProperty('java.specification.version') == '1.8' })
+class LibertyServletClassloaderNamingForkedTest extends Liberty23V0ForkedTest {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    // will not set the service name according to the servlet context value
+    injectSysConfig("trace.experimental.jee.split-by-deployment", "true")
+  }
+}

dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java

@@ -9,6 +9,7 @@
 import datadog.trace.api.DDTags;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.naming.ClassloaderServiceNames;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -39,6 +40,8 @@ public static boolean onEnter(
     Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE);
     final boolean hasServletTrace = spanAttr instanceof AgentSpan;
     if (hasServletTrace) {
+      final AgentSpan span = (AgentSpan) spanAttr;
+      ClassloaderServiceNames.maybeSetToSpan(span);
       // Tracing might already be applied by the FilterChain or a parent request (forward/include).
       return false;
     }