[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?

1299

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

1311

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.

1299

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

1311

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.

1299

I think this deserves a comment.

1309

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.

1311

Sounds good.

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

1309

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.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 getUsedBytes() {
+    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;
ASTWorkerHandle Worker =
                           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);
+
+  // Update memory usage metrics. Note that these are just estimates.
+  size_t ASTCacheBytes = IdleASTs->getUsedBytes();
+  trace::recordMemoryUsage(ASTCacheBytes, "ast_cache");
+
+  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.
+    PreambleBytes -= ASTCacheBytes;
+  }
+  trace::recordMemoryUsage(PreambleBytes, "preambles");
   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);

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