Add BRPC_SPAN_ENABLE_SHARED_PTR_API flag for backward-compatible Span lifecycle management (#3068)

Author: lhh

src/brpc/channel.cpp

@@ -510,7 +510,11 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
             span->set_base_cid(correlation_id);
             span->set_protocol(_options.protocol);
             span->set_start_send_us(start_send_us);
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
             accessor.set_span(span);
+#else
+            accessor.set_span(span.get());
+#endif
         }
     }
     // Override some options if they haven't been set by Controller

src/brpc/controller.cpp

@@ -1727,4 +1727,31 @@ void Controller::DoPrintLogPrefix(std::ostream& os) const {
     }
 }
 
+
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
+ControllerPrivateAccessor& ControllerPrivateAccessor::set_span(
+    std::shared_ptr<Span> span) {
+    _cntl->_span = span;
+    return *this;
+}
+#else
+ControllerPrivateAccessor& ControllerPrivateAccessor::set_span(Span* span) {
+    if (span) {
+        _cntl->_span = span->shared_from_this();
+    } else {
+        _cntl->_span.reset();
+    }
+    return *this;
+}
+#endif
+
+SpanPtr ControllerPrivateAccessor::span() const {
+    std::shared_ptr<Span> span_internal = _cntl->_span.lock();
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
+    return span_internal;
+#else
+    return span_internal.get();
+#endif
+}
+
 } // namespace brpc

src/brpc/details/controller_private_accessor.h

@@ -30,9 +30,20 @@ class Message;
 }
 }
 
-
 namespace brpc {
 
+class Span;
+
+#ifndef BRPC_SPAN_ENABLE_SHARED_PTR_API
+#define BRPC_SPAN_ENABLE_SHARED_PTR_API 0
+#endif
+
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
+using SpanPtr = std::shared_ptr<Span>;
+#else
+using SpanPtr = Span*;
+#endif
+
 class AuthContext;
 
 // A wrapper to access some private methods/fields of `Controller'
@@ -90,17 +101,18 @@ class ControllerPrivateAccessor {
         return *this;
     }
 
-    ControllerPrivateAccessor &set_span(std::shared_ptr<Span> span) {
-        _cntl->_span = span;
-        return *this;
-    }
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
+    ControllerPrivateAccessor &set_span(std::shared_ptr<Span> span);
+#else
+    ControllerPrivateAccessor &set_span(Span* span);
+#endif
 
     ControllerPrivateAccessor &set_request_protocol(ProtocolType protocol) {
         _cntl->_request_protocol = protocol;
         return *this;
     }
 
-    std::shared_ptr<Span> span() const { return _cntl->_span.lock(); }
+    SpanPtr 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/baidu_rpc_protocol.cpp

@@ -647,7 +647,11 @@ void ProcessRpcRequest(InputMessageBase* msg_base) {
         span = Span::CreateServerSpan(
             request_meta.trace_id(), request_meta.span_id(),
             request_meta.parent_span_id(), msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_log_id(request_meta.log_id());
         span->set_remote_side(cntl->remote_side());
         span->set_protocol(PROTOCOL_BAIDU_STD);

src/brpc/policy/http_rpc_protocol.cpp

@@ -1506,7 +1506,11 @@ void ProcessHttpRequest(InputMessageBase *msg) {
         }
         span = Span::CreateServerSpan(
             path, trace_id, span_id, parent_span_id, msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_log_id(cntl->log_id());
         span->set_remote_side(user_addr);
         span->set_received_us(msg->received_us());

src/brpc/policy/hulu_pbrpc_protocol.cpp

@@ -416,7 +416,11 @@ void ProcessHuluRequest(InputMessageBase* msg_base) {
         span = Span::CreateServerSpan(
             meta.trace_id(), meta.span_id(), meta.parent_span_id(),
             msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_log_id(meta.log_id());
         span->set_remote_side(cntl->remote_side());
         span->set_protocol(PROTOCOL_HULU_PBRPC);

src/brpc/policy/nshead_protocol.cpp

@@ -300,7 +300,11 @@ void ProcessNsheadRequest(InputMessageBase* msg_base) {
     std::shared_ptr<Span> span;
     if (IsTraceable(false)) {
         span = Span::CreateServerSpan(0, 0, 0, msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_log_id(req_head->log_id);
         span->set_remote_side(cntl->remote_side());
         span->set_protocol(PROTOCOL_NSHEAD);

src/brpc/policy/sofa_pbrpc_protocol.cpp

@@ -376,7 +376,11 @@ void ProcessSofaRequest(InputMessageBase* msg_base) {
         span = Span::CreateServerSpan(
             0/*meta.trace_id()*/, 0/*meta.span_id()*/,
             0/*meta.parent_span_id()*/, msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_remote_side(cntl->remote_side());
         span->set_protocol(PROTOCOL_SOFA_PBRPC);
         span->set_received_us(msg->received_us());

src/brpc/policy/thrift_protocol.cpp

@@ -518,7 +518,11 @@ void ProcessThriftRequest(InputMessageBase* msg_base) {
     std::shared_ptr<Span> span;
     if (IsTraceable(false)) {
         span = Span::CreateServerSpan(0, 0, 0, msg->base_real_us());
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
         accessor.set_span(span);
+#else
+        accessor.set_span(span.get());
+#endif
         span->set_log_id(seq_id);
         span->set_remote_side(cntl->remote_side());
         span->set_protocol(PROTOCOL_THRIFT);

src/brpc/span.h

@@ -42,8 +42,43 @@ extern __thread bthread::LocalStorage tls_bls;
 
 namespace brpc {
 
+// ============================================================================
+// Span Lifecycle Management API Compatibility Layer
+// ============================================================================
+//
+// COMPATIBILITY NOTE:
+// brpc uses std::shared_ptr<Span> internally
+// to prevent use-after-free bugs in async RPC callbacks.
+//
+// For backward compatibility with existing protocol extensions, external APIs
+// can return raw pointers (Span*) by default. To enable the modern shared_ptr
+// API, compile with -DBRPC_SPAN_ENABLE_SHARED_PTR_API=1.
+//
+// MIGRATION GUIDE:
+// - Legacy mode (default): SpanPtr = Span*
+//   Users must ensure the Span outlives their usage (typically by keeping
+//   the Controller alive).
+//
+// - Modern mode (recommended): SpanPtr = std::shared_ptr<Span>
+//   Automatic lifetime management, safer for async operations.
+//
+// ============================================================================
+
+#ifndef BRPC_SPAN_ENABLE_SHARED_PTR_API
+#define BRPC_SPAN_ENABLE_SHARED_PTR_API 0  // Default: legacy mode for compatibility
+#endif
+
 class Span;
 
+#if BRPC_SPAN_ENABLE_SHARED_PTR_API
+// Modern API: Return shared_ptr for safe lifecycle management
+using SpanPtr = std::shared_ptr<Span>;
+#else
+// Legacy API: Return raw pointer for backward compatibility
+// WARNING: Users must ensure the Span outlives their usage
+using SpanPtr = Span*;
+#endif
+
 void SetTlsParentSpan(std::shared_ptr<Span> span);
 std::shared_ptr<Span> GetTlsParentSpan();
 void ClearTlsParentSpan();
@@ -252,7 +287,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