Hot-Reloading Refactor

Signed-off-by: Jason Volk <jason@zemos.net>

src/router/run.rs

@@ -0,0 +1,185 @@
+use std::{sync::Arc, time::Duration};
+
+use axum_server::Handle as ServerHandle;
+use tokio::{
+	signal,
+	sync::oneshot::{self, Sender},
+};
+use tracing::{debug, info, warn};
+
+extern crate conduit_admin as admin;
+extern crate conduit_core as conduit;
+extern crate conduit_database as database;
+extern crate conduit_service as service;
+
+use std::sync::atomic::Ordering;
+
+use conduit::{debug_info, trace, Error, Result, Server};
+use database::KeyValueDatabase;
+use service::{services, Services};
+
+use crate::{layers, serve};
+
+/// Main loop base
+#[tracing::instrument(skip_all)]
+pub(crate) async fn run(server: Arc<Server>) -> Result<(), Error> {
+	let config = &server.config;
+	let app = layers::build(&server)?;
+	let addrs = config.get_bind_addrs();
+
+	// Install the admin room callback here for now
+	_ = services().admin.handle.lock().await.insert(admin::handle);
+
+	// Setup shutdown/signal handling
+	let handle = ServerHandle::new();
+	_ = server
+		.shutdown
+		.lock()
+		.expect("locked")
+		.insert(handle.clone());
+
+	server.interrupt.store(false, Ordering::Release);
+	let (tx, rx) = oneshot::channel::<()>();
+	let sigs = server.runtime().spawn(sighandle(server.clone(), tx));
+
+	// Prepare to serve http clients
+	let res;
+	// Serve clients
+	if cfg!(unix) && config.unix_socket_path.is_some() {
+		res = serve::unix_socket(&server, app, rx).await;
+	} else if config.tls.is_some() {
+		res = serve::tls(&server, app, handle.clone(), addrs).await;
+	} else {
+		res = serve::plain(&server, app, handle.clone(), addrs).await;
+	}
+
+	// Join the signal handler before we leave.
+	sigs.abort();
+	_ = sigs.await;
+
+	// Reset the axum handle instance; this should be reusable and might be
+	// reload-survivable but better to be safe than sorry.
+	_ = server.shutdown.lock().expect("locked").take();
+
+	// Remove the admin room callback
+	_ = services().admin.handle.lock().await.take();
+
+	debug_info!("Finished");
+	Ok(res?)
+}
+
+/// Async initializations
+#[tracing::instrument(skip_all)]
+pub(crate) async fn start(server: Arc<Server>) -> Result<(), Error> {
+	debug!("Starting...");
+	let d = Arc::new(KeyValueDatabase::load_or_create(&server).await?);
+	let s = Box::new(Services::build(server, d.clone()).await?);
+	_ = service::SERVICES
+		.write()
+		.expect("write locked")
+		.insert(Box::leak(s));
+	services().start().await?;
+
+	#[cfg(feature = "systemd")]
+	#[allow(clippy::let_underscore_untyped)] // error[E0658]: attributes on expressions are experimental
+	let _ = sd_notify::notify(true, &[sd_notify::NotifyState::Ready]);
+
+	debug!("Started");
+	Ok(())
+}
+
+/// Async destructions
+#[tracing::instrument(skip_all)]
+pub(crate) async fn stop(_server: Arc<Server>) -> Result<(), Error> {
+	debug!("Shutting down...");
+
+	// Wait for all completions before dropping or we'll lose them to the module
+	// unload and explode.
+	services().shutdown().await;
+
+	// Deactivate services(). Any further use will panic the caller.
+	let s = service::SERVICES
+		.write()
+		.expect("write locked")
+		.take()
+		.unwrap();
+
+	let s = std::ptr::from_ref(s) as *mut Services;
+	//SAFETY: Services was instantiated in start() and leaked into the SERVICES
+	// global perusing as 'static for the duration of run_server(). Now we reclaim
+	// it to drop it before unloading the module. If this is not done there will be
+	// multiple instances after module reload.
+	let s = unsafe { Box::from_raw(s) };
+	debug!("Cleaning up...");
+	// Drop it so we encounter any trouble before the infolog message
+	drop(s);
+
+	#[cfg(feature = "systemd")]
+	#[allow(clippy::let_underscore_untyped)] // error[E0658]: attributes on expressions are experimental
+	let _ = sd_notify::notify(true, &[sd_notify::NotifyState::Stopping]);
+
+	info!("Shutdown complete.");
+	Ok(())
+}
+
+#[tracing::instrument(skip_all)]
+async fn sighandle(server: Arc<Server>, tx: Sender<()>) -> Result<(), Error> {
+	let ctrl_c = async {
+		signal::ctrl_c()
+			.await
+			.expect("failed to install Ctrl+C handler");
+
+		let reload = cfg!(unix) && cfg!(debug_assertions);
+		server.reload.store(reload, Ordering::Release);
+	};
+
+	#[cfg(unix)]
+	let ctrl_bs = async {
+		signal::unix::signal(signal::unix::SignalKind::quit())
+			.expect("failed to install Ctrl+\\ handler")
+			.recv()
+			.await;
+	};
+
+	#[cfg(unix)]
+	let terminate = async {
+		signal::unix::signal(signal::unix::SignalKind::terminate())
+			.expect("failed to install SIGTERM handler")
+			.recv()
+			.await;
+	};
+
+	debug!("Installed signal handlers");
+	let sig: &str;
+	#[cfg(unix)]
+	tokio::select! {
+		() = ctrl_c => { sig = "Ctrl+C"; },
+		() = ctrl_bs => { sig = "Ctrl+\\"; },
+		() = terminate => { sig = "SIGTERM"; },
+	}
+
+	#[cfg(not(unix))]
+	tokio::select! {
+		_ = ctrl_c => { sig = "Ctrl+C"; },
+	}
+
+	warn!("Received {}", sig);
+	server.interrupt.store(true, Ordering::Release);
+	services().globals.rotate.fire();
+	tx.send(())
+		.expect("failed sending shutdown transaction to oneshot channel");
+
+	if let Some(handle) = server.shutdown.lock().expect("locked").as_ref() {
+		let pending = server.requests_spawn_active.load(Ordering::Relaxed);
+		if pending > 0 {
+			let timeout = Duration::from_secs(36);
+			trace!(pending, ?timeout, "Notifying for graceful shutdown");
+			handle.graceful_shutdown(Some(timeout));
+		} else {
+			debug!(pending, "Notifying for immediate shutdown");
+			handle.shutdown();
+		}
+	}
+
+	Ok(())
+}