Delay creation of advice remapper+cache until needed

DataDog · Jan 13, 2025 · ede74e3 · ede74e3
1 parent d402356
commit ede74e3
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 64 deletions.
diff --git a/...va-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/ShadedAdviceLocator.java b/...va-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/ShadedAdviceLocator.java
@@ -17,7 +17,7 @@ public ShadedAdviceLocator(ClassLoader adviceLoader, AdviceShader adviceShader)
   public Resolution locate(String className) throws IOException {
     final Resolution resolution = adviceLocator.locate(className);
     if (resolution.isResolved()) {
-      return new Resolution.Explicit(adviceShader.shade(resolution.resolve()));
+      return new Resolution.Explicit(adviceShader.shadeClass(resolution.resolve()));
     } else {
       return resolution;
     }

diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceShader.java
@@ -10,89 +10,105 @@
 import net.bytebuddy.jar.asm.commons.Remapper;
 
 /** Shades advice bytecode by applying relocations to all references. */
-public final class AdviceShader extends Remapper {
-  private final DDCache<String, String> cache = DDCaches.newFixedSizeCache(64);
+public final class AdviceShader {
+  private final Map<String, String> relocations;
 
-  /** Flattened sequence of old-prefix, new-prefix relocations. */
-  private final String[] prefixes;
+  private volatile Remapper remapper;
 
   public static AdviceShader with(Map<String, String> relocations) {
-    return relocations != null ? new AdviceShader(relocations) : null;
+    if (relocations != null) {
+      return new AdviceShader(relocations);
+    }
+    return null;
   }
 
-  AdviceShader(Map<String, String> relocations) {
-    // convert relocations to a flattened sequence: old-prefix, new-prefix, etc.
-    this.prefixes = new String[relocations.size() * 2];
-    int i = 0;
-    for (Map.Entry<String, String> e : relocations.entrySet()) {
-      String oldPrefix = e.getKey();
-      String newPrefix = e.getValue();
-      if (oldPrefix.indexOf('.') > 0) {
-        // accept dotted prefixes, but store them in their internal form
-        this.prefixes[i++] = oldPrefix.replace('.', '/');
-        this.prefixes[i++] = newPrefix.replace('.', '/');
-      } else {
-        this.prefixes[i++] = oldPrefix;
-        this.prefixes[i++] = newPrefix;
-      }
-    }
+  private AdviceShader(Map<String, String> relocations) {
+    this.relocations = relocations;
   }
 
   /** Applies shading before calling the given {@link ClassVisitor}. */
-  public ClassVisitor shade(ClassVisitor cv) {
-    return new ClassRemapper(cv, this);
+  public ClassVisitor shadeClass(ClassVisitor cv) {
+    if (null == remapper) {
+      remapper = new AdviceMapper();
+    }
+    return new ClassRemapper(cv, remapper);
   }
 
   /** Returns the result of shading the given bytecode. */
-  public byte[] shade(byte[] bytecode) {
+  public byte[] shadeClass(byte[] bytecode) {
     ClassReader cr = new ClassReader(bytecode);
     ClassWriter cw = new ClassWriter(null, 0);
-    cr.accept(shade(cw), 0);
+    cr.accept(shadeClass(cw), 0);
     return cw.toByteArray();
   }
 
-  @Override
-  public String map(String internalName) {
-    if (internalName.startsWith("java/")
-        || internalName.startsWith("datadog/")
-        || internalName.startsWith("net/bytebuddy/")) {
-      return internalName; // never shade these references
+  final class AdviceMapper extends Remapper {
+    private final DDCache<String, String> mappingCache = DDCaches.newFixedSizeCache(64);
+
+    /** Flattened sequence of old-prefix, new-prefix relocations. */
+    private final String[] prefixes;
+
+    AdviceMapper() {
+      // convert relocations to a flattened sequence: old-prefix, new-prefix, etc.
+      this.prefixes = new String[relocations.size() * 2];
+      int i = 0;
+      for (Map.Entry<String, String> e : relocations.entrySet()) {
+        String oldPrefix = e.getKey();
+        String newPrefix = e.getValue();
+        if (oldPrefix.indexOf('.') > 0) {
+          // accept dotted prefixes, but store them in their internal form
+          this.prefixes[i++] = oldPrefix.replace('.', '/');
+          this.prefixes[i++] = newPrefix.replace('.', '/');
+        } else {
+          this.prefixes[i++] = oldPrefix;
+          this.prefixes[i++] = newPrefix;
+        }
+      }
     }
-    return cache.computeIfAbsent(internalName, this::shade);
-  }
 
-  @Override
-  public Object mapValue(Object value) {
-    if (value instanceof String) {
-      String text = (String) value;
-      if (text.isEmpty()) {
-        return text;
-      } else if (text.indexOf('.') > 0) {
-        return shadeDottedName(text);
+    @Override
+    public String map(String internalName) {
+      if (internalName.startsWith("java/")
+          || internalName.startsWith("datadog/")
+          || internalName.startsWith("net/bytebuddy/")) {
+        return internalName; // never shade these references
+      }
+      return mappingCache.computeIfAbsent(internalName, this::shadeInternalName);
+    }
+
+    @Override
+    public Object mapValue(Object value) {
+      if (value instanceof String) {
+        String text = (String) value;
+        if (text.isEmpty()) {
+          return text;
+        } else if (text.indexOf('.') > 0) {
+          return shadeDottedName(text);
+        } else {
+          return shadeInternalName(text);
+        }
       } else {
-        return shade(text);
+        return super.mapValue(value);
       }
-    } else {
-      return super.mapValue(value);
     }
-  }
 
-  private String shade(String internalName) {
-    for (int i = 0; i < prefixes.length; i += 2) {
-      if (internalName.startsWith(prefixes[i])) {
-        return prefixes[i + 1] + internalName.substring(prefixes[i].length());
+    private String shadeInternalName(String internalName) {
+      for (int i = 0; i < prefixes.length; i += 2) {
+        if (internalName.startsWith(prefixes[i])) {
+          return prefixes[i + 1] + internalName.substring(prefixes[i].length());
+        }
       }
+      return internalName;
     }
-    return internalName;
-  }
 
-  private String shadeDottedName(String name) {
-    String internalName = name.replace('.', '/');
-    for (int i = 0; i < prefixes.length; i += 2) {
-      if (internalName.startsWith(prefixes[i])) {
-        return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length());
+    private String shadeDottedName(String name) {
+      String internalName = name.replace('.', '/');
+      for (int i = 0; i < prefixes.length; i += 2) {
+        if (internalName.startsWith(prefixes[i])) {
+          return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length());
+        }
       }
+      return name;
     }
-    return name;
   }
 }
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -89,12 +89,12 @@ public HelperInjector(
   private Map<String, byte[]> getHelperMap() throws IOException {
     if (dynamicTypeMap.isEmpty()) {
       final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
-      for (final String helperClassName : helperClassNames) {
-        byte[] classBytes = classFileLocator.locate(helperClassName).resolve();
+      for (String helperName : helperClassNames) {
+        byte[] classBytes = classFileLocator.locate(helperName).resolve();
         if (adviceShader != null) {
-          classBytes = adviceShader.shade(classBytes);
+          classBytes = adviceShader.shadeClass(classBytes);
         }
-        classnameToBytes.put(helperClassName, classBytes);
+        classnameToBytes.put(helperName, classBytes);
       }
 
       return classnameToBytes;

diff --git a/...ent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java b/...ent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
@@ -132,7 +132,7 @@ private static Map<String, byte[]> createHelperMap(final InstrumenterModule modu
           ClassFileLocator.ForClassLoader.of(module.getClass().getClassLoader());
       byte[] classBytes = locator.locate(helperName).resolve();
       if (null != adviceShader) {
-        classBytes = adviceShader.shade(classBytes);
+        classBytes = adviceShader.shadeClass(classBytes);
       }
       helperMap.put(helperName, classBytes);
     }

diff --git a/...gent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java b/...gent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceCreator.java
@@ -70,7 +70,7 @@ public static Map<String, Reference> createReferencesFrom(
         if (null == adviceShader) {
           reader.accept(cv, ClassReader.SKIP_FRAMES);
         } else {
-          reader.accept(adviceShader.shade(cv), ClassReader.SKIP_FRAMES);
+          reader.accept(adviceShader.shadeClass(cv), ClassReader.SKIP_FRAMES);
         }
 
         final Map<String, Reference> instrumentationReferences = cv.getReferences();