[clangd] Add a way for exporting memory usage metrics

[kadircet]

Summary

Introduces a helper to collect memory usage metrics from multiple
components.
Adds measurements for dynamic index, ast cache and preambles.

[adamcz]

clang-tools-extra/clangd/TUScheduler.cpp

172

Any idea how expensive this is? I suppose TUScheduler::update() is rare enough that it's not a big deal?

I ask because of bad experience with estimating memory usage inside critical path on another project ;-)

Maybe we can add a span around this, so we can see if it's expensive in traces? Or maybe that will turn out to be more expensive than memory estimation?

1300

Isn't FD->Worker->update() above async (with correct options)? If so, this is still reading the old data, right?

1312

Would it make sense to collect the number of preambles and ASTs cached at this point as well? Exported under a different metric, but at the same time, to be able to join the data together?

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

172

TUScheduler::update() is rare enough

actually it can potentially be triggered after every keystroke. I suppose it might make sense to move this out of the main thread completely. I suppose runWithPreambles action might be a better place to record, which is still quite often btw (every signature help and code completion request goes through it), but runs on a separate thread.

i would expect that will turn out to be more expensive than memory estimation to be the case, will do some checks though.

Most of the calculations happening in there is just addition of some members. the source manager makes me afraid a little bit tho, as it seems to be going over all the buffers.

1300

yes that's true. i was aiming for this to be an estimate, rather than an accurate view.

1312

yeah that definitely makes sense, but I would rather do that on a separate patch.

[adamcz]

clang-tools-extra/clangd/TUScheduler.cpp

168

nit: This is the same name as getUsedBytes(Key K). Maybe rename to getTotalUsedBytes()?

172

+1 to moving call to this to runWithPreamble or something like that. If a request gets cancelled, for example, there's no need to update the memory usage.

1300

I think this deserves a comment.

1310

So technically this is incorrect.
IdleASTs might contain AST for a file that is no longer tracked (e.g. had removeDocument() called on it). ASTCacheBytes will include it, but PreambleBytes will not, since it only asks for size of the tracked files. It's probably fine, considering how rare this will be and that these are all estimates, but I would recommend a comment, in case someone sees PreambleBytes be negative and wonders what's going on.

1312

Sounds good.

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

1310

right, thanks for pointing that out. added a comment, and clamped for non-negativity.

[kadircet]

• Rename the overload
• Add comments around possible caveats that might result in inaccuracies.
• Move the metric recording itself into another thread.
• Keep the calculations in the main thread, as they seemed to have <1ms latency, even with huge preambles/asts.

[sammccall]

Just for the record, as I think Kadir and Adam are both aware...

We discussed making this a bit richer and less reliant on static state.
We'd build a tree-shaped profile by passing a tree-builder recursively into various components.
Then the metrics would be exported at the top level, but we'd also want to expose it via debugging actions. (The tree edges from e.g. dynamic index to individual files would probably be optional - collapsed for metrics but present for debugging)

LMK if I have this wrong and you want to move forward with a simpler approch.

[kadircet]

In D86077#2281218, @sammccall wrote:

Just for the record, as I think Kadir and Adam are both aware...

We discussed making this a bit richer and less reliant on static state.
We'd build a tree-shaped profile by passing a tree-builder recursively into various components.
Then the metrics would be exported at the top level, but we'd also want to expose it via debugging actions. (The tree edges from e.g. dynamic index to individual files would probably be optional - collapsed for metrics but present for debugging)

LMK if I have this wrong and you want to move forward with a simpler approch.

No that aligns perfectly with what I have in my mind as well. I just didn't have time to get back to this yet. Thanks a lot for summarizing offline discussions here!

clang-tools-extra/clangd/TUScheduler.h

public:
   /// Wait until there are no scheduled or running tasks.
   /// Mostly useful for synchronizing tests.
   bool blockUntilIdle(Deadline D) const;
 
 private:
   /// This class stores per-file data in the Files map.
   struct FileData;
 
+  /// Records estimated memory usage for the TUScheduler internals. These are
+  /// best-effort estimates as TUScheduler itself is multi-threaded and internal
+  /// state is likely to change during the calculations.
+  void recordMemoryUsage();
+
 public:
   /// Responsible for retaining and rebuilding idle ASTs. An implementation is
   /// an LRU cache.
   class ASTCache;
 
   // 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.

clang-tools-extra/clangd/TUScheduler.cpp

 #include "llvm/Support/Errc.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Threading.h"
 #include <algorithm>
 #include <chrono>
 #include <condition_variable>
+#include <cstddef>
 #include <functional>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <string>
 #include <thread>
 #include <type_traits>
 #include <utility>
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));
   }
 
+  size_t getTotalUsedBytes() {
+    size_t TotalBytes = 0;
+    std::lock_guard<std::mutex> Lock(Mut);
+    for (const auto &Elem : LRU)
+      TotalBytes += Elem.second->getUsedBytes();
+    return TotalBytes;
+  }
+
 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;
bool TUScheduler::blockUntilIdle(Deadline D) const {
   if (PreambleTasks)
     if (!PreambleTasks->wait(D))
       return false;
   return true;
 }
 
 bool TUScheduler::update(PathRef File, ParseInputs Inputs,
                          WantDiagnostics WantDiags) {
+  recordMemoryUsage();
   std::unique_ptr<FileData> &FD = Files[File];
   bool NewFile = FD == nullptr;
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker =
         ASTWorker::create(File, CDB, *IdleASTs,
                           WorkerThreads ? WorkerThreads.getPointer() : nullptr,
                           Barrier, Opts, *Callbacks);
     FD = std::unique_ptr<FileData>(
         new FileData{Inputs.Contents, std::move(Worker)});
   } else {
     FD->Contents = Inputs.Contents;
   }
   FD->Worker->update(std::move(Inputs), WantDiags);
   return NewFile;
 }
 
 void TUScheduler::remove(PathRef File) {
   bool Removed = Files.erase(File);
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
 }
 
 llvm::StringMap<std::string> TUScheduler::getAllFileContents() const {
   llvm::StringMap<std::string> Results;
   for (auto &It : Files)
     Results.try_emplace(It.getKey(), It.getValue()->Contents);
   return Results;
 }
 
 void TUScheduler::run(llvm::StringRef Name, llvm::StringRef Path,
                       llvm::unique_function<void()> Action) {
   if (!PreambleTasks) {
     WithContext WithProvidedContext(Opts.ContextProvider(Path));
     return Action();

 }
 
 DebouncePolicy DebouncePolicy::fixed(clock::duration T) {
   DebouncePolicy P;
   P.Min = P.Max = T;
   return P;
 }
 
+void TUScheduler::recordMemoryUsage() {
+  // We defer recording of the measurements to a worker thread, as it might
+  // perform IO and we don't want to block main thread. We perform the
+  // calculation itself here, as it accesses TUScheduler state, which is not
+  // threadsafe. It is cheap enough though.
+  size_t ASTCacheBytes = IdleASTs->getTotalUsedBytes();
+  size_t PreambleBytes = 0;
+  // Otherwise preambles are stored on disk and we only keep filename in memory.
+  if (Opts.StorePreamblesInMemory) {
+    for (const auto &Elem : fileStats())
+      PreambleBytes += Elem.second.UsedBytes;
+    // fileStats results include ast cache sizes too, subtract them. Note that
+    // the size of the ast cache seen by fileStats() and here might differ in
+    // some circumstances. E.g. user closes a file that still has some pending
+    // action (like preamble build), the cached ast size will be included in
+    // ASTCacheBytes, but not in fileStats() as file is no longer active. This
+    // might also overshoot, e.g. worker might cache an AST between the
+    // calculation of ASTCacheBytes and fileStats().
+    // We clamp to guarantee non-negative values.
+    PreambleBytes -= std::min(PreambleBytes, ASTCacheBytes);
+  }
+
+  run("RecordMemoryUsage", /*Path=*/"", [ASTCacheBytes, PreambleBytes] {
+    trace::recordMemoryUsage(ASTCacheBytes, "ast_cache");
+    trace::recordMemoryUsage(PreambleBytes, "preambles");
+  });
+}
 } // namespace clangd
 } // namespace clang

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

 #include "index/Relation.h"
 #include "index/Serialization.h"
 #include "index/Symbol.h"
 #include "index/SymbolID.h"
 #include "index/SymbolOrigin.h"
 #include "index/dex/Dex.h"
 #include "support/Logger.h"
 #include "support/Path.h"
+#include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
PreambleSymbols.update(
         std::make_unique<RefSlab>(),
         std::make_unique<RelationSlab>(std::move(*IF->Relations)),
         /*CountReferences=*/false);
   }
   size_t IndexVersion = 0;
   auto NewIndex =
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
                                  DuplicateHandling::PickOne, &IndexVersion);
+  trace::recordMemoryUsage(NewIndex->estimateMemoryUsage(), "preamble_index");
   {
     std::lock_guard<std::mutex> Lock(UpdateIndexMu);
     if (IndexVersion <= PreambleIndexVersion) {
       // We lost the race, some other thread built a later version.
       return;
     }
     PreambleIndexVersion = IndexVersion;
     PreambleIndex.reset(std::move(NewIndex));
void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
   MainFileSymbols.update(
       Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Contents))),
       std::make_unique<RefSlab>(std::move(std::get<1>(Contents))),
       std::make_unique<RelationSlab>(std::move(std::get<2>(Contents))),
       /*CountReferences=*/true);
   size_t IndexVersion = 0;
   auto NewIndex = MainFileSymbols.buildIndex(
       IndexType::Light, DuplicateHandling::Merge, &IndexVersion);
+  trace::recordMemoryUsage(NewIndex->estimateMemoryUsage(), "main_file_index");
   {
     std::lock_guard<std::mutex> Lock(UpdateIndexMu);
     if (IndexVersion <= MainIndexVersion) {
       // We lost the race, some other thread built a later version.
       return;
     }
     MainIndexVersion = IndexVersion;
     MainFileIndex.reset(std::move(NewIndex));

clang-tools-extra/clangd/support/Trace.h

struct Metric {
   const llvm::StringLiteral Name;
   const MetricType Type;
   /// Indicates what measurement labels represent, e.g. "operation_name" for a
   /// metric tracking latencies. If non empty all measurements must also have a
   /// non-empty label.
   const llvm::StringLiteral LabelName;
 };
 
+/// Convenient helper for collecting memory usage metrics from across multiple
+/// components. Results are recorded under a metric named "memory_usage".
+void recordMemoryUsage(double Value, llvm::StringRef ComponentName);
+
 /// A consumer of trace events and measurements. The events are produced by
 /// Spans and trace::log, the measurements are produced by Metrics::record.
 /// Implementations of this interface must be thread-safe.
 class EventTracer {
 public:
   virtual ~EventTracer() = default;
 
   /// Called when event that has a duration starts. \p Name describes the event.

clang-tools-extra/clangd/support/Trace.cpp

void Metric::record(double Value, llvm::StringRef Label) const {
   assert((LabelName.empty() == Label.empty()) &&
          "recording a measurement with inconsistent labeling");
   T->record(*this, Value, Label);
 }
 
 Context EventTracer::beginSpan(llvm::StringRef Name, llvm::json::Object *Args) {
   return Context::current().clone();
 }
+
+void recordMemoryUsage(double Value, llvm::StringRef ComponentName) {
+  static constexpr Metric MemoryUsage("memory_usage", Metric::Value,
+                                      "component_name");
+
+  MemoryUsage.record(Value, ComponentName);
+}
 } // namespace trace
 } // namespace clangd
 } // namespace clang