diff --git a/src/brpc/input_message_base.h b/src/brpc/input_message_base.h index b117eb99c3..0d441ee05d 100644 --- a/src/brpc/input_message_base.h +++ b/src/brpc/input_message_base.h @@ -53,6 +53,7 @@ class InputMessageBase : public Destroyable { private: friend class InputMessenger; +friend class InputMessengerProcessor; friend void* ProcessInputMessage(void*); friend class Stream; friend class Transport; diff --git a/src/brpc/input_messenger.cpp b/src/brpc/input_messenger.cpp index 26aac06378..50dc1821f7 100644 --- a/src/brpc/input_messenger.cpp +++ b/src/brpc/input_messenger.cpp @@ -74,98 +74,6 @@ DEFINE_int32(socket_tcp_user_timeout_ms, -1, DECLARE_bool(usercode_in_pthread); DECLARE_bool(usercode_in_coroutine); -DECLARE_uint64(max_body_size); - -const size_t MSG_SIZE_WINDOW = 10; // Take last so many message into stat. -const size_t MIN_ONCE_READ = 4096; -const size_t MAX_ONCE_READ = 524288; - -ParseResult InputMessenger::CutInputMessage( - Socket* m, size_t* index, bool read_eof) { - const int preferred = m->preferred_index(); - const int max_index = (int)_max_index.load(butil::memory_order_acquire); - // Try preferred handler first. The preferred_index is set on last - // selection or by client. - if (preferred >= 0 && preferred <= max_index - && _handlers[preferred].parse != nullptr) { - int cur_index = preferred; - do { - ParseResult result = - _handlers[cur_index].parse(&m->_read_buf, m, read_eof, _handlers[cur_index].arg); - if (result.is_ok() || - result.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { - m->set_preferred_index(cur_index); - *index = cur_index; - return result; - } else if (result.error() != PARSE_ERROR_TRY_OTHERS) { - // Critical error, return directly. - LOG_IF(ERROR, result.error() == PARSE_ERROR_TOO_BIG_DATA) - << "A message from " << m->remote_side() - << "(protocol=" << _handlers[cur_index].name - << ") is bigger than " << FLAGS_max_body_size - << " bytes, the connection will be closed." - " Set max_body_size to allow bigger messages"; - return result; - } - - if (m->CreatedByConnect()) { - if((ProtocolType)cur_index == PROTOCOL_BAIDU_STD && cur_index == preferred) { - // baidu_std may fall to streaming_rpc. - cur_index = (int)PROTOCOL_STREAMING_RPC; - continue; - } else if((ProtocolType)cur_index == PROTOCOL_STREAMING_RPC && cur_index == preferred) { - // streaming_rpc may fall to baidu_std. - cur_index = (int)PROTOCOL_BAIDU_STD; - continue; - } else { - // The protocol is fixed at client-side, no need to try others. - LOG(ERROR) << "Fail to parse response from " << m->remote_side() - << " by " << _handlers[preferred].name - << " at client-side"; - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } - } else { - // Try other protocols. - break; - } - } while (true); - // Clear context before trying next protocol which probably has - // an incompatible context with the current one. - if (m->parsing_context()) { - m->reset_parsing_context(nullptr); - } - m->set_preferred_index(-1); - } - for (int i = 0; i <= max_index; ++i) { - if (i == preferred || _handlers[i].parse == nullptr) { - // Don't try preferred handler(already tried) or invalid handler - continue; - } - ParseResult result = _handlers[i].parse(&m->_read_buf, m, read_eof, _handlers[i].arg); - if (result.is_ok() || - result.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { - m->set_preferred_index(i); - *index = i; - return result; - } else if (result.error() != PARSE_ERROR_TRY_OTHERS) { - // Critical error, return directly. - LOG_IF(ERROR, result.error() == PARSE_ERROR_TOO_BIG_DATA) - << "A message from " << m->remote_side() - << "(protocol=" << _handlers[i].name - << ") is bigger than " << FLAGS_max_body_size - << " bytes, the connection will be closed." - " Set max_body_size to allow bigger messages"; - return result; - } - // Clear context before trying next protocol which definitely has - // an incompatible context with the current one. - if (m->parsing_context()) { - m->reset_parsing_context(nullptr); - } - // Try other protocols. - } - return MakeParseError(PARSE_ERROR_TRY_OTHERS); -} void* ProcessInputMessage(void* void_arg) { InputMessageBase* msg = static_cast(void_arg); @@ -192,124 +100,6 @@ void InputMessageClosure::reset(InputMessageBase* m) { _msg = m; } -int InputMessenger::ProcessNewMessage( - Socket* m, ssize_t bytes, bool read_eof, - const uint64_t received_us, const uint64_t base_realtime, - InputMessageClosure& last_msg) { - m->AddInputBytes(bytes); - - // Avoid this socket to be closed due to idle_timeout_s - m->_last_readtime_us.store(received_us, butil::memory_order_relaxed); - - size_t last_size = m->_read_buf.length(); - int num_bthread_created = 0; - while (1) { - size_t index = 8888; - ParseResult pr = CutInputMessage(m, &index, read_eof); - if (!pr.is_ok()) { - if (pr.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { - // incomplete message, re-read. - // However, some buffer may have been consumed - // under protocols like HTTP. Record this size - m->_last_msg_size += (last_size - m->_read_buf.length()); - break; - } else if (pr.error() == PARSE_ERROR_TRY_OTHERS) { - LOG(WARNING) - << "Close " << *m << " due to unknown message: " - << butil::ToPrintable(m->_read_buf); - m->SetFailed(EINVAL, "Close %s due to unknown message", - m->description().c_str()); - return -1; - } else { - LOG(WARNING) << "Close " << *m << ": " << pr.error_str(); - m->SetFailed(EINVAL, "Close %s: %s", - m->description().c_str(), pr.error_str()); - return -1; - } - } - - m->AddInputMessages(1); - // Calculate average size of messages - const size_t cur_size = m->_read_buf.length(); - if (cur_size == 0) { - // _read_buf is consumed, it's good timing to return blocks - // cached internally back to TLS, otherwise the memory is not - // reused until next message arrives which is quite uncertain - // in situations that most connections are idle. - m->_read_buf.return_cached_blocks(); - } - m->_last_msg_size += (last_size - cur_size); - last_size = cur_size; - const size_t old_avg = m->_avg_msg_size; - if (old_avg != 0) { - m->_avg_msg_size = (old_avg * (MSG_SIZE_WINDOW - 1) + m->_last_msg_size) - / MSG_SIZE_WINDOW; - } else { - m->_avg_msg_size = m->_last_msg_size; - } - m->_last_msg_size = 0; - - if (pr.message() == nullptr) { // the Process() step can be skipped. - continue; - } - pr.message()->_received_us = received_us; - pr.message()->_base_real_us = base_realtime; - - // This unique_ptr prevents msg to be lost before transfering - // ownership to last_msg - DestroyingPtr msg(pr.message()); - m->_transport->QueueMessage(last_msg, &num_bthread_created, false); - if (_handlers[index].process == nullptr) { - LOG(ERROR) << "process of index=" << index << " is NULL"; - continue; - } - m->ReAddress(&msg->_socket); - m->PostponeEOF(); - msg->_process = _handlers[index].process; - msg->_arg = _handlers[index].arg; - - if (_handlers[index].verify != nullptr) { - int auth_error = 0; - if (0 == m->FightAuthentication(&auth_error)) { - // Get the right to authenticate - if (_handlers[index].verify(msg.get())) { - m->SetAuthentication(0); - } else { - m->SetAuthentication(ERPCAUTH); - LOG(WARNING) << "Fail to authenticate " << *m; - m->SetFailed(ERPCAUTH, "Fail to authenticate %s", - m->description().c_str()); - return -1; - } - } else { - LOG_IF(FATAL, auth_error != 0) << - "Impossible! Socket should have been " - "destroyed when authentication failed"; - } - } - if (!m->is_read_progressive()) { - // Transfer ownership to last_msg - last_msg.reset(msg.release()); - } else { - last_msg.reset(msg.release()); - m->_transport->QueueMessage(last_msg, &num_bthread_created, false); - bthread_flush(); - num_bthread_created = 0; - } - } - // In RDMA polling mode, all messages must be executed in a new bthread and - // not in the bthread where the polling bthread is located, because the - // method for processing messages may call synchronization primitives, - // causing the polling bthread to be scheduled out. - if (m->_socket_mode == SOCKET_MODE_RDMA || m->_socket_mode == SOCKET_MODE_UBRING) { - m->_transport->QueueMessage(last_msg, &num_bthread_created, true); - } - if (num_bthread_created) { - bthread_flush(); - } - return 0; -} - void InputMessenger::OnNewMessages(Socket* m) { // Notes: // - If the socket has only one message, the message will be parsed and @@ -321,7 +111,7 @@ void InputMessenger::OnNewMessages(Socket* m) { // is batched(notice the BTHREAD_NOSIGNAL and bthread_flush). // - Verify will always be called in this bthread at most once and before // any process. - InputMessenger* messenger = static_cast(m->user()); + InputMessengerProcessor& processor = m->fd_input_processor(); int progress = Socket::PROGRESS_INIT; // Notice that all *return* no matter successful or not will run last @@ -333,21 +123,14 @@ void InputMessenger::OnNewMessages(Socket* m) { const int64_t received_us = butil::cpuwide_time_us(); const int64_t base_realtime = butil::gettimeofday_us() - received_us; - // Calculate bytes to be read. - size_t once_read = m->_avg_msg_size * 16; - if (once_read < MIN_ONCE_READ) { - once_read = MIN_ONCE_READ; - } else if (once_read > MAX_ONCE_READ) { - once_read = MAX_ONCE_READ; - } - // Read. - const ssize_t nr = m->DoRead(once_read); + const ssize_t nr = m->DoRead(&processor.read_buf(), + processor.OnceReadSize()); if (nr <= 0) { if (0 == nr) { // Set `read_eof' flag and proceed to feed EOF into `Protocol' - // (implied by m->_read_buf.empty), which may produce a new - // `InputMessageBase' under some protocols such as HTTP + // (implied by an empty processor.read_buf()), which may produce + // a new `InputMessageBase' under some protocols such as HTTP LOG_IF(WARNING, FLAGS_log_connection_close) << *m << " was closed by remote side"; read_eof = true; } else if (errno != EAGAIN) { @@ -366,10 +149,10 @@ void InputMessenger::OnNewMessages(Socket* m) { } } - if (messenger->ProcessNewMessage(m, nr, read_eof, received_us, - base_realtime, last_msg) < 0) { + if (processor.ProcessNewMessage(nr, read_eof, received_us, + base_realtime, last_msg) < 0) { return; - } + } } if (read_eof) { diff --git a/src/brpc/input_messenger.h b/src/brpc/input_messenger.h index d056263e2e..8a4b53d8d5 100644 --- a/src/brpc/input_messenger.h +++ b/src/brpc/input_messenger.h @@ -23,6 +23,7 @@ #include "brpc/socket.h" // SocketId, SocketUser #include "brpc/parse_result.h" // ParseResult #include "brpc/input_message_base.h" // InputMessageBase +#include "brpc/input_messenger_processor.h" // InputMessengerProcessor namespace brpc { @@ -94,11 +95,11 @@ class InputMessageClosure { // Process messages from connections. // `Message' corresponds to a client's request or a server's response. class InputMessenger : public SocketUser { -friend class Socket; friend class TcpTransport; friend class RdmaTransport; friend class rdma::RdmaEndpoint; friend class ubring::UBShmEndpoint; +friend class InputMessengerProcessor; public: explicit InputMessenger(size_t capacity = 128); ~InputMessenger(); @@ -137,17 +138,6 @@ friend class ubring::UBShmEndpoint; private: - // Find a valid scissor from `handlers' to cut off `header' and `payload' - // from m->read_buf, save index of the scissor into `index'. - ParseResult CutInputMessage(Socket* m, size_t* index, bool read_eof); - - // Process a new message just received in OnNewMessages - // Return value >= 0 means success - int ProcessNewMessage( - Socket* m, ssize_t bytes, bool read_eof, - const uint64_t received_us, const uint64_t base_realtime, - InputMessageClosure& last_msg); - // User-supplied scissors and handlers. // the index of handler is exactly the same as the protocol InputMessageHandler* _handlers; diff --git a/src/brpc/input_messenger_processor.cpp b/src/brpc/input_messenger_processor.cpp new file mode 100644 index 0000000000..d7541dd60d --- /dev/null +++ b/src/brpc/input_messenger_processor.cpp @@ -0,0 +1,291 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "butil/logging.h" +#include "butil/binary_printer.h" +#include "bthread/unstable.h" +#include "brpc/options.pb.h" +#include "brpc/transport.h" +#include "brpc/input_messenger_processor.h" +#include "brpc/input_messenger.h" + +namespace brpc { + +DECLARE_uint64(max_body_size); + +const size_t MSG_SIZE_WINDOW = 10; // Take last so many message into stat. +const size_t MIN_ONCE_READ = 4096; +const size_t MAX_ONCE_READ = 524288; + +static const char* StreamTypeName(InputMessengerProcessor::StreamType type) { + switch (type) { + case InputMessengerProcessor::STREAM_NONE: return "none"; + case InputMessengerProcessor::STREAM_TCP_FD: return "tcp_fd"; + case InputMessengerProcessor::STREAM_RDMA_QP: return "rdma_qp"; + } + return "unknown"; +} + +InputMessengerProcessor::ParsingStreamGuard::ParsingStreamGuard(Socket* socket, StreamType type) + : _socket(socket) { + CHECK_NE(STREAM_NONE, type) + << "Parsing through a processor that was never Init()ed, " << *socket; + CHECK_EQ(STREAM_NONE, socket->parsing_stream_type()) + << "Two input streams of " << *socket << " are parsing at the same time: " + << StreamTypeName(socket->parsing_stream_type()) << " and " + << StreamTypeName(type); + socket->set_parsing_stream_type(type); +} + +InputMessengerProcessor::ParsingStreamGuard::~ParsingStreamGuard() { + _socket->set_parsing_stream_type(STREAM_NONE); +} + +ParseResult InputMessengerProcessor::CutInputMessage(InputMessenger* messenger, + size_t* index, bool read_eof) { + ParsingStreamGuard parsing_stream_guard(_socket, _stream_type); + InputMessageHandler* handlers = messenger->_handlers; + int preferred = _socket->preferred_index(); + int max_index = (int)messenger->_max_index.load(butil::memory_order_acquire); + // Try preferred handler first. The preferred_index is set on last + // selection or by client. + if (preferred >= 0 && preferred <= max_index + && handlers[preferred].parse != nullptr) { + int cur_index = preferred; + do { + ParseResult result = + handlers[cur_index].parse(&_read_buf, _socket, read_eof, + handlers[cur_index].arg); + if (result.is_ok() || + result.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { + _socket->set_preferred_index(cur_index); + *index = cur_index; + return result; + } else if (result.error() != PARSE_ERROR_TRY_OTHERS) { + // Critical error, return directly. + LOG_IF(ERROR, result.error() == PARSE_ERROR_TOO_BIG_DATA) + << "A message from " << _socket->remote_side() + << "(protocol=" << handlers[cur_index].name + << ") is bigger than " << FLAGS_max_body_size + << " bytes, the connection will be closed." + " Set max_body_size to allow bigger messages"; + return result; + } + + if (_socket->CreatedByConnect()) { + if((ProtocolType)cur_index == PROTOCOL_BAIDU_STD && cur_index == preferred) { + // baidu_std may fall to streaming_rpc. + cur_index = (int)PROTOCOL_STREAMING_RPC; + continue; + } else if((ProtocolType)cur_index == PROTOCOL_STREAMING_RPC && + cur_index == preferred) { + // streaming_rpc may fall to baidu_std. + cur_index = (int)PROTOCOL_BAIDU_STD; + continue; + } else { + // The protocol is fixed at client-side, no need to try others. + LOG(ERROR) << "Fail to parse response from " << _socket->remote_side() + << " by " << handlers[preferred].name + << " at client-side"; + return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); + } + } else { + // Try other protocols. + // + // A handler may lean on this: returning PARSE_ERROR_NOT_ENOUGH_DATA + // above keeps `preferred_index' pinned on it, TRY_OTHERS here gives + // it up. RdmaEndpoint::ExecuteServerHandshake() pins itself that way + // to keep the last handshake read from being taken for a protocol + // detection. See the tail of its phase 2 before changing what + // happens to `preferred_index' here. + break; + } + } while (true); + // Clear context before trying next protocol which probably has + // an incompatible context with the current one. + if (_socket->parsing_context()) { + _socket->reset_parsing_context(nullptr); + } + _socket->set_preferred_index(-1); + } + for (int i = 0; i <= max_index; ++i) { + if (i == preferred || handlers[i].parse == nullptr) { + // Don't try preferred handler(already tried) or invalid handler + continue; + } + ParseResult result = handlers[i].parse(&_read_buf, _socket, read_eof, handlers[i].arg); + if (result.is_ok() || + result.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { + _socket->set_preferred_index(i); + *index = i; + return result; + } else if (result.error() != PARSE_ERROR_TRY_OTHERS) { + // Critical error, return directly. + LOG_IF(ERROR, result.error() == PARSE_ERROR_TOO_BIG_DATA) + << "A message from " << _socket->remote_side() + << "(protocol=" << handlers[i].name + << ") is bigger than " << FLAGS_max_body_size + << " bytes, the connection will be closed." + " Set max_body_size to allow bigger messages"; + return result; + } + // Clear context before trying next protocol which definitely has + // an incompatible context with the current one. + if (_socket->parsing_context()) { + _socket->reset_parsing_context(nullptr); + } + // Try other protocols. + } + return MakeParseError(PARSE_ERROR_TRY_OTHERS); +} + +size_t InputMessengerProcessor::OnceReadSize() const { + size_t once_read = _avg_msg_size * 16; + if (once_read < MIN_ONCE_READ) { + once_read = MIN_ONCE_READ; + } else if (once_read > MAX_ONCE_READ) { + once_read = MAX_ONCE_READ; + } + return once_read; +} + +void InputMessengerProcessor::Reset() { + _read_buf.clear(); + _last_msg_size = 0; + _avg_msg_size = 0; +} + +int InputMessengerProcessor::ProcessNewMessage(ssize_t bytes, bool read_eof, + uint64_t received_us, + uint64_t base_realtime, + InputMessageClosure& last_msg) { + auto messenger = static_cast(_socket->user()); + const InputMessageHandler* handlers = messenger->_handlers; + _socket->AddInputBytes(bytes); + + // Avoid this socket to be closed due to idle_timeout_s + _socket->_last_readtime_us.store(received_us, butil::memory_order_relaxed); + + size_t last_size = _read_buf.length(); + int num_bthread_created = 0; + while (true) { + size_t index = 8888; + ParseResult pr = CutInputMessage(messenger, &index, read_eof); + if (!pr.is_ok()) { + if (pr.error() == PARSE_ERROR_NOT_ENOUGH_DATA) { + // incomplete message, re-read. + // However, some buffer may have been consumed + // under protocols like HTTP. Record this size + _last_msg_size += (last_size - _read_buf.length()); + break; + } else if (pr.error() == PARSE_ERROR_TRY_OTHERS) { + LOG(WARNING) << "Close " << *_socket << " due to unknown message: " + << butil::ToPrintable(_read_buf); + _socket->SetFailed(EINVAL, "Close %s due to unknown message", + _socket->description().c_str()); + return -1; + } else { + LOG(WARNING) << "Close " << *_socket << ": " << pr.error_str(); + _socket->SetFailed(EINVAL, "Close %s: %s", + _socket->description().c_str(), pr.error_str()); + return -1; + } + } + + _socket->AddInputMessages(1); + // Calculate average size of messages + const size_t cur_size = _read_buf.length(); + if (cur_size == 0) { + // _read_buf is consumed, it's good timing to return blocks + // cached internally back to TLS, otherwise the memory is not + // reused until next message arrives which is quite uncertain + // in situations that most connections are idle. + _read_buf.return_cached_blocks(); + } + _last_msg_size += (last_size - cur_size); + last_size = cur_size; + const size_t old_avg = _avg_msg_size; + if (old_avg != 0) { + _avg_msg_size = (old_avg * (MSG_SIZE_WINDOW - 1) + _last_msg_size) / MSG_SIZE_WINDOW; + } else { + _avg_msg_size = _last_msg_size; + } + _last_msg_size = 0; + + if (pr.message() == nullptr) { // the Process() step can be skipped. + continue; + } + pr.message()->_received_us = received_us; + pr.message()->_base_real_us = base_realtime; + + // This unique_ptr prevents msg to be lost before transfering + // ownership to last_msg + DestroyingPtr msg(pr.message()); + _socket->_transport->QueueMessage(last_msg, &num_bthread_created, false); + if (handlers[index].process == nullptr) { + LOG(ERROR) << "process of index=" << index << " is NULL"; + continue; + } + _socket->ReAddress(&msg->_socket); + _socket->PostponeEOF(); + msg->_process = handlers[index].process; + msg->_arg = handlers[index].arg; + + if (handlers[index].verify != nullptr) { + int auth_error = 0; + if (0 == _socket->FightAuthentication(&auth_error)) { + // Get the right to authenticate + if (handlers[index].verify(msg.get())) { + _socket->SetAuthentication(0); + } else { + _socket->SetAuthentication(ERPCAUTH); + LOG(WARNING) << "Fail to authenticate " << *_socket; + _socket->SetFailed(ERPCAUTH, "Fail to authenticate %s", + _socket->description().c_str()); + return -1; + } + } else { + LOG_IF(FATAL, auth_error != 0) << + "Impossible! Socket should have been " + "destroyed when authentication failed"; + } + } + if (!_socket->is_read_progressive()) { + // Transfer ownership to last_msg + last_msg.reset(msg.release()); + } else { + last_msg.reset(msg.release()); + _socket->_transport->QueueMessage(last_msg, &num_bthread_created, false); + bthread_flush(); + num_bthread_created = 0; + } + } + // In RDMA polling mode, all messages must be executed in a new bthread and + // not in the bthread where the polling bthread is located, because the + // method for processing messages may call synchronization primitives, + // causing the polling bthread to be scheduled out. + if (_socket->_socket_mode == SOCKET_MODE_RDMA || + _socket->_socket_mode == SOCKET_MODE_UBRING) { + _socket->_transport->QueueMessage(last_msg, &num_bthread_created, true); + } + if (num_bthread_created) { + bthread_flush(); + } + return 0; +} + +} // namespace brpc diff --git a/src/brpc/input_messenger_processor.h b/src/brpc/input_messenger_processor.h new file mode 100644 index 0000000000..7b94087d65 --- /dev/null +++ b/src/brpc/input_messenger_processor.h @@ -0,0 +1,123 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + + +#ifndef BRPC_INPUT_MESSENGER_PROCESSOR_H_ +#define BRPC_INPUT_MESSENGER_PROCESSOR_H_ + +#include "butil/iobuf.h" // butil::IOPortal +#include "butil/logging.h" // DCHECK +#include "butil/macros.h" // DISALLOW_COPY_AND_ASSIGN +#include "brpc/parse_result.h" // ParseResult + +namespace brpc { + +class Socket; +class InputMessenger; +class InputMessageClosure; + +// The state of one input stream: the data read but not cut off yet, and the +// message-size statistics used to size the next read. +class InputMessengerProcessor { +public: + // Which stream of a Socket a processor drains. + enum StreamType { + STREAM_NONE, + STREAM_TCP_FD, + STREAM_RDMA_QP, + }; + + InputMessengerProcessor() + : _socket(nullptr), _stream_type(STREAM_NONE) + , _last_msg_size(0), _avg_msg_size(0) {} + + DISALLOW_COPY_AND_ASSIGN(InputMessengerProcessor); + + void Init(Socket* socket, StreamType type) { + DCHECK(socket != nullptr); + DCHECK_NE(STREAM_NONE, type) << "STREAM_NONE is not a stream"; + _socket = socket; + _stream_type = type; + } + + // Cut off and process all complete messages in read_buf(), which just + // grew by `bytes` bytes. + // Returns 0 on success, -1 otherwise. + int ProcessNewMessage(ssize_t bytes, bool read_eof, + uint64_t received_us, + uint64_t base_realtime, + InputMessageClosure& last_msg); + + // Data read from the stream but not cut off yet. Only the bthread + // draining this particular stream may touch it. + butil::IOPortal& read_buf() { return _read_buf; } + const butil::IOPortal& read_buf() const { return _read_buf; } + + // How many bytes to ask for on the next read, derived from the sizes of + // the messages seen recently. + size_t OnceReadSize() const; + + uint32_t avg_msg_size() const { return _avg_msg_size; } + + // Drop buffered data and reset the statistics. + void Reset(); + + // Reset the message-size statistics only, keeping buffered data. + void ResetMsgSizeStats() { _last_msg_size = 0; _avg_msg_size = 0; } + +private: + + // Publishes on the Socket which stream the parse callbacks are cutting + // from and clears it when they return, which is what makes + // Socket::parsing_stream_type() mean "the stream being parsed right now". + // + // Entering also DCHECKs that no other stream of that Socket is parsing -- + // the one-stream-at-a-time rule above, checked instead of assumed -- and + // that the processor was Init()ed. + // + // Nested, with Socket::set_parsing_stream_type() private to the enclosing + // class, so nothing else can publish and nothing can forget to unpublish. + class ParsingStreamGuard { + public: + ParsingStreamGuard(Socket* socket, StreamType type); + DISALLOW_COPY_AND_ASSIGN(ParsingStreamGuard); + ~ParsingStreamGuard(); + private: + Socket* _socket; + }; + + // Find a valid scissor among messenger's handlers to cut off one message + // from `_read_buf`, save the index of the scissor into `index`. + ParseResult CutInputMessage(InputMessenger* messenger, size_t* index, bool read_eof); + + // The Socket this stream belongs to. Not owned. + Socket* _socket; + + // Which of that Socket's streams this is, never STREAM_NONE once Init()ed. + StreamType _stream_type; + + butil::IOPortal _read_buf; + + // Size of current incomplete message, set to 0 on complete. + uint32_t _last_msg_size; + // Average message size of last #MSG_SIZE_WINDOW messages (roughly) + uint32_t _avg_msg_size; +}; + +} // namespace brpc + +#endif // BRPC_INPUT_MESSENGER_PROCESSOR_H_ diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index fd70fb5c2b..70a679f8d7 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -147,6 +147,7 @@ RdmaEndpoint::RdmaEndpoint(Socket* s) _rq_size = MAX_QP_SIZE; } _read_butex = bthread::butex_create_checked >(); + _input_processor.Init(s, InputMessengerProcessor::STREAM_RDMA_QP); } RdmaEndpoint::~RdmaEndpoint() { @@ -167,6 +168,7 @@ void RdmaEndpoint::Reset() { _sbuf.clear(); _rbuf.clear(); _rbuf_data.clear(); + _input_processor.Reset(); _remote_recv_block_size = 0; _accumulated_ack = 0; _unsolicited = 0; @@ -220,8 +222,16 @@ void RdmaConnect::Run() { _done(errno, _data); } -void RdmaEndpoint::OnNewDataFromTcp(Socket* m) { - auto* rdma_transport = static_cast(m->_transport.get()); +void RdmaEndpoint::OnNewDataFromTcp(Socket* s) { + if (s->CreatedByConnect()) { + OnNewDataFromTcpAtClient(s); + } else { + OnNewDataFromTcpAtServer(s); + } +} + +void RdmaEndpoint::OnNewDataFromTcpAtClient(Socket* s) { + auto* rdma_transport = static_cast(s->_transport.get()); RdmaEndpoint* ep = rdma_transport->GetRdmaEp(); CHECK(ep != nullptr); @@ -237,36 +247,88 @@ void RdmaEndpoint::OnNewDataFromTcp(Socket* m) { ep->_read_butex->fetch_add(1, butil::memory_order_release); bthread::butex_wake(ep->_read_butex); } else if (state == FALLBACK_TCP){ // handshake finishes - InputMessenger::OnNewMessages(m); + InputMessenger::OnNewMessages(s); return; } else if (state == ESTABLISHED) { - uint8_t tmp; - ssize_t nr = read(ep->_socket->fd(), &tmp, 1); - if (nr == 0) { - ep->_socket->SetEOF(); - return; - } - if (nr > 0) { - LOG(WARNING) << "Read unexpected data from " << ep->_socket; - ep->_socket->SetFailed(EPROTO, "Read unexpected data from %s", - ep->_socket->description().c_str()); + if (!ep->HandleTcpEventAfterEstablished()) { return; } + } + if (!s->MoreReadEvents(&progress)) { + break; + } + } +} - if (errno != EAGAIN) { +void RdmaEndpoint::OnNewDataFromTcpAtServer(Socket* _socket) { + auto* rdma_transport = static_cast(_socket->_transport.get()); + RdmaEndpoint* ep = rdma_transport->GetRdmaEp(); + CHECK(ep != nullptr); + + int progress = Socket::PROGRESS_INIT; + while (true) { + if (_socket->Failed()) { + return; + } + + // Pair with the release stores of ESTABLISHED / FALLBACK_TCP. + if (ep->_state.load(butil::memory_order_acquire) != ESTABLISHED) { + InputMessenger::OnNewMessages(_socket); + // That call may have just finished the handshake and turned RDMA + // on. Start consuming CQ events here rather than inside the parse + // callback: by now OnNewMessages is done with the Socket's + // `parsing_context` / `preferred_index`, so the QP stream can take + // them over without ever overlapping with the fd stream. This is + // the ordering StartCqEvents() asks for. + if (!_socket->Failed() && + ep->_state.load(butil::memory_order_acquire) == ESTABLISHED && + ep->StartCqEvents() < 0) { const int saved_errno = errno; - PLOG(WARNING) << "Fail to read from " << ep->_socket; - ep->_socket->SetFailed(saved_errno, "Fail to read from %s: %s", - ep->_socket->description().c_str(), - berror(saved_errno)); + PLOG(WARNING) << "Fail to start cq events on " << *_socket; + ep->_state.store(FAILED, butil::memory_order_relaxed); + _socket->SetFailed(saved_errno, "Fail to start cq events on %s: %s", + _socket->description().c_str(), berror(saved_errno)); } + return; + } + // RDMA carries the RPCs now, so the fd is watched for EOF only and must + // not be parsed: `preferred_index' / `parsing_context' live on the Socket + // and the QP stream is driving them (https://github.com/apache/brpc/issues/3479). + if (!ep->HandleTcpEventAfterEstablished()) { + return; } - if (!m->MoreReadEvents(&progress)) { + if (!_socket->MoreReadEvents(&progress)) { break; } } } +bool RdmaEndpoint::HandleTcpEventAfterEstablished() { + uint8_t tmp; + ssize_t nr = read(_socket->fd(), &tmp, 1); + if (nr == 0) { + _socket->SetEOF(); + return false; + } + if (nr > 0) { + LOG(WARNING) << "Read unexpected data from " << *_socket; + _socket->SetFailed(EPROTO, "Read unexpected data from %s", + _socket->description().c_str()); + return false; + } + + if (errno != EAGAIN) { + const int saved_errno = errno; + PLOG(WARNING) << "Fail to read from " << *_socket; + _socket->SetFailed(saved_errno, "Fail to read from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + // The socket is dead now, so do not come back for another read of it. + return false; + } + return true; +} + static const int WAIT_TIMEOUT_MS = 50; // Drive an EAGAIN-aware read loop to completion (exactly `len` bytes). @@ -500,7 +562,16 @@ void* RdmaEndpoint::ProcessHandshakeAtClient(void* arg) { } if (rdma_transport->_rdma_state == RdmaTransport::RDMA_ON) { - ep->_state.store(ESTABLISHED, butil::memory_order_relaxed); + ep->_state.store(ESTABLISHED, butil::memory_order_release); + // The handshake is over, so the QP stream may start parsing now. + if (ep->StartCqEvents() < 0) { + const int saved_errno = errno; + PLOG(WARNING) << "Fail to start cq events on " << s->description(); + s->SetFailed(saved_errno, "Fail to complete rdma handshake from %s: %s", + s->description().c_str(), berror(saved_errno)); + ep->_state.store(FAILED, butil::memory_order_relaxed); + return nullptr; + } LOG_IF(INFO, FLAGS_rdma_trace_verbose) << "Client handshake ends (use rdma v" << ep->_handshake_version << ") on " << s->description(); @@ -535,6 +606,36 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s RdmaEndpoint* ep = rdma_transport->_rdma_ep; CHECK(ep != nullptr); + const State state = ep->_state.load(butil::memory_order_acquire); + if (state >= ESTABLISHED) { + // The handshake is over (ESTABLISHED / FALLBACK_TCP / FAILED). Data + // arriving now belongs to a real protocol, yet CutInputMessage() still + // reaches us. + if (state == ESTABLISHED && + s->parsing_stream_type() == InputMessengerProcessor::STREAM_TCP_FD) { + // RDMA is on, so the fd is not an RPC channel any more and whatever + // shows up on it is a protocol error. Reached even though + // OnNewDataFromTcpAtServer() stops handing the fd to OnNewMessages() + // once RDMA is on, because the handshake completes inside OnNewMessages(): + // that round keeps reading the fd until it goes quiet. + if (source->empty()) { + // Nothing to reject yet. Asking for more data keeps the pin, so + // the rest of this round comes back here rather than reaching a + // real protocol, and lets OnNewMessages() report EOF as usual. + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); + } + LOG(WARNING) << "Unexpected " << source->size() << " bytes on the tcp " + "fd of an RDMA connection, drop connection: " + << s->description(); + ep->_state.store(FAILED, butil::memory_order_relaxed); + return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); + } + // Anything else, for the real protocol to parse: the stream carried by + // the QP, or an fd that stayed a normal RPC stream because the handshake + // fell back or failed. + return MakeParseError(PARSE_ERROR_TRY_OTHERS); + } + if (s->parsing_context() == nullptr) { // Phase 1: read the client hello, negotiate, reply server hello. if (source->size() < HELLO_MAGIC_LEN) { @@ -607,13 +708,6 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s if (source->size() < HELLO_ACK_LEN) { return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } - if (source->size() > HELLO_ACK_LEN) { - LOG(WARNING) << "Too many bytes in handshake ACK, drop connection: " - << s->description(); - ep->_state.store(FAILED, butil::memory_order_relaxed); - s->reset_parsing_context(nullptr); - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } uint32_t flags_be = 0; CHECK_EQ(source->cutn(&flags_be, HELLO_ACK_LEN), HELLO_ACK_LEN); @@ -636,13 +730,38 @@ ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); } + if (!source->empty()) { + // RDMA is on, so the TCP fd is no longer an RPC channel. Anything + // trailing the ACK on it can only be a protocol error. This catches what + // arrived in the same read as the ACK. + LOG(WARNING) << "Unexpected " << source->size() << " bytes after the " + "handshake ACK of an RDMA connection, drop connection: " + << s->description(); + ep->_state.store(FAILED, butil::memory_order_relaxed); + s->reset_parsing_context(nullptr); + return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); + } + LOG_IF(INFO, FLAGS_rdma_trace_verbose) << "Server handshake ends (use rdma v" << ep->_handshake_version << ") on " << s->description(); rdma_transport->_rdma_state = RdmaTransport::RDMA_ON; - ep->_state.store(ESTABLISHED, butil::memory_order_relaxed); + ep->_state.store(ESTABLISHED, butil::memory_order_release); s->reset_parsing_context(nullptr); - return MakeParseError(PARSE_ERROR_TRY_OTHERS); + + // Two things are deliberately not done here. + // + // The CQ events are not started: this runs inside CutInputMessage, which + // keeps touching `preferred_index` / `parsing_context` after we return, + // and PollCq would race it for those. OnNewDataFromTcpAtServer() starts + // them once that is over. + // + // TRY_OTHERS is not returned: it would hand `preferred_index` to the real + // protocol, and the remaining reads of this OnNewMessages() round would + // parse the fd as an RPC stream although RDMA has just taken over. Asking + // for more data keeps this handler pinned, so those reads come back to the + // guard at the top of this function and are rejected there. + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } bool RdmaEndpoint::IsWritable() const { @@ -921,11 +1040,12 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) { zerocopy = false; } CHECK_NE(_state.load(butil::memory_order_relaxed), FALLBACK_TCP); + butil::IOPortal& read_buf = _input_processor.read_buf(); if (zerocopy) { - _rbuf[_rq_received].cutn(&_socket->_read_buf, wc.byte_len); + _rbuf[_rq_received].cutn(&read_buf, wc.byte_len); } else { // Copy data when the receive data is really small - _socket->_read_buf.append(_rbuf_data[_rq_received], wc.byte_len); + read_buf.append(_rbuf_data[_rq_received], wc.byte_len); } } if (0 != (wc.wc_flags & IBV_WC_WITH_IMM) && wc.imm_data > 0) { @@ -1111,12 +1231,12 @@ int RdmaEndpoint::DoAllocateResources() { g_rdma_resource_list = g_rdma_resource_list->next; } } - if (!_resource) { + if (_resource == nullptr) { _resource = AllocateQpCq(_sq_size, _rq_size); } else { _resource->next = nullptr; } - if (!_resource) { + if (_resource == nullptr) { return -1; } @@ -1127,25 +1247,6 @@ int RdmaEndpoint::DoAllocateResources() { if (0 != ReqNotifyCq(false, false)) { return -1; } - - SocketOptions options; - options.user = this; - options.keytable_pool = _socket->_keytable_pool; - options.fd = _resource->comp_channel->fd; - options.on_edge_triggered_events = PollCq; - if (Socket::Create(options, &_cq_sid) < 0) { - PLOG(WARNING) << "Fail to create socket for cq"; - return -1; - } - } else { - SocketOptions options; - options.user = this; - options.keytable_pool = _socket->_keytable_pool; - if (Socket::Create(options, &_cq_sid) < 0) { - PLOG(WARNING) << "Fail to create socket for cq"; - return -1; - } - PollerAddCqSid(); } _sbuf.resize(_sq_size - RESERVED_WR_NUM); @@ -1164,6 +1265,43 @@ int RdmaEndpoint::DoAllocateResources() { return 0; } +int RdmaEndpoint::StartCqEvents() { + CHECK_EQ(InputMessengerProcessor::STREAM_NONE, _socket->parsing_stream_type()) + << "StartCqEvents() called while " << *_socket << " is parsing"; + if (_cq_sid != INVALID_SOCKET_ID) { + // Already started. + return 0; + } + if (_resource == nullptr) { + if (BAIDU_UNLIKELY(g_skip_rdma_init)) { + // For UT: AllocateResources() succeeds without allocating anything. + return 0; + } + + LOG(WARNING) << "No RDMA resource to start CQ events on, " << *_socket; + errno = ERDMA; + return -1; + } + + SocketOptions options; + options.user = this; + options.keytable_pool = _socket->_keytable_pool; + if (!FLAGS_rdma_use_polling) { + options.fd = _resource->comp_channel->fd; + options.on_edge_triggered_events = PollCq; + } + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(WARNING) << "Fail to create socket for cq"; + return -1; + } + + if (FLAGS_rdma_use_polling) { + PollerAddCqSid(); + } + + return 0; +} + int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { if (BAIDU_UNLIKELY(g_skip_rdma_init)) { // For UT @@ -1301,7 +1439,7 @@ static int DrainCq(ibv_cq* cq) { } void RdmaEndpoint::DeallocateResources() { - if (!_resource) { + if (_resource == nullptr) { return; } if (FLAGS_rdma_use_polling) { @@ -1328,7 +1466,7 @@ void RdmaEndpoint::DeallocateResources() { bool remove_consumer = true; _reclaim: if (!move_to_rdma_resource_list) { - if (nullptr != _resource->qp) { + if (_resource->qp != nullptr) { int err = IbvDestroyQp(_resource->qp); LOG_IF(WARNING, 0 != err) << "Fail to destroy QP: " << berror(err); _resource->qp = nullptr; @@ -1338,15 +1476,16 @@ void RdmaEndpoint::DeallocateResources() { DeallocateCq(_resource->send_cq); DeallocateCq(_resource->recv_cq); - if (nullptr != _resource->comp_channel) { - // Destroy send_comp_channel will destroy this fd, - // so that we should remove it from epoll fd first - int fd = _resource->comp_channel->fd; - GetGlobalEventDispatcher(fd, _socket->_io_event.bthread_tag()).RemoveConsumer(fd); - remove_consumer = false; + if (_resource->comp_channel != nullptr) { + if (_cq_sid != INVALID_SOCKET_ID) { + // Destroy send_comp_channel will destroy this fd, + // so that we should remove it from epoll fd first + int fd = _resource->comp_channel->fd; + GetGlobalEventDispatcher(fd, _socket->_io_event.bthread_tag()).RemoveConsumer(fd); + remove_consumer = false; + } int err = IbvDestroyCompChannel(_resource->comp_channel); LOG_IF(WARNING, 0 != err) << "Fail to destroy CQ channel: " << berror(err); - } _resource->polling_cq = nullptr; @@ -1478,6 +1617,7 @@ void RdmaEndpoint::PollCq(Socket* m) { } auto* rdma_transport = static_cast(s->_transport.get()); CHECK(ep == rdma_transport->_rdma_ep); + CHECK_GE(ep->_state.load(butil::memory_order_acquire), ESTABLISHED); bool send = false; ibv_cq* cq = ep->_resource->recv_cq; @@ -1595,9 +1735,8 @@ void RdmaEndpoint::PollCq(Socket* m) { // Otherwise it may call too many bthread_flush to affect performance. const int64_t received_us = butil::cpuwide_time_us(); const int64_t base_realtime = butil::gettimeofday_us() - received_us; - InputMessenger* messenger = static_cast(s->user()); - if (messenger->ProcessNewMessage( - s.get(), bytes, false, received_us, base_realtime, last_msg) < 0) { + if (ep->_input_processor.ProcessNewMessage(bytes, false, received_us, + base_realtime, last_msg) < 0) { return; } } @@ -1638,7 +1777,8 @@ void RdmaEndpoint::DebugInfo(std::ostream& os, butil::StringPiece connector) con << connector << "rdma_unacked_rq_wr=" << _new_rq_wrs.load(butil::memory_order_relaxed) << connector << "rdma_received_ack=" << _accumulated_ack << connector << "rdma_unsolicited_sent=" << _unsolicited - << connector << "rdma_unsignaled_sq_wr=" << _sq_unsignaled; + << connector << "rdma_unsignaled_sq_wr=" << _sq_unsignaled + << connector << "rdma_read_buf=" << _input_processor.read_buf().size(); } int RdmaEndpoint::GlobalInitialize() { @@ -1778,23 +1918,27 @@ void RdmaEndpoint::PollingModeRelease(bthread_tag_t tag) { } void RdmaEndpoint::PollerAddCqSid() { + if (_cq_sid == INVALID_SOCKET_ID) { + return; + } + auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num; auto& group = _poller_groups[bthread_self_tag()]; auto& pollers = group.pollers; auto& poller = pollers[index]; - if (INVALID_SOCKET_ID != _cq_sid) { - poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::ADD}); - } + poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::ADD}); } void RdmaEndpoint::PollerRemoveCqSid() { + if (INVALID_SOCKET_ID == _cq_sid) { + return; + } + auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num; auto& group = _poller_groups[bthread_self_tag()]; auto& pollers = group.pollers; auto& poller = pollers[index]; - if (INVALID_SOCKET_ID != _cq_sid) { - poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::REMOVE}); - } + poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::REMOVE}); } } // namespace rdma diff --git a/src/brpc/rdma/rdma_endpoint.h b/src/brpc/rdma/rdma_endpoint.h index 388e31d78e..6d6ac391cd 100644 --- a/src/brpc/rdma/rdma_endpoint.h +++ b/src/brpc/rdma/rdma_endpoint.h @@ -129,7 +129,6 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); butil::StringPiece connector = "\n") const; // Callback when there is new epollin event on TCP fd. - // Only used by client-side RDMA sockets. static void OnNewDataFromTcp(Socket* m); // Real handshake for RDMA-mode sockets. @@ -164,6 +163,11 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Process handshake at the client static void* ProcessHandshakeAtClient(void* arg); + static void OnNewDataFromTcpAtClient(Socket* m); + static void OnNewDataFromTcpAtServer(Socket* m); + + bool HandleTcpEventAfterEstablished(); + // Allocate resources. On failure the endpoint is left with no RDMA // resource attached, so that the handshake can safely fall back to TCP. // Return 0 if success, -1 if failed and errno set @@ -177,6 +181,27 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Release resources void DeallocateResources(); + // Create the Socket wrapping the CQ (and register it with the poller in + // polling mode), which is what makes CQ events reachable and thus starts + // PollCq. + // + // Must not be called before the handshake has reached ESTABLISHED, nor + // from within the fd stream's parsing path: PollCq() parses the input + // stream carried by the QP, and the Socket's `parsing_context` and + // `preferred_index` belong to the fd stream until the handshake is over + // and CutInputMessage has returned. It keeps writing both after the + // handshake handler hands the stream back. Those two are per-Socket, + // so letting PollCq in early makes two streams parse through one context. + // The server therefore calls this from OnNewDataFromTcpAtServer(), after + // OnNewMessages() returns, not from ExecuteServerHandshake(). + // + // No CQE is lost by deferring: BringUpQp() fills the RQ before the QP + // reaches RTS, both CQs are armed by DoAllocateResources(), and adding an + // already readable fd to an edge-triggered epoll reports it immediately. + // + // Return 0 if success, -1 if failed and errno set + int StartCqEvents(); + // Send Imm data to the remote side // Arguments: // imm: imm data in the WR @@ -268,7 +293,7 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); Socket* _socket; // State of Handshake. FALLBACK_TCP publishes RdmaTransport::_rdma_state - // with release ordering and is consumed by OnNewDataFromTcp with acquire + // with release ordering and is consumed by OnNewDataFromTcpAtClient with acquire // ordering. Other state accesses do not publish data and use relaxed // ordering. butil::atomic _state; @@ -301,6 +326,9 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); uint16_t _sq_size; uint16_t _rq_size; + // The input stream carried by the QP. + InputMessengerProcessor _input_processor; + // Act as sendbuf and recvbuf, but requires no memcpy std::vector _sbuf; std::vector _rbuf; diff --git a/src/brpc/rdma/rdma_handshake_server.cpp b/src/brpc/rdma/rdma_handshake_server.cpp index 6dfb0c91f5..1b0ca4226a 100644 --- a/src/brpc/rdma/rdma_handshake_server.cpp +++ b/src/brpc/rdma/rdma_handshake_server.cpp @@ -153,6 +153,11 @@ static int SendUnnegotiableHello(Socket* socket, int version) { } // Fallback handshake for connections that are NOT in RDMA mode. +// +// Unlike the RDMA-mode path, which turns handshake bytes away once its endpoint +// has left the handshake (the state >= ESTABLISHED guard in +// RdmaEndpoint::ExecuteServerHandshake), this one keeps no record of having +// run. See the tail of phase 2. static ParseResult FallbackServerHandshake(butil::IOBuf* source, Socket* socket) { if (socket->parsing_context() == nullptr) { if (source->size() < HELLO_MAGIC_LEN) { @@ -192,8 +197,8 @@ static ParseResult FallbackServerHandshake(butil::IOBuf* source, Socket* socket) return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } CHECK_EQ(source->pop_front(HELLO_ACK_LEN), HELLO_ACK_LEN); - // Handshake done (downgraded to TCP); drop the context and let - // InputMessenger parse the following real RPC. + // Handshake done. + // Drop the context and let InputMessenger parse the following real RPC. socket->reset_parsing_context(nullptr); return MakeParseError(PARSE_ERROR_TRY_OTHERS); } diff --git a/src/brpc/rdma_transport.cpp b/src/brpc/rdma_transport.cpp index ee5151c3a5..a2037aa725 100644 --- a/src/brpc/rdma_transport.cpp +++ b/src/brpc/rdma_transport.cpp @@ -43,12 +43,7 @@ void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { _default_connect = options.app_connect; _on_edge_trigger = options.on_edge_triggered_events; if (options.need_on_edge_trigger && _on_edge_trigger == nullptr) { - // Server-side RDMA sockets drive the handshake through the standard - // InputMessenger path (ParseRdmaHandshake), so they use OnNewMessages - // just like TCP sockets. Only client-side sockets, whose handshake - // (ProcessHandshakeAtClient) is an active blocking bthread relying on - // _read_butex woken by OnNewDataFromTcp, still need OnNewDataFromTcp. - if (options.user == static_cast(get_client_side_messenger())) { + if (_rdma_ep != nullptr) { _on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp; } else { _on_edge_trigger = InputMessenger::OnNewMessages; diff --git a/src/brpc/socket.cpp b/src/brpc/socket.cpp index becb094486..339018744d 100644 --- a/src/brpc/socket.cpp +++ b/src/brpc/socket.cpp @@ -465,10 +465,9 @@ Socket::Socket(Forbidden f) , _conn(nullptr) , _preferred_index(-1) , _hc_count(0) - , _last_msg_size(0) - , _avg_msg_size(0) , _last_readtime_us(0) , _parsing_context(nullptr) + , _parsing_stream_type(InputMessengerProcessor::STREAM_NONE) , _correlation_id(0) , _health_check_interval_s(-1) , _is_hc_related_ref_held(false) @@ -499,6 +498,7 @@ Socket::Socket(Forbidden f) CreateVarsOnce(); pthread_mutex_init(&_id_wait_list_mutex, nullptr); _epollout_butex = bthread::butex_create_checked >(); + _fd_input_processor.Init(this, InputMessengerProcessor::STREAM_TCP_FD); } Socket::~Socket() { @@ -572,8 +572,7 @@ void Socket::ReleaseAllFailedWriteRequests(Socket::WriteRequest* req) { int Socket::ResetFileDescriptor(int fd) { // Reset message sizes when fd is changed. - _last_msg_size = 0; - _avg_msg_size = 0; + _fd_input_processor.ResetMsgSizeStats(); // MUST store `_fd' before adding itself into epoll device to avoid // race conditions with the callback function inside epoll static butil::atomic BAIDU_CACHELINE_ALIGNMENT fd_version(0); @@ -751,7 +750,9 @@ int Socket::OnCreated(const SocketOptions& options) { _app_connect = _transport->Connect(); _preferred_index = -1; _hc_count = 0; - CHECK(_read_buf.empty()); + CHECK(_fd_input_processor.read_buf().empty()); + _parsing_stream_type.store(InputMessengerProcessor::STREAM_NONE, + butil::memory_order_relaxed); const int64_t cpuwide_now = butil::cpuwide_time_us(); _last_readtime_us.store(cpuwide_now, butil::memory_order_relaxed); reset_parsing_context(options.initial_parsing_context); @@ -861,7 +862,7 @@ void Socket::BeforeRecycled() { } _transport->Release(); reset_parsing_context(nullptr); - _read_buf.clear(); + _fd_input_processor.read_buf().clear(); _auth_flag_error.store(0, butil::memory_order_relaxed); bthread_id_error(_auth_id, 0); @@ -1023,9 +1024,11 @@ int Socket::WaitAndReset(int32_t expected_nref) { // parsing_context is very likely to be associated with the fd, // removing it is a safer choice and required by http2. reset_parsing_context(nullptr); - // Must clear _read_buf otehrwise even if the connections is recovered, - // the kept old data is likely to make parsing fail. - _read_buf.clear(); + _parsing_stream_type.store(InputMessengerProcessor::STREAM_NONE, + butil::memory_order_relaxed); + // Must clear the read buffer otehrwise even if the connections is + // recovered, the kept old data is likely to make parsing fail. + _fd_input_processor.read_buf().clear(); _ninprocess.store(1, butil::memory_order_relaxed); _auth_flag_error.store(0, butil::memory_order_relaxed); bthread_id_error(_auth_id, 0); @@ -2080,7 +2083,7 @@ int Socket::SSLHandshake(int fd, bool server_mode) { } } -ssize_t Socket::DoRead(size_t size_hint) { +ssize_t Socket::DoRead(butil::IOPortal* read_buf, size_t size_hint) { if (ssl_state() == SSL_UNKNOWN) { int error_code = 0; _ssl_state = DetectSSLState(fd(), &error_code); @@ -2114,7 +2117,7 @@ ssize_t Socket::DoRead(size_t size_hint) { errno = ESSL; return -1; } - return _read_buf.append_from_file_descriptor(fd(), size_hint); + return read_buf->append_from_file_descriptor(fd(), size_hint); } CHECK_EQ(SSL_CONNECTED, ssl_state()); @@ -2122,7 +2125,7 @@ ssize_t Socket::DoRead(size_t size_hint) { ssize_t nr = 0; { BAIDU_SCOPED_LOCK(_ssl_session_mutex); - nr = _read_buf.append_from_SSL_channel(_ssl_session, &ssl_error, size_hint); + nr = read_buf->append_from_SSL_channel(_ssl_session, &ssl_error, size_hint); } switch (ssl_error) { case SSL_ERROR_NONE: // `nr' > 0 @@ -2382,11 +2385,12 @@ void Socket::DebugSocket(std::ostream& os, SocketId id) { os << " (" << messenger->NameOfProtocol(preferred_index) << ')'; } const int64_t cpuwide_now = butil::cpuwide_time_us(); + InputMessengerProcessor& input_processor = ptr->fd_input_processor(); os << "\nhc_count=" << ptr->_hc_count - << "\navg_input_msg_size=" << ptr->_avg_msg_size + << "\navg_input_msg_size=" << input_processor.avg_msg_size() // NOTE: We're assuming that butil::IOBuf.size() is thread-safe, it is now // however it's not guaranteed. - << "\nread_buf=" << ptr->_read_buf.size() + << "\nread_buf=" << input_processor.read_buf().size() << "\nlast_read_to_now=" << cpuwide_now - ptr->_last_readtime_us << "us" << "\nlast_write_to_now=" << cpuwide_now - ptr->_last_writetime_us << "us" << "\novercrowded=" << ptr->_overcrowded; diff --git a/src/brpc/socket.h b/src/brpc/socket.h index bbde3ecda9..6f6f52fb2b 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -43,6 +43,7 @@ #include "brpc/versioned_ref_with_id.h" #include "brpc/health_check_option.h" #include "brpc/socket_mode.h" +#include "brpc/input_messenger_processor.h" // InputMessengerProcessor namespace brpc { namespace policy { @@ -317,6 +318,7 @@ struct SocketOptions { class BAIDU_CACHELINE_ALIGNMENT/*note*/ Socket : public VersionedRefWithId { friend class EventDispatcher; friend class InputMessenger; +friend class InputMessengerProcessor; friend class Acceptor; friend class ConnectionsService; friend class SocketUser; @@ -722,11 +724,24 @@ friend class TransportFactory; // Returns 0 on success, -1 otherwise int SSLHandshake(int fd, bool server_mode); + // The input stream carried by `_fd`. + InputMessengerProcessor& fd_input_processor() { return _fd_input_processor; } + void set_parsing_stream_type(InputMessengerProcessor::StreamType type) { + _parsing_stream_type.store(type, butil::memory_order_relaxed); + } + + // Which of this Socket's input streams the parse callbacks running right + // now were called for, STREAM_NONE while none are. Only meaningful with a + // parse callback on the stack. + InputMessengerProcessor::StreamType parsing_stream_type() const { + return _parsing_stream_type.load(butil::memory_order_relaxed); + } + // Based upon whether the underlying channel is using SSL (if // SSLState is SSL_UNKNOWN, try to detect at first), read data - // using the corresponding method into `_read_buf'. Returns read + // using the corresponding method into `read_buf`. Returns read // bytes on success, 0 on EOF, -1 otherwise and errno is set - ssize_t DoRead(size_t size_hint); + ssize_t DoRead(butil::IOPortal* read_buf, size_t size_hint); // Based upon whether the underlying channel is using SSL, write // `req' using the corresponding method. Returns written bytes on @@ -890,7 +905,7 @@ friend class TransportFactory; IOEvent _io_event; - // last chosen index of the protocol as a heuristic value to avoid + // Last chosen index of the protocol as a heuristic value to avoid // iterating all protocol handlers each time. int _preferred_index; @@ -898,13 +913,10 @@ friend class TransportFactory; // socket is revived. Only set in HealthCheckTask::OnTriggeringTask() int _hc_count; - // Size of current incomplete message, set to 0 on complete. - uint32_t _last_msg_size; - // Average message size of last #MSG_SIZE_WINDOW messages (roughly) - uint32_t _avg_msg_size; - - // Storing data read from `_fd' but cut-off yet. - butil::IOPortal _read_buf; + // The input stream carried by `_fd`, holding the data read from it but not cut off + // yet. Only the bthread draining the fd (InputMessenger:: OnNewMessages) may touch + // it, see InputMessengerProcessor. + InputMessengerProcessor _fd_input_processor; // Set with cpuwide_time_us() at last read operation butil::atomic _last_readtime_us; @@ -912,6 +924,9 @@ friend class TransportFactory; // Saved context for parsing, reset before trying other protocols. butil::atomic _parsing_context; + // The stream whose data the handlers are currently cutting. + butil::atomic _parsing_stream_type; + // Saving the correlation_id of RPC on protocols that cannot put // correlation_id on-wire and do not send multiple requests on one // connection simultaneously. diff --git a/src/brpc/ubshm/ub_endpoint.cpp b/src/brpc/ubshm/ub_endpoint.cpp index 31539fda85..19bca4f964 100644 --- a/src/brpc/ubshm/ub_endpoint.cpp +++ b/src/brpc/ubshm/ub_endpoint.cpp @@ -52,8 +52,6 @@ DEFINE_bool(ub_poller_yield, false, "Yield thread in UBRing polling mode."); DEFINE_bool(ub_edisp_unsched, false, "Disable event dispatcher schedule"); DEFINE_bool(ub_disable_bthread, false, "Disable bthread in UBRing polling mode."); -static const size_t MIN_ONCE_READ = 4096; -static const size_t MAX_ONCE_READ = 524288; static const size_t IOBUF_IOV_MAX = 256; static const char* MAGIC_STR = "UB"; @@ -480,7 +478,7 @@ void* UBShmEndpoint::ProcessHandshakeAtServer(void* arg) { LOG_IF(INFO, FLAGS_ub_trace_verbose) << "It seems that the " << "client does not use RDMA, fallback to TCP:" << s->description(); - s->_read_buf.append(data, MAGIC_STR_LEN); + s->fd_input_processor().read_buf().append(data, MAGIC_STR_LEN); ep->_state = FALLBACK_TCP; ub_transport->_ub_state = UBShmTransport::UB_OFF; ep->TryReadOnTcp(); @@ -746,24 +744,19 @@ void UBShmEndpoint::PollIn(UBShmEndpoint* ep, uint32_t ep_event) { return; } + InputMessengerProcessor& processor = s->fd_input_processor(); bool read_eof = false; while (!read_eof) { const int64_t received_us = butil::cpuwide_time_us(); const int64_t base_realtime = butil::gettimeofday_us() - received_us; - size_t once_read = s->_avg_msg_size * 16; - if (once_read < MIN_ONCE_READ) { - once_read = MIN_ONCE_READ; - } else if (once_read > MAX_ONCE_READ) { - once_read = MAX_ONCE_READ; - } - - const ssize_t nr = s->_read_buf.append_from_reader(ep->_ub_ring, once_read); + const ssize_t nr = processor.read_buf().append_from_reader( + ep->_ub_ring, processor.OnceReadSize()); if (nr <= 0) { if (0 == nr) { // Set `read_eof' flag and proceed to feed EOF into `Protocol' - // (implied by m->_read_buf.empty), which may produce a new - // `InputMessageBase' under some protocols such as HTTP + // (implied by an empty processor.read_buf()), which may produce + // a new `InputMessageBase' under some protocols such as HTTP LOG_IF(WARNING, FLAGS_log_connection_close) << *s << " was closed by remote side"; read_eof = true; } else if (errno != EAGAIN) { @@ -780,11 +773,10 @@ void UBShmEndpoint::PollIn(UBShmEndpoint* ep, uint32_t ep_event) { } } - InputMessenger* messenger = static_cast(s->user()); - if (messenger->ProcessNewMessage(s.get(), nr, read_eof, received_us, - base_realtime, last_msg) < 0) { + if (processor.ProcessNewMessage(nr, read_eof, received_us, + base_realtime, last_msg) < 0) { return; - } + } } if (read_eof) { diff --git a/test/brpc_rdma_unittest.cpp b/test/brpc_rdma_unittest.cpp index 9c52acb797..ce40fa59f1 100644 --- a/test/brpc_rdma_unittest.cpp +++ b/test/brpc_rdma_unittest.cpp @@ -21,6 +21,8 @@ #include #include #if BRPC_WITH_RDMA +#include +#include #include #include "butil/endpoint.h" #include "butil/fd_guard.h" @@ -79,6 +81,14 @@ extern bool g_fail_resource_alloc_for_test; static std::string g_ip = "127.0.0.1"; static butil::EndPoint g_ep; +// Number of Echo requests the server has actually served. The churn tests below +// cut the connection while requests are in flight, and from the client side +// "the server never saw this request" is indistinguishable from "the server +// answered it and the reply died with the connection" -- both just look like a +// failed RPC. Only this counter tells the two apart, and it is the server having +// real work in flight that makes the race those tests hunt reachable at all. +static butil::atomic g_echo_served(0); + class MyEchoService : public ::test::EchoService { void Echo(google::protobuf::RpcController* cntl_base, const ::test::EchoRequest* req, @@ -86,6 +96,7 @@ class MyEchoService : public ::test::EchoService { google::protobuf::Closure* done) { Controller* cntl = static_cast(cntl_base); ClosureGuard done_guard(done); + g_echo_served.fetch_add(1, butil::memory_order_relaxed); if (req->server_fail()) { cntl->SetFailed(req->server_fail(), "Server fail1"); cntl->SetFailed(req->server_fail(), "Server fail2"); @@ -719,6 +730,234 @@ TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { StopServer(); } +// Build a well-formed v2 client hello: "RDMA" followed by the 36B body. +static void MakeV2ClientHello(uint8_t (&data)[rdma::HELLO_V2_MSG_LEN_MIN]) { + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); +} + +// Connect, push a well-formed v2 hello and read back the server's reply, which +// leaves the server in S_ACK_WAIT waiting for the 4B ACK. +static void HandshakeUntilAckWait(butil::fd_guard* sockfd) { + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + sockfd->reset(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(*sockfd >= 0); + ASSERT_EQ(0, connect(*sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + + uint8_t hello[rdma::HELLO_V2_MSG_LEN_MIN]; + MakeV2ClientHello(hello); + ASSERT_EQ((ssize_t)sizeof(hello), write(*sockfd, hello, sizeof(hello))); + usleep(100000); // wait for server to handle the msg + uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ((ssize_t)sizeof(reply), read(*sockfd, reply, sizeof(reply))); +} + +// A client is free to pipeline its first request right behind the handshake +// ACK. Only the 4B ACK belongs to the handshake. Whatever follows it must be +// handed over to the real protocol instead of dropping the connection. +TEST_F(RdmaTest, server_accepts_data_pipelined_behind_fallback_ack) { + StartServer(); + + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + ASSERT_EQ(rdma::RdmaEndpoint::S_ACK_WAIT, transport->_rdma_ep->_state); + + // An ACK asking for TCP, plus the first 4 bytes of a baidu_std request. One + // write, so that both end up in the same read on the server. + uint8_t ack_and_data[rdma::HELLO_ACK_LEN + 4]; + const uint32_t flags = butil::HostToNet32(0); + memcpy(ack_and_data, &flags, rdma::HELLO_ACK_LEN); + memcpy(ack_and_data + rdma::HELLO_ACK_LEN, "PRPC", 4); + ASSERT_EQ((ssize_t)sizeof(ack_and_data), + write(sockfd, ack_and_data, sizeof(ack_and_data))); + usleep(100000); // wait for server to handle the msg + + // The handshake took the ACK only and left "PRPC" to baidu_std, which is + // now waiting for the rest of its 12B header. So the connection lives on + // with those 4 bytes still buffered. + ASSERT_EQ(rdma::RdmaEndpoint::FALLBACK_TCP, transport->_rdma_ep->_state); + ASSERT_EQ(RdmaTransport::RDMA_OFF, transport->_rdma_state); + ASSERT_TRUE(GetSocketFromServer(0) != nullptr); + ASSERT_EQ(4u, s->fd_input_processor().read_buf().size()); + + sockfd.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +// Once RDMA is on, the TCP fd is no longer an RPC channel, so bytes trailing +// the ACK can only be a protocol error. +TEST_F(RdmaTest, server_rejects_data_pipelined_behind_rdma_ack) { + StartServer(); + + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + ASSERT_EQ(rdma::RdmaEndpoint::S_ACK_WAIT, transport->_rdma_ep->_state); + + uint8_t ack_and_data[rdma::HELLO_ACK_LEN + 4]; + const uint32_t flags = butil::HostToNet32(rdma::HELLO_ACK_RDMA_OK); + memcpy(ack_and_data, &flags, rdma::HELLO_ACK_LEN); + memcpy(ack_and_data + rdma::HELLO_ACK_LEN, "PRPC", 4); + ASSERT_EQ((ssize_t)sizeof(ack_and_data), + write(sockfd, ack_and_data, sizeof(ack_and_data))); + usleep(100000); // wait for server to handle the msg + + // Note that `transport->_rdma_ep` is gone by now: dropping the connection + // recycles the Socket, and RdmaTransport::Release() deletes the endpoint. + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +// Once RDMA is on, the server must stop parsing its TCP fd altogether. +TEST_F(RdmaTest, server_stops_parsing_tcp_fd_once_rdma_is_on) { + StartServer(); + + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + ASSERT_EQ(rdma::RdmaEndpoint::S_ACK_WAIT, transport->_rdma_ep->_state); + + // A bare ACK asking for RDMA. Nothing trails it, so the handshake ends in + // ESTABLISHED instead of being rejected (see the test above). + const uint32_t flags = butil::HostToNet32(rdma::HELLO_ACK_RDMA_OK); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(rdma::RdmaEndpoint::ESTABLISHED, transport->_rdma_ep->_state); + ASSERT_EQ(RdmaTransport::RDMA_ON, transport->_rdma_state); + ASSERT_TRUE(GetSocketFromServer(0) != nullptr); + + ASSERT_EQ(4, write(sockfd, "PRPC", 4)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +// The same bytes on the stream carried by the QP are a real RPC, and the handler +// must decline so that CutInputMessage() moves on to the protocol handlers. +TEST_F(RdmaTest, server_parses_qp_stream_after_rdma_is_on) { + StartServer(); + + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + ASSERT_EQ(rdma::RdmaEndpoint::S_ACK_WAIT, transport->_rdma_ep->_state); + + const uint32_t flags = butil::HostToNet32(rdma::HELLO_ACK_RDMA_OK); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(rdma::RdmaEndpoint::ESTABLISHED, transport->_rdma_ep->_state); + + InputMessengerProcessor& qp_stream = transport->_rdma_ep->_input_processor; + ASSERT_TRUE(qp_stream.read_buf().empty()); + qp_stream.read_buf().append("PRPC"); + InputMessageClosure last_msg; + ASSERT_EQ(0, qp_stream.ProcessNewMessage(4, false, butil::gettimeofday_us(), 0, last_msg)); + // baidu_std claimed the stream and is waiting for the rest of its header. + ASSERT_EQ((int)PROTOCOL_BAIDU_STD, s->preferred_index()); + ASSERT_EQ(4u, qp_stream.read_buf().size()); + ASSERT_TRUE(s->fd_input_processor().read_buf().empty()); + ASSERT_EQ(rdma::RdmaEndpoint::ESTABLISHED, transport->_rdma_ep->_state); + ASSERT_FALSE(s->Failed()); + + StopServer(); +} + +// After the handshake is over, CutInputMessage() still offers the data to every +// registered handler, this one included. It must decline instead of reading the +// data as a fresh client hello. +TEST_F(RdmaTest, server_declines_handshake_bytes_after_fallback) { + StartServer(); + + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + + const uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(rdma::RdmaEndpoint::FALLBACK_TCP, transport->_rdma_ep->_state); + + // Replay a valid hello. baidu_std rejects it and no other protocol claims + // it, so the connection is dropped. What must NOT happen is a second + // handshake: that would answer with another server hello. + uint8_t hello[rdma::HELLO_V2_MSG_LEN_MIN]; + MakeV2ClientHello(hello); + ASSERT_EQ((ssize_t)sizeof(hello), write(sockfd, hello, sizeof(hello))); + usleep(100000); // wait for server to handle the msg + + // Note that `transport->_rdma_ep` is gone by now: dropping the connection + // recycles the Socket, and RdmaTransport::Release() deletes the endpoint. + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_LE(recv(sockfd, reply, sizeof(reply), MSG_DONTWAIT), 0); + + StopServer(); +} + +TEST_F(RdmaTest, fd_and_qp_input_streams_are_separate) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + auto* transport = static_cast(s->_transport.get()); + + // The fd stream belongs to the Socket, the QP stream to the endpoint. + InputMessengerProcessor& fd_stream = s->fd_input_processor(); + InputMessengerProcessor& qp_stream = transport->_rdma_ep->_input_processor; + ASSERT_NE(&fd_stream, &qp_stream); + ASSERT_NE(&fd_stream.read_buf(), &qp_stream.read_buf()); + + // Two magic bytes are too few to dispatch on, so they stay buffered. In the + // fd stream, and only there. + ASSERT_EQ(2, write(sockfd, "RD", 2)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, transport->_rdma_ep->_state); + ASSERT_EQ(2u, fd_stream.read_buf().size()); + ASSERT_TRUE(qp_stream.read_buf().empty()); + + sockfd.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + TEST_F(RdmaTest, server_miss_before_hello_send) { butil::fd_guard sockfd(butil::tcp_listen(g_ep)); EXPECT_TRUE(sockfd >= 0); @@ -2145,6 +2384,73 @@ TEST_F(RdmaTest, channel_option_invalid) { ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); } +// Rounds, per-round RPC count and attachment sizes shared by the end-to-end +// tests below. One RPC per test leaves everything that only shows up on the +// second message untouched -- buffer reuse, an EOF read racing another writer +// of the same input stream, resource recycling. +static const int E2E_ROUND_NUM = 3; +static const int E2E_RPC_NUM = 32; +static const size_t E2E_ATTACH_SIZE[] = { 0, 4096, 128 * 1024 }; + +static void ShutdownClientConnection(Controller& cntl) { + SocketUniquePtr s; + if (Socket::Address(cntl._single_server_id, &s) == 0) { + ::shutdown(s->fd(), SHUT_WR); + } +} + +// Returns the number of RPCs that succeeded. A test that severs the connection +// cannot predict which ones make it, but the ones that do must still be right, +// so failures are tolerated here and the caller decides how many it demands. +static int SendEchoRpcs(Channel& channel, int rpc_num, size_t attach_size, + const std::function& disturb = nullptr, + int disturb_at = 0) { + std::vector cntl(rpc_num); + std::vector req(rpc_num); + std::vector res(rpc_num); + std::vector attach(rpc_num); + for (int i = 0; i < rpc_num; ++i) { + req[i].set_message("hello"); + req[i].set_code(i + 1); + if (attach_size > 0) { + EXPECT_EQ(0, attach[i].resize( + attach_size, static_cast('a' + i % 26))); + cntl[i].request_attachment().append(attach[i]); + } + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], DoNothing()); + if (disturb && i == disturb_at) { + disturb(cntl[i]); + } + } + int succeeded = 0; + for (int i = 0; i < rpc_num; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].Failed()) { + continue; + } + ++succeeded; + EXPECT_EQ("MyEchoService", res[i].message()) << "rpc[" << i << "]"; + EXPECT_EQ(1, res[i].code_list_size()) << "rpc[" << i << "]"; + if (res[i].code_list_size() == 1) { + EXPECT_EQ(i + 1, res[i].code_list(0)) << "rpc[" << i << "]"; + } + EXPECT_EQ(attach_size, cntl[i].response_attachment().size()) << "rpc[" << i << "]"; + EXPECT_TRUE(attach[i].equals(cntl[i].response_attachment())) << "rpc[" << i << "]"; + } + return succeeded; +} + +static void SendEchoRpcsInRounds(Channel& channel) { + for (int round = 0; round < E2E_ROUND_NUM; ++round) { + for (size_t i = 0; i < arraysize(E2E_ATTACH_SIZE); ++i) { + ASSERT_EQ(E2E_RPC_NUM, + SendEchoRpcs(channel, E2E_RPC_NUM, E2E_ATTACH_SIZE[i])) + << "round=" << round + << " attach_size=" << E2E_ATTACH_SIZE[i]; + } + } +} + TEST_P(RdmaRpcTest, rdma_client_to_rdma_server) { if (!FLAGS_rdma_test_enable) { return; @@ -2156,18 +2462,10 @@ TEST_P(RdmaRpcTest, rdma_client_to_rdma_server) { ChannelOptions chan_options; chan_options.socket_mode = SOCKET_MODE_RDMA; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; + chan_options.timeout_ms = 5000; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure* done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - // usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); + ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); StopServer(); } @@ -2178,18 +2476,10 @@ TEST_P(RdmaRpcTest, tcp_client_to_tcp_server) { Channel channel; ChannelOptions chan_options; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; + chan_options.timeout_ms = 5000; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure* done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); + ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); StopServer(); } @@ -2200,18 +2490,10 @@ TEST_P(RdmaRpcTest, tcp_client_to_rdma_server) { Channel channel; ChannelOptions chan_options; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; + chan_options.timeout_ms = 5000; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure* done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); + ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); StopServer(); } @@ -2223,18 +2505,73 @@ TEST_P(RdmaRpcTest, rdma_client_to_tcp_server) { ChannelOptions chan_options; chan_options.socket_mode = SOCKET_MODE_RDMA; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; + chan_options.timeout_ms = 5000; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure* done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_FALSE(cntl.Failed()); + ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); + + StopServer(); +} + +TEST_P(RdmaRpcTest, tcp_client_to_rdma_server_short_connection) { + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.connect_timeout_ms = 1000; + chan_options.timeout_ms = 10000; + chan_options.max_retry = 0; + chan_options.connection_type = "short"; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + for (int round = 0; round < 8; ++round) { + ASSERT_EQ(E2E_RPC_NUM, SendEchoRpcs(channel, E2E_RPC_NUM, 4096)) + << "round=" << round; + } + + StopServer(); +} + +// Rounds of connection churn: a race needs attempts, not one well-timed shot. +static const int CHURN_ROUND_NUM = 16; +static const int CHURN_RPC_NUM = 64; +static const size_t CHURN_ATTACH_SIZE = 32 * 1024; + +TEST_P(RdmaRpcTest, rdma_server_survives_connection_churn) { + StartServer(); + + ChannelOptions chan_options; + chan_options.connect_timeout_ms = 1000; + chan_options.timeout_ms = 3000; + chan_options.max_retry = 0; + const int served_before = g_echo_served.load(butil::memory_order_relaxed); + int succeeded = 0; + for (int round = 0; round < CHURN_ROUND_NUM; ++round) { + // A fresh Channel per round so the Socket is dropped from the socket + // map when the Channel dies, instead of the next round inheriting it. + Channel channel; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + // Warm up before cutting. Without this the cut of round 0 lands on a + // connection that has not finished connecting yet, where fd() is still + // -1 and shutdown() is a silent no-op. + ASSERT_EQ(1, SendEchoRpcs(channel, 1, 0)) << "round=" << round; + succeeded += SendEchoRpcs(channel, CHURN_RPC_NUM, CHURN_ATTACH_SIZE, + ShutdownClientConnection, + round * CHURN_RPC_NUM / CHURN_ROUND_NUM); + ASSERT_FALSE(HasFailure()) << "round=" << round; + } + + const int served = g_echo_served.load(butil::memory_order_relaxed) - + served_before - CHURN_ROUND_NUM; + LOG(INFO) << "server served " << served << " of " + << CHURN_ROUND_NUM * CHURN_RPC_NUM << " requests during the churn, " + << succeeded << " replies made it back"; + ASSERT_GT(served, 0); + + // The churn must leave the server able to serve a fresh connection, and + // serve it correctly. + Channel channel; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + ASSERT_EQ(CHURN_RPC_NUM, SendEchoRpcs(channel, CHURN_RPC_NUM, CHURN_ATTACH_SIZE)); StopServer(); } @@ -2602,6 +2939,46 @@ TEST_P(RdmaRpcTest, client_close_during_rpc) { StopServer(); } +TEST_P(RdmaRpcTest, rdma_client_close_during_rpc_repeatedly) { + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 1000; + chan_options.timeout_ms = 3000; + chan_options.max_retry = 0; + const int served_before = g_echo_served.load(butil::memory_order_relaxed); + int succeeded = 0; + for (int round = 0; round < CHURN_ROUND_NUM; ++round) { + Channel channel; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + // Warm up so the cut lands on an established RDMA connection. + ASSERT_EQ(1, SendEchoRpcs(channel, 1, 0)) << "round=" << round; + succeeded += SendEchoRpcs(channel, CHURN_RPC_NUM, CHURN_ATTACH_SIZE, + ShutdownClientConnection, + round * CHURN_RPC_NUM / CHURN_ROUND_NUM); + ASSERT_FALSE(HasFailure()) << "round=" << round; + } + + + const int served = g_echo_served.load(butil::memory_order_relaxed) - + served_before - CHURN_ROUND_NUM; + LOG(INFO) << "server served " << served << " of " + << CHURN_ROUND_NUM * CHURN_RPC_NUM << " requests during the churn, " + << succeeded << " replies made it back"; + ASSERT_GT(served, 0); + + Channel channel; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + ASSERT_EQ(CHURN_RPC_NUM, SendEchoRpcs(channel, CHURN_RPC_NUM, CHURN_ATTACH_SIZE)); + + StopServer(); +} + TEST_P(RdmaRpcTest, verbs_error_handling) { if (!FLAGS_rdma_test_enable) { return;