pipedash_plugin_api/
utils.rs1use std::future::Future;
2use std::time::Duration;
3
4use futures::StreamExt;
5use reqwest::Url;
6
7use crate::http::map_send_error;
8use crate::{
9 PluginError,
10 PluginResult,
11};
12
13pub struct RetryPolicy {
14 pub max_retries: usize,
15 pub initial_delay: Duration,
16 pub exponential_backoff: bool,
17}
18
19impl Default for RetryPolicy {
20 fn default() -> Self {
21 Self {
22 max_retries: 3,
23 initial_delay: Duration::from_millis(100),
24 exponential_backoff: true,
25 }
26 }
27}
28
29impl RetryPolicy {
30 pub fn new(max_retries: usize, initial_delay: Duration, exponential_backoff: bool) -> Self {
31 Self {
32 max_retries,
33 initial_delay,
34 exponential_backoff,
35 }
36 }
37
38 pub async fn retry<F, Fut, T>(&self, operation: F) -> PluginResult<T>
39 where
40 F: Fn() -> Fut,
41 Fut: Future<Output = PluginResult<T>>,
42 {
43 let mut delay = self.initial_delay;
44 let mut attempt = 1;
45
46 loop {
47 match operation().await {
48 Ok(result) => return Ok(result),
49 Err(e) if attempt < self.max_retries && Self::should_retry(&e, delay) => {
50 tokio::time::sleep(delay).await;
51 delay = self.next_delay(delay);
52 attempt += 1;
53 }
54 Err(e) => return Err(e),
55 }
56 }
57 }
58
59 pub async fn retry_unsent<F, Fut>(&self, url: &Url, send: F) -> PluginResult<reqwest::Response>
60 where
61 F: Fn() -> Fut,
62 Fut: Future<Output = Result<reqwest::Response, reqwest::Error>>,
63 {
64 let mut delay = self.initial_delay;
65 let mut attempt = 1;
66
67 loop {
68 match send().await {
69 Ok(response) => return Ok(response),
70 Err(e) if attempt < self.max_retries && e.is_connect() && !e.is_timeout() => {
71 tokio::time::sleep(delay).await;
72 delay = self.next_delay(delay);
73 attempt += 1;
74 }
75 Err(e) => return Err(map_send_error(e, url)),
76 }
77 }
78 }
79
80 fn should_retry(error: &PluginError, delay: Duration) -> bool {
81 error.is_transient() && error.retry_after().is_none_or(|hint| hint <= delay)
82 }
83
84 fn next_delay(&self, delay: Duration) -> Duration {
85 if self.exponential_backoff {
86 delay.saturating_mul(2)
87 } else {
88 delay
89 }
90 }
91}
92
93pub async fn retain_with_ci<F, Fut>(ids: Vec<String>, has_ci: F) -> Vec<String>
94where
95 F: Fn(String) -> Fut,
96 Fut: Future<Output = PluginResult<bool>>,
97{
98 futures::stream::iter(ids)
99 .map(|id| {
100 let check = has_ci(id.clone());
101 async move { check.await.unwrap_or(true).then_some(id) }
102 })
103 .buffered(8)
104 .filter_map(std::future::ready)
105 .collect()
106 .await
107}
108
109#[cfg(test)]
110mod tests {
111 use std::io::{
112 Read,
113 Write,
114 };
115 use std::net::TcpListener;
116 use std::sync::atomic::{
117 AtomicUsize,
118 Ordering,
119 };
120 use std::sync::Arc;
121
122 use super::*;
123
124 fn fast_policy() -> RetryPolicy {
125 RetryPolicy::new(3, Duration::from_millis(5), false)
126 }
127
128 fn client(timeout: Duration) -> reqwest::Client {
129 let _ = rustls::crypto::ring::default_provider().install_default();
130 reqwest::Client::builder().timeout(timeout).build().unwrap()
131 }
132
133 async fn attempts_for(error: fn() -> PluginError) -> usize {
134 let attempts = AtomicUsize::new(0);
135 let result: PluginResult<()> = fast_policy()
136 .retry(|| async {
137 attempts.fetch_add(1, Ordering::SeqCst);
138 Err(error())
139 })
140 .await;
141 assert!(result.is_err());
142 attempts.load(Ordering::SeqCst)
143 }
144
145 #[tokio::test]
146 async fn test_retry_success() {
147 let policy = RetryPolicy::default();
148 let result = policy.retry(|| async { Ok::<_, PluginError>(42) }).await;
149 assert_eq!(result.unwrap(), 42);
150 }
151
152 #[tokio::test]
153 async fn test_retry_eventual_success() {
154 let policy = RetryPolicy::new(3, Duration::from_millis(10), false);
155 let attempts = std::cell::Cell::new(0);
156
157 let result = policy
158 .retry(|| async {
159 let count = attempts.get() + 1;
160 attempts.set(count);
161 if count < 2 {
162 Err(PluginError::NetworkError("Temporary failure".to_string()))
163 } else {
164 Ok(42)
165 }
166 })
167 .await;
168
169 assert_eq!(result.unwrap(), 42);
170 }
171
172 #[tokio::test]
173 async fn retries_network_and_short_hint_unavailable() {
174 assert_eq!(
175 attempts_for(|| PluginError::NetworkError("x".into())).await,
176 3
177 );
178 assert_eq!(
179 attempts_for(|| PluginError::ServiceUnavailable {
180 message: "x".into(),
181 retry_after: None,
182 })
183 .await,
184 3
185 );
186 assert_eq!(
187 attempts_for(|| PluginError::ServiceUnavailable {
188 message: "x".into(),
189 retry_after: Some(Duration::from_millis(1)),
190 })
191 .await,
192 3
193 );
194 }
195
196 #[tokio::test]
197 async fn does_not_retry_permanent_errors_or_long_hints() {
198 assert_eq!(
199 attempts_for(|| PluginError::AuthenticationFailed("x".into())).await,
200 1
201 );
202 assert_eq!(
203 attempts_for(|| PluginError::PipelineNotFound("x".into())).await,
204 1
205 );
206 assert_eq!(attempts_for(|| PluginError::ApiError("x".into())).await, 1);
207 assert_eq!(
208 attempts_for(|| PluginError::RateLimited {
209 message: "x".into(),
210 retry_after: Some(Duration::from_secs(120)),
211 })
212 .await,
213 1
214 );
215 assert_eq!(
216 attempts_for(|| PluginError::ServiceUnavailable {
217 message: "x".into(),
218 retry_after: Some(Duration::from_secs(120)),
219 })
220 .await,
221 1
222 );
223 }
224
225 #[tokio::test]
226 async fn retry_unsent_retries_connect_errors() {
227 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
228 let addr = listener.local_addr().unwrap();
229 drop(listener);
230 let url = Url::parse(&format!("http://{addr}/trigger")).unwrap();
231 let client = client(Duration::from_secs(5));
232 let attempts = AtomicUsize::new(0);
233
234 let result = fast_policy()
235 .retry_unsent(&url, || {
236 attempts.fetch_add(1, Ordering::SeqCst);
237 client.post(url.clone()).send()
238 })
239 .await;
240
241 assert!(matches!(result, Err(PluginError::NetworkError(_))));
242 assert_eq!(attempts.load(Ordering::SeqCst), 3);
243 }
244
245 #[tokio::test]
246 async fn retry_unsent_never_resends_after_timeout() {
247 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
248 let addr = listener.local_addr().unwrap();
249 let received = Arc::new(AtomicUsize::new(0));
250 let counter = received.clone();
251 std::thread::spawn(move || {
252 for mut stream in listener.incoming().flatten() {
253 let mut buf = [0u8; 1024];
254 if stream.read(&mut buf).unwrap_or(0) > 0 {
255 counter.fetch_add(1, Ordering::SeqCst);
256 }
257 std::thread::sleep(Duration::from_millis(300));
258 let _ = stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n");
259 }
260 });
261 let url = Url::parse(&format!("http://{addr}/trigger")).unwrap();
262 let client = client(Duration::from_millis(50));
263
264 let result = fast_policy()
265 .retry_unsent(&url, || client.post(url.clone()).send())
266 .await;
267
268 assert!(matches!(result, Err(PluginError::NetworkError(_))));
269 tokio::time::sleep(Duration::from_millis(100)).await;
270 assert_eq!(received.load(Ordering::SeqCst), 1);
271 }
272
273 #[tokio::test]
274 async fn ids_without_ci_are_dropped_and_unchecked_ones_kept_in_order() {
275 let ids = ["empty", "ci", "unreachable", "ci-too"]
276 .map(String::from)
277 .to_vec();
278
279 let kept = retain_with_ci(ids, |id| async move {
280 match id.as_str() {
281 "empty" => Ok(false),
282 "unreachable" => Err(PluginError::NetworkError("reset".into())),
283 _ => Ok(true),
284 }
285 })
286 .await;
287
288 assert_eq!(kept, ["ci", "unreachable", "ci-too"]);
289 }
290}