# rs-mq

**Repository Path**: fastkit/rs-mq

## Basic Information

- **Project Name**: rs-mq
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-10-15
- **Last Updated**: 2026-11-06

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# rs-mq

基于 Tokio 的单中心内存 MQ。Broker 和 SDK 使用同一份代码同步部署，协议版本保持 `1`。

## 链路上下文

rs-mq 不依赖 OpenTelemetry，也不解析、生成或覆盖链路字段。发布方可将 `traceparent`、`tracestate` 写入任务 metadata，rs-mq 会像传输其他 metadata 一样将其原样交给 Worker；上下文的注入、提取和导出由业务应用负责。

## 发布与消费

Rust 应用先创建一个只保存 Broker 地址的顶层 `Client`，再从中取得发布或订阅角色。发布连接由同一个 `Publisher` 复用，每次发布指定目标服务；每次 `subscriber()` 都创建一条独立的单服务消费连接：

```rust
use bytes::Bytes;
use rs_mq::client::{Client, Metadata};

let client = Client::start("rs-mq", 9000)?;
let publisher = client.publisher();
let task = publisher
    .publish("phc-det", Bytes::from_static(b"image"), Metadata::new())
    .await?;
let result = task.result().await?;

let subscriber = client.subscriber("phc-det").await?;
let task = subscriber.next_task().await?;
```

发布目标属于每条 `PUBLISH`，消费目标属于 `Subscriber` 连接。一个顶层 Client 可以创建任意多个订阅者，相同 service 的多个订阅者也各自占用一个独立消费槽位。Broker、Rust SDK 和 Python wheel 必须使用同一份协议代码。

Python 使用相同的角色模型：

```python
from rs_mq import RsMQClient

client = RsMQClient(host="rs-mq", port=9000)
publisher = client.publisher().client(timeout=5)
task = publisher.publish("phc-det", b"image", {"content-type": "application/msgpack"})
result = task.result(timeout=30)

subscriber = client.subscriber("phc-det", timeout=5)
task = subscriber.next_task()
print(task.task_id, task.attempt_id)
```

## Broker 资源配置

环境变量在启动监听前读取并校验，非法配置导致启动失败。

| 环境变量 | 默认值 | 含义 |
| --- | --- | --- |
| `RSMQ_MAX_MEMORY_BYTES` | `4g` | Broker 管理的协议帧数据预算 |
| `RSMQ_MAX_TASKS` | `100000` | 全局排队与在途任务总数 |
| `RSMQ_MAX_CONNECTIONS` | `1024` | 包含握手中的连接总数 |

内存值会去除首尾空白、统一转小写，支持正整数加 `b`、`k`、`m`、`g`；不带单位表示字节。按 1024 换算：`2g`/`2G` 为 2 GiB，`100m`/`100M` 为 100 MiB。拒绝零、符号、小数、内部空白、未知单位及溢出。任务数和连接数仅接受正十进制整数。

```powershell
$env:RSMQ_MAX_MEMORY_BYTES = "2g"
$env:RSMQ_MAX_TASKS = "100000"
$env:RSMQ_MAX_CONNECTIONS = "1024"
$env:RSMQ_BIND = "0.0.0.0:9000"
rs-mq-server
```

Compose 示例（小预算开发配置）：

```yaml
services:
  rs-mq:
    image: rs-mq-server:latest
    environment:
      RSMQ_MAX_MEMORY_BYTES: "100m"
      RSMQ_MAX_TASKS: "10000"
      RSMQ_MAX_CONNECTIONS: "128"
    ports:
      - "9000:9000"
```

75% 数据预算供发布输入使用，25% 留给结果、握手和其他网络帧。单帧超过对应预算立即拒绝，仍受现有 256 MiB 协议帧大小限制。完整帧（含长度前缀）在接收前计费，共享负载不重复计费；转发生成的协议头另行计费。任务派发、重投保留输入预算，慢写入完成前也保留相应凭证。

这个预算不是进程 RSS 上限：解码对象、metadata 副本、连接、运行时及分配器都有额外开销。容器内存限制必须高于数据预算并按实际负载留余量；不要把 `4g` 数据预算等同于 `4g` 容器内存。

## 过载处理

- 任务数满：拒绝本次发布，返回 `BrokerBusyError`，保持连接。
- 发布字节预算不足：读取固定请求头并尝试写回关联的 `broker_busy` 错误，随后关闭连接。网络故障或控制帧预算也已耗尽时，调用方可能只收到连接错误。
- 该连接已有任务按断线规则处理：尚未派发的任务删除，已派发任务废弃；SDK 自动重连但不会自动重试发布。
- 结果或其他帧预算不足：关闭连接，worker 任务按已有规则重投。过大的结果可能持续触发重投，应调整预算或缩小结果。
- 连接数满：新 socket 直接关闭，已连接客户端继续运行。

```python
from rs_mq import BrokerBusyError

try:
    task = publisher.publish("phc-det", b"payload", ttl=10)
except BrokerBusyError:
    # 由业务方决定拒绝上游请求或稍后重新发布。
    pass
```

SDK 不在本地缓存或自动重试任务；`result(timeout=...)` 超时仅结束本次等待，不取消远端任务。

## 运行状态与关键日志

`RSMQ_STATUS_BIND` 默认 `127.0.0.1:9001`，提供只读 `GET /status`，返回 UTF-8 纯文本；设为 `off` 关闭。配置或端口绑定失败会阻止服务启动。接口没有认证，应仅在可信运维网络使用。

```text
curl http://127.0.0.1:9001/status
```

容器内可设置 `RSMQ_STATUS_BIND: "0.0.0.0:9001"`，并通过 Compose 的 `"127.0.0.1:9001:9001"` 端口映射让宿主机本地查询。

接口使用分组中文文本，内存同时显示易读单位和原始字节数，例如：

```text
rs-mq 运行状态

【运行与连接】
运行时长：120 秒
当前连接数：2
已用连接槽位：2
连接上限：1024
连接拒绝总数：0

【任务与内存】
当前任务数：1
任务上限：100000
输入数据内存：1.25 MiB / 3.00 GiB（1310720 / 3221225472 字节）
其他数据内存：0 B / 1.00 GiB（0 / 1073741824 字节）

【服务状态】
服务："paper-det"
  绑定连接数：2
  排队任务数：1
  排队数据：1.25 MiB（1310720 字节）
  空闲 Worker：1
  执行中任务数：0
```

全局状态包含运行时间、活动连接、包含握手的连接槽位及上限、任务数量及上限、输入和其他帧内存预算。累计统计包含发布、成功、失败、过期、重投、任务或内存拒绝、发布者断线废弃和日志丢失次数。

每个活动 service 输出绑定连接数、排队任务数、排队数据、空闲 Worker 和执行中任务数。绑定连接数只包含绑定该 service 的 Subscriber，不包含不绑定服务的 Publisher 连接。排队数据只计算 payload 与 UTF-8 metadata 原文，不含协议头或对象开销，不能与全局预算使用量直接相加。service 名称转义并最多显示 128 个字符；服务回收后该分组消失，全局累计计数保留至 broker 重启。

成功/失败计数表示 broker 接受了 worker 的结果，包括发布者已离线的情况，不代表发布者已经收到结果。废弃计数包括排队删除和在途标记废弃，因此可能与之后的完成计数重叠。`other_budget_rejections_total` 是资源预占失败次数，不是去重任务数。

状态由调度器生成，最多缓存 1 秒，仅遍历服务；原子预算值是读取时的近似快照。最多同时处理 8 个 HTTP 请求，请求头最多 1024 字节、处理限时 2 秒，不支持 keep-alive。调度通道满时返回 503，超过并发、头大小或时限时关闭状态连接。状态 HTTP 接口不影响业务 TCP；本次 `PUBLISH` 帧已增加目标 service，broker 与 Rust SDK 必须同步部署，Python Worker 接口保持不变。

关键日志写入 stderr，可使用 `docker compose logs` 查看：监听地址、握手成功/失败、连接关闭原因、过载拒绝、发送失败、worker 断线重投（含连接和任务 ID）、业务失败以及进程退出。图片、metadata 和模型错误正文不写入日志；正常发布、成功完成及逐条 TTL 过期不打印。

运行日志通过容量 256 的后台队列写入，队列满时丢弃并累计到状态页的“日志丢失数”，不会等待 stderr。日志不是审计存储，进程退出时未写出的后台日志可能丢失；进程退出行由入口直接输出。

### 发送背压与压测诊断

Broker 每连接保留 32 帧写通道，并额外提供最多 256 帧的 FIFO 待发送缓冲（不含正在写入的帧）。短时写通道满不再立即断开：写端释放容量后通知调度器续发，调度器不等待 socket。所有缓冲帧继续持有原有内存预算；额外缓冲满、预算不足或连接关闭时仍可断开慢连接，不提供无限缓存。

全局调度事件通道满时，每连接最多保留一个待提交事件并暂停领取后续网络帧，等待可被关闭信号打断；心跳继续发送，本地等待时间不计入未响应超时。超过转发头部总预算的 metadata 在发布确认前拒绝。

`/status` 的“发送缓冲”分组显示额外待发送帧总数、历史单连接峰值、256 帧上限和溢出次数。发送失败日志分别使用 `pending_write_limit`、`write_queue_closed`、`memory_budget` 和 `frame_encoding`，便于区分持续拥塞与编码错误。

现有 `rs-mq-bench` 的 CSV 保留前四列，并追加 `connection_lost,unavailable,broker_busy,timeouts,other,pull_errors,completion_errors`。全部计数仅统计正式测量窗口内返回的结果；`errors` 是前五类发布/结果错误的总和，消费者错误单列展示，避免同一次故障重复合计。与旧版相比，`Unavailable` 现在计入错误，不再把预热阶段错误混入结果。压测不逐条打印错误。

## 自动重连抖动

首次连接仍立即发起。失败后的重试采用每客户端独立的随机抖动：初始等待 100–200ms，随后指数增大，最终等待区间为 2.5–5 秒；成功握手后重置退避阶段。随机性来自标准库 `RandomState`，不增加依赖。关闭或排空信号可以打断连接、握手与退避等待。

## Worker 排空退出

Python 和 Rust SDK 均提供两个阶段：

- `start_draining()`：立即发起不可逆排空，不等待当前任务，适合控制线程调用。
- `drain()`：发起排空并等待完成后关闭。Python 顶层 Client 的阻塞调用会释放 GIL；Rust 使用 `client.drain().await`。顶层 Client 统一处理已创建的 Publisher 和 Subscriber。
- Python `client.draining` / Rust `client.is_draining()` 查询是否已发起排空。现有 `close()` 仍立即强制关闭，可中断排空等待。

排空后拒绝新发布和新领取，Python 抛出 `ClientDrainingError`（继承 `RsMQError`），Rust 返回 `Error::Draining`。已经发出的 PULL 由 broker 撤销；与排空发生竞态的最后一项 ASSIGN 仍会返回给原先等待的 `next_task()`，必须照常完成。不要在取得任务后仅因为 `draining=True` 就丢弃它。

SDK 等待 broker 的 `DRAIN_ACK`、当前任务的 `COMPLETION_ACK`，以及同连接上已经发布任务的最终结果，期间保持心跳。排空完成后停止自动重连。无在线任务时，连接/握手/退避阶段的排空可直接完成。排空期间在线连接中断返回连接错误，不会重连后宣称成功；重复 `drain()` 返回同一个最终结果。协议增加 `DRAIN`/`DRAIN_ACK`，版本仍为 `1`，broker 和 wheel 必须同步更新。

排空没有超时，也不会中断 GPU 推理；永久阻塞的推理或没有最终结果的已发布任务会阻止排空完成。不要在尚未完成当前任务的执行线程里直接调用阻塞式 `drain()`；先调用 `start_draining()`，完成当前任务后再 `drain()`。发布者与订阅者可以通过各自角色句柄独立排空，也可以由顶层 Client 统一排空。

Linux / Docker worker 示例（`infer` 替换为实际推理函数）：

```python
import signal
import threading
from rs_mq import RsMQClient, ClientDrainingError

# 必须在线程和 SDK runtime 启动前设置信号掩码。
stop_signals = {signal.SIGTERM, signal.SIGINT}
signal.pthread_sigmask(signal.SIG_BLOCK, stop_signals)
client = RsMQClient(host="rs-mq", port=9000)

def request_stop():
    signal.sigwait(stop_signals)
    client.start_draining()

threading.Thread(target=request_stop, daemon=True).start()
subscriber = client.subscriber("signature")

try:
    while True:
        try:
            task = subscriber.next_task()
        except ClientDrainingError:
            break
        try:
            payload = infer(task.payload)
        except Exception as exc:
            task.fail(str(exc))
        else:
            task.success(payload)
    client.drain()
finally:
    client.close()
```

使用专门的信号等待线程，避免 Python 主线程阻塞在原生 `next_task()` 时无法及时执行普通 Python 信号处理器。此示例的信号 API 仅适用于 Unix；其他平台可由应用的控制线程调用 `start_draining()`。容器应直接向 Python worker 进程发送 SIGTERM，并将 `stop_grace_period` 配置为足够完成当前推理的时长；宽限期后 SIGKILL 仍会导致重投。