Disconnect worker on shutdown

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,27 +129,36 @@ 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(
     router: &mut Router,
     mut pool_addr: Option<std::net::SocketAddr>,
     epsilon: Duration,
     signature: String,
+    shutdown_signal: watch::Receiver<bool>,
 ) {
     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<bool>,
 ) -> 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 {

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())
+            }
+        }
+    }
 }

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<Mutex<Vec<WorkerActivity>>> = 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(())
+    }
 }

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<bool> {
+    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
+        }
+    }
+}