Fix span lifecycle with smart pointers to prevent use-after-free in async RPC callbacks (#3140)

* Fix span lifecycle with smart pointers to prevent use-after-free in async RPC callbacks (#3068)

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

---------

Co-authored-by: lhh <lhh>

Author: lh2debug

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

@@ -185,16 +185,43 @@ static void PrintElapse(std::ostream& os, int64_t cur_time,
 
 static void PrintAnnotations(
     std::ostream& os, int64_t cur_time, int64_t* last_time,
-    SpanInfoExtractor** extractors, int num_extr) {
+    SpanInfoExtractor** extractors, int num_extr, const RpczSpan* span) {
     int64_t anno_time;
     std::string a;
+    const char* span_type_str = "Span";
+    if (span) {
+        switch (span->type()) {
+        case SPAN_TYPE_SERVER:
+            span_type_str = "ServerSpan";
+            break;
+        case SPAN_TYPE_CLIENT:
+            span_type_str = "ClientSpan";
+            break;
+        case SPAN_TYPE_BTHREAD:
+            span_type_str = "BthreadSpan";
+            break;
+        }
+    }
+
     // TODO: Going through all extractors is not strictly correct because 
     // later extractors may have earlier annotations.
     for (int i = 0; i < num_extr; ++i) {
         while (extractors[i]->PopAnnotation(cur_time, &anno_time, &a)) {
             PrintRealTime(os, anno_time);
             PrintElapse(os, anno_time, last_time);
-            os << ' ' << WebEscape(a);
+            os << ' ';
+            if (span) {
+                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') {
                 os << '\n';
             }
@@ -204,12 +231,12 @@ static void PrintAnnotations(
 
 static bool PrintAnnotationsAndRealTimeSpan(
     std::ostream& os, int64_t cur_time, int64_t* last_time,
-    SpanInfoExtractor** extr, int num_extr) {
+    SpanInfoExtractor** extr, int num_extr, const RpczSpan* span) {
     if (cur_time == 0) {
         // the field was not set.
         return false;
     }
-    PrintAnnotations(os, cur_time, last_time, extr, num_extr);
+    PrintAnnotations(os, cur_time, last_time, extr, num_extr, span);
     PrintRealTime(os, cur_time);
     PrintElapse(os, cur_time, last_time);
     return true;
@@ -239,9 +266,10 @@ static void PrintClientSpan(
         extr[num_extr++] = server_extr;
     }
     extr[num_extr++] = &client_extr;
-    // start_send_us is always set for client spans.
-    CHECK(PrintAnnotationsAndRealTimeSpan(os, span.start_send_real_us(),
-                                          last_time, extr, num_extr));
+    if (!PrintAnnotationsAndRealTimeSpan(os, span.start_send_real_us(),
+                                           last_time, extr, num_extr, &span)) {
+        os << " start_send_real_us:not-set";
+    }
     const Protocol* protocol = FindProtocol(span.protocol());
     const char* protocol_name = (protocol ? protocol->name : "Unknown");
     const butil::EndPoint remote_side(butil::int2ip(span.remote_ip()), span.remote_port());
@@ -271,12 +299,12 @@ static void PrintClientSpan(
     os << std::endl;
 
     if (PrintAnnotationsAndRealTimeSpan(os, span.sent_real_us(),
-                                        last_time, extr, num_extr)) {
-        os << " Requested(" << span.request_size() << ") [1]" << std::endl;
+                                        last_time, extr, num_extr, &span)) {
+        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)) {
-        os << " Received response(" << span.response_size() << ")";
+                                        last_time, extr, num_extr, &span)) {
+        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) {
@@ -289,18 +317,18 @@ static void PrintClientSpan(
     }
 
     if (PrintAnnotationsAndRealTimeSpan(os, span.start_parse_real_us(),
-                                        last_time, extr, num_extr)) {
-        os << " Processing the response in a new bthread" << std::endl;
+                                        last_time, extr, num_extr, &span)) {
+        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)) {
-        os << (span.async() ? " Enter user's done" : " Back to user's callsite") << std::endl;
+            last_time, extr, num_extr, &span)) {
+        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(),
-                     last_time, extr, num_extr);
+                     last_time, extr, num_extr, &span);
 }
 
 static void PrintClientSpan(std::ostream& os,const RpczSpan& span,
@@ -318,7 +346,15 @@ static void PrintBthreadSpan(std::ostream& os, const RpczSpan& span, int64_t* la
         extr[num_extr++] = server_extr;
     }
     extr[num_extr++] = &client_extr;
-    PrintAnnotations(os, std::numeric_limits<int64_t>::max(), last_time, extr, num_extr);
+
+    // Print span id for bthread span context identification
+    os << " [Bthread SPAN#" << Hex(span.span_id());
+    if (span.parent_span_id() != 0) {
+        os << " parent#" << Hex(span.parent_span_id());
+    }
+    os << "] ";
+
+    PrintAnnotations(os, std::numeric_limits<int64_t>::max(), last_time, extr, num_extr, &span);
 }
 
 static void PrintServerSpan(std::ostream& os, const RpczSpan& span,
@@ -348,16 +384,16 @@ static void PrintServerSpan(std::ostream& os, const RpczSpan& span,
     os << std::endl;
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.start_parse_real_us(),
-            &last_time, extr, ARRAY_SIZE(extr))) {
-        os << " Processing the request in a new bthread" << std::endl;
+            &last_time, extr, ARRAY_SIZE(extr), &span)) {
+        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))) {
+            &last_time, extr, ARRAY_SIZE(extr), &span)) {
         entered_user_method = true;
-        os << " 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();
@@ -372,22 +408,22 @@ static void PrintServerSpan(std::ostream& os, const RpczSpan& span,
 
     if (PrintAnnotationsAndRealTimeSpan(
             os, span.start_send_real_us(),
-            &last_time, extr, ARRAY_SIZE(extr))) {
+            &last_time, extr, ARRAY_SIZE(extr), &span)) {
         if (entered_user_method) {
-            os << " 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 << " 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))) {
-        os << " Responded(" << span.response_size() << ')' << std::endl;
+            &last_time, extr, ARRAY_SIZE(extr), &span)) {
+        os << " [Server SPAN#" << Hex(span.span_id()) << "] Responded(" << span.response_size() << ')' << std::endl;
     }
 
     PrintAnnotations(os, std::numeric_limits<int64_t>::max(),
-                     &last_time, extr, ARRAY_SIZE(extr));
+                      &last_time, extr, ARRAY_SIZE(extr), &span);
 }
 
 class RpczSpanFilter : public SpanFilter {

src/brpc/channel.cpp

@@ -38,6 +38,7 @@
 #include "brpc/rdma/rdma_helper.h"
 #include "brpc/policy/esp_authenticator.h"
 #include "brpc/transport_factory.h"
+#include "brpc/details/controller_private_accessor.h"
 
 namespace brpc {
 
@@ -502,7 +503,7 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
     }
     cntl->set_used_by_rpc();
 
-    if (cntl->_sender == NULL && IsTraceable(Span::tls_parent())) {
+    if (cntl->_sender == NULL && IsTraceable(Span::tls_parent().get())) {
         const int64_t start_send_us = butil::cpuwide_time_us();
         std::string method_name;
         if (_get_method_name) {
@@ -513,13 +514,16 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
             const static std::string NULL_METHOD_STR = "null-method";
             method_name = NULL_METHOD_STR;
         }
-        Span* span = Span::CreateClientSpan(
+        std::shared_ptr<Span> span = Span::CreateClientSpan(
             method_name, start_send_real_us - start_send_us);
-        span->set_log_id(cntl->log_id());
-        span->set_base_cid(correlation_id);
-        span->set_protocol(_options.protocol);
-        span->set_start_send_us(start_send_us);
-        cntl->_span = span;
+        if (span) {
+            ControllerPrivateAccessor accessor(cntl);
+            span->set_log_id(cntl->log_id());
+            span->set_base_cid(correlation_id);
+            span->set_protocol(_options.protocol);
+            span->set_start_send_us(start_send_us);
+            accessor.set_span(span);
+        }
     }
     // Override some options if they haven't been set by Controller
     if (cntl->timeout_ms() == UNSET_MAGIC_NUM) {
@@ -620,9 +624,7 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
         // be woken up by callback when RPC finishes (succeeds or still
         // fails after retry)
         Join(correlation_id);
-        if (cntl->_span) {
-            cntl->SubmitSpan();
-        }
+        cntl->SubmitSpan();
         cntl->OnRPCEnd(butil::gettimeofday_us());
     }
 }