admin: add allocation profiler

Change-Id: I79629537ab83c54b655de7ef1010b29665d30541
Signed-off-by: Kuat Yessenov <[email protected]>

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

source/common/profiler/profiler.cc

@@ -68,11 +68,33 @@ bool Heap::stopProfiler() { return false; }
 namespace Envoy {
 namespace Profiler {
 
+static absl::optional<tcmalloc::MallocExtension::AllocationProfilingToken> alloc_profiler =
+    absl::nullopt;
+
 absl::StatusOr<std::string> TcmallocProfiler::tcmallocHeapProfile() {
   auto profile = tcmalloc::MallocExtension::SnapshotCurrent(tcmalloc::ProfileType::kHeap);
   return tcmalloc::Marshal(profile);
 }
 
+bool TcmallocProfiler::startAllocationProfile() {
+  if (alloc_profiler) {
+    return false;
+  }
+  alloc_profiler = tcmalloc::MallocExtension::StartAllocationProfiling();
+  return true;
+}
+
+absl::StatusOr<std::string> TcmallocProfiler::stopAllocationProfile() {
+  if (!alloc_profiler) {
+    return absl::Status(absl::StatusCode::kFailedPrecondition,
+                        "Allocation profiler is not started");
+  }
+  const auto profile = std::move(alloc_profiler.value()).Stop();
+  const auto result = tcmalloc::Marshal(profile);
+  alloc_profiler = absl::nullopt;
+  return result;
+}
+
 } // namespace Profiler
 } // namespace Envoy
 
@@ -86,6 +108,13 @@ absl::StatusOr<std::string> TcmallocProfiler::tcmallocHeapProfile() {
                       "Heap profile is not implemented in current build");
 }
 
+bool TcmallocProfiler::startAllocationProfile() { return false; }
+
+absl::StatusOr<std::string> TcmallocProfiler::stopAllocationProfile() {
+  return absl::Status(absl::StatusCode::kUnimplemented,
+                      "Allocation profile is not implemented in current build");
+}
+
 } // namespace Profiler
 } // namespace Envoy
 

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 bool startAllocationProfile();
+  static absl::StatusOr<std::string> stopAllocationProfile();
 };
 
 } // namespace Profiler

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",

source/server/admin/profiling_handler.cc

@@ -93,5 +93,33 @@ 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 bool started = Profiler::TcmallocProfiler::startAllocationProfile();
+    if (!started) {
+      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

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

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