[clangd] Add symbol slab size to index memory consumption estimates

[kbobyrev]

Summary

Currently, SymbolIndex::estimateMemoryUsage() returns the "overhead" estimate, i.e. the estimate of the Index data structure excluding SymbolSlab size. This patch propagates information about paired SymbolSlab size where necessary.

[kbobyrev]

I should use SymbolSlab.bytes() instead of SymbolSlab.size(). Either way, something seems to be wrong, because manually checking the PairedSlabSize made me suspicious.

[kbobyrev]

Resolved the issues. Measurements show that static Dex index for LLVM takes ~140 MB, ~80 MB of which is the size of SymbolSlab.

[sammccall]

This looks reasonable!
This is going to conflict with D51422, you might want to rebase.

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

23

could you make these names somewhat more abstract (e.g. BackingMemory rather than SlabSize)?
The index shouldn't know/care which configuration of slabs it's built from.

[kbobyrev]

I'll rebase it on top of D51422 and address the last comment before proceeding to the review.

[kadircet]

clang-tools-extra/clangd/index/dex/DexIndex.cpp

231

Why not use LookupTable.getMemorySize() directly? Is it intentionally for not counting non-used but allocated buckets? Same for the below two as well.

[kadircet]

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

144

Why not bytes here as well?

[kbobyrev]

Rebase on top of recent revisions.

[ioeric]

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

157

This can be a bit tricky. Generally, the size of a FileIndex is the size of FSymbols plus the overhead of the underlying index, but there can be multiple snapshots of symbols in SwapIndex. I think we can treat them as burst and only consider the size of the latest snapshot.

FSymbols and the index have shared ownership to the symbol corpus. We can make either of them bookkeep the size, but I think it might be more straightforward to let the index "own" the size. You could set the payload size when creating the index, and you wouldn't need to expose estimateMemoryUsage from FSymbols.

clang-tools-extra/clangd/index/dex/DexIndex.h

45

For this constructor, the index doesn't own the backing data. We should probably not include it in the estimation?

(same for MemIndex)

60

nit: maybe call BackingDataSize? BackingMemory sounds like the actual memory. Maybe also put the parameter right after BackingData, as they are closely related?

(same for MemIndex)

106

nit: =0?

[sammccall]

LG with a few nits

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

140

this isn't safe as it doesn't acquire the mutex

157

(I had trouble parsing this at first - If I understand right, you want the *memIndex* to own the size, and FileIndex to just inherit SwapIndex::estimateMemoryUsage which delegates to MemIndex::estimateMemoryUsage?)

I like this idea, it's simpler and also means one less place you need to lock, which is always nice.

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

24

I don't think defaulting this parameter makes sense, seems like an easy source of bugs
(and in 3 other cases)

63

just // Size of memory retained by KeepAlive?
Don't mention Slab, this class doesn't know about it.

(similar in Dex)

[ioeric]

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

157

Yes, that's correct. Thanks for the clarification!

[kbobyrev]

Address a round of comments.

[sammccall]

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

138

this size should be calculated from the slabs above

[sammccall]

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

131

also the refslabs and refsstorage

[kbobyrev]

Oh, I thought they're empty.

[kbobyrev]

Sync with HEAD.

[kbobyrev]

Remove accidental formatting changes.

[ioeric]

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

512

nit: I'd avoid irrelevant changes.

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

20

Include Refs slab size?

76

Include Refs size?

[kbobyrev]

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

512

This one is out of sync; fixed in the previous update :)

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

for (auto &Sym : MergedRefs) {
       std::copy(SymRefs.begin(), SymRefs.end(), back_inserter(RefsStorage));
       AllRefs.try_emplace(
           Sym.first,
           ArrayRef<Ref>(&RefsStorage[RefsStorage.size() - SymRefs.size()],
                         SymRefs.size()));
     }
   }
 
+  size_t StorageSize = 0;
+  for (const auto &Slab : SymbolSlabs)
+    StorageSize += Slab->bytes();
   // Index must keep the slabs and contiguous ranges alive.
   return llvm::make_unique<MemIndex>(
       llvm::make_pointee_range(AllSymbols), std::move(AllRefs),
       std::make_tuple(std::move(SymbolSlabs), std::move(RefSlabs),
-                      std::move(RefsStorage)));
+                      std::move(RefsStorage)),
+      StorageSize);
 }
 
 void FileIndex::update(PathRef Path, ASTContext *AST,
                        std::shared_ptr<Preprocessor> PP,
                        llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls) {
   if (!AST) {
     FSymbols.update(Path, nullptr, nullptr);
   } else {
     assert(PP);
     auto Contents = indexAST(*AST, PP, TopLevelDecls, URISchemes);
     FSymbols.update(Path,
                     llvm::make_unique<SymbolSlab>(std::move(Contents.first)),
                     llvm::make_unique<RefSlab>(std::move(Contents.second)));
   }
   reset(FSymbols.buildMemIndex());
 }
 
 } // namespace clangd
 } // namespace clang

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

public:
   SwapIndex(std::unique_ptr<SymbolIndex> Index = nullptr)
       : Index(std::move(Index)) {}
   void reset(std::unique_ptr<SymbolIndex>);
 
   // SymbolIndex methods delegate to the current index, which is kept alive
   // until the call returns (even if reset() is called).
   bool fuzzyFind(const FuzzyFindRequest &,
                  llvm::function_ref<void(const Symbol &)>) const override;
+
   void lookup(const LookupRequest &,
               llvm::function_ref<void(const Symbol &)>) const override;
+
   void refs(const RefsRequest &,
             llvm::function_ref<void(const Ref &)>) const override;
+
   size_t estimateMemoryUsage() const override;
 
 private:
   std::shared_ptr<SymbolIndex> snapshot() const;
   mutable std::mutex Mutex;
   std::shared_ptr<SymbolIndex> Index;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_INDEX_H

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

std::shared_ptr<SymbolIndex> SwapIndex::snapshot() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return Index;
 }
 
 bool SwapIndex::fuzzyFind(const FuzzyFindRequest &R,
                           llvm::function_ref<void(const Symbol &)> CB) const {
   return snapshot()->fuzzyFind(R, CB);
 }
+
 void SwapIndex::lookup(const LookupRequest &R,
                        llvm::function_ref<void(const Symbol &)> CB) const {
   return snapshot()->lookup(R, CB);
 }
+
 void SwapIndex::refs(const RefsRequest &R,
                      llvm::function_ref<void(const Ref &)> CB) const {
   return snapshot()->refs(R, CB);
 }
+
 size_t SwapIndex::estimateMemoryUsage() const {
   return snapshot()->estimateMemoryUsage();
 }
 
 } // namespace clangd
 } // namespace clang

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

 
 namespace clang {
 namespace clangd {
 
 /// MemIndex is a naive in-memory index suitable for a small set of symbols.
 class MemIndex : public SymbolIndex {
 public:
   MemIndex() = default;
   // All symbols and refs must outlive this index.
   template <typename SymbolRange, typename RefRange>
   MemIndex(SymbolRange &&Symbols, RefRange &&Refs) {
     for (const Symbol &S : Symbols)
       Index[S.ID] = &S;
     for (const std::pair<SymbolID, llvm::ArrayRef<Ref>> &R : Refs)
       this->Refs.try_emplace(R.first, R.second.begin(), R.second.end());
   }
   // Symbols are owned by BackingData, Index takes ownership.
   template <typename SymbolRange, typename RefRange, typename Payload>
-  MemIndex(SymbolRange &&Symbols, RefRange &&Refs, Payload &&BackingData)
+  MemIndex(SymbolRange &&Symbols, RefRange &&Refs, Payload &&BackingData,
+           size_t BackingDataSize)
       : MemIndex(std::forward<SymbolRange>(Symbols),
                  std::forward<RefRange>(Refs)) {
     KeepAlive = std::shared_ptr<void>(
         std::make_shared<Payload>(std::move(BackingData)), nullptr);
+    this->BackingDataSize = BackingDataSize;
   }
 
   /// Builds an index from slabs. The index takes ownership of the data.
   static std::unique_ptr<SymbolIndex> build(SymbolSlab Symbols, RefSlab Refs);
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void lookup(const LookupRequest &Req,
               llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void refs(const RefsRequest &Req,
             llvm::function_ref<void(const Ref &)> Callback) const override;
 
   size_t estimateMemoryUsage() const override;
 
 private:
   // Index is a set of symbols that are deduplicated by symbol IDs.
   llvm::DenseMap<SymbolID, const Symbol *> Index;
   // A map from symbol ID to symbol refs, support query by IDs.
   llvm::DenseMap<SymbolID, llvm::ArrayRef<Ref>> Refs;
   std::shared_ptr<void> KeepAlive; // poor man's move-only std::any
+  // Size of memory retained by KeepAlive.
+  size_t BackingDataSize = 0;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_MEMINDEX_H

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

 #include "FuzzyMatch.h"
 #include "Logger.h"
 #include "Quality.h"
 
 namespace clang {
 namespace clangd {
 
 std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Slab, RefSlab Refs) {
+  // Store Slab size before it is moved.
+  const auto BackingDataSize = Slab.bytes();
   auto Data = std::make_pair(std::move(Slab), std::move(Refs));
-  return llvm::make_unique<MemIndex>(Data.first, Data.second, std::move(Data));
+  return llvm::make_unique<MemIndex>(Data.first, Data.second, std::move(Data),
+                                     BackingDataSize);
 }
 
 bool MemIndex::fuzzyFind(
     const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   assert(!StringRef(Req.Query).contains("::") &&
          "There must be no :: in query.");
 
if (SymRefs == Refs.end())
       continue;
     for (const auto &O : SymRefs->second)
       if (static_cast<int>(Req.Filter & O.Kind))
         Callback(O);
   }
 }
 
 size_t MemIndex::estimateMemoryUsage() const {
-  return Index.getMemorySize();
+  return Index.getMemorySize() + BackingDataSize;
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/dex/DexIndex.h

 // changed frequently, it's safe to assume that it has to be built only once
 // (when the clangd process starts). Therefore, it can be easier to store built
 // index on disk and then load it if available.
 class DexIndex : public SymbolIndex {
 public:
   // All symbols must outlive this index.
   template <typename Range>
   DexIndex(Range &&Symbols, llvm::ArrayRef<std::string> Schemes)
       : URISchemes(Schemes) {
     // If Schemes don't contain any items, fall back to SymbolCollector's
     // default URI schemes.
     if (URISchemes.empty()) {
       SymbolCollector::Options Opts;
       URISchemes = Opts.URISchemes;
     }
     for (auto &&Sym : Symbols)
       this->Symbols.push_back(&Sym);
     buildIndex();
   }
   // Symbols are owned by BackingData, Index takes ownership.
   template <typename Range, typename Payload>
-  DexIndex(Range &&Symbols, Payload &&BackingData,
+  DexIndex(Range &&Symbols, Payload &&BackingData, size_t BackingDataSize,
            llvm::ArrayRef<std::string> URISchemes)
       : DexIndex(std::forward<Range>(Symbols), URISchemes) {
     KeepAlive = std::shared_ptr<void>(
         std::make_shared<Payload>(std::move(BackingData)), nullptr);
+    this->BackingDataSize = BackingDataSize;
   }
 
   /// Builds an index from a slab. The index takes ownership of the slab.
   static std::unique_ptr<SymbolIndex>
   build(SymbolSlab Slab, llvm::ArrayRef<std::string> URISchemes) {
-    return llvm::make_unique<DexIndex>(Slab, std::move(Slab), URISchemes);
+    // Store Slab size before it is moved.
+    const auto BackingDataSize = Slab.bytes();
+    return llvm::make_unique<DexIndex>(Slab, std::move(Slab), BackingDataSize,
+                                       URISchemes);
   }
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void lookup(const LookupRequest &Req,
               llvm::function_ref<void(const Symbol &)> Callback) const override;
private:
   /// Inverted index is a mapping from the search token to the posting list,
   /// which contains all items which can be characterized by such search token.
   /// For example, if the search token is scope "std::", the corresponding
   /// posting list would contain all indices of symbols defined in namespace
   /// std. Inverted index is used to retrieve posting lists which are processed
   /// during the fuzzyFind process.
   llvm::DenseMap<Token, PostingList> InvertedIndex;
   std::shared_ptr<void> KeepAlive; // poor man's move-only std::any
+  // Size of memory retained by KeepAlive.
+  size_t BackingDataSize = 0;
 
   std::vector<std::string> URISchemes;
 };
 
 /// Returns Search Token for a number of parent directories of given Path.
 /// Should be used within the index build process.
 ///
 /// This function is exposed for testing only.
 std::vector<std::string> generateProximityURIs(llvm::StringRef URIPath);
 
 } // namespace dex
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_DEX_DEXINDEX_H

clang-tools-extra/clangd/index/dex/DexIndex.cpp

 }
 
 void DexIndex::refs(const RefsRequest &Req,
                     llvm::function_ref<void(const Ref &)> Callback) const {
   log("refs is not implemented.");
 }
 
 size_t DexIndex::estimateMemoryUsage() const {
-  size_t Bytes =
-      LookupTable.size() * sizeof(std::pair<SymbolID, const Symbol *>);
-  Bytes += SymbolQuality.size() * sizeof(std::pair<const Symbol *, float>);
-  Bytes += InvertedIndex.size() * sizeof(Token);
-
-  for (const auto &P : InvertedIndex) {
+  size_t Bytes = Symbols.size() * sizeof(const Symbol *);
+  Bytes += SymbolQuality.size() * sizeof(float);
+  Bytes += LookupTable.getMemorySize();
+  Bytes += InvertedIndex.getMemorySize();
+  for (const auto &P : InvertedIndex)
     Bytes += P.second.size() * sizeof(DocID);
-  }
-  return Bytes;
+  return Bytes + BackingDataSize;
 }
 
 std::vector<std::string> generateProximityURIs(llvm::StringRef URIPath) {
   std::vector<std::string> Result;
   auto ParsedURI = URI::parse(URIPath);
   assert(ParsedURI &&
          "Non-empty argument of generateProximityURIs() should be a valid "
          "URI.");

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

 //===-- IndexTests.cpp  -------------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
 #include "TestIndex.h"
 #include "TestTU.h"
-#include "gmock/gmock.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::_;
 using testing::AllOf;
 using testing::ElementsAre;
 using testing::Pair;
 using testing::Pointee;
 using testing::UnorderedElementsAre;
TEST(SymbolSlab, FindAndIterate) {
   for (const char *Sym : {"X", "Y", "Z"})
     EXPECT_THAT(*S.find(SymbolID(Sym)), Named(Sym));
 }
 
 TEST(SwapIndexTest, OldIndexRecycled) {
   auto Token = std::make_shared<int>();
   std::weak_ptr<int> WeakToken = Token;
 
-  SwapIndex S(
-      llvm::make_unique<MemIndex>(SymbolSlab(), RefSlab(), std::move(Token)));
-  EXPECT_FALSE(WeakToken.expired());     // Current MemIndex keeps it alive.
+  SwapIndex S(llvm::make_unique<MemIndex>(
+      SymbolSlab(), RefSlab(), std::move(Token), /*BackingDataSize=*/0));
+  EXPECT_FALSE(WeakToken.expired());      // Current MemIndex keeps it alive.
   S.reset(llvm::make_unique<MemIndex>()); // Now the MemIndex is destroyed.
-  EXPECT_TRUE(WeakToken.expired());      // So the token is too.
+  EXPECT_TRUE(WeakToken.expired());       // So the token is too.
 }
 
 TEST(MemIndexTest, MemIndexDeduplicate) {
   std::vector<Symbol> Symbols = {symbol("1"), symbol("2"), symbol("3"),
                                  symbol("2") /* duplicate */};
   FuzzyFindRequest Req;
   Req.Query = "2";
   MemIndex I(Symbols, RefSlab());
EXPECT_THAT(
       std::move(Results).build(),
       ElementsAre(Pair(
           _, UnorderedElementsAre(AllOf(RefRange(Test1Code.range("Foo")),
                                         FileURI("unittest:///test.cc")),
                                   AllOf(RefRange(Test2Code.range("Foo")),
                                         FileURI("unittest:///test2.cc"))))));
 }
 
-MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References,  "") {
+MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") {
   return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
 }
 
 TEST(MergeTest, MergeIncludesOnDifferentDefinitions) {
   Symbol L, R;
   L.Name = "left";
   R.Name = "right";
   L.ID = R.ID = SymbolID("hello");