Throw exception when using repo index to resolve source path for classes with identical names

nikita-tkachenko-datadog · nikita-tkachenko-datadog · commit e5f5e07d4164 · 2024-10-18T14:26:20.000+02:00
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/ConcurrentCoverageStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/ConcurrentCoverageStore.java
@@ -4,6 +4,7 @@
 import datadog.trace.api.civisibility.coverage.CoverageProbes;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.coverage.TestReport;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -36,13 +37,18 @@ private T create(Thread thread) {
 
   @Override
   public boolean report(DDTraceId testSessionId, Long testSuiteId, long testSpanId) {
-    report = report(testSessionId, testSuiteId, testSpanId, probes.values());
-    return report != null && report.isNotEmpty();
+    try {
+      report = report(testSessionId, testSuiteId, testSpanId, probes.values());
+      return report != null && report.isNotEmpty();
+    } catch (SourceResolutionException e) {
+      return false;
+    }
   }
 
   @Nullable
   protected abstract TestReport report(
-      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes);
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes)
+      throws SourceResolutionException;
 
   @Nullable
   @Override
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/file/FileCoverageStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/file/FileCoverageStore.java
@@ -11,6 +11,7 @@
 import datadog.trace.api.civisibility.telemetry.tag.CoverageErrorType;
 import datadog.trace.civisibility.coverage.ConcurrentCoverageStore;
 import datadog.trace.civisibility.source.SourcePathResolver;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -46,7 +47,8 @@ private FileCoverageStore(
   @Nullable
   @Override
   protected TestReport report(
-      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<FileProbes> probes) {
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<FileProbes> probes)
+      throws SourceResolutionException {
     try {
       Set<Class<?>> combinedClasses = Collections.newSetFromMap(new IdentityHashMap<>());
       Collection<String> combinedNonCodeResources = new HashSet<>();
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/line/LineCoverageStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/line/LineCoverageStore.java
@@ -11,6 +11,7 @@
 import datadog.trace.api.civisibility.telemetry.tag.CoverageErrorType;
 import datadog.trace.civisibility.coverage.ConcurrentCoverageStore;
 import datadog.trace.civisibility.source.SourcePathResolver;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import datadog.trace.civisibility.source.Utils;
 import java.io.InputStream;
 import java.util.ArrayList;
@@ -52,7 +53,8 @@ private LineCoverageStore(
   @Nullable
   @Override
   protected TestReport report(
-      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<LineProbes> probes) {
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<LineProbes> probes)
+      throws SourceResolutionException {
     try {
       Map<Class<?>, ExecutionDataAdapter> combinedExecutionData = new IdentityHashMap<>();
       Collection<String> combinedNonCodeResources = new HashSet<>();
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/percentage/JacocoCoverageCalculator.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/percentage/JacocoCoverageCalculator.java
@@ -9,6 +9,7 @@
 import datadog.trace.civisibility.ipc.ModuleCoverageDataJacoco;
 import datadog.trace.civisibility.ipc.SignalResponse;
 import datadog.trace.civisibility.ipc.SignalType;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import datadog.trace.civisibility.source.index.RepoIndex;
 import datadog.trace.civisibility.source.index.RepoIndexProvider;
 import datadog.trace.util.Strings;
@@ -312,13 +313,19 @@ private RepoIndexFileLocator(RepoIndex repoIndex, String repoRoot) {
 
     @Override
     protected InputStream getSourceStream(String path) throws IOException {
-      String relativePath = repoIndex.getSourcePath(path);
-      if (relativePath == null) {
+      try {
+        String relativePath = repoIndex.getSourcePath(path);
+        if (relativePath == null) {
+          return null;
+        }
+        String absolutePath =
+            repoRoot + (!repoRoot.endsWith(File.separator) ? File.separator : "") + relativePath;
+        return new BufferedInputStream(Files.newInputStream(Paths.get(absolutePath)));
+
+      } catch (SourceResolutionException e) {
+        LOGGER.debug("Could not resolve source for path {}", path, e);
         return null;
       }
-      String absolutePath =
-          repoRoot + (!repoRoot.endsWith(File.separator) ? File.separator : "") + relativePath;
-      return new BufferedInputStream(Files.newInputStream(Paths.get(absolutePath)));
     }
   }
 
@@ -354,7 +361,13 @@ private long getMergedCoveragePercentage(IBundleCoverage coverageBundle) {
         String fileName = sourceFile.getName();
         String pathRelativeToSourceRoot =
             (Strings.isNotBlank(packageName) ? packageName + "/" : "") + fileName;
-        String pathRelativeToIndexRoot = repoIndex.getSourcePath(pathRelativeToSourceRoot);
+        String pathRelativeToIndexRoot;
+        try {
+          pathRelativeToIndexRoot = repoIndex.getSourcePath(pathRelativeToSourceRoot);
+        } catch (SourceResolutionException e) {
+          LOGGER.debug("Could not resolve source for path {}", pathRelativeToSourceRoot, e);
+          continue;
+        }
 
         BitSet sourceFileCoveredLines = getCoveredLines(sourceFile);
         // backendCoverageData contains data for all modules in the repo,
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestImpl.java
@@ -32,6 +32,7 @@
 import datadog.trace.civisibility.decorator.TestDecorator;
 import datadog.trace.civisibility.source.MethodLinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import java.lang.reflect.Method;
 import java.util.Collection;
 import java.util.function.Consumer;
@@ -147,8 +148,14 @@ private void populateSourceDataTags(
       return;
     }
 
-    String sourcePath = sourcePathResolver.getSourcePath(testClass);
-    if (sourcePath == null || sourcePath.isEmpty()) {
+    String sourcePath;
+    try {
+      sourcePath = sourcePathResolver.getSourcePath(testClass);
+      if (sourcePath == null || sourcePath.isEmpty()) {
+        return;
+      }
+    } catch (SourceResolutionException e) {
+      log.debug("Could not populate source path for {}", testClass, e);
       return;
     }
 
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java
@@ -19,13 +19,18 @@
 import datadog.trace.civisibility.decorator.TestDecorator;
 import datadog.trace.civisibility.source.MethodLinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import datadog.trace.civisibility.utils.SpanUtils;
 import java.lang.reflect.Method;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class TestSuiteImpl implements DDTestSuite {
 
+  private static final Logger log = LoggerFactory.getLogger(TestSuiteImpl.class);
+
   private final AgentSpan.Context moduleSpanContext;
   private final AgentSpan span;
   private final String moduleName;
@@ -106,13 +111,9 @@ public TestSuiteImpl(
     span.setTag(Tags.TEST_STATUS, TestStatus.skip);
 
     this.testClass = testClass;
-    if (this.testClass != null) {
-      if (config.isCiVisibilitySourceDataEnabled()) {
-        String sourcePath = sourcePathResolver.getSourcePath(testClass);
-        if (sourcePath != null && !sourcePath.isEmpty()) {
-          span.setTag(Tags.TEST_SOURCE_FILE, sourcePath);
-        }
-      }
+
+    if (config.isCiVisibilitySourceDataEnabled()) {
+      populateSourceDataTags(testClass, sourcePathResolver);
     }
 
     testDecorator.afterStart(span);
@@ -129,6 +130,20 @@ public TestSuiteImpl(
     }
   }
 
+  private void populateSourceDataTags(Class<?> testClass, SourcePathResolver sourcePathResolver) {
+    if (this.testClass == null) {
+      return;
+    }
+    try {
+      String sourcePath = sourcePathResolver.getSourcePath(testClass);
+      if (sourcePath != null && !sourcePath.isEmpty()) {
+        span.setTag(Tags.TEST_SOURCE_FILE, sourcePath);
+      }
+    } catch (SourceResolutionException e) {
+      log.debug("Could not populate source path for {}", testClass, e);
+    }
+  }
+
   @Override
   public void setTag(String key, Object value) {
     span.setTag(key, value);
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/BestEffortSourcePathResolver.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/BestEffortSourcePathResolver.java
@@ -13,7 +13,7 @@ public BestEffortSourcePathResolver(SourcePathResolver... delegates) {
 
   @Nullable
   @Override
-  public String getSourcePath(@Nonnull Class<?> c) {
+  public String getSourcePath(@Nonnull Class<?> c) throws SourceResolutionException {
     for (SourcePathResolver delegate : delegates) {
       String sourcePath = delegate.getSourcePath(c);
       if (sourcePath != null) {
@@ -25,7 +25,7 @@ public String getSourcePath(@Nonnull Class<?> c) {
 
   @Nullable
   @Override
-  public String getResourcePath(@Nullable String relativePath) {
+  public String getResourcePath(@Nullable String relativePath) throws SourceResolutionException {
     for (SourcePathResolver delegate : delegates) {
       String resourcePath = delegate.getResourcePath(relativePath);
       if (resourcePath != null) {
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/SourcePathResolver.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/SourcePathResolver.java
@@ -9,12 +9,12 @@ public interface SourcePathResolver {
    *     root. {@code null} is returned if the path could not be resolved
    */
   @Nullable
-  String getSourcePath(@Nonnull Class<?> c);
+  String getSourcePath(@Nonnull Class<?> c) throws SourceResolutionException;
 
   /**
    * @param relativePath Path to a resource in current run's repository, relative to a resource root
    * @return Path relative to repository root
    */
   @Nullable
-  String getResourcePath(@Nullable String relativePath);
+  String getResourcePath(@Nullable String relativePath) throws SourceResolutionException;
 }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/SourceResolutionException.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/SourceResolutionException.java
@@ -0,0 +1,8 @@
+package datadog.trace.civisibility.source;
+
+public class SourceResolutionException extends Exception {
+
+  public SourceResolutionException(String message) {
+    super(message);
+  }
+}
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java
@@ -1,7 +1,9 @@
 package datadog.trace.civisibility.source.index;
 
+import datadog.trace.api.Config;
 import datadog.trace.api.civisibility.domain.Language;
 import datadog.trace.civisibility.ipc.Serializer;
+import datadog.trace.civisibility.source.SourceResolutionException;
 import datadog.trace.civisibility.source.Utils;
 import datadog.trace.util.ClassNameTrie;
 import java.io.ByteArrayInputStream;
@@ -11,6 +13,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -23,16 +26,25 @@ public class RepoIndex {
 
   static final RepoIndex EMPTY =
       new RepoIndex(
-          ClassNameTrie.Builder.EMPTY_TRIE, Collections.emptyList(), Collections.emptyList());
+          ClassNameTrie.Builder.EMPTY_TRIE,
+          Collections.emptyList(),
+          Collections.emptyList(),
+          Collections.emptyList());
 
   private static final Logger log = LoggerFactory.getLogger(RepoIndex.class);
 
   private final ClassNameTrie trie;
+  private final Collection<String> duplicateTrieKeys;
   private final List<SourceRoot> sourceRoots;
   private final List<String> rootPackages;
 
-  RepoIndex(ClassNameTrie trie, List<SourceRoot> sourceRoots, List<String> rootPackages) {
+  RepoIndex(
+      ClassNameTrie trie,
+      Collection<String> duplicateTrieKeys,
+      List<SourceRoot> sourceRoots,
+      List<String> rootPackages) {
     this.trie = trie;
+    this.duplicateTrieKeys = duplicateTrieKeys;
     this.sourceRoots = sourceRoots;
     this.rootPackages = rootPackages;
   }
@@ -42,7 +54,7 @@ public List<String> getRootPackages() {
   }
 
   @Nullable
-  public String getSourcePath(@Nonnull Class<?> c) {
+  public String getSourcePath(@Nonnull Class<?> c) throws SourceResolutionException {
     String topLevelClassName = Utils.stripNestedClassNames(c.getName());
     String sourcePath = doGetSourcePath(topLevelClassName);
     return sourcePath != null ? sourcePath : getFallbackSourcePath(c);
@@ -54,7 +66,7 @@ public String getSourcePath(@Nonnull Class<?> c) {
    * retrieved from the bytecode.
    */
   @Nullable
-  private String getFallbackSourcePath(@Nonnull Class<?> c) {
+  private String getFallbackSourcePath(@Nonnull Class<?> c) throws SourceResolutionException {
     try {
       String fileName = Utils.getFileName(c);
       if (fileName == null) {
@@ -75,7 +87,8 @@ private String getFallbackSourcePath(@Nonnull Class<?> c) {
   }
 
   @Nullable
-  public String getSourcePath(@Nullable String pathRelativeToSourceRoot) {
+  public String getSourcePath(@Nullable String pathRelativeToSourceRoot)
+      throws SourceResolutionException {
     if (pathRelativeToSourceRoot == null) {
       return null;
     }
@@ -84,7 +97,13 @@ public String getSourcePath(@Nullable String pathRelativeToSourceRoot) {
   }
 
   @Nullable
-  private String doGetSourcePath(String key) {
+  private String doGetSourcePath(String key) throws SourceResolutionException {
+    if (Config.get().isCiVisibilityRepoIndexDuplicateKeyCheckEnabled()) {
+      if (!duplicateTrieKeys.isEmpty() && duplicateTrieKeys.contains(key)) {
+        throw new SourceResolutionException("There are multiple repo index entries for " + key);
+      }
+    }
+
     int sourceRootIdx = trie.apply(key);
     if (sourceRootIdx < 0) {
       log.debug("Could not find source root for {}", key);
@@ -110,6 +129,7 @@ public ByteBuffer serialize() {
 
     Serializer s = new Serializer();
     s.write(serializedTrie);
+    s.write(duplicateTrieKeys);
     s.write(sourceRoots, SourceRoot::serialize);
     s.write(rootPackages);
     return s.flush();
@@ -132,9 +152,10 @@ public static RepoIndex deserialize(ByteBuffer buffer) {
       }
     }
 
+    Collection<String> duplicateTrieKeys = Serializer.readSet(buffer, Serializer::readString);
     List<SourceRoot> sourceRoots = Serializer.readList(buffer, SourceRoot::deserialize);
     List<String> rootPackages = Serializer.readStringList(buffer);
-    return new RepoIndex(trie, sourceRoots, rootPackages);
+    return new RepoIndex(trie, duplicateTrieKeys, sourceRoots, rootPackages);
   }
 
   static final class SourceRoot {
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexBuilder.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexBuilder.java
@@ -13,8 +13,10 @@
 import java.nio.file.Path;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.EnumSet;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.slf4j.Logger;
@@ -94,6 +96,8 @@ private static final class RepoIndexingFileVisitor implements FileVisitor<Path>
     private final PackageResolver packageResolver;
     private final ResourceResolver resourceResolver;
     private final ClassNameTrie.Builder trieBuilder;
+    private final Map<String, String> trieKeyToPath;
+    private final Collection<String> duplicateTrieKeys;
     private final Map<RepoIndex.SourceRoot, Integer> sourceRoots;
     private final PackageTree packageTree;
     private final RepoIndexingStats indexingStats;
@@ -109,6 +113,8 @@ private RepoIndexingFileVisitor(
       this.resourceResolver = resourceResolver;
       this.repoRoot = repoRoot;
       trieBuilder = new ClassNameTrie.Builder();
+      trieKeyToPath = new HashMap<>();
+      duplicateTrieKeys = new HashSet<>();
       sourceRoots = new HashMap<>();
       packageTree = new PackageTree(config);
       indexingStats = new RepoIndexingStats();
@@ -161,6 +167,12 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
           if (!relativePath.isEmpty()) {
             String key = Utils.toTrieKey(relativePath);
             trieBuilder.put(key, sourceRootIdx);
+
+            String existingPath = trieKeyToPath.put(key, file.toString());
+            if (existingPath != null) {
+              log.warn("Duplicate repo index key: {} - {}", existingPath, file);
+              duplicateTrieKeys.add(key);
+            }
           }
         }
       } catch (Exception e) {
@@ -221,7 +233,8 @@ public RepoIndex getIndex() {
         roots[e.getValue()] = e.getKey();
       }
 
-      return new RepoIndex(trieBuilder.buildTrie(), Arrays.asList(roots), packageTree.asList());
+      return new RepoIndex(
+          trieBuilder.buildTrie(), duplicateTrieKeys, Arrays.asList(roots), packageTree.asList());
     }
   }
 
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolver.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolver.java
diff --git a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy
diff --git a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexTest.groovy b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexTest.groovy
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/CiVisibilityConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/CiVisibilityConfig.java
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
