Enable configuring retryable status codes (#3243)

Author: chlowell

sdk/core/azure_core/CHANGELOG.md

@@ -21,6 +21,7 @@
 - Removed `Poller::wait()` function. Call `await` on a `Poller` to wait for it to complete and, upon success, return the final model.
 - Removed `xml::read_xml_str()`.
 - Renamed `BearerTokenCredentialPolicy` to `BearerTokenAuthorizationPolicy`.
+- Renamed `RetryPolicy::get_retry_headers()` to `RetryPolicy::retry_headers()`
 - Renamed `xml::read_xml()` to `xml::from_xml()` congruent with `json::from_json()`.
 
 ### Bugs Fixed

sdk/core/azure_core/src/http/pipeline.rs

@@ -174,6 +174,7 @@ impl Pipeline {
             retry_headers: RetryHeaders {
                 retry_headers: vec![X_MS_RETRY_AFTER_MS, RETRY_AFTER_MS, RETRY_AFTER],
             },
+            ..PipelineOptions::default()
         });
 
         Self(http::Pipeline::new(

sdk/core/typespec_client_core/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 
+- Added `PipelineOptions::retry_status_codes` for configuring which status codes should trigger a retry.
 - Added `UrlExt::append_path()`.
 
 ### Breaking Changes

sdk/core/typespec_client_core/src/http/options/mod.rs

@@ -12,7 +12,7 @@ pub use transport::*;
 use crate::http::{
     headers::RETRY_AFTER,
     policies::{Policy, RetryHeaders},
-    Context,
+    Context, StatusCode,
 };
 use std::borrow::Cow;
 use std::fmt::Debug;
@@ -64,7 +64,15 @@ pub struct ClientMethodOptions<'a> {
 pub struct PipelineOptions {
     /// The set of headers which should be considered when
     /// determining the interval to wait for retry attempts.
+    /// This field doesn't apply to custom retry policies.
     pub retry_headers: RetryHeaders,
+
+    /// The status codes that should trigger retries. This
+    /// field doesn't apply to custom retry policies.
+    ///
+    /// When empty, the default retry status codes are used as
+    /// described by [`crate::http::policies::RetryPolicy::retry_status_codes`].
+    pub retry_status_codes: Vec<StatusCode>,
 }
 
 impl Default for PipelineOptions {
@@ -73,6 +81,7 @@ impl Default for PipelineOptions {
             retry_headers: RetryHeaders {
                 retry_headers: vec![RETRY_AFTER],
             },
+            retry_status_codes: Vec::new(),
         }
     }
 }

sdk/core/typespec_client_core/src/http/options/retry.rs

@@ -2,8 +2,12 @@
 // Licensed under the MIT License.
 
 use crate::{
-    http::policies::{
-        ExponentialRetryPolicy, FixedRetryPolicy, NoRetryPolicy, Policy, RetryHeaders, RetryPolicy,
+    http::{
+        policies::{
+            ExponentialRetryPolicy, FixedRetryPolicy, NoRetryPolicy, Policy, RetryHeaders,
+            RetryPolicy,
+        },
+        StatusCode,
     },
     time::Duration,
 };
@@ -87,20 +91,26 @@ impl RetryOptions {
         }
     }
 
-    pub(crate) fn to_policy(&self, retry_headers: RetryHeaders) -> Arc<dyn Policy> {
+    pub(crate) fn to_policy(
+        &self,
+        retry_headers: RetryHeaders,
+        retry_status_codes: &[StatusCode],
+    ) -> Arc<dyn Policy> {
         match &self.mode {
             RetryMode::Exponential(options) => Arc::new(ExponentialRetryPolicy::new(
                 options.initial_delay,
                 options.max_retries,
                 options.max_total_elapsed,
                 options.max_delay,
                 retry_headers,
+                retry_status_codes.to_vec(),
             )),
             RetryMode::Fixed(options) => Arc::new(FixedRetryPolicy::new(
                 options.delay,
                 options.max_retries,
                 options.max_total_elapsed,
                 retry_headers,
+                retry_status_codes.to_vec(),
             )),
             RetryMode::Custom(c) => c.clone(),
             RetryMode::None => Arc::new(NoRetryPolicy::new(retry_headers)),

sdk/core/typespec_client_core/src/http/pipeline.rs

@@ -73,7 +73,10 @@ impl Pipeline {
 
         let pipeline_options = pipeline_options.unwrap_or_default();
 
-        let retry_policy = options.retry.to_policy(pipeline_options.retry_headers);
+        let retry_policy = options.retry.to_policy(
+            pipeline_options.retry_headers.clone(),
+            &pipeline_options.retry_status_codes,
+        );
         pipeline.push(retry_policy);
 
         pipeline.extend_from_slice(&per_try_policies);
@@ -134,10 +137,8 @@ mod tests {
     use crate::{
         error::{Error, ErrorKind},
         http::{
-            headers::{Headers, RETRY_AFTER},
-            policies::{PolicyResult, RetryHeaders},
-            BufResponse, FixedRetryOptions, JsonFormat, Method, Response, RetryOptions, StatusCode,
-            Transport,
+            headers::Headers, policies::PolicyResult, BufResponse, FixedRetryOptions, JsonFormat,
+            Method, Response, RetryOptions, StatusCode, Transport,
         },
         stream::BytesStream,
         Bytes,
@@ -180,12 +181,7 @@ mod tests {
                 transport: Some(Transport::with_policy(Arc::new(Responder {}))),
                 ..Default::default()
             };
-            let pipeline_options = PipelineOptions {
-                retry_headers: RetryHeaders {
-                    retry_headers: vec![RETRY_AFTER],
-                },
-            };
-            let pipeline = Pipeline::new(options, Vec::new(), Vec::new(), Some(pipeline_options));
+            let pipeline = Pipeline::new(options, Vec::new(), Vec::new(), None);
             let mut request = Request::new("http://localhost".parse().unwrap(), Method::Get);
             let raw_response = pipeline
                 .send(&Context::default(), &mut request, None)

sdk/core/typespec_client_core/src/http/policies/retry/exponential.rs

@@ -2,7 +2,10 @@
 // Licensed under the MIT License.
 
 use super::RetryPolicy;
-use crate::{http::policies::RetryHeaders, time::Duration};
+use crate::{
+    http::{policies::RetryHeaders, StatusCode},
+    time::Duration,
+};
 
 /// Retry policy with exponential back-off.
 ///
@@ -18,6 +21,7 @@ pub(crate) struct ExponentialRetryPolicy {
     max_elapsed: Duration,
     max_delay: Duration,
     retry_headers: RetryHeaders,
+    retry_status_codes: Vec<StatusCode>,
 }
 
 impl ExponentialRetryPolicy {
@@ -27,13 +31,15 @@ impl ExponentialRetryPolicy {
         max_elapsed: Duration,
         max_delay: Duration,
         retry_headers: RetryHeaders,
+        retry_status_codes: Vec<StatusCode>,
     ) -> Self {
         Self {
             initial_delay: initial_delay.max(Duration::milliseconds(1)),
             max_retries,
             max_elapsed,
             max_delay: max_delay.max(Duration::seconds(1)),
             retry_headers,
+            retry_status_codes,
         }
     }
 }
@@ -43,10 +49,14 @@ impl RetryPolicy for ExponentialRetryPolicy {
         retry_count >= self.max_retries || time_since_start >= self.max_elapsed
     }
 
-    fn get_retry_headers(&self) -> Option<&RetryHeaders> {
+    fn retry_headers(&self) -> Option<&RetryHeaders> {
         Some(&self.retry_headers)
     }
 
+    fn retry_status_codes(&self) -> &[StatusCode] {
+        &self.retry_status_codes
+    }
+
     fn sleep_duration(&self, retry_count: u32) -> Duration {
         let sleep_ms = self.initial_delay.whole_milliseconds() as u64 * 2u64.pow(retry_count)
             + u64::from(rand::random::<u8>());
@@ -83,6 +93,7 @@ mod tests {
                     RETRY_AFTER,
                 ],
             },
+            vec![],
         );
 
         let mut elapsed_time = Duration::seconds(0);

sdk/core/typespec_client_core/src/http/policies/retry/fixed.rs

@@ -1,7 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-use crate::{http::policies::RetryHeaders, time::Duration};
+use crate::{
+    http::{policies::RetryHeaders, StatusCode},
+    time::Duration,
+};
 
 /// Retry policy with a fixed back-off.
 ///
@@ -15,6 +18,7 @@ pub(crate) struct FixedRetryPolicy {
     max_retries: u32,
     max_elapsed: Duration,
     retry_headers: RetryHeaders,
+    retry_status_codes: Vec<StatusCode>,
 }
 
 impl FixedRetryPolicy {
@@ -23,12 +27,14 @@ impl FixedRetryPolicy {
         max_retries: u32,
         max_elapsed: Duration,
         retry_headers: RetryHeaders,
+        retry_status_codes: Vec<StatusCode>,
     ) -> Self {
         Self {
             delay: delay.max(Duration::milliseconds(10)),
             max_retries,
             max_elapsed,
             retry_headers,
+            retry_status_codes,
         }
     }
 }
@@ -38,10 +44,14 @@ impl super::RetryPolicy for FixedRetryPolicy {
         retry_count >= self.max_retries || time_since_start >= self.max_elapsed
     }
 
-    fn get_retry_headers(&self) -> Option<&RetryHeaders> {
+    fn retry_headers(&self) -> Option<&RetryHeaders> {
         Some(&self.retry_headers)
     }
 
+    fn retry_status_codes(&self) -> &[StatusCode] {
+        &self.retry_status_codes
+    }
+
     fn sleep_duration(&self, _retry_count: u32) -> Duration {
         let sleep_ms = self.delay.whole_milliseconds() as u64 + u64::from(rand::random::<u8>());
         Duration::milliseconds(sleep_ms as i64)