review

Author: amarziali

dd-java-agent/instrumentation/wildfly-9/src/main/java/datadog/trace/instrumentation/wildfly/ResourceReferenceProcessorInstrumentation.java

@@ -51,8 +51,7 @@ public static class GetEnvironmentEntriesAdvice {
     public static void onExit(
         @Advice.Argument(value = 1) final ClassLoader classLoader,
         @Advice.Return final List<BindingConfiguration> configurations) {
-      ClassloaderConfigurationOverrides.ContextualInfo info =
-          ClassloaderConfigurationOverrides.getOrAddEmpty(classLoader);
+      ClassloaderConfigurationOverrides.ContextualInfo info = null;
       ContextStore<EnvEntryInjectionSource, Object> contextStore =
           InstrumentationContext.get(EnvEntryInjectionSource.class, Object.class);
       for (BindingConfiguration bindingConfiguration : configurations) {
@@ -63,6 +62,9 @@ public static void onExit(
               && bindingConfiguration
                   .getName()
                   .startsWith(ClassloaderConfigurationOverrides.DATADOG_TAGS_JNDI_PREFIX)) {
+            if (info == null) {
+              info = ClassloaderConfigurationOverrides.maybeCreateContextualInfo(classLoader);
+            }
             info.addTag(
                 bindingConfiguration
                     .getName()

internal-api/src/main/java/datadog/trace/api/ClassloaderConfigurationOverrides.java

@@ -19,7 +19,7 @@ public class ClassloaderConfigurationOverrides {
   static boolean CAN_SPLIT_SERVICE_NAME_BY_DEPLOYMENT =
       Config.get().isJeeSplitByDeployment() && !Config.get().isServiceNameSetByUser();
 
- static class Lazy {
+  static class Lazy {
     static final ClassloaderConfigurationOverrides INSTANCE =
         new ClassloaderConfigurationOverrides();
   }
@@ -54,13 +54,12 @@ public Map<String, Object> getTags() {
 
   private ClassloaderConfigurationOverrides() {}
 
-  public static void addContextualInfo(
-      @Nonnull ClassLoader classLoader, @Nonnull ContextualInfo contextualInfo) {
+  public static void addContextualInfo(ClassLoader classLoader, ContextualInfo contextualInfo) {
     Lazy.INSTANCE.weakCache.put(classLoader, contextualInfo);
   }
 
   @Nullable
-  public static ContextualInfo maybeGetContextualInfo(@Nonnull ClassLoader classLoader) {
+  public static ContextualInfo maybeGetContextualInfo(ClassLoader classLoader) {
     return Lazy.INSTANCE.weakCache.get(classLoader);
   }
 
@@ -74,8 +73,7 @@ public static ContextualInfo maybeGetContextualInfo() {
     return maybeGetContextualInfo(Thread.currentThread().getContextClassLoader());
   }
 
-  @Nonnull
-  public static ContextualInfo maybeCreateContextualInfo(@Nonnull ClassLoader classLoader) {
+  public static ContextualInfo maybeCreateContextualInfo(ClassLoader classLoader) {
     return Lazy.INSTANCE.weakCache.computeIfAbsent(classLoader, EMPTY_CONTEXTUAL_INFO_ADDER);
   }
 
@@ -95,14 +93,11 @@ public static void maybeEnrichSpan(
     if (contextualInfo == null) {
       return;
     }
-    if (CAN_SPLIT_SERVICE_NAME_BY_DEPLOYMENT
-        && contextualInfo.serviceName != null
-        && !contextualInfo.getServiceName().isEmpty()) {
+    final String serviceName = contextualInfo.getServiceName();
+    if (CAN_SPLIT_SERVICE_NAME_BY_DEPLOYMENT && serviceName != null && !serviceName.isEmpty()) {
       final String currentServiceName = span.getServiceName();
       if (currentServiceName == null
           || currentServiceName.equals(Lazy.INSTANCE.inferredServiceName)) {
-        final String serviceName = contextualInfo.getServiceName();
-
         span.setServiceName(serviceName);
         ServiceNameCollector.get().addService(serviceName);
       }

internal-api/src/test/groovy/datadog/trace/api/ClassloaderConfigurationOverridesTest.groovy

@@ -45,7 +45,7 @@ class ClassloaderConfigurationOverridesTest extends DDSpecification {
   def "enrich should set tags when present"() {
     def span = Mock(AgentSpan)
     when:
-    ClassloaderConfigurationOverrides.getOrAddEmpty(Thread.currentThread().getContextClassLoader()).addTag("key", "value")
+    ClassloaderConfigurationOverrides.maybeCreateContextualInfo(Thread.currentThread().getContextClassLoader()).addTag("key", "value")
     ClassloaderConfigurationOverrides.maybeEnrichSpan(span)
     then:
     1 * span.setTag("key", "value")