diff --git a/BUILD.bazel b/BUILD.bazel index 5882a034da..3e33178ffb 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -545,7 +545,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/Makefile b/Makefile index 08df1f04e9..3d8938ded6 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_FLATBUFFERS),1) BRPC_DIRS += src/brpc/flatbuffers endif diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md new file mode 100644 index 0000000000..930806e8fa --- /dev/null +++ b/docs/cn/handshake_common_design.md @@ -0,0 +1,42 @@ +# RDMA 与 UBSHM 公共握手 + +RDMA 和 UBSHM 共用握手流程，以 TCP 交换控制帧。升级成功后使用高速 +数据面，协商不可用时继续使用 TCP。 + +## 分工 + +| 组件 | 职责 | +| --- | --- | +| `AdapterTransport` | 持有 TCP 与高速 Transport，启动握手、分发输入并选择数据面 | +| `HandshakeSession` | 编排客户端和服务端流程，管理帧收发、增量解析和终态发布 | +| `HandshakeProtocol` | 编解码 hello、ACK 和扩展帧，校验协议版本及字段 | +| `HandshakeTransport` | 准备和协商资源，处理升级、回退及失败清理 | +| `RdmaEndpoint` / `UBShmEndpoint` | 管理传输资源和数据面 | + +客户端在 TCP 连接建立后启动握手 bthread；服务端通过 `InputMessenger` +增量解析控制帧。在 `EXTENSION_WAIT` 和 `ACK_WAIT`，服务端保留解析上下文， +按已选协议版本继续处理后续帧。 + +## 状态与回退 + +- `ESTABLISHED`：业务数据走 RDMA/UBSHM，TCP 保留为控制连接。服务端 + 完成最终 ACK 解析并退出 TCP 解析循环后启动高速接收；额外 TCP 数据触发连接错误。 +- `FALLBACK_TCP`：释放升级资源后以 release 顺序发布终态，事件线程以 + acquire 顺序读取并继续处理已缓存及后续的 TCP 业务数据。 +- `FAILED`：握手 I/O 错误或畸形帧导致失败，清理资源并关闭连接。 + +完整帧按边界消费；分片输入返回 `NOT_ENOUGH_DATA` 等待补齐， +与握手 magic 不匹配的前缀返回 `TRY_OTHERS`。 + +## 协议兼容 + +- 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 回退、资源释放、 +终态发布、最终 ACK 与首包并发，以及 RDMA 和 UBSHM 回归测试。 diff --git a/example/ubring_performance/client.cpp b/example/ubring_performance/client.cpp index 24a786f8f1..0001781fa5 100644 --- a/example/ubring_performance/client.cpp +++ b/example/ubring_performance/client.cpp @@ -17,6 +17,8 @@ #include #include +#include +#include #include #include #include @@ -55,23 +57,31 @@ butil::atomic g_total_bytes; butil::atomic g_total_cnt; std::vector g_servers; int rr_index = 0; -volatile bool g_stop = false; +std::atomic g_failed{false}; +std::atomic g_next_request_ns{0}; -butil::atomic g_token(10000); - -static void* GenerateToken(void* arg) { - int64_t start_time = butil::monotonic_time_ns(); - int64_t accumulative_token = g_token.load(butil::memory_order_relaxed); - while (!g_stop) { - bthread_usleep(100000); - int64_t now = butil::monotonic_time_ns(); - if (accumulative_token * 1000000000 / (now - start_time) < FLAGS_expected_qps) { - int64_t delta = FLAGS_expected_qps * (now - start_time) / 1000000000 - accumulative_token; - g_token.fetch_add(delta, butil::memory_order_relaxed); - accumulative_token += delta; +static bool WaitForRequest(int64_t deadline) { + int64_t next = butil::monotonic_time_ns(); + if (FLAGS_expected_qps > 0) { + const int64_t interval = 1000000000L / FLAGS_expected_qps; + int64_t previous = g_next_request_ns.load(); + do { + next = std::max(previous, butil::monotonic_time_ns()); + } while (!g_next_request_ns.compare_exchange_weak( + previous, next + interval)); + } + while (!g_failed.load()) { + const int64_t now = butil::monotonic_time_ns(); + if (now >= deadline) { + return false; } + if (now >= next) { + return true; + } + bthread_usleep(std::min( + (std::min(next, deadline) - now) / 1000 + 1, 10000)); } - return nullptr; + return false; } class PerformanceTest { @@ -79,9 +89,6 @@ class PerformanceTest { PerformanceTest(int attachment_size, bool echo_attachment) : _addr(nullptr) , _channel(nullptr) - , _start_time(0) - , _iterations(0) - , _stop(false) { if (attachment_size > 0) { _addr = malloc(attachment_size); @@ -98,8 +105,6 @@ class PerformanceTest { delete _channel; } - inline bool IsStop() { return _stop; } - int Init() { brpc::ChannelOptions options; options.socket_mode = FLAGS_use_ubring? brpc::SOCKET_MODE_UBRING : brpc::SOCKET_MODE_TCP; @@ -116,7 +121,7 @@ class PerformanceTest { return -1; } - // Add retry mechanism for RPC call + // Warm up the connection before measuring requests. int retry = 3; while (retry > 0) { brpc::Controller cntl; @@ -130,107 +135,79 @@ class PerformanceTest { } LOG(WARNING) << "RPC call failed, retrying... (" << retry << " left): " << cntl.ErrorText(); retry--; - bthread_usleep(1000000); // 100ms delay before retry + bthread_usleep(1000000); } LOG(ERROR) << "RPC call failed after multiple retries"; return -1; } - struct RespClosure { - brpc::Controller* cntl; - test::PerfTestResponse* resp; - PerformanceTest* test; - }; - - void SendRequest() { - if (FLAGS_expected_qps > 0) { - while (g_token.load(butil::memory_order_relaxed) <= 0) { - bthread_usleep(10); + static void* RunTest(void* arg) { + auto* self = static_cast(arg); + const int64_t deadline = butil::monotonic_time_ns() + + static_cast(FLAGS_test_seconds) * 1000000000L; + uint64_t sent = 0; + test::PerfTestService_Stub stub(self->_channel); + while (!g_failed.load() && + (FLAGS_test_iterations == 0 || + sent < static_cast(FLAGS_test_iterations))) { + std::vector cntls(FLAGS_queue_depth); + std::vector requests(FLAGS_queue_depth); + std::vector responses(FLAGS_queue_depth); + int pending = 0; + for (; pending < FLAGS_queue_depth; ++pending) { + if ((FLAGS_test_iterations > 0 && + sent >= static_cast(FLAGS_test_iterations)) || + !WaitForRequest(deadline)) { + break; + } + requests[pending].set_echo_attachment(self->_echo_attachment); + cntls[pending].request_attachment().append(self->_attachment); + stub.Test(&cntls[pending], &requests[pending], + &responses[pending], brpc::DoNothing()); + ++sent; } - g_token.fetch_sub(1, butil::memory_order_relaxed); - } - RespClosure* closure = new RespClosure; - test::PerfTestRequest request; - closure->resp = new test::PerfTestResponse(); - closure->cntl = new brpc::Controller(); - request.set_echo_attachment(_echo_attachment); - closure->cntl->request_attachment().append(_attachment); - closure->test = this; - google::protobuf::Closure* done = brpc::NewCallback(&HandleResponse, closure); - test::PerfTestService_Stub stub(_channel); - stub.Test(closure->cntl, &request, closure->resp, done); - } - - static void HandleResponse(RespClosure* closure) { - std::unique_ptr cntl_guard(closure->cntl); - std::unique_ptr response_guard(closure->resp); - if (closure->cntl->Failed()) { - LOG(ERROR) << "RPC call failed: " << closure->cntl->ErrorText(); - // RPCs in this example are expected to succeed. Silently ignoring - // failures would hide problems and invalidate the performance result. - closure->test->_stop = true; - return; - } - - g_latency_recorder << closure->cntl->latency_us(); - if (closure->resp->cpu_usage().size() > 0) { - g_server_cpu_recorder << atof(closure->resp->cpu_usage().c_str()) * 100; - } - g_total_bytes.fetch_add(closure->cntl->request_attachment().size(), butil::memory_order_relaxed); - g_total_cnt.fetch_add(1, butil::memory_order_relaxed); - - cntl_guard.reset(nullptr); - response_guard.reset(nullptr); - - if (closure->test->_iterations == 0 && FLAGS_test_iterations > 0) { - closure->test->_stop = true; - return; - } - --closure->test->_iterations; - uint64_t last = g_last_time.load(butil::memory_order_relaxed); - uint64_t now = butil::gettimeofday_us(); - if (now > last && now - last > 100000) { - if (g_last_time.exchange(now, butil::memory_order_relaxed) == last) { - g_client_cpu_recorder << - atof(bvar::Variable::describe_exposed("process_cpu_usage").c_str()) * 100; + for (int i = 0; i < pending; ++i) { + brpc::Join(cntls[i].call_id()); + if (cntls[i].Failed()) { + LOG(ERROR) << "RPC call failed: " << cntls[i].ErrorText(); + g_failed.store(true); + } else if (self->_echo_attachment && + cntls[i].response_attachment() != self->_attachment) { + LOG(ERROR) << "Echo attachment does not match request"; + g_failed.store(true); + } else { + g_latency_recorder << cntls[i].latency_us(); + g_total_bytes.fetch_add(self->_attachment.size(), + butil::memory_order_relaxed); + g_total_cnt.fetch_add(1, butil::memory_order_relaxed); + if (!responses[i].cpu_usage().empty()) { + g_server_cpu_recorder << + atof(responses[i].cpu_usage().c_str()) * 100; + } + } + } + if (pending == 0) { + break; + } + uint64_t last = g_last_time.load(butil::memory_order_relaxed); + const uint64_t now = butil::gettimeofday_us(); + if (now > last && now - last > 100000 && + g_last_time.exchange(now, butil::memory_order_relaxed) == last) { + g_client_cpu_recorder << atof(bvar::Variable::describe_exposed( + "process_cpu_usage").c_str()) * 100; } } - if (now - closure->test->_start_time > FLAGS_test_seconds * 1000000u) { - closure->test->_stop = true; - return; - } - closure->test->SendRequest(); - } - - static void* RunTest(void* arg) { - PerformanceTest* test = (PerformanceTest*)arg; - test->_start_time = butil::gettimeofday_us(); - test->_iterations = FLAGS_test_iterations; - - for (int i = 0; i < FLAGS_queue_depth; ++i) { - test->SendRequest(); - } - return nullptr; } private: void* _addr; brpc::Channel* _channel; - uint64_t _start_time; - uint32_t _iterations; - volatile bool _stop; butil::IOBuf _attachment; bool _echo_attachment; }; -static void* DeleteTest(void* arg) { - PerformanceTest* test = (PerformanceTest*)arg; - delete test; - return nullptr; -} - -void Test(int thread_num, int attachment_size) { +bool Test(int thread_num, int attachment_size) { std::cout << "[Threads: " << thread_num << ", Depth: " << FLAGS_queue_depth << ", Attachment: " << attachment_size << "B" @@ -248,19 +225,14 @@ void Test(int thread_num, int attachment_size) { tests.push_back(t); } uint64_t start_time = butil::gettimeofday_us(); - bthread_t tid[thread_num]; - if (FLAGS_expected_qps > 0) { - bthread_t tid; - bthread_start_background(&tid, &BTHREAD_ATTR_NORMAL, GenerateToken, nullptr); - } + g_next_request_ns.store(butil::monotonic_time_ns()); + std::vector tid(thread_num); for (int k = 0; k < thread_num; ++k) { - bthread_start_background(&tid[k], &BTHREAD_ATTR_NORMAL, - PerformanceTest::RunTest, tests[k]); + CHECK_EQ(0, bthread_start_background( + &tid[k], &BTHREAD_ATTR_NORMAL, PerformanceTest::RunTest, tests[k])); } for (int k = 0; k < thread_num; ++k) { - while (!tests[k]->IsStop()) { - bthread_usleep(10000); - } + bthread_join(tid[k], nullptr); } uint64_t end_time = butil::gettimeofday_us(); double throughput = g_total_bytes / 1.048576 / (end_time - start_time); @@ -270,25 +242,29 @@ void Test(int thread_num, int attachment_size) { << ", 99th-Latency: " << g_latency_recorder.latency_percentile(0.99) << ", 99.9th-Latency: " << g_latency_recorder.latency_percentile(0.999) << ", Throughput: " << throughput << "MB/s" - << ", QPS: " << (g_total_cnt.load(butil::memory_order_relaxed) * 1000 / (end_time - start_time)) << "k" + << ", QPS: " << (g_total_cnt.load(butil::memory_order_relaxed) * + 1000000.0 / (end_time - start_time)) << ", Server CPU-utilization: " << g_server_cpu_recorder.latency(10) << "%" << ", Client CPU-utilization: " << g_client_cpu_recorder.latency(10) << "%" << std::endl; } else { std::cout << " Throughput: " << throughput << "MB/s" << std::endl; } - g_stop = true; - for (int k = 0; k < thread_num; ++k) { - bthread_start_background(&tid[k], &BTHREAD_ATTR_NORMAL, DeleteTest, tests[k]); - } - for (int k = 0; k < thread_num; ++k) { - bthread_join(tid[k], nullptr); + for (auto* test : tests) { + delete test; } + return !g_failed.load(); } int main(int argc, char* argv[]) { GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true); + if (FLAGS_queue_depth <= 0 || FLAGS_test_seconds <= 0 || + FLAGS_expected_qps < 0 || FLAGS_test_iterations < 0) { + LOG(ERROR) << "queue_depth and test_seconds must be positive; " + "expected_qps and test_iterations must be nonnegative"; + return 1; + } brpc::StartDummyServerAt(FLAGS_dummy_port); std::string::size_type pos1 = 0; @@ -301,19 +277,27 @@ int main(int argc, char* argv[]) { g_servers.push_back(FLAGS_servers.substr(pos1)); if (FLAGS_thread_num > 0 && FLAGS_attachment_size >= 0) { - Test(FLAGS_thread_num, FLAGS_attachment_size); + if (!Test(FLAGS_thread_num, FLAGS_attachment_size)) { + return 1; + } } else if (FLAGS_thread_num <= 0 && FLAGS_attachment_size >= 0) { for (int i = 1; i <= FLAGS_max_thread_num; i *= 2) { - Test(i, FLAGS_attachment_size); + if (!Test(i, FLAGS_attachment_size)) { + return 1; + } } } else if (FLAGS_thread_num > 0 && FLAGS_attachment_size < 0) { for (int i = 1; i <= 1024; i *= 4) { - Test(FLAGS_thread_num, i); + if (!Test(FLAGS_thread_num, i)) { + return 1; + } } } else { for (int j = 1; j <= 1024; j *= 4) { for (int i = 1; i <= FLAGS_max_thread_num; i *= 2) { - Test(i, j); + if (!Test(i, j)) { + return 1; + } } } } @@ -325,7 +309,7 @@ int main(int argc, char* argv[]) { int main(int argc, char* argv[]) { LOG(ERROR) << " brpc is not compiled with ubring. To enable it, please refer to the ubring documentation"; - return 0; + return 1; } #endif diff --git a/src/brpc/CMakeLists.txt b/src/brpc/CMakeLists.txt index 552e3ce99f..af77478789 100644 --- a/src/brpc/CMakeLists.txt +++ b/src/brpc/CMakeLists.txt @@ -291,7 +291,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/acceptor.h b/src/brpc/acceptor.h index 01a2b26a9d..db443da17b 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -135,7 +135,7 @@ friend class Server; bool _force_ssl; std::shared_ptr _ssl_ctx; - // Choose to use a certain socket: 0 TCP, 1 RDMA + // Requested transport for accepted connections. SocketMode _socket_mode; // Acceptor belongs to this tag diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp new file mode 100644 index 0000000000..36fd043ad0 --- /dev/null +++ b/src/brpc/adapter_transport.cpp @@ -0,0 +1,718 @@ +// 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/shm/shm_mgr.h" +#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) { + const int error = errno != 0 ? errno : EAGAIN; + AdapterTransport::Get(const_cast(socket))->CompleteConnection( + handshake::FAILED); + delete task; + done(error, data); + } + } + + void StopConnect(Socket* socket) override { + if (_app_connect) { + _app_connect->StopConnect(socket); + } + } + +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 || BRPC_WITH_UBRING +template +void FailClientHandshake(T* transport, Socket* socket, + int* connect_error, const char* name) { + const int saved_errno = errno != 0 ? errno : EPROTO; + transport->DeactivateUpgrade(); + *connect_error = saved_errno; + socket->SetFailed(saved_errno, "Fail to complete %s handshake from %s: %s", + name, socket->description().c_str(), berror(saved_errno)); +} +#endif + +#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 { + FailClientHandshake(_transport, _socket, _connect_error, "rdma"); + } + +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(), _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 { + FailClientHandshake(_transport, _socket, _connect_error, "ubring"); + } + +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 != nullptr); + return static_cast(socket->_transport.get()); +} + +const AdapterTransport* AdapterTransport::Get(const Socket* socket) { + CHECK(socket != nullptr); + return static_cast(socket->_transport.get()); +} + +bool AdapterTransport::upgrade_capable(SocketMode mode) const { + if (_mode != mode || _high_speed_transport == nullptr) { + 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) { + 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"); + const int rc = start_background(&tid, &attr, ProcessClientHandshake, task); + if (rc != 0) { + delete task; + errno = rc; + return -1; + } + return 0; +} + +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 = + static_cast( + _socket->parsing_context()); + CHECK(context->adapter() != nullptr); + 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 nullptr; + } + + std::unique_ptr protocol = + transport->CreateClientHandshakeAdapter(); + CHECK(protocol != nullptr); + 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 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 nullptr; + } + + const size_t local_shm_len = + static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + ubring::SHM local_trx_shm = { + nullptr, 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 nullptr; + } +#endif + + socket->SetFailed(EPROTO, "Unsupported client transport handshake"); + adapter->CompleteConnection(handshake::FAILED); + task->done(EPROTO, task->data); + return nullptr; +} + +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 == nullptr) { + 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 != nullptr && + _socket->user() == client_messenger) { + // This socket keeps the unwrapped AppConnect across health-check + // resets, so its read path must not depend on another handshake. + _on_edge_trigger = _tcp_transport->GetOnEdgeTrigger(); + 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..686754c1f1 --- /dev/null +++ b/src/brpc/adapter_transport.h @@ -0,0 +1,107 @@ +// 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 a high-speed 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); + } + 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); + 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; + 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/channel.cpp b/src/brpc/channel.cpp index 7b5f0a0da6..6111030edf 100644 --- a/src/brpc/channel.cpp +++ b/src/brpc/channel.cpp @@ -35,7 +35,6 @@ #include "brpc/serialized_request.h" #include "brpc/serialized_response.h" #include "brpc/details/usercode_backup_pool.h" // TooManyUserCode -#include "brpc/rdma/rdma_helper.h" #include "brpc/policy/esp_authenticator.h" #include "brpc/transport_factory.h" #include "brpc/details/controller_private_accessor.h" @@ -78,7 +77,8 @@ ChannelSSLOptions* ChannelOptions::mutable_ssl_options() { } static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { - if (opt.auth == nullptr && + if (opt.socket_mode == SOCKET_MODE_TCP && + opt.auth == nullptr && !opt.has_ssl_options() && opt.client_host.empty() && opt.device_name.empty() && @@ -87,6 +87,7 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { // Returning zeroized result by default is more intuitive for users. return ChannelSignature(); } + const int mode = static_cast(opt.socket_mode); uint32_t seed = 0; std::string buf; buf.reserve(1024); @@ -94,6 +95,7 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { do { buf.clear(); butil::MurmurHash3_x64_128_Init(&mm_ctx, seed); + butil::MurmurHash3_x64_128_Update(&mm_ctx, &mode, sizeof(mode)); if (!opt.connection_group.empty()) { buf.append("|conng="); @@ -137,11 +139,6 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { } else { // All disabled ChannelSSLOptions are the same } - if (opt.socket_mode == SOCKET_MODE_RDMA) { - buf.append("|rdma"); - } else if (opt.socket_mode == SOCKET_MODE_URMA) { - buf.append("|urma"); - } butil::MurmurHash3_x64_128_Update(&mm_ctx, buf.data(), buf.size()); buf.clear(); diff --git a/src/brpc/global.cpp b/src/brpc/global.cpp index 4cce93ed0c..5703eeeadd 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,13 @@ 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) { + CONNECTION_TYPE_ALL, "transport_handshake" }; + if (RegisterProtocol(PROTOCOL_TRANSPORT_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..d22252b7d2 --- /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 HandshakeAdapter::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/rdma/rdma_handshake_server.h b/src/brpc/handshake/handshake_adapter.h similarity index 52% rename from src/brpc/rdma/rdma_handshake_server.h rename to src/brpc/handshake/handshake_adapter.h index 705aa893ef..8081480162 100644 --- a/src/brpc/rdma/rdma_handshake_server.h +++ b/src/brpc/handshake/handshake_adapter.h @@ -15,34 +15,42 @@ // specific language governing permissions and limitations // under the License. -#ifndef BRPC_RDMA_RDMA_HANDSHAKE_SERVER_H -#define BRPC_RDMA_RDMA_HANDSHAKE_SERVER_H +#ifndef BRPC_HANDSHAKE_HANDSHAKE_ADAPTER_H +#define BRPC_HANDSHAKE_HANDSHAKE_ADAPTER_H -#include "brpc/destroyable.h" +#include "butil/macros.h" #include "brpc/parse_result.h" +#include "brpc/transport_handshake.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; -}; +namespace handshake { + +// Common InputMessenger parser lifecycle. Protocol adapters provide the server +// step and connection session. +class HandshakeAdapter { +public: + virtual ~HandshakeAdapter() = default; + virtual ParseResult ExecuteServerHandshake( + butil::IOBuf* source, Socket* socket); -// 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); +protected: + HandshakeAdapter() = default; + virtual StepResult RunServerStep( + butil::IOBuf* source, Socket* socket) = 0; + virtual HandshakeSession* GetSession(Socket* socket) const = 0; + +private: + DISALLOW_COPY_AND_ASSIGN(HandshakeAdapter); +}; -} // namespace rdma +} // namespace handshake } // namespace brpc -#endif // BRPC_RDMA_RDMA_HANDSHAKE_SERVER_H +#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..f2a551928b --- /dev/null +++ b/src/brpc/handshake/handshake_frame.cpp @@ -0,0 +1,226 @@ +// 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 == nullptr || (spec.magic_len != 0 && spec.magic == nullptr)) { + 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 == nullptr || payload == nullptr || + (spec.magic_len != 0 && spec.magic == nullptr)) { + 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(butil::IOBuf* input, + const FrameSpec& spec, + std::string* payload, + bool* magic_matched) { + if (magic_matched != nullptr) { + *magic_matched = false; + } + 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); + if (input->size() < spec.magic_len) { + return FRAME_NEED_MORE; + } + std::string header(header_len, '\0'); + if (spec.magic_len != 0) { + input->copy_to(&header[0], spec.magic_len); + } + if (spec.magic_len != 0 && + memcmp(header.data(), spec.magic, spec.magic_len) != 0) { + return FRAME_NOT_MINE; + } + if (magic_matched != nullptr) { + *magic_matched = true; + } + if (input->size() < header_len) { + return FRAME_NEED_MORE; + } + if (header_len != 0) { + input->copy_to(&header[0], header_len); + } + 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'); + input->cutn(&frame[0], frame_len); + 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 == nullptr) { + 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; +} + +} // 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..651ab6a8ab --- /dev/null +++ b/src/brpc/handshake/handshake_frame.h @@ -0,0 +1,91 @@ +// 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(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, + 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(butil::IOBuf* input, + const FrameSpec& spec, + std::string* payload, + bool* magic_matched = nullptr); + static FrameResult WriteFrame(HandshakeIO* io, const FrameSpec& spec, + const std::string& payload); + +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..ede2a96b28 --- /dev/null +++ b/src/brpc/handshake/handshake_io.cpp @@ -0,0 +1,141 @@ +// 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 { + +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() { + 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); +} + +int SocketHandshakeIO::ReadExact(void* data, size_t len) { + CHECK(data != nullptr); + CHECK(_socket != nullptr); + const int fd = _socket->fd(); + size_t received = 0; + while (received < len) { + if (_socket->Failed()) { + 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( + fd, static_cast(data) + received, len - received); + if (nr < 0) { + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN) { + 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::WriteAll(const void* data, size_t len) { + CHECK(data != nullptr); + CHECK(_socket != nullptr); + const int fd = _socket->fd(); + size_t written = 0; + while (written < len) { + if (_socket->Failed()) { + errno = EFAILEDSOCKET; + return -1; + } + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const ssize_t nw = write( + fd, static_cast(data) + 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 (_socket->WaitEpollOut(_socket->fd(), true, &duetime) < 0 && + errno != ETIMEDOUT) { + if (_socket->Failed()) { + errno = EFAILEDSOCKET; + } + return -1; + } + } + return 0; +} + +int SocketHandshakeIO::PushBack(const void* data, size_t len) { + CHECK(_socket != nullptr); + if (len != 0) { + 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; +} + +} // 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..5a3fe1c0f4 --- /dev/null +++ b/src/brpc/handshake/handshake_io.h @@ -0,0 +1,68 @@ +// 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; + // Restore data before any buffered input. + virtual int PushBack(const void* data, size_t len) = 0; +}; + +// The caller retains the Socket for the duration of each I/O operation. +class SocketHandshakeIO : public HandshakeIO { +public: + explicit SocketHandshakeIO(Socket* socket = nullptr); + ~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..ed15f7e2a3 --- /dev/null +++ b/src/brpc/handshake/rdma_handshake.cpp @@ -0,0 +1,516 @@ +// 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" +#include "brpc/rdma/rdma_helper.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(); +} + +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 + +handshake::StepResult RdmaClientHandshakeAdapterV2::BuildHello( + 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 RdmaServerHandshakeAdapterV2::BuildHello( + 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); +} + +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; + } + if (msg.has_mtu() && + (msg.mtu() < IBV_MTU_256 || msg.mtu() > IBV_MTU_4096)) { + 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 (local.path_mtu.has_value()) { + msg->set_mtu(*local.path_mtu); + } + 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 (msg.has_mtu()) { + out->path_mtu = msg.mtu(); + } + 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 + +handshake::StepResult RdmaClientHandshakeAdapterV3::BuildHello( + bool enabled, std::string* payload) { + CHECK(enabled); + PrepareClientEce(); + ParsedHello local{}; + FillLocalHello(&local); + local.path_mtu = GetRdmaActiveMtu(); + RdmaHello msg; + v3_wire::FillLocalRdmaHello(local, &msg); + return v3_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaServerHandshakeAdapterV3::BuildHello( + 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 RdmaHandshakeAdapter::ParseHello( + const std::string& payload) { + return _version == 2 ? v2_wire::ParsePayload(payload, &_remote) + : 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 HandshakeAdapter { +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) { + RdmaFallbackProtocol v2(2); + RdmaFallbackProtocol v3(3); + std::vector protocols; + protocols.push_back(&v2); + protocols.push_back(&v3); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, source, &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(nullptr) {} + + 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 != nullptr); + 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() != nullptr); + + 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()); + } + RdmaServerHandshakeTransport participant(transport, socket, source); + return GetSession(socket)->RunServer( + protocol_ptrs, source, &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..040df9e98e --- /dev/null +++ b/src/brpc/handshake/rdma_handshake.h @@ -0,0 +1,144 @@ +// 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 ParseHello( + const std::string& payload) override; + const ParsedHello& remote() const { return _remote; } + + const handshake::FrameSpec& HelloFrameSpec() const override { + return RdmaHelloFrameSpec(_version); + } + +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) {} + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; +}; + +class RdmaServerHandshakeAdapterV2 : public RdmaHandshakeAdapter { +public: + explicit RdmaServerHandshakeAdapterV2(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 2) {} + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; +}; + +class RdmaClientHandshakeAdapterV3 : public RdmaHandshakeAdapter { +public: + explicit RdmaClientHandshakeAdapterV3(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 3) {} + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; +}; + +class RdmaServerHandshakeAdapterV3 : public RdmaHandshakeAdapter { +public: + explicit RdmaServerHandshakeAdapterV3(RdmaEndpoint* ep) + : RdmaHandshakeAdapter(ep, 3) {} + handshake::StepResult BuildHello( + bool enabled, std::string* payload) 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..f67a93ab9d 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( + nullptr, 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..5024a0c7a9 --- /dev/null +++ b/src/brpc/handshake/ubshm_handshake.cpp @@ -0,0 +1,484 @@ +// 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( + nullptr, 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( + nullptr, 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 != nullptr ? shm_name : ""; + _server_reply = false; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + const char* name = nullptr; + 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 == nullptr) { + 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) == nullptr) { + 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 HandshakeAdapter { +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) { + UBShmFallbackProtocol protocol; + std::vector protocols(1, &protocol); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, source, &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 = { + 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 = { + 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 != nullptr) { + *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() != nullptr); + + 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); + std::vector protocols(1, &wire); + UBShmServerHandshakeTransport participant( + transport, &wire, socket, source); + const StepResult result = GetSession(socket)->RunServer( + protocols, source, &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/input_messenger.cpp b/src/brpc/input_messenger.cpp index 0f141682bd..5caec2e1ed 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/input_messenger_processor.cpp b/src/brpc/input_messenger_processor.cpp index febd1b580c..2620b02aeb 100644 --- a/src/brpc/input_messenger_processor.cpp +++ b/src/brpc/input_messenger_processor.cpp @@ -277,15 +277,8 @@ int InputMessengerProcessor::ProcessNewMessage(ssize_t bytes, bool read_eof, num_bthread_created = 0; } } - // In RDMA polling mode, all messages must be executed in a new bthread and - // not in the bthread where the polling bthread is located, because the - // method for processing messages may call synchronization primitives, - // causing the polling bthread to be scheduled out. - if (_socket->_socket_mode == SOCKET_MODE_RDMA || - _socket->_socket_mode == SOCKET_MODE_UBRING || - _socket->_socket_mode == SOCKET_MODE_URMA) { - _socket->_transport->QueueMessage(last_msg, &num_bthread_created, true); - } + // The active transport decides whether to dispatch or retain the last message. + _socket->_transport->QueueMessage(last_msg, &num_bthread_created, true); if (num_bthread_created) { bthread_flush(); } 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/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 110279b701..7984653085 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -17,23 +17,22 @@ #if BRPC_WITH_RDMA -#include // std::min -#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 #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/scoped_lock.h" +#include "butil/logging.h" // CHECK, LOG +#include "butil/sys_byteorder.h" // HostToNet,NetToHost +#include DECLARE_int32(task_group_ntags); @@ -130,8 +129,6 @@ RdmaResource::~RdmaResource() { RdmaEndpoint::RdmaEndpoint(Socket* s) : _socket(s) - , _state(UNINIT) - , _handshake_version(0) , _resource(nullptr) , _send_cq_events(0) , _recv_cq_events(0) @@ -165,29 +162,26 @@ 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(); + _outgoing_mtu.reset(); _resource = nullptr; _send_cq_events = 0; _recv_cq_events = 0; _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; @@ -204,583 +198,42 @@ 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; +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->path_mtu = _outgoing_mtu; + local->ece.reset(); + if (_outgoing_ece.has_value()) { + local->ece = _outgoing_ece; } - - // 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(); - } - - 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 { @@ -797,6 +250,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 @@ -873,8 +327,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]; @@ -933,7 +386,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; @@ -977,7 +431,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; } @@ -994,14 +449,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)); } @@ -1027,7 +484,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; } @@ -1073,13 +531,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) { @@ -1139,7 +595,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 @@ -1252,15 +709,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; @@ -1273,8 +732,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; @@ -1326,12 +785,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; } @@ -1361,54 +820,50 @@ 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; } SocketOptions options; - options.user = this; options.keytable_pool = _socket->_keytable_pool; if (!FLAGS_rdma_use_polling) { - options.fd = _resource->comp_channel->fd; options.on_edge_triggered_events = PollCq; } if (Socket::Create(options, &_cq_sid) < 0) { PLOG(WARNING) << "Fail to create socket for cq"; return -1; } - + SocketUniquePtr cq_socket; + CHECK_EQ(0, Socket::Address(_cq_sid, &cq_socket)); + cq_socket->set_user(new CqSocketUser(_socket->id())); if (FLAGS_rdma_use_polling) { PollerAddCqSid(); + } else if (cq_socket->ResetFileDescriptor(_resource->comp_channel->fd) != 0) { + return -1; } - return 0; } -int RdmaEndpoint::BringUpQp(const ParsedHello& remote, bool is_server) { - // MTU negotiation (server side): compute min(local, client) and store it - // so FillLocalRdmaHello can advertise the negotiated value in the server - // hello reply. This is done before the g_skip_rdma_init early-return so - // that UT (which skips real QP bring-up) still sets _outgoing_mtu. +int RdmaEndpoint::BringUpQp(const RdmaConnectionInfo& remote, bool is_server) { if (is_server && remote.path_mtu.has_value()) { - uint32_t local_mtu = GetRdmaActiveMtu(); - _outgoing_mtu = std::min(local_mtu, *remote.path_mtu); + _outgoing_mtu = std::min(GetRdmaActiveMtu(), *remote.path_mtu); } if (BAIDU_UNLIKELY(g_skip_rdma_init)) { @@ -1422,11 +877,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; @@ -1438,7 +891,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; @@ -1487,14 +940,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; @@ -1506,29 +956,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); } } @@ -1560,7 +1010,7 @@ static int DrainCq(ibv_cq* cq) { } void RdmaEndpoint::DeallocateResources() { - if (_resource == nullptr) { + if (!_resource) { return; } if (FLAGS_rdma_use_polling) { @@ -1587,7 +1037,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; @@ -1602,11 +1052,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; @@ -1623,7 +1076,6 @@ void RdmaEndpoint::DeallocateResources() { if (remove_consumer) { s->_io_event.RemoveConsumer(s->_fd); } - s->_user = nullptr; // Do not release user (this RdmaEndpoint). s->_fd = -1; // Already remove fd from epoll fd. s->SetFailed(); } @@ -1687,8 +1139,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); @@ -1707,16 +1159,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; @@ -1727,24 +1178,19 @@ int RdmaEndpoint::ReqNotifyCq(bool send_cq, bool fatal_on_error) { } void RdmaEndpoint::PollCq(Socket* m) { - RdmaEndpoint* ep = static_cast(m->user()); - if (!ep) { - return; - } - + const auto* user = static_cast(m->user()); SocketUniquePtr s; - if (Socket::Address(ep->_socket->id(), &s) < 0) { + if (Socket::Address(user->socket_id, &s) < 0) { return; } + // The main socket pin keeps the endpoint and its resources alive. + RdmaTransport* rdma_transport = RdmaTransport::Get(s.get()); + RdmaEndpoint* ep = rdma_transport->_rdma_ep; // A queued callback may outlive Reset() and see the main Socket after // it has been revived with another CQ. if (m->id() != ep->_cq_sid) { return; } - auto* rdma_transport = static_cast(s->_transport.get()); - CHECK(ep == rdma_transport->_rdma_ep); - CHECK_GE(ep->_state.load(butil::memory_order_acquire), ESTABLISHED); - bool send = false; ibv_cq* cq = ep->_resource->recv_cq; @@ -1861,50 +1307,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() { @@ -1918,8 +1344,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; } @@ -1954,6 +1380,30 @@ void RdmaEndpoint::GlobalRelease() { std::vector RdmaEndpoint::_poller_groups; +// Thread creation is replaceable in startup failure tests. +decltype(&bthread_start_background) StartPollingThread = bthread_start_background; + +void RdmaEndpoint::PollerGroup::Stop(std::unique_lock& lock) { + if (stopping) { + stopped.wait(lock, [this] { return !stopping; }); + return; + } + if (!running.exchange(false, std::memory_order_relaxed)) { + return; + } + stopping = true; + lock.unlock(); + for (auto& poller : pollers) { + if (poller.tid != INVALID_BTHREAD) { + bthread_join(poller.tid, nullptr); + poller.tid = INVALID_BTHREAD; + } + } + lock.lock(); + stopping = false; + stopped.notify_all(); +} + int RdmaEndpoint::PollingModeInitialize(bthread_tag_t tag, std::function callback, std::function init_fn, @@ -1964,10 +1414,15 @@ int RdmaEndpoint::PollingModeInitialize(bthread_tag_t tag, auto& group = _poller_groups[tag]; auto& pollers = group.pollers; auto& running = group.running; - bool expected = false; - if (!running.compare_exchange_strong(expected, true)) { + std::unique_lock lock(group.mutex); + if (group.stopping) { + errno = EBUSY; + return -1; + } + if (running.load(std::memory_order_relaxed)) { return 0; } + running.store(true, std::memory_order_relaxed); struct FnArgs { Poller* poller; std::atomic* running; @@ -2012,18 +1467,22 @@ int RdmaEndpoint::PollingModeInitialize(bthread_tag_t tag, return nullptr; }; - for (int i = 0; i < FLAGS_rdma_poller_num; ++i) { + for (size_t i = 0; i < pollers.size(); ++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; pollers[i].init_fn = init_fn; pollers[i].release_fn = release_fn; - auto rc = bthread_start_background(&pollers[i].tid, &attr, fn, args); + auto rc = StartPollingThread(&pollers[i].tid, &attr, fn, args); if (rc != 0) { - LOG(ERROR) << "Fail to start rdma polling bthread"; + delete args; + group.Stop(lock); + LOG(ERROR) << "Fail to start rdma polling bthread: " + << berror(rc); + errno = rc; return -1; } } @@ -2035,36 +1494,28 @@ void RdmaEndpoint::PollingModeRelease(bthread_tag_t tag) { return; } auto& group = _poller_groups[tag]; - auto& pollers = group.pollers; - auto& running = group.running; - running.store(false, std::memory_order_relaxed); - for (int i = 0; i < FLAGS_rdma_poller_num; ++i) { - bthread_join(pollers[i].tid, nullptr); - } + std::unique_lock lock(group.mutex); + group.Stop(lock); } 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 16dd782d86..f26da2534b 100644 --- a/src/brpc/rdma/rdma_endpoint.h +++ b/src/brpc/rdma/rdma_endpoint.h @@ -31,51 +31,31 @@ #include "butil/containers/mpsc_queue.h" #include "butil/containers/optional.h" #include "brpc/socket.h" -#include "brpc/rdma/rdma_handshake_server.h" +#include "bthread/mutex.h" +#include "bthread/condition_variable.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; + // Active MTU in a client hello, negotiated MTU in a server hello (v3). + butil::optional path_mtu; }; struct RdmaResource { @@ -95,17 +75,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 +97,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 +116,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 +125,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 +168,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); @@ -279,12 +187,15 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Request completion notification on a send/recv CQ. int ReqNotifyCq(bool send_cq, bool fatal_on_error); + struct CqSocketUser : SocketUser { + explicit CqSocketUser(SocketId id) : socket_id(id) {} + void BeforeRecycle(Socket*) override { delete this; } + const SocketId socket_id; + }; + // 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,32 +204,13 @@ 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; - // MTU payload to advertise in the next local hello: - // Client: the locally queried active MTU (filled - // before C_HELLO_SEND); - // Server: the negotiated MTU = min(local_active_mtu, client_mtu) - // (filled in BringUpQp). + // Negotiated MTU returned by server-side resource setup. butil::optional _outgoing_mtu; // rdma resource @@ -335,9 +227,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; @@ -373,9 +262,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 @@ -400,8 +286,12 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // Poller group struct BAIDU_CACHELINE_ALIGNMENT PollerGroup { PollerGroup() : pollers(FLAGS_rdma_poller_num), running(false) {} + bthread::Mutex mutex; std::vector pollers; std::atomic running; + bool stopping{false}; + bthread::ConditionVariable stopped; + void Stop(std::unique_lock& lock); }; static std::vector _poller_groups; }; diff --git a/src/brpc/rdma/rdma_handshake.cpp b/src/brpc/rdma/rdma_handshake.cpp deleted file mode 100644 index 63681999bf..0000000000 --- a/src/brpc/rdma/rdma_handshake.cpp +++ /dev/null @@ -1,540 +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; - } - // Validate MTU: must be a valid ibv_mtu enum value (IBV_MTU_256..IBV_MTU_4096) - // to prevent invalid values from being passed to ibv_modify_qp. - if (msg.has_mtu() && - (msg.mtu() < IBV_MTU_256 || msg.mtu() > IBV_MTU_4096)) { - 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); - } - - // Advertise MTU if the endpoint has a value to advertise. - // Client side: queried local active MTU (filled before C_HELLO_SEND). - // Server side: negotiated MTU = min(local_mtu, client_mtu) - // (filled in BringUpQp). - // nullopt -> omit the field (peer falls back to IBV_MTU_1024). - if (ep->_outgoing_mtu.has_value()) { - msg->set_mtu(*ep->_outgoing_mtu); - } -} - -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; - } - if (msg.has_mtu()) { - out->path_mtu = msg.mtu(); - } -} - -} // 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(); - } - } - - // Query local active MTU so it can be advertised in the client hello. - // Best-effort: any failure just means we won't advertise MTU - // (the peer falls back to IBV_MTU_1024). - // In UT mode (g_skip_rdma_init=true), GetRdmaActiveMtu() returns the - // default IBV_MTU_1024 (g_active_mtu is never overwritten by OpenDevice), - // so the client hello still includes a valid MTU. - _ep->_outgoing_mtu = GetRdmaActiveMtu(); - - 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 b0d4cfa49d..0000000000 --- a/src/brpc/rdma/rdma_handshake.h +++ /dev/null @@ -1,193 +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; - - // MTU negotiation, v3 handshake only. - // nullopt means the peer did not advertise an MTU (v2 peer or older v3 - // peer that predates MTU negotiation). When engaged: - // - on the server side: the client's active MTU; - // - on the client side: the server's negotiated MTU - // (= min(local_active_mtu, client_mtu)). - butil::optional path_mtu; -}; - -// 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.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..d60ebccda7 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,20 @@ DECLARE_bool(usercode_in_pthread); extern SocketVarsCollector *g_vars; -void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { +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; - } - } - _tcp_transport = std::make_shared(); - _tcp_transport->Init(socket, options); + _on_edge_trigger = nullptr; + _rdma_state = RDMA_UNKNOWN; + _rdma_ep = new rdma::RdmaEndpoint(socket); } void RdmaTransport::Release() { @@ -70,31 +62,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 +125,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 +138,6 @@ int RdmaTransport::WaitEpollOut(butil::atomic *_epollout_butex, } } } - } else { - return _tcp_transport->WaitEpollOut(_epollout_butex, pollin, duetime); } return 0; } @@ -163,15 +177,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 +201,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 +231,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/server.cpp b/src/brpc/server.cpp index 3dcacb96f8..9d8a6e36b3 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -80,7 +80,6 @@ #include "brpc/rtmp.h" #include "brpc/builtin/common.h" // GetProgramName #include "brpc/details/tcmalloc_extension.h" -#include "brpc/rdma/rdma_helper.h" #include "brpc/baidu_master_service.h" #include "brpc/transport_factory.h" @@ -619,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() { @@ -641,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; @@ -865,14 +864,7 @@ int Server::StartInternal(const butil::EndPoint& endpoint, return -1; } - // Validate the user-provided ServerOptions BEFORE - // copy_and_fill_server_options below. This is important: - // copy_and_fill_server_options unconditionally transfers ownership of - // user-provided pointers (nshead_service, thrift_service, ...) into - // _options. If we instead validated against _options after the copy, - // a failed Start() would leave fake/invalid pointers behind in - // _options, and the NEXT Start() would attempt to `delete` them via - // FREE_PTR_IF_NOT_REUSED, crashing (see RdmaTest.server_option_invalid). + // Validate options before copying owned service pointers into the server. const ServerOptions default_opt; const ServerOptions& real_opt = opt ? *opt : default_opt; 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/src/brpc/socket.cpp b/src/brpc/socket.cpp index c4411c9606..4222f4efe4 100644 --- a/src/brpc/socket.cpp +++ b/src/brpc/socket.cpp @@ -2594,7 +2594,7 @@ void Socket::DebugSocket(std::ostream& os, SocketId id) { } #endif - os << "\nrdma={\n"; + os << "\ntransport={\n"; ptr->_transport->Debug(os); os << "}\n"; diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 8a1d372294..0e5d603907 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -57,11 +57,6 @@ class ChannelBalancer; } namespace rdma { class RdmaEndpoint; -class RdmaConnect; -class RdmaHandshakeClientV2; -class RdmaHandshakeServerV2; -class RdmaHandshakeClientV3; -class RdmaHandshakeServerV3; } namespace urma { @@ -75,14 +70,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); @@ -343,14 +340,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; @@ -365,6 +356,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; @@ -1015,8 +1008,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 transport for this connection. 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/tcp_transport.cpp b/src/brpc/tcp_transport.cpp index 98fea81674..6a97cfc72e 100644 --- a/src/brpc/tcp_transport.cpp +++ b/src/brpc/tcp_transport.cpp @@ -81,7 +81,10 @@ void TcpTransport::ProcessEvent(bthread_attr_t attr) { } } void TcpTransport::QueueMessage(InputMessageClosure& input_msg, - int* num_bthread_created, bool) { + int* num_bthread_created, bool last_msg) { + if (last_msg) { + return; + } InputMessageBase* to_run_msg = input_msg.release(); if (!to_run_msg) { return; 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..7f6e9560f3 100644 --- a/src/brpc/transport_factory.h +++ b/src/brpc/transport_factory.h @@ -22,8 +22,7 @@ #include "brpc/transport.h" namespace brpc { - -// Creates transport instances for a SocketMode. +// Creates transports and initializes their shared context. 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..bfab03b888 --- /dev/null +++ b/src/brpc/transport_handshake.cpp @@ -0,0 +1,427 @@ +// 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 != nullptr); + 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 != nullptr); + 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( + nullptr, 0, 0, 0, FrameSpec::FIXED); + return empty_spec; +} + +ServerHandshakeContext* ServerHandshakeContext::Create( + HandshakeAdapter* adapter) { + ServerHandshakeContext* context = + butil::get_object(); + if (context != nullptr) { + context->_adapter = adapter; + } + return context; +} + +void ServerHandshakeContext::Destroy() { + _adapter = nullptr; + butil::return_object(this); +} + +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; +} + +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, + butil::IOBuf* input, + bool push_back_on_not_mine, + bool* magic_matched) { + std::string payload; + const FrameResult frame_result = input != nullptr + ? 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, + butil::IOBuf* input, + bool* enabled) { + std::string payload; + const FrameResult frame_result = input != nullptr + ? 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, + butil::IOBuf* input) { + std::string payload; + const FrameResult frame_result = input != nullptr + ? 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, butil::IOBuf* input, + bool push_back_on_not_mine, HandshakeProtocol** selected) { + CHECK(!protocols.empty()); + 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, nullptr, 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 != 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. + 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, nullptr, 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, nullptr); + 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, butil::IOBuf* input, + HandshakeTransport* transport, bool fallback_on_not_mine) { + CHECK(!protocols.empty()); + CHECK(transport != nullptr); + + // 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 = 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); + _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 == nullptr) { + SetPhase(previous_phase); + } + return STEP_NEED_MORE; + } + if (result == STEP_ERROR) { + return FinishWithFailure(this, transport); + } + CHECK(selected != nullptr); + 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 != nullptr); + } + + 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..ef09ea439e --- /dev/null +++ b/src/brpc/transport_handshake.h @@ -0,0 +1,236 @@ +// 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(nullptr) {} + 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 let AdapterTransport read the handshake result uniformly. +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 = nullptr) + : _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.store(0, butil::memory_order_relaxed); + _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.load(butil::memory_order_relaxed); + } + void set_protocol_version(int version) { + _protocol_version.store(version, butil::memory_order_relaxed); + } + + 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, + butil::IOBuf* input, + HandshakeTransport* transport, + bool fallback_on_not_mine); + +private: + StepResult SendHello(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveHello(HandshakeProtocol* protocol, + butil::IOBuf* input, + bool push_back_on_not_mine, + bool* magic_matched = nullptr); + StepResult SendAck(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveAck(HandshakeProtocol* protocol, + butil::IOBuf* input, bool* enabled); + StepResult SendExtension(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveExtension(HandshakeProtocol* protocol, + butil::IOBuf* input); + StepResult SelectAndReceiveHello( + const std::vector& protocols, + butil::IOBuf* input, bool push_back_on_not_mine, + HandshakeProtocol** selected); + + SocketHandshakeIO _socket_io; + HandshakeIO* _io; + butil::atomic _phase; + butil::atomic _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/shm/shm_def.h b/src/brpc/ubshm/shm/shm_def.h index f237ca939f..ce3437cc8e 100644 --- a/src/brpc/ubshm/shm/shm_def.h +++ b/src/brpc/ubshm/shm/shm_def.h @@ -48,6 +48,8 @@ typedef struct { typedef struct ShmListNode { SHM shm; + bool unmap; + bool deallocate; struct ShmListNode *next; struct ShmListNode *prev; } ShmListNode; diff --git a/src/brpc/ubshm/shm/shm_ubs.cpp b/src/brpc/ubshm/shm/shm_ubs.cpp index ccbd325923..ef6027173e 100644 --- a/src/brpc/ubshm/shm/shm_ubs.cpp +++ b/src/brpc/ubshm/shm/shm_ubs.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "brpc/ubshm/timer/timer_mgr.h" #include "brpc/ubshm/common/thread_lock.h" @@ -33,6 +34,7 @@ #include "ubs_mem.h" #endif #include "shm_ubs.h" +#include "brpc/reloadable_flags.h" namespace brpc { namespace ubring { @@ -46,6 +48,7 @@ DEFINE_bool(shm_wr_delay_comp, true, "Indicates whether to enable the write rela DEFINE_int32(ub_flying_io_timeout_s, 5, "Time in seconds to wait for stopping data sending and receiving " "when the link is disconnected."); +BRPC_VALIDATE_GFLAG(ub_flying_io_timeout_s, brpc::PositiveInteger); char g_region_name[MAX_REGION_NAME_DESC_LENGTH] = {0}; int g_shm_timer_fd = 0; ShmList *g_shm_list = nullptr; @@ -133,68 +136,132 @@ static uint64_t AquireFlagIfWrDelayComp(const uint64_t flag) return flag | UBSM_FLAG_WR_DELAY_COMP; } -RETURN_CODE UbsShmLocalMalloc(SHM *shm) -{ - int ret = ubsmem_shmem_allocate(g_region_name, shm->name, shm->len, SHM_RIGHT_MODE, - AquireFlagIfWrDelayComp(UBSM_FLAG_ONLY_IMPORT_NONCACHE | UBSM_FLAG_MEM_ANONYMOUS)); -do { - if (ret == UBSM_ERR_ALREADY_EXIST) { - if (ubsmem_shmem_deallocate(shm->name) != UBSM_OK) { - LOG(ERROR) << "Ubs create shm name=" << shm->name << " failed, shm exists, ret=" << ret; - return SHM_ERR_EXIST; - } - LOG(INFO) << "Ubs delete shm name=" << shm->name << " success, try to recreate."; - ret = ubsmem_shmem_allocate(g_region_name, shm->name, shm->len, SHM_RIGHT_MODE, - AquireFlagIfWrDelayComp(UBSM_FLAG_ONLY_IMPORT_NONCACHE | UBSM_FLAG_MEM_ANONYMOUS)); +// Each task owns a copy of the mapping and its remaining operations. +static bool TryUbsCleanup(ShmListNode* task) { + if (task->unmap) { + const int ret = ubsmem_shmem_unmap(task->shm.addr, task->shm.len); if (ret != UBSM_OK) { - LOG(ERROR) << "Ubs recreate shm name=" << shm->name << " failed, ret=" << ret; - return SHM_ERR; + LOG_EVERY_SECOND(WARNING) << "Ubs unmap shm=" << task->shm.name + << " failed, ret=" << ret; + return false; } - } else if (ret != UBSM_OK) { - LOG(ERROR) << "Ubs create shm name=" << shm->name << " failed, ret=" << ret; - return SHM_ERR; + task->unmap = false; + } + if (task->deallocate) { + const int ret = ubsmem_shmem_deallocate(task->shm.name); + if (ret != UBSM_OK && ret != UBSM_ERR_NOT_FOUND) { + LOG_EVERY_SECOND(WARNING) << "Ubs deallocate shm=" << task->shm.name + << " failed, ret=" << ret; + return false; + } + task->deallocate = false; } -} while (0); + return true; +} - ret = ubsmem_shmem_map(nullptr, shm->len, PROT_READ | PROT_WRITE, MAP_SHARED, shm->name, 0, (void**)&(shm->addr)); - if (ret != UBSM_OK) { - LOG(ERROR) << "Ubs map shm=" << shm->name << " failed, ret=" << ret; - if (ret == UBSM_ERR_NOT_FOUND) { - return SHM_ERR_NOT_FOUND; +static RETURN_CODE ReleaseUbsShmLocked(ShmList* list, SHM* shm, + ShmListNode* task) { + if (list != nullptr) { + for (ShmListNode* node = list->head; node != nullptr; node = node->next) { + if (node->shm.addr == task->shm.addr && + node->shm.len == task->shm.len && + strcmp(node->shm.name, task->shm.name) == 0) { + node->deallocate |= task->deallocate; + free(task); + shm->addr = nullptr; + return UBRING_OK; + } } - ubsmem_shmem_deallocate(shm->name); - return SHM_ERR; } - - // Obtain memid via MXE - shm->memid = 1; // temporarily stubbed - LOG(INFO) << "Ubs malloc local shm=" << shm->name << " length=" << shm->len << " memid=" << shm->memid << " success."; + if (TryUbsCleanup(task)) { + free(task); + } else if (list != nullptr) { + task->next = nullptr; + task->prev = list->tail; + if (list->tail != nullptr) { + list->tail->next = task; + } else { + list->head = task; + } + list->tail = task; + ++list->size; + } else { + if (!task->unmap) { + shm->addr = nullptr; + } + free(task); + return UBRING_ERR; + } + shm->addr = nullptr; return UBRING_OK; } -RETURN_CODE UbsShmMunmap(SHM *shm) -{ - // unmap - if (shm->addr == nullptr) { - LOG(ERROR) << "Ubs input shm param is invalid, addr is NULL."; - return SHM_ERR_INPUT_INVALID; +// Success means cleanup completed or its remaining operations were queued. +static RETURN_CODE ReleaseUbsShm(SHM* shm, bool deallocate) { + // Reserve retry storage before changing the SDK mapping. + auto* task = static_cast(malloc(sizeof(ShmListNode))); + if (task == nullptr) { + return UBRING_ERR; } + task->shm = *shm; + task->unmap = shm->addr != nullptr; + task->deallocate = deallocate; + ShmList* list = g_shm_list; + if (list != nullptr) { + LOCK_GUARD(list->shm_lock); + return ReleaseUbsShmLocked(list, shm, task); + } + return ReleaseUbsShmLocked(nullptr, shm, task); +} - int ret = ubsmem_shmem_unmap(shm->addr, shm->len); - if (ret != UBSM_OK) { - if (ret == UBSM_ERR_NET) { - LOG(ERROR) << "Ubs unmap shm=" << shm->name << " failed, ubsm net err=" << ret; - AddShmToList(g_shm_list, shm); - return SHM_ERR_UBSM_NET_ERR; +RETURN_CODE UbsShmLocalMalloc(SHM* shm) { + ShmList* list = g_shm_list; + if (list == nullptr) { + return UBRING_ERR; + } + LOCK_GUARD(list->shm_lock); + for (ShmListNode* task = list->head; task != nullptr; task = task->next) { + if (strcmp(task->shm.name, shm->name) == 0) { + return SHM_ERR_EXIST; } - LOG(ERROR) << "Ubs unmap shm=" << shm->name << " length=" << shm->len << " failed, ret=" << ret; + } + // Reserve rollback storage before creating a named SDK resource. + std::unique_ptr rollback( + static_cast(malloc(sizeof(ShmListNode))), &free); + if (!rollback) { + return UBRING_ERR; + } + const uint64_t flags = AquireFlagIfWrDelayComp( + UBSM_FLAG_ONLY_IMPORT_NONCACHE | UBSM_FLAG_MEM_ANONYMOUS); + int ret = ubsmem_shmem_allocate( + g_region_name, shm->name, shm->len, SHM_RIGHT_MODE, flags); + if (ret == UBSM_ERR_ALREADY_EXIST) { + return SHM_ERR_EXIST; + } + if (ret != UBSM_OK) { + LOG(ERROR) << "Ubs allocate shm=" << shm->name << " failed, ret=" << ret; return SHM_ERR; } - - LOG(INFO) << "Ubs unmap shm=" << shm->name << " length=" << shm->len << " success."; + shm->addr = nullptr; + ret = ubsmem_shmem_map(nullptr, shm->len, PROT_READ | PROT_WRITE, + MAP_SHARED, shm->name, 0, + reinterpret_cast(&shm->addr)); + if (ret != UBSM_OK) { + LOG(ERROR) << "Ubs map shm=" << shm->name << " failed, ret=" << ret; + rollback->shm = *shm; + rollback->unmap = false; + rollback->deallocate = true; + ReleaseUbsShmLocked(list, shm, rollback.release()); + return ret == UBSM_ERR_NOT_FOUND ? SHM_ERR_NOT_FOUND : SHM_ERR; + } + shm->memid = 1; return UBRING_OK; } +RETURN_CODE UbsShmMunmap(SHM *shm) { + return ReleaseUbsShm(shm, false); +} + RETURN_CODE UbsShmFree(SHM *shm) { if (shm->addr == nullptr) { @@ -220,37 +287,8 @@ RETURN_CODE UbsShmFree(SHM *shm) return UBRING_OK; } -RETURN_CODE UbsShmLocalFree(SHM *shm) -{ - // unmap - if (shm->addr == nullptr) { - LOG(ERROR) << "Ubs input shm param is invalid, addr is NULL."; - return SHM_ERR_INPUT_INVALID; - } - - int ret = ubsmem_shmem_unmap(shm->addr, shm->len); - if (ret != UBSM_OK) { - if (ret == UBSM_ERR_NET) { - LOG(ERROR) << "Ubs unmap shm=" << shm->name << " failed, ubsm net err=" << ret; - AddShmToList(g_shm_list, shm); - return SHM_ERR_UBSM_NET_ERR; - } - LOG(WARNING) << "Ubs unmap shm=" << shm->name << " length=" << shm->len << " failed, ret=" << ret; - } - - // free - ret = ubsmem_shmem_deallocate(shm->name); - if (ret != UBSM_OK) { - if (ret == UBSM_ERR_IN_USING) { - LOG_EVERY_SECOND(INFO) << "Ubs delete shm=" << shm->name << " failed, resource attached=" << ret; - return SHM_ERR_RESOURCE_ATTACHED; - } - LOG(ERROR) << "Ubs delete shm=" << shm->name << " failed, ret=" << ret; - return SHM_ERR; - } - shm->addr = nullptr; - LOG(INFO) << "Ubs free local shm=" << shm->name << " length=" << shm->len << " success."; - return UBRING_OK; +RETURN_CODE UbsShmLocalFree(SHM *shm) { + return ReleaseUbsShm(shm, true); } RETURN_CODE UbsShmRemoteMalloc(SHM *shm) @@ -277,27 +315,8 @@ RETURN_CODE UbsShmLocalMmap(SHM *shm, int prot) return UBRING_OK; } -RETURN_CODE UbsShmRemoteFree(SHM *shm) -{ - // unmap - if (shm->addr == nullptr) { - LOG(ERROR) << "Ubs input shm param is invalid, addr is NULL."; - return SHM_ERR_INPUT_INVALID; - } - - int ret = ubsmem_shmem_unmap(shm->addr, shm->len); - if (ret != UBSM_OK) { - if (ret == UBSM_ERR_NET) { - LOG(ERROR) << "Ubs unmap shm=" << shm->name << " failed, ubsm net err=" << ret; - AddShmToList(g_shm_list, shm); - return SHM_ERR_UBSM_NET_ERR; - } - LOG(ERROR) << "Ubs unmap shm=" << shm->name << " length=" << shm->len << " failed, ret=" << ret; - return SHM_ERR; - } - - LOG(INFO) << "Ubs free Remote shm=" << shm->name << " length=" << shm->len << " success."; - return UBRING_OK; +RETURN_CODE UbsShmRemoteFree(SHM *shm) { + return ReleaseUbsShm(shm, false); } void UbsMemLoggerPrint(int level, const char *msg) @@ -375,77 +394,52 @@ RETURN_CODE UbsShmInit(void) return UBRING_OK; } -RETURN_CODE UbsShmFini(void) -{ - int ret = ubsmem_finalize(); +RETURN_CODE UbsShmFini(void) { + DeleteTimerSafe(g_shm_timer_fd); + g_shm_timer_fd = 0; + if (g_shm_list != nullptr) { + UbsShmCallback(g_shm_list); + if (g_shm_list->size != 0) { + UbsShmAddTimer(g_shm_list); + return UBRING_ERR; + } + } + const int ret = ubsmem_finalize(); if (ret != UBSM_OK) { LOG(ERROR) << "Ubs shm finalize fail, ret=" << ret; return UBRING_ERR; } - - if (UNLIKELY(DestroyShmTimer(g_shm_list) != UBRING_OK)) { - LOG(ERROR) << "Ubs shm list finalize failed."; - return UBRING_ERR; + if (g_shm_list != nullptr) { + DestroyShmTimer(g_shm_list); + g_shm_list = nullptr; } - - LOG(INFO) << "Ubs shm finalize success."; return UBRING_OK; } -static void DeleteShmToList(ShmList* shm_list) -{ - if (shm_list == nullptr || shm_list->head == nullptr) { - return; - } - - ShmListNode *cur_node = shm_list->head; - shm_list->head = cur_node->next; - if (shm_list->head != nullptr) { - shm_list->head->prev = nullptr; - } else { - shm_list->tail = nullptr; - } - LOG(INFO) << "Delete shm to list, name=" << cur_node->shm.name << " size=" << shm_list->size; - FREE_PTR(cur_node); - shm_list->size--; -} - -void *UbsShmCallback(void* args) -{ - ShmList *shm_list = (ShmList*)args; - if (UNLIKELY(shm_list == nullptr)) { - LOG(ERROR) << "Shm list is null."; +void *UbsShmCallback(void* args) { + auto* list = static_cast(args); + if (list == nullptr) { return nullptr; } - - LOCK_GUARD(shm_list->shm_lock); - while (shm_list->head != nullptr) { - SHM shm = shm_list->head->shm; - if (shm.addr == nullptr) { - LOG(ERROR) << "Ubs input shm param is invalid, addr is NULL."; - return nullptr; - } - - int ret = ubsmem_shmem_unmap(shm.addr, shm.len); - if (ret != UBSM_OK) { - if (ret == UBSM_ERR_NET) { - return nullptr; + LOCK_GUARD(list->shm_lock); + for (ShmListNode* task = list->head; task != nullptr;) { + ShmListNode* next = task->next; + if (TryUbsCleanup(task)) { + if (task->prev != nullptr) { + task->prev->next = next; + } else { + list->head = next; } - LOG(ERROR) << "Ubs unmap shm=" << shm.name << " length=" << shm.len << " failed, ret=" << ret; - return nullptr; - } - LOG(INFO) << "Ubs unmap shm=" << shm.name << " length=" << shm.len << " success."; - - ret = ubsmem_shmem_deallocate(shm.name); - if (ret != UBSM_OK) { - DeleteShmToList(shm_list); - LOG(ERROR) << "Ubs delete shm=" << shm.name << " failed, ret=" << ret; - return nullptr; + if (next != nullptr) { + next->prev = task->prev; + } else { + list->tail = task->prev; + } + --list->size; + free(task); } - DeleteShmToList(shm_list); - LOG(INFO) << "Ubs free local shm=" << shm.name << " length=" << shm.len << " success."; + task = next; } - return nullptr; } @@ -485,6 +479,7 @@ RETURN_CODE InitShmTimer(ShmList **shm_list) if (UbsShmAddTimer(*shm_list) == UBRING_ERR) { LOG(ERROR) << "Ubs add timer failed."; + pthread_mutex_destroy(&(*shm_list)->shm_lock); FREE_PTR(*shm_list); return UBRING_ERR; } @@ -494,6 +489,7 @@ RETURN_CODE InitShmTimer(ShmList **shm_list) RETURN_CODE DestroyShmTimer(ShmList *shm_list) { DeleteTimerSafe((uint32_t)g_shm_timer_fd); + g_shm_timer_fd = 0; if (shm_list == nullptr) { LOG(WARNING) << "Shm list is null."; return UBRING_ERR; @@ -511,56 +507,5 @@ RETURN_CODE DestroyShmTimer(ShmList *shm_list) return UBRING_OK; } -RETURN_CODE IsExistInShmList(ShmList *shm_list, const SHM *shm) -{ - if (UNLIKELY(shm_list == nullptr || shm == nullptr)) { - LOG(ERROR) << "Shm list or shm is null."; - return UBRING_ERR; - } - LOCK_GUARD(shm_list->shm_lock); - - ShmListNode *cur_node = shm_list->head; - while (cur_node != nullptr) { - if (strcmp(cur_node->shm.name, shm->name) == 0 && cur_node->shm.len == shm->len) { - return UBRING_OK; - } - cur_node = cur_node->next; - } - return UBRING_ERR; -} - -RETURN_CODE AddShmToList(ShmList *shm_list, SHM *shm) -{ - if (shm_list == nullptr || shm == nullptr) { - LOG(ERROR) << "Shm list or shm is null."; - return UBRING_ERR; - } - - if (IsExistInShmList(shm_list, shm) == UBRING_OK) { - LOG(ERROR) << "Shm name=" << shm->name << " is exist in shm list."; - return UBRING_ERR; - } - - ShmListNode *new_shm_node = (ShmListNode *)malloc(sizeof(ShmListNode)); - if (new_shm_node == nullptr) { - LOG(ERROR) << "Malloc shm node failed."; - return UBRING_ERR; - } - - memcpy(&new_shm_node->shm, shm, sizeof(SHM)); - LOCK_GUARD(shm_list->shm_lock); - new_shm_node->next = nullptr; - new_shm_node->prev = shm_list->tail; - if (shm_list->tail) { - shm_list->tail->next = new_shm_node; - shm_list->tail = new_shm_node; - } else { - shm_list->head = new_shm_node; - shm_list->tail = new_shm_node; - } - shm_list->size++; - LOG(INFO) << "Add shm to list success, shm name=" << shm->name << " size=" << shm_list->size; - return UBRING_OK; -} } } diff --git a/src/brpc/ubshm/shm/shm_ubs.h b/src/brpc/ubshm/shm/shm_ubs.h index ee936c0a55..2b908cd1fc 100644 --- a/src/brpc/ubshm/shm/shm_ubs.h +++ b/src/brpc/ubshm/shm/shm_ubs.h @@ -44,8 +44,6 @@ void *UbsShmCallback(void* args); RETURN_CODE UbsShmAddTimer(ShmList *shm_list); RETURN_CODE InitShmTimer(ShmList **shm_list); RETURN_CODE DestroyShmTimer(ShmList *shm_list); -RETURN_CODE AddShmToList(ShmList *shm_list, SHM *shm); -RETURN_CODE IsExistInShmList(ShmList *shm_list, const SHM *shm); } } #endif //BRPC_SHM_UBS_H diff --git a/src/brpc/ubshm/timer/timer_mgr.cpp b/src/brpc/ubshm/timer/timer_mgr.cpp index b5e0c9ef3b..a91c0f2149 100644 --- a/src/brpc/ubshm/timer/timer_mgr.cpp +++ b/src/brpc/ubshm/timer/timer_mgr.cpp @@ -25,6 +25,7 @@ #include #include #include "brpc/ubshm/timer/timer_mgr.h" +#include "butil/scoped_lock.h" namespace brpc { namespace ubring { @@ -34,7 +35,10 @@ std::atomic g_total_timer_num(0); TimerFdCtx *g_timer_fd_ctx_map = nullptr; uint32_t g_max_system_fd = 0; static pthread_t g_epoll_execute_thread = 0; -static int32_t g_timer_module_initialized = 0; +static std::atomic g_timer_module_initialized(0); +static pthread_mutex_t g_timer_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_timer_idle = PTHREAD_COND_INITIALIZER; +static thread_local TimerFdCtx* g_running_timer = nullptr; #if defined(OS_MACOSX) static int timerfd_create_macosx(int clockid, int flags); @@ -43,41 +47,47 @@ static int timerfd_settime_macosx(int fd, int flags, itimerspec *old_value); #endif -static RETURN_CODE DeleteTimerInner(uint32_t fd) { - if (g_timer_fd_ctx_map == nullptr) { - return UBRING_OK; - } - - if (pthread_spin_lock(&g_timer_fd_ctx_map[fd].spin_lock) != 0) { - return UBRING_ERR; - } - - if (g_timer_fd_ctx_map[fd].status == TIMER_CONTEXT_NOT_USING) { - pthread_spin_unlock(&g_timer_fd_ctx_map[fd].spin_lock); - return UBRING_OK; - } - - g_timer_fd_ctx_map[fd].status = TIMER_CONTEXT_NOT_USING; - g_timer_fd_ctx_map[fd].cb = nullptr; - g_timer_fd_ctx_map[fd].args = nullptr; - g_timer_fd_ctx_map[fd].periodical = 0; - g_timer_fd_ctx_map[fd].fd = 0; - - pthread_spin_unlock(&g_timer_fd_ctx_map[fd].spin_lock); - +// Called with g_timer_mutex held, after the callback has finished. +static void CloseTimerLocked(uint32_t fd) { + TimerFdCtx& ctx = g_timer_fd_ctx_map[fd]; #if defined(OS_LINUX) - epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, (int)fd, nullptr); + epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, fd, nullptr); #elif defined(OS_MACOSX) struct kevent evt; EV_SET(&evt, fd, EVFILT_TIMER, EV_DELETE, 0, 0, nullptr); kevent(g_epoll_fd, &evt, 1, nullptr, 0, nullptr); #endif + close(fd); + ctx.status = TIMER_CONTEXT_NOT_USING; + ctx.cb = nullptr; + ctx.args = nullptr; + ctx.periodical = 0; + g_total_timer_num.fetch_sub(1); + pthread_cond_broadcast(&g_timer_idle); +} - uint64_t exp = 0; - read((int)fd, &exp, sizeof(exp)); - - close((int)fd); - std::atomic_fetch_sub(&g_total_timer_num, 1U); +static RETURN_CODE DeleteTimerInner(uint32_t fd) { + pthread_mutex_lock(&g_timer_mutex); + if (g_timer_fd_ctx_map == nullptr || fd >= g_max_system_fd) { + pthread_mutex_unlock(&g_timer_mutex); + return UBRING_OK; + } + TimerFdCtx& ctx = g_timer_fd_ctx_map[fd]; + const uint32_t generation = ctx.generation; + if (ctx.status != TIMER_CONTEXT_NOT_USING) { + ctx.canceled = true; + if (ctx.status == TIMER_CONTEXT_CALLBACK_ONGOING) { + // A callback can cancel itself. Its epilogue closes the descriptor. + while (g_running_timer != &ctx && + ctx.generation == generation && + ctx.status == TIMER_CONTEXT_CALLBACK_ONGOING) { + pthread_cond_wait(&g_timer_idle, &g_timer_mutex); + } + } else { + CloseTimerLocked(fd); + } + } + pthread_mutex_unlock(&g_timer_mutex); return UBRING_OK; } @@ -92,47 +102,18 @@ static RETURN_CODE StartTimeEpoll(void) { return UBRING_ERR; } + g_timer_module_initialized = 1; int ret = pthread_create(&g_epoll_execute_thread, nullptr, TimerEpoll, nullptr); if (UNLIKELY(ret != 0)) { + g_timer_module_initialized = 0; LOG(ERROR) << "Failed to create thread err=" << ret; return UBRING_ERR; } return UBRING_OK; } -static RETURN_CODE TimerSpinLocksInit(void) { - if (g_timer_fd_ctx_map == nullptr) { - LOG(ERROR) << "Timer module is not fully initialized."; - return UBRING_ERR; - } - - for (uint32_t fd = 0; fd < g_max_system_fd; fd++) { - int ret = pthread_spin_init(&g_timer_fd_ctx_map[fd].spin_lock, - PTHREAD_PROCESS_PRIVATE); - if (ret != EOK) { - LOG(ERROR) << "Failed to initialize spin lock for fd=" << fd; - for (uint32_t cleanup_fd = 0; cleanup_fd < fd; cleanup_fd++) { - pthread_spin_destroy(&g_timer_fd_ctx_map[cleanup_fd].spin_lock); - } - return UBRING_ERR; - } - } - return UBRING_OK; -} - -static RETURN_CODE ExecuteCallback(int32_t timer_fd) { - UnifiedCallback((void *)(&g_timer_fd_ctx_map[timer_fd])); - return UBRING_OK; -} - static RETURN_CODE TimerCtxMapCompletion(void) { memset(g_timer_fd_ctx_map, 0, sizeof(TimerFdCtx) * g_max_system_fd); - - RETURN_CODE ret = TimerSpinLocksInit(); - if (ret != UBRING_OK) { - LOG(ERROR) << "Failed to init spin locks for timer module."; - return UBRING_ERR; - } return UBRING_OK; } @@ -178,31 +159,45 @@ RETURN_CODE TimerInit(void) { return UBRING_OK; } -void *UnifiedCallback(void *args) { - TimerFdCtx *ctx = (TimerFdCtx *)args; - if (pthread_spin_lock(&ctx->spin_lock) != 0) { - return nullptr; +static void DispatchTimer(uint32_t fd, uint32_t generation, + bool consume_event) { + pthread_mutex_lock(&g_timer_mutex); + if (fd >= g_max_system_fd) { + pthread_mutex_unlock(&g_timer_mutex); + return; } - - if (ctx->status == TIMER_CONTEXT_NOT_USING) { - pthread_spin_unlock(&ctx->spin_lock); - return nullptr; + TimerFdCtx& ctx = g_timer_fd_ctx_map[fd]; + if (ctx.generation != generation || + ctx.status != TIMER_CONTEXT_EPOLL_WAITING || ctx.canceled) { + pthread_mutex_unlock(&g_timer_mutex); + return; } +#if defined(OS_LINUX) + if (consume_event) { + uint64_t expirations; + if (read(fd, &expirations, sizeof(expirations)) != sizeof(expirations)) { + pthread_mutex_unlock(&g_timer_mutex); + return; + } + } +#endif + auto cb = ctx.cb; + void* cb_args = ctx.args; + ctx.status = TIMER_CONTEXT_CALLBACK_ONGOING; + pthread_mutex_unlock(&g_timer_mutex); - void *(*cb)(void *) = ctx->cb; - void *cb_args = ctx->args; - uint32_t fd = ctx->fd; - int is_periodical = ctx->periodical; - ctx->status = TIMER_CONTEXT_CALLBACK_ONGOING; - - pthread_spin_unlock(&ctx->spin_lock); - + g_running_timer = &ctx; cb(cb_args); + g_running_timer = nullptr; - if (!is_periodical) { - DeleteTimerInner(fd); + pthread_mutex_lock(&g_timer_mutex); + if (ctx.canceled || !ctx.periodical) { + CloseTimerLocked(fd); + } else { + ctx.status = TIMER_CONTEXT_EPOLL_WAITING; + pthread_cond_broadcast(&g_timer_idle); } - return nullptr; + pthread_mutex_unlock(&g_timer_mutex); } void *TimerEpoll(void *args) { @@ -243,97 +238,56 @@ void *TimerEpoll(void *args) { for (int32_t i = 0; i < ready_num; i++) { #if defined(OS_LINUX) struct epoll_event *event = &ready_events[i]; - int32_t timer_fd = event->data.fd; + int32_t timer_fd = static_cast(event->data.u64); #elif defined(OS_MACOSX) struct kevent *event = &ready_events[i]; int32_t timer_fd = event->ident; #endif - uint64_t exp = 0; - if (read(timer_fd, &exp, sizeof(exp)) < 0) { - if (errno != EBADF) { - LOG(ERROR) << "Failed to read timerfd=" << timer_fd << " errno=" << errno; - } - continue; - } - if (TimerFdCtxValidate((uint32_t)timer_fd) != UBRING_OK) { - continue; - } - - RETURN_CODE ret = ExecuteCallback(timer_fd); - if (ret != UBRING_OK) { - LOG(ERROR) << "Failed execute callback ret=" << ret; - DeleteTimerInner((uint32_t)timer_fd); - continue; - } +#if defined(OS_LINUX) + const uint32_t generation = event->data.u64 >> 32; +#elif defined(OS_MACOSX) + const uint32_t generation = + reinterpret_cast(event->udata); +#endif + DispatchTimer(timer_fd, generation, true); } } return nullptr; } void DeleteTimerSafe(uint32_t fd) { - if (g_timer_fd_ctx_map == nullptr) { - return; - } - - if (pthread_spin_lock(&g_timer_fd_ctx_map[fd].spin_lock) != 0) { - return; - } - - if (g_timer_fd_ctx_map[fd].status == TIMER_CONTEXT_NOT_USING) { - pthread_spin_unlock(&g_timer_fd_ctx_map[fd].spin_lock); - return; - } - - g_timer_fd_ctx_map[fd].status = TIMER_CONTEXT_NOT_USING; - g_timer_fd_ctx_map[fd].cb = nullptr; - g_timer_fd_ctx_map[fd].args = nullptr; - g_timer_fd_ctx_map[fd].periodical = 0; - g_timer_fd_ctx_map[fd].fd = 0; - - pthread_spin_unlock(&g_timer_fd_ctx_map[fd].spin_lock); - -#if defined(OS_LINUX) - epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, (int)fd, nullptr); -#elif defined(OS_MACOSX) - struct kevent evt; - EV_SET(&evt, fd, EVFILT_TIMER, EV_DELETE, 0, 0, nullptr); - kevent(g_epoll_fd, &evt, 1, nullptr, 0, nullptr); -#endif - - uint64_t exp = 0; - read((int)fd, &exp, sizeof(exp)); - - close((int)fd); - std::atomic_fetch_sub(&g_total_timer_num, 1U); + DeleteTimerInner(fd); } void DeleteTimer(uint32_t fd) { - if (g_timer_fd_ctx_map == nullptr) { - LOG(WARNING) << "The timer is not initialized."; - return; - } - - g_timer_fd_ctx_map[fd].periodical = 0; + DeleteTimerInner(fd); } int32_t TimerStart(const itimerspec *time, void *(*cb)(void *), void *args) { + BAIDU_SCOPED_LOCK(g_timer_mutex); if (g_epoll_fd == -1) { LOG(ERROR) << "Timer epoll/kqueue encountered internal error."; return -1; } #if defined(OS_LINUX) - int timer_fd = timerfd_create(CLOCK_MONOTONIC, 0); + int timer_fd = timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC); #elif defined(OS_MACOSX) int timer_fd = timerfd_create_macosx(CLOCK_MONOTONIC, 0); #endif if (UNLIKELY(timer_fd >= (int)g_max_system_fd || timer_fd == -1)) { + if (timer_fd >= 0) { + close(timer_fd); + } LOG(ERROR) << "Failed to create timerfd=" << timer_fd << " errno=" << errno; return -1; } + ++g_timer_fd_ctx_map[timer_fd].generation; + g_timer_fd_ctx_map[timer_fd].canceled = false; + g_timer_fd_ctx_map[timer_fd].periodical = 0; g_timer_fd_ctx_map[timer_fd].status = TIMER_CONTEXT_EPOLL_WAITING; g_timer_fd_ctx_map[timer_fd].cb = cb; g_timer_fd_ctx_map[timer_fd].args = args; @@ -344,9 +298,12 @@ int32_t TimerStart(const itimerspec *time, void *(*cb)(void *), void *args) { } #if defined(OS_LINUX) + const uint64_t token = + (static_cast(g_timer_fd_ctx_map[timer_fd].generation) << 32) | + static_cast(timer_fd); struct epoll_event event = { .events = EPOLLIN, - .data = {.fd = timer_fd} + .data = {.u64 = token} }; int32_t ret = epoll_ctl(g_epoll_fd, EPOLL_CTL_ADD, timer_fd, &event); @@ -355,7 +312,8 @@ int32_t TimerStart(const itimerspec *time, void *(*cb)(void *), void *args) { uint64_t timeout_nsec = time->it_value.tv_sec * 1000000000ULL + time->it_value.tv_nsec; uint64_t interval_nsec = time->it_interval.tv_sec * 1000000000ULL + time->it_interval.tv_nsec; EV_SET(&event, timer_fd, EVFILT_TIMER, EV_ADD | EV_ENABLE, 0, - timeout_nsec / 1000000, nullptr); + timeout_nsec / 1000000, + reinterpret_cast(static_cast(g_timer_fd_ctx_map[timer_fd].generation))); int32_t ret = kevent(g_epoll_fd, &event, 1, nullptr, 0, nullptr); #endif @@ -412,9 +370,7 @@ void TimerModuleDestroy(void) { uint32_t max_fd = g_max_system_fd; if (g_timer_fd_ctx_map) { for (uint32_t fd = 0; fd < max_fd; fd++) { - if (g_timer_fd_ctx_map[fd].status != TIMER_CONTEXT_NOT_USING) { - DeleteTimerSafe(fd); - } + DeleteTimerSafe(fd); } } close(g_epoll_fd); @@ -428,23 +384,6 @@ void TimerModuleDestroy(void) { } } -RETURN_CODE TimerFdCtxValidate(uint32_t fd) { - if (fd >= g_max_system_fd) { - LOG(ERROR) << "TimerFd=" << fd << " is out of range=" << g_max_system_fd; - return UBRING_ERR; - } - if (g_timer_fd_ctx_map[fd].status == TIMER_CONTEXT_NOT_USING) { - LOG(ERROR) << "TimerFd=" << fd << " has wrong status=" << g_timer_fd_ctx_map[fd].status; - return UBRING_ERR; - } - if (g_timer_fd_ctx_map[fd].cb == nullptr) { - LOG(ERROR) << "The callback is not set."; - return UBRING_ERR; - } - - return UBRING_OK; -} - #if defined(OS_MACOSX) static int timerfd_create_macosx(int clockid, int flags) { int pipefd[2]; @@ -465,4 +404,4 @@ static int timerfd_settime_macosx(int fd, int flags, #endif } // namespace ubring -} // namespace brpc \ No newline at end of file +} // namespace brpc diff --git a/src/brpc/ubshm/timer/timer_mgr.h b/src/brpc/ubshm/timer/timer_mgr.h index 1b42caef04..455fda0beb 100644 --- a/src/brpc/ubshm/timer/timer_mgr.h +++ b/src/brpc/ubshm/timer/timer_mgr.h @@ -54,20 +54,21 @@ typedef struct { uint32_t fd; TimerFdCtxStatus status; uint32_t periodical; - pthread_spinlock_t spin_lock; + uint32_t generation; + bool canceled; } TimerFdCtx; RETURN_CODE TimerInit(void); void TimerModuleDestroy(void); -void *UnifiedCallback(void *args); void *TimerEpoll(void *args); int32_t TimerStart(const itimerspec *time, void *(*cb)(void *), void *args); uint32_t GetActiveTimerNum(void); void CloseTimerFd(int fd); +// Wait for an active callback before closing its descriptor. Self-cancellation +// closes the descriptor when the callback returns. void DeleteTimerSafe(uint32_t fd); void DeleteTimer(uint32_t fd); -RETURN_CODE TimerFdCtxValidate(uint32_t fd); } } -#endif //BRPC_TIMER_MGR_H \ No newline at end of file +#endif //BRPC_TIMER_MGR_H diff --git a/src/brpc/ubshm/ub_endpoint.cpp b/src/brpc/ubshm/ub_endpoint.cpp index 97e39e5d77..0ef884aa7c 100644 --- a/src/brpc/ubshm/ub_endpoint.cpp +++ b/src/brpc/ubshm/ub_endpoint.cpp @@ -18,25 +18,22 @@ #if BRPC_WITH_UBRING #include +#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/scoped_lock.h" +#include "butil/logging.h" // CHECK, LOG +#include DECLARE_int32(task_group_ntags); @@ -45,9 +42,6 @@ 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"); @@ -55,686 +49,23 @@ DEFINE_bool(ub_disable_bthread, false, "Disable bthread in UBRing polling mode." 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); -} - -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() { - DeallocateResources(); +void UBShmEndpoint::Reset(bool abort_upgrade) { + _receive_events_started.store(false, butil::memory_order_release); + DeallocateResources(abort_upgrade); 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 +106,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 +130,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 +145,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,24 +182,54 @@ 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::DeallocateResources() { +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(bool abort_upgrade) { if (!_ub_ring) { return; } PollerRegisterEvent(PollerSidOp::REMOVE); - _ub_ring->UbrTrxClose(); + if (abort_upgrade) { + _ub_ring->UbrTrxAbort(); + } else { + _ub_ring->UbrTrxClose(); + } if (INVALID_SOCKET_ID != _poller_sid) { SocketUniquePtr s; if (Socket::Address(_poller_sid, &s) == 0) { @@ -867,11 +241,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 +269,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,28 +305,69 @@ 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(); } +} +void UBShmEndpoint::Poll(const PollerSidOp& op) { + SocketUniquePtr owner; + if (Socket::Address(op.owner_sid, &owner) < 0) { + return; + } + UBShmEndpoint* ep = UBShmTransport::Get(owner.get())->_ub_ep; + if (ep == nullptr || ep->_poller_sid != op.sid) { + return; + } + if (op.event & EPOLLIN) { + PollIn(ep, op.event); + } + if (op.event & EPOLLOUT) { + PollOut(ep, op.event); + } } int UBShmEndpoint::GlobalInitialize() { - g_ubring_resource_mutex = new butil::Mutex; _poller_groups = std::vector(FLAGS_task_group_ntags); - return 0; + // Stop polling before the static poller storage is destroyed. + static const int registered = std::atexit(GlobalRelease); + return registered == 0 ? 0 : -1; } void UBShmEndpoint::GlobalRelease() { - for (int i = 0; i < FLAGS_task_group_ntags; ++i) { + for (size_t i = 0; i < _poller_groups.size(); ++i) { PollingModeRelease(i); } } std::vector UBShmEndpoint::_poller_groups; +// Thread creation is replaceable in startup failure tests. +decltype(&bthread_start_background) StartPollingThread = bthread_start_background; + +void UBShmEndpoint::PollerGroup::Stop(std::unique_lock& lock) { + if (stopping) { + stopped.wait(lock, [this] { return !stopping; }); + return; + } + if (!running.exchange(false, std::memory_order_relaxed)) { + return; + } + stopping = true; + lock.unlock(); + for (auto& poller : pollers) { + if (poller.tid != INVALID_BTHREAD) { + bthread_join(poller.tid, nullptr); + poller.tid = INVALID_BTHREAD; + } + } + lock.lock(); + stopping = false; + stopped.notify_all(); +} + int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, std::function callback, std::function init_fn, @@ -956,10 +375,15 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, auto& group = _poller_groups[tag]; auto& pollers = group.pollers; auto& running = group.running; - bool expected = false; - if (!running.compare_exchange_strong(expected, true)) { + std::unique_lock lock(group.mutex); + if (group.stopping) { + errno = EBUSY; + return -1; + } + if (running.load(std::memory_order_relaxed)) { return 0; } + running.store(true, std::memory_order_relaxed); struct FnArgs { Poller* poller; std::atomic* running; @@ -968,7 +392,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,31 +401,17 @@ 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) { - SocketUniquePtr s; - if (Socket::Address(poller_sid.sid, &s) < 0) { - continue; - } - UBShmEndpoint* ep = static_cast(s->user()); - if (!ep) { - continue; - } - - if (poller_sid.events & EPOLLIN) { - PollIn(ep, poller_sid.events); - } - - if (poller_sid.events & EPOLLOUT) { - PollOut(ep, poller_sid.events); - } + for (auto cq : cq_sids) { + Poll(cq); } if (poller->callback) { poller->callback(); @@ -1017,18 +427,22 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, return nullptr; }; - for (int i = 0; i < FLAGS_ub_poller_num; ++i) { + for (size_t i = 0; i < pollers.size(); ++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; pollers[i].init_fn = init_fn; pollers[i].release_fn = release_fn; - auto rc = bthread_start_background(&pollers[i].tid, &attr, fn, args); + auto rc = StartPollingThread(&pollers[i].tid, &attr, fn, args); if (rc != 0) { - LOG(ERROR) << "Fail to start ubring polling bthread"; + delete args; + group.Stop(lock); + LOG(ERROR) << "Fail to start ubring polling bthread: " + << berror(rc); + errno = rc; return -1; } } @@ -1037,22 +451,18 @@ int UBShmEndpoint::PollingModeInitialize(bthread_tag_t tag, void UBShmEndpoint::PollingModeRelease(bthread_tag_t tag) { auto& group = _poller_groups[tag]; - auto& pollers = group.pollers; - auto& running = group.running; - running.store(false, std::memory_order_relaxed); - for (int i = 0; i < FLAGS_ub_poller_num; ++i) { - bthread_join(pollers[i].tid, nullptr); - } + std::unique_lock lock(group.mutex); + group.Stop(lock); } -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; auto& poller = pollers[index]; if (INVALID_SOCKET_ID != _poller_sid) { - poller.op_queue.Enqueue(PollerSidOp{_poller_sid, events, op}); + poller.op_queue.Enqueue( + PollerSidOp{_poller_sid, _socket_id, events, op}); } } diff --git a/src/brpc/ubshm/ub_endpoint.h b/src/brpc/ubshm/ub_endpoint.h index 526966c9f7..d83da7846f 100644 --- a/src/brpc/ubshm/ub_endpoint.h +++ b/src/brpc/ubshm/ub_endpoint.h @@ -20,78 +20,38 @@ #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 "bthread/mutex.h" +#include "bthread/condition_variable.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; @@ -103,7 +63,13 @@ friend class Socket; static void GlobalRelease(); // Reset the endpoint (for next use) - void Reset(); + void Reset(bool abort_upgrade = false); + 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 +78,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 +99,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 +107,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); @@ -180,60 +116,37 @@ friend class Socket; UbrDataFormat format); // 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); + void DeallocateResources(bool abort_upgrade = false); - // 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; + SocketId owner_sid; + uint32_t event; OpType type; }; + static void Poll(const PollerSidOp& op); + 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 +158,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 @@ -256,13 +168,18 @@ friend class Socket; // Poller group struct BAIDU_CACHELINE_ALIGNMENT PollerGroup { PollerGroup() : pollers(FLAGS_ub_poller_num), running(false) {} + bthread::Mutex mutex; std::vector pollers; std::atomic running; + bool stopping{false}; + bthread::ConditionVariable stopped; + void Stop(std::unique_lock& lock); }; 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/ub_ring.cpp b/src/brpc/ubshm/ub_ring.cpp index 595b96d03b..1dada487bf 100644 --- a/src/brpc/ubshm/ub_ring.cpp +++ b/src/brpc/ubshm/ub_ring.cpp @@ -22,6 +22,7 @@ #include #include "bthread/bthread.h" #include "butil/logging.h" +#include "butil/scoped_lock.h" #include "brpc/ubshm/ub_ring.h" #include "brpc/ubshm/ub_ring_manager.h" #include "brpc/ubshm/shm/shm_ipc.h" @@ -46,6 +47,20 @@ UBRing::UBRing() UBRing::~UBRing() {} +RETURN_CODE UBRing::RegisterShmName(SHM* target, const char* name) { + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } + memcpy(target->name, name, sizeof(target->name)); + return UBRING_OK; +} + +void UBRing::RecordShm(SHM* target, const SHM& source) { + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + *target = source; +} + RETURN_CODE UBRing::UbrTrxMapShm(SHM *local_shm, SHM *remote_shm) { RETURN_CODE rc = UbrTrxMapLocalShm(local_shm); @@ -61,7 +76,52 @@ RETURN_CODE UBRing::UbrTrxMapShm(SHM *local_shm, SHM *remote_shm) return UBRING_OK; } +RETURN_CODE UBRing::UbrTrxAbort() { + if (_trx == nullptr) { + return UBRING_OK; + } + if (ATOMIC_LOAD(_trx->close_state) != UBR_CLOSE_FIRST) { + return ReleaseTrx(); + } + if (_trx->ubr_tx.local_tx_event_q.addr != nullptr) { + const RETURN_CODE close_rc = UbrTrxCloseCheck(_trx); + if (close_rc != UBRING_OK) { + return close_rc; + } + } + int expected = UBR_CLOSE_FIRST; + if (!ATOMIC_COMPARE_EXCHANGE_STRONG( + _trx->close_state, expected, UBR_CLOSE_SECOND)) { + return ReleaseTrx(); + } + // Drain timer callbacks before unmapping the queues they access. + DeleteTimerSafe(_trx->timer_fd); + DeleteTimerSafe(_trx->hb_timer_fd); + DeleteTimerSafe(_trx->clear_timer_fd); + if (UbrTrxFreeShm(_trx) == UBRING_OK) { + return ReleaseTrx(); + } + // Transfer the ring reference when retry storage could not be acquired. + UbrTrx* trx = _trx; + _trx = nullptr; + trx->clear_timer_fd = 0; + if (UbrAddAsynClearTimer(trx) != UBRING_OK) { + UbrAsynClearCallback(trx); + } + return UBRING_OK; +} + RETURN_CODE UBRing::UbrTrxClose() { + if (_trx == nullptr) { + return UBRING_OK; + } + if (ATOMIC_LOAD(_trx->close_state) != UBR_CLOSE_FIRST) { + return ReleaseTrx(); + } + // Incomplete upgrades do not participate in the peer-close exchange. + if (ATOMIC_LOAD(_trx->initializing) || _trx->remote_shm.addr == nullptr) { + return UbrTrxAbort(); + } RETURN_CODE close_check_rc = UbrTrxCloseCheck(_trx); if (UNLIKELY(close_check_rc != UBRING_OK)) { if (close_check_rc == UBRING_REENTRY) { @@ -89,41 +149,24 @@ RETURN_CODE UBRing::UbrTrxClose() { UbrSetSleepTask(UBR_TASK_CLOSE); if (HasTimedOut(start_time, disconnect_timeout_s) != UBRING_OK) { LOG(WARNING) << "Local shm " << _trx->local_shm.name - << " wait for the peer to close timed out, force cleanup."; + << " wait for the peer to close timed out, schedule cleanup."; _trx->ubr_rx.trx_state = UBR_STATE_CLOSED; - // Force synchronous cleanup instead of relying on async timer - DeleteTimerSafe((uint32_t)_trx->timer_fd); - DeleteTimerSafe((uint32_t)_trx->hb_timer_fd); - if (_trx->ubr_tx.remote_rx_event_q.addr != nullptr) { - ((UbrEventQMsg *)_trx->ubr_tx.remote_rx_event_q.addr)->flag = UBR_STATE_CLOSED; - } - if (UNLIKELY(UbrTrxFreeShm(_trx) != UBRING_OK)) { - LOG(WARNING) << "Force close, local shm " << _trx->local_shm.name << " free failed."; - } - if (UNLIKELY(UBRingManager::ReleaseUbrTrxFromMgr(_trx) != UBRING_OK)) { - LOG(WARNING) << "Force close, release trx " << _trx->local_shm.name << " failed."; - } + ClearTrxResource(_trx, UBR_SEND_CLOSE); + ReleaseTrx(); return UBRING_ERR_TIMEOUT; } bthread_usleep(1000); // 1ms, yield to other bthreads } _trx->ubr_rx.trx_state = UBR_STATE_CLOSED; - RETURN_CODE rc; - if (UNLIKELY((rc = ClearTrxResource(_trx, start_time, UBR_SEND_CLOSE)) != UBRING_OK)) { - if (rc == UBRING_REENTRY) { - LOG(INFO) << "Trx close, peer is closing, trx local name=" << _trx->local_shm.name; - return UBRING_OK; - } - LOG(ERROR) << "Trx close, clear trx resource failed, trx local name=" << _trx->local_shm.name; - return UBRING_ERR; - } - // Unlink local shm name immediately so process exit does not leave visible leftovers. - RETURN_CODE unlink_rc = ShmFree(&_trx->local_shm); - if (unlink_rc != UBRING_OK && unlink_rc != SHM_ERR_NOT_FOUND && unlink_rc != SHM_ERR_RESOURCE_ATTACHED) { - LOG(WARNING) << "Trx close, unlink local shm failed, trx local name=" << _trx->local_shm.name - << ", rc=" << unlink_rc; - } - return UBRING_OK; + const RETURN_CODE rc = ClearTrxResource(_trx, UBR_SEND_CLOSE); + ReleaseTrx(); + return rc == UBRING_REENTRY ? UBRING_OK : rc; +} + +RETURN_CODE UBRing::ReleaseTrx() { + UbrTrx* trx = _trx; + _trx = nullptr; + return UBRingManager::ReleaseUbrTrxFromMgr(trx); } RETURN_CODE UBRing::UbrAddCloseTimer() { @@ -155,6 +198,7 @@ RETURN_CODE UBRing::UbrAddTimer() { if (UNLIKELY(UbrAddHBTimer() != UBRING_OK)) { DeleteTimerSafe((uint32_t)_trx->timer_fd); + _trx->timer_fd = 0; LOG(ERROR) << "Ubr " << _trx->local_shm.name << " add heartbeat timer failed."; return UBRING_ERR; } @@ -172,32 +216,12 @@ void* UBRing::UbrTrxCloseCallback(void* args) { if (local_rx_event_q->flag != UBR_STATE_CLOSED || local_tx_event_q->flag == UBR_STATE_CLOSED) { return nullptr; } - trx->ubr_rx.trx_state = UBR_STATE_CLOSED; - int fd = (int)trx->local_shm.fd; - do { - if (ATOMIC_LOAD(trx->close_cnt) == 0) { - break; - } - ATOMIC_SUB(trx->close_cnt, 1); - - uint64_t start_time = GetCurNanoSeconds(); - - if (local_tx_event_q->flag == UBR_STATE_CONNECTED || ATOMIC_LOAD(trx->close_cnt) == 1) { - local_tx_event_q->flag = UBR_STATE_CLOSED; - trx->ubr_tx.trx_state = UBR_STATE_CLOSED; - } - UbrEventQMsg* remote_rx_event_q = (UbrEventQMsg *)trx->ubr_tx.remote_rx_event_q.addr; - if (remote_rx_event_q == nullptr) { - LOG(ERROR) << "Trx close callback failed, " << trx->local_shm.name << " remote_rx_event_q is NULL."; - break; - } + UbrEventQMsg* remote_rx_event_q = + (UbrEventQMsg*)trx->ubr_tx.remote_rx_event_q.addr; + if (remote_rx_event_q != nullptr) { remote_rx_event_q->flag = UBR_STATE_CLOSED; - RETURN_CODE clear_rc = ClearTrxResource(trx, start_time, UBR_CALL_BACK_CLOSE, 1); - if (UNLIKELY(clear_rc != UBRING_OK && clear_rc != UBRING_REENTRY)) { - LOG(ERROR) << "Trx close callback failed, " << trx->local_shm.name << " clear trx resource failed."; - break; - } - } while (0); + } + ClearTrxResource(trx, UBR_CALL_BACK_CLOSE); return nullptr; } @@ -220,41 +244,9 @@ RETURN_CODE UBRing::UbrAddHBTimer() { return UBRING_OK; } -RETURN_CODE UBRing::UbrPassiveClearTrx(UbrTrx *trx, int fd, PASSIVE_DISC_TYPE type) { - RETURN_CODE passive_close_check_rc = UbrTrxCloseCheck(trx); - if (UNLIKELY(passive_close_check_rc != UBRING_OK)) { - if (passive_close_check_rc == UBRING_REENTRY) { - LOG(INFO) << "Passive close skipped, active close in progress, name=" << trx->local_shm.name; - uint64_t start_time = GetCurNanoSeconds(); - return ClearTrxResource(trx, start_time, UBR_CALL_BACK_CLOSE); - } - return UBRING_ERR; - } - trx->ubr_tx.trx_state = UBR_STATE_CLOSED; - trx->ubr_rx.trx_state = UBR_STATE_CLOSED; - DeleteTimerSafe((uint32_t)trx->timer_fd); - const char *type_name = nullptr; - if (type == UBR_HEARTBEAT) { - DeleteTimer((uint32_t)trx->hb_timer_fd); - type_name = "Trx heartbeat"; - } else if (type == UBR_UB_EVENT) { - DeleteTimerSafe((uint32_t)trx->hb_timer_fd); - type_name = "Ub event callback"; - } - constexpr int64_t kMicrosecondsPerSecond = 1000000LL; - bthread_usleep(FLAGS_ub_flying_io_timeout_s * kMicrosecondsPerSecond); - - int rc = ShmLocalFree(&trx->remote_shm); - if (rc != UBRING_OK) { - LOG(ERROR) << type_name << ", delete remote shm failed. ret=" << rc; - } - rc = ShmLocalFree(&trx->local_shm); - if (rc != UBRING_OK) { - LOG(ERROR) << type_name << ", delete local shm failed. ret=" << rc; - } - - UBRingManager::ReleaseUbrTrxFromMgr(trx); - return UBRING_OK; +RETURN_CODE UBRing::UbrPassiveClearTrx(UbrTrx* trx) { + const RETURN_CODE rc = ClearTrxResource(trx, UBR_PASSIVE_CLOSE); + return rc == UBRING_REENTRY ? UBRING_OK : rc; } void* UBRing::UbrTrxHBCallback(void* args) { @@ -289,7 +281,7 @@ void* UBRing::UbrTrxHBCallback(void* args) { int fd = (int)trx->local_shm.fd; LOG(INFO) << "Hlc heartbeat, start to clear trx resource. hb_timer_fd=" << fd << ", shm_name=" << trx->local_shm.name; - UbrPassiveClearTrx(trx, fd, UBR_HEARTBEAT); + UbrPassiveClearTrx(trx); LOG(INFO) << "Hlc heartbeat clear trx resource finish."; return nullptr; } @@ -311,7 +303,8 @@ RETURN_CODE UBRing::UbrAddAsynClearTimer(UbrTrx *trx) { int timer_fd = TimerStart(&time_spec, UbrAsynClearCallback, (void*)trx); if (UNLIKELY(timer_fd == -1)) { - LOG(ERROR) << "Start ubr close timer failed, trx name=%s.", trx->local_shm.name; + LOG(ERROR) << "Start ubr close timer failed, trx name=" + << trx->local_shm.name; return UBRING_ERR; } trx->clear_timer_fd = timer_fd; @@ -326,8 +319,15 @@ void *UBRing::UbrAsynClearCallback(void *args) return nullptr; } - if (UNLIKELY(UbrTrxFreeShm(trx) != UBRING_OK)) { - LOG(ERROR) << "Trx close, wait for local shm " << trx->local_shm.name << " free fail."; + while (UbrTrxFreeShm(trx) != UBRING_OK) { + LOG_EVERY_SECOND(ERROR) << "Trx close, retry cleanup for " + << trx->local_shm.name; + trx->clear_timer_fd = 0; + if (UbrAddAsynClearTimer(trx) == UBRING_OK) { + return nullptr; + } + bthread_usleep(static_cast(FLAGS_ub_flying_io_timeout_s) * + SEC_TO_NSEC / USEC_TO_NSEC); } if (UNLIKELY(UBRingManager::ReleaseUbrTrxFromMgr(trx) != UBRING_OK)) { @@ -848,46 +848,53 @@ RETURN_CODE UBRing::UbrTrxFreeShm(UbrTrx *trx) return UBRING_ERR; } - RETURN_CODE rc = UBRING_OK; - rc = ShmMunmap(&trx->local_shm); - if (UNLIKELY(rc != UBRING_OK)) { - LOG(ERROR) << "Trx close, local unmap " << trx->local_shm.name << " shm fail."; - return UBRING_ERR; - } - - rc = ShmFree(&trx->local_shm); - if (UNLIKELY(rc != UBRING_OK)) { - if (rc != SHM_ERR_RESOURCE_ATTACHED && rc != SHM_ERR_NOT_FOUND) { - LOG(ERROR) << "Wait for " << trx->local_shm.name << " local shm free fail."; - return UBRING_ERR; + if (!trx->local_shm_released) { + int state = ATOMIC_LOAD(trx->local_name_state); + if (state != UBR_SHM_NAME_RELEASING && + ATOMIC_COMPARE_EXCHANGE_STRONG(trx->local_name_state, state, + UBR_SHM_NAME_RELEASING)) { + const RETURN_CODE rc = trx->local_shm.len == 0 ? UBRING_OK + : state == UBR_SHM_NAME_RELEASED + ? ShmRemoteFree(&trx->local_shm) + : ShmLocalFree(&trx->local_shm); + trx->local_shm_released = rc == UBRING_OK || + rc == SHM_ERR_NOT_FOUND || rc == SHM_ERR_RESOURCE_ATTACHED; + ATOMIC_STORE(trx->local_name_state, trx->local_shm_released + ? UBR_SHM_NAME_RELEASED : state); } - LOG(INFO) << "Local shm " << trx->local_shm.name << " already freed, continue to free remote shm."; } - - RETURN_CODE remote_rc = UBRING_OK; - if (trx->remote_shm.addr != nullptr) { - remote_rc = ShmRemoteFree(&trx->remote_shm); - } - if (remote_rc != UBRING_OK) { - LOG(WARNING) << "Free remote shm " << trx->remote_shm.name << " failed, rc=" << remote_rc; + if (!trx->remote_shm_released) { + // The peer owns its name. Releasing a mapping never deletes it. + const RETURN_CODE rc = trx->remote_shm.len == 0 + ? UBRING_OK : ShmRemoteFree(&trx->remote_shm); + trx->remote_shm_released = rc == UBRING_OK || rc == SHM_ERR_NOT_FOUND; } - - return UBRING_OK; + return trx->local_shm_released && trx->remote_shm_released + ? UBRING_OK : UBRING_ERR; } -RETURN_CODE UBRing::UbrUnlinkLocalShm() -{ - if (UNLIKELY(_trx == nullptr)) { - return UBRING_ERR; - } - RETURN_CODE rc = ShmFree(&_trx->local_shm); - if (rc != UBRING_OK && rc != SHM_ERR_NOT_FOUND && rc != SHM_ERR_RESOURCE_ATTACHED) { - LOG(WARNING) << "Unlink local shm " << _trx->local_shm.name << " failed, rc=" << rc; +RETURN_CODE UBRing::UnlinkLocalShm(UbrTrx* trx) { + int expected = UBR_SHM_NAME_OWNED; + if (!ATOMIC_COMPARE_EXCHANGE_STRONG(trx->local_name_state, expected, + UBR_SHM_NAME_RELEASING)) { + return expected == UBR_SHM_NAME_RELEASED ? UBRING_OK : UBRING_RETRY; + } + SHM local_shm = trx->local_shm; + const RETURN_CODE rc = ShmFree(&local_shm); + const bool released = rc == UBRING_OK || rc == SHM_ERR_NOT_FOUND; + ATOMIC_STORE(trx->local_name_state, + released ? UBR_SHM_NAME_RELEASED : UBR_SHM_NAME_OWNED); + if (!released && rc != SHM_ERR_RESOURCE_ATTACHED) { + LOG(WARNING) << "Unlink local shm " << local_shm.name << " failed, rc=" << rc; return rc; } return UBRING_OK; } +RETURN_CODE UBRing::UbrUnlinkLocalShm() { + return _trx != nullptr ? UnlinkLocalShm(_trx) : UBRING_ERR; +} + void UBRing::PreWriteAddr(uint8_t *addr, size_t len) { if (addr == nullptr) { @@ -938,7 +945,7 @@ RETURN_CODE UBRing::UbrTrxMapLocalShm(SHM *local_shm) LOG(ERROR) << "Trx map Shared memory failed, local_shm is null or addr is NULL."; return UBRING_ERR; } - _trx->local_shm = *local_shm; + RecordShm(&_trx->local_shm, *local_shm); _trx->ubr_tx.local_tx_event_q.addr = local_shm->addr + TX_EVENTQ_ADDR_OFFSET; _trx->ubr_tx.local_tx_event_q.len = UBR_EVENTQ_LEN; _trx->ubr_rx.local_rx_event_q.addr = local_shm->addr + RX_EVENTQ_ADDR_OFFSET; @@ -961,7 +968,7 @@ RETURN_CODE UBRing::UbrTrxMapRemoteShm(SHM *remote_shm) LOG(ERROR) << "Trx map Shared memory failed, remote_shm is null or addr is NULL."; return UBRING_ERR; } - _trx->remote_shm = *remote_shm; + RecordShm(&_trx->remote_shm, *remote_shm); _trx->ubr_rx.remote_tx_event_q.addr = remote_shm->addr + TX_EVENTQ_ADDR_OFFSET; _trx->ubr_rx.remote_tx_event_q.len = UBR_EVENTQ_LEN; _trx->ubr_tx.remote_rx_event_q.addr = remote_shm->addr + RX_EVENTQ_ADDR_OFFSET; @@ -997,6 +1004,9 @@ RETURN_CODE UBRing::InitializeServerLegacyFormat() RETURN_CODE UBRing::UbrServerTrxInit(SHM* local_shm, SHM* remote_shm, UbrDataFormat format) { + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } RETURN_CODE rc = UbrTrxMapShm(local_shm, remote_shm); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx map shared memory failed."; @@ -1019,6 +1029,10 @@ RETURN_CODE UBRing::UbrServerTrxInit(SHM* local_shm, SHM* remote_shm, return UBRING_ERR; } + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } ((UbrDataStatusQMsg *)(_trx->ubr_tx.local_data_status_q.addr))->timeout = FLAGS_ub_connect_timeout_s; ((UbrDataStatusQMsg *)(_trx->ubr_rx.remote_data_status_q.addr))->timeout = @@ -1029,6 +1043,7 @@ RETURN_CODE UBRing::UbrServerTrxInit(SHM* local_shm, SHM* remote_shm, _trx->ubr_tx.trx_state = UBR_STATE_CONNECTED; _trx->ubr_rx.trx_state = UBR_STATE_CONNECTED; _trx->data_format = format; + ATOMIC_STORE(_trx->initializing, false); return UBRING_OK; } @@ -1080,32 +1095,32 @@ int UBRing::UbrPrepareIpcV2Format() int UBRing::UbrAllocateServerShm(SHM* remote_trx_shm, SHM* local_trx_shm, UbrDataFormat format) { + if (UBRingManager::AcquireUbrTrxFromMgr(&_trx) != UBRING_OK) { + return -1; + } + _trx->type = TCP_TRX; + if (RegisterShmName(&_trx->remote_shm, remote_trx_shm->name) != UBRING_OK || + RegisterShmName(&_trx->local_shm, local_trx_shm->name) != UBRING_OK) { + UbrTrxAbort(); + return -1; + } UbrSetSleepTask(UBR_TASK_ACCEPT_MAP_FRONT); - if (UNLIKELY((ShmRemoteMalloc(remote_trx_shm)) != UBRING_OK)) { + if (ShmRemoteMalloc(remote_trx_shm) != UBRING_OK) { LOG(ERROR) << "Trx apply remote shared memory failed."; + UbrTrxAbort(); return -1; } - - if (UNLIKELY((ShmLocalCalloc(local_trx_shm)) != UBRING_OK)) { + RecordShm(&_trx->remote_shm, *remote_trx_shm); + if (ATOMIC_LOAD(_trx->device_fault) || + ShmLocalCalloc(local_trx_shm) != UBRING_OK) { LOG(ERROR) << "Trx apply local shared memory failed."; - ShmRemoteFree(remote_trx_shm); + UbrTrxAbort(); return -1; } - - UbrTrx **ubr_trx_ptr = &_trx; - if (UNLIKELY((UBRingManager::AcquireUbrTrxFromMgr(ubr_trx_ptr)) != UBRING_OK)) { - LOG(ERROR) << "Acquire ubrtrx failed."; - ShmRemoteFree(remote_trx_shm); - ShmLocalFree(local_trx_shm); - return -1; - } - _trx->type = TCP_TRX; - if (UNLIKELY((UbrServerTrxInit( - local_trx_shm, remote_trx_shm, format)) != UBRING_OK)) { + RecordShm(&_trx->local_shm, *local_trx_shm); + if (UbrServerTrxInit(local_trx_shm, remote_trx_shm, format) != UBRING_OK) { LOG(ERROR) << "Server trx init failed."; - UbrTrxFreeShm(_trx); - UBRingManager::ReleaseUbrTrxFromMgr(_trx); - _trx = nullptr; + UbrTrxAbort(); return -1; } return 0; @@ -1122,7 +1137,7 @@ int UBRing::UbrAllocateLocalShm(SHM *local_trx_shm, const char *shm_name) RETURN_CODE rc = ApplyAndMapLocalShm(local_trx_shm, shm_name); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx apply or map local shared memory failed, local_name=" << shm_name; - _trx = nullptr; + UbrTrxAbort(); return -1; } return 0; @@ -1167,14 +1182,20 @@ void UBRing::InitializeClientRemoteLegacyFormat() (uint32_t)(_trx->ubr_tx.remote_data_q.len / UBR_MSG_LEN); } -void UBRing::ActivateClientFormat() +RETURN_CODE UBRing::ActivateClientFormat() { + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } // Only publish the connected states here. UbrMapRemoteShmAddTimer has // already started the timers; preserve that legacy ordering. ((UbrEventQMsg *)_trx->ubr_rx.remote_tx_event_q.addr)->flag = UBR_STATE_CONNECTED; ((UbrEventQMsg *)_trx->ubr_rx.local_rx_event_q.addr)->flag = UBR_STATE_CONNECTED; _trx->ubr_tx.trx_state = UBR_STATE_CONNECTED; _trx->ubr_rx.trx_state = UBR_STATE_CONNECTED; + ATOMIC_STORE(_trx->initializing, false); + return UBRING_OK; } RETURN_CODE UBRing::UbrMapRemoteShmAddTimer(SHM *local_trx_shm, const char *local_name, @@ -1209,6 +1230,9 @@ RETURN_CODE UBRing::UbrMapRemoteShmAddTimer(SHM *local_trx_shm, const char *loca LOG(ERROR) << "Connect Trx map shared memory failed, remote shm=" << remote_trx_shm.name; return rc; } + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } if (format == UBR_DATA_FORMAT_LEGACY_64) { InitializeClientRemoteLegacyFormat(); rc = InitializeClientLocalLegacyFormat(); @@ -1228,7 +1252,6 @@ RETURN_CODE UBRing::UbrMapRemoteShmAddTimer(SHM *local_trx_shm, const char *loca if (UNLIKELY(UbrAddTimer() != UBRING_OK)) { LOG(ERROR) << "Ubr add timer failed, local_name=" << local_name; - ShmRemoteFree(&_trx->remote_shm); return UBRING_ERR; } @@ -1239,13 +1262,13 @@ RETURN_CODE UBRing::UbrMapRemoteShmAddTimer(SHM *local_trx_shm, const char *loca LOG(ERROR) << "Local shm " << local_trx_shm->name << " wait for connect remote map timeout."; DeleteTimerSafe((uint32_t)_trx->hb_timer_fd); DeleteTimerSafe((uint32_t)_trx->timer_fd); - ShmRemoteFree(&_trx->remote_shm); + _trx->hb_timer_fd = 0; + _trx->timer_fd = 0; return UBRING_ERR_TIMEOUT; } - ActivateClientFormat(); _trx->data_format = format; - return UBRING_OK; + return ActivateClientFormat(); } RETURN_CODE UBRing::ApplyAndMapLocalShm(SHM *local_trx_shm, const char *local_name) @@ -1265,20 +1288,24 @@ RETURN_CODE UBRing::ApplyAndMapLocalShm(SHM *local_trx_shm, const char *local_na return UBRING_ERR; } + if (RegisterShmName(&_trx->local_shm, local_trx_shm->name) != UBRING_OK) { + return UBRING_ERR; + } RETURN_CODE rc = ShmLocalCalloc(local_trx_shm); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx apply local shared memory failed, local shm name=" << local_trx_shm->name << ", rc=" << rc; if (rc == SHM_ERR_EXIST || rc == SHM_ERR_NOT_FOUND) { rc = UBR_ERR_ADDR_IN_USE; } - UBRingManager::ReleaseUbrTrxFromMgr(_trx); return rc; } + RecordShm(&_trx->local_shm, *local_trx_shm); + if (ATOMIC_LOAD(_trx->device_fault)) { + return UBRING_ERR; + } rc = UbrTrxMapLocalShm(local_trx_shm); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx map local shared memory failed, local shm name=" << local_trx_shm->name; - ShmLocalFree(local_trx_shm); - UBRingManager::ReleaseUbrTrxFromMgr(_trx); return rc; } ((UbrDataStatusQMsg *)_trx->ubr_tx.local_data_status_q.addr)->timeout = @@ -1288,15 +1315,18 @@ RETURN_CODE UBRing::ApplyAndMapLocalShm(SHM *local_trx_shm, const char *local_na RETURN_CODE UBRing::ApplyAndMapRemoteShm(SHM *remote_trx_shm) { + if (RegisterShmName(&_trx->remote_shm, remote_trx_shm->name) != UBRING_OK) { + return UBRING_ERR; + } RETURN_CODE rc = ShmRemoteMalloc(remote_trx_shm); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx apply remote shared memory failed."; return rc; } + RecordShm(&_trx->remote_shm, *remote_trx_shm); rc = UbrTrxMapRemoteShm(remote_trx_shm); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Trx map shared memory failed."; - ShmRemoteFree(remote_trx_shm); return rc; } return UBRING_OK; @@ -1322,44 +1352,72 @@ RETURN_CODE UBRing::WritevHasEnoughSpace(size_t buf_len) return UBRING_OK; } -RETURN_CODE UBRing::UbrClearResourceCheck(UbrTrx *trx, uint64_t start_time, UbrCloseType close_type) +RETURN_CODE UBRing::UbrClearResourceCheck(UbrTrx* trx, UbrCloseType close_type) { if (UNLIKELY(trx == nullptr)) { LOG(ERROR) << "Trx close failed, trx is null."; return UBRING_ERR; } + { + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + if (ATOMIC_LOAD(trx->initializing)) { + ATOMIC_STORE(trx->device_fault, true); + return UBRING_REENTRY; + } + } + // Pin before publishing cleanup ownership to a closing ring. + ATOMIC_ADD(trx->ref_count, 1); + int expected = UBR_CLOSE_FIRST; + if (!ATOMIC_COMPARE_EXCHANGE_STRONG( + trx->close_state, expected, UBR_CLOSE_SECOND)) { + UBRingManager::ReleaseUbrTrxFromMgr(trx); + return UBRING_REENTRY; + } + UbrEventQMsg* local_tx_event_q = (UbrEventQMsg *)trx->ubr_tx.local_tx_event_q.addr; if (local_tx_event_q->flag == UBR_STATE_CONNECTED) { local_tx_event_q->flag = UBR_STATE_CLOSING; } - if (close_type == UBR_SEND_CLOSE) { - DeleteTimerSafe((uint32_t)trx->timer_fd); - } else { - DeleteTimer((uint32_t)trx->timer_fd); - } + DeleteTimerSafe((uint32_t)trx->timer_fd); DeleteTimerSafe((uint32_t)trx->hb_timer_fd); + trx->timer_fd = 0; + trx->hb_timer_fd = 0; if (local_tx_event_q->flag == UBR_STATE_CLOSING) { local_tx_event_q->flag = UBR_STATE_CLOSED; - trx->ubr_tx.trx_state = UBR_STATE_CLOSED; } + if (trx->ubr_rx.local_rx_event_q.addr != nullptr) { + ((UbrEventQMsg*)trx->ubr_rx.local_rx_event_q.addr)->flag = UBR_STATE_CLOSED; + } + trx->ubr_tx.trx_state = UBR_STATE_CLOSED; + trx->ubr_rx.trx_state = UBR_STATE_CLOSED; return UBRING_OK; } -RETURN_CODE UBRing::ClearTrxResource(UbrTrx *trx, uint64_t start_time, UbrCloseType close_type, int op) +RETURN_CODE UBRing::ClearTrxResource(UbrTrx* trx, UbrCloseType close_type) { - RETURN_CODE rc = UbrClearResourceCheck(trx, start_time, close_type); + RETURN_CODE rc = UbrClearResourceCheck(trx, close_type); if (rc != UBRING_OK) { return rc; } + if (close_type == UBR_SEND_CLOSE) { + UnlinkLocalShm(trx); + } rc = UbrAddAsynClearTimer(trx); if (rc != UBRING_OK) { - LOG(ERROR) << "Trx close, add " << trx->local_shm.name << " close clear timer failed."; - return UBRING_ERR; + LOG(WARNING) << "Trx close, clean up " << trx->local_shm.name + << " synchronously after timer creation failed."; + const uint64_t deadline = GetCurNanoSeconds() + + static_cast(FLAGS_ub_flying_io_timeout_s) * SEC_TO_NSEC; + for (uint64_t now = GetCurNanoSeconds(); now < deadline; + now = GetCurNanoSeconds()) { + bthread_usleep((deadline - now) / USEC_TO_NSEC + 1); + } + UbrAsynClearCallback(trx); } return UBRING_OK; diff --git a/src/brpc/ubshm/ub_ring.h b/src/brpc/ubshm/ub_ring.h index bd667d1f3f..c19d152846 100644 --- a/src/brpc/ubshm/ub_ring.h +++ b/src/brpc/ubshm/ub_ring.h @@ -44,6 +44,8 @@ class UBRing : public butil::IReader { RETURN_CODE UbrTrxMapShm(SHM *local_shm, SHM *remote_shm); RETURN_CODE UbrTrxClose(); + // Release an upgrade that has not reached the established state. + RETURN_CODE UbrTrxAbort(); RETURN_CODE UbrAddCloseTimer(); @@ -55,7 +57,7 @@ class UBRing : public butil::IReader { static void *UbrTrxHBCallback(void *args); - static RETURN_CODE UbrPassiveClearTrx(UbrTrx *trx, int fd, PASSIVE_DISC_TYPE type); + static RETURN_CODE UbrPassiveClearTrx(UbrTrx* trx); static RETURN_CODE UbrAddAsynClearTimer(UbrTrx *trx); @@ -206,6 +208,9 @@ class UBRing : public butil::IReader { LOG(ERROR) << "Trx close callback failed, trx is null."; return UBRING_ERR; } + if (ATOMIC_LOAD(trx->initializing)) { + return UBRING_REENTRY; + } if (UNLIKELY(trx->local_shm.addr == nullptr)) { LOG(ERROR) << "Trx close failed, local_shm addr is NULL."; return UBRING_ERR; @@ -241,9 +246,13 @@ class UBRing : public butil::IReader { RETURN_CODE InitializeClientLocalLegacyFormat(); void InitializeClientRemoteLegacyFormat(); void InitializeClientLegacyFormat(); - void ActivateClientFormat(); - static RETURN_CODE UbrClearResourceCheck(UbrTrx *trx, uint64_t start_time, UbrCloseType close_type); - static RETURN_CODE ClearTrxResource(UbrTrx *trx, uint64_t start_time, UbrCloseType close_type, int op=0); + RETURN_CODE ActivateClientFormat(); + RETURN_CODE RegisterShmName(SHM* target, const char* name); + void RecordShm(SHM* target, const SHM& source); + static RETURN_CODE UnlinkLocalShm(UbrTrx* trx); + static RETURN_CODE UbrClearResourceCheck(UbrTrx* trx, UbrCloseType close_type); + static RETURN_CODE ClearTrxResource(UbrTrx* trx, UbrCloseType close_type); + RETURN_CODE ReleaseTrx(); UbrTrx* _trx{nullptr}; }; diff --git a/src/brpc/ubshm/ub_ring_manager.cpp b/src/brpc/ubshm/ub_ring_manager.cpp index 64f2434eba..c351cf1c8d 100644 --- a/src/brpc/ubshm/ub_ring_manager.cpp +++ b/src/brpc/ubshm/ub_ring_manager.cpp @@ -30,6 +30,7 @@ UbrLinkInfoMgr UBRingManager::g_link_info_mgr; pthread_mutex_t UBRingManager::g_ubr_trx_mgr_mtx = PTHREAD_MUTEX_INITIALIZER; pthread_mutex_t UBRingManager::g_ubr_listener_mgr_mtx = PTHREAD_MUTEX_INITIALIZER; pthread_mutex_t UBRingManager::g_link_info_mgr_mtx = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_ubr_trx_idle = PTHREAD_COND_INITIALIZER; uint64_t g_ubr_trx_num = 0; uint64_t g_ub_event_cnt = 0; @@ -58,6 +59,10 @@ RETURN_CODE UBRingManager::UbrMgrDefault() } RETURN_CODE UBRingManager::UbrMgrInit() { + LOCK_GUARD(g_ubr_trx_mgr_mtx); + if (g_ubr_mgr.trx_mgr != nullptr) { + return UBRING_OK; + } RETURN_CODE rc = UbrMgrDefault(); if (UNLIKELY(rc != UBRING_OK)) { LOG(ERROR) << "Ubr manager set default values failed."; @@ -71,7 +76,9 @@ RETURN_CODE UBRingManager::UbrMgrInit() { if (UNLIKELY(g_ubr_mgr.trx_mgr == nullptr || g_ubr_mgr.trx_mgr_unit_status == nullptr)) { LOG(ERROR) << "Ubr manager memory allocation failed."; - UbrMgrFini(); + FREE_PTR(g_ubr_mgr.trx_mgr); + FREE_PTR(g_ubr_mgr.trx_mgr_unit_status); + g_ubr_mgr.trx_cap = 0; return UBRING_ERR; } @@ -84,14 +91,17 @@ RETURN_CODE UBRingManager::UbrMgrInit() { void UBRingManager::UbrMgrFini() { { LOCK_GUARD(g_ubr_trx_mgr_mtx); + // Connections and cleanup callbacks retain transaction slots. + while (g_ubr_mgr.trx_num != 0) { + pthread_cond_wait(&g_ubr_trx_idle, &g_ubr_trx_mgr_mtx); + } FREE_PTR(g_ubr_mgr.trx_mgr); FREE_PTR(g_ubr_mgr.trx_mgr_unit_status); + g_ubr_mgr.trx_cap = 0; } { LOCK_GUARD(g_ubr_listener_mgr_mtx); } - g_ubr_mgr.trx_num = 0; - g_ubr_mgr.trx_cap = 0; LinkInfoFini(); } @@ -101,12 +111,12 @@ RETURN_CODE UBRingManager::AcquireUbrTrxFromMgr(UbrTrx **trx) { return UBRING_ERR; } + LOCK_GUARD(g_ubr_trx_mgr_mtx); if (UNLIKELY(g_ubr_mgr.trx_mgr == nullptr)) { LOG(ERROR) << "Acquire trx failed, trx_mgr is null."; return UBRING_ERR; } - LOCK_GUARD(g_ubr_trx_mgr_mtx); if (g_ubr_mgr.trx_num >= g_ubr_mgr.trx_cap) { LOG(ERROR) << "Acquire trx failed, trx number is full."; return UBRING_ERR; @@ -119,6 +129,8 @@ RETURN_CODE UBRingManager::AcquireUbrTrxFromMgr(UbrTrx **trx) { *trx = &g_ubr_mgr.trx_mgr[i]; (*trx)->trx_mgr_index = i; (*trx)->ubr_id = g_ubr_trx_num; + (*trx)->ref_count = 1; + (*trx)->initializing = true; (*trx)->close_state = UBR_CLOSE_FIRST; (*trx)->close_cnt = MAX_CLOSE_COUNT; ++g_ubr_mgr.trx_num; @@ -136,17 +148,12 @@ RETURN_CODE UBRingManager::ReleaseUbrTrxFromMgr(UbrTrx *trx) { return UBRING_ERR; } - trx->local_shm.addr = nullptr; - trx->ubr_tx.local_tx_event_q.addr = nullptr; - trx->ubr_tx.local_data_status_q.addr = nullptr; - trx->ubr_rx.local_rx_event_q.addr = nullptr; - trx->ubr_rx.remote_data_status_q.addr = nullptr; + LOCK_GUARD(g_ubr_trx_mgr_mtx); if (UNLIKELY(g_ubr_mgr.trx_mgr == nullptr)) { LOG(ERROR) << "Release trx failed, trx_mgr is null."; return UBRING_ERR; } - LOCK_GUARD(g_ubr_trx_mgr_mtx); uint32_t idx = trx->trx_mgr_index; if (g_ubr_mgr.trx_mgr_unit_status[idx] == UBR_MGR_UNIT_FREE) { LOG(INFO) << "Release trx already freed, name=" << trx->local_shm.name; @@ -158,8 +165,19 @@ RETURN_CODE UBRingManager::ReleaseUbrTrxFromMgr(UbrTrx *trx) { return UBRING_ERR; } + if (ATOMIC_SUB(trx->ref_count, 1) != 1) { + return UBRING_OK; + } + trx->local_shm.addr = nullptr; + trx->ubr_tx.local_tx_event_q.addr = nullptr; + trx->ubr_tx.local_data_status_q.addr = nullptr; + trx->ubr_rx.local_rx_event_q.addr = nullptr; + trx->ubr_rx.remote_data_status_q.addr = nullptr; g_ubr_mgr.trx_mgr_unit_status[idx] = UBR_MGR_UNIT_FREE; --g_ubr_mgr.trx_num; + if (g_ubr_mgr.trx_num == 0) { + pthread_cond_broadcast(&g_ubr_trx_idle); + } return UBRING_OK; } @@ -239,26 +257,34 @@ int32_t UBRingManager::UbEventCallback(const char *shm_name) LOG(ERROR) << "Ub event callback failed, shm name is null."; return UBRING_ERR; } - if (UNLIKELY(g_ubr_mgr.trx_mgr == nullptr)) { - LOG(ERROR) << "Ub event callback failed, trx mgr is null."; - return UBRING_ERR; - } LOG(INFO) << "Ub event callback is processing. shm_name=" << shm_name; - - for (uint32_t i = 0; i < g_ubr_mgr.trx_cap; ++i) { - if (g_ubr_mgr.trx_mgr_unit_status[i] == UBR_MGR_UNIT_FREE) { - continue; - } - - if (strcmp(g_ubr_mgr.trx_mgr[i].local_shm.name, shm_name) == 0 || // the failed link is this trx's local shm - strcmp(g_ubr_mgr.trx_mgr[i].remote_shm.name, shm_name) == 0) { // the failed link is this trx's remote shm - ++g_ub_event_cnt; - int fd = (int)g_ubr_mgr.trx_mgr[i].local_shm.fd; - LOG(WARNING) << "Ub event callback, the fd of the faulty link is " << fd; - return UBRing::UbrPassiveClearTrx(&g_ubr_mgr.trx_mgr[i], fd, UBR_UB_EVENT); + UbrTrx* trx = nullptr; + { + LOCK_GUARD(g_ubr_trx_mgr_mtx); + for (uint32_t i = 0; i < g_ubr_mgr.trx_cap; ++i) { + if (g_ubr_mgr.trx_mgr_unit_status[i] == UBR_MGR_UNIT_FREE) { + continue; + } + UbrTrx* candidate = &g_ubr_mgr.trx_mgr[i]; + if (strcmp(candidate->local_shm.name, shm_name) == 0 || + strcmp(candidate->remote_shm.name, shm_name) == 0) { + ATOMIC_STORE(candidate->device_fault, true); + if (ATOMIC_LOAD(candidate->initializing)) { + return UBRING_OK; + } + trx = candidate; + ATOMIC_ADD(trx->ref_count, 1); + ++g_ub_event_cnt; + break; + } } } - return UBRING_ERR; + if (trx == nullptr) { + return UBRING_ERR; + } + const RETURN_CODE rc = UBRing::UbrPassiveClearTrx(trx); + ReleaseUbrTrxFromMgr(trx); + return rc; } } } diff --git a/src/brpc/ubshm/ub_ring_manager.h b/src/brpc/ubshm/ub_ring_manager.h index e2afa324c7..b08bba8e3d 100644 --- a/src/brpc/ubshm/ub_ring_manager.h +++ b/src/brpc/ubshm/ub_ring_manager.h @@ -48,6 +48,7 @@ typedef struct TagUbrLinkInfoMgr { } UbrLinkInfoMgr; class UBRingManager { +friend class UBRing; public: ~UBRingManager(){ UbrMgrFini(); @@ -59,10 +60,12 @@ class UBRingManager { static RETURN_CODE UbrMgrInit(); + // Wait for acquired transactions to finish cleanup before freeing the pool. static void UbrMgrFini(); static RETURN_CODE AcquireUbrTrxFromMgr(UbrTrx **trx); + // Drop a reference and return the slot after the last owner releases it. static RETURN_CODE ReleaseUbrTrxFromMgr(UbrTrx *trx); static void LinkInfoInit(void); @@ -84,4 +87,4 @@ class UBRingManager { } } -#endif //BRPC_UB_RING_MANAGER_H \ No newline at end of file +#endif //BRPC_UB_RING_MANAGER_H diff --git a/src/brpc/ubshm/ubr_trx.h b/src/brpc/ubshm/ubr_trx.h index 4111b8ee87..712774f1ba 100644 --- a/src/brpc/ubshm/ubr_trx.h +++ b/src/brpc/ubshm/ubr_trx.h @@ -60,7 +60,8 @@ typedef enum { typedef enum { UBR_SEND_CLOSE, - UBR_CALL_BACK_CLOSE + UBR_CALL_BACK_CLOSE, + UBR_PASSIVE_CLOSE } UbrCloseType; typedef enum { @@ -129,6 +130,12 @@ typedef struct TagUbrRx { volatile EventQState trx_state; } UbrRx; +enum UbrShmNameState { + UBR_SHM_NAME_OWNED, + UBR_SHM_NAME_RELEASING, + UBR_SHM_NAME_RELEASED // Unlinked or owned by the SHM retry queue. +}; + typedef struct TagUbrTrx { UbrTx ubr_tx; UbrRx ubr_rx; @@ -141,6 +148,14 @@ typedef struct TagUbrTrx { int timer_fd; int hb_timer_fd; int clear_timer_fd; + // The ring, cleanup callback and device-event callback hold references. + AtomicInt ref_count; + // SDK events defer cleanup until queue initialization has completed. + AtomicBool initializing; + AtomicBool device_fault; + AtomicInt local_name_state; + bool local_shm_released; + bool remote_shm_released; AtomicInt close_cnt; AtomicInt close_state; } UbrTrx; @@ -155,11 +170,6 @@ typedef struct TagUbrLinkLock { FileLock* file_lock; } UbrLinkLock; -typedef enum { - UBR_UB_EVENT, - UBR_HEARTBEAT, -}PASSIVE_DISC_TYPE; - } } #endif //BRPC_UBR_TRX_H diff --git a/src/brpc/ubshm_transport.cpp b/src/brpc/ubshm_transport.cpp index 45f6a61fa6..b96515694e 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,20 @@ DECLARE_bool(usercode_in_pthread); extern SocketVarsCollector *g_vars; -void UBShmTransport::Init(Socket *socket, const SocketOptions &options) { +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; - } - _tcp_transport = std::make_shared(); - _tcp_transport->Init(socket, options); + _on_edge_trigger = nullptr; + _ub_state = UB_UNKNOWN; + _ub_ep = new ubring::UBShmEndpoint(socket); } void UBShmTransport::Release() { @@ -64,26 +66,59 @@ 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() { + const bool abort_upgrade = !UpgradeActive(); + SetHighSpeedAvailable(false); + if (_ub_ep != nullptr) { + _ub_ep->Reset(abort_upgrade); } } -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; } @@ -139,9 +171,6 @@ void UBShmTransport::ProcessEvent(bthread_attr_t attr) { void UBShmTransport::QueueMessage(InputMessageClosure& input_msg, int* num_bthread_created, bool last_msg) { - if (last_msg) { - return; - } InputMessageBase* to_run_msg = input_msg.release(); if (!to_run_msg) { return; @@ -156,15 +185,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 +209,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 +233,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/test/brpc_rdma_unittest.cpp b/test/brpc_rdma_unittest.cpp index 59bd9da55b..c61a0c24ca 100644 --- a/test/brpc_rdma_unittest.cpp +++ b/test/brpc_rdma_unittest.cpp @@ -15,41 +15,39 @@ // 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 // std::min - -#include "brpc/rdma/rdma_handshake.h" -#include "brpc/rdma/rdma_handshake_constants.h" -#include "brpc/rdma/rdma_handshake.pb.h" +#include #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; @@ -64,10 +62,11 @@ 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); +DECLARE_int32(rdma_prepared_qp_cnt); DECLARE_int32(rdma_client_handshake_version); DECLARE_bool(rdma_ece); @@ -126,7 +125,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) { @@ -143,7 +142,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) { @@ -180,8 +179,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); @@ -220,11 +218,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); } @@ -241,6 +240,7 @@ class RdmaTest : public ::testing::Test { void StopServer() { _server.Stop(0); + _wait_socket.reset(); _server.Join(); } @@ -257,19 +257,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; @@ -277,41 +287,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; @@ -332,8 +324,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(); @@ -349,6 +340,21 @@ class RdmaRpcTest : public RdmaTest, int _saved_handshake_version = 2; }; +TEST_F(RdmaTest, channels_share_connections_only_with_the_same_transport) { + for (const char* group : {"", "transport-signature"}) { + Channel channels[4]; + for (int i = 0; i < 4; ++i) { + ChannelOptions options; + options.socket_mode = i < 2 ? SOCKET_MODE_TCP : SOCKET_MODE_RDMA; + options.connection_group = group; + ASSERT_EQ(0, channels[i].Init(g_ep, &options)); + } + EXPECT_EQ(channels[0]._server_id, channels[1]._server_id); + EXPECT_EQ(channels[2]._server_id, channels[3]._server_id); + EXPECT_NE(channels[0]._server_id, channels[2]._server_id); + } +} + TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { SocketOptions main_options; main_options.socket_mode = SOCKET_MODE_RDMA; @@ -357,12 +363,12 @@ 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; SocketId stale_cq_sid; SocketId current_cq_sid; ASSERT_EQ(0, Socket::Create(cq_options, &stale_cq_sid)); @@ -372,25 +378,53 @@ TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { SocketUniquePtr current_cq_socket; ASSERT_EQ(0, Socket::Address(stale_cq_sid, &stale_cq_socket)); ASSERT_EQ(0, Socket::Address(current_cq_sid, ¤t_cq_socket)); + stale_cq_socket->_user = new rdma::RdmaEndpoint::CqSocketUser(main_sid); + current_cq_socket->_user = new rdma::RdmaEndpoint::CqSocketUser(main_sid); ep->_cq_sid = current_cq_sid; rdma::RdmaEndpoint::PollCq(stale_cq_socket.get()); - stale_cq_socket->_user = NULL; - current_cq_socket->_user = NULL; stale_cq_socket->SetFailed(); current_cq_socket->SetFailed(); main_socket->SetFailed(); } +TEST_F(RdmaTest, cq_callback_outlives_main_socket) { + SocketOptions options; + options.socket_mode = SOCKET_MODE_RDMA; + SocketId main_sid; + ASSERT_EQ(0, Socket::Create(options, &main_sid)); + SocketUniquePtr main_socket; + ASSERT_EQ(0, Socket::Address(main_sid, &main_socket)); + auto* ep = RdmaTransportOf(main_socket.get())->_rdma_ep; + + SocketId cq_sid; + ASSERT_EQ(0, Socket::Create(SocketOptions(), &cq_sid)); + SocketUniquePtr cq_socket; + ASSERT_EQ(0, Socket::Address(cq_sid, &cq_socket)); + cq_socket->_user = new rdma::RdmaEndpoint::CqSocketUser(main_sid); + ep->_cq_sid = cq_sid; + ep->_resource = new rdma::RdmaResource; + const int saved = rdma::FLAGS_rdma_prepared_qp_cnt; + rdma::FLAGS_rdma_prepared_qp_cnt = 0; + main_socket->SetFailed(); + main_socket.reset(); + rdma::FLAGS_rdma_prepared_qp_cnt = saved; + + // A dispatched callback holds the CQ socket after main connection teardown. + EXPECT_NE(nullptr, cq_socket->user()); + EXPECT_NE(0, Socket::Address(main_sid, &main_socket)); + rdma::RdmaEndpoint::PollCq(cq_socket.get()); +} + TEST_F(RdmaTest, client_close_before_hello_send) { StartServer(); butil::fd_guard sockfd; ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -403,20 +437,12 @@ TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { butil::fd_guard sockfd; ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + 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_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); StopServer(); } @@ -430,35 +456,31 @@ TEST_F(RdmaTest, client_close_during_hello_send) { 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_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); 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); + // 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)); s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); 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_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); 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_NE(nullptr, s); + 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.) @@ -466,7 +488,7 @@ TEST_F(RdmaTest, client_close_during_hello_send) { 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_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); sockfd3.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -482,11 +504,11 @@ TEST_F(RdmaTest, client_hello_msg_invalid_len) { 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_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - 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)); // Write magic string. + ASSERT_HANDSHAKE_PHASE(handshake::HELLO_WAIT, s); memset(data + 4, 0, 36); ASSERT_TRUE(WriteAll(sockfd1, data + 4, 36)); // Write invalid length. ASSERT_TRUE(WaitForServerSocketGone()); @@ -494,11 +516,11 @@ TEST_F(RdmaTest, client_hello_msg_invalid_len) { 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_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - 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)); // 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); @@ -519,11 +541,11 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { 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_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - 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)); // 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 @@ -535,32 +557,32 @@ TEST_F(RdmaTest, client_hello_msg_invalid_version) { // 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_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + 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_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); 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_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); memcpy(data, "RDMA", 4); - 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)); // 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 `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_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -586,15 +608,15 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd1.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -606,15 +628,15 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -626,15 +648,15 @@ TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); ASSERT_TRUE(WriteAll(sockfd3, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd3.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -645,16 +667,26 @@ TEST_F(RdmaTest, client_close_after_qp_build) { StartServer(); 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)); + 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()); @@ -665,21 +697,31 @@ TEST_F(RdmaTest, client_close_during_ack_send) { StartServer(); 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); uint32_t flags = butil::HostToNet32(1); ASSERT_TRUE(WriteAll(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); sockfd1.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -690,37 +732,47 @@ TEST_F(RdmaTest, client_close_after_ack_send) { StartServer(); 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); 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); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); flags = butil::HostToNet32(1); ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); sockfd2.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -731,47 +783,52 @@ TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { StartServer(); 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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); 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_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_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_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_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); 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_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); ASSERT_TRUE(WriteAll(sockfd2, &flags, sizeof(flags))); ASSERT_TRUE(WaitForServerSocketGone()); 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)); @@ -784,9 +841,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(); @@ -795,7 +849,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. @@ -805,14 +859,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()); @@ -820,8 +870,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(); @@ -829,8 +877,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); @@ -845,7 +892,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(); @@ -854,24 +900,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(); @@ -880,11 +998,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()); @@ -901,15 +1019,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_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); 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(); @@ -917,11 +1032,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 @@ -954,13 +1068,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()); @@ -988,8 +1098,10 @@ TEST_F(RdmaTest, server_miss_before_hello_send) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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); @@ -1018,15 +1130,17 @@ TEST_F(RdmaTest, server_close_before_hello_send) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + 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()); @@ -1052,17 +1166,16 @@ TEST_F(RdmaTest, server_miss_during_magic_str) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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)); 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. bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1088,18 +1201,18 @@ TEST_F(RdmaTest, server_close_during_magic_str) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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)); - // 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(ReadAll(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_TRUE(WriteAll(acc_fd, "RD", 2)); acc_fd.reset(-1); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + ASSERT_HANDSHAKE_PHASE(handshake::FAILED, s.get()); bthread_id_join(cntl.call_id()); ASSERT_EQ(EEOF, cntl.ErrorCode()); @@ -1125,15 +1238,17 @@ TEST_F(RdmaTest, server_hello_invalid_magic_str) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + 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()); @@ -1159,15 +1274,19 @@ TEST_F(RdmaTest, server_miss_during_hello_msg) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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_EQ(2, write(acc_fd, "00", 2)); + 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_TRUE(WriteAll(acc_fd, &msg_len, sizeof(msg_len))); bthread_id_join(cntl.call_id()); ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); @@ -1193,17 +1312,21 @@ TEST_F(RdmaTest, server_close_during_hello_msg) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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_EQ(2, write(acc_fd, "00", 2)); - close(acc_fd); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + 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_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()); @@ -1229,19 +1352,21 @@ TEST_F(RdmaTest, server_hello_invalid_msg_len) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FAILED, RdmaTransportOf(s)); + 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()); @@ -1267,20 +1392,22 @@ TEST_F(RdmaTest, server_hello_invalid_version) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - 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()); @@ -1308,13 +1435,15 @@ TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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; @@ -1327,10 +1456,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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - 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()); @@ -1358,13 +1487,15 @@ TEST_F(RdmaTest, server_miss_after_ack) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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; @@ -1377,10 +1508,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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - 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()); @@ -1408,13 +1539,15 @@ TEST_F(RdmaTest, server_close_after_ack) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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; @@ -1427,13 +1560,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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::ESTABLISHED, RdmaTransportOf(s)); - 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()); @@ -1459,13 +1592,15 @@ TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::C_HELLO_WAIT, RdmaTransportOf(s)); + 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; @@ -1478,16 +1613,15 @@ 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)); - 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)); + 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()); } - TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { butil::fd_guard sockfd(butil::tcp_listen(g_ep)); EXPECT_TRUE(sockfd >= 0); @@ -1508,13 +1642,15 @@ TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); 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)); @@ -1532,7 +1668,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()); } @@ -1540,18 +1676,13 @@ 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); - // 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; @@ -1565,12 +1696,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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + 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, @@ -1584,13 +1715,13 @@ 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)); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -1598,19 +1729,15 @@ TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { 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; - 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_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + ASSERT_NE(nullptr, s); + ASSERT_TRUE(s != NULL); + ASSERT_HANDSHAKE_PHASE(handshake::UNINITIALIZED, s); // Build a v2 hello with msg_len = 48 (40 base + 8B zero tail). rdma::v2_wire::HelloMessage msg{}; @@ -1623,20 +1750,17 @@ 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_TRUE(WriteAll(sockfd, buf, sizeof(buf))); + + ASSERT_HANDSHAKE_PHASE(handshake::ESTABLISHED, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -1647,16 +1771,11 @@ TEST_F(RdmaTest, v2_server_drains_tail_then_reads_ack) { 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + 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 @@ -1674,7 +1793,7 @@ 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(WriteAll(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); ASSERT_TRUE(WaitForServerSocketGone()); sockfd.reset(-1); @@ -1694,6 +1813,7 @@ class HandshakeVersionFlag { ~HandshakeVersionFlag() { rdma::FLAGS_rdma_client_handshake_version = _saved; } + private: int _saved; }; @@ -1705,8 +1825,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; @@ -1725,13 +1844,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); @@ -1758,12 +1876,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); @@ -1771,7 +1889,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)); @@ -1793,37 +1911,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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + 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())); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + 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()); @@ -1835,9 +1947,8 @@ 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_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -1848,37 +1959,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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); + 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)); + 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); + ASSERT_NE(nullptr, s); uint8_t buf[8]; memcpy(buf, "RDM3", 4); @@ -1886,25 +1988,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)); + 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); + ASSERT_NE(nullptr, s); // "RDM3" + pb_size = 8 + 8 bytes of 0xff (invalid protobuf body). uint8_t buf[16]; @@ -1912,25 +2010,21 @@ 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)); + 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3Hello(); msg.set_sq_size(0); // invalid: < MIN_QP_SIZE (16) @@ -1938,26 +2032,24 @@ TEST_F(RdmaTest, v3_server_invalid_sq_size_falls_back) { ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); // 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_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)); + 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))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -1971,16 +2063,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,46 +2084,42 @@ rdma::RdmaHello MakeValidV3HelloWithEce(uint32_t vendor_id, // 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)); + 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)); } // 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(); - 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_TRUE(s != nullptr); + 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())); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + + 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))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -2046,28 +2132,21 @@ 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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); - // 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)); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -2076,33 +2155,27 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_when_disabled) { // 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(); - 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_TRUE(s != nullptr); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); std::string packet = MakeV3Packet(msg); ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); - // 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)); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -2115,13 +2188,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); @@ -2141,11 +2218,12 @@ 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); - 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_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()); @@ -2160,16 +2238,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))); + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(rdma::RdmaEndpoint::UNINIT, RdmaTransportOf(s)->_rdma_ep->_state); + 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. @@ -2186,16 +2259,15 @@ 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)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + 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))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); ASSERT_FALSE(s->Failed()); sockfd.reset(-1); @@ -2212,34 +2284,28 @@ rdma::RdmaHello MakeValidV3HelloWithMtu(uint32_t mtu) { } // A client hello carrying MTU must not break the server handshake: -// the server still parses the hello and advances to S_ACK_WAIT. +// the server parses the hello and advances to ACK_WAIT. TEST_F(RdmaTest, v3_server_accepts_client_hello_with_mtu) { 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_TRUE(s != NULL); + ASSERT_NE(nullptr, s); rdma::RdmaHello msg = MakeValidV3HelloWithMtu(IBV_MTU_4096); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); + 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))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -2253,21 +2319,15 @@ TEST_F(RdmaTest, v3_server_accepts_client_hello_with_mtu) { TEST_F(RdmaTest, v3_server_reply_has_negotiated_mtu) { 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_TRUE(s != NULL); + ASSERT_NE(nullptr, s); // Advertise a non-default MTU so we can verify negotiation. rdma::RdmaHello msg = MakeValidV3HelloWithMtu(IBV_MTU_4096); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); @@ -2285,7 +2345,7 @@ TEST_F(RdmaTest, v3_server_reply_has_negotiated_mtu) { EXPECT_EQ(expected_mtu, reply.mtu()); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); @@ -2298,21 +2358,15 @@ TEST_F(RdmaTest, v3_server_reply_has_negotiated_mtu) { TEST_F(RdmaTest, v3_server_reply_has_no_mtu_without_client_mtu) { 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_TRUE(s != NULL); + ASSERT_NE(nullptr, s); // Client hello without MTU field. rdma::RdmaHello msg = MakeValidV3Hello(); std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); rdma::RdmaHello reply; ReadServerV3Reply(sockfd, &reply); @@ -2321,88 +2375,34 @@ TEST_F(RdmaTest, v3_server_reply_has_no_mtu_without_client_mtu) { EXPECT_FALSE(reply.has_mtu()); uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); sockfd.reset(-1); ASSERT_TRUE(WaitForServerSocketGone()); StopServer(); } -// A client hello with an out-of-range MTU (e.g., 0 or > IBV_MTU_4096) must be -// rejected by ValidRdmaHello. Per the handshake-fallback design, the server -// does NOT close the connection; instead it degrades to TCP fallback: the -// endpoint enters S_ACK_WAIT with RDMA_OFF, the reply carries no MTU, and a -// subsequent ACK-without-RDMA finalizes FALLBACK_TCP. TEST_F(RdmaTest, v3_server_rejects_invalid_mtu) { StartServer(); - - sockaddr_in addr; - bzero((char*)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - - // Test with MTU = 0 (below IBV_MTU_256). - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); - Socket* s = WaitForServerSocket(); - ASSERT_TRUE(s != NULL); - - rdma::RdmaHello msg = MakeValidV3HelloWithMtu(0); - std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd1, packet.data(), packet.size())); - - // The server rejects the hello but falls back to TCP instead of failing. - s = GetSocketFromServer(0); - ASSERT_TRUE(s != NULL); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_FALSE(s->Failed()); - - // The fallback reply must not carry an MTU. - rdma::RdmaHello reply; - ReadServerV3Reply(sockfd1, &reply); - EXPECT_FALSE(reply.has_mtu()); - - // Ack without RDMA so the server finalizes FALLBACK_TCP. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - ASSERT_FALSE(s->Failed()); - - sockfd1.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); - - // Test with MTU = 6 (above IBV_MTU_4096). - 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 != NULL); - - msg = MakeValidV3HelloWithMtu(6); - packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd2, packet.data(), packet.size())); - - s = GetSocketFromServer(0); - ASSERT_TRUE(s != NULL); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::S_ACK_WAIT, RdmaTransportOf(s)); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); - ASSERT_FALSE(s->Failed()); - - ReadServerV3Reply(sockfd2, &reply); - EXPECT_FALSE(reply.has_mtu()); - - flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); - ASSERT_FALSE(s->Failed()); - - sockfd2.reset(-1); - ASSERT_TRUE(WaitForServerSocketGone()); - + for (uint32_t mtu : {0u, 6u}) { + butil::fd_guard sockfd; + ASSERT_NO_FATAL_FAILURE(ConnectToServer(&sockfd)); + Socket* s = WaitForServerSocket(); + ASSERT_NE(nullptr, s); + const std::string packet = MakeV3Packet(MakeValidV3HelloWithMtu(mtu)); + ASSERT_TRUE(WriteAll(sockfd, packet.data(), packet.size())); + rdma::RdmaHello reply; + ASSERT_NO_FATAL_FAILURE(ReadServerV3Reply(sockfd, &reply)); + EXPECT_FALSE(reply.has_mtu()); + ASSERT_HANDSHAKE_PHASE(handshake::ACK_WAIT, s); + EXPECT_EQ(RdmaTransport::RDMA_OFF, RdmaTransportOf(s)->_rdma_state); + EXPECT_FALSE(s->Failed()); + const uint32_t flags = butil::HostToNet32(0); + ASSERT_TRUE(WriteAll(sockfd, &flags, sizeof(flags))); + ASSERT_HANDSHAKE_PHASE(handshake::FALLBACK_TCP, s); + sockfd.reset(-1); + ASSERT_TRUE(WaitForServerSocketGone()); + } StopServer(); } @@ -2428,21 +2428,11 @@ TEST_F(RdmaTest, v3_client_hello_includes_mtu) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - butil::fd_guard acc_fd(accept(sockfd, NULL, NULL)); + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); ASSERT_TRUE(acc_fd >= 0); - // Read 4B magic + 4B pb_size + body from the client hello. - uint8_t hdr[8]; - ASSERT_EQ(8, read(acc_fd, hdr, 8)); - ASSERT_EQ(0, memcmp(hdr, "RDM3", 4)); - uint32_t pb_size = butil::NetToHost32(*reinterpret_cast(hdr + 4)); - ASSERT_GT(pb_size, 0u); - ASSERT_LE(pb_size, 4096u); - std::string body(pb_size, '\0'); - ASSERT_EQ((ssize_t)pb_size, read(acc_fd, &body[0], pb_size)); - rdma::RdmaHello hello; - ASSERT_TRUE(hello.ParseFromString(body)); + ASSERT_NO_FATAL_FAILURE(ReadServerV3Reply(acc_fd, &hello)); // The client hello must carry the locally queried active MTU. // In UT mode (g_skip_rdma_init=true), OpenDevice is never called so @@ -2475,8 +2465,10 @@ TEST_F(RdmaTest, try_global_disable_rdma) { google::protobuf::Closure* done = DoNothing(); ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_RDMA_STATE(rdma::RdmaEndpoint::FALLBACK_TCP, RdmaTransportOf(s)); + 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()); @@ -2573,13 +2565,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; @@ -2603,7 +2589,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()); @@ -2629,17 +2615,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; @@ -2651,10 +2626,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(); } @@ -2665,10 +2648,17 @@ 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); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); StopServer(); } @@ -2679,10 +2669,17 @@ 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); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); StopServer(); } @@ -2694,10 +2691,17 @@ 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); + bthread_id_join(cntl.call_id()); + ASSERT_FALSE(cntl.Failed()); StopServer(); } @@ -2714,13 +2718,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; @@ -2768,10 +2771,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) { @@ -2845,8 +2850,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; @@ -2899,7 +2904,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()); @@ -3037,12 +3043,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; } } @@ -3079,10 +3085,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(); @@ -3120,9 +3125,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(); @@ -3153,7 +3158,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 " @@ -3192,10 +3196,11 @@ TEST_P(RdmaRpcTest, verbs_error_handling) { ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); 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)); + 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); @@ -3205,7 +3210,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()); @@ -3228,9 +3233,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)); @@ -3276,7 +3280,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) { @@ -3297,7 +3301,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); @@ -3316,12 +3321,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); @@ -3387,12 +3394,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_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 new file mode 100644 index 0000000000..3db05446e2 --- /dev/null +++ b/test/brpc_transport_handshake_unittest.cpp @@ -0,0 +1,1407 @@ +// 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 + +// GoogleTest 1.11 and earlier expose flags without getter/setter macros. +#ifndef GTEST_FLAG_GET +#define GTEST_FLAG_GET(name) ::testing::GTEST_FLAG(name) +#define GTEST_FLAG_SET(name, value) (::testing::GTEST_FLAG(name) = (value)) +#endif + +#include +#include +#include +#include +#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/sys_byteorder.h" +#include "butil/time.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" +#include "brpc/rdma/rdma_endpoint.h" +#endif +#if BRPC_WITH_UBRING +#include "brpc/ubshm/ub_endpoint.h" +#endif +#include "brpc/socket.h" +#include "brpc/tcp_transport.h" +#include "brpc/transport_handshake.h" + +namespace brpc { +#if BRPC_WITH_RDMA +namespace rdma { +extern decltype(&bthread_start_background) StartPollingThread; +} +#endif +#if BRPC_WITH_UBRING +namespace ubring { +extern decltype(&bthread_start_background) StartPollingThread; +} +#endif +namespace handshake { + +#if BRPC_WITH_RDMA || BRPC_WITH_UBRING +struct PollerStartFailure { + int attempts = 0; + int fail_on = 0; + bthread::CountdownEvent failed; + bthread::CountdownEvent resume; +}; +static PollerStartFailure* g_poller_start_failure; + +static int StartPollerWithFailure(bthread_t* tid, const bthread_attr_t* attr, + void* (*fn)(void*), void* arg) { + auto* failure = g_poller_start_failure; + if (failure->attempts++ == failure->fail_on) { + failure->failed.signal(); + failure->resume.wait(); + return EAGAIN; + } + return bthread_start_background(tid, attr, fn, arg); +} + +template +static void CheckPollerStartup(decltype(&bthread_start_background)& start) { + auto saved_groups = std::move(Endpoint::_poller_groups); + Endpoint::_poller_groups = std::vector(1); + auto& group = Endpoint::_poller_groups[0]; + group.pollers = std::vector(2); + auto saved_start = start; + start = StartPollerWithFailure; + for (int fail_on : {0, 1}) { + PollerStartFailure failure; + failure.fail_on = fail_on; + g_poller_start_failure = &failure; + std::atomic initialized{0}; + std::atomic released{0}; + std::atomic rejected_reentry{0}; + std::atomic pause_release{false}; + bthread::CountdownEvent release_entered(2), release_resume; + auto initialize = [&] { + return Endpoint::PollingModeInitialize( + 0, nullptr, [&] { ++initialized; }, [&] { + if (pause_release.load()) { + release_entered.signal(); + release_resume.wait(); + } + if (Endpoint::PollingModeInitialize( + 0, nullptr, nullptr, nullptr) == -1 && + errno == EBUSY) { + ++rejected_reentry; + } + ++released; + }); + }; + int first_rc = 0; + std::thread first([&] { first_rc = initialize(); }); + failure.failed.wait(); + // Another caller must wait until startup or rollback completes. + const bool unlocked = group.mutex.try_lock(); + EXPECT_FALSE(unlocked); + if (unlocked) { + group.mutex.unlock(); + } + failure.resume.signal(); + first.join(); + EXPECT_EQ(-1, first_rc); + EXPECT_FALSE(group.running.load()); + EXPECT_EQ(fail_on, released.load()); + for (const auto& poller : group.pollers) { + EXPECT_EQ(INVALID_BTHREAD, poller.tid); + } + EXPECT_EQ(0, initialize()); + EXPECT_EQ(fail_on + 3, failure.attempts); + pause_release.store(true); + std::thread first_release([] { Endpoint::PollingModeRelease(0); }); + release_entered.wait(); + EXPECT_EQ(-1, initialize()); + EXPECT_EQ(EBUSY, errno); + std::thread second_release([] { Endpoint::PollingModeRelease(0); }); + release_resume.signal(); + first_release.join(); + second_release.join(); + EXPECT_EQ(fail_on + 2, initialized.load()); + EXPECT_EQ(initialized.load(), released.load()); + EXPECT_EQ(released.load(), rejected_reentry.load()); + } + start = saved_start; + Endpoint::_poller_groups = std::move(saved_groups); + g_poller_start_failure = nullptr; +} +#endif + +#if BRPC_WITH_RDMA +TEST(PollerStartupTest, rdma_rolls_back_and_retries) { + const std::string saved = GTEST_FLAG_GET(death_test_style); + GTEST_FLAG_SET(death_test_style, "threadsafe"); + EXPECT_EXIT({ + alarm(5); + rdma::FLAGS_rdma_use_polling = true; + CheckPollerStartup(rdma::StartPollingThread); + _exit(::testing::Test::HasFailure() ? 1 : 0); + }, ::testing::ExitedWithCode(0), ""); + GTEST_FLAG_SET(death_test_style, saved); +} +#endif + +#if BRPC_WITH_UBRING +TEST(PollerStartupTest, ubshm_rolls_back_and_retries) { + const std::string saved = GTEST_FLAG_GET(death_test_style); + GTEST_FLAG_SET(death_test_style, "threadsafe"); + EXPECT_EXIT({ + alarm(5); + CheckPollerStartup(ubring::StartPollingThread); + _exit(::testing::Test::HasFailure() ? 1 : 0); + }, ::testing::ExitedWithCode(0), ""); + GTEST_FLAG_SET(death_test_style, saved); +} +#endif + +class HandshakeSocketPair { +public: + void Init(SocketOptions options = SocketOptions()) { + int fds[2]; + ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + peer_fd.reset(fds[1]); + options.fd = fds[0]; + SocketId id; + ASSERT_EQ(0, Socket::Create(options, &id)); + ASSERT_EQ(0, Socket::Address(id, &socket)); + } + + ~HandshakeSocketPair() { + if (socket) { + socket->SetFailed(); + } + } + + butil::fd_guard peer_fd; + SocketUniquePtr socket; +}; + +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; + bthread::CountdownEvent completed; + int result = 0; + int error = 0; + + explicit BlockingHandshakeRead(SocketHandshakeIO* value) : io(value) {} +}; + +static void* RunBlockingHandshakeRead(void* arg) { + BlockingHandshakeRead* read = + static_cast(arg); + char byte = 0; + read->result = read->io->ReadExact(&byte, sizeof(byte)); + read->error = errno; + 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 { +public: + void StartConnect(const Socket*, void (*done)(int, void*), + void* data) override { + done(0, data); + } + + 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; + 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, + FrameSpec::FIXED); +} + +class TestHandshakeProtocol : public HandshakeProtocol { +public: + explicit TestHandshakeProtocol(std::string* calls = nullptr) + : _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(nullptr, 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(nullptr, 0, size); + } + +private: + void Append(const char* value) { + if (_calls != nullptr) { + *_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 = nullptr) + : 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 != nullptr) { + *_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 HandshakeAdapter { +public: + TestServerHandshakeAdapter() { session.SetIOForTest(&io); } + MemoryHandshakeIO io; + mutable HandshakeSession session; + TestHandshakeProtocol protocol; + TestServerHandshakeTransport transport; + +protected: + StepResult RunServerStep(butil::IOBuf* source, Socket*) override { + std::vector protocols(1, &protocol); + transport.source = source; + return session.RunServer(protocols, source, &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, 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); + 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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& socket = pair.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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& socket = pair.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, 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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& peer_fd = pair.peer_fd; + auto& socket = pair.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); + 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, 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, 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)); + 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); + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + ASSERT_EQ(0, bthread_start_background( + &tid, &attr, RunBlockingHandshakeRead, &read)); + + 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); +} + +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); + std::string payload; + ASSERT_EQ(FRAME_NEED_MORE, + FrameCodec::ParseBufferedFrame(&source, spec, &payload)); + ASSERT_EQ(frame.size() - 1, source.size()); + + source.append(frame.data() + frame.size() - 1, 1); + ASSERT_EQ(FRAME_OK, + FrameCodec::ParseBufferedFrame(&source, 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); + std::string payload; + ASSERT_EQ(FRAME_NOT_MINE, + FrameCodec::ParseBufferedFrame(&source, 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); + std::string payload; + ASSERT_EQ(FRAME_PROTOCOL_ERROR, + FrameCodec::ParseBufferedFrame(&source, 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, tcp_keeps_last_message_for_inline_processing) { + struct Message : InputMessageBase { + void DestroyImpl() override {} + ~Message() override = default; + } message; + TcpTransport transport; + InputMessageClosure input; + input.reset(&message); + int queued = 0; + transport.QueueMessage(input, &queued, true); + EXPECT_EQ(&message, input.release()); + EXPECT_EQ(0, queued); +} + +TEST(TransportHandshakeTest, debug_reads_protocol_version_during_negotiation) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto* adapter = AdapterTransport::Get(pair.socket.get()); + auto* session = adapter->handshake_session(); + session->set_protocol_version(2); + session->SetPhase(HELLO_WAIT); + std::atomic stop{false}; + std::thread writer([&] { + while (!stop.load(std::memory_order_relaxed)) { + session->set_protocol_version(2); + session->set_protocol_version(3); + } + }); + for (int i = 0; i < 64; ++i) { + std::ostringstream output; + adapter->Debug(output); + const std::string snapshot = output.str(); + EXPECT_TRUE(snapshot.find("handshake_version=2") != std::string::npos || + snapshot.find("handshake_version=3") != std::string::npos); + } + stop.store(true, std::memory_order_relaxed); + writer.join(); +} + +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, 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)); +} + +class HandshakeSessionTest : public ::testing::Test { +protected: + HandshakeSessionTest() : io("HSOK"), protocol(&calls), transport(&calls) { + session.SetIOForTest(&io); + } + MemoryHandshakeIO io; + HandshakeSession session; + std::string calls; + TestHandshakeProtocol protocol; + TestHandshakeTransport transport; + butil::IOBuf source; + std::vector protocols{&protocol}; +}; + +TEST_F(HandshakeSessionTest, client_runs_codec_and_resource_sequence) { + + 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_F(HandshakeSessionTest, client_exchanges_extension_before_ack) { + io = MemoryHandshakeIO("HSOKE"); + protocol.EnableExtension(1); + + ASSERT_EQ(STEP_OK, session.RunClient(&protocol, &transport)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST_F(HandshakeSessionTest, server_resumes_fragmented_extension_then_ack) { + source.append("HSOK", 4); + protocol.EnableExtension(2); + + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &source, &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, &source, &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, &source, &transport, false)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_TRUE(source.empty()); + EXPECT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST_F(HandshakeSessionTest, client_resource_failure_falls_back_before_io) { + 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_F(HandshakeSessionTest, server_resumes_at_buffered_ack) { + source.append("HSOK", 4); + + + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &source, &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, &source, &transport, false)); + ASSERT_EQ("parse prepare negotiate build parse_ack validate activate", + calls); + ASSERT_EQ(ESTABLISHED, session.phase(butil::memory_order_acquire)); +} + +TEST_F(HandshakeSessionTest, server_resource_failure_falls_back_after_ack) { + source.append("HSOK0", 5); + transport.prepare_result = STEP_FALLBACK; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &source, &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_F(HandshakeSessionTest, server_falls_back_without_consuming_other_magic) { + source.append("XXok", 4); + transport.prepare_result = STEP_ERROR; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &source, &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, &source, &transport, true)); + ASSERT_EQ(4UL, source.size()); + ASSERT_EQ(FALLBACK_TCP, session.phase(butil::memory_order_acquire)); +} + +TEST_F(HandshakeSessionTest, server_enters_hello_phase_after_magic_matches) { + + source.append("H", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &source, &transport, false)); + ASSERT_EQ(UNINITIALIZED, session.phase(butil::memory_order_acquire)); + + source.append("S", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &source, &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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& socket = pair.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, nullptr); + 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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& socket = pair.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, nullptr); + 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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& peer_fd = pair.peer_fd; + auto& socket = pair.socket; + + const std::string hello = MakeUBShmHello(); + butil::IOBuf source; + source.append(hello.data(), 1); + ParseResult result = policy::ParseTransportHandshake( + &source, socket.get(), false, nullptr); + 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, nullptr); + 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, nullptr); + 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) { + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init()); + auto& socket = pair.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, nullptr); + 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 || 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) { + auto original = std::make_shared(); + SocketOptions options; + options.user = get_or_new_client_side_messenger(); + options.need_on_edge_trigger = true; + options.socket_mode = SOCKET_MODE_RDMA; + options.app_connect = original; + HandshakeSocketPair first; + ASSERT_NO_FATAL_FAILURE(first.Init(options)); + const auto wrapped = AdapterTransport::Get(first.socket.get())->Connect(); + ASSERT_NE(nullptr, wrapped); + EXPECT_NE(original.get(), wrapped.get()); + + options.app_connect = wrapped; + HandshakeSocketPair second; + ASSERT_NO_FATAL_FAILURE(second.Init(options)); + EXPECT_EQ(wrapped.get(), AdapterTransport::Get(second.socket.get())->Connect().get()); +} + +#endif + +#if BRPC_WITH_RDMA || BRPC_WITH_UBRING +TEST(TransportHandshakeTest, unavailable_upgrade_keeps_tcp_read_callback_after_reset) { + const SocketMode modes[] = { +#if BRPC_WITH_RDMA + SOCKET_MODE_RDMA, +#endif +#if BRPC_WITH_UBRING + SOCKET_MODE_UBRING, +#endif + }; + auto custom_read = [](Socket*) {}; + for (SocketMode mode : modes) { + for (bool custom : {false, true}) { + SCOPED_TRACE(::testing::Message() << mode << "/" << custom); + auto original = std::make_shared(); + SocketOptions options; + options.user = get_or_new_client_side_messenger(); + options.need_on_edge_trigger = true; + options.socket_mode = mode; + options.app_connect = original; + if (custom) { + options.on_edge_triggered_events = +custom_read; + } + HandshakeSocketPair pair; + ASSERT_NO_FATAL_FAILURE(pair.Init(options)); + auto* adapter = AdapterTransport::Get(pair.socket.get()); + adapter->high_speed_transport()->Release(); + EXPECT_EQ(original.get(), adapter->Connect().get()); + EXPECT_EQ(FALLBACK_TCP, adapter->handshake_phase()); + auto expected = custom ? +custom_read : InputMessenger::OnNewMessages; + ASSERT_NE(nullptr, adapter->_on_edge_trigger.target()); + EXPECT_EQ(expected, *adapter->_on_edge_trigger.target()); + ASSERT_EQ(0, adapter->Reset(2)); + ASSERT_NE(nullptr, adapter->_on_edge_trigger.target()); + EXPECT_EQ(expected, *adapter->_on_edge_trigger.target()); + } + } +} + +#endif + +} // namespace handshake +} // namespace brpc diff --git a/test/brpc_ubring_unittest.cpp b/test/brpc_ubring_unittest.cpp index 558c696649..634094ea58 100644 --- a/test/brpc_ubring_unittest.cpp +++ b/test/brpc_ubring_unittest.cpp @@ -15,10 +15,18 @@ // limitations under the License. #include + +// GoogleTest 1.11 and earlier expose flags without getter/setter macros. +#ifndef GTEST_FLAG_GET +#define GTEST_FLAG_GET(name) ::testing::GTEST_FLAG(name) +#define GTEST_FLAG_SET(name, value) (::testing::GTEST_FLAG(name) = (value)) +#endif + #include #include #include #include +#include #include #include #include @@ -28,12 +36,44 @@ #include "brpc/socket.h" #if BRPC_WITH_UBRING +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "bthread/bthread.h" +#include "bthread/countdown_event.h" +#include "bthread/unstable.h" +#include "brpc/acceptor.h" +#include "brpc/adapter_transport.h" +#include "brpc/channel.h" +#include "brpc/controller.h" +#include "brpc/input_message_base.h" +#include "brpc/input_messenger.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/scoped_lock.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" +#include "brpc/ubshm/shm/shm_ubs.h" +#include "brpc/ubshm/ubs_mem/ubs_mem_def.h" #include "brpc/ubshm/ub_ring_manager.h" #include "brpc/ubshm/ub_ring.h" +#include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/timer/timer_mgr.h" #include "brpc/ubshm/ubr_msg.h" #include "brpc/ubshm/ubr_msg_v2.h" @@ -45,7 +85,14 @@ DECLARE_int32(ub_hb_timer_interval_s); DECLARE_int32(ub_event_queue_timer_interval_us); DECLARE_int32(ub_flying_io_timeout_s); +extern ShmList* g_shm_list; +extern int g_shm_timer_fd; +extern int (*ubsmem_finalize)(); +extern int (*ubsmem_shmem_allocate)(const char*, const char*, size_t, mode_t, uint64_t); +extern int (*ubsmem_shmem_map)(void*, size_t, int, int, const char*, off_t, void**); extern bool g_skip_ub_init; +extern int (*ubsmem_shmem_deallocate)(const char*); +extern int (*ubsmem_shmem_unmap)(void*, size_t); } // namespace ubring } // namespace brpc @@ -586,129 +633,92 @@ TEST(IpcV2UbrDispatchTest, rejects_uninitialized_remote_ring) { EXPECT_EQ(0u, context.remote_rx_event.io_id); } -struct __attribute__((packed)) HelloMessageLayout { - uint16_t msg_len; - uint16_t hello_ver; - uint16_t impl_ver; - uint64_t len; - char shm_name[SHM_MAX_NAME_BUFF_LEN]; -}; -} +constexpr size_t UB_HELLO_PAYLOAD_SIZE = 64 - 2; +} // namespace class HelloMessageTest : public ::testing::Test { protected: - void SetUp() override { - memset(&msg, 0, sizeof(msg)); - buffer.resize(256, 0); - } - - brpc::ubring::HelloMessage msg; - std::string buffer; + brpc::ubring::HelloMessage msg{64, 2, 1, 4 * 1024 * 1024, {0}}; + char buffer[UB_HELLO_PAYLOAD_SIZE]{}; }; -TEST(HelloFormatExtensionTest, serialize_deserialize_roundtrip) { - brpc::ubring::HelloFormatExtension extension = { - brpc::ubring::HelloFormatExtension::WIRE_SIZE, - brpc::ubring::UBR_DATA_FORMAT_LEGACY_64}; - char buffer[brpc::ubring::HelloFormatExtension::WIRE_SIZE] = {}; - - extension.Serialize(buffer); - - brpc::ubring::HelloFormatExtension decoded = {}; - decoded.Deserialize(buffer); - EXPECT_EQ(extension.extension_len, decoded.extension_len); - EXPECT_EQ(extension.format_id, decoded.format_id); - - extension.format_id = brpc::ubring::UBR_DATA_FORMAT_IPC_V2; - extension.Serialize(buffer); - decoded.Deserialize(buffer); - EXPECT_EQ(extension.extension_len, decoded.extension_len); - EXPECT_EQ(extension.format_id, decoded.format_id); -} - -TEST(HelloFormatExtensionTest, serialize_uses_network_byte_order) { - brpc::ubring::HelloFormatExtension extension = {0x0102, 0x0304}; - char buffer[brpc::ubring::HelloFormatExtension::WIRE_SIZE] = {}; - const unsigned char expected[] = {0x01, 0x02, 0x03, 0x04}; - - extension.Serialize(buffer); - - EXPECT_EQ(0, memcmp(expected, buffer, sizeof(expected))); -} - -TEST(HelloFormatExtensionTest, deserialize_none_format) { - const unsigned char buffer[] = {0x00, 0x04, 0x00, 0x00}; - brpc::ubring::HelloFormatExtension extension = {}; - - extension.Deserialize(buffer); - - EXPECT_EQ(4, extension.extension_len); - EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, extension.format_id); -} - -TEST(HelloFormatExtensionTest, deserialize_unknown_format) { - const unsigned char buffer[] = {0x00, 0x04, 0x12, 0x34}; - brpc::ubring::HelloFormatExtension extension = {}; - - extension.Deserialize(buffer); - - EXPECT_EQ(4, extension.extension_len); - EXPECT_EQ(0x1234, extension.format_id); -} - -TEST(HelloFormatExtensionTest, selects_format_for_local_backend) { +TEST(HelloFormatExtensionTest, wire_encoding_and_format_selection) { using namespace brpc::ubring; - EXPECT_EQ(UBR_DATA_FORMAT_IPC_V2, - PreferredDataFormatForShmType(SHM_TYPE_IPC)); - EXPECT_EQ(UBR_DATA_FORMAT_LEGACY_64, - PreferredDataFormatForShmType(SHM_TYPE_UBS)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - PreferredDataFormatForShmType(SHM_TYPE_UB)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - PreferredDataFormatForShmType(SHM_TYPE_UNSUPPORT)); + const struct { uint16_t length; uint16_t format; } cases[] = { + {4, UBR_DATA_FORMAT_LEGACY_64}, {4, UBR_DATA_FORMAT_IPC_V2}, + {4, UBR_DATA_FORMAT_NONE}, {4, 0x1234}, {0x0102, 0x0304}, + }; + for (const auto& c : cases) { + SCOPED_TRACE(c.format); + HelloFormatExtension extension = {c.length, c.format}; + char wire[4]; + extension.Serialize(wire); + const unsigned char expected[] = { + static_cast(c.length >> 8), + static_cast(c.length), + static_cast(c.format >> 8), + static_cast(c.format)}; + EXPECT_EQ(0, memcmp(expected, wire, sizeof(wire))); + HelloFormatExtension decoded{}; + decoded.Deserialize(expected); + EXPECT_EQ(c.length, decoded.extension_len); + EXPECT_EQ(c.format, decoded.format_id); + } + EXPECT_EQ(UBR_DATA_FORMAT_IPC_V2, PreferredDataFormatForShmType(SHM_TYPE_IPC)); + EXPECT_EQ(UBR_DATA_FORMAT_LEGACY_64, PreferredDataFormatForShmType(SHM_TYPE_UBS)); + EXPECT_EQ(UBR_DATA_FORMAT_NONE, PreferredDataFormatForShmType(SHM_TYPE_UB)); + EXPECT_EQ(UBR_DATA_FORMAT_NONE, PreferredDataFormatForShmType(SHM_TYPE_UNSUPPORT)); + for (uint16_t local : {0, 1, 2, 0x1234}) { + for (uint16_t remote : {0, 1, 2, 0x1234}) { + SCOPED_TRACE(::testing::Message() << local << "/" << remote); + const auto expected = local == remote && (local == 1 || local == 2) + ? static_cast(local) : UBR_DATA_FORMAT_NONE; + EXPECT_EQ(expected, SelectDataFormat( + static_cast(local), remote)); + } + } } -TEST(HelloFormatExtensionTest, selects_only_matching_supported_format) { +TEST(UBShmHandshakeAdapterTest, extension_negotiation) { using namespace brpc::ubring; - EXPECT_EQ(UBR_DATA_FORMAT_IPC_V2, - SelectDataFormat(UBR_DATA_FORMAT_IPC_V2, - UBR_DATA_FORMAT_IPC_V2)); - EXPECT_EQ(UBR_DATA_FORMAT_LEGACY_64, - SelectDataFormat(UBR_DATA_FORMAT_LEGACY_64, - UBR_DATA_FORMAT_LEGACY_64)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - SelectDataFormat(UBR_DATA_FORMAT_IPC_V2, - UBR_DATA_FORMAT_LEGACY_64)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - SelectDataFormat(UBR_DATA_FORMAT_LEGACY_64, - UBR_DATA_FORMAT_IPC_V2)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - SelectDataFormat(UBR_DATA_FORMAT_IPC_V2, 0x1234)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - SelectDataFormat(UBR_DATA_FORMAT_NONE, - UBR_DATA_FORMAT_NONE)); - EXPECT_EQ(UBR_DATA_FORMAT_NONE, - SelectDataFormat(static_cast(0x1234), 0x1234)); + for (auto format : {UBR_DATA_FORMAT_LEGACY_64, UBR_DATA_FORMAT_IPC_V2}) { + SCOPED_TRACE(format); + UBShmHandshakeAdapter protocol(format); + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, protocol.BuildExtension(true, &payload)); + const char wire[] = {0, 4, 0, static_cast(format)}; + EXPECT_EQ(std::string(wire, sizeof(wire)), payload); + EXPECT_EQ(brpc::handshake::STEP_OK, protocol.ParseExtension(payload)); + for (char peer : {char(0), char(3 - format)}) { + payload[3] = peer; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, protocol.ParseExtension(payload)); + } + 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))); + ASSERT_EQ(brpc::handshake::STEP_OK, protocol.BuildExtension(false, &payload)); + EXPECT_EQ(std::string("\0\4\0\0", 4), payload); + } } TEST_F(HelloMessageTest, serialize_deserialize_roundtrip) { - msg.msg_len = 64; - msg.hello_ver = 2; - msg.impl_ver = 1; - msg.len = 4 * 1024 * 1024; - memcpy(msg.shm_name, "UBRING_test_C", 14); - - msg.Serialize(&buffer[0]); - - brpc::ubring::HelloMessage decoded; - memset(&decoded, 0, sizeof(decoded)); - decoded.Deserialize(&buffer[0]); - - EXPECT_EQ(msg.msg_len, decoded.msg_len); - EXPECT_EQ(msg.hello_ver, decoded.hello_ver); - EXPECT_EQ(msg.impl_ver, decoded.impl_ver); - EXPECT_EQ(msg.len, decoded.len); - EXPECT_EQ(0, memcmp(msg.shm_name, decoded.shm_name, SHM_MAX_NAME_BUFF_LEN)); + const uint64_t lengths[] = {4 * 1024 * 1024, 0xFFFFFFFFFFFFFFFFULL}; + for (uint64_t len : lengths) { + SCOPED_TRACE(len); + msg.len = len; + memcpy(msg.shm_name, "UBRING_test_C", 14); + msg.Serialize(buffer); + brpc::ubring::HelloMessage decoded{}; + decoded.Deserialize(buffer); + EXPECT_EQ(msg.msg_len, decoded.msg_len); + EXPECT_EQ(msg.hello_ver, decoded.hello_ver); + EXPECT_EQ(msg.impl_ver, decoded.impl_ver); + EXPECT_EQ(msg.len, decoded.len); + EXPECT_EQ(0, memcmp(msg.shm_name, decoded.shm_name, SHM_MAX_NAME_BUFF_LEN)); + } } TEST_F(HelloMessageTest, serialize_uses_network_byte_order) { @@ -720,39 +730,14 @@ TEST_F(HelloMessageTest, serialize_uses_network_byte_order) { msg.Serialize(&buffer[0]); - HelloMessageLayout* raw = reinterpret_cast(&buffer[0]); - EXPECT_EQ(butil::HostToNet16(0x0102), raw->msg_len); - EXPECT_EQ(butil::HostToNet16(0x0304), raw->hello_ver); - EXPECT_EQ(butil::HostToNet16(0x0506), raw->impl_ver); - EXPECT_EQ(butil::HostToNet64(0x0102030405060708ULL), raw->len); -} - -TEST_F(HelloMessageTest, large_len_value) { - msg.msg_len = 64; - msg.hello_ver = 2; - msg.impl_ver = 1; - msg.len = 0xFFFFFFFFFFFFFFFFULL; - memset(msg.shm_name, 0, SHM_MAX_NAME_BUFF_LEN); - - msg.Serialize(&buffer[0]); - - brpc::ubring::HelloMessage decoded; - memset(&decoded, 0, sizeof(decoded)); - decoded.Deserialize(&buffer[0]); - - EXPECT_EQ(0xFFFFFFFFFFFFFFFFULL, decoded.len); + const unsigned char expected[] = { + 1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6, 7, 8}; + EXPECT_EQ(0, memcmp(expected, buffer, sizeof(expected))); } TEST_F(HelloMessageTest, full_shm_name) { - // A peer that fills every byte of shm_name leaves no terminator on the - // wire. Deserialize must still hand back a valid C-string, because the - // handshake feeds the decoded name straight into strncpy/strrchr/ - // snprintf("%s") and the shm allocate/map calls; without a terminator - // those scan past shm_name. + // A full wire name is terminated before use by the SHM mapping functions. memset(msg.shm_name, 'A', SHM_MAX_NAME_BUFF_LEN); - msg.msg_len = 64; - msg.hello_ver = 2; - msg.impl_ver = 1; msg.len = 0; msg.Serialize(&buffer[0]); @@ -768,10 +753,6 @@ TEST_F(HelloMessageTest, full_shm_name) { } TEST_F(HelloMessageTest, toString_contains_fields) { - msg.msg_len = 64; - msg.hello_ver = 2; - msg.impl_ver = 1; - msg.len = 4194304; memcpy(msg.shm_name, "UBRING_test", 12); std::string s = msg.toString(); @@ -781,6 +762,117 @@ 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(UB_HELLO_PAYLOAD_SIZE, '\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(UB_HELLO_PAYLOAD_SIZE, '\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; @@ -819,6 +911,18 @@ TEST(UBRingConfigurationTest, time_flags_include_units_and_expected_defaults) { USEC_TO_NSEC); } +TEST(UBRingConfigurationTest, cleanup_delay_requires_positive_seconds) { + const int saved = brpc::ubring::FLAGS_ub_flying_io_timeout_s; + BRPC_SCOPE_EXIT { brpc::ubring::FLAGS_ub_flying_io_timeout_s = saved; }; + for (const char* value : {"0", "-1"}) { + EXPECT_TRUE(GFLAGS_NAMESPACE::SetCommandLineOption( + "ub_flying_io_timeout_s", value).empty()); + EXPECT_EQ(saved, brpc::ubring::FLAGS_ub_flying_io_timeout_s); + } + EXPECT_FALSE(GFLAGS_NAMESPACE::SetCommandLineOption( + "ub_flying_io_timeout_s", "1").empty()); +} + namespace brpc { namespace ubring { class UBShmEndpointTest : public ::testing::Test { @@ -826,6 +930,7 @@ class UBShmEndpointTest : public ::testing::Test { void SetUp() override { _saved_skip = g_skip_ub_init; g_skip_ub_init = false; + GlobalUBInitializeOrDie(); ShmMgrInit(); UBRingManager::UbrMgrInit(); UBShmEndpoint::GlobalInitialize(); @@ -860,19 +965,89 @@ class UBShmEndpointTest : public ::testing::Test { using brpc::ubring::UBShmEndpointTest; +TEST(UBShmExitDeathTest, stops_pollers_before_static_storage_is_destroyed) { + const std::string saved = GTEST_FLAG_GET(death_test_style); + GTEST_FLAG_SET(death_test_style, "threadsafe"); + BRPC_SCOPE_EXIT { GTEST_FLAG_SET(death_test_style, saved); }; + EXPECT_EXIT({ + static std::atomic stopped{false}; + std::atexit([] { if (!stopped.load()) { _exit(1); } }); + brpc::ubring::g_skip_ub_init = false; + brpc::ubring::GlobalUBInitializeOrDie(); + if (!brpc::ubring::InitPollingModeWithTag( + 0, nullptr, nullptr, [] { stopped.store(true); })) { + _exit(2); + } + std::exit(0); + }, ::testing::ExitedWithCode(0), ""); +} + 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, last_message_is_dispatched_outside_the_poller) { + struct Message : brpc::InputMessageBase { + void DestroyImpl() override {} + ~Message() override = default; + bthread::CountdownEvent done; + bthread_t worker = 0; + } message; + message._process = [](brpc::InputMessageBase* base) { + auto* msg = static_cast(base); + msg->worker = bthread_self(); + msg->done.signal(); + }; + brpc::UBShmTransport transport; + brpc::SocketOptions options; + transport.Init(_socket, options); + BRPC_SCOPE_EXIT { transport.Release(); }; + brpc::InputMessageClosure input; + input.reset(&message); + int queued = 0; + transport.QueueMessage(input, &queued, true); + auto* remaining = input.release(); + EXPECT_EQ(nullptr, remaining); + if (remaining != nullptr) { + return; + } + bthread_flush(); + message.done.wait(); + EXPECT_NE(bthread_self(), message.worker); +} + +TEST_F(UBShmEndpointTest, queued_poll_does_not_retain_a_destroyed_endpoint) { + using namespace brpc::ubring; + brpc::SocketOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + brpc::SocketUniquePtr owner; + ASSERT_EQ(0, brpc::Socket::Address(id, &owner)); + auto* ep = brpc::UBShmTransport::Get(owner.get())->GetUBShmEp(); + SHM local = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + ASSERT_EQ(0, ep->AllocateClientResources(&local, "queued_poll")); + brpc::SocketUniquePtr poller; + ASSERT_EQ(0, brpc::Socket::Address(ep->_poller_sid, &poller)); + const UBShmEndpoint::PollerSidOp op = { + ep->_poller_sid, id, EPOLLIN | EPOLLOUT, + UBShmEndpoint::PollerSidOp::ADD}; + owner->SetFailed(); + owner.reset(); + EXPECT_EQ(nullptr, poller->user()); + UBShmEndpoint::Poll(op); + EXPECT_EQ(0u, UBRingManager::g_ubr_mgr.trx_num); } 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 +1069,861 @@ TEST_F(UBShmEndpointTest, reset_is_idempotent) { _ep->Reset(); } +#if defined(OS_LINUX) +TEST_F(UBShmEndpointTest, cancel_unexpired_timer_does_not_wait_for_expiration) { + using namespace brpc::ubring; + bthread::CountdownEvent canceled; + std::atomic callbacks(0); + itimerspec spec{}; + spec.it_value.tv_sec = 10; + const int fd = TimerStart(&spec, [](void* arg) -> void* { + static_cast*>(arg)->fetch_add(1); + return nullptr; + }, &callbacks); + ASSERT_GE(fd, 0); + std::thread cancel([&] { + DeleteTimerSafe(fd); + canceled.signal(); + }); + const int result = canceled.timed_wait(butil::milliseconds_from_now(1000)); + if (result != 0) { + // Wake a blocking delete so a failure can still join the thread. + spec.it_value.tv_sec = 0; + spec.it_value.tv_nsec = 1; + timerfd_settime(fd, 0, &spec, nullptr); + } + cancel.join(); + EXPECT_EQ(0, result); + EXPECT_EQ(0, callbacks.load()); +} + +TEST_F(UBShmEndpointTest, expired_one_shot_timer_releases_its_descriptor) { + using namespace brpc::ubring; + const uint32_t timers = GetActiveTimerNum(); + bthread::CountdownEvent called; + itimerspec spec{}; + spec.it_value.tv_nsec = 1; + const int fd = TimerStart(&spec, [](void* arg) -> void* { + static_cast(arg)->signal(); + return nullptr; + }, &called); + ASSERT_GE(fd, 0); + EXPECT_EQ(0, called.timed_wait(butil::milliseconds_from_now(1000))); + const int64_t deadline = butil::gettimeofday_us() + 1000000; + while (GetActiveTimerNum() != timers && + butil::gettimeofday_us() < deadline) { + bthread_usleep(1000); + } + const bool released = GetActiveTimerNum() == timers; + if (!released) { + timerfd_settime(fd, 0, &spec, nullptr); + } + EXPECT_TRUE(released); +} + +TEST_F(UBShmEndpointTest, repeated_close_preserves_deferred_cleanup_timer) { + using namespace brpc::ubring; + for (UbrCloseType type : {UBR_SEND_CLOSE, UBR_CALL_BACK_CLOSE}) { + UBRing ring; + SHM local = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + ASSERT_EQ(0, ring.UbrAllocateLocalShm(&local, "close_once")); + UbrTrx& trx = *ring._trx; + reinterpret_cast(trx.ubr_tx.local_tx_event_q.addr) + ->flag = UBR_STATE_CONNECTED; + trx.initializing = false; + itimerspec spec{}; + spec.it_interval.tv_sec = 10; + spec.it_value.tv_sec = 10; + auto noop = [](void*) -> void* { return nullptr; }; + trx.timer_fd = TimerStart(&spec, noop, nullptr); + trx.hb_timer_fd = TimerStart(&spec, noop, nullptr); + ASSERT_GE(trx.timer_fd, 0); + ASSERT_GE(trx.hb_timer_fd, 0); + BRPC_SCOPE_EXIT { + DeleteTimerSafe(trx.clear_timer_fd); + if (ATOMIC_LOAD(trx.ref_count) == 2) { + UBRing::UbrAsynClearCallback(&trx); + } + ring.UbrTrxAbort(); + }; + ASSERT_EQ(UBRING_OK, UBRing::ClearTrxResource(&trx, type)); + EXPECT_EQ(UBRING_REENTRY, + UBRing::ClearTrxResource(&trx, UBR_SEND_CLOSE)); + EXPECT_GE(fcntl(trx.clear_timer_fd, F_GETFD), 0); + } +} + +TEST_F(UBShmEndpointTest, concurrent_close_schedules_one_cleanup_timer) { + using namespace brpc::ubring; + UBRing ring; + SHM local = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + ASSERT_EQ(0, ring.UbrAllocateLocalShm(&local, "close_concurrently")); + UbrTrx& trx = *ring._trx; + reinterpret_cast(trx.ubr_tx.local_tx_event_q.addr) + ->flag = UBR_STATE_CONNECTED; + trx.initializing = false; + itimerspec spec{}; + spec.it_value.tv_sec = 10; + auto noop = [](void*) -> void* { return nullptr; }; + trx.timer_fd = TimerStart(&spec, noop, nullptr); + trx.hb_timer_fd = TimerStart(&spec, noop, nullptr); + ASSERT_GE(trx.timer_fd, 0); + ASSERT_GE(trx.hb_timer_fd, 0); + BRPC_SCOPE_EXIT { + DeleteTimerSafe(trx.clear_timer_fd); + if (ATOMIC_LOAD(trx.ref_count) == 2) { + UBRing::UbrAsynClearCallback(&trx); + } + ring.UbrTrxAbort(); + }; + bthread::CountdownEvent start; + int results[2]; + auto close = [&](int i) { + start.wait(); + results[i] = UBRing::ClearTrxResource(&trx, UBR_SEND_CLOSE); + }; + std::thread first(close, 0); + std::thread second(close, 1); + start.signal(); + first.join(); + second.join(); + std::sort(results, results + 2); + int expected[] = {UBRING_OK, UBRING_REENTRY}; + std::sort(expected, expected + 2); + EXPECT_EQ(expected[0], results[0]); + EXPECT_EQ(expected[1], results[1]); + EXPECT_GE(fcntl(trx.clear_timer_fd, F_GETFD), 0); +} + +TEST_F(UBShmEndpointTest, failed_cleanup_timer_releases_transaction_after_io_wait) { + using namespace brpc::ubring; + const int saved_timeout = FLAGS_ub_flying_io_timeout_s; + FLAGS_ub_flying_io_timeout_s = 1; + BRPC_SCOPE_EXIT { FLAGS_ub_flying_io_timeout_s = saved_timeout; }; + struct sigaction action{}; + struct sigaction saved_action; + action.sa_handler = [](int) {}; + sigemptyset(&action.sa_mask); + ASSERT_EQ(0, sigaction(SIGUSR1, &action, &saved_action)); + BRPC_SCOPE_EXIT { sigaction(SIGUSR1, &saved_action, nullptr); }; + struct rlimit limit; + ASSERT_EQ(0, getrlimit(RLIMIT_NOFILE, &limit)); + for (UbrCloseType type : {UBR_SEND_CLOSE, UBR_CALL_BACK_CLOSE}) { + const uint32_t allocated = UBRingManager::g_ubr_mgr.trx_num; + UBRing ring; + SHM local = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + const std::string name = "timer_failure_" + std::to_string(getpid()); + ASSERT_EQ(UBRING_OK, ring.UbrAllocateLocalShm(&local, name.c_str())); + UbrTrx* trx = ring._trx; + BRPC_SCOPE_EXIT { + if (UBRingManager::g_ubr_mgr.trx_num != allocated) { + DeleteTimerSafe(trx->timer_fd); + DeleteTimerSafe(trx->hb_timer_fd); + UBRing::UbrTrxFreeShm(trx); + UBRingManager::ReleaseUbrTrxFromMgr(trx); + } + }; + SHM remote = local; + ASSERT_EQ(UBRING_OK, ShmRemoteMalloc(&remote)); + ASSERT_EQ(UBRING_OK, ring.UbrTrxMapRemoteShm(&remote)); + auto* tx = reinterpret_cast( + trx->ubr_tx.local_tx_event_q.addr); + auto* rx = reinterpret_cast( + trx->ubr_rx.local_rx_event_q.addr); + tx->flag = UBR_STATE_CONNECTED; + rx->flag = UBR_STATE_CLOSED; + trx->initializing = false; + itimerspec spec{}; + spec.it_value.tv_sec = 10; + auto noop = [](void*) -> void* { return nullptr; }; + trx->timer_fd = TimerStart(&spec, noop, nullptr); + trx->hb_timer_fd = TimerStart(&spec, noop, nullptr); + ASSERT_GE(trx->timer_fd, 0); + ASSERT_GE(trx->hb_timer_fd, 0); + + const pthread_t closing_thread = pthread_self(); + std::thread interrupt([closing_thread] { + usleep(100000); + pthread_kill(closing_thread, SIGUSR1); + }); + BRPC_SCOPE_EXIT { interrupt.join(); }; + // Closing existing descriptors still works, but timerfd_create fails. + struct rlimit exhausted = limit; + exhausted.rlim_cur = 0; + ASSERT_EQ(0, setrlimit(RLIMIT_NOFILE, &exhausted)); + BRPC_SCOPE_EXIT { setrlimit(RLIMIT_NOFILE, &limit); }; + const int64_t start = butil::gettimeofday_us(); + const RETURN_CODE result = type == UBR_SEND_CLOSE + ? ring.UbrTrxClose() + : UBRing::ClearTrxResource(trx, type); + const int64_t elapsed = butil::gettimeofday_us() - start; + ASSERT_EQ(0, setrlimit(RLIMIT_NOFILE, &limit)); + if (type == UBR_CALL_BACK_CLOSE) { + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + } + EXPECT_EQ(UBRING_OK, result); + EXPECT_GE(elapsed, 1000000); + EXPECT_EQ(allocated, UBRingManager::g_ubr_mgr.trx_num); + EXPECT_EQ(nullptr, trx->local_shm.addr); + EXPECT_EQ(nullptr, trx->remote_shm.addr); + if (type == UBR_SEND_CLOSE) { + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + } + butil::fd_guard removed(shm_open(local.name, O_RDONLY, 0)); + EXPECT_LT(removed, 0); + EXPECT_EQ(ENOENT, errno); + } +} +#endif + +TEST_F(UBShmEndpointTest, initialize_preserves_acquired_transactions) { + using namespace brpc::ubring; + UbrTrx* trx = nullptr; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&trx)); + auto* pool = UBRingManager::g_ubr_mgr.trx_mgr; + EXPECT_EQ(UBRING_OK, UBRingManager::UbrMgrInit()); + EXPECT_EQ(pool, UBRingManager::g_ubr_mgr.trx_mgr); + EXPECT_EQ(1u, UBRingManager::g_ubr_mgr.trx_num); + if (pool == UBRingManager::g_ubr_mgr.trx_mgr) { + EXPECT_EQ(UBRING_OK, UBRingManager::ReleaseUbrTrxFromMgr(trx)); + } +} + +TEST_F(UBShmEndpointTest, finalize_waits_for_pending_cleanup) { + using namespace brpc::ubring; + UbrTrx* trx = nullptr; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&trx)); + bthread::CountdownEvent finalized; + std::thread finalize([&] { + UBRingManager::UbrMgrFini(); + finalized.signal(); + }); + const int result = finalized.timed_wait(butil::milliseconds_from_now(100)); + if (result == ETIMEDOUT) { + EXPECT_EQ(UBRING_OK, UBRingManager::ReleaseUbrTrxFromMgr(trx)); + } + finalize.join(); + EXPECT_EQ(ETIMEDOUT, result); + EXPECT_EQ(nullptr, UBRingManager::g_ubr_mgr.trx_mgr); + EXPECT_EQ(UBRING_OK, UBRingManager::UbrMgrInit()); +} + +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; + 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 + +class UBRingLifetimeTest : public UBShmEndpointTest { +protected: + void SetUp() override { + UBShmEndpointTest::SetUp(); + using namespace brpc::ubring; + saved_delay = brpc::ubring::FLAGS_ub_flying_io_timeout_s; + saved_timeout = brpc::ubring::FLAGS_ub_disconnect_timeout_s; + brpc::ubring::FLAGS_ub_flying_io_timeout_s = 1; + local.len = peer.len = 4 * 1024 * 1024; + const std::string name = "lifetime_" + std::to_string(getpid()); + ASSERT_EQ(0, ring.UbrAllocateLocalShm(&local, name.c_str())); + snprintf(peer.name, sizeof(peer.name), "lifetime_peer_%d", getpid()); + ASSERT_EQ(UBRING_OK, ShmLocalCalloc(&peer)); + ASSERT_EQ(UBRING_OK, ring.UbrTrxMapRemoteShm(&peer)); + auto* trx = ring._trx; + reinterpret_cast(trx->ubr_tx.local_tx_event_q.addr) + ->flag = UBR_STATE_CONNECTED; + reinterpret_cast(trx->ubr_rx.local_rx_event_q.addr) + ->flag = UBR_STATE_CONNECTED; + trx->initializing = false; + timers = brpc::ubring::GetActiveTimerNum(); + } + + void TearDown() override { + ring.UbrTrxAbort(); + shm_unlink(peer.name); + UBShmEndpointTest::TearDown(); + brpc::ubring::FLAGS_ub_flying_io_timeout_s = saved_delay; + brpc::ubring::FLAGS_ub_disconnect_timeout_s = saved_timeout; + } + + brpc::ubring::UBRing ring; + brpc::ubring::SHM local{}; + brpc::ubring::SHM peer{}; + int saved_delay = 0; + int saved_timeout = 0; + uint32_t timers = 0; +}; + +TEST_F(UBRingLifetimeTest, timed_out_close_detaches_before_pool_slot_reuse) { + using namespace brpc::ubring; + FLAGS_ub_disconnect_timeout_s = 0; + UbrTrx* old_trx = ring._trx; + ASSERT_EQ(UBRING_ERR_TIMEOUT, ring.UbrTrxClose()); + EXPECT_EQ(nullptr, ring._trx); + ASSERT_TRUE(WaitForUBCondition([&] { return GetActiveTimerNum() == timers; })); + EXPECT_EQ(0u, UBRingManager::g_ubr_mgr.trx_num); + + UBRing replacement; + SHM next = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + ASSERT_EQ(0, replacement.UbrAllocateLocalShm(&next, "timeout_replacement")); + BRPC_SCOPE_EXIT { replacement.UbrTrxAbort(); }; + EXPECT_EQ(old_trx, replacement._trx); + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + EXPECT_EQ(1u, UBRingManager::g_ubr_mgr.trx_num); + butil::fd_guard fd(shm_open(next.name, O_RDONLY, 0)); + EXPECT_GE(fd, 0); +} + +TEST_F(UBRingLifetimeTest, concurrent_device_events_share_cleanup_and_keep_ring_reference) { + using namespace brpc::ubring; + bthread::CountdownEvent start; + RETURN_CODE results[2]; + auto close = [&](int i) { + start.wait(); + results[i] = static_cast( + UBRingManager::UbEventCallback(local.name)); + }; + std::thread first(close, 0); + std::thread second(close, 1); + start.signal(); + first.join(); + second.join(); + EXPECT_EQ(UBRING_OK, results[0]); + EXPECT_EQ(UBRING_OK, results[1]); + ASSERT_TRUE(WaitForUBCondition([&] { return GetActiveTimerNum() == timers; })); + EXPECT_EQ(nullptr, ring._trx->local_shm.addr); + EXPECT_EQ(1u, UBRingManager::g_ubr_mgr.trx_num); + EXPECT_EQ(1, ATOMIC_LOAD(ring._trx->ref_count)); + + UBRing replacement; + SHM next = {nullptr, 4 * 1024 * 1024, 0, {0}, 0}; + ASSERT_EQ(0, replacement.UbrAllocateLocalShm(&next, "passive_replacement")); + BRPC_SCOPE_EXIT { replacement.UbrTrxAbort(); }; + EXPECT_NE(ring._trx, replacement._trx); + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + EXPECT_EQ(1u, UBRingManager::g_ubr_mgr.trx_num); + butil::fd_guard fd(shm_open(peer.name, O_RDONLY, 0)); + EXPECT_GE(fd, 0); +} + +TEST_F(UBRingLifetimeTest, device_event_completes_an_active_close_once) { + using namespace brpc::ubring; + UbrTrx* trx = ring._trx; + ATOMIC_ADD(trx->ref_count, 1); + BRPC_SCOPE_EXIT { UBRingManager::ReleaseUbrTrxFromMgr(trx); }; + RETURN_CODE result = UBRING_ERR; + std::thread close([&] { result = ring.UbrTrxClose(); }); + const bool closing = WaitForUBCondition([&] { + return ATOMIC_LOAD(trx->close_cnt) == MAX_CLOSE_COUNT - 1; + }); + const int event_result = UBRingManager::UbEventCallback(local.name); + close.join(); + EXPECT_TRUE(closing); + EXPECT_EQ(UBRING_OK, event_result); + EXPECT_EQ(UBRING_OK, result); + EXPECT_EQ(nullptr, ring._trx); + ASSERT_TRUE(WaitForUBCondition([&] { return GetActiveTimerNum() == timers; })); + EXPECT_EQ(1, ATOMIC_LOAD(trx->ref_count)); +} + +TEST_F(UBRingLifetimeTest, released_ipc_name_survives_old_mapping_cleanup) { + using namespace brpc::ubring; + ASSERT_EQ(UBRING_OK, ring.UbrUnlinkLocalShm()); + SHM replacement = local; + replacement.addr = nullptr; + ASSERT_EQ(UBRING_OK, ShmLocalCalloc(&replacement)); + BRPC_SCOPE_EXIT { ShmLocalFree(&replacement); }; + ASSERT_EQ(UBRING_OK, UBRing::UbrTrxFreeShm(ring._trx)); + EXPECT_EQ(UBRING_OK, ring.ReleaseTrx()); + butil::fd_guard fd(shm_open(replacement.name, O_RDONLY, 0)); + EXPECT_GE(fd, 0); +} + +TEST_F(UBShmEndpointTest, ubs_unlink_preserves_mapping_until_unmap) { + using namespace brpc::ubring; + const SHM_TYPE saved_type = GetShmType(); + const auto saved_deallocate = ubsmem_shmem_deallocate; + const auto saved_unmap = ubsmem_shmem_unmap; + BRPC_SCOPE_EXIT { + SetShmType(saved_type); + ubsmem_shmem_deallocate = saved_deallocate; + ubsmem_shmem_unmap = saved_unmap; + }; + SetShmType(SHM_TYPE_UBS); + static int unmaps = 0; + unmaps = 0; + ubsmem_shmem_deallocate = [](const char*) { return 0; }; + ubsmem_shmem_unmap = [](void*, size_t) { ++unmaps; return 0; }; + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + BRPC_SCOPE_EXIT { ring.ReleaseTrx(); }; + uint8_t mapped_byte = 0; + ring._trx->local_shm = {&mapped_byte, 4 * 1024 * 1024, 0, {0}, 0}; + strcpy(ring._trx->local_shm.name, "ubs_mapping"); + ASSERT_EQ(UBRING_OK, ring.UbrUnlinkLocalShm()); + EXPECT_EQ(&mapped_byte, ring._trx->local_shm.addr); + EXPECT_EQ(UBRING_OK, UBRing::UbrTrxFreeShm(ring._trx)); + EXPECT_EQ(1, unmaps); +} + +class UBShmServerTest : public UBShmEndpointTest { +protected: + void StartServer(brpc::SocketMode mode = brpc::SOCKET_MODE_UBRING) { + ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions options; + options.socket_mode = mode; + options.enabled_protocols = "baidu_std"; + options.internal_port = -1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &options)); + } + + void ConnectPeer(int timeout_s = 5) { + peer.reset(::socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_GE(peer, 0); + timeval timeout = {timeout_s, 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; + })); + } + + void TearDown() override { + peer.reset(-1); + server_socket.reset(); + server.Stop(0); + server.Join(); + UBShmEndpointTest::TearDown(); + } + + UBFallbackEchoService service; + brpc::Server server; + brpc::SocketUniquePtr server_socket; + butil::fd_guard peer; +}; + +TEST_F(UBShmServerTest, tcp_server_fallback_completes_with_short_rpc_timeout) { + ASSERT_NO_FATAL_FAILURE(StartServer(brpc::SOCKET_MODE_TCP)); + + 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(UBShmServerTest, unavailable_endpoint_serves_rpc_after_health_check_revival) { + ASSERT_NO_FATAL_FAILURE(StartServer(brpc::SOCKET_MODE_TCP)); + brpc::Channel channel; + brpc::ChannelOptions options; + options.socket_mode = brpc::SOCKET_MODE_UBRING; + options.timeout_ms = 1000; + options.max_retry = 0; + ASSERT_EQ(0, channel.Init(server.listen_address(), &options)); + const brpc::SocketId id = channel._server_id; + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(id, &socket)); + auto* adapter = brpc::AdapterTransport::Get(socket.get()); + // Reproduce endpoint allocation failure before selecting AppConnect. + adapter->high_speed_transport()->Release(); + socket->_app_connect = adapter->Connect(); + const auto app_connect = socket->_app_connect; + + test::EchoService_Stub stub(&channel); + auto echo = [&] { + brpc::Controller cntl; + test::EchoRequest request; + test::EchoResponse response; + request.set_message("TCP fallback across health-check revival"); + stub.Echo(&cntl, &request, &response, nullptr); + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ(request.message(), response.message()); + EXPECT_EQ(id, cntl._single_server_id); + }; + ASSERT_NO_FATAL_FAILURE(echo()); + ASSERT_GE(socket->fd(), 0); + ASSERT_EQ(0, socket->SetFailed()); + socket.reset(); // The health check waits for outstanding pins to drain. + ASSERT_TRUE(WaitForUBCondition([&] { + return brpc::Socket::Address(id, &socket) == 0; + })); + EXPECT_EQ(-1, socket->fd()); + EXPECT_EQ(app_connect.get(), socket->_app_connect.get()); + ASSERT_NO_FATAL_FAILURE(echo()); + EXPECT_GE(socket->fd(), 0); + socket->SetFailed(); +} + +TEST_F(UBShmServerTest, failed_server_allocation_cleans_resources_and_serves_tcp) { + ASSERT_NO_FATAL_FAILURE(StartServer()); + ASSERT_NO_FATAL_FAILURE(ConnectPeer()); + + // 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(server_socket.get()); + ASSERT_TRUE(WaitForUBCondition([&] { + return adapter->handshake_phase() == brpc::handshake::FALLBACK_TCP; + })); + auto* transport = brpc::UBShmTransport::Get(server_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(server_socket->Failed()); +} + +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); +} + +class UBShmHandshakeTest : public UBShmServerTest { +protected: + void PrepareUntilAck(const char* prefix, int timeout_s) { + ASSERT_NO_FATAL_FAILURE(StartServer()); + ASSERT_NO_FATAL_FAILURE(ConnectPeer(timeout_s)); + + 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)); + client_transport = brpc::UBShmTransport::Get(client_socket.get()); + local.len = 4 * 1024 * 1024; + local.fd = peer; + name = std::string(prefix) + 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)); + 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)); + adapter = brpc::AdapterTransport::Get(server_socket.get()); + ASSERT_TRUE(WaitForUBCondition([&] { + return adapter->handshake_phase() == brpc::handshake::ACK_WAIT; + })); + format = protocol.local_format(); + } + + void TearDown() override { + if (client_socket) { + client_socket->SetFailed(); + } + client_socket.reset(); + UBShmServerTest::TearDown(); + } + + brpc::SocketUniquePtr client_socket; + brpc::ubring::SHM local{}; + brpc::ubring::HelloMessage server_hello{}; + std::string name; + brpc::ubring::UbrDataFormat format; + brpc::UBShmTransport* client_transport = nullptr; + brpc::AdapterTransport* adapter = nullptr; +}; + +TEST_F(UBShmHandshakeTest, shared_memory_rpc_waits_for_final_tcp_ack) { + ASSERT_NO_FATAL_FAILURE(PrepareUntilAck("ack_race_", 5)); + ASSERT_EQ(0, client_transport->NegotiateUpgradeResources( + &local, name.c_str(), server_hello.len, 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(); +} + +TEST_F(UBShmServerTest, upgraded_client_receives_shared_memory_rpc_responses) { + ASSERT_NO_FATAL_FAILURE(StartServer()); + { + 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 @@ -934,6 +1964,697 @@ TEST(UBRingRecvChunkHeaderTest, crafted_chunk_is_rejected) { chunk.header[UBR_MSG_LEN_INDEX], chunk.header[UBR_MSG_CUR_INDEX])); } +TEST_F(UBShmHandshakeTest, negative_ack_after_server_allocation_serves_tcp_without_delay) { + ASSERT_NO_FATAL_FAILURE(PrepareUntilAck("negative_ack_", 1)); + + ASSERT_NO_FATAL_FAILURE(SendAndCheckUBFallbackRPC(peer, true, 301)); + EXPECT_EQ(brpc::handshake::FALLBACK_TCP, adapter->handshake_phase()); + auto* ep = brpc::UBShmTransport::Get(server_socket.get())->GetUBShmEp(); + EXPECT_EQ(nullptr, ep->_ub_ring); + EXPECT_EQ(brpc::INVALID_SOCKET_ID, ep->_poller_sid); + ASSERT_NO_FATAL_FAILURE(SendAndCheckUBFallbackRPC(peer, false, 302)); + client_transport->DeactivateUpgrade(); +} + +#if defined(OS_LINUX) +TEST_F(UBShmEndpointTest, cancel_waits_for_callback_before_releasing_descriptor) { + using namespace brpc::ubring; + struct Context { + bthread::CountdownEvent entered; + bthread::CountdownEvent release; + } context; + itimerspec spec{}; + spec.it_value.tv_nsec = 1; + const int fd = TimerStart(&spec, [](void* arg) -> void* { + auto* ctx = static_cast(arg); + ctx->entered.signal(); + ctx->release.wait(); + return nullptr; + }, &context); + ASSERT_GE(fd, 0); + bool released = false; + BRPC_SCOPE_EXIT { + if (!released) { context.release.signal(); } + DeleteTimerSafe(fd); + }; + ASSERT_EQ(0, context.entered.timed_wait(butil::milliseconds_from_now(1000))); + bthread::CountdownEvent cancel_started; + bthread::CountdownEvent canceled; + std::thread cancel([&] { + cancel_started.signal(); + DeleteTimerSafe(fd); + canceled.signal(); + }); + cancel_started.wait(); + const int wait_result = canceled.timed_wait(butil::milliseconds_from_now(100)); + spec.it_value.tv_nsec = 0; + spec.it_value.tv_sec = 10; + const int replacement = TimerStart(&spec, [](void*) -> void* { + return nullptr; + }, nullptr); + context.release.signal(); + released = true; + cancel.join(); + EXPECT_EQ(ETIMEDOUT, wait_result); + EXPECT_NE(fd, replacement); + EXPECT_GE(fcntl(replacement, F_GETFD), 0); + DeleteTimerSafe(replacement); +} + +TEST_F(UBShmEndpointTest, callback_can_cancel_its_own_periodic_timer) { + using namespace brpc::ubring; + struct Context { + bthread::CountdownEvent ready; + bthread::CountdownEvent done; + int fd; + } context; + itimerspec spec{}; + spec.it_value.tv_nsec = 1; + spec.it_interval.tv_sec = 10; + context.fd = TimerStart(&spec, [](void* arg) -> void* { + auto* ctx = static_cast(arg); + ctx->ready.wait(); + DeleteTimerSafe(ctx->fd); + ctx->done.signal(); + return nullptr; + }, &context); + ASSERT_GE(context.fd, 0); + context.ready.signal(); + EXPECT_EQ(0, context.done.timed_wait(butil::milliseconds_from_now(1000))); + DeleteTimerSafe(context.fd); + EXPECT_EQ(-1, fcntl(context.fd, F_GETFD)); +} +#endif + + +namespace { +struct UbsCleanupState { + uint8_t local = 0; + uint8_t remote = 0; + std::atomic unmap_error{UBSM_ERR_NET}; + std::atomic deallocate_error{0}; + std::atomic unmaps{0}; + std::atomic remote_unmaps{0}; + std::atomic deallocates{0}; + std::atomic finalizes{0}; + int generation = 1; + bool shm_exists = true; + int allocations = 0; + int map_error = 0; + uint8_t* mappings[2] = {}; + int map_count = 0; + int fault_stage = -1; + std::atomic block_unmap{false}; + std::atomic in_unmap{false}; + std::atomic finalized_during_unmap{false}; + bthread::CountdownEvent entered; + bthread::CountdownEvent resume; +}; +UbsCleanupState* ubs_cleanup = nullptr; +} + +class UbsCleanupTest : public UBShmEndpointTest { +protected: + void SetUp() override { + using namespace brpc::ubring; + UBShmEndpointTest::SetUp(); + saved_type = GetShmType(); + saved_list = g_shm_list; + saved_fd = g_shm_timer_fd; + saved_unmap = ubsmem_shmem_unmap; + saved_deallocate = ubsmem_shmem_deallocate; + saved_finalize = ubsmem_finalize; + saved_allocate = ubsmem_shmem_allocate; + saved_map = ubsmem_shmem_map; + saved_delay = FLAGS_ub_flying_io_timeout_s; + list = static_cast(calloc(1, sizeof(ShmList))); + ASSERT_NE(nullptr, list); + ASSERT_EQ(0, pthread_mutex_init(&list->shm_lock, nullptr)); + g_shm_list = list; + g_shm_timer_fd = 0; + FLAGS_ub_flying_io_timeout_s = 1; + ubs_cleanup = &state; + SetShmType(SHM_TYPE_UBS); + ubsmem_shmem_unmap = [](void* address, size_t) { + auto& state = *ubs_cleanup; + ++state.unmaps; + if (address == &state.remote) { + ++state.remote_unmaps; + } + if (state.block_unmap.load()) { + state.in_unmap.store(true); + state.entered.signal(); + state.resume.wait(); + state.in_unmap.store(false); + } + return state.unmap_error.load(); + }; + ubsmem_shmem_deallocate = [](const char*) { + ++ubs_cleanup->deallocates; + const int error = ubs_cleanup->deallocate_error.load(); + if (error == 0) { + ubs_cleanup->shm_exists = false; + } + return error; + }; + ubsmem_shmem_allocate = [](const char*, const char*, size_t, mode_t, uint64_t) { + auto& state = *ubs_cleanup; + ++state.allocations; + if (state.shm_exists) { + return static_cast(UBSM_ERR_ALREADY_EXIST); + } + state.shm_exists = true; + ++state.generation; + return 0; + }; + ubsmem_shmem_map = [](void*, size_t, int, int, const char*, off_t, void** address) { + *address = &ubs_cleanup->remote; + return ubs_cleanup->map_error; + }; + ubsmem_finalize = []() { + ++ubs_cleanup->finalizes; + ubs_cleanup->finalized_during_unmap.store( + ubs_cleanup->in_unmap.load()); + return 0; + }; + } + void TearDown() override { + using namespace brpc::ubring; + if (g_shm_list != nullptr) { + DestroyShmTimer(g_shm_list); + } + g_shm_list = saved_list; + g_shm_timer_fd = saved_fd; + FLAGS_ub_flying_io_timeout_s = saved_delay; + ubsmem_shmem_unmap = saved_unmap; + ubsmem_shmem_deallocate = saved_deallocate; + ubsmem_finalize = saved_finalize; + ubsmem_shmem_allocate = saved_allocate; + ubsmem_shmem_map = saved_map; + SetShmType(saved_type); + ubs_cleanup = nullptr; + UBShmEndpointTest::TearDown(); + } + brpc::ubring::SHM Mapping(bool remote = false) { + brpc::ubring::SHM shm{}; + shm.addr = remote ? &state.remote : &state.local; + shm.len = 4 * 1024 * 1024; + strcpy(shm.name, remote ? "ubs_remote" : "ubs_local"); + return shm; + } + uint32_t TransactionCount() { + using namespace brpc::ubring; + std::lock_guard lock(UBRingManager::g_ubr_trx_mgr_mtx); + return UBRingManager::g_ubr_mgr.trx_num; + } + UbsCleanupState state; + brpc::ubring::ShmList* list = nullptr; + brpc::ubring::ShmList* saved_list = nullptr; + int saved_fd = 0; + int saved_delay = 0; + brpc::ubring::SHM_TYPE saved_type; + decltype(brpc::ubring::ubsmem_shmem_unmap) saved_unmap; + decltype(brpc::ubring::ubsmem_shmem_deallocate) saved_deallocate; + decltype(brpc::ubring::ubsmem_finalize) saved_finalize; + decltype(brpc::ubring::ubsmem_shmem_allocate) saved_allocate; + decltype(brpc::ubring::ubsmem_shmem_map) saved_map; +}; + +TEST_F(UbsCleanupTest, abort_hands_off_resources_before_reusing_transaction_slot) { + using namespace brpc::ubring; + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + UbrTrx* old_trx = ring._trx; + BRPC_SCOPE_EXIT { + if (ring._trx != nullptr) { + ring.ReleaseTrx(); + } + }; + old_trx->local_shm = Mapping(); + UbrEventQMsg event{}; + event.flag = UBR_STATE_CONNECTED; + old_trx->ubr_tx.local_tx_event_q.addr = reinterpret_cast(&event); + ASSERT_EQ(UBRING_OK, ring.UbrTrxAbort()); + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(0u, UBRingManager::g_ubr_mgr.trx_num); + ASSERT_EQ(1u, list->size); + + UBRing replacement; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&replacement._trx)); + BRPC_SCOPE_EXIT { replacement.ReleaseTrx(); }; + EXPECT_EQ(old_trx, replacement._trx); + state.unmap_error.store(0); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(1, state.deallocates.load()); + EXPECT_EQ(1u, UBRingManager::g_ubr_mgr.trx_num); + EXPECT_EQ(1, ATOMIC_LOAD(replacement._trx->ref_count)); +} + +TEST_F(UbsCleanupTest, passive_cleanup_retries_both_mappings_independently) { + using namespace brpc::ubring; + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + ring._trx->local_shm = Mapping(); + ring._trx->remote_shm = Mapping(true); + ring._trx->close_state = UBR_CLOSE_SECOND; + ring._trx->ref_count = 2; + UBRing::UbrAsynClearCallback(ring._trx); + ASSERT_EQ(UBRING_OK, ring.UbrTrxClose()); + EXPECT_EQ(0u, UBRingManager::g_ubr_mgr.trx_num); + EXPECT_EQ(1, state.remote_unmaps.load()); + ASSERT_EQ(2u, list->size); + // A completed remote task is removed even while the local task still fails. + state.unmap_error.store(0); + state.deallocate_error.store(UBSM_ERR_NET); + ubsmem_shmem_deallocate = [](const char* name) { + ++ubs_cleanup->deallocates; + return strcmp(name, "ubs_local") == 0 + ? ubs_cleanup->deallocate_error.load() : 0; + }; + UbsShmCallback(list); + EXPECT_EQ(1u, list->size); + state.deallocate_error.store(0); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(4, state.unmaps.load()); + EXPECT_EQ(2, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, retry_resumes_at_deallocate_and_accepts_already_deleted_names) { + using namespace brpc::ubring; + state.unmap_error.store(0); + state.deallocate_error.store(UBSM_ERR_IN_USING); + SHM shm = Mapping(); + ASSERT_EQ(UBRING_OK, UbsShmLocalFree(&shm)); + EXPECT_EQ(nullptr, shm.addr); + UbsShmCallback(list); + EXPECT_EQ(1u, list->size); + state.deallocate_error.store(UBSM_ERR_NOT_FOUND); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(1, state.unmaps.load()); + EXPECT_EQ(3, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, remote_retry_keeps_distinct_mappings_and_never_deallocates) { + using namespace brpc::ubring; + SHM first = Mapping(); + SHM second = Mapping(true); + strcpy(second.name, first.name); + ASSERT_EQ(UBRING_OK, UbsShmRemoteFree(&first)); + ASSERT_EQ(UBRING_OK, UbsShmRemoteFree(&second)); + ASSERT_EQ(2u, list->size); + state.unmap_error.store(0); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(4, state.unmaps.load()); + EXPECT_EQ(0, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, duplicate_handoff_preserves_completed_operations) { + using namespace brpc::ubring; + SHM remote = Mapping(); + SHM local = remote; + ASSERT_EQ(UBRING_OK, UbsShmRemoteFree(&remote)); + state.unmap_error.store(0); + ASSERT_EQ(UBRING_OK, UbsShmLocalFree(&local)); + EXPECT_EQ(1u, list->size); + EXPECT_EQ(1, state.unmaps.load()); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(2, state.unmaps.load()); + EXPECT_EQ(1, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, rejected_handoff_keeps_a_cleanup_reference_and_retries) { + using namespace brpc::ubring; + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + UbrTrx* trx = ring._trx; + trx->local_shm = Mapping(); + UbrEventQMsg event{}; + event.flag = UBR_STATE_CONNECTED; + trx->ubr_tx.local_tx_event_q.addr = reinterpret_cast(&event); + g_shm_list = nullptr; + const RETURN_CODE rc = ring.UbrTrxAbort(); + state.block_unmap.store(true); + bool resumed = false; + BRPC_SCOPE_EXIT { + g_shm_list = list; + state.unmap_error.store(0); + if (!resumed) { + state.resume.signal(); + } + EXPECT_TRUE(WaitForUBCondition([&] { return TransactionCount() == 0; })); + }; + ASSERT_EQ(0, state.entered.timed_wait(butil::milliseconds_from_now(3000))); + EXPECT_EQ(UBRING_OK, rc); + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(1u, TransactionCount()); + const int retry_fd = trx->clear_timer_fd; + g_shm_list = list; + state.unmap_error.store(0); + state.resume.signal(); + resumed = true; + DeleteTimerSafe(retry_fd); + EXPECT_EQ(0u, TransactionCount()); + EXPECT_EQ(1, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, allocation_waits_for_old_cleanup_before_reusing_a_name) { + using namespace brpc::ubring; + state.unmap_error.store(0); + state.deallocate_error.store(UBSM_ERR_NET); + SHM old_shm = Mapping(); + ASSERT_EQ(UBRING_OK, UbsShmLocalFree(&old_shm)); + ASSERT_EQ(1u, list->size); + SHM new_shm = Mapping(); + state.deallocate_error.store(0); + EXPECT_EQ(SHM_ERR_EXIST, UbsShmLocalMalloc(&new_shm)); + EXPECT_EQ(0, state.allocations); + EXPECT_EQ(1, state.generation); + EXPECT_TRUE(state.shm_exists); + UbsShmCallback(list); + ASSERT_EQ(0u, list->size); + ASSERT_EQ(UBRING_OK, UbsShmLocalMalloc(&new_shm)); + EXPECT_EQ(2, state.generation); + UbsShmCallback(list); + EXPECT_TRUE(state.shm_exists); + EXPECT_EQ(UBRING_OK, UbsShmLocalFree(&new_shm)); +} + +TEST_F(UbsCleanupTest, allocation_does_not_replace_an_existing_sdk_object) { + using namespace brpc::ubring; + SHM shm = Mapping(); + EXPECT_EQ(SHM_ERR_EXIST, UbsShmLocalMalloc(&shm)); + EXPECT_EQ(1, state.allocations); + EXPECT_EQ(0, state.deallocates.load()); + EXPECT_EQ(1, state.generation); + EXPECT_TRUE(state.shm_exists); +} + +TEST_F(UbsCleanupTest, failed_mapping_keeps_its_name_reserved_until_deallocated) { + using namespace brpc::ubring; + state.shm_exists = false; + state.map_error = UBSM_ERR_NET; + state.deallocate_error.store(UBSM_ERR_NET); + SHM shm = Mapping(); + EXPECT_EQ(SHM_ERR, UbsShmLocalMalloc(&shm)); + ASSERT_EQ(1u, list->size); + EXPECT_FALSE(list->head->unmap); + EXPECT_TRUE(list->head->deallocate); + state.deallocate_error.store(0); + EXPECT_EQ(SHM_ERR_EXIST, UbsShmLocalMalloc(&shm)); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_FALSE(state.shm_exists); + EXPECT_EQ(0, state.unmaps.load()); +} + +TEST_F(UbsCleanupTest, failed_server_allocation_hands_off_the_remote_mapping) { + using namespace brpc::ubring; + ubsmem_shmem_allocate = [](const char*, const char*, size_t, mode_t, uint64_t) { + return static_cast(UBSM_ERR_NET); + }; + UBRing ring; + SHM remote = Mapping(true); + SHM local = Mapping(); + EXPECT_EQ(-1, ring.UbrAllocateServerShm(&remote, &local, UBR_DATA_FORMAT_LEGACY_64)); + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(0u, TransactionCount()); + ASSERT_EQ(1u, list->size); + EXPECT_FALSE(list->head->deallocate); + state.unmap_error.store(0); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(0, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, failed_server_initialization_retains_rejected_cleanup) { + using namespace brpc::ubring; + rlimit limit{}; + ASSERT_EQ(0, getrlimit(RLIMIT_NOFILE, &limit)); + BRPC_SCOPE_EXIT { setrlimit(RLIMIT_NOFILE, &limit); }; + std::vector remote_bytes(4 * 1024 * 1024); + std::vector local_bytes(4 * 1024 * 1024); + state.mappings[0] = remote_bytes.data(); + state.mappings[1] = local_bytes.data(); + state.shm_exists = false; + ubsmem_shmem_map = [](void*, size_t, int, int, const char*, off_t, void** address) { + *address = ubs_cleanup->mappings[ubs_cleanup->map_count++]; + if (ubs_cleanup->map_count == 2) { + brpc::ubring::g_shm_list = nullptr; + } + return 0; + }; + ubsmem_shmem_unmap = [](void*, size_t) { + auto& state = *ubs_cleanup; + if (++state.unmaps == 3) { + state.entered.signal(); + state.resume.wait(); + } + return state.unmap_error.load(); + }; + UBRing ring; + SHM remote = Mapping(true); + SHM local = Mapping(); + rlimit exhausted = limit; + exhausted.rlim_cur = 0; + ASSERT_EQ(0, setrlimit(RLIMIT_NOFILE, &exhausted)); + int result = 0; + std::thread allocate([&] { + result = ring.UbrAllocateServerShm(&remote, &local, UBR_DATA_FORMAT_LEGACY_64); + }); + const int entered = state.entered.timed_wait(butil::milliseconds_from_now(3000)); + EXPECT_EQ(0, entered); + EXPECT_EQ(1u, TransactionCount()); + g_shm_list = list; + state.unmap_error.store(0); + setrlimit(RLIMIT_NOFILE, &limit); + state.resume.signal(); + allocate.join(); + EXPECT_EQ(-1, result); + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(0u, TransactionCount()); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(4, state.unmaps.load()); + EXPECT_EQ(1, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, successful_unlink_never_deletes_a_reused_name) { + using namespace brpc::ubring; + state.unmap_error.store(0); + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + BRPC_SCOPE_EXIT { ring.ReleaseTrx(); }; + ring._trx->local_shm = Mapping(); + ASSERT_EQ(UBRING_OK, ring.UbrUnlinkLocalShm()); + SHM replacement = Mapping(); + ASSERT_EQ(UBRING_OK, UbsShmLocalMalloc(&replacement)); + ASSERT_EQ(2, state.generation); + EXPECT_EQ(UBRING_OK, ring.UbrUnlinkLocalShm()); + EXPECT_EQ(UBRING_OK, UBRing::UbrTrxFreeShm(ring._trx)); + EXPECT_TRUE(state.shm_exists); + EXPECT_EQ(1, state.deallocates.load()); + EXPECT_EQ(1, state.unmaps.load()); + EXPECT_EQ(UBRING_OK, UbsShmLocalFree(&replacement)); +} + +TEST_F(UbsCleanupTest, cleanup_waits_for_in_progress_unlink_without_unmapping) { + using namespace brpc::ubring; + state.unmap_error.store(0); + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + BRPC_SCOPE_EXIT { ring.ReleaseTrx(); }; + ring._trx->local_shm = Mapping(); + ubsmem_shmem_deallocate = [](const char*) { + ++ubs_cleanup->deallocates; + ubs_cleanup->entered.signal(); + ubs_cleanup->resume.wait(); + ubs_cleanup->shm_exists = false; + return 0; + }; + RETURN_CODE rc = UBRING_ERR; + std::thread unlink([&] { rc = ring.UbrUnlinkLocalShm(); }); + const int entered = state.entered.timed_wait(butil::milliseconds_from_now(3000)); + EXPECT_EQ(0, entered); + EXPECT_EQ(UBRING_ERR, UBRing::UbrTrxFreeShm(ring._trx)); + EXPECT_EQ(0, state.unmaps.load()); + state.resume.signal(); + unlink.join(); + EXPECT_EQ(UBRING_OK, rc); + EXPECT_EQ(UBRING_OK, UBRing::UbrTrxFreeShm(ring._trx)); + EXPECT_EQ(1, state.deallocates.load()); + EXPECT_EQ(1, state.unmaps.load()); +} + +TEST_F(UbsCleanupTest, failed_unlink_retains_name_cleanup_for_retry) { + using namespace brpc::ubring; + state.unmap_error.store(0); + state.deallocate_error.store(UBSM_ERR_NET); + UBRing ring; + ASSERT_EQ(UBRING_OK, UBRingManager::AcquireUbrTrxFromMgr(&ring._trx)); + BRPC_SCOPE_EXIT { ring.ReleaseTrx(); }; + ring._trx->local_shm = Mapping(); + EXPECT_EQ(SHM_ERR, ring.UbrUnlinkLocalShm()); + EXPECT_EQ(UBR_SHM_NAME_OWNED, ATOMIC_LOAD(ring._trx->local_name_state)); + EXPECT_EQ(UBRING_OK, UBRing::UbrTrxFreeShm(ring._trx)); + ASSERT_EQ(1u, list->size); + EXPECT_TRUE(list->head->deallocate); + state.deallocate_error.store(0); + UbsShmCallback(list); + EXPECT_EQ(0u, list->size); + EXPECT_EQ(1, state.unmaps.load()); +} + +TEST_F(UbsCleanupTest, device_faults_during_server_construction_abort_the_upgrade) { + using namespace brpc::ubring; + std::vector remote_bytes(4 * 1024 * 1024); + std::vector local_bytes(4 * 1024 * 1024); + state.mappings[0] = remote_bytes.data(); + state.mappings[1] = local_bytes.data(); + state.unmap_error.store(0); + ubsmem_shmem_allocate = [](const char*, const char*, size_t, mode_t, uint64_t) { + ubs_cleanup->shm_exists = true; + if (ubs_cleanup->fault_stage == 1) { + std::thread event([] { + EXPECT_EQ(UBRING_OK, brpc::ubring::UBRingManager::UbEventCallback("ubs_remote")); + }); + event.join(); + EXPECT_EQ(0, ubs_cleanup->unmaps.load()); + } + return 0; + }; + ubsmem_shmem_map = [](void*, size_t, int, int, const char* name, off_t, void** address) { + auto& state = *ubs_cleanup; + const int stage = state.map_count++; + *address = state.mappings[stage]; + if ((stage == 0 && state.fault_stage == 0) || + (stage == 1 && state.fault_stage == 2)) { + std::thread event([name] { + EXPECT_EQ(UBRING_OK, brpc::ubring::UBRingManager::UbEventCallback(name)); + }); + event.join(); + EXPECT_EQ(0, state.unmaps.load()); + } + return 0; + }; + for (int stage = 0; stage != 3; ++stage) { + SCOPED_TRACE(stage); + state.fault_stage = stage; + state.map_count = 0; + state.unmaps.store(0); + state.deallocates.store(0); + UBRing ring; + SHM remote = Mapping(true); + SHM local = Mapping(); + EXPECT_EQ(-1, ring.UbrAllocateServerShm(&remote, &local, UBR_DATA_FORMAT_LEGACY_64)); + EXPECT_EQ(nullptr, ring._trx); + EXPECT_EQ(0u, TransactionCount()); + EXPECT_EQ(stage == 0 ? 1 : 2, state.unmaps.load()); + EXPECT_EQ(stage == 0 ? 0 : 1, state.deallocates.load()); + } +} + +TEST_F(UbsCleanupTest, device_fault_before_client_negotiation_prevents_activation) { + using namespace brpc::ubring; + std::vector bytes(4 * 1024 * 1024); + state.mappings[0] = bytes.data(); + state.shm_exists = false; + state.unmap_error.store(0); + ubsmem_shmem_map = [](void*, size_t, int, int, const char*, off_t, void** address) { + *address = ubs_cleanup->mappings[0]; + return 0; + }; + UBRing ring; + SHM local = Mapping(); + ASSERT_EQ(0, ring.UbrAllocateLocalShm(&local, "client_fault")); + EXPECT_EQ(UBRING_OK, UBRingManager::UbEventCallback(local.name)); + EXPECT_TRUE(ATOMIC_LOAD(ring._trx->device_fault)); + EXPECT_EQ(0, state.unmaps.load()); + EXPECT_EQ(-1, ring.UbrMapRemoteShm(&local, "client_fault", local.len, + UBR_DATA_FORMAT_LEGACY_64)); + EXPECT_EQ(UBRING_OK, ring.UbrTrxAbort()); + EXPECT_EQ(0u, TransactionCount()); + EXPECT_EQ(1, state.unmaps.load()); +} + +TEST_F(UbsCleanupTest, device_fault_after_initialization_cleans_mapping_without_peer_name) { + using namespace brpc::ubring; + std::vector remote_bytes(4 * 1024 * 1024); + std::vector local_bytes(4 * 1024 * 1024); + state.mappings[0] = remote_bytes.data(); + state.mappings[1] = local_bytes.data(); + state.shm_exists = false; + state.unmap_error.store(0); + ubsmem_shmem_map = [](void*, size_t, int, int, const char*, off_t, void** address) { + *address = ubs_cleanup->mappings[ubs_cleanup->map_count++]; + return 0; + }; + UBRing ring; + SHM remote = Mapping(true); + SHM local = Mapping(); + ASSERT_EQ(0, ring.UbrAllocateServerShm(&remote, &local, UBR_DATA_FORMAT_LEGACY_64)); + EXPECT_FALSE(ATOMIC_LOAD(ring._trx->initializing)); + EXPECT_EQ(UBRING_OK, UBRingManager::UbEventCallback(remote.name)); + EXPECT_EQ(UBRING_OK, ring.UbrTrxClose()); + EXPECT_TRUE(WaitForUBCondition([&] { return TransactionCount() == 0; })); + EXPECT_EQ(2, state.unmaps.load()); + EXPECT_EQ(1, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, finalize_retains_failed_tasks_until_sdk_recovery) { + using namespace brpc::ubring; + SHM shm = Mapping(); + ASSERT_EQ(UBRING_OK, UbsShmLocalFree(&shm)); + EXPECT_EQ(UBRING_ERR, UbsShmFini()); + EXPECT_EQ(0, state.finalizes.load()); + EXPECT_NE(nullptr, g_shm_list); + DeleteTimerSafe(g_shm_timer_fd); + g_shm_timer_fd = 0; + state.unmap_error.store(0); + EXPECT_EQ(UBRING_OK, UbsShmFini()); + EXPECT_EQ(nullptr, g_shm_list); + EXPECT_EQ(1, state.finalizes.load()); + EXPECT_EQ(1, state.deallocates.load()); +} + +TEST_F(UbsCleanupTest, finalize_waits_for_an_active_retry_callback) { + using namespace brpc::ubring; + SHM shm = Mapping(); + ASSERT_EQ(UBRING_OK, UbsShmLocalFree(&shm)); + state.unmap_error.store(0); + state.block_unmap.store(true); + itimerspec spec{}; + spec.it_interval.tv_sec = 10; + spec.it_value.tv_nsec = 1; + g_shm_timer_fd = TimerStart(&spec, UbsShmCallback, list); + ASSERT_GE(g_shm_timer_fd, 0); + const int entered = state.entered.timed_wait(butil::milliseconds_from_now(1000)); + if (entered != 0) { + state.resume.signal(); + FAIL() << "Retry callback did not start"; + } + bthread::CountdownEvent started; + bthread::CountdownEvent finished; + RETURN_CODE rc = UBRING_ERR; + std::thread close([&] { + started.signal(); + rc = UbsShmFini(); + finished.signal(); + }); + started.wait(); + EXPECT_NE(0, finished.timed_wait(butil::milliseconds_from_now(30))); + EXPECT_EQ(0, state.finalizes.load()); + state.resume.signal(); + close.join(); + EXPECT_EQ(UBRING_OK, rc); + EXPECT_FALSE(state.finalized_during_unmap.load()); + EXPECT_EQ(1, state.finalizes.load()); + EXPECT_EQ(nullptr, g_shm_list); +} + #else TEST(UbringDisabledTest, skip) {