Merge pull request dmnd-pool#96 from Priceless-P/fix/handle-downstream-disconnect

handle downstream disconnection

Author: jbesraa

src/translator/downstream/notify.rs

@@ -123,16 +123,13 @@ pub async fn start_notify(
                     Downstream::send_message_downstream(downstream.clone(), message).await;
                 }
             }
-            // TODO here we want to be sure that on drop this is called
-            let _ = Downstream::remove_downstream_hashrate_from_channel(&downstream);
-            // TODO here we want to kill the tasks
             warn!(
                 "Downstream: Shutting down sv1 downstream job notifier for {}",
                 &host
             );
         })
     };
-    TaskManager::add_notify(task_manager, handle.into())
+    TaskManager::add_notify(task_manager, handle.into(), connection_id)
         .await
         .map_err(|_| Error::TranslatorTaskManagerFailed)
 }
@@ -166,7 +163,7 @@ async fn start_update(
             };
         }
     });
-    TaskManager::add_update(task_manager, handle.into())
+    TaskManager::add_update(task_manager, handle.into(), connection_id)
         .await
         .map_err(|_| Error::TranslatorTaskManagerFailed)
 }

src/translator/downstream/receive_from_downstream.rs

@@ -16,6 +16,7 @@ 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);
@@ -53,8 +54,18 @@ pub async fn start_receive_downstream(
             "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")
+        };
     });
-    TaskManager::add_receive_downstream(task_manager, handle.into())
+    TaskManager::add_receive_downstream(task_manager, handle.into(), connection_id)
         .await
         .map_err(|_| Error::TranslatorTaskManagerFailed)
 }

src/translator/downstream/send_to_downstream.rs

@@ -33,7 +33,7 @@ pub async fn start_send_to_downstream(
             connection_id
         );
     });
-    TaskManager::add_send_downstream(task_manager, handle.into())
+    TaskManager::add_send_downstream(task_manager, handle.into(), connection_id)
         .await
         .map_err(|_| Error::TranslatorTaskManagerFailed)
 }

src/translator/downstream/task_manager.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 use crate::shared::utils::AbortOnDrop;
 use roles_logic_sv2::utils::Mutex;
@@ -16,17 +16,32 @@ enum Task {
 }
 
 pub struct TaskManager {
-    send_task: mpsc::Sender<Task>,
+    send_task: mpsc::Sender<(Option<u32>, Task)>,
     abort: Option<AbortOnDrop>,
+    tasks: Arc<Mutex<HashMap<Option<u32>, Vec<AbortOnDrop>>>>, // Track tasks by connection_id
 }
 
 impl TaskManager {
     pub fn initialize() -> Arc<Mutex<Self>> {
-        let (sender, mut receiver) = mpsc::channel(10);
+        type TaskMessage = (Option<u32>, Task);
+
+        let (sender, mut receiver): (mpsc::Sender<TaskMessage>, mpsc::Receiver<TaskMessage>) =
+            mpsc::channel(10);
+
+        let tasks = Arc::new(Mutex::new(HashMap::new()));
+        let task_clone = tasks.clone();
         let handle = tokio::task::spawn(async move {
-            let mut tasks = vec![];
-            while let Some(task) = receiver.recv().await {
-                tasks.push(task);
+            while let Some((connection_id, task)) = receiver.recv().await {
+                if task_clone
+                    .safe_lock(|tasks| {
+                        let tasks_list: &mut Vec<AbortOnDrop> =
+                            tasks.entry(connection_id).or_default();
+                        tasks_list.push(task.into());
+                    })
+                    .is_err()
+                {
+                    tracing::error!("TasKManager Mutex Poisoned")
+                };
             }
             warn!("Translator downstream task manager stopped, keep alive tasks");
             loop {
@@ -36,6 +51,7 @@ impl TaskManager {
         Arc::new(Mutex::new(Self {
             send_task: sender,
             abort: Some(handle.into()),
+            tasks,
         }))
     }
 
@@ -46,34 +62,44 @@ impl TaskManager {
     pub async fn add_receive_downstream(
         self_: Arc<Mutex<Self>>,
         abortable: AbortOnDrop,
+        connection_id: u32,
     ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::ReceiveDownstream(abortable))
+            .send((Some(connection_id), Task::ReceiveDownstream(abortable)))
             .await
             .map_err(|_| ())
     }
-    pub async fn add_update(self_: Arc<Mutex<Self>>, abortable: AbortOnDrop) -> Result<(), ()> {
+    pub async fn add_update(
+        self_: Arc<Mutex<Self>>,
+        abortable: AbortOnDrop,
+        connection_id: u32,
+    ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::Update(abortable))
+            .send((Some(connection_id), Task::Update(abortable)))
             .await
             .map_err(|_| ())
     }
-    pub async fn add_notify(self_: Arc<Mutex<Self>>, abortable: AbortOnDrop) -> Result<(), ()> {
+    pub async fn add_notify(
+        self_: Arc<Mutex<Self>>,
+        abortable: AbortOnDrop,
+        connection_id: u32,
+    ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::Notify(abortable))
+            .send((Some(connection_id), Task::Notify(abortable)))
             .await
             .map_err(|_| ())
     }
     pub async fn add_send_downstream(
         self_: Arc<Mutex<Self>>,
         abortable: AbortOnDrop,
+        connection_id: u32,
     ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::SendDownstream(abortable))
+            .send((Some(connection_id), Task::SendDownstream(abortable)))
             .await
             .map_err(|_| ())
     }
@@ -83,7 +109,7 @@ impl TaskManager {
     ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::AcceptConnection(abortable))
+            .send((None, Task::AcceptConnection(abortable)))
             .await
             .map_err(|_| ())
     }
@@ -94,8 +120,43 @@ impl TaskManager {
     ) -> Result<(), ()> {
         let send_task = self_.safe_lock(|s| s.send_task.clone()).unwrap();
         send_task
-            .send(Task::SharesMonitor(abortable))
+            .send((None, Task::SharesMonitor(abortable)))
             .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 {
+        match task {
+            Task::AcceptConnection(handle) => handle,
+            Task::ReceiveDownstream(handle) => handle,
+            Task::SendDownstream(handle) => handle,
+            Task::Notify(handle) => handle,
+            Task::Update(handle) => handle,
+            Task::SharesMonitor(handle) => handle,
+        }
+    }
 }