dwh110 commented on code in PR #3428: URL: https://github.com/apache/brpc/pull/3428#discussion_r3930753680
########## src/brpc/urma/urma_endpoint.cpp: ########## @@ -0,0 +1,1846 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "brpc/urma/urma_endpoint.h" + +#if BRPC_WITH_URMA + +#include <sys/resource.h> +#include <unistd.h> + +#include <algorithm> +#include <cstdio> +#include <cstring> +#include <memory> +#include <mutex> +#include <string> +#include <unordered_set> +#include <utility> +#include <vector> + +#include <gflags/gflags.h> + +#include "butil/atomicops.h" +#include "butil/iobuf.h" +#include "butil/logging.h" +#include "butil/macros.h" +#include "butil/sys_byteorder.h" +#include "butil/time.h" +#include "bthread/bthread.h" +#include "bthread/butex.h" + +#include "urma_api.h" + +#include "brpc/input_messenger.h" +#include "brpc/socket.h" +#include "brpc/urma/urma_bonding.h" +#include "brpc/urma/urma_handshake.h" +#include "brpc/urma/urma_handshake.pb.h" +#include "brpc/urma/urma_helper.h" +#include "brpc/urma_transport.h" + +DECLARE_int32(task_group_ntags); + +namespace brpc { +namespace urma { + +// Flags used here are declared in urma_endpoint.h (urma_use_polling, +// urma_poller_num, urma_disable_bthread). Declare the rest here. +DECLARE_int32(urma_sq_size); +DECLARE_int32(urma_rq_size); +DECLARE_int32(urma_cqe_poll_once); +DECLARE_bool(urma_recv_zerocopy); +DECLARE_int32(urma_zerocopy_min_size); +DECLARE_int32(urma_prepared_jetty_cnt); +DECLARE_bool(urma_poller_yield); + +// ---- Constants shared with the handshake module ---- +static const int WAIT_TIMEOUT_MS = 50; +static const size_t HELLO_ACK_LEN = 4; +static const uint32_t HELLO_ACK_URMA_OK = 0x1; +static const size_t IOBUF_BLOCK_HEADER_LEN = sizeof(butil::IOBuf::Block); + +// ---- Globals: prepared jetty pool + poller groups ---- +struct PreparedJetty { + UrmaResource* res; +}; +static butil::Mutex g_prepared_mutex; +static UrmaResource* g_prepared_list = nullptr; // singly-linked +static int g_prepared_cnt = 0; + +static int PreparedJettyCount() { + const int requested = + std::max(0, std::min(FLAGS_urma_prepared_jetty_cnt, 1024)); + if (requested == 0) { + return 0; + } + + struct rlimit nofile; + if (getrlimit(RLIMIT_NOFILE, &nofile) != 0 || + nofile.rlim_cur == RLIM_INFINITY) { + return requested; + } + + // In event mode each prepared JFCE consumes a file descriptor. Keep room + // for one TCP fd per future URMA connection and for brpc/system internals. + static const rlim_t kReservedFdCount = 64; + const rlim_t max_prepared = + nofile.rlim_cur > kReservedFdCount + ? (nofile.rlim_cur - kReservedFdCount) / 2 + : 0; + if (max_prepared >= static_cast<rlim_t>(requested)) { + return requested; + } + + LOG(WARNING) << "Cap URMA prepared jetty count from " << requested + << " to " << max_prepared + << " due to RLIMIT_NOFILE=" << nofile.rlim_cur; + return static_cast<int>(max_prepared); +} + +std::vector<UrmaEndpoint::PollerGroup> UrmaEndpoint::_poller_groups; + +// ============================================================================ +// UrmaResource lifecycle. +// ============================================================================ + +UrmaResource::~UrmaResource() { + if (remote_jetty) { + urma_unimport_jetty(remote_jetty); + } + if (remote_seg) { + urma_unimport_seg(remote_seg); + } + if (jetty) { + urma_delete_jetty(jetty); + } + if (jfr) { + urma_delete_jfr(jfr); + } + if (jfc) { + urma_delete_jfc(jfc); + } + if (jfce) { + urma_delete_jfce(jfce); + } +} + +// ============================================================================ +// Constructor / destructor / Reset. +// ============================================================================ + +UrmaEndpoint::UrmaEndpoint(Socket* s) + : _socket(s), + _state(UNINIT), + _handshake_version(0), + _resource(nullptr) { + _sq_size = static_cast<uint16_t>( + std::max(16, std::min(4096, static_cast<int>(FLAGS_urma_sq_size)))); + _rq_size = static_cast<uint16_t>( + std::max(16, std::min(4096, static_cast<int>(FLAGS_urma_rq_size)))); + _read_butex = bthread::butex_create_checked<butil::atomic<int>>(); + _read_butex->store(0, butil::memory_order_relaxed); +} + +UrmaEndpoint::~UrmaEndpoint() { + DeallocateResources(); + if (_read_butex) { + bthread::butex_destroy(_read_butex); + _read_butex = nullptr; + } +} + +void UrmaEndpoint::Reset() { + DeallocateResources(); + _state = UNINIT; + _handshake_version = 0; + _remote_recv_block_size = 0; + _local_window_capacity = 0; + _remote_window_capacity = 0; + _remote_rq_window_size.store(0, butil::memory_order_relaxed); + _sq_window_size.store(0, butil::memory_order_relaxed); + _new_rq_wrs.store(0, butil::memory_order_relaxed); + _sq_imm_window_size = 0; + _sq_current = 0; + _sq_sent = 0; + _rq_received = 0; + _pending_received_bytes.store(0, butil::memory_order_relaxed); + _sbuf.clear(); + _rbuf.clear(); + _rbuf_data.clear(); + _read_butex->store(0, butil::memory_order_relaxed); +} + +// ============================================================================ +// Handshake IO helpers (ReadFromFd / WriteToFd / PushBackToReadBuf). +// Modeled on RdmaEndpoint::ReadFromFdLoop / WriteToFdLoop. +// ============================================================================ + +int UrmaEndpoint::ReadFromFd(void* data, size_t len) { + char* p = static_cast<char*>(data); + size_t received = 0; + while (received < len) { + const int expected_val = _read_butex->load(butil::memory_order_acquire); + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nr = read(fd, p + received, len - received); + if (nr < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) { + int rc = bthread::butex_wait(_read_butex, expected_val, &duetime); + if (rc < 0 && errno != EWOULDBLOCK && errno != ETIMEDOUT) { + return -1; + } + continue; + } + return -1; + } + if (nr == 0) { + errno = EEOF; + return -1; + } + received += nr; + } + return 0; +} + +void UrmaEndpoint::PushBackToReadBuf(const void* data, size_t len) { + _socket->_read_buf.append(data, len); +} + +int UrmaEndpoint::WriteToFd(void* data, size_t len) { + char* p = static_cast<char*>(data); + size_t written = 0; + while (written < len) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = write(fd, p + written, len - written); + if (nw >= 0) { + written += nw; + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +// ============================================================================ +// Hello builders / parsers. +// ============================================================================ + +void UrmaEndpoint::MakeLocalParsedHello(ParsedHello* out) const { + *out = ParsedHello{}; // value-initialize (avoids memset on non-trivial type) + out->buffer_size = static_cast<uint32_t>(GetUrmaRecvBlockSize()); + out->recv_buffer_cnt = _rq_size - 1; + if (_resource && _resource->jetty) { + out->jetty_id = _resource->jetty->jetty_id.id; + out->uasid = _resource->jetty->jetty_id.uasid; + const urma_eid_t* local_eid = GetUrmaLocalEid(); + const uint8_t* advertised_eid = + local_eid != nullptr + ? local_eid->raw + : _resource->jetty->jetty_id.eid.raw; + std::memcpy(out->eid, advertised_eid, 16); + } + out->tp_type = static_cast<uint8_t>(URMA_CTP); + // Pool segment: flatten g_pool_seg's seg fields. + urma_target_seg_t* pool = GetPoolSegFor(nullptr); + if (pool) { + std::memcpy(out->seg_eid, pool->seg.ubva.eid.raw, 16); + out->seg_uasid = pool->seg.ubva.uasid; + out->seg_va = pool->seg.ubva.va; + out->seg_len = pool->seg.len; + out->seg_token_id = pool->seg.token_id; + } +} + +void UrmaEndpoint::FillLocalHelloV2(v2_wire::HelloMessage* out) const { + *out = v2_wire::HelloMessage{}; // value-initialize + out->msg_len = v2_wire::HELLO_PACKET_LEN; + out->hello_ver = v2_wire::HELLO_V2_VERSION; + out->impl_ver = v2_wire::IMPL_V2_VERSION; + ParsedHello p; + MakeLocalParsedHello(&p); + out->buffer_size = p.buffer_size; + out->recv_buffer_cnt = p.recv_buffer_cnt; + out->jetty_id = p.jetty_id; + std::memcpy(out->eid, p.eid, 16); + out->uasid = p.uasid; + out->tp_type = p.tp_type; + std::memcpy(out->seg_eid, p.seg_eid, 16); + out->seg_uasid = p.seg_uasid; + out->seg_va = p.seg_va; + out->seg_len = p.seg_len; + out->seg_token_id = p.seg_token_id; +} + +void UrmaEndpoint::FillLocalHelloV3(UrmaHello* out) const { + ParsedHello p; + MakeLocalParsedHello(&p); + out->set_buffer_size(p.buffer_size); + out->set_recv_buffer_cnt(p.recv_buffer_cnt); + out->set_jetty_id(p.jetty_id); + out->set_eid(p.eid, 16); + out->set_uasid(p.uasid); + out->set_tp_type(p.tp_type); + out->set_seg_eid(p.seg_eid, 16); + out->set_seg_uasid(p.seg_uasid); + out->set_seg_va(p.seg_va); + out->set_seg_len(p.seg_len); + out->set_seg_token_id(p.seg_token_id); +} + +int UrmaEndpoint::WriteHelloV3(const UrmaHello& msg) { + butil::IOBuf packet; + packet.append("URM3", 4); + std::string body; + if (!msg.SerializeToString(&body)) { + LOG(ERROR) << "Fail to serialize UrmaHello"; + return -1; + } + uint32_t pb_size_be = butil::HostToNet32(static_cast<uint32_t>(body.size())); + packet.append(&pb_size_be, sizeof(pb_size_be)); + packet.append(body); + return WriteToFd(packet); +} + +int UrmaEndpoint::WriteToFd(butil::IOBuf& data) { + // Write out the IOBuf in a single WriteToFd-style loop. + while (!data.empty()) { + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const int fd = _socket->fd(); + const ssize_t nw = data.cut_into_file_descriptor(fd); + if (nw >= 0) { + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return -1; + } + if (_socket->WaitEpollOut(fd, true, &duetime) != 0 && errno != ETIMEDOUT) { + return -1; + } + } + return 0; +} + +int UrmaEndpoint::ReadAndParseHelloV3(ParsedHello* out, bool* negotiated) { + *negotiated = false; + uint32_t pb_size_be = 0; + if (ReadFromFd(&pb_size_be, sizeof(pb_size_be)) < 0) { + return -1; + } + const uint32_t pb_size = butil::NetToHost32(pb_size_be); + if (pb_size == 0 || pb_size > 4096) { + return 0; + } + std::string body(pb_size, '\0'); + if (ReadFromFd(&body[0], pb_size) < 0) { + return -1; + } + UrmaHello msg; + if (!msg.ParseFromArray(body.data(), static_cast<int>(body.size()))) { + return 0; + } + if (msg.eid().size() != 16 || msg.seg_eid().size() != 16) { + return 0; + } + out->buffer_size = msg.buffer_size(); + out->recv_buffer_cnt = msg.recv_buffer_cnt(); + out->jetty_id = msg.jetty_id(); + std::memcpy(out->eid, msg.eid().data(), 16); + out->uasid = msg.uasid(); + out->tp_type = static_cast<uint8_t>(msg.tp_type()); + std::memcpy(out->seg_eid, msg.seg_eid().data(), 16); + out->seg_uasid = msg.seg_uasid(); + out->seg_va = msg.seg_va(); + out->seg_len = msg.seg_len(); + out->seg_token_id = msg.seg_token_id(); + if (!ValidHello(*out)) { + return 0; + } + *negotiated = true; + return 0; +} + +// ============================================================================ +// Allocate / deallocate per-connection resources. +// ============================================================================ + +int UrmaEndpoint::AllocateResources() { + if (_resource) { + return 0; + } + urma_context_t* ctx = GetUrmaContext(); + if (!ctx) { + errno = ENODEV; + return -1; + } + + _resource = new (std::nothrow) UrmaResource(); + if (!_resource) { + return -1; + } + + // Try the prepared pool first (sized sq/rq match). + if (_sq_size <= static_cast<uint16_t>(FLAGS_urma_sq_size) && + _rq_size <= static_cast<uint16_t>(FLAGS_urma_rq_size)) { + BAIDU_SCOPED_LOCK(g_prepared_mutex); + if (g_prepared_list) { + UrmaResource* next = g_prepared_list->next; + delete _resource; + _resource = g_prepared_list; + g_prepared_list = next; + _resource->next = nullptr; + --g_prepared_cnt; + } + } + + if (!_resource->jfc) { + // The SDK requires every JFC to reference a JFCE. Polling mode does + // not arm or consume it, but still supplies the required object. + _resource->jfce = urma_create_jfce(ctx); + if (!_resource->jfce || + (!FLAGS_urma_use_polling && _resource->jfce->fd < 0)) { + LOG(ERROR) << "Fail to create a usable URMA JFCE"; + errno = ENODEV; + return -1; + } + + urma_jfc_cfg_t jfc_cfg{}; + jfc_cfg.depth = static_cast<uint32_t>(_sq_size + _rq_size); + jfc_cfg.jfce = _resource->jfce; + _resource->jfc = urma_create_jfc(ctx, &jfc_cfg); + if (!_resource->jfc) { + PLOG(ERROR) << "urma_create_jfc"; + return -1; + } + + urma_jfr_cfg_t jfr_cfg{}; + jfr_cfg.depth = static_cast<uint32_t>(_rq_size); + jfr_cfg.trans_mode = URMA_TM_RM; + jfr_cfg.max_sge = 1; + jfr_cfg.min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + jfr_cfg.jfc = _resource->jfc; + _resource->jfr = urma_create_jfr(ctx, &jfr_cfg); + if (!_resource->jfr) { + PLOG(ERROR) << "urma_create_jfr"; + return -1; + } + + urma_jetty_cfg_t jetty_cfg{}; + jetty_cfg.flag.bs.share_jfr = 1; + jetty_cfg.jfs_cfg.depth = static_cast<uint32_t>(_sq_size); + jetty_cfg.jfs_cfg.trans_mode = URMA_TM_RM; + jetty_cfg.jfs_cfg.priority = GetUrmaJettyPriority(); + jetty_cfg.jfs_cfg.max_sge = + static_cast<uint8_t>(GetUrmaMaxSge()); + jetty_cfg.jfs_cfg.rnr_retry = URMA_TYPICAL_RNR_RETRY; + jetty_cfg.jfs_cfg.err_timeout = URMA_TYPICAL_ERR_TIMEOUT; + jetty_cfg.jfs_cfg.jfc = _resource->jfc; + jetty_cfg.shared.jfr = _resource->jfr; + jetty_cfg.shared.jfc = _resource->jfc; + _resource->jetty = urma_create_jetty(ctx, &jetty_cfg); + if (!_resource->jetty) { + PLOG(ERROR) << "urma_create_jetty"; + return -1; + } + } + + _sbuf.resize(_sq_size - RESERVED_WR_NUM); + _rbuf.resize(_rq_size); + _rbuf_data.resize(_rq_size, nullptr); + + // Wrap the JFCE fd in a brpc Socket so PollCq is driven by epoll. + if (!FLAGS_urma_use_polling) { + if (!_resource->jfce || _resource->jfce->fd < 0) { + LOG(ERROR) << "Prepared URMA resource has no usable JFCE"; + errno = ENODEV; + return -1; + } + if (ReqNotifyCq() != 0) { + return -1; + } + SocketOptions options; + options.user = this; + options.keytable_pool = _socket->keytable_pool(); + options.fd = _resource->jfce->fd; + options.on_edge_triggered_events = PollCq; + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(ERROR) << "Fail to create CQ socket"; + return -1; + } + } else { + // Polling mode: synthetic carrier socket (no fd). + SocketOptions options; + options.user = this; + options.keytable_pool = _socket->keytable_pool(); + options.on_edge_triggered_events = PollCq; + if (Socket::Create(options, &_cq_sid) < 0) { + PLOG(ERROR) << "Fail to create CQ socket (polling)"; + return -1; + } + PollerAddCqSid(); + } + return 0; +} + +void UrmaEndpoint::DeallocateResources() { + if (!_resource) { + return; + } + + if (FLAGS_urma_use_polling) { + PollerRemoveCqSid(); + } + + // Tear down the CQ socket so the EventDispatcher stops calling PollCq. + if (_cq_sid != INVALID_SOCKET_ID) { + SocketUniquePtr s; + if (Socket::Address(_cq_sid, &s) == 0) { + if (s->fd() >= 0) { + s->_io_event.RemoveConsumer(s->_fd); + } + s->_user = nullptr; // Do not release user (this UrmaEndpoint). + s->_fd = -1; // Already removed fd from epoll. + s->SetFailed(); + } + _cq_sid = INVALID_SOCKET_ID; + } + + // Reusing a Jetty requires a driver-supported RESET plus a complete JFC + // drain. Until that lifecycle is implemented, prepared resources are + // one-shot: they accelerate connection setup but are destroyed on close. + delete _resource; + _resource = nullptr; +} + +// ============================================================================ +// ImportPeer: the critical import_seg-before-import_jetty sequence. +// ============================================================================ + +int UrmaEndpoint::ImportPeer(const ParsedHello& peer) { + urma_context_t* ctx = GetUrmaContext(); + if (!ctx) { + errno = ENODEV; + return -1; + } + + // 1. urma_import_seg FIRST so the kernel establishes TP routing for the + // remote EID. Without this the first SEND is rejected by hardware with + // URMA_CR_RNR_RETRY_CNT_EXC_ERR. + urma_seg_t peer_seg{}; + std::memcpy(peer_seg.ubva.eid.raw, peer.seg_eid, 16); + peer_seg.ubva.uasid = peer.seg_uasid; + peer_seg.ubva.va = peer.seg_va; + peer_seg.len = peer.seg_len; + peer_seg.token_id = peer.seg_token_id; + urma_token_t seg_token{}; + urma_import_seg_flag_t seg_flag{}; + seg_flag.bs.cacheable = URMA_NON_CACHEABLE; + seg_flag.bs.access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC; + seg_flag.bs.mapping = URMA_SEG_NOMAP; + _resource->remote_seg = urma_import_seg(ctx, &peer_seg, &seg_token, 0, seg_flag); + if (!_resource->remote_seg) { + PLOG(ERROR) << "urma_import_seg failed"; + return -1; + } + + // 2. urma_import_jetty. + urma_rjetty_t remote{}; + std::memcpy(remote.jetty_id.eid.raw, peer.eid, 16); + remote.jetty_id.uasid = peer.uasid; + remote.jetty_id.id = peer.jetty_id; + remote.trans_mode = URMA_TM_RM; + remote.type = URMA_JETTY; + if (peer.tp_type > static_cast<uint8_t>(URMA_UTP)) { + errno = EPROTO; + return -1; + } + remote.tp_type = static_cast<urma_tp_type_t>(peer.tp_type); + + urma_token_t token{}; + const bool use_bonding_extension = + IsUrmaBondingDevice() && remote.trans_mode == URMA_TM_RM; + errno = 0; + if (use_bonding_extension) { +#if BRPC_URMA_HAS_BONDING_EXT + // The bonding provider needs the local jetty to associate its send + // path with the imported target. A plain import may return success + // without setting that association, leaving traffic one-way only. + bondp_rjetty_t bonding_remote{}; + bonding_remote.base = remote; + bonding_remote.base.flag.bs.has_drv_ext = 1; + bonding_remote.jetty = _resource->jetty; + _resource->remote_jetty = + urma_import_jetty(ctx, &bonding_remote.base, &token); +#else + LOG(ERROR) << "Bonding remote jetty import requires provider header " + "urma_ubagg.h"; + errno = ENOTSUP; +#endif + } else { + _resource->remote_jetty = urma_import_jetty(ctx, &remote, &token); + } + if (!_resource->remote_jetty) { + if (errno == 0) { + errno = EIO; + } + char remote_eid[URMA_EID_STR_LEN + 1] = {}; + std::snprintf(remote_eid, sizeof(remote_eid), EID_FMT, + EID_RAW_ARGS(peer.eid)); + PLOG(ERROR) << "urma_import_jetty failed" + << " remote_eid=" << remote_eid + << " remote_uasid=" << peer.uasid + << " remote_jetty_id=" << peer.jetty_id + << " trans_mode=" << remote.trans_mode + << " tp_type=" << remote.tp_type + << " bonding_extension=" << use_bonding_extension; + return -1; + } + return 0; +} + +// ============================================================================ +// Send / recv data path. +// ============================================================================ + +// Private IOBuf accessor mirroring RdmaIOBuf: reach into IOBuf block refs to +// build a urma_sge_t directly, without memcpy. +class UrmaIOBuf : private butil::IOBuf { + friend class ::brpc::urma::UrmaEndpoint; +public: + using butil::IOBuf::_ref_num; + using butil::IOBuf::_ref_at; + using butil::IOBuf::fetch1; + using butil::IOBuf::get_first_data_meta; + using butil::IOBuf::cutn; + // Build the SGE for the current head block. + // Returns bytes added, or -1 (errno set). + ssize_t cut_into_sglist(urma_sge_t* sglist, size_t* sge_index, + butil::IOBuf* to, size_t max_sge, + size_t max_len) { + size_t len = 0; + while (*sge_index < max_sge && len < max_len && _ref_num() != 0) { + butil::IOBuf::BlockRef const& r = _ref_at(0); + const void* start = fetch1(); + urma_target_seg_t* tseg = + GetPoolSegFor(const_cast<void*>(start)); + if (!tseg) { + // User-registered memory: look up the seg handle. + uint64_t meta = get_first_data_meta(); + if (meta != 0) { + tseg = reinterpret_cast<urma_target_seg_t*>( + static_cast<uintptr_t>(meta)); + } + } + if (!tseg) { + errno = ERDMAMEM; + return -1; + } + size_t this_len = r.length; + if (len + this_len > max_len) { + this_len = max_len - len; + } + sglist[*sge_index].addr = reinterpret_cast<uint64_t>(start); + sglist[*sge_index].len = static_cast<uint32_t>(this_len); + sglist[*sge_index].tseg = tseg; + cutn(to, this_len); + len += this_len; + (*sge_index)++; + } + return static_cast<ssize_t>(len); + } +}; + +ssize_t UrmaEndpoint::CutFromIOBufList(butil::IOBuf** from, size_t ndata) { + if (!_resource || !_resource->jetty || !_resource->remote_jetty) { + errno = ENOTCONN; + return -1; + } + int max_sge = GetUrmaMaxSge(); + if (max_sge < 1) { + max_sge = 1; + } + + urma_sge_t* sglist = static_cast<urma_sge_t*>( + alloca(sizeof(urma_sge_t) * max_sge)); + if (!sglist) { + errno = ENOMEM; + return -1; + } + + size_t current = 0; + ssize_t total_len = 0; + while (current < ndata) { + uint16_t remote_wnd = _remote_rq_window_size.load(butil::memory_order_relaxed); + uint16_t sq_wnd = _sq_window_size.load(butil::memory_order_relaxed); + if (remote_wnd == 0 || sq_wnd == 0) { + if (total_len > 0) { + break; + } + errno = EAGAIN; + return -1; + } + butil::IOBuf* to = &_sbuf[_sq_current]; + size_t sge_index = 0; + size_t this_len = 0; + size_t max_len = _remote_recv_block_size > 0 + ? _remote_recv_block_size + : GetUrmaRecvBlockSize(); + while (sge_index < static_cast<size_t>(max_sge) && + this_len < max_len && current < ndata) { + auto* data = reinterpret_cast<UrmaIOBuf*>(from[current]); + if (data->empty()) { + ++current; + continue; + } + ssize_t n = data->cut_into_sglist(sglist, &sge_index, to, + max_sge, max_len - this_len); + if (n < 0) { + return -1; + } + this_len += n; + } + if (sge_index == 0) { + break; + } + + urma_sg_t sg{sglist, static_cast<uint32_t>(sge_index)}; + urma_jfs_wr_t wr{}; + std::memset(&wr, 0, sizeof(wr)); + // Send payload with URMA_OPC_SEND. Receive credits are flushed + // separately by SendImm() after SendAck() reaches its threshold. + // Piggybacking credits turns every payload into SEND_IMM and can + // produce asymmetric completions with the bonding provider. + wr.opcode = URMA_OPC_SEND; + wr.flag.bs.complete_enable = 1; + wr.tjetty = _resource->remote_jetty; + wr.send.src = sg; + wr.user_ctx = 1; + urma_jfs_wr_t* bad_wr = nullptr; + const uint16_t sq_slot = _sq_current; + const uint32_t local_jetty_id = _resource->jetty->jetty_id.id; + const uint32_t remote_jetty_id = _resource->remote_jetty->id.id; + + // Reserve both credits before making the WR visible to the provider. + // In polling mode a completion (and even the peer's receive-credit + // ACK) can be processed by another thread before post_send returns. + // Decrementing after post therefore creates a transient capacity + 1 + // window and makes the strict credit check tear down a healthy + // connection. + _remote_rq_window_size.fetch_sub(1, butil::memory_order_relaxed); + _sq_window_size.fetch_sub(1, butil::memory_order_relaxed); + int rc = urma_post_jetty_send_wr(_resource->jetty, &wr, &bad_wr); + if (rc != URMA_SUCCESS) { + const int provider_errno = errno; + _remote_rq_window_size.fetch_add(1, butil::memory_order_relaxed); + _sq_window_size.fetch_add(1, butil::memory_order_relaxed); + LOG(WARNING) << "urma_post_jetty_send_wr failed: " << rc + << ", provider_errno=" << provider_errno + << " (" << berror(provider_errno) << ')' + << ", bad_wr=" << static_cast<const void*>(bad_wr) + << ", bad_is_current=" << (bad_wr == &wr) + << ", sq_slot=" << sq_slot + << ", local_jetty_id=" << local_jetty_id + << ", remote_jetty_id=" << remote_jetty_id + << ", state=" << GetStateStr() + << ", sq_window=" << sq_wnd + << ", remote_rq_window=" << remote_wnd + << ", num_sge=" << sge_index + << ", configured_max_sge=" << GetUrmaMaxSge() + << ", payload_size=" << this_len + << " on " << _socket->description(); + errno = rc; + return -1; + } + _sq_current = (_sq_current + 1) % (_sq_size - RESERVED_WR_NUM); + total_len += static_cast<ssize_t>(this_len); + } + return total_len; +} + +bool UrmaEndpoint::IsWritable() const { + return _remote_rq_window_size.load(butil::memory_order_relaxed) > 0 && + _sq_window_size.load(butil::memory_order_relaxed) > 0; +} + +// ============================================================================ +// Recv path. +// ============================================================================ + +int UrmaEndpoint::DoPostRecv(void* block, size_t block_size) { + urma_target_seg_t* tseg = GetPoolSegFor(block); + if (!tseg) { + errno = ERDMAMEM; + return -1; + } + urma_sge_t sge{reinterpret_cast<uint64_t>(block), + static_cast<uint32_t>(block_size), tseg, nullptr}; + urma_sg_t sg{&sge, 1}; + urma_jfr_wr_t wr{sg, 0, nullptr}; + urma_jfr_wr_t* bad = nullptr; + // Use the shared-JFR path on every device, including bonding. The bonding + // provider owns physical receive scheduling for the JFR; a local + // jetty-to-target association is not part of the RM receive API. + const urma_status_t status = + urma_post_jfr_wr(_resource->jfr, &wr, &bad); + if (status != URMA_SUCCESS) { + LOG(WARNING) << "Failed to post URMA receive WR: status=" << status + << " bonding=" << IsUrmaBondingDevice() + << " bad_wr=" << static_cast<const void*>(bad) + << " bad_is_current=" << (bad == &wr) + << " local_jetty_id=" << _resource->jetty->jetty_id.id + << " provider_associated_remote=" + << static_cast<const void*>( + _resource->jetty->remote_jetty) + << " state=" << GetStateStr() + << " on " << _socket->description(); + errno = status; + return -1; + } + return 0; +} + +int UrmaEndpoint::PostRecv(uint32_t num, bool zerocopy) { + for (uint32_t i = 0; i < num; ++i) { + size_t block_size = GetUrmaRecvBlockSize(); + if (zerocopy) { + _rbuf[_rq_received].clear(); + butil::IOBufAsZeroCopyOutputStream zcis( + &_rbuf[_rq_received], block_size + IOBUF_BLOCK_HEADER_LEN); + void* data = nullptr; + int size = 0; + if (!zcis.Next(&data, &size) || !data || + size < static_cast<int>(block_size)) { + errno = ENOMEM; + return -1; + } + _rbuf_data[_rq_received] = data; + if (DoPostRecv(data, block_size) < 0) { + return -1; + } + } else { + if (_rbuf_data[_rq_received] == nullptr) { + _rbuf[_rq_received].clear(); + butil::IOBufAsZeroCopyOutputStream zcos( + &_rbuf[_rq_received], + block_size + IOBUF_BLOCK_HEADER_LEN); + void* data = nullptr; + int size = 0; + if (!zcos.Next(&data, &size) || !data || + size < static_cast<int>(block_size)) { + errno = ENOMEM; + return -1; + } + _rbuf_data[_rq_received] = data; + } + if (DoPostRecv(_rbuf_data[_rq_received], block_size) < 0) { + return -1; + } + } + _rq_received = (_rq_received + 1) % _rq_size; + } + return 0; +} + +int UrmaEndpoint::SendImm(uint32_t imm) { + if (imm == 0) { + return 0; + } + if (!_resource || !_resource->jetty || !_resource->remote_jetty) { + errno = ENOTCONN; + return -1; + } + if (_sq_imm_window_size == 0) { + errno = EAGAIN; + return -1; + } + // Empty-payload SEND_IMM flushes peer-side receive credit. Connection + // lifetime is owned by the TCP fd, so this is not an EOF marker. + urma_jfs_wr_t wr{}; + std::memset(&wr, 0, sizeof(wr)); + wr.opcode = URMA_OPC_SEND_IMM; + wr.flag.bs.complete_enable = 1; + wr.flag.bs.solicited_enable = 1; + wr.tjetty = _resource->remote_jetty; + wr.send.imm_data = imm; + wr.user_ctx = 0; // 0 == pure ack (HandleCompletion reuses budget). + urma_jfs_wr_t* bad = nullptr; + // Reserve the ACK-only SQ slot before posting for the same reason as the + // data windows in CutFromIOBufList: polling may observe its completion as + // soon as the provider accepts the WR. + --_sq_imm_window_size; + const urma_status_t status = + urma_post_jetty_send_wr(_resource->jetty, &wr, &bad); + if (status != URMA_SUCCESS) { + const int provider_errno = errno; + ++_sq_imm_window_size; + _new_rq_wrs.fetch_add(imm, butil::memory_order_relaxed); + LOG(WARNING) << "Failed to post URMA credit ACK: status=" << status + << " provider_errno=" << provider_errno + << " (" << berror(provider_errno) << ')' + << " bad_wr=" << static_cast<const void*>(bad) + << " bad_is_current=" << (bad == &wr) + << " imm=" << imm + << " local_jetty_id=" + << _resource->jetty->jetty_id.id + << " remote_jetty_id=" + << _resource->remote_jetty->id.id + << " state=" << GetStateStr() + << " on " << _socket->description(); + errno = status; + return -1; + } + return 0; +} + +int UrmaEndpoint::SendAck(int num) { + const uint16_t old = + _new_rq_wrs.fetch_add(num, butil::memory_order_relaxed); + if (old + num > _remote_window_capacity / 2 && + _sq_imm_window_size > 0) { + return SendImm(_new_rq_wrs.exchange(0, butil::memory_order_relaxed)); + } + return 0; +} + +ssize_t UrmaEndpoint::HandleCompletion(const urma_cr_t& cr) { + bool zerocopy = FLAGS_urma_recv_zerocopy; + if (cr.status != URMA_CR_SUCCESS) { + LOG(WARNING) << "URMA completion failed, status=" << cr.status; + errno = EIO; + return -1; + } + if (cr.flag.bs.s_r == 0) { + // Send completion: reclaim SQ window and wake the writer. + if (cr.user_ctx == 0) { + // Pure-ack WR: just replenish the imm budget. + if (_sq_imm_window_size >= RESERVED_WR_NUM) { + LOG(WARNING) + << "URMA credit-ACK completion exceeds reserved SQ " + "window: current=" + << _sq_imm_window_size + << " capacity=" << RESERVED_WR_NUM + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + _sq_imm_window_size += 1; + SendAck(0); + return 0; + } + uint16_t wnd = 1; // We signal every WR (complete_enable=1). + uint16_t old = + _sq_window_size.load(butil::memory_order_relaxed); + while (true) { + if (old >= _local_window_capacity) { + LOG(WARNING) + << "URMA send completion exceeds SQ window: old=" << old + << " increment=" << wnd + << " capacity=" << _local_window_capacity + << " user_ctx=" << cr.user_ctx + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + if (_sq_window_size.compare_exchange_weak( + old, static_cast<uint16_t>(old + wnd), + butil::memory_order_relaxed)) { + break; + } + } + for (uint16_t i = 0; i < wnd; ++i) { + _sbuf[_sq_sent].clear(); + _sq_sent = (_sq_sent + 1) % (_sq_size - RESERVED_WR_NUM); + } + butil::subtle::MemoryBarrier(); + if (_remote_rq_window_size.load(butil::memory_order_relaxed) >= + _local_window_capacity / 8) { + _socket->WakeAsEpollOut(); + } + return 0; + } + // Recv completion. + if (cr.opcode == URMA_CR_OPC_SEND_WITH_IMM && cr.imm_data > 0) { + if (cr.imm_data > _local_window_capacity) { + LOG(WARNING) << "Invalid URMA receive credit: " << cr.imm_data; + errno = EPROTO; + return -1; + } + const uint16_t acks = static_cast<uint16_t>(cr.imm_data); + uint16_t old = + _remote_rq_window_size.load(butil::memory_order_relaxed); + while (true) { + if (old > _local_window_capacity - acks) { + LOG(WARNING) + << "URMA receive credit exceeds window: old=" << old + << " credit=" << acks + << " capacity=" << _local_window_capacity + << " imm=" << cr.imm_data + << " remote_window_capacity=" + << _remote_window_capacity + << " on " << _socket->description(); + errno = EPROTO; + return -1; + } + if (_remote_rq_window_size.compare_exchange_weak( + old, static_cast<uint16_t>(old + acks), + butil::memory_order_relaxed)) { + break; + } + } + if (_sq_window_size.load(butil::memory_order_relaxed) > 0) { + _socket->WakeAsEpollOut(); + } + } else if (cr.completion_len == 0) { + LOG(WARNING) << "Zero-length URMA receive without immediate credit"; + errno = EPROTO; + return -1; + } + if (cr.completion_len > GetUrmaRecvBlockSize()) { + LOG(WARNING) << "URMA completion exceeds receive buffer: " + << cr.completion_len; + errno = EPROTO; + return -1; + } + if (cr.completion_len < static_cast<uint32_t>(FLAGS_urma_zerocopy_min_size)) { + zerocopy = false; + } + if (zerocopy) { + _rbuf[_rq_received].cutn(&_socket->_read_buf, cr.completion_len); + } else { + _socket->_read_buf.append(_rbuf_data[_rq_received], cr.completion_len); + } + if (PostRecv(1, zerocopy) < 0) { + return -1; + } + if (cr.completion_len > 0) { + SendAck(1); + } + return static_cast<ssize_t>(cr.completion_len); +} + +void UrmaEndpoint::DispatchReceivedBytes(SocketUniquePtr& s, ssize_t bytes) { + int64_t pending = _pending_received_bytes.load(butil::memory_order_relaxed); + if (bytes > 0) { + pending = _pending_received_bytes.fetch_add( + bytes, butil::memory_order_acq_rel) + bytes; + } + + const State state = _state.load(butil::memory_order_acquire); + if (state != ESTABLISHED) { + return; + } + + // PollCq and the handshake bthread can both reach this method when the + // state changes to ESTABLISHED. Serialize them so each byte added to + // _socket->_read_buf is reported to InputMessenger exactly once. + std::unique_lock<butil::Mutex> dispatch_lock(_dispatch_mutex); + if (_state.load(butil::memory_order_acquire) != ESTABLISHED) { + return; + } + pending = _pending_received_bytes.exchange( + 0, butil::memory_order_acq_rel); + if (pending <= 0 || s->Failed()) { + return; + } + + auto* messenger = static_cast<InputMessenger*>(s->user()); + if (!messenger) { + LOG(ERROR) << "URMA socket has no InputMessenger: " + << s->description(); + return; + } + + const int64_t received_us = butil::cpuwide_time_us(); + const int64_t base_realtime = butil::gettimeofday_us() - received_us; + InputMessageClosure last_msg; + messenger->ProcessNewMessage(s.get(), static_cast<ssize_t>(pending), + false, received_us, base_realtime, last_msg); +} + +void UrmaEndpoint::PollCq(Socket* m) { + auto* ep = static_cast<UrmaEndpoint*>(m->user()); + if (!ep || !ep->_resource || !ep->_resource->jfc) { + return; + } + SocketUniquePtr s; + if (Socket::Address(ep->_socket->id(), &s) != 0) { + return; + } + if (s->Failed()) { + return; + } + + const bool event_mode = !FLAGS_urma_use_polling; + int progress = Socket::PROGRESS_INIT; + while (true) { + urma_jfc_t* event_jfc = nullptr; + if (event_mode) { + const int event_count = ep->WaitCqEvent(s, &event_jfc); + if (event_count < 0) { + return; + } + if (event_count == 0) { + if (!m->MoreReadEvents(&progress)) { + return; + } + continue; + } + } + + ssize_t bytes = 0; + auto drain_cq = [&]() -> int { + while (true) { + const int n = + std::max(1, std::min<int>(FLAGS_urma_cqe_poll_once, 32)); + urma_cr_t crs[32]; + const int cnt = + urma_poll_jfc(ep->_resource->jfc, n, crs); + if (cnt < 0) { + return EIO; + } + if (cnt == 0) { + return 0; + } + for (int i = 0; i < cnt; ++i) { + if (s->Failed()) { + return ECANCELED; + } + const ssize_t nr = ep->HandleCompletion(crs[i]); + if (nr < 0) { + return errno ? errno : EIO; + } + bytes += nr; + } + } + }; + + int completion_error = drain_cq(); Review Comment: Should it be activated only in polling modeļ¼ not in event_mode? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
