Aborting a Task with Spawn_Blocking

[neomaitre]

Why does using this method to attempt to abort a task that is running in the spawn_blocking thread pool not working as expected, but works otherwise?

Works as intended:

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use tokio::sync::Mutex;
use tokio::task::{AbortHandle, JoinHandle};

#[derive(Clone)]
struct TaskManager {
    task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
    abort_handle: Arc<Mutex<Option<AbortHandle>>>,
    task_counter: Arc<AtomicUsize>,
    task_running: Arc<Mutex<bool>>,
}

impl TaskManager {
    pub fn new() -> Self {
        TaskManager {
            task_handle: Arc::new(Mutex::new(None)),
            abort_handle: Arc::new(Mutex::new(None)),
            task_counter: Arc::new(AtomicUsize::new(0)),
            task_running: Arc::new(Mutex::new(false)),
        }
    }

    pub async fn start_task(self: Arc<Self>) {
        // Cancel any existing task
        {
            let mut task_handle = self.task_handle.lock().await;
            let mut abort_handle = self.abort_handle.lock().await;
            if let Some(handle) = abort_handle.take() {
                println!("Aborting task.");
                handle.abort();
                println!("Task aborted.");
                let mut task_running = self.task_running.lock().await;
                *task_running = false;
            } else {
                println!("No existing task to abort.");
            }
        }

        // Increment task counter
        let task_id = self.task_counter.fetch_add(1, Ordering::SeqCst);
        let task_running = Arc::clone(&self.task_running);

        // Spawn a new task
        let self_clone = Arc::clone(&self);
        let task = tokio::spawn(async move {
            println!("Spawning new task with ID: {}", task_id);

            {
                let mut task_running = task_running.lock().await;
                *task_running = true;
            }

            for i in 0..10 {
                {
                    let task_running = task_running.lock().await;
                    if !*task_running {
                        println!("Task {} detected it was aborted.", task_id);
                        return;
                    }
                }
                println!("Task {} running: iteration {}", task_id, i);
                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
            }
            println!("Task {} completed.", task_id);
        });

        // Create an AbortHandle for the new task
        let abort_handle = task.abort_handle();

        println!("New task spawned with ID: {}", task_id);

        // Store the handle to the new task
        {
            let mut task_handle = self.task_handle.lock().await;
            let mut abort_handle_guard = self.abort_handle.lock().await;
            *task_handle = Some(task);
            *abort_handle_guard = Some(abort_handle);
            println!("New task handle stored.");
        }
    }
}

#[tokio::main]
async fn main() {
    let manager = Arc::new(TaskManager::new());

    // Start initial task
    manager.clone().start_task().await;

    // Simulate waiting and then starting a new task, aborting the previous one
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Simulate another wait and start another new task
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Start a fourth task after another wait
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Keep the main task alive long enough to see the outputs
    tokio::time::sleep(tokio::time::Duration::from_secs(15)).await;
}

Does Not Work (Only Difference is using spawn_blocking):

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use tokio::runtime::Handle;
use tokio::sync::Mutex;
use tokio::task::{AbortHandle, JoinHandle};

#[derive(Clone)]
struct TaskManager {
    task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
    abort_handle: Arc<Mutex<Option<AbortHandle>>>,
    task_counter: Arc<AtomicUsize>,
    task_running: Arc<Mutex<bool>>,
}

impl TaskManager {
    pub fn new() -> Self {
        TaskManager {
            task_handle: Arc::new(Mutex::new(None)),
            abort_handle: Arc::new(Mutex::new(None)),
            task_counter: Arc::new(AtomicUsize::new(0)),
            task_running: Arc::new(Mutex::new(false)),
        }
    }

    pub async fn start_task(self: Arc<Self>) {
        // Cancel any existing task
        {
            let mut task_handle = self.task_handle.lock().await;
            let mut abort_handle = self.abort_handle.lock().await;
            if let Some(handle) = abort_handle.take() {
                println!("Aborting task.");
                handle.abort();
                println!("Task aborted.");
                let mut task_running = self.task_running.lock().await;
                *task_running = false;
            } else {
                println!("No existing task to abort.");
            }
        }

        // Increment task counter
        let task_id = self.task_counter.fetch_add(1, Ordering::SeqCst);
        let task_running = Arc::clone(&self.task_running);

        // Spawn a new task
        let self_clone = Arc::clone(&self);
        let task = tokio::task::spawn_blocking(move || {
            println!("Spawning new task with ID: {}", task_id);
            let rt = Handle::current();
            rt.block_on(async {
                {
                    let mut task_running = task_running.lock().await;
                    *task_running = true;
                }

                for i in 0..10 {
                    {
                        let task_running = task_running.lock().await;
                        if !*task_running {
                            println!("Task {} detected it was aborted.", task_id);
                            return;
                        }
                    }
                    println!("Task {} running: iteration {}", task_id, i);
                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                }
                println!("Task {} completed.", task_id);
            })
        });

        // Create an AbortHandle for the new task
        let abort_handle = task.abort_handle();

        println!("New task spawned with ID: {}", task_id);

        // Store the handle to the new task
        {
            let mut task_handle = self.task_handle.lock().await;
            let mut abort_handle_guard = self.abort_handle.lock().await;
            *task_handle = Some(task);
            *abort_handle_guard = Some(abort_handle);
            println!("New task handle stored.");
        }
    }
}

#[tokio::main]
async fn main() {
    let manager = Arc::new(TaskManager::new());

    // Start initial task
    manager.clone().start_task().await;

    // Simulate waiting and then starting a new task, aborting the previous one
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Simulate another wait and start another new task
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Start a fourth task after another wait
    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    manager.clone().start_task().await;

    // Keep the main task alive long enough to see the outputs
    tokio::time::sleep(tokio::time::Duration::from_secs(15)).await;
}

[Darksonn]

We discussed this on discord, and it lead to this PR: #6571