Support memory service

apache · Jan 16, 2024 · f7b22c0 · f7b22c0
1 parent d131b0d
commit f7b22c0
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 39 deletions.
diff --git a/src/brpc/builtin/index_service.cpp b/src/brpc/builtin/index_service.cpp
@@ -164,8 +164,7 @@ void IndexService::default_method(::google::protobuf::RpcController* controller,
        << (!FLAGS_enable_threads_service ? " (disabled)" : "") << NL
        << Path("/dir", html_addr) << " : Browse directories and files"
        << (!FLAGS_enable_dir_service ? " (disabled)" : "") << NL
-       << Path("/tcmalloc", html_addr) << " : Get TCMalloc information"
-       << (!IsTCMallocEnabled() ? " (disabled)" : "") << NL;
+       << Path("/memory", html_addr) << " : Get malloc allocator information" << NL;
     if (use_html) {
         os << "</body></html>";
     }

diff --git a/src/brpc/builtin/tcmalloc_service.cpp → src/brpc/builtin/memory_service.cpp b/src/brpc/builtin/tcmalloc_service.cpp → src/brpc/builtin/memory_service.cpp
@@ -19,7 +19,7 @@
 #include "butil/logging.h"
 #include "brpc/controller.h"           // Controller
 #include "brpc/closure_guard.h"        // ClosureGuard
-#include "brpc/builtin/tcmalloc_service.h"
+#include "brpc/builtin/memory_service.h"
 #include "brpc/details/tcmalloc_extension.h"
 
 namespace brpc {
@@ -36,24 +36,9 @@ static inline void get_tcmalloc_num_prop(MallocExtension* malloc_ext,
     }
 }
 
-void TcmallocService::default_method(::google::protobuf::RpcController* cntl_base,
-                                    const ::brpc::TCMallocRequest*,
-                                    ::brpc::TCMallocResponse*,
-                                    ::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();
-
+static void get_tcmalloc_memory_info(butil::IOBuf& out) {
     MallocExtension* malloc_ext = MallocExtension::instance();
-    if (!malloc_ext) {
-        resp.append("tcmalloc is not enabled");
-        cntl->http_response().set_status_code(HTTP_STATUS_FORBIDDEN);
-        return;
-    }
-
     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);
@@ -70,7 +55,26 @@ void TcmallocService::default_method(::google::protobuf::RpcController* cntl_bas
     malloc_ext->GetStats(buf.get(), len);
     os << buf.get();
 
-    os.move_to(resp);
+    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
diff --git a/src/brpc/builtin/tcmalloc_service.h → src/brpc/builtin/memory_service.h b/src/brpc/builtin/tcmalloc_service.h → src/brpc/builtin/memory_service.h
@@ -15,22 +15,22 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#ifndef  BRPC_TCMALLOC_SERVICE_H
-#define  BRPC_TCMALLOC_SERVICE_H
+#ifndef  BRPC_MALLOC_SERVICE_H
+#define  BRPC_MALLOC_SERVICE_H
 
 #include "brpc/builtin_service.pb.h"
 
 namespace brpc {
 
-class TcmallocService : public tcmalloc {
+class MemoryService : public memory {
 public:
     void default_method(::google::protobuf::RpcController* cntl_base,
-                        const ::brpc::TCMallocRequest* request,
-                        ::brpc::TCMallocResponse* response,
+                        const ::brpc::MemoryRequest* request,
+                        ::brpc::MemoryResponse* response,
                         ::google::protobuf::Closure* done) override;
 };
 
 } // namespace brpc
 
 
-#endif  //BRPC_TCMALLOC_SERVICE_H
+#endif  // BRPC_MALLOC_SERVICE_H
diff --git a/src/brpc/builtin_service.proto b/src/brpc/builtin_service.proto
@@ -61,8 +61,8 @@ message VLogRequest {}
 message VLogResponse {}
 message MetricsRequest {}
 message MetricsResponse {}
-message TCMallocRequest {}
-message TCMallocResponse {}
+message MemoryRequest {}
+message MemoryResponse {}
 message BadMethodRequest {
     required string service_name = 1;
 }
@@ -172,6 +172,6 @@ service dir {
     rpc default_method(DirRequest) returns (DirResponse);
 }
 
-service tcmalloc {
-    rpc default_method(TCMallocRequest) returns (TCMallocResponse);
+service memory {
+    rpc default_method(MemoryRequest) returns (MemoryResponse);
 }
diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp
@@ -64,7 +64,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/tcmalloc_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,9 +528,8 @@ int Server::AddBuiltinServices() {
         LOG(ERROR) << "Fail to add ThreadsService";
         return -1;
     }
-    if (IsTCMallocEnabled() &&
-        AddBuiltinService(new (std::nothrow) TcmallocService)) {
-        LOG(ERROR) << "Fail to add TcmallocService";
+    if (AddBuiltinService(new (std::nothrow) MemoryService)) {
+        LOG(ERROR) << "Fail to add MemoryService";
         return -1;
     }
 

diff --git a/test/brpc_builtin_service_unittest.cpp b/test/brpc_builtin_service_unittest.cpp
@@ -50,7 +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/tcmalloc_service.h"
+#include "brpc/builtin/memory_service.h"
 #include "brpc/builtin/common.h"
 #include "brpc/builtin/bad_method_service.h"
 #include "echo.pb.h"
@@ -835,10 +835,10 @@ TEST_F(BuiltinServiceTest, sockets) {
     }    
 }
 
-TEST_F(BuiltinServiceTest, tcmalloc) {
-    brpc::TcmallocService service;
-    brpc::TCMallocRequest req;
-    brpc::TCMallocResponse res;
+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);