reqwest_builder_retry/
lib.rs

1use crate::error::Error;
2use reqwest::{RequestBuilder, Response};
3use std::future::Future;
4use std::time::Duration;
5
6pub mod error;
7pub use reqwest;
8
9#[cfg(feature = "convenience")]
10pub mod convenience;
11
12pub enum RetryType {
13    Stop,
14    Retry,
15    RetryAfter(Duration),
16}
17
18pub async fn execute<
19    SuccessResponse,
20    ErrorResponse,
21    MakerBuilder,
22    CheckDone,
23    JITTER,
24    SLEEPER,
25    FutCheckDone,
26    FutSLEEPER,
27>(
28    make_builder: MakerBuilder,
29    check_done: CheckDone,
30    try_count: u8,
31    retry_duration: Duration,
32    jitter: JITTER,
33    sleeper: SLEEPER,
34) -> Result<SuccessResponse, Error<ErrorResponse>>
35where
36    MakerBuilder: Fn(u8) -> RequestBuilder,
37    CheckDone: Fn(Result<Response, reqwest::Error>) -> FutCheckDone,
38    JITTER: Fn() -> Duration,
39    SLEEPER: Fn(Duration) -> FutSLEEPER,
40    FutCheckDone: Future<Output = Result<SuccessResponse, (RetryType, ErrorResponse)>>,
41    FutSLEEPER: Future<Output = ()>,
42{
43    // リトライ回数が0の時はエラーを返す
44    if try_count == 0 {
45        return Err(Error::NoTry);
46    }
47
48    // 指定回数実行する
49    for i in 0..try_count {
50        let builder = make_builder(i);
51        let response = builder.send().await;
52        let (next_retry_duration, error_response) = match check_done(response).await {
53            Ok(result) => return Ok(result),
54            Err((retry_type, err)) => {
55                match retry_type {
56                    RetryType::Retry => {
57                        // リトライなので停止時間を計算する
58                        (calc_retry_duration(retry_duration, jitter(), i as u32), err)
59                    }
60                    RetryType::RetryAfter(target_duration) => {
61                        // リトライ後の待機時間を指定された時間に設定する
62                        (target_duration, err)
63                    }
64                    RetryType::Stop => {
65                        // 停止指示が来たのでエラーを返す
66                        return Err(Error::Stop(err));
67                    }
68                }
69            }
70        };
71
72        if i >= try_count - 1 {
73            // 最後の実行であれば、エラーを返す
74            return Err(Error::TryOver(error_response));
75        }
76
77        // リトライ時間が0以上なら待機する
78        if next_retry_duration > Duration::ZERO {
79            sleeper(next_retry_duration).await;
80        }
81    }
82
83    // ここに到達することはないが、コンパイラの警告を避けるためにエラーを返す
84    Err(Error::NoTry)
85}
86
87fn calc_retry_duration(
88    retry_duration: Duration,
89    jitter_duration: Duration,
90    try_count: u32,
91) -> Duration {
92    // exponential backoff
93    // 0の時1回、1の時2回、2の時4回、3の時8回
94    let retry_count = 2u64.pow(try_count) as u32;
95    retry_duration * retry_count + jitter_duration
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    fn make_builder_for_test(i: u8) -> RequestBuilder {
103        reqwest::Client::new()
104            .get("https://httpbin.org/get")
105            .header("Try-Count", i.to_string())
106    }
107
108    #[tokio::test]
109    async fn test_stop() {
110        match execute(
111            make_builder_for_test,
112            |_| async move { Err::<serde_json::Value, (RetryType, ())>((RetryType::Stop, ())) },
113            3,
114            Duration::from_secs(1),
115            || Duration::from_millis(100),
116            |_| async move {},
117        )
118        .await
119        {
120            Err(Error::Stop(_)) => {}
121            _ => {
122                panic!("Test failed: Expected TryOver error.");
123            }
124        }
125    }
126
127    #[tokio::test]
128    async fn test_over_try() {
129        match execute(
130            make_builder_for_test,
131            |_| async move { Err::<serde_json::Value, (RetryType, ())>((RetryType::Retry, ())) },
132            4,
133            Duration::from_secs(2),
134            || Duration::from_millis(100),
135            |duration| async move { println!("Sleeping for {:?}", duration) },
136        )
137        .await
138        {
139            Err(Error::TryOver(_)) => {}
140            _ => {
141                panic!("Test failed: Expected TryOver error.");
142            }
143        }
144    }
145
146    #[tokio::test]
147    async fn test_success() {
148        let check_done = |response: Result<Response, _>| async move {
149            let Ok(response) = response else {
150                return Err((RetryType::Retry, ())); // Retry on failure
151            };
152            if !response.status().is_success() {
153                return Err((RetryType::Retry, ())); // Retry on failure
154            }
155            let Ok(json) = response.json::<serde_json::Value>().await else {
156                return Err((RetryType::Retry, ()));
157            };
158            Ok(json)
159        };
160
161        match execute(
162            make_builder_for_test,
163            check_done,
164            3,
165            Duration::from_secs(1),
166            || Duration::from_millis(100),
167            |_| async move {},
168        )
169        .await
170        {
171            Ok(res) => {
172                assert_eq!(res.is_object(), true);
173            }
174            Err(e) => {
175                panic!("Test failed: {:?}", e);
176            }
177        }
178    }
179}