aquatic_ws: use PanicSentinel, simplify lib.rs

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 {
+        let executor = builder
-            workers::socket::run_socket_worker(
+            .spawn(move || async move {
-                config,
+                workers::socket::run_socket_worker(
-                state,
+                    sentinel,
-                tls_config,
+                    config,
-                request_mesh_builder,
+                    state,
-                response_mesh_builder,
+                    tls_config,
-                priv_dropper,
+                    request_mesh_builder,
-            )
+                    response_mesh_builder,
-            .await
+                    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 {
+        let executor = builder
-            workers::request::run_request_worker(
+            .spawn(move || async move {
-                config,
+                workers::request::run_request_worker(
-                state,
+                    sentinel,
-                request_mesh_builder,
+                    config,
-                response_mesh_builder,
+                    state,
-            )
+                    request_mesh_builder,
-            .await
+                    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 {
+    for signal in &mut signals {
-        executor
+        match signal {
-            .expect("failed to spawn local executor")
+            SIGUSR1 => {
-            .join()
+                let _ = update_access_list(&config.access_list, &state.access_list);
-            .unwrap();
+            }
+            SIGTERM => {
+                if sentinel_watcher.panic_was_triggered() {
+                    return Err(anyhow::anyhow!("worker thread panicked"));
+                } else {
+                    return Ok(());
+                }
+            }
+            _ => unreachable!(),
+        }
     }
 
     Ok(())

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>,

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<RustlsConfig>,