use std::future::Future;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
pub fn spawn_with_token<F, T>(token: CancellationToken, future: F) -> tokio::task::JoinHandle<T>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
tokio::spawn(async move {
let _ = token; future.await
})
}
pub fn spawn_after<F, T>(delay: Duration, future: F) -> tokio::task::JoinHandle<T>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
tokio::spawn(async move {
tokio::time::sleep(delay).await;
future.await
})
}
pub fn spawn_tick<F>(
interval_ms: u64,
token: CancellationToken,
mut future: F,
) -> tokio::task::JoinHandle<()>
where
F: FnMut() + Send + 'static,
{
tokio::spawn(async move {
let mut ticker = tokio::time::interval(Duration::from_millis(interval_ms));
loop {
tokio::select! {
_ = token.cancelled() => break,
_ = ticker.tick() => future(),
}
}
})
}
pub fn spawn_with_timeout<F, T>(
timeout: Duration,
future: F,
) -> tokio::task::JoinHandle<Result<T, TimeoutError>>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
tokio::spawn(async move {
match tokio::time::timeout(timeout, future).await {
Ok(result) => Ok(result),
Err(_) => Err(TimeoutError),
}
})
}
pub async fn with_timeout<F, T>(future: F) -> Result<T, TimeoutError>
where
F: Future<Output = T>,
{
tokio::time::timeout(DEFAULT_IO_TIMEOUT, future)
.await
.map_err(|_| TimeoutError)
}
const DEFAULT_IO_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("task timed out")]
pub struct TimeoutError;
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Instant;
#[tokio::test]
async fn test_spawn_with_token_basic() {
let token = CancellationToken::new();
let handle = spawn_with_token(token, async { 42 });
assert_eq!(handle.await.unwrap(), 42);
}
#[tokio::test]
async fn test_spawn_with_token_cancellation() {
let token = CancellationToken::new();
let token_clone = token.clone();
let token_for_closure = token.clone();
let handle = spawn_with_token(token_clone, async move {
token_for_closure.cancelled().await;
99
});
token.cancel();
assert_eq!(handle.await.unwrap(), 99);
}
#[tokio::test]
async fn test_spawn_after_basic() {
let start = Instant::now();
let handle = spawn_after(Duration::from_millis(50), async { 7 });
let result = handle.await.unwrap();
assert_eq!(result, 7);
assert!(start.elapsed() >= Duration::from_millis(40));
}
#[tokio::test]
async fn test_spawn_after_zero_delay() {
let handle = spawn_after(Duration::from_millis(0), async { 1 });
assert_eq!(handle.await.unwrap(), 1);
}
#[tokio::test]
async fn test_spawn_tick_fires_multiple_times() {
let counter = Arc::new(AtomicUsize::new(0));
let token = CancellationToken::new();
let counter_clone = counter.clone();
let handle = spawn_tick(10, token.clone(), move || {
counter_clone.fetch_add(1, Ordering::SeqCst);
});
tokio::time::sleep(Duration::from_millis(100)).await;
token.cancel();
let _ = handle.await;
assert!(counter.load(Ordering::SeqCst) >= 1);
}
#[tokio::test]
async fn test_spawn_tick_stops_on_cancel() {
let counter = Arc::new(AtomicUsize::new(0));
let token = CancellationToken::new();
let counter_clone = counter.clone();
let handle = spawn_tick(10, token.clone(), move || {
counter_clone.fetch_add(1, Ordering::SeqCst);
});
tokio::time::sleep(Duration::from_millis(30)).await;
token.cancel();
let _ = handle.await;
let count_after_cancel = counter.load(Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(counter.load(Ordering::SeqCst), count_after_cancel);
}
#[tokio::test]
async fn test_spawn_with_timeout_success() {
let handle = spawn_with_timeout(Duration::from_millis(100), async { 42 });
assert_eq!(handle.await.unwrap().unwrap(), 42);
}
#[tokio::test]
async fn test_spawn_with_timeout_failure() {
let handle = spawn_with_timeout(Duration::from_millis(10), async {
tokio::time::sleep(Duration::from_millis(100)).await;
42
});
let result = handle.await.unwrap();
assert!(result.is_err());
assert_eq!(result.unwrap_err(), TimeoutError);
}
#[tokio::test]
async fn test_spawn_with_timeout_zero_timeout() {
let handle = spawn_with_timeout(Duration::from_millis(0), async { 1 });
let result = handle.await.unwrap();
let _ = result;
}
#[tokio::test]
async fn test_timeout_error_display() {
let err = TimeoutError;
assert_eq!(format!("{}", err), "task timed out");
}
#[tokio::test]
async fn test_spawn_multiple_concurrent_tasks() {
let token1 = CancellationToken::new();
let token2 = CancellationToken::new();
let h1 = spawn_with_token(token1, async { 1 });
let h2 = spawn_with_token(token2, async { 2 });
assert_eq!(h1.await.unwrap(), 1);
assert_eq!(h2.await.unwrap(), 2);
}
#[tokio::test]
async fn test_spawn_tick_immediate_first_fire() {
let counter = Arc::new(AtomicUsize::new(0));
let token = CancellationToken::new();
let counter_clone = counter.clone();
let handle = spawn_tick(1, token.clone(), move || {
counter_clone.fetch_add(1, Ordering::SeqCst);
});
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while counter.load(Ordering::SeqCst) < 1 {
tokio::time::sleep(Duration::from_millis(5)).await;
if tokio::time::Instant::now() > deadline {
token.cancel();
let _ = handle.await;
panic!("spawn_tick 首次 tick 未在 5s 内触发");
}
}
token.cancel();
let _ = handle.await;
assert!(
counter.load(Ordering::SeqCst) >= 1,
"spawn_tick 首次 tick 应在创建后立即触发,实际计数为 {}",
counter.load(Ordering::SeqCst)
);
}
}