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);