Support memory builtin service (#2505)

* Support TCMalloc builtin service

* Support memory service

src/brpc/builtin/index_service.cpp

@@ -165,7 +165,8 @@ void IndexService::default_method(::google::protobuf::RpcController* controller,
        << Path("/threads", html_addr) << " : Check pstack"
        << (!FLAGS_enable_threads_service ? " (disabled)" : "") << NL
        << Path("/dir", html_addr) << " : Browse directories and files"
-       << (!FLAGS_enable_dir_service ? " (disabled)" : "") << NL;
+       << (!FLAGS_enable_dir_service ? " (disabled)" : "") << NL
+       << Path("/memory", html_addr) << " : Get malloc allocator information" << NL;
     if (use_html) {
         os << "</body></html>";
     }

src/brpc/builtin/memory_service.cpp

@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "butil/time.h"
+#include "butil/logging.h"
+#include "brpc/controller.h"           // Controller
+#include "brpc/closure_guard.h"        // ClosureGuard
+#include "brpc/builtin/memory_service.h"
+#include "brpc/details/tcmalloc_extension.h"
+
+namespace brpc {
+
+DEFINE_int32(max_tc_stats_buf_len, 32 * 1024, "max length of TCMalloc stats");
+BRPC_VALIDATE_GFLAG(max_tc_stats_buf_len, PositiveInteger);
+
+static inline void get_tcmalloc_num_prop(MallocExtension* malloc_ext,
+                                         const char* prop_name,
+                                         butil::IOBufBuilder& os) {
+    size_t value;
+    if (malloc_ext->GetNumericProperty(prop_name, &value)) {
+        os << prop_name << ": " << value << "\n";
+    }
+}
+
+static void get_tcmalloc_memory_info(butil::IOBuf& out) {
+    MallocExtension* malloc_ext = MallocExtension::instance();
+    butil::IOBufBuilder os;
+    os << "------------------------------------------------\n";
+    get_tcmalloc_num_prop(malloc_ext, "generic.total_physical_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "generic.current_allocated_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "generic.heap_size", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.current_total_thread_cache_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.central_cache_free_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.transfer_cache_free_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.thread_cache_free_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.pageheap_free_bytes", os);
+    get_tcmalloc_num_prop(malloc_ext, "tcmalloc.pageheap_unmapped_bytes", os);
+
+    int32_t len = FLAGS_max_tc_stats_buf_len;
+    std::unique_ptr<char[]> buf(new char[len]);
+    malloc_ext->GetStats(buf.get(), len);
+    os << buf.get();
+
+    os.move_to(out);
+}
+
+void MemoryService::default_method(::google::protobuf::RpcController* cntl_base,
+                                    const ::brpc::MemoryRequest*,
+                                    ::brpc::MemoryResponse*,
+                                    ::google::protobuf::Closure* done) {
+    ClosureGuard done_guard(done);
+    auto cntl = static_cast<Controller*>(cntl_base);
+    cntl->http_response().set_content_type("text/plain");
+    butil::IOBuf& resp = cntl->response_attachment();
+
+    if (IsTCMallocEnabled()) {
+        butil::IOBufBuilder os;
+        get_tcmalloc_memory_info(resp);
+    } else {
+        resp.append("tcmalloc is not enabled");
+        cntl->http_response().set_status_code(HTTP_STATUS_FORBIDDEN);
+        return;
+    }
+}
+
+} // namespace brpc

src/brpc/builtin/memory_service.h

@@ -0,0 +1,36 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#ifndef  BRPC_MALLOC_SERVICE_H
+#define  BRPC_MALLOC_SERVICE_H
+
+#include "brpc/builtin_service.pb.h"
+
+namespace brpc {
+
+class MemoryService : public memory {
+public:
+    void default_method(::google::protobuf::RpcController* cntl_base,
+                        const ::brpc::MemoryRequest* request,
+                        ::brpc::MemoryResponse* response,
+                        ::google::protobuf::Closure* done) override;
+};
+
+} // namespace brpc
+
+
+#endif  // BRPC_MALLOC_SERVICE_H

src/brpc/builtin_service.proto

@@ -61,6 +61,8 @@ message VLogRequest {}
 message VLogResponse {}
 message MetricsRequest {}
 message MetricsResponse {}
+message MemoryRequest {}
+message MemoryResponse {}
 message BadMethodRequest {
     required string service_name = 1;
 }
@@ -169,3 +171,7 @@ service rpcz {
 service dir {
     rpc default_method(DirRequest) returns (DirResponse);
 }
+
+service memory {
+    rpc default_method(MemoryRequest) returns (MemoryResponse);
+}

src/brpc/details/tcmalloc_extension.cpp

@@ -50,6 +50,10 @@ bool IsHeapProfilerEnabled() {
     return MallocExtension::instance() != NULL;
 }
 
+bool IsTCMallocEnabled() {
+    return IsHeapProfilerEnabled();
+}
+
 static bool check_TCMALLOC_SAMPLE_PARAMETER() {
     char* str = getenv("TCMALLOC_SAMPLE_PARAMETER");
     if (str == NULL) {

src/brpc/details/tcmalloc_extension.h

@@ -315,5 +315,7 @@ class PERFTOOLS_DLL_DECL MallocExtension {
 // True iff heap profiler is enabled.
 bool IsHeapProfilerEnabled();
 
+bool IsTCMallocEnabled();
+
 // True iff TCMALLOC_SAMPLE_PARAMETER is set in environment.
 bool has_TCMALLOC_SAMPLE_PARAMETER();

src/brpc/server.cpp

@@ -65,6 +65,7 @@
 #include "brpc/builtin/sockets_service.h"      // SocketsService
 #include "brpc/builtin/hotspots_service.h"     // HotspotsService
 #include "brpc/builtin/prometheus_metrics_service.h"
+#include "brpc/builtin/memory_service.h"
 #include "brpc/details/method_status.h"
 #include "brpc/load_balancer.h"
 #include "brpc/naming_service.h"
@@ -528,6 +529,10 @@ int Server::AddBuiltinServices() {
         LOG(ERROR) << "Fail to add ThreadsService";
         return -1;
     }
+    if (AddBuiltinService(new (std::nothrow) MemoryService)) {
+        LOG(ERROR) << "Fail to add MemoryService";
+        return -1;
+    }
 
 #if !BRPC_WITH_GLOG
     if (AddBuiltinService(new (std::nothrow) VLogService)) {

test/brpc_builtin_service_unittest.cpp

@@ -50,6 +50,7 @@
 #include "brpc/builtin/bthreads_service.h"     // BthreadsService
 #include "brpc/builtin/ids_service.h"          // IdsService
 #include "brpc/builtin/sockets_service.h"      // SocketsService
+#include "brpc/builtin/memory_service.h"
 #include "brpc/builtin/common.h"
 #include "brpc/builtin/bad_method_service.h"
 #include "echo.pb.h"
@@ -902,3 +903,21 @@ TEST_F(BuiltinServiceTest, sockets) {
         CheckContent(cntl, "fd=-1");
     }    
 }
+
+TEST_F(BuiltinServiceTest, memory) {
+    brpc::MemoryService service;
+    brpc::MemoryRequest req;
+    brpc::MemoryResponse res;
+    brpc::Controller cntl;
+    ClosureChecker done;
+    service.default_method(&cntl, &req, &res, &done);
+    EXPECT_FALSE(cntl.Failed());
+    CheckContent(cntl, "generic.current_allocated_bytes");
+    CheckContent(cntl, "generic.heap_size");
+    CheckContent(cntl, "tcmalloc.current_total_thread_cache_bytes");
+    CheckContent(cntl, "tcmalloc.central_cache_free_bytes");
+    CheckContent(cntl, "tcmalloc.transfer_cache_free_bytes");
+    CheckContent(cntl, "tcmalloc.thread_cache_free_bytes");
+    CheckContent(cntl, "tcmalloc.pageheap_free_bytes");
+    CheckContent(cntl, "tcmalloc.pageheap_unmapped_bytes");
+}