如何获取Tokio框架中活跃任务的数量？

我需要在Tokio中获取当前运行的活跃任务数量，类似Python里用len(asyncio.all_tasks())获取事件循环中未完成任务数的功能。以下是示例代码，期望实现active_tasks()函数来得到预期输出：

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());
}

预期输出：

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

注：active_tasks()可以是异步函数。

Tokio本身没有公开的API直接获取所有活跃任务（这是由于其多线程任务调度的内部封装设计），但可以通过手动维护任务计数器的方式实现需求：

实现思路 #

使用线程安全的原子计数器，在每次启动任务时递增计数，任务执行完毕后递减计数；同时将main任务本身计入计数。

完整代码实现 #

use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tokio; // 1.24.1
use tokio::time::sleep;

// 全局线程安全计数器，跟踪活跃任务数
static ACTIVE_TASKS: AtomicUsize = AtomicUsize::new(0);

/// 封装Tokio的spawn，自动维护任务计数
fn spawn_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
where
    F: std::future::Future + Send + 'static,
    F::Output: Send + 'static,
{
    ACTIVE_TASKS.fetch_add(1, Ordering::SeqCst);
    tokio::spawn(async move {
        let result = future.await;
        ACTIVE_TASKS.fetch_sub(1, Ordering::SeqCst);
        result
    })
}

fn active_tasks() -> usize {
    ACTIVE_TASKS.load(Ordering::SeqCst)
}

#[tokio::main]
async fn main() {
    // 将main任务计入活跃任务计数
    ACTIVE_TASKS.fetch_add(1, Ordering::SeqCst);

    spawn_task(async { sleep(Duration::from_secs(5)).await });
    spawn_task(async { sleep(Duration::from_secs(1)).await });
    spawn_task(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());

    // main任务结束前移除计数
    ACTIVE_TASKS.fetch_sub(1, Ordering::SeqCst);
}

代码说明 #

1. 使用AtomicUsize保证多线程环境下计数操作的原子性和正确性。
2. 自定义spawn_task函数替代原生tokio::spawn，自动处理任务启动和结束时的计数变更。
3. 手动维护main任务的计数，因为它本身也是Tokio调度的任务之一。
4. 该实现兼容Tokio默认的多线程runtime配置。

如果需要异步版本的active_tasks()，可以直接修改函数签名（无需额外异步操作，仅保持接口一致性）：

async fn active_tasks() -> usize {
    ACTIVE_TASKS.load(Ordering::SeqCst)
}

内容的提问来源于stack exchange，提问作者coder3101