From c20ed5fad272dbe5e26e3de5b47e6205ec057d16 Mon Sep 17 00:00:00 2001
From: Sonee <sone102911@gmail.com>
Date: Sun, 29 Mar 2026 13:53:16 +0100
Subject: [PATCH] Fix task unregistration to consistently remove tasks from
 both DB and DelayQueue
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Context

`unregister_task` previously removed a task only from the database, and from the `DelayQueue` only via `process_cancel_request`.
This made the API incomplete and could lead to desynchronization and “zombie tasks.”

Changes
`unregister_task:`
now accepts `&mut self`
removes the task from both the DelayQueue and the database
the duplicate `remove_task_from_queue` has been removed from process_cancel_request
the fallback (if the task is not in the database) has been preserved
---
 magicblock-task-scheduler/src/service.rs | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/magicblock-task-scheduler/src/service.rs b/magicblock-task-scheduler/src/service.rs
index 84b4032d3..3ecb950b0 100644
--- a/magicblock-task-scheduler/src/service.rs
+++ b/magicblock-task-scheduler/src/service.rs
@@ -221,9 +221,7 @@ impl TaskSchedulerService {
             return Ok(());
         }
 
-        self.remove_task_from_queue(cancel_request.task_id);
-
-        // Remove task from database
+        // Remove task from queue and database
         self.unregister_task(cancel_request.task_id).await?;
 
         Ok(())
@@ -278,17 +276,21 @@ impl TaskSchedulerService {
         }
 
         self.db.insert_task(&task).await?;
-        self.task_queue
+        self.remove_task_from_queue(task.id);
+        let key = self
+            .task_queue
             .insert(task.clone(), Duration::from_millis(0));
+        self.task_queue_keys.insert(task.id, key);
         debug!("Registered task {} from context", task.id);
 
         Ok(())
     }
 
     pub async fn unregister_task(
-        &self,
+        &mut self,
         task_id: i64,
     ) -> TaskSchedulerResult<()> {
+        self.remove_task_from_queue(task_id);
         self.db.remove_task(task_id).await?;
         debug!("Removed task {} from database", task_id);