How to find number of active tokio task? - asynchronous

Question

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.

Answer 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.