From a10de421e2057268170f8df511cdbfe3dd4b954a Mon Sep 17 00:00:00 2001 From: zchuango Date: Wed, 5 Aug 2026 14:56:06 +0800 Subject: [PATCH] Add URMA transport support --- .gitignore | 4 + BUILD.bazel | 13 +- CMakeLists.txt | 84 +- MODULE.bazel | 11 + Makefile | 6 + WORKSPACE | 9 +- bazel/config/BUILD.bazel | 8 +- bazel/third_party/umdk/umdk.BUILD | 28 + config_brpc.sh | 25 +- docs/cn/urma.md | 149 ++ docs/en/urma.md | 161 ++ example/cmake/BrpcExample.cmake | 59 +- example/urma_performance/CMakeLists.txt | 43 + example/urma_performance/client.cpp | 156 ++ example/urma_performance/server.cpp | 96 ++ example/urma_performance/test.proto | 34 + src/brpc/channel.cpp | 2 + src/brpc/channel.h | 3 +- src/brpc/input_messenger.cpp | 10 +- src/brpc/input_messenger.h | 6 + src/brpc/server.h | 2 +- src/brpc/socket.h | 18 + src/brpc/socket_mode.h | 8 +- src/brpc/transport_factory.cpp | 44 +- src/brpc/transport_factory.h | 11 +- src/brpc/urma/mock_urma.cpp | 713 +++++++++ src/brpc/urma/urma_bonding.h | 35 + src/brpc/urma/urma_endpoint.cpp | 1846 +++++++++++++++++++++++ src/brpc/urma/urma_endpoint.h | 355 +++++ src/brpc/urma/urma_handshake.cpp | 356 +++++ src/brpc/urma/urma_handshake.h | 180 +++ src/brpc/urma/urma_handshake.proto | 47 + src/brpc/urma/urma_helper.cpp | 788 ++++++++++ src/brpc/urma/urma_helper.h | 119 ++ src/brpc/urma_transport.cpp | 248 +++ src/brpc/urma_transport.h | 84 ++ test/brpc_urma_unittest.cpp | 609 ++++++++ 37 files changed, 6331 insertions(+), 39 deletions(-) create mode 100644 bazel/third_party/umdk/umdk.BUILD create mode 100644 docs/cn/urma.md create mode 100644 docs/en/urma.md create mode 100644 example/urma_performance/CMakeLists.txt create mode 100644 example/urma_performance/client.cpp create mode 100644 example/urma_performance/server.cpp create mode 100644 example/urma_performance/test.proto create mode 100644 src/brpc/urma/mock_urma.cpp create mode 100644 src/brpc/urma/urma_bonding.h create mode 100644 src/brpc/urma/urma_endpoint.cpp create mode 100644 src/brpc/urma/urma_endpoint.h create mode 100644 src/brpc/urma/urma_handshake.cpp create mode 100644 src/brpc/urma/urma_handshake.h create mode 100644 src/brpc/urma/urma_handshake.proto create mode 100644 src/brpc/urma/urma_helper.cpp create mode 100644 src/brpc/urma/urma_helper.h create mode 100644 src/brpc/urma_transport.cpp create mode 100644 src/brpc/urma_transport.h create mode 100644 test/brpc_urma_unittest.cpp diff --git a/.gitignore b/.gitignore index c7b21b9350..739963a26c 100644 --- a/.gitignore +++ b/.gitignore @@ -45,6 +45,10 @@ CTestTestfile.cmake /test/out.txt /test/recordio_ref.io +# Local design notes and Graphify artifacts. +docs/cn/urma_proposal.md +graphify-out/ + # Ignore protoc-gen-mcpack files /protoc-gen-mcpack*/ diff --git a/BUILD.bazel b/BUILD.bazel index 5dc5fcf726..042883273d 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -51,6 +51,9 @@ DEFINES = [ }) + select({ "//bazel/config:brpc_with_rdma": ["BRPC_WITH_RDMA=1"], "//conditions:default": [], + }) + select({ + "//bazel/config:brpc_with_urma": ["BRPC_WITH_URMA=1"], + "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_debug_bthread_sche_safety": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1"], "//conditions:default": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=0"], @@ -509,6 +512,7 @@ filegroup( "src/brpc/*.proto", "src/brpc/policy/*.proto", "src/brpc/rdma/*.proto", + "src/brpc/urma/*.proto", ]), visibility = ["//visibility:public"], ) @@ -547,9 +551,7 @@ cc_library( "src/brpc/event_dispatcher_kqueue.cpp", ]), copts = COPTS, - includes = [ - "src/", - ], + includes = ["src/"], linkopts = LINKOPTS, visibility = ["//visibility:public"], deps = [ @@ -561,6 +563,11 @@ cc_library( ":mcpack2pb", "@com_github_google_leveldb//:leveldb", ] + select({ + "//bazel/config:brpc_with_urma": [ + "@umdk//:urma_headers", + ], + "//conditions:default": [], + }) + select({ "//bazel/config:brpc_with_thrift": [ "@org_apache_thrift//:thrift", ], diff --git a/CMakeLists.txt b/CMakeLists.txt index 915b7d2977..3f5838205f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,6 +27,10 @@ option(WITH_THRIFT "With thrift framed protocol supported" OFF) option(WITH_BTHREAD_TRACER "With bthread tracer supported" OFF) option(WITH_SNAPPY "With snappy" OFF) option(WITH_RDMA "With RDMA" OFF) +option(WITH_URMA "With URMA (openEuler Unified Remote Memory Access)" OFF) +option(DOWNLOAD_URMA_HEADERS + "Download UMDK headers when WITH_URMA is enabled and headers are absent" + ON) option(WITH_UBRING "With UB" OFF) option(WITH_DEBUG_BTHREAD_SCHE_SAFETY "With debugging bthread sche safety" OFF) option(WITH_DEBUG_LOCK "With debugging lock" OFF) @@ -114,6 +118,11 @@ if(WITH_RDMA) set(WITH_RDMA_VAL "1") endif() +set(WITH_URMA_VAL "0") +if(WITH_URMA) + set(WITH_URMA_VAL "1") +endif() + set(WITH_UBRING_VAL "0") if(WITH_UBRING) set(WITH_UBRING_VAL "1") @@ -153,6 +162,7 @@ endif() list(APPEND BRPC_COMMON_DEFINITIONS BRPC_WITH_GLOG=${WITH_GLOG_VAL} BRPC_WITH_RDMA=${WITH_RDMA_VAL} + BRPC_WITH_URMA=${WITH_URMA_VAL} BRPC_WITH_UBRING=${WITH_UBRING_VAL} BRPC_DEBUG_BTHREAD_SCHE_SAFETY=${WITH_DEBUG_BTHREAD_SCHE_SAFETY_VAL} BRPC_DEBUG_LOCK=${WITH_DEBUG_LOCK_VAL} @@ -309,6 +319,56 @@ if(WITH_RDMA) list(APPEND BRPC_COMMON_INCLUDE_DIRS ${RDMA_INCLUDE_PATH}) endif() +if(WITH_URMA) + # UrmaTransport and its link-time mock both compile against the upstream + # UMDK API. Prefer installed headers; otherwise fetch a pinned upstream + # release, following Mooncake's URMA mock setup. + find_path(URMA_INCLUDE_PATH NAMES urma_api.h + HINTS ENV URMA_ROOT + PATHS /usr/include /usr/local/include + PATH_SUFFIXES ub/umdk/urma umdk/urma urma + src/urma/lib/urma/core/include) + find_path(URMA_BOND_INCLUDE_PATH NAMES urma_ubagg.h + HINTS ENV URMA_ROOT + PATHS /usr/include /usr/local/include + PATH_SUFFIXES ub/umdk/urma umdk/urma urma + src/urma/lib/urma/bond/include) + if(NOT URMA_INCLUDE_PATH AND DOWNLOAD_URMA_HEADERS) + include(FetchContent) + FetchContent_Declare( + urma_headers + GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git + GIT_TAG v26.06.0_CAM + GIT_SHALLOW TRUE) + FetchContent_GetProperties(urma_headers) + if(NOT urma_headers_POPULATED) + FetchContent_Populate(urma_headers) + endif() + set(URMA_INCLUDE_PATH + "${urma_headers_SOURCE_DIR}/src/urma/lib/urma/core/include") + set(URMA_BOND_INCLUDE_PATH + "${urma_headers_SOURCE_DIR}/src/urma/lib/urma/bond/include") + message(STATUS "Using downloaded UMDK headers: ${URMA_INCLUDE_PATH}") + endif() + if(NOT URMA_INCLUDE_PATH) + message(FATAL_ERROR + "Fail to find urma_api.h. Install UMDK headers, set URMA_ROOT, " + "or enable DOWNLOAD_URMA_HEADERS.") + endif() + + find_library(URMA_LIB NAMES urma + HINTS ENV URMA_ROOT + PATH_SUFFIXES lib lib64) + if(URMA_LIB) + message(STATUS "Found URMA library: ${URMA_LIB}") + set(URMA_USE_MOCK 0) + else() + message(STATUS + "liburma not found; building with the URMA link-time mock") + set(URMA_USE_MOCK 1) + endif() +endif() + find_library(PROTOC_LIB NAMES protoc) if(NOT PROTOC_LIB) message(FATAL_ERROR "Fail to find protoc lib") @@ -333,6 +393,12 @@ list(APPEND BRPC_COMMON_INCLUDE_DIRS ${PROTOBUF_INCLUDE_DIRS} ${LEVELDB_INCLUDE_PATH} ) +if(WITH_URMA) + list(APPEND BRPC_COMMON_INCLUDE_DIRS ${URMA_INCLUDE_PATH}) + if(URMA_BOND_INCLUDE_PATH) + list(APPEND BRPC_COMMON_INCLUDE_DIRS ${URMA_BOND_INCLUDE_PATH}) + endif() +endif() set(DYNAMIC_LIB ${GFLAGS_LIBRARY} @@ -361,6 +427,13 @@ if(WITH_RDMA) list(APPEND DYNAMIC_LIB ${RDMA_LIB}) endif() +if(WITH_URMA) + message(STATUS "brpc compile with URMA (mock=${URMA_USE_MOCK})") + if(NOT URMA_USE_MOCK) + list(APPEND DYNAMIC_LIB ${URMA_LIB}) + endif() +endif() + if(WITH_UBRING) message(STATUS "brpc compile with ubring") list(APPEND DYNAMIC_LIB ${UB_LIB}) @@ -566,6 +639,14 @@ file(GLOB_RECURSE BRPC_SOURCES CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/src/brpc file(GLOB_RECURSE THRIFT_SOURCES CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/src/brpc/thrift*.cpp") file(GLOB_RECURSE EXCLUDE_SOURCES CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/src/brpc/event_dispatcher_*.cpp") +# When building with the real liburma, exclude the link-time mock so its urma_* +# symbols do not clash with the library. When liburma is absent, keep it so CI +# can build and test UrmaTransport without URMA hardware. +if(WITH_URMA AND NOT URMA_USE_MOCK) + list(REMOVE_ITEM BRPC_SOURCES + "${PROJECT_SOURCE_DIR}/src/brpc/urma/mock_urma.cpp") +endif() + if(WITH_THRIFT) message("brpc compile with thrift protocol") else() @@ -606,7 +687,8 @@ set(PROTO_FILES idl_options.proto brpc/trackme.proto brpc/streaming_rpc_meta.proto brpc/proto_base.proto - brpc/rdma/rdma_handshake.proto) + brpc/rdma/rdma_handshake.proto + brpc/urma/urma_handshake.proto) file(MAKE_DIRECTORY ${PROJECT_BINARY_DIR}/output/include/brpc) set(PROTOC_FLAGS ${PROTOC_FLAGS} -I${PROTOBUF_INCLUDE_DIR}) compile_proto(PROTO_HDRS PROTO_SRCS ${PROJECT_BINARY_DIR} diff --git a/MODULE.bazel b/MODULE.bazel index 6f7b01a86e..657b7ceb43 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -52,3 +52,14 @@ git_override( remote = 'https://github.com/hedronvision/bazel-compile-commands-extractor.git', commit = '1e08f8e0507b6b6b1f4416a9a22cf5c28beaba93', # Jun 28, 2024 ) + +git_repository = use_repo_rule( + '@bazel_tools//tools/build_defs/repo:git.bzl', + 'git_repository', +) +git_repository( + name = 'umdk', + build_file = '//bazel/third_party/umdk:umdk.BUILD', + remote = 'https://atomgit.com/openeuler/umdk.git', + tag = 'v26.06.0_CAM', +) diff --git a/Makefile b/Makefile index 86de388448..271b518ae6 100644 --- a/Makefile +++ b/Makefile @@ -204,9 +204,15 @@ JSON2PB_SOURCES = $(foreach d,$(JSON2PB_DIRS),$(wildcard $(addprefix $(d)/*,$(SR JSON2PB_OBJS = $(addsuffix .o, $(basename $(JSON2PB_SOURCES))) BRPC_DIRS = src/brpc src/brpc/details src/brpc/builtin src/brpc/policy src/brpc/policy/mysql src/brpc/rdma +ifeq ($(WITH_URMA),1) +BRPC_DIRS += src/brpc/urma +endif THRIFT_SOURCES = $(foreach d,$(BRPC_DIRS),$(wildcard $(addprefix $(d)/thrift*,$(SRCEXTS)))) EXCLUDE_SOURCES = $(foreach d,$(BRPC_DIRS),$(wildcard $(addprefix $(d)/event_dispatcher_*,$(SRCEXTS)))) BRPC_SOURCES_ALL = $(foreach d,$(BRPC_DIRS),$(wildcard $(addprefix $(d)/*,$(SRCEXTS)))) +ifeq ($(URMA_USE_MOCK),0) +BRPC_SOURCES_ALL := $(filter-out src/brpc/urma/mock_urma.cpp,$(BRPC_SOURCES_ALL)) +endif BRPC_SOURCES = $(filter-out $(THRIFT_SOURCES) $(EXCLUDE_SOURCES), $(BRPC_SOURCES_ALL)) BRPC_PROTOS = $(filter %.proto,$(BRPC_SOURCES)) BRPC_CFAMILIES = $(filter-out %.proto %.pb.cc,$(BRPC_SOURCES)) diff --git a/WORKSPACE b/WORKSPACE index 78a6c2836a..b197d666f6 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -279,6 +279,13 @@ http_archive( urls = ["https://archive.apache.org/dist/thrift/0.15.0/thrift-0.15.0.tar.gz"], ) +git_repository( + name = "umdk", + build_file = "//bazel/third_party/umdk:umdk.BUILD", + remote = "https://atomgit.com/openeuler/umdk.git", + tag = "v26.06.0_CAM", +) + # Header-only JSON library used by iobuf_unittest's IOBuf<->std::iostream # adapter tests. Keep version in sync with MODULE.bazel. http_archive( @@ -317,4 +324,4 @@ http_archive( sha256 = "3cd0e49f0f4a6d406c1d74b53b7616f5e24f5fd319eafc1bf8eee6e14124d115", ) load("@hedron_compile_commands//:workspace_setup.bzl", "hedron_compile_commands_setup") -hedron_compile_commands_setup() \ No newline at end of file +hedron_compile_commands_setup() diff --git a/bazel/config/BUILD.bazel b/bazel/config/BUILD.bazel index eec551da8b..e801d096ee 100644 --- a/bazel/config/BUILD.bazel +++ b/bazel/config/BUILD.bazel @@ -104,6 +104,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "brpc_with_urma", + define_values = {"BRPC_WITH_URMA": "true"}, + visibility = ["//visibility:public"], +) + config_setting( name = "brpc_with_boringssl", define_values = {"BRPC_WITH_BORINGSSL": "true"}, @@ -149,4 +155,4 @@ config_setting( name = "with_babylon_counter", define_values = {"with_babylon_counter": "true"}, visibility = ["//visibility:public"], -) \ No newline at end of file +) diff --git a/bazel/third_party/umdk/umdk.BUILD b/bazel/third_party/umdk/umdk.BUILD new file mode 100644 index 0000000000..410e6a2736 --- /dev/null +++ b/bazel/third_party/umdk/umdk.BUILD @@ -0,0 +1,28 @@ +# 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. + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "urma_headers", + hdrs = glob([ + "src/urma/lib/urma/bond/include/*.h", + "src/urma/lib/urma/core/include/*.h", + ]), + includes = [ + "src/urma/lib/urma/bond/include", + "src/urma/lib/urma/core/include", + ], +) diff --git a/config_brpc.sh b/config_brpc.sh index 1c05942068..ce46bf7a8e 100755 --- a/config_brpc.sh +++ b/config_brpc.sh @@ -54,10 +54,11 @@ else LDD=ldd fi -TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,nodebugsymbols,werror -n 'config_brpc' -- "$@"` +TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-urma,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,nodebugsymbols,werror -n 'config_brpc' -- "$@"` WITH_GLOG=0 WITH_THRIFT=0 WITH_RDMA=0 +WITH_URMA=0 WITH_MESALINK=0 WITH_BTHREAD_TRACER=0 WITH_ASAN=0 @@ -89,6 +90,7 @@ while true; do --with-glog ) WITH_GLOG=1; shift 1 ;; --with-thrift) WITH_THRIFT=1; shift 1 ;; --with-rdma) WITH_RDMA=1; shift 1 ;; + --with-urma) WITH_URMA=1; shift 1 ;; --with-mesalink) WITH_MESALINK=1; shift 1 ;; --with-bthread-tracer) WITH_BTHREAD_TRACER=1; shift 1 ;; --with-debug-bthread-sche-safety ) BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1; shift 1 ;; @@ -536,6 +538,26 @@ if [ $WITH_RDMA != 0 ]; then append_to_output "WITH_RDMA=1" fi +if [ $WITH_URMA != 0 ]; then + URMA_LIB=$(find_dir_of_lib urma) + URMA_HDR=$(find_dir_of_header_or_die urma_api.h) + URMA_BOND_HDR=$(find_dir_of_header urma_ubagg.h) + CPPFLAGS="${CPPFLAGS} -DBRPC_WITH_URMA=1" + append_to_output "WITH_URMA=1" + append_to_output_headers "$URMA_HDR" + if [ -n "$URMA_BOND_HDR" ]; then + append_to_output_headers "$URMA_BOND_HDR" + fi + if [ -n "$URMA_LIB" ]; then + append_to_output_libs "$URMA_LIB" + append_to_output "DYNAMIC_LINKINGS+=-lurma" + append_to_output "URMA_USE_MOCK=0" + else + append_to_output "URMA_USE_MOCK=1" + print_info "liburma not found; using URMA link-time mock" + fi +fi + if [ $WITH_MESALINK != 0 ]; then CPPFLAGS="${CPPFLAGS} -DUSE_MESALINK" fi @@ -667,6 +689,7 @@ print_info "System: $SYSTEM" if [ $WITH_GLOG -ne 0 ]; then print_info "With glog: yes"; fi if [ $WITH_THRIFT -ne 0 ]; then print_info "With thrift: yes"; fi if [ $WITH_RDMA -ne 0 ]; then print_info "With RDMA: yes"; fi +if [ $WITH_URMA -ne 0 ]; then print_info "With URMA: yes"; fi if [ $WITH_MESALINK -ne 0 ]; then print_info "With MesaLink: yes"; fi if [ $WITH_BTHREAD_TRACER -ne 0 ]; then print_info "With bthread tracer: yes"; fi if [ $WITH_ASAN -ne 0 ]; then print_info "With ASAN: yes"; fi diff --git a/docs/cn/urma.md b/docs/cn/urma.md new file mode 100644 index 0000000000..a9be634ec7 --- /dev/null +++ b/docs/cn/urma.md @@ -0,0 +1,149 @@ +# UrmaTransport：基于 URMA 的远程内存 RPC + +UrmaTransport 是使用 openEuler +[URMA](https://atomgit.com/openeuler/umdk)（Unified Remote Memory Access）SDK +实现的传输层。它是 +[#3217](https://github.com/apache/brpc/discussions/3217) 中提出的路线 B， +与基于 OBMM 的 UBRing 传输（路线 A， +[#3226](https://github.com/apache/brpc/issues/3226)）互补，承担大包/跨节点 +高吞吐场景，共同构成路线 C（双后端）。 + +## 技术背景 + +URMA 在 UMDK 支持的设备上提供 verbs 风格接口。当前实现创建可靠消息 +（`URMA_TM_RM`）Jetty，并使用 CTP 传输路径；通过 +`urma_post_jetty_send_wr` 提交发送 WR，通过 `urma_post_jfr_wr` 提交接收 +WR。完成事件既可由 JFC 忙轮询获取，也可通过 JFCE 事件 fd 获取。 + +## 编译配置 + +### CMake 编译 + +```bash +# 带 URMA 支持编译 brpc +cmake -B build -DWITH_URMA=ON +make -C build -j$(nproc) + +# 编译 urma_performance 示例 +cd example/urma_performance +cmake -B build +make -C build -j$(nproc) +``` + +`WITH_URMA=ON` 使用上游 UMDK 头文件进行编译。CMake 优先使用系统安装的 +SDK；找不到头文件时，会参照 Mooncake 的 mock 构建方式下载固定版本的 +UMDK，可通过 `DOWNLOAD_URMA_HEADERS=OFF` 禁止下载。找到 `liburma` 时使用 +真实硬件数据通路，否则链接 brpc 的 mock，使 URMA 代码和测试仍可在无硬件 +环境编译。 + +## 使用 + +通过在 channel / server 上设置 `socket_mode` 选择传输层： + +```cpp +// 客户端 +brpc::ChannelOptions opt; +opt.socket_mode = brpc::SOCKET_MODE_URMA; +opt.protocol = "baidu_std"; // URMA 仅支持 baidu_std +brpc::Channel channel; +channel.Init("127.0.0.1:8003", &opt); + +// 服务端 +brpc::ServerOptions sopt; +sopt.socket_mode = brpc::SOCKET_MODE_URMA; +server.Start(port, &sopt); +``` + +若对端不支持 URMA（例如 TCP 客户端连接 URMA 服务端），在 4 字节 magic +握手后透明回退到 TCP，应用代码无需改动。 + +## 架构 + +UrmaTransport 沿用与 `RdmaTransport` / `UBShmTransport` 一致的两层设计： + +``` +UrmaTransport : public Transport (urma_transport.{h,cpp}) + +-- std::shared_ptr (回退路径) + +-- urma::UrmaEndpoint* (URMA 数据路径) + +-- UrmaState { URMA_ON, URMA_OFF, URMA_UNKNOWN } + +urma::UrmaEndpoint : public SocketUser (urma/urma_endpoint.{h,cpp}) + +-- UrmaResource { jfc, jfce, jfr, jetty, remote_jetty, remote_seg } + +-- 握手状态机（C/S 对称，在 TCP fd 上驱动） + +-- 发送路径：urma_post_jetty_send_wr(URMA_OPC_SEND) + +-- 接收路径：urma_poll_jfc -> HandleCompletion -> InputMessenger + +-- 双窗口信用流控（_remote_rq_window_size / _sq_window_size） +``` + +### 建链流程（双平面） + +与 RDMA / UBRing 一致，控制面为 TCP，数据面为 URMA： + +1. TCP 连接建立。 +2. `UrmaConnect::StartConnect` 起客户端握手 bthread。 +3. 双方在 TCP fd 上交换 `UrmaHello` 消息（v2 二进制 magic `URMA`， + v3 protobuf magic `URM3`），携带本地 EID、jetty id、recv buffer 数量、 + 以及扁平化的 buffer 池 segment。 +4. 双方先调用 `urma_import_seg` **再**调用 `urma_import_jetty`，为远端 EID + 建立传输路径（TP）路由。跳过 `import_seg` 会导致首个 SEND 被硬件以 + `URMA_CR_RNR_RETRY_CNT_EXC_ERR` 拒绝。 +5. 4 字节 ACK（`HELLO_ACK_URMA_OK = 0x1`）确认双方均要 URMA。 +6. 成功后 TCP fd 仅保留用于 epoll 生命周期和回退，数据走 URMA。 + +### 内存管理 + +申请一大段 `mmap` 内存，用 `urma_register_seg` 一次性注册，再切成固定大小 +buffer（默认 8KB）。劫持 `butil::iobuf::blockmem_allocate` 使每个 IOBuf +block 都由注册 segment 支撑，发送路径可直接从 IOBuf block refs 构建 +`urma_sge_t`，无需逐消息注册（与 RDMA `block_pool` 设计一致）。用户注册 +内存通过 `urma::RegisterMemoryForUrma` / `DeregisterMemoryForUrma` 支持。 + +## 配置 + +所有 flag 使用 `urma_` 前缀（对标 RDMA 的 `rdma_` 前缀）： + +| Flag | 默认 | 用途 | +|------|------|------| +| `--urma_use_polling` | false | 轮询 JFC 而非事件模式 | +| `--urma_poller_num` | 1 | 每 bthread tag 的轮询器数（轮询模式） | +| `--urma_disable_bthread` | false | 内联处理消息（不起 bthread） | +| `--urma_sq_size` | 128 | 本地 JFS 深度 [16, 4096] | +| `--urma_rq_size` | 128 | 本地 JFR 深度 [16, 4096] | +| `--urma_cqe_poll_once` | 32 | 每次 `urma_poll_jfc` 的上限 | +| `--urma_recv_zerocopy` | true | 大于 `--urma_zerocopy_min_size` 的接收零拷贝 | +| `--urma_zerocopy_min_size` | 512 | 小于此值的接收拷贝 | +| `--urma_device` | "" | URMA 设备名（空=首个） | +| `--urma_max_sge` | 0 | 每 WR SGE 上限（0=设备上限） | +| `--urma_bonding_mode` | 0 | bonding 模式：0=standalone，1=active-backup，2=balance | +| `--urma_bonding_level` | 0 | bonding 层级：0=IODIE，1=port | +| `--urma_prepared_jetty_cnt` | 8 | 预连接 Jetty+CQ 请求数量；会根据 `RLIMIT_NOFILE` 自动限制 | +| `--urma_buffer_size` | 8192 | 池中每个 buffer 大小（字节） | +| `--urma_buffer_count` | 65536 | 池中 buffer 数量 | +| `--urma_poller_yield` | false | 忙轮询循环中主动让出 bthread | +| `--urma_client_handshake_version` | 2 | 客户端握手版本（2=二进制，3=protobuf） | + +设备名以 `bonding` 开头时，brpc 会在创建 context 后、创建 segment 和队列 +前配置 provider。默认 standalone+IODIE 配置与 UMDK 性能工具保持一致。 +bonding 支持需要 provider 扩展头文件 `urma_ubagg.h`。 + +## 与 UBRing 协同 + +UrmaTransport 推荐用于**大包和跨节点**高吞吐路径，而 UBRing +（`SOCKET_MODE_UBRING`）对**小包和同机 IPC**最优（亚微秒、零系统调用）。 +混合负载可按服务选择传输层： + +| 场景 | 建议 `socket_mode` | +|------|---------------------| +| 同机 IPC | `SOCKET_MODE_UBRING` | +| 跨节点小包（< 64KB） | `SOCKET_MODE_UBRING`（UBS-Mem）或 `SOCKET_MODE_URMA` | +| 跨节点大包（>= 64KB） | `SOCKET_MODE_URMA` | +| 传统 RoCE/IB 数据中心 | `SOCKET_MODE_URMA` 或 `SOCKET_MODE_RDMA` | + +单连接内按包大小自动分流的方案（路线 C 方案 B）作为后续演进方向。 + +## 限制 + +- 仅支持 `baidu_std` 协议（与 RDMA 一致）。SSL、RTMP、NSHEAD、MONGO 在 + `ContextInitOrDie` 阶段拒绝。 +- 硬件数据通路需要受支持的 UMDK provider 和 `liburma`。 +- 当前实现面向 Linux。 diff --git a/docs/en/urma.md b/docs/en/urma.md new file mode 100644 index 0000000000..964f648bd2 --- /dev/null +++ b/docs/en/urma.md @@ -0,0 +1,161 @@ +# UrmaTransport: URMA-based Remote Memory RPC + +UrmaTransport is a transport implementation that uses openEuler's +[URMA](https://atomgit.com/openeuler/umdk) (Unified Remote Memory Access) SDK +for remote-memory RPC. It is the Route B transport proposed in +[#3217](https://github.com/apache/brpc/discussions/3217) and complements the +OBMM-based UBRing transport (Route A, [#3226](https://github.com/apache/brpc/issues/3226)) +for large-packet / cross-node scenarios (Route C, "double backend"). + +## Technical Background + +URMA exposes a verbs-style API over devices supported by UMDK. This +implementation creates a reliable-message (`URMA_TM_RM`) Jetty with a CTP +transport path, posts send work requests with `urma_post_jetty_send_wr`, and +posts receive work requests with `urma_post_jfr_wr`. Completions are consumed +from a JFC either by busy polling or through a JFCE event fd. + +## Build Configuration + +### Build with CMake + +```bash +# Build brpc with URMA support +cmake -B build -DWITH_URMA=ON +make -C build -j$(nproc) + +# Build the urma_performance example +cd example/urma_performance +cmake -B build +make -C build -j$(nproc) +``` + +`WITH_URMA=ON` compiles against upstream UMDK headers. CMake prefers an +installed SDK and, following Mooncake's mock setup, downloads a pinned UMDK +release when the headers are unavailable. Set `DOWNLOAD_URMA_HEADERS=OFF` to +disable downloading. +When `liburma` is found it is linked for the hardware data path. Otherwise, +brpc uses its link-time mock so URMA code and tests can still be built without +hardware. + +## Usage + +Select the transport by setting `socket_mode` on the channel / server: + +```cpp +// Client +brpc::ChannelOptions opt; +opt.socket_mode = brpc::SOCKET_MODE_URMA; +opt.protocol = "baidu_std"; // URMA supports baidu_std only +brpc::Channel channel; +channel.Init("127.0.0.1:8003", &opt); + +// Server +brpc::ServerOptions sopt; +sopt.socket_mode = brpc::SOCKET_MODE_URMA; +server.Start(port, &sopt); +``` + +If the peer does not speak URMA (e.g. a TCP-only client connecting to a +URMA-enabled server), the transport transparently falls back to TCP after +the 4-byte magic handshake. No application code change is required. + +## Architecture + +UrmaTransport follows the same two-layer design as `RdmaTransport` and +`UBShmTransport`: + +``` +UrmaTransport : public Transport (urma_transport.{h,cpp}) + +-- std::shared_ptr (fallback path) + +-- urma::UrmaEndpoint* (URMA data path) + +-- UrmaState { URMA_ON, URMA_OFF, URMA_UNKNOWN } + +urma::UrmaEndpoint : public SocketUser (urma/urma_endpoint.{h,cpp}) + +-- UrmaResource { jfc, jfce, jfr, jetty, remote_jetty, remote_seg } + +-- handshake state machine (C/S symmetric, driven over the TCP fd) + +-- send path: urma_post_jetty_send_wr(URMA_OPC_SEND) + +-- recv path: urma_poll_jfc -> HandleCompletion -> InputMessenger + +-- two-window credit flow control + (_remote_rq_window_size / _sq_window_size) +``` + +### Connection establishment (dual-plane) + +Like RDMA / UBRing, the control plane is TCP and the data plane is URMA: + +1. TCP connect completes. +2. `UrmaConnect::StartConnect` spawns the client handshake bthread. +3. Both sides exchange a `UrmaHello` message (magic `URMA` for v2 binary, + `URM3` for v3 protobuf) over the TCP fd. The message carries the local + EID, jetty id, recv buffer count, and the flattened buffer-pool segment. +4. Each side calls `urma_import_seg` **before** `urma_import_jetty` to + establish transport-path (TP) routing for the remote EID. Skipping the + `import_seg` step causes the first SEND to be rejected by hardware with + `URMA_CR_RNR_RETRY_CNT_EXC_ERR`. +5. A 4-byte ACK (`HELLO_ACK_URMA_OK = 0x1`) confirms both sides want URMA. +6. On success, the TCP fd is kept only for the epoll lifecycle and fallback; + payloads flow through URMA. + +### Memory management + +A single large region is `mmap`-ed and registered once with +`urma_register_seg`, then sliced into fixed-size buffers (default 8 KB). +`butil::iobuf::blockmem_allocate` is hijacked so every IOBuf block is backed +by the registered segment, allowing the send path to build `urma_sge_t` +directly from IOBuf block refs without per-message registration (mirroring the +RDMA `block_pool` design). User-registered memory is supported via +`urma::RegisterMemoryForUrma` / `DeregisterMemoryForUrma`. + +## Configuration + +All flags use the `urma_` prefix (mirroring RDMA's `rdma_` prefix): + +| Flag | Default | Purpose | +|------|---------|---------| +| `--urma_use_polling` | false | Busy-poll the JFC instead of event mode | +| `--urma_poller_num` | 1 | Poller bthreads per bthread tag (polling mode) | +| `--urma_disable_bthread` | false | Run message processing inline | +| `--urma_sq_size` | 128 | Local JFS depth [16, 4096] | +| `--urma_rq_size` | 128 | Local JFR depth [16, 4096] | +| `--urma_cqe_poll_once` | 32 | Max CQEs per `urma_poll_jfc` | +| `--urma_recv_zerocopy` | true | Zero-copy receives above `--urma_zerocopy_min_size` | +| `--urma_zerocopy_min_size` | 512 | Receives smaller than this are copied | +| `--urma_device` | "" | URMA device name (empty = first) | +| `--urma_max_sge` | 0 | Max SGEs per WR (0 = device max) | +| `--urma_bonding_mode` | 0 | Bonding mode: 0=standalone, 1=active-backup, 2=balance | +| `--urma_bonding_level` | 0 | Bonding level: 0=IODIE, 1=port | +| `--urma_prepared_jetty_cnt` | 8 | Requested pre-allocated Jetty+CQ sets; automatically capped according to `RLIMIT_NOFILE` | +| `--urma_buffer_size` | 8192 | Per-buffer size in the pool (bytes) | +| `--urma_buffer_count` | 65536 | Number of buffers in the pool | +| `--urma_poller_yield` | false | Yield in the busy-poll loop | +| `--urma_client_handshake_version` | 2 | Client wire version (2=binary, 3=protobuf) | + +For a device whose name starts with `bonding`, brpc configures the provider +immediately after context creation and before creating segments or queues. +The default standalone+IODIE combination matches the UMDK performance tool. +Bonding support requires the provider extension header `urma_ubagg.h`. + +## Coexistence with UBRing + +UrmaTransport is the recommended transport for **large packets and +cross-node** high-throughput paths, while UBRing (`SOCKET_MODE_UBRING`) is +optimal for **small packets and same-node IPC** (sub-microsecond, zero +syscalls). For a mixed workload, select the transport per service: + +| Scenario | Recommended `socket_mode` | +|----------|---------------------------| +| Same-host IPC | `SOCKET_MODE_UBRING` | +| Cross-node small packets (< 64 KB) | `SOCKET_MODE_UBRING` (UBS-Mem) or `SOCKET_MODE_URMA` | +| Cross-node large packets (>= 64 KB) | `SOCKET_MODE_URMA` | +| Traditional RoCE / IB datacenter | `SOCKET_MODE_URMA` or `SOCKET_MODE_RDMA` | + +A single-connection hybrid that auto-routes by packet size (Route C, scheme B) +is tracked as a future enhancement. + +## Limitations + +- `baidu_std` protocol only (same as RDMA). SSL, RTMP, NSHEAD, MONGO are + rejected at `ContextInitOrDie`. +- A hardware data path requires a supported UMDK provider and `liburma`. +- The current implementation targets Linux. diff --git a/example/cmake/BrpcExample.cmake b/example/cmake/BrpcExample.cmake index 7afb3bd3a9..bd003aa600 100644 --- a/example/cmake/BrpcExample.cmake +++ b/example/cmake/BrpcExample.cmake @@ -39,6 +39,52 @@ macro(brpc_example_find_common_deps out_libs) find_package(Threads REQUIRED) find_package(Protobuf REQUIRED) + set(BRPC_EXAMPLE_CXX_STANDARD 11) + set(_protobuf_absl_targets) + if(Protobuf_VERSION VERSION_GREATER 4.21) + # Protobuf 5+ exposes Abseil types from generated code and requires + # C++17. Keep this list in sync with the top-level CMake build. + set(BRPC_EXAMPLE_CXX_STANDARD 17) + find_package(absl REQUIRED CONFIG) + set(_protobuf_absl_targets + absl::absl_check + absl::absl_log + absl::algorithm + absl::base + absl::bind_front + absl::bits + absl::btree + absl::cleanup + absl::cord + absl::core_headers + absl::debugging + absl::die_if_null + absl::dynamic_annotations + absl::flags + absl::flat_hash_map + absl::flat_hash_set + absl::function_ref + absl::hash + absl::layout + absl::log_initialize + absl::log_globals + absl::log_severity + absl::memory + absl::node_hash_map + absl::node_hash_set + absl::random_distributions + absl::random_random + absl::span + absl::status + absl::statusor + absl::strings + absl::synchronization + absl::time + absl::type_traits + absl::utility + absl::variant + ) + endif() # Search for libthrift* by best effort. If it is not found and brpc is # compiled with thrift protocol enabled, a link error would be reported. @@ -81,6 +127,7 @@ macro(brpc_example_find_common_deps out_libs) Threads::Threads ${GFLAGS_LIBRARY} ${PROTOBUF_LIBRARIES} + ${_protobuf_absl_targets} ${LEVELDB_LIB} ${OPENSSL_CRYPTO_LIBRARY} ${OPENSSL_SSL_LIBRARY} @@ -118,13 +165,19 @@ function(brpc_example_configure_target target_name) ${OPENSSL_INCLUDE_DIR} ${GPERFTOOLS_INCLUDE_DIR} ${RDMA_INCLUDE_PATH} + ${URMA_INCLUDE_PATH} ) if(_include_dirs) target_include_directories(${target_name} PRIVATE ${_include_dirs}) endif() - target_compile_features(${target_name} PRIVATE cxx_std_11) + if(NOT BRPC_EXAMPLE_CXX_STANDARD) + set(BRPC_EXAMPLE_CXX_STANDARD 11) + endif() + target_compile_features(${target_name} PRIVATE + cxx_std_${BRPC_EXAMPLE_CXX_STANDARD} + ) target_compile_definitions(${target_name} PRIVATE NDEBUG __const__=__unused__ @@ -147,6 +200,10 @@ function(brpc_example_configure_target target_name) target_compile_definitions(${target_name} PRIVATE BRPC_WITH_RDMA=1) endif() + if(BRPC_EXAMPLE_WITH_URMA) + target_compile_definitions(${target_name} PRIVATE BRPC_WITH_URMA=1) + endif() + if(CMAKE_SYSTEM_NAME STREQUAL "Darwin") include(CheckFunctionExists) check_function_exists(clock_gettime BRPC_EXAMPLE_HAVE_CLOCK_GETTIME) diff --git a/example/urma_performance/CMakeLists.txt b/example/urma_performance/CMakeLists.txt new file mode 100644 index 0000000000..154970fbd3 --- /dev/null +++ b/example/urma_performance/CMakeLists.txt @@ -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. + +cmake_minimum_required(VERSION 3.16...3.28) +project(urma_performance C CXX) + +include(${CMAKE_CURRENT_LIST_DIR}/../cmake/BrpcExample.cmake) + +option(LINK_SO "Whether examples are linked dynamically" OFF) + +brpc_example_find_common_deps(DYNAMIC_LIB) + +protobuf_generate_cpp(PROTO_SRC PROTO_HEADER test.proto) +set(BRPC_EXAMPLE_WITH_URMA ON) +find_library(URMA_LIB NAMES urma) +if(URMA_LIB) + list(APPEND DYNAMIC_LIB ${URMA_LIB}) +else() + message(STATUS + "liburma not found; using the URMA implementation linked into brpc") +endif() + +add_executable(urma_performance_client client.cpp ${PROTO_SRC} ${PROTO_HEADER}) +brpc_example_configure_target(urma_performance_client) +add_executable(urma_performance_server server.cpp ${PROTO_SRC} ${PROTO_HEADER}) +brpc_example_configure_target(urma_performance_server) + +target_link_libraries(urma_performance_client PRIVATE ${BRPC_LIB} ${DYNAMIC_LIB}) +target_link_libraries(urma_performance_server PRIVATE ${BRPC_LIB} ${DYNAMIC_LIB}) diff --git a/example/urma_performance/client.cpp b/example/urma_performance/client.cpp new file mode 100644 index 0000000000..9449dbe960 --- /dev/null +++ b/example/urma_performance/client.cpp @@ -0,0 +1,156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#include + +#include +#include + +#include + +#include "butil/atomicops.h" +#include "butil/fast_rand.h" +#include "butil/logging.h" +#include "brpc/channel.h" +#include "brpc/controller.h" +#include "bthread/bthread.h" +#include "bvar/latency_recorder.h" +#include "bvar/variable.h" +#include "test.pb.h" + +#if BRPC_WITH_URMA + +DEFINE_string(server, "127.0.0.1:8003", "IP Port of urma performance server"); +DEFINE_int32(thread_num, 0, "How many threads are used"); +DEFINE_int32(queue_depth, 1, "How many requests can be pending in the queue"); +DEFINE_int32(expected_qps, 0, "The expected QPS"); +DEFINE_int32(max_thread_num, 16, "The max number of threads are used"); +DEFINE_int32(attachment_size, -1, "Attachment size is used (in Bytes)"); +DEFINE_int32(rpc_timeout_ms, 5000, "Timeout for each RPC in milliseconds"); +DEFINE_bool(echo_attachment, false, "Select whether attachment should be echo"); +DEFINE_bool(use_urma, true, "Use URMA transport (true) or TCP (false)"); + +bvar::LatencyRecorder g_latency("client"); +bvar::Adder g_error_count("client_error_count"); + +static void* worker(void* arg) { + test::PerfTestService_Stub* stub = + static_cast(arg); + int qps = FLAGS_expected_qps; + while (!brpc::IsAskedToQuit()) { + butil::FastRandSeed seed; + butil::init_fast_rand_seed(&seed); + std::vector cntls(FLAGS_queue_depth); + std::vector reqs(FLAGS_queue_depth); + std::vector resps(FLAGS_queue_depth); + std::vector ids(FLAGS_queue_depth); + for (int i = 0; i < FLAGS_queue_depth; ++i) { + cntls[i].set_log_id(butil::fast_rand(&seed) & 0x7fffffff); + reqs[i].set_echo_attachment(FLAGS_echo_attachment); + if (FLAGS_attachment_size >= 0) { + cntls[i].request_attachment().resize(FLAGS_attachment_size, 'a'); + } + ids[i] = cntls[i].call_id(); + stub->Test(&cntls[i], &reqs[i], &resps[i], brpc::DoNothing()); + } + for (int i = 0; i < FLAGS_queue_depth; ++i) { + brpc::Join(ids[i]); + if (cntls[i].Failed()) { + g_error_count << 1; + LOG_EVERY_SECOND(WARNING) + << "RPC failed: " << cntls[i].ErrorText(); + } else { + g_latency << cntls[i].latency_us(); + } + } + if (qps > 0) { + usleep(FLAGS_queue_depth * 1000000 / qps); + } + } + return nullptr; +} + +int main(int argc, char* argv[]) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + brpc::ChannelOptions options; + options.socket_mode = FLAGS_use_urma ? brpc::SOCKET_MODE_URMA + : brpc::SOCKET_MODE_TCP; + options.connect_timeout_ms = FLAGS_rpc_timeout_ms; + options.timeout_ms = FLAGS_rpc_timeout_ms; + options.max_retry = 0; + brpc::Channel channel; + if (channel.Init(FLAGS_server.c_str(), &options) != 0) { + LOG(ERROR) << "Fail to init channel to " << FLAGS_server; + return -1; + } + test::PerfTestService_Stub stub(&channel); + + // Complete one RPC before starting all workers. This makes handshake and + // data-path failures visible instead of looking like a hung benchmark. + brpc::Controller warmup_cntl; + warmup_cntl.set_timeout_ms(FLAGS_rpc_timeout_ms); + test::PerfTestRequest warmup_req; + test::PerfTestResponse warmup_resp; + warmup_req.set_echo_attachment(false); + stub.Test(&warmup_cntl, &warmup_req, &warmup_resp, nullptr); + if (warmup_cntl.Failed()) { + LOG(ERROR) << "Warm-up RPC failed after timeout_ms=" + << FLAGS_rpc_timeout_ms << ": " + << warmup_cntl.ErrorText(); + return -1; + } + LOG(INFO) << "Warm-up RPC to " << FLAGS_server + << " succeeded, latency=" << warmup_cntl.latency_us() << "us"; + + int thread_num = FLAGS_thread_num; + if (thread_num == 0) { + thread_num = FLAGS_max_thread_num; + } + if (thread_num <= 0 || FLAGS_queue_depth <= 0) { + LOG(ERROR) << "thread_num and queue_depth must be positive"; + return -1; + } + std::vector tids(thread_num); + for (int i = 0; i < thread_num; ++i) { + bthread_start_background(&tids[i], nullptr, worker, &stub); + } + LOG(INFO) << "URMA performance client started (server=" << FLAGS_server + << ", use_urma=" << FLAGS_use_urma + << ", threads=" << thread_num + << ", rpc_timeout_ms=" << FLAGS_rpc_timeout_ms << ")"; + while (!brpc::IsAskedToQuit()) { + sleep(1); + LOG(INFO) << "qps=" << g_latency.qps(1) + << " latency=" << g_latency.latency(1) << "us" + << " errors=" << g_error_count.get_value(); + } + for (int i = 0; i < thread_num; ++i) { + bthread_join(tids[i], nullptr); + } + return 0; +} + +#else + +#include +int main() { + printf("This example requires brpc built with -DWITH_URMA=ON.\n"); + return 0; +} + +#endif // BRPC_WITH_URMA diff --git a/example/urma_performance/server.cpp b/example/urma_performance/server.cpp new file mode 100644 index 0000000000..baaff323b6 --- /dev/null +++ b/example/urma_performance/server.cpp @@ -0,0 +1,96 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "butil/atomicops.h" +#include "butil/logging.h" +#include "butil/time.h" +#include "brpc/closure_guard.h" +#include "brpc/controller.h" +#include "brpc/server.h" +#include "bvar/variable.h" +#include "test.pb.h" + +#if BRPC_WITH_URMA + +DEFINE_int32(port, 8003, "TCP Port of this server"); +DEFINE_bool(use_urma, true, "Use URMA transport (true) or TCP (false)"); + +butil::atomic g_last_time(0); + +namespace test { +class PerfTestServiceImpl : public PerfTestService { +public: + void Test(google::protobuf::RpcController* cntl_base, + const PerfTestRequest* request, + PerfTestResponse* response, + google::protobuf::Closure* done) { + brpc::ClosureGuard done_guard(done); + const uint64_t last = + g_last_time.load(butil::memory_order_relaxed); + const uint64_t now = butil::monotonic_time_us(); + if (now > last && now - last > 100000) { + if (g_last_time.exchange(now, butil::memory_order_relaxed) == last) { + response->set_cpu_usage( + bvar::Variable::describe_exposed("process_cpu_usage")); + } else { + response->set_cpu_usage(""); + } + } else { + response->set_cpu_usage(""); + } + if (request->echo_attachment()) { + brpc::Controller* cntl = static_cast(cntl_base); + cntl->response_attachment().append(cntl->request_attachment()); + } + } +}; +} // namespace test + +int main(int argc, char* argv[]) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + brpc::Server server; + test::PerfTestServiceImpl service; + + if (server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE) != 0) { + LOG(ERROR) << "Fail to add PerfTestService"; + return -1; + } + + brpc::ServerOptions options; + options.socket_mode = FLAGS_use_urma ? brpc::SOCKET_MODE_URMA + : brpc::SOCKET_MODE_TCP; + if (server.Start(FLAGS_port, &options) != 0) { + LOG(ERROR) << "Fail to start server"; + return -1; + } + LOG(INFO) << "URMA performance server started on port " << FLAGS_port + << " (use_urma=" << FLAGS_use_urma << ")"; + server.RunUntilAskedToQuit(); + return 0; +} + +#else + +#include +int main() { + printf("This example requires brpc built with -DWITH_URMA=ON.\n"); + return 0; +} + +#endif // BRPC_WITH_URMA diff --git a/example/urma_performance/test.proto b/example/urma_performance/test.proto new file mode 100644 index 0000000000..10b41c8fe5 --- /dev/null +++ b/example/urma_performance/test.proto @@ -0,0 +1,34 @@ +// 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. + +syntax = "proto2"; + +option cc_generic_services = true; + +package test; + +message PerfTestRequest { + required bool echo_attachment = 1; +} + +message PerfTestResponse { + required string cpu_usage = 1; +} + +service PerfTestService { + rpc Test(PerfTestRequest) returns (PerfTestResponse); +} diff --git a/src/brpc/channel.cpp b/src/brpc/channel.cpp index 83fc37b077..d578c2a3ac 100644 --- a/src/brpc/channel.cpp +++ b/src/brpc/channel.cpp @@ -139,6 +139,8 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { } 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/channel.h b/src/brpc/channel.h index 28a17ac8ea..62bc6bfb49 100644 --- a/src/brpc/channel.h +++ b/src/brpc/channel.h @@ -106,7 +106,8 @@ struct ChannelOptions { const ChannelSSLOptions& ssl_options() const { return *_ssl_options; } ChannelSSLOptions* mutable_ssl_options(); - // Let this channel Choose to use a certain socket: 0 SOCKET_MODE_TCP, 1 SOCKET_MODE_RDMA. + // Let this channel choose a transport. + // See SocketMode for supported values. // Default: SOCKET_MODE_TCP SocketMode socket_mode; diff --git a/src/brpc/input_messenger.cpp b/src/brpc/input_messenger.cpp index 1e42f92351..8ed0246e90 100644 --- a/src/brpc/input_messenger.cpp +++ b/src/brpc/input_messenger.cpp @@ -297,11 +297,11 @@ int InputMessenger::ProcessNewMessage( num_bthread_created = 0; } } - // In RDMA polling mode, all messages must be executed in a new bthread and - // not in the bthread where the polling bthread is located, because the - // method for processing messages may call synchronization primitives, - // causing the polling bthread to be scheduled out. - if (m->_socket_mode == SOCKET_MODE_RDMA || m->_socket_mode == SOCKET_MODE_UBRING) { + // These transports may deliver messages from completion pollers. Process + // the messages in another bthread so user code cannot block the poller. + if (m->_socket_mode == SOCKET_MODE_RDMA || + m->_socket_mode == SOCKET_MODE_UBRING || + m->_socket_mode == SOCKET_MODE_URMA) { m->_transport->QueueMessage(last_msg, &num_bthread_created, true); } if (num_bthread_created) { diff --git a/src/brpc/input_messenger.h b/src/brpc/input_messenger.h index e82ecd4a66..a0163d0698 100644 --- a/src/brpc/input_messenger.h +++ b/src/brpc/input_messenger.h @@ -29,6 +29,11 @@ namespace brpc { namespace rdma { class RdmaEndpoint; } + +namespace urma { +class UrmaEndpoint; +} + namespace ubring { class UBShmEndpoint; } @@ -98,6 +103,7 @@ friend class Socket; friend class TcpTransport; friend class RdmaTransport; friend class rdma::RdmaEndpoint; +friend class urma::UrmaEndpoint; friend class ubring::UBShmEndpoint; public: explicit InputMessenger(size_t capacity = 128); diff --git a/src/brpc/server.h b/src/brpc/server.h index 4fbe304fde..08083037ac 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -224,7 +224,7 @@ struct ServerOptions { // Force ssl for all connections of the port to Start(). bool force_ssl; - // the server socket mode uses tcp or rdma or other + // Transport used by accepted sockets. // Default: SOCKET_MODE_TCP SocketMode socket_mode; diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 7c53058941..547f14e512 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -61,10 +61,21 @@ class RdmaHandshakeServerV2; class RdmaHandshakeClientV3; class RdmaHandshakeServerV3; } + +namespace urma { +class UrmaEndpoint; +class UrmaConnect; +class UrmaHandshakeClientV2; +class UrmaHandshakeServerV2; +class UrmaHandshakeClientV3; +class UrmaHandshakeServerV3; +} + namespace ubring { class UBShmEndpoint; class UBConnect; } + class Socket; class AuthContext; class EventDispatcher; @@ -334,6 +345,12 @@ 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; +friend class urma::UrmaHandshakeServerV2; +friend class urma::UrmaHandshakeClientV3; +friend class urma::UrmaHandshakeServerV3; friend class HealthCheckTask; friend class OnAppHealthCheckDone; friend class HealthCheckManager; @@ -344,6 +361,7 @@ friend void DereferenceSocket(Socket*); friend class Transport; friend class TcpTransport; friend class RdmaTransport; +friend class UrmaTransport; friend class TransportFactory; class SharedPart; struct WriteRequest; diff --git a/src/brpc/socket_mode.h b/src/brpc/socket_mode.h index b4ac7dfbca..1ea12a57e3 100644 --- a/src/brpc/socket_mode.h +++ b/src/brpc/socket_mode.h @@ -21,7 +21,9 @@ namespace brpc { enum SocketMode { SOCKET_MODE_TCP = 0, SOCKET_MODE_RDMA = 1, - SOCKET_MODE_UBRING = 2 + SOCKET_MODE_UBRING = 2, + SOCKET_MODE_URMA = 3 }; -} // namespace brpc -#endif //BRPC_SOCKET_MODE_H \ No newline at end of file +} // namespace brpc + +#endif // BRPC_SOCKET_MODE_H diff --git a/src/brpc/transport_factory.cpp b/src/brpc/transport_factory.cpp index 36fdaaed05..31bb168801 100644 --- a/src/brpc/transport_factory.cpp +++ b/src/brpc/transport_factory.cpp @@ -16,29 +16,35 @@ // under the License. #include "brpc/transport_factory.h" -#include "brpc/tcp_transport.h" #include "brpc/rdma_transport.h" +#include "brpc/tcp_transport.h" #include "brpc/ubshm_transport.h" +#include "brpc/urma_transport.h" namespace brpc { -int TransportFactory::ContextInitOrDie(SocketMode mode, bool serverOrNot, const void* _options) { + +int TransportFactory::ContextInitOrDie( + SocketMode mode, bool server_or_not, const void* options) { if (mode == SOCKET_MODE_TCP) { return 0; } #if BRPC_WITH_RDMA - else if (mode == SOCKET_MODE_RDMA) { - return RdmaTransport::ContextInitOrDie(serverOrNot, _options); + if (mode == SOCKET_MODE_RDMA) { + return RdmaTransport::ContextInitOrDie(server_or_not, options); } #endif -#if BRPC_WITH_UBRING - else if (mode == SOCKET_MODE_UBRING) { - return UBShmTransport::ContextInitOrDie(serverOrNot, _options); +#if BRPC_WITH_URMA + if (mode == SOCKET_MODE_URMA) { + return UrmaTransport::ContextInitOrDie(server_or_not, options); } #endif - else { - LOG(ERROR) << "unknown transport type " << mode; - return 1; +#if BRPC_WITH_UBRING + if (mode == SOCKET_MODE_UBRING) { + return UBShmTransport::ContextInitOrDie(server_or_not, options); } +#endif + LOG(ERROR) << "Unknown transport type " << mode; + return 1; } std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { @@ -46,18 +52,22 @@ std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { return std::unique_ptr(new TcpTransport()); } #if BRPC_WITH_RDMA - else if (mode == SOCKET_MODE_RDMA) { + if (mode == SOCKET_MODE_RDMA) { return std::unique_ptr(new RdmaTransport()); } #endif +#if BRPC_WITH_URMA + if (mode == SOCKET_MODE_URMA) { + return std::unique_ptr(new UrmaTransport()); + } +#endif #if BRPC_WITH_UBRING - else if (mode == SOCKET_MODE_UBRING) { + if (mode == SOCKET_MODE_UBRING) { return std::unique_ptr(new UBShmTransport()); } #endif - else { - LOG(ERROR) << "socket_mode set error"; - return nullptr; - } + LOG(ERROR) << "Unknown transport type " << mode; + return nullptr; } -} // namespace brpc \ No newline at end of file + +} // namespace brpc diff --git a/src/brpc/transport_factory.h b/src/brpc/transport_factory.h index d933a130e1..84b047daac 100644 --- a/src/brpc/transport_factory.h +++ b/src/brpc/transport_factory.h @@ -22,13 +22,16 @@ #include "brpc/transport.h" namespace brpc { -// TransportFactory to create transport instance with socket_mode {TCP, RDMA} + +// Creates transport instances for a SocketMode. class TransportFactory { public: - static int ContextInitOrDie(SocketMode mode, bool serverOrNot, const void* _options); + static int ContextInitOrDie(SocketMode mode, bool server_or_not, + const void* options); // Create transport instance with socket mode. static std::unique_ptr CreateTransport(SocketMode mode); }; -} // namespace brpc -#endif //BRPC_TRANSPORT_FACTORY_H \ No newline at end of file +} // namespace brpc + +#endif // BRPC_TRANSPORT_FACTORY_H diff --git a/src/brpc/urma/mock_urma.cpp b/src/brpc/urma/mock_urma.cpp new file mode 100644 index 0000000000..266c08b19a --- /dev/null +++ b/src/brpc/urma/mock_urma.cpp @@ -0,0 +1,713 @@ +// 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. + +// Link-time mock of the URMA user-level C API, modeled on Mooncake's +// mock_urma.cpp. It compiles against upstream UMDK headers and is linked into +// brpc when liburma.so is unavailable. It lets CI run UrmaTransport unit tests +// without URMA hardware. +// +// Design (same as Mooncake): +// - Link-time substitution: the symbols are literally named urma_create_jfc +// etc.; the linker picks this TU when liburma is absent. +// - Per-object state is held in anonymous-namespace maps keyed on the opaque +// pointer returned to the caller. Membership is checked on delete so misuse +// returns URMA_EINVAL rather than crashes. +// - Completion path: posted recv WRs remain pending until a SEND targets the +// corresponding mock jetty. Payload and immediate data are copied into the +// remote recv WR and both local-send and remote-recv completions are queued. +// - Device-name contract: device->name == "mock_urma_device" so tests can +// match it with --urma_device=mock_urma_device. + +#if BRPC_WITH_URMA + +#include "urma_api.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct JfcState { + std::mutex mutex; + std::deque completions; + bool event_pending{false}; +}; + +struct PendingRecv { + uint64_t addr; + uint32_t len; + uint64_t user_ctx; +}; + +std::shared_mutex g_rw_mutex; +bool initialized = false; +std::vector device_list; +std::map context_map; +std::map jfce_map; +std::map jfc_state_map; +std::map jfr_map; +// Side-table: JFR -> the JFC it was created with (used to route recv +// completions to the right JfcState, since the JFR is an opaque handle). +std::map jfr_jfc_map; +std::map> jfr_recv_map; +std::map seg_map; +std::map jetty_map; +std::map jetty_id_map; +std::map target_jetty_map; +std::atomic next_jetty_id{1}; + +void PushCompletion(urma_jfc_t* jfc, JfcState* state, + const urma_cr_t& completion) { + bool signal = false; + { + std::lock_guard lock(state->mutex); + state->completions.push_back(completion); + if (!state->event_pending) { + state->event_pending = true; + signal = true; + } + } + if (signal && jfc && jfc->jfc_cfg.jfce && + jfc->jfc_cfg.jfce->fd >= 0) { + uint64_t one = 1; + (void)write(jfc->jfc_cfg.jfce->fd, &one, sizeof(one)); + } +} + +urma_device_attr_t mock_device_attr = { + .guid = {.raw = {10}}, + .dev_cap = {}, + .port_cnt = 1, + .port_attr = {{.max_mtu = URMA_MTU_4096, + .state = URMA_PORT_ACTIVE, + .active_width = URMA_LINK_X1, + .active_speed = URMA_SP_100G, + .active_mtu = URMA_MTU_4096}}, + .reserved_jetty_id_min = 0, + .reserved_jetty_id_max = 1024}; + +urma_eid_info_t mock_eid_info = { + .eid = {{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, + 0x0C, 0x0D, 0x0E, 0x0F, 0x10}}, + .eid_index = 0}; + +} // namespace + +extern "C" { + +urma_status_t urma_init(urma_init_attr_t *init_attr) { + std::unique_lock lock(g_rw_mutex); + if (initialized) { + return URMA_EEXIST; + } + initialized = true; + return URMA_SUCCESS; +} + +urma_status_t urma_uninit(void) { + std::unique_lock lock(g_rw_mutex); + initialized = false; + for (auto device : device_list) { + delete device; + } + device_list.clear(); + context_map.clear(); + jfce_map.clear(); + for (auto &kv : jfc_state_map) { + delete kv.second; + } + jfc_state_map.clear(); + jfr_map.clear(); + jfr_jfc_map.clear(); + jfr_recv_map.clear(); + seg_map.clear(); + jetty_map.clear(); + jetty_id_map.clear(); + target_jetty_map.clear(); + next_jetty_id.store(1); + return URMA_SUCCESS; +} + +urma_device_t **urma_get_device_list(int *num_devices) { + { + std::shared_lock lock(g_rw_mutex); + if (!initialized) { + *num_devices = 0; + return nullptr; + } + if (!device_list.empty()) { + *num_devices = device_list.size(); + urma_device_t **devices = new urma_device_t *[device_list.size()]; + for (size_t i = 0; i < device_list.size(); ++i) { + devices[i] = device_list[i]; + } + return devices; + } + } + { + std::unique_lock write_lock(g_rw_mutex); + if (!initialized) { + *num_devices = 0; + return nullptr; + } + if (device_list.empty()) { + urma_device_t *device = new urma_device_t; + strcpy(device->name, "mock_urma_device"); + strcpy(device->path, "/sys/class/infiniband/mock_device"); + device->type = URMA_TRANSPORT_UB; + device->ops = nullptr; + device->sysfs_dev = nullptr; + device_list.push_back(device); + } + *num_devices = device_list.size(); + urma_device_t **devices = new urma_device_t *[device_list.size()]; + for (size_t i = 0; i < device_list.size(); ++i) { + devices[i] = device_list[i]; + } + return devices; + } +} + +urma_device_t *urma_get_device_by_name(char *dev_name) { + { + std::shared_lock lock(g_rw_mutex); + if (!initialized) { + return nullptr; + } + if (!device_list.empty()) { + for (auto device : device_list) { + if (strcmp(device->name, dev_name) == 0) { + return device; + } + } + return device_list[0]; + } + } + { + std::unique_lock write_lock(g_rw_mutex); + if (!initialized) { + return nullptr; + } + if (device_list.empty()) { + auto *device = new urma_device_t; + strcpy(device->name, "mock_urma_device"); + strcpy(device->path, "/sys/class/infiniband/mock_device"); + device->type = URMA_TRANSPORT_UB; + device->ops = nullptr; + device->sysfs_dev = nullptr; + device_list.push_back(device); + } + for (auto device : device_list) { + if (strcmp(device->name, dev_name) == 0) { + return device; + } + } + return device_list.empty() ? nullptr : device_list[0]; + } +} + +void urma_free_device_list(urma_device_t **device_list) { + if (device_list) { + delete[] device_list; + } +} + +urma_status_t urma_query_device(urma_device_t *device, + urma_device_attr_t *attr) { + if (!device || !attr) { + return URMA_EINVAL; + } + mock_device_attr.dev_cap.max_jfc = 1024; + mock_device_attr.dev_cap.max_jetty = 1024; + mock_device_attr.dev_cap.max_jfs_sge = 8; + mock_device_attr.dev_cap.max_jfr_sge = 8; + memcpy(attr, &mock_device_attr, sizeof(urma_device_attr_t)); + return URMA_SUCCESS; +} + +urma_eid_info_t *urma_get_eid_list(urma_device_t *device, uint32_t *eid_cnt) { + if (!device || !eid_cnt) { + return nullptr; + } + *eid_cnt = 1; + auto *eid_list = new urma_eid_info_t[1]; + memcpy(eid_list, &mock_eid_info, sizeof(urma_eid_info_t)); + return eid_list; +} + +void urma_free_eid_list(urma_eid_info_t *eid_list) { + if (eid_list) { + delete[] eid_list; + } +} + +urma_context_t *urma_create_context(urma_device_t *device, uint32_t eid_index) { + std::unique_lock lock(g_rw_mutex); + if (!device) { + return nullptr; + } + urma_context_t *ctx = new urma_context_t; + ctx->async_fd = 0; + ctx->dev = device; + context_map[ctx] = 1; + return ctx; +} + +urma_status_t urma_delete_context(urma_context_t *ctx) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || context_map.find(ctx) == context_map.end()) { + return URMA_EINVAL; + } + context_map.erase(ctx); + delete ctx; + return URMA_SUCCESS; +} + +urma_jfce_t *urma_create_jfce(urma_context_t *ctx) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || context_map.find(ctx) == context_map.end()) { + return nullptr; + } + // Allocate a real nonblocking eventfd so brpc's event-mode CQ socket can + // be exercised by the mock as well. + urma_jfce_t *jfce = new urma_jfce_t{}; + jfce->urma_ctx = ctx; + jfce->fd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + if (jfce->fd < 0) { + delete jfce; + return nullptr; + } + jfce_map[jfce] = 1; + return jfce; +} + +urma_status_t urma_delete_jfce(urma_jfce_t *jfce) { + std::unique_lock lock(g_rw_mutex); + if (!jfce || jfce_map.find(jfce) == jfce_map.end()) { + return URMA_EINVAL; + } + jfce_map.erase(jfce); + close(jfce->fd); + delete jfce; + return URMA_SUCCESS; +} + +urma_jfc_t *urma_create_jfc(urma_context_t *ctx, urma_jfc_cfg_t *cfg) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) { + return nullptr; + } + urma_jfc_t *jfc = new urma_jfc_t; + memset(&jfc->jfc_id.eid, 0, sizeof(urma_eid_t)); + jfc->jfc_id.eid.raw[0] = 1; + jfc->jfc_id.uasid = 0; + jfc->jfc_id.id = 1; + jfc->handle = cfg->user_ctx; + jfc->comp_events_acked = 0; + jfc->async_events_acked = 0; + jfc->jfc_cfg = *cfg; + jfc_state_map[jfc] = new JfcState(); + return jfc; +} + +urma_status_t urma_delete_jfc(urma_jfc_t *jfc) { + std::unique_lock lock(g_rw_mutex); + if (!jfc || jfc_state_map.find(jfc) == jfc_state_map.end()) { + return URMA_EINVAL; + } + delete jfc_state_map[jfc]; + jfc_state_map.erase(jfc); + delete jfc; + return URMA_SUCCESS; +} + +urma_jfr_t *urma_create_jfr(urma_context_t *ctx, urma_jfr_cfg_t *cfg) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) { + return nullptr; + } + // Opaque handle (avoids partially initializing urma_jfr_t's pthread + // members). brpc's endpoint reads back jfr_cfg.jfc from the stored cfg + // pointer below, so we keep a side-table mapping jfr -> cfg.jfc. + urma_jfr_t *jfr = reinterpret_cast(new int(1)); + jfr_map[jfr] = 1; + // Stash the JFC for urma_post_jfr_wr's completion routing. + jfr_jfc_map[jfr] = cfg->jfc; + jfr_recv_map[jfr] = {}; + return jfr; +} + +urma_status_t urma_delete_jfr(urma_jfr_t *jfr) { + std::unique_lock lock(g_rw_mutex); + if (!jfr || jfr_map.find(jfr) == jfr_map.end()) { + return URMA_EINVAL; + } + jfr_map.erase(jfr); + jfr_jfc_map.erase(jfr); + jfr_recv_map.erase(jfr); + delete reinterpret_cast(jfr); + return URMA_SUCCESS; +} + +urma_target_seg_t *urma_register_seg(urma_context_t *ctx, urma_seg_cfg_t *cfg) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) { + return nullptr; + } + urma_target_seg_t *seg = new urma_target_seg_t; + memset(&seg->seg.ubva.eid, 0, sizeof(urma_eid_t)); + seg->seg.ubva.eid.raw[0] = 1; + seg->seg.ubva.uasid = 0; + seg->seg.ubva.va = cfg->va; + seg->seg.len = cfg->len; + seg->seg.token_id = cfg->token_value.token; + seg_map[seg] = 1; + return seg; +} + +urma_status_t urma_unregister_seg(urma_target_seg_t *seg) { + std::unique_lock lock(g_rw_mutex); + if (!seg || seg_map.find(seg) == seg_map.end()) { + return URMA_EINVAL; + } + seg_map.erase(seg); + delete seg; + return URMA_SUCCESS; +} + +urma_target_seg_t *urma_import_seg(urma_context_t *ctx, urma_seg_t *seg, + urma_token_t *token_value, uint64_t addr, + urma_import_seg_flag_t flag) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !seg || !token_value || + context_map.find(ctx) == context_map.end()) { + return nullptr; + } + urma_target_seg_t *tseg = new urma_target_seg_t; + tseg->seg = *seg; + *token_value = {.token = seg->token_id}; + seg_map[tseg] = 1; + return tseg; +} + +urma_status_t urma_unimport_seg(urma_target_seg_t *tseg) { + std::unique_lock lock(g_rw_mutex); + if (!tseg || seg_map.find(tseg) == seg_map.end()) { + return URMA_EINVAL; + } + seg_map.erase(tseg); + delete tseg; + return URMA_SUCCESS; +} + +urma_status_t urma_get_async_event(urma_context_t *ctx, + urma_async_event_t *event) { + if (!ctx || !event) { + return URMA_EINVAL; + } + std::shared_lock lock(g_rw_mutex); + if (context_map.find(ctx) == context_map.end()) { + return URMA_EINVAL; + } + return URMA_ETIMEOUT; +} + +void urma_ack_async_event(urma_async_event_t *event) {} + +urma_jetty_t *urma_create_jetty(urma_context_t *ctx, urma_jetty_cfg_t *cfg) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !cfg || context_map.find(ctx) == context_map.end()) { + return nullptr; + } + urma_jetty_t *jetty = new urma_jetty_t; + memset(&jetty->jetty_id.eid, 0, sizeof(urma_eid_t)); + jetty->jetty_id.eid.raw[0] = 1; + jetty->jetty_id.uasid = 0; + jetty->jetty_id.id = next_jetty_id.fetch_add(1); + jetty->jetty_cfg = *cfg; + jetty->remote_jetty = nullptr; + jetty_map[jetty] = 1; + jetty_id_map[jetty->jetty_id.id] = jetty; + return jetty; +} + +urma_status_t urma_delete_jetty(urma_jetty_t *jetty) { + std::unique_lock lock(g_rw_mutex); + if (!jetty || jetty_map.find(jetty) == jetty_map.end()) { + return URMA_EINVAL; + } + jetty_id_map.erase(jetty->jetty_id.id); + jetty_map.erase(jetty); + delete jetty; + return URMA_SUCCESS; +} + +urma_status_t urma_unbind_jetty(urma_jetty_t *jetty) { + std::unique_lock lock(g_rw_mutex); + if (!jetty || jetty_map.find(jetty) == jetty_map.end()) { + return URMA_EINVAL; + } + jetty->remote_jetty = nullptr; + return URMA_SUCCESS; +} + +urma_target_jetty_t *urma_import_jetty(urma_context_t *ctx, + urma_rjetty_t *rjetty, + urma_token_t *token_value) { + std::unique_lock lock(g_rw_mutex); + if (!ctx || !rjetty || !token_value || + context_map.find(ctx) == context_map.end()) { + return nullptr; + } + urma_target_jetty_t *tjetty = new urma_target_jetty_t; + tjetty->id = rjetty->jetty_id; + target_jetty_map[tjetty] = 1; + *token_value = {.token = 1}; + return tjetty; +} + +urma_status_t urma_unimport_jetty(urma_target_jetty_t *tjetty) { + std::unique_lock lock(g_rw_mutex); + if (!tjetty || target_jetty_map.find(tjetty) == target_jetty_map.end()) { + return URMA_EINVAL; + } + target_jetty_map.erase(tjetty); + delete tjetty; + return URMA_SUCCESS; +} + +urma_status_t urma_bind_jetty(urma_jetty_t *jetty, + urma_target_jetty_t *tjetty) { + std::unique_lock lock(g_rw_mutex); + if (!jetty || !tjetty || jetty_map.find(jetty) == jetty_map.end() || + target_jetty_map.find(tjetty) == target_jetty_map.end()) { + return URMA_EINVAL; + } + jetty->remote_jetty = tjetty; + return URMA_SUCCESS; +} + +urma_status_t urma_modify_jetty(urma_jetty_t *jetty, urma_jetty_attr_t *attr) { + std::shared_lock lock(g_rw_mutex); + if (!jetty || !attr || jetty_map.find(jetty) == jetty_map.end()) { + return URMA_EINVAL; + } + return URMA_SUCCESS; +} + +urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr, + urma_jfs_wr_t **bad_wr) { + std::shared_lock read_lock(g_rw_mutex); + auto local_it = jetty_map.find(jetty); + auto local_jfc_it = + jetty ? jfc_state_map.find(jetty->jetty_cfg.jfs_cfg.jfc) + : jfc_state_map.end(); + if (!jetty || !wr || local_it == jetty_map.end() || + local_jfc_it == jfc_state_map.end()) { + if (bad_wr) { + *bad_wr = wr; + } + return URMA_EINVAL; + } + JfcState* local_state = local_jfc_it->second; + read_lock.unlock(); + + for (urma_jfs_wr_t* current = wr; current; current = current->next) { + if (!current->tjetty) { + if (bad_wr) { + *bad_wr = current; + } + return URMA_EINVAL; + } + + urma_cr_t send_cr{}; + send_cr.status = URMA_CR_SUCCESS; + send_cr.user_ctx = current->user_ctx; + send_cr.flag.bs.s_r = 0; + if (current->flag.bs.complete_enable) { + PushCompletion(jetty->jetty_cfg.jfs_cfg.jfc, + local_state, send_cr); + } + + PendingRecv recv{}; + urma_jfc_t* remote_jfc = nullptr; + { + std::unique_lock lock(g_rw_mutex); + auto remote_it = jetty_id_map.find(current->tjetty->id.id); + if (remote_it == jetty_id_map.end()) { + continue; + } + urma_jfr_t* remote_jfr = + remote_it->second->jetty_cfg.shared.jfr; + auto recv_it = jfr_recv_map.find(remote_jfr); + auto jfc_it = jfr_jfc_map.find(remote_jfr); + if (recv_it == jfr_recv_map.end() || recv_it->second.empty() || + jfc_it == jfr_jfc_map.end()) { + continue; + } + recv = recv_it->second.front(); + recv_it->second.pop_front(); + remote_jfc = jfc_it->second; + } + + uint32_t copied = 0; + for (uint32_t i = 0; + i < current->send.src.num_sge && copied < recv.len; ++i) { + const urma_sge_t& sge = current->send.src.sge[i]; + const uint32_t n = std::min(sge.len, recv.len - copied); + std::memcpy(reinterpret_cast(recv.addr + copied), + reinterpret_cast(sge.addr), n); + copied += n; + } + + JfcState* remote_state = nullptr; + { + std::shared_lock lock(g_rw_mutex); + auto state_it = jfc_state_map.find(remote_jfc); + if (state_it != jfc_state_map.end()) { + remote_state = state_it->second; + } + } + if (remote_state) { + urma_cr_t recv_cr{}; + recv_cr.status = URMA_CR_SUCCESS; + recv_cr.user_ctx = recv.user_ctx; + recv_cr.flag.bs.s_r = 1; + recv_cr.completion_len = copied; + recv_cr.opcode = + current->opcode == URMA_OPC_SEND_IMM + ? URMA_CR_OPC_SEND_WITH_IMM + : URMA_CR_OPC_SEND; + recv_cr.imm_data = current->send.imm_data; + PushCompletion(remote_jfc, remote_state, recv_cr); + } + } + if (bad_wr) { + *bad_wr = nullptr; + } + return URMA_SUCCESS; +} + +urma_status_t urma_post_jfr_wr(urma_jfr_t *jfr, urma_jfr_wr_t *wr, + urma_jfr_wr_t **bad_wr) { + std::unique_lock lock(g_rw_mutex); + auto recv_it = jfr_recv_map.find(jfr); + if (!jfr || !wr || recv_it == jfr_recv_map.end()) { + if (bad_wr) { + *bad_wr = wr; + } + return URMA_EINVAL; + } + for (urma_jfr_wr_t* current = wr; current; current = current->next) { + if (current->src.num_sge == 0 || !current->src.sge) { + if (bad_wr) { + *bad_wr = current; + } + return URMA_EINVAL; + } + recv_it->second.push_back(PendingRecv{ + current->src.sge[0].addr, + current->src.sge[0].len, + current->user_ctx}); + } + if (bad_wr) { + *bad_wr = nullptr; + } + return URMA_SUCCESS; +} + +urma_status_t urma_post_jetty_recv_wr(urma_jetty_t *jetty, + urma_jfr_wr_t *wr, + urma_jfr_wr_t **bad_wr) { + urma_jfr_t* shared_jfr = nullptr; + { + std::shared_lock lock(g_rw_mutex); + if (!jetty || jetty_map.find(jetty) == jetty_map.end()) { + if (bad_wr) { + *bad_wr = wr; + } + return URMA_EINVAL; + } + shared_jfr = jetty->jetty_cfg.shared.jfr; + } + return urma_post_jfr_wr(shared_jfr, wr, bad_wr); +} + +int urma_poll_jfc(urma_jfc_t *jfc, int num_entries, urma_cr_t *cr_list) { + JfcState* state = nullptr; + { + std::shared_lock lock(g_rw_mutex); + auto it = jfc_state_map.find(jfc); + if (it == jfc_state_map.end()) { + return -1; + } + state = it->second; + } + std::lock_guard lock(state->mutex); + int count = 0; + while (count < num_entries && !state->completions.empty()) { + cr_list[count++] = state->completions.front(); + state->completions.pop_front(); + } + if (state->completions.empty()) { + state->event_pending = false; + } + return count; +} + +urma_status_t urma_rearm_jfc(urma_jfc_t*, bool) { + return URMA_SUCCESS; +} + +int urma_wait_jfc(urma_jfce_t* jfce, uint32_t jfc_cnt, int, + urma_jfc_t* jfcs[]) { + if (!jfce || !jfcs || jfc_cnt == 0) { + errno = EINVAL; + return -1; + } + uint64_t value = 0; + (void)read(jfce->fd, &value, sizeof(value)); + std::shared_lock lock(g_rw_mutex); + uint32_t count = 0; + for (const auto& item : jfc_state_map) { + if (count >= jfc_cnt || item.first->jfc_cfg.jfce != jfce) { + continue; + } + std::lock_guard state_lock(item.second->mutex); + if (item.second->event_pending) { + jfcs[count++] = item.first; + item.second->event_pending = false; + } + } + return static_cast(count); +} + +void urma_ack_jfc(urma_jfc_t*[], uint32_t[], uint32_t) { +} + +} // extern "C" + +#endif // BRPC_WITH_URMA diff --git a/src/brpc/urma/urma_bonding.h b/src/brpc/urma/urma_bonding.h new file mode 100644 index 0000000000..6332db8495 --- /dev/null +++ b/src/brpc/urma/urma_bonding.h @@ -0,0 +1,35 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef BRPC_URMA_URMA_BONDING_H +#define BRPC_URMA_URMA_BONDING_H + +// urma_ubagg.h is a provider-private extension and is not shipped by every +// UMDK installation. Keep the dependency optional so non-bonding devices and +// mock builds continue to work. +#if defined(__has_include) +#if __has_include("urma_ubagg.h") +#include "urma_ubagg.h" +#define BRPC_URMA_HAS_BONDING_EXT 1 +#endif +#endif + +#ifndef BRPC_URMA_HAS_BONDING_EXT +#define BRPC_URMA_HAS_BONDING_EXT 0 +#endif + +#endif // BRPC_URMA_URMA_BONDING_H diff --git a/src/brpc/urma/urma_endpoint.cpp b/src/brpc/urma/urma_endpoint.cpp new file mode 100644 index 0000000000..bb8914d7d2 --- /dev/null +++ b/src/brpc/urma/urma_endpoint.cpp @@ -0,0 +1,1846 @@ +// 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/urma/urma_endpoint.h" + +#if BRPC_WITH_URMA + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "butil/atomicops.h" +#include "butil/iobuf.h" +#include "butil/logging.h" +#include "butil/macros.h" +#include "butil/sys_byteorder.h" +#include "butil/time.h" +#include "bthread/bthread.h" +#include "bthread/butex.h" + +#include "urma_api.h" + +#include "brpc/input_messenger.h" +#include "brpc/socket.h" +#include "brpc/urma/urma_bonding.h" +#include "brpc/urma/urma_handshake.h" +#include "brpc/urma/urma_handshake.pb.h" +#include "brpc/urma/urma_helper.h" +#include "brpc/urma_transport.h" + +DECLARE_int32(task_group_ntags); + +namespace brpc { +namespace urma { + +// Flags used here are declared in urma_endpoint.h (urma_use_polling, +// urma_poller_num, urma_disable_bthread). Declare the rest here. +DECLARE_int32(urma_sq_size); +DECLARE_int32(urma_rq_size); +DECLARE_int32(urma_cqe_poll_once); +DECLARE_bool(urma_recv_zerocopy); +DECLARE_int32(urma_zerocopy_min_size); +DECLARE_int32(urma_prepared_jetty_cnt); +DECLARE_bool(urma_poller_yield); + +// ---- Constants shared with the handshake module ---- +static const int WAIT_TIMEOUT_MS = 50; +static const size_t HELLO_ACK_LEN = 4; +static const uint32_t HELLO_ACK_URMA_OK = 0x1; +static const size_t IOBUF_BLOCK_HEADER_LEN = 32; // matches butil IOBuf + +// ---- Globals: prepared jetty pool + poller groups ---- +struct PreparedJetty { + UrmaResource* res; +}; +static butil::Mutex g_prepared_mutex; +static UrmaResource* g_prepared_list = nullptr; // singly-linked +static int g_prepared_cnt = 0; + +static int PreparedJettyCount() { + const int requested = + std::max(0, std::min(FLAGS_urma_prepared_jetty_cnt, 1024)); + if (requested == 0) { + return 0; + } + + struct rlimit nofile; + if (getrlimit(RLIMIT_NOFILE, &nofile) != 0 || + nofile.rlim_cur == RLIM_INFINITY) { + return requested; + } + + // In event mode each prepared JFCE consumes a file descriptor. Keep room + // for one TCP fd per future URMA connection and for brpc/system internals. + static const rlim_t kReservedFdCount = 64; + const rlim_t max_prepared = + nofile.rlim_cur > kReservedFdCount + ? (nofile.rlim_cur - kReservedFdCount) / 2 + : 0; + if (max_prepared >= static_cast(requested)) { + return requested; + } + + LOG(WARNING) << "Cap URMA prepared jetty count from " << requested + << " to " << max_prepared + << " due to RLIMIT_NOFILE=" << nofile.rlim_cur; + return static_cast(max_prepared); +} + +std::vector UrmaEndpoint::_poller_groups; + +// ============================================================================ +// UrmaResource lifecycle. +// ============================================================================ + +UrmaResource::~UrmaResource() { + if (remote_jetty) { + urma_unimport_jetty(remote_jetty); + } + if (remote_seg) { + urma_unimport_seg(remote_seg); + } + if (jetty) { + urma_delete_jetty(jetty); + } + if (jfr) { + urma_delete_jfr(jfr); + } + if (jfc) { + urma_delete_jfc(jfc); + } + if (jfce) { + urma_delete_jfce(jfce); + } +} + +// ============================================================================ +// Constructor / destructor / Reset. +// ============================================================================ + +UrmaEndpoint::UrmaEndpoint(Socket* s) + : _socket(s), + _state(UNINIT), + _handshake_version(0), + _resource(nullptr) { + _sq_size = static_cast( + std::max(16, std::min(4096, static_cast(FLAGS_urma_sq_size)))); + _rq_size = static_cast( + std::max(16, std::min(4096, static_cast(FLAGS_urma_rq_size)))); + _read_butex = bthread::butex_create_checked>(); + _read_butex->store(0, butil::memory_order_relaxed); +} + +UrmaEndpoint::~UrmaEndpoint() { + DeallocateResources(); + if (_read_butex) { + bthread::butex_destroy(_read_butex); + _read_butex = nullptr; + } +} + +void UrmaEndpoint::Reset() { + DeallocateResources(); + _state = UNINIT; + _handshake_version = 0; + _remote_recv_block_size = 0; + _local_window_capacity = 0; + _remote_window_capacity = 0; + _remote_rq_window_size.store(0, butil::memory_order_relaxed); + _sq_window_size.store(0, butil::memory_order_relaxed); + _new_rq_wrs.store(0, butil::memory_order_relaxed); + _sq_imm_window_size = 0; + _sq_current = 0; + _sq_sent = 0; + _rq_received = 0; + _pending_received_bytes.store(0, butil::memory_order_relaxed); + _sbuf.clear(); + _rbuf.clear(); + _rbuf_data.clear(); + _read_butex->store(0, butil::memory_order_relaxed); +} + +// ============================================================================ +// Handshake IO helpers (ReadFromFd / WriteToFd / PushBackToReadBuf). +// Modeled on RdmaEndpoint::ReadFromFdLoop / WriteToFdLoop. +// ============================================================================ + +int UrmaEndpoint::ReadFromFd(void* data, size_t len) { + char* p = static_cast(data); + size_t received = 0; + while (received < len) { + const int expected_val = _read_butex->load(butil::memory_order_acquire); + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nr = read(fd, p + received, len - received); + if (nr < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) { + int rc = bthread::butex_wait(_read_butex, expected_val, &duetime); + if (rc < 0 && errno != EWOULDBLOCK && errno != ETIMEDOUT) { + return -1; + } + continue; + } + return -1; + } + if (nr == 0) { + errno = EEOF; + return -1; + } + received += nr; + } + return 0; +} + +void UrmaEndpoint::PushBackToReadBuf(const void* data, size_t len) { + _socket->_read_buf.append(data, len); +} + +int UrmaEndpoint::WriteToFd(void* data, size_t len) { + char* p = static_cast(data); + size_t written = 0; + while (written < len) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = write(fd, p + written, len - written); + if (nw >= 0) { + written += nw; + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +// ============================================================================ +// Hello builders / parsers. +// ============================================================================ + +void UrmaEndpoint::MakeLocalParsedHello(ParsedHello* out) const { + *out = ParsedHello{}; // value-initialize (avoids memset on non-trivial type) + out->buffer_size = static_cast(GetUrmaRecvBlockSize()); + out->recv_buffer_cnt = _rq_size - 1; + if (_resource && _resource->jetty) { + out->jetty_id = _resource->jetty->jetty_id.id; + out->uasid = _resource->jetty->jetty_id.uasid; + const urma_eid_t* local_eid = GetUrmaLocalEid(); + const uint8_t* advertised_eid = + local_eid != nullptr + ? local_eid->raw + : _resource->jetty->jetty_id.eid.raw; + std::memcpy(out->eid, advertised_eid, 16); + } + out->tp_type = static_cast(URMA_CTP); + // Pool segment: flatten g_pool_seg's seg fields. + urma_target_seg_t* pool = GetPoolSegFor(nullptr); + if (pool) { + std::memcpy(out->seg_eid, pool->seg.ubva.eid.raw, 16); + out->seg_uasid = pool->seg.ubva.uasid; + out->seg_va = pool->seg.ubva.va; + out->seg_len = pool->seg.len; + out->seg_token_id = pool->seg.token_id; + } +} + +void UrmaEndpoint::FillLocalHelloV2(v2_wire::HelloMessage* out) const { + *out = v2_wire::HelloMessage{}; // value-initialize + out->msg_len = v2_wire::HELLO_PACKET_LEN; + out->hello_ver = v2_wire::HELLO_V2_VERSION; + out->impl_ver = v2_wire::IMPL_V2_VERSION; + ParsedHello p; + MakeLocalParsedHello(&p); + out->buffer_size = p.buffer_size; + out->recv_buffer_cnt = p.recv_buffer_cnt; + out->jetty_id = p.jetty_id; + std::memcpy(out->eid, p.eid, 16); + out->uasid = p.uasid; + out->tp_type = p.tp_type; + std::memcpy(out->seg_eid, p.seg_eid, 16); + out->seg_uasid = p.seg_uasid; + out->seg_va = p.seg_va; + out->seg_len = p.seg_len; + out->seg_token_id = p.seg_token_id; +} + +void UrmaEndpoint::FillLocalHelloV3(UrmaHello* out) const { + ParsedHello p; + MakeLocalParsedHello(&p); + out->set_buffer_size(p.buffer_size); + out->set_recv_buffer_cnt(p.recv_buffer_cnt); + out->set_jetty_id(p.jetty_id); + out->set_eid(p.eid, 16); + out->set_uasid(p.uasid); + out->set_tp_type(p.tp_type); + out->set_seg_eid(p.seg_eid, 16); + out->set_seg_uasid(p.seg_uasid); + out->set_seg_va(p.seg_va); + out->set_seg_len(p.seg_len); + out->set_seg_token_id(p.seg_token_id); +} + +int UrmaEndpoint::WriteHelloV3(const UrmaHello& msg) { + butil::IOBuf packet; + packet.append("URM3", 4); + std::string body; + if (!msg.SerializeToString(&body)) { + LOG(ERROR) << "Fail to serialize UrmaHello"; + return -1; + } + uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); + packet.append(&pb_size_be, sizeof(pb_size_be)); + packet.append(body); + return WriteToFd(packet); +} + +int UrmaEndpoint::WriteToFd(butil::IOBuf& data) { + // Write out the IOBuf in a single WriteToFd-style loop. + while (!data.empty()) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = data.cut_into_file_descriptor(fd); + if (nw >= 0) { + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +int UrmaEndpoint::ReadAndParseHelloV3(ParsedHello* out, bool* negotiated) { + *negotiated = false; + uint32_t pb_size_be = 0; + if (ReadFromFd(&pb_size_be, sizeof(pb_size_be)) < 0) { + return -1; + } + const uint32_t pb_size = butil::NetToHost32(pb_size_be); + if (pb_size == 0 || pb_size > 4096) { + return 0; + } + std::string body(pb_size, '\0'); + if (ReadFromFd(&body[0], pb_size) < 0) { + return -1; + } + UrmaHello msg; + if (!msg.ParseFromArray(body.data(), static_cast(body.size()))) { + return 0; + } + if (msg.eid().size() != 16 || msg.seg_eid().size() != 16) { + return 0; + } + out->buffer_size = msg.buffer_size(); + out->recv_buffer_cnt = msg.recv_buffer_cnt(); + out->jetty_id = msg.jetty_id(); + std::memcpy(out->eid, msg.eid().data(), 16); + out->uasid = msg.uasid(); + out->tp_type = static_cast(msg.tp_type()); + std::memcpy(out->seg_eid, msg.seg_eid().data(), 16); + out->seg_uasid = msg.seg_uasid(); + out->seg_va = msg.seg_va(); + out->seg_len = msg.seg_len(); + out->seg_token_id = msg.seg_token_id(); + if (!ValidHello(*out)) { + return 0; + } + *negotiated = true; + return 0; +} + +// ============================================================================ +// Allocate / deallocate per-connection resources. +// ============================================================================ + +int UrmaEndpoint::AllocateResources() { + if (_resource) { + return 0; + } + urma_context_t* ctx = GetUrmaContext(); + if (!ctx) { + errno = ENODEV; + return -1; + } + + _resource = new (std::nothrow) UrmaResource(); + if (!_resource) { + return -1; + } + + // Try the prepared pool first (sized sq/rq match). + if (_sq_size <= static_cast(FLAGS_urma_sq_size) && + _rq_size <= static_cast(FLAGS_urma_rq_size)) { + BAIDU_SCOPED_LOCK(g_prepared_mutex); + if (g_prepared_list) { + UrmaResource* next = g_prepared_list->next; + delete _resource; + _resource = g_prepared_list; + g_prepared_list = next; + _resource->next = nullptr; + --g_prepared_cnt; + } + } + + if (!_resource->jfc) { + // The SDK requires every JFC to reference a JFCE. Polling mode does + // not arm or consume it, but still supplies the required object. + _resource->jfce = urma_create_jfce(ctx); + if (!_resource->jfce || + (!FLAGS_urma_use_polling && _resource->jfce->fd < 0)) { + LOG(ERROR) << "Fail to create a usable URMA JFCE"; + errno = ENODEV; + return -1; + } + + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = static_cast(_sq_size + _rq_size); + jfc_cfg.jfce = _resource->jfce; + _resource->jfc = urma_create_jfc(ctx, &jfc_cfg); + if (!_resource->jfc) { + PLOG(ERROR) << "urma_create_jfc"; + return -1; + } + + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = static_cast(_rq_size); + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + jfr_cfg.jfc = _resource->jfc; + _resource->jfr = urma_create_jfr(ctx, &jfr_cfg); + if (!_resource->jfr) { + PLOG(ERROR) << "urma_create_jfr"; + return -1; + } + + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = static_cast(_sq_size); + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.priority = GetUrmaJettyPriority(); + jetty_cfg.jfs_cfg.max_sge = + static_cast(GetUrmaMaxSge()); + jetty_cfg.jfs_cfg.rnr_retry = URMA_TYPICAL_RNR_RETRY; + jetty_cfg.jfs_cfg.err_timeout = URMA_TYPICAL_ERR_TIMEOUT; + jetty_cfg.jfs_cfg.jfc = _resource->jfc; + jetty_cfg.shared.jfr = _resource->jfr; + jetty_cfg.shared.jfc = _resource->jfc; + _resource->jetty = urma_create_jetty(ctx, &jetty_cfg); + if (!_resource->jetty) { + PLOG(ERROR) << "urma_create_jetty"; + return -1; + } + } + + _sbuf.resize(_sq_size - RESERVED_WR_NUM); + _rbuf.resize(_rq_size); + _rbuf_data.resize(_rq_size, nullptr); + + // Wrap the JFCE fd in a brpc Socket so PollCq is driven by epoll. + if (!FLAGS_urma_use_polling) { + if (!_resource->jfce || _resource->jfce->fd < 0) { + LOG(ERROR) << "Prepared URMA resource has no usable JFCE"; + errno = ENODEV; + return -1; + } + if (ReqNotifyCq() != 0) { + return -1; + } + SocketOptions options; + options.user = this; + options.keytable_pool = _socket->keytable_pool(); + options.fd = _resource->jfce->fd; + options.on_edge_triggered_events = PollCq; + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(ERROR) << "Fail to create CQ socket"; + return -1; + } + } else { + // Polling mode: synthetic carrier socket (no fd). + SocketOptions options; + options.user = this; + options.keytable_pool = _socket->keytable_pool(); + options.on_edge_triggered_events = PollCq; + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(ERROR) << "Fail to create CQ socket (polling)"; + return -1; + } + PollerAddCqSid(); + } + return 0; +} + +void UrmaEndpoint::DeallocateResources() { + if (!_resource) { + return; + } + + if (FLAGS_urma_use_polling) { + PollerRemoveCqSid(); + } + + // Tear down the CQ socket so the EventDispatcher stops calling PollCq. + if (_cq_sid != INVALID_SOCKET_ID) { + SocketUniquePtr s; + if (Socket::Address(_cq_sid, &s) == 0) { + if (s->fd() >= 0) { + s->_io_event.RemoveConsumer(s->_fd); + } + s->_user = nullptr; // Do not release user (this UrmaEndpoint). + s->_fd = -1; // Already removed fd from epoll. + s->SetFailed(); + } + _cq_sid = INVALID_SOCKET_ID; + } + + // Reusing a Jetty requires a driver-supported RESET plus a complete JFC + // drain. Until that lifecycle is implemented, prepared resources are + // one-shot: they accelerate connection setup but are destroyed on close. + delete _resource; + _resource = nullptr; +} + +// ============================================================================ +// ImportPeer: the critical import_seg-before-import_jetty sequence. +// ============================================================================ + +int UrmaEndpoint::ImportPeer(const ParsedHello& peer) { + urma_context_t* ctx = GetUrmaContext(); + if (!ctx) { + errno = ENODEV; + return -1; + } + + // 1. urma_import_seg FIRST so the kernel establishes TP routing for the + // remote EID. Without this the first SEND is rejected by hardware with + // URMA_CR_RNR_RETRY_CNT_EXC_ERR. + urma_seg_t peer_seg{}; + std::memcpy(peer_seg.ubva.eid.raw, peer.seg_eid, 16); + peer_seg.ubva.uasid = peer.seg_uasid; + peer_seg.ubva.va = peer.seg_va; + peer_seg.len = peer.seg_len; + peer_seg.token_id = peer.seg_token_id; + urma_token_t seg_token{}; + urma_import_seg_flag_t seg_flag{}; + seg_flag.bs.cacheable = URMA_NON_CACHEABLE; + seg_flag.bs.access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC; + seg_flag.bs.mapping = URMA_SEG_NOMAP; + _resource->remote_seg = urma_import_seg(ctx, &peer_seg, &seg_token, 0, seg_flag); + if (!_resource->remote_seg) { + PLOG(ERROR) << "urma_import_seg failed"; + return -1; + } + + // 2. urma_import_jetty. + urma_rjetty_t remote{}; + std::memcpy(remote.jetty_id.eid.raw, peer.eid, 16); + remote.jetty_id.uasid = peer.uasid; + remote.jetty_id.id = peer.jetty_id; + remote.trans_mode = URMA_TM_RM; + remote.type = URMA_JETTY; + if (peer.tp_type > static_cast(URMA_UTP)) { + errno = EPROTO; + return -1; + } + remote.tp_type = static_cast(peer.tp_type); + + urma_token_t token{}; + const bool use_bonding_extension = + IsUrmaBondingDevice() && remote.trans_mode == URMA_TM_RM; + errno = 0; + if (use_bonding_extension) { +#if BRPC_URMA_HAS_BONDING_EXT + // The bonding provider needs the local jetty to associate its send + // path with the imported target. A plain import may return success + // without setting that association, leaving traffic one-way only. + bondp_rjetty_t bonding_remote{}; + bonding_remote.base = remote; + bonding_remote.base.flag.bs.has_drv_ext = 1; + bonding_remote.jetty = _resource->jetty; + _resource->remote_jetty = + urma_import_jetty(ctx, &bonding_remote.base, &token); +#else + LOG(ERROR) << "Bonding remote jetty import requires provider header " + "urma_ubagg.h"; + errno = ENOTSUP; +#endif + } else { + _resource->remote_jetty = urma_import_jetty(ctx, &remote, &token); + } + if (!_resource->remote_jetty) { + if (errno == 0) { + errno = EIO; + } + char remote_eid[URMA_EID_STR_LEN + 1] = {}; + std::snprintf(remote_eid, sizeof(remote_eid), EID_FMT, + EID_RAW_ARGS(peer.eid)); + PLOG(ERROR) << "urma_import_jetty failed" + << " remote_eid=" << remote_eid + << " remote_uasid=" << peer.uasid + << " remote_jetty_id=" << peer.jetty_id + << " trans_mode=" << remote.trans_mode + << " tp_type=" << remote.tp_type + << " bonding_extension=" << use_bonding_extension; + return -1; + } + return 0; +} + +// ============================================================================ +// Send / recv data path. +// ============================================================================ + +// Private IOBuf accessor mirroring RdmaIOBuf: reach into IOBuf block refs to +// build a urma_sge_t directly, without memcpy. +class UrmaIOBuf : private butil::IOBuf { + friend class ::brpc::urma::UrmaEndpoint; +public: + using butil::IOBuf::_ref_num; + using butil::IOBuf::_ref_at; + using butil::IOBuf::fetch1; + using butil::IOBuf::get_first_data_meta; + using butil::IOBuf::cutn; + // Build the SGE for the current head block. + // Returns bytes added, or -1 (errno set). + ssize_t cut_into_sglist(urma_sge_t* sglist, size_t* sge_index, + butil::IOBuf* to, size_t max_sge, + size_t max_len) { + size_t len = 0; + while (*sge_index < max_sge && len < max_len && _ref_num() != 0) { + butil::IOBuf::BlockRef const& r = _ref_at(0); + const void* start = fetch1(); + urma_target_seg_t* tseg = + GetPoolSegFor(const_cast(start)); + if (!tseg) { + // User-registered memory: look up the seg handle. + uint64_t meta = get_first_data_meta(); + if (meta != 0) { + tseg = reinterpret_cast( + static_cast(meta)); + } + } + if (!tseg) { + errno = ERDMAMEM; + return -1; + } + size_t this_len = r.length; + if (len + this_len > max_len) { + this_len = max_len - len; + } + sglist[*sge_index].addr = reinterpret_cast(start); + sglist[*sge_index].len = static_cast(this_len); + sglist[*sge_index].tseg = tseg; + cutn(to, this_len); + len += this_len; + (*sge_index)++; + } + return static_cast(len); + } +}; + +ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { + if (!_resource || !_resource->jetty || !_resource->remote_jetty) { + errno = ENOTCONN; + return -1; + } + int max_sge = GetUrmaMaxSge(); + if (max_sge < 1) { + max_sge = 1; + } + + urma_sge_t* sglist = static_cast( + alloca(sizeof(urma_sge_t) * max_sge)); + if (!sglist) { + errno = ENOMEM; + return -1; + } + + size_t current = 0; + ssize_t total_len = 0; + while (current < ndata) { + uint16_t remote_wnd = _remote_rq_window_size.load(butil::memory_order_relaxed); + uint16_t sq_wnd = _sq_window_size.load(butil::memory_order_relaxed); + if (remote_wnd == 0 || sq_wnd == 0) { + if (total_len > 0) { + break; + } + errno = EAGAIN; + return -1; + } + butil::IOBuf* to = &_sbuf[_sq_current]; + size_t sge_index = 0; + size_t this_len = 0; + size_t max_len = _remote_recv_block_size > 0 + ? _remote_recv_block_size + : GetUrmaRecvBlockSize(); + while (sge_index < static_cast(max_sge) && + this_len < max_len && current < ndata) { + auto* data = reinterpret_cast(from[current]); + if (data->empty()) { + ++current; + continue; + } + ssize_t n = data->cut_into_sglist(sglist, &sge_index, to, + max_sge, max_len - this_len); + if (n < 0) { + return -1; + } + this_len += n; + } + if (sge_index == 0) { + break; + } + + urma_sg_t sg{sglist, static_cast(sge_index)}; + urma_jfs_wr_t wr{}; + std::memset(&wr, 0, sizeof(wr)); + // Send payload with URMA_OPC_SEND. Receive credits are flushed + // separately by SendImm() after SendAck() reaches its threshold. + // Piggybacking credits turns every payload into SEND_IMM and can + // produce asymmetric completions with the bonding provider. + wr.opcode = URMA_OPC_SEND; + wr.flag.bs.complete_enable = 1; + wr.tjetty = _resource->remote_jetty; + wr.send.src = sg; + wr.user_ctx = 1; + urma_jfs_wr_t* bad_wr = nullptr; + const uint16_t sq_slot = _sq_current; + const uint32_t local_jetty_id = _resource->jetty->jetty_id.id; + const uint32_t remote_jetty_id = _resource->remote_jetty->id.id; + + // Reserve both credits before making the WR visible to the provider. + // In polling mode a completion (and even the peer's receive-credit + // ACK) can be processed by another thread before post_send returns. + // Decrementing after post therefore creates a transient capacity + 1 + // window and makes the strict credit check tear down a healthy + // connection. + _remote_rq_window_size.fetch_sub(1, butil::memory_order_relaxed); + _sq_window_size.fetch_sub(1, butil::memory_order_relaxed); + int rc = urma_post_jetty_send_wr(_resource->jetty, &wr, &bad_wr); + if (rc != URMA_SUCCESS) { + const int provider_errno = errno; + _remote_rq_window_size.fetch_add(1, butil::memory_order_relaxed); + _sq_window_size.fetch_add(1, butil::memory_order_relaxed); + LOG(WARNING) << "urma_post_jetty_send_wr failed: " << rc + << ", provider_errno=" << provider_errno + << " (" << berror(provider_errno) << ')' + << ", bad_wr=" << static_cast(bad_wr) + << ", bad_is_current=" << (bad_wr == &wr) + << ", sq_slot=" << sq_slot + << ", local_jetty_id=" << local_jetty_id + << ", remote_jetty_id=" << remote_jetty_id + << ", state=" << GetStateStr() + << ", sq_window=" << sq_wnd + << ", remote_rq_window=" << remote_wnd + << ", num_sge=" << sge_index + << ", configured_max_sge=" << GetUrmaMaxSge() + << ", payload_size=" << this_len + << " on " << _socket->description(); + errno = rc; + return -1; + } + _sq_current = (_sq_current + 1) % (_sq_size - RESERVED_WR_NUM); + total_len += static_cast(this_len); + } + return total_len; +} + +bool UrmaEndpoint::IsWritable() const { + return _remote_rq_window_size.load(butil::memory_order_relaxed) > 0 && + _sq_window_size.load(butil::memory_order_relaxed) > 0; +} + +// ============================================================================ +// Recv path. +// ============================================================================ + +int UrmaEndpoint::DoPostRecv(void* block, size_t block_size) { + urma_target_seg_t* tseg = GetPoolSegFor(block); + if (!tseg) { + errno = ERDMAMEM; + return -1; + } + urma_sge_t sge{reinterpret_cast(block), + static_cast(block_size), tseg, nullptr}; + urma_sg_t sg{&sge, 1}; + urma_jfr_wr_t wr{sg, 0, nullptr}; + urma_jfr_wr_t* bad = nullptr; + // Use the shared-JFR path on every device, including bonding. The bonding + // provider owns physical receive scheduling for the JFR; a local + // jetty-to-target association is not part of the RM receive API. + const urma_status_t status = + urma_post_jfr_wr(_resource->jfr, &wr, &bad); + if (status != URMA_SUCCESS) { + LOG(WARNING) << "Failed to post URMA receive WR: status=" << status + << " bonding=" << IsUrmaBondingDevice() + << " bad_wr=" << static_cast(bad) + << " bad_is_current=" << (bad == &wr) + << " local_jetty_id=" << _resource->jetty->jetty_id.id + << " provider_associated_remote=" + << static_cast( + _resource->jetty->remote_jetty) + << " state=" << GetStateStr() + << " on " << _socket->description(); + errno = status; + return -1; + } + return 0; +} + +int UrmaEndpoint::PostRecv(uint32_t num, bool zerocopy) { + for (uint32_t i = 0; i < num; ++i) { + size_t block_size = GetUrmaRecvBlockSize(); + if (zerocopy) { + _rbuf[_rq_received].clear(); + butil::IOBufAsZeroCopyOutputStream zcis( + &_rbuf[_rq_received], block_size + IOBUF_BLOCK_HEADER_LEN); + void* data = nullptr; + int size = 0; + if (!zcis.Next(&data, &size) || !data || + size < static_cast(block_size)) { + errno = ENOMEM; + return -1; + } + _rbuf_data[_rq_received] = data; + if (DoPostRecv(data, block_size) < 0) { + return -1; + } + } else { + if (_rbuf_data[_rq_received] == nullptr) { + _rbuf[_rq_received].clear(); + butil::IOBufAsZeroCopyOutputStream zcos( + &_rbuf[_rq_received], + block_size + IOBUF_BLOCK_HEADER_LEN); + void* data = nullptr; + int size = 0; + if (!zcos.Next(&data, &size) || !data || + size < static_cast(block_size)) { + errno = ENOMEM; + return -1; + } + _rbuf_data[_rq_received] = data; + } + if (DoPostRecv(_rbuf_data[_rq_received], block_size) < 0) { + return -1; + } + } + _rq_received = (_rq_received + 1) % _rq_size; + } + return 0; +} + +int UrmaEndpoint::SendImm(uint32_t imm) { + if (imm == 0) { + return 0; + } + if (!_resource || !_resource->jetty || !_resource->remote_jetty) { + errno = ENOTCONN; + return -1; + } + if (_sq_imm_window_size == 0) { + errno = EAGAIN; + return -1; + } + // Empty-payload SEND_IMM flushes peer-side receive credit. Connection + // lifetime is owned by the TCP fd, so this is not an EOF marker. + urma_jfs_wr_t wr{}; + std::memset(&wr, 0, sizeof(wr)); + wr.opcode = URMA_OPC_SEND_IMM; + wr.flag.bs.complete_enable = 1; + wr.flag.bs.solicited_enable = 1; + wr.tjetty = _resource->remote_jetty; + wr.send.imm_data = imm; + wr.user_ctx = 0; // 0 == pure ack (HandleCompletion reuses budget). + urma_jfs_wr_t* bad = nullptr; + // Reserve the ACK-only SQ slot before posting for the same reason as the + // data windows in CutFromIOBufList: polling may observe its completion as + // soon as the provider accepts the WR. + --_sq_imm_window_size; + const urma_status_t status = + urma_post_jetty_send_wr(_resource->jetty, &wr, &bad); + if (status != URMA_SUCCESS) { + const int provider_errno = errno; + ++_sq_imm_window_size; + _new_rq_wrs.fetch_add(imm, butil::memory_order_relaxed); + LOG(WARNING) << "Failed to post URMA credit ACK: status=" << status + << " provider_errno=" << provider_errno + << " (" << berror(provider_errno) << ')' + << " bad_wr=" << static_cast(bad) + << " bad_is_current=" << (bad == &wr) + << " imm=" << imm + << " local_jetty_id=" + << _resource->jetty->jetty_id.id + << " remote_jetty_id=" + << _resource->remote_jetty->id.id + << " state=" << GetStateStr() + << " on " << _socket->description(); + errno = status; + return -1; + } + return 0; +} + +int UrmaEndpoint::SendAck(int num) { + const uint16_t old = + _new_rq_wrs.fetch_add(num, butil::memory_order_relaxed); + if (old + num > _remote_window_capacity / 2 && + _sq_imm_window_size > 0) { + return SendImm(_new_rq_wrs.exchange(0, butil::memory_order_relaxed)); + } + return 0; +} + +ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) { + bool zerocopy = FLAGS_urma_recv_zerocopy; + if (cr.status != URMA_CR_SUCCESS) { + LOG(WARNING) << "URMA completion failed, status=" << cr.status; + errno = EIO; + return -1; + } + if (cr.flag.bs.s_r == 0) { + // Send completion: reclaim SQ window and wake the writer. + if (cr.user_ctx == 0) { + // Pure-ack WR: just replenish the imm budget. + if (_sq_imm_window_size >= RESERVED_WR_NUM) { + LOG(WARNING) + << "URMA credit-ACK completion exceeds reserved SQ " + "window: current=" + << _sq_imm_window_size + << " capacity=" << RESERVED_WR_NUM + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + _sq_imm_window_size += 1; + SendAck(0); + return 0; + } + uint16_t wnd = 1; // We signal every WR (complete_enable=1). + uint16_t old = + _sq_window_size.load(butil::memory_order_relaxed); + while (true) { + if (old >= _local_window_capacity) { + LOG(WARNING) + << "URMA send completion exceeds SQ window: old=" << old + << " increment=" << wnd + << " capacity=" << _local_window_capacity + << " user_ctx=" << cr.user_ctx + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + if (_sq_window_size.compare_exchange_weak( + old, static_cast(old + wnd), + butil::memory_order_relaxed)) { + break; + } + } + for (uint16_t i = 0; i < wnd; ++i) { + _sbuf[_sq_sent].clear(); + _sq_sent = (_sq_sent + 1) % (_sq_size - RESERVED_WR_NUM); + } + butil::subtle::MemoryBarrier(); + if (_remote_rq_window_size.load(butil::memory_order_relaxed) >= + _local_window_capacity / 8) { + _socket->WakeAsEpollOut(); + } + return 0; + } + // Recv completion. + if (cr.opcode == URMA_CR_OPC_SEND_WITH_IMM && cr.imm_data > 0) { + if (cr.imm_data > _local_window_capacity) { + LOG(WARNING) << "Invalid URMA receive credit: " << cr.imm_data; + errno = EPROTO; + return -1; + } + const uint16_t acks = static_cast(cr.imm_data); + uint16_t old = + _remote_rq_window_size.load(butil::memory_order_relaxed); + while (true) { + if (old > _local_window_capacity - acks) { + LOG(WARNING) + << "URMA receive credit exceeds window: old=" << old + << " credit=" << acks + << " capacity=" << _local_window_capacity + << " imm=" << cr.imm_data + << " remote_window_capacity=" + << _remote_window_capacity + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + if (_remote_rq_window_size.compare_exchange_weak( + old, static_cast(old + acks), + butil::memory_order_relaxed)) { + break; + } + } + if (_sq_window_size.load(butil::memory_order_relaxed) > 0) { + _socket->WakeAsEpollOut(); + } + } else if (cr.completion_len == 0) { + LOG(WARNING) << "Zero-length URMA receive without immediate credit"; + errno = EPROTO; + return -1; + } + if (cr.completion_len > GetUrmaRecvBlockSize()) { + LOG(WARNING) << "URMA completion exceeds receive buffer: " + << cr.completion_len; + errno = EPROTO; + return -1; + } + if (cr.completion_len < static_cast(FLAGS_urma_zerocopy_min_size)) { + zerocopy = false; + } + if (zerocopy) { + _rbuf[_rq_received].cutn(&_socket->_read_buf, cr.completion_len); + } else { + _socket->_read_buf.append(_rbuf_data[_rq_received], cr.completion_len); + } + if (PostRecv(1, zerocopy) < 0) { + return -1; + } + if (cr.completion_len > 0) { + SendAck(1); + } + return static_cast(cr.completion_len); +} + +void UrmaEndpoint::DispatchReceivedBytes(SocketUniquePtr& s, ssize_t bytes) { + int64_t pending = _pending_received_bytes.load(butil::memory_order_relaxed); + if (bytes > 0) { + pending = _pending_received_bytes.fetch_add( + bytes, butil::memory_order_acq_rel) + bytes; + } + + const State state = _state.load(butil::memory_order_acquire); + if (state != ESTABLISHED) { + return; + } + + // PollCq and the handshake bthread can both reach this method when the + // state changes to ESTABLISHED. Serialize them so each byte added to + // _socket->_read_buf is reported to InputMessenger exactly once. + std::unique_lock dispatch_lock(_dispatch_mutex); + if (_state.load(butil::memory_order_acquire) != ESTABLISHED) { + return; + } + pending = _pending_received_bytes.exchange( + 0, butil::memory_order_acq_rel); + if (pending <= 0 || s->Failed()) { + return; + } + + auto* messenger = static_cast(s->user()); + if (!messenger) { + LOG(ERROR) << "URMA socket has no InputMessenger: " + << s->description(); + return; + } + + const int64_t received_us = butil::cpuwide_time_us(); + const int64_t base_realtime = butil::gettimeofday_us() - received_us; + InputMessageClosure last_msg; + messenger->ProcessNewMessage(s.get(), static_cast(pending), + false, received_us, base_realtime, last_msg); +} + +void UrmaEndpoint::PollCq(Socket* m) { + auto* ep = static_cast(m->user()); + if (!ep || !ep->_resource || !ep->_resource->jfc) { + return; + } + SocketUniquePtr s; + if (Socket::Address(ep->_socket->id(), &s) != 0) { + return; + } + if (s->Failed()) { + return; + } + + const bool event_mode = !FLAGS_urma_use_polling; + int progress = Socket::PROGRESS_INIT; + while (true) { + urma_jfc_t* event_jfc = nullptr; + if (event_mode) { + const int event_count = ep->WaitCqEvent(s, &event_jfc); + if (event_count < 0) { + return; + } + if (event_count == 0) { + if (!m->MoreReadEvents(&progress)) { + return; + } + continue; + } + } + + ssize_t bytes = 0; + auto drain_cq = [&]() -> int { + while (true) { + const int n = + std::max(1, std::min(FLAGS_urma_cqe_poll_once, 32)); + urma_cr_t crs[32]; + const int cnt = + urma_poll_jfc(ep->_resource->jfc, n, crs); + if (cnt < 0) { + return EIO; + } + if (cnt == 0) { + return 0; + } + for (int i = 0; i < cnt; ++i) { + if (s->Failed()) { + return ECANCELED; + } + const ssize_t nr = ep->HandleCompletion(crs[i]); + if (nr < 0) { + return errno ? errno : EIO; + } + bytes += nr; + } + } + }; + + int completion_error = drain_cq(); + if (event_mode) { + // The bonding provider records which physical JFCs produced CRs + // while bondp_poll_jfc drains the virtual JFC. + // bondp_rearm_jfc consumes that mask, so rearming before the drain + // leaves those physical JFCs unarmed. + uint32_t nevents = 1; + urma_ack_jfc(&event_jfc, &nevents, 1); + if (completion_error == 0) { + if (ep->ReqNotifyCq() != 0) { + return; + } + + // Close the drain/rearm race. A completion that arrived while + // the JFC was unarmed may not produce an edge on every + // provider. The JFC is armed now, so a final nonblocking drain + // is safe. + completion_error = drain_cq(); + } + } + + if (completion_error != 0) { + if (!s->Failed()) { + s->SetFailed(completion_error, "URMA completion error"); + } + return; + } + ep->DispatchReceivedBytes(s, bytes); + + if (!event_mode) { + return; + } + // The bonding JFCE fd is itself an epoll fd aggregating physical + // JFCEs, while brpc watches it with EPOLLET. urma_wait_jfc(..., 1, ...) + // consumes only one aggregated event. Keep draining the inner JFCE + // until it reports no event; otherwise another physical event can + // leave the fd continuously readable and never create a new outer + // edge. The event_count == 0 branch above resets _nevent only after + // the inner queue is empty. + } +} + +// ============================================================================ +// ApplyRemoteHello: size the send/recv windows from the peer's hello. +// ============================================================================ + +void UrmaEndpoint::ApplyRemoteHello(const ParsedHello& remote) { + _remote_recv_block_size = remote.buffer_size; + const uint32_t peer_rq_size = remote.recv_buffer_cnt + 1; + const uint32_t local_capacity = + std::min(_sq_size, peer_rq_size); + _local_window_capacity = static_cast( + local_capacity > RESERVED_WR_NUM + ? local_capacity - RESERVED_WR_NUM + : 0); + _remote_window_capacity = + _rq_size > RESERVED_WR_NUM ? _rq_size - RESERVED_WR_NUM : 0; + _sq_imm_window_size = RESERVED_WR_NUM; + _remote_rq_window_size.store(_local_window_capacity, + butil::memory_order_relaxed); + _sq_window_size.store(_local_window_capacity, butil::memory_order_relaxed); +} + +// ============================================================================ +// OnNewDataFromTcp: edge-triggered dispatcher. +// ============================================================================ + +static void TryReadOnTcpDuringUrmaEst(Socket* socket); + +void UrmaEndpoint::OnNewDataFromTcp(Socket* m) { + auto* tp = static_cast(m->_transport.get()); + if (!tp) { + return; + } + // Access _urma_ep directly (OnNewDataFromTcp is a friend of UrmaTransport); + // GetUrmaEp() CHECKs non-null which would crash on TCP-fallback sockets. + UrmaEndpoint* ep = tp->_urma_ep; + if (!ep) { + // No URMA endpoint: pure TCP path. + InputMessenger::OnNewMessages(m); + return; + } + int progress = 0; + while (true) { + const State state = + ep->_state.load(butil::memory_order_acquire); + if (state == UNINIT) { + if (!m->CreatedByConnect()) { + // Server side: kick off the handshake bthread. + if (!IsUrmaAvailable()) { + ep->_state = FALLBACK_TCP; + tp->_urma_state = UrmaTransport::URMA_OFF; + InputMessenger::OnNewMessages(m); + return; + } + SocketUniquePtr s; + m->ReAddress(&s); + ep->_state = S_HELLO_WAIT; + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + bthread_attr_set_name(&attr, "UrmaServerHandshake"); + if (bthread_start_background(&tid, &attr, + ProcessHandshakeAtServer, ep) != 0) { + ep->_state = UNINIT; + LOG(FATAL) << "Fail to start UrmaServerHandshake bthread"; + } else { + s.release(); + } + return; + } + // Client side: handled by ProcessHandshakeAtClient. + return; + } else if (state < ESTABLISHED) { + // During handshake: wake the handshake bthread parked in ReadFromFd. + ep->_read_butex->fetch_add(1, butil::memory_order_release); + bthread::butex_wake(ep->_read_butex); + return; + } else if (state == FALLBACK_TCP) { + InputMessenger::OnNewMessages(m); + return; + } else if (state == ESTABLISHED) { + TryReadOnTcpDuringUrmaEst(m); + return; + } + if (!m->MoreReadEvents(&progress)) { + break; + } + } +} + +inline void UrmaEndpoint::TryReadOnTcp() { + if (_state.load(butil::memory_order_acquire) == FALLBACK_TCP) { + InputMessenger::OnNewMessages(_socket); + } +} + +static void TryReadOnTcpDuringUrmaEst(Socket* socket) { + int progress = Socket::PROGRESS_INIT; + while (true) { + uint8_t byte = 0; + const ssize_t nr = read(socket->fd(), &byte, 1); + if (nr < 0) { + if (errno != EAGAIN) { + const int saved_errno = errno; + socket->SetFailed(saved_errno, "Fail to read URMA TCP fd: %s", + berror(saved_errno)); + return; + } + if (!socket->MoreReadEvents(&progress)) { + return; + } + } else if (nr == 0) { + socket->SetEOF(); + return; + } else { + socket->SetFailed( + EPROTO, "Unexpected TCP data after URMA was established"); + return; + } + } +} + +void UrmaEndpoint::FallbackToTcp(UrmaTransport* transport, bool process_tcp) { + transport->_urma_state = UrmaTransport::URMA_OFF; + _state.store(FALLBACK_TCP, butil::memory_order_release); + DeallocateResources(); + if (process_tcp) { + TryReadOnTcp(); + } +} + +void UrmaEndpoint::FailHandshake(UrmaTransport* transport, int error, + const char* reason) { + LOG(ERROR) << "URMA handshake failed in state=" << GetStateStr() + << " on " << _socket->description() + << ": " << reason << ", error=" << error + << " (" << berror(error) << ')'; + transport->_urma_state = UrmaTransport::URMA_OFF; + _state.store(FAILED, butil::memory_order_release); + DeallocateResources(); + auto* connect = + static_cast(_socket->_app_connect.get()); + if (connect) { + connect->_error = error; + } + _socket->SetFailed(error, "URMA handshake failed: %s", reason); +} + +// ============================================================================ +// Handshake state machines (client / server). Run in a background bthread. +// ============================================================================ + +void* UrmaEndpoint::ProcessHandshakeAtClient(void* arg) { + auto* ep = static_cast(arg); + SocketUniquePtr s(ep->_socket); + auto* tp = static_cast(s->_transport.get()); + UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); + if (!IsUrmaAvailable()) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_state = C_ALLOC_RES; + if (ep->AllocateResources() < 0) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + // Prepost the shared JFR before sending the client hello so the peer sees + // a ready receive queue as soon as its import completes. + if (ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_state = C_HELLO_SEND; + std::unique_ptr hs(CreateClientHandshake(ep)); + ep->_handshake_version = hs->ProtocolVersion(); + if (hs->SendLocalHello() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send client hello"); + return nullptr; + } + ep->_state = C_HELLO_WAIT; + ParsedHello remote; + bool negotiated = false; + if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read server hello"); + return nullptr; + } + if (!negotiated) { + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->ApplyRemoteHello(remote); + ep->_state = C_IMPORT_PEER; + if (ep->ImportPeer(remote) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "import server resources"); + return nullptr; + } + ep->_state = C_ACK_SEND; + uint32_t flags = HELLO_ACK_URMA_OK; + uint32_t flags_be = butil::HostToNet32(flags); + if (ep->WriteToFd(&flags_be, HELLO_ACK_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send client ack"); + return nullptr; + } + tp->_urma_state = UrmaTransport::URMA_ON; + ep->_state = ESTABLISHED; + ep->DispatchReceivedBytes(s, 0); + return nullptr; +} + +void* UrmaEndpoint::ProcessHandshakeAtServer(void* arg) { + auto* ep = static_cast(arg); + SocketUniquePtr s(ep->_socket); + auto* tp = static_cast(s->_transport.get()); + UrmaConnect::RunGuard guard(static_cast(s->_app_connect.get())); + ep->_state = S_HELLO_WAIT; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client magic"); + return nullptr; + } + std::unique_ptr hs(CreateServerHandshakeByMagic(ep, magic)); + if (!hs) { + // Not an URMA peer: push the magic back and fall back to TCP. + ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + ep->FallbackToTcp(tp, true); + return nullptr; + } + ep->_handshake_version = hs->ProtocolVersion(); + ParsedHello remote; + bool negotiated = false; + if (hs->ReceiveAndParseRemoteHello(&remote, &negotiated) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client hello"); + return nullptr; + } + if (!negotiated) { + ep->FailHandshake(tp, EPROTO, "invalid client hello"); + return nullptr; + } + ep->_state = S_ALLOC_RES; + if (ep->AllocateResources() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "allocate server resources"); + return nullptr; + } + const bool bonding = IsUrmaBondingDevice(); + if (!bonding && + ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "post server receives"); + return nullptr; + } + ep->ApplyRemoteHello(remote); + ep->_state = S_IMPORT_PEER; + if (ep->ImportPeer(remote) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "import client resources"); + return nullptr; + } + if (bonding && + ep->PostRecv(ep->_rq_size, FLAGS_urma_recv_zerocopy) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "post server receives"); + return nullptr; + } + ep->_state = S_HELLO_SEND; + if (hs->SendLocalHello() < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "send server hello"); + return nullptr; + } + ep->_state = S_ACK_WAIT; + uint32_t flags_be = 0; + if (ep->ReadFromFd(&flags_be, HELLO_ACK_LEN) < 0) { + const int saved_errno = errno ? errno : EIO; + ep->FailHandshake(tp, saved_errno, "read client ack"); + return nullptr; + } + uint32_t flags = butil::NetToHost32(flags_be); + bool client_ack_ok = (flags & HELLO_ACK_URMA_OK) != 0; + if (client_ack_ok) { + if (tp->_urma_state.load(butil::memory_order_acquire) == + UrmaTransport::URMA_OFF) { + // Protocol breakdown: client wants URMA but we already fell back. + ep->FailHandshake(tp, EPROTO, "client ack mismatch"); + return nullptr; + } + tp->_urma_state = UrmaTransport::URMA_ON; + ep->_state = ESTABLISHED; + ep->DispatchReceivedBytes(s, 0); + } else { + ep->FallbackToTcp(tp, true); + } + return nullptr; +} + +// ============================================================================ +// UrmaConnect: drives the client handshake bthread. +// ============================================================================ + +void UrmaConnect::StartConnect(const Socket* socket, + void (*done)(int, void*), void* data) { + SocketUniquePtr s; + if (Socket::Address(socket->id(), &s) != 0) { + return; + } + _done = done; + _data = data; + _error = 0; + auto* tp = static_cast(socket->_transport.get()); + if (!tp) { + Run(); + return; + } + if (!tp->_urma_ep || !IsUrmaAvailable()) { + // Fall back to TCP immediately. + if (tp->_urma_ep) { + tp->_urma_ep->_state = UrmaEndpoint::FALLBACK_TCP; + } + tp->_urma_state = UrmaTransport::URMA_OFF; + Run(); + return; + } + bthread_t tid; + bthread_attr_t attr = BTHREAD_ATTR_NORMAL; + bthread_attr_set_name(&attr, "UrmaClientHandshake"); + if (bthread_start_background(&tid, &attr, + UrmaEndpoint::ProcessHandshakeAtClient, + tp->_urma_ep) != 0) { + tp->_urma_ep->_state = UrmaEndpoint::FALLBACK_TCP; + tp->_urma_state = UrmaTransport::URMA_OFF; + Run(); + } else { + // ProcessHandshakeAtClient adopts this reference in its + // SocketUniquePtr constructor. + s.release(); + } +} + +void UrmaConnect::StopConnect(Socket*) {} + +void UrmaConnect::Run() { + if (_done) { + auto cb = _done; + _done = nullptr; + cb(_error, _data); + } +} + +// ============================================================================ +// Debug / polling-mode stubs. +// ============================================================================ + +std::string UrmaEndpoint::GetStateStr() const { + switch (_state.load(butil::memory_order_acquire)) { + case UNINIT: return "UNINIT"; + case C_ALLOC_RES: return "C_ALLOC_RES"; + case C_HELLO_SEND: return "C_HELLO_SEND"; + case C_HELLO_WAIT: return "C_HELLO_WAIT"; + case C_IMPORT_PEER: return "C_IMPORT_PEER"; + case C_ACK_SEND: return "C_ACK_SEND"; + case S_HELLO_WAIT: return "S_HELLO_WAIT"; + case S_ALLOC_RES: return "S_ALLOC_RES"; + case S_IMPORT_PEER: return "S_IMPORT_PEER"; + case S_HELLO_SEND: return "S_HELLO_SEND"; + case S_ACK_WAIT: return "S_ACK_WAIT"; + case ESTABLISHED: return "ESTABLISHED"; + case FALLBACK_TCP: return "FALLBACK_TCP"; + case FAILED: return "FAILED"; + } + return "UNKNOWN"; +} + +void UrmaEndpoint::DebugInfo(std::ostream& os, butil::StringPiece) const { + os << "state=" << GetStateStr() + << " sq_size=" << _sq_size << " rq_size=" << _rq_size + << " remote_recv_block_size=" << _remote_recv_block_size + << " sq_window=" << _sq_window_size.load(butil::memory_order_relaxed) + << " remote_rq_window=" << _remote_rq_window_size.load(butil::memory_order_relaxed) + << " handshake_version=" << _handshake_version; +} + +int UrmaEndpoint::WaitCqEvent(SocketUniquePtr& s, + urma_jfc_t** event_jfc) { + if (!_resource || !_resource->jfce || !_resource->jfc) { + errno = ENODEV; + return -1; + } + *event_jfc = nullptr; + int count = urma_wait_jfc(_resource->jfce, 1, 0, event_jfc); + if (count < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) { + return 0; + } + const int saved_errno = errno; + PLOG(ERROR) << "Fail to wait URMA JFC event from " + << s->description(); + s->SetFailed(saved_errno, "Fail to wait URMA JFC event: %s", + berror(saved_errno)); + return -1; + } + if (count == 0) { + return 0; + } + if (*event_jfc != _resource->jfc) { + LOG(ERROR) << "Unexpected URMA JFC event on " << s->description(); + errno = EPROTO; + s->SetFailed(EPROTO, "Unexpected URMA JFC event"); + return -1; + } + return 1; +} + +int UrmaEndpoint::ReqNotifyCq() { + if (!_resource || !_resource->jfc) { + errno = ENODEV; + return -1; + } + const int rc = urma_rearm_jfc(_resource->jfc, false); + if (rc != URMA_SUCCESS) { + errno = rc; + PLOG(WARNING) << "Fail to rearm URMA JFC"; + _socket->SetFailed(rc, "Fail to rearm URMA JFC: %s", berror(rc)); + return -1; + } + return 0; +} + +int UrmaEndpoint::PollingModeInitialize( + bthread_tag_t tag, std::function callback, + std::function init_fn, std::function release_fn) { + if (!FLAGS_urma_use_polling) { + return 0; + } + if (tag >= _poller_groups.size() || + _poller_groups[tag].pollers.empty()) { + errno = EINVAL; + return -1; + } + auto& group = _poller_groups[tag]; + bool expected = false; + if (!group.running.compare_exchange_strong(expected, true)) { + return 0; + } + struct FnArgs { + Poller* poller; + butil::atomic* running; + }; + auto fn = [](void* p) -> void* { + std::unique_ptr args(static_cast(p)); + Poller* poller = args->poller; + butil::atomic* running = args->running; + std::unordered_set cq_sids; + CqSidOp op; + + if (poller->init_fn) { + poller->init_fn(); + } + while (running->load(butil::memory_order_relaxed)) { + while (poller->op_queue.Dequeue(op)) { + if (op.type == CqSidOp::ADD) { + cq_sids.emplace(op.sid); + } else { + cq_sids.erase(op.sid); + } + } + for (SocketId sid : cq_sids) { + SocketUniquePtr s; + if (Socket::Address(sid, &s) == 0) { + PollCq(s.get()); + } + } + if (poller->callback) { + poller->callback(); + } + if (FLAGS_urma_poller_yield || cq_sids.empty()) { + bthread_yield(); + } + } + if (poller->release_fn) { + poller->release_fn(); + } + return nullptr; + }; + + auto& pollers = group.pollers; + for (size_t i = 0; i < pollers.size(); ++i) { + pollers[i].callback = callback; + pollers[i].init_fn = init_fn; + pollers[i].release_fn = release_fn; + std::unique_ptr args(new (std::nothrow) + FnArgs{&pollers[i], &group.running}); + if (!args) { + group.running.store(false, butil::memory_order_relaxed); + for (size_t j = 0; j < i; ++j) { + bthread_join(pollers[j].tid, nullptr); + pollers[j].tid = INVALID_BTHREAD; + } + errno = ENOMEM; + return -1; + } + bthread_attr_t attr = FLAGS_urma_disable_bthread + ? BTHREAD_ATTR_PTHREAD + : BTHREAD_ATTR_NORMAL; + attr.tag = tag; + bthread_attr_set_name(&attr, "UrmaPolling"); + const int rc = bthread_start_background( + &pollers[i].tid, &attr, fn, args.get()); + if (rc != 0) { + group.running.store(false, butil::memory_order_relaxed); + for (size_t j = 0; j < i; ++j) { + bthread_join(pollers[j].tid, nullptr); + pollers[j].tid = INVALID_BTHREAD; + } + errno = rc; + return -1; + } + args.release(); + } + return 0; +} + +void UrmaEndpoint::PollingModeRelease(bthread_tag_t tag) { + if (!FLAGS_urma_use_polling || tag >= _poller_groups.size()) { + return; + } + auto& group = _poller_groups[tag]; + group.running.store(false, butil::memory_order_relaxed); + for (auto& poller : group.pollers) { + if (poller.tid != INVALID_BTHREAD) { + bthread_join(poller.tid, nullptr); + poller.tid = INVALID_BTHREAD; + } + } +} + +void UrmaEndpoint::PollerAddCqSid() { + if (_cq_sid == INVALID_SOCKET_ID || _poller_groups.empty()) { + return; + } + _poller_tag = bthread_self_tag(); + if (_poller_tag >= _poller_groups.size()) { + return; + } + auto& pollers = _poller_groups[_poller_tag].pollers; + if (pollers.empty()) { + return; + } + const size_t index = + butil::fmix32(_cq_sid) % pollers.size(); + pollers[index].op_queue.Enqueue( + CqSidOp{CqSidOp::ADD, _cq_sid}); +} + +void UrmaEndpoint::PollerRemoveCqSid() { + if (_cq_sid == INVALID_SOCKET_ID || _poller_groups.empty() || + _poller_tag >= _poller_groups.size()) { + return; + } + auto& pollers = _poller_groups[_poller_tag].pollers; + if (pollers.empty()) { + return; + } + const size_t index = + butil::fmix32(_cq_sid) % pollers.size(); + pollers[index].op_queue.Enqueue( + CqSidOp{CqSidOp::REMOVE, _cq_sid}); +} + +int UrmaEndpoint::GlobalInitialize() { + // Pre-allocate the prepared jetty pool. Skipped if URMA init is skipped + // (unit-test mode). + if (FLAGS_urma_use_polling && _poller_groups.empty()) { + if (FLAGS_urma_poller_num <= 0) { + LOG(ERROR) << "urma_poller_num must be positive"; + errno = EINVAL; + return -1; + } + size_t ntags = static_cast(FLAGS_task_group_ntags); + if (ntags == 0) { + ntags = 1; + } + _poller_groups = std::vector(ntags); + } + if (g_prepared_cnt > 0) { + return 0; + } + urma_context_t* ctx = GetUrmaContext(); + if (!ctx) { + return 0; + } + const int prepared_jetty_count = PreparedJettyCount(); + for (int i = 0; i < prepared_jetty_count; ++i) { + auto* r = new (std::nothrow) UrmaResource(); + if (!r) { + break; + } + r->jfce = urma_create_jfce(ctx); + if (!r->jfce || + (!FLAGS_urma_use_polling && r->jfce->fd < 0)) { + delete r; + break; + } + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = static_cast(FLAGS_urma_sq_size + FLAGS_urma_rq_size); + jfc_cfg.jfce = r->jfce; + r->jfc = urma_create_jfc(ctx, &jfc_cfg); + if (!r->jfc) { + delete r; + break; + } + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = static_cast(FLAGS_urma_rq_size); + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + jfr_cfg.jfc = r->jfc; + r->jfr = urma_create_jfr(ctx, &jfr_cfg); + if (!r->jfr) { + delete r; + break; + } + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = static_cast(FLAGS_urma_sq_size); + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.priority = GetUrmaJettyPriority(); + jetty_cfg.jfs_cfg.max_sge = + static_cast(GetUrmaMaxSge()); + jetty_cfg.jfs_cfg.rnr_retry = URMA_TYPICAL_RNR_RETRY; + jetty_cfg.jfs_cfg.err_timeout = URMA_TYPICAL_ERR_TIMEOUT; + jetty_cfg.jfs_cfg.jfc = r->jfc; + jetty_cfg.shared.jfr = r->jfr; + jetty_cfg.shared.jfc = r->jfc; + r->jetty = urma_create_jetty(ctx, &jetty_cfg); + if (!r->jetty) { + delete r; + break; + } + r->next = g_prepared_list; + g_prepared_list = r; + ++g_prepared_cnt; + } + return 0; +} + +void UrmaEndpoint::GlobalRelease() { + { + BAIDU_SCOPED_LOCK(g_prepared_mutex); + while (g_prepared_list) { + UrmaResource* next = g_prepared_list->next; + delete g_prepared_list; + g_prepared_list = next; + } + g_prepared_cnt = 0; + } + for (size_t tag = 0; tag < _poller_groups.size(); ++tag) { + PollingModeRelease(static_cast(tag)); + } +} + +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA diff --git a/src/brpc/urma/urma_endpoint.h b/src/brpc/urma/urma_endpoint.h new file mode 100644 index 0000000000..0b87d06d24 --- /dev/null +++ b/src/brpc/urma/urma_endpoint.h @@ -0,0 +1,355 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef BRPC_URMA_ENDPOINT_H +#define BRPC_URMA_ENDPOINT_H + +#include +#include +#include +#include + +#include "butil/atomicops.h" +#include "butil/containers/mpsc_queue.h" +#include "butil/iobuf.h" +#include "butil/macros.h" +#include "butil/synchronization/lock.h" +#include "bthread/types.h" + +#include "brpc/socket.h" + +#if BRPC_WITH_URMA + +#include "urma_api.h" +#include "urma_types.h" +#include "brpc/urma/urma_handshake.h" +#include "brpc/urma/urma_handshake.pb.h" + +namespace brpc { + +class UrmaTransport; + +namespace urma { + +class UrmaHandshake; +struct ParsedHello; + +DECLARE_bool(urma_use_polling); +DECLARE_int32(urma_poller_num); +DECLARE_bool(urma_disable_bthread); + +// Per-connection application-level connect object. Returned by +// UrmaTransport::Connect(); its StartConnect spawns the client-side handshake +// bthread that drives the URMA negotiation over the already-connected TCP fd. +class UrmaConnect : public AppConnect { + friend class UrmaEndpoint; +public: + void StartConnect(const Socket* socket, + void (*done)(int err, void* data), void* data) override; + void StopConnect(Socket*) override; + + struct RunGuard { + explicit RunGuard(UrmaConnect* rc) : this_rc(rc) {} + ~RunGuard() { if (this_rc) this_rc->Run(); } + UrmaConnect* this_rc; + }; + +private: + void Run(); + void (*_done)(int, void*){nullptr}; + void* _data{nullptr}; + int _error{0}; +}; + +// POD holder for the URMA kernel objects backing one connection: +// - the send jetty (JFS), shared-JFR jetty, JFR, JFC, and (event mode) JFCE +// - the imported peer jetty and peer segment +// Mirrors the role of rdma::RdmaResource, but the URMA object graph is a +// little different (no separate QP/CQ split). +struct UrmaResource { + UrmaResource* next{nullptr}; // singly-linked list for the prepared pool + + urma_jfc_t* jfc{nullptr}; + urma_jfce_t* jfce{nullptr}; // event mode only; null in polling mode + urma_jfr_t* jfr{nullptr}; + urma_jetty_t* jetty{nullptr}; + // Imported peer objects (created per-connection, not pooled). + urma_target_jetty_t* remote_jetty{nullptr}; + urma_target_seg_t* remote_seg{nullptr}; + + UrmaResource() = default; + ~UrmaResource(); + DISALLOW_COPY_AND_ASSIGN(UrmaResource); +}; + +// One per Socket. Carries the handshake state machine, the send/recv +// windows, and the registered buffer bookkeeping. Cache-line padded to avoid +// false sharing between connections. +class BAIDU_CACHELINE_ALIGNMENT UrmaEndpoint : public SocketUser { + friend class UrmaConnect; + friend class Socket; + friend class UrmaTransport; + friend class UrmaHandshakeClientV2; + friend class UrmaHandshakeServerV2; + friend class UrmaHandshakeClientV3; + friend class UrmaHandshakeServerV3; + friend int DrainBytes(UrmaEndpoint*, size_t); + friend int ReadBodyAndNegotiate(UrmaEndpoint*, ParsedHello*, bool*); + +public: + explicit UrmaEndpoint(Socket* s); + ~UrmaEndpoint() override; + + // ---- Global initialization / release ---- + // Pre-allocate the prepared Jetty pool. Called once at GlobalUrmaInitializeOrDie. + static int GlobalInitialize(); + static void GlobalRelease(); + + // ---- Per-connection lifecycle ---- + // Reset the endpoint back to UNINIT for reuse. + void Reset(); + + // ---- Data path (called by UrmaTransport) ---- + // Cut data from the IOBuf list and post it via URMA SEND. + // Returns bytes sent, or -1 (errno set; EAGAIN when windows are full). + ssize_t CutFromIOBufList(butil::IOBuf** data, size_t ndata); + + // Whether the endpoint can post more sends (both windows non-empty). + bool IsWritable() const; + + // For debug: dump endpoint state to @os. + void DebugInfo(std::ostream& os, + butil::StringPiece connector = "\n") const; + + // Edge-triggered callback installed on the TCP socket when the transport + // is URMA. Dispatches on _state (see the .cpp for the full state machine): + // UNINIT (server) -> start the server handshake bthread + // < ESTABLISHED -> wake the handshake bthread parked in ReadFromFd + // FALLBACK_TCP -> hand off to InputMessenger::OnNewMessages + // ESTABLISHED -> drain stray TCP bytes + static void OnNewDataFromTcp(Socket* m); + + // CQ completion poller: the edge-triggered callback for the CQ socket + // (whose user is this endpoint). Drains urma_poll_jfc and feeds the + // completions into HandleCompletion, then calls InputMessenger. Event mode + // also drains the aggregated JFCE until urma_wait_jfc reports no event, + // then calls Socket::MoreReadEvents so later JFCE edges can start the + // callback again. + static void PollCq(Socket* m); + + // Initialize the per-tag polling infrastructure (polling mode). + static int PollingModeInitialize(bthread_tag_t tag, + std::function callback, + std::function init_fn, + std::function release_fn); + static void PollingModeRelease(bthread_tag_t tag); + + // ---- Handshake IO helpers (also used by urma_handshake.cpp) ---- + // Read at most @len bytes from the TCP fd into @data; waits on _read_butex + // on EAGAIN. Returns 0 on success, -1 on IO error (errno set). + int ReadFromFd(void* data, size_t len); + // Push @len bytes back into _socket->_read_buf so the TCP input messenger + // can re-parse them (used on fallback when the magic is not "URMA"). + void PushBackToReadBuf(const void* data, size_t len); + // Write at most @len bytes from @data to the TCP fd; waits on + // _epollout_butex on EAGAIN. Returns 0 on success, -1 on IO error. + int WriteToFd(void* data, size_t len); + // Write an IOBuf to the TCP fd (used by v3 protobuf handshake). + int WriteToFd(butil::IOBuf& data); + + // ---- Hello builders/parsers (used by urma_handshake.cpp) ---- + // Fill the v2 binary HelloMessage with this endpoint's local params. + void FillLocalHelloV2(v2_wire::HelloMessage* out) const; + // Fill the v3 protobuf UrmaHello with this endpoint's local params. + void FillLocalHelloV3(UrmaHello* out) const; + // Write "URM3" + 4B big-endian pb_size + protobuf bytes. + int WriteHelloV3(const UrmaHello& msg); + // Read pb_size (4B) + protobuf bytes; parse; validate. Sets *negotiated + // to false (returns 0) if the message is invalid. + int ReadAndParseHelloV3(ParsedHello* out, bool* negotiated); + + // Apply the peer's negotiated parameters (window sizes, peer jetty/seg). + void ApplyRemoteHello(const ParsedHello& remote); + +private: + enum State { + UNINIT = 0x0, + C_ALLOC_RES = 0x1, // client: allocate Jetty/JFC/JFR + C_HELLO_SEND = 0x2, + C_HELLO_WAIT = 0x3, + C_IMPORT_PEER = 0x4, // urma_import_seg then urma_import_jetty + C_ACK_SEND = 0x5, + S_HELLO_WAIT = 0x11, + S_ALLOC_RES = 0x12, + S_IMPORT_PEER = 0x13, + S_HELLO_SEND = 0x14, + S_ACK_WAIT = 0x15, + ESTABLISHED = 0x100, + FALLBACK_TCP = 0x200, + FAILED = 0x300 + }; + + // Process handshake at the client / server (run in a background bthread). + static void* ProcessHandshakeAtClient(void* arg); + static void* ProcessHandshakeAtServer(void* arg); + + // Allocate / deallocate the per-connection URMA resources (JFCE/JFC/JFR/ + // Jetty, plus the CQ socket). Returns 0 on success. + int AllocateResources(); + void DeallocateResources(); + + // Import the peer's jetty and buffer-pool segment. CRITICAL: + // urma_import_seg is called BEFORE urma_import_jetty, otherwise the kernel + // does not establish the transport-path routing for the remote EID and + // the first SEND is rejected with URMA_CR_RNR_RETRY_CNT_EXC_ERR. + // Returns 0 on success. + int ImportPeer(const ParsedHello& peer); + + // Post @num recv WRs into the local JFR. If @zerocopy, use a fresh pool + // buffer; otherwise reuse the fixed _rbuf slot. Returns 0 on success. + int PostRecv(uint32_t num, bool zerocopy); + + // Post a single recv WR pointing at @block of @block_size. + int DoPostRecv(void* block, size_t block_size); + + // Send a pure-ACK URMA_OPC_SEND_IMM WR with no payload. @imm carries the + // number of receive WRs reposted for peer-side flow-control credit. + int SendImm(uint32_t imm); + // Batched credit ack: if _new_rq_wrs is above the threshold, flush it via + // a standalone SendImm. @num is the number of new recv WRs to add. + int SendAck(int num); + + // Handle one completion record. For SEND completions: reclaim the SQ + // window and wake the writer. For RECV completions: cut the payload into + // _socket->_read_buf, repost the recv WR, and ack. Returns bytes received + // (0 for send completions), or -1 on error (errno set). + ssize_t HandleCompletion(const urma_cr_t& cr); + + // Queue received bytes for InputMessenger. CQ receive completions can + // arrive while the server is still waiting for the final TCP handshake + // ACK. Keep those bytes in _socket->_read_buf and dispatch them only after + // ESTABLISHED, matching the connection-ready boundary seen by user code. + void DispatchReceivedBytes(SocketUniquePtr& s, ssize_t bytes); + + // Consume one async event from the JFCE fd (event mode only). The caller + // drains completions before acknowledge/rearm because the bonding provider + // uses poll_jfc to record which physical JFCs must be rearmed. + int WaitCqEvent(SocketUniquePtr& s, urma_jfc_t** event_jfc); + + // Request completion notification (event mode only). + int ReqNotifyCq(); + + // Add/remove the CQ socket id to/from the poller (polling mode only). + void PollerAddCqSid(); + void PollerRemoveCqSid(); + + inline void TryReadOnTcp(); + void FallbackToTcp(UrmaTransport* transport, bool process_tcp); + void FailHandshake(UrmaTransport* transport, int error, + const char* reason); + + // Construct a ParsedHello from local state (used by SendLocalHello). + void MakeLocalParsedHello(ParsedHello* out) const; + + std::string GetStateStr() const; + + // Not owned. + Socket* _socket; + + // Handshake state. + butil::atomic _state{UNINIT}; + int _handshake_version{0}; // 0 = unnegotiated; 2 = v2; 3 = v3 + butil::atomic _pending_received_bytes{0}; + butil::Mutex _dispatch_mutex; + + // The URMA resources (jetty / jfc / jfr / jfce / imported peer objects). + UrmaResource* _resource{nullptr}; + + // The SocketId wrapping the JFCE fd (event mode) or a synthetic carrier + // (polling mode). PollCq is the edge-triggered callback on this socket. + SocketId _cq_sid{INVALID_SOCKET_ID}; + + // ---- Send / recv window bookkeeping ---- + uint16_t _sq_size{0}; // local JFS depth + uint16_t _rq_size{0}; // local JFR depth + + // Per-WR send buffers (own the IOBuf until the SEND completes). + std::vector _sbuf; + // Per-WR recv buffers (zero-copy targets). + std::vector _rbuf; + std::vector _rbuf_data; + // Peer's advertised recv buffer size (caps each WR's payload). + uint32_t _remote_recv_block_size{0}; + + // Flow-control windows (atomic; producer decrements, completion handler + // increments). + uint16_t _local_window_capacity{0}; + uint16_t _remote_window_capacity{0}; + butil::atomic _remote_rq_window_size{0}; // WRs we can send + butil::atomic _sq_window_size{0}; // WRs we can post + butil::atomic _new_rq_wrs{0}; // new recv WRs (to ack) + uint16_t _sq_imm_window_size{0}; // budget for pure-ack WRs + + // SQ producer / consumer indices. + uint16_t _sq_current{0}; + uint16_t _sq_sent{0}; + // RQ consumer index. + uint16_t _rq_received{0}; + + // Butex for waking the handshake bthread parked in ReadFromFd. + butil::atomic* _read_butex{nullptr}; + + // Reserved WR slots for pure-ack IMM WRs. + static constexpr uint16_t RESERVED_WR_NUM = 3; + + // Cq socket id operation (polling mode registration queue). + struct CqSidOp { + enum OpType { ADD, REMOVE } type; + SocketId sid; + }; + struct BAIDU_CACHELINE_ALIGNMENT Poller { + bthread_t tid{INVALID_BTHREAD}; + butil::MPSCQueue> op_queue; + std::function callback; + std::function init_fn; + std::function release_fn; + }; + struct BAIDU_CACHELINE_ALIGNMENT PollerGroup { + PollerGroup() : pollers(FLAGS_urma_poller_num), running(false) {} + std::vector pollers; + butil::atomic running; + }; + static std::vector _poller_groups; + bthread_tag_t _poller_tag{0}; + + DISALLOW_COPY_AND_ASSIGN(UrmaEndpoint); +}; + +} // namespace urma +} // namespace brpc + +#else // BRPC_WITH_URMA + +namespace brpc { +namespace urma { +class UrmaEndpoint {}; +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA + +#endif // BRPC_URMA_ENDPOINT_H diff --git a/src/brpc/urma/urma_handshake.cpp b/src/brpc/urma/urma_handshake.cpp new file mode 100644 index 0000000000..07a02f3bc1 --- /dev/null +++ b/src/brpc/urma/urma_handshake.cpp @@ -0,0 +1,356 @@ +// 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/urma/urma_handshake.h" + +#if BRPC_WITH_URMA + +#include +#include +#include + +#include + +#include "butil/atomicops.h" +#include "butil/iobuf.h" // IOBuf, IOPortal, IOBufAsZeroCopy* +#include "butil/logging.h" +#include "butil/sys_byteorder.h" + +#include "brpc/socket.h" +#include "brpc/urma/urma_endpoint.h" +#include "brpc/urma/urma_helper.h" +#include "brpc/urma/urma_handshake.pb.h" +#include "brpc/urma_transport.h" + +namespace brpc { +namespace urma { + +DEFINE_int32(urma_client_handshake_version, 2, + "Client handshake version: 2 = binary, 3 = protobuf"); + +// ============================================================================ +// v2 binary HelloMessage. +// On-wire layout (network byte order, tightly packed body of 82 bytes): +// +// offset field size +// 0 msg_len 2B (full packet length incl. magic) +// 2 hello_ver 2B +// 4 impl_ver 2B +// 6 buffer_size 4B +// 10 recv_buffer_cnt 4B +// 14 jetty_id 4B +// 18 eid 16B (raw, no swap) +// 34 uasid 4B +// 38 tp_type 1B +// 39 pad 3B +// 42 seg_eid 16B (raw) +// 58 seg_uasid 4B +// 62 seg_va 8B +// 70 seg_len 8B +// 78 seg_token_id 4B +// total = 82 bytes body. +// +// Full packet = magic "URMA" (4B) + body (82B) = 86 bytes. +// ============================================================================ + +namespace v2_wire { + +void HelloMessage::Serialize(void* buf) const { + uint8_t* p = static_cast(buf); + auto write16 = [&p](uint16_t value) { + value = butil::HostToNet16(value); + std::memcpy(p, &value, sizeof(value)); + p += sizeof(value); + }; + auto write32 = [&p](uint32_t value) { + value = butil::HostToNet32(value); + std::memcpy(p, &value, sizeof(value)); + p += sizeof(value); + }; + auto write64 = [&p](uint64_t value) { + value = butil::HostToNet64(value); + std::memcpy(p, &value, sizeof(value)); + p += sizeof(value); + }; + + write16(msg_len); + write16(hello_ver); + write16(impl_ver); + write32(buffer_size); + write32(recv_buffer_cnt); + write32(jetty_id); + std::memcpy(p, eid, sizeof(eid)); + p += sizeof(eid); + write32(uasid); + *p++ = tp_type; + std::memset(p, 0, sizeof(pad)); + p += sizeof(pad); + std::memcpy(p, seg_eid, sizeof(seg_eid)); + p += sizeof(seg_eid); + write32(seg_uasid); + write64(seg_va); + write64(seg_len); + write32(seg_token_id); +} + +void HelloMessage::Deserialize(const void* buf) { + const uint8_t* p = static_cast(buf); + auto read16 = [&p]() { + uint16_t value; + std::memcpy(&value, p, sizeof(value)); + p += sizeof(value); + return butil::NetToHost16(value); + }; + auto read32 = [&p]() { + uint32_t value; + std::memcpy(&value, p, sizeof(value)); + p += sizeof(value); + return butil::NetToHost32(value); + }; + auto read64 = [&p]() { + uint64_t value; + std::memcpy(&value, p, sizeof(value)); + p += sizeof(value); + return butil::NetToHost64(value); + }; + + msg_len = read16(); + hello_ver = read16(); + impl_ver = read16(); + buffer_size = read32(); + recv_buffer_cnt = read32(); + jetty_id = read32(); + std::memcpy(eid, p, sizeof(eid)); + p += sizeof(eid); + uasid = read32(); + tp_type = *p++; + p += sizeof(pad); + std::memcpy(seg_eid, p, sizeof(seg_eid)); + p += sizeof(seg_eid); + seg_uasid = read32(); + seg_va = read64(); + seg_len = read64(); + seg_token_id = read32(); +} + +} // namespace v2_wire + +// ============================================================================ +// Shared helpers. +// ============================================================================ + +namespace { + +constexpr uint32_t MIN_BUFFER_SIZE = 1024; +// Three SQ entries are reserved for flow-control messages. The advertised +// receive count must leave at least one data WR after those reservations. +constexpr uint32_t MIN_BUFFER_CNT = 3; +constexpr uint32_t MAX_BUFFER_CNT = 65535; +constexpr uint32_t MAX_V3_PB_SIZE = 4096; + +} // namespace + +bool ValidHello(const ParsedHello& h) { + if (h.buffer_size < MIN_BUFFER_SIZE) { + return false; + } + if (h.recv_buffer_cnt < MIN_BUFFER_CNT || h.recv_buffer_cnt > MAX_BUFFER_CNT) { + return false; + } + if (h.jetty_id == 0) { + return false; + } + if (h.tp_type > static_cast(URMA_UTP)) { + return false; + } + if (h.seg_len == 0 || h.seg_va == 0) { + return false; + } + return true; +} + +// File-local (not in the anonymous namespace so it can be friend-declared +// from urma_endpoint.h's UrmaEndpoint). Reads the body following the magic +// and translates it into ParsedHello. +int ReadBodyAndNegotiate(UrmaEndpoint* ep, ParsedHello* out, bool* negotiated) { + *negotiated = false; + uint8_t body[v2_wire::HELLO_BODY_LEN]; + if (ep->ReadFromFd(body, v2_wire::HELLO_BODY_LEN) < 0) { + return -1; + } + v2_wire::HelloMessage m; + m.Deserialize(body); + if (m.msg_len < v2_wire::HELLO_MSG_LEN_MIN || + m.msg_len > v2_wire::HELLO_MSG_LEN_MAX || + m.hello_ver != v2_wire::HELLO_V2_VERSION || + m.impl_ver != v2_wire::IMPL_V2_VERSION) { return 0; } + ParsedHello p; + p.buffer_size = m.buffer_size; + p.recv_buffer_cnt = m.recv_buffer_cnt; + p.jetty_id = m.jetty_id; + std::memcpy(p.eid, m.eid, 16); + p.uasid = m.uasid; + p.tp_type = m.tp_type; + std::memcpy(p.seg_eid, m.seg_eid, 16); + p.seg_uasid = m.seg_uasid; + p.seg_va = m.seg_va; + p.seg_len = m.seg_len; + p.seg_token_id = m.seg_token_id; + if (!ValidHello(p)) { + return 0; + } + // Drain trailing bytes if msg_len advertises more than the fixed body. + if (m.msg_len > v2_wire::HELLO_PACKET_LEN) { + if (DrainBytes(ep, m.msg_len - v2_wire::HELLO_PACKET_LEN) < 0) { + return -1; + } + } + *out = p; + *negotiated = true; + return 0; +} + +int DrainBytes(UrmaEndpoint* ep, size_t n) { + char buf[4096]; + while (n > 0) { + size_t want = std::min(n, sizeof(buf)); + if (ep->ReadFromFd(buf, want) < 0) { + return -1; + } + n -= want; + } + return 0; +} + +// ============================================================================ +// v2 client / server. +// ============================================================================ + +int UrmaHandshakeClientV2::SendLocalHello() { + v2_wire::HelloMessage m; + _ep->FillLocalHelloV2(&m); + uint8_t packet[v2_wire::HELLO_PACKET_LEN]; + std::memcpy(packet, "URMA", 4); + m.Serialize(packet + 4); + return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); +} + +int UrmaHandshakeClientV2::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + *negotiated = false; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + return -1; + } + if (std::memcmp(magic, "URMA", 4) != 0) { + // Peer is not URMA-capable; push the magic back so the TCP input + // messenger can re-parse it. + _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + return 0; + } + return ReadBodyAndNegotiate(_ep, out, negotiated); +} + +int UrmaHandshakeServerV2::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + return ReadBodyAndNegotiate(_ep, out, negotiated); +} + +int UrmaHandshakeServerV2::SendLocalHello() { + v2_wire::HelloMessage m; + _ep->FillLocalHelloV2(&m); + auto* tp = static_cast(_ep->_socket->_transport.get()); + if (tp->_urma_state.load(butil::memory_order_acquire) == + UrmaTransport::URMA_OFF) { + // Tell the client we are not URMA-capable: zero the version fields so + // the client's version check fails and it falls back to TCP. + m.hello_ver = 0; + m.impl_ver = 0; + m.jetty_id = 0; + m.buffer_size = 0; + } + uint8_t packet[v2_wire::HELLO_PACKET_LEN]; + std::memcpy(packet, "URMA", 4); + m.Serialize(packet + 4); + return _ep->WriteToFd(packet, v2_wire::HELLO_PACKET_LEN); +} + +// ============================================================================ +// v3 protobuf ("URM3"). The protobuf (de)serialization lives on the endpoint +// because it touches UrmaEndpoint private state; the classes here just call +// FillLocalHelloV3 / WriteHelloV3 / ReadAndParseHelloV3. +// ============================================================================ + +int UrmaHandshakeClientV3::SendLocalHello() { + UrmaHello msg; + _ep->FillLocalHelloV3(&msg); + return _ep->WriteHelloV3(msg); +} + +int UrmaHandshakeClientV3::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + *negotiated = false; + uint8_t magic[v2_wire::MAGIC_STR_LEN]; + if (_ep->ReadFromFd(magic, v2_wire::MAGIC_STR_LEN) < 0) { + return -1; + } + if (std::memcmp(magic, "URM3", 4) != 0) { + _ep->PushBackToReadBuf(magic, v2_wire::MAGIC_STR_LEN); + return 0; + } + return _ep->ReadAndParseHelloV3(out, negotiated); +} + +int UrmaHandshakeServerV3::ReceiveAndParseRemoteHello(ParsedHello* out, + bool* negotiated) { + return _ep->ReadAndParseHelloV3(out, negotiated); +} + +int UrmaHandshakeServerV3::SendLocalHello() { + UrmaHello msg; + _ep->FillLocalHelloV3(&msg); + // v3 has no zero-out path; the client rejects jetty_id==0 via ValidHello. + return _ep->WriteHelloV3(msg); +} + +// ============================================================================ +// Factories. +// ============================================================================ + +UrmaHandshake* CreateClientHandshake(UrmaEndpoint* ep) { + switch (FLAGS_urma_client_handshake_version) { + case 3: return new UrmaHandshakeClientV3(ep); + case 2: + default: return new UrmaHandshakeClientV2(ep); + } +} + +UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, + const uint8_t magic[v2_wire::MAGIC_STR_LEN]) { + if (std::memcmp(magic, "URMA", 4) == 0) { + return new UrmaHandshakeServerV2(ep); + } + if (std::memcmp(magic, "URM3", 4) == 0) { + return new UrmaHandshakeServerV3(ep); + } + return nullptr; +} + +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA diff --git a/src/brpc/urma/urma_handshake.h b/src/brpc/urma/urma_handshake.h new file mode 100644 index 0000000000..f406bfa1b9 --- /dev/null +++ b/src/brpc/urma/urma_handshake.h @@ -0,0 +1,180 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef BRPC_URMA_HANDSHAKE_H +#define BRPC_URMA_HANDSHAKE_H + +#include +#include + +#if BRPC_WITH_URMA + +#include "urma_types.h" + +namespace brpc { +namespace urma { + +class UrmaEndpoint; + +// Wire-format-agnostic view of a peer's hello message. Both v2 (binary) and +// v3 (protobuf) translate the bytes on the wire into this struct. The endpoint +// then uses ApplyRemoteHello to size its send/recv windows and to drive +// urma_import_seg / urma_import_jetty. +struct ParsedHello { + uint32_t buffer_size = 0; // peer recv buffer size in bytes (per WR) + uint32_t recv_buffer_cnt = 0; // peer recv buffer count (RQ depth - 1) + uint32_t jetty_id = 0; // peer jetty id + uint8_t eid[16] = {0}; // peer EID (network order) + uint32_t uasid = 0; // peer uasid + uint8_t tp_type = 0; // urma_tp_type_t: URMA_RTP / URMA_CTP / URMA_UTP + + // Flattened peer buffer-pool segment. + uint8_t seg_eid[16] = {0}; // EID owning the peer segment + uint32_t seg_uasid = 0; // uasid owning the peer segment + uint64_t seg_va = 0; // segment virtual address + uint64_t seg_len = 0; // segment length in bytes + uint32_t seg_token_id = 0; // segment token id +}; + +// Validate all values that influence resource import and queue/window sizing. +// Both v2 and v3 handshakes must pass this check before negotiation succeeds. +bool ValidHello(const ParsedHello& hello); + +// v2 binary wire layout. The full on-wire packet is: +// [ "URMA" 4B ][ HelloMessage body 82B ] => 86 bytes total. +// (msg_len = 86 includes the magic prefix and the body, matching RDMA's +// convention where msg_len covers the whole packet.) +namespace v2_wire { + +constexpr size_t MAGIC_STR_LEN = 4; +constexpr size_t HELLO_BODY_LEN = 82; +constexpr size_t HELLO_PACKET_LEN = MAGIC_STR_LEN + HELLO_BODY_LEN; // 86 +constexpr size_t HELLO_MSG_LEN_MIN = HELLO_PACKET_LEN; +constexpr size_t HELLO_MSG_LEN_MAX = 4096; +constexpr uint16_t HELLO_V2_VERSION = 2; +constexpr uint16_t IMPL_V2_VERSION = 1; + +// The serializable struct. Aligned so it can be reinterpreted as raw bytes. +struct HelloMessage { + uint16_t msg_len; // total packet length (incl. magic) + uint16_t hello_ver; + uint16_t impl_ver; + uint32_t buffer_size; + uint32_t recv_buffer_cnt; + uint32_t jetty_id; + uint8_t eid[16]; + uint32_t uasid; + uint8_t tp_type; + uint8_t pad[3]; // keep the struct 4-byte aligned + uint8_t seg_eid[16]; + uint32_t seg_uasid; + uint64_t seg_va; + uint64_t seg_len; + uint32_t seg_token_id; + + void Serialize(void* buf) const; // host -> network order, write to buf + void Deserialize(const void* buf); // network -> host order, read from buf +}; + +} // namespace v2_wire + +// Abstract handshake strategy. v2 speaks binary; v3 speaks protobuf ("URM3"). +class UrmaHandshake { +public: + virtual ~UrmaHandshake() = default; + virtual int ProtocolVersion() const = 0; + + // Send our local hello over the TCP fd held by the endpoint. + // Returns 0 on success, -1 on failure (errno set). + virtual int SendLocalHello() = 0; + + // Read the peer's hello and parse it into @out. Sets @negotiated to true + // if the peer is URMA-capable and the message validates; false (with + // return 0) to signal a graceful fall-back to TCP. + // Returns -1 on IO error (errno set). + virtual int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) = 0; +}; + +// v2 binary handshake (magic "URMA"). +class UrmaHandshakeClientV2 : public UrmaHandshake { +public: + explicit UrmaHandshakeClientV2(UrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 2; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + UrmaEndpoint* _ep; +}; + +class UrmaHandshakeServerV2 : public UrmaHandshake { +public: + explicit UrmaHandshakeServerV2(UrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 2; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + UrmaEndpoint* _ep; +}; + +// v3 protobuf handshake (magic "URM3"). +class UrmaHandshakeClientV3 : public UrmaHandshake { +public: + explicit UrmaHandshakeClientV3(UrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 3; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + UrmaEndpoint* _ep; +}; + +class UrmaHandshakeServerV3 : public UrmaHandshake { +public: + explicit UrmaHandshakeServerV3(UrmaEndpoint* ep) : _ep(ep) {} + int ProtocolVersion() const override { return 3; } + int SendLocalHello() override; + int ReceiveAndParseRemoteHello(ParsedHello* out, bool* negotiated) override; + +private: + UrmaEndpoint* _ep; +}; + +// Client-side factory: picks v2/v3 based on --urma_client_handshake_version. +UrmaHandshake* CreateClientHandshake(UrmaEndpoint* ep); + +// Server-side factory: dispatches on the first 4 bytes read from the TCP fd. +// Returns nullptr if the magic is not "URMA"/"URM3" (caller falls back to TCP). +// @magic is the first MAGIC_STR_LEN bytes already read from the fd. +UrmaHandshake* CreateServerHandshakeByMagic(UrmaEndpoint* ep, + const uint8_t magic[v2_wire::MAGIC_STR_LEN]); + +// Drain @n bytes from the TCP fd (used when msg_len > HELLO_MSG_LEN_MIN). +int DrainBytes(UrmaEndpoint* ep, size_t n); + +// Read the body following the magic and translate to ParsedHello. Friend of +// UrmaEndpoint. Returns -1 on IO error; 0 with *negotiated=false on invalid +// (peer not URMA-capable); 0 with *negotiated=true on success. +int ReadBodyAndNegotiate(UrmaEndpoint* ep, ParsedHello* out, bool* negotiated); + +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA + +#endif // BRPC_URMA_HANDSHAKE_H diff --git a/src/brpc/urma/urma_handshake.proto b/src/brpc/urma/urma_handshake.proto new file mode 100644 index 0000000000..f6d4a6d1ea --- /dev/null +++ b/src/brpc/urma/urma_handshake.proto @@ -0,0 +1,47 @@ +// 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. + +syntax = "proto2"; + +package brpc.urma; + +option java_package = "com.brpc.urma"; +option java_outer_classname = "UrmaHandshakeProto"; + +// Wire-level handshake message exchanged between two UrmaTransport peers +// after the TCP connection is established. The first 4 bytes on the wire +// are the magic "URM3" (v3) followed by a 4-byte big-endian length prefix +// and then the protobuf-encoded UrmaHello. +message UrmaHello { + required uint32 buffer_size = 1; // recv buffer size in bytes (per WR) + required uint32 recv_buffer_cnt = 2; // number of recv buffers posted + required uint32 jetty_id = 3; // local jetty id + required bytes eid = 4; // 16-byte local EID (network order) + required uint32 uasid = 5; // local uasid + required uint32 tp_type = 6; // urma_tp_type_t + + // Flattened peer buffer-pool segment (urma_seg_t + token). The receiver + // calls urma_import_seg with these fields BEFORE urma_import_jetty, so + // the kernel establishes the transport path (TP) routing for the remote + // EID. Otherwise the first SEND is rejected by hardware with + // URMA_CR_RNR_RETRY_CNT_EXC_ERR (status=10). + required bytes seg_eid = 7; // 16-byte EID owning the segment + required uint32 seg_uasid = 8; // uasid owning the segment + required uint64 seg_va = 9; // segment virtual address + required uint64 seg_len = 10; // segment length in bytes + required uint32 seg_token_id = 11; // segment token id +} diff --git a/src/brpc/urma/urma_helper.cpp b/src/brpc/urma/urma_helper.cpp new file mode 100644 index 0000000000..ceeeae9a1b --- /dev/null +++ b/src/brpc/urma/urma_helper.cpp @@ -0,0 +1,788 @@ +// 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/urma/urma_helper.h" + +#if BRPC_WITH_URMA + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +#include "butil/atomicops.h" +#include "butil/containers/flat_map.h" +#include "butil/iobuf.h" +#include "butil/logging.h" +#include "butil/macros.h" +#include "butil/scoped_lock.h" +#include "butil/synchronization/lock.h" + +#include "urma_api.h" +#include "urma_types.h" +#include "brpc/urma/urma_bonding.h" +#include "brpc/urma/urma_endpoint.h" + +DECLARE_int32(task_group_ntags); + +namespace butil { +namespace iobuf { +// declared in iobuf.cpp +extern void* (*blockmem_allocate)(size_t); +extern void (*blockmem_deallocate)(void*); +} +} + +namespace brpc { +namespace urma { + +DEFINE_bool(urma_use_polling, false, + "Use busy polling to poll JFC, instead of event mode"); +DEFINE_int32(urma_poller_num, 1, + "Number of poller bthreads per bthread tag (polling mode only)"); +DEFINE_bool(urma_disable_bthread, false, + "Run the message-processing callback inline (no bthread spawned)"); + +DEFINE_int32(urma_sq_size, 128, + "Depth of the local send jetty (JFS). [16, 4096]"); +DEFINE_int32(urma_rq_size, 128, + "Depth of the local recv jetty (JFR). [16, 4096]"); +DEFINE_int32(urma_cqe_poll_once, 32, + "Max completion entries polled per urma_poll_jfc call"); +DEFINE_bool(urma_recv_zerocopy, true, + "Use zero-copy for receives larger than --urma_zerocopy_min_size"); +DEFINE_int32(urma_zerocopy_min_size, 512, + "Receives smaller than this many bytes are copied (not zero-copy)"); + +DEFINE_string(urma_device, "", + "The name of the URMA device to use. Empty means the first one."); +DEFINE_int32(urma_max_sge, 0, + "Max SGEs per WR. 0 means the device maximum."); +DEFINE_int32(urma_bonding_mode, 0, + "Bonding mode for bonding devices: 0=standalone, " + "1=active-backup, 2=balance."); +DEFINE_int32(urma_bonding_level, 0, + "Bonding level for bonding devices: 0=IODIE, 1=port."); +DEFINE_int32(urma_prepared_jetty_cnt, 8, + "Requested number of pre-allocated Jetty+CQ sets for fast " + "connect; capped automatically according to RLIMIT_NOFILE"); + +DEFINE_int32(urma_buffer_size, 8 * 1024, + "Per-buffer size in the URMA buffer pool (bytes). " + "Must match IOBuf block size to keep zero-copy working."); +DEFINE_int32(urma_buffer_count, 65536, + "Number of buffers in the URMA buffer pool."); + +DEFINE_bool(urma_poller_yield, false, + "Yield (bthread_yield) in the busy poll loop to let other " + "bthreads run"); + +constexpr size_t kIOBufBlockHeaderLen = 32; + +// Set to true to skip real URMA hardware initialization (unit tests). When +// true, GlobalUrmaInitializeOrDie() returns without touching liburma and the +// endpoint builds its state machine without posting real WRs. +bool g_skip_urma_init = false; +butil::atomic g_urma_available(false); + +// ============================================================================ +// Global URMA state (single device / single context chosen at init time). +// ============================================================================ + +static urma_device_t* g_device = nullptr; +static urma_context_t* g_context = nullptr; +static urma_eid_t g_local_eid{}; +static bool g_has_local_eid = false; +static urma_device_attr_t g_device_attr{}; +static int g_max_sge = 1; +static size_t g_recv_block_size = 8 * 1024; +static bool g_is_bonding_device = false; +// Prefer the device capability table and retain priority 6 as a compatibility +// fallback for CTP providers that do not report a priority. +static uint8_t g_jetty_priority = 6; +// urma_init/urma_uninit manage process-global liburma state. Only uninitialize +// it when this helper performed the successful initialization; URMA_EEXIST +// means another component owns that state. +static bool g_owns_urma_init = false; + +// The single registered segment backing the buffer pool. The whole pool is +// one urma_register_seg call, sliced into fixed-size buffers. urma_target_seg_t +// is the per-buffer handle carried by urma_sge_t.tseg on the send/recv path. +static urma_target_seg_t* g_pool_seg = nullptr; +static void* g_pool_base = nullptr; +static size_t g_pool_size = 0; +static size_t g_pool_buffer_size = 0; + +// User-registered segments (RegisterMemoryForUrma). Keyed by buffer address. +struct UserSeg { + urma_target_seg_t* tseg = nullptr; + void* base = nullptr; + size_t len = 0; +}; +static butil::FlatMap* g_user_segs = nullptr; +static butil::Mutex* g_user_segs_lock = nullptr; + +// Original IOBuf allocator (saved so we can restore it on release). +static void* (*g_mem_alloc_orig)(size_t) = nullptr; +static void (*g_mem_dealloc_orig)(void*) = nullptr; +static size_t g_default_block_size_orig = 0; + +namespace { + +// Round up to the page size. +size_t PageSize() { + long ps = sysconf(_SC_PAGESIZE); + return ps > 0 ? static_cast(ps) : 4096; +} + +size_t AlignUp(size_t v, size_t align) { + return (v + align - 1) / align * align; +} + +bool IsBondingDeviceName(const char* name) { + return name != nullptr && strncmp(name, "bonding", 7) == 0; +} + +union urma_tp_type_en TpTypeCapability(urma_tp_type_t tp_type) { + union urma_tp_type_en capability{}; + switch (tp_type) { + case URMA_RTP: + capability.bs.rtp = 1; + break; + case URMA_CTP: + capability.bs.ctp = 1; + break; + case URMA_UTP: + capability.bs.utp = 1; + break; + } + return capability; +} + +bool ConfigureBondingMode(const std::string& device_name) { + const char* context_device_name = + g_context != nullptr && g_context->dev != nullptr + ? g_context->dev->name + : nullptr; + g_is_bonding_device = + IsBondingDeviceName(device_name.c_str()) || + IsBondingDeviceName(context_device_name); + if (!g_is_bonding_device) { + return true; + } + +#if BRPC_URMA_HAS_BONDING_EXT + if (FLAGS_urma_bonding_mode < 0 || + FLAGS_urma_bonding_mode >= BONDP_BONDING_MODE_MAX || + FLAGS_urma_bonding_level < 0 || + FLAGS_urma_bonding_level >= BONDP_BONDING_LEVEL_MAX) { + LOG(ERROR) << "Invalid URMA bonding configuration: mode=" + << FLAGS_urma_bonding_mode + << " level=" << FLAGS_urma_bonding_level; + errno = EINVAL; + return false; + } + + bondp_set_bonding_mode_in_t bond_in{}; + bond_in.bonding_mode = + static_cast(FLAGS_urma_bonding_mode); + bond_in.bonding_level = + static_cast(FLAGS_urma_bonding_level); + + urma_user_ctl_in_t ctl_in{}; + ctl_in.addr = reinterpret_cast(&bond_in); + ctl_in.len = static_cast(sizeof(bond_in)); + ctl_in.opcode = BONDP_USER_CTL_SET_BONDING_MODE; + urma_user_ctl_out_t ctl_out{}; + const urma_status_t status = urma_user_ctl(g_context, &ctl_in, &ctl_out); + if (status != URMA_SUCCESS) { + LOG(ERROR) << "urma_user_ctl(SET_BONDING_MODE) failed: status=" + << status << " device=" << device_name + << " mode=" << FLAGS_urma_bonding_mode + << " level=" << FLAGS_urma_bonding_level + << ". It must run before segment/JFC/JFR creation"; + errno = status > 0 ? status : EIO; + return false; + } + return true; +#else + LOG(ERROR) << "URMA bonding device " << device_name + << " requires provider header urma_ubagg.h"; + errno = ENOTSUP; + return false; +#endif +} + +} // namespace + +// ============================================================================ +// Buffer pool: one registered segment, sliced into fixed-size buffers. +// ============================================================================ + +namespace { + +// Shard the free list to reduce contention between allocator threads. +constexpr size_t kShardCount = 64; +struct BufferPool { + butil::Mutex mutexes[kShardCount]; + std::vector free_lists[kShardCount]; + std::vector in_use; // 0/1 per buffer + butil::atomic outstanding{0}; + + size_t buffer_count() const { + return in_use.size(); + } +}; +BufferPool* g_pool = nullptr; + +size_t ShardFor(void* buf) { + auto* base = static_cast(buf); + auto offset = static_cast(base - static_cast(g_pool_base)); + auto idx = offset / g_pool_buffer_size; + return idx % kShardCount; +} + +size_t PreferredShard() { + // Hash the current thread id across shards. pthread_self() returns an + // opaque pthread_t; cast through uintptr_t to get a hashable value. + auto tid = static_cast(reinterpret_cast(pthread_self())); + return tid % kShardCount; +} + +void* PoolAllocate(size_t size) { + if (BAIDU_UNLIKELY(g_skip_urma_init)) { + return g_mem_alloc_orig ? g_mem_alloc_orig(size) : malloc(size); + } + // Only serve the configured buffer size; callers always ask for that. + if (size > g_pool_buffer_size) { + // Larger than a single buffer -- fall back to the system allocator. + return g_mem_alloc_orig ? g_mem_alloc_orig(size) : malloc(size); + } + auto start = PreferredShard(); + for (size_t i = 0; i < kShardCount; ++i) { + auto shard = (start + i) % kShardCount; + BAIDU_SCOPED_LOCK(g_pool->mutexes[shard]); + auto& fl = g_pool->free_lists[shard]; + if (fl.empty()) { + continue; + } + void* buf = fl.back(); + fl.pop_back(); + auto idx = (static_cast(buf) - + static_cast(g_pool_base)) / g_pool_buffer_size; + if (idx < g_pool->in_use.size()) { + g_pool->in_use[idx] = 1; + } + g_pool->outstanding.fetch_add(1, butil::memory_order_relaxed); + return buf; + } + LOG_EVERY_SECOND(WARNING) + << "URMA buffer pool exhausted; falling back to malloc"; + return g_mem_alloc_orig ? g_mem_alloc_orig(size) : malloc(size); +} + +void PoolDeallocate(void* buf) { + if (BAIDU_UNLIKELY(g_skip_urma_init)) { + if (g_mem_dealloc_orig) { + g_mem_dealloc_orig(buf); + } else { + free(buf); + } + return; + } + auto* base = static_cast(g_pool_base); + auto* p = static_cast(buf); + if (!base || p < base || p >= base + g_pool_size || + static_cast(p - base) % g_pool_buffer_size != 0) { + // Not a pool buffer -- hand back to the original allocator. + if (g_mem_dealloc_orig) { + g_mem_dealloc_orig(buf); + } else { + free(buf); + } + return; + } + auto shard = ShardFor(buf); + BAIDU_SCOPED_LOCK(g_pool->mutexes[shard]); + auto idx = static_cast(p - base) / g_pool_buffer_size; + if (idx < g_pool->in_use.size()) { + if (!g_pool->in_use[idx]) { + LOG(WARNING) << "double-free of URMA pool buffer " << buf; + return; + } + g_pool->in_use[idx] = 0; + } + g_pool->free_lists[shard].push_back(buf); + if (g_pool->outstanding.load(butil::memory_order_relaxed) > 0) { + g_pool->outstanding.fetch_sub(1, butil::memory_order_relaxed); + } +} + +// Register the pool: mmap one large region and urma_register_seg it. +bool InitPool() { + if (g_pool_buffer_size == 0 || g_pool == nullptr) { + return false; + } + size_t count = g_pool->buffer_count(); + if (count == 0) { + return false; + } + size_t raw = g_pool_buffer_size * count; + size_t page = PageSize(); + g_pool_size = AlignUp(raw, page); + + g_pool_base = mmap(nullptr, g_pool_size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (g_pool_base == MAP_FAILED) { + PLOG(WARNING) << "Fail to mmap URMA buffer pool"; + g_pool_base = nullptr; + return false; + } + + urma_reg_seg_flag_t flag{}; + flag.bs.token_policy = URMA_TOKEN_NONE; + flag.bs.cacheable = URMA_NON_CACHEABLE; + flag.bs.access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC; + + urma_seg_cfg_t cfg{}; + cfg.va = reinterpret_cast(g_pool_base); + cfg.len = g_pool_size; + cfg.token_id = nullptr; + cfg.token_value = {}; + cfg.flag = flag; + cfg.user_ctx = reinterpret_cast(g_pool_base); + cfg.iova = 0; + + errno = 0; + g_pool_seg = urma_register_seg(g_context, &cfg); + if (!g_pool_seg) { + PLOG(WARNING) << "Fail to urma_register_seg"; + munmap(g_pool_base, g_pool_size); + g_pool_base = nullptr; + return false; + } + g_pool->in_use.assign(count, 0); + for (size_t i = 0; i < count; ++i) { + auto shard = i % kShardCount; + g_pool->free_lists[shard].push_back( + static_cast(g_pool_base) + i * g_pool_buffer_size); + } + return true; +} + +} // namespace + +// Exposed to urma_endpoint.cpp: return the per-buffer target_seg pointer. +// All pool buffers share the same segment, so we return g_pool_seg for any +// pool address; the per-WR length selects the slice. +urma_target_seg_t* GetPoolSegFor(void* buf) { + if (g_skip_urma_init || !g_pool_seg || !g_pool_base) { + return nullptr; + } + if (buf == nullptr) { + return g_pool_seg; + } + const uintptr_t base = reinterpret_cast(g_pool_base); + const uintptr_t p = reinterpret_cast(buf); + if (p >= base && p - base < g_pool_size) { + return g_pool_seg; + } + return nullptr; +} + +static void GlobalRelease() { + g_urma_available.store(false, butil::memory_order_release); + if (g_mem_alloc_orig) { + butil::iobuf::blockmem_allocate = g_mem_alloc_orig; + g_mem_alloc_orig = nullptr; + } + if (g_mem_dealloc_orig) { + butil::iobuf::blockmem_deallocate = g_mem_dealloc_orig; + g_mem_dealloc_orig = nullptr; + } + if (g_default_block_size_orig != 0) { + butil::SetDefaultBlockSize(g_default_block_size_orig); + g_default_block_size_orig = 0; + } + UrmaEndpoint::GlobalRelease(); + if (g_pool_seg) { + urma_unregister_seg(g_pool_seg); + g_pool_seg = nullptr; + } + if (g_pool_base) { + munmap(g_pool_base, g_pool_size); + g_pool_base = nullptr; + g_pool_size = 0; + } + delete g_pool; + g_pool = nullptr; + delete g_user_segs; + g_user_segs = nullptr; + delete g_user_segs_lock; + g_user_segs_lock = nullptr; + if (g_context) { + urma_delete_context(g_context); + g_context = nullptr; + } + g_local_eid = urma_eid_t{}; + g_has_local_eid = false; + g_is_bonding_device = false; + g_jetty_priority = 6; + if (g_owns_urma_init) { + const urma_status_t status = urma_uninit(); + if (status != URMA_SUCCESS) { + LOG(WARNING) << "Fail to urma_uninit: " << status; + } + g_owns_urma_init = false; + } + g_device = nullptr; +} + +// ============================================================================ +// Global initialization. +// ============================================================================ + +static bool GlobalUrmaInitializeImpl() { + if (BAIDU_UNLIKELY(g_skip_urma_init)) { + g_urma_available.store(true, butil::memory_order_release); + return true; + } + if (FLAGS_urma_sq_size < 16 || FLAGS_urma_sq_size > 4096 || + FLAGS_urma_rq_size < 16 || FLAGS_urma_rq_size > 4096 || + FLAGS_urma_buffer_size < 1024 || FLAGS_urma_buffer_count <= 0 || + FLAGS_urma_poller_num <= 0) { + LOG(ERROR) << "Invalid URMA queue, buffer, or poller configuration"; + errno = EINVAL; + return false; + } + + urma_init_attr_t init_attr{}; + const urma_status_t status = urma_init(&init_attr); + if (status != URMA_SUCCESS && status != URMA_EEXIST) { + if (status == URMA_FAIL) { + LOG(ERROR) << "Fail to urma_init: " << status + << " (URMA_FAIL). liburma returns URMA_FAIL when it " + "cannot load a provider, or when URMA was already " + "initialized by another component. Verify readable " + "provider libraries under /usr/lib64/urma, loaded " + "URMA kernel drivers, and that urma_init is called " + "only once per process"; + } else { + LOG(ERROR) << "Fail to urma_init: " << status; + } + return false; + } + g_owns_urma_init = (status == URMA_SUCCESS); + + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + if (!devices || num_devices <= 0) { + LOG(ERROR) << "No URMA device found"; + urma_free_device_list(devices); + return false; + } + urma_device_t* found = nullptr; + for (int i = 0; i < num_devices; ++i) { + if (FLAGS_urma_device.empty() || + std::string(devices[i]->name) == FLAGS_urma_device) { + found = devices[i]; + break; + } + } + if (!found) { + LOG(ERROR) << "URMA device not found: " << FLAGS_urma_device; + urma_free_device_list(devices); + return false; + } + g_device = found; + + const std::string device_name = found->name; + if (urma_query_device(found, &g_device_attr) != URMA_SUCCESS) { + LOG(ERROR) << "Fail to urma_query_device"; + urma_free_device_list(devices); + g_device = nullptr; + return false; + } + const int ctp_priority = + FindUrmaPriorityForTpType(g_device_attr, URMA_CTP); + if (ctp_priority >= 0) { + g_jetty_priority = static_cast(ctp_priority); + } else { + LOG(WARNING) << "URMA device does not report a CTP priority; " + "falling back to compatibility priority " + << static_cast(g_jetty_priority); + } + uint32_t eid_cnt = 0; + urma_eid_info_t* eids = urma_get_eid_list(found, &eid_cnt); + if (!eids || eid_cnt == 0) { + LOG(ERROR) << "Fail to urma_get_eid_list"; + urma_free_eid_list(eids); + urma_free_device_list(devices); + g_device = nullptr; + return false; + } + // Retain the exact EID used to create the context and advertise it in the + // handshake. This matters for a bonding virtual device whose jetty can + // carry a provider-selected physical EID. + g_local_eid = eids[0].eid; + g_has_local_eid = true; + g_context = urma_create_context(found, eids[0].eid_index); + urma_free_eid_list(eids); + urma_free_device_list(devices); + g_device = nullptr; + if (!g_context) { + LOG(ERROR) << "Fail to urma_create_context"; + return false; + } + // The bonding provider only accepts SET_BONDING_MODE while the context + // has no dependent resource. This must precede register_seg/JFC/JFR. + if (!ConfigureBondingMode(device_name)) { + return false; + } + uint32_t device_max_sge = g_device_attr.dev_cap.max_jfs_sge; + if (device_max_sge == 0) { + device_max_sge = 1; + } + // urma_jfs_cfg_t::max_sge is uint8_t. + if (device_max_sge > 255) { + device_max_sge = 255; + } + g_max_sge = static_cast(device_max_sge); + if (FLAGS_urma_max_sge > 0) { + if (FLAGS_urma_max_sge > g_max_sge) { + LOG(WARNING) << "Cap urma_max_sge from " << FLAGS_urma_max_sge + << " to device/config limit " << g_max_sge; + } else { + g_max_sge = FLAGS_urma_max_sge; + } + } + g_recv_block_size = + static_cast(FLAGS_urma_buffer_size) - + kIOBufBlockHeaderLen; + + // User-segment table. + g_user_segs_lock = new (std::nothrow) butil::Mutex; + g_user_segs = new (std::nothrow) butil::FlatMap(); + if (!g_user_segs_lock || !g_user_segs || + g_user_segs->init(65536) < 0) { + LOG(ERROR) << "Fail to init g_user_segs"; + return false; + } + + // Buffer pool. + g_pool = new BufferPool(); + g_pool_buffer_size = static_cast(FLAGS_urma_buffer_size); + // Resize the pool's per-shard vectors to hold the configured count. + size_t count = static_cast(FLAGS_urma_buffer_count); + g_pool->in_use.assign(count, 0); + for (size_t s = 0; s < kShardCount; ++s) { + g_pool->free_lists[s].reserve(count / kShardCount + 1); + } + if (!InitPool()) { + LOG(ERROR) << "Fail to init URMA buffer pool"; + return false; + } + + // Hijack IOBuf allocation so every IOBuf block is backed by a registered + // segment. This makes the send path trivial: any IOBuf can be posted + // directly as an urma_sge_t pointing at g_pool_seg. + g_mem_alloc_orig = butil::iobuf::blockmem_allocate; + g_mem_dealloc_orig = butil::iobuf::blockmem_deallocate; + g_default_block_size_orig = butil::GetDefaultBlockSize(); + butil::iobuf::blockmem_allocate = PoolAllocate; + butil::iobuf::blockmem_deallocate = PoolDeallocate; + butil::SetDefaultBlockSize(g_pool_buffer_size); + + if (UrmaEndpoint::GlobalInitialize() != 0) { + LOG(ERROR) << "Fail to initialize URMA endpoint resources"; + return false; + } + + g_urma_available.store(true, butil::memory_order_release); + // Do not register GlobalRelease with atexit. IOBuf keeps blocks in + // thread-local chains whose destructors may run after atexit handlers. + // Unmapping the registered pool here would leave those TLS chains + // pointing into unmapped memory. The process reclaims global URMA + // resources on exit; GlobalRelease remains available for init rollback. + LOG(INFO) << "URMA initialized: device=" << device_name + << " bonding=" << g_is_bonding_device + << " max_sge=" << g_max_sge + << " buffer_size=" << g_pool_buffer_size + << " buffer_count=" << g_pool->buffer_count(); + return true; +} + +static butil::atomic g_init_once{0}; +static butil::Mutex g_init_mutex; + +void GlobalUrmaInitializeOrDie() { + int expected = 0; + if (g_init_once.load(butil::memory_order_acquire) == 2) { + return; + } + if (g_init_once.compare_exchange_strong(expected, 1, + butil::memory_order_acq_rel)) { + BAIDU_SCOPED_LOCK(g_init_mutex); + if (!GlobalUrmaInitializeImpl()) { + LOG(WARNING) << "URMA initialization failed; falling back to TCP"; + GlobalRelease(); + } + g_init_once.store(2, butil::memory_order_release); + } else { + // Wait for the other thread to finish init. + while (g_init_once.load(butil::memory_order_acquire) != 2) { + // spin briefly + } + } +} + +bool IsUrmaAvailable() { + return g_urma_available.load(butil::memory_order_acquire); +} + +void GlobalDisableUrma() { + g_urma_available.store(false, butil::memory_order_release); +} + +bool SupportedByUrma(const std::string& protocol) { + return protocol == "baidu_std"; +} + +urma_context_t* GetUrmaContext() { return g_context; } +const urma_eid_t* GetUrmaLocalEid() { + return g_has_local_eid ? &g_local_eid : nullptr; +} +bool IsUrmaBondingDevice() { return g_is_bonding_device; } +int FindUrmaPriorityForTpType(const urma_device_attr_t& attr, + urma_tp_type_t tp_type) { + const union urma_tp_type_en expected = TpTypeCapability(tp_type); + for (int priority = 0; priority <= URMA_MAX_PRIORITY; ++priority) { + if (attr.dev_cap.priority_info[priority].tp_type.value == + expected.value) { + return priority; + } + } + return -1; +} +uint8_t GetUrmaJettyPriority() { return g_jetty_priority; } +int GetUrmaMaxSge() { return g_max_sge; } +size_t GetUrmaRecvBlockSize() { return g_recv_block_size; } + +// ============================================================================ +// Polling mode (per bthread tag). +// ============================================================================ + +bool InitPollingModeWithTag(bthread_tag_t tag, + std::function callback, + std::function init_fn, + std::function release_fn) { + if (BAIDU_UNLIKELY(g_skip_urma_init)) { + return true; + } + return UrmaEndpoint::PollingModeInitialize( + tag, std::move(callback), std::move(init_fn), + std::move(release_fn)) == 0; +} + +void ReleasePollingModeWithTag(bthread_tag_t tag) { + UrmaEndpoint::PollingModeRelease(tag); +} + +// ============================================================================ +// User memory registration. +// ============================================================================ + +uint64_t RegisterMemoryForUrma(void* buf, size_t len) { + if (BAIDU_UNLIKELY(g_skip_urma_init) || !g_context) { + return 0; + } + urma_reg_seg_flag_t flag{}; + flag.bs.token_policy = URMA_TOKEN_NONE; + flag.bs.cacheable = URMA_NON_CACHEABLE; + flag.bs.access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC; + + urma_seg_cfg_t cfg{}; + cfg.va = reinterpret_cast(buf); + cfg.len = len; + cfg.token_id = nullptr; + cfg.token_value = {}; + cfg.flag = flag; + cfg.user_ctx = reinterpret_cast(buf); + cfg.iova = 0; + + errno = 0; + urma_target_seg_t* tseg = urma_register_seg(g_context, &cfg); + if (!tseg) { + PLOG(WARNING) << "Fail to urma_register_seg for user memory"; + return 0; + } + BAIDU_SCOPED_LOCK(*g_user_segs_lock); + UserSeg us; + us.tseg = tseg; + us.base = buf; + us.len = len; + if (!g_user_segs->insert(buf, us)) { + LOG(WARNING) << "Fail to insert user seg (duplicate?)"; + urma_unregister_seg(tseg); + return 0; + } + return static_cast(reinterpret_cast(tseg)); +} + +void DeregisterMemoryForUrma(void* buf) { + if (BAIDU_UNLIKELY(g_skip_urma_init) || !g_user_segs) { + return; + } + BAIDU_SCOPED_LOCK(*g_user_segs_lock); + UserSeg* us = g_user_segs->seek(buf); + if (!us) { + return; + } + urma_unregister_seg(us->tseg); + g_user_segs->erase(buf); +} + +} // namespace urma +} // namespace brpc + +#else // BRPC_WITH_URMA + +#include + +#include "butil/logging.h" + +namespace brpc { +namespace urma { + +void GlobalUrmaInitializeOrDie() { + LOG(FATAL) << "URMA is not compiled in. Rebuild with -DWITH_URMA=ON."; + exit(1); +} + +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA diff --git a/src/brpc/urma/urma_helper.h b/src/brpc/urma/urma_helper.h new file mode 100644 index 0000000000..871f1a6354 --- /dev/null +++ b/src/brpc/urma/urma_helper.h @@ -0,0 +1,119 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef BRPC_URMA_HELPER_H +#define BRPC_URMA_HELPER_H + +#include +#include +#include +#include + +#include "bthread/types.h" +#include "butil/atomicops.h" + +#if BRPC_WITH_URMA + +#include "urma_api.h" +#include "urma_types.h" + +namespace brpc { +DECLARE_bool(usercode_in_coroutine); +DECLARE_bool(usercode_in_pthread); +namespace urma { + +// Initialize the URMA environment. Failure disables URMA globally so +// individual sockets transparently fall back to TCP. +void GlobalUrmaInitializeOrDie(); + +// Initialize URMA polling mode for a given bthread tag. +// Returns false on failure. +bool InitPollingModeWithTag(bthread_tag_t tag, + std::function callback = nullptr, + std::function init_fn = nullptr, + std::function release_fn = nullptr); + +void ReleasePollingModeWithTag(bthread_tag_t tag); + +// Register the given user buffer for URMA access. +// Returns the (opaque, non-zero) target segment handle stored in the user-mr +// table; 0 on failure. To use the memory in an IOBuf, append it via +// append_user_data_with_meta and pass the returned handle as the data meta. +uint64_t RegisterMemoryForUrma(void* buf, size_t len); + +// Deregister a previously registered user buffer. +void DeregisterMemoryForUrma(void* buf); + +// Return the target segment for a buffer-pool address. Passing nullptr returns +// the segment backing the whole pool. Returns nullptr for any other address. +urma_target_seg_t* GetPoolSegFor(void* buf); + +// Get the global URMA context (the urma_context_t created on the selected +// device / EID). Returns nullptr if URMA is not initialized. +urma_context_t* GetUrmaContext(); + +// Get the EID selected when the global context was created. This is the +// device EID that must be advertised to peers, especially for bonding +// devices where a created jetty may expose a provider-specific physical EID. +// Returns nullptr if URMA is not initialized. +const urma_eid_t* GetUrmaLocalEid(); + +// Return true when the selected URMA device is a bonding provider device. +bool IsUrmaBondingDevice(); + +// Find the priority whose advertised transport-path capability exactly +// matches @tp_type. Returns -1 when the device does not report one. +int FindUrmaPriorityForTpType(const urma_device_attr_t& attr, + urma_tp_type_t tp_type); + +// Return the priority selected for the CTP jettys created by brpc. +uint8_t GetUrmaJettyPriority(); + +// If the URMA environment is available. +bool IsUrmaAvailable(); + +// Disable URMA for the remaining lifetime of the process. +void GlobalDisableUrma(); + +// If the given protocol is supported by UrmaTransport. +// Currently only "baidu_std" is supported. +bool SupportedByUrma(const std::string& protocol); + +// Return the configured recv buffer size (one URMA recv WR's payload size). +size_t GetUrmaRecvBlockSize(); + +// Return max_sge supported by the device. +int GetUrmaMaxSge(); + +} // namespace urma +} // namespace brpc + +#else // BRPC_WITH_URMA + +namespace brpc { +namespace urma { + +// Initialize the URMA environment. +// Exit the process if initialization fails. +void GlobalUrmaInitializeOrDie(); + +} // namespace urma +} // namespace brpc + +#endif // BRPC_WITH_URMA + +#endif // BRPC_URMA_HELPER_H diff --git a/src/brpc/urma_transport.cpp b/src/brpc/urma_transport.cpp new file mode 100644 index 0000000000..311786e8eb --- /dev/null +++ b/src/brpc/urma_transport.cpp @@ -0,0 +1,248 @@ +// 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/urma_transport.h" + +#if BRPC_WITH_URMA + +#include + +#include "butil/iobuf.h" +#include "butil/logging.h" +#include "bthread/bthread.h" +#include "bthread/types.h" + +#include "brpc/event_dispatcher.h" +#include "brpc/input_messenger.h" +#include "brpc/socket.h" +#include "brpc/tcp_transport.h" +#include "brpc/urma/urma_endpoint.h" +#include "brpc/urma/urma_helper.h" + +namespace brpc { + +// Defined in urma_helper.cpp. +DECLARE_bool(urma_use_polling); +DECLARE_bool(urma_disable_bthread); + +void UrmaTransport::Init(Socket* socket, const SocketOptions& options) { + CHECK(_urma_ep == nullptr); + if (options.socket_mode == SOCKET_MODE_URMA) { + _urma_ep = new (std::nothrow) urma::UrmaEndpoint(socket); + if (!_urma_ep) { + const int saved_errno = errno; + PLOG(ERROR) << "Fail to create UrmaEndpoint"; + socket->SetFailed(saved_errno, "Fail to create UrmaEndpoint: %s", + berror(saved_errno)); + } + _urma_state = URMA_UNKNOWN; + } else { + _urma_state = URMA_OFF; + socket->_socket_mode = SOCKET_MODE_TCP; + } + _socket = socket; + _default_connect = options.app_connect; + _on_edge_trigger = options.on_edge_triggered_events; + if (options.need_on_edge_trigger && _on_edge_trigger == nullptr) { + _on_edge_trigger = urma::UrmaEndpoint::OnNewDataFromTcp; + } + _tcp_transport = std::make_shared(); + _tcp_transport->Init(socket, options); +} + +void UrmaTransport::Release() { + if (_urma_ep) { + delete _urma_ep; + _urma_ep = nullptr; + } +} + +int UrmaTransport::Reset(int32_t /*expected_nref*/) { + if (_urma_ep) { + _urma_ep->Reset(); + } + _urma_state = URMA_UNKNOWN; + return 0; +} + +std::shared_ptr UrmaTransport::Connect() { + if (_default_connect == nullptr) { + return std::make_shared(); + } + return _default_connect; +} + +int UrmaTransport::CutFromIOBuf(butil::IOBuf* buf) { + if (_urma_ep && + _urma_state.load(butil::memory_order_acquire) != URMA_OFF) { + butil::IOBuf* data_arr[1] = {buf}; + return _urma_ep->CutFromIOBufList(data_arr, 1); + } else { + return _tcp_transport->CutFromIOBuf(buf); + } +} + +ssize_t UrmaTransport::CutFromIOBufList(butil::IOBuf** buf, size_t ndata) { + if (_urma_ep && + _urma_state.load(butil::memory_order_acquire) != URMA_OFF) { + return _urma_ep->CutFromIOBufList(buf, ndata); + } else { + return _tcp_transport->CutFromIOBufList(buf, ndata); + } +} + +int UrmaTransport::WaitEpollOut(butil::atomic* epollout_butex, + bool pollin, const timespec duetime) { + if (_urma_state.load(butil::memory_order_acquire) == URMA_ON) { + const int expected_val = + epollout_butex->load(butil::memory_order_acquire); + CHECK(_urma_ep != nullptr); + if (!_urma_ep->IsWritable()) { + // Same caveat as RDMA: URMA cannot detect failure by writing, so + // after a failed butex wait we must re-check _socket->Failed(). + int rc = + bthread::butex_wait(epollout_butex, expected_val, &duetime); + if (rc < 0 && errno != EWOULDBLOCK && errno != ETIMEDOUT) { + const int saved_errno = errno; + PLOG(ERROR) << "Fail to wait epollout butex"; + if (_socket->SetFailed(saved_errno, "Fail to wait epollout butex: %s", + berror(saved_errno))) { + return 1; + } + } + if (_socket->Failed()) { + return 1; + } + return 0; + } + return 0; + } + return _tcp_transport->WaitEpollOut(epollout_butex, pollin, duetime); +} + +void UrmaTransport::ProcessEvent(bthread_attr_t attr) { + // Identical to TcpTransport/RdmaTransport: dispatch OnEdge(_socket) on a + // bthread, falling back to inline invocation on bthread_start failure. + bthread_t tid; + if (FLAGS_usercode_in_coroutine) { + OnEdge(_socket); + } else if (!EventDispatcherUnsched()) { + auto rc = bthread_start_urgent(&tid, &attr, OnEdge, _socket); + if (rc != 0) { + LOG(FATAL) << "Fail to start ProcessEvent"; + OnEdge(_socket); + } + } else if (bthread_start_background(&tid, &attr, OnEdge, _socket) != 0) { + LOG(FATAL) << "Fail to start ProcessEvent"; + OnEdge(_socket); + } +} + +void UrmaTransport::QueueMessage(InputMessageClosure& input_msg, + int* num_bthread_created, bool last_msg) { + if (last_msg && !urma::FLAGS_urma_use_polling) { + return; + } + InputMessageBase* to_run_msg = input_msg.release(); + if (!to_run_msg) { + return; + } + if (urma::FLAGS_urma_disable_bthread) { + Transport::ProcessInputMessage(to_run_msg); + return; + } + bthread_t th; + bthread_attr_t tmp = + (FLAGS_usercode_in_pthread ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL) | + BTHREAD_NOSIGNAL; + tmp.keytable_pool = _socket->keytable_pool(); + tmp.tag = bthread_self_tag(); + bthread_attr_set_name(&tmp, "ProcessInputMessage"); + if (!FLAGS_usercode_in_coroutine && bthread_start_background( + &th, &tmp, Transport::ProcessInputMessage, to_run_msg) == 0) { + ++*num_bthread_created; + } else { + Transport::ProcessInputMessage(to_run_msg); + } +} + +void UrmaTransport::Debug(std::ostream& os) { + if (_urma_state.load(butil::memory_order_acquire) == URMA_ON && + _urma_ep) { + _urma_ep->DebugInfo(os); + } +} + +int UrmaTransport::ContextInitOrDie(bool server_or_not, const void* options) { + if (server_or_not) { + if (!OptionsAvailableOverUrma( + static_cast(options))) { + return -1; + } + urma::GlobalUrmaInitializeOrDie(); + if (!urma::InitPollingModeWithTag( + static_cast(options)->bthread_tag)) { + return -1; + } + } else { + if (!OptionsAvailableForUrma( + static_cast(options))) { + return -1; + } + urma::GlobalUrmaInitializeOrDie(); + if (!urma::InitPollingModeWithTag(bthread_self_tag())) { + return -1; + } + } + return 0; +} + +bool UrmaTransport::OptionsAvailableForUrma(const ChannelOptions* opt) { + if (opt->has_ssl_options()) { + LOG(WARNING) << "Cannot use SSL and URMA at the same time"; + return false; + } + if (!urma::SupportedByUrma(opt->protocol.name())) { + LOG(WARNING) << "Cannot use " << opt->protocol.name() << " over URMA"; + return false; + } + return true; +} + +bool UrmaTransport::OptionsAvailableOverUrma(const ServerOptions* opt) { + if (opt->rtmp_service) { + LOG(WARNING) << "RTMP is not supported by URMA"; + return false; + } + if (opt->has_ssl_options()) { + LOG(WARNING) << "SSL is not supported by URMA"; + return false; + } + if (opt->nshead_service) { + LOG(WARNING) << "NSHEAD is not supported by URMA"; + return false; + } + if (opt->mongo_service_adaptor) { + LOG(WARNING) << "MONGO is not supported by URMA"; + return false; + } + return true; +} + +} // namespace brpc + +#endif // BRPC_WITH_URMA diff --git a/src/brpc/urma_transport.h b/src/brpc/urma_transport.h new file mode 100644 index 0000000000..9a496d56be --- /dev/null +++ b/src/brpc/urma_transport.h @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef BRPC_URMA_TRANSPORT_H +#define BRPC_URMA_TRANSPORT_H + +#if BRPC_WITH_URMA + +#include "brpc/channel.h" +#include "brpc/socket.h" +#include "brpc/transport.h" +#include "brpc/urma/urma_endpoint.h" + +namespace brpc { + +// UrmaTransport is the Transport subclass for URMA (openEuler Unified Remote +// Memory Access). It composes a TcpTransport for fallback (mirroring the +// RdmaTransport / UBShmTransport design) and delegates the URMA data path to +// the per-connection urma::UrmaEndpoint. Negotiation runs over the TCP fd and +// resolves _urma_state to URMA_ON or URMA_OFF. +class UrmaTransport : public Transport { + friend class TransportFactory; + friend class urma::UrmaEndpoint; + friend class urma::UrmaConnect; + friend class urma::UrmaHandshakeServerV2; + friend class urma::UrmaHandshakeClientV2; + friend class urma::UrmaHandshakeServerV3; + friend class urma::UrmaHandshakeClientV3; + +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, + const 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; + + urma::UrmaEndpoint* GetUrmaEp() { + CHECK(_urma_ep != nullptr); + return _urma_ep; + } + + static int ContextInitOrDie(bool server_or_not, const void* options); + +private: + static bool OptionsAvailableForUrma(const ChannelOptions* opt); + static bool OptionsAvailableOverUrma(const ServerOptions* opt); + + // The on/off state of URMA. UNKNOWN until the handshake resolves. + enum UrmaState { + URMA_ON, + URMA_OFF, + URMA_UNKNOWN + }; + + urma::UrmaEndpoint* _urma_ep = nullptr; + butil::atomic _urma_state{URMA_UNKNOWN}; + std::shared_ptr _tcp_transport; +}; + +} // namespace brpc + +#endif // BRPC_WITH_URMA +#endif // BRPC_URMA_TRANSPORT_H diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp new file mode 100644 index 0000000000..de3ce31228 --- /dev/null +++ b/test/brpc_urma_unittest.cpp @@ -0,0 +1,609 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#include +#include + +#if BRPC_WITH_URMA +#include "butil/atomicops.h" +#include "butil/sys_byteorder.h" +#include "urma_api.h" +#include "brpc/urma/urma_handshake.h" +#include "brpc/urma/urma_handshake.pb.h" +#include "brpc/urma/urma_helper.h" +#include "urma_types.h" + +using namespace brpc; + +namespace brpc { +namespace urma { + +DECLARE_int32(urma_client_handshake_version); +extern bool g_skip_urma_init; +extern butil::atomic g_urma_available; + +} // namespace urma +} // namespace brpc + +// --------------------------------------------------------------------------- +// v2 binary HelloMessage: serialize + deserialize round-trips. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, v2_serialize_deserialize_roundtrip) { + urma::v2_wire::HelloMessage m; + m.msg_len = urma::v2_wire::HELLO_PACKET_LEN; + m.hello_ver = urma::v2_wire::HELLO_V2_VERSION; + m.impl_ver = urma::v2_wire::IMPL_V2_VERSION; + m.buffer_size = 8192; + m.recv_buffer_cnt = 127; + m.jetty_id = 0x12345678; + for (int i = 0; i < 16; ++i) { + m.eid[i] = static_cast(i + 1); + } + m.uasid = 0xdeadbeef; + m.tp_type = 1; // URMA_CTP + for (int i = 0; i < 16; ++i) { + m.seg_eid[i] = static_cast(16 - i); + } + m.seg_uasid = 0xcafebabe; + m.seg_va = 0x1122334455667788ULL; + m.seg_len = 1ULL << 20; + m.seg_token_id = 0x42424242; + + uint8_t buf[urma::v2_wire::HELLO_BODY_LEN]; + m.Serialize(buf); + + urma::v2_wire::HelloMessage m2; + m2.Deserialize(buf); + EXPECT_EQ(m.msg_len, m2.msg_len); + EXPECT_EQ(m.hello_ver, m2.hello_ver); + EXPECT_EQ(m.impl_ver, m2.impl_ver); + EXPECT_EQ(m.buffer_size, m2.buffer_size); + EXPECT_EQ(m.recv_buffer_cnt, m2.recv_buffer_cnt); + EXPECT_EQ(m.jetty_id, m2.jetty_id); + EXPECT_EQ(0, memcmp(m.eid, m2.eid, 16)); + EXPECT_EQ(m.uasid, m2.uasid); + EXPECT_EQ(m.tp_type, m2.tp_type); + EXPECT_EQ(0, memcmp(m.seg_eid, m2.seg_eid, 16)); + EXPECT_EQ(m.seg_uasid, m2.seg_uasid); + EXPECT_EQ(m.seg_va, m2.seg_va); + EXPECT_EQ(m.seg_len, m2.seg_len); + EXPECT_EQ(m.seg_token_id, m2.seg_token_id); +} + +// --------------------------------------------------------------------------- +// v2 packet on the wire: "URMA" magic + body. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, v2_packet_magic_is_urma) { + EXPECT_EQ(4u, urma::v2_wire::MAGIC_STR_LEN); + char magic[4] = {'U', 'R', 'M', 'A'}; + EXPECT_EQ(0, memcmp(magic, "URMA", 4)); + EXPECT_EQ(4u + 82u, urma::v2_wire::HELLO_PACKET_LEN); +} + +// --------------------------------------------------------------------------- +// v3 protobuf UrmaHello: serialize + parse round-trips. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, v3_protobuf_roundtrip) { + urma::UrmaHello msg; + msg.set_buffer_size(8192); + msg.set_recv_buffer_cnt(127); + msg.set_jetty_id(0x12345678); + uint8_t eid[16]; + for (int i = 0; i < 16; ++i) { + eid[i] = static_cast(i + 1); + } + msg.set_eid(eid, 16); + msg.set_uasid(0xdeadbeef); + msg.set_tp_type(1); + uint8_t seg_eid[16]; + for (int i = 0; i < 16; ++i) { + seg_eid[i] = static_cast(16 - i); + } + msg.set_seg_eid(seg_eid, 16); + msg.set_seg_uasid(0xcafebabe); + msg.set_seg_va(0x1122334455667788ULL); + msg.set_seg_len(1ULL << 20); + msg.set_seg_token_id(0x42424242); + + std::string body; + ASSERT_TRUE(msg.SerializeToString(&body)); + urma::UrmaHello msg2; + ASSERT_TRUE(msg2.ParseFromString(body)); + EXPECT_EQ(msg.buffer_size(), msg2.buffer_size()); + EXPECT_EQ(msg.recv_buffer_cnt(), msg2.recv_buffer_cnt()); + EXPECT_EQ(msg.jetty_id(), msg2.jetty_id()); + EXPECT_EQ(16, msg2.eid().size()); + EXPECT_EQ(0, memcmp(msg.eid().data(), msg2.eid().data(), 16)); + EXPECT_EQ(msg.uasid(), msg2.uasid()); + EXPECT_EQ(msg.tp_type(), msg2.tp_type()); + EXPECT_EQ(16, msg2.seg_eid().size()); + EXPECT_EQ(msg.seg_uasid(), msg2.seg_uasid()); + EXPECT_EQ(msg.seg_va(), msg2.seg_va()); + EXPECT_EQ(msg.seg_len(), msg2.seg_len()); + EXPECT_EQ(msg.seg_token_id(), msg2.seg_token_id()); +} + +// --------------------------------------------------------------------------- +// CreateServerHandshakeByMagic dispatches on the magic bytes. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, server_handshake_factory_dispatches_on_magic) { + // We cannot fully exercise the server handshake without a real socket + + // endpoint, but we can verify the factory returns the right protocol + // version for each magic, and nullptr for an unknown magic. + uint8_t magic_v2[4] = {'U', 'R', 'M', 'A'}; + uint8_t magic_v3[4] = {'U', 'R', 'M', '3'}; + uint8_t magic_bad[4] = {'P', 'R', 'P', 'C'}; + + // v2 magic -> protocol version 2 + urma::UrmaHandshake* hs2 = + urma::CreateServerHandshakeByMagic(nullptr, magic_v2); + // Note: the factory dereferences the endpoint only inside SendLocalHello / + // ReceiveAndParseRemoteHello; passing nullptr is safe for the version query. + // (We delete immediately to avoid touching the endpoint.) + if (hs2) { + EXPECT_EQ(2, hs2->ProtocolVersion()); + delete hs2; + } + // v3 magic -> protocol version 3 + urma::UrmaHandshake* hs3 = + urma::CreateServerHandshakeByMagic(nullptr, magic_v3); + if (hs3) { + EXPECT_EQ(3, hs3->ProtocolVersion()); + delete hs3; + } + // unknown magic -> nullptr (caller falls back to TCP) + urma::UrmaHandshake* hsb = + urma::CreateServerHandshakeByMagic(nullptr, magic_bad); + EXPECT_EQ(nullptr, hsb); +} + +// --------------------------------------------------------------------------- +// CreateClientHandshake picks the version from the gflag. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, client_handshake_factory_respects_flag) { + const int saved = urma::FLAGS_urma_client_handshake_version; + + urma::FLAGS_urma_client_handshake_version = 2; + urma::UrmaHandshake* hs2 = urma::CreateClientHandshake(nullptr); + if (hs2) { + EXPECT_EQ(2, hs2->ProtocolVersion()); + delete hs2; + } + + urma::FLAGS_urma_client_handshake_version = 3; + urma::UrmaHandshake* hs3 = urma::CreateClientHandshake(nullptr); + if (hs3) { + EXPECT_EQ(3, hs3->ProtocolVersion()); + delete hs3; + } + + urma::FLAGS_urma_client_handshake_version = saved; +} + +// --------------------------------------------------------------------------- +// 4-byte ACK: HELLO_ACK_URMA_OK bit. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, ack_bit_is_rdma_ok) { + // The ACK is a 4-byte big-endian flags word; bit 0 means "I want URMA". + // Verify the round-trip: host -> net -> host preserves the bit. + uint32_t flags = 0x1; // HELLO_ACK_URMA_OK + uint32_t flags_be = butil::HostToNet32(flags); + uint32_t flags_back = butil::NetToHost32(flags_be); + EXPECT_EQ(flags, flags_back); + EXPECT_NE(0u, flags_back & 0x1); +} + +// --------------------------------------------------------------------------- +// ParsedHello field layout: covers the flattened segment (seg_* fields). +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, parsed_hello_segment_fields) { + urma::ParsedHello p; + std::memset(&p, 0, sizeof(p)); + p.buffer_size = 8192; + p.recv_buffer_cnt = 127; + p.jetty_id = 42; + p.tp_type = 1; + p.seg_va = 0x1000; + p.seg_len = 0x100000; + p.seg_token_id = 7; + EXPECT_EQ(8192u, p.buffer_size); + EXPECT_EQ(127u, p.recv_buffer_cnt); + EXPECT_EQ(42u, p.jetty_id); + EXPECT_EQ(1u, p.tp_type); + EXPECT_EQ(0x1000u, p.seg_va); + EXPECT_EQ(0x100000u, p.seg_len); + EXPECT_EQ(7u, p.seg_token_id); +} + +TEST(UrmaHandshakeTest, rejects_invalid_resource_and_window_values) { + urma::ParsedHello hello; + hello.buffer_size = 8192; + hello.recv_buffer_cnt = 127; + hello.jetty_id = 1; + hello.tp_type = URMA_CTP; + hello.seg_va = 0x1000; + hello.seg_len = 8192; + EXPECT_TRUE(urma::ValidHello(hello)); + + hello.recv_buffer_cnt = 2; + EXPECT_FALSE(urma::ValidHello(hello)); + hello.recv_buffer_cnt = 127; + hello.seg_len = 0; + EXPECT_FALSE(urma::ValidHello(hello)); + hello.seg_len = 8192; + hello.jetty_id = 0; + EXPECT_FALSE(urma::ValidHello(hello)); +} + +TEST(UrmaHelperTest, selects_priority_matching_transport_path_type) { + urma_device_attr_t attr{}; + attr.dev_cap.priority_info[3].tp_type.bs.rtp = 1; + attr.dev_cap.priority_info[6].tp_type.bs.ctp = 1; + + EXPECT_EQ(3, urma::FindUrmaPriorityForTpType(attr, URMA_RTP)); + EXPECT_EQ(6, urma::FindUrmaPriorityForTpType(attr, URMA_CTP)); + EXPECT_EQ(-1, urma::FindUrmaPriorityForTpType(attr, URMA_UTP)); +} + +// --------------------------------------------------------------------------- +// SupportedByUrma: only baidu_std. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, supported_by_urma_protocol_allowlist) { + EXPECT_TRUE(urma::SupportedByUrma("baidu_std")); + EXPECT_FALSE(urma::SupportedByUrma("http")); + EXPECT_FALSE(urma::SupportedByUrma("hulu_pbrpc")); + EXPECT_FALSE(urma::SupportedByUrma("nshead")); +} + +// --------------------------------------------------------------------------- +// URMA mock smoke test: drives urma_init / device enumeration / context / +// jetty / post / poll. These tests rely on mock semantics and are skipped +// when the test binary links a real liburma provider. +// --------------------------------------------------------------------------- +class UrmaMockTest : public ::testing::Test { +protected: + void SetUp() override { + urma::g_skip_urma_init = false; + + urma_init_attr_t init_attr{}; + const urma_status_t status = urma_init(&init_attr); + ASSERT_TRUE(status == URMA_SUCCESS || status == URMA_EEXIST); + _owns_urma_init = status == URMA_SUCCESS; + + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + ASSERT_GT(num_devices, 0); + const bool using_mock = + strcmp(devices[0]->name, "mock_urma_device") == 0; + urma_free_device_list(devices); + if (!using_mock) { + if (_owns_urma_init) { + EXPECT_EQ(URMA_SUCCESS, urma_uninit()); + _owns_urma_init = false; + } + GTEST_SKIP() << "UrmaMockTest requires the URMA link-time mock"; + } + } + + void TearDown() override { + if (_owns_urma_init) { + EXPECT_EQ(URMA_SUCCESS, urma_uninit()); + } + urma::g_skip_urma_init = true; + urma::g_urma_available.store(true, butil::memory_order_relaxed); + } + +private: + bool _owns_urma_init{false}; +}; + +TEST_F(UrmaMockTest, init_and_enumerate_device) { + urma_init_attr_t init_attr{}; + EXPECT_EQ(URMA_EEXIST, urma_init(&init_attr)); + + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + ASSERT_GE(num_devices, 1); + EXPECT_STREQ("mock_urma_device", devices[0]->name); + urma_free_device_list(devices); +} + +TEST_F(UrmaMockTest, create_context_and_query_device) { + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + ASSERT_GE(num_devices, 1); + + uint32_t eid_cnt = 0; + urma_eid_info_t* eids = urma_get_eid_list(devices[0], &eid_cnt); + ASSERT_NE(nullptr, eids); + ASSERT_GE(eid_cnt, 1u); + urma_free_eid_list(eids); + + urma_context_t* ctx = urma_create_context(devices[0], 0); + ASSERT_NE(nullptr, ctx); + + urma_device_attr_t attr{}; + ASSERT_EQ(URMA_SUCCESS, urma_query_device(devices[0], &attr)); + EXPECT_GE(attr.dev_cap.max_jfc, 1u); + EXPECT_GE(attr.dev_cap.max_jetty, 1u); + + EXPECT_EQ(URMA_SUCCESS, urma_delete_context(ctx)); + urma_free_device_list(devices); +} + +TEST_F(UrmaMockTest, rejects_send_without_target_jetty) { + // A SEND without a target jetty is invalid. Keeping the mock strict here + // prevents tests from relying on input that a real provider cannot post. + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + ASSERT_GE(num_devices, 1); + uint32_t eid_cnt = 0; + urma_eid_info_t* eids = urma_get_eid_list(devices[0], &eid_cnt); + urma_free_eid_list(eids); + urma_context_t* ctx = urma_create_context(devices[0], 0); + ASSERT_NE(nullptr, ctx); + + urma_jfce_t* jfce = urma_create_jfce(ctx); + ASSERT_NE(nullptr, jfce); + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = 16; + jfc_cfg.jfce = jfce; + urma_jfc_t* jfc = urma_create_jfc(ctx, &jfc_cfg); + ASSERT_NE(nullptr, jfc); + + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = 16; + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + jfr_cfg.jfc = jfc; + urma_jfr_t* jfr = urma_create_jfr(ctx, &jfr_cfg); + ASSERT_NE(nullptr, jfr); + + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = 16; + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.priority = URMA_MAX_PRIORITY; + jetty_cfg.jfs_cfg.max_sge = 1; + jetty_cfg.jfs_cfg.rnr_retry = URMA_TYPICAL_RNR_RETRY; + jetty_cfg.jfs_cfg.err_timeout = URMA_TYPICAL_ERR_TIMEOUT; + jetty_cfg.jfs_cfg.jfc = jfc; + jetty_cfg.shared.jfr = jfr; + jetty_cfg.shared.jfc = jfc; + urma_jetty_t* jetty = urma_create_jetty(ctx, &jetty_cfg); + ASSERT_NE(nullptr, jetty); + + urma_jfs_wr_t wr{}; + memset(&wr, 0, sizeof(wr)); + wr.opcode = URMA_OPC_SEND; + wr.flag.bs.complete_enable = 1; + wr.user_ctx = 0xABCD; + wr.next = nullptr; + urma_jfs_wr_t* bad = nullptr; + EXPECT_EQ(URMA_EINVAL, urma_post_jetty_send_wr(jetty, &wr, &bad)); + EXPECT_EQ(&wr, bad); + + urma_cr_t crs[4]; + EXPECT_EQ(0, urma_poll_jfc(jfc, 4, crs)); + + urma_delete_jetty(jetty); + urma_delete_jfr(jfr); + urma_delete_jfc(jfc); + urma_delete_jfce(jfce); + urma_delete_context(ctx); + urma_free_device_list(devices); +} + +TEST_F(UrmaMockTest, + paired_send_is_bidirectional_and_separates_immediate_credit) { + int num_devices = 0; + urma_device_t** devices = urma_get_device_list(&num_devices); + ASSERT_NE(nullptr, devices); + ASSERT_GT(num_devices, 0); + urma_context_t* ctx = urma_create_context(devices[0], 0); + ASSERT_NE(nullptr, ctx); + + urma_jfce_t* sender_jfce = urma_create_jfce(ctx); + urma_jfce_t* receiver_jfce = urma_create_jfce(ctx); + ASSERT_NE(nullptr, sender_jfce); + ASSERT_NE(nullptr, receiver_jfce); + urma_jfc_cfg_t sender_jfc_cfg{}; + sender_jfc_cfg.depth = 8; + sender_jfc_cfg.jfce = sender_jfce; + urma_jfc_t* sender_jfc = urma_create_jfc(ctx, &sender_jfc_cfg); + ASSERT_NE(nullptr, sender_jfc); + urma_jfc_cfg_t receiver_jfc_cfg{}; + receiver_jfc_cfg.depth = 8; + receiver_jfc_cfg.jfce = receiver_jfce; + urma_jfc_t* receiver_jfc = urma_create_jfc(ctx, &receiver_jfc_cfg); + ASSERT_NE(nullptr, receiver_jfc); + + urma_jfr_cfg_t sender_jfr_cfg{}; + sender_jfr_cfg.depth = 4; + sender_jfr_cfg.trans_mode = URMA_TM_RM; + sender_jfr_cfg.max_sge = 1; + sender_jfr_cfg.jfc = sender_jfc; + urma_jfr_t* sender_jfr = urma_create_jfr(ctx, &sender_jfr_cfg); + ASSERT_NE(nullptr, sender_jfr); + urma_jfr_cfg_t receiver_jfr_cfg = sender_jfr_cfg; + receiver_jfr_cfg.jfc = receiver_jfc; + urma_jfr_t* receiver_jfr = urma_create_jfr(ctx, &receiver_jfr_cfg); + ASSERT_NE(nullptr, receiver_jfr); + + auto create_jetty = [&](urma_jfc_t* jfc, urma_jfr_t* jfr) { + urma_jetty_cfg_t cfg{}; + cfg.flag.bs.share_jfr = 1; + cfg.jfs_cfg.depth = 4; + cfg.jfs_cfg.trans_mode = URMA_TM_RM; + cfg.jfs_cfg.max_sge = 1; + cfg.jfs_cfg.jfc = jfc; + cfg.shared.jfr = jfr; + cfg.shared.jfc = jfc; + return urma_create_jetty(ctx, &cfg); + }; + urma_jetty_t* sender = create_jetty(sender_jfc, sender_jfr); + urma_jetty_t* receiver = create_jetty(receiver_jfc, receiver_jfr); + ASSERT_NE(nullptr, sender); + ASSERT_NE(nullptr, receiver); + + urma_rjetty_t remote{}; + remote.jetty_id = receiver->jetty_id; + remote.trans_mode = URMA_TM_RM; + remote.type = URMA_JETTY; + remote.tp_type = URMA_CTP; + urma_token_t token{}; + urma_target_jetty_t* target = + urma_import_jetty(ctx, &remote, &token); + ASSERT_NE(nullptr, target); + + char recv_buf[64]{}; + urma_sge_t recv_sge{ + reinterpret_cast(recv_buf), sizeof(recv_buf), nullptr, + nullptr}; + urma_sg_t recv_sg{&recv_sge, 1}; + urma_jfr_wr_t recv_wr{recv_sg, 99, nullptr}; + char credit_recv_buf[1]{}; + urma_sge_t credit_recv_sge{ + reinterpret_cast(credit_recv_buf), + sizeof(credit_recv_buf), nullptr, nullptr}; + urma_sg_t credit_recv_sg{&credit_recv_sge, 1}; + urma_jfr_wr_t credit_recv_wr{credit_recv_sg, 100, nullptr}; + recv_wr.next = &credit_recv_wr; + urma_jfr_wr_t* bad_recv = nullptr; + ASSERT_EQ(URMA_SUCCESS, + urma_post_jfr_wr(receiver_jfr, &recv_wr, &bad_recv)); + + const char payload[] = "urma-payload"; + urma_sge_t send_sge{ + reinterpret_cast(payload), sizeof(payload), nullptr, + nullptr}; + urma_sg_t send_sg{&send_sge, 1}; + urma_jfs_wr_t send_wr{}; + send_wr.opcode = URMA_OPC_SEND; + send_wr.flag.bs.complete_enable = 1; + send_wr.tjetty = target; + send_wr.user_ctx = 7; + send_wr.send.src = send_sg; + urma_jfs_wr_t credit_wr{}; + credit_wr.opcode = URMA_OPC_SEND_IMM; + credit_wr.flag.bs.complete_enable = 1; + credit_wr.tjetty = target; + credit_wr.user_ctx = 8; + credit_wr.send.imm_data = 13; + send_wr.next = &credit_wr; + urma_jfs_wr_t* bad_send = nullptr; + ASSERT_EQ(URMA_SUCCESS, + urma_post_jetty_send_wr(sender, &send_wr, &bad_send)); + + urma_cr_t sender_cr[2]{}; + ASSERT_EQ(2, urma_poll_jfc(sender_jfc, 2, sender_cr)); + EXPECT_EQ(0, sender_cr[0].flag.bs.s_r); + EXPECT_EQ(7u, sender_cr[0].user_ctx); + EXPECT_EQ(0, sender_cr[1].flag.bs.s_r); + EXPECT_EQ(8u, sender_cr[1].user_ctx); + + urma_cr_t receiver_cr[2]{}; + ASSERT_EQ(2, urma_poll_jfc(receiver_jfc, 2, receiver_cr)); + EXPECT_EQ(1, receiver_cr[0].flag.bs.s_r); + EXPECT_EQ(URMA_CR_OPC_SEND, receiver_cr[0].opcode); + EXPECT_EQ(0u, receiver_cr[0].imm_data); + EXPECT_EQ(sizeof(payload), receiver_cr[0].completion_len); + EXPECT_EQ(0, memcmp(payload, recv_buf, sizeof(payload))); + EXPECT_EQ(1, receiver_cr[1].flag.bs.s_r); + EXPECT_EQ(URMA_CR_OPC_SEND_WITH_IMM, receiver_cr[1].opcode); + EXPECT_EQ(13u, receiver_cr[1].imm_data); + EXPECT_EQ(0u, receiver_cr[1].completion_len); + + // Exercise the response direction as well. Production posts all receive + // WRs through the shared JFR. + urma_rjetty_t sender_remote = remote; + sender_remote.jetty_id = sender->jetty_id; + urma_target_jetty_t* sender_target = + urma_import_jetty(ctx, &sender_remote, &token); + ASSERT_NE(nullptr, sender_target); + char response_buf[64]{}; + urma_sge_t response_recv_sge{ + reinterpret_cast(response_buf), sizeof(response_buf), + nullptr, nullptr}; + urma_sg_t response_recv_sg{&response_recv_sge, 1}; + urma_jfr_wr_t response_recv_wr{response_recv_sg, 101, nullptr}; + ASSERT_EQ(URMA_SUCCESS, + urma_post_jfr_wr(sender_jfr, &response_recv_wr, &bad_recv)); + + const char response[] = "urma-response"; + urma_sge_t response_send_sge{ + reinterpret_cast(response), sizeof(response), nullptr, + nullptr}; + urma_sg_t response_send_sg{&response_send_sge, 1}; + urma_jfs_wr_t response_send_wr{}; + response_send_wr.opcode = URMA_OPC_SEND; + response_send_wr.flag.bs.complete_enable = 1; + response_send_wr.tjetty = sender_target; + response_send_wr.user_ctx = 9; + response_send_wr.send.src = response_send_sg; + ASSERT_EQ(URMA_SUCCESS, + urma_post_jetty_send_wr(receiver, &response_send_wr, &bad_send)); + + urma_cr_t response_send_cr{}; + ASSERT_EQ(1, urma_poll_jfc(receiver_jfc, 1, &response_send_cr)); + EXPECT_EQ(0, response_send_cr.flag.bs.s_r); + EXPECT_EQ(9u, response_send_cr.user_ctx); + + urma_cr_t response_recv_cr{}; + ASSERT_EQ(1, urma_poll_jfc(sender_jfc, 1, &response_recv_cr)); + EXPECT_EQ(1, response_recv_cr.flag.bs.s_r); + EXPECT_EQ(URMA_CR_OPC_SEND, response_recv_cr.opcode); + EXPECT_EQ(sizeof(response), response_recv_cr.completion_len); + EXPECT_EQ(0, memcmp(response, response_buf, sizeof(response))); + + urma_unimport_jetty(sender_target); + urma_unimport_jetty(target); + urma_delete_jetty(receiver); + urma_delete_jetty(sender); + urma_delete_jfr(receiver_jfr); + urma_delete_jfr(sender_jfr); + urma_delete_jfc(receiver_jfc); + urma_delete_jfc(sender_jfc); + urma_delete_jfce(receiver_jfce); + urma_delete_jfce(sender_jfce); + urma_delete_context(ctx); + urma_free_device_list(devices); +} + +#else // BRPC_WITH_URMA + +// When URMA is not compiled in, the test file is a no-op so the build stays +// clean. The brpc_urma_unittest target still links (against brpc-shared which +// provides the empty stubs). + +#endif // BRPC_WITH_URMA + +int main(int argc, char** argv) { + testing::InitGoogleTest(&argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, true); +#if BRPC_WITH_URMA + urma::g_skip_urma_init = true; + urma::g_urma_available.store(true, butil::memory_order_relaxed); +#endif + return RUN_ALL_TESTS(); +}