195 lines
7.8 KiB
Rust
195 lines
7.8 KiB
Rust
//! 参考插件（多语言演进规划阶段 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<usize>,
|
||
/// 全局配置文件路径（读取 [zenoh] 节）
|
||
#[arg(long, default_value = "config.toml")]
|
||
config: String,
|
||
/// 实例 id（多实例全局分摊用；缺省 {hostname}:{pid}）
|
||
#[arg(long)]
|
||
instance: Option<String>,
|
||
}
|
||
|
||
fn now_ts() -> f64 {
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.map(|d| d.as_secs_f64())
|
||
.unwrap_or(0.0)
|
||
}
|
||
|
||
fn deser<T: serde::de::DeserializeOwned>(data: &Value) -> Result<T, String> {
|
||
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<Reply, String> {
|
||
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<Reply, String> {
|
||
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<Reply, String> {
|
||
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<Reply, String> {
|
||
Err("模拟业务异常（rsdemo.error）".to_string())
|
||
}
|
||
|
||
// ==================== demo 端点（使用 demo 生成模型，与 Python 混布对拍） ====================
|
||
|
||
/// demo.hello：与 Python workers/demo 对位（强类型模型来自 demo 生成模型）
|
||
async fn hello_demo(request: Request) -> Result<Reply, String> {
|
||
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<Reply, String> {
|
||
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 ↔ 生成模型）");
|
||
}
|
||
}
|