Add controller flag to manage ConcurrencyRemover lifecycle for CallAfterRpcResp

In SendRpcResponse, ConcurrencyRemover was destroyed before CallAfterRpcResp
was called, meaning concurrency control didn't cover the after-response callback.
This could lead to inaccurate concurrency tracking and latency measurements.

This change adds a controller-level flag that is automatically set when
set_after_rpc_resp_fn() is called. When the flag is true, ConcurrencyRemover's
lifetime is extended to cover CallAfterRpcResp.

Implementation:
- Add _concurrency_remover_manages_after_rpc_resp flag to Controller
- Automatically set flag to true in set_after_rpc_resp_fn()
- Use unique_ptr with explicit reset() for clear control flow
- Apply to both baidu_rpc_protocol and http_rpc_protocol

When false (default): Original behavior - ConcurrencyRemover is released before
CallAfterRpcResp via explicit reset().

When true (set_after_rpc_resp_fn called): ConcurrencyRemover lives until the
end of BRPC_SCOPE_EXIT, covering the entire response lifecycle including
after-response callbacks.

Author: feng-y

src/brpc/controller.cpp

@@ -298,6 +298,7 @@ void Controller::ResetPods() {
     _response_streams.clear();
     _remote_stream_settings = NULL;
     _auth_flags = 0;
+    _concurrency_remover_manages_after_rpc_resp = false;
     _rpc_received_us = 0;
 }
 

src/brpc/controller.h

@@ -621,10 +621,19 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
                                                const google::protobuf::Message* req,
                                                const google::protobuf::Message* res)>;
 
-    void set_after_rpc_resp_fn(AfterRpcRespFnType&& fn) { _after_rpc_resp_fn = fn; }
+    void set_after_rpc_resp_fn(AfterRpcRespFnType&& fn) {
+        _after_rpc_resp_fn = fn;
+        // When after_rpc_resp_fn is set, enable concurrency management for it
+        _concurrency_remover_manages_after_rpc_resp = true;
+    }
 
     void CallAfterRpcResp(const google::protobuf::Message* req, const google::protobuf::Message* res);
 
+    // Check whether ConcurrencyRemover should manage the lifecycle of CallAfterRpcResp.
+    bool concurrency_remover_manages_after_rpc_resp() const {
+        return _concurrency_remover_manages_after_rpc_resp;
+    }
+
     void set_request_content_type(ContentType type) {
         _request_content_type = type;
     }
@@ -921,6 +930,7 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     uint32_t _auth_flags;
 
     AfterRpcRespFnType _after_rpc_resp_fn;
+    bool _concurrency_remover_manages_after_rpc_resp;
 
     // The point in time when the rpc is read from the socket
     int64_t _rpc_received_us;

src/brpc/policy/baidu_rpc_protocol.cpp

@@ -285,9 +285,12 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
 
     // Recycle resources at the end of this function.
     BRPC_SCOPE_EXIT {
-        {
-            // Remove concurrency and record latency at first.
-            ConcurrencyRemover concurrency_remover(method_status, cntl, received_us);
+        std::unique_ptr<ConcurrencyRemover> concurrency_remover_ptr(
+            new ConcurrencyRemover(method_status, cntl, received_us));
+
+        if (!cntl->concurrency_remover_manages_after_rpc_resp()) {
+            // Original behavior: remove concurrency before CallAfterRpcResp
+            concurrency_remover_ptr.reset();
         }
 
         std::unique_ptr<Controller, LogErrorTextAndDelete> recycle_cntl(cntl);
@@ -302,6 +305,8 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
         } else {
             BaiduProxyPBMessages::Return(static_cast<BaiduProxyPBMessages*>(messages));
         }
+        // If concurrency_remover_manages_after_rpc_resp is true,
+        // concurrency_remover_ptr will be destroyed here
     };
 
     StreamIds response_stream_ids = accessor.response_streams();

src/brpc/policy/http_rpc_protocol.cpp

@@ -813,10 +813,7 @@ class HttpResponseSenderAsDone : public google::protobuf::Closure {
 public:
     explicit HttpResponseSenderAsDone(HttpResponseSender* s) : _sender(std::move(*s)) {}
     void Run() override {
-        if (NULL != _sender._messages) {
-            _sender._cntl->CallAfterRpcResp(_sender._messages->Request(),
-                                            _sender._messages->Response());
-        }
+        // CallAfterRpcResp is now called in HttpResponseSender::~HttpResponseSender()
         delete this;
     }
 
@@ -827,9 +824,23 @@ class HttpResponseSenderAsDone : public google::protobuf::Closure {
 HttpResponseSender::~HttpResponseSender() {
     // Return messages to factory at the end.
     BRPC_SCOPE_EXIT {
+        std::unique_ptr<ConcurrencyRemover> concurrency_remover_ptr;
+        Controller* cntl = _cntl.get();
+        if (cntl != NULL) {
+            concurrency_remover_ptr.reset(new ConcurrencyRemover(_method_status, cntl, _received_us));
+
+            if (!cntl->concurrency_remover_manages_after_rpc_resp()) {
+                // Original behavior: remove concurrency before CallAfterRpcResp
+                concurrency_remover_ptr.reset();
+            }
+        }
+
         if (NULL != _messages) {
+            _cntl->CallAfterRpcResp(_messages->Request(), _messages->Response());
             _cntl->server()->options().rpc_pb_message_factory->Return(_messages);
         }
+        // If concurrency_remover_manages_after_rpc_resp is true,
+        // concurrency_remover_ptr will be destroyed here
     };
     Controller* cntl = _cntl.get();
     if (cntl == NULL) {
@@ -840,7 +851,6 @@ HttpResponseSender::~HttpResponseSender() {
     if (span) {
         span->set_start_send_us(butil::cpuwide_time_us());
     }
-    ConcurrencyRemover concurrency_remover(_method_status, cntl, _received_us);
     Socket* socket = accessor.get_sending_socket();
     const google::protobuf::Message* res = NULL != _messages ? _messages->Response() : NULL;