refactor according to the review

DataDog · Dec 12, 2024 · c2cb1c4 · c2cb1c4
1 parent 4b43be2
commit c2cb1c4
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 51 deletions.
diff --git a/...dules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleInstrumentation.java b/...dules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleInstrumentation.java
@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.jbossmodules;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.instrumentation.jbossmodules.ModuleNameHelper.NAME_EXTRACTOR;
+import static datadog.trace.instrumentation.jbossmodules.ModuleNameHelper.extractDeploymentName;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -163,7 +163,10 @@ 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(), NAME_EXTRACTOR);
+      final String name = extractDeploymentName(module.getClassLoader());
+      if (name != null && !name.isEmpty()) {
+        ClassloaderServiceNames.addServiceName(module.getClassLoader(), name);
+      }
     }
   }
 }
diff --git a/...ss-modules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleNameHelper.java b/...ss-modules/src/main/java/datadog/trace/instrumentation/jbossmodules/ModuleNameHelper.java
@@ -1,25 +1,23 @@
 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 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 final Function<ClassLoader, Supplier<String>> NAME_EXTRACTOR =
-      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;
-      };
+
+  public static String extractDeploymentName(@Nonnull final ClassLoader classLoader) {
+    final Matcher matcher =
+        SUBDEPLOYMENT_MATCH.matcher(
+            ((ModuleClassLoader) classLoader).getModule().getIdentifier().getName());
+    if (matcher.matches()) {
+      return matcher.group(1);
+    }
+    return null;
+  };
 }
diff --git a/...on/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/BundleNameHelper.java b/...on/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/BundleNameHelper.java
@@ -1,25 +1,22 @@
 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;
-      };
+  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('#', head);
+    if (tail < 0) {
+      return null;
+    }
+    final String name = id.substring(head + 1, tail);
+    return null;
+  }
 }
diff --git a/...java/datadog/trace/instrumentation/liberty20/ThreadContextClassloaderInstrumentation.java b/...java/datadog/trace/instrumentation/liberty20/ThreadContextClassloaderInstrumentation.java
@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.liberty20;
 
+import static datadog.trace.instrumentation.liberty20.BundleNameHelper.extractDeploymentName;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
 import com.google.auto.service.AutoService;
@@ -38,7 +39,10 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class ThreadContextClassloaderAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void afterConstruct(@Advice.This ThreadContextClassLoader self) {
-      ClassloaderServiceNames.addIfMissing(self, BundleNameHelper.EXTRACTOR);
+      final String name = extractDeploymentName(self);
+      if (name != null && !name.isEmpty()) {
+        ClassloaderServiceNames.addServiceName(self, name);
+      }
     }
   }
 }
diff --git a/...5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java b/...5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java
@@ -2,6 +2,7 @@
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -56,9 +57,7 @@ public static AgentSpan before(@Advice.Argument(0) final ServletRequest request)
       if (!(request instanceof HttpServletRequest)) {
         return null;
       }
-      Object span =
-          request.getAttribute(
-              "datadog.span"); // hardcode to avoid injecting HttpServiceDecorator just for this
+      Object span = request.getAttribute(DD_SPAN_ATTRIBUTE);
       if (span instanceof AgentSpan
           && CallDepthThreadLocalMap.incrementCallDepth(HttpServletRequest.class) == 0) {
         final AgentSpan agentSpan = (AgentSpan) span;

diff --git a/...src/main/java/datadog/trace/instrumentation/tomcat9/WebappClassLoaderInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/tomcat9/WebappClassLoaderInstrumentation.java
@@ -1,12 +1,15 @@
 package datadog.trace.instrumentation.tomcat9;
 
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 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 net.bytebuddy.asm.Advice;
+import org.apache.catalina.Context;
+import org.apache.catalina.WebResourceRoot;
 import org.apache.catalina.loader.WebappClassLoaderBase;
 
 @AutoService(InstrumenterModule.class)
@@ -30,13 +33,23 @@ public String[] helperClassNames() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureWebappNameAdvice");
+    transformer.applyAdvice(
+        isMethod().and(named("setResources")), getClass().getName() + "$CaptureWebappNameAdvice");
   }
 
   public static class CaptureWebappNameAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void afterConstruct(@Advice.This final WebappClassLoaderBase classLoader) {
-      ClassloaderServiceNames.addIfMissing(classLoader, ContextNameHelper.ADDER);
+    public static void onContextAvailable(
+        @Advice.This final WebappClassLoaderBase classLoader,
+        @Advice.Argument(0) final WebResourceRoot webResourceRoot) {
+      // at this moment we have the context set in this classloader, hence its name
+      final Context context = webResourceRoot.getContext();
+      if (context != null) {
+        final String contextName = context.getBaseName();
+        if (contextName != null && !contextName.isEmpty()) {
+          ClassloaderServiceNames.addServiceName(classLoader, contextName);
+        }
+      }
     }
   }
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/naming/ClassloaderServiceNames.java b/internal-api/src/main/java/datadog/trace/api/naming/ClassloaderServiceNames.java
@@ -6,8 +6,6 @@
 import datadog.trace.api.remoteconfig.ServiceNameCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.WeakHashMap;
-import java.util.function.Function;
-import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -16,29 +14,24 @@ private static class Lazy {
     private static final ClassloaderServiceNames INSTANCE = new ClassloaderServiceNames();
   }
 
-  private final WeakHashMap<ClassLoader, Supplier<String>> weakCache = new WeakHashMap<>();
+  private final WeakHashMap<ClassLoader, String> weakCache = new WeakHashMap<>();
   private final String inferredServiceName =
       CapturedEnvironment.get().getProperties().get(GeneralConfig.SERVICE_NAME);
   private final boolean enabled =
       Config.get().isJeeSplitByDeployment() && !Config.get().isServiceNameSetByUser();
 
   private ClassloaderServiceNames() {}
 
-  public static void addIfMissing(
-      @Nonnull ClassLoader classLoader,
-      @Nonnull Function<? super ClassLoader, Supplier<String>> adder) {
+  public static void addServiceName(@Nonnull ClassLoader classLoader, @Nonnull String serviceName) {
     if (Lazy.INSTANCE.enabled) {
-      Lazy.INSTANCE.weakCache.computeIfAbsent(classLoader, adder);
+      Lazy.INSTANCE.weakCache.put(classLoader, serviceName);
     }
   }
 
   @Nullable
   public static String maybeGet(@Nonnull ClassLoader classLoader) {
     if (Lazy.INSTANCE.enabled) {
-      final Supplier<String> supplier = Lazy.INSTANCE.weakCache.get(classLoader);
-      if (supplier != null) {
-        return supplier.get();
-      }
+      return Lazy.INSTANCE.weakCache.get(classLoader);
     }
     return null;
   }