[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

26 ↗

(On Diff #163519)

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

184 ↗

(On Diff #163523)

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

91 ↗

(On Diff #163523)

Why not bytes here as well?

[kbobyrev]

Rebase on top of recent revisions.

[ioeric]

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

144 ↗

(On Diff #163786)

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

44 ↗

(On Diff #163786)

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

(same for MemIndex)

53 ↗

(On Diff #163786)

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)

100 ↗

(On Diff #163786)

nit: =0?

[sammccall]

LG with a few nits

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

121 ↗

(On Diff #163786)

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

144 ↗

(On Diff #163786)

(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

32 ↗

(On Diff #163786)

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

70 ↗

(On Diff #163786)

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

144 ↗

(On Diff #163786)

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

[kbobyrev]

Address a round of comments.

[sammccall]

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

135 ↗

(On Diff #164369)

this size should be calculated from the slabs above

[sammccall]

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

131 ↗

(On Diff #164377)

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 ↗

(On Diff #164643)

nit: I'd avoid irrelevant changes.

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

20 ↗

(On Diff #164643)

Include Refs slab size?

76 ↗

(On Diff #164643)

Include Refs size?

[kbobyrev]

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

512 ↗

(On Diff #164643)

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

clang-tools-extra/trunk/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 = RefsStorage.size() * sizeof(Ref);
+  for (const auto &Slab : SymbolSlabs)
+    StorageSize += Slab->bytes();
+  for (const auto &RefSlab : RefSlabs)
+    StorageSize += RefSlab->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 {

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

public:
   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/trunk/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() + Refs.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() + Refs.getMemorySize() + BackingDataSize;
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/clangd/index/dex/Dex.h

if (URISchemes.empty()) {
       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>
-  Dex(Range &&Symbols, Payload &&BackingData,
+  Dex(Range &&Symbols, Payload &&BackingData, size_t BackingDataSize,
       llvm::ArrayRef<std::string> URISchemes)
       : Dex(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<Dex>(Slab, std::move(Slab), URISchemes);
+    // Store Slab size before it is moved.
+    const auto BackingDataSize = Slab.bytes();
+    return llvm::make_unique<Dex>(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_DEX_H

clang-tools-extra/trunk/clangd/index/dex/Dex.cpp

 }
 
 void Dex::refs(const RefsRequest &Req,
                llvm::function_ref<void(const Ref &)> Callback) const {
   log("refs is not implemented.");
 }
 
 size_t Dex::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/trunk/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());