Axum how to save spawn task JoinHandle to state so I can abort in future?

Question

I am new to rust, I would like to spawn a task that can be aborted when running but am having trouble figuring out how to achieve what I want.

What I have tried (does not compile):

use std::time::Duration;
use axum::extract::State;
use tokio::time::sleep;
use axum::{
  routing::get,
  Router,
};
use tokio::task::JoinHandle;

#[derive(Clone)]
struct AppState {
  join_handle: Option<JoinHandle<()>>,
}

#[tokio::main]
async fn main() {
  let state = AppState { join_handle: None };

  let app = Router::new()
      .route("/spawn_task", get(spawn_task))
      .route("/abort_task", get(abort_task))
      .with_state(state);

  axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
      .serve(app.into_make_service())
      .await
      .unwrap();
}

pub async fn spawn_task(State(mut state): State<AppState>) {
  state.join_handle = Some(tokio::spawn(async move {
    println!("spawn task");
    sleep(Duration::from_millis(10000)).await;
    println!("spawn task after 10s");
  }));
}

pub async fn abort_task(State(state): State<AppState>) {
  state.join_handle.unwrap().abort();
}

Answer 1

I got it to work with Arc and Mutex. Open to other more idiomatic rust answers.

use std::sync::Arc;
use std::time::Duration;
use axum::extract::State;
use tokio::time::sleep;
use axum::{
  routing::get,
  Router,
};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;


#[derive(Clone)]
pub struct AppState {
  join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
}

#[tokio::main]
async fn main() {
  let state = AppState { join_handle: Arc::new(Mutex::new(None)) };

  let app = Router::new()
      .route("/spawn_task", get(spawn_task))
      .route("/abort_task", get(abort_task))
      .with_state(state);

  axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
      .serve(app.into_make_service())
      .await
      .unwrap();
}

pub async fn spawn_task(State(state): State<AppState>) {
  let mut lock = state.join_handle.lock().await;
  *lock = Some(tokio::spawn(async move {
    println!("spawn task");
    sleep(Duration::from_millis(10000)).await;
    println!("spawn task after 10s");
  }));
}

pub async fn abort_task(State(state): State<AppState>) {
  state.join_handle.as_ref().lock().await.as_ref().unwrap().abort();
}