Skip to content

Commit

Permalink
✨ Add task::CloseAndWait::close_and_wait()
Browse files Browse the repository at this point in the history
  • Loading branch information
czy-29 committed Dec 12, 2024
1 parent 672d2cf commit 0ac7790
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 1 deletion.
3 changes: 2 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,10 @@ ron = "0.8.1"
serde = { version = "1.0.216", features = ["derive"] }
thiserror = "2.0.6"
tokio = { version = "1.42.0", features = ["rt"] }
tokio-util = { version = "0.7.13", features = ["rt"] }

[dev-dependencies]
tokio = { version = "1.42.0", features = ["macros"] }
tokio = { version = "1.42.0", features = ["macros", "time"] }

[build-dependencies]
version_check = "0.9.5"
1 change: 1 addition & 0 deletions README-CN.md
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
部分条目如下:
- [`AnyRes`](https://docs.rs/est/latest/est/result/type.AnyRes.html)
- [`collections::MapExt::replace_key()`](https://docs.rs/est/latest/est/collections/trait.MapExt.html#tymethod.replace_key)
- [`task::CloseAndWait::close_and_wait()`](https://docs.rs/est/latest/est/task/trait.CloseAndWait.html#tymethod.close_and_wait)
- [`task::TaskId`](https://docs.rs/est/latest/est/task/struct.TaskId.html)
- [`thread::ThreadId`](https://docs.rs/est/latest/est/thread/struct.ThreadId.html)

Expand Down
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ Latest version: [v0.5.0](https://github.com/opensound-org/est/releases/tag/v0.5.
Some of the items are as follows:
- [`AnyRes`](https://docs.rs/est/latest/est/result/type.AnyRes.html)
- [`collections::MapExt::replace_key()`](https://docs.rs/est/latest/est/collections/trait.MapExt.html#tymethod.replace_key)
- [`task::CloseAndWait::close_and_wait()`](https://docs.rs/est/latest/est/task/trait.CloseAndWait.html#tymethod.close_and_wait)
- [`task::TaskId`](https://docs.rs/est/latest/est/task/struct.TaskId.html)
- [`thread::ThreadId`](https://docs.rs/est/latest/est/thread/struct.ThreadId.html)

Expand Down
49 changes: 49 additions & 0 deletions src/task.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use derive_more::Display;
use serde::{Deserialize, Serialize};
use std::num::NonZeroU64;
use tokio_util::task::{task_tracker::TaskTrackerWaitFuture, TaskTracker};

/// A [`TaskId`](https://docs.rs/tokio/latest/tokio/task/struct.Id.html) that can be `serde`.
#[derive(Debug, Display, Serialize, Deserialize, Copy, Clone, Eq, PartialEq, Hash)]
Expand All @@ -13,6 +14,20 @@ impl From<tokio::task::Id> for TaskId {
}
}

/// Execute [`close`](https://docs.rs/tokio-util/latest/tokio_util/task/task_tracker/struct.TaskTracker.html#method.close)
/// and [`wait`](https://docs.rs/tokio-util/latest/tokio_util/task/task_tracker/struct.TaskTracker.html#method.wait)
/// for [`TaskTracker`](https://docs.rs/tokio-util/latest/tokio_util/task/task_tracker/struct.TaskTracker.html) at once.
pub trait CloseAndWait {
fn close_and_wait(&self) -> TaskTrackerWaitFuture;
}

impl CloseAndWait for TaskTracker {
fn close_and_wait(&self) -> TaskTrackerWaitFuture {
self.close();
self.wait()
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand All @@ -22,4 +37,38 @@ mod tests {
let id = tokio::spawn(async { tokio::task::id() }).await.unwrap();
assert_eq!(id.to_string(), TaskId::from(id).to_string());
}

fn tracker_spawn() -> TaskTracker {
let tracker = TaskTracker::new();

for i in 0..3 {
tracker.spawn(async move { i });
}

tracker
}

#[tokio::test]
async fn close_and_wait() {
use std::time::Duration;
use tokio::time::timeout;

let tracker = tracker_spawn();
assert!(timeout(Duration::from_secs_f64(1.5), tracker.wait())
.await
.is_err());

let tracker = tracker_spawn();
tracker.close();
assert!(timeout(Duration::from_secs_f64(1.5), tracker.wait())
.await
.is_ok());

let tracker = tracker_spawn();
assert!(
timeout(Duration::from_secs_f64(1.5), tracker.close_and_wait())
.await
.is_ok()
);
}
}

0 comments on commit 0ac7790

Please sign in to comment.