diff --git a/aquatic_ws/src/lib.rs b/aquatic_ws/src/lib.rs index 6671b00..b6ffc22 100644 --- a/aquatic_ws/src/lib.rs +++ b/aquatic_ws/src/lib.rs @@ -7,8 +7,12 @@ use std::sync::Arc; use aquatic_common::cpu_pinning::glommio::{get_worker_placement, set_affinity_for_util_worker}; use aquatic_common::cpu_pinning::WorkerIndex; use aquatic_common::rustls_config::create_rustls_config; +use aquatic_common::PanicSentinelWatcher; use glommio::{channels::channel_mesh::MeshBuilder, prelude::*}; -use signal_hook::{consts::SIGUSR1, iterator::Signals}; +use signal_hook::{ + consts::{SIGTERM, SIGUSR1}, + iterator::Signals, +}; use aquatic_common::access_list::update_access_list; use aquatic_common::privileges::PrivilegeDropper; @@ -22,45 +26,18 @@ pub const APP_VERSION: &str = env!("CARGO_PKG_VERSION"); pub const SHARED_IN_CHANNEL_SIZE: usize = 1024; pub fn run(config: Config) -> ::anyhow::Result<()> { + let mut signals = Signals::new([SIGUSR1, SIGTERM])?; + let state = State::default(); update_access_list(&config.access_list, &state.access_list)?; - let mut signals = Signals::new(::std::iter::once(SIGUSR1))?; - - { - let config = config.clone(); - let state = state.clone(); - - ::std::thread::spawn(move || run_workers(config, state)); - } - - if config.cpu_pinning.active { - set_affinity_for_util_worker( - &config.cpu_pinning, - config.socket_workers, - config.request_workers, - )?; - } - - for signal in &mut signals { - match signal { - SIGUSR1 => { - let _ = update_access_list(&config.access_list, &state.access_list); - } - _ => unreachable!(), - } - } - - Ok(()) -} - -fn run_workers(config: Config, state: State) -> anyhow::Result<()> { let num_peers = config.socket_workers + config.request_workers; let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE); let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE * 16); + let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel(); let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers); let tls_config = Arc::new(create_rustls_config( @@ -71,6 +48,7 @@ fn run_workers(config: Config, state: State) -> anyhow::Result<()> { let mut executors = Vec::new(); for i in 0..(config.socket_workers) { + let sentinel = sentinel.clone(); let config = config.clone(); let state = state.clone(); let tls_config = tls_config.clone(); @@ -86,22 +64,26 @@ fn run_workers(config: Config, state: State) -> anyhow::Result<()> { )?; let builder = LocalExecutorBuilder::new(placement).name("socket"); - let executor = builder.spawn(move || async move { - workers::socket::run_socket_worker( - config, - state, - tls_config, - request_mesh_builder, - response_mesh_builder, - priv_dropper, - ) - .await - }); + let executor = builder + .spawn(move || async move { + workers::socket::run_socket_worker( + sentinel, + config, + state, + tls_config, + request_mesh_builder, + response_mesh_builder, + priv_dropper, + ) + .await + }) + .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?; executors.push(executor); } for i in 0..(config.request_workers) { + let sentinel = sentinel.clone(); let config = config.clone(); let state = state.clone(); let request_mesh_builder = request_mesh_builder.clone(); @@ -115,15 +97,18 @@ fn run_workers(config: Config, state: State) -> anyhow::Result<()> { )?; let builder = LocalExecutorBuilder::new(placement).name("request"); - let executor = builder.spawn(move || async move { - workers::request::run_request_worker( - config, - state, - request_mesh_builder, - response_mesh_builder, - ) - .await - }); + let executor = builder + .spawn(move || async move { + workers::request::run_request_worker( + sentinel, + config, + state, + request_mesh_builder, + response_mesh_builder, + ) + .await + }) + .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?; executors.push(executor); } @@ -136,11 +121,20 @@ fn run_workers(config: Config, state: State) -> anyhow::Result<()> { )?; } - for executor in executors { - executor - .expect("failed to spawn local executor") - .join() - .unwrap(); + for signal in &mut signals { + match signal { + SIGUSR1 => { + let _ = update_access_list(&config.access_list, &state.access_list); + } + SIGTERM => { + if sentinel_watcher.panic_was_triggered() { + return Err(anyhow::anyhow!("worker thread panicked")); + } else { + return Ok(()); + } + } + _ => unreachable!(), + } } Ok(()) diff --git a/aquatic_ws/src/workers/request.rs b/aquatic_ws/src/workers/request.rs index 693ac7d..92f07ad 100644 --- a/aquatic_ws/src/workers/request.rs +++ b/aquatic_ws/src/workers/request.rs @@ -12,7 +12,7 @@ use glommio::timer::TimerActionRepeat; use hashbrown::HashMap; use rand::{rngs::SmallRng, SeedableRng}; -use aquatic_common::{extract_response_peers, AmortizedIndexMap}; +use aquatic_common::{extract_response_peers, AmortizedIndexMap, PanicSentinel}; use aquatic_ws_protocol::*; use crate::common::*; @@ -128,6 +128,7 @@ impl TorrentMaps { } pub async fn run_request_worker( + _sentinel: PanicSentinel, config: Config, state: State, in_message_mesh_builder: MeshBuilder<(ConnectionMeta, InMessage), Partial>, diff --git a/aquatic_ws/src/workers/socket.rs b/aquatic_ws/src/workers/socket.rs index 2da8d05..cfda49a 100644 --- a/aquatic_ws/src/workers/socket.rs +++ b/aquatic_ws/src/workers/socket.rs @@ -10,7 +10,7 @@ use anyhow::Context; use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; use aquatic_common::privileges::PrivilegeDropper; use aquatic_common::rustls_config::RustlsConfig; -use aquatic_common::CanonicalSocketAddr; +use aquatic_common::{CanonicalSocketAddr, PanicSentinel}; use aquatic_ws_protocol::*; use async_tungstenite::WebSocketStream; use futures::stream::{SplitSink, SplitStream}; @@ -49,6 +49,7 @@ struct ConnectionReference { } pub async fn run_socket_worker( + _sentinel: PanicSentinel, config: Config, state: State, tls_config: Arc,