[clangd] Log memory usage of DexIndex and MemIndex

This patch prints information about built index size estimation to
verbose logs. This is useful for optimizing memory usage of DexIndex and
comparisons with MemIndex.

Reviewed by: sammccall

Differential Revision: https://reviews.llvm.org/D51154

llvm-svn: 340601

9 files changed, 52 insertions, 0 deletions

diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 3d8e17bfb2f..48d05959dae 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -118,5 +118,9 @@ void FileIndex::findOccurrences(
   log("findOccurrences is not implemented.");
 }
 
+size_t FileIndex::estimateMemoryUsage() const {
+  return Index.estimateMemoryUsage();
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index 6f24f9d2c23..58fad2f62e2 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -81,6 +81,9 @@ public:
   void findOccurrences(const OccurrencesRequest &Req,
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
+
+  size_t estimateMemoryUsage() const override;
+
 private:
   FileSymbols FSymbols;
   MemIndex Index;
diff --git a/clang-tools-extra/clangd/index/Index.h b/clang-tools-extra/clangd/index/Index.h
index eac4c5e7d85..95b081acec7 100644
--- a/clang-tools-extra/clangd/index/Index.h
+++ b/clang-tools-extra/clangd/index/Index.h
@@ -385,6 +385,12 @@ public:
   virtual void findOccurrences(
       const OccurrencesRequest &Req,
       llvm::function_ref<void(const SymbolOccurrence &)> Callback) const = 0;
+
+  /// Returns estimated size of index (in bytes).
+  // FIXME(kbobyrev): Currently, this only returns the size of index itself
+  // excluding the size of actual symbol slab index refers to. We should include
+  // both.
+  virtual size_t estimateMemoryUsage() const = 0;
 };
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/index/MemIndex.cpp b/clang-tools-extra/clangd/index/MemIndex.cpp
index db11244d8ac..19a64ad9525 100644
--- a/clang-tools-extra/clangd/index/MemIndex.cpp
+++ b/clang-tools-extra/clangd/index/MemIndex.cpp
@@ -26,6 +26,9 @@ void MemIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
     Index = std::move(TempIndex);
     Symbols = std::move(Syms); // Relase old symbols.
   }
+
+  vlog("Built MemIndex with estimated memory usage {0} bytes.",
+       estimateMemoryUsage());
 }
 
 std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Slab) {
@@ -98,5 +101,10 @@ getSymbolsFromSlab(SymbolSlab Slab) {
                                                       &Snap->Pointers);
 }
 
+size_t MemIndex::estimateMemoryUsage() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return Index.getMemorySize();
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/MemIndex.h b/clang-tools-extra/clangd/index/MemIndex.h
index 8b12f8c7921..feca87af5e4 100644
--- a/clang-tools-extra/clangd/index/MemIndex.h
+++ b/clang-tools-extra/clangd/index/MemIndex.h
@@ -39,7 +39,10 @@ public:
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 
+  size_t estimateMemoryUsage() const override;
+
 private:
+
   std::shared_ptr<std::vector<const Symbol *>> Symbols;
   // Index is a set of symbols that are deduplicated by symbol IDs.
   // FIXME: build smarter index structure.
diff --git a/clang-tools-extra/clangd/index/Merge.cpp b/clang-tools-extra/clangd/index/Merge.cpp
index 4834bfba69f..84928ca1a5e 100644
--- a/clang-tools-extra/clangd/index/Merge.cpp
+++ b/clang-tools-extra/clangd/index/Merge.cpp
@@ -84,6 +84,10 @@ class MergedIndex : public SymbolIndex {
     log("findOccurrences is not implemented.");
   }
 
+  size_t estimateMemoryUsage() const override {
+    return Dynamic->estimateMemoryUsage() + Static->estimateMemoryUsage();
+  }
+
 private:
   const SymbolIndex *Dynamic, *Static;
 };
diff --git a/clang-tools-extra/clangd/index/dex/DexIndex.cpp b/clang-tools-extra/clangd/index/dex/DexIndex.cpp
index 534f51e0c3d..9280cc8fdfc 100644
--- a/clang-tools-extra/clangd/index/dex/DexIndex.cpp
+++ b/clang-tools-extra/clangd/index/dex/DexIndex.cpp
@@ -67,6 +67,9 @@ void DexIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
     InvertedIndex = std::move(TempInvertedIndex);
     SymbolQuality = std::move(TempSymbolQuality);
   }
+
+  vlog("Built DexIndex with estimated memory usage {0} bytes.",
+       estimateMemoryUsage());
 }
 
 std::unique_ptr<SymbolIndex> DexIndex::build(SymbolSlab Slab) {
@@ -171,6 +174,20 @@ void DexIndex::findOccurrences(
   log("findOccurrences is not implemented.");
 }
 
+size_t DexIndex::estimateMemoryUsage() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  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) {
+    Bytes += P.second.size() * sizeof(DocID);
+  }
+  return Bytes;
+}
+
 } // namespace dex
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/dex/DexIndex.h b/clang-tools-extra/clangd/index/dex/DexIndex.h
index 4485150b132..d9b2cd14389 100644
--- a/clang-tools-extra/clangd/index/dex/DexIndex.h
+++ b/clang-tools-extra/clangd/index/dex/DexIndex.h
@@ -57,7 +57,10 @@ public:
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 
+  size_t estimateMemoryUsage() const override;
+
 private:
+
   mutable std::mutex Mutex;
 
   std::shared_ptr<std::vector<const Symbol *>> Symbols /*GUARDED_BY(Mutex)*/;
diff --git a/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp b/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
index 8f7f85377b6..3f855844aed 100644
--- a/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
+++ b/clang-tools-extra/unittests/clangd/CodeCompleteTests.cpp
@@ -923,6 +923,10 @@ public:
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override {}
 
+  // This is incorrect, but IndexRequestCollector is not an actual index and it
+  // isn't used in production code.
+  size_t estimateMemoryUsage() const override { return 0; }
+
   const std::vector<FuzzyFindRequest> allRequests() const { return Requests; }
 
 private: