diff --git a/src/main.rs b/src/main.rs index dd7defc2..ca4d5bad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,16 +6,19 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer}; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; -use crate::{monitor::logs::SendLogLayer, shared::utils::AbortOnDrop}; +use crate::{ + monitor::logs::SendLogLayer, shared::utils::AbortOnDrop, shutdown::update_node_status, +}; use config::Configuration; use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; use self_update::{backends, cargo_crate_version, update::UpdateStatus, TempDir}; use std::{net::SocketAddr, time::Duration}; -use tokio::sync::mpsc::channel; +use tokio::sync::{mpsc::channel, watch}; use tracing::{debug, error, info, warn}; mod api; +mod shutdown; mod config; mod ingress; @@ -126,18 +129,23 @@ async fn main() { _ => unreachable!(), }); + let shutdown_signal = shutdown::handle_shutdown(); let mut router = router::Router::new(pool_addresses, auth_pub_k, None, None); let epsilon = Duration::from_millis(30_000); let best_upstream = router.select_pool_connect().await; + // Register node on startup + update_node_status("register").await; + initialize_proxy( &mut router, best_upstream, epsilon, Configuration::signature(), + shutdown_signal, ) .await; info!("exiting"); - tokio::time::sleep(tokio::time::Duration::from_secs(60)).await; + tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; } async fn initialize_proxy( @@ -145,8 +153,12 @@ async fn initialize_proxy( mut pool_addr: Option, epsilon: Duration, signature: String, + shutdown_signal: watch::Receiver, ) { loop { + if *shutdown_signal.borrow() { + break; + } let stats_sender = api::stats::StatsSender::new(); let (send_to_pool, recv_from_pool, pool_connection_abortable) = match router.connect_pool(pool_addr).await { @@ -258,7 +270,7 @@ async fn initialize_proxy( } let server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); abort_handles.push((server_handle.into(), "api_server".to_string())); - match monitor(router, abort_handles, epsilon).await { + match monitor(router, abort_handles, epsilon, shutdown_signal.clone()).await { Reconnect::NewUpstream(new_pool_addr) => { ProxyState::update_proxy_state_up(); pool_addr = Some(new_pool_addr); @@ -277,9 +289,15 @@ async fn monitor( router: &mut Router, abort_handles: Vec<(AbortOnDrop, std::string::String)>, epsilon: Duration, + shutdown_signal: watch::Receiver, ) -> Reconnect { let mut should_check_upstreams_latency = 0; loop { + if *shutdown_signal.borrow() { + info!("Dropping all abort handles"); + drop(abort_handles); + return Reconnect::NoUpstream; + } if Configuration::monitor() { // Check if a better upstream exist every 100 seconds if should_check_upstreams_latency == 10 * 100 { diff --git a/src/monitor/mod.rs b/src/monitor/mod.rs index 05538b97..7271ae53 100644 --- a/src/monitor/mod.rs +++ b/src/monitor/mod.rs @@ -48,6 +48,26 @@ fn worker_activity_server_endpoint() -> String { } } +pub fn node_register_endpoint() -> String { + match Configuration::environment().as_str() { + "staging" => format!("{}/api/nodes/add", STAGING_URL), + "testnet3" => format!("{}/api/nodes/add", TESTNET3_URL), + "local" => format!("{}/api/nodes/add", LOCAL_URL), + "production" => format!("{}/api/nodes/add", PRODUCTION_URL), + _ => unreachable!(), + } +} + +pub fn node_unregister_endpoint() -> String { + match Configuration::environment().as_str() { + "staging" => format!("{}/api/nodes/remove", STAGING_URL), + "testnet3" => format!("{}/api/nodes/remove", TESTNET3_URL), + "local" => format!("{}/api/nodes/remove", LOCAL_URL), + "production" => format!("{}/api/nodes/remove", PRODUCTION_URL), + _ => unreachable!(), + } +} + impl MonitorAPI { pub fn new(url: String) -> Self { let client = reqwest::Client::new(); @@ -118,4 +138,42 @@ impl MonitorAPI { } } } + + pub async fn register_bitcoin_node(&self) -> Result<(), Error> { + let token = crate::config::Configuration::token().expect("Token is not set"); + debug!("Registering bitcoin node"); + let response = self + .client + .post(self.url.clone()) + .json(&json!({ "token": token })) + .send() + .await?; + + match response.error_for_status() { + Ok(_) => Ok(()), + Err(err) => { + error!("Failed to register bitcoin node: {}", err); + Err(err.into()) + } + } + } + + pub async fn unregister_bitcoin_node(&self) -> Result<(), Error> { + let token = crate::config::Configuration::token().expect("Token is not set"); + debug!("Unregistering bitcoin node"); + let response = self + .client + .post(self.url.clone()) + .json(&json!({ "token": token })) + .send() + .await?; + + match response.error_for_status() { + Ok(_) => Ok(()), + Err(err) => { + error!("Failed to unregister bitcoin node: {}", err); + Err(err.into()) + } + } + } } diff --git a/src/monitor/worker_activity.rs b/src/monitor/worker_activity.rs index d2c6e71c..959dac98 100644 --- a/src/monitor/worker_activity.rs +++ b/src/monitor/worker_activity.rs @@ -1,12 +1,20 @@ +use std::sync::Arc; + use crate::monitor::{worker_activity_server_endpoint, MonitorAPI}; +use lazy_static::lazy_static; +use roles_logic_sv2::utils::Mutex; +use tracing::{debug, error}; +lazy_static! { + static ref ACTIVE_WORKERS: Arc>> = Arc::new(Mutex::new(Vec::new())); // To keep track of all active workers, so we can disconnect them on shutdown -#[derive(serde::Serialize, Debug)] +} +#[derive(serde::Serialize, Debug, PartialEq, Eq, Clone)] pub enum WorkerActivityType { Connected, Disconnected, } -#[derive(serde::Serialize, Debug)] +#[derive(serde::Serialize, Debug, PartialEq, Eq, Clone)] pub struct WorkerActivity { user_agent: String, worker_name: String, @@ -15,6 +23,26 @@ pub struct WorkerActivity { impl WorkerActivity { pub fn new(user_agent: String, worker_name: String, activity: WorkerActivityType) -> Self { + if ACTIVE_WORKERS + .safe_lock(|workers| match activity { + WorkerActivityType::Connected => { + workers.push(WorkerActivity { + user_agent: user_agent.clone(), + worker_name: worker_name.clone(), + activity: activity.clone(), + }); + } + WorkerActivityType::Disconnected => { + workers + .retain(|w| !(w.user_agent == user_agent && w.worker_name == worker_name)); + } + }) + .is_err() + { + error!("Failed to acquire lock on ACTIVE_WORKERS"); + std::process::exit(1) + } + WorkerActivity { user_agent, worker_name, @@ -25,4 +53,28 @@ impl WorkerActivity { pub fn monitor_api(&self) -> MonitorAPI { MonitorAPI::new(worker_activity_server_endpoint()) } + + pub async fn disconnect_all_active_workers() -> Result<(), crate::shared::error::Error> { + let active_workers = ACTIVE_WORKERS.safe_lock(|workers| workers.clone()).expect( + "Failed to acquire lock on ACTIVE_WORKERS when disconnecting all active workers", + ); + + for worker in &active_workers { + let worker_activity = WorkerActivity::new( + worker.user_agent.clone(), + worker.worker_name.clone(), + WorkerActivityType::Disconnected, + ); + + worker_activity + .monitor_api() + .send_worker_activity(worker_activity) + .await?; + debug!( + "Disconnected worker: {} ({})", + worker.worker_name, worker.user_agent + ); + } + Ok(()) + } } diff --git a/src/shutdown.rs b/src/shutdown.rs new file mode 100644 index 00000000..a4ff21ed --- /dev/null +++ b/src/shutdown.rs @@ -0,0 +1,111 @@ +use tokio::net::TcpStream; +use tokio::signal; +use tokio::signal::unix::{signal as unix_signal, SignalKind}; +use tokio::sync::watch; +use tokio::time::timeout; +use tracing::{debug, error, info}; + +use crate::monitor::worker_activity::WorkerActivity; +use crate::monitor::{self, node_register_endpoint, node_unregister_endpoint}; + +/// Spawns a background task that listens for SIGINT (Ctrl+C) and SIGTERM, +/// unregisters the node, and then notifies all listeners via a watch channel. +pub fn handle_shutdown() -> watch::Receiver { + let (tx, rx) = watch::channel(false); + + tokio::spawn(async move { + // Set up SIGTERM stream (Unix only) + let mut sigterm = match unix_signal(SignalKind::terminate()) { + Ok(s) => s, + Err(e) => { + error!("Failed to listen for SIGTERM: {:?}", e); + return; + } + }; + + tokio::select! { + res = signal::ctrl_c() => { + match res { + Ok(()) => info!("Received SIGINT (Ctrl+C), shutting down..."), + Err(e) => { + error!("Failed to listen for Ctrl+C: {:?}", e); + return; + } + } + } + _ = sigterm.recv() => { + info!("Received SIGTERM, shutting down..."); + } + } + + update_node_status("unregister").await; + WorkerActivity::disconnect_all_active_workers() + .await + .unwrap_or_else(|e| { + error!("Failed to disconnect all active workers: {}", e); + }); + + // Notify all listeners; ignore error if there are no receivers left. + let _ = tx.send(true); + }); + + rx +} + +/// Updates the node status by either registering or unregistering it only if TP_ADDRESS is set and reachable. +pub async fn update_node_status(action: &str) { + let tp_address = match crate::config::Configuration::tp_address() { + Some(addr) => addr, + None => { + debug!("TP_ADDRESS not set, skipping node {}.", action); + return; + } + }; + + // Test if TP_ADDRESS is reachable + if !is_reachable(&tp_address).await { + debug!( + "TP_ADDRESS {} is unreachable, skipping node {}.", + tp_address, action + ); + return; + } + match action { + "unregister" => { + let monitor_api = monitor::MonitorAPI::new(node_unregister_endpoint()); + monitor_api.unregister_bitcoin_node().await + } + "register" => { + let monitor_api = monitor::MonitorAPI::new(node_register_endpoint()); + monitor_api.register_bitcoin_node().await + } + _ => { + error!("Unknown action: {}", action); + return; + } + } + .unwrap_or_else(|e| { + error!("Failed to {} bitcoin node: {}", action, e); + }); +} + +async fn is_reachable(address: &str) -> bool { + // Try to establish a TCP connection + debug!("Checking reachability of {}", address); + match timeout( + std::time::Duration::from_secs(5), + TcpStream::connect(address), + ) + .await + { + Ok(Ok(_)) => true, + Ok(Err(e)) => { + debug!("Failed to connect to {}: {}", address, e); + false + } + Err(_) => { + debug!("Connection to {} timed out", address); + false + } + } +}