[clangd] Refactor background-index shard loading

[sammccall]

So I've stared at this refactoring for a while, and I still don't totally get it.
It seems like a class that really wants to be a function. It's a complicated function though - maybe a separate cpp file does make sense.

We discussed moving ShardVersions and Writes into the loader, splitting ShardVersions from IndexSymbols seems pretty odd though.

What's the advantage of making this stateful?

[kadircet]

• As discussed offline, moved the "staleness" detection back to background index
• A renaming within Rebuilder, "LoadedTUs -> LoadedShards"

[sammccall]

Refactoring generally looks good.
You're replacing a lot of documented code with new undocumented code, can we add some high-level comments?

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

130

nit: unclear what's changed with respect to what.

I'd suggest shardIsStale() and swap the param order?

133

While here: this message lacks context and is unnecessarily jargony: "Background-indexing: couldn't read {0} to validate stored index"?

455

This function now has almost no comments, can you explain the major steps?

484

move this outside the loop and use a counter? it's an extra lock/unlock for no reason

520

nit: again, this patch is adding Path/PathRef aliases to parts of the code that don't use them

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

177

and their dependencies

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

78

you can't call this an (opaque?) ShardIdentifier and then use it as a path - maybe SourceFile?

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

2

header is out of date. The new file/structs/functions need documentation

26

LoadedShard?

31

IndexFileIn is just a struct - optional<>? (And need a comment explaining when it can be null)

37

no, they're not :-)

we talked about making the values refs to ShardInfo.AbsolutePath.

38

why are these outside the ShardInfo?

43

clangd::load seems like an insufficiently precise name for this function

loadIndexShards?

[kadircet]

• Address comments

[kadircet]

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

520

I believe they represent the intention better then string/StringRefthat was the reason I was using them.

[sammccall]

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

133

{0 is missing its }

497

// We'll accept data from stale shards, but ensure the files get reindexed soon.

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

94

This handles only one main file, I can't see where you skip loading shards for headers that were loaded for a previous file

105

ShardIdentifier still here

127

I don't think this works, the storage for these StringRefs is the LoadedShards you just moved from.
(In practice, maybe you get away with it if the strings are longer than the SSO?)

[kadircet]

• Address comments

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

94

The real loading happens down the line at const CachedShard &CS = loadShard(ShardIdentifier, Storage);, which returns shard from cache if it was already loaded.

127

Yes you are right, I don't see a good way out of it with current data model. I suppose we can keep an "intern table :)" for the strings, but then LoadedShards would be meaningless without that intern table so we would end up returning another struct from loadIndexShards that will wrap std::vector<LoadedShards> and the intern table(but we tried to mitigate from that struct in previous iteration).

So instead I am changing the model to return only the first translation unit and actually this makes a lot of sense when we think about shards for standard library, in a code base like chromium vector could easily have ~40k dependent TUs and even if we just returned a vector<StringRef> it would still be around 640kB, and we have a lot of standard library headers. this can quite easily add hundreds of MB to clangd's memory footprint while loading shards.

[kadircet]

• Add comments

[sammccall]

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

47

I think we only need to traverse this when *missing* the cache, so it can be a return value of loadShard rather than stored in CachedShard (so CachedShard becomes just LoadedShard)

[sammccall]

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

94

OK, then we still traverse the whole subtree but hit every node.

I think it'd be nicer to avoid the traversal (as noted above) if that works well, up to you.

[kadircet]

• Get rid of CachedShard

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/BackgroundIndexLoader.cpp
   index/BackgroundIndexStorage.cpp
   index/BackgroundQueue.cpp
   index/BackgroundRebuild.cpp
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/IndexAction.cpp
   index/MemIndex.cpp

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 "Path.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"
llvm::Error index(tooling::CompileCommand,
                     BackgroundIndexStorage *IndexStorage);
 
   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.
+  // Tries to load shards for the MainFiles and their dependencies.
   std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
-  loadShards(std::vector<std::string> ChangedFiles);
+  loadProject(std::vector<std::string> MainFiles);
 
   BackgroundQueue::Task
   changedFilesTask(const std::vector<std::string> &ChangedFiles);
   BackgroundQueue::Task indexFileTask(tooling::CompileCommand Cmd,
                                       BackgroundIndexStorage *Storage);
 
   // from lowest to highest priority
   enum QueuePriority {

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

 //===-- Background.cpp - Build an index in a background thread ------------===//
 //
 // 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/Background.h"
 #include "ClangdUnit.h"
 #include "Compiler.h"
 #include "Context.h"
+#include "FSProvider.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/BackgroundIndexLoader.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 "clang/Driver/Types.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.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/Path.h"
 #include "llvm/Support/Threading.h"
 
 #include <algorithm>
 #include <atomic>
 #include <chrono>
 #include <condition_variable>
+#include <cstddef>
 #include <memory>
 #include <mutex>
 #include <numeric>
 #include <queue>
 #include <random>
 #include <string>
 #include <thread>
+#include <utility>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 // Resolves URI to file paths with cache.
 class URIToFileCache {
 public:
if (llvm::sys::path::is_absolute(Cmd.Filename)) {
     AbsolutePath = Cmd.Filename;
   } else {
     AbsolutePath = Cmd.Directory;
     llvm::sys::path::append(AbsolutePath, Cmd.Filename);
     llvm::sys::path::remove_dots(AbsolutePath, true);
   }
   return AbsolutePath;
 }
+
+bool shardIsStale(const LoadedShard &LS, llvm::vfs::FileSystem *FS) {
+  auto Buf = FS->getBufferForFile(LS.AbsolutePath);
+  if (!Buf) {
+    elog("Background-index: Couldn't read {0} to validate stored index: {1}",
+         LS.AbsolutePath, Buf.getError().message());
+    // There is no point in indexing an unreadable file.
+    return false;
+  }
+  return digest(Buf->get()->getBuffer()) != LS.Digest;
+}
+
 } // namespace
 
 BackgroundIndex::BackgroundIndex(
     Context BackgroundContext, const FileSystemProvider &FSProvider,
     const GlobalCompilationDatabase &CDB,
     BackgroundIndexStorage::Factory IndexStorageFactory, size_t ThreadPoolSize)
     : SwapIndex(llvm::make_unique<MemIndex>()), FSProvider(FSProvider),
       CDB(CDB), BackgroundContext(std::move(BackgroundContext)),

 BackgroundQueue::Task BackgroundIndex::changedFilesTask(
     const std::vector<std::string> &ChangedFiles) {
   BackgroundQueue::Task T([this, ChangedFiles] {
     trace::Span Tracer("BackgroundIndexEnqueue");
     // We're doing this asynchronously, because we'll read shards here too.
     log("Enqueueing {0} commands for indexing", ChangedFiles.size());
     SPAN_ATTACH(Tracer, "files", int64_t(ChangedFiles.size()));
 
-    auto NeedsReIndexing = loadShards(std::move(ChangedFiles));
+    auto NeedsReIndexing = loadProject(std::move(ChangedFiles));
     // Run indexing for files that need to be updated.
     std::shuffle(NeedsReIndexing.begin(), NeedsReIndexing.end(),
                  std::mt19937(std::random_device{}()));
     std::vector<BackgroundQueue::Task> Tasks;
     Tasks.reserve(NeedsReIndexing.size());
     for (auto &Elem : NeedsReIndexing)
       Tasks.push_back(indexFileTask(std::move(Elem.first), Elem.second));
     Queue.append(std::move(Tasks));
llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd,
   }
   update(AbsolutePath, std::move(Index), ShardVersionsSnapshot, IndexStorage,
          HadErrors);
 
   Rebuilder.indexedTU();
   return llvm::Error::success();
 }
 
-std::vector<BackgroundIndex::Source>
-BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
-                           BackgroundIndexStorage *IndexStorage,
-                           llvm::StringSet<> &LoadedShards) {
-  struct ShardInfo {
-    std::string AbsolutePath;
-    std::unique_ptr<IndexFileIn> Shard;
-    FileDigest Digest = {};
-    bool CountReferences = false;
-    bool HadErrors = false;
-  };
-  std::vector<ShardInfo> IntermediateSymbols;
-  // Make sure we don't have duplicate elements in the queue. Keys are absolute
-  // paths.
-  llvm::StringSet<> InQueue;
-  auto FS = FSProvider.getFileSystem();
-  // Dependencies of this TU, paired with the information about whether they
-  // need to be re-indexed or not.
-  std::vector<Source> Dependencies;
-  std::queue<Source> ToVisit;
-  std::string AbsolutePath = getAbsolutePath(Cmd).str();
-  // Up until we load the shard related to a dependency it needs to be
-  // re-indexed.
-  ToVisit.emplace(AbsolutePath, true);
-  InQueue.insert(AbsolutePath);
-  // Goes over each dependency.
-  while (!ToVisit.empty()) {
-    Dependencies.push_back(std::move(ToVisit.front()));
-    // Dependencies is not modified during the rest of the loop, so it is safe
-    // to keep the reference.
-    auto &CurDependency = Dependencies.back();
-    ToVisit.pop();
-    // If we have already seen this shard before(either loaded or failed) don't
-    // re-try again. Since the information in the shard won't change from one TU
-    // to another.
-    if (!LoadedShards.try_emplace(CurDependency.Path).second) {
-      // If the dependency needs to be re-indexed, first occurence would already
-      // have detected that, so we don't need to issue it again.
-      CurDependency.NeedsReIndexing = false;
-      continue;
-    }
-
-    auto Shard = IndexStorage->loadShard(CurDependency.Path);
-    if (!Shard || !Shard->Sources) {
-      // File will be returned as requiring re-indexing to caller.
-      vlog("Failed to load shard: {0}", CurDependency.Path);
-      continue;
-    }
-    // These are the edges in the include graph for current dependency.
-    for (const auto &I : *Shard->Sources) {
-      auto U = URI::parse(I.getKey());
-      if (!U)
-        continue;
-      auto AbsolutePath = URI::resolve(*U, CurDependency.Path);
-      if (!AbsolutePath)
-        continue;
-      // Add file as dependency if haven't seen before.
-      if (InQueue.try_emplace(*AbsolutePath).second)
-        ToVisit.emplace(*AbsolutePath, true);
-      // The node contains symbol information only for current file, the rest is
-      // just edges.
-      if (*AbsolutePath != CurDependency.Path)
-        continue;
+// Restores shards for \p MainFiles from index storage. Then checks staleness of
+// those shards and returns a list of TUs that needs to be indexed to update
+// staleness.
+std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
+BackgroundIndex::loadProject(std::vector<std::string> MainFiles) {
+  std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
+      NeedsReIndexing;
 
-      // We found source file info for current dependency.
-      assert(I.getValue().Digest != FileDigest{{0}} && "Digest is empty?");
-      ShardInfo SI;
-      SI.AbsolutePath = CurDependency.Path;
-      SI.Shard = std::move(Shard);
-      SI.Digest = I.getValue().Digest;
-      SI.CountReferences =
-          I.getValue().Flags & IncludeGraphNode::SourceFlag::IsTU;
-      SI.HadErrors =
-          I.getValue().Flags & IncludeGraphNode::SourceFlag::HadErrors;
-      IntermediateSymbols.push_back(std::move(SI));
-      // Check if the source needs re-indexing.
-      // Get the digest, skip it if file doesn't exist.
-      auto Buf = FS->getBufferForFile(CurDependency.Path);
-      if (!Buf) {
-        elog("Couldn't get buffer for file: {0}: {1}", CurDependency.Path,
-             Buf.getError().message());
-        continue;
-      }
-      // If digests match then dependency doesn't need re-indexing.
-      // FIXME: Also check for dependencies(sources) of this shard and compile
-      // commands for cache invalidation.
-      CurDependency.NeedsReIndexing =
-          digest(Buf->get()->getBuffer()) != I.getValue().Digest;
-    }
-  }
-  // Load shard information into background-index.
+  Rebuilder.startLoading();
+  // Load shards for all of the mainfiles.
+  const std::vector<LoadedShard> Result =
+      loadIndexShards(MainFiles, IndexStorageFactory, CDB);
+  size_t LoadedShards = 0;
   {
+    // Update in-memory state.
     std::lock_guard<std::mutex> Lock(ShardVersionsMu);
-    // This can override a newer version that is added in another thread,
-    // if this thread sees the older version but finishes later. This
-    // should be rare in practice.
-    for (const ShardInfo &SI : IntermediateSymbols) {
+    for (auto &LS : Result) {
+      if (!LS.Shard)
+        continue;
       auto SS =
-          SI.Shard->Symbols
-              ? llvm::make_unique<SymbolSlab>(std::move(*SI.Shard->Symbols))
+          LS.Shard->Symbols
+              ? llvm::make_unique<SymbolSlab>(std::move(*LS.Shard->Symbols))
               : nullptr;
-      auto RS = SI.Shard->Refs
-                    ? llvm::make_unique<RefSlab>(std::move(*SI.Shard->Refs))
+      auto RS = LS.Shard->Refs
+                    ? llvm::make_unique<RefSlab>(std::move(*LS.Shard->Refs))
                     : nullptr;
       auto RelS =
-          SI.Shard->Relations
-              ? llvm::make_unique<RelationSlab>(std::move(*SI.Shard->Relations))
+          LS.Shard->Relations
+              ? llvm::make_unique<RelationSlab>(std::move(*LS.Shard->Relations))
               : nullptr;
-      ShardVersion &SV = ShardVersions[SI.AbsolutePath];
-      SV.Digest = SI.Digest;
-      SV.HadErrors = SI.HadErrors;
+      ShardVersion &SV = ShardVersions[LS.AbsolutePath];
+      SV.Digest = LS.Digest;
+      SV.HadErrors = LS.HadErrors;
+      ++LoadedShards;
 
-      IndexedSymbols.update(SI.AbsolutePath, std::move(SS), std::move(RS),
-                            std::move(RelS), SI.CountReferences);
+      IndexedSymbols.update(LS.AbsolutePath, std::move(SS), std::move(RS),
+                            std::move(RelS), LS.CountReferences);
     }
   }
-  if (!IntermediateSymbols.empty())
-    Rebuilder.loadedTU();
-
-  return Dependencies;
-}
+  Rebuilder.loadedShard(LoadedShards);
+  Rebuilder.doneLoading();
 
-// 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) {
-    auto Cmd = CDB.getCompileCommand(File);
-    if (!Cmd)
+  auto FS = FSProvider.getFileSystem();
+  llvm::DenseSet<PathRef> TUsToIndex;
+  // We'll accept data from stale shards, but ensure the files get reindexed
+  // soon.
+  for (auto &LS : Result) {
+    if (!shardIsStale(LS, FS.get()))
       continue;
+    PathRef TUForFile = LS.DependentTU;
+    assert(!TUForFile.empty() && "File without a TU!");
 
-    std::string ProjectRoot;
-    if (auto PI = CDB.getProjectInfo(File))
-      ProjectRoot = std::move(PI->SourceRoot);
-
-    BackgroundIndexStorage *IndexStorage = IndexStorageFactory(ProjectRoot);
-    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;
-    }
+    // FIXME: Try looking at other TUs if no compile commands are available
+    // for this TU, i.e TU was deleted after we performed indexing.
+    TUsToIndex.insert(TUForFile);
+
+  for (PathRef TU : TUsToIndex) {
+    auto Cmd = CDB.getCompileCommand(TU);
+    if (!Cmd)
+      continue;
+    std::string ProjectRoot;
+    if (auto PI = CDB.getProjectInfo(TU))
+      ProjectRoot = std::move(PI->SourceRoot);
+    BackgroundIndexStorage *Storage = IndexStorageFactory(ProjectRoot);
+    NeedsReIndexing.emplace_back(std::move(*Cmd), Storage);
   }
-  Rebuilder.doneLoading();
+
   return NeedsReIndexing;
 }
 
 } // namespace clangd
 } // namespace clang

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

+//===--- BackgroundIndexLoader.h - Load shards from index storage-*- 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_INDEX_LOADER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_LOADER_H
+
+#include "Path.h"
+#include "index/Background.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/VirtualFileSystem.h"
+#include <memory>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+
+/// Represents a shard loaded from storage, stores contents in \p Shard and
+/// metadata about the source file that generated this shard.
+struct LoadedShard {
+  /// Path of the source file that produced this shard.
+  Path AbsolutePath;
+  /// Digest of the source file contents that produced this shard.
+  FileDigest Digest = {};
+  /// Whether the RefSlab in Shard should be used for updating symbol reference
+  /// counts when building an index.
+  bool CountReferences = false;
+  /// Whether the indexing action producing that shard had errors.
+  bool HadErrors = false;
+  /// Path to a TU that is depending on this shard.
+  Path DependentTU;
+  /// Will be nullptr when index storage couldn't provide a valid shard for
+  /// AbsolutePath.
+  std::unique_ptr<IndexFileIn> Shard;
+};
+
+/// Loads all shards for the TU \p MainFile from \p Storage.
+std::vector<LoadedShard>
+loadIndexShards(llvm::ArrayRef<Path> MainFiles,
+                BackgroundIndexStorage::Factory &IndexStorageFactory,
+                const GlobalCompilationDatabase &CDB);
+
+} // namespace clangd
+} // namespace clang
+
+#endif

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

+//===-- BackgroundIndexLoader.cpp - ---------------------------------------===//
+//
+// 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/BackgroundIndexLoader.h"
+#include "GlobalCompilationDatabase.h"
+#include "Logger.h"
+#include "Path.h"
+#include "index/Background.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Path.h"
+#include <string>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+namespace {
+
+llvm::Optional<Path> uriToAbsolutePath(llvm::StringRef URI, PathRef HintPath) {
+  auto U = URI::parse(URI);
+  if (!U)
+    return llvm::None;
+  auto AbsolutePath = URI::resolve(*U, HintPath);
+  if (!AbsolutePath)
+    return llvm::None;
+  return *AbsolutePath;
+}
+
+class BackgroundIndexLoader {
+public:
+  void load(PathRef MainFile, BackgroundIndexStorage *Storage);
+
+  /// Consumes the loader and returns all shards.
+  std::vector<LoadedShard> takeResult() &&;
+
+private:
+  struct CachedShard {
+    LoadedShard LS;
+    // Stores the absolute paths for the direct dependencies of this file.
+    std::vector<Path> Edges;
+  };
+  const CachedShard &loadShard(PathRef StartSourceFile,
+                               BackgroundIndexStorage *Storage);
+
+  // Cache for Storage lookups.
+  llvm::StringMap<CachedShard> LoadedShards;
+
+  // References are into the AbsolutePaths in LoadedShards.
+  llvm::DenseMap<PathRef, PathRef> FileToTU;
+};
+
+const BackgroundIndexLoader::CachedShard &
+BackgroundIndexLoader::loadShard(PathRef StartSourceFile,
+                                 BackgroundIndexStorage *Storage) {
+  auto It = LoadedShards.try_emplace(StartSourceFile);
+  CachedShard &CS = It.first->getValue();
+  if (!It.second)
+    return CS;
+
+  LoadedShard &LS = CS.LS;
+  LS.AbsolutePath = StartSourceFile.str();
+  auto Shard = Storage->loadShard(StartSourceFile);
+  if (!Shard || !Shard->Sources) {
+    vlog("Failed to load shard: {0}", StartSourceFile);
+    LS.HadErrors = true;
+    return CS;
+  }
+
+  LS.Shard = std::move(Shard);
+  for (const auto &It : *LS.Shard->Sources) {
+    auto AbsPath = uriToAbsolutePath(It.getKey(), StartSourceFile);
+    if (!AbsPath)
+      continue;
+    // A shard contains only edges for non main-file sources.
+    if (*AbsPath != StartSourceFile) {
+      CS.Edges.push_back(*AbsPath);
+      continue;
+    }
+
+    const IncludeGraphNode &IGN = It.getValue();
+    LS.Digest = IGN.Digest;
+    LS.CountReferences = IGN.Flags & IncludeGraphNode::SourceFlag::IsTU;
+    LS.HadErrors = IGN.Flags & IncludeGraphNode::SourceFlag::HadErrors;
+  }
+  assert(LS.Digest != FileDigest{{0}} && "Digest is empty?");
+  return CS;
+}
+
+void BackgroundIndexLoader::load(PathRef MainFile,
+                                 BackgroundIndexStorage *Storage) {
+  // Following containers points to strings inside LoadedShards.
+  llvm::DenseSet<PathRef> InQueue;
+  std::queue<PathRef> ToVisit;
+
+  const CachedShard &MainShard = loadShard(MainFile, Storage);
+  ToVisit.push(MainShard.LS.AbsolutePath);
+  InQueue.insert(MainShard.LS.AbsolutePath);
+
+  while (!ToVisit.empty()) {
+    PathRef SourceFile = ToVisit.front();
+    ToVisit.pop();
+
+    const CachedShard &CS = loadShard(SourceFile, Storage);
+    FileToTU.try_emplace(CS.LS.AbsolutePath).first->second =
+        MainShard.LS.AbsolutePath;
+
+    for (PathRef Edge : CS.Edges) {
+      if (InQueue.insert(Edge).second)
+        ToVisit.push(Edge);
+    }
+  }
+}
+
+std::vector<LoadedShard> BackgroundIndexLoader::takeResult() && {
+  std::vector<LoadedShard> Result;
+  Result.reserve(LoadedShards.size());
+  for (auto &It : LoadedShards) {
+    Result.push_back(std::move(It.getValue().LS));
+    LoadedShard &LS = Result.back();
+    auto TUsIt = FileToTU.find(LS.AbsolutePath);
+    assert(TUsIt != FileToTU.end() && "No TU registered for the shard");
+    Result.back().DependentTU = TUsIt->second;
+  }
+  return Result;
+}
+} // namespace
+
+std::vector<LoadedShard>
+loadIndexShards(llvm::ArrayRef<Path> MainFiles,
+                BackgroundIndexStorage::Factory &IndexStorageFactory,
+                const GlobalCompilationDatabase &CDB) {
+  BackgroundIndexLoader Loader;
+  for (llvm::StringRef MainFile : MainFiles) {
+    assert(llvm::sys::path::is_absolute(MainFile));
+
+    std::string ProjectRoot;
+    if (auto PI = CDB.getProjectInfo(MainFile))
+      ProjectRoot = std::move(PI->SourceRoot);
+    BackgroundIndexStorage *Storage = IndexStorageFactory(ProjectRoot);
+    Loader.load(MainFile, Storage);
+  }
+  return std::move(Loader).takeResult();
+}
+
+} // namespace clangd
+} // namespace clang

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

 //
 //===----------------------------------------------------------------------===//
 
 #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"
+#include <cstddef>
 
 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.
public:
   // 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();
+  void loadedShard(size_t ShardCount);
   // 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.
+  unsigned LoadedShards; // In the current loading session.
 
   SwapIndex *Target;
   FileSymbols *Source;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

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

maybeRebuild("when background indexer is idle", [this] {
     // (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;
+    LoadedShards = 0;
   ++Loading;
 }
-void BackgroundIndexRebuilder::loadedTU() {
+void BackgroundIndexRebuilder::loadedShard(size_t ShardCount) {
   std::lock_guard<std::mutex> Lock(Mu);
   assert(Loading);
-  ++LoadedTUs;
+  LoadedShards += ShardCount;
 }
 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();
+    return LoadedShards > 0 || enoughTUsToRebuild();
   });
 }
 
 void BackgroundIndexRebuilder::shutdown() {
   std::lock_guard<std::mutex> Lock(Mu);
   ShouldStop = true;
 }
 

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

TEST_F(BackgroundIndexTest, ShardStorageTest) {
   }
   EXPECT_EQ(CacheHits, 0U);
   EXPECT_EQ(Storage.size(), 2U);
 
   {
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
   EXPECT_EQ(Storage.size(), 2U);
 
   auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(
FS.Files[testPath("root/A.h")] = R"cpp(
       void f_b();
       class A_CC {};
       class A_CCnew {};
       )cpp";
   {
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
 
   // Check if the new symbol has arrived.
   auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(*ShardHeader->Symbols, Contains(Named("A_CCnew")));
 
   // Change source.
   FS.Files[testPath("root/A.cc")] =
       "#include \"A.h\"\nvoid g() { (void)common; }\nvoid f_b() {}";
   {
     CacheHits = 0;
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
 
   // Check if the new symbol has arrived.
   ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(*ShardHeader->Symbols, Contains(Named("A_CCnew")));
TEST_F(BackgroundIndexRebuilderTest, IndexingTUs) {
   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();
+  Rebuilder.loadedShard(10);
+  Rebuilder.loadedShard(20);
   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.loadedShard(1);
   Rebuilder.startLoading();
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(1);
   EXPECT_FALSE(checkRebuild([&] { Rebuilder.doneLoading(); }));
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(1);
   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(); }));