feat(network): use Retry-After header for HTTP 429 responses

Author: arlosi

src/cargo/util/network/retry.rs

@@ -47,6 +47,7 @@ use anyhow::Error;
 use rand::Rng;
 use std::cmp::min;
 use std::time::Duration;
+use time::OffsetDateTime;
 
 /// State for managing retrying a network operation.
 pub struct Retry<'a> {
@@ -104,8 +105,8 @@ impl<'a> Retry<'a> {
     pub fn r#try<T>(&mut self, f: impl FnOnce() -> CargoResult<T>) -> RetryResult<T> {
         match f() {
             Err(ref e) if maybe_spurious(e) && self.retries < self.max_retries => {
-                let err_msg = e
-                    .downcast_ref::<HttpNotSuccessful>()
+                let err = e.downcast_ref::<HttpNotSuccessful>();
+                let err_msg = err
                     .map(|http_err| http_err.display_short())
                     .unwrap_or_else(|| e.root_cause().to_string());
                 let left_retries = self.max_retries - self.retries;
@@ -118,7 +119,12 @@ impl<'a> Retry<'a> {
                     return RetryResult::Err(e);
                 }
                 self.retries += 1;
-                RetryResult::Retry(self.next_sleep_ms())
+                let sleep = err
+                    .and_then(Self::parse_retry_after)
+                    // Limit the Retry-After to a maximum value to avoid waiting too long.
+                    .map(|retry_after| retry_after.min(MAX_RETRY_SLEEP_MS))
+                    .unwrap_or_else(|| self.next_sleep_ms());
+                RetryResult::Retry(sleep)
             }
             Err(e) => RetryResult::Err(e),
             Ok(r) => RetryResult::Success(r),
@@ -141,6 +147,40 @@ impl<'a> Retry<'a> {
             )
         }
     }
+
+    /// Parse the HTTP `Retry-After` header.
+    /// Returns the number of milliseconds to wait before retrying according to the header.
+    fn parse_retry_after(response: &HttpNotSuccessful) -> Option<u64> {
+        // Only applies to HTTP 429 (too many requests) and 503 (service unavailable).
+        if !matches!(response.code, 429 | 503) {
+            return None;
+        }
+
+        // Extract the Retry-After header value.
+        let retry_after = response
+            .headers
+            .iter()
+            .filter_map(|h| h.split_once(':'))
+            .map(|(k, v)| (k.trim(), v.trim()))
+            .find(|(k, _)| k.eq_ignore_ascii_case("retry-after"))?
+            .1;
+
+        // First option: Retry-After is a positive integer of seconds to wait.
+        if let Ok(delay_secs) = retry_after.parse::<u32>() {
+            return Some(delay_secs as u64 * 1000);
+        }
+
+        // Second option: Retry-After is a future HTTP date string that tells us when to retry.
+        if let Ok(retry_time) =
+            OffsetDateTime::parse(retry_after, &time::format_description::well_known::Rfc2822)
+        {
+            let now = OffsetDateTime::now_utc();
+            if retry_time > now {
+                return Some((retry_time - now).whole_milliseconds() as u64);
+            }
+        }
+        None
+    }
 }
 
 fn maybe_spurious(err: &Error) -> bool {
@@ -169,7 +209,7 @@ fn maybe_spurious(err: &Error) -> bool {
         }
     }
     if let Some(not_200) = err.downcast_ref::<HttpNotSuccessful>() {
-        if 500 <= not_200.code && not_200.code < 600 {
+        if 500 <= not_200.code && not_200.code < 600 || not_200.code == 429 {
             return true;
         }
     }
@@ -317,3 +357,51 @@ fn curle_http2_stream_is_spurious() {
     let err = curl::Error::new(code);
     assert!(maybe_spurious(&err.into()));
 }
+
+#[test]
+fn retry_after_parsing() {
+    use crate::core::Shell;
+    fn spurious(code: u32, header: &str) -> HttpNotSuccessful {
+        HttpNotSuccessful {
+            code,
+            url: "Uri".to_string(),
+            ip: None,
+            body: Vec::new(),
+            headers: vec![header.to_string()],
+        }
+    }
+
+    let headers = spurious(429, "Retry-After: 10");
+    assert_eq!(Retry::parse_retry_after(&headers), Some(10_000));
+
+    let headers = spurious(429, "retry-after: Fri, 01 Jan 2100 00:00:00 GMT");
+    let expected = (OffsetDateTime::new_utc(
+        time::Date::from_calendar_date(2100, time::Month::January, 1).unwrap(),
+        time::Time::from_hms(0, 0, 0).unwrap(),
+    ) - OffsetDateTime::now_utc())
+    .whole_milliseconds();
+    let actual = Retry::parse_retry_after(&headers).unwrap();
+    assert!((expected.abs_diff(actual.into()) < 1000));
+
+    let headers = spurious(429, "Content-Type: text/html");
+    assert_eq!(Retry::parse_retry_after(&headers), None);
+
+    let headers = spurious(429, "retry-after: Fri, 01 Jan 2000 00:00:00 GMT");
+    assert_eq!(Retry::parse_retry_after(&headers), None);
+
+    let headers = spurious(429, "retry-after: -1");
+    assert_eq!(Retry::parse_retry_after(&headers), None);
+
+    let headers = spurious(400, "retry-after: 1");
+    assert_eq!(Retry::parse_retry_after(&headers), None);
+
+    let gctx = GlobalContext::default().unwrap();
+    *gctx.shell() = Shell::from_write(Box::new(Vec::new()));
+    let mut retry = Retry::new(&gctx).unwrap();
+    match retry
+        .r#try(|| -> CargoResult<()> { Err(anyhow::Error::from(spurious(429, "Retry-After: 7"))) })
+    {
+        RetryResult::Retry(sleep) => assert_eq!(sleep, 7_000),
+        _ => panic!("unexpected non-retry"),
+    }
+}