* Abstract stream-level data as StreamUserData

* Apply StreamUserData to h2 and rtmp

src/brpc/controller.cpp

@@ -261,13 +261,13 @@ Controller::Call::Call(Controller::Call* rhs)
     , peer_id(rhs->peer_id)
     , begin_time_us(rhs->begin_time_us)
     , sending_sock(rhs->sending_sock.release())
-    , stream_creator(rhs->stream_creator) {
+    , stream_user_data(rhs->stream_user_data) {
     // NOTE: fields in rhs should be reset because RPC could fail before
     // setting all the fields to next call and _current_call.OnComplete
     // will behave incorrectly.
     rhs->need_feedback = false;
     rhs->peer_id = (SocketId)-1;
-    rhs->stream_creator = NULL;
+    rhs->stream_user_data = NULL;
 }
 
 Controller::Call::~Call() {
@@ -282,7 +282,7 @@ void Controller::Call::Reset() {
     peer_id = (SocketId)-1;
     begin_time_us = 0;
     sending_sock.reset(NULL);
-    stream_creator = NULL;
+    stream_user_data = NULL;
 }
 
 void Controller::set_timeout_ms(int64_t timeout_ms) {
@@ -769,13 +769,13 @@ void Controller::Call::OnComplete(
         c->_lb->Feedback(info);
     }
 
-    if (stream_creator) {
-        stream_creator->OnDestroyingStream(
-            sending_sock, c, error_code, end_of_rpc);
+    if (c->stream_creator()) {
+        c->stream_creator()->OnDestroyingStream(
+            sending_sock, c, error_code, end_of_rpc, stream_user_data);
     }
-
     // Release the `Socket' we used to send/receive data
     sending_sock.reset(NULL);
+    stream_user_data = NULL;
 }
 
 void Controller::EndRPC(const CompletionInfo& info) {

src/brpc/controller.h

@@ -248,16 +248,11 @@ public:
     void reset_rpc_dump_meta(RpcDumpMeta* meta);
     const RpcDumpMeta* rpc_dump_meta() { return _rpc_dump_meta; }
 
-    // Attach a global StreamCreator to this RPC. Notice that controller never
-    // deletes the StreamCreator, you can do the deletion inside OnDestroyingStream.
+    // Attach a StreamCreator to this RPC. Notice that controller never deletes
+    // the StreamCreator, you can do the deletion inside OnDestroyingStream.
     void set_stream_creator(StreamCreator* sc) { _stream_creator = sc; }
     StreamCreator* stream_creator() const { return _stream_creator; }
 
-    // Attach a StreamCreator to current call. In some protocols(such as h2), each
-    // call has a specific StreamCreator, use this function to set.
-    void set_current_stream_creator(StreamCreator* sc) { _current_call.stream_creator = sc; }
-    StreamCreator* current_stream_creator() const { return _current_call.stream_creator; }
-
     // Make the RPC end when the HTTP response has complete headers and let
     // user read the remaining body by using ReadProgressiveAttachmentBy().
     void response_will_be_read_progressively() { add_flag(FLAGS_READ_PROGRESSIVELY); }
@@ -586,7 +581,7 @@ private:
         // CONNECTION_TYPE_SINGLE. Otherwise, it may be a temporary
         // socket fetched from socket pool
         SocketUniquePtr sending_sock;
-        StreamCreator* stream_creator;
+        StreamUserData* stream_user_data;
     };
 
     void HandleStreamConnection(Socket *host_socket);

src/brpc/details/controller_private_accessor.h

@@ -66,6 +66,11 @@ public:
         CHECK(_cntl->_current_call.sending_sock == NULL);
         _cntl->_current_call.sending_sock.reset(ptr.release());
     }
+
+    ControllerPrivateAccessor &set_stream_user_data(StreamUserData* ptr){
+        _cntl->_current_call.stream_user_data = ptr;
+        return *this;
+    }
     
     ControllerPrivateAccessor &set_security_mode(bool security_mode) {
         _cntl->set_flag(Controller::FLAGS_SECURITY_MODE, security_mode);

src/brpc/policy/http2_rpc_protocol.cpp

@@ -1323,39 +1323,27 @@ void H2UnsentRequest::Destroy() {
     free(this);
 }
 
-void H2UnsentRequest::OnCreatingStream(SocketUniquePtr*, Controller*) {
-    CHECK(false) << "H2UnsentRequest::OnCreatingStream should not be called";
-}
+struct RemoveRefOnQuit {
+    RemoveRefOnQuit(H2UnsentRequest* msg) : _msg(msg) {}
+    ~RemoveRefOnQuit() { _msg->RemoveRefManually(); }
+private:
+    DISALLOW_COPY_AND_ASSIGN(RemoveRefOnQuit);
+    H2UnsentRequest* _msg;
+};
 
-void H2UnsentRequest::OnDestroyingStream(
-    SocketUniquePtr& sending_sock, Controller* cntl, int error_code, bool end_of_rpc) {
-    if (!end_of_rpc) {
-        return;
-    }
-    // If cntl->ErrorCode == 0, then it is a normal response and stream has
-    // already been removed in EndRemoteStream.
+void H2UnsentRequest::OnDestroy(SocketUniquePtr& sending_sock, Controller* cntl) {
+    RemoveRefOnQuit deref_self(this);
     if (sending_sock != NULL && cntl->ErrorCode() != 0) {
-        CHECK_EQ(_cntl, cntl);
-        _mutex.lock();
+        CHECK_EQ(cntl, _cntl);
+        std::unique_lock<butil::Mutex> mu(_mutex);
         _cntl = NULL;
         if (_stream_id != 0) {
-            H2Context* ctx =
-                static_cast<H2Context*>(sending_sock->parsing_context());
+            H2Context* ctx = static_cast<H2Context*>(sending_sock->parsing_context());
             ctx->AddAbandonedStream(_stream_id);
         }
-        _mutex.unlock();
     }
-    RemoveRefManually();
 }
 
-struct RemoveRefOnQuit {
-    RemoveRefOnQuit(H2UnsentRequest* msg) : _msg(msg) {}
-    ~RemoveRefOnQuit() { _msg->RemoveRefManually(); }
-private:
-    DISALLOW_COPY_AND_ASSIGN(RemoveRefOnQuit);
-    H2UnsentRequest* _msg;
-};
-
 // bvar::Adder<int64_t> g_append_request_time;
 // bvar::PerSecond<bvar::Adder<int64_t> > g_append_request_time_per_second(
 //     "h2_append_request_second",     &g_append_request_time);
@@ -1649,11 +1637,11 @@ void PackH2Request(butil::IOBuf*,
 
     // Serialize http2 request
     H2UnsentRequest* h2_req = H2UnsentRequest::New(cntl, correlation_id);
-    h2_req->AddRefManually();   // add for OnDestroyingStream
-    if (cntl->current_stream_creator()) {
-        dynamic_cast<H2UnsentRequest*>(cntl->current_stream_creator())->RemoveRefManually();
+    if (!h2_req) {
+        return cntl->SetFailed(ENOMEM, "Fail to create H2UnsentRequest");
     }
-    cntl->set_current_stream_creator(h2_req);
+    h2_req->AddRefManually();   // add ref for H2UnsentRequest::OnDestroy
+    accessor.set_stream_user_data(h2_req);
     *user_message = h2_req;
     
     if (FLAGS_http_verbose) {
@@ -1706,10 +1694,15 @@ void H2GlobalStreamCreator::OnCreatingStream(
     }
 }
 
-void H2GlobalStreamCreator::OnDestroyingStream(
-    SocketUniquePtr& sending_sock, Controller* cntl, int error_code, bool end_of_rpc) {
-    // If any error happens during the time of sending rpc, this function
-    // would be called. Currently just do nothing.
+void H2GlobalStreamCreator::OnDestroyingStream(SocketUniquePtr& sending_sock,
+                                               Controller* cntl,
+                                               int error_code,
+                                               bool end_of_rpc,
+                                               StreamUserData* stream_user_data) {
+    // [stream_user_data == NULL] means that RPC has failed before it is created.
+    if (stream_user_data) {
+        stream_user_data->OnDestroy(sending_sock, cntl);
+    }
 }
 
 StreamCreator* get_h2_global_stream_creator() {

src/brpc/policy/http2_rpc_protocol.h

@@ -88,7 +88,7 @@ inline Http2Bvars* get_http2_bvars() {
     return butil::get_leaky_singleton<Http2Bvars>();
 }
 
-class H2UnsentRequest : public SocketMessage, public StreamCreator {
+class H2UnsentRequest : public SocketMessage, public StreamUserData {
 public:
     static H2UnsentRequest* New(Controller* c, uint64_t correlation_id);
     void Describe(butil::IOBuf*) const;
@@ -103,13 +103,12 @@ public:
         }
     }
 
-    // @StreamCreator
-    void OnCreatingStream(SocketUniquePtr* inout, Controller* cntl) override;
-    void OnDestroyingStream(SocketUniquePtr& sending_sock, Controller* cntl, int error_code, bool end_of_rpc) override;
-    
     // @SocketMessage
     butil::Status AppendAndDestroySelf(butil::IOBuf* out, Socket*);
     size_t EstimatedByteSize();
+
+    // @StreamUserData
+    void OnDestroy(SocketUniquePtr& sending_sock, Controller* cntl);
     
 private:
     std::string& push(const std::string& name)
@@ -235,7 +234,11 @@ void PackH2Request(butil::IOBuf* buf,
 class H2GlobalStreamCreator : public StreamCreator {
 protected:
     void OnCreatingStream(SocketUniquePtr* inout, Controller* cntl) override;
-    void OnDestroyingStream(SocketUniquePtr& sending_sock, Controller* cntl, int error_code, bool end_of_rpc) override;
+    void OnDestroyingStream(SocketUniquePtr& sending_sock,
+                            Controller* cntl,
+                            int error_code,
+                            bool end_of_rpc,
+                            StreamUserData* stream_user_data) override;
 private:
     butil::Mutex _mutex;
 };

src/brpc/rtmp.cpp

@@ -1748,7 +1748,8 @@ void RtmpClientStream::OnFailedToCreateStream() {
 void RtmpClientStream::OnDestroyingStream(SocketUniquePtr& sending_sock,
                                           Controller* cntl,
                                           int /*error_code*/,
-                                          bool end_of_rpc) {
+                                          bool end_of_rpc,
+                                          StreamUserData* /*stream_data*/) {
     if (!end_of_rpc) {
         if (sending_sock) {
             if (_from_socketmap) {
@@ -2137,7 +2138,6 @@ void RtmpClientStream::Init(const RtmpClient* client,
     // In RTMP, stream_creator and current stream_creator is always
     // this RtmpClientStream.
     done->cntl.set_stream_creator(this);
-    done->cntl.set_current_stream_creator(this);
     done->cntl.set_connection_type(_options.share_connection ?
                                    CONNECTION_TYPE_SINGLE :
                                    CONNECTION_TYPE_SHORT);

src/brpc/rtmp.h

@@ -822,7 +822,8 @@ friend class RtmpRetryingClientStream;
 
     // @StreamCreator
     void OnCreatingStream(SocketUniquePtr* inout, Controller* cntl) override;
-    void OnDestroyingStream(SocketUniquePtr&, Controller*, int error_code, bool end_of_rpc) override;
+    void OnDestroyingStream(SocketUniquePtr&, Controller*, int error_code,
+                            bool end_of_rpc, StreamUserData*) override;
 
     void OnFailedToCreateStream();
     

src/brpc/stream_creator.h

@@ -19,10 +19,17 @@
 
 #include "brpc/socket_id.h"
 
-
 namespace brpc {
 class Controller;
 
+// The stream user data on a specific Call
+class StreamUserData {
+public:
+    virtual ~StreamUserData() {}
+
+    virtual void OnDestroy(SocketUniquePtr& sending_sock, Controller* cntl) = 0;
+};
+
 // Abstract creation of "user-level connection" over a RPC-like process.
 // Lifetime of this object should be guaranteed by user during the RPC,
 // generally this object is created before RPC and destroyed after RPC.
@@ -52,10 +59,12 @@ public:
     //   cntl: contexts of the RPC
     //   error_code: Use this instead of cntl->ErrorCode()
     //   end_of_rpc: true if the RPC is about to destroyed.
+    //   stream_user_data: the corresponding user data of this very stream
     virtual void OnDestroyingStream(SocketUniquePtr& sending_sock,
                                     Controller* cntl,
                                     int error_code,
-                                    bool end_of_rpc) = 0;
+                                    bool end_of_rpc,
+                                    StreamUserData* stream_user_data) = 0;
 };
 
 } // namespace brpc