From f318a68ecbdca11f7980c08b30c0c64994008b46 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Thu, 8 Oct 2026 23:58:52 -0700 Subject: [PATCH 01/10] Unify transport handshakes on current upstream master --- BUILD.bazel | 1 - CMakeLists.txt | 2 +- Makefile | 2 +- docs/cn/handshake_common_design.md | 52 + src/brpc/adapter_transport.cpp | 694 +++++++++ src/brpc/adapter_transport.h | 103 ++ src/brpc/global.cpp | 14 +- src/brpc/handshake/handshake_adapter.cpp | 61 + src/brpc/handshake/handshake_adapter.h | 73 + src/brpc/handshake/handshake_frame.cpp | 235 +++ src/brpc/handshake/handshake_frame.h | 92 ++ src/brpc/handshake/handshake_io.cpp | 162 ++ src/brpc/handshake/handshake_io.h | 89 ++ src/brpc/handshake/rdma_handshake.cpp | 550 +++++++ src/brpc/handshake/rdma_handshake.h | 160 ++ .../rdma_handshake_constants.h | 31 +- src/brpc/handshake/ubshm_handshake.cpp | 454 ++++++ src/brpc/handshake/ubshm_handshake.h | 126 ++ src/brpc/input_messenger.cpp | 8 + src/brpc/input_messenger.h | 4 + src/brpc/policy/rdma_handshake_protocol.cpp | 14 +- src/brpc/policy/rdma_handshake_protocol.h | 27 +- .../policy/transport_handshake_protocol.cpp | 43 + .../policy/transport_handshake_protocol.h | 44 + src/brpc/rdma/rdma_endpoint.cpp | 835 ++-------- src/brpc/rdma/rdma_endpoint.h | 178 +-- src/brpc/rdma/rdma_handshake.cpp | 514 ------ src/brpc/rdma/rdma_handshake.h | 185 --- src/brpc/rdma/rdma_handshake_server.cpp | 219 --- src/brpc/rdma/rdma_handshake_server.h | 48 - src/brpc/{rdma => }/rdma_handshake.proto | 0 src/brpc/rdma_transport.cpp | 123 +- src/brpc/rdma_transport.h | 35 +- src/brpc/socket.h | 23 +- src/brpc/transport_factory.cpp | 7 +- src/brpc/transport_factory.h | 4 +- src/brpc/transport_handshake.cpp | 425 +++++ src/brpc/transport_handshake.h | 233 +++ src/brpc/ubshm/ub_endpoint.cpp | 753 +-------- src/brpc/ubshm/ub_endpoint.h | 157 +- src/brpc/ubshm_transport.cpp | 115 +- src/brpc/ubshm_transport.h | 22 +- test/brpc_rdma_unittest.cpp | 1376 ++++++++++------- test/brpc_transport_handshake_unittest.cpp | 960 ++++++++++++ test/brpc_ubring_unittest.cpp | 503 +++++- 45 files changed, 6449 insertions(+), 3307 deletions(-) create mode 100644 docs/cn/handshake_common_design.md create mode 100644 src/brpc/adapter_transport.cpp create mode 100644 src/brpc/adapter_transport.h create mode 100644 src/brpc/handshake/handshake_adapter.cpp create mode 100644 src/brpc/handshake/handshake_adapter.h create mode 100644 src/brpc/handshake/handshake_frame.cpp create mode 100644 src/brpc/handshake/handshake_frame.h create mode 100644 src/brpc/handshake/handshake_io.cpp create mode 100644 src/brpc/handshake/handshake_io.h create mode 100644 src/brpc/handshake/rdma_handshake.cpp create mode 100644 src/brpc/handshake/rdma_handshake.h rename src/brpc/{rdma => handshake}/rdma_handshake_constants.h (69%) create mode 100644 src/brpc/handshake/ubshm_handshake.cpp create mode 100644 src/brpc/handshake/ubshm_handshake.h create mode 100644 src/brpc/policy/transport_handshake_protocol.cpp create mode 100644 src/brpc/policy/transport_handshake_protocol.h delete mode 100644 src/brpc/rdma/rdma_handshake.cpp delete mode 100644 src/brpc/rdma/rdma_handshake.h delete mode 100644 src/brpc/rdma/rdma_handshake_server.cpp delete mode 100644 src/brpc/rdma/rdma_handshake_server.h rename src/brpc/{rdma => }/rdma_handshake.proto (100%) create mode 100644 src/brpc/transport_handshake.cpp create mode 100644 src/brpc/transport_handshake.h create mode 100644 test/brpc_transport_handshake_unittest.cpp diff --git a/BUILD.bazel b/BUILD.bazel index 23e433c76e..46ff01b63c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -528,7 +528,6 @@ filegroup( srcs = glob([ "src/brpc/*.proto", "src/brpc/policy/*.proto", - "src/brpc/rdma/*.proto", "src/brpc/urma/*.proto", ]), visibility = ["//visibility:public"], diff --git a/CMakeLists.txt b/CMakeLists.txt index e98e0c4fa0..ac83f1bc1e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -719,7 +719,7 @@ set(PROTO_FILES idl_options.proto brpc/trackme.proto brpc/streaming_rpc_meta.proto brpc/proto_base.proto - brpc/rdma/rdma_handshake.proto + brpc/rdma_handshake.proto brpc/urma/urma_handshake.proto) file(MAKE_DIRECTORY ${PROJECT_BINARY_DIR}/output/include/brpc) set(PROTOC_FLAGS ${PROTOC_FLAGS} -I${PROTOBUF_INCLUDE_DIR}) diff --git a/Makefile b/Makefile index 271b518ae6..dd33e0dc6a 100644 --- a/Makefile +++ b/Makefile @@ -203,7 +203,7 @@ JSON2PB_DIRS = src/json2pb JSON2PB_SOURCES = $(foreach d,$(JSON2PB_DIRS),$(wildcard $(addprefix $(d)/*,$(SRCEXTS)))) JSON2PB_OBJS = $(addsuffix .o, $(basename $(JSON2PB_SOURCES))) -BRPC_DIRS = src/brpc src/brpc/details src/brpc/builtin src/brpc/policy src/brpc/policy/mysql src/brpc/rdma +BRPC_DIRS = src/brpc src/brpc/details src/brpc/builtin src/brpc/handshake src/brpc/policy src/brpc/policy/mysql src/brpc/rdma ifeq ($(WITH_URMA),1) BRPC_DIRS += src/brpc/urma endif diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md new file mode 100644 index 0000000000..ca53ae53ec --- /dev/null +++ b/docs/cn/handshake_common_design.md @@ -0,0 +1,52 @@ +# RDMA 与 UBSHM 公共握手 + +## 范围 + +本次重构将 RDMA 和 UBSHM 的握手流程收敛到公共层。`Socket` 持有 +`AdapterTransport`；后者以 TCP 作为握手控制通道，在升级成功后选择高速 +数据面，协商不可用时继续使用 TCP。URMA 目前仍使用独立的 `UrmaTransport` +和握手实现，不在本次迁移范围内。 + +## 分工 + +| 组件 | 职责 | +| --- | --- | +| `AdapterTransport` | 持有 TCP 与候选高速 Transport，启动客户端握手、分发服务端输入，并选择数据面 | +| `HandshakeSession` | 执行公共客户端和服务端状态机，管理帧收发、增量解析和终态发布 | +| `HandshakeProtocol` | 定义 hello、ACK 和扩展帧，编解码协议字段并保存协议状态 | +| `HandshakeTransport` | 准备和协商资源，处理升级成功、TCP 回退及失败清理 | +| RDMA / UBSHM protocol | 实现各自的 wire format、版本兼容和字段校验 | +| `RdmaEndpoint` / `UBShmEndpoint` | 管理各自的资源和数据面，不驱动完整握手流程 | + +客户端在 TCP 连接建立后启动握手 bthread；服务端通过 `InputMessenger` +增量解析 hello、扩展字段（若协议需要）和 ACK。`AdapterTransport` 选择具体的 +protocol 和 transport participant，`HandshakeSession` 只负责调用参与者并编排流程。 +服务端在 `EXTENSION_WAIT` 和 `ACK_WAIT` 保留解析上下文，并通过协议版本 +恢复选中的 protocol，因为扩展帧和 ACK 本身不带 magic。 + +## 状态与回退 + +握手终态为 `ESTABLISHED`、`FALLBACK_TCP` 或 `FAILED`。升级成功后，TCP +仍是控制连接，业务数据走 RDMA/UBSHM；控制连接上出现额外业务数据视为 +协议错误。升级完成时停止当前 TCP 业务解析循环，由控制连接处理额外输入。 +资源不可用或协商拒绝时，释放本次升级资源，设置 TCP 为数据面， +再以 release 顺序发布 `FALLBACK_TCP`。事件线程以 acquire 顺序读取终态， +继续解析已缓存及后续的 TCP 业务数据。已确认属于握手协议的畸形帧则失败， +不作为普通 RPC 数据回放。 + +增量解析中，不完整帧返回 `NOT_ENOUGH_DATA` 且不消费输入；不可能匹配 +握手 magic 的前缀返回 `TRY_OTHERS`。TCP 分片、握手帧与后续数据粘连都由 +帧边界处理，不能假设一次 read 恰好得到一帧。 + +## 协议兼容 + +- RDMA 保持既有 v2/v3 wire format、协议 ID 和注册名称。 +- UBSHM 使用 v3 hello。64 字节 hello 后，双方交换 4 字节网络序格式扩展； + 只有双方选择 `LEGACY_64` 才能升级，否则回退 TCP。ACK 仍为 4 字节。 +- 公共 framing 仅负责帧边界，协议字段由各自 `HandshakeProtocol` 解释。 + +## 验证重点 + +测试覆盖分片和粘连输入、magic 前缀分流、协议版本与格式校验、TCP 回退、 +资源释放及终态发布，并分别运行 RDMA 与 UBSHM 回归测试。URMA 的公共握手 +迁移需要单独设计和测试，不应仅凭本次重构推断其已完成。 diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp new file mode 100644 index 0000000000..0ed26406db --- /dev/null +++ b/src/brpc/adapter_transport.cpp @@ -0,0 +1,694 @@ +// 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/adapter_transport.h" + +#include +#include +#include +#include +#include + +#include "brpc/input_messenger.h" +#include "brpc/destroyable.h" +#include "brpc/handshake/rdma_handshake.h" +#include "brpc/handshake/ubshm_handshake.h" +#if BRPC_WITH_RDMA +#include "brpc/rdma/rdma_helper.h" +#endif +#if BRPC_WITH_UBRING +#include "brpc/ubshm/ub_endpoint.h" +#include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/ubr_trx.h" +#endif +#include "brpc/rdma_transport.h" +#include "brpc/tcp_transport.h" +#include "brpc/ubshm_transport.h" + +namespace brpc { + +namespace { + +bool MatchesMagicPrefix(const char *prefix, size_t prefix_len, + const char *magic, size_t magic_len) { + const size_t compare_len = std::min(prefix_len, magic_len); + return memcmp(prefix, magic, compare_len) == 0; +} + +class AdapterConnect : public AppConnect { +public: + explicit AdapterConnect(const std::shared_ptr& app_connect) + : _app_connect(app_connect) {} + + static std::shared_ptr Wrap( + const std::shared_ptr& app_connect) { + if (std::dynamic_pointer_cast(app_connect)) { + return app_connect; + } + return std::make_shared(app_connect); + } + + static std::shared_ptr Unwrap( + const std::shared_ptr& app_connect) { + const std::shared_ptr adapter = + std::dynamic_pointer_cast(app_connect); + return adapter ? adapter->_app_connect : app_connect; + } + + void StartConnect(const Socket* socket, + void (*done)(int, void*), void* data) override { + ApplicationConnectTask* task = new ApplicationConnectTask{ + socket, _app_connect, done, data}; + if (AdapterTransport::StartClientUpgrade( + socket, OnUpgradeComplete, task) != 0) { + AdapterTransport::Get(const_cast(socket))->CompleteConnection( + handshake::FAILED); + const int error = errno != 0 ? errno : EAGAIN; + delete task; + done(error, data); + } + } + + void StopConnect(Socket*) override {} + +private: + struct ApplicationConnectTask { + const Socket* socket; + std::shared_ptr app_connect; + void (*done)(int, void*); + void* data; + }; + + static void OnApplicationComplete(int error, void* arg) { + std::unique_ptr task( + static_cast(arg)); + task->done(error, task->data); + } + + static void OnUpgradeComplete(int error, void* arg) { + ApplicationConnectTask* task = + static_cast(arg); + if (error != 0 || !task->app_connect) { + std::unique_ptr owned(task); + task->done(error, task->data); + return; + } + task->app_connect->StartConnect( + task->socket, OnApplicationComplete, task); + } + + std::shared_ptr _app_connect; +}; + +struct ClientHandshakeTask { + AdapterTransport* adapter; + void (*done)(int, void*); + void* data; + SocketUniquePtr socket; +}; + +#if BRPC_WITH_RDMA +class RdmaClientHandshakeTransport : public handshake::HandshakeTransport { +public: + RdmaClientHandshakeTransport( + RdmaTransport* transport, rdma::RdmaHandshakeAdapter* protocol, + Socket* socket, int* connect_error) + : _transport(transport), _protocol(protocol), _socket(socket), + _connect_error(connect_error) {} + + handshake::StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources() == 0) { + return handshake::STEP_OK; + } + errno = 0; + return handshake::STEP_FALLBACK; + } + + handshake::StepResult NegotiateResources() override { + return _transport->NegotiateUpgradeResources( + _protocol->remote(), false) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + + void OnFallback() override { _transport->DeactivateUpgrade(); } + + void OnFailed() override { + _transport->DeactivateUpgrade(); + const int saved_errno = errno != 0 ? errno : EPROTO; + *_connect_error = saved_errno; + _socket->SetFailed(saved_errno, + "Fail to complete rdma handshake from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + } + +private: + RdmaTransport* _transport; + rdma::RdmaHandshakeAdapter* _protocol; + Socket* _socket; + int* _connect_error; +}; +#endif + +#if BRPC_WITH_UBRING +class UBShmClientHandshakeTransport : public handshake::HandshakeTransport { +public: + UBShmClientHandshakeTransport( + UBShmTransport* transport, ubring::SHM* local_shm, + const std::string& shm_name, ubring::UBShmHandshakeAdapter* protocol, + Socket* socket, int* connect_error) + : _transport(transport), _local_shm(local_shm), + _shm_name(shm_name), _protocol(protocol), _socket(socket), + _connect_error(connect_error) {} + + handshake::StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources( + _local_shm, _shm_name.c_str()) != 0) { + return handshake::STEP_FALLBACK; + } + _protocol->ConfigureClientHello(*_local_shm); + return handshake::STEP_OK; + } + + handshake::StepResult NegotiateResources() override { + return _transport->NegotiateUpgradeResources( + _local_shm, _shm_name.c_str()) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + + void OnFailed() override { + _transport->DeactivateUpgrade(); + const int saved_errno = errno != 0 ? errno : EPROTO; + *_connect_error = saved_errno; + _socket->SetFailed(saved_errno, + "Fail to complete ubring handshake from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + } + +private: + UBShmTransport* _transport; + ubring::SHM* _local_shm; + std::string _shm_name; + ubring::UBShmHandshakeAdapter* _protocol; + Socket* _socket; + int* _connect_error; +}; +#endif + +} // namespace + +AdapterTransport::AdapterTransport(SocketMode mode) + : _mode(mode), _connection_completed(0) {} +AdapterTransport::~AdapterTransport() = default; + +AdapterTransport* AdapterTransport::Get(Socket* socket) { + CHECK(socket != NULL); + return static_cast(socket->_transport.get()); +} + +const AdapterTransport* AdapterTransport::Get(const Socket* socket) { + CHECK(socket != NULL); + return static_cast(socket->_transport.get()); +} + +bool AdapterTransport::upgrade_capable(SocketMode mode) const { + if (_mode != mode || _high_speed_transport == NULL) { + return false; + } + switch (mode) { +#if BRPC_WITH_RDMA + case SOCKET_MODE_RDMA: + return static_cast( + _high_speed_transport.get())->UpgradeReady(); +#endif +#if BRPC_WITH_UBRING + case SOCKET_MODE_UBRING: + return static_cast( + _high_speed_transport.get())->UpgradeReady(); +#endif + default: + return false; + } +} + +int AdapterTransport::StartClientUpgrade(const Socket* socket, + void (*done)(int, void*), + void* data) { + AdapterTransport* adapter = Get(const_cast(socket)); + ClientHandshakeTask* task = new ClientHandshakeTask{adapter, done, data, SocketUniquePtr()}; + if (Socket::Address(socket->id(), &task->socket) != 0) { + delete task; + return -1; + } + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + bthread_attr_set_name(&attr, "StartClientUpgrade"); + if (bthread_start_background(&tid, &attr, + ProcessClientHandshake, task) < 0) { + delete task; + return -1; + } + return 0; +} + +ParseResult AdapterTransport::ProcessUpgradeReadable(butil::IOBuf* source) { + ParseResult result(PARSE_ERROR_NOT_ENOUGH_DATA); + if (_socket->parsing_context() != NULL) { + handshake::ServerHandshakeContext* context = + static_cast( + _socket->parsing_context()); + CHECK(context->adapter() != NULL); + result = context->adapter()->ExecuteServerHandshake(source, _socket); + } else if (!source->empty()) { + static const size_t MAX_MAGIC_LEN = 4; + char prefix[MAX_MAGIC_LEN] = {}; + const size_t prefix_len = std::min(source->size(), MAX_MAGIC_LEN); + source->copy_to(prefix, prefix_len); + + const bool matches_ub = + MatchesMagicPrefix(prefix, prefix_len, "UB", 2); + const bool matches_rdma = + MatchesMagicPrefix(prefix, prefix_len, "RDMA", 4) || + MatchesMagicPrefix(prefix, prefix_len, "RDM3", 4); + if (!matches_ub && !matches_rdma) { + result = ParseResult(PARSE_ERROR_TRY_OTHERS); + } else { + handshake::HandshakeAdapter* adapter = + matches_ub + ? handshake::GetUBShmServerHandshakeAdapter() + : handshake::GetRdmaServerHandshakeAdapter(); + result = adapter->ExecuteServerHandshake(source, _socket); + } + } + const int phase = _handshake.phase(butil::memory_order_acquire); + if (!connection_completed() && + (phase == handshake::ESTABLISHED || + phase == handshake::FALLBACK_TCP || phase == handshake::FAILED)) { + CompleteConnection(static_cast(phase)); + } + return result; +} + +void AdapterTransport::CompleteConnection(handshake::Phase terminal_phase) { + CHECK(terminal_phase == handshake::ESTABLISHED || + terminal_phase == handshake::FALLBACK_TCP || + terminal_phase == handshake::FAILED); + if (terminal_phase == handshake::FAILED && + _handshake.phase(butil::memory_order_acquire) != handshake::FAILED) { + _handshake.MarkFailed(); + } + int expected = 0; + _connection_completed.compare_exchange_strong( + expected, 1, butil::memory_order_release, + butil::memory_order_relaxed); +} + +void* AdapterTransport::ProcessClientHandshake(void* arg) { + std::unique_ptr task( + static_cast(arg)); + AdapterTransport* adapter = task->adapter; + Socket* socket = task->socket.get(); + int connect_error = 0; + (void)connect_error; + +#if BRPC_WITH_RDMA + if (adapter->_mode == SOCKET_MODE_RDMA) { + RdmaTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + if (!rdma::IsRdmaAvailable()) { + adapter->FallbackToTcp(); + adapter->CompleteConnection(handshake::FALLBACK_TCP); + task->done(0, task->data); + return NULL; + } + + std::unique_ptr protocol = + transport->CreateClientHandshakeAdapter(); + CHECK(protocol != NULL); + RdmaClientHandshakeTransport participant( + transport, protocol.get(), socket, &connect_error); + const handshake::StepResult result = adapter->_handshake.RunClient( + protocol.get(), &participant); + if (result == handshake::STEP_OK && + transport->StartUpgradeEvents() < 0) { + const int saved_errno = errno != 0 ? errno : ERDMA; + transport->DeactivateUpgrade(); + adapter->_handshake.MarkFailed(); + socket->SetFailed( + saved_errno, + "Fail to start RDMA CQ events from %s: %s", + socket->description().c_str(), berror(saved_errno)); + connect_error = saved_errno; + } + if (result == handshake::STEP_ERROR && connect_error == 0) { + connect_error = errno != 0 ? errno : EPROTO; + } + adapter->CompleteConnection(static_cast( + adapter->_handshake.phase(butil::memory_order_acquire))); + task->done(connect_error, task->data); + return NULL; + } +#endif + +#if BRPC_WITH_UBRING + if (adapter->_mode == SOCKET_MODE_UBRING) { + UBShmTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + if (!ubring::IsUBAvailable()) { + adapter->FallbackToTcp(); + adapter->CompleteConnection(handshake::FALLBACK_TCP); + task->done(0, task->data); + return NULL; + } + + const size_t local_shm_len = + static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + ubring::SHM local_trx_shm = { + NULL, local_shm_len, 0, {0}, static_cast(socket->fd())}; + const auto shm_name_str = + butil::endpoint2str(socket->local_side()); + ubring::UBShmHandshakeAdapter wire; + UBShmClientHandshakeTransport participant( + transport, &local_trx_shm, shm_name_str.c_str(), &wire, socket, + &connect_error); + const handshake::StepResult result = adapter->_handshake.RunClient( + &wire, &participant); + if (result == handshake::STEP_OK) { + transport->GetUBShmEp()->SetNegotiatedDataFormat( + ubring::UBR_DATA_FORMAT_LEGACY_64); + transport->FinishUpgrade(); + transport->GetUBShmEp()->StartReceiveEvents(); + } + if (result == handshake::STEP_ERROR && connect_error == 0) { + connect_error = errno != 0 ? errno : EPROTO; + } + adapter->CompleteConnection(static_cast( + adapter->_handshake.phase(butil::memory_order_acquire))); + task->done(connect_error, task->data); + return NULL; + } +#endif + + socket->SetFailed(EPROTO, "Unsupported client transport handshake"); + adapter->CompleteConnection(handshake::FAILED); + task->done(EPROTO, task->data); + return NULL; +} + +void AdapterTransport::Init(Socket* socket, const SocketOptions& options) { + CHECK_EQ(_mode, options.socket_mode); + _socket = socket; + _default_connect = options.app_connect; + _on_edge_trigger = options.on_edge_triggered_events; + if (options.need_on_edge_trigger && _on_edge_trigger == NULL) { + if (_mode == SOCKET_MODE_TCP) { + _on_edge_trigger = InputMessenger::OnNewMessages; +#if BRPC_WITH_RDMA + } else if (_mode == SOCKET_MODE_RDMA && + options.user != static_cast( + get_client_side_messenger())) { + // RDMA server handshake is parsed by InputMessenger. + _on_edge_trigger = OnNewMessagesAfterUpgrade; +#endif +#if BRPC_WITH_UBRING + } else if (_mode == SOCKET_MODE_UBRING && + options.user != static_cast( + get_client_side_messenger())) { + // UBSHM server handshake is parsed by InputMessenger. + _on_edge_trigger = OnNewMessagesAfterUpgrade; +#endif + } else { + _on_edge_trigger = OnNewDataFromTcp; + } + } + _handshake.Reset(socket); + _connection_completed.store(0, butil::memory_order_relaxed); + _tcp_transport.reset(new TcpTransport); + _tcp_transport->Init(socket, options); + + switch (_mode) { +#if BRPC_WITH_RDMA + case SOCKET_MODE_RDMA: + _high_speed_transport.reset(new RdmaTransport); + break; +#endif +#if BRPC_WITH_UBRING + case SOCKET_MODE_UBRING: + _high_speed_transport.reset(new UBShmTransport); + break; +#endif + default: + break; + } + if (_high_speed_transport) { + _high_speed_transport->Init(socket, options); + } +} + +void AdapterTransport::Release() { + if (_high_speed_transport) { + _high_speed_transport->Release(); + } + _tcp_transport->Release(); +} + +int AdapterTransport::Reset(int32_t expected_nref) { + if (_high_speed_transport) { + _high_speed_transport->Reset(expected_nref); + } + _tcp_transport->Reset(expected_nref); + _handshake.Reset(_socket); + _connection_completed.store(0, butil::memory_order_relaxed); + return 0; +} + +std::shared_ptr AdapterTransport::Connect() { + if (upgrade_capable(_mode)) { + return AdapterConnect::Wrap(_default_connect); + } + SocketUser* const client_messenger = + static_cast(get_client_side_messenger()); + if (client_messenger != NULL && + _socket->user() == client_messenger) { + FallbackToTcp(); + return AdapterConnect::Unwrap(_tcp_transport->Connect()); + } + return _tcp_transport->Connect(); +} + +Transport* AdapterTransport::ActiveTransport() const { + if (_high_speed_transport && + _handshake.phase(butil::memory_order_acquire) == + handshake::ESTABLISHED) { + return _high_speed_transport.get(); + } + return _tcp_transport.get(); +} + +int AdapterTransport::CutFromIOBuf(butil::IOBuf* buf) { + return ActiveTransport()->CutFromIOBuf(buf); +} + +ssize_t AdapterTransport::CutFromIOBufList( + butil::IOBuf** buf, size_t ndata) { + return ActiveTransport()->CutFromIOBufList(buf, ndata); +} + +int AdapterTransport::WaitEpollOut(butil::atomic* epollout_butex, + bool pollin, timespec duetime) { + return ActiveTransport()->WaitEpollOut( + epollout_butex, pollin, duetime); +} + +void AdapterTransport::ProcessEvent(bthread_attr_t attr) { + ActiveTransport()->ProcessEvent(attr); +} + +void AdapterTransport::QueueMessage(InputMessageClosure& input_msg, + int* num_bthread_created, + bool last_msg) { + ActiveTransport()->QueueMessage( + input_msg, num_bthread_created, last_msg); +} + +void AdapterTransport::Debug(std::ostream& os) { + if (_high_speed_transport) { + _high_speed_transport->Debug(os); + } + const char* state = "UNKNOWN"; + switch (_handshake.phase(butil::memory_order_acquire)) { + case handshake::UNINITIALIZED: state = "UNINITIALIZED"; break; + case handshake::PREPARING: state = "PREPARING"; break; + case handshake::HELLO_SEND: state = "HELLO_SEND"; break; + case handshake::HELLO_WAIT: state = "HELLO_WAIT"; break; + case handshake::NEGOTIATING: state = "NEGOTIATING"; break; + case handshake::ACK_SEND: state = "ACK_SEND"; break; + case handshake::ACK_WAIT: state = "ACK_WAIT"; break; + case handshake::EXTENSION_SEND: state = "EXTENSION_SEND"; break; + case handshake::EXTENSION_WAIT: state = "EXTENSION_WAIT"; break; + case handshake::ESTABLISHED: state = "ESTABLISHED"; break; + case handshake::FALLBACK_TCP: state = "FALLBACK_TCP"; break; + case handshake::FAILED: state = "FAILED"; break; + } + os << "\nhandshake_state=" << state + << "\nhandshake_version=" << _handshake.protocol_version(); +} + +void AdapterTransport::FallbackToTcp() { + _handshake.PublishFallback([this]() { + SetHighSpeedAvailable(false); + }); +} + +void AdapterTransport::SetHighSpeedAvailable(bool available) { + if (!_high_speed_transport) { + return; + } + switch (_mode) { +#if BRPC_WITH_RDMA + case SOCKET_MODE_RDMA: + static_cast(_high_speed_transport.get()) + ->SetHighSpeedAvailable(available); + break; +#endif +#if BRPC_WITH_UBRING + case SOCKET_MODE_UBRING: + static_cast(_high_speed_transport.get()) + ->SetHighSpeedAvailable(available); + break; +#endif + default: + break; + } +} + +void AdapterTransport::OnNewMessagesAfterUpgrade(Socket* socket) { + AdapterTransport* adapter = Get(socket); + if (adapter->_handshake.phase(butil::memory_order_acquire) == + handshake::ESTABLISHED) { + adapter->CheckUnexpectedTcpData(); + return; + } + + InputMessenger::OnNewMessagesUntil(socket, [](Socket* s) { + return Get(s)->handshake_phase() == handshake::ESTABLISHED; + }); + +#if BRPC_WITH_RDMA + if (adapter->_mode == SOCKET_MODE_RDMA && + adapter->_handshake.phase(butil::memory_order_acquire) == + handshake::ESTABLISHED) { + RdmaTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + if (transport->StartUpgradeEvents() < 0) { + const int saved_errno = errno != 0 ? errno : ERDMA; + transport->DeactivateUpgrade(); + adapter->_handshake.MarkFailed(); + adapter->CompleteConnection(handshake::FAILED); + socket->SetFailed( + saved_errno, + "Fail to start RDMA CQ events from %s: %s", + socket->description().c_str(), berror(saved_errno)); + } + } +#endif +#if BRPC_WITH_UBRING + if (adapter->_mode == SOCKET_MODE_UBRING && + adapter->handshake_phase() == handshake::ESTABLISHED) { + UBShmTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + transport->GetUBShmEp()->StartReceiveEvents(); + } +#endif + if (adapter->handshake_phase() == handshake::ESTABLISHED) { + adapter->CheckUnexpectedTcpData(); + } +} + +void AdapterTransport::OnNewDataFromTcp(Socket* socket) { + static_cast(socket->_transport.get())->ProcessTcpEvent(); +} + +void AdapterTransport::ProcessTcpEvent() { + int progress = Socket::PROGRESS_INIT; + while (true) { + const int phase = _handshake.phase(butil::memory_order_acquire); + if (phase != handshake::UNINITIALIZED && + phase < handshake::ESTABLISHED) { + _handshake.NotifyReadable(); + } else if (phase == handshake::FALLBACK_TCP) { + InputMessenger::OnNewMessages(_socket); + return; + } else if (phase == handshake::ESTABLISHED) { + CheckUnexpectedTcpData(); + return; + } + if (!_socket->MoreReadEvents(&progress)) { + break; + } + } +} + +void AdapterTransport::CheckUnexpectedTcpData() { + int progress = Socket::PROGRESS_INIT; + while (true) { + uint8_t byte; + const ssize_t nr = read(_socket->fd(), &byte, 1); + if (nr == 0) { + _socket->SetEOF(); + return; + } + if (nr > 0) { + _socket->SetFailed(EPROTO, "Read unexpected data from %s", + _socket->description().c_str()); + return; + } + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN) { + const int saved_errno = errno; + _socket->SetFailed(saved_errno, "Fail to read from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + return; + } + if (!_socket->MoreReadEvents(&progress)) { + return; + } + } +} + +void AdapterTransport::TryReadOnTcp() { + if (_socket->_nevent.fetch_add(1, butil::memory_order_acq_rel) != 0) { + return; + } + const int phase = _handshake.phase(butil::memory_order_acquire); + if (phase == handshake::FALLBACK_TCP) { + InputMessenger::OnNewMessages(_socket); + } else if (phase == handshake::ESTABLISHED) { + CheckUnexpectedTcpData(); + } +} + +} // namespace brpc diff --git a/src/brpc/adapter_transport.h b/src/brpc/adapter_transport.h new file mode 100644 index 0000000000..aed6ee5bb9 --- /dev/null +++ b/src/brpc/adapter_transport.h @@ -0,0 +1,103 @@ +// 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_ADAPTER_TRANSPORT_H +#define BRPC_ADAPTER_TRANSPORT_H + +#include + +#include "brpc/socket_mode.h" +#include "brpc/transport.h" +#include "brpc/transport_handshake.h" +#include "brpc/parse_result.h" + +namespace brpc { + +class TcpTransport; +class RdmaTransport; +class UBShmTransport; + +// The top-level Transport installed in Socket. It starts on TcpTransport and +// may switch to an independent RDMA/URMA/UBSHM Transport after a successful +// handshake. TCP remains usable before negotiation and after fallback. +class AdapterTransport : public Transport { +friend class TransportFactory; +friend class RdmaTransport; +friend class UBShmTransport; +public: + void Init(Socket* socket, const SocketOptions& options) override; + void Release() override; + int Reset(int32_t expected_nref) override; + std::shared_ptr Connect() override; + int CutFromIOBuf(butil::IOBuf* buf) override; + ssize_t CutFromIOBufList(butil::IOBuf** buf, size_t ndata) override; + 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 Debug(std::ostream& os) override; + + int handshake_phase() const { + return _handshake.phase(butil::memory_order_acquire); + } + int handshake_version() const { return _handshake.protocol_version(); } + handshake::HandshakeSession* handshake_session() { return &_handshake; } + Transport* high_speed_transport() const { + return _high_speed_transport.get(); + } + bool upgrade_capable(SocketMode mode) const; + + static AdapterTransport* Get(Socket* socket); + static const AdapterTransport* Get(const Socket* socket); + + // The only client-side upgrade entry point. Concrete transports provide + // resources; AdapterTransport owns the handshake orchestration. + static int StartClientUpgrade(const Socket* socket, + void (*done)(int, void*), void* data); + + ParseResult ProcessUpgradeReadable(butil::IOBuf* source); + void CompleteConnection(handshake::Phase terminal_phase); + bool connection_completed() const { + return _connection_completed.load(butil::memory_order_acquire) != 0; + } + + static void OnNewDataFromTcp(Socket* socket); + +private: + explicit AdapterTransport(SocketMode mode); + ~AdapterTransport() override; + + Transport* ActiveTransport() const; + void SetHighSpeedAvailable(bool available); + void FallbackToTcp(); + void TryReadOnTcp(); + void ProcessTcpEvent(); + void CheckUnexpectedTcpData(); + static void OnNewMessagesAfterUpgrade(Socket* socket); + static void* ProcessClientHandshake(void* arg); + + SocketMode _mode; + handshake::HandshakeSession _handshake; + std::unique_ptr _tcp_transport; + std::unique_ptr _high_speed_transport; + butil::atomic _connection_completed; +}; + +} // namespace brpc + +#endif // BRPC_ADAPTER_TRANSPORT_H diff --git a/src/brpc/global.cpp b/src/brpc/global.cpp index 0a0837e096..d64ae9f60a 100644 --- a/src/brpc/global.cpp +++ b/src/brpc/global.cpp @@ -69,7 +69,7 @@ // Protocols #include "brpc/protocol.h" -#include "brpc/policy/rdma_handshake_protocol.h" +#include "brpc/policy/transport_handshake_protocol.h" #include "brpc/policy/baidu_rpc_protocol.h" #include "brpc/policy/http_rpc_protocol.h" #include "brpc/policy/http2_rpc_protocol.h" @@ -438,12 +438,16 @@ static void GlobalInitializeOrDieImpl() { } // Protocols - Protocol rdma_handshake_protocol = { - ParseRdmaHandshake, nullptr, nullptr, - ProcessRdmaHandshake, nullptr, + Protocol transport_handshake_protocol = { + ParseTransportHandshake, nullptr, nullptr, + ProcessTransportHandshake, nullptr, nullptr, nullptr, nullptr, CONNECTION_TYPE_ALL, "rdma_handshake" }; - if (RegisterProtocol(PROTOCOL_RDMA_HANDSHAKE, rdma_handshake_protocol) != 0) { + // Retain the existing enum value and registered name to avoid changing + // public protocol identifiers while widening the implementation from RDMA + // to all transport upgrades. + if (RegisterProtocol(PROTOCOL_RDMA_HANDSHAKE, + transport_handshake_protocol) != 0) { exit(1); } diff --git a/src/brpc/handshake/handshake_adapter.cpp b/src/brpc/handshake/handshake_adapter.cpp new file mode 100644 index 0000000000..8678ca0d4a --- /dev/null +++ b/src/brpc/handshake/handshake_adapter.cpp @@ -0,0 +1,61 @@ +// 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/handshake/handshake_adapter.h" + +#include "brpc/socket.h" + +namespace brpc { +namespace handshake { + +// InputMessenger may call this entry repeatedly while bytes arrive. Keep all +// connection state in HandshakeSession and Socket rather than in the adapter, +// so a single stateless adapter can serve every connection. The parsing +// context retains that selected adapter while waiting for the peer extension +// or ACK, whose frames have no protocol magic of their own. +ParseResult StandardHandshakeAdapter::ExecuteServerHandshake( + butil::IOBuf* source, Socket* socket) { + const StepResult result = RunServerStep(source, socket); + if (result == STEP_NEED_MORE) { + const int phase = GetSession(socket)->phase(butil::memory_order_acquire); + if ((phase == ACK_WAIT || phase == EXTENSION_WAIT) && + socket->parsing_context() == nullptr) { + ServerHandshakeContext* context = + ServerHandshakeContext::Create(this); + if (context == nullptr) { + GetSession(socket)->MarkFailed(); + return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); + } + socket->reset_parsing_context(context); + } + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); + } + + socket->reset_parsing_context(nullptr); + if (result == STEP_ERROR) { + return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); + } + if (GetSession(socket)->phase(butil::memory_order_acquire) == ESTABLISHED) { + // Do not hand this TCP control stream to ordinary RPC parsers. The + // outer transport event handler stops reading once upgrade completes. + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); + } + return MakeParseError(PARSE_ERROR_TRY_OTHERS); +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/handshake/handshake_adapter.h b/src/brpc/handshake/handshake_adapter.h new file mode 100644 index 0000000000..2a344c9ed3 --- /dev/null +++ b/src/brpc/handshake/handshake_adapter.h @@ -0,0 +1,73 @@ +// 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_HANDSHAKE_HANDSHAKE_ADAPTER_H +#define BRPC_HANDSHAKE_HANDSHAKE_ADAPTER_H + +#include "butil/macros.h" +#include "brpc/parse_result.h" +#include "brpc/transport_handshake.h" + +namespace butil { +class IOBuf; +} + +namespace brpc { + +class Socket; + +namespace handshake { + +// The minimal seam between an upgrade protocol and InputMessenger. Protocols +// that cannot use the standard parser lifecycle implement this interface +// directly. +class HandshakeAdapter { +public: + virtual ~HandshakeAdapter() = default; + + virtual ParseResult ExecuteServerHandshake( + butil::IOBuf* source, Socket* socket) = 0; + +protected: + HandshakeAdapter() = default; + +private: + DISALLOW_COPY_AND_ASSIGN(HandshakeAdapter); +}; + +// Reusable InputMessenger implementation. Protocol adapters provide only the +// protocol-specific server step; the common session owns all phases. +class StandardHandshakeAdapter : public HandshakeAdapter { +public: + ParseResult ExecuteServerHandshake( + butil::IOBuf* source, Socket* socket) override; + +protected: + StandardHandshakeAdapter() = default; + + virtual StepResult RunServerStep( + butil::IOBuf* source, Socket* socket) = 0; + virtual HandshakeSession* GetSession(Socket* socket) const = 0; + +private: + DISALLOW_COPY_AND_ASSIGN(StandardHandshakeAdapter); +}; + +} // namespace handshake +} // namespace brpc + +#endif // BRPC_HANDSHAKE_HANDSHAKE_ADAPTER_H diff --git a/src/brpc/handshake/handshake_frame.cpp b/src/brpc/handshake/handshake_frame.cpp new file mode 100644 index 0000000000..f25e50c15c --- /dev/null +++ b/src/brpc/handshake/handshake_frame.cpp @@ -0,0 +1,235 @@ +// 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/handshake/handshake_frame.h" + +#include +#include +#include + +#include "butil/sys_byteorder.h" + +namespace brpc { +namespace handshake { + +size_t FrameCodec::LengthFieldSize(const FrameSpec& spec) { + switch (spec.length_encoding) { + case FrameSpec::FIXED: return 0; + case FrameSpec::U16_TOTAL_LENGTH: return sizeof(uint16_t); + case FrameSpec::U32_BODY_LENGTH: return sizeof(uint32_t); + } + return 0; +} + +FrameResult FrameCodec::DecodeLength(const FrameSpec& spec, + const void* header, + size_t* frame_len) { + const size_t length_size = LengthFieldSize(spec); + const size_t header_len = spec.magic_len + length_size; + if (spec.min_frame_len < header_len || + spec.max_frame_len < spec.min_frame_len) { + return FRAME_PROTOCOL_ERROR; + } + + if (spec.length_encoding == FrameSpec::FIXED) { + if (spec.min_frame_len != spec.max_frame_len) { + return FRAME_PROTOCOL_ERROR; + } + *frame_len = spec.min_frame_len; + } else if (spec.length_encoding == FrameSpec::U16_TOTAL_LENGTH) { + uint16_t total_be = 0; + memcpy(&total_be, static_cast(header) + spec.magic_len, + sizeof(total_be)); + *frame_len = butil::NetToHost16(total_be); + } else { + uint32_t body_be = 0; + memcpy(&body_be, static_cast(header) + spec.magic_len, + sizeof(body_be)); + const size_t body_len = butil::NetToHost32(body_be); + if (body_len > std::numeric_limits::max() - header_len) { + return FRAME_PROTOCOL_ERROR; + } + *frame_len = header_len + body_len; + } + + if (*frame_len < spec.min_frame_len || + *frame_len > spec.max_frame_len) { + return FRAME_PROTOCOL_ERROR; + } + return FRAME_OK; +} + +FrameResult FrameCodec::Encode(const FrameSpec& spec, + const std::string& payload, + std::string* frame) { + if (frame == NULL || (spec.magic_len != 0 && spec.magic == NULL)) { + return FRAME_PROTOCOL_ERROR; + } + const size_t length_size = LengthFieldSize(spec); + const size_t header_len = spec.magic_len + length_size; + if (payload.size() > std::numeric_limits::max() - header_len) { + return FRAME_PROTOCOL_ERROR; + } + const size_t total_len = header_len + payload.size(); + if (total_len < spec.min_frame_len || total_len > spec.max_frame_len) { + return FRAME_PROTOCOL_ERROR; + } + if (spec.length_encoding == FrameSpec::FIXED && + spec.min_frame_len != spec.max_frame_len) { + return FRAME_PROTOCOL_ERROR; + } + if (spec.length_encoding == FrameSpec::U16_TOTAL_LENGTH && + total_len > std::numeric_limits::max()) { + return FRAME_PROTOCOL_ERROR; + } + if (spec.length_encoding == FrameSpec::U32_BODY_LENGTH && + payload.size() > std::numeric_limits::max()) { + return FRAME_PROTOCOL_ERROR; + } + + frame->clear(); + frame->reserve(total_len); + if (spec.magic_len != 0) { + frame->append(spec.magic, spec.magic_len); + } + if (spec.length_encoding == FrameSpec::U16_TOTAL_LENGTH) { + const uint16_t total_be = + butil::HostToNet16(static_cast(total_len)); + frame->append(reinterpret_cast(&total_be), + sizeof(total_be)); + } else if (spec.length_encoding == FrameSpec::U32_BODY_LENGTH) { + const uint32_t body_be = + butil::HostToNet32(static_cast(payload.size())); + frame->append(reinterpret_cast(&body_be), + sizeof(body_be)); + } + frame->append(payload); + return FRAME_OK; +} + +FrameResult FrameCodec::ReadFrame(HandshakeIO* io, const FrameSpec& spec, + bool push_back_on_not_mine, + std::string* payload) { + if (io == NULL || payload == NULL || + (spec.magic_len != 0 && spec.magic == NULL)) { + return FRAME_PROTOCOL_ERROR; + } + + std::string header(spec.magic_len + LengthFieldSize(spec), '\0'); + if (spec.magic_len != 0 && + io->ReadExact(&header[0], spec.magic_len) < 0) { + return FRAME_IO_ERROR; + } + if (spec.magic_len != 0 && + memcmp(header.data(), spec.magic, spec.magic_len) != 0) { + if (push_back_on_not_mine && + io->PushBack(header.data(), spec.magic_len) < 0) { + return FRAME_IO_ERROR; + } + return FRAME_NOT_MINE; + } + + const size_t length_size = LengthFieldSize(spec); + if (length_size != 0 && + io->ReadExact(&header[spec.magic_len], length_size) < 0) { + return FRAME_IO_ERROR; + } + size_t frame_len = 0; + FrameResult result = DecodeLength(spec, header.data(), &frame_len); + if (result != FRAME_OK) { + return result; + } + const size_t body_len = frame_len - header.size(); + payload->assign(body_len, '\0'); + if (body_len != 0 && io->ReadExact(&(*payload)[0], body_len) < 0) { + return FRAME_IO_ERROR; + } + return FRAME_OK; +} + +FrameResult FrameCodec::ParseBufferedFrame(HandshakeInput* input, + const FrameSpec& spec, + std::string* payload, + bool* magic_matched) { + if (magic_matched != NULL) { + *magic_matched = false; + } + if (input == NULL || payload == NULL || + (spec.magic_len != 0 && spec.magic == NULL)) { + return FRAME_PROTOCOL_ERROR; + } + const size_t header_len = spec.magic_len + LengthFieldSize(spec); + if (input->Size() < spec.magic_len) { + return FRAME_NEED_MORE; + } + std::string header(header_len, '\0'); + if (spec.magic_len != 0 && + !input->CopyTo(&header[0], spec.magic_len)) { + return FRAME_NEED_MORE; + } + if (spec.magic_len != 0 && + memcmp(header.data(), spec.magic, spec.magic_len) != 0) { + return FRAME_NOT_MINE; + } + if (magic_matched != NULL) { + *magic_matched = true; + } + if (input->Size() < header_len || + (header_len != 0 && !input->CopyTo(&header[0], header_len))) { + return FRAME_NEED_MORE; + } + size_t frame_len = 0; + FrameResult result = DecodeLength(spec, header.data(), &frame_len); + if (result != FRAME_OK) { + return result; + } + if (input->Size() < frame_len) { + return FRAME_NEED_MORE; + } + + std::string frame(frame_len, '\0'); + if (!input->CopyTo(&frame[0], frame_len)) { + return FRAME_NEED_MORE; + } + if (!input->Consume(frame_len)) { + return FRAME_PROTOCOL_ERROR; + } + payload->assign(frame.data() + header_len, frame_len - header_len); + return FRAME_OK; +} + +FrameResult FrameCodec::WriteFrame(HandshakeIO* io, const FrameSpec& spec, + const std::string& payload) { + if (io == NULL) { + return FRAME_PROTOCOL_ERROR; + } + std::string frame; + const FrameResult result = Encode(spec, payload, &frame); + if (result != FRAME_OK) { + return result; + } + return io->WriteAll(frame.data(), frame.size()) == 0 + ? FRAME_OK : FRAME_IO_ERROR; +} + +FrameResult FrameCodec::DrainFrame(HandshakeIO* io, const FrameSpec& spec) { + std::string ignored; + return ReadFrame(io, spec, false, &ignored); +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/handshake/handshake_frame.h b/src/brpc/handshake/handshake_frame.h new file mode 100644 index 0000000000..dbe399f738 --- /dev/null +++ b/src/brpc/handshake/handshake_frame.h @@ -0,0 +1,92 @@ +// 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_HANDSHAKE_HANDSHAKE_FRAME_H +#define BRPC_HANDSHAKE_HANDSHAKE_FRAME_H + +#include +#include + +#include "brpc/handshake/handshake_io.h" + +namespace brpc { +namespace handshake { + +enum FrameResult { + FRAME_OK = 0, + FRAME_NOT_MINE, + FRAME_NEED_MORE, + FRAME_IO_ERROR, + FRAME_PROTOCOL_ERROR, +}; + +struct FrameSpec { + enum LengthEncoding { + FIXED, + U16_TOTAL_LENGTH, + U32_BODY_LENGTH, + }; + + FrameSpec() + : magic(NULL), magic_len(0), min_frame_len(0), max_frame_len(0), + length_encoding(FIXED) {} + + FrameSpec(const char* magic_in, size_t magic_len_in, + size_t min_frame_len_in, size_t max_frame_len_in, + LengthEncoding length_encoding_in) + : magic(magic_in), magic_len(magic_len_in), + min_frame_len(min_frame_len_in), + max_frame_len(max_frame_len_in), + length_encoding(length_encoding_in) {} + + const char* magic; + size_t magic_len; + size_t min_frame_len; + size_t max_frame_len; + LengthEncoding length_encoding; +}; + +// Handles only framing. Protocol implementations receive and produce payloads +// after magic/length fields and remain responsible for their own business +// fields and version semantics. +class FrameCodec { +public: + static FrameResult Encode(const FrameSpec& spec, + const std::string& payload, + std::string* frame); + static FrameResult ReadFrame(HandshakeIO* io, const FrameSpec& spec, + bool push_back_on_not_mine, + std::string* payload); + static FrameResult ParseBufferedFrame(HandshakeInput* input, + const FrameSpec& spec, + std::string* payload, + bool* magic_matched = NULL); + static FrameResult WriteFrame(HandshakeIO* io, const FrameSpec& spec, + const std::string& payload); + static FrameResult DrainFrame(HandshakeIO* io, const FrameSpec& spec); + +private: + static size_t LengthFieldSize(const FrameSpec& spec); + static FrameResult DecodeLength(const FrameSpec& spec, + const void* header, + size_t* frame_len); +}; + +} // namespace handshake +} // namespace brpc + +#endif // BRPC_HANDSHAKE_HANDSHAKE_FRAME_H diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp new file mode 100644 index 0000000000..67cc95e37c --- /dev/null +++ b/src/brpc/handshake/handshake_io.cpp @@ -0,0 +1,162 @@ +// 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/handshake/handshake_io.h" + +#include +#include +#include + +#include "bthread/butex.h" +#include "butil/time.h" +#include "brpc/errno.pb.h" +#include "brpc/socket.h" + +namespace brpc { +namespace handshake { + +size_t IOBufHandshakeInput::Size() const { + return _source != NULL ? _source->size() : 0; +} + +bool IOBufHandshakeInput::CopyTo(void* data, size_t len) const { + return _source != NULL && _source->copy_to(data, len) == len; +} + +bool IOBufHandshakeInput::Consume(size_t len) { + return _source != NULL && _source->pop_front(len) == len; +} + +static const int WAIT_TIMEOUT_MS = 50; + +SocketHandshakeIO::SocketHandshakeIO(Socket* socket) + : _socket(socket) + , _read_butex(bthread::butex_create_checked >()) { +} + +SocketHandshakeIO::~SocketHandshakeIO() { + bthread::butex_destroy(_read_butex); +} + +void SocketHandshakeIO::Reset(Socket* socket) { + _socket = socket; +} + +void SocketHandshakeIO::NotifyReadable() { + _read_butex->fetch_add(1, butil::memory_order_release); + bthread::butex_wake(_read_butex); +} + +template +static int ReadExactLoop(butil::atomic* read_butex, + SocketId socket_id, + size_t len, ReadOnce read_once) { + size_t received = 0; + while (received < len) { + const int expected = read_butex->load(butil::memory_order_acquire); + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const ssize_t nr = read_once(received, len - received); + if (nr < 0) { + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN) { + return -1; + } + SocketUniquePtr alive; + if (Socket::Address(socket_id, &alive) != 0) { + errno = EFAILEDSOCKET; + return -1; + } + if (bthread::butex_wait(read_butex, expected, &duetime) < 0 && + errno != EWOULDBLOCK && errno != ETIMEDOUT) { + return -1; + } + } else if (nr == 0) { + errno = EEOF; + return -1; + } else { + received += nr; + } + } + return 0; +} + +int SocketHandshakeIO::ReadExact(void* data, size_t len) { + CHECK(data != NULL); + CHECK(_socket != NULL); + const int fd = _socket->fd(); + return ReadExactLoop(_read_butex, _socket->id(), len, + [data, fd](size_t offset, size_t remaining) { + return read(fd, static_cast(data) + offset, remaining); + }); +} + +template +static int WriteAllLoop(size_t len, WriteOnce write_once, + WaitWritable wait_writable) { + size_t written = 0; + while (written < len) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const ssize_t nw = write_once(written, len - written); + if (nw > 0) { + written += nw; + continue; + } + if (nw == 0) { + errno = EPIPE; + return -1; + } + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN) { + return -1; + } + if (wait_writable(&duetime) < 0 && + errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +int SocketHandshakeIO::WriteAll(const void* data, size_t len) { + CHECK(data != NULL); + CHECK(_socket != NULL); + const int fd = _socket->fd(); + return WriteAllLoop(len, + [data, fd](size_t offset, size_t remaining) { + return write(fd, static_cast(data) + offset, + remaining); + }, + [this](const timespec* duetime) { + return _socket->WaitEpollOut( + _socket->fd(), true, duetime); + }); +} + +int SocketHandshakeIO::PushBack(const void* data, size_t len) { + CHECK(_socket != NULL); + if (len != 0) { + return _socket->fd_input_processor().read_buf().append(data, len) == 0 ? 0 : -1; + } + return 0; +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/handshake/handshake_io.h b/src/brpc/handshake/handshake_io.h new file mode 100644 index 0000000000..82f83ab5cd --- /dev/null +++ b/src/brpc/handshake/handshake_io.h @@ -0,0 +1,89 @@ +// 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_HANDSHAKE_HANDSHAKE_IO_H +#define BRPC_HANDSHAKE_HANDSHAKE_IO_H + +#include + +#include "butil/atomicops.h" +#include "butil/iobuf.h" +#include "butil/macros.h" + +namespace brpc { + +class Socket; + +namespace handshake { + +// Blocking byte-stream interface used by client handshakes and by protocols +// whose server handshake still runs in a dedicated bthread. +class HandshakeIO { +public: + virtual ~HandshakeIO() = default; + + virtual int ReadExact(void* data, size_t len) = 0; + virtual int WriteAll(const void* data, size_t len) = 0; + virtual int PushBack(const void* data, size_t len) = 0; +}; + +// Non-blocking input used by the standard InputMessenger parser path. +// Consume is called only after a complete frame has been validated. +class HandshakeInput { +public: + virtual ~HandshakeInput() = default; + + virtual size_t Size() const = 0; + virtual bool CopyTo(void* data, size_t len) const = 0; + virtual bool Consume(size_t len) = 0; +}; + +class IOBufHandshakeInput : public HandshakeInput { +public: + explicit IOBufHandshakeInput(butil::IOBuf* source) : _source(source) {} + + size_t Size() const override; + bool CopyTo(void* data, size_t len) const override; + bool Consume(size_t len) override; + +private: + butil::IOBuf* _source; +}; + +class SocketHandshakeIO : public HandshakeIO { +public: + explicit SocketHandshakeIO(Socket* socket = NULL); + ~SocketHandshakeIO() override; + + void Reset(Socket* socket); + void NotifyReadable(); + + int ReadExact(void* data, size_t len) override; + int WriteAll(const void* data, size_t len) override; + int PushBack(const void* data, size_t len) override; + +private: + Socket* _socket; + butil::atomic* _read_butex; + + DISALLOW_COPY_AND_ASSIGN(SocketHandshakeIO); +}; + +} // namespace handshake +} // namespace brpc + +#endif // BRPC_HANDSHAKE_HANDSHAKE_IO_H diff --git a/src/brpc/handshake/rdma_handshake.cpp b/src/brpc/handshake/rdma_handshake.cpp new file mode 100644 index 0000000000..26c622a1ac --- /dev/null +++ b/src/brpc/handshake/rdma_handshake.cpp @@ -0,0 +1,550 @@ +// 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/handshake/rdma_handshake.h" + +#include +#include +#include + +#include "butil/logging.h" +#include "butil/raw_pack.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/handshake/rdma_handshake_constants.h" +#include "brpc/rdma_handshake.pb.h" +#include "brpc/socket.h" + +#if BRPC_WITH_RDMA + +#include + +#include + +#include "brpc/rdma_transport.h" + +namespace brpc { +namespace rdma { + +DEFINE_int32(rdma_client_handshake_version, 2, + "RDMA handshake protocol version used by client. " + "2 = legacy 'RDMA' magic (default, compatible with all servers); " + "3 = new 'RDM3' protobuf-based handshake " + "(MUST only be enabled after target servers support v3)."); +DECLARE_bool(rdma_trace_verbose); + +extern const uint16_t MIN_QP_SIZE; +extern const uint16_t MIN_BLOCK_SIZE; +extern bool g_skip_rdma_init; + +DEFINE_bool(rdma_ece, false, + "Enable end-to-end ECE negotiation in the RDMA v3 handshake"); + +void RdmaHandshakeAdapter::FillLocalHello(ParsedHello* local) const { + _ep->GetLocalConnectionInfo(local); +} + +void RdmaHandshakeAdapter::PrepareClientEce() { + if (!FLAGS_rdma_ece) { + return; + } + ibv_ece ece; + const int rc = _ep->QueryLocalEce(&ece); + if (rc == 0) { + _ep->SetOutgoingEce(ece); + } else if (rc < 0) { + LOG_IF(WARNING, FLAGS_rdma_trace_verbose) + << "Fail to IbvQueryEce on client, ECE not advertised"; + } +} + +const handshake::FrameSpec& RdmaHandshakeAdapter::AckFrameSpec() const { + return RdmaAckFrameSpec(); +} + +handshake::StepResult RdmaHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + return BuildLocalHello(enabled, payload); +} + +handshake::StepResult RdmaHandshakeAdapter::ParseHello( + const std::string& payload) { + return ParseRemoteHello(payload, &_remote); +} + +namespace v2_wire { + +void HelloMessage::Serialize(void* data) const { + butil::RawPacker(data) + .pack16(msg_len) + .pack16(hello_ver) + .pack16(impl_ver) + .pack32(block_size) + .pack16(sq_size) + .pack16(rq_size) + .pack16(lid) + .pack_bytes(gid.raw, sizeof(gid.raw)) + .pack32(qp_num); +} + +void HelloMessage::Deserialize(const void* data) { + butil::RawUnpacker(data) + .unpack16(msg_len) + .unpack16(hello_ver) + .unpack16(impl_ver) + .unpack32(block_size) + .unpack16(sq_size) + .unpack16(rq_size) + .unpack16(lid) + .unpack_bytes(gid.raw, sizeof(gid.raw)) + .unpack32(qp_num); +} + +static bool ValidHelloMessage(const HelloMessage& msg) { + return msg.hello_ver == HELLO_V2_VERSION && + msg.impl_ver == IMPL_V2_VERSION && + msg.block_size >= MIN_BLOCK_SIZE && + msg.sq_size >= MIN_QP_SIZE && + msg.rq_size >= MIN_QP_SIZE; +} + +static void TranslateHello(const HelloMessage& msg, ParsedHello* out) { + out->block_size = msg.block_size; + out->sq_size = msg.sq_size; + out->rq_size = msg.rq_size; + out->lid = msg.lid; + out->gid = msg.gid; + out->qp_num = msg.qp_num; +} + +static void FillMessage(const ParsedHello& local, HelloMessage* msg) { + msg->msg_len = HELLO_V2_MSG_LEN_MIN; + msg->hello_ver = HELLO_V2_VERSION; + msg->impl_ver = IMPL_V2_VERSION; + msg->block_size = local.block_size; + msg->sq_size = local.sq_size; + msg->rq_size = local.rq_size; + msg->lid = local.lid; + msg->gid = local.gid; + msg->qp_num = local.qp_num; +} + +static handshake::StepResult SerializePayload( + const HelloMessage& msg, std::string* payload) { + uint8_t body[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]; + msg.Serialize(body); + // FrameCodec owns msg_len, so the protocol payload starts after it. + payload->assign(reinterpret_cast(body + sizeof(uint16_t)), + sizeof(body) - sizeof(uint16_t)); + return handshake::STEP_OK; +} + +static handshake::StepResult ParsePayload( + const std::string& payload, ParsedHello* remote) { + const size_t base_payload_len = + HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN - sizeof(uint16_t); + if (payload.size() < base_payload_len) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint8_t body[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]; + const uint16_t total_be = butil::HostToNet16( + static_cast(HELLO_MAGIC_LEN + sizeof(uint16_t) + + payload.size())); + memcpy(body, &total_be, sizeof(total_be)); + memcpy(body + sizeof(total_be), payload.data(), base_payload_len); + + HelloMessage msg{}; + msg.Deserialize(body); + if (!ValidHelloMessage(msg)) { + return handshake::STEP_FALLBACK; + } + TranslateHello(msg, remote); + return handshake::STEP_OK; +} + +} // namespace v2_wire + +const handshake::FrameSpec& +RdmaClientHandshakeAdapterV2::HelloFrameSpec() const { + return RdmaHelloFrameSpec(2); +} + +handshake::StepResult RdmaClientHandshakeAdapterV2::BuildLocalHello( + bool enabled, std::string* payload) { + CHECK(enabled); + ParsedHello local{}; + FillLocalHello(&local); + v2_wire::HelloMessage msg{}; + v2_wire::FillMessage(local, &msg); + return v2_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaClientHandshakeAdapterV2::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v2_wire::ParsePayload(payload, remote); +} + +const handshake::FrameSpec& +RdmaServerHandshakeAdapterV2::HelloFrameSpec() const { + return RdmaHelloFrameSpec(2); +} + +handshake::StepResult RdmaServerHandshakeAdapterV2::BuildLocalHello( + bool enabled, std::string* payload) { + v2_wire::HelloMessage msg{}; + msg.msg_len = HELLO_V2_MSG_LEN_MIN; + if (enabled) { + ParsedHello local{}; + FillLocalHello(&local); + v2_wire::FillMessage(local, &msg); + } + return v2_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaServerHandshakeAdapterV2::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v2_wire::ParsePayload(payload, remote); +} + +namespace v3_wire { + +static bool ValidRdmaHello(const RdmaHello& msg) { + if (msg.gid().size() != sizeof(ibv_gid)) { + return false; + } + const uint16_t max_uint16 = std::numeric_limits::max(); + if (msg.sq_size() > max_uint16 || msg.rq_size() > max_uint16 || + msg.lid() > max_uint16) { + return false; + } + if (msg.block_size() < MIN_BLOCK_SIZE || msg.sq_size() < MIN_QP_SIZE || + msg.rq_size() < MIN_QP_SIZE) { + return false; + } + return msg.qp_num() != 0 || g_skip_rdma_init; +} + +static void FillLocalRdmaHello(const ParsedHello& local, RdmaHello* msg) { + msg->set_block_size(local.block_size); + msg->set_sq_size(local.sq_size); + msg->set_rq_size(local.rq_size); + msg->set_lid(local.lid); + msg->set_gid(reinterpret_cast(local.gid.raw), + sizeof(local.gid.raw)); + msg->set_qp_num(local.qp_num); + if (FLAGS_rdma_ece && local.ece.has_value()) { + RdmaEce* ece = msg->mutable_ece(); + ece->set_vendor_id(local.ece->vendor_id); + ece->set_options(local.ece->options); + ece->set_comp_mask(local.ece->comp_mask); + } +} + +static void TranslateHello(const RdmaHello& msg, ParsedHello* out) { + out->block_size = msg.block_size(); + out->sq_size = static_cast(msg.sq_size()); + out->rq_size = static_cast(msg.rq_size()); + out->lid = static_cast(msg.lid()); + fast_memcpy(out->gid.raw, msg.gid().data(), sizeof(out->gid.raw)); + out->qp_num = msg.qp_num(); + if (FLAGS_rdma_ece && msg.has_ece()) { + ibv_ece ece; + ece.vendor_id = msg.ece().vendor_id(); + ece.options = msg.ece().options(); + ece.comp_mask = msg.ece().comp_mask(); + out->ece = ece; + } +} + +static handshake::StepResult SerializePayload( + const RdmaHello& msg, std::string* payload) { + if (!msg.SerializeToString(payload) || + payload->size() > HELLO_V3_MAX_PB_SIZE) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + return handshake::STEP_OK; +} + +static handshake::StepResult ParsePayload( + const std::string& payload, ParsedHello* remote) { + RdmaHello msg; + if (!msg.ParseFromArray(payload.data(), static_cast(payload.size()))) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + if (!ValidRdmaHello(msg)) { + return handshake::STEP_FALLBACK; + } + TranslateHello(msg, remote); + return handshake::STEP_OK; +} + +static void FillDisabledHello(RdmaHello* msg) { + msg->set_block_size(0); + msg->set_sq_size(0); + msg->set_rq_size(0); + msg->set_lid(0); + msg->set_gid(std::string(sizeof(ibv_gid), '\0')); + msg->set_qp_num(0); +} + +} // namespace v3_wire + +const handshake::FrameSpec& +RdmaClientHandshakeAdapterV3::HelloFrameSpec() const { + return RdmaHelloFrameSpec(3); +} + +handshake::StepResult RdmaClientHandshakeAdapterV3::BuildLocalHello( + bool enabled, std::string* payload) { + CHECK(enabled); + PrepareClientEce(); + ParsedHello local{}; + FillLocalHello(&local); + RdmaHello msg; + v3_wire::FillLocalRdmaHello(local, &msg); + return v3_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaClientHandshakeAdapterV3::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v3_wire::ParsePayload(payload, remote); +} + +const handshake::FrameSpec& +RdmaServerHandshakeAdapterV3::HelloFrameSpec() const { + return RdmaHelloFrameSpec(3); +} + +handshake::StepResult RdmaServerHandshakeAdapterV3::BuildLocalHello( + bool enabled, std::string* payload) { + RdmaHello msg; + if (enabled) { + ParsedHello local{}; + FillLocalHello(&local); + v3_wire::FillLocalRdmaHello(local, &msg); + } else { + v3_wire::FillDisabledHello(&msg); + } + return v3_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaServerHandshakeAdapterV3::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v3_wire::ParsePayload(payload, remote); +} + +std::unique_ptr CreateClientHandshakeAdapter( + RdmaEndpoint* ep) { + if (FLAGS_rdma_client_handshake_version == 3) { + return std::unique_ptr( + new RdmaClientHandshakeAdapterV3(ep)); + } + return std::unique_ptr( + new RdmaClientHandshakeAdapterV2(ep)); +} + +std::vector > +CreateServerHandshakeAdapters(RdmaEndpoint* ep) { + std::vector > adapters; + adapters.emplace_back(new RdmaServerHandshakeAdapterV2(ep)); + adapters.emplace_back(new RdmaServerHandshakeAdapterV3(ep)); + return adapters; +} + +} // namespace rdma +} // namespace brpc + +#endif // BRPC_WITH_RDMA + +namespace brpc { +namespace handshake { + +class RdmaServerHandshakeAdapter : public StandardHandshakeAdapter { +public: + RdmaServerHandshakeAdapter() = default; + +protected: + StepResult RunServerStep( + butil::IOBuf* source, Socket* socket) override; + HandshakeSession* GetSession(Socket* socket) const override; + +private: + StepResult RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket); +#if BRPC_WITH_RDMA + StepResult RunRdmaServerHandshake( + butil::IOBuf* source, Socket* socket); +#endif + + DISALLOW_COPY_AND_ASSIGN(RdmaServerHandshakeAdapter); +}; + +static constexpr uint16_t V2_HELLO_VERSION_INVALID = + std::numeric_limits::max(); +static constexpr size_t V3_GID_LEN = 16; + +class RdmaFallbackProtocol : public FallbackHandshakeProtocol { +public: + explicit RdmaFallbackProtocol(int version) : _version(version) {} + + int ProtocolVersion() const override { return _version; } + const FrameSpec& HelloFrameSpec() const override { + return rdma::RdmaHelloFrameSpec(_version); + } + const FrameSpec& AckFrameSpec() const override { + return rdma::RdmaAckFrameSpec(); + } + StepResult BuildHello(bool enabled, std::string* payload) override { + if (enabled) { + errno = EPROTO; + return STEP_ERROR; + } + if (_version == 2) { + payload->assign( + rdma::HELLO_V2_MSG_LEN_MIN - rdma::HELLO_MAGIC_LEN - + sizeof(uint16_t), + '\0'); + butil::RawPacker(&(*payload)[0]) + .pack16(V2_HELLO_VERSION_INVALID); + return STEP_OK; + } + + rdma::RdmaHello reply; + reply.set_block_size(0); + reply.set_sq_size(0); + reply.set_rq_size(0); + reply.set_lid(0); + reply.set_gid(std::string(V3_GID_LEN, '\0')); + reply.set_qp_num(0); + if (!reply.SerializeToString(payload)) { + errno = EPROTO; + return STEP_ERROR; + } + return STEP_OK; + } + +private: + int _version; +}; + +HandshakeAdapter* GetRdmaServerHandshakeAdapter() { + static RdmaServerHandshakeAdapter adapter; + return &adapter; +} + +HandshakeSession* RdmaServerHandshakeAdapter::GetSession( + Socket* socket) const { + return AdapterTransport::Get(socket)->handshake_session(); +} + +StepResult RdmaServerHandshakeAdapter::RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket) { + IOBufHandshakeInput input(source); + RdmaFallbackProtocol v2(2); + RdmaFallbackProtocol v3(3); + std::vector protocols; + protocols.push_back(&v2); + protocols.push_back(&v3); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, &input, &transport, false); +} + +#if BRPC_WITH_RDMA +class RdmaServerHandshakeTransport : public HandshakeTransport { +public: + RdmaServerHandshakeTransport( + RdmaTransport* transport, Socket* socket, butil::IOBuf* source) + : _transport(transport), _socket(socket), _source(source), + _protocol(NULL) {} + + void OnProtocolSelected(HandshakeProtocol* protocol) override { + _protocol = static_cast(protocol); + } + + StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources() == 0) { + return STEP_OK; + } + PLOG(WARNING) << "Fail to allocate rdma resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + + StepResult NegotiateResources() override { + CHECK(_protocol != NULL); + if (_transport->NegotiateUpgradeResources( + _protocol->remote(), true) == 0) { + return STEP_OK; + } + PLOG(WARNING) << "Fail to negotiate rdma resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + void OnFailed() override { _transport->DeactivateUpgrade(); } + + StepResult ValidateEstablished() override { + return _source->empty() ? STEP_OK : STEP_ERROR; + } + +private: + RdmaTransport* _transport; + Socket* _socket; + butil::IOBuf* _source; + rdma::RdmaHandshakeAdapter* _protocol; +}; + +StepResult RdmaServerHandshakeAdapter::RunRdmaServerHandshake( + butil::IOBuf* source, Socket* socket) { + RdmaTransport* transport = RdmaTransport::Get(socket); + CHECK(transport->GetRdmaEp() != NULL); + + std::vector > protocols = + transport->CreateServerHandshakeAdapters(); + std::vector protocol_ptrs; + protocol_ptrs.reserve(protocols.size()); + for (size_t i = 0; i < protocols.size(); ++i) { + protocol_ptrs.push_back(protocols[i].get()); + } + IOBufHandshakeInput input(source); + RdmaServerHandshakeTransport participant(transport, socket, source); + return GetSession(socket)->RunServer( + protocol_ptrs, &input, &participant, false); +} +#endif + +StepResult RdmaServerHandshakeAdapter::RunServerStep( + butil::IOBuf* source, Socket* socket) { +#if BRPC_WITH_RDMA + if (AdapterTransport::Get(socket)->upgrade_capable( + SOCKET_MODE_RDMA)) { + return RunRdmaServerHandshake(source, socket); + } +#endif + return RunFallbackServerHandshake(source, socket); +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/handshake/rdma_handshake.h b/src/brpc/handshake/rdma_handshake.h new file mode 100644 index 0000000000..80057e687e --- /dev/null +++ b/src/brpc/handshake/rdma_handshake.h @@ -0,0 +1,160 @@ +// 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_HANDSHAKE_RDMA_HANDSHAKE_H +#define BRPC_HANDSHAKE_RDMA_HANDSHAKE_H + +#include "brpc/handshake/handshake_adapter.h" + +namespace brpc { +namespace handshake { + +// Returns the RDMA adapter used by policy::ParseTransportHandshake. The +// concrete type is private to the implementation; callers only learn the +// common HandshakeAdapter interface. +HandshakeAdapter* GetRdmaServerHandshakeAdapter(); + +} // namespace handshake +} // namespace brpc + +#if BRPC_WITH_RDMA + +#include +#include +#include + +#include + +#include "butil/containers/optional.h" +#include "butil/macros.h" +#include "brpc/rdma/rdma_endpoint.h" +#include "brpc/handshake/rdma_handshake_constants.h" +#include "brpc/transport_handshake.h" + +namespace brpc { +namespace rdma { + +using ParsedHello = RdmaConnectionInfo; + +namespace v2_wire { + +struct HelloMessage { + void Serialize(void* data) const; + void Deserialize(const void* data); + + uint16_t msg_len; + uint16_t hello_ver; + uint16_t impl_ver; + uint32_t block_size; + uint16_t sq_size; + uint16_t rq_size; + uint16_t lid; + ibv_gid gid; + uint32_t qp_num; +}; + +} // namespace v2_wire + +// RDMA adapters implement only protocol fields. HandshakeSession owns frame +// I/O, length validation, ACK exchange and resource callback ordering. +class RdmaHandshakeAdapter : public handshake::HandshakeProtocol { +public: + RdmaHandshakeAdapter(RdmaEndpoint* ep, int version) + : _ep(ep), _version(version), _remote() {} + ~RdmaHandshakeAdapter() override = default; + + int ProtocolVersion() const override { return _version; } + const handshake::FrameSpec& AckFrameSpec() const override; + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseHello( + const std::string& payload) override; + const ParsedHello& remote() const { return _remote; } + + const handshake::FrameSpec& HelloFrameSpec() const override = 0; + virtual handshake::StepResult BuildLocalHello( + bool enabled, std::string* payload) = 0; + virtual handshake::StepResult ParseRemoteHello( + const std::string& payload, ParsedHello* remote) = 0; + +protected: + void FillLocalHello(ParsedHello* local) const; + void PrepareClientEce(); + + RdmaEndpoint* _ep; + int _version; + ParsedHello _remote; + +private: + DISALLOW_COPY_AND_ASSIGN(RdmaHandshakeAdapter); +}; + +class RdmaClientHandshakeAdapterV2 : public RdmaHandshakeAdapter { +public: + explicit RdmaClientHandshakeAdapterV2(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 2) {} + const handshake::FrameSpec& HelloFrameSpec() const override; + handshake::StepResult BuildLocalHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseRemoteHello( + const std::string& payload, ParsedHello* remote) override; +}; + +class RdmaServerHandshakeAdapterV2 : public RdmaHandshakeAdapter { +public: + explicit RdmaServerHandshakeAdapterV2(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 2) {} + const handshake::FrameSpec& HelloFrameSpec() const override; + handshake::StepResult BuildLocalHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseRemoteHello( + const std::string& payload, ParsedHello* remote) override; +}; + +class RdmaClientHandshakeAdapterV3 : public RdmaHandshakeAdapter { +public: + explicit RdmaClientHandshakeAdapterV3(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 3) {} + const handshake::FrameSpec& HelloFrameSpec() const override; + handshake::StepResult BuildLocalHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseRemoteHello( + const std::string& payload, ParsedHello* remote) override; +}; + +class RdmaServerHandshakeAdapterV3 : public RdmaHandshakeAdapter { +public: + explicit RdmaServerHandshakeAdapterV3(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 3) {} + const handshake::FrameSpec& HelloFrameSpec() const override; + handshake::StepResult BuildLocalHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseRemoteHello( + const std::string& payload, ParsedHello* remote) override; +}; + +std::unique_ptr CreateClientHandshakeAdapter( + RdmaEndpoint* ep); + +std::vector > +CreateServerHandshakeAdapters(RdmaEndpoint* ep); + +} // namespace rdma +} // namespace brpc + +#endif // BRPC_WITH_RDMA +#endif // BRPC_HANDSHAKE_RDMA_HANDSHAKE_H diff --git a/src/brpc/rdma/rdma_handshake_constants.h b/src/brpc/handshake/rdma_handshake_constants.h similarity index 69% rename from src/brpc/rdma/rdma_handshake_constants.h rename to src/brpc/handshake/rdma_handshake_constants.h index aa9811b98e..af60d2e9d3 100644 --- a/src/brpc/rdma/rdma_handshake_constants.h +++ b/src/brpc/handshake/rdma_handshake_constants.h @@ -15,8 +15,13 @@ // specific language governing permissions and limitations // under the License. -#ifndef BRPC_RDMA_RDMA_HANDSHAKE_CONSTANTS_H -#define BRPC_RDMA_RDMA_HANDSHAKE_CONSTANTS_H +#ifndef BRPC_HANDSHAKE_RDMA_HANDSHAKE_CONSTANTS_H +#define BRPC_HANDSHAKE_RDMA_HANDSHAKE_CONSTANTS_H + +#include +#include + +#include "brpc/handshake/handshake_frame.h" namespace brpc { namespace rdma { @@ -50,7 +55,27 @@ constexpr size_t HELLO_V3_MAX_PB_SIZE = 8192; constexpr size_t HELLO_ACK_LEN = 4; constexpr uint32_t HELLO_ACK_RDMA_OK = 0x1; +inline const handshake::FrameSpec& RdmaHelloFrameSpec(int version) { + static const handshake::FrameSpec v2( + HELLO_MAGIC, HELLO_MAGIC_LEN, + HELLO_V2_MSG_LEN_MIN, HELLO_V2_MSG_LEN_MAX, + handshake::FrameSpec::U16_TOTAL_LENGTH); + static const handshake::FrameSpec v3( + HELLO_MAGIC_V3, HELLO_MAGIC_LEN, + HELLO_MAGIC_LEN + HELLO_V3_PB_SIZE_LEN + 1, + HELLO_MAGIC_LEN + HELLO_V3_PB_SIZE_LEN + HELLO_V3_MAX_PB_SIZE, + handshake::FrameSpec::U32_BODY_LENGTH); + return version == 2 ? v2 : v3; +} + +inline const handshake::FrameSpec& RdmaAckFrameSpec() { + static const handshake::FrameSpec spec( + NULL, 0, HELLO_ACK_LEN, HELLO_ACK_LEN, + handshake::FrameSpec::FIXED); + return spec; +} + } // namespace rdma } // namespace brpc -#endif // BRPC_RDMA_RDMA_HANDSHAKE_CONSTANTS_H +#endif // BRPC_HANDSHAKE_RDMA_HANDSHAKE_CONSTANTS_H diff --git a/src/brpc/handshake/ubshm_handshake.cpp b/src/brpc/handshake/ubshm_handshake.cpp new file mode 100644 index 0000000000..a409915d0f --- /dev/null +++ b/src/brpc/handshake/ubshm_handshake.cpp @@ -0,0 +1,454 @@ +// 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/handshake/ubshm_handshake.h" + +#include +#include + +#include "butil/raw_pack.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/socket.h" + +#if BRPC_WITH_UBRING + +#include +#include + +#include "butil/logging.h" +#include "brpc/reloadable_flags.h" +#include "brpc/ubshm/common/common.h" +#include "brpc/ubshm/ub_endpoint.h" +#include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/ubr_trx.h" +#include "brpc/ubshm_transport.h" + +#endif + +namespace brpc { +namespace handshake { +namespace ubshm_wire { + +static const char* const MAGIC = "UB"; +static const size_t MAGIC_LEN = 2; +static const size_t HELLO_LEN = 64; +static const size_t ACK_LEN = 4; +#if BRPC_WITH_UBRING +static const uint16_t HELLO_VERSION = 3; +static const uint16_t IMPL_VERSION = 1; +#endif // BRPC_WITH_UBRING +static const FrameSpec& HelloFrameSpec() { + static const FrameSpec spec( + MAGIC, MAGIC_LEN, HELLO_LEN, HELLO_LEN, FrameSpec::FIXED); + return spec; +} + +static const FrameSpec& AckFrameSpec() { + static const FrameSpec spec( + NULL, 0, ACK_LEN, ACK_LEN, FrameSpec::FIXED); + return spec; +} + +} // namespace ubshm_wire +} // namespace handshake +} // namespace brpc + +#if BRPC_WITH_UBRING + +namespace brpc { +namespace ubring { + +DEFINE_int32(data_queue_size, 4, "data queue size for UB"); +DEFINE_bool(ub_trace_verbose, false, "Print log message verbosely"); +BRPC_VALIDATE_GFLAG(ub_trace_verbose, brpc::PassValidate); + +void HelloMessage::Serialize(void* data) const { + char* current_pos = static_cast(data); + const uint16_t net_msg_len = butil::HostToNet16(msg_len); + memcpy(current_pos, &net_msg_len, sizeof(net_msg_len)); + current_pos += sizeof(net_msg_len); + const uint16_t net_hello_ver = butil::HostToNet16(hello_ver); + memcpy(current_pos, &net_hello_ver, sizeof(net_hello_ver)); + current_pos += sizeof(net_hello_ver); + const uint16_t net_impl_ver = butil::HostToNet16(impl_ver); + memcpy(current_pos, &net_impl_ver, sizeof(net_impl_ver)); + current_pos += sizeof(net_impl_ver); + const uint64_t net_len = butil::HostToNet64(len); + memcpy(current_pos, &net_len, sizeof(net_len)); + current_pos += sizeof(net_len); + memcpy(current_pos, shm_name, SHM_MAX_NAME_BUFF_LEN); +} + +void HelloFormatExtension::Serialize(void* data) const { + char* current = static_cast(data); + const uint16_t length = butil::HostToNet16(extension_len); + const uint16_t format = butil::HostToNet16(format_id); + memcpy(current, &length, sizeof(length)); + memcpy(current + sizeof(length), &format, sizeof(format)); +} + +void HelloFormatExtension::Deserialize(const void* data) { + const char* current = static_cast(data); + uint16_t length; + uint16_t format; + memcpy(&length, current, sizeof(length)); + memcpy(&format, current + sizeof(length), sizeof(format)); + extension_len = butil::NetToHost16(length); + format_id = butil::NetToHost16(format); +} + +void HelloMessage::Deserialize(const void* data) { + const char* current_pos = static_cast(data); + uint16_t net_msg_len; + memcpy(&net_msg_len, current_pos, sizeof(net_msg_len)); + msg_len = butil::NetToHost16(net_msg_len); + current_pos += sizeof(net_msg_len); + uint16_t net_hello_ver; + memcpy(&net_hello_ver, current_pos, sizeof(net_hello_ver)); + hello_ver = butil::NetToHost16(net_hello_ver); + current_pos += sizeof(net_hello_ver); + uint16_t net_impl_ver; + memcpy(&net_impl_ver, current_pos, sizeof(net_impl_ver)); + impl_ver = butil::NetToHost16(net_impl_ver); + current_pos += sizeof(net_impl_ver); + uint64_t net_len; + memcpy(&net_len, current_pos, sizeof(net_len)); + len = butil::NetToHost64(net_len); + current_pos += sizeof(net_len); + memcpy(shm_name, current_pos, SHM_MAX_NAME_BUFF_LEN); +} + +std::string HelloMessage::toString() const { + constexpr size_t MAX_LEN = + 16 + 6 + 16 + 6 + 16 + 6 + 20 + 6 + SHM_MAX_NAME_BUFF_LEN + 32; + std::array buf; + const int n = snprintf( + buf.data(), buf.size(), + "msg_len=%u, hello_ver=%u, impl_ver=%u, len=%lu, shm_name=%.*s", + msg_len, hello_ver, impl_ver, + static_cast(len), + static_cast(SHM_MAX_NAME_BUFF_LEN), shm_name); + return std::string(buf.data(), static_cast(n)); +} + +const handshake::FrameSpec& UBShmHandshakeAdapter::HelloFrameSpec() const { + return handshake::ubshm_wire::HelloFrameSpec(); +} + +const handshake::FrameSpec& UBShmHandshakeAdapter::AckFrameSpec() const { + return handshake::ubshm_wire::AckFrameSpec(); +} + +const handshake::FrameSpec& +UBShmHandshakeAdapter::ExtensionFrameSpec() const { + static const handshake::FrameSpec spec( + NULL, 0, HelloFormatExtension::WIRE_SIZE, + HelloFormatExtension::WIRE_SIZE, handshake::FrameSpec::FIXED); + return spec; +} + +void UBShmHandshakeAdapter::ConfigureClientHello( + uint64_t len, const char* shm_name) { + _local_len = len; + _local_name = shm_name != NULL ? shm_name : ""; + _server_reply = false; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + const char* name = NULL; + if (enabled) { + name = _server_reply ? _remote.shm_name : _local_name.c_str(); + } + return BuildHello(enabled, enabled ? _local_len : 0, name, payload); +} + +handshake::StepResult UBShmHandshakeAdapter::ParseHello( + const std::string& payload) { + return ParseHello(payload, &_remote); +} + +handshake::StepResult UBShmHandshakeAdapter::BuildExtension( + bool enabled, std::string* payload) { + const HelloFormatExtension extension = { + HelloFormatExtension::WIRE_SIZE, + static_cast(enabled ? UBR_DATA_FORMAT_LEGACY_64 + : UBR_DATA_FORMAT_NONE)}; + payload->resize(HelloFormatExtension::WIRE_SIZE); + extension.Serialize(&(*payload)[0]); + return handshake::STEP_OK; +} + +handshake::StepResult UBShmHandshakeAdapter::ParseExtension( + const std::string& payload) { + if (payload.size() != HelloFormatExtension::WIRE_SIZE) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + HelloFormatExtension extension{}; + extension.Deserialize(payload.data()); + return extension.extension_len == HelloFormatExtension::WIRE_SIZE && + extension.format_id == UBR_DATA_FORMAT_LEGACY_64 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, uint64_t len, const char* shm_name, + std::string* payload) const { + HelloMessage message{}; + message.msg_len = static_cast( + handshake::ubshm_wire::HELLO_LEN); + if (enabled) { + message.hello_ver = handshake::ubshm_wire::HELLO_VERSION; + message.impl_ver = handshake::ubshm_wire::IMPL_VERSION; + message.len = len; + if (shm_name == NULL) { + errno = EINVAL; + return handshake::STEP_ERROR; + } + const size_t shm_name_len = + strnlen(shm_name, SHM_MAX_NAME_LEN); + memcpy(message.shm_name, shm_name, shm_name_len); + } + payload->assign( + handshake::ubshm_wire::HELLO_LEN - + handshake::ubshm_wire::MAGIC_LEN, + '\0'); + message.Serialize(&(*payload)[0]); + return handshake::STEP_OK; +} + +handshake::StepResult UBShmHandshakeAdapter::ParseHello( + const std::string& payload, HelloMessage* message) const { + if (payload.size() != handshake::ubshm_wire::HELLO_LEN - + handshake::ubshm_wire::MAGIC_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + message->Deserialize(payload.data()); + if (message->msg_len != handshake::ubshm_wire::HELLO_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + if (!NegotiationValid(*message)) { + return handshake::STEP_FALLBACK; + } + if (strnlen(message->shm_name, SHM_MAX_NAME_BUFF_LEN) == + SHM_MAX_NAME_BUFF_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + return handshake::STEP_OK; +} + +bool UBShmHandshakeAdapter::NegotiationValid( + const HelloMessage& message) const { + return message.hello_ver == handshake::ubshm_wire::HELLO_VERSION && + message.impl_ver == handshake::ubshm_wire::IMPL_VERSION; +} + +} // namespace ubring +} // namespace brpc + +#endif // BRPC_WITH_UBRING + +namespace brpc { +namespace handshake { + +class UBShmServerHandshakeAdapter : public StandardHandshakeAdapter { +public: + UBShmServerHandshakeAdapter() = default; + +protected: + StepResult RunServerStep( + butil::IOBuf* source, Socket* socket) override; + HandshakeSession* GetSession(Socket* socket) const override; + +private: + StepResult RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket); +#if BRPC_WITH_UBRING + StepResult RunUBShmServerHandshake( + butil::IOBuf* source, Socket* socket); +#endif + + DISALLOW_COPY_AND_ASSIGN(UBShmServerHandshakeAdapter); +}; + + +class UBShmFallbackProtocol : public FallbackHandshakeProtocol { +public: + int ProtocolVersion() const override { return 2; } + const FrameSpec& HelloFrameSpec() const override { + return ubshm_wire::HelloFrameSpec(); + } + const FrameSpec& AckFrameSpec() const override { + return ubshm_wire::AckFrameSpec(); + } + StepResult BuildHello(bool enabled, std::string* payload) override { + if (enabled) { + errno = EPROTO; + return STEP_ERROR; + } + payload->assign( + ubshm_wire::HELLO_LEN - ubshm_wire::MAGIC_LEN, '\0'); + butil::RawPacker(&(*payload)[0]) + .pack16(static_cast(ubshm_wire::HELLO_LEN)); + return STEP_OK; + } +}; + +HandshakeAdapter* GetUBShmServerHandshakeAdapter() { + static UBShmServerHandshakeAdapter adapter; + return &adapter; +} + +HandshakeSession* UBShmServerHandshakeAdapter::GetSession( + Socket* socket) const { + return AdapterTransport::Get(socket)->handshake_session(); +} + +StepResult UBShmServerHandshakeAdapter::RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket) { + IOBufHandshakeInput input(source); + UBShmFallbackProtocol protocol; + std::vector protocols(1, &protocol); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, &input, &transport, false); +} + +#if BRPC_WITH_UBRING +class UBShmServerHandshakeTransport : public HandshakeTransport { +public: + UBShmServerHandshakeTransport( + UBShmTransport* transport, ubring::UBShmHandshakeAdapter* protocol, + Socket* socket, butil::IOBuf* source) + : _transport(transport), _protocol(protocol), _socket(socket), + _source(source) {} + + void OnProtocolSelected(HandshakeProtocol*) override { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "server receive handshake message : " + << _protocol->remote().toString(); + } + + StepResult PrepareResources() override { + if (!ubring::IsUBAvailable()) { + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + const ubring::HelloMessage& remote = _protocol->remote(); + const size_t remote_name_len = + strnlen(remote.shm_name, SHM_MAX_NAME_BUFF_LEN); + + ubring::SHM remote_trx_shm = { + NULL, remote.len, 0, {0}, + static_cast(_socket->fd())}; + memcpy(remote_trx_shm.name, remote.shm_name, remote_name_len + 1); + + const size_t local_shm_len = + static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + ubring::SHM local_trx_shm = { + NULL, local_shm_len, 0, {0}, + static_cast(_socket->fd())}; + char client_name[SHM_MAX_NAME_BUFF_LEN + 1]; + memcpy(client_name, remote.shm_name, SHM_MAX_NAME_BUFF_LEN); + client_name[SHM_MAX_NAME_BUFF_LEN] = '\0'; + char* client_ip_port = strrchr(client_name, '_'); + if (client_ip_port != NULL) { + *client_ip_port = '\0'; + } + const int result = snprintf( + local_trx_shm.name, SHM_MAX_NAME_BUFF_LEN, "%s_%s", + client_name, SERVER_SHM_NAME_SUFFIX); + if (UNLIKELY(result < 0)) { + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + if (_transport->PrepareServerUpgradeResources( + &remote_trx_shm, &local_trx_shm) < 0) { + LOG(WARNING) + << "Fail to allocate ub resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + return STEP_OK; + } + + StepResult NegotiateResources() override { return STEP_OK; } + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + void OnFailed() override { _transport->DeactivateUpgrade(); } + + StepResult ValidateEstablished() override { + return _source->empty() ? STEP_OK : STEP_ERROR; + } + +private: + UBShmTransport* _transport; + ubring::UBShmHandshakeAdapter* _protocol; + Socket* _socket; + butil::IOBuf* _source; +}; + +StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( + butil::IOBuf* source, Socket* socket) { + UBShmTransport* transport = UBShmTransport::Get(socket); + CHECK(transport->GetUBShmEp() != NULL); + + ubring::UBShmHandshakeAdapter wire; + const uint64_t local_shm_len = + static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + wire.ConfigureServerReply(local_shm_len); + IOBufHandshakeInput input(source); + std::vector protocols(1, &wire); + UBShmServerHandshakeTransport participant( + transport, &wire, socket, source); + const StepResult result = GetSession(socket)->RunServer( + protocols, &input, &participant, false); + if (result == STEP_OK) { + transport->GetUBShmEp()->SetNegotiatedDataFormat( + ubring::UBR_DATA_FORMAT_LEGACY_64); + transport->FinishUpgrade(); + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "Server handshake ends (use ubring) on " + << socket->description(); + } else if (result == STEP_FALLBACK) { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "Server handshake ends (use tcp) on " + << socket->description(); + } + return result; +} +#endif + +StepResult UBShmServerHandshakeAdapter::RunServerStep( + butil::IOBuf* source, Socket* socket) { +#if BRPC_WITH_UBRING + if (AdapterTransport::Get(socket)->upgrade_capable( + SOCKET_MODE_UBRING)) { + return RunUBShmServerHandshake(source, socket); + } +#endif + return RunFallbackServerHandshake(source, socket); +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/handshake/ubshm_handshake.h b/src/brpc/handshake/ubshm_handshake.h new file mode 100644 index 0000000000..f303c28fd1 --- /dev/null +++ b/src/brpc/handshake/ubshm_handshake.h @@ -0,0 +1,126 @@ +// 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_HANDSHAKE_UBSHM_HANDSHAKE_H +#define BRPC_HANDSHAKE_UBSHM_HANDSHAKE_H + +#include "brpc/handshake/handshake_adapter.h" + +namespace brpc { +namespace handshake { + +// Returns the adapter used by the common transport-handshake policy parser. +// The concrete server executor is private to the implementation. +HandshakeAdapter* GetUBShmServerHandshakeAdapter(); + +} // namespace handshake +} // namespace brpc + +#if BRPC_WITH_UBRING + +#include +#include + +#include + +#include "butil/macros.h" +#include "brpc/transport_handshake.h" +#include "brpc/ubshm/shm/shm_def.h" + +namespace brpc { +namespace ubring { + +DECLARE_int32(data_queue_size); +DECLARE_bool(ub_trace_verbose); + +// UBSHM v3 wire payload. HandshakeSession owns framing and ACK exchange; +// this type and UBShmHandshakeAdapter only handle protocol fields. +struct HelloMessage { + void Serialize(void* data) const; + void Deserialize(const void* data); + std::string toString() const; + + uint16_t msg_len; + uint16_t hello_ver; + uint16_t impl_ver; + uint64_t len; + char shm_name[SHM_MAX_NAME_BUFF_LEN]; +}; + +enum UbrDataFormat { + UBR_DATA_FORMAT_NONE = 0, + UBR_DATA_FORMAT_LEGACY_64 = 1, +}; + +struct HelloFormatExtension { + static const uint16_t WIRE_SIZE = 4; + uint16_t extension_len; + uint16_t format_id; + + void Serialize(void* data) const; + void Deserialize(const void* data); +}; + +class UBShmHandshakeAdapter : public handshake::HandshakeProtocol { +public: + UBShmHandshakeAdapter() + : _local_len(0), _server_reply(false), _remote() {} + + int ProtocolVersion() const override { return 3; } + const handshake::FrameSpec& HelloFrameSpec() const override; + const handshake::FrameSpec& AckFrameSpec() const override; + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseHello( + const std::string& payload) override; + bool HasExtension() const override { return true; } + const handshake::FrameSpec& ExtensionFrameSpec() const override; + handshake::StepResult BuildExtension( + bool enabled, std::string* payload) override; + handshake::StepResult ParseExtension( + const std::string& payload) override; + + void ConfigureClientHello(uint64_t len, const char* shm_name); + void ConfigureClientHello(const SHM& local_shm) { + ConfigureClientHello(local_shm.len, local_shm.name); + } + void ConfigureServerReply(uint64_t len) { + _local_len = len; + _server_reply = true; + } + const HelloMessage& remote() const { return _remote; } + + handshake::StepResult BuildHello( + bool enabled, uint64_t len, const char* shm_name, + std::string* payload) const; + handshake::StepResult ParseHello( + const std::string& payload, HelloMessage* message) const; + +private: + bool NegotiationValid(const HelloMessage& message) const; + uint64_t _local_len; + std::string _local_name; + bool _server_reply; + HelloMessage _remote; + DISALLOW_COPY_AND_ASSIGN(UBShmHandshakeAdapter); +}; + +} // namespace ubring +} // namespace brpc + +#endif // BRPC_WITH_UBRING +#endif // BRPC_HANDSHAKE_UBSHM_HANDSHAKE_H diff --git a/src/brpc/input_messenger.cpp b/src/brpc/input_messenger.cpp index 50dc1821f7..7b0544503b 100644 --- a/src/brpc/input_messenger.cpp +++ b/src/brpc/input_messenger.cpp @@ -101,6 +101,11 @@ void InputMessageClosure::reset(InputMessageBase* m) { } void InputMessenger::OnNewMessages(Socket* m) { + OnNewMessagesUntil(m, nullptr); +} + +void InputMessenger::OnNewMessagesUntil( + Socket* m, bool (*stop_reading)(Socket*)) { // Notes: // - If the socket has only one message, the message will be parsed and // processed in this bthread. nova-pbrpc and http works in this way. @@ -120,6 +125,9 @@ void InputMessenger::OnNewMessages(Socket* m) { InputMessageClosure last_msg; bool read_eof = false; while (!read_eof) { + if (stop_reading != nullptr && stop_reading(m)) { + return; + } const int64_t received_us = butil::cpuwide_time_us(); const int64_t base_realtime = butil::gettimeofday_us() - received_us; diff --git a/src/brpc/input_messenger.h b/src/brpc/input_messenger.h index c0041d6f8a..b4fea80f95 100644 --- a/src/brpc/input_messenger.h +++ b/src/brpc/input_messenger.h @@ -40,6 +40,7 @@ class UBShmEndpoint; } class TcpTransport; class RdmaTransport; +class AdapterTransport; struct InputMessageHandler { // The callback to cut a message from `source'. // Returned message will be passed to process_request or process_response @@ -102,6 +103,7 @@ class InputMessageClosure { class InputMessenger : public SocketUser { friend class TcpTransport; friend class RdmaTransport; +friend class AdapterTransport; friend class rdma::RdmaEndpoint; friend class urma::UrmaEndpoint; friend class ubring::UBShmEndpoint; @@ -141,6 +143,8 @@ friend class InputMessengerProcessor; // Load data from m->fd() into m->read_buf, cut off new messages and // call callbacks. static void OnNewMessages(Socket* m); + // Stop draining the TCP fd when a parser changes the transport state. + static void OnNewMessagesUntil(Socket* m, bool (*stop_reading)(Socket*)); private: diff --git a/src/brpc/policy/rdma_handshake_protocol.cpp b/src/brpc/policy/rdma_handshake_protocol.cpp index 580abdda5d..fb1e183f1e 100644 --- a/src/brpc/policy/rdma_handshake_protocol.cpp +++ b/src/brpc/policy/rdma_handshake_protocol.cpp @@ -17,24 +17,16 @@ #include "brpc/policy/rdma_handshake_protocol.h" -#include "butil/logging.h" -#include "brpc/destroyable.h" -#include "brpc/rdma/rdma_handshake_server.h" - namespace brpc { namespace policy { ParseResult ParseRdmaHandshake(butil::IOBuf* source, Socket* socket, - bool /*read_eof*/, const void* /*arg*/) { - return rdma::ExecuteServerHandshake(source, socket); + bool read_eof, const void* arg) { + return ParseTransportHandshake(source, socket, read_eof, arg); } void ProcessRdmaHandshake(InputMessageBase* msg) { - // ParseRdmaHandshake replies inline and only ever returns - // NOT_ENOUGH_DATA / TRY_OTHERS / hard errors, never a real message, so this - // must never run. Keep a placeholder (required for server registration). - DestroyingPtr destroying_msg(msg); - CHECK(false) << "ProcessRdmaHandshake should never be called"; + ProcessTransportHandshake(msg); } } // namespace policy diff --git a/src/brpc/policy/rdma_handshake_protocol.h b/src/brpc/policy/rdma_handshake_protocol.h index e569a4c333..51416e2eee 100644 --- a/src/brpc/policy/rdma_handshake_protocol.h +++ b/src/brpc/policy/rdma_handshake_protocol.h @@ -18,36 +18,15 @@ #ifndef BRPC_POLICY_RDMA_HANDSHAKE_PROTOCOL_H #define BRPC_POLICY_RDMA_HANDSHAKE_PROTOCOL_H -// NOTE: This file is intentionally INDEPENDENT of BRPC_WITH_RDMA. A server may -// run in TCP mode either because it was built without RDMA, or because RDMA was -// not enabled at runtime. In both cases an RDMA client that connects to it will -// send an RDMA handshake magic ("RDMA" for v2, "RDM3" for v3) first. Without -// special handling the server treats those bytes as an unknown protocol and -// closes the connection, so the client (blocked reading the server hello) only -// sees EOF and cannot fall back to TCP. -// -// To let the client fall back on the SAME connection, the server recognizes the -// RDMA handshake as a "first-class" protocol (magic in the first 4 bytes, -// PROTOCOL_RDMA_HANDSHAKE is ordered before PROTOCOL_HTTP), replies a hello with -// an incompatible version so the client rejects it and downgrades to TCP, then -// drains the client's subsequent ACK and lets normal RPC parsing continue. - -#include "butil/iobuf.h" -#include "brpc/input_message_base.h" -#include "brpc/parse_result.h" -#include "brpc/socket.h" +// Compatibility facade. New code should include +// transport_handshake_protocol.h and use ParseTransportHandshake. +#include "brpc/policy/transport_handshake_protocol.h" namespace brpc { namespace policy { -// Parse binary format of rdma handshake. ParseResult ParseRdmaHandshake(butil::IOBuf* source, Socket* socket, bool read_eof, const void* arg); - -// Actions to a rdma handshake request, which is left unimplemented. -// All requests are processed in the parsing process. This function -// must be declared since server only enables rdma handshake as a -// server-side protocol when this function is declared. void ProcessRdmaHandshake(InputMessageBase* msg); } // namespace policy diff --git a/src/brpc/policy/transport_handshake_protocol.cpp b/src/brpc/policy/transport_handshake_protocol.cpp new file mode 100644 index 0000000000..bdbd7da6cb --- /dev/null +++ b/src/brpc/policy/transport_handshake_protocol.cpp @@ -0,0 +1,43 @@ +// 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/policy/transport_handshake_protocol.h" + +#include "butil/logging.h" +#include "brpc/adapter_transport.h" +#include "brpc/socket.h" + +namespace brpc { +namespace policy { + +ParseResult ParseTransportHandshake(butil::IOBuf* source, Socket* socket, + bool /*read_eof*/, const void* /*arg*/) { + // URMA still uses its own Transport. Do not cast it to AdapterTransport + // when this globally registered parser is tried during protocol detection. + if (socket->socket_mode() == SOCKET_MODE_URMA) { + return MakeParseError(PARSE_ERROR_TRY_OTHERS); + } + return AdapterTransport::Get(socket)->ProcessUpgradeReadable(source); +} + +void ProcessTransportHandshake(InputMessageBase* msg) { + DestroyingPtr destroying_msg(msg); + CHECK(false) << "ProcessTransportHandshake should never be called"; +} + +} // namespace policy +} // namespace brpc diff --git a/src/brpc/policy/transport_handshake_protocol.h b/src/brpc/policy/transport_handshake_protocol.h new file mode 100644 index 0000000000..f4a1fe9acc --- /dev/null +++ b/src/brpc/policy/transport_handshake_protocol.h @@ -0,0 +1,44 @@ +// 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_POLICY_TRANSPORT_HANDSHAKE_PROTOCOL_H +#define BRPC_POLICY_TRANSPORT_HANDSHAKE_PROTOCOL_H + +// This policy is intentionally independent of BRPC_WITH_RDMA and +// BRPC_WITH_UBRING. A plain TCP server must recognize an upgrade hello and +// return a disabled hello so the client can continue with TCP on the same +// connection. + +#include "butil/iobuf.h" +#include "brpc/input_message_base.h" +#include "brpc/parse_result.h" +#include "brpc/socket.h" + +namespace brpc { +namespace policy { + +ParseResult ParseTransportHandshake(butil::IOBuf* source, Socket* socket, + bool read_eof, const void* arg); + +// Upgrade handshakes are completed inline by the parser. This placeholder is +// required for server-side protocol registration and must never be invoked. +void ProcessTransportHandshake(InputMessageBase* msg); + +} // namespace policy +} // namespace brpc + +#endif // BRPC_POLICY_TRANSPORT_HANDSHAKE_PROTOCOL_H diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index 0dc41c9dd1..1711ee752c 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -17,22 +17,20 @@ #if BRPC_WITH_RDMA -#include -#include "butil/fd_utility.h" -#include "butil/logging.h" // CHECK, LOG -#include "butil/sys_byteorder.h" // HostToNet,NetToHost -#include "bthread/bthread.h" +#include "brpc/rdma/rdma_endpoint.h" #include "brpc/errno.pb.h" #include "brpc/event_dispatcher.h" #include "brpc/input_messenger.h" -#include "brpc/socket.h" -#include "brpc/reloadable_flags.h" #include "brpc/rdma/block_pool.h" #include "brpc/rdma/rdma_helper.h" -#include "brpc/rdma/rdma_endpoint.h" #include "brpc/rdma_transport.h" -#include "brpc/rdma/rdma_handshake.h" -#include "brpc/rdma/rdma_handshake_constants.h" +#include "brpc/reloadable_flags.h" +#include "brpc/socket.h" +#include "bthread/bthread.h" +#include "butil/fd_utility.h" +#include "butil/logging.h" // CHECK, LOG +#include "butil/sys_byteorder.h" // HostToNet,NetToHost +#include DECLARE_int32(task_group_ntags); @@ -129,8 +127,6 @@ RdmaResource::~RdmaResource() { RdmaEndpoint::RdmaEndpoint(Socket* s) : _socket(s) - , _state(UNINIT) - , _handshake_version(0) , _resource(nullptr) , _send_cq_events(0) , _recv_cq_events(0) @@ -164,20 +160,16 @@ RdmaEndpoint::RdmaEndpoint(Socket* s) if (_rq_size > MAX_QP_SIZE) { _rq_size = MAX_QP_SIZE; } - _read_butex = bthread::butex_create_checked >(); _input_processor.Init(s, InputMessengerProcessor::STREAM_RDMA_QP); } RdmaEndpoint::~RdmaEndpoint() { Reset(); - bthread::butex_destroy(_read_butex); } void RdmaEndpoint::Reset() { DeallocateResources(); - _state.store(UNINIT, butil::memory_order_relaxed); - _handshake_version = 0; _outgoing_ece.reset(); _resource = nullptr; _send_cq_events = 0; @@ -185,8 +177,8 @@ void RdmaEndpoint::Reset() { _cq_sid = INVALID_SOCKET_ID; _sbuf.clear(); _rbuf.clear(); - _rbuf_data.clear(); _input_processor.Reset(); + _rbuf_data.clear(); _remote_recv_block_size = 0; _accumulated_ack = 0; _unsolicited = 0; @@ -203,583 +195,41 @@ void RdmaEndpoint::Reset() { _new_rq_wrs.store(0, butil::memory_order_relaxed); } -void RdmaConnect::StartConnect(const Socket* socket, - void (*done)(int err, void* data), - void* data) { - auto* rdma_transport = static_cast(socket->_transport.get()); - CHECK(rdma_transport->_rdma_ep != nullptr); - SocketUniquePtr s; - if (Socket::Address(socket->id(), &s) != 0) { - return; - } - if (!IsRdmaAvailable()) { - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - rdma_transport->_rdma_ep->_state.store( - RdmaEndpoint::FALLBACK_TCP, butil::memory_order_release); - done(0, data); - return; - } - _done = done; - _data = data; - bthread_t tid; - bthread_attr_t attr = BTHREAD_ATTR_NORMAL; - bthread_attr_set_name(&attr, "RdmaProcessHandshakeAtClient"); - if (bthread_start_background(&tid, &attr, - RdmaEndpoint::ProcessHandshakeAtClient, - rdma_transport->_rdma_ep) < 0) { - LOG(FATAL) << "Fail to start handshake bthread"; - Run(); - } else { - s.release(); - } -} - -void RdmaConnect::StopConnect(Socket* socket) { } - -void RdmaConnect::Run() { - _done(errno, _data); -} - -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); - - int progress = Socket::PROGRESS_INIT; - while (true) { - // Pair with release stores of FALLBACK_TCP so RDMA_OFF is visible - // before normal TCP message processing starts. - const State state = ep->_state.load(butil::memory_order_acquire); - if (state == UNINIT) { - // The connection may be closed or reset before the client starts - // handshake. This will be handled by client handshake. Ignore here. - } else if (state < ESTABLISHED) { // during handshake - 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(s); - return; - } else if (state == ESTABLISHED) { - if (!ep->HandleTcpEventAfterEstablished()) { - return; - } - } - if (!s->MoreReadEvents(&progress)) { - break; - } - } -} - -void RdmaEndpoint::OnNewDataFromTcpAtServer(Socket* s) { - auto* rdma_transport = static_cast(s->_transport.get()); - RdmaEndpoint* ep = rdma_transport->GetRdmaEp(); - CHECK(ep != nullptr); - - int progress = Socket::PROGRESS_INIT; - while (true) { - if (s->Failed()) { - return; - } - - // Pair with the release stores of ESTABLISHED / FALLBACK_TCP. - if (ep->_state.load(butil::memory_order_acquire) != ESTABLISHED) { - InputMessenger::OnNewMessages(s); - // 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 (!s->Failed() && - ep->_state.load(butil::memory_order_acquire) == ESTABLISHED && - ep->StartCqEvents() < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to start cq events on " << *s; - ep->_state.store(FAILED, butil::memory_order_relaxed); - s->SetFailed(saved_errno, "Fail to start cq events on %s: %s", - s->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 (!s->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). -// `read_once(offset, remaining)` performs ONE underlying read attempt: -// returns > 0 : number of bytes consumed (added to running total); -// returns = 0 : end-of-stream (the loop fails with EEOF); -// returns < 0 : errno set; EAGAIN is handled here via butex_wait, -// any other errno bubbles up. -// `offset` is bytes already received in THIS call (initially 0); the -// callable uses it to choose the next write target (e.g. `(char*)buf -// + offset`). Callables that don't need offset (e.g. IOPortal append) -// can ignore it. -// -// Centralizes the EAGAIN/butex/EOF loop so the two ReadFromFd -// overloads below stay one-liners; any future read source (memory- -// mapped, scatter-vector, etc.) can plug in by passing its own -// `read_once`. -template -static int ReadFromFdLoop(butil::atomic* read_butex, - size_t len, ReadOnce&& read_once) { - size_t received = 0; - while (received < len) { - const int expected_val = read_butex->load(butil::memory_order_acquire); - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - ssize_t nr = read_once(received, len - received); - if (nr < 0) { - if (errno == EAGAIN) { - if (bthread::butex_wait(read_butex, expected_val, &duetime) < 0) { - if (errno != EWOULDBLOCK && errno != ETIMEDOUT) { - return -1; - } - } - } else { - return -1; - } - } else if (nr == 0) { // Got EOF - errno = EEOF; - return -1; - } else { - received += nr; - } - } - return 0; -} - -int RdmaEndpoint::ReadFromFd(void* data, size_t len) { - CHECK(data != nullptr); - const int fd = _socket->fd(); - return ReadFromFdLoop(_read_butex, len, - [data, fd](size_t offset, size_t remaining) { - return read(fd, (uint8_t*)data + offset, remaining); - }); -} - -int RdmaEndpoint::ReadFromFd(butil::IOPortal* data, size_t len) { - CHECK(data != nullptr); - const int fd = _socket->fd(); - return ReadFromFdLoop(_read_butex, len, - [data, fd](size_t /*offset*/, size_t remaining) { - return data->append_from_file_descriptor(fd, remaining); - }); -} - -// Drive an EAGAIN-aware write loop to completion (exactly `len` bytes). -// -// `write_once(offset, remaining)` performs ONE underlying write attempt: -// - returns >= 0 : number of bytes consumed (added to running total); -// - returns < 0 : errno set; EAGAIN triggers `wait_writable(duetime)`, -// any other errno bubbles up. -// `offset` is bytes already written in THIS call (initially 0); the -// callable uses it to choose the next read source (e.g. `(char*)buf -// + offset`). Callables that drain a self-tracking sink (e.g. -// IOBuf::cut_into_file_descriptor) can ignore both args. -// -// `wait_writable(duetime)` is invoked on EAGAIN to park until the fd -// becomes writable again. It returns 0 on wake-up (or ETIMEDOUT), -// non-zero on hard failure. -template -static int WriteToFdLoop(size_t len, WriteOnce&& write_once, WaitWritable&& wait_writable) { - size_t written = 0; - while (written < len) { - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - ssize_t nw = write_once(written, len - written); - if (nw >= 0) { - written += nw; - continue; - } - - if (errno != EAGAIN) { - return -1; - } - if (!wait_writable(&duetime)) { - return -1; - } - } - return 0; -} - -int RdmaEndpoint::WriteToFd(void* data, size_t len) { - CHECK(data != nullptr); - Socket* s = _socket; - const int fd = s->fd(); - return WriteToFdLoop(len, - [data, fd](size_t offset, size_t remaining) { - return write(fd, (uint8_t*)data + offset, remaining); - }, - [s, fd](const timespec* duetime) { - return s->WaitEpollOut(fd, true, duetime) == 0 || errno == ETIMEDOUT; - }); -} - -int RdmaEndpoint::WriteToFd(butil::IOBuf* data) { - CHECK(data != nullptr); - Socket* s = _socket; - const int fd = s->fd(); - return WriteToFdLoop(data->size(), - [data, fd](size_t /*offset*/, size_t /*remaining*/) { - return data->cut_into_file_descriptor(fd); - }, - [s, fd](const timespec* duetime) { - return s->WaitEpollOut(fd, true, duetime) == 0 || errno == ETIMEDOUT; - }); -} - -void RdmaEndpoint::ApplyRemoteHello(const ParsedHello& remote) { +void RdmaEndpoint::ApplyRemoteInfo(const RdmaConnectionInfo &remote) { _remote_recv_block_size = remote.block_size; _local_window_capacity = std::min(_sq_size, remote.rq_size) - RESERVED_WR_NUM; - _remote_window_capacity = std::min(_rq_size, remote.sq_size) - RESERVED_WR_NUM; + _remote_window_capacity = + std::min(_rq_size, remote.sq_size) - RESERVED_WR_NUM; _sq_imm_window_size = RESERVED_WR_NUM; - _remote_rq_window_size.store(_local_window_capacity, butil::memory_order_relaxed); + _remote_rq_window_size.store(_local_window_capacity, + butil::memory_order_relaxed); _sq_window_size.store(_local_window_capacity, butil::memory_order_relaxed); } -// Client-side handshake entry: the state machine. -// -// C_ALLOC_QPCQ -// | -// v -// C_HELLO_SEND (hs->SendLocalHello) -// | -// v -// C_HELLO_WAIT (hs->ReceiveAndParseRemoteHello) -// | -// v -// [negotiation: ApplyRemoteHello + C_BRINGUP_QP] -// | -// v -// C_ACK_SEND -// | -// v -// ESTABLISHED / FALLBACK_TCP -void* RdmaEndpoint::ProcessHandshakeAtClient(void* arg) { - auto ep = static_cast(arg); - SocketUniquePtr s(ep->_socket); - RdmaConnect::RunGuard rg((RdmaConnect*)s->_app_connect.get()); - auto rdma_transport = static_cast(s->_transport.get()); - - LOG_IF(INFO, FLAGS_rdma_trace_verbose) - << "Start handshake on " << s->description(); - - std::unique_ptr handshake = CreateClientHandshake(ep); - CHECK(handshake != nullptr); - ep->_handshake_version = handshake->ProtocolVersion(); - - // First initialize CQ and QP resources. - ep->_state.store(C_ALLOC_QPCQ, butil::memory_order_relaxed); - if (ep->AllocateResources() < 0) { - PLOG(WARNING) << "Fail to allocate rdma resources, fallback to tcp:" - << s->description(); - errno = 0; - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - ep->_state.store(FALLBACK_TCP, butil::memory_order_release); - return nullptr; - } - - // Send hello message to server - ep->_state.store(C_HELLO_SEND, butil::memory_order_relaxed); - if (handshake->SendLocalHello() < 0) { - int saved_errno = errno; - PLOG(WARNING) << "Fail to send hello message to server:" - << 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; - } - - // Receive and parse remote hello. - ep->_state.store(C_HELLO_WAIT, butil::memory_order_relaxed); - ParsedHello remote{}; - const RemoteHelloResult r = handshake->ReceiveAndParseRemoteHello(&remote); - if (r == RemoteHelloResult::ERROR) { - int saved_errno = errno; - PLOG(WARNING) << "Fail to receive hello from server:" - << 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; - } - - if (r != RemoteHelloResult::NEGOTIATED) { - LOG(WARNING) << "Fail to negotiate with server, fallback to tcp:" - << s->description(); - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - } else { - ep->ApplyRemoteHello(remote); - ep->_state.store(C_BRINGUP_QP, butil::memory_order_relaxed); - if (ep->BringUpQp(remote, /*is_server=*/false) < 0) { - LOG(WARNING) << "Fail to bringup QP, fallback to tcp:" - << s->description(); - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - } else { - rdma_transport->_rdma_state = RdmaTransport::RDMA_ON; - } - } - - // Send ACK message to server - ep->_state.store(C_ACK_SEND, butil::memory_order_relaxed); - bool rdma_on = rdma_transport->_rdma_state == RdmaTransport::RDMA_ON; - uint32_t flags = rdma_on ? HELLO_ACK_RDMA_OK : 0; - uint32_t flags_be = butil::HostToNet32(flags); - if (ep->WriteToFd(&flags_be, HELLO_ACK_LEN) < 0) { - int saved_errno = errno; - PLOG(WARNING) << "Fail to send Ack Message to server:" - << 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; - } - - if (rdma_transport->_rdma_state == RdmaTransport::RDMA_ON) { - 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(); - } else { - ep->_state.store(FALLBACK_TCP, butil::memory_order_release); - LOG_IF(INFO, FLAGS_rdma_trace_verbose) - << "Client handshake ends (use tcp) on " << s->description(); +void RdmaEndpoint::GetLocalConnectionInfo(RdmaConnectionInfo *local) const { + CHECK(local != NULL); + local->block_size = g_rdma_recv_block_size; + local->sq_size = _sq_size; + local->rq_size = _rq_size; + local->lid = GetRdmaLid(); + local->gid = GetRdmaGid(); + local->qp_num = BAIDU_LIKELY(_resource) ? _resource->qp->qp_num : 0; + local->ece.reset(); + if (_outgoing_ece.has_value()) { + local->ece = _outgoing_ece; } - - errno = 0; - - return nullptr; } -// Server-side handshake entry: the state machine. -// -// S_HELLO_WAIT (read magic + dispatch + hs->ReceiveAndParseRemoteHello) -// | -// v -// [negotiation: ApplyRemoteHello + S_ALLOC_QPCQ + S_BRINGUP_QP] -// | -// v -// S_HELLO_SEND (hs->SendLocalHello) -// | -// v -// S_ACK_WAIT -// | -// v -// ESTABLISHED / FALLBACK_TCP -ParseResult RdmaEndpoint::ExecuteServerHandshake(butil::IOBuf* source, Socket* s) { - RdmaTransport* rdma_transport = static_cast(s->_transport.get()); - 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) { - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - uint8_t magic[HELLO_MAGIC_LEN]; - CHECK_EQ(source->copy_to(magic, HELLO_MAGIC_LEN), HELLO_MAGIC_LEN); - - // Pick the version-specific server handshake from the peeked magic (the - // magic is NOT consumed; ReceiveAndParseRemoteHello() reads it again - // from `source`). - std::unique_ptr hs = CreateServerHandshakeByMagic(ep, source, magic); - if (hs == nullptr) { - return MakeParseError(PARSE_ERROR_TRY_OTHERS); - } - ep->_handshake_version = hs->ProtocolVersion(); - ep->_state.store(S_HELLO_WAIT, butil::memory_order_relaxed); - - ParsedHello remote{}; - const RemoteHelloResult r = hs->ReceiveAndParseRemoteHello(&remote); - if (r == RemoteHelloResult::NEED_MORE) { - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - if (r == RemoteHelloResult::ERROR) { - ep->_state.store(FAILED, butil::memory_order_relaxed); - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } - - // Negotiate + allocate resources. - bool negotiated = r == RemoteHelloResult::NEGOTIATED; - if (negotiated) { - ep->ApplyRemoteHello(remote); - ep->_state.store(S_ALLOC_QPCQ, butil::memory_order_relaxed); - if (ep->AllocateResources() < 0) { - PLOG(WARNING) << "Fail to allocate rdma resources, fallback to tcp:" - << s->description(); - negotiated = false; - } else { - ep->_state.store(S_BRINGUP_QP, butil::memory_order_relaxed); - if (ep->BringUpQp(remote, /*is_server=*/true) < 0) { - LOG(WARNING) << "Fail to bringup QP, fallback to tcp:" - << s->description(); - negotiated = false; - } - } - } - if (!negotiated) { - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - } - - // Reply the server hello. - // Emits a real hello when _rdma_state != RDMA_OFF; - // an un-negotiable one otherwise. - ep->_state.store(S_HELLO_SEND, butil::memory_order_relaxed); - if (hs->SendLocalHello() < 0) { - PLOG(WARNING) << "Fail to send server hello to " << s->description(); - ep->_state.store(FAILED, butil::memory_order_relaxed); - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } - - // Enter the wait-ACK phase. Whether negotiation succeeded is already - // recorded in rdma_transport->_rdma_state (RDMA_OFF iff negotiation - // failed), so the context itself needs no extra flag. - s->reset_parsing_context(ServerHandshakeContext::Create()); - ep->_state.store(S_ACK_WAIT, butil::memory_order_relaxed); - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - - // Phase 2: drain the 4B ACK and finalize. - if (source->size() < HELLO_ACK_LEN) { - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); +int RdmaEndpoint::QueryLocalEce(ibv_ece *ece) const { + if (ece == NULL || IbvQueryEce == NULL || _resource == NULL || + _resource->qp == NULL) { + return 1; } + return IbvQueryEce(_resource->qp, ece) == 0 ? 0 : -1; +} - uint32_t flags_be = 0; - CHECK_EQ(source->cutn(&flags_be, HELLO_ACK_LEN), HELLO_ACK_LEN); - uint32_t flags = butil::NetToHost32(flags_be); - bool client_ack_ok = (flags & HELLO_ACK_RDMA_OK) != 0; - if (!client_ack_ok) { - LOG_IF(INFO, FLAGS_rdma_trace_verbose) - << "Server handshake ends (use tcp) on " << s->description(); - rdma_transport->_rdma_state = RdmaTransport::RDMA_OFF; - ep->_state.store(FALLBACK_TCP, butil::memory_order_release); - s->reset_parsing_context(nullptr); - return MakeParseError(PARSE_ERROR_TRY_OTHERS); - } - - if (rdma_transport->_rdma_state == RdmaTransport::RDMA_OFF) { - LOG(WARNING) << "Client wants RDMA in ACK but server fell back: " - << s->description(); - ep->_state.store(FAILED, butil::memory_order_relaxed); - s->reset_parsing_context(nullptr); - 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_release); - s->reset_parsing_context(nullptr); - - // 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); +void RdmaEndpoint::SetOutgoingEce(const ibv_ece &ece) { + _outgoing_ece = ece; } bool RdmaEndpoint::IsWritable() const { @@ -796,6 +246,7 @@ bool RdmaEndpoint::IsWritable() const { // The reason is that we need to use some protected member function of IOBuf. class RdmaIOBuf : public butil::IOBuf { friend class RdmaEndpoint; + private: // Cut the current IOBuf to ibv_sge list and `to' for at most first max_sge // blocks or first max_len bytes. `*in_pool' is cleared if any block is @@ -820,7 +271,8 @@ friend class RdmaEndpoint; lkey = (uint32_t)meta; } } - if (BAIDU_UNLIKELY(lkey == 0)) { // only happens when meta is not specified + if (BAIDU_UNLIKELY(lkey == + 0)) { // only happens when meta is not specified lkey = GetLKey((char*)start - r.offset); } if (lkey == 0) { @@ -872,8 +324,7 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { size_t current = 0; uint32_t remote_rq_window_size = _remote_rq_window_size.load(butil::memory_order_relaxed); - uint32_t sq_window_size = - _sq_window_size.load(butil::memory_order_relaxed); + uint32_t sq_window_size = _sq_window_size.load(butil::memory_order_relaxed); ibv_send_wr wr; int max_sge = GetRdmaMaxSge(); ibv_sge sglist[max_sge]; @@ -932,7 +383,8 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { wr.imm_data = butil::HostToNet32(imm); // Avoid too much recv completion event to reduce the cpu overhead bool solicited = false; - if (remote_rq_window_size == 1 || sq_window_size == 1 || current + 1 >= ndata) { + if (remote_rq_window_size == 1 || sq_window_size == 1 || + current + 1 >= ndata) { // Only last message in the write queue or last message in the // current window will be flagged as solicited. solicited = true; @@ -976,7 +428,8 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { // So we just consider this error as an unrecoverable error. std::ostringstream oss; DebugInfo(oss, ", "); - LOG(WARNING) << "Fail to ibv_post_send: " << berror(err) << " " << oss.str(); + LOG(WARNING) << "Fail to ibv_post_send: " << berror(err) << " " + << oss.str(); errno = err; return -1; } @@ -993,14 +446,16 @@ ssize_t RdmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { // counters. remote_rq_window_size = _remote_rq_window_size.fetch_sub(1, butil::memory_order_relaxed) - 1; - sq_window_size = _sq_window_size.fetch_sub(1, butil::memory_order_relaxed) - 1; + sq_window_size = + _sq_window_size.fetch_sub(1, butil::memory_order_relaxed) - 1; } return total_len; } int RdmaEndpoint::SendAck(int num) { - if (_new_rq_wrs.fetch_add(num, butil::memory_order_relaxed) > _remote_window_capacity / 2 && + if (_new_rq_wrs.fetch_add(num, butil::memory_order_relaxed) > + _remote_window_capacity / 2 && _sq_imm_window_size > 0) { return SendImm(_new_rq_wrs.exchange(0, butil::memory_order_relaxed)); } @@ -1026,7 +481,8 @@ int RdmaEndpoint::SendImm(uint32_t imm) { DebugInfo(oss, ", "); // We use other way to guarantee the Send Queue is not full. // So we just consider this error as an unrecoverable error. - LOG(WARNING) << "Fail to ibv_post_send: " << berror(err) << " " << oss.str(); + LOG(WARNING) << "Fail to ibv_post_send: " << berror(err) << " " + << oss.str(); return -1; } @@ -1072,13 +528,11 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) { if (wc.byte_len < (uint32_t)FLAGS_rdma_zerocopy_min_size) { 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(&read_buf, wc.byte_len); + _rbuf[_rq_received].cutn(&_input_processor.read_buf(), wc.byte_len); } else { // Copy data when the receive data is really small - read_buf.append(_rbuf_data[_rq_received], wc.byte_len); + _input_processor.read_buf().append(_rbuf_data[_rq_received], wc.byte_len); } } if (0 != (wc.wc_flags & IBV_WC_WITH_IMM) && wc.imm_data > 0) { @@ -1138,7 +592,8 @@ int RdmaEndpoint::PostRecv(uint32_t num, bool zerocopy) { if (zerocopy) { _rbuf[_rq_received].clear(); butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received], - g_rdma_recv_block_size + IOBUF_BLOCK_HEADER_LEN); + g_rdma_recv_block_size + + IOBUF_BLOCK_HEADER_LEN); int size = 0; if (!os.Next(&_rbuf_data[_rq_received], &size)) { // Memory is not enough for preparing a block @@ -1251,15 +706,17 @@ static RdmaResource* AllocateQpCq(uint16_t sq_size, uint16_t rq_size) { return nullptr; } - resource->send_cq = IbvCreateCq(GetRdmaContext(), FLAGS_rdma_prepared_qp_size, - nullptr, resource->comp_channel, GetRdmaCompVector()); + resource->send_cq = + IbvCreateCq(GetRdmaContext(), FLAGS_rdma_prepared_qp_size, nullptr, + resource->comp_channel, GetRdmaCompVector()); if (nullptr == resource->send_cq) { PLOG(WARNING) << "Fail to create send CQ"; return nullptr; } - resource->recv_cq = IbvCreateCq(GetRdmaContext(), FLAGS_rdma_prepared_qp_size, - nullptr, resource->comp_channel, GetRdmaCompVector()); + resource->recv_cq = + IbvCreateCq(GetRdmaContext(), FLAGS_rdma_prepared_qp_size, nullptr, + resource->comp_channel, GetRdmaCompVector()); if (nullptr == resource->recv_cq) { PLOG(WARNING) << "Fail to create recv CQ"; return nullptr; @@ -1272,8 +729,8 @@ static RdmaResource* AllocateQpCq(uint16_t sq_size, uint16_t rq_size) { return nullptr; } } else { - resource->polling_cq = - IbvCreateCq(GetRdmaContext(), 2 * FLAGS_rdma_prepared_qp_size, nullptr, nullptr, 0); + resource->polling_cq = IbvCreateCq( + GetRdmaContext(), 2 * FLAGS_rdma_prepared_qp_size, nullptr, nullptr, 0); if (nullptr == resource->polling_cq) { PLOG(WARNING) << "Fail to create polling CQ"; return nullptr; @@ -1325,12 +782,12 @@ int RdmaEndpoint::DoAllocateResources() { g_rdma_resource_list = g_rdma_resource_list->next; } } - if (_resource == nullptr) { + if (!_resource) { _resource = AllocateQpCq(_sq_size, _rq_size); } else { _resource->next = nullptr; } - if (_resource == nullptr) { + if (!_resource) { return -1; } @@ -1360,23 +817,23 @@ int RdmaEndpoint::DoAllocateResources() { } int RdmaEndpoint::StartCqEvents() { - if (InputMessengerProcessor::STREAM_NONE != _socket->parsing_stream_type()) { - LOG(WARNING) << "StartCqEvents() called while " << *_socket << " is parsing"; + if (InputMessengerProcessor::STREAM_NONE != + _socket->parsing_stream_type()) { + LOG(WARNING) << "StartCqEvents() called while " << *_socket + << " is parsing"; errno = ERDMA; return -1; } 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; + LOG(WARNING) << "No RDMA resource to start CQ events on, " + << *_socket; errno = ERDMA; return -1; } @@ -1392,15 +849,13 @@ int RdmaEndpoint::StartCqEvents() { 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) { +int RdmaEndpoint::BringUpQp(const RdmaConnectionInfo &remote, bool is_server) { if (BAIDU_UNLIKELY(g_skip_rdma_init)) { // For UT return 0; @@ -1412,11 +867,9 @@ int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { attr.pkey_index = 0; // TODO: support more pkey use in future attr.port_num = GetRdmaPortNum(); attr.qp_access_flags = IBV_ACCESS_REMOTE_WRITE; - int err = IbvModifyQp(_resource->qp, &attr, (ibv_qp_attr_mask)( - IBV_QP_STATE | - IBV_QP_PKEY_INDEX | - IBV_QP_PORT | - IBV_QP_ACCESS_FLAGS)); + int err = IbvModifyQp(_resource->qp, &attr, + (ibv_qp_attr_mask)(IBV_QP_STATE | IBV_QP_PKEY_INDEX | + IBV_QP_PORT | IBV_QP_ACCESS_FLAGS)); if (err != 0) { LOG(WARNING) << "Fail to modify QP from RESET to INIT: " << berror(err); return -1; @@ -1428,7 +881,7 @@ int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { // End-to-end model: // Server: `remote->ece' is the client's queried ECE; set it here, // then after RTS we query the reduced/negotiated ECE and - // send it back in the server hello. + // return it in the server negotiation response. // Client: `remote->ece' is the server's reduced ECE; // just set it here. bool use_ece = true; @@ -1464,14 +917,11 @@ int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { attr.rq_psn = 0; attr.max_dest_rd_atomic = 0; attr.min_rnr_timer = 0; // We do not allow rnr error - err = IbvModifyQp(_resource->qp, &attr, (ibv_qp_attr_mask)( - IBV_QP_STATE | - IBV_QP_PATH_MTU | - IBV_QP_MIN_RNR_TIMER | - IBV_QP_AV | - IBV_QP_MAX_DEST_RD_ATOMIC | - IBV_QP_DEST_QPN | - IBV_QP_RQ_PSN)); + err = IbvModifyQp(_resource->qp, &attr, + (ibv_qp_attr_mask)(IBV_QP_STATE | IBV_QP_PATH_MTU | + IBV_QP_MIN_RNR_TIMER | IBV_QP_AV | + IBV_QP_MAX_DEST_RD_ATOMIC | + IBV_QP_DEST_QPN | IBV_QP_RQ_PSN)); if (err != 0) { LOG(WARNING) << "Fail to modify QP from INIT to RTR: " << berror(err); return -1; @@ -1483,29 +933,29 @@ int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { attr.rnr_retry = 0; // We do not allow rnr error attr.sq_psn = 0; attr.max_rd_atomic = 0; - err = IbvModifyQp(_resource->qp, &attr, (ibv_qp_attr_mask)( - IBV_QP_STATE | - IBV_QP_RNR_RETRY | - IBV_QP_RETRY_CNT | - IBV_QP_TIMEOUT | - IBV_QP_SQ_PSN | - IBV_QP_MAX_QP_RD_ATOMIC)); + err = + IbvModifyQp(_resource->qp, &attr, + (ibv_qp_attr_mask)(IBV_QP_STATE | IBV_QP_RNR_RETRY | + IBV_QP_RETRY_CNT | IBV_QP_TIMEOUT | + IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC)); if (err != 0) { LOG(WARNING) << "Fail to modify QP from RTR to RTS: " << berror(err); return -1; } - // On the server side, now that the QP reached RTS, query the reduced/negotiated - // ECE (the subset of enhancements supported by both peers) so it can be returned - // to the client in the server hello. - if (is_server && use_ece && IbvQueryEce != nullptr && remote.ece.has_value()) { + // On the server side, now that the QP reached RTS, query the + // reduced/negotiated ECE (the subset of enhancements supported by both peers) + // so it can be returned to the client in the server hello. + if (is_server && use_ece && IbvQueryEce != nullptr && + remote.ece.has_value()) { ibv_ece ece; int qerr = IbvQueryEce(_resource->qp, &ece); if (qerr == 0) { _outgoing_ece = ece; } else { LOG(WARNING) << "Fail to IbvQueryEce(negotiated), " - "continue without ECE: " << berror(qerr); + "continue without ECE: " + << berror(qerr); } } @@ -1537,7 +987,7 @@ static int DrainCq(ibv_cq* cq) { } void RdmaEndpoint::DeallocateResources() { - if (_resource == nullptr) { + if (!_resource) { return; } if (FLAGS_rdma_use_polling) { @@ -1564,7 +1014,7 @@ void RdmaEndpoint::DeallocateResources() { bool remove_consumer = true; _reclaim: if (!move_to_rdma_resource_list) { - if (_resource->qp != nullptr) { + if (nullptr != _resource->qp) { int err = IbvDestroyQp(_resource->qp); LOG_IF(WARNING, 0 != err) << "Fail to destroy QP: " << berror(err); _resource->qp = nullptr; @@ -1579,11 +1029,14 @@ void RdmaEndpoint::DeallocateResources() { // 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); + 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); + LOG_IF(WARNING, 0 != err) + << "Fail to destroy CQ channel: " << berror(err); } _resource->polling_cq = nullptr; @@ -1664,8 +1117,8 @@ int RdmaEndpoint::GetAndAckEvents(SocketUniquePtr& s) { } else { // Unexpected CQ event that does not belong to // this endpoint's send/recv CQs. - LOG(WARNING) << "Unexpected CQ event from cq=" << cq - << " of " << s->description(); + LOG(WARNING) << "Unexpected CQ event from cq=" << cq << " of " + << s->description(); // Acknowledge this single event immediately // to avoid leaking unacknowledged events. IbvAckCqEvents(cq, 1); @@ -1684,16 +1137,15 @@ int RdmaEndpoint::GetAndAckEvents(SocketUniquePtr& s) { int RdmaEndpoint::ReqNotifyCq(bool send_cq, bool fatal_on_error) { const int err = ibv_req_notify_cq( - send_cq ? _resource->send_cq : _resource->recv_cq, - send_cq ? 0 : 1); + send_cq ? _resource->send_cq : _resource->recv_cq, send_cq ? 0 : 1); if (0 != err) { errno = err; PLOG(WARNING) << "Fail to arm " << (send_cq ? "send" : "recv") << " CQ comp channel from " << _socket->description(); if (fatal_on_error) { _socket->SetFailed(err, "Fail to arm %s CQ channel from %s: %s", - send_cq ? "send" : "recv", _socket->description().c_str(), - berror(err)); + send_cq ? "send" : "recv", + _socket->description().c_str(), berror(err)); } // The logging and SetFailed() above may clobber errno. errno = err; @@ -1718,9 +1170,8 @@ void RdmaEndpoint::PollCq(Socket* m) { if (m->id() != ep->_cq_sid) { return; } - auto* rdma_transport = static_cast(s->_transport.get()); + RdmaTransport *rdma_transport = RdmaTransport::Get(s.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; @@ -1838,50 +1289,30 @@ 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; - if (ep->_input_processor.ProcessNewMessage(bytes, false, received_us, - base_realtime, last_msg) < 0) { + if (ep->_input_processor.ProcessNewMessage( + bytes, false, received_us, base_realtime, last_msg) < 0) { return; } } } -std::string RdmaEndpoint::GetStateStr() const { - switch (_state.load(butil::memory_order_relaxed)) { - case UNINIT: return "UNINIT"; - case C_ALLOC_QPCQ: return "C_ALLOC_QPCQ"; - case C_HELLO_SEND: return "C_HELLO_SEND"; - case C_HELLO_WAIT: return "C_HELLO_WAIT"; - case C_BRINGUP_QP: return "C_BRINGUP_QP"; - case C_ACK_SEND: return "C_ACK_SEND"; - case S_HELLO_WAIT: return "S_HELLO_WAIT"; - case S_ALLOC_QPCQ: return "S_ALLOC_QPCQ"; - case S_BRINGUP_QP: return "S_BRINGUP_QP"; - case S_HELLO_SEND: return "S_HELLO_SEND"; - case S_ACK_WAIT: return "S_ACK_WAIT"; - case ESTABLISHED: return "ESTABLISHED"; - case FALLBACK_TCP: return "FALLBACK_TCP"; - case FAILED: return "FAILED"; - default: return "UNKNOWN"; - } -} - -void RdmaEndpoint::DebugInfo(std::ostream& os, butil::StringPiece connector) const { - os << "rdma_state=ON" - << connector << "handshake_state=" << GetStateStr() - << connector << "handshake_version=" << static_cast(_handshake_version) - << connector << "rdma_sq_imm_window_size=" << _sq_imm_window_size - << connector << "rdma_remote_rq_window_size=" << _remote_rq_window_size.load(butil::memory_order_relaxed) - << connector << "rdma_sq_window_size=" << _sq_window_size.load(butil::memory_order_relaxed) - << connector << "rdma_local_window_capacity=" << _local_window_capacity - << connector << "rdma_remote_window_capacity=" << _remote_window_capacity - << connector << "rdma_sbuf_head=" << _sq_current - << connector << "rdma_sbuf_tail=" << _sq_sent - << connector << "rdma_rbuf_head=" << _rq_received - << 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_read_buf=" << _input_processor.read_buf().size(); +void RdmaEndpoint::DebugInfo(std::ostream &os, + butil::StringPiece connector) const { + os << "rdma_state=ON" << connector + << "rdma_sq_imm_window_size=" << _sq_imm_window_size << connector + << "rdma_remote_rq_window_size=" + << _remote_rq_window_size.load(butil::memory_order_relaxed) << connector + << "rdma_sq_window_size=" + << _sq_window_size.load(butil::memory_order_relaxed) << connector + << "rdma_local_window_capacity=" << _local_window_capacity << connector + << "rdma_remote_window_capacity=" << _remote_window_capacity << connector + << "rdma_sbuf_head=" << _sq_current << connector + << "rdma_sbuf_tail=" << _sq_sent << connector + << "rdma_rbuf_head=" << _rq_received << 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; } int RdmaEndpoint::GlobalInitialize() { @@ -1895,8 +1326,8 @@ int RdmaEndpoint::GlobalInitialize() { g_rdma_resource_mutex = new butil::Mutex; for (int i = 0; i < FLAGS_rdma_prepared_qp_cnt; ++i) { - RdmaResource* res = AllocateQpCq(FLAGS_rdma_prepared_qp_size, - FLAGS_rdma_prepared_qp_size); + RdmaResource *res = + AllocateQpCq(FLAGS_rdma_prepared_qp_size, FLAGS_rdma_prepared_qp_size); if (!res) { return -1; } @@ -1991,8 +1422,8 @@ int RdmaEndpoint::PollingModeInitialize(bthread_tag_t tag, }; for (int i = 0; i < FLAGS_rdma_poller_num; ++i) { auto args = new FnArgs{&pollers[i], &running}; - auto attr = FLAGS_rdma_disable_bthread ? BTHREAD_ATTR_PTHREAD - : BTHREAD_ATTR_NORMAL; + auto attr = + FLAGS_rdma_disable_bthread ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL; attr.tag = tag; bthread_attr_set_name(&attr, "RdmaPolling"); pollers[i].callback = callback; @@ -2021,27 +1452,23 @@ 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]; - poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::ADD}); + if (INVALID_SOCKET_ID != _cq_sid) { + 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]; - poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::REMOVE}); + if (INVALID_SOCKET_ID != _cq_sid) { + 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 eda5e558ef..065ee06926 100644 --- a/src/brpc/rdma/rdma_endpoint.h +++ b/src/brpc/rdma/rdma_endpoint.h @@ -31,51 +31,27 @@ #include "butil/containers/mpsc_queue.h" #include "butil/containers/optional.h" #include "brpc/socket.h" -#include "brpc/rdma/rdma_handshake_server.h" namespace brpc { class Socket; +class RdmaTransport; namespace rdma { DECLARE_bool(rdma_use_polling); DECLARE_int32(rdma_poller_num); DECLARE_bool(rdma_disable_bthread); -class RdmaHandshakeClientV2; -class RdmaHandshakeServerV2; -class RdmaHandshakeClientV3; -class RdmaHandshakeServerV3; -struct ParsedHello; -enum class RemoteHelloResult; -class RdmaHello; -class RdmaEndpoint; -namespace v2_wire { - RemoteHelloResult ReadBodyAndNegotiate(RdmaEndpoint* ep, ParsedHello* remote); - int DrainBytes(RdmaEndpoint* ep, size_t n); -} // namespace v2_wire - -namespace v3_wire { - void FillLocalRdmaHello(const RdmaEndpoint* ep, RdmaHello* msg); - int ReadAndParseV3Hello(RdmaEndpoint* ep, RdmaHello* out); - int WriteV3Hello(RdmaEndpoint* ep, const RdmaHello& msg); -} // namespace v3_wire - -class RdmaConnect : public AppConnect { -public: - void StartConnect(const Socket* socket, - void (*done)(int err, void* data), void* data) override; - void StopConnect(Socket*) override; - struct RunGuard { - RunGuard(RdmaConnect* rc) { this_rc = rc; } - ~RunGuard() { if (this_rc) this_rc->Run(); } - RdmaConnect* this_rc; - }; - -private: - void Run(); - void (*_done)(int, void*){nullptr}; - void* _data{nullptr}; +// Wire-independent RDMA connection parameters consumed by resource setup. +// Transport adapters translate their protocol-specific payload into this DTO. +struct RdmaConnectionInfo { + uint32_t block_size; + uint16_t sq_size; + uint16_t rq_size; + uint16_t lid; + ibv_gid gid; + uint32_t qp_num; + butil::optional ece; }; struct RdmaResource { @@ -95,17 +71,8 @@ struct RdmaResource { }; class BAIDU_CACHELINE_ALIGNMENT RdmaEndpoint : public SocketUser { -friend class RdmaConnect; friend class Socket; -friend class RdmaHandshakeClientV2; -friend class RdmaHandshakeServerV2; -friend class RdmaHandshakeClientV3; -friend class RdmaHandshakeServerV3; -friend RemoteHelloResult v2_wire::ReadBodyAndNegotiate(RdmaEndpoint*, ParsedHello*); -friend int v2_wire::DrainBytes(RdmaEndpoint*, size_t); -friend void v3_wire::FillLocalRdmaHello(const RdmaEndpoint*, RdmaHello*); -friend int v3_wire::ReadAndParseV3Hello(RdmaEndpoint*, RdmaHello*); -friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); +friend class ::brpc::RdmaTransport; public: explicit RdmaEndpoint(Socket* s); ~RdmaEndpoint() override; @@ -126,16 +93,16 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Whether the endpoint can send more data bool IsWritable() const; + // Resource information consumed by the transport-level RDMA adapter. + void GetLocalConnectionInfo(RdmaConnectionInfo* local) const; + // Returns 0 on success, 1 when ECE query is unavailable, -1 on error. + int QueryLocalEce(ibv_ece* ece) const; + void SetOutgoingEce(const ibv_ece& ece); + // For debug void DebugInfo(std::ostream& os, butil::StringPiece connector = "\n") const; - // Callback when there is new epollin event on TCP fd. - static void OnNewDataFromTcp(Socket* m); - - // Real handshake for RDMA-mode sockets. - static ParseResult ExecuteServerHandshake(butil::IOBuf* source, Socket* socket); - // Initialize polling mode static int PollingModeInitialize(bthread_tag_t tag, std::function callback, @@ -145,33 +112,8 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); static void PollingModeRelease(bthread_tag_t tag); private: - enum State { - UNINIT = 0x0, - C_ALLOC_QPCQ = 0x1, - C_HELLO_SEND = 0x2, - C_HELLO_WAIT = 0x3, - C_BRINGUP_QP = 0x4, - C_ACK_SEND = 0x5, - S_HELLO_WAIT = 0x11, - S_ALLOC_QPCQ = 0x12, - S_BRINGUP_QP = 0x13, - S_HELLO_SEND = 0x14, - S_ACK_WAIT = 0x15, - ESTABLISHED = 0x100, - FALLBACK_TCP = 0x200, - FAILED = 0x300 - }; - - // 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. + // resource attached, so the caller can safely continue without RDMA. // Return 0 if success, -1 if failed and errno set int AllocateResources(); @@ -179,31 +121,12 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // in the middle with resources partially allocated. // Return 0 if success, -1 if failed and errno set int DoAllocateResources(); + // Start consuming CQ events only after handshake parsing has finished. + int StartCqEvents(); // 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 @@ -241,37 +164,18 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // -1: failed, errno set int DoPostRecv(void* block, size_t block_size); - // Read at most len bytes from fd in _socket to data - // wait for _read_butex if encounter EAGAIN - // return -1 if encounter other errno (including EOF) - int ReadFromFd(void* data, size_t len); - int ReadFromFd(butil::IOPortal* data, size_t len); - - - // Write at most len bytes from data to fd in _socket - // wait for _epollout_butex if encounter EAGAIN - // return -1 if encounter other errno - int WriteToFd(void* data, size_t len); - - // Write data to fd in _socket. - // wait for _epollout_butex if encounter EAGAIN. - // return -1 if encounter other errno. - int WriteToFd(butil::IOBuf* data); - - // Copy negotiated remote parameters into the endpoint and compute - // the SQ/RQ window capacities. Called by both - // ProcessHandshakeAtClient and ProcessHandshakeAtServer after the - // peer's hello has been validated. - void ApplyRemoteHello(const ParsedHello& remote); + // Copy negotiated remote parameters into the endpoint and compute the + // SQ/RQ window capacities. + void ApplyRemoteInfo(const RdmaConnectionInfo& remote); // Bringup the QP from RESET state to RTS state. // Arguments: - // remote: parsed remote hello. Provides the remote LID/GID/QP + // remote: negotiated peer parameters. Provides the remote LID/GID/QP // number for the RTR transition, and (on v3) the peer's // ECE to set during the INIT->RTR transition. // is_server: true on the server side, false on the client side. // Returns 0 on success, -1 on failed and errno set. - int BringUpQp(const ParsedHello& remote, bool is_server); + int BringUpQp(const RdmaConnectionInfo& remote, bool is_server); // Get event from comp channel and ack the events int GetAndAckEvents(SocketUniquePtr& s); @@ -282,9 +186,6 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Poll CQ and get the work completion static void PollCq(Socket* m); - // Get the description of current handshake state - std::string GetStateStr() const; - // Add cq socket id to poller void PollerAddCqSid(); @@ -293,25 +194,10 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Not owner Socket* _socket; + // Input state dedicated to the stream carried by the RDMA QP. + InputMessengerProcessor _input_processor; - // State of Handshake. FALLBACK_TCP publishes RdmaTransport::_rdma_state - // 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; - - // Wire-level handshake protocol version (set by dispatch in - // ProcessHandshakeAtClient/Server). Aligned with the protocol code: - // 0 = unnegotiated - // 2 = v2 "RDMA" - // 3 = v3 "RDM3" - int _handshake_version; - - // ECE payload to advertise in the next local hello: - // Client: the locally queried ECE capabilities (filled - // before C_HELLO_SEND); - // Server: the reduced/negotiated ECE queried after the - // QP reached RTS (filled in BringUpQp). + // ECE payload prepared by resource setup and consumed by the RDMA adapter. butil::optional _outgoing_ece; // rdma resource @@ -328,9 +214,6 @@ 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; @@ -366,9 +249,6 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // The number of new WRs posted in the local Recv Queue butil::atomic _new_rq_wrs; - // butex for inform read events on TCP fd during handshake - butil::atomic *_read_butex; - DISALLOW_COPY_AND_ASSIGN(RdmaEndpoint); // Cq socket id operation type diff --git a/src/brpc/rdma/rdma_handshake.cpp b/src/brpc/rdma/rdma_handshake.cpp deleted file mode 100644 index 17c6715562..0000000000 --- a/src/brpc/rdma/rdma_handshake.cpp +++ /dev/null @@ -1,514 +0,0 @@ -// 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. - -#if BRPC_WITH_RDMA - -#include "brpc/rdma/rdma_handshake.h" -#include "brpc/rdma/rdma_handshake_constants.h" - -#include -#include // std::min -#include -#include -#include -#include "butil/iobuf.h" // IOBuf, IOPortal, IOBufAsZeroCopy*Stream -#include "butil/sys_byteorder.h" -#include "butil/raw_pack.h" // RawPacker, RawUnpacker -#include "brpc/socket.h" -#include "brpc/rdma/rdma_endpoint.h" -#include "brpc/rdma/rdma_helper.h" -#include "brpc/rdma_transport.h" -#include "brpc/rdma/rdma_handshake.pb.h" -#include "butil/object_pool.h" - -namespace brpc { -namespace rdma { - -DEFINE_int32(rdma_client_handshake_version, 2, - "RDMA handshake protocol version used by client. " - "2 = legacy 'RDMA' magic (default, compatible with all servers); " - "3 = new 'RDM3' protobuf-based handshake " - "(MUST only be enabled after target servers support v3)."); -DECLARE_bool(rdma_trace_verbose); - -extern const uint16_t MIN_QP_SIZE; -extern const uint16_t MIN_BLOCK_SIZE; -extern uint32_t g_rdma_recv_block_size; -extern bool g_skip_rdma_init; - -extern int (*IbvQueryEce)(ibv_qp*, ibv_ece*); - -DEFINE_bool(rdma_ece, false, "Enable end-to-end ECE (Enhanced Connection Establishment) " - "negotiation in the RDMA v3 handshake. Automatically degrades " - "to no-ECE when the peer, the local libibverbs, or set_ece " - "does not support it. Acts as a kill switch (default off)."); - -DECLARE_bool(rdma_trace_verbose); - -namespace v2_wire { - -void HelloMessage::Serialize(void* data) const { - butil::RawPacker(data) - .pack16(msg_len) - .pack16(hello_ver) - .pack16(impl_ver) - .pack32(block_size) - .pack16(sq_size) - .pack16(rq_size) - .pack16(lid) - // gid is a raw 16-byte identifier and must NOT be byte-swapped. - .pack_bytes(gid.raw, sizeof(gid.raw)) - .pack32(qp_num); -} - -void HelloMessage::Deserialize(void* data) { - butil::RawUnpacker(data) - .unpack16(msg_len) - .unpack16(hello_ver) - .unpack16(impl_ver) - .unpack32(block_size) - .unpack16(sq_size) - .unpack16(rq_size) - .unpack16(lid) - // gid is a raw 16-byte identifier and must NOT be byte-swapped. - .unpack_bytes(gid.raw, sizeof(gid.raw)) - .unpack32(qp_num); -} - -static bool ValidHelloMessage(const HelloMessage& msg) { - return msg.hello_ver == HELLO_V2_VERSION && - msg.impl_ver == IMPL_V2_VERSION && - msg.block_size >= MIN_BLOCK_SIZE && - msg.sq_size >= MIN_QP_SIZE && - msg.rq_size >= MIN_QP_SIZE; -} - -static void TranslateV2Hello(const HelloMessage& msg, ParsedHello* out) { - out->block_size = msg.block_size; - out->sq_size = msg.sq_size; - out->rq_size = msg.rq_size; - out->lid = msg.lid; - out->gid = msg.gid; - out->qp_num = msg.qp_num; -} - -RemoteHelloResult ReadBodyAndNegotiate(RdmaEndpoint* ep, ParsedHello* remote) { - uint8_t data[HELLO_V2_MSG_LEN_MIN]; - if (ep->ReadFromFd(data, HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN) < 0) { - return RemoteHelloResult::ERROR; - } - HelloMessage remote_msg{}; - remote_msg.Deserialize(data); - if (remote_msg.msg_len < HELLO_V2_MSG_LEN_MIN || - remote_msg.msg_len > HELLO_V2_MSG_LEN_MAX) { - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - if (remote_msg.msg_len > HELLO_V2_MSG_LEN_MIN) { - // Drain unknown trailing bytes so they don't pollute subsequent - // reads (e.g. the upcoming ACK message). v2 base fields already - // carry enough information for negotiation; unknown trailing - // bytes are treated as optional hints that v2 safely ignores. - size_t ext_len = remote_msg.msg_len - HELLO_V2_MSG_LEN_MIN; - if (DrainBytes(ep, ext_len) < 0) { - return RemoteHelloResult::ERROR; - } - } - if (!ValidHelloMessage(remote_msg)) { - return RemoteHelloResult::FALLBACK; - } - TranslateV2Hello(remote_msg, remote); - return RemoteHelloResult::NEGOTIATED; -} - -int DrainBytes(RdmaEndpoint* ep, size_t n) { - uint8_t scratch[64]; - while (n > 0) { - size_t chunk = std::min(n, sizeof(scratch)); - if (ep->ReadFromFd(scratch, chunk) < 0) { - return -1; - } - n -= chunk; - } - return 0; -} - -} // namespace v2_wire - -int RdmaHandshakeClientV2::SendLocalHello() { - RdmaEndpoint* ep = _ep; - uint8_t data[HELLO_V2_MSG_LEN_MIN]; - - v2_wire::HelloMessage local_msg{}; - local_msg.msg_len = HELLO_V2_MSG_LEN_MIN; - local_msg.hello_ver = HELLO_V2_VERSION; - local_msg.impl_ver = IMPL_V2_VERSION; - local_msg.block_size = g_rdma_recv_block_size; - local_msg.sq_size = ep->_sq_size; - local_msg.rq_size = ep->_rq_size; - local_msg.lid = GetRdmaLid(); - local_msg.gid = GetRdmaGid(); - if (BAIDU_LIKELY(ep->_resource)) { - local_msg.qp_num = ep->_resource->qp->qp_num; - } else { - // Only happens in UT - local_msg.qp_num = 0; - } - fast_memcpy(data, HELLO_MAGIC, 4); - local_msg.Serialize((char*)data + 4); - return ep->WriteToFd(data, HELLO_V2_MSG_LEN_MIN); -} - -RemoteHelloResult RdmaHandshakeClientV2::ReceiveAndParseRemoteHello(ParsedHello* remote) { - uint8_t magic[HELLO_MAGIC_LEN]; - if (_ep->ReadFromFd(magic, HELLO_MAGIC_LEN) < 0) { - return RemoteHelloResult::ERROR; - } - if (memcmp(magic, HELLO_MAGIC, HELLO_MAGIC_LEN) != 0) { - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - - return v2_wire::ReadBodyAndNegotiate(_ep, remote); -} - -// Parse one complete v2 client hello out of `_source` (non-blocking). -// v2 hello: [ "RDMA" 4B ][ msg_len 2B ][ 34B ... ], base total = 40B. -RemoteHelloResult RdmaHandshakeServerV2::ReceiveAndParseRemoteHello(ParsedHello* remote) { - butil::IOBuf* source = _source; - constexpr size_t HDR_LEN = HELLO_MAGIC_LEN + 2; - if (source->size() < HDR_LEN) { - // msg_len has not fully arrived yet. - return RemoteHelloResult::NEED_MORE; - } - - uint8_t hdr[HDR_LEN]; - CHECK_EQ(source->copy_to(hdr, sizeof(hdr)), sizeof(hdr)); - - uint16_t msg_len = 0; - butil::RawUnpacker(hdr + HELLO_MAGIC_LEN).unpack16(msg_len); - if (msg_len < HELLO_V2_MSG_LEN_MIN || msg_len > HELLO_V2_MSG_LEN_MAX) { - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - if (source->size() < msg_len) { - // Full message has not fully arrived yet. - return RemoteHelloResult::NEED_MORE; - } - - // Consume the whole hello: magic + 36B base body + optional extension. - CHECK_EQ(source->pop_front(HELLO_MAGIC_LEN), HELLO_MAGIC_LEN); - uint8_t body[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]; // 36B - CHECK_EQ(source->cutn(body, sizeof(body)), sizeof(body)); - if (!source->empty()) { - // Drain unknown trailing bytes. - source->clear(); - } - - v2_wire::HelloMessage remote_msg{}; - remote_msg.Deserialize(body); - if (!v2_wire::ValidHelloMessage(remote_msg)) { - return RemoteHelloResult::FALLBACK; - } - v2_wire::TranslateV2Hello(remote_msg, remote); - return RemoteHelloResult::NEGOTIATED; -} - -int RdmaHandshakeServerV2::SendLocalHello() { - uint8_t data[HELLO_V2_MSG_LEN_MIN]; - v2_wire::HelloMessage local_msg{}; - local_msg.msg_len = HELLO_V2_MSG_LEN_MIN; - auto rdma_transport = static_cast(_ep->_socket->_transport.get()); - if (rdma_transport->_rdma_state == RdmaTransport::RDMA_OFF) { - local_msg.hello_ver = 0; - local_msg.impl_ver = 0; - local_msg.block_size = 0; - local_msg.sq_size = 0; - local_msg.rq_size = 0; - local_msg.lid = 0; - memset(local_msg.gid.raw, 0, sizeof(local_msg.gid.raw)); - local_msg.qp_num = 0; - } else { - local_msg.hello_ver = HELLO_V2_VERSION; - local_msg.impl_ver = IMPL_V2_VERSION; - local_msg.block_size = g_rdma_recv_block_size; - local_msg.sq_size = _ep->_sq_size; - local_msg.rq_size = _ep->_rq_size; - local_msg.lid = GetRdmaLid(); - local_msg.gid = GetRdmaGid(); - if (BAIDU_LIKELY(_ep->_resource)) { - local_msg.qp_num = _ep->_resource->qp->qp_num; - } else { - // Only happens in UT - local_msg.qp_num = 0; - } - } - fast_memcpy(data, HELLO_MAGIC, 4); - local_msg.Serialize((char*)data + 4); - return _ep->WriteToFd(data, HELLO_V2_MSG_LEN_MIN); -} - -namespace v3_wire { - -bool ValidRdmaHello(const RdmaHello& msg) { - if (msg.gid().size() != sizeof(ibv_gid)) { - return false; - } - // ParsedHello stores these as uint16_t; reject values that would truncate. - constexpr uint16_t MAX_UINT16 = std::numeric_limits::max(); - if (msg.sq_size() > MAX_UINT16 || msg.rq_size() > MAX_UINT16 || msg.lid() > MAX_UINT16) { - return false; - } - if (msg.block_size() < MIN_BLOCK_SIZE) { - return false; - } - if (msg.sq_size() < MIN_QP_SIZE) { - return false; - } - if (msg.rq_size() < MIN_QP_SIZE) { - return false; - } - // qp_num == 0 only happens in UT (no real QP allocated). - if (msg.qp_num() == 0 && !g_skip_rdma_init) { - return false; - } - return true; -} - -void FillLocalRdmaHello(const RdmaEndpoint* ep, RdmaHello* msg) { - msg->set_block_size(g_rdma_recv_block_size); - msg->set_sq_size(ep->_sq_size); - msg->set_rq_size(ep->_rq_size); - msg->set_lid(GetRdmaLid()); - ibv_gid gid = GetRdmaGid(); - msg->set_gid(reinterpret_cast(gid.raw), sizeof(gid.raw)); - if (BAIDU_LIKELY(ep->_resource)) { - msg->set_qp_num(ep->_resource->qp->qp_num); - } else { - // Only happens in UT - msg->set_qp_num(0); - } - - // Advertise ECE only when enabled. Role-dependent payload: - // Client hello: the locally queried ECE capabilities; - // Server hello: the reduced/negotiated ECE queried after RTS. - // When the relevant ECE is not valid (disabled, unsupported, or query - // failed) the field is simply omitted and the peer degrades to no-ECE. - // Advertise ECE if there is anything to advertise. The endpoint pre-fills - // _outgoing_ece in a role-specific way: client side stores its locally - // queried capabilities; server side stores the reduced/negotiated ECE - // after RTS. nullopt -> omit the field (peer degrades to no-ECE). - if (FLAGS_rdma_ece && ep->_outgoing_ece.has_value()) { - RdmaEce* ece = msg->mutable_ece(); - ece->set_vendor_id(ep->_outgoing_ece->vendor_id); - ece->set_options(ep->_outgoing_ece->options); - ece->set_comp_mask(ep->_outgoing_ece->comp_mask); - } -} - -int ReadAndParseV3Hello(RdmaEndpoint* ep, RdmaHello* out) { - uint8_t size_buf[HELLO_V3_PB_SIZE_LEN]; - if (ep->ReadFromFd(size_buf, HELLO_V3_PB_SIZE_LEN) < 0) { - return -1; - } - uint32_t pb_size = butil::NetToHost32( - *reinterpret_cast(size_buf)); - if (pb_size == 0 || pb_size > HELLO_V3_MAX_PB_SIZE) { - errno = EPROTO; - return -1; - } - butil::IOPortal body; - if (ep->ReadFromFd(&body, pb_size) < 0) { - return -1; - } - - butil::IOBufAsZeroCopyInputStream input(body); - if (!out->ParseFromZeroCopyStream(&input)) { - LOG(ERROR) << "Failed to parse RdmaHello"; - errno = EPROTO; - return -1; - } - return 0; -} - -int WriteV3Hello(RdmaEndpoint* ep, const RdmaHello& msg) { - uint32_t pb_size = static_cast(msg.ByteSizeLong()); - if (pb_size > HELLO_V3_MAX_PB_SIZE) { - errno = EPROTO; - return -1; - } - - // [ "RDM3" 4B ][ pb_size 4B (big-endian) ][ RdmaHello protobuf bytes ] - butil::IOBuf packet; - packet.append(HELLO_MAGIC_V3, HELLO_MAGIC_LEN); - uint32_t pb_size_be = butil::HostToNet32(pb_size); - packet.append(&pb_size_be, HELLO_V3_PB_SIZE_LEN); - butil::IOBufAsZeroCopyOutputStream output(&packet); - if (!msg.SerializeToZeroCopyStream(&output)) { - LOG(ERROR) << "Failed to serialize RdmaHello"; - errno = EPROTO; - return -1; - } - return ep->WriteToFd(&packet); -} - -void TranslateHello(const RdmaHello& msg, ParsedHello* out) { - out->block_size = msg.block_size(); - out->sq_size = static_cast(msg.sq_size()); - out->rq_size = static_cast(msg.rq_size()); - out->lid = static_cast(msg.lid()); - fast_memcpy(out->gid.raw, msg.gid().data(), sizeof(out->gid.raw)); - out->qp_num = msg.qp_num(); - if (FLAGS_rdma_ece && msg.has_ece()) { - ibv_ece ece; - ece.vendor_id = msg.ece().vendor_id(); - ece.options = msg.ece().options(); - ece.comp_mask = msg.ece().comp_mask(); - out->ece = ece; - } -} - -} // namespace v3_wire - -int RdmaHandshakeClientV3::SendLocalHello() { - // Query local ECE capabilities so they can be advertised in the client - // hello. v3-only. Best-effort: any failure or missing API just means we - // won't advertise ECE (the peer then degrades to no-ECE establishment). - if (FLAGS_rdma_ece && IbvQueryEce != nullptr && - _ep->_resource && _ep->_resource->qp) { - ibv_ece ece; - if (IbvQueryEce(_ep->_resource->qp, &ece) == 0) { - _ep->_outgoing_ece = ece; - } else { - LOG_IF(WARNING, FLAGS_rdma_trace_verbose) - << "Fail to IbvQueryEce on client, ECE not advertised: " - << _ep->_socket->description(); - } - } - - RdmaHello local_msg{}; - v3_wire::FillLocalRdmaHello(_ep, &local_msg); - return v3_wire::WriteV3Hello(_ep, local_msg); -} - -RemoteHelloResult RdmaHandshakeClientV3::ReceiveAndParseRemoteHello(ParsedHello* remote) { - uint8_t magic[HELLO_MAGIC_LEN]; - if (_ep->ReadFromFd(magic, HELLO_MAGIC_LEN) < 0) { - return RemoteHelloResult::ERROR; - } - if (memcmp(magic, HELLO_MAGIC_V3, HELLO_MAGIC_LEN) != 0) { - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - - RdmaHello remote_msg{}; - if (v3_wire::ReadAndParseV3Hello(_ep, &remote_msg) < 0) { - return RemoteHelloResult::ERROR; - } - if (!v3_wire::ValidRdmaHello(remote_msg)) { - return RemoteHelloResult::FALLBACK; - } - v3_wire::TranslateHello(remote_msg, remote); - return RemoteHelloResult::NEGOTIATED; -} - -// Parse one complete v3 client hello out of `_source` (non-blocking). -// v3 hello: [ "RDM3" 4B ][ pb_size 4B (big-endian) ][ RdmaHello ] -RemoteHelloResult RdmaHandshakeServerV3::ReceiveAndParseRemoteHello(ParsedHello* remote) { - constexpr size_t HDR_LEN = HELLO_MAGIC_LEN + HELLO_V3_PB_SIZE_LEN; - if (_source->size() < HDR_LEN) { - // pb_size has not fully arrived yet. - return RemoteHelloResult::NEED_MORE; - } - - uint8_t hdr[HDR_LEN]; - CHECK_EQ(_source->copy_to(hdr, sizeof(hdr)), sizeof(hdr)); - - uint32_t pb_size = butil::NetToHost32( - *reinterpret_cast(hdr + HELLO_MAGIC_LEN)); - if (pb_size == 0 || pb_size > HELLO_V3_MAX_PB_SIZE) { - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - size_t total = HDR_LEN + pb_size; - if (_source->size() < total) { - // Full message has not fully arrived yet. - return RemoteHelloResult::NEED_MORE; - } - - CHECK_EQ(_source->cutn(hdr, HDR_LEN), HDR_LEN); - butil::IOBuf pb; - CHECK_EQ(_source->cutn(&pb, pb_size), pb_size); - RdmaHello remote_msg; - butil::IOBufAsZeroCopyInputStream input(pb); - if (!remote_msg.ParseFromZeroCopyStream(&input)) { - LOG(ERROR) << "Failed to parse RdmaHello"; - errno = EPROTO; - return RemoteHelloResult::ERROR; - } - if (!v3_wire::ValidRdmaHello(remote_msg)) { - return RemoteHelloResult::FALLBACK; - } - v3_wire::TranslateHello(remote_msg, remote); - return RemoteHelloResult::NEGOTIATED; -} - -int RdmaHandshakeServerV3::SendLocalHello() { - RdmaHello local_msg{}; - auto rdma_transport = static_cast(_ep->_socket->_transport.get()); - if (rdma_transport->_rdma_state == RdmaTransport::RDMA_OFF) { - // Un-negotiable hello: all body fields are zero so the client's - // rejects it and downgrades to TCP on the same connection. - local_msg.set_block_size(0); - local_msg.set_sq_size(0); - local_msg.set_rq_size(0); - local_msg.set_lid(0); - local_msg.set_gid(std::string(sizeof(ibv_gid), '\0')); - local_msg.set_qp_num(0); - } else { - v3_wire::FillLocalRdmaHello(_ep, &local_msg); - } - return v3_wire::WriteV3Hello(_ep, local_msg); -} - -std::unique_ptr CreateClientHandshake(RdmaEndpoint* ep) { - switch (FLAGS_rdma_client_handshake_version) { - case 3: - return std::unique_ptr(new RdmaHandshakeClientV3(ep)); - case 2: - default: - return std::unique_ptr(new RdmaHandshakeClientV2(ep)); - } -} - -std::unique_ptr CreateServerHandshakeByMagic( - RdmaEndpoint* ep, butil::IOBuf* source, const uint8_t magic[HELLO_MAGIC_LEN]) { - if (memcmp(magic, HELLO_MAGIC, HELLO_MAGIC_LEN) == 0) { - return std::unique_ptr( - new RdmaHandshakeServerV2(ep, source)); - } - if (memcmp(magic, HELLO_MAGIC_V3, HELLO_MAGIC_LEN) == 0) { - return std::unique_ptr( - new RdmaHandshakeServerV3(ep, source)); - } - return nullptr; -} - -} // namespace rdma -} // namespace brpc - -#endif // BRPC_WITH_RDMA diff --git a/src/brpc/rdma/rdma_handshake.h b/src/brpc/rdma/rdma_handshake.h deleted file mode 100644 index 2d10220ab9..0000000000 --- a/src/brpc/rdma/rdma_handshake.h +++ /dev/null @@ -1,185 +0,0 @@ -// 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_RDMA_HANDSHAKE_H -#define BRPC_RDMA_HANDSHAKE_H - -#if BRPC_WITH_RDMA - -#include -#include -#include "butil/macros.h" -#include "butil/containers/optional.h" -#include "brpc/rdma/rdma_handshake_constants.h" - -namespace butil { -class IOBuf; -} - -namespace brpc { -namespace rdma { - -class RdmaEndpoint; - -// Wire-format-agnostic representation of a peer's hello message. -// Each protocol version (v2 binary, v3 protobuf) translates its own -// wire format into this struct so the state-machine driver in -// RdmaEndpoint::ProcessHandshakeAt{Client,Server} stays free of any -// wire-format details. -struct ParsedHello { - uint32_t block_size; - uint16_t sq_size; - uint16_t rq_size; - uint16_t lid; - ibv_gid gid; - uint32_t qp_num; - // ECE (Enhanced Connection Establishment), v3 handshake only. - // nullopt means the peer did not advertise any ECE (either it disabled - // ECE, its lib does not support ECE, or it is a v2 peer). When engaged: - // - on the server side: the client's queried ECE capabilities; - // - on the client side: the server's reduced/negotiated ECE. - butil::optional ece; -}; - -// Result of reading/parsing a peer's hello (see ReceiveAndParseRemoteHello). -enum class RemoteHelloResult { - // A full hello was read and negotiation succeeded. - NEGOTIATED, - // A full hello was read but negotiation failed. - FALLBACK, - // (server only) Not enough data yet. - NEED_MORE, - // IO/protocol error (errno set). - ERROR, -}; - -namespace v2_wire { - -// v2 binary HelloMessage. -struct HelloMessage { - void Serialize(void* data) const; - void Deserialize(void* data); - - uint16_t msg_len; - uint16_t hello_ver; - uint16_t impl_ver; - uint32_t block_size; - uint16_t sq_size; - uint16_t rq_size; - uint16_t lid; - ibv_gid gid; - uint32_t qp_num; -}; - -} // namespace v2_wire - -// Base class of an RDMA handshake, shared by both roles. -class RdmaHandshake { -public: - RdmaHandshake(RdmaEndpoint* ep, int version) : _ep(ep), _version(version) {} - virtual ~RdmaHandshake() = default; - - DISALLOW_COPY_AND_ASSIGN(RdmaHandshake); - - // Wire-level protocol version (2 for "RDMA", 3 for "RDM3"). - int ProtocolVersion() const { return _version; } - - // Build and send the local hello (including the protocol magic). - // Returns 0 on success, -1 on IO error (errno set). - // - // For a server in fallback state, implementations MUST still - // produce a sendable message; each version uses its own wire - // convention to signal "I am falling back" to the peer: - // - v2: zero hello_ver/impl_ver so the peer's HelloNegotiationValid - // rejects it; - // - v3: qp_num==0 so the peer's ValidRdmaHello rejects it. - virtual int SendLocalHello() = 0; - - // Read and parse the peer's hello into *remote. - virtual RemoteHelloResult ReceiveAndParseRemoteHello(ParsedHello* remote) = 0; - -protected: - RdmaEndpoint* _ep; - int _version; -}; - -// Server-side handshake base: parses the remote hello non-blockingly out of -// `_source` (an IOBuf filled by InputMessenger), never touching the fd. -class ServerRdmaHandshake : public RdmaHandshake { -public: - ServerRdmaHandshake(RdmaEndpoint* ep, butil::IOBuf* source, int version) - : RdmaHandshake(ep, version), _source(source) {} - -protected: - butil::IOBuf* _source; -}; - -// v2 handshake (legacy "RDMA" magic, 36B binary HelloMessage). -class RdmaHandshakeClientV2 : public RdmaHandshake { -public: - explicit RdmaHandshakeClientV2(RdmaEndpoint* ep) : RdmaHandshake(ep, 2) {} - int SendLocalHello() override; - RemoteHelloResult ReceiveAndParseRemoteHello(ParsedHello* remote) override; -}; - -class RdmaHandshakeServerV2 : public ServerRdmaHandshake { -public: - RdmaHandshakeServerV2(RdmaEndpoint* ep, butil::IOBuf* source) - : ServerRdmaHandshake(ep, source, 2) {} - int SendLocalHello() override; - RemoteHelloResult ReceiveAndParseRemoteHello(ParsedHello* remote) override; -}; - -// v3 handshake (new "RDM3" magic, protobuf RdmaHello). -// [ "RDM3" 4B ][ pb_size 4B (big-endian) ][ RdmaHello protobuf bytes ] -class RdmaHandshakeClientV3 : public RdmaHandshake { -public: - explicit RdmaHandshakeClientV3(RdmaEndpoint* ep) : RdmaHandshake(ep, 3) {} - int SendLocalHello() override; - RemoteHelloResult ReceiveAndParseRemoteHello(ParsedHello* remote) override; -}; - -class RdmaHandshakeServerV3 : public ServerRdmaHandshake { -public: - RdmaHandshakeServerV3(RdmaEndpoint* ep, butil::IOBuf* source) - : ServerRdmaHandshake(ep, source, 3) {} - int SendLocalHello() override; - RemoteHelloResult ReceiveAndParseRemoteHello(ParsedHello* remote) override; -}; - -// Factory methods -// -// Pick the client-side handshake based on -// FLAGS_rdma_client_handshake_version: -// 2 (default) -> RdmaHandshakeClientV2 -// 3 -> RdmaHandshakeClientV3 -// Other values fall back to V2. -std::unique_ptr CreateClientHandshake(RdmaEndpoint* ep); - -// Pick the server-side handshake based on the 4B magic already read. -// Returns nullptr if `magic` is not a recognized RDMA magic -// (the caller should then fallback to TCP). -// "RDMA" -> RdmaHandshakeServerV2 -// "RDM3" -> RdmaHandshakeServerV3 -std::unique_ptr CreateServerHandshakeByMagic( - RdmaEndpoint* ep, butil::IOBuf* source, const uint8_t magic[HELLO_MAGIC_LEN]); - -} // namespace rdma -} // namespace brpc - -#endif // BRPC_WITH_RDMA -#endif // BRPC_RDMA_HANDSHAKE_H diff --git a/src/brpc/rdma/rdma_handshake_server.cpp b/src/brpc/rdma/rdma_handshake_server.cpp deleted file mode 100644 index 1b0ca4226a..0000000000 --- a/src/brpc/rdma/rdma_handshake_server.cpp +++ /dev/null @@ -1,219 +0,0 @@ -// 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/rdma/rdma_handshake_server.h" - -#include -#include -#include -#include "butil/iobuf.h" -#include "butil/logging.h" -#include "butil/object_pool.h" -#include "butil/raw_pack.h" -#include "butil/sys_byteorder.h" -#include "brpc/socket.h" -#include "brpc/rdma/rdma_handshake.pb.h" -#include "brpc/rdma/rdma_handshake_constants.h" -#if BRPC_WITH_RDMA -#include "brpc/rdma/rdma_endpoint.h" -#endif - -namespace brpc { -namespace rdma { - -ServerHandshakeContext* ServerHandshakeContext::Create() { - return butil::get_object(); -} - -void ServerHandshakeContext::Destroy() { - butil::return_object(this); -} - -// Fallback-only server handshake. Used for any connection that is NOT in RDMA -// mode: builds without RDMA (no RdmaEndpoint exists at all), and RDMA-enabled -// builds where this particular connection is plain TCP. Every RDMA client that -// reaches here is answered with an un-negotiable hello and asked to downgrade -// to TCP on the same connection, then its ACK is drained. - -// An intentionally-invalid v2 hello_ver: the client's ValidHelloMessage() -// requires hello_ver==2, so it rejects this and falls back. -static constexpr uint16_t V2_HELLO_VERSION_INVALID = std::numeric_limits::max(); -// Length of the gid field (== sizeof(ibv_gid)), spelled as a literal so this -// compile-switch-independent file stays free of . -static constexpr size_t V3_GID_LEN = 16; - -// Consume one complete v2 client hello ("RDMA" + 2B msg_len + body) from -// `source` without interpreting its content (fallback does not negotiate). -// Returns 1 (consumed), 0 (not enough data yet, nothing consumed) or -1 (error). -static int DrainClientHelloV2(butil::IOBuf* source) { - constexpr size_t HDR_LEN = HELLO_MAGIC_LEN + 2; - if (source->size() < HDR_LEN) { - return 0; - } - - uint8_t hdr[HDR_LEN]; - CHECK_EQ(source->copy_to(hdr, sizeof(hdr)), sizeof(hdr)); - - uint16_t msg_len = 0; - butil::RawUnpacker(hdr + HELLO_MAGIC_LEN).unpack16(msg_len); - if (msg_len < HELLO_V2_MSG_LEN_MIN || msg_len > HELLO_V2_MSG_LEN_MAX) { - return -1; - } - - if (source->size() < msg_len) { - return 0; - } - - CHECK_EQ(source->pop_front(msg_len), msg_len); - return 1; -} - -// Consume one complete v3 client hello ("RDM3" + 4B pb_size + protobuf body) -// from `source` without interpreting its content. -// Returns 1 (consumed), 0 (not enough data yet, nothing consumed) or -1 (error). -static int DrainClientHelloV3(butil::IOBuf* source) { - constexpr size_t HDR_LEN = HELLO_MAGIC_LEN + HELLO_V3_PB_SIZE_LEN; - if (source->size()< HDR_LEN) { - return 0; - } - - uint8_t hdr[HDR_LEN]; - CHECK_EQ(source->copy_to(hdr, sizeof(hdr)), sizeof(hdr)); - - uint32_t pb_size = butil::NetToHost32( - *reinterpret_cast(hdr + HELLO_MAGIC_LEN)); - if (pb_size == 0 || pb_size > HELLO_V3_MAX_PB_SIZE) { - return -1; - } - - const size_t total = HDR_LEN + pb_size; - if (source->size() < total) { - return 0; - } - - CHECK_EQ(source->pop_front(total), total); - return 1; -} - -// Reply an un-negotiable hello so the client downgrades to TCP. All body fields -// are zero/invalid so the client's validity check rejects it. -// Returns 0 on success, -1 otherwise. -static int SendUnnegotiableHello(Socket* socket, int version) { - butil::IOBuf packet; - if (version == 2) { - // magic "RDMA" + msg_len(=40) + invalid hello_ver; rest stays zero. - // NOTE: msg_len is the length of the WHOLE hello INCLUDING the 4B magic - // (== HELLO_V2_MSG_LEN_MIN), not just the body; the client rejects any - // msg_len < HELLO_V2_MSG_LEN_MIN as a protocol error. - packet.append(HELLO_MAGIC, HELLO_MAGIC_LEN); - char reply[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]{}; - butil::RawPacker(reply).pack16(HELLO_V2_MSG_LEN_MIN) - .pack16(V2_HELLO_VERSION_INVALID); - packet.append(reply, sizeof(reply)); - } else { - // "RDM3" + pb_size + RdmaHello with block_size==0 & qp_num==0 so the - // client's ValidRdmaHello() returns false. - RdmaHello reply_msg; - reply_msg.set_block_size(0); - reply_msg.set_sq_size(0); - reply_msg.set_rq_size(0); - reply_msg.set_lid(0); - reply_msg.set_gid(std::string(V3_GID_LEN, '\0')); - reply_msg.set_qp_num(0); - packet.append(HELLO_MAGIC_V3, HELLO_MAGIC_LEN); - uint32_t pb_size_be = - butil::HostToNet32(static_cast(reply_msg.ByteSizeLong())); - packet.append(&pb_size_be, sizeof(pb_size_be)); - butil::IOBufAsZeroCopyOutputStream output(&packet); - if (!reply_msg.SerializeToZeroCopyStream(&output)) { - LOG(WARNING) << "Fail to serialize RDMA v3 fallback hello"; - return -1; - } - } - - if (socket->Write(&packet) != 0) { - PLOG(WARNING) << "Fail to send RDMA fallback hello to " << socket->description(); - return -1; - } - return 0; -} - -// 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) { - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - // Phase 1: consume the client hello and reply an un-negotiable hello. - uint8_t magic[HELLO_MAGIC_LEN]; - CHECK_EQ(source->copy_to(magic, HELLO_MAGIC_LEN), HELLO_MAGIC_LEN); - - int version; - if (memcmp(magic, HELLO_MAGIC, HELLO_MAGIC_LEN) == 0) { - version = 2; - } else if (memcmp(magic, HELLO_MAGIC_V3, HELLO_MAGIC_LEN) == 0) { - version = 3; - } else { - return MakeParseError(PARSE_ERROR_TRY_OTHERS); - } - - const int r = version == 2 ? DrainClientHelloV2(source) : DrainClientHelloV3(source); - if (r == 0) { - // Hello not complete yet; keep the buffer intact and retry later. - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - if (r < 0) { - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } - if (SendUnnegotiableHello(socket, version) < 0) { - return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG); - } - // Wait for the client ACK across subsequent reads. - socket->reset_parsing_context(ServerHandshakeContext::Create()); - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - - // Phase 2: drain the 4B ACK. - if (source->size() < HELLO_ACK_LEN) { - return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); - } - CHECK_EQ(source->pop_front(HELLO_ACK_LEN), HELLO_ACK_LEN); - // Handshake done. - // Drop the context and let InputMessenger parse the following real RPC. - socket->reset_parsing_context(nullptr); - return MakeParseError(PARSE_ERROR_TRY_OTHERS); -} - -ParseResult ExecuteServerHandshake(butil::IOBuf* source, Socket* socket) { - // Only RDMA-mode connections carry a live RdmaEndpoint and run the real - // handshake. A connection that is not in RDMA mode (RDMA compiled in but - // this connection is plain TCP, or RDMA not compiled at all) falls back. -#if BRPC_WITH_RDMA - if (socket->socket_mode() == SOCKET_MODE_RDMA) { - return RdmaEndpoint::ExecuteServerHandshake(source, socket); - } -#endif - return FallbackServerHandshake(source, socket); -} - -} // namespace rdma -} // namespace brpc diff --git a/src/brpc/rdma/rdma_handshake_server.h b/src/brpc/rdma/rdma_handshake_server.h deleted file mode 100644 index 705aa893ef..0000000000 --- a/src/brpc/rdma/rdma_handshake_server.h +++ /dev/null @@ -1,48 +0,0 @@ -// 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_RDMA_RDMA_HANDSHAKE_SERVER_H -#define BRPC_RDMA_RDMA_HANDSHAKE_SERVER_H - -#include "brpc/destroyable.h" -#include "brpc/parse_result.h" - -namespace butil { -class IOBuf; -} - -namespace brpc { -class Socket; -namespace rdma { - -// State kept across multiple parse calls of the server handshake. -struct ServerHandshakeContext : public Destroyable { - static ServerHandshakeContext* Create(); - void Destroy() override; -}; - -// The single server-side RDMA handshake entry for policy::ParseRdmaHandshake. -// Returns a ParseResult ready to be handed back from the protocol parser: -// - not an RDMA magic / handshake finished -> PARSE_ERROR_TRY_OTHERS; -// - an RDMA magic but not enough bytes yet -> PARSE_ERROR_NOT_ENOUGH_DATA; -// - IO/protocol error -> PARSE_ERROR_ABSOLUTELY_WRONG. -ParseResult ExecuteServerHandshake(butil::IOBuf* source, Socket* socket); - -} // namespace rdma -} // namespace brpc - -#endif // BRPC_RDMA_RDMA_HANDSHAKE_SERVER_H diff --git a/src/brpc/rdma/rdma_handshake.proto b/src/brpc/rdma_handshake.proto similarity index 100% rename from src/brpc/rdma/rdma_handshake.proto rename to src/brpc/rdma_handshake.proto diff --git a/src/brpc/rdma_transport.cpp b/src/brpc/rdma_transport.cpp index a2037aa725..f85c1af243 100644 --- a/src/brpc/rdma_transport.cpp +++ b/src/brpc/rdma_transport.cpp @@ -18,9 +18,8 @@ #if BRPC_WITH_RDMA #include "brpc/rdma_transport.h" +#include "brpc/adapter_transport.h" #include "brpc/event_dispatcher.h" -#include "brpc/tcp_transport.h" -#include "brpc/input_messenger.h" #include "brpc/rdma/rdma_endpoint.h" #include "brpc/rdma/rdma_helper.h" @@ -30,27 +29,26 @@ DECLARE_bool(usercode_in_pthread); extern SocketVarsCollector *g_vars; +RdmaTransport* RdmaTransport::Get(const Socket* socket) { + const AdapterTransport* adapter = AdapterTransport::Get(socket); + Transport* transport = adapter->high_speed_transport(); + CHECK(transport != nullptr); + return static_cast(transport); +} + void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { CHECK(_rdma_ep == nullptr); - if (options.socket_mode == SOCKET_MODE_RDMA) { - _rdma_ep = new rdma::RdmaEndpoint(socket); - _rdma_state = RDMA_UNKNOWN; - } else { - _rdma_state = RDMA_OFF; - socket->_socket_mode = SOCKET_MODE_TCP; - } _socket = socket; _default_connect = options.app_connect; - _on_edge_trigger = options.on_edge_triggered_events; - if (options.need_on_edge_trigger && _on_edge_trigger == nullptr) { - if (_rdma_ep != nullptr) { - _on_edge_trigger = rdma::RdmaEndpoint::OnNewDataFromTcp; - } else { - _on_edge_trigger = InputMessenger::OnNewMessages; - } + _on_edge_trigger = nullptr; + _rdma_state = RDMA_UNKNOWN; + _rdma_ep = new (std::nothrow) rdma::RdmaEndpoint(socket); + if (!_rdma_ep) { + const int saved_errno = errno != 0 ? errno : ENOMEM; + errno = saved_errno; + PLOG(WARNING) << "Fail to create RdmaEndpoint, disable RDMA upgrade"; + _rdma_state = RDMA_OFF; } - _tcp_transport = std::make_shared(); - _tcp_transport->Init(socket, options); } void RdmaTransport::Release() { @@ -70,31 +68,56 @@ int RdmaTransport::Reset(int32_t expected_nref) { } std::shared_ptr RdmaTransport::Connect() { - if (_default_connect == nullptr) { - return std::make_shared(); - } return _default_connect; } -int RdmaTransport::CutFromIOBuf(butil::IOBuf *buf) { - // Only send over the RDMA channel once the handshake has NEGOTIATED it - // (RDMA_ON). While the state is still RDMA_UNKNOWN (handshake in progress, - // or a server connection that turned out to be plain TCP and never - // handshook) or RDMA_OFF (fell back), the QP is not usable and everything - // must go over the TCP fd. Mirrors the RDMA_ON check in WaitEpollOut(). - if (_rdma_ep && _rdma_state == RDMA_ON) { - butil::IOBuf *data_arr[1] = {buf}; - return _rdma_ep->CutFromIOBufList(data_arr, 1); - } else { - return _tcp_transport->CutFromIOBuf(buf); +void RdmaTransport::SetHighSpeedAvailable(bool available) { + _rdma_state = available ? RDMA_ON : RDMA_OFF; +} + +int RdmaTransport::PrepareUpgradeResources() { + return _rdma_ep->AllocateResources(); +} + +int RdmaTransport::NegotiateUpgradeResources( + const rdma::RdmaConnectionInfo &remote, bool server) { + _rdma_ep->ApplyRemoteInfo(remote); + return _rdma_ep->BringUpQp(remote, server); +} + +int RdmaTransport::StartUpgradeEvents() { + return _rdma_ep->StartCqEvents(); +} + +std::unique_ptr +RdmaTransport::CreateClientHandshakeAdapter() { + return rdma::CreateClientHandshakeAdapter(_rdma_ep); +} + +std::vector> +RdmaTransport::CreateServerHandshakeAdapters() { + return rdma::CreateServerHandshakeAdapters(_rdma_ep); +} + +void RdmaTransport::ActivateUpgrade() { + SetHighSpeedAvailable(true); +} + +void RdmaTransport::DeactivateUpgrade() { + SetHighSpeedAvailable(false); + if (_rdma_ep != nullptr) { + _rdma_ep->Reset(); } } +int RdmaTransport::CutFromIOBuf(butil::IOBuf *buf) { + butil::IOBuf *data[1] = {buf}; + return static_cast(CutFromIOBufList(data, 1)); +} + ssize_t RdmaTransport::CutFromIOBufList(butil::IOBuf **buf, size_t ndata) { - if (_rdma_ep && _rdma_state == RDMA_ON) { - return _rdma_ep->CutFromIOBufList(buf, ndata); - } - return _tcp_transport->CutFromIOBufList(buf, ndata); + CHECK(_rdma_ep != nullptr); + return _rdma_ep->CutFromIOBufList(buf, ndata); } int RdmaTransport::WaitEpollOut(butil::atomic *_epollout_butex, @@ -108,8 +131,7 @@ int RdmaTransport::WaitEpollOut(butil::atomic *_epollout_butex, if (errno != EAGAIN && errno != ETIMEDOUT) { const int saved_errno = errno; PLOG(WARNING) << "Fail to wait rdma window of " << _socket; - _socket->SetFailed(saved_errno, - "Fail to wait rdma window of %s: %s", + _socket->SetFailed(saved_errno, "Fail to wait rdma window of %s: %s", _socket->description().c_str(), berror(saved_errno)); } @@ -122,8 +144,6 @@ int RdmaTransport::WaitEpollOut(butil::atomic *_epollout_butex, } } } - } else { - return _tcp_transport->WaitEpollOut(_epollout_butex, pollin, duetime); } return 0; } @@ -163,15 +183,16 @@ void RdmaTransport::QueueMessage(InputMessageClosure& input_msg, // TODO(gejun): Join threads. bthread_t th; - bthread_attr_t tmp = (FLAGS_usercode_in_pthread ? - BTHREAD_ATTR_PTHREAD : - BTHREAD_ATTR_NORMAL) | BTHREAD_NOSIGNAL; + 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, "ProcessInputMessage"); - if (!FLAGS_usercode_in_coroutine && bthread_start_background( - &th, &tmp, ProcessInputMessage, to_run_msg) == 0) { + if (!FLAGS_usercode_in_coroutine && + bthread_start_background(&th, &tmp, ProcessInputMessage, to_run_msg) == + 0) { ++*num_bthread_created; } else { ProcessInputMessage(to_run_msg); @@ -186,15 +207,18 @@ void RdmaTransport::Debug(std::ostream &os) { int RdmaTransport::ContextInitOrDie(bool serverOrNot, const void* _options) { if (serverOrNot) { - if (!OptionsAvailableOverRdma(static_cast(_options))) { + if (!OptionsAvailableOverRdma( + static_cast(_options))) { return -1; } rdma::GlobalRdmaInitializeOrDie(); - if (!rdma::InitPollingModeWithTag(static_cast(_options)->bthread_tag)) { + if (!rdma::InitPollingModeWithTag( + static_cast(_options)->bthread_tag)) { return -1; } } else { - if (!OptionsAvailableForRdma(static_cast(_options))) { + if (!OptionsAvailableForRdma( + static_cast(_options))) { return -1; } rdma::GlobalRdmaInitializeOrDie(); @@ -213,8 +237,7 @@ bool RdmaTransport::OptionsAvailableForRdma(const ChannelOptions* opt) { return false; } if (!rdma::SupportedByRdma(opt->protocol.name())) { - LOG(WARNING) << "Cannot use " << opt->protocol.name() - << " over RDMA"; + LOG(WARNING) << "Cannot use " << opt->protocol.name() << " over RDMA"; return false; } return true; diff --git a/src/brpc/rdma_transport.h b/src/brpc/rdma_transport.h index 1d78fbb430..2aaf5fadab 100644 --- a/src/brpc/rdma_transport.h +++ b/src/brpc/rdma_transport.h @@ -22,14 +22,15 @@ #include "brpc/socket.h" #include "brpc/channel.h" #include "brpc/transport.h" +#include "brpc/rdma/rdma_endpoint.h" +#include "brpc/handshake/rdma_handshake.h" namespace brpc { +class AdapterTransport; class RdmaTransport : public Transport { -friend class TransportFactory; -friend class rdma::RdmaEndpoint; -friend class rdma::RdmaConnect; -friend class rdma::RdmaHandshakeServerV2; -friend class rdma::RdmaHandshakeServerV3; + friend class TransportFactory; + friend class AdapterTransport; + friend class rdma::RdmaEndpoint; public: void Init(Socket* socket, const SocketOptions& options) override; void Release() override; @@ -37,7 +38,8 @@ friend class rdma::RdmaHandshakeServerV3; std::shared_ptr Connect() override; int CutFromIOBuf(butil::IOBuf* buf) override; ssize_t CutFromIOBufList(butil::IOBuf** buf, size_t ndata) override; - int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, const timespec duetime) override; + int WaitEpollOut(butil::atomic* epollout_butex, + bool pollin, timespec duetime) override; void ProcessEvent(bthread_attr_t attr) override; void QueueMessage(InputMessageClosure& inputMsg, int* num_bthread_created, bool last_msg) override; void Debug(std::ostream &os) override; @@ -45,8 +47,28 @@ friend class rdma::RdmaHandshakeServerV3; CHECK(_rdma_ep != nullptr); return _rdma_ep; } + static RdmaTransport* Get(const Socket* socket); + static RdmaTransport* Get(const SocketUniquePtr& socket) { + return Get(socket.get()); + } static int ContextInitOrDie(bool serverOrNot, const void* _options); + + // Resource operations consumed by the upper-level handshake coordinator. + int PrepareUpgradeResources(); + int StartUpgradeEvents(); + int NegotiateUpgradeResources(const rdma::RdmaConnectionInfo& remote, + bool server); + std::unique_ptr + CreateClientHandshakeAdapter(); + std::vector> + CreateServerHandshakeAdapters(); + void ActivateUpgrade(); + void DeactivateUpgrade(); + bool UpgradeActive() const { return _rdma_state == RDMA_ON; } + bool UpgradeReady() const { return _rdma_ep != nullptr; } private: + void SetHighSpeedAvailable(bool available); + static bool OptionsAvailableForRdma(const ChannelOptions* opt); static bool OptionsAvailableOverRdma(const ServerOptions* opt); @@ -60,7 +82,6 @@ friend class rdma::RdmaHandshakeServerV3; rdma::RdmaEndpoint* _rdma_ep = nullptr; // Should use RDMA or not RdmaState _rdma_state; - std::shared_ptr _tcp_transport; }; } // namespace brpc #endif // BRPC_WITH_RDMA diff --git a/src/brpc/socket.h b/src/brpc/socket.h index da701e9ad7..194b8fd2ab 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -56,11 +56,6 @@ class ChannelBalancer; } namespace rdma { class RdmaEndpoint; -class RdmaConnect; -class RdmaHandshakeClientV2; -class RdmaHandshakeServerV2; -class RdmaHandshakeClientV3; -class RdmaHandshakeServerV3; } namespace urma { @@ -74,14 +69,16 @@ class UrmaHandshakeServerV3; namespace ubring { class UBShmEndpoint; - class UBConnect; } - +namespace handshake { +class SocketHandshakeIO; +} class Socket; class AuthContext; class EventDispatcher; class Stream; class Transport; +class AdapterTransport; // Set SO_SNDBUF/SO_RCVBUF according to socket_*_buffer_size flags. void SetSocketBufferOptions(int fd); @@ -342,14 +339,8 @@ friend class policy::ConsistentHashingLoadBalancer; friend class policy::RtmpContext; friend class schan::ChannelBalancer; friend class rdma::RdmaEndpoint; -friend class rdma::RdmaConnect; friend class ubring::UBShmEndpoint; -friend class ubring::UBConnect; friend class UBShmTransport; -friend class rdma::RdmaHandshakeClientV2; -friend class rdma::RdmaHandshakeServerV2; -friend class rdma::RdmaHandshakeClientV3; -friend class rdma::RdmaHandshakeServerV3; friend class urma::UrmaEndpoint; friend class urma::UrmaConnect; friend class urma::UrmaHandshakeClientV2; @@ -364,6 +355,8 @@ friend class VersionedRefWithId; friend class IOEvent; friend void DereferenceSocket(Socket*); friend class Transport; +friend class AdapterTransport; +friend class handshake::SocketHandshakeIO; friend class TcpTransport; friend class RdmaTransport; friend class UrmaTransport; @@ -997,8 +990,10 @@ friend class TransportFactory; SSL* _ssl_session; // owner std::shared_ptr _ssl_ctx; - // Should use SOCKET_MODE_RDMA or SOCKET_MODE_TCP or Other, default is SOCKET_MODE_TCP Transport + // Requested provider: SOCKET_MODE_TCP, SOCKET_MODE_RDMA or another mode. SocketMode _socket_mode; + // The top-level AdapterTransport, which selects TCP or the requested + // accelerated Transport. std::unique_ptr _transport; // Pass from controller, for progressive reading. diff --git a/src/brpc/transport_factory.cpp b/src/brpc/transport_factory.cpp index 31bb168801..4c450bca85 100644 --- a/src/brpc/transport_factory.cpp +++ b/src/brpc/transport_factory.cpp @@ -16,6 +16,7 @@ // under the License. #include "brpc/transport_factory.h" +#include "brpc/adapter_transport.h" #include "brpc/rdma_transport.h" #include "brpc/tcp_transport.h" #include "brpc/ubshm_transport.h" @@ -49,11 +50,11 @@ int TransportFactory::ContextInitOrDie( std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { if (mode == SOCKET_MODE_TCP) { - return std::unique_ptr(new TcpTransport()); + return std::unique_ptr(new AdapterTransport(mode)); } #if BRPC_WITH_RDMA if (mode == SOCKET_MODE_RDMA) { - return std::unique_ptr(new RdmaTransport()); + return std::unique_ptr(new AdapterTransport(mode)); } #endif #if BRPC_WITH_URMA @@ -63,7 +64,7 @@ std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { #endif #if BRPC_WITH_UBRING if (mode == SOCKET_MODE_UBRING) { - return std::unique_ptr(new UBShmTransport()); + return std::unique_ptr(new AdapterTransport(mode)); } #endif LOG(ERROR) << "Unknown transport type " << mode; diff --git a/src/brpc/transport_factory.h b/src/brpc/transport_factory.h index 84b047daac..0cb6f355b2 100644 --- a/src/brpc/transport_factory.h +++ b/src/brpc/transport_factory.h @@ -22,8 +22,8 @@ #include "brpc/transport.h" namespace brpc { - -// Creates transport instances for a SocketMode. +// Creates AdapterTransport for TCP, RDMA, and UBSHM sockets. URMA currently +// uses its concrete transport directly. class TransportFactory { public: static int ContextInitOrDie(SocketMode mode, bool server_or_not, diff --git a/src/brpc/transport_handshake.cpp b/src/brpc/transport_handshake.cpp new file mode 100644 index 0000000000..e7fd9580c5 --- /dev/null +++ b/src/brpc/transport_handshake.cpp @@ -0,0 +1,425 @@ +// 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_handshake.h" + +#include +#include + +#include "butil/logging.h" +#include "butil/object_pool.h" +#include "butil/sys_byteorder.h" + +namespace brpc { +namespace handshake { + +namespace { + +const size_t COMMON_ACK_SIZE = sizeof(uint32_t); +const uint32_t COMMON_ACK_OK = 0x1; + +} // namespace + +StepResult HandshakeProtocol::BuildAck(bool enabled, std::string* payload) { + CHECK(payload != NULL); + const uint32_t flags = butil::HostToNet32(enabled ? COMMON_ACK_OK : 0); + payload->assign(reinterpret_cast(&flags), sizeof(flags)); + return STEP_OK; +} + +StepResult HandshakeProtocol::ParseAck(const std::string& payload, + bool* enabled) { + CHECK(enabled != NULL); + if (payload.size() != COMMON_ACK_SIZE) { + errno = EPROTO; + return STEP_ERROR; + } + uint32_t flags = 0; + memcpy(&flags, payload.data(), sizeof(flags)); + *enabled = (butil::NetToHost32(flags) & COMMON_ACK_OK) != 0; + return STEP_OK; +} + +const FrameSpec& HandshakeProtocol::ExtensionFrameSpec() const { + static const FrameSpec empty_spec( + NULL, 0, 0, 0, FrameSpec::FIXED); + return empty_spec; +} + +ServerHandshakeContext* ServerHandshakeContext::Create( + HandshakeAdapter* adapter) { + ServerHandshakeContext* context = + butil::get_object(); + if (context != NULL) { + context->_adapter = adapter; + } + return context; +} + +void ServerHandshakeContext::Destroy() { + _adapter = NULL; + butil::return_object(this); +} + +static StepResult FinishWithFailure(HandshakeSession* session, + HandshakeTransport* transport) { + transport->OnFailed(); + session->MarkFailed(); + return STEP_ERROR; +} + +static StepResult FinishWithFallback(HandshakeSession* session, + HandshakeTransport* transport) { + session->PublishFallback([transport]() { transport->OnFallback(); }); + return STEP_FALLBACK; +} + +static StepResult ConvertFrameResult(FrameResult result) { + switch (result) { + case FRAME_OK: return STEP_OK; + case FRAME_NOT_MINE: return STEP_NOT_MINE; + case FRAME_NEED_MORE: return STEP_NEED_MORE; + case FRAME_IO_ERROR: return STEP_ERROR; + case FRAME_PROTOCOL_ERROR: + errno = EPROTO; + return STEP_ERROR; + } + errno = EPROTO; + return STEP_ERROR; +} + +StepResult HandshakeSession::SendHello(HandshakeProtocol* protocol, + bool enabled) { + std::string payload; + const StepResult result = protocol->BuildHello(enabled, &payload); + if (result != STEP_OK) { + return result; + } + return ConvertFrameResult( + FrameCodec::WriteFrame(_io, protocol->HelloFrameSpec(), payload)); +} + +StepResult HandshakeSession::ReceiveHello(HandshakeProtocol* protocol, + HandshakeInput* input, + bool push_back_on_not_mine, + bool* magic_matched) { + std::string payload; + const FrameResult frame_result = input != NULL + ? FrameCodec::ParseBufferedFrame( + input, protocol->HelloFrameSpec(), &payload, magic_matched) + : FrameCodec::ReadFrame( + _io, protocol->HelloFrameSpec(), push_back_on_not_mine, + &payload); + const StepResult result = ConvertFrameResult(frame_result); + if (result != STEP_OK) { + return result; + } + set_protocol_version(protocol->ProtocolVersion()); + return protocol->ParseHello(payload); +} + +StepResult HandshakeSession::SendAck(HandshakeProtocol* protocol, + bool enabled) { + std::string payload; + const StepResult result = protocol->BuildAck(enabled, &payload); + if (result != STEP_OK) { + return result; + } + return ConvertFrameResult( + FrameCodec::WriteFrame(_io, protocol->AckFrameSpec(), payload)); +} + +StepResult HandshakeSession::ReceiveAck(HandshakeProtocol* protocol, + HandshakeInput* input, + bool* enabled) { + std::string payload; + const FrameResult frame_result = input != NULL + ? FrameCodec::ParseBufferedFrame( + input, protocol->AckFrameSpec(), &payload) + : FrameCodec::ReadFrame( + _io, protocol->AckFrameSpec(), false, &payload); + const StepResult result = ConvertFrameResult(frame_result); + if (result != STEP_OK) { + return result; + } + return protocol->ParseAck(payload, enabled); +} + +StepResult HandshakeSession::SendExtension(HandshakeProtocol* protocol, + bool enabled) { + std::string payload; + const StepResult result = protocol->BuildExtension(enabled, &payload); + if (result != STEP_OK) { + return result; + } + return ConvertFrameResult( + FrameCodec::WriteFrame( + _io, protocol->ExtensionFrameSpec(), payload)); +} + +StepResult HandshakeSession::ReceiveExtension(HandshakeProtocol* protocol, + HandshakeInput* input) { + std::string payload; + const FrameResult frame_result = input != NULL + ? FrameCodec::ParseBufferedFrame( + input, protocol->ExtensionFrameSpec(), &payload) + : FrameCodec::ReadFrame( + _io, protocol->ExtensionFrameSpec(), false, &payload); + const StepResult result = ConvertFrameResult(frame_result); + return result == STEP_OK ? protocol->ParseExtension(payload) : result; +} + +StepResult HandshakeSession::SelectAndReceiveHello( + const std::vector& protocols, HandshakeInput* input, + bool push_back_on_not_mine, HandshakeProtocol** selected) { + CHECK(!protocols.empty()); + CHECK(selected != NULL); + if (input == NULL) { + // A blocking byte stream cannot try a second codec after consuming + // bytes from the fd. Such protocols must select a single codec before + // entering the common session. + CHECK_EQ(1UL, protocols.size()); + *selected = protocols.front(); + return ReceiveHello(*selected, NULL, push_back_on_not_mine); + } + + bool need_more = false; + for (size_t i = 0; i < protocols.size(); ++i) { + bool magic_matched = false; + const StepResult result = ReceiveHello( + protocols[i], input, false, &magic_matched); + if (result == STEP_NOT_MINE) { + continue; + } + if (result == STEP_NEED_MORE) { + if (magic_matched) { + *selected = protocols[i]; + set_protocol_version(protocols[i]->ProtocolVersion()); + return STEP_NEED_MORE; + } + need_more = true; + continue; + } + *selected = protocols[i]; + return result; + } + return need_more ? STEP_NEED_MORE : STEP_NOT_MINE; +} + +StepResult HandshakeSession::RunClient(HandshakeProtocol* protocol, + HandshakeTransport* transport) { + CHECK(protocol != NULL); + CHECK(transport != NULL); + transport->OnProtocolSelected(protocol); + // A client handshake runs once on a potentially reused bthread. Do not + // let an errno left by earlier work override this handshake's result. + errno = 0; + + SetPhase(PREPARING); + StepResult result = transport->PrepareResources(); + if (result == STEP_FALLBACK) { + return FinishWithFallback(this, transport); + } + if (result != STEP_OK) { + return FinishWithFailure(this, transport); + } + + SetPhase(HELLO_SEND); + if (SendHello(protocol, true) != STEP_OK) { + return FinishWithFailure(this, transport); + } + + SetPhase(HELLO_WAIT); + result = ReceiveHello(protocol, NULL, false); + if (result == STEP_NOT_MINE || result == STEP_NEED_MORE) { + errno = EPROTO; + } + if (result == STEP_ERROR || result == STEP_NOT_MINE || + result == STEP_NEED_MORE) { + return FinishWithFailure(this, transport); + } + bool enabled = result == STEP_OK; + + if (enabled && protocol->HasExtension()) { + SetPhase(EXTENSION_SEND); + if (SendExtension(protocol, true) != STEP_OK) { + return FinishWithFailure(this, transport); + } + SetPhase(EXTENSION_WAIT); + result = ReceiveExtension(protocol, NULL); + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + enabled = result == STEP_OK; + } + + if (enabled) { + SetPhase(NEGOTIATING); + result = transport->NegotiateResources(); + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + enabled = result == STEP_OK; + } + + SetPhase(ACK_SEND); + if (SendAck(protocol, enabled) != STEP_OK) { + return FinishWithFailure(this, transport); + } + + if (enabled) { + transport->OnEstablished(); + MarkEstablished(); + return STEP_OK; + } + return FinishWithFallback(this, transport); +} + +StepResult HandshakeSession::RunServer( + const std::vector& protocols, HandshakeInput* input, + HandshakeTransport* transport, bool fallback_on_not_mine) { + CHECK(!protocols.empty()); + CHECK(transport != NULL); + + // Once TCP fallback has been published, subsequent bytes are application + // protocol data and must bypass every upgrade codec without changing the + // terminal state. + if (phase(butil::memory_order_acquire) == FALLBACK_TCP) { + return STEP_NOT_MINE; + } + + HandshakeProtocol* selected = NULL; + if (phase(butil::memory_order_acquire) != ACK_WAIT && + phase(butil::memory_order_acquire) != EXTENSION_WAIT) { + const int previous_phase = phase(butil::memory_order_acquire); + _local_enabled = false; + SetPhase(HELLO_WAIT); + StepResult result = SelectAndReceiveHello( + protocols, input, fallback_on_not_mine, &selected); + if (result == STEP_NOT_MINE) { + if (fallback_on_not_mine) { + return FinishWithFallback(this, transport); + } + SetPhase(UNINITIALIZED); + return STEP_NOT_MINE; + } + if (result == STEP_NEED_MORE) { + if (selected == NULL) { + SetPhase(previous_phase); + } + return STEP_NEED_MORE; + } + if (result == STEP_ERROR) { + return FinishWithFailure(this, transport); + } + CHECK(selected != NULL); + transport->OnProtocolSelected(selected); + if (result == STEP_FALLBACK) { + // Publish the disabled transport state immediately. The server + // still sends a disabled hello and consumes the peer ACK before + // publishing the terminal FALLBACK_TCP phase. + transport->OnFallback(); + } + bool enabled = result == STEP_OK; + + if (enabled) { + SetPhase(PREPARING); + result = transport->PrepareResources(); + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + enabled = result == STEP_OK; + } + + if (enabled) { + SetPhase(NEGOTIATING); + result = transport->NegotiateResources(); + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + enabled = result == STEP_OK; + } + + SetPhase(HELLO_SEND); + _local_enabled = enabled; + if (SendHello(selected, enabled) != STEP_OK) { + return FinishWithFailure(this, transport); + } + SetPhase(enabled && selected->HasExtension() + ? EXTENSION_WAIT + : ACK_WAIT); + } else { + for (size_t i = 0; i < protocols.size(); ++i) { + if (protocols[i]->ProtocolVersion() == protocol_version()) { + selected = protocols[i]; + break; + } + } + CHECK(selected != NULL); + } + + if (phase(butil::memory_order_acquire) == EXTENSION_WAIT) { + StepResult result = ReceiveExtension(selected, input); + if (result == STEP_NEED_MORE) { + return STEP_NEED_MORE; + } + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + if (result == STEP_FALLBACK) { + _local_enabled = false; + } + SetPhase(EXTENSION_SEND); + if (SendExtension(selected, _local_enabled) != STEP_OK) { + return FinishWithFailure(this, transport); + } + SetPhase(ACK_WAIT); + } + + // Always try the ACK callback once. For a non-blocking server it returns + // STEP_NEED_MORE when the ACK has not arrived; when Hello and ACK are + // coalesced in the input buffer this consumes the ACK without waiting for + // another socket edge. + bool peer_enabled = false; + StepResult result = ReceiveAck(selected, input, &peer_enabled); + if (result == STEP_NEED_MORE) { + return STEP_NEED_MORE; + } + if (result == STEP_ERROR || result == STEP_NOT_MINE) { + return FinishWithFailure(this, transport); + } + if (result == STEP_FALLBACK) { + return FinishWithFallback(this, transport); + } + if (!peer_enabled) { + return FinishWithFallback(this, transport); + } + if (!_local_enabled) { + errno = EPROTO; + return FinishWithFailure(this, transport); + } + if (transport->ValidateEstablished() != STEP_OK) { + return FinishWithFailure(this, transport); + } + + transport->OnEstablished(); + MarkEstablished(); + return STEP_OK; +} + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/transport_handshake.h b/src/brpc/transport_handshake.h new file mode 100644 index 0000000000..3228448449 --- /dev/null +++ b/src/brpc/transport_handshake.h @@ -0,0 +1,233 @@ +// 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_TRANSPORT_HANDSHAKE_H +#define BRPC_TRANSPORT_HANDSHAKE_H + +#include +#include +#include + +#include "butil/atomicops.h" +#include "butil/macros.h" +#include "brpc/destroyable.h" +#include "brpc/handshake/handshake_frame.h" + +namespace brpc { + +class Socket; + +namespace handshake { + +class HandshakeAdapter; + +// Context retained by InputMessenger between the hello and ACK parse calls. +// Remembering the selected stateless adapter is necessary because ACK frames +// have no magic and cannot be dispatched from their bytes alone. +struct ServerHandshakeContext : public Destroyable { + ServerHandshakeContext() : _adapter(NULL) {} + static ServerHandshakeContext* Create(HandshakeAdapter* adapter); + HandshakeAdapter* adapter() const { return _adapter; } + void Destroy() override; + +private: + HandshakeAdapter* _adapter; +}; + +// Protocol adapters may use transport-specific intermediate values, but the +// terminal values are shared so that AdapterTransport can make the same +// acquire-side decision for RDMA, URMA and UBSHM. +enum Phase { + UNINITIALIZED = 0, + PREPARING = 1, + HELLO_SEND = 2, + HELLO_WAIT = 3, + NEGOTIATING = 4, + ACK_SEND = 5, + ACK_WAIT = 6, + EXTENSION_SEND = 7, + EXTENSION_WAIT = 8, + ESTABLISHED = 0x100, + FALLBACK_TCP = 0x200, + FAILED = 0x300, +}; + +enum StepResult { + STEP_OK = 0, + STEP_FALLBACK, + STEP_NEED_MORE, + STEP_NOT_MINE, + STEP_ERROR, +}; + +// Wire-level participant in a transport upgrade. Implementations own parsed +// protocol state; HandshakeSession owns framing and phase orchestration. +class HandshakeProtocol { +public: + virtual ~HandshakeProtocol() = default; + + virtual int ProtocolVersion() const = 0; + virtual const FrameSpec& HelloFrameSpec() const = 0; + virtual const FrameSpec& AckFrameSpec() const = 0; + virtual StepResult BuildHello(bool enabled, std::string* payload) = 0; + virtual StepResult ParseHello(const std::string& payload) = 0; + + // RDMA and UBSHM use the same four-byte, network-order ACK. Protocols + // with a different ACK format may override these methods. + virtual StepResult BuildAck(bool enabled, std::string* payload); + virtual StepResult ParseAck(const std::string& payload, bool* enabled); + + virtual bool HasExtension() const { return false; } + virtual const FrameSpec& ExtensionFrameSpec() const; + virtual StepResult BuildExtension(bool, std::string*) { + return STEP_ERROR; + } + virtual StepResult ParseExtension(const std::string&) { + return STEP_ERROR; + } +}; + +// Resource-level participant in a transport upgrade. Cleanup remains part of +// this contract: resources may already exist when negotiation falls back or +// a framing/I/O error terminates the handshake. +class HandshakeTransport { +public: + virtual ~HandshakeTransport() = default; + + virtual void OnProtocolSelected(HandshakeProtocol*) {} + virtual StepResult PrepareResources() = 0; + virtual StepResult NegotiateResources() = 0; + virtual void OnEstablished() = 0; + virtual void OnFallback() = 0; + virtual void OnFailed() = 0; + virtual StepResult ValidateEstablished() { return STEP_OK; } +}; + +// Participant for a server that recognizes an upgrade protocol only to +// negotiate TCP fallback. It owns no high-speed resources. +class FallbackHandshakeTransport : public HandshakeTransport { +public: + StepResult PrepareResources() override { return STEP_OK; } + StepResult NegotiateResources() override { return STEP_OK; } + void OnEstablished() override {} + void OnFallback() override {} + void OnFailed() override {} +}; + +class FallbackHandshakeProtocol : public HandshakeProtocol { +public: + StepResult ParseHello(const std::string&) override { + return STEP_FALLBACK; + } + StepResult ParseAck(const std::string& payload, bool* enabled) override { + bool ignored = false; + const StepResult result = HandshakeProtocol::ParseAck( + payload, &ignored); + *enabled = false; + return result; + } +}; + +// Owns one connection-upgrade attempt, invokes the protocol field codec and +// resource callbacks, and provides common framing, TCP control-plane I/O, +// lifecycle and publication ordering. +class HandshakeSession { +public: + explicit HandshakeSession(Socket* socket = NULL) + : _socket_io(socket), _io(&_socket_io), _phase(UNINITIALIZED), + _protocol_version(0), _local_enabled(false) {} + + void Reset(Socket* socket) { + _socket_io.Reset(socket); + _io = &_socket_io; + _protocol_version = 0; + _local_enabled = false; + _phase.store(UNINITIALIZED, butil::memory_order_relaxed); + } + + int phase(butil::memory_order order) const { + return _phase.load(order); + } + + void SetPhase(int phase) { + _phase.store(phase, butil::memory_order_release); + } + + int protocol_version() const { return _protocol_version; } + void set_protocol_version(int version) { _protocol_version = version; } + + void MarkEstablished() { + _phase.store(ESTABLISHED, butil::memory_order_release); + } + + void MarkFailed() { + _phase.store(FAILED, butil::memory_order_release); + } + + // The callback MUST publish the transport's TCP-active state. The release + // store then makes that state and any pushed-back bytes visible to the + // event thread that observes FALLBACK_TCP with an acquire load. This is + // the common form of the ordering fixes from #3347 and #3406. + template + void PublishFallback(PublishTcpActive publish_tcp_active) { + publish_tcp_active(); + _phase.store(FALLBACK_TCP, butil::memory_order_release); + } + + void NotifyReadable() { _socket_io.NotifyReadable(); } + + // Injects an in-memory stream in common-component unit tests. Reset() + // restores the Socket-backed implementation. + void SetIOForTest(HandshakeIO* io) { _io = io; } + + StepResult RunClient(HandshakeProtocol* protocol, + HandshakeTransport* transport); + StepResult RunServer(const std::vector& protocols, + HandshakeInput* input, + HandshakeTransport* transport, + bool fallback_on_not_mine); + +private: + StepResult SendHello(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveHello(HandshakeProtocol* protocol, + HandshakeInput* input, + bool push_back_on_not_mine, + bool* magic_matched = NULL); + StepResult SendAck(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveAck(HandshakeProtocol* protocol, + HandshakeInput* input, bool* enabled); + StepResult SendExtension(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveExtension(HandshakeProtocol* protocol, + HandshakeInput* input); + StepResult SelectAndReceiveHello( + const std::vector& protocols, + HandshakeInput* input, bool push_back_on_not_mine, + HandshakeProtocol** selected); + + SocketHandshakeIO _socket_io; + HandshakeIO* _io; + butil::atomic _phase; + int _protocol_version; + bool _local_enabled; + + DISALLOW_COPY_AND_ASSIGN(HandshakeSession); +}; + +} // namespace handshake +} // namespace brpc + +#endif // BRPC_TRANSPORT_HANDSHAKE_H diff --git a/src/brpc/ubshm/ub_endpoint.cpp b/src/brpc/ubshm/ub_endpoint.cpp index 44685eb302..b52ff16af0 100644 --- a/src/brpc/ubshm/ub_endpoint.cpp +++ b/src/brpc/ubshm/ub_endpoint.cpp @@ -19,23 +19,19 @@ #include -#include -#include -#include "butil/fd_utility.h" -#include "butil/logging.h" // CHECK, LOG -#include "butil/sys_byteorder.h" // HostToNet,NetToHost -#include "bthread/bthread.h" #include "brpc/errno.pb.h" #include "brpc/event_dispatcher.h" #include "brpc/input_messenger.h" #include "brpc/socket.h" -#include "brpc/reloadable_flags.h" -#include "brpc/ubshm/ub_helper.h" -#include "brpc/ubshm/ub_endpoint.h" -#include "brpc/ubshm/shm/shm_def.h" #include "brpc/ubshm/common/common.h" -#include "brpc/ubshm_transport.h" +#include "brpc/ubshm/shm/shm_def.h" +#include "brpc/ubshm/ub_endpoint.h" +#include "brpc/ubshm/ub_helper.h" #include "brpc/ubshm/ubr_trx.h" +#include "brpc/ubshm_transport.h" +#include "bthread/bthread.h" +#include "butil/logging.h" // CHECK, LOG +#include DECLARE_int32(task_group_ntags); @@ -44,658 +40,34 @@ DECLARE_bool(log_connection_close); namespace ubring { extern bool g_skip_ub_init; -DEFINE_int32(data_queue_size, 4, "data queue size for UB"); -DEFINE_bool(ub_trace_verbose, false, "Print log message verbosely"); -BRPC_VALIDATE_GFLAG(ub_trace_verbose, brpc::PassValidate); DEFINE_int32(ub_poller_num, 1, "Poller number in ub polling mode."); 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"; -static const size_t MAGIC_STR_LEN = 2; -static const size_t HELLO_MSG_LEN_MIN = 64; -static const size_t ACK_MSG_LEN = 4; -static uint16_t g_ub_hello_msg_len = 64; -static uint16_t g_ub_hello_version = 3; -static uint16_t g_ub_impl_version = 1; - -static const uint32_t ACK_MSG_UB_OK = 0x1; - -static butil::Mutex* g_ubring_resource_mutex = nullptr; - -void HelloFormatExtension::Serialize(void* data) const { - char* current_pos = static_cast(data); - const uint16_t net_extension_len = butil::HostToNet16(extension_len); - memcpy(current_pos, &net_extension_len, sizeof(net_extension_len)); - current_pos += sizeof(net_extension_len); - const uint16_t net_format_id = butil::HostToNet16(format_id); - memcpy(current_pos, &net_format_id, sizeof(net_format_id)); -} - -void HelloFormatExtension::Deserialize(const void* data) { - const char* current_pos = static_cast(data); - uint16_t net_extension_len; - memcpy(&net_extension_len, current_pos, sizeof(net_extension_len)); - extension_len = butil::NetToHost16(net_extension_len); - current_pos += sizeof(net_extension_len); - uint16_t net_format_id; - memcpy(&net_format_id, current_pos, sizeof(net_format_id)); - format_id = butil::NetToHost16(net_format_id); -} +static butil::Mutex *g_ubring_resource_mutex = NULL; -void HelloMessage::Serialize(void* data) const { - char* current_pos = static_cast(data); - const uint16_t net_msg_len = butil::HostToNet16(msg_len); - memcpy(current_pos, &net_msg_len, sizeof(net_msg_len)); - current_pos += sizeof(net_msg_len); - const uint16_t net_hello_ver = butil::HostToNet16(hello_ver); - memcpy(current_pos, &net_hello_ver, sizeof(net_hello_ver)); - current_pos += sizeof(net_hello_ver); - const uint16_t net_impl_ver = butil::HostToNet16(impl_ver); - memcpy(current_pos, &net_impl_ver, sizeof(net_impl_ver)); - current_pos += sizeof(net_impl_ver); - const uint64_t net_len = butil::HostToNet64(len); - memcpy(current_pos, &net_len, sizeof(net_len)); - current_pos += sizeof(net_len); - memcpy(current_pos, shm_name, SHM_MAX_NAME_BUFF_LEN); -} - -void HelloMessage::Deserialize(void* data) { - char* current_pos = static_cast(data); - uint16_t net_msg_len; - memcpy(&net_msg_len, current_pos, sizeof(net_msg_len)); - msg_len = butil::NetToHost16(net_msg_len); - current_pos += sizeof(net_msg_len); - uint16_t net_hello_ver; - memcpy(&net_hello_ver, current_pos, sizeof(net_hello_ver)); - hello_ver = butil::NetToHost16(net_hello_ver); - current_pos += sizeof(net_hello_ver); - uint16_t net_impl_ver; - memcpy(&net_impl_ver, current_pos, sizeof(net_impl_ver)); - impl_ver = butil::NetToHost16(net_impl_ver); - current_pos += sizeof(net_impl_ver); - uint64_t net_len; - memcpy(&net_len, current_pos, sizeof(net_len)); - len = butil::NetToHost64(net_len); - current_pos += sizeof(net_len); - memcpy(shm_name, current_pos, SHM_MAX_NAME_BUFF_LEN); -} - -std::string HelloMessage::toString() const { - constexpr size_t MAX_LEN = 16 + 6 + 16 + 6 + 16 + 6 + 20 + 6 + SHM_MAX_NAME_BUFF_LEN + 32; - std::array buf; - int n = snprintf(buf.data(), buf.size(), - "msg_len=%u, hello_ver=%u, impl_ver=%u, len=%lu, shm_name=%.*s", - msg_len, - hello_ver, - impl_ver, - static_cast(len), // compatible with 32/64-bit - static_cast(SHM_MAX_NAME_BUFF_LEN), // limit max output length - shm_name - ); - return std::string(buf.data(), static_cast(n)); -} - -UBShmEndpoint::UBShmEndpoint(Socket* s) - : _socket(s) - , _socket_id(s ? s->id() : INVALID_SOCKET_ID) - , _state(UNINIT) - , _ub_ring(nullptr) - , _poller_sid(INVALID_SOCKET_ID) -{ - _read_butex = bthread::butex_create_checked>(); +UBShmEndpoint::UBShmEndpoint(Socket *s) + : _socket(s), _socket_id(s ? s->id() : INVALID_SOCKET_ID), + _ub_ring(nullptr), _poller_sid(INVALID_SOCKET_ID) { } UBShmEndpoint::~UBShmEndpoint() { Reset(); - bthread::butex_destroy(_read_butex); } void UBShmEndpoint::Reset() { + _receive_events_started.store(false, butil::memory_order_release); DeallocateResources(); delete _ub_ring; _ub_ring = nullptr; _poller_sid = INVALID_SOCKET_ID; _negotiated_data_format = UBR_DATA_FORMAT_NONE; - _state = UNINIT; -} - -void UBConnect::StartConnect(const Socket* socket, - void (*done)(int err, void* data), - void* data) { - auto* ub_transport = static_cast(socket->_transport.get()); - CHECK(ub_transport->_ub_ep != nullptr); - SocketUniquePtr s; - if (Socket::Address(socket->id(), &s) != 0) { - return; - } - if (!IsUBAvailable()) { - ub_transport->_ub_ep->_state = UBShmEndpoint::FALLBACK_TCP; - ub_transport->_ub_state = UBShmTransport::UB_OFF; - done(0, data); - return; - } - _done = done; - _data = data; - bthread_t tid; - bthread_attr_t attr = BTHREAD_ATTR_NORMAL; - bthread_attr_set_name(&attr, "UBProcessHandshakeAtClient"); - if (bthread_start_background(&tid, &attr, - UBShmEndpoint::ProcessHandshakeAtClient, ub_transport->_ub_ep) < 0) { - LOG(FATAL) << "Fail to start handshake bthread"; - Run(); - } else { - s.release(); - } -} - -void UBConnect::StopConnect(Socket* socket) { } - -void UBConnect::Run() { - _done(errno, _data); -} - -static void TryReadOnTcpDuringRdmaEst(Socket* s) { - int progress = Socket::PROGRESS_INIT; - while (true) { - uint8_t tmp; - ssize_t nr = read(s->fd(), &tmp, 1); - if (nr < 0) { - if (errno != EAGAIN) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read from " << s; - s->SetFailed(saved_errno, "Fail to read from %s: %s", - s->description().c_str(), berror(saved_errno)); - return; - } - if (!s->MoreReadEvents(&progress)) { - break; - } - } else if (nr == 0) { - s->SetEOF(); - return; - } else { - LOG(WARNING) << "Read unexpected data from " << s; - s->SetFailed(EPROTO, "Read unexpected data from %s", - s->description().c_str()); - return; - } - } -} - -void UBShmEndpoint::OnNewDataFromTcp(Socket* m) { - auto* ub_transport = static_cast(m->_transport.get()); - UBShmEndpoint* ep = ub_transport->GetUBShmEp(); - CHECK(ep != nullptr); - - int progress = Socket::PROGRESS_INIT; - while (true) { - if (ep->_state == UNINIT) { - if (!m->CreatedByConnect()) { - if (!IsUBAvailable()) { - ep->_state = FALLBACK_TCP; - ub_transport->_ub_state = UBShmTransport::UB_OFF; - continue; - } - bthread_t tid; - ep->_state = S_HELLO_WAIT; - SocketUniquePtr s; - m->ReAddress(&s); - bthread_attr_t attr = BTHREAD_ATTR_NORMAL; - bthread_attr_set_name(&attr, "UBProcessHandshakeAtServer"); - if (bthread_start_background(&tid, &attr, - ProcessHandshakeAtServer, ep) < 0) { - ep->_state = UNINIT; - LOG(FATAL) << "Fail to start handshake bthread"; - } else { - s.release(); - } - } else { - // The connection may be closed or reset before the client - // starts handshake. This will be handled by client handshake. - // Ignore the exception here. - } - } else if (ep->_state < ESTABLISHED) { // during handshake - ep->_read_butex->fetch_add(1, butil::memory_order_release); - bthread::butex_wake(ep->_read_butex); - } else if (ep->_state == FALLBACK_TCP){ // handshake finishes - InputMessenger::OnNewMessages(m); - return; - } else if (ep->_state == ESTABLISHED) { - TryReadOnTcpDuringRdmaEst(ep->_socket); - return; - } - if (!m->MoreReadEvents(&progress)) { - break; - } - } -} -bool HelloNegotiationValid(HelloMessage& msg) { - if (msg.hello_ver == g_ub_hello_version && - msg.impl_ver == g_ub_impl_version) { - // This can be modified for future compatibility - return true; - } - return false; -} - -static const int WAIT_TIMEOUT_MS = 50; - -int UBShmEndpoint::ReadFromFd(void* data, size_t len) { - CHECK(data != nullptr); - int nr = 0; - size_t received = 0; - do { - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - nr = read(_socket->fd(), (uint8_t*)data + received, len - received); - if (nr < 0) { - if (errno == EAGAIN) { - const int expected_val = _read_butex->load(butil::memory_order_acquire); - if (bthread::butex_wait(_read_butex, expected_val, &duetime) < 0) { - if (errno != EWOULDBLOCK && errno != ETIMEDOUT) { - return -1; - } - } - } else { - return -1; - } - } else if (nr == 0) { - errno = EEOF; - return -1; - } else { - received += nr; - } - } while (received < len); - return 0; -} - -int UBShmEndpoint::WriteToFd(void* data, size_t len) { - CHECK(data != nullptr); - int nw = 0; - size_t written = 0; - do { - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - nw = write(_socket->fd(), (uint8_t*)data + written, len - written); - if (nw < 0) { - if (errno == EAGAIN) { - if (_socket->WaitEpollOut(_socket->fd(), true, &duetime) < 0) { - if (errno != ETIMEDOUT) { - return -1; - } - } - } else { - return -1; - } - } else { - written += nw; - } - } while (written < len); - return 0; -} - -inline void UBShmEndpoint::TryReadOnTcp() { - if (_socket->_nevent.fetch_add(1, butil::memory_order_acq_rel) == 0) { - if (_state == FALLBACK_TCP) { - InputMessenger::OnNewMessages(_socket); - } else if (_state == ESTABLISHED) { - TryReadOnTcpDuringRdmaEst(_socket); - } - } -} - -void* UBShmEndpoint::ProcessHandshakeAtClient(void* arg) { - UBShmEndpoint* ep = static_cast(arg); - ep->_negotiated_data_format = UBR_DATA_FORMAT_NONE; - SocketUniquePtr s(ep->_socket); - UBConnect::RunGuard rg((UBConnect*)s->_app_connect.get()); - - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Start handshake on " << s->_local_side; - - uint8_t data[g_ub_hello_msg_len]; - - ep->_state = C_ALLOC_SHM; - auto* ub_transport = static_cast(s->_transport.get()); - size_t local_shm_len = (size_t)(FLAGS_data_queue_size) * MB_TO_BYTE; - SHM local_trx_shm = {nullptr, local_shm_len, 0, {0}, (uint32_t)s->fd()}; - auto shm_name_str = butil::endpoint2str(s->local_side()); - const char* shm_name = shm_name_str.c_str(); - if (ep->AllocateClientResources(&local_trx_shm, shm_name) < 0) { - LOG(WARNING) << "Fallback to tcp:" << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - ep->_state = FALLBACK_TCP; - return nullptr; - } - - ep->_state = C_HELLO_SEND; - HelloMessage local_msg{}; - local_msg.msg_len = g_ub_hello_msg_len; - local_msg.hello_ver = g_ub_hello_version; - local_msg.impl_ver = g_ub_impl_version; - local_msg.len = local_shm_len; - memcpy(local_msg.shm_name, local_trx_shm.name, SHM_MAX_NAME_BUFF_LEN); - memcpy(data, MAGIC_STR, MAGIC_STR_LEN); - local_msg.Serialize((char*)data + MAGIC_STR_LEN); - if (ep->WriteToFd(data, g_ub_hello_msg_len) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to send hello message to server:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - LOG_IF(INFO, FLAGS_ub_trace_verbose) << "client handshake message : " << local_msg.toString(); - - ep->_state = C_HELLO_WAIT; - if (ep->ReadFromFd(data, MAGIC_STR_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to get hello message from server:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - if (memcmp(data, MAGIC_STR, MAGIC_STR_LEN) != 0) { - LOG(WARNING) << "Read unexpected data during handshake:" << s->description(); - s->SetFailed(EPROTO, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(EPROTO)); - ep->_state = FAILED; - return nullptr; - } - - if (ep->ReadFromFd(data, HELLO_MSG_LEN_MIN - MAGIC_STR_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to get Hello Message from server:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - HelloMessage remote_msg; - remote_msg.Deserialize(data); - if (remote_msg.msg_len != HELLO_MSG_LEN_MIN) { - LOG(WARNING) << "Fail to parse Hello Message length from server:" - << s->description(); - s->SetFailed(EPROTO, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(EPROTO)); - ep->_state = FAILED; - return nullptr; - } - - UbrDataFormat selected_format = UBR_DATA_FORMAT_NONE; - if (!HelloNegotiationValid(remote_msg)) { - LOG(WARNING) << "Fail to negotiate with server, fallback to tcp:" - << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } else { - HelloFormatExtension local_extension = { - HelloFormatExtension::WIRE_SIZE, UBR_DATA_FORMAT_LEGACY_64}; - local_extension.Serialize(data); - ep->_state = C_FORMAT_SEND; - if (ep->WriteToFd(data, HelloFormatExtension::WIRE_SIZE) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to send format extension to server:" - << s->description(); - s->SetFailed(saved_errno, - "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - - ep->_state = C_FORMAT_WAIT; - if (ep->ReadFromFd(data, HelloFormatExtension::WIRE_SIZE) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read format extension from server:" - << s->description(); - s->SetFailed(saved_errno, - "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - HelloFormatExtension remote_extension; - remote_extension.Deserialize(data); - if (remote_extension.extension_len != HelloFormatExtension::WIRE_SIZE || - remote_extension.format_id == UBR_DATA_FORMAT_NONE || - remote_extension.format_id != local_extension.format_id) { - LOG(WARNING) << "Fail to negotiate data format with server, " - << "fallback to tcp:" << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } else { - selected_format = UBR_DATA_FORMAT_LEGACY_64; - ep->_state = C_MAP_REMOTE_SHM; - if (ep->_ub_ring->UbrMapRemoteShm(&local_trx_shm, shm_name) < 0) { - LOG(WARNING) << "Fail to map the remote shm, fallback to tcp:" - << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } else { - ub_transport->_ub_state = UBShmTransport::UB_ON; - } - } - } - - ep->_state = C_ACK_SEND; - uint32_t flags = 0; - if (ub_transport->_ub_state != UBShmTransport::UB_OFF) { - flags |= ACK_MSG_UB_OK; - } - const uint32_t net_flags = butil::HostToNet32(flags); - memcpy(data, &net_flags, sizeof(net_flags)); - if (ep->WriteToFd(data, ACK_MSG_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to send Ack Message to server:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - - if (ub_transport->_ub_state == UBShmTransport::UB_ON) { - ep->_negotiated_data_format = selected_format; - ep->_state = ESTABLISHED; - ep->_ub_ring->UbrUnlinkLocalShm(); - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Client handshake ends (use ubring) on " << s->description(); - } else { - ep->_state = FALLBACK_TCP; - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Client handshake ends (use tcp) on " << s->description(); - } - - errno = 0; - - return nullptr; -} - -void* UBShmEndpoint::ProcessHandshakeAtServer(void* arg) { - UBShmEndpoint* ep = static_cast(arg); - ep->_negotiated_data_format = UBR_DATA_FORMAT_NONE; - SocketUniquePtr s(ep->_socket); - - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Start handshake on " << s->description(); - - uint8_t data[g_ub_hello_msg_len]; - - ep->_state = S_HELLO_WAIT; - if (ep->ReadFromFd(data, MAGIC_STR_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read Hello Message from client:" << s->description() << " " << s->_remote_side; - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - auto* ub_transport = static_cast(s->_transport.get()); - if (memcmp(data, MAGIC_STR, MAGIC_STR_LEN) != 0) { - LOG_IF(INFO, FLAGS_ub_trace_verbose) << "It seems that the " - << "client does not use RDMA, fallback to TCP:" - << s->description(); - s->fd_input_processor().read_buf().append(data, MAGIC_STR_LEN); - ep->_state = FALLBACK_TCP; - ub_transport->_ub_state = UBShmTransport::UB_OFF; - ep->TryReadOnTcp(); - return nullptr; - } - - if (ep->ReadFromFd(data, g_ub_hello_msg_len - MAGIC_STR_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read Hello Message from client:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - - HelloMessage remote_msg; - remote_msg.Deserialize(data); - LOG_IF(INFO, FLAGS_ub_trace_verbose) << "server receive handshake message : " << remote_msg.toString(); - if (remote_msg.msg_len != HELLO_MSG_LEN_MIN) { - LOG(WARNING) << "Fail to parse Hello Message length from client:" - << s->description(); - s->SetFailed(EPROTO, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(EPROTO)); - ep->_state = FAILED; - return nullptr; - } - if (!HelloNegotiationValid(remote_msg)) { - LOG(WARNING) << "Fail to negotiate with client, fallback to tcp:" - << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } else { - ep->_state = S_ALLOC_SHM; - ubring::SHM remote_trx_shm = {nullptr, remote_msg.len, 0, {0}, (uint32_t)ep->_socket->fd()}; - strncpy(remote_trx_shm.name, remote_msg.shm_name, SHM_MAX_NAME_BUFF_LEN); - - size_t local_shm_len = (size_t)(FLAGS_data_queue_size) * MB_TO_BYTE; - // server-side shared memory name - ubring::SHM local_trx_shm = {nullptr, local_shm_len, 0, {0}, (uint32_t)ep->_socket->fd()}; - char client_name[SHM_MAX_NAME_BUFF_LEN]; - strncpy(client_name, remote_msg.shm_name, SHM_MAX_NAME_BUFF_LEN); - - char *client_ip_port = strrchr(client_name, '_'); - if (client_ip_port != nullptr) { - *client_ip_port = '\0'; - } - int result = snprintf(local_trx_shm.name, SHM_MAX_NAME_BUFF_LEN, "%s_%s", - client_name, SERVER_SHM_NAME_SUFFIX); - if (UNLIKELY(result < 0)) { - LOG(WARNING) << "Copy client shared memory name failed, ret=" << result; - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } - if (result >= 0 && ep->AllocateServerResources(&remote_trx_shm, &local_trx_shm) < 0) { - LOG(WARNING) << "Fail to allocate ub resources, fallback to tcp:" - << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } - } - - ep->_state = S_HELLO_SEND; - HelloMessage local_msg{}; - local_msg.msg_len = g_ub_hello_msg_len; - if (ub_transport->_ub_state == UBShmTransport::UB_OFF) { - local_msg.impl_ver = 0; - local_msg.hello_ver = 0; - } else { - local_msg.hello_ver = g_ub_hello_version; - local_msg.impl_ver = g_ub_impl_version; - local_msg.len = (FLAGS_data_queue_size) * MB_TO_BYTE; - memcpy(local_msg.shm_name, remote_msg.shm_name, SHM_MAX_NAME_BUFF_LEN); - } - memcpy(data, MAGIC_STR, MAGIC_STR_LEN); - local_msg.Serialize((char*)data + MAGIC_STR_LEN); - if (ep->WriteToFd(data, g_ub_hello_msg_len) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to send Hello Message to client:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ub handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - - UbrDataFormat selected_format = UBR_DATA_FORMAT_NONE; - if (HelloNegotiationValid(remote_msg) && - HelloNegotiationValid(local_msg)) { - ep->_state = S_FORMAT_WAIT; - if (ep->ReadFromFd(data, HelloFormatExtension::WIRE_SIZE) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read format extension from client:" - << s->description(); - s->SetFailed(saved_errno, - "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - HelloFormatExtension remote_extension; - remote_extension.Deserialize(data); - HelloFormatExtension local_extension = { - HelloFormatExtension::WIRE_SIZE, UBR_DATA_FORMAT_NONE}; - if (remote_extension.extension_len == HelloFormatExtension::WIRE_SIZE && - remote_extension.format_id == UBR_DATA_FORMAT_LEGACY_64) { - local_extension.format_id = UBR_DATA_FORMAT_LEGACY_64; - selected_format = UBR_DATA_FORMAT_LEGACY_64; - } else { - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } - local_extension.Serialize(data); - ep->_state = S_FORMAT_SEND; - if (ep->WriteToFd(data, HelloFormatExtension::WIRE_SIZE) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to send format extension to client:" - << s->description(); - s->SetFailed(saved_errno, - "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - } - - ep->_state = S_ACK_WAIT; - if (ep->ReadFromFd(data, ACK_MSG_LEN) < 0) { - const int saved_errno = errno; - PLOG(WARNING) << "Fail to read ack message from client:" << s->description(); - s->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", - s->description().c_str(), berror(saved_errno)); - ep->_state = FAILED; - return nullptr; - } - - uint32_t net_flags; - memcpy(&net_flags, data, sizeof(net_flags)); - const uint32_t flags = butil::NetToHost32(net_flags); - if (flags & ACK_MSG_UB_OK) { - if (ub_transport->_ub_state == UBShmTransport::UB_OFF || - selected_format == UBR_DATA_FORMAT_NONE) { - LOG(WARNING) << "Invalid successful ACK from client:" - << s->description(); - s->SetFailed(EPROTO, "Fail to complete ub handshake from %s: %s", - s->description().c_str(), berror(EPROTO)); - ep->_state = FAILED; - return nullptr; - } else { - ub_transport->_ub_state = UBShmTransport::UB_ON; - ep->_negotiated_data_format = selected_format; - ep->_state = ESTABLISHED; - ep->_ub_ring->UbrUnlinkLocalShm(); - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Server handshake ends (use ubring) on " << s->description(); - } - } else { - ub_transport->_ub_state = UBShmTransport::UB_OFF; - ep->_state = FALLBACK_TCP; - LOG_IF(INFO, FLAGS_ub_trace_verbose) - << "Server handshake ends (use tcp) on " << s->description(); - } - ep->TryReadOnTcp(); - - return nullptr; } bool UBShmEndpoint::IsWritable() const { @@ -736,9 +108,11 @@ ssize_t UBShmEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { nw = _ub_ring->UbrTrxWritev(vec, nvec); if (UNLIKELY(nw == -1)) { if (errno == EMSGSIZE) { - LOG(ERROR) << "Non-blocking send msg failed, message is larger than ubring capacity."; + LOG(ERROR) << "Non-blocking send msg failed, message is larger than " + "ubring capacity."; } else { - LOG(ERROR) << "Non-blocking send msg in failed, connection has been closed."; + LOG(ERROR) + << "Non-blocking send msg in failed, connection has been closed."; errno = EPIPE; } } else if (UNLIKELY(nw == UBRING_RETRY)) { @@ -758,7 +132,8 @@ ssize_t UBShmEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { return nw; } -int UBShmEndpoint::AllocateClientResources(ubring::SHM* local_trx_shm, const char* shm_name) { +int UBShmEndpoint::AllocateClientResources(ubring::SHM *local_trx_shm, + const char *shm_name) { if (BAIDU_UNLIKELY(g_skip_ub_init)) { // For UT return 0; @@ -772,18 +147,29 @@ int UBShmEndpoint::AllocateClientResources(ubring::SHM* local_trx_shm, const cha options.user = this; options.keytable_pool = _socket->_keytable_pool; if (Socket::Create(options, &_poller_sid) < 0) { + const int saved_errno = errno; PLOG(WARNING) << "Fail to create socket for UBRing poller"; + delete _ub_ring; + _ub_ring = NULL; + _poller_sid = INVALID_SOCKET_ID; + errno = saved_errno; return -1; } int ret = _ub_ring->UbrAllocateLocalShm(local_trx_shm, shm_name); if (ret != 0) { + const int saved_errno = errno; + DeallocateResources(); + delete _ub_ring; + _ub_ring = NULL; + _poller_sid = INVALID_SOCKET_ID; + errno = saved_errno; return ret; } - PollerRegisterEvent(PollerSidOp::ADD, EPOLLIN); return 0; } -int UBShmEndpoint::AllocateServerResources(ubring::SHM* remote_trx_shm, ubring::SHM* local_trx_shm) { +int UBShmEndpoint::AllocateServerResources(ubring::SHM *remote_trx_shm, + ubring::SHM *local_trx_shm) { if (BAIDU_UNLIKELY(g_skip_ub_init)) { // For UT return 0; @@ -797,17 +183,43 @@ int UBShmEndpoint::AllocateServerResources(ubring::SHM* remote_trx_shm, ubring:: options.user = this; options.keytable_pool = _socket->_keytable_pool; if (Socket::Create(options, &_poller_sid) < 0) { + const int saved_errno = errno; PLOG(WARNING) << "Fail to create socket for UBRing poller"; + delete _ub_ring; + _ub_ring = NULL; + _poller_sid = INVALID_SOCKET_ID; + errno = saved_errno; return -1; } int ret = _ub_ring->UbrAllocateServerShm(remote_trx_shm, local_trx_shm); if (ret != 0) { + const int saved_errno = errno; + DeallocateResources(); + delete _ub_ring; + _ub_ring = NULL; + _poller_sid = INVALID_SOCKET_ID; + errno = saved_errno; return ret; } - PollerRegisterEvent(PollerSidOp::ADD, EPOLLIN); return ret; } +void UBShmEndpoint::StartReceiveEvents() { + if (g_skip_ub_init || _ub_ring == nullptr || + _poller_sid == INVALID_SOCKET_ID || + _receive_events_started.load(butil::memory_order_acquire)) { + return; + } + // The last TCP parse pins the handshake handler on NOT_ENOUGH_DATA. + // Release that preference before the first shared-memory RPC is parsed. + _socket->set_preferred_index(-1); + if (!_receive_events_started.exchange(true, butil::memory_order_acq_rel)) { + // The poller is level-triggered: bytes queued before registration are + // picked up on its first scan, without needing another peer write. + PollerRegisterEvent(PollerSidOp::ADD, EPOLLIN); + } +} + void UBShmEndpoint::DeallocateResources() { if (!_ub_ring) { return; @@ -825,11 +237,14 @@ void UBShmEndpoint::DeallocateResources() { } void UBShmEndpoint::PollIn(UBShmEndpoint* ep, uint32_t ep_event) { + if (!ep->_receive_events_started.load(butil::memory_order_acquire)) { + return; + } SocketUniquePtr s; if (Socket::Address(ep->_socket_id, &s) < 0) { return; } - auto* ub_transport = static_cast(s->_transport.get()); + UBShmTransport *ub_transport = UBShmTransport::Get(s.get()); CHECK(ep == ub_transport->_ub_ep); InputMessageClosure last_msg; @@ -850,9 +265,10 @@ void UBShmEndpoint::PollIn(UBShmEndpoint* ep, uint32_t ep_event) { if (nr <= 0) { if (0 == nr) { // Set `read_eof' flag and proceed to feed EOF into `Protocol' - // (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"; + // (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) { if (errno == EINTR) { @@ -885,12 +301,11 @@ void UBShmEndpoint::PollOut(UBShmEndpoint* ep, uint32_t ep_event) { if (Socket::Address(ep->_socket_id, &s) < 0) { return; } - auto* ub_transport = static_cast(s->_transport.get()); + UBShmTransport *ub_transport = UBShmTransport::Get(s.get()); CHECK(ep == ub_transport->_ub_ep); if (ep->IsWritable()) { s->WakeAsEpollOut(); } - } int UBShmEndpoint::GlobalInitialize() { @@ -926,7 +341,7 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, std::unique_ptr args(static_cast(p)); auto poller = args->poller; auto running = args->running; - std::unordered_set poller_sids; + std::unordered_set cq_sids; PollerSidOp op; if (poller->init_fn) { @@ -935,17 +350,18 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, while (running->load(std::memory_order_relaxed)) { while (poller->op_queue.Dequeue(op)) { if (op.type == PollerSidOp::ADD) { - poller_sids.emplace(op); + cq_sids.emplace(op); } else if (op.type == PollerSidOp::REMOVE) { - poller_sids.erase(op); + cq_sids.erase(op); + } else if (op.type == PollerSidOp::MOD) { - poller_sids.erase(op); - poller_sids.emplace(op); + cq_sids.erase(op); + cq_sids.emplace(op); } } - for (const auto& poller_sid : poller_sids) { + for (auto cq : cq_sids) { SocketUniquePtr s; - if (Socket::Address(poller_sid.sid, &s) < 0) { + if (Socket::Address(cq.sid, &s) < 0) { continue; } UBShmEndpoint* ep = static_cast(s->user()); @@ -953,12 +369,12 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, continue; } - if (poller_sid.events & EPOLLIN) { - PollIn(ep, poller_sid.events); + if (cq.event & EPOLLIN) { + PollIn(ep, cq.event); } - if (poller_sid.events & EPOLLOUT) { - PollOut(ep, poller_sid.events); + if (cq.event & EPOLLOUT) { + PollOut(ep, cq.event); } } if (poller->callback) { @@ -977,8 +393,8 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, }; for (int i = 0; i < FLAGS_ub_poller_num; ++i) { auto args = new FnArgs{&pollers[i], &running}; - auto attr = FLAGS_ub_disable_bthread ? BTHREAD_ATTR_PTHREAD - : BTHREAD_ATTR_NORMAL; + auto attr = + FLAGS_ub_disable_bthread ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL; attr.tag = tag; bthread_attr_set_name(&attr, "UBPolling"); pollers[i].callback = callback; @@ -1003,8 +419,7 @@ void UBShmEndpoint::PollingModeRelease(bthread_tag_t tag) { } } -void UBShmEndpoint::PollerRegisterEvent(PollerSidOp::OpType op, - uint32_t events) { +void UBShmEndpoint::PollerRegisterEvent(PollerSidOp::OpType op, uint32_t events) { auto index = butil::fmix32(_poller_sid) % FLAGS_ub_poller_num; auto& group = _poller_groups[bthread_self_tag()]; auto& pollers = group.pollers; diff --git a/src/brpc/ubshm/ub_endpoint.h b/src/brpc/ubshm/ub_endpoint.h index 2d7b1f5512..d30535805e 100644 --- a/src/brpc/ubshm/ub_endpoint.h +++ b/src/brpc/ubshm/ub_endpoint.h @@ -20,78 +20,35 @@ #if BRPC_WITH_UBRING -#include -#include -#include -#include -#include -#include "butil/atomicops.h" -#include "butil/iobuf.h" -#include "butil/macros.h" -#include "butil/containers/mpsc_queue.h" +#include "brpc/handshake/ubshm_handshake.h" #include "brpc/socket.h" +#include "brpc/ubshm/shm/shm_def.h" #include "brpc/ubshm/ub_helper.h" #include "brpc/ubshm/ub_ring.h" -#include "brpc/ubshm/shm/shm_def.h" - +#include "butil/atomicops.h" +#include "butil/containers/mpsc_queue.h" +#include "butil/iobuf.h" +#include "butil/macros.h" +#include +#include namespace brpc { class Socket; +class UBShmTransport; +namespace handshake { +class UBShmServerHandshakeAdapter; +} namespace ubring { DECLARE_int32(ub_poller_num); DECLARE_bool(ub_edisp_unsched); DECLARE_bool(ub_disable_bthread); -enum UbrDataFormat { - UBR_DATA_FORMAT_NONE = 0, - UBR_DATA_FORMAT_LEGACY_64 = 1, -}; - -struct HelloFormatExtension { - // The V3 format extension is a fixed-size frame. A different wire size - // requires negotiation through a new hello version. - static const uint16_t WIRE_SIZE = 4; - - uint16_t extension_len; - uint16_t format_id; - - void Serialize(void* data) const; - void Deserialize(const void* data); -}; - -struct HelloMessage { - void Serialize(void* data) const; - void Deserialize(void* data); - std::string toString() const; - - uint16_t msg_len; - uint16_t hello_ver; - uint16_t impl_ver; - uint64_t len; - char shm_name[SHM_MAX_NAME_BUFF_LEN]; -}; - -class UBConnect : public AppConnect { -public: - void StartConnect(const Socket* socket, - void (*done)(int err, void* data), void* data) override; - void StopConnect(Socket*) override; - struct RunGuard { - RunGuard(UBConnect* rc) { this_rc = rc; } - ~RunGuard() { if (this_rc) this_rc->Run(); } - UBConnect* this_rc; - }; - -private: - void Run(); - void (*_done)(int, void*){nullptr}; - void* _data{nullptr}; -}; - class BAIDU_CACHELINE_ALIGNMENT UBShmEndpoint : public SocketUser { -friend class UBConnect; friend class Socket; +friend class ::brpc::UBShmTransport; +friend class ::brpc::handshake::UBShmServerHandshakeAdapter; + public: explicit UBShmEndpoint(Socket* s); ~UBShmEndpoint() override; @@ -104,6 +61,12 @@ friend class Socket; // Reset the endpoint (for next use) void Reset(); + void SetNegotiatedDataFormat(UbrDataFormat format) { + _negotiated_data_format = format; + } + UbrDataFormat negotiated_data_format() const { + return _negotiated_data_format; + } // Cut data from the given IOBuf list and use UBRING to send // Return bytes cut if success, -1 if failed and errno set @@ -112,6 +75,9 @@ friend class Socket; // Whether the endpoint can send more data bool IsWritable() const; + // Start receiving only after the TCP handshake parser has returned. + void StartReceiveEvents(); + void PollerRegisterEpollOut(bool pollin) { uint32_t events = EPOLLOUT | EPOLLET; if (pollin) { @@ -130,9 +96,6 @@ friend class Socket; PollerRegisterEvent(PollerSidOp::REMOVE); } - // Callback when there is new epollin event on TCP fd - static void OnNewDataFromTcp(Socket* m); - // Initialize polling mode static int PollingModeInitialize(bthread_tag_t tag, std::function callback, @@ -141,37 +104,7 @@ friend class Socket; static void PollingModeRelease(bthread_tag_t tag); -#ifdef UNIT_TEST -public: -#else private: -#endif - enum State { - UNINIT = 0x0, - C_ALLOC_SHM = 0x1, - C_HELLO_SEND = 0x2, - C_HELLO_WAIT = 0x3, - C_FORMAT_SEND = 0x4, - C_FORMAT_WAIT = 0x5, - C_MAP_REMOTE_SHM = 0x6, - C_ACK_SEND = 0x7, - S_HELLO_WAIT = 0x11, - S_ALLOC_SHM = 0x12, - S_HELLO_SEND = 0x13, - S_FORMAT_WAIT = 0x14, - S_FORMAT_SEND = 0x15, - S_ACK_WAIT = 0x16, - ESTABLISHED = 0x100, - FALLBACK_TCP = 0x200, - FAILED = 0x300 - }; - - // Process handshake at the client - static void* ProcessHandshakeAtClient(void* arg); - - // Process handshake at the server - static void* ProcessHandshakeAtServer(void* arg); - // Allocate resources // Return 0 if success, -1 if failed and errno set int AllocateClientResources(SHM* local_trx_shm, const char* shm_name); @@ -181,58 +114,32 @@ friend class Socket; // Release resources void DeallocateResources(); - // Read at most len bytes from fd in _socket to data - // wait for _read_butex if encounter EAGAIN - // return -1 if encounter other errno (including EOF) - int ReadFromFd(void* data, size_t len); - - - // Write at most len bytes from data to fd in _socket - // wait for _epollout_butex if encounter EAGAIN - // return -1 if encounter other errno - int WriteToFd(void* data, size_t len); - - // Poll inbound and outbound UBRing events. + // Poll CQ and get the work completion static void PollIn(UBShmEndpoint* ep, uint32_t ep_event); static void PollOut(UBShmEndpoint* ep, uint32_t ep_event); - // Try to read data on TCP fd in _socket - inline void TryReadOnTcp(); - // Not owner Socket* _socket; SocketId _socket_id; - - State _state; UbrDataFormat _negotiated_data_format{UBR_DATA_FORMAT_NONE}; // ub resource ubring::UBRing* _ub_ring{nullptr}; - // Synthetic SocketId registered with the UBRing poller. SocketId _poller_sid; - // butex for inform read events on TCP fd during handshake - butil::atomic *_read_butex; - DISALLOW_COPY_AND_ASSIGN(UBShmEndpoint); struct PollerSidOp { - enum OpType { - ADD, - REMOVE, - MOD - }; + enum OpType { ADD, REMOVE, MOD }; SocketId sid; - uint32_t events; + uint32_t event; OpType type; }; struct PollerSidOpHash { - std::size_t operator()(const PollerSidOp& op) const { - return op.sid; - } + std::size_t operator()(const PollerSidOp& op) const { return op.sid; } }; struct PollerSidOpEqual { @@ -244,8 +151,7 @@ friend class Socket; // Poller instance struct BAIDU_CACHELINE_ALIGNMENT Poller { bthread_t tid{INVALID_BTHREAD}; - butil::MPSCQueue< - PollerSidOp, butil::ObjectPoolAllocator> op_queue; + butil::MPSCQueue> op_queue; // Callback used for io_uring/spdk etc std::function callback; // Init and Destroy function @@ -260,8 +166,9 @@ friend class Socket; }; static std::vector _poller_groups; - void PollerRegisterEvent(PollerSidOp::OpType op, - uint32_t events = EPOLLET); + butil::atomic _receive_events_started{false}; + + void PollerRegisterEvent(PollerSidOp::OpType op, uint32_t events = EPOLLET); }; } // namespace ubring diff --git a/src/brpc/ubshm_transport.cpp b/src/brpc/ubshm_transport.cpp index 45f6a61fa6..9b0a0203e0 100644 --- a/src/brpc/ubshm_transport.cpp +++ b/src/brpc/ubshm_transport.cpp @@ -17,10 +17,15 @@ #if BRPC_WITH_UBRING -#include "brpc/ubshm_transport.h" -#include "brpc/tcp_transport.h" +#include + +#include "brpc/adapter_transport.h" +#include "brpc/errno.pb.h" +#include "brpc/ubshm/common/common.h" #include "brpc/ubshm/ub_endpoint.h" #include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/ubr_trx.h" +#include "brpc/ubshm_transport.h" namespace brpc { DECLARE_bool(usercode_in_coroutine); @@ -28,23 +33,26 @@ DECLARE_bool(usercode_in_pthread); extern SocketVarsCollector *g_vars; +UBShmTransport *UBShmTransport::Get(const Socket *socket) { + const AdapterTransport *adapter = AdapterTransport::Get(socket); + Transport *transport = adapter->high_speed_transport(); + CHECK(transport != NULL); + return static_cast(transport); +} + void UBShmTransport::Init(Socket *socket, const SocketOptions &options) { CHECK(_ub_ep == nullptr); - if (options.socket_mode == SOCKET_MODE_UBRING) { - _ub_ep = new ubring::UBShmEndpoint(socket); - _ub_state = UB_UNKNOWN; - } else { - _ub_state = UB_OFF; - socket->_socket_mode = SOCKET_MODE_TCP; - } _socket = socket; _default_connect = options.app_connect; - _on_edge_trigger = options.on_edge_triggered_events; - if (options.need_on_edge_trigger && _on_edge_trigger == nullptr) { - _on_edge_trigger = ubring::UBShmEndpoint::OnNewDataFromTcp; + _on_edge_trigger = nullptr; + _ub_state = UB_UNKNOWN; + _ub_ep = new (std::nothrow) ubring::UBShmEndpoint(socket); + if (!_ub_ep) { + const int saved_errno = errno != 0 ? errno : ENOMEM; + errno = saved_errno; + PLOG(WARNING) << "Fail to create UBShmEndpoint, disable UBSHM upgrade"; + _ub_state = UB_OFF; } - _tcp_transport = std::make_shared(); - _tcp_transport->Init(socket, options); } void UBShmTransport::Release() { @@ -64,26 +72,53 @@ int UBShmTransport::Reset(int32_t expected_nref) { } std::shared_ptr UBShmTransport::Connect() { - if (_default_connect == nullptr) { - return std::make_shared(); - } return _default_connect; } -int UBShmTransport::CutFromIOBuf(butil::IOBuf *buf) { - if (_ub_ep && _ub_state != UB_OFF) { - butil::IOBuf *data_arr[1] = {buf}; - return _ub_ep->CutFromIOBufList(data_arr, 1); - } else { - return _tcp_transport->CutFromIOBuf(buf); +void UBShmTransport::SetHighSpeedAvailable(bool available) { + _ub_state = available ? UB_ON : UB_OFF; +} + +int UBShmTransport::PrepareUpgradeResources(ubring::SHM *local_trx_shm, + const char *shm_name) { + return _ub_ep->AllocateClientResources(local_trx_shm, shm_name); +} + +int UBShmTransport::NegotiateUpgradeResources(ubring::SHM *local_trx_shm, + const char *shm_name) { + return _ub_ep->_ub_ring->UbrMapRemoteShm(local_trx_shm, shm_name); +} + +int UBShmTransport::PrepareServerUpgradeResources(ubring::SHM *remote_trx_shm, + ubring::SHM *local_trx_shm) { + return _ub_ep->AllocateServerResources(remote_trx_shm, local_trx_shm); +} + +void UBShmTransport::ActivateUpgrade() { + SetHighSpeedAvailable(true); +} + +void UBShmTransport::DeactivateUpgrade() { + SetHighSpeedAvailable(false); + if (_ub_ep != nullptr) { + _ub_ep->Reset(); } } -ssize_t UBShmTransport::CutFromIOBufList(butil::IOBuf **buf, size_t ndata) { - if (_ub_ep && _ub_state != UB_OFF) { - return _ub_ep->CutFromIOBufList(buf, ndata); +void UBShmTransport::FinishUpgrade() { + if (_ub_ep != NULL && _ub_ep->_ub_ring != NULL) { + _ub_ep->_ub_ring->UbrUnlinkLocalShm(); } - return _tcp_transport->CutFromIOBufList(buf, ndata); +} + +int UBShmTransport::CutFromIOBuf(butil::IOBuf *buf) { + butil::IOBuf *data[1] = {buf}; + return static_cast(CutFromIOBufList(data, 1)); +} + +ssize_t UBShmTransport::CutFromIOBufList(butil::IOBuf **buf, size_t ndata) { + CHECK(_ub_ep != NULL); + return _ub_ep->CutFromIOBufList(buf, ndata); } int UBShmTransport::WaitEpollOut(butil::atomic *_epollout_butex, @@ -94,14 +129,13 @@ int UBShmTransport::WaitEpollOut(butil::atomic *_epollout_butex, if (!_ub_ep->IsWritable()) { g_vars->nwaitepollout << 1; _ub_ep->PollerRegisterEpollOut(pollin); - const int wait_rc = bthread::butex_wait( - _epollout_butex, expected_val, &duetime); + const int wait_rc = + bthread::butex_wait(_epollout_butex, expected_val, &duetime); if (wait_rc < 0) { if (errno != EAGAIN && errno != ETIMEDOUT) { const int saved_errno = errno; PLOG(WARNING) << "Fail to wait ub window of " << _socket; - _socket->SetFailed(saved_errno, - "Fail to wait ub window of %s: %s", + _socket->SetFailed(saved_errno, "Fail to wait ub window of %s: %s", _socket->description().c_str(), berror(saved_errno)); } @@ -115,8 +149,6 @@ int UBShmTransport::WaitEpollOut(butil::atomic *_epollout_butex, } _ub_ep->PollerUnRegisterEpollOut(pollin); } - } else { - return _tcp_transport->WaitEpollOut(_epollout_butex, pollin, duetime); } return 0; } @@ -156,15 +188,16 @@ void UBShmTransport::QueueMessage(InputMessageClosure& input_msg, // TODO(gejun): Join threads. bthread_t th; - bthread_attr_t tmp = (FLAGS_usercode_in_pthread ? - BTHREAD_ATTR_PTHREAD : - BTHREAD_ATTR_NORMAL) | BTHREAD_NOSIGNAL; + 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, "ProcessInputMessage"); - if (!FLAGS_usercode_in_coroutine && bthread_start_background( - &th, &tmp, ProcessInputMessage, to_run_msg) == 0) { + if (!FLAGS_usercode_in_coroutine && + bthread_start_background(&th, &tmp, ProcessInputMessage, to_run_msg) == + 0) { ++*num_bthread_created; } else { ProcessInputMessage(to_run_msg); @@ -179,7 +212,8 @@ int UBShmTransport::ContextInitOrDie(bool serverOrNot, const void* _options) { return -1; } ubring::GlobalUBInitializeOrDie(); - if (!ubring::InitPollingModeWithTag(static_cast(_options)->bthread_tag)) { + if (!ubring::InitPollingModeWithTag( + static_cast(_options)->bthread_tag)) { return -1; } } else { @@ -202,8 +236,7 @@ bool UBShmTransport::OptionsAvailableForUB(const ChannelOptions* opt) { return false; } if (!ubring::SupportedByUB(opt->protocol.name())) { - LOG(WARNING) << "Cannot use " << opt->protocol.name() - << " over UB"; + LOG(WARNING) << "Cannot use " << opt->protocol.name() << " over UB"; return false; } return true; diff --git a/src/brpc/ubshm_transport.h b/src/brpc/ubshm_transport.h index b3d1e7c518..b8d840595e 100644 --- a/src/brpc/ubshm_transport.h +++ b/src/brpc/ubshm_transport.h @@ -21,12 +21,14 @@ #include "brpc/socket.h" #include "brpc/channel.h" #include "brpc/transport.h" +#include "brpc/ubshm/shm/shm_def.h" namespace brpc { +class AdapterTransport; class UBShmTransport : public Transport { friend class TransportFactory; + friend class AdapterTransport; friend class ubring::UBShmEndpoint; -friend class ubring::UBConnect; public: void Init(Socket* socket, const SocketOptions& options) override; void Release() override; @@ -34,7 +36,8 @@ friend class ubring::UBConnect; std::shared_ptr Connect() override; int CutFromIOBuf(butil::IOBuf* buf) override; ssize_t CutFromIOBufList(butil::IOBuf** buf, size_t ndata) override; - int WaitEpollOut(butil::atomic* _epollout_butex, bool pollin, const timespec duetime) override; + int WaitEpollOut(butil::atomic* epollout_butex, + bool pollin, timespec duetime) override; void ProcessEvent(bthread_attr_t attr) override; void QueueMessage(InputMessageClosure& inputMsg, int* num_bthread_created, bool last_msg) override; void Debug(std::ostream &os) override; @@ -42,8 +45,22 @@ friend class ubring::UBConnect; CHECK(_ub_ep != nullptr); return _ub_ep; } + static UBShmTransport* Get(const Socket* socket); static int ContextInitOrDie(bool serverOrNot, const void* _options); + int PrepareUpgradeResources(ubring::SHM* local_trx_shm, + const char* shm_name); + int NegotiateUpgradeResources(ubring::SHM* local_trx_shm, + const char* shm_name); + int PrepareServerUpgradeResources(ubring::SHM* remote_trx_shm, + ubring::SHM* local_trx_shm); + void ActivateUpgrade(); + void DeactivateUpgrade(); + void FinishUpgrade(); + bool UpgradeActive() const { return _ub_state == UB_ON; } + bool UpgradeReady() const { return _ub_ep != nullptr; } private: + void SetHighSpeedAvailable(bool available); + static bool OptionsAvailableForUB(const ChannelOptions* opt); static bool OptionsAvailableOverUB(const ServerOptions* opt); private: @@ -57,7 +74,6 @@ friend class ubring::UBConnect; ubring::UBShmEndpoint* _ub_ep = nullptr; // Should use UB or not UBState _ub_state; - std::shared_ptr _tcp_transport; }; } // namespace brpc #endif // BRPC_WITH_UBRING diff --git a/test/brpc_rdma_unittest.cpp b/test/brpc_rdma_unittest.cpp index 6ec54ed90b..ed0ae7262e 100644 --- a/test/brpc_rdma_unittest.cpp +++ b/test/brpc_rdma_unittest.cpp @@ -15,39 +15,38 @@ // specific language governing permissions and limitations // under the License. - +#include +#include +#include #include #include -#include -#include + #if BRPC_WITH_RDMA -#include -#include -#include -#include -#include -#include "butil/endpoint.h" -#include "butil/fd_guard.h" -#include "butil/iobuf.h" -#include "butil/sys_byteorder.h" -#include "butil/time.h" -#include "butil/files/temp_file.h" #include "brpc/acceptor.h" +#include "brpc/adapter_transport.h" #include "brpc/channel.h" #include "brpc/controller.h" -#include "brpc/server.h" -#include "brpc/socket.h" #include "brpc/errno.pb.h" +#include "brpc/handshake/rdma_handshake.h" +#include "brpc/handshake/rdma_handshake_constants.h" #include "brpc/parallel_channel.h" -#include "brpc/selective_channel.h" -#include "brpc/rdma_transport.h" +#include "brpc/policy/baidu_rpc_meta.pb.h" +#include "brpc/policy/baidu_rpc_protocol.h" #include "brpc/rdma/block_pool.h" #include "brpc/rdma/rdma_endpoint.h" -#include "brpc/rdma/rdma_handshake.h" -#include "brpc/rdma/rdma_handshake_constants.h" -#include "brpc/rdma/rdma_handshake.pb.h" #include "brpc/rdma/rdma_helper.h" +#include "brpc/rdma_handshake.pb.h" +#include "brpc/rdma_transport.h" +#include "brpc/selective_channel.h" +#include "brpc/server.h" +#include "brpc/socket.h" +#include "butil/endpoint.h" +#include "butil/fd_guard.h" +#include "butil/files/temp_file.h" +#include "butil/iobuf.h" +#include "butil/sys_byteorder.h" #include "echo.pb.h" +#include static const int PORT = 8713; @@ -62,7 +61,7 @@ DEFINE_bool(rdma_test_enable, false, "Enable tests requring rdma runtime."); namespace rdma { // HELLO_V2_VERSION / IMPL_V2_VERSION come from -// brpc/rdma/rdma_handshake_constants.h (shared wire constants). +// brpc/handshake/rdma_handshake_constants.h (shared wire constants). DECLARE_bool(rdma_trace_verbose); DECLARE_int32(rdma_memory_pool_max_regions); @@ -124,7 +123,7 @@ static bool WaitUntil(const std::function& pred, // whole buffer reaching the peer has to loop. static bool WriteAll(int fd, const void* buf, size_t len) { const uint8_t* p = (const uint8_t*)buf; - for (size_t done = 0; done < len; ) { + for (size_t done = 0; done < len;) { ssize_t n = write(fd, p + done, len - done); if (n < 0) { if (errno == EINTR) { @@ -141,7 +140,7 @@ static bool WriteAll(int fd, const void* buf, size_t len) { // and on EOF before the whole buffer arrived. static bool ReadAll(int fd, void* buf, size_t len) { uint8_t* p = (uint8_t*)buf; - for (size_t done = 0; done < len; ) { + for (size_t done = 0; done < len;) { ssize_t n = read(fd, p + done, len - done); if (n < 0) { if (errno == EINTR) { @@ -178,8 +177,7 @@ static void ConnectToServer(butil::fd_guard* sockfd) { class MyEchoService : public ::test::EchoService { void Echo(google::protobuf::RpcController* cntl_base, - const ::test::EchoRequest* req, - ::test::EchoResponse* res, + const ::test::EchoRequest* req, ::test::EchoResponse* res, google::protobuf::Closure* done) { Controller* cntl = static_cast(cntl_base); ClosureGuard done_guard(done); @@ -218,11 +216,12 @@ class RdmaTest : public ::testing::Test { _naming_url = std::string("File://") + _server_list.fname(); _server.AddService(&_svc, SERVER_DOESNT_OWN_SERVICE); } - ~RdmaTest() { } + ~RdmaTest() {} - virtual void SetUp() { } + virtual void SetUp() {} virtual void TearDown() { + _wait_socket.reset(); rdma::DumpMemoryPoolInfo(std::cout); } @@ -239,6 +238,7 @@ class RdmaTest : public ::testing::Test { void StopServer() { _server.Stop(0); + _wait_socket.reset(); _server.Join(); } @@ -255,19 +255,29 @@ class RdmaTest : public ::testing::Test { return nullptr; } - // Accepting the connection happens in the server threads, so poll for it - // rather than sleeping. Returns nullptr if it never showed up. + // Server-side connection creation and teardown are asynchronous. Socket* WaitForServerSocket() { - Socket* s = nullptr; - WaitUntil([this, &s] { return (s = GetSocketFromServer(0)) != nullptr; }); - return s; + _wait_socket.reset(); + WaitUntil([this] { + std::vector sids; + _server._am->ListConnections(&sids); + return !sids.empty() && + Socket::Address(sids[0], &_wait_socket) == 0; + }); + return _wait_socket.get(); } - // Ditto for the connection going away. bool WaitForServerSocketGone() { + // Acceptor removes the connection in BeforeRecycle; release our pin + // before waiting for removal, and do not use the raw pointer afterwards. + _wait_socket.reset(); return WaitUntil([this] { return GetSocketFromServer(0) == nullptr; }); } + SocketUniquePtr _wait_socket; + + void VerifyCompleteFallbackRPC(); + butil::TempFile _server_list; std::string _naming_url; @@ -275,41 +285,23 @@ class RdmaTest : public ::testing::Test { MyEchoService _svc; }; -// Shorthand for the RDMA transport behind a Socket, which every endpoint state -// check below has to go through. +// Shorthand for the RDMA transport behind a Socket. static RdmaTransport* RdmaTransportOf(Socket* s) { - return static_cast(s->_transport.get()); -} -static RdmaTransport* RdmaTransportOf(const SocketUniquePtr& s) { - return RdmaTransportOf(s.get()); + return RdmaTransport::Get(s); } -// Polls until the endpoint reaches `expected` and returns the last state seen, -// so that ASSERT_RDMA_STATE() reports what the endpoint actually settled on. -static rdma::RdmaEndpoint::State WaitForRdmaState( - RdmaTransport* transport, rdma::RdmaEndpoint::State expected) { - rdma::RdmaEndpoint::State state = transport->_rdma_ep->_state; - WaitUntil([transport, expected, &state] { - state = transport->_rdma_ep->_state; - return state == expected; +static int WaitForHandshakePhase(Socket* s, handshake::Phase expected) { + int phase = AdapterTransport::Get(s)->handshake_phase(); + WaitUntil([s, expected, &phase] { + phase = AdapterTransport::Get(s)->handshake_phase(); + return phase == expected; }); - return state; + return phase; } -// Waits for `transport` to reach `expected`, failing the test if it does not. -#define ASSERT_RDMA_STATE(expected, transport) \ - ASSERT_EQ(expected, WaitForRdmaState(transport, expected)) +#define ASSERT_HANDSHAKE_PHASE(expected, socket) \ + ASSERT_EQ(expected, WaitForHandshakePhase(socket, expected)) -// Polls until the fd stream of `s` holds exactly `size` bytes. Tests asserting -// that a state did NOT change need this: waiting for the state itself would -// return before the peer had read anything at all. -static bool WaitForFdReadBuf(Socket* s, size_t size) { - return WaitUntil([s, size] { - return s->fd_input_processor().read_buf().size() == size; - }); -} - -// 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; @@ -330,8 +322,7 @@ static void MakeV2ClientHello(uint8_t (&data)[rdma::HELLO_V2_MSG_LEN_MIN]) { // so every TEST_P below is automatically executed once per supported // version. Add a new version to INSTANTIATE_TEST_SUITE_P at the bottom // of this file and these RPC tests will gain coverage for free. -class RdmaRpcTest : public RdmaTest, - public ::testing::WithParamInterface { +class RdmaRpcTest : public RdmaTest, public ::testing::WithParamInterface { protected: void SetUp() override { RdmaTest::SetUp(); @@ -355,9 +346,10 @@ TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { SocketUniquePtr main_socket; ASSERT_EQ(0, Socket::Address(main_sid, &main_socket)); - RdmaTransport* transport = - static_cast(main_socket->_transport.get()); + RdmaTransport* transport = RdmaTransportOf(main_socket.get()); + ASSERT_NE(nullptr, transport); rdma::RdmaEndpoint* ep = transport->_rdma_ep; + ASSERT_NE(nullptr, ep); SocketOptions cq_options; cq_options.user = ep; @@ -384,11 +376,18 @@ TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { TEST_F(RdmaTest, client_close_before_hello_send) { StartServer(); - butil::fd_guard sockfd; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + 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))); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -398,23 +397,22 @@ TEST_F(RdmaTest, client_close_before_hello_send) { TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { StartServer(); - butil::fd_guard sockfd; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + 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))); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(data, "PRPC", 4); // send as normal baidu_std protocol - ASSERT_TRUE(WriteAll(sockfd, data, 4)); - // Wait for the bytes to show up in the fd stream (baidu_std wants 12B of - // header, so they stay buffered). Waiting on the state instead would prove - // nothing: it is already UNINIT before the server has read anything. - ASSERT_TRUE(WaitForFdReadBuf(s, 4)); - // A non-RDMA magic makes ParseRdmaHandshake return TRY_OTHERS and hand the - // bytes to other protocols; it does not touch the endpoint state, so it - // stays UNINIT (the old blocking handshake used to set FALLBACK_TCP here). - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_EQ(4, write(sockfd, data, 4)); StopServer(); } @@ -422,49 +420,55 @@ TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { TEST_F(RdmaTest, client_close_during_hello_send) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[8]; - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RD", 2); - ASSERT_TRUE(WriteAll(sockfd1, data, 2)); // break in magic str - // Fewer than 4 magic bytes: ParseRdmaHandshake can't tell yet, returns - // NOT_ENOUGH_DATA and leaves the endpoint UNINIT (the old blocking - // handshake used to set S_HELLO_WAIT before reading the magic). Wait for - // the bytes to be buffered, the state alone would prove nothing. - ASSERT_TRUE(WaitForFdReadBuf(s, 2)); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_EQ(2, write(sockfd1, data, 2)); // break in magic str + // Close and wait for teardown. Parser phase assertions for incomplete + // magic belong to the synchronous common-handshake regression test. sockfd1.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RDMA", 4); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // break after magic str - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(4, write(sockfd2, data, 4)); // break after magic str + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); - butil::fd_guard sockfd3; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd3)); + butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd3 >= 0); + ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); // Send the 4B magic plus a valid msg_len (=40) but no body, so the server // recognizes an RDMA v2 hello and waits for the remaining bytes. (A zero // msg_len would now be rejected up-front as a protocol error.) memcpy(data, "RDMA", 4); uint16_t v2_len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); memcpy(data + 4, &v2_len, sizeof(v2_len)); - ASSERT_TRUE(WriteAll(sockfd3, data, 6)); // magic + msg_len, body missing - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(6, write(sockfd3, data, 6)); // magic + msg_len, body missing + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd3.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -474,34 +478,46 @@ TEST_F(RdmaTest, client_close_during_hello_send) { TEST_F(RdmaTest, client_hello_msg_invalid_len) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RDMA", 4); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); memset(data + 4, 0, 36); - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); // Write invalid length. - ASSERT_TRUE(WaitForServerSocketGone()); - - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); // Write invalid length. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RDMA", 4); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); uint16_t len = butil::HostToNet16(35); memcpy(data + 4, &len, sizeof(len)); memset(data + 6, 0, 34); - ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); // write invalid length - ASSERT_TRUE(WaitForServerSocketGone()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); // write invalid length + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -509,19 +525,26 @@ TEST_F(RdmaTest, client_hello_msg_invalid_len) { TEST_F(RdmaTest, client_hello_msg_invalid_version) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); uint16_t ver = butil::HostToNet16(1); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RDMA", 4); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); memcpy(data + 4, &len, 2); memset(data + 6, 0, 34); memcpy(data + 6, &ver, 2); // hello_ver == 1, impl_ver == 0 @@ -532,35 +555,46 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { // hello_ver). Now that Step 1 enforces a HELLO_V2_MSG_LEN_MAX upper bound, // such an oversized msg_len would be rejected before reaching the // version check, breaking the intent of this UT. - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); uint32_t flags = 0; - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); - - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); memcpy(data, "RDMA", 4); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); memcpy(data + 8, &ver, 2); // hello_ver == 0, impl_ver == 1 - // See comment above on `WriteAll(sockfd1, data + 4, 36)` for why we + // See comment above on `write(sockfd1, data + 4, 36)` for why we // write from data + 4 instead of data. - ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd2.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -568,6 +602,10 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; uint32_t flags = butil::HostToNet32(0); rdma::v2_wire::HelloMessage msg{}; @@ -581,60 +619,81 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { msg.block_size = 8192; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); msg.sq_size = 16; msg.rq_size = 10; msg.block_size = 8192; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd2.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); msg.sq_size = 16; msg.rq_size = 16; msg.block_size = 1000; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd3; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd3)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd3, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd3, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_TRUE(WriteAll(sockfd3, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd3 >= 0); + ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd3, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd3, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd3, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd3.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -642,19 +701,37 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { TEST_F(RdmaTest, client_close_after_qp_build) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - MakeV2ClientHello(data); + 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); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd1, data, sizeof(data))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(40, write(sockfd1, data, 40)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -662,24 +739,45 @@ TEST_F(RdmaTest, client_close_after_qp_build) { TEST_F(RdmaTest, client_close_during_ack_send) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - MakeV2ClientHello(data); + 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); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); uint32_t flags = butil::HostToNet32(1); - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -687,40 +785,68 @@ TEST_F(RdmaTest, client_close_during_ack_send) { TEST_F(RdmaTest, client_close_after_ack_send) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - MakeV2ClientHello(data); + 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); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); uint32_t flags = butil::HostToNet32(0); - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); - - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); flags = butil::HostToNet32(1); - ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - sockfd2.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd2); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -728,48 +854,71 @@ TEST_F(RdmaTest, client_close_after_ack_send) { TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { StartServer(); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - MakeV2ClientHello(data); + 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); - butil::fd_guard sockfd1; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); uint32_t flags = butil::HostToNet32(0); - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - // 4 more bytes on a fd that fell back to TCP are not a protocol baidu_std - // knows, so the connection is dropped. - ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_TRUE(WaitForServerSocketGone()); - - butil::fd_guard sockfd2; - ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); - s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); - ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_HELLO_WAIT, RdmaTransportOf(s)); - ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); flags = butil::HostToNet32(1); - ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - // Once RDMA is on the fd carries no RPC data at all, so this is an error. - ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_TRUE(WaitForServerSocketGone()); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } -// 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) { ASSERT_NO_FATAL_FAILURE(ConnectToServer(sockfd)); @@ -782,9 +931,6 @@ static void HandshakeUntilAckWait(butil::fd_guard* sockfd) { ASSERT_TRUE(ReadAll(*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(); @@ -793,7 +939,7 @@ TEST_F(RdmaTest, server_accepts_data_pipelined_behind_fallback_ack) { Socket* s = WaitForServerSocket(); ASSERT_TRUE(s != nullptr); auto* transport = RdmaTransportOf(s); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); // 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. @@ -803,14 +949,10 @@ TEST_F(RdmaTest, server_accepts_data_pipelined_behind_fallback_ack) { memcpy(ack_and_data + rdma::HELLO_ACK_LEN, "PRPC", 4); ASSERT_TRUE(WriteAll(sockfd, ack_and_data, sizeof(ack_and_data))); - // 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. Note that baidu_std gets them a moment - // after the handshake gave up the stream, hence the wait. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, transport); + // The complete-RPC test below verifies parser handoff. + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, transport->_rdma_state); ASSERT_TRUE(GetSocketFromServer(0) != nullptr); - ASSERT_TRUE(WaitForFdReadBuf(s, 4)); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -818,8 +960,6 @@ TEST_F(RdmaTest, server_accepts_data_pipelined_behind_fallback_ack) { 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(); @@ -827,8 +967,7 @@ TEST_F(RdmaTest, server_rejects_data_pipelined_behind_rdma_ack) { ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); Socket* s = WaitForServerSocket(); ASSERT_TRUE(s != nullptr); - auto* transport = RdmaTransportOf(s); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); uint8_t ack_and_data[rdma::HELLO_ACK_LEN + 4]; const uint32_t flags = butil::HostToNet32(rdma::HELLO_ACK_RDMA_OK); @@ -843,7 +982,6 @@ TEST_F(RdmaTest, server_rejects_data_pipelined_behind_rdma_ack) { 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(); @@ -852,24 +990,96 @@ TEST_F(RdmaTest, server_stops_parsing_tcp_fd_once_rdma_is_on) { Socket* s = WaitForServerSocket(); ASSERT_TRUE(s != nullptr); auto* transport = RdmaTransportOf(s); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); // 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_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); ASSERT_EQ(RdmaTransport::RDMA_ON, transport->_rdma_state); ASSERT_TRUE(GetSocketFromServer(0) != nullptr); ASSERT_TRUE(WriteAll(sockfd, "PRPC", 4)); + ASSERT_TRUE(WaitUntil([s] { return s->Failed(); })); + // OnFailed publishes error details under this mutex. Keep the Socket pin + // until the event callback has reported the protocol error. + ASSERT_TRUE(WaitUntil([s] { + BAIDU_SCOPED_LOCK(s->_id_wait_list_mutex); + return s->_error_code == EPROTO; + })); ASSERT_TRUE(WaitForServerSocketGone()); 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. +void RdmaTest::VerifyCompleteFallbackRPC() { + StartServer(); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + + test::EchoRequest request; + request.set_message("fallback"); + request.set_code(42); + Controller cntl; + butil::IOBuf body; + policy::SerializeRpcRequest(&body, &cntl, &request); + butil::IOBuf packet; + policy::PackRpcRequest(&packet, nullptr, 123, + test::EchoService::descriptor()->FindMethodByName("Echo"), + &cntl, body, nullptr); + ASSERT_FALSE(cntl.Failed()); + + const uint32_t disabled_ack = 0; + std::string bytes(reinterpret_cast(&disabled_ack), + sizeof(disabled_ack)); + bytes += packet.to_string(); + ASSERT_TRUE(WriteAll(sockfd, bytes.data(), bytes.size())); + + uint8_t header[12]; + ASSERT_TRUE(ReadAll(sockfd, header, sizeof(header))); + ASSERT_EQ(0, memcmp(header, "PRPC", 4)); + uint32_t total_size; + uint32_t meta_size; + memcpy(&total_size, header + 4, sizeof(total_size)); + memcpy(&meta_size, header + 8, sizeof(meta_size)); + total_size = butil::NetToHost32(total_size); + meta_size = butil::NetToHost32(meta_size); + ASSERT_LE(meta_size, total_size); + ASSERT_GT(total_size, 0u); + ASSERT_LT(total_size, 65536u); + std::string response(total_size, '\0'); + ASSERT_TRUE(ReadAll(sockfd, &response[0], response.size())); + policy::RpcMeta meta; + ASSERT_TRUE(meta.ParseFromArray(response.data(), meta_size)); + ASSERT_EQ(123u, meta.correlation_id()); + ASSERT_EQ(0, meta.response().error_code()); + test::EchoResponse echo; + ASSERT_TRUE(echo.ParseFromArray(response.data() + meta_size, + total_size - meta_size)); + ASSERT_EQ("MyEchoService", echo.message()); + ASSERT_EQ(1, echo.code_list_size()); + ASSERT_EQ(42, echo.code_list(0)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); + auto* transport = RdmaTransportOf(s); + EXPECT_EQ(RdmaTransport::RDMA_OFF, transport->_rdma_state); + EXPECT_EQ(nullptr, transport->_rdma_ep->_resource); + EXPECT_EQ(INVALID_SOCKET_ID, transport->_rdma_ep->_cq_sid); + EXPECT_EQ(0, transport->_rdma_ep->_remote_rq_window_size.load( + butil::memory_order_relaxed)); + ASSERT_FALSE(s->Failed()); + sockfd.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); + StopServer(); +} + +TEST_F(RdmaTest, server_serves_complete_rpc_after_fallback_ack) { + ASSERT_NO_FATAL_FAILURE(VerifyCompleteFallbackRPC()); +} + TEST_F(RdmaTest, server_parses_qp_stream_after_rdma_is_on) { StartServer(); @@ -878,11 +1088,11 @@ TEST_F(RdmaTest, server_parses_qp_stream_after_rdma_is_on) { Socket* s = WaitForServerSocket(); ASSERT_TRUE(s != nullptr); auto* transport = RdmaTransportOf(s); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); const uint32_t flags = butil::HostToNet32(rdma::HELLO_ACK_RDMA_OK); ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, transport); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); InputMessengerProcessor& qp_stream = transport->_rdma_ep->_input_processor; ASSERT_TRUE(qp_stream.read_buf().empty()); @@ -899,15 +1109,12 @@ TEST_F(RdmaTest, server_parses_qp_stream_after_rdma_is_on) { 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_EQ(handshake::ESTABLISHED, AdapterTransport::Get(s)->handshake_phase()); 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(); @@ -915,11 +1122,10 @@ TEST_F(RdmaTest, server_declines_handshake_bytes_after_fallback) { ASSERT_NO_FATAL_FAILURE(HandshakeUntilAckWait(&sockfd)); Socket* s = WaitForServerSocket(); ASSERT_TRUE(s != nullptr); - auto* transport = RdmaTransportOf(s); const uint32_t flags = butil::HostToNet32(0); ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, transport); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); // 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 @@ -952,13 +1158,9 @@ TEST_F(RdmaTest, fd_and_qp_input_streams_are_separate) { 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. + // Exercise close with an incomplete prefix without inspecting the IO + // thread's buffer or treating a read timestamp as parser completion. ASSERT_TRUE(WriteAll(sockfd, "RD", 2)); - ASSERT_TRUE(WaitForFdReadBuf(s, 2)); - 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); ASSERT_TRUE(WaitForServerSocketGone()); @@ -985,9 +1187,10 @@ TEST_F(RdmaTest, server_miss_before_hello_send) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); @@ -1015,16 +1218,19 @@ TEST_F(RdmaTest, server_close_before_hello_send) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); close(acc_fd); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1049,18 +1255,18 @@ TEST_F(RdmaTest, server_miss_during_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_TRUE(WriteAll(acc_fd, "RD", 2)); - // Half a magic is not enough to decide anything, so the client stays stuck - // in the handshake read and the RPC runs into its timeout. Joining below - // waits for exactly that, no sleeping needed. + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(2, write(acc_fd, "RD", 2)); + usleep(100000); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1085,19 +1291,21 @@ TEST_F(RdmaTest, server_close_during_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - // Half a magic and then EOF. TCP keeps the order, so the client always sees - // the two bytes first and then the close, which is what this test is about. - ASSERT_TRUE(WriteAll(acc_fd, "RD", 2)); - acc_fd.reset(-1); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(2, write(acc_fd, "RD", 2)); + usleep(100000); + close(acc_fd); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1122,16 +1330,19 @@ TEST_F(RdmaTest, server_hello_invalid_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_EQ(4, write(acc_fd, "ABCD", 4)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); @@ -1156,16 +1367,21 @@ TEST_F(RdmaTest, server_miss_during_hello_msg) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); - ASSERT_EQ(2, write(acc_fd, "00", 2)); + const uint16_t msg_len = butil::HostToNet16( + static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(static_cast(sizeof(msg_len)), + write(acc_fd, &msg_len, sizeof(msg_len))); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1190,18 +1406,24 @@ TEST_F(RdmaTest, server_close_during_hello_msg) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); - ASSERT_EQ(2, write(acc_fd, "00", 2)); + const uint16_t msg_len = butil::HostToNet16( + static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(static_cast(sizeof(msg_len)), + write(acc_fd, &msg_len, sizeof(msg_len))); close(acc_fd); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1226,20 +1448,24 @@ TEST_F(RdmaTest, server_hello_invalid_msg_len) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); memcpy(data, "RDMA", 4); uint16_t len = butil::HostToNet16(35); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); @@ -1264,20 +1490,25 @@ TEST_F(RdmaTest, server_hello_invalid_version) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); memcpy(data, "RDMA", 4); uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); ASSERT_EQ(4, read(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(0, butil::NetToHost32(*tmp)); @@ -1305,14 +1536,16 @@ TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1325,9 +1558,12 @@ TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); ASSERT_EQ(4, read(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(0, butil::NetToHost32(*tmp)); @@ -1355,14 +1591,16 @@ TEST_F(RdmaTest, server_miss_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1375,9 +1613,12 @@ TEST_F(RdmaTest, server_miss_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); ASSERT_EQ(4, read(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(1, butil::NetToHost32(*tmp)); @@ -1405,14 +1646,16 @@ TEST_F(RdmaTest, server_close_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1425,9 +1668,12 @@ TEST_F(RdmaTest, server_close_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); ASSERT_EQ(4, read(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(1, butil::NetToHost32(*tmp)); @@ -1456,14 +1702,16 @@ TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1476,16 +1724,19 @@ TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); } - TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { butil::fd_guard sockfd(butil::tcp_listen(g_ep)); EXPECT_TRUE(sockfd >= 0); @@ -1505,6 +1756,7 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); @@ -1512,7 +1764,8 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); // [0..4) magic ASSERT_EQ(0, memcmp(data, "RDMA", 4)); @@ -1530,7 +1783,7 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { msg.Deserialize(data + 4); ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, msg.msg_len); ASSERT_EQ(rdma::HELLO_V2_VERSION, msg.hello_ver); - ASSERT_EQ(rdma::IMPL_V2_VERSION, msg.impl_ver); + ASSERT_EQ(rdma::IMPL_V2_VERSION, msg.impl_ver); bthread_id_join(cntl.call_id()); } @@ -1546,10 +1799,12 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); - // Send a well-formed v2 hello so the server enters S_ACK_WAIT. + // Send a well-formed v2 hello so the server enters the common ACK_WAIT. rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; msg.hello_ver = rdma::HELLO_V2_VERSION; @@ -1563,12 +1818,15 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); // Read server's reply hello and assert its byte-level layout. uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, read(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_EQ(0, memcmp(reply, "RDMA", 4)); ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, @@ -1582,21 +1840,24 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { reply_msg.Deserialize(reply + 4); ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, reply_msg.msg_len); ASSERT_EQ(rdma::HELLO_V2_VERSION, reply_msg.hello_ver); - ASSERT_EQ(rdma::IMPL_V2_VERSION, reply_msg.impl_ver); + ASSERT_EQ(rdma::IMPL_V2_VERSION, reply_msg.impl_ver); // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends // cleanly without requiring real RDMA hardware. uint32_t flags = butil::HostToNet32(0); ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } -TEST_F(RdmaTest, v2_server_drains_tail_then_reads_ack) { +TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { StartServer(); sockaddr_in addr; @@ -1606,9 +1867,11 @@ TEST_F(RdmaTest, v2_server_drains_tail_then_reads_ack) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != NULL); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); // Build a v2 hello with msg_len = 48 (40 base + 8B zero tail). rdma::v2_wire::HelloMessage msg{}; @@ -1621,23 +1884,23 @@ TEST_F(RdmaTest, v2_server_drains_tail_then_reads_ack) { msg.qp_num = 0; msg.gid = rdma::GetRdmaGid(); - uint8_t buf[48]; + uint8_t buf[52]; memcpy(buf, "RDMA", 4); msg.Serialize(buf + 4); memset(buf + 40, 0x00, 8); // 8B zero tail - ASSERT_TRUE(WriteAll(sockfd, buf, 48)); - // The tail is drained as part of the hello, so the server ends up waiting - // for the ACK. Wait for that before sending it, otherwise the ACK could - // ride along in the same read and this would no longer test the drain. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - - // Send the real ACK (flags=1 = ACK_MSG_RDMA_OK). + // Coalesce the real ACK with the hello. The v2 parser must consume only + // msg_len bytes and preserve the ACK for the next handshake step. uint32_t flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); + memcpy(buf + 48, &flags, sizeof(flags)); + ASSERT_EQ(sizeof(buf), write(sockfd, buf, sizeof(buf))); + usleep(100000); + + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -1652,9 +1915,11 @@ TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); // Build a v2 hello with msg_len = 4097 (HELLO_V2_MSG_LEN_MAX + 1). // We only send the 40B base; the server must reject before reading @@ -1672,10 +1937,14 @@ TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { uint8_t buf[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(buf, "RDMA", 4); msg.Serialize(buf + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_TRUE(WaitForServerSocketGone()); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); sockfd.reset(-1); + usleep(100000); StopServer(); } @@ -1692,6 +1961,7 @@ class HandshakeVersionFlag { ~HandshakeVersionFlag() { rdma::FLAGS_rdma_client_handshake_version = _saved; } + private: int _saved; }; @@ -1703,8 +1973,7 @@ std::string MakeV3Packet(const rdma::RdmaHello& msg) { std::string packet; packet.reserve(4 + 4 + body.size()); packet.append("RDM3", 4); - uint32_t pb_size_be = - butil::HostToNet32(static_cast(body.size())); + uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); packet.append(reinterpret_cast(&pb_size_be), 4); packet.append(body); return packet; @@ -1723,13 +1992,12 @@ rdma::RdmaHello MakeValidV3Hello() { msg.set_rq_size(16); msg.set_lid(0); ibv_gid gid = rdma::GetRdmaGid(); - msg.set_gid(std::string(reinterpret_cast(gid.raw), - sizeof(gid.raw))); + msg.set_gid( + std::string(reinterpret_cast(gid.raw), sizeof(gid.raw))); msg.set_qp_num(0); return msg; } - TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { HandshakeVersionFlag _hsv(3); @@ -1799,14 +2067,17 @@ TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); // Send a valid v3 hello. std::string packet = MakeV3Packet(MakeValidV3Hello()); ASSERT_EQ((ssize_t)packet.size(), write(sockfd, packet.data(), packet.size())); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); // Read server's reply hello: 4B magic + 4B pb_size + body. uint8_t reply_magic[4]; @@ -1833,12 +2104,14 @@ TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends // cleanly without requiring real RDMA hardware. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), - write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -1853,13 +2126,16 @@ TEST_F(RdmaTest, v3_server_rejects_zero_pb_size) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); // "RDM3" + pb_size = 0 (4B big-endian zero). uint8_t buf[8] = {'R', 'D', 'M', '3', 0, 0, 0, 0}; ASSERT_EQ(8, write(sockfd, buf, 8)); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); sockfd.reset(-1); StopServer(); @@ -1875,7 +2151,8 @@ TEST_F(RdmaTest, v3_server_rejects_oversized_pb_size) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); uint8_t buf[8]; @@ -1885,7 +2162,9 @@ TEST_F(RdmaTest, v3_server_rejects_oversized_pb_size) { butil::HostToNet32(static_cast(rdma::HELLO_V3_MAX_PB_SIZE + 1)); memcpy(buf + 4, &pb_size_be, 4); ASSERT_EQ(8, write(sockfd, buf, 8)); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); sockfd.reset(-1); StopServer(); @@ -1901,7 +2180,8 @@ TEST_F(RdmaTest, v3_server_rejects_invalid_pb_bytes) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); // "RDM3" + pb_size = 8 + 8 bytes of 0xff (invalid protobuf body). @@ -1911,7 +2191,9 @@ TEST_F(RdmaTest, v3_server_rejects_invalid_pb_bytes) { memcpy(buf + 4, &pb_size_be, 4); memset(buf + 8, 0xff, 8); ASSERT_EQ(16, write(sockfd, buf, 16)); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); sockfd.reset(-1); StopServer(); @@ -1927,38 +2209,43 @@ TEST_F(RdmaTest, v3_server_invalid_sq_size_falls_back) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); rdma::RdmaHello msg = MakeValidV3Hello(); msg.set_sq_size(0); // invalid: < MIN_QP_SIZE (16) std::string packet = MakeV3Packet(msg); - ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); // Server validated the hello as invalid -> _rdma_state = RDMA_OFF, - // but still proceeds to S_ACK_WAIT (sends its own reply hello). - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + // but still proceeds to the common ACK_WAIT (sends its own reply hello). + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); // Drain server's reply hello (content not asserted here; covered // by v3_server_hello_bytes_baseline). uint8_t reply_hdr[8]; ASSERT_EQ(8, read(sockfd, reply_hdr, 8)); ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); - uint32_t reply_pb_size = butil::NetToHost32( - *reinterpret_cast(reply_hdr + 4)); + uint32_t reply_pb_size = + butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); std::string reply_body(reply_pb_size, '\0'); ASSERT_EQ((ssize_t)reply_pb_size, read(sockfd, &reply_body[0], reply_pb_size)); // Client ACK flags=0 -> server settles into FALLBACK_TCP. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), - write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -1969,16 +2256,14 @@ class EceFlagGuard { explicit EceFlagGuard(bool v) : _saved(rdma::FLAGS_rdma_ece) { rdma::FLAGS_rdma_ece = v; } - ~EceFlagGuard() { - rdma::FLAGS_rdma_ece = _saved; - } + ~EceFlagGuard() { rdma::FLAGS_rdma_ece = _saved; } + private: bool _saved; }; // Build a valid v3 hello that also carries an ECE block. -rdma::RdmaHello MakeValidV3HelloWithEce(uint32_t vendor_id, - uint32_t options, +rdma::RdmaHello MakeValidV3HelloWithEce(uint32_t vendor_id, uint32_t options, uint32_t comp_mask) { rdma::RdmaHello msg = MakeValidV3Hello(); rdma::RdmaEce* ece = msg.mutable_ece(); @@ -1994,17 +2279,18 @@ static void ReadServerV3Reply(int fd, rdma::RdmaHello* reply) { uint8_t reply_hdr[8]; ASSERT_EQ(8, read(fd, reply_hdr, 8)); ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); - uint32_t reply_pb_size = butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); + uint32_t reply_pb_size = + butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); ASSERT_GT(reply_pb_size, 0u); ASSERT_LE(reply_pb_size, 4096u); std::string reply_body(reply_pb_size, '\0'); - ASSERT_EQ((ssize_t)reply_pb_size, - read(fd, &reply_body[0], reply_pb_size)); + ASSERT_EQ((ssize_t)reply_pb_size, read(fd, &reply_body[0], reply_pb_size)); ASSERT_TRUE(reply->ParseFromString(reply_body)); } // A client hello carrying ECE must not break the server handshake: with ECE -// enabled the server still parses the hello and advances to S_ACK_WAIT. +// enabled the server still parses the hello and advances to the common +// ACK_WAIT. TEST_F(RdmaTest, v3_server_accepts_client_hello_with_ece) { EceFlagGuard ece_flag_guard(true); StartServer(); @@ -2016,14 +2302,17 @@ TEST_F(RdmaTest, v3_server_accepts_client_hello_with_ece) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); ASSERT_EQ((ssize_t)packet.size(), write(sockfd, packet.data(), packet.size())); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + usleep(100000); + + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); @@ -2031,10 +2320,13 @@ TEST_F(RdmaTest, v3_server_accepts_client_hello_with_ece) { // ACK flags=0 -> clean FALLBACK_TCP so the test ends without hardware. uint32_t flags = butil::HostToNet32(0); ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -2051,30 +2343,33 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_when_disabled) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); - ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); - // Reading the reply in full doubles as the synchronization point. rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); EXPECT_FALSE(reply.has_ece()); uint32_t flags = butil::HostToNet32(0); - ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); StopServer(); } // When ECE is enabled but there is no negotiated result (UT skips the real QP // bring-up, so the server never fills _outgoing_ece), the server reply must -// still NOT advertise ECE (FillLocalRdmaHello degrade branch #2 -> degrade-safe). +// still NOT advertise ECE (FillLocalRdmaHello degrade branch #2 -> +// degrade-safe). TEST_F(RdmaTest, v3_server_reply_has_no_ece_without_hw_negotiation) { EceFlagGuard ece_flag_guard(true); StartServer(); @@ -2086,24 +2381,26 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_without_hw_negotiation) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); - ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); - // Reading the reply in full doubles as the synchronization point. rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); EXPECT_FALSE(reply.has_ece()); uint32_t flags = butil::HostToNet32(0); - ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); StopServer(); } @@ -2113,13 +2410,17 @@ class ResourceAllocFailGuard { : _saved(rdma::g_fail_resource_alloc_for_test) { rdma::g_fail_resource_alloc_for_test = v; } - ~ResourceAllocFailGuard() { - rdma::g_fail_resource_alloc_for_test = _saved; - } + ~ResourceAllocFailGuard() { rdma::g_fail_resource_alloc_for_test = _saved; } + private: bool _saved; }; +TEST_F(RdmaTest, allocation_failure_cleans_endpoint_and_serves_complete_rpc) { + ResourceAllocFailGuard alloc_fail_guard(true); + ASSERT_NO_FATAL_FAILURE(VerifyCompleteFallbackRPC()); +} + TEST_F(RdmaTest, client_alloc_resource_fail_fallback_tcp) { StartServer(); ResourceAllocFailGuard alloc_fail_guard(true); @@ -2139,11 +2440,13 @@ TEST_F(RdmaTest, client_alloc_resource_fail_fallback_tcp) { req.set_sleep_us(200000); google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); // The socket must not be failed, otherwise it can no longer carry TCP. ASSERT_FALSE(s->Failed()); @@ -2165,9 +2468,11 @@ TEST_F(RdmaTest, server_alloc_resource_fail_fallback_tcp) { butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); ASSERT_TRUE(sockfd >= 0); ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); + usleep(100000); // wait for server to handle the msg + Socket* s = GetSocketFromServer(0); ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); // Send a well-formed v2 hello: the negotiation succeeds // but the resource allocation does not. @@ -2186,18 +2491,22 @@ TEST_F(RdmaTest, server_alloc_resource_fail_fallback_tcp) { msg.Serialize(data + 4); ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_FALSE(s->Failed()); // Ack without RDMA so that the server finishes the handshake in TCP mode. uint32_t flags = butil::HostToNet32(0); ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); ASSERT_FALSE(s->Failed()); sockfd.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); StopServer(); } @@ -2221,9 +2530,11 @@ TEST_F(RdmaTest, try_global_disable_rdma) { req.set_sleep_us(200000); google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); bthread_id_join(cntl.call_id()); ASSERT_EQ(0, cntl.ErrorCode()); @@ -2320,13 +2631,7 @@ 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; @@ -2350,7 +2655,7 @@ static int SendEchoRpcs(Channel& channel, int rpc_num, size_t attach_size, req[i].set_code(i + 1); if (attach_size > 0) { EXPECT_EQ(0, attach[i].resize( - attach_size, static_cast('a' + i % 26))); + 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()); @@ -2376,17 +2681,6 @@ static int SendEchoRpcs(Channel& channel, int rpc_num, size_t attach_size, 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; @@ -2398,10 +2692,18 @@ 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 = 5000; + chan_options.timeout_ms = 500; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); + 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()); StopServer(); } @@ -2412,10 +2714,18 @@ TEST_P(RdmaRpcTest, tcp_client_to_tcp_server) { Channel channel; ChannelOptions chan_options; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 5000; + chan_options.timeout_ms = 500; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); + 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()); StopServer(); } @@ -2426,10 +2736,18 @@ TEST_P(RdmaRpcTest, tcp_client_to_rdma_server) { Channel channel; ChannelOptions chan_options; chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 5000; + chan_options.timeout_ms = 500; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); + 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()); StopServer(); } @@ -2441,10 +2759,18 @@ 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 = 5000; + chan_options.timeout_ms = 500; chan_options.max_retry = 0; ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - ASSERT_NO_FATAL_FAILURE(SendEchoRpcsInRounds(channel)); + 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()); StopServer(); } @@ -2461,13 +2787,12 @@ TEST_P(RdmaRpcTest, tcp_client_to_rdma_server_short_connection) { 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; + << "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; @@ -2515,10 +2840,12 @@ TEST_P(RdmaRpcTest, rdma_server_survives_connection_churn) { static const int RPC_NUM = 1024; void DumpRdmaEndpointInfo(Socket* client, Socket* server) { - std::cout << std::endl << "client:"; - static_cast(client->_transport.get())->_rdma_ep->DebugInfo(std::cout); - std::cout << std::endl << "server:"; - static_cast(server->_transport.get())->_rdma_ep->DebugInfo(std::cout); + std::cout << std::endl + << "client:"; + RdmaTransport::Get(client)->_rdma_ep->DebugInfo(std::cout); + std::cout << std::endl + << "server:"; + RdmaTransport::Get(server)->_rdma_ep->DebugInfo(std::cout); } TEST_P(RdmaRpcTest, send_rpcs_in_one_qp) { @@ -2592,8 +2919,8 @@ TEST_P(RdmaRpcTest, send_rpcs_in_one_qp) { Socket* m = GetSocketFromServer(0); DumpRdmaEndpointInfo(s.get(), m); } - ASSERT_TRUE(0 == cntl[i].ErrorCode() || - EOVERCROWDED == cntl[i].ErrorCode()) << "req[" << i << "] " << berror(cntl[i].ErrorCode()); + ASSERT_TRUE(0 == cntl[i].ErrorCode() || EOVERCROWDED == cntl[i].ErrorCode()) + << "req[" << i << "] " << berror(cntl[i].ErrorCode()); } SocketUniquePtr s; @@ -2646,7 +2973,8 @@ TEST_P(RdmaRpcTest, send_rpc_in_many_qp) { req[i].set_message(__FUNCTION__); cntl[i].request_attachment().append(attach); google::protobuf::Closure* done = DoNothing(); - ::test::EchoService::Stub(&channel[i]).Echo(&cntl[i], &req[i], &res[i], done); + ::test::EchoService::Stub(&channel[i]) + .Echo(&cntl[i], &req[i], &res[i], done); } for (int i = 0; i < RPC_NUM; ++i) { bthread_id_join(cntl[i].call_id()); @@ -2784,12 +3112,12 @@ TEST_P(RdmaRpcTest, server_stop_during_rpc) { for (int i = 0; i < RPC_NUM; ++i) { bthread_id_join(cntl[i].call_id()); - if (i == 0) StopServer(); + if (i == 0) + StopServer(); int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || - error_code == EEOF || - error_code == ELOGOFF || - error_code == EHOSTDOWN) << "req[" << i << "]: " << error_code; + ASSERT_TRUE(error_code == 0 || error_code == EEOF || + error_code == ELOGOFF || error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; } } @@ -2826,10 +3154,9 @@ TEST_P(RdmaRpcTest, server_close_during_rpc) { for (int i = 0; i < RPC_NUM; ++i) { bthread_id_join(cntl[i].call_id()); int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || - error_code == EEOF || - error_code == EFAILEDSOCKET || - error_code == EHOSTDOWN) << "req[" << i << "]: " << error_code; + ASSERT_TRUE(error_code == 0 || error_code == EEOF || + error_code == EFAILEDSOCKET || error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; } StopServer(); @@ -2867,9 +3194,9 @@ TEST_P(RdmaRpcTest, client_close_during_rpc) { for (int i = 0; i < RPC_NUM; ++i) { bthread_id_join(cntl[i].call_id()); int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || - error_code == ECLOSE || - error_code == EHOSTDOWN) << "req[" << i << "]: " << error_code; + ASSERT_TRUE(error_code == 0 || error_code == ECLOSE || + error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; } StopServer(); @@ -2900,7 +3227,6 @@ TEST_P(RdmaRpcTest, rdma_client_close_during_rpc_repeatedly) { 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 " @@ -2938,10 +3264,10 @@ TEST_P(RdmaRpcTest, verbs_error_handling) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); // wait for rdma handshake complete + SocketUniquePtr s; ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - // The QP below only exists once the handshake is over. - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); ibv_send_wr wr; memset(&wr, 0, sizeof(wr)); ibv_sge sge; @@ -2952,7 +3278,7 @@ TEST_P(RdmaRpcTest, verbs_error_handling) { wr.sg_list = &sge; wr.num_sge = 1; ibv_send_wr* bad = nullptr; - auto rdma_transport = RdmaTransportOf(s); + auto rdma_transport = RdmaTransport::Get(s); ibv_post_send(rdma_transport->_rdma_ep->_resource->qp, &wr, &bad); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERDMA, cntl.ErrorCode()); @@ -2975,9 +3301,8 @@ TEST_P(RdmaRpcTest, rdma_use_parallel_channel) { opts.socket_mode = SOCKET_MODE_RDMA; for (size_t i = 0; i < NCHANS; ++i) { ASSERT_EQ(0, subchans[i].Init(_naming_url.c_str(), "rR", &opts)); - ASSERT_EQ(0, channel.AddChannel( - &subchans[i], DOESNT_OWN_CHANNEL, - nullptr, nullptr)); + ASSERT_EQ(0, channel.AddChannel(&subchans[i], DOESNT_OWN_CHANNEL, nullptr, + nullptr)); } ASSERT_EQ(0, channel.Init(nullptr)); @@ -3023,7 +3348,7 @@ TEST_P(RdmaRpcTest, rdma_use_selective_channel) { StopServer(); } -static void MockFree(void* buf) { } +static void MockFree(void* buf) {} TEST_P(RdmaRpcTest, send_rpcs_with_user_defined_iobuf) { if (!FLAGS_rdma_test_enable) { @@ -3044,7 +3369,8 @@ TEST_P(RdmaRpcTest, send_rpcs_with_user_defined_iobuf) { test::EchoResponse res[RPC_NUM]; butil::IOBuf attach; - void* data = malloc(4096);; + void* data = malloc(4096); + ; attach.append_user_data(data, 4096, nullptr); req[0].set_message(__FUNCTION__); cntl[0].request_attachment().append(attach); @@ -3063,12 +3389,14 @@ TEST_P(RdmaRpcTest, send_rpcs_with_user_defined_iobuf) { memset(mr[2 * i], i % 100, 4096); lkey[2 * i] = rdma::RegisterMemoryForRdma(mr[2 * i], 4096); ASSERT_TRUE(lkey[2 * i] != 0); - cntl[i].request_attachment().append_user_data_with_meta(mr[2 * i] + i, 4096 - i, MockFree, lkey[2 * i]); + cntl[i].request_attachment().append_user_data_with_meta( + mr[2 * i] + i, 4096 - i, MockFree, lkey[2 * i]); mr[2 * i + 1] = (char*)malloc(4096); memset(mr[2 * i + 1], i % 100, 4096); lkey[2 * i + 1] = rdma::RegisterMemoryForRdma(mr[2 * i + 1], 4096); ASSERT_TRUE(lkey[2 * i + 1] != 0); - cntl[i].request_attachment().append_user_data_with_meta(mr[2 * i + 1] + i, 4096 - i, MockFree, lkey[2 * i + 1]); + cntl[i].request_attachment().append_user_data_with_meta( + mr[2 * i + 1] + i, 4096 - i, MockFree, lkey[2 * i + 1]); req[i].set_message(__FUNCTION__); google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); @@ -3134,12 +3462,10 @@ TEST_P(RdmaRpcTest, try_memory_pool_empty) { // The server always accepts both via magic-byte dispatch, so this // proves the upper-layer RPC paths behave identically under either // wire format. -INSTANTIATE_TEST_SUITE_P( - HandshakeVersion, RdmaRpcTest, - ::testing::Values(2, 3), - [](const ::testing::TestParamInfo& info) { - return std::string("v") + std::to_string(info.param); - }); +INSTANTIATE_TEST_SUITE_P(HandshakeVersion, RdmaRpcTest, ::testing::Values(2, 3), + [](const ::testing::TestParamInfo& info) { + return std::string("v") + std::to_string(info.param); + }); static const uint32_t MELLANOX = 0x02c9; static const uint32_t INTEL = 0x8086; diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp new file mode 100644 index 0000000000..8c5ad480e6 --- /dev/null +++ b/test/brpc_transport_handshake_unittest.cpp @@ -0,0 +1,960 @@ +// 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 + +#include +#include +#include +#include +#include + +#include "bthread/bthread.h" +#include "butil/fd_guard.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/handshake/handshake_io.h" +#include "brpc/handshake/handshake_adapter.h" +#include "brpc/input_messenger.h" +#include "brpc/policy/transport_handshake_protocol.h" +#if BRPC_WITH_RDMA +#include "brpc/rdma_transport.h" +#endif +#include "brpc/socket.h" +#include "brpc/transport_handshake.h" + +namespace brpc { +namespace handshake { + +class MemoryHandshakeIO : public HandshakeIO { +public: + explicit MemoryHandshakeIO(const std::string& input = std::string()) + : _input(input), _offset(0) {} + + int ReadExact(void* data, size_t len) override { + if (_input.size() - _offset < len) { + errno = EIO; + return -1; + } + memcpy(data, _input.data() + _offset, len); + _offset += len; + return 0; + } + + int WriteAll(const void* data, size_t len) override { + _output.append(static_cast(data), len); + return 0; + } + + int PushBack(const void* data, size_t len) override { + _pushed_back.append(static_cast(data), len); + return 0; + } + + const std::string& output() const { return _output; } + const std::string& pushed_back() const { return _pushed_back; } + +private: + std::string _input; + size_t _offset; + std::string _output; + std::string _pushed_back; +}; + +struct BlockingHandshakeRead { + SocketHandshakeIO* io; + int result; + int error; +}; + +static void* RunBlockingHandshakeRead(void* arg) { + BlockingHandshakeRead* read = + static_cast(arg); + char byte = 0; + read->result = read->io->ReadExact(&byte, sizeof(byte)); + read->error = errno; + return NULL; +} + +class TestAppConnect : public AppConnect { +public: + void StartConnect(const Socket*, void (*done)(int, void*), + void* data) override { + done(0, data); + } + + void StopConnect(Socket*) override {} +}; + +static FrameSpec FixedSpec(const char* magic, size_t magic_len, + size_t total_len) { + return FrameSpec(magic, magic_len, total_len, total_len, + FrameSpec::FIXED); +} + +class TestHandshakeProtocol : public HandshakeProtocol { +public: + explicit TestHandshakeProtocol(std::string* calls = NULL) + : _calls(calls), _with_extension(false), + _extension_size(0) {} + + int ProtocolVersion() const override { return 7; } + const FrameSpec& HelloFrameSpec() const override { + static const FrameSpec spec = FixedSpec("HS", 2, 4); + return spec; + } + const FrameSpec& AckFrameSpec() const override { + static const FrameSpec spec = FixedSpec(NULL, 0, 1); + return spec; + } + StepResult BuildHello(bool enabled, std::string* payload) override { + Append("build "); + *payload = enabled ? "LO" : "NO"; + return STEP_OK; + } + StepResult ParseHello(const std::string& payload) override { + Append("parse "); + return payload == "OK" ? STEP_OK : STEP_FALLBACK; + } + StepResult BuildAck(bool enabled, std::string* payload) override { + Append(enabled ? "ack1 " : "ack0 "); + *payload = enabled ? "1" : "0"; + return STEP_OK; + } + StepResult ParseAck( + const std::string& payload, bool* enabled) override { + Append("parse_ack "); + *enabled = payload == "1"; + return payload == "1" || payload == "0" ? STEP_OK : STEP_ERROR; + } + bool HasExtension() const override { return _with_extension; } + const FrameSpec& ExtensionFrameSpec() const override { + return _extension_spec; + } + StepResult BuildExtension( + bool enabled, std::string* payload) override { + *payload = enabled + ? std::string("E1", _extension_size) + : std::string("E0", _extension_size); + return STEP_OK; + } + StepResult ParseExtension(const std::string& payload) override { + return payload == std::string("E1", _extension_size) + ? STEP_OK : STEP_FALLBACK; + } + void EnableExtension(size_t size) { + _with_extension = true; + _extension_size = size; + _extension_spec = FixedSpec(NULL, 0, size); + } + +private: + void Append(const char* value) { + if (_calls != NULL) { + *_calls += value; + } + } + + std::string* _calls; + bool _with_extension; + size_t _extension_size; + FrameSpec _extension_spec; +}; + +class TestHandshakeTransport : public HandshakeTransport { +public: + explicit TestHandshakeTransport(std::string* calls = NULL) + : prepare_result(STEP_OK), negotiate_result(STEP_OK), + validate_result(STEP_OK), tcp_active(false), + high_speed_active(false), failed(false), _calls(calls) {} + + StepResult PrepareResources() override { + Append("prepare "); + return prepare_result; + } + StepResult NegotiateResources() override { + Append("negotiate "); + return negotiate_result; + } + void OnEstablished() override { + high_speed_active = true; + Append("activate"); + } + void OnFallback() override { tcp_active = true; } + void OnFailed() override { failed = true; } + StepResult ValidateEstablished() override { + Append("validate "); + return validate_result; + } + + StepResult prepare_result; + StepResult negotiate_result; + StepResult validate_result; + bool tcp_active; + bool high_speed_active; + bool failed; + +private: + void Append(const char* value) { + if (_calls != NULL) { + *_calls += value; + } + } + std::string* _calls; +}; + +class TestServerHandshakeTransport : public TestHandshakeTransport { +public: + butil::IOBuf* source = nullptr; + StepResult ValidateEstablished() override { + return source->empty() ? STEP_OK : STEP_ERROR; + } +}; + +class TestServerHandshakeAdapter : public StandardHandshakeAdapter { +public: + TestServerHandshakeAdapter() { session.SetIOForTest(&io); } + MemoryHandshakeIO io; + mutable HandshakeSession session; + TestHandshakeProtocol protocol; + TestServerHandshakeTransport transport; + +protected: + StepResult RunServerStep(butil::IOBuf* source, Socket*) override { + IOBufHandshakeInput input(source); + std::vector protocols(1, &protocol); + transport.source = source; + return session.RunServer(protocols, &input, &transport, false); + } + HandshakeSession* GetSession(Socket*) const override { + return &session; + } +}; + +class TestUpgradeMessenger : public InputMessenger { +public: + using InputMessenger::OnNewMessagesUntil; + TestServerHandshakeAdapter adapter; + int peer_fd = -1; + int parse_calls = 0; + int write_result = -1; + + static ParseResult Parse(butil::IOBuf* source, Socket* socket, + bool, const void*) { + TestUpgradeMessenger* messenger = + static_cast(socket->user()); + ++messenger->parse_calls; + ParseResult result = messenger->adapter.ExecuteServerHandshake( + source, socket); + if (messenger->adapter.session.phase(butil::memory_order_acquire) == + ESTABLISHED) { + // Make bytes readable after the ACK was consumed, but before the + // outer event handler has regained control. + messenger->write_result = write(messenger->peer_fd, "RPC", 3); + } + return result; + } + static void Process(InputMessageBase*) {} + static bool StopReading(Socket* socket) { + return static_cast(socket->user()) + ->adapter.session.phase(butil::memory_order_acquire) == ESTABLISHED; + } +}; + +TEST(TransportHandshakeTest, adapter_retains_context_for_fragmented_extension) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + TestServerHandshakeAdapter adapter; + adapter.protocol.EnableExtension(2); + butil::IOBuf source; + source.append("HSOK", 4); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + ASSERT_NE(nullptr, socket->parsing_context()); + ServerHandshakeContext* context = static_cast( + socket->parsing_context()); + ASSERT_EQ(&adapter, context->adapter()); + source.append("E", 1); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + context->adapter()->ExecuteServerHandshake(&source, socket.get()) + .error()); + source.append("11", 2); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + context->adapter()->ExecuteServerHandshake(&source, socket.get()) + .error()); + EXPECT_EQ(ESTABLISHED, adapter.session.phase(butil::memory_order_acquire)); + EXPECT_TRUE(source.empty()); + EXPECT_EQ(nullptr, socket->parsing_context()); + socket->SetFailed(); +} + +TEST(TransportHandshakeTest, adapter_rejects_tcp_data_coalesced_with_upgrade_ack) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + TestServerHandshakeAdapter adapter; + butil::IOBuf source; + source.append("HSOK1RPC", 8); + EXPECT_EQ(PARSE_ERROR_ABSOLUTELY_WRONG, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_EQ(FAILED, adapter.session.phase(butil::memory_order_acquire)); + EXPECT_FALSE(adapter.transport.high_speed_active); + EXPECT_EQ(3UL, source.size()); + socket->SetFailed(); +} + +TEST(TransportHandshakeTest, tcp_read_loop_stops_when_upgrade_completes) { + TestUpgradeMessenger messenger; + InputMessageHandler handler{}; + handler.parse = TestUpgradeMessenger::Parse; + handler.process = TestUpgradeMessenger::Process; + handler.name = "test_upgrade"; + ASSERT_EQ(0, messenger.AddNonProtocolHandler(handler)); + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + messenger.peer_fd = peer_fd; + SocketOptions options; + options.fd = fds[0]; + options.user = &messenger; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + ASSERT_EQ(5, write(peer_fd, "HSOK1", 5)); + messenger.OnNewMessagesUntil(socket.get(), TestUpgradeMessenger::StopReading); + EXPECT_EQ(1, messenger.parse_calls); + EXPECT_EQ(3, messenger.write_result); + EXPECT_EQ(ESTABLISHED, + messenger.adapter.session.phase(butil::memory_order_acquire)); + char remaining[3]; + ASSERT_EQ(3, recv(socket->fd(), remaining, sizeof(remaining), MSG_DONTWAIT)); + EXPECT_EQ(0, memcmp(remaining, "RPC", 3)); + socket->SetFailed(); +} + +static std::string MakeUBShmHello() { + std::string frame(64, '\0'); + memcpy(&frame[0], "UB", 2); + const uint16_t msg_len = butil::HostToNet16(64); + const uint16_t hello_ver = butil::HostToNet16(2); + const uint16_t impl_ver = butil::HostToNet16(1); + memcpy(&frame[2], &msg_len, sizeof(msg_len)); + memcpy(&frame[4], &hello_ver, sizeof(hello_ver)); + memcpy(&frame[6], &impl_ver, sizeof(impl_ver)); + return frame; +} + +TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + SocketHandshakeIO io(socket.get()); + BlockingHandshakeRead read = {&io, 0, 0}; + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + ASSERT_EQ(0, bthread_start_background( + &tid, &attr, RunBlockingHandshakeRead, &read)); + + bthread_usleep(10000); + ASSERT_EQ(0, socket->SetFailed( + EFAILEDSOCKET, "cancel blocked handshake read")); + ASSERT_EQ(0, bthread_join(tid, NULL)); + + EXPECT_EQ(-1, read.result); + EXPECT_EQ(EFAILEDSOCKET, read.error); +} + +TEST(HandshakeFrameTest, supports_two_byte_fixed_magic) { + const FrameSpec spec = FixedSpec("UB", 2, 6); + std::string frame; + ASSERT_EQ(FRAME_OK, FrameCodec::Encode(spec, "data", &frame)); + ASSERT_EQ("UBdata", frame); +} + +TEST(HandshakeFrameTest, encodes_u16_total_length) { + const FrameSpec spec( + "RDMA", 4, 8, 64, FrameSpec::U16_TOTAL_LENGTH); + std::string frame; + ASSERT_EQ(FRAME_OK, FrameCodec::Encode(spec, "xy", &frame)); + ASSERT_EQ(8UL, frame.size()); + uint16_t total_be = 0; + memcpy(&total_be, frame.data() + 4, sizeof(total_be)); + ASSERT_EQ(8, butil::NetToHost16(total_be)); + ASSERT_EQ("xy", frame.substr(6)); +} + +TEST(HandshakeFrameTest, encodes_u32_body_length) { + const FrameSpec spec( + "RDM3", 4, 9, 32, FrameSpec::U32_BODY_LENGTH); + std::string frame; + ASSERT_EQ(FRAME_OK, FrameCodec::Encode(spec, "abc", &frame)); + uint32_t body_be = 0; + memcpy(&body_be, frame.data() + 4, sizeof(body_be)); + ASSERT_EQ(3U, butil::NetToHost32(body_be)); + ASSERT_EQ("abc", frame.substr(8)); +} + +TEST(HandshakeFrameTest, buffered_partial_frame_is_not_consumed) { + const FrameSpec spec( + "RDMA", 4, 8, 64, FrameSpec::U16_TOTAL_LENGTH); + std::string frame; + ASSERT_EQ(FRAME_OK, FrameCodec::Encode(spec, "payload", &frame)); + butil::IOBuf source; + source.append(frame.data(), frame.size() - 1); + IOBufHandshakeInput input(&source); + std::string payload; + ASSERT_EQ(FRAME_NEED_MORE, + FrameCodec::ParseBufferedFrame(&input, spec, &payload)); + ASSERT_EQ(frame.size() - 1, source.size()); + + source.append(frame.data() + frame.size() - 1, 1); + ASSERT_EQ(FRAME_OK, + FrameCodec::ParseBufferedFrame(&input, spec, &payload)); + ASSERT_EQ("payload", payload); + ASSERT_TRUE(source.empty()); +} + +TEST(HandshakeFrameTest, buffered_magic_mismatch_is_not_consumed) { + const FrameSpec spec = FixedSpec("UB", 2, 6); + butil::IOBuf source; + source.append("XXdata", 6); + IOBufHandshakeInput input(&source); + std::string payload; + ASSERT_EQ(FRAME_NOT_MINE, + FrameCodec::ParseBufferedFrame(&input, spec, &payload)); + ASSERT_EQ(6UL, source.size()); +} + +TEST(HandshakeFrameTest, blocking_magic_mismatch_is_pushed_back) { + MemoryHandshakeIO io("XX"); + const FrameSpec spec = FixedSpec("UB", 2, 6); + std::string payload; + ASSERT_EQ(FRAME_NOT_MINE, + FrameCodec::ReadFrame(&io, spec, true, &payload)); + ASSERT_EQ("XX", io.pushed_back()); +} + +TEST(HandshakeFrameTest, rejects_lengths_outside_bounds) { + const FrameSpec spec( + "RDMA", 4, 8, 16, FrameSpec::U16_TOTAL_LENGTH); + std::string header("RDMA", 4); + const uint16_t total_be = butil::HostToNet16(17); + header.append(reinterpret_cast(&total_be), sizeof(total_be)); + butil::IOBuf source; + source.append(header); + IOBufHandshakeInput input(&source); + std::string payload; + ASSERT_EQ(FRAME_PROTOCOL_ERROR, + FrameCodec::ParseBufferedFrame(&input, spec, &payload)); + ASSERT_EQ(header.size(), source.size()); +} + +#if BRPC_WITH_RDMA && BRPC_WITH_UBRING +TEST(TransportHandshakeTest, upgrade_capability_is_mode_specific) { + SocketOptions options; + SocketId id; + SocketUniquePtr socket; + + options.socket_mode = SOCKET_MODE_RDMA; + ASSERT_EQ(0, Socket::Create(options, &id)); + ASSERT_EQ(0, Socket::Address(id, &socket)); + AdapterTransport* rdma_adapter = AdapterTransport::Get(socket.get()); + ASSERT_TRUE(rdma_adapter->upgrade_capable(SOCKET_MODE_RDMA)); + ASSERT_FALSE(rdma_adapter->upgrade_capable(SOCKET_MODE_UBRING)); + socket->SetFailed(); + socket.reset(); + + options.socket_mode = SOCKET_MODE_UBRING; + ASSERT_EQ(0, Socket::Create(options, &id)); + ASSERT_EQ(0, Socket::Address(id, &socket)); + AdapterTransport* ubshm_adapter = AdapterTransport::Get(socket.get()); + ASSERT_TRUE(ubshm_adapter->upgrade_capable(SOCKET_MODE_UBRING)); + ASSERT_FALSE(ubshm_adapter->upgrade_capable(SOCKET_MODE_RDMA)); + socket->SetFailed(); +} +#endif + +TEST(TransportHandshakeTest, publish_fallback_after_tcp_state) { + HandshakeSession session; + int tcp_active = 0; + session.SetPhase(NEGOTIATING); + session.PublishFallback([&tcp_active]() { tcp_active = 1; }); + ASSERT_EQ(1, tcp_active); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, client_runs_codec_and_resource_sequence) { + MemoryHandshakeIO io("HSOK"); + HandshakeSession session; + session.SetIOForTest(&io); + std::string calls; + + TestHandshakeProtocol protocol(&calls); + TestHandshakeTransport transport(&calls); + + ASSERT_EQ(STEP_OK, session.RunClient(&protocol, &transport)); + ASSERT_EQ("prepare build parse negotiate ack1 activate", calls); + ASSERT_EQ("HSLO1", io.output()); + ASSERT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); + ASSERT_EQ(7, session.protocol_version()); +} + +TEST(TransportHandshakeTest, client_exchanges_extension_before_ack) { + MemoryHandshakeIO io("HSOKE"); + HandshakeSession session; + session.SetIOForTest(&io); + TestHandshakeProtocol protocol; + protocol.EnableExtension(1); + TestHandshakeTransport transport; + + ASSERT_EQ(STEP_OK, session.RunClient(&protocol, &transport)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, server_resumes_fragmented_extension_then_ack) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + source.append("HSOK", 4); + IOBufHandshakeInput input(&source); + TestHandshakeProtocol protocol; + protocol.EnableExtension(2); + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ(EXTENSION_WAIT, session.phase(butil::memory_order_acquire)); + EXPECT_EQ("HSLO", io.output()); + source.append("E", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ(EXTENSION_WAIT, session.phase(butil::memory_order_acquire)); + source.append("11", 2); // Remainder of extension, then ACK. + ASSERT_EQ(STEP_OK, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_TRUE(source.empty()); + EXPECT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, client_resource_failure_falls_back_before_io) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + TestHandshakeProtocol protocol; + TestHandshakeTransport transport; + transport.prepare_result = STEP_FALLBACK; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, session.RunClient(&protocol, &transport)); + ASSERT_TRUE(transport.tcp_active); + ASSERT_TRUE(io.output().empty()); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, server_resumes_at_buffered_ack) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + source.append("HSOK", 4); + IOBufHandshakeInput input(&source); + std::string calls; + + TestHandshakeProtocol protocol(&calls); + std::vector protocols(1, &protocol); + TestHandshakeTransport transport(&calls); + + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + ASSERT_EQ(ACK_WAIT, session.phase(butil::memory_order_acquire)); + ASSERT_EQ("HSLO", io.output()); + ASSERT_TRUE(source.empty()); + + source.append("1", 1); + ASSERT_EQ(STEP_OK, + session.RunServer(protocols, &input, &transport, false)); + ASSERT_EQ("parse prepare negotiate build parse_ack validate activate", + calls); + ASSERT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, server_resource_failure_falls_back_after_ack) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + source.append("HSOK0", 5); + IOBufHandshakeInput input(&source); + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + transport.prepare_result = STEP_FALLBACK; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &input, &transport, false)); + ASSERT_EQ("HSNO", io.output()); + ASSERT_TRUE(source.empty()); + ASSERT_TRUE(transport.tcp_active); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, server_falls_back_without_consuming_other_magic) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + source.append("XXok", 4); + IOBufHandshakeInput input(&source); + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + transport.prepare_result = STEP_ERROR; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &input, &transport, true)); + ASSERT_TRUE(transport.tcp_active); + ASSERT_EQ(4UL, source.size()); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); + + ASSERT_EQ(STEP_NOT_MINE, + session.RunServer(protocols, &input, &transport, true)); + ASSERT_EQ(4UL, source.size()); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); +} + +TEST(TransportHandshakeTest, server_enters_hello_phase_after_magic_matches) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + IOBufHandshakeInput input(&source); + + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + + source.append("H", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + ASSERT_EQ(UNINITIALIZED, session.phase(butil::memory_order_acquire)); + + source.append("S", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + ASSERT_EQ(HELLO_WAIT, session.phase(butil::memory_order_acquire)); + ASSERT_EQ(7, session.protocol_version()); + ASSERT_EQ(2UL, source.size()); +} + +TEST(TransportHandshakeTest, + impossible_magic_prefixes_try_other_protocols_without_consuming) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + const char* const invalid_prefixes[] = { + "X", "UX", "RX", "RDN", "RDMB", + }; + for (size_t i = 0; i < arraysize(invalid_prefixes); ++i) { + butil::IOBuf source; + source.append(invalid_prefixes[i]); + const size_t original_size = source.size(); + const ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()) + << invalid_prefixes[i]; + EXPECT_EQ(original_size, source.size()) + << invalid_prefixes[i]; + EXPECT_EQ(UNINITIALIZED, + AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_EQ(nullptr, socket->parsing_context()); + } + socket->SetFailed(); +} + +TEST(TransportHandshakeTest, + partial_rdma_magic_waits_for_more_data_without_consuming) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + const char* const partial_prefixes[] = { + "R", "RD", "RDM", + }; + for (size_t i = 0; i < arraysize(partial_prefixes); ++i) { + butil::IOBuf source; + source.append(partial_prefixes[i]); + const size_t original_size = source.size(); + const ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + EXPECT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()) + << partial_prefixes[i]; + EXPECT_EQ(original_size, source.size()) + << partial_prefixes[i]; + EXPECT_EQ(UNINITIALIZED, + AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_EQ(nullptr, socket->parsing_context()); + } + socket->SetFailed(); +} + +TEST(TransportHandshakeTest, + plain_tcp_server_incrementally_rejects_ubshm_upgrade) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + const std::string hello = MakeUBShmHello(); + butil::IOBuf source; + source.append(hello.data(), 1); + ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); + ASSERT_EQ(1UL, source.size()); + ASSERT_EQ(UNINITIALIZED, + AdapterTransport::Get(socket.get())->handshake_phase()); + + source.append(hello.data() + 1, hello.size() - 1); + result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); + ASSERT_TRUE(source.empty()); + ASSERT_NE(nullptr, socket->parsing_context()); + + char reply[64]; + ASSERT_EQ(sizeof(reply), read(peer_fd, reply, sizeof(reply))); + EXPECT_EQ(0, memcmp(reply, "UB", 2)); + uint16_t reply_len = 0; + memcpy(&reply_len, reply + 2, sizeof(reply_len)); + EXPECT_EQ(64, butil::NetToHost16(reply_len)); + EXPECT_EQ(0, reply[4]); + EXPECT_EQ(0, reply[5]); + EXPECT_EQ(0, reply[6]); + EXPECT_EQ(0, reply[7]); + + const uint32_t ack = 0; + source.append(&ack, sizeof(ack)); + result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + ASSERT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); + ASSERT_TRUE(source.empty()); + ASSERT_EQ(FALLBACK_TCP, + AdapterTransport::Get(socket.get())->handshake_phase()); + ASSERT_EQ(nullptr, socket->parsing_context()); + socket->SetFailed(); +} + +TEST(TransportHandshakeTest, + plain_tcp_server_preserves_data_coalesced_behind_ubshm_ack) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + butil::IOBuf source; + source.append(MakeUBShmHello()); + const uint32_t ack = 0; + source.append(&ack, sizeof(ack)); + const char application_data[] = "application-data"; + source.append(application_data, sizeof(application_data) - 1); + + const ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, NULL); + ASSERT_FALSE(result.is_ok()); + ASSERT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); + ASSERT_EQ(sizeof(application_data) - 1, source.size()); + + char remaining[sizeof(application_data) - 1] = {}; + ASSERT_EQ(sizeof(remaining), + source.copy_to(remaining, sizeof(remaining))); + EXPECT_EQ(0, memcmp(application_data, remaining, sizeof(remaining))); + ASSERT_EQ(FALLBACK_TCP, + AdapterTransport::Get(socket.get())->handshake_phase()); + ASSERT_EQ(nullptr, socket->parsing_context()); + socket->SetFailed(); +} + +#if BRPC_WITH_URMA +TEST(TransportHandshakeTest, urma_socket_skips_adapter_handshake_parser) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + + SocketOptions options; + options.fd = fds[0]; + options.socket_mode = SOCKET_MODE_URMA; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + butil::IOBuf source; + source.append("PRPC", 4); + const ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, nullptr); + ASSERT_FALSE(result.is_ok()); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); + EXPECT_EQ(4UL, source.size()); + EXPECT_EQ(nullptr, socket->parsing_context()); + socket->SetFailed(); +} +#endif + +#if BRPC_WITH_RDMA +TEST(TransportHandshakeTest, + inherited_adapter_connect_is_not_wrapped_twice) { + int first_fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, first_fds)); + butil::fd_guard first_peer(first_fds[1]); + + const std::shared_ptr original = + std::make_shared(); + + SocketOptions first_options; + first_options.fd = first_fds[0]; + first_options.user = + static_cast(get_or_new_client_side_messenger()); + first_options.need_on_edge_trigger = true; + first_options.socket_mode = SOCKET_MODE_RDMA; + first_options.app_connect = original; + + SocketId first_id; + ASSERT_EQ(0, Socket::Create(first_options, &first_id)); + SocketUniquePtr first_socket; + ASSERT_EQ(0, Socket::Address(first_id, &first_socket)); + + const std::shared_ptr wrapped = + AdapterTransport::Get(first_socket.get())->Connect(); + ASSERT_NE(nullptr, wrapped); + EXPECT_NE(original.get(), wrapped.get()); + + int second_fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, second_fds)); + butil::fd_guard second_peer(second_fds[1]); + + SocketOptions second_options = first_options; + second_options.fd = second_fds[0]; + second_options.app_connect = wrapped; + + SocketId second_id; + ASSERT_EQ(0, Socket::Create(second_options, &second_id)); + SocketUniquePtr second_socket; + ASSERT_EQ(0, Socket::Address(second_id, &second_socket)); + + const std::shared_ptr inherited = + AdapterTransport::Get(second_socket.get())->Connect(); + EXPECT_EQ(wrapped.get(), inherited.get()); + + first_socket->SetFailed(); + second_socket->SetFailed(); +} + +TEST(TransportHandshakeTest, + unavailable_client_upgrade_publishes_tcp_fallback) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + + const std::shared_ptr original = + std::make_shared(); + + SocketOptions options; + options.fd = fds[0]; + options.user = + static_cast(get_or_new_client_side_messenger()); + options.need_on_edge_trigger = true; + options.socket_mode = SOCKET_MODE_RDMA; + options.app_connect = original; + + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + + AdapterTransport* adapter = AdapterTransport::Get(socket.get()); + RdmaTransport* rdma_transport = + static_cast(adapter->high_speed_transport()); + ASSERT_NE(nullptr, rdma_transport); + rdma_transport->Release(); + + const std::shared_ptr connect = adapter->Connect(); + EXPECT_EQ(original.get(), connect.get()); + EXPECT_EQ(FALLBACK_TCP, adapter->handshake_phase()); + + socket->SetFailed(); +} +#endif + +} // namespace handshake +} // namespace brpc diff --git a/test/brpc_ubring_unittest.cpp b/test/brpc_ubring_unittest.cpp index 295d3f1289..f151a6492e 100644 --- a/test/brpc_ubring_unittest.cpp +++ b/test/brpc_ubring_unittest.cpp @@ -23,7 +23,28 @@ #include "brpc/socket.h" #if BRPC_WITH_UBRING +#include +#include +#include +#include +#include +#include +#include +#include "bthread/bthread.h" +#include "brpc/acceptor.h" +#include "brpc/adapter_transport.h" +#include "brpc/channel.h" +#include "brpc/controller.h" +#include "brpc/policy/baidu_rpc_meta.pb.h" +#include "brpc/policy/baidu_rpc_protocol.h" +#include "brpc/server.h" +#include "brpc/ubshm_transport.h" +#include "butil/fd_guard.h" +#include "butil/memory/scope_guard.h" +#include "butil/time.h" +#include "echo.pb.h" #include "brpc/ubshm/common/common.h" +#include "brpc/handshake/ubshm_handshake.h" #include "brpc/ubshm/ub_endpoint.h" #include "brpc/ubshm/shm/shm_def.h" #include "brpc/ubshm/shm/shm_mgr.h" @@ -108,6 +129,26 @@ TEST(HelloFormatExtensionTest, deserialize_unknown_format) { EXPECT_EQ(0x1234, extension.format_id); } +TEST(UBShmHandshakeAdapterTest, rejects_unsupported_format_extension) { + brpc::ubring::UBShmHandshakeAdapter adapter; + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.BuildExtension(true, &payload)); + EXPECT_EQ(std::string("\0\4\0\1", 4), payload); + EXPECT_EQ(brpc::handshake::STEP_OK, + adapter.ParseExtension(payload)); + + payload[3] = 2; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); + payload[3] = 0; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); + payload[1] = 3; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); +} + TEST_F(HelloMessageTest, serialize_deserialize_roundtrip) { msg.msg_len = 64; msg.hello_ver = 2; @@ -189,6 +230,98 @@ TEST_F(HelloMessageTest, toString_contains_fields) { EXPECT_NE(std::string::npos, s.find("UBRING_test")); } +TEST(UBShmHandshakeAdapterTest, codec_uses_v3_wire_format) { + brpc::ubring::UBShmHandshakeAdapter adapter; + char shm_name[SHM_MAX_NAME_BUFF_LEN] = {0}; + memcpy(shm_name, "UBRING_test_C", 14); + + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.BuildHello(true, 4 * 1024 * 1024, + shm_name, &payload)); + brpc::ubring::HelloMessage decoded{}; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.ParseHello(payload, &decoded)); + EXPECT_EQ(64, decoded.msg_len); + EXPECT_EQ(3, decoded.hello_ver); + EXPECT_EQ(1, decoded.impl_ver); + EXPECT_EQ(4 * 1024 * 1024, decoded.len); + EXPECT_EQ(0, memcmp(shm_name, decoded.shm_name, + SHM_MAX_NAME_BUFF_LEN)); + + std::string frame; + ASSERT_EQ(brpc::handshake::FRAME_OK, + brpc::handshake::FrameCodec::Encode( + adapter.HelloFrameSpec(), payload, &frame)); + ASSERT_EQ(64, frame.size()); + EXPECT_EQ("UB", frame.substr(0, 2)); +} + +TEST(UBShmHandshakeAdapterTest, short_name_is_zero_padded) { + brpc::ubring::UBShmHandshakeAdapter adapter; + char short_name[SHM_MAX_NAME_BUFF_LEN]; + memset(short_name, 0x5a, sizeof(short_name)); + short_name[0] = 'x'; + short_name[1] = '\0'; + + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.BuildHello(true, 4096, short_name, &payload)); + + brpc::ubring::HelloMessage decoded{}; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.ParseHello(payload, &decoded)); + EXPECT_EQ('x', decoded.shm_name[0]); + for (size_t i = 1; i < SHM_MAX_NAME_BUFF_LEN; ++i) { + EXPECT_EQ('\0', decoded.shm_name[i]) << "index=" << i; + } +} + +TEST(UBShmHandshakeAdapterTest, client_hello_advertises_allocated_shm_name) { + brpc::ubring::SHM local_shm{}; + local_shm.len = 4096; + strcpy(local_shm.name, "UBRING_127.0.0.1:8000_C"); + brpc::ubring::UBShmHandshakeAdapter adapter; + adapter.ConfigureClientHello(local_shm); + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, adapter.BuildHello(true, &payload)); + brpc::ubring::HelloMessage decoded{}; + ASSERT_EQ(brpc::handshake::STEP_OK, adapter.ParseHello(payload, &decoded)); + EXPECT_EQ(local_shm.len, decoded.len); + EXPECT_STREQ(local_shm.name, decoded.shm_name); +} + +TEST(UBShmHandshakeAdapterTest, rejects_unterminated_remote_name) { + brpc::ubring::HelloMessage message{}; + message.msg_len = 64; + message.hello_ver = 3; + message.impl_ver = 1; + message.len = 4096; + memset(message.shm_name, 'A', SHM_MAX_NAME_BUFF_LEN); + + std::string payload( + sizeof(HelloMessageLayout), '\0'); + message.Serialize(&payload[0]); + + brpc::ubring::UBShmHandshakeAdapter adapter; + brpc::ubring::HelloMessage decoded{}; + errno = 0; + EXPECT_EQ(brpc::handshake::STEP_ERROR, + adapter.ParseHello(payload, &decoded)); + EXPECT_EQ(EPROTO, errno); +} + +TEST(UBShmHandshakeAdapterTest, disabled_hello_requests_tcp_fallback) { + brpc::ubring::UBShmHandshakeAdapter adapter; + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.BuildHello(false, 0, NULL, &payload)); + brpc::ubring::HelloMessage decoded{}; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseHello(payload, &decoded)); + EXPECT_EQ(64, decoded.msg_len); +} + TEST(UBRingConfigurationTest, time_flags_include_units_and_expected_defaults) { struct TimeFlagExpectation { const char* name; @@ -271,16 +404,16 @@ using brpc::ubring::UBShmEndpointTest; TEST_F(UBShmEndpointTest, construct_initial_state) { ASSERT_NE(nullptr, _ep); EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, - _ep->_negotiated_data_format); + _ep->negotiated_data_format()); } TEST_F(UBShmEndpointTest, reset_clears_negotiated_data_format) { - _ep->_negotiated_data_format = brpc::ubring::UBR_DATA_FORMAT_LEGACY_64; + _ep->SetNegotiatedDataFormat(brpc::ubring::UBR_DATA_FORMAT_LEGACY_64); _ep->Reset(); EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, - _ep->_negotiated_data_format); + _ep->negotiated_data_format()); } TEST_F(UBShmEndpointTest, allocate_client_resources_real_shm) { @@ -302,6 +435,370 @@ TEST_F(UBShmEndpointTest, reset_is_idempotent) { _ep->Reset(); } +namespace { +bool WaitForUBCondition(const std::function& condition) { + const int64_t deadline = butil::gettimeofday_us() + 5000000; + while (!condition()) { + if (butil::gettimeofday_us() >= deadline) { + return false; + } + bthread_usleep(1000); + } + return true; +} + +bool TransferUBTestBytes(int fd, void* bytes, size_t size, bool sending) { + size_t offset = 0; + while (offset < size) { + const ssize_t n = sending + ? send(fd, static_cast(bytes) + offset, size - offset, + MSG_NOSIGNAL) + : recv(fd, static_cast(bytes) + offset, size - offset, 0); + if (n < 0 && errno == EINTR) { + continue; + } + if (n <= 0) { + return false; + } + offset += n; + } + return true; +} + +class UBFallbackEchoService : public test::EchoService { +public: + void Echo(google::protobuf::RpcController*, const test::EchoRequest* request, + test::EchoResponse* response, google::protobuf::Closure* done) + override { + response->set_message(request->message()); + done->Run(); + } +}; + +void SendAndCheckUBFallbackRPC(int fd, bool include_ack, uint64_t id) { + test::EchoRequest request; + request.set_message("UBSHM fallback RPC"); + brpc::Controller cntl; + butil::IOBuf body; + brpc::policy::SerializeRpcRequest(&body, &cntl, &request); + butil::IOBuf packet; + brpc::policy::PackRpcRequest(&packet, nullptr, id, + test::EchoService::descriptor()->FindMethodByName("Echo"), + &cntl, body, nullptr); + ASSERT_FALSE(cntl.Failed()); + std::string bytes(include_ack ? 4 : 0, '\0'); + bytes += packet.to_string(); + ASSERT_TRUE(TransferUBTestBytes(fd, &bytes[0], bytes.size(), true)); + uint8_t header[12]; + ASSERT_TRUE(TransferUBTestBytes(fd, header, sizeof(header), false)); + ASSERT_EQ(0, memcmp(header, "PRPC", 4)); + uint32_t total_size; + uint32_t meta_size; + memcpy(&total_size, header + 4, sizeof(total_size)); + memcpy(&meta_size, header + 8, sizeof(meta_size)); + total_size = butil::NetToHost32(total_size); + meta_size = butil::NetToHost32(meta_size); + ASSERT_GT(total_size, 0u); + ASSERT_LT(total_size, 65536u); + ASSERT_LE(meta_size, total_size); + std::string response(total_size, '\0'); + ASSERT_TRUE(TransferUBTestBytes(fd, &response[0], response.size(), false)); + brpc::policy::RpcMeta meta; + ASSERT_TRUE(meta.ParseFromArray(response.data(), meta_size)); + ASSERT_EQ(id, meta.correlation_id()); + ASSERT_EQ(0, meta.response().error_code()); + test::EchoResponse echo; + ASSERT_TRUE(echo.ParseFromArray(response.data() + meta_size, + total_size - meta_size)); + ASSERT_EQ(request.message(), echo.message()); +} +} // namespace + +TEST_F(UBShmEndpointTest, failed_server_allocation_cleans_resources_and_serves_tcp) { + UBFallbackEchoService service; + brpc::Server server; + ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + options.enabled_protocols = "baidu_std"; + options.internal_port = -1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &options)); + + // The pin is destroyed before Server on any assertion failure. + brpc::SocketUniquePtr socket; + butil::fd_guard peer(::socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_GE(peer, 0); + timeval timeout = {5, 0}; + ASSERT_EQ(0, setsockopt(peer, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + addr.sin_port = htons(server.listen_address().port); + ASSERT_EQ(0, connect(peer, reinterpret_cast(&addr), sizeof(addr))); + ASSERT_TRUE(WaitForUBCondition([&] { + std::vector ids; + server._am->ListConnections(&ids); + return !ids.empty() && brpc::Socket::Address(ids[0], &socket) == 0; + })); + + // A valid Hello references absent remote SHM. Allocation creates a ring + // and poller socket before mapping fails, exercising partial cleanup. + std::string missing_name = "UBRING_missing_" + std::to_string(getpid()); + butil::fd_guard absent(shm_open(missing_name.c_str(), O_RDONLY, 0)); + ASSERT_LT(absent, 0); + ASSERT_EQ(ENOENT, errno); + brpc::ubring::UBShmHandshakeAdapter protocol; + protocol.ConfigureClientHello(4 * 1024 * 1024, missing_name.c_str()); + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, protocol.BuildHello(true, &payload)); + std::string hello; + ASSERT_EQ(brpc::handshake::FRAME_OK, brpc::handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), payload, &hello)); + ASSERT_TRUE(TransferUBTestBytes(peer, &hello[0], hello.size(), true)); + char reply[64]; + ASSERT_TRUE(TransferUBTestBytes(peer, reply, sizeof(reply), false)); + brpc::ubring::HelloMessage disabled{}; + disabled.Deserialize(reply + 2); + ASSERT_EQ(0u, disabled.len); + + ASSERT_NO_FATAL_FAILURE(SendAndCheckUBFallbackRPC(peer, true, 101)); + auto* adapter = brpc::AdapterTransport::Get(socket.get()); + ASSERT_TRUE(WaitForUBCondition([&] { + return adapter->handshake_phase() == brpc::handshake::FALLBACK_TCP; + })); + auto* transport = brpc::UBShmTransport::Get(socket.get()); + EXPECT_FALSE(transport->UpgradeActive()); + EXPECT_EQ(nullptr, transport->GetUBShmEp()->_ub_ring); + EXPECT_EQ(brpc::INVALID_SOCKET_ID, transport->GetUBShmEp()->_poller_sid); + EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, + transport->GetUBShmEp()->negotiated_data_format()); + ASSERT_NO_FATAL_FAILURE(SendAndCheckUBFallbackRPC(peer, false, 102)); + EXPECT_EQ(brpc::handshake::FALLBACK_TCP, adapter->handshake_phase()); + EXPECT_FALSE(socket->Failed()); + peer.reset(-1); + socket.reset(); + server.Stop(0); + server.Join(); +} + +TEST_F(UBShmEndpointTest, malformed_hello_releases_allocated_ring_and_poller) { + brpc::SocketOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(id, &socket)); + BRPC_SCOPE_EXIT { socket->SetFailed(); }; + auto* transport = brpc::UBShmTransport::Get(socket.get()); + brpc::ubring::SHM local{}; + local.len = 4 * 1024 * 1024; + const std::string name = "cleanup_" + std::to_string(getpid()); + ASSERT_EQ(0, transport->PrepareUpgradeResources(&local, name.c_str())); + auto* ep = transport->GetUBShmEp(); + ASSERT_NE(nullptr, ep->_ub_ring); + const brpc::SocketId poller = ep->_poller_sid; + ASSERT_NE(brpc::INVALID_SOCKET_ID, poller); + + brpc::ubring::HelloMessage hello{}; + hello.msg_len = 64; + hello.hello_ver = 3; + hello.impl_ver = 1; + hello.len = 4 * 1024 * 1024; + memset(hello.shm_name, 'x', sizeof(hello.shm_name)); + char frame[64] = {'U', 'B'}; + hello.Serialize(frame + 2); + butil::IOBuf source; + source.append(frame, sizeof(frame)); + EXPECT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, + brpc::handshake::GetUBShmServerHandshakeAdapter() + ->ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_EQ(brpc::handshake::FAILED, + brpc::AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_FALSE(transport->UpgradeActive()); + EXPECT_EQ(nullptr, ep->_ub_ring); + EXPECT_EQ(brpc::INVALID_SOCKET_ID, ep->_poller_sid); + brpc::SocketUniquePtr old_poller; + EXPECT_NE(0, brpc::Socket::Address(poller, &old_poller)); + butil::fd_guard removed(shm_open(local.name, O_RDONLY, 0)); + EXPECT_LT(removed, 0); + EXPECT_EQ(ENOENT, errno); +} + +TEST_F(UBShmEndpointTest, shared_memory_rpc_waits_for_final_tcp_ack) { + UBFallbackEchoService service; + brpc::Server server; + ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + options.enabled_protocols = "baidu_std"; + options.internal_port = -1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &options)); + + brpc::SocketUniquePtr server_socket; + brpc::SocketUniquePtr client_socket; + butil::fd_guard peer(::socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_GE(peer, 0); + timeval timeout = {5, 0}; + ASSERT_EQ(0, setsockopt(peer, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + addr.sin_port = htons(server.listen_address().port); + ASSERT_EQ(0, connect(peer, reinterpret_cast(&addr), sizeof(addr))); + ASSERT_TRUE(WaitForUBCondition([&] { + std::vector ids; + server._am->ListConnections(&ids); + return !ids.empty() && brpc::Socket::Address(ids[0], &server_socket) == 0; + })); + + brpc::SocketOptions client_options; + client_options.socket_mode = brpc::SOCKET_MODE_UBRING; + brpc::SocketId client_id; + ASSERT_EQ(0, brpc::Socket::Create(client_options, &client_id)); + ASSERT_EQ(0, brpc::Socket::Address(client_id, &client_socket)); + BRPC_SCOPE_EXIT { client_socket->SetFailed(); }; + auto* client_transport = brpc::UBShmTransport::Get(client_socket.get()); + brpc::ubring::SHM local{}; + local.len = 4 * 1024 * 1024; + local.fd = peer; + const std::string name = "ack_race_" + std::to_string(getpid()); + ASSERT_EQ(0, client_transport->PrepareUpgradeResources(&local, name.c_str())); + brpc::ubring::UBShmHandshakeAdapter protocol; + protocol.ConfigureClientHello(local); + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, protocol.BuildHello(true, &payload)); + std::string hello; + ASSERT_EQ(brpc::handshake::FRAME_OK, brpc::handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), payload, &hello)); + ASSERT_TRUE(TransferUBTestBytes(peer, &hello[0], hello.size(), true)); + char reply[64]; + ASSERT_TRUE(TransferUBTestBytes(peer, reply, sizeof(reply), false)); + brpc::ubring::HelloMessage server_hello{}; + server_hello.Deserialize(reply + 2); + ASSERT_GT(server_hello.len, 0u); + std::string extension_payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + protocol.BuildExtension(true, &extension_payload)); + std::string extension; + ASSERT_EQ(brpc::handshake::FRAME_OK, brpc::handshake::FrameCodec::Encode( + protocol.ExtensionFrameSpec(), extension_payload, &extension)); + ASSERT_TRUE(TransferUBTestBytes(peer, &extension[0], extension.size(), true)); + char extension_reply[brpc::ubring::HelloFormatExtension::WIRE_SIZE]; + ASSERT_TRUE(TransferUBTestBytes(peer, extension_reply, + sizeof(extension_reply), false)); + auto* adapter = brpc::AdapterTransport::Get(server_socket.get()); + ASSERT_TRUE(WaitForUBCondition([&] { + return adapter->handshake_phase() == brpc::handshake::ACK_WAIT; + })); + ASSERT_EQ(0, client_transport->NegotiateUpgradeResources(&local, name.c_str())); + + test::EchoRequest request; + request.set_message("RPC queued before the final ACK"); + brpc::Controller cntl; + butil::IOBuf body; + brpc::policy::SerializeRpcRequest(&body, &cntl, &request); + butil::IOBuf packet; + brpc::policy::PackRpcRequest(&packet, nullptr, 201, + test::EchoService::descriptor()->FindMethodByName("Echo"), + &cntl, body, nullptr); + ASSERT_FALSE(cntl.Failed()); + const std::string bytes = packet.to_string(); + iovec data{const_cast(bytes.data()), bytes.size()}; + auto* client_ring = client_transport->GetUBShmEp()->_ub_ring; + ASSERT_EQ(static_cast(bytes.size()), client_ring->UbrTrxWritev(&data, 1)); + + auto* server_ep = brpc::UBShmTransport::Get(server_socket.get())->GetUBShmEp(); + EXPECT_FALSE(server_ep->_receive_events_started.load(butil::memory_order_acquire)); + // Invoke a would-be early poll synchronously while ACK is deliberately + // withheld. It must not feed PRPC bytes into the TCP handshake context. + ASSERT_NE(nullptr, server_socket->parsing_context()); + brpc::ubring::UBShmEndpoint::PollIn(server_ep, EPOLLIN); + EXPECT_EQ(brpc::handshake::ACK_WAIT, adapter->handshake_phase()); + EXPECT_FALSE(server_socket->Failed()); + + uint32_t ack = butil::HostToNet32(1); + ASSERT_TRUE(TransferUBTestBytes(peer, &ack, sizeof(ack), true)); + // No second SHM write: registering level-triggered receive polling must + // pick up the request that was already queued before the ACK. + butil::IOPortal response; + ASSERT_TRUE(WaitForUBCondition([&] { + const ssize_t n = response.append_from_reader(client_ring, 65536); + if (n < 0 && errno != EAGAIN && errno != EINTR) { + return true; + } + if (response.size() < 12) { + return false; + } + uint8_t header[12]; + response.copy_to(header, sizeof(header)); + uint32_t total_size; + memcpy(&total_size, header + 4, sizeof(total_size)); + return response.size() >= 12 + butil::NetToHost32(total_size); + })); + const std::string received = response.to_string(); + ASSERT_GE(received.size(), 12u); + ASSERT_EQ(0, memcmp(received.data(), "PRPC", 4)); + uint32_t meta_size; + memcpy(&meta_size, received.data() + 8, sizeof(meta_size)); + meta_size = butil::NetToHost32(meta_size); + ASSERT_LE(meta_size, received.size() - 12); + brpc::policy::RpcMeta meta; + ASSERT_TRUE(meta.ParseFromArray(received.data() + 12, meta_size)); + EXPECT_EQ(201u, meta.correlation_id()); + EXPECT_EQ(0, meta.response().error_code()); + test::EchoResponse echo; + ASSERT_TRUE(echo.ParseFromArray(received.data() + 12 + meta_size, + received.size() - 12 - meta_size)); + EXPECT_EQ(request.message(), echo.message()); + EXPECT_EQ(brpc::handshake::ESTABLISHED, adapter->handshake_phase()); + EXPECT_TRUE(server_ep->_receive_events_started.load(butil::memory_order_acquire)); + client_transport->DeactivateUpgrade(); + peer.reset(-1); + server_socket.reset(); + server.Stop(0); + server.Join(); +} + +TEST_F(UBShmEndpointTest, upgraded_client_receives_shared_memory_rpc_responses) { + UBFallbackEchoService service; + brpc::Server server; + ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions server_options; + server_options.socket_mode = brpc::SOCKET_MODE_UBRING; + server_options.enabled_protocols = "baidu_std"; + server_options.internal_port = -1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &server_options)); + { + brpc::Channel channel; + brpc::ChannelOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + options.timeout_ms = 5000; + ASSERT_EQ(0, channel.Init(server.listen_address(), &options)); + test::EchoService_Stub stub(&channel); + for (int i = 0; i < 2; ++i) { + brpc::Controller cntl; + test::EchoRequest request; + request.set_message(std::string(1024, 'a' + i)); + test::EchoResponse response; + stub.Echo(&cntl, &request, &response, nullptr); + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ(request.message(), response.message()); + } + std::vector ids; + server._am->ListConnections(&ids); + ASSERT_EQ(1u, ids.size()); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(ids[0], &socket)); + EXPECT_EQ(brpc::handshake::ESTABLISHED, + brpc::AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_TRUE(brpc::UBShmTransport::Get(socket.get())->UpgradeActive()); + } + server.Stop(0); + server.Join(); +} + // The receive paths (UbrTrxRecvBlockMode / StartReadv) read `msg_len' and // `cur_index' out of a chunk header the remote peer writes into the ring, then // copy `msg_len - cur_index' bytes from the 60-byte `payload.inner'. A peer From 2e5136150a7ca3bc7f35c2d61abeede363f6da5e Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Thu, 15 Oct 2026 19:58:53 -0700 Subject: [PATCH 02/10] Fix certificate name extraction with const OpenSSL accessors --- src/brpc/details/ssl_helper.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/brpc/details/ssl_helper.cpp b/src/brpc/details/ssl_helper.cpp index 1673db1d87..25ee4c6a83 100644 --- a/src/brpc/details/ssl_helper.cpp +++ b/src/brpc/details/ssl_helper.cpp @@ -287,10 +287,10 @@ void ExtractHostnames(X509* x, std::vector* hostnames) { #endif // SSL_CTRL_SET_TLSEXT_HOSTNAME int i = -1; - X509_NAME* xname = X509_get_subject_name(x); + auto* xname = X509_get_subject_name(x); while ((i = X509_NAME_get_index_by_NID(xname, NID_commonName, i)) != -1) { char* str = nullptr; - X509_NAME_ENTRY* entry = X509_NAME_get_entry(xname, i); + auto* entry = X509_NAME_get_entry(xname, i); const int len = ASN1_STRING_to_UTF8((unsigned char**)&str, X509_NAME_ENTRY_get_data(entry)); if (len >= 0) { From d1217db88845b308767af082a0e661071eae0ed8 Mon Sep 17 00:00:00 2001 From: Xu Lei <98076253+AnDiXL@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:18:32 +0800 Subject: [PATCH 03/10] Add optional FlatBuffers messages and code generation (#3558) * Add FlatBuffers message and service support Add IOBuf-backed messages, allocator-aware builders, and stable service descriptors. Preserve shared-buffer copy semantics and validate allocator, parsing, descriptor, and ownership contracts with the core unit suite. Keep FlatBuffers 25.2.10 support disabled by default across CMake, Make, and Bazel. Isolate generated test headers and select matching GoogleTest headers and libraries. Network transport and Channel/Server integration remain in the follow-up PR. * Add FlatBuffers service generator and acceptance tests Generate service and stub bindings with the upstream FlatBuffers parser. Match flatc type escaping, preserve schema service identities, and reject collisions with schema declarations, generated API names, and reserved C++ identifiers. Publish header/source pairs with locking and rollback. Cover imported and nested schemas, compiler modes, relocation, concurrent generation, failure recovery, and real runtime dispatch. Pass compiler flags through acceptance checks and pin validation to flatc 25.2.10. * Verify FlatBuffers support in Linux and macOS CI Run the core unit suite, codegen acceptance, and runtime tests through a shared verification entry. Pin FlatBuffers and GoogleTest sources and reject empty or skipped core test runs. Exercise GCC 11 and Clang 14 on Linux and Protobuf 21 and 29 on macOS. Build checksum-pinned Gperftools in the Linux runner temporary directory without libunwind to avoid conflicts with hosted runner packages. * Document FlatBuffers support Add Chinese and English guides for messages, builders, descriptors, and service generation. Describe shared-buffer ownership, verification, allocation failures, supported dependency versions, and build options. Document generator naming rules, failure recovery, acceptance tests, and the shared CI verification entry. Keep network transport and Channel/Server integration explicitly outside this PR. * Document using other FlatBuffers versions Describe how to move off the default 25.2.10: build the target version into a private prefix, update the runtime, generator, and codegen test version checks, and rerun verify.sh against that prefix. Use 23.5.26 as the worked example and limit it to CMake builds. List FlatBuffers among the supported dependencies in the getting started guides, following the format used for protobuf and libunwind. --------- Co-authored-by: xulei25 --- .github/workflows/ci-linux.yml | 36 + .github/workflows/ci-macos.yml | 28 + BUILD.bazel | 27 +- CMakeLists.txt | 1 + MODULE.bazel | 16 + Makefile | 3 + WORKSPACE | 8 + config_brpc.sh | 19 +- docs/cn/flatbuffers.md | 116 +++ docs/cn/getting_started.md | 4 + docs/en/flatbuffers.md | 137 ++++ docs/en/getting_started.md | 4 + src/brpc/CMakeLists.txt | 11 + src/brpc/flatbuffers/message.cpp | 310 ++++++++ src/brpc/flatbuffers/message.h | 187 +++++ src/brpc/flatbuffers/service.cpp | 158 ++++ src/brpc/flatbuffers/service.h | 133 ++++ src/butil/CMakeLists.txt | 6 + src/butil/config.h.in | 5 + test/BUILD.bazel | 42 +- test/CMakeLists.txt | 51 +- test/Makefile | 28 +- test/brpc_flatbuffers_unittest.cpp | 565 ++++++++++++++ test/flatbuffers_codegen/CMakeLists.txt | 148 ++++ test/flatbuffers_codegen/acceptance.cmake | 809 ++++++++++++++++++++ test/flatbuffers_codegen/echo.fbs | 50 ++ test/flatbuffers_codegen/runtime.cpp | 339 ++++++++ test/flatbuffers_message.fbs | 26 + tools/flatbuffers/CMakeLists.txt | 46 ++ tools/flatbuffers/README.md | 243 ++++++ tools/flatbuffers/brpc_flatc.cpp | 894 ++++++++++++++++++++++ tools/flatbuffers/verify.sh | 138 ++++ 32 files changed, 4572 insertions(+), 16 deletions(-) create mode 100644 docs/cn/flatbuffers.md create mode 100644 docs/en/flatbuffers.md create mode 100644 src/brpc/flatbuffers/message.cpp create mode 100644 src/brpc/flatbuffers/message.h create mode 100644 src/brpc/flatbuffers/service.cpp create mode 100644 src/brpc/flatbuffers/service.h create mode 100644 test/brpc_flatbuffers_unittest.cpp create mode 100644 test/flatbuffers_codegen/CMakeLists.txt create mode 100644 test/flatbuffers_codegen/acceptance.cmake create mode 100644 test/flatbuffers_codegen/echo.fbs create mode 100644 test/flatbuffers_codegen/runtime.cpp create mode 100644 test/flatbuffers_message.fbs create mode 100644 tools/flatbuffers/CMakeLists.txt create mode 100644 tools/flatbuffers/README.md create mode 100644 tools/flatbuffers/brpc_flatc.cpp create mode 100644 tools/flatbuffers/verify.sh diff --git a/.github/workflows/ci-linux.yml b/.github/workflows/ci-linux.yml index 341c53d392..9dcadae328 100644 --- a/.github/workflows/ci-linux.yml +++ b/.github/workflows/ci-linux.yml @@ -19,6 +19,42 @@ concurrency: # https://github.com/actions/runner-images jobs: + flatbuffers-codegen: + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + cxx: [g++-11, clang++-14] + steps: + - uses: actions/checkout@v7.0.1 + - name: Install codegen dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential cmake g++-11 clang-14 \ + libprotobuf-dev libprotoc-dev protobuf-compiler \ + libgflags-dev libleveldb-dev libssl-dev zlib1g-dev + - name: Build isolated Gperftools + env: + CXX: ${{ matrix.cxx }} + run: | + archive="$RUNNER_TEMP/gperftools-2.16.tar.gz" + curl --fail --location --retry 3 \ + https://github.com/gperftools/gperftools/releases/download/gperftools-2.16/gperftools-2.16.tar.gz \ + --output "$archive" + printf '%s %s\n' f12624af5c5987f2cc830ee534f754c3c5961eec08004c26a8b80de015cf056f "$archive" | sha256sum --check + tar -xzf "$archive" -C "$RUNNER_TEMP" + cd "$RUNNER_TEMP/gperftools-2.16" + ./configure --prefix="$RUNNER_TEMP/gperftools-install" \ + --enable-frame-pointers --disable-libunwind + make -j2 + make install + printf 'CMAKE_PREFIX_PATH=%s\n' "$RUNNER_TEMP/gperftools-install" >> "$GITHUB_ENV" + - name: Verify FlatBuffers generation and runtime + env: + CXX: ${{ matrix.cxx }} + JOBS: 2 + run: bash tools/flatbuffers/verify.sh "$RUNNER_TEMP/flatbuffers-codegen" + compile-with-make: name: compile-with-make (${{ matrix.name }}) runs-on: ubuntu-22.04 diff --git a/.github/workflows/ci-macos.yml b/.github/workflows/ci-macos.yml index a4e4272afe..0de86024b4 100644 --- a/.github/workflows/ci-macos.yml +++ b/.github/workflows/ci-macos.yml @@ -18,6 +18,34 @@ concurrency: cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: + flatbuffers-codegen: + runs-on: macos-latest + strategy: + fail-fast: false + matrix: + protobuf: [protobuf@21, protobuf@29] + steps: + - uses: actions/checkout@v7.0.1 + - name: Install codegen dependencies + run: | + brew install cmake openssl gflags leveldb snappy gperftools + if [ "${{ matrix.protobuf }}" = protobuf@21 ]; then + # Same local tap as compile-with-make-cmake-protobuf21. + brew tap-new --no-git local/pins + cp homebrew-formula/protobuf@21.rb "$(brew --repository local/pins)/Formula/" + brew install local/pins/protobuf@21 + else + brew install ${{ matrix.protobuf }} + fi + - name: Verify FlatBuffers generation and runtime + env: + CXX: clang++ + JOBS: 2 + run: | + export CMAKE_PREFIX_PATH="$(brew --prefix ${{ matrix.protobuf }});$(brew --prefix)" + export OPENSSL_ROOT_DIR="$(brew --prefix openssl)" + bash tools/flatbuffers/verify.sh "$RUNNER_TEMP/flatbuffers-codegen" + compile-with-make-cmake-protobuf21: runs-on: macos-latest # https://github.com/actions/runner-images diff --git a/BUILD.bazel b/BUILD.bazel index 23e433c76e..5882a034da 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -20,6 +20,12 @@ licenses(["notice"]) # Apache v2 exports_files(["LICENSE"]) +config_setting( + name = "brpc_with_flatbuffers", + define_values = {"BRPC_WITH_FLATBUFFERS": "true"}, + visibility = ["//visibility:public"], +) + COPTS = [ "-fno-omit-frame-pointer", ] + select({ @@ -45,6 +51,9 @@ DEFINES = [ }) + select({ "//bazel/config:brpc_with_thrift": ["ENABLE_THRIFT_FRAMED_PROTOCOL=1"], "//conditions:default": [], + }) + select({ + ":brpc_with_flatbuffers": ["BRPC_WITH_FLATBUFFERS=1"], + "//conditions:default": ["BRPC_WITH_FLATBUFFERS=0"], }) + select({ "//bazel/config:brpc_with_thrift_legacy_version": [], "//conditions:default": ["THRIFT_STDCXX=std"], @@ -125,6 +134,14 @@ genrule( "//conditions:default": "0", }) + """ +#ifdef BRPC_WITH_FLATBUFFERS +#undef BRPC_WITH_FLATBUFFERS +#endif +#define BRPC_WITH_FLATBUFFERS """ + select({ + ":brpc_with_flatbuffers": "1", + "//conditions:default": "0", + }) + + """ #ifdef BUTIL_USE_CPU_FREQUENCY #undef BUTIL_USE_CPU_FREQUENCY #endif @@ -542,6 +559,8 @@ brpc_proto_library( visibility = ["//visibility:public"], ) +FLATBUFFERS_SRC_PATTERNS = ["src/brpc/flatbuffers/*.cpp"] + URMA_SRC_PATTERNS = [ "src/brpc/urma/*.cpp", "src/brpc/urma/**/*.cpp", @@ -565,7 +584,7 @@ BRPC_BASE_SRCS = glob( "src/brpc/policy/thrift_protocol.cpp", "src/brpc/event_dispatcher_epoll.cpp", "src/brpc/event_dispatcher_kqueue.cpp", - ] + URMA_SRC_PATTERNS, + ] + URMA_SRC_PATTERNS + FLATBUFFERS_SRC_PATTERNS, ) cc_library( @@ -576,6 +595,9 @@ cc_library( "src/brpc/**/thrift*.cpp", ]), "//conditions:default": [], + }) + select({ + ":brpc_with_flatbuffers": glob(FLATBUFFERS_SRC_PATTERNS), + "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_urma_use_real": URMA_SRCS, "//bazel/config:brpc_with_urma": URMA_SRCS + [ @@ -615,6 +637,9 @@ cc_library( "@org_apache_thrift//:thrift", ], "//conditions:default": [], + }) + select({ + ":brpc_with_flatbuffers": ["@com_github_google_flatbuffers//:runtime_cc"], + "//conditions:default": [], }), ) diff --git a/CMakeLists.txt b/CMakeLists.txt index 59acdfe3bc..0966595d13 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -23,6 +23,7 @@ option(WITH_MESALINK "With MesaLink" OFF) option(WITH_BORINGSSL "With BoringSSL" OFF) option(WITH_DEBUG_SYMBOLS "With debug symbols" ON) option(WITH_THRIFT "With thrift framed protocol supported" OFF) +option(WITH_FLATBUFFERS "With FlatBuffers message support (headers only)" OFF) option(WITH_BTHREAD_TRACER "With bthread tracer supported" OFF) option(WITH_SNAPPY "With snappy" OFF) option(WITH_RDMA "With RDMA" OFF) diff --git a/MODULE.bazel b/MODULE.bazel index 97862d36f1..abe905655d 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -79,3 +79,19 @@ git_repository( remote = 'https://atomgit.com/openeuler/umdk.git', commit = '564ee727a55523d4351a8fb3c94292b388ebb924', # v26.06.0_CAM ) + +# Do not use `bazel_dep(name = 'flatbuffers', version = '25.2.10')` here. +# The BCR module pulls gRPC and JS/Go/Swift rule dependencies for FlatBuffers' +# full upstream build. bRPC only needs runtime_cc and flatc, and gRPC would +# otherwise introduce another BoringSSL version even when BRPC_WITH_FLATBUFFERS +# is false. Keep this archive and checksum in sync with WORKSPACE. +flatbuffers_http_archive = use_repo_rule( + '@bazel_tools//tools/build_defs/repo:http.bzl', + 'http_archive', +) +flatbuffers_http_archive( + name = 'com_github_google_flatbuffers', + sha256 = 'b9c2df49707c57a48fc0923d52b8c73beb72d675f9d44b2211e4569be40a7421', + strip_prefix = 'flatbuffers-25.2.10', + urls = ['https://github.com/google/flatbuffers/archive/refs/tags/v25.2.10.tar.gz'], +) diff --git a/Makefile b/Makefile index 271b518ae6..08df1f04e9 100644 --- a/Makefile +++ b/Makefile @@ -204,6 +204,9 @@ JSON2PB_SOURCES = $(foreach d,$(JSON2PB_DIRS),$(wildcard $(addprefix $(d)/*,$(SR JSON2PB_OBJS = $(addsuffix .o, $(basename $(JSON2PB_SOURCES))) BRPC_DIRS = src/brpc src/brpc/details src/brpc/builtin src/brpc/policy src/brpc/policy/mysql src/brpc/rdma +ifeq ($(WITH_FLATBUFFERS),1) +BRPC_DIRS += src/brpc/flatbuffers +endif ifeq ($(WITH_URMA),1) BRPC_DIRS += src/brpc/urma endif diff --git a/WORKSPACE b/WORKSPACE index 22fc411b32..9778a704ba 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -150,6 +150,14 @@ http_archive( urls = ["https://github.com/google/crc32c/archive/1.1.2.tar.gz"], ) +# Optional FlatBuffers support uses runtime_cc; keep this version in sync with MODULE.bazel. +http_archive( + name = "com_github_google_flatbuffers", + integrity = "sha256-ucLfSXB8V6SPwJI9UrjHO+ty1nX51EsiEeRWm+QKdCE=", + strip_prefix = "flatbuffers-25.2.10", + urls = ["https://github.com/google/flatbuffers/archive/refs/tags/v25.2.10.tar.gz"], +) + http_archive( name = "com_github_google_glog", # 2021-05-07T23:06:39Z patch_args = ["-p1"], diff --git a/config_brpc.sh b/config_brpc.sh index 0efa4d374d..e940e486f4 100755 --- a/config_brpc.sh +++ b/config_brpc.sh @@ -54,9 +54,10 @@ else LDD=ldd fi -TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-urma,with-urma-mock,without-urma-mock,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"` +TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-flatbuffers,with-rdma,with-urma,with-urma-mock,without-urma-mock,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"` WITH_GLOG=0 WITH_THRIFT=0 +WITH_FLATBUFFERS=0 WITH_RDMA=0 WITH_URMA=0 URMA_MOCK_MODE=auto @@ -91,6 +92,7 @@ while true; do --cxx ) CXX=$2; shift 2 ;; --with-glog ) WITH_GLOG=1; shift 1 ;; --with-thrift) WITH_THRIFT=1; shift 1 ;; + --with-flatbuffers) WITH_FLATBUFFERS=1; shift 1 ;; --with-rdma) WITH_RDMA=1; shift 1 ;; --with-urma) WITH_URMA=1; shift 1 ;; --with-urma-mock) URMA_MOCK_MODE=on; shift 1 ;; @@ -479,6 +481,7 @@ append_to_output "HDRS=$($ECHO $HDRS)" append_to_output "LIBS=$($ECHO $LIBS)" append_to_output "PROTOC=$PROTOC" append_to_output "PROTOBUF_HDR=$PROTOBUF_HDR" +append_to_output "WITH_FLATBUFFERS=$WITH_FLATBUFFERS" append_to_output "CC=$CC" append_to_output "CXX=$CXX" append_to_output "GCC_VERSION=$GCC_VERSION" @@ -486,7 +489,7 @@ append_to_output "STATIC_LINKINGS=$STATIC_LINKINGS" append_to_output "DYNAMIC_LINKINGS=$DYNAMIC_LINKINGS" # CPP means C PreProcessing, not C PlusPlus -CPPFLAGS="${CPPFLAGS} -DBRPC_WITH_GLOG=$WITH_GLOG -DBRPC_DEBUG_BTHREAD_SCHE_SAFETY=$BRPC_DEBUG_BTHREAD_SCHE_SAFETY -DBRPC_DEBUG_LOCK=$BRPC_DEBUG_LOCK -DBUTIL_USE_CPU_FREQUENCY=$WITH_CPU_FREQUENCY" +CPPFLAGS="${CPPFLAGS} -DBRPC_WITH_GLOG=$WITH_GLOG -DBRPC_WITH_FLATBUFFERS=$WITH_FLATBUFFERS -DBRPC_DEBUG_BTHREAD_SCHE_SAFETY=$BRPC_DEBUG_BTHREAD_SCHE_SAFETY -DBRPC_DEBUG_LOCK=$BRPC_DEBUG_LOCK -DBUTIL_USE_CPU_FREQUENCY=$WITH_CPU_FREQUENCY" # Avoid over-optimizations of TLS variables by GCC>=4.8 # See: https://github.com/apache/brpc/issues/1693 @@ -506,6 +509,12 @@ if [ "$SYSTEM" = "Darwin" ]; then fi fi +if [ $WITH_FLATBUFFERS != 0 ]; then + FLATBUFFERS_HDR=$(find_dir_of_header_or_die flatbuffers/flatbuffers.h) || exit 1 + append_to_output_headers "$FLATBUFFERS_HDR" + print_success "Found FlatBuffers headers: $FLATBUFFERS_HDR" +fi + if [ $WITH_THRIFT != 0 ]; then THRIFT_LIB=$(find_dir_of_lib_or_die thriftnb) THRIFT_HDR=$(find_dir_of_header_or_die thrift/Thrift.h) @@ -692,6 +701,11 @@ cat << EOF > src/butil/config.h #endif #define BRPC_WITH_GLOG $WITH_GLOG +#ifdef BRPC_WITH_FLATBUFFERS +#undef BRPC_WITH_FLATBUFFERS +#endif +#define BRPC_WITH_FLATBUFFERS $WITH_FLATBUFFERS + #ifdef BUTIL_USE_CPU_FREQUENCY #undef BUTIL_USE_CPU_FREQUENCY #endif @@ -714,6 +728,7 @@ print_info "C++ std: $CXXFLAGS" print_info "System: $SYSTEM" if [ $WITH_GLOG -ne 0 ]; then print_info "With glog: yes"; fi if [ $WITH_THRIFT -ne 0 ]; then print_info "With thrift: yes"; fi +if [ $WITH_FLATBUFFERS -ne 0 ]; then print_info "With FlatBuffers: yes (headers only)"; fi if [ $WITH_RDMA -ne 0 ]; then print_info "With RDMA: yes"; fi if [ $WITH_URMA -ne 0 ]; then print_info "With URMA: yes"; fi if [ $WITH_MESALINK -ne 0 ]; then print_info "With MesaLink: yes"; fi diff --git a/docs/cn/flatbuffers.md b/docs/cn/flatbuffers.md new file mode 100644 index 0000000000..5e58801594 --- /dev/null +++ b/docs/cn/flatbuffers.md @@ -0,0 +1,116 @@ +# FlatBuffers 消息 + +[English version](../en/flatbuffers.md) + +bRPC 支持基于 IOBuf 的 FlatBuffers 消息、构造器和服务描述符，默认不启用。 +消息构造方案源自 [apache/brpc#3196](https://github.com/apache/brpc/pull/3196)。 + +本组件不包含 `fb_rpc` 传输协议，也不改动 `brpc::Channel` 或 `brpc::Server`。 +服务代码生成和进程内分派测试不涉及网络 RPC，也不是性能测试。 + +## 构建 + +运行库只需要 FlatBuffers 头文件，不链接其库；测试还需要版本匹配的 `flatc`。 +版本不一致时应重新生成头文件，不要删除或放宽上游生成代码中的版本断言。 + +假设 GoogleTest 源码位于 `/usr/src/googletest`： + +```sh +cmake -S . -B build -DWITH_FLATBUFFERS=ON -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BRPC_TOOLS=OFF -DDOWNLOAD_GTEST=OFF \ + -DBRPC_SYSTEM_GTEST_SOURCE_DIR=/usr/src/googletest +cmake --build build --target brpc_flatbuffers_unittest -j6 +ctest --test-dir build -R '^brpc_flatbuffers_unittest$' --output-on-failure +``` + +非标准安装可设置 `FLATBUFFERS_INCLUDE_DIR`、`FLATBUFFERS_FLATC_EXECUTABLE` +和 `BRPC_SYSTEM_GTEST_SOURCE_DIR`。其他测试依赖与 bRPC 原有配置相同。 + +Make 使用 `config_brpc.sh --with-flatbuffers`，测试时可通过 +`FLATC=/path/to/flatc` 指定生成器。Bazel 使用 +`--define=BRPC_WITH_FLATBUFFERS=true`。 + +运行库使用了 FlatBuffers builder 的内部结构，目前只支持 25.2.10，其他版本会在 +编译期报错。Bzlmod 和 WORKSPACE 使用同一个源码归档，并用校验和锁定内容。 +未采用 `bazel_dep`，是为了避免引入 gRPC、多语言工具以及与 bRPC 冲突的 BoringSSL +依赖。 + +公共头文件位于 `brpc/flatbuffers/`，命名空间为 `brpc::flatbuffers`。 +构造消息时包含 `message.h`；使用服务描述符和接口时包含 `service.h`。`butil/config.h` 中的 +`BRPC_WITH_FLATBUFFERS` 始终是 0 或 1，因此应使用 `#if` 判断。 + +旧版 flatc 2.0.x 可能生成未全限定名称。业务 schema 不要使用 `brpc` 命名空间，也不要依赖 include 顺序。 + +如需使用其他 FlatBuffers 版本，需要修改运行库、生成器和 codegen 测试中的版本检查， +再用该版本完整运行 `tools/flatbuffers/verify.sh`。步骤见生成器 +[README](../../tools/flatbuffers/README.md#using-another-flatbuffers-version)， +其中以 23.5.26 为例，仅适用于 CMake 构建。 + +## 消息构造和所有权 + +先用上游 `flatc --cpp` 生成 `*_generated.h`，再将 +`brpc::flatbuffers::MessageBuilder` 传给生成的 `Create...` 函数。 +调用 `Finish(root)` 后，通过 `ReleaseMessage()` 取得消息。 + +* `ReleaseMessage()` 不复制 payload。返回的 Message 是 move-only 对象，持有 IOBuf block 引用；builder 复用或析构后，消息仍然有效。 +* Message 和 builder 被移动后，源对象仍可复用。移动启用了 shared string 的 builder 会清空去重缓存，但已生成的 offset 不受影响。 +* `Message::CopyFrom` 和 `MergeFrom` 共享同一个引用计数 IOBuf block，不复制 + payload 或 metadata。通过任一别名修改数据，其他别名也会看到变化。有并发读者时， + 写操作必须由外部同步保护。 +* 从普通 `::flatbuffers::FlatBufferBuilder` 导入时，会复制 payload 和 scratch， + 保留未完成 table 的状态。旧缓冲区由原分配器释放；若源 builder 拥有该分配器， + 导入后也会销毁它。导入过程不猜测应使用 `free` 还是 `delete[]`，也不要求 payload + 前有预留空间。 +* 只使用 MessageBuilder 自身提供的 move、swap 和 release 操作。不要通过基类 + cast 转移对象，也不要调用继承来的 raw-buffer release；这类 detached buffer + 会保留成员 allocator 的地址。 +* payload 前预留 64 字节并初始化为零。`reduce_meta_size_and_get_buf` 可以缩短 + 这段空间；扩大时返回失败，原对象保持不变。缩短 metadata 不会移动或改写 payload。 +* 序列化 const Message 时，输出 IOBuf 会继续引用原存储。它不是 copy-on-write；仍有读者或序列化结果存活时，不要修改 payload 或 metadata。 +* 分配长度在转成 SingleIOBuf 的 uint32_t 长度前会检查。SlabAllocator/builder + 的分配失败会终止进程，release build 也一样，不受 `crash_on_fatal_log` 影响。 + 这是因为 FlatBuffers 的 `vector_downward` 无法在分配器返回 null 后继续工作。 + 解析时复制数据所需的内存若分配失败，则返回 false，保留原 Message。 + +`ParseFbFromIOBUF` 负责检查长度和 framing，解析后的 Message 持有自己的存储引用， +但该接口不负责限制接收消息大小。若 `msg_size` 来自不可信对端，调用方必须先按 +自己的 max-message-size 策略检查，再调用解析接口。分片或未对齐的数据会在 +schema 校验前复制到新存储，因此大小限制必须放在解析之前。 + +连续输入的 payload 地址按 64 字节对齐时，解析可以共享原存储。MessageBuilder +的分配按 64 字节对齐，但最终 payload 只保证满足 schema 的对齐要求，所以本地 +构造的消息未必都能在接收侧零拷贝。当前不支持超过 64 字节的对齐要求。 + +**Framing 检查不等于 schema 校验。** 读取不可信数据前，先调用 +`msg.Verify()`，成功后再使用 `GetRoot()` 或 +`GetMutableRoot()`。合法消息中的可选 string/vector 仍可能为 null。 +Framing 检查失败不会修改原 Message。 + +## Service ID 和代码生成 + +`BrpcDescriptorTable` 保存 namespace、service 名称、按空白分隔的方法名和显式 +方法 ID。ID 必须是唯一的非负 int32。手写 descriptor 可以传空 ID 列表，此时按 +声明顺序分配 ordinal ID；生成服务必须显式声明 ID： + +```fbs +rpc_service BenchmarkService { + First(Request):Response (id: 2); + Second(Request):Response (id: 5); +} +``` + +* `descriptor.method(position)` 按声明顺序取方法。 +* `method.index()` 是稳定的 wire ID，不是数组下标。 +* 稀疏 ID 必须通过 `descriptor.FindMethodByIndex(id)` 查找。传输层不能直接用 wire ID 索引稠密数组。 +* 删除方法后，不要把原 ID 分配给其他方法。调整声明顺序也不应改变已有方法的 ID。 +* namespace `a.b` 和 `a.b.` 会规范化为同一个 service 名称；空 namespace 表示 + 全局作用域。方法全名包含 service 名称。service hash 使用规范化全名和 + MurmurHash3 seed 1。若 ID 需要持久化或在线路上传输，service 名称必须保持稳定。 +* Descriptor 初始化成功后不能再次初始化。它通过 RAII 持有 method；生成的 accessor 使用函数局部静态对象，初始化过程是线程安全的。 + +`tools/flatbuffers/` 中的生成器使用上游 parser 生成服务绑定，并单独链接 +`libflatbuffers`。命令和限制见其 [README](../../tools/flatbuffers/README.md)。 + +生成的 dispatch 会验证请求，拒绝未知、不属于当前 service 或尚未实现的方法。 +失败时，它会执行非空 completion callback；成功分派后，completion 由业务实现负责， +并且必须恰好执行一次。 diff --git a/docs/cn/getting_started.md b/docs/cn/getting_started.md index 2a3f98d166..391873127d 100644 --- a/docs/cn/getting_started.md +++ b/docs/cn/getting_started.md @@ -453,6 +453,10 @@ bRPC默认**不**链接 [libunwind](https://github.com/libunwind/libunwind)。 建议使用最新版本的libunwind。 +## flatbuffers: 25.2.10 + +可选依赖，默认不启用。bRPC 使用了 FlatBuffers builder 的内部接口，只接受该版本；使用其他版本见 [FlatBuffers 消息](flatbuffers.md)。 + # 实例追踪 我们提供了一个程序去帮助你追踪和监控所有brpc实例。 只需要在某处运行 [trackme_server](https://github.com/apache/brpc/tree/master/tools/trackme_server/) 然后再带着 -trackme_server=SERVER参数启动需要被追踪的实例。trackme_server将从实例周期性地收到ping消息然后打印日志。您可以从日志中聚合实例地址，并调用实例的内置服务以获取更多信息。 diff --git a/docs/en/flatbuffers.md b/docs/en/flatbuffers.md new file mode 100644 index 0000000000..50748900ad --- /dev/null +++ b/docs/en/flatbuffers.md @@ -0,0 +1,137 @@ +# FlatBuffers messages + +[中文版](../cn/flatbuffers.md) + +bRPC supports IOBuf-backed FlatBuffers messages, builders, and service +descriptors, disabled by default. Message construction is based on +[apache/brpc#3196](https://github.com/apache/brpc/pull/3196). + +This component does not include `fb_rpc` transport or change `brpc::Channel` +or `brpc::Server`. Service generation and in-process dispatch tests do not +exercise network RPC or measure performance. + +## Build + +The runtime needs only FlatBuffers headers; tests also need a matching `flatc`. +Regenerate incompatible headers rather than removing or weakening upstream +version assertions. + +With GoogleTest sources at `/usr/src/googletest`: + +```sh +cmake -S . -B build -DWITH_FLATBUFFERS=ON -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BRPC_TOOLS=OFF -DDOWNLOAD_GTEST=OFF \ + -DBRPC_SYSTEM_GTEST_SOURCE_DIR=/usr/src/googletest +cmake --build build --target brpc_flatbuffers_unittest -j6 +ctest --test-dir build -R '^brpc_flatbuffers_unittest$' --output-on-failure +``` + +For nonstandard installations, set `FLATBUFFERS_INCLUDE_DIR`, +`FLATBUFFERS_FLATC_EXECUTABLE`, and `BRPC_SYSTEM_GTEST_SOURCE_DIR` as needed. +The usual bRPC test dependencies still apply. + +For Make, use `config_brpc.sh --with-flatbuffers` and set `FLATC=/path/to/flatc` +when testing. For Bazel, use `--define=BRPC_WITH_FLATBUFFERS=true`. + +The runtime depends on FlatBuffers builder internals and supports only 25.2.10; +other versions fail at compile time. Bzlmod and WORKSPACE use the same +checksum-pinned archive. Using this instead of `bazel_dep` avoids gRPC, +other language tools, and conflicting BoringSSL dependencies. + +Public headers are in `brpc/flatbuffers/`, with namespace `brpc::flatbuffers`. +Include `message.h` for message construction and `service.h` for service +descriptors and interfaces. `BRPC_WITH_FLATBUFFERS` in `butil/config.h` is +0 or 1; test it with `#if`. + +Older flatc 2.0.x may emit unqualified names. Keep business schemas outside +the `brpc` namespace and do not rely on include order. + +To use another FlatBuffers version, update the version checks in the runtime, +the generator and the codegen tests, then run `tools/flatbuffers/verify.sh` +against that version. See the generator +[README](../../tools/flatbuffers/README.md#using-another-flatbuffers-version) +for the steps, using 23.5.26 as an example. They apply to CMake builds only. + +## Message construction and ownership + +Generate `*_generated.h` with upstream `flatc --cpp`. Pass a +`brpc::flatbuffers::MessageBuilder` to the generated `Create...` functions, +then call `Finish(root)` and `ReleaseMessage()`. + +* `ReleaseMessage()` returns a move-only Message without copying payload bytes. + Its IOBuf block reference survives builder reuse or destruction. +* Moves leave the source Message or builder reusable. Moving a shared-string + builder clears its deduplication cache but preserves existing offsets. +* `Message::CopyFrom` and `MergeFrom` share a ref-counted IOBuf block without + copying payload or metadata. Changes through one alias affect all others; + concurrent reads and writes require external synchronization. +* Importing a regular `::flatbuffers::FlatBufferBuilder` copies payload and + scratch, preserving unfinished tables. The original allocator frees the old + storage and is destroyed if owned by the source builder. Import neither + guesses between `free` and `delete[]` nor requires a payload prefix. +* Use only MessageBuilder's own move, swap, and release operations. Base-class + transfers and inherited raw-buffer releases can leave a detached buffer + pointing to the member allocator. +* Released payloads have a 64-byte zeroed prefix. `reduce_meta_size_and_get_buf` + can shrink it without changing the payload address or bytes. Attempts to grow + it fail without changing the message. +* Serializing a const Message shares its storage with the output IOBuf; it is + not copy-on-write. Do not modify payload or metadata while readers or + serialized buffers still use it. +* Sizes are checked before conversion to SingleIOBuf's uint32_t length. + Allocation failure in SlabAllocator or a builder terminates the process even + in release builds, regardless of `crash_on_fatal_log`: FlatBuffers' + `vector_downward` cannot continue after a null allocation. Allocation failure + when copying during parsing returns false and leaves the Message unchanged. + +`ParseFbFromIOBUF` checks lengths and framing and retains a storage reference; +it does not limit received message sizes. Bound untrusted `msg_size` with your +max-message-size policy before parsing: fragmented or unaligned input is copied +to new storage before schema validation. + +Parsing shares contiguous input when its payload is 64-byte aligned. Builder +allocations have this alignment, but finished payloads need only meet their +schema's alignment, so local messages may still require a copy. Alignment +requirements above 64 bytes are unsupported. + +**Framing is not schema verification.** Before reading untrusted data, call +`msg.Verify()` and use `GetRoot()` or +`GetMutableRoot()` only if it succeeds. Optional strings/vectors can +still be null in valid messages. Failed framing checks leave the Message intact. + +## Service IDs and generation + +`BrpcDescriptorTable` holds a namespace, service name, whitespace-separated +method names, and unique nonnegative int32 IDs. Handwritten descriptors may +use an empty ID list to assign IDs in declaration order; generated services +require explicit IDs: + +```fbs +rpc_service BenchmarkService { + First(Request):Response (id: 2); + Second(Request):Response (id: 5); +} +``` + +* `descriptor.method(position)` returns methods in declaration order. +* `method.index()` is the stable wire ID, not an array index. +* Look up sparse IDs with `descriptor.FindMethodByIndex(id)`, not by indexing + a dense array with a wire ID. +* Do not reuse removed method IDs or change existing IDs when reordering + methods. +* Namespaces `a.b` and `a.b.` normalize to the same service name; an empty + namespace means global scope. Method full names include the service name. + Service hashes use the normalized full name and MurmurHash3 seed 1. Keep + service names stable when persisting or transmitting IDs. +* Descriptors own methods through RAII and cannot be reinitialized after + success. + Generated accessors use thread-safe function-local static initialization. + +The generator in `tools/flatbuffers/` uses the upstream parser and links +`libflatbuffers` separately. See its [README](../../tools/flatbuffers/README.md) +for commands and limitations. + +Generated dispatch validates requests and rejects unknown, foreign, or +unimplemented methods. On failure it runs a non-null completion callback. After +successful dispatch, the application owns completion and must run it exactly +once. diff --git a/docs/en/getting_started.md b/docs/en/getting_started.md index cabd5ecc3c..2ac22c3a75 100644 --- a/docs/en/getting_started.md +++ b/docs/en/getting_started.md @@ -454,6 +454,10 @@ It is recommended to use the latest possible version of libunwind. no known issues. +## flatbuffers: 25.2.10 + +Optional and disabled by default. bRPC uses FlatBuffers builder internals, so only this version is accepted; see [FlatBuffers messages](flatbuffers.md) for using another version. + # Track instances We provide a program to help you to track and monitor all brpc instances. Just run [trackme_server](https://github.com/apache/brpc/tree/master/tools/trackme_server/) somewhere and launch need-to-be-tracked instances with -trackme_server=SERVER. The trackme_server will receive pings from instances periodically and print logs when it does. You can aggregate instance addresses from the log and call builtin services of the instances for further information. diff --git a/src/brpc/CMakeLists.txt b/src/brpc/CMakeLists.txt index 4efd059bfd..03e45da465 100644 --- a/src/brpc/CMakeLists.txt +++ b/src/brpc/CMakeLists.txt @@ -43,6 +43,14 @@ set(BRPC_DEFINITIONS BRPC_WITH_URMA=${WITH_URMA_VAL} BRPC_WITH_UBRING=${WITH_UBRING_VAL}) set(BRPC_INCLUDE_DIRS) +if(WITH_FLATBUFFERS) + find_path(FLATBUFFERS_INCLUDE_DIR NAMES flatbuffers/flatbuffers.h) + if(NOT FLATBUFFERS_INCLUDE_DIR) + message(FATAL_ERROR + "WITH_FLATBUFFERS requires FlatBuffers headers; set FLATBUFFERS_INCLUDE_DIR.") + endif() + list(APPEND BRPC_INCLUDE_DIRS ${FLATBUFFERS_INCLUDE_DIR}) +endif() if(WITH_THRIFT) list(APPEND BRPC_DEFINITIONS ENABLE_THRIFT_FRAMED_PROTOCOL) find_library(THRIFT_LIB NAMES thrift) @@ -241,6 +249,9 @@ endif() file(GLOB_RECURSE BRPC_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp") file(GLOB_RECURSE THRIFT_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/thrift*.cpp") +if(NOT WITH_FLATBUFFERS) + list(FILTER BRPC_SOURCES EXCLUDE REGEX "/brpc/flatbuffers/.*\\.cpp$") +endif() file(GLOB_RECURSE EXCLUDE_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/event_dispatcher_*.cpp") # When building with the real liburma, exclude the link-time mock so its urma_* diff --git a/src/brpc/flatbuffers/message.cpp b/src/brpc/flatbuffers/message.cpp new file mode 100644 index 0000000000..d3116d9a75 --- /dev/null +++ b/src/brpc/flatbuffers/message.cpp @@ -0,0 +1,310 @@ +// 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/flatbuffers/message.h" + +#if BRPC_WITH_FLATBUFFERS +#include +#include +#include "butil/logging.h" + +namespace brpc { +namespace flatbuffers { +namespace { + +uint8_t* AlignPayload(uint8_t* data) { + const uintptr_t address = reinterpret_cast(data); + const size_t padding = (kBufferAlignment - address % kBufferAlignment) % + kBufferAlignment; + return data + padding; +} + +} // namespace + +SlabAllocator& SlabAllocator::operator=(SlabAllocator&& other) noexcept { + if (this != &other) { + SlabAllocator tmp(std::move(other)); + swap(tmp); + } + return *this; +} + +void SlabAllocator::swap(SlabAllocator& other) noexcept { + _iobuf.swap(other._iobuf); + std::swap(_data, other._data); + std::swap(_capacity, other._capacity); +} + +uint8_t* SlabAllocator::allocate(size_t size) { + RELEASE_ASSERT_VERBOSE(size > 0 && size < FLATBUFFERS_MAX_BUFFER_SIZE, + "Invalid FlatBuffers allocation size"); + RELEASE_ASSERT_VERBOSE(_data == nullptr, + "Allocator already has a live buffer"); + butil::SingleIOBuf storage; + const uint32_t allocation_size = static_cast( + size + kDefaultMetaSize + kBufferAlignment - 1); + uint8_t* raw = static_cast(storage.allocate(allocation_size)); + RELEASE_ASSERT_VERBOSE(raw != nullptr, + "Fail to allocate FlatBuffers storage"); + _data = AlignPayload(raw + kDefaultMetaSize); + _capacity = size; + _iobuf.swap(storage); + return _data; +} + +void SlabAllocator::deallocate(uint8_t* p, size_t /*size*/) { + if (!p) { + return; + } + RELEASE_ASSERT_VERBOSE(p == _data, + "Invalid FlatBuffers deallocation pointer"); + _iobuf.reset(); + _data = nullptr; + _capacity = 0; +} + +uint8_t* SlabAllocator::reallocate_downward( + uint8_t* old_p, size_t old_size, size_t new_size, + size_t in_use_back, size_t in_use_front) { + RELEASE_ASSERT_VERBOSE(old_p != nullptr && old_p == _data && + old_size == _capacity, + "Invalid FlatBuffers reallocation buffer"); + RELEASE_ASSERT_VERBOSE(new_size > old_size, + "FlatBuffers reallocation must grow"); + RELEASE_ASSERT_VERBOSE(in_use_back <= old_size && + in_use_front <= old_size - in_use_back, + "Invalid FlatBuffers scratch or payload size"); + // Keep the old allocation alive until BOTH data and scratch are copied. + SlabAllocator replacement; + uint8_t* data = replacement.allocate(new_size); + if (in_use_back) { + memcpy(data + new_size - in_use_back, + old_p + old_size - in_use_back, in_use_back); + } + if (in_use_front) { + memcpy(data, old_p, in_use_front); + } + swap(replacement); + return _data; +} + +Message::Message(const butil::IOBuf::BlockRef& ref, uint32_t meta_size, + uint32_t msg_size) + : _iobuf(ref), _meta_size(meta_size), _msg_size(msg_size) {} + +Message& Message::operator=(Message&& other) noexcept { + if (this != &other) { + Clear(); + Swap(other); + } + return *this; +} + +void Message::Swap(Message& other) noexcept { + _iobuf.swap(other._iobuf); + std::swap(_meta_size, other._meta_size); + std::swap(_msg_size, other._msg_size); +} + +void Message::MergeFrom(const Message& other) { + if (&other == this) { + return; + } + _iobuf = other._iobuf; + _meta_size = other._meta_size; + _msg_size = other._msg_size; +} + +void Message::Clear() { + _iobuf.reset(); + _meta_size = 0; + _msg_size = 0; +} + +const uint8_t* Message::data() const { + const uint8_t* raw = static_cast(_iobuf.get_begin()); + return raw ? raw + _meta_size : nullptr; +} + +bool Message::parse_msg_from_iobuf(const butil::IOBuf& buf, size_t msg_size, + size_t meta_size) { + const size_t total = buf.size(); + // Subtraction avoids overflow, and the limits cover SingleIOBuf's uint32_t + // sizes, including allocation overhead, before any narrowing conversion. + if (msg_size == 0 || total >= FLATBUFFERS_MAX_BUFFER_SIZE || + meta_size > total || msg_size != total - meta_size) { + return false; + } + butil::SingleIOBuf storage; + const butil::StringPiece first = buf.backing_block(0); + if (first.size() == total && + reinterpret_cast(first.data() + meta_size) % + kBufferAlignment == 0) { + if (!storage.assign(buf, static_cast(total))) { + return false; + } + } else { + uint8_t* raw = static_cast(storage.allocate( + static_cast(total + kBufferAlignment - 1))); + if (!raw) { + return false; + } + uint8_t* begin = AlignPayload(raw + meta_size) - meta_size; + if (buf.copy_to(begin, total) != total) { + return false; + } + const butil::IOBuf::BlockRef& ref = storage.get_cur_ref(); + butil::IOBuf::BlockRef aligned_ref = { + ref.offset + static_cast(begin - raw), + static_cast(total), ref.block}; + butil::SingleIOBuf aligned(aligned_ref); + storage.swap(aligned); + } + _iobuf.swap(storage); + _meta_size = static_cast(meta_size); + _msg_size = static_cast(msg_size); + return true; +} + +bool Message::append_msg_to_iobuf(butil::IOBuf& buf) const { + if (!data() || !_msg_size) { + return false; + } + _iobuf.append_to(&buf); + return true; +} + +void* Message::reduce_meta_size_and_get_buf(uint32_t new_size) { + if (!data() || new_size > _meta_size) { + errno = EINVAL; + return nullptr; + } + if (new_size != _meta_size) { + const uint32_t offset = _meta_size - new_size; + const butil::IOBuf::BlockRef& ref = _iobuf.get_cur_ref(); + butil::IOBuf::BlockRef sub_ref = { + ref.offset + offset, ref.length - offset, ref.block}; + butil::SingleIOBuf slice(sub_ref); + _iobuf.swap(slice); + _meta_size = new_size; + } + return mutable_buf_begin(); +} + +MessageBuilder::MessageBuilder(size_t initial_size) + : ::flatbuffers::FlatBufferBuilder( + initial_size, &slab_allocator_, false, kBufferAlignment) { + RELEASE_ASSERT_VERBOSE(initial_size > 0 && + initial_size < FLATBUFFERS_MAX_BUFFER_SIZE, + "Invalid FlatBuffers initial size"); +} + +MessageBuilder::MessageBuilder(MessageBuilder&& other) : MessageBuilder() { + Swap(other); +} + +MessageBuilder& MessageBuilder::operator=(MessageBuilder&& other) { + if (this != &other) { + MessageBuilder tmp(std::move(other)); + Swap(tmp); + } + return *this; +} + +void MessageBuilder::ClearStringPool() { + // FlatBuffers' shared-string comparator stores a vector_downward pointer. + // Discard only this optional cache when moving between vector objects. + delete string_pool; + string_pool = nullptr; +} + +void MessageBuilder::Swap(MessageBuilder& other) { + if (this == &other) { + return; + } + ClearStringPool(); + other.ClearStringPool(); + slab_allocator_.swap(other.slab_allocator_); + ::flatbuffers::FlatBufferBuilder::Swap(other); + // Each builder must continue to refer to its OWN allocator member. + buf_.swap_allocator(other.buf_); +} + +MessageBuilder::MessageBuilder(::flatbuffers::FlatBufferBuilder&& src) + : ::flatbuffers::FlatBufferBuilder(std::move(src)) { + ClearStringPool(); + RELEASE_ASSERT_VERBOSE(minalign_ <= kBufferAlignment, + "Unsupported FlatBuffers alignment"); + decltype(buf_) replacement( + buf_.capacity() ? buf_.capacity() : 1024, + &slab_allocator_, false, kBufferAlignment); + if (buf_.capacity()) { + replacement.push(buf_.data(), buf_.size()); + for (::flatbuffers::uoffset_t i = 0; i < buf_.scratch_size(); ++i) { + replacement.scratch_push_small(buf_.scratch_data()[i]); + } + } + // The temporary returns the original buffer to its actual allocator. + buf_.swap(replacement); +} + +MessageBuilder& MessageBuilder::operator=( + ::flatbuffers::FlatBufferBuilder&& src) { + if (static_cast<::flatbuffers::FlatBufferBuilder*>(this) != &src) { + MessageBuilder tmp(std::move(src)); + Swap(tmp); + } + return *this; +} + +Message MessageBuilder::ReleaseMessage() { + RELEASE_ASSERT_VERBOSE(finished && buf_.size() > 0, + "Finish the FlatBuffer before releasing a message"); + const uint32_t msg_size = static_cast(buf_.size()); + const uint8_t* msg_data = buf_.data(); + const butil::SingleIOBuf& storage = slab_allocator_._iobuf; + const uint8_t* raw = static_cast(storage.get_begin()); + RELEASE_ASSERT_VERBOSE(raw != nullptr, "Missing FlatBuffers storage"); + RELEASE_ASSERT_VERBOSE(msg_data >= raw + kDefaultMetaSize, + "Missing FlatBuffers metadata prefix"); + const size_t begin = msg_data - raw - kDefaultMetaSize; + RELEASE_ASSERT_VERBOSE( + begin + kDefaultMetaSize + msg_size <= storage.get_length(), + "FlatBuffers message exceeds storage"); + const butil::IOBuf::BlockRef& ref = storage.get_cur_ref(); + const butil::IOBuf::BlockRef sub_ref = { + ref.offset + static_cast(begin), + kDefaultMetaSize + msg_size, ref.block}; + Message msg(sub_ref, kDefaultMetaSize, msg_size); + // Never expose stale scratch or allocator bytes in the reserved prefix. + memset(msg.mutable_buf_begin(), 0, kDefaultMetaSize); + Reset(); + return msg; +} + +bool ParseFbFromIOBUF(Message* msg, size_t msg_size, const butil::IOBuf& buf, + size_t meta_size) { + return msg && msg->parse_msg_from_iobuf(buf, msg_size, meta_size); +} + +bool SerializeFbToIOBUF(const Message* msg, butil::IOBuf& buf) { + return msg && msg->append_msg_to_iobuf(buf); +} + +} // namespace flatbuffers +} // namespace brpc +#endif // BRPC_WITH_FLATBUFFERS diff --git a/src/brpc/flatbuffers/message.h b/src/brpc/flatbuffers/message.h new file mode 100644 index 0000000000..7ecbb6c9fc --- /dev/null +++ b/src/brpc/flatbuffers/message.h @@ -0,0 +1,187 @@ +// 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_FLATBUFFERS_MESSAGE_H +#define BRPC_FLATBUFFERS_MESSAGE_H + +#include "butil/config.h" + +#if BRPC_WITH_FLATBUFFERS +#include +#include +#include +#include + +#if !defined(FLATBUFFERS_VERSION_MAJOR) || \ + !defined(FLATBUFFERS_VERSION_MINOR) || \ + !defined(FLATBUFFERS_VERSION_REVISION) || \ + FLATBUFFERS_VERSION_MAJOR != 25 || \ + FLATBUFFERS_VERSION_MINOR != 2 || \ + FLATBUFFERS_VERSION_REVISION != 10 +#error "bRPC FlatBuffers support requires FlatBuffers 25.2.10" +#endif + +#include "butil/macros.h" +#include "butil/single_iobuf.h" +#include "brpc/nonreflectable_message.h" + +namespace brpc { +namespace flatbuffers { + +// Space preceding a payload, available to a future RPC transport. +constexpr uint32_t kDefaultMetaSize = 64; +// IOBuf slices need not start at an aligned address. The allocator corrects it. +constexpr size_t kBufferAlignment = 64; + +class MessageBuilder; + +// FlatBufferBuilder cannot recover from a null allocation. Allocation failure +// and sizes outside FlatBuffers' offset range are fatal, even in release +// builds. +class SlabAllocator : public ::flatbuffers::Allocator { +public: + SlabAllocator() : _data(nullptr), _capacity(0) {} + SlabAllocator(SlabAllocator&& other) noexcept : SlabAllocator() { + swap(other); + } + SlabAllocator& operator=(SlabAllocator&& other) noexcept; + + uint8_t* allocate(size_t size) override; + void deallocate(uint8_t* p, size_t size) override; + uint8_t* reallocate_downward(uint8_t* old_p, size_t old_size, + size_t new_size, size_t in_use_back, + size_t in_use_front) override; + void swap(SlabAllocator& other) noexcept; + +private: + DISALLOW_COPY_AND_ASSIGN(SlabAllocator); +friend class MessageBuilder; + butil::SingleIOBuf _iobuf; + uint8_t* _data; + size_t _capacity; +}; + +// Construct the allocator before the FlatBufferBuilder base uses it, and +// destroy it after that base has returned its buffer. +struct SlabAllocatorMember { + SlabAllocator slab_allocator_; +}; + +// A move-only message for application code. Protobuf CopyFrom/MergeFrom retain +// the same ref-counted IOBuf instead of deep-copying it. Mutable access through +// any alias changes every message sharing that storage and must not race +// readers. +// Parsing checks framing, not the schema: Verify() must succeed before +// reading data received from an untrusted peer. +class Message : public NonreflectableMessage { +public: + Message() : _meta_size(0), _msg_size(0) {} + Message(Message&& other) noexcept : Message() { Swap(other); } + Message& operator=(Message&& other) noexcept; + + void MergeFrom(const Message&) override; + void Clear() override; + void Swap(Message& other) noexcept; + + const uint8_t* data() const; + // These mutable pointers may refer to storage shared by CopyFrom/MergeFrom. + void* mutable_data() { return const_cast(data()); } + void* mutable_buf_begin() { + return const_cast(_iobuf.get_begin()); + } + void* reduce_meta_size_and_get_buf(uint32_t new_size); + uint32_t get_meta_size() const { return _meta_size; } + size_t size() const { return _msg_size; } + + template + bool Verify() const { + if (!data() || size() < sizeof(::flatbuffers::uoffset_t) || + size() >= FLATBUFFERS_MAX_BUFFER_SIZE) { + return false; + } + ::flatbuffers::Verifier verifier(data(), size()); + return verifier.VerifyBuffer(nullptr); + } + + // These accessors require a schema-verified buffer. GetMutableRoot may + // mutate all Message aliases that retain the same IOBuf block. + template const T* GetRoot() const { + return data() ? ::flatbuffers::GetRoot(data()) : nullptr; + } + template T* GetMutableRoot() { + return data() + ? ::flatbuffers::GetMutableRoot(mutable_data()) : nullptr; + } + + // Failure leaves the old message unchanged. A fragmented or unaligned + // payload is copied into aligned storage; aligned contiguous input is + // shared. + // When msg_size comes from an untrusted peer, callers must bound it with + // their own max-message-size policy before parsing. This method does not + // verify the schema; call Verify() before reading the root. + bool parse_msg_from_iobuf(const butil::IOBuf& buf, size_t msg_size, + size_t meta_size); + bool append_msg_to_iobuf(butil::IOBuf& buf) const; + +private: + DISALLOW_COPY_AND_ASSIGN(Message); +friend class MessageBuilder; + Message(const butil::IOBuf::BlockRef& ref, uint32_t meta_size, + uint32_t msg_size); + butil::SingleIOBuf _iobuf; + uint32_t _meta_size; + uint32_t _msg_size; +}; + +class MessageBuilder : private SlabAllocatorMember, + public ::flatbuffers::FlatBufferBuilder { +public: + explicit MessageBuilder(size_t initial_size = 1024); + MessageBuilder(MessageBuilder&& other); + MessageBuilder& operator=(MessageBuilder&& other); + + // Importing a foreign builder copies its payload and scratch data, + // retaining build state. Its original allocator frees the source + // allocation correctly. No free()/delete[] guess or spare-prefix + // assumption is made. + explicit MessageBuilder(::flatbuffers::FlatBufferBuilder&& src); + MessageBuilder& operator=(::flatbuffers::FlatBufferBuilder&& src); + + void Swap(MessageBuilder& other); + // Requires Finish(). The returned message shares storage without copying; + // both this builder and a moved-from builder may immediately be reused. + Message ReleaseMessage(); + +private: + DISALLOW_COPY_AND_ASSIGN(MessageBuilder); + void ClearStringPool(); + // The inherited release methods would retain a pointer to our member + // allocator after this builder dies. Only ReleaseMessage is supported. + using ::flatbuffers::FlatBufferBuilder::Release; + using ::flatbuffers::FlatBufferBuilder::ReleaseRaw; + void ReleaseBufferPointer() = delete; + using ::flatbuffers::FlatBufferBuilder::SwapBufAllocator; +}; + +bool ParseFbFromIOBUF(Message* msg, size_t msg_size, const butil::IOBuf& buf, + size_t meta_size = 0); +bool SerializeFbToIOBUF(const Message* msg, butil::IOBuf& buf); + +} // namespace flatbuffers +} // namespace brpc +#endif // BRPC_WITH_FLATBUFFERS +#endif // BRPC_FLATBUFFERS_MESSAGE_H diff --git a/src/brpc/flatbuffers/service.cpp b/src/brpc/flatbuffers/service.cpp new file mode 100644 index 0000000000..eb4f44f4e9 --- /dev/null +++ b/src/brpc/flatbuffers/service.cpp @@ -0,0 +1,158 @@ +// 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/flatbuffers/service.h" + +#if BRPC_WITH_FLATBUFFERS +#include +#include +#include +#include +#include "butil/third_party/murmurhash3/murmurhash3.h" + +namespace brpc { +namespace flatbuffers { +namespace { + +bool IsIdentifier(const std::string& name) { + if (name.empty()) { + return false; + } + for (size_t i = 0; i < name.size(); ++i) { + const char c = name[i]; + if (c != '_' && !(c >= 'a' && c <= 'z') && + !(c >= 'A' && c <= 'Z') && + !(i != 0 && c >= '0' && c <= '9')) { + return false; + } + } + return true; +} + +bool IsNamespace(const std::string& name) { + if (name.empty()) { + return true; + } + size_t begin = 0; + do { + const size_t end = name.find('.', begin); + if (!IsIdentifier(name.substr(begin, end - begin))) { + return false; + } + if (end == std::string::npos) { + return true; + } + begin = end + 1; + } while (begin < name.size()); + return false; +} + +} // namespace + +int ServiceDescriptor::init(const BrpcDescriptorTable& table) { + if (!_methods.empty()) { + errno = EALREADY; + return -1; + } + std::string prefix = table.prefix; + if (!prefix.empty() && prefix.back() == '.') { + prefix.pop_back(); + if (prefix.empty()) { + errno = EINVAL; + return -1; + } + } + if (!IsIdentifier(table.service_name) || !IsNamespace(prefix)) { + errno = EINVAL; + return -1; + } + std::istringstream input(table.method_name_list); + std::vector names; + std::set unique_names; + std::string name; + while (input >> name) { + if (!IsIdentifier(name) || !unique_names.insert(name).second) { + errno = EINVAL; + return -1; + } + names.push_back(name); + } + if (names.empty() || names.size() > static_cast( + std::numeric_limits::max()) || + (!table.method_ids.empty() && table.method_ids.size() != names.size())) { + errno = EINVAL; + return -1; + } + std::set unique_ids; + for (int id : table.method_ids) { + if (id < 0 || !unique_ids.insert(id).second) { + errno = EINVAL; + return -1; + } + } + _name = table.service_name; + _full_name = prefix.empty() ? _name : prefix + "." + _name; + butil::MurmurHash3_x86_32(_full_name.data(), _full_name.size(), 1, &_index); + _methods.reserve(names.size()); + for (size_t i = 0; i < names.size(); ++i) { + const int id = table.method_ids.empty() ? static_cast(i) + : table.method_ids[i]; + _methods.emplace_back(new MethodDescriptor(names[i], this, id)); + } + return 0; +} + +const MethodDescriptor* ServiceDescriptor::method(int position) const { + if (position < 0 || static_cast(position) >= _methods.size()) { + errno = EINVAL; + return nullptr; + } + return _methods[position].get(); +} + +const MethodDescriptor* ServiceDescriptor::FindMethodByIndex(int id) const { + for (const auto& method : _methods) { + if (method->index() == id) { + return method.get(); + } + } + errno = EINVAL; + return nullptr; +} + +MethodDescriptor::MethodDescriptor(const std::string& name, + const ServiceDescriptor* service, int index) + : _name(name), _full_name(service->full_name() + "." + name), + _service(service), _index(index) {} + +int parse_service_descriptors(const BrpcDescriptorTable& table, + ServiceDescriptor** out) { + if (!out) { + errno = EINVAL; + return -1; + } + auto descriptor = std::make_unique(); + if (descriptor->init(table) != 0) { + return -1; + } + *out = descriptor.release(); + return 0; +} + +} // namespace flatbuffers +} // namespace brpc +#endif // BRPC_WITH_FLATBUFFERS diff --git a/src/brpc/flatbuffers/service.h b/src/brpc/flatbuffers/service.h new file mode 100644 index 0000000000..56345f63dd --- /dev/null +++ b/src/brpc/flatbuffers/service.h @@ -0,0 +1,133 @@ +// 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_FLATBUFFERS_SERVICE_H +#define BRPC_FLATBUFFERS_SERVICE_H + +#include "butil/config.h" +#include "butil/macros.h" + +#if BRPC_WITH_FLATBUFFERS +#include +#include +#include +#include + +namespace google { +namespace protobuf { +class Closure; +class RpcController; +} // namespace protobuf +} // namespace google + +namespace brpc { +namespace flatbuffers { +class Message; +class ServiceDescriptor; + +struct BrpcDescriptorTable { + // A namespace, optionally followed by one period. Empty means global scope. + std::string prefix; + std::string service_name; + // Whitespace-separated method names in declaration order. + std::string method_name_list; + // Stable wire IDs, NOT array positions. Empty retains legacy ordinal IDs. + // New schemas should specify every ID; IDs must be nonnegative and unique. + std::vector method_ids; +}; + +class MethodDescriptor { +public: + MethodDescriptor(const std::string& name, const ServiceDescriptor* service, + int index); + const std::string& name() const { return _name; } + const std::string& full_name() const { return _full_name; } + int index() const { return _index; } + const ServiceDescriptor* service() const { return _service; } + +private: + std::string _name; + std::string _full_name; + const ServiceDescriptor* _service; + int _index; +}; + +// Immutable after successful initialization. Owns all method descriptors. +class ServiceDescriptor { +public: + ServiceDescriptor() : _index(0) {} + int init(const BrpcDescriptorTable& table); + const std::string& name() const { return _name; } + const std::string& full_name() const { return _full_name; } + // MurmurHash3 of the canonical fully-qualified service name, seed 1. + uint32_t index() const { return _index; } + int method_count() const { return static_cast(_methods.size()); } + // Dense declaration-order lookup, for enumeration and generated stubs. + const MethodDescriptor* method(int position) const; + // Sparse stable-ID lookup, for wire dispatch. The fb_rpc transport uses + // this method to resolve stable wire IDs; missing IDs return nullptr. + const MethodDescriptor* FindMethodByIndex(int id) const; + +private: + DISALLOW_COPY_AND_ASSIGN(ServiceDescriptor); + std::string _name; + std::string _full_name; + uint32_t _index; + std::vector > _methods; +}; + +// On success the caller owns *out; on failure *out is unchanged. +int parse_service_descriptors(const BrpcDescriptorTable& table, + ServiceDescriptor** out); + +class RpcChannel { +public: + RpcChannel() = default; + virtual ~RpcChannel() = default; + virtual void FBCallMethod(const MethodDescriptor* method, + google::protobuf::RpcController* controller, + const Message* request, Message* response, + google::protobuf::Closure* done) = 0; + +private: + DISALLOW_COPY_AND_ASSIGN(RpcChannel); +}; + +class Service { +public: + Service() = default; + virtual ~Service() = default; + enum ChannelOwnership { STUB_OWNS_CHANNEL, STUB_DOESNT_OWN_CHANNEL }; + virtual const ServiceDescriptor* GetDescriptor() = 0; + // Public dispatch contract consumed by generated services and fb_rpc: + // generated dispatch validates the method and request schema, reports + // failures through controller, and runs a non-null done exactly once. + // Successful implementations may complete asynchronously; request, + // response and controller must remain alive until done is run. + virtual void FBCallMethod(const MethodDescriptor* method, + google::protobuf::RpcController* controller, + const Message* request, Message* response, + google::protobuf::Closure* done) = 0; + +private: + DISALLOW_COPY_AND_ASSIGN(Service); +}; + +} // namespace flatbuffers +} // namespace brpc +#endif // BRPC_WITH_FLATBUFFERS +#endif // BRPC_FLATBUFFERS_SERVICE_H diff --git a/src/butil/CMakeLists.txt b/src/butil/CMakeLists.txt index 1e57710a19..901bc1954a 100644 --- a/src/butil/CMakeLists.txt +++ b/src/butil/CMakeLists.txt @@ -21,6 +21,12 @@ if(WITH_GLOG) set(BRPC_WITH_GLOG 1) endif() +# Recorded in config.h so installed brpc/flatbuffers headers see the choice. +set(WITH_FLATBUFFERS_VAL "0") +if(WITH_FLATBUFFERS) + set(WITH_FLATBUFFERS_VAL "1") +endif() + set(WITH_CPU_FREQUENCY_VAL "0") if(WITH_CPU_FREQUENCY) set(WITH_CPU_FREQUENCY_VAL "1") diff --git a/src/butil/config.h.in b/src/butil/config.h.in index d8de111be9..ece54504ce 100644 --- a/src/butil/config.h.in +++ b/src/butil/config.h.in @@ -21,6 +21,11 @@ #endif #cmakedefine BRPC_WITH_GLOG @WITH_GLOG_VAL@ +#ifdef BRPC_WITH_FLATBUFFERS +#undef BRPC_WITH_FLATBUFFERS +#endif +#define BRPC_WITH_FLATBUFFERS @WITH_FLATBUFFERS_VAL@ + #ifdef BUTIL_USE_CPU_FREQUENCY #undef BUTIL_USE_CPU_FREQUENCY #endif diff --git a/test/BUILD.bazel b/test/BUILD.bazel index 1aecc39987..a2c291eb7e 100644 --- a/test/BUILD.bazel +++ b/test/BUILD.bazel @@ -229,11 +229,47 @@ root_runfiles( ], ) +FLATBUFFERS_COMPATIBILITY = select({ + "//:brpc_with_flatbuffers": [], + "//conditions:default": ["@platforms//:incompatible"], +}) + +genrule( + name = "flatbuffers_test_messages", + srcs = ["flatbuffers_message.fbs"], + outs = ["flatbuffers_message_generated.h"], + cmd = select({ + "//:brpc_with_flatbuffers": "$(location @com_github_google_flatbuffers//:flatc) --cpp -o $(@D) $(location flatbuffers_message.fbs)", + "//conditions:default": "exit 1", + }), + target_compatible_with = FLATBUFFERS_COMPATIBILITY, + tools = select({ + "//:brpc_with_flatbuffers": ["@com_github_google_flatbuffers//:flatc"], + "//conditions:default": [], + }), +) + +cc_test( + name = "brpc_flatbuffers_unittest", + srcs = [ + "brpc_flatbuffers_unittest.cpp", + ":flatbuffers_test_messages", + ], + copts = COPTS, + includes = ["."], + target_compatible_with = FLATBUFFERS_COMPATIBILITY, + deps = [ + "//:brpc", + "@com_google_googletest//:gtest_main", + ], +) + generate_unittests( name = "brpc_unittests", - srcs = glob([ - "brpc_*_unittest.cpp", - ]), + srcs = glob( + ["brpc_*_unittest.cpp"], + exclude = ["brpc_flatbuffers_unittest.cpp"], + ), deps = [ ":gperftools_helper", "//:brpc", diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index a0fc2d0536..e2e9a84757 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -59,21 +59,25 @@ set(BRPC_SYSTEM_GTEST_SOURCE_DIR "" CACHE PATH "System googletest source directo if(DOWNLOAD_GTEST) include(SetupGtest) - set(BRPC_DOWNLOADED_GTEST_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/downloaded-googletest-include) - file(REMOVE_RECURSE ${BRPC_DOWNLOADED_GTEST_INCLUDE_DIR}) - file(MAKE_DIRECTORY ${BRPC_DOWNLOADED_GTEST_INCLUDE_DIR}) - file(COPY - ${PROJECT_BINARY_DIR}/googletest-src/googletest/include/ - ${PROJECT_BINARY_DIR}/googletest-src/googlemock/include/ - DESTINATION ${BRPC_DOWNLOADED_GTEST_INCLUDE_DIR}) - target_include_directories(brpc_test_config BEFORE INTERFACE - ${BRPC_DOWNLOADED_GTEST_INCLUDE_DIR}) + set(BRPC_GTEST_SOURCE_DIR "${PROJECT_BINARY_DIR}/googletest-src") elseif(BRPC_SYSTEM_GTEST_SOURCE_DIR) add_subdirectory("${BRPC_SYSTEM_GTEST_SOURCE_DIR}" "${PROJECT_BINARY_DIR}/system-googletest-build") + set(BRPC_GTEST_SOURCE_DIR "${BRPC_SYSTEM_GTEST_SOURCE_DIR}") else() message(FATAL_ERROR "Googletest is not available") endif() +# Keep the selected headers ahead of unrelated -I paths. The original GTest +# include directories can be marked SYSTEM by its own targets. +set(BRPC_GTEST_INCLUDE_DIR "${CMAKE_CURRENT_BINARY_DIR}/selected-googletest-include") +file(REMOVE_RECURSE "${BRPC_GTEST_INCLUDE_DIR}") +file(MAKE_DIRECTORY "${BRPC_GTEST_INCLUDE_DIR}") +file(COPY + "${BRPC_GTEST_SOURCE_DIR}/googletest/include/" + "${BRPC_GTEST_SOURCE_DIR}/googlemock/include/" + DESTINATION "${BRPC_GTEST_INCLUDE_DIR}") +target_include_directories(brpc_test_config BEFORE INTERFACE "${BRPC_GTEST_INCLUDE_DIR}") + file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/cert1.key DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/cert2.key DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/cert1.crt DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) @@ -284,8 +288,37 @@ foreach(BTHREAD_UT ${BTHREAD_UNITTESTS}) add_test(NAME ${BTHREAD_UT_WE} COMMAND ${BTHREAD_UT_WE}) endforeach() +# FlatBuffers tests use the release library, independently of the debug suite. +if(WITH_FLATBUFFERS) + find_program(FLATBUFFERS_FLATC_EXECUTABLE NAMES flatc) + if(NOT FLATBUFFERS_FLATC_EXECUTABLE) + message(FATAL_ERROR + "FlatBuffers tests require an installed flatc; set FLATBUFFERS_FLATC_EXECUTABLE.") + endif() + set(FLATBUFFERS_TEST_HEADER + "${CMAKE_CURRENT_BINARY_DIR}/flatbuffers_message_generated.h") + add_custom_command( + OUTPUT "${FLATBUFFERS_TEST_HEADER}" + COMMAND "${FLATBUFFERS_FLATC_EXECUTABLE}" --cpp + -o "${CMAKE_CURRENT_BINARY_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/flatbuffers_message.fbs" + DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/flatbuffers_message.fbs" + "${FLATBUFFERS_FLATC_EXECUTABLE}" + COMMENT "Generating FlatBuffers test messages" + VERBATIM) + add_executable(brpc_flatbuffers_unittest + brpc_flatbuffers_unittest.cpp "${FLATBUFFERS_TEST_HEADER}") + target_include_directories(brpc_flatbuffers_unittest BEFORE PRIVATE + "${BRPC_GTEST_INCLUDE_DIR}") + target_link_libraries(brpc_flatbuffers_unittest PRIVATE + brpc-static brpc_test_config gtest_main ${DYNAMIC_LIB}) + add_test(NAME brpc_flatbuffers_unittest COMMAND brpc_flatbuffers_unittest) +endif() + # brpc tests file(GLOB BRPC_UNITTESTS "brpc_*_unittest.cpp") +list(REMOVE_ITEM BRPC_UNITTESTS + "${CMAKE_CURRENT_SOURCE_DIR}/brpc_flatbuffers_unittest.cpp") foreach(BRPC_UT ${BRPC_UNITTESTS}) get_filename_component(BRPC_UT_WE ${BRPC_UT} NAME_WE) add_executable(${BRPC_UT_WE} ${BRPC_UT} $) diff --git a/test/Makefile b/test/Makefile index de4d566bcd..f4d5462e63 100644 --- a/test/Makefile +++ b/test/Makefile @@ -181,12 +181,16 @@ TEST_BTHREAD_SOURCES = $(wildcard bthread_*unittest.cpp) TEST_BTHREAD_OBJS = $(addsuffix .o, $(basename $(TEST_BTHREAD_SOURCES))) TEST_BRPC_SOURCES = $(wildcard brpc_*unittest.cpp) +ifneq ($(WITH_FLATBUFFERS),1) +TEST_BRPC_SOURCES := $(filter-out brpc_flatbuffers_unittest.cpp,$(TEST_BRPC_SOURCES)) +endif TEST_BRPC_OBJS = $(addsuffix .o, $(basename $(TEST_BRPC_SOURCES))) TEST_PROTO_SOURCES = $(wildcard *.proto) TEST_PROTO_OBJS = $(TEST_PROTO_SOURCES:.proto=.pb.o) TEST_BINS = test_butil test_bvar $(TEST_BTHREAD_SOURCES:.cpp=) $(TEST_BRPC_SOURCES:.cpp=) +FLATBUFFERS_TEST_GEN_DIR = ../output/test/flatbuffers-generated .PHONY:all all: $(TEST_BINS) @@ -194,7 +198,7 @@ all: $(TEST_BINS) .PHONY:clean clean:clean_bins @echo "> Cleaning" - rm -rf $(TEST_BUTIL_OBJS) $(TEST_BVAR_OBJS) $(TEST_BTHREAD_OBJS) $(TEST_BRPC_OBJS) $(TEST_PROTO_OBJS) $(TEST_PROTO_SOURCES:.proto=.pb.h) $(TEST_PROTO_SOURCES:.proto=.pb.cc) + rm -rf $(TEST_BUTIL_OBJS) $(TEST_BVAR_OBJS) $(TEST_BTHREAD_OBJS) $(TEST_BRPC_OBJS) $(TEST_PROTO_OBJS) $(TEST_PROTO_SOURCES:.proto=.pb.h) $(TEST_PROTO_SOURCES:.proto=.pb.cc) flatbuffers_message_generated.h $(FLATBUFFERS_TEST_GEN_DIR) $(MAKE) -C.. clean_debug .PHONY:clean_bins @@ -235,6 +239,28 @@ else ifeq ($(SYSTEM),Darwin) $(CXX) -o $@ $(LIBPATHS) $(SOPATHS) $^ $(GTEST_STATIC_LINKINGS) $(UT_DYNAMIC_LINKINGS) endif +ifeq ($(WITH_FLATBUFFERS),1) +FLATC ?= flatc +FLATBUFFERS_TEST_HEADER = $(FLATBUFFERS_TEST_GEN_DIR)/flatbuffers_message_generated.h +FLATBUFFERS_TEST_HDRPATHS = -I$(FLATBUFFERS_TEST_GEN_DIR) $(HDRPATHS) + +$(FLATBUFFERS_TEST_HEADER):flatbuffers_message.fbs + @mkdir -p $(FLATBUFFERS_TEST_GEN_DIR) + $(FLATC) --cpp -o $(FLATBUFFERS_TEST_GEN_DIR) $< + +brpc_flatbuffers_unittest.o:brpc_flatbuffers_unittest.cpp $(FLATBUFFERS_TEST_HEADER) | libbrpc.dbg.$(SOEXT) + @echo "> Compiling $@" + $(CXX) -c $(FLATBUFFERS_TEST_HDRPATHS) $(CXXFLAGS) $< -o $@ + +brpc_flatbuffers_unittest:brpc_flatbuffers_unittest.o | libbrpc.dbg.$(SOEXT) + @echo "> Linking $@" +ifeq ($(SYSTEM),Linux) + $(CXX) -o $@ $(LIBPATHS) $(SOPATHS) -Xlinker "-(" $^ -Xlinker "-)" $(STATIC_LINKINGS) $(UT_DYNAMIC_LINKINGS) +else ifeq ($(SYSTEM),Darwin) + $(CXX) -o $@ $(LIBPATHS) $(SOPATHS) $^ $(GTEST_STATIC_LINKINGS) $(UT_DYNAMIC_LINKINGS) +endif +endif + brpc_%_unittest:$(TEST_PROTO_OBJS) brpc_%_unittest.o | libbrpc.dbg.$(SOEXT) @echo "> Linking $@" ifeq ($(SYSTEM),Linux) diff --git a/test/brpc_flatbuffers_unittest.cpp b/test/brpc_flatbuffers_unittest.cpp new file mode 100644 index 0000000000..b2b2dd74e0 --- /dev/null +++ b/test/brpc_flatbuffers_unittest.cpp @@ -0,0 +1,565 @@ +// 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/config.h" +#if BRPC_WITH_FLATBUFFERS +#include +#include +#include +#include +#include +#include +#include +#include +#include "brpc/flatbuffers/message.h" +#include "brpc/flatbuffers/service.h" +#include + +#if !BRPC_WITH_GLOG +namespace logging { +DECLARE_bool(crash_on_fatal_log); +} +#endif + +namespace { +using brpc::flatbuffers::Message; +using brpc::flatbuffers::MessageBuilder; +using brpc::flatbuffers::SlabAllocator; +using brpc::flatbuffers::ServiceDescriptor; +using brpc::flatbuffers::BrpcDescriptorTable; +using brpc_fbtest::Payload; + +void FinishPayload(::flatbuffers::FlatBufferBuilder& builder, + const std::string& text, int64_t value = 42) { + auto str = builder.CreateString(text); + std::vector numbers = {1, -2, 3, 4000}; + auto values = builder.CreateVector(numbers); + builder.Finish(brpc_fbtest::CreatePayload(builder, value, str, values)); +} + +Message MakeMessage(size_t length = 16, int64_t value = 42) { + MessageBuilder builder(8); + FinishPayload(builder, std::string(length, 'x'), value); + return builder.ReleaseMessage(); +} + +void ExpectPayload(const Message& msg, size_t length, int64_t value = 42) { + ASSERT_TRUE(msg.Verify()); + const Payload* root = msg.GetRoot(); + ASSERT_NE(nullptr, root); + EXPECT_EQ(value, root->value()); + ASSERT_NE(nullptr, root->message()); + EXPECT_EQ(std::string(length, 'x'), root->message()->str()); + ASSERT_NE(nullptr, root->values()); + ASSERT_EQ(4u, root->values()->size()); + EXPECT_EQ(-2, root->values()->Get(1)); +} + +TEST(FlatbuffersTest, EmptyAndClear) { + Message msg; + EXPECT_EQ(nullptr, msg.data()); + EXPECT_EQ(nullptr, msg.GetRoot()); + EXPECT_EQ(nullptr, msg.GetMutableRoot()); + EXPECT_EQ(nullptr, msg.reduce_meta_size_and_get_buf(0)); + EXPECT_FALSE(msg.Verify()); + EXPECT_EQ(0u, msg.size()); + butil::IOBuf wire; + EXPECT_FALSE(brpc::flatbuffers::SerializeFbToIOBUF(&msg, wire)); + EXPECT_FALSE(brpc::flatbuffers::SerializeFbToIOBUF(nullptr, wire)); + EXPECT_FALSE(brpc::flatbuffers::ParseFbFromIOBUF(nullptr, 0, wire)); + msg = MakeMessage(); + msg.Clear(); + msg.Clear(); + EXPECT_EQ(nullptr, msg.data()); + EXPECT_EQ(0u, msg.get_meta_size()); + EXPECT_EQ(0u, msg.size()); +} + +TEST(FlatbuffersTest, ReleaseLifetimeAndGrowth) { + for (size_t length : {0u, 16u, 63u, 1024u, 8192u, 32768u}) { + SCOPED_TRACE(length); + Message msg; + const uint8_t* payload = nullptr; + { + MessageBuilder builder(8); + FinishPayload(builder, std::string(length, 'x')); + payload = builder.GetBufferPointer(); + msg = builder.ReleaseMessage(); + EXPECT_EQ(payload, msg.data()); + EXPECT_EQ(0u, builder.GetSize()); + FinishPayload(builder, "replacement"); + } + ExpectPayload(msg, length); + EXPECT_EQ(std::string(brpc::flatbuffers::kDefaultMetaSize, '\0'), + std::string(static_cast(msg.mutable_buf_begin()), + msg.get_meta_size())); + } +} + +TEST(FlatbuffersTest, MessageMoveAndSwap) { + Message first = MakeMessage(16, 1); + const uint8_t* ptr = first.data(); + Message second(std::move(first)); + EXPECT_EQ(ptr, second.data()); + EXPECT_EQ(nullptr, first.data()); + EXPECT_EQ(0u, first.size()); + EXPECT_EQ(0u, first.get_meta_size()); + EXPECT_FALSE(first.Verify()); + first = MakeMessage(8192, 2); + second = std::move(first); + EXPECT_EQ(nullptr, first.data()); + ExpectPayload(second, 8192, 2); + Message* alias = &second; + second = std::move(*alias); + ExpectPayload(second, 8192, 2); + first = MakeMessage(16, 3); + first.Swap(second); + ExpectPayload(first, 8192, 2); + ExpectPayload(second, 16, 3); + Message empty; + second = std::move(empty); + EXPECT_EQ(nullptr, second.data()); + EXPECT_EQ(nullptr, empty.data()); +} + +TEST(FlatbuffersTest, MergeFromAndCopyFromShareStorage) { + Message source = MakeMessage(128, 7); + Message merged = MakeMessage(16, 1); + merged.MergeFrom(source); + ExpectPayload(merged, 128, 7); + EXPECT_EQ(source.data(), merged.data()); + merged.MergeFrom(merged); + ExpectPayload(merged, 128, 7); + + Message copied = MakeMessage(32, 2); + copied.CopyFrom(source); + ExpectPayload(copied, 128, 7); + EXPECT_EQ(source.data(), copied.data()); + + uint8_t* buffer = static_cast(copied.mutable_data()); + uint8_t* table = buffer + ::flatbuffers::ReadScalar<::flatbuffers::uoffset_t>(buffer); + const ::flatbuffers::soffset_t vtable_offset = + ::flatbuffers::ReadScalar<::flatbuffers::soffset_t>(table); + uint8_t* vtable = table - vtable_offset; + const ::flatbuffers::voffset_t value_offset = + ::flatbuffers::ReadScalar<::flatbuffers::voffset_t>( + vtable + Payload::VT_VALUE); + ASSERT_NE(0, value_offset); + ::flatbuffers::WriteScalar(table + value_offset, 9); + ExpectPayload(source, 128, 9); + ExpectPayload(merged, 128, 9); + ExpectPayload(copied, 128, 9); + + source.Clear(); + ExpectPayload(merged, 128, 9); + ExpectPayload(copied, 128, 9); + + Message empty; + merged.MergeFrom(empty); + EXPECT_EQ(nullptr, merged.data()); + EXPECT_EQ(0u, merged.size()); +} + +TEST(FlatbuffersTest, BuilderMoveAndReuse) { + MessageBuilder first(8); + auto str = first.CreateString(std::string(8192, 'x')); + MessageBuilder second(std::move(first)); + second.Finish(brpc_fbtest::CreatePayload(second, 8, str)); + Message msg = second.ReleaseMessage(); + ASSERT_TRUE(msg.Verify()); + EXPECT_EQ(8192u, msg.GetRoot()->message()->size()); + FinishPayload(first, std::string(16, 'x')); + ExpectPayload(first.ReleaseMessage(), 16); + FinishPayload(second, std::string(63, 'x')); + first = std::move(second); + MessageBuilder* alias = &first; + first = std::move(*alias); + ExpectPayload(first.ReleaseMessage(), 63); + FinishPayload(second, std::string(1024, 'x')); + ExpectPayload(second.ReleaseMessage(), 1024); +} + +TEST(FlatbuffersTest, BuilderSwapFinishedAndUnfinished) { + MessageBuilder first; + FinishPayload(first, std::string(16, 'x')); + MessageBuilder second; + auto str = second.CreateString(std::string(63, 'x')); + first.Swap(second); + ExpectPayload(second.ReleaseMessage(), 16); + first.Finish(brpc_fbtest::CreatePayload(first, 9, str)); + Message msg = first.ReleaseMessage(); + ASSERT_TRUE(msg.Verify()); + EXPECT_EQ(9, msg.GetRoot()->value()); + EXPECT_EQ(63u, msg.GetRoot()->message()->size()); +} + +TEST(FlatbuffersTest, SharedStringsAfterMoveAndImport) { + MessageBuilder first(8); + auto old = first.CreateSharedString("shared"); + MessageBuilder second(std::move(first)); + auto same = second.CreateSharedString("shared"); + auto other = second.CreateSharedString("other"); + std::vector<::flatbuffers::Offset<::flatbuffers::String> > strings = { + old, same, other}; + second.Finish(second.CreateVector(strings)); + Message msg = second.ReleaseMessage(); + auto root = ::flatbuffers::GetRoot< + ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String> > >( + msg.data()); + EXPECT_EQ("shared", root->Get(0)->str()); + EXPECT_EQ("shared", root->Get(1)->str()); + EXPECT_EQ("other", root->Get(2)->str()); + ::flatbuffers::FlatBufferBuilder foreign; + auto text = foreign.CreateSharedString("shared"); + MessageBuilder imported(std::move(foreign)); + imported.CreateSharedString("after import"); + imported.Finish(brpc_fbtest::CreatePayload(imported, 1, text)); + ASSERT_TRUE(imported.ReleaseMessage().Verify()); +} + +class CountingAllocator : public ::flatbuffers::Allocator { +public: + CountingAllocator(int* allocations, int* frees, int* destructors) + : _allocations(allocations), _frees(frees), _destructors(destructors) {} + ~CountingAllocator() override { ++*_destructors; } + uint8_t* allocate(size_t n) override { + ++*_allocations; + return new uint8_t[n]; + } + void deallocate(uint8_t* p, size_t) override { + ++*_frees; + delete[] p; + } +private: + int* _allocations; + int* _frees; + int* _destructors; +}; + +TEST(FlatbuffersTest, ImportForeignAllocatorAndScratch) { + int allocations = 0; + int frees = 0; + int destructors = 0; + Message msg; + { + ::flatbuffers::FlatBufferBuilder foreign(8, + new CountingAllocator(&allocations, &frees, &destructors), true); + auto text = foreign.CreateString(std::string(8192, 'x')); + const auto table = foreign.StartTable(); + foreign.AddOffset(Payload::VT_MESSAGE, text); + foreign.AddElement(Payload::VT_VALUE, 123, 0); + MessageBuilder imported(std::move(foreign)); + EXPECT_EQ(allocations, frees); + EXPECT_EQ(1, destructors); + imported.Finish(::flatbuffers::Offset(imported.EndTable(table))); + msg = imported.ReleaseMessage(); + FinishPayload(foreign, "reuse source"); + } + ASSERT_TRUE(msg.Verify()); + EXPECT_EQ(123, msg.GetRoot()->value()); + EXPECT_EQ(8192u, msg.GetRoot()->message()->size()); + EXPECT_EQ(allocations, frees); + EXPECT_EQ(1, destructors); +} + +TEST(FlatbuffersTest, ImportBorrowedAllocatorDoesNotOwnAllocator) { + int allocations = 0; + int frees = 0; + int destructors = 0; + Message msg; + { + CountingAllocator allocator(&allocations, &frees, &destructors); + ::flatbuffers::FlatBufferBuilder foreign(8, &allocator, false); + FinishPayload(foreign, std::string(1024, 'x')); + MessageBuilder imported(std::move(foreign)); + EXPECT_EQ(allocations, frees); + EXPECT_EQ(0, destructors); + msg = imported.ReleaseMessage(); + } + EXPECT_EQ(1, destructors); + ExpectPayload(msg, 1024); +} + +TEST(FlatbuffersTest, ImportFinishedDefaultBuilderAndEmpty) { + ::flatbuffers::FlatBufferBuilder foreign(8); + FinishPayload(foreign, std::string(8192, 'x')); + MessageBuilder imported(std::move(foreign)); + ExpectPayload(imported.ReleaseMessage(), 8192); + ::flatbuffers::FlatBufferBuilder empty; + imported = std::move(empty); + FinishPayload(imported, std::string(16, 'x')); + ExpectPayload(imported.ReleaseMessage(), 16); + FinishPayload(imported, std::string(63, 'x')); + ::flatbuffers::FlatBufferBuilder& alias = imported; + imported = std::move(alias); + ExpectPayload(imported.ReleaseMessage(), 63); +} + +TEST(FlatbuffersTest, AllocatorFrontBackAndMove) { + SlabAllocator first; + uint8_t* data = first.allocate(64); + EXPECT_EQ(0u, reinterpret_cast(data) % + brpc::flatbuffers::kBufferAlignment); + memset(data, 'f', 16); + memset(data + 48, 'b', 16); + SlabAllocator second(std::move(first)); + size_t previous = 64; + for (size_t next : {128u, 512u, 8192u, 16384u}) { + data = second.reallocate_downward(data, previous, next, 16, 16); + EXPECT_EQ(std::string(16, 'f'), std::string((char*)data, 16)); + EXPECT_EQ(std::string(16, 'b'), std::string((char*)data + next - 16, 16)); + previous = next; + } + first = std::move(second); + SlabAllocator* alias = &first; + first = std::move(*alias); + first.deallocate(data, 16384); + second.deallocate(nullptr, 0); + data = second.allocate(32); + second.deallocate(data, 32); +} + +TEST(FlatbuffersTest, SerializeConstAndRetainWireStorage) { + butil::IOBuf wire; + { + const Message msg = MakeMessage(8192); + ASSERT_TRUE(brpc::flatbuffers::SerializeFbToIOBUF(&msg, wire)); + EXPECT_EQ(msg.size() + msg.get_meta_size(), wire.size()); + } + Message parsed; + const size_t meta = brpc::flatbuffers::kDefaultMetaSize; + ASSERT_TRUE(brpc::flatbuffers::ParseFbFromIOBUF( + &parsed, wire.size() - meta, wire, meta)); + wire.clear(); + ExpectPayload(parsed, 8192); +} + +TEST(FlatbuffersTest, ParseAlignedStorageIsSharedAndRetained) { + Message source = MakeMessage(); + void* raw = nullptr; + ASSERT_EQ(0, posix_memalign(&raw, brpc::flatbuffers::kBufferAlignment, + source.size())); + memcpy(raw, source.data(), source.size()); + int frees = 0; + butil::IOBuf wire; + ASSERT_EQ(0, wire.append_user_data(raw, source.size(), [&frees](void* p) { + ++frees; + free(p); + })); + Message parsed; + ASSERT_TRUE(parsed.parse_msg_from_iobuf(wire, source.size(), 0)); + EXPECT_EQ(raw, parsed.data()); + wire.clear(); + EXPECT_EQ(0, frees); + ExpectPayload(parsed, 16); + parsed.Clear(); + EXPECT_EQ(1, frees); +} + +TEST(FlatbuffersTest, ParseFragmentedUnalignedAndReplace) { + for (size_t length : {16u, 8192u, 32768u}) { + Message source = MakeMessage(length); + butil::IOBuf wire; + const size_t split = source.size() / 2; + void* first = malloc(split); + void* second = malloc(source.size() - split); + memcpy(first, source.data(), split); + memcpy(second, source.data() + split, source.size() - split); + ASSERT_EQ(0, wire.append_user_data(first, split, free)); + ASSERT_EQ(0, wire.append_user_data(second, source.size() - split, free)); + ASSERT_EQ(2u, wire.backing_block_num()); + Message parsed = MakeMessage(1); + ASSERT_TRUE(parsed.parse_msg_from_iobuf(wire, source.size(), 0)); + ASSERT_TRUE(parsed.parse_msg_from_iobuf(wire, source.size(), 0)); + wire.clear(); + ExpectPayload(parsed, length); + EXPECT_EQ(0u, reinterpret_cast(parsed.data()) % + brpc::flatbuffers::kBufferAlignment); + + char* bytes = static_cast(malloc(source.size() + 1)); + bytes[0] = 'h'; + memcpy(bytes + 1, source.data(), source.size()); + ASSERT_EQ(0, wire.append_user_data(bytes, source.size() + 1, free)); + ASSERT_TRUE(parsed.parse_msg_from_iobuf(wire, source.size(), 1)); + EXPECT_NE(bytes + 1, reinterpret_cast(parsed.data())); + wire.clear(); + ExpectPayload(parsed, length); + EXPECT_EQ(1u, parsed.get_meta_size()); + } +} + +TEST(FlatbuffersTest, ParseRejectsBadFramingWithoutChangingMessage) { + Message msg = MakeMessage(16); + const uint8_t* data = msg.data(); + butil::IOBuf buf; + buf.append("abcd"); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, 0, 4)); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, 4, 1)); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, 5, 0)); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, 1, 5)); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, std::numeric_limits::max(), 5)); + EXPECT_FALSE(msg.parse_msg_from_iobuf(buf, 5, std::numeric_limits::max())); + EXPECT_EQ(data, msg.data()); + ExpectPayload(msg, 16); +} + +TEST(FlatbuffersTest, InvalidPayloadAndOptionalString) { + Message msg; + for (size_t size : {1u, 3u, 4u, 16u, 64u}) { + butil::IOBuf buf; + buf.append(std::string(size, '\xff')); + ASSERT_TRUE(msg.parse_msg_from_iobuf(buf, size, 0)); + EXPECT_FALSE(msg.Verify()); + } + MessageBuilder builder; + builder.Finish(brpc_fbtest::CreatePayload(builder)); + msg = builder.ReleaseMessage(); + ASSERT_TRUE(msg.Verify()); + EXPECT_EQ(nullptr, msg.GetRoot()->message()); + // The schema allows an absent string; callers must not dereference it. + msg = MakeMessage(); + memset(msg.mutable_data(), 0xff, sizeof(::flatbuffers::uoffset_t)); + EXPECT_FALSE(msg.Verify()); +} + +TEST(FlatbuffersTest, ShrinkMetadataDoesNotMovePayload) { + Message msg = MakeMessage(16); + const uint8_t* data = msg.data(); + EXPECT_EQ(nullptr, msg.reduce_meta_size_and_get_buf(65)); + ASSERT_NE(nullptr, msg.reduce_meta_size_and_get_buf(36)); + memset(msg.mutable_buf_begin(), 'm', 36); + EXPECT_EQ(data, msg.data()); + ExpectPayload(msg, 16); + EXPECT_EQ(data, msg.reduce_meta_size_and_get_buf(0)); + EXPECT_EQ(data, msg.reduce_meta_size_and_get_buf(0)); + EXPECT_EQ(nullptr, msg.reduce_meta_size_and_get_buf(1)); + EXPECT_EQ(0u, msg.get_meta_size()); + butil::IOBuf wire; + ASSERT_TRUE(msg.append_msg_to_iobuf(wire)); + EXPECT_EQ(msg.size(), wire.size()); +} + +void* FailBlockAllocation(size_t) { return nullptr; } + +class FlatbuffersDeathTest : public ::testing::Test { +protected: + void SetUp() override { + _old_style = GTEST_FLAG_GET(death_test_style); + GTEST_FLAG_SET(death_test_style, "threadsafe"); +#if !BRPC_WITH_GLOG + _old = logging::FLAGS_crash_on_fatal_log; + logging::FLAGS_crash_on_fatal_log = false; +#endif + } + void TearDown() override { + GTEST_FLAG_SET(death_test_style, _old_style); +#if !BRPC_WITH_GLOG + logging::FLAGS_crash_on_fatal_log = _old; +#endif + } +private: + bool _old = false; + std::string _old_style; +}; + +TEST_F(FlatbuffersDeathTest, AllocationFailureIsChecked) { + EXPECT_DEATH({ + butil::iobuf::blockmem_allocate = FailBlockAllocation; + SlabAllocator allocator; + allocator.allocate(1024 * 1024); + }, "Fail to allocate"); + EXPECT_DEATH({ + SlabAllocator allocator; + uint8_t* data = allocator.allocate(1024 * 1024); + butil::iobuf::blockmem_allocate = FailBlockAllocation; + allocator.reallocate_downward(data, 1024 * 1024, 2 * 1024 * 1024, 8, 8); + }, "Fail to allocate"); +} + +TEST_F(FlatbuffersDeathTest, RejectsUnsafeAllocatorInputsAndUnfinishedRelease) { + EXPECT_DEATH({ SlabAllocator a; a.allocate(0); }, ""); + EXPECT_DEATH({ SlabAllocator a; a.allocate(std::numeric_limits::max()); }, ""); + EXPECT_DEATH({ + SlabAllocator a; + uint8_t* p = a.allocate(8); + a.reallocate_downward(p, 8, 16, 8, 1); + }, ""); + EXPECT_DEATH({ MessageBuilder builder; builder.ReleaseMessage(); }, "Finish"); +} + +TEST(FlatbuffersDescriptorTest, SparseStableIDsAndCanonicalNames) { + ServiceDescriptor first; + ASSERT_EQ(0, first.init({"test.", "Benchmark", "One Two Three", {2, 5, 1}})); + EXPECT_EQ("test.Benchmark", first.full_name()); + EXPECT_EQ(3, first.method_count()); + EXPECT_EQ(2, first.method(0)->index()); + EXPECT_EQ(5, first.method(1)->index()); + EXPECT_EQ("test.Benchmark.Two", first.method(1)->full_name()); + EXPECT_EQ(&first, first.method(0)->service()); + EXPECT_EQ(first.method(1), first.FindMethodByIndex(5)); + EXPECT_EQ(nullptr, first.FindMethodByIndex(3)); + EXPECT_EQ(nullptr, first.method(-1)); + EXPECT_EQ(nullptr, first.method(3)); + ServiceDescriptor reduced; + ASSERT_EQ(0, reduced.init({"test", "Benchmark", "Three Two", {1, 5}})); + EXPECT_EQ(first.index(), reduced.index()); + EXPECT_EQ(first.FindMethodByIndex(5)->index(), reduced.method(1)->index()); + EXPECT_EQ(first.FindMethodByIndex(5)->full_name(), reduced.method(1)->full_name()); + EXPECT_NE(0, first.init({"test", "Other", "Call", {3}})); + EXPECT_EQ("test.Benchmark", first.full_name()); +} + +TEST(FlatbuffersDescriptorTest, LegacyGlobalAndWhitespace) { + ServiceDescriptor desc; + ASSERT_EQ(0, desc.init({"", "Global", " First\tSecond\n Third ", {}})); + EXPECT_EQ("Global", desc.full_name()); + EXPECT_EQ(3, desc.method_count()); + EXPECT_EQ(0, desc.method(0)->index()); + EXPECT_EQ(2, desc.method(2)->index()); + EXPECT_EQ("Global.First", desc.method(0)->full_name()); +} + +TEST(FlatbuffersDescriptorTest, InvalidTablesAndFailureOwnership) { + const std::vector invalid = { + {"test", "", "A", {}}, {"test", "Service", "", {}}, + {"test", "Service", "A A", {1, 2}}, + {"test", "Service", "A B", {1, 1}}, + {"test", "Service", "A B", {1}}, + {"test", "Service", "A B", {-1, 2}}, + {"test..", "Service", "A", {}}, {".", "Service", "A", {}}, + {"test", "Bad.Name", "A", {}}, {"test", "Service", "1Bad", {}} + }; + ServiceDescriptor desc; + for (const auto& table : invalid) { + EXPECT_NE(0, desc.init(table)); + EXPECT_EQ(0, desc.method_count()); + ServiceDescriptor* output = &desc; + EXPECT_NE(0, brpc::flatbuffers::parse_service_descriptors(table, &output)); + EXPECT_EQ(&desc, output); + } + ASSERT_EQ(0, desc.init({"test", "Service", "A", {0}})); + EXPECT_NE(0, brpc::flatbuffers::parse_service_descriptors( + {"test", "Service", "A", {0}}, nullptr)); + ServiceDescriptor* output = nullptr; + ASSERT_EQ(0, brpc::flatbuffers::parse_service_descriptors( + {"test", "Service", "A", {0}}, &output)); + std::unique_ptr owner(output); + EXPECT_EQ("test.Service", owner->full_name()); +} + +} // namespace +#endif // BRPC_WITH_FLATBUFFERS diff --git a/test/flatbuffers_codegen/CMakeLists.txt b/test/flatbuffers_codegen/CMakeLists.txt new file mode 100644 index 0000000000..cd3436e469 --- /dev/null +++ b/test/flatbuffers_codegen/CMakeLists.txt @@ -0,0 +1,148 @@ +# 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. + +set(protobuf_MODULE_COMPATIBLE ON CACHE BOOL "Expose Protobuf compatibility variables for codegen tests" FORCE) +find_package(Protobuf CONFIG QUIET) +set(CODEGEN_PROTOBUF_CONFIG ${Protobuf_FOUND}) +if(NOT Protobuf_FOUND) + find_package(Protobuf REQUIRED) +endif() +set(CODEGEN_CXX_STANDARD 14) +if(Protobuf_VERSION VERSION_GREATER 4.21) + set(CODEGEN_CXX_STANDARD 17) + if(NOT CODEGEN_PROTOBUF_CONFIG) + message(FATAL_ERROR "Protobuf 5+ needs its CMake config package for codegen tests") + endif() +endif() +message(STATUS "Codegen tests: Protobuf ${Protobuf_VERSION}, C++${CODEGEN_CXX_STANDARD}") +find_program(FLATC_EXECUTABLE flatc) +if(NOT FLATC_EXECUTABLE) + message(FATAL_ERROR "Codegen tests need the official flatc executable") +endif() +execute_process(COMMAND "${FLATC_EXECUTABLE}" --version + OUTPUT_VARIABLE CODEGEN_FLATC_VERSION RESULT_VARIABLE CODEGEN_FLATC_RESULT + OUTPUT_STRIP_TRAILING_WHITESPACE) +if(NOT CODEGEN_FLATC_RESULT EQUAL 0 OR + NOT CODEGEN_FLATC_VERSION STREQUAL "flatc version 25.2.10") + message(FATAL_ERROR "Codegen tests require official flatc 25.2.10: ${CODEGEN_FLATC_VERSION}") +endif() +find_path(CODEGEN_GFLAGS_INCLUDE_DIR gflags/gflags_declare.h) +if(NOT CODEGEN_GFLAGS_INCLUDE_DIR) + message(FATAL_ERROR "Codegen tests need gflags development headers") +endif() +# execute_process does not inherit a target's compiler flags or macOS SDK. +set(CODEGEN_CXX_FLAGS "${CMAKE_CXX_COMPILER_ARG1} ${CMAKE_CXX_FLAGS}") +foreach(config DEBUG RELEASE RELWITHDEBINFO MINSIZEREL) + string(APPEND CODEGEN_CXX_FLAGS + " $<$:${CMAKE_CXX_FLAGS_${config}}>") +endforeach() +if(CMAKE_SYSROOT) + string(APPEND CODEGEN_CXX_FLAGS " --sysroot=\"${CMAKE_SYSROOT}\"") +endif() +if(APPLE) + if(CMAKE_OSX_SYSROOT) + set(CODEGEN_OSX_SYSROOT "${CMAKE_OSX_SYSROOT}") + if(NOT IS_ABSOLUTE "${CODEGEN_OSX_SYSROOT}") + execute_process(COMMAND xcrun --sdk "${CMAKE_OSX_SYSROOT}" --show-sdk-path + OUTPUT_VARIABLE CODEGEN_OSX_SYSROOT RESULT_VARIABLE sdk_result + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT sdk_result EQUAL 0 OR NOT IS_DIRECTORY "${CODEGEN_OSX_SYSROOT}") + message(FATAL_ERROR "Cannot resolve codegen SDK: ${CMAKE_OSX_SYSROOT}") + endif() + endif() + string(APPEND CODEGEN_CXX_FLAGS " -isysroot \"${CODEGEN_OSX_SYSROOT}\"") + endif() + foreach(arch IN LISTS CMAKE_OSX_ARCHITECTURES) + string(APPEND CODEGEN_CXX_FLAGS " -arch ${arch}") + endforeach() + if(CMAKE_OSX_DEPLOYMENT_TARGET) + string(APPEND CODEGEN_CXX_FLAGS + " -mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") + endif() +endif() +message(STATUS "Codegen flatc: ${FLATC_EXECUTABLE} (${CODEGEN_FLATC_VERSION})") +get_filename_component(BRPC_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../.." ABSOLUTE) +set(GENERATED_DIR "${CMAKE_CURRENT_BINARY_DIR}/generated") +file(MAKE_DIRECTORY "${GENERATED_DIR}" "${CMAKE_CURRENT_BINARY_DIR}/include/butil") +# Keep test-only configuration in this build tree, never in the repository. +set(WITH_FLATBUFFERS_VAL 1) +configure_file("${BRPC_SOURCE_DIR}/src/butil/config.h.in" + "${CMAKE_CURRENT_BINARY_DIR}/include/butil/config.h") +set(CODEGEN_INCLUDE_DIRS + "${CMAKE_CURRENT_BINARY_DIR}/include" + "${BRPC_SOURCE_DIR}/src" + "${GENERATED_DIR}" + "${FLATBUFFERS_INCLUDE_DIR}" + "${CODEGEN_GFLAGS_INCLUDE_DIR}" + ${Protobuf_INCLUDE_DIRS}) +get_target_property(CODEGEN_PROTOBUF_INCLUDE_DIRS protobuf::libprotobuf INTERFACE_INCLUDE_DIRECTORIES) +if(CODEGEN_PROTOBUF_INCLUDE_DIRS) + list(APPEND CODEGEN_INCLUDE_DIRS ${CODEGEN_PROTOBUF_INCLUDE_DIRS}) +endif() +if(TARGET absl::base) + get_target_property(CODEGEN_ABSL_INCLUDE_DIRS absl::base INTERFACE_INCLUDE_DIRECTORIES) + if(CODEGEN_ABSL_INCLUDE_DIRS) + list(APPEND CODEGEN_INCLUDE_DIRS ${CODEGEN_ABSL_INCLUDE_DIRS}) + endif() +endif() +add_custom_command( + OUTPUT "${GENERATED_DIR}/echo.brpc.fb.cpp" + "${GENERATED_DIR}/echo.brpc.fb.h" + "${GENERATED_DIR}/echo_generated.h" + COMMAND "${FLATC_EXECUTABLE}" --cpp -o "${GENERATED_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/echo.fbs" + COMMAND brpc_flatc -o "${GENERATED_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/echo.fbs" + DEPENDS brpc_flatc "${CMAKE_CURRENT_SOURCE_DIR}/echo.fbs" + VERBATIM) +add_library(brpc_codegen_compile OBJECT "${GENERATED_DIR}/echo.brpc.fb.cpp") +target_include_directories(brpc_codegen_compile PRIVATE ${CODEGEN_INCLUDE_DIRS}) +target_compile_features(brpc_codegen_compile PRIVATE cxx_std_${CODEGEN_CXX_STANDARD}) + +add_test(NAME flatbuffers_codegen_acceptance + COMMAND "${CMAKE_COMMAND}" + "-DGENERATOR=$" + "-DFLATC=${FLATC_EXECUTABLE}" + "-DSCHEMA=${CMAKE_CURRENT_SOURCE_DIR}/echo.fbs" + "-DWORK=${CMAKE_CURRENT_BINARY_DIR}/acceptance" + "-DCXX=${CMAKE_CXX_COMPILER}" + "-DCXX_FLAGS=${CODEGEN_CXX_FLAGS}" + "-DCXX_ID=${CMAKE_CXX_COMPILER_ID}" + "-DIMPLICIT_INCLUDE_DIRS=${CMAKE_CXX_IMPLICIT_INCLUDE_DIRECTORIES}" + "-DCXX_STANDARD=${CODEGEN_CXX_STANDARD}" + "-DINCLUDE_DIRS=${CODEGEN_INCLUDE_DIRS}" + -P "${CMAKE_CURRENT_SOURCE_DIR}/acceptance.cmake") + +set(BRPC_CODEGEN_BRPC_LIBRARY "" CACHE FILEPATH + "Optional existing FlatBuffers-enabled libbrpc for runtime acceptance") +set(BRPC_CODEGEN_EXTRA_LIBRARIES "" CACHE STRING + "Additional link dependencies of the supplied libbrpc") +if(BRPC_CODEGEN_BRPC_LIBRARY) + find_package(Threads REQUIRED) + find_package(OpenSSL REQUIRED) + find_package(ZLIB REQUIRED) + find_library(CODEGEN_GFLAGS_LIBRARY gflags) + find_library(CODEGEN_LEVELDB_LIBRARY leveldb) + add_executable(brpc_codegen_runtime runtime.cpp $) + target_include_directories(brpc_codegen_runtime PRIVATE ${CODEGEN_INCLUDE_DIRS}) + target_compile_features(brpc_codegen_runtime PRIVATE cxx_std_${CODEGEN_CXX_STANDARD}) + target_link_libraries(brpc_codegen_runtime PRIVATE + "${BRPC_CODEGEN_BRPC_LIBRARY}" protobuf::libprotobuf + ${CODEGEN_GFLAGS_LIBRARY} ${CODEGEN_LEVELDB_LIBRARY} + OpenSSL::SSL OpenSSL::Crypto ZLIB::ZLIB Threads::Threads + ${CMAKE_DL_LIBS} ${BRPC_CODEGEN_EXTRA_LIBRARIES}) + add_test(NAME flatbuffers_codegen_runtime COMMAND brpc_codegen_runtime) +endif() diff --git a/test/flatbuffers_codegen/acceptance.cmake b/test/flatbuffers_codegen/acceptance.cmake new file mode 100644 index 0000000000..ee8d363699 --- /dev/null +++ b/test/flatbuffers_codegen/acceptance.cmake @@ -0,0 +1,809 @@ +# 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. + +if(NOT DEFINED CXX_STANDARD) + set(CXX_STANDARD 14) +endif() +separate_arguments(compiler_flags NATIVE_COMMAND "${CXX_FLAGS}") +message(STATUS "Codegen compiler: ${CXX}; flags: ${CXX_FLAGS}; C++${CXX_STANDARD}") +file(READ "${SCHEMA}" schema_text) +file(MAKE_DIRECTORY "${WORK}") + +function(expect_rejected name replacement) + set(directory "${WORK}/${name}") + file(MAKE_DIRECTORY "${directory}") + string(REPLACE "(id: 7)" "${replacement}" text "${schema_text}") + file(WRITE "${directory}/echo.fbs" "${text}") + execute_process(COMMAND "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if("${result}" STREQUAL "0" OR error STREQUAL "") + message(FATAL_ERROR "${name}: invalid schema was accepted or lacked diagnostic") + endif() + if(EXISTS "${directory}/echo.brpc.fb.h" OR EXISTS "${directory}/echo.brpc.fb.cpp") + message(FATAL_ERROR "${name}: invalid schema emitted service files") + endif() + message(STATUS "Rejected ${name}: ${error}") +endfunction() + +function(expect_schema_rejected name text diagnostic) + set(directory "${WORK}/${name}") + file(REMOVE_RECURSE "${directory}") + file(MAKE_DIRECTORY "${directory}") + file(WRITE "${directory}/echo.fbs" "${text}") + execute_process(COMMAND "${FLATC}" --cpp -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE flatc_result ERROR_VARIABLE flatc_error) + if(NOT "${flatc_result}" STREQUAL "0") + message(FATAL_ERROR "${name}: not a valid official flatc input: ${flatc_error}") + endif() + execute_process(COMMAND "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + string(FIND "${error}" "${diagnostic}" diagnostic_position) + if("${result}" STREQUAL "0" OR diagnostic_position LESS 0) + message(FATAL_ERROR "${name}: schema was accepted or lacked diagnostic: ${error}") + endif() + if(EXISTS "${directory}/echo.brpc.fb.h" OR EXISTS "${directory}/echo.brpc.fb.cpp") + message(FATAL_ERROR "${name}: rejected schema emitted service files") + endif() + file(WRITE "${directory}/echo.brpc.fb.h" "previous header\n") + file(WRITE "${directory}/echo.brpc.fb.cpp" "previous source\n") + execute_process(COMMAND "${GENERATOR}" -o "${directory}" + "${directory}/echo.fbs" RESULT_VARIABLE result ERROR_VARIABLE error) + string(FIND "${error}" "${diagnostic}" diagnostic_position) + file(READ "${directory}/echo.brpc.fb.h" old_header) + file(READ "${directory}/echo.brpc.fb.cpp" old_source) + if("${result}" STREQUAL "0" OR diagnostic_position LESS 0 OR + NOT "${old_header}" STREQUAL "previous header\n" OR + NOT "${old_source}" STREQUAL "previous source\n") + message(FATAL_ERROR "${name}: rejection changed existing output files") + endif() + message(STATUS "Rejected ${name} without changing outputs: ${error}") +endfunction() + +function(prepare_output_pair directory) + file(REMOVE_RECURSE "${directory}") + file(MAKE_DIRECTORY "${directory}") + file(WRITE "${directory}/echo.fbs" "${schema_text}") + file(WRITE "${directory}/echo.brpc.fb.h" "previous header\n") + file(WRITE "${directory}/echo.brpc.fb.cpp" "previous source\n") +endfunction() + +function(expect_output_pair_restored name failure) + set(directory "${WORK}/${name}") + prepare_output_pair("${directory}") + execute_process(COMMAND "${CMAKE_COMMAND}" -E env + "BRPC_FLATC_TEST_FAILURE=${failure}" + "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if("${result}" STREQUAL "0" OR error STREQUAL "") + message(FATAL_ERROR "${name}: injected failure was not reported") + endif() + file(READ "${directory}/echo.brpc.fb.h" current_header) + file(READ "${directory}/echo.brpc.fb.cpp" current_source) + if(NOT current_header STREQUAL "previous header\n" OR + NOT current_source STREQUAL "previous source\n") + message(FATAL_ERROR "${name}: previous output pair was not restored") + endif() + file(GLOB transaction_files "${directory}/echo.brpc.fb.*.tmp" + "${directory}/echo.brpc.fb.*.bak") + if(transaction_files) + message(FATAL_ERROR "${name}: transaction files remained after rollback") + endif() + message(STATUS "Restored output pair after ${failure}: ${error}") +endfunction() + +function(expect_output_pair_removed name failure remaining_backup) + set(directory "${WORK}/${name}") + prepare_output_pair("${directory}") + execute_process(COMMAND "${CMAKE_COMMAND}" -E env + "BRPC_FLATC_TEST_FAILURE=${failure}" + "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if("${result}" STREQUAL "0" OR error STREQUAL "") + message(FATAL_ERROR "${name}: injected restore failure was not reported") + endif() + if(EXISTS "${directory}/echo.brpc.fb.h" OR EXISTS "${directory}/echo.brpc.fb.cpp") + message(FATAL_ERROR "${name}: partial final output remained") + endif() + if(NOT EXISTS "${directory}/echo.brpc.fb.${remaining_backup}.bak") + message(FATAL_ERROR "${name}: recoverable backup was not retained") + endif() + file(GLOB temporary_outputs "${directory}/echo.brpc.fb.*.tmp") + if(temporary_outputs) + message(FATAL_ERROR "${name}: temporary files remained") + endif() + message(STATUS "Removed final pair after ${failure}: ${error}") +endfunction() + +function(expect_concurrent_output_pair) + set(directory "${WORK}/concurrent_output_pair") + set(output "${directory}/output") + file(REMOVE_RECURSE "${directory}") + file(MAKE_DIRECTORY "${directory}/first" "${directory}/second" "${output}") + set(first_schema + "namespace codegen.concurrent;\ntable Request {}\ntable Response {}\nrpc_service Echo { First(Request):Response (id: 1); }\n") + set(second_schema + "namespace codegen.concurrent;\ntable Request {}\ntable Response {}\nrpc_service Echo { Second(Request):Response (id: 2); }\n") + file(WRITE "${directory}/first/echo.fbs" "${first_schema}") + file(WRITE "${directory}/second/echo.fbs" "${second_schema}") + file(WRITE "${directory}/run.sh" + "#!/bin/sh\n\"${GENERATOR}\" -o \"${output}\" \"${directory}/first/echo.fbs\" &\nfirst=\$!\n\"${GENERATOR}\" -o \"${output}\" \"${directory}/second/echo.fbs\" &\nsecond=\$!\nwait \$first\nfirst_result=\$?\nwait \$second\nsecond_result=\$?\ntest \$first_result -eq 0 -a \$second_result -eq 0\n") + execute_process(COMMAND /bin/sh "${directory}/run.sh" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "concurrent_output_pair: generation failed: ${error}") + endif() + execute_process(COMMAND "${FLATC}" --cpp -o "${output}" + "${directory}/first/echo.fbs" RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "concurrent_output_pair: official flatc failed: ${error}") + endif() + set(includes "-I${output}") + foreach(include_dir IN LISTS INCLUDE_DIRS) + list(APPEND includes "-I${include_dir}") + endforeach() + execute_process(COMMAND "${CXX}" ${compiler_flags} "-std=c++${CXX_STANDARD}" ${includes} + -c "${output}/echo.brpc.fb.cpp" -o "${directory}/echo.o" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "concurrent_output_pair: mixed generated files: ${error}") + endif() + file(GLOB transaction_files "${output}/echo.brpc.fb.*.tmp" + "${output}/echo.brpc.fb.*.bak") + if(transaction_files) + message(FATAL_ERROR "concurrent_output_pair: transaction files remained") + endif() + message(STATUS "Serialized concurrent output publication") +endfunction() + +function(expect_rejected_rpc_name name) + set(directory "${WORK}/reserved_${name}") + file(MAKE_DIRECTORY "${directory}") + file(REMOVE "${directory}/echo.brpc.fb.h" "${directory}/echo.brpc.fb.cpp") + string(REPLACE "Inspect(" "${name}(" text "${schema_text}") + file(WRITE "${directory}/echo.fbs" "${text}") + execute_process(COMMAND "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if("${result}" STREQUAL "0") + message(SEND_ERROR "RPC ${name}: generator incorrectly accepted a fixed Stub field name") + return() + endif() + if(NOT error MATCHES "method name collides with generated API: ${name}") + message(FATAL_ERROR "RPC ${name}: expected a name-collision diagnostic: ${error}") + endif() + if(EXISTS "${directory}/echo.brpc.fb.h" OR EXISTS "${directory}/echo.brpc.fb.cpp") + message(FATAL_ERROR "RPC ${name}: rejected schema emitted service files") + endif() + message(STATUS "Rejected RPC ${name}: ${error}") +endfunction() + +function(expect_compiles name text) + set(standard "${CXX_STANDARD}") + if(ARGC GREATER 2) + set(standard "${ARGV2}") + endif() + set(directory "${WORK}/${name}") + file(MAKE_DIRECTORY "${directory}") + file(WRITE "${directory}/echo.fbs" "${text}") + execute_process(COMMAND "${FLATC}" --cpp -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: official flatc failed: ${error}") + endif() + execute_process(COMMAND "${GENERATOR}" -o "${directory}" "${directory}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: brpc_flatc failed: ${error}") + endif() + set(includes "-I${directory}") + set(diagnostics) + foreach(include_dir IN LISTS INCLUDE_DIRS) + if(ARGC GREATER 3) + list(FIND IMPLICIT_INCLUDE_DIRS "${include_dir}" implicit_index) + if(implicit_index GREATER_EQUAL 0) + continue() + endif() + list(APPEND includes -isystem "${include_dir}") + else() + list(APPEND includes "-I${include_dir}") + endif() + endforeach() + if(ARGC GREATER 3) + set(diagnostics ${ARGV3}) + endif() + execute_process(COMMAND "${CXX}" ${compiler_flags} ${diagnostics} "-std=c++${standard}" ${includes} + -c "${directory}/echo.brpc.fb.cpp" -o "${directory}/echo.o" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: generated code did not compile: ${error}") + endif() + # The generated header must also compile without incidental prior includes. + file(WRITE "${directory}/header.cpp" "#include \"echo.brpc.fb.h\"\n") + execute_process(COMMAND "${CXX}" ${compiler_flags} ${diagnostics} "-std=c++${standard}" ${includes} + -c "${directory}/header.cpp" -o "${directory}/header.o" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: generated header is not self-contained: ${error}") + endif() + foreach(suffix brpc.fb.h brpc.fb.cpp) + file(READ "${directory}/echo.${suffix}" generated) + if(NOT generated MATCHES "Licensed to the Apache Software Foundation") + message(FATAL_ERROR "${name}: ${suffix} lacks Apache license") + endif() + if(generated MATCHES "FLATBUFFERS_VERSION") + message(FATAL_ERROR "${name}: service output pins FlatBuffers version") + endif() + endforeach() + message(STATUS "Compiled ${name}") +endfunction() + +function(expect_distinct_headers name first_stem second_stem) + set(directory "${WORK}/header_guards_${name}") + set(includes) + foreach(include_dir IN LISTS INCLUDE_DIRS) + list(APPEND includes "-I${include_dir}") + endforeach() + foreach(side first second) + set(stem "${${side}_stem}") + set(output "${directory}/${side}") + file(MAKE_DIRECTORY "${output}") + string(REPLACE "namespace codegen.example;" "namespace guard_${name}.${side};" + text "${schema_text}") + # Isolate the fixture's global helper table as well as its services. + string(PREPEND text "namespace guard_${name}.${side};\n") + if(side STREQUAL "first") + set(first_text "${text}") + endif() + file(WRITE "${output}/${stem}.fbs" "${text}") + execute_process(COMMAND "${FLATC}" --cpp -o "${output}" "${output}/${stem}.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}/${side}: official flatc failed: ${error}") + endif() + execute_process(COMMAND "${GENERATOR}" -o "${output}" "${output}/${stem}.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}/${side}: brpc_flatc failed: ${error}") + endif() + file(WRITE "${output}/single.cpp" + "#include \"${stem}.brpc.fb.h\"\n::guard_${name}::${side}::Echo* service = nullptr;\n") + execute_process(COMMAND "${CXX}" ${compiler_flags} "-std=c++${CXX_STANDARD}" ${includes} + -c "${output}/single.cpp" -o "${output}/single.o" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}/${side}: standalone header failed: ${error}") + endif() + endforeach() + foreach(reverse FALSE TRUE) + set(first "#include \"first/${first_stem}.brpc.fb.h\"\n") + set(second "#include \"second/${second_stem}.brpc.fb.h\"\n") + if(reverse) + set(headers "${second}${first}") + else() + set(headers "${first}${second}") + endif() + file(WRITE "${directory}/combined_${reverse}.cpp" + "${headers}::guard_${name}::first::Echo* first_echo = nullptr;\n::guard_${name}::second::Echo* second_echo = nullptr;\n") + execute_process(COMMAND "${CXX}" ${compiler_flags} "-std=c++${CXX_STANDARD}" ${includes} + -c "${directory}/combined_${reverse}.cpp" + -o "${directory}/combined_${reverse}.o" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: combined headers (reverse=${reverse}) failed: ${error}") + endif() + endforeach() + # A checkout/output directory change must not alter generated identifiers. + set(relocated "${directory}/relocated") + file(MAKE_DIRECTORY "${relocated}") + file(WRITE "${relocated}/${first_stem}.fbs" "${first_text}") + execute_process(COMMAND "${GENERATOR}" -o . "${first_stem}.fbs" + WORKING_DIRECTORY "${relocated}" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${name}: relocated generation failed: ${error}") + endif() + foreach(suffix brpc.fb.h brpc.fb.cpp) + file(READ "${directory}/first/${first_stem}.${suffix}" original) + file(READ "${relocated}/${first_stem}.${suffix}" regenerated) + if(NOT "${original}" STREQUAL "${regenerated}") + message(FATAL_ERROR "${name}: ${suffix} depends on checkout/output paths") + endif() + endforeach() + message(STATUS "Distinct, relocatable header guards: ${name}") +endfunction() + +# Exercise escaped table names as RPC types, not only collision candidates. +foreach(type_name module import synchronized atomic_cancel atomic_commit atomic_noexcept) + set(type_schema "table ${type_name} { value:int; }\ntable Response {}\n") + set(service_schema "rpc_service Echo { Call(${type_name}):Response (id: 1); Reverse(Response):${type_name} (id: 2); }\n") + expect_compiles(escaped_rpc_global_${type_name} + "${type_schema}${service_schema}") + expect_compiles(escaped_rpc_nested_${type_name} + "namespace api.module;\n${type_schema}${service_schema}") + + set(imported_name escaped_rpc_imported_${type_name}) + set(imported_dir "${WORK}/${imported_name}") + file(MAKE_DIRECTORY "${imported_dir}") + file(WRITE "${imported_dir}/types.fbs" "namespace payload;\n${type_schema}") + execute_process(COMMAND "${FLATC}" --cpp -o "${imported_dir}" + "${imported_dir}/types.fbs" RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "${imported_name}: official flatc failed (${result}): ${error}") + endif() + expect_compiles(${imported_name} + "include \"types.fbs\"; namespace api; rpc_service Echo { Call(payload.${type_name}):payload.Response (id: 1); Reverse(payload.Response):payload.${type_name} (id: 2); }") +endforeach() +expect_compiles(escaped_rpc_near_name [=[ +namespace api; +table module_ { value:int; } +rpc_service Echo { Call(module_):module_ (id: 1); } +]=]) +expect_compiles(escaped_type_preserves_service_name [=[ +namespace api; +table module { value:int; } +rpc_service module { Call(module):module (id: 1); } +]=]) + +# Exercise the names emitted by official flatc, not only schema type names. +foreach(scope "" "api.") + if(scope STREQUAL "") + set(namespace_decl "") + set(qualified "::") + set(scope_id global) + else() + set(namespace_decl "namespace api;\n") + set(qualified "::api::") + set(scope_id nested) + endif() + set(tables "table Request {}\ntable Response {}\n") + foreach(kind enum_value enum_min enum_any enum_values enum_names enum_name + union_traits union_verify union_vector table_create table_direct) + if(kind MATCHES "^enum_") + set(definitions "enum Collision: byte { Echo_Stub = 0 }\n") + if(kind STREQUAL enum_value) + set(service Collision_Echo) + set(symbol Collision_Echo_Stub) + elseif(kind STREQUAL enum_min) + set(service Collision_MIN) + set(symbol Collision_MIN) + elseif(kind STREQUAL enum_any) + set(definitions "enum Collision: byte (bit_flags) { Flag = 0 }\n") + set(service Collision_ANY) + set(symbol Collision_ANY) + elseif(kind STREQUAL enum_values) + set(service EnumValuesCollision) + set(symbol EnumValuesCollision) + elseif(kind STREQUAL enum_names) + set(service EnumNamesCollision) + set(symbol EnumNamesCollision) + else() + set(service EnumNameCollision) + set(symbol EnumNameCollision) + endif() + elseif(kind MATCHES "^union_") + set(definitions "union Collision { Request }\n") + if(kind STREQUAL union_traits) + set(service CollisionTraits) + elseif(kind STREQUAL union_verify) + set(service VerifyCollision) + else() + set(service VerifyCollisionVector) + endif() + set(symbol "${service}") + elseif(kind STREQUAL table_create) + set(definitions "table Echo_Stub {}\n") + set(service CreateEcho) + set(symbol CreateEcho_Stub) + else() + set(definitions "table Echo { text:string; }\n") + set(service CreateEchoDirect) + set(symbol CreateEchoDirect) + endif() + expect_schema_rejected(${scope_id}_${kind} + "${namespace_decl}${tables}${definitions}rpc_service ${service} { Call(Request):Response (id: 1); }\n" + "generated service class name collides: ${qualified}${symbol}") + endforeach() +endforeach() + +# Pair each root/helper declaration with actual flatc output and legal names. +foreach(symbol GetRoot_Stub GetSizePrefixedRoot_Stub VerifyRoot_StubBuffer + VerifySizePrefixedRoot_StubBuffer FinishRoot_StubBuffer + FinishSizePrefixedRoot_StubBuffer Root_StubIdentifier + Root_StubBufferHasIdentifier SizePrefixedRoot_StubBufferHasIdentifier + Root_StubExtension) + expect_schema_rejected(root_helper_${symbol} + "table Request {} table Response {} table Root_Stub {} root_type Root_Stub; file_identifier \"TEST\"; file_extension \"bin\"; rpc_service ${symbol} { Call(Request):Response (id: 1); }" + "generated service class name collides: ::${symbol}") +endforeach() +file(WRITE "${WORK}/root_types.fbs" + "namespace data; table Request {} table Response {} table Root_Stub {} root_type Root_Stub; file_identifier \"TEST\"; file_extension \"bin\";") +expect_schema_rejected(imported_root_helper + "include \"../root_types.fbs\"; namespace data; rpc_service GetRoot { Call(Request):Response (id: 1); }" + "generated service class name collides: ::data::GetRoot_Stub") +expect_schema_rejected(escaped_table_builder + "table Request {} table Response {} table class {} rpc_service class_Builder { Call(Request):Response (id: 1); }" + "generated service class name collides: ::class_Builder") +expect_compiles(unescaped_table_builder_available + "table Request {} table Response {} table class {} rpc_service classBuilder { Call(Request):Response (id: 1); }") +expect_compiles(escaped_enum_value_available + "table Request {} table Response {} enum E:byte { module = 0 } rpc_service E_module { Call(Request):Response (id: 1); }") +expect_compiles(sparse_enum_no_names_table + "table Request {} table Response {} enum E:int { A = 0, B = 10 } rpc_service EnumNamesE { Call(Request):Response (id: 1); }") +expect_schema_rejected(dense_enum_names_table + "table Request {} table Response {} enum E:int { A = 0, B = 9 } rpc_service EnumNamesE { Call(Request):Response (id: 1); }" + "generated service class name collides: ::EnumNamesE") +expect_compiles(deprecated_string_no_direct + "table Request {} table Response {} table Echo { text:string (deprecated); } rpc_service CreateEchoDirect { Call(Request):Response (id: 1); }") +expect_compiles(union_alias_no_traits + "table Request {} table Response {} union E { A:Request, B:Request } rpc_service ETraits { Call(Request):Response (id: 1); }") +expect_compiles(nonroot_helper_available + "table Request {} table Response {} table Root_Stub {} rpc_service GetRoot { Call(Request):Response (id: 1); }") + +foreach(name __Echo Echo__Name _Echo _echo Echo_ channel_ owned_channel_) + expect_schema_rejected(reserved_service_${name} + "table Request {} table Response {} rpc_service ${name} { Call(Request):Response (id: 1); }" + "reserved C++ identifier") +endforeach() +foreach(name __Call _Call Call__Name) + expect_schema_rejected(reserved_method_${name} + "table Request {} table Response {} rpc_service Echo { ${name}(Request):Response (id: 1); }" + "reserved C++ identifier") +endforeach() +foreach(name __api _Api _api _) + expect_schema_rejected(reserved_namespace_${name} + "namespace ${name}; table Request {} table Response {} rpc_service Echo { Call(Request):Response (id: 1); }" + "reserved C++ identifier") +endforeach() +expect_compiles(global_helper_table [=[ +table BrpcFlatbuffersFail { value:int; } +table Request {} table Response {} +rpc_service Echo { Call(Request):Response (id: 1); } +]=]) +expect_compiles(global_helper_service [=[ +table Request {} table Response {} +rpc_service BrpcFlatbuffersFail { Call(Request):Response (id: 1); } +]=]) +expect_compiles(legal_underscore_names [=[ +namespace api._internal; +table Request {} table Response {} +rpc_service _echo { _call(Request):Response (id: 1); call_(Request):Response (id: 2); } +]=]) +if(CXX_ID MATCHES "Clang") + # Integer-only tables avoid unrelated reserved locals in upstream flatc. + expect_compiles(strict_generated_identifiers + "namespace api; table Request { value:int; } table Response { value:int; } rpc_service Echo { Call(Request):Response (id: 1); }" + "${CXX_STANDARD}" "-Wreserved-identifier;-Werror=reserved-identifier") +else() + message(STATUS "Reserved-name rejection tested; strict diagnostics require Clang") +endif() +expect_compiles(absent_direct_helper [=[ +namespace api; +table Request {} table Response {} table Echo { value:int; } +rpc_service CreateEchoDirect { Call(Request):Response (id: 1); } +]=]) +expect_compiles(cross_namespace_flatc_symbols [=[ +namespace data; +table Request {} table Response {} enum Collision:byte { Echo_Stub = 0 } +namespace api; +rpc_service Collision_Echo { Call(data.Request):data.Response (id: 1); } +]=]) + +foreach(name Stub descriptor GetDescriptor FBCallMethod) + string(CONCAT reserved_service_schema + "table Request {}\ntable Response {}\n" + "rpc_service ${name} { Call(Request):Response (id: 1); }\n") + expect_schema_rejected(service_member_${name} "${reserved_service_schema}" + "service name collides with generated API: ${name}") + expect_schema_rejected(nested_service_member_${name} + "namespace api;\n${reserved_service_schema}" + "service name collides with generated API: ${name}") +endforeach() + +foreach(name Echo Echo_Stub) + foreach(depth "" ".nested") + string(CONCAT namespace_collision_schema + "namespace api.${name}${depth};\n" + "table Request {}\ntable Response {}\n" + "namespace api;\n" + "rpc_service Echo {\n" + " Call(api.${name}${depth}.Request):" + "api.${name}${depth}.Response (id: 1);\n}\n") + expect_schema_rejected(namespace_collision_${name}${depth} + "${namespace_collision_schema}" + "generated service class name collides: ::api::${name}") + endforeach() +endforeach() +file(WRITE "${WORK}/namespace_types.fbs" + "namespace api.Echo; table Request {} table Response {}\n") +expect_schema_rejected(imported_namespace_collision + "include \"../namespace_types.fbs\"; namespace api; + rpc_service Echo { + Call(api.Echo.Request):api.Echo.Response (id: 1); + }" + "generated service class name collides: ::api::Echo") + +expect_distinct_headers(punctuation foo-bar foo_bar) +expect_distinct_headers(letter_case FooBar foobar) +expect_distinct_headers(same_basename echo echo) + +expect_rejected_rpc_name(channel_) +expect_rejected_rpc_name(owned_channel_) +string(REPLACE "rpc_service Echo {" "rpc_service char8_t {" + char8_t_schema "${schema_text}") +expect_schema_rejected(cpp20_char8_t "${char8_t_schema}" + "C++ keyword is not supported: char8_t") +foreach(name brpc butil flatbuffers google std) + set(global_namespace_collision_schema + "table Request {}\ntable Response {}\nrpc_service ${name} { Call(Request):Response (id: 1); }\n") + expect_schema_rejected(global_namespace_${name} + "${global_namespace_collision_schema}" + "generated service class name collides: ::${name}") +endforeach() +set(service_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\nrpc_service Echo { Call(Request):Response (id: 1); }\nrpc_service Echo_Stub { Call(Request):Response (id: 2); }\n") +expect_schema_rejected(service_stub_collision "${service_collision_schema}" + "generated service class name collides: ::codegen::collision::Echo_Stub") +set(stub_service_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\nrpc_service Echo_Stub { Call(Request):Response (id: 2); }\nrpc_service Echo { Call(Request):Response (id: 1); }\n") +expect_schema_rejected(stub_service_collision "${stub_service_collision_schema}" + "generated service class name collides: ::codegen::collision::Echo_Stub") +set(table_service_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\ntable Echo {}\nrpc_service Echo { Call(Request):Response (id: 1); }\n") +expect_schema_rejected(table_service_collision "${table_service_collision_schema}" + "generated service class name collides: ::codegen::collision::Echo") +set(table_stub_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\ntable Echo_Stub {}\nrpc_service Echo { Call(Request):Response (id: 1); }\n") +expect_schema_rejected(table_stub_collision "${table_stub_collision_schema}" + "generated service class name collides: ::codegen::collision::Echo_Stub") +set(table_builder_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\ntable Echo {}\nrpc_service EchoBuilder { Call(Request):Response (id: 1); }\n") +expect_schema_rejected(table_builder_collision "${table_builder_collision_schema}" + "generated service class name collides: ::codegen::collision::EchoBuilder") +set(enum_service_collision_schema + "namespace codegen.collision;\ntable Request {}\ntable Response {}\nenum Echo : byte { Value = 0 }\nrpc_service Echo { Call(Request):Response (id: 1); }\n") +expect_schema_rejected(enum_service_collision "${enum_service_collision_schema}" + "generated service class name collides: ::codegen::collision::Echo") +expect_output_pair_restored(stage_source_failure stage_source) +expect_output_pair_restored(preserve_source_failure preserve_source) +expect_output_pair_restored(publish_header_failure publish_header) +expect_output_pair_restored(publish_source_failure publish_source) +expect_output_pair_removed(restore_header_failure + "publish_source,restore_header" h) +expect_output_pair_removed(restore_source_failure + "publish_source,restore_source" cpp) +expect_concurrent_output_pair() +expect_rejected(missing_id "") +expect_rejected(duplicate_id "(id: 41)") +expect_rejected(negative_id "(id: -1)") +expect_rejected(overflow_id "(id: 2147483648)") +expect_rejected(string_id "(id: \"7\")") +expect_rejected(streaming "(id: 7, streaming: \"server\")") +expect_compiles(sparse_ids "${schema_text}") +file(WRITE "${WORK}/cxx20_probe.cpp" "int main() { return 0; }\n") +execute_process(COMMAND "${CXX}" ${compiler_flags} -std=c++20 -fsyntax-only "${WORK}/cxx20_probe.cpp" + RESULT_VARIABLE cxx20_result ERROR_QUIET) +if("${cxx20_result}" STREQUAL "0") + expect_compiles(cxx20 "${schema_text}" 20) +else() + message(STATUS "Skipping C++20 compile smoke test: compiler lacks -std=c++20") +endif() +set(cross_namespace_schema + "namespace codegen.common;\ntable Request {}\ntable Response {}\nnamespace codegen.first;\nrpc_service Echo { Call(codegen.common.Request):codegen.common.Response (id: 1); }\nnamespace codegen.second;\nrpc_service Echo_Stub { Call(codegen.common.Request):codegen.common.Response (id: 2); }\n") +expect_compiles(cross_namespace_names "${cross_namespace_schema}") +set(cross_namespace_type_schema + "namespace codegen.common;\ntable Request {}\ntable Response {}\ntable Echo {}\nnamespace codegen.api;\nrpc_service Echo { Call(codegen.common.Request):codegen.common.Response (id: 1); }\n") +expect_compiles(cross_namespace_type_names "${cross_namespace_type_schema}") +expect_compiles(namespace_sibling [=[ +namespace api.other.Echo; +table Request {} +table Response {} +namespace api; +rpc_service Echo { + Call(api.other.Echo.Request):api.other.Echo.Response (id: 1); +} +]=]) +expect_compiles(service_member_near_names [=[ +namespace api; +table Request {} +table Response {} +rpc_service channel { Call(Request):Response (id: 1); } +rpc_service channel_handle { Call(Request):Response (id: 1); } +rpc_service owned_channel_handle { Call(Request):Response (id: 1); } +rpc_service StubService { Call(Request):Response (id: 1); } +rpc_service ChannelOwnership { Call(Request):Response (id: 1); } +]=]) + +set(shadow_names request response controller done method std fail BrpcFlatbuffersFail + STUB_OWNS_CHANNEL STUB_DOESNT_OWN_CHANNEL ChannelOwnership) +set(shadow_schema "namespace codegen.names;\ntable Request { text:string; }\ntable Response { text:string; }\nrpc_service Names {\n") +set(method_id 0) +foreach(name IN LISTS shadow_names) + math(EXPR method_id "${method_id} + 7") + string(APPEND shadow_schema " ${name}(Request):Response (id: ${method_id});\n") +endforeach() +string(APPEND shadow_schema "}\nrpc_service ChannelOwnership { Ping(Request):Response (id: 3); }\n") +expect_compiles(shadowed_names "${shadow_schema}") + +# Supply RUNTIME_LIBRARIES when invoking this script directly to exercise the +# new names against an existing FlatBuffers-enabled bRPC library as well. +if(RUNTIME_LIBRARIES) + set(directory "${WORK}/shadowed_names") + set(runtime_source [=[ +#include "echo.brpc.fb.h" +#include + +using ::brpc::flatbuffers::Message; +using ::google::protobuf::Closure; +using ::google::protobuf::RpcController; + +class Completion : public Closure { +public: + void Run() override { ++runs; } + int runs = 0; +}; + +class Implementation : public ::codegen::names::Names { +public: + int selected = 0; +]=]) + set(method_id 0) + foreach(name IN LISTS shadow_names) + math(EXPR method_id "${method_id} + 7") + string(APPEND runtime_source + " void ${name}(RpcController*, const Message*, Message*, Closure* completion) override { selected = ${method_id}; completion->Run(); }\n") + endforeach() + string(APPEND runtime_source [=[ +}; + +class LocalChannel : public ::brpc::flatbuffers::RpcChannel { +public: + explicit LocalChannel(::brpc::flatbuffers::Service* service) : service_(service) {} + void FBCallMethod(const ::brpc::flatbuffers::MethodDescriptor* method, + RpcController* controller, const Message* request, + Message* response, Closure* done) override { + service_->FBCallMethod(method, controller, request, response, done); + } +private: + ::brpc::flatbuffers::Service* service_; +}; + +int main() { + ::brpc::flatbuffers::MessageBuilder builder; + builder.Finish(::codegen::names::CreateRequest(builder)); + Message request = builder.ReleaseMessage(); + Message response; + Implementation implementation; + LocalChannel channel(&implementation); + ::codegen::names::Names::Stub stub(&channel); + ::codegen::names::Names defaults; + LocalChannel default_channel(&defaults); + ::codegen::names::Names::Stub default_stub(&default_channel); + ::codegen::names::Names::Stub disconnected(nullptr); + Completion completion; + int expected_runs = 0; +]=]) + set(method_id 0) + foreach(name IN LISTS shadow_names) + math(EXPR method_id "${method_id} + 7") + string(APPEND runtime_source + " stub.${name}(nullptr, &request, &response, &completion);\n" + " if (implementation.selected != ${method_id} || completion.runs != ++expected_runs) return 1;\n" + " default_stub.${name}(nullptr, &request, &response, &completion);\n" + " if (completion.runs != ++expected_runs) return 2;\n" + " disconnected.${name}(nullptr, &request, &response, &completion);\n" + " if (completion.runs != ++expected_runs) return 3;\n") + endforeach() + string(APPEND runtime_source [=[ + ::codegen::names::ChannelOwnership service; + LocalChannel service_channel(&service); + ::codegen::names::ChannelOwnership::Stub named_stub(&service_channel, + ::brpc::flatbuffers::Service::STUB_DOESNT_OWN_CHANNEL); + named_stub.Ping(nullptr, &request, &response, &completion); + if (completion.runs != ++expected_runs) return 4; + ::codegen::names::Names::Stub owned_stub(new LocalChannel(&implementation), + ::brpc::flatbuffers::Service::STUB_OWNS_CHANNEL); + owned_stub.request(nullptr, &request, &response, &completion); + if (implementation.selected != 7 || completion.runs != ++expected_runs) return 5; + std::cout << "Shadowed-name dispatch, failures, and constructors passed\n"; + return 0; +} +]=]) + file(WRITE "${directory}/runtime.cpp" "${runtime_source}") + set(includes "-I${directory}") + foreach(include_dir IN LISTS INCLUDE_DIRS) + list(APPEND includes "-I${include_dir}") + endforeach() + execute_process(COMMAND "${CXX}" ${compiler_flags} "-std=c++${CXX_STANDARD}" ${includes} + "${directory}/runtime.cpp" "${directory}/echo.o" ${RUNTIME_LIBRARIES} + -o "${directory}/runtime" + RESULT_VARIABLE result ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "Shadowed-name runtime link failed: ${error}") + endif() + execute_process(COMMAND "${directory}/runtime" + RESULT_VARIABLE result OUTPUT_VARIABLE output ERROR_VARIABLE error) + if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "Shadowed-name runtime failed: ${output}${error}") + endif() + message(STATUS "${output}") +endif() + +string(REPLACE + " Repeat(Request):Response (id: 41);\n Inspect(Request):Response (id: 7);" + " Inspect(Request):Response (id: 7);\n Repeat(Request):Response (id: 41);" + reordered "${schema_text}") +expect_compiles(reordered "${reordered}") +file(READ "${WORK}/reordered/echo.brpc.fb.cpp" reordered_source) +if(NOT reordered_source MATCHES "case 41:" OR NOT reordered_source MATCHES "case 7:") + message(FATAL_ERROR "Reordering changed wire IDs") +endif() +string(FIND "${reordered_source}" "void Echo_Stub::Repeat(" repeat_begin) +if(repeat_begin LESS 0) + message(FATAL_ERROR "Reordered stub lacks Repeat") +endif() +string(SUBSTRING "${reordered_source}" ${repeat_begin} -1 repeat_body) +string(FIND "${repeat_body}" "\n}\n" repeat_end) +string(SUBSTRING "${repeat_body}" 0 ${repeat_end} repeat_body) +if(NOT repeat_body MATCHES "method\\(1\\)") + message(FATAL_ERROR "Reordered stub does not use its dense method position") +endif() + +string(REPLACE "namespace codegen.example;" "" global_schema "${schema_text}") +expect_compiles(global_namespace "${global_schema}") +string(REPLACE "(id: 7)" "(id: 0)" zero_schema "${schema_text}") +expect_compiles(zero_id "${zero_schema}") +string(REPLACE "(id: 7)" "(id: 2147483647)" maximum_schema "${schema_text}") +expect_compiles(maximum_id "${maximum_schema}") + +# Include schemas are read through the official Parser and not re-emitted. +file(MAKE_DIRECTORY "${WORK}/included") +file(WRITE "${WORK}/included/types.fbs" + "namespace shared; table Input { note:string; } table Output { note:string; }\n" + "rpc_service Imported { Ping(Input):Output (id: 3); }\n") +execute_process(COMMAND "${FLATC}" --cpp -o "${WORK}/included" + "${WORK}/included/types.fbs" RESULT_VARIABLE result) +if(NOT "${result}" STREQUAL "0") + message(FATAL_ERROR "Cannot generate included table definitions") +endif() +expect_compiles(included + "include \"types.fbs\"; namespace api; rpc_service Local { Call(shared.Input):shared.Output (id: 17); }") +file(READ "${WORK}/included/echo.brpc.fb.h" included_header) +if(included_header MATCHES "class Imported") + message(FATAL_ERROR "Included service was emitted twice") +endif() + +set(imported_collision_dir "${WORK}/imported_service_collision") +file(REMOVE_RECURSE "${imported_collision_dir}") +file(MAKE_DIRECTORY "${imported_collision_dir}") +file(WRITE "${imported_collision_dir}/types.fbs" + "namespace shared; table Input {} table Output {}\n" + "rpc_service Imported { Ping(Input):Output (id: 3); }\n") +file(WRITE "${imported_collision_dir}/echo.fbs" + "include \"types.fbs\"; namespace shared;\n" + "rpc_service Imported_Stub { Call(Input):Output (id: 7); }\n") +execute_process(COMMAND "${GENERATOR}" -o "${imported_collision_dir}" + "${imported_collision_dir}/echo.fbs" + RESULT_VARIABLE result ERROR_VARIABLE error) +if("${result}" STREQUAL "0" OR + NOT error MATCHES "generated service class name collides: ::shared::Imported_Stub") + message(FATAL_ERROR + "imported_service_collision: schema was accepted or lacked diagnostic: ${error}") +endif() +if(EXISTS "${imported_collision_dir}/echo.brpc.fb.h" OR + EXISTS "${imported_collision_dir}/echo.brpc.fb.cpp") + message(FATAL_ERROR "imported_service_collision: rejected schema emitted files") +endif() +message(STATUS "Rejected imported service collision: ${error}") diff --git a/test/flatbuffers_codegen/echo.fbs b/test/flatbuffers_codegen/echo.fbs new file mode 100644 index 0000000000..8a59fa8cb9 --- /dev/null +++ b/test/flatbuffers_codegen/echo.fbs @@ -0,0 +1,50 @@ +// 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. + +// Exercise helper-name collisions in the normal compile and runtime targets. +table BrpcFlatbuffersFail { + marker:int; +} + +namespace codegen.example; + +table Request { + text:string; +} + +table Response { + text:string; +} + +rpc_service Echo { + Repeat(Request):Response (id: 41); + Inspect(Request):Response (id: 7); +} + +rpc_service Other { + Repeat(Request):Response (id: 41); +} + +table module { + value:int; +} + +rpc_service module { + RoundTrip(module):module (id: 17); +} + +root_type Request; diff --git a/test/flatbuffers_codegen/runtime.cpp b/test/flatbuffers_codegen/runtime.cpp new file mode 100644 index 0000000000..133565c536 --- /dev/null +++ b/test/flatbuffers_codegen/runtime.cpp @@ -0,0 +1,339 @@ +// 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 "echo.brpc.fb.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +#define CODEGEN_CHECK(condition) \ + do { \ + if (!(condition)) { \ + throw std::runtime_error("check failed: " #condition); \ + } \ + } while (0) + +using brpc::flatbuffers::Message; +using brpc::flatbuffers::MessageBuilder; +using brpc::flatbuffers::MethodDescriptor; +using brpc::flatbuffers::ServiceDescriptor; +using codegen::example::Echo; +using codegen::example::Other; +using codegen::example::Request; +using codegen::example::Response; +using google::protobuf::Closure; +using google::protobuf::RpcController; + +class Controller : public RpcController { +public: + void Reset() override { failures = 0; error.clear(); } + bool Failed() const override { return failures != 0; } + std::string ErrorText() const override { return error; } + void StartCancel() override {} + void SetFailed(const std::string& reason) override { ++failures; error = reason; } + bool IsCanceled() const override { return false; } + void NotifyOnCancel(Closure*) override {} + int failures = 0; + std::string error; +}; + +class Done : public Closure { +public: + void Run() override { ++runs; } + int runs = 0; +}; + +class DeletingDone : public Closure { +public: + explicit DeletingDone(int* runs) : runs_(runs) {} + void Run() override { ++*runs_; delete this; } +private: + int* runs_; +}; + +Message MakeRequest(const char* text) { + MessageBuilder builder; + const auto value = text ? builder.CreateString(text) : + ::flatbuffers::Offset<::flatbuffers::String>(); + builder.Finish(codegen::example::CreateRequest(builder, value)); + return builder.ReleaseMessage(); +} + +class Implementation : public Echo { +public: + void Repeat(RpcController*, const Message* request, + Message* response, Closure* done) override { + ++calls; + selected = 41; + Reply(request, response, done); + } + void Inspect(RpcController*, const Message* request, + Message* response, Closure* done) override { + ++calls; + selected = 7; + Reply(request, response, done); + } + void Reply(const Message* request, Message* response, Closure* done) { + const auto* input = request->GetRoot(); + MessageBuilder builder; + const auto value = input->text() ? builder.CreateString(input->text()->str()) : + ::flatbuffers::Offset<::flatbuffers::String>(); + builder.Finish(codegen::example::CreateResponse(builder, value)); + *response = builder.ReleaseMessage(); + if (done) { + done->Run(); + } + } + int calls = 0; + int selected = -1; +}; + +class LocalChannel : public brpc::flatbuffers::RpcChannel { +public: + explicit LocalChannel(Echo* service, int* destroyed = nullptr) + : service_(service), destroyed_(destroyed) {} + ~LocalChannel() override { if (destroyed_) { ++*destroyed_; } } + void FBCallMethod(const MethodDescriptor* method, RpcController* controller, + const Message* request, Message* response, Closure* done) override { + last_method = method; + service_->FBCallMethod(method, controller, request, response, done); + } + const MethodDescriptor* last_method = nullptr; +private: + Echo* service_; + int* destroyed_; +}; + +void TestConcurrentDescriptors() { + const int thread_count = 24; + std::atomic start(false); + std::atomic invalid(0); + std::vector descriptors(thread_count); + std::vector threads; + for (int i = 0; i < thread_count; ++i) { + threads.emplace_back([&, i] { + while (!start.load()) { + std::this_thread::yield(); + } + for (int n = 0; n < 1000; ++n) { + descriptors[i] = Echo::descriptor(); + if (descriptors[i]->method_count() != 2 || + descriptors[i]->method(0)->index() != 41 || + descriptors[i]->method(1)->index() != 7) { + ++invalid; + } + } + }); + } + start.store(true); + for (auto& thread : threads) { + thread.join(); + } + CODEGEN_CHECK(invalid.load() == 0); + for (const auto* descriptor : descriptors) { + CODEGEN_CHECK(descriptor == descriptors.front()); + } + const auto* descriptor = descriptors.front(); + CODEGEN_CHECK(descriptor->full_name() == "codegen.example.Echo"); + CODEGEN_CHECK(descriptor->FindMethodByIndex(41) == descriptor->method(0)); + CODEGEN_CHECK(descriptor->FindMethodByIndex(7) == descriptor->method(1)); + CODEGEN_CHECK(descriptor->FindMethodByIndex(0) == nullptr); + CODEGEN_CHECK(descriptor->method(0)->service() == descriptor); +} + +void ExpectFailure(const std::function& call) { + Controller controller; + Done done; + call(&controller, &done); + CODEGEN_CHECK(controller.Failed()); + CODEGEN_CHECK(controller.failures == 1); + CODEGEN_CHECK(!controller.ErrorText().empty()); + CODEGEN_CHECK(done.runs == 1); +} + +void TestDispatchAndErrors() { + Implementation service; + LocalChannel channel(&service); + Echo::Stub stub(&channel); + const auto* descriptor = Echo::descriptor(); + CODEGEN_CHECK(stub.GetDescriptor() == descriptor); + CODEGEN_CHECK(stub.channel() == &channel); + for (const char* text : {static_cast(nullptr), "hello"}) { + Message request = MakeRequest(text); + for (int position = 0; position < 2; ++position) { + Controller controller; + Done done; + Message response; + if (position == 0) { + stub.Repeat(&controller, &request, &response, &done); + } else { + stub.Inspect(&controller, &request, &response, &done); + } + CODEGEN_CHECK(!controller.Failed()); + CODEGEN_CHECK(done.runs == 1); + CODEGEN_CHECK(channel.last_method == descriptor->method(position)); + CODEGEN_CHECK(service.selected == (position == 0 ? 41 : 7)); + CODEGEN_CHECK(response.Verify()); + const auto* output = response.GetRoot(); + CODEGEN_CHECK(text ? output->text() && output->text()->str() == text : + output->text() == nullptr); + } + } + const int calls = service.calls; + Message request = MakeRequest(nullptr); + Message response; + Message invalid; + const MethodDescriptor unknown("Unknown", descriptor, 99); + const MethodDescriptor forged("Forged", descriptor, 41); + for (const auto* method : {static_cast(nullptr), + Other::descriptor()->method(0), &unknown, &forged}) { + ExpectFailure([&](Controller* controller, Done* done) { + service.FBCallMethod(method, controller, &request, &response, done); + }); + } + for (int position = 0; position < 2; ++position) { + ExpectFailure([&](Controller* controller, Done* done) { + service.FBCallMethod(descriptor->method(position), controller, + &invalid, &response, done); + }); + } + ExpectFailure([&](Controller* controller, Done* done) { + service.FBCallMethod(descriptor->method(0), controller, nullptr, &response, done); + }); + ExpectFailure([&](Controller* controller, Done* done) { + service.FBCallMethod(descriptor->method(0), controller, &request, nullptr, done); + }); + CODEGEN_CHECK(service.calls == calls); + + Echo unimplemented; + for (int position = 0; position < 2; ++position) { + ExpectFailure([&](Controller* controller, Done* done) { + unimplemented.FBCallMethod(descriptor->method(position), controller, + &request, &response, done); + }); + } + ExpectFailure([&](Controller* controller, Done* done) { + unimplemented.Repeat(controller, &request, &response, done); + }); + Echo::Stub disconnected(nullptr); + ExpectFailure([&](Controller* controller, Done* done) { + disconnected.Repeat(controller, &request, &response, done); + }); + ExpectFailure([&](Controller* controller, Done* done) { + disconnected.Inspect(controller, &request, &response, done); + }); + Controller controller; + service.FBCallMethod(nullptr, &controller, &request, &response, nullptr); + CODEGEN_CHECK(controller.failures == 1); + Done done; + service.FBCallMethod(nullptr, nullptr, &request, &response, &done); + CODEGEN_CHECK(done.runs == 1); + int deleted_runs = 0; + service.FBCallMethod(nullptr, nullptr, &request, &response, + new DeletingDone(&deleted_runs)); + CODEGEN_CHECK(deleted_runs == 1); +} + +void TestOwnershipAndAsyncCompletion() { + Implementation service; + int destroyed = 0; + { + Echo::Stub owner(new LocalChannel(&service, &destroyed), + brpc::flatbuffers::Service::STUB_OWNS_CHANNEL); + } + CODEGEN_CHECK(destroyed == 1); + { + LocalChannel borrowed(&service, &destroyed); + { + Echo::Stub stub(&borrowed); + } + CODEGEN_CHECK(destroyed == 1); + } + CODEGEN_CHECK(destroyed == 2); + + class Deferred : public Echo { + public: + void Repeat(RpcController*, const Message*, Message*, Closure* done) override { + pending = done; + } + Closure* pending = nullptr; + } deferred; + Message request = MakeRequest(nullptr); + Message response; + Controller controller; + Done done; + deferred.FBCallMethod(Echo::descriptor()->method(0), &controller, + &request, &response, &done); + CODEGEN_CHECK(!controller.Failed()); + CODEGEN_CHECK(done.runs == 0); + CODEGEN_CHECK(deferred.pending == &done); + deferred.pending->Run(); + CODEGEN_CHECK(done.runs == 1); +} + +void TestEscapedTableTypes() { + class EscapedImplementation : public codegen::example::module { + public: + void RoundTrip(RpcController*, const Message* request, + Message* response, Closure* done) override { + value = request->GetRoot()->value(); + response->CopyFrom(*request); + if (done) { + done->Run(); + } + } + int value = 0; + } service; + MessageBuilder builder; + builder.Finish(codegen::example::Createmodule_(builder, 73)); + Message request = builder.ReleaseMessage(); + Message response; + Controller controller; + Done done; + const auto* descriptor = codegen::example::module::descriptor(); + CODEGEN_CHECK(descriptor->full_name() == "codegen.example.module"); + service.FBCallMethod(descriptor->FindMethodByIndex(17), &controller, + &request, &response, &done); + CODEGEN_CHECK(!controller.Failed()); + CODEGEN_CHECK(service.value == 73); + CODEGEN_CHECK(done.runs == 1); + CODEGEN_CHECK(response.Verify()); + CODEGEN_CHECK(response.GetRoot()->value() == 73); +} + +} // namespace + +int main() { + try { + TestConcurrentDescriptors(); + TestDispatchAndErrors(); + TestOwnershipAndAsyncCompletion(); + TestEscapedTableTypes(); + std::cout << "Descriptor, sparse dispatch, verifier, callback and ownership checks passed\n"; + return 0; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/test/flatbuffers_message.fbs b/test/flatbuffers_message.fbs new file mode 100644 index 0000000000..0288499125 --- /dev/null +++ b/test/flatbuffers_message.fbs @@ -0,0 +1,26 @@ +// 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. + +namespace brpc_fbtest; + +table Payload { + value:long; + message:string; + values:[int]; +} + +root_type Payload; diff --git a/tools/flatbuffers/CMakeLists.txt b/tools/flatbuffers/CMakeLists.txt new file mode 100644 index 0000000000..df4f578526 --- /dev/null +++ b/tools/flatbuffers/CMakeLists.txt @@ -0,0 +1,46 @@ +# 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. + +cmake_minimum_required(VERSION 3.10) +project(brpc_flatc LANGUAGES CXX) + +find_path(FLATBUFFERS_INCLUDE_DIR flatbuffers/idl.h) +find_library(FLATBUFFERS_LIBRARY NAMES flatbuffers) +if(NOT FLATBUFFERS_INCLUDE_DIR OR NOT FLATBUFFERS_LIBRARY) + message(FATAL_ERROR "Install official FlatBuffers development headers and libflatbuffers") +endif() + +function(configure_brpc_flatc target) + target_compile_features(${target} PRIVATE cxx_std_11) + target_include_directories(${target} SYSTEM PRIVATE "${FLATBUFFERS_INCLUDE_DIR}") + target_link_libraries(${target} PRIVATE "${FLATBUFFERS_LIBRARY}") + if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_compile_options(${target} PRIVATE -Wall -Wextra -Wpedantic) + endif() +endfunction() + +add_executable(brpc_flatc brpc_flatc.cpp) +configure_brpc_flatc(brpc_flatc) + +include(CTest) +if(BUILD_TESTING) + add_executable(brpc_flatc_test brpc_flatc.cpp) + configure_brpc_flatc(brpc_flatc_test) + target_compile_definitions(brpc_flatc_test PRIVATE BRPC_FLATC_TESTING=1) + add_subdirectory("${CMAKE_CURRENT_SOURCE_DIR}/../../test/flatbuffers_codegen" + "${CMAKE_CURRENT_BINARY_DIR}/tests") +endif() diff --git a/tools/flatbuffers/README.md b/tools/flatbuffers/README.md new file mode 100644 index 0000000000..93bc4f988d --- /dev/null +++ b/tools/flatbuffers/README.md @@ -0,0 +1,243 @@ +# Standalone FlatBuffers service generator + +`brpc_flatc` generates the bRPC **service abstraction**, not a network protocol or +an adapter for `brpc::Channel`. It uses the installed, official +`flatbuffers/idl.h` Parser and `libflatbuffers`; it does not download or modify a +FlatBuffers fork. The official `flatc --cpp` remains responsible for table types. + +## Build and generate + +The tool needs a C++11 compiler, CMake, and official FlatBuffers **25.2.10** +headers/library. Use the matching `flatc` with default `--cpp` generation; +options that change C++ naming, such as `--scoped-enums`, are not modeled by +this generator. To build only the generator: + +```sh +cmake -S tools/flatbuffers -B /tmp/brpc-codegen-build -DBUILD_TESTING=OFF +cmake --build /tmp/brpc-codegen-build -j2 +mkdir -p /tmp/brpc-generated +flatc --cpp -o /tmp/brpc-generated test/flatbuffers_codegen/echo.fbs +/tmp/brpc-codegen-build/brpc_flatc -o /tmp/brpc-generated test/flatbuffers_codegen/echo.fbs +``` + +The output directory must already exist. The command line is: + +```text +brpc_flatc [-I include_dir]... [-o existing_output_dir] schema.fbs +``` + +For `echo.fbs`, the two tools produce: + +- Official `flatc`: `echo_generated.h`. +- `brpc_flatc`: `echo.brpc.fb.h` and `echo.brpc.fb.cpp`, both Apache-licensed. + +Compile the generated `.cpp` into the application with C++14 or newer and bRPC +configured with FlatBuffers support. Keep the generated header, official table header, and bRPC +headers on the include path. The generated service files do not pin a +FlatBuffers version; use an official `flatc` matching your installed table +headers, as required by official generated code. + +Schema includes are resolved relative to the input file and through repeated +`-I` options (`-Idir` also works). Run official `flatc --cpp` for each included +schema to obtain its table header. Run `brpc_flatc` separately for each schema +that declares services; imported services are not emitted again. Pass the same +include directories to both tools. This tool processes one `.fbs` per invocation +and follows the official default `*_generated.h` naming convention. + +## Explicit, stable IDs + +```fbs +namespace example.api; + +table Request { text:string; } +table Response { text:string; } + +rpc_service Echo { + Repeat(Request):Response (id: 41); + Inspect(Request):Response (id: 7); +} +``` + +Every RPC method must have an explicit integer `(id: N)` in `[0, 2147483647]`. +IDs must be unique **within a service**. Missing, negative, duplicate, quoted, +or out-of-range IDs fail generation with a diagnostic and nonzero exit status. +Reordering declarations never changes their wire IDs. This does not require +explicit IDs on FlatBuffers table fields. + +The generated names are `example::api::Echo`, +`example::api::Echo_Stub`, and the alias `example::api::Echo::Stub`. +They are **not** the old `EchoStub` spelling. Methods take protobuf +`RpcController`/`Closure` and `brpc::flatbuffers::Message` parameters. Derive from +`Echo` and override the schema methods; supply an implementation of +`brpc::flatbuffers::RpcChannel` to a stub. Channel ownership defaults to borrowed; +`Service::STUB_OWNS_CHANNEL` transfers ownership to the stub. + +This intentionally bounded generator supports unary table RPCs, multiple +services/methods, namespaces, included request/response tables, and absent +optional strings. Streaming and C++ keyword names in service/type/namespace +positions are rejected rather than silently misgenerated. Method names that +collide with generated service APIs are also rejected. Services cannot be named +`Stub`, `descriptor`, `GetDescriptor`, or `FBCallMethod`, because these names +are also members of the generated service class. Generated service and stub +class names must be distinct from other services, schema types, table builder +classes, unscoped enum values, union traits, table creation functions, and +namespace prefixes, including those in imported schemas. C++ reserved +identifiers are rejected in the scopes where they are reserved. A service name +ending in `_` is also rejected because its `_Stub` suffix would produce `__`. +Global service names also cannot shadow +namespaces introduced by the generated header (`brpc`, `butil`, `flatbuffers`, +`google`, or `std`). Schema file basenames may contain ASCII letters, digits, +underscores, dots, and hyphens. The header and source are staged under an +output-directory lock before publication. A failure restores the previous pair +when possible; if rollback also fails, the generator +attempts to remove both final outputs and reports that manual cleanup is needed +when the filesystem refuses. The standalone generator currently requires POSIX +file-locking semantics, including on Linux and macOS. + +## Descriptor and completion contracts + +- `Echo::descriptor()` owns a `ServiceDescriptor` in a function-local static + RAII holder. C++11 initialization is thread-safe; the holder and its owned + method descriptors are destroyed normally at process exit. There is no + leaked singleton allocation or unsynchronized lazy initialization. +- The descriptor table contains the namespace prefix, service name, ordered + method names, and explicit IDs. Stubs use `method(position)`; service dispatch + switches on the stable `method->index()`. +- Service dispatch checks descriptor ownership and identity, non-null request + and response, and `request->Verify()` before calling user code. +- Invalid calls, null stub channels, and default unimplemented methods call + `controller->SetFailed()` when a controller exists, then run non-null `done` + exactly once. Null controllers and callbacks are tolerated on failure. +- Successful dispatch transfers completion responsibility to the user method. + The generated dispatcher does not run `done` again, so asynchronous service + implementations remain possible. User methods and channel implementations + must themselves honor the exactly-once completion contract. + +## Using another FlatBuffers version + +The runtime uses FlatBuffers builder internals that may change between upstream +releases, so the source accepts only 25.2.10. To use another version, build it +into a private prefix, update the version checks, and run the full +verification against that prefix. The steps below use 23.5.26 and were tested +on macOS ARM64 (AppleClang 15) and Ubuntu 22.04 (GCC 11.4, Clang 14). They +cover CMake only; the 23.5.26 Bazel files load `@aspect_rules_js` and `@npm`, +which bRPC does not configure. + +Build FlatBuffers. `CMAKE_INSTALL_LIBDIR=lib` keeps the library where +`verify.sh` looks for it: + +```sh +git clone --depth 1 --branch v23.5.26 https://github.com/google/flatbuffers.git /tmp/flatbuffers-23.5.26 +cmake -S /tmp/flatbuffers-23.5.26 -B /tmp/flatbuffers-23.5.26/build \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/flatbuffers-23.5.26-install \ + -DCMAKE_INSTALL_LIBDIR=lib -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ + -DFLATBUFFERS_BUILD_TESTS=OFF +cmake --build /tmp/flatbuffers-23.5.26/build --target install -j4 +``` + +In the bRPC source tree, update the version checks in +`src/brpc/flatbuffers/message.h`, `tools/flatbuffers/brpc_flatc.cpp` and +`test/flatbuffers_codegen/CMakeLists.txt`: + +```sh +perl -pi -e 's/25\.2\.10/23.5.26/g; + s/(FLATBUFFERS_VERSION_MAJOR != )25\b/${1}23/; + s/(FLATBUFFERS_VERSION_MINOR != )2\b/${1}5/; + s/(FLATBUFFERS_VERSION_REVISION != )10\b/${1}26/' \ + src/brpc/flatbuffers/message.h tools/flatbuffers/brpc_flatc.cpp \ + test/flatbuffers_codegen/CMakeLists.txt +git diff --stat +``` + +The diff should touch 13 lines in these three files. If it does not, the source +has changed since this example was written; review the remaining version +checks by hand. + +Then run the verification with the new prefix: + +```sh +FLATBUFFERS_PREFIX=/tmp/flatbuffers-23.5.26-install \ + bash tools/flatbuffers/verify.sh /tmp/brpc-fb23-verify +``` + +Without `FLATBUFFERS_PREFIX`, the script downloads 25.2.10, which the updated +checks reject. Test your own schemas as well. Other versions may need more than +new version numbers, for example changes for Builder internals or the +generator's naming rules. + +## Pre-submit verification + +Local runs and the Linux/macOS codegen CI jobs use the same entry point: + +```sh +CXX=clang++ bash tools/flatbuffers/verify.sh /tmp/brpc-codegen-verify +``` + +Install CMake, a C++ compiler, protobuf, gflags, LevelDB, OpenSSL, zlib, and +Gperftools first. The root unit-test configuration requires Gperftools even +when only the FlatBuffers target is built. +Use `CMAKE_PREFIX_PATH` for non-system dependencies (on macOS, select the desired +Homebrew protobuf prefix). The script fetches SHA-256-pinned official FlatBuffers +25.2.10 and GTest 1.12.0 archives (URLs and hashes are in `verify.sh`). It enables +`BUILD_UNIT_TESTS`, selects the matching `flatc`, and supplies GTest source through +`BRPC_SYSTEM_GTEST_SOURCE_DIR` with `DOWNLOAD_GTEST=OFF`; no system GTest package +is needed. It builds and runs `brpc_flatbuffers_unittest`, rejecting zero-case, +disabled, or skipped runs, then builds the generator and runs both codegen +acceptance and runtime tests. Missing tests or failures stop the script. +`FLATBUFFERS_PREFIX` may reuse an existing matching installation. +`CODEGEN_RUNTIME=OFF` only builds the generator and runs compile-only acceptance; +it skips GTest, the core bRPC build/unit tests, and runtime tests, even when reusing +a build directory. This is not the full CI check. + +For generator changes, add both rejected and accepted schemas, compile the +actual output of both generators, and demonstrate that a new regression fails +before the fix. Check generated symbol families and global/nested/imported +scopes rather than only the reported spelling. Record compiler, dependency, +SDK, flags, and source hashes; the script prints tool versions, build directories, +archive hashes, verbose CTest results, and the core GTest XML report path. +The same script does not make different runner/compiler versions equivalent. +Ordinary FlatBuffers-OFF jobs are not codegen evidence. + +## Independent acceptance tests + +No root build file needs modification. The default standalone build enables +CTest and additionally needs matching official `flatc`, protobuf, and gflags +development headers: + +```sh +cmake -S tools/flatbuffers -B /tmp/brpc-codegen-build -DBUILD_TESTING=ON +cmake --build /tmp/brpc-codegen-build -j2 +ctest --test-dir /tmp/brpc-codegen-build --output-on-failure +``` + +The build compiles the generated service against the real bRPC headers. CTest +rejects missing/duplicate/negative/overflow/string IDs, streaming RPCs, C++20 +keywords, and service/stub names that collide with generated C++ types. +Test-only fault injection covers staging, backup, publication, and rollback +failures; concurrent publishers must also leave a consistent pair. The suite +then compiles the generated service and self-contained header for C++14-or-newer +compatibility, sparse IDs, reordered methods, global/nested +namespaces, zero/max IDs, and included table types. Test configuration headers +and generated artifacts stay in the standalone build directory, not in the +source tree. The name-shadowing cases are compile-only by default; pass +`RUNTIME_LIBRARIES` when invoking `acceptance.cmake` directly to exercise their +runtime checks as well. + +To enable additional behavior tests against an **already built** +FlatBuffers-enabled bRPC library: + +```sh +cmake -S tools/flatbuffers -B /tmp/brpc-codegen-build \ + -DBUILD_TESTING=ON \ + -DBRPC_CODEGEN_BRPC_LIBRARY=/path/to/libbrpc.a +cmake --build /tmp/brpc-codegen-build -j2 +ctest --test-dir /tmp/brpc-codegen-build --output-on-failure +``` + +Use `BRPC_CODEGEN_EXTRA_LIBRARIES` for additional dependencies of a custom bRPC +build. The runtime test covers concurrent first descriptor access, sparse +stub/dispatch mapping, absent and present strings, invalid/foreign/forged +methods, failed request verification, null arguments, unimplemented methods, +exactly-once error callbacks (including self-deleting closures), deferred +completion, and channel ownership. It uses a local abstract channel; no network +RPC compatibility is claimed. diff --git a/tools/flatbuffers/brpc_flatc.cpp b/tools/flatbuffers/brpc_flatc.cpp new file mode 100644 index 0000000000..9c4665387f --- /dev/null +++ b/tools/flatbuffers/brpc_flatc.cpp @@ -0,0 +1,894 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if FLATBUFFERS_VERSION_MAJOR != 25 || FLATBUFFERS_VERSION_MINOR != 2 || \ + FLATBUFFERS_VERSION_REVISION != 10 +#error "brpc_flatc requires FlatBuffers 25.2.10 headers and library" +#endif + +namespace { + +const char kLicense[] = R"(// 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. + +// Generated by brpc_flatc. Do not edit. + +)"; + +const char kArguments[] = + "::google::protobuf::RpcController* controller,\n" + " const ::brpc::flatbuffers::Message* request,\n" + " ::brpc::flatbuffers::Message* response,\n" + " ::google::protobuf::Closure* done"; + +struct Service { + const flatbuffers::ServiceDef* definition; + std::vector ids; +}; + +std::string JoinNamespace(const flatbuffers::Definition& definition, + const std::string& separator) { + std::string result; + if (definition.defined_namespace) { + for (const auto& component : definition.defined_namespace->components) { + if (!result.empty()) { + result += separator; + } + result += component; + } + } + return result; +} + +std::string Qualified(const flatbuffers::Definition& definition) { + const std::string ns = JoinNamespace(definition, "::"); + return "::" + (ns.empty() ? "" : ns + "::") + definition.name; +} + +std::string GuardComponent(const std::string& value) { + static const char digits[] = "0123456789ABCDEF"; + std::string result; + for (unsigned char c : value) { + result += digits[c >> 4]; + result += digits[c & 15]; + } + return result; +} + +bool IsCppIdentifier(const std::string& name) { + static const std::set keywords = { + "alignas", "alignof", "and", "and_eq", "asm", "auto", "bitand", + "bitor", "bool", "break", "case", "catch", "char", "char8_t", + "char16_t", "char32_t", "class", "compl", "concept", "const", + "const_cast", "consteval", "constexpr", "constinit", "continue", + "co_await", "co_return", "co_yield", "decltype", "default", + "delete", "do", + "double", "dynamic_cast", "else", "enum", "explicit", "export", + "extern", "false", "float", "for", "friend", "goto", "if", + "inline", "int", "long", "mutable", "namespace", "new", + "noexcept", "not", "not_eq", "nullptr", "operator", "or", "or_eq", + "private", "protected", "public", "register", "reinterpret_cast", + "requires", "return", "short", "signed", "sizeof", "static", + "static_assert", "static_cast", "struct", "switch", "template", + "this", "thread_local", "throw", "true", "try", "typedef", + "typeid", "typename", "union", "unsigned", "using", "virtual", + "void", "volatile", "wchar_t", "while", "xor", "xor_eq" + }; + return !name.empty() && keywords.count(name) == 0; +} + +std::string FlatcName(const std::string& name) { + // flatc 25.2.10 escapes a different set than the C++20 rejection list. + static const std::set extra = { + "atomic_cancel", "atomic_commit", "atomic_noexcept", "import", + "module", "synchronized" + }; + const bool newer_keyword = name == "char8_t" || name == "consteval" || + name == "constinit"; + return ((!IsCppIdentifier(name) && !newer_keyword) || extra.count(name)) + ? name + "_" : name; +} + +std::string FlatcNamespace(const flatbuffers::Definition& definition) { + const std::string ns = JoinNamespace(definition, "::"); + return ns.empty() ? "::" : "::" + ns + "::"; +} + +std::string FlatcQualifiedType(const flatbuffers::Definition& definition) { + return FlatcNamespace(definition) + FlatcName(definition.name); +} + +bool IsReservedIdentifier(const std::string& name, bool global) { + return name.find("__") != std::string::npos || + (!name.empty() && name[0] == '_' && + (global || (name.size() > 1 && name[1] >= 'A' && name[1] <= 'Z'))); +} + +bool ValidateName(const flatbuffers::Definition& definition, + std::string* error, bool member = false) { + if (!IsCppIdentifier(definition.name)) { + *error = "C++ keyword is not supported: " + definition.name; + return false; + } + const bool global = !member && (!definition.defined_namespace || + definition.defined_namespace->components.empty()); + if (IsReservedIdentifier(definition.name, global)) { + *error = "reserved C++ identifier: " + definition.name; + return false; + } + if (definition.defined_namespace) { + bool first = true; + for (const auto& component : definition.defined_namespace->components) { + if (IsReservedIdentifier(component, first)) { + *error = "reserved C++ identifier in namespace: " + component; + return false; + } + first = false; + if (!IsCppIdentifier(component)) { + *error = "C++ keyword namespace is not supported: " + component; + return false; + } + } + } + return true; +} + +bool ParseId(const flatbuffers::Value& value, int32_t* id) { + if (value.constant.empty() || + !flatbuffers::IsInteger(value.type.base_type)) { + return false; + } + int64_t number = 0; + for (char c : value.constant) { + if (c < '0' || c > '9') { + return false; + } + number = number * 10 + (c - '0'); + if (number > std::numeric_limits::max()) { + return false; + } + } + *id = static_cast(number); + return true; +} + +void CollectRootNames(const flatbuffers::Parser& parser, + std::set* names) { + if (!parser.root_struct_def_) { + return; + } + const auto& root = *parser.root_struct_def_; + const std::string ns = FlatcNamespace(root); + const std::string name = FlatcName(root.name); + names->insert(ns + "Get" + name); + names->insert(ns + "GetSizePrefixed" + name); + names->insert(ns + "Verify" + name + "Buffer"); + names->insert(ns + "VerifySizePrefixed" + name + "Buffer"); + names->insert(ns + "Finish" + name + "Buffer"); + names->insert(ns + "FinishSizePrefixed" + name + "Buffer"); + if (!parser.file_identifier_.empty()) { + names->insert(ns + name + "Identifier"); + names->insert(ns + name + "BufferHasIdentifier"); + names->insert(ns + "SizePrefixed" + name + "BufferHasIdentifier"); + } + if (!parser.file_extension_.empty()) { + names->insert(ns + name + "Extension"); + } +} + +bool CollectServices(const flatbuffers::Parser& parser, + const std::string& input, const char** include_paths, + std::vector* services, std::string* error) { + std::set generated_class_names = { + "::brpc", "::butil", "::flatbuffers", "::google", "::std" + }; + CollectRootNames(parser, &generated_class_names); + // The parser resets root_type after each include; flatc emits those helpers + // when each imported schema is generated separately. + for (const auto& file : parser.GetIncludedFilesRecursive(input)) { + if (file == input) { + continue; + } + std::string schema; + flatbuffers::Parser included; + if (!flatbuffers::LoadFile(file.c_str(), false, &schema) || + !included.Parse(schema.c_str(), include_paths, file.c_str())) { + *error = "cannot inspect imported root_type: " + file; + return false; + } + CollectRootNames(included, &generated_class_names); + } + // Reserve every namespace prefix, including those from imported schemas. + for (const auto* ns : parser.namespaces_) { + std::string prefix; + for (const auto& component : ns->components) { + prefix += "::" + component; + generated_class_names.insert(prefix); + } + } + for (const auto* definition : parser.services_.vec) { + if (definition->generated) { + const std::string service_name = Qualified(*definition); + generated_class_names.insert(service_name); + generated_class_names.insert(service_name + "_Stub"); + } + } + for (const auto* definition : parser.structs_.vec) { + const std::string ns = FlatcNamespace(*definition); + const std::string name = FlatcName(definition->name); + const std::string type_name = ns + name; + generated_class_names.insert(type_name); + if (!definition->fixed) { + generated_class_names.insert(type_name + "Builder"); + generated_class_names.insert(ns + "Create" + name); + for (const auto* field : definition->fields.vec) { + if (!field->deprecated && + (flatbuffers::IsString(field->value.type) || + flatbuffers::IsVector(field->value.type))) { + generated_class_names.insert( + ns + "Create" + name + "Direct"); + break; + } + } + } + } + // Match the default flatc --cpp API, including non-type declarations. + for (const auto* definition : parser.enums_.vec) { + const std::string ns = FlatcNamespace(*definition); + const std::string name = FlatcName(definition->name); + const std::string enum_name = ns + name; + generated_class_names.insert(enum_name); + for (const auto* value : definition->Vals()) { + generated_class_names.insert(enum_name + "_" + FlatcName(value->name)); + } + if (definition->attributes.Lookup("bit_flags")) { + generated_class_names.insert(enum_name + "_NONE"); + generated_class_names.insert(enum_name + "_ANY"); + } else { + generated_class_names.insert(enum_name + "_MIN"); + generated_class_names.insert(enum_name + "_MAX"); + } + generated_class_names.insert(ns + "EnumValues" + name); + generated_class_names.insert(ns + "EnumName" + name); + if (definition->Distance() / definition->size() < 5) { + generated_class_names.insert(ns + "EnumNames" + name); + } + if (definition->is_union) { + if (!definition->uses_multiple_type_instances) { + generated_class_names.insert(enum_name + "Traits"); + } + generated_class_names.insert(ns + "Verify" + name); + generated_class_names.insert( + ns + "Verify" + name + "Vector"); + } + } + for (const auto* definition : parser.services_.vec) { + // Included schemas are generated separately, just as with flatc --cpp. + if (definition->generated) { + continue; + } + if (!ValidateName(*definition, error)) { + return false; + } + if (definition->name == "Stub" || definition->name == "descriptor" || + definition->name == "GetDescriptor" || + definition->name == "FBCallMethod") { + *error = "service name collides with generated API: " + + definition->name; + return false; + } + const std::string service_name = Qualified(*definition); + const std::string stub_name = service_name + "_Stub"; + if (IsReservedIdentifier(definition->name + "_Stub", false)) { + *error = "generated stub uses reserved C++ identifier: " + stub_name; + return false; + } + if (!generated_class_names.insert(service_name).second) { + *error = "generated service class name collides: " + service_name; + return false; + } + if (!generated_class_names.insert(stub_name).second) { + *error = "generated service class name collides: " + stub_name; + return false; + } + Service service = {definition, {}}; + std::set ids; + if (definition->calls.vec.empty()) { + *error = "rpc_service must contain at least one method: " + + definition->name; + return false; + } + for (const auto* call : definition->calls.vec) { + if (!ValidateName(*call, error, true) || + !ValidateName(*call->request, error) || + !ValidateName(*call->response, error)) { + return false; + } + if (call->name == definition->name || + call->name == definition->name + "_Stub" || + call->name == "descriptor" || call->name == "GetDescriptor" || + call->name == "FBCallMethod" || call->name == "Stub" || + call->name == "channel" || call->name == "channel_" || + call->name == "owned_channel_") { + *error = "method name collides with generated API: " + + call->name; + return false; + } + const auto* attribute = call->attributes.Lookup("id"); + int32_t id = 0; + if (!attribute || !ParseId(*attribute, &id)) { + *error = definition->name + "." + call->name + + " requires an explicit nonnegative int32 (id: N)"; + return false; + } + if (!ids.insert(id).second) { + *error = "duplicate method id " + std::to_string(id) + + " in " + definition->name; + return false; + } + if (call->attributes.Lookup("streaming")) { + *error = "streaming RPC is not supported: " + call->name; + return false; + } + service.ids.push_back(id); + } + services->push_back(service); + } + if (services->empty()) { + *error = "input schema contains no rpc_service to generate"; + return false; + } + return true; +} + +void OpenNamespace(const flatbuffers::Definition& definition, + std::ostream& out) { + if (definition.defined_namespace) { + for (const auto& component : definition.defined_namespace->components) { + out << "namespace " << component << " {\n"; + } + } + out << '\n'; +} + +void CloseNamespace(const flatbuffers::Definition& definition, + std::ostream& out) { + if (definition.defined_namespace) { + const auto& components = definition.defined_namespace->components; + for (auto it = components.rbegin(); it != components.rend(); ++it) { + out << "} // namespace " << *it << '\n'; + } + } + out << '\n'; +} + +void GenerateHeader(const Service& service, std::ostream& out) { + const auto& definition = *service.definition; + const std::string& name = definition.name; + OpenNamespace(definition, out); + out << "class " << name << "_Stub;\n\n" + << "class " << name << " : public ::brpc::flatbuffers::Service {\n" + << "public:\n" + << " typedef " << name << "_Stub Stub;\n" + << " static const ::brpc::flatbuffers::ServiceDescriptor* " + "descriptor();\n" + << " const ::brpc::flatbuffers::ServiceDescriptor* " + "GetDescriptor() override;\n" + << " void FBCallMethod(\n" + << " const ::brpc::flatbuffers::MethodDescriptor* method,\n" + << " " << kArguments << ") override;\n"; + for (const auto* call : definition.calls.vec) { + out << " virtual void " << call->name << "(\n" + << " " << kArguments << ");\n"; + } + out << "};\n\n" + << "class " << name << "_Stub : public " << name << " {\n" + << "public:\n" + << " explicit " << name << "_Stub(\n" + << " ::brpc::flatbuffers::RpcChannel* channel,\n" + << " ::brpc::flatbuffers::Service::ChannelOwnership " + "ownership =\n" + << " ::brpc::flatbuffers::Service::" + "STUB_DOESNT_OWN_CHANNEL);\n" + << " ::brpc::flatbuffers::RpcChannel* channel() const {\n" + << " return channel_;\n" + << " }\n"; + for (const auto* call : definition.calls.vec) { + out << " void " << call->name << "(\n" + << " " << kArguments << ") override;\n"; + } + out << "\nprivate:\n" + << " ::brpc::flatbuffers::RpcChannel* channel_;\n" + << " ::std::unique_ptr<::brpc::flatbuffers::RpcChannel> " + "owned_channel_;\n" + << "};\n\n"; + CloseNamespace(definition, out); +} + +void GenerateFailureHandler(std::ostream& out) { + // A local callable cannot be hidden by schema declarations or RPC names. + out << " const auto fail = [controller, done](const char* reason) {\n" + << " if (controller) {\n" + << " controller->SetFailed(reason);\n" + << " }\n" + << " if (done) {\n" + << " done->Run();\n" + << " }\n" + << " };\n"; +} + +void GenerateSource(const Service& service, std::ostream& out) { + const auto& definition = *service.definition; + const std::string& name = definition.name; + OpenNamespace(definition, out); + out << "const ::brpc::flatbuffers::ServiceDescriptor* " << name + << "::descriptor() {\n" + << " struct Holder {\n" + << " ::brpc::flatbuffers::ServiceDescriptor value;\n" + << " Holder() {\n" + << " const ::brpc::flatbuffers::BrpcDescriptorTable " + "table = {\n" + << " \"" << JoinNamespace(definition, ".") << "\", \"" + << name << "\",\n \""; + for (size_t i = 0; i < definition.calls.vec.size(); ++i) { + out << (i == 0 ? "" : " ") << definition.calls.vec[i]->name; + } + out << "\", {"; + for (size_t i = 0; i < service.ids.size(); ++i) { + out << (i == 0 ? "" : ", ") << service.ids[i]; + } + out << "}\n };\n" + << " if (value.init(table) != 0) {\n" + << " throw ::std::runtime_error(\n" + << " \"invalid generated service descriptor\");\n" + << " }\n" + << " }\n" + << " };\n" + << " static const Holder holder;\n" + << " return &holder.value;\n" + << "}\n\n" + << "const ::brpc::flatbuffers::ServiceDescriptor* " << name + << "::GetDescriptor() {\n" + << " return descriptor();\n" + << "}\n\n" + << "void " << name << "::FBCallMethod(\n" + << " const ::brpc::flatbuffers::MethodDescriptor* method,\n" + << " " << kArguments << ") {\n"; + GenerateFailureHandler(out); + out << " if (!method || method->service() != descriptor() ||\n" + << " descriptor()->FindMethodByIndex(method->index()) != " + "method) {\n" + << " fail(\"invalid service method\");\n" + << " return;\n" + << " }\n" + << " if (!request || !response) {\n" + << " fail(\"null request or response\");\n" + << " return;\n" + << " }\n" + << " switch (method->index()) {\n"; + for (size_t i = 0; i < definition.calls.vec.size(); ++i) { + const auto& call = *definition.calls.vec[i]; + out << " case " << service.ids[i] << ":\n" + << " if (!request->Verify<" << FlatcQualifiedType(*call.request) + << ">()) {\n" + << " fail(\"invalid " + << call.request->name << " request\");\n" + << " return;\n" + << " }\n" + << " this->" << call.name + << "(controller, request, response, done);\n" + << " return;\n"; + } + out << " default:\n" + << " fail(\"unknown method id\");\n" + << " return;\n" + << " }\n" + << "}\n\n"; + for (const auto* call : definition.calls.vec) { + out << "void " << name << "::" << call->name << "(\n" + << " " << kArguments << ") {\n" + << " (void)request;\n" + << " (void)response;\n"; + GenerateFailureHandler(out); + out << " fail(\"method not implemented: " + << name << "." << call->name << "\");\n" + << "}\n\n"; + } + out << name << "_Stub::" << name << "_Stub(\n" + << " ::brpc::flatbuffers::RpcChannel* channel,\n" + << " ::brpc::flatbuffers::Service::ChannelOwnership ownership)\n" + << " : channel_(channel),\n" + << " owned_channel_(\n" + << " ownership == " + "::brpc::flatbuffers::Service::STUB_OWNS_CHANNEL\n" + << " ? channel : nullptr) {}\n\n"; + for (size_t i = 0; i < definition.calls.vec.size(); ++i) { + out << "void " << name << "_Stub::" << definition.calls.vec[i]->name + << "(\n " << kArguments << ") {\n"; + GenerateFailureHandler(out); + out << " if (!channel_) {\n" + << " fail(\"null RPC channel\");\n" + << " return;\n" + << " }\n" + << " channel_->FBCallMethod(\n" + << " descriptor()->method(" << i + << "), controller, request, response, done);\n" + << "}\n\n"; + } + CloseNamespace(definition, out); +} + +bool TestFailureEnabled(const char* step) { +#if defined(BRPC_FLATC_TESTING) + const char* value = std::getenv("BRPC_FLATC_TEST_FAILURE"); + if (value) { + const std::string failures = "," + std::string(value) + ","; + return failures.find("," + std::string(step) + ",") != + std::string::npos; + } +#else + (void)step; +#endif + return false; +} + +bool WriteFile(const std::string& path, const std::string& contents, + const char* step) { + if (TestFailureEnabled(step)) { + errno = EIO; + return false; + } + std::ofstream stream(path.c_str(), std::ios::binary | std::ios::trunc); + stream << contents; + stream.close(); + return !stream.fail(); +} + +bool RenameFile(const std::string& from, const std::string& to, + const char* step) { + if (TestFailureEnabled(step)) { + errno = EIO; + return false; + } + return std::rename(from.c_str(), to.c_str()) == 0; +} + +class OutputDirectoryLock { +public: + explicit OutputDirectoryLock(const std::string& directory) + : fd_(open(directory.c_str(), O_RDONLY)) { + if (fd_ < 0) { + return; + } + while (flock(fd_, LOCK_EX) != 0) { + if (errno != EINTR) { + close(fd_); + fd_ = -1; + return; + } + } + } + + ~OutputDirectoryLock() { + if (fd_ >= 0) { + close(fd_); + } + } + + bool acquired() const { return fd_ >= 0; } + +private: + OutputDirectoryLock(const OutputDirectoryLock&) = delete; + OutputDirectoryLock& operator=(const OutputDirectoryLock&) = delete; + + int fd_; +}; + +bool RemoveFile(const std::string& path) { + return std::remove(path.c_str()) == 0 || errno == ENOENT; +} + +bool IsRegularFileOrMissing(const std::string& path) { + struct stat file_status; + if (lstat(path.c_str(), &file_status) == 0) { + return S_ISREG(file_status.st_mode); + } + return errno == ENOENT; +} + +bool MoveExistingFile(const std::string& path, const std::string& backup, + const char* step, bool* moved) { + *moved = false; + if (RenameFile(path, backup, step)) { + *moved = true; + return true; + } + return errno == ENOENT; +} + +bool RestoreFile(const std::string& path, const std::string& backup, + const char* step, bool had_original) { + if (!RemoveFile(path)) { + return false; + } + return !had_original || RenameFile(backup, path, step); +} + +bool RemoveOutputPair(const std::string& header_path, + const std::string& source_path) { + const bool header_removed = RemoveFile(header_path); + const bool source_removed = RemoveFile(source_path); + return header_removed && source_removed; +} + +bool PublishFiles(const std::string& output_dir, + const std::string& header_path, + const std::string& header_contents, + const std::string& source_path, + const std::string& source_contents, + std::string* error) { + OutputDirectoryLock lock(output_dir); + if (!lock.acquired()) { + *error = "cannot lock output directory"; + return false; + } + const std::string header_tmp = header_path + ".tmp"; + const std::string source_tmp = source_path + ".tmp"; + const std::string header_backup = header_path + ".bak"; + const std::string source_backup = source_path + ".bak"; + if (!RemoveFile(header_tmp) || + !WriteFile(header_tmp, header_contents, "stage_header")) { + RemoveFile(header_tmp); + *error = "cannot stage header output"; + return false; + } + if (!RemoveFile(source_tmp) || + !WriteFile(source_tmp, source_contents, "stage_source")) { + RemoveFile(header_tmp); + RemoveFile(source_tmp); + *error = "cannot stage source output"; + return false; + } + if (!IsRegularFileOrMissing(header_path) || + !IsRegularFileOrMissing(source_path) || + !RemoveFile(header_backup) || !RemoveFile(source_backup)) { + RemoveFile(header_tmp); + RemoveFile(source_tmp); + *error = "output or backup path is not replaceable"; + return false; + } + + bool had_header = false; + bool had_source = false; + if (!MoveExistingFile(header_path, header_backup, + "preserve_header", &had_header)) { + RemoveFile(header_tmp); + RemoveFile(source_tmp); + *error = "cannot preserve previous header"; + return false; + } + if (!MoveExistingFile(source_path, source_backup, + "preserve_source", &had_source)) { + const bool restored = RestoreFile( + header_path, header_backup, "restore_header", had_header); + RemoveFile(header_tmp); + RemoveFile(source_tmp); + if (!restored) { + const bool cleared = RemoveOutputPair(header_path, source_path); + *error = cleared + ? "cannot preserve source or restore header" + : "cannot preserve source; manual cleanup required"; + } else { + *error = "cannot preserve previous source"; + } + return false; + } + + if (!RenameFile(header_tmp, header_path, "publish_header") || + !RenameFile(source_tmp, source_path, "publish_source")) { + const bool header_restored = RestoreFile( + header_path, header_backup, "restore_header", had_header); + const bool source_restored = RestoreFile( + source_path, source_backup, "restore_source", had_source); + RemoveFile(header_tmp); + RemoveFile(source_tmp); + if (!header_restored || !source_restored) { + const bool cleared = RemoveOutputPair(header_path, source_path); + *error = cleared + ? "cannot publish or restore output files" + : "cannot publish output files; manual cleanup required"; + } else { + *error = "cannot publish output files"; + } + return false; + } + + RemoveFile(header_backup); + RemoveFile(source_backup); + return true; +} + +void Usage(std::ostream& out) { + out << "Usage: brpc_flatc [-I include_dir]... " + "[-o existing_output_dir] schema.fbs\n" + << "Run official flatc --cpp separately to produce " + "schema_generated.h.\n"; +} + +int Run(int argc, char** argv) { + if (std::string(flatbuffers::FLATBUFFERS_VERSION()) != "25.2.10") { + std::cerr << "brpc_flatc: requires FlatBuffers 25.2.10 library\n"; + return 1; + } + if (argc == 2 && std::string(argv[1]) == "--version") { + std::cout << "brpc_flatc (FlatBuffers 25.2.10)\n"; + return 0; + } + std::string input; + std::string output_dir = "."; + std::vector include_dirs; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--help" || arg == "-h") { + Usage(std::cout); + return 0; + } + if (arg == "-o" || arg == "-I") { + if (++i == argc) { + Usage(std::cerr); + return 1; + } + if (arg == "-o") { + output_dir = argv[i]; + } else { + include_dirs.push_back(argv[i]); + } + } else if (arg.compare(0, 2, "-I") == 0 && arg.size() > 2) { + include_dirs.push_back(arg.substr(2)); + } else if (arg.empty() || arg[0] == '-' || !input.empty()) { + Usage(std::cerr); + return 1; + } else { + input = arg; + } + } + if (input.size() < 5 || input.substr(input.size() - 4) != ".fbs" || + output_dir.empty()) { + Usage(std::cerr); + return 1; + } + const size_t slash = input.find_last_of("/\\"); + const std::string basename = input.substr( + slash == std::string::npos ? 0 : slash + 1); + const std::string stem = basename.substr(0, basename.size() - 4); + if (stem.empty() || stem.find_first_not_of( + "abcdefghijklmnopqrstuvwxyz" + "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_.-") != + std::string::npos) { + std::cerr << "brpc_flatc: unsupported schema filename\n"; + return 1; + } + include_dirs.push_back(slash == std::string::npos ? "." : + input.substr(0, slash == 0 ? 1 : slash)); + std::vector include_paths; + for (const auto& directory : include_dirs) { + include_paths.push_back(directory.c_str()); + } + include_paths.push_back(nullptr); + std::ifstream stream(input.c_str(), std::ios::binary); + if (!stream) { + std::cerr << "brpc_flatc: cannot read " << input << '\n'; + return 1; + } + const std::string schema((std::istreambuf_iterator(stream)), + std::istreambuf_iterator()); + if (stream.bad() || schema.find('\0') != std::string::npos) { + std::cerr << "brpc_flatc: invalid schema input\n"; + return 1; + } + flatbuffers::Parser parser; + if (!parser.Parse(schema.c_str(), include_paths.data(), input.c_str())) { + std::cerr << parser.error_ << '\n'; + return 1; + } + std::vector services; + std::string error; + if (!CollectServices(parser, input, include_paths.data(), &services, &error)) { + std::cerr << "brpc_flatc: " << error << '\n'; + return 1; + } + // Coexisting service headers cannot define the same qualified service. + // Preserve filename bytes and that service identity, without tying output + // to checkout paths or folding punctuation and case into the same guard. + const std::string guard = + "BRPC_FLATBUFFERS_GENERATED_" + GuardComponent(stem) + "_" + + GuardComponent(Qualified(*services.front().definition)) + "_H_"; + std::ostringstream header; + header << kLicense << "#ifndef " << guard << "\n#define " << guard << "\n\n" + << "#include \n" + << "#include \n" + << "#include \"brpc/flatbuffers/message.h\"\n" + << "#include \"brpc/flatbuffers/service.h\"\n" + << "#include \"" << stem << "_generated.h\"\n\n" + << "#if !BRPC_WITH_FLATBUFFERS\n" + << "#error \"Generated services require BRPC_WITH_FLATBUFFERS\"\n" + << "#endif\n\n"; + std::ostringstream source; + source << kLicense << "#include \"" << stem << ".brpc.fb.h\"\n\n" + << "#include \n\n"; + for (const auto& service : services) { + GenerateHeader(service, header); + GenerateSource(service, source); + } + header << "#endif // " << guard << '\n'; + if (output_dir.back() != '/') { + output_dir += '/'; + } + if (!PublishFiles(output_dir, + output_dir + stem + ".brpc.fb.h", header.str(), + output_dir + stem + ".brpc.fb.cpp", source.str(), + &error)) { + std::cerr << "brpc_flatc: cannot write output in " << output_dir + << ": " << error << '\n'; + return 1; + } + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + try { + return Run(argc, argv); + } catch (const std::exception& error) { + std::cerr << "brpc_flatc: " << error.what() << '\n'; + return 1; + } +} diff --git a/tools/flatbuffers/verify.sh b/tools/flatbuffers/verify.sh new file mode 100644 index 0000000000..a0ce7ae7f4 --- /dev/null +++ b/tools/flatbuffers/verify.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash +# 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. + +set -euo pipefail +if [[ $# != 1 ]]; then + printf 'Usage: bash tools/flatbuffers/verify.sh \n' >&2 + exit 2 +fi +root=$(cd "$(dirname "$0")/../.." && pwd) +mkdir -p "$1" +work=$(cd "$1" && pwd) +cmake=${CMAKE:-cmake} +ctest=${CTEST:-ctest} +cxx=${CXX:-c++} +jobs=${JOBS:-2} +runtime=${CODEGEN_RUNTIME:-ON} +case "$runtime" in + ON|OFF) ;; + *) printf 'CODEGEN_RUNTIME must be ON or OFF\n' >&2; exit 2 ;; +esac + +printf 'Source: %s\nWork: %s\nCODEGEN_RUNTIME: %s\n' "$root" "$work" "$runtime" +uname -sm +"$cxx" --version +"$cmake" --version +"$ctest" --version +"$cmake" -E sha256sum "$root/tools/flatbuffers/brpc_flatc.cpp" \ + "$root/test/flatbuffers_codegen/acceptance.cmake" \ + "$root/test/flatbuffers_codegen/CMakeLists.txt" \ + "$root/test/flatbuffers_codegen/echo.fbs" \ + "$root/test/flatbuffers_codegen/runtime.cpp" + +# Keep the archive and checksum in sync with MODULE.bazel and WORKSPACE. +if [[ -z ${FLATBUFFERS_PREFIX:-} ]]; then + archive="$work/flatbuffers-25.2.10.tar.gz" + curl --fail --location --retry 3 \ + https://github.com/google/flatbuffers/archive/refs/tags/v25.2.10.tar.gz \ + --output "$archive" + checksum=$("$cmake" -E sha256sum "$archive") + printf '%s\n' "$checksum" + if [[ ${checksum%% *} != b9c2df49707c57a48fc0923d52b8c73beb72d675f9d44b2211e4569be40a7421 ]]; then + printf 'FlatBuffers archive checksum mismatch\n' >&2 + exit 1 + fi + tar -xzf "$archive" -C "$work" + fb="$work/flatbuffers-install" + printf 'FlatBuffers build: %s\n' "$work/flatbuffers-build" + "$cmake" -S "$work/flatbuffers-25.2.10" -B "$work/flatbuffers-build" \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_CXX_COMPILER="$cxx" \ + -DCMAKE_INSTALL_PREFIX="$fb" -DFLATBUFFERS_BUILD_TESTS=OFF \ + -DFLATBUFFERS_BUILD_FLATC=ON + "$cmake" --build "$work/flatbuffers-build" --target install --parallel "$jobs" +else + fb=$FLATBUFFERS_PREFIX +fi +printf 'FlatBuffers prefix: %s\n' "$fb" +"$fb/bin/flatc" --version + +runtime_library= +codegen_tests=(flatbuffers_codegen_acceptance) +if [[ $runtime == ON ]]; then + # Match the root GTest version without its unverified Git download path. + archive="$work/googletest-release-1.12.0.tar.gz" + curl --fail --location --retry 3 \ + https://github.com/google/googletest/archive/refs/tags/release-1.12.0.tar.gz \ + --output "$archive" + checksum=$("$cmake" -E sha256sum "$archive") + printf '%s\n' "$checksum" + if [[ ${checksum%% *} != 2a4f11dce6188b256f3650061525d0fe352069e5c162452818efbbf8d0b5fe1c ]]; then + printf 'GTest archive checksum mismatch\n' >&2 + exit 1 + fi + tar -xzf "$archive" -C "$work" + gtest="$work/googletest-release-1.12.0" + printf 'GTest source (1.12.0): %s\nbRPC build: %s\n' "$gtest" "$work/brpc" + "$cmake" -S "$root" -B "$work/brpc" \ + -DCMAKE_BUILD_TYPE=Release -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \ + -DCMAKE_CXX_COMPILER="$cxx" -DBUILD_SHARED_LIBS=ON \ + -DBUILD_UNIT_TESTS=ON -DDOWNLOAD_GTEST=OFF \ + -DBRPC_SYSTEM_GTEST_SOURCE_DIR="$gtest" \ + -DWITH_FLATBUFFERS=ON -DFLATBUFFERS_INCLUDE_DIR="$fb/include" \ + -DFLATBUFFERS_FLATC_EXECUTABLE="$fb/bin/flatc" \ + -DOPENSSL_ROOT_DIR="${OPENSSL_ROOT_DIR:-}" \ + -DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" + "$cmake" --build "$work/brpc" --target brpc_flatbuffers_unittest --parallel "$jobs" + # A registered CTest can still pass with zero GTest cases or all cases skipped. + report=$(mktemp "$work/brpc/flatbuffers-gtest.XXXXXX") + printf 'Core GTest XML: %s\n' "$report" + ( + unset GTEST_TOTAL_SHARDS GTEST_SHARD_INDEX GTEST_SHARD_STATUS_FILE + GTEST_FILTER='*' GTEST_REPEAT=1 GTEST_OUTPUT="xml:$report" \ + "$ctest" --test-dir "$work/brpc" --no-tests=error \ + -R '^brpc_flatbuffers_unittest$' --output-on-failure --verbose + ) + if ! grep -Eq '&2 + exit 1 + fi + "$cmake" --build "$work/brpc" --target brpc-shared --parallel "$jobs" + case $(uname -s) in + Darwin) extension=dylib ;; + *) extension=so ;; + esac + runtime_library="$work/brpc/output/lib/libbrpc.$extension" + codegen_tests+=(flatbuffers_codegen_runtime) +else + printf 'Compile-only codegen verification: skipping core and runtime tests\n' +fi +printf 'Codegen build: %s\n' "$work/codegen" +"$cmake" -S "$root/tools/flatbuffers" -B "$work/codegen" \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON \ + -DCMAKE_CXX_COMPILER="$cxx" -DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}" \ + -DFLATBUFFERS_INCLUDE_DIR="$fb/include" \ + -DFLATBUFFERS_LIBRARY="$fb/lib/libflatbuffers.a" \ + -DFLATC_EXECUTABLE="$fb/bin/flatc" -DOPENSSL_ROOT_DIR="${OPENSSL_ROOT_DIR:-}" \ + -DBRPC_CODEGEN_BRPC_LIBRARY="$runtime_library" +"$cmake" --build "$work/codegen" --parallel "$jobs" +"$work/codegen/brpc_flatc" --version +for test_name in "${codegen_tests[@]}"; do + "$ctest" --test-dir "$work/codegen" --no-tests=error \ + -R "^${test_name}$" --output-on-failure --verbose +done From 8d4aa518352d7dc51e5dbf9f37c12dd1482e3e2c Mon Sep 17 00:00:00 2001 From: Bright Chen Date: Thu, 22 Oct 2026 21:21:24 +0800 Subject: [PATCH 04/10] Fix macOS futex wake accounting with per-waiter notifications (#3637) * Fix macOS futex wake accounting with per-waiter notifications * Fix ut * Guard macOS futex waiter cleanup and fix timeout test diagnostic --- src/bthread/sys_futex.cpp | 76 ++++++++++++++++-------- test/bthread_futex_unittest.cpp | 101 ++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+), 25 deletions(-) diff --git a/src/bthread/sys_futex.cpp b/src/bthread/sys_futex.cpp index 010ddcada2..b6ebdd6013 100644 --- a/src/bthread/sys_futex.cpp +++ b/src/bthread/sys_futex.cpp @@ -22,6 +22,7 @@ #include "bthread/sys_futex.h" #include "butil/scoped_lock.h" #include "butil/atomicops.h" +#include "butil/containers/linked_list.h" #include #include @@ -29,22 +30,35 @@ namespace bthread { -class SimuFutex { -public: - SimuFutex() : counts(0) - , ref(0) { - pthread_mutex_init(&lock, nullptr); +struct SimuFutexWaiter : public butil::LinkNode { + SimuFutexWaiter() : signaled(false) { pthread_cond_init(&cond, nullptr); } + ~SimuFutexWaiter() { + // The futex mutex is still held when the stack waiter is destroyed. + // Notified waiters were detached by futex_wake_private. Remove only + // waiters that leave on timeout or error. + if (!signaled) { + RemoveFromList(); + } + pthread_cond_destroy(&cond); + } + + pthread_cond_t cond; + bool signaled; +}; + +struct SimuFutex { + SimuFutex() : ref(0) { + pthread_mutex_init(&lock, nullptr); + } ~SimuFutex() { pthread_mutex_destroy(&lock); - pthread_cond_destroy(&cond); } -public: pthread_mutex_t lock; - pthread_cond_t cond; - int32_t counts; + // Only waiters that have not been notified are in this queue. + butil::LinkedList waiters; int32_t ref; }; @@ -70,20 +84,28 @@ int futex_wait_private(void* addr1, int expected, const timespec* timeout) { { std::unique_lock mu1(simu_futex.lock); if (static_cast*>(addr1)->load() == expected) { - ++simu_futex.counts; + SimuFutexWaiter waiter; + simu_futex.waiters.Append(&waiter); + timespec timeout_abs = {}; if (timeout) { - timespec timeout_abs = butil::timespec_from_now(*timeout); - if ((rc = pthread_cond_timedwait(&simu_futex.cond, &simu_futex.lock, &timeout_abs)) != 0) { - errno = rc; - rc = -1; - } - } else { - if ((rc = pthread_cond_wait(&simu_futex.cond, &simu_futex.lock)) != 0) { - errno = rc; - rc = -1; + timeout_abs = butil::timespec_from_now(*timeout); + } + do { + if (timeout) { + rc = pthread_cond_timedwait( + &waiter.cond, &simu_futex.lock, &timeout_abs); + } else { + rc = pthread_cond_wait(&waiter.cond, &simu_futex.lock); } + } while (rc == 0 && !waiter.signaled); + // A wake may win the mutex while a timed-out waiter is reacquiring + // it. Honor that notification so it cannot be counted and then lost. + if (waiter.signaled) { + rc = 0; + } else if (rc != 0) { + errno = rc; + rc = -1; } - --simu_futex.counts; } else { errno = EAGAIN; rc = -1; @@ -117,14 +139,18 @@ int futex_wake_private(void* addr1, int nwake) { int rc = 0; { std::unique_lock mu1(simu_futex.lock); - nwake = (nwake < simu_futex.counts)? nwake: simu_futex.counts; - for (int i = 0; i < nwake; ++i) { - if ((rc = pthread_cond_signal(&simu_futex.cond)) != 0) { + while (nwakedup < nwake && !simu_futex.waiters.empty()) { + SimuFutexWaiter* waiter = simu_futex.waiters.head()->value(); + if ((rc = pthread_cond_signal(&waiter->cond)) != 0) { errno = rc; break; - } else { - ++nwakedup; } + // Commit the notification to this waiter before releasing the + // mutex. Later wakes cannot count it again while it is rescheduled, + // and another waiter's timeout cannot consume its notification. + waiter->signaled = true; + waiter->RemoveFromList(); + ++nwakedup; } } diff --git a/test/bthread_futex_unittest.cpp b/test/bthread_futex_unittest.cpp index ce32356e52..d08982b62d 100644 --- a/test/bthread_futex_unittest.cpp +++ b/test/bthread_futex_unittest.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include "butil/time.h" @@ -115,6 +116,106 @@ TEST(FutexTest, futex_wake_before_wait) { ASSERT_EQ(ETIMEDOUT, errno); } +const int kWakeCountIterations = 256; +const int64_t kWakeCountTimeoutUs = 5 * 1000 * 1000L; +const useconds_t kWakeCountRetryIntervalUs = 50; + +struct WakeCountWaiterArg { + butil::atomic* word; + butil::atomic* finished; + int64_t deadline_us; + int rc; + int error; +}; + +void* wake_count_waiter(void* void_arg) { + WakeCountWaiterArg* arg = static_cast(void_arg); + const int64_t remaining_us = arg->deadline_us - butil::cpuwide_time_us(); + const timespec timeout = butil::microseconds_to_timespec( + remaining_us > 0 ? remaining_us : 0); + arg->rc = bthread::futex_wait_private(arg->word, 0, &timeout); + arg->error = errno; + arg->finished->store(true, butil::memory_order_release); + return nullptr; +} + +// A successful wake confirms that the waiter was registered. The waiter never +// waits again, so every subsequent wake must return zero, even before it runs. +// Repeat to exercise both orders of reacquiring the simulated futex's mutex. +void check_wake_count_after_notification(bool expire_another_waiter) { + // Share one deadline across all iterations and the waiter itself, so slow + // scheduling cannot renew the timeout on every iteration. + const int64_t deadline_us = butil::cpuwide_time_us() + kWakeCountTimeoutUs; + for (int i = 0; i < kWakeCountIterations; ++i) { + ASSERT_LT(butil::cpuwide_time_us(), deadline_us) + << "Timed out after " << i << " iterations"; + butil::atomic word(0); + butil::atomic finished(false); + WakeCountWaiterArg arg = { &word, &finished, deadline_us, -1, 0 }; + pthread_t waiter; + ASSERT_EQ(0, pthread_create(&waiter, nullptr, wake_count_waiter, &arg)); + + int first_wake = 0; + // A readiness flag published before futex_wait does not prove that the + // waiter is queued. Use the wake result, with a pause between retries. + while (butil::cpuwide_time_us() < deadline_us) { + first_wake = bthread::futex_wake_private(&word, 1); + if (first_wake != 0 || finished.load(butil::memory_order_acquire)) { + break; + } + usleep(kWakeCountRetryIntervalUs); + } + + int expired_rc = 0; + int expired_errno = 0; + int second_wake = 0; + if (first_wake == 1) { + if (expire_another_waiter) { + const timespec timeout = { 0, 0 }; + expired_rc = bthread::futex_wait_private(&word, 0, &timeout); + expired_errno = errno; + } + second_wake = bthread::futex_wake_private(&word, 1); + } + // The waiter has a timeout even if no wake succeeds. Join before + // checking results so a failed assertion cannot leave it referring to + // this iteration's stack objects. + ASSERT_EQ(0, pthread_join(waiter, nullptr)); + ASSERT_EQ(1, first_wake) + << "No waiter was notified before the deadline: iteration=" << i + << " waiter_rc=" << arg.rc << " errno=" << arg.error; + ASSERT_EQ(0, arg.rc) << "errno=" << arg.error; + if (expire_another_waiter) { + ASSERT_EQ(-1, expired_rc); + ASSERT_EQ(ETIMEDOUT, expired_errno); + } + ASSERT_EQ(0, second_wake) << "iteration=" << i; + } +} + +TEST(FutexTest, notified_waiter_is_not_counted_twice) { + check_wake_count_after_notification(false); +} + +TEST(FutexTest, timeout_does_not_consume_another_waiters_notification) { + check_wake_count_after_notification(true); +} + +TEST(FutexTest, timed_out_waiter_is_removed) { + butil::atomic word(0); + const timespec timeout = { 0, 0 }; + ASSERT_EQ(-1, bthread::futex_wait_private(&word, 0, &timeout)); + ASSERT_EQ(ETIMEDOUT, errno); + ASSERT_EQ(0, bthread::futex_wake_private(&word, INT_MAX)); +} + +TEST(FutexTest, mismatched_value_does_not_queue_a_waiter) { + butil::atomic word(1); + ASSERT_EQ(-1, bthread::futex_wait_private(&word, 0, nullptr)); + ASSERT_EQ(EAGAIN, errno); + ASSERT_EQ(0, bthread::futex_wake_private(&word, INT_MAX)); +} + struct DummyWaiterArg { butil::atomic* lock; butil::atomic* registered; From e6fd4c5a803bfe4c803456081f064632cc883539 Mon Sep 17 00:00:00 2001 From: Bright Chen Date: Thu, 22 Oct 2026 21:28:56 +0800 Subject: [PATCH 05/10] Fix JSON parsing offsets at EOF (#3635) * Fix JSON parsing offsets at EOF * Opt * Fix ut --- src/butil/third_party/rapidjson/reader.h | 7 +- test/brpc_protobuf_json_unittest.cpp | 96 ++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 4 deletions(-) diff --git a/src/butil/third_party/rapidjson/reader.h b/src/butil/third_party/rapidjson/reader.h index 26e53c2a63..938a082725 100644 --- a/src/butil/third_party/rapidjson/reader.h +++ b/src/butil/third_party/rapidjson/reader.h @@ -1539,10 +1539,9 @@ class GenericReader { if ((parseFlags & kParseStopWhenDoneFlag) && state == IterativeParsingFinishState) { // wwb: Update parseResult_.Offset() when kParseStopWhenDoneFlag // is set which means the user needs to know where to resume - // parsing in next calls to JsonToProtoMessage() - if (is.Peek() != '\0') { - SetParseError(kParseErrorNone, is.Tell()); - } + // parsing in next calls to JsonToProtoMessage(), including + // when the root ends exactly at EOF. + SetParseError(kParseErrorNone, is.Tell()); break; } SkipWhitespace(is); diff --git a/test/brpc_protobuf_json_unittest.cpp b/test/brpc_protobuf_json_unittest.cpp index eeadc4f7d2..e38f3b9ccc 100644 --- a/test/brpc_protobuf_json_unittest.cpp +++ b/test/brpc_protobuf_json_unittest.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include "butil/iobuf.h" @@ -1689,6 +1690,101 @@ TEST_F(ProtobufJsonTest, string_to_int64) { ASSERT_EQ(person.data(), 1234567); } +TEST_F(ProtobufJsonTest, parse_json_offset_at_eof) { + const struct { + const char* json; + size_t expected_offset; + bool array_to_single_repeated; + } cases[] = { + {"{}", 2, false}, + {" \t{}", 4, false}, + {R"({"person":[]})", 13, false}, + {"[]", 2, true}, + {"{} \n", 2, false}, + {"{}{}", 2, false}, + }; + + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.json); + json2pb::Json2PbOptions options; + options.allow_remaining_bytes_after_parsing = true; + options.array_to_single_repeated = test_case.array_to_single_repeated; + AddressBook message; + std::string error; + size_t offset = 0; + ASSERT_TRUE(json2pb::JsonToProtoMessage( + test_case.json, &message, options, &error, &offset)) << error; + EXPECT_EQ(test_case.expected_offset, offset); + EXPECT_TRUE(error.empty()); + + butil::IOBuf buf; + buf.append(test_case.json); + butil::IOBufAsZeroCopyInputStream stream(buf); + message.Clear(); + offset = 0; + ASSERT_TRUE(json2pb::JsonToProtoMessage( + &stream, &message, options, &error, &offset)) << error; + EXPECT_EQ(test_case.expected_offset, offset); + EXPECT_TRUE(error.empty()); + } +} + +TEST_F(ProtobufJsonTest, parse_multiple_json_at_eof) { + const std::string jsons[] = { + R"({"message":"first"})", + R"({"message":"second"})", + }; + const char* expected_messages[] = {"first", "second"}; + std::string remaining = jsons[0] + jsons[1]; + butil::IOBuf buf; + buf.append(remaining); + butil::IOBufAsZeroCopyInputStream stream(buf); + json2pb::ZeroCopyStreamReader reader(&stream); + json2pb::Json2PbOptions options; + options.allow_remaining_bytes_after_parsing = true; + test::EchoRequest message; + std::string error; + size_t string_offset = 0; + size_t stream_offset = 0; + size_t total_offset = 0; + + for (size_t i = 0; i < 2; ++i) { + SCOPED_TRACE(i); + message.Clear(); + ASSERT_TRUE(json2pb::JsonToProtoMessage( + remaining, &message, options, &error, &string_offset)) << error; + EXPECT_EQ(expected_messages[i], message.message()); + EXPECT_EQ(jsons[i].size(), string_offset); + remaining = remaining.substr(string_offset); + + message.Clear(); + ASSERT_TRUE(json2pb::JsonToProtoMessage( + &reader, &message, options, &error, &stream_offset)) << error; + EXPECT_EQ(expected_messages[i], message.message()); + total_offset += jsons[i].size(); + EXPECT_EQ(total_offset, stream_offset); + EXPECT_EQ(total_offset, reader.Tell()); + } + ASSERT_TRUE(remaining.empty()); + // Seed outputs to verify that EOF parsing overwrites them. + message.Clear(); + error = "previous error"; + string_offset = std::numeric_limits::max(); + EXPECT_FALSE(json2pb::JsonToProtoMessage( + remaining, &message, options, &error, &string_offset)); + EXPECT_TRUE(error.empty()); + EXPECT_EQ(0ul, string_offset); + + message.Clear(); + error = "previous error"; + stream_offset = std::numeric_limits::max(); + EXPECT_FALSE(json2pb::JsonToProtoMessage( + &reader, &message, options, &error, &stream_offset)); + EXPECT_TRUE(error.empty()); + EXPECT_EQ(total_offset, stream_offset); + EXPECT_EQ(total_offset, reader.Tell()); +} + TEST_F(ProtobufJsonTest, parse_multiple_json) { const int COUNT = 4; std::vector expectedNames = { "tom", "bob", "jerry", "lucy" }; From 81c15259750a9659865a1fae307aaf7614459630 Mon Sep 17 00:00:00 2001 From: nkysg Date: Fri, 23 Oct 2026 10:37:01 +0800 Subject: [PATCH 06/10] Fix ResourcePool block item count race (#3640) * Fix ResourcePool block item count race * Optimize ResourcePool nitem publication --- src/butil/resource_pool_inl.h | 36 +++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/src/butil/resource_pool_inl.h b/src/butil/resource_pool_inl.h index 163b033443..57134c9011 100644 --- a/src/butil/resource_pool_inl.h +++ b/src/butil/resource_pool_inl.h @@ -116,7 +116,10 @@ class BAIDU_CACHELINE_ALIGNMENT ResourcePool { // To support cache-aligned objects, align Block.items by cacheline. struct BAIDU_CACHELINE_ALIGNMENT Block { BlockItem items[BLOCK_NITEM]; - size_t nitem; + // Written by the owning LocalPool and read by address_resource() from + // other threads. Publishing the increment after construction keeps the + // bounds check data-race-free. + butil::atomic nitem; Block() : nitem(0) {} }; @@ -184,28 +187,32 @@ class BAIDU_CACHELINE_ALIGNMENT ResourcePool { } \ T* p = nullptr; \ /* Fetch memory from local block */ \ - if (_cur_block && _cur_block->nitem < BLOCK_NITEM) { \ - id->value = _cur_block_index * BLOCK_NITEM + _cur_block->nitem; \ - auto item = _cur_block->items + _cur_block->nitem; \ - p = new (item->void_data()) T CTOR_ARGS; \ - if (!ResourcePoolValidator::validate(p)) { \ - p->~T(); \ - return nullptr; \ + if (_cur_block) { \ + const size_t nitem = _cur_block->nitem.load(butil::memory_order_relaxed); \ + if (nitem < BLOCK_NITEM) { \ + id->value = _cur_block_index * BLOCK_NITEM + nitem; \ + auto item = _cur_block->items + nitem; \ + p = new (item->void_data()) T CTOR_ARGS; \ + if (!ResourcePoolValidator::validate(p)) { \ + p->~T(); \ + return nullptr; \ + } \ + _cur_block->nitem.store(nitem + 1, butil::memory_order_release); \ + return p; \ } \ - ++_cur_block->nitem; \ - return p; \ } \ /* Fetch a Block from global */ \ _cur_block = add_block(&_cur_block_index); \ if (_cur_block != nullptr) { \ - id->value = _cur_block_index * BLOCK_NITEM + _cur_block->nitem; \ - auto item = _cur_block->items + _cur_block->nitem; \ + const size_t nitem = _cur_block->nitem.load(butil::memory_order_relaxed); \ + id->value = _cur_block_index * BLOCK_NITEM + nitem; \ + auto item = _cur_block->items + nitem; \ p = new (item->void_data()) T CTOR_ARGS; \ if (!ResourcePoolValidator::validate(p)) { \ p->~T(); \ return nullptr; \ } \ - ++_cur_block->nitem; \ + _cur_block->nitem.store(nitem + 1, butil::memory_order_release); \ return p; \ } \ return nullptr; \ @@ -267,7 +274,8 @@ class BAIDU_CACHELINE_ALIGNMENT ResourcePool { .load(butil::memory_order_consume); if (__builtin_expect(b != nullptr, 1)) { const size_t offset = id.value - block_index * BLOCK_NITEM; - if (__builtin_expect(offset < b->nitem, 1)) { + if (__builtin_expect( + offset < b->nitem.load(butil::memory_order_acquire), 1)) { return (T*)b->items + offset; } } From 6ff47acf554262f6fc2d0b2c6fe538b8208a6675 Mon Sep 17 00:00:00 2001 From: darion-yaphet Date: Fri, 23 Oct 2026 12:33:42 +0800 Subject: [PATCH 07/10] fix(bthread): Synchronize keytable pool pointer access (#3638) * fix(bthread): synchronize keytable pool access Read pool pointers and the destroyed state under the existing rwlock. Keep only a boolean availability hint between the read and write lock sections, then recheck the pool state under the write lock. Add coverage for concurrent borrowing, returning, and pool destruction, including reclamation of tables returned after destruction. * test(bthread): ensure keytable return after pool destruction * test(bthread): isolate keytable pool test flag --- src/bthread/key.cpp | 10 ++++- test/bthread_key_unittest.cpp | 75 +++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/src/bthread/key.cpp b/src/bthread/key.cpp index 8db23e0816..74e718a440 100644 --- a/src/bthread/key.cpp +++ b/src/bthread/key.cpp @@ -322,10 +322,15 @@ class BAIDU_CACHELINE_ALIGNMENT KeyTableList { }; KeyTable* borrow_keytable(bthread_keytable_pool_t* pool) { - if (pool != nullptr && (pool->list || pool->free_keytables)) { + if (pool != nullptr) { KeyTable* p; + bool has_free_keytables; { pthread_rwlock_rdlock(&pool->rwlock); + if (pool->destroyed) { + pthread_rwlock_unlock(&pool->rwlock); + return nullptr; + } auto list = (butil::ThreadLocal*)pool->list; if (list) { p = list->get()->remove_front(); @@ -334,9 +339,10 @@ KeyTable* borrow_keytable(bthread_keytable_pool_t* pool) { return p; } } + has_free_keytables = pool->free_keytables != nullptr; pthread_rwlock_unlock(&pool->rwlock); } - if (pool->free_keytables) { + if (has_free_keytables) { pthread_rwlock_wrlock(&pool->rwlock); if (pool->destroyed) { pthread_rwlock_unlock(&pool->rwlock); diff --git a/test/bthread_key_unittest.cpp b/test/bthread_key_unittest.cpp index f0bfadec42..a883ad1ec0 100644 --- a/test/bthread_key_unittest.cpp +++ b/test/bthread_key_unittest.cpp @@ -26,6 +26,7 @@ #include "butil/scoped_lock.h" #include "butil/logging.h" #include "bthread/bthread.h" +#include "bthread/countdown_event.h" #include "bthread/unstable.h" using namespace bthread; namespace bthread { @@ -566,6 +567,80 @@ TEST(KeyTest, borrow_and_return_keytable_when_using_pool) { ASSERT_EQ(0, bthread_keytable_pool_destroy(&test_pool)); } +struct ConcurrentPoolData { + bthread_keytable_pool_t pool; + CountdownEvent ready; + CountdownEvent start; + CountdownEvent destroyed; + Counters counters; + + explicit ConcurrentPoolData(int nworkers) : ready(nworkers) {} +}; + +static void* create_concurrent_pool_data(const void* arg) { + Counters* counters = const_cast(static_cast(arg)); + counters->ncreate.fetch_add(1, butil::memory_order_relaxed); + return counters; +} + +static void destroy_concurrent_pool_data(void* arg) { + static_cast(arg)->ndestroy.fetch_add(1, butil::memory_order_relaxed); +} + +static void* borrow_and_return_during_pool_destruction(void* arg) { + ConcurrentPoolData* data = static_cast(arg); + // Keep a borrowed table across destruction as well as racing new borrows + // with destruction. The pool and key remain valid until all workers exit. + KeyTable* held = borrow_keytable(&data->pool); + EXPECT_NE(nullptr, held); + data->ready.signal(); + data->start.wait(); + for (int i = 0; i < 4096; ++i) { + KeyTable* kt = borrow_keytable(&data->pool); + return_keytable(&data->pool, kt); + } + data->destroyed.wait(); + return_keytable(&data->pool, held); + return nullptr; +} + +TEST(KeyTest, borrow_and_return_keytable_during_pool_destruction) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + bthread::FLAGS_borrow_from_globle_size = 20; + const size_t NWORKERS = 8; + const size_t NTABLES = NWORKERS * bthread::FLAGS_borrow_from_globle_size; + ConcurrentPoolData data(NWORKERS); + ASSERT_EQ(0, bthread_keytable_pool_init(&data.pool)); + bthread_key_t key; + ASSERT_EQ(0, bthread_key_create(&key, destroy_concurrent_pool_data)); + bthread_keytable_pool_reserve(&data.pool, NTABLES, key, + create_concurrent_pool_data, &data.counters); + ASSERT_EQ(NTABLES, data.counters.ncreate.load(butil::memory_order_relaxed)); + + pthread_t workers[NWORKERS]; + size_t nstarted = 0; + for (; nstarted < NWORKERS; ++nstarted) { + int rc = pthread_create(&workers[nstarted], nullptr, + borrow_and_return_during_pool_destruction, &data); + EXPECT_EQ(0, rc); + if (rc != 0) { + data.ready.signal(NWORKERS - nstarted); + break; + } + } + data.ready.wait(); + data.start.signal(); + EXPECT_EQ(0, bthread_keytable_pool_destroy(&data.pool)); + data.destroyed.signal(); + for (size_t i = 0; i < nstarted; ++i) { + EXPECT_EQ(0, pthread_join(workers[i], nullptr)); + } + EXPECT_EQ(nullptr, borrow_keytable(&data.pool)); + EXPECT_EQ(NTABLES, data.counters.ndestroy.load(butil::memory_order_relaxed)); + EXPECT_EQ(0, bthread_key_delete(key)); + EXPECT_EQ(0, pthread_rwlock_destroy(&data.pool.rwlock)); +} + // NOTE: lid is short for 'lock in dtor'. butil::atomic lid_seq(1); std::vector lid_seqs; From 96babbd25550b4b431ebd56ef86178ca84bbe47e Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Fri, 23 Oct 2026 01:58:43 -0700 Subject: [PATCH 08/10] Retrigger community CI From 9ae366f8bd125379d42611fdf1243531a656eb5a Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Fri, 23 Oct 2026 02:17:35 -0700 Subject: [PATCH 09/10] Rename the common transport handshake protocol --- src/brpc/global.cpp | 7 ++----- src/brpc/options.proto | 2 +- src/brpc/server.cpp | 8 ++++---- src/brpc/server.h | 4 ++-- test/brpc_server_unittest.cpp | 6 +++--- test/brpc_transport_handshake_unittest.cpp | 16 ++++++++++++++++ 6 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/brpc/global.cpp b/src/brpc/global.cpp index 7aa6a91cdf..5703eeeadd 100644 --- a/src/brpc/global.cpp +++ b/src/brpc/global.cpp @@ -442,11 +442,8 @@ static void GlobalInitializeOrDieImpl() { ParseTransportHandshake, nullptr, nullptr, ProcessTransportHandshake, nullptr, nullptr, nullptr, nullptr, - CONNECTION_TYPE_ALL, "rdma_handshake" }; - // Retain the existing enum value and registered name to avoid changing - // public protocol identifiers while widening the implementation from RDMA - // to all transport upgrades. - if (RegisterProtocol(PROTOCOL_RDMA_HANDSHAKE, + CONNECTION_TYPE_ALL, "transport_handshake" }; + if (RegisterProtocol(PROTOCOL_TRANSPORT_HANDSHAKE, transport_handshake_protocol) != 0) { exit(1); } diff --git a/src/brpc/options.proto b/src/brpc/options.proto index 13b8b682e6..450e08629c 100644 --- a/src/brpc/options.proto +++ b/src/brpc/options.proto @@ -37,7 +37,7 @@ enum ConnectionType { enum ProtocolType { PROTOCOL_UNKNOWN = 0; - PROTOCOL_RDMA_HANDSHAKE = 1; + PROTOCOL_TRANSPORT_HANDSHAKE = 1; PROTOCOL_BAIDU_STD = 2; PROTOCOL_STREAMING_RPC = 3; PROTOCOL_HULU_PBRPC = 4; diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp index 19d58ac422..571fa8d5d5 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -618,8 +618,8 @@ BUTIL_FORCE_INLINE bool is_http_protocol(const char* name) { return strcmp(name, "http") == 0 || strcmp(name, "h2") == 0; } -BUTIL_FORCE_INLINE bool is_rdma_handshake_protocol(const char* name) { - return strcmp(name, "rdma_handshake") == 0; +BUTIL_FORCE_INLINE bool is_transport_handshake_protocol(const char* name) { + return strcmp(name, "transport_handshake") == 0; } Acceptor* Server::BuildAcceptor() { @@ -640,12 +640,12 @@ Acceptor* Server::BuildAcceptor() { continue; } // Erase whatever the exemptions below say. http, h2 and - // rdma_handshake are always served, but they are still + // transport_handshake are always served, but they are still // valid names for the whitelist. bool in_whitelist = (whitelist.erase(protocols[i].name) != 0); if (has_whitelist && !in_whitelist && !is_http_protocol(protocols[i].name) && - !is_rdma_handshake_protocol(protocols[i].name)) { + !is_transport_handshake_protocol(protocols[i].name)) { // the protocol is not allowed to serve. RPC_VLOG << "Skip protocol=" << protocols[i].name; continue; diff --git a/src/brpc/server.h b/src/brpc/server.h index 31dc69af2f..117774d1d0 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -274,9 +274,9 @@ struct ServerOptions { // Only enable these protocols, separated by spaces. // All names inside must be valid, check protocols name in global.cpp - // http/h2 and rdma_handshake are served whatever this field says: + // http/h2 and transport_handshake are served whatever this field says: // the builtin services are only reachable over http/h2, and - // rdma_handshake is a transport level handshake dispatching no request. + // transport_handshake is a transport level handshake dispatching no request. // Naming them here is allowed and changes nothing. // Default: empty (all protocols) std::string enabled_protocols; diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index 8073208692..c5b195cb99 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -584,7 +584,7 @@ TEST_F(ServerTest, only_allow_protocols_in_enabled_protocols) { ASSERT_EQ(0, server.Join()); } -// http, h2 and rdma_handshake are served whatever enabled_protocols says, but +// http, h2 and transport_handshake are served whatever enabled_protocols says, but // naming them in it used to make Start() fail: the exemption was tested // before whitelist.erase() and `&&' short-circuited it, so the names survived // into the leftover check and came back as "unknown protocols=`http '". @@ -592,8 +592,8 @@ TEST_F(ServerTest, enabled_protocols_can_name_always_enabled_protocols) { std::string cases[] = { "baidu_std http", "baidu_std h2", - "baidu_std rdma_handshake", - "baidu_std http h2 rdma_handshake", + "baidu_std transport_handshake", + "baidu_std http h2 transport_handshake", }; for (size_t i = 0; i < arraysize(cases); ++i) { brpc::Server server; diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 8c5ad480e6..d021758224 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -27,10 +27,12 @@ #include "butil/fd_guard.h" #include "butil/sys_byteorder.h" #include "brpc/adapter_transport.h" +#include "brpc/global.h" #include "brpc/handshake/handshake_io.h" #include "brpc/handshake/handshake_adapter.h" #include "brpc/input_messenger.h" #include "brpc/policy/transport_handshake_protocol.h" +#include "brpc/protocol.h" #if BRPC_WITH_RDMA #include "brpc/rdma_transport.h" #endif @@ -275,6 +277,20 @@ class TestUpgradeMessenger : public InputMessenger { } }; +TEST(TransportHandshakeTest, registered_protocol_uses_transport_name) { + GlobalInitializeOrDie(); + EXPECT_EQ(1, PROTOCOL_TRANSPORT_HANDSHAKE); + EXPECT_EQ(PROTOCOL_TRANSPORT_HANDSHAKE, + StringToProtocolType("transport_handshake")); + EXPECT_STREQ("transport_handshake", + ProtocolTypeToString(PROTOCOL_TRANSPORT_HANDSHAKE)); + + const Protocol* protocol = FindProtocol(PROTOCOL_TRANSPORT_HANDSHAKE); + ASSERT_NE(nullptr, protocol); + EXPECT_EQ(policy::ParseTransportHandshake, protocol->parse); + EXPECT_EQ(policy::ProcessTransportHandshake, protocol->process_request); +} + TEST(TransportHandshakeTest, adapter_retains_context_for_fragmented_extension) { int fds[2]; ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); From 0b9a6fd93f063cf3444229ba743d27f0d6f5ba3c Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Fri, 23 Oct 2026 02:39:29 -0700 Subject: [PATCH 10/10] Apply handshake review style and test cleanup --- src/brpc/adapter_transport.cpp | 28 +- src/brpc/handshake/handshake_frame.cpp | 16 +- src/brpc/handshake/handshake_frame.h | 4 +- src/brpc/handshake/handshake_io.cpp | 16 +- src/brpc/handshake/handshake_io.h | 2 +- src/brpc/handshake/rdma_handshake.cpp | 6 +- src/brpc/handshake/rdma_handshake_constants.h | 2 +- src/brpc/handshake/ubshm_handshake.cpp | 20 +- src/brpc/transport_handshake.cpp | 40 +- src/brpc/transport_handshake.h | 6 +- test/brpc_rdma_unittest.cpp | 1075 ++++++----------- test/brpc_transport_handshake_unittest.cpp | 28 +- 12 files changed, 481 insertions(+), 762 deletions(-) diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index 93e6656d8e..49558efa07 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -228,17 +228,17 @@ AdapterTransport::AdapterTransport(SocketMode mode) AdapterTransport::~AdapterTransport() = default; AdapterTransport* AdapterTransport::Get(Socket* socket) { - CHECK(socket != NULL); + CHECK(socket != nullptr); return static_cast(socket->_transport.get()); } const AdapterTransport* AdapterTransport::Get(const Socket* socket) { - CHECK(socket != NULL); + CHECK(socket != nullptr); return static_cast(socket->_transport.get()); } bool AdapterTransport::upgrade_capable(SocketMode mode) const { - if (_mode != mode || _high_speed_transport == NULL) { + if (_mode != mode || _high_speed_transport == nullptr) { return false; } switch (mode) { @@ -279,11 +279,11 @@ int AdapterTransport::StartClientUpgrade(const Socket* socket, ParseResult AdapterTransport::ProcessUpgradeReadable(butil::IOBuf* source) { ParseResult result(PARSE_ERROR_NOT_ENOUGH_DATA); - if (_socket->parsing_context() != NULL) { + if (_socket->parsing_context() != nullptr) { handshake::ServerHandshakeContext* context = static_cast( _socket->parsing_context()); - CHECK(context->adapter() != NULL); + CHECK(context->adapter() != nullptr); result = context->adapter()->ExecuteServerHandshake(source, _socket); } else if (!source->empty()) { static const size_t MAX_MAGIC_LEN = 4; @@ -345,12 +345,12 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { adapter->FallbackToTcp(); adapter->CompleteConnection(handshake::FALLBACK_TCP); task->done(0, task->data); - return NULL; + return nullptr; } std::unique_ptr protocol = transport->CreateClientHandshakeAdapter(); - CHECK(protocol != NULL); + CHECK(protocol != nullptr); RdmaClientHandshakeTransport participant( transport, protocol.get(), socket, &connect_error); const handshake::StepResult result = adapter->_handshake.RunClient( @@ -372,7 +372,7 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { adapter->CompleteConnection(static_cast( adapter->_handshake.phase(butil::memory_order_acquire))); task->done(connect_error, task->data); - return NULL; + return nullptr; } #endif @@ -384,13 +384,13 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { adapter->FallbackToTcp(); adapter->CompleteConnection(handshake::FALLBACK_TCP); task->done(0, task->data); - return NULL; + return nullptr; } const size_t local_shm_len = static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; ubring::SHM local_trx_shm = { - NULL, local_shm_len, 0, {0}, static_cast(socket->fd())}; + nullptr, local_shm_len, 0, {0}, static_cast(socket->fd())}; const auto shm_name_str = butil::endpoint2str(socket->local_side()); ubring::UBShmHandshakeAdapter wire( @@ -410,14 +410,14 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { adapter->CompleteConnection(static_cast( adapter->_handshake.phase(butil::memory_order_acquire))); task->done(connect_error, task->data); - return NULL; + return nullptr; } #endif socket->SetFailed(EPROTO, "Unsupported client transport handshake"); adapter->CompleteConnection(handshake::FAILED); task->done(EPROTO, task->data); - return NULL; + return nullptr; } void AdapterTransport::Init(Socket* socket, const SocketOptions& options) { @@ -425,7 +425,7 @@ void AdapterTransport::Init(Socket* socket, const SocketOptions& options) { _socket = socket; _default_connect = options.app_connect; _on_edge_trigger = options.on_edge_triggered_events; - if (options.need_on_edge_trigger && _on_edge_trigger == NULL) { + if (options.need_on_edge_trigger && _on_edge_trigger == nullptr) { if (_mode == SOCKET_MODE_TCP) { _on_edge_trigger = InputMessenger::OnNewMessages; #if BRPC_WITH_RDMA @@ -493,7 +493,7 @@ std::shared_ptr AdapterTransport::Connect() { } SocketUser* const client_messenger = static_cast(get_client_side_messenger()); - if (client_messenger != NULL && + if (client_messenger != nullptr && _socket->user() == client_messenger) { FallbackToTcp(); return AdapterConnect::Unwrap(_tcp_transport->Connect()); diff --git a/src/brpc/handshake/handshake_frame.cpp b/src/brpc/handshake/handshake_frame.cpp index f25e50c15c..d185f8110e 100644 --- a/src/brpc/handshake/handshake_frame.cpp +++ b/src/brpc/handshake/handshake_frame.cpp @@ -76,7 +76,7 @@ FrameResult FrameCodec::DecodeLength(const FrameSpec& spec, FrameResult FrameCodec::Encode(const FrameSpec& spec, const std::string& payload, std::string* frame) { - if (frame == NULL || (spec.magic_len != 0 && spec.magic == NULL)) { + if (frame == nullptr || (spec.magic_len != 0 && spec.magic == nullptr)) { return FRAME_PROTOCOL_ERROR; } const size_t length_size = LengthFieldSize(spec); @@ -124,8 +124,8 @@ FrameResult FrameCodec::Encode(const FrameSpec& spec, FrameResult FrameCodec::ReadFrame(HandshakeIO* io, const FrameSpec& spec, bool push_back_on_not_mine, std::string* payload) { - if (io == NULL || payload == NULL || - (spec.magic_len != 0 && spec.magic == NULL)) { + if (io == nullptr || payload == nullptr || + (spec.magic_len != 0 && spec.magic == nullptr)) { return FRAME_PROTOCOL_ERROR; } @@ -165,11 +165,11 @@ FrameResult FrameCodec::ParseBufferedFrame(HandshakeInput* input, const FrameSpec& spec, std::string* payload, bool* magic_matched) { - if (magic_matched != NULL) { + if (magic_matched != nullptr) { *magic_matched = false; } - if (input == NULL || payload == NULL || - (spec.magic_len != 0 && spec.magic == NULL)) { + if (input == nullptr || payload == nullptr || + (spec.magic_len != 0 && spec.magic == nullptr)) { return FRAME_PROTOCOL_ERROR; } const size_t header_len = spec.magic_len + LengthFieldSize(spec); @@ -185,7 +185,7 @@ FrameResult FrameCodec::ParseBufferedFrame(HandshakeInput* input, memcmp(header.data(), spec.magic, spec.magic_len) != 0) { return FRAME_NOT_MINE; } - if (magic_matched != NULL) { + if (magic_matched != nullptr) { *magic_matched = true; } if (input->Size() < header_len || @@ -214,7 +214,7 @@ FrameResult FrameCodec::ParseBufferedFrame(HandshakeInput* input, FrameResult FrameCodec::WriteFrame(HandshakeIO* io, const FrameSpec& spec, const std::string& payload) { - if (io == NULL) { + if (io == nullptr) { return FRAME_PROTOCOL_ERROR; } std::string frame; diff --git a/src/brpc/handshake/handshake_frame.h b/src/brpc/handshake/handshake_frame.h index dbe399f738..fb99494fee 100644 --- a/src/brpc/handshake/handshake_frame.h +++ b/src/brpc/handshake/handshake_frame.h @@ -42,7 +42,7 @@ struct FrameSpec { }; FrameSpec() - : magic(NULL), magic_len(0), min_frame_len(0), max_frame_len(0), + : magic(nullptr), magic_len(0), min_frame_len(0), max_frame_len(0), length_encoding(FIXED) {} FrameSpec(const char* magic_in, size_t magic_len_in, @@ -74,7 +74,7 @@ class FrameCodec { static FrameResult ParseBufferedFrame(HandshakeInput* input, const FrameSpec& spec, std::string* payload, - bool* magic_matched = NULL); + bool* magic_matched = nullptr); static FrameResult WriteFrame(HandshakeIO* io, const FrameSpec& spec, const std::string& payload); static FrameResult DrainFrame(HandshakeIO* io, const FrameSpec& spec); diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp index 67cc95e37c..86ef81026d 100644 --- a/src/brpc/handshake/handshake_io.cpp +++ b/src/brpc/handshake/handshake_io.cpp @@ -30,15 +30,15 @@ namespace brpc { namespace handshake { size_t IOBufHandshakeInput::Size() const { - return _source != NULL ? _source->size() : 0; + return _source != nullptr ? _source->size() : 0; } bool IOBufHandshakeInput::CopyTo(void* data, size_t len) const { - return _source != NULL && _source->copy_to(data, len) == len; + return _source != nullptr && _source->copy_to(data, len) == len; } bool IOBufHandshakeInput::Consume(size_t len) { - return _source != NULL && _source->pop_front(len) == len; + return _source != nullptr && _source->pop_front(len) == len; } static const int WAIT_TIMEOUT_MS = 50; @@ -97,8 +97,8 @@ static int ReadExactLoop(butil::atomic* read_butex, } int SocketHandshakeIO::ReadExact(void* data, size_t len) { - CHECK(data != NULL); - CHECK(_socket != NULL); + CHECK(data != nullptr); + CHECK(_socket != nullptr); const int fd = _socket->fd(); return ReadExactLoop(_read_butex, _socket->id(), len, [data, fd](size_t offset, size_t remaining) { @@ -136,8 +136,8 @@ static int WriteAllLoop(size_t len, WriteOnce write_once, } int SocketHandshakeIO::WriteAll(const void* data, size_t len) { - CHECK(data != NULL); - CHECK(_socket != NULL); + CHECK(data != nullptr); + CHECK(_socket != nullptr); const int fd = _socket->fd(); return WriteAllLoop(len, [data, fd](size_t offset, size_t remaining) { @@ -151,7 +151,7 @@ int SocketHandshakeIO::WriteAll(const void* data, size_t len) { } int SocketHandshakeIO::PushBack(const void* data, size_t len) { - CHECK(_socket != NULL); + CHECK(_socket != nullptr); if (len != 0) { return _socket->fd_input_processor().read_buf().append(data, len) == 0 ? 0 : -1; } diff --git a/src/brpc/handshake/handshake_io.h b/src/brpc/handshake/handshake_io.h index 82f83ab5cd..28340ab437 100644 --- a/src/brpc/handshake/handshake_io.h +++ b/src/brpc/handshake/handshake_io.h @@ -66,7 +66,7 @@ class IOBufHandshakeInput : public HandshakeInput { class SocketHandshakeIO : public HandshakeIO { public: - explicit SocketHandshakeIO(Socket* socket = NULL); + explicit SocketHandshakeIO(Socket* socket = nullptr); ~SocketHandshakeIO() override; void Reset(Socket* socket); diff --git a/src/brpc/handshake/rdma_handshake.cpp b/src/brpc/handshake/rdma_handshake.cpp index 26c622a1ac..dcde3f47ab 100644 --- a/src/brpc/handshake/rdma_handshake.cpp +++ b/src/brpc/handshake/rdma_handshake.cpp @@ -473,7 +473,7 @@ class RdmaServerHandshakeTransport : public HandshakeTransport { RdmaServerHandshakeTransport( RdmaTransport* transport, Socket* socket, butil::IOBuf* source) : _transport(transport), _socket(socket), _source(source), - _protocol(NULL) {} + _protocol(nullptr) {} void OnProtocolSelected(HandshakeProtocol* protocol) override { _protocol = static_cast(protocol); @@ -490,7 +490,7 @@ class RdmaServerHandshakeTransport : public HandshakeTransport { } StepResult NegotiateResources() override { - CHECK(_protocol != NULL); + CHECK(_protocol != nullptr); if (_transport->NegotiateUpgradeResources( _protocol->remote(), true) == 0) { return STEP_OK; @@ -519,7 +519,7 @@ class RdmaServerHandshakeTransport : public HandshakeTransport { StepResult RdmaServerHandshakeAdapter::RunRdmaServerHandshake( butil::IOBuf* source, Socket* socket) { RdmaTransport* transport = RdmaTransport::Get(socket); - CHECK(transport->GetRdmaEp() != NULL); + CHECK(transport->GetRdmaEp() != nullptr); std::vector > protocols = transport->CreateServerHandshakeAdapters(); diff --git a/src/brpc/handshake/rdma_handshake_constants.h b/src/brpc/handshake/rdma_handshake_constants.h index af60d2e9d3..f67a93ab9d 100644 --- a/src/brpc/handshake/rdma_handshake_constants.h +++ b/src/brpc/handshake/rdma_handshake_constants.h @@ -70,7 +70,7 @@ inline const handshake::FrameSpec& RdmaHelloFrameSpec(int version) { inline const handshake::FrameSpec& RdmaAckFrameSpec() { static const handshake::FrameSpec spec( - NULL, 0, HELLO_ACK_LEN, HELLO_ACK_LEN, + nullptr, 0, HELLO_ACK_LEN, HELLO_ACK_LEN, handshake::FrameSpec::FIXED); return spec; } diff --git a/src/brpc/handshake/ubshm_handshake.cpp b/src/brpc/handshake/ubshm_handshake.cpp index 28c7f84e62..e2b6af1d33 100644 --- a/src/brpc/handshake/ubshm_handshake.cpp +++ b/src/brpc/handshake/ubshm_handshake.cpp @@ -61,7 +61,7 @@ static const FrameSpec& HelloFrameSpec() { static const FrameSpec& AckFrameSpec() { static const FrameSpec spec( - NULL, 0, ACK_LEN, ACK_LEN, FrameSpec::FIXED); + nullptr, 0, ACK_LEN, ACK_LEN, FrameSpec::FIXED); return spec; } @@ -181,7 +181,7 @@ const handshake::FrameSpec& UBShmHandshakeAdapter::AckFrameSpec() const { const handshake::FrameSpec& UBShmHandshakeAdapter::ExtensionFrameSpec() const { static const handshake::FrameSpec spec( - NULL, 0, HelloFormatExtension::WIRE_SIZE, + nullptr, 0, HelloFormatExtension::WIRE_SIZE, HelloFormatExtension::WIRE_SIZE, handshake::FrameSpec::FIXED); return spec; } @@ -189,13 +189,13 @@ UBShmHandshakeAdapter::ExtensionFrameSpec() const { void UBShmHandshakeAdapter::ConfigureClientHello( uint64_t len, const char* shm_name) { _local_len = len; - _local_name = shm_name != NULL ? shm_name : ""; + _local_name = shm_name != nullptr ? shm_name : ""; _server_reply = false; } handshake::StepResult UBShmHandshakeAdapter::BuildHello( bool enabled, std::string* payload) { - const char* name = NULL; + const char* name = nullptr; if (enabled) { name = _server_reply ? _remote.shm_name : _local_name.c_str(); } @@ -242,7 +242,7 @@ handshake::StepResult UBShmHandshakeAdapter::BuildHello( message.hello_ver = handshake::ubshm_wire::HELLO_VERSION; message.impl_ver = handshake::ubshm_wire::IMPL_VERSION; message.len = len; - if (shm_name == NULL) { + if (shm_name == nullptr) { errno = EINVAL; return handshake::STEP_ERROR; } @@ -277,7 +277,7 @@ handshake::StepResult UBShmHandshakeAdapter::ParseHello( // the original wire bytes so that truncation cannot hide a malformed name. const char* wire_name = payload.data() + payload.size() - SHM_MAX_NAME_BUFF_LEN; - if (memchr(wire_name, '\0', SHM_MAX_NAME_BUFF_LEN) == NULL) { + if (memchr(wire_name, '\0', SHM_MAX_NAME_BUFF_LEN) == nullptr) { errno = EPROTO; return handshake::STEP_ERROR; } @@ -386,20 +386,20 @@ class UBShmServerHandshakeTransport : public HandshakeTransport { strnlen(remote.shm_name, SHM_MAX_NAME_BUFF_LEN); ubring::SHM remote_trx_shm = { - NULL, remote.len, 0, {0}, + nullptr, remote.len, 0, {0}, static_cast(_socket->fd())}; memcpy(remote_trx_shm.name, remote.shm_name, remote_name_len + 1); const size_t local_shm_len = static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; ubring::SHM local_trx_shm = { - NULL, local_shm_len, 0, {0}, + nullptr, local_shm_len, 0, {0}, static_cast(_socket->fd())}; char client_name[SHM_MAX_NAME_BUFF_LEN + 1]; memcpy(client_name, remote.shm_name, SHM_MAX_NAME_BUFF_LEN); client_name[SHM_MAX_NAME_BUFF_LEN] = '\0'; char* client_ip_port = strrchr(client_name, '_'); - if (client_ip_port != NULL) { + if (client_ip_port != nullptr) { *client_ip_port = '\0'; } const int result = snprintf( @@ -444,7 +444,7 @@ class UBShmServerHandshakeTransport : public HandshakeTransport { StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( butil::IOBuf* source, Socket* socket) { UBShmTransport* transport = UBShmTransport::Get(socket); - CHECK(transport->GetUBShmEp() != NULL); + CHECK(transport->GetUBShmEp() != nullptr); ubring::UBShmHandshakeAdapter wire( ubring::PreferredDataFormatForShmType(ubring::GetShmType())); diff --git a/src/brpc/transport_handshake.cpp b/src/brpc/transport_handshake.cpp index e7fd9580c5..ab07971c2b 100644 --- a/src/brpc/transport_handshake.cpp +++ b/src/brpc/transport_handshake.cpp @@ -35,7 +35,7 @@ const uint32_t COMMON_ACK_OK = 0x1; } // namespace StepResult HandshakeProtocol::BuildAck(bool enabled, std::string* payload) { - CHECK(payload != NULL); + CHECK(payload != nullptr); const uint32_t flags = butil::HostToNet32(enabled ? COMMON_ACK_OK : 0); payload->assign(reinterpret_cast(&flags), sizeof(flags)); return STEP_OK; @@ -43,7 +43,7 @@ StepResult HandshakeProtocol::BuildAck(bool enabled, std::string* payload) { StepResult HandshakeProtocol::ParseAck(const std::string& payload, bool* enabled) { - CHECK(enabled != NULL); + CHECK(enabled != nullptr); if (payload.size() != COMMON_ACK_SIZE) { errno = EPROTO; return STEP_ERROR; @@ -56,7 +56,7 @@ StepResult HandshakeProtocol::ParseAck(const std::string& payload, const FrameSpec& HandshakeProtocol::ExtensionFrameSpec() const { static const FrameSpec empty_spec( - NULL, 0, 0, 0, FrameSpec::FIXED); + nullptr, 0, 0, 0, FrameSpec::FIXED); return empty_spec; } @@ -64,14 +64,14 @@ ServerHandshakeContext* ServerHandshakeContext::Create( HandshakeAdapter* adapter) { ServerHandshakeContext* context = butil::get_object(); - if (context != NULL) { + if (context != nullptr) { context->_adapter = adapter; } return context; } void ServerHandshakeContext::Destroy() { - _adapter = NULL; + _adapter = nullptr; butil::return_object(this); } @@ -118,7 +118,7 @@ StepResult HandshakeSession::ReceiveHello(HandshakeProtocol* protocol, bool push_back_on_not_mine, bool* magic_matched) { std::string payload; - const FrameResult frame_result = input != NULL + const FrameResult frame_result = input != nullptr ? FrameCodec::ParseBufferedFrame( input, protocol->HelloFrameSpec(), &payload, magic_matched) : FrameCodec::ReadFrame( @@ -147,7 +147,7 @@ StepResult HandshakeSession::ReceiveAck(HandshakeProtocol* protocol, HandshakeInput* input, bool* enabled) { std::string payload; - const FrameResult frame_result = input != NULL + const FrameResult frame_result = input != nullptr ? FrameCodec::ParseBufferedFrame( input, protocol->AckFrameSpec(), &payload) : FrameCodec::ReadFrame( @@ -174,7 +174,7 @@ StepResult HandshakeSession::SendExtension(HandshakeProtocol* protocol, StepResult HandshakeSession::ReceiveExtension(HandshakeProtocol* protocol, HandshakeInput* input) { std::string payload; - const FrameResult frame_result = input != NULL + const FrameResult frame_result = input != nullptr ? FrameCodec::ParseBufferedFrame( input, protocol->ExtensionFrameSpec(), &payload) : FrameCodec::ReadFrame( @@ -187,14 +187,14 @@ StepResult HandshakeSession::SelectAndReceiveHello( const std::vector& protocols, HandshakeInput* input, bool push_back_on_not_mine, HandshakeProtocol** selected) { CHECK(!protocols.empty()); - CHECK(selected != NULL); - if (input == NULL) { + CHECK(selected != nullptr); + if (input == nullptr) { // A blocking byte stream cannot try a second codec after consuming // bytes from the fd. Such protocols must select a single codec before // entering the common session. CHECK_EQ(1UL, protocols.size()); *selected = protocols.front(); - return ReceiveHello(*selected, NULL, push_back_on_not_mine); + return ReceiveHello(*selected, nullptr, push_back_on_not_mine); } bool need_more = false; @@ -222,8 +222,8 @@ StepResult HandshakeSession::SelectAndReceiveHello( StepResult HandshakeSession::RunClient(HandshakeProtocol* protocol, HandshakeTransport* transport) { - CHECK(protocol != NULL); - CHECK(transport != NULL); + CHECK(protocol != nullptr); + CHECK(transport != nullptr); transport->OnProtocolSelected(protocol); // A client handshake runs once on a potentially reused bthread. Do not // let an errno left by earlier work override this handshake's result. @@ -244,7 +244,7 @@ StepResult HandshakeSession::RunClient(HandshakeProtocol* protocol, } SetPhase(HELLO_WAIT); - result = ReceiveHello(protocol, NULL, false); + result = ReceiveHello(protocol, nullptr, false); if (result == STEP_NOT_MINE || result == STEP_NEED_MORE) { errno = EPROTO; } @@ -260,7 +260,7 @@ StepResult HandshakeSession::RunClient(HandshakeProtocol* protocol, return FinishWithFailure(this, transport); } SetPhase(EXTENSION_WAIT); - result = ReceiveExtension(protocol, NULL); + result = ReceiveExtension(protocol, nullptr); if (result != STEP_OK && result != STEP_FALLBACK) { return FinishWithFailure(this, transport); } @@ -293,7 +293,7 @@ StepResult HandshakeSession::RunServer( const std::vector& protocols, HandshakeInput* input, HandshakeTransport* transport, bool fallback_on_not_mine) { CHECK(!protocols.empty()); - CHECK(transport != NULL); + CHECK(transport != nullptr); // Once TCP fallback has been published, subsequent bytes are application // protocol data and must bypass every upgrade codec without changing the @@ -302,7 +302,7 @@ StepResult HandshakeSession::RunServer( return STEP_NOT_MINE; } - HandshakeProtocol* selected = NULL; + HandshakeProtocol* selected = nullptr; if (phase(butil::memory_order_acquire) != ACK_WAIT && phase(butil::memory_order_acquire) != EXTENSION_WAIT) { const int previous_phase = phase(butil::memory_order_acquire); @@ -318,7 +318,7 @@ StepResult HandshakeSession::RunServer( return STEP_NOT_MINE; } if (result == STEP_NEED_MORE) { - if (selected == NULL) { + if (selected == nullptr) { SetPhase(previous_phase); } return STEP_NEED_MORE; @@ -326,7 +326,7 @@ StepResult HandshakeSession::RunServer( if (result == STEP_ERROR) { return FinishWithFailure(this, transport); } - CHECK(selected != NULL); + CHECK(selected != nullptr); transport->OnProtocolSelected(selected); if (result == STEP_FALLBACK) { // Publish the disabled transport state immediately. The server @@ -369,7 +369,7 @@ StepResult HandshakeSession::RunServer( break; } } - CHECK(selected != NULL); + CHECK(selected != nullptr); } if (phase(butil::memory_order_acquire) == EXTENSION_WAIT) { diff --git a/src/brpc/transport_handshake.h b/src/brpc/transport_handshake.h index 3228448449..52281554ee 100644 --- a/src/brpc/transport_handshake.h +++ b/src/brpc/transport_handshake.h @@ -39,7 +39,7 @@ class HandshakeAdapter; // Remembering the selected stateless adapter is necessary because ACK frames // have no magic and cannot be dispatched from their bytes alone. struct ServerHandshakeContext : public Destroyable { - ServerHandshakeContext() : _adapter(NULL) {} + ServerHandshakeContext() : _adapter(nullptr) {} static ServerHandshakeContext* Create(HandshakeAdapter* adapter); HandshakeAdapter* adapter() const { return _adapter; } void Destroy() override; @@ -147,7 +147,7 @@ class FallbackHandshakeProtocol : public HandshakeProtocol { // lifecycle and publication ordering. class HandshakeSession { public: - explicit HandshakeSession(Socket* socket = NULL) + explicit HandshakeSession(Socket* socket = nullptr) : _socket_io(socket), _io(&_socket_io), _phase(UNINITIALIZED), _protocol_version(0), _local_enabled(false) {} @@ -206,7 +206,7 @@ class HandshakeSession { StepResult ReceiveHello(HandshakeProtocol* protocol, HandshakeInput* input, bool push_back_on_not_mine, - bool* magic_matched = NULL); + bool* magic_matched = nullptr); StepResult SendAck(HandshakeProtocol* protocol, bool enabled); StepResult ReceiveAck(HandshakeProtocol* protocol, HandshakeInput* input, bool* enabled); diff --git a/test/brpc_rdma_unittest.cpp b/test/brpc_rdma_unittest.cpp index ed0ae7262e..2f36494c2d 100644 --- a/test/brpc_rdma_unittest.cpp +++ b/test/brpc_rdma_unittest.cpp @@ -376,18 +376,11 @@ TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { TEST_F(RdmaTest, client_close_before_hello_send) { 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); ASSERT_NE(nullptr, s); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -397,22 +390,15 @@ TEST_F(RdmaTest, client_close_before_hello_send) { TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); ASSERT_NE(nullptr, s); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(data, "PRPC", 4); // send as normal baidu_std protocol - ASSERT_EQ(4, write(sockfd, data, 4)); + ASSERT_TRUE(WriteAll(sockfd, data, 4)); StopServer(); } @@ -420,54 +406,44 @@ TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { TEST_F(RdmaTest, client_close_during_hello_send) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[8]; - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); s = WaitForServerSocket(); ASSERT_NE(nullptr, s); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RD", 2); - ASSERT_EQ(2, write(sockfd1, data, 2)); // break in magic str + ASSERT_TRUE(WriteAll(sockfd1, data, 2)); // break in magic str // Close and wait for teardown. Parser phase assertions for incomplete // magic belong to the synchronous common-handshake regression test. sockfd1.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); s = WaitForServerSocket(); ASSERT_NE(nullptr, s); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // break after magic str + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // break after magic str ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); - butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd3 >= 0); - ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); + butil::fd_guard sockfd3; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd3)); s = WaitForServerSocket(); ASSERT_NE(nullptr, s); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Send the 4B magic plus a valid msg_len (=40) but no body, so the server // recognizes an RDMA v2 hello and waits for the remaining bytes. (A zero // msg_len would now be rejected up-front as a protocol error.) memcpy(data, "RDMA", 4); uint16_t v2_len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); memcpy(data + 4, &v2_len, sizeof(v2_len)); - ASSERT_EQ(6, write(sockfd3, data, 6)); // magic + msg_len, body missing + ASSERT_TRUE(WriteAll(sockfd3, data, 6)); // magic + msg_len, body missing ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd3.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -478,46 +454,34 @@ TEST_F(RdmaTest, client_close_during_hello_send) { TEST_F(RdmaTest, client_hello_msg_invalid_len) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); memset(data + 4, 0, 36); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); // Write invalid length. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); // Write invalid length. + ASSERT_TRUE(WaitForServerSocketGone()); + + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); uint16_t len = butil::HostToNet16(35); memcpy(data + 4, &len, sizeof(len)); memset(data + 6, 0, 34); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); // write invalid length - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); // write invalid length + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -525,26 +489,19 @@ TEST_F(RdmaTest, client_hello_msg_invalid_len) { TEST_F(RdmaTest, client_hello_msg_invalid_version) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); uint16_t ver = butil::HostToNet16(1); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); memcpy(data + 4, &len, 2); memset(data + 6, 0, 34); memcpy(data + 6, &ver, 2); // hello_ver == 1, impl_ver == 0 @@ -555,46 +512,35 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { // hello_ver). Now that Step 1 enforces a HELLO_V2_MSG_LEN_MAX upper bound, // such an oversized msg_len would be rejected before reaching the // version check, breaking the intent of this UT. - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); uint32_t flags = 0; - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd1.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WaitForServerSocketGone()); + + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); memcpy(data + 8, &ver, 2); // hello_ver == 0, impl_ver == 1 // See comment above on `write(sockfd1, data + 4, 36)` for why we // write from data + 4 instead of data. - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd2.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -602,10 +548,6 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; uint32_t flags = butil::HostToNet32(0); rdma::v2_wire::HelloMessage msg{}; @@ -619,81 +561,60 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { msg.block_size = 8192; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd1.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); msg.sq_size = 16; msg.rq_size = 10; msg.block_size = 8192; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd2.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); msg.sq_size = 16; msg.rq_size = 16; msg.block_size = 1000; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd3 >= 0); - ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd3, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd3, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd3; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd3)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd3, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd3, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd3, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd3, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd3.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -701,10 +622,6 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { TEST_F(RdmaTest, client_close_after_qp_build) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; @@ -719,19 +636,15 @@ TEST_F(RdmaTest, client_close_after_qp_build) { memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(40, write(sockfd1, data, 40)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd1, data, 40)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + sockfd1.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -739,10 +652,6 @@ TEST_F(RdmaTest, client_close_after_qp_build) { TEST_F(RdmaTest, client_close_during_ack_send) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; @@ -757,27 +666,20 @@ TEST_F(RdmaTest, client_close_during_ack_send) { memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); uint32_t flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); + sockfd1.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -785,10 +687,6 @@ TEST_F(RdmaTest, client_close_during_ack_send) { TEST_F(RdmaTest, client_close_after_ack_send) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; @@ -803,50 +701,36 @@ TEST_F(RdmaTest, client_close_after_ack_send) { memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + sockfd1.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); + + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd2); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); + sockfd2.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -854,10 +738,6 @@ TEST_F(RdmaTest, client_close_after_ack_send) { TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { StartServer(); - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); Socket* s = nullptr; rdma::v2_wire::HelloMessage msg{}; uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; @@ -872,49 +752,35 @@ TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd1; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd1)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd1, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); + ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); + ASSERT_TRUE(WaitForServerSocketGone()); + + butil::fd_guard sockfd2; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd2)); + s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); + ASSERT_TRUE(WriteAll(sockfd2, data, 4)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); + ASSERT_TRUE(WriteAll(sockfd2, data + 4, 36)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); + ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -1109,7 +975,7 @@ TEST_F(RdmaTest, server_parses_qp_stream_after_rdma_is_on) { 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(handshake::ESTABLISHED, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); ASSERT_FALSE(s->Failed()); StopServer(); @@ -1187,10 +1053,11 @@ TEST_F(RdmaTest, server_miss_before_hello_send) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); @@ -1218,19 +1085,18 @@ TEST_F(RdmaTest, server_close_before_hello_send) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + acc_fd.reset(-1); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1255,18 +1121,17 @@ TEST_F(RdmaTest, server_miss_during_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(2, write(acc_fd, "RD", 2)); - usleep(100000); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, "RD", 2)); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1291,21 +1156,19 @@ TEST_F(RdmaTest, server_close_during_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(2, write(acc_fd, "RD", 2)); - usleep(100000); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, "RD", 2)); + acc_fd.reset(-1); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1330,19 +1193,18 @@ TEST_F(RdmaTest, server_hello_invalid_magic_str) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "ABCD", 4)); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, "ABCD", 4)); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); @@ -1367,21 +1229,20 @@ TEST_F(RdmaTest, server_miss_during_hello_msg) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, "RDMA", 4)); const uint16_t msg_len = butil::HostToNet16( static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(static_cast(sizeof(msg_len)), - write(acc_fd, &msg_len, sizeof(msg_len))); + ASSERT_TRUE(WriteAll(acc_fd, &msg_len, sizeof(msg_len))); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1406,24 +1267,22 @@ TEST_F(RdmaTest, server_close_during_hello_msg) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, "RDMA", 4)); const uint16_t msg_len = butil::HostToNet16( static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(static_cast(sizeof(msg_len)), - write(acc_fd, &msg_len, sizeof(msg_len))); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WriteAll(acc_fd, &msg_len, sizeof(msg_len))); + acc_fd.reset(-1); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1448,24 +1307,22 @@ TEST_F(RdmaTest, server_hello_invalid_msg_len) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); memcpy(data, "RDMA", 4); uint16_t len = butil::HostToNet16(35); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); @@ -1490,26 +1347,23 @@ TEST_F(RdmaTest, server_hello_invalid_version) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); memcpy(data, "RDMA", 4); uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); memcpy(data + 4, &len, 2); memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s.get()); + ASSERT_TRUE(ReadAll(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(0, butil::NetToHost32(*tmp)); bthread_id_join(cntl.call_id()); @@ -1536,16 +1390,16 @@ TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1558,13 +1412,10 @@ TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s.get()); + ASSERT_TRUE(ReadAll(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(0, butil::NetToHost32(*tmp)); bthread_id_join(cntl.call_id()); @@ -1591,16 +1442,16 @@ TEST_F(RdmaTest, server_miss_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1613,13 +1464,10 @@ TEST_F(RdmaTest, server_miss_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s.get()); + ASSERT_TRUE(ReadAll(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(1, butil::NetToHost32(*tmp)); bthread_id_join(cntl.call_id()); @@ -1646,16 +1494,16 @@ TEST_F(RdmaTest, server_close_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1668,16 +1516,13 @@ TEST_F(RdmaTest, server_close_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s.get()); + ASSERT_TRUE(ReadAll(acc_fd, data, 4)); uint32_t* tmp = (uint32_t*)data; ASSERT_EQ(1, butil::NetToHost32(*tmp)); - close(acc_fd); + acc_fd.reset(-1); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1702,16 +1547,16 @@ TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s.get()); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); rdma::v2_wire::HelloMessage msg{}; msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; @@ -1724,14 +1569,10 @@ TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { msg.gid = rdma::GetRdmaGid(); memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s.get()); + ASSERT_TRUE(WriteAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); bthread_id_join(cntl.call_id()); ASSERT_EQ(EPROTO, cntl.ErrorCode()); @@ -1756,16 +1597,16 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); // [0..4) magic ASSERT_EQ(0, memcmp(data, "RDMA", 4)); @@ -1791,18 +1632,11 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Send a well-formed v2 hello so the server enters the common ACK_WAIT. rdma::v2_wire::HelloMessage msg{}; @@ -1818,15 +1652,12 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); // Read server's reply hello and assert its byte-level layout. uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(ReadAll(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_EQ(0, memcmp(reply, "RDMA", 4)); ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, @@ -1845,14 +1676,11 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends // cleanly without requiring real RDMA hardware. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -1860,18 +1688,12 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { 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); - Socket* s = GetSocketFromServer(0); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); ASSERT_TRUE(s != NULL); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Build a v2 hello with msg_len = 48 (40 base + 8B zero tail). rdma::v2_wire::HelloMessage msg{}; @@ -1892,15 +1714,12 @@ TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { // msg_len bytes and preserve the ACK for the next handshake step. uint32_t flags = butil::HostToNet32(1); memcpy(buf + 48, &flags, sizeof(flags)); - ASSERT_EQ(sizeof(buf), write(sockfd, buf, sizeof(buf))); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, buf, sizeof(buf))); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -1908,18 +1727,11 @@ TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Build a v2 hello with msg_len = 4097 (HELLO_V2_MSG_LEN_MAX + 1). // We only send the 40B base; the server must reject before reading @@ -1937,14 +1749,10 @@ TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { uint8_t buf[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(buf, "RDMA", 4); msg.Serialize(buf + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_TRUE(WaitForServerSocketGone()); sockfd.reset(-1); - usleep(100000); StopServer(); } @@ -2024,12 +1832,12 @@ TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { // [0..4) magic "RDM3" uint8_t magic[4]; - ASSERT_EQ(4, read(acc_fd, magic, 4)); + ASSERT_TRUE(ReadAll(acc_fd, magic, 4)); ASSERT_EQ(0, memcmp(magic, "RDM3", 4)); // [4..8) pb_size, big-endian uint32, must be in (0, 4096] uint8_t size_buf[4]; - ASSERT_EQ(4, read(acc_fd, size_buf, 4)); + ASSERT_TRUE(ReadAll(acc_fd, size_buf, 4)); uint32_t pb_size = butil::NetToHost32(*reinterpret_cast(size_buf)); ASSERT_GT(pb_size, 0u); @@ -2037,7 +1845,7 @@ TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { // [8..8+pb_size) RdmaHello protobuf body. std::string body(pb_size, '\0'); - ASSERT_EQ((ssize_t)pb_size, read(acc_fd, &body[0], pb_size)); + ASSERT_TRUE(ReadAll(acc_fd, &body[0], pb_size)); rdma::RdmaHello msg; ASSERT_TRUE(msg.ParseFromString(body)); @@ -2059,40 +1867,31 @@ TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Send a valid v3 hello. std::string packet = MakeV3Packet(MakeValidV3Hello()); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); // Read server's reply hello: 4B magic + 4B pb_size + body. uint8_t reply_magic[4]; - ASSERT_EQ(4, read(sockfd, reply_magic, 4)); + ASSERT_TRUE(ReadAll(sockfd, reply_magic, 4)); ASSERT_EQ(0, memcmp(reply_magic, "RDM3", 4)); uint8_t size_buf[4]; - ASSERT_EQ(4, read(sockfd, size_buf, 4)); + ASSERT_TRUE(ReadAll(sockfd, size_buf, 4)); uint32_t pb_size = butil::NetToHost32(*reinterpret_cast(size_buf)); ASSERT_GT(pb_size, 0u); ASSERT_LE(pb_size, 4096u); std::string body(pb_size, '\0'); - ASSERT_EQ((ssize_t)pb_size, read(sockfd, &body[0], pb_size)); + ASSERT_TRUE(ReadAll(sockfd, &body[0], pb_size)); rdma::RdmaHello reply; ASSERT_TRUE(reply.ParseFromString(body)); ASSERT_TRUE(reply.has_block_size()); @@ -2104,14 +1903,11 @@ TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends // cleanly without requiring real RDMA hardware. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2119,41 +1915,28 @@ TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { TEST_F(RdmaTest, v3_server_rejects_zero_pb_size) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); // "RDM3" + pb_size = 0 (4B big-endian zero). uint8_t buf[8] = {'R', 'D', 'M', '3', 0, 0, 0, 0}; - ASSERT_EQ(8, write(sockfd, buf, 8)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd, buf, 8)); + ASSERT_TRUE(WaitForServerSocketGone()); sockfd.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } TEST_F(RdmaTest, v3_server_rejects_oversized_pb_size) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); uint8_t buf[8]; memcpy(buf, "RDM3", 4); @@ -2161,28 +1944,21 @@ TEST_F(RdmaTest, v3_server_rejects_oversized_pb_size) { uint32_t pb_size_be = butil::HostToNet32(static_cast(rdma::HELLO_V3_MAX_PB_SIZE + 1)); memcpy(buf + 4, &pb_size_be, 4); - ASSERT_EQ(8, write(sockfd, buf, 8)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd, buf, 8)); + ASSERT_TRUE(WaitForServerSocketGone()); sockfd.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } TEST_F(RdmaTest, v3_server_rejects_invalid_pb_bytes) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); // "RDM3" + pb_size = 8 + 8 bytes of 0xff (invalid protobuf body). uint8_t buf[16]; @@ -2190,62 +1966,49 @@ TEST_F(RdmaTest, v3_server_rejects_invalid_pb_bytes) { uint32_t pb_size_be = butil::HostToNet32(8); memcpy(buf + 4, &pb_size_be, 4); memset(buf + 8, 0xff, 8); - ASSERT_EQ(16, write(sockfd, buf, 16)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WriteAll(sockfd, buf, 16)); + ASSERT_TRUE(WaitForServerSocketGone()); sockfd.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } TEST_F(RdmaTest, v3_server_invalid_sq_size_falls_back) { 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3Hello(); msg.set_sq_size(0); // invalid: < MIN_QP_SIZE (16) std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); // Server validated the hello as invalid -> _rdma_state = RDMA_OFF, // but still proceeds to the common ACK_WAIT (sends its own reply hello). - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); // Drain server's reply hello (content not asserted here; covered // by v3_server_hello_bytes_baseline). uint8_t reply_hdr[8]; - ASSERT_EQ(8, read(sockfd, reply_hdr, 8)); + ASSERT_TRUE(ReadAll(sockfd, reply_hdr, 8)); ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); uint32_t reply_pb_size = butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); std::string reply_body(reply_pb_size, '\0'); - ASSERT_EQ((ssize_t)reply_pb_size, - read(sockfd, &reply_body[0], reply_pb_size)); + ASSERT_TRUE(ReadAll(sockfd, &reply_body[0], reply_pb_size)); // Client ACK flags=0 -> server settles into FALLBACK_TCP. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2277,14 +2040,14 @@ rdma::RdmaHello MakeValidV3HelloWithEce(uint32_t vendor_id, uint32_t options, // it into `reply`. Asserts the framing along the way. static void ReadServerV3Reply(int fd, rdma::RdmaHello* reply) { uint8_t reply_hdr[8]; - ASSERT_EQ(8, read(fd, reply_hdr, 8)); + ASSERT_TRUE(ReadAll(fd, reply_hdr, 8)); ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); uint32_t reply_pb_size = butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); ASSERT_GT(reply_pb_size, 0u); ASSERT_LE(reply_pb_size, 4096u); std::string reply_body(reply_pb_size, '\0'); - ASSERT_EQ((ssize_t)reply_pb_size, read(fd, &reply_body[0], reply_pb_size)); + ASSERT_TRUE(ReadAll(fd, &reply_body[0], reply_pb_size)); ASSERT_TRUE(reply->ParseFromString(reply_body)); } @@ -2295,38 +2058,27 @@ TEST_F(RdmaTest, v3_server_accepts_client_hello_with_ece) { EceFlagGuard ece_flag_guard(true); 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); // ACK flags=0 -> clean FALLBACK_TCP so the test ends without hardware. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2336,33 +2088,24 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_when_disabled) { EceFlagGuard ece_flag_guard(false); 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); EXPECT_FALSE(reply.has_ece()); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); sockfd.reset(-1); - usleep(100000); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2374,33 +2117,24 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_without_hw_negotiation) { EceFlagGuard ece_flag_guard(true); 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); - Socket* s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); EXPECT_FALSE(reply.has_ece()); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); sockfd.reset(-1); - usleep(100000); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2440,12 +2174,11 @@ TEST_F(RdmaTest, client_alloc_resource_fail_fallback_tcp) { req.set_sleep_us(200000); google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s.get()); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); // The socket must not be failed, otherwise it can no longer carry TCP. ASSERT_FALSE(s->Failed()); @@ -2461,18 +2194,11 @@ TEST_F(RdmaTest, server_alloc_resource_fail_fallback_tcp) { StartServer(); ResourceAllocFailGuard alloc_fail_guard(true); - 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); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Send a well-formed v2 hello: the negotiation succeeds // but the resource allocation does not. @@ -2489,24 +2215,19 @@ TEST_F(RdmaTest, server_alloc_resource_fail_fallback_tcp) { uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; memcpy(data, "RDMA", 4); msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_FALSE(s->Failed()); // Ack without RDMA so that the server finishes the handshake in TCP mode. uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); ASSERT_FALSE(s->Failed()); sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } @@ -2530,11 +2251,11 @@ TEST_F(RdmaTest, try_global_disable_rdma) { req.set_sleep_us(200000); google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(0, cntl.ErrorCode()); @@ -2723,7 +2444,6 @@ TEST_P(RdmaRpcTest, tcp_client_to_tcp_server) { 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()); @@ -2745,7 +2465,6 @@ TEST_P(RdmaRpcTest, tcp_client_to_rdma_server) { 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()); @@ -2768,7 +2487,6 @@ TEST_P(RdmaRpcTest, rdma_client_to_tcp_server) { 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()); @@ -3264,11 +2982,12 @@ TEST_P(RdmaRpcTest, verbs_error_handling) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); // wait for rdma handshake complete - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_TRUE(WaitUntil([&] { + return Socket::Address(cntl._single_server_id, &s) == 0; + })); ibv_send_wr wr; + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s.get()); memset(&wr, 0, sizeof(wr)); ibv_sge sge; void* buf = malloc(8192); diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index d021758224..e70f397448 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -89,7 +89,7 @@ static void* RunBlockingHandshakeRead(void* arg) { char byte = 0; read->result = read->io->ReadExact(&byte, sizeof(byte)); read->error = errno; - return NULL; + return nullptr; } class TestAppConnect : public AppConnect { @@ -110,7 +110,7 @@ static FrameSpec FixedSpec(const char* magic, size_t magic_len, class TestHandshakeProtocol : public HandshakeProtocol { public: - explicit TestHandshakeProtocol(std::string* calls = NULL) + explicit TestHandshakeProtocol(std::string* calls = nullptr) : _calls(calls), _with_extension(false), _extension_size(0) {} @@ -120,7 +120,7 @@ class TestHandshakeProtocol : public HandshakeProtocol { return spec; } const FrameSpec& AckFrameSpec() const override { - static const FrameSpec spec = FixedSpec(NULL, 0, 1); + static const FrameSpec spec = FixedSpec(nullptr, 0, 1); return spec; } StepResult BuildHello(bool enabled, std::string* payload) override { @@ -161,12 +161,12 @@ class TestHandshakeProtocol : public HandshakeProtocol { void EnableExtension(size_t size) { _with_extension = true; _extension_size = size; - _extension_spec = FixedSpec(NULL, 0, size); + _extension_spec = FixedSpec(nullptr, 0, size); } private: void Append(const char* value) { - if (_calls != NULL) { + if (_calls != nullptr) { *_calls += value; } } @@ -179,7 +179,7 @@ class TestHandshakeProtocol : public HandshakeProtocol { class TestHandshakeTransport : public HandshakeTransport { public: - explicit TestHandshakeTransport(std::string* calls = NULL) + explicit TestHandshakeTransport(std::string* calls = nullptr) : prepare_result(STEP_OK), negotiate_result(STEP_OK), validate_result(STEP_OK), tcp_active(false), high_speed_active(false), failed(false), _calls(calls) {} @@ -212,7 +212,7 @@ class TestHandshakeTransport : public HandshakeTransport { private: void Append(const char* value) { - if (_calls != NULL) { + if (_calls != nullptr) { *_calls += value; } } @@ -411,7 +411,7 @@ TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { bthread_usleep(10000); ASSERT_EQ(0, socket->SetFailed( EFAILEDSOCKET, "cancel blocked handshake read")); - ASSERT_EQ(0, bthread_join(tid, NULL)); + ASSERT_EQ(0, bthread_join(tid, nullptr)); EXPECT_EQ(-1, read.result); EXPECT_EQ(EFAILEDSOCKET, read.error); @@ -725,7 +725,7 @@ TEST(TransportHandshakeTest, source.append(invalid_prefixes[i]); const size_t original_size = source.size(); const ParseResult result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()) << invalid_prefixes[i]; @@ -758,7 +758,7 @@ TEST(TransportHandshakeTest, source.append(partial_prefixes[i]); const size_t original_size = source.size(); const ParseResult result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); EXPECT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()) << partial_prefixes[i]; @@ -787,7 +787,7 @@ TEST(TransportHandshakeTest, butil::IOBuf source; source.append(hello.data(), 1); ParseResult result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); ASSERT_EQ(1UL, source.size()); @@ -796,7 +796,7 @@ TEST(TransportHandshakeTest, source.append(hello.data() + 1, hello.size() - 1); result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); ASSERT_TRUE(source.empty()); @@ -816,7 +816,7 @@ TEST(TransportHandshakeTest, const uint32_t ack = 0; source.append(&ack, sizeof(ack)); result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); ASSERT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); ASSERT_TRUE(source.empty()); @@ -846,7 +846,7 @@ TEST(TransportHandshakeTest, source.append(application_data, sizeof(application_data) - 1); const ParseResult result = policy::ParseTransportHandshake( - &source, socket.get(), false, NULL); + &source, socket.get(), false, nullptr); ASSERT_FALSE(result.is_ok()); ASSERT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); ASSERT_EQ(sizeof(application_data) - 1, source.size());