[clangd] Avoid relations being overwritten in a header shard

[nridge]

Summary

Fixes https://github.com/clangd/clangd/issues/510

[hokein]

+@kadircet, he is tracking on this -- we had some discussion internally last week, but we don't reply on that thread yet (sorry).

[nridge]

Address review comments

[kadircet]

clang-tools-extra/clangd/index/FileIndex.cpp

152

instead of (or in addition to) providing the link, can we give a short summary? e.g.

Subject and/or Object files might be part of multiple TUs. In such cases there will be a race
and last TU to write the shard will win and all the other relations will be lost. We are storing
relations in both places, as we do for symbols, to reduce such races. Note that they might still
happen if same translation unit produces different relations under different configurations
but that's something clangd doesn't handle in general.

354

can you move this outside of the for loop, so that we do it only once.

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

232

do we still need this test?

396

s/ShardSource/ShardHeader

clang-tools-extra/clangd/unittests/FileIndexTests.cpp

571–573

maybe mention that Sym1 belongs to a.h in the comment. stored in a.h (where Sym1 is stored) even tho the reference is dangling as Sym3 is unknown

[nridge]

Address review comments

[kadircet]

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

232

this one was marked as resolved but i still don't see the reasoning. can we test this in fileindextests instead?

we already test the sharding logic, we need to test the merging logic now. can we rather test it at FileSymbols layer directly? or is there something extra that i am misisng?

[nridge]

Sorry, I had a response to that comment but accidentally left it as unsubmitted...

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

232

I think it's still useful, yes. If someone makes a change to the way relations are stored in the future, they could regress this use case without a test specifically for it.

[kadircet]

thanks! (and sorry for taking too long on this one.)

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

232

okay, i still think it is better to test on FileSymbols layer but not that important.

[nridge]

Slightly simplified test (base class does not need to be a template to trigger the issue)

[nridge]

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

232

The reason I'd like to have a test at this layer is so that we have a test that closely approximates the steps to reproduce: create files with certain contents and inclusion relationships between them, index them, and perform a particular query. I feel like internal abstractions like FileSymbols are liable to change over time through refactorings, and so tests written against them are less robust.

Note, I don't think using BackgroundIndex is important for the kind of test I'd like; FileIndex would be fine too. However, I could not see how to use FileIndex with files that contain #include statements (didn't see a way to provide a MockFS or another way to get the includes to be resolved).

That said, I did try to write an additional test using FileSymbols, but I found that:

• it only exercises the merging logic, not the sharding logic
• I couldn't get the test to fail even if I omit the merging logic of this patch, because both MemIndex and Dex build a DenseMap for relations and therefore end up implicitly deduplicating anyways

clang-tools-extra/clangd/index/FileIndex.cpp

 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
+#include <algorithm>
 #include <memory>
 #include <tuple>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
if (Index.Refs) {
     for (const auto &SymRefs : *Index.Refs) {
       for (const auto &R : SymRefs.second) {
         const auto It = Shards.try_emplace(R.Location.FileURI);
         It.first->getValue().Refs.insert(&R);
         RefToSymID[&R] = SymRefs.first;
       }
     }
   }
-  // Attribute relations to the file declaraing their Subject as Object might
-  // not have been indexed, see SymbolCollector::processRelations for details.
+  // The Subject and/or Object shards might be part of multiple TUs. In
+  // such cases there will be a race and the last TU to write the shard
+  // will win and all the other relations will be lost. To avoid this,
+  // we store relations in both shards. A race might still happen if the
+  // same translation unit produces different relations under different
+  // configurations, but that's something clangd doesn't handle in general.
   if (Index.Relations) {
     for (const auto &R : *Index.Relations) {
       // FIXME: RelationSlab shouldn't contain dangling relations.
-      if (auto *File = SymbolIDToFile.lookup(R.Subject))
-        File->Relations.insert(&R);
+      FileShard *SubjectFile = SymbolIDToFile.lookup(R.Subject);
+      FileShard *ObjectFile = SymbolIDToFile.lookup(R.Object);
+      if (SubjectFile)
+        SubjectFile->Relations.insert(&R);
+      if (ObjectFile && ObjectFile != SubjectFile)
+        ObjectFile->Relations.insert(&R);
     }
   }
   // Store only the direct includes of a file in a shard.
   if (Index.Sources) {
     const auto &FullGraph = *Index.Sources;
     for (const auto &It : FullGraph) {
       auto ShardIt = Shards.try_emplace(It.first());
       ShardIt.first->getValue().IG = getSubGraph(It.first(), FullGraph);
for (auto &Sym : MergedRefs) {
                               SymRefs.size()));
     }
   }
 
   std::vector<Relation> AllRelations;
   for (const auto &RelationSlab : RelationSlabs) {
     for (const auto &R : *RelationSlab)
       AllRelations.push_back(R);
   }
+  // Sort relations and remove duplicates that could arise due to
+  // relations being stored in both the shards containing their
+  // subject and object.
+  llvm::sort(AllRelations);
+  AllRelations.erase(std::unique(AllRelations.begin(), AllRelations.end()),
+                     AllRelations.end());
 
   size_t StorageSize =
       RefsStorage.size() * sizeof(Ref) + SymsStorage.size() * sizeof(Symbol);
   for (const auto &Slab : SymbolSlabs)
     StorageSize += Slab->bytes();
   for (const auto &RefSlab : RefSlabs)
     StorageSize += RefSlab->bytes();
   for (const auto &RelationSlab : RelationSlabs)

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

EXPECT_THAT(getRefs(Idx, Common.ID),
               RefsAre({FileURI("unittest:///root/A.h"),
                        FileURI("unittest:///root/A.cc"),
                        FileURI("unittest:///root/B.cc"),
                        FileURI("unittest:///root/B.cc"),
                        FileURI("unittest:///root/B.cc"),
                        FileURI("unittest:///root/B.cc")}));
 }
 
+TEST_F(BackgroundIndexTest, RelationsMultiFile) {
+  MockFS FS;
+  FS.Files[testPath("root/Base.h")] = "class Base {};";
+  FS.Files[testPath("root/A.cc")] = R"cpp(
+    #include "Base.h"
+    class A : public Base {};
+  )cpp";
+  FS.Files[testPath("root/B.cc")] = R"cpp(
+    #include "Base.h"
+    class B : public Base {};
+  )cpp";
+
+  llvm::StringMap<std::string> Storage;
+  size_t CacheHits = 0;
+  MemoryShardStorage MSS(Storage, CacheHits);
+  OverlayCDB CDB(/*Base=*/nullptr);
+  BackgroundIndex Index(FS, CDB, [&](llvm::StringRef) { return &MSS; },
+                        /*Opts=*/{});
+
+  tooling::CompileCommand Cmd;
+  Cmd.Filename = testPath("root/A.cc");
+  Cmd.Directory = testPath("root");
+  Cmd.CommandLine = {"clang++", Cmd.Filename};
+  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
+  ASSERT_TRUE(Index.blockUntilIdleForTest());
+
+  Cmd.Filename = testPath("root/B.cc");
+  Cmd.CommandLine = {"clang++", Cmd.Filename};
+  CDB.setCompileCommand(testPath("root/B.cc"), Cmd);
+  ASSERT_TRUE(Index.blockUntilIdleForTest());
+
+  auto HeaderShard = MSS.loadShard(testPath("root/Base.h"));
+  EXPECT_NE(HeaderShard, nullptr);
+  SymbolID Base = findSymbol(*HeaderShard->Symbols, "Base").ID;
+
+  RelationsRequest Req;
+  Req.Subjects.insert(Base);
+  Req.Predicate = RelationKind::BaseOf;
+  uint32_t Results = 0;
+  Index.relations(Req, [&](const SymbolID &, const Symbol &) { ++Results; });
+  EXPECT_EQ(Results, 2u);
+}
+
 TEST_F(BackgroundIndexTest, MainFileRefs) {
   MockFS FS;
   FS.Files[testPath("root/A.h")] = R"cpp(
       void header_sym();
       )cpp";
   FS.Files[testPath("root/A.cc")] =
       "#include \"A.h\"\nstatic void main_sym() { (void)header_sym; }";
 
TEST_F(BackgroundIndexTest, ShardStorageTest) {
   auto ShardSource = MSS.loadShard(testPath("root/A.cc"));
   EXPECT_NE(ShardSource, nullptr);
   EXPECT_THAT(*ShardSource->Symbols,
               UnorderedElementsAre(Named("g"), Named("B_CC")));
   for (const auto &Ref : *ShardSource->Refs)
     EXPECT_THAT(Ref.second,
                 UnorderedElementsAre(FileURI("unittest:///root/A.cc")));
 
-  // The BaseOf relationship between A_CC and B_CC is stored in the file
-  // containing the definition of the subject (A_CC)
+  // The BaseOf relationship between A_CC and B_CC is stored in both the file
+  // containing the definition of the subject (A_CC) and the file containing
+  // the definition of the object (B_CC).
   SymbolID A = findSymbol(*ShardHeader->Symbols, "A_CC").ID;
   SymbolID B = findSymbol(*ShardSource->Symbols, "B_CC").ID;
   EXPECT_THAT(*ShardHeader->Relations,
               UnorderedElementsAre(Relation{A, RelationKind::BaseOf, B}));
-  // (and not in the file containing the definition of the object (B_CC)).
-  EXPECT_EQ(ShardSource->Relations->size(), 0u);
+  EXPECT_THAT(*ShardSource->Relations,
+              UnorderedElementsAre(Relation{A, RelationKind::BaseOf, B}));
 }
 
 TEST_F(BackgroundIndexTest, DirectIncludesTest) {
   MockFS FS;
   FS.Files[testPath("root/B.h")] = "";
   FS.Files[testPath("root/A.h")] = R"cpp(
       #include "B.h"
       void common();

clang-tools-extra/clangd/unittests/FileIndexTests.cpp

TEST(FileShardedIndexTest, Sharding) {
 
   auto Sym1 = symbol("1");
   Sym1.CanonicalDeclaration.FileURI = AHeaderUri.c_str();
 
   auto Sym2 = symbol("2");
   Sym2.CanonicalDeclaration.FileURI = BHeaderUri.c_str();
   Sym2.Definition.FileURI = BSourceUri.c_str();
 
+  auto Sym3 = symbol("3"); // not stored
+
   IndexFileIn IF;
   {
     SymbolSlab::Builder B;
     // Should be stored in only a.h
     B.insert(Sym1);
     // Should be stored in both b.h and b.cc
     B.insert(Sym2);
     IF.Symbols.emplace(std::move(B).build());
   }
   {
     // Should be stored in b.cc
     IF.Refs.emplace(std::move(*refSlab(Sym1.ID, BSourceUri.c_str())));
   }
   {
     RelationSlab::Builder B;
-    // Should be stored in a.h
+    // Should be stored in a.h and b.h
     B.insert(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID});
-    // Should be stored in b.h
+    // Should be stored in a.h and b.h
     B.insert(Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID});
-    // Dangling relation should be dropped.
-    B.insert(Relation{symbol("3").ID, RelationKind::BaseOf, Sym1.ID});
+    // Should be stored in a.h (where Sym1 is stored) even though
+    // the relation is dangling as Sym3 is unknown.
+    B.insert(Relation{Sym3.ID, RelationKind::BaseOf, Sym1.ID});
     IF.Relations.emplace(std::move(B).build());
   }
 
   IF.Sources.emplace();
   IncludeGraph &IG = *IF.Sources;
   {
     // b.cc includes b.h
     auto &Node = IG[BSourceUri];
TEST(FileShardedIndexTest, Sharding) {
 
   {
     auto Shard = ShardedIndex.getShard(AHeaderUri);
     ASSERT_TRUE(Shard);
     EXPECT_THAT(*Shard->Symbols, UnorderedElementsAre(QName("1")));
     EXPECT_THAT(*Shard->Refs, IsEmpty());
     EXPECT_THAT(
         *Shard->Relations,
-        UnorderedElementsAre(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID}));
+        UnorderedElementsAre(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID},
+                             Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID},
+                             Relation{Sym3.ID, RelationKind::BaseOf, Sym1.ID}));
     ASSERT_THAT(Shard->Sources->keys(), UnorderedElementsAre(AHeaderUri));
     EXPECT_THAT(Shard->Sources->lookup(AHeaderUri).DirectIncludes, IsEmpty());
     EXPECT_TRUE(Shard->Cmd.hasValue());
   }
   {
     auto Shard = ShardedIndex.getShard(BHeaderUri);
     ASSERT_TRUE(Shard);
     EXPECT_THAT(*Shard->Symbols, UnorderedElementsAre(QName("2")));
     EXPECT_THAT(*Shard->Refs, IsEmpty());
     EXPECT_THAT(
         *Shard->Relations,
-        UnorderedElementsAre(Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID}));
+        UnorderedElementsAre(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID},
+                             Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID}));
     ASSERT_THAT(Shard->Sources->keys(),
                 UnorderedElementsAre(BHeaderUri, AHeaderUri));
     EXPECT_THAT(Shard->Sources->lookup(BHeaderUri).DirectIncludes,
                 UnorderedElementsAre(AHeaderUri));
     EXPECT_TRUE(Shard->Cmd.hasValue());
   }
   {
     auto Shard = ShardedIndex.getShard(BSourceUri);