[clangd] Add a metric for tracking memory usage

[kadircet]

• Rebase

[kadircet]

• Rebase and introduce helper for recording a MemoryTree.

[kadircet]

• Move MemoryTree recording from Trace into MemoryTree as suggested in D88417.

[sammccall]

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

68

can we make this a non-member function and take the metric to write into as a parameter?

Putting this here for reuse seems OK to me but it's not an essential part of MemoryTree itself.
In particular I'd like to avoid baking magic metrics into support/.

(I know we have one already for fallback latency recording, but that seems much costlier to avoid as there are many callsites. Here we expect... 2?)

[kadircet]

• Make traversal a free function and take metric to populate as a parameter.

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

public:
   const llvm::DenseMap<llvm::StringRef, MemoryTree> &children() const;
 
   /// Returns total number of bytes used by this sub-tree. Performs a traversal.
   size_t total() const;
 
   /// Returns total number of bytes used by this node only.
   size_t self() const { return Size; }
 
+  /// Records total memory usage of each node under "memory_usage" metric.
+  /// Labels are edges on the path joined with ".", starting with \p RootName.
+  void record(std::string RootName) const;
+
 private:
   /// Adds a child with an edge labeled as \p Name. Multiple calls to this
   /// function returns the same node.
   MemoryTree &createChild(llvm::StringRef Name);
 
   /// Allocator to use for detailed edge names.
   llvm::BumpPtrAllocator *DetailAlloc = nullptr;
 

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

 #include "support/MemoryTree.h"
+#include "Trace.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include <cstddef>
 
 namespace clang {
 namespace clangd {
 
+namespace {
+
+size_t traverseTree(const MemoryTree &MT, std::string &ComponentName) {
+  size_t OriginalLen = ComponentName.size();
+  if (!ComponentName.empty())
+    ComponentName += '.';
+  size_t Total = MT.self();
+  for (const auto &Entry : MT.children()) {
+    ComponentName += Entry.first;
+    Total += traverseTree(Entry.getSecond(), ComponentName);
+    ComponentName.resize(OriginalLen + 1);
+  }
+  ComponentName.resize(OriginalLen);
+  trace::recordMemoryUsage(Total, ComponentName);
+  return Total;
+}
+} // namespace
+
 MemoryTree &MemoryTree::createChild(llvm::StringRef Name) {
   auto &Child = Children.try_emplace(Name, DetailAlloc).first->getSecond();
   return Child;
 }
 
 const llvm::DenseMap<llvm::StringRef, MemoryTree> &
 MemoryTree::children() const {
   return Children;
 }
 
 size_t MemoryTree::total() const {
   size_t Total = Size;
   for (const auto &Entry : Children)
     Total += Entry.getSecond().total();
   return Total;
 }
+
+void MemoryTree::record(std::string RootName) const {
+  traverseTree(*this, RootName);
+}
 } // namespace clangd
 } // namespace clang

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

 #include "llvm/Support/Chrono.h"
 #include "llvm/Support/FormatProviders.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Threading.h"
 #include <atomic>
 #include <chrono>
 #include <memory>
 #include <mutex>
+#include <string>
 
 namespace clang {
 namespace clangd {
 namespace trace {
 
 namespace {
 // The current implementation is naive: each thread writes to Out guarded by Mu.
 // Perhaps we should replace this by something that disturbs performance less.
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

clang-tools-extra/clangd/unittests/support/MemoryTreeTests.cpp

 //===-- MemoryTreeTests.cpp -------------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "support/MemoryTree.h"
+#include "support/TestTracer.h"
 #include "llvm/Support/Allocator.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <ostream>
 
 namespace clang {
 namespace clangd {
 namespace {
 using testing::Contains;
+using testing::ElementsAre;
 using testing::IsEmpty;
 using testing::UnorderedElementsAre;
 
 MATCHER_P2(WithNameAndSize, Name, Size, "") {
   return arg.first == Name &&
          arg.getSecond().total() == static_cast<size_t>(Size);
 }
 
TEST(MemoryTree, DetailedNodesWithDetails) {
 
   {
     auto &Detail = MT.detail("second_detail");
     Detail.child("leaf").addUsage(1);
     EXPECT_THAT(MT.children(), Contains(WithNameAndSize("second_detail", 1)));
     EXPECT_THAT(Detail.children(), Contains(WithNameAndSize("leaf", 1)));
   }
 }
+
+TEST(MemoryTree, Record) {
+  trace::TestTracer Tracer;
+  constexpr llvm::StringLiteral MetricName = "memory_usage";
+  auto AddNodes = [](MemoryTree Root) {
+    Root.child("leaf").addUsage(1);
+
+    {
+      auto &Detail = Root.detail("detail");
+      Detail.addUsage(1);
+      Detail.child("leaf").addUsage(1);
+      auto &Child = Detail.child("child");
+      Child.addUsage(1);
+      Child.child("leaf").addUsage(1);
+    }
+
+    {
+      auto &Child = Root.child("child");
+      Child.addUsage(1);
+      Child.child("leaf").addUsage(1);
+    }
+    return Root;
+  };
+
+  llvm::BumpPtrAllocator Alloc;
+  AddNodes(MemoryTree(&Alloc)).record("root");
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root"), ElementsAre(7));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.leaf"), ElementsAre(1));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.detail"), ElementsAre(4));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.detail.leaf"),
+              ElementsAre(1));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.detail.child"),
+              ElementsAre(2));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.detail.child.leaf"),
+              ElementsAre(1));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.child"), ElementsAre(2));
+  EXPECT_THAT(Tracer.takeMetric(MetricName, "root.child.leaf"), ElementsAre(1));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang