[clangd] Fix race in FileIndex that sometimes temporarily lost updates.

[sammccall]

Summary

FileIndex was built out of threadsafe components, so update() didn't have data
races, but wasn't actually correct.

[kadircet]

LGTM, thanks!

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

270

Bumping the version in update and just exposing it in buildIndex feels more natural, WDYT?

Might even enable us to get rid of some redundant resets (not that it is expensive ..), imagine a sequencing like:

• update1
• update2
• buildIndex1
• buildIndex2

even though both build the same index, they got different version numbers hence both will reset the current index.

clang-tools-extra/clangd/index/Symbol.h

190

i suppose this is for testing::SizeIs right?

[kadircet]

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

97

oh and also maybe use size_t ?

as there's only one (well three actually, if you count both preamble and main-file index separately, and one in background-index) instance of FileSymbols for a single clangd instance, hence we could in theory bump this number with every keystroke :/

I suppose 4 bytes is also far from any practical limits, but pushing it as far as possible wouldn't hurt?

[sammccall]

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

270

Yeah, that makes sense.

clang-tools-extra/clangd/index/Symbol.h

190

Right, gtest is weird. Could live without it, seems harmless though.

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

public:
   /// during merging.
   void update(llvm::StringRef Key, std::unique_ptr<SymbolSlab> Symbols,
               std::unique_ptr<RefSlab> Refs,
               std::unique_ptr<RelationSlab> Relations, bool CountReferences);
 
   /// The index keeps the slabs alive.
   /// Will count Symbol::References based on number of references in the main
   /// files, while building the index with DuplicateHandling::Merge option.
+  /// Version is populated with an increasing sequence counter.
   std::unique_ptr<SymbolIndex>
   buildIndex(IndexType,
-             DuplicateHandling DuplicateHandle = DuplicateHandling::PickOne);
+             DuplicateHandling DuplicateHandle = DuplicateHandling::PickOne,
+             size_t *Version = nullptr);
 
 private:
   struct RefSlabAndCountReferences {
     std::shared_ptr<RefSlab> Slab;
     bool CountReferences = false;
   };
   mutable std::mutex Mutex;
 
+  size_t Version = 0;
   llvm::StringMap<std::shared_ptr<SymbolSlab>> SymbolsSnapshot;
   llvm::StringMap<RefSlabAndCountReferences> RefsSnapshot;
   llvm::StringMap<std::shared_ptr<RelationSlab>> RelatiosSnapshot;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
private:
   // Contains information from each file's main AST.
   // These are updated frequently (on file change), but are relatively small.
   // Mostly contains:
   //  - refs to symbols declared in the preamble and referenced from main
   //  - symbols declared both in the main file and the preamble
   // (Note that symbols *only* in the main file are not indexed).
   FileSymbols MainFileSymbols;
   SwapIndex MainFileIndex;
+
+  // While both the FileIndex and SwapIndex are threadsafe, we need to track
+  // versions to ensure that we don't overwrite newer indexes with older ones.
+  std::mutex UpdateIndexMu;
+  unsigned MainIndexVersion = 0;
+  unsigned PreambleIndexVersion = 0;
 };
 
 using SlabTuple = std::tuple<SymbolSlab, RefSlab, RelationSlab>;
 
 /// Retrieves symbols and refs of local top level decls in \p AST (i.e.
 /// `AST.getLocalTopLevelDecls()`).
 /// Exposed to assist in unit tests.
 SlabTuple indexMainDecls(ParsedAST &AST);

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

 }
 
 void FileSymbols::update(llvm::StringRef Key,
                          std::unique_ptr<SymbolSlab> Symbols,
                          std::unique_ptr<RefSlab> Refs,
                          std::unique_ptr<RelationSlab> Relations,
                          bool CountReferences) {
   std::lock_guard<std::mutex> Lock(Mutex);
+  ++Version;
   if (!Symbols)
     SymbolsSnapshot.erase(Key);
   else
     SymbolsSnapshot[Key] = std::move(Symbols);
   if (!Refs) {
     RefsSnapshot.erase(Key);
   } else {
     RefSlabAndCountReferences Item;
     Item.CountReferences = CountReferences;
     Item.Slab = std::move(Refs);
     RefsSnapshot[Key] = std::move(Item);
   }
   if (!Relations)
     RelatiosSnapshot.erase(Key);
   else
     RelatiosSnapshot[Key] = std::move(Relations);
 }
 
 std::unique_ptr<SymbolIndex>
-FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle) {
+FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle,
+                        size_t *Version) {
   std::vector<std::shared_ptr<SymbolSlab>> SymbolSlabs;
   std::vector<std::shared_ptr<RefSlab>> RefSlabs;
   std::vector<std::shared_ptr<RelationSlab>> RelationSlabs;
   std::vector<RefSlab *> MainFileRefs;
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     for (const auto &FileAndSymbols : SymbolsSnapshot)
       SymbolSlabs.push_back(FileAndSymbols.second);
     for (const auto &FileAndRefs : RefsSnapshot) {
       RefSlabs.push_back(FileAndRefs.second.Slab);
       if (FileAndRefs.second.CountReferences)
         MainFileRefs.push_back(RefSlabs.back().get());
     }
     for (const auto &FileAndRelations : RelatiosSnapshot)
       RelationSlabs.push_back(FileAndRelations.second);
+
+    if (Version)
+      *Version = this->Version;
   }
   std::vector<const Symbol *> AllSymbols;
   std::vector<Symbol> SymsStorage;
   switch (DuplicateHandle) {
   case DuplicateHandling::Merge: {
     llvm::DenseMap<SymbolID, Symbol> Merged;
     for (const auto &Slab : SymbolSlabs) {
       for (const auto &Sym : *Slab) {
for (auto Uri : ShardedIndex.getAllSources()) {
     // exist.
     assert(IF);
     PreambleSymbols.update(
         Uri, std::make_unique<SymbolSlab>(std::move(*IF->Symbols)),
         std::make_unique<RefSlab>(),
         std::make_unique<RelationSlab>(std::move(*IF->Relations)),
         /*CountReferences=*/false);
   }
-  PreambleIndex.reset(
+  size_t IndexVersion = 0;
+  auto NewIndex =
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
-                                 DuplicateHandling::PickOne));
-  vlog("Build dynamic index for header symbols with estimated memory usage of "
+                                 DuplicateHandling::PickOne, &IndexVersion);
+  {
+    std::lock_guard<std::mutex> Lock(UpdateIndexMu);
+    if (IndexVersion <= PreambleIndexVersion) {
+      // We lost the race, some other thread built a later version.
+      return;
+    }
+    PreambleIndexVersion = IndexVersion;
+    PreambleIndex.reset(std::move(NewIndex));
+    vlog(
+        "Build dynamic index for header symbols with estimated memory usage of "
         "{0} bytes",
         PreambleIndex.estimateMemoryUsage());
   }
+}
 
 void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
   auto Contents = indexMainDecls(AST);
   MainFileSymbols.update(
       Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Contents))),
       std::make_unique<RefSlab>(std::move(std::get<1>(Contents))),
       std::make_unique<RelationSlab>(std::move(std::get<2>(Contents))),
       /*CountReferences=*/true);
-  MainFileIndex.reset(
-      MainFileSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge));
-  vlog("Build dynamic index for main-file symbols with estimated memory usage "
+  size_t IndexVersion = 0;
+  auto NewIndex = MainFileSymbols.buildIndex(
+      IndexType::Light, DuplicateHandling::Merge, &IndexVersion);
+  {
+    std::lock_guard<std::mutex> Lock(UpdateIndexMu);
+    if (IndexVersion <= MainIndexVersion) {
+      // We lost the race, some other thread built a later version.
+      return;
+    }
+    MainIndexVersion = IndexVersion;
+    MainFileIndex.reset(std::move(NewIndex));
+    vlog(
+        "Build dynamic index for main-file symbols with estimated memory usage "
         "of {0} bytes",
         MainFileIndex.estimateMemoryUsage());
   }
+}
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/Symbol.h

public:
   using value_type = Symbol;
 
   SymbolSlab() = default;
 
   const_iterator begin() const { return Symbols.begin(); }
   const_iterator end() const { return Symbols.end(); }
   const_iterator find(const SymbolID &SymID) const;
 
-  size_t size() const { return Symbols.size(); }
+  using size_type = size_t;
+  size_type size() const { return Symbols.size(); }
   bool empty() const { return Symbols.empty(); }
   // Estimates the total memory usage.
   size_t bytes() const {
     return sizeof(*this) + Arena.getTotalMemory() +
            Symbols.capacity() * sizeof(Symbol);
   }
 
   /// SymbolSlab::Builder is a mutable container that can 'freeze' to

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

 #include "URI.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "index/Ref.h"
 #include "index/Relation.h"
 #include "index/Serialization.h"
 #include "index/Symbol.h"
+#include "support/Threading.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <utility>
TEST(FileIndexTest, StalePreambleSymbolsDeleted) {
   File.HeaderCode = "int b;";
   AST = File.build();
   M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
                    AST.getASTContext(), AST.getPreprocessorPtr(),
                    AST.getCanonicalIncludes());
   EXPECT_THAT(runFuzzyFind(M, ""), UnorderedElementsAre(QName("b")));
 }
 
+// Verifies that concurrent calls to updateMain don't "lose" any updates.
+TEST(FileIndexTest, Threadsafety) {
+  FileIndex M;
+  Notification Go;
+
+  constexpr int Count = 10;
+  {
+    // Set up workers to concurrently call updateMain() with separate files.
+    AsyncTaskRunner Pool;
+    for (unsigned I = 0; I < Count; ++I) {
+      auto TU = TestTU::withCode(llvm::formatv("int xxx{0};", I).str());
+      TU.Filename = llvm::formatv("x{0}.c", I).str();
+      Pool.runAsync(TU.Filename, [&, Filename(testPath(TU.Filename)),
+                                  AST(TU.build())]() mutable {
+        Go.wait();
+        M.updateMain(Filename, AST);
+      });
+    }
+    // On your marks, get set...
+    Go.notify();
+  }
+
+  EXPECT_THAT(runFuzzyFind(M, "xxx"), ::testing::SizeIs(Count));
+}
+
 TEST(FileShardedIndexTest, Sharding) {
   auto AHeaderUri = URI::create(testPath("a.h")).toString();
   auto BHeaderUri = URI::create(testPath("b.h")).toString();
   auto BSourceUri = URI::create(testPath("b.cc")).toString();
 
   auto Sym1 = symbol("1");
   Sym1.CanonicalDeclaration.FileURI = AHeaderUri.c_str();