From c96a4dd10ec940d08e99e802b32cbc6692cd89f7 Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Thu, 20 Jul 2023 13:04:23 -0700 Subject: [PATCH] refactor: move startup/lifecycle to server-common --- Cargo.lock | 4 ++ crates/integration-tests/src/kumod.rs | 2 +- crates/kumo-server-common/Cargo.toml | 4 ++ crates/kumo-server-common/src/lib.rs | 1 + crates/kumo-server-common/src/start.rs | 68 ++++++++++++++++++++ crates/kumod/Cargo.toml | 4 +- crates/kumod/src/logging.rs | 32 +++++---- crates/kumod/src/main.rs | 89 +++++++++----------------- 8 files changed, 130 insertions(+), 74 deletions(-) create mode 100644 crates/kumo-server-common/src/start.rs diff --git a/Cargo.lock b/Cargo.lock index 0c771354..3153802c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2258,6 +2258,9 @@ dependencies = [ "data-loader", "domain-map", "json_comments", + "kumo-server-lifecycle", + "kumo-server-memory", + "kumo-server-runtime", "metrics", "metrics-prometheus", "metrics-tracing-context", @@ -2276,6 +2279,7 @@ dependencies = [ "throttle", "tokio", "toml", + "tracing", "tracing-appender", "tracing-subscriber", ] diff --git a/crates/integration-tests/src/kumod.rs b/crates/integration-tests/src/kumod.rs index ff9f0289..67ebe64d 100644 --- a/crates/integration-tests/src/kumod.rs +++ b/crates/integration-tests/src/kumod.rs @@ -285,7 +285,7 @@ impl KumoDaemon { let mut child = Command::new(&path) .args(["--policy", &args.policy_file, "--user", &user.name]) - .env("KUMOD_LOG", "kumod=trace") + .env("KUMOD_LOG", "kumod=trace,kumo_server_common=info") .env("KUMOD_TEST_DIR", dir.path()) .envs(args.env.iter().cloned()) .stdout(Stdio::piped()) diff --git a/crates/kumo-server-common/Cargo.toml b/crates/kumo-server-common/Cargo.toml index 1da38a46..ce1bee4b 100644 --- a/crates/kumo-server-common/Cargo.toml +++ b/crates/kumo-server-common/Cargo.toml @@ -13,6 +13,9 @@ console-subscriber = "0.1" data-loader = {path="../data-loader"} domain-map = {path="../domain-map"} json_comments = "0.2" +kumo-server-lifecycle = {path="../kumo-server-lifecycle"} +kumo-server-memory = {path="../kumo-server-memory"} +kumo-server-runtime = {path="../kumo-server-runtime"} metrics = "0.20" metrics-prometheus = "0.3" metrics-tracing-context = "0.13" @@ -31,5 +34,6 @@ serde_json = "1.0" throttle = {path="../throttle"} tokio = {version="1.25", features=["full", "tracing"]} toml = "0.7" +tracing = "0.1" tracing-appender = "0.2" tracing-subscriber = {version="0.3", features=["env-filter", "std", "fmt", "json"]} diff --git a/crates/kumo-server-common/src/lib.rs b/crates/kumo-server-common/src/lib.rs index 4a0d732d..bed79507 100644 --- a/crates/kumo-server-common/src/lib.rs +++ b/crates/kumo-server-common/src/lib.rs @@ -4,6 +4,7 @@ use mlua::{Function, Lua, LuaSerdeExt, Value}; use mod_redis::RedisConnKey; pub mod diagnostic_logging; +pub mod start; pub mod tls_helpers; pub fn register(lua: &Lua) -> anyhow::Result<()> { diff --git a/crates/kumo-server-common/src/start.rs b/crates/kumo-server-common/src/start.rs new file mode 100644 index 00000000..02b9ee71 --- /dev/null +++ b/crates/kumo-server-common/src/start.rs @@ -0,0 +1,68 @@ +use crate::diagnostic_logging::LoggingConfig; +use config::RegisterFunc; +use kumo_server_lifecycle::LifeCycle; +use kumo_server_runtime::rt_spawn; +use std::future::Future; +use std::path::Path; +use std::pin::Pin; + +pub struct StartConfig<'a> { + pub logging: LoggingConfig<'a>, + pub lua_funcs: &'a [RegisterFunc], + pub policy: &'a Path, +} + +impl<'a> StartConfig<'a> { + pub async fn run( + self, + perform_init: INIT, + broadcast_shutdown: FINI, + ) -> anyhow::Result<()> + where + INIT: FnOnce() -> Pin>>> + Send + 'static, + FINI: FnOnce() -> Pin>> + Send + 'static, + { + self.logging.init()?; + + kumo_server_memory::setup_memory_limit()?; + + for &func in self.lua_funcs { + config::register(func); + } + + config::set_policy_path(self.policy.to_path_buf()).await?; + + let mut life_cycle = LifeCycle::new(); + + let init_handle = rt_spawn("initialize".to_string(), move || { + Ok(async move { + let mut ok = true; + let init_future = (perform_init)(); + if let Err(err) = init_future.await { + tracing::error!("problem initializing: {err:#}"); + LifeCycle::request_shutdown().await; + ok = false; + } + // This log line is depended upon by the integration + // test harness. Do not change or remove it without + // making appropriate adjustments over there! + tracing::info!("initialization complete"); + ok + }) + }) + .await?; + + life_cycle.wait_for_shutdown().await; + + // after waiting for those to idle out, shut down logging + let shutdown_future = (broadcast_shutdown)(); + shutdown_future.await; + + tracing::info!("Shutdown completed OK!"); + + if !init_handle.await? { + anyhow::bail!("Initialization raised an error"); + } + Ok(()) + } +} diff --git a/crates/kumod/Cargo.toml b/crates/kumod/Cargo.toml index 45a46dda..c052f524 100644 --- a/crates/kumod/Cargo.toml +++ b/crates/kumod/Cargo.toml @@ -28,6 +28,7 @@ flume = "0.10" gcd = "2.3" gethostname = "0.4" humansize = "2.1" # for printing +humantime-serde = "1.1" kumo-api-types = {path="../kumo-api-types"} kumo-log-types = {path="../kumo-log-types"} kumo-server-common = {path="../kumo-server-common"} @@ -51,9 +52,8 @@ prometheus = "0.13" rand = "0.8" rfc5321 = {path="../rfc5321"} rustls = "0.20" -serde = {version="1.0", features=["derive"]} -humantime-serde = "1.1" self_cell = "1.0" +serde = {version="1.0", features=["derive"]} serde_json = "1.0" socksv5 = {version="0.3", default-features=false, features=["tokio"]} spool = {path="../spool"} diff --git a/crates/kumod/src/logging.rs b/crates/kumod/src/logging.rs index be6418f9..227aac36 100644 --- a/crates/kumod/src/logging.rs +++ b/crates/kumod/src/logging.rs @@ -18,9 +18,11 @@ use serde_json::Value; use spool::SpoolId; use std::collections::HashMap; use std::fs::File; +use std::future::Future; use std::io::Write; use std::net::Ipv4Addr; use std::path::PathBuf; +use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::thread::JoinHandle; use std::time::{Duration, Instant}; @@ -305,20 +307,22 @@ impl Logger { Ok(self.sender.send(LogCommand::Record(record)).await?) } - pub async fn signal_shutdown() { - let loggers = Self::get_loggers(); - for logger in loggers.iter() { - tracing::debug!("Terminating a logger"); - logger.sender.send(LogCommand::Terminate).await.ok(); - tracing::debug!("Joining that logger"); - let res = logger - .thread - .lock() - .await - .take() - .map(|thread| thread.join()); - tracing::debug!("Joined -> {res:?}"); - } + pub fn signal_shutdown() -> Pin>> { + Box::pin(async move { + let loggers = Self::get_loggers(); + for logger in loggers.iter() { + tracing::debug!("Terminating a logger"); + logger.sender.send(LogCommand::Terminate).await.ok(); + tracing::debug!("Joining that logger"); + let res = logger + .thread + .lock() + .await + .take() + .map(|thread| thread.join()); + tracing::debug!("Joined -> {res:?}"); + } + }) } pub async fn extract_fields( diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index 3faae664..2fc44861 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -2,12 +2,14 @@ use anyhow::Context; use caps::{CapSet, Capability, CapsHashSet}; use clap::Parser; use kumo_server_common::diagnostic_logging::{DiagnosticFormat, LoggingConfig}; -use kumo_server_lifecycle::LifeCycle; +use kumo_server_common::start::StartConfig; use kumo_server_runtime::rt_spawn; use nix::sys::resource::{getrlimit, setrlimit, Resource}; use nix::sys::signal::{kill, SIGQUIT}; use nix::unistd::{Pid, Uid, User}; +use std::future::Future; use std::path::PathBuf; +use std::pin::Pin; mod delivery_metrics; mod egress_path; @@ -129,67 +131,40 @@ fn main() -> anyhow::Result<()> { .block_on(async move { run(opts).await }) } -async fn perform_init() -> anyhow::Result<()> { - let mut config = config::load_config().await?; - config.async_call_callback("init", ()).await?; +fn perform_init() -> Pin>>> { + Box::pin(async move { + let mut config = config::load_config().await?; + config.async_call_callback("init", ()).await?; - crate::spool::SpoolManager::get().await.start_spool().await + crate::spool::SpoolManager::get() + .await + .start_spool() + .await?; + + Ok(()) + }) } async fn run(opts: Opt) -> anyhow::Result<()> { - LoggingConfig { - log_dir: opts.diag_log_dir.clone(), - diag_format: opts.diag_format, - tokio_console: opts.tokio_console, - filter_env_var: "KUMOD_LOG", - default_filter: "kumod=info", + StartConfig { + logging: LoggingConfig { + log_dir: opts.diag_log_dir.clone(), + diag_format: opts.diag_format, + tokio_console: opts.tokio_console, + filter_env_var: "KUMOD_LOG", + default_filter: "kumod=info,kumo_server_common=info", + }, + lua_funcs: &[ + kumo_server_common::register, + crate::mod_kumo::register, + crate::spool::register, + crate::logging::register, + message::dkim::register, + ], + policy: &opts.policy, } - .init()?; - - kumo_server_memory::setup_memory_limit()?; - - for func in [ - kumo_server_common::register, - crate::mod_kumo::register, - crate::spool::register, - crate::logging::register, - message::dkim::register, - ] { - config::register(func); - } - - config::set_policy_path(opts.policy.clone()).await?; - - let mut life_cycle = LifeCycle::new(); - - let init_handle = rt_spawn("initialize".to_string(), move || { - Ok(async move { - let mut ok = true; - if let Err(err) = perform_init().await { - tracing::error!("problem initializing: {err:#}"); - LifeCycle::request_shutdown().await; - ok = false; - } - // This log line is depended upon by the integration - // test harness. Do not change or remove it without - // making appropriate adjustments over there! - tracing::info!("initialization complete"); - ok - }) - }) - .await?; - - life_cycle.wait_for_shutdown().await; - - // after waiting for those to idle out, shut down logging - crate::logging::Logger::signal_shutdown().await; - - tracing::info!("Shutdown completed OK!"); - - if !init_handle.await? { - anyhow::bail!("Initialization raised an error"); - } - Ok(()) + .run(perform_init, crate::logging::Logger::signal_shutdown) + .await } fn register_panic_hook() {