diff --git a/cmake/config.cmake b/cmake/config.cmake index 3dc0ee673..42ca4a255 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -78,6 +78,43 @@ if (YLT_ENABLE_IBV) target_link_libraries(${ylt_target_name} INTERFACE -libverbs -lmlx5) endif () endif () +option(YLT_ENABLE_URMA "Enable URMA support" OFF) +if (YLT_ENABLE_URMA) + message(STATUS "Enable URMA support") + find_path(URMA_INCLUDE_PATH NAMES urma_api.h + HINTS ${URMA_ROOT} 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 ${URMA_ROOT} 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 OR NOT URMA_BOND_INCLUDE_PATH) + message(FATAL_ERROR + "Fail to find URMA headers. Install UMDK headers or set URMA_ROOT.") + endif() + message(STATUS "Found URMA headers: ${URMA_INCLUDE_PATH}") + if (URMA_BOND_INCLUDE_PATH) + message(STATUS "Found URMA bonding headers: ${URMA_BOND_INCLUDE_PATH}") + endif() + if(CMAKE_PROJECT_NAME STREQUAL "yaLanTingLibs") + add_compile_definitions("YLT_ENABLE_URMA") + include_directories(${URMA_INCLUDE_PATH}) + if (URMA_BOND_INCLUDE_PATH) + include_directories(${URMA_BOND_INCLUDE_PATH}) + endif() + link_libraries(-lurma) + else () + target_compile_definitions(${ylt_target_name} INTERFACE "YLT_ENABLE_URMA") + target_include_directories(${ylt_target_name} INTERFACE ${URMA_INCLUDE_PATH}) + if (URMA_BOND_INCLUDE_PATH) + target_include_directories(${ylt_target_name} INTERFACE ${URMA_BOND_INCLUDE_PATH}) + endif() + target_link_libraries(${ylt_target_name} INTERFACE -lurma) + endif () +endif () if (YLT_ENABLE_CUDA) message(STATUS "Enable cuda support") find_package(CUDAToolkit REQUIRED) diff --git a/docs/urma_rpc.md b/docs/urma_rpc.md new file mode 100644 index 000000000..b2a91b074 --- /dev/null +++ b/docs/urma_rpc.md @@ -0,0 +1,157 @@ +# URMA RPC + +本文介绍 `coro_rpc` 的 URMA 传输层、自动升级机制、配置方式和底层接口。URMA 是可选功能；未启用或初始化失败时，`coro_rpc` 继续使用 TCP。 + +## 1. 架构与连接流程 + +URMA RPC 由三层组成： + +1. `coro_rpc`：负责 RPC 编解码、handler、连接管理和 attachment。 +2. `coro_io::urma_socket_t`：负责 TCP 握手、URMA Jetty 建链、收发队列和完成事件。 +3. URMA 资源层：负责设备、context、JFC/JFCE、Jetty、Segment 和 work request。 + +连接建立先使用 TCP 交换 EID、UASID、Jetty ID 和 buffer pool segment；随后双方导入对端资源，数据面切换到 URMA。TCP 只用于握手。 + +```mermaid +flowchart TD + RPC["coro_rpc client/server"] --> SOCKET["coro_io::urma_socket_t"] + SOCKET --> HANDSHAKE["TCP handshake"] + HANDSHAKE --> META["Peer metadata: EID, UASID, Jetty ID, Segment"] + SOCKET --> RESOURCES["URMA resources"] + RESOURCES --> CONTEXT["Device and context"] + RESOURCES --> QUEUES["JFC/JFCE and Jetty"] + RESOURCES --> SEGMENT["Registered buffer pool segment"] + META --> IMPORT["Import peer resources"] + QUEUES --> WR["Send/receive work requests"] + SEGMENT --> WR + IMPORT --> WR + WR --> COMPLETION["Completion polling or event loop"] +``` + +buffer pool 将一块连续内存注册为 Segment，再切分为固定大小的 buffer，减少重复注册开销。发送窗口受本地发送 buffer 数和对端接收 buffer 数共同限制。 + +## 2. 构建与显式启用 + +```bash +cmake -S . -B build -DYLT_ENABLE_URMA=ON -DBUILD_EXAMPLES=ON +cmake --build build --target coro_rpc_urma_example -j +``` + +CMake 从 `URMA_ROOT`、`/usr/include` 和 `/usr/local/include` 查找 UMDK +头文件，支持 `urma/`、`umdk/urma/` 和 UMDK 源码目录布局。例如： + +```bash +cmake -S . -B build -DYLT_ENABLE_URMA=ON \ + -DURMA_ROOT=/opt/umdk +``` + +也可以通过环境变量指定： + +```bash +export URMA_ROOT=/opt/umdk +``` + +配置必须同时包含 `urma_api.h` 和 bonding 扩展头 `urma_ubagg.h`，并在启用 +URMA 时链接系统 `liburma`。 + +```cpp +coro_rpc_server server(std::thread::hardware_concurrency(), 9000); +server.init_urma(); +server.register_handler(); +server.start(); +``` + +```cpp +coro_rpc_client client; +if (!client.init_urma()) { + co_return; +} +auto ec = co_await client.connect("127.0.0.1:9000"); +auto result = co_await client.call_for(30s, "hello urma"); +``` + +## 3. 配置 + +```cpp +coro_io::urma_socket_t::config_t config{ + .cq_size = 128, + .recv_buffer_cnt = 64, + .send_buffer_cnt = 64, + .buffer_size = 4096, + .max_memory_usage = 256ull * 1024 * 1024, + .device_name = "bonding_dev_0", + .eid_index = 0, + .tp_type = URMA_CTP, +}; + +server.init_urma(config); +client.init_urma(config); +``` + +## 4. 环境变量自动升级 + +默认 TCP 配置下，设置 `URMA_RPC_ENABLE=1` 可自动探测 URMA 设备并升级： + +```bash +export URMA_RPC_ENABLE=1 +export URMA_RPC_DEVICE=bonding_dev_0 +export URMA_RPC_EID_INDEX=0 +``` + +`URMA_RPC_ENABLE` 接受 `1`、`on`、`true`、`yes`。未设置、关闭或设备探测失败时回退 TCP。显式调用 `init_urma(config)` 的配置优先。 + +| 变量 | 作用 | +| --- | --- | +| `URMA_RPC_ENABLE` | 启用自动升级 | +| `URMA_RPC_DEVICE` | 设备名；为空时自动选择 | +| `URMA_RPC_EID_INDEX` | EID 下标，默认 `0` | +| `URMA_RPC_CQ_SIZE` | completion queue 大小 | +| `URMA_RPC_RECV_BUFFER_CNT` / `URMA_RPC_SEND_BUFFER_CNT` | 收发 buffer 数 | +| `URMA_RPC_BUFFER_SIZE` | 单个 buffer 大小 | +| `URMA_RPC_MAX_MEMORY_USAGE` | buffer pool 上限，单位 byte | +| `URMA_RPC_TP_TYPE` | `ctp`、`rtp` 或其他支持的类型 | +| `URMA_RPC_EVENT_MODE` | JFCE 事件模式开关 | +| `URMA_RPC_BUSY_POLL_BUDGET` | 事件唤醒后的忙轮询次数 | +| `URMA_RPC_POLL_INTERVAL` | 轮询间隔，单位微秒 | + +实现位于 `include/ylt/coro_io/urma/urma_rpc_env.hpp`，行为测试位于 `src/coro_rpc/tests/test_urma_rpc_env.cpp`。 + +## 5. URMA 接口 + +| 层次 | 主要接口 | 作用 | +| --- | --- | --- | +| 生命周期 | `urma_init`、`urma_uninit` | 初始化和释放 runtime | +| 设备 | `urma_get_device_list`、`urma_get_eid_list`、`urma_query_device` | 查询设备、EID 和能力 | +| Context | `urma_create_context`、`urma_delete_context` | 创建设备上下文 | +| 完成队列 | `urma_create_jfc`、`urma_create_jfce`、`urma_poll_jfc`、`urma_wait_jfc` | 等待并读取完成记录 | +| 数据队列 | `urma_create_jfs`、`urma_create_jfr`、`urma_create_jetty` | 创建收发资源 | +| 远端资源 | `urma_import_seg`、`urma_import_jetty` | 导入对端 Segment 和 Jetty | +| 本地内存 | `urma_register_seg` | 注册 URMA 可访问内存 | +| 数据收发 | `urma_post_jetty_send_wr`、`urma_post_jetty_recv_wr` | 提交 work request | + +典型顺序是：创建 context → 创建 JFC/JFCE 和 Jetty → 注册 Segment → 交换元数据 → 导入对端资源 → post WR → poll/wait completion → 释放资源。声明来自外部 UMDK 的 `urma_api.h` 和 `urma_types.h`。 + +## 6. 大数据与压测 + +大 payload 建议使用 attachment： + +```cpp +client.set_req_attachment(payload); +auto result = co_await client.call_for(30s, payload.size()); +``` + +```bash +cmake --build build --target coro_rpc_urma_benchmark -j +``` + +详细参数见 [`urma_benchmark/README.md`](../src/coro_rpc/examples/urma_benchmark/README.md)。`--transport raw` 用于排除 RPC 编解码开销，`--rpc attach_sink` 用于测试 attachment 快路径。 + +## 7. 排障 + +- 没有 URMA 目标：确认使用 `-DYLT_ENABLE_URMA=ON`，并安装匹配的库、头文件和驱动。 +- 自动升级未生效：确认 `URMA_RPC_ENABLE` 有效，且应用走默认 TCP 配置路径。 +- 初始化或连接失败：检查设备名、EID、TP 类型和 TCP 握手端口。 +- 高并发下出现 RNR 或 `WR_FLUSH_ERR`：降低连接数、pipeline depth、队列深度，或增加内存上限。 +- 定位延迟：启用 benchmark 的 `--profile`，观察握手、post send、completion wait、RPC dispatch 和 attachment 阶段。 + +示例：`src/coro_rpc/examples/urma_example/urma_example.cpp`；实现：`include/ylt/coro_io/urma/`；公共 API 来自 UMDK 的 `urma_api.h`、`urma_types.h` 和 `urma_ubagg.h`。 diff --git a/include/ylt/coro_io/client_pool.hpp b/include/ylt/coro_io/client_pool.hpp index b12cbf75d..02f93ff9a 100644 --- a/include/ylt/coro_io/client_pool.hpp +++ b/include/ylt/coro_io/client_pool.hpp @@ -49,6 +49,7 @@ #include "coro_io.hpp" #include "detail/client_queue.hpp" #include "io_context_pool.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" #include "ylt/easylog.hpp" #include "ylt/util/atomic_shared_ptr.hpp" #ifdef YLT_ENABLE_IBV @@ -316,7 +317,13 @@ class client_pool : public std::enable_shared_from_this< ELOG_ERROR << "init client config failed."; co_return nullptr; } + auto connect_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; co_await reconnect(client, this->weak_from_this()); + coro_io::urma_benchmark_profile::record_since( + coro_io::urma_benchmark_profile::stage::client_connect_total, + connect_begin); } else { ELOG_TRACE << "get free client{" << client.get() << "}. from queue"; diff --git a/include/ylt/coro_io/coro_io.hpp b/include/ylt/coro_io/coro_io.hpp index ed639921b..54834f5ff 100644 --- a/include/ylt/coro_io/coro_io.hpp +++ b/include/ylt/coro_io/coro_io.hpp @@ -612,6 +612,7 @@ inline async_simple::coro::Lazy async_connect( co_return ec; } #endif + class period_timer : public asio::steady_timer { public: using asio::steady_timer::steady_timer; diff --git a/include/ylt/coro_io/data_view.hpp b/include/ylt/coro_io/data_view.hpp index a01332e8b..1a37ac19e 100644 --- a/include/ylt/coro_io/data_view.hpp +++ b/include/ylt/coro_io/data_view.hpp @@ -16,6 +16,7 @@ #pragma once #include #include +#include #include #include @@ -69,4 +70,20 @@ class data_view : public std::string_view { private: int gpu_id_; // GPU ID (-1 for CPU memory, >=0 for GPU memory) }; -} // namespace coro_io \ No newline at end of file + +struct owned_data_view { + data_view view; + std::shared_ptr owner; + + owned_data_view() = default; + owned_data_view(data_view view, std::shared_ptr owner) + : view(view), owner(std::move(owner)) {} + + bool empty() const noexcept { return view.empty(); } + const char* data() const noexcept { return view.data(); } + std::size_t size() const noexcept { return view.size(); } + int gpu_id() const noexcept { return view.gpu_id(); } + operator data_view() const noexcept { return view; } + operator std::string_view() const noexcept { return view; } +}; +} // namespace coro_io diff --git a/include/ylt/coro_io/detail/circle_buffer.hpp b/include/ylt/coro_io/detail/circle_buffer.hpp new file mode 100644 index 000000000..0d3db600c --- /dev/null +++ b/include/ylt/coro_io/detail/circle_buffer.hpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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 YLT_CORO_IO_DETAIL_CIRCLE_BUFFER_HPP +#define YLT_CORO_IO_DETAIL_CIRCLE_BUFFER_HPP + +#include +#include +#include + +namespace coro_io { +namespace detail { + +template +struct circle_buffer { + std::vector queue; + uint32_t front_ = 0, end_ = 0; + bool may_empty = true; + circle_buffer() = default; + circle_buffer(uint32_t size) { + assert(size > 0); + queue.resize(size); + } + void push(T&& elem) { + assert(!full()); + may_empty = false; + end_ = (end_ + 1) % queue.size(); + queue[end_] = std::move(elem); + } + T pop() { + assert(!empty()); + front_ = (front_ + 1) % queue.size(); + may_empty = true; + return std::move(queue[front_]); + } + T& back() { return queue[end_]; } + T& front() { return queue[(front_ + 1) % queue.size()]; } + bool full() const noexcept { return end_ == front_ && !may_empty; } + bool empty() const noexcept { return end_ == front_ && may_empty; } + std::size_t size() const noexcept { + if (front_ > end_) { + return queue.size() + end_ - front_; + } + else if (front_ == end_) { + return empty() ? 0 : queue.size(); + } + else { + return end_ - front_; + } + } +}; + +} // namespace detail +} // namespace coro_io + +#endif // YLT_CORO_IO_DETAIL_CIRCLE_BUFFER_HPP \ No newline at end of file diff --git a/include/ylt/coro_io/ibverbs/ib_socket.hpp b/include/ylt/coro_io/ibverbs/ib_socket.hpp index 1df71bd4e..4cb1482c8 100644 --- a/include/ylt/coro_io/ibverbs/ib_socket.hpp +++ b/include/ylt/coro_io/ibverbs/ib_socket.hpp @@ -37,6 +37,7 @@ #include "asio/ip/tcp.hpp" #include "asio/posix/stream_descriptor.hpp" #include "async_simple/Future.h" +#include "async_simple/Promise.h" #include "async_simple/Signal.h" #include "async_simple/coro/FutureAwaiter.h" #include "async_simple/coro/Lazy.h" @@ -45,6 +46,7 @@ #include "ib_device.hpp" #include "ib_error.hpp" #include "ylt/coro_io/coro_io.hpp" +#include "ylt/coro_io/detail/circle_buffer.hpp" #include "ylt/coro_io/ibverbs/ib_buffer.hpp" #include "ylt/coro_io/io_context_pool.hpp" #include "ylt/easylog.hpp" @@ -54,43 +56,6 @@ namespace coro_io { namespace detail { struct ib_socket_shared_state_t; -template -struct circle_buffer { - std::vector queue; - uint32_t front_ = 0, end_ = 0; - bool may_empty = true; - circle_buffer(uint32_t size) { - assert(size > 0); - queue.resize(size); - } - void push(T&& elem) { - assert(!full()); - may_empty = false; - end_ = (end_ + 1) % queue.size(); - queue[end_] = std::move(elem); - } - T pop() { - assert(!empty()); - front_ = (front_ + 1) % queue.size(); - may_empty = true; - return std::move(queue[front_]); - } - T& back() { return queue[end_]; } - T& front() { return queue[(front_ + 1) % queue.size()]; } - bool full() const noexcept { return end_ == front_ && !may_empty; } - bool empty() const noexcept { return end_ == front_ && may_empty; } - std::size_t size() const noexcept { - if (front_ > end_) { - return queue.size() + end_ - front_; - } - else if (front_ == end_) { - return empty() ? 0 : queue.size(); - } - else { - return end_ - front_; - } - } -}; struct ib_buffer_queue : public circle_buffer { using circle_buffer::circle_buffer; std::error_code post_recv_real(ibv_sge buffer, diff --git a/include/ylt/coro_io/server_acceptor.hpp b/include/ylt/coro_io/server_acceptor.hpp index 7e46bb5d0..f91f9eb98 100644 --- a/include/ylt/coro_io/server_acceptor.hpp +++ b/include/ylt/coro_io/server_acceptor.hpp @@ -126,10 +126,15 @@ struct tcp_server_acceptor : public server_acceptor_base { else { ELOG_ERROR << "accept error: " << error.message(); } + if (error == asio::error::operation_aborted || error == asio::error::bad_descriptor) { + ELOG_DEBUG << "accept stopped: " << error.message(); acceptor_close_waiter_.set_value(); } + else { + ELOG_ERROR << "accept error: " << error.message(); + } co_return ylt::expected{ ylt::unexpected{error}}; } diff --git a/include/ylt/coro_io/socket_wrapper.hpp b/include/ylt/coro_io/socket_wrapper.hpp index 97d7b306d..8c55a72e2 100644 --- a/include/ylt/coro_io/socket_wrapper.hpp +++ b/include/ylt/coro_io/socket_wrapper.hpp @@ -20,6 +20,11 @@ #include "ibverbs/ib_io.hpp" #include "ibverbs/ib_socket.hpp" #endif + +#ifdef YLT_ENABLE_URMA +#include "urma/urma_io.hpp" +#include "urma/urma_socket.hpp" +#endif #ifdef YLT_ENABLE_ND #include "ylt/coro_io/networkdirect/nd_io.hpp" #include "ylt/coro_io/networkdirect/nd_use_device.hpp" @@ -55,6 +60,16 @@ struct socket_wrapper_t { ib_socket_->prepare_accpet(std::move(soc)); } #endif + +#ifdef YLT_ENABLE_URMA + socket_wrapper_t(asio::ip::tcp::socket &&soc, + coro_io::ExecutorWrapper<> *executor, + const coro_io::urma_socket_t::config_t &config) + : executor_(executor), + urma_socket_(std::make_unique(executor_, config)) { + urma_socket_->prepare_accept(std::move(soc)); + } +#endif #ifdef YLT_ENABLE_ND socket_wrapper_t(coro_io::nd_socket_t &&soc, coro_io::ExecutorWrapper<> *executor) @@ -149,6 +164,44 @@ struct socket_wrapper_t { return true; } #endif +#ifdef YLT_ENABLE_URMA + bool init_client(const coro_io::urma_socket_t::config_t &config) { + ELOG_INFO << "URMA init_client: executor=" << executor_ + << ", device=" << config.device_name + << ", eid_index=" << config.eid_index + << ", tp_type=" << static_cast(config.tp_type) + << ", cq_size=" << config.cq_size + << ", recv_buffer_cnt=" << config.recv_buffer_cnt + << ", send_buffer_cnt=" << config.send_buffer_cnt + << ", buffer_size=" << config.buffer_size; + try { + init_tcp_socket(); + ELOG_DEBUG << "URMA init_client: TCP socket initialized"; + if (urma_socket_) { + ELOG_DEBUG << "URMA init_client: replacing existing URMA socket"; + *urma_socket_ = urma_socket_t(executor_, config); + ELOG_DEBUG << "URMA init_client: existing URMA socket replaced"; + } + else { + ELOG_DEBUG << "URMA init_client: constructing URMA socket"; + urma_socket_ = std::make_unique(executor_, config); + ELOG_DEBUG << "URMA init_client: URMA socket constructed"; + } + } catch (const std::system_error &e) { + ELOG_WARN << "init urma client failed: code=" << e.code().value() + << ", category=" << e.code().category().name() + << ", message=" << e.code().message() + << ", what=" << e.what(); + init_ok_ = false; + return false; + } catch (const std::exception &e) { + ELOG_WARN << "init urma client failed: " << e.what(); + init_ok_ = false; + return false; + } + return true; + } +#endif void set_local_ip(const std::string &local_ip) { local_ip_ = local_ip; } @@ -189,6 +242,10 @@ struct socket_wrapper_t { #ifdef YLT_ENABLE_IBV std::unique_ptr ib_socket_; #endif + +#ifdef YLT_ENABLE_URMA + std::unique_ptr urma_socket_; +#endif #ifdef YLT_ENABLE_ND std::unique_ptr nd_socket_; #endif @@ -209,6 +266,11 @@ struct socket_wrapper_t { return op(*ib_socket_); } #endif +#ifdef YLT_ENABLE_URMA + if (urma_socket_) { + return op(*urma_socket_); + } +#endif #ifdef YLT_ENABLE_ND if (nd_socket_) { return op(*nd_socket_); @@ -232,6 +294,12 @@ struct socket_wrapper_t { return; } #endif +#ifdef YLT_ENABLE_URMA + if (urma_socket_) { + urma_socket_->close(); + return; + } +#endif #ifdef YLT_ENABLE_ND if (nd_socket_) { nd_socket_->close(); @@ -260,6 +328,14 @@ struct socket_wrapper_t { coro_io::endpoint::rdma}; } #endif +#ifdef YLT_ENABLE_URMA + if (urma_socket_) { + return {urma_socket_->get_remote_address(), + urma_socket_->get_remote_qp_num(), + coro_io::endpoint::rdma}; + } +#endif + #ifdef YLT_ENABLE_ND if (nd_socket_) { return {nd_socket_->get_remote_address(), nd_socket_->get_remote_qp_num(), @@ -276,6 +352,14 @@ struct socket_wrapper_t { coro_io::endpoint::rdma}; } #endif + +#ifdef YLT_ENABLE_URMA + if (urma_socket_) { + return {urma_socket_->get_local_address(), + urma_socket_->get_local_qp_num(), + coro_io::endpoint::rdma}; + } +#endif #ifdef YLT_ENABLE_ND if (nd_socket_) { return {nd_socket_->get_local_address(), nd_socket_->get_local_qp_num(), @@ -311,6 +395,9 @@ struct socket_wrapper_t { #ifdef YLT_ENABLE_IBV using ibv_socket_t = coro_io::ib_socket_t; #endif +#ifdef YLT_ENABLE_URMA + using urma_socket_type = coro_io::urma_socket_t; +#endif #ifdef YLT_ENABLE_ND using nd_socket_t = coro_io::nd_socket_t; #endif diff --git a/include/ylt/coro_io/urma/urma_benchmark_profile.hpp b/include/ylt/coro_io/urma/urma_benchmark_profile.hpp new file mode 100644 index 000000000..8a9519008 --- /dev/null +++ b/include/ylt/coro_io/urma/urma_benchmark_profile.hpp @@ -0,0 +1,499 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace coro_io::urma_benchmark_profile { + +// Payload size buckets in bytes. +// 0: [0, 100) +// 1: [100, 500) +// 2: [500, 1K) +// 3..N: [1K, 2K), [2K, 3K), ... each 1K wide +inline constexpr std::size_t bucket_count = 16; +inline constexpr std::array bucket_upper = { + 100, 500, 1000, 2000, 3000, 4000, 5000, 6000, + 7000, 8000, 9000, 10000, 12000, 16000, 32000, UINT32_MAX, +}; +inline constexpr std::array bucket_names = { + "0-100", "100-500", "500-1K", "1K-2K", "2K-3K", + "3K-4K", "4K-5K", "5K-6K", "6K-7K", "7K-8K", + "8K-9K", "9K-10K", "10K-12K", "12K-16K", "16K-32K", + "32K+", +}; + +inline std::size_t size_to_bucket(std::size_t payload_size) { + for (std::size_t i = 0; i < bucket_count; ++i) { + if (payload_size < bucket_upper[i]) return i; + } + return bucket_count - 1; +} + +enum class stage : uint8_t { + benchmark_rpc_call = 0, + client_prepare_request, + client_send_request, + client_recv_header, + client_recv_payload, + client_deserialize_response, + server_read_header, + server_read_payload, + server_deserialize_request, + server_handler_execute, + server_dispatch, + server_serialize_response, + server_serialize_result, + server_response_queue, + server_send_response, + client_connect_total, + client_connect_tcp, + client_connect_handshake, + urma_write_total, + urma_write_copy, + urma_post_send, + urma_wait_send_completion, + urma_read_wait_completion, + urma_read_copy, + urma_read_view, + raw_client_write, + count +}; + +inline constexpr std::array(stage::count)> + stage_names = { + "benchmark.rpc_call", + "client.prepare_request", + "client.send_request", + "client.recv_header", + "client.recv_payload", + "client.deserialize_response", + "server.read_header", + "server.read_payload", + "server.deserialize_request", + "server.handler_execute", + "server.dispatch", + "server.serialize_response", + "server.serialize_result", + "server.response_queue", + "server.send_response", + "client.connect_total", + "client.connect_tcp", + "client.connect_handshake", + "urma.write_total", + "urma.write_copy", + "urma.post_send", + "urma.wait_send_completion", + "urma.read_wait_completion", + "urma.read_copy", + "urma.read_view", + "raw.client_write", +}; + +inline std::atomic& enabled_flag() { + static std::atomic value{false}; + return value; +} + +inline std::atomic& sample_rate_value() { + static std::atomic value{1}; + return value; +} + +inline void print(std::ostream& os); + +inline void init_from_env() { + static std::once_flag flag; + std::call_once(flag, [] { + const char* env = std::getenv("YLT_RPC_PROFILE_ENABLE"); + if (env && (std::string_view(env) == "1" || + std::string_view(env) == "on" || + std::string_view(env) == "true")) { + enabled_flag().store(true, std::memory_order_relaxed); + } + const char* rate = std::getenv("YLT_RPC_PROFILE_SAMPLE_RATE"); + if (rate && *rate) { + uint32_t r = 0; + for (const char* p = rate; *p >= '0' && *p <= '9'; ++p) + r = r * 10 + (*p - '0'); + if (r > 0) sample_rate_value().store(r, std::memory_order_relaxed); + } + if (enabled_flag().load(std::memory_order_relaxed)) { + std::atexit([]() { + std::fprintf(stderr, "[rpc_profile] atexit: printing profile\n"); + print(std::cerr); + }); + } + }); +} + +inline bool enabled() noexcept { + if (enabled_flag().load(std::memory_order_relaxed)) return true; + init_from_env(); + return enabled_flag().load(std::memory_order_relaxed); +} + +inline void configure(bool enabled, uint32_t sample_rate) noexcept { + enabled_flag().store(enabled, std::memory_order_relaxed); + sample_rate_value().store(std::max(sample_rate, 1), + std::memory_order_relaxed); +} + +inline uint64_t now_ns() noexcept { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +// Per-thread, per-stage, per-bucket samples. +struct thread_samples { + std::array, bucket_count>, + static_cast(stage::count)> + samples; + std::array(stage::count)> counters{}; + ~thread_samples(); +}; + +inline std::mutex& registry_mutex() { + static std::mutex value; + return value; +} + +inline std::array, bucket_count>, + static_cast(stage::count)>& +merged_samples() { + static std::array, bucket_count>, + static_cast(stage::count)> + value; + return value; +} + +inline std::array(stage::count)>& +merged_counters() { + static std::array(stage::count)> value{}; + return value; +} + +inline std::vector& registry() { + static std::vector value; + return value; +} + +inline thread_samples& local_samples() { + thread_local thread_samples value; + thread_local bool registered = [] { + std::lock_guard lock(registry_mutex()); + registry().push_back(&value); + return true; + }(); + (void)registered; + return value; +} + +inline thread_samples::~thread_samples() { + std::lock_guard lock(registry_mutex()); + auto& dst = merged_samples(); + auto& dst_cnt = merged_counters(); + for (std::size_t s = 0; s < samples.size(); ++s) { + dst_cnt[s] += counters[s]; + for (std::size_t b = 0; b < bucket_count; ++b) { + dst[s][b].insert(dst[s][b].end(), samples[s][b].begin(), + samples[s][b].end()); + samples[s][b].clear(); + } + } + auto& reg = registry(); + reg.erase(std::remove(reg.begin(), reg.end(), this), reg.end()); +} + +inline void record_with_size(stage stage_id, uint64_t duration_ns, + std::size_t payload_size) { + if (!enabled()) return; + auto& local = local_samples(); + auto s = static_cast(stage_id); + auto b = size_to_bucket(payload_size); + auto rate = sample_rate_value().load(std::memory_order_relaxed); + auto counter = ++local.counters[s]; + if ((counter % rate) != 0) return; + local.samples[s][b].push_back(duration_ns); +} + +inline void record(stage stage_id, uint64_t duration_ns) { + record_with_size(stage_id, duration_ns, 0); +} + +inline void record_since(stage stage_id, uint64_t begin_ns) { + if (!enabled()) return; + auto end_ns = now_ns(); + if (end_ns >= begin_ns) record(stage_id, end_ns - begin_ns); +} + +inline void record_since_with_size(stage stage_id, uint64_t begin_ns, + std::size_t payload_size) { + if (!enabled()) return; + auto end_ns = now_ns(); + if (end_ns >= begin_ns) + record_with_size(stage_id, end_ns - begin_ns, payload_size); +} + +inline void reserve_per_stage(std::size_t count) { + if (!enabled()) return; + auto& local = local_samples(); + for (auto& stage_buckets : local.samples) + for (auto& bucket : stage_buckets) bucket.reserve(count); +} + +inline void print(std::ostream& os) { + // merged[stage][bucket] -> samples + std::array, bucket_count>, + static_cast(stage::count)> + merged; + std::array(stage::count)> total_counters{}; + { + std::lock_guard lock(registry_mutex()); + for (std::size_t s = 0; s < merged.size(); ++s) { + total_counters[s] += merged_counters()[s]; + for (std::size_t b = 0; b < bucket_count; ++b) { + auto& src = merged_samples()[s][b]; + merged[s][b].insert(merged[s][b].end(), src.begin(), src.end()); + } + } + for (auto* thread : registry()) { + if (thread == nullptr) continue; + for (std::size_t s = 0; s < merged.size(); ++s) { + total_counters[s] += thread->counters[s]; + for (std::size_t b = 0; b < bucket_count; ++b) { + auto& src = thread->samples[s][b]; + merged[s][b].insert(merged[s][b].end(), src.begin(), src.end()); + } + } + } + } + + os << "rpc_profile sample_rate=" + << sample_rate_value().load(std::memory_order_relaxed) + << " unit=us\n"; + os << std::fixed << std::setprecision(2); + + // stats for one stage (merge all buckets) + struct stats { + std::size_t calls = 0, sampled = 0; + double avg = 0, p50 = 0, p90 = 0, p99 = 0, p999 = 0, max = 0; + }; + auto compute = [&](std::size_t s) -> stats { + stats st; + st.calls = total_counters[s]; + std::vector all; + for (std::size_t b = 0; b < bucket_count; ++b) { + auto& v = merged[s][b]; + all.insert(all.end(), v.begin(), v.end()); + } + if (all.empty()) return st; + st.sampled = all.size(); + std::sort(all.begin(), all.end()); + auto pct = [&](double p) { + return static_cast(all[static_cast((all.size()-1)*p)]) / 1000.0; + }; + auto sum = std::accumulate(all.begin(), all.end(), 0LL); + st.avg = static_cast(sum) / all.size() / 1000.0; + st.p50 = pct(0.50); st.p90 = pct(0.90); st.p99 = pct(0.99); + st.p999 = pct(0.999); st.max = static_cast(all.back()) / 1000.0; + return st; + }; + // stats for one bucket of one stage + auto compute_b = [&](std::size_t s, std::size_t b) -> stats { + stats st; + st.calls = total_counters[s]; + auto& v = merged[s][b]; + if (v.empty()) return st; + st.sampled = v.size(); + std::sort(v.begin(), v.end()); + auto pct = [&](double p) { + return static_cast(v[static_cast((v.size()-1)*p)]) / 1000.0; + }; + auto sum = std::accumulate(v.begin(), v.end(), 0LL); + st.avg = static_cast(sum) / v.size() / 1000.0; + st.p50 = pct(0.50); st.p90 = pct(0.90); st.p99 = pct(0.99); + st.p999 = pct(0.999); st.max = static_cast(v.back()) / 1000.0; + return st; + }; + (void)compute_b; // bucket detail available if needed later + + // Print a tree node: indent + branch char + label + stats + auto node = [&](int indent, bool last, std::string_view label, const stats& st) { + for (int i = 0; i < indent; ++i) os << (i < indent - 1 ? "│ " : " "); + if (indent > 0) os << (last ? "└─ " : "├─ "); + os << std::left << std::setw(28) << label + << " n=" << std::right << std::setw(8) << st.sampled + << " avg=" << std::setw(10) << st.avg + << " p50=" << std::setw(10) << st.p50 + << " p90=" << std::setw(10) << st.p90 + << " p99=" << std::setw(10) << st.p99 + << " p999=" << std::setw(10) << st.p999 + << " max=" << std::setw(10) << st.max + << "\n"; + }; + // Print a parent node (with children) + auto parent = [&](int indent, bool last, std::string_view label, const stats& st) { + for (int i = 0; i < indent; ++i) os << (i < indent - 1 ? "│ " : " "); + if (indent > 0) os << (last ? "└─ " : "├─ "); + os << std::left << std::setw(28) << label + << " n=" << std::right << std::setw(8) << st.sampled + << " avg=" << std::setw(10) << st.avg + << " p50=" << std::setw(10) << st.p50 + << " p99=" << std::setw(10) << st.p99 + << " max=" << std::setw(10) << st.max + << "\n"; + }; + auto has = [&](std::size_t s) { return total_counters[s] > 0; }; + auto st = [&](stage s) { return compute(static_cast(s)); }; + + // Print bucket breakdown for a stage under the current parent + auto buckets = [&](int indent, stage s) { + for (std::size_t b = 0; b < bucket_count; ++b) { + auto& v = merged[static_cast(s)][b]; + if (v.empty()) continue; + auto bst = compute_b(static_cast(s), b); + for (int i = 0; i < indent; ++i) os << " "; + os << " [" << bucket_names[b] << "]" + << " n=" << std::right << std::setw(8) << bst.sampled + << " avg=" << std::setw(10) << bst.avg + << " p50=" << std::setw(10) << bst.p50 + << " p99=" << std::setw(10) << bst.p99 + << " max=" << std::setw(10) << bst.max + << "\n"; + } + }; + + // ── Benchmark ── + if (has(static_cast(stage::benchmark_rpc_call))) { + os << "\n[benchmark]\n"; + parent(0, true, "rpc_call (total)", st(stage::benchmark_rpc_call)); + buckets(1, stage::benchmark_rpc_call); + } + + // ── Client ── + os << "\n[client]\n"; + if (has(static_cast(stage::client_send_request))) { + parent(0, false, "send_request", st(stage::client_send_request)); + buckets(1, stage::client_send_request); + if (has(static_cast(stage::client_prepare_request))) + node(1, !has(static_cast(stage::urma_write_total)), + "prepare_request", st(stage::client_prepare_request)); + if (has(static_cast(stage::urma_write_total))) { + parent(1, true, "urma.write_total", st(stage::urma_write_total)); + bool has_wc = has(static_cast(stage::urma_write_copy)); + bool has_ps = has(static_cast(stage::urma_post_send)); + bool has_wsc = has(static_cast(stage::urma_wait_send_completion)); + int cnt = has_wc + has_ps + has_wsc; + if (has_wc) node(2, --cnt == 0, "write_copy", st(stage::urma_write_copy)); + if (has_ps) node(2, --cnt == 0, "post_send", st(stage::urma_post_send)); + if (has_wsc) node(2, --cnt == 0, "wait_send_completion", st(stage::urma_wait_send_completion)); + } + } + if (has(static_cast(stage::client_recv_header))) { + parent(0, false, "recv_header", st(stage::client_recv_header)); + buckets(1, stage::client_recv_header); + bool has_rwc = has(static_cast(stage::urma_read_wait_completion)); + bool has_rc = has(static_cast(stage::urma_read_copy)); + int cnt = has_rwc + has_rc; + if (has_rwc) node(1, --cnt == 0, "urma.read_wait", st(stage::urma_read_wait_completion)); + if (has_rc) node(1, --cnt == 0, "urma.read_copy", st(stage::urma_read_copy)); + } + if (has(static_cast(stage::client_recv_payload))) { + parent(0, false, "recv_payload", st(stage::client_recv_payload)); + buckets(1, stage::client_recv_payload); + } + if (has(static_cast(stage::client_deserialize_response))) + parent(0, false, "deserialize_response", st(stage::client_deserialize_response)); + if (has(static_cast(stage::client_connect_total))) { + parent(0, true, "connect_total", st(stage::client_connect_total)); + bool has_tcp = has(static_cast(stage::client_connect_tcp)); + bool has_hs = has(static_cast(stage::client_connect_handshake)); + int cnt = has_tcp + has_hs; + if (has_tcp) node(1, --cnt == 0, "connect_tcp", st(stage::client_connect_tcp)); + if (has_hs) node(1, --cnt == 0, "connect_handshake", st(stage::client_connect_handshake)); + } + + // ── Server ── + bool has_server = has(static_cast(stage::server_read_header)) || + has(static_cast(stage::server_dispatch)); + if (has_server) { + os << "\n[server]\n"; + if (has(static_cast(stage::server_read_header))) { + parent(0, false, "read_header", st(stage::server_read_header)); + buckets(1, stage::server_read_header); + } + if (has(static_cast(stage::server_read_payload))) { + parent(0, false, "read_payload", st(stage::server_read_payload)); + buckets(1, stage::server_read_payload); + bool has_rwc = has(static_cast(stage::urma_read_wait_completion)); + bool has_rc = has(static_cast(stage::urma_read_copy)); + int cnt = has_rwc + has_rc; + if (has_rwc) node(1, --cnt == 0, "urma.read_wait", st(stage::urma_read_wait_completion)); + if (has_rc) node(1, --cnt == 0, "urma.read_copy", st(stage::urma_read_copy)); + } + if (has(static_cast(stage::server_dispatch))) { + parent(0, false, "dispatch", st(stage::server_dispatch)); + buckets(1, stage::server_dispatch); + bool has_deser = has(static_cast(stage::server_deserialize_request)); + bool has_exec = has(static_cast(stage::server_handler_execute)); + bool has_sres = has(static_cast(stage::server_serialize_result)); + int cnt = has_deser + has_exec + has_sres; + if (has_deser) node(1, --cnt == 0, "deserialize_request", st(stage::server_deserialize_request)); + if (has_exec) node(1, --cnt == 0, "handler_execute", st(stage::server_handler_execute)); + if (has_sres) node(1, --cnt == 0, "serialize_result", st(stage::server_serialize_result)); + } + if (has(static_cast(stage::server_serialize_response))) + parent(0, false, "serialize_response", st(stage::server_serialize_response)); + if (has(static_cast(stage::server_response_queue))) + parent(0, false, "response_queue", st(stage::server_response_queue)); + if (has(static_cast(stage::server_send_response))) { + parent(0, true, "send_response", st(stage::server_send_response)); + if (has(static_cast(stage::urma_write_total))) { + parent(1, true, "urma.write_total", st(stage::urma_write_total)); + bool has_wc = has(static_cast(stage::urma_write_copy)); + bool has_ps = has(static_cast(stage::urma_post_send)); + bool has_wsc = has(static_cast(stage::urma_wait_send_completion)); + int cnt = has_wc + has_ps + has_wsc; + if (has_wc) node(2, --cnt == 0, "write_copy", st(stage::urma_write_copy)); + if (has_ps) node(2, --cnt == 0, "post_send", st(stage::urma_post_send)); + if (has_wsc) node(2, --cnt == 0, "wait_send_completion", st(stage::urma_wait_send_completion)); + } + } + } + + if (has(static_cast(stage::urma_read_view))) { + os << "\n[urma]\n"; + parent(0, true, "read_view", st(stage::urma_read_view)); + } +} + +} // namespace coro_io::urma_benchmark_profile diff --git a/include/ylt/coro_io/urma/urma_buffer.hpp b/include/ylt/coro_io/urma/urma_buffer.hpp new file mode 100644 index 000000000..ac2930797 --- /dev/null +++ b/include/ylt/coro_io/urma/urma_buffer.hpp @@ -0,0 +1,307 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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 OF CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ylt/easylog.hpp" + +#ifdef YLT_ENABLE_URMA +#include +#endif + +namespace coro_io { + +struct urma_buffer_t { + void* addr = nullptr; + size_t length = 0; + void* seg = nullptr; // urma_target_seg_t* + + urma_buffer_t() = default; + urma_buffer_t(void* a, size_t l, void* s) : addr(a), length(l), seg(s) {} + explicit operator bool() const { return addr != nullptr && length > 0; } +}; + +class urma_buffer_pool_t { + public: + urma_buffer_pool_t(void* ctx, size_t buffer_size, size_t buffer_count); + ~urma_buffer_pool_t(); + + urma_buffer_pool_t(const urma_buffer_pool_t&) = delete; + urma_buffer_pool_t& operator=(const urma_buffer_pool_t&) = delete; + + urma_buffer_t get_buffer(int gpu_id = -1); + void return_buffer(urma_buffer_t& buffer); + + size_t buffer_size() const { return config_.buffer_size; } + size_t total_buffer_count() const { return config_.buffer_count; } + size_t total_memory_size() const { return allocation_size_; } + size_t free_buffer_count() const; +#ifdef YLT_ENABLE_URMA + urma_seg_t seg() const { + return seg_ ? reinterpret_cast(seg_)->seg : urma_seg_t{}; + } +#endif + size_t outstanding_buffer_count() const { + return outstanding_buffers_.load(std::memory_order_relaxed); + } + bool memory_out_of_limit() const { return free_buffer_count() == 0; } + void* context() const { return ctx_; } + + struct Config { + size_t buffer_size = 4 * 1024; + size_t buffer_count = 8; + int gpu_id = -1; + }; + const Config& get_config() const { return config_; } + + private: + bool init_buffers(); + + private: + void* ctx_ = nullptr; + Config config_; + void* base_addr_ = nullptr; + size_t allocation_size_ = 0; + void* seg_ = nullptr; // urma_target_seg_t* + std::vector buffers_; +#ifdef YLT_ENABLE_URMA + std::vector slice_segs_; +#endif + std::vector in_use_; + static constexpr size_t shard_count_ = 64; + size_t shard_for_index(size_t index) const noexcept { + return index % shard_count_; + } + size_t preferred_shard() const noexcept { + auto value = std::hash{}(std::this_thread::get_id()); + return value % shard_count_; + } + std::array, shard_count_> free_indices_; + mutable std::array mutexes_; + std::atomic outstanding_buffers_{0}; +}; + +// ============= Implementation (inline in header) ============= + +inline urma_buffer_pool_t::urma_buffer_pool_t(void* ctx, size_t buffer_size, + size_t buffer_count) + : ctx_(ctx), config_({buffer_size, buffer_count, -1}) { + init_buffers(); +} + +inline bool urma_buffer_pool_t::init_buffers() { +#ifdef YLT_ENABLE_URMA + auto init_begin = std::chrono::steady_clock::now(); + buffers_.reserve(config_.buffer_count); + + long page_size_value = ::sysconf(_SC_PAGESIZE); + size_t page_size = + page_size_value > 0 ? static_cast(page_size_value) : 4096; + if (config_.buffer_size == 0 || config_.buffer_count == 0) return false; + if (config_.buffer_count > + std::numeric_limits::max() / config_.buffer_size) { + ELOG_ERROR << "URMA buffer pool size overflow: buffer_size=" + << config_.buffer_size + << ", buffer_count=" << config_.buffer_count; + return false; + } + auto requested_size = config_.buffer_size * config_.buffer_count; + allocation_size_ = + (requested_size + page_size - 1) / page_size * page_size; + base_addr_ = ::mmap(nullptr, allocation_size_, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (base_addr_ == MAP_FAILED) { + auto error = std::error_code(errno, std::generic_category()); + ELOG_ERROR << "Failed to mmap page-aligned URMA buffer pool: errno=" + << errno << " (" << error.message() << ")" + << ", size=" << allocation_size_; + base_addr_ = nullptr; + allocation_size_ = 0; + 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; + flag.bs.token_id_valid = 0; + + urma_seg_cfg_t seg_cfg = {}; + seg_cfg.va = reinterpret_cast(base_addr_); + seg_cfg.len = allocation_size_; + seg_cfg.token_id = nullptr; + seg_cfg.token_value = {}; + seg_cfg.flag = flag; + seg_cfg.user_ctx = reinterpret_cast(base_addr_); + seg_cfg.iova = 0; + + errno = 0; + seg_ = urma_register_seg(reinterpret_cast(ctx_), &seg_cfg); + if (!seg_) { + auto error = std::error_code(errno, std::generic_category()); + ELOG_ERROR << "urma_register_seg failed: errno=" << errno << " (" + << error.message() << ")" + << ", address=" << base_addr_ << ", length=" << seg_cfg.len + << ", alignment=" << page_size << ", access=" << flag.bs.access; + ::munmap(base_addr_, allocation_size_); + base_addr_ = nullptr; + allocation_size_ = 0; + return false; + } + + auto* base = static_cast(base_addr_); + in_use_.assign(config_.buffer_count, 0); +#ifdef YLT_ENABLE_URMA + slice_segs_.resize(config_.buffer_count); +#endif + for (size_t i = 0; i < config_.buffer_count; ++i) { + auto* addr = base + i * config_.buffer_size; +#ifdef YLT_ENABLE_URMA + slice_segs_[i] = *static_cast(seg_); + slice_segs_[i].seg.ubva.va = reinterpret_cast(addr); + slice_segs_[i].seg.len = config_.buffer_size; + slice_segs_[i].user_ctx = reinterpret_cast(addr); + auto* buffer_seg = &slice_segs_[i]; +#else + auto* buffer_seg = seg_; +#endif + buffers_.push_back( + urma_buffer_t(addr, config_.buffer_size, buffer_seg)); + free_indices_[shard_for_index(i)].push(i); + } + + ELOG_INFO << "URMA buffer pool: " << buffers_.size() + << " buffers of " << config_.buffer_size + << " bytes, one registered segment of " << allocation_size_ + << " bytes, init_cost_us=" + << (std::chrono::steady_clock::now() - init_begin) / + std::chrono::microseconds(1); + return !buffers_.empty(); +#else + return false; +#endif +} + +inline urma_buffer_pool_t::~urma_buffer_pool_t() { +#ifdef YLT_ENABLE_URMA + std::array, shard_count_> locks; + for (size_t i = 0; i < shard_count_; ++i) { + locks[i] = std::unique_lock(mutexes_[i]); + } + if (seg_) { + urma_unregister_seg(static_cast(seg_)); + seg_ = nullptr; + } + if (base_addr_) { + ::munmap(base_addr_, allocation_size_); + base_addr_ = nullptr; + } + allocation_size_ = 0; + buffers_.clear(); +#ifdef YLT_ENABLE_URMA + slice_segs_.clear(); +#endif + in_use_.clear(); +#endif +} + +inline urma_buffer_t urma_buffer_pool_t::get_buffer(int gpu_id) { +#ifdef YLT_ENABLE_URMA + auto start = preferred_shard(); + for (size_t i = 0; i < shard_count_; ++i) { + auto shard = (start + i) % shard_count_; + std::lock_guard lock(mutexes_[shard]); + if (free_indices_[shard].empty()) continue; + size_t idx = free_indices_[shard].front(); + free_indices_[shard].pop(); + if (idx < in_use_.size()) in_use_[idx] = 1; + outstanding_buffers_++; + return buffers_[idx]; + } + ELOG_WARN << "URMA buffer pool out of buffers: total=" << buffers_.size() + << ", outstanding=" + << outstanding_buffers_.load(std::memory_order_relaxed) + << ", buffer_size=" << config_.buffer_size + << ", total_memory=" << allocation_size_; + return urma_buffer_t{}; +#else + return urma_buffer_t{}; +#endif +} + +inline void urma_buffer_pool_t::return_buffer(urma_buffer_t& buffer) { +#ifdef YLT_ENABLE_URMA + if (!buffer) return; + auto* base = static_cast(base_addr_); + auto* addr = static_cast(buffer.addr); + if (base && addr >= base && addr < base + allocation_size_) { + auto offset = static_cast(addr - base); + if (offset % config_.buffer_size == 0) { + auto index = offset / config_.buffer_size; + if (index < buffers_.size()) { + auto shard = shard_for_index(index); + std::lock_guard lock(mutexes_[shard]); + if (index < in_use_.size() && !in_use_[index]) { + ELOG_WARN << "return duplicated URMA buffer: " << buffer.addr + << ", index=" << index; + buffer = urma_buffer_t{}; + return; + } + if (index < in_use_.size()) in_use_[index] = 0; + free_indices_[shard].push(index); + if (outstanding_buffers_.load(std::memory_order_relaxed) > 0) + outstanding_buffers_--; + buffer = urma_buffer_t{}; + return; + } + } + } + ELOG_WARN << "return unknown URMA buffer: " << buffer.addr; +#endif +} + +inline size_t urma_buffer_pool_t::free_buffer_count() const { +#ifdef YLT_ENABLE_URMA + size_t total = 0; + for (size_t i = 0; i < shard_count_; ++i) { + std::lock_guard lock(mutexes_[i]); + total += free_indices_[i].size(); + } + return total; +#else + return 0; +#endif +} + +} // namespace coro_io diff --git a/include/ylt/coro_io/urma/urma_device.hpp b/include/ylt/coro_io/urma/urma_device.hpp new file mode 100644 index 000000000..cda91a5cb --- /dev/null +++ b/include/ylt/coro_io/urma/urma_device.hpp @@ -0,0 +1,431 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "ylt/easylog.hpp" +#include "ylt/coro_io/urma/urma_buffer.hpp" + +#include + +#ifdef YLT_ENABLE_URMA +#include +#endif + +namespace coro_io { + +class urma_buffer_pool_t; + +class urma_device_wrapper_t { + public: + urma_device_wrapper_t(); + ~urma_device_wrapper_t(); + + urma_device_wrapper_t(const urma_device_wrapper_t&) = delete; + urma_device_wrapper_t& operator=(const urma_device_wrapper_t&) = delete; + + bool init(const std::string& device_name, int eid_index = 0); + bool configure_buffer_pool(std::size_t buffer_size, + std::size_t max_memory_usage); + void close(); + + urma_context_t* context() const { return context_; } + urma_device_t* device() const { return device_ptr_; } + const std::string& name() const { return name_; } + int eid_index() const { return eid_index_; } + const urma_eid_t& eid() const { return eid_; } + const urma_device_attr_t& attr() const { return device_attr_; } + uint32_t uasid() const { return context_ ? context_->uasid : 0; } + uint32_t max_jetty() const { return device_attr_.dev_cap.max_jetty; } + uint32_t max_jfc() const { return device_attr_.dev_cap.max_jfc; } + bool supports_rm_rtp() const { + return device_attr_.dev_cap.rm_tp_cap.bs.rtp != 0; + } + bool supports_rm_ctp() const { + return device_attr_.dev_cap.rm_tp_cap.bs.ctp != 0; + } + + std::string eid_string() const; + asio::ip::address gid_address() const; + bool is_valid() const { return context_ != nullptr && device_ptr_ != nullptr; } + std::shared_ptr get_buffer_pool() const { return buffer_pool_; } + + void set_bonding_config(uint32_t mode, uint32_t level) { + bond_mode_ = mode; + bond_level_ = level; + } + + private: + std::string name_; + int eid_index_ = -1; + ::urma_device_t* device_ptr_ = nullptr; + urma_context_t* context_ = nullptr; + urma_eid_t eid_{}; + urma_device_attr_t device_attr_{}; + std::shared_ptr buffer_pool_; + uint32_t bond_mode_ = BONDP_BONDING_MODE_STANDALONE; + uint32_t bond_level_ = BONDP_BONDING_LEVEL_IODIE; +}; + +using urma_device_t = urma_device_wrapper_t; + +struct urma_buffer_pool_config_t { + uint32_t buffer_size = 4 * 1024; // buffer size + uint64_t max_memory_usage = 20 * 1024 * 1024; // max memory usage + std::chrono::seconds idle_timeout{5}; // idle timeout +}; + +struct urma_init_config_t { + std::string dev_name; // device name, empty for auto-select + urma_buffer_pool_config_t buffer_pool_config; // buffer pool config + int eid_index = 0; // EID index to use + uint32_t bond_mode = BONDP_BONDING_MODE_STANDALONE; + uint32_t bond_level = BONDP_BONDING_LEVEL_IODIE; +}; + +class urma_device_manager { + public: + static urma_device_manager& instance(); + bool init(); + std::shared_ptr get_device(const std::string& device_name = "", int eid_index = 0); + std::shared_ptr get_device(const urma_init_config_t& config); + std::vector> get_all_devices(); + std::shared_ptr get_global_device(); + + private: + urma_device_manager() = default; + ~urma_device_manager(); + bool initialized_ = false; + std::vector> devices_; + std::shared_ptr global_device_; +}; + +inline std::shared_ptr get_global_urma_device() { + ELOG_DEBUG << "get_global_urma_device() called"; + return urma_device_manager::instance().get_global_device(); +} + +inline std::shared_ptr get_global_urma_device( + const urma_init_config_t& config) { + ELOG_DEBUG << "get_global_urma_device(dev_name=" << config.dev_name + << ", eid_index=" << config.eid_index + << ", bond_mode=" << config.bond_mode + << ", bond_level=" << config.bond_level << ") called"; + auto device = urma_device_manager::instance().get_device(config); + if (device) { + device->configure_buffer_pool(config.buffer_pool_config.buffer_size, + config.buffer_pool_config.max_memory_usage); + } + return device; +} + +inline urma_device_wrapper_t::urma_device_wrapper_t() = default; + +inline urma_device_wrapper_t::~urma_device_wrapper_t() { close(); } + +inline bool urma_device_wrapper_t::configure_buffer_pool( + std::size_t buffer_size, std::size_t max_memory_usage) { + if (!context_ || buffer_size == 0) return false; + auto requested_buffer_count = + std::max(1, max_memory_usage / buffer_size); + if (buffer_pool_) { + if (buffer_pool_->buffer_size() == buffer_size && + buffer_pool_->total_buffer_count() >= requested_buffer_count) { + return true; + } + if (buffer_pool_->free_buffer_count() != + buffer_pool_->total_buffer_count()) { + ELOG_WARN << "cannot resize an in-use URMA buffer pool: current_size=" + << buffer_pool_->buffer_size() + << ", current_count=" << buffer_pool_->total_buffer_count() + << ", requested_size=" << buffer_size + << ", requested_count=" << requested_buffer_count; + return false; + } + } + buffer_pool_ = + std::make_shared(context_, buffer_size, + requested_buffer_count); + return buffer_pool_->total_buffer_count() != 0; +} + +inline bool urma_device_wrapper_t::init(const std::string& device_name, int eid_index) { +#ifdef YLT_ENABLE_URMA + name_ = device_name; + eid_index_ = eid_index; + + int num_devices = 0; + ::urma_device_t** devices = urma_get_device_list(&num_devices); + if (!devices || num_devices <= 0) { + ELOG_ERROR << "urma_get_device_list failed"; + return false; + } + + ::urma_device_t* found_device = nullptr; + for (int i = 0; i < num_devices; ++i) { + if (device_name.empty() || std::string(devices[i]->name) == device_name) { + found_device = devices[i]; + break; + } + } + + if (!found_device) { + ELOG_ERROR << "URMA device not found: " << device_name; + urma_free_device_list(devices); + return false; + } + + device_ptr_ = found_device; + name_ = found_device->name; + + uint32_t eid_cnt = 0; + urma_eid_info_t* eid_list = urma_get_eid_list(device_ptr_, &eid_cnt); + if (!eid_list || eid_cnt == 0) { + ELOG_ERROR << "urma_get_eid_list failed"; + urma_free_device_list(devices); + return false; + } + + int eid_slot = 0; + bool found_eid_index = false; + for (uint32_t i = 0; i < eid_cnt; ++i) { + if (eid_list[i].eid_index == static_cast(eid_index)) { + eid_slot = static_cast(i); + found_eid_index = true; + break; + } + } + if (!found_eid_index) { + ELOG_WARN << "URMA EID index " << eid_index + << " was not found; using index " << eid_list[0].eid_index; + } + eid_ = eid_list[eid_slot].eid; + eid_index_ = static_cast(eid_list[eid_slot].eid_index); + urma_free_eid_list(eid_list); + + context_ = urma_create_context(device_ptr_, eid_index_); + if (!context_) { + ELOG_ERROR << "urma_create_context failed"; + urma_free_device_list(devices); + return false; + } + + // Bonding devices default to STANDALONE+PORT level (bondp_provider_ops.c:514), + // which enables multiple port-EID devices. Under CTP the hardware sprays sends + // across the bonding group, but schedule_recv_standalone posts recv WRs to a + // single port, causing first-SEND RNR (status=10). Set STANDALONE+IODIE (single + // primary-EID device, matching urma_perftest) to avoid the spray/recv mismatch. + // Must run while context refcount==1 (before urma_register_seg in + // configure_buffer_pool), else bondp_set_bonding_mode returns URMA_EAGAIN. + if (name_.compare(0, 7, "bonding") == 0) { + bondp_set_bonding_mode_in_t bond_in = { + .bonding_mode = static_cast(bond_mode_), + .bonding_level = static_cast(bond_level_), + }; + urma_user_ctl_in_t ctl_in = { + .addr = reinterpret_cast(&bond_in), + .len = static_cast(sizeof(bond_in)), + .opcode = BONDP_USER_CTL_SET_BONDING_MODE, + }; + urma_user_ctl_out_t ctl_out = {}; + urma_status_t st = urma_user_ctl(context_, &ctl_in, &ctl_out); + if (st != URMA_SUCCESS) { + ELOG_ERROR << "urma_user_ctl(SET_BONDING_MODE) failed: status=" << st + << ", dev=" << name_ + << " (requires context refcount==1, before any resource " + "creation; check that no seg/jfc/jfr is created first)"; + urma_delete_context(context_); + context_ = nullptr; + urma_free_device_list(devices); + return false; + } + ELOG_INFO << "bonding mode set: mode=" << bond_mode_ + << ", level=" << bond_level_ << ", dev=" << name_; + } + + if (urma_query_device(device_ptr_, &device_attr_) != 0) { + ELOG_ERROR << "urma_query_device failed"; + urma_delete_context(context_); + context_ = nullptr; + urma_free_device_list(devices); + return false; + } + + urma_free_device_list(devices); + auto default_pool_config = urma_buffer_pool_config_t{}; + if (!configure_buffer_pool(default_pool_config.buffer_size, + default_pool_config.max_memory_usage)) { + close(); + return false; + } + ELOG_INFO << "URMA device: " << name_ << ", EID: " << eid_string(); + return true; +#else + ELOG_WARN << "URMA not enabled"; + return false; +#endif +} + +inline void urma_device_wrapper_t::close() { +#ifdef YLT_ENABLE_URMA + if (context_) { + buffer_pool_.reset(); + urma_delete_context(context_); + context_ = nullptr; + } +#endif +} + +inline std::string urma_device_wrapper_t::eid_string() const { + char buf[64] = {0}; + snprintf(buf, sizeof(buf), + "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:" + "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x", + eid_.raw[0], eid_.raw[1], eid_.raw[2], eid_.raw[3], + eid_.raw[4], eid_.raw[5], eid_.raw[6], eid_.raw[7], + eid_.raw[8], eid_.raw[9], eid_.raw[10], eid_.raw[11], + eid_.raw[12], eid_.raw[13], eid_.raw[14], eid_.raw[15]); + return std::string(buf); +} + +inline asio::ip::address urma_device_wrapper_t::gid_address() const { + char buf[64]; + snprintf(buf, sizeof(buf), "%d.%d.%d.%d", + eid_.raw[0], eid_.raw[1], eid_.raw[2], eid_.raw[3]); + std::error_code ec; + auto addr = asio::ip::make_address(buf, ec); + if (ec) return asio::ip::make_address_v4(0x7F000001); + return addr; +} + +// urma_device_manager +inline urma_device_manager& urma_device_manager::instance() { + static urma_device_manager inst; + return inst; +} + +inline urma_device_manager::~urma_device_manager() { + devices_.clear(); + global_device_.reset(); +} + +inline bool urma_device_manager::init() { +#ifdef YLT_ENABLE_URMA + if (initialized_) return true; + urma_init_attr_t init_attr = {}; + auto status = urma_init(&init_attr); + if (status != URMA_SUCCESS && status != URMA_EEXIST) { + ELOG_WARN << "urma_init returned " << status + << ", attempting to continue (provider may already be loaded)"; + } + // Even if urma_init returns URMA_FAIL (no provider opened), the device list + // may already be populated by a prior caller (e.g. Mooncake TransferEngine) + // or by liburma's constructor. Try to proceed rather than failing hard. + initialized_ = true; + return true; +#else + return false; +#endif +} + +inline std::shared_ptr urma_device_manager::get_device( + const std::string& device_name, int eid_index) { +#ifdef YLT_ENABLE_URMA + if (!initialized_ && !init()) return nullptr; + + for (auto& dev : devices_) { + if ((device_name.empty() || dev->name() == device_name) && + dev->eid_index() == eid_index) { + return dev; + } + } + + auto dev = std::make_shared(); + if (!dev->init(device_name, eid_index)) return nullptr; + + devices_.push_back(dev); + if (!global_device_) global_device_ = dev; + return dev; +#else + return nullptr; +#endif +} + +inline std::shared_ptr urma_device_manager::get_device( + const urma_init_config_t& config) { +#ifdef YLT_ENABLE_URMA + if (!initialized_ && !init()) return nullptr; + + // Reuse an existing device matching dev_name + eid_index if one was already + // created. Note: the bond_mode/bond_level of the existing device are kept + // (they were applied at its init() time and cannot be changed post-hoc). + for (auto& dev : devices_) { + if ((config.dev_name.empty() || dev->name() == config.dev_name) && + dev->eid_index() == config.eid_index) { + return dev; + } + } + + auto dev = std::make_shared(); + // Bonding config must be set before init() so the SET_BONDING_MODE ioctl + // runs while the context refcount is still 1. + dev->set_bonding_config(config.bond_mode, config.bond_level); + if (!dev->init(config.dev_name, config.eid_index)) return nullptr; + + devices_.push_back(dev); + if (!global_device_) global_device_ = dev; + return dev; +#else + return nullptr; +#endif +} + +inline std::vector> +urma_device_manager::get_all_devices() { +#ifdef YLT_ENABLE_URMA + if (!devices_.empty()) return devices_; + if (!initialized_) init(); + + int num_devices = 0; + ::urma_device_t** urma_dev_list = urma_get_device_list(&num_devices); + if (!urma_dev_list || num_devices <= 0) return devices_; + + for (int i = 0; i < num_devices; ++i) { + auto dev = std::make_shared(); + if (dev->init(urma_dev_list[i]->name)) { + devices_.push_back(dev); + if (!global_device_) global_device_ = dev; + } + } + urma_free_device_list(urma_dev_list); + return devices_; +#else + return {}; +#endif +} + +inline std::shared_ptr urma_device_manager::get_global_device() { + if (!global_device_) get_all_devices(); + return global_device_; +} + +} // namespace coro_io diff --git a/include/ylt/coro_io/urma/urma_io.hpp b/include/ylt/coro_io/urma/urma_io.hpp new file mode 100644 index 000000000..27763f9c0 --- /dev/null +++ b/include/ylt/coro_io/urma/urma_io.hpp @@ -0,0 +1,268 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "async_simple/Promise.h" +#include "async_simple/util/move_only_function.h" +#include "ylt/coro_io/coro_io.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" +#include "ylt/coro_io/urma/urma_socket.hpp" + +namespace coro_io { +namespace detail { + +template +void make_urma_buffers(std::vector& result, + BufferSequence& buffers) { + auto first = asio::buffer_sequence_begin(buffers); + auto last = asio::buffer_sequence_end(buffers); + for (; first != last; ++first) result.emplace_back(*first); +} + +struct urma_write_completion_state { + std::queue> completions; + async_simple::util::move_only_function resume_handler; + + void push(std::pair result) { + completions.push(result); + if (resume_handler) { + auto h = std::move(resume_handler); + resume_handler = nullptr; + h(); + } + } +}; + +inline async_simple::coro::Lazy> +wait_urma_write_completion( + const std::shared_ptr& state, + urma_socket_t& socket) { + // Suspend until event_loop polls the CQ and calls state->push. + // Do NOT call socket.poll_completion_once() here - it would race with + // event_loop's poll_completion on the same shared state (send_callbacks_, + // recv_queue_, etc.), causing double-free / SIGBUS. + while (state->completions.empty()) { + callback_awaitor awaitor; + co_await awaitor.await_resume([&state](auto handler) { + state->resume_handler = [handler]() mutable { + handler.resume(); + }; + }); + } + auto result = state->completions.front(); + state->completions.pop(); + co_return result; +} + +template +async_simple::coro::Lazy> +async_urma_read(urma_socket_t& socket, Buffer&& raw_buffer, bool read_some) { + if (!socket.get_executor().running_in_this_thread()) + co_await dispatch(socket.get_executor()); + + std::vector buffers; + make_urma_buffers(buffers, raw_buffer); + std::size_t completed = 0; + for (auto& buffer : buffers) { + if (socket.remain_read_buffer_size()) { + auto count = + socket.consume(static_cast(buffer.data()), buffer.size()); + buffer += count; + completed += count; + } + while (buffer.size()) { + auto wait_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + auto [ec, length] = + co_await async_io>( + [&](auto&& callback) { + socket.post_recv(std::move(callback)); + }, + socket); + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_read_wait_completion, + wait_begin, length); + if (ec) co_return std::pair{ec, completed}; + auto recv = socket.get_recv_buffer(); + auto count = std::min(length, buffer.size()); + auto copy_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + std::memcpy(buffer.data(), reinterpret_cast(recv.addr), count); + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_read_copy, copy_begin, count); + buffer += count; + completed += count; + socket.set_read_buffer_len(count, length - count); + if (read_some) co_return std::pair{std::error_code{}, completed}; + } + } + co_return std::pair{std::error_code{}, completed}; +} + +inline async_simple::coro::Lazy< + std::pair>> +async_urma_read_views(urma_socket_t& socket, std::size_t size) { + if (!socket.get_executor().running_in_this_thread()) + co_await dispatch(socket.get_executor()); + + std::vector views; + std::size_t completed = 0; + if (socket.remain_read_buffer_size()) { + auto view = socket.detach_remain_data_view(); + if (!view.empty()) { + completed += view.size(); + views.push_back(std::move(view)); + } + } + while (completed < size) { + auto wait_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + auto [ec, length] = + co_await async_io>( + [&](auto&& callback) { + socket.post_recv(std::move(callback)); + }, + socket); + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_read_wait_completion, wait_begin, length); + if (ec) co_return std::pair{ec, std::move(views)}; + if (completed + length > size) { + ELOG_ERROR << "URMA read view received more data than requested: " + << "requested=" << size << ", completed=" << completed + << ", incoming=" << length; + co_return std::pair{std::make_error_code(std::errc::protocol_error), + std::move(views)}; + } + auto view_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + auto view = socket.detach_recv_buffer_view(length); + if (!view.empty()) { + completed += view.size(); + views.push_back(std::move(view)); + } + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_read_view, view_begin, length); + } + co_return std::pair{std::error_code{}, std::move(views)}; +} + +} // namespace detail + +template +async_simple::coro::Lazy> async_write( + urma_socket_t& socket, Buffer&& raw_buffer) { + if (!socket.get_executor().running_in_this_thread()) + co_await dispatch(socket.get_executor()); + + std::vector buffers; + detail::make_urma_buffers(buffers, raw_buffer); + std::size_t total_size = 0; + for (auto& item : buffers) total_size += item.size(); + std::size_t completed = 0; + auto total_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + ELOG_DEBUG << "URMA async_write start: total_size=" << total_size + << ", chunk_size=" << socket.get_buffer_size() + << ", send_window=" << socket.get_send_window_size(); + auto state = std::make_shared(); + std::size_t in_flight = 0; + auto post_next = [&]() -> std::pair { + if (buffers.empty()) return {{}, false}; + auto buffer = socket.get_send_buffer(); + if (!buffer) + return {std::make_error_code(std::errc::no_buffer_space), false}; + std::size_t length = 0; + auto copy_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + while (!buffers.empty() && length < socket.get_buffer_size()) { + auto count = std::min( + buffers.front().size(), socket.get_buffer_size() - length); + std::memcpy(static_cast(buffer.addr) + length, + buffers.front().data(), count); + length += count; + buffers.front() += count; + if (buffers.front().size() == 0) buffers.erase(buffers.begin()); + } + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_write_copy, copy_begin, length); + auto post_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + socket.post_send(std::move(buffer), length, + [state](std::pair result) { + state->push(result); + }); + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_post_send, post_begin, length); + ++in_flight; + return {{}, true}; + }; + + const auto send_window = std::max(socket.get_send_window_size(), 1); + while (!buffers.empty() || in_flight != 0) { + while (!buffers.empty() && in_flight < send_window) { + auto [ec, posted] = post_next(); + if (ec) { + if (in_flight == 0) co_return std::pair{ec, completed}; + break; + } + if (!posted) break; + } + if (in_flight == 0) continue; + auto wait_begin = urma_benchmark_profile::enabled() + ? urma_benchmark_profile::now_ns() + : 0; + auto result = co_await detail::wait_urma_write_completion(state, socket); + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_wait_send_completion, + wait_begin, result.second); + --in_flight; + if (result.first) co_return std::pair{result.first, completed}; + completed += result.second; + } + ELOG_DEBUG << "URMA async_write done: total_size=" << total_size + << ", completed=" << completed; + urma_benchmark_profile::record_since_with_size( + urma_benchmark_profile::stage::urma_write_total, total_begin, total_size); + co_return std::pair{std::error_code{}, completed}; +} + +template +async_simple::coro::Lazy> async_read( + urma_socket_t& socket, Buffer&& buffer) { + return detail::async_urma_read(socket, std::forward(buffer), false); +} + +template +async_simple::coro::Lazy> +async_read_some(urma_socket_t& socket, Buffer&& buffer) { + return detail::async_urma_read(socket, std::forward(buffer), true); +} + +} // namespace coro_io diff --git a/include/ylt/coro_io/urma/urma_rpc_env.hpp b/include/ylt/coro_io/urma/urma_rpc_env.hpp new file mode 100644 index 000000000..1dfdf7fb3 --- /dev/null +++ b/include/ylt/coro_io/urma/urma_rpc_env.hpp @@ -0,0 +1,194 @@ +/* + * Copyright (c) 2026, Alibaba Group Holding Limited; + * + * Licensed 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. + */ +#pragma once + +#ifdef YLT_ENABLE_URMA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ylt/coro_io/urma/urma_socket.hpp" +#include "ylt/easylog.hpp" + +namespace coro_io::detail { + +inline std::string urma_rpc_lower_ascii(std::string_view value) { + std::string result(value); + std::transform(result.begin(), result.end(), result.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return result; +} + +inline const char* urma_rpc_getenv(const char* name) { return std::getenv(name); } + +inline bool urma_rpc_env_enabled() { + const char* value = urma_rpc_getenv("URMA_RPC_ENABLE"); + if (value == nullptr || *value == '\0') return false; + + auto normalized = urma_rpc_lower_ascii(value); + return normalized == "1" || normalized == "on" || normalized == "true" || + normalized == "yes"; +} + +inline bool urma_rpc_env_flag(const char* name, bool default_value) { + const char* value = urma_rpc_getenv(name); + if (value == nullptr || *value == '\0') return default_value; + auto normalized = urma_rpc_lower_ascii(value); + if (normalized == "1" || normalized == "on" || normalized == "true" || + normalized == "yes") + return true; + if (normalized == "0" || normalized == "off" || normalized == "false" || + normalized == "no") + return false; + ELOG_WARN << "invalid " << name << " value: " << value + << "; use default " << (default_value ? "true" : "false"); + return default_value; +} + +template +inline bool urma_rpc_parse_integer(std::string_view value, T& output) { + static_assert(std::is_integral_v); + T parsed{}; + const auto* first = value.data(); + const auto* last = value.data() + value.size(); + auto [ptr, ec] = std::from_chars(first, last, parsed); + if (ec != std::errc{} || ptr != last) return false; + output = parsed; + return true; +} + +template +inline void urma_rpc_parse_env_integer(const char* name, T& field) { + const char* value = urma_rpc_getenv(name); + if (value == nullptr || *value == '\0') return; + + T parsed{}; + if (!urma_rpc_parse_integer(std::string_view(value), parsed)) { + ELOG_WARN << "invalid " << name << " value: " << value + << "; use default " << field; + return; + } + field = parsed; +} + +inline void urma_rpc_parse_env_tp_type(urma_tp_type_t& tp_type) { + const char* value = urma_rpc_getenv("URMA_RPC_TP_TYPE"); + if (value == nullptr || *value == '\0') return; + + auto normalized = urma_rpc_lower_ascii(value); + if (normalized == "ctp" || normalized == "0") { + tp_type = URMA_CTP; + return; + } + if (normalized == "rtp" || normalized == "1") { + tp_type = URMA_RTP; + return; + } + + ELOG_WARN << "invalid URMA_RPC_TP_TYPE value: " << value + << "; use default " << static_cast(tp_type); +} + +inline coro_io::urma_socket_t::config_t make_urma_rpc_config_from_env() { + coro_io::urma_socket_t::config_t config{}; + + if (const char* device = urma_rpc_getenv("URMA_RPC_DEVICE"); + device != nullptr) { + config.device_name = device; + } + + urma_rpc_parse_env_integer("URMA_RPC_EID_INDEX", config.eid_index); + urma_rpc_parse_env_integer("URMA_RPC_CQ_SIZE", config.cq_size); + urma_rpc_parse_env_integer("URMA_RPC_RECV_BUFFER_CNT", + config.recv_buffer_cnt); + urma_rpc_parse_env_integer("URMA_RPC_SEND_BUFFER_CNT", + config.send_buffer_cnt); + urma_rpc_parse_env_integer("URMA_RPC_BUFFER_SIZE", config.buffer_size); + urma_rpc_parse_env_integer("URMA_RPC_MAX_MEMORY_USAGE", + config.max_memory_usage); + urma_rpc_parse_env_tp_type(config.tp_type); + config.event_mode = + urma_rpc_env_flag("URMA_RPC_EVENT_MODE", /*default*/ true); + urma_rpc_parse_env_integer("URMA_RPC_BUSY_POLL_BUDGET", + config.busy_poll_budget); + { + uint64_t interval_us = 5; + urma_rpc_parse_env_integer("URMA_RPC_POLL_INTERVAL", interval_us); + config.poll_interval = std::chrono::microseconds(interval_us); + } + + return config; +} + +inline std::optional +probe_urma_rpc_config(coro_io::urma_socket_t::config_t config) { + try { + auto device = coro_io::get_global_urma_device( + {.dev_name = config.device_name, + .buffer_pool_config = {.buffer_size = config.buffer_size, + .max_memory_usage = config.max_memory_usage}, + .eid_index = config.eid_index}); + if (!device || !device->is_valid() || !device->get_buffer_pool()) { + ELOG_WARN << "URMA_RPC_ENABLE is enabled but no usable URMA device was " + "found; fall back to TCP"; + return std::nullopt; + } + + config.device_name = device->name(); + config.eid_index = device->eid_index(); + ELOG_INFO << "URMA RPC auto enable succeeded: device=" + << config.device_name << ", eid_index=" << config.eid_index + << ", tp_type=" << static_cast(config.tp_type) + << ", cq_size=" << config.cq_size + << ", recv_buffer_cnt=" << config.recv_buffer_cnt + << ", send_buffer_cnt=" << config.send_buffer_cnt + << ", buffer_size=" << config.buffer_size + << ", max_memory_usage=" << config.max_memory_usage + << ", event_mode=" << (config.event_mode ? "on" : "off") + << ", busy_poll_budget=" << config.busy_poll_budget; + return config; + } catch (const std::exception& e) { + ELOG_WARN << "URMA RPC auto enable failed: " << e.what() + << "; fall back to TCP"; + return std::nullopt; + } +} + +} // namespace coro_io::detail + +namespace coro_io { + +inline std::optional +try_make_urma_rpc_config() { + if (!detail::urma_rpc_env_enabled()) return std::nullopt; + + static const auto cached_config = []() { + return detail::probe_urma_rpc_config( + detail::make_urma_rpc_config_from_env()); + }(); + return cached_config; +} + +} // namespace coro_io +#endif diff --git a/include/ylt/coro_io/urma/urma_socket.hpp b/include/ylt/coro_io/urma/urma_socket.hpp new file mode 100644 index 000000000..a9c473804 --- /dev/null +++ b/include/ylt/coro_io/urma/urma_socket.hpp @@ -0,0 +1,1159 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "asio/dispatch.hpp" +#include "asio/ip/address.hpp" +#include "asio/ip/tcp.hpp" +#include "asio/posix/stream_descriptor.hpp" +#include "asio/steady_timer.hpp" +#include "async_simple/Future.h" +#include "async_simple/Promise.h" +#include "async_simple/coro/FutureAwaiter.h" +#include "async_simple/coro/Lazy.h" +#include "async_simple/util/move_only_function.h" +#include "ylt/coro_io/coro_io.hpp" +#include "ylt/coro_io/data_view.hpp" +#include "ylt/coro_io/detail/circle_buffer.hpp" +#include "ylt/coro_io/urma/urma_buffer.hpp" +#include "ylt/coro_io/urma/urma_device.hpp" +#include "ylt/easylog.hpp" +#include "ylt/struct_pack.hpp" +#include +#include + +#include +#include + +namespace coro_io { +namespace detail { + +struct urma_recv_buffer_owner { + urma_recv_buffer_owner(std::shared_ptr pool, + urma_buffer_t buffer) + : pool(std::move(pool)), buffer(std::move(buffer)) {} + ~urma_recv_buffer_owner() { + if (pool && buffer) pool->return_buffer(buffer); + } + + std::shared_ptr pool; + urma_buffer_t buffer; +}; + +inline std::error_code make_urma_error(int status) { + if (status == URMA_SUCCESS) return {}; + return std::error_code(std::abs(status), std::generic_category()); +} + +struct urma_deleter { + void operator()(urma_jfc_t* value) const { + if (value) urma_delete_jfc(value); + } + void operator()(urma_jfr_t* value) const { + if (value) urma_delete_jfr(value); + } + void operator()(urma_jetty_t* value) const { + if (value) urma_delete_jetty(value); + } + void operator()(urma_jfce_t* value) const { + if (value) urma_delete_jfce(value); + } + void operator()(urma_target_jetty_t* value) const { + if (value) urma_unimport_jetty(value); + } + void operator()(urma_target_seg_t* value) const { + if (value) urma_unimport_seg(value); + } +}; + +struct urma_socket_shared_state_t + : std::enable_shared_from_this { + using callback_t = async_simple::util::move_only_function)>; + + struct pending_send { + urma_buffer_t buffer; + std::size_t length; + callback_t callback; + }; + + struct pending_recv { + std::pair result; + urma_buffer_t buffer; + }; + + urma_socket_shared_state_t(std::shared_ptr device, + ExecutorWrapper<>* executor, + std::size_t recv_buffer_cnt, + std::size_t send_buffer_cnt, + std::size_t cq_size) + : executor_(executor), + device_(std::move(device)), + socket_(executor->get_asio_executor()), + poll_timer_(executor->get_asio_executor()), + recv_buffer_cnt_(recv_buffer_cnt), + recv_queue_(recv_buffer_cnt + 1), + recv_result_(cq_size), + send_callbacks_(send_buffer_cnt + 2) {} + + ~urma_socket_shared_state_t() { release_resources(); } + + static void resume(std::pair result, + callback_t&& callback) { + if (callback) { + auto cb = std::move(callback); + cb(std::move(result)); + } + } + + bool init(std::size_t cq_size, std::size_t send_buffer_cnt, + bool event_mode) { + const auto& cap = device_->attr().dev_cap; + ELOG_INFO << "URMA resource init: device=" << device_->name() + << ", eid=" << device_->eid_string() + << ", eid_index=" << device_->eid_index() + << ", uasid=" << device_->uasid() + << ", jfc_depth=" << cq_size + << ", jfr_depth=" << recv_buffer_cnt_ + 1 + << ", jfs_depth=" << send_buffer_cnt + 2 + << ", max_jfc_depth=" << cap.max_jfc_depth + << ", max_jfr_depth=" << cap.max_jfr_depth + << ", max_jfs_depth=" << cap.max_jfs_depth + << ", event_mode=" << event_mode; + + // Create JFCE first so it can be bound to the JFC at creation time. + if (event_mode) { + errno = 0; + jfce_.reset(urma_create_jfce(device_->context())); + if (!jfce_) { + ELOG_WARN << "urma_create_jfce failed: errno=" << errno + << ", event_mode disabled, fall back to busy polling"; + event_mode_enabled_ = false; + } else { + ELOG_INFO << "urma_create_jfce succeeded: fd=" << jfce_->fd; + event_mode_enabled_ = true; + } + } else { + event_mode_enabled_ = false; + } + + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = static_cast(cq_size); + if (event_mode_enabled_) jfc_cfg.jfce = jfce_.get(); + errno = 0; + jfc_.reset(urma_create_jfc(device_->context(), &jfc_cfg)); + if (!jfc_) { + set_init_error("urma_create_jfc", errno); + ELOG_ERROR << "urma_create_jfc failed: depth=" << jfc_cfg.depth + << ", context=" << device_->context() + << ", errno=" << init_error_.value() + << ", error=" << init_error_.message(); + return false; + } + ELOG_INFO << "urma_create_jfc succeeded: jfc_id=" + << jfc_->jfc_id.id << ", depth=" << jfc_cfg.depth + << ", jfce=" << (event_mode_enabled_ ? "bound" : "null"); + + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = static_cast(recv_buffer_cnt_ + 1); + 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_.get(); + errno = 0; + jfr_.reset(urma_create_jfr(device_->context(), &jfr_cfg)); + if (!jfr_) { + set_init_error("urma_create_jfr", errno); + ELOG_ERROR << "urma_create_jfr failed: depth=" << jfr_cfg.depth + << ", trans_mode=" << static_cast(jfr_cfg.trans_mode) + << ", max_sge=" << static_cast(jfr_cfg.max_sge) + << ", jfc=" << jfr_cfg.jfc + << ", errno=" << init_error_.value() + << ", error=" << init_error_.message(); + return false; + } + ELOG_INFO << "urma_create_jfr succeeded: jfr_id=" + << jfr_->jfr_id.id << ", depth=" << jfr_cfg.depth; + + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = + static_cast(send_buffer_cnt + 2); + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.priority = 6; + 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_.get(); + jetty_cfg.shared.jfr = jfr_.get(); + jetty_cfg.shared.jfc = jfc_.get(); + errno = 0; + jetty_.reset(urma_create_jetty(device_->context(), &jetty_cfg)); + if (!jetty_) { + set_init_error("urma_create_jetty", errno); + ELOG_ERROR << "urma_create_jetty failed: jfs_depth=" + << jetty_cfg.jfs_cfg.depth + << ", trans_mode=" + << static_cast(jetty_cfg.jfs_cfg.trans_mode) + << ", priority=" + << static_cast(jetty_cfg.jfs_cfg.priority) + << ", shared_jfr=" << jetty_cfg.shared.jfr + << ", jfc=" << jetty_cfg.jfs_cfg.jfc + << ", errno=" << init_error_.value() + << ", error=" << init_error_.message(); + return false; + } + ELOG_INFO << "urma_create_jetty succeeded: jetty_id=" + << jetty_->jetty_id.id << ", uasid=" + << jetty_->jetty_id.uasid; + return true; + } + + void set_init_error(std::string_view stage, int error) { + init_stage_ = stage; + init_error_ = + error != 0 ? std::error_code(error, std::generic_category()) + : std::make_error_code(std::errc::io_error); + } + + std::error_code post_recv(urma_buffer_t buffer) { + urma_sge_t sge{reinterpret_cast(buffer.addr), + static_cast(buffer.length), + static_cast(buffer.seg), nullptr}; + urma_sg_t sg{&sge, 1}; + urma_jfr_wr_t wr{sg, 0, nullptr}; + urma_jfr_wr_t* bad_wr = nullptr; + auto ec = make_urma_error(urma_post_jfr_wr(jfr_.get(), &wr, &bad_wr)); + if (!ec) recv_queue_.push(std::move(buffer)); + return ec; + } + + std::error_code fill_recv_queue() { + while (recv_queue_.size() < recv_buffer_cnt_) { + auto buffer = device_->get_buffer_pool()->get_buffer(); + if (!buffer) return std::make_error_code(std::errc::no_buffer_space); + auto ec = post_recv(std::move(buffer)); + if (ec) return ec; + } + return {}; + } + + void post_send(urma_buffer_t buffer, std::size_t length, + callback_t&& callback) { + if (has_close_ || !remote_jetty_) { + if (buffer) device_->get_buffer_pool()->return_buffer(buffer); + resume({std::make_error_code(std::errc::not_connected), 0}, + std::move(callback)); + return; + } + + urma_sge_t sge{reinterpret_cast(buffer.addr), + static_cast(length), + static_cast(buffer.seg), nullptr}; + urma_sg_t sg{length ? &sge : nullptr, length ? 1u : 0u}; + urma_send_wr_t send_wr{}; + send_wr.src = sg; + urma_jfs_wr_t wr{}; + wr.opcode = URMA_OPC_SEND; + wr.flag.bs.complete_enable = 1; + wr.tjetty = remote_jetty_.get(); + wr.user_ctx = 1; + wr.send = send_wr; + urma_jfs_wr_t* bad_wr = nullptr; + auto ec = + make_urma_error(urma_post_jetty_send_wr(jetty_.get(), &wr, &bad_wr)); + if (ec) { + if (buffer) device_->get_buffer_pool()->return_buffer(buffer); + resume({ec, 0}, std::move(callback)); + return; + } + send_callbacks_.push( + pending_send{std::move(buffer), length, std::move(callback)}); + } + + void async_receive(callback_t&& callback) { + if (!recv_result_.empty()) { + auto pending = recv_result_.pop(); + recv_buffer_ = std::move(pending.buffer); + resume(std::move(pending.result), std::move(callback)); + } + else if (has_close_) { + resume({std::make_error_code(std::errc::operation_canceled), 0}, + std::move(callback)); + } + else { + recv_callback_ = std::move(callback); + } + } + + std::pair poll_completion() { + std::array completions{}; + int count = 0; + std::size_t polled = 0; + do { + count = urma_poll_jfc(jfc_.get(), static_cast(completions.size()), + completions.data()); + if (count < 0) + return {std::make_error_code(std::errc::io_error), polled}; + polled += static_cast(count); + for (int i = 0; i < count; ++i) { + auto& cr = completions[i]; + auto ec = cr.status == URMA_CR_SUCCESS + ? std::error_code{} + : std::make_error_code(std::errc::io_error); + if (ec) { + ELOG_ERROR << "URMA completion failed: status=" + << static_cast(cr.status) + << ", direction=" << (cr.flag.bs.s_r ? "recv" : "send") + << ", opcode=" << static_cast(cr.opcode) + << ", completion_len=" << cr.completion_len + << ", user_ctx=" << cr.user_ctx + << ", local_id=" << cr.local_id; + } + if (cr.flag.bs.s_r == 0) { + if (send_callbacks_.empty()) continue; + auto pending = send_callbacks_.pop(); + if (pending.buffer) + device_->get_buffer_pool()->return_buffer(pending.buffer); + resume({ec, pending.length}, std::move(pending.callback)); + wake_writer(ec); + continue; + } + + if (recv_queue_.empty()) + return {std::make_error_code(std::errc::protocol_error), polled}; + if (cr.completion_len == 0) { + peer_close_ = true; + has_close_ = true; + if (recv_callback_) { + resume({std::make_error_code(std::errc::connection_reset), 0}, + std::move(recv_callback_)); + } + continue; + } + auto completed_buffer = recv_queue_.pop(); + auto refill_ec = fill_recv_queue(); + if (refill_ec) { + ELOG_ERROR << "URMA refill recv queue failed: " + << refill_ec.message() + << ", recv_queue_size=" << recv_queue_.size() + << ", recv_buffer_cnt=" << recv_buffer_cnt_; + } + if (recv_callback_) { + recv_buffer_ = std::move(completed_buffer); + resume({ec, cr.completion_len}, std::move(recv_callback_)); + } + else { + if (recv_result_.full()) { + ELOG_ERROR << "URMA recv result queue is full; cannot cache " + "completed recv buffer"; + device_->get_buffer_pool()->return_buffer(completed_buffer); + return {std::make_error_code(std::errc::no_buffer_space), polled}; + } + recv_result_.push( + pending_recv{{ec, cr.completion_len}, std::move(completed_buffer)}); + } + } + } while (count == static_cast(completions.size())); + return {{}, polled}; + } + + void start_polling() { + auto self = shared_from_this(); + poll_timer_.expires_after(idle_poll_interval_); + poll_timer_.async_wait([self](const std::error_code& ec) { + if (ec || self->has_close_) return; + self->poll_once(); + }); + } + + void poll_once() { + if (has_close_) return; + auto [poll_ec, completion_count] = poll_completion(); + if (poll_ec) { + fail_pending(poll_ec); + close(); + return; + } + if (completion_count == 0) { + active_poll_budget_ = max_active_poll_budget_; + start_polling(); + return; + } + if (active_poll_budget_ == 0) { + active_poll_budget_ = max_active_poll_budget_; + start_polling(); + return; + } + --active_poll_budget_; + auto self = shared_from_this(); + asio::post(executor_->get_asio_executor(), [self] { + if (self->has_close_) return; + self->poll_once(); + }); + } + + // Wrap jfce_->fd in an asio stream_descriptor for async event waiting. + bool init_event_fd() { + if (!event_mode_enabled_ || !jfce_ || jfce_->fd < 0) return false; + int flags = fcntl(jfce_->fd, F_GETFL); + if (flags < 0) { + ELOG_WARN << "fcntl(F_GETFL) on jfce fd=" << jfce_->fd + << " failed: errno=" << errno << ", fall back to busy polling"; + event_mode_enabled_ = false; + return false; + } + if (fcntl(jfce_->fd, F_SETFL, flags | O_NONBLOCK) < 0) { + ELOG_WARN << "fcntl(F_SETFL, O_NONBLOCK) on jfce fd=" << jfce_->fd + << " failed: errno=" << errno << ", fall back to busy polling"; + event_mode_enabled_ = false; + return false; + } + try { + event_fd_ = std::make_unique( + executor_->get_asio_executor(), jfce_->fd); + } catch (const std::exception& e) { + ELOG_WARN << "create stream_descriptor for jfce fd=" << jfce_->fd + << " failed: " << e.what() << ", fall back to busy polling"; + event_mode_enabled_ = false; + return false; + } + return true; + } + + // Event-driven completion loop: rearm -> wait -> ack -> poll drain -> rearm. + // After an event wakeup, spin briefly (poll without sleeping) to catch the + // burst of completions that typically follow, avoiding repeated event + // wakeup latency under request/response workloads. Only when the spin + // budget is exhausted without new completions do we rearm and sleep. + async_simple::coro::Lazy event_loop() { + auto self = shared_from_this(); + std::error_code ec; + int consecutive_rearm_failures = 0; + while (!has_close_) { + if (urma_rearm_jfc(jfc_.get(), false) != URMA_SUCCESS) { + auto [drain_ec, drained] = poll_completion(); + if (drain_ec) { + fail_pending(drain_ec); + close(); + break; + } + if (urma_rearm_jfc(jfc_.get(), false) != URMA_SUCCESS) { + if (++consecutive_rearm_failures > 16) { + ELOG_WARN << "URMA rearm_jfc keeps failing after drain; yielding"; + consecutive_rearm_failures = 0; + coro_io::callback_awaitor yield_awaitor; + co_await yield_awaitor.await_resume([&self](auto handler) { + asio::post(self->executor_->get_asio_executor(), + [handler]() mutable { handler.resume(); }); + }); + } + continue; + } + consecutive_rearm_failures = 0; + } + coro_io::callback_awaitor awaitor; + ec = co_await awaitor.await_resume([&self](auto handler) { + self->event_fd_->async_wait( + asio::posix::stream_descriptor::wait_read, + [handler](const std::error_code& wait_ec) mutable { + handler.set_value_then_resume(wait_ec); + }); + }); + if (has_close_) break; + if (ec) { + ELOG_INFO << "URMA event fd wait ended with error: " << ec.message(); + break; + } + urma_jfc_t* ev_jfc = nullptr; + int ev_cnt = urma_wait_jfc(jfce_.get(), 1, 0, &ev_jfc); + if (ev_cnt > 0 && ev_jfc) { + uint32_t ack_cnt = 1; + urma_ack_jfc(&ev_jfc, &ack_cnt, 1); + } + // Drain all completions from this event, then decide: + // - If there are pending callbacks (send/recv waiting), skip spin and + // go straight to rearm+sleep so other event_loop coroutines (and + // the resumed callback coroutines) get CPU time. This prevents + // a short spin (4 polls) to catch an imminent CQE, then sleep. + // - If no pending callbacks, spin up to busy_poll_budget_ to catch a + // burst of idle traffic. + bool has_pending = !send_callbacks_.empty() || recv_callback_; + std::size_t idle_spins = 0; + std::size_t pending_budget = has_pending ? 64 : busy_poll_budget_; + while (!has_close_) { + auto [poll_ec, n] = poll_completion(); + if (poll_ec) { + fail_pending(poll_ec); + close(); + goto loop_end; + } + if (n == 0) { + if (++idle_spins >= pending_budget) break; + continue; + } + idle_spins = 0; + // After poll_completion resumes callbacks, has_pending may now be true + // (new sends/recvs posted by the resumed coroutines). Re-check and + // switch to the shorter pending budget so we don't spin too long. + bool now_pending = !send_callbacks_.empty() || recv_callback_; + if (now_pending && !has_pending) { + has_pending = true; + pending_budget = 4; + } + } + } + loop_end:; + } + + // Start the completion watcher; event-driven loop or legacy busy poll. + void start_completion_watch() { + if (event_mode_enabled_ && init_event_fd()) { + ELOG_INFO << "URMA starting event-driven completion loop (jfce fd=" + << jfce_->fd << ")"; + auto self = shared_from_this(); + event_loop().start([self](auto&&) { + ELOG_INFO << "URMA event_loop exited"; + }); + } else { + ELOG_INFO << "URMA starting timer-based busy poller (fallback)"; + poll_once(); + start_polling(); + } + } + + async_simple::coro::Lazy wait_for_send_slot( + std::size_t limit) { + if (send_callbacks_.size() < limit) co_return std::error_code{}; + write_promise_.emplace(); + co_return co_await write_promise_->getFuture(); + } + + void wake_writer(std::error_code ec) { + if (!write_promise_) return; + auto promise = std::move(*write_promise_); + write_promise_.reset(); + promise.setValue(ec); + } + + void fail_pending(std::error_code ec) { + if (recv_callback_) resume({ec, 0}, std::move(recv_callback_)); + while (!send_callbacks_.empty()) { + auto pending = send_callbacks_.pop(); + if (pending.buffer) + device_->get_buffer_pool()->return_buffer(pending.buffer); + resume({ec, 0}, std::move(pending.callback)); + } + wake_writer(ec); + } + + void close() { + if (has_close_.exchange(true)) return; + std::error_code ignored; + poll_timer_.cancel(ignored); + if (event_fd_) event_fd_->cancel(ignored); + socket_.cancel(ignored); + socket_.close(ignored); + fail_pending(std::make_error_code(std::errc::operation_canceled)); + } + + void release_resources() { + close(); + if (recv_buffer_) device_->get_buffer_pool()->return_buffer(recv_buffer_); + while (!recv_queue_.empty()) { + auto buffer = recv_queue_.pop(); + device_->get_buffer_pool()->return_buffer(buffer); + } + // Release stream_descriptor before closing the jfce fd it wraps. + event_fd_.reset(); + remote_seg_.reset(); + remote_jetty_.reset(); + jetty_.reset(); + jfce_.reset(); + jfr_.reset(); + jfc_.reset(); + } + + ExecutorWrapper<>* executor_; + std::shared_ptr device_; + asio::ip::tcp::socket socket_; + asio::steady_timer poll_timer_; + std::unique_ptr jfc_; + std::unique_ptr jfr_; + std::unique_ptr jetty_; + std::unique_ptr jfce_; + std::unique_ptr event_fd_; + std::unique_ptr remote_jetty_; + std::unique_ptr remote_seg_; + std::size_t recv_buffer_cnt_; + circle_buffer recv_queue_; + circle_buffer recv_result_; + circle_buffer send_callbacks_; + callback_t recv_callback_; + urma_buffer_t recv_buffer_; + std::optional> write_promise_; + std::atomic has_close_{false}; + bool peer_close_ = false; + bool event_mode_enabled_ = false; + std::size_t busy_poll_budget_ = 16; + std::chrono::microseconds idle_poll_interval_{5}; + static constexpr std::size_t max_active_poll_budget_ = 64; + std::size_t active_poll_budget_ = max_active_poll_budget_; + std::string init_stage_; + std::error_code init_error_; +}; + +} // namespace detail + +class urma_socket_t { + public: + struct config_t { + uint32_t cq_size = 128; + uint16_t recv_buffer_cnt = 8; + uint16_t send_buffer_cnt = 4; + uint32_t buffer_size = 4 * 1024; + uint64_t max_memory_usage = 256ull * 1024 * 1024; + std::string device_name; + int eid_index = 0; + urma_tp_type_t tp_type = URMA_CTP; + bool event_mode = true; + std::size_t busy_poll_budget = 16; + std::chrono::microseconds poll_interval{5}; + }; + + enum io_type { recv = 0, send = 1 }; + using callback_t = detail::urma_socket_shared_state_t::callback_t; + + struct urma_socket_info { + uint8_t eid[16]; + uint32_t uasid; + uint32_t jetty_id; + uint32_t buffer_size; + uint16_t recv_buffer_cnt; + uint8_t tp_type; + // Flattened from urma_seg_t which contains unions/bitfields not + // trivially serializable by struct_pack. + uint8_t seg_eid[16]; + uint32_t seg_uasid; + uint64_t seg_va; + uint64_t seg_len; + uint32_t seg_token_id; + constexpr static auto struct_pack_config = struct_pack::DISABLE_TYPE_INFO; + }; + + urma_socket_t(ExecutorWrapper<>* executor, const config_t& config) + : executor_(executor) { + init(config); + } + explicit urma_socket_t( + ExecutorWrapper<>* executor = coro_io::get_global_executor()) + : executor_(executor) { + init(config_t{}); + } + explicit urma_socket_t(const config_t& config) + : executor_(coro_io::get_global_executor()) { + init(config); + } + urma_socket_t(urma_socket_t&&) = default; + urma_socket_t& operator=(urma_socket_t&& other) { + if (this == &other) return *this; + close(); + executor_ = other.executor_; + conf_ = std::move(other.conf_); + state_ = std::move(other.state_); + remote_address_ = std::move(other.remote_address_); + handshake_remote_address_ = std::move(other.handshake_remote_address_); + handshake_local_address_ = std::move(other.handshake_local_address_); + remote_jetty_id_ = other.remote_jetty_id_; + handshake_remote_port_ = other.handshake_remote_port_; + handshake_local_port_ = other.handshake_local_port_; + buffer_size_ = other.buffer_size_; + send_window_size_ = other.send_window_size_; + remain_data_ = other.remain_data_; + return *this; + } + ~urma_socket_t() { close(); } + + bool is_open() const noexcept { + return state_ && !state_->has_close_ && state_->remote_jetty_; + } + auto get_executor() const { return executor_->get_asio_executor(); } + auto get_coro_executor() const { return executor_; } + const config_t& get_config() const noexcept { return conf_; } + uint32_t get_buffer_size() const noexcept { return buffer_size_; } + std::shared_ptr buffer_pool() const { + return state_->device_->get_buffer_pool(); + } + + async_simple::coro::Lazy connect( + const std::string& host, const std::string& port) noexcept { + auto tcp_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto ec = + co_await coro_io::async_connect(executor_, state_->socket_, host, port); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_connect_tcp, tcp_begin, 0); + if (!ec) ec = co_await connect_impl(); + if (ec) close(); + co_return ec; + } + + template + async_simple::coro::Lazy connect( + const EndPointSeq& endpoint) noexcept { + auto tcp_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto ec = co_await coro_io::async_connect(state_->socket_, endpoint); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_connect_tcp, tcp_begin, 0); + if (!ec) ec = co_await connect_impl(); + if (ec) close(); + co_return ec; + } + + async_simple::coro::Lazy accept( + std::string_view magic = "") noexcept { + urma_socket_info peer_info{}; + constexpr auto size = struct_pack::get_needed_size(peer_info); + if (magic.size() >= size.size()) + co_return std::make_error_code(std::errc::protocol_error); + std::array bytes{}; + std::memcpy(bytes.data(), magic.data(), magic.size()); + auto [ec, ignored] = co_await async_read( + state_->socket_, + asio::buffer(bytes.data() + magic.size(), bytes.size() - magic.size())); + if (ec) co_return ec; + if (struct_pack::deserialize_to(peer_info, std::span{bytes})) + co_return std::make_error_code(std::errc::protocol_error); + + ec = import_peer(peer_info); + if (ec) co_return ec; + ec = state_->fill_recv_queue(); + if (ec) co_return ec; + + auto local_info = make_local_info(); + struct_pack::serialize_to(bytes.data(), size, local_info); + auto [write_ec, ignored_write] = + co_await async_write(state_->socket_, asio::buffer(bytes)); + if (write_ec) co_return write_ec; + record_handshake_endpoints(); + close_handshake_socket(); + state_->start_completion_watch(); + co_return std::error_code{}; + } + + async_simple::coro::Lazy accept( + asio::ip::tcp::socket socket) noexcept { + prepare_accept(std::move(socket)); + return accept(); + } + + void prepare_accept(asio::ip::tcp::socket socket) noexcept { + state_->socket_ = std::move(socket); + } + + void close() const noexcept { + if (!state_) return; + asio::dispatch(executor_->get_asio_executor(), + [state = state_] { state->close(); }); + } + + auto cancel() const { + close(); + return std::error_code{}; + } + + void post_recv(callback_t&& callback) { + state_->async_receive(std::move(callback)); + } + + void post_send(urma_buffer_t buffer, std::size_t length, + callback_t&& callback) { + state_->post_send(std::move(buffer), length, std::move(callback)); + } + + async_simple::coro::Lazy waiting_write_over() { + return state_->wait_for_send_slot(send_window_size_); + } + + void poll_completion_once() { state_->poll_completion(); } + + std::size_t sent_request_count() const noexcept { + return state_->send_callbacks_.size(); + } + + std::size_t get_send_window_size() const noexcept { + return send_window_size_; + } + + urma_buffer_t get_send_buffer() { return buffer_pool()->get_buffer(); } + + urma_sge_t get_recv_buffer() const { + return {reinterpret_cast(state_->recv_buffer_.addr), + static_cast(state_->recv_buffer_.length), + static_cast(state_->recv_buffer_.seg), nullptr}; + } + + std::size_t consume(char* destination, std::size_t size) { + auto length = std::min(size, remain_data_.size()); + std::memcpy(destination, remain_data_.data(), length); + remain_data_.remove_prefix(length); + if (remain_data_.empty()) release_recv_buffer(); + return length; + } + + std::size_t remain_read_buffer_size() const { return remain_data_.size(); } + + owned_data_view detach_remain_data_view() { + if (remain_data_.empty() || !state_->recv_buffer_) return {}; + auto owner = std::make_shared( + buffer_pool(), std::move(state_->recv_buffer_)); + owned_data_view view{data_view{remain_data_, -1}, std::move(owner)}; + remain_data_ = {}; + return view; + } + + owned_data_view detach_recv_buffer_view(std::size_t length) { + if (!state_->recv_buffer_ || length == 0) return {}; + auto size = std::min(length, state_->recv_buffer_.length); + auto data = data_view{state_->recv_buffer_.addr, size, -1}; + auto owner = std::make_shared( + buffer_pool(), std::move(state_->recv_buffer_)); + return owned_data_view{data, std::move(owner)}; + } + + void set_read_buffer_len(std::size_t consumed, std::size_t remaining) { + remain_data_ = std::string_view( + static_cast(state_->recv_buffer_.addr) + consumed, remaining); + if (remaining == 0) release_recv_buffer(); + } + + asio::ip::address get_remote_address() const noexcept { + return handshake_remote_port_ != 0 ? handshake_remote_address_ + : remote_address_; + } + uint32_t get_remote_qp_num() const noexcept { + return handshake_remote_port_ != 0 ? handshake_remote_port_ + : remote_jetty_id_; + } + asio::ip::address get_local_address() const noexcept { + return handshake_local_port_ != 0 ? handshake_local_address_ + : state_->device_->gid_address(); + } + uint32_t get_local_qp_num() const noexcept { + return handshake_local_port_ != 0 ? handshake_local_port_ + : state_->jetty_->jetty_id.id; + } + + constexpr static uint32_t urma_md5_header = + struct_pack::get_type_code(); + constexpr static char urma_md5_first_header = urma_md5_header % 256; + + private: + void init(config_t config) { + ELOG_INFO << "URMA socket init requested: device=" << config.device_name + << ", eid_index=" << config.eid_index + << ", tp_type=" << static_cast(config.tp_type) + << ", cq_size=" << config.cq_size + << ", recv_buffer_cnt=" << config.recv_buffer_cnt + << ", send_buffer_cnt=" << config.send_buffer_cnt + << ", buffer_size=" << config.buffer_size + << ", max_memory_usage=" << config.max_memory_usage + << ", executor=" << executor_; + constexpr uint32_t ctp_max_send_size = 4 * 1024; + if (config.tp_type == URMA_CTP && config.buffer_size > ctp_max_send_size) { + ELOG_WARN << "URMA CTP buffer_size " << config.buffer_size + << " is larger than the documented bonding CTP max send packet " + "size; clamp to " + << ctp_max_send_size; + config.buffer_size = ctp_max_send_size; + } + config.recv_buffer_cnt = std::max(config.recv_buffer_cnt, 1); + config.send_buffer_cnt = std::max(config.send_buffer_cnt, 1); + config.cq_size = + std::max(config.cq_size, config.recv_buffer_cnt + + config.send_buffer_cnt + 2); + conf_ = std::move(config); + auto device = get_global_urma_device( + {.dev_name = conf_.device_name, + .buffer_pool_config = {.buffer_size = conf_.buffer_size, + .max_memory_usage = + conf_.max_memory_usage}, + .eid_index = conf_.eid_index}); + if (!device || !device->is_valid() || !device->get_buffer_pool()) + throw std::system_error( + std::make_error_code(std::errc::no_such_device)); + const auto& cap = device->attr().dev_cap; + ELOG_INFO << "URMA device capabilities: device=" << device->name() + << ", rm_tp_cap=" << cap.rm_tp_cap.value + << ", rtp=" << cap.rm_tp_cap.bs.rtp + << ", ctp=" << cap.rm_tp_cap.bs.ctp + << ", ctp_en=" << cap.feature.bs.ctp_en + << ", trans_mode=" << cap.trans_mode + << ", max_jfc=" << cap.max_jfc + << ", max_jfr=" << cap.max_jfr + << ", max_jetty=" << cap.max_jetty; + if (conf_.tp_type == URMA_CTP && !device->supports_rm_ctp()) { + ELOG_WARN << "URMA device capability does not report RM CTP support; " + "continuing because some providers expose CTP through " + "feature.ctp_en or resource creation"; + } + if (conf_.tp_type == URMA_RTP && !device->supports_rm_rtp()) { + ELOG_WARN << "URMA device capability does not report RM RTP support; " + "continuing and relying on resource creation"; + } + buffer_size_ = std::min( + conf_.buffer_size, device->get_buffer_pool()->buffer_size()); + state_ = std::make_shared( + std::move(device), executor_, conf_.recv_buffer_cnt, + conf_.send_buffer_cnt, conf_.cq_size); + state_->busy_poll_budget_ = conf_.busy_poll_budget; + state_->idle_poll_interval_ = conf_.poll_interval; + if (!state_->init(conf_.cq_size, conf_.send_buffer_cnt, conf_.event_mode)) { + auto stage = state_->init_stage_; + auto error = state_->init_error_; + ELOG_ERROR << "URMA socket resource initialization failed: stage=" + << stage << ", errno=" << error.value() + << ", error=" << error.message(); + state_.reset(); + throw std::system_error(error, stage); + } + ELOG_INFO << "URMA socket init succeeded: device=" + << state_->device_->name() << ", jetty_id=" + << state_->jetty_->jetty_id.id; + } + + urma_socket_info make_local_info() const { + urma_socket_info info{}; + std::memcpy(info.eid, state_->device_->eid().raw, sizeof(info.eid)); + info.uasid = state_->jetty_->jetty_id.uasid; + info.jetty_id = state_->jetty_->jetty_id.id; + info.buffer_size = buffer_pool()->buffer_size(); + info.recv_buffer_cnt = conf_.recv_buffer_cnt; + info.tp_type = static_cast(conf_.tp_type); + auto pool_seg = buffer_pool()->seg(); + std::memcpy(info.seg_eid, pool_seg.ubva.eid.raw, sizeof(info.seg_eid)); + info.seg_uasid = pool_seg.ubva.uasid; + info.seg_va = pool_seg.ubva.va; + info.seg_len = pool_seg.len; + info.seg_token_id = pool_seg.token_id; + return info; + } + + std::error_code import_peer(const urma_socket_info& peer) { + urma_rjetty_t remote{}; + std::memcpy(remote.jetty_id.eid.raw, peer.eid, sizeof(peer.eid)); + 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)) { + ELOG_ERROR << "invalid remote URMA TP type: " + << static_cast(peer.tp_type); + return std::make_error_code(std::errc::protocol_error); + } + remote.tp_type = static_cast(peer.tp_type); + + // Import the peer's buffer pool segment BEFORE importing the jetty. + // The URMA perftest reference implementation calls urma_import_seg + // before urma_import_jetty/urma_import_jetty_ex. Without this step, + // the kernel may not establish the transport path (TP) routing for + // the remote EID, causing the first SEND to be immediately rejected + // by hardware with URMA_CR_RNR_RETRY_CNT_EXC_ERR (status=10). + 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; + urma_seg_t peer_seg{}; + std::memcpy(peer_seg.ubva.eid.raw, peer.seg_eid, sizeof(peer.seg_eid)); + 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; + state_->remote_seg_.reset( + urma_import_seg(state_->device_->context(), &peer_seg, + &seg_token, 0, seg_flag)); + if (!state_->remote_seg_) { + ELOG_WARN << "urma_import_seg failed: errno=" << errno + << ", continuing with urma_import_jetty"; + } else { + ELOG_INFO << "urma_import_seg succeeded for peer EID=" + << eid_to_address(peer.eid).to_string(); + } + + urma_token_t token{}; + errno = 0; + state_->remote_jetty_.reset( + urma_import_jetty(state_->device_->context(), &remote, &token)); + // If plain import fails on a bonding device, retry with the bonding + // extension (has_drv_ext=1 + local jetty). + if (!state_->remote_jetty_ && + state_->device_->name().compare(0, 7, "bonding") == 0 && + remote.trans_mode == URMA_TM_RM) { + ELOG_WARN << "plain urma_import_jetty failed: errno=" << errno + << ", retrying with bonding extension"; + bondp_rjetty_t bondp_rjetty{}; + bondp_rjetty.base = remote; + bondp_rjetty.base.flag.bs.has_drv_ext = 1; + bondp_rjetty.jetty = state_->jetty_.get(); + errno = 0; + state_->remote_jetty_.reset( + urma_import_jetty(state_->device_->context(), + &bondp_rjetty.base, &token)); + } + if (!state_->remote_jetty_) { + auto error = errno != 0 + ? std::error_code(errno, std::generic_category()) + : std::make_error_code(std::errc::connection_refused); + ELOG_ERROR << "urma_import_jetty failed: " << error.message() + << ", errno=" << errno << ", remote_eid=" + << eid_to_address(peer.eid).to_string() + << ", remote_uasid=" << peer.uasid + << ", remote_jetty_id=" << peer.jetty_id + << ", trans_mode=" << remote.trans_mode + << ", tp_type=" << remote.tp_type; + return error; + } + remote_jetty_id_ = peer.jetty_id; + buffer_size_ = + std::min(peer.buffer_size, buffer_pool()->buffer_size()); + const auto remote_recv_capacity = + std::max(peer.recv_buffer_cnt, 1); + send_window_size_ = std::min( + conf_.send_buffer_cnt, + remote_recv_capacity > 1 ? remote_recv_capacity - 1 : 1); + ELOG_INFO << "URMA peer imported: remote_recv_buffer_cnt=" + << peer.recv_buffer_cnt + << ", local_send_buffer_cnt=" << conf_.send_buffer_cnt + << ", effective_send_window=" << send_window_size_; + remote_address_ = eid_to_address(peer.eid); + return {}; + } + + async_simple::coro::Lazy connect_impl() { + auto handshake_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto ec = state_->fill_recv_queue(); + if (ec) co_return ec; + auto local_info = make_local_info(); + constexpr auto size = struct_pack::get_needed_size(local_info); + std::array bytes{}; + struct_pack::serialize_to(bytes.data(), size, local_info); + auto [write_ec, ignored_write] = + co_await async_write(state_->socket_, asio::buffer(bytes)); + if (write_ec) co_return write_ec; + auto [read_ec, ignored_read] = + co_await async_read(state_->socket_, asio::buffer(bytes)); + if (read_ec) co_return read_ec; + urma_socket_info peer{}; + if (struct_pack::deserialize_to(peer, std::span{bytes})) + co_return std::make_error_code(std::errc::protocol_error); + ec = import_peer(peer); + if (ec) co_return ec; + record_handshake_endpoints(); + close_handshake_socket(); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_connect_handshake, + handshake_begin, 0); + state_->start_completion_watch(); + co_return std::error_code{}; + } + + void record_handshake_endpoints() { + std::error_code remote_ec; + auto remote_ep = state_->socket_.remote_endpoint(remote_ec); + if (!remote_ec) { + handshake_remote_address_ = remote_ep.address(); + handshake_remote_port_ = remote_ep.port(); + } + std::error_code local_ec; + auto local_ep = state_->socket_.local_endpoint(local_ec); + if (!local_ec) { + handshake_local_address_ = local_ep.address(); + handshake_local_port_ = local_ep.port(); + } + ELOG_INFO << "URMA handshake TCP endpoint: remote=" + << handshake_remote_address_.to_string() << ":" + << handshake_remote_port_ << ", local=" + << handshake_local_address_.to_string() << ":" + << handshake_local_port_ << ", remote_jetty_id=" + << remote_jetty_id_ << ", local_jetty_id=" + << state_->jetty_->jetty_id.id; + } + + void release_recv_buffer() { + if (state_->recv_buffer_) + state_->device_->get_buffer_pool()->return_buffer(state_->recv_buffer_); + } + + void close_handshake_socket() { + std::error_code ignored; + state_->socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored); + state_->socket_.close(ignored); + } + + static asio::ip::address eid_to_address(const uint8_t* eid) { + asio::ip::address_v6::bytes_type bytes{}; + std::memcpy(bytes.data(), eid, bytes.size()); + return asio::ip::address_v6(bytes); + } + + ExecutorWrapper<>* executor_; + config_t conf_; + std::shared_ptr state_; + asio::ip::address remote_address_; + asio::ip::address handshake_remote_address_; + asio::ip::address handshake_local_address_; + uint32_t remote_jetty_id_ = 0; + uint32_t handshake_remote_port_ = 0; + uint32_t handshake_local_port_ = 0; + uint32_t buffer_size_ = 0; + std::size_t send_window_size_ = 1; + std::string_view remain_data_; +}; + +template +inline async_simple::coro::Lazy async_connect( + urma_socket_t& socket, const EndPointSeq& endpoint) noexcept { + return socket.connect(endpoint); +} + +inline async_simple::coro::Lazy async_connect( + urma_socket_t& socket, const std::string& host, + const std::string& port) noexcept { + return socket.connect(host, port); +} + +} // namespace coro_io diff --git a/include/ylt/coro_rpc/impl/coro_connection.hpp b/include/ylt/coro_rpc/impl/coro_connection.hpp index f68142f86..8c7040dee 100644 --- a/include/ylt/coro_rpc/impl/coro_connection.hpp +++ b/include/ylt/coro_rpc/impl/coro_connection.hpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include "asio/dispatch.hpp" @@ -39,6 +40,7 @@ #include "ylt/coro_io/data_view.hpp" #include "ylt/coro_io/heterogeneous_buffer.hpp" #include "ylt/coro_io/socket_wrapper.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" #include "ylt/coro_rpc/impl/errno.h" #include "ylt/util/utils.hpp" #ifdef UNIT_TEST_INJECT @@ -60,6 +62,7 @@ struct context_info_t { typename rpc_protocol::req_header req_head_; std::string req_body_; coro_io::heterogeneous_buffer req_attachment_; + std::vector req_attachment_views_; std::function resp_attachment_ = [] { return coro_io::data_view{std::string_view{}, -1}; }; @@ -101,6 +104,12 @@ struct context_info_t { } std::string_view get_request_attachment() const; coro_io::data_view get_request_attachment2() const; + std::span get_request_attachment_views() + const noexcept; + void set_request_attachment_views( + std::vector views) noexcept { + req_attachment_views_ = std::move(views); + } std::string release_request_attachment(); coro_io::heterogeneous_buffer release_request_attachment2(); std::any &tag() noexcept; @@ -210,6 +219,21 @@ class coro_connection : public std::enable_shared_from_this { co_return; } } +#endif +#ifdef YLT_ENABLE_URMA + if constexpr (std::is_same_v< + Socket, coro_io::socket_wrapper_t::urma_socket_type>) { + reset_timer(0, "urma handshake"); + auto ec = co_await socket.accept(magic_number); + magic_number = ""; + cancel_timer(0, "urma handshake"); + if (ec) [[unlikely]] { + ELOG_ERROR << "urma handshake failed: " << ec.message() + << " conn_id " << conn_id_; + close(); + co_return; + } + } #endif auto context_info = std::make_shared>( router, shared_from_this()); @@ -219,6 +243,9 @@ class coro_connection : public std::enable_shared_from_this { typename rpc_protocol::req_header req_head_tmp{}; std::error_code ec; auto tp = std::chrono::steady_clock::now(); + auto profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; // timer will be reset after rpc call response if (req_id == 0) { ec = co_await rpc_protocol::read_first_head(socket, req_head_tmp, @@ -227,6 +254,9 @@ class coro_connection : public std::enable_shared_from_this { else { ec = co_await rpc_protocol::read_head(socket, req_head_tmp); } + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_read_header, + profile_begin, req_head_tmp.length + req_head_tmp.attach_length); // `co_await async_read` uses asio::async_read underlying. // If eof occurred, the bytes_transferred of `co_await async_read` must // less than RPC_HEAD_LEN. Incomplete data will be discarded. @@ -268,6 +298,7 @@ class coro_connection : public std::enable_shared_from_this { } else { // reuse string buffer + context_info->req_attachment_views_.clear(); context_info = std::make_shared>( router, shared_from_this(), std::move(context_info->req_body_), std::move(context_info->req_attachment_)); @@ -293,8 +324,16 @@ class coro_connection : public std::enable_shared_from_this { std::string_view payload; // rpc_protocol::buffer_type maybe from user, default from framework. + profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; ec = co_await rpc_protocol::read_payload(socket, req_head, body, - req_attachment); + req_attachment, + context_info.get()); + auto req_payload_size = body.size() + req_attachment.size(); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_read_payload, + profile_begin, req_payload_size); cancel_timer(req_id, "recv client data"); payload = std::string_view{body}; @@ -323,6 +362,9 @@ class coro_connection : public std::enable_shared_from_this { if (!handler) { auto coro_handler = router.get_coro_handler(key); set_rpc_return_by_callback(); + auto dispatch_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; router .route_coro(conn_id_, req_id, coro_handler, payload, serialize_proto.value(), key) @@ -355,12 +397,21 @@ class coro_connection : public std::enable_shared_from_this { }); }, socket_wrapper_.get_executor()); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_dispatch, + dispatch_begin, req_payload_size); } else { coro_rpc::detail::set_context() = context_info.get(); + auto dispatch_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; auto &&[resp_err, resp_buf] = router.route(conn_id_, req_id, handler, payload, context_info, serialize_proto.value(), key); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_dispatch, + dispatch_begin, req_payload_size); if (is_rpc_return_by_callback_) { if (!resp_err) { continue; @@ -427,8 +478,15 @@ class coro_connection : public std::enable_shared_from_this { ELOG_WARN << "rpc route/execute error, error msg: " << resp_error_msg << ", conn_id = " << conn_id_; } + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; std::string header_buf = rpc_protocol::prepare_response( resp_buf, req_head, attachment().length(), resp_err, resp_error_msg); + auto resp_payload_size = resp_buf.size() + attachment().length(); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_response, + ser_begin, resp_payload_size); response(start_tp, req_id, std::move(header_buf), std::move(resp_buf), std::move(attachment), std::move(complete_handler), nullptr) @@ -443,8 +501,14 @@ class coro_connection : public std::enable_shared_from_this { const typename rpc_protocol::req_header &req_head, std::function &&complete_handler) { + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; std::string header_buf = rpc_protocol::prepare_response( body_buf, req_head, resp_attachment().size()); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_response, + ser_begin, body_buf.size() + resp_attachment().size()); asio::dispatch( socket_wrapper_.get_executor()->get_asio_executor(), [watcher = weak_from_this(), header_buf = std::move(header_buf), @@ -469,8 +533,14 @@ class coro_connection : public std::enable_shared_from_this { std::function &&complete_handler) { std::string body_buf; + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; std::string header_buf = rpc_protocol::prepare_response(body_buf, req_head, 0, ec, error_msg); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_response, + ser_begin, body_buf.size()); asio::dispatch( socket_wrapper_.get_executor()->get_asio_executor(), [watcher = weak_from_this(), header_buf = std::move(header_buf), @@ -585,6 +655,11 @@ class coro_connection : public std::enable_shared_from_this { } #endif coro_io::data_view attachment = std::get<2>(msg)(); + auto send_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto send_payload_size = std::get<0>(msg).size() + + std::get<1>(msg).size() + attachment.size(); if (attachment.empty()) { std::array buffers{ asio::buffer(std::get<0>(msg)), asio::buffer(std::get<1>(msg))}; @@ -605,6 +680,9 @@ class coro_connection : public std::enable_shared_from_this { ret = co_await coro_io::async_write(socket, buffers); } } + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_send_response, + send_begin, send_payload_size); auto &complete_handler = std::get<3>(msg); if (complete_handler) { complete_handler(ret.first, ret.second); @@ -636,6 +714,9 @@ class coro_connection : public std::enable_shared_from_this { std::function resp_attachment, std::function complete_handler, rpc_conn self) noexcept { + auto response_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; if (has_closed()) AS_UNLIKELY { ELOG_DEBUG << "response_msg failed: connection has been closed" @@ -652,6 +733,9 @@ class coro_connection : public std::enable_shared_from_this { write_queue_.emplace_back(std::move(header_buf), std::move(body_buf), std::move(resp_attachment), std::move(complete_handler)); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_response_queue, + response_begin, 0); --rpc_processing_cnt_; assert(rpc_processing_cnt_ >= 0); ELOG_INFO << "finish rpc function execution, conn_id = " << conn_id_ @@ -821,6 +905,12 @@ coro_io::data_view context_info_t::get_request_attachment2() return req_attachment_; } +template +std::span +context_info_t::get_request_attachment_views() const noexcept { + return req_attachment_views_; +} + template std::string context_info_t::release_request_attachment() { auto str = req_attachment_.get_string(); diff --git a/include/ylt/coro_rpc/impl/coro_rpc_client.hpp b/include/ylt/coro_rpc/impl/coro_rpc_client.hpp index 207e4a18f..221f0840d 100644 --- a/include/ylt/coro_rpc/impl/coro_rpc_client.hpp +++ b/include/ylt/coro_rpc/impl/coro_rpc_client.hpp @@ -61,12 +61,18 @@ #include "ylt/coro_io/ibverbs/ib_buffer.hpp" #include "ylt/coro_io/ibverbs/ib_socket.hpp" #endif +#ifdef YLT_ENABLE_URMA +#include "ylt/coro_io/urma/urma_rpc_env.hpp" +#include "ylt/coro_io/urma/urma_socket.hpp" +#endif +#include "ylt/coro_io/data_view.hpp" #ifdef YLT_ENABLE_ND #include "ylt/coro_io/networkdirect/nd_socket.hpp" #endif #include "ylt/coro_io/heterogeneous_buffer.hpp" #include "ylt/coro_io/io_context_pool.hpp" #include "ylt/coro_io/socket_wrapper.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" #include "ylt/coro_rpc/impl/errno.h" #include "ylt/struct_pack.hpp" #include "ylt/struct_pack/reflection.hpp" @@ -255,6 +261,10 @@ class coro_rpc_client { , coro_io::ib_socket_t::config_t #endif +#ifdef YLT_ENABLE_URMA + , + coro_io::urma_socket_t::config_t +#endif #ifdef YLT_ENABLE_ND , coro_io::nd_socket_t::config_t @@ -285,8 +295,11 @@ class coro_rpc_client { coro_io::ExecutorWrapper<> *executor = coro_io::get_global_executor(), config conf = {}) : timer_(std::make_unique( - executor->get_asio_executor())), + executor ? executor->get_asio_executor() : asio::io_context{}.get_executor())), control_(std::make_shared(executor, false, conf.local_ip)) { + if (!executor) [[unlikely]] { + ELOG_ERROR << "coro_rpc_client: executor is nullptr"; + } if (!init_config(conf)) [[unlikely]] { close(); } @@ -310,6 +323,15 @@ class coro_rpc_client { return control_->socket_wrapper_.init_client(config); } #endif +#ifdef YLT_ENABLE_URMA + [[nodiscard]] bool init_socket_wrapper( + const coro_io::urma_socket_t::config_t &config) { + ELOG_INFO << "URMA init_socket_wrapper: buffer_size=" << config.buffer_size + << " recv_buffer_cnt=" << config.recv_buffer_cnt + << " send_buffer_cnt=" << config.send_buffer_cnt; + return control_->socket_wrapper_.init_client(config); + } +#endif #ifdef YLT_ENABLE_ND [[nodiscard]] bool init_socket_wrapper( const coro_io::nd_socket_t::config_t &config) { @@ -608,13 +630,20 @@ class coro_rpc_client { [[nodiscard]] bool init_config(const config &conf) { create_tp_ = std::chrono::steady_clock::now(); config_ = conf; +#ifdef YLT_ENABLE_URMA + if (std::holds_alternative(config_.socket_config)) { + if (auto urma_config = coro_io::try_make_urma_rpc_config()) { + config_.socket_config = *urma_config; + } + } +#endif control_->socket_wrapper_.set_local_ip(config_.local_ip); - control_->client_id = conf.client_id; + control_->client_id = config_.client_id; return std::visit( [this](auto &socket_config) { return init_socket_wrapper(socket_config); }, - conf.socket_config); + config_.socket_config); }; auto get_create_time_point() const noexcept { return create_tp_; } @@ -917,6 +946,15 @@ class coro_rpc_client { std::get(config_.socket_config)); } #endif +#ifdef YLT_ENABLE_URMA + [[nodiscard]] bool init_urma( + const coro_io::urma_socket_t::config_t &config = {}) { + ELOG_DEBUG << "URMA init_urma: buffer_size=" << config.buffer_size; + config_.socket_config = config; + return init_socket_wrapper( + std::get(config_.socket_config)); + } +#endif #ifdef YLT_ENABLE_ND [[nodiscard]] bool init_nd( const coro_io::nd_socket_t::config_t &config = {}) { @@ -1223,6 +1261,8 @@ class coro_rpc_client { co_return false; } if (auto self = socket_watcher.lock()) { + ELOG_WARN << err_msg << ", close socket by timeout" + << ", client_id: " << config_.client_id; self->is_timeout_ = is_timeout; close_socket_async(self); co_return true; @@ -1467,7 +1507,9 @@ class coro_rpc_client { std::unordered_map response_handler_table_; resp_body resp_buffer_; std::atomic recving_cnt_ = 0; + std::atomic recv_running_ = false; uint64_t client_id = 0; + std::size_t last_req_payload_size = 0; control_t(coro_io::ExecutorWrapper<> *executor, bool is_timeout, const std::string &local_ip) : is_timeout_(is_timeout), @@ -1544,8 +1586,10 @@ class coro_rpc_client { static_check(); if (config.request_timeout_duration->count() >= 0) { - timeout(timer, *config.request_timeout_duration, - "rpc call timer canceled") + ELOG_DEBUG << "rpc call timer start, timeout_ms=" + << config.request_timeout_duration->count() + << ", client_id: " << config_.client_id; + timeout(timer, *config.request_timeout_duration, "rpc call timeout") .start([](auto &&) { }); } @@ -1573,6 +1617,9 @@ class coro_rpc_client { coro_rpc_protocol::resp_header header; char buffer[coro_rpc_protocol::RESP_HEAD_LEN]; auto tp = std::chrono::steady_clock::now(); + auto profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; ret = co_await coro_io::async_read(socket, asio::buffer(buffer)); [[maybe_unused]] auto ec = struct_pack::deserialize_to< struct_pack::sp_config::DISABLE_ALL_META_INFO>( @@ -1601,6 +1648,10 @@ class coro_rpc_client { << ". start notify response handler" << ", client_id: " << controller->client_id; uint32_t body_len = header.length; + auto resp_payload_size = body_len + header.attach_length; + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_recv_header, + profile_begin, resp_payload_size); struct_pack::detail::resize( controller->resp_buffer_.read_buf_, std::max(body_len, sizeof(std::string))); @@ -1610,9 +1661,15 @@ class coro_rpc_client { controller->resp_buffer_.read_buf_.resize(body_len); } if (header.attach_length == 0) { + profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; ret = co_await coro_io::async_read( socket, asio::buffer(controller->resp_buffer_.read_buf_.data(), body_len)); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_recv_payload, + profile_begin, resp_payload_size); controller->resp_buffer_.resp_attachment_buf_.clear(); } else { @@ -1657,7 +1714,13 @@ class coro_rpc_client { body_len}, -1}, attachment_buffer}; + profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; ret = co_await coro_io::async_read(socket, iov); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_recv_payload, + profile_begin, resp_payload_size); } else { std::array iov{ @@ -1665,7 +1728,13 @@ class coro_rpc_client { body_len}, asio::mutable_buffer{attachment_buffer.mutable_data(), header.attach_length}}; + profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; ret = co_await coro_io::async_read(socket, iov); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_recv_payload, + profile_begin, resp_payload_size); } } auto cost_time = (std::chrono::steady_clock::now() - tp) / @@ -1692,9 +1761,17 @@ class coro_rpc_client { iter->second(std::move(controller->resp_buffer_), header.err_code); controller->response_handler_table_.erase(iter); if (controller->response_handler_table_.empty()) { + controller->recv_running_.store(false, std::memory_order_release); + // Re-check: a new send_request may have inserted a handler between + // the empty() check and the store. If so, restart the recv loop. + if (!controller->response_handler_table_.empty()) { + controller->recv_running_.store(true, std::memory_order_release); + continue; + } co_return; } } while (true); + controller->recv_running_.store(false, std::memory_order_release); close_socket_async(controller); send_err_response(controller.get(), ret.first); co_return; @@ -1723,11 +1800,19 @@ class coro_rpc_client { static async_simple::coro::Lazy> deserialize_rpc_result( async_simple::Future future, std::weak_ptr watcher, recving_guard guard, - uint64_t client_id) { + uint64_t client_id, uint64_t rpc_begin = 0, + std::size_t req_payload_size = 0) { + auto record_rpc = [&]() { + if (rpc_begin) + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::benchmark_rpc_call, + rpc_begin, req_payload_size); + }; auto ret_ = co_await std::move(future); guard.release(); if (ret_.index() == 1) [[unlikely]] { // local error auto &ret = std::get<1>(ret_); + record_rpc(); if (ret.value() == static_cast(std::errc::operation_canceled) || ret.value() == static_cast(std::errc::timed_out)) { co_return coro_rpc::unexpected{ @@ -1741,8 +1826,15 @@ class coro_rpc_client { bool has_error = false; auto &ret = std::get<0>(ret_); + auto deser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; auto result = handle_response_buffer(ret.buffer_.read_buf_, ret.errc_, has_error, client_id); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_deserialize_response, + deser_begin, ret.buffer_.read_buf_.size()); + record_rpc(); if (has_error) { if (auto w = watcher.lock(); w) { close_socket_async(std::move(w)); @@ -1802,6 +1894,9 @@ class coro_rpc_client { async_rpc_result())>>> send_request(request_config_t config, Args &&...args) { using rpc_return_t = decltype(get_return_type()); + auto rpc_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; recving_guard guard(control_.get()); uint32_t id; if (!config.request_timeout_duration) { @@ -1818,7 +1913,6 @@ class coro_rpc_client { if (!result) { async_simple::Promise promise; auto future = promise.getFuture(); - bool is_empty = control_->response_handler_table_.empty(); auto &&[_, is_ok] = control_->response_handler_table_.try_emplace( id, std::move(timer), std::move(promise), coro_io::data_view{config.resp_attachment_buf, @@ -1829,15 +1923,25 @@ class coro_rpc_client { rpc_error{coro_rpc::errc::serial_number_conflict}); } else { - if (is_empty) { + // Ensure at most one recv coroutine runs per connection. + // The previous is_empty check on response_handler_table_ was racy: + // multiple send_request coroutines could observe is_empty == true + // concurrently and each start a recv coroutine, causing multiple + // concurrent readers on the same URMA socket. + if (!control_->recv_running_.exchange(true)) { control_->socket_wrapper_.visit([control_ = control_](auto &socket) { recv(control_, socket).start([](auto &&) { + // recv_running_ is cleared inside the recv coroutine itself, + // so we don't reset it here. This avoids a race where the + // recv coroutine has decided to exit but a new handler was + // inserted before recv_running_ was cleared. }); }); } co_return deserialize_rpc_result( std::move(future), std::weak_ptr{control_}, - std::move(guard), config_.client_id); + std::move(guard), config_.client_id, rpc_begin, + control_->last_req_payload_size); } } else { @@ -1854,8 +1958,16 @@ class coro_rpc_client { async_simple::coro::Lazy send_impl( Socket &socket, uint32_t &id, coro_io::data_view req_attachment, Args &&...args) { + auto prepare_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; auto buffer = prepare_buffer(id, req_attachment.size(), std::forward(args)...); + auto payload_size = buffer.size() + req_attachment.size(); + control_->last_req_payload_size = payload_size; + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_prepare_request, + prepare_begin, payload_size); if (buffer.empty()) { co_return rpc_error{errc::message_too_large}; } @@ -1868,8 +1980,9 @@ class coro_rpc_client { #endif std::pair ret; auto tp = std::chrono::steady_clock::now(); - ELOG_TRACE << "rpc request send start, client_id: " << config_.client_id - << ", request ID: " << id; + ELOG_DEBUG << "rpc request send start, client_id: " << config_.client_id + << ", request ID: " << id << ", body_size=" << buffer.size() + << ", attachment_size=" << req_attachment.size(); #ifdef UNIT_TEST_INJECT if (g_action == inject_action::client_send_bad_header) { buffer[0] = (std::byte)(uint8_t(buffer[0]) + 1); @@ -1925,6 +2038,23 @@ class coro_rpc_client { }, control_->executor_); } +#ifdef YLT_ENABLE_URMA + // URMA backpressure: wait for send slots before posting new WRs. + // Without this, consecutive async_write calls (even serialized by + // write_mutex_) can overrun the remote JFR's pre-posted recv buffers, + // causing RNR retry exhaustion and WR_FLUSH_ERR. + if constexpr (std::is_same_v) { + auto slot_ec = co_await socket.waiting_write_over(); + if (slot_ec) { + write_mutex_ = false; + close(); + co_return rpc_error{errc::io_error, slot_ec.message()}; + } + } +#endif + auto send_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; if (req_attachment.empty()) { ret = co_await coro_io::async_write( socket, asio::buffer(buffer.data(), buffer.size())); @@ -1944,6 +2074,9 @@ class coro_rpc_client { ret = co_await coro_io::async_write(socket, iov); } } + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::client_send_request, + send_begin, payload_size); write_mutex_ = false; #ifdef UNIT_TEST_INJECT } @@ -1977,11 +2110,12 @@ class coro_rpc_client { co_return rpc_error{errc::io_error, ret.first.message()}; } } - ELOG_TRACE << "rpc request send over, client_id: " << config_.client_id + ELOG_DEBUG << "rpc request send over, client_id: " << config_.client_id << ", cost time = " << (std::chrono::steady_clock::now() - tp) / std::chrono::microseconds(1) - << "us" << ", request ID: " << id; + << "us" + << ", request ID: " << id << ", bytes=" << ret.second; co_return rpc_error{}; } diff --git a/include/ylt/coro_rpc/impl/coro_rpc_server.hpp b/include/ylt/coro_rpc/impl/coro_rpc_server.hpp index e0a189361..3dff0431b 100644 --- a/include/ylt/coro_rpc/impl/coro_rpc_server.hpp +++ b/include/ylt/coro_rpc/impl/coro_rpc_server.hpp @@ -48,6 +48,9 @@ #include "ylt/coro_rpc/impl/nd_server_acceptor.hpp" #endif #include "ylt/coro_io/server_acceptor.hpp" +#ifdef YLT_ENABLE_URMA +#include "ylt/coro_io/urma/urma_rpc_env.hpp" +#endif #include "ylt/coro_rpc/impl/errno.h" #include "ylt/coro_rpc/impl/expected.hpp" namespace coro_rpc { @@ -183,6 +186,11 @@ class coro_rpc_server_base { else { init_acceptors(config.address, config.port); } +#ifdef YLT_ENABLE_URMA + if (config.urma_config) { + init_urma(config.urma_config.value()); + } +#endif #ifdef YLT_ENABLE_ND if constexpr (requires { config.nd_config; @@ -219,6 +227,24 @@ class coro_rpc_server_base { ibv_dev_lists_ = std::move(ibv_dev_lists); } #endif +#ifdef YLT_ENABLE_URMA + void init_urma(const coro_io::urma_socket_t::config_t &conf = {}) { + urma_config_ = conf; + } + + void init_urma_from_env_if_default() { + if (urma_config_.has_value()) return; +#ifdef YLT_ENABLE_IBV + if (ibv_config_.has_value()) return; +#endif +#ifdef YLT_ENABLE_SSL + if (use_ssl_) return; +#endif + if (auto urma_config = coro_io::try_make_urma_rpc_config()) { + urma_config_ = *urma_config; + } + } +#endif #ifdef YLT_ENABLE_ND void init_nd(const coro_io::nd_socket_t::config_t& conf = {}, uint16_t nd_port = 0, std::string nd_address = {}) { @@ -278,6 +304,9 @@ class coro_rpc_server_base { return make_error_future( coro_rpc::err_code{coro_rpc::errc::server_has_ran}); } +#ifdef YLT_ENABLE_URMA + init_urma_from_env_if_default(); +#endif for (size_t i = 0; i < acceptors_.size(); ++i) { auto& acceptor = acceptors_[i]; acceptor->set_io_threads_pool(&pool_); @@ -550,16 +579,14 @@ class coro_rpc_server_base { #endif if (!result.has_value()) { auto error = result.error(); - if (error == asio::error::operation_aborted) { - ELOG_INFO << "server was canceled:" << error.message(); - } - else { - ELOG_ERROR << "server accept failed:" << error.message(); - } if (error == asio::error::operation_aborted || error == asio::error::bad_descriptor) { + ELOG_DEBUG << "server accept stopped: " << error.message(); co_return coro_rpc::errc::operation_canceled; } + else { + ELOG_ERROR << "server accept failed:" << error.message(); + } continue; } coro_io::socket_wrapper_t& wrapper = result.value(); @@ -644,6 +671,23 @@ class coro_rpc_server_base { co_return init_ok; } #endif +#ifdef YLT_ENABLE_URMA + async_simple::coro::Lazy update_to_urma(coro_connection *conn) { + bool init_ok = true; + auto &wrapper = conn->socket_wrapper(); + ELOG_DEBUG << "URMA update_to_urma: conn_id=" << conn->get_connection_id() + << " remote=" << conn->get_remote_endpoint(); + try { + wrapper = {std::move(*wrapper.socket()), wrapper.get_executor(), + urma_config_.value_or(coro_io::urma_socket_t::config_t{})}; + ELOG_INFO << "URMA socket created for conn_id=" << conn->get_connection_id(); + } catch (...) { + ELOG_WARN << "URMA init urma connection failed, conn_id=" << conn->get_connection_id(); + init_ok = false; + } + co_return init_ok; + } +#endif async_simple::coro::Lazy start_one( std::shared_ptr conn) noexcept { @@ -674,6 +718,19 @@ class coro_rpc_server_base { } break; } +#endif +#ifdef YLT_ENABLE_URMA + if (urma_config_.has_value() && result.magic_number.size() == 1 && + result.magic_number[0] == + coro_io::urma_socket_t::urma_md5_first_header) { + ELOG_TRACE << "protocol is urma, try to update"; + auto result = co_await update_to_urma(conn.get()); + if (!result) { + ELOG_WARN << "urma init failed"; + co_return; + } + break; + } #endif if (connection_transfer_) { ELOG_TRACE @@ -717,6 +774,9 @@ class coro_rpc_server_base { std::vector> ibv_dev_lists_; std::atomic rr_index_ = 0; #endif +#ifdef YLT_ENABLE_URMA + std::optional urma_config_; +#endif std::function client_filter_; }; diff --git a/include/ylt/coro_rpc/impl/default_config/coro_rpc_config.hpp b/include/ylt/coro_rpc/impl/default_config/coro_rpc_config.hpp index 140ac0cf2..840c2e861 100644 --- a/include/ylt/coro_rpc/impl/default_config/coro_rpc_config.hpp +++ b/include/ylt/coro_rpc/impl/default_config/coro_rpc_config.hpp @@ -53,6 +53,9 @@ struct config_t { std::optional ibv_config = std::nullopt; std::vector> ibv_dev_lists; #endif +#ifdef YLT_ENABLE_URMA + std::optional urma_config = std::nullopt; +#endif #ifdef YLT_ENABLE_ND std::optional nd_config = std::nullopt; uint16_t nd_port = 0; diff --git a/include/ylt/coro_rpc/impl/protocol/coro_rpc_protocol.hpp b/include/ylt/coro_rpc/impl/protocol/coro_rpc_protocol.hpp index f566d5f01..9ba89f5fa 100644 --- a/include/ylt/coro_rpc/impl/protocol/coro_rpc_protocol.hpp +++ b/include/ylt/coro_rpc/impl/protocol/coro_rpc_protocol.hpp @@ -31,6 +31,9 @@ #include "struct_pack_protocol.hpp" #include "ylt/coro_io/coro_io.hpp" #include "ylt/coro_io/data_view.hpp" +#ifdef YLT_ENABLE_URMA +#include "ylt/coro_io/urma/urma_io.hpp" +#endif #include "ylt/coro_rpc/impl/context.hpp" #include "ylt/coro_rpc/impl/errno.h" #include "ylt/coro_rpc/impl/expected.hpp" @@ -137,8 +140,37 @@ struct coro_rpc_protocol { static async_simple::coro::Lazy read_payload( Socket& socket, req_header& req_head, std::string& buffer, coro_io::heterogeneous_buffer& attachment) { + co_return co_await read_payload(socket, req_head, buffer, attachment, + static_cast*>(nullptr)); + } + + template + static async_simple::coro::Lazy read_payload( + Socket& socket, req_header& req_head, std::string& buffer, + coro_io::heterogeneous_buffer& attachment, + context_info_t* context_info) { struct_pack::detail::resize(buffer, req_head.length); if (req_head.attach_length > 0) { +#ifdef YLT_ENABLE_URMA + if constexpr (std::is_same_v, + coro_io::urma_socket_t>) { + if (context_info != nullptr) { + if (!buffer.empty()) { + auto [body_ec, ignored] = + co_await coro_io::async_read(socket, asio::buffer(buffer)); + if (body_ec) co_return body_ec; + } + attachment = coro_io::heterogeneous_buffer{}; + auto [attachment_ec, views] = + co_await coro_io::detail::async_urma_read_views( + socket, req_head.attach_length); + if (attachment_ec) co_return attachment_ec; + context_info->set_request_attachment_views(std::move(views)); + co_return std::error_code{}; + } + } +#endif if constexpr (requires { socket.get_gpu_id(); }) { if (auto id = socket.get_gpu_id(); id >= 0) { if (attachment.size() < req_head.attach_length || @@ -293,4 +325,4 @@ context_info_t* get_context() { return detail::set_context(); } -} // namespace coro_rpc \ No newline at end of file +} // namespace coro_rpc diff --git a/include/ylt/coro_rpc/impl/rpc_execute.hpp b/include/ylt/coro_rpc/impl/rpc_execute.hpp index af97a2f03..8c41dd31d 100644 --- a/include/ylt/coro_rpc/impl/rpc_execute.hpp +++ b/include/ylt/coro_rpc/impl/rpc_execute.hpp @@ -25,6 +25,7 @@ #include "context.hpp" #include "coro_connection.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" #include "ylt/coro_rpc/impl/errno.h" #include "ylt/easylog.hpp" #include "ylt/struct_pack/compatible.hpp" @@ -79,6 +80,9 @@ inline std::pair execute( bool is_ok = true; constexpr size_t size = std::tuple_size_v; + auto deser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; if constexpr (size > 0) { is_ok = serialize_proto::deserialize_to(args, data); } @@ -105,11 +109,16 @@ inline std::pair execute( } } } + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_deserialize_request, + deser_begin, data.size()); + auto exec_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; if constexpr (std::is_void_v) { if constexpr (std::is_void_v) { if constexpr (has_coro_conn_v) { - // call void func(coro_conn, args...) std::apply(func, std::tuple_cat( std::forward_as_tuple( context_base( @@ -117,13 +126,11 @@ inline std::pair execute( std::move(args))); } else { - // call void func(args...) std::apply(func, std::move(args)); } } else { if constexpr (has_coro_conn_v) { - // call void self->func(coro_conn, args...) std::apply( func, std::tuple_cat( std::forward_as_tuple( @@ -132,27 +139,51 @@ inline std::pair execute( std::move(args))); } else { - // call void self->func(args...) std::apply(func, std::tuple_cat(std::forward_as_tuple(*self), std::move(args))); } } - return std::pair{err_code{}, serialize_proto::serialize()}; + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_handler_execute, + exec_begin, data.size()); + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto result = std::pair{err_code{}, serialize_proto::serialize()}; + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_result, + ser_begin, 0); + return result; } else { if constexpr (std::is_void_v) { - // call return_type func(args...) - - return std::pair{err_code{}, serialize_proto::serialize( - std::apply(func, std::move(args)))}; + auto ret = std::apply(func, std::move(args)); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_handler_execute, + exec_begin, data.size()); + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto buf = serialize_proto::serialize(std::move(ret)); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_result, + ser_begin, buf.size()); + return std::pair{err_code{}, std::move(buf)}; } else { - // call return_type self->func(args...) - - return std::pair{err_code{}, - serialize_proto::serialize(std::apply( - func, std::tuple_cat(std::forward_as_tuple(*self), - std::move(args))))}; + auto ret = std::apply(func, std::tuple_cat(std::forward_as_tuple(*self), + std::move(args))); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_handler_execute, + exec_begin, data.size()); + auto ser_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto buf = serialize_proto::serialize(std::move(ret)); + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::server_serialize_result, + ser_begin, buf.size()); + return std::pair{err_code{}, std::move(buf)}; } } } diff --git a/src/coro_rpc/examples/CMakeLists.txt b/src/coro_rpc/examples/CMakeLists.txt index 7180da6b6..24938fa1b 100644 --- a/src/coro_rpc/examples/CMakeLists.txt +++ b/src/coro_rpc/examples/CMakeLists.txt @@ -4,7 +4,11 @@ add_subdirectory(file_transfer) if (YLT_HAVE_IBVERBS) add_subdirectory(rdma_example) endif() +if (YLT_ENABLE_URMA) + add_subdirectory(urma_example) +endif() +add_subdirectory(urma_benchmark) if (CORO_RPC_USE_OTHER_RPC) add_subdirectory(user_defined_rpc_protocol/rest_rpc) -endif() \ No newline at end of file +endif() diff --git a/src/coro_rpc/examples/urma_benchmark/CMakeLists.txt b/src/coro_rpc/examples/urma_benchmark/CMakeLists.txt new file mode 100644 index 000000000..6b1c11469 --- /dev/null +++ b/src/coro_rpc/examples/urma_benchmark/CMakeLists.txt @@ -0,0 +1,12 @@ +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/output/examples/coro_rpc) +if (YLT_ENABLE_URMA) + add_executable(coro_rpc_urma_benchmark + urma_benchmark.cpp + ) +else() + add_custom_target(coro_rpc_urma_benchmark + COMMAND ${CMAKE_COMMAND} -E echo + "coro_rpc_urma_benchmark requires -DYLT_ENABLE_URMA=ON. Re-run CMake configure with URMA enabled." + COMMAND ${CMAKE_COMMAND} -E false + ) +endif() diff --git a/src/coro_rpc/examples/urma_benchmark/README.md b/src/coro_rpc/examples/urma_benchmark/README.md new file mode 100644 index 000000000..103744d5d --- /dev/null +++ b/src/coro_rpc/examples/urma_benchmark/README.md @@ -0,0 +1,214 @@ +# URMA RPC Benchmark + +`coro_rpc_urma_benchmark` is a small benchmark tool for URMA-based coro_rpc. +It can run on the same node or across two nodes, and supports low-load latency +and throughput tests. + +## Build + +Configure with URMA enabled and examples enabled: + +```bash +cmake -S . -B build -DYLT_ENABLE_URMA=ON -DBUILD_EXAMPLES=ON +``` + +Build only this benchmark target: + +```bash +cmake --build build --target coro_rpc_urma_benchmark -j +``` + +The binary is generated at: + +```bash +build/output/examples/coro_rpc/coro_rpc_urma_benchmark +``` + +If the target prints this message: + +```text +coro_rpc_urma_benchmark requires -DYLT_ENABLE_URMA=ON. Re-run CMake configure with URMA enabled. +``` + +re-run the CMake configure command above. If CMake was already configured +before this example was added, reconfigure the build directory once. + +## Same-Node Test + +Start the server: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark server \ + --host 127.0.0.1 \ + --port 9001 \ + --device bonding_dev_0 +``` + +Run both latency and throughput tests: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark client \ + --host 127.0.0.1 \ + --port 9001 \ + --device bonding_dev_0 \ + --mode both \ + --payload 64 \ + --latency-iters 10000 \ + --connections 64 \ + --duration 10 +``` + +## Cross-Node Test + +On the server node: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark server \ + --host 0.0.0.0 \ + --port 9001 \ + --device bonding_dev_0 +``` + +On the client node: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark client \ + --host \ + --port 9001 \ + --device bonding_dev_0 \ + --mode both \ + --payload 64 \ + --latency-iters 10000 \ + --connections 128 \ + --duration 30 +``` + +Replace `` with the server address reachable from the client. + +## Test Modes + +Latency only: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark client \ + --host \ + --port 9001 \ + --mode latency \ + --payload 64 \ + --latency-iters 10000 +``` + +Throughput only: + +```bash +./build/output/examples/coro_rpc/coro_rpc_urma_benchmark client \ + --host \ + --port 9001 \ + --mode throughput \ + --payload 4096 \ + --connections 128 \ + --duration 30 +``` + +## Common Options + +```text +--host Server listen/connect host. Server default 0.0.0.0, + client default 127.0.0.1 +--port Server port. Default 9001 +--transport rpc uses coro_rpc. raw uses urma_socket directly. + Default rpc +--device URMA device. Default bonding_dev_0 +--eid-index URMA EID index. Default 0 +--payload Echo payload size. Default 64 +--buffer-size URMA SEND chunk size. Default 4096 for CTP +--queue-depth URMA send/recv queue depth. Default 64 +--max-memory-mib URMA buffer pool memory per process. Default 256, + auto-raised when payload/connections need more +--profile Enable in-memory stage latency profiling. Default off. +--profile-sample-rate Record one sample every n events per stage/thread. + Default 1. +--log Default info +``` + +Client-only options: + +```text +--mode +--rpc Echo returns the payload. Sink returns only + payload size. Attach_sink sends payload as a + request attachment and uses the URMA segmented + attachment fast path on the server. +--latency-iters +--warmup-iters +--connections +--pipeline-depth Outstanding RPC calls per connection in throughput mode. + Default 1. +--concurrency Compatibility option; URMA throughput uses one worker per connection. +--duration +--raw-report-interval Raw server report interval. Default 0 disables + periodic server-side reports. +--client-threads +``` + +Server-only option: + +```text +--server-threads +``` + +## Notes + +- The default URMA transport type is CTP. +- Server mode defaults to `--host 0.0.0.0`, so cross-node tests work without + explicitly overriding the listen address. Client mode defaults to + `127.0.0.1`. +- The default `--buffer-size` is `4096`, matching the documented 4KB max send + packet size for `bonding_dev_0` CTP. +- The URMA buffer pool allocates one large contiguous memory block and registers + it as one segment, then splits it into fixed-size buffers. This avoids doing + one `urma_register_seg` call per 4KB buffer during startup. +- The benchmark sizes the pool from `--max-memory-mib`, `--connections`, + `--queue-depth`, and `--payload`. If the explicit memory value is too small, + it is raised automatically for the benchmark process. +- `--connections` controls the number of RPC client connections and throughput + workers. Each throughput worker owns one `coro_rpc_client`. +- `--pipeline-depth` keeps multiple outstanding RPC calls on each connection in + throughput mode. Use it with `--rpc attach_sink` to test whether one + request/response at a time is limiting throughput. +- `--concurrency` is retained as a compatibility option. The URMA benchmark does + not run multiple throughput coroutines on the same connection because that can + overrun the current URMA send/recv credit model and produce `WR_FLUSH_ERR`. +- Latency output is in microseconds and includes avg/min/p50/p90/p99/p999/max. +- Throughput output includes request rate and payload MiB/s. +- With `--profile`, the client prints one final `urma_profile` block after the + selected benchmark modes finish. Each stage is sampled in thread-local memory + and reports avg/p99/p9999/max in microseconds. Use + `--profile-sample-rate` for long throughput runs to reduce memory and timing + overhead. +- Use `--rpc sink` to remove the large echo response from the server side. If + sink throughput is much higher than echo, the bottleneck is response + serialization/sending. If sink is also low, focus on request read, RPC + dispatch, and URMA receive/polling. +- Use `--rpc attach_sink` to test the URMA RPC fast path. The request payload is + sent as attachment data, and the server handler reads segmented + `owned_data_view`s instead of copying the attachment into a contiguous + `std::string`. +- Use `--transport raw` to bypass coro_rpc and struct_pack. Raw mode sends the + payload from client to server without a response, so it measures URMA ingress + throughput more directly. The raw server does not print periodic throughput by + default so stdout does not affect the measurement: + +```bash +./coro_rpc_urma_benchmark server --transport raw --host 0.0.0.0 --port 9001 +./coro_rpc_urma_benchmark client --transport raw --host \ + --payload 1048576 --connections 64 --queue-depth 128 --duration 30 +``` + +Example RPC fast-path throughput test: + +```bash +./coro_rpc_urma_benchmark client --host --mode throughput \ + --rpc attach_sink --payload 1048576 --connections 64 --pipeline-depth 8 \ + --queue-depth 128 --duration 30 +``` diff --git a/src/coro_rpc/examples/urma_benchmark/urma_benchmark.cpp b/src/coro_rpc/examples/urma_benchmark/urma_benchmark.cpp new file mode 100644 index 000000000..72022ae43 --- /dev/null +++ b/src/coro_rpc/examples/urma_benchmark/urma_benchmark.cpp @@ -0,0 +1,846 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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 +#include +#include +#include +#include +#include +#include + +#include "asio/ip/address.hpp" +#include "async_simple/Promise.h" +#include "async_simple/coro/Collect.h" +#include "async_simple/coro/Lazy.h" +#include "async_simple/coro/SyncAwait.h" +#include "ylt/coro_io/coro_io.hpp" +#include "ylt/coro_io/io_context_pool.hpp" +#include "ylt/coro_io/urma/urma_device.hpp" +#include "ylt/coro_io/urma/urma_benchmark_profile.hpp" +#include "ylt/coro_io/urma/urma_io.hpp" +#include "ylt/coro_io/urma/urma_socket.hpp" +#include "ylt/coro_rpc/coro_rpc_client.hpp" +#include "ylt/coro_rpc/coro_rpc_server.hpp" +#include "ylt/easylog.hpp" +#include "ylt/easylog/record.hpp" + +using namespace async_simple::coro; +using namespace coro_rpc; +using namespace std::chrono_literals; + +std::string_view bench_echo(std::string_view data) { return data; } +uint64_t bench_sink(std::string_view data) { return data.size(); } +void bench_attachment_sink(coro_rpc::context ctx) { + uint64_t size = 0; + auto views = ctx.get_context_info()->get_request_attachment_views(); + if (!views.empty()) { + for (auto& view : views) size += view.size(); + } + else { + size = ctx.get_context_info()->get_request_attachment2().size(); + } + ctx.response_msg(size); +} + +struct options_t { + std::string role = "client"; + std::string transport = "rpc"; + std::string mode = "both"; + std::string rpc = "echo"; + std::string host = "127.0.0.1"; + uint16_t port = 9001; + std::string device = "bonding_dev_0"; + int eid_index = 0; + uint32_t payload_size = 64; + uint32_t latency_iters = 10000; + uint32_t warmup_iters = 1000; + uint32_t concurrency = 64; + uint32_t connections = 64; + uint32_t pipeline_depth = 1; + uint32_t duration_seconds = 10; + uint32_t raw_report_interval_seconds = 0; + uint32_t server_threads = std::max(1u, std::thread::hardware_concurrency()); + uint32_t client_threads = std::max(1u, std::thread::hardware_concurrency()); + uint16_t queue_depth = 64; + uint32_t buffer_size = 4 * 1024; + uint64_t max_memory_usage = 256ull * 1024 * 1024; + bool profile = false; + bool use_urma = true; + uint32_t profile_sample_rate = 1; + easylog::Severity log_level = easylog::Severity::WARNING; +}; + +void status(std::string_view message) { + std::cout << "[urma_benchmark] " << message << std::endl; +} + +uint64_t effective_pool_memory_usage(const options_t& opt); + +void set_process_env(const char* name, const std::string& value) { +#ifdef _WIN32 + _putenv_s(name, value.c_str()); +#else + setenv(name, value.c_str(), 1); +#endif +} + +void configure_urma_rpc_env(const options_t& opt) { + set_process_env("URMA_RPC_ENABLE", "1"); + set_process_env("URMA_RPC_DEVICE", opt.device); + set_process_env("URMA_RPC_EID_INDEX", std::to_string(opt.eid_index)); + set_process_env("URMA_RPC_CQ_SIZE", std::to_string(opt.queue_depth * 2 + 8)); + set_process_env("URMA_RPC_RECV_BUFFER_CNT", std::to_string(opt.queue_depth)); + set_process_env("URMA_RPC_SEND_BUFFER_CNT", std::to_string(opt.queue_depth)); + set_process_env("URMA_RPC_BUFFER_SIZE", std::to_string(opt.buffer_size)); + set_process_env("URMA_RPC_MAX_MEMORY_USAGE", + std::to_string(effective_pool_memory_usage(opt))); + set_process_env("URMA_RPC_TP_TYPE", "ctp"); +} + +void print_usage(const char* program) { + std::cout + << "Usage:\n" + << " " << program << " server [options]\n" + << " " << program << " client [options]\n\n" + << "Common options:\n" + << " --host Server listen/connect host. Default 127.0.0.1\n" + << " --port Server port. Default 9001\n" + << " --transport rpc uses coro_rpc; raw uses urma_socket directly. Default rpc\n" + << " --device URMA device. Default bonding_dev_0\n" + << " --eid-index URMA EID index. Default 0\n" + << " --payload Echo payload size. Default 64\n" + << " --buffer-size URMA SEND chunk size. Default 4096 for CTP\n" + << " --queue-depth URMA send/recv queue depth. Default 64\n" + << " --max-memory-mib URMA buffer pool memory per process. Default 256, auto-raised when needed\n" + << " --no-urma Use TCP instead of URMA RPC. Default URMA\n" + << " --profile Enable in-memory stage latency profiling. Default off\n" + << " --profile-sample-rate Record one sample every n events per stage/thread. Default 1\n" + << " --log Default info\n\n" + << "Client options:\n" + << " --mode Default both\n" + << " --rpc echo returns payload; sink returns payload size; attach_sink uses request attachment. Default echo\n" + << " --latency-iters Low-load serial requests. Default 10000\n" + << " --warmup-iters Warmup requests per client. Default 1000\n" + << " --concurrency Compatibility option; URMA throughput uses one worker per connection\n" + << " --connections URMA RPC connections and throughput workers. Default 64\n" + << " --pipeline-depth Outstanding RPC calls per connection in throughput mode. Default 1\n" + << " --duration Throughput duration. Default 10\n" + << " --raw-report-interval Raw server report interval. Default 0 disables periodic reports\n" + << " --client-threads Client executor threads. Default hardware\n\n" + << "Server options:\n" + << " --server-threads Server threads. Default hardware\n\n" + << "Examples:\n" + << " same node server: " << program + << " server --host 127.0.0.1 --port 9001\n" + << " same node client: " << program + << " client --host 127.0.0.1 --port 9001 --mode both\n" + << " cross node server: " << program + << " server --host 0.0.0.0 --port 9001\n" + << " cross node client: " << program + << " client --host --port 9001 --mode throughput\n"; +} + +uint64_t parse_u64(std::string_view value, std::string_view name) { + char* end = nullptr; + auto str = std::string(value); + auto result = std::strtoull(str.c_str(), &end, 10); + if (end == str.c_str() || *end != '\0') { + throw std::invalid_argument("invalid numeric option: " + std::string(name)); + } + return result; +} + +easylog::Severity parse_log_level(std::string_view value) { + if (value == "trace") return easylog::Severity::TRACE; + if (value == "debug") return easylog::Severity::DEBUG; + if (value == "info") return easylog::Severity::INFO; + if (value == "warn") return easylog::Severity::WARN; + if (value == "error") return easylog::Severity::ERROR; + throw std::invalid_argument("invalid --log value"); +} + +options_t parse_options(int argc, char** argv) { + options_t opt; + if (argc >= 2) opt.role = argv[1]; + if (opt.role != "server" && opt.role != "client") { + print_usage(argv[0]); + throw std::invalid_argument("role must be server or client"); + } + + bool host_was_set = false; + for (int i = 2; i < argc; ++i) { + std::string_view key = argv[i]; + auto require_value = [&]() -> std::string_view { + if (i + 1 >= argc) { + throw std::invalid_argument("missing value for " + std::string(key)); + } + return argv[++i]; + }; + + if (key == "--help" || key == "-h") { + print_usage(argv[0]); + std::exit(0); + } + else if (key == "--host") { + opt.host = require_value(); + host_was_set = true; + } + else if (key == "--port") { + opt.port = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--transport") { + opt.transport = require_value(); + } + else if (key == "--device") { + opt.device = require_value(); + } + else if (key == "--eid-index") { + opt.eid_index = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--payload") { + opt.payload_size = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--buffer-size") { + opt.buffer_size = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--queue-depth") { + opt.queue_depth = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--max-memory-mib") { + opt.max_memory_usage = parse_u64(require_value(), key) * 1024 * 1024; + } + else if (key == "--profile") { + opt.profile = true; + } + else if (key == "--no-urma") { + opt.use_urma = false; + } + else if (key == "--profile-sample-rate") { + opt.profile_sample_rate = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--mode") { + opt.mode = require_value(); + } + else if (key == "--rpc") { + opt.rpc = require_value(); + } + else if (key == "--latency-iters") { + opt.latency_iters = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--warmup-iters") { + opt.warmup_iters = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--concurrency") { + opt.concurrency = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--connections") { + opt.connections = static_cast(parse_u64(require_value(), key)); + } + else if (key == "--pipeline-depth") { + opt.pipeline_depth = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--duration") { + opt.duration_seconds = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--raw-report-interval") { + opt.raw_report_interval_seconds = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--server-threads") { + opt.server_threads = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--client-threads") { + opt.client_threads = + static_cast(parse_u64(require_value(), key)); + } + else if (key == "--log") { + opt.log_level = parse_log_level(require_value()); + } + else { + throw std::invalid_argument("unknown option: " + std::string(key)); + } + } + + if (opt.mode != "latency" && opt.mode != "throughput" && + opt.mode != "both") { + throw std::invalid_argument("--mode must be latency, throughput, or both"); + } + if (opt.transport != "rpc" && opt.transport != "raw") { + throw std::invalid_argument("--transport must be rpc or raw"); + } + if (opt.rpc != "echo" && opt.rpc != "sink" && opt.rpc != "attach_sink") { + throw std::invalid_argument("--rpc must be echo, sink, or attach_sink"); + } + opt.connections = std::max(opt.connections, 1); + opt.pipeline_depth = std::max(opt.pipeline_depth, 1); + opt.concurrency = std::max(opt.concurrency, opt.connections); + opt.queue_depth = std::max(opt.queue_depth, 1); + opt.profile_sample_rate = std::max(opt.profile_sample_rate, 1); + if (opt.role == "server" && !host_was_set) { + opt.host = "0.0.0.0"; + } + return opt; +} + +uint64_t saturated_mul(uint64_t lhs, uint64_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return std::numeric_limits::max(); + } + return lhs * rhs; +} + +uint64_t ceil_div(uint64_t value, uint64_t divisor) { + return divisor == 0 ? value : (value + divisor - 1) / divisor; +} + +uint64_t estimated_pool_memory_usage(const options_t& opt) { + auto buffer_size = std::max(opt.buffer_size, 1); + auto connections = std::max(opt.connections, 1); + auto queue_depth = std::max(opt.queue_depth, 1); + auto payload_chunks = + ceil_div(static_cast(opt.payload_size) + 1024, buffer_size); + + // Each connection keeps recv WRs posted and can also cache completed recv + // buffers while the RPC layer consumes a large message. Keep enough slack for + // send completions and protocol framing without requiring users to guess. + auto payload_slack = std::min(payload_chunks, queue_depth); + auto buffers_per_connection = + queue_depth * 2 + queue_depth + payload_slack + 16; + auto required_buffers = + saturated_mul(connections, buffers_per_connection) + 1024; + return saturated_mul(required_buffers, buffer_size); +} + +uint64_t effective_pool_memory_usage(const options_t& opt) { + return std::max(opt.max_memory_usage, estimated_pool_memory_usage(opt)); +} + +coro_io::urma_socket_t::config_t make_urma_config(const options_t& opt) { + return coro_io::urma_socket_t::config_t{ + .cq_size = static_cast(opt.queue_depth * 2 + 8), + .recv_buffer_cnt = opt.queue_depth, + .send_buffer_cnt = opt.queue_depth, + .buffer_size = opt.buffer_size, + .max_memory_usage = effective_pool_memory_usage(opt), + .device_name = opt.device, + .eid_index = opt.eid_index, + .tp_type = URMA_CTP}; +} + +bool init_global_urma(const options_t& opt) { + status("initializing global URMA device"); + auto pool_memory = effective_pool_memory_usage(opt); + if (pool_memory > opt.max_memory_usage) { + std::cout << "[urma_benchmark] auto raise URMA buffer pool memory from " + << opt.max_memory_usage / 1024 / 1024 << " MiB to " + << pool_memory / 1024 / 1024 + << " MiB for payload/connections/queue-depth" << std::endl; + } + auto device = coro_io::get_global_urma_device(coro_io::urma_init_config_t{ + .dev_name = opt.device, + .buffer_pool_config = + { + .buffer_size = opt.buffer_size, + .max_memory_usage = pool_memory, + .idle_timeout = 5s, + }, + .eid_index = opt.eid_index}); + if (!device || !device->is_valid() || !device->get_buffer_pool()) { + std::cerr << "[urma_benchmark] failed to initialize global URMA device" + << std::endl; + return false; + } + auto pool = device->get_buffer_pool(); + std::cout << "[urma_benchmark] global URMA device initialized, pool_buffers=" + << pool->total_buffer_count() + << ", pool_free=" << pool->free_buffer_count() + << ", pool_memory_mib=" + << pool->total_memory_size() / 1024 / 1024 << std::endl; + return true; +} + +Lazy connect_client(coro_rpc_client& client, const options_t& opt, + uint32_t client_index = 0) { + std::cout << "[urma_benchmark] client " << client_index + << " using URMA RPC auto configuration" << std::endl; + std::cout << "[urma_benchmark] client " << client_index << " connecting to " + << opt.host << ":" << opt.port << std::endl; + auto ec = co_await client.connect(opt.host, std::to_string(opt.port), 30s); + if (ec) { + ELOG_ERROR << "connect failed: " << ec.message() + << ". Check that the server process is running, listening on " + << "0.0.0.0 or the requested NIC address, and that the TCP " + << "handshake port is reachable: " << opt.host << ":" + << opt.port; + co_return false; + } + std::cout << "[urma_benchmark] client " << client_index << " connected" + << std::endl; + co_return true; +} + +Lazy issue_rpc_call(coro_rpc_client& client, const std::string& payload, + std::string_view rpc, bool profile_call = true); + +Lazy warmup(coro_rpc_client& client, const std::string& payload, + uint32_t count, std::string_view rpc, + uint32_t client_index = 0) { + if (count == 0) co_return true; + std::cout << "[urma_benchmark] client " << client_index + << " warmup start, iterations=" << count << std::endl; + for (uint32_t i = 0; i < count; ++i) { + if (!(co_await issue_rpc_call(client, payload, rpc, false))) { + ELOG_ERROR << "warmup failed at iteration " << i; + co_return false; + } + } + std::cout << "[urma_benchmark] client " << client_index << " warmup done" + << std::endl; + co_return true; +} + +Lazy issue_rpc_call(coro_rpc_client& client, const std::string& payload, + std::string_view rpc, bool profile_call) { + auto profile_begin = profile_call && coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto payload_size = payload.size(); + bool ok = false; + if (rpc == "attach_sink") { + auto result = co_await client.call( + request_config_t{30s, payload, {}, -1, -1}); + ok = result && result.value() == payload.size(); + if (profile_call) { + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::benchmark_rpc_call, + profile_begin, payload_size); + } + co_return ok; + } + if (rpc == "sink") { + auto result = co_await client.call_for(30s, payload); + ok = result && result.value() == payload.size(); + if (profile_call) { + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::benchmark_rpc_call, + profile_begin, payload_size); + } + co_return ok; + } + auto result = co_await client.call_for(30s, payload); + ok = result && result.value().size() == payload.size(); + if (profile_call) { + coro_io::urma_benchmark_profile::record_since_with_size( + coro_io::urma_benchmark_profile::stage::benchmark_rpc_call, + profile_begin, payload_size); + } + co_return ok; +} + +void print_latency_result(std::vector& samples) { + if (samples.empty()) { + std::cout << "latency: no samples\n"; + return; + } + std::sort(samples.begin(), samples.end()); + auto percentile = [&](double p) { + auto index = static_cast((samples.size() - 1) * p); + return samples[index]; + }; + auto sum = std::accumulate(samples.begin(), samples.end(), uint64_t{0}); + auto avg = static_cast(sum) / static_cast(samples.size()); + std::cout << std::fixed << std::setprecision(2) + << "latency_us count=" << samples.size() << " avg=" << avg + << " min=" << samples.front() << " p50=" << percentile(0.50) + << " p90=" << percentile(0.90) << " p99=" << percentile(0.99) + << " p999=" << percentile(0.999) << " max=" << samples.back() + << "\n"; +} + +Lazy run_latency(const options_t& opt, const std::string& payload) { + status("latency test starting"); + coro_rpc_client client; + if (!(co_await connect_client(client, opt))) co_return; + if (!(co_await warmup(client, payload, opt.warmup_iters, opt.rpc))) + co_return; + + std::vector samples; + samples.reserve(opt.latency_iters); + std::cout << "[urma_benchmark] latency measurement start, iterations=" + << opt.latency_iters << std::endl; + coro_io::urma_benchmark_profile::configure(opt.profile, + opt.profile_sample_rate); + for (uint32_t i = 0; i < opt.latency_iters; ++i) { + auto begin = std::chrono::steady_clock::now(); + bool ok = co_await issue_rpc_call(client, payload, opt.rpc); + auto end = std::chrono::steady_clock::now(); + if (!ok) { + ELOG_ERROR << "latency call failed at iteration " << i; + break; + } + samples.push_back(static_cast( + std::chrono::duration_cast(end - begin) + .count())); + } + coro_io::urma_benchmark_profile::configure(false, opt.profile_sample_rate); + print_latency_result(samples); + status("latency test finished"); +} + +struct worker_result_t { + uint64_t requests = 0; + uint64_t errors = 0; + uint64_t bytes = 0; +}; + +Lazy throughput_worker(coro_rpc_client& client, + const std::string& payload, + std::string_view rpc, + uint32_t pipeline_depth, + std::chrono::steady_clock::time_point + deadline) { + worker_result_t stat; + while (std::chrono::steady_clock::now() < deadline) { + std::vector> calls; + calls.reserve(pipeline_depth); + for (uint32_t i = 0; i < pipeline_depth && + std::chrono::steady_clock::now() < deadline; + ++i) { + calls.push_back(issue_rpc_call(client, payload, rpc)); + } + if (calls.empty()) break; + auto results = co_await collectAll(std::move(calls)); + for (auto& item : results) { + bool ok = item.value(); + if (!ok) { + ++stat.errors; + continue; + } + ++stat.requests; + stat.bytes += payload.size(); + } + } + co_return stat; +} + +Lazy run_throughput(const options_t& opt, const std::string& payload) { + status("throughput test preparing clients"); + std::vector> clients; + clients.reserve(opt.connections); + for (uint32_t i = 0; i < opt.connections; ++i) { + auto client = std::make_unique(); + if (!(co_await connect_client(*client, opt, i))) co_return; + if (!(co_await warmup(*client, payload, opt.warmup_iters, opt.rpc, i))) + co_return; + clients.push_back(std::move(client)); + } + + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(opt.duration_seconds); + std::vector> workers; + workers.reserve(clients.size()); + for (auto& client : clients) { + workers.push_back(throughput_worker(*client, payload, opt.rpc, + opt.pipeline_depth, deadline)); + } + + std::cout << "[urma_benchmark] throughput measurement start, duration_s=" + << opt.duration_seconds << ", workers=" << workers.size() + << ", connections=" << opt.connections + << ", pipeline_depth=" << opt.pipeline_depth + << ", requested_concurrency=" << opt.concurrency + << ". URMA throughput uses one serial worker per connection." + << std::endl; + coro_io::urma_benchmark_profile::configure(opt.profile, + opt.profile_sample_rate); + auto begin = std::chrono::steady_clock::now(); + auto results = co_await collectAll(std::move(workers)); + auto end = std::chrono::steady_clock::now(); + coro_io::urma_benchmark_profile::configure(false, opt.profile_sample_rate); + + worker_result_t total; + for (auto& item : results) { + auto result = item.value(); + total.requests += result.requests; + total.errors += result.errors; + total.bytes += result.bytes; + } + + auto seconds = + std::chrono::duration_cast>(end - begin) + .count(); + auto rps = static_cast(total.requests) / seconds; + auto mbps = static_cast(total.bytes) / seconds / 1024.0 / 1024.0; + std::cout << std::fixed << std::setprecision(2) + << "throughput duration_s=" << seconds + << " requests=" << total.requests << " errors=" << total.errors + << " rps=" << rps << " payload_mib_per_s=" << mbps + << " workers=" << clients.size() + << " connections=" << opt.connections << "\n"; + status("throughput test finished"); +} + +struct raw_result_t { + uint64_t messages = 0; + uint64_t errors = 0; + uint64_t bytes = 0; +}; + +Lazy raw_server_session(asio::ip::tcp::socket tcp_socket, + const options_t& opt) { + coro_io::urma_socket_t socket(coro_io::get_global_executor(), + make_urma_config(opt)); + auto ec = co_await socket.accept(std::move(tcp_socket)); + if (ec) { + ELOG_ERROR << "raw URMA accept failed: " << ec.message(); + co_return; + } + + std::vector buffer(opt.payload_size); + uint64_t messages = 0; + uint64_t bytes = 0; + auto last = std::chrono::steady_clock::now(); + for (;;) { + auto [read_ec, read_size] = + co_await coro_io::async_read(socket, asio::buffer(buffer)); + if (read_ec) { + ELOG_DEBUG << "raw URMA session closed: " << read_ec.message() + << ", messages=" << messages << ", bytes=" << bytes; + co_return; + } + ++messages; + bytes += read_size; + if (opt.raw_report_interval_seconds == 0) continue; + auto now = std::chrono::steady_clock::now(); + if (now - last >= std::chrono::seconds(opt.raw_report_interval_seconds)) { + auto seconds = + std::chrono::duration_cast>(now - last) + .count(); + std::cout << "[urma_benchmark] raw server ingress payload_mib_per_s=" + << static_cast(bytes) / seconds / 1024.0 / 1024.0 + << ", messages=" << messages << std::endl; + messages = 0; + bytes = 0; + last = now; + } + } +} + +Lazy run_raw_server(const options_t& opt) { + status("raw URMA server starting"); + if (!init_global_urma(opt)) co_return; + auto executor = coro_io::get_global_executor(opt.server_threads); + asio::ip::tcp::endpoint endpoint(asio::ip::make_address(opt.host), opt.port); + asio::ip::tcp::acceptor acceptor(executor->context(), endpoint); + std::cout << "raw URMA server listening on " << opt.host << ":" << opt.port + << ", payload=" << opt.payload_size + << ", buffer_size=" << opt.buffer_size + << ", queue_depth=" << opt.queue_depth << std::endl; + for (;;) { + asio::ip::tcp::socket tcp_socket(executor->context()); + auto ec = co_await coro_io::async_accept(acceptor, tcp_socket); + if (ec) { + ELOG_ERROR << "raw accept failed: " << ec.message(); + co_return; + } + raw_server_session(std::move(tcp_socket), opt).start([](auto&&) { + }); + } +} + +Lazy raw_client_worker(const options_t& opt, + const std::string& payload, + std::chrono::steady_clock::time_point + deadline, + uint32_t client_index) { + raw_result_t stat; + coro_io::urma_socket_t socket(coro_io::get_global_executor(), + make_urma_config(opt)); + auto ec = co_await socket.connect(opt.host, std::to_string(opt.port)); + if (ec) { + ELOG_ERROR << "raw client " << client_index << " connect failed: " + << ec.message(); + stat.errors++; + co_return stat; + } + while (std::chrono::steady_clock::now() < deadline) { + auto profile_begin = coro_io::urma_benchmark_profile::enabled() + ? coro_io::urma_benchmark_profile::now_ns() + : 0; + auto [write_ec, written] = + co_await coro_io::async_write(socket, asio::buffer(payload)); + coro_io::urma_benchmark_profile::record_since( + coro_io::urma_benchmark_profile::stage::raw_client_write, + profile_begin); + if (write_ec || written != payload.size()) { + ++stat.errors; + continue; + } + ++stat.messages; + stat.bytes += written; + } + co_return stat; +} + +Lazy run_raw_client(const options_t& opt, const std::string& payload) { + status("raw URMA throughput test starting"); + if (!init_global_urma(opt)) co_return; + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(opt.duration_seconds); + std::vector> workers; + workers.reserve(opt.connections); + for (uint32_t i = 0; i < opt.connections; ++i) { + workers.push_back(raw_client_worker(opt, payload, deadline, i)); + } + + auto begin = std::chrono::steady_clock::now(); + coro_io::urma_benchmark_profile::configure(opt.profile, + opt.profile_sample_rate); + auto results = co_await collectAll(std::move(workers)); + auto end = std::chrono::steady_clock::now(); + coro_io::urma_benchmark_profile::configure(false, opt.profile_sample_rate); + raw_result_t total; + for (auto& item : results) { + auto result = item.value(); + total.messages += result.messages; + total.errors += result.errors; + total.bytes += result.bytes; + } + auto seconds = + std::chrono::duration_cast>(end - begin) + .count(); + std::cout << std::fixed << std::setprecision(2) + << "raw_throughput duration_s=" << seconds + << " messages=" << total.messages << " errors=" << total.errors + << " payload_mib_per_s=" + << static_cast(total.bytes) / seconds / 1024.0 / 1024.0 + << " connections=" << opt.connections << "\n"; + status("raw URMA throughput test finished"); +} + +int run_server(const options_t& opt) { + if (opt.transport == "raw") { + syncAwait(run_raw_server(opt)); + return 0; + } + std::cout << "[urma_benchmark] server starting, listen=" << opt.host << ":" + << opt.port << ", device=" << opt.device + << ", eid_index=" << opt.eid_index + << ", buffer_size=" << opt.buffer_size + << ", queue_depth=" << opt.queue_depth + << ", max_memory_mib=" + << effective_pool_memory_usage(opt) / 1024 / 1024 << std::endl; + if (opt.use_urma) { + configure_urma_rpc_env(opt); + if (!init_global_urma(opt)) return 1; + status("constructing RPC server with URMA RPC auto configuration"); + } else { + status("constructing RPC server with TCP (URMA disabled)"); + } + coro_rpc_server server(opt.server_threads, opt.port, opt.host); + server.register_handler(); + server.register_handler(); + server.register_handler(); + std::cout << (opt.use_urma ? "URMA" : "TCP") + << " RPC benchmark server listening on " << opt.host << ":" + << opt.port << ", device=" << opt.device + << ", eid_index=" << opt.eid_index + << ", buffer_size=" << opt.buffer_size + << ", queue_depth=" << opt.queue_depth + << ", max_memory_mib=" + << effective_pool_memory_usage(opt) / 1024 / 1024 << std::endl; + status("entering server event loop"); + return !server.start(); +} + +int run_client(const options_t& opt) { + std::cout << "[urma_benchmark] client starting, target=" << opt.host << ":" + << opt.port << ", mode=" << opt.mode + << ", transport=" << opt.transport + << ", rpc=" << opt.rpc + << ", payload=" << opt.payload_size + << ", buffer_size=" << opt.buffer_size + << ", queue_depth=" << opt.queue_depth + << ", max_memory_mib=" + << effective_pool_memory_usage(opt) / 1024 / 1024 + << ", connections=" << opt.connections + << ", pipeline_depth=" << opt.pipeline_depth + << ", total_outstanding=" + << static_cast(opt.connections) * opt.pipeline_depth + << ", concurrency=" << opt.concurrency + << ", use_urma=" << (opt.use_urma ? "on" : "off") + << ", profile=" << (opt.profile ? "on" : "off") + << ", profile_sample_rate=" << opt.profile_sample_rate + << std::endl; + coro_io::get_global_executor(opt.client_threads); + std::string payload(opt.payload_size, 'x'); + if (opt.transport == "raw") { + syncAwait(run_raw_client(opt, payload)); + if (opt.profile) coro_io::urma_benchmark_profile::print(std::cout); + return 0; + } + if (opt.use_urma) { + configure_urma_rpc_env(opt); + if (!init_global_urma(opt)) return 1; + } + std::cout << (opt.use_urma ? "URMA" : "TCP") + << " RPC benchmark client target " << opt.host << ":" + << opt.port << ", mode=" << opt.mode + << ", rpc=" << opt.rpc + << ", payload=" << opt.payload_size + << ", buffer_size=" << opt.buffer_size + << ", queue_depth=" << opt.queue_depth << std::endl; + + if (opt.mode == "latency" || opt.mode == "both") { + syncAwait(run_latency(opt, payload)); + } + if (opt.mode == "throughput" || opt.mode == "both") { + syncAwait(run_throughput(opt, payload)); + } + if (opt.profile) coro_io::urma_benchmark_profile::print(std::cout); + return 0; +} + +int main(int argc, char** argv) { + try { + std::cout.setf(std::ios::unitbuf); + std::cerr.setf(std::ios::unitbuf); + auto opt = parse_options(argc, argv); + coro_io::urma_benchmark_profile::configure(false, + opt.profile_sample_rate); + easylog::logger<>::instance().set_min_severity(opt.log_level); + easylog::logger<>::instance().set_async(false); + if (opt.role == "server") return run_server(opt); + return run_client(opt); + } catch (const std::exception& e) { + std::cerr << "error: " << e.what() << "\n"; + return 1; + } +} diff --git a/src/coro_rpc/examples/urma_example/CMakeLists.txt b/src/coro_rpc/examples/urma_example/CMakeLists.txt new file mode 100644 index 000000000..6a83f30ce --- /dev/null +++ b/src/coro_rpc/examples/urma_example/CMakeLists.txt @@ -0,0 +1,4 @@ +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/output/examples/coro_rpc) +add_executable(coro_rpc_urma_example + urma_example.cpp + ) diff --git a/src/coro_rpc/examples/urma_example/urma_example.cpp b/src/coro_rpc/examples/urma_example/urma_example.cpp new file mode 100644 index 000000000..eaab6d655 --- /dev/null +++ b/src/coro_rpc/examples/urma_example/urma_example.cpp @@ -0,0 +1,206 @@ +/* + * Copyright (c) 2025, Alibaba Group Holding Limited; + * + * Licensed 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 "async_simple/coro/SyncAwait.h" +#include "ylt/coro_io/client_pool.hpp" +#include "ylt/coro_io/io_context_pool.hpp" +#include "ylt/coro_io/urma/urma_device.hpp" +#include "ylt/coro_io/urma/urma_socket.hpp" +#include "ylt/coro_rpc/coro_rpc_client.hpp" +#include "ylt/coro_rpc/coro_rpc_server.hpp" +#include "ylt/coro_rpc/impl/coro_rpc_client.hpp" +#include "ylt/easylog.hpp" +#include "ylt/easylog/record.hpp" + +using namespace coro_rpc; +using namespace async_simple::coro; +using namespace std::chrono_literals; + +void set_process_env(const char* name, const std::string& value) { +#ifdef _WIN32 + _putenv_s(name, value.c_str()); +#else + setenv(name, value.c_str(), 1); +#endif +} + +void configure_urma_rpc_auto_env() { + set_process_env("URMA_RPC_ENABLE", "1"); + set_process_env("URMA_RPC_DEVICE", "bonding_dev_0"); + set_process_env("URMA_RPC_EID_INDEX", "0"); + set_process_env("URMA_RPC_RECV_BUFFER_CNT", "64"); + set_process_env("URMA_RPC_SEND_BUFFER_CNT", "64"); + set_process_env("URMA_RPC_BUFFER_SIZE", std::to_string(4 * 1024)); + set_process_env("URMA_RPC_MAX_MEMORY_USAGE", std::to_string(20 * 1024 * 1024)); + set_process_env("URMA_RPC_TP_TYPE", "ctp"); +} + +std::string_view echo(std::string_view data) { return data; } + +bool check_echo_result(const coro_rpc::rpc_result& result, + std::string_view expected) { + if (!result) { + ELOG_ERROR << "echo RPC failed: code=" << result.error().val() + << ", message=" << result.error().msg; + return false; + } + if (result.value() == expected) { + ELOG_INFO << "echo ok!"; + return true; + } + + auto actual = result.value(); + auto mismatch = std::mismatch(actual.begin(), actual.end(), expected.begin(), + expected.end()); + auto mismatch_offset = + static_cast(mismatch.first - actual.begin()); + ELOG_ERROR << "echo data err: expected_size=" << expected.size() + << ", actual_size=" << actual.size() + << ", first_mismatch_offset=" << mismatch_offset; + return false; +} + +bool warmup_urma_rpc(coro_rpc_client& client) { + std::string warmup = "urma warmup"; + ELOG_INFO << "running URMA warmup RPC"; + auto result = syncAwait(client.call_for(30s, warmup)); + return check_echo_result(result, warmup); +} + +// The basic example about how to start a rpc connection over URMA. +void basic_example() { + ELOG_INFO << "basic_example: starting"; + configure_urma_rpc_auto_env(); + coro_rpc_client client; + coro_rpc_server server; + + ELOG_INFO << "set_option: registering handler"; + server.register_handler(); + ELOG_INFO << "set_option: handler registered"; + ELOG_INFO << "set_option: calling server.async_start"; + auto future = server.async_start(); + ELOG_INFO << "set_option: server.async_start returned"; + if (future.hasResult()) { + ELOG_ERROR << future.result().value().message(); + return; + } + + // Client and server keep default TCP config; URMA_RPC_* enables the URMA upgrade. + ELOG_INFO << "set_option: server address=" << server.address() << " port=" << server.port(); + auto ec = syncAwait(client.connect(std::string{server.address()} + ":" + + std::to_string(server.port()))); + if (ec) { + ELOG_ERROR << ec.message(); + return; + } + + if (!warmup_urma_rpc(client)) { + server.stop(); + return; + } + + std::string data(1024 * 1024 * 10, 'A'); + auto result = syncAwait(client.call_for(120s, data)); + check_echo_result(result, data); + server.stop(); + return; +} + +// This example is about how to configure the detail URMA option. +void set_option() { + ELOG_INFO << "set_option: starting"; + /* init global device, should call before any other call */ + ELOG_INFO << "set_option: initializing global urma device"; + coro_io::get_global_urma_device(coro_io::urma_init_config_t{ + .dev_name = "bonding_dev_0", /*URMA device name, default is empty, which means choice + the first URMA device*/ + .buffer_pool_config = + { + .buffer_size = 4 * 1024, /*CTP send packet size*/ + .max_memory_usage = 20 * 1024 * 1024, /*max memory usage*/ + .idle_timeout = 5s, + }, + .eid_index = 0 /*EID index to use*/}); + + ELOG_INFO << "set_option: creating client"; + coro_rpc_client client; + coro_rpc_client::config conf; + ELOG_INFO << "set_option: client created, configuring urma"; + auto urma_config = coro_io::urma_socket_t::config_t{ + .recv_buffer_cnt = 64, // buffer cnt of recv queue + .send_buffer_cnt = 64, // buffer cnt of send queue + .buffer_size = 4 * 1024, // CTP max send packet size on bonding_dev_0 + .device_name = "bonding_dev_0", // empty means auto-select + .eid_index = 0 // EID index + }; + ELOG_INFO << "set_option: calling client.init_urma"; + if (!client.init_urma(urma_config)) { + ELOG_ERROR << "URMA client init failed"; + return; + } + ELOG_INFO << "set_option: client.init_urma succeeded"; + + ELOG_INFO << "set_option: creating server"; + coro_rpc_server server; + ELOG_INFO << "set_option: calling server.init_urma"; + server.init_urma(urma_config); + ELOG_INFO << "set_option: server.init_urma done"; + + ELOG_INFO << "set_option: registering handler"; + server.register_handler(); + ELOG_INFO << "set_option: handler registered"; + ELOG_INFO << "set_option: calling server.async_start"; + auto future = server.async_start(); + ELOG_INFO << "set_option: server.async_start returned"; + if (future.hasResult()) { + ELOG_ERROR << future.result().value().message(); + return; + } + + ELOG_INFO << "set_option: server address=" << server.address() << " port=" << server.port(); + auto ec = syncAwait(client.connect(std::string{server.address()} + ":" + + std::to_string(server.port()))); + if (ec) { + ELOG_ERROR << ec.message(); + return; + } + + if (!warmup_urma_rpc(client)) { + server.stop(); + return; + } + + std::string data(1024 * 1024 * 10, 'A'); + auto result = syncAwait(client.call_for(120s, data)); + check_echo_result(result, data); + server.stop(); +} + +int main() { + easylog::logger<>::instance().set_min_severity(easylog::Severity::DEBUG); + easylog::logger<>::instance().set_async(false); + ELOG_INFO << "URMA example main started"; + set_option(); + ELOG_INFO << "set_option completed, now basic_example"; + basic_example(); + ELOG_INFO << "basic_example completed"; + return 0; +} diff --git a/src/coro_rpc/tests/CMakeLists.txt b/src/coro_rpc/tests/CMakeLists.txt index 8e6593a42..13b6db427 100644 --- a/src/coro_rpc/tests/CMakeLists.txt +++ b/src/coro_rpc/tests/CMakeLists.txt @@ -12,6 +12,7 @@ set(TEST_SRCS test_parallel.cpp test_client_filter.cpp test_abi_compatible.cpp + test_urma_rpc_env.cpp ) if(YLT_ENABLE_ND) list(APPEND TEST_SRCS test_networkdirect_rpc.cpp) diff --git a/src/coro_rpc/tests/test_urma_rpc_env.cpp b/src/coro_rpc/tests/test_urma_rpc_env.cpp new file mode 100644 index 000000000..8f05a3b7c --- /dev/null +++ b/src/coro_rpc/tests/test_urma_rpc_env.cpp @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2026, Alibaba Group Holding Limited; + * + * Licensed 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 "doctest.h" + +#include +#include +#include +#include + +#include + +#ifdef _WIN32 +inline void set_test_env(const char* name, const char* value) { + _putenv_s(name, value); +} +inline void unset_test_env(const char* name) { _putenv_s(name, ""); } +#else +inline void set_test_env(const char* name, const char* value) { + setenv(name, value, 1); +} +inline void unset_test_env(const char* name) { unsetenv(name); } +#endif + +class scoped_env_var { + public: + scoped_env_var(const char* name, const char* value) : name_(name) { + const char* old = std::getenv(name); + if (old) { + old_value_ = old; + had_value_ = true; + } + set_test_env(name, value); + } + + ~scoped_env_var() { + if (had_value_) { + set_test_env(name_.c_str(), old_value_.c_str()); + } + else { + unset_test_env(name_.c_str()); + } + } + + private: + std::string name_; + std::string old_value_; + bool had_value_ = false; +}; + +#ifdef YLT_ENABLE_URMA +#include + +TEST_CASE("urma rpc env enable parsing") { + { + scoped_env_var env("URMA_RPC_ENABLE", "1"); + CHECK(coro_io::detail::urma_rpc_env_enabled()); + } + { + scoped_env_var env("URMA_RPC_ENABLE", "ON"); + CHECK(coro_io::detail::urma_rpc_env_enabled()); + } + { + scoped_env_var env("URMA_RPC_ENABLE", "true"); + CHECK(coro_io::detail::urma_rpc_env_enabled()); + } + { + scoped_env_var env("URMA_RPC_ENABLE", "yes"); + CHECK(coro_io::detail::urma_rpc_env_enabled()); + } + { + scoped_env_var env("URMA_RPC_ENABLE", "0"); + CHECK_FALSE(coro_io::detail::urma_rpc_env_enabled()); + } + { + scoped_env_var env("URMA_RPC_ENABLE", "abc"); + CHECK_FALSE(coro_io::detail::urma_rpc_env_enabled()); + } +} + +TEST_CASE("urma rpc env config parsing uses defaults for invalid values") { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + scoped_env_var device("URMA_RPC_DEVICE", "test_dev"); + scoped_env_var eid("URMA_RPC_EID_INDEX", "7"); + scoped_env_var cq("URMA_RPC_CQ_SIZE", "abc"); + scoped_env_var recv("URMA_RPC_RECV_BUFFER_CNT", "9"); + scoped_env_var send("URMA_RPC_SEND_BUFFER_CNT", "10"); + scoped_env_var buffer("URMA_RPC_BUFFER_SIZE", "8192"); + scoped_env_var memory("URMA_RPC_MAX_MEMORY_USAGE", "16777216"); + scoped_env_var tp("URMA_RPC_TP_TYPE", "rtp"); + + auto config = coro_io::detail::make_urma_rpc_config_from_env(); + CHECK(config.device_name == "test_dev"); + CHECK(config.eid_index == 7); + CHECK(config.cq_size == coro_io::urma_socket_t::config_t{}.cq_size); + CHECK(config.recv_buffer_cnt == 9); + CHECK(config.send_buffer_cnt == 10); + CHECK(config.buffer_size == 8192); + CHECK(config.max_memory_usage == 16777216); + CHECK(config.tp_type == URMA_RTP); +} + +TEST_CASE("urma rpc env event mode defaults to on and can be overridden") { + { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + auto config = coro_io::detail::make_urma_rpc_config_from_env(); + CHECK(config.event_mode == true); + CHECK(config.busy_poll_budget == + coro_io::urma_socket_t::config_t{}.busy_poll_budget); + } + { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + scoped_env_var mode("URMA_RPC_EVENT_MODE", "0"); + auto config = coro_io::detail::make_urma_rpc_config_from_env(); + CHECK(config.event_mode == false); + } + { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + scoped_env_var mode("URMA_RPC_EVENT_MODE", "off"); + auto config = coro_io::detail::make_urma_rpc_config_from_env(); + CHECK(config.event_mode == false); + } + { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + scoped_env_var budget("URMA_RPC_BUSY_POLL_BUDGET", "32"); + auto config = coro_io::detail::make_urma_rpc_config_from_env(); + CHECK(config.busy_poll_budget == 32); + } +} + +TEST_CASE("urma rpc env disabled keeps default client tcp config") { + scoped_env_var enable("URMA_RPC_ENABLE", "0"); + + coro_rpc::coro_rpc_client client; + CHECK(std::holds_alternative( + client.get_config().socket_config)); +} + +TEST_CASE("urma rpc enabled without usable device keeps client constructible") { + scoped_env_var enable("URMA_RPC_ENABLE", "1"); + scoped_env_var device("URMA_RPC_DEVICE", "device_that_should_not_exist_for_test"); + + coro_rpc::coro_rpc_client client; + CHECK(std::holds_alternative( + client.get_config().socket_config)); +} + +#else +TEST_CASE("urma rpc env tests compile without urma support") { + coro_rpc::coro_rpc_client client; + CHECK(std::holds_alternative( + client.get_config().socket_config)); +} +#endif