rust 如何找到当前时雄任务的数量?

iszxjhcz  于 2023-01-30  发布在  其他
关注(0)|答案(1)|浏览(121)

我想得到当前运行的tokio任务的计数。在python中,我可以使用len(asyncio.all_tasks())返回当前运行循环中未完成的任务。我想知道tokio中的任何等价物。
下面是一个示例代码:

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

我期望上面程序的输出打印活动任务的数量,因为main本身就是一个tokio任务,所以我不会惊讶地发现下面的输出:

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

如果需要,active_tasks()可以是异步函数。

8ulbf1ek

8ulbf1ek1#

我希望不稳定的RuntimeMetrics能够为你解决这个问题,但是它似乎是为不同的目的而设计的。我不相信东京能够为你处理这个问题。
话虽如此,这里有一个潜在的解决方案,以实现类似的结果:

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

结果是:

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

这将大致完成您所描述的内容。为了更接近您所寻找的内容,您可以将管理器与lazy_static组合在一起,并将其 Package 在名为spawn的函数中。您还可以从1开始计数器以考虑主线程。

相关问题