commit e59ecfc0e342c3a412096027dd7203bbfc1ffe84 Author: chakcy <947105045@qq.com> Date: Mon Aug 31 17:18:03 2026 +0800 modify: rust 插件 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7e5fa64 --- /dev/null +++ b/.gitignore @@ -0,0 +1,85 @@ +# ============================================ +# 1. Cargo 构建产物 +# ============================================ +# 编译生成的目标文件目录，必须忽略[reference:0] +**/target/ +# Rust 编译器发生内部错误（ICE）时生成的报告文件 +rustc-ice-*.txt + +# ============================================ +# 2. 依赖锁文件 (Cargo.lock) +# ============================================ +# 对于最终可执行程序（如 CLI 工具、应用），建议**提交** Cargo.lock 以保证构建确定性[reference:3][reference:4]。 +# 对于库项目（library），建议**忽略** Cargo.lock，因为库的依赖由使用它的项目决定[reference:5][reference:6]。 +# 这里默认忽略，请根据你的项目类型决定是否保留或注释掉下面这行[reference:7][reference:8]。 +**/Cargo.lock + +# ============================================ +# 3. 开发工具与临时文件 +# ============================================ +# rustfmt 生成的备份文件[reference:9][reference:10] +**/*.rs.bk +# 突变测试（mutagen/mutants）的输出目录[reference:11] +**/mutants.out*/ +# 各种临时文件 +*.swp +*.swo +*~ +\#*# +.#* + +# ============================================ +# 4. 集成开发环境（IDE）配置 +# ============================================ +# Visual Studio Code[reference:13] +.vscode/ +# JetBrains 系列 (IntelliJ, RustRover, CLion)[reference:15][reference:16] +.idea/ +*.iml +# 其他编辑器 +.helix/ +.zed/ + +# ============================================ +# 5. 操作系统文件 +# ============================================ +# macOS +.DS_Store +.AppleDouble +.LSOverride +._* +__MACOSX/ +# Windows +Thumbs.db +ehthumbs.db +Desktop.ini +$RECYCLE.BIN/ + +# ============================================ +# 6. 环境变量与配置文件 +# ============================================ +# 包含本地敏感信息的文件，务必忽略 +.env +.env.* +# 但可以保留示例文件 +!.env.example +!.env.sample +# 其他工具产生的环境配置 +.direnv/ +.envrc + +# ============================================ +# 7. 其他构建与测试产物 +# ============================================ +# Windows 调试符号文件[reference:21] +*.pdb +# 代码覆盖率报告[reference:23] +*.profraw +*.profdata +tarpaulin-report.html +lcov.info +coverage/ +# 软件物料清单（SBOM）文件 +sbom.cyclonedx* + +__pycache__/ diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..b6c71e0 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,44 @@ +[package] +name = "chongming-worker-rs" +version = "0.1.0" +edition = "2021" +description = "chongming 进程插件 Rust SDK（语言无关线协议的 Rust 实现，见 docs/plugin-wire-protocol.md）" +license = "MIT" + +[dependencies] +zenoh = "1" +tokio = { version = "1", features = ["full"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +toml = "0.8" +clap = { version = "4", features = ["derive"] } +tracing = "0.1" +tracing-subscriber = "0.3" +async-stream = "0.3" +futures-util = "0.3" +sha1 = "0.10" +# 数据库封装（SQLite / PostgreSQL 异步连接池，对齐 Python chongming_database）。 +# 注：默认不启用 TLS 后端；PostgreSQL 走 SSL 时需追加 tls-rustls-ring 等 feature。 +sqlx = { version = "0.8", default-features = false, features = [ + "runtime-tokio", + "sqlite", + "postgres", + "any", +], optional = true } +redis = { version = "0.27", default-features = false, features = ["tokio-comp", "connection-manager"], optional = true } + +[features] +default = [] +# 启用共享缓存（Redis 后端），用于多实例全局收敛配置与副本数 +redis-cache = ["redis"] +# 启用数据库封装（SQLite/PostgreSQL，sqlx）。开启后脚本打包用: +# bash build.sh dist "database" +database = ["sqlx"] + +[[bin]] +name = "gen_models" +path = "src/bin/gen_models.rs" + +[[example]] +name = "demo" +path = "examples/demo.rs" diff --git a/README.md b/README.md new file mode 100644 index 0000000..c5d6be1 --- /dev/null +++ b/README.md @@ -0,0 +1,126 @@ +# chongming-worker-rs —— 进程插件 Rust SDK + +chongming 进程插件线协议（[docs/plugin-wire-protocol.md](../../docs/plugin-wire-protocol.md)）的 **Rust 参考实现** +（多语言演进规划阶段 1，见 [docs/multi-language-plugin-plan.md](../../docs/multi-language-plugin-plan.md)）。 + +与 Python SDK（`utils/worker/`）行为等价：同一 `` 下可混布部署，网关负载均衡在 +两实现间透明分发。**已通过协议级 15/15 与混布对拍 6/6 端到端验证。** + +## 模块结构 + +| 模块 | 职责 | 对位（Python SDK） | +| ---- | ---- | ---- | +| `config.rs` | 全局 config.toml（`[zenoh].json5_path`）+ 插件 `[workers]` 节展平 + mqtt_topics | `process.py: load_plugin_config` | +| `descriptor.rs` | 端点 TOML 解析 → 路由描述符（defaults 合并 / subject / path / route_id） | `route_config.py` + `routes_publisher.py` | +| `envelope.rs` | 请求信封解析（meta.method / user_info / data，非信封载荷按 GET） | `runtime.py` | +| `error.rs` | 非流式错误格式 / 流式 error 收尾块 | `runtime.py` | +| `publisher.rs` | 描述符 PUT / DELETE / 10s 周期重发布 | `routes_publisher.py` | +| `serve.rs` | worker_serve 主流程：liveliness token + Queryable + 信号处理 | `process.py: worker_serve` | + +## 插件开发三步 + +```rust +use chongming_worker_rs::prelude::*; + +let mut registry = HandlerRegistry::new(); +// 普通端点 +registry.register("rsdemo.hello", "GET", |req: Request| async move { + Ok(Reply::Json(json!({"greeting": "Hello!"}))) +}); +// 流式端点（SSE 逐块） +registry.register("rsdemo.stream", "POST", |req: Request| async move { + let stream = async_stream::stream! { + for i in 0..5 { yield Ok(json!({"index": i})); } + }; + Ok(Reply::Stream(Box::pin(stream))) +}); + +serve::serve("rsdemo", Arc::new(registry), ServeOptions::default()).await?; +``` + +端点 TOML 定义与 Python 插件共用同一套格式（`config/workers//*.toml`）， +描述符生成规则逐字段对位（`docs/plugin-wire-protocol.md` §2.2）。 + +## 模型生成（TOML 是唯一事实来源） + +请求/响应模型由端点 TOML 的 `params` / `response_model` 声明**自动生成** +（对齐 Python `gen_models`，命名约定一致：`hello` + `GET` → `HelloGetRequest`）： + +```bash +# 项目根目录 +cargo run --bin gen_models -- --types rsdemo,demo --out utils/worker-rs/generated_models.rs +``` + +生成文件（crate 根 `generated_models.rs`，勿手动修改）每个插件类型一个 mod， +结构体带 serde derive，字段类型映射对齐网关（str/int/float/bool/any/list/dict）。 + +**一致性保证（比 Python 更强）**：生成文件头部写入 `FINGERPRINT`（端点 TOML 内容 hash）。 +插件启动时（`check_models_fingerprint`）重算指纹比对——**TOML 变更而未重新生成模型 +即告警并提示重跑 gen_models**，杜绝模型与配置漂移。 + +```rust +// 插件内使用生成模型（强类型 handler） +let req: generated_models::rsdemo::HelloGetRequest = serde_json::from_value(request.data)?; +let resp = generated_models::rsdemo::HelloGetResponse { + greeting: format!("Hello, {}!", req.name), + worker_id: Some(request.worker_id), + timestamp: now_ts(), +}; +``` + +类型映射规则（对齐网关 `route_generator` 的 TYPE_MAP）： + +| TOML 声明 | Rust 类型 | 说明 | +| ---- | ---- | ---- | +| `str` | `String` | | +| `int` | `i64` | | +| `float` | `f64` | | +| `bool` | `bool` | | +| `any` | `serde_json::Value` | | +| `list[X]` / `List[X]` | `Vec` | 递归映射 | +| `dict[K,V]` / `Dict[K,V]` | `HashMap` | 递归映射 | +| 其他 | `serde_json::Value` | 宽松兜底 | + +语义对齐：`params` 全部必填；`response_model` 中 `required = false` 的字段生成 +`Option` + `skip_serializing_if`（空值不输出）。 + +## 本地验证 + +前置：本地 Zenoh Router（替代 zenohd，`mode=router` 监听 7447）： + +```bash +.venv/bin/python utils/worker-rs/scripts/zenoh_router.py & +``` + +协议级端到端验证（路由注册 / RPC / 流式 / 错误 / 摘除 / 自愈，15 项断言）： + +```bash +cargo build --examples +.venv/bin/python utils/worker-rs/scripts/verify_e2e.py --type rsdemo +``` + +Python + Rust 同 type 混布对拍（描述符 hash 一致 / 双副本发现 / RPC 等价 / 互不干扰）： + +```bash +.venv/bin/python utils/worker-rs/scripts/verify_mixed.py +``` + +参考插件（`examples/demo.rs`）支持两个 type： + +```bash +# 项目根目录（config.toml 相对路径） +utils/worker-rs/target/debug/examples/demo --type rsdemo # 独立演示插件 +utils/worker-rs/target/debug/examples/demo --type demo # 与 Python demo 混布对拍 +``` + +**独立插件形态**（插件开发者实际使用的方式）：`workers/rsdemo/` 是一个完整独立的 +插件 crate（path 依赖本 SDK），与 Python 示例 `workers/demo/` 对位，含端点实现、 +生成模型与说明文档——建议新插件以此为模板。 + +## 与 Python SDK 的差异（有意为之） + +| 差异 | 说明 | +| ---- | ---- | +| 契约校验 | Python 在 worker 进程内做 TOML↔注册表↔模型三方校验（硬失败退出）；Rust 侧降级为**发布前自检**：无 handler 的端点告警并跳过，不发布残缺路由 | +| 模型生成 | Rust 有 gen_models 等价物（`--bin gen_models`）：从端点 TOML 生成强类型结构体 + **启动时指纹校验**（TOML 变更未重生成即告警，比 Python 只查模型存在性更强） | +| handler 注入 | worker_id 经 `Request.worker_id` 注入（Python 为函数参数）；user_info 经 `Request.user_info` | diff --git a/examples/demo.rs b/examples/demo.rs new file mode 100644 index 0000000..7cdfca4 --- /dev/null +++ b/examples/demo.rs @@ -0,0 +1,194 @@ +//! 参考插件（多语言演进规划阶段 1 的参考实现） +//! +//! 演示 Rust SDK 的完整能力，按 `--type` 对位不同插件： +//! - `--type rsdemo`（默认）：独立演示插件，覆盖普通/流式/错误/user_info 全形态 +//! - `--type demo`：与 Python demo 插件同 type 混布对拍（storage 未配置时降级 +//! 返回与 Python 一致的错误语义，验证负载均衡在两实现间透明分发） +//! +//! 模型一致性：请求/响应模型由端点 TOML 生成（`generated_models.rs`）， +//! 启动时做指纹校验——TOML 变更而未重新生成模型会告警。 +//! +//! 用法（项目根目录）: +//! cargo run --example demo -- --type rsdemo +//! cargo run --example demo -- --type demo + +// 模型由 gen_models 生成（crate 根，勿手动修改）： +// cargo run --bin gen_models -- --types rsdemo,demo --out generated_models.rs +#[path = "../generated_models.rs"] +mod generated_models; + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use chongming_worker_rs::gen; +use chongming_worker_rs::prelude::*; +use chongming_worker_rs::{HandlerRegistry, ServeOptions}; +use clap::Parser; +use serde_json::{json, Value}; +use tracing_subscriber::EnvFilter; + +/// 参考插件：Rust 实现 chongming 进程插件（对位 Python SDK 行为） +#[derive(Parser, Debug)] +#[command(name = "chongming-worker-rs-demo", about = "Rust 进程插件参考实现")] +struct Args { + /// 插件类型（rsdemo / demo） + #[arg(long, default_value = "rsdemo")] + r#type: String, + /// 副本数（覆盖插件配置中的 count） + #[arg(long)] + count: Option, + /// 全局配置文件路径（读取 [zenoh] 节） + #[arg(long, default_value = "config.toml")] + config: String, + /// 实例 id（多实例全局分摊用；缺省 {hostname}:{pid}） + #[arg(long)] + instance: Option, +} + +fn now_ts() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +fn deser(data: &Value) -> Result { + serde_json::from_value(data.clone()).map_err(|e| format!("参数解析失败: {e}")) +} + +// ==================== rsdemo 端点（使用 rsdemo 生成模型） ==================== + +/// rsdemo.hello：普通 GET 端点（对位 Python demo.hello，强类型请求/响应） +async fn hello_rsdemo(request: Request) -> Result { + let req: generated_models::rsdemo::HelloGetRequest = deser(&request.data)?; + let resp = generated_models::rsdemo::HelloGetResponse { + greeting: format!("Hello, {}!", req.name), + worker_id: Some(request.worker_id), + timestamp: now_ts(), + }; + Ok(Reply::Json(serde_json::to_value(&resp).map_err(|e| e.to_string())?)) +} + +/// rsdemo.echo：POST 端点，演示 user_info 透传（网关认证后注入 meta.user_info） +async fn echo(request: Request) -> Result { + let req: generated_models::rsdemo::EchoPostRequest = deser(&request.data)?; + let resp = generated_models::rsdemo::EchoPostResponse { + echo: req.msg, + user_info: request.user_info, + worker_id: Some(request.worker_id), + timestamp: now_ts(), + }; + Ok(Reply::Json(serde_json::to_value(&resp).map_err(|e| e.to_string())?)) +} + +/// rsdemo.stream：流式端点（SSE）。tokens 控制块数；tokens > 50 时中途抛错 +/// 演示 `{"error": ...}` 收尾块（网关收到后终止流）。 +async fn stream_chat(request: Request) -> Result { + let req: generated_models::rsdemo::StreamPostRequest = deser(&request.data)?; + let tokens = req.tokens.clamp(0, 100) as usize; + let stream = async_stream::stream! { + for i in 0..tokens { + if i == 25 && tokens > 50 { + yield Err("模拟流中异常（第 26 块）".to_string()); + return; + } + yield Ok(json!({"index": i, "chunk": format!("chunk-{i}")})); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + }; + Ok(Reply::Stream(Box::pin(stream))) +} + +/// rsdemo.error：普通端点，恒抛业务异常，演示非流式错误回复格式 +/// （{"result":0.0,"operation":"error","error":...,"worker":...}） +async fn always_error(_request: Request) -> Result { + Err("模拟业务异常（rsdemo.error）".to_string()) +} + +// ==================== demo 端点（使用 demo 生成模型，与 Python 混布对拍） ==================== + +/// demo.hello：与 Python workers/demo 对位（强类型模型来自 demo 生成模型） +async fn hello_demo(request: Request) -> Result { + let req: generated_models::demo::HelloGetRequest = deser(&request.data)?; + let resp = generated_models::demo::HelloGetResponse { + greeting: format!("Hello, {}!", req.name), + worker_id: Some(request.worker_id), + timestamp: now_ts(), + }; + Ok(Reply::Json(serde_json::to_value(&resp).map_err(|e| e.to_string())?)) +} + +/// storage 未配置降级（演示：Rust 插件无需 Minio 依赖即可对位 demo type） +async fn storage_not_configured(_request: Request) -> Result { + Err("storage 未配置（本 Rust 参考插件未实现 S3/Minio 客户端）".to_string()) +} + +fn main() -> Result<(), String> { + let args = Args::parse(); + + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("info")), + ) + .init(); + + // 模型一致性校验：端点 TOML 是唯一事实来源，TOML 变更后须重新运行 gen_models + check_models_fingerprint(&args.r#type); + + let mut registry = HandlerRegistry::new(); + + // ---- rsdemo 独立演示插件 ---- + registry.register("rsdemo.hello", "GET", hello_rsdemo); + registry.register("rsdemo.echo", "POST", echo); + registry.register("rsdemo.stream", "POST", stream_chat); + registry.register("rsdemo.error", "GET", always_error); + + // ---- demo 混布对拍：与 Python demo 插件同 type（同配置、同描述符）---- + registry.register("demo.hello", "GET", hello_demo); + registry.register("demo.storage.presign_upload", "POST", storage_not_configured); + registry.register("demo.storage.presign_download", "GET", storage_not_configured); + registry.register("demo.storage.upload_bytes", "POST", storage_not_configured); + registry.register("demo.storage.download_bytes", "GET", storage_not_configured); + + let opts = ServeOptions { + config_path: args.config, + config_dir: "config/workers".into(), + count: args.count, + instance: args.instance, + }; + + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| format!("创建 tokio runtime 失败: {e}"))? + .block_on(chongming_worker_rs::serve::serve( + &args.r#type, + Arc::new(registry), + opts, + )) +} + +/// TOML ↔ 生成模型一致性校验：指纹不匹配时告警并提示重新生成 +fn check_models_fingerprint(worker_type: &str) { + let fp = match gen::fingerprint(worker_type, "config/workers") { + Ok(f) => f, + Err(e) => { + tracing::warn!("[{worker_type}] 计算模型指纹失败（跳过一致性校验）: {e}"); + return; + } + }; + let expected = match worker_type { + "rsdemo" => generated_models::rsdemo::FINGERPRINT, + "demo" => generated_models::demo::FINGERPRINT, + _ => "", + }; + if fp != expected { + tracing::warn!( + "[{worker_type}] 模型指纹不匹配：端点 TOML 已变更但未重新生成模型。\ + 运行: cargo run --bin gen_models -- --types rsdemo,demo --out generated_models.rs" + ); + } else { + tracing::info!("[{worker_type}] 模型一致性校验通过（TOML ↔ 生成模型）"); + } +} diff --git a/generated_models.rs b/generated_models.rs new file mode 100644 index 0000000..f311c54 --- /dev/null +++ b/generated_models.rs @@ -0,0 +1,98 @@ +// 自动生成 —— 勿手动修改（chongming-worker-rs gen_models） +// 唯一事实来源: config/workers//*.toml 端点定义 +// 重新生成: cargo run --bin gen_models -- --types --out + +pub mod demo { + #![allow(dead_code)] + pub const FINGERPRINT: &str = "71c0a54c04bf6a8692166c71e2b6c0109c9c011e"; + #[derive(Debug, Clone, serde::Deserialize)] + pub struct DownloadBytesGetRequest { + pub object_key: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct DownloadBytesGetResponse { + pub data_b64: String, + pub object_key: String, + pub size: i64, + pub timestamp: f64, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct HelloGetRequest { + pub name: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct HelloGetResponse { + pub greeting: String, + pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_id: Option, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct PresignDownloadGetRequest { + pub object_key: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct PresignDownloadGetResponse { + pub expires: i64, + pub method: String, + pub object_key: String, + pub presigned_url: String, + pub timestamp: f64, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct PresignUploadPostRequest { + pub filename: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct PresignUploadPostResponse { + pub expires: i64, + pub method: String, + pub object_key: String, + pub presigned_url: String, + pub timestamp: f64, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct UploadBytesPostRequest { + pub object_key: String, + pub data_b64: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct UploadBytesPostResponse { + pub object_key: String, + pub size: i64, + pub timestamp: f64, + } +} +pub mod rsdemo { + #![allow(dead_code)] + pub const FINGERPRINT: &str = "d7249b91188e422a0fd3966cfc243d051cbb53c0"; + #[derive(Debug, Clone, serde::Deserialize)] + pub struct EchoPostRequest { + pub msg: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct EchoPostResponse { + pub echo: String, + pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub user_info: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_id: Option, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct HelloGetRequest { + pub name: String, + } + #[derive(Debug, Clone, serde::Serialize)] + pub struct HelloGetResponse { + pub greeting: String, + pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_id: Option, + } + #[derive(Debug, Clone, serde::Deserialize)] + pub struct StreamPostRequest { + pub prompt: String, + pub tokens: i64, + } +} diff --git a/scripts/cleanup.sh b/scripts/cleanup.sh new file mode 100755 index 0000000..ef8df5e --- /dev/null +++ b/scripts/cleanup.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# 清理本地验证进程（脚本化避免 pkill -f 自匹配问题） +set -u +for pat in "target/debug/examples/demo" ".venv/bin/worker_serve" "scripts/zenoh_router.py" ".venv/bin/serve" "target/debug/gen_models"; do + pids=$(pgrep -f "$pat" || true) + if [ -n "$pids" ]; then + kill $pids 2>/dev/null + echo "killed [$pat]: $(echo $pids | tr '\n' ' ')" + fi +done +sleep 1 +remaining=$(pgrep -f "target/debug/examples/demo|worker_serve|zenoh_router|\.venv/bin/serve" || true) +if [ -n "$remaining" ]; then + echo "WARN 仍存活: $remaining" + kill -9 $remaining 2>/dev/null +else + echo "清理完成" +fi diff --git a/scripts/verify_e2e.py b/scripts/verify_e2e.py new file mode 100644 index 0000000..e32eb6d --- /dev/null +++ b/scripts/verify_e2e.py @@ -0,0 +1,294 @@ +"""Rust 进程插件端到端协议验证（多语言演进规划阶段 1 验收） + +前置条件: + 1. 本地 Zenoh Router 已运行（scripts/zenoh_router.py） + 2. Rust 参考插件已编译（cargo build --examples 或 workers/rsdemo） + +验证内容（docs/plugin-wire-protocol.md 全协议面）: + 1. 路由描述符发布（keyexpr 编码、字段、schema_version/type） + 2. liveliness 服务发现（worker/rsdemo/*） + 3. 普通 RPC（GET 信封 / POST user_info 透传） + 4. 非信封载荷兼容（method 默认 GET） + 5. 流式多回复（逐块）+ 流中异常 {"error": ...} 收尾块 + 6. 非流式异常回复格式（result/operation/error/worker） + 7. SIGTERM 优雅退出 → 描述符 DELETE 摘除 + 8. 重启 → 10s 内重发布自愈 + +用法（项目根目录）: + .venv/bin/python utils/worker-rs/scripts/verify_e2e.py [--type rsdemo] [--count 1] [--bin <路径>] +""" +import argparse +import orjson +import os +import signal +import subprocess +import sys +import time + +import zenoh + +PASS = 0 +FAIL = 0 + +DEFAULT_BIN = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "target", "debug", "examples", "demo") +) +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +ROUTES_PREFIX = "chongming/routes" + + +def report(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f"✅ {name}" + (f" — {detail}" if detail else "")) + else: + FAIL += 1 + print(f"❌ {name}" + (f" — {detail}" if detail else "")) + + +def start_worker(worker_type, count=1, bin_path=None): + """启动 Rust 插件进程（CWD=项目根，读取全局 config.toml）""" + binary = bin_path or DEFAULT_BIN + proc = subprocess.Popen( + [binary, "--type", worker_type, "--count", str(count)], + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + # 等待 Queryable 就绪（日志出现"进程就绪"） + deadline = time.monotonic() + 20 + while time.monotonic() < deadline: + if proc.poll() is not None: + out = proc.stdout.read().decode() if proc.stdout else "" + raise RuntimeError(f"Worker 启动失败，退出码 {proc.returncode}:\n{out}") + time.sleep(0.2) + line = proc.stdout.readline() if proc.stdout else b"" + if "等待退出信号" in line.decode("utf-8", "ignore"): + return proc + raise RuntimeError("Worker 20s 内未就绪") + + +def wait_worker_exit(proc, timeout=10): + try: + proc.wait(timeout=timeout) + return True + except subprocess.TimeoutExpired: + return False + + +def make_subscriber(session, worker_type): + """建立常驻订阅，持续记录 PUT 描述符与 DELETE 事件（不重放历史值， + 调用方须在事件发生前建立订阅）""" + samples = {} + + def cb(sample): + key = str(sample.key_expr) + if sample.kind == zenoh.SampleKind.DELETE: + samples[key] = None # 摘除标记 + else: + samples[key] = orjson.loads(bytes(sample.payload)) + + sub = session.declare_subscriber(f"{ROUTES_PREFIX}/{worker_type}/**", cb) + return sub, samples + + +def rpc(session, worker_expr, subject, payload, timeout=10.0): + """按信封协议发起 Zenoh Query，返回所有回复块 JSON 列表""" + full_key = f"{worker_expr}/{subject}" + replies = session.get( + full_key, + payload=orjson.dumps(payload), + consolidation=zenoh.ConsolidationMode.NONE, + timeout=timeout, + ) + result = [] + for reply in replies: + if reply.ok: + result.append(orjson.loads(bytes(reply.ok.payload))) + else: + result.append({"__zenoh_err__": str(reply.err.payload)}) + return result + + +def get_worker_exprs(session, worker_type): + """liveliness 发现 worker/worker_type/* 副本列表""" + exprs = [] + for reply in session.liveliness().get(f"worker/{worker_type}/*"): + if reply.ok: + exprs.append(str(reply.ok.key_expr)) + return exprs + + +def main(): + global PASS, FAIL + parser = argparse.ArgumentParser() + parser.add_argument("--type", default="rsdemo") + parser.add_argument("--count", type=int, default=1) + parser.add_argument("--bin", default=None, help="插件二进制路径（默认 SDK 示例 demo）") + args = parser.parse_args() + wt = args.type + + session = zenoh.open(zenoh.Config.from_file("config/zenoh_config.json5")) + print(f"=== Rust 进程插件协议验证 (type={wt}, bin={args.bin or DEFAULT_BIN}) ===") + + # ========== 1. 路由描述符发布（先订阅再启动，zenoh 不重放历史值） ========== + sub1, samples1 = make_subscriber(session, wt) + proc = start_worker(wt, args.count, args.bin) + time.sleep(1.5) + + expect_keys = { + "rsdemo": { + "chongming/routes/rsdemo/rsdemo_hello__GET", + "chongming/routes/rsdemo/rsdemo_echo__POST", + "chongming/routes/rsdemo/rsdemo_stream__POST", + "chongming/routes/rsdemo/rsdemo_error__GET", + }, + "demo": { + "chongming/routes/demo/demo_hello__GET", + "chongming/routes/demo/demo_storage_upload_bytes__POST", + "chongming/routes/demo/demo_storage_download_bytes__GET", + "chongming/routes/demo/demo_storage_presign_upload__POST", + "chongming/routes/demo/demo_storage_presign_download__GET", + }, + }[wt] + missing = expect_keys - set(samples1.keys()) + report( + "路由描述符 keyexpr 编码正确", + not missing, + f"缺失: {missing or '无'}", + ) + + if "chongming/routes/rsdemo/rsdemo_hello__GET" in samples1: + d = samples1["chongming/routes/rsdemo/rsdemo_hello__GET"] + report("描述符 schema_version=1", d.get("schema_version") == 1, f"实际: {d.get('schema_version')}") + report("描述符 type=插件名", d.get("type") == wt, f"实际: {d.get('type')}") + report( + "描述符字段完整（subject/method/path/params/response_model）", + d.get("subject") == "rsdemo.hello" + and d.get("method") == "GET" + and d.get("path") == "/rsdemo/hello" + and d.get("params") == ["name: str"] + and isinstance(d.get("response_model"), dict), + ) + report( + "defaults 合并（timeout=30.0/auth_required=false）", + d.get("timeout") == 30.0 and d.get("auth_required") is False, + ) + sub1.undeclare() + + # ========== 2. liveliness 服务发现 ========== + exprs = get_worker_exprs(session, wt) + report( + "liveliness 服务发现副本", + len(exprs) == args.count, + f"实际: {exprs}", + ) + worker_expr = exprs[0] if exprs else f"worker/{wt}/missing" + + # ========== 3. 普通 RPC：GET 信封 ========== + replies = rpc(session, worker_expr, "rsdemo.hello", {"meta": {"method": "GET"}, "data": {"name": "rust"}}) + ok = ( + len(replies) == 1 + and replies[0].get("greeting") == "Hello, rust!" + and replies[0].get("worker_id") == worker_expr.split("/")[-1] + and isinstance(replies[0].get("timestamp"), (int, float)) + ) + report("RPC GET（信封分派 + worker_id 注入）", ok, f"实际: {replies}") + + # ========== 4. 非信封载荷兼容（method 默认 GET） ========== + replies = rpc(session, worker_expr, "rsdemo.hello", {"name": "legacy"}) + report( + "非信封载荷兼容（按 GET 处理）", + len(replies) == 1 and replies[0].get("greeting") == "Hello, legacy!", + f"实际: {replies}", + ) + + # ========== 5. POST + user_info 透传 ========== + replies = rpc( + session, worker_expr, "rsdemo.echo", + {"meta": {"method": "POST", "user_info": {"username": "alice"}}, "data": {"msg": "hi"}}, + ) + ok = ( + len(replies) == 1 + and replies[0].get("echo") == "hi" + and replies[0].get("user_info") == {"username": "alice"} + ) + report("RPC POST（user_info 信封透传）", ok, f"实际: {replies}") + + # ========== 6. 流式多回复 ========== + replies = rpc( + session, worker_expr, "rsdemo.stream", + {"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 5}}, + timeout=15.0, + ) + report( + "流式多回复（5 块逐块推送）", + len(replies) == 5 + and all(r.get("index") == i for i, r in enumerate(replies)) + and all("error" not in r for r in replies), + f"块数: {len(replies)}", + ) + + # 流中异常 → {"error": ...} 收尾块（第 26 块抛错，共 26 块） + replies = rpc( + session, worker_expr, "rsdemo.stream", + {"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 60}}, + timeout=20.0, + ) + ok = ( + len(replies) == 26 + and all("error" not in r for r in replies[:25]) + and isinstance(replies[25].get("error"), str) + and "模拟流中异常" in replies[25]["error"] + ) + report("流中异常 → error 收尾块终止", ok, f"块数: {len(replies)}, 末块: {replies[-1] if replies else None}") + + # ========== 7. 非流式异常回复格式 ========== + replies = rpc(session, worker_expr, "rsdemo.error", {"meta": {"method": "GET"}, "data": {}}) + ok = ( + len(replies) == 1 + and replies[0].get("result") == 0.0 + and replies[0].get("operation") == "error" + and "模拟业务异常" in replies[0].get("error", "") + and replies[0].get("worker") == worker_expr.split("/")[-1] + ) + report("非流式异常回复格式（result/operation/error/worker）", ok, f"实际: {replies}") + + # ========== 8. SIGTERM 优雅退出 → 描述符摘除（先订阅再 kill） ========== + sub2, samples2 = make_subscriber(session, wt) + os.kill(proc.pid, signal.SIGTERM) + exited = wait_worker_exit(proc, timeout=10) + report("SIGTERM 优雅退出", exited, f"退出码: {proc.returncode if exited else '超时'}") + time.sleep(1.5) + removed = [k for k, v in samples2.items() if v is None] + report( + "退出时 DELETE 摘除全部描述符", + len(removed) == len(expect_keys), + f"摘除: {sorted(removed)}", + ) + sub2.undeclare() + + # ========== 9. 重启 → 10s 内重发布自愈 ========== + sub3, samples3 = make_subscriber(session, wt) + proc = start_worker(wt, args.count, args.bin) + time.sleep(11.5) # 覆盖一个完整重发布周期 + ok = all(k in samples3 and samples3[k] is not None for k in expect_keys) + report( + "重启后周期重发布自愈（≤10s）", + ok, + f"恢复: {sorted(k for k in samples3 if samples3[k] is not None)}", + ) + sub3.undeclare() + + # 清理 + os.kill(proc.pid, signal.SIGTERM) + wait_worker_exit(proc, timeout=10) + session.close() + + print(f"\n=== 结果: {PASS} 通过, {FAIL} 失败 ===") + sys.exit(1 if FAIL else 0) + + +if __name__ == "__main__": + main() diff --git a/scripts/verify_mixed.py b/scripts/verify_mixed.py new file mode 100644 index 0000000..22a5b15 --- /dev/null +++ b/scripts/verify_mixed.py @@ -0,0 +1,188 @@ +"""Python 与 Rust 进程插件同 type 混布对拍（多语言演进规划 §3.4 场景 5） + +前置条件: + 1. 本地 Zenoh Router 已运行（scripts/zenoh_router.py） + 2. Rust 参考插件已编译；venv 已安装（worker_serve console script） + +验证内容: + 1. Python 与 Rust 对同一 TOML 配置发布的描述符内容 hash 一致 + （网关据此幂等去重，内容不一致会告警冲突） + 2. 同一 type 下两个实现副本均被 liveliness 发现（负载均衡透明） + 3. 两个实现的 RPC 均返回等价结果（worker_id 区分副本） + 4. SIGTERM 后各自摘除（不互相干扰） + +用法（项目根目录）: + .venv/bin/python utils/worker-rs/scripts/verify_mixed.py +""" +import hashlib +import orjson +import os +import signal +import subprocess +import sys +import time + +import zenoh + +PASS = 0 +FAIL = 0 + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +RUST_BIN = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "target", "debug", "examples", "demo") +) +PY_WORKER = os.path.join(ROOT, ".venv", "bin", "worker_serve") + + +def report(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f"✅ {name}" + (f" — {detail}" if detail else "")) + else: + FAIL += 1 + print(f"❌ {name}" + (f" — {detail}" if detail else "")) + + +def start(cmd, worker_type): + proc = subprocess.Popen( + cmd + ["--type", worker_type, "--count", "1"], + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + deadline = time.monotonic() + 25 + while time.monotonic() < deadline: + if proc.poll() is not None: + out = proc.stdout.read().decode() if proc.stdout else "" + raise RuntimeError(f"Worker 启动失败 (退出码 {proc.returncode}):\n{out}") + line = proc.stdout.readline() if proc.stdout else b"" + if "等待退出信号" in line.decode("utf-8", "ignore"): + return proc + raise RuntimeError("Worker 25s 未就绪") + + +def rpc(session, worker_expr, subject, payload, timeout=10.0): + full_key = f"{worker_expr}/{subject}" + replies = session.get( + full_key, + payload=orjson.dumps(payload), + consolidation=zenoh.ConsolidationMode.NONE, + timeout=timeout, + ) + return [orjson.loads(bytes(r.ok.payload)) for r in replies if r.ok] + + +def main(): + session = zenoh.open(zenoh.Config.from_file("config/zenoh_config.json5")) + print("=== Python + Rust 同 type(demo) 混布对拍 ===") + + # 先建订阅（zenoh 不重放历史值） + samples = {} + + def cb(sample): + if not sample.payload: + return # 空载荷样本（如 DELETE 预发布期），忽略 + key = str(sample.key_expr) + samples[key] = orjson.loads(bytes(sample.payload)) + + sub = session.declare_subscriber("chongming/routes/demo/**", cb) + + # 同时启动 Python 与 Rust 两个实现 + py_proc = start([PY_WORKER], "demo") + rs_proc = start([RUST_BIN], "demo") + time.sleep(2.0) + + # 1. 描述符内容一致性：同一 route_id 多次发布（Python/Rust 各一次）内容 hash 必须一致 + expect = { + "demo_hello__GET", + "demo_storage_upload_bytes__POST", + "demo_storage_download_bytes__GET", + "demo_storage_presign_upload__POST", + "demo_storage_presign_download__GET", + } + seen = {k: [] for k in expect} + for key, desc in samples.items(): + rid = key.split("/")[-1] + if rid in seen: + seen[rid].append(orjson.dumps(desc, option=orjson.OPT_SORT_KEYS).decode()) + + all_same = True + for rid, versions in seen.items(): + if len(versions) < 1: + continue + hashes = {hashlib.sha1(v.encode("utf-8")).hexdigest()[:8] for v in versions} + if len(hashes) > 1: + all_same = False + print(f" ⚠ {rid} 描述符内容不一致: {hashes}") + report( + "Python/Rust 描述符逐字段一致（内容 hash 唯一，无冲突告警）", + all_same and all(seen.values()), + f"各 route_id 版本数: { {k: len(v) for k, v in seen.items()} }", + ) + + # 2. liveliness：同 type 两个实现副本均被发现 + exprs = [ + str(r.ok.key_expr) + for r in session.liveliness().get("worker/demo/*") + if r.ok + ] + report( + "liveliness 发现 Python + Rust 副本（共 2 个）", + len(exprs) == 2, + f"实际: {exprs}", + ) + if len(exprs) < 2: + print("混布副本不足，跳过 RPC 对拍") + else: + # 区分 Python 副本与 Rust 副本：Rust worker_id 格式同为 worker-0-， + # 无法从 key 区分实现，但可分别对两个副本调用并验证都响应 + for expr in exprs: + replies = rpc( + session, expr, "demo.hello", + {"meta": {"method": "GET"}, "data": {"name": "mixed"}}, + ) + ok = ( + len(replies) == 1 + and replies[0].get("greeting") == "Hello, mixed!" + and replies[0].get("worker_id") == expr.split("/")[-1] + ) + report(f"RPC 对拍 {expr}", ok, f"实际: {replies}") + + # storage 端点：两实现均应返回错误（Python 未配置 storage 降级；Rust 未实现降级） + replies = rpc( + session, exprs[0], "demo.storage.upload_bytes", + {"meta": {"method": "POST"}, "data": {"object_key": "a.txt", "data_b64": "eA=="}}, + ) + report( + "storage 端点两实现均降级为错误", + len(replies) == 1 + and replies[0].get("operation") == "error", + f"实际: {replies}", + ) + + # 3. SIGTERM 顺序退出：互不干扰 + os.kill(rs_proc.pid, signal.SIGTERM) + rs_proc.wait(timeout=10) + time.sleep(1.5) + remain = [ + str(r.ok.key_expr) + for r in session.liveliness().get("worker/demo/*") + if r.ok + ] + report( + "Rust 退出后 Python 副本不受影响", + len(remain) == 1, + f"剩余: {remain}", + ) + os.kill(py_proc.pid, signal.SIGTERM) + py_proc.wait(timeout=10) + sub.undeclare() + session.close() + + print(f"\n=== 结果: {PASS} 通过, {FAIL} 失败 ===") + sys.exit(1 if FAIL else 0) + + +if __name__ == "__main__": + main() diff --git a/scripts/zenoh_router.py b/scripts/zenoh_router.py new file mode 100644 index 0000000..24252a1 --- /dev/null +++ b/scripts/zenoh_router.py @@ -0,0 +1,20 @@ +"""本地测试用 Zenoh Router（mode=router，监听 tcp/127.0.0.1:7447）。 + +替代 zenohd：环境无独立 router 二进制时，用 Python zenoh 开一个 router 节点， +供网关（peer + connect 7447）与 worker 进程共同接入。仅用于本地联调。 +""" +import time + +import zenoh + +config = zenoh.Config() +config.insert_json5("mode", '"router"') +session = zenoh.open(config) +print(">>> zenoh router listening on tcp/127.0.0.1:7447", flush=True) +try: + while True: + time.sleep(3600) +except KeyboardInterrupt: + pass +finally: + session.close() diff --git a/src/admin.rs b/src/admin.rs new file mode 100644 index 0000000..d706ad6 --- /dev/null +++ b/src/admin.rs @@ -0,0 +1,284 @@ +//! admin Queryable 指令处理（对齐 Python runtime.py `_handle_admin_command`） +//! +//! 载荷（信封或裸 JSON，优先信封 data 字段）形如 `{"cmd": "restart"}`。 +//! 支持指令: +//! - （无/其他）→ 状态查询 +//! - restart → 平滑退出并以新 id 重建 +//! - stop → 平滑退出且不重建 +//! - drain → 优雅排空（隔离新请求，等在途完成，数量-1） +//! - get_config → 读 [workers.config]（缓存优先，回退 toml） +//! - config_set → 合并写入 [workers.config]（仅 index0 落盘）+ 热更 + 重启 +//! - rescale → 调整本进程副本池目标数（逐实例绝对 / 全局分摊） + +use std::sync::Arc; + +use serde_json::{json, Value}; +use tokio::sync::Mutex; + +use crate::cache; +use crate::config; +use crate::runtime::{ + ReplicaController, WorkerPool, ADMIN_CONFIG_SET, ADMIN_DRAIN, ADMIN_GET_CONFIG, + ADMIN_RESTART, ADMIN_STOP, +}; + +/// 处理单个副本的 admin 指令 +pub async fn handle_admin( + query: zenoh::query::Query, + controller: Arc, + pool: Arc>, +) { + let key_expr = query.key_expr().to_string(); + let payload = query + .payload() + .map(|p| p.to_bytes().to_vec()) + .unwrap_or_default(); + let worker_id = controller.worker_id(); + + let raw: Value = serde_json::from_slice(&payload).unwrap_or(Value::Null); + // 解析：信封（meta/data）优先取 data；指令 cmd 从 data 或整包取 + let is_envelope = matches!(&raw, Value::Object(m) if m.contains_key("meta") || m.contains_key("data")); + let data = if is_envelope { + raw.get("data").cloned().unwrap_or(Value::Null) + } else { + raw.clone() + }; + let cmd = data.get("cmd").and_then(Value::as_str).unwrap_or("").to_lowercase(); + + match cmd.as_str() { + ADMIN_GET_CONFIG => reply_get_config(&query, &key_expr, &controller, &worker_id).await, + ADMIN_CONFIG_SET => config_set(&query, &key_expr, &controller, &worker_id, &data).await, + ADMIN_RESTART => { + controller.request_restart(); + reply_accepted(&query, &key_expr, &cmd, &controller, &worker_id, "").await; + } + ADMIN_STOP => { + controller.request_stop(); + reply_accepted(&query, &key_expr, &cmd, &controller, &worker_id, "").await; + } + ADMIN_DRAIN => { + controller.request_drain(); + reply_accepted(&query, &key_expr, &cmd, &controller, &worker_id, "优雅排空").await; + } + "rescale" => { + match rescale(&controller, &pool, &data).await { + Ok(note) => { + reply_accepted(&query, &key_expr, &cmd, &controller, &worker_id, ¬e).await + } + Err(e) => reply_error(&query, &key_expr, &cmd, &controller, &worker_id, &e).await, + } + } + // 无指令 / 未知指令 → 状态查询 + _ => reply_status(&query, &key_expr, &controller, &worker_id).await, + } +} + +/// 状态查询回复 +async fn reply_status( + query: &zenoh::query::Query, + key_expr: &str, + controller: &Arc, + worker_id: &str, +) { + let _ = query + .reply( + key_expr.to_string(), + json!({ + "worker_type": controller.worker_type, + "worker_id": worker_id, + "index": controller.index, + "epoch": controller.epoch(), + "instance_id": controller.instance_id, + "status": "running", + }) + .to_string(), + ) + .await; +} + +/// get_config：读 [workers.config]（缓存优先，回退 toml） +async fn reply_get_config( + query: &zenoh::query::Query, + key_expr: &str, + controller: &Arc, + worker_id: &str, +) { + let (cfg, source) = if controller.cache.enabled() { + let v = controller + .cache + .get(&cache::key(cache::_CONFIGS_KEY, &controller.worker_type, None)) + .await + .unwrap_or_else(|| Value::Object(Default::default())); + (v, "cache".to_string()) + } else { + ( + config::read_worker_config_section(&controller.worker_type, &controller.config_dir), + "toml".to_string(), + ) + }; + let _ = query + .reply( + key_expr.to_string(), + json!({ + "status": "ok", "cmd": "get_config", "worker_type": controller.worker_type, + "worker_id": worker_id, "config": cfg, "source": source, + }) + .to_string(), + ) + .await; +} + +/// config_set：合并写入 [workers.config]（仅 index0 持久化）+ 热更重启 +async fn config_set( + query: &zenoh::query::Query, + key_expr: &str, + controller: &Arc, + worker_id: &str, + data: &Value, +) { + let updates = data.get("config").cloned().unwrap_or(Value::Null); + if !updates.is_object() || updates.as_object().unwrap().is_empty() { + reply_error(query, key_expr, "config_set", controller, worker_id, "config_set 需要 data.config 字段（非空对象）").await; + return; + } + let allowed = config::count_allowed(&updates); + + if controller.index == 0 { + // 仅 index0 持久化（缓存或 toml），并热更 + 重启 + let merged = if controller.cache.enabled() { + let cur = controller + .cache + .get(&cache::key(cache::_CONFIGS_KEY, &controller.worker_type, None)) + .await + .unwrap_or_else(|| Value::Object(Default::default())); + let merged = config::merge_confirmed(cur, &updates); + if let Err(e) = controller + .cache + .set(&cache::key(cache::_CONFIGS_KEY, &controller.worker_type, None), &merged) + .await + { + reply_error(query, key_expr, "config_set", controller, worker_id, &format!("配置写入缓存失败: {e}")).await; + return; + } + merged + } else { + match config::persist_worker_config_section( + &controller.worker_type, + &controller.config_dir, + &updates, + ) { + Ok(m) => m, + Err(e) => { + reply_error(query, key_expr, "config_set", controller, worker_id, &format!("配置写入失败: {e}")).await; + return; + } + } + }; + *controller.pending_config.lock().unwrap() = Some(merged); + controller.signal(ADMIN_RESTART); + let note = if controller.cache.enabled() { + format!("{allowed} 键已写缓存并触发重启/热更") + } else { + format!("{allowed} 键已写 config.toml 并触发重启/热更") + }; + reply_accepted(query, key_expr, "config_set", controller, worker_id, ¬e).await; + } else { + // 非 index0：仅合并待应用，随重启生效（不落盘，避免并发写冲突） + let cur = config::read_worker_config_section(&controller.worker_type, &controller.config_dir); + let merged = config::merge_confirmed(cur, &updates); + *controller.pending_config.lock().unwrap() = Some(merged); + controller.signal(ADMIN_RESTART); + reply_accepted(query, key_expr, "config_set", controller, worker_id, "非 index0 副本：仅应用配置变更").await; + } +} + +/// rescale：调整本进程副本池目标数（逐实例绝对 / 全局分摊） +async fn rescale( + controller: &Arc, + pool: &Arc>, + data: &Value, +) -> Result { + let target = data.get("scale").and_then(Value::as_i64).unwrap_or(-1); + if target < 0 { + return Err("rescale 需要 data.scale（非负整数目标副本数）".into()); + } + let total = target.max(0) as usize; + let mut guard = pool.lock().await; + + let target_inst = data.get("instance_id").and_then(Value::as_str); + let absolute = data.get("absolute").and_then(Value::as_bool).unwrap_or(false); + if let Some(iid) = target_inst { + if absolute && iid == controller.instance_id { + cache::write_own_instance_count(&controller.cache, &controller.worker_type, iid, target.max(0)).await; + guard.set_target(total, pool).await; + return Ok(format!("实例 {iid} 绝对 count={total}")); + } else if absolute { + return Ok(format!("实例级缩放目标 {iid} 非本实例，忽略")); + } + } + + // 全局对称：清逐实例覆盖，按存活实例分摊 + cache::clear_instance_counts(&controller.cache, &controller.worker_type).await; + if controller.cache.enabled() { + let _ = controller + .cache + .set( + &cache::key(cache::_SCALE_COUNTS_KEY, &controller.worker_type, None), + &json!(target.max(0)), + ) + .await; + let app = cache::distribute_quota( + &controller.cache, + &controller.worker_type, + target.max(0), + None, + ) + .await + .max(1); + guard.set_target(app as usize, pool).await; + Ok(format!("全局 count={total}，本进程应用={app} 已写缓存")) + } else { + guard.set_target(total, pool).await; + Ok("缓存未启用，count 仅内存缩放".to_string()) + } +} + +async fn reply_accepted( + query: &zenoh::query::Query, + key_expr: &str, + cmd: &str, + controller: &Arc, + worker_id: &str, + note: &str, +) { + let _ = query + .reply( + key_expr.to_string(), + json!({ + "status": "accepted", "cmd": cmd, "worker_type": controller.worker_type, + "worker_id": worker_id, "persist": note, + }) + .to_string(), + ) + .await; +} + +async fn reply_error( + query: &zenoh::query::Query, + key_expr: &str, + cmd: &str, + controller: &Arc, + worker_id: &str, + error: &str, +) { + let _ = query + .reply( + key_expr.to_string(), + json!({ + "status": "error", "cmd": cmd, "worker_type": controller.worker_type, + "worker_id": worker_id, "error": error, + }) + .to_string(), + ) + .await; +} \ No newline at end of file diff --git a/src/bin/gen_models.rs b/src/bin/gen_models.rs new file mode 100644 index 0000000..1ef4c0b --- /dev/null +++ b/src/bin/gen_models.rs @@ -0,0 +1,54 @@ +//! gen_models —— 从端点 TOML 生成 Rust 模型（对齐 Python gen_models） +//! +//! 用法（项目根目录）: +//! cargo run --bin gen_models -- --types rsdemo,demo \ +//! --out utils/worker-rs/examples/generated_models.rs +//! +//! 生成文件头部含 FINGERPRINT（TOML 内容 hash），插件启动时校验， +//! 保证 TOML 与模型代码不漂移。 + +use std::path::Path; + +use chongming_worker_rs::gen; +use clap::Parser; + +#[derive(Parser, Debug)] +#[command(name = "gen_models", about = "从 config/workers//*.toml 生成 Rust 模型")] +struct Args { + /// 插件类型列表（逗号分隔，可重复） + #[arg(long, value_delimiter = ',')] + types: Vec, + /// 插件配置目录 + #[arg(long, default_value = "config/workers")] + config_dir: String, + /// 输出文件路径 + #[arg(long)] + out: String, +} + +fn main() { + let args = Args::parse(); + if args.types.is_empty() { + eprintln!("错误: 至少需要一个 --types"); + std::process::exit(2); + } + let type_refs: Vec<&str> = args.types.iter().map(|s| s.as_str()).collect(); + match gen::generate_all(&type_refs, &args.config_dir) { + Ok(source) => { + if let Some(parent) = Path::new(&args.out).parent() { + let _ = std::fs::create_dir_all(parent); + } + std::fs::write(&args.out, source) + .map_err(|e| format!("写入 {} 失败: {e}", args.out)) + .unwrap_or_else(|e| { + eprintln!("错误: {e}"); + std::process::exit(1); + }); + println!("已生成 {} 个插件类型模型 -> {}", args.types.join(", "), args.out); + } + Err(e) => { + eprintln!("错误: {e}"); + std::process::exit(1); + } + } +} diff --git a/src/cache.rs b/src/cache.rs new file mode 100644 index 0000000..7d7cd52 --- /dev/null +++ b/src/cache.rs @@ -0,0 +1,354 @@ +//! 配置缓存抽象（对齐 Python 的 chongming-cache 语义） +//! +//! 进程插件可通过 `[workers.cache]` 段启用共享缓存： +//! - 默认无缓存（`Cache::None`）：配置回退 toml 种子、实例分摊退化为单实例全额， +//! 行为与纯 Python 单实例一致； +//! - `backend = "redis"`（feature `redis-cache`）：跨实例收敛配置与副本数所需的 +//! 共享键（`worker:cm:*`）均与网关共享同一套协议。 +//! +//! 说明：Python 默认的本地 `diskcache` 后端本身是 Python 库，Rust 无法复用其本地存储 +//! 文件格式；Rust 侧分布式协同走 Redis（共享后端），本地单实例走 `Cache::None`。 +//! 这与 Python 的「多实例必须 redis/共享卷 diskcache」约束一致——Redis 是跨机收敛的推荐路径。 + +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde_json::{json, Value}; + +// ---- 共享缓存键（与网关/ Python 共享同一套协议，见 _root.py / runtime.py）---- +pub const _SCALE_COUNTS_KEY: &str = "worker:cm:counts:{type}"; +pub const _SCALE_GATEWAYS_KEY: &str = "worker:cm:gateways"; +pub const _SCALE_QUOTA_KEY: &str = "worker:cm:quota:{type}:{gw}"; +pub const _SCALE_INSTANCES_KEY: &str = "worker:cm:instances:{type}"; +pub const _SCALE_INSTCOUNTS_KEY: &str = "worker:cm:instcounts:{type}"; +pub const _CONFIGS_KEY: &str = "worker:cm:configs:{type}"; +pub const _LEASE: u64 = 60; + +pub fn now_ts() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// 按类型展开缓存键模板 +pub fn key(template: &str, worker_type: &str, gw: Option<&str>) -> String { + template + .replace("{type}", worker_type) + .replace("{gw}", gw.unwrap_or("")) +} + +/// 共享缓存后端（None = 无缓存） +#[derive(Clone)] +pub enum Cache { + /// 未启用共享缓存（默认）。get 恒 None、set 恒成功（no-op）。 + None, + /// Redis 共享缓存（feature `redis-cache` 编译）。JSON 值存储与网关一致。 + #[cfg(feature = "redis-cache")] + Redis(redis::aio::ConnectionManager), +} + +impl std::fmt::Debug for Cache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Cache::None => write!(f, "Cache::None"), + #[cfg(feature = "redis-cache")] + Cache::Redis(_) => write!(f, "Cache::Redis"), + } + } +} + +impl Cache { + pub fn enabled(&self) -> bool { + !matches!(self, Cache::None) + } + + /// 读取键 → JSON 值（chongming-cache 存储 JSON 字符串） + pub async fn get(&self, _key: &str) -> Option { + match self { + Cache::None => None, + #[cfg(feature = "redis-cache")] + Cache::Redis(ref con) => { + let mut con = con.clone(); + let raw: Option = redis::cmd("GET") + .arg(_key) + .query_async(&mut con) + .await + .ok() + .flatten(); + raw.and_then(|s| serde_json::from_str(&s).ok()) + } + } + } + + /// 写入键（JSON 序列化存储） + pub async fn set(&self, _key: &str, _value: &Value) -> Result<(), String> { + match self { + Cache::None => Ok(()), + #[cfg(feature = "redis-cache")] + Cache::Redis(ref con) => { + let mut con = con.clone(); + let s = serde_json::to_string(_value) + .map_err(|e| format!("缓存值序列化失败: {e}"))?; + redis::cmd("SET") + .arg(_key) + .arg(s) + .query_async::<()>(&mut con) + .await + .map_err(|e| format!("缓存写入失败: {e}")) + } + } + } + + /// 删除键 + pub async fn delete(&self, _key: &str) -> Result<(), String> { + match self { + Cache::None => Ok(()), + #[cfg(feature = "redis-cache")] + Cache::Redis(ref con) => { + let mut con = con.clone(); + redis::cmd("DEL") + .arg(_key) + .query_async::<()>(&mut con) + .await + .map_err(|e| format!("缓存删除失败: {e}")) + } + } + } + + /// 按 `[workers.cache]` 配置建立共享缓存（无法连接的后端降级为 None + 告警）。 + pub async fn from_config(cfg: &Value) -> (Cache, String) { + let backend = cfg.get("backend").and_then(Value::as_str).unwrap_or("diskcache"); + // 仅 redis 后端可跨实例收敛；diskcache/缺失 → None（single-instance 语义） + if backend != "redis" { + return (Cache::None, backend.to_string()); + } + #[cfg(not(feature = "redis-cache"))] + { + tracing::warn!( + "[cache] backend=redis 需要 cargo feature `redis-cache`（未编译），降级为无缓存" + ); + return (Cache::None, backend.to_string()); + } + #[cfg(feature = "redis-cache")] + { + let url = redis_url(cfg); + if url.is_empty() { + return (Cache::None, backend.to_string()); + } + let client = match redis::Client::open(url.as_str()) { + Ok(c) => c, + Err(e) => { + tracing::warn!("[cache] 连接 Redis {} 失败: {e}（降级为无缓存）", url); + return (Cache::None, backend.to_string()); + } + }; + match client.get_connection_manager().await { + Ok(con) => { + tracing::info!("[cache] 已连接 Redis（{}）", url); + (Cache::Redis(con), backend.to_string()) + } + Err(e) => { + tracing::warn!("[cache] 连接 Redis {} 失败: {e}（降级为无缓存）", url); + (Cache::None, backend.to_string()) + } + } + } + } +} + +#[cfg(feature = "redis-cache")] +fn redis_url(cfg: &Value) -> String { + if let Some(u) = cfg.get("url").and_then(Value::as_str) { + return u.to_string(); + } + let host = cfg.get("host").and_then(Value::as_str).unwrap_or("127.0.0.1"); + let port = cfg.get("port").and_then(Value::as_u64).unwrap_or(6379); + let db = cfg.get("db").and_then(Value::as_i64).unwrap_or(0); + match cfg.get("password").and_then(Value::as_str) { + Some(pw) => format!("redis://:{}@{}:{}/{}", pw, host, port, db), + None => format!("redis://{}:{}/{}", host, port, db), + } +} + +// ============================================================ +// 实例参与方登记 / 心跳 / 分摊（对齐 Python runtime.py） +// ============================================================ + +/// 读取并清理某类型实例参与方列表（[[inst_id, ts], ...]），返回存活实例 id（排序） +pub async fn read_instances(cache: &Cache, worker_type: &str) -> Vec { + let Some(raw) = cache.get(&key(_SCALE_INSTANCES_KEY, worker_type, None)).await else { + return Vec::new(); + }; + let now = now_ts(); + let mut out: Vec = Vec::new(); + if let Value::Array(arr) = raw { + for item in arr { + let inst = item.get(0).and_then(Value::as_str); + let ts = item.get(1).and_then(Value::as_u64).unwrap_or(0); + if item.as_array().map(|a| a.first().is_none()).unwrap_or(false) { + continue; + } + if let Some(id) = inst { + if id.is_empty() { + continue; + } + if now - ts > _LEASE { + continue; // 租约过期：视为死实例 + } + out.push(id.to_string()); + } + } + } + out.sort(); + out +} + +/// 登记当前实例为参与方（写心跳） +pub async fn heartbeat_instance(cache: &Cache, worker_type: &str, instance_id: &str) { + if !cache.enabled() { + return; + } + let mut instances = read_instances(cache, worker_type).await; + if !instances.iter().any(|s| s == instance_id) { + instances.push(instance_id.to_string()); + instances.sort(); + } + let arr: Value = Value::Array( + instances + .iter() + .map(|id| json!([id, now_ts()])) + .collect(), + ); + let _ = cache.set(&key(_SCALE_INSTANCES_KEY, worker_type, None), &arr).await; +} + +/// 摘除实例（优雅退出时） +pub async fn unregister_instance(cache: &Cache, worker_type: &str, instance_id: &str) { + if !cache.enabled() { + return; + } + let mut instances = read_instances(cache, worker_type).await; + instances.retain(|s| s != instance_id); + let arr: Value = Value::Array(instances.iter().map(|id| json!([id, now_ts()])).collect()); + let _ = cache.set(&key(_SCALE_INSTANCES_KEY, worker_type, None), &arr).await; +} + +/// 读逐实例绝对副本数覆盖表 {instance_id: count} +pub async fn read_instance_counts(cache: &Cache, worker_type: &str) -> std::collections::HashMap { + let mut out = std::collections::HashMap::new(); + if let Some(Value::Object(map)) = cache.get(&key(_SCALE_INSTCOUNTS_KEY, worker_type, None)).await { + for (k, v) in map { + if let Some(n) = v.as_i64() { + out.insert(k, n); + } + } + } + out +} + +/// 清除全部逐实例覆盖（全局对称扩缩时调用） +pub async fn clear_instance_counts(cache: &Cache, worker_type: &str) { + let _ = cache.set(&key(_SCALE_INSTCOUNTS_KEY, worker_type, None), &Value::Object(Default::default())).await; +} + +/// 写当前实例的绝对副本数覆盖（从缓存读-改-写） +pub async fn write_own_instance_count(cache: &Cache, worker_type: &str, instance_id: &str, count: i64) { + let mut counts = read_instance_counts(cache, worker_type).await; + counts.insert(instance_id.to_string(), count); + let obj: Value = Value::Object( + counts.into_iter().map(|(k, v)| (k, json!(v))).collect(), + ); + let _ = cache.set(&key(_SCALE_INSTCOUNTS_KEY, worker_type, None), &obj).await; +} + +/// 当前实例的绝对覆盖（有则返回） +pub async fn own_instance_count(cache: &Cache, worker_type: &str, instance_id: &str) -> Option { + read_instance_counts(cache, worker_type) + .await + .get(instance_id) + .copied() +} + +/// 读在线网关列表（[[gw, ts], ...] 含过期清理），排序返回 id +pub async fn read_online_gateways(cache: &Cache) -> Vec { + let Some(Value::Array(arr)) = cache.get(_SCALE_GATEWAYS_KEY).await else { + return Vec::new(); + }; + let now = now_ts(); + let mut out: Vec = Vec::new(); + for item in arr { + let gw = item.get(0).and_then(Value::as_str).unwrap_or(""); + let ts = item.get(1).and_then(Value::as_u64).unwrap_or(0); + if gw.is_empty() || now - ts > _LEASE { + continue; + } + out.push(gw.to_string()); + } + out.sort(); + out +} + +/// 全局 count 按分担方均分：网关身份按在线网关分摊；无网关身份（外部进程） +/// 按存活实例分摊。返回本分担方应承载的副本数。 +pub async fn distribute_quota( + cache: &Cache, + worker_type: &str, + total: i64, + gw_id: Option<&str>, +) -> i64 { + if !cache.enabled() { + return total; + } + let participants: Vec = match gw_id { + Some(gw) if !gw.is_empty() => read_online_gateways(cache).await, + _ => read_instances(cache, worker_type).await, + }; + if participants.is_empty() { + return total; + } + let base = total / participants.len() as i64; + // 余数分配给排序较早的分担方（与 Python 一致） + let my_index = match gw_id { + Some(gw) if !gw.is_empty() => participants.iter().position(|s| s == gw), + _ => None, + }; + let my_index = my_index.unwrap_or(0); + let extra = (total % participants.len() as i64) as usize; + let remainder_mine = if my_index < extra { 1 } else { 0 }; + base + remainder_mine +} + +/// 启动时解析本进程应承载的副本数（对齐 Python resolve_worker_count） +pub async fn resolve_worker_count( + cache: &Cache, + worker_type: &str, + default: usize, + mode: &str, + gw_id: Option<&str>, + instance_id: Option<&str>, +) -> usize { + if !cache.enabled() { + return default; + } + // 逐实例绝对覆盖优先（精细扩缩） + if gw_id.is_none() { + if let Some(iid) = instance_id { + if let Some(own) = own_instance_count(cache, worker_type, iid).await { + return own.max(1) as usize; + } + } + } + let mut total = default as i64; + if mode == "cache" { + if let Some(v) = cache.get(&key(_SCALE_COUNTS_KEY, worker_type, None)).await { + if let Some(n) = v.as_i64() { + if n >= 0 { + total = n; + } + } + } + } + let _ = cache.set(&key(_SCALE_COUNTS_KEY, worker_type, None), &json!(total)).await; + let app = distribute_quota(cache, worker_type, total, gw_id).await; + app.max(1) as usize +} \ No newline at end of file diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..8c4089e --- /dev/null +++ b/src/config.rs @@ -0,0 +1,310 @@ +//! 配置读取（与 Python SDK process.py 完全对位） +//! +//! - 全局配置 `config.toml`：仅读取 `[zenoh].json5_path`（默认 `config/zenoh_config.json5`） +//! - 插件配置 `config/workers//config.toml` 的 `[workers]` 节：展平 + mqtt_topics + +use std::path::Path; + +use serde_json::Value; + +/// 全局配置（仅 zenoh 会话配置路径） +pub struct GlobalConfig { + pub zenoh_json5_path: String, +} + +/// 插件配置（[workers] 节，语义对齐 Python load_plugin_config） +pub struct PluginConfig { + pub worker_type: String, + pub module: String, + pub factory: String, + pub cleanup: String, + pub count: usize, + pub load_mode: String, + pub mqtt_topics: Vec, + /// 展平后的运行时配置（`[workers.config]` 等，供插件自用） + pub runtime: Value, +} + +/// 读取全局 config.toml 的 `[zenoh].json5_path` +/// +/// `json5_path` 若为相对路径，则基于 config.toml 所在目录解析为绝对路径——这样 +/// worker 从任何 CWD 启动（含容器部署）都能定位到与之同目录的 zenoh 会话配置， +/// 与 Python 进程插件"配置自包含在插件根"的约定一致。 +pub fn load_global_config(path: &str) -> Result { + let cfg_path = std::path::Path::new(path); + let text = + std::fs::read_to_string(cfg_path).map_err(|e| format!("读取全局配置 {path} 失败: {e}"))?; + let cfg: toml::Value = + toml::from_str(&text).map_err(|e| format!("解析全局配置 {path} 失败: {e}"))?; + let json5_path = cfg + .get("zenoh") + .and_then(|z| z.get("json5_path")) + .and_then(toml::Value::as_str) + .unwrap_or("config/zenoh_config.json5"); + let base_dir = cfg_path + .parent() + .unwrap_or_else(|| std::path::Path::new(".")); + let resolved = std::path::Path::new(json5_path); + let zenoh_json5_path = if resolved.is_absolute() { + resolved.to_path_buf() + } else { + base_dir.join(resolved) + }; + Ok(GlobalConfig { + zenoh_json5_path: zenoh_json5_path.to_string_lossy().into_owned(), + }) +} + +/// 读取 `config/workers//config.toml` 的 `[workers]` 节并展平 +/// （展平规则与 Python SDK load_plugin_config 一致，worker 侧独立实现） +pub fn load_plugin_config(worker_type: &str, config_dir: &str) -> Result { + let path = Path::new(config_dir).join(worker_type).join("config.toml"); + let text = std::fs::read_to_string(&path) + .map_err(|e| format!("未找到插件配置 {}: {e}", path.display()))?; + let plugin_conf: toml::Value = + toml::from_str(&text).map_err(|e| format!("解析插件配置 {} 失败: {e}", path.display()))?; + + let workers = plugin_conf + .get("workers") + .and_then(toml::Value::as_table) + .ok_or_else(|| format!("{} 缺少 [workers] 节", path.display()))?; + + let plugin_type = workers + .get("type") + .and_then(toml::Value::as_str) + .unwrap_or(worker_type) + .to_string(); + let load_mode = workers + .get("load_mode") + .and_then(toml::Value::as_str) + .unwrap_or("inline") + .to_string(); + + // mqtt_topics：仅保留带 topic 的 dict，并补 worker 字段（与 Python 一致） + let mut mqtt_topics: Vec = Vec::new(); + if let Some(topics) = workers.get("mqtt_topics").and_then(toml::Value::as_array) { + for t in topics { + let has_topic = t.get("topic").map(toml::Value::is_str).unwrap_or(false); + if !has_topic { + continue; + } + let mut t = t.clone(); + let has_worker = t.get("worker").map(toml::Value::is_str).unwrap_or(false); + if !has_worker { + if let Some(table) = t.as_table_mut() { + table.insert( + "worker".to_string(), + toml::Value::String(plugin_type.clone()), + ); + } + } + mqtt_topics.push(toml_to_json(&t)); + } + } + + // 展平：插件全量配置 + [workers] 节（除 config/mqtt_topics/load_mode）+ [workers.config] + let mut runtime: toml::Value = plugin_conf.clone(); + if let Some(flat) = runtime.as_table_mut() { + for (k, v) in workers { + if k == "config" || k == "mqtt_topics" || k == "load_mode" { + continue; + } + flat.insert(k.clone(), v.clone()); + } + } + if let Some(cfg) = workers.get("config") { + if let Some(table) = runtime.as_table_mut() { + table.insert("config".to_string(), cfg.clone()); + } + } + // 顶层保留完整 [workers]（含 config），与 Python runtime_config 结构一致 + let runtime_json = toml_to_json(&runtime); + + Ok(PluginConfig { + worker_type: plugin_type, + module: workers + .get("module") + .and_then(toml::Value::as_str) + .unwrap_or(&format!("plugins.{worker_type}")) + .to_string(), + factory: workers + .get("factory") + .and_then(toml::Value::as_str) + .unwrap_or("create_workers") + .to_string(), + cleanup: workers + .get("cleanup") + .and_then(toml::Value::as_str) + .unwrap_or("cleanup") + .to_string(), + count: workers + .get("count") + .and_then(toml::Value::as_integer) + .map(|c| c.max(1) as usize) + .unwrap_or(1), + load_mode, + mqtt_topics, + runtime: runtime_json, + }) +} + +/// toml::Value → serde_json::Value（插件运行时配置直接以 JSON 暴露） +pub fn toml_to_json(v: &toml::Value) -> Value { + match v { + toml::Value::String(s) => Value::String(s.clone()), + toml::Value::Integer(i) => Value::from(*i), + toml::Value::Float(f) => Value::from(*f), + toml::Value::Boolean(b) => Value::Bool(*b), + toml::Value::Datetime(d) => Value::String(d.to_string()), + toml::Value::Array(a) => Value::Array(a.iter().map(toml_to_json).collect()), + toml::Value::Table(t) => Value::Object( + t.iter() + .map(|(k, v)| (k.clone(), toml_to_json(v))) + .collect(), + ), + } +} + +// ============================================================ +// [workers.config] 运行时读写（对齐 Python runtime.py 的 config 管理） +// ============================================================ + +/// 禁止通过 config_set 写入的保留键（破坏插件加载/路由语义） +pub const CONFIG_FORBIDDEN_KEYS: &[&str] = &[ + "type", + "module", + "factory", + "cleanup", + "count", + "load_mode", + "mqtt_topics", + "workers", + "worker", + "defaults", + "registration", +]; + +/// 插件 config.toml 路径（config/workers//config.toml） +pub fn plugin_config_path(worker_type: &str, config_dir: &str) -> std::path::PathBuf { + std::path::Path::new(config_dir) + .join(worker_type) + .join("config.toml") +} + +/// 读该类型 `[workers].config` 段（对象），缺失/损坏返回空对象 +pub fn read_worker_config_section(worker_type: &str, config_dir: &str) -> Value { + let path = plugin_config_path(worker_type, config_dir); + let text = match std::fs::read_to_string(&path) { + Ok(t) => t, + Err(_) => return Value::Object(Default::default()), + }; + let data: toml::Value = toml::from_str(&text).unwrap_or(toml::Value::Table(Default::default())); + match data.get("workers").and_then(toml::Value::as_table) { + Some(w) => match w.get("config") { + Some(cfg) => toml_to_json(cfg), + None => Value::Object(Default::default()), + }, + None => Value::Object(Default::default()), + } +} + +/// 把 updates 合并写回该类型 `[workers].config`（仅 index0 落盘）。 +/// 返回合并后的段；校验保留键失败抛 Err。 +pub fn persist_worker_config_section( + worker_type: &str, + config_dir: &str, + updates: &Value, +) -> Result { + if !updates.is_object() { + return Err("config_set 需要 data.config 字段（非空对象）".into()); + } + let bad: Vec = updates + .as_object() + .unwrap() + .keys() + .filter(|k| CONFIG_FORBIDDEN_KEYS.contains(&k.as_str())) + .cloned() + .collect(); + if !bad.is_empty() { + return Err(format!("禁止修改保留配置键: {bad:?}")); + } + + let path = plugin_config_path(worker_type, config_dir); + let current = merge_confirmed(read_worker_config_section(worker_type, config_dir), updates); + // 写回 toml 文件：[workers].config = merged（缺 [workers] 节则创建） + let text = std::fs::read_to_string(&path).unwrap_or_default(); + let parsed: toml::Value = + toml::from_str(&text).unwrap_or_else(|_| toml::Value::Table(toml::value::Table::new())); + let mut data: toml::Value = if parsed.as_table().is_some() { + parsed + } else { + toml::Value::Table(toml::value::Table::new()) + }; + // 定位/创建 [workers] 表 + if data + .get("workers") + .and_then(toml::Value::as_table) + .is_none() + { + let t = data.as_table_mut().unwrap(); + t.insert( + "workers".into(), + toml::Value::Table(toml::value::Table::new()), + ); + } + let workers_tbl = data.get_mut("workers").unwrap().as_table_mut().unwrap(); + workers_tbl.insert("config".into(), json_to_toml(¤t)); + let out = toml::to_string(&data).map_err(|e| format!("序列化 config.toml 失败: {e}"))?; + // 原子写：临时文件 + rename + let tmp = path.with_extension("toml.tmp"); + std::fs::write(&tmp, out).map_err(|e| format!("写入配置失败: {e}"))?; + std::fs::rename(&tmp, &path).map_err(|e| format!("原子替换配置失败: {e}"))?; + Ok(current) +} + +/// json::Value → toml::Value +pub fn json_to_toml(v: &Value) -> toml::Value { + match v { + Value::Null => toml::Value::String(String::new()), + Value::Bool(b) => toml::Value::Boolean(*b), + Value::Number(n) => { + if let Some(i) = n.as_i64() { + toml::Value::Integer(i) + } else if let Some(f) = n.as_f64() { + toml::Value::Float(f) + } else { + toml::Value::String(n.to_string()) + } + } + Value::String(s) => toml::Value::String(s.clone()), + Value::Array(a) => toml::Value::Array(a.iter().map(json_to_toml).collect()), + Value::Object(m) => toml::Value::Table( + m.iter() + .map(|(k, v)| (k.clone(), json_to_toml(v))) + .collect(), + ), + } +} + +/// 基于 current 合并 updates（过滤保留键） +pub fn merge_confirmed(current: Value, updates: &Value) -> Value { + let mut obj = current.as_object().cloned().unwrap_or_default(); + for (k, v) in updates.as_object().unwrap_or(&Default::default()) { + if !CONFIG_FORBIDDEN_KEYS.contains(&k.as_str()) { + obj.insert(k.clone(), v.clone()); + } + } + Value::Object(obj) +} + +/// 计算 updates 中被允许（非保留键）的键个数 +pub fn count_allowed(updates: &Value) -> usize { + updates + .as_object() + .map(|m| { + m.keys() + .filter(|k| !CONFIG_FORBIDDEN_KEYS.contains(&k.as_str())) + .count() + }) + .unwrap_or(0) +} diff --git a/src/database.rs b/src/database.rs new file mode 100644 index 0000000..e30e08a --- /dev/null +++ b/src/database.rs @@ -0,0 +1,190 @@ +//! 数据库封装（对齐 Python `chongming_database.Database`） +//! +//! 提供 SQLite / PostgreSQL 两种后端的异步连接池 + 生命周期： +//! - `from_config`（读插件 `[workers.config].database`）或 `Database::sqlite / from_url` +//! - `initialize()` 建立连接池并（可选）建表 +//! - `pool()` 取共享连接池（`AnyPool` 内部 `Arc`，`Clone` 即派发引用）——各 handler 直接用它跑 sqlx 查询 +//! - `dispose()` 关闭连接池（对齐 Python `dispose`；幂等） +//! - `database_for(worker_type, config_dir)` 便捷入口：按插件配置直接构建（未配置返回 `None`， +//! 对齐 Python `db_name = None` 的降级语义） +//! +//! 后端差异： +//! - SQLite: 连接串 `sqlite:<路径>?mode=rwc`（`mode=rwc` 自动建库，同 aiosqlite 行为） +//! - PostgreSQL: `postgres://user:pwd@host:5432/db` +//! SQL 占位符遵循各自方言（SQLite `?`，PostgreSQL `$1`），AnyPool 不做 SQL 重写。 + +use std::path::Path; +use std::str::FromStr; +use std::sync::Arc; + +use serde_json::Value; +use sqlx::any::{Any, AnyConnectOptions, AnyPoolOptions}; +use sqlx::Pool; +use tokio::sync::Mutex; + +const DEFAULT_SQLITE_PATH: &str = "./data/plugin.db"; +const DEFAULT_MAX_CONNECTIONS: u32 = 5; + +/// 插件数据库句柄。可 `Clone` 给各 endpoint handler 共享（内部连接池线程安全）。 +#[derive(Clone)] +pub struct Database { + url: String, + pool: Arc>>>, +} + +impl std::fmt::Debug for Database { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Database").field("url", &self.url).finish_non_exhaustive() + } +} + +impl Database { + /// 按 SQLite 文件路径构建（对齐 Python `Database(Path(...))`）。 + /// `mode=rwc`：文件不存在则自动创建（对齐 aiosqlite 建库行为）。 + pub fn sqlite(db_path: &str) -> Self { + Self::new(format!("sqlite:{db_path}?mode=rwc")) + } + + /// 按 SQL 连接串构建（推荐 PostgreSQL: `postgres://...`）。 + pub fn from_url(url: &str) -> Self { + Self::new(url.trim().to_string()) + } + + fn new(url: String) -> Self { + Self { url, pool: Arc::new(Mutex::new(None)) } + } + + /// 从 `[workers.config].database` 配置段构建。 + /// + /// 支持的键: + /// type "sqlite"（默认） / "postgresql"（或 "postgres"） + /// database_path SQLite 文件路径（type=sqlite 时必填） + /// url PostgreSQL 连接 URL（type=postgresql 时必填） + pub fn from_config(cfg: &Value) -> Result { + let db_type = cfg.get("type").and_then(Value::as_str).unwrap_or("sqlite"); + match db_type { + "sqlite" => { + let p = cfg.get("database_path") + .and_then(Value::as_str) + .unwrap_or(DEFAULT_SQLITE_PATH); + Ok(Self::sqlite(p)) + } + "postgresql" | "postgres" => { + let url = cfg.get("url") + .and_then(Value::as_str) + .ok_or_else(|| "postgresql 数据库配置缺少 url".to_string())?; + Ok(Self::from_url(url)) + } + other => Err(format!("暂不支持的数据库类型: {other}")), + } + } + + /// 便捷入口：按插件配置目录（`config/workers/`）读取 `[workers.config].database` 构建。 + /// 未配置 database 段时返回 `None`（插件可自行降级，对齐 Python `db_name = None`）。 + pub fn database_for(worker_type: &str, config_dir: &str) -> Result, String> { + let plugin = crate::config::load_plugin_config(worker_type, config_dir)?; + let db_cfg = plugin.runtime.get("config").and_then(|c| c.get("database")); + match db_cfg { + None => Ok(None), + Some(cfg) => Ok(Some(Self::from_config(cfg)?)), + } + } + + /// 连接串（日志/诊断用） + pub fn url(&self) -> &str { + &self.url + } + + /// 是否已初始化（存在活跃连接池） + pub async fn initialized(&self) -> bool { + self.pool.lock().await.is_some() + } + + /// SQLite 文件连接串的父目录是否缺失（`sqlite:?params`）。 + /// SQLite 的 `mode=rwc` 只会建文件、不会建目录，缺目录时连接必失败， + /// 这里主动补齐父目录（对齐 Python aiosqlite 在开库前 `mkdir` 目录的行为）。 + fn ensure_sqlite_parent_dir(url: &str) -> Result<(), String> { + let path_part = url + .strip_prefix("sqlite:") + .and_then(|s| s.split('?').next()) + .unwrap_or_default(); + if path_part.is_empty() || path_part == ":memory:" { + return Ok(()); + } + let path = Path::new(path_part); + let dir = path + .parent() + .filter(|d| !d.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + if !dir.exists() { + std::fs::create_dir_all(dir) + .map_err(|e| format!("创建数据库目录失败（{}）: {e}", dir.display()))?; + } + Ok(()) + } + + /// 建立连接池（对齐 Python `initialize`）。幂等；可覆盖连接池大小。 + pub async fn initialize(&self, max_connections: u32) -> Result<(), String> { + if self.initialized().await { + return Ok(()); + } + // SQLite 建库前确保父目录存在（mode=rwc 不负责建目录） + if self.url.starts_with("sqlite:") { + Self::ensure_sqlite_parent_dir(&self.url)?; + } + // 装载当前编译进去的驱动（sqlite/postgres）；同一进程只安装一次 + sqlx::any::install_default_drivers(); + let opts = AnyConnectOptions::from_str(&self.url) + .map_err(|e| format!("解析数据库连接串失败（{}）: {e}", self.url))?; + let pool = AnyPoolOptions::new() + .max_connections(max_connections.max(1)) + .connect_with(opts) + .await + .map_err(|e| format!("连接数据库失败（{}）: {e}", self.url))?; + *self.pool.lock().await = Some(pool); + tracing::info!("[db] 数据库已初始化: {}", self.url); + Ok(()) + } + + /// 取共享连接池（需先 `initialize`）。`AnyPool` 内部共享引用，可 Clone 后到处传。 + pub async fn pool(&self) -> Result, String> { + self.pool.lock().await.clone() + .ok_or_else(|| "Database 未初始化（请先调用 initialize()）".to_string()) + } + + /// 直接执行无参数 SQL/DDL（建表等），对齐 Python `create_all` 前的建表步骤。 + pub async fn raw_execute(&self, sql: &str) -> Result<(), String> { + let pool = self.pool().await?; + sqlx::query(sql) + .execute(&pool) + .await + .map(|_| ()) + .map_err(|e| format!("执行 SQL 失败（{sql}）: {e}")) + } + + /// 关闭连接池（对齐 Python `dispose`；幂等）。 + pub async fn dispose(&self) { + let mut guard = self.pool.lock().await; + if let Some(p) = guard.take() { + p.close().await; + } + tracing::info!("[db] 数据库已关闭: {}", self.url); + } +} + +/// 便捷初始化入口：从插件配置构建并按需建表。返回 `None` 表示未配置 database 段。 +/// `create_table_sql` 为可选的建表语句（如 `CREATE TABLE IF NOT EXISTS t(...)`）。 +pub async fn init_database( + worker_type: &str, + config_dir: &str, + create_table_sql: Option<&str>, +) -> Result, String> { + let Some(db) = Database::database_for(worker_type, config_dir)? else { + return Ok(None); + }; + db.initialize(DEFAULT_MAX_CONNECTIONS).await?; + if let Some(sql) = create_table_sql { + db.raw_execute(sql).await?; + } + Ok(Some(db)) +} \ No newline at end of file diff --git a/src/descriptor.rs b/src/descriptor.rs new file mode 100644 index 0000000..d62cf03 --- /dev/null +++ b/src/descriptor.rs @@ -0,0 +1,157 @@ +//! 路由描述符（docs/plugin-wire-protocol.md §2） +//! +//! 描述符来源为 `config/workers//*.toml` 端点定义。解析规则与 Python SDK +//! route_config.py 完全对位（TOML 兼容 dict），保证与 Python 插件发布的内容 +//! 逐字段一致 —— 网关按描述符内容 hash 做多副本幂等去重与冲突判定， +//! 混布部署时内容不一致会触发告警。 + +use std::collections::BTreeMap; +use std::path::Path; + +use serde_json::{Map, Value}; + +pub const ROUTES_PREFIX: &str = "chongming/routes"; +pub const META_KEY: &str = "__meta__"; +pub const SCHEMA_VERSION: u8 = 1; + +/// route_id = subject 的 "." 替换为 "_" + "__" + METHOD 大写 +/// （与 Python routes_publisher.route_id 一致） +pub fn route_id(subject: &str, method: &str) -> String { + format!("{}__{}", subject.replace('.', "_"), method.to_uppercase()) +} + +/// 从 `config/workers//` 目录解析全部路由描述符（TOML 兼容 dict）。 +/// +/// 规则（对齐 Python load_worker_routes_from_dir）： +/// - 每个端点 TOML 的每个方法节（GET/POST/MQTT...）生成一个独立描述符 +/// - `[worker.defaults]` 作为字段缺省值合并（显式字段优先） +/// - method 恒取节键大写；path 缺省为 `router_prefix + "/" + 端点名`； +/// subject 缺省为 `.<端点名>`；route_name 记录端点文件名 +pub fn load_routes_from_dir( + worker_type: &str, + config_dir: &str, +) -> Result, String> { + let base = Path::new(config_dir).join(worker_type); + if !base.is_dir() { + return Err(format!("Worker 配置目录不存在: {}", base.display())); + } + + // 1. 分组配置：[worker].defaults + [worker].registration.router_prefix + let group_path = base.join("config.toml"); + let mut group_defaults: Map = Map::new(); + let mut router_prefix = String::new(); + if group_path.exists() { + let text = std::fs::read_to_string(&group_path) + .map_err(|e| format!("读取 {} 失败: {e}", group_path.display()))?; + let cfg: toml::Value = toml::from_str(&text) + .map_err(|e| format!("解析 {} 失败: {e}", group_path.display()))?; + if let Some(defaults) = cfg + .get("worker") + .and_then(|w| w.get("defaults")) + .and_then(toml::Value::as_table) + { + group_defaults = defaults + .iter() + .map(|(k, v)| (k.clone(), crate::config::toml_to_json(v))) + .collect(); + } + if let Some(prefix) = cfg + .get("worker") + .and_then(|w| w.get("registration")) + .and_then(|r| r.get("router_prefix")) + .and_then(toml::Value::as_str) + { + router_prefix = prefix.to_string(); + } + } + + // 2. 遍历所有端点 TOML（config.toml 除外） + let mut routes: Vec = Vec::new(); + let mut files: Vec = std::fs::read_dir(&base) + .map_err(|e| format!("读取目录 {} 失败: {e}", base.display()))? + .filter_map(|e| e.ok()) + .filter(|e| { + e.file_name() + .to_string_lossy() + .ends_with(".toml") + }) + .map(|e| e.file_name().to_string_lossy().to_string()) + .collect(); + files.sort(); + + for file in files { + if file == "config.toml" { + continue; + } + // 端点名 = 文件名（不含扩展名） + let route_name = file.trim_end_matches(".toml").to_string(); + let base_path = format!("{}/{}", router_prefix, route_name); + + let text = std::fs::read_to_string(base.join(&file)) + .map_err(|e| format!("读取 {} 失败: {e}", file))?; + let route_data: toml::Value = toml::from_str(&text) + .map_err(|e| format!("解析 {} 失败: {e}", file))?; + let route_table = match route_data.as_table() { + Some(t) => t, + None => continue, + }; + + // 3. 每个方法节生成一个独立路由 + for (method, spec) in route_table { + let spec_table = match spec.as_table() { + Some(t) => t, + None => continue, + }; + let mut route: Map = Map::new(); + // 显式字段优先，defaults 兜底 + for (k, v) in spec_table { + route.insert(k.clone(), crate::config::toml_to_json(v)); + } + for (k, v) in &group_defaults { + if !route.contains_key(k) { + route.insert(k.clone(), v.clone()); + } + } + // 覆盖关键字段（与 Python 一致） + route.insert("method".into(), Value::String(method.to_uppercase())); + if !route.contains_key("path") { + route.insert("path".into(), Value::String(base_path.clone())); + } + if !route.contains_key("subject") { + route.insert( + "subject".into(), + Value::String(format!("{worker_type}.{route_name}")), + ); + } + route.insert("route_name".into(), Value::String(route_name.clone())); + routes.push(Value::Object(route)); + } + } + + Ok(routes) +} + +/// 构造发布用描述符：附加 schema_version 与 type（对齐 Python publish_routes） +pub fn finalize_descriptor(route: &Value, worker_type: &str) -> Value { + let mut map = match route.as_object() { + Some(m) => m.clone(), + None => return route.clone(), + }; + map.insert("schema_version".into(), Value::from(SCHEMA_VERSION)); + map.insert("type".into(), Value::String(worker_type.to_string())); + Value::Object(map) +} + +/// 元信息描述符（[[workers.mqtt_topics]] 时发布） +pub fn meta_descriptor(mqtt_topics: &[Value]) -> Value { + let mut map = BTreeMap::new(); + map.insert( + "schema_version".to_string(), + Value::from(SCHEMA_VERSION), + ); + map.insert( + "mqtt_topics".to_string(), + Value::Array(mqtt_topics.to_vec()), + ); + Value::Object(map.into_iter().collect()) +} diff --git a/src/envelope.rs b/src/envelope.rs new file mode 100644 index 0000000..8b40237 --- /dev/null +++ b/src/envelope.rs @@ -0,0 +1,59 @@ +//! 请求信封协议（docs/plugin-wire-protocol.md §4.2） +//! +//! 请求载荷: `{"meta": {"method": "GET", "user_info": {...}}, "data": {...业务参数...}}` +//! 非信封载荷（不含 meta/data 键）整包视为业务 data，method 默认 GET（兼容旧调用方）。 + +use serde_json::Value; + +/// 解析后的请求 +#[derive(Debug, Clone)] +pub struct Request { + /// 信封 meta 中的 HTTP 方法（None = 非信封载荷，按 GET 处理） + pub method: Option, + /// 网关注入的认证用户信息（auth_required 时存在） + pub user_info: Option, + /// 业务参数（信封的 data；非信封载荷时为整包） + pub data: Value, + /// 处理该请求的 Worker 副本 ID（由 serve 层注入，对应 liveliness id） + pub worker_id: String, +} + +impl Request { + /// 解析请求载荷（UTF-8 JSON） + pub fn parse(payload: &[u8]) -> Request { + let raw: Value = serde_json::from_slice(payload).unwrap_or(Value::Null); + + if let Value::Object(map) = &raw { + if map.contains_key("meta") || map.contains_key("data") { + let meta = map.get("meta").cloned().unwrap_or(Value::Null); + let method = meta + .get("method") + .and_then(Value::as_str) + .map(|s| s.to_uppercase()); + let user_info = meta + .get("user_info") + .cloned() + .filter(|v| !v.is_null()); + let data = map.get("data").cloned().unwrap_or(Value::Null); + return Request { + method, + user_info, + data, + worker_id: String::new(), + }; + } + } + // 非信封载荷：整包视为业务数据，method 默认 GET + Request { + method: None, + user_info: None, + data: raw, + worker_id: String::new(), + } + } + + /// 分派用 HTTP 方法（缺失默认 GET，与 Python runtime.py 一致） + pub fn method(&self) -> String { + self.method.clone().unwrap_or_else(|| "GET".to_string()) + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..0c075f6 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,19 @@ +//! 错误响应格式（与 Python SDK runtime.py 完全对位） + +use serde_json::{json, Value}; + +/// 非流式端点异常回复（Python runtime.py: `{"result":0.0,"operation":"error", +/// "error":..., "worker":...}`） +pub fn error_response(worker_id: &str, msg: &str) -> Value { + json!({ + "result": 0.0, + "operation": "error", + "error": msg, + "worker": worker_id, + }) +} + +/// 流式端点异常收尾块（`{"error": ...}`，网关收到含 error 键的块即终止流） +pub fn stream_error_block(msg: &str) -> Value { + json!({"error": msg}) +} diff --git a/src/gen.rs b/src/gen.rs new file mode 100644 index 0000000..d1f9100 --- /dev/null +++ b/src/gen.rs @@ -0,0 +1,222 @@ +//! 模型生成（对齐 Python gen_models：TOML 端点定义是模型的唯一事实来源） +//! +//! 从 `config/workers//*.toml` 的 params / response_model 声明生成 Rust +//! 结构体（serde derive），命名约定与 Python 完全一致： +//! `` +//! 如 hello + GET → `HelloGetRequest` / `HelloGetResponse`。 +//! +//! 一致性保证：生成文件头部写入 `FINGERPRINT`（当前端点 TOML 内容 hash）， +//! 插件启动时用 [`fingerprint`] 重算并比对——TOML 变更而未重新生成模型即告警。 + +use std::collections::BTreeMap; + +use serde_json::{Map, Value}; + +use crate::descriptor::load_routes_from_dir; + +/// 计算某插件类型当前端点定义的指纹（生成与校验同源，确保模型不漂移） +pub fn fingerprint(worker_type: &str, config_dir: &str) -> Result { + let routes = load_routes_from_dir(worker_type, config_dir)?; + // serde_json Map 保持插入序；同一 TOML 解析插入序稳定，可作指纹 + let canonical = serde_json::to_string(&routes).map_err(|e| e.to_string())?; + Ok(sha1_hex(canonical.as_bytes())) +} + +fn sha1_hex(data: &[u8]) -> String { + use sha1::Digest; + let digest = sha1::Sha1::digest(data); + digest.iter().map(|b| format!("{b:02x}")).collect() +} + +/// 生成单个插件类型的完整模型模块源码 +pub fn generate_type(worker_type: &str, config_dir: &str) -> Result { + let routes = load_routes_from_dir(worker_type, config_dir)?; + let fp = fingerprint(worker_type, config_dir)?; + + let mut structs: Vec = Vec::new(); + for r in &routes { + let route_name = r.get("route_name").and_then(Value::as_str).unwrap_or(""); + let method = r.get("method").and_then(Value::as_str).unwrap_or("GET"); + let prefix = model_base(route_name, method); + + // 请求模型（声明了 params 才生成） + if let Some(params) = r.get("params").and_then(Value::as_array) { + if !params.is_empty() { + structs.push(render_request_struct(&format!("{prefix}Request"), params)?); + } + } + // 响应模型（声明了 response_model 才生成） + if let Some(resp) = r.get("response_model") { + if let Value::Object(fields) = resp { + if !fields.is_empty() { + structs.push(render_response_struct(&format!("{prefix}Response"), fields)?); + } + } + } + } + + if structs.is_empty() { + return Ok(format!( + "pub mod {worker_type} {{\n pub const FINGERPRINT: &str = \"{fp}\";\n}}\n" + )); + } + Ok(format!( + "pub mod {worker_type} {{\n #![allow(dead_code)]\n pub const FINGERPRINT: &str = \"{fp}\";\n{}\n}}", + structs.join("\n") + )) +} + +/// 模型基名：（对齐 Python model_name） +fn model_base(route_name: &str, method: &str) -> String { + let base: String = route_name + .split(|c: char| c == '-' || c == '_' || c == ' ') + .filter(|s| !s.is_empty()) + .map(|s| { + let mut c = s.chars(); + match c.next() { + Some(f) => f.to_uppercase().collect::() + c.as_str(), + None => String::new(), + } + }) + .collect(); + let m = if method.eq_ignore_ascii_case("MQTT") { + "Mqtt".to_string() + } else { + // title-case：GET → Get（对齐 Python str.title()） + let lower = method.to_lowercase(); + let mut c = lower.chars(); + match c.next() { + Some(f) => f.to_uppercase().collect::() + c.as_str(), + None => String::new(), + } + }; + format!("{base}{m}") +} + +/// 类型映射（对齐网关 route_generator 的 TYPE_MAP + list/dict 泛型） +fn rust_type(t: &str) -> String { + let t = t.trim(); + for p in ["list[", "List["] { + if let Some(inner) = t.strip_prefix(p) { + let inner = inner.trim_end_matches(']').trim(); + return format!("Vec<{}>", rust_type(inner)); + } + } + for p in ["dict[", "Dict["] { + if let Some(inner) = t.strip_prefix(p) { + let inner = inner.trim_end_matches(']').trim(); + if let Some((k, v)) = inner.split_once(',') { + return format!( + "std::collections::HashMap<{}, {}>", + rust_type(k), + rust_type(v) + ); + } + return format!("std::collections::HashMap", rust_type(inner)); + } + } + match t { + "str" => "String".to_string(), + "int" => "i64".to_string(), + "float" => "f64".to_string(), + "bool" => "bool".to_string(), + "any" => "serde_json::Value".to_string(), + other => format!("serde_json::Value /* 未识别类型: {other} */"), + } +} + +/// 字段名合法化：非法 Rust 标识符时用 raw identifier，并加 serde rename 保持线协议原名 +fn field_ident(name: &str) -> (String, bool) { + if name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + && !name.is_empty() + { + (name.to_string(), false) + } else { + (format!("r#{}", name.replace('-', "_")), true) + } +} + +/// 请求模型：params 全部必填（对齐网关 generate_request_model） +fn render_request_struct(name: &str, params: &[Value]) -> Result { + let mut fields: Vec = Vec::new(); + for p in params { + let spec = p.as_str().unwrap_or(""); + let (field, ty) = spec + .split_once(':') + .map(|(f, t)| (f.trim(), t.trim())) + .unwrap_or((spec.trim(), "str")); + let (ident, rename) = field_ident(field); + let attr = if rename { + format!(" #[serde(rename = \"{field}\")]\n") + } else { + String::new() + }; + fields.push(format!("{attr} pub {ident}: {},", rust_type(ty))); + } + Ok(format!( + " #[derive(Debug, Clone, serde::Deserialize)]\n pub struct {name} {{\n{}\n }}", + fields.join("\n") + )) +} + +/// 响应模型：required=false → Option + skip 空值（对齐网关 generate_response_model） +fn render_response_struct(name: &str, fields: &Map) -> Result { + let mut out: Vec = Vec::new(); + for (field, spec) in fields { + let (ident, rename) = field_ident(field); + let rename_attr = if rename { + format!(" #[serde(rename = \"{field}\")]\n") + } else { + String::new() + }; + if let Value::Object(spec_map) = spec { + let ty = spec_map + .get("type") + .and_then(Value::as_str) + .unwrap_or("str"); + let required = spec_map + .get("required") + .and_then(Value::as_bool) + .unwrap_or(true); + if required { + out.push(format!( + "{rename_attr} pub {ident}: {},", + rust_type(ty) + )); + } else { + out.push(format!( + "{rename_attr} #[serde(skip_serializing_if = \"Option::is_none\")]\n pub {ident}: Option<{}>,", + rust_type(ty) + )); + } + } else { + // 简写形式：{"field": "type"} + let ty = spec.as_str().unwrap_or("str"); + out.push(format!("{rename_attr} pub {ident}: {},", rust_type(ty))); + } + } + Ok(format!( + " #[derive(Debug, Clone, serde::Serialize)]\n pub struct {name} {{\n{}\n }}", + out.join("\n") + )) +} + +/// 生成多个插件类型到一个文件（含文件头说明） +pub fn generate_all(types: &[&str], config_dir: &str) -> Result { + let mut body = String::from( + "// 自动生成 —— 勿手动修改（chongming-worker-rs gen_models）\n\ + // 唯一事实来源: config/workers//*.toml 端点定义\n\ + // 重新生成: cargo run --bin gen_models -- --types --out \n\n", + ); + let mut mods: BTreeMap = BTreeMap::new(); + for t in types { + mods.insert((*t).to_string(), generate_type(t, config_dir)?); + } + for (_name, src) in mods { + body.push_str(&src); + body.push('\n'); + } + Ok(body) +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..9dfef6b --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,106 @@ +//! chongming-worker-rs —— 进程插件 Rust SDK +//! +//! 语言无关线协议（docs/plugin-wire-protocol.md）的 Rust 实现： +//! 1. 解析 config.toml 与端点 TOML，构造路由描述符并发布到 `chongming/routes/**` +//! 2. liveliness token 注册 `worker//`（服务发现） +//! 3. Queryable 声明 `worker///`，解析请求信封并分派 handler +//! 4. 每 10s 重发布描述符（自愈），SIGTERM/SIGINT 优雅退出时 DELETE 摘除 +//! +//! 与 Python SDK（utils/worker/）行为等价，可同 type 混布参与同一套负载均衡。 + +pub mod admin; +pub mod cache; +pub mod config; +#[cfg(feature = "database")] +pub mod database; +pub mod descriptor; +pub mod envelope; +pub mod error; +pub mod gen; +pub mod publisher; +pub mod runtime; +pub mod serve; + +pub use runtime::WorkerPool; +pub use serve::{serve, ServeOptions}; + +/// 常用类型统一导出 +pub mod prelude { + pub use crate::envelope::Request; + pub use crate::error::{error_response, stream_error_block}; + pub use crate::publisher::{announce_loop, RoutePublisher, ANNOUNCE_INTERVAL_SECS}; + #[cfg(feature = "database")] + pub use crate::database::Database; + pub use crate::{HandlerRegistry, Reply}; +} + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use futures_util::stream::BoxStream; +use serde_json::Value; + +pub use envelope::Request; + +/// 异步函数指针（handler 与 announce 循环等共用） +pub type BoxFuture = Pin + Send>>; + +/// 端点响应：单块 JSON，或流式多块（SSE 逐块转发） +pub enum Reply { + /// 普通响应：单个回复，载荷为业务结果 JSON 对象 + Json(Value), + /// 流式响应：逐块多回复；每块 `Ok(v)` 作为独立回复块， + /// `Err(msg)` 以 `{"error": msg}` 收尾块结束流（与 Python SDK 一致） + Stream(BoxStream<'static, Result>), +} + +/// 端点 handler：解析后的请求信封 → 业务结果 +pub type Handler = dyn Fn(Request) -> BoxFuture> + Send + Sync; + +/// RPC 端点注册表（按 subject + HTTP 方法分派） +#[derive(Default)] +pub struct HandlerRegistry { + /// subject -> [(method, handler)]（同一 subject 可挂多个方法） + map: HashMap)>>, +} + +impl HandlerRegistry { + pub fn new() -> Self { + Self::default() + } + + /// 注册端点。`f` 为 async handler：`fn(Request) -> Result`。 + pub fn register(&mut self, subject: &str, method: &str, f: F) + where + F: Fn(Request) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let entry = self + .map + .entry(subject.to_string()) + .or_default(); + entry.push(( + method.to_uppercase(), + Arc::new(move |req: Request| { + Box::pin(f(req)) as Pin> + Send>> + }) as Arc, + )); + } + + /// 该 subject+method 是否有 handler（发布前的契约自检用） + pub fn has(&self, subject: &str, method: &str) -> bool { + self.lookup(subject, method).is_some() + } + + /// 按 subject+method 查 handler；无精确匹配时降级到该 subject 的第一个方法 + /// （与 Python runtime.py 的匹配逻辑一致）。 + pub fn lookup(&self, subject: &str, method: &str) -> Option> { + let methods = self.map.get(subject)?; + if let Some((_, h)) = methods.iter().find(|(m, _)| m == &method.to_uppercase()) { + return Some(h.clone()); + } + methods.first().map(|(_, h)| h.clone()) + } +} diff --git a/src/publisher.rs b/src/publisher.rs new file mode 100644 index 0000000..96e5617 --- /dev/null +++ b/src/publisher.rs @@ -0,0 +1,141 @@ +//! 路由描述符发布器（对齐 Python routes_publisher.py） +//! +//! PUT 到 `chongming/routes//...`，优雅退出时 DELETE 摘除。 +//! 调用方每 10s 周期重发布（自愈）。 + +use std::sync::Mutex; + +use serde_json::Value; +use tokio::sync::RwLock; +use zenoh::Session; + +use crate::descriptor::{meta_descriptor, route_id, ROUTES_PREFIX}; + +pub struct RoutePublisher { + session: Session, + worker_type: String, + /// 已发布 key（幂等去重后用于摘除） + keys: Mutex>, + /// 重发布并发保护（防止两个 publish 交叠乱序） + _lock: RwLock<()>, +} + +impl RoutePublisher { + pub fn new(session: Session, worker_type: &str) -> Self { + Self { + session, + worker_type: worker_type.to_string(), + keys: Mutex::new(Vec::new()), + _lock: RwLock::new(()), + } + } + + /// 发布全部路由描述符（可重复调用，幂等）。 + /// `routes` 为 load_routes_from_dir 的原始描述符，此处附加 schema_version/type。 + pub async fn publish_routes( + &self, + routes: &[Value], + mqtt_topics: &[Value], + ) -> Result<(), String> { + let _g = self._lock.write().await; + + // 1. 收集本次要发布的 (key, 描述符)（先 await 后更新 keys，避免持锁跨 await） + let mut entries: Vec<(String, Value)> = Vec::new(); + for r in routes { + let subject = r.get("subject").and_then(Value::as_str).unwrap_or(""); + if subject.is_empty() { + tracing::warn!("跳过无 subject 的路由描述符: {r}"); + continue; + } + let method = r.get("method").and_then(Value::as_str).unwrap_or("GET"); + let descriptor = crate::descriptor::finalize_descriptor(r, &self.worker_type); + let key = format!( + "{}/{}/{}", + ROUTES_PREFIX, + self.worker_type, + route_id(subject, method), + ); + entries.push((key, descriptor)); + tracing::info!( + "已发布路由: {} {} (subject={})", + method.to_uppercase(), + r.get("path") + .and_then(serde_json::Value::as_str) + .unwrap_or("?"), + subject, + ); + } + if !mqtt_topics.is_empty() { + let meta_key = format!("{}/{}/{}", ROUTES_PREFIX, self.worker_type, "__meta__"); + entries.push((meta_key, meta_descriptor(mqtt_topics))); + tracing::info!("已发布插件元信息: {} 个自定义 MQTT 主题", mqtt_topics.len()); + } + + // 2. 发布 + for (key, descriptor) in &entries { + self.put(key, descriptor).await?; + } + + // 3. 记录已发布 key（幂等去重） + { + let mut keys = self.keys.lock().unwrap(); + for (key, _) in &entries { + if !keys.contains(key) { + keys.push(key.clone()); + } + } + } + + tracing::info!( + "[{}] 路由发布完成: {} 端点 + {}", + self.worker_type, + entries.len(), + if mqtt_topics.is_empty() { "无元信息" } else { "元信息" }, + ); + Ok(()) + } + + /// 摘除全部已发布的路由描述符（优雅退出时调用） + pub async fn unpublish(&self) { + let keys = { + let mut k = self.keys.lock().unwrap(); + std::mem::take(&mut *k) + }; + for key in &keys { + if let Err(e) = self.session.delete(key.clone()).await { + tracing::warn!("摘除路由失败 [{key}]: {e}"); + } + } + tracing::info!("[{}] 已摘除 {} 个路由描述符", self.worker_type, keys.len()); + } + + async fn put(&self, key: &str, payload: &Value) -> Result<(), String> { + let payload_str = serde_json::to_string(payload) + .map_err(|e| format!("序列化描述符失败: {e}"))?; + self.session + .put(key.to_string(), payload_str) + .await + .map_err(|e| format!("发布路由 {key} 失败: {e}")) + } +} + +/// 描述符周期性重发布间隔（秒，与 Python ANNOUNCE_INTERVAL 一致） +pub const ANNOUNCE_INTERVAL_SECS: u64 = 10; + +/// 描述符周期性重发布（网关后启动/样本丢失时自愈）。 +/// 由 serve 层 spawn，进程退出时 abort。 +pub async fn announce_loop( + publisher: std::sync::Arc, + routes: Vec, + mqtt_topics: Vec, +) { + let mut interval = + tokio::time::interval(std::time::Duration::from_secs(ANNOUNCE_INTERVAL_SECS)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + interval.tick().await; + if let Err(e) = publisher.publish_routes(&routes, &mqtt_topics).await { + tracing::error!("路由重发布失败: {e}"); + } + } +} diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 0000000..d5b4b97 --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,499 @@ +//! 进程插件运行时 — 托管副本生命周期（对齐 Python runtime.py 的 +//! `WorkerController` / `create_worker` / `create_workers` / WorkerPool） +//! +//! 每个副本槽位（index）由一个 `ReplicaController` 串行驱动多次运行： +//! - restart → 以递增 epoch 的新 worker_id 重建该副本 +//! - stop / drain / config_set → 结束该次数运行并移交管理循环决定是否重建 +//! 支持 per-slot 管理（`__admin__` Queryable）、优雅排空（隔离新请求 → 等在途完成）。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; + +use futures_util::StreamExt; +use serde_json::Value; +use tokio::sync::Notify; +use zenoh::Session; + +use crate::admin; +use crate::cache::Cache; +use crate::envelope::Request; +use crate::error::{error_response, stream_error_block}; +use crate::{HandlerRegistry, Reply}; + +/// 管理指令常量（与 Python runtime.py 一致） +pub const ADMIN_RESTART: &str = "restart"; +pub const ADMIN_STOP: &str = "stop"; +pub const ADMIN_DRAIN: &str = "drain"; +pub const ADMIN_GET_CONFIG: &str = "get_config"; +pub const ADMIN_CONFIG_SET: &str = "config_set"; + +/// 优雅排空等待在途请求的最大时长（秒），超时强制关闭 +pub const DRAIN_INFLIGHT_TIMEOUT: u64 = 30; + +/// 单个副本槽位的跨重启控制对象（对齐 Python WorkerController） +pub struct ReplicaController { + pub worker_type: String, + pub config_dir: String, + /// 副本槽位序号（0..count） + pub index: usize, + /// 重启代数（make_worker_id 自增） + epoch: AtomicI64, + /// 当前副本被要求退出时的原因（"" = 正常/被取消；restart/stop） + exit_request: Mutex, + /// 当前运行副本 id + current_worker_id: Mutex, + /// 优雅排空标记 + pub draining: AtomicBool, + /// 副本主阻塞（admin 指令 / 池缩容 / 进程关闭唤醒） + pub control: Notify, + /// config_set 设置的待应用新配置（重启时生效） + pub pending_config: Mutex>, + /// 启动时解析的基础运行时配置（[workers.config] 等） + pub base_config: Value, + /// 该类型的配置缓存句柄 + pub cache: Cache, + /// 进程实例 id（多实例协同） + pub instance_id: String, + /// 在途请求计数（优雅排空用，各业务 Queryable 共享） + pub inflight: Arc, +} + +impl ReplicaController { + /// 创建槽位控制器（epoch 从 1 开始） + pub fn new( + worker_type: &str, + config_dir: &str, + index: usize, + base_config: Value, + cache: Cache, + instance_id: String, + ) -> Self { + Self { + worker_type: worker_type.to_string(), + config_dir: config_dir.to_string(), + index, + epoch: AtomicI64::new(0), + exit_request: Mutex::new(String::new()), + current_worker_id: Mutex::new(String::new()), + draining: AtomicBool::new(false), + control: Notify::new(), + pending_config: Mutex::new(None), + base_config, + cache, + instance_id, + inflight: Arc::new(AtomicI64::new(0)), + } + } + + /// 生成下一个副本 id（epoch 自增，保证重启后 id 变化） + pub fn make_worker_id(&self) -> String { + let epoch = self.epoch.fetch_add(1, Ordering::SeqCst) + 1; + format!("worker-{}-{}-{}", self.index, epoch, std::process::id()) + } + + pub fn worker_id(&self) -> String { + self.current_worker_id.lock().unwrap().clone() + } + + pub fn exit_request(&self) -> String { + self.exit_request.lock().unwrap().clone() + } + + pub fn epoch(&self) -> i64 { + self.epoch.load(Ordering::SeqCst) + } + + /// 记录退出原因并唤醒副本主阻塞 + pub fn signal(&self, reason: &str) { + *self.exit_request.lock().unwrap() = reason.to_string(); + self.control.notify_one(); + } + + /// 请求重启（以新 id 重建） + pub fn request_restart(&self) { + self.draining.store(false, Ordering::Relaxed); + self.signal(ADMIN_RESTART); + } + + /// 请求停止（不重建） + pub fn request_stop(&self) { + self.draining.store(false, Ordering::Relaxed); + self.signal(ADMIN_STOP); + } + + /// 请求优雅排空（隔离新请求，等在途完成，不重建） + pub fn request_drain(&self) { + self.draining.store(true, Ordering::Relaxed); + self.signal(ADMIN_STOP); + } +} + +/// 在途请求归零等待（对齐 Python _wait_inflight_zero） +async fn wait_inflight_zero(inflight: &AtomicI64) { + let deadline = + tokio::time::Instant::now() + tokio::time::Duration::from_secs(DRAIN_INFLIGHT_TIMEOUT); + loop { + let remain = inflight.load(Ordering::Relaxed); + if remain <= 0 { + tracing::info!("[drain] 在途请求已清空，完成排空"); + return; + } + if tokio::time::Instant::now() >= deadline { + tracing::warn!("[drain] 优雅排空超时（in-flight={remain}），强制关闭"); + return; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } +} + +/// 单个副本单次运行：声明 liveliness + 业务 Queryable + admin Queryable， +/// 等待管理唤醒 →（drain 时等在途）→ 清理退出。 +async fn run_replica( + session: Session, + controller: Arc, + registry: Arc, + routes: &[Value], + pool: Arc>, +) -> Result<(), String> { + let worker_id = controller.worker_id(); + let worker_type = controller.worker_type.clone(); + tracing::info!("[{worker_type}] Worker {worker_id} 启动"); + + // 1. 存活令牌 + let token = session + .liveliness() + .declare_token(format!("worker/{worker_type}/{worker_id}")) + .await + .map_err(|e| format!("声明 liveliness token 失败: {e}"))?; + + let mut tasks: Vec> = Vec::new(); + + // 2. 业务 Queryable（逐 subject） + for r in routes { + let subject = r.get("subject").and_then(Value::as_str).unwrap_or("").to_string(); + if subject.is_empty() { + continue; + } + let key = format!("worker/{worker_type}/{worker_id}/{subject}"); + let h = tokio::spawn(queryable_loop( + session.clone(), + key, + registry.clone(), + worker_id.clone(), + controller.inflight.clone(), + )); + tasks.push(h); + } + + // 3. admin Queryable（管理指令） + let admin_key = format!("worker/{worker_type}/{worker_id}/__admin__"); + tasks.push(tokio::spawn(admin_loop( + session.clone(), + admin_key, + controller.clone(), + pool, + ))); + tracing::info!("[{worker_type}] Worker {worker_id} 就绪（{count} 业务 + 1 admin Queryable）", + count = routes.len()); + + // 4. 等待管理唤醒 + controller.control.notified().await; + let exit_reason = controller.exit_request(); + + // 5. 优雅排空：摘除 liveliness/business/admin（隔绝新请求）后才等在途 + // （先 abort 处理循环让 queryable 释放，再等 inflight 归零） + for t in &tasks { + t.abort(); + } + for t in tasks { + let _ = t.await; + } + drop(token); // undeclare liveliness + + if controller.draining.load(Ordering::Relaxed) { + wait_inflight_zero(&controller.inflight).await; + } + tracing::info!("[{worker_type}] Worker {worker_id} 已关闭（exit={exit_reason}）"); + Ok(()) +} + +/// 单个副本槽位的管理循环：串行驱动多次运行，支持自愈重启。 +/// 返回时该槽位已结束（stop / 正常退出）。 +async fn managed_replica( + session: Session, + controller: Arc, + registry: Arc, + routes: Vec, + pool: Arc>, +) { + loop { + *controller.exit_request.lock().unwrap() = String::new(); + controller.draining.store(false, Ordering::Relaxed); + controller.inflight.store(0, Ordering::Relaxed); + let worker_id = controller.make_worker_id(); + *controller.current_worker_id.lock().unwrap() = worker_id.clone(); + tracing::info!( + "[{}] 槽位 #{} 启动管理循环（id={}）", + controller.worker_type, controller.index, worker_id, + ); + + let run_config = controller.pending_config.lock().unwrap().take() + .unwrap_or_else(|| controller.base_config.clone()); + // 新配置会随 epoll 重启/后续 run_replica 由调用方读取；此处仅记日志 + tracing::trace!("[{}] 槽位 #{} 运行配置: {run_config}", controller.worker_type, controller.index); + + if let Err(e) = run_replica( + session.clone(), + controller.clone(), + registry.clone(), + &routes, + pool.clone(), + ) + .await + { + tracing::error!("[{}] 槽位 #{} 运行异常: {e}", controller.worker_type, controller.index); + } + + let reason = controller.exit_request(); + if reason == ADMIN_RESTART { + tracing::info!( + "[{}] 槽位 #{} 收到重启指令，以新 id 重建", + controller.worker_type, controller.index, + ); + continue; + } + tracing::info!( + "[{}] 槽位 #{} 退出（{}）", + controller.worker_type, controller.index, if reason.is_empty() { "normal" } else { &reason }, + ); + break; + } +} + +/// 业务 Queryable 请求处理循环 +async fn queryable_loop( + session: Session, + key: String, + registry: Arc, + worker_id: String, + inflight: Arc, +) { + let subject = key.rsplit('/').next().unwrap_or("").to_string(); + let Ok(queryable) = session.declare_queryable(key.to_string()).await else { + tracing::error!("声明 Queryable {key} 失败"); + return; + }; + tracing::info!("[{worker_id}] Queryable: '{subject}'"); + while let Ok(query) = queryable.recv_async().await { + inflight.fetch_add(1, Ordering::SeqCst); + handle_query(query, registry.clone(), subject.clone(), worker_id.clone()).await; + inflight.fetch_sub(1, Ordering::SeqCst); + } +} + +/// 处理单个业务 Query（普通 / 流式 / 错误） +async fn handle_query( + query: zenoh::query::Query, + registry: Arc, + subject: String, + worker_id: String, +) { + let key_expr = query.key_expr().to_string(); + let payload = query + .payload() + .map(|p| p.to_bytes().to_vec()) + .unwrap_or_default(); + let mut request = Request::parse(&payload); + request.worker_id = worker_id.clone(); + + let method = request.method(); + let Some(handler) = registry.lookup(&subject, &method) else { + let _ = query + .reply_err(stream_error_block(&format!("No method for {subject}")).to_string()) + .await; + return; + }; + + let result = handler(request).await; + match result { + Ok(Reply::Json(v)) => { + if let Err(e) = query.reply(key_expr.clone(), v.to_string()).await { + tracing::debug!("回复失败（查询可能已终止）: {e}"); + } + } + Ok(Reply::Stream(mut stream)) => { + while let Some(chunk) = stream.next().await { + let (payload, is_error) = match chunk { + Ok(v) => (v.to_string(), false), + Err(msg) => { + tracing::error!("[{worker_id}] 流式生成器异常，发送错误块后停止: {msg}"); + (stream_error_block(&msg).to_string(), true) + } + }; + if query.reply(key_expr.clone(), payload).await.is_err() { + break; + } + if is_error { + break; + } + } + } + Err(msg) => { + tracing::error!("[{worker_id}] {subject} 异常: {msg}"); + // 非流式端点错误：以正常 reply 携带业务错误 JSON（对位 Python runtime.py）， + // 而非 Zenoh 错误回复——网关据此按业务错误处理（result=0/operation=error）。 + let resp = error_response(&worker_id, &msg); + if let Err(e) = query.reply(key_expr.clone(), resp.to_string()).await { + tracing::debug!("错误回复失败（查询可能已终止）: {e}"); + } + } + } +} + +/// admin Queryable 请求处理循环 +async fn admin_loop( + session: Session, + key: String, + controller: Arc, + pool: Arc>, +) { + let Ok(queryable) = session.declare_queryable(key.to_string()).await else { + tracing::error!("声明 admin Queryable {key} 失败"); + return; + }; + tracing::info!("[{}] admin Queryable: '{key}'", controller.worker_type); + while let Ok(query) = queryable.recv_async().await { + admin::handle_admin(query, controller.clone(), pool.clone()).await; + } +} + +// ============================================================ +// Worker 池（对齐 Python WorkerPool，动态扩缩） +// ============================================================ + +struct RunningReplica { + controller: Arc, + handle: tokio::task::JoinHandle<()>, +} + +/// 进程内的副本池：按 target 动态启动/停止槽位 +pub struct WorkerPool { + worker_type: String, + session: Session, + registry: Arc, + routes: Vec, + replicas: HashMap, + base_config: Value, + config_dir: String, + cache: Cache, + instance_id: String, +} + +impl WorkerPool { + pub fn new( + worker_type: String, + session: Session, + registry: Arc, + routes: Vec, + base_config: Value, + config_dir: String, + cache: Cache, + instance_id: String, + ) -> Self { + Self { + worker_type, + session, + registry, + routes, + replicas: HashMap::new(), + base_config, + config_dir, + cache, + instance_id, + } + } + + pub fn len(&self) -> usize { + self.replicas.len() + } + + /// 启动一个槽位（追加到末尾索引） + fn start_slot(&mut self, index: usize, pool: &Arc>) { + let controller = Arc::new(ReplicaController::new( + &self.worker_type, + &self.config_dir, + index, + self.base_config.clone(), + self.cache.clone(), + self.instance_id.clone(), + )); + let handle = tokio::spawn(managed_replica( + self.session.clone(), + controller.clone(), + self.registry.clone(), + self.routes.clone(), + Arc::clone(pool), + )); + self.replicas.insert(index, RunningReplica { controller, handle }); + } + + fn stop_slot(&mut self, index: usize, drain: bool) { + if let Some(rr) = self.replicas.remove(&index) { + if drain { + rr.controller.request_drain(); + } else { + rr.controller.request_stop(); + } + // 放弃跟踪：槽位按信号自行优雅退出（JoinHandle drop 即分离） + let _ = rr.handle; + } + } + + /// 调整目标副本数（保留槽位 0..target）。缩容走优雅排空/停止，扩容追加新槽位。 + pub async fn set_target( + &mut self, + target: usize, + pool: &Arc>, + ) { + let mut idx = self.replicas.len(); + if idx < target { + while idx < target { + self.start_slot(idx, pool); + idx += 1; + } + } + // 缩容：让超出 target 的槽位停止（index >= target） + let overflow: Vec = self.replicas.keys().cloned() + .filter(|i| *i >= target) + .collect(); + for idx in overflow { + self.stop_slot(idx, false); + } + tracing::info!( + "[{}] Worker 池目标={target}，当前槽位={current}", + self.worker_type, + current = self.replicas.len(), + ); + } + + /// 全部停止并等待回收（进程关闭） + pub async fn shutdown(&mut self) { + let indices: Vec = self.replicas.keys().cloned().collect(); + // 先信号全部停止 + for idx in &indices { + if let Some(rr) = self.replicas.get(idx) { + rr.controller.request_stop(); + } + } + // 等待各槽位退出 + let handles: Vec<_> = indices + .iter() + .filter_map(|idx| self.replicas.remove(idx)) + .map(|rr| rr.handle) + .collect(); + for h in handles { + let _ = h.await; + } + } +} \ No newline at end of file diff --git a/src/serve.rs b/src/serve.rs new file mode 100644 index 0000000..bc658f5 --- /dev/null +++ b/src/serve.rs @@ -0,0 +1,227 @@ +//! Worker 进程主流程（对齐 Python process.py `worker_serve`） +//! +//! 流程: +//! 1. 读全局 config.toml（[zenoh].json5_path）→ 打开 Zenoh 会话 +//! 2. 自读 config/workers//config.toml（[workers] 节展平 + mqtt_topics + cache） +//! 3. 加载路由描述符，按 handler 注册表做契约自检（跳过无 handler 的端点） +//! 4. （可选）建立共享缓存，登记实例参与方并启动心跳（多实例全局分摊） +//! 5. 发布路由描述符 → 网关 RouteRegistry 动态挂载 HTTP 路由 +//! 6. 建 WorkerPool：每副本槽位 = liveliness token + 业务 Queryable + __admin__ Queryable +//! 7. 每 10s 重发布描述符（自愈） +//! 8. SIGTERM/SIGINT → 摘除路由 → 优雅关停全部副本 → 摘除实例 → 关闭会话 + +use std::sync::Arc; + +use serde_json::Value; +use tokio::signal::unix::{signal, SignalKind}; +use tokio::sync::Mutex; + +use crate::cache::Cache; +use crate::config::{load_global_config, load_plugin_config}; +use crate::descriptor::load_routes_from_dir; +use crate::publisher::{announce_loop, RoutePublisher}; +use crate::runtime::WorkerPool; +use crate::HandlerRegistry; + +/// serve 入口参数 +pub struct ServeOptions { + /// 全局配置文件路径（读取 [zenoh] 节） + pub config_path: String, + /// 插件配置目录（默认 "config/workers"） + pub config_dir: String, + /// 副本数覆盖（None = 使用插件配置中的 count） + pub count: Option, + /// 外部进程实例标识（多实例全局分摊用；缺省 {hostname}:{pid}） + pub instance: Option, +} + +impl Default for ServeOptions { + fn default() -> Self { + Self { + config_path: "config.toml".into(), + config_dir: "config/workers".into(), + count: None, + instance: None, + } + } +} + +/// 实例参与者心跳间隔（秒，与网关/ Python 一致） +const INSTANCE_HEARTBEAT_SECS: u64 = 15; + +/// 启动独立 Worker 进程（阻塞直到收到 SIGTERM/SIGINT 并完成优雅退出） +pub async fn serve( + worker_type: &str, + registry: Arc, + opts: ServeOptions, +) -> Result<(), String> { + // 1. 全局配置 → Zenoh 会话 + let global = load_global_config(&opts.config_path)?; + tracing::info!("[{}] 读取 Zenoh 配置: {}", worker_type, global.zenoh_json5_path); + let zconfig = zenoh::Config::from_file(&global.zenoh_json5_path) + .map_err(|e| format!("加载 Zenoh 配置 {} 失败: {e}", global.zenoh_json5_path))?; + let session = zenoh::open(zconfig) + .await + .map_err(|e| format!("打开 Zenoh 会话失败: {e}"))?; + tracing::info!("[{}] Zenoh 会话已建立", worker_type); + + // 2. 插件配置 + let plugin = load_plugin_config(worker_type, &opts.config_dir)?; + if plugin.load_mode != "external" { + tracing::warn!( + "[{}] 插件配置 load_mode 不是 \"external\"：网关也会在进程内加载该插件，\ + 外部进程将与网关内副本重复服务。", + worker_type, + ); + } + + // 2.5 共享缓存（[workers.cache]，进程插件自行管理）+ 实例参与方登记 + let cache_cfg = plugin + .runtime + .get("workers") + .and_then(|w| w.get("cache")) + .cloned(); + let (cache, backend) = match &cache_cfg { + Some(cfg) => Cache::from_config(cfg).await, + None => (Cache::None, "none".to_string()), + }; + let mode = cache_cfg + .as_ref() + .and_then(|c| c.get("launch_worker_count").and_then(Value::as_str)) + .unwrap_or("toml") + .to_lowercase(); + let instance_id = opts.instance.clone().unwrap_or_else(default_instance_id); + if cache.enabled() { + crate::cache::heartbeat_instance(&cache, worker_type, &instance_id).await; + tracing::info!( + "[{}] 已启用配置缓存（backend={backend}，mode={mode}）并登记实例参与方 {}", + worker_type, instance_id, + ); + } + + // 3. 路由描述符 + 契约自检 + let all_routes = load_routes_from_dir(worker_type, &opts.config_dir)?; + let routes = contract_self_check(worker_type, &all_routes, ®istry)?; + + // 4. 发布路由描述符 + let publisher = Arc::new(RoutePublisher::new(session.clone(), worker_type)); + publisher + .publish_routes(&routes, &plugin.mqtt_topics) + .await?; + + // 5. 解析本进程应承载的副本数（多实例全局分摊） + let default_count = opts.count.unwrap_or(plugin.count).max(1); + let count = crate::cache::resolve_worker_count( + &cache, + worker_type, + default_count, + &mode, + None, + Some(&instance_id), + ) + .await + .max(1); + + // 6. 建 WorkerPool + let pool: Arc> = Arc::new(Mutex::new(WorkerPool::new( + worker_type.to_string(), + session.clone(), + registry, + routes.clone(), + plugin.runtime.clone(), + opts.config_dir, + cache.clone(), + instance_id.clone(), + ))); + { + let mut p = pool.lock().await; + p.set_target(count, &pool).await; + } + tracing::info!("[{worker_type}] Worker 池就绪（{count} 副本）"); + + // 7. 周期重发布（自愈）+ 实例心跳 + let announce = tokio::spawn(announce_loop( + publisher.clone(), + routes.clone(), + plugin.mqtt_topics.clone(), + )); + let hb_cache = cache.clone(); + let hb_type = worker_type.to_string(); + let hb_instance = instance_id.clone(); + let heartbeat = tokio::spawn(async move { + let mut interval = + tokio::time::interval(std::time::Duration::from_secs(INSTANCE_HEARTBEAT_SECS)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + interval.tick().await; + crate::cache::heartbeat_instance(&hb_cache, &hb_type, &hb_instance).await; + } + }); + + tracing::info!("[{worker_type}] Worker 进程就绪，等待退出信号..."); + + // 8. 等待退出信号 + let mut sigterm = signal(SignalKind::terminate()) + .map_err(|e| format!("注册 SIGTERM 失败: {e}"))?; + let mut sigint = signal(SignalKind::interrupt()) + .map_err(|e| format!("注册 SIGINT 失败: {e}"))?; + tokio::select! { + _ = sigterm.recv() => tracing::info!("[{worker_type}] 收到 SIGTERM，优雅退出"), + _ = sigint.recv() => tracing::info!("[{worker_type}] 收到 SIGINT，优雅退出"), + } + + // 9. 优雅退出：停周期任务 → 摘除路由 → 关停全部副本 → 摘除实例 → 关闭会话 + announce.abort(); + let _ = announce.await; + heartbeat.abort(); + let _ = heartbeat.await; + publisher.unpublish().await; + { + let mut p = pool.lock().await; + p.shutdown().await; + } + if cache.enabled() { + crate::cache::unregister_instance(&cache, worker_type, &instance_id).await; + } + session + .close() + .await + .map_err(|e| format!("关闭 Zenoh 会话失败: {e}"))?; + tracing::info!("[{worker_type}] Worker 进程已退出"); + Ok(()) +} + +/// 契约自检：跳过无 handler 的端点，无可用端点则报错 +fn contract_self_check( + worker_type: &str, + all_routes: &[Value], + registry: &HandlerRegistry, +) -> Result, String> { + let mut routes: Vec = Vec::new(); + for r in all_routes { + let subject = r.get("subject").and_then(Value::as_str).unwrap_or(""); + let method = r.get("method").and_then(Value::as_str).unwrap_or("GET"); + let path = r.get("path").and_then(Value::as_str).unwrap_or("?"); + if registry.has(subject, method) { + routes.push(r.clone()); + } else { + tracing::warn!( + "[{}] 跳过 {method} {path} (subject={subject}): 无对应 handler 注册", + worker_type, + ); + } + } + if routes.is_empty() { + return Err(format!( + "[{worker_type}] 没有任何可发布的端点（handler 注册表为空或全部被跳过）" + )); + } + tracing::info!("[{worker_type}] 契约自检通过，可发布 {} 个路由描述符", routes.len()); + Ok(routes) +} + +/// 默认实例 id：{hostname}:{pid}（同机多进程自动区分） +fn default_instance_id() -> String { + let host = std::env::var("HOSTNAME").unwrap_or_else(|_| "localhost".to_string()); + format!("{host}:{}", std::process::id()) +} \ No newline at end of file