bugfix: fix memory growth caused by brpc arena configuration.

Author: wly-115

third_party/Mooncake

@@ -148,8 +148,8 @@ void ChatCompletionsImpl(std::unique_ptr<Service>& service,
     return;
   }
 
-  auto call = std::make_shared<ChatCall>(
-      ctrl, guard.release(), req_pb, resp_pb, arena != nullptr /*use_arena*/);
+  auto call =
+      std::make_shared<ChatCall>(ctrl, guard.release(), req_pb, resp_pb);
   service->process_async(call);
 }
 }  // namespace
@@ -167,19 +167,17 @@ void APIService::ChatCompletionsHttp(
     LOG(ERROR) << "brpc request | respose | controller is null";
     return;
   }
-
+  auto arena = response->GetArena();
   auto ctrl = reinterpret_cast<brpc::Controller*>(controller);
 
   if (FLAGS_backend == "llm") {
-    auto arena = response->GetArena();
     CHECK(chat_service_impl_) << " chat service is invalid.";
     ChatCompletionsImpl<ChatCall, ChatServiceImpl>(
         chat_service_impl_, done_guard, arena, ctrl);
   } else if (FLAGS_backend == "vlm") {
     CHECK(mm_chat_service_impl_) << " mm chat service is invalid.";
-    // TODO: fix me - temporarily using heap allocation instead of arena
     ChatCompletionsImpl<MMChatCall, MMChatServiceImpl>(
-        mm_chat_service_impl_, done_guard, nullptr, ctrl);
+        mm_chat_service_impl_, done_guard, arena, ctrl);
   }
 }
 

xllm/api_service/api_service.cpp

@@ -39,13 +39,8 @@ class StreamCall : public Call {
   StreamCall(brpc::Controller* controller,
              ::google::protobuf::Closure* done,
              Request* request,
-             Response* response,
-             bool use_arena = true)
-      : Call(controller),
-        done_(done),
-        request_(request),
-        response_(response),
-        use_arena_(use_arena) {
+             Response* response)
+      : Call(controller), done_(done), request_(request), response_(response) {
     stream_ = request_->stream();
     if (stream_) {
       pa_ = controller_->CreateProgressiveAttachment();
@@ -72,10 +67,6 @@ class StreamCall : public Call {
     if (!stream_) {
       done_->Run();
     }
-    if (!use_arena_) {
-      delete request_;
-      delete response_;
-    }
   }
 
   bool write_and_finish(Response& response) {
@@ -151,7 +142,6 @@ class StreamCall : public Call {
   Response* response_;
 
   bool stream_ = false;
-  bool use_arena_ = true;
   butil::intrusive_ptr<brpc::ProgressiveAttachment> pa_;
   butil::IOBuf io_buf_;
 

xllm/api_service/stream_call.h

@@ -50,6 +50,8 @@ bool XllmServer::start(std::unique_ptr<APIService> service) {
   }
 
   brpc::ServerOptions options;
+  options.rpc_pb_message_factory =
+      brpc::GetArenaRpcPBMessageFactory<1024 * 1024, 1024 * 1024 * 100>();
   options.idle_timeout_sec = FLAGS_rpc_idle_timeout_s;
   options.num_threads = FLAGS_num_threads;
   if (server_->Start(FLAGS_port, &options) != 0) {