Skip to main content

vtcode_commons/
async_utils.rs

1//! Async utility functions
2
3use anyhow::{Context, Result};
4use std::future::Future;
5use std::time::Duration;
6use tokio::time::timeout;
7
8pub const DEFAULT_ASYNC_TIMEOUT: Duration = Duration::from_secs(30);
9pub const SHORT_ASYNC_TIMEOUT: Duration = Duration::from_secs(5);
10pub const LONG_ASYNC_TIMEOUT: Duration = Duration::from_secs(300);
11
12/// Execute a future with a timeout and context
13pub async fn with_timeout<F, T>(fut: F, duration: Duration, context: &str) -> Result<T>
14where
15    F: Future<Output = T>,
16{
17    match timeout(duration, fut).await {
18        Ok(result) => Ok(result),
19        Err(_) => anyhow::bail!("Operation timed out after {:?}: {}", duration, context),
20    }
21}
22
23/// Execute a future with the default timeout
24pub async fn with_default_timeout<F, T>(fut: F, context: &str) -> Result<T>
25where
26    F: Future<Output = T>,
27{
28    with_timeout(fut, DEFAULT_ASYNC_TIMEOUT, context).await
29}
30
31/// Execute a future with a short timeout
32pub async fn with_short_timeout<F, T>(fut: F, context: &str) -> Result<T>
33where
34    F: Future<Output = T>,
35{
36    with_timeout(fut, SHORT_ASYNC_TIMEOUT, context).await
37}
38
39/// Execute a future with a long timeout
40pub async fn with_long_timeout<F, T>(fut: F, context: &str) -> Result<T>
41where
42    F: Future<Output = T>,
43{
44    with_timeout(fut, LONG_ASYNC_TIMEOUT, context).await
45}
46
47/// Retry an operation with exponential backoff
48pub async fn retry_with_backoff<F, Fut, T>(
49    mut op: F,
50    max_retries: usize,
51    initial_delay: Duration,
52    context: &str,
53) -> Result<T>
54where
55    F: FnMut() -> Fut,
56    Fut: Future<Output = Result<T>>,
57{
58    let mut delay = initial_delay;
59    let mut last_error = None;
60
61    for i in 0..=max_retries {
62        match op().await {
63            Ok(result) => return Ok(result),
64            Err(e) => {
65                last_error = Some(e);
66                if i < max_retries {
67                    tokio::time::sleep(delay).await;
68                    delay *= 2;
69                }
70            }
71        }
72    }
73
74    let err = last_error.unwrap_or_else(|| anyhow::anyhow!("Retry failed without error"));
75    Err(err).with_context(|| {
76        format!(
77            "Operation failed after {} retries: {}",
78            max_retries, context
79        )
80    })
81}
82
83/// Sleep with context
84pub async fn sleep_with_context(duration: Duration, _context: &str) {
85    tokio::time::sleep(duration).await;
86}
87
88/// Run multiple futures and wait for all with a timeout
89pub async fn join_all_with_timeout<F, T>(
90    futs: Vec<F>,
91    duration: Duration,
92    context: &str,
93) -> Result<Vec<T>>
94where
95    F: Future<Output = T>,
96{
97    with_timeout(futures::future::join_all(futs), duration, context).await
98}