[clangd] Introduce memory usage dumping to TUScheduler, for Preambles and ASTCache

[kadircet]

Summary

File-granular information is considered details.

Depends on D88411

[kadircet]

• Rebase

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

1447

should we group by file instead? The fact that the AST-cache is owned by TUScheduler and not ASTWorker seems like a confusing detail that it might be better to omit

1452

why aren't we also using the UsedBytesAST estimate here, and instead walking through the cache directly?

clang-tools-extra/clangd/TUScheduler.h

316

(raised naming of these functions in another patch, however that ends up let's ensure they match)

[kadircet]

• Rename attachMemoryUsage to profile
• Group by filename, rather than ast_cache vs preamble
• Make use of UsedBytesAST rather than profiling IdleASTs
• Rebase

clang-tools-extra/clangd/TUScheduler.h

 //===--- TUScheduler.h -------------------------------------------*-C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Function.h"
+#include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include <chrono>
 
struct Options {
     std::function<Context(PathRef)> ContextProvider;
   };
 
   TUScheduler(const GlobalCompilationDatabase &CDB, const Options &Opts,
               std::unique_ptr<ParsingCallbacks> ASTCallbacks = nullptr);
   ~TUScheduler();
 
   struct FileStats {
-    std::size_t UsedBytes = 0;
+    std::size_t UsedBytesAST = 0;
+    std::size_t UsedBytesPreamble = 0;
     unsigned PreambleBuilds = 0;
     unsigned ASTBuilds = 0;
   };
   /// Returns resources used for each of the currently open files.
   /// Results are inherently racy as they measure activity of other threads.
   llvm::StringMap<FileStats> fileStats() const;
 
   /// Returns a list of files with ASTs currently stored in memory. This method
public:
   // The file being built/processed in the current thread. This is a hack in
   // order to get the file name into the index implementations. Do not depend on
   // this inside clangd.
   // FIXME: remove this when there is proper index support via build system
   // integration.
   // FIXME: move to ClangdServer via createProcessingContext.
   static llvm::Optional<llvm::StringRef> getFileBeingProcessedInContext();
 
+  MemoryTree getMemoryUsage() const;
+
 private:
   const GlobalCompilationDatabase &CDB;
   Options Opts;
   std::unique_ptr<ParsingCallbacks> Callbacks; // not nullptr
   Semaphore Barrier;
   llvm::StringMap<std::unique_ptr<FileData>> Files;
   std::unique_ptr<ASTCache> IdleASTs;
   // None when running tasks synchronously and non-None when running tasks

clang-tools-extra/clangd/TUScheduler.cpp

 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
 #include "ParsedAST.h"
 #include "Preamble.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Logger.h"
+#include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "support/Trace.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
take(Key K, const trace::Metric *AccessMetric = nullptr) {
     std::unique_ptr<ParsedAST> V = std::move(Existing->second);
     LRU.erase(Existing);
     // GCC 4.8 fails to compile `return V;`, as it tries to call the copy
     // constructor of unique_ptr, so we call the move ctor explicitly to avoid
     // this miscompile.
     return llvm::Optional<std::unique_ptr<ParsedAST>>(std::move(V));
   }
 
+  MemoryTree getMemoryUsage() const {
+    MemoryTree MT;
+    std::lock_guard<std::mutex> Lock(Mut);
+    for (const auto &Elem : LRU) {
+      auto &SM = Elem.second->getSourceManager();
+      MT.addChild(
+          // FIXME: Store the filename directly in the cache or ParsedAST.
+          SM.getFileEntryForID(SM.getMainFileID())->tryGetRealPathName(),
+          Elem.second->getUsedBytes());
+    }
+    return MT;
+  }
+
 private:
   using KVPair = std::pair<Key, std::unique_ptr<ParsedAST>>;
 
   std::vector<KVPair>::iterator findByKey(Key K) {
     return llvm::find_if(LRU, [K](const KVPair &P) { return P.first == K; });
   }
 
-  std::mutex Mut;
+  mutable std::mutex Mut;
   unsigned MaxRetainedASTs;
   /// Items sorted in LRU order, i.e. first item is the most recently accessed
   /// one.
   std::vector<KVPair> LRU; /* GUARDED_BY(Mut) */
 };
 
 namespace {
 /// Threadsafe manager for updating a TUStatus and emitting it after each

 
 TUScheduler::FileStats ASTWorker::stats() const {
   TUScheduler::FileStats Result;
   Result.ASTBuilds = ASTBuildCount;
   Result.PreambleBuilds = PreambleBuildCount;
   // Note that we don't report the size of ASTs currently used for processing
   // the in-flight requests. We used this information for debugging purposes
   // only, so this should be fine.
-  Result.UsedBytes = IdleASTs.getUsedBytes(this);
+  Result.UsedBytesAST = IdleASTs.getUsedBytes(this);
   if (auto Preamble = getPossiblyStalePreamble())
-    Result.UsedBytes += Preamble->Preamble.getSize();
+    Result.UsedBytesPreamble = Preamble->Preamble.getSize();
   return Result;
 }
 
 bool ASTWorker::isASTCached() const { return IdleASTs.getUsedBytes(this) != 0; }
 
 void ASTWorker::stop() {
   {
     std::lock_guard<std::mutex> Lock(PublishMu);

 }
 
 DebouncePolicy DebouncePolicy::fixed(clock::duration T) {
   DebouncePolicy P;
   P.Min = P.Max = T;
   return P;
 }
 
+MemoryTree TUScheduler::getMemoryUsage() const {
+  MemoryTree MT;
+  MT.addChild("ast_cahe", IdleASTs->getMemoryUsage(), true);
+  // Otherwise preambles are stored on disk and we only keep filename in memory.
+  if (Opts.StorePreamblesInMemory) {
+    MemoryTree Preambles;
+    for (const auto &Elem : fileStats())
+      Preambles.addChild(Elem.first(), Elem.second.UsedBytesPreamble);
+    MT.addChild("preambles", std::move(Preambles), true);
+  }
+  return MT;
+}
 } // namespace clangd
 } // namespace clang

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

EXPECT_THAT(DiagConsumer.filesWithDiags(),
               UnorderedElementsAre(Pair(FooCpp, false), Pair(BarCpp, false)));
 
   Locations = runLocateSymbolAt(Server, FooCpp, FooSource.point());
   EXPECT_TRUE(bool(Locations));
   EXPECT_THAT(*Locations, ElementsAre(DeclAt(FooCpp, FooSource.range("two"))));
 }
 
 MATCHER_P4(Stats, Name, UsesMemory, PreambleBuilds, ASTBuilds, "") {
-  return arg.first() == Name && (arg.second.UsedBytes != 0) == UsesMemory &&
+  return arg.first() == Name &&
+         (arg.second.UsedBytesAST + arg.second.UsedBytesPreamble != 0) ==
+             UsesMemory &&
          std::tie(arg.second.PreambleBuilds, ASTBuilds) ==
              std::tie(PreambleBuilds, ASTBuilds);
 }
 
 TEST(ClangdServerTest, FileStats) {
   MockFS FS;
   ErrorCheckingCallbacks DiagConsumer;
   MockCompilationDatabase CDB;