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

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

Reviewers: kadircet

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D82891

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

@@ -229,6 +229,7 @@ void FileSymbols::update(llvm::StringRef Key,
                          std::unique_ptr<RelationSlab> Relations,
                          bool CountReferences) {
   std::lock_guard<std::mutex> Lock(Mutex);
+  ++Version;
   if (!Symbols)
     SymbolsSnapshot.erase(Key);
   else
@@ -248,7 +249,8 @@ void FileSymbols::update(llvm::StringRef Key,
 }
 
 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;
@@ -264,6 +266,9 @@ FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle) {
     }
     for (const auto &FileAndRelations : RelatiosSnapshot)
       RelationSlabs.push_back(FileAndRelations.second);
+
+    if (Version)
+      *Version = this->Version;
   }
   std::vector<const Symbol *> AllSymbols;
   std::vector<Symbol> SymsStorage;
@@ -390,12 +395,23 @@ void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
         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 "
-       "{0} bytes",
-       PreambleIndex.estimateMemoryUsage());
+                                 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) {
@@ -405,11 +421,22 @@ void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
       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 "
-       "of {0} bytes",
-       MainFileIndex.estimateMemoryUsage());
+  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

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

@@ -81,9 +81,11 @@ public:
   /// 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 {
@@ -92,6 +94,7 @@ private:
   };
   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;
@@ -136,6 +139,12 @@ private:
   // (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>;

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

@@ -186,7 +186,8 @@ public:
   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 {

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

@@ -22,6 +22,7 @@
 #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"
@@ -509,6 +510,31 @@ TEST(FileIndexTest, StalePreambleSymbolsDeleted) {
   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();