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/17] 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/17] 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 0d5aeb55fc381dbbc92359a5ff3fbe243610e470 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Fri, 23 Oct 2026 00:47:23 -0700 Subject: [PATCH 03/17] Integrate URMA with the common transport handshake (#88) * Unify transport handshakes on current upstream master * Fix certificate name extraction with const OpenSSL accessors * Integrate URMA with the common transport handshake Port URMA v2 and v3 negotiation into AdapterTransport while preserving the existing RDMA and UBSHM implementations and regression tests. Include standalone bonding import, ordered SEND work requests, the 4096-byte CTP payload limit, bounded diagnostics, mock regression coverage and concise documentation. * Run Linux CI for transport handshake pull requests * Fix URMA handshake fallback compatibility * Use nullptr and refresh URMA handshake factory tests * Test legacy URMA interoperability and final ACK handoff * Fix URMA CQ callback lifetime during teardown Pin the main Socket through a carrier-owned context before endpoint access. Serialize CQ dispatch with teardown and reject callbacks from an earlier reset generation. Add barrier-controlled event/polling and copy/zero-copy regressions for concurrent close, last-reference recycling, and stale callbacks after reset. Keep this lifecycle fix separate from handshake wire compatibility. * Simplify URMA documentation * Restore upstream CI and URMA documentation * Clarify URMA handshake comments and align local style --------- Co-authored-by: Chuang Zhang --- BUILD.bazel | 2 +- Makefile | 7 +- docs/cn/handshake_common_design.md | 38 + src/brpc/CMakeLists.txt | 2 +- src/brpc/adapter_transport.cpp | 821 +++++++ src/brpc/adapter_transport.h | 105 + 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 | 486 ++++ src/brpc/handshake/ubshm_handshake.h | 130 ++ .../handshake/urma_fallback_handshake.cpp | 115 + .../{urma => handshake}/urma_handshake.cpp | 235 +- src/brpc/{urma => handshake}/urma_handshake.h | 61 +- 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 | 37 + .../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 | 9 +- src/brpc/transport_factory.h | 4 +- src/brpc/transport_handshake.cpp | 431 ++++ src/brpc/transport_handshake.h | 237 ++ src/brpc/ubshm/ub_endpoint.cpp | 789 +------ src/brpc/ubshm/ub_endpoint.h | 156 +- src/brpc/ubshm_transport.cpp | 120 +- src/brpc/ubshm_transport.h | 28 +- src/brpc/urma/mock_urma.cpp | 46 +- src/brpc/urma/urma_endpoint.cpp | 316 ++- src/brpc/urma/urma_endpoint.h | 73 +- src/brpc/urma/urma_helper.cpp | 3 + src/brpc/urma/urma_helper.h | 4 + src/brpc/urma/urma_stream_trace.h | 142 ++ src/brpc/urma_transport.cpp | 112 +- src/brpc/urma_transport.h | 32 +- test/CMakeLists.txt | 3 + test/brpc_rdma_unittest.cpp | 1376 ++++++----- test/brpc_transport_handshake_unittest.cpp | 960 ++++++++ test/brpc_ubring_unittest.cpp | 559 ++++- test/brpc_urma_fallback_unittest.cpp | 177 ++ test/brpc_urma_unittest.cpp | 2023 ++++++++++++++++- 59 files changed, 9824 insertions(+), 3539 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/handshake/urma_fallback_handshake.cpp rename src/brpc/{urma => handshake}/urma_handshake.cpp (58%) rename src/brpc/{urma => handshake}/urma_handshake.h (74%) 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 src/brpc/urma/urma_stream_trace.h create mode 100644 test/brpc_transport_handshake_unittest.cpp create mode 100644 test/brpc_urma_fallback_unittest.cpp diff --git a/BUILD.bazel b/BUILD.bazel index 23e433c76e..7c32ff82d0 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"], @@ -543,6 +542,7 @@ brpc_proto_library( ) URMA_SRC_PATTERNS = [ + "src/brpc/handshake/urma_handshake.cpp", "src/brpc/urma/*.cpp", "src/brpc/urma/**/*.cpp", ] diff --git a/Makefile b/Makefile index 271b518ae6..f31564f693 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 @@ -214,7 +214,8 @@ ifeq ($(URMA_USE_MOCK),0) BRPC_SOURCES_ALL := $(filter-out src/brpc/urma/mock_urma.cpp,$(BRPC_SOURCES_ALL)) endif BRPC_SOURCES = $(filter-out $(THRIFT_SOURCES) $(EXCLUDE_SOURCES), $(BRPC_SOURCES_ALL)) -BRPC_PROTOS = $(filter %.proto,$(BRPC_SOURCES)) +# TCP-only servers also encode the disabled URMA Hello. +BRPC_PROTOS = $(sort $(filter %.proto,$(BRPC_SOURCES)) src/brpc/urma/urma_handshake.proto) BRPC_CFAMILIES = $(filter-out %.proto %.pb.cc,$(BRPC_SOURCES)) BRPC_OBJS = $(BRPC_PROTOS:.proto=.pb.o) $(addsuffix .o, $(basename $(BRPC_CFAMILIES))) @@ -312,6 +313,8 @@ output/bin:protoc-gen-mcpack src/mcpack2pb/generator.o:src/mcpack2pb/generator.cpp src/idl_options.pb.h $(CXX) -c $(HDRPATHS) $(CXXFLAGS) $< -o $@ +src/brpc/handshake/urma_fallback_handshake.o src/brpc/handshake/urma_fallback_handshake.dbg.o: src/brpc/urma/urma_handshake.pb.h + %.o:%.cpp @echo "> Compiling $@" $(CXX) -c $(HDRPATHS) $(CXXFLAGS) $< -o $@ diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md new file mode 100644 index 0000000000..dba2808398 --- /dev/null +++ b/docs/cn/handshake_common_design.md @@ -0,0 +1,38 @@ +# 公共传输握手 + +`AdapterTransport` 统一 RDMA、UBSHM 和 URMA 的握手入口。TCP 承担控制面， +升级成功后使用高速数据面，协商不可用时继续使用 TCP。 + +## 组件分工 + +| 组件 | 职责 | +| --- | --- | +| `AdapterTransport` | 持有 TCP 与候选 Transport，选择协议及数据面 | +| `HandshakeSession` | 编排客户端、服务端状态机和增量帧解析 | +| `HandshakeProtocol` | 编解码 hello、ACK 和扩展帧，处理版本兼容 | +| `HandshakeTransport` | 准备资源，执行升级、回退和失败清理 | +| 各传输的 Endpoint | 管理资源及数据收发 | + +客户端在 TCP 建连后启动握手 bthread；服务端通过 `InputMessenger` 增量解析 +hello、扩展字段和 ACK。ACK 不带 magic，服务端使用已选协议继续解析。 + +## 状态与回退 + +- `ESTABLISHED`：使用高速数据面，TCP 保留为控制连接；额外 TCP 业务数据按协议错误处理。 +- `FALLBACK_TCP`：资源不可用或协商拒绝时释放升级资源，继续解析已缓存及后续 TCP 数据。 +- `FAILED`：握手帧畸形或处理失败，清理资源并关闭连接。 + +终态以 release/acquire 顺序发布和读取。未收齐的帧返回 `NOT_ENOUGH_DATA`， +不消费输入；非握手 magic 返回 `TRY_OTHERS`。帧边界处理支持 TCP 分片和粘连。 + +## 协议与数据面衔接 + +- RDMA 保持既有 v2/v3 wire format、协议 ID 和注册名称。 +- UBSHM 使用 v3 hello：64 字节 hello 后交换 4 字节网络序格式扩展，IPC 选择 + `IPC_V2`，UBS 选择 `LEGACY_64`；格式一致后升级，ACK 为 4 字节。 +- URMA 保持 v2 二进制和 v3 protobuf 格式；资源导入和收发信用由 + `UrmaEndpoint` 管理，详见 [UrmaTransport](urma.md)。 +- 激活数据面后，等待 TCP 握手解析退出再注册 URMA CQ 回调，避免两条输入路径并行解析。 + +测试覆盖分片与粘连、magic 分流、版本协商、终态发布、TCP 回退、 +资源清理及各传输的数据面回归。 diff --git a/src/brpc/CMakeLists.txt b/src/brpc/CMakeLists.txt index 4efd059bfd..c5174d3108 100644 --- a/src/brpc/CMakeLists.txt +++ b/src/brpc/CMakeLists.txt @@ -279,7 +279,7 @@ set(PROTO_FILES brpc/builtin_service.proto brpc/policy/public_pbrpc_meta.proto brpc/policy/sofa_pbrpc_meta.proto brpc/proto_base.proto - brpc/rdma/rdma_handshake.proto + brpc/rdma_handshake.proto brpc/rpc_dump.proto brpc/rtmp.proto brpc/span.proto diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp new file mode 100644 index 0000000000..8ac60da2be --- /dev/null +++ b/src/brpc/adapter_transport.cpp @@ -0,0 +1,821 @@ +// 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" +#include "brpc/handshake/urma_handshake.h" +#if BRPC_WITH_RDMA +#include "brpc/rdma/rdma_helper.h" +#endif +#if BRPC_WITH_UBRING +#include "brpc/ubshm/shm/shm_mgr.h" +#include "brpc/ubshm/ub_endpoint.h" +#include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/ubr_trx.h" +#endif +#if BRPC_WITH_URMA +#include "brpc/urma/urma_helper.h" +#endif +#include "brpc/rdma_transport.h" +#include "brpc/tcp_transport.h" +#include "brpc/ubshm_transport.h" +#include "brpc/urma_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_URMA +class UrmaClientHandshakeTransport : public handshake::HandshakeTransport { +public: + UrmaClientHandshakeTransport( + UrmaTransport* transport, urma::UrmaHandshakeAdapter* protocol, + Socket* socket, int* connect_error) + : _transport(transport), _protocol(protocol), _socket(socket), + _connect_error(connect_error) {} + + handshake::StepResult PrepareResources() override { + return _transport->PrepareUpgradeResources(false) == 0 + ? handshake::STEP_OK : 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 URMA handshake from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + } + +private: + UrmaTransport* _transport; + urma::UrmaHandshakeAdapter* _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(), _protocol->remote().len, + _protocol->local_format()) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + void OnEstablished() override { + _transport->GetUBShmEp()->SetNegotiatedDataFormat( + _protocol->local_format()); + _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_URMA + case SOCKET_MODE_URMA: + 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); + const bool matches_urma = + MatchesMagicPrefix(prefix, prefix_len, "URMA", 4) || + MatchesMagicPrefix(prefix, prefix_len, "URM3", 4); + if (!matches_ub && !matches_rdma && !matches_urma) { + result = ParseResult(PARSE_ERROR_TRY_OTHERS); + } else { + handshake::HandshakeAdapter* adapter = nullptr; + if (matches_ub) { + adapter = handshake::GetUBShmServerHandshakeAdapter(); + } else if (matches_rdma) { + adapter = handshake::GetRdmaServerHandshakeAdapter(); + } + else { + adapter = handshake::GetUrmaServerHandshakeAdapter(); + } + CHECK(adapter != nullptr); + 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_URMA + if (adapter->_mode == SOCKET_MODE_URMA) { + UrmaTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + if (!urma::IsUrmaAvailable()) { + adapter->FallbackToTcp(); + adapter->CompleteConnection(handshake::FALLBACK_TCP); + task->done(0, task->data); + return nullptr; + } + std::unique_ptr protocol = + urma::CreateClientHandshakeAdapter(transport->GetUrmaEp()); + UrmaClientHandshakeTransport 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) { + connect_error = errno != 0 ? errno : EIO; + transport->DeactivateUpgrade(); + adapter->_handshake.MarkFailed(); + socket->SetFailed(connect_error, "Fail to start URMA CQ events"); + } + 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 nullptr; + } +#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( + ubring::PreferredDataFormatForShmType(ubring::GetShmType())); + 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->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 +#if BRPC_WITH_URMA + } else if (_mode == SOCKET_MODE_URMA && + options.user != static_cast( + get_client_side_messenger())) { + _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 +#if BRPC_WITH_URMA + case SOCKET_MODE_URMA: + _high_speed_transport.reset(new UrmaTransport); + 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 +#if BRPC_WITH_URMA + case SOCKET_MODE_URMA: + 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 BRPC_WITH_URMA + if (adapter->_mode == SOCKET_MODE_URMA && + adapter->_handshake.phase(butil::memory_order_acquire) == + handshake::ESTABLISHED) { + UrmaTransport* transport = static_cast( + adapter->_high_speed_transport.get()); + if (transport->StartUpgradeEvents() < 0) { + const int saved_errno = errno != 0 ? errno : EIO; + transport->DeactivateUpgrade(); + adapter->_handshake.MarkFailed(); + adapter->CompleteConnection(handshake::FAILED); + socket->SetFailed(saved_errno, "Fail to start URMA CQ events"); + } + } +#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..f0b8ae60c4 --- /dev/null +++ b/src/brpc/adapter_transport.h @@ -0,0 +1,105 @@ +// 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; +class UrmaTransport; + +// 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; +friend class UrmaTransport; +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 4cce93ed0c..7aa6a91cdf 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..28c7f84e62 --- /dev/null +++ b/src/brpc/handshake/ubshm_handshake.cpp @@ -0,0 +1,486 @@ +// 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/shm/shm_mgr.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); + +UbrDataFormat PreferredDataFormatForShmType(SHM_TYPE shm_type) { + switch (shm_type) { + case SHM_TYPE_IPC: + return UBR_DATA_FORMAT_IPC_V2; + case SHM_TYPE_UBS: + return UBR_DATA_FORMAT_LEGACY_64; + default: + return UBR_DATA_FORMAT_NONE; + } +} + +UbrDataFormat SelectDataFormat(UbrDataFormat local_format, + uint16_t remote_format_id) { + if ((local_format != UBR_DATA_FORMAT_LEGACY_64 && + local_format != UBR_DATA_FORMAT_IPC_V2) || + remote_format_id != static_cast(local_format)) { + return UBR_DATA_FORMAT_NONE; + } + return local_format; +} + +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); + // The decoded name is also used directly as a C-string by callers. + shm_name[SHM_MAX_NAME_BUFF_LEN - 1] = '\0'; +} + +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 ? _local_format + : 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 && + SelectDataFormat(_local_format, extension.format_id) != + UBR_DATA_FORMAT_NONE + ? 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; + } + // Deserialize terminates the decoded name for safe C-string use. Inspect + // 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) { + 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, + _protocol->local_format()) < 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->GetUBShmEp()->SetNegotiatedDataFormat( + _protocol->local_format()); + _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( + ubring::PreferredDataFormatForShmType(ubring::GetShmType())); + 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->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..d487f78e99 --- /dev/null +++ b/src/brpc/handshake/ubshm_handshake.h @@ -0,0 +1,130 @@ +// 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" +#include "brpc/ubshm/ubr_data_format.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]; +}; + +UbrDataFormat PreferredDataFormatForShmType(SHM_TYPE shm_type); +UbrDataFormat SelectDataFormat(UbrDataFormat local_format, + uint16_t remote_format_id); + +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: + explicit UBShmHandshakeAdapter( + UbrDataFormat local_format = UBR_DATA_FORMAT_LEGACY_64) + : _local_len(0), _server_reply(false), _remote(), + _local_format(local_format) {} + + 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; } + UbrDataFormat local_format() const { return _local_format; } + + 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; + UbrDataFormat _local_format; + DISALLOW_COPY_AND_ASSIGN(UBShmHandshakeAdapter); +}; + +} // namespace ubring +} // namespace brpc + +#endif // BRPC_WITH_UBRING +#endif // BRPC_HANDSHAKE_UBSHM_HANDSHAKE_H diff --git a/src/brpc/handshake/urma_fallback_handshake.cpp b/src/brpc/handshake/urma_fallback_handshake.cpp new file mode 100644 index 0000000000..1461d46927 --- /dev/null +++ b/src/brpc/handshake/urma_fallback_handshake.cpp @@ -0,0 +1,115 @@ +// 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/urma_handshake.h" + +#include + +#include "brpc/adapter_transport.h" +#include "brpc/urma/urma_handshake.pb.h" + +namespace brpc { +namespace handshake { + +// This codec needs neither liburma nor its SDK headers. A TCP-only server +// still recognizes the URMA control frames and replies with a disabled Hello. +class UrmaFallbackProtocol : public FallbackHandshakeProtocol { +public: + explicit UrmaFallbackProtocol(int version) : _version(version) {} + + int ProtocolVersion() const override { return _version; } + bool AckDisabledHello() const override { return false; } + const FrameSpec& HelloFrameSpec() const override { + static const FrameSpec v2("URMA", 4, 86, 4096, + FrameSpec::U16_TOTAL_LENGTH); + static const FrameSpec v3("URM3", 4, 8, 4104, + FrameSpec::U32_BODY_LENGTH); + return _version == 3 ? v3 : v2; + } + const FrameSpec& AckFrameSpec() const override { + static const FrameSpec spec(nullptr, 0, 4, 4, FrameSpec::FIXED); + return spec; + } + StepResult ParseHello(const std::string& payload) override { + if (_version == 3) { + urma::UrmaHello hello; + if (!hello.ParseFromString(payload) || hello.eid().size() != 16 || + hello.seg_eid().size() != 16) { + errno = EPROTO; + return STEP_ERROR; + } + } + return STEP_FALLBACK; + } + StepResult BuildHello(bool enabled, std::string* payload) override { + if (enabled) { + errno = EPROTO; + return STEP_ERROR; + } + if (_version == 2) { + // FrameCodec supplies the magic and two-byte total length. + payload->assign(80, '\0'); + return STEP_OK; + } + urma::UrmaHello hello; + hello.set_buffer_size(0); + hello.set_recv_buffer_cnt(0); + hello.set_jetty_id(0); + hello.set_eid(std::string(16, '\0')); + hello.set_uasid(0); + hello.set_tp_type(0); + hello.set_seg_eid(std::string(16, '\0')); + hello.set_seg_uasid(0); + hello.set_seg_va(0); + hello.set_seg_len(0); + hello.set_seg_token_id(0); + return hello.SerializeToString(payload) ? STEP_OK : STEP_ERROR; + } + +private: + int _version; +}; + +StepResult RunUrmaFallbackServerHandshake(butil::IOBuf* source, Socket* socket) { + UrmaFallbackProtocol v2(2); + UrmaFallbackProtocol v3(3); + std::vector protocols{&v2, &v3}; + IOBufHandshakeInput input(source); + FallbackHandshakeTransport transport; + return AdapterTransport::Get(socket)->handshake_session()->RunServer( + protocols, &input, &transport, false); +} + +#if !BRPC_WITH_URMA +class UrmaFallbackServerAdapter : public StandardHandshakeAdapter { +protected: + StepResult RunServerStep(butil::IOBuf* source, Socket* socket) override { + return RunUrmaFallbackServerHandshake(source, socket); + } + HandshakeSession* GetSession(Socket* socket) const override { + return AdapterTransport::Get(socket)->handshake_session(); + } +}; + +HandshakeAdapter* GetUrmaServerHandshakeAdapter() { + static UrmaFallbackServerAdapter adapter; + return &adapter; +} +#endif + +} // namespace handshake +} // namespace brpc diff --git a/src/brpc/urma/urma_handshake.cpp b/src/brpc/handshake/urma_handshake.cpp similarity index 58% rename from src/brpc/urma/urma_handshake.cpp rename to src/brpc/handshake/urma_handshake.cpp index b7018606e0..0cd7e778c0 100644 --- a/src/brpc/urma/urma_handshake.cpp +++ b/src/brpc/handshake/urma_handshake.cpp @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -#include "brpc/urma/urma_handshake.h" +#include "brpc/handshake/urma_handshake.h" + +#include "brpc/adapter_transport.h" #if BRPC_WITH_URMA @@ -188,6 +190,164 @@ bool ValidHello(const ParsedHello& h) { return true; } +namespace { + +const handshake::FrameSpec& UrmaV2FrameSpec() { + static const handshake::FrameSpec spec( + "URMA", 4, v2_wire::HELLO_PACKET_LEN, + v2_wire::HELLO_MSG_LEN_MAX, + handshake::FrameSpec::U16_TOTAL_LENGTH); + return spec; +} + +const handshake::FrameSpec& UrmaV3FrameSpec() { + static const handshake::FrameSpec spec( + "URM3", 4, 8, 8 + MAX_V3_PB_SIZE, + handshake::FrameSpec::U32_BODY_LENGTH); + return spec; +} + +const handshake::FrameSpec& UrmaAckFrameSpec() { + static const handshake::FrameSpec spec( + nullptr, 0, sizeof(uint32_t), sizeof(uint32_t), + handshake::FrameSpec::FIXED); + return spec; +} + +void TranslateV2Hello(const v2_wire::HelloMessage& message, + ParsedHello* out) { + out->buffer_size = message.buffer_size; + out->recv_buffer_cnt = message.recv_buffer_cnt; + out->jetty_id = message.jetty_id; + memcpy(out->eid, message.eid, sizeof(out->eid)); + out->uasid = message.uasid; + out->tp_type = message.tp_type; + memcpy(out->seg_eid, message.seg_eid, sizeof(out->seg_eid)); + out->seg_uasid = message.seg_uasid; + out->seg_va = message.seg_va; + out->seg_len = message.seg_len; + out->seg_token_id = message.seg_token_id; +} + +void TranslateV3Hello(const UrmaHello& message, ParsedHello* out) { + out->buffer_size = message.buffer_size(); + out->recv_buffer_cnt = message.recv_buffer_cnt(); + out->jetty_id = message.jetty_id(); + memcpy(out->eid, message.eid().data(), sizeof(out->eid)); + out->uasid = message.uasid(); + out->tp_type = static_cast(message.tp_type()); + memcpy(out->seg_eid, message.seg_eid().data(), sizeof(out->seg_eid)); + out->seg_uasid = message.seg_uasid(); + out->seg_va = message.seg_va(); + out->seg_len = message.seg_len(); + out->seg_token_id = message.seg_token_id(); +} + +} // namespace + +const handshake::FrameSpec& UrmaHandshakeAdapter::HelloFrameSpec() const { + return _version == 3 ? UrmaV3FrameSpec() : UrmaV2FrameSpec(); +} + +const handshake::FrameSpec& UrmaHandshakeAdapter::AckFrameSpec() const { + return UrmaAckFrameSpec(); +} + +handshake::StepResult UrmaHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + if (_version == 3) { + UrmaHello message; + if (enabled) { + CHECK(_ep != nullptr); + _ep->FillLocalHelloV3(&message); + } else { + message.set_buffer_size(0); + message.set_recv_buffer_cnt(0); + message.set_jetty_id(0); + message.set_eid(std::string(16, '\0')); + message.set_uasid(0); + message.set_tp_type(0); + message.set_seg_eid(std::string(16, '\0')); + message.set_seg_uasid(0); + message.set_seg_va(0); + message.set_seg_len(0); + message.set_seg_token_id(0); + } + return message.SerializeToString(payload) + ? handshake::STEP_OK : handshake::STEP_ERROR; + } + v2_wire::HelloMessage message{}; + if (enabled) { + CHECK(_ep != nullptr); + _ep->FillLocalHelloV2(&message); + } else { + message.msg_len = v2_wire::HELLO_PACKET_LEN; + } + uint8_t body[v2_wire::HELLO_BODY_LEN]; + message.Serialize(body); + payload->assign(reinterpret_cast(body + sizeof(uint16_t)), + sizeof(body) - sizeof(uint16_t)); + return handshake::STEP_OK; +} + +handshake::StepResult UrmaHandshakeAdapter::ParseHello( + const std::string& payload) { + ParsedHello remote{}; + if (_version == 3) { + UrmaHello message; + if (!message.ParseFromString(payload) || + message.eid().size() != sizeof(remote.eid) || + message.seg_eid().size() != sizeof(remote.seg_eid)) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + // ParsedHello stores this as a byte. Reject an out-of-range protobuf + // value before narrowing it, otherwise e.g. 256 becomes a valid 0. + if (message.tp_type() > static_cast(URMA_UTP)) { + return handshake::STEP_FALLBACK; + } + TranslateV3Hello(message, &remote); + } else { + if (payload.size() < v2_wire::HELLO_BODY_LEN - sizeof(uint16_t)) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint8_t body[v2_wire::HELLO_BODY_LEN]; + const uint16_t total_be = butil::HostToNet16( + static_cast(4 + sizeof(uint16_t) + payload.size())); + memcpy(body, &total_be, sizeof(total_be)); + memcpy(body + sizeof(total_be), payload.data(), + sizeof(body) - sizeof(total_be)); + v2_wire::HelloMessage message{}; + message.Deserialize(body); + if (message.hello_ver != v2_wire::HELLO_V2_VERSION || + message.impl_ver != v2_wire::IMPL_V2_VERSION) { + return handshake::STEP_FALLBACK; + } + TranslateV2Hello(message, &remote); + } + if (!ValidHello(remote)) { + return handshake::STEP_FALLBACK; + } + _remote = remote; + return handshake::STEP_OK; +} + +std::unique_ptr CreateClientHandshakeAdapter( + UrmaEndpoint* ep) { + return std::unique_ptr( + new UrmaHandshakeAdapter( + ep, FLAGS_urma_client_handshake_version == 3 ? 3 : 2)); +} + +std::vector > +CreateServerHandshakeAdapters(UrmaEndpoint* ep) { + std::vector > adapters; + adapters.emplace_back(new UrmaHandshakeAdapter(ep, 2)); + adapters.emplace_back(new UrmaHandshakeAdapter(ep, 3)); + return adapters; +} + // File-local (not in the anonymous namespace so it can be friend-declared // from urma_endpoint.h's UrmaEndpoint). Reads the body following the magic // and translates it into ParsedHello. @@ -356,6 +516,79 @@ UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, } } // namespace urma + +namespace handshake { + +class UrmaServerHandshakeTransport : public HandshakeTransport { +public: + UrmaServerHandshakeTransport(UrmaTransport* transport, + butil::IOBuf* source) + : _transport(transport), _source(source), _protocol(nullptr) {} + + void OnProtocolSelected(HandshakeProtocol* protocol) override { + _protocol = static_cast(protocol); + } + StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources(true) == 0) { + return STEP_OK; + } + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + StepResult NegotiateResources() override { + CHECK(_protocol != nullptr); + if (_transport->NegotiateUpgradeResources( + _protocol->remote(), true) == 0) { + return STEP_OK; + } + _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: + UrmaTransport* _transport; + butil::IOBuf* _source; + urma::UrmaHandshakeAdapter* _protocol; +}; + +class UrmaServerHandshakeAdapter : public StandardHandshakeAdapter { +protected: + StepResult RunServerStep(butil::IOBuf* source, Socket* socket) override { + AdapterTransport* adapter = AdapterTransport::Get(socket); + if (!adapter->upgrade_capable(SOCKET_MODE_URMA)) { + return RunUrmaFallbackServerHandshake(source, socket); + } + UrmaTransport* transport = static_cast( + adapter->high_speed_transport()); + std::vector > owned = + urma::CreateServerHandshakeAdapters(transport->GetUrmaEp()); + std::vector protocols; + for (size_t i = 0; i < owned.size(); ++i) { + protocols.push_back(owned[i].get()); + } + IOBufHandshakeInput input(source); + UrmaServerHandshakeTransport participant(transport, source); + return adapter->handshake_session()->RunServer( + protocols, &input, &participant, false); + } + + HandshakeSession* GetSession(Socket* socket) const override { + return AdapterTransport::Get(socket)->handshake_session(); + } +}; + +HandshakeAdapter* GetUrmaServerHandshakeAdapter() { + static UrmaServerHandshakeAdapter adapter; + return &adapter; +} + +} // namespace handshake } // namespace brpc #endif // BRPC_WITH_URMA diff --git a/src/brpc/urma/urma_handshake.h b/src/brpc/handshake/urma_handshake.h similarity index 74% rename from src/brpc/urma/urma_handshake.h rename to src/brpc/handshake/urma_handshake.h index f406bfa1b9..a6ca94acea 100644 --- a/src/brpc/urma/urma_handshake.h +++ b/src/brpc/handshake/urma_handshake.h @@ -15,15 +15,30 @@ // specific language governing permissions and limitations // under the License. -#ifndef BRPC_URMA_HANDSHAKE_H -#define BRPC_URMA_HANDSHAKE_H +#ifndef BRPC_HANDSHAKE_URMA_HANDSHAKE_H +#define BRPC_HANDSHAKE_URMA_HANDSHAKE_H #include #include +#include +#include +#include + +#include "brpc/handshake/handshake_adapter.h" + +namespace brpc { +namespace handshake { + +HandshakeAdapter* GetUrmaServerHandshakeAdapter(); +StepResult RunUrmaFallbackServerHandshake(butil::IOBuf* source, Socket* socket); + +} // namespace handshake +} // namespace brpc #if BRPC_WITH_URMA #include "urma_types.h" +#include "brpc/transport_handshake.h" namespace brpc { namespace urma { @@ -68,7 +83,8 @@ constexpr size_t HELLO_MSG_LEN_MAX = 4096; constexpr uint16_t HELLO_V2_VERSION = 2; constexpr uint16_t IMPL_V2_VERSION = 1; -// The serializable struct. Aligned so it can be reinterpreted as raw bytes. +// The v2 hello fields are serialized explicitly in network byte order. +// The wire format is independent of the in-memory struct layout. struct HelloMessage { uint16_t msg_len; // total packet length (incl. magic) uint16_t hello_ver; @@ -79,7 +95,7 @@ struct HelloMessage { uint8_t eid[16]; uint32_t uasid; uint8_t tp_type; - uint8_t pad[3]; // keep the struct 4-byte aligned + uint8_t pad[3]; // reserved wire bytes, serialized as zeros uint8_t seg_eid[16]; uint32_t seg_uasid; uint64_t seg_va; @@ -92,7 +108,8 @@ struct HelloMessage { } // namespace v2_wire -// Abstract handshake strategy. v2 speaks binary; v3 speaks protobuf ("URM3"). +// Legacy TCP handshake strategy used by interoperability tests. +// v2 speaks binary; v3 speaks protobuf ("URM3"). class UrmaHandshake { public: virtual ~UrmaHandshake() = default; @@ -109,6 +126,34 @@ class UrmaHandshake { virtual int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) = 0; }; +// Wire-format adapter used by the common transport handshake. +// Stores the selected protocol version and parsed peer hello; +// UrmaEndpoint owns the local and imported URMA resources. +class UrmaHandshakeAdapter : public handshake::HandshakeProtocol { +public: + UrmaHandshakeAdapter(UrmaEndpoint* ep, int version) + : _ep(ep), _version(version), _remote() {} + + int ProtocolVersion() const override { return _version; } + bool AckDisabledHello() const override { return false; } + 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; + const ParsedHello& remote() const { return _remote; } + +private: + UrmaEndpoint* _ep; + int _version; + ParsedHello _remote; +}; + +std::unique_ptr CreateClientHandshakeAdapter( + UrmaEndpoint* ep); +std::vector > +CreateServerHandshakeAdapters(UrmaEndpoint* ep); + // v2 binary handshake (magic "URMA"). class UrmaHandshakeClientV2 : public UrmaHandshake { public: @@ -155,10 +200,10 @@ class UrmaHandshakeServerV3 : public UrmaHandshake { UrmaEndpoint* _ep; }; -// Client-side factory: picks v2/v3 based on --urma_client_handshake_version. +// Legacy client factory: picks v2/v3 via --urma_client_handshake_version. UrmaHandshake* CreateClientHandshake(UrmaEndpoint* ep); -// Server-side factory: dispatches on the first 4 bytes read from the TCP fd. +// Legacy server factory: dispatches on the first 4 bytes read from the TCP fd. // Returns nullptr if the magic is not "URMA"/"URM3" (caller falls back to TCP). // @magic is the first MAGIC_STR_LEN bytes already read from the fd. UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, @@ -177,4 +222,4 @@ int ReadBodyAndNegotiate(UrmaEndpoint* ep, ParsedHello* out, bool* negotiated); #endif // BRPC_WITH_URMA -#endif // BRPC_URMA_HANDSHAKE_H +#endif // BRPC_HANDSHAKE_URMA_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..5d343e0da7 --- /dev/null +++ b/src/brpc/policy/transport_handshake_protocol.cpp @@ -0,0 +1,37 @@ +// 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" + +namespace brpc { +namespace policy { + +ParseResult ParseTransportHandshake(butil::IOBuf* source, Socket* socket, + bool /*read_eof*/, const void* /*arg*/) { + 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..838a875fe0 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,21 +50,21 @@ 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 if (mode == SOCKET_MODE_URMA) { - return std::unique_ptr(new UrmaTransport()); + return std::unique_ptr(new AdapterTransport(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..2d93522059 --- /dev/null +++ b/src/brpc/transport_handshake.cpp @@ -0,0 +1,431 @@ +// 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->AckDisabledHello()) { + return FinishWithFallback(this, transport); + } + + 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, if the protocol requires it, + // consumes the peer ACK before publishing FALLBACK_TCP. + 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); + } + if (!enabled && !selected->AckDisabledHello()) { + return FinishWithFallback(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..72cfd24f24 --- /dev/null +++ b/src/brpc/transport_handshake.h @@ -0,0 +1,237 @@ +// 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); + + // Legacy URMA peers switch directly to TCP after a disabled Hello. + // Other protocols complete their rejection with an ACK. + virtual bool AckDisabledHello() const { return true; } + + 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 97e39e5d77..dd93a8f0a1 100644 --- a/src/brpc/ubshm/ub_endpoint.cpp +++ b/src/brpc/ubshm/ub_endpoint.cpp @@ -19,24 +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/shm/shm_mgr.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); @@ -45,696 +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 = NULL; -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); -} - -UbrDataFormat PreferredDataFormatForShmType(SHM_TYPE shm_type) { - switch (shm_type) { - case SHM_TYPE_IPC: - return UBR_DATA_FORMAT_IPC_V2; - case SHM_TYPE_UBS: - return UBR_DATA_FORMAT_LEGACY_64; - default: - return UBR_DATA_FORMAT_NONE; - } -} - -UbrDataFormat SelectDataFormat(UbrDataFormat local_format, - uint16_t remote_format_id) { - if ((local_format != UBR_DATA_FORMAT_LEGACY_64 && - local_format != UBR_DATA_FORMAT_IPC_V2) || - remote_format_id != static_cast(local_format)) { - return UBR_DATA_FORMAT_NONE; - } - return local_format; -} - -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); - // shm_name arrives verbatim from the peer and is used downstream as a - // C-string (strncpy/strrchr/snprintf("%s")/shm allocate+map). A peer that - // fills all SHM_MAX_NAME_BUFF_LEN bytes without a terminator would make - // those scans run past the buffer, so force termination here. - shm_name[SHM_MAX_NAME_BUFF_LEN - 1] = '\0'; -} - -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; - } - - const UbrDataFormat local_format = - PreferredDataFormatForShmType(GetShmType()); - 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, - static_cast(local_format)}; - 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) { - selected_format = - SelectDataFormat(local_format, remote_extension.format_id); - } - if (selected_format == UBR_DATA_FORMAT_NONE) { - LOG(WARNING) << "Fail to negotiate data format with server, " - << "fallback to tcp:" << s->description(); - ub_transport->_ub_state = UBShmTransport::UB_OFF; - } else { - ep->_state = C_MAP_REMOTE_SHM; - if (ep->_ub_ring->UbrMapRemoteShm( - &local_trx_shm, shm_name, remote_msg.len, - selected_format) < 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; - } - const UbrDataFormat local_format = - PreferredDataFormatForShmType(GetShmType()); - 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, local_format) < 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) { - selected_format = - SelectDataFormat(local_format, remote_extension.format_id); - } - if (selected_format != UBR_DATA_FORMAT_NONE) { - local_extension.format_id = - static_cast(selected_format); - } 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 { @@ -775,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)) { @@ -797,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; @@ -811,14 +147,24 @@ 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; } @@ -838,18 +184,44 @@ int UBShmEndpoint::AllocateServerResources(ubring::SHM* remote_trx_shm, 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, format); 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; @@ -867,11 +239,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; @@ -892,9 +267,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) { @@ -927,12 +303,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() { @@ -968,7 +343,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) { @@ -977,17 +352,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()); @@ -995,12 +371,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) { @@ -1019,8 +395,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; @@ -1045,8 +421,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 526966c9f7..ed4d0246eb 100644 --- a/src/brpc/ubshm/ub_endpoint.h +++ b/src/brpc/ubshm/ub_endpoint.h @@ -20,78 +20,36 @@ #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/ubr_data_format.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); -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); -}; - -UbrDataFormat PreferredDataFormatForShmType(SHM_TYPE shm_type); -UbrDataFormat SelectDataFormat(UbrDataFormat local_format, - uint16_t remote_format_id); - -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 +62,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 +76,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 +97,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 +105,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); @@ -182,58 +116,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 { @@ -245,8 +153,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 @@ -261,8 +168,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..3f8db23fc2 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,58 @@ 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, + uint64_t remote_shm_len, + ubring::UbrDataFormat format) { + return _ub_ep->_ub_ring->UbrMapRemoteShm( + local_trx_shm, shm_name, remote_shm_len, format); +} + +int UBShmTransport::PrepareServerUpgradeResources(ubring::SHM *remote_trx_shm, + ubring::SHM *local_trx_shm, + ubring::UbrDataFormat format) { + return _ub_ep->AllocateServerResources( + remote_trx_shm, local_trx_shm, format); +} + +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 +134,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 +154,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 +193,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 +217,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 +241,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..dc3ebd6383 100644 --- a/src/brpc/ubshm_transport.h +++ b/src/brpc/ubshm_transport.h @@ -21,12 +21,15 @@ #include "brpc/socket.h" #include "brpc/channel.h" #include "brpc/transport.h" +#include "brpc/ubshm/shm/shm_def.h" +#include "brpc/ubshm/ubr_data_format.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 +37,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 +46,25 @@ 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, + uint64_t remote_shm_len, + ubring::UbrDataFormat format); + int PrepareServerUpgradeResources(ubring::SHM* remote_trx_shm, + ubring::SHM* local_trx_shm, + ubring::UbrDataFormat format); + 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,8 +78,7 @@ 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 -#endif //BRPC_UB_TRANSPORT_H \ No newline at end of file +#endif //BRPC_UB_TRANSPORT_H diff --git a/src/brpc/urma/mock_urma.cpp b/src/brpc/urma/mock_urma.cpp index b0572c2e86..0adc00e27d 100644 --- a/src/brpc/urma/mock_urma.cpp +++ b/src/brpc/urma/mock_urma.cpp @@ -36,6 +36,8 @@ #include "urma_api.h" +#include + #include #include #include @@ -48,6 +50,13 @@ #include #include +// Opt-in contract checking for stream tests. The normal mock also serves raw +// message API tests, where unordered WRs are legal. +DEFINE_bool(urma_mock_require_send_order, false, + "Reject mock SEND WRs without stream placement/completion ordering"); +DEFINE_uint32(urma_mock_max_send_size, 0, + "Reject mock SEND WRs exceeding this total SGE size (0 disables)"); + namespace { struct JfcState { @@ -549,6 +558,27 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, read_lock.unlock(); for (urma_jfs_wr_t* current = wr; current; current = current->next) { + if (FLAGS_urma_mock_max_send_size != 0) { + uint64_t payload_size = 0; + for (uint32_t i = 0; i < current->send.src.num_sge; ++i) { + payload_size += current->send.src.sge[i].len; + } + if (payload_size > FLAGS_urma_mock_max_send_size) { + if (bad_wr) { + *bad_wr = current; + } + return URMA_EINVAL; + } + } + if (FLAGS_urma_mock_require_send_order && + (jetty->jetty_cfg.jfs_cfg.flag.bs.outorder_comp != 0 || + current->flag.bs.place_order != URMA_STRONG_ORDER || + current->flag.bs.comp_order != 1)) { + if (bad_wr) { + *bad_wr = current; + } + return URMA_EINVAL; + } if (!current->tjetty) { if (bad_wr) { *bad_wr = current; @@ -560,10 +590,12 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, send_cr.status = URMA_CR_SUCCESS; send_cr.user_ctx = current->user_ctx; send_cr.flag.bs.s_r = 0; - if (current->flag.bs.complete_enable) { - PushCompletion(jetty->jetty_cfg.jfs_cfg.jfc, - local_state, send_cr); - } + auto complete_send = [&]() { + if (current->flag.bs.complete_enable) { + PushCompletion(jetty->jetty_cfg.jfs_cfg.jfc, + local_state, send_cr); + } + }; PendingRecv recv{}; urma_jfc_t* remote_jfc = nullptr; @@ -571,6 +603,7 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, std::unique_lock lock(g_rw_mutex); auto remote_it = jetty_id_map.find(current->tjetty->id.id); if (remote_it == jetty_id_map.end()) { + complete_send(); continue; } urma_jfr_t* remote_jfr = @@ -579,6 +612,7 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, auto jfc_it = jfr_jfc_map.find(remote_jfr); if (recv_it == jfr_recv_map.end() || recv_it->second.empty() || jfc_it == jfr_jfc_map.end()) { + complete_send(); continue; } recv = recv_it->second.front(); @@ -596,6 +630,10 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, copied += n; } + // A send completion transfers buffer ownership back to the caller. + // Do not let a polling thread free/reuse it before memcpy finishes. + complete_send(); + JfcState* remote_state = nullptr; { std::shared_lock lock(g_rw_mutex); diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 26c7a1b9dc..9f8593a1ae 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -48,9 +48,10 @@ #include "brpc/input_messenger.h" #include "brpc/socket.h" #include "brpc/urma/urma_bonding.h" -#include "brpc/urma/urma_handshake.h" +#include "brpc/handshake/urma_handshake.h" #include "brpc/urma/urma_handshake.pb.h" #include "brpc/urma/urma_helper.h" +#include "brpc/urma/urma_stream_trace.h" #include "brpc/urma_transport.h" DECLARE_int32(task_group_ntags); @@ -66,13 +67,21 @@ DECLARE_int32(urma_cqe_poll_once); DECLARE_bool(urma_recv_zerocopy); DECLARE_int32(urma_zerocopy_min_size); DECLARE_int32(urma_prepared_jetty_cnt); +DECLARE_int32(urma_bonding_mode); DECLARE_bool(urma_poller_yield); +DEFINE_bool(urma_trace_stream, false, + "Capture the last 32 URMA stream events per connection and dump on " + "parse/completion/post failure or teardown. Includes payload prefixes; " + "enable before creating connections only for controlled diagnostics."); // ---- Constants shared with the handshake module ---- static const int WAIT_TIMEOUT_MS = 50; static const size_t HELLO_ACK_LEN = 4; static const uint32_t HELLO_ACK_URMA_OK = 0x1; static const size_t IOBUF_BLOCK_HEADER_LEN = sizeof(butil::IOBuf::Block); +// Match UMDK perftest's CTP SEND ceiling. This is a per-WR payload limit, +// independent of the pool block size and the device's generic max_msg_size. +static const size_t MAX_CTP_SEND_SIZE = 4096; // ---- Globals: prepared jetty pool + poller groups ---- struct PreparedJetty { @@ -147,7 +156,8 @@ UrmaEndpoint::UrmaEndpoint(Socket* s) : _socket(s), _state(UNINIT), _handshake_version(0), - _resource(nullptr) { + _resource(nullptr), + _cq_lifetime(std::make_shared(s->id())) { _sq_size = static_cast( std::max(16, std::min(4096, static_cast(FLAGS_urma_sq_size)))); _rq_size = static_cast( @@ -155,6 +165,9 @@ UrmaEndpoint::UrmaEndpoint(Socket* s) _read_butex = bthread::butex_create_checked>(); _read_butex->store(0, butil::memory_order_relaxed); _input_processor.Init(s, InputMessengerProcessor::STREAM_URMA_JETTY); + if (FLAGS_urma_trace_stream) { + _stream_trace.reset(new UrmaStreamTrace); + } } UrmaEndpoint::~UrmaEndpoint() { @@ -167,6 +180,7 @@ UrmaEndpoint::~UrmaEndpoint() { void UrmaEndpoint::Reset() { DeallocateResources(); + _stream_trace.reset(FLAGS_urma_trace_stream ? new UrmaStreamTrace : nullptr); _state = UNINIT; _handshake_version = 0; _remote_recv_block_size = 0; @@ -184,6 +198,8 @@ void UrmaEndpoint::Reset() { _sbuf.clear(); _rbuf.clear(); _rbuf_data.clear(); + _recv_completions.clear(); + _recv_completed.clear(); _read_butex->store(0, butil::memory_order_relaxed); } @@ -451,6 +467,8 @@ int UrmaEndpoint::AllocateResources() { jetty_cfg.flag.bs.share_jfr = 1; jetty_cfg.jfs_cfg.depth = static_cast(_sq_size); jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + // WR placement/completion ordering requires an in-order JFS. + jetty_cfg.jfs_cfg.flag.bs.outorder_comp = 0; jetty_cfg.jfs_cfg.priority = GetUrmaJettyPriority(); jetty_cfg.jfs_cfg.max_sge = static_cast(GetUrmaMaxSge()); @@ -469,58 +487,85 @@ int UrmaEndpoint::AllocateResources() { _sbuf.resize(_sq_size - RESERVED_WR_NUM); _rbuf.resize(_rq_size); _rbuf_data.resize(_rq_size, nullptr); + _recv_completions.resize(_rq_size); + _recv_completed.assign(_rq_size, false); - // Wrap the JFCE fd in a brpc Socket so PollCq is driven by epoll. - if (!FLAGS_urma_use_polling) { + // Arm before advertising receives. Completions may queue during the TCP + // handshake, but must not be dispatched until its parser has returned. + if (!FLAGS_urma_use_polling && ReqNotifyCq() != 0) { + return -1; + } + return 0; +} + +int UrmaEndpoint::StartCqEvents() { + std::unique_lock lock(_cq_lifetime->mutex); + if (_socket->parsing_stream_type() != InputMessengerProcessor::STREAM_NONE) { + errno = EPROTO; + return -1; + } + if (_cq_sid != INVALID_SOCKET_ID) { + return 0; + } + if (!_resource || !_resource->jfc) { + errno = ENODEV; + return -1; + } + const bool polling = FLAGS_urma_use_polling; + if (!polling) { if (!_resource->jfce || _resource->jfce->fd < 0) { LOG(ERROR) << "Prepared URMA resource has no usable JFCE"; errno = ENODEV; return -1; } - if (ReqNotifyCq() != 0) { - return -1; - } - SocketOptions options; - options.user = this; - options.keytable_pool = _socket->keytable_pool(); - options.fd = _resource->jfce->fd; - options.on_edge_triggered_events = PollCq; - if (Socket::Create(options, &_cq_sid) < 0) { - PLOG(ERROR) << "Fail to create CQ socket"; - return -1; - } - } else { - // Polling mode: synthetic carrier socket (no fd). - SocketOptions options; - options.user = this; - options.keytable_pool = _socket->keytable_pool(); - options.on_edge_triggered_events = PollCq; - if (Socket::Create(options, &_cq_sid) < 0) { - PLOG(ERROR) << "Fail to create CQ socket (polling)"; - return -1; - } + } + // Install the owned user only after Create succeeds, and before publishing + // the fd/id to a dispatcher. Create failures then cannot leak/double-delete + // the user or close the provider-owned JFCE fd. + SocketOptions options; + options.keytable_pool = _socket->keytable_pool(); + options.on_edge_triggered_events = PollCq; + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(ERROR) << "Fail to create CQ socket"; + return -1; + } + SocketUniquePtr carrier; + CHECK_EQ(0, Socket::Address(_cq_sid, &carrier)); + carrier->_user = new CqSocketUser(_cq_lifetime, this, polling); + if (!polling && carrier->ResetFileDescriptor(_resource->jfce->fd) != 0) { + const int saved_errno = errno; + carrier->SetFailed(); + _cq_sid = INVALID_SOCKET_ID; + errno = saved_errno; + return -1; + } + _cq_lifetime->closed = false; + if (polling) { PollerAddCqSid(); } return 0; } void UrmaEndpoint::DeallocateResources() { + std::unique_lock lock(_cq_lifetime->mutex); + _cq_lifetime->closed = true; + ++_cq_lifetime->generation; if (!_resource) { return; } + DumpStreamTrace("connection teardown"); - if (FLAGS_urma_use_polling) { - PollerRemoveCqSid(); - } - - // Tear down the CQ socket so the EventDispatcher stops calling PollCq. + // epoll DEL and poller REMOVE do not drain callbacks. The lifetime mutex + // waits for an active callback; the closed generation rejects queued ones. if (_cq_sid != INVALID_SOCKET_ID) { SocketUniquePtr s; - if (Socket::Address(_cq_sid, &s) == 0) { + if (Socket::AddressFailedAsWell(_cq_sid, &s) >= 0) { + if (static_cast(s->user())->polling) { + PollerRemoveCqSid(); + } if (s->fd() >= 0) { s->_io_event.RemoveConsumer(s->_fd); } - s->_user = nullptr; // Do not release user (this UrmaEndpoint). s->_fd = -1; // Already removed fd from epoll. s->SetFailed(); } @@ -580,11 +625,13 @@ int UrmaEndpoint::ImportPeer(const ParsedHello& peer) { urma_token_t token{}; const bool use_bonding_extension = - IsUrmaBondingDevice() && remote.trans_mode == URMA_TM_RM; + NeedsUrmaBondingJettyExtension(IsUrmaBondingDevice(), + FLAGS_urma_bonding_mode) && + remote.trans_mode == URMA_TM_RM; errno = 0; if (use_bonding_extension) { #if BRPC_URMA_HAS_BONDING_EXT - // The bonding provider needs the local jetty to associate its send + // Non-standalone bonding needs the local jetty to associate its send // path with the imported target. A plain import may return success // without setting that association, leaving traffic one-way only. bondp_rjetty_t bonding_remote{}; @@ -664,6 +711,7 @@ class UrmaIOBuf : private butil::IOBuf { sglist[*sge_index].addr = reinterpret_cast(start); sglist[*sge_index].len = static_cast(this_len); sglist[*sge_index].tseg = tseg; + sglist[*sge_index].user_tseg = nullptr; cutn(to, this_len); len += this_len; (*sge_index)++; @@ -693,7 +741,7 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { ssize_t total_len = 0; while (current < ndata) { uint16_t remote_wnd = _remote_rq_window_size.load(butil::memory_order_relaxed); - uint16_t sq_wnd = _sq_window_size.load(butil::memory_order_relaxed); + uint16_t sq_wnd = _sq_window_size.load(butil::memory_order_acquire); if (remote_wnd == 0 || sq_wnd == 0) { if (total_len > 0) { break; @@ -707,6 +755,9 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { size_t max_len = _remote_recv_block_size > 0 ? _remote_recv_block_size : GetUrmaRecvBlockSize(); + // UrmaTransport advertises CTP in both hello versions. Limit the sum + // of all SGEs, not each SGE, while retaining larger receive buffers. + max_len = std::min(max_len, MAX_CTP_SEND_SIZE); while (sge_index < static_cast(max_sge) && this_len < max_len && current < ndata) { auto* data = reinterpret_cast(from[current]); @@ -733,6 +784,13 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { // Piggybacking credits turns every payload into SEND_IMM and can // produce asymmetric completions with the bonding provider. wr.opcode = URMA_OPC_SEND; + // A WR may end in the middle of an RPC (e.g. a 229-byte prefix + // followed by a 2-byte payload). RM reliability alone is not a byte + // stream ordering contract: zero place_order/comp_order permits the + // later WR to be observed first. Keep placement and completions in + // posting order, including the intervening credit-only SEND_IMM WRs. + wr.flag.bs.place_order = URMA_STRONG_ORDER; + wr.flag.bs.comp_order = 1; wr.flag.bs.complete_enable = 1; wr.tjetty = _resource->remote_jetty; wr.send.src = sg; @@ -741,6 +799,26 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { const uint16_t sq_slot = _sq_current; const uint32_t local_jetty_id = _resource->jetty->jetty_id.id; const uint32_t remote_jetty_id = _resource->remote_jetty->id.id; + for (size_t i = 0; i < sge_index; ++i) { + const urma_sge_t& sge = sglist[i]; + const uint64_t seg_base = sge.tseg->seg.ubva.va; + const uint64_t seg_len = sge.tseg->seg.len; + const bool in_range = sge.addr >= seg_base && + sge.len <= seg_len && sge.addr - seg_base <= seg_len - sge.len; + VLOG(1) << "URMA send SGE: index=" << i + << " addr=" << reinterpret_cast(sge.addr) + << " len=" << sge.len + << " tseg=" << static_cast(sge.tseg) + << " seg_base=" << reinterpret_cast(seg_base) + << " seg_len=" << seg_len + << " in_range=" << in_range + << " from_pool=" << (GetPoolSegFor( + reinterpret_cast(sge.addr)) == sge.tseg) + << " context_matches=" + << (sge.tseg->urma_ctx == GetUrmaContext()) + << " local_jetty_id=" << local_jetty_id + << " sq_slot=" << sq_slot; + } // Reserve both credits before making the WR visible to the provider. // In polling mode a completion (and even the peer's receive-credit @@ -750,6 +828,10 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { // connection. _remote_rq_window_size.fetch_sub(1, butil::memory_order_relaxed); _sq_window_size.fetch_sub(1, butil::memory_order_relaxed); + if (_stream_trace) { + _stream_trace->Record(UrmaStreamTrace::TX, sq_slot, wr.user_ctx, 0, + sglist, sge_index); + } int rc = urma_post_jetty_send_wr(_resource->jetty, &wr, &bad_wr); if (rc != URMA_SUCCESS) { const int provider_errno = errno; @@ -771,6 +853,7 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { << ", payload_size=" << this_len << " on " << _socket->description(); errno = rc; + DumpStreamTrace("data post failed"); return -1; } _sq_current = (_sq_current + 1) % (_sq_size - RESERVED_WR_NUM); @@ -781,7 +864,7 @@ ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { bool UrmaEndpoint::IsWritable() const { return _remote_rq_window_size.load(butil::memory_order_relaxed) > 0 && - _sq_window_size.load(butil::memory_order_relaxed) > 0; + _sq_window_size.load(butil::memory_order_acquire) > 0; } // ============================================================================ @@ -797,7 +880,7 @@ int UrmaEndpoint::DoPostRecv(void* block, size_t block_size) { urma_sge_t sge{reinterpret_cast(block), static_cast(block_size), tseg, nullptr}; urma_sg_t sg{&sge, 1}; - urma_jfr_wr_t wr{sg, 0, nullptr}; + urma_jfr_wr_t wr{sg, static_cast(_rq_received) + 1, nullptr}; urma_jfr_wr_t* bad = nullptr; // Use the shared-JFR path on every device, including bonding. The bonding // provider owns physical receive scheduling for the JFR; a local @@ -880,6 +963,10 @@ int UrmaEndpoint::SendImm(uint32_t imm) { urma_jfs_wr_t wr{}; std::memset(&wr, 0, sizeof(wr)); wr.opcode = URMA_OPC_SEND_IMM; + // ACKs consume the same receive queue as data. They must participate in + // its ordering contract rather than allowing a later payload to overtake. + wr.flag.bs.place_order = URMA_STRONG_ORDER; + wr.flag.bs.comp_order = 1; wr.flag.bs.complete_enable = 1; wr.flag.bs.solicited_enable = 1; wr.tjetty = _resource->remote_jetty; @@ -890,6 +977,9 @@ int UrmaEndpoint::SendImm(uint32_t imm) { // data windows in CutFromIOBufList: polling may observe its completion as // soon as the provider accepts the WR. --_sq_imm_window_size; + if (_stream_trace) { + _stream_trace->Record(UrmaStreamTrace::TX, 0, wr.user_ctx, imm, nullptr, 0); + } const urma_status_t status = urma_post_jetty_send_wr(_resource->jetty, &wr, &bad); if (status != URMA_SUCCESS) { @@ -909,6 +999,7 @@ int UrmaEndpoint::SendImm(uint32_t imm) { << " state=" << GetStateStr() << " on " << _socket->description(); errno = status; + DumpStreamTrace("credit ACK post failed"); return -1; } return 0; @@ -925,9 +1016,14 @@ int UrmaEndpoint::SendAck(int num) { } ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) { - bool zerocopy = FLAGS_urma_recv_zerocopy; if (cr.status != URMA_CR_SUCCESS) { - LOG(WARNING) << "URMA completion failed, status=" << cr.status; + LOG(WARNING) << "URMA completion failed, status=" << cr.status + << " kind=" << (cr.flag.bs.s_r == 0 ? "send" : "recv") + << " user_ctx=" << cr.user_ctx + << " completion_len=" << cr.completion_len + << " local_id=" << cr.local_id + << " state=" << GetStateStr() + << " on " << _socket->description(); errno = EIO; return -1; } @@ -950,37 +1046,76 @@ ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) { return 0; } uint16_t wnd = 1; // We signal every WR (complete_enable=1). - uint16_t old = + const uint16_t old = _sq_window_size.load(butil::memory_order_relaxed); - while (true) { - if (old >= _local_window_capacity) { - LOG(WARNING) - << "URMA send completion exceeds SQ window: old=" << old - << " increment=" << wnd - << " capacity=" << _local_window_capacity - << " user_ctx=" << cr.user_ctx - << " on " << _socket->description(); - errno = EPROTO; - return -1; - } - if (_sq_window_size.compare_exchange_weak( - old, static_cast(old + wnd), - butil::memory_order_relaxed)) { - break; - } + if (old >= _local_window_capacity) { + LOG(WARNING) + << "URMA send completion exceeds SQ window: old=" << old + << " increment=" << wnd + << " capacity=" << _local_window_capacity + << " user_ctx=" << cr.user_ctx + << " on " << _socket->description(); + errno = EPROTO; + return -1; } for (uint16_t i = 0; i < wnd; ++i) { _sbuf[_sq_sent].clear(); _sq_sent = (_sq_sent + 1) % (_sq_size - RESERVED_WR_NUM); } - butil::subtle::MemoryBarrier(); + // Only the serialized CQ consumer returns credits. Reclaim the slot + // before publishing it to the writer; its acquire load must observe + // clear() before appending another payload to the same IOBuf. + _sq_window_size.fetch_add(wnd, butil::memory_order_release); if (_remote_rq_window_size.load(butil::memory_order_relaxed) >= _local_window_capacity / 8) { _socket->WakeAsEpollOut(); } return 0; } - // Recv completion. + // Identify the posted buffer rather than assuming CQ order equals RQ + // order. Otherwise an ACK completion can advance the ring past a data WR + // whose completion is still pending, corrupting the RPC byte stream. + if (cr.user_ctx == 0 || cr.user_ctx > _rq_size) { + LOG(WARNING) << "Invalid URMA receive slot: " << cr.user_ctx; + errno = EPROTO; + return -1; + } + const size_t slot = cr.user_ctx - 1; + if (_stream_trace && cr.completion_len <= GetUrmaRecvBlockSize()) { + urma_sge_t sge{reinterpret_cast(_rbuf_data[slot]), + cr.completion_len, nullptr, nullptr}; + _stream_trace->Record(UrmaStreamTrace::RX_CQ, slot, cr.user_ctx, + cr.opcode == URMA_CR_OPC_SEND_WITH_IMM ? cr.imm_data : 0, + &sge, 1); + } + if (_recv_completed[slot]) { + LOG(WARNING) << "Duplicate URMA receive completion: slot=" << slot; + errno = EPROTO; + return -1; + } + _recv_completions[slot] = cr; + _recv_completed[slot] = true; + if (slot != _rq_received) { + VLOG(1) << "URMA deferred receive completion: slot=" << slot + << " expected_slot=" << _rq_received + << " completion_len=" << cr.completion_len + << " on " << _socket->description(); + } + ssize_t bytes = 0; + while (_recv_completed[_rq_received]) { + const urma_cr_t next = _recv_completions[_rq_received]; + _recv_completed[_rq_received] = false; + const ssize_t received = HandleRecvCompletion(next); + if (received < 0) { + return -1; + } + bytes += received; + } + return bytes; +} + +ssize_t UrmaEndpoint::HandleRecvCompletion(const urma_cr_t& cr) { + bool zerocopy = FLAGS_urma_recv_zerocopy; if (cr.opcode == URMA_CR_OPC_SEND_WITH_IMM && cr.imm_data > 0) { if (cr.imm_data > _local_window_capacity) { LOG(WARNING) << "Invalid URMA receive credit: " << cr.imm_data; @@ -1027,6 +1162,14 @@ ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) { zerocopy = false; } butil::IOPortal& read_buf = _input_processor.read_buf(); + if (_stream_trace) { + urma_sge_t sge{reinterpret_cast(_rbuf_data[_rq_received]), + cr.completion_len, nullptr, nullptr}; + _stream_trace->Record(UrmaStreamTrace::RX_USE, _rq_received, + cr.user_ctx, + cr.opcode == URMA_CR_OPC_SEND_WITH_IMM ? cr.imm_data : 0, + &sge, 1); + } if (zerocopy) { _rbuf[_rq_received].cutn(&read_buf, cr.completion_len); } else { @@ -1075,24 +1218,53 @@ void UrmaEndpoint::DispatchReceivedBytes(SocketUniquePtr& s, ssize_t bytes) { const int64_t received_us = butil::cpuwide_time_us(); const int64_t base_realtime = butil::gettimeofday_us() - received_us; InputMessageClosure last_msg; - _input_processor.ProcessNewMessage(static_cast(pending), false, - received_us, base_realtime, last_msg); + const int result = _input_processor.ProcessNewMessage( + static_cast(pending), false, received_us, base_realtime, last_msg); + if (result < 0 && _stream_trace) { + char prefix[UrmaStreamTrace::PREFIX_SIZE]; + const size_t copied = _input_processor.read_buf().copy_to(prefix, sizeof(prefix)); + DumpStreamTrace("RPC parse failed", UrmaStreamTrace::HexPrefix(prefix, copied), + _input_processor.read_buf().size()); + } +} + +void UrmaEndpoint::DumpStreamTrace(const char* reason, const std::string& prefix, + size_t buffered) { + if (!_stream_trace) { + return; + } + const int saved_errno = errno; + const std::string records = _stream_trace->DumpOnce(); + if (!records.empty()) { + LOG(WARNING) << "URMA_TRACE reason=" << reason + << " buffered=" << buffered << " parse_prefix=" << prefix + << " on " << _socket->description() << '\n' << records; + } + errno = saved_errno; } void UrmaEndpoint::PollCq(Socket* m) { - auto* ep = static_cast(m->user()); - if (!ep || !ep->_resource || !ep->_resource->jfc) { + // The caller pins the carrier, not the endpoint. Its user/context remains + // alive even when the main Socket has already been recycled. + auto* user = static_cast(m->user()); + if (!user) { return; } + const auto lifetime = user->lifetime; SocketUniquePtr s; - if (Socket::Address(ep->_socket->id(), &s) != 0) { + if (Socket::Address(lifetime->socket_id, &s) != 0) { return; } - if (s->Failed()) { + // Declare the lock after s: unlock before releasing the last main Socket + // reference, which can synchronously call DeallocateResources again. + std::unique_lock lock(lifetime->mutex); + if (lifetime->closed || user->generation != lifetime->generation || + s->Failed()) { return; } + auto* ep = user->endpoint; - const bool event_mode = !FLAGS_urma_use_polling; + const bool event_mode = !user->polling; int progress = Socket::PROGRESS_INIT; while (true) { urma_jfc_t* event_jfc = nullptr; @@ -1158,6 +1330,7 @@ void UrmaEndpoint::PollCq(Socket* m) { } if (completion_error != 0) { + ep->DumpStreamTrace("completion processing failed"); if (!s->Failed()) { s->SetFailed(completion_error, "URMA completion error"); } @@ -1385,7 +1558,10 @@ void* UrmaEndpoint::ProcessHandshakeAtClient(void* arg) { } tp->_urma_state = UrmaTransport::URMA_ON; ep->_state = ESTABLISHED; - ep->DispatchReceivedBytes(s, 0); + if (ep->StartCqEvents() < 0) { + ep->FailHandshake(tp, errno, "start client CQ events"); + return nullptr; + } return nullptr; } @@ -1470,7 +1646,10 @@ void* UrmaEndpoint::ProcessHandshakeAtServer(void* arg) { } tp->_urma_state = UrmaTransport::URMA_ON; ep->_state = ESTABLISHED; - ep->DispatchReceivedBytes(s, 0); + if (ep->StartCqEvents() < 0) { + ep->FailHandshake(tp, errno, "start server CQ events"); + return nullptr; + } } else { ep->FallbackToTcp(tp, true); } @@ -1808,6 +1987,7 @@ int UrmaEndpoint::GlobalInitialize() { jetty_cfg.flag.bs.share_jfr = 1; jetty_cfg.jfs_cfg.depth = static_cast(FLAGS_urma_sq_size); jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.flag.bs.outorder_comp = 0; jetty_cfg.jfs_cfg.priority = GetUrmaJettyPriority(); jetty_cfg.jfs_cfg.max_sge = static_cast(GetUrmaMaxSge()); diff --git a/src/brpc/urma/urma_endpoint.h b/src/brpc/urma/urma_endpoint.h index e5aed5ac66..1deb641d75 100644 --- a/src/brpc/urma/urma_endpoint.h +++ b/src/brpc/urma/urma_endpoint.h @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -28,6 +29,7 @@ #include "butil/iobuf.h" #include "butil/macros.h" #include "butil/synchronization/lock.h" +#include "bthread/mutex.h" #include "bthread/types.h" #include "brpc/socket.h" @@ -36,7 +38,7 @@ #include "urma_api.h" #include "urma_types.h" -#include "brpc/urma/urma_handshake.h" +#include "brpc/handshake/urma_handshake.h" #include "brpc/urma/urma_handshake.pb.h" namespace brpc { @@ -46,15 +48,15 @@ class UrmaTransport; namespace urma { class UrmaHandshake; +class UrmaStreamTrace; struct ParsedHello; DECLARE_bool(urma_use_polling); DECLARE_int32(urma_poller_num); DECLARE_bool(urma_disable_bthread); -// Per-connection application-level connect object. Returned by -// UrmaTransport::Connect(); its StartConnect spawns the client-side handshake -// bthread that drives the URMA negotiation over the already-connected TCP fd. +// Legacy URMA connection handshake helper, retained for interoperability tests. +// AdapterTransport manages the normal client upgrade path. class UrmaConnect : public AppConnect { friend class UrmaEndpoint; public: @@ -96,13 +98,13 @@ struct UrmaResource { DISALLOW_COPY_AND_ASSIGN(UrmaResource); }; -// One per Socket. Carries the handshake state machine, the send/recv -// windows, and the registered buffer bookkeeping. Cache-line padded to avoid -// false sharing between connections. +// Per-Socket URMA resources, send/recv windows and registered buffers. +// AdapterTransport drives the common handshake; the legacy handshake state +// and helpers remain here for interoperability tests. Cache-line aligned. class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { friend class UrmaConnect; friend class Socket; - friend class UrmaTransport; + friend class ::brpc::UrmaTransport; friend class UrmaHandshakeClientV2; friend class UrmaHandshakeServerV2; friend class UrmaHandshakeClientV3; @@ -135,8 +137,8 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { void DebugInfo(std::ostream& os, butil::StringPiece connector = "\n") const; - // Edge-triggered callback installed on the TCP socket when the transport - // is URMA. Dispatches on _state (see the .cpp for the full state machine): + // Legacy TCP handshake callback used by interoperability tests. + // Dispatches on _state (see the .cpp for the full state machine): // UNINIT (server) -> start the server handshake bthread // < ESTABLISHED -> wake the handshake bthread parked in ReadFromFd // FALLBACK_TCP -> hand off to InputMessenger::OnNewMessages @@ -144,7 +146,7 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { static void OnNewDataFromTcp(Socket* m); // CQ completion poller: the edge-triggered callback for the CQ socket - // (whose user is this endpoint). Drains urma_poll_jfc and feeds the + // (whose user owns a CQ lifetime context). Drains urma_poll_jfc and feeds the // completions into HandleCompletion, then calls InputMessenger. Event mode // also drains the aggregated JFCE until urma_wait_jfc reports no event, // then calls Socket::MoreReadEvents so later JFCE edges can start the @@ -158,7 +160,7 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { std::function release_fn); static void PollingModeRelease(bthread_tag_t tag); - // ---- Handshake IO helpers (also used by urma_handshake.cpp) ---- + // ---- Legacy handshake IO helpers ---- // Read at most @len bytes from the TCP fd into @data; waits on _read_butex // on EAGAIN. Returns 0 on success, -1 on IO error (errno set). int ReadFromFd(void* data, size_t len); @@ -203,13 +205,15 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { FAILED = 0x300 }; - // Process handshake at the client / server (run in a background bthread). + // Legacy client/server handshake bthreads. static void* ProcessHandshakeAtClient(void* arg); static void* ProcessHandshakeAtServer(void* arg); // Allocate / deallocate the per-connection URMA resources (JFCE/JFC/JFR/ // Jetty, plus the CQ socket). Returns 0 on success. int AllocateResources(); + // Register CQ callbacks only after the TCP handshake parser has returned. + int StartCqEvents(); void DeallocateResources(); // Import the peer's jetty and buffer-pool segment. CRITICAL: @@ -239,13 +243,14 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { // received // (0 for send completions), or -1 on error (errno set). ssize_t HandleCompletion(const urma_cr_t& cr); + ssize_t HandleRecvCompletion(const urma_cr_t& cr); - // Queue received bytes for InputMessenger. CQ receive completions can - // arrive while the server is still waiting for the final TCP handshake - // ACK. Keep those bytes in the URMA input processor and dispatch them only - // after ESTABLISHED, matching the connection-ready boundary seen by user - // code. + // Queue received bytes for InputMessenger after ESTABLISHED. Hardware can + // complete receives during the handshake, but StartCqEvents defers their + // processing until the TCP parser has returned. void DispatchReceivedBytes(SocketUniquePtr& s, ssize_t bytes); + void DumpStreamTrace(const char* reason, const std::string& prefix = "", + size_t buffered = 0); // Consume one async event from the JFCE fd (event mode only). The caller // drains completions before acknowledge/rearm because the bonding provider @@ -280,9 +285,30 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { // The URMA resources (jetty / jfc / jfr / jfce / imported peer objects). UrmaResource* _resource{nullptr}; - - // The SocketId wrapping the JFCE fd (event mode) or a synthetic carrier - // (polling mode). PollCq is the edge-triggered callback on this socket. + std::unique_ptr _stream_trace; + + struct CqLifetime { + explicit CqLifetime(SocketId id) : socket_id(id) {} + const SocketId socket_id; + bthread::Mutex mutex; // Serializes CQ dispatch, start and teardown. + uint64_t generation{0}; + bool closed{true}; + }; + struct CqSocketUser : SocketUser { + CqSocketUser(const std::shared_ptr& lifetime, + UrmaEndpoint* ep, bool use_polling) + : lifetime(lifetime), endpoint(ep), + generation(lifetime->generation), polling(use_polling) {} + void BeforeRecycle(Socket*) override { delete this; } + const std::shared_ptr lifetime; + UrmaEndpoint* const endpoint; // Use only with main Socket pinned + mutex. + const uint64_t generation; + const bool polling; + }; + // Shared with old carriers so queued callbacks can reject a closed/reset + // generation without dereferencing an endpoint that may already be gone. + const std::shared_ptr _cq_lifetime; + // Wraps the JFCE fd (event mode) or a synthetic carrier (polling mode). SocketId _cq_sid{INVALID_SOCKET_ID}; // ---- Send / recv window bookkeeping ---- @@ -297,6 +323,11 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { // Per-WR recv buffers (zero-copy targets). std::vector _rbuf; std::vector _rbuf_data; + // A provider may report completions in a different order from posted + // receive WRs. Keep each completion with its owning slot and consume the + // stream only once all preceding receive WRs have completed. + std::vector _recv_completions; + std::vector _recv_completed; // Peer's advertised recv buffer size (caps each WR's payload). uint32_t _remote_recv_block_size{0}; diff --git a/src/brpc/urma/urma_helper.cpp b/src/brpc/urma/urma_helper.cpp index ded662fb22..883067aa75 100644 --- a/src/brpc/urma/urma_helper.cpp +++ b/src/brpc/urma/urma_helper.cpp @@ -675,6 +675,9 @@ const urma_eid_t* GetUrmaLocalEid() { return g_has_local_eid ? &g_local_eid : nullptr; } bool IsUrmaBondingDevice() { return g_is_bonding_device; } +bool NeedsUrmaBondingJettyExtension(bool bonding_device, int bonding_mode) { + return bonding_device && (bonding_mode == 1 || bonding_mode == 2); +} int FindUrmaPriorityForTpType(const urma_device_attr_t& attr, urma_tp_type_t tp_type) { const union urma_tp_type_en expected = TpTypeCapability(tp_type); diff --git a/src/brpc/urma/urma_helper.h b/src/brpc/urma/urma_helper.h index 871f1a6354..6df790d73d 100644 --- a/src/brpc/urma/urma_helper.h +++ b/src/brpc/urma/urma_helper.h @@ -75,6 +75,10 @@ const urma_eid_t* GetUrmaLocalEid(); // Return true when the selected URMA device is a bonding provider device. bool IsUrmaBondingDevice(); +// Standalone bonding uses the regular remote-jetty import layout. +// Mode values match urma_bonding_mode: 0=standalone, 1=active_backup, 2=balance. +bool NeedsUrmaBondingJettyExtension(bool bonding_device, int bonding_mode); + // Find the priority whose advertised transport-path capability exactly // matches @tp_type. Returns -1 when the device does not report one. int FindUrmaPriorityForTpType(const urma_device_attr_t& attr, diff --git a/src/brpc/urma/urma_stream_trace.h b/src/brpc/urma/urma_stream_trace.h new file mode 100644 index 0000000000..97ad0ad3e5 --- /dev/null +++ b/src/brpc/urma/urma_stream_trace.h @@ -0,0 +1,142 @@ +// 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_URMA_STREAM_TRACE_H +#define BRPC_URMA_STREAM_TRACE_H + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "butil/synchronization/lock.h" +#include "urma_types.h" + +namespace brpc { +namespace urma { + +// Diagnostic-only, non-cryptographic fingerprints. No protocol bytes change. +class UrmaStreamTrace { +public: + enum Stage { TX, RX_CQ, RX_USE, STAGE_COUNT }; + static constexpr size_t CAPACITY = 32; + static constexpr size_t PREFIX_SIZE = 16; + + struct Summary { + uint64_t hash = UINT64_C(14695981039346656037); + size_t length = 0; + std::string prefix; + }; + + static std::string HexPrefix(const void* data, size_t size) { + const auto* bytes = static_cast(data); + static const char digits[] = "0123456789abcdef"; + std::string hex; + const size_t count = size < PREFIX_SIZE ? size : PREFIX_SIZE; + hex.reserve(count * 2); + for (size_t i = 0; i < count; ++i) { + hex.push_back(digits[bytes[i] >> 4]); + hex.push_back(digits[bytes[i] & 15]); + } + return hex; + } + + static Summary Summarize(const urma_sge_t* sges, size_t count) { + Summary summary; + for (size_t i = 0; i < count; ++i) { + const auto* bytes = + reinterpret_cast(sges[i].addr); + for (uint32_t j = 0; j < sges[i].len; ++j) { + summary.hash ^= bytes[j]; + summary.hash *= UINT64_C(1099511628211); + } + if (summary.length < PREFIX_SIZE) { + const size_t room = PREFIX_SIZE - summary.length; + const size_t prefix_len = std::min(sges[i].len, room); + summary.prefix += HexPrefix(bytes, prefix_len); + } + summary.length += sges[i].len; + } + return summary; + } + + void Record(Stage stage, size_t slot, uint64_t user_ctx, uint64_t imm, + const urma_sge_t* sges, size_t count) { + Entry entry; + entry.stage = stage; + entry.slot = slot; + entry.user_ctx = user_ctx; + entry.imm = imm; + entry.sges = count; + entry.summary = Summarize(sges, count); + std::lock_guard lock(_mutex); + entry.sequence = _sequences[stage]++; + entry.offset = _offsets[stage]; + _offsets[stage] += entry.summary.length; + _entries[_written++ % CAPACITY] = entry; + } + + // One bounded dump per connection; Record is safe for concurrent TX/CQ. + std::string DumpOnce() { + std::lock_guard lock(_mutex); + if (_dumped || _written == 0) { + return std::string(); + } + _dumped = true; + const char* names[] = {"TX", "RX_CQ", "RX_USE"}; + std::ostringstream out; + const size_t first = _written > CAPACITY ? _written - CAPACITY : 0; + out << "records=" << _written - first << " total=" << _written; + for (size_t i = first; i < _written; ++i) { + const Entry& entry = _entries[i % CAPACITY]; + out << "\nURMA_TRACE stage=" << names[entry.stage] + << " seq=" << entry.sequence << " offset=" << entry.offset + << " len=" << entry.summary.length << " slot=" << entry.slot + << " user_ctx=" << entry.user_ctx << " imm=" << entry.imm + << " sges=" << entry.sges << " hash=" << std::hex + << entry.summary.hash << std::dec + << " prefix=" << entry.summary.prefix; + } + return out.str(); + } + +private: + struct Entry { + Stage stage = TX; + size_t slot = 0; + uint64_t user_ctx = 0; + uint64_t imm = 0; + size_t sges = 0; + uint64_t sequence = 0; + uint64_t offset = 0; + Summary summary; + }; + butil::Mutex _mutex; + std::array _entries; + std::array _sequences{}; + std::array _offsets{}; + size_t _written = 0; + bool _dumped = false; +}; + +} // namespace urma +} // namespace brpc +#endif // BRPC_URMA_STREAM_TRACE_H diff --git a/src/brpc/urma_transport.cpp b/src/brpc/urma_transport.cpp index c5cc163ffa..de485ab0d4 100644 --- a/src/brpc/urma_transport.cpp +++ b/src/brpc/urma_transport.cpp @@ -39,29 +39,24 @@ namespace brpc { DECLARE_bool(urma_use_polling); DECLARE_bool(urma_disable_bthread); +namespace urma { +DECLARE_bool(urma_recv_zerocopy); +} // namespace urma + void UrmaTransport::Init(Socket* socket, const SocketOptions& options) { CHECK(_urma_ep == nullptr); if (options.socket_mode == SOCKET_MODE_URMA) { _urma_ep = new (std::nothrow) urma::UrmaEndpoint(socket); if (!_urma_ep) { - const int saved_errno = errno; - PLOG(ERROR) << "Fail to create UrmaEndpoint"; - socket->SetFailed(saved_errno, "Fail to create UrmaEndpoint: %s", - berror(saved_errno)); + PLOG(WARNING) << "Fail to create UrmaEndpoint, fallback to TCP"; + _urma_state = URMA_OFF; + } else { + _urma_state = URMA_UNKNOWN; } - _urma_state = URMA_UNKNOWN; } else { _urma_state = URMA_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 = urma::UrmaEndpoint::OnNewDataFromTcp; - } - _tcp_transport = std::make_shared(); - _tcp_transport->Init(socket, options); } void UrmaTransport::Release() { @@ -80,36 +75,32 @@ int UrmaTransport::Reset(int32_t /*expected_nref*/) { } std::shared_ptr UrmaTransport::Connect() { - if (_default_connect == nullptr) { - return std::make_shared(); - } return _default_connect; } int UrmaTransport::CutFromIOBuf(butil::IOBuf* buf) { - // The URMA endpoint is not usable until negotiation and peer-resource - // import have completed. Keep using TCP while the state is UNKNOWN, and - // switch the data path only after the handshake publishes URMA_ON. - if (_urma_ep && - _urma_state.load(butil::memory_order_acquire) == URMA_ON) { - butil::IOBuf* data_arr[1] = {buf}; - return _urma_ep->CutFromIOBufList(data_arr, 1); - } else { - return _tcp_transport->CutFromIOBuf(buf); + // AdapterTransport selects TcpTransport until the handshake is established. + // Reaching this transport earlier is an error, not a second TCP fallback. + if (!_urma_ep || + _urma_state.load(butil::memory_order_acquire) != URMA_ON) { + errno = ENOTCONN; + return -1; } + butil::IOBuf* data_arr[1] = {buf}; + return _urma_ep->CutFromIOBufList(data_arr, 1); } ssize_t UrmaTransport::CutFromIOBufList(butil::IOBuf** buf, size_t ndata) { - if (_urma_ep && - _urma_state.load(butil::memory_order_acquire) == URMA_ON) { - return _urma_ep->CutFromIOBufList(buf, ndata); - } else { - return _tcp_transport->CutFromIOBufList(buf, ndata); + if (!_urma_ep || + _urma_state.load(butil::memory_order_acquire) != URMA_ON) { + errno = ENOTCONN; + return -1; } + return _urma_ep->CutFromIOBufList(buf, ndata); } int UrmaTransport::WaitEpollOut(butil::atomic* epollout_butex, - bool pollin, const timespec duetime) { + bool /*pollin*/, const timespec duetime) { if (_urma_state.load(butil::memory_order_acquire) == URMA_ON) { const int expected_val = epollout_butex->load(butil::memory_order_acquire); @@ -134,7 +125,8 @@ int UrmaTransport::WaitEpollOut(butil::atomic* epollout_butex, } return 0; } - return _tcp_transport->WaitEpollOut(epollout_butex, pollin, duetime); + errno = ENOTCONN; + return 1; } void UrmaTransport::ProcessEvent(bthread_attr_t attr) { @@ -190,6 +182,62 @@ void UrmaTransport::Debug(std::ostream& os) { } } +int UrmaTransport::PrepareUpgradeResources(bool server_side) { + if (!_urma_ep || _urma_ep->AllocateResources() < 0) { + return -1; + } + // A bonding provider must import the peer before receives are posted. + if ((!server_side || !urma::IsUrmaBondingDevice()) && + _urma_ep->PostRecv(_urma_ep->_rq_size, + urma::FLAGS_urma_recv_zerocopy) < 0) { + return -1; + } + return 0; +} + +int UrmaTransport::NegotiateUpgradeResources( + const urma::ParsedHello& remote, bool server_side) { + CHECK(_urma_ep != nullptr); + _urma_ep->ApplyRemoteHello(remote); + if (_urma_ep->ImportPeer(remote) < 0) { + return -1; + } + if (server_side && urma::IsUrmaBondingDevice() && + _urma_ep->PostRecv(_urma_ep->_rq_size, + urma::FLAGS_urma_recv_zerocopy) < 0) { + return -1; + } + return 0; +} + +void UrmaTransport::ActivateUpgrade() { + CHECK(_urma_ep != nullptr); + _urma_state.store(URMA_ON, butil::memory_order_release); + _urma_ep->_state.store(urma::UrmaEndpoint::ESTABLISHED, + butil::memory_order_release); +} + +int UrmaTransport::StartUpgradeEvents() { + return _urma_ep ? _urma_ep->StartCqEvents() : -1; +} + +void UrmaTransport::DeactivateUpgrade() { + _urma_state.store(URMA_OFF, butil::memory_order_release); + if (_urma_ep) { + _urma_ep->_state.store(urma::UrmaEndpoint::FALLBACK_TCP, + butil::memory_order_release); + _urma_ep->DeallocateResources(); + } +} + +void UrmaTransport::SetHighSpeedAvailable(bool available) { + if (available) { + ActivateUpgrade(); + } else { + DeactivateUpgrade(); + } +} + int UrmaTransport::ContextInitOrDie(bool server_or_not, const void* options) { if (server_or_not) { if (!OptionsAvailableOverUrma( diff --git a/src/brpc/urma_transport.h b/src/brpc/urma_transport.h index 9a496d56be..1c22a6ae24 100644 --- a/src/brpc/urma_transport.h +++ b/src/brpc/urma_transport.h @@ -27,19 +27,17 @@ namespace brpc { -// UrmaTransport is the Transport subclass for URMA (openEuler Unified Remote -// Memory Access). It composes a TcpTransport for fallback (mirroring the -// RdmaTransport / UBShmTransport design) and delegates the URMA data path to -// the per-connection urma::UrmaEndpoint. Negotiation runs over the TCP fd and -// resolves _urma_state to URMA_ON or URMA_OFF. +// UrmaTransport owns the URMA data path. AdapterTransport owns TCP fallback +// and selects this transport only after the handshake establishes URMA. class UrmaTransport : public Transport { - friend class TransportFactory; - friend class urma::UrmaEndpoint; - friend class urma::UrmaConnect; - friend class urma::UrmaHandshakeServerV2; - friend class urma::UrmaHandshakeClientV2; - friend class urma::UrmaHandshakeServerV3; - friend class urma::UrmaHandshakeClientV3; +friend class TransportFactory; +friend class AdapterTransport; +friend class urma::UrmaEndpoint; +friend class urma::UrmaConnect; +friend class urma::UrmaHandshakeServerV2; +friend class urma::UrmaHandshakeClientV2; +friend class urma::UrmaHandshakeServerV3; +friend class urma::UrmaHandshakeClientV3; public: void Init(Socket* socket, const SocketOptions& options) override; @@ -62,6 +60,15 @@ class UrmaTransport : public Transport { static int ContextInitOrDie(bool server_or_not, const void* options); + bool UpgradeReady() const { return _urma_ep != nullptr; } + int PrepareUpgradeResources(bool server_side); + int NegotiateUpgradeResources(const urma::ParsedHello& remote, + bool server_side); + void ActivateUpgrade(); + int StartUpgradeEvents(); + void DeactivateUpgrade(); + void SetHighSpeedAvailable(bool available); + private: static bool OptionsAvailableForUrma(const ChannelOptions* opt); static bool OptionsAvailableOverUrma(const ServerOptions* opt); @@ -75,7 +82,6 @@ class UrmaTransport : public Transport { urma::UrmaEndpoint* _urma_ep = nullptr; butil::atomic _urma_state{URMA_UNKNOWN}; - std::shared_ptr _tcp_transport; }; } // namespace brpc diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index a0fc2d0536..262a16e100 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -53,6 +53,9 @@ compile_proto(PROTO_HDRS PROTO_SRCS ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR} "${TEST_PROTO_FILES}") add_library(TEST_PROTO_LIB OBJECT ${PROTO_SRCS} ${PROTO_HDRS}) +# echo.proto imports src/idl_options.proto. Its generated header must exist +# before compiling the test protobuf sources, including in a clean build. +add_dependencies(TEST_PROTO_LIB PROTO_LIB) target_link_libraries(TEST_PROTO_LIB PRIVATE brpc_test_config) set(BRPC_SYSTEM_GTEST_SOURCE_DIR "" CACHE PATH "System googletest source directory.") 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 558c696649..e0ae27fb4d 100644 --- a/test/brpc_ubring_unittest.cpp +++ b/test/brpc_ubring_unittest.cpp @@ -28,7 +28,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" @@ -691,6 +712,57 @@ TEST(HelloFormatExtensionTest, selects_only_matching_supported_format) { SelectDataFormat(static_cast(0x1234), 0x1234)); } +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(UBShmHandshakeAdapterTest, ipc_v2_extension_matches_backend_format) { + using namespace brpc::ubring; + UBShmHandshakeAdapter protocol(UBR_DATA_FORMAT_IPC_V2); + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + protocol.BuildExtension(true, &payload)); + EXPECT_EQ(std::string("\0\4\0\2", 4), payload); + EXPECT_EQ(brpc::handshake::STEP_OK, protocol.ParseExtension(payload)); + + payload[3] = 1; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + protocol.ParseExtension(payload)); + payload[3] = 0; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + protocol.ParseExtension(payload)); + ASSERT_EQ(brpc::handshake::STEP_OK, + protocol.BuildExtension(false, &payload)); + EXPECT_EQ(std::string("\0\4\0\0", 4), payload); +} + +TEST(UBShmHandshakeAdapterTest, ipc_v2_extension_rejects_malformed_frame) { + using namespace brpc::ubring; + UBShmHandshakeAdapter protocol(UBR_DATA_FORMAT_IPC_V2); + EXPECT_EQ(brpc::handshake::STEP_ERROR, + protocol.ParseExtension(std::string(3, '\0'))); + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + protocol.ParseExtension(std::string("\0\3\0\2", 4))); + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + protocol.ParseExtension(std::string("\0\4\x12\x34", 4))); +} + TEST_F(HelloMessageTest, serialize_deserialize_roundtrip) { msg.msg_len = 64; msg.hello_ver = 2; @@ -781,6 +853,118 @@ 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, accepts_name_terminated_at_wire_boundary) { + 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 - 1); + std::string payload(sizeof(HelloMessageLayout), '\0'); + message.Serialize(&payload[0]); + + brpc::ubring::UBShmHandshakeAdapter adapter; + brpc::ubring::HelloMessage decoded{}; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.ParseHello(payload, &decoded)); + EXPECT_EQ(static_cast(SHM_MAX_NAME_BUFF_LEN - 1), + strlen(decoded.shm_name)); + EXPECT_EQ(0, memcmp(message.shm_name, decoded.shm_name, + SHM_MAX_NAME_BUFF_LEN)); +} + +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; @@ -863,16 +1047,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) { @@ -894,6 +1078,375 @@ 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( + brpc::ubring::PreferredDataFormatForShmType(brpc::ubring::GetShmType())); + 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(), server_hello.len, protocol.local_format())); + + 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()); + EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_IPC_V2, + brpc::UBShmTransport::Get(socket.get())->GetUBShmEp() + ->negotiated_data_format()); + } + 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 diff --git a/test/brpc_urma_fallback_unittest.cpp b/test/brpc_urma_fallback_unittest.cpp new file mode 100644 index 0000000000..0b19d56ea4 --- /dev/null +++ b/test/brpc_urma_fallback_unittest.cpp @@ -0,0 +1,177 @@ +// 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 "butil/fd_guard.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/policy/transport_handshake_protocol.h" +#include "brpc/socket.h" +#include "brpc/urma/urma_handshake.pb.h" + +using namespace brpc; + +// Keep these tests runnable without URMA SDK headers or a device. +class UrmaTcpFallbackTest : public ::testing::TestWithParam { +protected: + void SetUp() override { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + peer.reset(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + ASSERT_EQ(0, Socket::Address(id, &socket)); + } + void TearDown() override { socket->SetFailed(); } + + std::string Hello() const { + // A known legacy v2 packet, including its original version fields. + if (GetParam() == 2) { + std::string frame(86, '\0'); + frame.replace(0, 4, "URMA"); + const uint16_t fields[] = {butil::HostToNet16(86), + butil::HostToNet16(2), + butil::HostToNet16(1)}; + memcpy(&frame[4], fields, sizeof(fields)); + return frame; + } + urma::UrmaHello hello; + hello.set_buffer_size(8192); + hello.set_recv_buffer_cnt(127); + hello.set_jetty_id(7); + hello.set_eid(std::string(16, '\0')); + hello.set_uasid(0); + hello.set_tp_type(1); + hello.set_seg_eid(std::string(16, '\0')); + hello.set_seg_uasid(0); + hello.set_seg_va(4096); + hello.set_seg_len(8192); + hello.set_seg_token_id(0); + std::string payload; + EXPECT_TRUE(hello.SerializeToString(&payload)); + const uint32_t size = butil::HostToNet32(payload.size()); + std::string frame("URM3", 4); + frame.append(reinterpret_cast(&size), sizeof(size)); + return frame + payload; + } + + void CheckDisabledReply() { + char reply[256]; + const ssize_t nr = recv(peer, reply, sizeof(reply), MSG_DONTWAIT); + ASSERT_GT(nr, 8); + if (GetParam() == 2) { + ASSERT_EQ(86, nr); + EXPECT_EQ(0, memcmp(reply, "URMA", 4)); + EXPECT_EQ(std::string(80, '\0'), std::string(reply + 6, nr - 6)); + } else { + EXPECT_EQ(0, memcmp(reply, "URM3", 4)); + urma::UrmaHello hello; + ASSERT_TRUE(hello.ParseFromArray(reply + 8, nr - 8)); + EXPECT_EQ(0u, hello.jetty_id()); + EXPECT_EQ(0u, hello.buffer_size()); + EXPECT_EQ(16u, hello.eid().size()); + } + } + + SocketUniquePtr socket; + butil::fd_guard peer; +}; + +TEST_P(UrmaTcpFallbackTest, fragmented_legacy_hello_needs_no_ack) { + const std::string hello = Hello(); + butil::IOBuf source; + source.append(hello.data(), 3); + EXPECT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_EQ(3u, source.size()); + source.append(hello.data() + 3, hello.size() - 3); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_TRUE(source.empty()); + EXPECT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_EQ(nullptr, socket->parsing_context()); + CheckDisabledReply(); + // Legacy clients send RPC bytes immediately, without a four-byte ACK. + source.append("PRPC", 4); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_EQ("PRPC", source.to_string()); +} + +TEST_P(UrmaTcpFallbackTest, preserves_rpc_coalesced_after_rejected_hello) { + butil::IOBuf source; + source.append(Hello()); + source.append("PRPC", 4); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_EQ("PRPC", source.to_string()); + EXPECT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(socket.get())->handshake_phase()); + CheckDisabledReply(); +} + +TEST_P(UrmaTcpFallbackTest, rejects_bad_frame_length) { + std::string hello = Hello(); + memset(&hello[4], 0, GetParam() == 2 ? 2 : 4); + if (GetParam() == 3) { + // A 4097-byte body exceeds the maximum accepted v3 frame. + const uint32_t length = butil::HostToNet32(4097); + memcpy(&hello[4], &length, sizeof(length)); + } + butil::IOBuf source; + source.append(hello); + EXPECT_EQ(PARSE_ERROR_ABSOLUTELY_WRONG, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_EQ(handshake::FAILED, + AdapterTransport::Get(socket.get())->handshake_phase()); +} + +TEST_P(UrmaTcpFallbackTest, unsupported_hello_never_establishes_urma) { + std::string hello = Hello(); + if (GetParam() == 2) { + const uint16_t version = butil::HostToNet16(99); + memcpy(&hello[6], &version, sizeof(version)); + } else { + hello[8] = 0; // Invalid protobuf tag inside an otherwise valid frame. + } + butil::IOBuf source; + source.append(hello); + EXPECT_EQ(GetParam() == 2 ? PARSE_ERROR_TRY_OTHERS + : PARSE_ERROR_ABSOLUTELY_WRONG, + policy::ParseTransportHandshake(&source, socket.get(), false, + nullptr).error()); + EXPECT_EQ(GetParam() == 2 ? handshake::FALLBACK_TCP : handshake::FAILED, + AdapterTransport::Get(socket.get())->handshake_phase()); + if (GetParam() == 2) { + CheckDisabledReply(); + } +} + +INSTANTIATE_TEST_SUITE_P(WireVersions, UrmaTcpFallbackTest, + ::testing::Values(2, 3)); diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index d27367a504..dfcfa4730e 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -15,19 +15,45 @@ // 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 BRPC_WITH_URMA #include "butil/atomicops.h" +#include "butil/fd_guard.h" +#include "brpc/input_messenger.h" #include "butil/sys_byteorder.h" #include "urma_api.h" -#include "brpc/urma/urma_handshake.h" +#include "brpc/handshake/urma_handshake.h" +#include "brpc/handshake/handshake_frame.h" +#include "brpc/handshake/handshake_io.h" #include "brpc/urma/urma_handshake.pb.h" #include "brpc/urma/urma_helper.h" +#include "brpc/urma_transport.h" #include "urma_types.h" +#include "echo.pb.h" +#include "brpc/closure_guard.h" +#include "brpc/adapter_transport.h" +#include "brpc/acceptor.h" +#include "brpc/policy/transport_handshake_protocol.h" +#include "brpc/urma/urma_stream_trace.h" +#include "brpc/global.h" using namespace brpc; @@ -35,6 +61,7 @@ namespace brpc { namespace urma { DECLARE_int32(urma_client_handshake_version); +DECLARE_bool(urma_trace_stream); extern bool g_skip_urma_init; extern butil::atomic g_urma_available; @@ -139,58 +166,665 @@ TEST(UrmaHandshakeTest, v3_protobuf_roundtrip) { EXPECT_EQ(msg.seg_token_id(), msg2.seg_token_id()); } +TEST(UrmaHandshakeTest, common_adapter_parses_v2_and_v3_hello) { + urma::v2_wire::HelloMessage v2{}; + v2.msg_len = urma::v2_wire::HELLO_PACKET_LEN; + v2.hello_ver = urma::v2_wire::HELLO_V2_VERSION; + v2.impl_ver = urma::v2_wire::IMPL_V2_VERSION; + v2.buffer_size = 8192; + v2.recv_buffer_cnt = 127; + v2.jetty_id = 7; + v2.tp_type = URMA_CTP; + v2.seg_va = 0x1000; + v2.seg_len = 8192; + uint8_t v2_body[urma::v2_wire::HELLO_BODY_LEN]; + v2.Serialize(v2_body); + const std::string v2_payload( + reinterpret_cast(v2_body + sizeof(uint16_t)), + sizeof(v2_body) - sizeof(uint16_t)); + + urma::UrmaHandshakeAdapter v2_adapter(nullptr, 2); + ASSERT_EQ(handshake::STEP_OK, v2_adapter.ParseHello(v2_payload)); + EXPECT_EQ(7u, v2_adapter.remote().jetty_id); + EXPECT_EQ(urma::v2_wire::HELLO_PACKET_LEN, + v2_adapter.HelloFrameSpec().min_frame_len); + + urma::UrmaHello v3; + v3.set_buffer_size(8192); + v3.set_recv_buffer_cnt(127); + v3.set_jetty_id(9); + v3.set_eid(std::string(16, '\1')); + v3.set_uasid(1); + v3.set_tp_type(URMA_CTP); + v3.set_seg_eid(std::string(16, '\2')); + v3.set_seg_uasid(2); + v3.set_seg_va(0x2000); + v3.set_seg_len(8192); + v3.set_seg_token_id(3); + std::string v3_payload; + ASSERT_TRUE(v3.SerializeToString(&v3_payload)); + + urma::UrmaHandshakeAdapter v3_adapter(nullptr, 3); + ASSERT_EQ(handshake::STEP_OK, v3_adapter.ParseHello(v3_payload)); + EXPECT_EQ(9u, v3_adapter.remote().jetty_id); + EXPECT_EQ(handshake::FrameSpec::U32_BODY_LENGTH, + v3_adapter.HelloFrameSpec().length_encoding); +} + +TEST(UrmaHandshakeTest, common_adapter_rejects_oversized_v3_tp_type) { + urma::UrmaHello message; + message.set_buffer_size(8192); + message.set_recv_buffer_cnt(127); + message.set_jetty_id(9); + message.set_eid(std::string(16, '\1')); + message.set_uasid(1); + message.set_tp_type(256); + message.set_seg_eid(std::string(16, '\2')); + message.set_seg_uasid(2); + message.set_seg_va(0x2000); + message.set_seg_len(8192); + message.set_seg_token_id(3); + std::string payload; + ASSERT_TRUE(message.SerializeToString(&payload)); + + urma::UrmaHandshakeAdapter adapter(nullptr, 3); + EXPECT_EQ(handshake::STEP_FALLBACK, adapter.ParseHello(payload)); +} + +TEST(UrmaHandshakeTest, common_frames_preserve_fragmented_and_coalesced_data) { + for (int version = 2; version <= 3; ++version) { + urma::UrmaHandshakeAdapter adapter(nullptr, version); + std::string hello; + ASSERT_EQ(handshake::STEP_OK, adapter.BuildHello(false, &hello)); + std::string frame; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + adapter.HelloFrameSpec(), hello, &frame)); + + butil::IOBuf source; + source.append(frame.data(), 3); + handshake::IOBufHandshakeInput input(&source); + std::string parsed; + EXPECT_EQ(handshake::FRAME_NEED_MORE, + handshake::FrameCodec::ParseBufferedFrame( + &input, adapter.HelloFrameSpec(), &parsed)); + EXPECT_EQ(3u, source.size()); + + source.append(frame.data() + 3, frame.size() - 3); + source.append("RPC", 3); + ASSERT_EQ(handshake::FRAME_OK, + handshake::FrameCodec::ParseBufferedFrame( + &input, adapter.HelloFrameSpec(), &parsed)); + EXPECT_EQ(hello, parsed); + EXPECT_EQ(3u, source.size()); + EXPECT_EQ("RPC", source.to_string()); + } +} + +TEST(UrmaTransportTest, inactive_data_path_fails_without_tcp_delegate) { + UrmaTransport transport; + butil::IOBuf data; + butil::IOBuf* list[] = {&data}; + errno = 0; + EXPECT_EQ(-1, transport.CutFromIOBuf(&data)); + EXPECT_EQ(ENOTCONN, errno); + errno = 0; + EXPECT_EQ(-1, transport.CutFromIOBufList(list, 1)); + EXPECT_EQ(ENOTCONN, errno); + + butil::atomic epollout_butex(0); + const timespec duetime = {}; + errno = 0; + EXPECT_EQ(1, transport.WaitEpollOut(&epollout_butex, false, duetime)); + EXPECT_EQ(ENOTCONN, errno); +} + +class CapturingUrmaHandshakeIO : public handshake::HandshakeIO { +public: + int ReadExact(void*, size_t) override { return -1; } + int WriteAll(const void* data, size_t len) override { + output.append(static_cast(data), len); + return 0; + } + int PushBack(const void*, size_t) override { return -1; } + + std::string output; +}; + +class FailingUrmaResources : public handshake::HandshakeTransport { +public: + FailingUrmaResources() : version(0), fallback(false), failed(false) {} + void OnProtocolSelected(handshake::HandshakeProtocol* protocol) override { + version = protocol->ProtocolVersion(); + } + handshake::StepResult PrepareResources() override { + return handshake::STEP_FALLBACK; + } + handshake::StepResult NegotiateResources() override { + return handshake::STEP_ERROR; + } + void OnEstablished() override { ADD_FAILURE() << "Unexpected upgrade"; } + void OnFallback() override { fallback = true; } + void OnFailed() override { failed = true; } + + int version; + bool fallback; + bool failed; +}; + +TEST(UrmaHandshakeTest, common_server_fallback_preserves_coalesced_rpc) { + urma::v2_wire::HelloMessage hello{}; + hello.hello_ver = urma::v2_wire::HELLO_V2_VERSION; + hello.impl_ver = urma::v2_wire::IMPL_V2_VERSION; + hello.buffer_size = 8192; + hello.recv_buffer_cnt = 127; + hello.jetty_id = 7; + hello.tp_type = URMA_CTP; + hello.seg_va = 0x1000; + hello.seg_len = 8192; + uint8_t body[urma::v2_wire::HELLO_BODY_LEN]; + hello.Serialize(body); + const std::string payload( + reinterpret_cast(body + sizeof(uint16_t)), + sizeof(body) - sizeof(uint16_t)); + urma::UrmaHandshakeAdapter protocol(nullptr, 2); + std::string frame; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), payload, &frame)); + + butil::IOBuf source; + source.append(frame); + source.append("RPC", 3); + handshake::IOBufHandshakeInput input(&source); + CapturingUrmaHandshakeIO io; + handshake::HandshakeSession session; + session.SetIOForTest(&io); + FailingUrmaResources transport; + std::vector protocols(1, &protocol); + + EXPECT_EQ(handshake::STEP_FALLBACK, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ(handshake::FALLBACK_TCP, session.phase(butil::memory_order_acquire)); + EXPECT_EQ(2, transport.version); + EXPECT_TRUE(transport.fallback); + EXPECT_FALSE(transport.failed); + EXPECT_EQ("RPC", source.to_string()); + EXPECT_EQ(urma::v2_wire::HELLO_PACKET_LEN, io.output.size()); + EXPECT_EQ("URMA", io.output.substr(0, 4)); +} + +// Exercise the real URMA field codec and common server driver without device +// allocation. Only the local Hello and resource provider are test doubles. +class BufferedUrmaTestProtocol : public urma::UrmaHandshakeAdapter { +public: + explicit BufferedUrmaTestProtocol(int version) + : urma::UrmaHandshakeAdapter(nullptr, version) { + if (version == 2) { + urma::v2_wire::HelloMessage hello{}; + hello.msg_len = urma::v2_wire::HELLO_PACKET_LEN; + hello.hello_ver = urma::v2_wire::HELLO_V2_VERSION; + hello.impl_ver = urma::v2_wire::IMPL_V2_VERSION; + hello.buffer_size = 8192; + hello.recv_buffer_cnt = 127; + hello.jetty_id = 7; + hello.tp_type = URMA_CTP; + hello.seg_va = 0x1000; + hello.seg_len = 8192; + uint8_t body[urma::v2_wire::HELLO_BODY_LEN]; + hello.Serialize(body); + local_payload.assign( + reinterpret_cast(body + sizeof(uint16_t)), + sizeof(body) - sizeof(uint16_t)); + } else { + urma::UrmaHello hello; + hello.set_buffer_size(8192); + hello.set_recv_buffer_cnt(127); + hello.set_jetty_id(7); + hello.set_tp_type(URMA_CTP); + hello.set_eid(std::string(16, '\0')); + hello.set_uasid(0); + hello.set_seg_eid(std::string(16, '\0')); + hello.set_seg_uasid(0); + hello.set_seg_token_id(0); + hello.set_seg_va(0x1000); + hello.set_seg_len(8192); + hello.SerializeToString(&local_payload); + } + } + handshake::StepResult BuildHello(bool enabled, std::string* payload) override { + if (!enabled) { + return urma::UrmaHandshakeAdapter::BuildHello(false, payload); + } + *payload = local_payload; + return handshake::STEP_OK; + } + std::string local_payload; +}; + +class TrackedUrmaTestResources : public handshake::HandshakeTransport { +public: + butil::IOBuf* source = nullptr; + bool allocated = false; + bool active = false; + int cleanup_count = 0; + handshake::StepResult PrepareResources() override { + allocated = true; + return handshake::STEP_OK; + } + handshake::StepResult NegotiateResources() override { + return handshake::STEP_OK; + } + handshake::StepResult ValidateEstablished() override { + return source->empty() ? handshake::STEP_OK : handshake::STEP_ERROR; + } + void OnEstablished() override { active = true; } + void OnFallback() override { Cleanup(); } + void OnFailed() override { Cleanup(); } +private: + void Cleanup() { + allocated = false; + active = false; + ++cleanup_count; + } +}; + +class ScriptedUrmaHandshakeIO : public CapturingUrmaHandshakeIO { +public: + int ReadExact(void* data, size_t len) override { + if (input.size() < len) { + return -1; + } + memcpy(data, input.data(), len); + input.erase(0, len); + return 0; + } + std::string input; +}; + +class RejectingUrmaImport : public TrackedUrmaTestResources { +public: + handshake::StepResult NegotiateResources() override { + return handshake::STEP_FALLBACK; + } +}; + +class RejectingUrmaPreparation : public TrackedUrmaTestResources { +public: + handshake::StepResult PrepareResources() override { + allocated = true; + return handshake::STEP_FALLBACK; + } +}; + +TEST(UrmaHandshakeTest, server_resource_failure_preserves_legacy_rpc_without_ack) { + for (int version : {2, 3}) { + for (bool reject_prepare : {false, true}) { + BufferedUrmaTestProtocol protocol(version); + std::string hello; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), protocol.local_payload, &hello)); + butil::IOBuf source; + source.append(hello); + source.append("PRPC", 4); + handshake::IOBufHandshakeInput input(&source); + handshake::HandshakeSession session; + CapturingUrmaHandshakeIO io; + session.SetIOForTest(&io); + RejectingUrmaPreparation prepare; + RejectingUrmaImport negotiate; + TrackedUrmaTestResources* resources = reject_prepare + ? static_cast(&prepare) : &negotiate; + std::vector protocols(1, &protocol); + EXPECT_EQ(handshake::STEP_FALLBACK, + session.RunServer(protocols, &input, resources, false)); + EXPECT_EQ(handshake::FALLBACK_TCP, + session.phase(butil::memory_order_acquire)); + EXPECT_EQ("PRPC", source.to_string()); + EXPECT_FALSE(resources->allocated); + EXPECT_EQ(1, resources->cleanup_count); + std::string disabled; + ASSERT_EQ(handshake::STEP_OK, protocol.BuildHello(false, &disabled)); + std::string expected; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), disabled, &expected)); + EXPECT_EQ(expected, io.output); + } + } +} + +TEST(UrmaHandshakeTest, disabled_server_hello_does_not_send_ack) { + for (int version : {2, 3}) { + BufferedUrmaTestProtocol protocol(version); + std::string disabled; + ASSERT_EQ(handshake::STEP_OK, protocol.BuildHello(false, &disabled)); + ScriptedUrmaHandshakeIO io; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), disabled, &io.input)); + std::string expected; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), protocol.local_payload, &expected)); + handshake::HandshakeSession session; + session.SetIOForTest(&io); + TrackedUrmaTestResources resources; + EXPECT_EQ(handshake::STEP_FALLBACK, + session.RunClient(&protocol, &resources)); + EXPECT_EQ(expected, io.output); // No ACK before ordinary RPC bytes. + EXPECT_FALSE(resources.allocated); + EXPECT_EQ(1, resources.cleanup_count); + EXPECT_EQ(handshake::FALLBACK_TCP, + session.phase(butil::memory_order_acquire)); + } +} + +TEST(UrmaHandshakeTest, client_import_failure_still_sends_disabled_ack) { + for (int version : {2, 3}) { + BufferedUrmaTestProtocol protocol(version); + ScriptedUrmaHandshakeIO io; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + protocol.HelloFrameSpec(), protocol.local_payload, &io.input)); + const std::string expected = io.input + std::string(4, '\0'); + handshake::HandshakeSession session; + session.SetIOForTest(&io); + RejectingUrmaImport resources; + EXPECT_EQ(handshake::STEP_FALLBACK, + session.RunClient(&protocol, &resources)); + EXPECT_EQ(expected, io.output); + EXPECT_FALSE(resources.allocated); + EXPECT_EQ(1, resources.cleanup_count); + } +} + +class BufferedUrmaTestAdapter : public handshake::StandardHandshakeAdapter { +public: + explicit BufferedUrmaTestAdapter(int version) : protocol(version) { + session.SetIOForTest(&io); + } + BufferedUrmaTestProtocol protocol; + TrackedUrmaTestResources resources; + CapturingUrmaHandshakeIO io; + mutable handshake::HandshakeSession session; +protected: + handshake::StepResult RunServerStep(butil::IOBuf* source, Socket*) override { + resources.source = source; + handshake::IOBufHandshakeInput input(source); + std::vector protocols(1, &protocol); + return session.RunServer(protocols, &input, &resources, false); + } + handshake::HandshakeSession* GetSession(Socket*) const override { + return &session; + } +}; + +class UrmaCommonServerTest : public ::testing::Test { +protected: + void SetUp() override { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + peer_fd.reset(fds[1]); + SocketOptions options; + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + ASSERT_EQ(0, Socket::Address(id, &socket)); + } + void TearDown() override { + if (socket) { + socket->reset_parsing_context(nullptr); + socket->SetFailed(); + } + } + SocketUniquePtr socket; + butil::fd_guard peer_fd; +}; + +TEST_F(UrmaCommonServerTest, reclaims_send_buffer_before_publishing_sq_credit) { + urma::UrmaEndpoint endpoint(socket.get()); + const uint16_t capacity = 8; + endpoint._sq_size = capacity + urma::UrmaEndpoint::RESERVED_WR_NUM; + endpoint._local_window_capacity = capacity; + endpoint._sbuf.resize(capacity); + endpoint._sq_window_size.store(0, butil::memory_order_relaxed); + int reclaimed = 0; + urma_cr_t completion{}; + completion.status = URMA_CR_SUCCESS; + completion.flag.bs.s_r = 0; + completion.user_ctx = 1; + for (int round = 0; round < 4 * capacity; ++round) { + const uint16_t slot = round % capacity; + ASSERT_EQ(slot, endpoint._sq_sent); + ASSERT_EQ(0, endpoint._sbuf[slot].append_user_data( + new char('x'), 1, [&](void* data) { + // Deterministically inspect the publication boundary from + // inside clear(), before the slot can be reused by a writer. + EXPECT_EQ(0u, endpoint._sq_window_size.load( + butil::memory_order_acquire)); + ++reclaimed; + delete static_cast(data); + })); + ASSERT_EQ(0, endpoint.HandleCompletion(completion)); + EXPECT_TRUE(endpoint._sbuf[slot].empty()); + EXPECT_EQ((slot + 1) % capacity, endpoint._sq_sent); + EXPECT_EQ(1u, endpoint._sq_window_size.fetch_sub( + 1, butil::memory_order_acquire)); + } + EXPECT_EQ(4 * capacity, reclaimed); +} + +TEST_F(UrmaCommonServerTest, excess_send_completion_keeps_buffer_owned) { + urma::UrmaEndpoint endpoint(socket.get()); + endpoint._local_window_capacity = 8; + endpoint._sq_window_size.store(8, butil::memory_order_relaxed); + endpoint._sbuf.resize(8); + endpoint._sbuf[0].append("still owned"); + urma_cr_t completion{}; + completion.status = URMA_CR_SUCCESS; + completion.flag.bs.s_r = 0; + completion.user_ctx = 1; + errno = 0; + EXPECT_EQ(-1, endpoint.HandleCompletion(completion)); + EXPECT_EQ(EPROTO, errno); + EXPECT_EQ("still owned", endpoint._sbuf[0].to_string()); + EXPECT_EQ(0u, endpoint._sq_sent); + EXPECT_EQ(8u, endpoint._sq_window_size.load(butil::memory_order_acquire)); +} + +TEST_F(UrmaCommonServerTest, cq_events_cannot_start_inside_tcp_parser) { + urma::UrmaEndpoint endpoint(socket.get()); + socket->set_parsing_stream_type(InputMessengerProcessor::STREAM_TCP_FD); + EXPECT_EQ(-1, endpoint.StartCqEvents()); + EXPECT_EQ(EPROTO, errno); + EXPECT_EQ(INVALID_SOCKET_ID, endpoint._cq_sid); + socket->set_parsing_stream_type(InputMessengerProcessor::STREAM_NONE); +} + +TEST_F(UrmaCommonServerTest, rejects_unidentified_receive_completion) { + urma::UrmaEndpoint endpoint(socket.get()); + endpoint._rq_size = 4; + endpoint._recv_completions.resize(4); + endpoint._recv_completed.assign(4, false); + urma_cr_t cr{}; + cr.status = URMA_CR_SUCCESS; + cr.flag.bs.s_r = 1; + cr.opcode = URMA_CR_OPC_SEND; + cr.completion_len = 4; + for (uint64_t slot : {0u, 5u}) { + cr.user_ctx = slot; + EXPECT_EQ(-1, endpoint.HandleCompletion(cr)); + EXPECT_EQ(EPROTO, errno); + EXPECT_EQ(0u, endpoint._rq_received); + EXPECT_TRUE(endpoint._input_processor.read_buf().empty()); + } +} + +TEST_F(UrmaCommonServerTest, resumes_fragmented_hello_and_ack_without_rpc_handoff) { + for (int version : {2, 3}) { + BufferedUrmaTestAdapter adapter(version); + std::string hello; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + adapter.protocol.HelloFrameSpec(), adapter.protocol.local_payload, + &hello)); + butil::IOBuf source; + source.append(hello.data(), 3); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_FALSE(adapter.resources.allocated); + source.append(hello.data() + 3, hello.size() - 3); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + ASSERT_NE(nullptr, socket->parsing_context()); + EXPECT_TRUE(adapter.resources.allocated); + const uint32_t ack = butil::HostToNet32(0x1); + source.append(&ack, 2); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_FALSE(adapter.resources.active); + source.append(reinterpret_cast(&ack) + 2, 2); + ASSERT_EQ(PARSE_ERROR_NOT_ENOUGH_DATA, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_EQ(handshake::ESTABLISHED, + adapter.session.phase(butil::memory_order_acquire)); + EXPECT_TRUE(adapter.resources.active); + EXPECT_EQ(0, adapter.resources.cleanup_count); + EXPECT_EQ(nullptr, socket->parsing_context()); + EXPECT_TRUE(source.empty()); + } +} + +TEST_F(UrmaCommonServerTest, rejects_coalesced_tcp_data_and_cleans_failed_resources) { + for (int version : {2, 3}) { + BufferedUrmaTestAdapter adapter(version); + std::string hello; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + adapter.protocol.HelloFrameSpec(), adapter.protocol.local_payload, + &hello)); + butil::IOBuf source; + source.append(hello); + const uint32_t ack = butil::HostToNet32(0x1); + source.append(&ack, sizeof(ack)); + source.append("RPC", 3); + EXPECT_EQ(PARSE_ERROR_ABSOLUTELY_WRONG, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_EQ(handshake::FAILED, + adapter.session.phase(butil::memory_order_acquire)); + EXPECT_FALSE(adapter.resources.allocated); + EXPECT_FALSE(adapter.resources.active); + EXPECT_EQ(1, adapter.resources.cleanup_count); + EXPECT_EQ("RPC", source.to_string()); + } +} + +TEST_F(UrmaCommonServerTest, disabled_ack_cleans_resources_and_preserves_rpc) { + for (int version : {2, 3}) { + BufferedUrmaTestAdapter adapter(version); + std::string hello; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + adapter.protocol.HelloFrameSpec(), adapter.protocol.local_payload, + &hello)); + butil::IOBuf source; + source.append(hello); + const uint32_t ack = 0; + source.append(&ack, sizeof(ack)); + source.append("RPC", 3); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, + adapter.ExecuteServerHandshake(&source, socket.get()).error()); + EXPECT_EQ(handshake::FALLBACK_TCP, + adapter.session.phase(butil::memory_order_acquire)); + EXPECT_FALSE(adapter.resources.allocated); + EXPECT_FALSE(adapter.resources.active); + EXPECT_EQ(1, adapter.resources.cleanup_count); + EXPECT_EQ("RPC", source.to_string()); + } +} + +class UrmaUpgradeTestMessenger : public InputMessenger { +public: + using InputMessenger::OnNewMessagesUntil; + BufferedUrmaTestAdapter adapter{3}; + int peer_fd = -1; + int parse_calls = 0; + int write_result = -1; + static ParseResult Parse(butil::IOBuf* source, Socket* socket, + bool, const void*) { + auto* messenger = static_cast(socket->user()); + ++messenger->parse_calls; + ParseResult result = messenger->adapter.ExecuteServerHandshake( + source, socket); + if (messenger->adapter.resources.active) { + 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) == + handshake::ESTABLISHED; + } +}; + +TEST(UrmaHandshakeTest, tcp_read_loop_stops_after_urma_upgrade) { + UrmaUpgradeTestMessenger messenger; + InputMessageHandler handler{}; + handler.parse = UrmaUpgradeTestMessenger::Parse; + handler.process = UrmaUpgradeTestMessenger::Process; + handler.name = "test_urma_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)); + std::string hello; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + messenger.adapter.protocol.HelloFrameSpec(), + messenger.adapter.protocol.local_payload, &hello)); + const uint32_t ack = butil::HostToNet32(0x1); + hello.append(reinterpret_cast(&ack), sizeof(ack)); + ASSERT_EQ(static_cast(hello.size()), + write(peer_fd, hello.data(), hello.size())); + messenger.OnNewMessagesUntil(socket.get(), + UrmaUpgradeTestMessenger::StopReading); + EXPECT_EQ(1, messenger.parse_calls); + EXPECT_EQ(3, messenger.write_result); + EXPECT_TRUE(messenger.adapter.resources.active); + char remaining[3]; + ASSERT_EQ(3, recv(socket->fd(), remaining, sizeof(remaining), MSG_DONTWAIT)); + EXPECT_EQ(0, memcmp(remaining, "RPC", 3)); + socket->SetFailed(); +} + // --------------------------------------------------------------------------- -// CreateServerHandshakeByMagic dispatches on the magic bytes. +// Common server adapters support both URMA wire versions. // --------------------------------------------------------------------------- -TEST(UrmaHandshakeTest, server_handshake_factory_dispatches_on_magic) { - // We cannot fully exercise the server handshake without a real socket + - // endpoint, but we can verify the factory returns the right protocol - // version for each magic, and nullptr for an unknown magic. - uint8_t magic_v2[4] = {'U', 'R', 'M', 'A'}; - uint8_t magic_v3[4] = {'U', 'R', 'M', '3'}; - uint8_t magic_bad[4] = {'P', 'R', 'P', 'C'}; - - // v2 magic -> protocol version 2 - urma::UrmaHandshake* hs2 = - urma::CreateServerHandshakeByMagic(nullptr, magic_v2); - // Note: the factory dereferences the endpoint only inside SendLocalHello / - // ReceiveAndParseRemoteHello; passing nullptr is safe for the version query. - // (We delete immediately to avoid touching the endpoint.) - if (hs2) { - EXPECT_EQ(2, hs2->ProtocolVersion()); - delete hs2; - } - // v3 magic -> protocol version 3 - urma::UrmaHandshake* hs3 = - urma::CreateServerHandshakeByMagic(nullptr, magic_v3); - if (hs3) { - EXPECT_EQ(3, hs3->ProtocolVersion()); - delete hs3; - } - // unknown magic -> nullptr (caller falls back to TCP) - urma::UrmaHandshake* hsb = - urma::CreateServerHandshakeByMagic(nullptr, magic_bad); - EXPECT_EQ(nullptr, hsb); +TEST(UrmaHandshakeTest, common_server_factory_supports_both_wire_versions) { + auto adapters = urma::CreateServerHandshakeAdapters(nullptr); + ASSERT_EQ(2u, adapters.size()); + for (size_t i = 0; i < adapters.size(); ++i) { + ASSERT_NE(nullptr, adapters[i]); + EXPECT_EQ(static_cast(i) + 2, adapters[i]->ProtocolVersion()); + const auto& spec = adapters[i]->HelloFrameSpec(); + ASSERT_EQ(4u, spec.magic_len); + EXPECT_EQ(i == 0 ? "URMA" : "URM3", + std::string(spec.magic, spec.magic_len)); + } } // --------------------------------------------------------------------------- -// CreateClientHandshake picks the version from the gflag. +// The common client factory picks the version from the gflag. // --------------------------------------------------------------------------- -TEST(UrmaHandshakeTest, client_handshake_factory_respects_flag) { +TEST(UrmaHandshakeTest, common_client_factory_respects_flag) { GFLAGS_NAMESPACE::FlagSaver flag_saver; - urma::FLAGS_urma_client_handshake_version = 2; - urma::UrmaHandshake* hs2 = urma::CreateClientHandshake(nullptr); - if (hs2) { - EXPECT_EQ(2, hs2->ProtocolVersion()); - delete hs2; - } - - urma::FLAGS_urma_client_handshake_version = 3; - urma::UrmaHandshake* hs3 = urma::CreateClientHandshake(nullptr); - if (hs3) { - EXPECT_EQ(3, hs3->ProtocolVersion()); - delete hs3; + for (int version : {2, 3, 0, 99}) { + urma::FLAGS_urma_client_handshake_version = version; + auto adapter = urma::CreateClientHandshakeAdapter(nullptr); + ASSERT_NE(nullptr, adapter); + EXPECT_EQ(version == 3 ? 3 : 2, adapter->ProtocolVersion()); } } @@ -253,6 +887,80 @@ TEST(UrmaHandshakeTest, rejects_invalid_resource_and_window_values) { EXPECT_FALSE(urma::ValidHello(hello)); } +TEST(UrmaTraceTest, fingerprints_ignore_sge_boundaries) { + const std::string payload = "hello"; + urma_sge_t whole{reinterpret_cast(payload.data()), 5, nullptr, nullptr}; + urma_sge_t parts[] = { + {reinterpret_cast(payload.data()), 2, nullptr, nullptr}, + {reinterpret_cast(payload.data() + 2), 3, nullptr, nullptr}}; + const auto a = urma::UrmaStreamTrace::Summarize(&whole, 1); + const auto b = urma::UrmaStreamTrace::Summarize(parts, 2); + EXPECT_EQ(UINT64_C(0xa430d84680aabd0b), a.hash); + EXPECT_EQ(a.hash, b.hash); + EXPECT_EQ(a.length, b.length); + EXPECT_EQ("68656c6c6f", a.prefix); + EXPECT_EQ(a.prefix, b.prefix); +} + +TEST(UrmaTraceTest, bounds_payload_prefix_and_dump_history) { + const std::string payload(64, 'a'); + urma_sge_t sge{reinterpret_cast(payload.data()), 64, nullptr, nullptr}; + const auto summary = urma::UrmaStreamTrace::Summarize(&sge, 1); + EXPECT_EQ(64u, summary.length); + EXPECT_EQ(2 * urma::UrmaStreamTrace::PREFIX_SIZE, summary.prefix.size()); + urma::UrmaStreamTrace trace; + EXPECT_TRUE(trace.DumpOnce().empty()); + for (size_t i = 0; i < urma::UrmaStreamTrace::CAPACITY + 8; ++i) { + trace.Record(urma::UrmaStreamTrace::TX, i, 1, 0, &sge, 1); + } + const std::string dump = trace.DumpOnce(); + EXPECT_NE(std::string::npos, dump.find("records=32 total=40")); + EXPECT_EQ(std::string::npos, dump.find(" seq=7 ")); + EXPECT_NE(std::string::npos, dump.find(" seq=8 offset=512 len=64")); + EXPECT_NE(std::string::npos, dump.find(" seq=39 offset=2496 len=64")); + EXPECT_TRUE(trace.DumpOnce().empty()); +} + +TEST(UrmaTraceTest, concurrent_directions_keep_independent_offsets) { + const std::string payload = "xy"; + urma_sge_t sge{reinterpret_cast(payload.data()), 2, nullptr, nullptr}; + urma::UrmaStreamTrace trace; + auto record = [&](urma::UrmaStreamTrace::Stage stage) { + for (size_t i = 0; i < 50; ++i) { + trace.Record(stage, i, 1, 0, &sge, 1); + } + }; + std::thread sender(record, urma::UrmaStreamTrace::TX); + std::thread receiver(record, urma::UrmaStreamTrace::RX_CQ); + sender.join(); + receiver.join(); + const std::string dump = trace.DumpOnce(); + EXPECT_NE(std::string::npos, dump.find("records=32 total=100")); + EXPECT_NE(std::string::npos, dump.find("seq=49 offset=98 len=2")); +} + +TEST_F(UrmaCommonServerTest, stream_trace_is_opt_in) { + const bool saved = urma::FLAGS_urma_trace_stream; + urma::FLAGS_urma_trace_stream = false; + urma::UrmaEndpoint disabled(socket.get()); + urma::FLAGS_urma_trace_stream = true; + urma::UrmaEndpoint enabled(socket.get()); + urma::FLAGS_urma_trace_stream = saved; + EXPECT_EQ(nullptr, disabled._stream_trace.get()); + EXPECT_NE(nullptr, enabled._stream_trace.get()); +} + +TEST(UrmaHelperTest, standalone_bonding_uses_regular_jetty_import) { + EXPECT_FALSE(urma::NeedsUrmaBondingJettyExtension(true, 0)); + EXPECT_TRUE(urma::NeedsUrmaBondingJettyExtension(true, 1)); + EXPECT_TRUE(urma::NeedsUrmaBondingJettyExtension(true, 2)); + for (int mode = 0; mode <= 2; ++mode) { + EXPECT_FALSE(urma::NeedsUrmaBondingJettyExtension(false, mode)); + } + EXPECT_FALSE(urma::NeedsUrmaBondingJettyExtension(true, -1)); + EXPECT_FALSE(urma::NeedsUrmaBondingJettyExtension(true, 3)); +} + TEST(UrmaHelperTest, selects_priority_matching_transport_path_type) { urma_device_attr_t attr{}; attr.dev_cap.priority_info[3].tp_type.bs.rtp = 1; @@ -316,6 +1024,1231 @@ class UrmaMockTest : public ::testing::Test { bool _owns_urma_init{false}; }; +TEST_F(UrmaMockTest, send_completion_releases_source_only_after_copy) { + int device_count = 0; + urma_device_t** devices = urma_get_device_list(&device_count); + ASSERT_NE(nullptr, devices); + urma_context_t* raw_context = urma_create_context(devices[0], 0); + urma_free_device_list(devices); + std::unique_ptr context( + raw_context, urma_delete_context); + ASSERT_NE(nullptr, context.get()); + urma::UrmaResource resource; + resource.jfce = urma_create_jfce(context.get()); + ASSERT_NE(nullptr, resource.jfce); + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = 8; + jfc_cfg.jfce = resource.jfce; + resource.jfc = urma_create_jfc(context.get(), &jfc_cfg); + ASSERT_NE(nullptr, resource.jfc); + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = 4; + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.jfc = resource.jfc; + resource.jfr = urma_create_jfr(context.get(), &jfr_cfg); + ASSERT_NE(nullptr, resource.jfr); + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = 4; + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.max_sge = 1; + jetty_cfg.jfs_cfg.jfc = resource.jfc; + jetty_cfg.shared.jfr = resource.jfr; + jetty_cfg.shared.jfc = resource.jfc; + resource.jetty = urma_create_jetty(context.get(), &jetty_cfg); + ASSERT_NE(nullptr, resource.jetty); + urma_rjetty_t remote{}; + remote.jetty_id = resource.jetty->jetty_id; + remote.trans_mode = URMA_TM_RM; + remote.type = URMA_JETTY; + remote.tp_type = URMA_CTP; + urma_token_t token{}; + resource.remote_jetty = urma_import_jetty(context.get(), &remote, &token); + ASSERT_NE(nullptr, resource.remote_jetty); + + const std::string expected(1024 * 1024, 'x'); + std::string payload(expected); + std::string received(expected.size(), '\0'); + for (int round = 0; round < 8; ++round) { + payload = expected; + urma_sge_t receive_sge{reinterpret_cast(&received[0]), + static_cast(received.size()), nullptr, nullptr}; + urma_jfr_wr_t receive_wr{{&receive_sge, 1}, 0, nullptr}; + urma_jfr_wr_t* bad_receive = nullptr; + ASSERT_EQ(URMA_SUCCESS, urma_post_jfr_wr( + resource.jfr, &receive_wr, &bad_receive)); + urma_sge_t send_sge{reinterpret_cast(&payload[0]), + static_cast(payload.size()), nullptr, nullptr}; + urma_jfs_wr_t send_wr{}; + send_wr.opcode = URMA_OPC_SEND; + send_wr.flag.bs.complete_enable = 1; + send_wr.tjetty = resource.remote_jetty; + send_wr.send.src = {&send_sge, 1}; + std::atomic ready(false); + std::atomic stop(false); + int completions = 0; + std::thread poller([&]() { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(5); + ready.store(true, std::memory_order_release); + while (completions < 2 && !stop.load(std::memory_order_relaxed) && + std::chrono::steady_clock::now() < deadline) { + urma_cr_t completion{}; + if (urma_poll_jfc(resource.jfc, 1, &completion) == 1) { + if (completion.flag.bs.s_r == 0) { + // A real caller may reuse the source immediately + // after observing the send completion. + memset(&payload[0], 'z', payload.size()); + } + ++completions; + } else { + std::this_thread::yield(); + } + } + }); + while (!ready.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + urma_jfs_wr_t* bad_send = nullptr; + const urma_status_t status = urma_post_jetty_send_wr( + resource.jetty, &send_wr, &bad_send); + if (status != URMA_SUCCESS) { + stop.store(true, std::memory_order_relaxed); + } + poller.join(); + ASSERT_EQ(URMA_SUCCESS, status); + ASSERT_EQ(2, completions); + EXPECT_EQ(0, memcmp(expected.data(), received.data(), expected.size())); + } +} + +TEST_F(UrmaMockTest, segmented_stream_survives_send_queue_wraparound) { + int device_count = 0; + urma_device_t** devices = urma_get_device_list(&device_count); + ASSERT_NE(nullptr, devices); + urma_context_t* raw_context = urma_create_context(devices[0], 0); + urma_free_device_list(devices); + std::unique_ptr context( + raw_context, urma_delete_context); + ASSERT_NE(nullptr, context.get()); + + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer(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)); + struct FailSocket { + Socket* socket; + ~FailSocket() { socket->SetFailed(); } + } fail_socket{socket.get()}; + urma::UrmaEndpoint endpoint(socket.get()); + endpoint._resource = new urma::UrmaResource; + auto* resource = endpoint._resource; + resource->jfce = urma_create_jfce(context.get()); + ASSERT_NE(nullptr, resource->jfce); + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = 32; + jfc_cfg.jfce = resource->jfce; + resource->jfc = urma_create_jfc(context.get(), &jfc_cfg); + ASSERT_NE(nullptr, resource->jfc); + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = 8; + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.jfc = resource->jfc; + resource->jfr = urma_create_jfr(context.get(), &jfr_cfg); + ASSERT_NE(nullptr, resource->jfr); + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = 8; + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.max_sge = 13; + jetty_cfg.jfs_cfg.jfc = resource->jfc; + jetty_cfg.shared.jfr = resource->jfr; + jetty_cfg.shared.jfc = resource->jfc; + resource->jetty = urma_create_jetty(context.get(), &jetty_cfg); + ASSERT_NE(nullptr, resource->jetty); + urma_rjetty_t remote{}; + remote.jetty_id = resource->jetty->jetty_id; + remote.trans_mode = URMA_TM_RM; + remote.type = URMA_JETTY; + urma_token_t token{}; + resource->remote_jetty = urma_import_jetty(context.get(), &remote, &token); + ASSERT_NE(nullptr, resource->remote_jetty); + + const uint16_t capacity = 5; + const size_t block_size = 128; + endpoint._sq_size = capacity + urma::UrmaEndpoint::RESERVED_WR_NUM; + endpoint._local_window_capacity = capacity; + endpoint._sbuf.resize(capacity); + endpoint._remote_recv_block_size = block_size; + std::vector> receive_buffers( + capacity, std::vector(block_size)); + auto post_receive = [&](size_t slot) { + urma_sge_t sge{reinterpret_cast(receive_buffers[slot].data()), + static_cast(block_size), nullptr, nullptr}; + urma_jfr_wr_t wr{{&sge, 1}, slot + 1, nullptr}; + urma_jfr_wr_t* bad = nullptr; + return urma_post_jfr_wr(resource->jfr, &wr, &bad); + }; + for (size_t slot = 0; slot < capacity; ++slot) { + ASSERT_EQ(URMA_SUCCESS, post_receive(slot)); + } + + // The link-time mock does not DMA, so a synthetic registered segment is + // sufficient. Each part owns separate storage to force real SGE boundaries. + urma_target_seg_t segment{}; + segment.seg.len = std::numeric_limits::max(); + for (int round = 0; round < 64; ++round) { + endpoint._sq_window_size.store(capacity, butil::memory_order_relaxed); + endpoint._remote_rq_window_size.store(capacity, butil::memory_order_relaxed); + std::vector messages(64); + std::vector list; + std::string expected; + for (size_t i = 0; i < messages.size(); ++i) { + for (const std::string& part : { + std::string("PRPC:") + std::to_string(round) + ":" + + std::to_string(i) + std::string(i % 37, 'x'), + std::string("\x08\0", 2)}) { + char* data = new char[part.size()]; + memcpy(data, part.data(), part.size()); + ASSERT_EQ(0, messages[i].append_user_data_with_meta( + data, part.size(), [](void* p) { delete[] static_cast(p); }, + reinterpret_cast(&segment))); + expected.append(part); + } + list.push_back(&messages[i]); + } + std::string received; + size_t sent = 0; + while (sent < expected.size()) { + const ssize_t bytes = endpoint.CutFromIOBufList(list.data(), list.size()); + ASSERT_GT(bytes, 0); + sent += bytes; + urma_cr_t cr{}; + while (urma_poll_jfc(resource->jfc, 1, &cr) == 1) { + if (cr.flag.bs.s_r == 0) { + ASSERT_EQ(0, endpoint.HandleCompletion(cr)); + } else { + ASSERT_GE(cr.user_ctx, 1u); + ASSERT_LE(cr.user_ctx, capacity); + const size_t slot = cr.user_ctx - 1; + received.append(receive_buffers[slot].data(), cr.completion_len); + ASSERT_EQ(URMA_SUCCESS, post_receive(slot)); + endpoint._remote_rq_window_size.fetch_add( + 1, butil::memory_order_relaxed); + } + } + } + ASSERT_EQ(expected, received) << "round=" << round; + } +} + +class UrmaStreamEchoService : public test::EchoService { +public: + void Echo(google::protobuf::RpcController* controller, + const test::EchoRequest* request, test::EchoResponse* response, + google::protobuf::Closure* done) override { + ClosureGuard guard(done); + response->set_message(request->message()); + auto* cntl = static_cast(controller); + cntl->response_attachment().append(cntl->request_attachment()); + calls.fetch_add(1, std::memory_order_release); + } + std::atomic calls{0}; +}; + +// Use the retained legacy state machines/codecs, not a second common adapter +// or a synthesized Hello. Only transport installation is reconstructed here: +// production TransportFactory now always creates AdapterTransport. +static void InstallLegacyUrmaTransport(Socket* socket, SocketUser* user) { + socket->_transport->Release(); + SocketOptions options; + options.socket_mode = SOCKET_MODE_URMA; + options.user = user; + options.need_on_edge_trigger = true; + auto transport = std::unique_ptr(new UrmaTransport); + transport->Init(socket, options); + transport->_on_edge_trigger = urma::UrmaEndpoint::OnNewDataFromTcp; + transport->_default_connect = std::make_shared(); + socket->_app_connect = transport->_default_connect; + socket->_transport = std::move(transport); +} + +static int CreateUrmaInteropTcpPair(int fds[2]) { + butil::fd_guard listener(socket(AF_INET, SOCK_STREAM, 0)); + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + socklen_t size = sizeof(address); + if (listener < 0 || bind(listener, reinterpret_cast(&address), size) != 0 || + listen(listener, 1) != 0 || + getsockname(listener, reinterpret_cast(&address), &size) != 0) { + return -1; + } + butil::fd_guard client(socket(AF_INET, SOCK_STREAM, 0)); + if (client < 0 || connect(client, reinterpret_cast(&address), size) != 0) { + return -1; + } + butil::fd_guard accepted(accept(listener, nullptr, nullptr)); + if (accepted < 0) { + return -1; + } + fds[0] = client.release(); + fds[1] = accepted.release(); + return 0; +} + +class UrmaFinalAckGate { +public: + void Pause(Socket* socket) { + std::unique_lock lock(mutex); + tcp_parser_active = socket->parsing_stream_type() == + InputMessengerProcessor::STREAM_TCP_FD; + entered = true; + cv.notify_all(); + if (!cv.wait_for(lock, std::chrono::seconds(3), [&] { return released; })) { + timed_out = true; + } + } + bool Wait() { + std::unique_lock lock(mutex); + return cv.wait_for(lock, std::chrono::seconds(3), [&] { return entered; }); + } + void Release() { + std::lock_guard lock(mutex); + released = true; + cv.notify_all(); + } + std::mutex mutex; + std::condition_variable cv; + bool entered = false; + bool released = false; + bool tcp_parser_active = false; + std::atomic timed_out{false}; +}; + +// Each scenario runs in an isolated subprocess; its server has one connection. +static UrmaFinalAckGate* g_urma_final_ack_gate = nullptr; + +static ParseResult GateUrmaFinalAck(butil::IOBuf* source, Socket* socket, + bool eof, const void* arg) { + ParseResult result = policy::ParseTransportHandshake(source, socket, eof, arg); + if (socket->parsing_stream_type() == InputMessengerProcessor::STREAM_TCP_FD && + AdapterTransport::Get(socket)->handshake_phase() == handshake::ESTABLISHED) { + g_urma_final_ack_gate->Pause(socket); + } + return result; +} + +static ParseResult SkipLegacyUrmaHandshake(butil::IOBuf*, Socket*, bool, + const void*) { + // The old dedicated transport already consumed the TCP handshake. This is + // the baseline parser's URMA exclusion, before the common migration. + return MakeParseError(PARSE_ERROR_TRY_OTHERS); +} + +class UrmaDelayedAckIO : public handshake::HandshakeIO { +public: + explicit UrmaDelayedAckIO(handshake::HandshakeIO* io) : delegate(io) {} + int ReadExact(void* data, size_t len) override { + return delegate->ReadExact(data, len); + } + int PushBack(const void* data, size_t len) override { + return delegate->PushBack(data, len); + } + int WriteAll(const void* data, size_t len) override { + if (len == sizeof(uint32_t)) { + const char* ack = static_cast(data); + held.assign(ack + len - 1, 1); + // Model cross-channel delay: TCP's final ACK byte arrives after + // the URMA SEND, although the client's write has completed. + return delegate->WriteAll(data, len - 1); + } + return delegate->WriteAll(data, len); + } + int FlushAck() { return delegate->WriteAll(held.data(), held.size()); } + handshake::HandshakeIO* delegate; + std::string held; +}; + +enum class UrmaInteropPeer { COMMON, LEGACY_CLIENT, LEGACY_SERVER }; +enum class UrmaFirstPacket { AFTER_ACK, BEFORE_LAST_ACK_BYTE, IN_ACK_PARSER }; + +static ParseResult PauseUrmaCq(butil::IOBuf* source, Socket* socket, bool, + const void* arg) { + auto* gate = static_cast(const_cast(arg)); + gate->Pause(socket); + source->clear(); + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); +} + +// A real endpoint/carrier and mock SEND-to-self give the CQ callback a receive +// to dispatch. The parser barrier keeps that callback in flight during close. +class UrmaCqCloseTest { +public: + int Init(bool polling, bool zerocopy, UrmaFinalAckGate* gate, + bool start_poller = true) { + urma::g_skip_urma_init = false; + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + gflags::SetCommandLineOption("urma_zerocopy_min_size", "0"); + gflags::SetCommandLineOption("urma_use_polling", polling ? "true" : "false"); + urma::GlobalUrmaInitializeOrDie(); + if (!urma::GetUrmaContext() || + (polling && start_poller && + !urma::InitPollingModeWithTag(bthread_self_tag()))) { + return 10; + } + InputMessageHandler handler{}; + handler.parse = PauseUrmaCq; + handler.process = [](InputMessageBase* msg) { msg->Destroy(); }; + handler.arg = gate; + handler.name = "urma_cq_close"; + if (messenger.AddNonProtocolHandler(handler) != 0) { + return 11; + } + SocketOptions options; + options.socket_mode = SOCKET_MODE_URMA; + options.user = &messenger; + if (Socket::Create(options, &id) != 0 || + Socket::Address(id, &socket) != 0) { + return 12; + } + endpoint = static_cast( + AdapterTransport::Get(socket.get())->high_speed_transport())->GetUrmaEp(); + return Prepare(zerocopy); + } + + int Prepare(bool zerocopy) { + if (endpoint->AllocateResources() != 0 || + endpoint->PostRecv(endpoint->_rq_size, zerocopy) != 0) { + return 13; + } + urma::ParsedHello hello; + endpoint->MakeLocalParsedHello(&hello); + endpoint->ApplyRemoteHello(hello); + if (endpoint->ImportPeer(hello) != 0) { + return 14; + } + endpoint->_state = urma::UrmaEndpoint::ESTABLISHED; + if (endpoint->StartCqEvents() != 0 || + Socket::Address(endpoint->_cq_sid, &carrier) != 0) { + return 15; + } + return 0; + } + + bool Send() { + butil::IOBuf data; + data.append("cq-lifetime"); + butil::IOBuf* list[] = {&data}; + return endpoint->CutFromIOBufList(list, 1) == 11; + } + + ~UrmaCqCloseTest() { + if (socket) { + socket->SetFailed(); + socket.reset(); + } + } + + InputMessenger messenger; + SocketId id{INVALID_SOCKET_ID}; + SocketUniquePtr socket; + SocketUniquePtr carrier; + urma::UrmaEndpoint* endpoint{nullptr}; +}; + +static int RunUrmaCqConcurrentClose(bool polling, bool zerocopy, + bool release_last_reference) { + alarm(10); + UrmaFinalAckGate gate; + UrmaCqCloseTest test; + const int init = test.Init(polling, zerocopy, &gate); + if (init != 0) { + return init; + } + auto lifetime = test.endpoint->_cq_lifetime; + if (!test.Send()) { + return 16; + } + if (!gate.Wait()) { + return 17; + } + std::future close; + std::promise started; + if (release_last_reference) { + test.socket->SetFailed(); + test.socket.reset(); + // PollCq's reference must still keep the endpoint/resources alive. + if (lifetime->mutex.try_lock()) { + lifetime->mutex.unlock(); + return 18; + } + } else { + auto ready = started.get_future(); + close = std::async(std::launch::async, [&] { + started.set_value(); + test.endpoint->DeallocateResources(); + }); + ready.wait(); + if (close.wait_for(std::chrono::milliseconds(20)) != + std::future_status::timeout) { + return 19; + } + } + gate.Release(); + if (close.valid()) { + close.get(); + } + // The event dispatcher / poller owns the callback. Wait for its final main + // reference to trigger endpoint teardown. + if (release_last_reference) { + for (;;) { + { + std::unique_lock lock(lifetime->mutex); + if (lifetime->closed) { + break; + } + } + std::this_thread::yield(); + } + } + { + std::unique_lock lock(lifetime->mutex); + if (!lifetime->closed || gate.timed_out.load()) { + return 20; + } + } + // Keep a queued carrier callback alive beyond close / endpoint deletion. + // It must not touch the old endpoint, even if the main Socket is gone. + urma::UrmaEndpoint::PollCq(test.carrier.get()); + if (!release_last_reference && + (test.endpoint->_resource != nullptr || + test.endpoint->_cq_sid != INVALID_SOCKET_ID)) { + return 21; + } + std::weak_ptr weak = lifetime; + lifetime.reset(); + test.carrier.reset(); + if (release_last_reference) { + while (!weak.expired()) { + std::this_thread::yield(); + } + } + return 0; +} + +static int RunUrmaCqCallbackAfterReset(bool polling, bool zerocopy) { + alarm(10); + UrmaFinalAckGate gate; + UrmaCqCloseTest test; + // Keep pollers stopped so the new CQ remains untouched by dispatchers. + const int init = test.Init(polling, zerocopy, &gate, false); + if (init != 0) { + return init; + } + SocketUniquePtr old_carrier = std::move(test.carrier); + const int old_fd = test.endpoint->_resource->jfce->fd; + test.endpoint->Reset(); + if (test.Prepare(zerocopy) != 0) { + return 23; + } + // No dispatcher may drain the new generation before this assertion. + if (!polling) { + test.carrier->_io_event.RemoveConsumer(test.carrier->fd()); + } + if (!test.Send()) { + return 24; + } + urma::UrmaEndpoint::PollCq(old_carrier.get()); + urma_cr_t completions[2]; + if (urma_poll_jfc(test.endpoint->_resource->jfc, 2, completions) != 2) { + return 25; + } + // Old carrier destruction must not close a JFCE fd reused by the reset. + old_carrier.reset(); + if (old_fd == test.endpoint->_resource->jfce->fd && + fcntl(old_fd, F_GETFD) < 0) { + return 26; + } + return 0; +} + +TEST_F(UrmaMockTest, cq_close_waits_for_running_callback) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << polling << ":" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaCqConcurrentClose(polling, zerocopy, false)), + ::testing::ExitedWithCode(0), ""); + } + } +} + +TEST_F(UrmaMockTest, cq_callback_releases_last_socket_reference_before_teardown) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << polling << ":" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaCqConcurrentClose(polling, zerocopy, true)), + ::testing::ExitedWithCode(0), ""); + } + } +} + +TEST_F(UrmaMockTest, cq_callback_rejects_reset_generation) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << polling << ":" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaCqCallbackAfterReset(polling, zerocopy)), + ::testing::ExitedWithCode(0), ""); + } + } +} + +static int RunUrmaHandshakeInterop(int version, bool polling, bool zerocopy, + UrmaInteropPeer peer, UrmaFirstPacket order) { + alarm(15); + urma::g_skip_urma_init = false; + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_use_polling", polling ? "true" : "false"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + urma::FLAGS_urma_client_handshake_version = version; + urma::GlobalUrmaInitializeOrDie(); + if (!urma::GetUrmaContext()) { + return 30; + } + UrmaFinalAckGate gate; + UrmaStreamEchoService service; + Server server; + ServerOptions server_options; + server_options.socket_mode = SOCKET_MODE_URMA; + if (server.AddService(&service, SERVER_DOESNT_OWN_SERVICE) != 0 || + server.Start(0, &server_options) != 0) { + return 31; + } + for (int i = 0; i <= server._am->_max_index.load(butil::memory_order_acquire); ++i) { + auto& handler = server._am->_handlers[i]; + if (handler.parse == policy::ParseTransportHandshake) { + if (peer == UrmaInteropPeer::LEGACY_SERVER) { + handler.parse = SkipLegacyUrmaHandshake; + } else if (order == UrmaFirstPacket::IN_ACK_PARSER) { + g_urma_final_ack_gate = &gate; + handler.parse = GateUrmaFinalAck; + } + } + } + Channel channel; + ChannelOptions options; + options.socket_mode = SOCKET_MODE_URMA; + options.protocol = "baidu_std"; + options.connection_type = "single"; + options.timeout_ms = 3000; + options.max_retry = 0; + if (channel.Init(server.listen_address(), &options) != 0) { + return 32; + } + SocketUniquePtr client; + if (Socket::Address(channel._server_id, &client) != 0) { + return 33; + } + SocketOptions accepted_options; + accepted_options.socket_mode = SOCKET_MODE_URMA; + accepted_options.user = server._am; + accepted_options.need_on_edge_trigger = true; + SocketId accepted_id; + if (Socket::Create(accepted_options, &accepted_id) != 0) { + return 34; + } + SocketUniquePtr accepted; + if (Socket::Address(accepted_id, &accepted) != 0) { + return 35; + } + // Pin both main sockets throughout the handoff/RPC checks. CQ teardown + // lifetime races are a separate test and fix, not part of this scenario. + if (peer == UrmaInteropPeer::LEGACY_CLIENT) { + InstallLegacyUrmaTransport(client.get(), client->user()); + } else if (peer == UrmaInteropPeer::LEGACY_SERVER) { + InstallLegacyUrmaTransport(accepted.get(), server._am); + } + int fds[2]; + if (CreateUrmaInteropTcpPair(fds) != 0 || + client->ResetFileDescriptor(fds[0]) != 0 || + accepted->ResetFileDescriptor(fds[1]) != 0) { + return 36; + } + urma::UrmaEndpoint* accepted_ep = peer == UrmaInteropPeer::LEGACY_SERVER + ? static_cast(accepted->_transport.get())->GetUrmaEp() + : static_cast(AdapterTransport::Get(accepted.get()) + ->high_speed_transport())->GetUrmaEp(); + handshake::HandshakeSession* session = peer == UrmaInteropPeer::LEGACY_CLIENT + ? nullptr : AdapterTransport::Get(client.get())->handshake_session(); + UrmaDelayedAckIO delayed(session ? &session->_socket_io : nullptr); + if (order == UrmaFirstPacket::BEFORE_LAST_ACK_BYTE) { + session->SetIOForTest(&delayed); + } + std::promise connected; + auto ready = connected.get_future(); + auto done = [](int error, void* data) { + static_cast*>(data)->set_value(error); + }; + if (peer == UrmaInteropPeer::LEGACY_CLIENT) { + client->_app_connect->StartConnect(client.get(), done, &connected); + } else if (AdapterTransport::StartClientUpgrade(client.get(), done, &connected) != 0) { + return 37; + } + if (ready.wait_for(std::chrono::seconds(3)) != std::future_status::ready || + ready.get() != 0) { + return 38; + } + if (peer == UrmaInteropPeer::LEGACY_CLIENT) { + auto* transport = static_cast(client->_transport.get()); + if (transport->_urma_state != UrmaTransport::URMA_ON || + transport->GetUrmaEp()->_handshake_version != version) { + return 39; + } + } else if (AdapterTransport::Get(client.get())->handshake_phase() != + handshake::ESTABLISHED || + AdapterTransport::Get(client.get())->handshake_version() != version) { + return 40; + } + if (order == UrmaFirstPacket::IN_ACK_PARSER && !gate.Wait()) { + return 41; + } + if (order == UrmaFirstPacket::BEFORE_LAST_ACK_BYTE) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (AdapterTransport::Get(accepted.get())->handshake_phase() != + handshake::ACK_WAIT) { + if (std::chrono::steady_clock::now() > deadline) { + return 42; + } + std::this_thread::yield(); + } + } + if (order == UrmaFirstPacket::AFTER_ACK) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (accepted_ep->_state.load(butil::memory_order_acquire) != + urma::UrmaEndpoint::ESTABLISHED) { + if (std::chrono::steady_clock::now() > deadline) { + return 42; + } + std::this_thread::yield(); + } + } + test::EchoService::Stub stub(&channel); + test::EchoRequest request; + test::EchoResponse response; + Controller cntl; + request.set_message("first URMA RPC after final TCP ACK"); + cntl.request_attachment().append(std::string(512, 'a')); + stub.Echo(&cntl, &request, &response, DoNothing()); + if (order != UrmaFirstPacket::AFTER_ACK) { + pollfd cq_event{accepted_ep->_resource->jfce->fd, POLLIN, 0}; + if (poll(&cq_event, 1, 3000) != 1 || !(cq_event.revents & POLLIN) || + service.calls.load(std::memory_order_acquire) != 0) { + return 43; + } + // A real receive completion is queued, but no CQ consumer may run + // while ACK is incomplete or its TCP parser is still on the stack. + if (accepted_ep->_cq_sid != INVALID_SOCKET_ID) { + return 44; + } + if (order == UrmaFirstPacket::BEFORE_LAST_ACK_BYTE) { + if (AdapterTransport::Get(accepted.get())->handshake_phase() != + handshake::ACK_WAIT || delayed.held.size() != 1 || + delayed.FlushAck() != 0) { + return 45; + } + } else { + if (!gate.tcp_parser_active) { + return 46; + } + gate.Release(); + } + } + Join(cntl.call_id()); + if (cntl.Failed() || response.message() != request.message() || + cntl.response_attachment().to_string() != std::string(512, 'a')) { + return 47; + } + for (int i = 0; i < 16; ++i) { + Controller next; + test::EchoResponse reply; + request.set_message(std::to_string(i)); + const std::string attachment(256 + i * 257, static_cast('b' + i)); + next.request_attachment().append(attachment); + stub.Echo(&next, &request, &reply, nullptr); + if (next.Failed() || reply.message() != request.message() || + next.response_attachment().to_string() != attachment) { + return 48; + } + } + if (service.calls.load(std::memory_order_acquire) != 17 || + gate.timed_out.load(std::memory_order_acquire)) { + return 49; + } + if (peer == UrmaInteropPeer::LEGACY_SERVER && + accepted_ep->_handshake_version != version) { + return 50; + } + if (peer != UrmaInteropPeer::LEGACY_SERVER && + (AdapterTransport::Get(accepted.get())->handshake_phase() != + handshake::ESTABLISHED || + AdapterTransport::Get(accepted.get())->handshake_version() != version)) { + return 51; + } + client->SetFailed(); + accepted->SetFailed(); + server.Stop(0); + server.Join(); + return 0; +} + +TEST_F(UrmaMockTest, legacy_client_exchanges_rpc_with_common_server) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (int version : {2, 3}) { + for (bool polling : {false, true}) { + SCOPED_TRACE(::testing::Message() << "version=" << version + << " polling=" << polling); + ASSERT_EXIT(_exit(RunUrmaHandshakeInterop(version, polling, false, + UrmaInteropPeer::LEGACY_CLIENT, UrmaFirstPacket::AFTER_ACK)), + ::testing::ExitedWithCode(0), ""); + } + } +} + +TEST_F(UrmaMockTest, common_client_exchanges_rpc_with_legacy_server) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (int version : {2, 3}) { + for (bool polling : {false, true}) { + SCOPED_TRACE(::testing::Message() << "version=" << version + << " polling=" << polling); + ASSERT_EXIT(_exit(RunUrmaHandshakeInterop(version, polling, false, + UrmaInteropPeer::LEGACY_SERVER, UrmaFirstPacket::AFTER_ACK)), + ::testing::ExitedWithCode(0), ""); + } + } +} + +TEST_F(UrmaMockTest, first_urma_packet_waits_for_final_ack_byte) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (int version : {2, 3}) { + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << "version=" << version + << " polling=" << polling << " zerocopy=" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaHandshakeInterop(version, polling, zerocopy, + UrmaInteropPeer::COMMON, UrmaFirstPacket::BEFORE_LAST_ACK_BYTE)), + ::testing::ExitedWithCode(0), ""); + } + } + } +} + +TEST_F(UrmaMockTest, first_urma_packet_after_ack_exchanges_rpc) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (int version : {2, 3}) { + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << "version=" << version + << " polling=" << polling << " zerocopy=" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaHandshakeInterop(version, polling, zerocopy, + UrmaInteropPeer::COMMON, UrmaFirstPacket::AFTER_ACK)), + ::testing::ExitedWithCode(0), ""); + } + } + } +} + +TEST_F(UrmaMockTest, first_urma_packet_waits_for_ack_parser_to_exit) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + for (int version : {2, 3}) { + for (bool polling : {false, true}) { + for (bool zerocopy : {false, true}) { + SCOPED_TRACE(::testing::Message() << "version=" << version + << " polling=" << polling << " zerocopy=" << zerocopy); + ASSERT_EXIT(_exit(RunUrmaHandshakeInterop(version, polling, zerocopy, + UrmaInteropPeer::COMMON, UrmaFirstPacket::IN_ACK_PARSER)), + ::testing::ExitedWithCode(0), ""); + } + } + } +} + +static int RunUrmaMockConcurrentRpc(bool zerocopy, bool tcp_server = false) { + urma::g_skip_urma_init = false; + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + urma::GlobalUrmaInitializeOrDie(); + if (!urma::GetUrmaContext()) { + return 10; + } + UrmaStreamEchoService service; + Server server; + ServerOptions server_options; + server_options.socket_mode = tcp_server ? SOCKET_MODE_TCP : SOCKET_MODE_URMA; + if (server.AddService(&service, SERVER_DOESNT_OWN_SERVICE) != 0 || + server.Start(0, &server_options) != 0) { + return 11; + } + Channel channel; + ChannelOptions options; + options.socket_mode = SOCKET_MODE_URMA; + options.protocol = "baidu_std"; + options.timeout_ms = 3000; + options.max_retry = 0; + options.connection_type = "single"; + if (channel.Init(server.listen_address(), &options) != 0) { + return 12; + } + std::atomic errors(0); + std::vector workers; + for (int worker = 0; worker < 4; ++worker) { + workers.emplace_back([&, worker]() { + test::EchoService::Stub stub(&channel); + for (int round = 0; round < 8; ++round) { + const size_t depth = 16; + std::vector cntls(depth); + std::vector requests(depth); + std::vector responses(depth); + for (size_t i = 0; i < depth; ++i) { + requests[i].set_message(std::to_string(worker) + ":" + + std::to_string(round) + ":" + std::to_string(i)); + cntls[i].request_attachment().append(std::string(i * 13, 'x')); + stub.Echo(&cntls[i], &requests[i], &responses[i], DoNothing()); + } + for (size_t i = 0; i < depth; ++i) { + Join(cntls[i].call_id()); + if (cntls[i].Failed() || + responses[i].message() != requests[i].message() || + cntls[i].response_attachment().to_string() != + cntls[i].request_attachment().to_string()) { + ++errors; + } + } + if (errors.load() != 0) { + break; + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + SocketUniquePtr main_socket; + if (Socket::Address(channel._server_id, &main_socket) != 0 || + AdapterTransport::Get(main_socket.get())->handshake_phase() != + (tcp_server ? handshake::FALLBACK_TCP : handshake::ESTABLISHED)) { + ++errors; + } + server.Stop(0); + server.Join(); + return errors.load() == 0 ? 0 : 13; +} + +TEST_F(UrmaMockTest, concurrent_rpc_preserves_stream_boundaries) { + // Isolate the process-global pool/IOBuf hooks from tests which intentionally + // bypass URMA initialization. Threadsafe death tests exec a fresh process. + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockConcurrentRpc(true)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, concurrent_copy_rpc_preserves_stream_boundaries) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockConcurrentRpc(false)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, v2_client_falls_back_to_plain_tcp_server) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT({ + alarm(15); + urma::FLAGS_urma_client_handshake_version = 2; + _exit(RunUrmaMockConcurrentRpc(false, true)); + }, ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, v3_client_falls_back_to_plain_tcp_server) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT({ + alarm(15); + urma::FLAGS_urma_client_handshake_version = 3; + _exit(RunUrmaMockConcurrentRpc(false, true)); + }, ::testing::ExitedWithCode(0), ""); +} + +static int RunUrmaMockTracedRpc() { + gflags::SetCommandLineOption("urma_trace_stream", "true"); + return RunUrmaMockConcurrentRpc(false); +} + +static int RunUrmaMockOrderedRpc(bool zerocopy) { + if (gflags::SetCommandLineOption( + "urma_mock_require_send_order", "true").empty()) { + return 20; + } + // Exercise the real posting paths, including credit-only SEND_IMM WRs, + // rather than just checking a separately constructed flag value. + // A rejected ACK can exhaust credits. Bound this isolated subprocess so + // a regression fails the test instead of hanging the test runner. + alarm(10); + const int result = RunUrmaMockConcurrentRpc(zerocopy); + alarm(0); + return result; +} + +TEST_F(UrmaMockTest, copy_rpc_requests_ordered_data_and_credit_sends) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockOrderedRpc(false)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, zerocopy_rpc_requests_ordered_data_and_credit_sends) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockOrderedRpc(true)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, traced_rpc_dumps_bounded_history_on_teardown) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockTracedRpc()), + ::testing::ExitedWithCode(0), "URMA_TRACE"); +} + +static int RunUrmaMockCtpLargeRpc(bool zerocopy, bool polling) { + // Isolated by ASSERT_EXIT: bound failures that stop credit progress. + alarm(15); + urma::g_skip_urma_init = false; + if (gflags::SetCommandLineOption("urma_mock_max_send_size", "4096").empty()) { + return 20; + } + gflags::SetCommandLineOption("urma_buffer_size", "8192"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + gflags::SetCommandLineOption("urma_use_polling", polling ? "true" : "false"); + urma::GlobalUrmaInitializeOrDie(); + // The fix must split sends, not silently shrink receive/IOBuf blocks. + if (!urma::GetUrmaContext() || urma::GetUrmaRecvBlockSize() <= 4096) { + return 21; + } + UrmaStreamEchoService service; + Server server; + ServerOptions server_options; + server_options.socket_mode = SOCKET_MODE_URMA; + if (server.AddService(&service, SERVER_DOESNT_OWN_SERVICE) != 0 || + server.Start(0, &server_options) != 0) { + return 22; + } + Channel channel; + ChannelOptions options; + options.socket_mode = SOCKET_MODE_URMA; + options.protocol = "baidu_std"; + options.connection_type = "single"; + options.timeout_ms = 3000; + options.max_retry = 0; + if (channel.Init(server.listen_address(), &options) != 0) { + return 23; + } + test::EchoService::Stub stub(&channel); + bool ok = true; + for (size_t size : {size_t(0), size_t(4095), size_t(4096), size_t(4097), + size_t(8192), size_t(1048576), size_t(1048577)}) { + // A varying pattern catches omissions/reordering hidden by all-'a' + // attachments. 1 MiB crosses the 125-WR send-credit window repeatedly. + std::string expected(size, '\0'); + for (size_t i = 0; i < size; ++i) { + expected[i] = static_cast((i * 37 + i / 4096) % 251); + } + Controller cntl; + test::EchoRequest request; + test::EchoResponse response; + request.set_message("ctp-large-" + std::to_string(size)); + cntl.request_attachment().append(expected); + stub.Echo(&cntl, &request, &response, nullptr); + if (cntl.Failed() || response.message() != request.message() || + cntl.response_attachment().to_string() != expected) { + ok = false; + break; + } + } + SocketUniquePtr socket; + if (Socket::Address(channel._server_id, &socket) != 0 || + AdapterTransport::Get(socket.get())->handshake_phase() != + handshake::ESTABLISHED) { + ok = false; + } + server.Stop(0); + server.Join(); + alarm(0); + return ok ? 0 : 24; +} + +TEST_F(UrmaMockTest, ctp_large_copy_rpc_with_default_buffers) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockCtpLargeRpc(false, false)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, ctp_large_zerocopy_rpc_with_default_buffers) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockCtpLargeRpc(true, false)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, ctp_large_copy_rpc_with_polling) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockCtpLargeRpc(false, true)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, ctp_large_zerocopy_rpc_with_polling) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockCtpLargeRpc(true, true)), + ::testing::ExitedWithCode(0), ""); +} + +static int RunUrmaMalformedResponseTrace() { + GlobalInitializeOrDie(); + gflags::SetCommandLineOption("urma_trace_stream", "true"); + SocketOptions options; + options.user = get_or_new_client_side_messenger(); + SocketId id; + if (Socket::Create(options, &id) != 0) { + return 10; + } + SocketUniquePtr socket; + if (Socket::Address(id, &socket) != 0) { + return 11; + } + struct FailSocket { + Socket* socket; + ~FailSocket() { socket->SetFailed(); } + } fail_socket{socket.get()}; + socket->set_preferred_index(PROTOCOL_BAIDU_STD); + urma::UrmaEndpoint endpoint(socket.get()); + endpoint._state.store(urma::UrmaEndpoint::ESTABLISHED, + butil::memory_order_release); + const std::string bad = "BAD!"; + urma_sge_t sge{reinterpret_cast(bad.data()), 4, nullptr, nullptr}; + endpoint._stream_trace->Record(urma::UrmaStreamTrace::RX_USE, 0, 1, 0, &sge, 1); + endpoint._input_processor.read_buf().append(bad); + endpoint.DispatchReceivedBytes(socket, bad.size()); + return socket->Failed() ? 0 : 12; +} + +TEST_F(UrmaMockTest, malformed_response_dumps_exact_bounded_prefix) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMalformedResponseTrace()), + ::testing::ExitedWithCode(0), + "reason=RPC parse failed buffered=4 parse_prefix=42414421"); +} + +static int RunUrmaMockReorderedReceives(bool zerocopy) { + urma::g_skip_urma_init = false; + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + gflags::SetCommandLineOption("urma_zerocopy_min_size", "0"); + urma::GlobalUrmaInitializeOrDie(); + if (!urma::GetUrmaContext()) { + return 10; + } + SocketId id; + SocketOptions options; + if (Socket::Create(options, &id) != 0) { + return 11; + } + SocketUniquePtr socket; + if (Socket::Address(id, &socket) != 0) { + return 12; + } + struct FailSocket { + Socket* socket; + ~FailSocket() { socket->SetFailed(); } + } fail_socket{socket.get()}; + urma::UrmaEndpoint endpoint(socket.get()); + endpoint._rq_size = 4; + endpoint._remote_window_capacity = 16; + endpoint._local_window_capacity = 16; + if (endpoint.AllocateResources() != 0 || + endpoint.PostRecv(endpoint._rq_size, zerocopy) != 0) { + return 13; + } + const std::string payloads[] = {"PRPC", std::string("\x08\0", 2)}; + for (int round = 0; round < 32; ++round) { + const size_t first = endpoint._rq_received; + const size_t ack_slot = (first + 1) % endpoint._rq_size; + const size_t second = (first + 2) % endpoint._rq_size; + memcpy(endpoint._rbuf_data[first], payloads[0].data(), payloads[0].size()); + memcpy(endpoint._rbuf_data[second], payloads[1].data(), payloads[1].size()); + urma_cr_t later{}; + later.status = URMA_CR_SUCCESS; + later.flag.bs.s_r = 1; + later.opcode = URMA_CR_OPC_SEND; + later.user_ctx = second + 1; + later.completion_len = payloads[1].size(); + // Let the pure credit ACK and the later payload complete before the + // first payload. Neither may consume the first payload's buffer. + urma_cr_t ack = later; + ack.opcode = URMA_CR_OPC_SEND_WITH_IMM; + ack.user_ctx = ack_slot + 1; + ack.completion_len = 0; + ack.imm_data = 1; + endpoint._remote_rq_window_size.store(15, butil::memory_order_relaxed); + if (endpoint.HandleCompletion(ack) != 0 || + endpoint._remote_rq_window_size.load(butil::memory_order_relaxed) != 15) { + return 17; + } + if (endpoint.HandleCompletion(later) != 0 || + !endpoint._input_processor.read_buf().empty() || + endpoint._rq_received != first) { + return 14; + } + // A duplicate must be rejected without consuming the queued payload. + if (endpoint.HandleCompletion(later) != -1 || errno != EPROTO) { + return 15; + } + urma_cr_t earlier = later; + earlier.user_ctx = first + 1; + earlier.completion_len = payloads[0].size(); + if (endpoint.HandleCompletion(earlier) != 6 || + endpoint._input_processor.read_buf().to_string() != + payloads[0] + payloads[1] || + endpoint._remote_rq_window_size.load(butil::memory_order_relaxed) != 16) { + return 16; + } + endpoint._input_processor.read_buf().clear(); + // This test does not connect a peer; suppress credit-ACK posting. + endpoint._new_rq_wrs.store(0, butil::memory_order_relaxed); + } + return 0; +} + +TEST_F(UrmaMockTest, reordered_receives_preserve_copy_stream) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockReorderedReceives(false)), + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, reordered_receives_preserve_zerocopy_stream) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(_exit(RunUrmaMockReorderedReceives(true)), + ::testing::ExitedWithCode(0), ""); +} + TEST_F(UrmaMockTest, init_and_enumerate_device) { urma_init_attr_t init_attr{}; EXPECT_EQ(URMA_EEXIST, urma_init(&init_attr)); 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 04/17] 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 05/17] 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 06/17] 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()); From 822d5ded5c024a75e1163e5181010588952a76f5 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Sun, 25 Oct 2026 20:20:35 -0700 Subject: [PATCH 07/17] Fix transport handshake cleanup and dispatch --- src/brpc/adapter_transport.cpp | 4 + src/brpc/handshake/handshake_io.cpp | 13 ++- src/brpc/ubshm/ub_ring.cpp | 9 ++ test/brpc_transport_handshake_unittest.cpp | 110 +++++++++++++++++++++ test/brpc_ubring_unittest.cpp | 56 +++++++++++ 5 files changed, 190 insertions(+), 2 deletions(-) diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index 49558efa07..00b19f2c76 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -278,6 +278,10 @@ int AdapterTransport::StartClientUpgrade(const Socket* socket, } ParseResult AdapterTransport::ProcessUpgradeReadable(butil::IOBuf* source) { + if (_handshake.phase(butil::memory_order_acquire) == + handshake::ESTABLISHED) { + return MakeParseError(PARSE_ERROR_TRY_OTHERS); + } ParseResult result(PARSE_ERROR_NOT_ENOUGH_DATA); if (_socket->parsing_context() != nullptr) { handshake::ServerHandshakeContext* context = diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp index 86ef81026d..67840e4de8 100644 --- a/src/brpc/handshake/handshake_io.cpp +++ b/src/brpc/handshake/handshake_io.cpp @@ -46,6 +46,7 @@ static const int WAIT_TIMEOUT_MS = 50; SocketHandshakeIO::SocketHandshakeIO(Socket* socket) : _socket(socket) , _read_butex(bthread::butex_create_checked >()) { + _read_butex->store(0, butil::memory_order_relaxed); } SocketHandshakeIO::~SocketHandshakeIO() { @@ -107,10 +108,15 @@ int SocketHandshakeIO::ReadExact(void* data, size_t len) { } template -static int WriteAllLoop(size_t len, WriteOnce write_once, +static int WriteAllLoop(SocketId socket_id, size_t len, WriteOnce write_once, WaitWritable wait_writable) { size_t written = 0; while (written < len) { + SocketUniquePtr alive; + if (Socket::Address(socket_id, &alive) != 0) { + errno = EFAILEDSOCKET; + return -1; + } const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); const ssize_t nw = write_once(written, len - written); if (nw > 0) { @@ -129,6 +135,9 @@ static int WriteAllLoop(size_t len, WriteOnce write_once, } if (wait_writable(&duetime) < 0 && errno != ETIMEDOUT) { + if (alive->Failed()) { + errno = EFAILEDSOCKET; + } return -1; } } @@ -139,7 +148,7 @@ int SocketHandshakeIO::WriteAll(const void* data, size_t len) { CHECK(data != nullptr); CHECK(_socket != nullptr); const int fd = _socket->fd(); - return WriteAllLoop(len, + return WriteAllLoop(_socket->id(), len, [data, fd](size_t offset, size_t remaining) { return write(fd, static_cast(data) + offset, remaining); diff --git a/src/brpc/ubshm/ub_ring.cpp b/src/brpc/ubshm/ub_ring.cpp index 595b96d03b..a5ff3d06a4 100644 --- a/src/brpc/ubshm/ub_ring.cpp +++ b/src/brpc/ubshm/ub_ring.cpp @@ -70,6 +70,15 @@ RETURN_CODE UBRing::UbrTrxClose() { } return UBRING_ERR; } + // No remote mapping means no peer-close exchange was established. + if (_trx->remote_shm.addr == nullptr) { + RETURN_CODE rc = UbrTrxFreeShm(_trx); + if (rc == UBRING_OK) { + rc = UBRingManager::ReleaseUbrTrxFromMgr(_trx); + _trx = nullptr; + } + return rc; + } if (_trx->ubr_rx.remote_tx_event_q.addr != nullptr) { ((UbrEventQMsg *)_trx->ubr_rx.remote_tx_event_q.addr)->flag = UBR_STATE_CLOSING; } diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index e70f397448..d9be08fe5e 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -18,14 +18,18 @@ #include #include +#include #include #include #include #include #include "bthread/bthread.h" +#include "bthread/butex.h" +#include "bthread/countdown_event.h" #include "butil/fd_guard.h" #include "butil/sys_byteorder.h" +#include "butil/time.h" #include "brpc/adapter_transport.h" #include "brpc/global.h" #include "brpc/handshake/handshake_io.h" @@ -92,6 +96,26 @@ static void* RunBlockingHandshakeRead(void* arg) { return nullptr; } +struct BlockingHandshakeWrite { + SocketHandshakeIO* io; + bthread::CountdownEvent started; + bthread::CountdownEvent completed; + int result = 0; + int error = 0; + + explicit BlockingHandshakeWrite(SocketHandshakeIO* value) : io(value) {} +}; + +static void* RunBlockingHandshakeWrite(void* arg) { + auto* write = static_cast(arg); + const char byte = 0; + write->started.signal(); + write->result = write->io->WriteAll(&byte, sizeof(byte)); + write->error = errno; + write->completed.signal(); + return nullptr; +} + class TestAppConnect : public AppConnect { public: void StartConnect(const Socket*, void (*done)(int, void*), @@ -388,6 +412,92 @@ static std::string MakeUBShmHello() { return frame; } +TEST(TransportHandshakeTest, socket_io_initializes_read_butex) { + auto* reused = bthread::butex_create_checked>(); + reused->store(37, butil::memory_order_relaxed); + bthread::butex_destroy(reused); + SocketHandshakeIO io; + EXPECT_EQ(0, io._read_butex->load(butil::memory_order_relaxed)); + io.NotifyReadable(); + EXPECT_EQ(1, io._read_butex->load(butil::memory_order_relaxed)); +} + +TEST(TransportHandshakeTest, blocked_write_stops_after_socket_failure) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + const int flags = fcntl(fds[0], F_GETFL, 0); + ASSERT_GE(flags, 0); + ASSERT_EQ(0, fcntl(fds[0], F_SETFL, flags | O_NONBLOCK)); + char data[4096] = {}; + ssize_t n; + do { + n = send(fds[0], data, sizeof(data), MSG_NOSIGNAL); + } while (n > 0 || (n < 0 && errno == EINTR)); + ASSERT_EQ(-1, n); + ASSERT_EQ(EAGAIN, errno); + + SocketOptions options; + options.fd = fds[0]; + options.need_on_edge_trigger = true; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + SocketHandshakeIO io(socket.get()); + BlockingHandshakeWrite write(&io); + bthread_t tid; + ASSERT_EQ(0, bthread_start_background( + &tid, nullptr, RunBlockingHandshakeWrite, &write)); + write.started.wait(); + socket->SetFailed(EFAILEDSOCKET, "cancel blocked handshake write"); + const int completed = write.completed.timed_wait( + butil::milliseconds_from_now(1000)); + if (completed != 0) { + shutdown(socket->fd(), SHUT_RDWR); + } + ASSERT_EQ(0, bthread_join(tid, nullptr)); + EXPECT_EQ(0, completed); + EXPECT_EQ(-1, write.result); + EXPECT_EQ(EFAILEDSOCKET, write.error); +} + +TEST(TransportHandshakeTest, established_socket_does_not_restart_handshake) { + const SocketMode modes[] = { + SOCKET_MODE_TCP, +#if BRPC_WITH_RDMA + SOCKET_MODE_RDMA, +#endif +#if BRPC_WITH_UBRING + SOCKET_MODE_UBRING, +#endif + }; + for (SocketMode mode : modes) { + SocketOptions options; + options.socket_mode = mode; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + AdapterTransport* adapter = AdapterTransport::Get(socket.get()); + adapter->handshake_session()->SetPhase(ESTABLISHED); + adapter->CompleteConnection(ESTABLISHED); + + const char* prefixes[] = {"RDMA", "RDM3", "UB"}; + for (const char* prefix : prefixes) { + butil::IOBuf input; + input.append(prefix); + const ParseResult result = adapter->ProcessUpgradeReadable(&input); + EXPECT_EQ(PARSE_ERROR_TRY_OTHERS, result.error()); + EXPECT_EQ(prefix, input.to_string()); + EXPECT_EQ(ESTABLISHED, adapter->handshake_phase()); + EXPECT_EQ(nullptr, socket->parsing_context()); + EXPECT_TRUE(adapter->connection_completed()); + } + socket->SetFailed(); + } +} + TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { int fds[2]; ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); diff --git a/test/brpc_ubring_unittest.cpp b/test/brpc_ubring_unittest.cpp index e0ae27fb4d..7c61e50efd 100644 --- a/test/brpc_ubring_unittest.cpp +++ b/test/brpc_ubring_unittest.cpp @@ -1078,6 +1078,28 @@ TEST_F(UBShmEndpointTest, reset_is_idempotent) { _ep->Reset(); } +TEST_F(UBShmEndpointTest, reset_unmapped_client_does_not_wait_for_peer) { + using namespace brpc::ubring; + const uint32_t allocated = UBRingManager::g_ubr_mgr.trx_num; + SHM local = {nullptr, 4 * 1024 * 1024, 0, {0}, + static_cast(_socket->fd())}; + const std::string name = "unmapped_" + std::to_string(getpid()); + ASSERT_EQ(0, _ep->AllocateClientResources(&local, name.c_str())); + const std::string shm_name(local.name); + EXPECT_EQ(allocated + 1, UBRingManager::g_ubr_mgr.trx_num); + + const int64_t start = butil::gettimeofday_us(); + _ep->Reset(); + EXPECT_LT(butil::gettimeofday_us() - start, 1000000); + EXPECT_EQ(allocated, UBRingManager::g_ubr_mgr.trx_num); + EXPECT_EQ(nullptr, _ep->_ub_ring); + EXPECT_EQ(brpc::INVALID_SOCKET_ID, _ep->_poller_sid); + butil::fd_guard missing(shm_open(shm_name.c_str(), O_RDONLY, 0)); + EXPECT_LT(missing, 0); + EXPECT_EQ(ENOENT, errno); + _ep->Reset(); +} + namespace { bool WaitForUBCondition(const std::function& condition) { const int64_t deadline = butil::gettimeofday_us() + 5000000; @@ -1157,6 +1179,40 @@ void SendAndCheckUBFallbackRPC(int fd, bool include_ack, uint64_t id) { } } // namespace +TEST_F(UBShmEndpointTest, tcp_server_fallback_completes_with_short_rpc_timeout) { + UBFallbackEchoService service; + brpc::Server server; + ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions options; + options.internal_port = -1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &options)); + + const uint32_t allocated = brpc::ubring::UBRingManager::g_ubr_mgr.trx_num; + brpc::Channel channel; + brpc::ChannelOptions channel_options; + channel_options.socket_mode = brpc::SOCKET_MODE_UBRING; + channel_options.connect_timeout_ms = 1000; + channel_options.timeout_ms = 1000; + channel_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(server.listen_address(), &channel_options)); + brpc::Controller cntl; + test::EchoRequest request; + test::EchoResponse response; + request.set_message("UBSHM client to TCP server"); + test::EchoService_Stub stub(&channel); + stub.Echo(&cntl, &request, &response, nullptr); + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ(request.message(), response.message()); + EXPECT_EQ(allocated, brpc::ubring::UBRingManager::g_ubr_mgr.trx_num); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(cntl._single_server_id, &socket)); + EXPECT_EQ(brpc::handshake::FALLBACK_TCP, + brpc::AdapterTransport::Get(socket.get())->handshake_phase()); + EXPECT_EQ(nullptr, brpc::UBShmTransport::Get(socket.get()) + ->GetUBShmEp()->_ub_ring); + socket->SetFailed(); +} + TEST_F(UBShmEndpointTest, failed_server_allocation_cleans_resources_and_serves_tcp) { UBFallbackEchoService service; brpc::Server server; From cebd46c6bfe98c2a0886def9a86a099df7c835ed Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Sun, 25 Oct 2026 20:49:32 -0700 Subject: [PATCH 08/17] Handle handshake startup errors and preserve failure causes --- src/brpc/adapter_transport.cpp | 18 +++++-- src/brpc/adapter_transport.h | 5 ++ src/brpc/transport_handshake.cpp | 2 + test/brpc_transport_handshake_unittest.cpp | 60 ++++++++++++++++++++-- 4 files changed, 77 insertions(+), 8 deletions(-) diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index 00b19f2c76..fbde15df3b 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -76,9 +76,9 @@ class AdapterConnect : public AppConnect { socket, _app_connect, done, data}; if (AdapterTransport::StartClientUpgrade( socket, OnUpgradeComplete, task) != 0) { + const int error = errno != 0 ? errno : EAGAIN; AdapterTransport::Get(const_cast(socket))->CompleteConnection( handshake::FAILED); - const int error = errno != 0 ? errno : EAGAIN; delete task; done(error, data); } @@ -150,8 +150,8 @@ class RdmaClientHandshakeTransport : public handshake::HandshakeTransport { void OnFallback() override { _transport->DeactivateUpgrade(); } void OnFailed() override { - _transport->DeactivateUpgrade(); const int saved_errno = errno != 0 ? errno : EPROTO; + _transport->DeactivateUpgrade(); *_connect_error = saved_errno; _socket->SetFailed(saved_errno, "Fail to complete rdma handshake from %s: %s", @@ -202,8 +202,8 @@ class UBShmClientHandshakeTransport : public handshake::HandshakeTransport { void OnFallback() override { _transport->DeactivateUpgrade(); } void OnFailed() override { - _transport->DeactivateUpgrade(); const int saved_errno = errno != 0 ? errno : EPROTO; + _transport->DeactivateUpgrade(); *_connect_error = saved_errno; _socket->SetFailed(saved_errno, "Fail to complete ubring handshake from %s: %s", @@ -260,18 +260,26 @@ bool AdapterTransport::upgrade_capable(SocketMode mode) const { int AdapterTransport::StartClientUpgrade(const Socket* socket, void (*done)(int, void*), void* data) { + return StartClientUpgradeImpl(socket, done, data, bthread_start_background); +} + +int AdapterTransport::StartClientUpgradeImpl( + const Socket* socket, void (*done)(int, void*), void* data, + StartBackground start_background) { AdapterTransport* adapter = Get(const_cast(socket)); ClientHandshakeTask* task = new ClientHandshakeTask{adapter, done, data, SocketUniquePtr()}; if (Socket::Address(socket->id(), &task->socket) != 0) { delete task; + errno = EFAILEDSOCKET; 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) { + const int rc = start_background(&tid, &attr, ProcessClientHandshake, task); + if (rc != 0) { delete task; + errno = rc; return -1; } return 0; diff --git a/src/brpc/adapter_transport.h b/src/brpc/adapter_transport.h index aed6ee5bb9..5ec1537890 100644 --- a/src/brpc/adapter_transport.h +++ b/src/brpc/adapter_transport.h @@ -90,6 +90,11 @@ friend class UBShmTransport; void CheckUnexpectedTcpData(); static void OnNewMessagesAfterUpgrade(Socket* socket); static void* ProcessClientHandshake(void* arg); + using StartBackground = int (*)(bthread_t*, const bthread_attr_t*, + void* (*)(void*), void*); + static int StartClientUpgradeImpl(const Socket* socket, + void (*done)(int, void*), void* data, + StartBackground start_background); SocketMode _mode; handshake::HandshakeSession _handshake; diff --git a/src/brpc/transport_handshake.cpp b/src/brpc/transport_handshake.cpp index ab07971c2b..1657eabd80 100644 --- a/src/brpc/transport_handshake.cpp +++ b/src/brpc/transport_handshake.cpp @@ -77,8 +77,10 @@ void ServerHandshakeContext::Destroy() { static StepResult FinishWithFailure(HandshakeSession* session, HandshakeTransport* transport) { + const int saved_errno = errno != 0 ? errno : EPROTO; transport->OnFailed(); session->MarkFailed(); + errno = saved_errno; return STEP_ERROR; } diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index d9be08fe5e..e815414c53 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -27,6 +27,7 @@ #include "bthread/bthread.h" #include "bthread/butex.h" #include "bthread/countdown_event.h" +#include "bthread/unstable.h" #include "butil/fd_guard.h" #include "butil/sys_byteorder.h" #include "butil/time.h" @@ -98,7 +99,6 @@ static void* RunBlockingHandshakeRead(void* arg) { struct BlockingHandshakeWrite { SocketHandshakeIO* io; - bthread::CountdownEvent started; bthread::CountdownEvent completed; int result = 0; int error = 0; @@ -109,7 +109,6 @@ struct BlockingHandshakeWrite { static void* RunBlockingHandshakeWrite(void* arg) { auto* write = static_cast(arg); const char byte = 0; - write->started.signal(); write->result = write->io->WriteAll(&byte, sizeof(byte)); write->error = errno; write->completed.signal(); @@ -126,6 +125,12 @@ class TestAppConnect : public AppConnect { void StopConnect(Socket*) override {} }; +static int RejectHandshakeStart(bthread_t*, const bthread_attr_t*, + void* (*)(void*), void*) { + errno = EINTR; + return EAGAIN; +} + static FrameSpec FixedSpec(const char* magic, size_t magic_len, size_t total_len) { return FrameSpec(magic, magic_len, total_len, total_len, @@ -301,6 +306,25 @@ class TestUpgradeMessenger : public InputMessenger { } }; +TEST(TransportHandshakeTest, rejected_handshake_start_releases_socket_reference) { + SocketOptions options; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + const int32_t references = socket->nref(); + int completions = 0; + auto done = [](int, void* data) { ++*static_cast(data); }; + + EXPECT_EQ(-1, AdapterTransport::StartClientUpgradeImpl( + socket.get(), done, &completions, RejectHandshakeStart)); + EXPECT_EQ(EAGAIN, errno); + EXPECT_EQ(references, socket->nref()); + // The caller owns the completion callback when startup fails. + EXPECT_EQ(0, completions); + socket->SetFailed(); +} + TEST(TransportHandshakeTest, registered_protocol_uses_transport_name) { GlobalInitializeOrDie(); EXPECT_EQ(1, PROTOCOL_TRANSPORT_HANDSHAKE); @@ -449,14 +473,25 @@ TEST(TransportHandshakeTest, blocked_write_stops_after_socket_failure) { bthread_t tid; ASSERT_EQ(0, bthread_start_background( &tid, nullptr, RunBlockingHandshakeWrite, &write)); - write.started.wait(); + bool was_waiting = false; + const int64_t wait_deadline = butil::gettimeofday_us() + 1000000; + while (butil::gettimeofday_us() < wait_deadline) { + // A nonzero wake count proves WriteAll reached WaitEpollOut's butex. + if (bthread::butex_wake(socket->_epollout_butex, true) > 0) { + was_waiting = true; + break; + } + bthread_usleep(1000); + } socket->SetFailed(EFAILEDSOCKET, "cancel blocked handshake write"); + bthread_flush(); const int completed = write.completed.timed_wait( butil::milliseconds_from_now(1000)); if (completed != 0) { shutdown(socket->fd(), SHUT_RDWR); } ASSERT_EQ(0, bthread_join(tid, nullptr)); + EXPECT_TRUE(was_waiting); EXPECT_EQ(0, completed); EXPECT_EQ(-1, write.result); EXPECT_EQ(EFAILEDSOCKET, write.error); @@ -646,6 +681,25 @@ TEST(TransportHandshakeTest, publish_fallback_after_tcp_state) { ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); } +TEST(TransportHandshakeTest, failed_handshake_preserves_io_error_after_cleanup) { + class CleanupChangesErrno : public TestHandshakeTransport { + public: + void OnFailed() override { + TestHandshakeTransport::OnFailed(); + errno = ENOENT; + } + } transport; + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + TestHandshakeProtocol protocol; + + EXPECT_EQ(STEP_ERROR, session.RunClient(&protocol, &transport)); + EXPECT_EQ(EIO, errno); + EXPECT_TRUE(transport.failed); + EXPECT_EQ(FAILED, session.phase(butil::memory_order_acquire)); +} + TEST(TransportHandshakeTest, client_runs_codec_and_resource_sequence) { MemoryHandshakeIO io("HSOK"); HandshakeSession session; From 9ab0f4dd64bd7390f7d3a99eed0864acf2b3e918 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Sun, 25 Oct 2026 23:54:59 -0700 Subject: [PATCH 09/17] Retrigger transport handshake CI From 0c36adbc2b011097f46bdb5483367869664d93d4 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Sun, 25 Oct 2026 23:56:27 -0700 Subject: [PATCH 10/17] Release URMA fallback buffers and preserve handshake errors --- src/brpc/adapter_transport.cpp | 2 +- src/brpc/urma_transport.cpp | 2 +- test/brpc_urma_unittest.cpp | 128 +++++++++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 2 deletions(-) diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index 8ac60da2be..03a6075e38 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -193,8 +193,8 @@ class UrmaClientHandshakeTransport : public handshake::HandshakeTransport { void OnEstablished() override { _transport->ActivateUpgrade(); } void OnFallback() override { _transport->DeactivateUpgrade(); } void OnFailed() override { - _transport->DeactivateUpgrade(); const int saved_errno = errno != 0 ? errno : EPROTO; + _transport->DeactivateUpgrade(); *_connect_error = saved_errno; _socket->SetFailed(saved_errno, "Fail to complete URMA handshake from %s: %s", diff --git a/src/brpc/urma_transport.cpp b/src/brpc/urma_transport.cpp index de485ab0d4..acd8d58cec 100644 --- a/src/brpc/urma_transport.cpp +++ b/src/brpc/urma_transport.cpp @@ -224,9 +224,9 @@ int UrmaTransport::StartUpgradeEvents() { void UrmaTransport::DeactivateUpgrade() { _urma_state.store(URMA_OFF, butil::memory_order_release); if (_urma_ep) { + _urma_ep->Reset(); _urma_ep->_state.store(urma::UrmaEndpoint::FALLBACK_TCP, butil::memory_order_release); - _urma_ep->DeallocateResources(); } } diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index dfcfa4730e..48c32a63a5 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -1024,6 +1024,134 @@ class UrmaMockTest : public ::testing::Test { bool _owns_urma_init{false}; }; +static void CheckUrmaFallbackCleanup(bool zerocopy) { + alarm(20); + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + gflags::SetCommandLineOption("urma_recv_zerocopy", zerocopy ? "true" : "false"); + urma::GlobalUrmaInitializeOrDie(); + ASSERT_NE(nullptr, urma::GetUrmaContext()); + for (int version : {2, 3}) { + urma::FLAGS_urma_client_handshake_version = version; + urma::UrmaHandshakeAdapter disabled(nullptr, version); + std::string payload; + ASSERT_EQ(handshake::STEP_OK, disabled.BuildHello(false, &payload)); + std::string frame; + ASSERT_EQ(handshake::FRAME_OK, handshake::FrameCodec::Encode( + disabled.HelloFrameSpec(), payload, &frame)); + for (int round = 0; round < 16; ++round) { + SocketOptions options; + options.socket_mode = SOCKET_MODE_URMA; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + auto* adapter = AdapterTransport::Get(socket.get()); + auto* transport = static_cast( + adapter->high_speed_transport()); + ScriptedUrmaHandshakeIO io; + io.input = frame; + adapter->handshake_session()->SetIOForTest(&io); + std::promise connected; + auto ready = connected.get_future(); + ASSERT_EQ(0, AdapterTransport::StartClientUpgrade(socket.get(), + [](int error, void* data) { + static_cast*>(data)->set_value(error); + }, &connected)); + ASSERT_EQ(std::future_status::ready, + ready.wait_for(std::chrono::seconds(3))); + EXPECT_EQ(0, ready.get()); + EXPECT_EQ(handshake::FALLBACK_TCP, adapter->handshake_phase()); + auto* endpoint = transport->GetUrmaEp(); + EXPECT_EQ(nullptr, endpoint->_resource); + EXPECT_TRUE(endpoint->_rbuf.empty()); + EXPECT_TRUE(endpoint->_sbuf.empty()); + EXPECT_TRUE(endpoint->_rbuf_data.empty()); + EXPECT_TRUE(endpoint->_recv_completions.empty()); + EXPECT_TRUE(endpoint->_recv_completed.empty()); + EXPECT_EQ(0, endpoint->_local_window_capacity); + EXPECT_EQ(0, endpoint->_remote_window_capacity); + EXPECT_EQ(urma::UrmaEndpoint::FALLBACK_TCP, + endpoint->_state.load(butil::memory_order_acquire)); + socket->SetFailed(); + } + } +} + +TEST_F(UrmaMockTest, repeated_copy_fallback_releases_preposted_buffers) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT({ + CheckUrmaFallbackCleanup(false); + _exit(::testing::Test::HasFailure() ? 1 : 0); + }, ::testing::ExitedWithCode(0), ""); +} + +TEST_F(UrmaMockTest, repeated_zerocopy_fallback_releases_preposted_buffers) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT({ + CheckUrmaFallbackCleanup(true); + _exit(::testing::Test::HasFailure() ? 1 : 0); + }, ::testing::ExitedWithCode(0), ""); +} + +class CleanupErrnoUrmaHandshakeIO : public CapturingUrmaHandshakeIO { +public: + explicit CleanupErrnoUrmaHandshakeIO(urma::UrmaEndpoint* ep) : endpoint(ep) {} + int ReadExact(void*, size_t) override { + errno = EIO; + return -1; + } + int WriteAll(const void* data, size_t len) override { + endpoint->_sbuf.emplace_back(); + endpoint->_sbuf.back().append_user_data(new int(0), sizeof(int), + [](void* data) { + delete static_cast(data); + errno = ENOENT; + }); + return CapturingUrmaHandshakeIO::WriteAll(data, len); + } + urma::UrmaEndpoint* endpoint; +}; + +static void CheckUrmaFailureErrno() { + alarm(10); + gflags::SetCommandLineOption("urma_prepared_jetty_cnt", "0"); + gflags::SetCommandLineOption("urma_buffer_count", "2048"); + urma::GlobalUrmaInitializeOrDie(); + SocketOptions options; + options.socket_mode = SOCKET_MODE_URMA; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + auto* adapter = AdapterTransport::Get(socket.get()); + auto* endpoint = static_cast( + adapter->high_speed_transport())->GetUrmaEp(); + CleanupErrnoUrmaHandshakeIO io(endpoint); + adapter->handshake_session()->SetIOForTest(&io); + std::promise connected; + auto ready = connected.get_future(); + ASSERT_EQ(0, AdapterTransport::StartClientUpgrade(socket.get(), + [](int error, void* data) { + static_cast*>(data)->set_value(error); + }, &connected)); + ASSERT_EQ(std::future_status::ready, + ready.wait_for(std::chrono::seconds(3))); + EXPECT_EQ(EIO, ready.get()); + EXPECT_EQ(handshake::FAILED, adapter->handshake_phase()); + EXPECT_TRUE(endpoint->_sbuf.empty()); + EXPECT_TRUE(endpoint->_rbuf.empty()); + socket->SetFailed(); +} + +TEST_F(UrmaMockTest, client_failure_preserves_errno_across_buffer_cleanup) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT({ + CheckUrmaFailureErrno(); + _exit(::testing::Test::HasFailure() ? 1 : 0); + }, ::testing::ExitedWithCode(0), ""); +} + TEST_F(UrmaMockTest, send_completion_releases_source_only_after_copy) { int device_count = 0; urma_device_t** devices = urma_get_device_list(&device_count); From 071ee1ba6f8732920144979f053b3be5b26cc8f3 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Sun, 25 Oct 2026 23:57:44 -0700 Subject: [PATCH 11/17] Move legacy URMA handshake peers into interoperability tests --- src/brpc/handshake/urma_handshake.cpp | 224 ++----- src/brpc/handshake/urma_handshake.h | 103 +--- src/brpc/socket.h | 10 - src/brpc/urma/urma_endpoint.cpp | 543 +---------------- src/brpc/urma/urma_endpoint.h | 117 +--- src/brpc/urma_transport.cpp | 4 +- src/brpc/urma_transport.h | 7 +- test/BUILD.bazel | 8 + test/brpc_urma_legacy_handshake.h | 843 ++++++++++++++++++++++++++ test/brpc_urma_unittest.cpp | 15 +- 10 files changed, 933 insertions(+), 941 deletions(-) create mode 100644 test/brpc_urma_legacy_handshake.h diff --git a/src/brpc/handshake/urma_handshake.cpp b/src/brpc/handshake/urma_handshake.cpp index 0cd7e778c0..0a4577f761 100644 --- a/src/brpc/handshake/urma_handshake.cpp +++ b/src/brpc/handshake/urma_handshake.cpp @@ -243,6 +243,38 @@ void TranslateV3Hello(const UrmaHello& message, ParsedHello* out) { out->seg_token_id = message.seg_token_id(); } +void FillV2Hello(const ParsedHello& p, v2_wire::HelloMessage* out) { + *out = v2_wire::HelloMessage{}; // value-initialize + out->msg_len = v2_wire::HELLO_PACKET_LEN; + out->hello_ver = v2_wire::HELLO_V2_VERSION; + out->impl_ver = v2_wire::IMPL_V2_VERSION; + out->buffer_size = p.buffer_size; + out->recv_buffer_cnt = p.recv_buffer_cnt; + out->jetty_id = p.jetty_id; + std::memcpy(out->eid, p.eid, 16); + out->uasid = p.uasid; + out->tp_type = p.tp_type; + std::memcpy(out->seg_eid, p.seg_eid, 16); + out->seg_uasid = p.seg_uasid; + out->seg_va = p.seg_va; + out->seg_len = p.seg_len; + out->seg_token_id = p.seg_token_id; +} + +void FillV3Hello(const ParsedHello& p, UrmaHello* out) { + out->set_buffer_size(p.buffer_size); + out->set_recv_buffer_cnt(p.recv_buffer_cnt); + out->set_jetty_id(p.jetty_id); + out->set_eid(p.eid, 16); + out->set_uasid(p.uasid); + out->set_tp_type(p.tp_type); + out->set_seg_eid(p.seg_eid, 16); + out->set_seg_uasid(p.seg_uasid); + out->set_seg_va(p.seg_va); + out->set_seg_len(p.seg_len); + out->set_seg_token_id(p.seg_token_id); +} + } // namespace const handshake::FrameSpec& UrmaHandshakeAdapter::HelloFrameSpec() const { @@ -255,31 +287,20 @@ const handshake::FrameSpec& UrmaHandshakeAdapter::AckFrameSpec() const { handshake::StepResult UrmaHandshakeAdapter::BuildHello( bool enabled, std::string* payload) { + ParsedHello local; + if (enabled) { + CHECK(_ep != nullptr); + _ep->GetLocalConnectionInfo(&local); + } if (_version == 3) { UrmaHello message; - if (enabled) { - CHECK(_ep != nullptr); - _ep->FillLocalHelloV3(&message); - } else { - message.set_buffer_size(0); - message.set_recv_buffer_cnt(0); - message.set_jetty_id(0); - message.set_eid(std::string(16, '\0')); - message.set_uasid(0); - message.set_tp_type(0); - message.set_seg_eid(std::string(16, '\0')); - message.set_seg_uasid(0); - message.set_seg_va(0); - message.set_seg_len(0); - message.set_seg_token_id(0); - } + FillV3Hello(local, &message); return message.SerializeToString(payload) ? handshake::STEP_OK : handshake::STEP_ERROR; } v2_wire::HelloMessage message{}; if (enabled) { - CHECK(_ep != nullptr); - _ep->FillLocalHelloV2(&message); + FillV2Hello(local, &message); } else { message.msg_len = v2_wire::HELLO_PACKET_LEN; } @@ -348,173 +369,6 @@ CreateServerHandshakeAdapters(UrmaEndpoint* ep) { return adapters; } -// File-local (not in the anonymous namespace so it can be friend-declared -// from urma_endpoint.h's UrmaEndpoint). Reads the body following the magic -// and translates it into ParsedHello. -int ReadBodyAndNegotiate(UrmaEndpoint* ep, ParsedHello* out, bool* negotiated) { - *negotiated = false; - uint8_t body[v2_wire::HELLO_BODY_LEN]; - if (ep->ReadFromFd(body, v2_wire::HELLO_BODY_LEN) < 0) { - return -1; - } - v2_wire::HelloMessage m; - m.Deserialize(body); - if (m.msg_len < v2_wire::HELLO_MSG_LEN_MIN || - m.msg_len > v2_wire::HELLO_MSG_LEN_MAX || - m.hello_ver != v2_wire::HELLO_V2_VERSION || - m.impl_ver != v2_wire::IMPL_V2_VERSION) { return 0; } - ParsedHello p; - p.buffer_size = m.buffer_size; - p.recv_buffer_cnt = m.recv_buffer_cnt; - p.jetty_id = m.jetty_id; - std::memcpy(p.eid, m.eid, 16); - p.uasid = m.uasid; - p.tp_type = m.tp_type; - std::memcpy(p.seg_eid, m.seg_eid, 16); - p.seg_uasid = m.seg_uasid; - p.seg_va = m.seg_va; - p.seg_len = m.seg_len; - p.seg_token_id = m.seg_token_id; - if (!ValidHello(p)) { - return 0; - } - // Drain trailing bytes if msg_len advertises more than the fixed body. - if (m.msg_len > v2_wire::HELLO_PACKET_LEN) { - if (DrainBytes(ep, m.msg_len - v2_wire::HELLO_PACKET_LEN) < 0) { - return -1; - } - } - *out = p; - *negotiated = true; - return 0; -} - -int DrainBytes(UrmaEndpoint* ep, size_t n) { - char buf[4096]; - while (n > 0) { - size_t want = std::min(n, sizeof(buf)); - if (ep->ReadFromFd(buf, want) < 0) { - return -1; - } - n -= want; - } - return 0; -} - -// ============================================================================ -// v2 client / server. -// ============================================================================ - -int UrmaHandshakeClientV2::SendLocalHello() { - v2_wire::HelloMessage m; - _ep->FillLocalHelloV2(&m); - uint8_t packet[v2_wire::HELLO_PACKET_LEN]; - std::memcpy(packet, "URMA", 4); - m.Serialize(packet + 4); - return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); -} - -int UrmaHandshakeClientV2::ReceiveAndParseRemoteHello(ParsedHello* out, - bool* negotiated) { - *negotiated = false; - uint8_t magic[v2_wire::MAGIC_STR_LEN]; - if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { - return -1; - } - if (std::memcmp(magic, "URMA", 4) != 0) { - // Peer is not URMA-capable; push the magic back so the TCP input - // messenger can re-parse it. - _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); - return 0; - } - return ReadBodyAndNegotiate(_ep, out, negotiated); -} - -int UrmaHandshakeServerV2::ReceiveAndParseRemoteHello(ParsedHello* out, - bool* negotiated) { - return ReadBodyAndNegotiate(_ep, out, negotiated); -} - -int UrmaHandshakeServerV2::SendLocalHello() { - v2_wire::HelloMessage m; - _ep->FillLocalHelloV2(&m); - auto* tp = static_cast(_ep->_socket->_transport.get()); - if (tp->_urma_state.load(butil::memory_order_acquire) == - UrmaTransport::URMA_OFF) { - // Tell the client we are not URMA-capable: zero the version fields so - // the client's version check fails and it falls back to TCP. - m.hello_ver = 0; - m.impl_ver = 0; - m.jetty_id = 0; - m.buffer_size = 0; - } - uint8_t packet[v2_wire::HELLO_PACKET_LEN]; - std::memcpy(packet, "URMA", 4); - m.Serialize(packet + 4); - return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); -} - -// ============================================================================ -// v3 protobuf ("URM3"). The protobuf (de)serialization lives on the endpoint -// because it touches UrmaEndpoint private state; the classes here just call -// FillLocalHelloV3 / WriteHelloV3 / ReadAndParseHelloV3. -// ============================================================================ - -int UrmaHandshakeClientV3::SendLocalHello() { - UrmaHello msg; - _ep->FillLocalHelloV3(&msg); - return _ep->WriteHelloV3(msg); -} - -int UrmaHandshakeClientV3::ReceiveAndParseRemoteHello(ParsedHello* out, - bool* negotiated) { - *negotiated = false; - uint8_t magic[v2_wire::MAGIC_STR_LEN]; - if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { - return -1; - } - if (std::memcmp(magic, "URM3", 4) != 0) { - _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); - return 0; - } - return _ep->ReadAndParseHelloV3(out, negotiated); -} - -int UrmaHandshakeServerV3::ReceiveAndParseRemoteHello(ParsedHello* out, - bool* negotiated) { - return _ep->ReadAndParseHelloV3(out, negotiated); -} - -int UrmaHandshakeServerV3::SendLocalHello() { - UrmaHello msg; - _ep->FillLocalHelloV3(&msg); - // v3 has no zero-out path; the client rejects jetty_id==0 via ValidHello. - return _ep->WriteHelloV3(msg); -} - -// ============================================================================ -// Factories. -// ============================================================================ - -UrmaHandshake* CreateClientHandshake(UrmaEndpoint* ep) { - switch (FLAGS_urma_client_handshake_version) { - case 3: return new UrmaHandshakeClientV3(ep); - case 2: - default: return new UrmaHandshakeClientV2(ep); - } -} - -UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, - const uint8_t magic[v2_wire::MAGIC_STR_LEN]) { - if (std::memcmp(magic, "URMA", 4) == 0) { - return new UrmaHandshakeServerV2(ep); - } - if (std::memcmp(magic, "URM3", 4) == 0) { - return new UrmaHandshakeServerV3(ep); - } - return nullptr; -} - } // namespace urma namespace handshake { diff --git a/src/brpc/handshake/urma_handshake.h b/src/brpc/handshake/urma_handshake.h index a6ca94acea..52256b1fc0 100644 --- a/src/brpc/handshake/urma_handshake.h +++ b/src/brpc/handshake/urma_handshake.h @@ -37,7 +37,7 @@ StepResult RunUrmaFallbackServerHandshake(butil::IOBuf* source, Socket* socket); #if BRPC_WITH_URMA -#include "urma_types.h" +#include "brpc/urma/urma_endpoint.h" #include "brpc/transport_handshake.h" namespace brpc { @@ -45,25 +45,7 @@ namespace urma { class UrmaEndpoint; -// Wire-format-agnostic view of a peer's hello message. Both v2 (binary) and -// v3 (protobuf) translate the bytes on the wire into this struct. The endpoint -// then uses ApplyRemoteHello to size its send/recv windows and to drive -// urma_import_seg / urma_import_jetty. -struct ParsedHello { - uint32_t buffer_size = 0; // peer recv buffer size in bytes (per WR) - uint32_t recv_buffer_cnt = 0; // peer recv buffer count (RQ depth - 1) - uint32_t jetty_id = 0; // peer jetty id - uint8_t eid[16] = {0}; // peer EID (network order) - uint32_t uasid = 0; // peer uasid - uint8_t tp_type = 0; // urma_tp_type_t: URMA_RTP / URMA_CTP / URMA_UTP - - // Flattened peer buffer-pool segment. - uint8_t seg_eid[16] = {0}; // EID owning the peer segment - uint32_t seg_uasid = 0; // uasid owning the peer segment - uint64_t seg_va = 0; // segment virtual address - uint64_t seg_len = 0; // segment length in bytes - uint32_t seg_token_id = 0; // segment token id -}; +using ParsedHello = UrmaConnectionInfo; // Validate all values that influence resource import and queue/window sizing. // Both v2 and v3 handshakes must pass this check before negotiation succeeds. @@ -108,24 +90,6 @@ struct HelloMessage { } // namespace v2_wire -// Legacy TCP handshake strategy used by interoperability tests. -// v2 speaks binary; v3 speaks protobuf ("URM3"). -class UrmaHandshake { -public: - virtual ~UrmaHandshake() = default; - virtual int ProtocolVersion() const = 0; - - // Send our local hello over the TCP fd held by the endpoint. - // Returns 0 on success, -1 on failure (errno set). - virtual int SendLocalHello() = 0; - - // Read the peer's hello and parse it into @out. Sets @negotiated to true - // if the peer is URMA-capable and the message validates; false (with - // return 0) to signal a graceful fall-back to TCP. - // Returns -1 on IO error (errno set). - virtual int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) = 0; -}; - // Wire-format adapter used by the common transport handshake. // Stores the selected protocol version and parsed peer hello; // UrmaEndpoint owns the local and imported URMA resources. @@ -154,69 +118,6 @@ std::unique_ptr CreateClientHandshakeAdapter( std::vector > CreateServerHandshakeAdapters(UrmaEndpoint* ep); -// v2 binary handshake (magic "URMA"). -class UrmaHandshakeClientV2 : public UrmaHandshake { -public: - explicit UrmaHandshakeClientV2(UrmaEndpoint* ep) : _ep(ep) {} - int ProtocolVersion() const override { return 2; } - int SendLocalHello() override; - int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; - -private: - UrmaEndpoint* _ep; -}; - -class UrmaHandshakeServerV2 : public UrmaHandshake { -public: - explicit UrmaHandshakeServerV2(UrmaEndpoint* ep) : _ep(ep) {} - int ProtocolVersion() const override { return 2; } - int SendLocalHello() override; - int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; - -private: - UrmaEndpoint* _ep; -}; - -// v3 protobuf handshake (magic "URM3"). -class UrmaHandshakeClientV3 : public UrmaHandshake { -public: - explicit UrmaHandshakeClientV3(UrmaEndpoint* ep) : _ep(ep) {} - int ProtocolVersion() const override { return 3; } - int SendLocalHello() override; - int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; - -private: - UrmaEndpoint* _ep; -}; - -class UrmaHandshakeServerV3 : public UrmaHandshake { -public: - explicit UrmaHandshakeServerV3(UrmaEndpoint* ep) : _ep(ep) {} - int ProtocolVersion() const override { return 3; } - int SendLocalHello() override; - int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; - -private: - UrmaEndpoint* _ep; -}; - -// Legacy client factory: picks v2/v3 via --urma_client_handshake_version. -UrmaHandshake* CreateClientHandshake(UrmaEndpoint* ep); - -// Legacy server factory: dispatches on the first 4 bytes read from the TCP fd. -// Returns nullptr if the magic is not "URMA"/"URM3" (caller falls back to TCP). -// @magic is the first MAGIC_STR_LEN bytes already read from the fd. -UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, - const uint8_t magic[v2_wire::MAGIC_STR_LEN]); - -// Drain @n bytes from the TCP fd (used when msg_len > HELLO_MSG_LEN_MIN). -int DrainBytes(UrmaEndpoint* ep, size_t n); - -// Read the body following the magic and translate to ParsedHello. Friend of -// UrmaEndpoint. Returns -1 on IO error; 0 with *negotiated=false on invalid -// (peer not URMA-capable); 0 with *negotiated=true on success. -int ReadBodyAndNegotiate(UrmaEndpoint* ep, ParsedHello* out, bool* negotiated); - } // namespace urma } // namespace brpc diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 194b8fd2ab..d7bb0db74f 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -60,11 +60,6 @@ class RdmaEndpoint; namespace urma { class UrmaEndpoint; -class UrmaConnect; -class UrmaHandshakeClientV2; -class UrmaHandshakeServerV2; -class UrmaHandshakeClientV3; -class UrmaHandshakeServerV3; } namespace ubring { @@ -342,11 +337,6 @@ friend class rdma::RdmaEndpoint; friend class ubring::UBShmEndpoint; friend class UBShmTransport; friend class urma::UrmaEndpoint; -friend class urma::UrmaConnect; -friend class urma::UrmaHandshakeClientV2; -friend class urma::UrmaHandshakeServerV2; -friend class urma::UrmaHandshakeClientV3; -friend class urma::UrmaHandshakeServerV3; friend class HealthCheckTask; friend class OnAppHealthCheckDone; friend class HealthCheckManager; diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp index 9f8593a1ae..37af9310ca 100644 --- a/src/brpc/urma/urma_endpoint.cpp +++ b/src/brpc/urma/urma_endpoint.cpp @@ -20,7 +20,6 @@ #if BRPC_WITH_URMA #include -#include #include #include @@ -38,18 +37,14 @@ #include "butil/iobuf.h" #include "butil/logging.h" #include "butil/macros.h" -#include "butil/sys_byteorder.h" #include "butil/time.h" #include "bthread/bthread.h" -#include "bthread/butex.h" #include "urma_api.h" #include "brpc/input_messenger.h" #include "brpc/socket.h" #include "brpc/urma/urma_bonding.h" -#include "brpc/handshake/urma_handshake.h" -#include "brpc/urma/urma_handshake.pb.h" #include "brpc/urma/urma_helper.h" #include "brpc/urma/urma_stream_trace.h" #include "brpc/urma_transport.h" @@ -74,10 +69,6 @@ DEFINE_bool(urma_trace_stream, false, "parse/completion/post failure or teardown. Includes payload prefixes; " "enable before creating connections only for controlled diagnostics."); -// ---- Constants shared with the handshake module ---- -static const int WAIT_TIMEOUT_MS = 50; -static const size_t HELLO_ACK_LEN = 4; -static const uint32_t HELLO_ACK_URMA_OK = 0x1; static const size_t IOBUF_BLOCK_HEADER_LEN = sizeof(butil::IOBuf::Block); // Match UMDK perftest's CTP SEND ceiling. This is a per-WR payload limit, // independent of the pool block size and the device's generic max_msg_size. @@ -155,15 +146,12 @@ UrmaResource::~UrmaResource() { UrmaEndpoint::UrmaEndpoint(Socket* s) : _socket(s), _state(UNINIT), - _handshake_version(0), _resource(nullptr), _cq_lifetime(std::make_shared(s->id())) { _sq_size = static_cast( std::max(16, std::min(4096, static_cast(FLAGS_urma_sq_size)))); _rq_size = static_cast( std::max(16, std::min(4096, static_cast(FLAGS_urma_rq_size)))); - _read_butex = bthread::butex_create_checked>(); - _read_butex->store(0, butil::memory_order_relaxed); _input_processor.Init(s, InputMessengerProcessor::STREAM_URMA_JETTY); if (FLAGS_urma_trace_stream) { _stream_trace.reset(new UrmaStreamTrace); @@ -172,17 +160,12 @@ UrmaEndpoint::UrmaEndpoint(Socket* s) UrmaEndpoint::~UrmaEndpoint() { DeallocateResources(); - if (_read_butex) { - bthread::butex_destroy(_read_butex); - _read_butex = nullptr; - } } void UrmaEndpoint::Reset() { DeallocateResources(); _stream_trace.reset(FLAGS_urma_trace_stream ? new UrmaStreamTrace : nullptr); _state = UNINIT; - _handshake_version = 0; _remote_recv_block_size = 0; _local_window_capacity = 0; _remote_window_capacity = 0; @@ -200,72 +183,14 @@ void UrmaEndpoint::Reset() { _rbuf_data.clear(); _recv_completions.clear(); _recv_completed.clear(); - _read_butex->store(0, butil::memory_order_relaxed); } // ============================================================================ -// Handshake IO helpers (ReadFromFd / WriteToFd / PushBackToReadBuf). -// Modeled on RdmaEndpoint::ReadFromFdLoop / WriteToFdLoop. +// Local connection parameters. // ============================================================================ -int UrmaEndpoint::ReadFromFd(void* data, size_t len) { - char* p = static_cast(data); - 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); - const int fd = _socket->fd(); - const ssize_t nr = read(fd, p + received, len - received); - if (nr < 0) { - if (errno == EAGAIN || errno == EWOULDBLOCK) { - int rc = bthread::butex_wait(_read_butex, expected_val, &duetime); - if (rc < 0 && errno != EWOULDBLOCK && errno != ETIMEDOUT) { - return -1; - } - continue; - } - return -1; - } - if (nr == 0) { - errno = EEOF; - return -1; - } - received += nr; - } - return 0; -} - -void UrmaEndpoint::PushBackToReadBuf(const void* data, size_t len) { - _socket->fd_input_processor().read_buf().append(data, len); -} - -int UrmaEndpoint::WriteToFd(void* data, size_t len) { - char* p = static_cast(data); - size_t written = 0; - while (written < len) { - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - const int fd = _socket->fd(); - const ssize_t nw = write(fd, p + written, len - written); - if (nw >= 0) { - written += nw; - continue; - } - if (errno != EAGAIN && errno != EWOULDBLOCK) { - return -1; - } - if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { - return -1; - } - } - return 0; -} - -// ============================================================================ -// Hello builders / parsers. -// ============================================================================ - -void UrmaEndpoint::MakeLocalParsedHello(ParsedHello* out) const { - *out = ParsedHello{}; // value-initialize (avoids memset on non-trivial type) +void UrmaEndpoint::GetLocalConnectionInfo(UrmaConnectionInfo* out) const { + *out = UrmaConnectionInfo{}; out->buffer_size = static_cast(GetUrmaRecvBlockSize()); out->recv_buffer_cnt = _rq_size - 1; if (_resource && _resource->jetty) { @@ -290,114 +215,6 @@ void UrmaEndpoint::MakeLocalParsedHello(ParsedHello* out) const { } } -void UrmaEndpoint::FillLocalHelloV2(v2_wire::HelloMessage* out) const { - *out = v2_wire::HelloMessage{}; // value-initialize - out->msg_len = v2_wire::HELLO_PACKET_LEN; - out->hello_ver = v2_wire::HELLO_V2_VERSION; - out->impl_ver = v2_wire::IMPL_V2_VERSION; - ParsedHello p; - MakeLocalParsedHello(&p); - out->buffer_size = p.buffer_size; - out->recv_buffer_cnt = p.recv_buffer_cnt; - out->jetty_id = p.jetty_id; - std::memcpy(out->eid, p.eid, 16); - out->uasid = p.uasid; - out->tp_type = p.tp_type; - std::memcpy(out->seg_eid, p.seg_eid, 16); - out->seg_uasid = p.seg_uasid; - out->seg_va = p.seg_va; - out->seg_len = p.seg_len; - out->seg_token_id = p.seg_token_id; -} - -void UrmaEndpoint::FillLocalHelloV3(UrmaHello* out) const { - ParsedHello p; - MakeLocalParsedHello(&p); - out->set_buffer_size(p.buffer_size); - out->set_recv_buffer_cnt(p.recv_buffer_cnt); - out->set_jetty_id(p.jetty_id); - out->set_eid(p.eid, 16); - out->set_uasid(p.uasid); - out->set_tp_type(p.tp_type); - out->set_seg_eid(p.seg_eid, 16); - out->set_seg_uasid(p.seg_uasid); - out->set_seg_va(p.seg_va); - out->set_seg_len(p.seg_len); - out->set_seg_token_id(p.seg_token_id); -} - -int UrmaEndpoint::WriteHelloV3(const UrmaHello& msg) { - butil::IOBuf packet; - packet.append("URM3", 4); - std::string body; - if (!msg.SerializeToString(&body)) { - LOG(ERROR) << "Fail to serialize UrmaHello"; - return -1; - } - uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); - packet.append(&pb_size_be, sizeof(pb_size_be)); - packet.append(body); - return WriteToFd(packet); -} - -int UrmaEndpoint::WriteToFd(butil::IOBuf& data) { - // Write out the IOBuf in a single WriteToFd-style loop. - while (!data.empty()) { - const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); - const int fd = _socket->fd(); - const ssize_t nw = data.cut_into_file_descriptor(fd); - if (nw >= 0) { - continue; - } - if (errno != EAGAIN && errno != EWOULDBLOCK) { - return -1; - } - if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { - return -1; - } - } - return 0; -} - -int UrmaEndpoint::ReadAndParseHelloV3(ParsedHello* out, bool* negotiated) { - *negotiated = false; - uint32_t pb_size_be = 0; - if (ReadFromFd(&pb_size_be, sizeof(pb_size_be)) < 0) { - return -1; - } - const uint32_t pb_size = butil::NetToHost32(pb_size_be); - if (pb_size == 0 || pb_size > 4096) { - return 0; - } - std::string body(pb_size, '\0'); - if (ReadFromFd(&body[0], pb_size) < 0) { - return -1; - } - UrmaHello msg; - if (!msg.ParseFromArray(body.data(), static_cast(body.size()))) { - return 0; - } - if (msg.eid().size() != 16 || msg.seg_eid().size() != 16) { - return 0; - } - out->buffer_size = msg.buffer_size(); - out->recv_buffer_cnt = msg.recv_buffer_cnt(); - out->jetty_id = msg.jetty_id(); - std::memcpy(out->eid, msg.eid().data(), 16); - out->uasid = msg.uasid(); - out->tp_type = static_cast(msg.tp_type()); - std::memcpy(out->seg_eid, msg.seg_eid().data(), 16); - out->seg_uasid = msg.seg_uasid(); - out->seg_va = msg.seg_va(); - out->seg_len = msg.seg_len(); - out->seg_token_id = msg.seg_token_id(); - if (!ValidHello(*out)) { - return 0; - } - *negotiated = true; - return 0; -} - // ============================================================================ // Allocate / deallocate per-connection resources. // ============================================================================ @@ -583,7 +400,7 @@ void UrmaEndpoint::DeallocateResources() { // ImportPeer: the critical import_seg-before-import_jetty sequence. // ============================================================================ -int UrmaEndpoint::ImportPeer(const ParsedHello& peer) { +int UrmaEndpoint::ImportPeer(const UrmaConnectionInfo& peer) { urma_context_t* ctx = GetUrmaContext(); if (!ctx) { errno = ENODEV; @@ -1352,10 +1169,10 @@ void UrmaEndpoint::PollCq(Socket* m) { } // ============================================================================ -// ApplyRemoteHello: size the send/recv windows from the peer's hello. +// ApplyRemoteInfo: size the send/recv windows from peer connection parameters. // ============================================================================ -void UrmaEndpoint::ApplyRemoteHello(const ParsedHello& remote) { +void UrmaEndpoint::ApplyRemoteInfo(const UrmaConnectionInfo& remote) { _remote_recv_block_size = remote.buffer_size; const uint32_t peer_rq_size = remote.recv_buffer_cnt + 1; const uint32_t local_capacity = @@ -1372,343 +1189,6 @@ void UrmaEndpoint::ApplyRemoteHello(const ParsedHello& remote) { _sq_window_size.store(_local_window_capacity, butil::memory_order_relaxed); } -// ============================================================================ -// OnNewDataFromTcp: edge-triggered dispatcher. -// ============================================================================ - -static void TryReadOnTcpDuringUrmaEst(Socket* socket); - -void UrmaEndpoint::OnNewDataFromTcp(Socket* m) { - auto* tp = static_cast(m->_transport.get()); - if (!tp) { - return; - } - // Access _urma_ep directly (OnNewDataFromTcp is a friend of UrmaTransport); - // GetUrmaEp() CHECKs non-null which would crash on TCP-fallback sockets. - UrmaEndpoint* ep = tp->_urma_ep; - if (!ep) { - // No URMA endpoint: pure TCP path. - InputMessenger::OnNewMessages(m); - return; - } - int progress = 0; - while (true) { - const State state = - ep->_state.load(butil::memory_order_acquire); - if (state == UNINIT) { - if (!m->CreatedByConnect()) { - // Server side: kick off the handshake bthread. - if (!IsUrmaAvailable()) { - ep->_state = FALLBACK_TCP; - tp->_urma_state = UrmaTransport::URMA_OFF; - InputMessenger::OnNewMessages(m); - return; - } - SocketUniquePtr s; - m->ReAddress(&s); - ep->_state = S_HELLO_WAIT; - bthread_t tid; - bthread_attr_t attr = BTHREAD_ATTR_NORMAL; - bthread_attr_set_name(&attr, "UrmaServerHandshake"); - if (bthread_start_background(&tid, &attr, - ProcessHandshakeAtServer, ep) != 0) { - ep->_state = UNINIT; - LOG(FATAL) << "Fail to start UrmaServerHandshake bthread"; - } else { - s.release(); - } - return; - } - // Client side: handled by ProcessHandshakeAtClient. - return; - } else if (state < ESTABLISHED) { - // During handshake: wake the handshake bthread parked in ReadFromFd. - ep->_read_butex->fetch_add(1, butil::memory_order_release); - bthread::butex_wake(ep->_read_butex); - return; - } else if (state == FALLBACK_TCP) { - InputMessenger::OnNewMessages(m); - return; - } else if (state == ESTABLISHED) { - TryReadOnTcpDuringUrmaEst(m); - return; - } - if (!m->MoreReadEvents(&progress)) { - break; - } - } -} - -inline void UrmaEndpoint::TryReadOnTcp() { - if (_state.load(butil::memory_order_acquire) == FALLBACK_TCP) { - InputMessenger::OnNewMessages(_socket); - } -} - -static void TryReadOnTcpDuringUrmaEst(Socket* socket) { - int progress = Socket::PROGRESS_INIT; - while (true) { - uint8_t byte = 0; - const ssize_t nr = read(socket->fd(), &byte, 1); - if (nr < 0) { - if (errno != EAGAIN) { - const int saved_errno = errno; - socket->SetFailed(saved_errno, "Fail to read URMA TCP fd: %s", - berror(saved_errno)); - return; - } - if (!socket->MoreReadEvents(&progress)) { - return; - } - } else if (nr == 0) { - socket->SetEOF(); - return; - } else { - socket->SetFailed( - EPROTO, "Unexpected TCP data after URMA was established"); - return; - } - } -} - -void UrmaEndpoint::FallbackToTcp(UrmaTransport* transport, bool process_tcp) { - transport->_urma_state = UrmaTransport::URMA_OFF; - _state.store(FALLBACK_TCP, butil::memory_order_release); - DeallocateResources(); - if (process_tcp) { - TryReadOnTcp(); - } -} - -void UrmaEndpoint::FailHandshake(UrmaTransport* transport, int error, - const char* reason) { - LOG(ERROR) << "URMA handshake failed in state=" << GetStateStr() - << " on " << _socket->description() - << ": " << reason << ", error=" << error - << " (" << berror(error) << ')'; - transport->_urma_state = UrmaTransport::URMA_OFF; - _state.store(FAILED, butil::memory_order_release); - DeallocateResources(); - auto* connect = - static_cast(_socket->_app_connect.get()); - if (connect) { - connect->_error = error; - } - _socket->SetFailed(error, "URMA handshake failed: %s", reason); -} - -// ============================================================================ -// Handshake state machines (client / server). Run in a background bthread. -// ============================================================================ - -void* UrmaEndpoint::ProcessHandshakeAtClient(void* arg) { - auto* ep = static_cast(arg); - SocketUniquePtr s(ep->_socket); - auto* tp = static_cast(s->_transport.get()); - UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); - if (!IsUrmaAvailable()) { - ep->FallbackToTcp(tp, true); - return nullptr; - } - ep->_state = C_ALLOC_RES; - if (ep->AllocateResources() < 0) { - ep->FallbackToTcp(tp, true); - return nullptr; - } - // Prepost the shared JFR before sending the client hello so the peer sees - // a ready receive queue as soon as its import completes. - if (ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { - ep->FallbackToTcp(tp, true); - return nullptr; - } - ep->_state = C_HELLO_SEND; - std::unique_ptr hs(CreateClientHandshake(ep)); - ep->_handshake_version = hs->ProtocolVersion(); - if (hs->SendLocalHello() < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "send client hello"); - return nullptr; - } - ep->_state = C_HELLO_WAIT; - ParsedHello remote; - bool negotiated = false; - if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "read server hello"); - return nullptr; - } - if (!negotiated) { - ep->FallbackToTcp(tp, true); - return nullptr; - } - ep->ApplyRemoteHello(remote); - ep->_state = C_IMPORT_PEER; - if (ep->ImportPeer(remote) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "import server resources"); - return nullptr; - } - ep->_state = C_ACK_SEND; - uint32_t flags = HELLO_ACK_URMA_OK; - uint32_t flags_be = butil::HostToNet32(flags); - if (ep->WriteToFd(&flags_be, HELLO_ACK_LEN) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "send client ack"); - return nullptr; - } - tp->_urma_state = UrmaTransport::URMA_ON; - ep->_state = ESTABLISHED; - if (ep->StartCqEvents() < 0) { - ep->FailHandshake(tp, errno, "start client CQ events"); - return nullptr; - } - return nullptr; -} - -void* UrmaEndpoint::ProcessHandshakeAtServer(void* arg) { - auto* ep = static_cast(arg); - SocketUniquePtr s(ep->_socket); - auto* tp = static_cast(s->_transport.get()); - UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); - ep->_state = S_HELLO_WAIT; - uint8_t magic[v2_wire::MAGIC_STR_LEN]; - if (ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "read client magic"); - return nullptr; - } - std::unique_ptr hs(CreateServerHandshakeByMagic(ep, magic)); - if (!hs) { - // Not an URMA peer: push the magic back and fall back to TCP. - ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); - ep->FallbackToTcp(tp, true); - return nullptr; - } - ep->_handshake_version = hs->ProtocolVersion(); - ParsedHello remote; - bool negotiated = false; - if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "read client hello"); - return nullptr; - } - if (!negotiated) { - ep->FailHandshake(tp, EPROTO, "invalid client hello"); - return nullptr; - } - ep->_state = S_ALLOC_RES; - if (ep->AllocateResources() < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "allocate server resources"); - return nullptr; - } - const bool bonding = IsUrmaBondingDevice(); - if (!bonding && - ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "post server receives"); - return nullptr; - } - ep->ApplyRemoteHello(remote); - ep->_state = S_IMPORT_PEER; - if (ep->ImportPeer(remote) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "import client resources"); - return nullptr; - } - if (bonding && - ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "post server receives"); - return nullptr; - } - ep->_state = S_HELLO_SEND; - if (hs->SendLocalHello() < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "send server hello"); - return nullptr; - } - ep->_state = S_ACK_WAIT; - uint32_t flags_be = 0; - if (ep->ReadFromFd(&flags_be, HELLO_ACK_LEN) < 0) { - const int saved_errno = errno ? errno : EIO; - ep->FailHandshake(tp, saved_errno, "read client ack"); - return nullptr; - } - uint32_t flags = butil::NetToHost32(flags_be); - bool client_ack_ok = (flags & HELLO_ACK_URMA_OK) != 0; - if (client_ack_ok) { - if (tp->_urma_state.load(butil::memory_order_acquire) == - UrmaTransport::URMA_OFF) { - // Protocol breakdown: client wants URMA but we already fell back. - ep->FailHandshake(tp, EPROTO, "client ack mismatch"); - return nullptr; - } - tp->_urma_state = UrmaTransport::URMA_ON; - ep->_state = ESTABLISHED; - if (ep->StartCqEvents() < 0) { - ep->FailHandshake(tp, errno, "start server CQ events"); - return nullptr; - } - } else { - ep->FallbackToTcp(tp, true); - } - return nullptr; -} - -// ============================================================================ -// UrmaConnect: drives the client handshake bthread. -// ============================================================================ - -void UrmaConnect::StartConnect(const Socket* socket, - void (*done)(int, void*), void* data) { - SocketUniquePtr s; - if (Socket::Address(socket->id(), &s) != 0) { - return; - } - _done = done; - _data = data; - _error = 0; - auto* tp = static_cast(socket->_transport.get()); - if (!tp) { - Run(); - return; - } - if (!tp->_urma_ep || !IsUrmaAvailable()) { - // Fall back to TCP immediately. - if (tp->_urma_ep) { - tp->_urma_ep->_state = UrmaEndpoint::FALLBACK_TCP; - } - tp->_urma_state = UrmaTransport::URMA_OFF; - Run(); - return; - } - bthread_t tid; - bthread_attr_t attr = BTHREAD_ATTR_NORMAL; - bthread_attr_set_name(&attr, "UrmaClientHandshake"); - if (bthread_start_background(&tid, &attr, - UrmaEndpoint::ProcessHandshakeAtClient, - tp->_urma_ep) != 0) { - tp->_urma_ep->_state = UrmaEndpoint::FALLBACK_TCP; - tp->_urma_state = UrmaTransport::URMA_OFF; - Run(); - } else { - // ProcessHandshakeAtClient adopts this reference in its - // SocketUniquePtr constructor. - s.release(); - } -} - -void UrmaConnect::StopConnect(Socket*) {} - -void UrmaConnect::Run() { - if (_done) { - auto cb = _done; - _done = nullptr; - cb(_error, _data); - } -} - // ============================================================================ // Debug / polling-mode stubs. // ============================================================================ @@ -1716,16 +1196,6 @@ void UrmaConnect::Run() { std::string UrmaEndpoint::GetStateStr() const { switch (_state.load(butil::memory_order_acquire)) { case UNINIT: return "UNINIT"; - case C_ALLOC_RES: return "C_ALLOC_RES"; - case C_HELLO_SEND: return "C_HELLO_SEND"; - case C_HELLO_WAIT: return "C_HELLO_WAIT"; - case C_IMPORT_PEER: return "C_IMPORT_PEER"; - case C_ACK_SEND: return "C_ACK_SEND"; - case S_HELLO_WAIT: return "S_HELLO_WAIT"; - case S_ALLOC_RES: return "S_ALLOC_RES"; - case S_IMPORT_PEER: return "S_IMPORT_PEER"; - 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"; @@ -1739,7 +1209,6 @@ void UrmaEndpoint::DebugInfo(std::ostream& os, butil::StringPiece) const { << " remote_recv_block_size=" << _remote_recv_block_size << " sq_window=" << _sq_window_size.load(butil::memory_order_relaxed) << " remote_rq_window=" << _remote_rq_window_size.load(butil::memory_order_relaxed) - << " handshake_version=" << _handshake_version << " read_buf=" << _input_processor.read_buf().size(); } diff --git a/src/brpc/urma/urma_endpoint.h b/src/brpc/urma/urma_endpoint.h index 1deb641d75..e48806ea16 100644 --- a/src/brpc/urma/urma_endpoint.h +++ b/src/brpc/urma/urma_endpoint.h @@ -38,8 +38,6 @@ #include "urma_api.h" #include "urma_types.h" -#include "brpc/handshake/urma_handshake.h" -#include "brpc/urma/urma_handshake.pb.h" namespace brpc { @@ -47,36 +45,28 @@ class UrmaTransport; namespace urma { -class UrmaHandshake; class UrmaStreamTrace; -struct ParsedHello; +// Connection parameters consumed by URMA resource setup. +struct UrmaConnectionInfo { + uint32_t buffer_size = 0; // peer recv buffer size in bytes (per WR) + uint32_t recv_buffer_cnt = 0; // peer recv buffer count (RQ depth - 1) + uint32_t jetty_id = 0; // peer jetty id + uint8_t eid[16] = {0}; // peer EID (network order) + uint32_t uasid = 0; // peer uasid + uint8_t tp_type = 0; // urma_tp_type_t: URMA_RTP / URMA_CTP / URMA_UTP + + // Flattened peer buffer-pool segment. + uint8_t seg_eid[16] = {0}; // EID owning the peer segment + uint32_t seg_uasid = 0; // uasid owning the peer segment + uint64_t seg_va = 0; // segment virtual address + uint64_t seg_len = 0; // segment length in bytes + uint32_t seg_token_id = 0; // segment token id +}; DECLARE_bool(urma_use_polling); DECLARE_int32(urma_poller_num); DECLARE_bool(urma_disable_bthread); -// Legacy URMA connection handshake helper, retained for interoperability tests. -// AdapterTransport manages the normal client upgrade path. -class UrmaConnect : public AppConnect { - friend class UrmaEndpoint; -public: - void StartConnect(const Socket* socket, - void (*done)(int err, void* data), void* data) override; - void StopConnect(Socket*) override; - - struct RunGuard { - explicit RunGuard(UrmaConnect* rc) : this_rc(rc) {} - ~RunGuard() { if (this_rc) this_rc->Run(); } - UrmaConnect* this_rc; - }; - -private: - void Run(); - void (*_done)(int, void*){nullptr}; - void* _data{nullptr}; - int _error{0}; -}; - // POD holder for the URMA kernel objects backing one connection: // - the send jetty (JFS), shared-JFR jetty, JFR, JFC, and (event mode) JFCE // - the imported peer jetty and peer segment @@ -99,18 +89,10 @@ struct UrmaResource { }; // Per-Socket URMA resources, send/recv windows and registered buffers. -// AdapterTransport drives the common handshake; the legacy handshake state -// and helpers remain here for interoperability tests. Cache-line aligned. +// AdapterTransport drives the connection handshake. Cache-line aligned. class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { - friend class UrmaConnect; friend class Socket; friend class ::brpc::UrmaTransport; - friend class UrmaHandshakeClientV2; - friend class UrmaHandshakeServerV2; - friend class UrmaHandshakeClientV3; - friend class UrmaHandshakeServerV3; - friend int DrainBytes(UrmaEndpoint*, size_t); - friend int ReadBodyAndNegotiate(UrmaEndpoint*, ParsedHello*, bool*); public: explicit UrmaEndpoint(Socket* s); @@ -137,14 +119,6 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { void DebugInfo(std::ostream& os, butil::StringPiece connector = "\n") const; - // Legacy TCP handshake callback used by interoperability tests. - // Dispatches on _state (see the .cpp for the full state machine): - // UNINIT (server) -> start the server handshake bthread - // < ESTABLISHED -> wake the handshake bthread parked in ReadFromFd - // FALLBACK_TCP -> hand off to InputMessenger::OnNewMessages - // ESTABLISHED -> drain stray TCP bytes - static void OnNewDataFromTcp(Socket* m); - // CQ completion poller: the edge-triggered callback for the CQ socket // (whose user owns a CQ lifetime context). Drains urma_poll_jfc and feeds the // completions into HandleCompletion, then calls InputMessenger. Event mode @@ -160,55 +134,20 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { std::function release_fn); static void PollingModeRelease(bthread_tag_t tag); - // ---- Legacy handshake IO helpers ---- - // Read at most @len bytes from the TCP fd into @data; waits on _read_butex - // on EAGAIN. Returns 0 on success, -1 on IO error (errno set). - int ReadFromFd(void* data, size_t len); - // Push @len bytes back into the TCP input processor so it can re-parse - // them (used on fallback when the magic is not "URMA"). - void PushBackToReadBuf(const void* data, size_t len); - // Write at most @len bytes from @data to the TCP fd; waits on - // _epollout_butex on EAGAIN. Returns 0 on success, -1 on IO error. - int WriteToFd(void* data, size_t len); - // Write an IOBuf to the TCP fd (used by v3 protobuf handshake). - int WriteToFd(butil::IOBuf& data); - - // ---- Hello builders/parsers (used by urma_handshake.cpp) ---- - // Fill the v2 binary HelloMessage with this endpoint's local params. - void FillLocalHelloV2(v2_wire::HelloMessage* out) const; - // Fill the v3 protobuf UrmaHello with this endpoint's local params. - void FillLocalHelloV3(UrmaHello* out) const; - // Write "URM3" + 4B big-endian pb_size + protobuf bytes. - int WriteHelloV3(const UrmaHello& msg); - // Read pb_size (4B) + protobuf bytes; parse; validate. Sets *negotiated - // to false (returns 0) if the message is invalid. - int ReadAndParseHelloV3(ParsedHello* out, bool* negotiated); + // Read local connection parameters independently of the wire version. + void GetLocalConnectionInfo(UrmaConnectionInfo* out) const; // Apply the peer's negotiated parameters (window sizes, peer jetty/seg). - void ApplyRemoteHello(const ParsedHello& remote); + void ApplyRemoteInfo(const UrmaConnectionInfo& remote); private: enum State { UNINIT = 0x0, - C_ALLOC_RES = 0x1, // client: allocate Jetty/JFC/JFR - C_HELLO_SEND = 0x2, - C_HELLO_WAIT = 0x3, - C_IMPORT_PEER = 0x4, // urma_import_seg then urma_import_jetty - C_ACK_SEND = 0x5, - S_HELLO_WAIT = 0x11, - S_ALLOC_RES = 0x12, - S_IMPORT_PEER = 0x13, - S_HELLO_SEND = 0x14, - S_ACK_WAIT = 0x15, ESTABLISHED = 0x100, FALLBACK_TCP = 0x200, FAILED = 0x300 }; - // Legacy client/server handshake bthreads. - static void* ProcessHandshakeAtClient(void* arg); - static void* ProcessHandshakeAtServer(void* arg); - // Allocate / deallocate the per-connection URMA resources (JFCE/JFC/JFR/ // Jetty, plus the CQ socket). Returns 0 on success. int AllocateResources(); @@ -221,7 +160,7 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { // does not establish the transport-path routing for the remote EID and // the first SEND is rejected with URMA_CR_RNR_RETRY_CNT_EXC_ERR. // Returns 0 on success. - int ImportPeer(const ParsedHello& peer); + int ImportPeer(const UrmaConnectionInfo& peer); // Post @num recv WRs into the local JFR. If @zerocopy, use a fresh pool // buffer; otherwise reuse the fixed _rbuf slot. Returns 0 on success. @@ -264,22 +203,13 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { void PollerAddCqSid(); void PollerRemoveCqSid(); - inline void TryReadOnTcp(); - void FallbackToTcp(UrmaTransport* transport, bool process_tcp); - void FailHandshake(UrmaTransport* transport, int error, - const char* reason); - - // Construct a ParsedHello from local state (used by SendLocalHello). - void MakeLocalParsedHello(ParsedHello* out) const; - std::string GetStateStr() const; // Not owned. Socket* _socket; - // Handshake state. + // Data-path availability. butil::atomic _state{UNINIT}; - int _handshake_version{0}; // 0 = unnegotiated; 2 = v2; 3 = v3 butil::atomic _pending_received_bytes{0}; butil::Mutex _dispatch_mutex; @@ -346,9 +276,6 @@ class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { // RQ consumer index. uint16_t _rq_received{0}; - // Butex for waking the handshake bthread parked in ReadFromFd. - butil::atomic* _read_butex{nullptr}; - // Reserved WR slots for pure-ack IMM WRs. static constexpr uint16_t RESERVED_WR_NUM = 3; diff --git a/src/brpc/urma_transport.cpp b/src/brpc/urma_transport.cpp index acd8d58cec..773989dc97 100644 --- a/src/brpc/urma_transport.cpp +++ b/src/brpc/urma_transport.cpp @@ -196,9 +196,9 @@ int UrmaTransport::PrepareUpgradeResources(bool server_side) { } int UrmaTransport::NegotiateUpgradeResources( - const urma::ParsedHello& remote, bool server_side) { + const urma::UrmaConnectionInfo& remote, bool server_side) { CHECK(_urma_ep != nullptr); - _urma_ep->ApplyRemoteHello(remote); + _urma_ep->ApplyRemoteInfo(remote); if (_urma_ep->ImportPeer(remote) < 0) { return -1; } diff --git a/src/brpc/urma_transport.h b/src/brpc/urma_transport.h index 1c22a6ae24..651351d66c 100644 --- a/src/brpc/urma_transport.h +++ b/src/brpc/urma_transport.h @@ -33,11 +33,6 @@ class UrmaTransport : public Transport { friend class TransportFactory; friend class AdapterTransport; friend class urma::UrmaEndpoint; -friend class urma::UrmaConnect; -friend class urma::UrmaHandshakeServerV2; -friend class urma::UrmaHandshakeClientV2; -friend class urma::UrmaHandshakeServerV3; -friend class urma::UrmaHandshakeClientV3; public: void Init(Socket* socket, const SocketOptions& options) override; @@ -62,7 +57,7 @@ friend class urma::UrmaHandshakeClientV3; bool UpgradeReady() const { return _urma_ep != nullptr; } int PrepareUpgradeResources(bool server_side); - int NegotiateUpgradeResources(const urma::ParsedHello& remote, + int NegotiateUpgradeResources(const urma::UrmaConnectionInfo& remote, bool server_side); void ActivateUpgrade(); int StartUpgradeEvents(); diff --git a/test/BUILD.bazel b/test/BUILD.bazel index a2c291eb7e..6e9bbd8b0f 100644 --- a/test/BUILD.bazel +++ b/test/BUILD.bazel @@ -151,6 +151,13 @@ cc_library( ], ) +cc_library( + name = "urma_legacy_handshake", + hdrs = ["brpc_urma_legacy_handshake.h"], + copts = COPTS, + deps = ["//:brpc"], +) + # tcmalloc and AddressSanitizer both intercept malloc/free, so linking # both produces undefined behaviour at runtime (typically a hard hang # at process start, or random allocator errors). When `brpc_with_asan` @@ -272,6 +279,7 @@ generate_unittests( ), deps = [ ":gperftools_helper", + ":urma_legacy_handshake", "//:brpc", ":cc_test_proto", "@com_google_googletest//:gtest", diff --git a/test/brpc_urma_legacy_handshake.h b/test/brpc_urma_legacy_handshake.h new file mode 100644 index 0000000000..c53e5f4550 --- /dev/null +++ b/test/brpc_urma_legacy_handshake.h @@ -0,0 +1,843 @@ +// 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_TEST_URMA_LEGACY_HANDSHAKE_H +#define BRPC_TEST_URMA_LEGACY_HANDSHAKE_H + +#include "brpc/urma_transport.h" +#include "brpc/handshake/urma_handshake.h" +#include "brpc/input_messenger.h" +#include "brpc/urma/urma_handshake.pb.h" +#include "brpc/urma/urma_helper.h" +#include "bthread/butex.h" +#include "bthread/bthread.h" +#include "butil/sys_byteorder.h" +#include "butil/time.h" + +// Independent legacy v2/v3 peer used by the interoperability tests. +namespace brpc { +namespace urma { + +DECLARE_bool(urma_recv_zerocopy); +DECLARE_int32(urma_client_handshake_version); +class LegacyUrmaEndpoint; + +class UrmaConnect : public AppConnect { + friend class LegacyUrmaEndpoint; +public: + void StartConnect(const Socket* socket, + void (*done)(int err, void* data), void* data) override; + void StopConnect(Socket*) override; + + struct RunGuard { + explicit RunGuard(UrmaConnect* rc) : this_rc(rc) {} + ~RunGuard() { if (this_rc) this_rc->Run(); } + UrmaConnect* this_rc; + }; + +private: + void Run(); + void (*_done)(int, void*){nullptr}; + void* _data{nullptr}; + int _error{0}; +}; + +// Legacy TCP handshake strategy used by interoperability tests. +// v2 speaks binary; v3 speaks protobuf ("URM3"). +class UrmaHandshake { +public: + virtual ~UrmaHandshake() = default; + virtual int ProtocolVersion() const = 0; + + // Send our local hello over the TCP fd held by the endpoint. + // Returns 0 on success, -1 on failure (errno set). + virtual int SendLocalHello() = 0; + + // Read the peer's hello and parse it into @out. Sets @negotiated to true + // if the peer is URMA-capable and the message validates; false (with + // return 0) to signal a graceful fall-back to TCP. + // Returns -1 on IO error (errno set). + virtual int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) = 0; +}; + +class LegacyUrmaEndpoint : public UrmaEndpoint { +public: + explicit LegacyUrmaEndpoint(Socket* socket) : UrmaEndpoint(socket) { + _read_butex = bthread::butex_create_checked>(); + _read_butex->store(0, butil::memory_order_relaxed); + } + ~LegacyUrmaEndpoint() override { bthread::butex_destroy(_read_butex); } + + int _handshake_version{0}; + + static constexpr State C_ALLOC_RES = static_cast(0x1); + static constexpr State C_HELLO_SEND = static_cast(0x2); + static constexpr State C_HELLO_WAIT = static_cast(0x3); + static constexpr State C_IMPORT_PEER = static_cast(0x4); + static constexpr State C_ACK_SEND = static_cast(0x5); + static constexpr State S_HELLO_WAIT = static_cast(0x11); + static constexpr State S_ALLOC_RES = static_cast(0x12); + static constexpr State S_IMPORT_PEER = static_cast(0x13); + static constexpr State S_HELLO_SEND = static_cast(0x14); + static constexpr State S_ACK_WAIT = static_cast(0x15); + static void OnNewDataFromTcp(Socket* socket); + static void* ProcessHandshakeAtClient(void* arg); + static void* ProcessHandshakeAtServer(void* arg); + int ReadFromFd(void* data, size_t len); + void PushBackToReadBuf(const void* data, size_t len); + int WriteToFd(void* data, size_t len); + int WriteToFd(butil::IOBuf& data); + void FillLocalHelloV2(v2_wire::HelloMessage* out) const; + void FillLocalHelloV3(UrmaHello* out) const; + int WriteHelloV3(const UrmaHello& message); + int ReadAndParseHelloV3(ParsedHello* out, bool* negotiated); + void TryReadOnTcp(); + void FallbackToTcp(UrmaTransport* transport, bool process_tcp); + void FailHandshake(UrmaTransport* transport, int error, const char* reason); +private: + butil::atomic* _read_butex; +}; + +// v2 binary handshake (magic "URMA"). +class UrmaHandshakeClientV2 : public UrmaHandshake { +public: + explicit UrmaHandshakeClientV2(LegacyUrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 2; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + LegacyUrmaEndpoint* _ep; +}; + +class UrmaHandshakeServerV2 : public UrmaHandshake { +public: + explicit UrmaHandshakeServerV2(LegacyUrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 2; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + LegacyUrmaEndpoint* _ep; +}; + +// v3 protobuf handshake (magic "URM3"). +class UrmaHandshakeClientV3 : public UrmaHandshake { +public: + explicit UrmaHandshakeClientV3(LegacyUrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 3; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + LegacyUrmaEndpoint* _ep; +}; + +class UrmaHandshakeServerV3 : public UrmaHandshake { +public: + explicit UrmaHandshakeServerV3(LegacyUrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 3; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + LegacyUrmaEndpoint* _ep; +}; + +// Legacy client factory: picks v2/v3 via --urma_client_handshake_version. +UrmaHandshake* CreateClientHandshake(LegacyUrmaEndpoint* ep); + +// Legacy server factory: dispatches on the first 4 bytes read from the TCP fd. +// Returns nullptr if the magic is not "URMA"/"URM3" (caller falls back to TCP). +// @magic is the first MAGIC_STR_LEN bytes already read from the fd. +UrmaHandshake* CreateServerHandshakeByMagic(LegacyUrmaEndpoint* ep, + const uint8_t magic[v2_wire::MAGIC_STR_LEN]); + +// Drain @n bytes from the TCP fd (used when msg_len > HELLO_MSG_LEN_MIN). +int DrainBytes(LegacyUrmaEndpoint* ep, size_t n); + +// Read the body following the magic and translate to ParsedHello. +// Returns -1 on IO error; 0 with *negotiated=false on invalid +// (peer not URMA-capable); 0 with *negotiated=true on success. +int ReadBodyAndNegotiate(LegacyUrmaEndpoint* ep, ParsedHello* out, bool* negotiated); + +static const int WAIT_TIMEOUT_MS = 50; +static const size_t HELLO_ACK_LEN = 4; +static const uint32_t HELLO_ACK_URMA_OK = 0x1; + +int LegacyUrmaEndpoint::ReadFromFd(void* data, size_t len) { + char* p = static_cast(data); + 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); + const int fd = _socket->fd(); + const ssize_t nr = read(fd, p + received, len - received); + if (nr < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) { + int rc = bthread::butex_wait(_read_butex, expected_val, &duetime); + if (rc < 0 && errno != EWOULDBLOCK && errno != ETIMEDOUT) { + return -1; + } + continue; + } + return -1; + } + if (nr == 0) { + errno = EEOF; + return -1; + } + received += nr; + } + return 0; +} + +void LegacyUrmaEndpoint::PushBackToReadBuf(const void* data, size_t len) { + _socket->fd_input_processor().read_buf().append(data, len); +} + +int LegacyUrmaEndpoint::WriteToFd(void* data, size_t len) { + char* p = static_cast(data); + size_t written = 0; + while (written < len) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = write(fd, p + written, len - written); + if (nw >= 0) { + written += nw; + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +void LegacyUrmaEndpoint::FillLocalHelloV2(v2_wire::HelloMessage* out) const { + *out = v2_wire::HelloMessage{}; // value-initialize + out->msg_len = v2_wire::HELLO_PACKET_LEN; + out->hello_ver = v2_wire::HELLO_V2_VERSION; + out->impl_ver = v2_wire::IMPL_V2_VERSION; + ParsedHello p; + GetLocalConnectionInfo(&p); + out->buffer_size = p.buffer_size; + out->recv_buffer_cnt = p.recv_buffer_cnt; + out->jetty_id = p.jetty_id; + std::memcpy(out->eid, p.eid, 16); + out->uasid = p.uasid; + out->tp_type = p.tp_type; + std::memcpy(out->seg_eid, p.seg_eid, 16); + out->seg_uasid = p.seg_uasid; + out->seg_va = p.seg_va; + out->seg_len = p.seg_len; + out->seg_token_id = p.seg_token_id; +} + +void LegacyUrmaEndpoint::FillLocalHelloV3(UrmaHello* out) const { + ParsedHello p; + GetLocalConnectionInfo(&p); + out->set_buffer_size(p.buffer_size); + out->set_recv_buffer_cnt(p.recv_buffer_cnt); + out->set_jetty_id(p.jetty_id); + out->set_eid(p.eid, 16); + out->set_uasid(p.uasid); + out->set_tp_type(p.tp_type); + out->set_seg_eid(p.seg_eid, 16); + out->set_seg_uasid(p.seg_uasid); + out->set_seg_va(p.seg_va); + out->set_seg_len(p.seg_len); + out->set_seg_token_id(p.seg_token_id); +} + +int LegacyUrmaEndpoint::WriteHelloV3(const UrmaHello& msg) { + butil::IOBuf packet; + packet.append("URM3", 4); + std::string body; + if (!msg.SerializeToString(&body)) { + LOG(ERROR) << "Fail to serialize UrmaHello"; + return -1; + } + uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); + packet.append(&pb_size_be, sizeof(pb_size_be)); + packet.append(body); + return WriteToFd(packet); +} + +int LegacyUrmaEndpoint::WriteToFd(butil::IOBuf& data) { + // Write out the IOBuf in a single WriteToFd-style loop. + while (!data.empty()) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = data.cut_into_file_descriptor(fd); + if (nw >= 0) { + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +int LegacyUrmaEndpoint::ReadAndParseHelloV3(ParsedHello* out, bool* negotiated) { + *negotiated = false; + uint32_t pb_size_be = 0; + if (ReadFromFd(&pb_size_be, sizeof(pb_size_be)) < 0) { + return -1; + } + const uint32_t pb_size = butil::NetToHost32(pb_size_be); + if (pb_size == 0 || pb_size > 4096) { + return 0; + } + std::string body(pb_size, '\0'); + if (ReadFromFd(&body[0], pb_size) < 0) { + return -1; + } + UrmaHello msg; + if (!msg.ParseFromArray(body.data(), static_cast(body.size()))) { + return 0; + } + if (msg.eid().size() != 16 || msg.seg_eid().size() != 16) { + return 0; + } + out->buffer_size = msg.buffer_size(); + out->recv_buffer_cnt = msg.recv_buffer_cnt(); + out->jetty_id = msg.jetty_id(); + std::memcpy(out->eid, msg.eid().data(), 16); + out->uasid = msg.uasid(); + out->tp_type = static_cast(msg.tp_type()); + std::memcpy(out->seg_eid, msg.seg_eid().data(), 16); + out->seg_uasid = msg.seg_uasid(); + out->seg_va = msg.seg_va(); + out->seg_len = msg.seg_len(); + out->seg_token_id = msg.seg_token_id(); + if (!ValidHello(*out)) { + return 0; + } + *negotiated = true; + return 0; +} + +int ReadBodyAndNegotiate(LegacyUrmaEndpoint* ep, ParsedHello* out, bool* negotiated) { + *negotiated = false; + uint8_t body[v2_wire::HELLO_BODY_LEN]; + if (ep->ReadFromFd(body, v2_wire::HELLO_BODY_LEN) < 0) { + return -1; + } + v2_wire::HelloMessage m; + m.Deserialize(body); + if (m.msg_len < v2_wire::HELLO_MSG_LEN_MIN || + m.msg_len > v2_wire::HELLO_MSG_LEN_MAX || + m.hello_ver != v2_wire::HELLO_V2_VERSION || + m.impl_ver != v2_wire::IMPL_V2_VERSION) { return 0; } + ParsedHello p; + p.buffer_size = m.buffer_size; + p.recv_buffer_cnt = m.recv_buffer_cnt; + p.jetty_id = m.jetty_id; + std::memcpy(p.eid, m.eid, 16); + p.uasid = m.uasid; + p.tp_type = m.tp_type; + std::memcpy(p.seg_eid, m.seg_eid, 16); + p.seg_uasid = m.seg_uasid; + p.seg_va = m.seg_va; + p.seg_len = m.seg_len; + p.seg_token_id = m.seg_token_id; + if (!ValidHello(p)) { + return 0; + } + // Drain trailing bytes if msg_len advertises more than the fixed body. + if (m.msg_len > v2_wire::HELLO_PACKET_LEN) { + if (DrainBytes(ep, m.msg_len - v2_wire::HELLO_PACKET_LEN) < 0) { + return -1; + } + } + *out = p; + *negotiated = true; + return 0; +} + +int DrainBytes(LegacyUrmaEndpoint* ep, size_t n) { + char buf[4096]; + while (n > 0) { + size_t want = std::min(n, sizeof(buf)); + if (ep->ReadFromFd(buf, want) < 0) { + return -1; + } + n -= want; + } + return 0; +} + +// ============================================================================ +// v2 client / server. +// ============================================================================ + +int UrmaHandshakeClientV2::SendLocalHello() { + v2_wire::HelloMessage m; + _ep->FillLocalHelloV2(&m); + uint8_t packet[v2_wire::HELLO_PACKET_LEN]; + std::memcpy(packet, "URMA", 4); + m.Serialize(packet + 4); + return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); +} + +int UrmaHandshakeClientV2::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + *negotiated = false; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + return -1; + } + if (std::memcmp(magic, "URMA", 4) != 0) { + // Peer is not URMA-capable; push the magic back so the TCP input + // messenger can re-parse it. + _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + return 0; + } + return ReadBodyAndNegotiate(_ep, out, negotiated); +} + +int UrmaHandshakeServerV2::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + return ReadBodyAndNegotiate(_ep, out, negotiated); +} + +int UrmaHandshakeServerV2::SendLocalHello() { + v2_wire::HelloMessage m; + _ep->FillLocalHelloV2(&m); + auto* tp = static_cast(_ep->_socket->_transport.get()); + if (tp->_urma_state.load(butil::memory_order_acquire) == + UrmaTransport::URMA_OFF) { + // Tell the client we are not URMA-capable: zero the version fields so + // the client's version check fails and it falls back to TCP. + m.hello_ver = 0; + m.impl_ver = 0; + m.jetty_id = 0; + m.buffer_size = 0; + } + uint8_t packet[v2_wire::HELLO_PACKET_LEN]; + std::memcpy(packet, "URMA", 4); + m.Serialize(packet + 4); + return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); +} + +// ============================================================================ +// v3 protobuf ("URM3"). The protobuf (de)serialization lives on the endpoint +// because it touches UrmaEndpoint private state; the classes here just call +// FillLocalHelloV3 / WriteHelloV3 / ReadAndParseHelloV3. +// ============================================================================ + +int UrmaHandshakeClientV3::SendLocalHello() { + UrmaHello msg; + _ep->FillLocalHelloV3(&msg); + return _ep->WriteHelloV3(msg); +} + +int UrmaHandshakeClientV3::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + *negotiated = false; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + return -1; + } + if (std::memcmp(magic, "URM3", 4) != 0) { + _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + return 0; + } + return _ep->ReadAndParseHelloV3(out, negotiated); +} + +int UrmaHandshakeServerV3::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + return _ep->ReadAndParseHelloV3(out, negotiated); +} + +int UrmaHandshakeServerV3::SendLocalHello() { + UrmaHello msg; + _ep->FillLocalHelloV3(&msg); + // v3 has no zero-out path; the client rejects jetty_id==0 via ValidHello. + return _ep->WriteHelloV3(msg); +} + +// ============================================================================ +// Factories. +// ============================================================================ + +UrmaHandshake* CreateClientHandshake(LegacyUrmaEndpoint* ep) { + switch (FLAGS_urma_client_handshake_version) { + case 3: return new UrmaHandshakeClientV3(ep); + case 2: + default: return new UrmaHandshakeClientV2(ep); + } +} + +UrmaHandshake* CreateServerHandshakeByMagic(LegacyUrmaEndpoint* ep, + const uint8_t magic[v2_wire::MAGIC_STR_LEN]) { + if (std::memcmp(magic, "URMA", 4) == 0) { + return new UrmaHandshakeServerV2(ep); + } + if (std::memcmp(magic, "URM3", 4) == 0) { + return new UrmaHandshakeServerV3(ep); + } + return nullptr; +} + +static void TryReadOnTcpDuringUrmaEst(Socket* socket); + +void LegacyUrmaEndpoint::OnNewDataFromTcp(Socket* m) { + auto* tp = static_cast(m->_transport.get()); + if (!tp) { + return; + } + // TCP fallback sockets may have no URMA endpoint. + auto* ep = static_cast(tp->_urma_ep); + if (!ep) { + // No URMA endpoint: pure TCP path. + InputMessenger::OnNewMessages(m); + return; + } + int progress = 0; + while (true) { + const State state = + ep->_state.load(butil::memory_order_acquire); + if (state == UNINIT) { + if (!m->CreatedByConnect()) { + // Server side: kick off the handshake bthread. + if (!IsUrmaAvailable()) { + ep->_state = FALLBACK_TCP; + tp->_urma_state = UrmaTransport::URMA_OFF; + InputMessenger::OnNewMessages(m); + return; + } + SocketUniquePtr s; + m->ReAddress(&s); + ep->_state = S_HELLO_WAIT; + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + bthread_attr_set_name(&attr, "UrmaServerHandshake"); + if (bthread_start_background(&tid, &attr, + ProcessHandshakeAtServer, ep) != 0) { + ep->_state = UNINIT; + LOG(FATAL) << "Fail to start UrmaServerHandshake bthread"; + } else { + s.release(); + } + return; + } + // Client side: handled by ProcessHandshakeAtClient. + return; + } else if (state < ESTABLISHED) { + // During handshake: wake the handshake bthread parked in ReadFromFd. + ep->_read_butex->fetch_add(1, butil::memory_order_release); + bthread::butex_wake(ep->_read_butex); + return; + } else if (state == FALLBACK_TCP) { + InputMessenger::OnNewMessages(m); + return; + } else if (state == ESTABLISHED) { + TryReadOnTcpDuringUrmaEst(m); + return; + } + if (!m->MoreReadEvents(&progress)) { + break; + } + } +} + +inline void LegacyUrmaEndpoint::TryReadOnTcp() { + if (_state.load(butil::memory_order_acquire) == FALLBACK_TCP) { + InputMessenger::OnNewMessages(_socket); + } +} + +static void TryReadOnTcpDuringUrmaEst(Socket* socket) { + int progress = Socket::PROGRESS_INIT; + while (true) { + uint8_t byte = 0; + const ssize_t nr = read(socket->fd(), &byte, 1); + if (nr < 0) { + if (errno != EAGAIN) { + const int saved_errno = errno; + socket->SetFailed(saved_errno, "Fail to read URMA TCP fd: %s", + berror(saved_errno)); + return; + } + if (!socket->MoreReadEvents(&progress)) { + return; + } + } else if (nr == 0) { + socket->SetEOF(); + return; + } else { + socket->SetFailed( + EPROTO, "Unexpected TCP data after URMA was established"); + return; + } + } +} + +void LegacyUrmaEndpoint::FallbackToTcp(UrmaTransport* transport, bool process_tcp) { + transport->_urma_state = UrmaTransport::URMA_OFF; + _state.store(FALLBACK_TCP, butil::memory_order_release); + Reset(); + _state.store(FALLBACK_TCP, butil::memory_order_release); + if (process_tcp) { + TryReadOnTcp(); + } +} + +void LegacyUrmaEndpoint::FailHandshake(UrmaTransport* transport, int error, + const char* reason) { + LOG(ERROR) << "URMA handshake failed in state=" << GetStateStr() + << " on " << _socket->description() + << ": " << reason << ", error=" << error + << " (" << berror(error) << ')'; + transport->_urma_state = UrmaTransport::URMA_OFF; + _state.store(FAILED, butil::memory_order_release); + DeallocateResources(); + auto* connect = + static_cast(_socket->_app_connect.get()); + if (connect) { + connect->_error = error; + } + _socket->SetFailed(error, "URMA handshake failed: %s", reason); +} + +// ============================================================================ +// Handshake state machines (client / server). Run in a background bthread. +// ============================================================================ + +void* LegacyUrmaEndpoint::ProcessHandshakeAtClient(void* arg) { + auto* ep = static_cast(arg); + SocketUniquePtr s(ep->_socket); + auto* tp = static_cast(s->_transport.get()); + UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); + if (!IsUrmaAvailable()) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_state = C_ALLOC_RES; + if (ep->AllocateResources() < 0) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + // Prepost the shared JFR before sending the client hello so the peer sees + // a ready receive queue as soon as its import completes. + if (ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_state = C_HELLO_SEND; + std::unique_ptr hs(CreateClientHandshake(ep)); + ep->_handshake_version = hs->ProtocolVersion(); + if (hs->SendLocalHello() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send client hello"); + return nullptr; + } + ep->_state = C_HELLO_WAIT; + ParsedHello remote; + bool negotiated = false; + if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read server hello"); + return nullptr; + } + if (!negotiated) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->ApplyRemoteInfo(remote); + ep->_state = C_IMPORT_PEER; + if (ep->ImportPeer(remote) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "import server resources"); + return nullptr; + } + ep->_state = C_ACK_SEND; + uint32_t flags = HELLO_ACK_URMA_OK; + uint32_t flags_be = butil::HostToNet32(flags); + if (ep->WriteToFd(&flags_be, HELLO_ACK_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send client ack"); + return nullptr; + } + tp->_urma_state = UrmaTransport::URMA_ON; + ep->_state = ESTABLISHED; + if (ep->StartCqEvents() < 0) { + ep->FailHandshake(tp, errno, "start client CQ events"); + return nullptr; + } + return nullptr; +} + +void* LegacyUrmaEndpoint::ProcessHandshakeAtServer(void* arg) { + auto* ep = static_cast(arg); + SocketUniquePtr s(ep->_socket); + auto* tp = static_cast(s->_transport.get()); + UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); + ep->_state = S_HELLO_WAIT; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client magic"); + return nullptr; + } + std::unique_ptr hs(CreateServerHandshakeByMagic(ep, magic)); + if (!hs) { + // Not an URMA peer: push the magic back and fall back to TCP. + ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_handshake_version = hs->ProtocolVersion(); + ParsedHello remote; + bool negotiated = false; + if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client hello"); + return nullptr; + } + if (!negotiated) { + ep->FailHandshake(tp, EPROTO, "invalid client hello"); + return nullptr; + } + ep->_state = S_ALLOC_RES; + if (ep->AllocateResources() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "allocate server resources"); + return nullptr; + } + const bool bonding = IsUrmaBondingDevice(); + if (!bonding && + ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "post server receives"); + return nullptr; + } + ep->ApplyRemoteInfo(remote); + ep->_state = S_IMPORT_PEER; + if (ep->ImportPeer(remote) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "import client resources"); + return nullptr; + } + if (bonding && + ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "post server receives"); + return nullptr; + } + ep->_state = S_HELLO_SEND; + if (hs->SendLocalHello() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send server hello"); + return nullptr; + } + ep->_state = S_ACK_WAIT; + uint32_t flags_be = 0; + if (ep->ReadFromFd(&flags_be, HELLO_ACK_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client ack"); + return nullptr; + } + uint32_t flags = butil::NetToHost32(flags_be); + bool client_ack_ok = (flags & HELLO_ACK_URMA_OK) != 0; + if (client_ack_ok) { + if (tp->_urma_state.load(butil::memory_order_acquire) == + UrmaTransport::URMA_OFF) { + // Protocol breakdown: client wants URMA but we already fell back. + ep->FailHandshake(tp, EPROTO, "client ack mismatch"); + return nullptr; + } + tp->_urma_state = UrmaTransport::URMA_ON; + ep->_state = ESTABLISHED; + if (ep->StartCqEvents() < 0) { + ep->FailHandshake(tp, errno, "start server CQ events"); + return nullptr; + } + } else { + ep->FallbackToTcp(tp, true); + } + return nullptr; +} + +// ============================================================================ +// UrmaConnect: drives the client handshake bthread. +// ============================================================================ + +void UrmaConnect::StartConnect(const Socket* socket, + void (*done)(int, void*), void* data) { + SocketUniquePtr s; + if (Socket::Address(socket->id(), &s) != 0) { + return; + } + _done = done; + _data = data; + _error = 0; + auto* tp = static_cast(socket->_transport.get()); + if (!tp) { + Run(); + return; + } + if (!tp->_urma_ep || !IsUrmaAvailable()) { + // Fall back to TCP immediately. + if (tp->_urma_ep) { + tp->_urma_ep->_state = LegacyUrmaEndpoint::FALLBACK_TCP; + } + tp->_urma_state = UrmaTransport::URMA_OFF; + Run(); + return; + } + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + bthread_attr_set_name(&attr, "UrmaClientHandshake"); + if (bthread_start_background(&tid, &attr, + LegacyUrmaEndpoint::ProcessHandshakeAtClient, + tp->_urma_ep) != 0) { + tp->_urma_ep->_state = LegacyUrmaEndpoint::FALLBACK_TCP; + tp->_urma_state = UrmaTransport::URMA_OFF; + Run(); + } else { + // ProcessHandshakeAtClient adopts this reference in its + // SocketUniquePtr constructor. + s.release(); + } +} + +void UrmaConnect::StopConnect(Socket*) {} + +void UrmaConnect::Run() { + if (_done) { + auto cb = _done; + _done = nullptr; + cb(_error, _data); + } +} + + +} // namespace urma +} // namespace brpc + +#endif // BRPC_TEST_URMA_LEGACY_HANDSHAKE_H diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp index 48c32a63a5..f37fd32b38 100644 --- a/test/brpc_urma_unittest.cpp +++ b/test/brpc_urma_unittest.cpp @@ -54,6 +54,7 @@ #include "brpc/policy/transport_handshake_protocol.h" #include "brpc/urma/urma_stream_trace.h" #include "brpc/global.h" +#include "brpc_urma_legacy_handshake.h" using namespace brpc; @@ -1402,7 +1403,9 @@ static void InstallLegacyUrmaTransport(Socket* socket, SocketUser* user) { options.need_on_edge_trigger = true; auto transport = std::unique_ptr(new UrmaTransport); transport->Init(socket, options); - transport->_on_edge_trigger = urma::UrmaEndpoint::OnNewDataFromTcp; + delete transport->_urma_ep; + transport->_urma_ep = new urma::LegacyUrmaEndpoint(socket); + transport->_on_edge_trigger = urma::LegacyUrmaEndpoint::OnNewDataFromTcp; transport->_default_connect = std::make_shared(); socket->_app_connect = transport->_default_connect; socket->_transport = std::move(transport); @@ -1560,8 +1563,8 @@ class UrmaCqCloseTest { return 13; } urma::ParsedHello hello; - endpoint->MakeLocalParsedHello(&hello); - endpoint->ApplyRemoteHello(hello); + endpoint->GetLocalConnectionInfo(&hello); + endpoint->ApplyRemoteInfo(hello); if (endpoint->ImportPeer(hello) != 0) { return 14; } @@ -1842,7 +1845,8 @@ static int RunUrmaHandshakeInterop(int version, bool polling, bool zerocopy, if (peer == UrmaInteropPeer::LEGACY_CLIENT) { auto* transport = static_cast(client->_transport.get()); if (transport->_urma_state != UrmaTransport::URMA_ON || - transport->GetUrmaEp()->_handshake_version != version) { + static_cast(transport->GetUrmaEp()) + ->_handshake_version != version) { return 39; } } else if (AdapterTransport::Get(client.get())->handshake_phase() != @@ -1926,7 +1930,8 @@ static int RunUrmaHandshakeInterop(int version, bool polling, bool zerocopy, return 49; } if (peer == UrmaInteropPeer::LEGACY_SERVER && - accepted_ep->_handshake_version != version) { + static_cast(accepted_ep) + ->_handshake_version != version) { return 50; } if (peer != UrmaInteropPeer::LEGACY_SERVER && From c0ef0ea8de6410901698e7be5b35d1124fb80363 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 00:19:39 -0700 Subject: [PATCH 12/17] Synchronize handshake read cancellation test and clarify protocol naming --- docs/cn/handshake_common_design.md | 3 +- test/brpc_transport_handshake_unittest.cpp | 32 ++++++++++++++++++---- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md index 53e5412936..c2534b37ae 100644 --- a/docs/cn/handshake_common_design.md +++ b/docs/cn/handshake_common_design.md @@ -40,7 +40,8 @@ protocol 和 transport participant，`HandshakeSession` 只负责调用参与者 ## 协议兼容 -- RDMA 保持既有 v2/v3 wire format、协议 ID 和注册名称。 +- RDMA 保持既有 v2/v3 wire format 和协议数值 ID `1`。公共握手的枚举和 + 注册名称分别调整为 `PROTOCOL_TRANSPORT_HANDSHAKE` 和 `transport_handshake`。 - UBSHM 使用 v3 hello。64 字节 hello 后，双方交换 4 字节网络序格式扩展； 本地 IPC 后端选择 `IPC_V2`，UBS 后端选择 `LEGACY_64`；只有双方选择相同的 受支持格式才能升级，否则回退 TCP。资源映射传递对端共享内存长度及协商格式， diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index e815414c53..9d818e4875 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -84,8 +84,11 @@ class MemoryHandshakeIO : public HandshakeIO { struct BlockingHandshakeRead { SocketHandshakeIO* io; - int result; - int error; + bthread::CountdownEvent completed; + int result = 0; + int error = 0; + + explicit BlockingHandshakeRead(SocketHandshakeIO* value) : io(value) {} }; static void* RunBlockingHandshakeRead(void* arg) { @@ -94,6 +97,7 @@ static void* RunBlockingHandshakeRead(void* arg) { char byte = 0; read->result = read->io->ReadExact(&byte, sizeof(byte)); read->error = errno; + read->completed.signal(); return nullptr; } @@ -547,17 +551,35 @@ TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { ASSERT_EQ(0, Socket::Address(id, &socket)); SocketHandshakeIO io(socket.get()); - BlockingHandshakeRead read = {&io, 0, 0}; + BlockingHandshakeRead read(&io); 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( + bool was_waiting = false; + const int64_t wait_deadline = butil::gettimeofday_us() + 1000000; + while (butil::gettimeofday_us() < wait_deadline) { + // A nonzero wake count proves ReadExact reached its read butex. + if (bthread::butex_wake(io._read_butex, true) > 0) { + was_waiting = true; + break; + } + bthread_usleep(1000); + } + EXPECT_EQ(0, socket->SetFailed( EFAILEDSOCKET, "cancel blocked handshake read")); + bthread_flush(); + const int completed = read.completed.timed_wait( + butil::milliseconds_from_now(1000)); + if (completed != 0) { + shutdown(socket->fd(), SHUT_RDWR); + io.NotifyReadable(); + } ASSERT_EQ(0, bthread_join(tid, nullptr)); + EXPECT_TRUE(was_waiting); + EXPECT_EQ(0, completed); EXPECT_EQ(-1, read.result); EXPECT_EQ(EFAILEDSOCKET, read.error); } From 82cfd7657f509b73c06b43e06885126c2811a0ff Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 00:24:45 -0700 Subject: [PATCH 13/17] Condense the common handshake design document --- docs/cn/handshake_common_design.md | 61 ++++++++++++------------------ 1 file changed, 24 insertions(+), 37 deletions(-) diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md index c2534b37ae..930806e8fa 100644 --- a/docs/cn/handshake_common_design.md +++ b/docs/cn/handshake_common_design.md @@ -1,55 +1,42 @@ # RDMA 与 UBSHM 公共握手 -## 范围 - -本次重构将 RDMA 和 UBSHM 的握手流程收敛到公共层。`Socket` 持有 -`AdapterTransport`；后者以 TCP 作为握手控制通道，在升级成功后选择高速 -数据面，协商不可用时继续使用 TCP。URMA 目前仍使用独立的 `UrmaTransport` -和握手实现，不在本次迁移范围内。 +RDMA 和 UBSHM 共用握手流程，以 TCP 交换控制帧。升级成功后使用高速 +数据面，协商不可用时继续使用 TCP。 ## 分工 | 组件 | 职责 | | --- | --- | -| `AdapterTransport` | 持有 TCP 与候选高速 Transport，启动客户端握手、分发服务端输入，并选择数据面 | -| `HandshakeSession` | 执行公共客户端和服务端状态机，管理帧收发、增量解析和终态发布 | -| `HandshakeProtocol` | 定义 hello、ACK 和扩展帧，编解码协议字段并保存协议状态 | -| `HandshakeTransport` | 准备和协商资源，处理升级成功、TCP 回退及失败清理 | -| RDMA / UBSHM protocol | 实现各自的 wire format、版本兼容和字段校验 | -| `RdmaEndpoint` / `UBShmEndpoint` | 管理各自的资源和数据面，不驱动完整握手流程 | +| `AdapterTransport` | 持有 TCP 与高速 Transport，启动握手、分发输入并选择数据面 | +| `HandshakeSession` | 编排客户端和服务端流程，管理帧收发、增量解析和终态发布 | +| `HandshakeProtocol` | 编解码 hello、ACK 和扩展帧，校验协议版本及字段 | +| `HandshakeTransport` | 准备和协商资源，处理升级、回退及失败清理 | +| `RdmaEndpoint` / `UBShmEndpoint` | 管理传输资源和数据面 | 客户端在 TCP 连接建立后启动握手 bthread；服务端通过 `InputMessenger` -增量解析 hello、扩展字段（若协议需要）和 ACK。`AdapterTransport` 选择具体的 -protocol 和 transport participant，`HandshakeSession` 只负责调用参与者并编排流程。 -服务端在 `EXTENSION_WAIT` 和 `ACK_WAIT` 保留解析上下文，并通过协议版本 -恢复选中的 protocol，因为扩展帧和 ACK 本身不带 magic。 +增量解析控制帧。在 `EXTENSION_WAIT` 和 `ACK_WAIT`，服务端保留解析上下文， +按已选协议版本继续处理后续帧。 ## 状态与回退 -握手终态为 `ESTABLISHED`、`FALLBACK_TCP` 或 `FAILED`。升级成功后，TCP -仍是控制连接，业务数据走 RDMA/UBSHM；控制连接上出现额外业务数据视为 -协议错误。升级完成时停止当前 TCP 业务解析循环，由控制连接处理额外输入。 -资源不可用或协商拒绝时，释放本次升级资源，设置 TCP 为数据面， -再以 release 顺序发布 `FALLBACK_TCP`。事件线程以 acquire 顺序读取终态， -继续解析已缓存及后续的 TCP 业务数据。已确认属于握手协议的畸形帧则失败， -不作为普通 RPC 数据回放。 +- `ESTABLISHED`：业务数据走 RDMA/UBSHM，TCP 保留为控制连接。服务端 + 完成最终 ACK 解析并退出 TCP 解析循环后启动高速接收；额外 TCP 数据触发连接错误。 +- `FALLBACK_TCP`：释放升级资源后以 release 顺序发布终态，事件线程以 + acquire 顺序读取并继续处理已缓存及后续的 TCP 业务数据。 +- `FAILED`：握手 I/O 错误或畸形帧导致失败，清理资源并关闭连接。 -增量解析中，不完整帧返回 `NOT_ENOUGH_DATA` 且不消费输入；不可能匹配 -握手 magic 的前缀返回 `TRY_OTHERS`。TCP 分片、握手帧与后续数据粘连都由 -帧边界处理，不能假设一次 read 恰好得到一帧。 +完整帧按边界消费；分片输入返回 `NOT_ENOUGH_DATA` 等待补齐， +与握手 magic 不匹配的前缀返回 `TRY_OTHERS`。 ## 协议兼容 -- RDMA 保持既有 v2/v3 wire format 和协议数值 ID `1`。公共握手的枚举和 - 注册名称分别调整为 `PROTOCOL_TRANSPORT_HANDSHAKE` 和 `transport_handshake`。 -- UBSHM 使用 v3 hello。64 字节 hello 后，双方交换 4 字节网络序格式扩展； - 本地 IPC 后端选择 `IPC_V2`，UBS 后端选择 `LEGACY_64`；只有双方选择相同的 - 受支持格式才能升级，否则回退 TCP。资源映射传递对端共享内存长度及协商格式， - 激活数据面前发布格式。ACK 仍为 4 字节。 -- 公共 framing 仅负责帧边界，协议字段由各自 `HandshakeProtocol` 解释。 +- RDMA 保持 v2/v3 wire format 和协议数值 ID `1`。公共握手使用枚举 + `PROTOCOL_TRANSPORT_HANDSHAKE`，注册名称为 `transport_handshake`。 +- UBSHM 使用 64 字节 v3 hello、4 字节网络序格式扩展和 4 字节 ACK。 + IPC 后端选择 `IPC_V2`，UBS 后端选择 `LEGACY_64`；格式相同则升级，否则 + 回退 TCP。资源映射使用对端共享内存长度，数据面激活前发布协商格式。 -## 验证重点 +## 测试 -测试覆盖分片和粘连输入、magic 前缀分流、协议版本与格式校验、TCP 回退、 -资源释放及终态发布，并分别运行 RDMA 与 UBSHM 回归测试。URMA 的公共握手 -迁移需要单独设计和测试，不应仅凭本次重构推断其已完成。 +覆盖分片与粘连输入、magic 分流、版本与格式校验、TCP 回退、资源释放、 +终态发布、最终 ACK 与首包并发，以及 RDMA 和 UBSHM 回归测试。 From f3012a3612afb9a2248c0fbefa1e9f942843070e Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 01:12:05 -0700 Subject: [PATCH 14/17] Forward application connection stop notifications through the adapter --- src/brpc/adapter_transport.cpp | 6 +- test/brpc_transport_handshake_unittest.cpp | 76 ++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index fbde15df3b..acdd2fd42c 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -84,7 +84,11 @@ class AdapterConnect : public AppConnect { } } - void StopConnect(Socket*) override {} + void StopConnect(Socket* socket) override { + if (_app_connect) { + _app_connect->StopConnect(socket); + } + } private: struct ApplicationConnectTask { diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 9d818e4875..9e6cd304be 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -22,6 +22,8 @@ #include #include #include +#include +#include #include #include "bthread/bthread.h" @@ -129,6 +131,29 @@ class TestAppConnect : public AppConnect { void StopConnect(Socket*) override {} }; +class RecordingAppConnect : public TestAppConnect { +public: + void StopConnect(Socket* socket) override { + { + std::lock_guard lock(_mutex); + ++_stopped[socket->id()]; + } + stopped.signal(); + } + + int StopCount(SocketId id) { + std::lock_guard lock(_mutex); + auto it = _stopped.find(id); + return it == _stopped.end() ? 0 : it->second; + } + + bthread::CountdownEvent stopped; + +private: + std::mutex _mutex; + std::map _stopped; +}; + static int RejectHandshakeStart(bthread_t*, const bthread_attr_t*, void* (*)(void*), void*) { errno = EINTR; @@ -1073,6 +1098,57 @@ TEST(TransportHandshakeTest, urma_socket_skips_adapter_handshake_parser) { } #endif +#if BRPC_WITH_RDMA || BRPC_WITH_UBRING +TEST(TransportHandshakeTest, shared_app_connect_receives_each_socket_stop) { + const SocketMode modes[] = { +#if BRPC_WITH_RDMA + SOCKET_MODE_RDMA, +#endif +#if BRPC_WITH_UBRING + SOCKET_MODE_UBRING, +#endif + }; + for (SocketMode mode : modes) { + const auto original = std::make_shared(); + SocketOptions options; + options.user = + static_cast(get_or_new_client_side_messenger()); + options.socket_mode = mode; + options.app_connect = original; + SocketId first_id; + ASSERT_EQ(0, Socket::Create(options, &first_id)); + SocketUniquePtr first; + ASSERT_EQ(0, Socket::Address(first_id, &first)); + ASSERT_NE(nullptr, first->_app_connect); + ASSERT_NE(original.get(), first->_app_connect.get()); + + options.app_connect = first->_app_connect; + SocketId second_id; + ASSERT_EQ(0, Socket::Create(options, &second_id)); + SocketUniquePtr second; + ASSERT_EQ(0, Socket::Address(second_id, &second)); + ASSERT_EQ(first->_app_connect.get(), second->_app_connect.get()); + + AdapterTransport::Get(first.get())->FallbackToTcp(); + ASSERT_EQ(0, first->SetFailed()); + first.reset(); + ASSERT_EQ(0, original->stopped.timed_wait( + butil::milliseconds_from_now(1000))); + EXPECT_EQ(1, original->StopCount(first_id)); + EXPECT_EQ(0, original->StopCount(second_id)); + + original->stopped.reset(1); + AdapterTransport::Get(second.get())->FallbackToTcp(); + ASSERT_EQ(0, second->SetFailed()); + second.reset(); + ASSERT_EQ(0, original->stopped.timed_wait( + butil::milliseconds_from_now(1000))); + EXPECT_EQ(1, original->StopCount(first_id)); + EXPECT_EQ(1, original->StopCount(second_id)); + } +} +#endif + #if BRPC_WITH_RDMA TEST(TransportHandshakeTest, inherited_adapter_connect_is_not_wrapped_twice) { From 0296dd46fb39aa7756c260ffac601d3499a67c67 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 01:13:59 -0700 Subject: [PATCH 15/17] Initialize handshake wait state and stop writes on socket failure --- src/brpc/handshake/handshake_io.cpp | 29 +++-- test/brpc_transport_handshake_unittest.cpp | 120 +++++++++++++++++++-- 2 files changed, 132 insertions(+), 17 deletions(-) diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp index 67cc95e37c..67840e4de8 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; @@ -46,6 +46,7 @@ static const int WAIT_TIMEOUT_MS = 50; SocketHandshakeIO::SocketHandshakeIO(Socket* socket) : _socket(socket) , _read_butex(bthread::butex_create_checked >()) { + _read_butex->store(0, butil::memory_order_relaxed); } SocketHandshakeIO::~SocketHandshakeIO() { @@ -97,8 +98,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) { @@ -107,10 +108,15 @@ int SocketHandshakeIO::ReadExact(void* data, size_t len) { } template -static int WriteAllLoop(size_t len, WriteOnce write_once, +static int WriteAllLoop(SocketId socket_id, size_t len, WriteOnce write_once, WaitWritable wait_writable) { size_t written = 0; while (written < len) { + SocketUniquePtr alive; + if (Socket::Address(socket_id, &alive) != 0) { + errno = EFAILEDSOCKET; + return -1; + } const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); const ssize_t nw = write_once(written, len - written); if (nw > 0) { @@ -129,6 +135,9 @@ static int WriteAllLoop(size_t len, WriteOnce write_once, } if (wait_writable(&duetime) < 0 && errno != ETIMEDOUT) { + if (alive->Failed()) { + errno = EFAILEDSOCKET; + } return -1; } } @@ -136,10 +145,10 @@ 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, + return WriteAllLoop(_socket->id(), len, [data, fd](size_t offset, size_t remaining) { return write(fd, static_cast(data) + offset, remaining); @@ -151,7 +160,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/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 8c5ad480e6..55562450f6 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -18,13 +18,18 @@ #include #include +#include #include #include #include #include #include "bthread/bthread.h" +#include "bthread/butex.h" +#include "bthread/countdown_event.h" +#include "bthread/unstable.h" #include "butil/fd_guard.h" +#include "butil/time.h" #include "butil/sys_byteorder.h" #include "brpc/adapter_transport.h" #include "brpc/handshake/handshake_io.h" @@ -77,8 +82,11 @@ class MemoryHandshakeIO : public HandshakeIO { struct BlockingHandshakeRead { SocketHandshakeIO* io; - int result; - int error; + bthread::CountdownEvent completed; + int result = 0; + int error = 0; + + explicit BlockingHandshakeRead(SocketHandshakeIO* value) : io(value) {} }; static void* RunBlockingHandshakeRead(void* arg) { @@ -87,7 +95,26 @@ static void* RunBlockingHandshakeRead(void* arg) { char byte = 0; read->result = read->io->ReadExact(&byte, sizeof(byte)); read->error = errno; - return NULL; + read->completed.signal(); + return nullptr; +} + +struct BlockingHandshakeWrite { + SocketHandshakeIO* io; + bthread::CountdownEvent completed; + int result = 0; + int error = 0; + + explicit BlockingHandshakeWrite(SocketHandshakeIO* value) : io(value) {} +}; + +static void* RunBlockingHandshakeWrite(void* arg) { + auto* write = static_cast(arg); + const char byte = 0; + write->result = write->io->WriteAll(&byte, sizeof(byte)); + write->error = errno; + write->completed.signal(); + return nullptr; } class TestAppConnect : public AppConnect { @@ -372,6 +399,67 @@ static std::string MakeUBShmHello() { return frame; } +TEST(TransportHandshakeTest, socket_io_initializes_read_butex) { + auto* reused = bthread::butex_create_checked>(); + reused->store(37, butil::memory_order_relaxed); + bthread::butex_destroy(reused); + SocketHandshakeIO io; + EXPECT_EQ(0, io._read_butex->load(butil::memory_order_relaxed)); + io.NotifyReadable(); + EXPECT_EQ(1, io._read_butex->load(butil::memory_order_relaxed)); +} + +TEST(TransportHandshakeTest, blocked_write_stops_after_socket_failure) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + butil::fd_guard peer_fd(fds[1]); + const int flags = fcntl(fds[0], F_GETFL, 0); + ASSERT_GE(flags, 0); + ASSERT_EQ(0, fcntl(fds[0], F_SETFL, flags | O_NONBLOCK)); + char data[4096] = {}; + ssize_t n; + do { + n = send(fds[0], data, sizeof(data), MSG_NOSIGNAL); + } while (n > 0 || (n < 0 && errno == EINTR)); + ASSERT_EQ(-1, n); + ASSERT_EQ(EAGAIN, errno); + + SocketOptions options; + options.fd = fds[0]; + options.need_on_edge_trigger = true; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + SocketHandshakeIO io(socket.get()); + BlockingHandshakeWrite write(&io); + bthread_t tid; + ASSERT_EQ(0, bthread_start_background( + &tid, nullptr, RunBlockingHandshakeWrite, &write)); + bool was_waiting = false; + const int64_t wait_deadline = butil::gettimeofday_us() + 1000000; + while (butil::gettimeofday_us() < wait_deadline) { + // A nonzero wake count proves WriteAll reached WaitEpollOut's butex. + if (bthread::butex_wake(socket->_epollout_butex, true) > 0) { + was_waiting = true; + break; + } + bthread_usleep(1000); + } + socket->SetFailed(EFAILEDSOCKET, "cancel blocked handshake write"); + bthread_flush(); + const int completed = write.completed.timed_wait( + butil::milliseconds_from_now(1000)); + if (completed != 0) { + shutdown(socket->fd(), SHUT_RDWR); + } + ASSERT_EQ(0, bthread_join(tid, nullptr)); + EXPECT_TRUE(was_waiting); + EXPECT_EQ(0, completed); + EXPECT_EQ(-1, write.result); + EXPECT_EQ(EFAILEDSOCKET, write.error); +} + TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { int fds[2]; ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); @@ -386,17 +474,35 @@ TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { ASSERT_EQ(0, Socket::Address(id, &socket)); SocketHandshakeIO io(socket.get()); - BlockingHandshakeRead read = {&io, 0, 0}; + BlockingHandshakeRead read(&io); 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( + bool was_waiting = false; + const int64_t wait_deadline = butil::gettimeofday_us() + 1000000; + while (butil::gettimeofday_us() < wait_deadline) { + // A nonzero wake count proves ReadExact reached its read butex. + if (bthread::butex_wake(io._read_butex, true) > 0) { + was_waiting = true; + break; + } + bthread_usleep(1000); + } + EXPECT_EQ(0, socket->SetFailed( EFAILEDSOCKET, "cancel blocked handshake read")); - ASSERT_EQ(0, bthread_join(tid, NULL)); + bthread_flush(); + const int completed = read.completed.timed_wait( + butil::milliseconds_from_now(1000)); + if (completed != 0) { + shutdown(socket->fd(), SHUT_RDWR); + io.NotifyReadable(); + } + ASSERT_EQ(0, bthread_join(tid, nullptr)); + EXPECT_TRUE(was_waiting); + EXPECT_EQ(0, completed); EXPECT_EQ(-1, read.result); EXPECT_EQ(EFAILEDSOCKET, read.error); } From baf2acaab57f9958d7ce44828f64a1a9c8178fb4 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 01:17:13 -0700 Subject: [PATCH 16/17] Stop handshake reads before consuming data from failed sockets --- src/brpc/handshake/handshake_io.cpp | 10 +++---- test/brpc_transport_handshake_unittest.cpp | 33 ++++++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp index 67840e4de8..0cb6634b8f 100644 --- a/src/brpc/handshake/handshake_io.cpp +++ b/src/brpc/handshake/handshake_io.cpp @@ -68,6 +68,11 @@ static int ReadExactLoop(butil::atomic* read_butex, size_t len, ReadOnce read_once) { size_t received = 0; while (received < len) { + SocketUniquePtr alive; + if (Socket::Address(socket_id, &alive) != 0) { + errno = EFAILEDSOCKET; + return -1; + } 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); @@ -78,11 +83,6 @@ static int ReadExactLoop(butil::atomic* read_butex, 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; diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 9e6cd304be..673acd2072 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -562,6 +562,39 @@ TEST(TransportHandshakeTest, established_socket_does_not_restart_handshake) { } } +TEST(TransportHandshakeTest, failed_socket_read_preserves_buffered_data) { + for (size_t prefix_len : {0u, 1u}) { + SCOPED_TRACE(prefix_len); + 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)); + ASSERT_EQ(2, write(peer_fd, "HS", 2)); + + SocketHandshakeIO io(socket.get()); + char data[2] = {}; + if (prefix_len != 0) { + ASSERT_EQ(0, io.ReadExact(data, prefix_len)); + EXPECT_EQ('H', data[0]); + } + ASSERT_EQ(0, socket->SetFailed(EFAILEDSOCKET, + "cancel buffered handshake read")); + errno = 0; + EXPECT_EQ(-1, io.ReadExact(data, 2 - prefix_len)); + EXPECT_EQ(EFAILEDSOCKET, errno); + + const ssize_t remaining = recv(socket->fd(), data, sizeof(data), + MSG_DONTWAIT); + ASSERT_EQ(static_cast(2 - prefix_len), remaining); + EXPECT_EQ(0, memcmp(data, "HS" + prefix_len, remaining)); + } +} + TEST(TransportHandshakeTest, blocked_read_stops_after_socket_failure) { int fds[2]; ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); From 2c380959dfd3f110b3c16e6e63de1b86c5164fe5 Mon Sep 17 00:00:00 2001 From: Spicy-cream <2667314914@qq.com> Date: Mon, 26 Oct 2026 18:51:49 -0700 Subject: [PATCH 17/17] Restore handshake pushback bytes before buffered input --- src/brpc/handshake/handshake_io.cpp | 8 +++- src/brpc/handshake/handshake_io.h | 1 + test/brpc_transport_handshake_unittest.cpp | 44 ++++++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/brpc/handshake/handshake_io.cpp b/src/brpc/handshake/handshake_io.cpp index 0cb6634b8f..7bc3e83ff5 100644 --- a/src/brpc/handshake/handshake_io.cpp +++ b/src/brpc/handshake/handshake_io.cpp @@ -162,7 +162,13 @@ int SocketHandshakeIO::WriteAll(const void* data, size_t len) { int SocketHandshakeIO::PushBack(const void* data, size_t len) { CHECK(_socket != nullptr); if (len != 0) { - return _socket->fd_input_processor().read_buf().append(data, len) == 0 ? 0 : -1; + butil::IOBuf restored; + if (restored.append(data, len) != 0) { + return -1; + } + auto& input = _socket->fd_input_processor().read_buf(); + restored.append(input); + input.swap(restored); } return 0; } diff --git a/src/brpc/handshake/handshake_io.h b/src/brpc/handshake/handshake_io.h index 28340ab437..bad7f3ee46 100644 --- a/src/brpc/handshake/handshake_io.h +++ b/src/brpc/handshake/handshake_io.h @@ -38,6 +38,7 @@ class HandshakeIO { virtual int ReadExact(void* data, size_t len) = 0; virtual int WriteAll(const void* data, size_t len) = 0; + // Restore data before any buffered input. virtual int PushBack(const void* data, size_t len) = 0; }; diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 64987a03ed..d3b544f890 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -465,6 +465,50 @@ static std::string MakeUBShmHello() { return frame; } +TEST(TransportHandshakeTest, socket_io_pushback_prepends_buffered_input) { + for (const std::string& suffix : {std::string(), std::string("tail")}) { + SocketOptions options; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + SocketUniquePtr socket; + ASSERT_EQ(0, Socket::Address(id, &socket)); + auto& input = socket->fd_input_processor().read_buf(); + input.append(suffix); + SocketHandshakeIO io(socket.get()); + + EXPECT_EQ(0, io.PushBack("PRPC", 4)); + EXPECT_EQ("PRPC" + suffix, input.to_string()); + EXPECT_EQ(0, io.PushBack("\0", 1)); + const std::string expected = std::string(1, '\0') + "PRPC" + suffix; + EXPECT_EQ(expected, input.to_string()); + EXPECT_EQ(0, io.PushBack(nullptr, 0)); + EXPECT_EQ(expected, input.to_string()); + socket->SetFailed(); + } +} + +TEST(HandshakeFrameTest, + socket_magic_mismatch_restores_prefix_before_buffered_input) { + 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)); + socket->fd_input_processor().read_buf().append("tail"); + ASSERT_EQ(4, write(peer_fd, "PRPC", 4)); + + SocketHandshakeIO io(socket.get()); + std::string payload; + EXPECT_EQ(FRAME_NOT_MINE, FrameCodec::ReadFrame( + &io, FixedSpec("URMA", 4, 4), true, &payload)); + EXPECT_EQ("PRPCtail", socket->fd_input_processor().read_buf().to_string()); + socket->SetFailed(); +} + TEST(TransportHandshakeTest, socket_io_initializes_read_butex) { auto* reused = bthread::butex_create_checked>(); reused->store(37, butil::memory_order_relaxed);