Introduce shard storage to auto-index.

[Eugene.Zelenko]

clangd/index/Background.cpp

31

Please run Clang-format. Headers in sections should be in alphabetical order.

[kadircet]

• clang-format

[sammccall]

clangd/index/Background.h

31

nit: prefer the other way around: doc on the interface, the disk implementation can defer to the interface when there's nothing special to say.
That way callers don't have to be confused about what the semantics are in the general case.

39

Why not use the constructor? what does "directory" mean in the general case?

[sammccall]

clangd/index/Background.cpp

351

nit: i'd suggest doing the writes *after* updating the index, as the latter is user-facing

422

I don't think this is the right place to be reading shards, vs on startup. It means we're doing extra IO whenever a file is reindexed, and more importantly we don't make the index contents available on startup (if any file needs reindexing, reads may get stuck behind it).

(as discussed, we can also leave reading out of this patch entirely)

clangd/index/Background.h

35

nit: we probably don't need "shard" both in the interface name and the method names.

Consider either BackgroundIndexStorage/storeShard or ShardStorage/store

38

this signature looks surprising to me, expected something like `Expected<pair<IndexFileIn, FileDigest>> retrieveShard(ShardID).

e.g. if we have a stale shard for a file when initializing, maybe we want to use it until we can reindex the file? Current interface doesn't allow this.

As discussed offline, it'd also be fine to leave retrieve out of the initial patch.

71

So we have several possible designs on a collision course:

1. One instance of BackgroundIndex for clangd, which has a ShardStorage, which is used for every CDB
2. One instance of BackgroundIndex, which has many ShardStorages, one-per-CDB
3. Many instances of BackgroundIndex, one per CDB, each with a ShardStorage
4. One BackgroundIndex for everything, and one ShardStorage for everything (storage in $HOME/.clangd-index)

What are we aiming for here?

The design of ShardStorage in this patch precludes 1, taking unique_ptr here precludes 2. 3 requires us to revisit some aspects of BackgroundIndex (which is possible). With 4 it isn't obvious how we do multi-config, or how users can manage the storage.

(I'd been assuming 1 or 2, but I don't think we ever discussed it)

132

This class can be hidden in the cpp file.

[kadircet]

clangd/index/Background.cpp

351

but updating index consumes slabs. Is it worth making a copy?

422

What exactly you mean by startup exactly from BackgroundIndex's perspective. Is it enqueueAll ?

clangd/index/Background.h

38

I thought it would be better to get rid of disk io if the file was out-of-date, but using the stale file and deferring indexing also seems like a good idea. Moving with it.

39

Directory refers to the one specified in CompilationDatabase(which is usually the build directory?), sorry for the inconvenience.
I wasn't sure about where we plan to instantiate BackgroundIndex. If you plan to do that initialization at a point in which we already know the build directory we can move that to constructor, especially only to the constructor of DiskBackedIndexStorage.

71

The aim was more like 2. But I forgot the fact that we could be working with multiple CDBs :D

So it might be better to just pass a factory to BackgroundIndex and let it create one ShardStorage for each CDB(which is distinguished by the build directory?) WDYT?

132

But it needs to be passed into the constructor of BackgroundIndex, possibly by the ClangdServer, do you suggest hiding inside ClangdServer?

[sammccall]

clangd/index/Background.cpp

351

Oh, true - probably not. Add a comment?

422

Yes. (Or rather, a task scheduled there)

clangd/index/Background.h

39

tooling::CompileCommand::WorkingDirectory? That doesn't seem especially relevant here.
Or the directory that the CDB was discovered in?

Yes, this seems to be only relevant to DiskBackedIndexStorage

71

Yep, that works for me.

132

We can expose a factory function (e.g. static on ShardStorage)

[kadircet]

clangd/index/Background.h

39

I suppose you meant tooling::CompileCommand::Directory rather than WorkingDirectory ? That is the one I was talking about, why do you think it is irrelevant ?

[kadircet]

• Address comments.
• Move loading from cache out of indexing.
• Use a factory for creation of index storage.

clangd/index/Background.h

39

Changed logic to use Directory passed to enqueueAll or enqueue, which I assume is the directory that the compilation database lives. And I assumed we will have at most one compilation database per directory.

[sammccall]

clangd/index/Background.cpp

52

nit: these micro-optimizations with SmallString seem unlikely to be worthwhile - maybe just take StringRef and return string?

109

absoluteFilePath? (otherwise it's not clear what the path is to)

237

this absolutely can't be done synchronously (it's much slower than calling getAllCompileCommands - this could block clangd for tens of seconds when opening the first file).

Can we drop shard-loading in BackgroundIndex from this patch?

256

you can capture indexstorage by value

clangd/index/Background.h

34

I don't think we should be using a static mutable factory here.

I think I worded the previous comment about exposing DiskBackedIndexStorage badly, what I meant was something like: static unique_ptr<BackgroundIndexStorage> BackgroundIndexStorage::createDiskStorage(StringRef dir)

This would be separate from how the storage strategy is configured on the background index.

34

do we need shared_ptr here?
ISTM either the BackgroundIndex is going to be responsible for caching (in which case it should receive unique_ptr) or the factory is responsible for caching (in which case it can own the functions, and the BackgroundIndex can get a non-owning raw pointer)

38

unused?

42

Prefer llvm::Error over bool, or no return value if you don't think we ever want to handle it.

45

doc stale

48

nit: I think "loadShard" would better emphasize the relationship with storeShard, as load/store is such a common pair. (read/write would similarly work).
Nothing wrong with "retrieve" per se, I just think the pairing is better.

[kadircet]

• Address comments.
• Deleted shard loading logic from backgroundindex.

[kadircet]

clangd/index/Background.cpp

52

It wasn't for optimization purposes actually, sys::path::append only accepts a SmallString and there is no direct conversion between SmallString and std::string(or stringref).

[sammccall]

Thanks, mostly interface tweaks now.

clangd/index/Background.cpp

189

I think it would be simpler to just disallow the storage factory from returning null, and instead have it return a no-op implementation that logs.

Again, this moves the complexity behind an abstraction, where the rest of the code is simpler.

clangd/index/Background.h

36

I like that the implementations treat the shard scheme as opaque, but it'd be good for humans to know that this is going to be a filename in practice.

I'd suggest either changing the param name or adding to the interface comment: "Shards of the index are stored and retrieved independently, keyed by shard identifier - in practice this is a source file name"

39

docs

39

Hmm, we're going to attempt to load the shard corresponding to every file that we have in the CDB, even if the index isn't built yet. So "file doesn't exist" is an expected case (where we don't want to log etc), vs e.g. "file was corrupted" is unexpected and should definitely be logged. How do we distinguish these with this API?

One defensible option would be to have implementations handle errors: return Optional<IndexFileIn>, and have the disk version log errors appropriately and return None.

Another would be Expected<Optional<IndexFileIn>> or so, which is a little messy.

43

docs

64

nit: I'd move the description of IndexStorageFactory semantics to above the Using declaration - generally splitting docs to the most specific decl should make them easier to match up to the code.

65

by the directory they are found in?

118

I do think having the factory own and cache the storages would be simpler:

• it moves more work out of this complicated class with a lot of responsibilities, and puts it somewhere simpler
• it would mean we can just use a single storage object in tests that don't use multiple CDBs.

118

don't accesses to this map need to be locked?

[kadircet]

• Address comments.

clangd/index/Background.h

39

used the former proposition, but with std::unique_ptr instead of llvm::Optional since IndexFileIn is move-only.

[kadircet]

• Get rid off getIndexStorage and use IndexStorageCreator directly.

[sammccall]

just some nits. Main thing is the LoggingStorage - again I think this may have been a misunderstanding.

clangd/index/Background.cpp

75

ADD_FAILURE() I think

93

nit: move constructor to the top?

97

create_directory returns success if the directory exists, so no need for this check

107

Hmm, I'm not really sure what this is for - does someone constructing BackgroundIndex ever not want to define storage *and* not want to control logging?

360

this doesn't give enough context - the logger is global to clangd.

maybe "Failed to write background-index shard for file {0}: {1}"

clangd/index/Background.h

39

This is fine, though Optional works fine with move-only types.

(your comment still says None)

48

we may also want to provide the standard factory - can go in a later patch though.

unittests/clangd/BackgroundIndexTests.cpp

127

nit: seems a little less unusual to avoid the static:

MemoryShardStorage Storage;
auto MemoryStorageFactory = [&](llvm::StringRef){ return &Storage; }
...
EXPECT_THAT(Storage.contains("root/A.h"));

[kadircet]

I fear we still need to expose createDiskStorage, because someone still needs to tell factory to which creator function to use.

[kadircet]

• Address comments.
• Move storage related things to new files.

[kadircet]

• Change factory design to use llvm::unique_function.

[sammccall]

Thanks! Looks good, rest of the nits are obvious or up to you.(

clangd/index/Background.cpp

126

nit: I think the moves of the helper functions (digest/digestFile/getAbsoluteFilePath) can be reverted now

clangd/index/Background.h

45

Some docs? Maybe add a comment like

// The factory provides storage for each CDB. 
// It keeps ownership of the storage instances, and should manage caching itself.
// Factory must be threadsafe and never returns nullptr.

50

nit: create rather than get would be clearer that this returns an independent object each time.

115

(nit: this comment seems redundant with the type/name)

clangd/index/BackgroundIndexStorage.cpp

30 ↗

(On Diff #174187)

nit: can we put a separator like between the path and digest?
I'd suggest . so the file extension is most obviously preserved.

e.g. Foo.cpp.12345.idx vs Foo.cpp_123456.idx vs `Foo.cpp123456.idx'

51 ↗

(On Diff #174187)

nit: just "for index storage"?
it's clearly disk backed if it's going into a directory!
(and the extra qualification may confuse users)

79 ↗

(On Diff #174187)

You're ignoring write errors here (and in fact errors may not have occurred yet due to buffering).

OS.close();
return OS.error();

unittests/clangd/BackgroundIndexTests.cpp

80

Consider pulling this class to the top level, and use it any time you don't care about the storage (to avoid the need for DummyStorage)

91

I think this can just be
Storage[ShardIdentifier] = llvm::to_string(Shard)

(need to include ScopedPrinters.h)

101

this one shouldn't be a failure I think - just return nullptr? (If the test cares about it, the test should depend on it more explicitly).

The corrupt index file is a more obvious "should never happen" case.

110

You never actually assert on this value.

(I'm not sure you need to, but if not just remove it)

126

seems clearer just to inline this in the constructor(, but up to you

142

(even if the StringMap was encapsulated in the MemoryStorage, you can still write this assertion through the public interface: ASSERT_NE(Storage.loadShard("root/A.h"), None)

(you could even assert the actual stored symbols, but may not be worth it)

[kadircet]

• Address comments.

clangd/index/Background.h

 //===--- Background.h - Build an index in a background thread ----*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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 "index/FileIndex.h"
 #include "index/Index.h"
+#include "index/Serialization.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/SHA1.h"
 #include "llvm/Support/Threading.h"
 #include <condition_variable>
 #include <deque>
 #include <string>
 #include <thread>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
+// Handles storage and retrieval of index shards.
+class BackgroundIndexStorage {
+private:
+  static std::function<std::shared_ptr<BackgroundIndexStorage>(llvm::StringRef)>
+      Factory;
+
+public:
+  using FileDigest = decltype(llvm::SHA1::hash({}));
+
+  // Stores given shard associationg with ShardIdentifier, which can be
+  // retrieved later on with the same identifier.
+  virtual bool storeShard(llvm::StringRef ShardIdentifier,
+                          IndexFileOut Shard) const = 0;
+
+  // Retrieves the shard if found, also returns hash for the source file that
+  // generated this shard.
+  virtual llvm::Expected<IndexFileIn>
+  retrieveShard(llvm::StringRef ShardIdentifier) const = 0;
+
+  template <typename T> static void setStorageFactory(T Factory) {
+    BackgroundIndexStorage::Factory = Factory;
+  }
+
+  static std::shared_ptr<BackgroundIndexStorage>
+  getForDirectory(llvm::StringRef Directory) {
+    if (!Factory)
+      return nullptr;
+    return Factory(Directory);
+  }
+};
+
 // Builds an in-memory index by by running the static indexer action over
 // all commands in a compilation database. Indexing happens in the background.
 // FIXME: it should also persist its state on disk for fast start.
 // FIXME: it should watch for changes to files on disk.
 class BackgroundIndex : public SwapIndex {
 public:
   // FIXME: resource-dir injection should be hoisted somewhere common.
-  BackgroundIndex(Context BackgroundContext, StringRef ResourceDir,
+  BackgroundIndex(Context BackgroundContext, llvm::StringRef ResourceDir,
                   const FileSystemProvider &, ArrayRef<std::string> URISchemes,
                   size_t ThreadPoolSize = llvm::hardware_concurrency());
   ~BackgroundIndex(); // Blocks while the current task finishes.
 
   // Enqueue a translation unit for indexing.
   // The indexing happens in a background thread, so the symbols will be
   // available sometime later.
   void enqueue(llvm::StringRef Directory, tooling::CompileCommand);
   // Index all TUs described in the compilation database.
   void enqueueAll(llvm::StringRef Directory,
                   const tooling::CompilationDatabase &);
 
   // Cause background threads to stop after ther current task, any remaining
   // tasks will be discarded.
   void stop();
 
   // Wait until the queue is empty, to allow deterministic testing.
   void blockUntilIdleForTest();
 
   using FileDigest = decltype(llvm::SHA1::hash({}));
 
 private:
   /// Given index results from a TU, only update files in \p FilesToUpdate.
   void update(llvm::StringRef MainFile, SymbolSlab Symbols, RefSlab Refs,
-              const llvm::StringMap<FileDigest> &FilesToUpdate);
+              const llvm::StringMap<FileDigest> &FilesToUpdate,
+              BackgroundIndexStorage *IndexStorage);
+  void loadShard(BackgroundIndexStorage *IndexStorage,
+                 const tooling::CompileCommand &Cmd);
+  void loadShards(BackgroundIndexStorage *IndexStorage,
+                  const std::vector<tooling::CompileCommand> &Cmds);
 
   // configuration
   std::string ResourceDir;
   const FileSystemProvider &FSProvider;
   Context BackgroundContext;
   std::vector<std::string> URISchemes;
 
   // index state
-  llvm::Error index(tooling::CompileCommand);
+  llvm::Error index(tooling::CompileCommand,
+                    BackgroundIndexStorage *IndexStorage);
 
   FileSymbols IndexedSymbols;
   llvm::StringMap<FileDigest> IndexedFileDigests; // Key is absolute file path.
   std::mutex DigestsMu;
 
   // queue management
   using Task = std::function<void()>;
   void run(); // Main loop executed by Thread. Runs tasks from Queue.
-  void enqueueLocked(tooling::CompileCommand Cmd);
+  void enqueueLocked(tooling::CompileCommand Cmd,
+                     std::shared_ptr<BackgroundIndexStorage> IndexStorage);
   std::mutex QueueMu;
   unsigned NumActiveTasks = 0; // Only idle when queue is empty *and* no tasks.
   std::condition_variable QueueCV;
   bool ShouldStop = false;
   std::deque<Task> Queue;
   std::vector<std::thread> ThreadPool; // FIXME: Abstract this away.
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/index/Background.cpp

 #include "index/Serialization.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/SHA1.h"
+#include <memory>
+#include <queue>
 #include <random>
 #include <string>
 
 using namespace llvm;
 namespace clang {
 namespace clangd {
 
+namespace {
+
+static BackgroundIndex::FileDigest digest(StringRef Content) {
+  return SHA1::hash({(const uint8_t *)Content.data(), Content.size()});
+}
+
+static Optional<BackgroundIndex::FileDigest> digestFile(const SourceManager &SM,
+                                                        FileID FID) {
+  bool Invalid = false;
+  StringRef Content = SM.getBufferData(FID, &Invalid);
+  if (Invalid)
+    return None;
+  return digest(Content);
+}
+
+llvm::SmallString<128>
+getShardPathFromFilePath(llvm::SmallString<128> ShardRoot,
+                         llvm::StringRef FilePath) {
+  sys::path::append(ShardRoot, sys::path::filename(FilePath) +
+                                   toHex(digest(FilePath)) + ".idx");
+  return ShardRoot;
+}
+
+// Uses disk as a storage for index shards. Creates a directory called
+// ".clangd-index/" under the path provided during initialize.
+// Note: Requires initialize to be called before storing or retrieval.
+// This class is thread-safe.
+class DiskBackedIndexStorage : public BackgroundIndexStorage {
+  llvm::SmallString<128> DiskShardRoot;
+
+public:
+  llvm::Expected<IndexFileIn>
+  retrieveShard(llvm::StringRef ShardIdentifier) const override {
+    auto ShardPath = getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
+    auto Buffer = MemoryBuffer::getFile(ShardPath);
+    if (!Buffer) {
+      elog("Couldn't retrieve {0}: {1}", ShardPath,
+           Buffer.getError().message());
+      return llvm::make_error<llvm::StringError>(Buffer.getError());
+    }
+    // FIXME: Change readIndexFile to also look at Hash of the source that
+    // generated index and skip if there is a mismatch.
+    return readIndexFile(Buffer->get()->getBuffer());
+  }
+
+  bool storeShard(llvm::StringRef ShardIdentifier,
+                  IndexFileOut Shard) const override {
+    auto ShardPath = getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
+    std::error_code EC;
+    llvm::raw_fd_ostream OS(ShardPath, EC);
+    if (EC) {
+      elog("Failed to open {0} for writing: {1}", ShardPath, EC.message());
+      return false;
+    }
+    OS << Shard;
+    return true;
+  }
+
+  // Initializes DiskShardRoot to (Directory + ".clangd-index/") which is the
+  // base directory for all shard files. After the initialization succeeds all
+  // subsequent calls are no-op.
+  DiskBackedIndexStorage(llvm::StringRef Directory) : DiskShardRoot(Directory) {
+    sys::path::append(DiskShardRoot, ".clangd-index/");
+    if (!llvm::sys::fs::exists(DiskShardRoot)) {
+      std::error_code OK;
+      std::error_code EC = llvm::sys::fs::create_directory(DiskShardRoot);
+      if (EC != OK) {
+        elog("Failed to create {0}: {1}", DiskShardRoot, EC.message());
+      }
+    }
+  }
+};
+
+SmallString<128> getAbsolutePath(const tooling::CompileCommand &Cmd) {
+  SmallString<128> AbsolutePath;
+  if (sys::path::is_absolute(Cmd.Filename)) {
+    AbsolutePath = Cmd.Filename;
+  } else {
+    AbsolutePath = Cmd.Directory;
+    sys::path::append(AbsolutePath, Cmd.Filename);
+  }
+  return AbsolutePath;
+}
+
+} // namespace
+
 BackgroundIndex::BackgroundIndex(Context BackgroundContext,
                                  StringRef ResourceDir,
                                  const FileSystemProvider &FSProvider,
                                  ArrayRef<std::string> URISchemes,
                                  size_t ThreadPoolSize)
     : SwapIndex(make_unique<MemIndex>()), ResourceDir(ResourceDir),
       FSProvider(FSProvider), BackgroundContext(std::move(BackgroundContext)),
       URISchemes(URISchemes) {
   assert(ThreadPoolSize > 0 && "Thread pool size can't be zero.");
   while (ThreadPoolSize--) {
     ThreadPool.emplace_back([this] { run(); });
     // Set priority to low, since background indexing is a long running task we
     // do not want to eat up cpu when there are any other high priority threads.
     // FIXME: In the future we might want a more general way of handling this to
-    // support a tasks with various priorities.
+    // support tasks with various priorities.
     setThreadPriority(ThreadPool.back(), ThreadPriority::Low);
   }
 }
 
 BackgroundIndex::~BackgroundIndex() {
   stop();
   for (auto &Thread : ThreadPool)
     Thread.join();

 
 void BackgroundIndex::blockUntilIdleForTest() {
   std::unique_lock<std::mutex> Lock(QueueMu);
   QueueCV.wait(Lock, [&] { return Queue.empty() && NumActiveTasks == 0; });
 }
 
 void BackgroundIndex::enqueue(StringRef Directory,
                               tooling::CompileCommand Cmd) {
+  auto IndexStorage = BackgroundIndexStorage::getForDirectory(Directory);
+  if (IndexStorage)
+    loadShard(IndexStorage.get(), Cmd);
+  else
+    elog("No index storage for: {0}", Directory);
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
-    enqueueLocked(std::move(Cmd));
+    enqueueLocked(std::move(Cmd), std::move(IndexStorage));
   }
   QueueCV.notify_all();
 }
 
+void BackgroundIndex::loadShard(BackgroundIndexStorage *IndexStorage,
+                                const tooling::CompileCommand &Cmd) {
+  assert(IndexStorage && "No index storage to load from?");
+  auto AbsolutePath = getAbsolutePath(Cmd);
+  auto Shard = IndexStorage->retrieveShard(AbsolutePath);
+  if (Shard) {
+    // FIXME: Updated hashes once we have them in serialized format.
+    // IndexedFileDigests[AbsolutePath] = Hash;
+    IndexedSymbols.update(AbsolutePath,
+                          make_unique<SymbolSlab>(std::move(*Shard->Symbols)),
+                          make_unique<RefSlab>(std::move(*Shard->Refs)));
+
+    vlog("Loaded {0} from storage", AbsolutePath);
+  }
+}
+
+void BackgroundIndex::loadShards(
+    BackgroundIndexStorage *IndexStorage,
+    const std::vector<tooling::CompileCommand> &Cmds) {
+  assert(IndexStorage && "No index storage to load from?");
+  for (const auto &Cmd : Cmds)
+    loadShard(IndexStorage, Cmd);
+  // FIXME: Maybe we should get rid of this one once we start building index
+  // periodically? Especially if we also offload this task onto the queue.
+  vlog("Rebuilding automatic index");
+  reset(IndexedSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge,
+                                  URISchemes));
+}
+
 void BackgroundIndex::enqueueAll(StringRef Directory,
                                  const tooling::CompilationDatabase &CDB) {
   trace::Span Tracer("BackgroundIndexEnqueueCDB");
   // FIXME: this function may be slow. Perhaps enqueue a task to re-read the CDB
   // from disk and enqueue the commands asynchronously?
   auto Cmds = CDB.getAllCompileCommands();
+  auto IndexStorage = BackgroundIndexStorage::getForDirectory(Directory);
+  if (IndexStorage)
+    loadShards(IndexStorage.get(), Cmds);
+  else
+    elog("No index storage for: {0}", Directory);
   SPAN_ATTACH(Tracer, "commands", int64_t(Cmds.size()));
   std::mt19937 Generator(std::random_device{}());
   std::shuffle(Cmds.begin(), Cmds.end(), Generator);
   log("Enqueueing {0} commands for indexing from {1}", Cmds.size(), Directory);
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
     for (auto &Cmd : Cmds)
-      enqueueLocked(std::move(Cmd));
+      enqueueLocked(std::move(Cmd), IndexStorage);
   }
   QueueCV.notify_all();
 }
 
-void BackgroundIndex::enqueueLocked(tooling::CompileCommand Cmd) {
+void BackgroundIndex::enqueueLocked(
+    tooling::CompileCommand Cmd,
+    std::shared_ptr<BackgroundIndexStorage> IndexStorage) {
   Queue.push_back(Bind(
-      [this](tooling::CompileCommand Cmd) {
+      [this](tooling::CompileCommand Cmd,
+             std::shared_ptr<BackgroundIndexStorage> IndexStorage) {
         std::string Filename = Cmd.Filename;
         Cmd.CommandLine.push_back("-resource-dir=" + ResourceDir);
-        if (auto Error = index(std::move(Cmd)))
+        if (auto Error = index(std::move(Cmd), IndexStorage.get()))
           log("Indexing {0} failed: {1}", Filename, std::move(Error));
       },
-      std::move(Cmd)));
-}
-
-static BackgroundIndex::FileDigest digest(StringRef Content) {
-  return SHA1::hash({(const uint8_t *)Content.data(), Content.size()});
-}
-
-static Optional<BackgroundIndex::FileDigest> digestFile(const SourceManager &SM,
-                                                        FileID FID) {
-  bool Invalid = false;
-  StringRef Content = SM.getBufferData(FID, &Invalid);
-  if (Invalid)
-    return None;
-  return digest(Content);
+      std::move(Cmd), std::move(IndexStorage)));
 }
 
 // Resolves URI to file paths with cache.
 class URIToFileCache {
 public:
   URIToFileCache(llvm::StringRef HintPath) : HintPath(HintPath) {}
 
   llvm::StringRef resolve(llvm::StringRef FileURI) {

 private:
   std::string HintPath;
   llvm::StringMap<std::string> URIToPathCache;
 };
 
 /// Given index results from a TU, only update files in \p FilesToUpdate.
 void BackgroundIndex::update(StringRef MainFile, SymbolSlab Symbols,
                              RefSlab Refs,
-                             const StringMap<FileDigest> &FilesToUpdate) {
+                             const StringMap<FileDigest> &FilesToUpdate,
+                             BackgroundIndexStorage *IndexStorage) {
   // Partition symbols/references into files.
   struct File {
     DenseSet<const Symbol *> Symbols;
     DenseSet<const Ref *> Refs;
   };
   StringMap<File> Files;
   URIToFileCache URICache(MainFile);
   for (const auto &Sym : Symbols) {
for (const auto &F : Files) {
     vlog("Update symbols in {0}", Path);
     SymbolSlab::Builder Syms;
     RefSlab::Builder Refs;
     for (const auto *S : F.second.Symbols)
       Syms.insert(*S);
     for (const auto *R : F.second.Refs)
       Refs.insert(RefToIDs[R], *R);
 
+    auto SS = llvm::make_unique<SymbolSlab>(std::move(Syms).build());
+    auto RS = llvm::make_unique<RefSlab>(std::move(Refs).build());
+
+    auto Hash = FilesToUpdate.lookup(Path);
+    // We need to store shards before updating the index, since the latter
+    // consumes slabs.
+    // FIXME: Store Hash in the Shard.
+    if (IndexStorage) {
+      IndexFileOut Shard;
+      Shard.Symbols = SS.get();
+      Shard.Refs = RS.get();
+      IndexStorage->storeShard(Path, Shard);
+    }
+
     std::lock_guard<std::mutex> Lock(DigestsMu);
     // 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.
-    IndexedFileDigests[Path] = FilesToUpdate.lookup(Path);
-    IndexedSymbols.update(Path,
-                          make_unique<SymbolSlab>(std::move(Syms).build()),
-                          make_unique<RefSlab>(std::move(Refs).build()));
+    IndexedFileDigests[Path] = Hash;
+    IndexedSymbols.update(Path, std::move(SS), std::move(RS));
   }
 }
 
 // Creates a filter to not collect index results from files with unchanged
 // digests.
-// \p FileDigests contains file digests for the current indexed files, and all changed files will be added to \p FilesToUpdate.
+// \p FileDigests contains file digests for the current indexed files, and all
+// changed files will be added to \p FilesToUpdate.
 decltype(SymbolCollector::Options::FileFilter) createFileFilter(
     const llvm::StringMap<BackgroundIndex::FileDigest> &FileDigests,
     llvm::StringMap<BackgroundIndex::FileDigest> &FilesToUpdate) {
   return [&FileDigests, &FilesToUpdate](const SourceManager &SM, FileID FID) {
     StringRef Path;
     if (const auto *F = SM.getFileEntryForID(FID))
       Path = F->getName();
     if (Path.empty())
return [&FileDigests, &FilesToUpdate](const SourceManager &SM, FileID FID) {
     if (D != FileDigests.end() && D->second == Digest)
       return false; // Skip files that haven't changed.
 
     FilesToUpdate[AbsPath] = *Digest;
     return true;
   };
 }
 
-Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
+Error BackgroundIndex::index(tooling::CompileCommand Cmd,
+                             BackgroundIndexStorage *IndexStorage) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
-  SmallString<128> AbsolutePath;
-  if (sys::path::is_absolute(Cmd.Filename)) {
-    AbsolutePath = Cmd.Filename;
-  } else {
-    AbsolutePath = Cmd.Directory;
-    sys::path::append(AbsolutePath, Cmd.Filename);
-  }
+  SmallString<128> AbsolutePath = getAbsolutePath(Cmd);
 
   auto FS = FSProvider.getFileSystem();
   auto Buf = FS->getBufferForFile(AbsolutePath);
   if (!Buf)
     return errorCodeToError(Buf.getError());
   auto Hash = digest(Buf->get()->getBuffer());
 
   // Take a snapshot of the digests to avoid locking for each file in the TU.
   llvm::StringMap<FileDigest> DigestsSnapshot;
   {
     std::lock_guard<std::mutex> Lock(DigestsMu);
     if (IndexedFileDigests.lookup(AbsolutePath) == Hash) {
       vlog("No need to index {0}, already up to date", AbsolutePath);
       return Error::success();
     }
 
     DigestsSnapshot = IndexedFileDigests;
   }
 
   log("Indexing {0}", Cmd.Filename, toHex(Hash));
   ParseInputs Inputs;
   Inputs.FS = std::move(FS);
   Inputs.FS->setCurrentWorkingDirectory(Cmd.Directory);
Error BackgroundIndex::index(tooling::CompileCommand Cmd,
   if (!Action->Execute())
     return createStringError(inconvertibleErrorCode(), "Execute() failed");
   Action->EndSourceFile();
 
   log("Indexed {0} ({1} symbols, {2} refs)", Inputs.CompileCommand.Filename,
       Symbols.size(), Refs.numRefs());
   SPAN_ATTACH(Tracer, "symbols", int(Symbols.size()));
   SPAN_ATTACH(Tracer, "refs", int(Refs.numRefs()));
-  update(AbsolutePath, std::move(Symbols), std::move(Refs), FilesToUpdate);
+  update(AbsolutePath, std::move(Symbols), std::move(Refs), FilesToUpdate,
+         IndexStorage);
   {
     // Make sure hash for the main file is always updated even if there is no
     // index data in it.
     std::lock_guard<std::mutex> Lock(DigestsMu);
     IndexedFileDigests[AbsolutePath] = Hash;
   }
 
   // FIXME: this should rebuild once-in-a-while, not after every file.
   //       At that point we should use Dex, too.
   vlog("Rebuilding automatic index");
   reset(IndexedSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge,
                                   URISchemes));
 
   return Error::success();
 }
 
+std::function<std::shared_ptr<BackgroundIndexStorage>(llvm::StringRef)>
+    BackgroundIndexStorage::Factory = nullptr;
+
 } // namespace clangd
 } // namespace clang

unittests/clangd/BackgroundIndexTests.cpp

TEST(BackgroundIndexTest, IndexTwoFiles) {
   EXPECT_THAT(Syms, UnorderedElementsAre(Named("common")));
   auto Common = *Syms.begin();
   EXPECT_THAT(getRefs(Idx, Common.ID),
               RefsAre({FileURI("unittest:///root/A.h"),
                        FileURI("unittest:///root/A.cc"),
                        FileURI("unittest:///root/B.cc")}));
 }
 
+TEST(BackgroundIndexTest, ShardStorageTest) {
+  class MemoryShardStorage : public BackgroundIndexStorage {
+    mutable std::mutex StorageMu;
+    llvm::StringMap<std::string> &Storage;
+    size_t &CacheHits;
+
+  public:
+    MemoryShardStorage(llvm::StringMap<std::string> &Storage, size_t &CacheHits)
+        : Storage(Storage), CacheHits(CacheHits) {}
+
+    bool storeShard(llvm::StringRef ShardIdentifier,
+                    IndexFileOut Shard) const override {
+      std::lock_guard<std::mutex> Lock(StorageMu);
+      std::string &str = Storage[ShardIdentifier];
+      llvm::raw_string_ostream OS(str);
+      OS << Shard;
+      OS.flush();
+      return true;
+    }
+    llvm::Expected<IndexFileIn>
+    retrieveShard(llvm::StringRef ShardIdentifier) const override {
+      std::lock_guard<std::mutex> Lock(StorageMu);
+      if (Storage.find(ShardIdentifier) == Storage.end())
+        return llvm::make_error<llvm::StringError>(
+            "Shard not found.", llvm::inconvertibleErrorCode());
+      auto IndexFile = readIndexFile(Storage[ShardIdentifier]);
+      if (!IndexFile)
+        return IndexFile;
+      CacheHits++;
+      return IndexFile;
+    }
+    bool initialize(llvm::StringRef Directory) { return true; }
+  };
+  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")] =
+      "#include \"A.h\"\nvoid g() { (void)common; }";
+
+  llvm::StringMap<std::string> Storage;
+  size_t CacheHits = 0;
+  BackgroundIndexStorage::setStorageFactory(
+      [&Storage, &CacheHits](llvm::StringRef) {
+        return std::make_shared<MemoryShardStorage>(Storage, CacheHits);
+      });
+
+  tooling::CompileCommand Cmd;
+  Cmd.Filename = testPath("root/A.cc");
+  Cmd.Directory = testPath("root");
+  Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+  // Check nothing is loaded from Storage, but A.cc and A.h has been stored.
+  {
+    BackgroundIndex Idx(Context::empty(), "", FS, /*URISchemes=*/{"unittest"});
+    Idx.enqueue(testPath("root"), Cmd);
+    Idx.blockUntilIdleForTest();
+  }
+  EXPECT_EQ(CacheHits, 0U);
+  EXPECT_EQ(Storage.size(), 2U);
+  EXPECT_NE(Storage.find(testPath("root/A.h")), Storage.end());
+  EXPECT_NE(Storage.find(testPath("root/A.cc")), Storage.end());
+
+  // Check A.cc has been loaded from cache.
+  {
+    BackgroundIndex Idx(Context::empty(), "", FS, /*URISchemes=*/{"unittest"});
+    Idx.enqueue(testPath("root"), Cmd);
+    Idx.blockUntilIdleForTest();
+  }
+  EXPECT_EQ(CacheHits, 1U);
+  EXPECT_EQ(Storage.size(), 2U);
+  EXPECT_NE(Storage.find(testPath("root/A.h")), Storage.end());
+  EXPECT_NE(Storage.find(testPath("root/A.cc")), Storage.end());
+}
+
 } // namespace clangd
 } // namespace clang