Merge pull request dmnd-pool#131 from Fi3/FixProxyRestart

Fix proxy restart

Author: jbesraa

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "demand-cli"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 
 [dependencies]

Cargo.toml

@@ -6,23 +6,34 @@ use tokio::task::JoinHandle;
 
 #[derive(Debug)]
 pub struct AbortOnDrop {
-    abort_handle: AbortHandle,
+    abort_handle: Vec<AbortHandle>,
 }
 
 impl AbortOnDrop {
     pub fn new<T: Send + 'static>(handle: JoinHandle<T>) -> Self {
-        let abort_handle = handle.abort_handle();
+        let abort_handle = vec![handle.abort_handle()];
         Self { abort_handle }
     }
 
     pub fn is_finished(&self) -> bool {
-        self.abort_handle.is_finished()
+        for task in &self.abort_handle {
+            if !task.is_finished() {
+                return false;
+            }
+        }
+        true
+    }
+
+    pub fn add_task<T: Send + 'static>(&mut self, handle: JoinHandle<T>) {
+        self.abort_handle.push(handle.abort_handle());
     }
 }
 
 impl core::ops::Drop for AbortOnDrop {
     fn drop(&mut self) {
-        self.abort_handle.abort()
+        for task in &self.abort_handle {
+            task.abort();
+        }
     }
 }
 

src/shared/utils.rs

@@ -16,55 +16,65 @@ pub async fn start_receive_downstream(
     mut recv_from_down: mpsc::Receiver<String>,
     connection_id: u32,
 ) -> Result<(), Error<'static>> {
-    let task_manager_clone = task_manager.clone();
-    let handle = task::spawn(async move {
-        while let Some(incoming) = recv_from_down.recv().await {
-            let incoming: Result<json_rpc::Message, _> = serde_json::from_str(&incoming);
-            if let Ok(incoming) = incoming {
-                // if message is Submit Shares update difficulty management
-                if let sv1_api::Message::StandardRequest(standard_req) = incoming.clone() {
-                    if let Ok(Submit { .. }) = standard_req.try_into() {
-                        if let Err(e) = Downstream::save_share(downstream.clone()) {
-                            error!("{}", e);
-                            break;
+    let handle = {
+        let task_manager = task_manager.clone();
+        task::spawn(async move {
+            while let Some(incoming) = recv_from_down.recv().await {
+                let incoming: Result<json_rpc::Message, _> = serde_json::from_str(&incoming);
+                if let Ok(incoming) = incoming {
+                    // if message is Submit Shares update difficulty management
+                    if let sv1_api::Message::StandardRequest(standard_req) = incoming.clone() {
+                        if let Ok(Submit { .. }) = standard_req.try_into() {
+                            if let Err(e) = Downstream::save_share(downstream.clone()) {
+                                error!("{}", e);
+                                break;
+                            }
                         }
                     }
-                }
 
-                if let Err(error) =
-                    Downstream::handle_incoming_sv1(downstream.clone(), incoming).await
-                {
-                    error!("Failed to handle incoming sv1 msg: {:?}", error);
-                    ProxyState::update_downstream_state(DownstreamType::TranslatorDownstream);
-                };
-            } else {
-                // Message received could not be converted to rpc message
-                error!(
-                    "{}",
-                    Error::V1Protocol(Box::new(sv1_api::error::Error::InvalidJsonRpcMessageKind))
-                );
-                return;
+                    if let Err(error) =
+                        Downstream::handle_incoming_sv1(downstream.clone(), incoming).await
+                    {
+                        error!("Failed to handle incoming sv1 msg: {:?}", error);
+                        ProxyState::update_downstream_state(DownstreamType::TranslatorDownstream);
+                    };
+                } else {
+                    // Message received could not be converted to rpc message
+                    error!(
+                        "{}",
+                        Error::V1Protocol(Box::new(
+                            sv1_api::error::Error::InvalidJsonRpcMessageKind
+                        ))
+                    );
+                    return;
+                }
             }
-        }
-        if let Ok(stats_sender) = downstream.safe_lock(|d| d.stats_sender.clone()) {
-            stats_sender.remove_stats(connection_id);
-        }
-        // No message to receive
-        warn!(
-            "Downstream: Shutting down sv1 downstream reader {}",
-            connection_id
-        );
+            if let Ok(stats_sender) = downstream.safe_lock(|d| d.stats_sender.clone()) {
+                stats_sender.remove_stats(connection_id);
+            }
+            // No message to receive
+            warn!(
+                "Downstream: Shutting down sv1 downstream reader {}",
+                connection_id
+            );
 
-        if let Err(e) = Downstream::remove_downstream_hashrate_from_channel(&downstream) {
-            error!("Failed to remove downstream hashrate from channel: {}", e)
-        };
-        if task_manager_clone
-            .safe_lock(|tm| tm.abort_tasks_for_connection_id(connection_id))
-            .is_err()
-        {
-            error!("TaskManager mutex poisoned")
-        };
-    });
+            if let Err(e) = Downstream::remove_downstream_hashrate_from_channel(&downstream) {
+                error!("Failed to remove downstream hashrate from channel: {}", e)
+            };
+            // Apparently there is no way to make the compiler happy without unwrapping here. But
+            // is not an issue since:
+            // 1. the mutex should never get poisioned and if it does will be very very rare
+            // 2. restarting the process after the unwrapping or restarting the all the tasks from
+            //    inside the process (that is what we should do here) is almost the same thing
+            let send_kill_signal = task_manager
+                .safe_lock(|tm| tm.send_kill_signal.clone())
+                .unwrap();
+            if send_kill_signal.send(connection_id).await.is_err() {
+                error!("Proxy can not abort downstreams tasks");
+                ProxyState::update_inconsistency(Some(1));
+            }
+        })
+    };
     TaskManager::add_receive_downstream(task_manager, handle.into(), connection_id)
         .await
         .map_err(|_| Error::TranslatorTaskManagerFailed)

src/translator/downstream/receive_from_downstream.rs

@@ -1,6 +1,6 @@
 use std::{collections::HashMap, sync::Arc};
 
-use crate::shared::utils::AbortOnDrop;
+use crate::{proxy_state::ProxyState, shared::utils::AbortOnDrop};
 use roles_logic_sv2::utils::Mutex;
 use tokio::sync::mpsc;
 use tracing::warn;
@@ -15,23 +15,31 @@ enum Task {
     SharesMonitor(AbortOnDrop),
 }
 
+type TaskMessage = (Option<u32>, Task);
+
 pub struct TaskManager {
-    send_task: mpsc::Sender<(Option<u32>, Task)>,
+    send_task: mpsc::Sender<TaskMessage>,
     abort: Option<AbortOnDrop>,
-    tasks: Arc<Mutex<HashMap<Option<u32>, Vec<AbortOnDrop>>>>, // Track tasks by connection_id
+    pub send_kill_signal: mpsc::Sender<u32>,
 }
 
 impl TaskManager {
     pub fn initialize() -> Arc<Mutex<Self>> {
-        type TaskMessage = (Option<u32>, Task);
-
         let (sender, mut receiver): (mpsc::Sender<TaskMessage>, mpsc::Receiver<TaskMessage>) =
             mpsc::channel(10);
+        let (send_kill_signal, mut receiver_kill_signal) = mpsc::channel(10);
 
         let tasks = Arc::new(Mutex::new(HashMap::new()));
         let task_clone = tasks.clone();
         let handle = tokio::task::spawn(async move {
             while let Some((connection_id, task)) = receiver.recv().await {
+                // The tasks map is used to save task related to downstream managment, some of them
+                // are "global" in the sense that live for all the life of the transalator (like
+                // the task that create new downstreams when a downstream connect) others are
+                // specific to a downstream like the one that receive messages from it. Specific
+                // task have an id that is the connnection id, "global" ones do not have one; for
+                // that TaskMessage is an (Option<u32>, Task) where u32 is the connection id.
+                // "Global" tasks are saved in the map under the None key.
                 if task_clone
                     .safe_lock(|tasks| {
                         let tasks_list: &mut Vec<AbortOnDrop> =
@@ -48,10 +56,33 @@ impl TaskManager {
                 tokio::time::sleep(std::time::Duration::from_secs(1000)).await;
             }
         });
+        let kill_tasks = tokio::task::spawn(async move {
+            while let Some(connection_id) = receiver_kill_signal.recv().await {
+                if tasks
+                    .safe_lock(|tasks| {
+                        if let Some(handles) = tasks.remove(&Some(connection_id)) {
+                            for handle in handles {
+                                drop(handle);
+                            }
+                        }
+                    })
+                    .is_err()
+                {
+                    tracing::error!("TasKManager Mutex Poisoned");
+                    ProxyState::update_inconsistency(Some(1));
+                };
+                tracing::info!(
+                    "Aborted all tasks for downstream connection ID {}",
+                    connection_id
+                );
+            }
+        });
+        let mut aborter: AbortOnDrop = handle.into();
+        aborter.add_task(kill_tasks);
         Arc::new(Mutex::new(Self {
             send_task: sender,
-            abort: Some(handle.into()),
-            tasks,
+            abort: Some(aborter),
+            send_kill_signal,
         }))
     }
 
@@ -124,29 +155,7 @@ impl TaskManager {
             .await
             .map_err(|_| ())
     }
-
-    /// Kills all tasks for a given `connection_id` and removes them from TaskManager.
-    pub fn abort_tasks_for_connection_id(&mut self, connection_id: u32) {
-        if self
-            .tasks
-            .safe_lock(|tasks| {
-                if let Some(handles) = tasks.remove(&Some(connection_id)) {
-                    for handle in handles {
-                        drop(handle);
-                    }
-                }
-            })
-            .is_err()
-        {
-            tracing::error!("TasKManager Mutex Poisoned")
-        };
-        tracing::info!(
-            "Aborted all tasks for downstream connection ID {}",
-            connection_id
-        );
-    }
 }
-
 /// Converts a `Task` into its `AbortHandle` for task management.
 impl From<Task> for AbortOnDrop {
     fn from(task: Task) -> Self {