Search code examples
rustrust-tokio

The proper method to get tokio runtime handle based on current running environment


What's the idiomatic way to get a tokio runtime handle based on the current running environment?

  • For methods already run in tokio runtime, I want to use Handle.try_current().unwrap() to get the current one.
  • For methods not run in tokio, I can create a new one: Runtime::new().unwrap().handle().

However, when I wrote my code as:

fn get_runtime_handle() -> Handle {
    match Handle::try_current() {
        Ok(h) => h,
        Err(_) => Runtime::new().unwrap().handle().clone(),
    }
}

async fn a_async() -> Result<()> {
    ....
}

fn a() -> Result<()> {
   let handle = get_runtime_handle();
   handle.block_one (async { a_async().await; })
}

fn main() -> Result<()> {
    a();

    Ok(())
}

and call tokio::fs::read_dir inside, the code crash with Error: Custom { kind: Other, error: "background task failed" }.

And when I substitute handle.block_on with Runtime::new().unwrap().handle().block_on in main, the code runs successfully.

I suppose my get_runtime_handle function is problematic, what's the right way to express this? The full runnable code is here.

Besides, when the method get_runtime_handle is running inside a tokio runtime, other unit tests from the project complains:

thread 'main' panicked at 'Cannot start a runtime from within a runtime. 
This happens because a function (like `block_on`) attempted to block the
current thread while the thread is being used to drive asynchronous tasks.

Solution

  • The problem is the lifetime of your new runtime, it is dropped at the end of the get_runtime_handle(). You should return the runtime if the function create one.

    use tokio::runtime::{Runtime, Handle};
    
    fn get_runtime_handle() -> (Handle, Option<Runtime>) {
        match Handle::try_current() {
            Ok(h) => (h, None),
            Err(_) => {
                  let rt = Runtime::new().unwrap();
                  (rt.handle().clone(), Some(rt))
                }
        }
    
    
    fn main() {
        // let _ = Runtime::new().unwrap();
        let (handle, rt) = get_runtime_handle();
    }