Avoid duplicate class symbol extraction

Same classes can be loaded in different classloader, but symbol-wise
they are the same. In some situation (like Groovy script) we can a lot
 of class loaded for same symbols which pollute SymDB backend.
We need to track all classes loaded and extracted into ClassNameTrie
that is space efficient prefix tree. That way we prevent any
duplicates reaching SymDB backend

Author: jpbempel

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SymDBEnablement.java

@@ -115,7 +115,6 @@ public void startSymbolExtraction() {
                 Instant.ofEpochMilli(lastUploadTimestamp), ZoneId.systemDefault()));
         return;
       }
-      symbolAggregator.loadedClassesProcessStarted();
       try {
         symbolExtractionTransformer =
             new SymbolExtractionTransformer(symbolAggregator, classNameFilter);
@@ -125,8 +124,6 @@ public void startSymbolExtraction() {
       } catch (Throwable ex) {
         // catch all Throwables because LinkageError is possible (duplicate class definition)
         LOGGER.debug("Error during symbol extraction: ", ex);
-      } finally {
-        symbolAggregator.loadedClassesProcessEnded();
       }
     } finally {
       starting.set(false);

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SymbolAggregator.java

@@ -4,6 +4,8 @@
 
 import com.datadog.debugger.sink.SymbolSink;
 import datadog.trace.util.AgentTaskScheduler;
+import datadog.trace.util.ClassNameTrie;
+import datadog.trace.util.Strings;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.security.ProtectionDomain;
@@ -12,8 +14,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -28,7 +28,8 @@ public class SymbolAggregator {
   private final AgentTaskScheduler.Scheduled<SymbolAggregator> scheduled;
   private final Object jarScopeLock = new Object();
   private int totalClasses;
-  private volatile Set<String> loadedClasses;
+  // ClassNameTrie is not thread safe, All accesses must be protected by a lock
+  private final ClassNameTrie.Builder loadedClasses = new ClassNameTrie.Builder();
 
   public SymbolAggregator(SymbolSink sink, int symbolFlushThreshold) {
     this.sink = sink;
@@ -61,10 +62,13 @@ public void parseClass(String className, byte[] classfileBuffer, String jarName)
     if (className.endsWith(CLASS_SUFFIX)) {
       className = className.substring(0, className.length() - CLASS_SUFFIX.length());
     }
-    Set<String> localLoadedClasses = loadedClasses;
-    if (localLoadedClasses != null && !localLoadedClasses.add(className)) {
-      // class already loaded and symbol extracted
-      return;
+    synchronized (loadedClasses) {
+      String fqn = Strings.getClassName(className); // ClassNameTrie expects Java class names ('.')
+      if (loadedClasses.apply(fqn) > 0) {
+        // class already loaded and symbol extracted
+        return;
+      }
+      loadedClasses.put(fqn, 1);
     }
     LOGGER.debug("Extracting Symbols from: {}, located in: {}", className, jarName);
     Scope jarScope = SymbolExtractor.extract(classfileBuffer, jarName);
@@ -110,14 +114,4 @@ private void addJarScope(Scope jarScope, boolean forceFlush) {
       }
     }
   }
-
-  void loadedClassesProcessStarted() {
-    // to avoid duplicate symbol extraction we keep track of loaded classes
-    // during the loaded class extraction phase
-    loadedClasses = ConcurrentHashMap.newKeySet();
-  }
-
-  void loadedClassesProcessEnded() {
-    loadedClasses = null;
-  }
 }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/symbol/SymbolExtractionTransformerTest.java

@@ -43,6 +43,7 @@ class SymbolExtractionTransformerTest {
               "javax.",
               "javaslang.",
               "org.omg.",
+              "org.joor.",
               "com.datadog.debugger.")
           .collect(Collectors.toSet());
 
@@ -896,6 +897,20 @@ public void filterOutClassesFromExcludedPackages() throws IOException, URISyntax
             .anyMatch(scope -> scope.getName().equals(CLASS_NAME)));
   }
 
+  @Test
+  public void duplicateClassThroughDifferentClassLoader() throws IOException, URISyntaxException {
+    final String CLASS_NAME = SYMBOL_PACKAGE + "SymbolExtraction01";
+    SymbolSinkMock symbolSinkMock = new SymbolSinkMock(config);
+    SymbolExtractionTransformer transformer = createTransformer(symbolSinkMock);
+    instr.addTransformer(transformer);
+    for (int i = 0; i < 10; i++) {
+      // compile and load the class in a specific ClassLoader each time
+      Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+      Reflect.on(testClass).call("main", "1").get();
+    }
+    assertEquals(1, symbolSinkMock.jarScopes.size());
+  }
+
   private void assertLangSpecifics(
       LanguageSpecifics languageSpecifics,
       List<String> expectedModifiers,