Report retry count on Ok results (#2)

Report retry count on `Ok` results that underwent retries through a `RetryCount` response extension, so that users can include the retry count in error messages.

We discovered in uv that we were not reporting retries for status code errors. This change does for `Ok` results what TrueLayer#159 did for `Err` results.

Author: konstin

reqwest-retry/CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Report retry count on `Ok` results that underwent retries through a `RetryCount` response extension.
+
 ## [0.7.0] - 2024-11-08
 
 ### Breaking changes

reqwest-retry/src/lib.rs

@@ -32,7 +32,7 @@ mod retryable_strategy;
 pub use retry_policies::{policies, Jitter, RetryDecision, RetryPolicy};
 use thiserror::Error;
 
-pub use middleware::RetryTransientMiddleware;
+pub use middleware::{RetryCount, RetryTransientMiddleware};
 pub use retryable::Retryable;
 pub use retryable_strategy::{
     default_on_request_failure, default_on_request_success, DefaultRetryableStrategy,

reqwest-retry/src/middleware.rs

@@ -184,20 +184,51 @@ where
                 }
             };
 
-            // Report whether we failed with or without retries.
             break if n_past_retries > 0 {
-                result.map_err(|err| {
-                    Error::Middleware(
+                // Both `Ok` results (e.g. status code errors) and `Err` results (e.g. an
+                // `io::Error` for from connection reset), and we want to inform the user about the
+                // retries in both cases.
+                match result {
+                    Ok(mut response) => {
+                        response
+                            .extensions_mut()
+                            .insert(RetryCount::new(n_past_retries));
+                        Ok(response)
+                    }
+                    Err(err) => Err(Error::Middleware(
                         RetryError::WithRetries {
                             retries: n_past_retries,
                             err,
                         }
                         .into(),
-                    )
-                })
+                    )),
+                }
             } else {
                 result.map_err(|err| Error::Middleware(RetryError::Error(err).into()))
             };
         }
     }
 }
+
+/// Extension type to store retry count in a response.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct RetryCount(u32);
+
+impl RetryCount {
+    /// Create a new retry count.
+    pub fn new(count: u32) -> Self {
+        Self(count)
+    }
+
+    pub fn value(self) -> u32 {
+        self.0
+    }
+}
+
+impl std::ops::Deref for RetryCount {
+    type Target = u32;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}