3

I would like to get the count of active running tokio tasks. In python, I can use len(asyncio.all_tasks()) which returns the unfinished tasks for the current running loop. I would like to know any equivalent in tokio.

Here is a sample code:

use std::time::Duration;
use tokio; // 1.24.1
use tokio::time::sleep;

fn active_tasks() -> usize {
    todo!("get active task somehow")
}

#[tokio::main]
async fn main() {
    tokio::spawn(async { sleep(Duration::from_secs(5)).await });
    tokio::spawn(async { sleep(Duration::from_secs(1)).await });
    tokio::spawn(async { sleep(Duration::from_secs(3)).await });

    println!("t = 0, running = {}", active_tasks());

    sleep(Duration::from_secs(2)).await;
    println!("t = 2, running = {}", active_tasks());

    sleep(Duration::from_secs(4)).await;
    println!("t = 6, running = {}", active_tasks());
}

I expect the output of the above program to print number of active task, since main itself is a tokio task, I would not be surprised to find the following output:

t = 0, running = 4
t = 2, running = 3
t = 6, running = 1

active_tasks() can be an async function if required.

coder3101
  • 3,920
  • 3
  • 24
  • 28
  • 1
    Just curious: what do you need this number for? – Chayim Friedman Jan 26 '23 at 15:30
  • I have a web server that spins up a long running task, which in turn spins many other long running task, when server gets abort request it should Ideally abort parent and all its child tasks. I would like to see the number of active task before spawn and after abort to be same. – coder3101 Jan 26 '23 at 15:34
  • 4
    So this is just for debugging purposes? Then you are better using something like https://github.com/tokio-rs/console instead. – Chayim Friedman Jan 26 '23 at 15:40
  • I don't think there is a way to retrieve this number. You can file a tokio [feature request](https://github.com/tokio-rs/tokio/issues/new?assignees=&labels=A-tokio%2C+C-feature-request&template=feature_request.md&title=). – Chayim Friedman Jan 26 '23 at 15:45
  • This surely helps for my use case but I would leave the question open for others because there could be other use case and people coming from another language might be looking for something like `tokio::active_tasks()` which could return the number. (maybe in future) – coder3101 Jan 26 '23 at 15:47
  • Now I'm thinking... If tokio-console can get the list of the task, remotely from the app by using the tracing API, surely we can do the same, no? This will require tokio_unstable but should be possible. When I will have time I'll try to work on that. If somebody wants to do that before me, great. – Chayim Friedman Jan 26 '23 at 15:49
  • I opened for feature request for a simple API. Here: https://github.com/tokio-rs/tokio/issues/5400 – coder3101 Jan 26 '23 at 15:56
  • @ChayimFriedman almost definitely from the tokio-unstable [`RuntimeMetrics`](https://docs.rs/tokio/latest/tokio/runtime/struct.RuntimeMetrics.html). Perhaps by iterating over workers and calling `worker_local_queue_depth()`, though the nuances escape me whether this would count everything or not. – kmdreko Jan 26 '23 at 16:22
  • @kmdreko I tried that, it doesn't work. I assume it only gives the number of finished tasks? Not sure. – Chayim Friedman Jan 26 '23 at 16:32
  • Tried to look into that but it is more complicated than I expected. – Chayim Friedman Jan 26 '23 at 18:04

2 Answers2

2

With tokio 1.29 RuntimeMetrics now has a method active_task_count() which returns the number of active tokio tasks.

use tokio::runtime::Handle;

#[tokio::main]
async fn main() {
    let metrics = Handle::current().metrics();

    let n = metrics.active_tasks_count();
    println!("Runtime has {} active tasks", n);
}
coder3101
  • 3,920
  • 3
  • 24
  • 28
1

I was hoping that the unstable RuntimeMetrics would be albe to solve this for you, but it seems designed for a different purpose. I don't believe Tokio will be able to handle this for you.

With that said, here's a potential solution to achieve a similar result:

use std::{
    future::Future,
    sync::{Arc, Mutex},
    time::Duration,
};
use tokio::time::sleep;

struct ThreadManager {
    thread_count: Arc<Mutex<usize>>,
}

impl ThreadManager {
    #[must_use]
    fn new() -> Self {
        Self {
            thread_count: Arc::new(Mutex::new(0)),
        }
    }

    fn spawn<T>(&self, future: T)
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static,
    {
        // Increment the internal count just before the thread starts.
        let count = Arc::clone(&self.thread_count);
        *count.lock().unwrap() += 1;

        tokio::spawn(async move {
            let result = future.await;
            
            // Once we've executed the future, let's decrement this thread.
            *count.lock().unwrap() -= 1;

            result
        });
    }

    fn thread_count(&self) -> usize {
        // Get a copy of the current thread count.
        *Arc::clone(&self.thread_count).lock().unwrap()
    }
}

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

    manager.spawn(async { sleep(Duration::from_secs(5)).await });
    manager.spawn(async { sleep(Duration::from_secs(1)).await });
    manager.spawn(async { sleep(Duration::from_secs(3)).await });

    println!("t = 0, running = {}", manager.thread_count());

    sleep(Duration::from_secs(2)).await;
    println!("t = 2, running = {}", manager.thread_count());

    sleep(Duration::from_secs(4)).await;
    println!("t = 6, running = {}", manager.thread_count());
}

And the result is:

t = 0, running = 3
t = 2, running = 2
t = 6, running = 0

This will do approximately what you're describing. To get a little closer to what you're looking for, you can combine the manager with lazy_static and wrap it in a function called spawn or something. You can also start the counter at 1 to account for the main thread.

Chayim Friedman
  • 47,971
  • 5
  • 48
  • 77
rburmorrison
  • 612
  • 4
  • 12