Add gflag 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 FLAGS_concurrency_remover_manages_after_rpc_resp (default: false)
and automatically sets it to controller when set_after_rpc_resp_fn is called.

Implementation:
- Add _concurrency_remover_manages_after_rpc_resp flag to Controller
- In set_after_rpc_resp_fn(), read gflag value and set to controller instance
- In baidu_rpc_protocol, use controller flag instead of global gflag
- Use unique_ptr with explicit reset() for clear control flow

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

When true (gflag enabled when callback set): ConcurrencyRemover lives until the
end of BRPC_SCOPE_EXIT, covering the entire response lifecycle.

Note: HTTP protocol not modified in this change due to its more complex async
flow. Can be addressed separately if needed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: feng-y

src/brpc/controller.cpp

@@ -127,6 +127,7 @@ const Controller* GetSubControllerOfSelectiveChannel(
 
 DECLARE_bool(usercode_in_pthread);
 DECLARE_bool(usercode_in_coroutine);
+DECLARE_bool(concurrency_remover_manages_after_rpc_resp);
 static const int MAX_RETRY_COUNT = 1000;
 static bvar::Adder<int64_t>* g_ncontroller = NULL;
 
@@ -298,6 +299,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;
 }
 
@@ -1593,6 +1595,12 @@ int Controller::GetSockOption(int level, int optname, void* optval, socklen_t* o
     }
 }
 
+void Controller::set_after_rpc_resp_fn(AfterRpcRespFnType&& fn) {
+    _after_rpc_resp_fn = fn;
+    // Set the flag from global gflag when after_rpc_resp_fn is set
+    _concurrency_remover_manages_after_rpc_resp = FLAGS_concurrency_remover_manages_after_rpc_resp;
+}
+
 void Controller::CallAfterRpcResp(const google::protobuf::Message* req, const google::protobuf::Message* res) {
     if (_after_rpc_resp_fn) {
         _after_rpc_resp_fn(this, req, res);

src/brpc/controller.h

@@ -621,10 +621,15 @@ 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);
 
     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 +926,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

@@ -62,6 +62,11 @@ DEFINE_bool(baidu_protocol_use_fullname, true,
 DEFINE_bool(baidu_std_protocol_deliver_timeout_ms, false,
             "If this flag is true, baidu_std puts timeout_ms in requests.");
 
+DEFINE_bool(concurrency_remover_manages_after_rpc_resp, false,
+            "If this flag is true, ConcurrencyRemover will manage the lifecycle "
+            "of CallAfterRpcResp, ensuring concurrency control covers the entire "
+            "response processing including after-response callbacks.");
+
 DECLARE_bool(pb_enum_as_number);
 
 // Notes:
@@ -285,9 +290,14 @@ 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));
+
+        // Only manage CallAfterRpcResp lifecycle if the flag is set
+        // (which happens when set_after_rpc_resp_fn is called with the gflag enabled)
+        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 +312,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();