Add URMA observability and diagnostics

docs/cn/urma.md

@@ -193,3 +193,36 @@ UrmaTransport 推荐用于**大包和跨节点**高吞吐路径，而 UBRing
   `ContextInitOrDie` 阶段拒绝。
 - 硬件数据通路需要受支持的 UMDK provider 和 `liburma`。
 - 当前实现面向 Linux。
+
+## 可观测性与诊断
+
+通过 `/vars/urma_*` 和 Prometheus 查看进程指标，通过 `/sockets/urma` 查看连接与
+注册 BufferPool 快照。配置 `internal_port` 时使用该端口。
+
+```sh
+curl http://127.0.0.1:8003/sockets/urma
+curl 'http://127.0.0.1:8003/sockets/123?format=json'
+python3 tools/urma_diag.py --server 127.0.0.1:8003 pool
+python3 tools/urma_diag.py --server 127.0.0.1:8003 --json watch --interval 1
+python3 tools/urma_diag.py --server 127.0.0.1:8003 trace 123 --trace-id 456
+```
+
+CLI 使用 Python 3 标准库，支持 `connections`、`pool`、`socket`、`trace` 和 `watch`。
+
+| 指标 | 内容 |
+| --- | --- |
+| 写入与传输 | bulk 写入、数据字节、WR 提交与本地完成、错误 |
+| 流控 | SQ/远端 credit 剩余量、阻塞次数与等待时间 |
+| 建链 | 握手成功、失败、TCP 回退及原因 |
+| 注册池 | 容量、占用、峰值、耗尽与回退分配次数 |
+
+`urma_bulk_min_bytes` 默认 65536，按包含协议头和附件的序列化写入大小分类。
+等待时间按进入 butex 等待时的窗口状态归因，SQ 与远端 credit 等待可重叠。
+
+开启 rpcz 并设置 `urma_trace_sample_every=N`，按 span ID 每 N 条采样一条；
+默认 0 关闭，1 跟踪全部可用 span。采样记录排队、提交、本地完成及资源等待，
+写入 rpcz 注释，并在每连接保留最近 64 条已结束记录。
+
+快照中的 ID 为十进制字符串。列表最多返回 1024 条连接，超出时标记 `truncated`；
+按 SocketId 可直接查询。连接重用时 `generation` 递增、连接计数清零，CLI 按
+同一 generation 计算速率。

docs/en/urma.md

@@ -207,3 +207,40 @@ is tracked as a future enhancement.
   rejected at `ContextInitOrDie`.
 - A hardware data path requires a supported UMDK provider and `liburma`.
 - The current implementation targets Linux.
+
+## Observability and diagnostics
+
+Use `/vars/urma_*` and Prometheus for process metrics, and `/sockets/urma` for
+connection and registered BufferPool snapshots. Use `internal_port` when configured.
+
+```sh
+curl http://127.0.0.1:8003/sockets/urma
+curl 'http://127.0.0.1:8003/sockets/123?format=json'
+python3 tools/urma_diag.py --server 127.0.0.1:8003 pool
+python3 tools/urma_diag.py --server 127.0.0.1:8003 --json watch --interval 1
+python3 tools/urma_diag.py --server 127.0.0.1:8003 trace 123 --trace-id 456
+```
+
+The CLI uses the Python 3 standard library and supports `connections`, `pool`,
+`socket`, `trace`, and `watch`.
+
+| Metrics | Contents |
+| --- | --- |
+| Writes and transfers | Bulk writes, data bytes, submitted and locally completed WRs, errors |
+| Flow control | Remaining SQ/remote credit, blocked attempts, wait time |
+| Connections | Handshake success, failure, TCP fallback and reasons |
+| Registered pool | Capacity, occupancy, peak occupancy, exhaustion, allocation fallback |
+
+`urma_bulk_min_bytes` defaults to 65536 and classifies serialized writes including
+protocol headers and attachments. Wait time is attributed to the window state
+at butex-wait entry; SQ and remote-credit waits can overlap.
+
+Enable rpcz and set `urma_trace_sample_every=N` to sample one in N span IDs.
+The default is 0 (disabled); 1 traces every available span. Samples record queueing,
+submission, local completion and resource waits in rpcz annotations. Each connection
+also retains its latest 64 finished samples.
+
+Snapshot IDs are decimal strings. The inventory returns up to 1024 connections
+and sets `truncated` for larger lists; individual sockets can be queried directly.
+Connection reuse increments `generation` and resets connection counters. The CLI
+calculates rates within the same generation.

src/brpc/builtin/sockets_service.cpp

@@ -17,12 +17,15 @@
 
 
 #include <ostream>
+#include <sstream>
 #include "brpc/closure_guard.h"        // ClosureGuard
 #include "brpc/controller.h"           // Controller
 #include "brpc/builtin/common.h"
 #include "brpc/builtin/sockets_service.h"
 #include "brpc/socket.h"
-
+#if BRPC_WITH_URMA
+#include "brpc/urma/urma_observability.h"
+#endif
 
 namespace brpc {
 
@@ -35,15 +38,56 @@ void SocketsService::default_method(::google::protobuf::RpcController* cntl_base
     cntl->http_response().set_content_type("text/plain");
     butil::IOBufBuilder os;
     const std::string& constraint = cntl->http_request().unresolved_path();
-
-    if (constraint.empty()) {
+
+#if BRPC_WITH_URMA
+    if (constraint == "urma") {
+        cntl->http_response().set_content_type("application/json");
+        os << "{\"schema_version\":1,\"pool\":";
+        urma::DescribePoolJson(os);
+        os << ",\"connections\":[";
+        bool first = true;
+        // Fetch one extra ID to report truncation, rather than silently hide it.
+        const auto ids = urma::ListDiagnosticSockets(1025);
+        for (size_t i = 0; i < ids.size() && i < 1024; ++i) {
+            SocketUniquePtr socket;
+            if (Socket::AddressFailedAsWell(ids[i], &socket) < 0) {
+                continue;
+            }
+            std::ostringstream snapshot;
+            if (!socket->DescribeUrmaJson(snapshot)) {
+                continue;
+            }
+            if (!first) {
+                os << ',';
+            }
+            first = false;
+            os << snapshot.str();
+        }
+        os << "],\"truncated\":" << (ids.size() > 1024 ? "true" : "false") << '}';
+    } else
+#endif
+        if (constraint.empty()) {
         os << "# Use /sockets/<SocketId>\n"
            << butil::describe_resources<Socket>() << '\n';
     } else {
         char* endptr = nullptr;
         SocketId sid = strtoull(constraint.c_str(), &endptr, 10);
         if (*endptr == '\0' || *endptr == '/') {
-            Socket::DebugSocket(os, sid);
+#if BRPC_WITH_URMA
+            const std::string* format = cntl->http_request().uri().GetQuery("format");
+            if (format && *format == "json") {
+                SocketUniquePtr socket;
+                if (Socket::AddressFailedAsWell(sid, &socket) < 0 ||
+                    !socket->DescribeUrmaJson(os)) {
+                    cntl->SetFailed(ENOMETHOD, "Socket has no URMA diagnostics");
+                } else {
+                    cntl->http_response().set_content_type("application/json");
+                }
+            } else
+#endif
+            {
+                Socket::DebugSocket(os, sid);
+            }
         } else {
             cntl->SetFailed(ENOMETHOD, "path=%s is not a SocketId",
                             constraint.c_str());

src/brpc/controller.cpp

@@ -1509,6 +1509,9 @@ void Controller::IssueRPC(int64_t start_realtime_us) {
     }
     Socket::WriteOptions wopt;
     wopt.id_wait = cid;
+#if BRPC_WITH_URMA
+    wopt.urma_trace_span = _span;
+#endif
     wopt.abstime = pabstime;
     wopt.pipelined_count = _pipelined_count;
     wopt.auth_flags = _auth_flags;

src/brpc/policy/baidu_rpc_protocol.cpp

@@ -467,6 +467,9 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
         // Have the risk of unlimited pending responses, in which case, tell
         // users to set max_concurrency.
         Socket::WriteOptions wopt;
+#if BRPC_WITH_URMA
+        wopt.urma_trace_span = span;
+#endif
         wopt.ignore_eovercrowded = true;
         if (INVALID_BTHREAD_ID != response_id) {
             wopt.id_wait = response_id;

src/brpc/socket.cpp

@@ -41,6 +41,9 @@
 #include "brpc/errno.pb.h"
 #include "brpc/event_dispatcher.h"          // RemoveConsumer
 #include "brpc/socket.h"
+#if BRPC_WITH_URMA
+#include "brpc/urma_transport.h"
+#endif
 #include "brpc/describable.h"               // Describable
 #include "brpc/circuit_breaker.h"           // CircuitBreaker
 #include "brpc/input_messenger.h"
@@ -399,6 +402,11 @@ void Socket::WriteRequest::Setup(Socket* s) {
                 return;
             }
         }
+#if BRPC_WITH_URMA
+        if (s->socket_mode() == SOCKET_MODE_URMA) {
+            static_cast<UrmaTransport*>(s->_transport.get())->ObserveWrite(data.size());
+        }
+#endif
         const int64_t before_write =
             s->_unwritten_bytes.fetch_add(data.size(), butil::memory_order_relaxed);
         if (before_write + (int64_t)data.size() >= FLAGS_socket_max_unwritten_bytes) {
@@ -508,6 +516,11 @@ Socket::~Socket() {
 
 void Socket::ReturnSuccessfulWriteRequest(Socket::WriteRequest* p) {
     DCHECK(p->data.empty());
+#if BRPC_WITH_URMA
+    if (socket_mode() == SOCKET_MODE_URMA) {
+        static_cast<UrmaTransport*>(_transport.get())->FinishWrite(&p->data, 0);
+    }
+#endif
     AddOutputMessages(1);
     const bthread_id_t id_wait = p->id_wait;
     bool is_notify_on_success = p->is_notify_on_success();
@@ -528,6 +541,11 @@ void Socket::ReturnFailedWriteRequest(Socket::WriteRequest* p, int error_code,
         CancelUnwrittenBytes(p->data.size());
     }
     p->data.clear();  // data is probably not written.
+#if BRPC_WITH_URMA
+    if (socket_mode() == SOCKET_MODE_URMA) {
+        static_cast<UrmaTransport*>(_transport.get())->FinishWrite(&p->data, error_code);
+    }
+#endif
     const bthread_id_t id_wait = p->id_wait;
     butil::return_object(p);
     if (id_wait != INVALID_BTHREAD_ID) {
@@ -1677,6 +1695,11 @@ int Socket::Write(butil::IOBuf* data, const WriteOptions* options_in) {
     // wait until it points to a valid WriteRequest or nullptr.
     req->next = WriteRequest::UNCONNECTED;
     req->id_wait = opt.id_wait;
+#if BRPC_WITH_URMA
+    if (socket_mode() == SOCKET_MODE_URMA) {
+        static_cast<UrmaTransport*>(_transport.get())->TrackWrite(&req->data, opt.urma_trace_span);
+    }
+#endif
     req->clear_and_set_control_bits(opt.notify_on_success, opt.shutdown_write);
     req->set_pipelined_count_and_user_message(
         opt.pipelined_count, DUMMY_USER_MESSAGE, opt.auth_flags);
@@ -1714,6 +1737,11 @@ int Socket::Write(SocketMessagePtr<>& msg, const WriteOptions* options_in) {
     // wait until it points to a valid WriteRequest or nullptr.
     req->next = WriteRequest::UNCONNECTED;
     req->id_wait = opt.id_wait;
+#if BRPC_WITH_URMA
+    if (socket_mode() == SOCKET_MODE_URMA) {
+        static_cast<UrmaTransport*>(_transport.get())->TrackWrite(&req->data, opt.urma_trace_span);
+    }
+#endif
     req->clear_and_set_control_bits(opt.notify_on_success, opt.shutdown_write);
     req->set_pipelined_count_and_user_message(
         opt.pipelined_count, msg.release(), opt.auth_flags);
@@ -1780,7 +1808,15 @@ int Socket::StartWrite(WriteRequest* req, const WriteOptions& opt) {
         butil::IOBuf* data_arr[1] = { &req->data };
         nw = _conn->CutMessageIntoFileDescriptor(fd(), data_arr, 1);
     } else {
-        nw = _transport->CutFromIOBuf(&req->data);
+#if BRPC_WITH_URMA
+        if (socket_mode() == SOCKET_MODE_URMA) {
+            butil::IOBuf* data_arr[] = {&req->data};
+            nw = static_cast<UrmaTransport*>(_transport.get())->CutWithTrace(data_arr, 1);
+        } else
+#endif
+        {
+            nw = _transport->CutFromIOBuf(&req->data);
+        }
     }
     if (nw < 0) {
         // RTMP may return EOVERCROWDED
@@ -1925,6 +1961,12 @@ ssize_t Socket::DoWrite(WriteRequest* req) {
         if (_conn) {
             return _conn->CutMessageIntoFileDescriptor(fd(), data_list, ndata);
         } else {
+#if BRPC_WITH_URMA
+            if (socket_mode() == SOCKET_MODE_URMA) {
+                return static_cast<UrmaTransport*>(_transport.get())
+                    ->CutWithTrace(data_list, ndata);
+            }
+#endif
             return _transport->CutFromIOBufList(data_list, ndata);
         }
     }
@@ -2327,6 +2369,15 @@ std::ostream& operator<<(std::ostream& os, const ObjectPtr<T>& obj) {
     return os << obj._obj;
 }
 
+#if BRPC_WITH_URMA
+bool Socket::DescribeUrmaJson(std::ostream& os) const {
+    if (socket_mode() != SOCKET_MODE_URMA || !_transport) {
+        return false;
+    }
+    return static_cast<UrmaTransport*>(_transport.get())->DescribeJson(os);
+}
+#endif
+
 void Socket::DebugSocket(std::ostream& os, SocketId id) {
     SocketUniquePtr ptr;
     int ret = AddressFailedAsWell(id, &ptr);

src/brpc/socket.h

@@ -46,6 +46,7 @@
 #include "brpc/input_messenger_processor.h"  // InputMessengerProcessor
 
 namespace brpc {
+class Span;
 namespace policy {
 class ConsistentHashingLoadBalancer;
 class RtmpContext;
@@ -393,6 +394,10 @@ friend class TransportFactory;
     // - Write once when uncontended(most cases).
     // - Wait-free when contended.
     struct WriteOptions {
+#if BRPC_WITH_URMA
+        // Explicit RPC span; never infer this from the KeepWrite bthread's TLS.
+        std::shared_ptr<Span> urma_trace_span;
+#endif
         // `id_wait' is signalled when this Socket is SetFailed or data is written
         // successfully with `notify_on_success=true'. To disable the signal, set
         // this field to INVALID_BTHREAD_ID. `on_reset' of `id_wait' is NOT called
@@ -632,6 +637,10 @@ friend class TransportFactory;
 
     // The socket mode decided at creation.
     SocketMode socket_mode() const { return _socket_mode; }
+#if BRPC_WITH_URMA
+    // Called with an addressed Socket, keeping its transport alive.
+    bool DescribeUrmaJson(std::ostream& os) const;
+#endif
 
     // Get an UNUSED socket connecting to the same place as this socket
     // from the SocketPool of this socket.

src/brpc/urma/mock_urma.cpp

@@ -560,17 +560,19 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
         send_cr.status = URMA_CR_SUCCESS;
         send_cr.user_ctx = current->user_ctx;
         send_cr.flag.bs.s_r = 0;
-        if (current->flag.bs.complete_enable) {
-            PushCompletion(jetty->jetty_cfg.jfs_cfg.jfc,
-                           local_state, send_cr);
-        }
+        auto complete_send = [&] {
+            if (current->flag.bs.complete_enable) {
+                PushCompletion(jetty->jetty_cfg.jfs_cfg.jfc, local_state, send_cr);
+            }
+        };
 
         PendingRecv recv{};
         urma_jfc_t* remote_jfc = nullptr;
         {
             std::unique_lock<std::shared_mutex> lock(g_rw_mutex);
             auto remote_it = jetty_id_map.find(current->tjetty->id.id);
             if (remote_it == jetty_id_map.end()) {
+                complete_send();
                 continue;
             }
             urma_jfr_t* remote_jfr =
@@ -579,6 +581,7 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
             auto jfc_it = jfr_jfc_map.find(remote_jfr);
             if (recv_it == jfr_recv_map.end() || recv_it->second.empty() ||
                 jfc_it == jfr_jfc_map.end()) {
+                complete_send();
                 continue;
             }
             recv = recv_it->second.front();
@@ -596,6 +599,11 @@ urma_status_t urma_post_jetty_send_wr(urma_jetty_t *jetty, urma_jfs_wr_t *wr,
             copied += n;
         }
 
+        // A local completion releases the sender's IOBuf references. Publish
+        // it only after the mock has consumed every source SGE, matching the
+        // lifetime guarantee made by real providers.
+        complete_send();
+
         JfcState* remote_state = nullptr;
         {
             std::shared_lock<std::shared_mutex> lock(g_rw_mutex);