Service naming: split by jee deployment

Author: amarziali

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

@@ -1,13 +1,16 @@
 package datadog.trace.instrumentation.jbossmodules;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.instrumentation.jbossmodules.ModuleNameHelper.STRIPPER;
+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 +36,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 +64,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 +159,11 @@ public static void onExit(
       }
     }
   }
+
+  public static class CaptureModuleNameAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void afterConstruct(@Advice.This final Module module) {
+      ClassloaderServiceNames.addIfMissing(module.getClassLoader(), STRIPPER);
+    }
+  }
 }

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

@@ -0,0 +1,25 @@
+package datadog.trace.instrumentation.jbossmodules;
+
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+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 final Function<ClassLoader, Supplier<String>> STRIPPER =
+      classLoader -> {
+        final Matcher matcher =
+            SUBDEPLOYMENT_MATCH.matcher(
+                ((ModuleClassLoader) classLoader).getModule().getIdentifier().getName());
+        if (matcher.matches()) {
+          final String result = matcher.group(1);
+          return () -> result;
+        }
+        return () -> null;
+      };
+}

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

@@ -0,0 +1,25 @@
+package datadog.trace.instrumentation.liberty20;
+
+import com.ibm.ws.classloading.internal.ThreadContextClassLoader;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+public class BundleNameHelper {
+  private BundleNameHelper() {}
+
+  public static final Function<ClassLoader, Supplier<String>> EXTRACTOR =
+      classLoader -> {
+        final String id = ((ThreadContextClassLoader) 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 < 0) {
+          return () -> null;
+        }
+        final String name = id.substring(head + 1, tail);
+        return () -> name;
+      };
+}

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

@@ -12,11 +12,14 @@
 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.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 +107,16 @@ public static class HandleRequestAdvice {
       request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
       final AgentSpan span = DECORATE.startSpan(request, extractedContext);
       scope = activateSpan(span, true);
-
+      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::setServiceName, span::getServiceName, 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,44 @@
+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) {
+      ClassloaderServiceNames.addIfMissing(self, BundleNameHelper.EXTRACTOR);
+    }
+  }
+}

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,14 @@
 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.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 +109,16 @@ public static class HandleRequestAdvice {
       request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
       final AgentSpan span = DECORATE.startSpan(request, extractedContext);
       scope = activateSpan(span, true);
-
+      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::setServiceName, span::getServiceName, 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,9 @@ 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::setServiceName, span::getServiceName, Thread.currentThread());
       // Tracing might already be applied by the FilterChain or a parent request (forward/include).
       return false;
     }