[clangd] Record memory usages after each notification

Depends on D88415

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

clang-tools-extra/clangd/ClangdLSPServer.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
+#include "ClangdServer.h"
 #include "CodeComplete.h"
 #include "Diagnostics.h"
 #include "DraftStore.h"
@@ -18,6 +19,7 @@
 #include "URI.h"
 #include "refactor/Tweak.h"
 #include "support/Context.h"
+#include "support/MemoryTree.h"
 #include "support/Trace.h"
 #include "clang/Basic/Version.h"
 #include "clang/Tooling/Core/Replacement.h"
@@ -26,6 +28,7 @@
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -33,6 +36,7 @@
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SHA1.h"
 #include "llvm/Support/ScopedPrinter.h"
+#include <chrono>
 #include <cstddef>
 #include <memory>
 #include <mutex>
@@ -144,7 +148,6 @@ llvm::Error validateEdits(const DraftStore &DraftMgr, const FileEdits &FE) {
   return error("Files must be saved first: {0} (and {1} others)",
                LastInvalidFile, InvalidFileCount - 1);
 }
-
 } // namespace
 
 // MessageHandler dispatches incoming LSP messages.
@@ -163,14 +166,16 @@ public:
     log("<-- {0}", Method);
     if (Method == "exit")
       return false;
-    if (!Server.Server)
+    if (!Server.Server) {
       elog("Notification {0} before initialization", Method);
-    else if (Method == "$/cancelRequest")
+    } else if (Method == "$/cancelRequest") {
       onCancel(std::move(Params));
-    else if (auto Handler = Notifications.lookup(Method))
+    } else if (auto Handler = Notifications.lookup(Method)) {
       Handler(std::move(Params));
-    else
+      Server.maybeExportMemoryProfile();
+    } else {
       log("unhandled notification {0}", Method);
+    }
     return true;
   }
 
@@ -1234,6 +1239,25 @@ void ClangdLSPServer::publishDiagnostics(
   notify("textDocument/publishDiagnostics", Params);
 }
 
+void ClangdLSPServer::maybeExportMemoryProfile() {
+  if (!trace::enabled())
+    return;
+  // Profiling might be expensive, so we throttle it to happen once every 5
+  // minutes.
+  static constexpr auto ProfileInterval = std::chrono::minutes(5);
+  auto Now = std::chrono::steady_clock::now();
+  if (Now < NextProfileTime)
+    return;
+
+  static constexpr trace::Metric MemoryUsage(
+      "memory_usage", trace::Metric::Value, "component_name");
+  trace::Span Tracer("ProfileBrief");
+  MemoryTree MT;
+  profile(MT);
+  record(MT, "clangd_lsp_server", MemoryUsage);
+  NextProfileTime = Now + ProfileInterval;
+}
+
 // FIXME: This function needs to be properly tested.
 void ClangdLSPServer::onChangeConfiguration(
     const DidChangeConfigurationParams &Params) {
@@ -1404,6 +1428,9 @@ ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
   if (Opts.FoldingRanges)
     MsgHandler->bind("textDocument/foldingRange", &ClangdLSPServer::onFoldingRange);
   // clang-format on
+
+  // Delay first profile until we've finished warming up.
+  NextProfileTime = std::chrono::steady_clock::now() + std::chrono::minutes(1);
 }
 
 ClangdLSPServer::~ClangdLSPServer() {
@@ -1424,6 +1451,11 @@ bool ClangdLSPServer::run() {
   return CleanExit && ShutdownRequestReceived;
 }
 
+void ClangdLSPServer::profile(MemoryTree &MT) const {
+  if (Server)
+    Server->profile(MT.child("clangd_server"));
+}
+
 std::vector<Fix> ClangdLSPServer::getFixes(llvm::StringRef File,
                                            const clangd::Diagnostic &D) {
   std::lock_guard<std::mutex> Lock(FixItsMutex);

clang-tools-extra/clangd/ClangdLSPServer.h

@@ -17,11 +17,13 @@
 #include "Protocol.h"
 #include "Transport.h"
 #include "support/Context.h"
+#include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/JSON.h"
+#include <chrono>
 #include <memory>
 
 namespace clang {
@@ -67,6 +69,9 @@ public:
   /// \return Whether we shut down cleanly with a 'shutdown' -> 'exit' sequence.
   bool run();
 
+  /// Profiles resource-usage.
+  void profile(MemoryTree &MT) const;
+
 private:
   // Implement ClangdServer::Callbacks.
   void onDiagnosticsReady(PathRef File, llvm::StringRef Version,
@@ -160,6 +165,14 @@ private:
   /// Sends a "publishDiagnostics" notification to the LSP client.
   void publishDiagnostics(const PublishDiagnosticsParams &);
 
+  /// Runs profiling and exports memory usage metrics if tracing is enabled and
+  /// profiling hasn't happened recently.
+  void maybeExportMemoryProfile();
+
+  /// Timepoint until which profiling is off. It is used to throttle profiling
+  /// requests.
+  std::chrono::steady_clock::time_point NextProfileTime;
+
   /// Since initialization of CDBs and ClangdServer is done lazily, the
   /// following context captures the one used while creating ClangdLSPServer and
   /// passes it to above mentioned object instances to make sure they share the

clang-tools-extra/clangd/ClangdServer.cpp

@@ -28,6 +28,7 @@
 #include "refactor/Tweak.h"
 #include "support/Logger.h"
 #include "support/Markup.h"
+#include "support/MemoryTree.h"
 #include "support/ThreadsafeFS.h"
 #include "support/Trace.h"
 #include "clang/Format/Format.h"
@@ -826,5 +827,12 @@ ClangdServer::blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds) {
           BackgroundIdx->blockUntilIdleForTest(TimeoutSeconds));
 }
 
+void ClangdServer::profile(MemoryTree &MT) const {
+  if (DynamicIdx)
+    DynamicIdx->profile(MT.child("dynamic_index"));
+  if (BackgroundIdx)
+    BackgroundIdx->profile(MT.child("background_index"));
+  WorkScheduler.profile(MT.child("tuscheduler"));
+}
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/ClangdServer.h

@@ -25,6 +25,7 @@
 #include "refactor/Tweak.h"
 #include "support/Cancellation.h"
 #include "support/Function.h"
+#include "support/MemoryTree.h"
 #include "support/ThreadsafeFS.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
@@ -337,6 +338,9 @@ public:
   LLVM_NODISCARD bool
   blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds = 10);
 
+  /// Builds a nested representation of memory used by components.
+  void profile(MemoryTree &MT) const;
+
 private:
   void formatCode(PathRef File, llvm::StringRef Code,
                   ArrayRef<tooling::Range> Ranges,

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

@@ -17,6 +17,7 @@
 #include "TestFS.h"
 #include "TestTU.h"
 #include "URI.h"
+#include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "clang/Config/config.h"
@@ -27,6 +28,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
@@ -48,6 +50,7 @@ namespace clangd {
 namespace {
 
 using ::testing::AllOf;
+using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::Field;
 using ::testing::Gt;
@@ -1236,6 +1239,21 @@ TEST(ClangdServer, TidyOverrideTest) {
   EXPECT_FALSE(DiagConsumer.HadDiagsInLastCallback);
 }
 
+TEST(ClangdServer, MemoryUsageTest) {
+  MockFS FS;
+  MockCompilationDatabase CDB;
+  ClangdServer Server(CDB, FS, ClangdServer::optsForTest());
+
+  auto FooCpp = testPath("foo.cpp");
+  Server.addDocument(FooCpp, "");
+  ASSERT_TRUE(Server.blockUntilIdleForTest());
+
+  llvm::BumpPtrAllocator Alloc;
+  MemoryTree MT(&Alloc);
+  Server.profile(MT);
+  ASSERT_TRUE(MT.children().count("tuscheduler"));
+  EXPECT_TRUE(MT.child("tuscheduler").children().count(FooCpp));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang