Introduce shard storage to auto-index.

[Eugene.Zelenko]

clangd/index/Background.cpp

32

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

335

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

407

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.

65

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)

127

This class can be hidden in the cpp file.

[kadircet]

clangd/index/Background.cpp

335

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

407

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.

65

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?

127

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

335

Oh, true - probably not. Add a comment?

407

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

65

Yep, that works for me.

127

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

53

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

110

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

227

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?

241

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

53

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

212

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

53

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.

54

by the directory they are found in?

109

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.

109

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

76

ADD_FAILURE() I think

94

nit: move constructor to the top?

98

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

108

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?

344

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

146

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.

106

(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. Both store and load
+// operations can be called from multiple-threads concurrently.
+class BackgroundIndexStorage {
+public:
+  // Shards of the index are stored and retrieved independently, keyed by shard
+  // identifier - in practice this is a source file name
+  virtual llvm::Error storeShard(llvm::StringRef ShardIdentifier,
+                                 IndexFileOut Shard) const = 0;
+
+  // Tries to load shard with given identifier, returns llvm::None if shard
+  // couldn't be loaded.
+  virtual std::unique_ptr<IndexFileIn>
+  loadShard(llvm::StringRef ShardIdentifier) const = 0;
+
+  // Creates an Index Storage that saves shards into disk. Index storage uses
+  // CDBDirectory + ".clangd-index/" as the folder to save shards.
+  static BackgroundIndexStorage *
+  createDiskStorage(llvm::StringRef CDBDirectory);
+};
+
 // 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:
+  // A factory function used to create IndexStorage units for each compilation
+  // database. Those databases are identified by the directory they are found
+  // in. The function must be thread-safe.
+  using IndexStorageFactory =
+      std::function<BackgroundIndexStorage *(llvm::StringRef)>;
   // 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,
+                  IndexStorageFactory IndexStorageCreator = nullptr,
                   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);
 
   // 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;
 
+  // index storage
+  BackgroundIndexStorage *getIndexStorage(llvm::StringRef CDBDirectory);
+  IndexStorageFactory IndexStorageCreator;
+
   // 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,
+                     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);
+}
+
+std::string getShardPathFromFilePath(llvm::StringRef ShardRoot,
+                                     llvm::StringRef FilePath) {
+  llvm::SmallString<128> ShardRootSS(ShardRoot);
+  sys::path::append(ShardRootSS, sys::path::filename(FilePath) +
+                                     toHex(digest(FilePath)) + ".idx");
+  return ShardRoot.str();
+}
+
+// Uses disk as a storage for index shards. Creates a directory called
+// ".clangd-index/" under the path provided during construction.
+class DiskBackedIndexStorage : public BackgroundIndexStorage {
+  std::string DiskShardRoot;
+
+public:
+  std::unique_ptr<IndexFileIn>
+  loadShard(llvm::StringRef ShardIdentifier) const override {
+    const std::string ShardPath =
+        getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
+    auto Buffer = MemoryBuffer::getFile(ShardPath);
+    if (!Buffer)
+      return nullptr;
+    if (auto I = readIndexFile(Buffer->get()->getBuffer()))
+      return llvm::make_unique<IndexFileIn>(std::move(*I));
+    else
+      elog("Error while reading shard {0}: {1}", ShardIdentifier,
+           I.takeError());
+    return nullptr;
+  }
+
+  llvm::Error 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)
+      return errorCodeToError(EC);
+    OS << Shard;
+    return llvm::Error::success();
+  }
+
+  // Sets DiskShardRoot to (Directory + ".clangd-index/") which is the base
+  // directory for all shard files.
+  DiskBackedIndexStorage(llvm::StringRef Directory) {
+    llvm::SmallString<128> CDBDirectory(Directory);
+    sys::path::append(CDBDirectory, ".clangd-index/");
+    DiskShardRoot = CDBDirectory.str();
+    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());
+      }
+    }
+  }
+};
+
+class LoggingIndexStorage : public BackgroundIndexStorage {
+public:
+  std::unique_ptr<IndexFileIn>
+  loadShard(llvm::StringRef ShardIdentifier) const override {
+    vlog("called loadShard({0})", ShardIdentifier);
+    return nullptr;
+  }
+  llvm::Error storeShard(llvm::StringRef ShardIdentifier,
+                         IndexFileOut Shard) const override {
+    vlog("Called storeShard({0}, ...)", ShardIdentifier);
+    return llvm::Error::success();
+  }
+  static BackgroundIndexStorage *
+  createLoggingStorage(llvm::StringRef CDBDirectory) {
+    static LoggingIndexStorage LIS;
+    vlog("Creating Logging Storage for: {0}", CDBDirectory);
+    return &LIS;
+  }
+};
+
+std::string getAbsoluteFilePath(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.str();
+}
+
+} // namespace
+
 BackgroundIndex::BackgroundIndex(Context BackgroundContext,
                                  StringRef ResourceDir,
                                  const FileSystemProvider &FSProvider,
                                  ArrayRef<std::string> URISchemes,
+                                 IndexStorageFactory IndexStorageCreator,
                                  size_t ThreadPoolSize)
     : SwapIndex(make_unique<MemIndex>()), ResourceDir(ResourceDir),
       FSProvider(FSProvider), BackgroundContext(std::move(BackgroundContext)),
-      URISchemes(URISchemes) {
+      URISchemes(URISchemes),
+      IndexStorageCreator(IndexStorageCreator
+                              ? std::move(IndexStorageCreator)
+                              : LoggingIndexStorage::createLoggingStorage) {
   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) {
+  BackgroundIndexStorage *IndexStorage = getIndexStorage(Directory);
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
-    enqueueLocked(std::move(Cmd));
+    enqueueLocked(std::move(Cmd), IndexStorage);
   }
   QueueCV.notify_all();
 }
 
 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();
+  BackgroundIndexStorage *IndexStorage = getIndexStorage(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,
+                                    BackgroundIndexStorage *IndexStorage) {
   Queue.push_back(Bind(
-      [this](tooling::CompileCommand Cmd) {
+      [this, IndexStorage](tooling::CompileCommand Cmd) {
         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))
           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);
-}
-
 // Resolves URI to file paths with cache.
 class URIToFileCache {
 public:
   URIToFileCache(llvm::StringRef HintPath) : HintPath(HintPath) {}
 
   llvm::StringRef resolve(llvm::StringRef FileURI) {
     auto I = URIToPathCache.try_emplace(FileURI);
     if (I.second) {

 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();
+      if (auto Error = IndexStorage->storeShard(Path, Shard))
+        elog("Failed to store shard for {0}: {1}", Path, std::move(Error));
+    }
+
     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);
-  }
+  const std::string AbsolutePath = getAbsoluteFilePath(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();
 }
 
+BackgroundIndexStorage *
+BackgroundIndex::getIndexStorage(llvm::StringRef CDBDirectory) {
+  return IndexStorageCreator(CDBDirectory);
+}
+
+BackgroundIndexStorage *
+BackgroundIndexStorage::createDiskStorage(llvm::StringRef CDBDirectory) {
+  // Maps CDB Directory to index storage.
+  static llvm::StringMap<std::unique_ptr<BackgroundIndexStorage>>
+      IndexStorageMap;
+  static std::mutex IndexStorageMapMu;
+
+  std::lock_guard<std::mutex> Lock(IndexStorageMapMu);
+  auto &IndexStorage = IndexStorageMap[CDBDirectory];
+  if (!IndexStorage)
+    IndexStorage = llvm::make_unique<DiskBackedIndexStorage>(CDBDirectory);
+  return IndexStorage.get();
+}
+
 } // 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, ShardStorageWriteTest) {
+  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) {}
+    llvm::Error 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 llvm::Error::success();
+    }
+    std::unique_ptr<IndexFileIn>
+    loadShard(llvm::StringRef ShardIdentifier) const override {
+      std::lock_guard<std::mutex> Lock(StorageMu);
+      if (Storage.find(ShardIdentifier) == Storage.end()) {
+        elog("Shard {0}: not found.", ShardIdentifier);
+        return nullptr;
+      }
+      auto IndexFile = readIndexFile(Storage[ShardIdentifier]);
+      if (!IndexFile) {
+        elog("Error while reading {0}: {1}.", ShardIdentifier,
+             IndexFile.takeError());
+        return nullptr;
+      }
+      CacheHits++;
+      return llvm::make_unique<IndexFileIn>(std::move(*IndexFile));
+    }
+  };
+  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;
+  BackgroundIndex::IndexStorageFactory MemoryStorageFactory =
+      [&Storage, &CacheHits](llvm::StringRef) {
+        static MemoryShardStorage MSS(Storage, CacheHits);
+        return &MSS;
+      };
+
+  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"},
+                        MemoryStorageFactory);
+    Idx.enqueue(testPath("root"), Cmd);
+    Idx.blockUntilIdleForTest();
+  }
+  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