Hugging Face's logo

ProprietaryLegal
/
dsv41-remote-engram-rdma

vllm
deepseek-v4
engram
rdma
roce
dgx-spark
gb10
offload
source-code
pli-labs
proprietarylegal
legal-ai
Model card Files
xet
 Community
dsv41-remote-engram-rdma / code /transport /ofl.cpp
ProprietaryLegal's picture
ProprietaryLegal
PLI Labs public snapshot 2026-10-28 (history squashed)
939ace7
Raw History Blame Contribute Delete
53.7 kB
// ofl transport implementation (see ofl.h). Builds on x86_64 (x86host) and aarch64 (DGX Spark, inside or outside the
// serve container) with g++ >= 9 and libibverbs >= 28. The QP setup follows the plugin's RDMA ring
// (vllm_exl3_moe/csrc/exl3_rdma.cu, itself from TensorFold): RC, path MTU = the ports' active MTU, timeout 14,
// retry 7, RNR retry 7, no relaxed ordering (the response header must land after the response bytes).
#include "ofl.h"
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <infiniband/verbs.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <time.h>
#include <unistd.h>
#include <algorithm>
#include <atomic>
#include <memory>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <utility>
#include <vector>
static_assert(sizeof(ofl_hdr) == 64, "ofl_hdr must be 64 bytes");
static_assert(sizeof(ofl_wire_seg) == 16, "ofl_wire_seg must be 16 bytes");
#if !defined(__BYTE_ORDER__) || __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
#error "ofl assumes little-endian hosts (x86_64, aarch64)"
#endif
namespace {
constexpr uint32_t kSegArea = 64 + 16 * OFL_MAX_SEGS; // request slot: header + room for every segment, then payload
constexpr uint64_t kWrSend = 1ull << 62; // client wr_id tags
constexpr uint64_t kWrHdr = 1ull << 62; // server wr_id tags: header write (low bits: buffer)
constexpr uint64_t kWrMark = 1ull << 61; // periodic signaled data write
constexpr int kSignalEvery = 32; // server: one signaled data write per this many
struct Hello { // bootstrap message payload, both directions
uint32_t backend;
uint32_t max_req_bytes, max_resp_bytes, max_inflight;
uint32_t qpn, psn;
uint8_t gid[16];
uint32_t mtu; // enum ibv_mtu
int32_t status;
char service[64];
char text[128]; // client: host name; ack: reject reason
};
static_assert(sizeof(Hello) == 240, "Hello layout");
inline int64_t now_ns() {
timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return ts.tv_sec * 1000000000LL + ts.tv_nsec;
}
inline void cpu_relax() {
#if defined(__aarch64__)
asm volatile("isb" ::: "memory");
#elif defined(__x86_64__)
__builtin_ia32_pause();
#endif
}
std::string fmt(const char* f, ...) __attribute__((format(printf, 1, 2)));
std::string fmt(const char* f, ...) {
char buf[1024];
va_list ap;
va_start(ap, f);
vsnprintf(buf, sizeof(buf), f, ap);
va_end(ap);
return buf;
}
void copy_err(char* err, size_t errlen, const std::string& msg) {
if (err && errlen) snprintf(err, errlen, "%s", msg.c_str());
}
bool parse_ipv4(const char* s, in_addr* out) { return s && *s && inet_pton(AF_INET, s, out) == 1; }
// ---- sockets: non-blocking, every operation against a deadline --------------------------------------------------
int wait_fd(int fd, short ev, int64_t deadline) { // 1 ready, 0 deadline passed, -1 error
for (;;) {
int64_t left = deadline - now_ns();
if (left <= 0) return 0;
timespec ts{left / 1000000000LL, left % 1000000000LL};
pollfd p{fd, ev, 0};
int r = ppoll(&p, 1, &ts, nullptr);
if (r > 0) return 1;
if (r == 0) return 0;
if (errno != EINTR) return -1;
}
}
// 0 ok, OFL_ETIMEDOUT, OFL_ESYS (msg), OFL_EBROKEN (peer closed)
int read_full(int fd, void* buf, size_t len, int64_t deadline, std::string& msg) {
uint8_t* p = static_cast<uint8_t*>(buf);
while (len) {
ssize_t r = ::recv(fd, p, len, MSG_DONTWAIT);
if (r > 0) { p += r; len -= (size_t)r; continue; }
if (r == 0) { msg = "peer closed the connection"; return OFL_EBROKEN; }
if (errno == EINTR) continue;
if (errno != EAGAIN && errno != EWOULDBLOCK) { msg = fmt("recv: %s", strerror(errno)); return OFL_ESYS; }
int w = wait_fd(fd, POLLIN, deadline);
if (w == 0) { msg = "timed out reading"; return OFL_ETIMEDOUT; }
if (w < 0) { msg = fmt("poll: %s", strerror(errno)); return OFL_ESYS; }
}
return 0;
}
int write_iov(int fd, iovec* iov, int n, int64_t deadline, std::string& msg) {
while (n > 0) {
msghdr mh{};
mh.msg_iov = iov;
mh.msg_iovlen = (size_t)n;
ssize_t r = ::sendmsg(fd, &mh, MSG_DONTWAIT | MSG_NOSIGNAL);
if (r < 0) {
if (errno == EINTR) continue;
if (errno != EAGAIN && errno != EWOULDBLOCK) { msg = fmt("send: %s", strerror(errno)); return OFL_ESYS; }
int w = wait_fd(fd, POLLOUT, deadline);
if (w == 0) { msg = "timed out writing"; return OFL_ETIMEDOUT; }
if (w < 0) { msg = fmt("poll: %s", strerror(errno)); return OFL_ESYS; }
continue;
}
size_t done = (size_t)r;
while (n > 0 && done >= iov->iov_len) { done -= iov->iov_len; iov++; n--; }
if (n > 0) { iov->iov_base = static_cast<uint8_t*>(iov->iov_base) + done; iov->iov_len -= done; }
}
return 0;
}
void tune_socket(int fd) {
int one = 1;
setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one));
int big = 8 << 20; // large socket buffers: the TCP backend moves multi-MB responses
setsockopt(fd, SOL_SOCKET, SO_SNDBUF, &big, sizeof(big));
setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &big, sizeof(big));
}
int tcp_connect(const char* local_ip, const char* server_ip, uint16_t port, int64_t deadline, std::string& msg) {
in_addr la{}, ra{};
if (!parse_ipv4(local_ip, &la)) { msg = fmt("bad local_ip '%s'", local_ip ? local_ip : ""); return -1; }
if (!parse_ipv4(server_ip, &ra)) { msg = fmt("bad server_ip '%s'", server_ip ? server_ip : ""); return -1; }
int fd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
if (fd < 0) { msg = fmt("socket: %s", strerror(errno)); return -1; }
sockaddr_in l{};
l.sin_family = AF_INET;
l.sin_addr = la;
if (::bind(fd, reinterpret_cast<sockaddr*>(&l), sizeof(l)) != 0) {
msg = fmt("bind %s: %s", local_ip, strerror(errno));
::close(fd);
return -1;
}
tune_socket(fd);
sockaddr_in r{};
r.sin_family = AF_INET;
r.sin_addr = ra;
r.sin_port = htons(port);
if (::connect(fd, reinterpret_cast<sockaddr*>(&r), sizeof(r)) != 0 && errno != EINPROGRESS) {
msg = fmt("connect %s:%u: %s", server_ip, port, strerror(errno));
::close(fd);
return -1;
}
int w = wait_fd(fd, POLLOUT, deadline);
int soerr = 0;
socklen_t sl = sizeof(soerr);
if (w <= 0 || getsockopt(fd, SOL_SOCKET, SO_ERROR, &soerr, &sl) != 0 || soerr != 0) {
msg = fmt("connect %s:%u from %s: %s", server_ip, port, local_ip,
w == 0 ? "timed out" : strerror(soerr ? soerr : errno));
::close(fd);
return -1;
}
return fd;
}
// ---- verbs -------------------------------------------------------------------------------------------------------
bool gid_is_ipv4(const ibv_gid& g, const in_addr& a) {
for (int i = 0; i < 10; i++)
if (g.raw[i]) return false;
return g.raw[10] == 0xff && g.raw[11] == 0xff && memcmp(&g.raw[12], &a, 4) == 0;
}
struct Dev {
ibv_context* ctx = nullptr;
ibv_pd* pd = nullptr;
std::string name;
int port = 1;
int gid_index = -1;
ibv_gid gid{};
ibv_mtu mtu = IBV_MTU_1024;
void close() {
if (pd) ibv_dealloc_pd(pd);
if (ctx) ibv_close_device(ctx);
pd = nullptr;
ctx = nullptr;
}
};
// Find (device, gid index) whose RoCE v2 GID is the IPv4-mapped `ip` (live discovery); check explicit choices.
bool open_dev(const char* ip, const char* want_dev, int port, int want_gid, Dev& d, std::string& msg) {
in_addr a{};
if (!parse_ipv4(ip, &a)) { msg = fmt("bad IPv4 address '%s'", ip ? ip : ""); return false; }
if (port <= 0) port = 1;
int n = 0;
ibv_device** list = ibv_get_device_list(&n);
if (!list || n <= 0) { msg = "no RDMA devices"; if (list) ibv_free_device_list(list); return false; }
std::string tried;
for (int i = 0; i < n && !d.ctx; i++) {
const char* name = ibv_get_device_name(list[i]);
if (want_dev && *want_dev && strcmp(want_dev, name) != 0) continue;
ibv_context* ctx = ibv_open_device(list[i]);
if (!ctx) { tried += fmt(" %s(open failed)", name); continue; }
ibv_port_attr pa{};
if (ibv_query_port(ctx, (uint8_t)port, &pa) != 0) { ibv_close_device(ctx); tried += fmt(" %s(no port %d)", name, port); continue; }
int found = -1;
for (int g = 0; g < pa.gid_tbl_len && found < 0; g++) {
if (want_gid >= 0 && g != want_gid) continue;
ibv_gid_entry e{};
if (ibv_query_gid_ex(ctx, (uint32_t)port, (uint32_t)g, &e, 0) != 0) continue;
if (e.gid_type == IBV_GID_TYPE_ROCE_V2 && gid_is_ipv4(e.gid, a)) found = g;
}
if (found < 0) {
tried += fmt(" %s", name);
ibv_close_device(ctx);
continue;
}
if (pa.state != IBV_PORT_ACTIVE) {
msg = fmt("%s port %d holds %s but is not active (state %s)", name, port, ip, ibv_port_state_str(pa.state));
ibv_close_device(ctx);
ibv_free_device_list(list);
return false;
}
d.ctx = ctx;
d.name = name;
d.port = port;
d.gid_index = found;
d.mtu = pa.active_mtu;
ibv_query_gid(ctx, (uint8_t)port, found, &d.gid);
}
ibv_free_device_list(list);
if (!d.ctx) {
msg = want_gid >= 0 ? fmt("no RoCE v2 GID index %d carries %s (checked:%s)", want_gid, ip, tried.c_str())
: fmt("no RDMA device has a RoCE v2 GID for %s (checked:%s)", ip, tried.c_str());
return false;
}
d.pd = ibv_alloc_pd(d.ctx);
if (!d.pd) { msg = fmt("ibv_alloc_pd on %s failed", d.name.c_str()); d.close(); return false; }
return true;
}
ibv_qp* make_qp(Dev& d, ibv_cq* scq, ibv_cq* rcq, uint32_t send_wr, uint32_t recv_wr, uint32_t send_sge,
uint32_t inline_bytes, uint32_t* inline_out, std::string& msg) {
ibv_qp_init_attr at{};
at.send_cq = scq;
at.recv_cq = rcq;
at.qp_type = IBV_QPT_RC;
at.cap.max_send_wr = send_wr;
at.cap.max_recv_wr = recv_wr;
at.cap.max_send_sge = send_sge;
at.cap.max_recv_sge = 1;
at.cap.max_inline_data = inline_bytes;
at.sq_sig_all = 0;
ibv_qp* qp = ibv_create_qp(d.pd, &at);
if (!qp) { msg = fmt("ibv_create_qp on %s: %s", d.name.c_str(), strerror(errno)); return nullptr; }
if (inline_out) *inline_out = at.cap.max_inline_data;
ibv_qp_attr a{};
a.qp_state = IBV_QPS_INIT;
a.pkey_index = 0;
a.port_num = (uint8_t)d.port;
a.qp_access_flags = IBV_ACCESS_REMOTE_WRITE;
if (ibv_modify_qp(qp, &a, IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) {
msg = fmt("QP to INIT: %s", strerror(errno));
ibv_destroy_qp(qp);
return nullptr;
}
return qp;
}
bool qp_connect(Dev& d, ibv_qp* qp, uint32_t my_psn, const Hello& peer, std::string& msg) {
ibv_qp_attr r{};
r.qp_state = IBV_QPS_RTR;
r.path_mtu = std::min<ibv_mtu>(d.mtu, static_cast<ibv_mtu>(peer.mtu));
r.dest_qp_num = peer.qpn;
r.rq_psn = peer.psn;
r.max_dest_rd_atomic = 1;
r.min_rnr_timer = 12;
r.ah_attr.is_global = 1;
memcpy(r.ah_attr.grh.dgid.raw, peer.gid, 16);
r.ah_attr.grh.sgid_index = (uint8_t)d.gid_index;
r.ah_attr.grh.hop_limit = 64;
r.ah_attr.port_num = (uint8_t)d.port;
if (ibv_modify_qp(qp, &r, IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | IBV_QP_RQ_PSN |
IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) != 0) {
msg = fmt("QP to RTR: %s", strerror(errno));
return false;
}
ibv_qp_attr s{};
s.qp_state = IBV_QPS_RTS;
s.timeout = 14;
s.retry_cnt = 7;
s.rnr_retry = 7;
s.sq_psn = my_psn;
s.max_rd_atomic = 1;
if (ibv_modify_qp(qp, &s, IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN |
IBV_QP_MAX_QP_RD_ATOMIC) != 0) {
msg = fmt("QP to RTS: %s", strerror(errno));
return false;
}
return true;
}
bool qp_error(ibv_qp* qp) { // from here on the NIC drops everything addressed to this QP; false if that failed
if (!qp) return true;
ibv_qp_attr a{};
a.qp_state = IBV_QPS_ERR;
return ibv_modify_qp(qp, &a, IBV_QP_STATE) == 0;
}
uint32_t rand_psn() {
static std::mt19937 g{std::random_device{}()};
static std::mutex mu;
std::lock_guard<std::mutex> l(mu);
return g() & 0xffffffu;
}
void* alloc_aligned(size_t n) {
void* p = nullptr;
if (posix_memalign(&p, 4096, std::max<size_t>(n, 4096)) != 0) return nullptr;
memset(p, 0, std::max<size_t>(n, 4096));
return p;
}
const char* mtu_str(ibv_mtu m) {
switch (m) {
case IBV_MTU_256: return "256";
case IBV_MTU_512: return "512";
case IBV_MTU_1024: return "1024";
case IBV_MTU_2048: return "2048";
case IBV_MTU_4096: return "4096";
default: return "?";
}
}
void fill_hdr(ofl_hdr& h, uint16_t kind) {
memset(&h, 0, sizeof(h));
h.magic = OFL_MAGIC;
h.version = OFL_VERSION;
h.kind = kind;
}
bool hdr_ok(const ofl_hdr& h, uint16_t kind) { return h.magic == OFL_MAGIC && h.version == OFL_VERSION && h.kind == kind; }
} // namespace
// =================================================================================================================
// client
struct ofl_client {
// config (owned copies)
int backend = 0;
std::string server_ip, local_ip, rdma_device, service;
uint16_t server_port = 0;
int rdma_port = 1, gid_index = -1;
uint32_t want_req = 0, want_resp = 0, want_inflight = 1;
int64_t connect_timeout_ns = 5000000000LL;
// connection
int fd = -1;
bool broken = true;
std::string err, desc;
uint64_t conn_gen = 0;
uint32_t max_req = 0, max_resp = 0, inflight = 1;
struct Region { uint8_t* addr; size_t len; ibv_mr* mr; };
std::vector<Region> regions;
enum { FREE = 0, PENDING = 1, DONE = 2 };
struct Slot {
int state = FREE;
uint64_t req_id = 0;
bool sent = false, answered = false;
int rc = 0;
int64_t t0 = 0;
uint32_t resp_max = 0;
std::vector<ofl_seg> segs;
ofl_result res{};
};
std::vector<Slot> slots;
uint64_t next_id = 1; // next request id
uint64_t done_upto = 0; // every id <= this has completed (or failed)
uint32_t idle_polls = 0;
// TCP payload staging (ofl_client_payload_buf)
std::vector<uint8_t> tcp_payload;
// RDMA
Dev dev;
ibv_cq* cq = nullptr;
ibv_qp* qp = nullptr;
uint32_t psn = 0;
uint8_t* send_mem = nullptr;
size_t send_slot = 0;
ibv_mr* send_mr = nullptr;
uint8_t* rhdr_mem = nullptr;
ibv_mr* rhdr_mr = nullptr;
uint32_t recv_posted = 0;
int fail(int code, const std::string& m) {
err = m;
return code;
}
// Break the connection: QP to error first (no late write can land), then drop everything.
void break_conn(const std::string& why) {
if (!err.empty() && broken) return;
err = why;
broken = true;
if (!qp_error(qp) && qp) { // could not stop the QP: destroy it (also ends delivery into our memory)
ibv_destroy_qp(qp);
qp = nullptr;
}
if (fd >= 0) { ::close(fd); fd = -1; }
for (auto& s : slots)
if (s.state == PENDING) { s.state = DONE; s.rc = OFL_EBROKEN; }
}
void teardown() {
qp_error(qp);
if (qp) ibv_destroy_qp(qp);
qp = nullptr;
if (cq) ibv_destroy_cq(cq);
cq = nullptr;
for (auto& r : regions)
if (r.mr) { ibv_dereg_mr(r.mr); r.mr = nullptr; }
if (send_mr) ibv_dereg_mr(send_mr);
if (rhdr_mr) ibv_dereg_mr(rhdr_mr);
send_mr = rhdr_mr = nullptr;
// send_mem / rhdr_mem stay allocated for the client's lifetime (stable ofl_client_payload_buf across
// reconnects); ofl_client_close frees them
dev.close();
if (fd >= 0) ::close(fd);
fd = -1;
broken = true;
}
int connect_now(std::string& m);
int progress(int64_t deadline, uint64_t want);
};
static bool post_client_recvs(ofl_client* c, uint32_t n) {
for (uint32_t i = 0; i < n; i++) {
ibv_recv_wr wr{}, *bad = nullptr;
wr.wr_id = 0;
if (ibv_post_recv(c->qp, &wr, &bad) != 0) return false;
c->recv_posted++;
}
return true;
}
int ofl_client::connect_now(std::string& m) {
const int64_t deadline = now_ns() + connect_timeout_ns;
err.clear();
fd = tcp_connect(local_ip.c_str(), server_ip.c_str(), server_port, deadline, m);
if (fd < 0) return OFL_ESYS;
Hello h{};
h.backend = (uint32_t)backend;
h.max_req_bytes = want_req;
h.max_resp_bytes = want_resp;
h.max_inflight = want_inflight;
snprintf(h.service, sizeof(h.service), "%s", service.c_str());
gethostname(h.text, sizeof(h.text) - 1);
if (backend == OFL_BACKEND_RDMA) {
if (!open_dev(local_ip.c_str(), rdma_device.c_str(), rdma_port, gid_index, dev, m)) return OFL_ERDMA;
cq = ibv_create_cq(dev.ctx, (int)(4 * want_inflight + 16), nullptr, nullptr, 0);
if (!cq) { m = "ibv_create_cq failed"; return OFL_ERDMA; }
qp = make_qp(dev, cq, cq, 2 * want_inflight + 8, 2 * want_inflight + 8, 2, 0, nullptr, m);
if (!qp) return OFL_ERDMA;
psn = rand_psn();
send_slot = (size_t)kSegArea + want_req;
send_slot = (send_slot + 63) & ~size_t(63);
if (!send_mem) send_mem = static_cast<uint8_t*>(alloc_aligned(send_slot * want_inflight)); // page-aligned
if (!rhdr_mem) rhdr_mem = static_cast<uint8_t*>(alloc_aligned(64 * (size_t)want_inflight));
if (!send_mem || !rhdr_mem) { m = "out of memory for RDMA slots"; return OFL_ENOMEM; }
send_mr = ibv_reg_mr(dev.pd, send_mem, send_slot * want_inflight, IBV_ACCESS_LOCAL_WRITE);
rhdr_mr = ibv_reg_mr(dev.pd, rhdr_mem, 64 * (size_t)want_inflight, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
if (!send_mr || !rhdr_mr) { m = fmt("ibv_reg_mr of the request slots failed: %s", strerror(errno)); return OFL_ERDMA; }
for (auto& r : regions) {
r.mr = ibv_reg_mr(dev.pd, r.addr, r.len, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
if (!r.mr) { m = fmt("ibv_reg_mr of a %zu-byte region failed: %s", r.len, strerror(errno)); return OFL_ERDMA; }
}
recv_posted = 0;
if (!post_client_recvs(this, 2 * want_inflight + 4)) { m = "ibv_post_recv failed"; return OFL_ERDMA; }
h.qpn = qp->qp_num;
h.psn = psn;
memcpy(h.gid, dev.gid.raw, 16);
h.mtu = dev.mtu;
}
ofl_hdr hd;
fill_hdr(hd, OFL_K_HELLO);
hd.payload_len = sizeof(Hello);
iovec v[2] = {{&hd, sizeof(hd)}, {&h, sizeof(h)}};
int rc = write_iov(fd, v, 2, deadline, m);
if (rc) return rc;
ofl_hdr ah;
Hello ack{};
if ((rc = read_full(fd, &ah, sizeof(ah), deadline, m))) return rc;
if (!hdr_ok(ah, OFL_K_HELLO_ACK) || ah.payload_len != sizeof(Hello)) { m = "bad handshake reply"; return OFL_EPROTO; }
if ((rc = read_full(fd, &ack, sizeof(ack), deadline, m))) return rc;
ack.text[sizeof(ack.text) - 1] = 0;
if (ack.status != 0) { m = fmt("server refused: %s", ack.text); return OFL_EREJECT; }
if (ack.max_req_bytes < want_req || ack.max_resp_bytes < want_resp || ack.max_inflight < 1 ||
ack.max_inflight > want_inflight) {
m = fmt("server granted req %u resp %u inflight %u for req %u resp %u inflight %u", ack.max_req_bytes,
ack.max_resp_bytes, ack.max_inflight, want_req, want_resp, want_inflight);
return OFL_EREJECT;
}
max_req = want_req;
max_resp = want_resp;
inflight = ack.max_inflight;
if (backend == OFL_BACKEND_RDMA) {
if (!qp_connect(dev, qp, psn, ack, m)) return OFL_ERDMA;
desc = fmt("rdma %s port %d gid %d mtu %s qpn %u -> %s:%u qpn %u (inflight %u)", dev.name.c_str(), dev.port,
dev.gid_index, mtu_str(std::min<ibv_mtu>(dev.mtu, (ibv_mtu)ack.mtu)), qp->qp_num, server_ip.c_str(),
server_port, ack.qpn, inflight);
} else {
desc = fmt("tcp %s -> %s:%u (inflight %u)", local_ip.c_str(), server_ip.c_str(), server_port, inflight);
if (tcp_payload.size() < want_req) tcp_payload.assign(want_req, 0); // kept across reconnects
}
slots.assign(want_inflight, Slot{});
done_upto = next_id - 1;
conn_gen++;
broken = false;
return OFL_OK;
}
extern "C" {
ofl_client* ofl_client_connect(const ofl_client_config* cfg, char* err, size_t errlen) {
if (!cfg || (cfg->backend != OFL_BACKEND_TCP && cfg->backend != OFL_BACKEND_RDMA) || !cfg->server_ip ||
!cfg->local_ip || !cfg->server_port) {
copy_err(err, errlen, "ofl_client_connect: backend, server_ip, server_port and local_ip are required");
return nullptr;
}
if (cfg->max_inflight > 64) { copy_err(err, errlen, "ofl_client_connect: max_inflight above 64"); return nullptr; }
auto* c = new ofl_client();
c->backend = cfg->backend;
c->server_ip = cfg->server_ip;
c->local_ip = cfg->local_ip;
c->rdma_device = cfg->rdma_device ? cfg->rdma_device : "";
c->service = cfg->service ? cfg->service : "";
c->server_port = cfg->server_port;
c->rdma_port = cfg->rdma_port > 0 ? cfg->rdma_port : 1;
c->gid_index = cfg->gid_index;
c->want_req = cfg->max_req_bytes;
c->want_resp = cfg->max_resp_bytes;
c->want_inflight = cfg->max_inflight ? cfg->max_inflight : 1;
if (cfg->connect_timeout_ms > 0) c->connect_timeout_ns = cfg->connect_timeout_ms * 1000000LL;
std::string m;
int rc = c->connect_now(m);
if (rc != OFL_OK) {
copy_err(err, errlen, m);
c->teardown();
free(c->send_mem);
free(c->rhdr_mem);
delete c;
return nullptr;
}
return c;
}
int ofl_client_reconnect(ofl_client* c, char* err, size_t errlen) {
if (!c) return OFL_EINVAL;
c->teardown();
std::string m;
int rc = c->connect_now(m);
if (rc != OFL_OK) {
c->err = m;
copy_err(err, errlen, m);
c->teardown();
}
return rc;
}
int ofl_client_reg(ofl_client* c, void* addr, size_t len, uint32_t* region) {
if (!c || !addr || !len || !region) return OFL_EINVAL;
ofl_client::Region r{static_cast<uint8_t*>(addr), len, nullptr};
if (c->backend == OFL_BACKEND_RDMA && !c->broken) {
r.mr = ibv_reg_mr(c->dev.pd, addr, len, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
if (!r.mr) return c->fail(OFL_ERDMA, fmt("ibv_reg_mr of %zu bytes at %p: %s", len, addr, strerror(errno)));
}
c->regions.push_back(r);
*region = (uint32_t)(c->regions.size() - 1);
return OFL_OK;
}
void* ofl_client_payload_buf(ofl_client* c, uint32_t* capacity) {
if (!c || c->broken) return nullptr;
if (capacity) *capacity = c->max_req;
if (c->backend == OFL_BACKEND_TCP) return c->tcp_payload.data();
const uint32_t slot = (uint32_t)((c->next_id - 1) % c->inflight);
return c->send_mem + (size_t)slot * c->send_slot + kSegArea;
}
int64_t ofl_client_submit(ofl_client* c, uint16_t op, uint64_t generation, const void* payload, uint32_t payload_len,
const ofl_seg* segs, uint32_t nseg) {
if (!c) return OFL_EINVAL;
if (c->broken) return c->fail(OFL_EBROKEN, c->err.empty() ? "not connected" : c->err);
if (payload_len > c->max_req) return c->fail(OFL_ETOOBIG, fmt("payload %u > max_req %u", payload_len, c->max_req));
if (nseg > OFL_MAX_SEGS || (nseg && !segs)) return c->fail(OFL_EINVAL, "too many destination ranges");
uint64_t cap = 0;
for (uint32_t i = 0; i < nseg; i++) {
if (segs[i].region >= c->regions.size()) return c->fail(OFL_EINVAL, fmt("unknown region %u", segs[i].region));
const auto& r = c->regions[segs[i].region];
if (segs[i].offset > r.len || segs[i].len > r.len - segs[i].offset)
return c->fail(OFL_EINVAL, fmt("range %u outside region %u", i, segs[i].region));
cap += segs[i].len;
}
if (cap > c->max_resp) return c->fail(OFL_ETOOBIG, fmt("destination %llu > max_resp %u", (unsigned long long)cap, c->max_resp));
const uint64_t id = c->next_id;
const uint32_t si = (uint32_t)((id - 1) % c->inflight);
auto& s = c->slots[si];
if (s.state != ofl_client::FREE) return c->fail(OFL_EBUSY, "too many requests in flight (wait first)");
s = ofl_client::Slot{};
s.state = ofl_client::PENDING;
s.req_id = id;
s.resp_max = (uint32_t)cap;
s.segs.assign(segs, segs + nseg);
s.t0 = now_ns();
ofl_hdr h;
fill_hdr(h, OFL_K_REQ);
h.op = op;
h.req_id = id;
h.generation = generation;
h.payload_len = payload_len;
h.resp_max = (uint32_t)cap;
std::string m;
if (c->backend == OFL_BACKEND_TCP) {
iovec v[2] = {{&h, sizeof(h)}, {const_cast<void*>(payload), payload_len}};
int rc = write_iov(c->fd, v, payload_len ? 2 : 1, now_ns() + 10000000000LL, m);
if (rc) { c->break_conn("send request: " + m); s.state = ofl_client::FREE; return rc == OFL_ETIMEDOUT ? OFL_ETIMEDOUT : OFL_EBROKEN; }
c->next_id++;
return (int64_t)id;
}
uint8_t* base = c->send_mem + (size_t)si * c->send_slot;
h.nseg = (uint16_t)nseg;
h.rhdr_addr = (uint64_t)(uintptr_t)(c->rhdr_mem + 64 * (size_t)si);
h.rhdr_rkey = c->rhdr_mr->rkey;
memset(c->rhdr_mem + 64 * (size_t)si, 0, 64);
memcpy(base, &h, sizeof(h));
auto* ws = reinterpret_cast<ofl_wire_seg*>(base + 64);
for (uint32_t i = 0; i < nseg; i++) {
const auto& r = c->regions[segs[i].region];
ws[i].addr = (uint64_t)(uintptr_t)(r.addr + segs[i].offset);
ws[i].rkey = r.mr->rkey;
ws[i].len = (uint32_t)segs[i].len;
}
uint8_t* pl = base + kSegArea;
if (payload_len && payload != pl) memcpy(pl, payload, payload_len);
ibv_sge sge[2] = {{(uint64_t)(uintptr_t)base, 64 + 16 * nseg, c->send_mr->lkey},
{(uint64_t)(uintptr_t)pl, payload_len, c->send_mr->lkey}};
ibv_send_wr wr{}, *bad = nullptr;
wr.wr_id = kWrSend | id;
wr.sg_list = sge;
wr.num_sge = payload_len ? 2 : 1;
wr.opcode = IBV_WR_SEND;
wr.send_flags = IBV_SEND_SIGNALED;
if (ibv_post_send(c->qp, &wr, &bad) != 0) {
s.state = ofl_client::FREE;
c->break_conn(fmt("ibv_post_send: %s", strerror(errno)));
return OFL_ERDMA;
}
c->next_id++;
return (int64_t)id;
}
} // extern "C"
// Drive completions until request `want` is done or the deadline passes. 0, or an error (connection broken).
int ofl_client::progress(int64_t deadline, uint64_t want) {
std::string m;
while (done_upto < want) {
const uint64_t id = done_upto + 1;
Slot& s = slots[(id - 1) % inflight];
if (s.state == DONE) { done_upto++; continue; }
if (s.state != PENDING || s.req_id != id) { break_conn("internal: slot out of order"); return OFL_EPROTO; }
if (backend == OFL_BACKEND_TCP) {
ofl_hdr h;
int rc = read_full(fd, &h, sizeof(h), deadline, m);
if (rc) { break_conn("read response: " + m); return rc == OFL_ETIMEDOUT ? OFL_ETIMEDOUT : rc; }
if (!hdr_ok(h, OFL_K_RESP) || h.req_id != id) {
break_conn(fmt("bad response header (req %llu)", (unsigned long long)id));
return OFL_EPROTO;
}
if (h.payload_len > s.resp_max) { break_conn("response larger than its destination"); return OFL_EPROTO; }
uint32_t left = h.payload_len;
for (const auto& g : s.segs) {
if (!left) break;
const uint32_t n = (uint32_t)std::min<uint64_t>(g.len, left);
rc = read_full(fd, regions[g.region].addr + g.offset, n, deadline, m);
if (rc) { break_conn("read response body: " + m); return rc == OFL_ETIMEDOUT ? OFL_ETIMEDOUT : rc; }
left -= n;
}
s.res.status = h.status;
s.res.resp_len = h.payload_len;
s.res.req_id = id;
s.res.generation = h.generation;
s.res.server_ns = (int64_t)h.t_server_ns;
s.res.elapsed_ns = now_ns() - s.t0;
s.rc = h.status < 0 ? OFL_EREMOTE : OFL_OK;
s.state = DONE;
done_upto++;
continue;
}
// RDMA: poll; every completion applies to some pending slot
ibv_wc wc[16];
const int n = ibv_poll_cq(cq, 16, wc);
if (n < 0) { break_conn("ibv_poll_cq failed"); return OFL_ERDMA; }
if (n == 0) {
if (now_ns() > deadline) {
break_conn(fmt("timed out waiting for response %llu", (unsigned long long)id));
return OFL_ETIMEDOUT;
}
// the server never writes to the bootstrap socket after the handshake: readable = it closed (died / restarted),
// so fail now instead of at the deadline (checked every 1024 empty polls, ~50-100 us)
if ((++idle_polls & 1023) == 0) {
pollfd p{fd, POLLIN, 0};
if (fd >= 0 && poll(&p, 1, 0) > 0) {
break_conn(fmt("server closed the connection while request %llu was pending", (unsigned long long)id));
return OFL_EBROKEN;
}
}
cpu_relax();
continue;
}
for (int i = 0; i < n; i++) {
if (wc[i].status != IBV_WC_SUCCESS) {
break_conn(fmt("work completion: %s (opcode %d)", ibv_wc_status_str(wc[i].status), (int)wc[i].opcode));
return OFL_ERDMA;
}
if (wc[i].opcode == IBV_WC_SEND) {
const uint64_t rid = wc[i].wr_id & ~kWrSend;
Slot& t = slots[(rid - 1) % inflight];
if (t.req_id == rid) t.sent = true;
} else if (wc[i].opcode == IBV_WC_RECV_RDMA_WITH_IMM) {
recv_posted--;
if (!post_client_recvs(this, 1)) { break_conn("ibv_post_recv failed"); return OFL_ERDMA; }
const uint32_t imm = ntohl(wc[i].imm_data);
// responses come in request order; find the pending slot whose id matches the immediate
Slot* t = nullptr;
for (auto& x : slots)
if (x.state == PENDING && (uint32_t)x.req_id == imm) t = &x;
if (!t) { break_conn(fmt("response for an unknown request (imm %u)", imm)); return OFL_EPROTO; }
const size_t si = (size_t)(t - slots.data());
ofl_hdr h;
memcpy(&h, rhdr_mem + 64 * si, sizeof(h));
if (!hdr_ok(h, OFL_K_RESP) || h.req_id != t->req_id || h.payload_len > t->resp_max) {
break_conn(fmt("bad response header for request %llu", (unsigned long long)t->req_id));
return OFL_EPROTO;
}
t->answered = true;
t->res.status = h.status;
t->res.resp_len = h.payload_len;
t->res.req_id = t->req_id;
t->res.generation = h.generation;
t->res.server_ns = (int64_t)h.t_server_ns;
t->res.elapsed_ns = now_ns() - t->t0;
t->rc = h.status < 0 ? OFL_EREMOTE : OFL_OK;
} else {
break_conn(fmt("unexpected completion opcode %d", (int)wc[i].opcode));
return OFL_EPROTO;
}
}
// a slot is complete once its request buffer is free again and its response is in
for (auto& x : slots)
if (x.state == PENDING && x.sent && x.answered) x.state = DONE;
}
return OFL_OK;
}
extern "C" {
int ofl_client_wait(ofl_client* c, int64_t ticket, ofl_result* res, int64_t timeout_ns) {
if (!c || ticket <= 0) return OFL_EINVAL;
const uint64_t id = (uint64_t)ticket;
if (id >= c->next_id) return c->fail(OFL_EINVAL, "unknown ticket");
auto& s = c->slots[(id - 1) % c->inflight];
if (s.req_id != id || s.state == ofl_client::FREE) return c->fail(OFL_EINVAL, "ticket already waited for");
if (s.state == ofl_client::PENDING) {
if (c->broken) { s.state = ofl_client::FREE; return OFL_EBROKEN; }
int rc = c->progress(now_ns() + timeout_ns, id);
if (rc != OFL_OK) {
if (s.req_id == id) s.state = ofl_client::FREE;
return rc;
}
}
if (res) *res = s.res;
const int rc = s.rc;
s.state = ofl_client::FREE;
if (rc == OFL_EREMOTE) c->err = fmt("server returned status %d", s.res.status);
return rc;
}
int ofl_client_call(ofl_client* c, uint16_t op, uint64_t generation, const void* payload, uint32_t payload_len,
const ofl_seg* segs, uint32_t nseg, ofl_result* res, int64_t timeout_ns) {
const int64_t t = ofl_client_submit(c, op, generation, payload, payload_len, segs, nseg);
if (t < 0) return (int)t;
return ofl_client_wait(c, t, res, timeout_ns);
}
int ofl_client_broken(const ofl_client* c) { return !c || c->broken; }
const char* ofl_client_error(const ofl_client* c) { return c ? c->err.c_str() : "null client"; }
const char* ofl_client_describe(const ofl_client* c) { return c ? c->desc.c_str() : ""; }
void ofl_client_close(ofl_client* c) {
if (!c) return;
c->teardown();
free(c->send_mem);
free(c->rhdr_mem);
delete c;
}
} // extern "C"
// =================================================================================================================
// server
struct ofl_server;
struct Conn {
ofl_server* s = nullptr;
uint64_t id = 0;
int fd = -1;
int backend = 0;
std::thread thr;
std::atomic<bool> finished{false};
std::string peer;
uint32_t max_req = 0, max_resp = 0, inflight = 1;
// TCP
std::vector<uint8_t> req_buf, resp_buf;
// RDMA
ibv_comp_channel* ch = nullptr;
ibv_cq* scq = nullptr;
ibv_cq* rcq = nullptr;
ibv_qp* qp = nullptr;
uint32_t psn = 0;
uint8_t* recv_mem = nullptr;
size_t recv_slot = 0;
ibv_mr* recv_mr = nullptr;
uint8_t* resp_mem = nullptr;
ibv_mr* resp_mr = nullptr;
std::vector<int> resp_free;
std::mutex post_mu; // posting on the SQ and polling the send CQ
uint64_t posted = 0; // send WRs posted
uint64_t retired = 0; // send WRs known complete
uint32_t sq_depth = 0;
uint32_t max_inline = 0;
bool sq_failed = false;
std::string fail_why;
~Conn();
void main();
bool handshake();
void run_tcp();
void run_rdma();
bool handle_recv(const ibv_wc& w); // false = the connection must end
bool drain_send(bool block); // false = a send completion failed
};
struct ofl_req {
Conn* c = nullptr;
ofl_hdr hdr{};
const ofl_wire_seg* segs = nullptr;
uint32_t nseg = 0;
const uint8_t* payload = nullptr;
uint32_t payload_len = 0;
uint8_t* resp = nullptr;
uint32_t cap = 0, len = 0;
int buf = -1; // RDMA response buffer index
std::mutex mu;
std::vector<std::pair<uint32_t, uint32_t>> sent; // ranges already posted
int err = 0;
};
struct ofl_server {
std::string bind_ip, rdma_device, service, desc;
uint16_t port = 0;
int backends = 0, rdma_port = 1, gid_index = -1;
uint32_t max_req = 0, max_resp = 0, max_inflight = 2, max_conns = 8;
int64_t spin_ns = 2000000000LL;
ofl_handler_fn fn = nullptr;
void* user = nullptr;
int lfd = -1;
Dev dev; // shared device + PD (RDMA)
bool have_dev = false;
std::atomic<bool> running{false};
std::thread acceptor;
std::mutex mu;
std::vector<std::unique_ptr<Conn>> conns;
uint64_t next_conn = 1;
std::atomic<uint64_t> n_conns{0}, n_reqs{0}, n_bytes{0}, n_errs{0};
void accept_loop();
void handshake(int fd);
};
Conn::~Conn() {
if (thr.joinable()) thr.join();
qp_error(qp);
if (qp) ibv_destroy_qp(qp);
if (scq) ibv_destroy_cq(scq);
if (rcq) ibv_destroy_cq(rcq);
if (ch) ibv_destroy_comp_channel(ch);
if (recv_mr) ibv_dereg_mr(recv_mr);
if (resp_mr) ibv_dereg_mr(resp_mr);
free(recv_mem);
free(resp_mem);
if (fd >= 0) ::close(fd);
}
namespace {
// Map response bytes [off, off + len) onto the request's destination ranges and post one RDMA WRITE per piece.
// Caller holds c->post_mu.
int post_range_locked(ofl_req* r, uint32_t off, uint32_t len) {
Conn* c = r->c;
uint64_t pos = 0;
for (uint32_t i = 0; i < r->nseg && len; i++) {
const uint64_t s0 = pos, s1 = pos + r->segs[i].len;
pos = s1;
if (off >= s1 || off + (uint64_t)len <= s0) continue;
const uint64_t a = std::max<uint64_t>(off, s0), b = std::min<uint64_t>(off + (uint64_t)len, s1);
while (c->posted - c->retired + 2 > c->sq_depth)
if (!c->drain_send(true)) return OFL_ERDMA;
ibv_sge sge{(uint64_t)(uintptr_t)(r->resp + a), (uint32_t)(b - a), c->resp_mr->lkey};
ibv_send_wr wr{}, *bad = nullptr;
const uint64_t seq = ++c->posted;
wr.wr_id = kWrMark | seq;
wr.sg_list = &sge;
wr.num_sge = 1;
wr.opcode = IBV_WR_RDMA_WRITE;
wr.send_flags = (seq % kSignalEvery == 0) ? IBV_SEND_SIGNALED : 0;
wr.wr.rdma.remote_addr = r->segs[i].addr + (a - s0);
wr.wr.rdma.rkey = r->segs[i].rkey;
if (ibv_post_send(c->qp, &wr, &bad) != 0) {
c->posted--;
return OFL_ERDMA;
}
}
return OFL_OK;
}
} // namespace
bool Conn::drain_send(bool block) {
ibv_wc wc[32];
for (;;) {
const int n = ibv_poll_cq(scq, 32, wc);
if (n < 0) { sq_failed = true; fail_why = "ibv_poll_cq (send) failed"; return false; }
for (int i = 0; i < n; i++) {
if (wc[i].status != IBV_WC_SUCCESS) {
sq_failed = true;
fail_why = fmt("send completion: %s", ibv_wc_status_str(wc[i].status));
return false;
}
const uint64_t w = wc[i].wr_id;
if (w & kWrHdr) {
const uint64_t seq = (w >> 16) & 0xffffffffffull;
resp_free.push_back((int)(w & 0xffff));
retired = std::max(retired, seq);
} else if (w & kWrMark) {
retired = std::max(retired, w & ~kWrMark);
}
}
if (n > 0 || !block) return true;
if (!s->running.load(std::memory_order_relaxed)) { fail_why = "server stopping"; return false; }
cpu_relax();
}
}
static int32_t run_handler(ofl_server* s, ofl_req* r) {
if (r->hdr.op == OFL_OP_ECHO) {
const uint32_t n = std::min(r->payload_len, r->cap);
if (n) memcpy(r->resp, r->payload, n);
r->len = n;
return 0;
}
return s->fn ? s->fn(s->user, r) : -1;
}
void Conn::run_tcp() {
std::string m;
while (s->running.load(std::memory_order_relaxed)) {
const int w = wait_fd(fd, POLLIN, now_ns() + 100000000LL);
if (w == 0) continue;
if (w < 0) break;
ofl_hdr h;
const int64_t dl = now_ns() + 10000000000LL;
if (read_full(fd, &h, sizeof(h), dl, m)) break;
const int64_t t0 = now_ns();
if (!hdr_ok(h, OFL_K_REQ) || h.payload_len > max_req) { s->n_errs++; break; }
if (read_full(fd, req_buf.data(), h.payload_len, dl, m)) break;
ofl_req r;
r.c = this;
r.hdr = h;
r.payload = req_buf.data();
r.payload_len = h.payload_len;
r.resp = resp_buf.data();
r.cap = std::min<uint32_t>(h.resp_max, max_resp);
int32_t st = run_handler(s, &r);
if (r.len > r.cap) st = OFL_ETOOBIG;
ofl_hdr o;
fill_hdr(o, OFL_K_RESP);
o.op = h.op;
o.req_id = h.req_id;
o.generation = h.generation;
o.status = st;
o.payload_len = st < 0 ? 0 : r.len;
o.t_server_ns = (uint64_t)(now_ns() - t0);
iovec v[2] = {{&o, sizeof(o)}, {resp_buf.data(), o.payload_len}};
if (write_iov(fd, v, o.payload_len ? 2 : 1, now_ns() + 10000000000LL, m)) break;
s->n_reqs++;
s->n_bytes += o.payload_len;
if (st < 0) s->n_errs++;
}
}
void Conn::run_rdma() {
int64_t last = now_ns();
uint32_t idle_iter = 0;
ibv_wc wc[8];
while (s->running.load(std::memory_order_relaxed) && !sq_failed) {
int n = ibv_poll_cq(rcq, 8, wc);
if (n < 0) { fail_why = "ibv_poll_cq (recv) failed"; break; }
if (n == 0) {
{
std::lock_guard<std::mutex> g(post_mu);
if (!drain_send(false)) break;
}
if (now_ns() - last < s->spin_ns) {
if ((++idle_iter & 0xffff) == 0) { // every few ms while spinning: is the client still there?
pollfd p{fd, POLLIN, 0};
if (poll(&p, 1, 0) > 0) { fail_why = "client closed the bootstrap socket"; break; }
}
cpu_relax();
continue;
}
// idle: arm the completion channel, re-check (a completion may have raced the arm), then sleep on the channel
// and the bootstrap socket (the client never writes to it after the handshake: readable = closed)
if (ibv_req_notify_cq(rcq, 0) != 0) { fail_why = "ibv_req_notify_cq failed"; break; }
n = ibv_poll_cq(rcq, 8, wc);
if (n < 0) { fail_why = "ibv_poll_cq (recv) failed"; break; }
if (n == 0) {
pollfd p[2] = {{ch->fd, POLLIN, 0}, {fd, POLLIN, 0}};
const int r = poll(p, 2, 100);
if (r > 0 && (p[1].revents & (POLLIN | POLLHUP | POLLERR))) { fail_why = "client closed the bootstrap socket"; break; }
if (r > 0 && (p[0].revents & POLLIN)) {
ibv_cq* evcq = nullptr;
void* evctx = nullptr;
if (ibv_get_cq_event(ch, &evcq, &evctx) == 0) ibv_ack_cq_events(evcq, 1);
}
continue;
}
}
bool ok = true;
for (int i = 0; i < n && ok; i++) ok = handle_recv(wc[i]);
if (!ok) break;
last = now_ns();
idle_iter = 0;
}
if (!fail_why.empty() && s->running.load() && fail_why != "client closed the bootstrap socket")
fprintf(stderr, "ofl server: connection %llu (%s) ends: %s\n",
(unsigned long long)id, peer.c_str(), fail_why.c_str());
}
bool Conn::handle_recv(const ibv_wc& w) {
if (w.status != IBV_WC_SUCCESS) { fail_why = fmt("recv completion: %s", ibv_wc_status_str(w.status)); return false; }
if (w.opcode != IBV_WC_RECV) { fail_why = fmt("unexpected recv opcode %d", (int)w.opcode); return false; }
const int64_t t0 = now_ns();
const uint32_t k = (uint32_t)w.wr_id;
if (k >= inflight) { fail_why = "bad receive slot"; return false; }
uint8_t* base = recv_mem + (size_t)k * recv_slot;
ofl_req r;
memcpy(&r.hdr, base, sizeof(ofl_hdr));
const ofl_hdr& h = r.hdr;
if (!hdr_ok(h, OFL_K_REQ) || h.nseg > OFL_MAX_SEGS || h.payload_len > max_req ||
w.byte_len != 64u + 16u * h.nseg + h.payload_len) {
s->n_errs++;
fail_why = fmt("malformed request (%u bytes)", w.byte_len);
return false;
}
r.c = this;
r.segs = reinterpret_cast<const ofl_wire_seg*>(base + 64);
r.nseg = h.nseg;
r.payload = base + 64 + 16 * h.nseg;
r.payload_len = h.payload_len;
uint64_t cap = 0;
for (uint32_t i = 0; i < r.nseg; i++) cap += r.segs[i].len;
if (cap != h.resp_max) { s->n_errs++; fail_why = "request ranges do not add up to resp_max"; return false; }
{
std::lock_guard<std::mutex> g(post_mu);
while (resp_free.empty())
if (!drain_send(true)) return false;
r.buf = resp_free.back();
resp_free.pop_back();
}
r.resp = resp_mem + (size_t)r.buf * ((size_t)max_resp + 64);
r.cap = (uint32_t)std::min<uint64_t>(cap, max_resp);
int32_t st = run_handler(s, &r);
if (st >= 0 && r.len > r.cap) st = OFL_ETOOBIG;
const uint32_t len = st < 0 ? 0 : r.len;
std::lock_guard<std::mutex> g(post_mu);
if (r.err) { fail_why = "posting response writes failed"; sq_failed = true; return false; }
if (len) { // everything the handler did not already send
std::sort(r.sent.begin(), r.sent.end());
uint32_t pos = 0;
for (const auto& x : r.sent) {
if (x.first > pos && post_range_locked(&r, pos, std::min(x.first, len) - pos)) { fail_why = "ibv_post_send failed"; return false; }
pos = std::max(pos, x.first + x.second);
if (pos >= len) break;
}
if (pos < len && post_range_locked(&r, pos, len - pos)) { fail_why = "ibv_post_send failed"; return false; }
}
// the receive slot is free again (the range table was read by the posts above; the header fields come from the
// copy in r.hdr): repost it before the header, so the client's next request never meets an empty receive queue
ibv_sge rs{(uint64_t)(uintptr_t)base, (uint32_t)recv_slot, recv_mr->lkey};
ibv_recv_wr rw{}, *rbad = nullptr;
rw.wr_id = k;
rw.sg_list = &rs;
rw.num_sge = 1;
if (ibv_post_recv(qp, &rw, &rbad) != 0) { fail_why = "ibv_post_recv failed"; return false; }
while (posted - retired + 2 > sq_depth)
if (!drain_send(true)) return false;
ofl_hdr* o = reinterpret_cast<ofl_hdr*>(r.resp + max_resp); // the buffer's header area (registered)
fill_hdr(*o, OFL_K_RESP);
o->op = h.op;
o->req_id = h.req_id;
o->generation = h.generation;
o->status = st;
o->payload_len = len;
o->t_server_ns = (uint64_t)(now_ns() - t0);
ibv_sge sge{(uint64_t)(uintptr_t)o, 64, resp_mr->lkey};
ibv_send_wr wr{}, *bad = nullptr;
const uint64_t seq = ++posted;
wr.wr_id = kWrHdr | (seq << 16) | (uint64_t)r.buf;
wr.sg_list = &sge;
wr.num_sge = 1;
wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM;
wr.send_flags = IBV_SEND_SIGNALED | (max_inline >= 64 ? IBV_SEND_INLINE : 0);
wr.imm_data = htonl((uint32_t)h.req_id);
wr.wr.rdma.remote_addr = h.rhdr_addr;
wr.wr.rdma.rkey = h.rhdr_rkey;
if (ibv_post_send(qp, &wr, &bad) != 0) { posted--; fail_why = fmt("ibv_post_send (header): %s", strerror(errno)); return false; }
s->n_reqs++;
s->n_bytes += len;
if (st < 0) s->n_errs++;
return true;
}
bool Conn::handshake() {
std::string m;
const int64_t dl = now_ns() + 5000000000LL;
ofl_hdr h;
Hello hi{};
if (read_full(fd, &h, sizeof(h), dl, m) || !hdr_ok(h, OFL_K_HELLO) || h.payload_len != sizeof(Hello) ||
read_full(fd, &hi, sizeof(hi), dl, m)) {
fail_why = "bad or missing handshake: " + m;
return false;
}
hi.service[sizeof(hi.service) - 1] = 0;
hi.text[sizeof(hi.text) - 1] = 0;
peer = hi.text;
Hello ack{};
auto reply = [&](int32_t status, const std::string& why) {
ack.status = status;
snprintf(ack.text, sizeof(ack.text), "%s", why.c_str());
snprintf(ack.service, sizeof(ack.service), "%s", s->service.c_str());
ofl_hdr o;
fill_hdr(o, OFL_K_HELLO_ACK);
o.payload_len = sizeof(Hello);
iovec v[2] = {{&o, sizeof(o)}, {&ack, sizeof(ack)}};
std::string m2;
return write_iov(fd, v, 2, dl, m2) == 0;
};
backend = (int)hi.backend;
std::string why;
if (s->service != hi.service) why = fmt("service '%s' asked, this is '%s'", hi.service, s->service.c_str());
else if ((backend != OFL_BACKEND_TCP && backend != OFL_BACKEND_RDMA) || !(s->backends & (1 << backend)))
why = fmt("backend %d not offered", backend);
else if (hi.max_req_bytes > s->max_req || hi.max_resp_bytes > s->max_resp)
why = fmt("limits req %u resp %u exceed the server's %u / %u", hi.max_req_bytes, hi.max_resp_bytes, s->max_req, s->max_resp);
else if (backend == OFL_BACKEND_RDMA && !s->have_dev) why = "no RDMA device on this server";
if (!why.empty()) {
reply(OFL_EREJECT, why);
fail_why = "refused " + peer + ": " + why;
return false;
}
max_req = hi.max_req_bytes;
max_resp = hi.max_resp_bytes;
inflight = std::max<uint32_t>(1, std::min(hi.max_inflight ? hi.max_inflight : 1, s->max_inflight));
if (backend == OFL_BACKEND_TCP) {
req_buf.assign(std::max<uint32_t>(max_req, 1), 0);
resp_buf.assign(std::max<uint32_t>(max_resp, 1), 0);
} else {
Dev& d = s->dev;
sq_depth = 1024;
ch = ibv_create_comp_channel(d.ctx);
rcq = ch ? ibv_create_cq(d.ctx, (int)inflight + 8, nullptr, ch, 0) : nullptr;
scq = ibv_create_cq(d.ctx, (int)sq_depth + 16, nullptr, nullptr, 0);
if (!ch || !rcq || !scq) { reply(OFL_ERDMA, "completion queue setup failed"); fail_why = "cq setup failed"; return false; }
fcntl(ch->fd, F_SETFL, fcntl(ch->fd, F_GETFL) | O_NONBLOCK);
qp = make_qp(d, scq, rcq, sq_depth, inflight + 2, 1, 64, &max_inline, m);
if (!qp) qp = make_qp(d, scq, rcq, sq_depth, inflight + 2, 1, 0, &max_inline, m); // no inline support
if (!qp) { reply(OFL_ERDMA, m); fail_why = m; return false; }
recv_slot = ((size_t)kSegArea + max_req + 63) & ~size_t(63);
recv_mem = static_cast<uint8_t*>(alloc_aligned(recv_slot * inflight));
resp_mem = static_cast<uint8_t*>(alloc_aligned(((size_t)max_resp + 64) * inflight));
if (!recv_mem || !resp_mem) { reply(OFL_ENOMEM, "out of memory"); fail_why = "out of memory"; return false; }
recv_mr = ibv_reg_mr(d.pd, recv_mem, recv_slot * inflight, IBV_ACCESS_LOCAL_WRITE);
resp_mr = ibv_reg_mr(d.pd, resp_mem, ((size_t)max_resp + 64) * inflight, IBV_ACCESS_LOCAL_WRITE);
if (!recv_mr || !resp_mr) { reply(OFL_ERDMA, "ibv_reg_mr failed"); fail_why = fmt("ibv_reg_mr: %s", strerror(errno)); return false; }
for (uint32_t k = 0; k < inflight; k++) {
resp_free.push_back((int)k);
ibv_sge rs{(uint64_t)(uintptr_t)(recv_mem + k * recv_slot), (uint32_t)recv_slot, recv_mr->lkey};
ibv_recv_wr rw{}, *bad = nullptr;
rw.wr_id = k;
rw.sg_list = &rs;
rw.num_sge = 1;
if (ibv_post_recv(qp, &rw, &bad) != 0) { reply(OFL_ERDMA, "ibv_post_recv failed"); fail_why = "ibv_post_recv failed"; return false; }
}
psn = rand_psn();
if (!qp_connect(d, qp, psn, hi, m)) { reply(OFL_ERDMA, m); fail_why = m; return false; }
ack.qpn = qp->qp_num;
ack.psn = psn;
memcpy(ack.gid, d.gid.raw, 16);
ack.mtu = d.mtu;
}
ack.backend = (uint32_t)backend;
ack.max_req_bytes = max_req;
ack.max_resp_bytes = max_resp;
ack.max_inflight = inflight;
if (!reply(0, "ok")) { fail_why = "could not send the handshake reply"; return false; }
s->n_conns++;
return true;
}
void Conn::main() {
if (handshake()) {
if (backend == OFL_BACKEND_TCP)
run_tcp();
else
run_rdma();
} else if (s->running.load()) {
fprintf(stderr, "ofl server: %s\n", fail_why.c_str());
}
finished.store(true);
}
void ofl_server::accept_loop() {
while (running.load()) {
pollfd p{lfd, POLLIN, 0};
const int r = poll(&p, 1, 100);
{
std::lock_guard<std::mutex> g(mu); // reap finished connections
conns.erase(std::remove_if(conns.begin(), conns.end(), [](const std::unique_ptr<Conn>& c) { return c->finished.load(); }),
conns.end());
}
if (r <= 0) continue;
sockaddr_in a{};
socklen_t al = sizeof(a);
const int fd = accept4(lfd, reinterpret_cast<sockaddr*>(&a), &al, SOCK_NONBLOCK | SOCK_CLOEXEC);
if (fd < 0) continue;
tune_socket(fd);
std::lock_guard<std::mutex> g(mu);
if (conns.size() >= max_conns) {
fprintf(stderr, "ofl server %s: connection limit %u reached, closing a new one\n", desc.c_str(), max_conns);
::close(fd);
continue;
}
auto c = std::make_unique<Conn>();
c->s = this;
c->id = next_conn++;
c->fd = fd;
Conn* cp = c.get();
conns.push_back(std::move(c));
cp->thr = std::thread([cp] { cp->main(); });
}
}
extern "C" {
ofl_server* ofl_server_create(const ofl_server_config* cfg, ofl_handler_fn fn, void* user, char* err, size_t errlen) {
in_addr a{};
if (!cfg || !parse_ipv4(cfg->bind_ip, &a) || !cfg->port || !cfg->backends) {
copy_err(err, errlen, "ofl_server_create: bind_ip, port and backends are required");
return nullptr;
}
if (a.s_addr == htonl(INADDR_ANY)) { copy_err(err, errlen, "ofl_server_create: refusing to bind 0.0.0.0"); return nullptr; }
auto* s = new ofl_server();
s->bind_ip = cfg->bind_ip;
s->port = cfg->port;
s->backends = cfg->backends;
s->rdma_device = cfg->rdma_device ? cfg->rdma_device : "";
s->rdma_port = cfg->rdma_port > 0 ? cfg->rdma_port : 1;
s->gid_index = cfg->gid_index;
s->service = cfg->service ? cfg->service : "";
s->max_req = cfg->max_req_bytes;
s->max_resp = cfg->max_resp_bytes;
s->max_inflight = cfg->max_inflight ? std::min<uint32_t>(cfg->max_inflight, 64) : 2;
s->max_conns = cfg->max_conns ? cfg->max_conns : 8;
if (cfg->spin_us > 0) s->spin_ns = cfg->spin_us * 1000LL;
s->fn = fn;
s->user = user;
std::string m;
if (s->backends & (1 << OFL_BACKEND_RDMA)) {
if (!open_dev(cfg->bind_ip, s->rdma_device.c_str(), s->rdma_port, s->gid_index, s->dev, m)) {
copy_err(err, errlen, "RDMA: " + m);
delete s;
return nullptr;
}
s->have_dev = true;
}
s->lfd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
int one = 1;
setsockopt(s->lfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
sockaddr_in l{};
l.sin_family = AF_INET;
l.sin_addr = a;
l.sin_port = htons(cfg->port);
if (s->lfd < 0 || ::bind(s->lfd, reinterpret_cast<sockaddr*>(&l), sizeof(l)) != 0 || ::listen(s->lfd, 16) != 0) {
copy_err(err, errlen, fmt("listen %s:%u: %s", cfg->bind_ip, cfg->port, strerror(errno)));
if (s->lfd >= 0) ::close(s->lfd);
s->dev.close();
delete s;
return nullptr;
}
s->desc = fmt("%s:%u service '%s' backends%s%s", cfg->bind_ip, cfg->port, s->service.c_str(),
(s->backends & (1 << OFL_BACKEND_TCP)) ? " tcp" : "", s->have_dev ? "" : "");
if (s->have_dev)
s->desc += fmt(" rdma(%s port %d gid %d mtu %s)", s->dev.name.c_str(), s->dev.port, s->dev.gid_index, mtu_str(s->dev.mtu));
return s;
}
int ofl_server_start(ofl_server* s) {
if (!s || s->running.load()) return OFL_EINVAL;
s->running.store(true);
s->acceptor = std::thread([s] { s->accept_loop(); });
return OFL_OK;
}
void ofl_server_stop(ofl_server* s) {
if (!s) return;
s->running.store(false);
if (s->acceptor.joinable()) s->acceptor.join();
std::vector<std::unique_ptr<Conn>> cs;
{
std::lock_guard<std::mutex> g(s->mu);
cs.swap(s->conns);
}
cs.clear(); // each destructor joins its thread (they see running == false within ~100 ms)
}
void ofl_server_destroy(ofl_server* s) {
if (!s) return;
ofl_server_stop(s);
if (s->lfd >= 0) ::close(s->lfd);
s->dev.close();
delete s;
}
const char* ofl_server_describe(const ofl_server* s) { return s ? s->desc.c_str() : ""; }
void ofl_server_stats(const ofl_server* s, uint64_t out[4]) {
out[0] = s->n_conns.load();
out[1] = s->n_reqs.load();
out[2] = s->n_bytes.load();
out[3] = s->n_errs.load();
}
uint16_t ofl_req_op(const ofl_req* r) { return r->hdr.op; }
uint64_t ofl_req_id(const ofl_req* r) { return r->hdr.req_id; }
uint64_t ofl_req_generation(const ofl_req* r) { return r->hdr.generation; }
int ofl_req_backend(const ofl_req* r) { return r->c->backend; }
uint64_t ofl_req_conn_id(const ofl_req* r) { return r->c->id; }
const void* ofl_req_payload(const ofl_req* r, uint32_t* len) {
if (len) *len = r->payload_len;
return r->payload;
}
void* ofl_req_resp_buf(ofl_req* r, uint32_t* capacity) {
if (capacity) *capacity = r->cap;
return r->resp;
}
void ofl_req_set_len(ofl_req* r, uint32_t len) { r->len = len; }
int ofl_req_ready(ofl_req* r, uint32_t off, uint32_t len) {
if (!r || (uint64_t)off + len > r->cap) return OFL_EINVAL;
if (r->c->backend != OFL_BACKEND_RDMA || !len) return OFL_OK;
std::lock_guard<std::mutex> g(r->c->post_mu);
if (r->err) return r->err;
const int rc = post_range_locked(r, off, len);
if (rc) { r->err = rc; return rc; }
r->sent.emplace_back(off, len);
return OFL_OK;
}
} // extern "C"