107 lines
3.7 KiB
Rust
107 lines
3.7 KiB
Rust
//! chongming-worker-rs —— 进程插件 Rust SDK
|
||
//!
|
||
//! 语言无关线协议（docs/plugin-wire-protocol.md）的 Rust 实现：
|
||
//! 1. 解析 config.toml 与端点 TOML，构造路由描述符并发布到 `chongming/routes/**`
|
||
//! 2. liveliness token 注册 `worker/<type>/<id>`（服务发现）
|
||
//! 3. Queryable 声明 `worker/<type>/<id>/<subject>`，解析请求信封并分派 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<T> = Pin<Box<dyn Future<Output = T> + Send>>;
|
||
|
||
/// 端点响应：单块 JSON，或流式多块（SSE 逐块转发）
|
||
pub enum Reply {
|
||
/// 普通响应：单个回复，载荷为业务结果 JSON 对象
|
||
Json(Value),
|
||
/// 流式响应：逐块多回复；每块 `Ok(v)` 作为独立回复块，
|
||
/// `Err(msg)` 以 `{"error": msg}` 收尾块结束流（与 Python SDK 一致）
|
||
Stream(BoxStream<'static, Result<Value, String>>),
|
||
}
|
||
|
||
/// 端点 handler：解析后的请求信封 → 业务结果
|
||
pub type Handler = dyn Fn(Request) -> BoxFuture<Result<Reply, String>> + Send + Sync;
|
||
|
||
/// RPC 端点注册表（按 subject + HTTP 方法分派）
|
||
#[derive(Default)]
|
||
pub struct HandlerRegistry {
|
||
/// subject -> [(method, handler)]（同一 subject 可挂多个方法）
|
||
map: HashMap<String, Vec<(String, Arc<Handler>)>>,
|
||
}
|
||
|
||
impl HandlerRegistry {
|
||
pub fn new() -> Self {
|
||
Self::default()
|
||
}
|
||
|
||
/// 注册端点。`f` 为 async handler：`fn(Request) -> Result<Reply, String>`。
|
||
pub fn register<F, Fut>(&mut self, subject: &str, method: &str, f: F)
|
||
where
|
||
F: Fn(Request) -> Fut + Send + Sync + 'static,
|
||
Fut: Future<Output = Result<Reply, String>> + 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<Box<dyn Future<Output = Result<Reply, String>> + Send>>
|
||
}) as Arc<Handler>,
|
||
));
|
||
}
|
||
|
||
/// 该 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<Arc<Handler>> {
|
||
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())
|
||
}
|
||
}
|