Add PoolSv2::shutdown

A simple shutdown function that exits the State loop and shutsdown the
Pool server.

Author: jbesraa

roles/pool/src/lib/mining_pool/mod.rs

@@ -270,47 +270,58 @@ impl Pool {
     async fn accept_incoming_connection(
         self_: Arc<Mutex<Pool>>,
         config: PoolConfig,
+        mut recv_stop_signal: tokio::sync::watch::Receiver<()>,
     ) -> PoolResult<()> {
         let status_tx = self_.safe_lock(|s| s.status_tx.clone())?;
         let listener = TcpListener::bind(&config.listen_address()).await?;
-        info!(
-            "Listening for encrypted connection on: {}",
-            config.listen_address()
-        );
-        while let Ok((stream, _)) = listener.accept().await {
-            let address = stream.peer_addr().unwrap();
-            debug!(
-                "New connection from {:?}",
-                stream.peer_addr().map_err(PoolError::Io)
-            );
-
-            let responder = Responder::from_authority_kp(
-                &config.authority_public_key().into_bytes(),
-                &config.authority_secret_key().into_bytes(),
-                std::time::Duration::from_secs(config.cert_validity_sec()),
-            );
-            match responder {
-                Ok(resp) => {
-                    if let Ok((receiver, sender, _, _)) =
-                        Connection::new(stream, HandshakeRole::Responder(resp)).await
-                    {
-                        handle_result!(
-                            status_tx,
-                            Self::accept_incoming_connection_(
-                                self_.clone(),
-                                receiver,
-                                sender,
-                                address
-                            )
-                            .await
-                        );
+        info!("Pool is running on: {}", config.listen_address());
+        // Run the listener in the background
+        task::spawn(async move {
+            loop {
+                tokio::select! {
+                    _ = recv_stop_signal.changed() => {
+                        info!("Pool is stopping the server after stop shutdown signal received");
+                        break;
+                    },
+
+                    result = listener.accept() => {
+                        match result {
+                            Ok((stream, _)) => {
+                                let address = stream.peer_addr().unwrap();
+                                info!("New connection from {:?}", stream.peer_addr().map_err(PoolError::Io));
+                                let responder = Responder::from_authority_kp(
+                                    &config.authority_public_key().into_bytes(),
+                                    &config.authority_secret_key().into_bytes(),
+                                    std::time::Duration::from_secs(config.cert_validity_sec()),
+                                );
+
+                                match responder {
+                                    Ok(resp) => {
+                                        if let Ok((receiver, sender, _, _)) = Connection::new(stream, HandshakeRole::Responder(resp)).await {
+                                            handle_result!(
+                                                status_tx,
+                                                Self::accept_incoming_connection_(
+                                                    self_.clone(),
+                                                    receiver,
+                                                    sender,
+                                                    address
+                                                ).await
+                                            );
+                                        }
+                                    }
+                                    Err(_) => {
+                                        return;
+                                    }
+                                }
+                            }
+                            Err(e) => {
+                                error!("Error accepting connection: {:?}", e);
+                            }
+                        }
                     }
                 }
-                Err(_e) => {
-                    todo!()
-                }
             }
-        }
+        });
         Ok(())
     }
 
@@ -444,15 +455,17 @@ impl Pool {
         Ok(())
     }
 
-    pub fn start(
+    #[allow(clippy::too_many_arguments)]
+    pub async fn start(
         config: PoolConfig,
         new_template_rx: Receiver<NewTemplate<'static>>,
         new_prev_hash_rx: Receiver<SetNewPrevHash<'static>>,
         solution_sender: Sender<SubmitSolution<'static>>,
         sender_message_received_signal: Sender<()>,
         status_tx: status::Sender,
         shares_per_minute: f32,
-    ) -> Arc<Mutex<Self>> {
+        recv_stop_signal: tokio::sync::watch::Receiver<()>,
+    ) -> Result<Arc<Mutex<Self>>, PoolError> {
         let extranonce_len = 32;
         let range_0 = std::ops::Range { start: 0, end: 0 };
         let range_1 = std::ops::Range { start: 0, end: 16 };
@@ -488,24 +501,20 @@ impl Pool {
         let cloned2 = pool.clone();
         let cloned3 = pool.clone();
 
-        info!("Starting up pool listener");
+        info!("Starting up Pool server");
         let status_tx_clone = status_tx.clone();
-        task::spawn(async move {
-            if let Err(e) = Self::accept_incoming_connection(cloned, config).await {
-                error!("{}", e);
-            }
-            if status_tx_clone
+        if let Err(e) = Self::accept_incoming_connection(cloned, config, recv_stop_signal).await {
+            error!("Pool stopped accepting connections due to: {}", &e);
+            let _ = status_tx_clone
                 .send(status::Status {
                     state: status::State::DownstreamShutdown(PoolError::ComponentShutdown(
-                        "Downstream no longer accepting incoming connections".to_string(),
+                        "Pool stopped accepting connections".to_string(),
                     )),
                 })
-                .await
-                .is_err()
-            {
-                error!("Downstream shutdown and Status Channel dropped");
-            }
-        });
+                .await;
+
+            return Err(e);
+        }
 
         let cloned = sender_message_received_signal.clone();
         let status_tx_clone = status_tx.clone();
@@ -547,7 +556,7 @@ impl Pool {
                 error!("Downstream shutdown and Status Channel dropped");
             }
         });
-        cloned3
+        Ok(cloned3)
     }
 
     /// This removes the downstream from the list of downstreams

roles/pool/src/lib/mod.rs

@@ -7,23 +7,38 @@ use async_channel::{bounded, unbounded};
 use config::PoolConfig;
 use error::PoolError;
 use mining_pool::{get_coinbase_output, Pool};
+use std::sync::{Arc, Mutex};
 use template_receiver::TemplateRx;
 use tokio::select;
 use tracing::{error, info, warn};
 
 #[derive(Debug, Clone)]
 pub struct PoolSv2 {
     config: PoolConfig,
+    status_tx: Arc<Mutex<Option<async_channel::Sender<status::Status>>>>,
 }
 
 impl PoolSv2 {
     pub fn new(config: PoolConfig) -> PoolSv2 {
-        PoolSv2 { config }
+        PoolSv2 {
+            config,
+            status_tx: Arc::new(Mutex::new(None)),
+        }
     }
 
     pub async fn start(&self) -> Result<(), PoolError> {
         let config = self.config.clone();
         let (status_tx, status_rx) = unbounded();
+
+        if let Ok(mut s_tx) = self.status_tx.lock() {
+            *s_tx = Some(status_tx.clone());
+        } else {
+            error!("Failed to access Pool status lock");
+            return Err(PoolError::Custom(
+                "Failed to access Pool status lock".to_string(),
+            ));
+        }
+        let (send_stop_signal, recv_stop_signal) = tokio::sync::watch::channel(());
         let (s_new_t, r_new_t) = bounded(10);
         let (s_prev_hash, r_prev_hash) = bounded(10);
         let (s_solution, r_solution) = bounded(10);
@@ -60,8 +75,9 @@ impl PoolSv2 {
             s_message_recv_signal,
             status::Sender::DownstreamListener(status_tx),
             config.shares_per_minute(),
-        );
-
+            recv_stop_signal,
+        )
+        .await?;
         // Start the error handling loop
         // See `./status.rs` and `utils/error_handling` for information on how this operates
         tokio::spawn(async move {
@@ -84,16 +100,23 @@ impl PoolSv2 {
                 let task_status: status::Status = task_status.unwrap();
 
                 match task_status.state {
+                    status::State::Shutdown => {
+                        info!("Received shutdown signal");
+                        let _ = send_stop_signal.send(());
+                        break;
+                    }
                     // Should only be sent by the downstream listener
                     status::State::DownstreamShutdown(err) => {
                         error!(
                             "SHUTDOWN from Downstream: {}\nTry to restart the downstream listener",
                             err
                         );
+                        let _ = send_stop_signal.send(());
                         break;
                     }
                     status::State::TemplateProviderShutdown(err) => {
                         error!("SHUTDOWN from Upstream: {}\nTry to reconnecting or connecting to a new upstream", err);
+                        let _ = send_stop_signal.send(());
                         break;
                     }
                     status::State::Healthy(msg) => {
@@ -105,6 +128,7 @@ impl PoolSv2 {
                             .safe_lock(|p| p.remove_downstream(downstream_id))
                             .is_err()
                         {
+                            let _ = send_stop_signal.send(());
                             break;
                         }
                     }
@@ -113,6 +137,31 @@ impl PoolSv2 {
         });
         Ok(())
     }
+
+    pub fn shutdown(&self) {
+        info!("Attempting to shutdown pool");
+        if let Ok(status_tx) = &self.status_tx.lock() {
+            if let Some(status_tx) = status_tx.as_ref().cloned() {
+                info!("Pool is running, sending shutdown signal");
+                tokio::spawn(async move {
+                    if let Err(e) = status_tx
+                        .send(status::Status {
+                            state: status::State::Shutdown,
+                        })
+                        .await
+                    {
+                        error!("Failed to send shutdown signal to status loop: {:?}", e);
+                    } else {
+                        info!("Sent shutdown signal to Pool");
+                    }
+                });
+            } else {
+                info!("Pool is not running.");
+            }
+        } else {
+            error!("Failed to access Pool status lock");
+        }
+    }
 }
 
 #[cfg(test)]
@@ -148,4 +197,33 @@ mod tests {
         let result = pool.start().await;
         assert!(result.is_err());
     }
+
+    #[tokio::test]
+    async fn shutdown_pool() {
+        let config_path = "config-examples/pool-config-hosted-tp-example.toml";
+        let config: PoolConfig = match Config::builder()
+            .add_source(File::new(config_path, FileFormat::Toml))
+            .build()
+        {
+            Ok(settings) => match settings.try_deserialize::<PoolConfig>() {
+                Ok(c) => c,
+                Err(e) => {
+                    error!("Failed to deserialize config: {}", e);
+                    return;
+                }
+            },
+            Err(e) => {
+                error!("Failed to build config: {}", e);
+                return;
+            }
+        };
+        let pool_0 = PoolSv2::new(config.clone());
+        let pool_1 = PoolSv2::new(config);
+        assert!(pool_0.start().await.is_ok());
+        tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+        assert!(pool_1.start().await.is_err());
+        pool_0.shutdown();
+        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        assert!(pool_1.start().await.is_ok());
+    }
 }

roles/pool/src/lib/status.rs

@@ -48,6 +48,7 @@ pub enum State {
     TemplateProviderShutdown(PoolError),
     DownstreamInstanceDropped(u32),
     Healthy(String),
+    Shutdown,
 }
 
 /// message to be sent to the status loop on the main thread

test/integration-tests/lib/mod.rs

@@ -89,10 +89,7 @@ pub async fn start_pool(template_provider_address: Option<SocketAddr>) -> (PoolS
         SHARES_PER_MINUTE,
     );
     let pool = PoolSv2::new(config);
-    let pool_clone = pool.clone();
-    tokio::task::spawn(async move {
-        assert!(pool_clone.start().await.is_ok());
-    });
+    assert!(pool.start().await.is_ok());
     // Wait a bit to let the pool exchange initial messages with the TP
     sleep(1).await;
     (pool, listening_address)