Introduce shard storage to auto-index.

[Eugene.Zelenko]

clangd/index/Background.cpp

29

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

252

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

326

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.

51

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)

107

This class can be hidden in the cpp file.

[kadircet]

clangd/index/Background.cpp

252

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

326

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.

51

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?

107

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

252

Oh, true - probably not. Add a comment?

326

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

51

Yep, that works for me.

107

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)

141

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?

159

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

127

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

44

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.

45

by the directory they are found in?

94

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.

94

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?

261

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

36–64

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.

91

(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 {
 
+// Base class for Shard Storage operations. See DiskShardStorage for more info.
+class ShardStorage {
+public:
+  using FileDigest = decltype(llvm::SHA1::hash({}));
+  virtual bool storeShard(llvm::StringRef ShardIdentifier,
+                          IndexFileOut Shard) const = 0;
+  virtual llvm::Expected<IndexFileIn>
+  retrieveShard(llvm::StringRef ShardIdentifier, FileDigest Hash) const = 0;
+  virtual bool initialize(llvm::StringRef Directory) = 0;
+};
+
 // 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,
+                  std::unique_ptr<ShardStorage> IndexShardStorage = 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.
private:
   void update(llvm::StringRef MainFile, SymbolSlab Symbols, RefSlab Refs,
               const llvm::StringMap<FileDigest> &FilesToUpdate);
 
   // configuration
   std::string ResourceDir;
   const FileSystemProvider &FSProvider;
   Context BackgroundContext;
   std::vector<std::string> URISchemes;
+  std::unique_ptr<ShardStorage> IndexShardStorage;
 
   // index state
   llvm::Error index(tooling::CompileCommand);
 
   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);
   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.
 };
 
+// Handles storage and retrieval of index shards into disk. Requires Initialize
+// to be called before storing or retrieval. Creates a directory called
+// ".clangd-index/" under the path provided during initialize. This class is
+// thread-safe.
+class DiskShardStorage : public ShardStorage {
+  mutable std::mutex DiskShardRootMu;
+  llvm::SmallString<128> DiskShardRoot;
+  bool Initialized;
+
+public:
+  // Retrieves the shard if found and contents are consistent with the provided
+  // Hash.
+  llvm::Expected<IndexFileIn> retrieveShard(llvm::StringRef ShardIdentifier,
+                                            FileDigest Hash) const;
+
+  // Stores given shard with name ShardIdentifier under initialized directory.
+  bool storeShard(llvm::StringRef ShardIdentifier, IndexFileOut Shard) const;
+
+  // Initializes DiskShardRoot to (Directory + ".clangd-index/") which is the
+  // base directory for all shard files. After the initialization succeeds all
+  // subsequent calls or no-op.
+  bool initialize(llvm::StringRef Directory);
+};
+
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/index/Background.cpp

 #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 <random>
 #include <string>
+#include <queue>
+#include <memory>
 
 using namespace llvm;
 namespace clang {
 namespace clangd {
 
-BackgroundIndex::BackgroundIndex(Context BackgroundContext,
-                                 StringRef ResourceDir,
-                                 const FileSystemProvider &FSProvider,
-                                 ArrayRef<std::string> URISchemes,
-                                 size_t ThreadPoolSize)
+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;
+}
+
+} // namespace
+
+BackgroundIndex::BackgroundIndex(
+    Context BackgroundContext, StringRef ResourceDir,
+    const FileSystemProvider &FSProvider, ArrayRef<std::string> URISchemes,
+    std::unique_ptr<ShardStorage> IndexShardStorage, size_t ThreadPoolSize)
     : SwapIndex(make_unique<MemIndex>()), ResourceDir(ResourceDir),
       FSProvider(FSProvider), BackgroundContext(std::move(BackgroundContext)),
-      URISchemes(URISchemes) {
+      URISchemes(URISchemes), IndexShardStorage(std::move(IndexShardStorage)) {
   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::stop() {
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
     ShouldStop = true;
   }
   QueueCV.notify_all();
 }
 
 void BackgroundIndex::run() {
   WithContext Background(BackgroundContext.clone());
   while (true) {
     Optional<Task> Task;
     {
       std::unique_lock<std::mutex> Lock(QueueMu);
       QueueCV.wait(Lock, [&] { return ShouldStop || !Queue.empty(); });
       if (ShouldStop) {
         Queue.clear();
         QueueCV.notify_all();
         return;
       }
       ++NumActiveTasks;
       Task = std::move(Queue.front());
       Queue.pop_front();
     }
     (*Task)();
     {
       std::unique_lock<std::mutex> Lock(QueueMu);
       assert(NumActiveTasks > 0 && "before decrementing");
       --NumActiveTasks;
     }
     QueueCV.notify_all();
   }
 }
 
 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) {
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
     enqueueLocked(std::move(Cmd));
   }
   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();
   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));
   }
   QueueCV.notify_all();
 }
 
 void BackgroundIndex::enqueueLocked(tooling::CompileCommand Cmd) {
+  // Initialize storage to project root. Since Initialize is no-op for multiple
+  // calls we can simply call it for each file.
+  if (IndexShardStorage && !IndexShardStorage->initialize(Cmd.Directory)) {
+    elog("Failed to initialize shard storage");
+    IndexShardStorage.reset();
+  }
   Queue.push_back(Bind(
       [this](tooling::CompileCommand Cmd) {
         std::string Filename = Cmd.Filename;
         Cmd.CommandLine.push_back("-resource-dir=" + ResourceDir);
         if (auto Error = index(std::move(Cmd)))
           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) {
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);
+    // Put shards into storage for subsequent use.
+    // FIXME: Store Hash in the Shard.
+    if (IndexShardStorage) {
+      IndexFileOut Shard;
+      Shard.Symbols = SS.get();
+      Shard.Refs = RS.get();
+      IndexShardStorage->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.
 decltype(SymbolCollector::Options::FileFilter) createFileFilter(
     const llvm::StringMap<BackgroundIndex::FileDigest> &FileDigests,
Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
 
   // 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();
+    } else if (IndexShardStorage) { // Check if shard storage has the index.
+      auto Shard = IndexShardStorage->retrieveShard(AbsolutePath, Hash);
+      if (Shard) {
+        // FIXME: We might still want to re-index headers.
+        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);
+        return Error::success();
+      }
     }
 
     DigestsSnapshot = IndexedFileDigests;
   }
 
   log("Indexing {0}", Cmd.Filename, toHex(Hash));
   ParseInputs Inputs;
   Inputs.FS = std::move(FS);
Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
   //       At that point we should use Dex, too.
   vlog("Rebuilding automatic index");
   reset(IndexedSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge,
                                   URISchemes));
 
   return Error::success();
 }
 
+llvm::Expected<IndexFileIn>
+DiskShardStorage::retrieveShard(llvm::StringRef ShardIdentifier,
+                                FileDigest Hash) const {
+  assert(Initialized && "Not initialized?");
+  llvm::SmallString<128> ShardPath;
+  {
+    std::lock_guard<std::mutex> Lock(DiskShardRootMu);
+    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 DiskShardStorage::storeShard(llvm::StringRef ShardIdentifier,
+                                  IndexFileOut Shard) const {
+  assert(Initialized && "Not initialized?");
+  llvm::SmallString<128> ShardPath;
+  {
+    std::lock_guard<std::mutex> Lock(DiskShardRootMu);
+    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;
+}
+
+bool DiskShardStorage::initialize(llvm::StringRef Directory) {
+  if (Initialized)
+    return true;
+  std::lock_guard<std::mutex> Lock(DiskShardRootMu);
+  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());
+      return Initialized = false;
+    }
+  }
+  return Initialized = true;
+}
+
 } // 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 ShardStorage {
+    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 {
+      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,
+                                              FileDigest Hash) const {
+      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;
+  tooling::CompileCommand Cmd;
+  Cmd.Filename = testPath("root/A.cc");
+  Cmd.Directory = testPath("root");
+  Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+  {
+    BackgroundIndex Idx(
+        Context::empty(), "", FS, /*URISchemes=*/{"unittest"},
+        /*IndexShardStorage=*/
+        llvm::make_unique<MemoryShardStorage>(Storage, CacheHits));
+    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());
+
+  {
+    BackgroundIndex Idx(
+        Context::empty(), "", FS, /*URISchemes=*/{"unittest"},
+        /*IndexShardStorage=*/
+        llvm::make_unique<MemoryShardStorage>(Storage, CacheHits));
+    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());
+  // B_CC is dropped as we don't collect symbols from A.h in this compilation.
+}
+
 } // namespace clangd
 } // namespace clang