Skip to content

Refactor streaming rpc #3422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenBright wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Refactor streaming rpc #3422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/brpc/controller.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1478,15 +1478,15 @@ void Controller::HandleStreamConnection(Socket *host_socket) {
return;
}
size_t stream_num = _request_streams.size();
std::vector<SocketUniquePtr> ptrs(stream_num);
std::vector<StreamUniquePtr> ptrs(stream_num);
if (!FailedInline()) {
if (_remote_stream_settings == NULL) {
if (!FailedInline()) {
SetFailed(EREQUEST, "The server didn't accept the stream");
}
} else {
for (size_t i = 0; i < stream_num; ++i) {
if (Socket::Address(_request_streams[i], &ptrs[i]) != 0) {
if (Stream::Address(_request_streams[i], &ptrs[i]) != 0) {
if (!FailedInline()) {
SetFailed(EREQUEST, "Request stream=%" PRIu64 " was closed before responded",
_request_streams[i]);
Expand All @@ -1509,14 +1509,14 @@ void Controller::HandleStreamConnection(Socket *host_socket) {
}
return;
}
Stream* s = (Stream*)ptrs[0]->conn();
Stream* s = ptrs[0].get();
s->SetConnected(_remote_stream_settings);
if (stream_num > 1) {
auto extra_stream_ids = std::move(*_remote_stream_settings->mutable_extra_stream_ids());
_remote_stream_settings->clear_extra_stream_ids();
for (size_t i = 1; i < stream_num; ++i) {
if(!ptrs[i]) continue;
Stream* extra_stream = (Stream *) ptrs[i]->conn();
Stream* extra_stream = ptrs[i].get();
_remote_stream_settings->set_stream_id(extra_stream_ids[i - 1]);
extra_stream->SetHostSocket(host_socket);
extra_stream->SetConnected(_remote_stream_settings);
Expand Down
20 changes: 10 additions & 10 deletions src/brpc/policy/baidu_rpc_protocol.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,11 +353,11 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
meta.set_attachment_size(attached_size);
}
StreamId response_stream_id = INVALID_STREAM_ID;
SocketUniquePtr stream_ptr;
StreamUniquePtr stream_ptr;
if (!response_stream_ids.empty()) {
response_stream_id = response_stream_ids[0];
if (Socket::Address(response_stream_id, &stream_ptr) == 0) {
Stream* s = (Stream *) stream_ptr->conn();
if (Stream::Address(response_stream_id, &stream_ptr) == 0) {
Stream* s = stream_ptr.get();
StreamSettings *stream_settings = meta.mutable_stream_settings();
s->FillSettings(stream_settings);
s->SetHostSocket(sock);
Expand Down Expand Up @@ -419,13 +419,13 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
// written user data would follower the RPC response.
// Reuse stream_ptr to avoid address first stream id again
if (stream_ptr) {
((Stream*)stream_ptr->conn())->SetConnected();
stream_ptr->SetConnected();
}
for (size_t i = 1; i < response_stream_ids.size(); ++i) {
StreamId extra_stream_id = response_stream_ids[i];
SocketUniquePtr extra_stream_ptr;
if (Socket::Address(extra_stream_id, &extra_stream_ptr) == 0) {
Stream* extra_stream = (Stream *) extra_stream_ptr->conn();
StreamUniquePtr extra_stream_ptr;
if (Stream::Address(extra_stream_id, &extra_stream_ptr) == 0) {
Stream* extra_stream = extra_stream_ptr.get();
extra_stream->SetHostSocket(sock);
extra_stream->SetConnected();
} else {
Expand Down Expand Up @@ -1086,12 +1086,12 @@ void PackRpcRequest(butil::IOBuf* req_buf,
if (!request_stream_ids.empty()) {
StreamSettings* stream_settings = meta.mutable_stream_settings();
StreamId request_stream_id = request_stream_ids[0];
SocketUniquePtr ptr;
if (Socket::Address(request_stream_id, &ptr) != 0) {
StreamUniquePtr ptr;
if (Stream::Address(request_stream_id, &ptr) != 0) {
return cntl->SetFailed(EREQUEST, "Stream=%" PRIu64 " was closed",
request_stream_id);
}
Stream* s = (Stream*) ptr->conn();
Stream* s = ptr.get();
s->FillSettings(stream_settings);
for (size_t i = 1; i < request_stream_ids.size(); ++i) {
stream_settings->mutable_extra_stream_ids()->Add(request_stream_ids[i]);
Expand Down
25 changes: 8 additions & 17 deletions src/brpc/policy/streaming_rpc_protocol.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,11 +102,11 @@ ParseResult ParseStreamingMessage(butil::IOBuf* source,
LOG(WARNING) << "Fail to Parse StreamFrameMeta from " << *socket;
break;
}
SocketUniquePtr ptr;
if (Socket::Address((SocketId)fm.stream_id(), &ptr) != 0) {
RPC_VLOG_IF(fm.frame_type() != FRAME_TYPE_RST
&& fm.frame_type() != FRAME_TYPE_CLOSE
&& fm.frame_type() != FRAME_TYPE_FEEDBACK)
StreamUniquePtr sptr;
if (Stream::Address((StreamId)fm.stream_id(), &sptr) != 0) {
RPC_VLOG_IF(fm.frame_type() != FRAME_TYPE_RST &&
fm.frame_type() != FRAME_TYPE_CLOSE &&
fm.frame_type() != FRAME_TYPE_FEEDBACK)
<< "Fail to find stream=" << fm.stream_id();
// It's normal that the stream is closed before receiving feedback frames from peer.
// In this case, RST frame should not be sent to peer, otherwise on-fly data can be lost.
Expand All @@ -116,16 +116,7 @@ ParseResult ParseStreamingMessage(butil::IOBuf* source,
break;
}
meta_buf.clear(); // to reduce memory resident
// ptr->conn() returns the connection-level context attached to the
// socket. It may be NULL when the socket was found by ID but has no
// Stream object associated (e.g. during protocol probing or fuzz
// testing). Calling OnReceived on a null pointer would crash.
Stream* stream_conn = (Stream*)ptr->conn();
if (stream_conn == NULL) {
LOG(FATAL) << "No stream object found";
break;
}
stream_conn->OnReceived(fm, &payload, socket);
sptr->OnReceived(fm, &payload, socket);
} while (0);

// Hack input messenger
Expand All @@ -136,7 +127,7 @@ void ProcessStreamingMessage(InputMessageBase* /*msg*/) {
CHECK(false) << "Should never be called";
}

void SendStreamRst(Socket *sock, int64_t remote_stream_id) {
void SendStreamRst(Socket* sock, int64_t remote_stream_id) {
CHECK(sock != NULL);
StreamFrameMeta fm;
fm.set_stream_id(remote_stream_id);
Expand All @@ -148,7 +139,7 @@ void SendStreamRst(Socket *sock, int64_t remote_stream_id) {
sock->Write(&out, &wopt);
}

void SendStreamClose(Socket *sock, int64_t remote_stream_id,
void SendStreamClose(Socket* sock, int64_t remote_stream_id,
int64_t source_stream_id) {
CHECK(sock != NULL);
StreamFrameMeta fm;
Expand Down
2 changes: 1 addition & 1 deletion src/brpc/socket.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,7 @@ friend class TransportFactory;
// NOTE: User cannot create Socket from constructor. Use Create()
// instead. It's public just because of requirement of ResourcePool.
explicit Socket(Forbidden);
~Socket() override;
~Socket();

// Write `msg' into this Socket and clear it. The `msg' should be an
// intact request or response. To prevent messages from interleaving
Expand Down
Loading
Loading