fix: How to gracefully join worker thread when rocket shuts down?

[gabyx]

I need a worker thread along rocket which should shutdown gracefully when rocket shuts down.
But the thing does not compile because Fairing::on_shutdown is taking &self instead of &mut self.

I come up with this

/*
[dependencies]
[dependencies.rocket]
version = "=0.5.0"
features = ['json']
*/

use rocket::fairing::{Fairing, Info, Kind};
use std::{
    sync::{
        mpsc::{self, Sender, TryRecvError},
    },
    thread::{self, JoinHandle},
    time::Duration,
};

enum Notification {
    Cancel,
}

struct StatusConsumerShutdown {
    handle: JoinHandle<()>,
    tx: Sender<Notification>,
}

#[rocket::async_trait]
impl Fairing for StatusConsumerShutdown {
    fn info(&self) -> Info {
        Info {
            name: "Status Consumer Shutdown Fairing",
            kind: Kind::Shutdown,
        }
    }

    /// Cannot change this trait to `self`...... 
    async fn on_shutdown(&self, _rocket: &rocket::Rocket<rocket::Orbit>) {
        if let Err(e) = self.tx.send(Notification::Cancel) {
            println!(
                "Could not send cancel notification: error:\n{}", e
            );
        }

        /// I know that the following loop is blocking the executor, 
        /// but that is not relevant for the question.
        while !self.handle.is_finished() {
            self.handle.join(); ///////////// DOES NOT COMPILE DUE TO `&self`.
        }
    }
}

pub fn spawn_status_consumer() -> impl Fairing {
    let (tx, rx) = mpsc::channel::<Notification>();

    let handle = thread::spawn(move || loop {
        println!("Working...");
        thread::sleep(Duration::from_millis(500));

        match rx.try_recv() {
            Ok(Notification::Cancel) => {
                println!("Received cancelation -> Terminating.");
                break;
            }
            Err(TryRecvError::Disconnected) => {
                println!( "Parent thread disconnected -> Terminating.");
                break;
            }
            Err(TryRecvError::Empty) => {}
        }
    });

    return StatusConsumerShutdown { handle, tx };
}

fn main() {
    let shutdown_fairing = spawn_status_consumer();
    rocket::build().attach(shutdown_fairing).launch();
}

This is a cross post, and the question is basically here:

https://users.rust-lang.org/t/how-to-have-joinhandle-inside-rocket-fairing-for-shutdown/105748/2

Thanks a lot for the help.

[SergioBenitez]

Is there a reason you're starting a blocking thread as opposed to a tokio task? If you did the latter, the problem would be simplified. Something like the following, in pseudocode:

impl Fairing for SomeFairing {
    async fn on_shutdown(&self) -> _ {
        self.rx.await;
    }
}

#[launch]
fn rocket() -> _ {
    let (tx, rx) = tokio::sync::oneshot::channel();
    let rocket = rocket::build()
        .mount("/", routes![...])
        .attach(SomeFairing { rx });

    let shutdown = rocket.shutdown();
    tokio::task::spawn(async move {
        loop {
            select! {
                _ = shutdown => break,
                msg = channel => do(msg),
            }
        }

        cleanup().await;
        tx.send(());
    });

    rocket
}

If you must use a regular thread, you can do almost the same thing, but keep your cancellation channel around and use it to send the shutdown message instead of using the Shutdown future.

[gabyx]

The solution troubles me since shutdown inside the loop is not allowed:

 |
65 |                 _ = shutdown => break,
   |                     ^^^^^^^^ value moved here, in previous iteration of loop
   |
   = note: move occurs because `shutdown` has type `rocket::Shutdown`, which does not implement the `Copy` trait

what I am doing wrong?

[gabyx]

pub struct WaitForShutdown {
    pub shutdown: Shutdown,
    pub tx_consumer_done: Sender<()>,
}


/// Spawn a status queue consumer thread.
pub async fn spawn_status_consumer(
    l: Arc<Logger>,
    queue: StatusQueue,
    database_url: &str,
    wait_for_shutdown: WaitForShutdown,
) {
    let db_conn = database::connect(&l, &database_url);

    let log = l.clone();
    let shutdown = wait_for_shutdown.shutdown;

    tokio::task::spawn(async move {
        loop {
            select! {
                _ = shutdown => break,
                _ = async {
                    info!(log, "Working ...");
                    tokio::time::sleep(Duration::from_millis(100)).await;
                } => ()
            }
        }

        // Notify that we finished.
        wait_for_shutdown.tx_consumer_done.send(());
    });
}

[gabyx]

Ok a shutdown.clone() has helped.
The other problem is:

[SergioBenitez]

It feels like we're trying to solve some problem some specific way because that's the way you tried to solve the problem initially, but there may be an entirely different approach at hand. What is the actual thing you're trying to do? Based on the names in your source code, it sounds like there's some "status consumer"? What does it do, and why does it need to be notified of shutdown?

[gabyx]

I had a status consumer (a green thread next to rocket which consumes from a RabbitMQ queue) and when Rocket gets killed in kubernetes by s Signal, I just want that this status consumer gracefully joins before Rocket asyn main is resolved and the main function exits, which tears down the status consumer (it might be in the middle of writing to a DB ...). So all I wanted it gracefull termination. I solved the above by doing this. Which installs a consumer in the RabbitMQ queue then spawns a async listener for shutdown and when it happens it will unsubscribe the consumer (not sure if this await also waits till any consumer callback is finished) and then notifies the over the transmitt channel that the consumer green-thread has finished... Not sure if everything is right, but it works...