chore: add signal entrypoint (#162)

ltitanb · web-flow · commit 546b32eec531 · 2024-10-18T11:18:37.000+01:00
* add signal entrypoint

* fixes
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bin/Cargo.toml b/bin/Cargo.toml
@@ -18,6 +18,9 @@ tokio.workspace = true
 tree_hash.workspace = true
 tree_hash_derive.workspace = true
 
+# telemetry
+tracing.workspace = true
+
 # misc
 clap.workspace = true
 eyre.workspace = true
diff --git a/bin/pbs.rs b/bin/pbs.rs
@@ -1,6 +1,10 @@
-use cb_common::{config::load_pbs_config, utils::initialize_pbs_tracing_log};
+use cb_common::{
+    config::load_pbs_config,
+    utils::{initialize_pbs_tracing_log, wait_for_signal},
+};
 use cb_pbs::{DefaultBuilderApi, PbsService, PbsState};
 use eyre::Result;
+use tracing::{error, info};
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -13,7 +17,21 @@ async fn main() -> Result<()> {
 
     let pbs_config = load_pbs_config()?;
     let _guard = initialize_pbs_tracing_log();
+
     let state = PbsState::new(pbs_config);
     PbsService::init_metrics()?;
-    PbsService::run::<_, DefaultBuilderApi>(state).await
+    let server = PbsService::run::<_, DefaultBuilderApi>(state);
+
+    tokio::select! {
+        maybe_err = server => {
+            if let Err(err) = maybe_err {
+                error!(%err, "PBS service unexpectedly stopped");
+            }
+        },
+        _ = wait_for_signal() => {
+            info!("shutting down");
+        }
+    }
+
+    Ok(())
 }
diff --git a/bin/signer.rs b/bin/signer.rs
@@ -1,9 +1,10 @@
 use cb_common::{
     config::{StartSignerConfig, SIGNER_MODULE_NAME},
-    utils::initialize_tracing_log,
+    utils::{initialize_tracing_log, wait_for_signal},
 };
 use cb_signer::service::SigningService;
 use eyre::Result;
+use tracing::{error, info};
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -16,5 +17,18 @@ async fn main() -> Result<()> {
 
     let config = StartSignerConfig::load_from_env()?;
     let _guard = initialize_tracing_log(SIGNER_MODULE_NAME);
-    SigningService::run(config).await
+    let server = SigningService::run(config);
+
+    tokio::select! {
+        maybe_err = server => {
+            if let Err(err) = maybe_err {
+                error!(%err, "signing server unexpectedly stopped");
+            }
+        },
+        _ = wait_for_signal() => {
+            info!("shutting down");
+        }
+    }
+
+    Ok(())
 }
diff --git a/bin/src/lib.rs b/bin/src/lib.rs
@@ -21,7 +21,7 @@ pub mod prelude {
         get_header, get_status, register_validator, submit_block, BuilderApi, BuilderApiState,
         DefaultBuilderApi, PbsService, PbsState,
     };
-    // The TreeHash derive macro requires tree_hash:: as import
+    // The TreeHash derive macro requires tree_hash as import
     pub mod tree_hash {
         pub use tree_hash::*;
     }
diff --git a/crates/common/src/utils.rs b/crates/common/src/utils.rs
@@ -246,3 +246,24 @@ pub fn get_user_agent_with_version(req_headers: &HeaderMap) -> eyre::Result<Head
     let ua = get_user_agent(req_headers);
     Ok(HeaderValue::from_str(&format!("commit-boost/{HEADER_VERSION_VALUE} {}", ua))?)
 }
+
+#[cfg(unix)]
+pub async fn wait_for_signal() -> eyre::Result<()> {
+    use tokio::signal::unix::{signal, SignalKind};
+
+    let mut sigint = signal(SignalKind::interrupt())?;
+    let mut sigterm = signal(SignalKind::terminate())?;
+
+    tokio::select! {
+        _ = sigint.recv() => {}
+        _ = sigterm.recv() => {}
+    }
+
+    Ok(())
+}
+
+#[cfg(windows)]
+pub async fn wait_for_signal() -> eyre::Result<()> {
+    tokio::signal::ctrl_c().await?;
+    Ok(())
+}
diff --git a/crates/pbs/src/service.rs b/crates/pbs/src/service.rs
@@ -21,10 +21,10 @@ impl PbsService {
         let address = SocketAddr::from(([0, 0, 0, 0], state.config.pbs_config.port));
         let events_subs =
             state.config.event_publisher.as_ref().map(|e| e.n_subscribers()).unwrap_or_default();
-        info!(version = COMMIT_BOOST_VERSION, ?address, events_subs, chain =? state.config.chain, "Starting PBS service");
+        info!(version = COMMIT_BOOST_VERSION, ?address, events_subs, chain =? state.config.chain, "starting PBS service");
 
         let app = create_app_router::<S, A>(state);
-        let listener = TcpListener::bind(address).await.expect("failed tcp binding");
+        let listener = TcpListener::bind(address).await?;
 
         axum::serve(listener, app).await.wrap_err("PBS server exited")
     }
diff --git a/crates/signer/src/service.rs b/crates/signer/src/service.rs
@@ -22,7 +22,7 @@ use cb_common::{
     constants::COMMIT_BOOST_VERSION,
     types::{Jwt, ModuleId},
 };
-use eyre::{Result, WrapErr};
+use eyre::{Context, Result};
 use headers::{authorization::Bearer, Authorization};
 use tokio::{net::TcpListener, sync::RwLock};
 use tracing::{debug, error, info, warn};
@@ -71,12 +71,9 @@ impl SigningService {
             .route_layer(middleware::from_fn_with_state(state.clone(), jwt_auth));
 
         let address = SocketAddr::from(([0, 0, 0, 0], config.server_port));
-        let listener = TcpListener::bind(address).await.wrap_err("failed tcp binding")?;
+        let listener = TcpListener::bind(address).await?;
 
-        if let Err(err) = axum::serve(listener, app).await {
-            error!(%err, "Signing server exited")
-        }
-        Ok(())
+        axum::serve(listener, app).await.wrap_err("signer server exited")
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ pub mod prelude {`
`21`	`21`	`get_header, get_status, register_validator, submit_block, BuilderApi, BuilderApiState,`
`22`	`22`	`DefaultBuilderApi, PbsService, PbsState,`
`23`	`23`	`};`
`24`		`- // The TreeHash derive macro requires tree_hash:: as import`
	`24`	`+ // The TreeHash derive macro requires tree_hash as import`
`25`	`25`	`pub mod tree_hash {`
`26`	`26`	`pub use tree_hash::*;`
`27`	`27`	`}`