feat: 拆分 worker 架构并增强 assistant 安全保护
- 新增 daemon/channel/assistant/function/scheduler/logging worker 架构和 JetStream 消息协议 - 修复邮件通知误入 assistant、迁移 checksum、日志 subject 和工具调用无限循环风险 - 新增 assistant 工具轮次熔断、工具结果截断和提示词软收敛规则 - 补充邮件推送关键字、可信发件人配置、升级日志和验证测试
This commit is contained in:
@@ -0,0 +1,26 @@
|
||||
[package]
|
||||
name = "ias-daemon"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
ias-common = { path = "../ias-common" }
|
||||
anyhow = "1"
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio = { version = "1.0", features = [
|
||||
"rt-multi-thread",
|
||||
"macros",
|
||||
"sync",
|
||||
"time",
|
||||
"process",
|
||||
"signal",
|
||||
"io-util",
|
||||
"net",
|
||||
] }
|
||||
tracing = "0.1"
|
||||
libc = "0.2"
|
||||
|
||||
[dev-dependencies]
|
||||
toml = "0.8"
|
||||
@@ -0,0 +1,645 @@
|
||||
pub mod supervisor;
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::Utc;
|
||||
use ias_common::control::{
|
||||
ControlResponse, MqStatus, RuntimePaths, RuntimeSnapshot, ServicePhase, StreamStatus,
|
||||
WorkerPhase, WorkerStatus,
|
||||
};
|
||||
use ias_common::mq::{
|
||||
JetStreamBus, JetStreamConfig, LogEventPayload, MessageEnvelope, canonical_stream_specs,
|
||||
log_subject, new_trace_id,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::{Mutex, watch};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct WorkerSpec {
|
||||
name: String,
|
||||
role: String,
|
||||
args: Vec<String>,
|
||||
}
|
||||
|
||||
struct ManagedWorker {
|
||||
spec: WorkerSpec,
|
||||
child: Child,
|
||||
}
|
||||
|
||||
struct DaemonInner {
|
||||
phase: ServicePhase,
|
||||
started_at: chrono::DateTime<Utc>,
|
||||
updated_at: chrono::DateTime<Utc>,
|
||||
http_addr: Option<String>,
|
||||
mq: MqStatus,
|
||||
workers: HashMap<String, WorkerStatus>,
|
||||
last_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DaemonState {
|
||||
inner: Arc<Mutex<DaemonInner>>,
|
||||
}
|
||||
|
||||
impl DaemonState {
|
||||
fn new(config: &JetStreamConfig, http_addr: Option<String>) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(DaemonInner {
|
||||
phase: ServicePhase::Starting,
|
||||
started_at: now,
|
||||
updated_at: now,
|
||||
http_addr,
|
||||
mq: mq_status(config, false, Some("尚未连接 JetStream".to_string())),
|
||||
workers: HashMap::new(),
|
||||
last_error: None,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
async fn mark_running(&self) {
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner.phase = if inner.mq.connected {
|
||||
ServicePhase::Running
|
||||
} else {
|
||||
ServicePhase::Degraded
|
||||
};
|
||||
inner.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
async fn mark_stopping(&self) {
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner.phase = ServicePhase::Stopping;
|
||||
inner.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
async fn set_mq(&self, mq: MqStatus) {
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner.mq = mq;
|
||||
inner.phase = match inner.phase {
|
||||
ServicePhase::Starting | ServicePhase::Running | ServicePhase::Degraded => {
|
||||
if inner.mq.connected {
|
||||
ServicePhase::Running
|
||||
} else {
|
||||
ServicePhase::Degraded
|
||||
}
|
||||
}
|
||||
_ => inner.phase.clone(),
|
||||
};
|
||||
inner.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
async fn set_worker(&self, status: WorkerStatus) {
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner.workers.insert(status.name.clone(), status);
|
||||
inner.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
async fn snapshot(&self) -> RuntimeSnapshot {
|
||||
let inner = self.inner.lock().await;
|
||||
let mut workers = inner.workers.values().cloned().collect::<Vec<_>>();
|
||||
workers.sort_by(|left, right| left.name.cmp(&right.name));
|
||||
RuntimeSnapshot {
|
||||
phase: inner.phase.clone(),
|
||||
pid: std::process::id(),
|
||||
started_at: inner.started_at,
|
||||
updated_at: inner.updated_at,
|
||||
http_addr: inner.http_addr.clone(),
|
||||
mq: inner.mq.clone(),
|
||||
workers,
|
||||
last_error: inner.last_error.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_snapshot(&self, path: &std::path::Path) -> anyhow::Result<()> {
|
||||
let snapshot = self.snapshot().await;
|
||||
std::fs::write(path, serde_json::to_vec_pretty(&snapshot)?)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_daemon() -> anyhow::Result<()> {
|
||||
let paths = RuntimePaths::new();
|
||||
paths.ensure_dir()?;
|
||||
if let Ok(ControlResponse::Status { snapshot }) = request_control(&paths, "status").await {
|
||||
anyhow::bail!("daemon 已在运行，PID: {}", snapshot.pid);
|
||||
}
|
||||
cleanup_stale_runtime(&paths)?;
|
||||
write_pid_file(&paths)?;
|
||||
|
||||
let config = JetStreamConfig::from_env();
|
||||
let host = read_host();
|
||||
let state = DaemonState::new(&config, Some(host.clone()));
|
||||
let (shutdown_tx, mut shutdown_rx) = start_control(paths.clone(), state.clone()).await?;
|
||||
let signal_task = tokio::spawn(wait_for_shutdown_signal(shutdown_tx));
|
||||
|
||||
let mut bus = match JetStreamBus::connect(config.clone()).await {
|
||||
Ok(bus) => {
|
||||
if let Err(error) = bus.ensure_streams().await {
|
||||
state
|
||||
.set_mq(mq_status(&config, false, Some(error.to_string())))
|
||||
.await;
|
||||
None
|
||||
} else {
|
||||
state.set_mq(connected_mq_status(&bus, &config).await).await;
|
||||
publish_daemon_log(&bus, "info", "JetStream 初始化完成").await;
|
||||
Some(bus)
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
state
|
||||
.set_mq(mq_status(&config, false, Some(error.to_string())))
|
||||
.await;
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
let exe = std::env::current_exe()?;
|
||||
let mut workers = Vec::new();
|
||||
let worker_specs = if autostart_workers() {
|
||||
default_worker_specs(&host)
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
if !worker_specs.is_empty() {
|
||||
for spec in &worker_specs {
|
||||
match spawn_worker(&exe, spec.clone(), &state).await {
|
||||
Ok(worker) => workers.push(worker),
|
||||
Err(error) => {
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: spec.name.clone(),
|
||||
role: spec.role.clone(),
|
||||
pid: None,
|
||||
phase: WorkerPhase::Failed,
|
||||
last_heartbeat_at: None,
|
||||
last_error: Some(error.to_string()),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mark_running().await;
|
||||
state.write_snapshot(&paths.state_file).await?;
|
||||
let mut tick = tokio::time::interval(Duration::from_secs(5));
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
changed = shutdown_rx.changed() => {
|
||||
if changed.is_err() || *shutdown_rx.borrow() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = tick.tick() => {
|
||||
if let Some(bus) = bus.as_ref() {
|
||||
state.set_mq(connected_mq_status(bus, &config).await).await;
|
||||
} else {
|
||||
match JetStreamBus::connect(config.clone()).await {
|
||||
Ok(candidate) => {
|
||||
if let Err(error) = candidate.ensure_streams().await {
|
||||
state
|
||||
.set_mq(mq_status(&config, false, Some(error.to_string())))
|
||||
.await;
|
||||
} else {
|
||||
state.set_mq(connected_mq_status(&candidate, &config).await).await;
|
||||
bus = Some(candidate);
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
state
|
||||
.set_mq(mq_status(&config, false, Some(error.to_string())))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
refresh_workers(&exe, &state, &mut workers, &worker_specs).await;
|
||||
state.write_snapshot(&paths.state_file).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.mark_stopping().await;
|
||||
for worker in &mut workers {
|
||||
let _ = worker.child.start_kill();
|
||||
let _ = worker.child.wait().await;
|
||||
}
|
||||
signal_task.abort();
|
||||
if let Some(bus) = bus {
|
||||
publish_daemon_log(&bus, "info", "daemon 已停止").await;
|
||||
}
|
||||
cleanup_runtime_files(&paths);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn request_control(
|
||||
paths: &RuntimePaths,
|
||||
command: &str,
|
||||
) -> anyhow::Result<ControlResponse> {
|
||||
let mut stream = UnixStream::connect(&paths.socket_file).await?;
|
||||
stream.write_all(command.as_bytes()).await?;
|
||||
stream.write_all(b"\n").await?;
|
||||
|
||||
let mut reader = BufReader::new(stream);
|
||||
let mut line = String::new();
|
||||
reader.read_line(&mut line).await?;
|
||||
Ok(serde_json::from_str(line.trim())?)
|
||||
}
|
||||
|
||||
async fn start_control(
|
||||
paths: RuntimePaths,
|
||||
state: DaemonState,
|
||||
) -> anyhow::Result<(watch::Sender<bool>, watch::Receiver<bool>)> {
|
||||
paths.ensure_dir()?;
|
||||
if paths.socket_file.exists() {
|
||||
if UnixStream::connect(&paths.socket_file).await.is_ok() {
|
||||
anyhow::bail!(
|
||||
"daemon control socket 已存在且可连接: {}",
|
||||
paths.socket_file.display()
|
||||
);
|
||||
}
|
||||
std::fs::remove_file(&paths.socket_file)?;
|
||||
}
|
||||
let listener = UnixListener::bind(&paths.socket_file)?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
std::fs::set_permissions(&paths.socket_file, std::fs::Permissions::from_mode(0o600))?;
|
||||
}
|
||||
let (shutdown_tx, shutdown_rx) = watch::channel(false);
|
||||
let mut control_shutdown_rx = shutdown_tx.subscribe();
|
||||
let control_shutdown_tx = shutdown_tx.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
accepted = listener.accept() => {
|
||||
match accepted {
|
||||
Ok((stream, _)) => {
|
||||
let state = state.clone();
|
||||
let shutdown_tx = control_shutdown_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = handle_control_stream(stream, state, shutdown_tx).await {
|
||||
eprintln!("处理 daemon 控制命令失败: {error:#}");
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(error) => {
|
||||
eprintln!("接收 daemon 控制连接失败: {error:#}");
|
||||
}
|
||||
}
|
||||
}
|
||||
changed = control_shutdown_rx.changed() => {
|
||||
if changed.is_err() || *control_shutdown_rx.borrow() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok((shutdown_tx, shutdown_rx))
|
||||
}
|
||||
|
||||
async fn handle_control_stream(
|
||||
stream: UnixStream,
|
||||
state: DaemonState,
|
||||
shutdown_tx: watch::Sender<bool>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut reader = BufReader::new(stream);
|
||||
let mut command = String::new();
|
||||
reader.read_line(&mut command).await?;
|
||||
let response = match command.trim() {
|
||||
"status" => ControlResponse::Status {
|
||||
snapshot: state.snapshot().await,
|
||||
},
|
||||
"shutdown" => {
|
||||
state.mark_stopping().await;
|
||||
let _ = shutdown_tx.send(true);
|
||||
ControlResponse::Ack {
|
||||
message: "daemon 正在停止".to_string(),
|
||||
}
|
||||
}
|
||||
command => ControlResponse::Error {
|
||||
message: format!("未知控制命令: {command}"),
|
||||
},
|
||||
};
|
||||
let stream = reader.get_mut();
|
||||
stream
|
||||
.write_all(serde_json::to_string(&response)?.as_bytes())
|
||||
.await?;
|
||||
stream.write_all(b"\n").await?;
|
||||
stream.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn spawn_worker(
|
||||
exe: &std::path::Path,
|
||||
spec: WorkerSpec,
|
||||
state: &DaemonState,
|
||||
) -> anyhow::Result<ManagedWorker> {
|
||||
let mut command = Command::new(exe);
|
||||
command
|
||||
.args(&spec.args)
|
||||
.env("IAS_MANAGED_WORKER", "1")
|
||||
.stdin(Stdio::null());
|
||||
let child = command.spawn()?;
|
||||
let pid = child.id();
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: spec.name.clone(),
|
||||
role: spec.role.clone(),
|
||||
pid,
|
||||
phase: WorkerPhase::Running,
|
||||
last_heartbeat_at: Some(Utc::now()),
|
||||
last_error: None,
|
||||
})
|
||||
.await;
|
||||
Ok(ManagedWorker { spec, child })
|
||||
}
|
||||
|
||||
async fn refresh_workers(
|
||||
exe: &std::path::Path,
|
||||
state: &DaemonState,
|
||||
workers: &mut Vec<ManagedWorker>,
|
||||
worker_specs: &[WorkerSpec],
|
||||
) {
|
||||
for worker in workers.iter_mut() {
|
||||
match worker.child.try_wait() {
|
||||
Ok(Some(status)) => {
|
||||
let message = if status.success() {
|
||||
None
|
||||
} else {
|
||||
Some(format!("退出状态: {status}"))
|
||||
};
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: worker.spec.name.clone(),
|
||||
role: worker.spec.role.clone(),
|
||||
pid: None,
|
||||
phase: if status.success() {
|
||||
WorkerPhase::Stopped
|
||||
} else {
|
||||
WorkerPhase::Failed
|
||||
},
|
||||
last_heartbeat_at: Some(Utc::now()),
|
||||
last_error: message,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Ok(None) => {
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: worker.spec.name.clone(),
|
||||
role: worker.spec.role.clone(),
|
||||
pid: worker.child.id(),
|
||||
phase: WorkerPhase::Running,
|
||||
last_heartbeat_at: Some(Utc::now()),
|
||||
last_error: None,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: worker.spec.name.clone(),
|
||||
role: worker.spec.role.clone(),
|
||||
pid: worker.child.id(),
|
||||
phase: WorkerPhase::Failed,
|
||||
last_heartbeat_at: Some(Utc::now()),
|
||||
last_error: Some(error.to_string()),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if restart_workers() {
|
||||
workers.retain(|worker| worker.child.id().is_some());
|
||||
let active_names = workers
|
||||
.iter()
|
||||
.map(|worker| worker.spec.name.clone())
|
||||
.collect::<HashSet<_>>();
|
||||
let mut replacements = Vec::new();
|
||||
for spec in worker_specs {
|
||||
if active_names.contains(&spec.name) {
|
||||
continue;
|
||||
}
|
||||
match spawn_worker(exe, spec.clone(), state).await {
|
||||
Ok(replacement) => replacements.push(replacement),
|
||||
Err(error) => {
|
||||
state
|
||||
.set_worker(WorkerStatus {
|
||||
name: spec.name.clone(),
|
||||
role: spec.role.clone(),
|
||||
pid: None,
|
||||
phase: WorkerPhase::Failed,
|
||||
last_heartbeat_at: Some(Utc::now()),
|
||||
last_error: Some(error.to_string()),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
workers.extend(replacements);
|
||||
}
|
||||
}
|
||||
|
||||
fn default_worker_specs(host: &str) -> Vec<WorkerSpec> {
|
||||
vec![
|
||||
WorkerSpec {
|
||||
name: "http".to_string(),
|
||||
role: "http".to_string(),
|
||||
args: vec!["http".to_string(), "--host".to_string(), host.to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "assistant".to_string(),
|
||||
role: "assistant".to_string(),
|
||||
args: vec!["worker".to_string(), "assistant".to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "wechat".to_string(),
|
||||
role: "channel".to_string(),
|
||||
args: vec!["worker".to_string(), "wechat".to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "mail".to_string(),
|
||||
role: "channel".to_string(),
|
||||
args: vec!["worker".to_string(), "mail".to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "function".to_string(),
|
||||
role: "function".to_string(),
|
||||
args: vec!["worker".to_string(), "function".to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "scheduler".to_string(),
|
||||
role: "scheduler".to_string(),
|
||||
args: vec!["worker".to_string(), "scheduler".to_string()],
|
||||
},
|
||||
WorkerSpec {
|
||||
name: "logging".to_string(),
|
||||
role: "logging".to_string(),
|
||||
args: vec!["worker".to_string(), "logging".to_string()],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
fn mq_status(config: &JetStreamConfig, connected: bool, last_error: Option<String>) -> MqStatus {
|
||||
MqStatus {
|
||||
connected,
|
||||
nats_url: config.nats_url.clone(),
|
||||
stream_prefix: config.stream_prefix.clone(),
|
||||
streams: canonical_stream_specs(config)
|
||||
.into_iter()
|
||||
.map(|spec| StreamStatus {
|
||||
name: spec.name,
|
||||
subjects: spec.subjects,
|
||||
messages: 0,
|
||||
bytes: 0,
|
||||
consumer_lag: None,
|
||||
dlq_messages: None,
|
||||
})
|
||||
.collect(),
|
||||
last_error,
|
||||
}
|
||||
}
|
||||
|
||||
async fn connected_mq_status(bus: &JetStreamBus, config: &JetStreamConfig) -> MqStatus {
|
||||
let mut status = mq_status(config, true, None);
|
||||
match bus.stream_statuses().await {
|
||||
Ok(streams) => status.streams = streams,
|
||||
Err(error) => status.last_error = Some(error.to_string()),
|
||||
}
|
||||
status
|
||||
}
|
||||
|
||||
async fn publish_daemon_log(bus: &JetStreamBus, level: &str, message: &str) {
|
||||
if !log_events_enabled() {
|
||||
return;
|
||||
}
|
||||
let trace_id = new_trace_id();
|
||||
let payload = LogEventPayload::new(
|
||||
level,
|
||||
"daemon.control",
|
||||
message,
|
||||
trace_id.clone(),
|
||||
json!({}),
|
||||
);
|
||||
let subject = log_subject("daemon.control");
|
||||
let envelope = MessageEnvelope::new(
|
||||
subject.as_str(),
|
||||
payload.event_id.clone(),
|
||||
trace_id,
|
||||
payload,
|
||||
);
|
||||
if let Err(error) = bus.publish_envelope(subject.as_str(), &envelope).await {
|
||||
eprintln!("发布 daemon 日志事件失败: {error:#}");
|
||||
}
|
||||
}
|
||||
|
||||
fn cleanup_stale_runtime(paths: &RuntimePaths) -> anyhow::Result<()> {
|
||||
if let Some(pid) = read_pid(paths)?
|
||||
&& process_is_alive(pid)
|
||||
{
|
||||
anyhow::bail!(
|
||||
"已有 daemon 进程存在但 control socket 不可用，PID: {}。请先 stop 或检查日志: {}",
|
||||
pid,
|
||||
paths.log_file.display()
|
||||
);
|
||||
}
|
||||
let _ = std::fs::remove_file(&paths.pid_file);
|
||||
let _ = std::fs::remove_file(&paths.socket_file);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn write_pid_file(paths: &RuntimePaths) -> anyhow::Result<()> {
|
||||
paths.ensure_dir()?;
|
||||
std::fs::write(&paths.pid_file, std::process::id().to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn cleanup_runtime_files(paths: &RuntimePaths) {
|
||||
if let Some(pid) = read_pid(paths).ok().flatten()
|
||||
&& pid != std::process::id()
|
||||
{
|
||||
return;
|
||||
}
|
||||
let _ = std::fs::remove_file(&paths.pid_file);
|
||||
let _ = std::fs::remove_file(&paths.socket_file);
|
||||
}
|
||||
|
||||
pub fn read_pid(paths: &RuntimePaths) -> anyhow::Result<Option<u32>> {
|
||||
if !paths.pid_file.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let pid = std::fs::read_to_string(&paths.pid_file)?
|
||||
.trim()
|
||||
.parse::<u32>()?;
|
||||
Ok(Some(pid))
|
||||
}
|
||||
|
||||
pub fn process_is_alive(pid: u32) -> bool {
|
||||
std::path::Path::new("/proc").join(pid.to_string()).exists()
|
||||
}
|
||||
|
||||
async fn wait_for_shutdown_signal(shutdown_tx: watch::Sender<bool>) {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
let mut terminate = match signal(SignalKind::terminate()) {
|
||||
Ok(signal) => signal,
|
||||
Err(error) => {
|
||||
eprintln!("监听 SIGTERM 失败: {error}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
tokio::select! {
|
||||
_ = tokio::signal::ctrl_c() => {}
|
||||
_ = terminate.recv() => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
}
|
||||
|
||||
let _ = shutdown_tx.send(true);
|
||||
}
|
||||
|
||||
fn read_host() -> String {
|
||||
std::env::var("HOST")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| "0.0.0.0:9003".to_string())
|
||||
}
|
||||
|
||||
fn autostart_workers() -> bool {
|
||||
read_bool_env("IAS_DAEMON_AUTOSTART_WORKERS", true)
|
||||
}
|
||||
|
||||
fn restart_workers() -> bool {
|
||||
read_bool_env("IAS_DAEMON_RESTART_WORKERS", true)
|
||||
}
|
||||
|
||||
fn log_events_enabled() -> bool {
|
||||
read_bool_env("IAS_LOG_EVENT_ENABLED", true)
|
||||
}
|
||||
|
||||
fn read_bool_env(name: &str, default: bool) -> bool {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.map(|value| {
|
||||
let value = value.trim().to_ascii_lowercase();
|
||||
!matches!(value.as_str(), "0" | "false" | "no" | "off")
|
||||
})
|
||||
.unwrap_or(default)
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Write};
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::process::CommandExt;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use ias_common::control::{ControlResponse, RuntimePaths, RuntimeSnapshot, ServicePhase};
|
||||
|
||||
use crate::{cleanup_runtime_files, process_is_alive, read_pid, request_control};
|
||||
|
||||
pub async fn start_service() -> anyhow::Result<()> {
|
||||
let paths = RuntimePaths::new();
|
||||
paths.ensure_dir()?;
|
||||
|
||||
if let Ok(ControlResponse::Status { snapshot }) = request_control(&paths, "status").await {
|
||||
println!("daemon 已经在运行，PID: {}", snapshot.pid);
|
||||
print_snapshot(&snapshot);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
cleanup_stale_files(&paths)?;
|
||||
let mut log = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&paths.log_file)?;
|
||||
writeln!(log, "\n===== 启动 iAs daemon =====")?;
|
||||
|
||||
let exe = std::env::current_exe()?;
|
||||
let stdout = log.try_clone()?;
|
||||
let stderr = log;
|
||||
let mut command = Command::new(exe);
|
||||
command
|
||||
.arg("daemon")
|
||||
.arg("run")
|
||||
.env("IAS_BACKGROUND", "1")
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::from(stdout))
|
||||
.stderr(Stdio::from(stderr));
|
||||
detach_background_child(&mut command);
|
||||
let child = command.spawn()?;
|
||||
|
||||
println!("daemon 启动中，PID: {}", child.id());
|
||||
wait_until_ready(&paths, Duration::from_secs(8)).await?;
|
||||
println!("日志文件: {}", paths.log_file.display());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn stop_service() -> anyhow::Result<()> {
|
||||
let paths = RuntimePaths::new();
|
||||
paths.ensure_dir()?;
|
||||
|
||||
match request_control(&paths, "shutdown").await {
|
||||
Ok(ControlResponse::Ack { message }) => println!("{message}"),
|
||||
Ok(ControlResponse::Error { message }) => println!("停止失败: {message}"),
|
||||
Ok(ControlResponse::Status { .. }) => println!("停止失败: control 返回了意外状态响应"),
|
||||
Err(_) => {
|
||||
if let Some(pid) = read_pid(&paths)? {
|
||||
if process_is_alive(pid) {
|
||||
println!("control socket 不可用，发送 SIGTERM 到 PID: {pid}");
|
||||
send_sigterm(pid)?;
|
||||
} else {
|
||||
println!("daemon 未运行，清理过期运行文件");
|
||||
cleanup_stale_files(&paths)?;
|
||||
return Ok(());
|
||||
}
|
||||
} else {
|
||||
println!("daemon 未运行");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
wait_until_stopped(&paths, Duration::from_secs(10)).await?;
|
||||
cleanup_runtime_files(&paths);
|
||||
println!("daemon 已停止");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn restart_service() -> anyhow::Result<()> {
|
||||
stop_service().await?;
|
||||
start_service().await
|
||||
}
|
||||
|
||||
pub async fn status_service() -> anyhow::Result<()> {
|
||||
let paths = RuntimePaths::new();
|
||||
paths.ensure_dir()?;
|
||||
|
||||
match request_control(&paths, "status").await {
|
||||
Ok(ControlResponse::Status { snapshot }) => print_snapshot(&snapshot),
|
||||
Ok(ControlResponse::Error { message }) => println!("状态查询失败: {message}"),
|
||||
Ok(ControlResponse::Ack { message }) => println!("{message}"),
|
||||
Err(_) => print_offline_status(&paths)?,
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn print_logs(lines: usize) -> anyhow::Result<()> {
|
||||
let paths = RuntimePaths::new();
|
||||
if !paths.log_file.exists() {
|
||||
println!("日志文件不存在: {}", paths.log_file.display());
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut file = File::open(&paths.log_file)?;
|
||||
let mut content = String::new();
|
||||
file.read_to_string(&mut content)?;
|
||||
let mut tail = content.lines().rev().take(lines).collect::<Vec<_>>();
|
||||
tail.reverse();
|
||||
for line in tail {
|
||||
println!("{line}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn wait_until_ready(paths: &RuntimePaths, timeout: Duration) -> anyhow::Result<()> {
|
||||
let started = Instant::now();
|
||||
loop {
|
||||
if let Ok(ControlResponse::Status { snapshot }) = request_control(paths, "status").await {
|
||||
println!("daemon 已启动");
|
||||
print_snapshot(&snapshot);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if started.elapsed() >= timeout {
|
||||
anyhow::bail!(
|
||||
"daemon 未在 {} 秒内就绪，请查看日志: {}",
|
||||
timeout.as_secs(),
|
||||
paths.log_file.display()
|
||||
);
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_until_stopped(paths: &RuntimePaths, timeout: Duration) -> anyhow::Result<()> {
|
||||
let started = Instant::now();
|
||||
loop {
|
||||
if request_control(paths, "status").await.is_err()
|
||||
&& read_pid(paths)?.is_none_or(|pid| !process_is_alive(pid))
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if started.elapsed() >= timeout {
|
||||
anyhow::bail!(
|
||||
"daemon 未在 {} 秒内停止，请检查 PID 文件: {}",
|
||||
timeout.as_secs(),
|
||||
paths.pid_file.display()
|
||||
);
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn cleanup_stale_files(paths: &RuntimePaths) -> anyhow::Result<()> {
|
||||
if let Some(pid) = read_pid(paths)?
|
||||
&& process_is_alive(pid)
|
||||
{
|
||||
anyhow::bail!(
|
||||
"已有 daemon 进程存在但 control socket 不可用，PID: {}。请先执行 stop 或检查日志: {}",
|
||||
pid,
|
||||
paths.log_file.display()
|
||||
);
|
||||
}
|
||||
|
||||
let _ = std::fs::remove_file(&paths.pid_file);
|
||||
let _ = std::fs::remove_file(&paths.socket_file);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn print_offline_status(paths: &RuntimePaths) -> anyhow::Result<()> {
|
||||
if let Some(pid) = read_pid(paths)?
|
||||
&& process_is_alive(pid)
|
||||
{
|
||||
println!("daemon 进程存在但 control socket 不可用，PID: {pid}");
|
||||
println!("socket: {}", paths.socket_file.display());
|
||||
println!("日志文件: {}", paths.log_file.display());
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("daemon 状态: stopped");
|
||||
if paths.state_file.exists() {
|
||||
println!("最后状态快照: {}", paths.state_file.display());
|
||||
}
|
||||
println!("日志文件: {}", paths.log_file.display());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn print_snapshot(snapshot: &RuntimeSnapshot) {
|
||||
println!("daemon 状态: {}", phase_name(&snapshot.phase));
|
||||
println!("PID: {}", snapshot.pid);
|
||||
if let Some(http_addr) = &snapshot.http_addr {
|
||||
println!("HTTP: {http_addr}");
|
||||
}
|
||||
println!("启动时间: {}", snapshot.started_at);
|
||||
println!(
|
||||
"MQ: {} ({})",
|
||||
if snapshot.mq.connected {
|
||||
"connected"
|
||||
} else {
|
||||
"degraded"
|
||||
},
|
||||
snapshot.mq.nats_url
|
||||
);
|
||||
if let Some(error) = &snapshot.mq.last_error {
|
||||
println!("MQ 最近错误: {error}");
|
||||
}
|
||||
println!("Streams:");
|
||||
for stream in &snapshot.mq.streams {
|
||||
println!(
|
||||
" - {} subjects={} lag={} dlq={}",
|
||||
stream.name,
|
||||
stream.subjects.join(","),
|
||||
stream
|
||||
.consumer_lag
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_else(|| "-".to_string()),
|
||||
stream
|
||||
.dlq_messages
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_else(|| "-".to_string())
|
||||
);
|
||||
}
|
||||
println!("Workers:");
|
||||
if snapshot.workers.is_empty() {
|
||||
println!(" - none");
|
||||
} else {
|
||||
for worker in &snapshot.workers {
|
||||
println!(
|
||||
" - {} ({}) {:?} pid={} last_error={}",
|
||||
worker.name,
|
||||
worker.role,
|
||||
worker.phase,
|
||||
worker
|
||||
.pid
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_else(|| "-".to_string()),
|
||||
worker.last_error.as_deref().unwrap_or("-")
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Some(error) = &snapshot.last_error {
|
||||
println!("最近错误: {error}");
|
||||
}
|
||||
}
|
||||
|
||||
fn phase_name(phase: &ServicePhase) -> &'static str {
|
||||
match phase {
|
||||
ServicePhase::Starting => "starting",
|
||||
ServicePhase::Running => "running",
|
||||
ServicePhase::Degraded => "degraded",
|
||||
ServicePhase::Stopping => "stopping",
|
||||
ServicePhase::Stopped => "stopped",
|
||||
ServicePhase::Failed => "failed",
|
||||
}
|
||||
}
|
||||
|
||||
fn send_sigterm(pid: u32) -> anyhow::Result<()> {
|
||||
let status = Command::new("kill")
|
||||
.arg("-TERM")
|
||||
.arg(pid.to_string())
|
||||
.status()?;
|
||||
if !status.success() {
|
||||
anyhow::bail!("发送 SIGTERM 失败，PID: {pid}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn detach_background_child(command: &mut Command) {
|
||||
unsafe {
|
||||
command.pre_exec(|| {
|
||||
if libc::setsid() == -1 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn detach_background_child(_command: &mut Command) {}
|
||||
@@ -0,0 +1,19 @@
|
||||
#[test]
|
||||
fn daemon_does_not_depend_on_business_crates() {
|
||||
let manifest = std::fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/Cargo.toml"))
|
||||
.expect("read ias-daemon Cargo.toml");
|
||||
// 解析 [dependencies] 表而不是字符串匹配，避免注释/路径里的同名子串误报，
|
||||
// 也能阻止以其他形式（如 build-dependencies）绕过业务 crate 依赖边界。
|
||||
let parsed: toml::Value = toml::from_str(&manifest).expect("parse ias-daemon Cargo.toml");
|
||||
let deps = parsed
|
||||
.get("dependencies")
|
||||
.and_then(|value| value.as_table())
|
||||
.expect("ias-daemon Cargo.toml 缺少 [dependencies] 表");
|
||||
|
||||
for forbidden in ["ias-ai", "ias-context", "ias-wechat", "ias-mail"] {
|
||||
assert!(
|
||||
!deps.contains_key(forbidden),
|
||||
"ias-daemon 不得依赖业务 crate: {forbidden}"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user