diff --git a/src/daemon.rs b/src/daemon.rs index efb8793..1c493ad 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -14,12 +14,37 @@ use std::fs::File; +use anyhow::Result; use daemonize::Daemonize; +use crate::config::{DaemonConfig, DaemonWasmConfig}; +/// run the process as daemon +pub fn spawn_daemon(cfg: DaemonConfig) -> Result<()> { + let stdout = File::create(cfg.stdout)?; + let stderr = File::create(cfg.stderr)?; + let daemonize = Daemonize::new() + .pid_file(cfg.pid_file) + .chown_pid_file(true) // is optional, see `Daemonize` documentation + .working_directory(cfg.workdir) + .user("root") + .group("root") + .umask(0o027) // Set umask, `0o027` by default. + .stdout(stdout) // by default, stdout is redirect to `/tmp/psh.stdout`. + .stderr(stderr); // by default, stderr is redirect to `/tmp/psh.stderr`. + daemonize.start()?; + Ok(()) +} +pub fn get_daemon_wasm_args(cfg: DaemonWasmConfig) -> Option> { + if !cfg.enable { + return None; } + let mut vec = Vec::with_capacity(cfg.args.len() + 1); + vec.push(cfg.path); + vec.extend(cfg.args); + Some(vec) } diff --git a/src/main.rs b/src/main.rs index 4452579..63293f2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -31,7 +31,8 @@ use args::Args; use chrono::offset::LocalResult; use chrono::{TimeZone, Utc}; use clap::Parser; -use config::{PshConfig, RemoteConfig}; +use config::RemoteConfig; +use daemon::{get_daemon_wasm_args, spawn_daemon}; use log::log_init; use opentelemetry_otlp::ExportConfig; use runtime::{Task, TaskRuntime}; @@ -48,20 +49,17 @@ fn main() -> Result<()> { } let mut args = Args::parse(); - let mut psh_config = PshConfig::read_config(&args.config).unwrap_or_else(|e| { - tracing::warn!("{e}, use default Psh config."); - PshConfig::default() - }); + let cfg = config::read_or_gen(args.config.clone())?; // When running as a daemon, it ignores all other cli arguments let component_args = if args.systemd() || args.daemon() { - psh_config.get_component_args() + get_daemon_wasm_args(cfg.daemon.wasm.clone()) } else { args.get_component_args() }; if args.daemon() { - daemon::Daemon::new(psh_config.daemon().clone()).daemon()?; + spawn_daemon(cfg.daemon)?; } let mut task_rt = TaskRuntime::new()?; @@ -78,7 +76,7 @@ fn main() -> Result<()> { thread::spawn(move || -> Result<()> { let rt = tokio::runtime::Runtime::new()?; - let tasks = async_tasks(psh_config.remote.clone(), psh_config.take_token(), task_rt); + let tasks = async_tasks(cfg.remote, task_rt); rt.block_on(tasks)?; Ok(()) }) @@ -88,13 +86,8 @@ fn main() -> Result<()> { Ok(()) } -async fn async_tasks( - remote_cfg: RemoteConfig, - token: String, - mut task_rt: TaskRuntime, -) -> Result<()> { - let token_ = token.clone(); - +async fn async_tasks(remote_cfg: RemoteConfig, mut task_rt: TaskRuntime) -> Result<()> { + let token_cloned = remote_cfg.token.clone(); let rpc_task = async move { if !remote_cfg.enable { let handle = task_rt.spawn(None)?; @@ -102,8 +95,8 @@ async fn async_tasks( return Ok(()); } - let duration = Duration::from_secs(remote_cfg.rpc.duration); - let mut client = RpcClient::new(remote_cfg.rpc, token_).await?; + let duration = Duration::from_secs(remote_cfg.rpc.heartbeat_interval); + let mut client = RpcClient::new(remote_cfg.rpc, token_cloned).await?; task_rt.spawn(Some(client.clone()))?; client.send_info().await?; loop { @@ -135,8 +128,11 @@ async fn async_tasks( return Ok(()); } - let export_conf: ExportConfig = remote_cfg.otlp_conf.into(); - otlp::otlp_tasks(export_conf, token).await?; + let export_conf = ExportConfig { + endpoint: Some(remote_cfg.otlp.addr), + ..Default::default() + }; + otlp::otlp_tasks(export_conf, remote_cfg.token).await?; Ok::<(), Error>(()) }; diff --git a/src/services/rpc.rs b/src/services/rpc.rs index 596d80b..380477f 100644 --- a/src/services/rpc.rs +++ b/src/services/rpc.rs @@ -16,8 +16,8 @@ use anyhow::Result; use tonic::transport::{Channel, ClientTlsConfig, Endpoint}; use tonic::Request; -use super::config::RpcConfig; use super::pb::{self, DataRequest}; +use crate::config::RpcConfig; use crate::services::host_info::RawInfo; use crate::services::pb::psh_service_client::PshServiceClient; use crate::services::pb::HostInfoRequest;