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

@@ -73,10 +73,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 +95,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();
     }
 }
 

src/brpc/span.h

@@ -252,7 +252,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) {

src/bthread/key.cpp

@@ -35,12 +35,6 @@
 
 namespace bthread {
 
-void* (*g_create_bthread_span)() = NULL;
-
-void (*g_rpcz_parent_span_dtor)(void*) = NULL;
-
-void (*g_end_bthread_span)() = NULL;
-
 DEFINE_uint32(key_table_list_size, 4000,
               "The maximum length of the KeyTableList. Once this value is "
               "exceeded, a portion of the KeyTables will be moved to the "
@@ -252,11 +246,6 @@ class BAIDU_CACHELINE_ALIGNMENT KeyTableList {
         if (g) {
             g->current_task()->local_storage.keytable = old_kt;
         }
-
-        if (tls_bls.rpcz_parent_span && g_rpcz_parent_span_dtor) {
-            g_rpcz_parent_span_dtor(tls_bls.rpcz_parent_span);
-            tls_bls.rpcz_parent_span = NULL;
-        }
     }
 
     void append(KeyTable* keytable) {
@@ -417,11 +406,6 @@ static void cleanup_pthread(void* arg) {
         delete kt;
         // After deletion: tls may be set during deletion.
         tls_bls.keytable = NULL;
-
-        if (tls_bls.rpcz_parent_span && g_rpcz_parent_span_dtor) {
-            g_rpcz_parent_span_dtor(tls_bls.rpcz_parent_span);
-            tls_bls.rpcz_parent_span = NULL;
-        }
     }
 }
 
@@ -500,11 +484,6 @@ int bthread_keytable_pool_destroy(bthread_keytable_pool_t* pool) {
     if (g) {
         g->current_task()->local_storage.keytable = old_kt;
     }
-
-    if (bthread::tls_bls.rpcz_parent_span && bthread::g_rpcz_parent_span_dtor) {
-        bthread::g_rpcz_parent_span_dtor(bthread::tls_bls.rpcz_parent_span);
-        bthread::tls_bls.rpcz_parent_span = NULL;
-    }
     // TODO: return_keytable may race with this function, we don't destroy
     // the mutex right now.
     // pthread_mutex_destroy(&pool->mutex);
@@ -694,19 +673,4 @@ void* bthread_get_assigned_data() {
     return bthread::tls_bls.assigned_data;
 }
 
-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"