Fix proxy restart

d07b7f5 introduced a way to kill
specific tasks that handle downstream connections: so when a downstream
will disconnect we make sure to kill all the tasks related with that
downstream. In doing that it saved the task aborter in a field of the
TaskManager making impossible to kill the tasks when the aborter was
called somewhere else. Now we make sure that all the aborter live only
under tasks that get killed when the main aborter of the TaskManager is
called.

Author: fi3

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 {