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 fd315ffa28..f6d495dbe4 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -52,9 +52,12 @@ DEFINES = [ "//bazel/config:brpc_with_rdma": ["BRPC_WITH_RDMA=1"], "//conditions:default": [], }) + select({ - "//bazel/config:brpc_with_ubring": ["BRPC_WITH_UBRING=1"], + "//bazel/config:brpc_with_urma": ["BRPC_WITH_URMA=1"], "//conditions:default": [], }) + select({ + "//bazel/config:brpc_with_ubring": ["BRPC_WITH_UBRING=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"], }) + select({ @@ -522,6 +525,7 @@ filegroup( srcs = glob([ "src/brpc/*.proto", "src/brpc/policy/*.proto", + "src/brpc/urma/*.proto", ]), visibility = ["//visibility:public"], ) @@ -534,9 +538,20 @@ brpc_proto_library( visibility = ["//visibility:public"], ) -cc_library( - name = "brpc", - srcs = glob([ +URMA_SRC_PATTERNS = [ + "src/brpc/urma/*.cpp", + "src/brpc/urma/**/*.cpp", +] + +# Keep URMA implementation sources out of BRPC_BASE_SRCS. They are added only +# by the select() below so the mock can never leak into a real-liburma link. +URMA_SRCS = glob( + URMA_SRC_PATTERNS, + exclude = ["src/brpc/urma/mock_urma.cpp"], +) + +BRPC_BASE_SRCS = glob( + [ "src/brpc/*.cpp", "src/brpc/**/*.cpp", ], @@ -546,12 +561,23 @@ cc_library( "src/brpc/policy/thrift_protocol.cpp", "src/brpc/event_dispatcher_epoll.cpp", "src/brpc/event_dispatcher_kqueue.cpp", - ]) + select({ + ] + URMA_SRC_PATTERNS, +) + +cc_library( + name = "brpc", + srcs = BRPC_BASE_SRCS + select({ "//bazel/config:brpc_with_thrift": glob([ "src/brpc/thrift*.cpp", "src/brpc/**/thrift*.cpp", ]), "//conditions:default": [], + }) + select({ + "//bazel/config:brpc_with_urma_use_real": URMA_SRCS, + "//bazel/config:brpc_with_urma": URMA_SRCS + [ + "src/brpc/urma/mock_urma.cpp", + ], + "//conditions:default": [], }), hdrs = glob([ "src/brpc/*.h", @@ -560,10 +586,11 @@ cc_library( "src/brpc/event_dispatcher_kqueue.cpp", ]), copts = COPTS, - includes = [ - "src/", - ], - linkopts = LINKOPTS, + includes = ["src/"], + linkopts = LINKOPTS + select({ + "//bazel/config:brpc_with_urma_use_real": ["-lurma"], + "//conditions:default": [], + }), visibility = ["//visibility:public"], deps = [ ":brpc_internal_cc_proto", @@ -574,6 +601,12 @@ cc_library( ":mcpack2pb", "@com_github_google_leveldb//:leveldb", ] + select({ + "//bazel/config:brpc_with_urma_download_headers_disabled": [], + "//bazel/config:brpc_with_urma": [ + "@umdk//:urma_headers", + ], + "//conditions:default": [], + }) + select({ "//bazel/config:brpc_with_thrift": [ "@org_apache_thrift//:thrift", ], diff --git a/CLAUDE.md b/CLAUDE.md index fccb197c7b..6104e5ff8d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,7 +4,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## Project Overview -Apache bRPC is an industrial-grade C++ RPC framework supporting multiple protocols (baidu_std, HTTP/H2, gRPC, thrift, redis, memcached, RTMP, RDMA) on the same port. Used in high-performance systems: search, storage, ML, ads, recommendations. Current version: 1.17.0. +Apache bRPC is an industrial-grade C++ RPC framework supporting multiple protocols (baidu_std, HTTP/H2, gRPC, thrift, redis, memcached, RTMP, RDMA) on the same port. Used in high-performance systems: search, storage, ML, ads, recommendations. Current version: 1.18.0. ## Build Commands diff --git a/CMakeLists.txt b/CMakeLists.txt index 622f587963..7b023ec23e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,12 +21,18 @@ project(brpc C CXX) option(WITH_GLOG "With glog" OFF) option(WITH_MESALINK "With MesaLink" OFF) option(WITH_BORINGSSL "With BoringSSL" OFF) -option(DEBUG "Print debug logs" OFF) option(WITH_DEBUG_SYMBOLS "With debug symbols" ON) 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) +set(URMA_USE_MOCK "AUTO" CACHE STRING + "URMA mock mode: AUTO, ON, or OFF (OFF requires liburma)") +set_property(CACHE URMA_USE_MOCK PROPERTY STRINGS AUTO ON 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) @@ -42,7 +48,7 @@ if(POLICY CMP0042) cmake_policy(SET CMP0042 NEW) endif() -set(BRPC_VERSION 1.17.0) +set(BRPC_VERSION 1.18.0) SET(CPACK_GENERATOR "DEB") SET(CPACK_DEBIAN_PACKAGE_MAINTAINER "brpc authors") @@ -120,6 +126,19 @@ if(WITH_RDMA) set(WITH_RDMA_VAL "1") endif() +set(WITH_URMA_VAL "0") +if(WITH_URMA) + set(WITH_URMA_VAL "1") +endif() + +string(TOUPPER "${URMA_USE_MOCK}" URMA_USE_MOCK_MODE) +if(NOT URMA_USE_MOCK_MODE STREQUAL "AUTO" AND + NOT URMA_USE_MOCK_MODE STREQUAL "ON" AND + NOT URMA_USE_MOCK_MODE STREQUAL "OFF") + message(FATAL_ERROR + "URMA_USE_MOCK must be AUTO, ON, or OFF, got: ${URMA_USE_MOCK}") +endif() + set(WITH_UBRING_VAL "0") if(WITH_UBRING) set(WITH_UBRING_VAL "1") @@ -159,6 +178,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} @@ -174,9 +194,6 @@ list(APPEND BRPC_COMMON_DEFINITIONS __STRICT_ANSI__ ) -if(NOT DEBUG) - list(APPEND BRPC_COMMON_DEFINITIONS NDEBUG) -endif() if(WITH_ASAN) list(APPEND BRPC_COMMON_COMPILE_OPTIONS -fsanitize=address) list(APPEND BRPC_COMMON_LINK_OPTIONS -fsanitize=address) @@ -318,6 +335,69 @@ 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 + # commit, 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 564ee727a55523d4351a8fb3c94292b388ebb924 + # brpc only consumes UMDK headers; avoid cloning unrelated UMDK + # submodules such as CAM dependencies during configuration. + GIT_SUBMODULES "" + 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_USE_MOCK_MODE STREQUAL "ON") + message(STATUS "URMA_USE_MOCK=ON; forcing the link-time URMA mock") + set(URMA_USE_MOCK_ENABLED 1) + elseif(URMA_USE_MOCK_MODE STREQUAL "OFF") + if(NOT URMA_LIB) + message(FATAL_ERROR + "URMA_USE_MOCK=OFF requires liburma; install it or use AUTO/ON.") + endif() + message(STATUS "URMA_USE_MOCK=OFF; using liburma") + set(URMA_USE_MOCK_ENABLED 0) + elseif(URMA_LIB) + message(STATUS "Found URMA library: ${URMA_LIB}") + set(URMA_USE_MOCK_ENABLED 0) + else() + message(STATUS + "liburma not found; building with the URMA link-time mock") + set(URMA_USE_MOCK_ENABLED 1) + endif() +endif() + find_library(PROTOC_LIB NAMES protoc) if(NOT PROTOC_LIB) message(FATAL_ERROR "Fail to find protoc lib") @@ -342,6 +422,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} @@ -370,6 +456,13 @@ if(WITH_RDMA) list(APPEND DYNAMIC_LIB ${RDMA_LIB}) endif() +if(WITH_URMA) + message(STATUS "brpc compile with URMA (mock=${URMA_USE_MOCK_ENABLED})") + if(NOT URMA_USE_MOCK_ENABLED) + list(APPEND DYNAMIC_LIB ${URMA_LIB}) + endif() +endif() + if(WITH_UBRING) message(STATUS "brpc compile with ubring") list(APPEND DYNAMIC_LIB ${UB_LIB}) @@ -579,6 +672,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_ENABLED) + 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() @@ -619,7 +720,8 @@ set(PROTO_FILES idl_options.proto brpc/trackme.proto brpc/streaming_rpc_meta.proto brpc/proto_base.proto - brpc/rdma_handshake.proto) + brpc/rdma_handshake.proto + brpc/urma/urma_handshake.proto) file(MAKE_DIRECTORY ${PROJECT_BINARY_DIR}/output/include/brpc) set(PROTOC_FLAGS ${PROTOC_FLAGS} -I${PROTOBUF_INCLUDE_DIR}) compile_proto(PROTO_HDRS PROTO_SRCS ${PROJECT_BINARY_DIR} diff --git a/MODULE.bazel b/MODULE.bazel index 1e71bfcb9d..97862d36f1 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -17,7 +17,7 @@ module( name = 'brpc', - version = '1.17.0', + version = '1.18.0', compatibility_level = 1, ) @@ -68,3 +68,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', + commit = '564ee727a55523d4351a8fb3c94292b388ebb924', # v26.06.0_CAM +) diff --git a/Makefile b/Makefile index 90be5516c3..dd33e0dc6a 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/handshake 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/RELEASE_VERSION b/RELEASE_VERSION index 092afa15df..84cc529467 100644 --- a/RELEASE_VERSION +++ b/RELEASE_VERSION @@ -1 +1 @@ -1.17.0 +1.18.0 diff --git a/WORKSPACE b/WORKSPACE index 78a6c2836a..22fc411b32 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", + commit = "564ee727a55523d4351a8fb3c94292b388ebb924", # 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 e9d07eb73a..be35f55bcf 100644 --- a/bazel/config/BUILD.bazel +++ b/bazel/config/BUILD.bazel @@ -110,6 +110,30 @@ 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_urma_use_real", + define_values = { + "BRPC_WITH_URMA": "true", + "BRPC_URMA_USE_MOCK": "false", + }, + visibility = ["//visibility:public"], +) + +config_setting( + name = "brpc_with_urma_download_headers_disabled", + define_values = { + "BRPC_WITH_URMA": "true", + "BRPC_DOWNLOAD_URMA_HEADERS": "false", + }, + visibility = ["//visibility:public"], +) + config_setting( name = "brpc_with_boringssl", define_values = {"BRPC_WITH_BORINGSSL": "true"}, @@ -161,4 +185,4 @@ config_setting( name = "brpc_with_ubring", define_values = {"BRPC_WITH_UBRING": "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/community/newcommitter_cn.md b/community/newcommitter_cn.md index 34bbf3ec4e..668b0bc68b 100644 --- a/community/newcommitter_cn.md +++ b/community/newcommitter_cn.md @@ -1,4 +1,4 @@ -# 这里记录发展Committer 和 PPMC成员的流程和参考网站信息 +# 这里记录发展Committer 和 PMC成员的流程和参考网站信息 ## 1. 如何发展committer @@ -10,7 +10,7 @@ ### 成为committer的路程 1. 提名者在private@brpc中发起讨论和投票，投票通过即OK （最少3+1， +1 > -1)，[邮件模版](https://community.apache.org/newcommitter.html#committer-vote-template) 2. 提名者发送close vote邮件给private@brpc, 标题可以为subject [RESULT][VOTE]，[邮件模版](https://community.apache.org/newcommitter.html#close-vote) -3. 提名者给被提名者发invite letter([邮件模板](https://community.apache.org/newcommitter.html#committer-invite-template))，并得到回复后再提示他提交ICLA([邮件模板](https://community.apache.org/newcommitter.html#committer-accept-template)) +3. 提名者给被提名者发invite letter并提示他提交ICLA([邮件模板](https://community.apache.org/templates/committer-invite.txt)) 4. 被提名者填写[CLA](https://www.apache.org/licenses/contributor-agreements.html), 个人贡献者需要下载[ICLA](https://www.apache.org/licenses/icla.pdf)填写个人信息并签名，发送电子版给 secretary@apache.org。（注意：1. ICLA需要填写信息完全，包括邮寄地址和签名，否则会被ASF的秘书打回, 2. ICLA中含有个人信息，所以应该单发给secretary@apache.org，不应该同时抄送PMC或其他邮件组 3. ICLA中的preffered Apache id和notify project虽然是optional，但最好填上，这样秘书会帮新commiter申请好apache id，否则就需要PMC chair或者ASF member去申请）个人信息填写项（除签名外）可以使用 PDF 阅读器或浏览器填写，填写后保存进行签名。签名方式支持： - 打印 该pdf 文件，手工填写表单（姓名、邮箱、邮寄地址），然后手写签名，最后扫描为电子版； - 使用支持手写的设备进行电子签名； @@ -23,13 +23,13 @@ ### 如何赋予committer在github上的权限 1. 加为committer -https://whimsy.apache.org/roster/ppmc/brpc + 2. 让他设置github id -https://id.apache.org/ + 3. 让他访问该网址，获得github的权限 -https://gitbox.apache.org/setup/ + ### Apache 官网new committer相关的文档 @@ -51,18 +51,16 @@ Doing these things will make everyone's job easier. -## 2. 如何把committer变成为PPMC +## 2. 如何把committer变成为PMC ### 流程参考：Apache官网文档 -* https://incubator.apache.org/guides/ppmc.html#voting_in_a_new_ppmc_member -* https://community.apache.org/newcommitter.html -* https://incubator.apache.org/guides/ppmc.html#podling_project_management_committee_ppmc +* https://community.apache.org/pmc/adding-pmc-members.html ### 实际流程 1. 在private@brpc中发起讨论，如果没有反对，则继续 2. 在private@brpc中发起投票 3. 在private@brpc中发邮件，结束投票，并通知private@incubator.apache.org -4. 在private@brpc中和dev中announce new PPMC -5. 设定他的权限，通过访问https://whimsy.apache.org/roster/ppmc/brpc -6. 帮他订阅private邮件组，参见https://whimsy.apache.org/committers/moderationhelper.cgi +4. 在private@brpc中和dev中announce new PMC +5. 设定他的权限，通过访问 https://whimsy.apache.org/roster/committee/brpc +6. 帮他订阅private邮件组，参见 https://whimsy.apache.org/committers/moderationhelper.cgi diff --git a/community/newcommitter_en.md b/community/newcommitter_en.md index 2e99334d4d..6fa2966ba3 100644 --- a/community/newcommitter_en.md +++ b/community/newcommitter_en.md @@ -1,4 +1,4 @@ -# Here is the doc for the process and reference website information of developing committee and PPMC members. +# Here is the doc for the process and reference website information of developing committee and PMC members. ## 1. How to develop committers @@ -13,7 +13,7 @@ 1. The nominator sends an email to private@brpc to initiate discussion and begins to voting. If the voting is passed, it is OK (at least 3 +1, +1 >- 1). (See the [Vote email template](https://community.apache.org/newcommitter.html#committer-vote-template)) 2. The nominator sends a close vote email to private@brpc, the title can be subject [RESULT] [VOTE]. (See the [close email template](https://community.apache.org/newcommitter.html#close-vote)) -3. The nominator sends an invitation letter to the nominee ([email template](https://community.apache.org/newcommitter.html#committer-invite-template)) and prompts him to submit ICLA after receiving a reply ([email template](https://community.apache.org/newcommitter.html#committer-accept-template)) +3. The nominator sends an invitation letter to the nominee and prompts him to submit ICLA([email template](https://community.apache.org/templates/committer-invite.txt)). 4. The nominee fills in [ICLA](https://www.apache.org/licenses/contributor-agreements.html), individual contributors need to download [ICLA](https://www.apache.org/licenses/icla.pdf) Fill in personal information and sign, and send the electronic version to secretary@apache.org 。 (Note: 1. ICLA needs to fill in complete information, including mailing address and signature, otherwise it will be returned by ASF's secretary. 2. ICLA contains personal information, it should be send to secretary@apache.org only, not to the PMC or other mail list) The personal information entry (except for the signature) can be filled in with a PDF reader or browser, and then saved for signature. Signature method support: * Print pdf documents and scan them into electronic version after handwritten signature; * Electronic signature using devices that support handwriting; @@ -23,7 +23,7 @@ ### How to grant a committer permission on github -1. Add as committer ([https://whimsy.apache.org/roster/ppmc/brpc](https://whimsy.apache.org/roster/ppmc/brpc)) +1. Add as committer () 2. Let him set github id ([https://id.apache.org/](https://id.apache.org/)) 3. Let him visit the website and get github permission([https://gitbox.apache.org/setup/](https://id.apache.org/)) @@ -42,19 +42,17 @@ Please do these things: Doing these things will make everyone's job easier. -## 2. How to upgrade a committer into a PPMC +## 2. How to upgrade a committer into a PMC ### Process reference: Apache official website document -* https://incubator.apache.org/guides/ppmc.html#voting_in_a_new_ppmc_member -* https://community.apache.org/newcommitter.html -* https://incubator.apache.org/guides/ppmc.html#podling_project_management_committee_ppmc +* https://community.apache.org/pmc/adding-pmc-members.html ### Actual process 1. Initiate discussion in the private@brpc . If there is no objection, continue 2. Open a Vote in private@brpc 3. Send an email to close the voting and notify private@incubator.apache.org -4. Announce new PPMC On private@brpc -5. Set his authority by visiting https://whimsy.apache.org/roster/ppmc/brpc +4. Announce new PMC On private@brpc +5. Set his authority by visiting https://whimsy.apache.org/roster/committee/brpc 6. Help him subscribe to a private email group. See https://whimsy.apache.org/committers/moderationhelper.cgi diff --git a/community/release_cn.md b/community/release_cn.md index ee24c1bc4a..4774eb815c 100644 --- a/community/release_cn.md +++ b/community/release_cn.md @@ -1,5 +1,7 @@ # brpc 发布 apache release 版本流程 step by step +> **推荐先使用发版 Skill**：[`community/skills/brpc-release/`](skills/brpc-release) 将本流程编排为可恢复的阶段任务，提供版本更新、Release Notes、打包和验包脚本，并在 GPG 签名、推 tag、SVN 提交、发邮件及 GitHub Release 等不可逆操作前停止并提示 Release Manager（RM）手工执行。安装方法、依赖检查、RM/校验者使用指南见 [`community/skills/README.md`](skills/README.md)。本文件仍是权威发布流程；Skill 与本文不一致时，以本文为准并同步修正 Skill。 + ## 准备工作 ### 1. 确认 Release Notes @@ -199,6 +201,14 @@ set(BRPC_VERSION 1.0.0) Version: 1.0.0 ``` +#### 更新 `CLAUDE.md` 文件 + +编辑项目根目录下 `CLAUDE.md` 文件，将项目概述中的 `Current version` 更新为本次发布版本： + +``` +Current version: 1.0.0. +``` + #### 更新 `MODULE.bazel` 文件 编辑项目根目录下 `MODULE.bazel` 文件，更新版本号，并提交至代码仓库，本文以 `1.0.0` 版本为例，修改 `version` 为： @@ -262,13 +272,13 @@ sha512sum --check apache-brpc-$BRPCVERSION-src.tar.gz.sha512 如无本地工作目录，则先创建本地工作目录。将 Apache SVN 仓库克隆下来，username 需要使用自己的 Apache LDAP 用户名。 ```bash -mkdir -p ~/brpc_svn/dev/ +mkdir -p ~/brpc_release/svn/dev/ -cd ~/brpc_svn/dev/ +cd ~/brpc_release/svn/dev/ svn --username=$BRPCUSERNAME co https://dist.apache.org/repos/dist/dev/brpc/ -cd ~/brpc_svn/dev/brpc +cd ~/brpc_release/svn/dev/brpc ``` ### 2. 添加 GPG 公钥 @@ -291,15 +301,15 @@ cd ~/brpc_svn/dev/brpc ### 3. 将待发布的代码包添加至 SVN 目录 ```bash -mkdir -p ~/brpc_svn/dev/brpc/$BRPCVERSION +mkdir -p ~/brpc_release/svn/dev/brpc/$BRPCVERSION -cd ~/brpc_svn/dev/brpc/$BRPCVERSION +cd ~/brpc_release/svn/dev/brpc/$BRPCVERSION -cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz ~/brpc_svn/dev/brpc/$BRPCVERSION +cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz ~/brpc_release/svn/dev/brpc/$BRPCVERSION -cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz.asc ~/brpc_svn/dev/brpc/$BRPCVERSION +cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz.asc ~/brpc_release/svn/dev/brpc/$BRPCVERSION -cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz.sha512 ~/brpc_svn/dev/brpc/$BRPCVERSION +cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz.sha512 ~/brpc_release/svn/dev/brpc/$BRPCVERSION ``` ### 4. 提交 SVN @@ -307,7 +317,7 @@ cp ~/brpc/apache-brpc-$BRPCVERSION-src.tar.gz.sha512 ~/brpc_svn/dev/brpc/$BRPCVE 退回到上级目录，使用 Apache LDAP 账号提交 SVN。 ```bash -cd ~/brpc_svn/dev/brpc +cd ~/brpc_release/svn/dev/brpc svn add * @@ -316,7 +326,7 @@ svn --username=$BRPCUSERNAME commit -m "release $BRPCVERSION" ## 检查发布结果 ```bash -cd ~/brpc_svn/dev/brpc/$BRPCVERSION +cd ~/brpc_release/svn/dev/brpc/$BRPCVERSION ``` ### 1. 检查 sha512 哈希 @@ -404,8 +414,8 @@ diff -r brpc-$BRPCVERSION apache-brpc-$BRPCVERSION-src ### 1. 投票阶段 -1. 发起投票邮件到 dev@brpc.apache.org。PMC 需要先按文档检查版本的正确性，然后再进行投票。经过至少 72 小时并统计到 3 个 +1 PMC member 票后，方可进入下一阶段。 -2. 宣布投票结果，发起投票结果邮件到 dev@brpc.apache.org。 +1. 发起投票邮件到 dev@brpc.apache.org。邮件发出后，在 [dev@brpc.apache.org 邮件归档](https://lists.apache.org/list.html?dev@brpc.apache.org) 找到该邮件并保存其永久链接。PMC 需要先按文档检查版本的正确性，然后再进行投票。经过至少 72 小时并统计到 3 个 +1 PMC member 票后，方可进入下一阶段。 +2. 宣布投票结果，发起投票结果邮件到 dev@brpc.apache.org；同样在邮件归档中保存结果邮件的永久链接。 ### 2. 投票邮件模板 @@ -487,7 +497,7 @@ Non-binding votes: - bbb - ccc -Vote thread: xxx (vote email link in https://lists.apache.org/) +Vote thread: {VOTE_THREAD_URL} Thank you to all the above members to help us to verify and vote for the 1.0.0 release. I will process to publish the release and send ANNOUNCE. diff --git a/community/release_en.md b/community/release_en.md index 8a314826d8..d48cf74363 100644 --- a/community/release_en.md +++ b/community/release_en.md @@ -1,5 +1,12 @@ # brpc apache release guide step by step +> **Start with the release Skill when possible**: [`community/skills/brpc-release/`](skills/brpc-release) orchestrates +> this process as resumable stages, provides scripts for version updates, Release Notes, packaging, and verification, +> and stops for the Release Manager (RM) before irreversible actions such as GPG signing, pushing tags, SVN commits, +> sending emails, and publishing a GitHub Release. See [`community/skills/README.md`](skills/README.md) for installation, +> dependency checks, and RM/verifier usage guidance. This document remains the authoritative release process; if the +> Skill differs from this guide, follow this guide and update the Skill accordingly. + ## Preparation ### 1. Confirm the release notes @@ -198,6 +205,14 @@ Edit the `/package/rpm/brpc.spec` file in the project root directory, update the Version: 1.0.0 ``` +#### Update the `CLAUDE.md` file + +Edit the `CLAUDE.md` file in the project root directory and update `Current version` in the project overview to the release version: + +``` +Current version: 1.0.0. +``` + #### Update the `MODULE.bazel` file Edit the `MODULE.bazel` file in the project root directory, update the version number, and submit it to the code repository. For example: @@ -258,13 +273,13 @@ sha512sum --check apache-brpc-1.0.0-src.tar.gz.sha512 If there is no local working directory, create a local working directory first. Checkout the Apache SVN repository, username needs to use your own Apache LDAP username: ```bash -mkdir -p ~/brpc_svn/dev/ +mkdir -p ~/brpc_release/svn/dev/ -cd ~/brpc_svn/dev/ +cd ~/brpc_release/svn/dev/ svn --username=lorinlee co https://dist.apache.org/repos/dist/dev/brpc/ -cd ~/brpc_svn/dev/brpc +cd ~/brpc_release/svn/dev/brpc ``` ### 2. Add GPG public key @@ -289,15 +304,15 @@ By fingerprint: ### 3. Add the releasing package to SVN directory ```bash -mkdir -p ~/brpc_svn/dev/brpc/1.0.0 +mkdir -p ~/brpc_release/svn/dev/brpc/1.0.0 -cd ~/brpc_svn/dev/brpc/1.0.0 +cd ~/brpc_release/svn/dev/brpc/1.0.0 -cp ~/brpc/apache-brpc-1.0.0-src.tar.gz ~/brpc_svn/dev/brpc/1.0.0 +cp ~/brpc/apache-brpc-1.0.0-src.tar.gz ~/brpc_release/svn/dev/brpc/1.0.0 -cp ~/brpc/apache-brpc-1.0.0-src.tar.gz.asc ~/brpc_svn/dev/brpc/1.0.0 +cp ~/brpc/apache-brpc-1.0.0-src.tar.gz.asc ~/brpc_release/svn/dev/brpc/1.0.0 -cp ~/brpc/apache-brpc-1.0.0-src.tar.gz.sha512 ~/brpc_svn/dev/brpc/1.0.0 +cp ~/brpc/apache-brpc-1.0.0-src.tar.gz.sha512 ~/brpc_release/svn/dev/brpc/1.0.0 ``` ### 4. Submit SVN @@ -305,7 +320,7 @@ cp ~/brpc/apache-brpc-1.0.0-src.tar.gz.sha512 ~/brpc_svn/dev/brpc/1.0.0 Return to the parent directory and use the Apache LDAP account to submit SVN ```bash -cd ~/brpc_svn/dev/brpc +cd ~/brpc_release/svn/dev/brpc svn add * @@ -314,7 +329,7 @@ svn --username=lorinlee commit -m "release 1.0.0" ## Verify release ```bash -cd ~/brpc_svn/dev/brpc/1.0.0 +cd ~/brpc_release/svn/dev/brpc/1.0.0 ``` ### 1. Verify SHA512 checksum @@ -406,8 +421,8 @@ This stage will cost 3+ days. ### 1. Vote stage -1. Send a voting email to `dev@brpc.apache.org`. PMC needs to check the correctness of the version according to the document before voting. After at least 72 hours and 3 +1 PMC member votes, you can move to the next stage. -2. Announce the voting result and send the voting result to dev@brpc.apache.org. +1. Send a voting email to `dev@brpc.apache.org`. After sending it, find the message in the [dev@brpc.apache.org mail archive](https://lists.apache.org/list.html?dev@brpc.apache.org) and save its permanent link. PMC needs to check the correctness of the version according to the document before voting. After at least 72 hours and 3 +1 PMC member votes, you can move to the next stage. +2. Announce the voting result and send the voting result to dev@brpc.apache.org; likewise save the permanent link of the result email from the mail archive. ### 2. Vote email template @@ -491,7 +506,7 @@ Non-binding votes: - bbb - ccc -Vote thread: xxx (vote email link in https://lists.apache.org/) +Vote thread: {VOTE_THREAD_URL} Thank you to all the above members to help us to verify and vote for the 1.0.0 release. I will process to publish the release and send ANNOUNCE. diff --git a/community/skills/README.md b/community/skills/README.md new file mode 100644 index 0000000000..3bec9ec7f7 --- /dev/null +++ b/community/skills/README.md @@ -0,0 +1,221 @@ +# 社区维护工作的 AI Agent Skills + +这里放给社区维护者使用的通用 AI agent skill，可供 GitHub Copilot、Claude Code、Codex、Cursor +等支持 skill 的编码助手使用。每个 skill 是一个自包含目录，权威流程仍然是 `community/` 下对应的 +`.md` 文档；skill 只负责把流程编排起来、自动化机械步骤、并在不可逆操作前刹车。 + +| Skill | 用途 | 对应文档 | +|---|---|---| +| `brpc-release` | Apache bRPC 发版全流程 | `community/release_cn.md`、`release_en.md`、`releasecheck.md` | + +## 安装指南 + +Skill 可以安装在项目级目录（仅当前仓库生效），也可以安装在用户级目录（所有项目生效）。推荐使用 +符号链接，仓库里的 skill 更新后无需重新复制。以下命令均在 bRPC 仓库根目录执行。 + +### Claude Code + +项目级目录为 `.claude/skills/`： + +```bash +mkdir -p .claude/skills +ln -sfn ../../community/skills/brpc-release .claude/skills/brpc-release +``` + +用户级安装： + +```bash +mkdir -p ~/.claude/skills +ln -sfn "$(pwd)/community/skills/brpc-release" ~/.claude/skills/brpc-release +``` + +### GitHub Copilot + +项目级目录为 `.github/skills/`： + +```bash +mkdir -p .github/skills +ln -sfn ../../community/skills/brpc-release .github/skills/brpc-release +``` + +### Codex + +项目级目录为 `.agents/skills/`，用户级目录为 `~/.agents/skills/`： + +```bash +mkdir -p .agents/skills +ln -sfn ../../community/skills/brpc-release .agents/skills/brpc-release + +# 或安装到用户级 +mkdir -p ~/.agents/skills +ln -sfn "$(pwd)/community/skills/brpc-release" ~/.agents/skills/brpc-release +``` + +### Cursor + +项目级目录为 `.cursor/skills/`，用户级目录为 `~/.cursor/skills/`： + +```bash +mkdir -p .cursor/skills +ln -sfn ../../community/skills/brpc-release .cursor/skills/brpc-release + +# 或安装到用户级 +mkdir -p ~/.cursor/skills +ln -sfn "$(pwd)/community/skills/brpc-release" ~/.cursor/skills/brpc-release +``` + +安装后重新打开会话，在 bRPC 仓库中输入「发版 」「验证一下这个发布包」等请求即可触发。 +支持斜杠命令的客户端也可以尝试 `/brpc-release`。如果客户端没有识别到 skill，请确认： + +1. 安装目录中的 `brpc-release/SKILL.md` 可访问； +2. 符号链接没有失效； +3. 客户端版本支持 Agent Skills，并已重新加载项目或会话； +4. 客户端采用了不同的 skill 目录时，以该客户端当前版本的配置说明为准。 + +## 依赖 + +| 依赖 | 用途 | 必需角色/阶段 | +|---|---|---| +| Git | 读取提交范围、创建 release 分支和本地 tag、生成源码包 | RM；阶段 1–4、9 | +| Bash（macOS/Linux） | 运行版本更新、Release Notes、打包和验包脚本 | RM、校验者 | +| GitHub CLI `gh` | 批量查询合并 PR 的标题、作者和编号以生成 PR 维度 Release Notes；首次使用需执行 `gh auth login` | RM；阶段 7 | +| GnuPG `gpg` | RM 交互式签名候选包，以及签名验证 | RM；阶段 2、4；校验者验证签名 | +| Subversion `svn` | 检出 Apache `dist/dev` 工作副本、提交候选包、移动正式发布包 | RM；阶段 5、8 | +| SHA-512 工具 | 生成和校验候选包哈希；脚本依次支持 `sha512sum`、`gsha512sum` 或 `shasum -a 512` | RM；阶段 4；校验者验证哈希 | +| 网络访问 | 访问 GitHub、Apache `dist/dev`、下载候选包和查询 PR | RM、校验者；按需 | + +`make_package.sh` 会调用 GPG 签名，必须由 RM 在交互式终端运行；`svn commit`、`svn mv`、推 tag、发邮件和发布 GitHub Release 也必须由 RM 手工执行。 + +### 执行前检查 + +每次启动 skill、恢复发版会话或执行依赖外部服务的步骤前，助手都必须先检查所需依赖和登录态；缺失时只说明 +安装、登录或授权方式，不能代填密码、token、私钥口令或其他凭据。建议按当前角色和计划步骤执行： + +```bash +# 基础依赖：所有脚本均需要 +command -v git +command -v bash + +# 按 PR 生成 Release Notes：需要 GitHub CLI 已登录且可访问目标仓库 +command -v gh +gh auth status --hostname github.com +gh repo view apache/brpc --json nameWithOwner >/dev/null + +# RM 打包签名：确认 GnuPG 及对应 Apache ID 的私钥可用 +command -v gpg +gpg --list-secret-keys --keyid-format=long "${APACHE_ID}@apache.org" + +# RM 上传候选包：确认 SVN 客户端及 Apache SVN 认证可用 +command -v svn +svn --username="$APACHE_ID" info https://dist.apache.org/repos/dist/dev/brpc/ >/dev/null + +# 打包或验包：确认至少存在一种 SHA-512 工具 +command -v sha512sum || command -v gsha512sum || command -v shasum +``` + +检查结果处理： + +- `gh auth status` 失败：提示用户在其终端执行 `gh auth login --hostname github.com`，完成后重新检查； +- `gpg --list-secret-keys` 没有目标私钥：回到 GPG 准备流程，不进入打包； +- `svn info` 要求认证或失败：由 RM 在交互式终端完成认证或确认权限后再继续； +- SHA-512 工具缺失：先安装可用实现，再执行打包或验包； +- 网络、仓库权限或登录态不满足：将该步骤标记为阻塞，汇报当前已完成和可选的其他步骤，等待 RM 选择。 + +## 使用指南 + +`brpc-release` 支持两种使用角色。开始新会话时应明确自己的角色，避免校验者误触发发版写操作。 + +### Release Manager（RM） + +RM 负责准备并发布候选版本。下面示例中的 ``、`` 和 `` 都是占位符；使用前请替换为本次版本、上一版本和自己的 Apache ID（例如 `your-id`），不要填写密码、token 或其他凭据。可以这样开始： + +```text +使用 brpc-release，以 RM 角色发布 ，我的 Apache ID 是 。 +先读取已有状态并告诉我下一步；遇到 GPG 签名、push tag、SVN commit/mv、发邮件或发布 GitHub Release时停止，给我命令，等我确认成功后再继续。 +``` + +助手会读取 `~/brpc_release//STATE.md`；若文件不存在则创建，然后按阶段继续。RM 需要准备： + +- 本次版本号、Apache ID 和上一版本号； +- 可访问 Apache SVN 的账号； +- 可用的 GPG 私钥； +- 对 `apache/brpc` 和发布网站所需的操作权限。 + +职责边界： + +| 助手可以执行 | 必须由 RM 执行 | +|---|---| +| 前置检查、创建 release 分支、更新版本文件、整理 SVN 工作副本、验证候选包 | `make_package.sh`（含本地 tag、打包、GPG 签名和校验） | +| 生成 Release Notes 和邮件草稿 | push tag、SVN commit/mv、发送邮件、发布 GitHub Release | + +阶段 4 由助手提供命令，RM 在交互式终端一次执行完成： + +```bash +export VERSION="" +export APACHE_ID="" +cd ~/brpc +BRPCUSERNAME="$APACHE_ID" community/skills/brpc-release/scripts/make_package.sh "$VERSION" +``` + +脚本执行期间由 GPG pinentry 向 RM 获取私钥口令。RM 明确回报脚本成功并提供 tag commit id 后，助手会 +更新状态文件，汇报已完成和未完成事项，并列出当前可选下一步；RM 选择后才继续，不会自动进入下一阶段。 + +中断数小时或数天后，可以输入以下指令恢复，不需要从头开始： + +```text +使用 brpc-release，以 RM 角色继续发布 。先读取状态文件，只执行尚未完成的步骤。 +``` + +### 候选包校验者（Verifier） + +校验者只检查候选包，不创建分支、不改版本号、不打 tag，也不向 SVN 或 GitHub 写入内容。可以直接使用 +公开的 `dist/dev` 地址： + +```text +使用 brpc-release，以校验者角色验证 Apache bRPC 候选包。只执行只读校验，不修改仓库、不签名、不上传或发布任何内容；最后给我校验结果和投票建议。 +``` + +也可以手工执行： + +```bash +export VERSION="" +community/skills/brpc-release/scripts/verify_package.sh \ + "https://dist.apache.org/repos/dist/dev/brpc/${VERSION}/" +``` + +校验范围包括下载链接、SHA512、GPG 签名、源码包命名、许可证文件、归档内容，以及源码包与 GitHub tag +的逐文件差异。校验者应根据实际结果独立投票，并在回复 `+1` 或 `-1` 时列出检查项；助手不能代发邮件。 +校验公开候选包通常不需要 Apache ID，只有访问受认证资源时才由用户在终端交互认证。 + +校验脚本只会创建临时工作目录，并在退出时自动删除该临时目录及其中下载、解压的校验副本；**不会**删除 +传入的本地产物目录、SVN 工作副本、Git tag、分支或任何正式发布文件。 + +## 脚本可以脱离 AI 助手单独用 + +`scripts/` 下的脚本都是普通 bash，不依赖任何 AI 助手，手工发版时照样能用： + +```bash +export VERSION="" +export PREV_VERSION="" +export APACHE_ID="" + +community/skills/brpc-release/scripts/bump_version.sh "$VERSION" +community/skills/brpc-release/scripts/bump_version.sh "$VERSION" --check +# 通过 GitHub PR（而不是 commit）生成草稿；终点使用本地 release tag，无需推送到 GitHub；需要先完成 gh auth login +community/skills/brpc-release/scripts/release_notes.sh "$PREV_VERSION" "$VERSION" +# 以下命令必须由 RM 在交互式终端执行： +BRPCUSERNAME="$APACHE_ID" community/skills/brpc-release/scripts/make_package.sh "$VERSION" +community/skills/brpc-release/scripts/verify_package.sh \ + "https://dist.apache.org/repos/dist/dev/brpc/${VERSION}/" +``` + +macOS 和 Linux 都验证过（sha512 工具会在 `sha512sum` / `gsha512sum` / `shasum -a 512` 之间自动选择）。 + +> `verify_package.sh` 是 `community/apache-package-validator.sh` 的可移植替代： +> 后者依赖 wget 和 GNU coreutils，在 macOS 上跑不起来。两者检查项基本一致， +> `verify_package.sh` 另外做了「源码包 vs GitHub tag」的逐文件 diff。 + +## 改动须知 + +发版流程变了，`community/release_cn.md`（权威）和这里的 skill 要一起改，否则 skill 会带偏 RM。 +版本号硬编码的位置变了，记得同步 `scripts/bump_version.sh` 里的文件表和 `references/checklist.md` 的清单。 diff --git a/community/skills/brpc-release/SKILL.md b/community/skills/brpc-release/SKILL.md new file mode 100644 index 0000000000..41cd3752a7 --- /dev/null +++ b/community/skills/brpc-release/SKILL.md @@ -0,0 +1,286 @@ +--- +name: brpc-release +description: Drive an Apache bRPC release end to end - confirm the release slot, draft Release Notes, bump version files, cut the tag, build/sign/checksum the source tarball, stage it to Apache SVN dist/dev, verify the candidate, and draft the VOTE/RESULT/ANNOUNCE mails. Use when cutting, resuming, or verifying a bRPC release ("发版 ", "release apache brpc", "准备 RC", "写投票邮件", "验证发布包", "check the release candidate"). +--- + +# Apache bRPC 发版 + +面向 Release Manager（RM）的发版助手。权威流程是 `community/release_cn.md`（英文 `community/release_en.md`）， +本 skill 不取代它，只做三件事：**编排顺序、自动化机械步骤、在不可逆操作前刹车**。 + +## 铁律：这些命令永远不要代 RM 执行 + +发版有大量对外且不可撤回的动作。遇到下面这些，**把命令打印出来让 RM 自己跑**，然后等待他明确回报成功； +未收到成功确认前不得进入下一步： + +| 动作 | 为什么不能代跑 | +|---|---| +| `git push origin --tags` | tag 推到 apache/brpc 后立刻被镜像和下游 CI 抓取，回退要惊动 ASF INFRA | +| `svn commit` / `svn mv` | 写入 Apache 官方分发仓库；`dist/release` 下的内容会同步到全球镜像 | +| `gpg --detach-sign` | 模型不能执行签名，也不能接触私钥口令；必须由 RM 在交互式终端执行 | +| `git config user.name` / `git config user.email` | 发版不要求修改 Git 身份；模型不得读取后重写、覆盖或临时修改用户配置 | +| 发送任何邮件 | `dev@` / `announce@` 是公开存档，发出即无法撤回 | +| 发布 GitHub Release | 对外可见 | + +改版本号、打本地 tag、打包、算校验和、生成草稿这些**本地可逆操作**可以直接做。 + +## 安全条款与强制检查 + +发版助手必须采用 fail-closed 原则：前置条件不明确或检查失败时立即停止，不猜测、不绕过，也不以 +`--force`、临时改脚本或跳过校验的方式继续。 + +在任何会修改仓库或生成候选包的步骤前，必须逐项检查： + +1. 当前目录是预期的 Apache bRPC Git 仓库，且远端指向可识别的 bRPC 仓库； +2. 目标版本符合 `X.Y.Z` 或 `X.Y.Z-rcNN`，目标分支严格为 `release-X.Y`； +3. 当前分支已是目标 release 分支；minor 首发分支必须从 `origin/master` 创建，不能先改版本号再切分支； +4. 工作区没有与本次步骤无关的未提交修改；发现用户修改时不得覆盖、丢弃或自动 stash； +5. tag 不存在，或已存在且明确指向预期 commit；不得静默覆盖或移动 tag； +6. 待发布 commit 来自目标 release 分支，且版本文件检查全部通过； +7. 即将执行的操作不在上方“永远不要代 RM 执行”的列表中。 + +所有脚本必须使用严格模式、校验外部命令返回值，并将错误输出到 stderr。涉及版本号、分支、tag、 +产物路径的参数不得直接拼接成未经验证的命令。日志和草稿不得包含私钥、密码、token、cookie 或其他凭据。 +发版沿用仓库已有的 Git 身份，不要求 Apache 邮箱；模型不得执行任何修改本地或全局 Git 用户名、邮箱的命令。 + +## 依赖与登录态检查 + +每次启动、恢复会话，以及执行依赖外部服务的步骤前，必须先检查当前步骤所需的命令、网络、账号权限和 +登录态。基础检查包括 `git`、`bash`；阶段 2/4 检查 `gpg` 和 `${APACHE_ID}@apache.org` 私钥；阶段 5/8 +检查 `svn` 与 Apache SVN 认证；阶段 6 检查网络和 SHA-512 工具；阶段 7 检查 `gh`、GitHub 登录态和对 +`apache/brpc` 的访问权限。 + +检查失败时，将对应步骤标记为阻塞，汇报已完成项和其他可选下一步，等待 RM 选择；不得跳过检查、代填密码、 +token 或私钥口令。具体命令和失败处理见 `community/skills/README.md` 的“执行前检查”。 + +## 每次开工前先做这三件事 + +1. **读状态文件** `~/brpc_release//STATE.md`。发版横跨数周，绝大多数会话是接着上次继续， + 不要默认从头开始。文件不存在时，用 `references/state_template.md` 建一个。 +2. **确认版本号、RM 和 Apache ID**：查 `community/release_schedule.md` 里的排期表，并向 RM 询问 + Apache ID（仅用户名，不询问密码、token 或其他凭据）。Apache ID 用于以 `--username` 参数检出和操作 + Apache SVN 仓库，并写入状态文件。补丁版本 `${MAJOR}.${MINOR}.${PATCH}` 在已有的 + `release-${MAJOR}.${MINOR}` 分支上发；minor 首发版本 `${MAJOR}.${MINOR}.0` 要从 master 拉新分支。 + **先创建或切换 release 分支，再做任何版本号修改**；不得在 master 或其他工作分支先改版本号。 +3. **确认目录约定**（与 `release_cn.md` 一致，脚本也按这个假设）： + + | 路径 | 用途 | + |---|---| + | `~/brpc` | release 分支的 clone，打包在这里做 | + | `~/brpc_release/svn/dev/brpc` | Apache SVN `dist/dev` 工作副本 | + | `~/brpc_release//` | 本次发版的草稿、邮件、状态文件 | + +每完成一个阶段，**立刻回写 STATE.md**（勾选 + 记下关键产物，如 tag 的 commit id、投票邮件链接）。 + +## 进度汇报与下一步选择 + +每完成一个步骤或收到 RM 对手工步骤的成功回报后，助手必须先读取并回写 `STATE.md`，再向 RM 汇报： + +1. **本次完成**：列出刚完成的阶段或子步骤，以及产物路径、tag commit id、候选包 URL 等关键结果； +2. **当前进度**：列出已完成项和未完成项；对于被跳过的可选项，明确标记为“跳过”及原因； +3. **前置条件/阻塞项**：说明下一步是否需要 RM 执行 GPG、push tag、SVN commit/mv、发邮件或其他外部操作； +4. **可选下一步**：只列出当前前置条件已满足的步骤，按推荐顺序编号；如果工作可并行，明确标为“可并行”； +5. **等待选择**：停止执行，询问 RM 选择哪一项。除非 RM 明确指示连续执行，否则助手不能自动进入下一阶段。 + +推荐使用以下固定格式，保证跨会话可恢复： + +```text +当前发版进度： +本次完成： +- [x] <刚完成的步骤与关键产物> + +已完成：<编号列表> +未完成：<编号列表> +阻塞项：<无 / 需要 RM 执行的动作> + +可选下一步： +1. <推荐步骤> +2. <可并行步骤> +3. <其他已满足前置条件的步骤> + +请回复序号或直接说明希望继续的工作。 +``` + +RM 可以随时选择任何**前置条件已满足**的未完成项，例如“先生成 Release Notes”“先验证本地产物”或“继续阶段 5”。 +助手应先核对依赖关系，满足则执行并更新状态；不满足则解释所缺前置条件，并继续等待 RM 选择。 + +## 阶段地图 + +流程分 9 个阶段，跨度约 2 周。详细检查项在 `references/checklist.md`，邮件模板在 +`references/mail_templates.md`——**按需读取，不要一次性全读进来**。 + +### 阶段 1 — 确认发布范围（约 1 周） + +和社区确认本次发布范围，合入“计划发布但还没进来”的 PR，并冻结候选 commit 范围。 +此时不生成 Release Notes；分支、版本和候选包尚未稳定，提前生成容易遗漏或反复返工。 + +### 阶段 2 — GPG 准备（非首次发版跳过） + +先检查有没有可用的密钥： + +```bash +gpg --list-secret-keys --keyid-format=long @apache.org +``` + +有输出就跳过。没有则照 `references/checklist.md` 的「GPG 首次设置」走完四步：建 4096 位 RSA 密钥 +（邮箱必须是 Apache 邮箱）→ 发布公钥到 keyserver → 把 fingerprint 填进 https://id.apache.org → +把公钥追加进 SVN 的 `KEYS` 文件。 + +### 阶段 3 — 拉发版分支 + 改版本号 + +顺序不可颠倒：**先拉/切换发版分支，确认当前分支正确，再改版本号**。脚本会强制检查目标 release +分支、仓库身份和工作区状态；任一检查不通过都不会写文件。 + +```bash +export VERSION="" +export APACHE_ID="" +export RELEASE_BRANCH="release-${VERSION%.*}" + +# minor 首发版本才需要从 master 拉新分支 +# 如工作区已有未提交修改，应先提交或暂存，不能为了切分支而提前改版本号 +git checkout -b "$RELEASE_BRANCH" origin/master +git branch --show-current # 必须输出 $RELEASE_BRANCH + +# 改发布所需的硬编码版本号（幂等，可重复跑） +scripts/bump_version.sh "$VERSION" + +# 只校验不修改 +scripts/bump_version.sh "$VERSION" --check +``` + +脚本只更新发布元数据及用户文档；`example/build_with_bazel_module/MODULE.bazel` 中示例模块自身的 +`version` 和示例依赖的 bRPC `version` 都不跟随本次发版修改。脚本结尾会扫描其他旧版本号残留。 +另外**检查 `NOTICE` 的年份**（`Copyright 2018-<当年>`），年初发版时尤其容易漏。 + +发版过程中发现问题，一律在 release 分支上改，不要回 master。 + +### 阶段 4 — 打 tag、打包、签名、校验和 + +模型不得执行本阶段脚本，只能根据已确认的 Apache ID 和版本号向 RM 提供完整命令： + +```bash +cd ~/brpc +BRPCUSERNAME="$APACHE_ID" community/skills/brpc-release/scripts/make_package.sh "$VERSION" +``` + +由 RM 在交互式终端执行。脚本会一次完成：校验工作区和版本号 → 创建**本地** tag → 生成源码包 → +调用 GPG 并等待 RM 输入私钥口令 → 验证签名 → 生成并验证 `.sha512`。脚本不会修改 Git 用户名或邮箱， +也不会推送 tag。 + +模型提供命令后必须停止。只有 RM 明确回报脚本成功，并提供输出中的 tag commit id 后，才把阶段 4 +标为完成并进入下一阶段；失败或结果不明确时不得继续。`git push origin --tags` 仍由 RM 单独执行。 + +### 阶段 5 — 上传到 Apache SVN dist/dev + +如果状态文件中没有 Apache ID，先向 RM 询问；只记录 ID，**不得询问、记录或代填密码**。首次使用时， +可代 RM 使用 Apache ID 检出 SVN 工作副本（认证需要密码时由 RM 在终端交互输入）： + +```bash +export BRPCUSERNAME="$APACHE_ID" +mkdir -p ~/brpc_release/svn/dev/ +svn --username="$BRPCUSERNAME" co https://dist.apache.org/repos/dist/dev/brpc/ ~/brpc_release/svn/dev/brpc +``` + +按 `references/checklist.md` 的「SVN 上传」小节，把三个文件（`.tar.gz` / `.asc` / `.sha512`） +放进 `~/brpc_release/svn/dev/brpc//`，`svn add` 之后**由 RM 执行 `svn commit`**。 +首次发版的 RM 还要先把公钥追加进 `KEYS` 并一起提交。 + +### 阶段 6 — 验证候选包 + +```bash +# 验证已上传到 dist/dev 的包（联网下载并全面校验） +scripts/verify_package.sh "https://dist.apache.org/repos/dist/dev/brpc/${VERSION}/" + +# 或验证本地产物 +scripts/verify_package.sh "${HOME}/brpc_release/${VERSION}" +``` + +对标仓库里已有的 `community/apache-package-validator.sh`（那个脚本依赖 wget 和 GNU coreutils）， +本脚本是可在 macOS 直接运行的独立替代实现，并额外做一项文档里要求但那个脚本没覆盖的 +检查：**源码包与 GitHub tag 逐文件 diff**。 + +自己验过一遍再发投票邮件——PMC 常见的 -1 原因见 `references/checklist.md`。 + +### 阶段 7 — 准备 Release Notes 并发起投票（至少 72 小时） + +候选包验证通过后，以最终 tag 为边界生成 Release Notes 草稿： + +```bash +export PREV_VERSION="" +# 使用本地 release tag 作为终点；tag 无需推送到 GitHub。 +scripts/release_notes.sh "$PREV_VERSION" "$VERSION" \ + > "${HOME}/brpc_release/${VERSION}/notes-draft.md" +``` + +脚本以最终 tag 范围内的 GitHub PR 为唯一条目来源：每个合并 PR 只输出一项，并从 GitHub 获取 PR 标题和 +作者；输出格式为 `@GitHub-ID #PR-id`，不生成显式 PR URL，GitHub Release 会自动识别链接。不会再按 commit +逐条生成。分类仍是 Feature / Bugfix / Enhancement / Other 的启发式猜测，必须人工核对完整性和准确性；没有 +关联 GitHub PR 的 commit 会在草稿注释中单独提示，供人工确认归属。脚本按每批最多 50 个 PR 的 GraphQL +请求加载元数据，避免逐 commit 查询导致长时间等待。按渠道准备： +- GitHub Release：完整 Release Notes，逐项保留 `by @GitHub-ID (#PR-id)`，不附显式链接；分类内按功能主题 + 排序。同一 GitHub ID 的不同功能或修复必须作为独立条目保留； +- 微信公众号：与 GitHub Release 保持相同的“新功能 / Bug 修复 / 功能增强”分类和 PR 条目粒度，每项保留 + `by @GitHub-ID (#PR-id)`，不写 PR URL。阅读相关 PR 的说明和必要上下文后，以自然中文润色功能描述； + 不逐字直译、不照搬生硬的提交标题，无法确认含义时不臆测； +- VOTE 邮件的 Release Note：只挑重要 PR，Feature 优先，不粘贴完整列表； +- ANNOUNCE 邮件：只写几条主要变化，不标贡献者和 PR 编号。 + +从 `references/mail_templates.md` 取 `[VOTE]` 模板，填好版本号、精选 Release Note、 +**tag 的 commit id**，草稿写到 `~/brpc_release//vote-mail.txt`，**由 RM 发到 +dev@brpc.apache.org**。邮件发出后，在 [dev@brpc.apache.org 邮件归档](https://lists.apache.org/list.html?dev@brpc.apache.org) +找到该 `[VOTE]` 邮件并复制其永久链接，立即写入 `STATE.md` 的“投票邮件链接”。 + +需要至少 72 小时 + 3 张 PMC binding +1。够票后从同一归档页面找到 `[RESULT][VOTE]` 邮件的永久链接， +回写到 `STATE.md` 的“已完成记录”，再用 `[RESULT][VOTE]` 模板宣布结果。 + +投票没过：在 release 分支修问题，回到阶段 3 重新打包，版本号不变但要重打 tag。 + +### 阶段 8 — 完成发布 + +依次（每一步都由 RM 执行）： + +1. **PMC 成员**把包从 `dist/dev` 移到 `dist/release`（`svn mv`） +2. 在 GitHub 对应 tag 上创建 Release，标题统一为 `Apache bRPC ${VERSION}` +3. 等包同步到 Apache 镜像后，更新 https://brpc.apache.org/docs/downloadbrpc/ + （在 `apache/brpc-website` 仓库，中英文都要改） + - 签名和哈希链接前缀：`https://downloads.apache.org/brpc/` + - 代码包链接前缀：`https://dlcdn.apache.org/brpc/` +4. 用 `[ANNOUNCE]` 模板发信到 `dev@brpc.apache.org` 和 `announce@apache.org` + —— 必须用**个人 apache 邮箱**、必须**纯文本格式**，announce@ 要人工审核约一天 + +### 阶段 9 — 收尾 + +- 把 release 分支合回 master +- 更新 `community/release_schedule.md` 里本次发版的实际日期 +- 微信公众号等外部渠道（可选） + +完成上述事项并确认正式发布可下载后，助手应汇报可清理的本地发版产物，并由 RM 自行决定是否执行。 +**不得自动删除**，也不得删除 Git 分支、tag、SVN 工作副本或 `STATE.md`。建议 RM 在确认无需保留本地 +归档、且 `STATE.md` 已记录 tag commit id、SHA512、投票链接等关键事实后，执行： + +```bash +export VERSION="" + +# 仅删除本次生成的本地源码包、签名、哈希和草稿；保留 STATE.md 作为发版记录。 +rm -f "${HOME}/brpc_release/${VERSION}"/apache-brpc-"${VERSION}"-src.tar.gz{,.asc,.sha512} +rm -f "${HOME}/brpc_release/${VERSION}"/notes-draft.md \ + "${HOME}/brpc_release/${VERSION}"/wechat-draft.md \ + "${HOME}/brpc_release/${VERSION}"/{vote-mail,result-mail,announce-mail}.txt + +# 若确认不再需要整个本次草稿目录，先检查内容，再由 RM 手工删除。 +ls -la "${HOME}/brpc_release/${VERSION}" +# rm -rf "${HOME}/brpc_release/${VERSION}" +``` + +SVN 工作副本 `~/brpc_release/svn/dev/brpc` 通常应保留供下次发版复用；若 RM 明确要清理，先确认不存在 +未提交修改（`svn status`），再由 RM 手工删除该工作副本。 + +## 常见坑 + +- **改了版本号却漏了文件**：一定跑 `bump_version.sh --check`，别手改。 +- **在 master 上打 tag**：所有发版操作都在 `release-` 分支。 +- **sha512 文件里带路径**：`sha512sum` 必须在 tarball 所在目录执行，否则校验方 `--check` 会失败。 + 脚本已处理。 +- **macOS 没有 `sha512sum`**：脚本会自动回退到 `gsha512sum` 或 `shasum -a 512`。 +- **投票邮件写错 commit id**：填 tag 指向的 commit，不是分支 HEAD。 +- **只回 `+1` 不附检查项**：PMC 投票必须列出检查了哪些项。 diff --git a/community/skills/brpc-release/references/checklist.md b/community/skills/brpc-release/references/checklist.md new file mode 100644 index 0000000000..17d23edf69 --- /dev/null +++ b/community/skills/brpc-release/references/checklist.md @@ -0,0 +1,212 @@ +# 发版检查清单 + +`SKILL.md` 里省略的细节都在这里。按需查阅对应小节，不必通读。 + +--- + +## GPG 首次设置 + +只有**从未发过版**的 RM 需要做。已有密钥的直接跳到「SVN 上传」。 + +> **口令安全铁律**：模型不得执行 `make_package.sh` 或任何签名命令，只能把完整的脚本命令提供给 RM。 +> RM 在交互式终端运行脚本，由 GPG pinentry 获取私钥口令。任何要求用 `--passphrase`、 +> `--pinentry-mode loopback` 或环境变量传递口令的做法都要**拒绝**——私钥口令绝不落盘、 +> 绝不进命令行历史、绝不交给 agent。脚本会完成签名、签名验证和 SHA512 校验；RM 明确回报成功前， +> 模型不得进入下一阶段。 + +### 1. 安装 + +```bash +brew install gnupg # macOS +# Linux 发行版通常自带 GnuPG +gpg --version +``` + +### 2. 创建密钥 + +```bash +gpg --full-gen-key +``` + +交互式提示的关键选择： + +| 提示 | 填什么 | +|---|---| +| kind of key | `1`（RSA and RSA） | +| keysize | `4096` | +| valid for | `0`（永不过期） | +| Real name | 姓名拼音 / Apache ID / GitHub ID 均可 | +| Email address | **必须是 `@apache.org`** | +| Passphrase | 设一个并记牢，后面每次签名都要输 | + +生成后会打印公钥 ID，形如 `C30F211F071894258497F46392E18A11B6585834`。 + +### 3. 发布公钥到 keyserver + +```bash +gpg --keyserver hkps://pgp.mit.edu --send-key <公钥ID> +``` + +`hkps://keys.openpgp.org` 和 `hkps://keyserver.ubuntu.com` 也可以，都提供 Web 查询界面。 + +### 4. 登记 fingerprint + +```bash +gpg --fingerprint <用户ID> +``` + +把输出里带空格的指纹（如 `C30F 211F 0718 9425 8497 F463 92E1 8A11 B658 5834`）粘贴到 +https://id.apache.org 的 `OpenPGP Public Key Primary Fingerprint:` 字段。 + +公钥服务器没有校验机制，任何人都能以你的名义上传公钥，所以必须在 Apache 官方渠道公布指纹供人核对。 + +### 5. 把公钥追加进 SVN 的 KEYS + +在 `~/brpc_release/svn/dev/brpc` 目录下： + +```bash +(gpg --list-sigs $BRPCUSERNAME && gpg -a --export $BRPCUSERNAME) >> KEYS +``` + +同名密钥有多个时，用完整邮箱或公钥 ID 指定： + +```bash +(gpg --list-sigs $BRPCUSERNAME@apache.org && gpg -a --export $BRPCUSERNAME@apache.org) >> KEYS +# 或 +(gpg --list-sigs <公钥ID> && gpg -a --export <公钥ID>) >> KEYS +``` + +--- + +## 版本号文件清单 + +`bump_version.sh` 覆盖下面全部 7 个文件。手改容易遗漏，应优先使用脚本并执行 `--check` 校验。 + +| 文件 | 形式 | +|---|---| +| `RELEASE_VERSION` | 整个文件就是版本号 | +| `CMakeLists.txt` | `set(BRPC_VERSION 1.18.0)` | +| `package/rpm/brpc.spec` | `Version:\t1.18.0` | +| `CLAUDE.md` | `Current version: 1.18.0.` | +| `MODULE.bazel` | `module(... version = '1.18.0' ...)` | +| `docs/cn/bazel_support.md` | `bazel_dep(name = "brpc", version = "1.18.0", ...)` | +| `docs/en/bazel_support.md` | 同上 | + +`example/build_with_bazel_module/MODULE.bazel` 是示例配置，其中的 module 版本和 bRPC 依赖版本不随本次发版修改。 + +另外手工确认： + +- `NOTICE` 的年份是 `Copyright 2018-<当前年份>`（年初发版必查） + +发版不要求修改 Git 身份。模型和脚本不得执行 `git config user.name` 或 `git config user.email`； +若创建 annotated tag 时发现身份未配置，只提示 RM 自行处理并停止。 + +--- + +## SVN 上传 + +先从状态文件读取 Apache ID；若尚未记录，则向 RM 询问。这里只需要 Apache ID（用户名），不要询问或保存 +Apache LDAP 密码。`svn` 需要认证时，由 RM 在终端中交互输入密码。 + +```bash +export BRPCVERSION=1.18.0 +export BRPCUSERNAME=<你的 apache id> + +# 首次：检出 dist/dev 工作副本 +mkdir -p ~/brpc_release/svn/dev/ +svn --username=$BRPCUSERNAME co https://dist.apache.org/repos/dist/dev/brpc/ ~/brpc_release/svn/dev/brpc + +# 放入三个产物（源目录是 make_package.sh 的输出目录） +mkdir -p ~/brpc_release/svn/dev/brpc/$BRPCVERSION +cp ~/brpc_release/$BRPCVERSION/apache-brpc-$BRPCVERSION-src.tar.gz{,.asc,.sha512} \ + ~/brpc_release/svn/dev/brpc/$BRPCVERSION/ + +cd ~/brpc_release/svn/dev/brpc +svn add --force . +``` + +最后一步**由 RM 亲自执行**： + +```bash +svn --username=$BRPCUSERNAME commit -m "release $BRPCVERSION" +``` + +首次发版的 RM 记得把「GPG 首次设置」第 5 步生成的 `KEYS` 变更一起提交。 + +--- + +## 候选包检查项 + +`verify_package.sh` 自动覆盖的： + +- [ ] 下载链接有效 +- [ ] sha512 哈希正确 +- [ ] GPG 签名正确 +- [ ] 包内 `RELEASE_VERSION` 和 `CMakeLists.txt` 的版本号与本次发布一致 +- [ ] 存在 `LICENSE` 和 `NOTICE` +- [ ] 不含编译产物 / 意外的二进制文件 +- [ ] 源码包内容与 GitHub tag 完全一致 + +需要**人工判断**的： + +- [ ] `NOTICE` 年份正确 +- [ ] 源码包体积合理，没有夹带无关文件 +- [ ] 所有源文件都有 ASF License 头（本仓库用 skywalking-eyes 在 CI 里查，配置见 `.licenserc.yaml`） +- [ ] 能正常编译，单测通过 +- [ ] 没有空目录等多余文件夹 +- [ ] 第三方依赖许可证： + - 许可证兼容（见下方分类） + - 所有第三方依赖都在 `LICENSE` 中声明 + - 依赖许可证全文都在 `licenses/` 目录 + - 依赖若是 Apache 许可证且带 `NOTICE`，其 NOTICE 内容要并入本项目 `NOTICE` + +### ASF 许可证分类 + +| 类别 | 含义 | 例子 | +|---|---|---| +| Category A | 允许 | Apache-2.0, BSD-3-Clause, MIT | +| Category B | 允许依赖，但不允许放进源码包 | EPL, MPL, CDDL | +| Category X | **禁止** | GPL, LGPL, CC Non-Commercial | + +### 导入他人公钥做验证 + +帮别人验包时需要先导入并信任发布人公钥（RM 验自己的包不需要）： + +```bash +curl https://dist.apache.org/repos/dist/dev/brpc/KEYS | gpg --import +gpg --edit-key <发布人用户名> +# gpg> trust → 选 5（ultimate）→ y → save +``` + +### 常见的 -1 原因 + +- 包名不对，和当前发布版本对不上 +- 签名或哈希校验失败 +- 源码包里混进了编译产物或 jar/so 等二进制 +- 缺 `LICENSE` / `NOTICE`，或 NOTICE 年份过期 +- 源文件缺 ASF License 头 +- 引入了 Category X 许可证的依赖 + +--- + +## PMC 投票回复格式 + +**不要只回 `+1`**，必须列出实际检查了哪些项： + +``` ++1 (binding) + +I checked: +- LICENSE and NOTICE are good +- signatures and hashes correct +- All ASF files have ASF headers +- no unexpected binary files +- source distribution matches the git tag +- builds and unit tests pass +``` + +投 `-1` 同样必须给出明确理由。 + +清单来源：Incubator PMC Chair Justin 在 ApacheCon North America 2019 的分享 +（https://training.apache.org/topics/ApacheWay/NavigatingASFIncubator/index.html）， +详见 `community/releasecheck.md`。 diff --git a/community/skills/brpc-release/references/mail_templates.md b/community/skills/brpc-release/references/mail_templates.md new file mode 100644 index 0000000000..f11942c0c7 --- /dev/null +++ b/community/skills/brpc-release/references/mail_templates.md @@ -0,0 +1,323 @@ +# 发版邮件模板 + +三封邮件，按顺序发。全部用**个人 Apache 邮箱**、**纯文本格式**（Gmail 里选「纯文本模式」）。 + +Apache 邮箱配置参考 https://shenyu.apache.org/zh/community/use-apache-email ， +注意 SMTP 服务器要填 `mail-relay.apache.org`。 + +占位符：`{VERSION}` 版本号、`{COMMIT}` tag 指向的 commit id、`{RM}` 发布人名字。 + +--- + +## 1. 投票邮件 → dev@brpc.apache.org + +**标题** + +``` +[VOTE] Release Apache bRPC {VERSION} +``` + +**正文** + +``` +Hi Apache bRPC Community, + +This is a call for vote to release Apache bRPC version {VERSION} + +[Release Note] + +Features: + +- {feature description} by @github-id (#1234) +- {feature description} by @github-id (#1235) + +Bugfixes: + +- {bugfix description} by @github-id (#1236) +- {bugfix description} by @github-id (#1237) + +Enhancements: + +- {enhancement description} by @github-id (#1238) + +The release candidates: +https://dist.apache.org/repos/dist/dev/brpc/{VERSION}/ + +Git tag for the release: +https://github.com/apache/brpc/releases/tag/{VERSION} + +Release Commit ID: +https://github.com/apache/brpc/commit/{COMMIT} + +Keys to verify the Release Candidate: +https://dist.apache.org/repos/dist/dev/brpc/KEYS + +The vote will be open for at least 72 hours or until the necessary number of +votes are reached. + +Please vote accordingly: +[ ] +1 approve +[ ] +0 no opinion +[ ] -1 disapprove with the reason + +PMC vote is +1 binding, all others are +1 non-binding. + +Checklist for reference: +[ ] Download links are valid. +[ ] Checksums and PGP signatures are valid. +[ ] Source code distributions have correct names matching the current release. +[ ] LICENSE and NOTICE files are correct for each brpc repo. +[ ] All files have license headers if necessary. +[ ] No compiled archives bundled in source archive. + +Regards, +{RM} +``` + +> `[Release Note]` 只挑重要 PR，Feature 优先，不要粘贴完整 Release Notes；完整列表留给 GitHub Release。 +> +> `Release Commit ID` 填 **tag 指向的 commit**，不是分支 HEAD。 +> 取值：`git rev-parse {VERSION}^{commit}` + +投票邮件发出后，在 [dev@brpc.apache.org 邮件归档](https://lists.apache.org/list.html?dev@brpc.apache.org) 找到 +对应 `[VOTE]` 邮件并复制其永久链接，填写到 `STATE.md` 的“投票邮件链接”。投票需开放至少 72 小时， +且收集到 **3 张 PMC binding +1** 才能进入下一步。 + +--- + +## 2. 投票回复模板 + +投票者应根据实际检查结果回复投票邮件。PMC 成员使用 `+1 (binding)`，非 PMC 成员使用 +`+1 (non-binding)`；不得在未完成检查时勾选通过项。 + +**标题** + +保持邮件客户端的回复标题，例如： + +``` +Re: [VOTE] Release Apache bRPC {VERSION} +``` + +**正文（PMC 成员）** + +```text ++1 (binding) + +I have checked: + +[x] Download links are valid. +[x] Checksums and PGP signatures are valid. +[x] Source code distributions have correct names matching the current release. +[x] LICENSE and NOTICE files are correct for each brpc repo. +[x] All files have license headers if necessary. +[x] No compiled archives bundled in source archive. + +Best regards, +{VOTER_NAME} +``` + +**正文（非 PMC 成员）** + +将首行替换为： + +```text ++1 (non-binding) +``` + +其余检查清单和签名保持不变。 + +> 按实际检查情况填写。若有未完成项，保留 `[ ]` 并说明原因；发现阻塞性问题时投 `-1` 并附可复现信息。 +> 校验命令见 `verify_package.sh`，其输出可帮助填写上述清单。 + +--- + +## 3. 结果邮件 → dev@brpc.apache.org + +**标题** + +``` +[RESULT] [VOTE] Release Apache bRPC {VERSION} +``` + +**正文** + +``` +Hi all, + +The vote to release Apache bRPC {VERSION} has passed. + +The vote PASSED with 3 binding +1, 3 non binding +1 and no -1 votes: + +Binding votes: +- xxx +- yyy +- zzz + +Non-binding votes: +- aaa +- bbb +- ccc + +Vote thread: {VOTE_THREAD_URL} + +Thank you to all the above members to help us to verify and vote for +the {VERSION} release. I will process to publish the release and send ANNOUNCE. + +Regards, +{RM} +``` + +> 票数要按实际统计填写。`Vote thread` 填 `[VOTE]` 邮件的永久链接：在 +> [dev@brpc.apache.org 邮件归档](https://lists.apache.org/list.html?dev@brpc.apache.org) 中找到该邮件后复制链接。 +> `[RESULT][VOTE]` 发出后，也应复制其永久链接并写入 `STATE.md` 的“已完成记录”。 + +--- + +## 4. 发布公告 → dev@brpc.apache.org + announce@apache.org + +在包已经从 `dist/dev` 移到 `dist/release`、GitHub Release 已发布之后再发。 + +**标题** + +``` +[ANNOUNCE] Apache bRPC {VERSION} released +``` + +**正文** + +``` +Hi all, + +The Apache bRPC community is glad to announce the new release +of Apache bRPC {VERSION}. + +Apache bRPC is an Industrial-grade RPC framework using C++ Language, +which is often used in high performance systems such as Search, Storage, +Machine learning, Advertisement, Recommendation etc. + +Brief notes of this release: +- xxx +- yyy +- zzz + +More details regarding Apache brpc can be found at: +https://brpc.apache.org/ + +The release is available for download at: +https://brpc.apache.org/download/ + +The release notes can be found here: +https://github.com/apache/brpc/releases/tag/{VERSION} + +Website: https://brpc.apache.org/ + +Apache bRPC Resources: +- Issue: https://github.com/apache/brpc/issues/ +- Mailing list: dev@brpc.apache.org +- Documents: https://brpc.apache.org/docs/ + +We would like to thank all contributors of the Apache bRPC community +who made this release possible! + + +Best Regards, +Apache bRPC Community +``` + +> `Brief notes of this release` 只列本次的**主要**变更，不要贴完整 Release Notes， +> 也不用标注贡献人和 PR 编号。建议先翻一下 lists.apache.org 上之前的 ANNOUNCE 邮件对齐风格。 +> +> `announce@apache.org` 需人工审核，发出后耐心等，一般一天内通过。 + +--- + +## Release Notes 模板 + +阶段 7 在候选包验证通过后，需分别准备 GitHub Release 英文版和微信公众号中文版。两者均以**合并 PR +为单位**：每个 PR 仅一项，不按该 PR 的多个 commit 分拆。`scripts/release_notes.sh` 生成的内容仅作为 +待人工打磨的初稿。 + +### GitHub Release（英文） + +参考 Apache bRPC 1.16.0 的结构：用一段简洁摘要概括本次版本主题，再按 `Features`、`Bugfixes`、 +`Enhancements` 分类列出完整变更，最后致谢全部贡献者。每项使用 `@GitHub-ID #PR-id`；不要写显式 PR +URL，GitHub 会自动将 `#PR-id` 识别为链接。 + +```markdown +Apache bRPC {VERSION} is a {feature/maintenance} release that includes {one-sentence summary of the most important improvements}. This release {user-visible value summary}. + +## Features + +- {feature description} by @github-id (#1234) +- {feature description} by @github-id (#1235) + +## Bugfixes + +- {bugfix description} by @github-id (#1236) +- {bugfix description} by @github-id (#1237) + +## Enhancements + +- {enhancement description} by @github-id (#1238) + +## Other + +- {documentation, test, build, or maintenance description} by @github-id (#1239) +``` + +写作要求： + +- 摘要面向用户说明版本价值，避免仅罗列内部实现； +- 变更描述使用动词开头的简洁英文，保留协议名、API、类名和配置项等必要技术术语； +- 完整列出所有应发布的合并 PR；文档、测试、CI、构建和维护类变更归入 `Other`； +- **变更条目**按 `Features`、`Bugfixes`、`Enhancements`、`Other` 分类；同一分类内按功能主题排序，并且仅合并重复描述的同一 PR； + 同一 GitHub ID 贡献了不同功能或修复时，必须保留为多条独立条目，不能按贡献者去重； +- Contributors 仅致谢所有贡献者，不重复罗列 GitHub ID； +- 无法确认 PR 含义时标记待确认，不得臆测。 + +### 微信公众号（中文） + +微信公众号稿沿用 GitHub Release 的分类和条目粒度：先用一段发布摘要说明本次版本主题，再按“新功能”、 +“Bug 修复”、“功能增强”和“其他”列出变更，最后统一致谢。每项保留 `by @GitHub-ID (#PR-id)`，不写显式 PR URL； +公众号不需要逐条展开链接。参考结构： + +```markdown +# Apache bRPC {VERSION} 版本已发布 + +很高兴地通知大家，Apache bRPC {VERSION} 版本已发布，这是一次{版本定位}的版本，在{改进方向一}、 +{改进方向二}与{改进方向三}方面均有显著改进。本次发布{总体价值总结}。 + +Apache bRPC 官网：https://brpc.apache.org +下载链接：https://brpc.apache.org/zh/download/ +GitHub Release Tag：https://github.com/apache/brpc/releases/tag/{VERSION} + +## {VERSION} 版本变更 + +### 新功能 + +- {新功能描述} by @github-id (#1234) +- {新功能描述} by @github-id (#1235) + +### Bug 修复 + +- {修复描述} by @github-id (#1236) +- {修复描述} by @github-id (#1237) + +### 功能增强 + +- {增强描述} by @github-id (#1238) + +### 其他 + +- {文档、测试、构建或维护类描述} by @github-id (#1239) + +感谢所有关心和为 Apache bRPC 做出贡献的开发者！ +``` + +写作要求： + +- 先阅读对应 PR 的标题、说明、讨论及必要代码上下文，理解变更动机、用户价值和影响范围后再写； +- 保持与 GitHub Release 一致的四类分类和 PR 条目粒度；同一 GitHub ID 的不同功能或修复仍单独成条； +- 中文条目以自然、准确的中文重新组织表达，不逐字直译、不照搬生硬的提交标题； +- 每项保留 `by @GitHub-ID (#PR-id)`，不写 PR URL； +- 保留协议名、API、类名和配置项等必要技术术语；无法确认含义时标记待确认，不得臆测。 diff --git a/community/skills/brpc-release/references/state_template.md b/community/skills/brpc-release/references/state_template.md new file mode 100644 index 0000000000..f0ca687ce1 --- /dev/null +++ b/community/skills/brpc-release/references/state_template.md @@ -0,0 +1,73 @@ +# bRPC {VERSION} 发版状态 + +复制到 `~/brpc_release/{VERSION}/STATE.md` 后逐项维护。每完成一步就回写—— +发版横跨数周，这个文件是跨会话唯一的事实来源。 + +``` +Release Manager : {RM} +Apache ID : {apache-id} +版本号 : {VERSION} +发版分支 : release-{MAJOR.MINOR} +起始日期 : {YYYY-MM-DD} +上一个版本 : {PREV_VERSION} +``` + +## 关键产物 + +发出去之后再改代价很大，确定一个记一个。 + +``` +Tag commit id : +Tarball sha512 : +GPG key id : +dist/dev URL : https://dist.apache.org/repos/dist/dev/brpc/{VERSION}/ +投票邮件链接 : +投票开始时间 : +``` + +## 当前状态 + +``` +最后更新 : {YYYY-MM-DD HH:MM TZ} +最近完成 : +当前阻塞项 : 无 +建议下一步 : +``` + +每完成一个步骤立即更新本节；“建议下一步”可列多项，并标注是否可并行。用户选择暂不处理的工作应保留在下方进度清单中，不能删除。 + +## 进度 + +- [ ] **1. 发布范围** — 待发 PR 已合入，候选 commit 范围已冻结 +- [ ] **2. GPG** — 密钥可用，公钥已在 KEYS 中 +- [ ] **3. 分支与版本号** — `bump_version.sh --check` 通过，NOTICE 年份已确认 +- [ ] **4. 打包** — 本地 tag 已建、tarball/asc/sha512 已生成并自检通过 +- [ ] **4b. 推 tag** — `git push origin --tags`（RM 执行） +- [ ] **5. 上传 SVN** — 三个文件已 `svn commit` 到 dist/dev（RM 执行） +- [ ] **6. 验包** — `verify_package.sh` 全绿，与 GitHub tag diff 无差异 +- [ ] **7. Release Notes** — 最终 tag 范围的英文完整版和微信公众号中文版已定稿 +- [ ] **7a. 发起投票** — 邮件已发到 dev@（RM 执行） +- [ ] **7b. 投票通过** — ≥72h 且 ≥3 张 PMC binding +1，RESULT 邮件已发 +- [ ] **8a. 移到 dist/release** — `svn mv`（需 PMC 成员执行） +- [ ] **8b. GitHub Release** — 标题为 `Apache bRPC {VERSION}` +- [ ] **8c. 更新官网下载页** — brpc-website 仓库，中英文都改 +- [ ] **8d. ANNOUNCE 邮件** — 发到 dev@ 和 announce@（RM 执行，纯文本） +- [ ] **9. 收尾** — release 分支合回 master，更新 release_schedule.md 实际日期；已向 RM 提示可选的本地产物清理 + +## 已完成记录 + +每次完成步骤后追加一行，记录关键产物或 RM 对外操作的确认结果。 + +| 时间 | 步骤 | 结果/产物 | 执行人 | +|---|---|---|---| +| | | | | + +## 投票记录 + +| 投票人 | Binding | 票 | 备注 | +|---|---|---|---| +| | | | | + +## 遇到的问题 + + diff --git a/community/skills/brpc-release/scripts/bump_version.sh b/community/skills/brpc-release/scripts/bump_version.sh new file mode 100755 index 0000000000..125f962ec0 --- /dev/null +++ b/community/skills/brpc-release/scripts/bump_version.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Bump every hard-coded bRPC version string, then verify none was missed. +# +# Usage: +# bump_version.sh # rewrite all version files +# bump_version.sh --check # verify only, never write +# +# Idempotent: safe to re-run. Exits non-zero if any file ends up inconsistent. + +set -euo pipefail + +NEW_VERSION=${1:-} +MODE=${2:-write} + +if [[ -z ${NEW_VERSION} || ${NEW_VERSION} == -* ]]; then + echo "usage: $(basename "$0") [--check]" >&2 + exit 2 +fi + +if ! [[ ${NEW_VERSION} =~ ^[0-9]+\.[0-9]+\.[0-9]+(-rc[0-9]+)?$ ]]; then + echo "error: '${NEW_VERSION}' is not a valid version (expect 1.18.0 or 1.18.0-rc01)" >&2 + exit 2 +fi + +case ${MODE} in + write|--check) ;; + *) echo "error: unknown option '${MODE}'" >&2; exit 2 ;; +esac + +ROOT=$(git rev-parse --show-toplevel) +cd "${ROOT}" + +if [[ ! -f RELEASE_VERSION || ! -f CMakeLists.txt || ! -d src/brpc ]]; then + echo "error: '${ROOT}' does not look like the Apache bRPC repository" >&2 + exit 1 +fi + +if ! git remote -v | grep -Eq '(^|[/:])([^/]+/)?brpc(\.git)?([[:space:]]|$)'; then + echo "error: no recognizable bRPC git remote is configured" >&2 + exit 1 +fi + +# Never rewrite version files before switching to the release branch. This is +# deliberately enforced here rather than relying only on the release guide. +if [[ ${MODE} == write ]]; then + EXPECTED_BRANCH="release-${NEW_VERSION%.*}" + CURRENT_BRANCH=$(git branch --show-current) + if [[ ${CURRENT_BRANCH} != "${EXPECTED_BRANCH}" ]]; then + echo "error: version files may only be rewritten on '${EXPECTED_BRANCH}' (current: '${CURRENT_BRANCH:-detached HEAD}')" >&2 + echo "create or switch to the release branch first" >&2 + exit 1 + fi + + # Refuse to mix a version bump with unrelated edits. Untracked files are + # allowed because release tooling itself may not have been committed yet; + # tracked modifications are never overwritten or stashed automatically. + TRACKED_CHANGES=$(git status --short --untracked-files=no) + if [[ -n ${TRACKED_CHANGES} ]]; then + echo "error: tracked files have uncommitted changes; commit or restore them before bumping the version:" >&2 + printf '%s\n' "${TRACKED_CHANGES}" >&2 + exit 1 + fi +fi + +OLD_VERSION=$(tr -d '[:space:]' < RELEASE_VERSION) +# Escape dots so the string is safe inside a regex. +OLD_RE=${OLD_VERSION//./\\.} +NEW_RE=${NEW_VERSION//./\\.} + +failed=0 +changed=0 + +note() { printf ' %s\n' "$*"; } +fail() { printf ' FAIL %s\n' "$*" >&2; failed=1; } + +# Rewrite one file in place, portably (GNU and BSD sed disagree about -i). +rewrite() { + local file=$1 expr=$2 tmp + [[ -f ${file} ]] || { fail "${file}: not found"; return; } + tmp=$(mktemp) + sed -E "${expr}" "${file}" > "${tmp}" + if cmp -s "${file}" "${tmp}"; then + rm -f "${tmp}" + else + cat "${tmp}" > "${file}" # preserve the original file mode + rm -f "${tmp}" + changed=$((changed + 1)) + fi +} + +# Assert the file now carries the new version. +expect() { + local file=$1 pattern=$2 + if grep -Eq -- "${pattern}" "${file}"; then + note "ok ${file}" + else + fail "${file}: no line matching /${pattern}/" + fi +} + +if [[ ${MODE} == write ]]; then + echo "Bumping ${OLD_VERSION} -> ${NEW_VERSION}" + + printf '%s\n' "${NEW_VERSION}" > RELEASE_VERSION + + rewrite CMakeLists.txt \ + "s/^set\(BRPC_VERSION[[:space:]]+[^)]*\)/set(BRPC_VERSION ${NEW_VERSION})/" + + rewrite package/rpm/brpc.spec \ + "s/^(Version:[[:space:]]*).*/\1${NEW_VERSION}/" + + rewrite CLAUDE.md \ + "s/(Current version:[[:space:]]*)[^.[:space:]]+(\.[^.[:space:]]+){2}/\1${NEW_VERSION}/" + + # Only the root module() block is the released module version. The example + # module deliberately keeps its own version and dependency declaration. + rewrite MODULE.bazel \ + "s/^([[:space:]]*version = ')[^']*(')/\1${NEW_VERSION}\2/" + + for doc in docs/cn/bazel_support.md docs/en/bazel_support.md; do + rewrite "${doc}" \ + "s/(bazel_dep\(name = \"brpc\", version = \")[^\"]*/\1${NEW_VERSION}/" + done + + echo "Rewrote ${changed} file(s)" +else + echo "Checking version files against ${NEW_VERSION}" +fi + +echo "Verifying:" +expect RELEASE_VERSION "^${NEW_RE}$" +expect CMakeLists.txt "^set\(BRPC_VERSION ${NEW_RE}\)" +expect package/rpm/brpc.spec "^Version:[[:space:]]*${NEW_RE}$" +expect CLAUDE.md "Current version:[[:space:]]*${NEW_RE}\." +expect MODULE.bazel "^[[:space:]]*version = '${NEW_RE}'," +expect docs/cn/bazel_support.md "bazel_dep\(name = \"brpc\", version = \"${NEW_RE}\"" +expect docs/en/bazel_support.md "bazel_dep\(name = \"brpc\", version = \"${NEW_RE}\"" + +# Catch anything the table above does not know about. +if [[ ${OLD_VERSION} != "${NEW_VERSION}" ]]; then + echo "Scanning for leftover '${OLD_VERSION}':" + leftover=$(git grep -n -E -- "${OLD_RE}" -- \ + ':!community/release_schedule.md' ':!community/skills' ':!registry' \ + ':!example/build_with_bazel_module/MODULE.bazel' || true) + if [[ -n ${leftover} ]]; then + printf '%s\n' "${leftover}" | sed 's/^/ /' + echo " ^ review these by hand; release_schedule.md and past release notes are expected to keep old versions" >&2 + else + note "none" + fi +fi + +# NOTICE year is easy to forget in January releases. +notice_year=$(sed -n 's/^Copyright [0-9]*-\([0-9]*\).*/\1/p' NOTICE | head -1) +this_year=$(date +%Y) +if [[ ${notice_year} != "${this_year}" ]]; then + echo "warning: NOTICE says 'Copyright ...-${notice_year}' but it is ${this_year}; update it if this release ships this year" >&2 +fi + +if (( failed )); then + echo "FAILED: version files are inconsistent" >&2 + exit 1 +fi + +echo "All version files are at ${NEW_VERSION}" diff --git a/community/skills/brpc-release/scripts/make_package.sh b/community/skills/brpc-release/scripts/make_package.sh new file mode 100755 index 0000000000..475618cd38 --- /dev/null +++ b/community/skills/brpc-release/scripts/make_package.sh @@ -0,0 +1,177 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Build and sign an Apache bRPC source release in one RM-operated command: +# local tag -> tarball -> GPG signature -> sha512 -> self-check. +# +# Usage: +# BRPCUSERNAME= make_package.sh [options] +# +# Options: +# --user signing identity (default: $BRPCUSERNAME) +# --repo release-branch checkout (default: current git root) +# --outdir artifact directory (default: ~/brpc_release/) +# +# The AI model must not execute this script because it invokes interactive GPG +# signing. The RM runs it in a terminal and enters the private-key passphrase. +# This script NEVER pushes a tag, commits SVN changes, or changes Git +# user.name/user.email. + +set -euo pipefail + +VERSION="" +APACHE_ID=${BRPCUSERNAME:-} +REPO="" +OUTDIR="" + +while (( $# )); do + case $1 in + --user) APACHE_ID=$2; shift 2 ;; + --repo) REPO=$2; shift 2 ;; + --outdir) OUTDIR=$2; shift 2 ;; + -h|--help) sed -n '20,35p' "$0"; exit 0 ;; + -*) echo "error: unknown option '$1'" >&2; exit 2 ;; + *) VERSION=$1; shift ;; + esac +done + +if [[ -z ${VERSION} ]]; then + echo "usage: BRPCUSERNAME= $(basename "$0") [--repo dir] [--outdir dir]" >&2 + exit 2 +fi +if ! [[ ${VERSION} =~ ^[0-9]+\.[0-9]+\.[0-9]+(-rc[0-9]+)?$ ]]; then + echo "error: '${VERSION}' is not a valid version" >&2 + exit 2 +fi +if [[ -z ${APACHE_ID} ]]; then + echo "error: no Apache ID; set BRPCUSERNAME or pass --user " >&2 + exit 2 +fi +if ! [[ ${APACHE_ID} =~ ^[A-Za-z0-9._-]+$ ]]; then + echo "error: invalid Apache ID '${APACHE_ID}'" >&2 + exit 2 +fi + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO=${REPO:-$(git rev-parse --show-toplevel)} +OUTDIR=${OUTDIR:-${HOME}/brpc_release/${VERSION}} +PREFIX="apache-brpc-${VERSION}-src" +TARBALL="${PREFIX}.tar.gz" + +# GNU coreutils is not a given on macOS. +if command -v sha512sum >/dev/null 2>&1; then SHA512=(sha512sum) +elif command -v gsha512sum >/dev/null 2>&1; then SHA512=(gsha512sum) +elif command -v shasum >/dev/null 2>&1; then SHA512=(shasum -a 512) +else echo "error: need one of sha512sum, gsha512sum, shasum" >&2; exit 1 +fi + +step() { printf '\n==> %s\n' "$*"; } + +cd "${REPO}" + +step "Checking the working tree" +echo " repo : ${REPO}" +echo " branch : $(git rev-parse --abbrev-ref HEAD)" +echo " head : $(git rev-parse --short HEAD)" +echo " outdir : ${OUTDIR}" + +if [[ -n $(git status --porcelain) ]]; then + echo "error: working tree is dirty; commit or restore changes first" >&2 + git status --short >&2 + exit 1 +fi + +branch=$(git rev-parse --abbrev-ref HEAD) +expected_branch="release-${VERSION%.*}" +if [[ ${branch} != "${expected_branch}" ]]; then + echo "error: on '${branch}', expected '${expected_branch}'" >&2 + exit 1 +fi + +# Do not change Git identity. An existing identity is needed only because the +# local release tag is annotated; the RM owns repository configuration. +if ! git config user.email >/dev/null || ! git config user.name >/dev/null; then + echo "error: git identity is not configured; the RM must resolve this if desired" >&2 + echo " this script will not modify git user.name or user.email" >&2 + exit 1 +fi + +step "Checking version files" +"${SCRIPT_DIR}/bump_version.sh" "${VERSION}" --check + +step "Tagging ${VERSION} (local only)" +if git rev-parse -q --verify "refs/tags/${VERSION}" >/dev/null; then + tagged=$(git rev-parse "${VERSION}^{commit}") + head=$(git rev-parse HEAD) + if [[ ${tagged} != "${head}" ]]; then + echo "error: tag ${VERSION} already exists at ${tagged:0:12} but HEAD is ${head:0:12}" >&2 + echo " if the tag was never pushed, the RM may drop it with: git tag -d ${VERSION}" >&2 + exit 1 + fi + echo " tag already exists at HEAD, reusing" +else + git tag -a "${VERSION}" -m "release ${VERSION}" + echo " created" +fi +COMMIT=$(git rev-parse "${VERSION}^{commit}") + +step "Building ${TARBALL}" +mkdir -p "${OUTDIR}" +git archive --format=tar.gz "${VERSION}" \ + --prefix="${PREFIX}/" --output="${OUTDIR}/${TARBALL}" +echo " $(cd "${OUTDIR}" && du -h "${TARBALL}" | cut -f1)" + +cd "${OUTDIR}" + +step "Signing as ${APACHE_ID}@apache.org" +echo " GPG will ask the RM for the private-key passphrase" +rm -f "${TARBALL}.asc" +gpg -u "${APACHE_ID}@apache.org" --armor \ + --output "${TARBALL}.asc" --detach-sign "${TARBALL}" + +step "Verifying GPG signature" +gpg --verify "${TARBALL}.asc" "${TARBALL}" + +# Run from the tarball directory so the checksum records a bare filename. +step "Generating ${TARBALL}.sha512" +"${SHA512[@]}" "${TARBALL}" > "${TARBALL}.sha512" + +step "Checking SHA512" +"${SHA512[@]}" --check "${TARBALL}.sha512" + +cat < Artifacts in ${OUTDIR} +$(ls -1 "${PREFIX}"* | sed 's/^/ /') + + Release Commit ID: ${COMMIT} + (this is what goes in the [VOTE] mail, not the branch HEAD) + +==> Next, run these yourself -- this script will not: + + cd "${REPO}" && git push origin --tags + + mkdir -p ~/brpc_release/svn/dev/brpc/${VERSION} + cp "${OUTDIR}/${TARBALL}"{,.asc,.sha512} ~/brpc_release/svn/dev/brpc/${VERSION}/ + cd ~/brpc_release/svn/dev/brpc && svn add --force . \ + && svn --username="${APACHE_ID}" commit -m "release ${VERSION}" + + Then verify the uploaded candidate: + "${SCRIPT_DIR}/verify_package.sh" https://dist.apache.org/repos/dist/dev/brpc/${VERSION}/ +EOF diff --git a/community/skills/brpc-release/scripts/release_notes.sh b/community/skills/brpc-release/scripts/release_notes.sh new file mode 100755 index 0000000000..aca3f1535c --- /dev/null +++ b/community/skills/brpc-release/scripts/release_notes.sh @@ -0,0 +1,240 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Draft Release Notes from GitHub pull requests represented by commits between +# two refs. One merged PR produces at most one entry, even when it has many +# commits. The script fetches the PR title and author through `gh api`. +# +# Usage: +# release_notes.sh [--repo ] +# +# Example: is the local release tag, which need not be pushed to GitHub. +# release_notes.sh "$PREV_VERSION" "$VERSION" > "$HOME/brpc_release/$VERSION/notes-draft.md" +# +# `gh auth login` must be completed before running this script. Buckets are +# title-keyword guesses only; review every PR and reword it for users. + +set -euo pipefail + +FROM="" +TO="" +REPO=apache/brpc + +while (( $# )); do + case $1 in + --repo) + [[ $# -ge 2 ]] || { echo "error: --repo requires " >&2; exit 2; } + REPO=$2 + shift 2 + ;; + -h|--help) + sed -n '20,31p' "$0" + exit 0 + ;; + -*) + echo "error: unknown option '$1'" >&2 + exit 2 + ;; + *) + if [[ -z ${FROM} ]]; then + FROM=$1 + elif [[ -z ${TO} ]]; then + TO=$1 + else + echo "error: too many positional arguments" >&2 + exit 2 + fi + shift + ;; + esac +done + +if [[ -z ${FROM} || -z ${TO} ]]; then + echo "usage: $(basename "$0") [--repo ]" >&2 + exit 2 +fi +if ! [[ ${REPO} =~ ^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$ ]]; then + echo "error: '${REPO}' is not a valid GitHub owner/repository" >&2 + exit 2 +fi +if ! command -v gh >/dev/null 2>&1; then + echo "error: GitHub CLI (gh) is required to generate PR-based Release Notes" >&2 + exit 1 +fi +if ! gh auth status --hostname github.com >/dev/null 2>&1; then + echo "error: authenticate GitHub CLI first: gh auth login --hostname github.com" >&2 + exit 1 +fi + +# The local release tag is the endpoint. It need not have been pushed to GitHub: +# GitHub is consulted only for PR metadata, while the commit range is local. +for ref in "${FROM}" "${TO}"; do + if ! git rev-parse -q --verify "${ref}^{commit}" >/dev/null; then + echo "error: '${ref}' is not a commit. Fetch tags first: git fetch --tags" >&2 + exit 1 + fi +done + +RANGE="${FROM}..${TO}" +commit_total=$(git rev-list --count "${RANGE}") +if (( commit_total == 0 )); then + echo "error: no commits in ${RANGE}" >&2 + exit 1 +fi + +feature=$(mktemp) +bugfix=$(mktemp) +enhance=$(mktemp) +other=$(mktemp) +unassociated=$(mktemp) +trap 'rm -f "${feature}" "${bugfix}" "${enhance}" "${other}" "${unassociated}"' EXIT + +# Extract PR numbers from commit subjects, then retrieve all closed PR metadata +# with GitHub REST pagination and select the required merged PRs locally. This +# avoids one request per commit and remains robust when a subject contains an +# upstream issue number that is not a bRPC PR. GitHub-generated merge/squash/ +# rebase subjects preserve the final merged PR number. Any commit without one is +# listed separately for human review and is never made into a commit-level note. +# Do not use Bash associative arrays: macOS still ships Bash 3.2. +prs=$(mktemp) +all_prs=$(mktemp) +pr_data_file=$(mktemp) +trap 'rm -f "${feature}" "${bugfix}" "${enhance}" "${other}" "${unassociated}" "${prs}" "${all_prs}" "${pr_data_file}"' EXIT +while IFS=$'\t' read -r sha subject; do + # Example: "... (#2793) (#3491)". The final #number is the bRPC PR. + if [[ ${subject} =~ .*\#([0-9]+) ]]; then + printf '%s\n' "${BASH_REMATCH[1]}" >> "${prs}" + else + printf '%s %s\n' "${sha:0:12}" "${subject}" >> "${unassociated}" + fi +done < <(git log --format='%H%x09%s' "${RANGE}") + +sort -u "${prs}" -o "${prs}" +pr_total=$(wc -l < "${prs}" | tr -d ' ') +if (( pr_total == 0 )); then + echo "error: no GitHub PR references found in ${RANGE}; Release Notes must be PR-based" >&2 + exit 1 +fi + +# GitHub returns up to 100 closed PRs per page. This is dozens of requests for +# the whole repository rather than hundreds of commit/PR requests per release. +echo "Loading merged PR metadata in pages from ${REPO}..." >&2 +if ! gh api --paginate "repos/${REPO}/pulls?state=closed&per_page=100" \ + --jq '.[] | select(.merged_at != null) | [.number, .title, .user.login, .merged_at] | @tsv' \ + > "${all_prs}"; then + echo "error: failed to retrieve merged PR metadata from ${REPO}" >&2 + exit 1 +fi + +# Both files use lexicographic PR-number ordering required by join, so it can +# select only the PRs in the release range. If a subject's final #number is an +# external issue rather than a bRPC PR, leave it in the review comment instead +# of failing the entire draft. +sort -k1,1 "${all_prs}" -o "${all_prs}" +join -t $'\t' -1 1 -2 1 "${prs}" "${all_prs}" > "${pr_data_file}" +metadata_total=$(grep -cve '^$' "${pr_data_file}" || true) +if (( metadata_total != pr_total )); then + missing_prs=$(mktemp) + trap 'rm -f "${feature}" "${bugfix}" "${enhance}" "${other}" "${unassociated}" "${prs}" "${all_prs}" "${pr_data_file}" "${missing_prs}"' EXIT + join -t $'\t' -v 1 -1 1 -2 1 "${prs}" "${all_prs}" > "${missing_prs}" + while IFS= read -r missing_pr; do + printf 'unresolved #PR reference: %s\n' "${missing_pr}" >> "${unassociated}" + done < "${missing_prs}" + echo "warning: retrieved ${metadata_total} PR record(s), expected ${pr_total}; unresolved references were left for review" >&2 + pr_total=${metadata_total} +fi +if (( pr_total == 0 )); then + echo "error: no merged bRPC PR metadata found for ${RANGE}" >&2 + exit 1 +fi + +while IFS=$'\t' read -r number title author merged_at; do + if [[ -z ${merged_at} || ${merged_at} == null ]]; then + echo "error: #${number} is not a merged PR; refusing to include it" >&2 + exit 1 + fi + + lower=$(printf '%s' "${title}" | tr '[:upper:]' '[:lower:]') + # GitHub Release automatically recognizes @user and #PR-id; do not emit + # explicit Markdown links. Keep one independent entry per PR: the same + # author may contribute multiple unrelated features or fixes. + entry="- ${title} by @${author} (#${number})" + case ${lower} in + fix*|*"fix "*|*fixes*|bug*|*"bugfix"*|revert*) + printf '%s\n' "${entry}" >> "${bugfix}" + ;; + feat*|add\ *|"support "*|*"add support"*|"new "*|"implement "*|"introduce "*) + printf '%s\n' "${entry}" >> "${feature}" + ;; + perf*|refactor*|"improve "*|"optimize "*|"speed up"*|"reduce "*|"enhance "*|"clean"*|"remove "*|"update "*|"upgrade "*) + printf '%s\n' "${entry}" >> "${enhance}" + ;; + doc*|test*|ci*|build*|chore*|style*) + printf '%s\n' "${entry}" >> "${other}" + ;; + *) + printf '%s\n' "${entry}" >> "${enhance}" + ;; + esac +done < "${pr_data_file}" + +emit() { + local title=$1 file=$2 + printf '\n%s:\n' "${title}" + if [[ -s ${file} ]]; then + cat "${file}" + else + printf -- '- (none)\n' + fi +} + +cat < + +[Release Notes] +EOF + +emit "Feature" "${feature}" +emit "Bugfix" "${bugfix}" +emit "Enhancement" "${enhance}" +emit "Other" "${other}" + +cat < +EOF diff --git a/community/skills/brpc-release/scripts/verify_package.sh b/community/skills/brpc-release/scripts/verify_package.sh new file mode 100755 index 0000000000..10cfb379ed --- /dev/null +++ b/community/skills/brpc-release/scripts/verify_package.sh @@ -0,0 +1,253 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# Verify an Apache bRPC release candidate and print a summary that can be +# pasted straight into a vote reply. +# +# Usage: +# verify_package.sh [--keys ] [--no-diff] +# +# Examples: +# verify_package.sh https://dist.apache.org/repos/dist/dev/brpc/1.18.0/ +# verify_package.sh ~/brpc_release/1.18.0 +# +# A portable companion to community/apache-package-validator.sh (which needs +# wget and GNU coreutils); this one runs on macOS too and additionally diffs +# the tarball against the GitHub tag. + +set -uo pipefail + +SOURCE="" +KEYS_URL="https://downloads.apache.org/brpc/KEYS" +DO_DIFF=1 + +while (( $# )); do + case $1 in + --keys) KEYS_URL=$2; shift 2 ;; + --no-diff) DO_DIFF=0; shift ;; + -h|--help) sed -n '20,33p' "$0"; exit 0 ;; + -*) echo "error: unknown option '$1'" >&2; exit 2 ;; + *) SOURCE=$1; shift ;; + esac +done + +if [[ -z ${SOURCE} ]]; then + echo "usage: $(basename "$0") [--keys ] [--no-diff]" >&2 + exit 2 +fi + +if command -v sha512sum >/dev/null 2>&1; then SHA512=(sha512sum) +elif command -v gsha512sum >/dev/null 2>&1; then SHA512=(gsha512sum) +elif command -v shasum >/dev/null 2>&1; then SHA512=(shasum -a 512) +else echo "error: need one of sha512sum, gsha512sum, shasum" >&2; exit 1 +fi + +r_link=' '; r_sum=' '; r_sig=' '; r_ver=' '; r_lic=' '; r_bin=' '; r_tag=' ' +failures=0 + +pass() { printf ' PASS %s\n' "$*"; } +fail() { printf ' FAIL %s\n' "$*" >&2; failures=$((failures + 1)); } +warn() { printf ' WARN %s\n' "$*" >&2; } +step() { printf '\n==> %s\n' "$*"; } + +WORK=$(mktemp -d) +trap 'summary; rm -rf "${WORK}"' EXIT + +summary() { + cat <&2 + fi +} + +# ---------------------------------------------------------------- fetch + +step "Collecting artifacts" +if [[ -d ${SOURCE} ]]; then + found=$(ls "${SOURCE}"/apache-brpc-*-src.tar.gz 2>/dev/null | head -1) + if [[ -z ${found} ]]; then + fail "no apache-brpc-*-src.tar.gz in ${SOURCE}" + exit 1 + fi + TARBALL=$(basename "${found}") + VERSION=${TARBALL#apache-brpc-}; VERSION=${VERSION%-src.tar.gz} + for suffix in "" .asc .sha512; do + if [[ -f ${SOURCE}/${TARBALL}${suffix} ]]; then + cp "${SOURCE}/${TARBALL}${suffix}" "${WORK}/" + else + fail "missing ${TARBALL}${suffix}" + fi + done + (( failures == 0 )) && r_link='x' + echo " local: ${SOURCE} (version ${VERSION})" +else + base=${SOURCE%/} + VERSION=${base##*/} + TARBALL="apache-brpc-${VERSION}-src.tar.gz" + ok=1 + echo " from ${base}" + for suffix in "" .asc .sha512; do + # dist.apache.org has no CDN in front of it and routinely crawls along + # at a few KB/s, so show a progress bar for the tarball rather than + # leaving the RM staring at a silent terminal for ten minutes. + if [[ -z ${suffix} ]]; then + curl_opts=(--progress-bar) + echo " downloading ${TARBALL} (dist.apache.org can be very slow)" + else + curl_opts=(-sS) + fi + if curl -fL "${curl_opts[@]}" --connect-timeout 20 --max-time 1800 \ + -o "${WORK}/${TARBALL}${suffix}" "${base}/${TARBALL}${suffix}"; then + echo " got ${TARBALL}${suffix}" + else + fail "cannot download ${base}/${TARBALL}${suffix}" + ok=0 + fi + done + (( ok )) && r_link='x' +fi + +cd "${WORK}" || exit 1 +[[ -f ${TARBALL} ]] || { fail "no tarball to verify"; exit 1; } + +# ---------------------------------------------------------------- checksum + +step "sha512" +if [[ -f ${TARBALL}.sha512 ]]; then + if grep -q '/' "${TARBALL}.sha512"; then + warn "${TARBALL}.sha512 records a path, not a bare filename; --check may fail for others" + fi + if "${SHA512[@]}" --check "${TARBALL}.sha512"; then + r_sum='x'; pass "checksum matches" + else + fail "checksum mismatch" + fi +else + fail "no .sha512 file" +fi + +# ---------------------------------------------------------------- signature + +step "GPG signature" +if [[ -f ${TARBALL}.asc ]]; then + if ! curl -fsSL "${KEYS_URL}" | gpg --import 2>/dev/null; then + warn "could not import KEYS from ${KEYS_URL}" + fi + # A candidate still in dist/dev may only have its key in the dev KEYS file. + if [[ ${KEYS_URL} == *"/release/"* || ${KEYS_URL} == *"downloads.apache.org"* ]]; then + curl -fsSL "https://dist.apache.org/repos/dist/dev/brpc/KEYS" \ + | gpg --import 2>/dev/null || true + fi + if gpg --verify "${TARBALL}.asc" "${TARBALL}"; then + r_sig='x'; pass "signature is valid" + echo " (a 'no ultimately trusted keys' warning is expected and fine --" + echo " it means the key is not in your web of trust, not that the signature is bad)" + else + fail "signature does not verify" + fi +else + fail "no .asc file" +fi + +# ---------------------------------------------------------------- contents + +step "Unpacking" +tar -xzf "${TARBALL}" || { fail "cannot unpack ${TARBALL}"; exit 1; } +SRC="apache-brpc-${VERSION}-src" +[[ -d ${SRC} ]] || { fail "tarball does not contain ${SRC}/"; exit 1; } + +step "Version consistency" +in_release=$(tr -d '[:space:]' < "${SRC}/RELEASE_VERSION" 2>/dev/null) +if [[ ${in_release} == "${VERSION}" ]] \ + && grep -q "set(BRPC_VERSION ${VERSION})" "${SRC}/CMakeLists.txt" 2>/dev/null; then + r_ver='x'; pass "RELEASE_VERSION and CMakeLists.txt both say ${VERSION}" +else + fail "version mismatch (RELEASE_VERSION='${in_release}', expected '${VERSION}')" +fi + +step "LICENSE and NOTICE" +if [[ -f ${SRC}/LICENSE && -f ${SRC}/NOTICE ]]; then + r_lic='x'; pass "both present" + year=$(sed -n 's/^Copyright [0-9]*-\([0-9]*\).*/\1/p' "${SRC}/NOTICE" | head -1) + [[ ${year} == "$(date +%Y)" ]] || warn "NOTICE copyright ends at ${year}, current year is $(date +%Y)" +else + fail "LICENSE and/or NOTICE missing" +fi + +step "Unexpected binaries" +# `certificate` covers test cert/key fixtures whose `file` output varies by +# platform (PEM on macOS, sometimes a bare "certificate"/DER blob elsewhere). +nontext=$(find "${SRC}" -type f -print0 \ + | xargs -0 file \ + | grep -v 'GIF\|JPEG\|PNG\|SVG\|PowerPoint\|Git\|JSON\|PEM\|certificate\|empty\|text\|XML' || true) + +# Fuzzing seed corpora are opaque by design and are allowlisted in +# .licenserc.yaml; flagging them sends the RM chasing a phantom -1. +expected=$(printf '%s\n' "${nontext}" | grep 'test/fuzzing/fuzz_[^/]*_seed_corpus/' || true) +suspicious=$(printf '%s\n' "${nontext}" | grep -v 'test/fuzzing/fuzz_[^/]*_seed_corpus/' | grep -v '^$' || true) + +if [[ -n ${expected} ]]; then + note_count=$(printf '%s\n' "${expected}" | grep -c . || true) + echo " ${note_count} fuzzing seed corpus file(s) are binary by design (allowlisted in .licenserc.yaml)" +fi +if [[ -z ${suspicious} ]]; then + r_bin='x'; pass "no compiled archives or stray binaries" +else + fail "suspicious files:" + printf '%s\n' "${suspicious}" | sed 's/^/ /' >&2 +fi + +# ---------------------------------------------------------------- vs tag + +if (( DO_DIFF )); then + step "Diff against the GitHub tag" + if curl -fsSL -o "tag-${VERSION}.tar.gz" \ + "https://github.com/apache/brpc/archive/refs/tags/${VERSION}.tar.gz"; then + tar -xzf "tag-${VERSION}.tar.gz" + if diff -r "brpc-${VERSION}" "${SRC}" > tag.diff 2>&1; then + r_tag='x'; pass "identical to the tag" + else + fail "differs from the tag:" + head -40 tag.diff | sed 's/^/ /' >&2 + fi + else + warn "tag ${VERSION} not published on GitHub yet; skipping" + fi +else + warn "tag diff skipped (--no-diff)" +fi + +exit $(( failures > 0 )) diff --git a/config_brpc.sh b/config_brpc.sh index 85692de3cb..0efa4d374d 100755 --- a/config_brpc.sh +++ b/config_brpc.sh @@ -54,10 +54,12 @@ 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,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"` +TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-urma,with-urma-mock,without-urma-mock,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"` WITH_GLOG=0 WITH_THRIFT=0 WITH_RDMA=0 +WITH_URMA=0 +URMA_MOCK_MODE=auto WITH_MESALINK=0 WITH_BTHREAD_TRACER=0 WITH_ASAN=0 @@ -90,6 +92,9 @@ 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-urma-mock) URMA_MOCK_MODE=on; shift 1 ;; + --without-urma-mock) URMA_MOCK_MODE=off; 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 ;; @@ -538,6 +543,42 @@ 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 urma_api.h) + if [ -z "$URMA_HDR" ]; then + >&2 $ECHO "Fail to find urma_api.h from --headers" + exit 1 + fi + 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 [ "$URMA_MOCK_MODE" = "on" ]; then + append_to_output "URMA_USE_MOCK=1" + print_info "URMA mock forced by --with-urma-mock" + elif [ "$URMA_MOCK_MODE" = "off" ]; then + if [ -z "$URMA_LIB" ]; then + >&2 $ECHO "--without-urma-mock requires liburma" + exit 1 + fi + append_to_output_libs "$URMA_LIB" + append_to_output "DYNAMIC_LINKINGS+=-lurma" + append_to_output "URMA_USE_MOCK=0" + print_info "URMA mock disabled by --without-urma-mock" + elif [ -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 @@ -674,6 +715,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/bazel_support.md b/docs/cn/bazel_support.md index 3a138cfbb7..459dd38ef0 100644 --- a/docs/cn/bazel_support.md +++ b/docs/cn/bazel_support.md @@ -27,7 +27,7 @@ module( ) bazel_dep(name = "protobuf", version = "27.3", repo_name = "com_google_protobuf") -bazel_dep(name = "brpc", version = "1.17.0", repo_name = "apache_brpc") +bazel_dep(name = "brpc", version = "1.18.0", repo_name = "apache_brpc") local_path_override( module_name = "brpc", diff --git a/docs/cn/getting_started.md b/docs/cn/getting_started.md index 51e9b9759e..b8fdb53675 100644 --- a/docs/cn/getting_started.md +++ b/docs/cn/getting_started.md @@ -283,6 +283,42 @@ Monterey中openssl的安装位置可能不再位于`/usr/local/opt/openssl`， * 先运行`brew link openssl --force`看看`/usr/local/opt/openssl`是否出现了 * 没有的话可以自行设置软链：`sudo ln -s /opt/homebrew/Cellar/openssl@3/3.0.3 /usr/local/opt/openssl`。请注意此命令中openssl的目录可能随环境变化而变化，可通过`brew info openssl`查看。 +### 使用 CMake 编译 Debug 版 brpc + +Apple Silicon 可以使用 Homebrew 安装的依赖编译 Debug 版本： + +```shell +cmake -S . -B build -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Debug \ + -DDEBUG=ON \ + -DCMAKE_OSX_ARCHITECTURES=arm64 \ + -DCMAKE_OSX_SYSROOT="$(xcrun --sdk macosx --show-sdk-path)" \ + -DCMAKE_PREFIX_PATH="$(brew --prefix)" \ + -DOPENSSL_ROOT_DIR="$(brew --prefix openssl@3)" +cmake --build build --parallel +``` + +该命令使用单配置的Unix Makefiles生成器，因此`CMAKE_BUILD_TYPE=Debug`会选择 +CMake的Debug构建配置。`DEBUG=ON`启用brpc的调试日志并保留断言。构建产物位于 +`build/output/`。 + +如需为 clangd 等工具生成编译数据库，请在配置命令中添加以下可选项： + +```shell +-DCMAKE_EXPORT_COMPILE_COMMANDS=ON +``` + +生成的编译数据库位于 `build/compile_commands.json`。 + +如果 CMake 报错 `tapi error: malformed file`，请确认 Command Line Tools 与 Xcode +版本一致，并选择当前安装的 Xcode： + +```shell +sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer +``` + +如果问题仍然存在，请检查`xcrun --sdk macosx --show-sdk-path`返回的SDK路径是否有效。 + ### 使用config_brpc.sh编译brpc git克隆brpc，进入到项目目录然后运行： ```shell diff --git a/docs/cn/handshake_common_design.md b/docs/cn/handshake_common_design.md index 169866dc6d..afc8e7a580 100644 --- a/docs/cn/handshake_common_design.md +++ b/docs/cn/handshake_common_design.md @@ -1,499 +1,50 @@ -# RDMA、URMA、UBSHM 公共握手设计（修订版） +# RDMA 与 UBSHM 公共握手 -## 1. 文档目的 +## 范围 -本文在现有公共 framing、`HandshakeSession` 和协议字段 adapter 的基础上，进一步明确四个职责边界： +本次重构将 RDMA 和 UBSHM 的握手流程收敛到公共层。`Socket` 持有 +`AdapterTransport`；后者以 TCP 作为握手控制通道，在升级成功后选择高速 +数据面，协商不可用时继续使用 TCP。URMA 目前仍使用独立的 `UrmaTransport` +和握手实现，不在本次迁移范围内。 -1. `Socket` 只感知一个顶层 `AdapterTransport`，不感知 TCP、RDMA、URMA、UBSHM，也不感知握手 phase。 -2. `AdapterTransport` 自动完成连接升级；升级成功或回退 TCP 后，统一通知 `Socket` 建链完成。 -3. 具体 Transport 只提供握手所需的资源操作接口；握手顺序、状态转换和 fallback 由上层统一编排。 -4. Transport 只负责数据传输和资源生命周期，不负责 TCP 建链、wire framing、握手状态机或握手流程编排。 +## 分工 -本文只定义架构和迁移方向，不改变 RDMA/URMA/UBSHM 当前 wire format。 +| 组件 | 职责 | +| --- | --- | +| `AdapterTransport` | 持有 TCP 与候选高速 Transport，启动客户端握手、分发服务端输入，并选择数据面 | +| `HandshakeSession` | 执行公共客户端和服务端状态机，管理帧收发、增量解析和终态发布 | +| `HandshakeProtocol` | 定义 hello、ACK 和扩展帧，编解码协议字段并保存协议状态 | +| `HandshakeTransport` | 准备和协商资源，处理升级成功、TCP 回退及失败清理 | +| RDMA / UBSHM protocol | 实现各自的 wire format、版本兼容和字段校验 | +| `RdmaEndpoint` / `UBShmEndpoint` | 管理各自的资源和数据面，不驱动完整握手流程 | -## 2. 当前问题 +客户端在 TCP 连接建立后启动握手 bthread；服务端通过 `InputMessenger` +增量解析 hello、扩展字段（若协议需要）和 ACK。`AdapterTransport` 选择具体的 +protocol 和 transport participant，`HandshakeSession` 只负责调用参与者并编排流程。 +服务端通过协议版本恢复选中的 protocol，因为 ACK 本身不带 magic。 -现有实现已经抽取了 `HandshakeSession`、`HandshakeCodec` 和公共 framing，但职责仍未完全收敛： +## 状态与回退 -- `ProcessHandshakeAtClient` 仍位于具体 Transport 或 endpoint 路径中，client 的握手入口不统一。 -- `RdmaTransport`、`UBShmTransport` 等仍通过各自的 `S_*` 常量暴露握手阶段；`HandshakePhases` 需要填入不同 transport 的 phase，公共 session 无法真正统一。 -- server handshake adapter 需要通过 `Socket` 找到 `AdapterTransport`，再找到具体 Transport，并且直接驱动资源创建、激活和 fallback。 -- `AdapterTransport` 虽然已经是 `Socket` 的顶层对象，但握手的 client/server 入口、连接完成通知和数据面切换仍分散在多个层次。 -- “握手成功”与“Socket 建链完成”不是同一个明确事件，导致 TCP fallback、升级成功和异常退出的发布顺序难以验证。 +握手终态为 `ESTABLISHED`、`FALLBACK_TCP` 或 `FAILED`。升级成功后，TCP +仍是控制连接，业务数据走 RDMA/UBSHM；控制连接上出现额外业务数据视为 +协议错误。资源不可用或协商拒绝时，释放本次升级资源，设置 TCP 为数据面， +再以 release 顺序发布 `FALLBACK_TCP`。事件线程以 acquire 顺序读取终态， +继续解析已缓存及后续的 TCP 业务数据。已确认属于握手协议的畸形帧则失败， +不作为普通 RPC 数据回放。 -根因是把“传输能力”和“建链流程”混在了一起。Transport 是被流程调用的参与者，不应成为流程的拥有者。 +增量解析中，不完整帧返回 `NOT_ENOUGH_DATA` 且不消费输入；不可能匹配 +握手 magic 的前缀返回 `TRY_OTHERS`。TCP 分片、握手帧与后续数据粘连都由 +帧边界处理，不能假设一次 read 恰好得到一帧。 -## 3. 目标架构 +## 协议兼容 -### 3.1 分层 +- RDMA 保持既有 v2/v3 wire format、协议 ID 和注册名称。 +- UBSHM 使用 v3 hello。64 字节 hello 后，双方交换 4 字节网络序格式扩展； + 只有双方选择 `LEGACY_64` 才能升级，否则回退 TCP。ACK 仍为 4 字节。 +- 公共 framing 仅负责帧边界，协议字段由各自 `HandshakeProtocol` 解释。 -```text -Socket - | - v -AdapterTransport Socket 唯一感知的 Transport - |-- TcpTransport TCP 控制面和 fallback 数据面 - |-- HighSpeedTransport RDMA / URMA / UBSHM 数据面 - |-- ConnectionUpgradeCoordinator 统一的建链与升级编排 - | |-- HandshakeSession 公共状态机、I/O、framing - | |-- HandshakeCodec 协议 wire 字段编解码 - | |-- TransportUpgradeOps 被调用的资源阶段接口 - | `-- SocketConnectionNotifier 建链完成/失败通知 - `-- ActiveTransport 根据统一状态选择数据面 -``` +## 验证重点 -这里的 `ConnectionUpgradeCoordinator` 可以先作为 `AdapterTransport` 的内部实现，不要求立即新增独立公开类；重要的是职责必须集中在该层，而不是分散到具体 Transport。 - -### 3.2 依赖方向 - -```text -Socket -> AdapterTransport -> TcpTransport / HighSpeedTransport -Socket -> AdapterTransport -> ConnectionUpgradeCoordinator -ConnectionUpgradeCoordinator -> TransportUpgradeOps -ConnectionUpgradeCoordinator -> HandshakeSession -TransportUpgradeOps -> transport-specific endpoint/resource -``` - -禁止以下反向依赖： - -- `Socket` 直接调用具体 Transport 或 handshake adapter。 -- `TcpTransport` 调用具体高速度 Transport 的握手函数。 -- `RdmaTransport`、`UrmaTransport`、`UBShmTransport` 调用 `ProcessHandshakeAtClient`、`RunServerHandshake` 等流程函数。 -- 公共 `HandshakeSession` 依赖 `ibv_*`、`urma_*`、UBRING 类型。 -- 具体 Transport 通过自定义 phase 影响 `Socket` 的建链判断。 - -## 4. 统一状态模型 - -### 4.1 握手状态由上层拥有 - -公共状态只描述连接升级流程，不描述某一种资源如何创建： - -```cpp -namespace brpc { -namespace handshake { - -enum class Phase { - kUninitialized, - kPreparing, - kHelloSending, - kHelloWaiting, - kNegotiating, - kAckSending, - kAckWaiting, - kEstablished, - kFallbackTcp, - kFailed, -}; - -enum class StepResult { - kOk, - kFallback, - kNeedMore, - kNotMine, - kError, -}; - -} // namespace handshake -} // namespace brpc -``` - -`Socket` 和 `AdapterTransport` 只读取 `Phase` 的终态：`kEstablished`、`kFallbackTcp`、`kFailed`。中间状态只由 coordinator/session 使用。 - -### 4.2 删除 `HandshakePhases` - -不再使用： - -```cpp -struct HandshakePhases { - int prepare_local; - int hello_send; - int hello_wait; - int negotiate; - int ack_send; - int ack_wait; -}; -``` - -原因是这些数字并不是协议字段，也不是数据面状态；它们只是历史实现中的日志/状态映射。公共 session 应使用固定的 `Phase`，具体 Transport 的资源子状态留在自身实现中，不向上暴露。 - -例如 RDMA 的 `S_ALLOC_QPCQ`、`S_BRINGUP_QP`、UBSHM 的 `S_ALLOC_SHM` 都只能是具体 Transport 的内部 debug 状态；它们不能再填入公共 `HandshakePhases`，也不能作为 `Socket` 判断建链完成的依据。 - -## 5. 模块职责 - -### 5.1 Socket - -`Socket` 只保存和调用一个 `Transport*`，实际对象始终为 `AdapterTransport`。它只关心以下事件： - -- TCP socket 是否连接成功； -- `AdapterTransport` 是否报告建链完成； -- 当前数据面读写是否可用； -- 连接是否失败或 EOF。 - -Socket 不直接读取 handshake phase，也不负责决定 TCP 还是高速度 Transport。 - -### 5.2 AdapterTransport - -`AdapterTransport` 是顶层 Transport 和连接升级协调器的宿主，负责： - -- 创建并持有 TCP Transport 和候选高速度 Transport； -- 建立 TCP 控制连接； -- 自动启动 client/server 的升级编排； -- 将 TCP 可读事件交给 coordinator，直到升级进入终态； -- 在 `kEstablished` 与 `kFallbackTcp` 之间选择 active data transport； -- 对升级结果执行一次性的连接完成通知； -- 保证 fallback 数据回放和状态发布的内存顺序。 - -建议的内部接口如下： - -```cpp -class AdapterTransport : public Transport { -public: - std::shared_ptr Connect() override; - void ProcessEvent(bthread_attr_t attr) override; - - // 由上层 Socket/连接流程使用，不暴露具体 Transport 类型。 - handshake::Phase connection_phase() const; - -private: - void StartClientUpgrade(); - void ProcessUpgradeReadable(); - void CompleteConnection(handshake::Phase terminal_phase); - void ActivateTcp(); - void ActivateHighSpeed(); - - handshake::HandshakeSession _handshake; - std::unique_ptr _tcp_transport; - std::unique_ptr _high_speed_transport; - std::unique_ptr _upgrade; -}; -``` - -`StartClientUpgrade` 和 `ProcessUpgradeReadable` 是上层编排入口。它们不能下沉到 `RdmaTransport`、`UBShmTransport` 或 endpoint。 - -### 5.3 ConnectionUpgradeCoordinator / HandshakeSession - -该模块拥有连接升级的完整流程： - -- 选择协议 codec； -- 通过 TCP 发送和接收 hello/ack； -- 增量 framing、半包和非本协议数据回放； -- 按固定顺序调用 Transport 能力接口； -- 将协议不支持、资源失败转换为 fallback； -- 发布 `kEstablished`、`kFallbackTcp` 或 `kFailed`； -- 在终态时通知 `AdapterTransport` 完成建链。 - -`HandshakeSession` 不知道 RDMA、URMA、UBSHM 的资源类型，只调用抽象的阶段接口。 - -### 5.4 具体 Transport - -具体 Transport 只负责： - -- 数据面读写、事件和 completion； -- 资源创建、导入、激活、停用和释放； -- 提供本端 hello 所需的只读能力/字段； -- 接收上层解析后的远端参数； -- 报告资源阶段成功、不可用或失败。 - -具体 Transport 不负责： - -- TCP fd 读写； -- magic/length framing； -- hello/ack 的收发顺序； -- server 增量解析； -- fallback 决策； -- `Socket` 建链完成通知； -- 公共 handshake phase 的推进。 - -## 6. Transport 能力接口 - -Transport 提供的是“被上层调用的资源能力”，不是 handshake driver。建议抽象为以下接口；具体命名可根据现有类调整： - -```cpp -class TransportUpgradeOps { -public: - virtual ~TransportUpgradeOps() = default; - - // 返回本端能力和协议 adapter 所需的字段来源。 - virtual handshake::StepResult PrepareLocal() = 0; - - // 上层完成 wire parse/validate 后交付远端参数。 - virtual handshake::StepResult ApplyRemote(const RemoteParameters&) = 0; - - // 创建、导入或建立数据面资源。 - virtual handshake::StepResult PrepareResources() = 0; - virtual handshake::StepResult NegotiateResources() = 0; - - // 握手 ACK 已发送/确认后切换数据面。 - virtual handshake::StepResult Activate() = 0; - - // fallback 或失败时关闭本次升级准备的资源。 - virtual void Deactivate() = 0; -}; -``` - -说明： - -- `PrepareLocal`、`ApplyRemote`、`PrepareResources`、`NegotiateResources` 的具体数量可以按现有资源模型合并，但调用顺序由 coordinator 固定。 -- 只有 `TransportUpgradeOps` 的实现可以访问 endpoint 和 transport-specific 类型。 -- 如果某种 Transport 不支持升级，应返回 `kFallback`，而不是创建一套假的握手状态机。 -- `Activate` 成功后，coordinator 才能发布 `kEstablished`。 -- 资源失败默认进入 `kFallbackTcp`；不可恢复的协议错误才进入 `kFailed`，具体策略由 coordinator 统一决定。 - -## 7. 协议 adapter 与公共编排的边界 - -每一种 wire protocol 保留自己的 `HandshakeCodec` 和字段 adapter： - -```cpp -struct HandshakeCodec { - int protocol_version; - FrameSpec hello_frame; - FrameSpec ack_frame; - std::function build_hello; - std::function parse_hello; - std::function build_ack; - std::function parse_ack; -}; -``` - -adapter 只处理以下内容： - -- magic、版本、字段序列化和反序列化； -- payload 合法性检查； -- 将字段转换为 `RemoteParameters`； -- 生成本端 hello 和 ack。 - -adapter 不再实现完整的 `RunServerHandshake` 或 `ProcessHandshakeAtClient`。这些函数中的流程代码应迁移到 coordinator；adapter 只提供 codec 和 `TransportUpgradeOps` 所需的字段转换。 - -server 端的 `HandshakeAdapter::ExecuteServerHandshake` 如果暂时需要保留以兼容 `InputMessenger`，其实现只能是薄适配层： - -```text -InputMessenger -> AdapterTransport/Coordinator -> HandshakeSession - | - `-> protocol codec + TransportUpgradeOps -``` - -它不应再根据 `SocketMode` 选择不同的 phase，也不应直接调用具体 Transport 的握手流程。 - -## 8. Client 建链时序 - -```mermaid -sequenceDiagram - participant S as Socket - participant A as AdapterTransport - participant C as Coordinator - participant T as TcpTransport - participant H as HandshakeSession - participant U as TransportUpgradeOps - - S->>A: Connect() - A->>T: 建立 TCP 控制连接 - T-->>A: TCP connected - A->>C: StartClientUpgrade() - C->>U: PrepareLocal() - C->>H: 发送 local hello - H->>T: WriteFrame() - T-->>H: 接收 remote hello - C->>U: ApplyRemote() / PrepareResources() - C->>U: NegotiateResources() - alt 升级成功 - C->>H: 发送 enabled ACK - C->>U: Activate() - C->>A: kEstablished - A->>S: ConnectionReady(high-speed) - else 对端不支持或资源失败 - C->>H: 发送 disabled ACK（如协议要求） - C->>U: Deactivate() - C->>A: kFallbackTcp - A->>S: ConnectionReady(tcp) - end -``` - -关键约束：`ConnectionReady` 只发送一次，并且必须发生在 active transport 已设置、fallback 缓冲区已回放、终态已 release-store 之后。 - -## 9. Server 建链时序 - -server 收到 TCP 数据后，由 `AdapterTransport` 的上层 coordinator 处理；具体 Transport 不参与入口选择： - -```text -TCP readable - -> AdapterTransport::ProcessUpgradeReadable - -> HandshakeSession::RunServer(input) - -> 根据 magic 选择 codec - -> 增量读取完整 hello - -> codec parse/validate - -> TransportUpgradeOps::ApplyRemote/NegotiateResources - -> 发送 ACK - -> Activate 或 Deactivate - -> AdapterTransport::CompleteConnection -``` - -对于 `IOBuf` 增量输入： - -- `kNeedMore` 时不得消费不完整 frame； -- magic 不匹配时必须把已检查的数据回放给 TCP parser； -- 已确认属于升级协议但资源失败时不能把完整握手 frame 当作普通 TCP 数据； -- ACK 没有 magic 时，选中的 codec 必须保存在 coordinator/session context 中，而不能依赖重新探测。 - -## 10. 连接终态与 Socket 通知 - -### 10.1 终态定义 - -| 终态 | active data transport | Socket 结果 | 说明 | -|---|---|---|---| -| `kEstablished` | 高速度 Transport | 建链完成 | `Activate()` 成功后发布 | -| `kFallbackTcp` | TCP Transport | 建链完成 | 协议不支持或资源不可用 | -| `kFailed` | 无 | 建链失败 | I/O、协议或不可恢复错误 | - -`kFallbackTcp` 不是失败。它表示控制连接成功且连接可继续使用 TCP。 - -### 10.2 发布顺序 - -统一采用以下顺序： - -```text -1. 设置 active transport -2. 回放 fallback 时已读但不属于握手的数据（仅 TCP fallback） -3. 发布终态 phase（release） -4. 通知 Socket::ConnectionReady / ConnectionFailed -``` - -事件线程观察到终态后，才能读取 active transport 和回放数据。禁止在 phase 发布后再修改 active transport。 - -### 10.3 TCP 数据保护 - -进入 `kEstablished` 后，TCP 只作为控制连接存在；如果收到额外 TCP 应用数据，应按协议错误处理，不能静默交给高速度数据面。进入 `kFallbackTcp` 后，后续数据全部交给 TCP Transport。 - -## 11. 现有实现迁移方案 - -### Phase 1：统一公共状态 - -- 用 `handshake::Phase` 替换 `HandshakePhases` 的六个 transport-specific 数值。 -- `HandshakeSession` 只写入统一 phase。 -- 保留具体 Transport 内部状态用于日志和资源调试，但不再通过 `handshake_phase()` 暴露给 Socket。 - -### Phase 2：收拢 client 入口 - -- 将所有 `ProcessHandshakeAtClient` 的调用点迁移到 `AdapterTransport::StartClientUpgrade`。 -- 删除具体 Transport 中的 client handshake driver。 -- 具体 Transport 改为实现 `TransportUpgradeOps`，只提供资源阶段操作。 -- `AdapterTransport::Connect` 在 TCP connected 后自动启动 coordinator。 - -### Phase 3：收拢 server 入口 - -- `InputMessenger` 只把 handshake 输入交给 `AdapterTransport`/coordinator。 -- `RdmaServerHandshakeAdapter`、`UBShmServerHandshakeAdapter` 等降级为 codec/字段 adapter 或薄兼容层。 -- 删除 adapter 内按 `SocketMode` 选择 `S_ACK_WAIT` 等 phase 的逻辑。 -- URMA 接入时直接实现统一的 codec 和 `TransportUpgradeOps`，不复制一套 server driver。 - -### Phase 4：统一连接完成通知 - -- 为 `AdapterTransport` 增加单一的 `CompleteConnection` 路径。 -- 升级成功、TCP fallback、失败分别通过统一终态通知 Socket。 -- 增加断言，确保 `ConnectionReady` 只发生一次，且发生在终态发布之后。 - -### Phase 5：删除旧路径 - -- 删除具体 Transport 的 `ProcessHandshakeAtClient`、`RunServerHandshake` 和握手 phase 映射。 -- 删除 Socket 对具体 Transport 类型和 transport-specific phase 的依赖。 -- 清理只为旧握手路径存在的 endpoint 回调。 - -## 12. 兼容性与测试要求 - -### 12.1 Wire compatibility - -重构不得改变： - -- RDMA v2/v3、URMA v2/v3、UBSHM v2 的 magic； -- frame length 的字节序和语义； -- hello/ack 的字段布局和 ACK bit； -- fallback hello 的兼容行为。 - -### 12.2 公共模块测试 - -至少覆盖： - -- 2 字节和 4 字节 magic； -- fixed、U16 total length、U32 body length； -- 半包、粘包和多余数据； -- `kNeedMore` 不消费不完整输入； -- magic 不匹配的 push-back； -- I/O error、EOF、超时和协议错误； -- 终态发布顺序和一次性连接通知。 - -### 12.3 集成矩阵 - -| 场景 | 预期结果 | -|---|---| -| RDMA v2 ↔ RDMA v2 | 高速度 Transport 建链 | -| RDMA v3 ↔ RDMA v3 | 高速度 Transport 建链 | -| URMA v2/v3 ↔ 对应版本 | 高速度 Transport 建链 | -| UBSHM ↔ UBSHM | UBSHM 建链 | -| 高速度 client ↔ TCP server | TCP fallback，Socket 建链成功 | -| TCP client ↔ 高速度 server | TCP 建链成功 | -| hello 合法但资源创建失败 | TCP fallback | -| 已识别握手协议后发生协议错误 | 建链失败，不回放为 TCP 数据 | - -## 13. 验收标准 - -设计落地后应满足： - -1. `Socket::_transport` 永远只指向 `AdapterTransport`，Socket 源码中没有具体 Transport 类型判断。 -2. client 和 server 都只有一个握手编排入口，代码库中不再存在具体 Transport 的 `ProcessHandshakeAtClient`。 -3. 公共握手状态只有一套 `handshake::Phase`，不存在 RDMA/UBSHM/URMA 到公共 phase 的数字映射。 -4. 具体 Transport 不读写 TCP fd，不解析 magic/length，不调用 `HandshakeSession::RunClient/RunServer`。 -5. 升级成功和 TCP fallback 都由 `AdapterTransport` 自动完成 active transport 切换，并向 Socket 发出一次建链完成通知。 -6. 在 wire compatibility 测试通过的前提下，握手流程、fallback 和连接完成通知可以通过公共 coordinator 单元测试验证。 - -## 14. 结论 - -最终边界如下： - -```text -Socket - 只感知 AdapterTransport 和 ConnectionReady - -AdapterTransport / Coordinator - 拥有 TCP-first、握手状态机、升级/fallback 编排、active transport 切换 - -HandshakeSession / Codec - 拥有公共 I/O、framing、协议字段编解码和统一状态 - -TransportUpgradeOps - 提供具体传输所需的资源阶段接口 - -RDMA / URMA / UBSHM Transport - 只拥有各自资源、数据面和 completion -``` - -一句话概括：**握手是上层连接编排，Transport 是被编排的数据传输能力；Socket 只通过 AdapterTransport 观察连接结果。** -## 15. 整体架构图 - -```mermaid -flowchart TB - S["Socket\n只感知 AdapterTransport"] - A["AdapterTransport\nTCP-first / 建链编排 / active transport"] - C["Connection Upgrade Coordinator\nStartClientUpgrade / ProcessUpgradeReadable\nCompleteConnection"] - H["HandshakeSession\n统一 framing / I/O / handshake::Phase"] - RA["Protocol Codec Adapter\nRDMA v2/v3 / URMA / UBSHM"] - O["TransportUpgradeOps\nPrepare / Negotiate / Activate / Deactivate"] - T["Concrete Transport\nTCP / RDMA / URMA / UBSHM"] - E["Endpoint / Resource\n只提供传输资源能力"] - IM["InputMessenger\nserver 增量输入"] - READY["ConnectionReady\nESTABLISHED / FALLBACK_TCP"] - FAIL["ConnectionFailed\nFAILED"] - - S -->|Connect| A - A -->|TCP connected| C - IM -->|ProcessUpgradeReadable| A - A --> C - C <--> H - C --> RA - C --> O - O --> T - T --> E - C -->|terminal phase| A - A --> READY - A --> FAIL -``` - -核心边界：`HandshakeSession` 负责公共握手流程，协议 Adapter 只负责 wire codec，`TransportUpgradeOps` 只负责资源能力；具体 Transport 不读取握手输入、不编排握手阶段，也不向 Socket 暴露 transport-specific phase。 +测试覆盖分片和粘连输入、magic 前缀分流、协议版本与格式校验、TCP 回退、 +资源释放及终态发布，并分别运行 RDMA 与 UBSHM 回归测试。URMA 的公共握手 +迁移需要单独设计和测试，不应仅凭本次重构推断其已完成。 diff --git a/docs/cn/server.md b/docs/cn/server.md index 565b7bb3d0..efb6e9f2e6 100644 --- a/docs/cn/server.md +++ b/docs/cn/server.md @@ -378,6 +378,16 @@ Server.set_version(...)可以为server设置一个名称+版本，可通过/vers | ------------------------- | ----- | ---------------------------------------- | ------------------- | | log_idle_connection_close | false | Print log when an idle connection is closed | src/brpc/socket.cpp | +## 限制服务端连接数 + +设置`ServerOptions.max_connections`可以限制Server公共监听端口上的并发连接数。默认值为0，表示不限制。该端口上的所有协议和服务共享同一个上限，包括RPC、HTTP、Redis和内置服务，无需配置专用协议。 + +Acceptor会在创建brpc Socket前预留连接名额，因此空闲连接和尚未完成TLS握手的连接也计入上限。超限连接会在协议解析或TLS认证前被直接关闭，不发送协议专用的错误响应。Socket回收时释放连接名额。 + +通过`ServerOptions.internal_port`配置的内部监听端口不受限制，也不占用公共端口的连接名额，因此公共端口满载时，内部端口上的内置服务仍可访问。同一个业务Service注册到多个Server、监听多个端口时，各公共监听端口分别计数和限流；即使共享同一个Service对象，也不会合并连接数。`ServerStatistics.connection_count`包含公共和内部端口的连接总数，因此可能超过`max_connections`；`ServerStatistics.rejected_connection_count`记录公共端口因超限而拒绝的累计连接数。 + +运行中的Server可以调用`Server::SetMaxConnections()`原子更新上限，也可以在以无限制值启动后动态开启限制。调高上限会影响后续连接准入；调低上限不会断开已有连接，公共端口的活跃连接数降到新上限以下后才会重新接受新连接。设置为0会关闭限制。Server未运行时该方法返回-1；该方法不会修改`options().max_connections`记录的启动值。每次`Start()`都使用该次传入选项中的上限。 + ## pid_file 如果设置了此字段，Server启动时会创建一个同名文件，内容为进程号。默认为空。 diff --git a/docs/cn/urma.md b/docs/cn/urma.md new file mode 100644 index 0000000000..62e10389a8 --- /dev/null +++ b/docs/cn/urma.md @@ -0,0 +1,195 @@ +# 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 commit，可通过 `DOWNLOAD_URMA_HEADERS=OFF` 禁止下载。找到 `liburma` 时使用 +真实硬件数据通路，否则链接 brpc 的 mock，使 URMA 代码和测试仍可在无硬件 +环境编译。`URMA_USE_MOCK` 默认为 `AUTO`；设为 `ON` 可强制使用 mock，设为 +`OFF` 则要求找到 `liburma` 并排除 mock。 + +### Make 编译 + +```bash +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma +make -j4 + +# 强制使用 mock，或强制使用真实 liburma +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma --with-urma-mock +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma --without-urma-mock +``` + +默认根据 `liburma` 是否存在自动选择；`--without-urma-mock` 要求配置路径中 +存在 `liburma`。 + +### Bazel 编译 + +```bash +# 编译带 URMA 支持的 brpc +bazel build --define=BRPC_WITH_URMA=true //:brpc + +# 仅展示关闭下载的开关；头文件仍须由工具链或本地 BUILD 目标声明 +bazel build --define=BRPC_WITH_URMA=true \ + --define=BRPC_DOWNLOAD_URMA_HEADERS=false //:brpc + +# 编译并运行 URMA 单元测试（使用 mock，不需要 URMA 硬件） +bazel test --define=BRPC_WITH_URMA=true \ + --define=BRPC_URMA_USE_MOCK=true //test:brpc_unittests + +# 使用真实 liburma；要求链接器可找到 -lurma +bazel build --define=BRPC_WITH_URMA=true \ + --define=BRPC_URMA_USE_MOCK=false //:brpc +``` + +Bazel 会自动获取固定 commit 的 UMDK 头文件，并通过 `@umdk//:urma_headers` +提供给 URMA 目标。设置 `BRPC_DOWNLOAD_URMA_HEADERS=false` 后，Bazel 不再依赖 +该外部仓库；实际构建仍需通过工具链或本地 BUILD 目标声明 +`urma_api.h`、`urma_types.h` 和 `urma_ubagg.h`。仅通过 +`--copt=-I/path` 指向工作区外目录，在默认 Bazel 沙箱下不会把未声明的头文件 +带入编译输入。设置 +`BRPC_URMA_USE_MOCK=true` 强制加入 mock；设置为 `false` 则排除 mock 并链接 +`liburma`。未指定时，URMA 构建默认使用 mock。 + +## 使用 + +通过在 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/bazel_support.md b/docs/en/bazel_support.md index d1497606ce..ed3cc448a1 100644 --- a/docs/en/bazel_support.md +++ b/docs/en/bazel_support.md @@ -28,7 +28,7 @@ module( ) bazel_dep(name = "protobuf", version = "27.3", repo_name = "com_google_protobuf") -bazel_dep(name = "brpc", version = "1.17.0", repo_name = "apache_brpc") +bazel_dep(name = "brpc", version = "1.18.0", repo_name = "apache_brpc") local_path_override( module_name = "brpc", diff --git a/docs/en/getting_started.md b/docs/en/getting_started.md index 9acadc75e0..d605975e20 100644 --- a/docs/en/getting_started.md +++ b/docs/en/getting_started.md @@ -291,6 +291,45 @@ openssl installed in Monterey may not be found at `/usr/local/opt/openssl`, inst * Run `brew link openssl --force` first and check if `/usr/local/opt/openssl` appears. * If above command does not work, consider making a soft link using `sudo ln -s /opt/homebrew/Cellar/openssl@3/3.0.3 /usr/local/opt/openssl`. Note that the installed openssl in above command may be put in different places in different environments, which could be revealed by running `brew info openssl`. +### Compile a Debug build with CMake + +Apple Silicon can build the Debug configuration against dependencies installed by Homebrew: + +```shell +cmake -S . -B build -G "Unix Makefiles" \ + -DCMAKE_BUILD_TYPE=Debug \ + -DDEBUG=ON \ + -DCMAKE_OSX_ARCHITECTURES=arm64 \ + -DCMAKE_OSX_SYSROOT="$(xcrun --sdk macosx --show-sdk-path)" \ + -DCMAKE_PREFIX_PATH="$(brew --prefix)" \ + -DOPENSSL_ROOT_DIR="$(brew --prefix openssl@3)" +cmake --build build --parallel +``` + +The command uses the single-config Unix Makefiles generator, so +`CMAKE_BUILD_TYPE=Debug` selects CMake's Debug configuration. `DEBUG=ON` +enables brpc's debug logs and keeps assertions enabled. Build artifacts are +written to `build/output/`. + +To generate a compilation database for tools such as clangd, add the following +optional argument to the configuration command: + +```shell +-DCMAKE_EXPORT_COMPILE_COMMANDS=ON +``` + +The compilation database is generated at `build/compile_commands.json`. + +If CMake reports `tapi error: malformed file`, make sure the Command Line Tools +and Xcode versions match, then select the installed Xcode: + +```shell +sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer +``` + +If the problem persists, verify that `xcrun --sdk macosx --show-sdk-path` returns +a valid SDK path. + ### Compile brpc with config_brpc.sh git clone brpc, cd into the repo and run ```shell diff --git a/docs/en/server.md b/docs/en/server.md index f13193a8e6..30c7a8677b 100644 --- a/docs/en/server.md +++ b/docs/en/server.md @@ -375,6 +375,16 @@ If [-log_idle_connection_close](http://brpc.baidu.com:8765/flags/log_idle_connec | ------------------------- | ----- | ---------------------------------------- | ------------------- | | log_idle_connection_close | false | Print log when an idle connection is closed | src/brpc/socket.cpp | +## Limit server connections + +Set `ServerOptions.max_connections` to limit simultaneous connections on the Server's public listener. The default value is 0 (unlimited). All protocols and services on that listener share the same limit, including RPC, HTTP, Redis, and builtin services. No dedicated protocol configuration is required. + +The acceptor reserves a slot before creating a brpc Socket, so idle connections and connections awaiting TLS handshakes count toward the limit. Connections over the limit are closed immediately without a protocol-specific response, before protocol parsing or TLS authentication. Connection slots are released when sockets are recycled. + +The internal listener configured by `ServerOptions.internal_port` is unlimited and does not consume public connection slots, so builtin services on that port remain available when the public listener is full. A business service registered with multiple Server instances on different ports has an independent limit on each public listener, even when those Servers share the same service object. `ServerStatistics.connection_count` includes both public and internal connections and may therefore exceed `max_connections`; `ServerStatistics.rejected_connection_count` reports the cumulative number of connections rejected by the public listener's limit. + +Call `Server::SetMaxConnections()` to atomically update the limit on a running Server, including one started without a limit. Raising the limit affects subsequent admission checks. Lowering it does not close existing connections; new connections are accepted again after the active public connection count falls below the limit. Set the limit to 0 to disable it. The setter returns -1 if the Server is not running and leaves `options().max_connections` at its startup value. Each `Start()` uses the limit in the options passed to that call. + ## pid_file If this field is non-empty, Server creates a file named so at start-up, with pid as the content. Empty by default. diff --git a/docs/en/urma.md b/docs/en/urma.md new file mode 100644 index 0000000000..a5b80f45b9 --- /dev/null +++ b/docs/en/urma.md @@ -0,0 +1,209 @@ +# 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 UMDK at a pinned +commit 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. `URMA_USE_MOCK` defaults to `AUTO`; set it to `ON` to force the mock, +or to `OFF` to require `liburma` and exclude the mock. + +### Build with Make + +```bash +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma +make -j4 + +# Force the mock, or require the real liburma +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma --with-urma-mock +sh config_brpc.sh --headers=/usr/include --libs=/usr/lib \ + --with-urma --without-urma-mock +``` + +Without either override, the Make configuration selects the implementation +automatically from whether `liburma` is available. `--without-urma-mock` +requires `liburma` in the configured library paths. + +### Build with Bazel + +```bash +# Build brpc with URMA support +bazel build --define=BRPC_WITH_URMA=true //:brpc + +# This only shows the download switch; headers must still be declared by the toolchain or a local BUILD target +bazel build --define=BRPC_WITH_URMA=true \ + --define=BRPC_DOWNLOAD_URMA_HEADERS=false //:brpc + +# Build and run the URMA unit tests (using the mock, without URMA hardware) +bazel test --define=BRPC_WITH_URMA=true \ + --define=BRPC_URMA_USE_MOCK=true //test:brpc_unittests + +# Use the real liburma; the linker must be able to find -lurma +bazel build --define=BRPC_WITH_URMA=true \ + --define=BRPC_URMA_USE_MOCK=false //:brpc +``` + +Bazel fetches the pinned UMDK commit and exposes its headers through +`@umdk//:urma_headers` by default. With +`BRPC_DOWNLOAD_URMA_HEADERS=false`, Bazel does not depend on that external +repository; the actual build must still declare `urma_api.h`, `urma_types.h`, +and `urma_ubagg.h` through the toolchain or a local BUILD target. Passing only +`--copt=-I/path` to a directory outside the workspace does not make undeclared +headers available inside Bazel's default sandbox. +`BRPC_URMA_USE_MOCK=true` forces the link-time mock, while `false` excludes it +and links `liburma`. If unspecified, URMA builds use the mock by default. + +## 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 9eec3c0f86..6b2c7850ff 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 14) + 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_14) + if(NOT BRPC_EXAMPLE_CXX_STANDARD) + set(BRPC_EXAMPLE_CXX_STANDARD 14) + 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..4dbc989a45 --- /dev/null +++ b/example/urma_performance/CMakeLists.txt @@ -0,0 +1,54 @@ +# 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_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) +if(NOT URMA_INCLUDE_PATH) + message(FATAL_ERROR + "Fail to find urma_api.h. Install UMDK headers or set URMA_ROOT.") +endif() +find_library(URMA_LIB NAMES urma + HINTS ENV URMA_ROOT + PATH_SUFFIXES lib lib64) +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/package/rpm/brpc.spec b/package/rpm/brpc.spec index 19c2d83b12..678ace746c 100644 --- a/package/rpm/brpc.spec +++ b/package/rpm/brpc.spec @@ -18,7 +18,7 @@ # Name: brpc -Version: 1.17.0 +Version: 1.18.0 Release: 1%{?dist} Summary: Industrial-grade RPC framework using C++ Language. diff --git a/src/brpc/acceptor.cpp b/src/brpc/acceptor.cpp index 8333a17127..99280037d5 100644 --- a/src/brpc/acceptor.cpp +++ b/src/brpc/acceptor.cpp @@ -16,7 +16,9 @@ // under the License. +#include #include +#include #include #include "butil/fd_guard.h" // fd_guard #include "butil/fd_utility.h" // make_close_on_exec @@ -38,6 +40,9 @@ Acceptor::Acceptor(bthread_keytable_pool_t* pool) , _listened_fd(-1) , _acception_id(0) , _empty_cond(&_map_mutex) + , _connection_count(0) + , _rejected_connection_count(0) + , _max_connections(0) , _force_ssl(false) , _ssl_ctx(nullptr) , _socket_mode(SOCKET_MODE_TCP) @@ -52,6 +57,13 @@ Acceptor::~Acceptor() { int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, const std::shared_ptr& ssl_ctx, bool force_ssl) { + return StartAccept(listened_fd, idle_timeout_sec, ssl_ctx, force_ssl, 0); +} + +int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, + const std::shared_ptr& ssl_ctx, + bool force_ssl, + size_t max_connections) { if (listened_fd < 0) { LOG(FATAL) << "Invalid listened_fd=" << listened_fd; return -1; @@ -87,6 +99,7 @@ int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec, _idle_timeout_sec = idle_timeout_sec; _force_ssl = force_ssl; _ssl_ctx = ssl_ctx; + SetMaxConnections(max_connections); // Creation of _acception_id is inside lock so that OnNewConnections // (which may run immediately) should see sane fields set below. @@ -200,9 +213,37 @@ void Acceptor::Join() { } size_t Acceptor::ConnectionCount() const { - // Notice that _socket_map may be modified concurrently. This actually - // assumes that size() is safe to call concurrently. - return _socket_map.size(); + return _connection_count.load(butil::memory_order_relaxed); +} + +size_t Acceptor::RejectedConnectionCount() const { + return _rejected_connection_count.load(butil::memory_order_relaxed); +} + +bool Acceptor::TryAcquireConnectionSlot() { + size_t count = _connection_count.load(butil::memory_order_relaxed); + do { + const size_t max_connections = + _max_connections.load(butil::memory_order_relaxed); + if (max_connections != 0 && count >= max_connections) { + return false; + } + } while (!_connection_count.compare_exchange_weak( + count, count + 1, butil::memory_order_relaxed)); + return true; +} + +void Acceptor::ReleaseConnectionSlot() { + const size_t previous = + _connection_count.fetch_sub(1, butil::memory_order_relaxed); + CHECK_GT(previous, 0u); +} + +void Acceptor::SetMaxConnections(size_t max_connections) { + // The limit controls only future numeric admission decisions and does not + // publish socket state, so a relaxed store is sufficient. + _max_connections.store( + max_connections, butil::memory_order_relaxed); } void Acceptor::ListConnections(std::vector* conn_list, @@ -275,7 +316,18 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { acception->SetFailed(EINVAL, "Impossible! acception->user() MUST be Acceptor"); return; } - + + if (!am->TryAcquireConnectionSlot()) { + am->_rejected_connection_count.fetch_add( + 1, butil::memory_order_relaxed); + LOG_EVERY_SECOND(WARNING) + << "Reject connection on " << acception->local_side() + << ": max_connections limit reached"; + // in_fd closes the connection before Socket::Create(), protocol + // parsing or TLS authentication, without a protocol-specific reply. + continue; + } + SocketId socket_id; SocketOptions options; options.keytable_pool = am->_keytable_pool; @@ -288,21 +340,20 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { options.socket_mode = am->_socket_mode; options.bthread_tag = am->_bthread_tag; if (Socket::Create(options, &socket_id) != 0) { + am->ReleaseConnectionSlot(); LOG(ERROR) << "Fail to create Socket"; continue; } in_fd.release(); // transfer ownership to socket_id - // There's a funny race condition here. After Socket::Create, messages - // from the socket are already handled and a RPC is possibly done - // before the socket is added into _socket_map below. This is found in - // ChannelTest.skip_parallel in test/brpc_channel_unittest.cpp (running - // on machines with few cores) where the _messenger.ConnectionCount() - // may surprisingly be 0 even if the RPC is already done. + // Socket::Create() may start processing messages immediately. The + // connection is already counted, but this loop owns its slot until + // the socket is pinned and inserted into _socket_map below. SocketUniquePtr sock; if (Socket::AddressFailedAsWell(socket_id, &sock) >= 0) { bool is_running = true; + bool registered = false; { BAIDU_SCOPED_LOCK(am->_map_mutex); is_running = (am->status() == RUNNING); @@ -311,7 +362,13 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { // running or not. Otherwise, `Acceptor::BeforeRecycle' // may be called (inside Socket::BeforeRecycled) after `Acceptor' // has been destroyed - am->_socket_map.insert(socket_id, ConnectStatistics()); + registered = (am->_socket_map.insert( + socket_id, ConnectStatistics()) != nullptr); + } + if (!registered) { + am->ReleaseConnectionSlot(); + sock->SetFailed(ENOMEM, "Fail to register accepted Socket"); + continue; } if (!is_running) { LOG(WARNING) << "Acceptor on fd=" << acception->fd() @@ -321,8 +378,11 @@ void Acceptor::OnNewConnectionsUntilEAGAIN(Socket* acception) { sock->description().c_str()); return; } - } // else: The socket has already been destroyed, Don't add its id - // into _socket_map + } else { + // The socket was recycled before insertion. BeforeRecycle() did + // not release its slot, which is still owned by this accept loop. + am->ReleaseConnectionSlot(); + } } } @@ -346,9 +406,11 @@ void Acceptor::BeforeRecycle(Socket* sock) { _empty_cond.Broadcast(); return; } - // If a Socket could not be addressed shortly after its creation, it - // was not added into `_socket_map'. - _socket_map.erase(sock->id()); + // Socket::Create() can fail or the socket can be recycled before it is + // inserted into the map. In those cases the accept loop releases the slot. + if (_socket_map.erase(sock->id()) != 0) { + ReleaseConnectionSlot(); + } if (_socket_map.empty()) { _empty_cond.Broadcast(); } diff --git a/src/brpc/acceptor.h b/src/brpc/acceptor.h index 2f138b30ea..01a2b26a9d 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -19,6 +19,7 @@ #define BRPC_ACCEPTOR_H #include "bthread/bthread.h" // bthread_t +#include "butil/atomicops.h" // butil::atomic #include "butil/synchronization/condition_variable.h" #include "butil/containers/flat_map.h" #include "brpc/input_messenger.h" @@ -58,6 +59,12 @@ friend class Server; int StartAccept(int listened_fd, int idle_timeout_sec, const std::shared_ptr& ssl_ctx, bool force_ssl); + // Limit simultaneous connections on this listener. 0 means unlimited. + // Excess connections are closed before protocol parsing or TLS. + int StartAccept(int listened_fd, int idle_timeout_sec, + const std::shared_ptr& ssl_ctx, + bool force_ssl, + size_t max_connections); // [thread-safe] Stop accepting connections. // `closewait_ms' is not used anymore. @@ -72,6 +79,10 @@ friend class Server; // Get number of existing connections. size_t ConnectionCount() const; + // Get the cumulative number of connections rejected by this listener's + // connection limit. + size_t RejectedConnectionCount() const; + // Clear `conn_list' and append all connections into it. void ListConnections(std::vector* conn_list); @@ -93,6 +104,10 @@ friend class Server; // Remove the accepted socket `sock' from inside void BeforeRecycle(Socket* sock) override; + bool TryAcquireConnectionSlot(); + void ReleaseConnectionSlot(); + void SetMaxConnections(size_t max_connections); + bthread_keytable_pool_t* _keytable_pool; // owned by Server Status _status; int _idle_timeout_sec; @@ -108,6 +123,15 @@ friend class Server; // The map containing all the accepted sockets SocketMap _socket_map; + // A slot is reserved before Socket::Create(), closing the race where a + // socket starts processing before it is inserted into _socket_map. Until + // insertion, the accept loop owns the slot; afterwards BeforeRecycle() + // releases it. These atomics publish no socket state, so relaxed memory + // ordering is sufficient. + butil::atomic _connection_count; + butil::atomic _rejected_connection_count; + butil::atomic _max_connections; + bool _force_ssl; std::shared_ptr _ssl_ctx; diff --git a/src/brpc/adapter_transport.cpp b/src/brpc/adapter_transport.cpp index 1fb093a65d..ed51326426 100644 --- a/src/brpc/adapter_transport.cpp +++ b/src/brpc/adapter_transport.cpp @@ -31,6 +31,7 @@ #include "brpc/rdma/rdma_helper.h" #endif #if BRPC_WITH_UBRING +#include "brpc/ubshm/ub_endpoint.h" #include "brpc/ubshm/ub_helper.h" #include "brpc/ubshm/ubr_trx.h" #endif @@ -120,6 +121,95 @@ struct ClientHandshakeTask { SocketUniquePtr socket; }; +#if BRPC_WITH_RDMA +class RdmaClientHandshakeTransport : public handshake::HandshakeTransport { +public: + RdmaClientHandshakeTransport( + RdmaTransport* transport, rdma::RdmaHandshakeAdapter* protocol, + Socket* socket, int* connect_error) + : _transport(transport), _protocol(protocol), _socket(socket), + _connect_error(connect_error) {} + + handshake::StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources() == 0) { + return handshake::STEP_OK; + } + errno = 0; + return handshake::STEP_FALLBACK; + } + + handshake::StepResult NegotiateResources() override { + return _transport->NegotiateUpgradeResources( + _protocol->remote(), false) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + + void OnFallback() override { _transport->DeactivateUpgrade(); } + + void OnFailed() override { + _transport->DeactivateUpgrade(); + const int saved_errno = errno != 0 ? errno : EPROTO; + *_connect_error = saved_errno; + _socket->SetFailed(saved_errno, + "Fail to complete rdma handshake from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + } + +private: + RdmaTransport* _transport; + rdma::RdmaHandshakeAdapter* _protocol; + Socket* _socket; + int* _connect_error; +}; +#endif + +#if BRPC_WITH_UBRING +class UBShmClientHandshakeTransport : public handshake::HandshakeTransport { +public: + UBShmClientHandshakeTransport( + UBShmTransport* transport, ubring::SHM* local_shm, + const std::string& shm_name, Socket* socket, int* connect_error) + : _transport(transport), _local_shm(local_shm), + _shm_name(shm_name), _socket(socket), + _connect_error(connect_error) {} + + handshake::StepResult PrepareResources() override { + return _transport->PrepareUpgradeResources( + _local_shm, _shm_name.c_str()) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + handshake::StepResult NegotiateResources() override { + return _transport->NegotiateUpgradeResources( + _local_shm, _shm_name.c_str()) == 0 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + + void OnFailed() override { + _transport->DeactivateUpgrade(); + const int saved_errno = errno != 0 ? errno : EPROTO; + *_connect_error = saved_errno; + _socket->SetFailed(saved_errno, + "Fail to complete ubring handshake from %s: %s", + _socket->description().c_str(), + berror(saved_errno)); + } + +private: + UBShmTransport* _transport; + ubring::SHM* _local_shm; + std::string _shm_name; + Socket* _socket; + int* _connect_error; +}; +#endif + } // namespace AdapterTransport::AdapterTransport(SocketMode mode) @@ -250,35 +340,10 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { std::unique_ptr protocol = transport->CreateClientHandshakeAdapter(); CHECK(protocol != NULL); - rdma::ParsedHello remote{}; - handshake::ClientHandshakeCallbacks callbacks{}; - callbacks.codec = protocol->MakeCodec(&remote); - callbacks.transport.prepare_resources = [&]() { - if (transport->PrepareUpgradeResources() == 0) { - return handshake::STEP_OK; - } - errno = 0; - return handshake::STEP_FALLBACK; - }; - callbacks.transport.negotiate_resources = [&]() { - return transport->NegotiateUpgradeResources(remote, false) == 0 - ? handshake::STEP_OK : handshake::STEP_FALLBACK; - }; - callbacks.transport.set_high_speed_active = [transport]() { - transport->ActivateUpgrade(); - }; - callbacks.transport.set_tcp_active = [transport]() { - transport->DeactivateUpgrade(); - }; - callbacks.transport.on_failed = [&]() { - const int saved_errno = errno != 0 ? errno : EPROTO; - connect_error = saved_errno; - socket->SetFailed(saved_errno, - "Fail to complete rdma handshake from %s: %s", - socket->description().c_str(), - berror(saved_errno)); - }; - const handshake::StepResult result = adapter->_handshake.RunClient(callbacks); + RdmaClientHandshakeTransport participant( + transport, protocol.get(), socket, &connect_error); + const handshake::StepResult result = adapter->_handshake.RunClient( + protocol.get(), &participant); if (result == handshake::STEP_OK && transport->StartUpgradeEvents() < 0) { const int saved_errno = errno != 0 ? errno : ERDMA; @@ -317,44 +382,16 @@ void* AdapterTransport::ProcessClientHandshake(void* arg) { NULL, local_shm_len, 0, {0}, static_cast(socket->fd())}; const auto shm_name_str = butil::endpoint2str(socket->local_side()); - ubring::HelloMessage remote{}; ubring::UBShmHandshakeAdapter wire; - handshake::ClientHandshakeCallbacks callbacks{}; - callbacks.codec = wire.MakeCodec(); - callbacks.codec.build_hello = [&](bool enabled, std::string* payload) { - CHECK(enabled); - return wire.BuildHello(true, local_shm_len, shm_name_str.c_str(), - payload); - }; - callbacks.codec.parse_hello = [&](const std::string& payload) { - return wire.ParseHello(payload, &remote); - }; - callbacks.transport.prepare_resources = [&]() { - return transport->PrepareUpgradeResources( - &local_trx_shm, shm_name_str.c_str()) == 0 - ? handshake::STEP_OK : handshake::STEP_FALLBACK; - }; - callbacks.transport.negotiate_resources = [&]() { - return transport->NegotiateUpgradeResources( - &local_trx_shm, shm_name_str.c_str()) == 0 - ? handshake::STEP_OK : handshake::STEP_FALLBACK; - }; - callbacks.transport.set_high_speed_active = [transport]() { - transport->ActivateUpgrade(); - }; - callbacks.transport.set_tcp_active = [transport]() { - transport->DeactivateUpgrade(); - }; - callbacks.transport.on_failed = [&]() { - const int saved_errno = errno != 0 ? errno : EPROTO; - connect_error = saved_errno; - socket->SetFailed(saved_errno, - "Fail to complete ubring handshake from %s: %s", - socket->description().c_str(), - berror(saved_errno)); - }; - const handshake::StepResult result = adapter->_handshake.RunClient(callbacks); + wire.ConfigureClientHello(local_shm_len, shm_name_str.c_str()); + UBShmClientHandshakeTransport participant( + transport, &local_trx_shm, shm_name_str.c_str(), socket, + &connect_error); + const handshake::StepResult result = adapter->_handshake.RunClient( + &wire, &participant); if (result == handshake::STEP_OK) { + transport->GetUBShmEp()->SetNegotiatedDataFormat( + ubring::UBR_DATA_FORMAT_LEGACY_64); transport->FinishUpgrade(); } if (result == handshake::STEP_ERROR && connect_error == 0) { @@ -501,6 +538,8 @@ void AdapterTransport::Debug(std::ostream& os) { case handshake::NEGOTIATING: state = "NEGOTIATING"; break; case handshake::ACK_SEND: state = "ACK_SEND"; break; case handshake::ACK_WAIT: state = "ACK_WAIT"; break; + case handshake::EXTENSION_SEND: state = "EXTENSION_SEND"; break; + case handshake::EXTENSION_WAIT: state = "EXTENSION_WAIT"; break; case handshake::ESTABLISHED: state = "ESTABLISHED"; break; case handshake::FALLBACK_TCP: state = "FALLBACK_TCP"; break; case handshake::FAILED: state = "FAILED"; break; diff --git a/src/brpc/builtin/flags_service.cpp b/src/brpc/builtin/flags_service.cpp index 18152baa86..569da31394 100644 --- a/src/brpc/builtin/flags_service.cpp +++ b/src/brpc/builtin/flags_service.cpp @@ -117,11 +117,11 @@ void FlagsService::set_value_page(Controller* cntl, const bool is_string = (info.type == "string"); os << "" "
" - " Set `" << name << "' from "; + " Set `" << WebEscape(name) << "' from "; if (is_string) { os << '"'; } - os << info.current_value; + os << WebEscape(info.current_value); if (is_string) { os << '"'; } @@ -177,9 +177,12 @@ void FlagsService::default_method(::google::protobuf::RpcController* cntl_base, return; } butil::IOBufBuilder os; - os << "Set `" << constraint << "' to " << *value_str; if (use_html) { + os << "Set `" << WebEscape(constraint) << "' to " + << WebEscape(*value_str); os << "
[back to flags]"; + } else { + os << "Set `" << constraint << "' to " << *value_str; } os.move_to(cntl->response_attachment()); return; diff --git a/src/brpc/builtin/prometheus_metrics_service.cpp b/src/brpc/builtin/prometheus_metrics_service.cpp index c02e78ddf4..16ffd89515 100644 --- a/src/brpc/builtin/prometheus_metrics_service.cpp +++ b/src/brpc/builtin/prometheus_metrics_service.cpp @@ -181,6 +181,12 @@ bool PrometheusMetricsDumper::DumpLatencyRecorderSuffix( if (!si->IsComplete()) { return true; } + // The average latency can not be a quantile series of the summary below, + // because the quantile label must be parsable as a float. Dump it as a + // separate gauge, which is the same as the multi dimension one does. + *_os << "# HELP " << si->metric_name << "_avg_latency" << '\n' + << "# TYPE " << si->metric_name << "_avg_latency gauge\n" + << si->metric_name << "_avg_latency " << si->latency_avg << '\n'; *_os << "# HELP " << si->metric_name << '\n' << "# TYPE " << si->metric_name << " summary\n" << si->metric_name << "{quantile=\"" @@ -198,8 +204,6 @@ bool PrometheusMetricsDumper::DumpLatencyRecorderSuffix( << si->latency_percentiles[4] << '\n' << si->metric_name << "{quantile=\"1\"} " << si->latency_percentiles[5] << '\n' - << si->metric_name << "{quantile=\"avg\"} " - << si->latency_avg << '\n' << si->metric_name << "_sum " // There is no sum of latency in bvar output, just use // average * count as approximation diff --git a/src/brpc/channel.cpp b/src/brpc/channel.cpp index ff81e521d5..de0a1bde49 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 47f262f627..a778a1bd2a 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/circuit_breaker.cpp b/src/brpc/circuit_breaker.cpp index 785ec77a5c..385c5a7bdc 100644 --- a/src/brpc/circuit_breaker.cpp +++ b/src/brpc/circuit_breaker.cpp @@ -18,6 +18,7 @@ #include "brpc/circuit_breaker.h" #include +#include #include #include "brpc/errno.pb.h" @@ -191,7 +192,9 @@ bool CircuitBreaker::OnCallEnd(int error_code, int64_t latency) { // since the latency corresponding to ELIMIT is usually very small, we // cannot handle it as a successful request. Here we simply ignore the requests // that returned ELIMIT. - if (error_code == ELIMIT) { + // Canceled requests do not indicate a server failure and should not + // contribute samples or count as successful half-open probes either. + if (error_code == ELIMIT || error_code == ECANCELED) { return true; } if (_broken.load(butil::memory_order_relaxed)) { diff --git a/src/brpc/details/http_message.cpp b/src/brpc/details/http_message.cpp index 0cb4f78378..14a81e553c 100644 --- a/src/brpc/details/http_message.cpp +++ b/src/brpc/details/http_message.cpp @@ -49,6 +49,9 @@ DEFINE_int32(http_verbose_max_body_length, 512, DEFINE_bool(http_check_outbound_header_crlf, true, "Skip outbound http header fields whose name or value contains " "CR/LF to prevent request/response splitting."); +DEFINE_uint32(http_max_header_count, 100, + "Reject a message carrying more than so many header fields. " + "0 lifts the limit."); DECLARE_int64(socket_max_unwritten_bytes); DECLARE_uint64(max_body_size); @@ -131,6 +134,14 @@ int HttpMessage::on_header_value(http_parser *parser, http_message->_cur_value = &header.AddHeader(http_message->_cur_header); } + + if (FLAGS_http_max_header_count > 0 && + header.HeaderCount() > FLAGS_http_max_header_count) { + LOG(ERROR) << "Too many headers, max=" + << FLAGS_http_max_header_count; + return -1; + } + if (http_message->_cur_value && !http_message->_cur_value->empty()) { http_message->_cur_value->append( header.HeaderValueDelimiter(http_message->_cur_header)); diff --git a/src/brpc/handshake/rdma_handshake.cpp b/src/brpc/handshake/rdma_handshake.cpp index 61bc5fc762..26c622a1ac 100644 --- a/src/brpc/handshake/rdma_handshake.cpp +++ b/src/brpc/handshake/rdma_handshake.cpp @@ -72,36 +72,18 @@ void RdmaHandshakeAdapter::PrepareClientEce() { } } -handshake::HandshakeCodec RdmaHandshakeAdapter::MakeCodec( - ParsedHello* remote) { - handshake::HandshakeCodec codec{}; - codec.protocol_version = ProtocolVersion(); - codec.hello_frame = HelloFrameSpec(); - codec.ack_frame = RdmaAckFrameSpec(); - codec.build_hello = [this](bool enabled, std::string* payload) { - return BuildLocalHello(enabled, payload); - }; - codec.parse_hello = [this, remote](const std::string& payload) { - return ParseRemoteHello(payload, remote); - }; - codec.build_ack = [](bool enabled, std::string* payload) { - const uint32_t flags_be = butil::HostToNet32( - enabled ? HELLO_ACK_RDMA_OK : 0); - payload->assign(reinterpret_cast(&flags_be), - sizeof(flags_be)); - return handshake::STEP_OK; - }; - codec.parse_ack = [](const std::string& payload, bool* enabled) { - if (payload.size() != HELLO_ACK_LEN) { - errno = EPROTO; - return handshake::STEP_ERROR; - } - uint32_t flags_be = 0; - memcpy(&flags_be, payload.data(), sizeof(flags_be)); - *enabled = (butil::NetToHost32(flags_be) & HELLO_ACK_RDMA_OK) != 0; - return handshake::STEP_OK; - }; - return codec; +const handshake::FrameSpec& RdmaHandshakeAdapter::AckFrameSpec() const { + return RdmaAckFrameSpec(); +} + +handshake::StepResult RdmaHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + return BuildLocalHello(enabled, payload); +} + +handshake::StepResult RdmaHandshakeAdapter::ParseHello( + const std::string& payload) { + return ParseRemoteHello(payload, &_remote); } namespace v2_wire { @@ -418,20 +400,23 @@ static constexpr uint16_t V2_HELLO_VERSION_INVALID = std::numeric_limits::max(); static constexpr size_t V3_GID_LEN = 16; -static HandshakeCodec MakeRdmaFallbackCodec(int version) { - HandshakeCodec codec{}; - codec.protocol_version = version; - codec.hello_frame = rdma::RdmaHelloFrameSpec(version); - codec.ack_frame = rdma::RdmaAckFrameSpec(); - codec.parse_hello = [](const std::string&) { - return STEP_FALLBACK; - }; - codec.build_hello = [version](bool enabled, std::string* payload) { +class RdmaFallbackProtocol : public FallbackHandshakeProtocol { +public: + explicit RdmaFallbackProtocol(int version) : _version(version) {} + + int ProtocolVersion() const override { return _version; } + const FrameSpec& HelloFrameSpec() const override { + return rdma::RdmaHelloFrameSpec(_version); + } + const FrameSpec& AckFrameSpec() const override { + return rdma::RdmaAckFrameSpec(); + } + StepResult BuildHello(bool enabled, std::string* payload) override { if (enabled) { errno = EPROTO; return STEP_ERROR; } - if (version == 2) { + if (_version == 2) { payload->assign( rdma::HELLO_V2_MSG_LEN_MIN - rdma::HELLO_MAGIC_LEN - sizeof(uint16_t), @@ -453,24 +438,11 @@ static HandshakeCodec MakeRdmaFallbackCodec(int version) { return STEP_ERROR; } return STEP_OK; - }; - codec.build_ack = [](bool enabled, std::string* payload) { - const uint32_t flags_be = butil::HostToNet32( - enabled ? rdma::HELLO_ACK_RDMA_OK : 0); - payload->assign(reinterpret_cast(&flags_be), - sizeof(flags_be)); - return STEP_OK; - }; - codec.parse_ack = [](const std::string& payload, bool* enabled) { - if (payload.size() != rdma::HELLO_ACK_LEN) { - errno = EPROTO; - return STEP_ERROR; - } - *enabled = false; - return STEP_OK; - }; - return codec; -} + } + +private: + int _version; +}; HandshakeAdapter* GetRdmaServerHandshakeAdapter() { static RdmaServerHandshakeAdapter adapter; @@ -485,82 +457,81 @@ HandshakeSession* RdmaServerHandshakeAdapter::GetSession( StepResult RdmaServerHandshakeAdapter::RunFallbackServerHandshake( butil::IOBuf* source, Socket* socket) { IOBufHandshakeInput input(source); - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.codecs.push_back(MakeRdmaFallbackCodec(2)); - callbacks.codecs.push_back(MakeRdmaFallbackCodec(3)); - callbacks.input = &input; - callbacks.transport.prepare_resources = []() { return STEP_OK; }; - callbacks.transport.negotiate_resources = []() { return STEP_OK; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = []() {}; - callbacks.transport.on_failed = []() {}; - return GetSession(socket)->RunServer(callbacks); + RdmaFallbackProtocol v2(2); + RdmaFallbackProtocol v3(3); + std::vector protocols; + protocols.push_back(&v2); + protocols.push_back(&v3); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, &input, &transport, false); } #if BRPC_WITH_RDMA +class RdmaServerHandshakeTransport : public HandshakeTransport { +public: + RdmaServerHandshakeTransport( + RdmaTransport* transport, Socket* socket, butil::IOBuf* source) + : _transport(transport), _socket(socket), _source(source), + _protocol(NULL) {} + + void OnProtocolSelected(HandshakeProtocol* protocol) override { + _protocol = static_cast(protocol); + } + + StepResult PrepareResources() override { + if (_transport->PrepareUpgradeResources() == 0) { + return STEP_OK; + } + PLOG(WARNING) << "Fail to allocate rdma resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + + StepResult NegotiateResources() override { + CHECK(_protocol != NULL); + if (_transport->NegotiateUpgradeResources( + _protocol->remote(), true) == 0) { + return STEP_OK; + } + PLOG(WARNING) << "Fail to negotiate rdma resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + void OnFailed() override { _transport->DeactivateUpgrade(); } + + StepResult ValidateEstablished() override { + return _source->empty() ? STEP_OK : STEP_ERROR; + } + +private: + RdmaTransport* _transport; + Socket* _socket; + butil::IOBuf* _source; + rdma::RdmaHandshakeAdapter* _protocol; +}; + StepResult RdmaServerHandshakeAdapter::RunRdmaServerHandshake( butil::IOBuf* source, Socket* socket) { RdmaTransport* transport = RdmaTransport::Get(socket); CHECK(transport->GetRdmaEp() != NULL); - rdma::ParsedHello remote{}; std::vector > protocols = transport->CreateServerHandshakeAdapters(); - IOBufHandshakeInput input(source); - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.input = &input; + std::vector protocol_ptrs; + protocol_ptrs.reserve(protocols.size()); for (size_t i = 0; i < protocols.size(); ++i) { - HandshakeCodec codec = protocols[i]->MakeCodec(&remote); - const std::function parse_hello = - codec.parse_hello; - codec.parse_hello = [transport, parse_hello]( - const std::string& payload) { - const StepResult result = parse_hello(payload); - if (result == STEP_FALLBACK) { - transport->DeactivateUpgrade(); - } - return result; - }; - callbacks.codecs.push_back(codec); + protocol_ptrs.push_back(protocols[i].get()); } - callbacks.transport.prepare_resources = [&]() { - if (transport->PrepareUpgradeResources() < 0) { - PLOG(WARNING) - << "Fail to allocate rdma resources, fallback to tcp:" - << socket->description(); - transport->DeactivateUpgrade(); - return STEP_FALLBACK; - } - return STEP_OK; - }; - callbacks.transport.negotiate_resources = [&]() { - if (transport->NegotiateUpgradeResources(remote, true) < 0) { - PLOG(WARNING) - << "Fail to negotiate rdma resources, fallback to tcp:" - << socket->description(); - transport->DeactivateUpgrade(); - return STEP_FALLBACK; - } - return STEP_OK; - }; - callbacks.validate_established = [&]() { - if (!source->empty()) { - return STEP_ERROR; - } - return STEP_OK; - }; - callbacks.transport.set_high_speed_active = [transport]() { - transport->ActivateUpgrade(); - }; - callbacks.transport.set_tcp_active = [transport]() { - transport->DeactivateUpgrade(); - }; - callbacks.transport.on_failed = [transport]() { - transport->DeactivateUpgrade(); - }; - return GetSession(socket)->RunServer(callbacks); + IOBufHandshakeInput input(source); + RdmaServerHandshakeTransport participant(transport, socket, source); + return GetSession(socket)->RunServer( + protocol_ptrs, &input, &participant, false); } #endif diff --git a/src/brpc/handshake/rdma_handshake.h b/src/brpc/handshake/rdma_handshake.h index 7552454c34..80057e687e 100644 --- a/src/brpc/handshake/rdma_handshake.h +++ b/src/brpc/handshake/rdma_handshake.h @@ -71,16 +71,21 @@ struct HelloMessage { // RDMA adapters implement only protocol fields. HandshakeSession owns frame // I/O, length validation, ACK exchange and resource callback ordering. -class RdmaHandshakeAdapter { +class RdmaHandshakeAdapter : public handshake::HandshakeProtocol { public: RdmaHandshakeAdapter(RdmaEndpoint* ep, int version) - : _ep(ep), _version(version) {} - virtual ~RdmaHandshakeAdapter() = default; + : _ep(ep), _version(version), _remote() {} + ~RdmaHandshakeAdapter() override = default; - int ProtocolVersion() const { return _version; } - handshake::HandshakeCodec MakeCodec(ParsedHello* remote); + int ProtocolVersion() const override { return _version; } + const handshake::FrameSpec& AckFrameSpec() const override; + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseHello( + const std::string& payload) override; + const ParsedHello& remote() const { return _remote; } - virtual const handshake::FrameSpec& HelloFrameSpec() const = 0; + const handshake::FrameSpec& HelloFrameSpec() const override = 0; virtual handshake::StepResult BuildLocalHello( bool enabled, std::string* payload) = 0; virtual handshake::StepResult ParseRemoteHello( @@ -92,6 +97,7 @@ class RdmaHandshakeAdapter { RdmaEndpoint* _ep; int _version; + ParsedHello _remote; private: DISALLOW_COPY_AND_ASSIGN(RdmaHandshakeAdapter); diff --git a/src/brpc/handshake/ubshm_handshake.cpp b/src/brpc/handshake/ubshm_handshake.cpp index cc33ac81c5..a409915d0f 100644 --- a/src/brpc/handshake/ubshm_handshake.cpp +++ b/src/brpc/handshake/ubshm_handshake.cpp @@ -49,11 +49,9 @@ static const size_t MAGIC_LEN = 2; static const size_t HELLO_LEN = 64; static const size_t ACK_LEN = 4; #if BRPC_WITH_UBRING -static const uint16_t HELLO_VERSION = 2; +static const uint16_t HELLO_VERSION = 3; static const uint16_t IMPL_VERSION = 1; #endif // BRPC_WITH_UBRING -static const uint32_t ACK_OK = 0x1; - static const FrameSpec& HelloFrameSpec() { static const FrameSpec spec( MAGIC, MAGIC_LEN, HELLO_LEN, HELLO_LEN, FrameSpec::FIXED); @@ -96,6 +94,24 @@ void HelloMessage::Serialize(void* data) const { memcpy(current_pos, shm_name, SHM_MAX_NAME_BUFF_LEN); } +void HelloFormatExtension::Serialize(void* data) const { + char* current = static_cast(data); + const uint16_t length = butil::HostToNet16(extension_len); + const uint16_t format = butil::HostToNet16(format_id); + memcpy(current, &length, sizeof(length)); + memcpy(current + sizeof(length), &format, sizeof(format)); +} + +void HelloFormatExtension::Deserialize(const void* data) { + const char* current = static_cast(data); + uint16_t length; + uint16_t format; + memcpy(&length, current, sizeof(length)); + memcpy(&format, current + sizeof(length), sizeof(format)); + extension_len = butil::NetToHost16(length); + format_id = butil::NetToHost16(format); +} + void HelloMessage::Deserialize(const void* data) { const char* current_pos = static_cast(data); uint16_t net_msg_len; @@ -130,30 +146,65 @@ std::string HelloMessage::toString() const { return std::string(buf.data(), static_cast(n)); } -handshake::HandshakeCodec UBShmHandshakeAdapter::MakeCodec() const { - handshake::HandshakeCodec codec{}; - codec.protocol_version = 2; - codec.hello_frame = handshake::ubshm_wire::HelloFrameSpec(); - codec.ack_frame = handshake::ubshm_wire::AckFrameSpec(); - codec.build_ack = [](bool enabled, std::string* payload) { - const uint32_t flags_be = butil::HostToNet32( - enabled ? handshake::ubshm_wire::ACK_OK : 0); - payload->assign(reinterpret_cast(&flags_be), - sizeof(flags_be)); - return handshake::STEP_OK; - }; - codec.parse_ack = [](const std::string& payload, bool* enabled) { - if (payload.size() != handshake::ubshm_wire::ACK_LEN) { - errno = EPROTO; - return handshake::STEP_ERROR; - } - uint32_t flags_be = 0; - memcpy(&flags_be, payload.data(), sizeof(flags_be)); - *enabled = (butil::NetToHost32(flags_be) & - handshake::ubshm_wire::ACK_OK) != 0; - return handshake::STEP_OK; - }; - return codec; +const handshake::FrameSpec& UBShmHandshakeAdapter::HelloFrameSpec() const { + return handshake::ubshm_wire::HelloFrameSpec(); +} + +const handshake::FrameSpec& UBShmHandshakeAdapter::AckFrameSpec() const { + return handshake::ubshm_wire::AckFrameSpec(); +} + +const handshake::FrameSpec& +UBShmHandshakeAdapter::ExtensionFrameSpec() const { + static const handshake::FrameSpec spec( + NULL, 0, HelloFormatExtension::WIRE_SIZE, + HelloFormatExtension::WIRE_SIZE, handshake::FrameSpec::FIXED); + return spec; +} + +void UBShmHandshakeAdapter::ConfigureClientHello( + uint64_t len, const char* shm_name) { + _local_len = len; + _local_name = shm_name != NULL ? shm_name : ""; + _server_reply = false; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + const char* name = NULL; + if (enabled) { + name = _server_reply ? _remote.shm_name : _local_name.c_str(); + } + return BuildHello(enabled, enabled ? _local_len : 0, name, payload); +} + +handshake::StepResult UBShmHandshakeAdapter::ParseHello( + const std::string& payload) { + return ParseHello(payload, &_remote); +} + +handshake::StepResult UBShmHandshakeAdapter::BuildExtension( + bool enabled, std::string* payload) { + const HelloFormatExtension extension = { + HelloFormatExtension::WIRE_SIZE, + static_cast(enabled ? UBR_DATA_FORMAT_LEGACY_64 + : UBR_DATA_FORMAT_NONE)}; + payload->resize(HelloFormatExtension::WIRE_SIZE); + extension.Serialize(&(*payload)[0]); + return handshake::STEP_OK; +} + +handshake::StepResult UBShmHandshakeAdapter::ParseExtension( + const std::string& payload) { + if (payload.size() != HelloFormatExtension::WIRE_SIZE) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + HelloFormatExtension extension{}; + extension.Deserialize(payload.data()); + return extension.extension_len == HelloFormatExtension::WIRE_SIZE && + extension.format_id == UBR_DATA_FORMAT_LEGACY_64 + ? handshake::STEP_OK : handshake::STEP_FALLBACK; } handshake::StepResult UBShmHandshakeAdapter::BuildHello( @@ -190,7 +241,7 @@ handshake::StepResult UBShmHandshakeAdapter::ParseHello( return handshake::STEP_ERROR; } message->Deserialize(payload.data()); - if (message->msg_len < handshake::ubshm_wire::HELLO_LEN) { + if (message->msg_len != handshake::ubshm_wire::HELLO_LEN) { errno = EPROTO; return handshake::STEP_ERROR; } @@ -240,15 +291,16 @@ class UBShmServerHandshakeAdapter : public StandardHandshakeAdapter { }; -static HandshakeCodec MakeUBShmFallbackCodec() { - HandshakeCodec codec{}; - codec.protocol_version = 2; - codec.hello_frame = ubshm_wire::HelloFrameSpec(); - codec.ack_frame = ubshm_wire::AckFrameSpec(); - codec.parse_hello = [](const std::string&) { - return STEP_FALLBACK; - }; - codec.build_hello = [](bool enabled, std::string* payload) { +class UBShmFallbackProtocol : public FallbackHandshakeProtocol { +public: + int ProtocolVersion() const override { return 2; } + const FrameSpec& HelloFrameSpec() const override { + return ubshm_wire::HelloFrameSpec(); + } + const FrameSpec& AckFrameSpec() const override { + return ubshm_wire::AckFrameSpec(); + } + StepResult BuildHello(bool enabled, std::string* payload) override { if (enabled) { errno = EPROTO; return STEP_ERROR; @@ -258,24 +310,8 @@ static HandshakeCodec MakeUBShmFallbackCodec() { butil::RawPacker(&(*payload)[0]) .pack16(static_cast(ubshm_wire::HELLO_LEN)); return STEP_OK; - }; - codec.build_ack = [](bool enabled, std::string* payload) { - const uint32_t flags_be = butil::HostToNet32( - enabled ? ubshm_wire::ACK_OK : 0); - payload->assign(reinterpret_cast(&flags_be), - sizeof(flags_be)); - return STEP_OK; - }; - codec.parse_ack = [](const std::string& payload, bool* enabled) { - if (payload.size() != ubshm_wire::ACK_LEN) { - errno = EPROTO; - return STEP_ERROR; - } - *enabled = false; - return STEP_OK; - }; - return codec; -} + } +}; HandshakeAdapter* GetUBShmServerHandshakeAdapter() { static UBShmServerHandshakeAdapter adapter; @@ -290,71 +326,47 @@ HandshakeSession* UBShmServerHandshakeAdapter::GetSession( StepResult UBShmServerHandshakeAdapter::RunFallbackServerHandshake( butil::IOBuf* source, Socket* socket) { IOBufHandshakeInput input(source); - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.codecs.push_back(MakeUBShmFallbackCodec()); - callbacks.input = &input; - callbacks.transport.prepare_resources = []() { return STEP_OK; }; - callbacks.transport.negotiate_resources = []() { return STEP_OK; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = []() {}; - callbacks.transport.on_failed = []() {}; - return GetSession(socket)->RunServer(callbacks); + UBShmFallbackProtocol protocol; + std::vector protocols(1, &protocol); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, &input, &transport, false); } #if BRPC_WITH_UBRING -StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( - butil::IOBuf* source, Socket* socket) { - UBShmTransport* transport = UBShmTransport::Get(socket); - CHECK(transport->GetUBShmEp() != NULL); +class UBShmServerHandshakeTransport : public HandshakeTransport { +public: + UBShmServerHandshakeTransport( + UBShmTransport* transport, ubring::UBShmHandshakeAdapter* protocol, + Socket* socket, butil::IOBuf* source) + : _transport(transport), _protocol(protocol), _socket(socket), + _source(source) {} - ubring::HelloMessage remote{}; - ubring::UBShmHandshakeAdapter wire; - IOBufHandshakeInput input(source); - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.input = &input; - HandshakeCodec codec = wire.MakeCodec(); - codec.parse_hello = [&](const std::string& payload) { - const StepResult result = wire.ParseHello(payload, &remote); - if (result == STEP_OK || result == STEP_FALLBACK) { - LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) - << "server receive handshake message : " - << remote.toString(); - } - if (result == STEP_FALLBACK) { - transport->DeactivateUpgrade(); - } - return result; - }; - codec.build_hello = [&](bool enabled, std::string* payload) { - const uint64_t len = enabled - ? static_cast(ubring::FLAGS_data_queue_size) * - MB_TO_BYTE - : 0; - return wire.BuildHello( - enabled, len, enabled ? remote.shm_name : NULL, payload); - }; - callbacks.codecs.push_back(codec); - callbacks.transport.prepare_resources = [&]() { + void OnProtocolSelected(HandshakeProtocol*) override { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "server receive handshake message : " + << _protocol->remote().toString(); + } + + StepResult PrepareResources() override { if (!ubring::IsUBAvailable()) { - transport->DeactivateUpgrade(); + _transport->DeactivateUpgrade(); return STEP_FALLBACK; } + const ubring::HelloMessage& remote = _protocol->remote(); const size_t remote_name_len = strnlen(remote.shm_name, SHM_MAX_NAME_BUFF_LEN); ubring::SHM remote_trx_shm = { NULL, remote.len, 0, {0}, - static_cast(socket->fd())}; - memcpy(remote_trx_shm.name, remote.shm_name, - remote_name_len + 1); + static_cast(_socket->fd())}; + memcpy(remote_trx_shm.name, remote.shm_name, remote_name_len + 1); const size_t local_shm_len = static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; ubring::SHM local_trx_shm = { NULL, local_shm_len, 0, {0}, - static_cast(socket->fd())}; + static_cast(_socket->fd())}; char client_name[SHM_MAX_NAME_BUFF_LEN + 1]; memcpy(client_name, remote.shm_name, SHM_MAX_NAME_BUFF_LEN); client_name[SHM_MAX_NAME_BUFF_LEN] = '\0'; @@ -366,37 +378,54 @@ StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( local_trx_shm.name, SHM_MAX_NAME_BUFF_LEN, "%s_%s", client_name, SERVER_SHM_NAME_SUFFIX); if (UNLIKELY(result < 0)) { - transport->DeactivateUpgrade(); + _transport->DeactivateUpgrade(); return STEP_FALLBACK; } - if (transport->PrepareServerUpgradeResources( + if (_transport->PrepareServerUpgradeResources( &remote_trx_shm, &local_trx_shm) < 0) { LOG(WARNING) << "Fail to allocate ub resources, fallback to tcp:" - << socket->description(); - transport->DeactivateUpgrade(); + << _socket->description(); + _transport->DeactivateUpgrade(); return STEP_FALLBACK; } return STEP_OK; - }; - callbacks.transport.negotiate_resources = []() { return STEP_OK; }; - callbacks.validate_established = [&]() { - if (!source->empty()) { - return STEP_ERROR; - } - return STEP_OK; - }; - callbacks.transport.set_high_speed_active = [transport]() { - transport->ActivateUpgrade(); - }; - callbacks.transport.set_tcp_active = [transport]() { - transport->DeactivateUpgrade(); - }; - callbacks.transport.on_failed = [transport]() { - transport->DeactivateUpgrade(); - }; - const StepResult result = GetSession(socket)->RunServer(callbacks); + } + + StepResult NegotiateResources() override { return STEP_OK; } + void OnEstablished() override { _transport->ActivateUpgrade(); } + void OnFallback() override { _transport->DeactivateUpgrade(); } + void OnFailed() override { _transport->DeactivateUpgrade(); } + + StepResult ValidateEstablished() override { + return _source->empty() ? STEP_OK : STEP_ERROR; + } + +private: + UBShmTransport* _transport; + ubring::UBShmHandshakeAdapter* _protocol; + Socket* _socket; + butil::IOBuf* _source; +}; + +StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( + butil::IOBuf* source, Socket* socket) { + UBShmTransport* transport = UBShmTransport::Get(socket); + CHECK(transport->GetUBShmEp() != NULL); + + ubring::UBShmHandshakeAdapter wire; + const uint64_t local_shm_len = + static_cast(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + wire.ConfigureServerReply(local_shm_len); + IOBufHandshakeInput input(source); + std::vector protocols(1, &wire); + UBShmServerHandshakeTransport participant( + transport, &wire, socket, source); + const StepResult result = GetSession(socket)->RunServer( + protocols, &input, &participant, false); if (result == STEP_OK) { + transport->GetUBShmEp()->SetNegotiatedDataFormat( + ubring::UBR_DATA_FORMAT_LEGACY_64); transport->FinishUpgrade(); LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) << "Server handshake ends (use ubring) on " diff --git a/src/brpc/handshake/ubshm_handshake.h b/src/brpc/handshake/ubshm_handshake.h index 1bf1bb9d9b..a41af31243 100644 --- a/src/brpc/handshake/ubshm_handshake.h +++ b/src/brpc/handshake/ubshm_handshake.h @@ -47,7 +47,7 @@ namespace ubring { DECLARE_int32(data_queue_size); DECLARE_bool(ub_trace_verbose); -// UBSHM v2 wire payload. HandshakeSession owns framing and ACK exchange; +// UBSHM v3 wire payload. HandshakeSession owns framing and ACK exchange; // this type and UBShmHandshakeAdapter only handle protocol fields. struct HelloMessage { void Serialize(void* data) const; @@ -61,11 +61,46 @@ struct HelloMessage { char shm_name[SHM_MAX_NAME_BUFF_LEN]; }; -class UBShmHandshakeAdapter { +enum UbrDataFormat { + UBR_DATA_FORMAT_NONE = 0, + UBR_DATA_FORMAT_LEGACY_64 = 1, +}; + +struct HelloFormatExtension { + static const uint16_t WIRE_SIZE = 4; + uint16_t extension_len; + uint16_t format_id; + + void Serialize(void* data) const; + void Deserialize(const void* data); +}; + +class UBShmHandshakeAdapter : public handshake::HandshakeProtocol { public: - UBShmHandshakeAdapter() = default; + UBShmHandshakeAdapter() + : _local_len(0), _server_reply(false), _remote() {} + + int ProtocolVersion() const override { return 3; } + const handshake::FrameSpec& HelloFrameSpec() const override; + const handshake::FrameSpec& AckFrameSpec() const override; + handshake::StepResult BuildHello( + bool enabled, std::string* payload) override; + handshake::StepResult ParseHello( + const std::string& payload) override; + bool HasExtension() const override { return true; } + const handshake::FrameSpec& ExtensionFrameSpec() const override; + handshake::StepResult BuildExtension( + bool enabled, std::string* payload) override; + handshake::StepResult ParseExtension( + const std::string& payload) override; + + void ConfigureClientHello(uint64_t len, const char* shm_name); + void ConfigureServerReply(uint64_t len) { + _local_len = len; + _server_reply = true; + } + const HelloMessage& remote() const { return _remote; } - handshake::HandshakeCodec MakeCodec() const; handshake::StepResult BuildHello( bool enabled, uint64_t len, const char* shm_name, std::string* payload) const; @@ -74,6 +109,10 @@ class UBShmHandshakeAdapter { private: bool NegotiationValid(const HelloMessage& message) const; + uint64_t _local_len; + std::string _local_name; + bool _server_reply; + HelloMessage _remote; DISALLOW_COPY_AND_ASSIGN(UBShmHandshakeAdapter); }; diff --git a/src/brpc/input_messenger.h b/src/brpc/input_messenger.h index 97b163365b..2464034885 100644 --- a/src/brpc/input_messenger.h +++ b/src/brpc/input_messenger.h @@ -30,6 +30,11 @@ namespace brpc { namespace rdma { class RdmaEndpoint; } + +namespace urma { +class UrmaEndpoint; +} + namespace ubring { class UBShmEndpoint; } @@ -100,6 +105,7 @@ friend class TcpTransport; friend class RdmaTransport; friend class AdapterTransport; friend class rdma::RdmaEndpoint; +friend class urma::UrmaEndpoint; friend class ubring::UBShmEndpoint; friend class InputMessengerProcessor; public: diff --git a/src/brpc/input_messenger_processor.cpp b/src/brpc/input_messenger_processor.cpp index f23a423c8c..febd1b580c 100644 --- a/src/brpc/input_messenger_processor.cpp +++ b/src/brpc/input_messenger_processor.cpp @@ -36,6 +36,7 @@ static const char* StreamTypeName(InputMessengerProcessor::StreamType type) { case InputMessengerProcessor::STREAM_NONE: return "none"; case InputMessengerProcessor::STREAM_TCP_FD: return "tcp_fd"; case InputMessengerProcessor::STREAM_RDMA_QP: return "rdma_qp"; + case InputMessengerProcessor::STREAM_URMA_JETTY: return "urma_jetty"; } return "unknown"; } @@ -281,7 +282,8 @@ int InputMessengerProcessor::ProcessNewMessage(ssize_t bytes, bool read_eof, // method for processing messages may call synchronization primitives, // causing the polling bthread to be scheduled out. if (_socket->_socket_mode == SOCKET_MODE_RDMA || - _socket->_socket_mode == SOCKET_MODE_UBRING) { + _socket->_socket_mode == SOCKET_MODE_UBRING || + _socket->_socket_mode == SOCKET_MODE_URMA) { _socket->_transport->QueueMessage(last_msg, &num_bthread_created, true); } if (num_bthread_created) { diff --git a/src/brpc/input_messenger_processor.h b/src/brpc/input_messenger_processor.h index 7b94087d65..cdf9e203b8 100644 --- a/src/brpc/input_messenger_processor.h +++ b/src/brpc/input_messenger_processor.h @@ -39,6 +39,7 @@ class InputMessengerProcessor { STREAM_NONE, STREAM_TCP_FD, STREAM_RDMA_QP, + STREAM_URMA_JETTY, }; InputMessengerProcessor() diff --git a/src/brpc/policy/http2_rpc_protocol.cpp b/src/brpc/policy/http2_rpc_protocol.cpp index b20a05727c..6d9b2b8714 100644 --- a/src/brpc/policy/http2_rpc_protocol.cpp +++ b/src/brpc/policy/http2_rpc_protocol.cpp @@ -29,6 +29,7 @@ DECLARE_int32(http_verbose_max_body_length); DECLARE_int32(health_check_interval); DECLARE_bool(usercode_in_pthread); DECLARE_int64(socket_max_unwritten_bytes); +DECLARE_uint32(http_max_header_count); namespace policy { @@ -729,6 +730,11 @@ H2ParseResult H2StreamContext::OnHeaders( << ", stream_id=" << frame_head.stream_id; return MakeH2Error(H2_PROTOCOL_ERROR); } + // The whole block went through the decoder, the connection is in a + // consistent state again and only this stream needs to be reset. + if (_rejected_error != H2_NO_ERROR) { + return MakeH2Error(_rejected_error, stream_id()); + } if (frame_head.flags & H2_FLAGS_END_STREAM) { return OnEndStream(); } @@ -791,6 +797,10 @@ H2ParseResult H2StreamContext::OnContinuation( << ", stream_id=" << frame_head.stream_id; return MakeH2Error(H2_PROTOCOL_ERROR); } + // See the same check in H2StreamContext::OnHeaders(). + if (_rejected_error != H2_NO_ERROR) { + return MakeH2Error(_rejected_error, stream_id()); + } if (_stream_ended) { return OnEndStream(); } @@ -1294,7 +1304,8 @@ H2StreamContext::H2StreamContext(bool read_body_progressively) , _remote_window_left(0) , _deferred_window_update(0) , _correlation_id(INVALID_BTHREAD_ID.value) - , _decoded_header_list_size(0) { + , _decoded_header_list_size(0) + , _rejected_error(H2_NO_ERROR) { header().set_version(2, 0); #ifndef NDEBUG get_h2_bvars()->h2_stream_context_count << 1; @@ -1349,6 +1360,14 @@ int H2StreamContext::ConsumeHeaders(butil::IOBufBytesIterator& it) { << max_header_list_size << ", stream_id=" << _stream_id; return -1; } + if (_rejected_error != H2_NO_ERROR) { + // The stream is already refused, keep feeding the decoder so that + // the dynamic table stays in sync with the peer, but stop spending + // memory on fields nobody is going to read. A peer that keeps + // piling them up still runs into max_header_list_size above, which + // escalates to a connection error as it has to. + continue; + } const char* const name = pair.name.c_str(); bool matched = false; if (name[0] == ':') { // reserved names @@ -1364,10 +1383,12 @@ int H2StreamContext::ConsumeHeaders(butil::IOBufBytesIterator& it) { matched = true; HttpMethod method; if (!Str2HttpMethod(pair.value.c_str(), &method)) { - LOG(ERROR) << "Invalid method=" << pair.value; - return -1; + LOG(ERROR) << "Invalid method=" << pair.value + << ", stream_id=" << _stream_id; + _rejected_error = H2_PROTOCOL_ERROR; + } else { + h.set_method(method); } - h.set_method(method); } break; case 'p': @@ -1380,11 +1401,16 @@ int H2StreamContext::ConsumeHeaders(butil::IOBufBytesIterator& it) { // would take the whole header block, since HPACK does not // order pseudo-headers and :method may not have arrived. if (pair.value != "*" && (pair.value.empty() || pair.value[0] != '/')) { - LOG(ERROR) << "Invalid path=" << pair.value; - return -1; + LOG(ERROR) << "Invalid path=" << pair.value + << ", stream_id=" << _stream_id; + _rejected_error = H2_PROTOCOL_ERROR; + } else if (h.uri().SetH2Path(pair.value) != 0) { + // Including path/query/fragment. The only way this + // fails is too many query parameters. + LOG(ERROR) << h.uri().status().error_cstr() + << ", stream_id=" << _stream_id; + _rejected_error = H2_ENHANCE_YOUR_CALM; } - // Including path/query/fragment - h.uri().SetH2Path(pair.value); } break; case 's': @@ -1396,24 +1422,34 @@ int H2StreamContext::ConsumeHeaders(butil::IOBufBytesIterator& it) { char* endptr = nullptr; const int sc = strtol(pair.value.c_str(), &endptr, 10); if (*endptr != '\0') { - LOG(ERROR) << "Invalid status=" << pair.value; - return -1; + LOG(ERROR) << "Invalid status=" << pair.value + << ", stream_id=" << _stream_id; + _rejected_error = H2_PROTOCOL_ERROR; + } else { + h.set_status_code(sc); } - h.set_status_code(sc); } break; default: break; } if (!matched) { - LOG(ERROR) << "Unknown name=`" << name << '\''; - return -1; + LOG(ERROR) << "Unknown pseudo-header=`" << name + << "', stream_id=" << _stream_id; + _rejected_error = H2_PROTOCOL_ERROR; } } else if (name[0] == 'c' && strcmp(name + 1, /*c*/"ontent-type") == 0) { h.set_content_type(pair.value); } else { h.AppendHeader(pair.name, pair.value); + if (FLAGS_http_max_header_count > 0 && + h.HeaderCount() > FLAGS_http_max_header_count) { + LOG(ERROR) << "Too many headers, max=" + << FLAGS_http_max_header_count + << ", stream_id=" << _stream_id; + _rejected_error = H2_ENHANCE_YOUR_CALM; + } } if (FLAGS_http_verbose) { diff --git a/src/brpc/policy/http2_rpc_protocol.h b/src/brpc/policy/http2_rpc_protocol.h index 6c58f71aa5..64c9864d68 100644 --- a/src/brpc/policy/http2_rpc_protocol.h +++ b/src/brpc/policy/http2_rpc_protocol.h @@ -234,7 +234,9 @@ class H2StreamContext : public HttpContext { // Decode headers in HPACK from *it and set into this->header(). The input // does not need to complete. - // Returns 0 on success, -1 otherwise. + // Returns 0 on success, -1 on a connection-level error. A message that is + // merely malformed or unacceptable does not fail here, it sets + // `_rejected_error` instead, see the comment on that field. int ConsumeHeaders(butil::IOBufBytesIterator& it); H2ParseResult OnEndStream(); @@ -281,6 +283,19 @@ friend class H2Context; // (name + value + 32 per field, RFC 7540 section 10.5.1), checked // against the local max_header_list_size in ConsumeHeaders(). uint64_t _decoded_header_list_size; + // Set when this message must be refused although the connection itself is + // still healthy: it is malformed (invalid or unknown pseudo-header, RFC + // 9113 section 8.1.1 mandates a stream error of type PROTOCOL_ERROR) or it + // violates a local limit (too many headers, too many query parameters in + // :path). Only the stream is reset so that the other streams keep working, + // but the error cannot be raised where it is detected: HPACK keeps a + // dynamic table per connection, so leaving the rest of the block undecoded + // would desynchronize it from the encoding table of the peer and corrupt + // every header block that follows. RFC 9113 section 10.5.1: "The field + // block MUST be processed to ensure a consistent connection state, unless + // the connection is closed." Hence the rejection is remembered here and + // turned into a RST_STREAM once END_HEADERS is reached. + H2Error _rejected_error; butil::IOBuf _remaining_header_fragment; // Request body which cannot be sent yet due to remote flow control. // Accessed under H2Context::_stream_mutex. diff --git a/src/brpc/policy/mysql/mysql_reply.cpp b/src/brpc/policy/mysql/mysql_reply.cpp index 2f3a97078c..1a277339f9 100644 --- a/src/brpc/policy/mysql/mysql_reply.cpp +++ b/src/brpc/policy/mysql/mysql_reply.cpp @@ -200,6 +200,11 @@ ParseError MysqlReply::ConsumePartialIOBuf(butil::IOBuf& buf, } uint8_t header[4 + 1]; // use the extra byte to judge message type const uint8_t* p = (const uint8_t*)buf.fetch(header, sizeof(header)); + if (_type == MYSQL_RSP_UNKNOWN && + (p == nullptr || mysql_uint3korr(p) == 0)) { + LOG(ERROR) << "Invalid mysql packet with empty payload"; + return PARSE_ERROR_ABSOLUTELY_WRONG; + } uint8_t type = (_type == MYSQL_RSP_UNKNOWN) ? p[4] : (uint8_t)_type; // During the connection (auth) phase the server may send an AuthMoreData // packet (first byte 0x01) as part of the caching_sha2_password exchange @@ -243,6 +248,11 @@ ParseError MysqlReply::ConsumePartialIOBuf(butil::IOBuf& buf, butil::IOBuf discard; buf.cutn(&discard, amd_total); const uint8_t* p2 = (const uint8_t*)buf.fetch(header, sizeof(header)); + if (p2 == nullptr || mysql_uint3korr(p2) == 0) { + LOG(ERROR) << "Invalid mysql packet with empty payload after " + "fast-auth marker"; + return PARSE_ERROR_ABSOLUTELY_WRONG; + } type = p2[4]; } else { _type = MYSQL_RSP_AUTH_MORE_DATA; diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index 1b4b802648..2db9179897 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -37,8 +37,7 @@ DECLARE_int32(task_group_ntags); namespace brpc { namespace rdma { -extern ibv_cq *(*IbvCreateCq)(ibv_context *, int, void *, ibv_comp_channel *, - int); +extern ibv_cq* (*IbvCreateCq)(ibv_context*, int, void*, ibv_comp_channel*, int); extern int (*IbvDestroyCq)(ibv_cq*); extern ibv_comp_channel* (*IbvCreateCompChannel)(ibv_context*); extern int (*IbvDestroyCompChannel)(ibv_comp_channel*); @@ -108,15 +107,29 @@ RdmaResource::~RdmaResource() { } } -RdmaEndpoint::RdmaEndpoint(Socket *s) - : _socket(s), _resource(nullptr), - _send_cq_events(0), _recv_cq_events(0), _cq_sid(INVALID_SOCKET_ID), - _sq_size(FLAGS_rdma_sq_size), _rq_size(FLAGS_rdma_rq_size), - _remote_recv_block_size(0), _accumulated_ack(0), _unsolicited(0), - _unsolicited_bytes(0), _sq_current(0), _sq_unsignaled(0), _sq_sent(0), - _rq_received(0), _local_window_capacity(0), _remote_window_capacity(0), - _sq_imm_window_size(0), _remote_rq_window_size(0), _sq_window_size(0), - _new_rq_wrs(0) { +RdmaEndpoint::RdmaEndpoint(Socket* s) + : _socket(s) + , _resource(nullptr) + , _send_cq_events(0) + , _recv_cq_events(0) + , _cq_sid(INVALID_SOCKET_ID) + , _sq_size(FLAGS_rdma_sq_size) + , _rq_size(FLAGS_rdma_rq_size) + , _remote_recv_block_size(0) + , _accumulated_ack(0) + , _unsolicited(0) + , _unsolicited_bytes(0) + , _sq_current(0) + , _sq_unsignaled(0) + , _sq_sent(0) + , _rq_received(0) + , _local_window_capacity(0) + , _remote_window_capacity(0) + , _sq_imm_window_size(0) + , _remote_rq_window_size(0) + , _sq_window_size(0) + , _new_rq_wrs(0) +{ if (_sq_size < MIN_QP_SIZE) { _sq_size = MIN_QP_SIZE; } @@ -1347,9 +1360,9 @@ void RdmaEndpoint::PollingModeRelease(bthread_tag_t tag) { void RdmaEndpoint::PollerAddCqSid() { auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num; - auto &group = _poller_groups[bthread_self_tag()]; - auto &pollers = group.pollers; - auto &poller = pollers[index]; + auto& group = _poller_groups[bthread_self_tag()]; + auto& pollers = group.pollers; + auto& poller = pollers[index]; if (INVALID_SOCKET_ID != _cq_sid) { poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::ADD}); } @@ -1357,9 +1370,9 @@ void RdmaEndpoint::PollerAddCqSid() { void RdmaEndpoint::PollerRemoveCqSid() { auto index = butil::fmix32(_cq_sid) % FLAGS_rdma_poller_num; - auto &group = _poller_groups[bthread_self_tag()]; - auto &pollers = group.pollers; - auto &poller = pollers[index]; + auto& group = _poller_groups[bthread_self_tag()]; + auto& pollers = group.pollers; + auto& poller = pollers[index]; if (INVALID_SOCKET_ID != _cq_sid) { poller.op_queue.Enqueue(CqSidOp{_cq_sid, CqSidOp::REMOVE}); } diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp index baf8b8fa94..5a29ae236b 100644 --- a/src/brpc/server.cpp +++ b/src/brpc/server.cpp @@ -138,6 +138,7 @@ ServerOptions::ServerOptions() , server_owns_interceptor(false) , num_threads(8) , max_concurrency(0) + , max_connections(0) , session_local_data_factory(nullptr) , reserved_session_local_data(0) , thread_local_data_factory(nullptr) @@ -1167,7 +1168,8 @@ int Server::StartInternal(const butil::EndPoint& endpoint, // Pass ownership of `sockfd' to `_am' if (_am->StartAccept(sockfd, _options.idle_timeout_sec, _default_ssl_ctx, - _options.force_ssl) != 0) { + _options.force_ssl, + _options.max_connections) != 0) { LOG(ERROR) << "Fail to start acceptor"; return -1; } @@ -1209,7 +1211,8 @@ int Server::StartInternal(const butil::EndPoint& endpoint, // Pass ownership of `sockfd' to `_internal_am' if (_internal_am->StartAccept(sockfd, _options.idle_timeout_sec, _default_ssl_ctx, - false) != 0) { + false, + 0) != 0) { LOG(ERROR) << "Fail to start internal_acceptor"; return -1; } @@ -1791,8 +1794,11 @@ google::protobuf::Service* Server::FindServiceByName( void Server::GetStat(ServerStatistics* stat) const { stat->connection_count = 0; + stat->rejected_connection_count = 0; if (_am) { stat->connection_count += _am->ConnectionCount(); + stat->rejected_connection_count += + _am->RejectedConnectionCount(); } if (_internal_am) { stat->connection_count += _internal_am->ConnectionCount(); @@ -1801,6 +1807,15 @@ void Server::GetStat(ServerStatistics* stat) const { stat->builtin_service_count = builtin_service_count(); } +int Server::SetMaxConnections(size_t max_connections) { + if (!IsRunning() || _am == nullptr) { + LOG(WARNING) << "SetMaxConnections requires a running Server"; + return -1; + } + _am->SetMaxConnections(max_connections); + return 0; +} + void Server::ListServices(std::vector *services) { if (!services) { return; diff --git a/src/brpc/server.h b/src/brpc/server.h index 6e7d2b2b17..3c68641be5 100644 --- a/src/brpc/server.h +++ b/src/brpc/server.h @@ -132,6 +132,17 @@ struct ServerOptions { // Default: 0 (unlimited) int max_concurrency; + // Maximum number of simultaneous connections on the public listener, + // shared by all protocols and services, including builtin services. + // Idle connections and connections awaiting TLS handshakes count too. + // Excess connections are closed before protocol parsing or TLS, without + // a response. The internal listener has a separate, unlimited count. + // Other Server instances have independent limits, even if they share + // the same service object on different ports. + // Use Server::SetMaxConnections() to update the limit at runtime. + // Default: 0 (unlimited) + size_t max_connections; + // Default value of method-level max concurrencies, // Overridable by Server.MaxConcurrencyOf(). AdaptiveMaxConcurrency method_max_concurrency; @@ -227,7 +238,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; @@ -309,7 +320,10 @@ struct ServerOptions { // This struct is originally designed to contain basic statistics of the // server. But bvar contains more stats and is more convenient. struct ServerStatistics { + // Total connections on the public and internal listeners. size_t connection_count; + // Cumulative connections rejected by the public listener's limit. + size_t rejected_connection_count; int user_service_count; int builtin_service_count; }; @@ -546,6 +560,13 @@ class Server { // Get statistics of this server void GetStat(ServerStatistics* stat) const; + // Atomically update the connection limit of the running public listener. + // Existing connections are not closed when the limit is lowered. + // Set to 0 to disable the limit. Does not change options().max_connections, + // which records the startup setting. + // Returns 0 on success, -1 if this Server is not running. + int SetMaxConnections(size_t max_connections); + // Get the options passed to Start(). const ServerOptions& options() const { return _options; } diff --git a/src/brpc/socket.cpp b/src/brpc/socket.cpp index 283473df9e..57cfd2296a 100644 --- a/src/brpc/socket.cpp +++ b/src/brpc/socket.cpp @@ -742,6 +742,7 @@ int Socket::OnCreated(const SocketOptions& options) { _tos = 0; _remote_side = options.remote_side; _local_side = options.local_side; + _bind_local_side = options.local_side; _device_name = options.device_name; _on_edge_triggered_events = options.on_edge_triggered_events; _need_on_edge_trigger = options.need_on_edge_trigger; @@ -794,8 +795,28 @@ int Socket::OnCreated(const SocketOptions& options) { _unwritten_bytes.store(0, butil::memory_order_relaxed); _keepalive_options = options.keepalive_options; _tcp_user_timeout_ms = options.tcp_user_timeout_ms; + _http_request_method = HTTP_METHOD_GET; CHECK(nullptr == _write_head.load(butil::memory_order_relaxed)); _is_write_shutdown = false; + // EndPoint::port is a type tag for IPv6/UDS, not a network port. + // Normalize only IP ports and leave Unix-domain addresses intact. + if (!butil::is_endpoint_extended(_bind_local_side)) { + _bind_local_side.port = 0; + } else if (butil::get_endpoint_type(_bind_local_side) == AF_INET6) { + sockaddr_storage addr{}; + socklen_t addr_size = 0; + if (butil::endpoint2sockaddr(_bind_local_side, &addr, &addr_size) != 0) { + SetFailed(EINVAL, "Fail to get client binding sockaddr from %s", + butil::endpoint2str(_bind_local_side).c_str()); + return -1; + } + reinterpret_cast(&addr)->sin6_port = 0; + if (butil::sockaddr2endpoint(&addr, addr_size, &_bind_local_side) != 0) { + SetFailed(ENOMEM, "Fail to create client binding endpoint from %s", + butil::endpoint2str(_bind_local_side).c_str()); + return -1; + } + } int fd = options.fd; if (!ValidFileDescriptor(fd) && options.connect_on_create) { // Connect on create. @@ -1014,6 +1035,7 @@ int Socket::WaitAndReset(int32_t expected_nref) { } _transport->Reset(expected_nref); + // Clear the runtime endpoint, keeping the configured binding intact. _local_side = butil::EndPoint(); if (_ssl_session) { SSL_free(_ssl_session); @@ -1264,8 +1286,8 @@ int Socket::Connect(const timespec* abstime, _ssl_state = SSL_OFF; } struct sockaddr_storage serv_addr; - socklen_t addr_size = 0; - if (butil::endpoint2sockaddr(remote_side(), &serv_addr, &addr_size) != 0) { + socklen_t serv_addr_size = 0; + if (butil::endpoint2sockaddr(remote_side(), &serv_addr, &serv_addr_size) != 0) { PLOG(ERROR) << "Fail to get sockaddr"; return -1; } @@ -1293,19 +1315,20 @@ int Socket::Connect(const timespec* abstime, return -1; #endif } - if (local_side().ip != butil::IP_ANY) { - struct sockaddr_storage cli_addr; - if (butil::endpoint2sockaddr(local_side(), &cli_addr, &addr_size) != 0) { + // Use the configured endpoint, not the runtime endpoint from getsockname(). + if (butil::is_endpoint_extended(_bind_local_side) || _bind_local_side.ip != butil::IP_ANY) { + struct sockaddr_storage cli_addr{}; + socklen_t cli_addr_size = 0; + if (butil::endpoint2sockaddr(_bind_local_side, &cli_addr, &cli_addr_size) != 0) { PLOG(ERROR) << "Fail to get client sockaddr"; return -1; } - if (::bind(sockfd, (struct sockaddr*)&cli_addr, addr_size) != 0) { - PLOG(ERROR) << "Fail to bind client socket, errno=" << strerror(errno); + if (::bind(sockfd, (struct sockaddr*)&cli_addr, cli_addr_size) != 0) { + PLOG(ERROR) << "Fail to bind client socket"; return -1; } } - const int rc = ::connect( - sockfd, (struct sockaddr*)&serv_addr, addr_size); + const int rc = ::connect(sockfd, (struct sockaddr*)&serv_addr, serv_addr_size); if (rc != 0 && errno != EINPROGRESS) { PLOG(WARNING) << "Fail to connect to " << remote_side(); return -1; @@ -2796,7 +2819,10 @@ int Socket::GetPooledSocket(SocketUniquePtr* pooled_socket) { if (socket_pool == nullptr) { SocketOptions opt; opt.remote_side = remote_side(); - opt.local_side = butil::EndPoint(local_side().ip, 0); + // Propagate the configured client binding (source IP + device) to + // pooled sub-sockets so that they keep the same binding policy. + opt.local_side = _bind_local_side; + opt.device_name = _device_name; opt.user = user(); opt.on_edge_triggered_events = _on_edge_triggered_events; opt.need_on_edge_trigger = _need_on_edge_trigger; @@ -2899,7 +2925,10 @@ int Socket::GetShortSocket(SocketUniquePtr* short_socket) { SocketId id; SocketOptions opt; opt.remote_side = remote_side(); - opt.local_side = butil::EndPoint(local_side().ip, 0); + // Propagate the configured client binding (source IP + device) to short + // sub-sockets so that they keep the same binding policy. + opt.local_side = _bind_local_side; + opt.device_name = _device_name; opt.user = user(); opt.on_edge_triggered_events = _on_edge_triggered_events; opt.need_on_edge_trigger = _need_on_edge_trigger; diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 579bc9cb79..194b8fd2ab 100644 --- a/src/brpc/socket.h +++ b/src/brpc/socket.h @@ -57,9 +57,18 @@ class ChannelBalancer; namespace rdma { class RdmaEndpoint; } + +namespace urma { +class UrmaEndpoint; +class UrmaConnect; +class UrmaHandshakeClientV2; +class UrmaHandshakeServerV2; +class UrmaHandshakeClientV3; +class UrmaHandshakeServerV3; +} + namespace ubring { class UBShmEndpoint; - class UBConnect; } namespace handshake { class SocketHandshakeIO; @@ -262,6 +271,9 @@ struct SocketOptions { // user->BeforeRecycle() before recycling. int fd{-1}; butil::EndPoint remote_side; + // Client source address. For IPv4/IPv6, the port is ignored for binding + // and the OS allocates a source port. IPv4 IP_ANY disables explicit binding. + // Unix-domain addresses are preserved, including their paths. butil::EndPoint local_side; std::string device_name; // If `connect_on_create' is true and `fd' is less than 0, @@ -328,8 +340,13 @@ friend class policy::RtmpContext; friend class schan::ChannelBalancer; friend class rdma::RdmaEndpoint; friend class ubring::UBShmEndpoint; -friend class ubring::UBConnect; friend class UBShmTransport; +friend class urma::UrmaEndpoint; +friend class urma::UrmaConnect; +friend class urma::UrmaHandshakeClientV2; +friend class urma::UrmaHandshakeServerV2; +friend class urma::UrmaHandshakeClientV3; +friend class urma::UrmaHandshakeServerV3; friend class HealthCheckTask; friend class OnAppHealthCheckDone; friend class HealthCheckManager; @@ -342,6 +359,7 @@ friend class AdapterTransport; friend class handshake::SocketHandshakeIO; friend class TcpTransport; friend class RdmaTransport; +friend class UrmaTransport; friend class TransportFactory; class SharedPart; struct WriteRequest; @@ -877,9 +895,15 @@ friend class TransportFactory; // Address of peer. Initialized by SocketOptions.remote_side. butil::EndPoint _remote_side; - // Address of self. Initialized in ResetFileDescriptor(). + // Runtime local endpoint. Updated in ResetFileDescriptor() and cleared + // in WaitAndReset(). butil::EndPoint _local_side; + // Client binding address from SocketOptions.local_side, preserved across + // health-check/revive. IPv4/IPv6 network ports are normalized to 0; + // Unix-domain addresses and extended endpoint type tags are preserved. + butil::EndPoint _bind_local_side; + // The device name of the client's network adapter. std::string _device_name; 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 3355e0764d..4c450bca85 100644 --- a/src/brpc/transport_factory.cpp +++ b/src/brpc/transport_factory.cpp @@ -18,27 +18,34 @@ #include "brpc/transport_factory.h" #include "brpc/adapter_transport.h" #include "brpc/rdma_transport.h" +#include "brpc/tcp_transport.h" #include "brpc/ubshm_transport.h" +#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 +53,22 @@ std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { return std::unique_ptr(new AdapterTransport(mode)); } #if BRPC_WITH_RDMA - else if (mode == SOCKET_MODE_RDMA) { + if (mode == SOCKET_MODE_RDMA) { return std::unique_ptr(new AdapterTransport(mode)); } #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 AdapterTransport(mode)); } #endif - else { - LOG(ERROR) << "socket_mode set error"; - return nullptr; - } + LOG(ERROR) << "Unknown transport type " << mode; + return nullptr; } -} // namespace brpc + +} // namespace brpc diff --git a/src/brpc/transport_factory.h b/src/brpc/transport_factory.h index add249c438..0cb6f355b2 100644 --- a/src/brpc/transport_factory.h +++ b/src/brpc/transport_factory.h @@ -22,14 +22,16 @@ #include "brpc/transport.h" namespace brpc { -// Creates the top-level AdapterTransport for all socket modes. The adapter -// selects TcpTransport or a concrete accelerated Transport internally. +// Creates AdapterTransport for TCP, RDMA, and UBSHM sockets. URMA currently +// uses its concrete transport directly. 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 +} // namespace brpc + +#endif // BRPC_TRANSPORT_FACTORY_H diff --git a/src/brpc/transport_handshake.cpp b/src/brpc/transport_handshake.cpp index f97d7b7741..c99a2148a6 100644 --- a/src/brpc/transport_handshake.cpp +++ b/src/brpc/transport_handshake.cpp @@ -18,13 +18,48 @@ #include "brpc/transport_handshake.h" #include +#include #include "butil/logging.h" #include "butil/object_pool.h" +#include "butil/sys_byteorder.h" namespace brpc { namespace handshake { +namespace { + +const size_t COMMON_ACK_SIZE = sizeof(uint32_t); +const uint32_t COMMON_ACK_OK = 0x1; + +} // namespace + +StepResult HandshakeProtocol::BuildAck(bool enabled, std::string* payload) { + CHECK(payload != NULL); + const uint32_t flags = butil::HostToNet32(enabled ? COMMON_ACK_OK : 0); + payload->assign(reinterpret_cast(&flags), sizeof(flags)); + return STEP_OK; +} + +StepResult HandshakeProtocol::ParseAck(const std::string& payload, + bool* enabled) { + CHECK(enabled != NULL); + if (payload.size() != COMMON_ACK_SIZE) { + errno = EPROTO; + return STEP_ERROR; + } + uint32_t flags = 0; + memcpy(&flags, payload.data(), sizeof(flags)); + *enabled = (butil::NetToHost32(flags) & COMMON_ACK_OK) != 0; + return STEP_OK; +} + +const FrameSpec& HandshakeProtocol::ExtensionFrameSpec() const { + static const FrameSpec empty_spec( + NULL, 0, 0, 0, FrameSpec::FIXED); + return empty_spec; +} + ServerHandshakeContext* ServerHandshakeContext::Create( HandshakeAdapter* adapter) { ServerHandshakeContext* context = @@ -40,22 +75,16 @@ void ServerHandshakeContext::Destroy() { butil::return_object(this); } -static StepResult FinishWithFailure( - HandshakeSession* session, const std::function& on_failed) { - if (on_failed) { - on_failed(); - } +static StepResult FinishWithFailure(HandshakeSession* session, + HandshakeTransport* transport) { + transport->OnFailed(); session->MarkFailed(); return STEP_ERROR; } -static StepResult FinishWithFallback( - HandshakeSession* session, const std::function& set_tcp_active) { - session->PublishFallback([&set_tcp_active]() { - if (set_tcp_active) { - set_tcp_active(); - } - }); +static StepResult FinishWithFallback(HandshakeSession* session, + HandshakeTransport* transport) { + session->PublishFallback([transport]() { transport->OnFallback(); }); return STEP_FALLBACK; } @@ -73,165 +102,198 @@ static StepResult ConvertFrameResult(FrameResult result) { return STEP_ERROR; } -StepResult HandshakeSession::SendHello(const HandshakeCodec& codec, +StepResult HandshakeSession::SendHello(HandshakeProtocol* protocol, bool enabled) { - CHECK(codec.build_hello); std::string payload; - const StepResult result = codec.build_hello(enabled, &payload); + const StepResult result = protocol->BuildHello(enabled, &payload); if (result != STEP_OK) { return result; } return ConvertFrameResult( - FrameCodec::WriteFrame(_io, codec.hello_frame, payload)); + FrameCodec::WriteFrame(_io, protocol->HelloFrameSpec(), payload)); } -StepResult HandshakeSession::ReceiveHello(const HandshakeCodec& codec, +StepResult HandshakeSession::ReceiveHello(HandshakeProtocol* protocol, HandshakeInput* input, bool push_back_on_not_mine, bool* magic_matched) { - CHECK(codec.parse_hello); std::string payload; const FrameResult frame_result = input != NULL ? FrameCodec::ParseBufferedFrame( - input, codec.hello_frame, &payload, magic_matched) + input, protocol->HelloFrameSpec(), &payload, magic_matched) : FrameCodec::ReadFrame( - _io, codec.hello_frame, push_back_on_not_mine, &payload); + _io, protocol->HelloFrameSpec(), push_back_on_not_mine, + &payload); const StepResult result = ConvertFrameResult(frame_result); if (result != STEP_OK) { return result; } - set_protocol_version(codec.protocol_version); - return codec.parse_hello(payload); + set_protocol_version(protocol->ProtocolVersion()); + return protocol->ParseHello(payload); } -StepResult HandshakeSession::SendAck(const HandshakeCodec& codec, +StepResult HandshakeSession::SendAck(HandshakeProtocol* protocol, bool enabled) { - CHECK(codec.build_ack); std::string payload; - const StepResult result = codec.build_ack(enabled, &payload); + const StepResult result = protocol->BuildAck(enabled, &payload); if (result != STEP_OK) { return result; } return ConvertFrameResult( - FrameCodec::WriteFrame(_io, codec.ack_frame, payload)); + FrameCodec::WriteFrame(_io, protocol->AckFrameSpec(), payload)); } -StepResult HandshakeSession::ReceiveAck(const HandshakeCodec& codec, +StepResult HandshakeSession::ReceiveAck(HandshakeProtocol* protocol, HandshakeInput* input, bool* enabled) { - CHECK(codec.parse_ack); std::string payload; const FrameResult frame_result = input != NULL - ? FrameCodec::ParseBufferedFrame(input, codec.ack_frame, &payload) - : FrameCodec::ReadFrame(_io, codec.ack_frame, false, &payload); + ? FrameCodec::ParseBufferedFrame( + input, protocol->AckFrameSpec(), &payload) + : FrameCodec::ReadFrame( + _io, protocol->AckFrameSpec(), false, &payload); const StepResult result = ConvertFrameResult(frame_result); if (result != STEP_OK) { return result; } - return codec.parse_ack(payload, enabled); + return protocol->ParseAck(payload, enabled); +} + +StepResult HandshakeSession::SendExtension(HandshakeProtocol* protocol, + bool enabled) { + std::string payload; + const StepResult result = protocol->BuildExtension(enabled, &payload); + if (result != STEP_OK) { + return result; + } + return ConvertFrameResult( + FrameCodec::WriteFrame( + _io, protocol->ExtensionFrameSpec(), payload)); +} + +StepResult HandshakeSession::ReceiveExtension(HandshakeProtocol* protocol, + HandshakeInput* input) { + std::string payload; + const FrameResult frame_result = input != NULL + ? FrameCodec::ParseBufferedFrame( + input, protocol->ExtensionFrameSpec(), &payload) + : FrameCodec::ReadFrame( + _io, protocol->ExtensionFrameSpec(), false, &payload); + const StepResult result = ConvertFrameResult(frame_result); + return result == STEP_OK ? protocol->ParseExtension(payload) : result; } StepResult HandshakeSession::SelectAndReceiveHello( - const std::vector& codecs, HandshakeInput* input, - bool push_back_on_not_mine, const HandshakeCodec** selected) { - CHECK(!codecs.empty()); + const std::vector& protocols, HandshakeInput* input, + bool push_back_on_not_mine, HandshakeProtocol** selected) { + CHECK(!protocols.empty()); CHECK(selected != NULL); if (input == NULL) { // A blocking byte stream cannot try a second codec after consuming // bytes from the fd. Such protocols must select a single codec before // entering the common session. - CHECK_EQ(1UL, codecs.size()); - *selected = &codecs.front(); - return ReceiveHello(**selected, NULL, push_back_on_not_mine); + CHECK_EQ(1UL, protocols.size()); + *selected = protocols.front(); + return ReceiveHello(*selected, NULL, push_back_on_not_mine); } bool need_more = false; - for (size_t i = 0; i < codecs.size(); ++i) { + for (size_t i = 0; i < protocols.size(); ++i) { bool magic_matched = false; const StepResult result = ReceiveHello( - codecs[i], input, false, &magic_matched); + protocols[i], input, false, &magic_matched); if (result == STEP_NOT_MINE) { continue; } if (result == STEP_NEED_MORE) { if (magic_matched) { - *selected = &codecs[i]; - set_protocol_version(codecs[i].protocol_version); + *selected = protocols[i]; + set_protocol_version(protocols[i]->ProtocolVersion()); return STEP_NEED_MORE; } need_more = true; continue; } - *selected = &codecs[i]; + *selected = protocols[i]; return result; } return need_more ? STEP_NEED_MORE : STEP_NOT_MINE; } -StepResult HandshakeSession::RunClient( - const ClientHandshakeCallbacks& callbacks) { - CHECK(callbacks.transport.prepare_resources); - CHECK(callbacks.transport.negotiate_resources); - CHECK(callbacks.transport.set_high_speed_active); - CHECK(callbacks.transport.set_tcp_active); +StepResult HandshakeSession::RunClient(HandshakeProtocol* protocol, + HandshakeTransport* transport) { + CHECK(protocol != NULL); + CHECK(transport != NULL); + transport->OnProtocolSelected(protocol); // A client handshake runs once on a potentially reused bthread. Do not // let an errno left by earlier work override this handshake's result. errno = 0; SetPhase(PREPARING); - StepResult result = callbacks.transport.prepare_resources(); + StepResult result = transport->PrepareResources(); if (result == STEP_FALLBACK) { - return FinishWithFallback(this, callbacks.transport.set_tcp_active); + return FinishWithFallback(this, transport); } if (result != STEP_OK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } SetPhase(HELLO_SEND); - if (SendHello(callbacks.codec, true) != STEP_OK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + if (SendHello(protocol, true) != STEP_OK) { + return FinishWithFailure(this, transport); } SetPhase(HELLO_WAIT); - result = ReceiveHello(callbacks.codec, NULL, false); + result = ReceiveHello(protocol, NULL, false); if (result == STEP_NOT_MINE || result == STEP_NEED_MORE) { errno = EPROTO; } if (result == STEP_ERROR || result == STEP_NOT_MINE || result == STEP_NEED_MORE) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } bool enabled = result == STEP_OK; + if (enabled && protocol->HasExtension()) { + SetPhase(EXTENSION_SEND); + if (SendExtension(protocol, true) != STEP_OK) { + return FinishWithFailure(this, transport); + } + SetPhase(EXTENSION_WAIT); + result = ReceiveExtension(protocol, NULL); + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + enabled = result == STEP_OK; + } + if (enabled) { SetPhase(NEGOTIATING); - result = callbacks.transport.negotiate_resources(); + result = transport->NegotiateResources(); if (result != STEP_OK && result != STEP_FALLBACK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } enabled = result == STEP_OK; } SetPhase(ACK_SEND); - if (SendAck(callbacks.codec, enabled) != STEP_OK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + if (SendAck(protocol, enabled) != STEP_OK) { + return FinishWithFailure(this, transport); } if (enabled) { - callbacks.transport.set_high_speed_active(); + transport->OnEstablished(); MarkEstablished(); return STEP_OK; } - return FinishWithFallback(this, callbacks.transport.set_tcp_active); + return FinishWithFallback(this, transport); } StepResult HandshakeSession::RunServer( - const ServerHandshakeCallbacks& callbacks) { - CHECK(!callbacks.codecs.empty()); - CHECK(callbacks.transport.prepare_resources); - CHECK(callbacks.transport.negotiate_resources); - CHECK(callbacks.transport.set_high_speed_active); - CHECK(callbacks.transport.set_tcp_active); + const std::vector& protocols, HandshakeInput* input, + HandshakeTransport* transport, bool fallback_on_not_mine) { + CHECK(!protocols.empty()); + CHECK(transport != NULL); // Once TCP fallback has been published, subsequent bytes are application // protocol data and must bypass every upgrade codec without changing the @@ -240,17 +302,16 @@ StepResult HandshakeSession::RunServer( return STEP_NOT_MINE; } - const HandshakeCodec* selected = NULL; - if (phase() != ACK_WAIT) { + HandshakeProtocol* selected = NULL; + if (phase() != ACK_WAIT && phase() != EXTENSION_WAIT) { const int previous_phase = phase(); _local_enabled = false; SetPhase(HELLO_WAIT); StepResult result = SelectAndReceiveHello( - callbacks.codecs, callbacks.input, - callbacks.fallback_on_not_mine, &selected); + protocols, input, fallback_on_not_mine, &selected); if (result == STEP_NOT_MINE) { - if (callbacks.fallback_on_not_mine) { - return FinishWithFallback(this, callbacks.transport.set_tcp_active); + if (fallback_on_not_mine) { + return FinishWithFallback(this, transport); } SetPhase(UNINITIALIZED); return STEP_NOT_MINE; @@ -262,74 +323,99 @@ StepResult HandshakeSession::RunServer( return STEP_NEED_MORE; } if (result == STEP_ERROR) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); + } + CHECK(selected != NULL); + transport->OnProtocolSelected(selected); + if (result == STEP_FALLBACK) { + // Publish the disabled transport state immediately. The server + // still sends a disabled hello and consumes the peer ACK before + // publishing the terminal FALLBACK_TCP phase. + transport->OnFallback(); } bool enabled = result == STEP_OK; if (enabled) { SetPhase(PREPARING); - result = callbacks.transport.prepare_resources(); + result = transport->PrepareResources(); if (result != STEP_OK && result != STEP_FALLBACK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } enabled = result == STEP_OK; } if (enabled) { SetPhase(NEGOTIATING); - result = callbacks.transport.negotiate_resources(); + result = transport->NegotiateResources(); if (result != STEP_OK && result != STEP_FALLBACK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } enabled = result == STEP_OK; } SetPhase(HELLO_SEND); - CHECK(selected != NULL); _local_enabled = enabled; - if (SendHello(*selected, enabled) != STEP_OK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + if (SendHello(selected, enabled) != STEP_OK) { + return FinishWithFailure(this, transport); } - SetPhase(ACK_WAIT); + SetPhase(enabled && selected->HasExtension() + ? EXTENSION_WAIT + : ACK_WAIT); } else { - for (size_t i = 0; i < callbacks.codecs.size(); ++i) { - if (callbacks.codecs[i].protocol_version == protocol_version()) { - selected = &callbacks.codecs[i]; + for (size_t i = 0; i < protocols.size(); ++i) { + if (protocols[i]->ProtocolVersion() == protocol_version()) { + selected = protocols[i]; break; } } CHECK(selected != NULL); } + if (phase() == EXTENSION_WAIT) { + StepResult result = ReceiveExtension(selected, input); + if (result == STEP_NEED_MORE) { + return STEP_NEED_MORE; + } + if (result != STEP_OK && result != STEP_FALLBACK) { + return FinishWithFailure(this, transport); + } + if (result == STEP_FALLBACK) { + _local_enabled = false; + } + SetPhase(EXTENSION_SEND); + if (SendExtension(selected, _local_enabled) != STEP_OK) { + return FinishWithFailure(this, transport); + } + SetPhase(ACK_WAIT); + } + // Always try the ACK callback once. For a non-blocking server it returns // STEP_NEED_MORE when the ACK has not arrived; when Hello and ACK are // coalesced in the input buffer this consumes the ACK without waiting for // another socket edge. bool peer_enabled = false; - StepResult result = ReceiveAck( - *selected, callbacks.input, &peer_enabled); + StepResult result = ReceiveAck(selected, input, &peer_enabled); if (result == STEP_NEED_MORE) { return STEP_NEED_MORE; } if (result == STEP_ERROR || result == STEP_NOT_MINE) { - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } if (result == STEP_FALLBACK) { - return FinishWithFallback(this, callbacks.transport.set_tcp_active); + return FinishWithFallback(this, transport); } if (!peer_enabled) { - return FinishWithFallback(this, callbacks.transport.set_tcp_active); + return FinishWithFallback(this, transport); } if (!_local_enabled) { errno = EPROTO; - return FinishWithFailure(this, callbacks.transport.on_failed); + return FinishWithFailure(this, transport); } - if (callbacks.validate_established && - callbacks.validate_established() != STEP_OK) { - return FinishWithFailure(this, callbacks.transport.on_failed); + if (transport->ValidateEstablished() != STEP_OK) { + return FinishWithFailure(this, transport); } - callbacks.transport.set_high_speed_active(); + transport->OnEstablished(); MarkEstablished(); return STEP_OK; } diff --git a/src/brpc/transport_handshake.h b/src/brpc/transport_handshake.h index 1794ac3bef..b558ee0ac6 100644 --- a/src/brpc/transport_handshake.h +++ b/src/brpc/transport_handshake.h @@ -19,7 +19,6 @@ #define BRPC_TRANSPORT_HANDSHAKE_H #include -#include #include #include @@ -60,6 +59,8 @@ enum Phase { NEGOTIATING = 4, ACK_SEND = 5, ACK_WAIT = 6, + EXTENSION_SEND = 7, + EXTENSION_WAIT = 8, ESTABLISHED = 0x100, FALLBACK_TCP = 0x200, FAILED = 0x300, @@ -73,45 +74,72 @@ enum StepResult { STEP_ERROR, }; -// A protocol describes only its fields and resource-independent wire values. -// HandshakeSession owns framing and I/O through FrameCodec. The callbacks may -// retain strongly typed parsed state in their protocol adapter. -struct HandshakeCodec { - int protocol_version; - FrameSpec hello_frame; - FrameSpec ack_frame; - std::function build_hello; - std::function parse_hello; - std::function build_ack; - std::function parse_ack; +// Wire-level participant in a transport upgrade. Implementations own parsed +// protocol state; HandshakeSession owns framing and phase orchestration. +class HandshakeProtocol { +public: + virtual ~HandshakeProtocol() = default; + + virtual int ProtocolVersion() const = 0; + virtual const FrameSpec& HelloFrameSpec() const = 0; + virtual const FrameSpec& AckFrameSpec() const = 0; + virtual StepResult BuildHello(bool enabled, std::string* payload) = 0; + virtual StepResult ParseHello(const std::string& payload) = 0; + + // RDMA and UBSHM use the same four-byte, network-order ACK. Protocols + // with a different ACK format may override these methods. + virtual StepResult BuildAck(bool enabled, std::string* payload); + virtual StepResult ParseAck(const std::string& payload, bool* enabled); + + virtual bool HasExtension() const { return false; } + virtual const FrameSpec& ExtensionFrameSpec() const; + virtual StepResult BuildExtension(bool, std::string*) { + return STEP_ERROR; + } + virtual StepResult ParseExtension(const std::string&) { + return STEP_ERROR; + } }; -// Resource-specific operations supplied by a Transport and invoked by the -// common coordinator. Wire I/O and field codec invocation remain owned by -// HandshakeSession. -struct TransportUpgradeOps { - std::function prepare_resources; - std::function negotiate_resources; - std::function set_high_speed_active; - std::function set_tcp_active; - std::function on_failed; +// Resource-level participant in a transport upgrade. Cleanup remains part of +// this contract: resources may already exist when negotiation falls back or +// a framing/I/O error terminates the handshake. +class HandshakeTransport { +public: + virtual ~HandshakeTransport() = default; + + virtual void OnProtocolSelected(HandshakeProtocol*) {} + virtual StepResult PrepareResources() = 0; + virtual StepResult NegotiateResources() = 0; + virtual void OnEstablished() = 0; + virtual void OnFallback() = 0; + virtual void OnFailed() = 0; + virtual StepResult ValidateEstablished() { return STEP_OK; } }; -struct ClientHandshakeCallbacks { - HandshakeCodec codec; - TransportUpgradeOps transport; +// Participant for a server that recognizes an upgrade protocol only to +// negotiate TCP fallback. It owns no high-speed resources. +class FallbackHandshakeTransport : public HandshakeTransport { +public: + StepResult PrepareResources() override { return STEP_OK; } + StepResult NegotiateResources() override { return STEP_OK; } + void OnEstablished() override {} + void OnFallback() override {} + void OnFailed() override {} }; -// The server driver is independent of the input mode. A parser callback can -// return STEP_NEED_MORE, while a blocking callback waits before returning. -struct ServerHandshakeCallbacks { - bool fallback_on_not_mine; - // Buffered parsers may offer multiple codecs (RDMA v2/v3). Blocking - // server handshakes currently provide exactly one codec. - std::vector codecs; - HandshakeInput* input; - TransportUpgradeOps transport; - std::function validate_established; +class FallbackHandshakeProtocol : public HandshakeProtocol { +public: + StepResult ParseHello(const std::string&) override { + return STEP_FALLBACK; + } + StepResult ParseAck(const std::string& payload, bool* enabled) override { + bool ignored = false; + const StepResult result = HandshakeProtocol::ParseAck( + payload, &ignored); + *enabled = false; + return result; + } }; // Owns one connection-upgrade attempt, invokes the protocol field codec and @@ -166,21 +194,29 @@ class HandshakeSession { // restores the Socket-backed implementation. void SetIOForTest(HandshakeIO* io) { _io = io; } - StepResult RunClient(const ClientHandshakeCallbacks& callbacks); - StepResult RunServer(const ServerHandshakeCallbacks& callbacks); + StepResult RunClient(HandshakeProtocol* protocol, + HandshakeTransport* transport); + StepResult RunServer(const std::vector& protocols, + HandshakeInput* input, + HandshakeTransport* transport, + bool fallback_on_not_mine); private: - StepResult SendHello(const HandshakeCodec& codec, bool enabled); - StepResult ReceiveHello(const HandshakeCodec& codec, + StepResult SendHello(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveHello(HandshakeProtocol* protocol, HandshakeInput* input, bool push_back_on_not_mine, bool* magic_matched = NULL); - StepResult SendAck(const HandshakeCodec& codec, bool enabled); - StepResult ReceiveAck(const HandshakeCodec& codec, + StepResult SendAck(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveAck(HandshakeProtocol* protocol, HandshakeInput* input, bool* enabled); + StepResult SendExtension(HandshakeProtocol* protocol, bool enabled); + StepResult ReceiveExtension(HandshakeProtocol* protocol, + HandshakeInput* input); StepResult SelectAndReceiveHello( - const std::vector& codecs, HandshakeInput* input, - bool push_back_on_not_mine, const HandshakeCodec** selected); + const std::vector& protocols, + HandshakeInput* input, bool push_back_on_not_mine, + HandshakeProtocol** selected); SocketHandshakeIO _socket_io; HandshakeIO* _io; diff --git a/src/brpc/ubshm/ub_endpoint.cpp b/src/brpc/ubshm/ub_endpoint.cpp index 131fbfcadf..3835376646 100644 --- a/src/brpc/ubshm/ub_endpoint.cpp +++ b/src/brpc/ubshm/ub_endpoint.cpp @@ -66,6 +66,7 @@ void UBShmEndpoint::Reset() { delete _ub_ring; _ub_ring = nullptr; _poller_sid = INVALID_SOCKET_ID; + _negotiated_data_format = UBR_DATA_FORMAT_NONE; } bool UBShmEndpoint::IsWritable() const { diff --git a/src/brpc/ubshm/ub_endpoint.h b/src/brpc/ubshm/ub_endpoint.h index 90ed94bad1..b9d6d196c8 100644 --- a/src/brpc/ubshm/ub_endpoint.h +++ b/src/brpc/ubshm/ub_endpoint.h @@ -46,8 +46,8 @@ DECLARE_bool(ub_disable_bthread); class BAIDU_CACHELINE_ALIGNMENT UBShmEndpoint : public SocketUser { friend class Socket; - friend class ::brpc::UBShmTransport; - friend class ::brpc::handshake::UBShmServerHandshakeAdapter; +friend class ::brpc::UBShmTransport; +friend class ::brpc::handshake::UBShmServerHandshakeAdapter; public: explicit UBShmEndpoint(Socket* s); @@ -61,6 +61,12 @@ friend class Socket; // Reset the endpoint (for next use) void Reset(); + void SetNegotiatedDataFormat(UbrDataFormat format) { + _negotiated_data_format = format; + } + UbrDataFormat negotiated_data_format() const { + return _negotiated_data_format; + } // Cut data from the given IOBuf list and use UBRING to send // Return bytes cut if success, -1 if failed and errno set @@ -105,7 +111,7 @@ friend class Socket; // Release resources void DeallocateResources(); - // Poll CQ and get the work completion + // Poll CQ and get the work completion static void PollIn(UBShmEndpoint* ep, uint32_t ep_event); static void PollOut(UBShmEndpoint* ep, uint32_t ep_event); @@ -113,6 +119,7 @@ friend class Socket; // Not owner Socket* _socket; SocketId _socket_id; + UbrDataFormat _negotiated_data_format{UBR_DATA_FORMAT_NONE}; // ub resource ubring::UBRing* _ub_ring{nullptr}; @@ -122,14 +129,14 @@ friend class Socket; DISALLOW_COPY_AND_ASSIGN(UBShmEndpoint); struct PollerSidOp { - enum OpType { ADD, REMOVE, MOD }; + enum OpType { ADD, REMOVE, MOD }; SocketId sid; - uint32_t event; + uint32_t event; OpType type; }; struct PollerSidOpHash { - std::size_t operator()(const PollerSidOp &op) const { return op.sid; } + std::size_t operator()(const PollerSidOp& op) const { return op.sid; } }; struct PollerSidOpEqual { @@ -141,7 +148,7 @@ friend class Socket; // Poller instance struct BAIDU_CACHELINE_ALIGNMENT Poller { bthread_t tid{INVALID_BTHREAD}; - butil::MPSCQueue> op_queue; + butil::MPSCQueue> op_queue; // Callback used for io_uring/spdk etc std::function callback; // Init and Destroy function @@ -156,7 +163,7 @@ friend class Socket; }; static std::vector _poller_groups; - void PollerRegisterEvent(PollerSidOp::OpType op, uint32_t events = EPOLLET); + void PollerRegisterEvent(PollerSidOp::OpType op, uint32_t events = EPOLLET); }; } // namespace ubring diff --git a/src/brpc/uri.cpp b/src/brpc/uri.cpp index 2881a8e54a..6b10464ba7 100644 --- a/src/brpc/uri.cpp +++ b/src/brpc/uri.cpp @@ -17,9 +17,8 @@ #include // isalnum - #include - +#include #include "brpc/log.h" #include "brpc/details/http_parser.h" // http_parser_parse_url #include "brpc/uri.h" // URI @@ -27,15 +26,16 @@ namespace brpc { +DEFINE_uint32(http_max_query_count, 1000, + "Reject a URL carrying more than so many query parameters. " + "0 lifts the limit."); + URI::URI() : _port(-1) , _query_was_modified(false) , _initialized_query_map(false) {} -URI::~URI() { -} - void URI::Clear() { _st.reset(); _port = -1; @@ -64,6 +64,22 @@ void URI::Swap(URI &rhs) { _query_map.swap(rhs._query_map); } +// Counting separators rather than map entries deliberately overestimates: the +// splitter walks every segment even when the keys repeat, and it is that walk, +// not the final map size, that the limit is meant to bound. +static bool TooManyQueries(const std::string& query) { + if (FLAGS_http_max_query_count == 0 || query.empty()) { + return false; + } + uint32_t count = 1; + for (char i : query) { + if (i == '&' && ++count > FLAGS_http_max_query_count) { + return true; + } + } + return false; +} + // Parse queries, which is case-sensitive static void ParseQueries(URI::QueryMap& query_map, const std::string &query) { query_map.clear(); @@ -238,6 +254,11 @@ int URI::SetHttpURL(const char* url) { } } _query.assign(start, p - start); + if (TooManyQueries(_query)) { + _st.set_error(EINVAL, "More than %u query parameters in url", + FLAGS_http_max_query_count); + return -1; + } } if (*p == '#') { start = ++p; @@ -411,7 +432,8 @@ void URI::SetHostAndPort(const std::string& host) { _host.assign(host_begin, host_end - host_begin); } -void URI::SetH2Path(const char* h2_path) { +int URI::SetH2Path(const char* h2_path) { + _st.reset(); _path.clear(); _query.clear(); _fragment.clear(); @@ -427,12 +449,18 @@ void URI::SetH2Path(const char* h2_path) { start = ++p; for (; *p && *p != '#'; ++p) {} _query.assign(start, p - start); + if (TooManyQueries(_query)) { + _st.set_error(EINVAL, "More than %u query parameters in :path", + FLAGS_http_max_query_count); + return -1; + } } if (*p == '#') { start = ++p; for (; *p; ++p) {} _fragment.assign(start, p - start); } + return 0; } QueryRemover::QueryRemover(const std::string* str) diff --git a/src/brpc/uri.h b/src/brpc/uri.h index 7edac4002d..a42cf88f12 100644 --- a/src/brpc/uri.h +++ b/src/brpc/uri.h @@ -56,7 +56,7 @@ class URI { // You can copy a URI. URI(); - ~URI(); + ~URI() = default; // Exchange internal fields with another URI. void Swap(URI &rhs); @@ -99,8 +99,9 @@ class URI { void set_port(int port) { _port = port; } void SetHostAndPort(const std::string& host_and_optional_port); // Set path/query/fragment with the input in form of "path?query#fragment" - void SetH2Path(const char* h2_path); - void SetH2Path(const std::string& path) { SetH2Path(path.c_str()); } + // Returns 0 on success, -1 otherwise and status() is set. + int SetH2Path(const char* h2_path); + int SetH2Path(const std::string& path) { return SetH2Path(path.c_str()); } // Get the value of a CASE-SENSITIVE key. // Returns pointer to the value, nullptr when the key does not exist. diff --git a/src/brpc/urma/mock_urma.cpp b/src/brpc/urma/mock_urma.cpp new file mode 100644 index 0000000000..b0572c2e86 --- /dev/null +++ b/src/brpc/urma/mock_urma.cpp @@ -0,0 +1,726 @@ +// 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_user_ctl(urma_context_t *ctx, urma_user_ctl_in_t *in, + urma_user_ctl_out_t *out) { + std::shared_lock lock(g_rw_mutex); + if (!ctx || context_map.find(ctx) == context_map.end() || !in || + (in->len != 0 && in->addr == 0)) { + return URMA_EINVAL; + } + // The mock has no provider-specific controls. Accept the command so the + // optional bonding setup remains linkable and harmless in mock builds. + (void)out; + return URMA_SUCCESS; +} + +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..26c7a1b9dc --- /dev/null +++ b/src/brpc/urma/urma_endpoint.cpp @@ -0,0 +1,1849 @@ +// 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 = sizeof(butil::IOBuf::Block); + +// ---- 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); + _input_processor.Init(s, InputMessengerProcessor::STREAM_URMA_JETTY); +} + +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); + _input_processor.Reset(); + _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->fd_input_processor().read_buf().append(data, len); +} + +int UrmaEndpoint::WriteToFd(void* data, size_t len) { + char* p = static_cast(data); + size_t written = 0; + while (written < len) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = write(fd, p + written, len - written); + if (nw >= 0) { + written += nw; + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +// ============================================================================ +// Hello builders / parsers. +// ============================================================================ + +void UrmaEndpoint::MakeLocalParsedHello(ParsedHello* out) const { + *out = ParsedHello{}; // value-initialize (avoids memset on non-trivial type) + 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; + } + butil::IOPortal& read_buf = _input_processor.read_buf(); + if (zerocopy) { + _rbuf[_rq_received].cutn(&read_buf, cr.completion_len); + } else { + 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 the + // URMA input processor 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; + } + + if (!s->user()) { + 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; + _input_processor.ProcessNewMessage(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 + << " read_buf=" << _input_processor.read_buf().size(); +} + +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..e5aed5ac66 --- /dev/null +++ b/src/brpc/urma/urma_endpoint.h @@ -0,0 +1,360 @@ +// 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 the TCP input processor so it can re-parse + // them (used on fallback when the magic is not "URMA"). + void PushBackToReadBuf(const void* data, size_t len); + // Write at most @len bytes from @data to the TCP fd; waits on + // _epollout_butex on EAGAIN. Returns 0 on success, -1 on IO error. + int WriteToFd(void* data, size_t len); + // Write an IOBuf to the TCP fd (used by v3 protobuf handshake). + int WriteToFd(butil::IOBuf& data); + + // ---- Hello builders/parsers (used by urma_handshake.cpp) ---- + // Fill the v2 binary HelloMessage with this endpoint's local params. + void FillLocalHelloV2(v2_wire::HelloMessage* out) const; + // Fill the v3 protobuf UrmaHello with this endpoint's local params. + void FillLocalHelloV3(UrmaHello* out) const; + // Write "URM3" + 4B big-endian pb_size + protobuf bytes. + int WriteHelloV3(const UrmaHello& msg); + // Read pb_size (4B) + protobuf bytes; parse; validate. Sets *negotiated + // to false (returns 0) if the message is invalid. + int ReadAndParseHelloV3(ParsedHello* out, bool* negotiated); + + // 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 + // the URMA input processor, 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 the URMA input processor 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 + + // The input stream carried by the URMA jetty. + InputMessengerProcessor _input_processor; + + // 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..b7018606e0 --- /dev/null +++ b/src/brpc/urma/urma_handshake.cpp @@ -0,0 +1,361 @@ +// 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; + } + // The provider consumes the advertised VA/length as a single remote + // segment. Reject a wrapped end address before it reaches the import + // API, which could otherwise turn a malformed hello into an invalid + // segment range. + if (h.seg_len == 0 || h.seg_va == 0 || + h.seg_len > std::numeric_limits::max() - h.seg_va) { + 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..ded662fb22 --- /dev/null +++ b/src/brpc/urma/urma_helper.cpp @@ -0,0 +1,787 @@ +// 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 + +#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"); + + +// 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() { + // pthread_t is an integer on Linux and a pointer on macOS. Hash its native + // representation instead of assuming either form with a cast. + return std::hash()(pthread_self()) % 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) - + sizeof(butil::IOBuf::Block); + + // 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..c5cc163ffa --- /dev/null +++ b/src/brpc/urma_transport.cpp @@ -0,0 +1,251 @@ +// 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) { + // The URMA endpoint is not usable until negotiation and peer-resource + // import have completed. Keep using TCP while the state is UNKNOWN, and + // switch the data path only after the handshake publishes URMA_ON. + if (_urma_ep && + _urma_state.load(butil::memory_order_acquire) == URMA_ON) { + butil::IOBuf* data_arr[1] = {buf}; + return _urma_ep->CutFromIOBufList(data_arr, 1); + } else { + return _tcp_transport->CutFromIOBuf(buf); + } +} + +ssize_t UrmaTransport::CutFromIOBufList(butil::IOBuf** buf, size_t ndata) { + if (_urma_ep && + _urma_state.load(butil::memory_order_acquire) == URMA_ON) { + return _urma_ep->CutFromIOBufList(buf, ndata); + } else { + return _tcp_transport->CutFromIOBufList(buf, ndata); + } +} + +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/src/bthread/bthread.cpp b/src/bthread/bthread.cpp index 727b4ebace..836dec9994 100644 --- a/src/bthread/bthread.cpp +++ b/src/bthread/bthread.cpp @@ -141,7 +141,10 @@ bthread_t init_for_pthread_stack_trace() { } pthread_fake_meta->attr = BTHREAD_ATTR_PTHREAD; - pthread_fake_meta->tid = make_tid(*pthread_fake_meta->version_butex, slot); + auto* version = reinterpret_cast*>( + pthread_fake_meta->version_butex); + pthread_fake_meta->tid = make_tid(static_cast( + version->load(butil::memory_order_relaxed)), slot); // Make TaskTracer use signal trace mode for pthread. c->_task_tracer.set_running_status(syscall(SYS_gettid), pthread_fake_meta); @@ -152,11 +155,17 @@ bthread_t init_for_pthread_stack_trace() { { BAIDU_SCOPED_LOCK(pthread_fake_meta->version_lock); tracing = TaskTracer::set_end_status_unsafe(pthread_fake_meta); - // If resulting version is 0, - // change it to 1 to make bthread_t never be 0. - if (0 == ++*pthread_fake_meta->version_butex) { - ++*pthread_fake_meta->version_butex; + // Publish the version atomically, just like task_runner(), since + // lock-free readers may still access this TaskMeta. + auto* version = reinterpret_cast*>( + pthread_fake_meta->version_butex); + uint32_t next_version = static_cast( + version->load(butil::memory_order_relaxed)) + 1; + // Skip zero to make bthread_t never be 0. + if (0 == next_version) { + ++next_version; } + version->store(static_cast(next_version), butil::memory_order_release); } if (tracing) { diff --git a/src/bthread/butex.cpp b/src/bthread/butex.cpp index 63920ca9eb..790da14584 100644 --- a/src/bthread/butex.cpp +++ b/src/bthread/butex.cpp @@ -573,26 +573,62 @@ void wait_for_butex(void* arg) { BAIDU_SCOPED_LOCK(b->waiter_lock); if (b->value.load(butil::memory_order_relaxed) != bw->expected_value) { bw->waiter_state = WAITER_STATE_UNMATCHEDVALUE; - } else if (bw->waiter_state == WAITER_STATE_READY/*1*/ && - !bw->task_meta->interrupted) { - if (args->prepend) { - b->waiters.Prepend(bw); - } else { - b->waiters.Append(bw); - } - bw->container.store(b, butil::memory_order_relaxed); + } else { + // Checking `interrupted` and publishing `bw->container` must be + // atomic with respect to TaskGroup::interrupt(), which sets + // `interrupted` and consumes `current_waiter` under the same + // `version_lock`. Otherwise interrupt() may consume `bw` in between + // and its erase_from_butex() does nothing because `container` is + // still nullptr, leaving this bthread queued but never woken up. + // `container` cannot be published upfront: it must stay nullptr until + // the bthread is off its stack, see the comment after this block. + bool suspend = false; + { + // Only the interrupted check, the enqueue and the container + // store belong under `version_lock', so that they are atomic + // w.r.t. TaskGroup::interrupt(). The tracer update is kept here + // too because set_status_unsafe() takes no lock and this makes + // the transition to SUSPENDED atomic w.r.t. TraceImpl(). + BAIDU_SCOPED_LOCK(bw->task_meta->version_lock); + if (bw->waiter_state == WAITER_STATE_READY/*1*/ && + !bw->task_meta->interrupted) { + if (args->prepend) { + b->waiters.Prepend(bw); + } else { + b->waiters.Append(bw); + } + bw->container.store(b, butil::memory_order_relaxed); #ifdef BRPC_BTHREAD_TRACER - bw->control->_task_tracer.set_status(TASK_STATUS_SUSPENDED, bw->task_meta); + TaskTracer::set_status_unsafe(TASK_STATUS_SUSPENDED, bw->task_meta); #endif // BRPC_BTHREAD_TRACER - if (bw->abstime != nullptr) { + suspend = true; + } + } + if (suspend && bw->abstime != nullptr) { bw->sleep_id = get_global_timer_thread()->schedule( erase_from_butex_and_wakeup, bw, *bw->abstime); if (!bw->sleep_id) { // TimerThread stopped. + // No timer will ever fire, so `bw` must not stay queued. + // CAUTION: erase_from_butex_and_wakeup() must NOT be called + // here. It takes `waiter_lock` (already held) and then, via + // TaskGroup::ready_to_run{,_remote}() -> + // TaskTracer::set_status(), `version_lock` as well. Both are + // non-recursive, so calling it self-deadlocks. errno = ESTOP; - erase_from_butex_and_wakeup(bw); + bw->RemoveFromList(); + bw->container.store(nullptr, butil::memory_order_relaxed); + bw->waiter_state = WAITER_STATE_TIMEDOUT; + suspend = false; } } - return; + if (suspend) { + return; + } + // Not suspended: the bthread has already been switched out by + // set_remained()/sched() in butex_wait(), so nobody else will run + // it. Fall through to re-schedule it below. This covers value + // unmatched, already timed out (waiter_state != READY), already + // interrupted, and the ESTOP path above. } } @@ -619,33 +655,64 @@ static int butex_wait_from_pthread(TaskGroup* g, Butex* b, int expected_value, TaskMeta* task = nullptr; ButexPthreadWaiter pw; pw.tid = 0; + // Initialize `container` to nullptr before `pw` is published to + // `current_waiter` below: a concurrent TaskGroup::interrupt() may consume + // `pw` and call erase_from_butex() on it, which must observe either nullptr + // (a no-op) or a valid Butex, never stack garbage. `container` stays nullptr + // until `pw` is queued. + pw.container.store(nullptr, butil::memory_order_relaxed); pw.sig.store(PTHREAD_NOT_SIGNALLED, butil::memory_order_relaxed); int rc = 0; - + if (g) { task = g->current_task(); task->current_waiter.store(&pw, butil::memory_order_release); } - b->waiter_lock.lock(); - if (b->value.load(butil::memory_order_relaxed) != expected_value) { - b->waiter_lock.unlock(); - errno = EWOULDBLOCK; - rc = -1; - } else if (task != nullptr && task->interrupted) { - b->waiter_lock.unlock(); - // Race with set and may consume multiple interruptions, which are OK. - task->interrupted = false; - errno = EINTR; - rc = -1; - } else { - if (prepend) { - b->waiters.Prepend(&pw); + bool queued = false; + { + BAIDU_SCOPED_LOCK(b->waiter_lock); + if (b->value.load(butil::memory_order_relaxed) != expected_value) { + errno = EWOULDBLOCK; + rc = -1; + } else if (task != nullptr) { + // Checking `interrupted` and publishing `container` must be atomic + // with respect to TaskGroup::interrupt(), which sets `interrupted` + // and consumes `current_waiter` under the same version_lock. + // Otherwise interrupt() may consume `pw' in between and its + // erase_from_butex() does nothing because `container' is still + // nullptr, leaving `pw` queued but never woken up. + BAIDU_SCOPED_LOCK(task->version_lock); + if (task->interrupted) { + // Already interrupted before queueing: consume it and return + // EINTR without blocking. An interruption that arrives after + // `pw` is queued is handled by the epilogue below instead. + // Race with set and may consume multiple interruptions, + // which are OK. + task->interrupted = false; + errno = EINTR; + rc = -1; + } else { + if (prepend) { + b->waiters.Prepend(&pw); + } else { + b->waiters.Append(&pw); + } + pw.container.store(b, butil::memory_order_relaxed); + queued = true; + } } else { - b->waiters.Append(&pw); + // A non-bthread pthread cannot be interrupted, so there is no race + // with interrupt() and no need to hold `version_lock` here. + if (prepend) { + b->waiters.Prepend(&pw); + } else { + b->waiters.Append(&pw); + } + pw.container.store(b, butil::memory_order_relaxed); + queued = true; } - pw.container.store(b, butil::memory_order_relaxed); - b->waiter_lock.unlock(); - + } + if (queued) { #ifdef SHOW_BTHREAD_BUTEX_WAITER_COUNT_IN_VARS bvar::Adder& num_waiters = butex_waiter_count(); num_waiters << 1; @@ -655,7 +722,8 @@ static int butex_wait_from_pthread(TaskGroup* g, Butex* b, int expected_value, num_waiters << -1; #endif } - if (task) { + + if (task != nullptr) { // If current_waiter is nullptr, TaskGroup::interrupt() is running and // using pw, spin until current_waiter != nullptr. BT_LOOP_WHEN(task->current_waiter.exchange( diff --git a/src/bthread/context.cpp b/src/bthread/context.cpp index 7f913adfc1..467445c964 100644 --- a/src/bthread/context.cpp +++ b/src/bthread/context.cpp @@ -347,20 +347,40 @@ __asm ( " pushq %r14 \n" " pushq %r13 \n" " pushq %r12 \n" -" leaq -0x8(%rsp), %rsp\n" +" leaq -0xa8(%rsp), %rsp\n" +" movups %xmm6, 0x00(%rsp)\n" +" movups %xmm7, 0x10(%rsp)\n" +" movups %xmm8, 0x20(%rsp)\n" +" movups %xmm9, 0x30(%rsp)\n" +" movups %xmm10, 0x40(%rsp)\n" +" movups %xmm11, 0x50(%rsp)\n" +" movups %xmm12, 0x60(%rsp)\n" +" movups %xmm13, 0x70(%rsp)\n" +" movups %xmm14, 0x80(%rsp)\n" +" movups %xmm15, 0x90(%rsp)\n" " cmp $0, %rcx\n" " je 1f\n" -" stmxcsr (%rsp)\n" -" fnstcw 0x4(%rsp)\n" +" stmxcsr 0xa0(%rsp)\n" +" fnstcw 0xa4(%rsp)\n" "1:\n" " movq %rsp, (%rdi)\n" " movq %rsi, %rsp\n" +" movups 0x00(%rsp), %xmm6\n" +" movups 0x10(%rsp), %xmm7\n" +" movups 0x20(%rsp), %xmm8\n" +" movups 0x30(%rsp), %xmm9\n" +" movups 0x40(%rsp), %xmm10\n" +" movups 0x50(%rsp), %xmm11\n" +" movups 0x60(%rsp), %xmm12\n" +" movups 0x70(%rsp), %xmm13\n" +" movups 0x80(%rsp), %xmm14\n" +" movups 0x90(%rsp), %xmm15\n" " cmp $0, %rcx\n" " je 2f\n" -" ldmxcsr (%rsp)\n" -" fldcw 0x4(%rsp)\n" +" ldmxcsr 0xa0(%rsp)\n" +" fldcw 0xa4(%rsp)\n" "2:\n" -" leaq 0x8(%rsp), %rsp\n" +" leaq 0xa8(%rsp), %rsp\n" " popq %r12 \n" " popq %r13 \n" " popq %r14 \n" @@ -387,12 +407,12 @@ __asm ( "bthread_make_fcontext:\n" " movq %rdi, %rax\n" " andq $-16, %rax\n" -" leaq -0x48(%rax), %rax\n" -" movq %rdx, 0x38(%rax)\n" -" stmxcsr (%rax)\n" -" fnstcw 0x4(%rax)\n" +" leaq -0xe8(%rax), %rax\n" +" movq %rdx, 0xd8(%rax)\n" +" stmxcsr 0xa0(%rax)\n" +" fnstcw 0xa4(%rax)\n" " leaq finish(%rip), %rcx\n" -" movq %rcx, 0x40(%rax)\n" +" movq %rcx, 0xe0(%rax)\n" " ret \n" "finish:\n" " xorq %rdi, %rdi\n" diff --git a/src/bthread/parking_lot.h b/src/bthread/parking_lot.h index bd8c2c998c..ffa365ae2a 100644 --- a/src/bthread/parking_lot.h +++ b/src/bthread/parking_lot.h @@ -49,10 +49,46 @@ class BAIDU_CACHELINE_ALIGNMENT ParkingLot { // Wake up at most `num_task' workers. // Returns #workers woken up. + // + // When _no_signal_when_no_waiter is enabled, signal() and wait() form a + // Dekker-style synchronization. Both sides MUST preserve StoreLoad order: + // + // signal(): _pending_signal += num_task * 2; load(_waiter_num) + // wait(): _waiter_num += 1; load(_pending_signal) + // + // memory_order_release is NOT enough for the signal side, since it does not + // prevent the following load from being reordered before the store. Once the + // order is broken, each side may observe the other side's old value, which + // ends up with a lost wakeup: signal() sees no waiter and skips futex_wake(), + // while the waiter sees no new signal and goes to sleep. + // + // With the matching seq_cst fences below, the StoreLoad order is enforced + // on both sides: + // + // signaler waiter + // -------- ------ + // _pending_signal += num_task * 2 + // seq_cst fence + // _waiter_num += 1 + // seq_cst fence + // load(_waiter_num) + // | + // +-- > 0: futex_wake() -------> waiter is woken, or has not slept yet + // | + // +-- == 0: return + // load(_pending_signal) + // -> sees the new value; does not sleep + // + // Thus, signal() either observes a published waiter and wakes it, or the + // waiter observes the published signal before it can sleep. int signal(int num_task) { _pending_signal.fetch_add((num_task << 1), butil::memory_order_release); - if (_no_signal_when_no_waiter && _waiter_num.load(butil::memory_order_relaxed) == 0) { - return 0; + if (_no_signal_when_no_waiter) { + // Matching StoreLoad fence for the waiter-side fence below. + butil::atomic_thread_fence(butil::memory_order_seq_cst); + if (_waiter_num.load(butil::memory_order_relaxed) == 0) { + return 0; + } } return futex_wake_private(&_pending_signal, num_task); } @@ -69,13 +105,23 @@ class BAIDU_CACHELINE_ALIGNMENT ParkingLot { // Fast path, no need to futex_wait. return; } - if (_no_signal_when_no_waiter) { - _waiter_num.fetch_add(1, butil::memory_order_relaxed); + if (!_no_signal_when_no_waiter) { + futex_wait_private(&_pending_signal, expected_state.val, NULL); + return; } - futex_wait_private(&_pending_signal, expected_state.val, nullptr); - if (_no_signal_when_no_waiter) { - _waiter_num.fetch_sub(1, butil::memory_order_relaxed); + + _waiter_num.fetch_add(1, butil::memory_order_relaxed); + // Matching StoreLoad fence for the signal-side fence above. It publishes + // this waiter before _pending_signal is checked below. + butil::atomic_thread_fence(butil::memory_order_seq_cst); + // Re-check _pending_signal with a real atomic load, which closes the Dekker + // pattern at the C++ memory model level (the check inside futex_wait is done + // by the kernel and is invisible to the compiler). It also saves a futex + // syscall that would return EAGAIN immediately. + if (_pending_signal.load(butil::memory_order_relaxed) == expected_state.val) { + futex_wait_private(&_pending_signal, expected_state.val, nullptr); } + _waiter_num.fetch_sub(1, butil::memory_order_relaxed); } // Wakeup suspended wait() and make them unwaitable ever. diff --git a/src/bthread/processor.h b/src/bthread/processor.h index 06e75641e3..74be9e25ef 100644 --- a/src/bthread/processor.h +++ b/src/bthread/processor.h @@ -22,33 +22,7 @@ #ifndef BTHREAD_PROCESSOR_H #define BTHREAD_PROCESSOR_H -#include "butil/build_config.h" - -// Pause instruction to prevent excess processor bus usage, only works in GCC -# ifndef cpu_relax -#if defined(ARCH_CPU_ARM_FAMILY) -# define cpu_relax() asm volatile("yield\n": : :"memory") -#elif defined(ARCH_CPU_RISCV_FAMILY) -// Use the pause hint (Zihintpause extension). Encoding 0x0100000F -// (fence 0, 1) is a HINT on all RISC-V implementations: it never traps -// and is ignored on CPUs without Zihintpause. On CPUs with Zihintpause -// it provides a multi-cycle stall hint that reduces power and improves -// resource fairness during spin-wait loops. Matches the Linux kernel's -// RISC-V cpu_relax() behavior. .word is used instead of .insn or the -// pause mnemonic for maximum assembler compatibility. -# define cpu_relax() asm volatile(".word 0x0100000f\n": : :"memory") -#elif defined(ARCH_CPU_LOONGARCH64_FAMILY) -# define cpu_relax() asm volatile("nop\n": : :"memory"); -#else -# define cpu_relax() asm volatile("pause\n": : :"memory") -#endif -# endif - -// Compile read-write barrier -# ifndef barrier -# define barrier() asm volatile("": : :"memory") -# endif - +#include "butil/processor.h" # define BT_LOOP_WHEN(expr, num_spins) \ do { \ diff --git a/src/bthread/task_group.cpp b/src/bthread/task_group.cpp index 679e52ef5f..3051c07a26 100644 --- a/src/bthread/task_group.cpp +++ b/src/bthread/task_group.cpp @@ -84,103 +84,6 @@ BAIDU_VOLATILE_THREAD_LOCAL(void*, tls_unique_user_ptr, nullptr); const TaskStatistics EMPTY_STAT = { 0, 0, 0 }; -AtomicInteger128::Value AtomicInteger128::load() const { -#ifdef __x86_64__ - (void)_mutex; - (void)_seq; - __m128i value = _mm_load_si128(reinterpret_cast(&_value)); - return {value[0], value[1]}; -#elif defined(__ARM_NEON) - (void)_mutex; - (void)_seq; - int64x2_t value = vld1q_s64(reinterpret_cast(&_value)); - return {value[0], value[1]}; -#elif defined(__riscv) && __riscv_xlen == 64 - (void)_mutex; - // RISC-V: Seqlock-based atomic 128-bit load. - int64_t v1, v2; - uint64_t seq0, seq1; - do { - __asm__ volatile( - "ld %0, %1\n\t" - : "=r"(seq0) - : "m"(_seq) - : "memory" - ); - if (seq0 & 1) continue; - __asm__ volatile("fence r, rw\n\t" ::: "memory"); - __asm__ volatile( - "ld %0, %2\n\t" - "ld %1, %3\n\t" - : "=r"(v1), "=r"(v2) - : "m"(_value.v1), "m"(_value.v2) - : "memory" - ); - __asm__ volatile("fence r, rw\n\t" ::: "memory"); - __asm__ volatile( - "ld %0, %1\n\t" - : "=r"(seq1) - : "m"(_seq) - : "memory" - ); - } while (seq0 != seq1); - return {v1, v2}; -#else - BAIDU_SCOPED_LOCK(const_cast(_mutex)); - return _value; -#endif -} - -void AtomicInteger128::store(Value value) { -#ifdef __x86_64__ - (void)_seq; - __m128i v = _mm_load_si128(reinterpret_cast<__m128i*>(&value)); - _mm_store_si128(reinterpret_cast<__m128i*>(&_value), v); -#elif defined(__ARM_NEON) - (void)_seq; - int64x2_t v = vld1q_s64(reinterpret_cast(&value)); - vst1q_s64(reinterpret_cast(&_value), v); -#elif defined(__riscv) && __riscv_xlen == 64 - (void)_mutex; - // RISC-V: Seqlock-based atomic 128-bit store. - uint64_t old_seq; - __asm__ volatile( - "ld %0, %1\n\t" - : "=r"(old_seq) - : "m"(_seq) - : "memory" - ); - uint64_t new_seq = old_seq + 1; - __asm__ volatile( - "fence w, w\n\t" - "sd %1, %0\n\t" - : "=m"(_seq) - : "r"(new_seq) - : "memory" - ); - __asm__ volatile("fence w, w\n\t" ::: "memory"); - __asm__ volatile( - "sd %2, %0\n\t" - "sd %3, %1\n\t" - : "=m"(_value.v1), "=m"(_value.v2) - : "r"(value.v1), "r"(value.v2) - : "memory" - ); - __asm__ volatile("fence w, w\n\t" ::: "memory"); - new_seq++; - __asm__ volatile( - "sd %1, %0\n\t" - : "=m"(_seq) - : "r"(new_seq) - : "memory" - ); -#else - BAIDU_SCOPED_LOCK(const_cast(_mutex)); - _value = value; -#endif -} - - int TaskGroup::get_attr(bthread_t tid, bthread_attr_t* out) { TaskMeta* const m = address_meta(tid); if (m != nullptr) { @@ -249,7 +152,9 @@ static double get_cumulated_cputime_from_this(void* arg) { int64_t TaskGroup::cumulated_cputime_ns() const { CPUTimeStat cpu_time_stat = _cpu_time_stat.load(); - // Add the elapsed time of running bthread. + // Add elapsed time only for a running non-main task. cpuwide_time_ns() + // advances while the worker is parked, so including the main task would + // count idle waiting as worker usage. int64_t cumulated_cputime_ns = cpu_time_stat.cumulated_cputime_ns(); if (!cpu_time_stat.is_main_task()) { cumulated_cputime_ns += butil::cpuwide_time_ns() - cpu_time_stat.last_run_ns(); @@ -286,7 +191,7 @@ void TaskGroup::run_main_task() { } // Don't forget to add elapse of last wait_task. current_task()->stat.cputime_ns += - butil::cpuwide_time_ns() - _cpu_time_stat.load_unsafe().last_run_ns(); + butil::cpuwide_time_ns() - _cpu_time_stat.load_for_writer().last_run_ns(); } TaskGroup::TaskGroup(TaskControl* c) @@ -380,7 +285,9 @@ int TaskGroup::init(size_t runqueue_capacity) { m->cpuwide_start_ns = butil::cpuwide_time_ns(); m->stat = EMPTY_STAT; m->attr = BTHREAD_ATTR_TASKGROUP; - m->tid = make_tid(*m->version_butex, slot); + auto version = reinterpret_cast*>(m->version_butex); + m->tid = make_tid(static_cast( + version->load(butil::memory_order_relaxed)), slot); m->set_stack(stk); #ifdef BUTIL_USE_ASAN @@ -520,9 +427,17 @@ void TaskGroup::task_runner(intptr_t skip_remained) { #ifdef BRPC_BTHREAD_TRACER tracing = TaskTracer::set_end_status_unsafe(m); #endif // BRPC_BTHREAD_TRACER - if (0 == ++*m->version_butex) { - ++*m->version_butex; + // Bump the version with a release store so that it pairs with the + // acquire load in TaskGroup::join(): all memory writes made by this + // bthread become visible to the joining thread. Atomic access also + // avoids data races with the lock-free reads in join() and exists(). + auto* version = reinterpret_cast*>(m->version_butex); + uint32_t next_version = static_cast( + version->load(butil::memory_order_relaxed)) + 1; + if (0 == next_version) { + ++next_version; } + version->store(static_cast(next_version), butil::memory_order_release); } butex_wake_except(m->version_butex, 0); @@ -590,7 +505,9 @@ int TaskGroup::start_foreground(TaskGroup** pg, } m->cpuwide_start_ns = start_ns; m->stat = EMPTY_STAT; - m->tid = make_tid(*m->version_butex, slot); + auto version = reinterpret_cast*>(m->version_butex); + m->tid = make_tid(static_cast( + version->load(butil::memory_order_relaxed)), slot); TaskGroup* g = *pg; m->priority_index = g->_cur_meta->priority_index; @@ -662,7 +579,9 @@ int TaskGroup::start_background(bthread_t* __restrict th, } m->cpuwide_start_ns = start_ns; m->stat = EMPTY_STAT; - m->tid = make_tid(*m->version_butex, slot); + auto* version = reinterpret_cast*>(m->version_butex); + m->tid = make_tid(static_cast( + version->load(butil::memory_order_relaxed)), slot); m->priority_index = _cur_meta->priority_index; *th = m->tid; if (using_attr.flags & BTHREAD_LOG_START_AND_FINISH) { @@ -709,16 +628,18 @@ int TaskGroup::join(bthread_t tid, void** return_value) { return EINVAL; } const uint32_t expected_version = get_version(tid); - while (*m->version_butex == expected_version) { - if (butex_wait(m->version_butex, expected_version, nullptr) < 0 && + // Acquire load pairs with the release store performed when the joined + // bthread ends (see the version bump above), ensuring all of its memory + // writes are visible after join() returns. This matches the semantic + // guarantee provided by pthread_join() across supported architectures. + auto* version = reinterpret_cast*>(m->version_butex); + const int expected_version_int = static_cast(expected_version); + while (version->load(butil::memory_order_acquire) == expected_version_int) { + if (butex_wait(m->version_butex, expected_version_int, nullptr) < 0 && errno != EWOULDBLOCK && errno != EINTR) { return errno; } } - // Ensure all memory writes made by the joined bthread are visible to - // the joining thread after join returns. This matches the semantic - // guarantee provided by pthread_join() across supported architectures. - butil::atomic_thread_fence(butil::memory_order_acquire); if (return_value) { *return_value = nullptr; } @@ -729,7 +650,10 @@ bool TaskGroup::exists(bthread_t tid) { if (tid != 0) { // tid of bthread is never 0. TaskMeta* m = address_meta(tid); if (m != nullptr) { - return (*m->version_butex == get_version(tid)); + auto version = reinterpret_cast*>(m->version_butex); + // Only check liveness; unlike join(), no user data is acquired. + return static_cast(version->load(butil::memory_order_relaxed)) + == get_version(tid); } } return false; @@ -821,7 +745,7 @@ void TaskGroup::sched_to(TaskGroup** pg, TaskMeta* next_meta) { TaskMeta* const cur_meta = g->_cur_meta; int64_t now = butil::cpuwide_time_ns(); - CPUTimeStat cpu_time_stat = g->_cpu_time_stat.load_unsafe(); + CPUTimeStat cpu_time_stat = g->_cpu_time_stat.load_for_writer(); int64_t elp_ns = now - cpu_time_stat.last_run_ns(); cur_meta->stat.cputime_ns += elp_ns; // Update cpu_time_stat. @@ -1124,7 +1048,7 @@ int TaskGroup::usleep(TaskGroup** pg, uint64_t timeout_us) { bool erase_from_butex_because_of_interruption(ButexWaiter* bw); static int interrupt_and_consume_waiters( - bthread_t tid, ButexWaiter** pw, uint64_t* sleep_id) { + bthread_t tid, ButexWaiter** bw, uint64_t* sleep_id) { TaskMeta* const m = TaskGroup::address_meta(tid); if (m == nullptr) { return EINVAL; @@ -1132,7 +1056,7 @@ static int interrupt_and_consume_waiters( const uint32_t given_ver = get_version(tid); BAIDU_SCOPED_LOCK(m->version_lock); if (given_ver == *m->version_butex) { - *pw = m->current_waiter.exchange(nullptr, butil::memory_order_acquire); + *bw = m->current_waiter.exchange(nullptr, butil::memory_order_acquire); *sleep_id = m->current_sleep; m->current_sleep = 0; // only one stopper gets the sleep_id m->interrupted = true; diff --git a/src/bthread/task_group.h b/src/bthread/task_group.h index f48e02f5e9..bbd76ae74b 100644 --- a/src/bthread/task_group.h +++ b/src/bthread/task_group.h @@ -22,12 +22,13 @@ #ifndef BTHREAD_TASK_GROUP_H #define BTHREAD_TASK_GROUP_H -#include "butil/time.h" // cpuwide_time_ns +#include "butil/time.h" +#include "butil/synchronization/seqlock.h" #include "bthread/task_control.h" -#include "bthread/task_meta.h" // bthread_t, TaskMeta -#include "bthread/work_stealing_queue.h" // WorkStealingQueue -#include "bthread/remote_task_queue.h" // RemoteTaskQueue -#include "butil/resource_pool.h" // ResourceId +#include "bthread/task_meta.h" +#include "bthread/work_stealing_queue.h" +#include "bthread/remote_task_queue.h" +#include "butil/resource_pool.h" #include "bthread/parking_lot.h" #include "bthread/prime_offset.h" @@ -48,37 +49,6 @@ class ExitException : public std::exception { void* _value; }; -// Refer to https://rigtorp.se/isatomic/, On the modern CPU microarchitectures -// (Skylake and Zen 2) AVX/AVX2 128b/256b aligned loads and stores are atomic -// even though Intel and AMD officially doesn’t guarantee this. -// On X86, SSE instructions can ensure atomic loads and stores. -// Starting from Armv8.4-A, neon can ensure atomic loads and stores. -// Otherwise, use mutex to guarantee atomicity. -class AtomicInteger128 { -public: - struct BAIDU_CACHELINE_ALIGNMENT Value { - int64_t v1; - int64_t v2; - }; - - AtomicInteger128() = default; - explicit AtomicInteger128(Value value) : _value(value) {} - - Value load() const; - Value load_unsafe() const { - return _value; - } - - void store(Value value); - -private: - Value _value{}; - // Used to protect `_cpu_time_stat' on architectures without lock-free 128-bit atomics. - FastPthreadMutex _mutex; - // Sequence counter for RISC-V seqlock implementation. - uint64_t _seq = 0; -}; - // Thread-local group of tasks. // Notice that most methods involving context switching are static otherwise // pointer `this' may change after wakeup. The **pg parameters in following @@ -237,66 +207,83 @@ friend class TaskControl; // Last scheduling time, task type and cumulated CPU time. class CPUTimeStat { - static constexpr int64_t LAST_SCHEDULING_TIME_MASK = 0x7FFFFFFFFFFFFFFFLL; - static constexpr int64_t TASK_TYPE_MASK = 0x8000000000000000LL; public: - CPUTimeStat() : _last_run_ns_and_type(0), _cumulated_cputime_ns(0) {} - CPUTimeStat(AtomicInteger128::Value value) - : _last_run_ns_and_type(value.v1), _cumulated_cputime_ns(value.v2) {} - - // Convert to AtomicInteger128::Value for atomic operations. - explicit operator AtomicInteger128::Value() const { - return {_last_run_ns_and_type, _cumulated_cputime_ns}; + CPUTimeStat() : CPUTimeStat(0, 0, false) {} + + CPUTimeStat(int64_t last_run_ns, int64_t cumulated_cputime_ns, bool main_task) + : _cumulated_cputime_ns(cumulated_cputime_ns) + , _last_run_ns(last_run_ns) + , _main_task(main_task) {} + + CPUTimeStat(const CPUTimeStat& other) + : CPUTimeStat(other.last_run_ns(), + other.cumulated_cputime_ns(), + other.is_main_task()) {} + + CPUTimeStat& operator=(const CPUTimeStat& other) { + if (this != &other) { + _last_run_ns.store(other.last_run_ns(), + butil::memory_order_relaxed); + _cumulated_cputime_ns.store(other.cumulated_cputime_ns(), + butil::memory_order_relaxed); + _main_task.store(other.is_main_task(), butil::memory_order_relaxed); + } + return *this; } void set_last_run_ns(int64_t last_run_ns, bool main_task) { - _last_run_ns_and_type = (last_run_ns & LAST_SCHEDULING_TIME_MASK) | - (static_cast(main_task) << 63); + _last_run_ns.store(last_run_ns, butil::memory_order_relaxed); + _main_task.store(main_task, butil::memory_order_relaxed); } int64_t last_run_ns() const { - return _last_run_ns_and_type & LAST_SCHEDULING_TIME_MASK; - } - int64_t last_run_ns_and_type() const { - return _last_run_ns_and_type; + return _last_run_ns.load(butil::memory_order_relaxed); } bool is_main_task() const { - return _last_run_ns_and_type & TASK_TYPE_MASK; + return _main_task.load(butil::memory_order_relaxed); } void add_cumulated_cputime_ns(int64_t cputime_ns, bool main_task) { if (main_task) { return; } - _cumulated_cputime_ns += cputime_ns; + + _cumulated_cputime_ns.store(cumulated_cputime_ns() + cputime_ns, + butil::memory_order_relaxed); } int64_t cumulated_cputime_ns() const { - return _cumulated_cputime_ns; + return _cumulated_cputime_ns.load(butil::memory_order_relaxed); } private: - // The higher bit for task type, main task is 1, otherwise 0. - // Lowest 63 bits for last scheduling time. - int64_t _last_run_ns_and_type; - // Cumulated CPU time in nanoseconds. - int64_t _cumulated_cputime_ns; + // Cumulated non-main-task elapsed time in nanoseconds. + butil::atomic _cumulated_cputime_ns; + butil::atomic _last_run_ns; + butil::atomic _main_task; }; class AtomicCPUTimeStat { public: CPUTimeStat load() const { - return _cpu_time_stat.load(); + return _seqlock.load([&]() -> CPUTimeStat { + return _stat; + }); } - CPUTimeStat load_unsafe() const { - return _cpu_time_stat.load_unsafe(); + // For the owning writer only, with no concurrent writes. Copies fields + // with relaxed atomic loads but skips sequence validation. + CPUTimeStat load_for_writer() const { + return _stat; } - void store(CPUTimeStat cpu_time_stat) { - _cpu_time_stat.store(AtomicInteger128::Value(cpu_time_stat)); + void store(const CPUTimeStat& stat) { + _seqlock.store([this, &stat]() { + _stat = stat; + }); } private: - AtomicInteger128 _cpu_time_stat; + CPUTimeStat _stat; + butil::Seqlock<> _seqlock; }; // You shall use TaskControl::create_group to create new instance. diff --git a/src/bthread/task_meta.h b/src/bthread/task_meta.h index 2dae2fea27..7ca52e04c8 100644 --- a/src/bthread/task_meta.h +++ b/src/bthread/task_meta.h @@ -85,10 +85,18 @@ struct TaskMeta { // Scheduling of the thread can be delayed. bool about_to_quit{false}; - // [Not Reset] guarantee visibility of version_butex. + // [Not Reset] Serializes the version bump at bthread end (in task_runner) + // with accessors that validate the version before touching other fields of + // this TaskMeta (get_attr/set_stopped/interrupt/set_butex_waiter/...). It + // makes their "check version then read/write field" sequence atomic w.r.t. + // the bump, so they never operate on a slot that got recycled in between. pthread_spinlock_t version_lock{}; - - // [Not Reset] only modified by one bthread at any time, no need to be atomic + + // [Not Reset] Backed by a butex (internally `butil::atomic`). The version + // bump at bthread end is published with a release store, and join() observes + // it with an acquire load so the joined bthread's prior writes are visible + // after join() returns. All lock-free accesses must be atomic; liveness + // checks and reads before publishing a new task only need relaxed loads. uint32_t* version_butex{nullptr}; // The identifier. It does not have to be here, however many code is diff --git a/src/bthread/task_tracer.cpp b/src/bthread/task_tracer.cpp index 1cc5c05a2e..afb1642d99 100644 --- a/src/bthread/task_tracer.cpp +++ b/src/bthread/task_tracer.cpp @@ -35,6 +35,23 @@ namespace bthread { +namespace { + +// These slots must match the Linux x86_64 frame built by +// bthread_jump_fcontext in context.cpp. XMM6-XMM15 add 10 128-bit values, +// or 20 uintptr_t slots, ahead of the previously saved registers. +constexpr size_t kSavedXmmRegisterCount = 10; +constexpr size_t kSavedXmmRegisterBytes = 16; +constexpr size_t kSavedXmmRegisterSlots = + kSavedXmmRegisterCount * kSavedXmmRegisterBytes / sizeof(uintptr_t); +constexpr size_t kRbpContextSlot = 6 + kSavedXmmRegisterSlots; +constexpr size_t kRipContextSlot = 7 + kSavedXmmRegisterSlots; +#if UNW_VERSION_MAJOR >= 1 && UNW_VERSION_MINOR >= 7 +constexpr size_t kRspContextSlot = 8 + kSavedXmmRegisterSlots; +#endif + +} // namespace + DEFINE_uint32(signal_trace_timeout_ms, 50, "Timeout for signal trace in ms"); BUTIL_VALIDATE_GFLAG(signal_trace_timeout_ms, butil::PositiveInteger); // Note that SIGURG handler may be registered by some library such as cgo @@ -144,27 +161,10 @@ bool TaskTracer::Init() { } void TaskTracer::set_status(TaskStatus s, TaskMeta* m) { - CHECK_NE(TASK_STATUS_RUNNING, s) << "Use `set_running_status' instead"; - CHECK_NE(TASK_STATUS_END, s) << "Use `set_end_status_unsafe' instead"; - bool tracing = false; { BAIDU_SCOPED_LOCK(m->version_lock); - if (TASK_STATUS_UNKNOWN == m->status && TASK_STATUS_JUMPING == s) { - // Do not update status for jumping when bthread is ending. - return; - } - - tracing = m->traced; - // bthread is scheduled for the first time. - if (TASK_STATUS_READY == s && nullptr == m->stack) { - m->status = TASK_STATUS_FIRST_READY; - } else { - m->status = s; - } - if (TASK_STATUS_CREATED == s) { - m->worker_tid = pthread_t{}; - } + tracing = set_status_unsafe(s, m); } // Make sure bthread does not jump stack when it is being traced. @@ -173,6 +173,27 @@ void TaskTracer::set_status(TaskStatus s, TaskMeta* m) { } } +bool TaskTracer::set_status_unsafe(TaskStatus s, TaskMeta* m) { + CHECK_NE(TASK_STATUS_RUNNING, s) << "Use `set_running_status' instead"; + CHECK_NE(TASK_STATUS_END, s) << "Use `set_end_status_unsafe' instead"; + + if (TASK_STATUS_UNKNOWN == m->status && TASK_STATUS_JUMPING == s) { + // Do not update status for jumping when bthread is ending. + return false; + } + + // A bthread is scheduled for the first time. + if (TASK_STATUS_READY == s && m->stack == nullptr) { + m->status = TASK_STATUS_FIRST_READY; + } else { + m->status = s; + } + if (TASK_STATUS_CREATED == s) { + m->worker_tid = pthread_t{}; + } + return m->traced; +} + void TaskTracer::set_running_status(pthread_t worker_tid, TaskMeta* m) { BAIDU_SCOPED_LOCK(m->version_lock); m->worker_tid = worker_tid; @@ -280,16 +301,16 @@ unw_cursor_t TaskTracer::MakeCursor(bthread_fcontext_t fcontext) { // Only need RBP, RIP, RSP on x86_64. // The base pointer (RBP). - if (unw_set_reg(&cursor, UNW_X86_64_RBP, regs[6]) != 0) { + if (unw_set_reg(&cursor, UNW_X86_64_RBP, regs[kRbpContextSlot]) != 0) { LOG(ERROR) << "Fail to set RBP"; } // The instruction pointer (RIP). - if (unw_set_reg(&cursor, UNW_REG_IP, regs[7]) != 0) { + if (unw_set_reg(&cursor, UNW_REG_IP, regs[kRipContextSlot]) != 0) { LOG(ERROR) << "Fail to set RIP"; } #if UNW_VERSION_MAJOR >= 1 && UNW_VERSION_MINOR >= 7 // The stack pointer (RSP). - if (unw_set_reg(&cursor, UNW_REG_SP, regs[8]) != 0) { + if (unw_set_reg(&cursor, UNW_REG_SP, regs[kRspContextSlot]) != 0) { LOG(ERROR) << "Fail to set RSP"; } #endif diff --git a/src/bthread/task_tracer.h b/src/bthread/task_tracer.h index 8844413af6..2bf4099b5f 100644 --- a/src/bthread/task_tracer.h +++ b/src/bthread/task_tracer.h @@ -39,6 +39,7 @@ class TaskTracer { bool Init(); // Set the status to `s'. void set_status(TaskStatus s, TaskMeta* meta); + static bool set_status_unsafe(TaskStatus s, TaskMeta* meta); static void set_running_status(pthread_t worker_tid, TaskMeta* meta); static bool set_end_status_unsafe(TaskMeta* m); diff --git a/src/butil/containers/optional.h b/src/butil/containers/optional.h index bde0ad2f3c..e196958496 100644 --- a/src/butil/containers/optional.h +++ b/src/butil/containers/optional.h @@ -22,7 +22,7 @@ #include "butil/memory/manual_constructor.h" // The `optional` for managing an optional contained value, -// i.e. a value that may or may not be present, is a C++11 +// i.e. a value that may or may not be present, is a C++14 // compatible version of the C++17 `std::optional` abstraction. // After C++17, `optional` is an alias for `std::optional`. @@ -145,33 +145,33 @@ class optional { optional(optional&& rhs) noexcept = default; - template ::value && std::is_constructible::value && !internal::is_constructible_convertible_from_optional::value && - std::is_convertible::value, bool>::type = false> + std::is_convertible::value, bool> = false> optional(const optional& rhs) : _engaged(rhs.has_value()) { if (_engaged) { _storage.Init(*rhs); } } - template ::value && std::is_constructible::value && !internal::is_constructible_convertible_from_optional::value && - !std::is_convertible::value, bool>::type = false> + !std::is_convertible::value, bool> = false> explicit optional(const optional& rhs) : _engaged(rhs.has_value()) { if (_engaged) { _storage.Init(*rhs); } } - template ::value && std::is_constructible::value && !internal::is_constructible_convertible_from_optional::value && - std::is_convertible::value, bool>::type = false> + std::is_convertible::value, bool> = false> optional(optional&& rhs) : _engaged(rhs.has_value()) { if (_engaged) { _storage.Init(std::move(*rhs)); @@ -179,11 +179,11 @@ class optional { } } - template ::value && std::is_constructible::value && !internal::is_constructible_convertible_from_optional::value && - !std::is_convertible::value, bool>::type = false> + !std::is_convertible::value, bool> = false> explicit optional(optional&& rhs) : _engaged(rhs.has_value()) { if (_engaged) { _storage.Init(std::move(*rhs)); @@ -199,33 +199,34 @@ class optional { _storage.Init(std::move(value)); } - template ::value>* = nullptr> - explicit optional(const in_place_t, Args&&... args) : _engaged(true) { + template ::value, bool> = false> + explicit optional(in_place_t, Args&&... args) : _engaged(true) { _storage.Init(std::forward(args)...); } - template &, Args&&...>::value>::type> + template &, Args&&...>::value, + bool> = false> optional(in_place_t, std::initializer_list il, Args&&... args) : _engaged(true) { _storage.Init(il, std::forward(args)...); } - template ::type>::value && !std::is_same, typename std::decay::type>::value && std::is_constructible::value && - std::is_convertible::value, bool>::type = false> + std::is_convertible::value, bool> = false> optional(U&& v) : _engaged(true) { _storage.Init(std::forward(v)); } - template ::type>::value && !std::is_same, typename std::decay::type>::value && std::is_constructible::value && - !std::is_convertible::value, bool>::type = false> + !std::is_convertible::value, bool> = false> explicit optional(U&& v) : _engaged(true) { _storage.Init(std::forward(v)); } @@ -244,11 +245,11 @@ class optional { optional& operator=(optional&& rhs) = default; // Value assignment operators - template , typename std::decay::type>::value && !std::is_same, typename remove_cvref::type>::value && std::is_constructible::value && std::is_assignable::value && - (!std::is_scalar::value || !std::is_same::type>::value)>::type> + (!std::is_scalar::value || !std::is_same::type>::value)>> optional& operator=(U&& v) { reset(); _storage.Init(std::forward(v)); @@ -256,11 +257,11 @@ class optional { return *this; } - template ::value && !internal::is_constructible_convertible_assignable_from_optional::value && std::is_constructible::value && - std::is_assignable::value>::type> + std::is_assignable::value>> optional& operator=(const optional& rhs) { if (rhs) { operator=(*rhs); @@ -270,11 +271,11 @@ class optional { return *this; } - template ::value && !internal::is_constructible_convertible_assignable_from_optional::value && std::is_constructible::value && - std::is_assignable::value>::type> + std::is_assignable::value>> optional& operator=(optional&& rhs) { if (rhs) { operator=(std::move(*rhs)); diff --git a/src/butil/memory/scope_guard.h b/src/butil/memory/scope_guard.h index 377819b5db..27fc2ba6de 100644 --- a/src/butil/memory/scope_guard.h +++ b/src/butil/memory/scope_guard.h @@ -24,7 +24,7 @@ namespace butil { template::value>> + typename = std::enable_if_t::value>> class ScopeGuard; template @@ -33,7 +33,7 @@ ScopeGuard MakeScopeGuard(Callback&& callback) noexcept; // ScopeGuard is a simple implementation to guarantee that // a function is executed upon leaving the current scope. template -class ScopeGuard { +class ScopeGuard { public: ScopeGuard(ScopeGuard&& other) noexcept : _callback(std::move(other._callback)) diff --git a/src/butil/processor.h b/src/butil/processor.h new file mode 100644 index 0000000000..9219d31630 --- /dev/null +++ b/src/butil/processor.h @@ -0,0 +1,48 @@ +// 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 BUTIL_PROCESSOR_H +#define BUTIL_PROCESSOR_H + +#include "butil/build_config.h" + +// Pause instruction to prevent excess processor bus usage, only works in GCC +#ifndef cpu_relax +#if defined(ARCH_CPU_ARM_FAMILY) +#define cpu_relax() asm volatile("yield\n": : :"memory") +#elif defined(ARCH_CPU_RISCV_FAMILY) +// Use the pause hint (Zihintpause extension). Encoding 0x0100000F +// (fence 0, 1) is a HINT on all RISC-V implementations: it never traps +// and is ignored on CPUs without Zihintpause. On CPUs with Zihintpause +// it provides a multi-cycle stall hint that reduces power and improves +// resource fairness during spin-wait loops. Matches the Linux kernel's +// RISC-V cpu_relax() behavior. .word is used instead of .insn or the +// pause mnemonic for maximum assembler compatibility. +# define cpu_relax() asm volatile(".word 0x0100000f\n": : :"memory") +#elif defined(ARCH_CPU_LOONGARCH64_FAMILY) +# define cpu_relax() asm volatile("nop\n": : :"memory") +#else +# define cpu_relax() asm volatile("pause\n": : :"memory") +#endif +#endif // cpu_relax + +// Compile read-write barrier +#ifndef barrier +#define barrier() asm volatile("": : :"memory") +#endif // barrier + +#endif // BUTIL_PROCESSOR_H \ No newline at end of file diff --git a/src/butil/synchronization/seqlock.h b/src/butil/synchronization/seqlock.h new file mode 100644 index 0000000000..3e91dcb55f --- /dev/null +++ b/src/butil/synchronization/seqlock.h @@ -0,0 +1,302 @@ +// 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 BUTIL_SYNCHRONIZATION_SEQLOCK_H +#define BUTIL_SYNCHRONIZATION_SEQLOCK_H + +#include +#include +#include +#include +#include + +#include "butil/atomicops.h" +#include "butil/compiler_specific.h" +#include "butil/macros.h" +#include "butil/processor.h" + +namespace butil { +namespace internal { + +// Detects std::reference_wrapper. +template +struct IsReferenceWrapper : std::false_type {}; +template +struct IsReferenceWrapper> : std::true_type {}; + +// A sequence counter for consistent reads without acquiring a reader mutex +// around caller-owned payloads. This is an implementation detail of butil::Seqlock; +// use butil::Seqlock instead. +// +// SeqCounter does not protect payload accesses from C++ data races. Callers must +// access shared payloads atomically (typically with relaxed ordering) or use +// another mechanism that makes concurrent accesses valid. +// +// +// Cache-line aligned so the sequence counter does not falsely share a line +// with adjacent data (e.g. the writer mutex in Seqlock or a neighbouring +// payload), which would bounce the line between readers and writers. +class BAIDU_CACHELINE_ALIGNMENT SeqCounter { +public: + SeqCounter() : _seq(0) {} + DISALLOW_COPY_AND_ASSIGN(SeqCounter); + + // Repeatedly invoke `load_payload` until it observes one consistent version. + // `load_payload` may run several times when it races with a writer, so it + // must be cheap and side-effect free: only read the shared payload and + // return a copy, never mutate observable state. It must also access the + // shared payload without causing a C++ data race (relaxed atomics). + // + // The callback MUST return an owning value (a copy of the data), never a + // pointer, reference, or view (string_view, span, reference_wrapper, ...) + // into the payload. The consistency guarantee covers only the bytes copied + // out before the validating load: once load() returns, a later writer may + // mutate the payload, so any handle that still points into it no longer + // refers to a consistent snapshot. + template + typename std::decay()())>::type + load(Load&& load_payload) const { + typedef typename std::decay()())>::type Result; + static_assert(!std::is_void::value, + "SeqCounter load callback must return a value"); + static_assert(!std::is_pointer::value, + "SeqCounter load callback must return an owning value, not " + "a pointer into the payload: the pointee can be mutated by " + "a later writer, so it is not a consistent snapshot"); + static_assert(!IsReferenceWrapper::value, + "SeqCounter load callback must return an owning value, not " + "a std::reference_wrapper into the payload: the referent " + "can be mutated by a later writer, so it is not a " + "consistent snapshot"); + + while (true) { + // Wait for any active writer, then read on an even sequence. + uint64_t seq; + while ((seq = _seq.load(memory_order_acquire)) & 1) { + cpu_relax(); + } + + Result result = load_payload(); + + // Keep the payload reads above before the sequence validation + // below; retry if a writer intervened. + atomic_thread_fence(memory_order_acquire); + if (_seq.load(memory_order_relaxed) == seq) { + return result; + } + } + } + + // Invoke a payload writer inside one write section. Writers must be serialized + // externally and must not nest write sections. The section is closed via RAII + // even if `store_payload` throws, so an exception does not leave the sequence + // permanently odd. A throwing writer may leave the payload partially updated, + // so callers must ensure any partial state is still safe (non-crashing) to read. + template + void store(Store&& store_payload) { + WriteGuard guard(*this); + store_payload(); + } + +private: + // RAII scope for a write section. + class WriteGuard { + public: + explicit WriteGuard(SeqCounter& sc) + : _seq_counter(&sc) + , _next_seq(sc._seq.load(memory_order_relaxed) + 2) { + // Writers are serialized, so no atomic read-modify-write is needed. + // Keep the next even sequence for publication when the guard exits. + _seq_counter->_seq.store(_next_seq - 1, memory_order_relaxed); + // Order the odd-sequence store before the payload stores that + // follow. A release fence is a StoreStore (+LoadStore) barrier: it + // keeps any store after the fence (the payload writes) from being + // reordered ahead of stores before it (the odd sequence), so no + // reader can observe a payload write without also observing the odd + // sequence. Equivalently, this release fence pairs -- through the + // payload atomics -- with the acquire fence in load(): if a reader + // reads a payload value published after this fence, that fence + // synchronizes-with the reader's acquire fence, so the reader is + // guaranteed to also see the odd sequence on its validating load + // and retry. The acquire half of an acq_rel fence would add + // nothing here (no prior load needs ordering), so release alone is + // sufficient and states the intent precisely. + atomic_thread_fence(memory_order_release); + } + + DISALLOW_COPY_AND_ASSIGN(WriteGuard); + + ~WriteGuard() { + // Publish the payload stores and leave the write section. + _seq_counter->_seq.store(_next_seq, memory_order_release); + } + + private: + SeqCounter* _seq_counter; + uint64_t _next_seq; + }; + + atomic _seq; +}; + +} // namespace internal + +// A sequence lock built on top of internal::SeqCounter. +// +// Seqlock<> : single-writer, no writer mutex. Backed directly by SeqCounter; +// the caller must serialize writers with appropriate synchronization +// if ownership is transferred between threads. +// +// Seqlock : multi-writer. Owns a writer Mutex so that concurrent +// store() calls are serialized automatically (this is the +// Linux seqlock_t = SeqCounter + lock). +// +// In both forms readers do not acquire a mutex or prevent writers from entering +// a write section. Reads are NOT lock-free: they spin while the sequence is odd +// and retry if a write intervenes. A suspended writer can stall all readers, and +// continuous writes can starve readers; there is no bounded completion time. +// +// REQUIRED: write sections must not nest. A store() callback must not call load() +// or store() on the same Seqlock, or wait for work that needs to read or write it. +// In particular, a signal handler must not access the same Seqlock when it has +// interrupted a writer: the interrupted write cannot finish until the handler +// returns, so a reader in the handler would spin forever. Keep write sections +// short and avoid blocking or yielding inside them. +// +// Example: publish a (x, y) pair atomically so readers never see a torn mix. +// +// // the caller-owned payload +// struct Point { +// butil::atomic x{0}; +// butil::atomic y{0}; +// }; +// Point point; +// // single writer; use Seqlock if several threads may write. +// butil::Seqlock<> seqlock; +// +// // Writer: the whole update is published as one consistent version. +// void set(int64_t x, int64_t y) { +// seqlock.store([&] { +// point.x.store(x, butil::memory_order_relaxed); +// point.y.store(y, butil::memory_order_relaxed); +// }); +// } +// +// // Reader: load() retries internally until it copies out one consistent +// // version, then returns whatever the callback returned. +// std::pair get() { +// return seqlock.load([&] { +// return std::make_pair(point.x.load(butil::memory_order_relaxed), +// point.y.load(butil::memory_order_relaxed)); +// }); +// } +// +// REQUIRED: the payload accessed inside the load/store callbacks MUST be +// atomic (e.g. butil::atomic fields, typically read/written with +// memory_order_relaxed). This is not just a style preference -- the reader +// deliberately reads the payload while a writer may be mutating it, so the +// accesses are concurrent by design. Correctness relies on it in two ways: +// +// 1. Data race / UB. A non-atomic object read while another thread writes it +// is a C++ data race, i.e. undefined behavior. The compiler is then free +// to tear or fuse the access, invent extra reads, or hoist/sink it across +// the fences below. The sequence-validation retry cannot rescue this: it +// only tells you *whether* to retry, it cannot un-corrupt a value the +// compiler already mangled, so you may return garbage even on a "clean" +// (seq unchanged) read. +// 2. Ordering. The release fence on the write side pairs with the acquire fence +// on the read side through the payload atomics (fence-fence synchronization). +// If the payload is not atomic that pairing does not hold, so "observe a +// payload write => observe the odd sequence and retry" is no longer guaranteed +// and a reader can silently accept a stale or half-written snapshot. +// +// If you cannot make the payload atomic, do not use Seqlock -- use a Mutex or +// an RWLock instead. +// +// REQUIRED: the load() callback must return an OWNING value (a copy of the +// data), never a pointer, reference, or view (string_view, span, +// reference_wrapper, ...) into the payload. load() only guarantees that the +// bytes copied out before its validating load are consistent; after load() +// returns a writer may mutate the payload again, so any handle still pointing +// into it is no longer a consistent snapshot. +// +// Best fit: a small payload that is read far more often than written. The read +// path copies the whole payload out and retries the copy whenever a write races +// it, so a large payload makes both the copy and the retries expensive. For a +// large or heap-owning payload prefer a RCU/RWLock scheme instead. +template +class Seqlock; + +// Single-writer specialization: no mutex, delegates straight to SeqCounter. +template <> +class Seqlock { +public: + Seqlock() = default; + DISALLOW_COPY_AND_ASSIGN(Seqlock); + + // Consistent read without acquiring a mutex; may spin/retry indefinitely. + template + typename std::decay()())>::type + load(Load&& load_payload) const { + return _seq.load(std::forward(load_payload)); + } + + // Single writer only: the caller MUST ensure there is no concurrent or nested + // store(). The write section is closed via RAII even if store_payload + // throws. + template + void store(Store&& store_payload) { + _seq.store(std::forward(store_payload)); + } + +private: + internal::SeqCounter _seq; +}; + +// Multi-writer specialization: SeqCounter + a writer Mutex. +// +// Mutex must be default-constructible and satisfy the C++ Lockable +// requirements (lock()/unlock()), e.g. butil::Mutex. +template +class Seqlock { +public: + Seqlock() = default; + DISALLOW_COPY_AND_ASSIGN(Seqlock); + + // Consistent read without acquiring a mutex; may spin/retry indefinitely. + template + typename std::decay()())>::type + load(Load&& load_payload) const { + return _seq.load(std::forward(load_payload)); + } + + // Serialized write: acquires the mutex, then runs the payload writer in + // one write section. + template + void store(Store&& store_payload) { + std::lock_guard lk(_mutex); + _seq.store(std::forward(store_payload)); + } + +private: + internal::SeqCounter _seq; + Mutex _mutex; +}; + +} // namespace butil + +#endif // BUTIL_SYNCHRONIZATION_SEQLOCK_H diff --git a/src/bvar/multi_dimension.h b/src/bvar/multi_dimension.h index 2eb31b805d..bcf46c08d2 100644 --- a/src/bvar/multi_dimension.h +++ b/src/bvar/multi_dimension.h @@ -190,10 +190,10 @@ class MultiDimension : public MVariable { dump_impl(Dumper* dumper, const DumpOptions* options); void make_dump_key(std::ostream& os, const key_type& labels_value, - const std::string& suffix = "", int quantile = 0); + const std::string& suffix = "", double quantile = 0); void make_labels_kvpair_string( - std::ostream& os, const key_type& labels_value, int quantile); + std::ostream& os, const key_type& labels_value, double quantile); template diff --git a/src/bvar/multi_dimension_inl.h b/src/bvar/multi_dimension_inl.h index 7ff15230cb..108a66b0a0 100644 --- a/src/bvar/multi_dimension_inl.h +++ b/src/bvar/multi_dimension_inl.h @@ -258,85 +258,103 @@ MultiDimension::dump_impl(Dumper* dumper, const DumpOptions* if (label_names.empty()) { return 0; } - size_t n = 0; - // To meet prometheus specification, we must guarantee no second TYPE line for one metric name - - // latency comment - dumper->dump_comment(this->name() + "_latency", METRIC_TYPE_GAUGE); - for (auto &label_name : label_names) { + // The latency of one quantile. The quantile must be a fraction to meet + // prometheus specification, e.g. 0.99 for p99. + struct LatencyPercentile { + double quantile; + int64_t latency; + }; + // All the values dumped for one label set. + struct DumpedStats { + const key_type* label_name; + LatencyPercentile latency_percentiles[5]; + int64_t avg_latency; + int64_t max_latency; + int64_t qps; + int64_t count; + }; + // Read all the values in one traversal, so that a LatencyRecorder is looked + // up only once no matter how many metrics are dumped for it. Keep the values + // instead of the LatencyRecorder pointers, which delete_stats() may free. + std::vector stats_list; + stats_list.reserve(label_names.size()); + for (const auto& label_name : label_names) { bvar::LatencyRecorder* bvar = get_stats_impl(label_name); if (!bvar) { continue; } - - // latency - std::ostringstream oss_latency_key; - make_dump_key(oss_latency_key, label_name, "_latency"); - if (dumper->dump_mvar(oss_latency_key.str(), std::to_string(bvar->latency()))) { - n++; + DumpedStats stats{}; + stats.label_name = &label_name; + stats.latency_percentiles[0].quantile = FLAGS_bvar_latency_p1 / 100.0; + stats.latency_percentiles[1].quantile = FLAGS_bvar_latency_p2 / 100.0; + stats.latency_percentiles[2].quantile = FLAGS_bvar_latency_p3 / 100.0; + stats.latency_percentiles[3].quantile = 0.999; + stats.latency_percentiles[4].quantile = 0.9999; + for (auto& lp : stats.latency_percentiles) { + lp.latency = bvar->latency_percentile(lp.quantile); } - // latency_percentiles - // p1/p2/p3 - int latency_percentiles[3] {FLAGS_bvar_latency_p1, FLAGS_bvar_latency_p2, FLAGS_bvar_latency_p3}; - for (auto lp : latency_percentiles) { - std::ostringstream oss_lp_key; - make_dump_key(oss_lp_key, label_name, "_latency", lp); - if (dumper->dump_mvar(oss_lp_key.str(), std::to_string(bvar->latency_percentile(lp / 100.0)))) { + stats.avg_latency = bvar->latency(); + stats.max_latency = bvar->max_latency(); + stats.qps = bvar->qps(); + stats.count = bvar->count(); + stats_list.push_back(stats); + } + + size_t n = 0; + + // To meet prometheus specification, we must guarantee no second TYPE line for one metric name + + // latency comment + dumper->dump_comment(this->name() + "_latency", METRIC_TYPE_GAUGE); + for (const auto& stats : stats_list) { + for (const auto& lp : stats.latency_percentiles) { + std::ostringstream oss_latency_key; + make_dump_key(oss_latency_key, *stats.label_name, "_latency", lp.quantile); + if (dumper->dump_mvar(oss_latency_key.str(), std::to_string(lp.latency))) { n++; } } - // 999 - std::ostringstream oss_p999_key; - make_dump_key(oss_p999_key, label_name, "_latency", 999); - if (dumper->dump_mvar(oss_p999_key.str(), std::to_string(bvar->latency_percentile(0.999)))) { - n++; - } - // 9999 - std::ostringstream oss_p9999_key; - make_dump_key(oss_p9999_key, label_name, "_latency", 9999); - if (dumper->dump_mvar(oss_p9999_key.str(), std::to_string(bvar->latency_percentile(0.9999)))) { + } + + // avg_latency comment + // The average latency has to be a separate metric rather than a series of + // `_latency` without a quantile label, otherwise an aggregation over + // `_latency` would silently mix the average into the percentiles. + dumper->dump_comment(this->name() + "_avg_latency", METRIC_TYPE_GAUGE); + for (const auto& stats : stats_list) { + std::ostringstream oss_avg_latency_key; + make_dump_key(oss_avg_latency_key, *stats.label_name, "_avg_latency"); + if (dumper->dump_mvar(oss_avg_latency_key.str(), std::to_string(stats.avg_latency))) { n++; } } // max_latency comment dumper->dump_comment(this->name() + "_max_latency", METRIC_TYPE_GAUGE); - for (auto &label_name : label_names) { - LatencyRecorder* bvar = get_stats_impl(label_name); - if (nullptr == bvar) { - continue; - } + for (const auto& stats : stats_list) { std::ostringstream oss_max_latency_key; - make_dump_key(oss_max_latency_key, label_name, "_max_latency"); - if (dumper->dump_mvar(oss_max_latency_key.str(), std::to_string(bvar->max_latency()))) { + make_dump_key(oss_max_latency_key, *stats.label_name, "_max_latency"); + if (dumper->dump_mvar(oss_max_latency_key.str(), std::to_string(stats.max_latency))) { n++; } } // qps comment dumper->dump_comment(this->name() + "_qps", METRIC_TYPE_GAUGE); - for (auto &label_name : label_names) { - LatencyRecorder* bvar = get_stats_impl(label_name); - if (nullptr == bvar) { - continue; - } + for (const auto& stats : stats_list) { std::ostringstream oss_qps_key; - make_dump_key(oss_qps_key, label_name, "_qps"); - if (dumper->dump_mvar(oss_qps_key.str(), std::to_string(bvar->qps()))) { + make_dump_key(oss_qps_key, *stats.label_name, "_qps"); + if (dumper->dump_mvar(oss_qps_key.str(), std::to_string(stats.qps))) { n++; } } // count comment dumper->dump_comment(this->name() + "_count", METRIC_TYPE_COUNTER); - for (auto &label_name : label_names) { - LatencyRecorder* bvar = get_stats_impl(label_name); - if (nullptr == bvar) { - continue; - } + for (const auto& stats : stats_list) { std::ostringstream oss_count_key; - make_dump_key(oss_count_key, label_name, "_count"); - if (dumper->dump_mvar(oss_count_key.str(), std::to_string(bvar->count()))) { + make_dump_key(oss_count_key, *stats.label_name, "_count"); + if (dumper->dump_mvar(oss_count_key.str(), std::to_string(stats.count))) { n++; } } @@ -345,7 +363,7 @@ MultiDimension::dump_impl(Dumper* dumper, const DumpOptions* template void MultiDimension::make_dump_key(std::ostream& os, const key_type& labels_value, - const std::string& suffix, int quantile) { + const std::string& suffix, double quantile) { os << this->name(); if (!suffix.empty()) { os << suffix; @@ -354,8 +372,9 @@ void MultiDimension::make_dump_key(std::ostream& os, const k } template -void MultiDimension::make_labels_kvpair_string( - std::ostream& os, const key_type& labels_value, int quantile) { +void MultiDimension::make_labels_kvpair_string(std::ostream& os, + const key_type& labels_value, + double quantile) { os << "{"; auto label_key = this->_labels.cbegin(); auto label_value = labels_value.cbegin(); @@ -365,6 +384,8 @@ void MultiDimension::make_labels_kvpair_string( os << comma << label_key->c_str() << "=\"" << label_value->c_str() << "\""; comma[0] = ','; } + // The `quantile` label must be parsable as a float, so a non-positive + // `quantile` means "this metric is not a quantile series". if (quantile > 0) { os << comma << "quantile=\"" << quantile << "\""; } diff --git a/src/bvar/reducer.h b/src/bvar/reducer.h index 9fdb1960b7..8b40fe3f81 100644 --- a/src/bvar/reducer.h +++ b/src/bvar/reducer.h @@ -30,6 +30,7 @@ #include "bvar/detail/series.h" #include "bvar/window.h" #if WITH_BABYLON_COUNTER +#include // std::max std::min #include "babylon/concurrent/counter.h" #endif // WITH_BABYLON_COUNTER @@ -50,6 +51,33 @@ class SeriesSamplerImpl : public Sampler { }; #if WITH_BABYLON_COUNTER +// babylon counters constrain their value type with a static_assert inside the +// class body, which is a hard error rather than a substitution failure. Probing +// them with std::is_constructible<> hence breaks the build for the types they +// reject, instead of falling back to the generic implementation. Mirror the +// constraint here (see babylon/concurrent/counter.h) so that a rejected type is +// never used to instantiate a babylon counter. +// NOTE: the constraint has to be checked in two steps rather than in a single +// expression, because sizeof() can not be applied to void or to an incomplete +// type, and `&&' does not help: short-circuiting is about evaluation, an +// ill-formed sizeof() is still an error. +template ::value || + std::is_floating_point::value> +struct IsBabylonCounterSupported : std::false_type {}; + +template +struct IsBabylonCounterSupported + : butil::integral_constant {}; + +// `void` if the babylon counter supports T, a substitution failure otherwise. +// Selects the babylon-backed partial specializations of Adder/Maxer/Miner below. +// NOTE: resolving to the *type* is a MUST. std::enable_if itself is a type +// no matter what `cond` is, and Adder means Adder, so specializing on +// std::enable_if instead of std::enable_if_t silently never matches. +template +using EnableIfBabylonCounter = + std::enable_if_t::value>; + template class BabylonVariable: public Variable { public: @@ -58,12 +86,12 @@ class BabylonVariable: public Variable { BabylonVariable() = default; - template::value, bool>::type = false> + template::value, bool> = false> BabylonVariable(U) {} // For Maxer. - template::value, bool>::type = false> + template::value, bool> = false> BabylonVariable(U default_value) : _counter(default_value) {} DISALLOW_COPY_AND_MOVE(BabylonVariable); @@ -93,17 +121,21 @@ class BabylonVariable: public Variable { } T get_value() const { + CHECK(!(butil::is_same::value) || nullptr == _sampler) + << "You should not call Reducer<" << butil::class_name_str() + << ", " << butil::class_name_str() << ">::get_value() when a" + << " Window<> is used because the operator does not have inverse."; return _counter.value(); } T reset() { - if (BAIDU_UNLIKELY((!butil::is_same::value))) { - CHECK(false) << "You should not call Reducer<" << butil::class_name_str() - << ", " << butil::class_name_str() << ">::get_value() when a" - << " Window<> is used because the operator does not have inverse."; - return get_value(); - } - + // Unlike AgentCombiner::reset_all_agents(), reading and clearing the babylon + // counter are two separate steps, so values added in between are lost. This + // affects both explicit reset() by users and periodic sampling by the sampler + // thread (e.g. Window/series sampling), where concurrent additions around each + // reset may be dropped. This is an accepted trade-off for the babylon backend: + // in statistics/monitoring scenarios such minor loss does not change the overall + // trend, so slightly inaccurate samples are acceptable. T result = _counter.value(); _counter.reset(); return result; @@ -370,15 +402,13 @@ class Adder : public Reducer, detail::MinusFrom > { #if WITH_BABYLON_COUNTER // Numerical types supported by babylon counter. template -class Adder>::value>> +class Adder> : public detail::BabylonVariable, detail::AddTo, detail::MinusFrom> { public: typedef T value_type; -private: typedef detail::BabylonVariable, detail::AddTo, detail::MinusFrom> Base; -public: typedef detail::AddTo Op; typedef detail::MinusFrom InvOp; typedef typename Base::sampler_type sampler_type; @@ -449,28 +479,25 @@ class ConcurrentMaxer : public babylon::GenericsConcurrentMaxer { ConcurrentMaxer(T default_value) : _default_value(default_value) {} T value() const { - T result; - if (!Base::value(result)) { - return _default_value; - } + // Base::value() leaves `result' untouched if nothing was counted. + T result = _default_value; + Base::value(result); return std::max(result, _default_value); } private: - T _default_value{0}; + T _default_value{std::numeric_limits::min()}; }; } // namespace detail // Numerical types supported by babylon counter. template -class Maxer>::value>> +class Maxer> : public detail::BabylonVariable, detail::MaxTo, detail::VoidOp> { public: typedef T value_type; -private: typedef detail::BabylonVariable, detail::MaxTo, detail::VoidOp> Base; -public: typedef detail::MaxTo Op; typedef detail::VoidOp InvOp; typedef typename Base::sampler_type sampler_type; @@ -527,17 +554,35 @@ class Miner : public Reducer > { }; #if WITH_BABYLON_COUNTER +namespace detail { +// The min counterpart of ConcurrentMaxer, see there for why the default value is +// needed. +template +class ConcurrentMiner : public babylon::GenericsConcurrentMiner { + typedef babylon::GenericsConcurrentMiner Base; +public: + ConcurrentMiner() = default; + explicit ConcurrentMiner(T default_value) : _default_value(default_value) {} + + T value() const { + T result = _default_value; + Base::value(result); + return std::min(result, _default_value); + } +private: + T _default_value{std::numeric_limits::max()}; +}; +} // namespace detail + // Numerical types supported by babylon counter. template -class Miner>::value>> - : public detail::BabylonVariable, +class Miner> + : public detail::BabylonVariable, detail::MinTo, detail::VoidOp> { public: typedef T value_type; -private: - typedef detail::BabylonVariable, + typedef detail::BabylonVariable, detail::MinTo, detail::VoidOp> Base; -public: typedef detail::MinTo Op; typedef detail::VoidOp InvOp; typedef typename Base::sampler_type sampler_type; diff --git a/test/BUILD.bazel b/test/BUILD.bazel index 8706f63f2a..1aecc39987 100644 --- a/test/BUILD.bazel +++ b/test/BUILD.bazel @@ -125,9 +125,10 @@ TEST_BUTIL_SOURCES = [ "recordio_unittest.cpp", #"popen_unittest.cpp", "bounded_queue_unittest.cc", - "butil_unittest_main.cpp", "scope_guard_unittest.cpp", "optional_unittest.cpp", + "seqlock_unittest.cpp", + "butil_unittest_main.cpp", ] + select({ "@bazel_tools//tools/osx:darwin_x86_64": [], "//conditions:default": [ diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 6aebe271f4..71baf85df4 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -181,6 +181,7 @@ SET(TEST_BUTIL_SOURCES ${PROJECT_SOURCE_DIR}/test/scoped_locale.cc ${PROJECT_SOURCE_DIR}/test/scope_guard_unittest.cpp ${PROJECT_SOURCE_DIR}/test/optional_unittest.cpp + ${PROJECT_SOURCE_DIR}/test/seqlock_unittest.cpp ${PROJECT_SOURCE_DIR}/test/butil_unittest_main.cpp ) diff --git a/test/Makefile b/test/Makefile index f2348e7fa3..de4d566bcd 100644 --- a/test/Makefile +++ b/test/Makefile @@ -150,9 +150,10 @@ TEST_BUTIL_SOURCES = \ scoped_locale.cc \ popen_unittest.cpp \ bounded_queue_unittest.cc \ - butil_unittest_main.cpp \ - scope_guard_unittest.cpp \ - optional_unittest.cpp + scope_guard_unittest.cpp \ + optional_unittest.cpp \ + seqlock_unittest.cpp \ + butil_unittest_main.cpp ifeq ($(SYSTEM), Linux) TEST_BUTIL_SOURCES += test_file_util_linux.cc \ diff --git a/test/brpc_builtin_service_unittest.cpp b/test/brpc_builtin_service_unittest.cpp index e496c05f57..0d1ca248dd 100644 --- a/test/brpc_builtin_service_unittest.cpp +++ b/test/brpc_builtin_service_unittest.cpp @@ -60,6 +60,17 @@ DEFINE_bool(foo, false, "Flags for UT"); BRPC_VALIDATE_GFLAG(foo, brpc::PassValidate); +// A reloadable string gflag so that FlagsService is able to modify its +// value via ?setvalue=. String flags must register the validator manually, +// see comments in butil/reloadable_flags.h. +DEFINE_string(reloadable_string_flag_for_ut, "", "Flags for UT"); +static bool PassValidateStringFlag(const char*, const std::string&) { + return true; +} +const bool ALLOW_UNUSED dummy_validate_reloadable_string_flag_for_ut = + GFLAGS_NAMESPACE::RegisterFlagValidator( + &FLAGS_reloadable_string_flag_for_ut, PassValidateStringFlag); + namespace brpc { DECLARE_bool(enable_rpcz); DECLARE_bool(rpcz_hex_log_id); @@ -686,6 +697,62 @@ TEST_F(BuiltinServiceTest, flags) { TestFlags(true); } +TEST_F(BuiltinServiceTest, flags_escaping) { + // Save all flags and restore them on any exit of this test, since the + // /flags service below modifies `reloadable_string_flag_for_ut'. + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FlagsService service; + const std::string payload = "&\"'"; + const std::string escaped = brpc::WebEscape(payload); + + // Reflected: the ?setvalue= value is echoed into the html page. + { + ClosureChecker done; + brpc::Controller cntl; + brpc::FlagsRequest req; + brpc::FlagsResponse res; + SetUpController(&cntl, true); + cntl.http_request()._unresolved_path = "reloadable_string_flag_for_ut"; + cntl.http_request().uri().SetQuery(brpc::SETVALUE_STR, payload); + service.default_method(&cntl, &req, &res, &done); + EXPECT_FALSE(cntl.Failed()); + const std::string& body = cntl.response_attachment().to_string(); + EXPECT_EQ(std::string::npos, body.find(payload)) + << "unescaped payload in html: " << body; + CheckContent(cntl, escaped.c_str()); + } + // Stored: ?setvalue&withform renders the flag value stored above. + { + ClosureChecker done; + brpc::Controller cntl; + brpc::FlagsRequest req; + brpc::FlagsResponse res; + SetUpController(&cntl, true); + cntl.http_request()._unresolved_path = "reloadable_string_flag_for_ut"; + cntl.http_request().uri().SetQuery(brpc::SETVALUE_STR, ""); + cntl.http_request().uri().SetQuery("withform", ""); + service.default_method(&cntl, &req, &res, &done); + EXPECT_FALSE(cntl.Failed()); + const std::string& body = cntl.response_attachment().to_string(); + EXPECT_EQ(std::string::npos, body.find(payload)) + << "unescaped payload in html: " << body; + CheckContent(cntl, escaped.c_str()); + } + // Plain text output is not html-escaped. + { + ClosureChecker done; + brpc::Controller cntl; + brpc::FlagsRequest req; + brpc::FlagsResponse res; + SetUpController(&cntl, false); + cntl.http_request()._unresolved_path = "reloadable_string_flag_for_ut"; + cntl.http_request().uri().SetQuery(brpc::SETVALUE_STR, payload); + service.default_method(&cntl, &req, &res, &done); + EXPECT_FALSE(cntl.Failed()); + CheckContent(cntl, payload.c_str()); + } +} + TEST_F(BuiltinServiceTest, bad_method) { TestBadMethod(false); TestBadMethod(true); diff --git a/test/brpc_circuit_breaker_unittest.cpp b/test/brpc_circuit_breaker_unittest.cpp index 607fede9ea..975f3c654a 100644 --- a/test/brpc_circuit_breaker_unittest.cpp +++ b/test/brpc_circuit_breaker_unittest.cpp @@ -19,6 +19,7 @@ // Date: 2018/09/19 14:51:06 +#include #include #include #include @@ -150,6 +151,60 @@ TEST_F(CircuitBreakerTest, should_not_isolate) { } } +TEST_F(CircuitBreakerTest, canceled_requests_during_initialization) { + brpc::CircuitBreaker baseline; + for (int i = 0; i < 2 * kLongWindowSize; ++i) { + ASSERT_TRUE(_circuit_breaker.OnCallEnd(ECANCELED, kLatency)); + } + EXPECT_EQ(0, _circuit_breaker.isolated_times()); + + // Cancellations must not advance initialization or consume its error budget. + bool healthy = true; + for (int i = 0; i < kLongWindowSize && healthy; ++i) { + healthy = baseline.OnCallEnd(kErrorCodeForFailed, kErrorCost); + ASSERT_EQ(healthy, + _circuit_breaker.OnCallEnd(kErrorCodeForFailed, kErrorCost)); + } + EXPECT_FALSE(healthy); + EXPECT_EQ(1, _circuit_breaker.isolated_times()); +} + +TEST_F(CircuitBreakerTest, canceled_requests_after_initialization) { + brpc::CircuitBreaker baseline; + for (int i = 0; i < 2 * kLongWindowSize; ++i) { + ASSERT_TRUE(baseline.OnCallEnd(0, kLatency)); + ASSERT_TRUE(_circuit_breaker.OnCallEnd(0, kLatency)); + } + + // Interleaved cancellations must neither add error cost nor decay it + // like successful requests, regardless of their latency. + bool healthy = true; + for (int i = 0; i < 2 * kLongWindowSize && healthy; ++i) { + ASSERT_TRUE(_circuit_breaker.OnCallEnd(ECANCELED, 1)); + ASSERT_TRUE(_circuit_breaker.OnCallEnd(ECANCELED, 100 * kLatency)); + healthy = baseline.OnCallEnd(kErrorCodeForFailed, kErrorCost); + ASSERT_EQ(healthy, + _circuit_breaker.OnCallEnd(kErrorCodeForFailed, kErrorCost)); + } + EXPECT_FALSE(healthy); + EXPECT_EQ(1, _circuit_breaker.isolated_times()); +} + +TEST_F(CircuitBreakerTest, canceled_requests_in_half_open) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_circuit_breaker_half_open_window_size = 2; + _circuit_breaker.Reset(); + ASSERT_TRUE(_circuit_breaker.OnCallEnd(0, kLatency)); + for (int i = 0; i < 2 * kLongWindowSize; ++i) { + ASSERT_TRUE(_circuit_breaker.OnCallEnd(ECANCELED, kLatency)); + } + EXPECT_EQ(0, _circuit_breaker.isolated_times()); + + // One successful probe is still missing: a real error must reopen it. + EXPECT_FALSE(_circuit_breaker.OnCallEnd(kErrorCodeForFailed, kErrorCost)); + EXPECT_EQ(1, _circuit_breaker.isolated_times()); +} + TEST_F(CircuitBreakerTest, should_isolate) { std::vector thread_list; std::vector> fc_list; diff --git a/test/brpc_http_message_unittest.cpp b/test/brpc_http_message_unittest.cpp index 57e98ccab0..90c9dbddae 100644 --- a/test/brpc_http_message_unittest.cpp +++ b/test/brpc_http_message_unittest.cpp @@ -32,6 +32,7 @@ DECLARE_bool(allow_chunked_length); DECLARE_bool(allow_http_1_1_request_without_host); DECLARE_bool(http_allow_obs_fold); DECLARE_bool(http_strict_header_token); +DECLARE_uint32(http_max_header_count); int main(int argc, char* argv[]) { testing::InitGoogleTest(&argc, argv); @@ -643,6 +644,65 @@ TEST(HttpMessageTest, htab_is_ows_in_header_values) { } } +TEST(HttpMessageTest, too_many_headers) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_http_max_header_count = 8; + + // Host counts as well, so 8 distinct names in total are accepted. + std::string at_limit = "GET / HTTP/1.1\r\nHost: a.com\r\n"; + for (int i = 1; i < 8; ++i) { + at_limit.append("h" + std::to_string(i) + ": v\r\n"); + } + std::string over_limit = at_limit + "last: v\r\n\r\n"; + at_limit.append("\r\n"); + { + brpc::HttpMessage http_message; + ASSERT_EQ((ssize_t)at_limit.size(), + http_message.ParseFromArray(at_limit.data(), at_limit.size())) + << http_message._parser; + ASSERT_EQ(8u, http_message.header().HeaderCount()); + } + { + brpc::HttpMessage http_message; + ASSERT_EQ(-1, http_message.ParseFromArray(over_limit.data(), + over_limit.size())); + } + + // Repeated names fold into one entry, so they occupy one bucket and are not + // what the limit is aimed at. + std::string folded = "GET / HTTP/1.1\r\nHost: a.com\r\n"; + for (int i = 0; i < 100; ++i) { + folded.append("dup: v\r\n"); + } + folded.append("\r\n"); + { + brpc::HttpMessage http_message; + ASSERT_EQ((ssize_t)folded.size(), + http_message.ParseFromArray(folded.data(), folded.size())) + << http_message._parser; + ASSERT_EQ(2u, http_message.header().HeaderCount()); + } + // Set-Cookie is the one name that does not fold, so each occurrence is its + // own entry and does count. + std::string cookies = "GET / HTTP/1.1\r\nHost: a.com\r\n"; + for (int i = 0; i < 100; ++i) { + cookies.append("Set-Cookie: a=b\r\n"); + } + cookies.append("\r\n"); + { + brpc::HttpMessage http_message; + ASSERT_EQ(-1, http_message.ParseFromArray(cookies.data(), cookies.size())); + } + + brpc::FLAGS_http_max_header_count = 0; + { + brpc::HttpMessage http_message; + ASSERT_EQ((ssize_t)over_limit.size(), + http_message.ParseFromArray(over_limit.data(), over_limit.size())) + << http_message._parser; + } +} + TEST(HttpMessageTest, find_method_property_by_uri) { brpc::Server server; ASSERT_EQ(0, server.AddService(new test::EchoService(), diff --git a/test/brpc_http_rpc_protocol_unittest.cpp b/test/brpc_http_rpc_protocol_unittest.cpp index 10fb5cc3a2..136c0f8b41 100644 --- a/test/brpc_http_rpc_protocol_unittest.cpp +++ b/test/brpc_http_rpc_protocol_unittest.cpp @@ -63,6 +63,8 @@ DECLARE_bool(allow_chunked_length); DECLARE_int32(max_connection_pool_size); DECLARE_uint64(max_body_size); DECLARE_int64(socket_max_unwritten_bytes); +DECLARE_uint32(http_max_header_count); +DECLARE_uint32(http_max_query_count); extern bvar::CollectorSpeedLimit g_rpc_dump_sl; } @@ -727,10 +729,10 @@ TEST_F(HttpTest, complete_flow) { } TEST_F(HttpTest, chunked_uploading) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; // Send request via curl using chunked encoding const std::string req = "{\"message\":\"hello\"}"; @@ -889,11 +891,11 @@ class DownloadServiceImpl : public ::test::DownloadService { }; TEST_F(HttpTest, read_chunked_response_normally) { - const int port = 8923; brpc::Server server; DownloadServiceImpl svc; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; for (int i = 0; i < 3; ++i) { svc.set_done_place((DonePlace)i); @@ -913,11 +915,11 @@ TEST_F(HttpTest, read_chunked_response_normally) { } TEST_F(HttpTest, read_failed_chunked_response) { - const int port = 8923; brpc::Server server; DownloadServiceImpl svc; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1036,10 +1038,10 @@ TEST_F(HttpTest, read_long_body_progressively) { std::numeric_limits::max()); butil::intrusive_ptr reader; { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; { brpc::Channel channel; brpc::ChannelOptions options; @@ -1082,12 +1084,12 @@ TEST_F(HttpTest, read_long_body_progressively) { TEST_F(HttpTest, read_short_body_progressively) { butil::intrusive_ptr reader; - const int port = 8923; brpc::Server server; const int NREP = 10000; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, NREP); - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; { brpc::Channel channel; brpc::ChannelOptions options; @@ -1119,11 +1121,11 @@ TEST_F(HttpTest, read_short_body_progressively) { } TEST_F(HttpTest, progressive_read_timeout_keeps_active_reader_alive) { - const int port = 8923; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, 8, 100000); brpc::Server server; ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1148,11 +1150,11 @@ TEST_F(HttpTest, progressive_read_timeout_keeps_active_reader_alive) { } TEST_F(HttpTest, progressive_read_timeout_closes_idle_http1_reader_once) { - const int port = 8923; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, 2, 300000); brpc::Server server; ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; butil::intrusive_ptr reader(new TimeoutReadBody); { @@ -1182,11 +1184,11 @@ TEST_F(HttpTest, progressive_read_timeout_closes_idle_http1_reader_once) { } TEST_F(HttpTest, progressive_read_timeout_before_first_body_part) { - const int port = 8923; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, 1, 0, 300000); brpc::Server server; ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; butil::intrusive_ptr reader(new TimeoutReadBody); { @@ -1215,11 +1217,11 @@ TEST_F(HttpTest, progressive_read_timeout_before_first_body_part) { } TEST_F(HttpTest, progressive_read_timeout_ignores_slow_user_callback) { - const int port = 8923; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, 3, 50000); brpc::Server server; ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1245,11 +1247,11 @@ TEST_F(HttpTest, progressive_read_timeout_ignores_slow_user_callback) { } TEST_F(HttpTest, progressive_read_timeout_preserves_reader_error) { - const int port = 8923; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, 10); brpc::Server server; ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1274,10 +1276,10 @@ TEST_F(HttpTest, progressive_read_timeout_preserves_reader_error) { } TEST_F(HttpTest, progressive_read_timeout_rejects_http2) { - const int port = 8923; brpc::Server server; ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - ASSERT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1305,10 +1307,10 @@ TEST_F(HttpTest, read_progressively_after_cntl_destroys) { std::numeric_limits::max()); butil::intrusive_ptr reader; { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; { brpc::Channel channel; brpc::ChannelOptions options; @@ -1351,10 +1353,10 @@ TEST_F(HttpTest, read_progressively_after_long_delay) { DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, std::numeric_limits::max()); { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; { brpc::Channel channel; brpc::ChannelOptions options; @@ -1399,10 +1401,10 @@ TEST_F(HttpTest, read_progressively_after_long_delay) { TEST_F(HttpTest, skip_progressive_reading) { DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, std::numeric_limits::max()); - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; options.protocol = brpc::PROTOCOL_HTTP; @@ -1438,12 +1440,12 @@ class AlwaysFailRead : public brpc::ProgressiveReader { }; TEST_F(HttpTest, failed_on_read_one_part) { - const int port = 8923; brpc::Server server; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, std::numeric_limits::max()); - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; options.protocol = brpc::PROTOCOL_HTTP; @@ -1464,12 +1466,12 @@ TEST_F(HttpTest, failed_on_read_one_part) { TEST_F(HttpTest, broken_socket_stops_progressive_reading) { butil::intrusive_ptr reader; - const int port = 8923; brpc::Server server; DownloadServiceImpl svc(DONE_BEFORE_CREATE_PA, std::numeric_limits::max()); - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1587,14 +1589,14 @@ class UploadServiceImpl : public ::test::UploadService { }; TEST_F(HttpTest, server_end_read_short_body_progressively) { - const int port = 8923; brpc::ServiceOptions opt; opt.enable_progressive_read = true; opt.ownership = brpc::SERVER_DOESNT_OWN_SERVICE; UploadServiceImpl upsvc; brpc::Server server; - EXPECT_EQ(0, server.AddService(&upsvc, opt)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&upsvc, opt)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1626,14 +1628,14 @@ TEST_F(HttpTest, server_end_read_short_body_progressively) { // Fixme!!! Server progressive reader has a heap-use-after-free bug detected by ASan. // For details, see https://github.com/apache/brpc/issues/2145#issuecomment-2329413363 TEST_F(HttpTest, server_end_read_failed) { - const int port = 8923; brpc::ServiceOptions opt; opt.enable_progressive_read = true; opt.ownership = brpc::SERVER_DOESNT_OWN_SERVICE; UploadServiceImpl upsvc; brpc::Server server; - EXPECT_EQ(0, server.AddService(&upsvc, opt)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&upsvc, opt)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1664,10 +1666,10 @@ TEST_F(HttpTest, server_end_read_failed) { #endif // BUTIL_USE_ASAN TEST_F(HttpTest, http2_sanity) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -1969,6 +1971,231 @@ TEST_F(HttpTest, h2_header_list_budget_resets_per_block) { delete sctx; } +// Literal header field with a new name, with both lengths in a single 7-bit +// prefix octet. `first_octet` selects the representation: 0x00 is "without +// indexing" (RFC 7541 6.2.2), 0x40 is "with incremental indexing" (6.2.1) +// which also adds the field to the dynamic table. +// 0x80 of a length octet is the Huffman flag and a length of 128 or more needs +// the multi-octet form, so refuse what does not fit instead of emitting a +// corrupt header block. +void AppendLiteralHeader(butil::IOBuf* out, const std::string& name, + const std::string& value, uint8_t first_octet = 0x00) { + ASSERT_LT(name.size(), 0x80u); + ASSERT_LT(value.size(), 0x80u); + uint8_t prefix[] = { first_octet, (uint8_t)name.size() }; + out->append(prefix, sizeof(prefix)); + out->append(name); + uint8_t value_len = (uint8_t)value.size(); + out->append(&value_len, 1); + out->append(value); +} + +// Feed `payload` to `sctx` as one complete HEADERS block, the way +// H2Context::Consume() would. A non-zero stream_id in the result means the +// frame handler asked for a RST_STREAM, a zero one means a GOAWAY that closes +// the whole connection. +brpc::policy::H2ParseResult ConsumeHeadersBlock( + brpc::policy::H2StreamContext* sctx, const butil::IOBuf& payload, + int stream_id) { + brpc::policy::H2FrameHead head; + head.payload_size = payload.size(); + head.type = brpc::policy::H2_FRAME_HEADERS; + head.flags = 0x4; // H2_FLAGS_END_HEADERS + head.stream_id = stream_id; + butil::IOBufBytesIterator it(payload); + return sctx->OnHeaders(it, head, payload.size(), 0); +} + +TEST_F(HttpTest, h2_too_many_headers) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_http_max_header_count = 8; + + brpc::policy::H2Context* ctx = + new brpc::policy::H2Context(_socket.get(), nullptr); + CHECK_EQ(ctx->Init(), 0); + _socket->initialize_parsing_context(&ctx); + + { + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, 1); + butil::IOBuf payload; + for (int i = 0; i < 8; ++i) { + AppendLiteralHeader(&payload, "h" + std::to_string(i), "v"); + } + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, 1); + ASSERT_TRUE(res.is_ok()) << brpc::H2ErrorToString(res.error()); + ASSERT_EQ(8u, sctx->header().HeaderCount()); + } + { + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, 3); + butil::IOBuf payload; + for (int i = 0; i < 9; ++i) { + AppendLiteralHeader(&payload, "h" + std::to_string(i), "v"); + } + // Refusing the request must not cost the connection its other + // streams, so the frame handler asks for a RST_STREAM (non-zero + // stream_id) rather than a GOAWAY. + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, 3); + ASSERT_FALSE(res.is_ok()); + ASSERT_EQ(brpc::H2_ENHANCE_YOUR_CALM, res.error()); + ASSERT_EQ(3, res.stream_id()); + } +} + +TEST_F(HttpTest, h2_too_many_queries_in_path) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_http_max_query_count = 4; + + brpc::policy::H2Context* ctx = + new brpc::policy::H2Context(_socket.get(), nullptr); + CHECK_EQ(ctx->Init(), 0); + _socket->initialize_parsing_context(&ctx); + + { + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, 1); + butil::IOBuf payload; + AppendLiteralHeader(&payload, ":path", "/s?a=1&b=2&c=3&d=4"); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, 1); + ASSERT_TRUE(res.is_ok()) << brpc::H2ErrorToString(res.error()); + } + { + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, 3); + butil::IOBuf payload; + AppendLiteralHeader(&payload, ":path", "/s?a=1&b=2&c=3&d=4&e=5"); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, 3); + ASSERT_FALSE(res.is_ok()); + ASSERT_EQ(brpc::H2_ENHANCE_YOUR_CALM, res.error()); + ASSERT_EQ(3, res.stream_id()); + } +} + +// A refused header block still has to be fed to the HPACK decoder in full. +// The dynamic table belongs to the connection, so dropping the tail of a block +// would leave it out of step with the encoding table of the peer and turn +// every later block into garbage, which is why RFC 9113 section 10.5.1 says +// the field block MUST be processed unless the connection is closed. +TEST_F(HttpTest, h2_refused_header_block_keeps_hpack_in_sync) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_http_max_header_count = 2; + + brpc::policy::H2Context* ctx = + new brpc::policy::H2Context(_socket.get(), nullptr); + CHECK_EQ(ctx->Init(), 0); + _socket->initialize_parsing_context(&ctx); + + // Four headers with incremental indexing, two of them past the limit. + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, 1); + butil::IOBuf payload; + AppendLiteralHeader(&payload, "a", "1", 0x40); + AppendLiteralHeader(&payload, "b", "2", 0x40); + AppendLiteralHeader(&payload, "c", "3", 0x40); + AppendLiteralHeader(&payload, "d", "4", 0x40); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, 1); + ASSERT_FALSE(res.is_ok()); + ASSERT_EQ(brpc::H2_ENHANCE_YOUR_CALM, res.error()); + ASSERT_EQ(1, res.stream_id()); + // Everything after the offending field is decoded but thrown away. + ASSERT_EQ(3u, sctx->header().HeaderCount()); + + // The static table ends at index 61, so 62 names the newest dynamic entry. + // That is "d" only because decoding ran to the end of the block; had it + // stopped at the limit, 62 would still be "c". + std::unique_ptr sctx2( + new brpc::policy::H2StreamContext(false)); + sctx2->Init(ctx, 3); + butil::IOBuf indexed; + const uint8_t indexed_field[] = { 0x80 | 62 }; // Indexed Header Field + indexed.append(indexed_field, sizeof(indexed_field)); + brpc::policy::H2ParseResult res2 = + ConsumeHeadersBlock(sctx2.get(), indexed, 3); + ASSERT_TRUE(res2.is_ok()) << brpc::H2ErrorToString(res2.error()); + const std::string* value = sctx2->header().GetHeader("d"); + ASSERT_TRUE(value != nullptr); + ASSERT_EQ("4", *value); +} + +// RFC 9113 section 8.1.1: "Malformed requests or responses that are detected +// MUST be treated as a stream error (Section 5.4.2) of type PROTOCOL_ERROR." +// A bad pseudo-header says nothing about the health of the connection, so it +// must not cost the other streams theirs. :path has its own case table in +// HttpTest.http2_reject_path_not_starting_with_slash. +TEST_F(HttpTest, h2_malformed_pseudo_header_resets_stream_only) { + struct MalformedField { + const char* name; + const char* value; + }; + MalformedField malformed[] = { + { ":method", "NOSUCH" }, + { ":status", "20x" }, + { ":nosuchheader", "1" }, // 8.3: undefined pseudo-header + }; + + brpc::policy::H2Context* ctx = + new brpc::policy::H2Context(_socket.get(), nullptr); + CHECK_EQ(ctx->Init(), 0); + _socket->initialize_parsing_context(&ctx); + + int stream_id = 1; + for (const auto& bad : malformed) { + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, stream_id); + butil::IOBuf payload; + AppendLiteralHeader(&payload, bad.name, bad.value); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, stream_id); + std::string desc = std::string(bad.name) + '=' + bad.value; + ASSERT_FALSE(res.is_ok()) << desc; + ASSERT_EQ(brpc::H2_PROTOCOL_ERROR, res.error()) + << desc << ": " << brpc::H2ErrorToString(res.error()); + ASSERT_EQ(stream_id, res.stream_id()) << desc; + stream_id += 2; + } + + // A malformed block is drained like any other refusal, so a field that + // follows the bad pseudo-header still reaches the dynamic table. RFC 9113 + // section 4.3 leaves no choice here: only a decoding error may take down + // the connection, so everything else has to be decoded to the end. + std::unique_ptr sctx( + new brpc::policy::H2StreamContext(false)); + sctx->Init(ctx, stream_id); + butil::IOBuf payload; + AppendLiteralHeader(&payload, ":path", "foo"); + AppendLiteralHeader(&payload, "after-the-bad-one", "1", 0x40); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(sctx.get(), payload, stream_id); + ASSERT_EQ(brpc::H2_PROTOCOL_ERROR, res.error()); + ASSERT_EQ(stream_id, res.stream_id()); + + stream_id += 2; + std::unique_ptr sctx2( + new brpc::policy::H2StreamContext(false)); + sctx2->Init(ctx, stream_id); + butil::IOBuf indexed; + const uint8_t indexed_field[] = { 0x80 | 62 }; // newest dynamic entry + indexed.append(indexed_field, sizeof(indexed_field)); + brpc::policy::H2ParseResult res2 = + ConsumeHeadersBlock(sctx2.get(), indexed, stream_id); + ASSERT_TRUE(res2.is_ok()) << brpc::H2ErrorToString(res2.error()); + const std::string* value = sctx2->header().GetHeader("after-the-bad-one"); + ASSERT_TRUE(value != nullptr); + ASSERT_EQ("1", *value); +} + TEST_F(HttpTest, h2_oversized_single_headers_block_rejected) { // A single HEADERS frame whose decoded header list exceeds // max_header_list_size must be rejected at the block boundary (before @@ -2127,29 +2354,29 @@ TEST_F(HttpTest, http2_invalid_settings) { brpc::Server server; brpc::ServerOptions options; options.h2_settings.stream_window_size = brpc::H2Settings::MAX_WINDOW_SIZE + 1; - ASSERT_EQ(-1, server.Start("127.0.0.1:8924", &options)); + ASSERT_EQ(-1, server.Start("127.0.0.1:0", &options)); } { brpc::Server server; brpc::ServerOptions options; options.h2_settings.max_frame_size = brpc::H2Settings::DEFAULT_MAX_FRAME_SIZE - 1; - ASSERT_EQ(-1, server.Start("127.0.0.1:8924", &options)); + ASSERT_EQ(-1, server.Start("127.0.0.1:0", &options)); } { brpc::Server server; brpc::ServerOptions options; options.h2_settings.max_frame_size = brpc::H2Settings::MAX_OF_MAX_FRAME_SIZE + 1; - ASSERT_EQ(-1, server.Start("127.0.0.1:8924", &options)); + ASSERT_EQ(-1, server.Start("127.0.0.1:0", &options)); } } TEST_F(HttpTest, http2_not_closing_socket_when_rpc_timeout) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; options.protocol = "h2"; @@ -2386,7 +2613,8 @@ TEST_F(HttpTest, http2_handle_goaway_streams) { } // RFC 9113 8.3.1: :path MUST NOT be empty and MUST begin with '/', the only -// exception being the asterisk-form that OPTIONS uses. +// exception being the asterisk-form that OPTIONS uses. Violating that makes +// the request malformed, which 8.1.1 turns into a stream error. TEST_F(HttpTest, http2_reject_path_not_starting_with_slash) { brpc::policy::H2Context* h2_ctx = new brpc::policy::H2Context(_socket.get(), &_server); @@ -2423,23 +2651,32 @@ TEST_F(HttpTest, http2_reject_path_not_starting_with_slash) { h2_ctx->hpacker().Encode(&appender, header, options); butil::IOBuf buf; appender.move_to(buf); - butil::IOBufBytesIterator it(buf); brpc::policy::H2StreamContext* h2_msg = new brpc::policy::H2StreamContext(false); h2_msg->Init(h2_ctx, stream_id); + brpc::policy::H2ParseResult res = + ConsumeHeadersBlock(h2_msg, buf, stream_id); + if (c.accepted) { + ASSERT_TRUE(res.is_ok()) << "path=`" << c.path << "': " + << brpc::H2ErrorToString(res.error()); + } else { + // A bad :path makes the request malformed, not the connection + // unusable, so only this stream is reset. + ASSERT_EQ(brpc::H2_PROTOCOL_ERROR, res.error()) + << "path=`" << c.path << '\''; + ASSERT_EQ(stream_id, res.stream_id()) << "path=`" << c.path << '\''; + } stream_id += 2; - ASSERT_EQ(c.accepted ? 0 : -1, h2_msg->ConsumeHeaders(it)) - << "path=`" << c.path << '\''; h2_msg->Destroy(); } } TEST_F(HttpTest, spring_protobuf_content_type) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -2483,10 +2720,10 @@ TEST_F(HttpTest, dump_http_request) { brpc::g_rpc_dump_sl.sampling_range = bvar::COLLECTOR_SAMPLING_BASE; // init channel - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -2555,10 +2792,10 @@ TEST_F(HttpTest, dump_http_request) { } TEST_F(HttpTest, proto_text_content_type) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -2593,10 +2830,10 @@ TEST_F(HttpTest, proto_text_content_type) { } TEST_F(HttpTest, proto_json_content_type) { - const int port = 8923; brpc::Server server; - EXPECT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&_svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -2670,11 +2907,11 @@ class HttpServiceImpl : public ::test::HttpService { }; TEST_F(HttpTest, http_head) { - const int port = 8923; brpc::Server server; HttpServiceImpl svc; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(port, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + int port = server.listen_address().port; brpc::Channel channel; brpc::ChannelOptions options; @@ -2798,8 +3035,8 @@ void ReadOneResponse(brpc::SocketUniquePtr& sock, TEST_F(HttpTest, http_expect) { brpc::Server server; HttpServiceImpl svc; - EXPECT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); - EXPECT_EQ(0, server.Start(0, nullptr)); + ASSERT_EQ(0, server.AddService(&svc, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); const butil::EndPoint ep = server.listen_address(); brpc::SocketOptions options; diff --git a/test/brpc_mysql_reply_parse_unittest.cpp b/test/brpc_mysql_reply_parse_unittest.cpp index 33ab95a519..f10773dbfd 100644 --- a/test/brpc_mysql_reply_parse_unittest.cpp +++ b/test/brpc_mysql_reply_parse_unittest.cpp @@ -101,6 +101,47 @@ TEST(MysqlReplyParseTest, RejectOversizedTextFieldLength) { ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, rc); } +TEST(MysqlReplyParseTest, RejectZeroPayloadPacket) { + butil::IOBuf buf; + buf.append(std::string("\x00\x00\x00\x01", 4)); + + brpc::MysqlReply reply; + butil::Arena arena; + bool more_results = false; + brpc::ParseError rc = reply.ConsumePartialIOBuf( + buf, &arena, false, brpc::MYSQL_NORMAL_STATEMENT, &more_results); + ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, rc); +} + +TEST(MysqlReplyParseTest, RejectZeroPayloadPacketWithTrailingBytes) { + std::string wire("\x00\x00\x00\x01", 4); + AppendPacket(&wire, 2, std::string("\x00\x00\x00\x00\x00\x00\x00", 7)); + butil::IOBuf buf; + buf.append(wire); + + brpc::MysqlReply reply; + butil::Arena arena; + bool more_results = false; + brpc::ParseError rc = reply.ConsumePartialIOBuf( + buf, &arena, false, brpc::MYSQL_NORMAL_STATEMENT, &more_results); + ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, rc); +} + +TEST(MysqlReplyParseTest, RejectZeroPayloadPacketAfterFastAuthMarker) { + std::string wire; + AppendPacket(&wire, 2, std::string("\x01\x03", 2)); + wire.append(std::string("\x00\x00\x00\x03", 4)); + butil::IOBuf buf; + buf.append(wire); + + brpc::MysqlReply reply; + butil::Arena arena; + bool more_results = false; + brpc::ParseError rc = reply.ConsumePartialIOBuf( + buf, &arena, true, brpc::MYSQL_NORMAL_STATEMENT, &more_results); + ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, rc); +} + // A well-formed field whose length matches the bytes present still parses, so // the guard does not reject legitimate result sets. TEST(MysqlReplyParseTest, AcceptWellFormedTextField) { diff --git a/test/brpc_prometheus_metrics_unittest.cpp b/test/brpc_prometheus_metrics_unittest.cpp index 166d6f80a6..2edd484ca5 100644 --- a/test/brpc_prometheus_metrics_unittest.cpp +++ b/test/brpc_prometheus_metrics_unittest.cpp @@ -78,6 +78,10 @@ TEST(PrometheusMetrics, sanity) { ASSERT_TRUE(my_lat2); *my_lat2 << 3 << 4; + // Only a bvar prefixed with the server prefix is folded into a summary. + bvar::LatencyRecorder my_lat3("rpc_server_lat_test"); + my_lat3 << 5 << 6; + brpc::Channel channel; brpc::ChannelOptions channel_opts; channel_opts.protocol = "http"; @@ -88,6 +92,36 @@ TEST(PrometheusMetrics, sanity) { ASSERT_FALSE(cntl.Failed()); std::string res = cntl.response_attachment().to_string(); LOG(INFO) << "output:\n" << res; + + // The average latency is a separate metric rather than a quantile series, + // because the quantile label must be parsable as a float. + ASSERT_EQ(std::string::npos, res.find("quantile=\"avg\"")); + ASSERT_NE(std::string::npos, res.find("# TYPE mlat_avg_latency gauge\n")); + ASSERT_NE(std::string::npos, res.find("mlat_avg_latency{label1=\"val1\"," + "label2=\"val2\"}")); + // The single dimension LatencyRecorder uses the same suffix. + ASSERT_NE(std::string::npos, res.find("_service_echo_avg_latency ")); + // Quantile is a fraction rather than an integer. + ASSERT_NE(std::string::npos, res.find("quantile=\"0.99\"")); + ASSERT_NE(std::string::npos, res.find("quantile=\"0.999\"")); + ASSERT_NE(std::string::npos, res.find("quantile=\"0.9999\"")); + ASSERT_EQ(std::string::npos, res.find("quantile=\"99\"")); + ASSERT_EQ(std::string::npos, res.find("quantile=\"999\"")); + ASSERT_EQ(std::string::npos, res.find("quantile=\"9999\"")); + ASSERT_NE(std::string::npos, res.find("mlat_latency{label1=\"val1\",label2=\"val2\"," + "quantile=\"0.99\"}")); + // The average must not be dumped as a series of `_latency` as well, otherwise + // an aggregation over `_latency` would still pick it up. + ASSERT_EQ(std::string::npos, res.find("mlat_latency{label1=\"val1\"," + "label2=\"val2\"} ")); + ASSERT_NE(std::string::npos, res.find("rpc_server_lat_test_count 2\n")); + // `_avg_latency` is dumped before the summary it belongs to. + size_t average_pos = res.find("# TYPE rpc_server_lat_test_avg_latency gauge\n"); + size_t summary_pos = res.find("# TYPE rpc_server_lat_test summary\n"); + ASSERT_NE(std::string::npos, average_pos); + ASSERT_NE(std::string::npos, summary_pos); + ASSERT_LT(average_pos, summary_pos); + size_t start_pos = 0; size_t end_pos = 0; size_t label_start = 0; diff --git a/test/brpc_rdma_unittest.cpp b/test/brpc_rdma_unittest.cpp index 69b1cd8222..97265ae0af 100644 --- a/test/brpc_rdma_unittest.cpp +++ b/test/brpc_rdma_unittest.cpp @@ -66,19 +66,19 @@ DECLARE_int32(rdma_memory_pool_max_regions); DECLARE_int32(rdma_client_handshake_version); DECLARE_bool(rdma_ece); -extern ibv_cq *(*IbvCreateCq)(ibv_context *, int, void *, ibv_comp_channel *, +extern ibv_cq* (*IbvCreateCq)(ibv_context*, int, void*, ibv_comp_channel*, int); -extern int (*IbvDestroyCq)(ibv_cq *); -extern ibv_qp *(*IbvCreateQp)(ibv_pd *, ibv_qp_init_attr *); -extern int (*IbvModifyQp)(ibv_qp *, ibv_qp_attr *, ibv_qp_attr_mask); -extern int (*IbvQueryQp)(ibv_qp *, ibv_qp_attr *, ibv_qp_attr_mask, - ibv_qp_init_attr *); -extern int (*IbvDestroyQp)(ibv_qp *); +extern int (*IbvDestroyCq)(ibv_cq*); +extern ibv_qp* (*IbvCreateQp)(ibv_pd*, ibv_qp_init_attr*); +extern int (*IbvModifyQp)(ibv_qp*, ibv_qp_attr*, ibv_qp_attr_mask); +extern int (*IbvQueryQp)(ibv_qp*, ibv_qp_attr*, ibv_qp_attr_mask, + ibv_qp_init_attr*); +extern int (*IbvDestroyQp)(ibv_qp*); extern butil::atomic g_rdma_available; extern bool g_skip_rdma_init; extern bool g_fail_resource_alloc_for_test; -} // namespace rdma -} // namespace brpc +} // namespace rdma +} // namespace brpc static std::string g_ip = "127.0.0.1"; static butil::EndPoint g_ep; @@ -115,7 +115,7 @@ static bool WaitUntil(const std::function& pred, // whole buffer reaching the peer has to loop. static bool WriteAll(int fd, const void* buf, size_t len) { const uint8_t* p = (const uint8_t*)buf; - for (size_t done = 0; done < len; ) { + for (size_t done = 0; done < len;) { ssize_t n = write(fd, p + done, len - done); if (n < 0) { if (errno == EINTR) { @@ -132,7 +132,7 @@ static bool WriteAll(int fd, const void* buf, size_t len) { // and on EOF before the whole buffer arrived. static bool ReadAll(int fd, void* buf, size_t len) { uint8_t* p = (uint8_t*)buf; - for (size_t done = 0; done < len; ) { + for (size_t done = 0; done < len;) { ssize_t n = read(fd, p + done, len - done); if (n < 0) { if (errno == EINTR) { @@ -168,120 +168,120 @@ static void ConnectToServer(butil::fd_guard* sockfd) { } class MyEchoService : public ::test::EchoService { - void Echo(google::protobuf::RpcController *cntl_base, - const ::test::EchoRequest *req, ::test::EchoResponse *res, - google::protobuf::Closure *done) { - Controller *cntl = static_cast(cntl_base); - ClosureGuard done_guard(done); - g_echo_served.fetch_add(1, butil::memory_order_relaxed); - if (req->server_fail()) { - cntl->SetFailed(req->server_fail(), "Server fail1"); - cntl->SetFailed(req->server_fail(), "Server fail2"); - return; - } - if (req->close_fd()) { - usleep(1); - LOG(INFO) << "close fd..."; - cntl->CloseConnection("Close connection according to request"); - return; - } - if (req->sleep_us() > 0) { - LOG(INFO) << "sleep " << req->sleep_us() << "us..."; - bthread_usleep(req->sleep_us()); - } - res->set_message("MyEchoService"); - if (req->code() != 0) { - res->add_code_list(req->code()); - } - cntl->response_attachment().append(cntl->request_attachment()); - } + void Echo(google::protobuf::RpcController* cntl_base, + const ::test::EchoRequest* req, ::test::EchoResponse* res, + google::protobuf::Closure* done) { + Controller* cntl = static_cast(cntl_base); + ClosureGuard done_guard(done); + g_echo_served.fetch_add(1, butil::memory_order_relaxed); + if (req->server_fail()) { + cntl->SetFailed(req->server_fail(), "Server fail1"); + cntl->SetFailed(req->server_fail(), "Server fail2"); + return; + } + if (req->close_fd()) { + usleep(1); + LOG(INFO) << "close fd..."; + cntl->CloseConnection("Close connection according to request"); + return; + } + if (req->sleep_us() > 0) { + LOG(INFO) << "sleep " << req->sleep_us() << "us..."; + bthread_usleep(req->sleep_us()); + } + res->set_message("MyEchoService"); + if (req->code() != 0) { + res->add_code_list(req->code()); + } + cntl->response_attachment().append(cntl->request_attachment()); + } }; class RdmaTest : public ::testing::Test { protected: - RdmaTest() { - butil::ip_t ip; - EXPECT_EQ(0, butil::str2ip(g_ip.c_str(), &ip)); - butil::EndPoint ep(ip, PORT); - g_ep = ep; - EXPECT_EQ(0, _server_list.save(butil::endpoint2str(g_ep).c_str())); - _naming_url = std::string("File://") + _server_list.fname(); - _server.AddService(&_svc, SERVER_DOESNT_OWN_SERVICE); - } - ~RdmaTest() {} + RdmaTest() { + butil::ip_t ip; + EXPECT_EQ(0, butil::str2ip(g_ip.c_str(), &ip)); + butil::EndPoint ep(ip, PORT); + g_ep = ep; + EXPECT_EQ(0, _server_list.save(butil::endpoint2str(g_ep).c_str())); + _naming_url = std::string("File://") + _server_list.fname(); + _server.AddService(&_svc, SERVER_DOESNT_OWN_SERVICE); + } + ~RdmaTest() {} - virtual void SetUp() {} + virtual void SetUp() {} - virtual void TearDown() { rdma::DumpMemoryPoolInfo(std::cout); } + virtual void TearDown() { rdma::DumpMemoryPoolInfo(std::cout); } protected: - void StartServer(bool use_rdma = true) { - ServerOptions options; - options.enabled_protocols = "baidu_std"; - options.socket_mode = use_rdma ? SOCKET_MODE_RDMA : SOCKET_MODE_TCP; - options.idle_timeout_sec = 5; - options.max_concurrency = 0; - options.internal_port = -1; - EXPECT_EQ(0, _server.Start(PORT, &options)); - } - - void StopServer() { - _server.Stop(0); - _server.Join(); - } - - Socket *GetSocketFromServer(size_t index) { - std::vector sids; - _server._am->ListConnections(&sids); - if (index >= sids.size()) { - return nullptr; + void StartServer(bool use_rdma = true) { + ServerOptions options; + options.enabled_protocols = "baidu_std"; + options.socket_mode = use_rdma ? SOCKET_MODE_RDMA : SOCKET_MODE_TCP; + options.idle_timeout_sec = 5; + options.max_concurrency = 0; + options.internal_port = -1; + EXPECT_EQ(0, _server.Start(PORT, &options)); } - SocketUniquePtr s; - if (Socket::Address(sids[index], &s) == 0) { - return s.get(); + + void StopServer() { + _server.Stop(0); + _server.Join(); + } + + Socket* GetSocketFromServer(size_t index) { + std::vector sids; + _server._am->ListConnections(&sids); + if (index >= sids.size()) { + return nullptr; + } + SocketUniquePtr s; + if (Socket::Address(sids[index], &s) == 0) { + return s.get(); + } + return nullptr; } - return nullptr; - } - // Server-side connection creation and teardown are asynchronous. - Socket *WaitForServerSocket() { - Socket *s = nullptr; - WaitUntil([this, &s] { return (s = GetSocketFromServer(0)) != nullptr; }); - return s; - } + // Server-side connection creation and teardown are asynchronous. + Socket* WaitForServerSocket() { + Socket* s = nullptr; + WaitUntil([this, &s] { return (s = GetSocketFromServer(0)) != nullptr; }); + return s; + } - bool WaitForServerSocketGone() { - return WaitUntil([this] { return GetSocketFromServer(0) == nullptr; }); - } + bool WaitForServerSocketGone() { + return WaitUntil([this] { return GetSocketFromServer(0) == nullptr; }); + } - butil::TempFile _server_list; - std::string _naming_url; + butil::TempFile _server_list; + std::string _naming_url; - Server _server; - MyEchoService _svc; + Server _server; + MyEchoService _svc; }; // Shorthand for the RDMA transport behind a Socket. -static RdmaTransport *RdmaTransportOf(Socket *s) { - return RdmaTransport::Get(s); +static RdmaTransport* RdmaTransportOf(Socket* s) { + return RdmaTransport::Get(s); } -static int WaitForHandshakePhase(Socket *s, handshake::Phase expected) { - int phase = AdapterTransport::Get(s)->handshake_phase(); - WaitUntil([s, expected, &phase] { - phase = AdapterTransport::Get(s)->handshake_phase(); - return phase == expected; - }); - return phase; +static int WaitForHandshakePhase(Socket* s, handshake::Phase expected) { + int phase = AdapterTransport::Get(s)->handshake_phase(); + WaitUntil([s, expected, &phase] { + phase = AdapterTransport::Get(s)->handshake_phase(); + return phase == expected; + }); + return phase; } #define ASSERT_HANDSHAKE_PHASE(expected, socket) \ - ASSERT_EQ(expected, WaitForHandshakePhase(socket, expected)) + ASSERT_EQ(expected, WaitForHandshakePhase(socket, expected)) -static bool WaitForFdReadBuf(Socket *s, size_t size) { - return WaitUntil([s, size] { - return s->fd_input_processor().read_buf().size() == size; - }); +static bool WaitForFdReadBuf(Socket* s, size_t size) { + return WaitUntil([s, size] { + return s->fd_input_processor().read_buf().size() == size; + }); } static void MakeV2ClientHello(uint8_t (&data)[rdma::HELLO_V2_MSG_LEN_MIN]) { @@ -306,18 +306,18 @@ static void MakeV2ClientHello(uint8_t (&data)[rdma::HELLO_V2_MSG_LEN_MIN]) { // of this file and these RPC tests will gain coverage for free. class RdmaRpcTest : public RdmaTest, public ::testing::WithParamInterface { protected: - void SetUp() override { - RdmaTest::SetUp(); - _saved_handshake_version = rdma::FLAGS_rdma_client_handshake_version; - rdma::FLAGS_rdma_client_handshake_version = GetParam(); - } - void TearDown() override { - rdma::FLAGS_rdma_client_handshake_version = _saved_handshake_version; - RdmaTest::TearDown(); - } + void SetUp() override { + RdmaTest::SetUp(); + _saved_handshake_version = rdma::FLAGS_rdma_client_handshake_version; + rdma::FLAGS_rdma_client_handshake_version = GetParam(); + } + void TearDown() override { + rdma::FLAGS_rdma_client_handshake_version = _saved_handshake_version; + RdmaTest::TearDown(); + } private: - int _saved_handshake_version = 2; + int _saved_handshake_version = 2; }; TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { @@ -356,566 +356,566 @@ TEST_F(RdmaTest, stale_cq_callback_does_not_poll_new_generation) { } TEST_F(RdmaTest, client_close_before_hello_send) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); + StartServer(); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - Socket *s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); - StopServer(); + StopServer(); } TEST_F(RdmaTest, client_hello_msg_invalid_magic_str) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - Socket *s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - memcpy(data, "PRPC", 4); // send as normal baidu_std protocol - ASSERT_EQ(4, write(sockfd, data, 4)); - usleep(100000); // wait for server to handle the msg - // A non-RDMA magic makes the transport-handshake parser return TRY_OTHERS - // and hand the bytes to other protocols; it does not touch the endpoint - // state, so it stays UNINIT (the old blocking handshake used to set - // FALLBACK_TCP here). - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + memcpy(data, "PRPC", 4); // send as normal baidu_std protocol + ASSERT_EQ(4, write(sockfd, data, 4)); + usleep(100000); // wait for server to handle the msg + // A non-RDMA magic makes the transport-handshake parser return TRY_OTHERS + // and hand the bytes to other protocols; it does not touch the endpoint + // state, so it stays UNINIT (the old blocking handshake used to set + // FALLBACK_TCP here). + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + StopServer(); } TEST_F(RdmaTest, client_close_during_hello_send) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - uint8_t data[8]; - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RD", 2); - ASSERT_EQ(2, write(sockfd1, data, 2)); // break in magic str - usleep(100000); // wait for server to handle the msg - // Fewer than 4 magic bytes: the transport-handshake parser can't tell yet, - // returns NOT_ENOUGH_DATA and leaves the endpoint UNINIT (the old blocking - // the common handshake state remains uninitialized before reading magic). - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // break after magic str - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - close(sockfd2); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd3 >= 0); - ASSERT_EQ(0, connect(sockfd3, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - // Send the 4B magic plus a valid msg_len (=40) but no body, so the server - // recognizes an RDMA v2 hello and waits for the remaining bytes. (A zero - // msg_len would now be rejected up-front as a protocol error.) - memcpy(data, "RDMA", 4); - uint16_t v2_len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); - memcpy(data + 4, &v2_len, sizeof(v2_len)); - ASSERT_EQ(6, write(sockfd3, data, 6)); // magic + msg_len, body missing - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - close(sockfd3); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + uint8_t data[8]; + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RD", 2); + ASSERT_EQ(2, write(sockfd1, data, 2)); // break in magic str + usleep(100000); // wait for server to handle the msg + // Fewer than 4 magic bytes: the transport-handshake parser can't tell yet, + // returns NOT_ENOUGH_DATA and leaves the endpoint UNINIT (the old blocking + // the common handshake state remains uninitialized before reading magic). + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RDMA", 4); + ASSERT_EQ(4, write(sockfd2, data, 4)); // break after magic str + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + close(sockfd2); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd3 >= 0); + ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + // Send the 4B magic plus a valid msg_len (=40) but no body, so the server + // recognizes an RDMA v2 hello and waits for the remaining bytes. (A zero + // msg_len would now be rejected up-front as a protocol error.) + memcpy(data, "RDMA", 4); + uint16_t v2_len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); + memcpy(data + 4, &v2_len, sizeof(v2_len)); + ASSERT_EQ(6, write(sockfd3, data, 6)); // magic + msg_len, body missing + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + close(sockfd3); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_hello_msg_invalid_len) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - memset(data + 4, 0, 36); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); // Write invalid length. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - uint16_t len = butil::HostToNet16(35); - memcpy(data + 4, &len, sizeof(len)); - memset(data + 6, 0, 34); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); // write invalid length - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RDMA", 4); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + memset(data + 4, 0, 36); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); // Write invalid length. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RDMA", 4); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + uint16_t len = butil::HostToNet16(35); + memcpy(data + 4, &len, sizeof(len)); + memset(data + 6, 0, 34); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); // write invalid length + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_hello_msg_invalid_version) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); - uint16_t ver = butil::HostToNet16(1); - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - memcpy(data + 4, &len, 2); - memset(data + 6, 0, 34); - memcpy(data + 6, &ver, 2); // hello_ver == 1, impl_ver == 0 - // Write the 36B base starting at data + 4 (NOT data). Pre-Step-1 this - // UT mistakenly wrote `data, 36` which included the leftover "RDMA" - // magic at data[0..4); the server parsed it as msg_len = 0x5244 and - // happened to fall through to NegotiationValid (which then failed on - // hello_ver). Now that Step 1 enforces a HELLO_V2_MSG_LEN_MAX upper bound, - // such an oversized msg_len would be rejected before reaching the - // version check, breaking the intent of this UT. - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - uint32_t flags = 0; - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - sockfd1.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - memcpy(data, "RDMA", 4); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - memcpy(data + 4, &len, 2); - memset(data + 6, 0, 32); - memcpy(data + 8, &ver, 2); // hello_ver == 0, impl_ver == 1 - // See comment above on `write(sockfd1, data + 4, 36)` for why we - // write from data + 4 instead of data. - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - sockfd2.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); + uint16_t ver = butil::HostToNet16(1); + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RDMA", 4); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + memcpy(data + 4, &len, 2); + memset(data + 6, 0, 34); + memcpy(data + 6, &ver, 2); // hello_ver == 1, impl_ver == 0 + // Write the 36B base starting at data + 4 (NOT data). Pre-Step-1 this + // UT mistakenly wrote `data, 36` which included the leftover "RDMA" + // magic at data[0..4); the server parsed it as msg_len = 0x5244 and + // happened to fall through to NegotiationValid (which then failed on + // hello_ver). Now that Step 1 enforces a HELLO_V2_MSG_LEN_MAX upper bound, + // such an oversized msg_len would be rejected before reaching the + // version check, breaking the intent of this UT. + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + uint32_t flags = 0; + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + sockfd1.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + memcpy(data, "RDMA", 4); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + memcpy(data + 4, &len, 2); + memset(data + 6, 0, 32); + memcpy(data + 8, &ver, 2); // hello_ver == 0, impl_ver == 1 + // See comment above on `write(sockfd1, data + 4, 36)` for why we + // write from data + 4 instead of data. + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + sockfd2.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_hello_msg_invalid_sq_rq_block_size) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - uint32_t flags = butil::HostToNet32(0); - rdma::v2_wire::HelloMessage msg{}; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - - msg.sq_size = 10; - msg.rq_size = 16; - msg.block_size = 8192; - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - sockfd1.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - msg.sq_size = 16; - msg.rq_size = 10; - msg.block_size = 8192; - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - sockfd2.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 1000; - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd3 >= 0); - ASSERT_EQ(0, connect(sockfd3, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd3, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd3, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_EQ(sizeof(flags), write(sockfd3, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - sockfd3.reset(-1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + uint32_t flags = butil::HostToNet32(0); + rdma::v2_wire::HelloMessage msg{}; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + + msg.sq_size = 10; + msg.rq_size = 16; + msg.block_size = 8192; + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + sockfd1.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + msg.sq_size = 16; + msg.rq_size = 10; + msg.block_size = 8192; + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + sockfd2.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 1000; + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + butil::fd_guard sockfd3(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd3 >= 0); + ASSERT_EQ(0, connect(sockfd3, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd3, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd3, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_EQ(sizeof(flags), write(sockfd3, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + sockfd3.reset(-1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_close_after_qp_build) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - rdma::v2_wire::HelloMessage msg{}; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(40, write(sockfd1, data, 40)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(40, write(sockfd1, data, 40)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_close_during_ack_send) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - rdma::v2_wire::HelloMessage msg{}; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - uint32_t flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + uint32_t flags = butil::HostToNet32(1); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_close_after_ack_send) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - rdma::v2_wire::HelloMessage msg{}; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - close(sockfd1); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - close(sockfd2); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + close(sockfd1); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + flags = butil::HostToNet32(1); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + close(sockfd2); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, client_send_data_on_tcp_after_ack_send) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - Socket *s = nullptr; - rdma::v2_wire::HelloMessage msg{}; - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - - butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd1 >= 0); - ASSERT_EQ(0, connect(sockfd1, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd1, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd2 >= 0); - ASSERT_EQ(0, connect(sockfd2, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(36, write(sockfd2, data + 4, 36)); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - flags = butil::HostToNet32(1); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); // wait for server to handle the msg - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + Socket* s = nullptr; + rdma::v2_wire::HelloMessage msg{}; + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + + butil::fd_guard sockfd1(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd1 >= 0); + ASSERT_EQ(0, connect(sockfd1, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd1, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd1, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(sizeof(flags), write(sockfd1, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + butil::fd_guard sockfd2(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd2 >= 0); + ASSERT_EQ(0, connect(sockfd2, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(4, write(sockfd2, data, 4)); // Write magic string. + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(36, write(sockfd2, data + 4, 36)); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + flags = butil::HostToNet32(1); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); // wait for server to handle the msg + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(sizeof(flags), write(sockfd2, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } static void HandshakeUntilAckWait(butil::fd_guard* sockfd) { @@ -1102,1201 +1102,1201 @@ TEST_F(RdmaTest, fd_and_qp_input_streams_are_separate) { } TEST_F(RdmaTest, server_miss_before_hello_send) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - bthread_id_join(cntl.call_id()); + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + bthread_id_join(cntl.call_id()); - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); } TEST_F(RdmaTest, server_close_before_hello_send) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EEOF, cntl.ErrorCode()); -} + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); -TEST_F(RdmaTest, server_miss_during_magic_str) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(2, write(acc_fd, "RD", 2)); - usleep(100000); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); -} + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); -TEST_F(RdmaTest, server_close_during_magic_str) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(2, write(acc_fd, "RD", 2)); - usleep(100000); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EEOF, cntl.ErrorCode()); -} + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); -TEST_F(RdmaTest, server_hello_invalid_magic_str) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "ABCD", 4)); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EPROTO, cntl.ErrorCode()); -} + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); -TEST_F(RdmaTest, server_miss_during_hello_msg) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); - const uint16_t msg_len = butil::HostToNet16( - static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(static_cast(sizeof(msg_len)), - write(acc_fd, &msg_len, sizeof(msg_len))); - bthread_id_join(cntl.call_id()); + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + close(acc_fd); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); + ASSERT_EQ(EEOF, cntl.ErrorCode()); } -TEST_F(RdmaTest, server_close_during_hello_msg) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); - const uint16_t msg_len = butil::HostToNet16( - static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); - ASSERT_EQ(static_cast(sizeof(msg_len)), - write(acc_fd, &msg_len, sizeof(msg_len))); - close(acc_fd); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); +TEST_F(RdmaTest, server_miss_during_magic_str) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); - ASSERT_EQ(EEOF, cntl.ErrorCode()); -} + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); -TEST_F(RdmaTest, server_hello_invalid_msg_len) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - memcpy(data, "RDMA", 4); - uint16_t len = butil::HostToNet16(35); - memcpy(data + 4, &len, 2); - memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EPROTO, cntl.ErrorCode()); -} + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); -TEST_F(RdmaTest, server_hello_invalid_version) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - memcpy(data, "RDMA", 4); - uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); - memcpy(data + 4, &len, 2); - memset(data + 6, 0, 32); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); - uint32_t *tmp = (uint32_t *)data; - ASSERT_EQ(0, butil::NetToHost32(*tmp)); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); -} + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); -TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = 1; - msg.impl_ver = 1; - msg.sq_size = 0; - msg.rq_size = 0; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); - uint32_t *tmp = (uint32_t *)data; - ASSERT_EQ(0, butil::NetToHost32(*tmp)); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); -} + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(2, write(acc_fd, "RD", 2)); + usleep(100000); + bthread_id_join(cntl.call_id()); -TEST_F(RdmaTest, server_miss_after_ack) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); - uint32_t *tmp = (uint32_t *)data; - ASSERT_EQ(1, butil::NetToHost32(*tmp)); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); } -TEST_F(RdmaTest, server_close_after_ack) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(4, read(acc_fd, data, 4)); - uint32_t *tmp = (uint32_t *)data; - ASSERT_EQ(1, butil::NetToHost32(*tmp)); - close(acc_fd); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EEOF, cntl.ErrorCode()); -} +TEST_F(RdmaTest, server_close_during_magic_str) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); -TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - usleep(100000); - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - bthread_id_join(cntl.call_id()); - - ASSERT_EQ(EPROTO, cntl.ErrorCode()); -} + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); -TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - - // [0..4) magic - ASSERT_EQ(0, memcmp(data, "RDMA", 4)); - // [4..6) msg_len, big-endian uint16 == 40 - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - (size_t)(((uint16_t)data[4] << 8) | (uint16_t)data[5])); - // [6..8) hello_ver, big-endian uint16 == rdma::HELLO_V2_VERSION - ASSERT_EQ(rdma::HELLO_V2_VERSION, - (uint16_t)(((uint16_t)data[6] << 8) | (uint16_t)data[7])); - // [8..10) impl_ver, big-endian uint16 == rdma::IMPL_V2_VERSION - ASSERT_EQ(rdma::IMPL_V2_VERSION, - (uint16_t)(((uint16_t)data[8] << 8) | (uint16_t)data[9])); - - rdma::v2_wire::HelloMessage msg{}; - msg.Deserialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, msg.msg_len); - ASSERT_EQ(rdma::HELLO_V2_VERSION, msg.hello_ver); - ASSERT_EQ(rdma::IMPL_V2_VERSION, msg.impl_ver); - - bthread_id_join(cntl.call_id()); -} + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); -TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - // Send a well-formed v2 hello so the server enters the common ACK_WAIT. - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - - // Read server's reply hello and assert its byte-level layout. - uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - read(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); - - ASSERT_EQ(0, memcmp(reply, "RDMA", 4)); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - (size_t)(((uint16_t)reply[4] << 8) | (uint16_t)reply[5])); - ASSERT_EQ(rdma::HELLO_V2_VERSION, - (uint16_t)(((uint16_t)reply[6] << 8) | (uint16_t)reply[7])); - ASSERT_EQ(rdma::IMPL_V2_VERSION, - (uint16_t)(((uint16_t)reply[8] << 8) | (uint16_t)reply[9])); - - rdma::v2_wire::HelloMessage reply_msg{}; - reply_msg.Deserialize(reply + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, reply_msg.msg_len); - ASSERT_EQ(rdma::HELLO_V2_VERSION, reply_msg.hello_ver); - ASSERT_EQ(rdma::IMPL_V2_VERSION, reply_msg.impl_ver); - - // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends - // cleanly without requiring real RDMA hardware. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); -} + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); -TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != NULL); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - // Build a v2 hello with msg_len = 48 (40 base + 8B zero tail). - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = 48; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - - uint8_t buf[52]; - memcpy(buf, "RDMA", 4); - msg.Serialize(buf + 4); - memset(buf + 40, 0x00, 8); // 8B zero tail - // Coalesce the real ACK with the hello. The v2 parser must consume only - // msg_len bytes and preserve the ACK for the next handshake step. - uint32_t flags = butil::HostToNet32(1); - memcpy(buf + 48, &flags, sizeof(flags)); - ASSERT_EQ(sizeof(buf), write(sockfd, buf, sizeof(buf))); - usleep(100000); - - ASSERT_EQ(handshake::ESTABLISHED, - AdapterTransport::Get(s)->handshake_phase()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); -} + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(2, write(acc_fd, "RD", 2)); + usleep(100000); + close(acc_fd); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); -TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - // Build a v2 hello with msg_len = 4097 (HELLO_V2_MSG_LEN_MAX + 1). - // We only send the 40B base; the server must reject before reading - // (and definitely before attempting to drain) any "tail". - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = 4097; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - - uint8_t buf[rdma::HELLO_V2_MSG_LEN_MIN]; - memcpy(buf, "RDMA", 4); - msg.Serialize(buf + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - sockfd.reset(-1); - usleep(100000); - - StopServer(); + ASSERT_EQ(EEOF, cntl.ErrorCode()); } -// RAII for FLAGS_rdma_client_handshake_version: lets us flip the -// client-side handshake version for a single test and restore it on -// scope exit so subsequent tests stay on the v2 default. -class HandshakeVersionFlag { -public: - explicit HandshakeVersionFlag(int v) - : _saved(rdma::FLAGS_rdma_client_handshake_version) { - rdma::FLAGS_rdma_client_handshake_version = v; - } - ~HandshakeVersionFlag() { - rdma::FLAGS_rdma_client_handshake_version = _saved; - } +TEST_F(RdmaTest, server_hello_invalid_magic_str) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); -private: - int _saved; -}; + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); -// Build a v3 wire packet from an RdmaHello: "RDM3" + pb_size_be + body. -std::string MakeV3Packet(const rdma::RdmaHello &msg) { - std::string body; - EXPECT_TRUE(msg.SerializeToString(&body)); - std::string packet; - packet.reserve(4 + 4 + body.size()); - packet.append("RDM3", 4); - uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); - packet.append(reinterpret_cast(&pb_size_be), 4); - packet.append(body); - return packet; -} + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); -// Build a fully-valid RdmaHello: all 6 required fields are set, with -// values that pass RdmaHelloV3Wire::RdmaHelloValid(). -// - block_size = 8192 (>= MIN_BLOCK_SIZE) -// - sq_size / rq_size = 16 (>= MIN_QP_SIZE) -// - gid = exactly 16B (sizeof(ibv_gid)) -// - qp_num = 0 (allowed because g_skip_rdma_init in UT) -rdma::RdmaHello MakeValidV3Hello() { - rdma::RdmaHello msg; - msg.set_block_size(8192); - msg.set_sq_size(16); - msg.set_rq_size(16); - msg.set_lid(0); - ibv_gid gid = rdma::GetRdmaGid(); - msg.set_gid( - std::string(reinterpret_cast(gid.raw), sizeof(gid.raw))); - msg.set_qp_num(0); - return msg; -} + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); -TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { - HandshakeVersionFlag _hsv(3); - - butil::fd_guard sockfd(butil::tcp_listen(g_ep)); - EXPECT_TRUE(sockfd >= 0); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); - ASSERT_TRUE(acc_fd >= 0); - - // [0..4) magic "RDM3" - uint8_t magic[4]; - ASSERT_EQ(4, read(acc_fd, magic, 4)); - ASSERT_EQ(0, memcmp(magic, "RDM3", 4)); - - // [4..8) pb_size, big-endian uint32, must be in (0, 4096] - uint8_t size_buf[4]; - ASSERT_EQ(4, read(acc_fd, size_buf, 4)); - uint32_t pb_size = - butil::NetToHost32(*reinterpret_cast(size_buf)); - ASSERT_GT(pb_size, 0u); - ASSERT_LE(pb_size, 4096u); - - // [8..8+pb_size) RdmaHello protobuf body. - std::string body(pb_size, '\0'); - ASSERT_EQ((ssize_t)pb_size, read(acc_fd, &body[0], pb_size)); - rdma::RdmaHello msg; - ASSERT_TRUE(msg.ParseFromString(body)); - - // All 6 required fields must be present (ParseFromString would - // have already returned false otherwise). - ASSERT_TRUE(msg.has_block_size()); - ASSERT_TRUE(msg.has_sq_size()); - ASSERT_TRUE(msg.has_rq_size()); - ASSERT_TRUE(msg.has_lid()); - ASSERT_TRUE(msg.has_gid()); - ASSERT_TRUE(msg.has_qp_num()); - // gid wire encoding must be exactly 16 bytes (sizeof(ibv_gid)). - ASSERT_EQ(sizeof(ibv_gid), msg.gid().size()); - - // Let the RPC time out and release resources. - bthread_id_join(cntl.call_id()); -} + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(4, write(acc_fd, "ABCD", 4)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); -TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - // Send a valid v3 hello. - std::string packet = MakeV3Packet(MakeValidV3Hello()); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - - // Read server's reply hello: 4B magic + 4B pb_size + body. - uint8_t reply_magic[4]; - ASSERT_EQ(4, read(sockfd, reply_magic, 4)); - ASSERT_EQ(0, memcmp(reply_magic, "RDM3", 4)); - - uint8_t size_buf[4]; - ASSERT_EQ(4, read(sockfd, size_buf, 4)); - uint32_t pb_size = - butil::NetToHost32(*reinterpret_cast(size_buf)); - ASSERT_GT(pb_size, 0u); - ASSERT_LE(pb_size, 4096u); - - std::string body(pb_size, '\0'); - ASSERT_EQ((ssize_t)pb_size, read(sockfd, &body[0], pb_size)); - rdma::RdmaHello reply; - ASSERT_TRUE(reply.ParseFromString(body)); - ASSERT_TRUE(reply.has_block_size()); - ASSERT_TRUE(reply.has_sq_size()); - ASSERT_TRUE(reply.has_rq_size()); - ASSERT_TRUE(reply.has_gid()); - ASSERT_EQ(sizeof(ibv_gid), reply.gid().size()); - - // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends - // cleanly without requiring real RDMA hardware. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + ASSERT_EQ(EPROTO, cntl.ErrorCode()); } -TEST_F(RdmaTest, v3_server_rejects_zero_pb_size) { - StartServer(); +TEST_F(RdmaTest, server_miss_during_hello_msg) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - // "RDM3" + pb_size = 0 (4B big-endian zero). - uint8_t buf[8] = {'R', 'D', 'M', '3', 0, 0, 0, 0}; - ASSERT_EQ(8, write(sockfd, buf, 8)); - usleep(100000); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); + const uint16_t msg_len = butil::HostToNet16( + static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(static_cast(sizeof(msg_len)), + write(acc_fd, &msg_len, sizeof(msg_len))); + bthread_id_join(cntl.call_id()); - sockfd.reset(-1); - StopServer(); + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_close_during_hello_msg) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(4, write(acc_fd, "RDMA", 4)); + const uint16_t msg_len = butil::HostToNet16( + static_cast(rdma::HELLO_V2_MSG_LEN_MIN)); + ASSERT_EQ(static_cast(sizeof(msg_len)), + write(acc_fd, &msg_len, sizeof(msg_len))); + close(acc_fd); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(EEOF, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_hello_invalid_msg_len) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + memcpy(data, "RDMA", 4); + uint16_t len = butil::HostToNet16(35); + memcpy(data + 4, &len, 2); + memset(data + 6, 0, 32); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::FAILED, AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(EPROTO, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_hello_invalid_version) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + memcpy(data, "RDMA", 4); + uint16_t len = butil::HostToNet16(rdma::HELLO_V2_MSG_LEN_MIN); + memcpy(data + 4, &len, 2); + memset(data + 6, 0, 32); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(4, read(acc_fd, data, 4)); + uint32_t* tmp = (uint32_t*)data; + ASSERT_EQ(0, butil::NetToHost32(*tmp)); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_hello_invalid_sq_rq_size) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = 1; + msg.impl_ver = 1; + msg.sq_size = 0; + msg.rq_size = 0; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(4, read(acc_fd, data, 4)); + uint32_t* tmp = (uint32_t*)data; + ASSERT_EQ(0, butil::NetToHost32(*tmp)); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_miss_after_ack) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(4, read(acc_fd, data, 4)); + uint32_t* tmp = (uint32_t*)data; + ASSERT_EQ(1, butil::NetToHost32(*tmp)); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(ERPCTIMEDOUT, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_close_after_ack) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(4, read(acc_fd, data, 4)); + uint32_t* tmp = (uint32_t*)data; + ASSERT_EQ(1, butil::NetToHost32(*tmp)); + close(acc_fd); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(EEOF, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, server_send_data_on_tcp_after_ack) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::HELLO_WAIT, AdapterTransport::Get(s.get())->handshake_phase()); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + usleep(100000); + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + bthread_id_join(cntl.call_id()); + + ASSERT_EQ(EPROTO, cntl.ErrorCode()); +} + +TEST_F(RdmaTest, v2_client_hello_bytes_baseline) { + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(acc_fd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + + // [0..4) magic + ASSERT_EQ(0, memcmp(data, "RDMA", 4)); + // [4..6) msg_len, big-endian uint16 == 40 + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + (size_t)(((uint16_t)data[4] << 8) | (uint16_t)data[5])); + // [6..8) hello_ver, big-endian uint16 == rdma::HELLO_V2_VERSION + ASSERT_EQ(rdma::HELLO_V2_VERSION, + (uint16_t)(((uint16_t)data[6] << 8) | (uint16_t)data[7])); + // [8..10) impl_ver, big-endian uint16 == rdma::IMPL_V2_VERSION + ASSERT_EQ(rdma::IMPL_V2_VERSION, + (uint16_t)(((uint16_t)data[8] << 8) | (uint16_t)data[9])); + + rdma::v2_wire::HelloMessage msg{}; + msg.Deserialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, msg.msg_len); + ASSERT_EQ(rdma::HELLO_V2_VERSION, msg.hello_ver); + ASSERT_EQ(rdma::IMPL_V2_VERSION, msg.impl_ver); + + bthread_id_join(cntl.call_id()); +} + +TEST_F(RdmaTest, v2_server_hello_bytes_baseline) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + // Send a well-formed v2 hello so the server enters the common ACK_WAIT. + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + + // Read server's reply hello and assert its byte-level layout. + uint8_t reply[rdma::HELLO_V2_MSG_LEN_MIN]; + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + read(sockfd, reply, rdma::HELLO_V2_MSG_LEN_MIN)); + + ASSERT_EQ(0, memcmp(reply, "RDMA", 4)); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + (size_t)(((uint16_t)reply[4] << 8) | (uint16_t)reply[5])); + ASSERT_EQ(rdma::HELLO_V2_VERSION, + (uint16_t)(((uint16_t)reply[6] << 8) | (uint16_t)reply[7])); + ASSERT_EQ(rdma::IMPL_V2_VERSION, + (uint16_t)(((uint16_t)reply[8] << 8) | (uint16_t)reply[9])); + + rdma::v2_wire::HelloMessage reply_msg{}; + reply_msg.Deserialize(reply + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, reply_msg.msg_len); + ASSERT_EQ(rdma::HELLO_V2_VERSION, reply_msg.hello_ver); + ASSERT_EQ(rdma::IMPL_V2_VERSION, reply_msg.impl_ver); + + // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends + // cleanly without requiring real RDMA hardware. + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +TEST_F(RdmaTest, v2_server_preserves_coalesced_ack_after_extension) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != NULL); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + // Build a v2 hello with msg_len = 48 (40 base + 8B zero tail). + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = 48; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + + uint8_t buf[52]; + memcpy(buf, "RDMA", 4); + msg.Serialize(buf + 4); + memset(buf + 40, 0x00, 8); // 8B zero tail + // Coalesce the real ACK with the hello. The v2 parser must consume only + // msg_len bytes and preserve the ACK for the next handshake step. + uint32_t flags = butil::HostToNet32(1); + memcpy(buf + 48, &flags, sizeof(flags)); + ASSERT_EQ(sizeof(buf), write(sockfd, buf, sizeof(buf))); + usleep(100000); + + ASSERT_EQ(handshake::ESTABLISHED, + AdapterTransport::Get(s)->handshake_phase()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +TEST_F(RdmaTest, v2_server_rejects_oversized_msg_len) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + // Build a v2 hello with msg_len = 4097 (HELLO_V2_MSG_LEN_MAX + 1). + // We only send the 40B base; the server must reject before reading + // (and definitely before attempting to drain) any "tail". + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = 4097; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + + uint8_t buf[rdma::HELLO_V2_MSG_LEN_MIN]; + memcpy(buf, "RDMA", 4); + msg.Serialize(buf + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(sockfd, buf, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + sockfd.reset(-1); + usleep(100000); + + StopServer(); +} + +// RAII for FLAGS_rdma_client_handshake_version: lets us flip the +// client-side handshake version for a single test and restore it on +// scope exit so subsequent tests stay on the v2 default. +class HandshakeVersionFlag { +public: + explicit HandshakeVersionFlag(int v) + : _saved(rdma::FLAGS_rdma_client_handshake_version) { + rdma::FLAGS_rdma_client_handshake_version = v; + } + ~HandshakeVersionFlag() { + rdma::FLAGS_rdma_client_handshake_version = _saved; + } + +private: + int _saved; +}; + +// Build a v3 wire packet from an RdmaHello: "RDM3" + pb_size_be + body. +std::string MakeV3Packet(const rdma::RdmaHello& msg) { + std::string body; + EXPECT_TRUE(msg.SerializeToString(&body)); + std::string packet; + packet.reserve(4 + 4 + body.size()); + packet.append("RDM3", 4); + uint32_t pb_size_be = butil::HostToNet32(static_cast(body.size())); + packet.append(reinterpret_cast(&pb_size_be), 4); + packet.append(body); + return packet; +} + +// Build a fully-valid RdmaHello: all 6 required fields are set, with +// values that pass RdmaHelloV3Wire::RdmaHelloValid(). +// - block_size = 8192 (>= MIN_BLOCK_SIZE) +// - sq_size / rq_size = 16 (>= MIN_QP_SIZE) +// - gid = exactly 16B (sizeof(ibv_gid)) +// - qp_num = 0 (allowed because g_skip_rdma_init in UT) +rdma::RdmaHello MakeValidV3Hello() { + rdma::RdmaHello msg; + msg.set_block_size(8192); + msg.set_sq_size(16); + msg.set_rq_size(16); + msg.set_lid(0); + ibv_gid gid = rdma::GetRdmaGid(); + msg.set_gid( + std::string(reinterpret_cast(gid.raw), sizeof(gid.raw))); + msg.set_qp_num(0); + return msg; +} + +TEST_F(RdmaTest, v3_client_hello_bytes_baseline) { + HandshakeVersionFlag _hsv(3); + + butil::fd_guard sockfd(butil::tcp_listen(g_ep)); + EXPECT_TRUE(sockfd >= 0); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + butil::fd_guard acc_fd(accept(sockfd, nullptr, nullptr)); + ASSERT_TRUE(acc_fd >= 0); + + // [0..4) magic "RDM3" + uint8_t magic[4]; + ASSERT_EQ(4, read(acc_fd, magic, 4)); + ASSERT_EQ(0, memcmp(magic, "RDM3", 4)); + + // [4..8) pb_size, big-endian uint32, must be in (0, 4096] + uint8_t size_buf[4]; + ASSERT_EQ(4, read(acc_fd, size_buf, 4)); + uint32_t pb_size = + butil::NetToHost32(*reinterpret_cast(size_buf)); + ASSERT_GT(pb_size, 0u); + ASSERT_LE(pb_size, 4096u); + + // [8..8+pb_size) RdmaHello protobuf body. + std::string body(pb_size, '\0'); + ASSERT_EQ((ssize_t)pb_size, read(acc_fd, &body[0], pb_size)); + rdma::RdmaHello msg; + ASSERT_TRUE(msg.ParseFromString(body)); + + // All 6 required fields must be present (ParseFromString would + // have already returned false otherwise). + ASSERT_TRUE(msg.has_block_size()); + ASSERT_TRUE(msg.has_sq_size()); + ASSERT_TRUE(msg.has_rq_size()); + ASSERT_TRUE(msg.has_lid()); + ASSERT_TRUE(msg.has_gid()); + ASSERT_TRUE(msg.has_qp_num()); + // gid wire encoding must be exactly 16 bytes (sizeof(ibv_gid)). + ASSERT_EQ(sizeof(ibv_gid), msg.gid().size()); + + // Let the RPC time out and release resources. + bthread_id_join(cntl.call_id()); +} + +TEST_F(RdmaTest, v3_server_hello_bytes_baseline) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + // Send a valid v3 hello. + std::string packet = MakeV3Packet(MakeValidV3Hello()); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + + // Read server's reply hello: 4B magic + 4B pb_size + body. + uint8_t reply_magic[4]; + ASSERT_EQ(4, read(sockfd, reply_magic, 4)); + ASSERT_EQ(0, memcmp(reply_magic, "RDM3", 4)); + + uint8_t size_buf[4]; + ASSERT_EQ(4, read(sockfd, size_buf, 4)); + uint32_t pb_size = + butil::NetToHost32(*reinterpret_cast(size_buf)); + ASSERT_GT(pb_size, 0u); + ASSERT_LE(pb_size, 4096u); + + std::string body(pb_size, '\0'); + ASSERT_EQ((ssize_t)pb_size, read(sockfd, &body[0], pb_size)); + rdma::RdmaHello reply; + ASSERT_TRUE(reply.ParseFromString(body)); + ASSERT_TRUE(reply.has_block_size()); + ASSERT_TRUE(reply.has_sq_size()); + ASSERT_TRUE(reply.has_rq_size()); + ASSERT_TRUE(reply.has_gid()); + ASSERT_EQ(sizeof(ibv_gid), reply.gid().size()); + + // Drive the server into FALLBACK_TCP via ACK flags=0 so the test ends + // cleanly without requiring real RDMA hardware. + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); +} + +TEST_F(RdmaTest, v3_server_rejects_zero_pb_size) { + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + // "RDM3" + pb_size = 0 (4B big-endian zero). + uint8_t buf[8] = {'R', 'D', 'M', '3', 0, 0, 0, 0}; + ASSERT_EQ(8, write(sockfd, buf, 8)); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + sockfd.reset(-1); + StopServer(); } TEST_F(RdmaTest, v3_server_rejects_oversized_pb_size) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - uint8_t buf[8]; - memcpy(buf, "RDM3", 4); - // pb_size just above the allowed maximum -> rejected. - uint32_t pb_size_be = - butil::HostToNet32(static_cast(rdma::HELLO_V3_MAX_PB_SIZE + 1)); - memcpy(buf + 4, &pb_size_be, 4); - ASSERT_EQ(8, write(sockfd, buf, 8)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - sockfd.reset(-1); - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + uint8_t buf[8]; + memcpy(buf, "RDM3", 4); + // pb_size just above the allowed maximum -> rejected. + uint32_t pb_size_be = + butil::HostToNet32(static_cast(rdma::HELLO_V3_MAX_PB_SIZE + 1)); + memcpy(buf + 4, &pb_size_be, 4); + ASSERT_EQ(8, write(sockfd, buf, 8)); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + sockfd.reset(-1); + StopServer(); } TEST_F(RdmaTest, v3_server_rejects_invalid_pb_bytes) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - // "RDM3" + pb_size = 8 + 8 bytes of 0xff (invalid protobuf body). - uint8_t buf[16]; - memcpy(buf, "RDM3", 4); - uint32_t pb_size_be = butil::HostToNet32(8); - memcpy(buf + 4, &pb_size_be, 4); - memset(buf + 8, 0xff, 8); - ASSERT_EQ(16, write(sockfd, buf, 16)); - usleep(100000); - - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - sockfd.reset(-1); - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + // "RDM3" + pb_size = 8 + 8 bytes of 0xff (invalid protobuf body). + uint8_t buf[16]; + memcpy(buf, "RDM3", 4); + uint32_t pb_size_be = butil::HostToNet32(8); + memcpy(buf + 4, &pb_size_be, 4); + memset(buf + 8, 0xff, 8); + ASSERT_EQ(16, write(sockfd, buf, 16)); + usleep(100000); + + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + sockfd.reset(-1); + StopServer(); } TEST_F(RdmaTest, v3_server_invalid_sq_size_falls_back) { - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - rdma::RdmaHello msg = MakeValidV3Hello(); - msg.set_sq_size(0); // invalid: < MIN_QP_SIZE (16) - std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - - // Server validated the hello as invalid -> _rdma_state = RDMA_OFF, - // but still proceeds to the common ACK_WAIT (sends its own reply hello). - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - - // Drain server's reply hello (content not asserted here; covered - // by v3_server_hello_bytes_baseline). - uint8_t reply_hdr[8]; - ASSERT_EQ(8, read(sockfd, reply_hdr, 8)); - ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); - uint32_t reply_pb_size = - butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); - std::string reply_body(reply_pb_size, '\0'); - ASSERT_EQ((ssize_t)reply_pb_size, - read(sockfd, &reply_body[0], reply_pb_size)); - - // Client ACK flags=0 -> server settles into FALLBACK_TCP. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + rdma::RdmaHello msg = MakeValidV3Hello(); + msg.set_sq_size(0); // invalid: < MIN_QP_SIZE (16) + std::string packet = MakeV3Packet(msg); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); + + // Server validated the hello as invalid -> _rdma_state = RDMA_OFF, + // but still proceeds to the common ACK_WAIT (sends its own reply hello). + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + + // Drain server's reply hello (content not asserted here; covered + // by v3_server_hello_bytes_baseline). + uint8_t reply_hdr[8]; + ASSERT_EQ(8, read(sockfd, reply_hdr, 8)); + ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); + uint32_t reply_pb_size = + butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); + std::string reply_body(reply_pb_size, '\0'); + ASSERT_EQ((ssize_t)reply_pb_size, + read(sockfd, &reply_body[0], reply_pb_size)); + + // Client ACK flags=0 -> server settles into FALLBACK_TCP. + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } // RAII guard to toggle FLAGS_rdma_ece for a single test and restore it. class EceFlagGuard { public: - explicit EceFlagGuard(bool v) : _saved(rdma::FLAGS_rdma_ece) { - rdma::FLAGS_rdma_ece = v; - } - ~EceFlagGuard() { rdma::FLAGS_rdma_ece = _saved; } + explicit EceFlagGuard(bool v) : _saved(rdma::FLAGS_rdma_ece) { + rdma::FLAGS_rdma_ece = v; + } + ~EceFlagGuard() { rdma::FLAGS_rdma_ece = _saved; } private: - bool _saved; + bool _saved; }; // Build a valid v3 hello that also carries an ECE block. rdma::RdmaHello MakeValidV3HelloWithEce(uint32_t vendor_id, uint32_t options, uint32_t comp_mask) { - rdma::RdmaHello msg = MakeValidV3Hello(); - rdma::RdmaEce *ece = msg.mutable_ece(); - ece->set_vendor_id(vendor_id); - ece->set_options(options); - ece->set_comp_mask(comp_mask); - return msg; + rdma::RdmaHello msg = MakeValidV3Hello(); + rdma::RdmaEce* ece = msg.mutable_ece(); + ece->set_vendor_id(vendor_id); + ece->set_options(options); + ece->set_comp_mask(comp_mask); + return msg; } // Read the server's v3 reply hello (4B magic + 4B pb_size + body) and parse // it into `reply`. Asserts the framing along the way. -static void ReadServerV3Reply(int fd, rdma::RdmaHello *reply) { - uint8_t reply_hdr[8]; - ASSERT_EQ(8, read(fd, reply_hdr, 8)); - ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); - uint32_t reply_pb_size = - butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); - ASSERT_GT(reply_pb_size, 0u); - ASSERT_LE(reply_pb_size, 4096u); - std::string reply_body(reply_pb_size, '\0'); - ASSERT_EQ((ssize_t)reply_pb_size, read(fd, &reply_body[0], reply_pb_size)); - ASSERT_TRUE(reply->ParseFromString(reply_body)); +static void ReadServerV3Reply(int fd, rdma::RdmaHello* reply) { + uint8_t reply_hdr[8]; + ASSERT_EQ(8, read(fd, reply_hdr, 8)); + ASSERT_EQ(0, memcmp(reply_hdr, "RDM3", 4)); + uint32_t reply_pb_size = + butil::NetToHost32(*reinterpret_cast(reply_hdr + 4)); + ASSERT_GT(reply_pb_size, 0u); + ASSERT_LE(reply_pb_size, 4096u); + std::string reply_body(reply_pb_size, '\0'); + ASSERT_EQ((ssize_t)reply_pb_size, read(fd, &reply_body[0], reply_pb_size)); + ASSERT_TRUE(reply->ParseFromString(reply_body)); } // A client hello carrying ECE must not break the server handshake: with ECE // enabled the server still parses the hello and advances to the common // ACK_WAIT. TEST_F(RdmaTest, v3_server_accepts_client_hello_with_ece) { - EceFlagGuard ece_flag_guard(true); - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); - std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - - rdma::RdmaHello reply; - ReadServerV3Reply(sockfd, &reply); - - // ACK flags=0 -> clean FALLBACK_TCP so the test ends without hardware. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - StopServer(); + EceFlagGuard ece_flag_guard(true); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); + std::string packet = MakeV3Packet(msg); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); + + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + + rdma::RdmaHello reply; + ReadServerV3Reply(sockfd, &reply); + + // ACK flags=0 -> clean FALLBACK_TCP so the test ends without hardware. + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + StopServer(); } // When ECE negotiation is disabled, the server reply must NOT advertise ECE, // even if the client advertised it (FillLocalRdmaHello degrade branch #1). TEST_F(RdmaTest, v3_server_reply_has_no_ece_when_disabled) { - EceFlagGuard ece_flag_guard(false); - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); - std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - - rdma::RdmaHello reply; - ReadServerV3Reply(sockfd, &reply); - EXPECT_FALSE(reply.has_ece()); - - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - - sockfd.reset(-1); - usleep(100000); - StopServer(); + EceFlagGuard ece_flag_guard(false); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); + std::string packet = MakeV3Packet(msg); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); + + rdma::RdmaHello reply; + ReadServerV3Reply(sockfd, &reply); + EXPECT_FALSE(reply.has_ece()); + + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + + sockfd.reset(-1); + usleep(100000); + StopServer(); } // When ECE is enabled but there is no negotiated result (UT skips the real QP @@ -2304,259 +2304,259 @@ TEST_F(RdmaTest, v3_server_reply_has_no_ece_when_disabled) { // still NOT advertise ECE (FillLocalRdmaHello degrade branch #2 -> // degrade-safe). TEST_F(RdmaTest, v3_server_reply_has_no_ece_without_hw_negotiation) { - EceFlagGuard ece_flag_guard(true); - StartServer(); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - - rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); - std::string packet = MakeV3Packet(msg); - ASSERT_EQ((ssize_t)packet.size(), - write(sockfd, packet.data(), packet.size())); - usleep(100000); - - rdma::RdmaHello reply; - ReadServerV3Reply(sockfd, &reply); - EXPECT_FALSE(reply.has_ece()); - - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - - sockfd.reset(-1); - usleep(100000); - StopServer(); + EceFlagGuard ece_flag_guard(true); + StartServer(); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + + rdma::RdmaHello msg = MakeValidV3HelloWithEce(0x02c9, 0x1, 0x0); + std::string packet = MakeV3Packet(msg); + ASSERT_EQ((ssize_t)packet.size(), + write(sockfd, packet.data(), packet.size())); + usleep(100000); + + rdma::RdmaHello reply; + ReadServerV3Reply(sockfd, &reply); + EXPECT_FALSE(reply.has_ece()); + + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ((ssize_t)sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + + sockfd.reset(-1); + usleep(100000); + StopServer(); } class ResourceAllocFailGuard { public: - explicit ResourceAllocFailGuard(bool v) - : _saved(rdma::g_fail_resource_alloc_for_test) { - rdma::g_fail_resource_alloc_for_test = v; - } - ~ResourceAllocFailGuard() { rdma::g_fail_resource_alloc_for_test = _saved; } + explicit ResourceAllocFailGuard(bool v) + : _saved(rdma::g_fail_resource_alloc_for_test) { + rdma::g_fail_resource_alloc_for_test = v; + } + ~ResourceAllocFailGuard() { rdma::g_fail_resource_alloc_for_test = _saved; } private: - bool _saved; + bool _saved; }; TEST_F(RdmaTest, client_alloc_resource_fail_fallback_tcp) { - StartServer(); - ResourceAllocFailGuard alloc_fail_guard(true); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - req.set_sleep_us(200000); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - // The socket must not be failed, otherwise it can no longer carry TCP. - ASSERT_FALSE(s->Failed()); - - // The RPC still completes over TCP. - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText(); - - StopServer(); + StartServer(); + ResourceAllocFailGuard alloc_fail_guard(true); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + req.set_sleep_us(200000); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); + + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + // The socket must not be failed, otherwise it can no longer carry TCP. + ASSERT_FALSE(s->Failed()); + + // The RPC still completes over TCP. + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText(); + + StopServer(); } TEST_F(RdmaTest, server_alloc_resource_fail_fallback_tcp) { - StartServer(); - ResourceAllocFailGuard alloc_fail_guard(true); - - sockaddr_in addr; - bzero((char *)&addr, sizeof(addr)); - addr.sin_family = AF_INET; - addr.sin_port = htons(PORT); - butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); - ASSERT_TRUE(sockfd >= 0); - ASSERT_EQ(0, connect(sockfd, (sockaddr *)&addr, sizeof(sockaddr))); - usleep(100000); // wait for server to handle the msg - Socket *s = GetSocketFromServer(0); - ASSERT_TRUE(s != nullptr); - ASSERT_EQ(handshake::UNINITIALIZED, - AdapterTransport::Get(s)->handshake_phase()); - - // Send a well-formed v2 hello: the negotiation succeeds - // but the resource allocation does not. - rdma::v2_wire::HelloMessage msg{}; - msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; - msg.hello_ver = rdma::HELLO_V2_VERSION; - msg.impl_ver = rdma::IMPL_V2_VERSION; - msg.sq_size = 16; - msg.rq_size = 16; - msg.block_size = 8192; - msg.qp_num = 0; - msg.gid = rdma::GetRdmaGid(); - - uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; - memcpy(data, "RDMA", 4); - msg.Serialize(data + 4); - ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, - write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); - usleep(100000); - ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); - ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); - ASSERT_FALSE(s->Failed()); - - // Ack without RDMA so that the server finishes the handshake in TCP mode. - uint32_t flags = butil::HostToNet32(0); - ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); - usleep(100000); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s)->handshake_phase()); - ASSERT_FALSE(s->Failed()); - - sockfd.reset(-1); - usleep(100000); - ASSERT_EQ(nullptr, GetSocketFromServer(0)); - - StopServer(); + StartServer(); + ResourceAllocFailGuard alloc_fail_guard(true); + + sockaddr_in addr; + bzero((char*)&addr, sizeof(addr)); + addr.sin_family = AF_INET; + addr.sin_port = htons(PORT); + butil::fd_guard sockfd(socket(AF_INET, SOCK_STREAM, 0)); + ASSERT_TRUE(sockfd >= 0); + ASSERT_EQ(0, connect(sockfd, (sockaddr*)&addr, sizeof(sockaddr))); + usleep(100000); // wait for server to handle the msg + Socket* s = GetSocketFromServer(0); + ASSERT_TRUE(s != nullptr); + ASSERT_EQ(handshake::UNINITIALIZED, + AdapterTransport::Get(s)->handshake_phase()); + + // Send a well-formed v2 hello: the negotiation succeeds + // but the resource allocation does not. + rdma::v2_wire::HelloMessage msg{}; + msg.msg_len = rdma::HELLO_V2_MSG_LEN_MIN; + msg.hello_ver = rdma::HELLO_V2_VERSION; + msg.impl_ver = rdma::IMPL_V2_VERSION; + msg.sq_size = 16; + msg.rq_size = 16; + msg.block_size = 8192; + msg.qp_num = 0; + msg.gid = rdma::GetRdmaGid(); + + uint8_t data[rdma::HELLO_V2_MSG_LEN_MIN]; + memcpy(data, "RDMA", 4); + msg.Serialize(data + 4); + ASSERT_EQ(rdma::HELLO_V2_MSG_LEN_MIN, + write(sockfd, data, rdma::HELLO_V2_MSG_LEN_MIN)); + usleep(100000); + ASSERT_EQ(handshake::ACK_WAIT, AdapterTransport::Get(s)->handshake_phase()); + ASSERT_EQ(RdmaTransport::RDMA_OFF, RdmaTransport::Get(s)->_rdma_state); + ASSERT_FALSE(s->Failed()); + + // Ack without RDMA so that the server finishes the handshake in TCP mode. + uint32_t flags = butil::HostToNet32(0); + ASSERT_EQ(sizeof(flags), write(sockfd, &flags, sizeof(flags))); + usleep(100000); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s)->handshake_phase()); + ASSERT_FALSE(s->Failed()); + + sockfd.reset(-1); + usleep(100000); + ASSERT_EQ(nullptr, GetSocketFromServer(0)); + + StopServer(); } TEST_F(RdmaTest, try_global_disable_rdma) { - StartServer(); - rdma::g_rdma_available.store(false, butil::memory_order_relaxed); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - - req.set_message(__FUNCTION__); - req.set_sleep_us(200000); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ASSERT_EQ(handshake::FALLBACK_TCP, - AdapterTransport::Get(s.get())->handshake_phase()); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); - - StopServer(); - rdma::g_rdma_available.store(true, butil::memory_order_relaxed); + StartServer(); + rdma::g_rdma_available.store(false, butil::memory_order_relaxed); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + + req.set_message(__FUNCTION__); + req.set_sleep_us(200000); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ASSERT_EQ(handshake::FALLBACK_TCP, + AdapterTransport::Get(s.get())->handshake_phase()); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); + + StopServer(); + rdma::g_rdma_available.store(true, butil::memory_order_relaxed); } TEST_F(RdmaTest, server_option_invalid) { - Server server; - ServerOptions options; - options.socket_mode = SOCKET_MODE_RDMA; + Server server; + ServerOptions options; + options.socket_mode = SOCKET_MODE_RDMA; - // rtmp and rdma are incompatible - options.rtmp_service = (RtmpService *)1; - ASSERT_EQ(-1, server.Start(PORT, &options)); + // rtmp and rdma are incompatible + options.rtmp_service = (RtmpService*)1; + ASSERT_EQ(-1, server.Start(PORT, &options)); - // nshead and rdma are incompatible - options.rtmp_service = nullptr; - options.nshead_service = (NsheadService *)1; - ASSERT_EQ(-1, server.Start(PORT, &options)); + // nshead and rdma are incompatible + options.rtmp_service = nullptr; + options.nshead_service = (NsheadService*)1; + ASSERT_EQ(-1, server.Start(PORT, &options)); - // mongo and rdma are incompatible - options.nshead_service = nullptr; - options.mongo_service_adaptor = (MongoServiceAdaptor *)1; - ASSERT_EQ(-1, server.Start(PORT, &options)); + // mongo and rdma are incompatible + options.nshead_service = nullptr; + options.mongo_service_adaptor = (MongoServiceAdaptor*)1; + ASSERT_EQ(-1, server.Start(PORT, &options)); - // ssl and rdma are incompatible - options.mongo_service_adaptor = nullptr; - options.mutable_ssl_options()->default_cert.certificate = "test"; - ASSERT_EQ(-1, server.Start(PORT, &options)); + // ssl and rdma are incompatible + options.mongo_service_adaptor = nullptr; + options.mutable_ssl_options()->default_cert.certificate = "test"; + ASSERT_EQ(-1, server.Start(PORT, &options)); } TEST_F(RdmaTest, channel_option_invalid) { - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; - // rtmp and rdma are incompatible - chan_options.protocol = "rtmp"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // rtmp and rdma are incompatible + chan_options.protocol = "rtmp"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - chan_options.protocol = "streaming_rpc"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + chan_options.protocol = "streaming_rpc"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // nshead and rdma are incompatible - chan_options.protocol = "nshead"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - chan_options.protocol = "nshead_mcpack"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // nshead and rdma are incompatible + chan_options.protocol = "nshead"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + chan_options.protocol = "nshead_mcpack"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // nova_pbrpc and rdma are incompatible - chan_options.protocol = "nova_pbrpc"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // nova_pbrpc and rdma are incompatible + chan_options.protocol = "nova_pbrpc"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // public_pbrpc and rdma are incompatible - chan_options.protocol = "public_pbrpc"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // public_pbrpc and rdma are incompatible + chan_options.protocol = "public_pbrpc"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // redis and rdma are incompatible - chan_options.protocol = "redis"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // redis and rdma are incompatible + chan_options.protocol = "redis"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // memcache and rdma are incompatible - chan_options.protocol = "memcache"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // memcache and rdma are incompatible + chan_options.protocol = "memcache"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // ubrpc and rdma are incompatible - chan_options.protocol = "ubrpc_compack"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // ubrpc and rdma are incompatible + chan_options.protocol = "ubrpc_compack"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // itp and rdma are incompatible - chan_options.protocol = "itp"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // itp and rdma are incompatible + chan_options.protocol = "itp"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // esp and rdma are incompatible - chan_options.protocol = "esp"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // esp and rdma are incompatible + chan_options.protocol = "esp"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // hulu_pbrpc and rdma are incompatible - chan_options.protocol = "hulu_pbrpc"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // hulu_pbrpc and rdma are incompatible + chan_options.protocol = "hulu_pbrpc"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // sofa_pbrpc and rdma are incompatible - chan_options.protocol = "sofa_pbrpc"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // sofa_pbrpc and rdma are incompatible + chan_options.protocol = "sofa_pbrpc"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // http and rdma are incompatible - chan_options.protocol = "http"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // http and rdma are incompatible + chan_options.protocol = "http"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); - // ssl and rdma are incompatible - chan_options.protocol = "baidu_std"; - chan_options.mutable_ssl_options()->sni_name = "test"; - ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); + // ssl and rdma are incompatible + chan_options.protocol = "baidu_std"; + chan_options.mutable_ssl_options()->sni_name = "test"; + ASSERT_EQ(-1, channel.Init(g_ep, &chan_options)); } static const int E2E_RPC_NUM = 32; @@ -2583,7 +2583,7 @@ static int SendEchoRpcs(Channel& channel, int rpc_num, size_t attach_size, req[i].set_code(i + 1); if (attach_size > 0) { EXPECT_EQ(0, attach[i].resize( - attach_size, static_cast('a' + i % 26))); + attach_size, static_cast('a' + i % 26))); cntl[i].request_attachment().append(attach[i]); } ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], DoNothing()); @@ -2609,99 +2609,98 @@ static int SendEchoRpcs(Channel& channel, int rpc_num, size_t attach_size, return succeeded; } - TEST_P(RdmaRpcTest, rdma_client_to_rdma_server) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - // usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + // usleep(100000); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); + + StopServer(); } TEST_P(RdmaRpcTest, tcp_client_to_tcp_server) { - StartServer(false); - - Channel channel; - ChannelOptions chan_options; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); - - StopServer(); + StartServer(false); + + Channel channel; + ChannelOptions chan_options; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); + + StopServer(); } TEST_P(RdmaRpcTest, tcp_client_to_rdma_server) { - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(0, cntl.ErrorCode()); - - StopServer(); + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(0, cntl.ErrorCode()); + + StopServer(); } TEST_P(RdmaRpcTest, rdma_client_to_tcp_server) { - StartServer(false); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - usleep(100000); - bthread_id_join(cntl.call_id()); - ASSERT_FALSE(cntl.Failed()); - - StopServer(); + StartServer(false); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + usleep(100000); + bthread_id_join(cntl.call_id()); + ASSERT_FALSE(cntl.Failed()); + + StopServer(); } TEST_P(RdmaRpcTest, tcp_client_to_rdma_server_short_connection) { @@ -2716,7 +2715,7 @@ TEST_P(RdmaRpcTest, tcp_client_to_rdma_server_short_connection) { ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); for (int round = 0; round < 8; ++round) { ASSERT_EQ(E2E_RPC_NUM, SendEchoRpcs(channel, E2E_RPC_NUM, 4096)) - << "round=" << round; + << "round=" << round; } StopServer(); @@ -2768,365 +2767,367 @@ TEST_P(RdmaRpcTest, rdma_server_survives_connection_churn) { static const int RPC_NUM = 1024; -void DumpRdmaEndpointInfo(Socket *client, Socket *server) { - std::cout << std::endl << "client:"; - RdmaTransport::Get(client)->_rdma_ep->DebugInfo(std::cout); - std::cout << std::endl << "server:"; - RdmaTransport::Get(server)->_rdma_ep->DebugInfo(std::cout); +void DumpRdmaEndpointInfo(Socket* client, Socket* server) { + std::cout << std::endl + << "client:"; + RdmaTransport::Get(client)->_rdma_ep->DebugInfo(std::cout); + std::cout << std::endl + << "server:"; + RdmaTransport::Get(server)->_rdma_ep->DebugInfo(std::cout); } TEST_P(RdmaRpcTest, send_rpcs_in_one_qp) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 50000; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - LOG(INFO) << "send 0 attachment"; - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - } - ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - } - - LOG(INFO) << "send 4KB attachment"; - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - cntl[i].Reset(); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - } - ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - } - - LOG(INFO) << "send 1MB attachment"; - attach.resize(1048576); - for (int i = 0; i < RPC_NUM; ++i) { - cntl[i].Reset(); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - } - ASSERT_TRUE(0 == cntl[i].ErrorCode() || EOVERCROWDED == cntl[i].ErrorCode()) - << "req[" << i << "] " << berror(cntl[i].ErrorCode()); - } - - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[0]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 50000; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + LOG(INFO) << "send 0 attachment"; + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + } + ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; + } + + LOG(INFO) << "send 4KB attachment"; + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + cntl[i].Reset(); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + } + ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; + } + + LOG(INFO) << "send 1MB attachment"; + attach.resize(1048576); + for (int i = 0; i < RPC_NUM; ++i) { + cntl[i].Reset(); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + } + ASSERT_TRUE(0 == cntl[i].ErrorCode() || EOVERCROWDED == cntl[i].ErrorCode()) + << "req[" << i << "] " << berror(cntl[i].ErrorCode()); + } + + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[0]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + + StopServer(); } TEST_P(RdmaRpcTest, send_rpc_in_many_qp) { - if (!FLAGS_rdma_test_enable) { - return; - } - - butil::ip_t ip; - ASSERT_EQ(0, butil::str2ip(g_ip.c_str(), &ip)); - - Server server[100]; - MyEchoService svc[100]; - int num = 100; - butil::EndPoint server_eps[100]; - for (int i = 0; i < num; ++i) { - ServerOptions options; - options.socket_mode = SOCKET_MODE_RDMA; - options.idle_timeout_sec = 1; - options.max_concurrency = 0; - options.internal_port = -1; - server[i].AddService(&svc[i], SERVER_DOESNT_OWN_SERVICE); - ASSERT_EQ(0, server[i].Start(0, &options)); - server_eps[i] = butil::EndPoint(ip, server[i].listen_address().port); - } - - int port = 0; - butil::IOBuf attach; - attach.resize(4096); - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 100000; - chan_options.max_retry = 0; - Channel channel[RPC_NUM]; - Server *svr[RPC_NUM]; - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - for (int i = 0; i < RPC_NUM; ++i) { - svr[i] = &server[i % num]; - ASSERT_EQ(0, channel[i].Init(server_eps[(port++) % num], &chan_options)); - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel[i]) - .Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - EXPECT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - if (s && svr[i] && svr[i]->_am) { - std::vector sids; - svr[i]->_am->ListConnections(&sids); - for (size_t j = 0; j < sids.size(); ++j) { - SocketUniquePtr m; - if (Socket::AddressFailedAsWell(sids[j], &m) == 0) { - DumpRdmaEndpointInfo(s.get(), m.get()); - } + if (!FLAGS_rdma_test_enable) { + return; + } + + butil::ip_t ip; + ASSERT_EQ(0, butil::str2ip(g_ip.c_str(), &ip)); + + Server server[100]; + MyEchoService svc[100]; + int num = 100; + butil::EndPoint server_eps[100]; + for (int i = 0; i < num; ++i) { + ServerOptions options; + options.socket_mode = SOCKET_MODE_RDMA; + options.idle_timeout_sec = 1; + options.max_concurrency = 0; + options.internal_port = -1; + server[i].AddService(&svc[i], SERVER_DOESNT_OWN_SERVICE); + ASSERT_EQ(0, server[i].Start(0, &options)); + server_eps[i] = butil::EndPoint(ip, server[i].listen_address().port); + } + + int port = 0; + butil::IOBuf attach; + attach.resize(4096); + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 100000; + chan_options.max_retry = 0; + Channel channel[RPC_NUM]; + Server* svr[RPC_NUM]; + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + for (int i = 0; i < RPC_NUM; ++i) { + svr[i] = &server[i % num]; + ASSERT_EQ(0, channel[i].Init(server_eps[(port++) % num], &chan_options)); + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel[i]) + .Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + EXPECT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + if (s && svr[i] && svr[i]->_am) { + std::vector sids; + svr[i]->_am->ListConnections(&sids); + for (size_t j = 0; j < sids.size(); ++j) { + SocketUniquePtr m; + if (Socket::AddressFailedAsWell(sids[j], &m) == 0) { + DumpRdmaEndpointInfo(s.get(), m.get()); + } + } + } } - } + EXPECT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; } - EXPECT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - } - for (int i = 0; i < num; ++i) { - server[i].Stop(0); - server[i].Join(); - } + for (int i = 0; i < num; ++i) { + server[i].Stop(0); + server[i].Join(); + } } TEST_P(RdmaRpcTest, send_rpcs_as_pooled_connection) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 30000; // it may very slow - chan_options.timeout_ms = 30000; - chan_options.max_retry = 0; - chan_options.connection_type = "pooled"; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - } - ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 30000; // it may very slow + chan_options.timeout_ms = 30000; + chan_options.max_retry = 0; + chan_options.connection_type = "pooled"; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + } + ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; + } + + StopServer(); } TEST_P(RdmaRpcTest, send_rpcs_as_short_connection) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 30000; // it may very slow - chan_options.timeout_ms = 30000; - chan_options.max_retry = 0; - chan_options.connection_type = "short"; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); - Socket *m = GetSocketFromServer(0); - DumpRdmaEndpointInfo(s.get(), m); - } - ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 30000; // it may very slow + chan_options.timeout_ms = 30000; + chan_options.max_retry = 0; + chan_options.connection_type = "short"; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (cntl[i].ErrorCode() == ERPCTIMEDOUT) { + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl[i]._single_server_id, &s)); + Socket* m = GetSocketFromServer(0); + DumpRdmaEndpointInfo(s.get(), m); + } + ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; + } + + StopServer(); } TEST_P(RdmaRpcTest, server_stop_during_rpc) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 3000; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - if (i == 0) - StopServer(); - int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || error_code == EEOF || - error_code == ELOGOFF || error_code == EHOSTDOWN) - << "req[" << i << "]: " << error_code; - } + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 3000; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + if (i == 0) + StopServer(); + int error_code = cntl[i].ErrorCode(); + ASSERT_TRUE(error_code == 0 || error_code == EEOF || + error_code == ELOGOFF || error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; + } } TEST_P(RdmaRpcTest, server_close_during_rpc) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 3000; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - if (i == RPC_NUM / 2) { - req[i].set_close_fd(true); - } - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || error_code == EEOF || - error_code == EFAILEDSOCKET || error_code == EHOSTDOWN) - << "req[" << i << "]: " << error_code; - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 3000; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + if (i == RPC_NUM / 2) { + req[i].set_close_fd(true); + } + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + int error_code = cntl[i].ErrorCode(); + ASSERT_TRUE(error_code == 0 || error_code == EEOF || + error_code == EFAILEDSOCKET || error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; + } + + StopServer(); } TEST_P(RdmaRpcTest, client_close_during_rpc) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 3000; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - attach.resize(4096); - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - - cntl[0].CloseConnection("Close connection"); - - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - int error_code = cntl[i].ErrorCode(); - ASSERT_TRUE(error_code == 0 || error_code == ECLOSE || - error_code == EHOSTDOWN) - << "req[" << i << "]: " << error_code; - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 3000; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + attach.resize(4096); + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + + cntl[0].CloseConnection("Close connection"); + + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + int error_code = cntl[i].ErrorCode(); + ASSERT_TRUE(error_code == 0 || error_code == ECLOSE || + error_code == EHOSTDOWN) + << "req[" << i << "]: " << error_code; + } + + StopServer(); } TEST_P(RdmaRpcTest, rdma_client_close_during_rpc_repeatedly) { @@ -3154,7 +3155,6 @@ TEST_P(RdmaRpcTest, rdma_client_close_during_rpc_repeatedly) { ASSERT_FALSE(HasFailure()) << "round=" << round; } - const int served = g_echo_served.load(butil::memory_order_relaxed) - served_before - CHURN_ROUND_NUM; LOG(INFO) << "server served " << served << " of " @@ -3170,218 +3170,218 @@ TEST_P(RdmaRpcTest, rdma_client_close_during_rpc_repeatedly) { } TEST_P(RdmaRpcTest, verbs_error_handling) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - req.set_sleep_us(200000); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); - - usleep(100000); // wait for rdma handshake complete - - SocketUniquePtr s; - ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); - ibv_send_wr wr; - memset(&wr, 0, sizeof(wr)); - ibv_sge sge; - void *buf = malloc(8192); - sge.addr = (uint64_t)buf; - sge.length = 8192; - sge.lkey = 1; // incorrect lkey - wr.sg_list = &sge; - wr.num_sge = 1; - ibv_send_wr *bad = nullptr; - auto rdma_transport = RdmaTransport::Get(s); - ibv_post_send(rdma_transport->_rdma_ep->_resource->qp, &wr, &bad); - bthread_id_join(cntl.call_id()); - ASSERT_EQ(ERDMA, cntl.ErrorCode()); - free(buf); - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + req.set_sleep_us(200000); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, done); + + usleep(100000); // wait for rdma handshake complete + + SocketUniquePtr s; + ASSERT_EQ(0, Socket::Address(cntl._single_server_id, &s)); + ibv_send_wr wr; + memset(&wr, 0, sizeof(wr)); + ibv_sge sge; + void* buf = malloc(8192); + sge.addr = (uint64_t)buf; + sge.length = 8192; + sge.lkey = 1; // incorrect lkey + wr.sg_list = &sge; + wr.num_sge = 1; + ibv_send_wr* bad = nullptr; + auto rdma_transport = RdmaTransport::Get(s); + ibv_post_send(rdma_transport->_rdma_ep->_resource->qp, &wr, &bad); + bthread_id_join(cntl.call_id()); + ASSERT_EQ(ERDMA, cntl.ErrorCode()); + free(buf); + + StopServer(); } TEST_P(RdmaRpcTest, rdma_use_parallel_channel) { - if (!FLAGS_rdma_test_enable) { - return; - } + if (!FLAGS_rdma_test_enable) { + return; + } - StartServer(); + StartServer(); - const size_t NCHANS = 8; - Channel subchans[NCHANS]; - ParallelChannel channel; - ChannelOptions opts; - opts.socket_mode = SOCKET_MODE_RDMA; - for (size_t i = 0; i < NCHANS; ++i) { - ASSERT_EQ(0, subchans[i].Init(_naming_url.c_str(), "rR", &opts)); - ASSERT_EQ(0, channel.AddChannel(&subchans[i], DOESNT_OWN_CHANNEL, nullptr, - nullptr)); - } - ASSERT_EQ(0, channel.Init(nullptr)); + const size_t NCHANS = 8; + Channel subchans[NCHANS]; + ParallelChannel channel; + ChannelOptions opts; + opts.socket_mode = SOCKET_MODE_RDMA; + for (size_t i = 0; i < NCHANS; ++i) { + ASSERT_EQ(0, subchans[i].Init(_naming_url.c_str(), "rR", &opts)); + ASSERT_EQ(0, channel.AddChannel(&subchans[i], DOESNT_OWN_CHANNEL, nullptr, + nullptr)); + } + ASSERT_EQ(0, channel.Init(nullptr)); - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, nullptr); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, nullptr); - ASSERT_EQ(0, cntl.ErrorCode()); - ASSERT_EQ(NCHANS, (size_t)cntl.sub_count()); + ASSERT_EQ(0, cntl.ErrorCode()); + ASSERT_EQ(NCHANS, (size_t)cntl.sub_count()); - StopServer(); + StopServer(); } TEST_P(RdmaRpcTest, rdma_use_selective_channel) { - if (!FLAGS_rdma_test_enable) { - return; - } + if (!FLAGS_rdma_test_enable) { + return; + } - StartServer(); + StartServer(); - const size_t NCHANS = 8; - SelectiveChannel channel; - ChannelOptions opts; - opts.socket_mode = SOCKET_MODE_RDMA; - ASSERT_EQ(0, channel.Init("rr", &opts)); - for (size_t i = 0; i < NCHANS; ++i) { - Channel *subchan = new Channel; - ASSERT_EQ(0, subchan->Init(_naming_url.c_str(), "rR", &opts)); - ASSERT_EQ(0, channel.AddChannel(subchan, nullptr)); - } + const size_t NCHANS = 8; + SelectiveChannel channel; + ChannelOptions opts; + opts.socket_mode = SOCKET_MODE_RDMA; + ASSERT_EQ(0, channel.Init("rr", &opts)); + for (size_t i = 0; i < NCHANS; ++i) { + Channel* subchan = new Channel; + ASSERT_EQ(0, subchan->Init(_naming_url.c_str(), "rR", &opts)); + ASSERT_EQ(0, channel.AddChannel(subchan, nullptr)); + } - Controller cntl; - test::EchoRequest req; - test::EchoResponse res; - req.set_message(__FUNCTION__); - ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, nullptr); + Controller cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + ::test::EchoService::Stub(&channel).Echo(&cntl, &req, &res, nullptr); - ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText(); - ASSERT_EQ(1, cntl.sub_count()); + ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText(); + ASSERT_EQ(1, cntl.sub_count()); - StopServer(); + StopServer(); } -static void MockFree(void *buf) {} +static void MockFree(void* buf) {} TEST_P(RdmaRpcTest, send_rpcs_with_user_defined_iobuf) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 500; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf attach; - void *data = malloc(4096); - ; - attach.append_user_data(data, 4096, nullptr); - req[0].set_message(__FUNCTION__); - cntl[0].request_attachment().append(attach); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[0], &req[0], &res[0], done); - bthread_id_join(cntl[0].call_id()); - ASSERT_EQ(ERDMAMEM, cntl[0].ErrorCode()); - attach.clear(); - sleep(2); // wait for client recover from EHOSTDOWN - cntl[0].Reset(); - - char *mr[2 * RPC_NUM]; - uint32_t lkey[2 * RPC_NUM]; - for (size_t i = 0; i < RPC_NUM; ++i) { - mr[2 * i] = (char *)malloc(4096); - memset(mr[2 * i], i % 100, 4096); - lkey[2 * i] = rdma::RegisterMemoryForRdma(mr[2 * i], 4096); - ASSERT_TRUE(lkey[2 * i] != 0); - cntl[i].request_attachment().append_user_data_with_meta( - mr[2 * i] + i, 4096 - i, MockFree, lkey[2 * i]); - mr[2 * i + 1] = (char *)malloc(4096); - memset(mr[2 * i + 1], i % 100, 4096); - lkey[2 * i + 1] = rdma::RegisterMemoryForRdma(mr[2 * i + 1], 4096); - ASSERT_TRUE(lkey[2 * i + 1] != 0); - cntl[i].request_attachment().append_user_data_with_meta( - mr[2 * i + 1] + i, 4096 - i, MockFree, lkey[2 * i + 1]); - req[i].set_message(__FUNCTION__); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (size_t i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; - rdma::DeregisterMemoryForRdma(mr[i]); - ASSERT_EQ(2 * (4096 - i), cntl[i].response_attachment().size()); - char tmp[8192]; - cntl[i].response_attachment().copy_to(tmp, 2 * (4096 - i)); - ASSERT_EQ(0, memcmp(mr[2 * i] + i, tmp, 4096 - i)); - ASSERT_EQ(0, memcmp(mr[2 * i + 1] + i, tmp + 4096 - i, 4096 - i)); - free(mr[2 * i]); - free(mr[2 * i + 1]); - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 500; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf attach; + void* data = malloc(4096); + ; + attach.append_user_data(data, 4096, nullptr); + req[0].set_message(__FUNCTION__); + cntl[0].request_attachment().append(attach); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[0], &req[0], &res[0], done); + bthread_id_join(cntl[0].call_id()); + ASSERT_EQ(ERDMAMEM, cntl[0].ErrorCode()); + attach.clear(); + sleep(2); // wait for client recover from EHOSTDOWN + cntl[0].Reset(); + + char* mr[2 * RPC_NUM]; + uint32_t lkey[2 * RPC_NUM]; + for (size_t i = 0; i < RPC_NUM; ++i) { + mr[2 * i] = (char*)malloc(4096); + memset(mr[2 * i], i % 100, 4096); + lkey[2 * i] = rdma::RegisterMemoryForRdma(mr[2 * i], 4096); + ASSERT_TRUE(lkey[2 * i] != 0); + cntl[i].request_attachment().append_user_data_with_meta( + mr[2 * i] + i, 4096 - i, MockFree, lkey[2 * i]); + mr[2 * i + 1] = (char*)malloc(4096); + memset(mr[2 * i + 1], i % 100, 4096); + lkey[2 * i + 1] = rdma::RegisterMemoryForRdma(mr[2 * i + 1], 4096); + ASSERT_TRUE(lkey[2 * i + 1] != 0); + cntl[i].request_attachment().append_user_data_with_meta( + mr[2 * i + 1] + i, 4096 - i, MockFree, lkey[2 * i + 1]); + req[i].set_message(__FUNCTION__); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (size_t i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + ASSERT_EQ(0, cntl[i].ErrorCode()) << "req[" << i << "]"; + rdma::DeregisterMemoryForRdma(mr[i]); + ASSERT_EQ(2 * (4096 - i), cntl[i].response_attachment().size()); + char tmp[8192]; + cntl[i].response_attachment().copy_to(tmp, 2 * (4096 - i)); + ASSERT_EQ(0, memcmp(mr[2 * i] + i, tmp, 4096 - i)); + ASSERT_EQ(0, memcmp(mr[2 * i + 1] + i, tmp + 4096 - i, 4096 - i)); + free(mr[2 * i]); + free(mr[2 * i + 1]); + } + + StopServer(); } TEST_P(RdmaRpcTest, try_memory_pool_empty) { - if (!FLAGS_rdma_test_enable) { - return; - } - - StartServer(); - - Channel channel; - ChannelOptions chan_options; - chan_options.socket_mode = SOCKET_MODE_RDMA; - chan_options.connect_timeout_ms = 500; - chan_options.timeout_ms = 60000; - chan_options.max_retry = 0; - ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); - Controller cntl[RPC_NUM]; - test::EchoRequest req[RPC_NUM]; - test::EchoResponse res[RPC_NUM]; - - butil::IOBuf iobuf[RPC_NUM]; - for (int i = 0; i < 1024; ++i) { - if (iobuf[i].resize(1048576 * 8)) { - // 8MB for each iobuf - break; - } - } - - for (int i = 0; i < RPC_NUM; ++i) { - req[i].set_message(__FUNCTION__); - cntl[i].request_attachment().append(iobuf[i]); - google::protobuf::Closure *done = DoNothing(); - ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); - } - for (int i = 0; i < RPC_NUM; ++i) { - bthread_id_join(cntl[i].call_id()); - } - - StopServer(); + if (!FLAGS_rdma_test_enable) { + return; + } + + StartServer(); + + Channel channel; + ChannelOptions chan_options; + chan_options.socket_mode = SOCKET_MODE_RDMA; + chan_options.connect_timeout_ms = 500; + chan_options.timeout_ms = 60000; + chan_options.max_retry = 0; + ASSERT_EQ(0, channel.Init(g_ep, &chan_options)); + Controller cntl[RPC_NUM]; + test::EchoRequest req[RPC_NUM]; + test::EchoResponse res[RPC_NUM]; + + butil::IOBuf iobuf[RPC_NUM]; + for (int i = 0; i < 1024; ++i) { + if (iobuf[i].resize(1048576 * 8)) { + // 8MB for each iobuf + break; + } + } + + for (int i = 0; i < RPC_NUM; ++i) { + req[i].set_message(__FUNCTION__); + cntl[i].request_attachment().append(iobuf[i]); + google::protobuf::Closure* done = DoNothing(); + ::test::EchoService::Stub(&channel).Echo(&cntl[i], &req[i], &res[i], done); + } + for (int i = 0; i < RPC_NUM; ++i) { + bthread_id_join(cntl[i].call_id()); + } + + StopServer(); } // Run every TEST_P(RdmaRpcTest, ...) above twice: once with the @@ -3391,24 +3391,24 @@ TEST_P(RdmaRpcTest, try_memory_pool_empty) { // proves the upper-layer RPC paths behave identically under either // wire format. INSTANTIATE_TEST_SUITE_P(HandshakeVersion, RdmaRpcTest, ::testing::Values(2, 3), - [](const ::testing::TestParamInfo &info) { - return std::string("v") + std::to_string(info.param); + [](const ::testing::TestParamInfo& info) { + return std::string("v") + std::to_string(info.param); }); -#endif // if BRPC_WITH_RDMA +#endif // if BRPC_WITH_RDMA -int main(int argc, char *argv[]) { - testing::InitGoogleTest(&argc, argv); - GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true); +int main(int argc, char* argv[]) { + testing::InitGoogleTest(&argc, argv); + GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true); #if BRPC_WITH_RDMA - rdma::FLAGS_rdma_trace_verbose = true; - rdma::FLAGS_rdma_memory_pool_max_regions = 2; - FLAGS_log_idle_connection_close = true; - if (!FLAGS_rdma_test_enable) { - // skip UT requiring rdma runtime environment - rdma::g_rdma_available.store(true, butil::memory_order_relaxed); - rdma::g_skip_rdma_init = true; - } -#endif // if BRPC_WITH_RDMA - return RUN_ALL_TESTS(); + rdma::FLAGS_rdma_trace_verbose = true; + rdma::FLAGS_rdma_memory_pool_max_regions = 2; + FLAGS_log_idle_connection_close = true; + if (!FLAGS_rdma_test_enable) { + // skip UT requiring rdma runtime environment + rdma::g_rdma_available.store(true, butil::memory_order_relaxed); + rdma::g_skip_rdma_init = true; + } +#endif // if BRPC_WITH_RDMA + return RUN_ALL_TESTS(); } diff --git a/test/brpc_rtmp_unittest.cpp b/test/brpc_rtmp_unittest.cpp index 4315031e81..397e5d3bcd 100644 --- a/test/brpc_rtmp_unittest.cpp +++ b/test/brpc_rtmp_unittest.cpp @@ -269,9 +269,16 @@ class PlayingDummyStream : public brpc::RtmpServerStream { << " ms before responding play request"; bthread_usleep(_sleep_ms * 1000L); } + // Reserve a reference for the sender bthread, keeping the stream alive + // until SendData() returns, even if a failed send synchronously runs + // OnStop() and releases the framework's references. The reference + // belongs to the sender from the moment it is acquired, so only the + // failed creation path gives it back. + AddRefManually(); int rc = bthread_start_background(&_play_thread, nullptr, RunSendData, this); if (rc) { + RemoveRefManually(); status->set_error(rc, "Fail to create thread"); return; } @@ -289,8 +296,17 @@ class PlayingDummyStream : public brpc::RtmpServerStream { void OnStop() { LOG(INFO) << "OnStop of PlayingDummyStream=" << this; if (_state.exchange(STATE_STOPPED) == STATE_PLAYING) { + // A send failure can invoke this callback on the sender bthread + // itself. bthread_stop() only sets the stop flag that SendData's + // loop checks, so it must be called on the current bthread too, + // otherwise the sender loops and sleeps forever while holding its + // stream reference. Only joining the current bthread is unsafe, so + // skip just the join on the self-callback path; the sender's own + // reference keeps the stream alive until SendData returns. bthread_stop(_play_thread); - bthread_join(_play_thread, nullptr); + if (_play_thread != bthread_self()) { + bthread_join(_play_thread, nullptr); + } } } @@ -298,7 +314,10 @@ class PlayingDummyStream : public brpc::RtmpServerStream { private: static void* RunSendData(void* arg) { - ((PlayingDummyStream*)arg)->SendData(); + // Adopt the reference acquired before starting this bthread. + butil::intrusive_ptr stream( + static_cast(arg), false); + stream->SendData(); return nullptr; } @@ -342,6 +361,90 @@ void PlayingDummyStream::SendData() { LOG(INFO) << "Quit SendData of PlayingDummyStream=" << this; } +// Regression test for a use-after-free where a send failure synchronously ran +// OnStop() on the sender bthread and dropped the framework's reference while +// SendData() was still using the stream. This reproduces the ordering +// deterministically without a socket: the sender bthread runs OnStop() itself, +// which releases the framework's reference, and then keeps touching the stream. +class SelfStopStream : public brpc::RtmpServerStream { +public: + SelfStopStream(butil::atomic* destroyed, + butil::atomic* alive_after_stop, + butil::atomic* stop_on_sender) + : _destroyed(destroyed) + , _alive_after_stop(alive_after_stop) + , _stop_on_sender(stop_on_sender) + , _sender_ready(false) {} + ~SelfStopStream() override { + _destroyed->store(true, butil::memory_order_relaxed); + } + + void OnStop() override { + _stop_on_sender->store(_sender == bthread_self(), + butil::memory_order_relaxed); + // Drop the framework's reference from within OnStop(), exactly like + // RtmpServerStream::RunOnFailed does on the sender bthread. + _framework_ref.reset(); + } + + // Starts the sender bthread and returns its id. The framework keeps one + // reference until OnStop() drops it; a second reference is reserved for + // the sender bthread, exactly like PlayingDummyStream::OnPlay does. + bthread_t Start() { + _framework_ref.reset(this); + AddRefManually(); + CHECK_EQ(0, bthread_start_background( + &_sender, nullptr, RunSender, this)); + // Publish _sender before the sender bthread reads it in OnStop(). + _sender_ready.store(true, butil::memory_order_release); + return _sender; + } + +private: + static void* RunSender(void* arg) { + butil::intrusive_ptr self( + static_cast(arg), false); + while (!self->_sender_ready.load(butil::memory_order_acquire)) { + bthread_usleep(100); + } + // Run OnStop() on this very bthread, releasing the framework reference. + self->OnStop(); + // The stream must still be alive here, kept by the sender's reference. + self->_alive_after_stop->store( + !self->_destroyed->load(butil::memory_order_relaxed) && + self->ref_count() >= 1, + butil::memory_order_relaxed); + return nullptr; + } + + butil::atomic* _destroyed; + butil::atomic* _alive_after_stop; + butil::atomic* _stop_on_sender; + butil::atomic _sender_ready; + bthread_t _sender; + butil::intrusive_ptr _framework_ref; +}; + +TEST(RtmpTest, on_stop_on_sender_bthread_keeps_stream_alive) { + butil::atomic destroyed(false); + butil::atomic alive_after_stop(false); + butil::atomic stop_on_sender(false); + bthread_t sender; + { + butil::intrusive_ptr stream( + new SelfStopStream(&destroyed, &alive_after_stop, &stop_on_sender)); + sender = stream->Start(); + ASSERT_EQ(0, bthread_join(sender, nullptr)); + // OnStop ran on the sender bthread and the stream survived it. + ASSERT_TRUE(stop_on_sender.load(butil::memory_order_relaxed)); + ASSERT_TRUE(alive_after_stop.load(butil::memory_order_relaxed)); + // The sender's reference is gone but the test still holds one. + ASSERT_FALSE(destroyed.load(butil::memory_order_relaxed)); + } + // The last reference is gone now, the stream must be destroyed. + ASSERT_TRUE(destroyed.load(butil::memory_order_relaxed)); +} + class PlayingDummyService : public brpc::RtmpService { public: PlayingDummyService(int64_t sleep_ms = 0) : _sleep_ms(sleep_ms) {} diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index e94d850815..898e3e7962 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -19,9 +19,12 @@ // Date: Sun Jul 13 15:04:18 CST 2014 +#include #include #include #include +#include +#include #include #include #include "butil/time.h" @@ -49,8 +52,10 @@ #include "brpc/builtin/bad_method_service.h" #include "brpc/server.h" #include "brpc/nshead_service.h" +#include "brpc/acceptor.h" #include "brpc/restful.h" #include "brpc/channel.h" +#include "brpc/redis.h" #include "brpc/socket_map.h" #include "brpc/controller.h" #include "brpc/compress.h" @@ -1426,6 +1431,410 @@ TEST_F(ServerTest, close_idle_connections) { ASSERT_EQ(0ul, stat.connection_count); } +// Returns a port nothing is listening on, or -1. `ServerOptions.internal_port` +// has to be an explicit number, Server::Start() rejects 0 because it stands +// for an ephemeral port, so ask the system for a free one rather than hardcode +// a port that another test may be listening on. +int PickUnusedPort() { + butil::fd_guard sockfd(butil::tcp_listen(butil::EndPoint(butil::IP_ANY, 0))); + if (sockfd < 0) { + return -1; + } + butil::EndPoint point; + if (butil::get_local_side(sockfd, &point) != 0) { + return -1; + } + return point.port; +} + +// Starts `server` on an ephemeral port and fills `options->internal_port` with +// another one. Both are released before Start() binds them and something else +// may take one in between, hence the retries. Returns 0 on success. +int StartWithInternalPort(brpc::Server* server, brpc::ServerOptions* options) { + for (int i = 0; i < 10; ++i) { + int internal_port = PickUnusedPort(); + if (internal_port < 0) { + continue; + } + options->internal_port = internal_port; + if (0 == server->Start("127.0.0.1:0", options)) { + return 0; + } + } + return -1; +} + +static testing::AssertionResult WaitForServerConnections( + const brpc::Server& server, size_t expected) { + brpc::ServerStatistics stat; + const int64_t deadline = butil::gettimeofday_us() + 1000000; + do { + server.GetStat(&stat); + if (stat.connection_count == expected) { + return testing::AssertionSuccess(); + } + usleep(1000); + } while (butil::gettimeofday_us() < deadline); + return testing::AssertionFailure() + << "Expected " << expected << " connections, got " + << stat.connection_count; +} + +static void ExpectConnectionClosed(int fd) { + const struct timeval timeout = {1, 0}; + ASSERT_EQ(0, setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, + &timeout, sizeof(timeout))); + char response; + ssize_t nr; + do { + nr = recv(fd, &response, sizeof(response), 0); + } while (nr < 0 && errno == EINTR); + // Fresh idle clients have sent no data: rejection must close the fd + // without sending a protocol-specific error or waiting for a request. + EXPECT_EQ(0, nr); +} + +class ConnectionLimitPingHandler : public brpc::RedisCommandHandler { +public: + brpc::RedisCommandHandlerResult Run( + brpc::RedisConnContext*, const std::vector&, + brpc::RedisReply* output, bool) override { + output->SetStatus("PONG"); + return brpc::REDIS_CMD_HANDLED; + } +}; + +TEST_F(ServerTest, connection_limit_with_mixed_protocols) { + ConnectionLimitPingHandler ping_handler; + EchoServiceImpl echo_service; + brpc::Server server; + ASSERT_EQ(0, server.AddService( + &echo_service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions opt; + opt.redis_service = new brpc::RedisService; + opt.redis_service->AddCommandHandler("ping", &ping_handler); + opt.max_connections = 3; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + + // Each protocol uses a separate persistent connection on the same port. + brpc::ChannelOptions copt; + copt.connection_type = "single"; + copt.timeout_ms = 1000; + copt.max_retry = 0; + copt.connection_group = "connection_limit_rpc"; + brpc::Channel rpc_channel; + ASSERT_EQ(0, rpc_channel.Init(server.listen_address(), &copt)); + brpc::Controller rpc_cntl; + test::EchoRequest req; + test::EchoResponse res; + req.set_message(EXP_REQUEST); + test::EchoService_Stub stub(&rpc_channel); + stub.Echo(&rpc_cntl, &req, &res, nullptr); + ASSERT_FALSE(rpc_cntl.Failed()) << rpc_cntl.ErrorText(); + + copt.protocol = "redis"; + copt.connection_group = "connection_limit_redis"; + brpc::Channel redis_channel; + ASSERT_EQ(0, redis_channel.Init(server.listen_address(), &copt)); + brpc::RedisRequest redis_req; + brpc::RedisResponse redis_res; + brpc::Controller redis_cntl; + ASSERT_TRUE(redis_req.AddCommand("ping")); + redis_channel.CallMethod( + nullptr, &redis_cntl, &redis_req, &redis_res, nullptr); + ASSERT_FALSE(redis_cntl.Failed()) << redis_cntl.ErrorText(); + ASSERT_EQ(1, redis_res.reply_size()); + ASSERT_STREQ("PONG", redis_res.reply(0).c_str()); + + copt.protocol = "http"; + copt.connection_type = "pooled"; + copt.connection_group = "connection_limit_http"; + brpc::Channel http_channel; + ASSERT_EQ(0, http_channel.Init(server.listen_address(), &copt)); + brpc::Controller http_cntl; + http_cntl.http_request().uri() = "/status"; + http_channel.CallMethod(nullptr, &http_cntl, nullptr, nullptr, nullptr); + ASSERT_FALSE(http_cntl.Failed()) << http_cntl.ErrorText(); + ASSERT_TRUE(WaitForServerConnections(server, 3)); + + butil::fd_guard rejected_client( + tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(rejected_client, 0); + ExpectConnectionClosed(rejected_client); + brpc::ServerStatistics stat; + server.GetStat(&stat); + EXPECT_EQ(3ul, stat.connection_count); + EXPECT_EQ(1ul, stat.rejected_connection_count); + + // Requests on admitted sockets still work when the listener is full. + rpc_cntl.Reset(); + stub.Echo(&rpc_cntl, &req, &res, nullptr); + ASSERT_FALSE(rpc_cntl.Failed()) << rpc_cntl.ErrorText(); +} + +TEST_F(ServerTest, connection_limit_runtime_updates) { + brpc::Server server; + EXPECT_EQ(-1, server.SetMaxConnections(1)); + ASSERT_EQ(0, server.Start("127.0.0.1:0", nullptr)); + const butil::EndPoint ep = server.listen_address(); + butil::fd_guard first_client(tcp_connect(ep, nullptr)); + butil::fd_guard second_client(tcp_connect(ep, nullptr)); + ASSERT_GE(first_client, 0); + ASSERT_GE(second_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 2)); + + // Enabling a limit must count clients admitted while it was unlimited. + ASSERT_EQ(0, server.SetMaxConnections(1)); + butil::fd_guard rejected_client(tcp_connect(ep, nullptr)); + ASSERT_GE(rejected_client, 0); + ExpectConnectionClosed(rejected_client); + ASSERT_TRUE(WaitForServerConnections(server, 2)); + + // A separate Server remains reachable while this public listener is full. + EchoServiceImpl rpc_service; + brpc::Server rpc_server; + ASSERT_EQ(0, rpc_server.AddService( + &rpc_service, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, rpc_server.Start("127.0.0.1:0", nullptr)); + SendSleepRPC(rpc_server.listen_address(), 0, true); + + ASSERT_EQ(0, server.SetMaxConnections(3)); + butil::fd_guard third_client(tcp_connect(ep, nullptr)); + ASSERT_GE(third_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 3)); + + // Lowering the limit leaves established clients connected, but rejects + // new clients until the active count drops strictly below the limit. + ASSERT_EQ(0, server.SetMaxConnections(1)); + ASSERT_TRUE(WaitForServerConnections(server, 3)); + butil::fd_guard lowered_limit_client(tcp_connect(ep, nullptr)); + ASSERT_GE(lowered_limit_client, 0); + ExpectConnectionClosed(lowered_limit_client); + first_client.reset(-1); + second_client.reset(-1); + ASSERT_TRUE(WaitForServerConnections(server, 1)); + butil::fd_guard at_limit_client(tcp_connect(ep, nullptr)); + ASSERT_GE(at_limit_client, 0); + ExpectConnectionClosed(at_limit_client); + + third_client.reset(-1); + ASSERT_TRUE(WaitForServerConnections(server, 0)); + butil::fd_guard recovered_client(tcp_connect(ep, nullptr)); + ASSERT_GE(recovered_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); + + ASSERT_EQ(0, server.SetMaxConnections(0)); + first_client.reset(tcp_connect(ep, nullptr)); + second_client.reset(tcp_connect(ep, nullptr)); + ASSERT_GE(first_client, 0); + ASSERT_GE(second_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 3)); + brpc::ServerStatistics stat; + server.GetStat(&stat); + EXPECT_EQ(3ul, stat.rejected_connection_count); + EXPECT_EQ(0ul, server.options().max_connections); + + ASSERT_EQ(0, server.Stop(0)); + EXPECT_EQ(-1, server.SetMaxConnections(1)); + ASSERT_EQ(0, server.Join()); + ASSERT_TRUE(WaitForServerConnections(server, 0)); +} + +TEST_F(ServerTest, connection_limit_keeps_internal_listener_available) { + brpc::Server server; + brpc::ServerOptions opt; + opt.max_connections = 1; + ASSERT_EQ(0, StartWithInternalPort(&server, &opt)); + butil::fd_guard public_client( + tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(public_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); + + butil::EndPoint internal_ep = server.listen_address(); + internal_ep.port = opt.internal_port; + butil::fd_guard first_internal(tcp_connect(internal_ep, nullptr)); + butil::fd_guard second_internal(tcp_connect(internal_ep, nullptr)); + ASSERT_GE(first_internal, 0); + ASSERT_GE(second_internal, 0); + // GetStat includes both listeners; only the public count is limited. + ASSERT_TRUE(WaitForServerConnections(server, 3)); + brpc::ChannelOptions copt; + copt.protocol = "http"; + copt.connection_type = "short"; + copt.timeout_ms = 1000; + copt.max_retry = 0; + brpc::Channel channel; + ASSERT_EQ(0, channel.Init(internal_ep, &copt)); + brpc::Controller cntl; + cntl.http_request().uri() = "/status"; + channel.CallMethod(nullptr, &cntl, nullptr, nullptr, nullptr); + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + ASSERT_TRUE(WaitForServerConnections(server, 3)); + + butil::fd_guard rejected_client( + tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(rejected_client, 0); + ExpectConnectionClosed(rejected_client); + brpc::ServerStatistics stat; + server.GetStat(&stat); + EXPECT_EQ(3ul, stat.connection_count); + EXPECT_EQ(1ul, stat.rejected_connection_count); + + // Internal connections must not prevent public connection slot recovery. + public_client.reset(-1); + ASSERT_TRUE(WaitForServerConnections(server, 2)); + public_client.reset(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(public_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 3)); +} + +TEST_F(ServerTest, connection_limit_is_per_listener) { + // One business service is shared by two Servers listening on different + // ports. Each port must have its own limit, count and runtime updates. + EchoServiceImpl shared_service; + brpc::Server servers[2]; + for (size_t i = 0; i < 2; ++i) { + ASSERT_EQ(0, servers[i].AddService( + &shared_service, brpc::SERVER_DOESNT_OWN_SERVICE)); + brpc::ServerOptions opt; + opt.max_connections = i + 1; + ASSERT_EQ(0, servers[i].Start("127.0.0.1:0", &opt)); + } + butil::fd_guard first_client( + tcp_connect(servers[0].listen_address(), nullptr)); + ASSERT_GE(first_client, 0); + ASSERT_TRUE(WaitForServerConnections(servers[0], 1)); + butil::fd_guard second_port_clients[2]; + for (auto& client : second_port_clients) { + client.reset(tcp_connect(servers[1].listen_address(), nullptr)); + ASSERT_GE(client, 0); + } + ASSERT_TRUE(WaitForServerConnections(servers[1], 2)); + + for (auto& server : servers) { + butil::fd_guard rejected_client( + tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(rejected_client, 0); + ExpectConnectionClosed(rejected_client); + brpc::ServerStatistics stat; + server.GetStat(&stat); + EXPECT_EQ(1ul, stat.rejected_connection_count); + } + + ASSERT_EQ(0, servers[0].SetMaxConnections(2)); + butil::fd_guard extra_client( + tcp_connect(servers[0].listen_address(), nullptr)); + ASSERT_GE(extra_client, 0); + ASSERT_TRUE(WaitForServerConnections(servers[0], 2)); + butil::fd_guard still_rejected( + tcp_connect(servers[1].listen_address(), nullptr)); + ASSERT_GE(still_rejected, 0); + ExpectConnectionClosed(still_rejected); + brpc::ServerStatistics stat; + servers[1].GetStat(&stat); + EXPECT_EQ(2ul, stat.connection_count); + EXPECT_EQ(2ul, stat.rejected_connection_count); +} + +TEST_F(ServerTest, connection_limit_restarts_with_startup_setting) { + brpc::Server server; + brpc::ServerOptions opt; + opt.max_connections = 1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + ASSERT_EQ(0, server.SetMaxConnections(0)); + EXPECT_EQ(1ul, server.options().max_connections); + butil::fd_guard clients[2]; + for (auto& client : clients) { + client.reset(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(client, 0); + } + ASSERT_TRUE(WaitForServerConnections(server, 2)); + ASSERT_EQ(0, server.Stop(0)); + ASSERT_EQ(0, server.Join()); + ASSERT_TRUE(WaitForServerConnections(server, 0)); + + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + clients[0].reset(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(clients[0], 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); + clients[1].reset(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(clients[1], 0); + ExpectConnectionClosed(clients[1]); +} + +TEST_F(ServerTest, connection_limit_rejects_before_tls_handshake) { + brpc::Server server; + brpc::ServerOptions opt; + opt.max_connections = 1; + opt.force_ssl = true; + brpc::CertInfo& cert = opt.mutable_ssl_options()->default_cert; + cert.certificate = "cert1.crt"; + cert.private_key = "cert1.key"; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + + const butil::EndPoint ep = server.listen_address(); + // An idle TCP socket consumes a slot without initiating a TLS handshake. + butil::fd_guard first_client(tcp_connect(ep, nullptr)); + ASSERT_GE(first_client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); + butil::fd_guard rejected_client(tcp_connect(ep, nullptr)); + ASSERT_GE(rejected_client, 0); + // EOF without a ClientHello verifies rejection before TLS authentication. + ExpectConnectionClosed(rejected_client); + brpc::ServerStatistics stat; + server.GetStat(&stat); + EXPECT_EQ(1ul, stat.connection_count); + EXPECT_EQ(1ul, stat.rejected_connection_count); +} + +TEST_F(ServerTest, connection_limit_socket_creation_failure) { + brpc::Server server; + brpc::ServerOptions opt; + opt.max_connections = 1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + brpc::Acceptor* acceptor = server._am; + ASSERT_TRUE(acceptor->TryAcquireConnectionSlot()); + + brpc::SocketOptions socket_opt; + socket_opt.user = acceptor; + // Avoid STREAM_FAKE_FD, which is intentionally accepted without an OS fd. + socket_opt.fd = std::numeric_limits::max() - 1; + brpc::SocketId id; + ASSERT_NE(0, brpc::Socket::Create(socket_opt, &id)); + // Socket::Create can call BeforeRecycle on failure. Since the socket was + // never inserted into the map, the accept loop must still own the slot. + ASSERT_EQ(1ul, acceptor->ConnectionCount()); + acceptor->ReleaseConnectionSlot(); + butil::fd_guard client(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); +} + +TEST_F(ServerTest, connection_limit_socket_recycled_before_registration) { + brpc::Server server; + brpc::ServerOptions opt; + opt.max_connections = 1; + ASSERT_EQ(0, server.Start("127.0.0.1:0", &opt)); + brpc::Acceptor* acceptor = server._am; + ASSERT_TRUE(acceptor->TryAcquireConnectionSlot()); + + brpc::SocketOptions socket_opt; + socket_opt.user = acceptor; + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(socket_opt, &id)); + // Exercise a socket recycled after successful creation, before the accept + // loop can address it and insert it into the connection map. + ASSERT_EQ(0, brpc::Socket::SetFailed(id)); + brpc::SocketUniquePtr socket; + ASSERT_EQ(-1, brpc::Socket::AddressFailedAsWell(id, &socket)); + ASSERT_EQ(1ul, acceptor->ConnectionCount()); + acceptor->ReleaseConnectionSlot(); + + butil::fd_guard client(tcp_connect(server.listen_address(), nullptr)); + ASSERT_GE(client, 0); + ASSERT_TRUE(WaitForServerConnections(server, 1)); +} + TEST_F(ServerTest, logoff_and_multiple_start) { butil::Timer timer; butil::EndPoint ep; @@ -1737,39 +2146,6 @@ void CallEchoByHttp(const butil::EndPoint& ep, brpc::Controller* cntl) { chan.CallMethod(nullptr, cntl, &req, &res, nullptr); } -// Returns a port nothing is listening on, or -1. `ServerOptions.internal_port` -// has to be an explicit number, Server::Start() rejects 0 because it stands -// for an ephemeral port, so ask the system for a free one rather than hardcode -// a port that another test may be listening on. -int PickUnusedPort() { - butil::fd_guard sockfd(butil::tcp_listen(butil::EndPoint(butil::IP_ANY, 0))); - if (sockfd < 0) { - return -1; - } - butil::EndPoint point; - if (butil::get_local_side(sockfd, &point) != 0) { - return -1; - } - return point.port; -} - -// Starts `server` on an ephemeral port and fills `options->internal_port` with -// another one. Both are released before Start() binds them and something else -// may take one in between, hence the retries. Returns 0 on success. -int StartWithInternalPort(brpc::Server* server, brpc::ServerOptions* options) { - for (int i = 0; i < 10; ++i) { - int internal_port = PickUnusedPort(); - if (internal_port < 0) { - continue; - } - options->internal_port = internal_port; - if (0 == server->Start("127.0.0.1:0", options)) { - return 0; - } - } - return -1; -} - TEST_F(ServerTest, ordinary_services_are_not_served_on_internal_port) { const struct { brpc::ProtocolType protocol; diff --git a/test/brpc_socket_unittest.cpp b/test/brpc_socket_unittest.cpp index 6e616cbc16..30e4e763ce 100644 --- a/test/brpc_socket_unittest.cpp +++ b/test/brpc_socket_unittest.cpp @@ -29,7 +29,10 @@ #include "butil/macros.h" #include "butil/fd_utility.h" #include "butil/debug/leak_annotations.h" +#include "butil/memory/scope_guard.h" +#include "butil/files/scoped_temp_dir.h" #include +#include "bthread/countdown_event.h" #include "bthread/unstable.h" #include "bthread/task_control.h" #include "brpc/socket.h" @@ -701,11 +704,11 @@ TEST_F(SocketTest, health_check) { brpc::SocketId id = 8888; butil::EndPoint point; ASSERT_NO_FATAL_FAILURE(PickUnusedEndPoint(&point)); - const int kCheckInteval = 1; + const int kCheckInterval = 1; brpc::SocketOptions options; options.remote_side = point; options.user = new CheckRecycle; - options.health_check_interval_s = kCheckInteval/*s*/; + options.health_check_interval_s = kCheckInterval/*s*/; ASSERT_EQ(0, brpc::Socket::Create(options, &id)); brpc::Socket* s = nullptr; { @@ -793,7 +796,7 @@ TEST_F(SocketTest, health_check) { while (brpc::Socket::Status(id, &nref) != 0) { bthread_usleep(1000); ASSERT_LT(butil::cpuwide_time_us(), - start_time + kCheckInteval * 1000000L + 100000L/*100ms*/); + start_time + kCheckInterval * 1000000L + 100000L/*100ms*/); } //ASSERT_EQ(2, nref); ASSERT_TRUE(global_sock); @@ -855,7 +858,7 @@ static void DoNothingOnEdgeTriggeredEvents(brpc::Socket*) {} // and get dispatched within that window. `OnInputEvent` must drop such a // stale event instead of crashing. TEST_F(SocketTest, input_event_on_revived_socket) { - const int kCheckInteval = 1; + const int kCheckInterval = 1; butil::EndPoint point; butil::fd_guard listening_fd; ASSERT_NO_FATAL_FAILURE(ListenOnFreePort(&point, &listening_fd)); @@ -863,7 +866,7 @@ TEST_F(SocketTest, input_event_on_revived_socket) { brpc::SocketId id = 8888; brpc::SocketOptions options; options.remote_side = point; - options.health_check_interval_s = kCheckInteval; + options.health_check_interval_s = kCheckInterval; // `OnInputEvent` returns early without an edge-triggered handler. options.on_edge_triggered_events = DoNothingOnEdgeTriggeredEvents; ASSERT_EQ(0, brpc::Socket::Create(options, &id)); @@ -896,7 +899,7 @@ TEST_F(SocketTest, input_event_on_revived_socket) { while (brpc::Socket::Status(id) != 0) { bthread_usleep(1000); ASSERT_LT(butil::cpuwide_time_us(), - start_time + kCheckInteval * 1000000L + 1000000L); + start_time + kCheckInterval * 1000000L + 1000000L); } ASSERT_EQ(-1, s->fd()); @@ -929,6 +932,189 @@ TEST_F(SocketTest, input_event_on_revived_socket) { ASSERT_EQ(0, brpc::Socket::SetFailed(id)); } +TEST_F(SocketTest, client_binding_endpoint_types) { + const char* inputs[] = {"127.0.0.1:12345", "[::1]:12345", + "unix:client-binding-test.sock"}; + const char* normalized[] = {"127.0.0.1:0", "[::1]:0", + "unix:client-binding-test.sock"}; + for (size_t i = 0; i < 3; ++i) { + SCOPED_TRACE(inputs[i]); + brpc::SocketOptions options; + // No connection is made: verify propagation without requiring a real + // network device or platform support for SO_BINDTODEVICE. + options.device_name = "test-device"; + ASSERT_EQ(0, butil::str2endpoint(inputs[i], &options.local_side)); + butil::EndPoint expected; + ASSERT_EQ(0, butil::str2endpoint(normalized[i], &expected)); + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); }; + brpc::SocketUniquePtr ptr; + ASSERT_EQ(0, brpc::Socket::Address(id, &ptr)); + ASSERT_EQ(expected, ptr->_bind_local_side); + ASSERT_EQ(options.device_name, ptr->_device_name); + // Normalization must not mutate a shared extended endpoint. + ASSERT_STREQ(inputs[i], butil::endpoint2str(options.local_side).c_str()); + + brpc::SocketUniquePtr short_socket; + ASSERT_EQ(0, ptr->GetShortSocket(&short_socket)); + BRPC_SCOPE_EXIT { short_socket->SetFailed(); }; + ASSERT_EQ(expected, short_socket->_bind_local_side); + ASSERT_EQ(options.device_name, short_socket->_device_name); + brpc::SocketUniquePtr pooled_socket; + ASSERT_EQ(0, ptr->GetPooledSocket(&pooled_socket)); + BRPC_SCOPE_EXIT { pooled_socket->SetFailed(); }; + ASSERT_EQ(expected, pooled_socket->_bind_local_side); + ASSERT_EQ(options.device_name, pooled_socket->_device_name); + + ASSERT_EQ(0, ptr->SetFailed()); + ptr->_is_hc_related_ref_held = true; + BRPC_SCOPE_EXIT { ptr->_is_hc_related_ref_held = false; }; + ASSERT_EQ(0, ptr->WaitAndReset(1)); + ASSERT_EQ(butil::EndPoint(), ptr->local_side()); + ASSERT_EQ(expected, ptr->_bind_local_side); + ASSERT_EQ(options.device_name, ptr->_device_name); + } +} + +TEST_F(SocketTest, client_binding_ipv6_connect) { + butil::EndPoint point; + ASSERT_EQ(0, butil::str2endpoint("[::1]:0", &point)); + butil::fd_guard listening_fd(butil::tcp_listen(point)); + ASSERT_GE(listening_fd, 0) << berror(); + ASSERT_EQ(0, butil::get_local_side(listening_fd, &point)); + brpc::SocketOptions options; + options.remote_side = point; + // This port is already occupied by the listener. Binding must use port 0. + options.local_side = point; + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); }; + brpc::SocketUniquePtr ptr; + ASSERT_EQ(0, brpc::Socket::Address(id, &ptr)); + for (int i = 0; i < 2; ++i) { + const timespec deadline = butil::milliseconds_from_now(1000); + butil::fd_guard fd(ptr->Connect(&deadline, nullptr, nullptr)); + ASSERT_GE(fd, 0) << berror(); + butil::EndPoint actual; + ASSERT_EQ(0, butil::get_local_side(fd, &actual)); + sockaddr_storage addr{}; + ASSERT_EQ(0, butil::endpoint2sockaddr(actual, &addr)); + ASSERT_EQ(AF_INET6, addr.ss_family); + const sockaddr_in6* in6 = reinterpret_cast(&addr); + ASSERT_TRUE(IN6_IS_ADDR_LOOPBACK(&in6->sin6_addr)); + ASSERT_NE(0, in6->sin6_port); + } +} + +TEST_F(SocketTest, client_binding_uds_connect) { + butil::ScopedTempDir dir; + ASSERT_TRUE(dir.CreateUniqueTempDir()); + // Deliberately use different path lengths for bind() and connect(). + const std::string server = "unix:" + dir.path().Append("server-long.sock").value(); + const std::string client = "unix:" + dir.path().Append("c.sock").value(); + brpc::SocketOptions options; + ASSERT_EQ(0, butil::str2endpoint(server.c_str(), &options.remote_side)); + ASSERT_EQ(0, butil::str2endpoint(client.c_str(), &options.local_side)); + butil::fd_guard listening_fd(butil::tcp_listen(options.remote_side)); + ASSERT_GE(listening_fd, 0) << berror(); + brpc::SocketId id; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + BRPC_SCOPE_EXIT { brpc::Socket::SetFailed(id); }; + brpc::SocketUniquePtr ptr; + ASSERT_EQ(0, brpc::Socket::Address(id, &ptr)); + const timespec deadline = butil::milliseconds_from_now(1000); + butil::fd_guard fd(ptr->Connect(&deadline, nullptr, nullptr)); + ASSERT_GE(fd, 0) << berror(); + butil::EndPoint actual; + ASSERT_EQ(0, butil::get_local_side(fd, &actual)); + ASSERT_EQ(AF_UNIX, butil::get_endpoint_type(actual)); + ASSERT_STREQ(client.c_str(), butil::endpoint2str(actual).c_str()); +} + +#if defined(OS_LINUX) +TEST_F(SocketTest, keep_client_bind_after_revive) { + int kCheckInterval = 1; + butil::EndPoint point; + butil::fd_guard listening_fd; + ASSERT_NO_FATAL_FAILURE(ListenOnFreePort(&point, &listening_fd)); + + butil::EndPoint bind_point; + // A distinct loopback source address (127.0.0.2) is used so that + // getsockname() can tell whether the explicit bind() actually happened: the + // kernel would pick 127.0.0.1 as the source for a 127.0.0.1 destination when + // no bind() is performed. This relies on the whole 127/8 being loopback, + // which is Linux specific. + ASSERT_EQ(0, str2endpoint("127.0.0.2:0", &bind_point)); + + brpc::SocketId id = 8888; + brpc::SocketOptions options; + options.remote_side = point; + // The explicitly configured client source address (ChannelOptions::client_host). + options.local_side = bind_point; + options.health_check_interval_s = kCheckInterval; + options.on_edge_triggered_events = DoNothingOnEdgeTriggeredEvents; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + + brpc::Socket* s = nullptr; + { + // `WaitAndReset' inside the health check waits until nref drops back + // to 2, thus no SocketUniquePtr may be held across `SetFailed'. + brpc::SocketUniquePtr ptr; + ASSERT_EQ(0, brpc::Socket::Address(id, &ptr)); + s = ptr.get(); + } + ASSERT_EQ(-1, s->fd()); + + // First connection: must bind to the configured source IP. + butil::IOBuf src; + src.append("hello"); + ASSERT_EQ(0, s->Write(&src)); + int64_t start_time = butil::cpuwide_time_us(); + while (s->fd() < 0) { + bthread_usleep(1000); + ASSERT_LT(butil::cpuwide_time_us(), start_time + 1000000L); + } + ASSERT_EQ(bind_point.ip, s->local_side().ip); + + // Fail the Socket and wait for the health check to revive it. This runs + // `WaitAndReset` which clears the runtime endpoint, not the binding config. + ASSERT_EQ(0, s->SetFailed()); + start_time = butil::cpuwide_time_us(); + while (brpc::Socket::Status(id) != 0) { + bthread_usleep(1000); + ASSERT_LT(butil::cpuwide_time_us(), + start_time + kCheckInterval * 1000000L + 1000000L); + } + + ASSERT_EQ(butil::EndPoint(), s->local_side()); + ASSERT_EQ(bind_point, s->_bind_local_side); + + // Reconnect on demand and verify the explicit source IP is still bound. + butil::EndPoint local_after_revive; + { + brpc::SocketUniquePtr ptr; + ASSERT_EQ(0, brpc::Socket::Address(id, &ptr)); + butil::IOBuf src2; + src2.append("world"); + ASSERT_EQ(0, ptr->Write(&src2)); + start_time = butil::cpuwide_time_us(); + while (ptr->fd() < 0) { + bthread_usleep(1000); + ASSERT_LT(butil::cpuwide_time_us(), start_time + 1000000L); + } + local_after_revive = ptr->local_side(); + } + ASSERT_EQ(bind_point.ip, local_after_revive.ip); + + s->ReleaseHCRelatedReference(); + // Must close the listening fd before SetFailed, otherwise the health + // check still has chance to get reconnected and revive the id. + listening_fd.reset(-1); + ASSERT_EQ(0, brpc::Socket::SetFailed(id)); +} +#endif // OS_LINUX + void* Writer(void* void_arg) { WriterArg* arg = static_cast(void_arg); brpc::SocketUniquePtr sock; @@ -1812,6 +1998,7 @@ TEST_F(SocketTest, notify_on_success) { struct ShutdownWriterArg { size_t times; brpc::SocketId socket_id; + bthread::CountdownEvent* done; butil::atomic total_count; butil::atomic success_count; }; @@ -1826,6 +2013,7 @@ int HandleSocketShutdownWrite(bthread_id_t id, void* data, int error_code, ++arg->success_count; } CHECK_EQ(0, bthread_id_unlock_and_destroy(id)); + arg->done->signal(); return 0; } @@ -1887,10 +2075,13 @@ void TestShutdownWrite() { pthread_create(&rth, nullptr, reader, &reader_arg); bthread_t th[3]; + const size_t expected_count = REP * ARRAY_SIZE(th); + bthread::CountdownEvent done(expected_count); ShutdownWriterArg args[ARRAY_SIZE(th)]; for (size_t i = 0; i < ARRAY_SIZE(th); ++i) { args[i].times = REP; args[i].socket_id = id; + args[i].done = &done; args[i].total_count = 0; args[i].success_count = 0; bthread_start_background(&th[i], nullptr, ShutdownWriter, &args[i]); @@ -1899,11 +2090,15 @@ void TestShutdownWrite() { for (size_t i = 0; i < ARRAY_SIZE(th); ++i) { ASSERT_EQ(0, bthread_join(th[i], nullptr)); } - bthread_usleep(50 * 1000); - - ASSERT_TRUE(s->IsWriteShutdown()); - ASSERT_FALSE(s->Failed()); - ASSERT_EQ(0, s->SetFailed()); + const int wait_result = done.timed_wait(butil::seconds_from_now(5)); + if (wait_result == 0) { + EXPECT_TRUE(s->IsWriteShutdown()); + EXPECT_FALSE(s->Failed()); + EXPECT_EQ(0, s->SetFailed()); + } else { + s->SetFailed(); + done.wait(); + } s.release()->Dereference(); pthread_join(rth, nullptr); ASSERT_EQ((brpc::Socket*)nullptr, global_sock); @@ -1911,11 +2106,12 @@ void TestShutdownWrite() { size_t total_count = 0; size_t success_count = 0; - for (auto & arg : args) { + for (const auto& arg : args) { total_count += arg.total_count; success_count += arg.success_count; } - ASSERT_EQ(REP * ARRAY_SIZE(th), total_count); + ASSERT_EQ(0, wait_result); + ASSERT_EQ(expected_count, total_count); EXPECT_EQ((size_t)1, reader_arg.nread); EXPECT_EQ((size_t)1, success_count); } @@ -1923,6 +2119,9 @@ void TestShutdownWrite() { TEST_F(SocketTest, shutdown_write) { for (int i = 0; i < 100; ++i) { TestShutdownWrite(); + if (::testing::Test::HasFailure()) { + return; + } } } diff --git a/test/brpc_transport_handshake_unittest.cpp b/test/brpc_transport_handshake_unittest.cpp index 6562ba90ab..2ed34445a1 100644 --- a/test/brpc_transport_handshake_unittest.cpp +++ b/test/brpc_transport_handshake_unittest.cpp @@ -105,32 +105,116 @@ static FrameSpec FixedSpec(const char* magic, size_t magic_len, FrameSpec::FIXED); } -static HandshakeCodec MakeTestCodec(std::string* calls = NULL) { - HandshakeCodec codec{}; - codec.protocol_version = 7; - codec.hello_frame = FixedSpec("HS", 2, 4); - codec.ack_frame = FixedSpec(NULL, 0, 1); - codec.build_hello = [calls](bool enabled, std::string* payload) { - if (calls) *calls += "build "; +class TestHandshakeProtocol : public HandshakeProtocol { +public: + explicit TestHandshakeProtocol(std::string* calls = NULL) + : _calls(calls), _with_extension(false), + _extension_size(0) {} + + int ProtocolVersion() const override { return 7; } + const FrameSpec& HelloFrameSpec() const override { + static const FrameSpec spec = FixedSpec("HS", 2, 4); + return spec; + } + const FrameSpec& AckFrameSpec() const override { + static const FrameSpec spec = FixedSpec(NULL, 0, 1); + return spec; + } + StepResult BuildHello(bool enabled, std::string* payload) override { + Append("build "); *payload = enabled ? "LO" : "NO"; return STEP_OK; - }; - codec.parse_hello = [calls](const std::string& payload) { - if (calls) *calls += "parse "; + } + StepResult ParseHello(const std::string& payload) override { + Append("parse "); return payload == "OK" ? STEP_OK : STEP_FALLBACK; - }; - codec.build_ack = [calls](bool enabled, std::string* payload) { - if (calls) *calls += enabled ? "ack1 " : "ack0 "; + } + StepResult BuildAck(bool enabled, std::string* payload) override { + Append(enabled ? "ack1 " : "ack0 "); *payload = enabled ? "1" : "0"; return STEP_OK; - }; - codec.parse_ack = [calls](const std::string& payload, bool* enabled) { - if (calls) *calls += "parse_ack "; + } + StepResult ParseAck( + const std::string& payload, bool* enabled) override { + Append("parse_ack "); *enabled = payload == "1"; return payload == "1" || payload == "0" ? STEP_OK : STEP_ERROR; - }; - return codec; -} + } + bool HasExtension() const override { return _with_extension; } + const FrameSpec& ExtensionFrameSpec() const override { + return _extension_spec; + } + StepResult BuildExtension( + bool enabled, std::string* payload) override { + *payload = enabled + ? std::string("E1", _extension_size) + : std::string("E0", _extension_size); + return STEP_OK; + } + StepResult ParseExtension(const std::string& payload) override { + return payload == std::string("E1", _extension_size) + ? STEP_OK : STEP_FALLBACK; + } + void EnableExtension(size_t size) { + _with_extension = true; + _extension_size = size; + _extension_spec = FixedSpec(NULL, 0, size); + } + +private: + void Append(const char* value) { + if (_calls != NULL) { + *_calls += value; + } + } + + std::string* _calls; + bool _with_extension; + size_t _extension_size; + FrameSpec _extension_spec; +}; + +class TestHandshakeTransport : public HandshakeTransport { +public: + explicit TestHandshakeTransport(std::string* calls = NULL) + : prepare_result(STEP_OK), negotiate_result(STEP_OK), + validate_result(STEP_OK), tcp_active(false), + high_speed_active(false), failed(false), _calls(calls) {} + + StepResult PrepareResources() override { + Append("prepare "); + return prepare_result; + } + StepResult NegotiateResources() override { + Append("negotiate "); + return negotiate_result; + } + void OnEstablished() override { + high_speed_active = true; + Append("activate"); + } + void OnFallback() override { tcp_active = true; } + void OnFailed() override { failed = true; } + StepResult ValidateEstablished() override { + Append("validate "); + return validate_result; + } + + StepResult prepare_result; + StepResult negotiate_result; + StepResult validate_result; + bool tcp_active; + bool high_speed_active; + bool failed; + +private: + void Append(const char* value) { + if (_calls != NULL) { + *_calls += value; + } + } + std::string* _calls; +}; static std::string MakeUBShmHello() { std::string frame(64, '\0'); @@ -298,43 +382,68 @@ TEST(TransportHandshakeTest, client_runs_codec_and_resource_sequence) { session.SetIOForTest(&io); std::string calls; - ClientHandshakeCallbacks callbacks{}; - callbacks.codec = MakeTestCodec(&calls); - callbacks.transport.prepare_resources = [&]() { - calls += "prepare "; - return STEP_OK; - }; - callbacks.transport.negotiate_resources = [&]() { - calls += "negotiate "; - return STEP_OK; - }; - callbacks.transport.set_high_speed_active = [&]() { calls += "activate"; }; - callbacks.transport.set_tcp_active = []() {}; - callbacks.transport.on_failed = []() {}; + TestHandshakeProtocol protocol(&calls); + TestHandshakeTransport transport(&calls); - ASSERT_EQ(STEP_OK, session.RunClient(callbacks)); + ASSERT_EQ(STEP_OK, session.RunClient(&protocol, &transport)); ASSERT_EQ("prepare build parse negotiate ack1 activate", calls); ASSERT_EQ("HSLO1", io.output()); ASSERT_EQ(ESTABLISHED, session.phase()); ASSERT_EQ(7, session.protocol_version()); } +TEST(TransportHandshakeTest, client_exchanges_extension_before_ack) { + MemoryHandshakeIO io("HSOKE"); + HandshakeSession session; + session.SetIOForTest(&io); + TestHandshakeProtocol protocol; + protocol.EnableExtension(1); + TestHandshakeTransport transport; + + ASSERT_EQ(STEP_OK, session.RunClient(&protocol, &transport)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_EQ(ESTABLISHED, session.phase()); +} + +TEST(TransportHandshakeTest, server_resumes_fragmented_extension_then_ack) { + MemoryHandshakeIO io; + HandshakeSession session; + session.SetIOForTest(&io); + butil::IOBuf source; + source.append("HSOK", 4); + IOBufHandshakeInput input(&source); + TestHandshakeProtocol protocol; + protocol.EnableExtension(2); + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ(EXTENSION_WAIT, session.phase()); + EXPECT_EQ("HSLO", io.output()); + source.append("E", 1); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ(EXTENSION_WAIT, session.phase()); + source.append("11", 2); // Remainder of extension, then ACK. + ASSERT_EQ(STEP_OK, + session.RunServer(protocols, &input, &transport, false)); + EXPECT_EQ("HSLOE1", io.output()); + EXPECT_TRUE(source.empty()); + EXPECT_EQ(ESTABLISHED, session.phase()); +} + TEST(TransportHandshakeTest, client_resource_failure_falls_back_before_io) { MemoryHandshakeIO io; HandshakeSession session; session.SetIOForTest(&io); - bool tcp_active = false; - - ClientHandshakeCallbacks callbacks{}; - callbacks.codec = MakeTestCodec(); - callbacks.transport.prepare_resources = []() { return STEP_FALLBACK; }; - callbacks.transport.negotiate_resources = []() { return STEP_ERROR; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = [&]() { tcp_active = true; }; - callbacks.transport.on_failed = []() {}; - - ASSERT_EQ(STEP_FALLBACK, session.RunClient(callbacks)); - ASSERT_TRUE(tcp_active); + TestHandshakeProtocol protocol; + TestHandshakeTransport transport; + transport.prepare_result = STEP_FALLBACK; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, session.RunClient(&protocol, &transport)); + ASSERT_TRUE(transport.tcp_active); ASSERT_TRUE(io.output().empty()); ASSERT_EQ(FALLBACK_TCP, session.phase()); } @@ -348,33 +457,19 @@ TEST(TransportHandshakeTest, server_resumes_at_buffered_ack) { IOBufHandshakeInput input(&source); std::string calls; - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.codecs.push_back(MakeTestCodec(&calls)); - callbacks.input = &input; - callbacks.transport.prepare_resources = [&]() { - calls += "prepare "; - return STEP_OK; - }; - callbacks.transport.negotiate_resources = [&]() { - calls += "negotiate "; - return STEP_OK; - }; - callbacks.validate_established = [&]() { - calls += "validate "; - return STEP_OK; - }; - callbacks.transport.set_high_speed_active = [&]() { calls += "activate"; }; - callbacks.transport.set_tcp_active = []() {}; - callbacks.transport.on_failed = []() {}; + TestHandshakeProtocol protocol(&calls); + std::vector protocols(1, &protocol); + TestHandshakeTransport transport(&calls); - ASSERT_EQ(STEP_NEED_MORE, session.RunServer(callbacks)); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); ASSERT_EQ(ACK_WAIT, session.phase()); ASSERT_EQ("HSLO", io.output()); ASSERT_TRUE(source.empty()); source.append("1", 1); - ASSERT_EQ(STEP_OK, session.RunServer(callbacks)); + ASSERT_EQ(STEP_OK, + session.RunServer(protocols, &input, &transport, false)); ASSERT_EQ("parse prepare negotiate build parse_ack validate activate", calls); ASSERT_EQ(ESTABLISHED, session.phase()); @@ -387,22 +482,17 @@ TEST(TransportHandshakeTest, server_resource_failure_falls_back_after_ack) { butil::IOBuf source; source.append("HSOK0", 5); IOBufHandshakeInput input(&source); - bool tcp_active = false; - - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.codecs.push_back(MakeTestCodec()); - callbacks.input = &input; - callbacks.transport.prepare_resources = []() { return STEP_FALLBACK; }; - callbacks.transport.negotiate_resources = []() { return STEP_ERROR; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = [&]() { tcp_active = true; }; - callbacks.transport.on_failed = []() {}; - - ASSERT_EQ(STEP_FALLBACK, session.RunServer(callbacks)); + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + transport.prepare_result = STEP_FALLBACK; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &input, &transport, false)); ASSERT_EQ("HSNO", io.output()); ASSERT_TRUE(source.empty()); - ASSERT_TRUE(tcp_active); + ASSERT_TRUE(transport.tcp_active); ASSERT_EQ(FALLBACK_TCP, session.phase()); } @@ -413,24 +503,20 @@ TEST(TransportHandshakeTest, server_falls_back_without_consuming_other_magic) { butil::IOBuf source; source.append("XXok", 4); IOBufHandshakeInput input(&source); - bool tcp_active = false; - - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = true; - callbacks.codecs.push_back(MakeTestCodec()); - callbacks.input = &input; - callbacks.transport.prepare_resources = []() { return STEP_ERROR; }; - callbacks.transport.negotiate_resources = []() { return STEP_ERROR; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = [&]() { tcp_active = true; }; - callbacks.transport.on_failed = []() {}; - - ASSERT_EQ(STEP_FALLBACK, session.RunServer(callbacks)); - ASSERT_TRUE(tcp_active); + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; + transport.prepare_result = STEP_ERROR; + transport.negotiate_result = STEP_ERROR; + + ASSERT_EQ(STEP_FALLBACK, + session.RunServer(protocols, &input, &transport, true)); + ASSERT_TRUE(transport.tcp_active); ASSERT_EQ(4UL, source.size()); ASSERT_EQ(FALLBACK_TCP, session.phase()); - ASSERT_EQ(STEP_NOT_MINE, session.RunServer(callbacks)); + ASSERT_EQ(STEP_NOT_MINE, + session.RunServer(protocols, &input, &transport, true)); ASSERT_EQ(4UL, source.size()); ASSERT_EQ(FALLBACK_TCP, session.phase()); } @@ -442,22 +528,18 @@ TEST(TransportHandshakeTest, server_enters_hello_phase_after_magic_matches) { butil::IOBuf source; IOBufHandshakeInput input(&source); - ServerHandshakeCallbacks callbacks{}; - callbacks.fallback_on_not_mine = false; - callbacks.codecs.push_back(MakeTestCodec()); - callbacks.input = &input; - callbacks.transport.prepare_resources = []() { return STEP_OK; }; - callbacks.transport.negotiate_resources = []() { return STEP_OK; }; - callbacks.transport.set_high_speed_active = []() {}; - callbacks.transport.set_tcp_active = []() {}; - callbacks.transport.on_failed = []() {}; + TestHandshakeProtocol protocol; + std::vector protocols(1, &protocol); + TestHandshakeTransport transport; source.append("H", 1); - ASSERT_EQ(STEP_NEED_MORE, session.RunServer(callbacks)); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); ASSERT_EQ(UNINITIALIZED, session.phase()); source.append("S", 1); - ASSERT_EQ(STEP_NEED_MORE, session.RunServer(callbacks)); + ASSERT_EQ(STEP_NEED_MORE, + session.RunServer(protocols, &input, &transport, false)); ASSERT_EQ(HELLO_WAIT, session.phase()); ASSERT_EQ(7, session.protocol_version()); ASSERT_EQ(2UL, source.size()); diff --git a/test/brpc_ubring_unittest.cpp b/test/brpc_ubring_unittest.cpp index f4b416303d..02ad6ecd6a 100644 --- a/test/brpc_ubring_unittest.cpp +++ b/test/brpc_ubring_unittest.cpp @@ -65,6 +65,70 @@ class HelloMessageTest : public ::testing::Test { std::string buffer; }; +TEST(HelloFormatExtensionTest, serialize_deserialize_roundtrip) { + brpc::ubring::HelloFormatExtension extension = { + brpc::ubring::HelloFormatExtension::WIRE_SIZE, + brpc::ubring::UBR_DATA_FORMAT_LEGACY_64}; + char buffer[brpc::ubring::HelloFormatExtension::WIRE_SIZE] = {}; + + extension.Serialize(buffer); + + brpc::ubring::HelloFormatExtension decoded = {}; + decoded.Deserialize(buffer); + EXPECT_EQ(extension.extension_len, decoded.extension_len); + EXPECT_EQ(extension.format_id, decoded.format_id); +} + +TEST(HelloFormatExtensionTest, serialize_uses_network_byte_order) { + brpc::ubring::HelloFormatExtension extension = {0x0102, 0x0304}; + char buffer[brpc::ubring::HelloFormatExtension::WIRE_SIZE] = {}; + const unsigned char expected[] = {0x01, 0x02, 0x03, 0x04}; + + extension.Serialize(buffer); + + EXPECT_EQ(0, memcmp(expected, buffer, sizeof(expected))); +} + +TEST(HelloFormatExtensionTest, deserialize_none_format) { + const unsigned char buffer[] = {0x00, 0x04, 0x00, 0x00}; + brpc::ubring::HelloFormatExtension extension = {}; + + extension.Deserialize(buffer); + + EXPECT_EQ(4, extension.extension_len); + EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, extension.format_id); +} + +TEST(HelloFormatExtensionTest, deserialize_unknown_format) { + const unsigned char buffer[] = {0x00, 0x04, 0x12, 0x34}; + brpc::ubring::HelloFormatExtension extension = {}; + + extension.Deserialize(buffer); + + EXPECT_EQ(4, extension.extension_len); + EXPECT_EQ(0x1234, extension.format_id); +} + +TEST(UBShmHandshakeAdapterTest, rejects_unsupported_format_extension) { + brpc::ubring::UBShmHandshakeAdapter adapter; + std::string payload; + ASSERT_EQ(brpc::handshake::STEP_OK, + adapter.BuildExtension(true, &payload)); + EXPECT_EQ(std::string("\0\4\0\1", 4), payload); + EXPECT_EQ(brpc::handshake::STEP_OK, + adapter.ParseExtension(payload)); + + payload[3] = 2; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); + payload[3] = 0; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); + payload[1] = 3; + EXPECT_EQ(brpc::handshake::STEP_FALLBACK, + adapter.ParseExtension(payload)); +} + TEST_F(HelloMessageTest, serialize_deserialize_roundtrip) { msg.msg_len = 64; msg.hello_ver = 2; @@ -146,7 +210,7 @@ TEST_F(HelloMessageTest, toString_contains_fields) { EXPECT_NE(std::string::npos, s.find("UBRING_test")); } -TEST(UBShmHandshakeAdapterTest, codec_preserves_v2_wire_format) { +TEST(UBShmHandshakeAdapterTest, codec_uses_v3_wire_format) { brpc::ubring::UBShmHandshakeAdapter adapter; char shm_name[SHM_MAX_NAME_BUFF_LEN] = {0}; memcpy(shm_name, "UBRING_test_C", 14); @@ -159,17 +223,16 @@ TEST(UBShmHandshakeAdapterTest, codec_preserves_v2_wire_format) { ASSERT_EQ(brpc::handshake::STEP_OK, adapter.ParseHello(payload, &decoded)); EXPECT_EQ(64, decoded.msg_len); - EXPECT_EQ(2, decoded.hello_ver); + EXPECT_EQ(3, decoded.hello_ver); EXPECT_EQ(1, decoded.impl_ver); EXPECT_EQ(4 * 1024 * 1024, decoded.len); EXPECT_EQ(0, memcmp(shm_name, decoded.shm_name, SHM_MAX_NAME_BUFF_LEN)); std::string frame; - const brpc::handshake::HandshakeCodec codec = adapter.MakeCodec(); ASSERT_EQ(brpc::handshake::FRAME_OK, brpc::handshake::FrameCodec::Encode( - codec.hello_frame, payload, &frame)); + adapter.HelloFrameSpec(), payload, &frame)); ASSERT_EQ(64, frame.size()); EXPECT_EQ("UB", frame.substr(0, 2)); } @@ -197,7 +260,7 @@ TEST(UBShmHandshakeAdapterTest, short_name_is_zero_padded) { TEST(UBShmHandshakeAdapterTest, rejects_unterminated_remote_name) { brpc::ubring::HelloMessage message{}; message.msg_len = 64; - message.hello_ver = 2; + message.hello_ver = 3; message.impl_ver = 1; message.len = 4096; memset(message.shm_name, 'A', SHM_MAX_NAME_BUFF_LEN); @@ -306,6 +369,17 @@ using brpc::ubring::UBShmEndpointTest; TEST_F(UBShmEndpointTest, construct_initial_state) { ASSERT_NE(nullptr, _ep); + EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, + _ep->negotiated_data_format()); +} + +TEST_F(UBShmEndpointTest, reset_clears_negotiated_data_format) { + _ep->SetNegotiatedDataFormat(brpc::ubring::UBR_DATA_FORMAT_LEGACY_64); + + _ep->Reset(); + + EXPECT_EQ(brpc::ubring::UBR_DATA_FORMAT_NONE, + _ep->negotiated_data_format()); } TEST_F(UBShmEndpointTest, allocate_client_resources_real_shm) { diff --git a/test/brpc_uri_unittest.cpp b/test/brpc_uri_unittest.cpp index b9d6b6508e..b1be2c55b3 100644 --- a/test/brpc_uri_unittest.cpp +++ b/test/brpc_uri_unittest.cpp @@ -15,10 +15,15 @@ // specific language governing permissions and limitations // under the License. +#include #include #include "brpc/uri.h" +namespace brpc { +DECLARE_uint32(http_max_query_count); +} + TEST(URITest, everything) { brpc::URI uri; std::string uri_str = " foobar://user:passwd@www.baidu.com:80/s?wd=uri#frag "; @@ -347,6 +352,33 @@ TEST(URITest, invalid_query) { ASSERT_EQ("a-b-c:def", uri.query()); } +TEST(URITest, too_many_queries) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_http_max_query_count = 4; + + brpc::URI uri; + ASSERT_EQ(0, uri.SetHttpURL("http://a.com/s?a=1&b=2&c=3&d=4")) << uri.status(); + ASSERT_EQ(-1, uri.SetHttpURL("http://a.com/s?a=1&b=2&c=3&d=4&e=5")); + ASSERT_STREQ("More than 4 query parameters in url", uri.status().error_cstr()); + // Repeated keys collapse into one map entry, but the splitter still walks + // every segment, so they count. + ASSERT_EQ(-1, uri.SetHttpURL("http://a.com/s?a=1&a=2&a=3&a=4&a=5")); + // An empty query is not one parameter. + brpc::FLAGS_http_max_query_count = 1; + ASSERT_EQ(0, uri.SetHttpURL("http://a.com/s?")) << uri.status(); + + brpc::FLAGS_http_max_query_count = 4; + ASSERT_EQ(0, uri.SetH2Path("/s?a=1&b=2&c=3&d=4")) << uri.status(); + ASSERT_EQ(-1, uri.SetH2Path("/s?a=1&b=2&c=3&d=4&e=5")); + ASSERT_STREQ("More than 4 query parameters in :path", uri.status().error_cstr()); + // The next path clears the failure rather than inheriting it. + ASSERT_EQ(0, uri.SetH2Path("/s?a=1")) << uri.status(); + + brpc::FLAGS_http_max_query_count = 0; + ASSERT_EQ(0, uri.SetHttpURL("http://a.com/s?a=1&b=2&c=3&d=4&e=5")) << uri.status(); + ASSERT_EQ(0, uri.SetH2Path("/s?a=1&b=2&c=3&d=4&e=5")) << uri.status(); +} + TEST(URITest, high_bit_bytes) { // Bytes >= 0x80 (e.g. UTF-8 in the host/path) index the +128-biased // action table. On unsigned-char platforms they would read past the diff --git a/test/brpc_urma_unittest.cpp b/test/brpc_urma_unittest.cpp new file mode 100644 index 0000000000..d27367a504 --- /dev/null +++ b/test/brpc_urma_unittest.cpp @@ -0,0 +1,612 @@ +// 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 + +#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) { + GFLAGS_NAMESPACE::FlagSaver flag_saver; + + urma::FLAGS_urma_client_handshake_version = 2; + urma::UrmaHandshake* hs2 = urma::CreateClientHandshake(nullptr); + if (hs2) { + EXPECT_EQ(2, hs2->ProtocolVersion()); + delete hs2; + } + + urma::FLAGS_urma_client_handshake_version = 3; + urma::UrmaHandshake* hs3 = urma::CreateClientHandshake(nullptr); + if (hs3) { + EXPECT_EQ(3, hs3->ProtocolVersion()); + delete hs3; + } +} + +// --------------------------------------------------------------------------- +// 4-byte ACK: HELLO_ACK_URMA_OK bit. +// --------------------------------------------------------------------------- +TEST(UrmaHandshakeTest, ack_bit_is_urma_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{}; + 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)); + + hello.jetty_id = 1; + hello.seg_va = std::numeric_limits::max() - 7; + hello.seg_len = 8; + 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(); +} diff --git a/test/bthread_butex_unittest.cpp b/test/bthread_butex_unittest.cpp index 730eac7b4e..a92f447d4b 100644 --- a/test/bthread_butex_unittest.cpp +++ b/test/bthread_butex_unittest.cpp @@ -45,25 +45,37 @@ TEST(ButexTest, wait_on_already_timedout_butex) { ASSERT_EQ(ETIMEDOUT, errno); } +struct JoinSleepArg { + bthread_t tid = 0; + uint64_t sleep_us = 0; + butil::atomic finished{false}; +}; + void* sleeper(void* arg) { bthread_usleep((uint64_t)arg); return nullptr; } +void* join_sleeper(void* arg) { + JoinSleepArg* a = static_cast(arg); + butil::Timer tm; + tm.start(); + EXPECT_EQ(0, bthread_usleep(a->sleep_us)); + tm.stop(); + // Timer::u_elapsed() floors nanoseconds to microseconds. + EXPECT_GE(tm.u_elapsed() + 1, static_cast(a->sleep_us)); + a->finished.store(true, butil::memory_order_release); + return nullptr; +} + void* joiner(void* arg) { - const long t1 = butil::gettimeofday_us(); - for (bthread_t* th = (bthread_t*)arg; *th; ++th) { - if (0 != bthread_join(*th, nullptr)) { - LOG(FATAL) << "fail to join thread_" << th - (bthread_t*)arg; - } - long elp = butil::gettimeofday_us() - t1; - EXPECT_LE(labs(elp - (th - (bthread_t*)arg + 1) * 100000L), 15000L) - << "timeout when joining thread_" << th - (bthread_t*)arg; - LOG(INFO) << "Joined thread " << *th << " at " << elp << "us [" - << bthread_self() << "]"; + JoinSleepArg* args = static_cast(arg); + for (JoinSleepArg* a = args; a->tid; ++a) { + EXPECT_EQ(0, bthread_join(a->tid, nullptr)); + EXPECT_TRUE(a->finished.load(butil::memory_order_acquire)); } - for (bthread_t* th = (bthread_t*)arg; *th; ++th) { - EXPECT_EQ(0, bthread_join(*th, nullptr)); + for (JoinSleepArg* a = args; a->tid; ++a) { + EXPECT_EQ(0, bthread_join(a->tid, nullptr)); } return nullptr; } @@ -86,21 +98,20 @@ TEST(ButexTest, with_or_without_array_zero) { TEST(ButexTest, join) { const size_t N = 6; const size_t M = 6; - bthread_t th[N+1]; + JoinSleepArg args[N+1]; bthread_t jth[M]; pthread_t pth[M]; for (size_t i = 0; i < N; ++i) { bthread_attr_t attr = (i == 0 ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL); + args[i].sleep_us = 100000L/*100ms*/ * (i + 1); ASSERT_EQ(0, bthread_start_urgent( - &th[i], &attr, sleeper, - (void*)(100000L/*100ms*/ * (i + 1)))); + &args[i].tid, &attr, join_sleeper, &args[i])); } - th[N] = 0; // joiner will join tids in `th' until seeing 0. for (size_t i = 0; i < M; ++i) { - ASSERT_EQ(0, bthread_start_urgent(&jth[i], nullptr, joiner, th)); + ASSERT_EQ(0, bthread_start_urgent(&jth[i], nullptr, joiner, args)); } for (size_t i = 0; i < M; ++i) { - ASSERT_EQ(0, pthread_create(&pth[i], nullptr, joiner, th)); + ASSERT_EQ(0, pthread_create(&pth[i], nullptr, joiner, args)); } for (size_t i = 0; i < M; ++i) { @@ -257,7 +268,6 @@ TEST(ButexTest, stop_after_running) { TEST(ButexTest, stop_before_running) { int* butex = bthread::butex_create_checked(); *butex = 7; - butil::Timer tm; const long WAIT_MSEC = 500; for (int i = 0; i < 2; ++i) { @@ -265,17 +275,11 @@ TEST(ButexTest, stop_before_running) { (i == 0 ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL) | BTHREAD_NOSIGNAL; bthread_t th; ButexWaitArg arg = { butex, *butex, WAIT_MSEC, EINTR }; - - tm.start(); + ASSERT_EQ(0, bthread_start_background(&th, &attr, wait_butex, &arg)); ASSERT_EQ(0, bthread_stop(th)); bthread_flush(); ASSERT_EQ(0, bthread_join(th, nullptr)); - tm.stop(); - - ASSERT_LT(tm.m_elapsed(), 5); - // ASSERT_TRUE(bthread::get_task_control()-> - // timer_thread()._idset.empty()); ASSERT_EQ(EINVAL, bthread_stop(th)); } bthread::butex_destroy(butex); diff --git a/test/bthread_context_unittest.cpp b/test/bthread_context_unittest.cpp new file mode 100644 index 0000000000..ff4a2f6917 --- /dev/null +++ b/test/bthread_context_unittest.cpp @@ -0,0 +1,112 @@ +// 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 "bthread/context.h" + +#if defined(BTHREAD_CONTEXT_PLATFORM_linux_x86_64) + +namespace { + +constexpr size_t kStackSize = 8192; +constexpr size_t kXmmRegisterCount = 10; + +struct XmmState { + uint8_t registers[kXmmRegisterCount][16]; +}; + +bthread_fcontext_t g_main_context; +bthread_fcontext_t g_test_context; +XmmState g_clobber_state; + +extern "C" intptr_t jump_with_xmm_state( + bthread_fcontext_t* old_context, bthread_fcontext_t new_context, + intptr_t value, const XmmState* input, XmmState* output); + +// Keep the register setup, context switch, and register capture in one +// assembly function so the compiler cannot use an XMM register between them. +__asm__( +".text\n" +".p2align 4,,15\n" +".type jump_with_xmm_state,@function\n" +"jump_with_xmm_state:\n" +" movups 0x00(%rcx), %xmm6\n" +" movups 0x10(%rcx), %xmm7\n" +" movups 0x20(%rcx), %xmm8\n" +" movups 0x30(%rcx), %xmm9\n" +" movups 0x40(%rcx), %xmm10\n" +" movups 0x50(%rcx), %xmm11\n" +" movups 0x60(%rcx), %xmm12\n" +" movups 0x70(%rcx), %xmm13\n" +" movups 0x80(%rcx), %xmm14\n" +" movups 0x90(%rcx), %xmm15\n" +" pushq %r8\n" +" xorl %ecx, %ecx\n" +" call bthread_jump_fcontext\n" +" popq %r8\n" +" movups %xmm6, 0x00(%r8)\n" +" movups %xmm7, 0x10(%r8)\n" +" movups %xmm8, 0x20(%r8)\n" +" movups %xmm9, 0x30(%r8)\n" +" movups %xmm10, 0x40(%r8)\n" +" movups %xmm11, 0x50(%r8)\n" +" movups %xmm12, 0x60(%r8)\n" +" movups %xmm13, 0x70(%r8)\n" +" movups %xmm14, 0x80(%r8)\n" +" movups %xmm15, 0x90(%r8)\n" +" ret\n" +".size jump_with_xmm_state,.-jump_with_xmm_state\n"); + +void overwrite_xmm_and_return(intptr_t) { + XmmState ignored = {}; + jump_with_xmm_state(&g_test_context, g_main_context, 0, + &g_clobber_state, &ignored); +} + +TEST(BthreadContextTest, preserves_xmm6_through_xmm15) { + XmmState expected; + XmmState actual = {}; + for (size_t reg = 0; reg < kXmmRegisterCount; ++reg) { + for (size_t byte = 0; byte < 16; ++byte) { + expected.registers[reg][byte] = + static_cast(reg * 16 + byte); + g_clobber_state.registers[reg][byte] = + static_cast(0xff - reg * 16 - byte); + } + } + + void* stack = std::malloc(kStackSize); + ASSERT_NE(nullptr, stack); + g_main_context = nullptr; + g_test_context = bthread_make_fcontext( + static_cast(stack) + kStackSize, kStackSize, + overwrite_xmm_and_return); + + jump_with_xmm_state(&g_main_context, g_test_context, 0, &expected, &actual); + + EXPECT_EQ(0, std::memcmp(&expected, &actual, sizeof(expected))); + std::free(stack); +} + +} // namespace + +#endif // BTHREAD_CONTEXT_PLATFORM_linux_x86_64 diff --git a/test/bthread_parking_lot_unittest.cpp b/test/bthread_parking_lot_unittest.cpp new file mode 100644 index 0000000000..d14902a968 --- /dev/null +++ b/test/bthread_parking_lot_unittest.cpp @@ -0,0 +1,473 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#include +#include +#include +#include +#include +#include +#include "butil/time.h" +#include "bthread/parking_lot.h" + +namespace { + +using bthread::ParkingLot; + +const int64_t kTimeoutUs = 5 * 1000 * 1000; + +template +bool wait_until(Predicate predicate) { + const int64_t deadline_us = butil::cpuwide_time_us() + kTimeoutUs; + int spins = 0; + while (!predicate()) { + if (butil::cpuwide_time_us() >= deadline_us) { + return false; + } + // Keep short handshakes responsive so the signaler can race with waiter + // registration instead of usually arriving after the waiter has slept. + if (++spins < 256) { + std::this_thread::yield(); + } else { + usleep(50); + } + } + return true; +} + +// ParkingLot blocks OS threads, so use native threads rather than bthreads. +// All assertions are made before cleanup: rescue wakeups must not make a lost +// wakeup test pass. The guard also cleans up on a fatal assertion/early return. +class ParkingLotThreads { +public: + explicit ParkingLotThreads(ParkingLot* lot) + : _lot(lot), _cancelled(false), _finished(0) {} + + ~ParkingLotThreads() { + _cancelled.store(true, butil::memory_order_release); + const bool joined = wait_until([this] { + if (_finished.load(butil::memory_order_acquire) == _threads.size()) { + return true; + } + // Do not depend on signal() or stop() working to rescue a failure. + // Changing the value also rescues a thread not yet in futex_wait. + _lot->_pending_signal.fetch_add(2, butil::memory_order_release); + bthread::futex_wake_private(&_lot->_pending_signal, INT_MAX); + return false; + }); + if (!joined) { + ADD_FAILURE() << "ParkingLot test threads could not be rescued"; + // Never detach threads referencing stack objects or hang the suite. + std::abort(); + } + for (auto& thread : _threads) { + thread.join(); + } + } + + template + void start(Function function) { + _threads.emplace_back([this, function] { + function(); + _finished.fetch_add(1, butil::memory_order_release); + }); + } + + bool cancelled() const { + return _cancelled.load(butil::memory_order_acquire); + } + + bool finished() const { + return _finished.load(butil::memory_order_acquire) == _threads.size(); + } + +private: + ParkingLot* _lot; + butil::atomic _cancelled; + butil::atomic _finished; + std::vector _threads; +}; + +class ParkingLotTest : public ::testing::TestWithParam { +protected: + void SetUp() override { + bthread::FLAGS_parking_lot_no_signal_when_no_waiter = GetParam(); + } + + // Test builds use -fno-access-control. Inspect state and waiter bookkeeping + // without adding test-only hooks to the production synchronization path. + void check_wake_count(int num_waiters, int num_task, int expected_woken) { + ParkingLot lot; + butil::atomic ready(0); + butil::atomic returned(0); + ParkingLotThreads threads(&lot); + for (int i = 0; i < num_waiters; ++i) { + threads.start([&] { + ready.fetch_add(1, butil::memory_order_release); + while (!threads.cancelled()) { + const auto state = lot.get_state(); + lot.wait(state); + returned.fetch_add(1, butil::memory_order_release); + } + }); + } + ASSERT_TRUE(wait_until([&] { return ready.load() == num_waiters; })); + if (GetParam()) { + ASSERT_TRUE(wait_until([&] { + return lot._waiter_num.load() == num_waiters; + })); + } + + // Neither ready nor _waiter_num proves that the kernel has enqueued a + // waiter. Retry until the actual futex return value confirms the desired + // batch. This avoids assuming that a sleep is a scheduling barrier. + int nwoken = -1; + ASSERT_TRUE(wait_until([&] { + nwoken = lot.signal(num_task); + return nwoken == expected_woken || nwoken < 0 || + nwoken > expected_woken; + })); + ASSERT_EQ(expected_woken, nwoken); + ASSERT_TRUE(wait_until([&] { return returned.load() >= nwoken; })); + } + +private: + GFLAGS_NAMESPACE::FlagSaver _flag_saver; +}; + +TEST_P(ParkingLotTest, initial_state) { + ParkingLot::State state; + ASSERT_EQ(0, state.val); + ASSERT_FALSE(state.stopped()); + + ParkingLot lot; + ASSERT_EQ(0, lot.get_state().val); + ASSERT_FALSE(lot.get_state().stopped()); + ASSERT_EQ(0, lot._waiter_num.load()); + ASSERT_EQ(GetParam(), lot._no_signal_when_no_waiter); +} + +TEST_P(ParkingLotTest, configuration_is_captured_at_construction) { + ParkingLot first; + bthread::FLAGS_parking_lot_no_signal_when_no_waiter = !GetParam(); + ParkingLot second; + ASSERT_EQ(GetParam(), first._no_signal_when_no_waiter); + ASSERT_EQ(!GetParam(), second._no_signal_when_no_waiter); + ASSERT_EQ(0, first.signal(1)); + ASSERT_EQ(0, second.signal(1)); + ASSERT_EQ(2, first.get_state().val); + ASSERT_EQ(2, second.get_state().val); +} + +TEST_P(ParkingLotTest, signal_without_waiters_updates_state) { + ParkingLot lot; + const auto initial = lot.get_state(); + ASSERT_EQ(0, lot.signal(1)); + const auto first = lot.get_state(); + ASSERT_EQ(2, first.val); + ASSERT_EQ(0, lot.signal(3)); + ASSERT_EQ(8, lot.get_state().val); + ASSERT_EQ(0, lot.signal(0)); + ASSERT_EQ(8, lot.get_state().val); + ASSERT_FALSE(lot.get_state().stopped()); + // Snapshots are values, not references to the live signal counter. + ASSERT_EQ(0, initial.val); + ASSERT_EQ(2, first.val); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, stale_state_does_not_wait_or_consume_signals) { + ParkingLot lot; + const auto initial = lot.get_state(); + lot.signal(1); + const auto first = lot.get_state(); + lot.signal(2); + ParkingLotThreads threads(&lot); + threads.start([&] { + for (int i = 0; i < 100 && !threads.cancelled(); ++i) { + lot.wait(initial); + lot.wait(first); + } + }); + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_EQ(6, lot.get_state().val); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, current_state_waits_even_after_previous_signals) { + ParkingLot lot; + lot.signal(3); + const auto current = lot.get_state(); + butil::atomic ready(false); + butil::atomic returned(0); + ParkingLotThreads threads(&lot); + threads.start([&] { + ready.store(true, butil::memory_order_release); + // Allow spurious wakeups without mistaking them for a new signal. + while (!threads.cancelled() && lot.get_state().val == current.val) { + lot.wait(current); + } + returned.store(1, butil::memory_order_release); + }); + ASSERT_TRUE(wait_until([&] { return ready.load(); })); + if (GetParam()) { + ASSERT_TRUE(wait_until([&] { return lot._waiter_num.load() == 1; })); + } + usleep(10 * 1000); + ASSERT_EQ(0, returned.load()); + const int nwoken = lot.signal(1); + ASSERT_GE(nwoken, 0); + ASSERT_LE(nwoken, 1); + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_EQ(1, returned.load()); + ASSERT_EQ(8, lot.get_state().val); + ASSERT_EQ(0, lot._waiter_num.load()); + ASSERT_EQ(0, lot.signal(1)); +} + +TEST_P(ParkingLotTest, signal_wakes_one_waiter) { + check_wake_count(1, 1, 1); +} + +TEST_P(ParkingLotTest, signal_wakes_a_limited_batch) { + check_wake_count(4, 2, 2); +} + +TEST_P(ParkingLotTest, signal_is_limited_by_available_waiters) { + check_wake_count(4, 10, 4); +} + +TEST_P(ParkingLotTest, stop_preserves_signal_count_and_is_idempotent) { + ParkingLot lot; + lot.signal(3); + const auto before_stop = lot.get_state(); + lot.stop(); + const auto stopped = lot.get_state(); + ASSERT_TRUE(stopped.stopped()); + ASSERT_EQ(7, stopped.val); + lot.stop(); + ASSERT_EQ(stopped.val, lot.get_state().val); + ASSERT_EQ(0, lot.signal(2)); + ASSERT_EQ(11, lot.get_state().val); + ASSERT_TRUE(lot.get_state().stopped()); + ASSERT_FALSE(before_stop.stopped()); + ASSERT_EQ(6, before_stop.val); + ASSERT_EQ(7, stopped.val); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, stop_before_wait_invalidates_old_snapshots) { + ParkingLot lot; + const auto initial = lot.get_state(); + lot.signal(2); + const auto before_stop = lot.get_state(); + lot.stop(); + ParkingLotThreads threads(&lot); + threads.start([&] { + // Callers must check State::stopped() before waiting. Only snapshots + // taken before stop(), not a fresh stopped snapshot, are waitable here. + lot.wait(initial); + lot.wait(before_stop); + }); + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_TRUE(lot.get_state().stopped()); + ASSERT_EQ(5, lot.get_state().val); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, stop_wakes_all_waiters) { + const int num_waiters = 8; + ParkingLot lot; + const auto initial = lot.get_state(); + butil::atomic ready(0); + ParkingLotThreads threads(&lot); + for (int i = 0; i < num_waiters; ++i) { + threads.start([&] { + ready.fetch_add(1, butil::memory_order_release); + while (!threads.cancelled() && !lot.get_state().stopped()) { + lot.wait(initial); + } + }); + } + ASSERT_TRUE(wait_until([&] { return ready.load() == num_waiters; })); + if (GetParam()) { + ASSERT_TRUE(wait_until([&] { + return lot._waiter_num.load() == num_waiters; + })); + } + lot.stop(); + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_TRUE(lot.get_state().stopped()); + ASSERT_EQ(1, lot.get_state().val); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, concurrent_signals_are_not_lost) { + const int num_signalers = 4; + const int iterations = 1000; + ParkingLot lot; + butil::atomic bad_returns(0); + ParkingLotThreads threads(&lot); + for (int i = 1; i <= num_signalers; ++i) { + threads.start([&, i] { + for (int j = 0; j < iterations && !threads.cancelled(); ++j) { + if (lot.signal(i) != 0) { + bad_returns.fetch_add(1); + } + } + }); + } + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_EQ(0, bad_returns.load()); + ASSERT_EQ(iterations * num_signalers * (num_signalers + 1), + lot.get_state().val); + ASSERT_FALSE(lot.get_state().stopped()); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, get_state_acquires_published_data) { + ParkingLot lot; + int payload = 0; + ParkingLotThreads threads(&lot); + threads.start([&] { + payload = 42; + lot.signal(1); + }); + // Observe the release in signal() through get_state(), not thread.join() + // or the helper's completion counter, before reading the non-atomic data. + ASSERT_TRUE(wait_until([&] { return lot.get_state().val != 0; })); + ASSERT_EQ(42, payload); +} + +TEST_P(ParkingLotTest, signal_counter_wraparound_preserves_stop_bit) { + ParkingLot lot; + // Atomic signed fetch_add wraps without signed-overflow UB. Keep num_task + // small enough that the separate signed left shift in signal() is valid. + ASSERT_EQ(0, lot.signal(INT_MAX / 2)); + ASSERT_EQ(INT_MAX - 1, lot.get_state().val); + ASSERT_EQ(0, lot.signal(1)); + ASSERT_EQ(INT_MIN, lot.get_state().val); + ASSERT_FALSE(lot.get_state().stopped()); + + lot.stop(); + ASSERT_EQ(INT_MIN + 1, lot.get_state().val); + ASSERT_EQ(0, lot.signal(INT_MAX / 2)); + ASSERT_EQ(-1, lot.get_state().val); + ASSERT_EQ(0, lot.signal(1)); + ASSERT_EQ(1, lot.get_state().val); + ASSERT_TRUE(lot.get_state().stopped()); +} + +TEST_P(ParkingLotTest, stop_races_with_signals_and_wait) { + const int iterations = 1000; + ParkingLot lot; + butil::atomic start(false); + ParkingLotThreads threads(&lot); + threads.start([&] { + while (!threads.cancelled()) { + const auto state = lot.get_state(); + if (state.stopped()) { + return; + } + lot.wait(state); + } + }); + threads.start([&] { + while (!start.load(butil::memory_order_acquire)) { + if (threads.cancelled()) { + return; + } + std::this_thread::yield(); + } + for (int i = 0; i < iterations && !threads.cancelled(); ++i) { + lot.signal(1); + } + }); + threads.start([&] { + while (!start.load(butil::memory_order_acquire)) { + if (threads.cancelled()) { + return; + } + std::this_thread::yield(); + } + for (int i = 0; i < iterations && !threads.cancelled(); ++i) { + lot.stop(); + } + }); + start.store(true, butil::memory_order_release); + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_EQ(iterations * 2 + 1, lot.get_state().val); + ASSERT_TRUE(lot.get_state().stopped()); + ASSERT_EQ(0, lot._waiter_num.load()); +} + +TEST_P(ParkingLotTest, racing_signal_and_wait_do_not_lose_wakeups) { + const int iterations = 2000; + ParkingLot lot; + butil::atomic start(0); + butil::atomic ready(0); + butil::atomic completed(0); + ParkingLotThreads threads(&lot); + threads.start([&] { + for (int round = 1; round <= iterations; ++round) { + while (start.load(butil::memory_order_acquire) != round) { + if (threads.cancelled()) { + return; + } + std::this_thread::yield(); + } + const auto state = lot.get_state(); + ready.store(round, butil::memory_order_release); + if (round % 3 == 0) { + std::this_thread::yield(); + } + lot.wait(state); + completed.store(round, butil::memory_order_release); + } + }); + + for (int round = 1; round <= iterations; ++round) { + SCOPED_TRACE(round); + start.store(round, butil::memory_order_release); + ASSERT_TRUE(wait_until([&] { + return ready.load(butil::memory_order_acquire) == round; + })); + if (round % 3 == 1) { + std::this_thread::yield(); + } + // The handshake precedes BOTH the waiter registration and signal(). + // Do not wait for _waiter_num here: that would remove the Dekker race. + // Exactly one signal per round; no later signal may rescue this round. + const int nwoken = lot.signal(1); + ASSERT_GE(nwoken, 0); + ASSERT_LE(nwoken, 1); + ASSERT_TRUE(wait_until([&] { + return completed.load(butil::memory_order_acquire) == round; + })) << "Possible lost wakeup"; + ASSERT_EQ(0, lot._waiter_num.load()); + } + ASSERT_TRUE(wait_until([&] { return threads.finished(); })); + ASSERT_EQ(iterations * 2, lot.get_state().val); + ASSERT_EQ(0, lot.signal(1)); + // This is a real-implementation stress regression, not a deterministic + // proof of weak-memory ordering. Removing a fence need not fail on x86. +} + +INSTANTIATE_TEST_SUITE_P(WaiterCheckModes, ParkingLotTest, ::testing::Bool()); + +} // namespace diff --git a/test/bthread_timer_thread_unittest.cpp b/test/bthread_timer_thread_unittest.cpp index c37c74cab5..c784cf252d 100644 --- a/test/bthread_timer_thread_unittest.cpp +++ b/test/bthread_timer_thread_unittest.cpp @@ -148,7 +148,9 @@ TEST(TimerThreadTest, RunTasks) { tm.start(); timer_thread.stop_and_join(); tm.stop(); - ASSERT_LE(tm.m_elapsed(), 15); + // stop_and_join() should wake the timer thread instead of waiting for the + // tasks scheduled 10 seconds later. Allow for CI runner scheduling delays. + ASSERT_LT(tm.m_elapsed(), 1000); // verify all runs in expected time range. keeper1.expect_first_run(); diff --git a/test/bthread_unittest.cpp b/test/bthread_unittest.cpp index ce2b368c91..c86bc83e8a 100644 --- a/test/bthread_unittest.cpp +++ b/test/bthread_unittest.cpp @@ -26,6 +26,7 @@ #include #include "bthread/bthread.h" #include "bthread/unstable.h" +#include "bthread/task_group.h" #include "bthread/task_meta.h" #include "bvar/bvar.h" @@ -265,6 +266,75 @@ TEST_F(BthreadTest, bthread_join) { ASSERT_EQ(0, bthread_start_urgent(&th, nullptr, join_self, nullptr)); } +struct JoinVisibilityData { + int seed; + bool delay_write; + int values[64]; +}; + +void* write_join_visibility_data(void* arg) { + auto data = static_cast(arg); + if (data->delay_write) { + // Give the caller a chance to enter the join wait path before writing. + bthread_usleep(1000); + } + for (size_t i = 0; i < ARRAY_SIZE(data->values); ++i) { + data->values[i] = data->seed + static_cast(i); + } + return nullptr; +} + +void check_join_visibility(bool join_after_exit) { + for (int round = 0; round < 1000; ++round) { + JoinVisibilityData data = {}; + data.seed = round + 1; + data.delay_write = !join_after_exit; + bthread_t tid; + ASSERT_EQ(0, bthread_start_background( + &tid, nullptr, write_join_visibility_data, &data)); + if (join_after_exit) { + // exists() uses a relaxed load, so observing completion here does + // not acquire the worker's writes. join() must still do so even + // when it returns without waiting on the butex. + while (bthread::TaskGroup::exists(tid)) { + bthread_usleep(10); + } + } + ASSERT_EQ(0, bthread_join(tid, nullptr)); + // The payload is deliberately non-atomic and is read only after join. + // Do not add a lock or a release/acquire completion flag to this test: + // that would provide an alternative way to publish the worker's data. + for (size_t i = 0; i < ARRAY_SIZE(data.values); ++i) { + ASSERT_EQ(data.seed + static_cast(i), data.values[i]) + << "round=" << round << " index=" << i + << " join_after_exit=" << join_after_exit; + } + } +} + +void* join_visibility_caller(void* arg) { + const bool is_bthread = *static_cast(arg); + EXPECT_EQ(is_bthread, bthread_self() != 0); + check_join_visibility(false); + check_join_visibility(true); + return nullptr; +} + +TEST_F(BthreadTest, join_visibility_from_pthread) { + bool is_bthread = false; + pthread_t caller; + ASSERT_EQ(0, pthread_create(&caller, nullptr, join_visibility_caller, &is_bthread)); + ASSERT_EQ(0, pthread_join(caller, nullptr)); +} + +TEST_F(BthreadTest, join_visibility_from_bthread) { + bool is_bthread = true; + bthread_t caller; + ASSERT_EQ(0, bthread_start_background( + &caller, nullptr, join_visibility_caller, &is_bthread)); + ASSERT_EQ(0, bthread_join(caller, nullptr)); +} + void* change_errno(void* arg) { errno = (intptr_t)arg; return nullptr; diff --git a/test/bvar_reducer_unittest.cpp b/test/bvar_reducer_unittest.cpp index 02923d92c5..f905e110d1 100644 --- a/test/bvar_reducer_unittest.cpp +++ b/test/bvar_reducer_unittest.cpp @@ -279,6 +279,72 @@ TEST_F(ReducerTest, non_primitive) { ASSERT_EQ(9, adder.get_value().x); } +// The generic Adder/Maxer/Miner keep their data in a shared AgentCombiner and +// hence expose share_combiner(), while the babylon-backed ones keep a value-type +// counter and do not. This tells the two implementations apart. +template +struct IsBabylonBacked + : std::integral_constant::value> {}; + +// Which implementation backs a given value type is decided by SFINAE at compile +// time, so a malformed condition silently makes the babylon-backed partial +// specializations unreachable without failing any runtime assertion. Assert the +// selection statically. +TEST_F(ReducerTest, babylon_counter_backend) { +#if WITH_BABYLON_COUNTER + static_assert(IsBabylonBacked >::value, + "Adder should be backed by a babylon counter"); + static_assert(IsBabylonBacked >::value, + "Adder should be backed by a babylon counter"); + static_assert(IsBabylonBacked >::value, + "Adder should be backed by a babylon counter"); + static_assert(IsBabylonBacked >::value, + "Maxer should be backed by a babylon counter"); + static_assert(IsBabylonBacked >::value, + "Miner should be backed by a babylon counter"); + // babylon counters only support arithmetic types not larger than 8 bytes. + // The types they reject must fall back to the generic implementation rather + // than break the build. + static_assert(IsBabylonBacked >::value == + (sizeof(long double) <= 8), + "Adder should follow the size constraint"); + static_assert(!IsBabylonBacked >::value, + "Adder should be backed by an AgentCombiner"); + static_assert(!IsBabylonBacked >::value, + "Adder should be backed by an AgentCombiner"); +#else + static_assert(!IsBabylonBacked >::value, + "Adder should be backed by an AgentCombiner"); + static_assert(!IsBabylonBacked >::value, + "Maxer should be backed by an AgentCombiner"); + static_assert(!IsBabylonBacked >::value, + "Miner should be backed by an AgentCombiner"); +#endif // WITH_BABYLON_COUNTER +} + +// reset() must return the accumulated value and restore the identity of the +// operator, which is how Window<> samples an operator without inverse. This runs +// against whichever implementation backs the value type, hence it covers the +// babylon counters as well when WITH_BABYLON_COUNTER is enabled. +TEST_F(ReducerTest, reset) { + bvar::Adder adder; + adder << 3 << 1; + ASSERT_EQ(4, adder.reset()); + ASSERT_EQ(0, adder.get_value()); + adder << 2; + ASSERT_EQ(2, adder.reset()); + + bvar::Maxer maxer; + maxer << 3 << 1; + ASSERT_EQ(3, maxer.reset()); + ASSERT_EQ(std::numeric_limits::min(), maxer.get_value()); + + bvar::Miner miner; + miner << 3 << 1; + ASSERT_EQ(1, miner.reset()); + ASSERT_EQ(std::numeric_limits::max(), miner.get_value()); +} + bool g_stop = false; struct StringAppenderResult { int count; diff --git a/test/seqlock_unittest.cpp b/test/seqlock_unittest.cpp new file mode 100644 index 0000000000..85ab4e9cb1 --- /dev/null +++ b/test/seqlock_unittest.cpp @@ -0,0 +1,391 @@ +// 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/time.h" +#include "butil/synchronization/lock.h" +#include "butil/synchronization/seqlock.h" + +namespace { + +// A multi-word payload. The seqlock's job is to make every reader observe a +// snapshot in which all words are equal; a torn read would see a mix of an old +// and a new value. Every field is a relaxed atomic, as Seqlock requires. +static const int kWords = 8; +struct Payload { + butil::atomic w[kWords]; + + void relaxed_set(uint64_t v) { + for (auto& i : w) { + // Store word by word (not as one atomic group) so that, without the + // seqlock, a concurrent reader could observe a torn value. + i.store(v, butil::memory_order_relaxed); + } + } + // Returns the first word and whether all words are equal to it. + uint64_t relaxed_get(bool* consistent) const { + uint64_t v0 = w[0].load(butil::memory_order_relaxed); + *consistent = true; + for (int i = 1; i < kWords; ++i) { + if (w[i].load(butil::memory_order_relaxed) != v0) { + *consistent = false; + } + } + return v0; + } +}; + +TEST(SeqlockTest, SingleThreadedReadWrite) { + butil::Seqlock<> seqlock; + Payload payload; + payload.relaxed_set(0); + + for (uint64_t v = 1; v <= 1000; ++v) { + seqlock.store([&] { payload.relaxed_set(v); }); + uint64_t got = seqlock.load([&] { + bool consistent = false; + uint64_t r = payload.relaxed_get(&consistent); + EXPECT_TRUE(consistent); + return r; + }); + ASSERT_EQ(v, got); + } +} + +TEST(SeqlockTest, LoadReturnsValueByType) { + butil::Seqlock<> seqlock; + butil::atomic payload(0); + seqlock.store([&] { + payload.store(42, butil::memory_order_relaxed); + }); + // load returns whatever the callback returns, by value. + int v = seqlock.load([&] { + return payload.load(butil::memory_order_relaxed); + }); + ASSERT_EQ(42, v); + // A different return type also works. + std::pair pr = seqlock.load([&] { + int v = payload.load(butil::memory_order_relaxed); + return std::make_pair(v, v + 1); + }); + ASSERT_EQ(42, pr.first); + ASSERT_EQ(43, pr.second); +} + +TEST(SeqlockTest, MutexSpecializationSingleThreaded) { + butil::Seqlock seqlock; + Payload payload; + payload.relaxed_set(0); + for (uint64_t v = 1; v <= 1000; ++v) { + seqlock.store([&] { payload.relaxed_set(v); }); + bool consistent = false; + uint64_t got = seqlock.load([&] { + bool c = false; + uint64_t r = payload.relaxed_get(&c); + consistent = c; + return r; + }); + ASSERT_TRUE(consistent); + ASSERT_EQ(v, got); + } +} + +// Single-threaded performance comparison: Seqlock vs Mutex. +// +// There is no contention here, so this measures the bare cost of the +// synchronization itself: +// * Seqlock<> : one relaxed counter load, two counter stores and a +// release fence on the write path; two counter loads and +// an acquire fence on the read path (no atomic RMW in +// either path). +// * Seqlock : the same counter work plus a lock/unlock pair on the +// write path; the read path is identical to Seqlock<>. +// * Mutex : a lock/unlock pair (atomic RMW + potential syscall on +// contention) on BOTH the write and the read path. +// +// Timings are reported, not asserted: absolute numbers depend on the machine, +// the compiler and the current load, so turning them into thresholds would +// make the test flaky. The test only asserts that the measurements ran and +// produced the expected values. +static const int kPerfRounds = 1000000; + +// Accumulates every value read so the compiler cannot optimize the read loops +// away as dead code. +static butil::atomic g_perf_sink(0); + +inline void perf_consume(uint64_t v) { + g_perf_sink.fetch_add(v, butil::memory_order_relaxed); +} + +// A pure compiler barrier: it emits no instruction, but forbids the compiler +// from moving memory accesses across it. Without one between the rounds below, +// a write body made purely of relaxed stores is legally collapsible into its +// last iteration -- which is exactly what happens to Seqlock<> once its counter +// increments are plain stores rather than fetch_add, and it measures 0 ns/op. +// The mutex bodies resist that on their own; barriering all of them keeps the +// comparison between the three honest. +inline void perf_barrier() { + butil::atomic_signal_fence(butil::memory_order_seq_cst); +} + +// Runs `body(round)` for `rounds` rounds and returns the elapsed nanoseconds. +// A short warmup runs first so that page faults, branch predictor and cache +// warmup are not charged to the measured loop. +template +int64_t TimeLoopNs(int rounds, Body&& body) { + int kWarmup = 1000; + for (int i = 1; i <= kWarmup; ++i) { + body(static_cast(i)); + perf_barrier(); + } + int64_t start_ns = butil::cpuwide_time_ns(); + for (int i = 1; i <= rounds; ++i) { + body(static_cast(i)); + perf_barrier(); + } + return butil::cpuwide_time_ns() - start_ns; +} + +inline double ns_per_op(int64_t elapsed_ns, int rounds) { + return static_cast(elapsed_ns) / rounds; +} + +TEST(SeqlockTest, SingleThreadedPerfVsMutex) { + // Seqlock<>: single-writer, no mutex. + butil::Seqlock<> seqlock; + Payload seqlock_payload; + seqlock_payload.relaxed_set(0); + int64_t seqlock_write_ns = TimeLoopNs(kPerfRounds, [&](uint64_t v) { + seqlock.store([&] { seqlock_payload.relaxed_set(v); }); + // Make both the counter and payload observable after every write so + // the compiler cannot eliminate stores to these local test objects. + asm volatile("" : : "m"(seqlock), "m"(seqlock_payload) : "memory"); + }); + int64_t seqlock_read_ns = TimeLoopNs(kPerfRounds, [&](uint64_t) { + perf_consume(seqlock.load([&] { + bool consistent = false; + return seqlock_payload.relaxed_get(&consistent); + })); + }); + + // Seqlock: same read path, writes additionally take a mutex. + butil::Seqlock mutex_seqlock; + Payload mutex_seqlock_payload; + mutex_seqlock_payload.relaxed_set(0); + int64_t mutex_seqlock_write_ns = TimeLoopNs(kPerfRounds, [&](uint64_t v) { + mutex_seqlock.store([&] { mutex_seqlock_payload.relaxed_set(v); }); + asm volatile("" : : "m"(mutex_seqlock), "m"(mutex_seqlock_payload) + : "memory"); + }); + int64_t mutex_seqlock_read_ns = TimeLoopNs(kPerfRounds, [&](uint64_t) { + perf_consume(mutex_seqlock.load([&] { + bool consistent = false; + return mutex_seqlock_payload.relaxed_get(&consistent); + })); + }); + + // Plain Mutex baseline: both sides take the lock. + butil::Mutex mutex; + Payload mutex_payload; + mutex_payload.relaxed_set(0); + int64_t mutex_write_ns = TimeLoopNs(kPerfRounds, [&](uint64_t v) { + { + std::lock_guard lk(mutex); + mutex_payload.relaxed_set(v); + } + asm volatile("" : : "m"(mutex), "m"(mutex_payload) : "memory"); + }); + int64_t mutex_read_ns = TimeLoopNs(kPerfRounds, [&](uint64_t) { + std::lock_guard lk(mutex); + bool consistent = false; + perf_consume(mutex_payload.relaxed_get(&consistent)); + }); + + printf("\n[seqlock perf] single thread, %d rounds, payload=%d words\n", + kPerfRounds, kWords); + printf("%-18s %14s %14s\n", "impl", "write(ns/op)", "read(ns/op)"); + printf("%-18s %14.2f %14.2f\n", "Seqlock<>", + ns_per_op(seqlock_write_ns, kPerfRounds), + ns_per_op(seqlock_read_ns, kPerfRounds)); + printf("%-18s %14.2f %14.2f\n", "Seqlock", + ns_per_op(mutex_seqlock_write_ns, kPerfRounds), + ns_per_op(mutex_seqlock_read_ns, kPerfRounds)); + printf("%-18s %14.2f %14.2f\n", "Mutex", + ns_per_op(mutex_write_ns, kPerfRounds), + ns_per_op(mutex_read_ns, kPerfRounds)); + printf("%-18s %13.2fx %13.2fx\n", "Mutex/Seqlock<>", + static_cast(mutex_write_ns) / seqlock_write_ns, + static_cast(mutex_read_ns) / seqlock_read_ns); + + // The loops really ran and every implementation published the last value. + ASSERT_GT(seqlock_write_ns, 0); + ASSERT_GT(mutex_seqlock_write_ns, 0); + ASSERT_GT(mutex_write_ns, 0); + ASSERT_GT(seqlock_read_ns, 0); + ASSERT_GT(mutex_seqlock_read_ns, 0); + ASSERT_GT(mutex_read_ns, 0); + ASSERT_GT(g_perf_sink.load(butil::memory_order_relaxed), 0u); + + bool consistent = false; + ASSERT_EQ(static_cast(kPerfRounds), + seqlock_payload.relaxed_get(&consistent)); + ASSERT_TRUE(consistent); + ASSERT_EQ(static_cast(kPerfRounds), + mutex_seqlock_payload.relaxed_get(&consistent)); + ASSERT_TRUE(consistent); + ASSERT_EQ(static_cast(kPerfRounds), + mutex_payload.relaxed_get(&consistent)); + ASSERT_TRUE(consistent); +} + +// Concurrent consistency tests +struct SharedState { + butil::Seqlock<>* single_writer_seqlock = nullptr; // single-writer lock + butil::Seqlock* multi_writer_seqlock = nullptr; // multi-writer lock + Payload payload; + butil::atomic stopped{false}; + butil::atomic version{0}; // source of the value written + butil::atomic reads{0}; // total reads performed + butil::atomic torn{0}; // inconsistent snapshots observed +}; + +void* SingleWriterThread(void* arg) { + SharedState* shared_state = static_cast(arg); + while (!shared_state->stopped.load(butil::memory_order_relaxed)) { + uint64_t version = shared_state->version.fetch_add(1, butil::memory_order_relaxed) + 1; + shared_state->single_writer_seqlock->store([&] { + shared_state->payload.relaxed_set(version); + }); + } + return nullptr; +} + +void* MultiWriterThread(void* arg) { + SharedState* shared_state = static_cast(arg); + while (!shared_state->stopped.load(butil::memory_order_relaxed)) { + uint64_t version = shared_state->version.fetch_add(1, butil::memory_order_relaxed) + 1; + shared_state->multi_writer_seqlock->store([&] { + shared_state->payload.relaxed_set(version); + }); + } + return nullptr; +} + +struct ReaderArg { + SharedState* shared_state; + bool multi_writer; +}; + +void* ReaderThread(void* arg) { + ReaderArg* reader_arg = static_cast(arg); + SharedState* shared_state = reader_arg->shared_state; + uint64_t local_reads = 0; + uint64_t local_torn = 0; + while (!shared_state->stopped.load(butil::memory_order_relaxed)) { + bool consistent = false; + auto load_body = [&] { + bool c = false; + uint64_t r = shared_state->payload.relaxed_get(&c); + consistent = c; + return r; + }; + if (reader_arg->multi_writer) { + shared_state->multi_writer_seqlock->load(load_body); + } else { + shared_state->single_writer_seqlock->load(load_body); + } + ++local_reads; + if (!consistent) { + ++local_torn; + } + } + shared_state->reads.fetch_add(local_reads, butil::memory_order_relaxed); + shared_state->torn.fetch_add(local_torn, butil::memory_order_relaxed); + return nullptr; +} + +TEST(SeqlockTest, SingleWriterManyReaders) { + SharedState shared_state; + butil::Seqlock<> sl; + shared_state.single_writer_seqlock = &sl; + shared_state.payload.relaxed_set(0); + + const int kReaders = 4; + pthread_t writer; + pthread_t readers[kReaders]; + ReaderArg args[kReaders]; + + ASSERT_EQ(0, pthread_create(&writer, nullptr, SingleWriterThread, &shared_state)); + for (int i = 0; i < kReaders; ++i) { + args[i].shared_state = &shared_state; + args[i].multi_writer = false; + ASSERT_EQ(0, pthread_create(&readers[i], nullptr, ReaderThread, &args[i])); + } + + usleep(500 * 1000); // 0.5s of hammering + shared_state.stopped.store(true, butil::memory_order_relaxed); + + pthread_join(writer, nullptr); + for (auto reader : readers) { + pthread_join(reader, nullptr); + } + + ASSERT_GT(shared_state.reads.load(), 0u); + ASSERT_EQ(0u, shared_state.torn.load()) << "readers observed torn snapshots"; +} + +TEST(SeqlockTest, MultiWriterManyReaders) { + SharedState shared_state; + butil::Seqlock seqlock; + shared_state.multi_writer_seqlock = &seqlock; + shared_state.payload.relaxed_set(0); + + const int kWriters = 3; + const int kReaders = 4; + pthread_t writers[kWriters]; + pthread_t readers[kReaders]; + ReaderArg args[kReaders]; + + for (auto& writer : writers) { + ASSERT_EQ(0, pthread_create(&writer, nullptr, MultiWriterThread, &shared_state)); + } + for (int i = 0; i < kReaders; ++i) { + args[i].shared_state = &shared_state; + args[i].multi_writer = true; + ASSERT_EQ(0, pthread_create(&readers[i], nullptr, ReaderThread, &args[i])); + } + + usleep(500 * 1000); + shared_state.stopped.store(true, butil::memory_order_relaxed); + + for (auto writer : writers) { + pthread_join(writer, nullptr); + } + for (auto reader : readers) { + pthread_join(reader, nullptr); + } + + ASSERT_GT(shared_state.reads.load(), 0u); + ASSERT_EQ(0u, shared_state.torn.load()) << "readers observed torn snapshots"; +} + +} // namespace