use futures::FutureExt;
use hydrant::config::Config;
use hydrant::control::{ApiBind, ApiBinds, DatabasePoisoned, Hydrant};
use miette::IntoDiagnostic;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::Notify;
use tracing::{error, info};
mod allocator;
const DEFAULT_API_PORT: u16 = 3000;
const DEFAULT_DEBUG_PORT: u16 = DEFAULT_API_PORT + 1;
fn default_api_binds() -> ApiBinds {
ApiBinds::try_from_iter([
SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), DEFAULT_API_PORT),
SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), DEFAULT_API_PORT),
])
.expect("two-element literal is non-empty")
}
const USAGE: &str = "\
usage: hydrant [--version | --help]
hydrant takes no other arguments, it reads its config from HYDRANT_* env vars.
see docs/configuration.md for the full list.
";
#[derive(Debug, PartialEq, Eq)]
enum Command {
Run,
Version,
Help,
}
fn parse_args(mut args: impl Iterator<Item = std::ffi::OsString>) -> miette::Result<Command> {
let Some(arg) = args.next() else {
return Ok(Command::Run);
};
let command = match arg.to_str() {
Some("--version" | "-V") => Command::Version,
Some("--help" | "-h") => Command::Help,
_ => miette::bail!(
help = "config goes in HYDRANT_* env vars, see `hydrant --help`",
"unexpected argument {}",
arg.to_string_lossy()
),
};
if let Some(extra) = args.next() {
miette::bail!("unexpected argument {}", extra.to_string_lossy());
}
Ok(command)
}
struct AppConfig {
api_binds: Option<ApiBinds>,
enable_debug: bool,
debug_port: u16,
}
impl AppConfig {
fn from_env() -> miette::Result<Self> {
use hydrant::__cfg as cfg;
let env = |key: &str| std::env::var(key).ok();
let api_binds = ApiBinds::from_lookup(env, Some(default_api_binds()))?;
let enable_debug = cfg!(env, "ENABLE_DEBUG", false);
let debug_port_default = api_binds
.iter()
.flat_map(ApiBinds::iter)
.find_map(ApiBind::tcp)
.and_then(|addr| addr.port().checked_add(1))
.unwrap_or(DEFAULT_DEBUG_PORT);
let debug_port = cfg!(env, "DEBUG_PORT", debug_port_default);
Ok(Self {
api_binds,
enable_debug,
debug_port,
})
}
}
// long enough for a busy database to finish what it's writing, short enough that a
// supervisor's own stop timeout (systemd's is 90s) doesn't kill it first
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);
const DRAIN_TIMEOUT: Duration = Duration::from_secs(2);
fn main() -> miette::Result<()> {
// config lives in env vars, so a stray argument (a `--version` probe, a typo) would
// otherwise boot a full default instance without anyone noticing
match parse_args(std::env::args_os().skip(1))? {
Command::Run => {}
Command::Version => {
println!("hydrant {}", env!("CARGO_PKG_VERSION"));
return Ok(());
}
Command::Help => {
print!("{USAGE}");
return Ok(());
}
}
rustls::crypto::aws_lc_rs::default_provider()
.install_default()
.ok();
// logging starts before the config is read, so what reading it warns about gets
// printed, and after `.env` is loaded, since RUST_LOG may be set there
hydrant::config::env::load_dotenv();
let directives = std::env::var("RUST_LOG")
.ok()
.filter(|v| !v.trim().is_empty());
let env_filter = tracing_subscriber::EnvFilter::builder()
.with_default_directive(tracing::Level::INFO.into())
.parse_lossy(directives.as_deref().unwrap_or(hydrant::DEFAULT_LOG_FILTER));
tracing_subscriber::fmt().with_env_filter(env_filter).init();
let cfg = Config::from_lookup(|key| std::env::var(key).ok())?;
let app = AppConfig::from_env()?;
let stop = watch_stop_signals()?;
let runtime = tokio::runtime::Runtime::new().into_diagnostic()?;
let hydrant = runtime.block_on(Hydrant::new(cfg))?;
let result = runtime.block_on(async {
let api_fut = app
.api_binds
.map(|binds| hydrant.serve(binds).boxed())
.unwrap_or_else(|| std::future::pending().boxed());
let debug_fut = app
.enable_debug
.then(|| hydrant.serve_debug(app.debug_port).boxed())
.unwrap_or_else(|| std::future::pending().boxed());
let result = tokio::select! {
r = hydrant.run()? => r,
r = api_fut => r,
r = debug_fut => r,
_ = stop.notified() => Ok(()),
};
// the select dropped the listeners. reconnects go to the next hydrant
hydrant.drain_streams(DRAIN_TIMEOUT).await;
result
});
// dropping the runtime would wait on blocking tasks that only `Hydrant::shutdown` ends
runtime.shutdown_background();
info!("stopping, waiting up to {SHUTDOWN_TIMEOUT:?} for the database to close");
let closed = hydrant.shutdown(SHUTDOWN_TIMEOUT);
if let (Err(_), Err(e)) = (&result, &closed) {
error!("{e}");
}
// a supervisor can tell this apart from other failures by the exit code
if let Err(e) = &result
&& e.downcast_ref::<DatabasePoisoned>().is_some()
{
eprintln!("{e:?}");
std::process::exit(10);
}
result?;
closed?;
info!("stopped");
Ok(())
}
/// the first SIGINT or SIGTERM fires the returned notify, a second one exits on the spot.
/// it gets its own thread and runtime, since the main runtime is shut down while the
/// database closes and a signal then still has to reach someone.
fn watch_stop_signals() -> miette::Result<Arc<Notify>> {
let stop = Arc::new(Notify::new());
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_io()
.build()
.into_diagnostic()?;
// registered before returning, so a signal right after startup isn't the default kill
let (mut int, mut term) = runtime
.block_on(async {
Ok::<_, std::io::Error>((
signal(SignalKind::interrupt())?,
signal(SignalKind::terminate())?,
))
})
.into_diagnostic()?;
std::thread::Builder::new()
.name("signals".into())
.spawn({
let stop = stop.clone();
move || {
runtime.block_on(async {
let mut next = async || {
tokio::select! {
_ = int.recv() => {}
_ = term.recv() => {}
}
};
next().await;
info!("got a stop signal, stopping (send another to exit right away)");
stop.notify_one();
next().await;
eprintln!("got a second stop signal, exiting without closing the database");
std::process::exit(1);
})
}
})
.into_diagnostic()?;
Ok(stop)
}