[clangd] Rewrite of logic to rebuild the background index serving structures.

[sammccall]

Summary

Previously it was rebuilding every 5s by default, which was much too frequent
in the long run - the goal was to provide an early build. There were also some
bugs. There were also some bugs, and a dedicated thread was used in production
but not tested.

• rebuilds are triggered by #TUs built, rather than time. This should scale more sensibly to fast vs slow machines.
• there are two separate indexed-TU thresholds to trigger index build: 5 TUs for the first build, 100 for subsequent rebuilds.
• rebuild is always done on the regular indexing threads, and is affected by blockUntilIdle. This means unit/lit tests run the production configuration.
• fixed a bug where we'd rebuild after attempting to load shards, even if there were no shards.
• the BackgroundIndexTests don't really test the subtleties of the rebuild policy (for determinism, we call blockUntilIdle, so rebuild-on-idle is enough to pass the tests). Instead, we expose the rebuilder as a separate class and have fine-grained tests for it.

[sammccall]

Fix broken format string

[kadircet]

LGTM

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

185

clang-format

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

142

• do we really want to make 5 and 100 hardcoded in the comments?
• what about moving this class and its implementation to different header/source files?

[sammccall]

Address comments

clang-tools-extra/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   QueryDriverDatabase.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp
 
   index/Background.cpp
+  index/BackgroundRebuild.cpp
   index/BackgroundIndexStorage.cpp
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/IndexAction.cpp
   index/MemIndex.cpp
   index/Merge.cpp
   index/Ref.cpp

clang-tools-extra/clangd/ClangdServer.h

struct Options {
     /// If true, ClangdServer builds a dynamic in-memory index for symbols in
     /// opened files and uses the index to augment code completion results.
     bool BuildDynamicSymbolIndex = false;
     /// Use a heavier and faster in-memory index implementation.
     bool HeavyweightDynamicSymbolIndex = true;
     /// If true, ClangdServer automatically indexes files in the current project
     /// on background threads. The index is stored in the project root.
     bool BackgroundIndex = false;
-    /// If set to non-zero, the background index rebuilds the symbol index
-    /// periodically every BuildIndexPeriodMs milliseconds; otherwise, the
-    /// symbol index will be updated for each indexed file.
-    size_t BackgroundIndexRebuildPeriodMs = 0;
 
     /// If set, use this index to augment code completion results.
     SymbolIndex *StaticIndex = nullptr;
 
     /// If set, enable clang-tidy in clangd and use to it get clang-tidy
     /// configurations for a particular file.
     /// Clangd supports only a small subset of ClangTidyOptions, these options
     /// (Checks, CheckOptions) are about which clang-tidy checks will be

clang-tools-extra/clangd/ClangdServer.cpp

if (this->Index != nullptr) {
       this->Index = Idx;
     }
   };
   if (Opts.StaticIndex)
     AddIndex(Opts.StaticIndex);
   if (Opts.BackgroundIndex) {
     BackgroundIdx = llvm::make_unique<BackgroundIndex>(
         Context::current().clone(), FSProvider, CDB,
-        BackgroundIndexStorage::createDiskBackedStorageFactory(),
-        Opts.BackgroundIndexRebuildPeriodMs);
+        BackgroundIndexStorage::createDiskBackedStorageFactory());
     AddIndex(BackgroundIdx.get());
   }
   if (DynamicIdx)
     AddIndex(DynamicIdx.get());
 }
 
 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
                                WantDiagnostics WantDiags) {

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

 //===--- Background.h - Build an index in a background thread ----*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_H
 
 #include "Context.h"
 #include "FSProvider.h"
 #include "GlobalCompilationDatabase.h"
 #include "SourceCode.h"
 #include "Threading.h"
+#include "index/BackgroundRebuild.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "index/Serialization.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Threading.h"
 #include <atomic>
 #include <condition_variable>

 public:
   /// If BuildIndexPeriodMs is greater than 0, the symbol index will only be
   /// rebuilt periodically (one per \p BuildIndexPeriodMs); otherwise, index is
   /// rebuilt for each indexed file.
   BackgroundIndex(
       Context BackgroundContext, const FileSystemProvider &,
       const GlobalCompilationDatabase &CDB,
       BackgroundIndexStorage::Factory IndexStorageFactory,
-      size_t BuildIndexPeriodMs = 0,
       size_t ThreadPoolSize = llvm::heavyweight_hardware_concurrency());
   ~BackgroundIndex(); // Blocks while the current task finishes.
 
   // Enqueue translation units for indexing.
   // The indexing happens in a background thread, so the symbols will be
   // available sometime later.
   void enqueue(const std::vector<std::string> &ChangedFiles);
 
private:
   // configuration
   const FileSystemProvider &FSProvider;
   const GlobalCompilationDatabase &CDB;
   Context BackgroundContext;
 
   // index state
   llvm::Error index(tooling::CompileCommand,
                     BackgroundIndexStorage *IndexStorage);
-  void buildIndex(); // Rebuild index periodically every BuildIndexPeriodMs.
-  const size_t BuildIndexPeriodMs;
-  std::atomic<bool> SymbolsUpdatedSinceLastIndex;
   std::mutex IndexMu;
   std::condition_variable IndexCV;
 
   FileSymbols IndexedSymbols;
+  BackgroundIndexRebuilder Rebuilder;
   llvm::StringMap<ShardVersion> ShardVersions; // Key is absolute file path.
   std::mutex ShardVersionsMu;
 
   BackgroundIndexStorage::Factory IndexStorageFactory;
   struct Source {
     std::string Path;
     bool NeedsReIndexing;
     Source(llvm::StringRef Path, bool NeedsReIndexing)
         : Path(Path), NeedsReIndexing(NeedsReIndexing) {}
   };
   // Loads the shards for a single TU and all of its dependencies. Returns the
   // list of sources and whether they need to be re-indexed.
   std::vector<Source> loadShard(const tooling::CompileCommand &Cmd,
                                 BackgroundIndexStorage *IndexStorage,
                                 llvm::StringSet<> &LoadedShards);
   // Tries to load shards for the ChangedFiles.
   std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
   loadShards(std::vector<std::string> ChangedFiles);
   void enqueue(tooling::CompileCommand Cmd, BackgroundIndexStorage *Storage);
 
   // queue management
   using Task = std::function<void()>;
   void run(); // Main loop executed by Thread. Runs tasks from Queue.
   void enqueueTask(Task T, llvm::ThreadPriority Prioirty);
   void enqueueLocked(tooling::CompileCommand Cmd,
                      BackgroundIndexStorage *IndexStorage);
   std::mutex QueueMu;
   unsigned NumActiveTasks = 0; // Only idle when queue is empty *and* no tasks.
   std::condition_variable QueueCV;

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

 #include "Headers.h"
 #include "Logger.h"
 #include "Path.h"
 #include "SourceCode.h"
 #include "Symbol.h"
 #include "Threading.h"
 #include "Trace.h"
 #include "URI.h"
+#include "index/FileIndex.h"
 #include "index/IndexAction.h"
 #include "index/MemIndex.h"
 #include "index/Ref.h"
 #include "index/Relation.h"
 #include "index/Serialization.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Threading.h"
 
 #include <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <memory>
+#include <mutex>
 #include <numeric>
 #include <queue>
 #include <random>
 #include <string>
 #include <thread>
 
 namespace clang {
 namespace clangd {
llvm::SmallString<128> getAbsolutePath(const tooling::CompileCommand &Cmd) {
   }
   return AbsolutePath;
 }
 } // namespace
 
 BackgroundIndex::BackgroundIndex(
     Context BackgroundContext, const FileSystemProvider &FSProvider,
     const GlobalCompilationDatabase &CDB,
-    BackgroundIndexStorage::Factory IndexStorageFactory,
-    size_t BuildIndexPeriodMs, size_t ThreadPoolSize)
+    BackgroundIndexStorage::Factory IndexStorageFactory, size_t ThreadPoolSize)
     : SwapIndex(llvm::make_unique<MemIndex>()), FSProvider(FSProvider),
       CDB(CDB), BackgroundContext(std::move(BackgroundContext)),
-      BuildIndexPeriodMs(BuildIndexPeriodMs),
-      SymbolsUpdatedSinceLastIndex(false),
+      Rebuilder(this, &IndexedSymbols),
       IndexStorageFactory(std::move(IndexStorageFactory)),
       CommandsChanged(
           CDB.watch([&](const std::vector<std::string> &ChangedFiles) {
             enqueue(ChangedFiles);
           })) {
   assert(ThreadPoolSize > 0 && "Thread pool size can't be zero.");
   assert(this->IndexStorageFactory && "Storage factory can not be null!");
   for (unsigned I = 0; I < ThreadPoolSize; ++I) {
     ThreadPool.runAsync("background-worker-" + llvm::Twine(I + 1),
                         [this] { run(); });
   }
-  if (BuildIndexPeriodMs > 0) {
-    log("BackgroundIndex: build symbol index periodically every {0} ms.",
-        BuildIndexPeriodMs);
-    ThreadPool.runAsync("background-index-builder", [this] { buildIndex(); });
-  }
 }
 
 BackgroundIndex::~BackgroundIndex() {
   stop();
   ThreadPool.wait();
 }
 
 void BackgroundIndex::stop() {
+  Rebuilder.shutdown();
   {
     std::lock_guard<std::mutex> QueueLock(QueueMu);
     std::lock_guard<std::mutex> IndexLock(IndexMu);
     ShouldStop = true;
   }
   QueueCV.notify_all();
   IndexCV.notify_all();
 }
while (true) {
     if (Priority != llvm::ThreadPriority::Default && !PreventStarvation.load())
       llvm::set_thread_priority(Priority);
     (*Task)();
     if (Priority != llvm::ThreadPriority::Default)
       llvm::set_thread_priority(llvm::ThreadPriority::Default);
 
     {
       std::unique_lock<std::mutex> Lock(QueueMu);
+      if (NumActiveTasks == 1 && Queue.empty()) {
+        // We just finished the last item, the queue is going idle.
+        Lock.unlock();
+        Rebuilder.idle();
+        Lock.lock();
+      }
       assert(NumActiveTasks > 0 && "before decrementing");
       --NumActiveTasks;
     }
     QueueCV.notify_all();
   }
 }
 
 bool BackgroundIndex::blockUntilIdleForTest(
if (IndexStorage) {
       // this thread sees the older version but finishes later. This should be
       // rare in practice.
       IndexedSymbols.update(Path, std::move(SS), std::move(RS), std::move(RelS),
                             Path == MainFile);
     }
   }
 }
 
-void BackgroundIndex::buildIndex() {
-  assert(BuildIndexPeriodMs > 0);
-  while (true) {
-    {
-      std::unique_lock<std::mutex> Lock(IndexMu);
-      if (ShouldStop) // Avoid waiting if stopped.
-        break;
-      // Wait until this is notified to stop or `BuildIndexPeriodMs` has past.
-      IndexCV.wait_for(Lock, std::chrono::milliseconds(BuildIndexPeriodMs));
-      if (ShouldStop) // Avoid rebuilding index if stopped.
-        break;
-    }
-    if (!SymbolsUpdatedSinceLastIndex.exchange(false))
-      continue;
-    // There can be symbol update right after the flag is reset above and before
-    // index is rebuilt below. The new index would contain the updated symbols
-    // but the flag would still be true. This is fine as we would simply run an
-    // extra index build.
-    reset(
-        IndexedSymbols.buildIndex(IndexType::Heavy, DuplicateHandling::Merge));
-    log("BackgroundIndex: rebuilt symbol index with estimated memory {0} "
-        "bytes.",
-        estimateMemoryUsage());
-  }
-}
-
 llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd,
                                    BackgroundIndexStorage *IndexStorage) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
   auto AbsolutePath = getAbsolutePath(Cmd);
 
   auto FS = FSProvider.getFileSystem();
   auto Buf = FS->getBufferForFile(AbsolutePath);
llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd,
   if (HadErrors) {
     log("Failed to compile {0}, index may be incomplete", AbsolutePath);
     for (auto &It : *Index.Sources)
       It.second.Flags |= IncludeGraphNode::SourceFlag::HadErrors;
   }
   update(AbsolutePath, std::move(Index), ShardVersionsSnapshot, IndexStorage,
          HadErrors);
 
-  if (BuildIndexPeriodMs > 0)
-    SymbolsUpdatedSinceLastIndex = true;
-  else
-    reset(
-        IndexedSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge));
-
+  Rebuilder.indexedTU();
   return llvm::Error::success();
 }
 
 std::vector<BackgroundIndex::Source>
 BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
                            BackgroundIndexStorage *IndexStorage,
                            llvm::StringSet<> &LoadedShards) {
   struct ShardInfo {
for (const ShardInfo &SI : IntermediateSymbols) {
       ShardVersion &SV = ShardVersions[SI.AbsolutePath];
       SV.Digest = SI.Digest;
       SV.HadErrors = SI.HadErrors;
 
       IndexedSymbols.update(SI.AbsolutePath, std::move(SS), std::move(RS),
                             std::move(RelS), SI.CountReferences);
     }
   }
+  if (!IntermediateSymbols.empty())
+    Rebuilder.loadedTU();
 
   return Dependencies;
 }
 
 // Goes over each changed file and loads them from index. Returns the list of
 // TUs that had out-of-date/no shards.
 std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
 BackgroundIndex::loadShards(std::vector<std::string> ChangedFiles) {
   std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
       NeedsReIndexing;
   // Keeps track of the files that will be reindexed, to make sure we won't
   // re-index same dependencies more than once. Keys are AbsolutePaths.
   llvm::StringSet<> FilesToIndex;
   // Keeps track of the loaded shards to make sure we don't perform redundant
   // disk IO. Keys are absolute paths.
   llvm::StringSet<> LoadedShards;
+  Rebuilder.startLoading();
   for (const auto &File : ChangedFiles) {
     ProjectInfo PI;
     auto Cmd = CDB.getCompileCommand(File, &PI);
     if (!Cmd)
       continue;
     BackgroundIndexStorage *IndexStorage = IndexStorageFactory(PI.SourceRoot);
     auto Dependencies = loadShard(*Cmd, IndexStorage, LoadedShards);
     for (const auto &Dependency : Dependencies) {
       if (!Dependency.NeedsReIndexing || FilesToIndex.count(Dependency.Path))
         continue;
       // FIXME: Currently, we simply schedule indexing on a TU whenever any of
       // its dependencies needs re-indexing. We might do it smarter by figuring
       // out a minimal set of TUs that will cover all the stale dependencies.
       vlog("Enqueueing TU {0} because its dependency {1} needs re-indexing.",
            Cmd->Filename, Dependency.Path);
       NeedsReIndexing.push_back({std::move(*Cmd), IndexStorage});
       // Mark all of this TU's dependencies as to-be-indexed so that we won't
       // try to re-index those.
       for (const auto &Dependency : Dependencies)
         FilesToIndex.insert(Dependency.Path);
       break;
     }
   }
-  vlog("Loaded all shards");
-  reset(IndexedSymbols.buildIndex(IndexType::Heavy, DuplicateHandling::Merge));
-  vlog("BackgroundIndex: built symbol index with estimated memory {0} "
-       "bytes.",
-       estimateMemoryUsage());
+  Rebuilder.doneLoading();
   return NeedsReIndexing;
 }
 
 void BackgroundIndex::preventThreadStarvationInTests() {
   PreventStarvation.store(true);
 }
 
 } // namespace clangd
 } // namespace clang

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

+//===--- BackgroundIndexRebuild.h - when to rebuild the bg index--*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains an implementation detail of the background indexer
+// (Background.h), which is exposed for testing.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_REBUILD_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_REBUILD_H
+
+#include "index/FileIndex.h"
+#include "index/Index.h"
+
+namespace clang {
+namespace clangd {
+
+// The BackgroundIndexRebuilder builds the serving data structures periodically
+// in response to events in the background indexer. The goal is to ensure the
+// served data stays fairly fresh, without wasting lots of CPU rebuilding it
+// often.
+//
+// The index is always built after a set of shards are loaded from disk.
+// This happens when clangd discovers a compilation database that we've
+// previously built an index for. It's a fairly fast process that yields lots
+// of data, so we wait to get all of it.
+//
+// The index is built after indexing a few translation units, if it wasn't built
+// already. This ensures quick startup if there's no existing index.
+// Waiting for a few random TUs yields coverage of the most common headers.
+//
+// The index is rebuilt every N TUs, to keep if fresh as files are indexed.
+//
+// The index is rebuilt every time the queue goes idle, if it's stale.
+//
+// All methods are threadsafe. They're called after FileSymbols is updated
+// etc. Without external locking, the rebuilt index may include more updates
+// than intended, which is fine.
+//
+// This class is exposed in the header so it can be tested.
+class BackgroundIndexRebuilder {
+public:
+  // Thresholds for rebuilding as TUs get indexed.
+  static constexpr unsigned TUsBeforeFirstBuild = 5;
+  static constexpr unsigned TUsBeforeRebuild = 100;
+
+  BackgroundIndexRebuilder(SwapIndex *Target, FileSymbols *Source);
+
+  // Called to indicate a TU has been indexed.
+  // May rebuild, if enough TUs have been indexed.
+  void indexedTU();
+  // Called to indicate that all worker threads are idle.
+  // May reindex, if the index is not up to date.
+  void idle();
+  // Called to indicate we're going to load a batch of shards from disk.
+  // startLoading() and doneLoading() must be paired, but multiple loading
+  // sessions may happen concurrently.
+  void startLoading();
+  // Called to indicate some shards were actually loaded from disk.
+  void loadedTU();
+  // Called to indicate we're finished loading shards from disk.
+  // May rebuild (if any were loaded).
+  void doneLoading();
+
+  // Ensures we won't start any more rebuilds.
+  void shutdown();
+
+private:
+  // Run Check under the lock, and rebuild if it returns true.
+  void maybeRebuild(const char *Reason, std::function<bool()> Check);
+  bool enoughTUsToRebuild() const;
+
+  // All transient state is guarded by the mutex.
+  std::mutex Mu;
+  bool ShouldStop = false;
+  // Index builds are versioned. ActiveVersion chases StartedVersion.
+  unsigned StartedVersion = 0;
+  unsigned ActiveVersion = 0;
+  // How many TUs have we indexed so far since startup?
+  unsigned IndexedTUs = 0;
+  unsigned IndexedTUsAtLastRebuild = 0;
+  // Are we loading shards? May be multiple concurrent sessions.
+  unsigned Loading = 0;
+  unsigned LoadedTUs; // In the current loading session.
+
+  SwapIndex *Target;
+  FileSymbols *Source;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

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

+//===-- BackgroundRebuild.cpp - when to rebuild thei background index -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "index/BackgroundRebuild.h"
+#include "ClangdUnit.h"
+#include "Compiler.h"
+#include "Headers.h"
+#include "Logger.h"
+#include "Path.h"
+#include "SourceCode.h"
+#include "Symbol.h"
+#include "Threading.h"
+#include "Trace.h"
+#include "URI.h"
+#include "index/FileIndex.h"
+#include "index/IndexAction.h"
+#include "index/MemIndex.h"
+#include "index/Ref.h"
+#include "index/Relation.h"
+#include "index/Serialization.h"
+#include "index/SymbolCollector.h"
+#include "clang/Basic/SourceLocation.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/Threading.h"
+
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <memory>
+#include <mutex>
+#include <numeric>
+#include <queue>
+#include <random>
+#include <string>
+#include <thread>
+
+namespace clang {
+namespace clangd {
+
+BackgroundIndexRebuilder::BackgroundIndexRebuilder(SwapIndex *Target,
+                                                   FileSymbols *Source)
+    : Target(Target), Source(Source) {}
+
+bool BackgroundIndexRebuilder::enoughTUsToRebuild() const {
+  if (!ActiveVersion)                         // never built
+    return IndexedTUs == TUsBeforeFirstBuild; // use low threshold
+  // rebuild if we've reached the (higher) threshold
+  return IndexedTUs >= IndexedTUsAtLastRebuild + TUsBeforeRebuild;
+}
+
+void BackgroundIndexRebuilder::indexedTU() {
+  maybeRebuild("after indexing enough files", [this] {
+    ++IndexedTUs;
+    if (Loading)
+      return false;                      // rebuild once loading finishes
+    if (ActiveVersion != StartedVersion) // currently building
+      return false;                      // no urgency, avoid overlapping builds
+    return enoughTUsToRebuild();
+  });
+}
+
+void BackgroundIndexRebuilder::idle() {
+  maybeRebuild("when background indexer is idle", [this] {
+    // rebuild if there's anything new in the index.
+    // (even if currently rebuilding! this ensures eventual completeness)
+    return IndexedTUs > IndexedTUsAtLastRebuild;
+  });
+}
+
+void BackgroundIndexRebuilder::startLoading() {
+  std::lock_guard<std::mutex> Lock(Mu);
+  if (!Loading)
+    LoadedTUs = 0;
+  ++Loading;
+}
+void BackgroundIndexRebuilder::loadedTU() {
+  std::lock_guard<std::mutex> Lock(Mu);
+  assert(Loading);
+  ++LoadedTUs;
+}
+void BackgroundIndexRebuilder::doneLoading() {
+  maybeRebuild("after loading index from disk", [this] {
+    assert(Loading);
+    --Loading;
+    if (Loading)    // was loading multiple batches concurrently
+      return false; // rebuild once the last batch is done.
+    // Rebuild if we loaded any shards, or if we stopped an indexedTU rebuild.
+    return LoadedTUs > 0 || enoughTUsToRebuild();
+  });
+}
+
+void BackgroundIndexRebuilder::shutdown() {
+  std::lock_guard<std::mutex> Lock(Mu);
+  ShouldStop = true;
+}
+
+void BackgroundIndexRebuilder::maybeRebuild(const char *Reason,
+                                            std::function<bool()> Check) {
+  unsigned BuildVersion = 0;
+  {
+    std::lock_guard<std::mutex> Lock(Mu);
+    if (!ShouldStop && Check()) {
+      BuildVersion = ++StartedVersion;
+      IndexedTUsAtLastRebuild = IndexedTUs;
+    }
+  }
+  if (BuildVersion) {
+    std::unique_ptr<SymbolIndex> NewIndex;
+    {
+      vlog("BackgroundIndex: building version {0} {1}", BuildVersion, Reason);
+      trace::Span Tracer("RebuildBackgroundIndex");
+      SPAN_ATTACH(Tracer, "reason", Reason);
+      NewIndex = Source->buildIndex(IndexType::Heavy, DuplicateHandling::Merge);
+    }
+    {
+      std::lock_guard<std::mutex> Lock(Mu);
+      // Guard against rebuild finishing in the wrong order.
+      if (BuildVersion > ActiveVersion) {
+        ActiveVersion = BuildVersion;
+        vlog("BackgroundIndex: serving version {0} ({1} bytes)", BuildVersion,
+             NewIndex->estimateMemoryUsage());
+        Target->reset(std::move(NewIndex));
+      }
+    }
+  }
+}
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/test/background-index.test

 # We need to splice paths into file:// URIs for this test.
 # UNSUPPORTED: win32
 
 # Use a copy of inputs, as we'll mutate it (as will the background index).
 # RUN: rm -rf %t
 # RUN: cp -r %S/Inputs/background-index %t
 # Need to embed the correct temp path in the actual JSON-RPC requests.
 # RUN: sed -i -e "s|DIRECTORY|%t|" %t/*
 
 # We're editing bar.cpp, which includes foo.h.
 # foo() is declared in foo.h and defined in foo.cpp.
 # The background index should allow us to go-to-definition on foo().
-# RUN: clangd -background-index -background-index-rebuild-period=0 -lit-test < %t/definition.jsonrpc | FileCheck %t/definition.jsonrpc
+# RUN: clangd -background-index -lit-test < %t/definition.jsonrpc | FileCheck %t/definition.jsonrpc
 
 # Test that the index is writing files in the expected location.
 # RUN: ls %t/.clangd/index/foo.cpp.*.idx
 
 # Test the index is read from disk: delete code and restart clangd.
 # RUN: rm %t/foo.cpp
 # RUN: clangd -background-index -lit-test < %t/definition.jsonrpc | FileCheck %t/definition.jsonrpc

clang-tools-extra/clangd/tool/ClangdMain.cpp

 
 static llvm::cl::opt<bool> EnableBackgroundIndex(
     "background-index",
     llvm::cl::desc(
         "Index project code in the background and persist index on disk. "
         "Experimental"),
     llvm::cl::init(true));
 
-static llvm::cl::opt<int> BackgroundIndexRebuildPeriod(
-    "background-index-rebuild-period",
-    llvm::cl::desc(
-        "If set to non-zero, the background index rebuilds the symbol index "
-        "periodically every X milliseconds; otherwise, the "
-        "symbol index will be updated for each indexed file"),
-    llvm::cl::init(5000), llvm::cl::Hidden);
-
 enum CompileArgsFrom { LSPCompileArgs, FilesystemCompileArgs };
 static llvm::cl::opt<CompileArgsFrom> CompileArgsFrom(
     "compile_args_from", llvm::cl::desc("The source of compile commands"),
     llvm::cl::values(clEnumValN(LSPCompileArgs, "lsp",
                                 "All compile commands come from LSP and "
                                 "'compile_commands.json' files are ignored"),
                      clEnumValN(FilesystemCompileArgs, "filesystem",
                                 "All compile commands come from the "
int main(int argc, char *argv[]) {
   case PCHStorageFlag::Disk:
     Opts.StorePreamblesInMemory = false;
     break;
   }
   if (!ResourceDir.empty())
     Opts.ResourceDir = ResourceDir;
   Opts.BuildDynamicSymbolIndex = EnableIndex;
   Opts.BackgroundIndex = EnableBackgroundIndex;
-  Opts.BackgroundIndexRebuildPeriodMs = BackgroundIndexRebuildPeriod;
   std::unique_ptr<SymbolIndex> StaticIdx;
   std::future<void> AsyncIndexLoad; // Block exit while loading the index.
   if (EnableIndex && !IndexFile.empty()) {
     // Load the index asynchronously. Meanwhile SwapIndex returns no results.
     SwapIndex *Placeholder;
     StaticIdx.reset(Placeholder = new SwapIndex(llvm::make_unique<MemIndex>()));
     AsyncIndexLoad = runAsync<void>([Placeholder] {
       if (auto Idx = loadIndex(IndexFile, /*UseDex=*/true))

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

 #include "Headers.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
+#include "TestIndex.h"
 #include "TestTU.h"
 #include "index/Background.h"
+#include "index/BackgroundRebuild.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/Threading.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <thread>
 
 using ::testing::_;
EXPECT_THAT(
       ShardHeader->Sources->lookup("unittest:///root/A.h").DirectIncludes,
       UnorderedElementsAre("unittest:///root/B.h"));
   EXPECT_NE(ShardHeader->Sources->lookup("unittest:///root/A.h").Digest,
             FileDigest{{0}});
   EXPECT_THAT(ShardHeader->Sources->lookup("unittest:///root/B.h"),
               EmptyIncludeNode());
 }
 
-// FIXME: figure out the right timeouts or rewrite to not use the timeouts and
-// re-enable.
-TEST_F(BackgroundIndexTest, DISABLED_PeriodicalIndex) {
-  MockFSProvider FS;
-  llvm::StringMap<std::string> Storage;
-  size_t CacheHits = 0;
-  MemoryShardStorage MSS(Storage, CacheHits);
-  OverlayCDB CDB(/*Base=*/nullptr);
-  BackgroundIndex Idx(
-      Context::empty(), FS, CDB, [&](llvm::StringRef) { return &MSS; },
-      /*BuildIndexPeriodMs=*/500);
-
-  FS.Files[testPath("root/A.cc")] = "#include \"A.h\"";
-
-  tooling::CompileCommand Cmd;
-  FS.Files[testPath("root/A.h")] = "class X {};";
-  Cmd.Filename = testPath("root/A.cc");
-  Cmd.CommandLine = {"clang++", Cmd.Filename};
-  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
-
-  ASSERT_TRUE(Idx.blockUntilIdleForTest());
-  EXPECT_THAT(runFuzzyFind(Idx, ""), ElementsAre());
-  std::this_thread::sleep_for(std::chrono::milliseconds(1000));
-  EXPECT_THAT(runFuzzyFind(Idx, ""), ElementsAre(Named("X")));
-
-  FS.Files[testPath("root/A.h")] = "class Y {};";
-  FS.Files[testPath("root/A.cc")] += " "; // Force reindex the file.
-  Cmd.CommandLine = {"clang++", "-DA=1", testPath("root/A.cc")};
-  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
-
-  ASSERT_TRUE(Idx.blockUntilIdleForTest());
-  EXPECT_THAT(runFuzzyFind(Idx, ""), ElementsAre(Named("X")));
-  std::this_thread::sleep_for(std::chrono::milliseconds(1000));
-  EXPECT_THAT(runFuzzyFind(Idx, ""), ElementsAre(Named("Y")));
-}
-
 TEST_F(BackgroundIndexTest, ShardStorageLoad) {
   MockFSProvider FS;
   FS.Files[testPath("root/A.h")] = R"cpp(
       void common();
       void f_b();
       class A_CC {};
       )cpp";
   FS.Files[testPath("root/A.cc")] =
TEST_F(BackgroundIndexTest, CmdLineHash) {
 
   {
     tooling::CompileCommand CmdStored = *MSS.loadShard(testPath("A.cc"))->Cmd;
     EXPECT_EQ(CmdStored.CommandLine, Cmd.CommandLine);
     EXPECT_EQ(CmdStored.Directory, Cmd.Directory);
   }
 }
 
+class BackgroundIndexRebuilderTest : public testing::Test {
+protected:
+  BackgroundIndexRebuilderTest()
+      : Target(llvm::make_unique<MemIndex>()), Rebuilder(&Target, &Source) {
+    // Prepare FileSymbols with TestSymbol in it, for checkRebuild.
+    TestSymbol.ID = SymbolID("foo");
+  }
+
+  // Perform Action and determine whether it rebuilt the index or not.
+  bool checkRebuild(std::function<void()> Action) {
+    // Update reference count so we can tell if the index updates.
+    ++TestSymbol.References;
+    SymbolSlab::Builder SB;
+    SB.insert(TestSymbol);
+    Source.update("", llvm::make_unique<SymbolSlab>(std::move(SB).build()),
+                  nullptr, nullptr, false);
+    // Now maybe update the index.
+    Action();
+    // Now query the index to get the reference count.
+    unsigned ReadReferences = 0;
+    LookupRequest Req;
+    Req.IDs.insert(TestSymbol.ID);
+    Target.lookup(Req, [&](const Symbol &S) { ReadReferences = S.References; });
+    // The index was rebuild if the reference count is up to date.
+    return ReadReferences == TestSymbol.References;
+  }
+
+  Symbol TestSymbol;
+  FileSymbols Source;
+  SwapIndex Target;
+  BackgroundIndexRebuilder Rebuilder;
+};
+
+TEST_F(BackgroundIndexRebuilderTest, IndexingTUs) {
+  for (unsigned I = 0; I < BackgroundIndexRebuilder::TUsBeforeFirstBuild - 1;
+       ++I)
+    EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
+  EXPECT_TRUE(checkRebuild([&] { Rebuilder.indexedTU(); }));
+  for (unsigned I = 0; I < BackgroundIndexRebuilder::TUsBeforeRebuild - 1; ++I)
+    EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
+  EXPECT_TRUE(checkRebuild([&] { Rebuilder.indexedTU(); }));
+}
+
+TEST_F(BackgroundIndexRebuilderTest, LoadingShards) {
+  Rebuilder.startLoading();
+  Rebuilder.loadedTU();
+  Rebuilder.loadedTU();
+  EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
+
+  // No rebuild for no shards.
+  Rebuilder.startLoading();
+  EXPECT_FALSE(checkRebuild([&] { Rebuilder.doneLoading(); }));
+
+  // Loads can overlap.
+  Rebuilder.startLoading();
+  Rebuilder.loadedTU();
+  Rebuilder.startLoading();
+  Rebuilder.loadedTU();
+  EXPECT_FALSE(checkRebuild([&] { Rebuilder.doneLoading(); }));
+  Rebuilder.loadedTU();
+  EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
+
+  // No rebuilding for indexed files while loading.
+  Rebuilder.startLoading();
+  for (unsigned I = 0; I < 3 * BackgroundIndexRebuilder::TUsBeforeRebuild; ++I)
+    EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
+  // But they get indexed when we're done, even if no shards were loaded.
+  EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
+}
+
 } // namespace clangd
 } // namespace clang