feat: add URMA transport support - #3428
Conversation
feat: add URMA transport support
|
I'm not familiar with URMA, but it seems that the code in |
There was a problem hiding this comment.
Pull request overview
This PR adds an optional URMA (UMDK Unified Remote Memory Access) transport backend to bRPC’s existing Socket/Transport architecture, including a TCP-based negotiation/handshake with transparent TCP fallback when URMA is unavailable or negotiation fails.
Changes:
- Introduces
SOCKET_MODE_URMA,UrmaTransport, and per-connectionurma::UrmaEndpoint, plus handshake (v2 binary + v3 protobuf) and a credit-based flow-control model. - Adds build-system integration for URMA across CMake/Make/Bazel (including fetching UMDK headers and a link-time URMA mock when
liburma/hardware is unavailable). - Adds URMA unit tests, documentation (EN/CN), and a new
urma_performanceexample.
Reviewed changes
Copilot reviewed 36 out of 37 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| WORKSPACE | Adds Bazel external umdk repo for URMA headers. |
| MODULE.bazel | Adds bzlmod umdk repo rule for URMA headers. |
| BUILD.bazel | Wires URMA define and exposes URMA protos/headers to Bazel build. |
| bazel/config/BUILD.bazel | Adds brpc_with_urma config setting. |
| bazel/third_party/umdk/umdk.BUILD | Defines urma_headers cc_library for UMDK header export. |
| CMakeLists.txt | Adds WITH_URMA build option, header discovery/download, library detection, mock handling, and proto compilation. |
| Makefile | Includes src/brpc/urma sources conditionally; excludes mock when linking real liburma. |
| config_brpc.sh | Adds --with-urma configuration, header/lib detection, and mock selection. |
| src/brpc/socket_mode.h | Adds SOCKET_MODE_URMA. |
| src/brpc/transport_factory.h | Minor API/comment cleanup and parameter naming. |
| src/brpc/transport_factory.cpp | Adds URMA transport creation/context-init plumbing. |
| src/brpc/socket.h | Adds URMA forward decls and friend access for new endpoint/handshake types. |
| src/brpc/channel.h | Updates socket_mode comments to reference SocketMode generally. |
| src/brpc/channel.cpp | Incorporates URMA into channel signature hashing. |
| src/brpc/server.h | Clarifies socket_mode comment for accepted sockets. |
| src/brpc/input_messenger.h | Adds friend access for URMA endpoint. |
| src/brpc/input_messenger.cpp | Ensures URMA-delivered messages are executed off poller threads. |
| src/brpc/urma_transport.h | Declares UrmaTransport composing TcpTransport fallback and managing URMA state. |
| src/brpc/urma_transport.cpp | Implements transport dispatch, epoll-out waiting, message queueing, and context init constraints. |
| src/brpc/urma/urma_endpoint.h | Declares per-socket endpoint state machine, resources, polling/event CQ handling, and data-path APIs. |
| src/brpc/urma/urma_endpoint.cpp | Implements endpoint handshake, CQ polling, SEND/RECV processing, and connect driver. |
| src/brpc/urma/urma_helper.h | Declares global URMA init, polling-mode setup, and memory registration APIs. |
| src/brpc/urma/urma_helper.cpp | Implements global URMA init, buffer-pool + IOBuf allocator hijack, and user memory registration. |
| src/brpc/urma/urma_handshake.proto | Defines protobuf v3 handshake message (UrmaHello). |
| src/brpc/urma/urma_handshake.h | Declares handshake strategy abstraction + v2/v3 implementations and negotiation helpers. |
| src/brpc/urma/urma_handshake.cpp | Implements v2 binary + v3 protobuf handshake serialization/parsing and validation. |
| src/brpc/urma/urma_bonding.h | Optional provider-extension include gate for URMA bonding. |
| src/brpc/urma/mock_urma.cpp | Provides link-time URMA C-API mock for builds without liburma/hardware. |
| test/brpc_urma_unittest.cpp | Adds URMA handshake/helper/mock unit coverage. |
| example/cmake/BrpcExample.cmake | Updates example toolchain to handle newer Protobuf/Abseil and URMA include/defines. |
| example/urma_performance/CMakeLists.txt | Adds URMA performance example build (links liburma if present). |
| example/urma_performance/test.proto | Proto definitions for performance example RPC service. |
| example/urma_performance/server.cpp | URMA/TCP selectable server for performance testing. |
| example/urma_performance/client.cpp | URMA/TCP selectable benchmark client with warmup + worker bthreads. |
| docs/en/urma.md | English documentation for building/using UrmaTransport and configuration flags. |
| docs/cn/urma.md | Chinese documentation for building/using UrmaTransport and configuration flags. |
| .gitignore | Ignores local URMA proposal notes and graph artifacts. |
Suppressed comments (1)
src/brpc/urma/urma_helper.cpp:103
- kIOBufBlockHeaderLen is currently hard-coded to 32, which assumes the current IOBuf::Block layout. Once butil/iobuf_inl.h is included, this can be derived from sizeof(butil::IOBuf::Block) to stay correct if the layout ever changes.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // --------------------------------------------------------------------------- | ||
| // 4-byte ACK: HELLO_ACK_URMA_OK bit. | ||
| // --------------------------------------------------------------------------- | ||
| TEST(UrmaHandshakeTest, ack_bit_is_rdma_ok) { |
| #include "butil/atomicops.h" | ||
| #include "butil/containers/flat_map.h" | ||
| #include "butil/iobuf.h" | ||
| #include "butil/logging.h" |
@wwbmmm Thanks for pointing this out. Yes, the current user-space liburma API largely aligns with the RDMA interface. The differences are mostly pushed down to the kernel-space uburma , ubcore library. |
|
The common handshake abstraction is being discussed in:#3432 |
- Rename ack_bit_is_rdma_ok to ack_bit_is_urma_ok. - Replace hard-coded IOBuf block header size with sizeof(butil::IOBuf::Block). - Pin UMDK dependency to a specific commit instead of a mutable tag.
|
|
||
| ## 编译配置 | ||
|
|
||
| ### CMake 编译 |
There was a problem hiding this comment.
Bazel‑specific build instructions may be added.
|
|
||
| `WITH_URMA=ON` 使用上游 UMDK 头文件进行编译。CMake 优先使用系统安装的 | ||
| SDK；找不到头文件时，会参照 Mooncake 的 mock 构建方式下载固定版本的 | ||
| UMDK，可通过 `DOWNLOAD_URMA_HEADERS=OFF` 禁止下载。找到 `liburma` 时使用 |
There was a problem hiding this comment.
DOWNLOAD_URMA_HEADERS lacks Bazel build support
|
|
||
| extern "C" { | ||
|
|
||
| urma_status_t urma_init(urma_init_attr_t *init_attr) { |
There was a problem hiding this comment.
The mock functionality lacks an independent feature switch, which can break normal URMA functionality.
|
@dwh110 Thanks for the review and suggestion. I have noted this concern and will evaluate the proposed changes. I will update the PR once the necessary modifications are implemented. |
|
I have addressed the main review feedback:
Please take another look when available. Kindly manually verify any potential issue before posting a review comment. please feel free ping me anytime. |
|
Will this issue : #3432 resolved in this PR, or in further PR? |
|
@wwbmmm We plan to handle #3432 in a follow-up PR. We're currently testing the common handshake approach internally with RDMA and UBShm. For URMA, we'd like to get this PR merged first, since the handshake work will depend on how this PR progresses. Once this pr merged, we plan to move the URMA handshake to the common handshake framework as well. If there are any other concerns with this PR, feel free to ping me. |
|
LGTM |
| } | ||
| }; | ||
|
|
||
| int completion_error = drain_cq(); |
There was a problem hiding this comment.
Should it be activated only in polling mode， not in event_mode?
There was a problem hiding this comment.
This is mainly an optimization for event_mode. Once we get an event notification, we drain all currently available events in one go, which helps reduce event handling latency.
|
Overall the new UrmaTransport subsystem is well-structured and defensively written (handshake parsing bounds-checks, flow-control credits with overflow validation, robust TCP fallback). One concern in the Bazel build is worth addressing before merge.
🤖 This reply was automatically generated by brpc-oncall |
|
@wwbmmm Addressed the URMA-related Copilot review findings in fe6ea6f:
|
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
Copilot reviewed 38 out of 39 changed files in this pull request and generated 1 comment.
Suppressed comments (4)
BUILD.bazel:1
- URMA sources are included unconditionally via the base
glob([... 'src/brpc/urma/*.cpp', 'src/brpc/urma/**/*.cpp', ...])and then included again via the URMAselect(...), which can cause duplicate-source errors and also defeats the intent to includemock_urma.cpponly in mock builds (and exclude it when linking the real-lurma). Remove URMA globs from the baseglob(...)and rely solely on theselect(...)(or exclude the URMA paths from the base glob) so URMA sources (and especiallymock_urma.cpp) are only compiled in the intended configurations.
# Licensed to the Apache Software Foundation (ASF) under one or more
src/brpc/urma/urma_handshake.cpp:1
- On invalid v2 hello (e.g., version mismatch or
ValidHellofailure), the function returns0without draining any trailing bytes whenmsg_len > HELLO_PACKET_LEN. This can leave the TCP stream misaligned and break the subsequent TCP fallback parsing on the same connection. Consider draining the advertised trailing bytes (bounded byHELLO_MSG_LEN_MAX) before returning0for invalid hellos, and consider treating out-of-rangemsg_len(or other structural corruption after seeing the URMA magic) as a handshake failure rather than a graceful fallback to avoid consuming bytes that TCP would need to parse.
src/brpc/urma/urma_handshake.cpp:1 - On invalid v2 hello (e.g., version mismatch or
ValidHellofailure), the function returns0without draining any trailing bytes whenmsg_len > HELLO_PACKET_LEN. This can leave the TCP stream misaligned and break the subsequent TCP fallback parsing on the same connection. Consider draining the advertised trailing bytes (bounded byHELLO_MSG_LEN_MAX) before returning0for invalid hellos, and consider treating out-of-rangemsg_len(or other structural corruption after seeing the URMA magic) as a handshake failure rather than a graceful fallback to avoid consuming bytes that TCP would need to parse.
src/brpc/urma/urma_helper.cpp:1 - This busy-spin loop can burn CPU if multiple threads race on URMA initialization (especially under contention or slow provider init). Prefer a blocking wait (e.g., butex/condition variable) or at least a backoff strategy (
sched_yield/sleep) to avoid pegging a core during initialization.
| namespace brpc { | ||
|
|
||
| InputMessenger* g_messenger = nullptr; | ||
| InputMessenger *g_messenger = nullptr; |
There was a problem hiding this comment.
There's no need to change these formats
| } | ||
|
|
||
| InputMessengerProcessor::ParsingStreamGuard::ParsingStreamGuard(Socket* socket, StreamType type) | ||
| InputMessengerProcessor::ParsingStreamGuard::ParsingStreamGuard(Socket *socket, |
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
Copilot reviewed 38 out of 39 changed files in this pull request and generated 1 comment.
Suppressed comments (4)
BUILD.bazel:1
- URMA sources are excluded from
BRPC_BASE_SRCSviaURMA_SRC_PATTERNS, but then added back via select(). This is good; however, earlier in the diffBRPC_BASE_SRCSoriginally included URMA patterns (andsrc/brpc/**/*.cppalready coverssrc/brpc/urma/**). If the final state accidentally double-includes URMA sources (base + select), Bazel will fail on duplicatesrcs. Please ensure URMA sources (includingmock_urma.cpp) are only introduced through the URMA select(), and are fully excluded from the base glob so duplicates cannot occur.
# Licensed to the Apache Software Foundation (ASF) under one or more
src/brpc/urma_transport.cpp:1
QueueMessagereturns early forlast_msgin event mode, which prevents the transport from honoring the caller’s request to queue/flush processing for the final message (notably whenInputMessengerProcessortreats URMA like other poller-driven transports). This can result in user message processing running on the poller/dispatcher thread (or being deferred in an unintended way). Consider removing this early-return and instead always enqueue whenlast_msgis true, matching the intent described inInputMessengerProcessorfor URMA.
src/brpc/urma/urma_helper.cpp:1- The init contended path spins in a tight loop waiting for
g_init_once==2, which can waste CPU (especially if URMA init is slow due to device/provider discovery). Add a small backoff (e.g.,bthread_yield()/sched_yield()or a short sleep) or use a waitable primitive to reduce unnecessary CPU burn under contention.
src/brpc/urma/urma_handshake.cpp:1 - On invalid v2 hello (version/length checks fail), the code returns without draining any extra bytes when
msg_lenadvertises a packet larger than the fixed body. If a peer sendsURMAmagic + an invalid header with a largemsg_len, the remaining bytes will stay in the TCP stream and can disrupt subsequent fallback parsing/processing. Consider drainingmin(msg_len, HELLO_MSG_LEN_MAX) - HELLO_PACKET_LENbytes (when msg_len is in-range) before returning negotiated=false, or explicitly failing/closing on malformed hello to keep the stream in a well-defined state.
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
Copilot reviewed 37 out of 38 changed files in this pull request and generated 2 comments.
Suppressed comments (3)
src/brpc/urma_transport.cpp:1
- This early return can drop the pending
input_msgwithout processing it (in non-polling mode), because the closure is neither released nor executed here. A concrete fix is to always consume/processinput_msginQueueMessage(either inline or by spawning a bthread) and remove thereturnpath; if URMA non-polling should behave like TCP, delegate to the same processing strategy rather than skipping the message.
BUILD.bazel:1 - The comment says URMA sources should be kept out of
BRPC_BASE_SRCS, butBRPC_BASE_SRCScurrently includesURMA_SRC_PATTERNS. This will (a) duplicate URMA sources when the laterselect()also addsURMA_SRCS, and (b) unintentionally includemock_urma.cppeven for thebrpc_with_urma_use_realconfiguration, risking symbol clashes with-lurma. Fix by removing+ URMA_SRC_PATTERNSfrom theBRPC_BASE_SRCSglob input, and rely solely on theselect()blocks to add URMA sources (and the mock only when appropriate).
# Licensed to the Apache Software Foundation (ASF) under one or more
src/brpc/urma/urma_helper.cpp:1
- This is a tight spin-wait and will burn CPU if URMA initialization takes non-trivial time (device enumeration, mmap+register, etc.). Prefer a blocking wait strategy (e.g., a condition variable/butex/
bthread_usleepbackoff) to reduce CPU usage while waiting forg_init_onceto reach the terminal state.
| 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 (nw >= 0) { | ||
| written += nw; | ||
| continue; | ||
| } |
be43c42 to
7fad174
Compare
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
Copilot reviewed 38 out of 39 changed files in this pull request and generated 2 comments.
Suppressed comments (2)
src/brpc/urma/urma_helper.cpp:1
- This check-then-decrement on
outstandingis racy: multiple concurrent deallocations can observe> 0and both decrement, underflowing thesize_tcounter. Use an atomic decrement that cannot underflow (e.g.,auto prev = fetch_sub(1)and clamp/restore ifprev == 0, or switch to a signed counter with assertions under the pool lock).
src/brpc/urma/urma_helper.cpp:1 - The empty spin loop can burn CPU if URMA initialization takes non-trivial time (e.g., provider load/device query). Add a backoff/yield (e.g.,
sched_yield()/bthread_yield()/ small sleep) or replace with a proper wait/notify mechanism to avoid busy-waiting.
| int InputMessenger::AddNonProtocolHandler(const InputMessageHandler &handler) { | ||
| if (handler.parse == nullptr || handler.process == nullptr || | ||
| handler.name == nullptr) { | ||
| CHECK(false) << "Invalid argument"; | ||
| return -1; | ||
| } | ||
| BAIDU_SCOPED_LOCK(_add_handler_mutex); | ||
| if (nullptr == _handlers) { | ||
| _handlers = new InputMessageHandler[_capacity]; | ||
| memset(_handlers, 0, sizeof(*_handlers) * _capacity); | ||
| _non_protocol = true; | ||
| } | ||
| if (!_non_protocol) { | ||
| CHECK(false) << "AddHandler was invoked"; | ||
| return -1; | ||
| } | ||
| const int index = _max_index.load(butil::memory_order_relaxed) + 1; | ||
| _handlers[index] = handler; | ||
| _max_index.store(index, butil::memory_order_release); | ||
| return 0; | ||
| } |
| // On dedicated transport pollers (RDMA, UBRING, and URMA), all messages must | ||
| // be executed in a new bthread. Processing user code on the poller may call | ||
| // synchronization primitives and prevent it from draining more events. | ||
| if (_socket->_socket_mode == SOCKET_MODE_RDMA || | ||
| _socket->_socket_mode == SOCKET_MODE_UBRING || | ||
| _socket->_socket_mode == SOCKET_MODE_URMA) { |
|
@wwbmmm I just noticed a conflict between the format and master of these two files: |
There was a problem hiding this comment.
🔵 Needs a closer look
The change introduces a large new transport/data-path and extensive build-system wiring, which warrants final human review for correctness, safety, and integration risk.
Review details
- Files reviewed: 37/38 changed files
- Comments generated: 1
- Review effort level: Lite
| // Wait for the other thread to finish init. | ||
| while (g_init_once.load(butil::memory_order_acquire) != 2) { | ||
| // spin briefly | ||
| } | ||
| } |
What problem does this PR solve?
About discussion：
Issue Number: Related to #3401
Problem Summary:
brpc does not currently provide a remote-memory transport based on the UMDK
URMA API.
This PR adds URMA (Unified Remote Memory Access) as an optional transport in
the existing
Transportframework. TCP is used for connection establishmentand URMA capability/resource negotiation. After both peers successfully
negotiate URMA support, the data path switches to URMA. When URMA is
unavailable or negotiation does not succeed, the connection falls back to TCP.
The initial implementation supports the
baidu_stdprotocol only and isdisabled by default.
What is changed and the side effects?
Changed:
SOCKET_MODE_URMAandUrmaTransport, and integrate them with theexisting transport, socket, and message-processing framework.
UrmaEndpointfor URMA resource management, SEND/RECV processing,completion handling, and credit-based flow control.
and v3 protobuf handshake formats, with TCP fallback.
memory registration.
is available.
a link-time URMA mock for build and test environments without
liburmaorURMA hardware.
urma_performanceexample.Side effects:
Performance effects:
default behavior remain unchanged.
default configuration is 65,536 buffers of 8 KiB each (512 MiB in total).
actively poll the JFC.
configuration. This PR does not claim a specific hardware performance
improvement.
Breaking backward compatibility:
SOCKET_MODE_URMAand the URMA memory-registration APIs are additive.Existing socket modes and their default behavior remain unchanged.
Check List: