diff --git a/src/brpc/input_messenger.cpp b/src/brpc/input_messenger.cpp index 81154be134..3c9a69689a 100644 --- a/src/brpc/input_messenger.cpp +++ b/src/brpc/input_messenger.cpp @@ -17,6 +17,8 @@ #include +#include +#include #include "butil/fd_guard.h" // fd_guard #include "butil/logging.h" // CHECK #include "butil/time.h" // cpuwide_time_us @@ -72,6 +74,17 @@ DEFINE_int32(socket_tcp_user_timeout_ms, -1, "connection and return ETIMEDOUT to the application. Only linux supports " "TCP_USER_TIMEOUT."); +DEFINE_int32(input_message_batch_process_size, 0, + "Experimental. -1 adaptively processes up to 16 parsed input " + "messages in one bthread based on the recent per-socket burst. " + "Values greater than 1 use a fixed batch size. 0 or 1 preserves " + "the original one-message-per-bthread behavior."); +static bool ValidateInputMessageBatchProcessSize(const char*, int32_t value) { + return value >= -1; +} +BRPC_VALIDATE_GFLAG(input_message_batch_process_size, + ValidateInputMessageBatchProcessSize); + DECLARE_bool(usercode_in_pthread); DECLARE_bool(usercode_in_coroutine); DECLARE_uint64(max_body_size); @@ -79,6 +92,10 @@ 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; +const uint32_t INPUT_BATCH_EMA_SCALE = 256; +const uint32_t MAX_ADAPTIVE_INPUT_BATCH_SIZE = 16; +const uint32_t MAX_ADAPTIVE_INPUT_BATCH_SAMPLE = + MAX_ADAPTIVE_INPUT_BATCH_SIZE * 2; ParseResult InputMessenger::CutInputMessage( Socket* m, size_t* index, bool read_eof) { @@ -173,6 +190,65 @@ void* ProcessInputMessage(void* void_arg) { return nullptr; } +void* ProcessInputMessageBatch(void* void_arg) { + std::unique_ptr batch( + static_cast(void_arg)); + try { + batch->Run(); + } catch (...) { + LOG(ERROR) << "An input message handler threw while processing a batch"; + } + return nullptr; +} + +InputMessageBatch::InputMessageBatch(size_t capacity) { + // Avoid a large upfront allocation from a user-controlled fixed batch size. + _msgs.reserve(std::min( + capacity, static_cast(MAX_ADAPTIVE_INPUT_BATCH_SIZE))); +} + +InputMessageBatch::~InputMessageBatch() noexcept { + try { + Run(); + } catch (...) { + LOG(ERROR) << "An input message handler threw during batch cleanup"; + DestroyRemainingMessages(); + } +} + +void InputMessageBatch::add(InputMessageBase* msg) { + if (msg) { + _msgs.push_back(msg); + } +} + +void InputMessageBatch::Run() { + for (size_t i = 0; i < _msgs.size(); ++i) { + InputMessageBase* msg = _msgs[i]; + _msgs[i] = nullptr; + if (msg == nullptr) { + continue; + } + ProcessInputMessage(msg); + } + _msgs.clear(); +} + +void InputMessageBatch::DestroyRemainingMessages() noexcept { + for (size_t i = 0; i < _msgs.size(); ++i) { + if (_msgs[i] == nullptr) { + continue; + } + try { + _msgs[i]->Destroy(); + } catch (...) { + LOG(ERROR) << "Failed to destroy an unprocessed input message"; + } + _msgs[i] = nullptr; + } + _msgs.clear(); +} + struct RunLastMessage { inline void operator()(InputMessageBase* last_msg) { ProcessInputMessage(last_msg); @@ -192,6 +268,77 @@ void InputMessageClosure::reset(InputMessageBase* m) { _msg = m; } +void InputMessenger::QueueInputMessageBatch( + Socket* m, std::unique_ptr* batch, + int* num_bthread_created) { + if (!batch->get() || (*batch)->empty()) { + return; + } + m->_transport->QueueMessages(batch->release(), num_bthread_created); +} + +void InputMessenger::QueueLastMessageOrBatch( + Socket* m, InputMessageClosure& last_msg, + std::unique_ptr* batch, + int* num_bthread_created, size_t batch_size) { + InputMessageBase* msg = last_msg.release(); + if (!msg) { + return; + } + if (!batch->get()) { + batch->reset(new (std::nothrow) InputMessageBatch(batch_size)); + } + if (!batch->get()) { + last_msg.reset(msg); + m->_transport->QueueMessage( + last_msg, num_bthread_created, false); + return; + } + (*batch)->add(msg); + if ((*batch)->size() >= batch_size) { + QueueInputMessageBatch(m, batch, num_bthread_created); + } +} + +uint32_t InputMessenger::UpdateAdaptiveBatchSize( + uint32_t* messages_per_read_ema_q8, + uint32_t current_batch_size, + size_t parsed_message_count) { + if (parsed_message_count == 0) { + return current_batch_size; + } + if (*messages_per_read_ema_q8 == 0 || current_batch_size == 0) { + *messages_per_read_ema_q8 = INPUT_BATCH_EMA_SCALE; + current_batch_size = 1; + } + + const uint32_t sample = static_cast( + std::min(parsed_message_count, + static_cast(MAX_ADAPTIVE_INPUT_BATCH_SAMPLE))); + const uint32_t sample_q8 = sample * INPUT_BATCH_EMA_SCALE; + uint32_t ema_q8 = *messages_per_read_ema_q8; + if (sample_q8 > ema_q8) { + // Increase slowly to avoid turning a short burst into persistent + // head-of-line blocking. + ema_q8 += (sample_q8 - ema_q8) / 8; + } else { + // Reduce quickly when the connection becomes sparse. + ema_q8 -= (ema_q8 - sample_q8 + 1) / 2; + } + *messages_per_read_ema_q8 = ema_q8; + + uint32_t desired_batch_size = 1; + while (desired_batch_size < MAX_ADAPTIVE_INPUT_BATCH_SIZE && + ema_q8 > desired_batch_size * INPUT_BATCH_EMA_SCALE) { + desired_batch_size *= 2; + } + if (desired_batch_size > current_batch_size) { + // Increase at most one level for each observation. + return std::min(current_batch_size * 2, desired_batch_size); + } + return desired_batch_size; +} + int InputMessenger::ProcessNewMessage( Socket* m, ssize_t bytes, bool read_eof, const uint64_t received_us, const uint64_t base_realtime, @@ -203,6 +350,27 @@ int InputMessenger::ProcessNewMessage( size_t last_size = m->_read_buf.length(); int num_bthread_created = 0; + const int configured_batch_size = + FLAGS_input_message_batch_process_size; + const bool adaptive_batch_process = + configured_batch_size == -1 && !FLAGS_usercode_in_coroutine; + size_t batch_size = configured_batch_size > 0 + ? static_cast(configured_batch_size) : 1; + if (adaptive_batch_process) { + if (m->_adaptive_input_message_batch_size == 0) { + m->_input_messages_per_read_ema_q8 = INPUT_BATCH_EMA_SCALE; + m->_adaptive_input_message_batch_size = 1; + } + batch_size = m->_adaptive_input_message_batch_size; + } else if (m->_adaptive_input_message_batch_size != 0) { + // Do not reuse history after switching away from adaptive mode. + m->_input_messages_per_read_ema_q8 = 0; + m->_adaptive_input_message_batch_size = 0; + } + const bool batch_process = + batch_size > 1 && !FLAGS_usercode_in_coroutine; + size_t batchable_message_count = 0; + std::unique_ptr input_batch; while (1) { size_t index = 8888; ParseResult pr = CutInputMessage(m, &index, read_eof); @@ -258,7 +426,14 @@ int InputMessenger::ProcessNewMessage( // 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 (batch_process) { + QueueLastMessageOrBatch( + m, last_msg, &input_batch, + &num_bthread_created, batch_size); + } else { + m->_transport->QueueMessage( + last_msg, &num_bthread_created, false); + } if (_handlers[index].process == nullptr) { LOG(ERROR) << "process of index=" << index << " is NULL"; continue; @@ -290,8 +465,15 @@ int InputMessenger::ProcessNewMessage( if (!m->is_read_progressive()) { // Transfer ownership to last_msg last_msg.reset(msg.release()); + if (adaptive_batch_process) { + ++batchable_message_count; + } } else { last_msg.reset(msg.release()); + if (batch_process) { + QueueInputMessageBatch( + m, &input_batch, &num_bthread_created); + } m->_transport->QueueMessage(last_msg, &num_bthread_created, false); bthread_flush(); num_bthread_created = 0; @@ -301,9 +483,21 @@ int InputMessenger::ProcessNewMessage( // 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) { + if (batch_process) { + QueueLastMessageOrBatch( + m, last_msg, &input_batch, &num_bthread_created, batch_size); + QueueInputMessageBatch(m, &input_batch, &num_bthread_created); + } else if (m->_socket_mode == SOCKET_MODE_RDMA || + m->_socket_mode == SOCKET_MODE_UBRING) { m->_transport->QueueMessage(last_msg, &num_bthread_created, true); } + if (adaptive_batch_process && batchable_message_count != 0) { + m->_adaptive_input_message_batch_size = + UpdateAdaptiveBatchSize( + &m->_input_messages_per_read_ema_q8, + m->_adaptive_input_message_batch_size, + batchable_message_count); + } if (num_bthread_created) { bthread_flush(); } @@ -317,8 +511,10 @@ void InputMessenger::OnNewMessages(Socket* m) { // - If the socket has several messages, all messages will be parsed ( // meaning cutting from butil::IOBuf. serializing from protobuf is part of // "process") in this bthread. All messages except the last one will be - // processed in separate bthreads. To minimize the overhead, scheduling - // is batched(notice the BTHREAD_NOSIGNAL and bthread_flush). + // processed in separate bthreads, or in batches when + // -input_message_batch_process_size is -1 or greater than 1. To minimize + // the overhead, scheduling 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()); diff --git a/src/brpc/input_messenger.h b/src/brpc/input_messenger.h index d056263e2e..e0e5c7e47d 100644 --- a/src/brpc/input_messenger.h +++ b/src/brpc/input_messenger.h @@ -19,6 +19,9 @@ #ifndef BRPC_INPUT_MESSENGER_H #define BRPC_INPUT_MESSENGER_H +#include +#include + #include "butil/iobuf.h" // butil::IOBuf #include "brpc/socket.h" // SocketId, SocketUser #include "brpc/parse_result.h" // ParseResult @@ -91,6 +94,26 @@ class InputMessageClosure { InputMessageBase* _msg; }; +class InputMessageBatch { +public: + InputMessageBatch() {} + explicit InputMessageBatch(size_t capacity); + ~InputMessageBatch() noexcept; + + void add(InputMessageBase* msg); + void Run(); + bool empty() const { return _msgs.empty(); } + size_t size() const { return _msgs.size(); } + +private: + void DestroyRemainingMessages() noexcept; + + std::vector _msgs; +}; + +void* ProcessInputMessage(void* void_arg); +void* ProcessInputMessageBatch(void* void_arg); + // Process messages from connections. // `Message' corresponds to a client's request or a server's response. class InputMessenger : public SocketUser { @@ -136,7 +159,6 @@ friend class ubring::UBShmEndpoint; static void OnNewMessages(Socket* m); 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); @@ -148,6 +170,20 @@ friend class ubring::UBShmEndpoint; const uint64_t received_us, const uint64_t base_realtime, InputMessageClosure& last_msg); + static void QueueInputMessageBatch( + Socket* m, std::unique_ptr* batch, + int* num_bthread_created); + + static void QueueLastMessageOrBatch( + Socket* m, InputMessageClosure& last_msg, + std::unique_ptr* batch, + int* num_bthread_created, size_t batch_size); + + static uint32_t UpdateAdaptiveBatchSize( + uint32_t* messages_per_read_ema_q8, + uint32_t current_batch_size, + size_t parsed_message_count); + // User-supplied scissors and handlers. // the index of handler is exactly the same as the protocol InputMessageHandler* _handlers; diff --git a/src/brpc/rdma_transport.cpp b/src/brpc/rdma_transport.cpp index ee5151c3a5..4de3218ee2 100644 --- a/src/brpc/rdma_transport.cpp +++ b/src/brpc/rdma_transport.cpp @@ -183,6 +183,12 @@ void RdmaTransport::QueueMessage(InputMessageClosure& input_msg, } } +void RdmaTransport::QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) { + QueueInputMessageBatch( + input_msgs, num_bthread_created, rdma::FLAGS_rdma_disable_bthread); +} + void RdmaTransport::Debug(std::ostream &os) { if (_rdma_state == RDMA_ON && _rdma_ep) { _rdma_ep->DebugInfo(os); diff --git a/src/brpc/rdma_transport.h b/src/brpc/rdma_transport.h index 1d78fbb430..12babc370c 100644 --- a/src/brpc/rdma_transport.h +++ b/src/brpc/rdma_transport.h @@ -40,6 +40,8 @@ friend class rdma::RdmaHandshakeServerV3; int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, const timespec duetime) override; void ProcessEvent(bthread_attr_t attr) override; void QueueMessage(InputMessageClosure& inputMsg, int* num_bthread_created, bool last_msg) override; + void QueueMessages(InputMessageBatch* inputMsgs, + int* num_bthread_created) override; void Debug(std::ostream &os) override; rdma::RdmaEndpoint* GetRdmaEp() { CHECK(_rdma_ep != nullptr); @@ -64,4 +66,4 @@ friend class rdma::RdmaHandshakeServerV3; }; } // namespace brpc #endif // BRPC_WITH_RDMA -#endif //BRPC_RDMA_TRANSPORT_H \ No newline at end of file +#endif //BRPC_RDMA_TRANSPORT_H diff --git a/src/brpc/socket.cpp b/src/brpc/socket.cpp index 38727571ba..e7118c49e8 100644 --- a/src/brpc/socket.cpp +++ b/src/brpc/socket.cpp @@ -467,6 +467,8 @@ Socket::Socket(Forbidden f) , _hc_count(0) , _last_msg_size(0) , _avg_msg_size(0) + , _input_messages_per_read_ema_q8(0) + , _adaptive_input_message_batch_size(0) , _last_readtime_us(0) , _parsing_context(nullptr) , _correlation_id(0) @@ -571,9 +573,11 @@ void Socket::ReleaseAllFailedWriteRequests(Socket::WriteRequest* req) { } int Socket::ResetFileDescriptor(int fd) { - // Reset message sizes when fd is changed. + // Reset input heuristics when fd is changed. _last_msg_size = 0; _avg_msg_size = 0; + _input_messages_per_read_ema_q8 = 0; + _adaptive_input_message_batch_size = 0; // 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); @@ -2377,6 +2381,10 @@ void Socket::DebugSocket(std::ostream& os, SocketId id) { const int64_t cpuwide_now = butil::cpuwide_time_us(); os << "\nhc_count=" << ptr->_hc_count << "\navg_input_msg_size=" << ptr->_avg_msg_size + << "\navg_input_messages_per_read=" + << ((ptr->_input_messages_per_read_ema_q8 + 128) >> 8) + << "\nadaptive_input_message_batch_size=" + << ptr->_adaptive_input_message_batch_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() diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 6d321f8bc3..e9d25775d5 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -898,6 +898,11 @@ friend class TransportFactory; uint32_t _last_msg_size; // Average message size of last #MSG_SIZE_WINDOW messages (roughly) uint32_t _avg_msg_size; + // Q8 EMA of processable messages parsed in one read. Accessed only from + // the serialized input callback. + uint32_t _input_messages_per_read_ema_q8; + // 0 when adaptive batching is inactive, otherwise one of 1/2/4/8/16. + uint32_t _adaptive_input_message_batch_size; // Storing data read from `_fd' but cut-off yet. butil::IOPortal _read_buf; diff --git a/src/brpc/tcp_transport.cpp b/src/brpc/tcp_transport.cpp index 98fea81674..d93db5d721 100644 --- a/src/brpc/tcp_transport.cpp +++ b/src/brpc/tcp_transport.cpp @@ -103,4 +103,9 @@ void TcpTransport::QueueMessage(InputMessageClosure& input_msg, } } +void TcpTransport::QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) { + QueueInputMessageBatch(input_msgs, num_bthread_created, false); +} + } // namespace brpc diff --git a/src/brpc/tcp_transport.h b/src/brpc/tcp_transport.h index 8a06a85d37..cfe44a251c 100644 --- a/src/brpc/tcp_transport.h +++ b/src/brpc/tcp_transport.h @@ -34,8 +34,10 @@ class TcpTransport : public Transport { int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, timespec duetime) override; void ProcessEvent(bthread_attr_t attr) override; void QueueMessage(InputMessageClosure& input_msg, int* num_bthread_created, bool last_msg) override; + void QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) override; void Debug(std::ostream &os) override {} }; } // namespace brpc -#endif //BRPC_TCP_TRANSPORT_H \ No newline at end of file +#endif //BRPC_TCP_TRANSPORT_H diff --git a/src/brpc/transport.cpp b/src/brpc/transport.cpp new file mode 100644 index 0000000000..34e35bdd2b --- /dev/null +++ b/src/brpc/transport.cpp @@ -0,0 +1,51 @@ +// 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 "brpc/transport.h" + +namespace brpc { +DECLARE_bool(usercode_in_coroutine); +DECLARE_bool(usercode_in_pthread); + +void Transport::QueueInputMessageBatch(InputMessageBatch* input_msgs, + int* num_bthread_created, + bool run_inline) { + if (input_msgs == nullptr || input_msgs->empty()) { + delete input_msgs; + return; + } + if (run_inline || FLAGS_usercode_in_coroutine) { + ProcessInputMessageBatch(input_msgs); + return; + } + + bthread_t th; + bthread_attr_t tmp = + (FLAGS_usercode_in_pthread ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL) | + BTHREAD_NOSIGNAL; + tmp.keytable_pool = _socket->keytable_pool(); + tmp.tag = bthread_self_tag(); + bthread_attr_set_name(&tmp, "ProcessInputMessageBatch"); + if (bthread_start_background( + &th, &tmp, ProcessInputMessageBatch, input_msgs) == 0) { + ++*num_bthread_created; + } else { + ProcessInputMessageBatch(input_msgs); + } +} + +} // namespace brpc diff --git a/src/brpc/transport.h b/src/brpc/transport.h index edef24879c..592fbc934f 100644 --- a/src/brpc/transport.h +++ b/src/brpc/transport.h @@ -49,6 +49,8 @@ class Transport { virtual int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, timespec duetime) = 0; virtual void ProcessEvent(bthread_attr_t attr) = 0; virtual void QueueMessage(InputMessageClosure& input_msg, int* num_bthread_created, bool last_msg) = 0; + virtual void QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) = 0; virtual void Debug(std::ostream &os) = 0; bool HasOnEdgeTrigger() { @@ -58,9 +60,13 @@ class Transport { return _on_edge_trigger; } protected: + void QueueInputMessageBatch(InputMessageBatch* input_msgs, + int* num_bthread_created, + bool run_inline); + Socket* _socket; std::shared_ptr _default_connect; OnEdgeTrigger _on_edge_trigger; }; } -#endif //BRPC_TRANSPORT_H \ No newline at end of file +#endif //BRPC_TRANSPORT_H diff --git a/src/brpc/ubshm_transport.cpp b/src/brpc/ubshm_transport.cpp index df4eb36bed..3a084084e8 100644 --- a/src/brpc/ubshm_transport.cpp +++ b/src/brpc/ubshm_transport.cpp @@ -174,6 +174,12 @@ void UBShmTransport::QueueMessage(InputMessageClosure& input_msg, } } +void UBShmTransport::QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) { + QueueInputMessageBatch( + input_msgs, num_bthread_created, ubring::FLAGS_ub_disable_bthread); +} + void UBShmTransport::Debug(std::ostream &os) {} int UBShmTransport::ContextInitOrDie(bool serverOrNot, const void* _options) { @@ -232,4 +238,4 @@ bool UBShmTransport::OptionsAvailableOverUB(const ServerOptions* opt) { return true; } } // namespace brpc -#endif \ No newline at end of file +#endif diff --git a/src/brpc/ubshm_transport.h b/src/brpc/ubshm_transport.h index b3d1e7c518..4af1b143bf 100644 --- a/src/brpc/ubshm_transport.h +++ b/src/brpc/ubshm_transport.h @@ -37,6 +37,8 @@ friend class ubring::UBConnect; int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, const timespec duetime) override; void ProcessEvent(bthread_attr_t attr) override; void QueueMessage(InputMessageClosure& inputMsg, int* num_bthread_created, bool last_msg) override; + void QueueMessages(InputMessageBatch* input_msgs, + int* num_bthread_created) override; void Debug(std::ostream &os) override; ubring::UBShmEndpoint* GetUBShmEp() { CHECK(_ub_ep != nullptr); @@ -61,4 +63,4 @@ friend class ubring::UBConnect; }; } // namespace brpc #endif // BRPC_WITH_UBRING -#endif //BRPC_UB_TRANSPORT_H \ No newline at end of file +#endif //BRPC_UB_TRANSPORT_H diff --git a/test/brpc_input_messenger_unittest.cpp b/test/brpc_input_messenger_unittest.cpp index fcf233b154..d97867d945 100644 --- a/test/brpc_input_messenger_unittest.cpp +++ b/test/brpc_input_messenger_unittest.cpp @@ -22,6 +22,16 @@ #include #include #include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include #include "gperftools_helper.h" #include "butil/time.h" @@ -29,8 +39,144 @@ #include "butil/fd_utility.h" #include "butil/fd_guard.h" #include "butil/unix_socket.h" +#include "bthread/unstable.h" #include "brpc/acceptor.h" +#include "brpc/input_messenger.h" #include "brpc/policy/hulu_pbrpc_protocol.h" +#include "brpc/transport.h" + +namespace brpc { +DECLARE_bool(usercode_in_coroutine); +DECLARE_int32(input_message_batch_process_size); +} + +namespace { + +struct BatchRecorder { + void Record(int value) { + std::lock_guard lock(mutex); + values.push_back(value); + condition.notify_all(); + } + + void RecordDestroy() { + destroyed.fetch_add(1, std::memory_order_relaxed); + condition.notify_all(); + } + + bool WaitForSize(size_t expected) { + std::unique_lock lock(mutex); + return condition.wait_for( + lock, std::chrono::seconds(5), + [this, expected] { return values.size() >= expected; }); + } + + bool WaitForDestroyed(int expected) { + std::unique_lock lock(mutex); + return condition.wait_for( + lock, std::chrono::seconds(5), + [this, expected] { + return destroyed.load(std::memory_order_relaxed) >= expected; + }); + } + + std::vector Snapshot() { + std::lock_guard lock(mutex); + return values; + } + + std::mutex mutex; + std::condition_variable condition; + std::vector values; + std::atomic destroyed{0}; +}; + +class BatchTestMessage : public brpc::InputMessageBase { +public: + BatchTestMessage(int value, BatchRecorder* recorder, bool throw_on_process) + : value(value) + , recorder(recorder) + , throw_on_process(throw_on_process) {} + + int value; + BatchRecorder* recorder; + bool throw_on_process; + +private: + void DestroyImpl() override { + recorder->RecordDestroy(); + delete this; + } +}; + +void ProcessBatchTestMessage(brpc::InputMessageBase* msg_base) { + brpc::DestroyingPtr guard(msg_base); + BatchTestMessage* msg = static_cast(msg_base); + msg->recorder->Record(msg->value); + if (msg->throw_on_process) { + throw std::runtime_error("injected handler failure"); + } +} + +BatchTestMessage* NewBatchTestMessage( + int value, BatchRecorder* recorder, bool throw_on_process = false) { + BatchTestMessage* msg = + new BatchTestMessage(value, recorder, throw_on_process); + msg->_process = ProcessBatchTestMessage; + return msg; +} + +brpc::ParseResult ParseBatchTestMessage( + butil::IOBuf* source, brpc::Socket*, bool, const void* arg) { + if (source->empty()) { + return brpc::MakeParseError(brpc::PARSE_ERROR_NOT_ENOUGH_DATA); + } + char value = '\0'; + source->copy_to(&value, 1); + source->pop_front(1); + return brpc::MakeMessage(new BatchTestMessage( + static_cast(value), + static_cast(const_cast(arg)), false)); +} + +brpc::InputMessageHandler MakeBatchTestHandler(BatchRecorder* recorder) { + const brpc::InputMessageHandler handler = { + ParseBatchTestMessage, + ProcessBatchTestMessage, + nullptr, + recorder, + "batch_test", + }; + return handler; +} + +brpc::SocketUniquePtr CreateBatchTestSocket( + brpc::InputMessenger* messenger, brpc::SocketId* id) { + brpc::SocketOptions options; + options.socket_mode = brpc::SOCKET_MODE_TCP; + EXPECT_EQ(0, messenger->Create(options, id)); + brpc::SocketUniquePtr socket; + EXPECT_EQ(0, brpc::Socket::Address(*id, &socket)); + return socket; +} + +class InputBatchFlagGuard { +public: + InputBatchFlagGuard() + : batch_size(brpc::FLAGS_input_message_batch_process_size) + , usercode_in_coroutine(brpc::FLAGS_usercode_in_coroutine) {} + + ~InputBatchFlagGuard() { + brpc::FLAGS_input_message_batch_process_size = batch_size; + brpc::FLAGS_usercode_in_coroutine = usercode_in_coroutine; + } + +private: + int batch_size; + bool usercode_in_coroutine; +}; + +} // namespace void EmptyProcessHuluRequest(brpc::InputMessageBase* msg_base) { brpc::DestroyingPtr a(msg_base); @@ -60,6 +206,208 @@ class MessengerTest : public ::testing::Test{ }; }; +TEST_F(MessengerTest, input_message_batch_runs_in_order_once) { + BatchRecorder recorder; + brpc::InputMessageBatch batch(8); + batch.add(NewBatchTestMessage(1, &recorder)); + batch.add(nullptr); + batch.add(NewBatchTestMessage(2, &recorder)); + batch.add(NewBatchTestMessage(3, &recorder)); + ASSERT_EQ(3u, batch.size()); + + batch.Run(); + EXPECT_TRUE(batch.empty()); + EXPECT_EQ((std::vector{1, 2, 3}), recorder.Snapshot()); + EXPECT_EQ(3, recorder.destroyed.load()); + + batch.Run(); + EXPECT_EQ((std::vector{1, 2, 3}), recorder.Snapshot()); + EXPECT_EQ(3, recorder.destroyed.load()); +} + +TEST_F(MessengerTest, input_message_batch_destructor_contains_exceptions) { + static_assert( + std::is_nothrow_destructible::value, + "InputMessageBatch must be nothrow destructible"); + + BatchRecorder recorder; + EXPECT_NO_THROW({ + brpc::InputMessageBatch batch(2); + batch.add(NewBatchTestMessage(1, &recorder, true)); + batch.add(NewBatchTestMessage(2, &recorder)); + }); + EXPECT_EQ((std::vector{1}), recorder.Snapshot()); + EXPECT_EQ(2, recorder.destroyed.load()); + + BatchRecorder worker_recorder; + brpc::InputMessageBatch* batch = new brpc::InputMessageBatch(2); + batch->add(NewBatchTestMessage(1, &worker_recorder, true)); + batch->add(NewBatchTestMessage(2, &worker_recorder)); + EXPECT_NO_THROW(brpc::ProcessInputMessageBatch(batch)); + EXPECT_EQ((std::vector{1, 2}), worker_recorder.Snapshot()); + EXPECT_EQ(2, worker_recorder.destroyed.load()); +} + +TEST_F(MessengerTest, input_message_batch_flag_validation) { + InputBatchFlagGuard flag_guard; + EXPECT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption( + "input_message_batch_process_size", "-1").empty()); + EXPECT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption( + "input_message_batch_process_size", "0").empty()); + EXPECT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption( + "input_message_batch_process_size", "1").empty()); + EXPECT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption( + "input_message_batch_process_size", "8").empty()); + EXPECT_TRUE(GFLAGS_NAMESPACE::SetCommandLineOption( + "input_message_batch_process_size", "-2").empty()); +} + +TEST_F(MessengerTest, adaptive_input_message_batch_rises_falls_and_caps) { + uint32_t ema_q8 = 256; + uint32_t batch_size = 1; + std::vector rising_levels(1, batch_size); + for (int i = 0; i < 32 && batch_size < 16; ++i) { + const uint32_t old_batch_size = batch_size; + batch_size = brpc::InputMessenger::UpdateAdaptiveBatchSize( + &ema_q8, batch_size, std::numeric_limits::max()); + if (batch_size != old_batch_size) { + rising_levels.push_back(batch_size); + } + } + EXPECT_EQ((std::vector{1, 2, 4, 8, 16}), rising_levels); + EXPECT_LE(ema_q8, 32u * 256); + + std::vector falling_levels(1, batch_size); + for (int i = 0; i < 32 && batch_size > 1; ++i) { + const uint32_t old_batch_size = batch_size; + batch_size = brpc::InputMessenger::UpdateAdaptiveBatchSize( + &ema_q8, batch_size, 1); + if (batch_size != old_batch_size) { + falling_levels.push_back(batch_size); + } + } + EXPECT_EQ((std::vector{16, 8, 4, 2, 1}), falling_levels); + + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(brpc::SocketOptions(), &id)); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(id, &socket)); + socket->_input_messages_per_read_ema_q8 = ema_q8; + socket->_adaptive_input_message_batch_size = batch_size; + ASSERT_EQ(0, socket->ResetFileDescriptor(-1)); + EXPECT_EQ(0u, socket->_input_messages_per_read_ema_q8); + EXPECT_EQ(0u, socket->_adaptive_input_message_batch_size); +} + +TEST_F(MessengerTest, batching_consumes_the_last_message) { + InputBatchFlagGuard flag_guard; + brpc::FLAGS_input_message_batch_process_size = 8; + brpc::FLAGS_usercode_in_coroutine = false; + + BatchRecorder recorder; + brpc::InputMessenger messenger(4); + ASSERT_EQ(0, messenger.AddNonProtocolHandler( + MakeBatchTestHandler(&recorder))); + brpc::SocketId id; + brpc::SocketUniquePtr socket = CreateBatchTestSocket(&messenger, &id); + ASSERT_TRUE(socket); + socket->_read_buf.append("abcd", 4); + + brpc::InputMessageClosure last_msg; + ASSERT_EQ(0, messenger.ProcessNewMessage( + socket.get(), 4, false, 123, 456, last_msg)); + brpc::DestroyingPtr leftover(last_msg.release()); + EXPECT_EQ(nullptr, leftover.get()); + ASSERT_TRUE(recorder.WaitForSize(4)); + EXPECT_EQ((std::vector{'a', 'b', 'c', 'd'}), recorder.Snapshot()); + + socket->SetFailed(); + ASSERT_TRUE(recorder.WaitForDestroyed(4)); +} + +TEST_F(MessengerTest, disabled_and_progressive_paths_remain_individual) { + InputBatchFlagGuard flag_guard; + brpc::FLAGS_usercode_in_coroutine = false; + + { + brpc::FLAGS_input_message_batch_process_size = 0; + BatchRecorder recorder; + brpc::InputMessenger messenger(4); + ASSERT_EQ(0, messenger.AddNonProtocolHandler( + MakeBatchTestHandler(&recorder))); + brpc::SocketId id; + brpc::SocketUniquePtr socket = CreateBatchTestSocket(&messenger, &id); + ASSERT_TRUE(socket); + socket->_read_buf.append("ab", 2); + brpc::InputMessageClosure last_msg; + ASSERT_EQ(0, messenger.ProcessNewMessage( + socket.get(), 2, false, 123, 456, last_msg)); + brpc::DestroyingPtr leftover(last_msg.release()); + EXPECT_NE(nullptr, leftover.get()); + leftover.reset(); + ASSERT_TRUE(recorder.WaitForDestroyed(2)); + socket->SetFailed(); + } + + { + brpc::FLAGS_input_message_batch_process_size = 8; + BatchRecorder recorder; + brpc::InputMessenger messenger(4); + ASSERT_EQ(0, messenger.AddNonProtocolHandler( + MakeBatchTestHandler(&recorder))); + brpc::SocketId id; + brpc::SocketUniquePtr socket = CreateBatchTestSocket(&messenger, &id); + ASSERT_TRUE(socket); + socket->read_will_be_progressive(brpc::CONNECTION_TYPE_SINGLE); + socket->_read_buf.append("abc", 3); + brpc::InputMessageClosure last_msg; + ASSERT_EQ(0, messenger.ProcessNewMessage( + socket.get(), 3, false, 123, 456, last_msg)); + EXPECT_EQ(nullptr, last_msg.release()); + ASSERT_TRUE(recorder.WaitForSize(3)); + std::vector values = recorder.Snapshot(); + std::sort(values.begin(), values.end()); + EXPECT_EQ((std::vector{'a', 'b', 'c'}), values); + socket->SetFailed(); + ASSERT_TRUE(recorder.WaitForDestroyed(3)); + } +} + +TEST_F(MessengerTest, transport_batch_helper_handles_inline_and_empty_batches) { + InputBatchFlagGuard flag_guard; + brpc::FLAGS_usercode_in_coroutine = true; + + brpc::InputMessenger messenger; + brpc::SocketId id; + brpc::SocketUniquePtr socket = CreateBatchTestSocket(&messenger, &id); + ASSERT_TRUE(socket); + + BatchRecorder recorder; + int num_bthread_created = 0; + brpc::InputMessageBatch* batch = new brpc::InputMessageBatch(2); + batch->add(NewBatchTestMessage(1, &recorder)); + batch->add(NewBatchTestMessage(2, &recorder)); + socket->_transport->QueueMessages(batch, &num_bthread_created); + EXPECT_EQ((std::vector{1, 2}), recorder.Snapshot()); + EXPECT_EQ(0, num_bthread_created); + + socket->_transport->QueueMessages( + new brpc::InputMessageBatch, &num_bthread_created); + EXPECT_EQ(0, num_bthread_created); + + brpc::FLAGS_usercode_in_coroutine = false; + batch = new brpc::InputMessageBatch(2); + batch->add(NewBatchTestMessage(3, &recorder)); + batch->add(NewBatchTestMessage(4, &recorder)); + socket->_transport->QueueMessages(batch, &num_bthread_created); + EXPECT_EQ(1, num_bthread_created); + bthread_flush(); + ASSERT_TRUE(recorder.WaitForSize(4)); + EXPECT_EQ((std::vector{1, 2, 3, 4}), recorder.Snapshot()); + + socket->SetFailed(); +} + #define USE_UNIX_DOMAIN_SOCKET 1 const size_t NEPOLL = 1;