#![warn(missing_docs)]
use std::path::Path;
use std::pin::pin;
use std::process;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::{env, mem};
use conjure_error::Error;
use conjure_http::server::{AsyncService, ConjureRuntime};
use conjure_runtime::{Agent, ClientFactory, HostMetricsRegistry, UserAgent};
use debug::endpoint::DebugResource;
use debug::endpoint::DebugServiceEndpoints;
#[cfg(feature = "jemalloc")]
use debug::heap_profile::{self, HeapProfileDiagnostic};
use futures_util::{stream, Stream, StreamExt};
use refreshable::Refreshable;
use serde::de::DeserializeOwned;
use status::StatusResource;
use status::StatusServiceEndpoints;
use tokio::runtime::{Handle, Runtime};
use tokio::signal::unix::{self, SignalKind};
use tokio::{runtime, select, time};
use witchcraft_log::{error, fatal, info};
use witchcraft_metrics::MetricRegistry;
pub use body::{RequestBody, ResponseWriter};
use config::install::InstallConfig;
use config::runtime::RuntimeConfig;
pub use witchcraft::Witchcraft;
#[doc(inline)]
pub use witchcraft_server_config as config;
#[doc(inline)]
pub use witchcraft_server_macros::main;
use crate::debug::diagnostic_types::DiagnosticTypesDiagnostic;
#[cfg(feature = "jemalloc")]
use crate::debug::heap_stats::HeapStatsDiagnostic;
use crate::debug::metric_names::MetricNamesDiagnostic;
#[cfg(target_os = "linux")]
use crate::debug::thread_dump::ThreadDumpDiagnostic;
use crate::debug::DiagnosticRegistry;
use crate::health::config_reload::ConfigReloadHealthCheck;
use crate::health::endpoint_500s::Endpoint500sHealthCheck;
use crate::health::minidump::MinidumpHealthCheck;
use crate::health::panics::PanicsHealthCheck;
use crate::health::service_dependency::ServiceDependencyHealthCheck;
use crate::health::HealthCheckRegistry;
use crate::readiness::ReadinessCheckRegistry;
use crate::server::Listener;
use crate::shutdown_hooks::ShutdownHooks;
pub mod blocking;
mod body;
mod configs;
pub mod debug;
mod endpoint;
pub mod extensions;
pub mod health;
pub mod logging;
mod metrics;
mod minidump;
pub mod readiness;
mod server;
mod service;
mod shutdown_hooks;
mod status;
pub mod tls;
mod witchcraft;
#[cfg(feature = "jemalloc")]
#[global_allocator]
static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
pub fn init<I, R, F>(init: F)
where
I: AsRef<InstallConfig> + DeserializeOwned,
R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
{
init_with_configs(init, configs::load_install::<I>, configs::load_runtime::<R>)
}
pub fn init_with_configs<I, R, F, LI, LR>(init: F, load_install: LI, load_runtime: LR)
where
I: AsRef<InstallConfig> + DeserializeOwned,
R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
LI: FnOnce() -> Result<I, Error>,
LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
{
logging::early_init();
let args = env::args_os().collect::<Vec<_>>();
if args.len() == 3 && args[1] == "minidump" {
let ret = match minidump::server(Path::new(&args[2])) {
Ok(()) => 0,
Err(e) => {
fatal!("error starting minidump server", error: e);
1
}
};
process::exit(ret);
} else {
let mut runtime_guard = None;
let ret = match init_inner(
init,
load_install,
load_runtime,
&mut runtime_guard,
InitOptions::default(),
) {
Ok(witchcraft) => {
match witchcraft.handle.block_on(shutdown(
witchcraft.shutdown_hooks,
witchcraft.install_config.server().shutdown_timeout(),
)) {
Ok(_) => 0,
Err(e) => {
fatal!("error after starting server", error: e);
1
}
}
}
Err(e) => {
fatal!("error starting server", error: e);
1
}
};
drop(runtime_guard);
process::exit(ret);
}
}
#[cfg(feature = "in-memory-testing")]
pub mod in_memory_testing {
use crate::{drain_shutdown_hooks, init_inner, logging, InitOptions, Witchcraft};
use conjure_error::Error;
use refreshable::Refreshable;
use serde::de::DeserializeOwned;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{mpsc, Arc, Once};
use std::thread::{self, JoinHandle};
use tokio::runtime::Handle;
use tokio::sync::oneshot;
use witchcraft_server_config::install::InstallConfig;
use witchcraft_server_config::runtime::RuntimeConfig;
static INIT: Once = Once::new();
static FIRST_INIT_SUCCESS: AtomicBool = AtomicBool::new(false);
pub fn init_with_configs_for_tests<I, R, F, LI, LR>(
init: F,
load_install: LI,
load_runtime: LR,
thread_prefix: Option<String>,
) -> Result<RunHandle, Error>
where
I: AsRef<InstallConfig> + DeserializeOwned,
R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error> + Send + 'static,
LI: FnOnce() -> Result<I, Error> + Send + 'static,
LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>
+ Send
+ 'static,
{
let init_fn = |init_globals: bool| {
if init_globals {
logging::early_init();
}
let (startup_result_sender, startup_result_receiver) =
mpsc::channel::<Result<(), Error>>();
let (shutdown_signal_sender, shutdown_signal_receiver) = oneshot::channel::<()>();
let server_thread = thread::spawn(move || {
let mut runtime_guard = None;
let witchcraft = match init_inner(
init,
load_install,
load_runtime,
&mut runtime_guard,
InitOptions {
init_globals,
thread_prefix,
},
) {
Ok(witchcraft) => witchcraft,
Err(e) => {
let _ = startup_result_sender.send(Err(e));
return;
}
};
let _ = startup_result_sender.send(Ok(()));
let handle = witchcraft.handle.clone();
let timeout = witchcraft.install_config.server().shutdown_timeout();
handle.block_on(async move {
let _ = shutdown_signal_receiver.await;
drain_shutdown_hooks(witchcraft.shutdown_hooks, timeout).await;
});
drop(runtime_guard);
});
match startup_result_receiver.recv() {
Ok(Ok(())) => Ok(RunHandle {
shutdown_signal_sender: Some(shutdown_signal_sender),
server_thread: Some(server_thread),
}),
Ok(Err(e)) => Err(e), Err(_) => Err(Error::internal_safe(
"in-memory server thread panicked during initialization",
)),
}
};
let mut init_fn = Some(init_fn);
let mut first_init_result: Option<Result<RunHandle, Error>> = None;
INIT.call_once(|| {
let res = init_fn.take().unwrap()(true);
FIRST_INIT_SUCCESS.store(res.is_ok(), Ordering::Relaxed);
first_init_result = Some(res);
});
first_init_result.unwrap_or_else(|| {
if !FIRST_INIT_SUCCESS.load(Ordering::Relaxed) {
return Err(Error::internal_safe("Initial init failed, not continuing"));
}
init_fn.take().unwrap()(false)
})
}
pub struct RunHandle {
shutdown_signal_sender: Option<oneshot::Sender<()>>,
server_thread: Option<JoinHandle<()>>,
}
impl Drop for RunHandle {
fn drop(&mut self) {
let _ = self.shutdown_signal_sender.take().unwrap().send(());
let _ = self.server_thread.take().unwrap().join();
}
}
}
struct InitOptions {
init_globals: bool,
thread_prefix: Option<String>,
}
impl Default for InitOptions {
fn default() -> Self {
Self {
init_globals: true,
thread_prefix: None,
}
}
}
fn init_inner<I, R, F, LI, LR>(
init: F,
load_install: LI,
load_runtime: LR,
runtime_guard: &mut Option<RuntimeGuard>,
init_options: InitOptions,
) -> Result<Witchcraft, Error>
where
I: AsRef<InstallConfig> + DeserializeOwned,
R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
LI: FnOnce() -> Result<I, Error>,
LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
{
let install_config = load_install()?;
let thread_id = AtomicUsize::new(0);
let thread_prefix = init_options
.thread_prefix
.clone()
.unwrap_or("runtime".to_string());
let runtime = runtime::Builder::new_multi_thread()
.enable_all()
.thread_name_fn(move || {
format!(
"{}-{}",
thread_prefix,
thread_id.fetch_add(1, Ordering::Relaxed)
)
})
.worker_threads(install_config.as_ref().server().io_threads())
.thread_keep_alive(install_config.as_ref().server().idle_thread_timeout())
.build()
.map_err(Error::internal_safe)?;
let handle = runtime.handle().clone();
let runtime = runtime_guard.insert(RuntimeGuard {
runtime: Some(runtime),
logger_shutdown: Some(ShutdownHooks::new()),
});
let runtime_config_ok = Arc::new(AtomicBool::new(true));
let runtime_config = load_runtime(&handle, &runtime_config_ok)?;
let metrics = Arc::new(MetricRegistry::new());
let host_metrics = Arc::new(HostMetricsRegistry::new());
let health_checks = Arc::new(HealthCheckRegistry::new(&handle));
let readiness_checks = Arc::new(ReadinessCheckRegistry::new());
let diagnostics = Arc::new(DiagnosticRegistry::new());
let loggers = if init_options.init_globals {
handle.block_on(logging::init(
&metrics,
install_config.as_ref(),
&runtime_config.map(|c| c.as_ref().logging().clone()),
runtime.logger_shutdown.as_mut().unwrap(),
))
} else {
logging::get_existing()
}?;
info!("server starting");
if init_options.init_globals && install_config.as_ref().minidump().enabled() {
let minidump_ok = Arc::new(AtomicBool::new(true));
let socket_dir = install_config.as_ref().minidump().socket_dir();
handle.spawn({
let minidump_ok = minidump_ok.clone();
let socket_dir = socket_dir.to_owned();
async move {
let result = minidump::init(&socket_dir).await;
minidump_ok.store(result.is_ok(), Ordering::Relaxed);
if let Err(e) = result {
error!("error during minidump init", error: e)
}
}
});
health_checks.register(MinidumpHealthCheck::new(minidump_ok));
#[cfg(target_os = "linux")]
diagnostics.register(ThreadDumpDiagnostic::new(socket_dir));
}
metrics::init(&metrics);
health_checks.register(ServiceDependencyHealthCheck::new(&host_metrics));
health_checks.register(PanicsHealthCheck::new());
health_checks.register(ConfigReloadHealthCheck::new(runtime_config_ok));
diagnostics.register(MetricNamesDiagnostic::new(&metrics));
#[cfg(feature = "jemalloc")]
{
if init_options.init_globals {
diagnostics.register(HeapStatsDiagnostic);
heap_profile::init(&runtime_config);
diagnostics.register(HeapProfileDiagnostic);
}
}
diagnostics.register(DiagnosticTypesDiagnostic::new(Arc::downgrade(&diagnostics)));
let client_factory = ClientFactory::builder()
.config(runtime_config.map(|c| c.as_ref().service_discovery().clone()))
.user_agent(UserAgent::new(Agent::new(
install_config.as_ref().product_name(),
install_config.as_ref().product_version(),
)))
.metrics(metrics.clone())
.host_metrics(host_metrics)
.blocking_handle(handle.clone());
let mut witchcraft = Witchcraft {
metrics,
health_checks,
readiness_checks,
client_factory,
diagnostics: diagnostics.clone(),
handle: handle.clone(),
install_config: install_config.as_ref().clone(),
thread_pool: None,
thread_prefix: init_options.thread_prefix,
endpoints: vec![],
shutdown_hooks: ShutdownHooks::new(),
conjure_runtime: Arc::new(ConjureRuntime::new()),
};
let status_endpoints = StatusServiceEndpoints::new(StatusResource::new(
&runtime_config,
&witchcraft.health_checks,
&witchcraft.readiness_checks,
));
witchcraft.endpoints(
None,
status_endpoints.endpoints(&witchcraft.conjure_runtime),
false,
);
let debug_endpoints =
DebugServiceEndpoints::new(DebugResource::new(&runtime_config, &diagnostics));
witchcraft.app(debug_endpoints);
let management_endpoints = mem::take(&mut witchcraft.endpoints);
init(install_config, runtime_config, &mut witchcraft)?;
witchcraft
.health_checks
.register(Endpoint500sHealthCheck::new(&witchcraft.endpoints));
let mut main_endpoints = mem::take(&mut witchcraft.endpoints);
match witchcraft.install_config.management_port() {
Some(management_port) if management_port != witchcraft.install_config.port() => {
handle.block_on(server::start(
&mut witchcraft,
management_endpoints,
&loggers,
Listener::Management,
management_port,
))?;
}
_ => main_endpoints.extend(management_endpoints),
}
let port = witchcraft.install_config.port();
handle.block_on(server::start(
&mut witchcraft,
main_endpoints,
&loggers,
Listener::Service,
port,
))?;
Ok(witchcraft)
}
async fn shutdown(shutdown_hooks: ShutdownHooks, timeout: Duration) -> Result<(), Error> {
let mut signals = pin!(signals()?);
signals.next().await;
select! {
_ = drain_shutdown_hooks(shutdown_hooks, timeout) => {}
_ = signals.next() => info!("graceful shutdown interrupted by signal"),
}
Ok(())
}
async fn drain_shutdown_hooks(shutdown_hooks: ShutdownHooks, timeout: Duration) {
info!("server shutting down");
select! {
_ = shutdown_hooks => {}
_ = time::sleep(timeout) => {
info!(
"graceful shutdown timed out",
safe: {
timeout: format_args!("{timeout:?}"),
},
);
}
}
}
fn signals() -> Result<impl Stream<Item = ()>, Error> {
let sigint = signal(SignalKind::interrupt())?;
let sigterm = signal(SignalKind::terminate())?;
Ok(stream::select(sigint, sigterm))
}
fn signal(kind: SignalKind) -> Result<impl Stream<Item = ()>, Error> {
let mut signal = unix::signal(kind).map_err(Error::internal_safe)?;
Ok(stream::poll_fn(move |cx| signal.poll_recv(cx)))
}
struct RuntimeGuard {
runtime: Option<Runtime>,
logger_shutdown: Option<ShutdownHooks>,
}
impl Drop for RuntimeGuard {
fn drop(&mut self) {
let runtime = self.runtime.take().unwrap();
runtime.block_on(self.logger_shutdown.take().unwrap());
runtime.shutdown_background()
}
}