admin: add allocation profiler (#36136)

add `/allocprofiler` admin handler to record the tcmalloc allocation profile. I used this to chase down some excessive heap memory use (which turned out to be tcmalloc caches). Risk Level: low Testing: done Docs Changes: yes Release Notes: none Change-Id: I79629537ab83c54b655de7ef1010b29665d30541 Signed-off-by: Kuat Yessenov <[email protected]>
envoyproxy · Sep 23, 2024 · 6307e4a · 6307e4a
1 parent 6bccf47
commit 6307e4a
Show file tree

Hide file tree

Showing 11 changed files with 110 additions and 0 deletions.
diff --git a/docs/root/operations/admin.rst b/docs/root/operations/admin.rst
@@ -250,6 +250,11 @@ modify different aspects of the server:
   Dump current heap profile of Envoy process. The output content is parsable binary by the ``pprof`` tool.
   Requires compiling with tcmalloc (default).
 
+.. http:post:: /allocprofiler
+
+  Enable or disable the allocation profiler. The output content is parsable binary by the ``pprof`` tool.
+  Requires compiling with tcmalloc (default).
+
 .. _operations_admin_interface_healthcheck_fail:
 
 .. http:post:: /healthcheck/fail

diff --git a/source/common/common/BUILD b/source/common/common/BUILD
@@ -413,6 +413,7 @@ envoy_cc_library(
     srcs = ["thread.cc"],
     hdrs = ["thread.h"],
     deps = envoy_cc_platform_dep("thread_impl_lib") + [
+        ":assert_lib",
         ":macros",
         ":non_copyable",
         "@com_google_absl//absl/synchronization",

diff --git a/source/common/common/thread.h b/source/common/common/thread.h
@@ -7,6 +7,7 @@
 
 #include "envoy/thread/thread.h"
 
+#include "source/common/common/assert.h"
 #include "source/common/common/non_copyable.h"
 
 #include "absl/synchronization/mutex.h"

diff --git a/source/common/profiler/BUILD b/source/common/profiler/BUILD
@@ -14,6 +14,7 @@ envoy_cc_library(
     hdrs = ["profiler.h"],
     tcmalloc_dep = 1,
     deps = [
+        "//source/common/common:thread_lib",
         "@com_google_absl//absl/status:statusor",
     ],
 )
diff --git a/source/common/profiler/profiler.cc b/source/common/profiler/profiler.cc
@@ -2,6 +2,8 @@
 
 #include <string>
 
+#include "source/common/common/thread.h"
+
 #ifdef PROFILER_AVAILABLE
 
 #include "gperftools/heap-profiler.h"
@@ -68,11 +70,37 @@ bool Heap::stopProfiler() { return false; }
 namespace Envoy {
 namespace Profiler {
 
+static tcmalloc::MallocExtension::AllocationProfilingToken* alloc_profiler = nullptr;
+
 absl::StatusOr<std::string> TcmallocProfiler::tcmallocHeapProfile() {
   auto profile = tcmalloc::MallocExtension::SnapshotCurrent(tcmalloc::ProfileType::kHeap);
   return tcmalloc::Marshal(profile);
 }
 
+absl::Status TcmallocProfiler::startAllocationProfile() {
+  ASSERT_IS_MAIN_OR_TEST_THREAD();
+  if (alloc_profiler != nullptr) {
+    return absl::Status(absl::StatusCode::kFailedPrecondition,
+                        "Allocation profiler has already started");
+  }
+  alloc_profiler = new tcmalloc::MallocExtension::AllocationProfilingToken(
+      tcmalloc::MallocExtension::StartAllocationProfiling());
+  return absl::OkStatus();
+}
+
+absl::StatusOr<std::string> TcmallocProfiler::stopAllocationProfile() {
+  ASSERT_IS_MAIN_OR_TEST_THREAD();
+  if (!alloc_profiler) {
+    return absl::Status(absl::StatusCode::kFailedPrecondition,
+                        "Allocation profiler is not started");
+  }
+  const auto profile = std::move(*alloc_profiler).Stop();
+  const auto result = tcmalloc::Marshal(profile);
+  delete alloc_profiler;
+  alloc_profiler = nullptr;
+  return result;
+}
+
 } // namespace Profiler
 } // namespace Envoy
 
@@ -86,6 +114,16 @@ absl::StatusOr<std::string> TcmallocProfiler::tcmallocHeapProfile() {
                       "Heap profile is not implemented in current build");
 }
 
+absl::Status TcmallocProfiler::startAllocationProfile() {
+  return absl::Status(absl::StatusCode::kUnimplemented,
+                      "Allocation profile is not implemented in current build");
+}
+
+absl::StatusOr<std::string> TcmallocProfiler::stopAllocationProfile() {
+  return absl::Status(absl::StatusCode::kUnimplemented,
+                      "Allocation profile is not implemented in current build");
+}
+
 } // namespace Profiler
 } // namespace Envoy
 

diff --git a/source/common/profiler/profiler.h b/source/common/profiler/profiler.h
@@ -66,12 +66,15 @@ class Heap {
 
 /**
  * Default profiler which will be enabled when tcmalloc (not `gperftools`) is used.
+ * This class is not thread-safe.
  */
 class TcmallocProfiler {
 public:
   TcmallocProfiler() = default;
 
   static absl::StatusOr<std::string> tcmallocHeapProfile();
+  static absl::Status startAllocationProfile();
+  static absl::StatusOr<std::string> stopAllocationProfile();
 };
 
 } // namespace Profiler

diff --git a/source/server/admin/admin.cc b/source/server/admin/admin.cc
@@ -166,6 +166,13 @@ AdminImpl::AdminImpl(const std::string& profile_path, Server::Instance& server,
           makeHandler("/heap_dump", "dump current Envoy heap (if supported)",
                       MAKE_ADMIN_HANDLER(tcmalloc_profiling_handler_.handlerHeapDump), false,
                       false),
+          makeHandler("/allocprofiler", "enable/disable the allocation profiler (if supported)",
+                      MAKE_ADMIN_HANDLER(tcmalloc_profiling_handler_.handlerAllocationProfiler),
+                      false, true,
+                      {{Admin::ParamDescriptor::Type::Enum,
+                        "enable",
+                        "enable/disable the allocation profiler",
+                        {"y", "n"}}}),
           makeHandler("/healthcheck/fail", "cause the server to fail health checks",
                       MAKE_ADMIN_HANDLER(server_cmd_handler_.handlerHealthcheckFail), false, true),
           makeHandler("/healthcheck/ok", "cause the server to pass health checks",

diff --git a/source/server/admin/profiling_handler.cc b/source/server/admin/profiling_handler.cc
@@ -93,5 +93,34 @@ Http::Code TcmallocProfilingHandler::handlerHeapDump(Http::ResponseHeaderMap&,
   return Http::Code::NotImplemented;
 }
 
+Http::Code TcmallocProfilingHandler::handlerAllocationProfiler(Http::ResponseHeaderMap&,
+                                                               Buffer::Instance& response,
+                                                               AdminStream& admin_stream) {
+  Http::Utility::QueryParamsMulti query_params = admin_stream.queryParams();
+  const auto enableVal = query_params.getFirstValue("enable");
+  if (query_params.data().size() != 1 || !enableVal.has_value() ||
+      (enableVal.value() != "y" && enableVal.value() != "n")) {
+    response.add("?enable=<y|n>\n");
+    return Http::Code::BadRequest;
+  }
+  const bool enable = enableVal.value() == "y";
+  if (enable) {
+    const auto started = Profiler::TcmallocProfiler::startAllocationProfile();
+    if (!started.ok()) {
+      response.add(started.message());
+      return Http::Code::BadRequest;
+    }
+    response.add("OK\n");
+    return Http::Code::OK;
+  }
+  const auto profile = Profiler::TcmallocProfiler::stopAllocationProfile();
+  if (!profile.ok()) {
+    response.add(profile.status().message());
+    return Http::Code::BadRequest;
+  }
+  response.add(profile.value());
+  return Http::Code::OK;
+}
+
 } // namespace Server
 } // namespace Envoy
diff --git a/source/server/admin/profiling_handler.h b/source/server/admin/profiling_handler.h
@@ -31,6 +31,9 @@ class TcmallocProfilingHandler {
 
   Http::Code handlerHeapDump(Http::ResponseHeaderMap& response_headers, Buffer::Instance& response,
                              AdminStream&);
+
+  Http::Code handlerAllocationProfiler(Http::ResponseHeaderMap& response_headers,
+                                       Buffer::Instance& response, AdminStream&);
 };
 
 } // namespace Server

diff --git a/test/server/admin/admin_test.cc b/test/server/admin/admin_test.cc
@@ -142,6 +142,8 @@ TEST_P(AdminInstanceTest, Help) {
   EXPECT_EQ(Http::Code::OK, getCallback("/help", header_map, response));
   const std::string expected = R"EOF(admin commands are:
   /: Admin home page
+  /allocprofiler (POST): enable/disable the allocation profiler (if supported)
+      enable: enable/disable the allocation profiler; One of (y, n)
   /certs: print certs on machine
   /clusters: upstream cluster status
   /config_dump: dump current Envoy configs (experimental)

diff --git a/test/server/admin/profiling_handler_test.cc b/test/server/admin/profiling_handler_test.cc
@@ -88,5 +88,25 @@ TEST_P(AdminInstanceTest, AdminHeapDump) {
 #endif
 }
 
+TEST_P(AdminInstanceTest, AdminAllocationDump) {
+  Buffer::OwnedImpl data;
+  Http::TestResponseHeaderMapImpl header_map;
+  EXPECT_EQ(Http::Code::BadRequest, postCallback("/allocprofiler", header_map, data));
+
+#ifdef TCMALLOC
+  EXPECT_EQ(Http::Code::OK, postCallback("/allocprofiler?enable=y", header_map, data));
+  EXPECT_EQ(Http::Code::OK, postCallback("/allocprofiler?enable=n", header_map, data));
+
+  // Out-of-order calls
+  EXPECT_EQ(Http::Code::BadRequest, postCallback("/allocprofiler?enable=n", header_map, data));
+  EXPECT_EQ(Http::Code::OK, postCallback("/allocprofiler?enable=y", header_map, data));
+  EXPECT_EQ(Http::Code::BadRequest, postCallback("/allocprofiler?enable=y", header_map, data));
+  EXPECT_EQ(Http::Code::OK, postCallback("/allocprofiler?enable=n", header_map, data));
+#else
+  EXPECT_EQ(Http::Code::BadRequest, postCallback("/allocprofiler?enable=y", header_map, data));
+  EXPECT_EQ(Http::Code::BadRequest, postCallback("/allocprofiler?enable=n", header_map, data));
+#endif
+}
+
 } // namespace Server
 } // namespace Envoy