Refactor bthread span lifecycle management and optimize span API with smart pointer reuse (#3068)

Author: lhh

docs/cn/rpcz.md

@@ -66,4 +66,14 @@ bthread_attr_t attr = { BTHREAD_STACKTYPE_NORMAL, BTHREAD_INHERIT_SPAN, NULL };
 bthread_start_urgent(&tid, &attr, thread_proc, arg);
 ```
 
-注意：使用这种方式创建子bthread来发送rpc，请确保rpc在server返回response之前完成，否则可能导致使用被释放的Span对象而出core。
+### Span生命周期管理
+
+brpc使用智能指针（`std::shared_ptr`/`std::weak_ptr`）管理Span对象的生命周期，并通过自旋锁保护并发访问，解决了以下问题：
+
+1. **Use-after-free防护**：父Span通过`shared_ptr`持有子Span的强引用，TLS中使用`weak_ptr`存储，确保Span对象在被访问时仍然有效。即使server在子bthread完成前返回response，也不会导致访问已释放的Span对象。
+
+2. **线程安全**：使用自旋锁保护`_client_list`和`_info`的并发修改，支持多个bthread同时创建子span或添加annotation。
+
+3. **自动生命周期管理**：当父Span销毁时，会自动清理所有子Span（通过`_client_list.clear()`），无需手动管理。
+
+使用`BTHREAD_INHERIT_SPAN`创建子bthread时，不再需要担心Span对象的生命周期问题，可以安全地在异步场景中使用。

src/brpc/builtin/rpcz_service.cpp

@@ -211,7 +211,15 @@ static void PrintAnnotations(
             PrintElapse(os, anno_time, last_time);
             os << ' ';
             if (span) {
-                os << '[' << span_type_str << ' ' << SPAN_ID_STR << '=' << Hex(span->span_id()) << "] ";
+                const char* short_type = "SPAN";
+                if (span->type() == SPAN_TYPE_SERVER) {
+                    short_type = "Server";
+                } else if (span->type() == SPAN_TYPE_CLIENT) {
+                    short_type = "Client";
+                } else if (span->type() == SPAN_TYPE_BTHREAD) {
+                    short_type = "Bthread";
+                }
+                os << '[' << short_type << " SPAN#" << Hex(span->span_id()) << "] ";
             }
             os << WebEscape(a);
             if (a.empty() || butil::back_char(a) != '\n') {
@@ -292,11 +300,11 @@ static void PrintClientSpan(
 
     if (PrintAnnotationsAndRealTimeSpan(os, span.sent_real_us(),
                                         last_time, extr, num_extr, &span)) {
-        os << " [ClientSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Requested(" << span.request_size() << ") [1]" << std::endl;
+        os << " [Client SPAN#" << Hex(span.span_id()) << "] Requested(" << span.request_size() << ") [1]" << std::endl;
     }
     if (PrintAnnotationsAndRealTimeSpan(os, span.received_real_us(),
                                         last_time, extr, num_extr, &span)) {
-        os << " [ClientSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Received response(" << span.response_size() << ")";
+        os << " [Client SPAN#" << Hex(span.span_id()) << "] Received response(" << span.response_size() << ")";
         if (span.base_cid() != 0 && span.ending_cid() != 0) {
             int64_t ver = span.ending_cid() - span.base_cid();
             if (ver >= 1) {
@@ -310,13 +318,13 @@ static void PrintClientSpan(
 
     if (PrintAnnotationsAndRealTimeSpan(os, span.start_parse_real_us(),
                                         last_time, extr, num_extr, &span)) {
-        os << " [ClientSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Processing the response in a new bthread" << std::endl;
+        os << " [Client SPAN#" << Hex(span.span_id()) << "] Processing the response in a new bthread" << std::endl;
     }
 
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.start_callback_real_us(),
             last_time, extr, num_extr, &span)) {
-        os << " [ClientSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] " << (span.async() ? " Enter user's done" : " Back to user's callsite") << std::endl;
+        os << " [Client SPAN#" << Hex(span.span_id()) << "] " << (span.async() ? " Enter user's done" : " Back to user's callsite") << std::endl;
     }
 
     PrintAnnotations(os, std::numeric_limits<int64_t>::max(),
@@ -340,9 +348,9 @@ static void PrintBthreadSpan(std::ostream& os, const RpczSpan& span, int64_t* la
     extr[num_extr++] = &client_extr;
 
     // Print span id for bthread span context identification
-    os << " [BthreadSpan " << SPAN_ID_STR << '=' << Hex(span.span_id());
+    os << " [Bthread SPAN#" << Hex(span.span_id());
     if (span.parent_span_id() != 0) {
-        os << " parent_span=" << Hex(span.parent_span_id());
+        os << " parent#" << Hex(span.parent_span_id());
     }
     os << "] ";
 
@@ -377,15 +385,15 @@ static void PrintServerSpan(std::ostream& os, const RpczSpan& span,
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.start_parse_real_us(),
             &last_time, extr, ARRAY_SIZE(extr), &span)) {
-        os << " [ServerSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Processing the request in a new bthread" << std::endl;
+        os << " [Server SPAN#" << Hex(span.span_id()) << "] Processing the request in a new bthread" << std::endl;
     }
 
     bool entered_user_method = false;
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.start_callback_real_us(),
             &last_time, extr, ARRAY_SIZE(extr), &span)) {
         entered_user_method = true;
-        os << " [ServerSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Enter " << WebEscape(span.full_method_name()) << std::endl;
+        os << " [Server SPAN#" << Hex(span.span_id()) << "] Enter " << WebEscape(span.full_method_name()) << std::endl;
     }
 
     const int nclient = span.client_spans_size();
@@ -402,16 +410,16 @@ static void PrintServerSpan(std::ostream& os, const RpczSpan& span,
             os, span.start_send_real_us(),
             &last_time, extr, ARRAY_SIZE(extr), &span)) {
         if (entered_user_method) {
-            os << " [ServerSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Leave " << WebEscape(span.full_method_name()) << std::endl;
+            os << " [Server SPAN#" << Hex(span.span_id()) << "] Leave " << WebEscape(span.full_method_name()) << std::endl;
         } else {
-            os << " [ServerSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Responding" << std::endl;
+            os << " [Server SPAN#" << Hex(span.span_id()) << "] Responding" << std::endl;
         }
     }
 
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.sent_real_us(),
             &last_time, extr, ARRAY_SIZE(extr), &span)) {
-        os << " [ServerSpan " << SPAN_ID_STR << '=' << Hex(span.span_id()) << "] Responded(" << span.response_size() << ')' << std::endl;
+        os << " [Server SPAN#" << Hex(span.span_id()) << "] Responded(" << span.response_size() << ')' << std::endl;
     }
 
     PrintAnnotations(os, std::numeric_limits<int64_t>::max(),

src/brpc/channel.cpp

@@ -515,7 +515,7 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
             method_name = NULL_METHOD_STR;
         }
         std::shared_ptr<Span> span = Span::CreateClientSpan(
-            *method_name, start_send_real_us - start_send_us);
+            method_name, start_send_real_us - start_send_us);
         if (span) {
             ControllerPrivateAccessor accessor(cntl);
             span->set_log_id(cntl->log_id());

src/brpc/controller.cpp

@@ -1405,6 +1405,7 @@ uint64_t Controller::trace_id() const {
     }
     return 0;
 }
+
 uint64_t Controller::span_id() const {
     if (auto span = _span.lock()) {
         return span->span_id();
@@ -1737,4 +1738,24 @@ void Controller::DoPrintLogPrefix(std::ostream& os) const {
     }
 }
 
+
+ControllerPrivateAccessor& ControllerPrivateAccessor::set_span(
+    const std::shared_ptr<Span>& span) {
+    _cntl->_span = span;
+    return *this;
+}
+
+ControllerPrivateAccessor& ControllerPrivateAccessor::set_span(Span* span) {
+    if (span) {
+        _cntl->_span = span->shared_from_this();
+    } else {
+        _cntl->_span.reset();
+    }
+    return *this;
+}
+
+std::shared_ptr<Span> ControllerPrivateAccessor::span() const {
+    return _cntl->_span.lock();
+}
+
 } // namespace brpc

src/brpc/details/controller_private_accessor.h

@@ -30,9 +30,10 @@ class Message;
 }
 }
 
-
 namespace brpc {
 
+class Span;
+
 class AuthContext;
 
 // A wrapper to access some private methods/fields of `Controller'
@@ -90,17 +91,16 @@ class ControllerPrivateAccessor {
         return *this;
     }
 
-    ControllerPrivateAccessor &set_span(std::shared_ptr<Span> span) {
-        _cntl->_span = span;
-        return *this;
-    }
+    // Overloaded set_span methods to support both shared_ptr and raw pointer
+    ControllerPrivateAccessor &set_span(const std::shared_ptr<Span>& span);
+    ControllerPrivateAccessor &set_span(Span* span);
 
     ControllerPrivateAccessor &set_request_protocol(ProtocolType protocol) {
         _cntl->_request_protocol = protocol;
         return *this;
     }
 
-    std::shared_ptr<Span> span() const { return _cntl->_span.lock(); }
+    std::shared_ptr<Span> span() const;
 
     uint32_t pipelined_count() const { return _cntl->_pipelined_count; }
     void set_pipelined_count(uint32_t count) {  _cntl->_pipelined_count = count; }

src/brpc/policy/couchbase_protocol.cpp

@@ -160,8 +160,7 @@ void ProcessCouchbaseResponse(InputMessageBase* msg_base) {
   }
 
   ControllerPrivateAccessor accessor(cntl);
-  Span* span = accessor.span();
-  if (span) {
+  if (auto span = accessor.span()) {
     span->set_base_real_us(msg->base_real_us());
     span->set_received_us(msg->received_us());
     span->set_response_size(msg->meta.length());

src/brpc/span.cpp

@@ -62,8 +62,7 @@ void DestroyRpczParentSpan(void* ptr) {
 void EndBthreadSpan() {
     std::shared_ptr<Span> span = GetTlsParentSpan();
     if (span) {
-        bthread_id_t id = {bthread_self()};
-        span->set_ending_cid(id);
+        span->set_ending_tid(bthread_self());
     }
 
     ClearTlsParentSpan();
@@ -73,10 +72,11 @@ void SetTlsParentSpan(std::shared_ptr<Span> span) {
     using namespace bthread;
     LocalStorage ls = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_bls);
     if (ls.rpcz_parent_span) {
-        delete static_cast<std::weak_ptr<Span>*>(ls.rpcz_parent_span);
+        *static_cast<std::weak_ptr<Span>*>(ls.rpcz_parent_span) = span;
+    } else {
+        ls.rpcz_parent_span = new std::weak_ptr<Span>(span);
+        BAIDU_SET_VOLATILE_THREAD_LOCAL(tls_bls, ls);
     }
-    ls.rpcz_parent_span = new std::weak_ptr<Span>(span);
-    BAIDU_SET_VOLATILE_THREAD_LOCAL(tls_bls, ls);
 }
 
 std::shared_ptr<Span> GetTlsParentSpan() {
@@ -94,9 +94,7 @@ void ClearTlsParentSpan() {
     using namespace bthread;
     LocalStorage ls = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_bls);
     if (ls.rpcz_parent_span) {
-        delete static_cast<std::weak_ptr<Span>*>(ls.rpcz_parent_span);
-        ls.rpcz_parent_span = nullptr;
-        BAIDU_SET_VOLATILE_THREAD_LOCAL(tls_bls, ls);
+        static_cast<std::weak_ptr<Span>*>(ls.rpcz_parent_span)->reset();
     }
 }
 
@@ -122,9 +120,6 @@ void SpanDeleter::operator()(Span* r) const {
     // children.
     r->_client_list.clear();
     r->_info.clear();
-    // Destroy the spinlocks, as the destructor might not be invoked.
-    pthread_spin_destroy(&r->_client_list_spinlock);
-    pthread_spin_destroy(&r->_info_spinlock);
     butil::return_object(r);
 }
 
@@ -200,7 +195,8 @@ Span::Span(Forbidden) {
 }
 
 Span::~Span() {
-    // The destruction of the spinlock has been handled in SpanDeleter.
+    pthread_spin_destroy(&_client_list_spinlock);
+    pthread_spin_destroy(&_info_spinlock);
 }
 
 std::shared_ptr<Span> Span::CreateClientSpan(const std::string& full_method_name,
@@ -212,7 +208,7 @@ std::shared_ptr<Span> Span::CreateClientSpan(const std::string& full_method_name
     std::shared_ptr<Span> span(span_raw, SpanDeleter());
     span->_log_id = 0;
     span->_base_cid = INVALID_BTHREAD_ID;
-    span->_ending_cid = INVALID_BTHREAD_ID;
+    span->_ending_cid = INVALID_BTHREAD_ID;  // Client Span uses ending_cid
     span->_type = SPAN_TYPE_CLIENT;
     span->_async = false;
     span->_protocol = PROTOCOL_UNKNOWN;
@@ -254,7 +250,7 @@ std::shared_ptr<Span> Span::CreateBthreadSpan(const std::string& full_method_nam
     std::shared_ptr<Span> span(span_raw, SpanDeleter());
     span->_log_id = 0;
     span->_base_cid = INVALID_BTHREAD_ID;
-    span->_ending_cid = INVALID_BTHREAD_ID;
+    span->_ending_tid = INVALID_BTHREAD;  // Bthread Span uses ending_tid
     span->_type = SPAN_TYPE_BTHREAD;
     span->_async = false;
     span->_protocol = PROTOCOL_UNKNOWN;
@@ -307,7 +303,7 @@ std::shared_ptr<Span> Span::CreateServerSpan(
     span->_parent_span_id = parent_span_id;
     span->_log_id = 0;
     span->_base_cid = INVALID_BTHREAD_ID;
-    span->_ending_cid = INVALID_BTHREAD_ID;
+    span->_ending_cid = INVALID_BTHREAD_ID;  // Server Span uses ending_cid
     span->_type = SPAN_TYPE_SERVER;
     span->_async = false;
     span->_protocol = PROTOCOL_UNKNOWN;

src/brpc/span.h

@@ -125,6 +125,7 @@ friend class SpanContainer;
     void set_log_id(uint64_t cid) { _log_id = cid; }
     void set_base_cid(bthread_id_t id) { _base_cid = id; }
     void set_ending_cid(bthread_id_t id) { _ending_cid = id; }
+    void set_ending_tid(bthread_t tid) { _ending_tid = tid; }
     void set_remote_side(const butil::EndPoint& pt) { _remote_side = pt; }
     void set_protocol(ProtocolType p) { _protocol = p; }
     void set_error_code(int error_code) { _error_code = error_code; }
@@ -144,7 +145,12 @@ friend class SpanContainer;
     void set_sent_us(int64_t tm)
     { _sent_real_us = tm + _base_real_us; }
 
-    bool is_active() const { return _ending_cid == INVALID_BTHREAD_ID; }
+    bool is_active() const {
+        if (_type == SPAN_TYPE_BTHREAD) {
+            return _ending_tid == INVALID_BTHREAD;
+        }
+        return _ending_cid == INVALID_BTHREAD_ID;
+    }
 
     std::weak_ptr<Span> local_parent() const { return _local_parent; }
     static std::shared_ptr<Span> tls_parent() {
@@ -161,6 +167,7 @@ friend class SpanContainer;
     uint64_t log_id() const { return _log_id; }
     bthread_id_t base_cid() const { return _base_cid; }
     bthread_id_t ending_cid() const { return _ending_cid; }
+    bthread_t ending_tid() const { return _ending_tid; }
     const butil::EndPoint& remote_side() const { return _remote_side; }
     SpanType type() const { return _type; }
     ProtocolType protocol() const { return _protocol; }
@@ -215,6 +222,7 @@ friend class SpanContainer;
     uint64_t _log_id;
     bthread_id_t _base_cid;
     bthread_id_t _ending_cid;
+    bthread_t _ending_tid;  // Used for bthread span to store the ending bthread tid
     butil::EndPoint _remote_side;
     SpanType _type;
     bool _async;
@@ -252,7 +260,7 @@ friend class SpanContainer;
 
 class SpanContainer : public bvar::Collected {
 public:
-    explicit SpanContainer(std::shared_ptr<Span> span) : _span(span) {}
+    explicit SpanContainer(const std::shared_ptr<Span>& span) : _span(span) {}
     ~SpanContainer() {}
 
     // Implementations of bvar::Collected

src/bthread/bthread.cpp

@@ -659,6 +659,21 @@ uint64_t bthread_cpu_clock_ns(void) {
     return 0;
 }
 
+int bthread_set_span_funcs(bthread_create_span_fn create_fn,
+                            bthread_destroy_span_fn destroy_fn,
+                            bthread_end_span_fn end_fn) {
+    if ((create_fn && destroy_fn && end_fn) ||
+        (!create_fn && !destroy_fn && !end_fn)) {
+        bthread::g_create_bthread_span = create_fn;
+        bthread::g_rpcz_parent_span_dtor = destroy_fn;
+        bthread::g_end_bthread_span = end_fn;
+        return 0;
+    }
+
+    errno = EINVAL;
+    return -1;
+}
+
 }  // extern "C"
 
 void bthread_attr_set_name(bthread_attr_t* attr, const char* name) {