retry download part failures (#49)

Author: aajtodd

aws-s3-transfer-manager/Cargo.toml

@@ -17,11 +17,13 @@ aws-smithy-runtime-api = "1.7.1"
 aws-smithy-types = "1.2.0"
 aws-types = "1.3.3"
 bytes = "1"
+futures-util = "0.3.30"
 # FIXME - upgrade to hyper 1.x
 hyper = { version = "0.14.29", features = ["client"] }
 path-clean = "1.0.1"
+pin-project-lite = "0.2.14"
 tokio = { version = "1.38.0", features = ["rt-multi-thread", "io-util", "sync", "fs", "macros"] }
-tower = { version = "0.5.0", features = ["limit", "util"] }
+tower = { version = "0.5.0", features = ["limit", "retry", "util"] }
 tracing = "0.1"
 
 [dev-dependencies]
@@ -31,10 +33,13 @@ aws-smithy-runtime = { version = "1.7.1", features = ["client", "connector-hyper
 clap = { version = "4.5.7", default-features = false, features = ["derive", "std", "help"] }
 console-subscriber = "0.3.0"
 http-02x = { package = "http", version = "0.2.9" }
+http-body-1x = { package = "http-body", version = "1" }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 tempfile = "3.10.1"
 fastrand = "2.1.0"
 walkdir = "2"
+tower-test = "0.4.0"
+tokio-test = "0.4.4"
 
 [target.'cfg(not(target_env = "msvc"))'.dev-dependencies]
 jemallocator = "0.5.4"

aws-s3-transfer-manager/src/lib.rs

@@ -43,5 +43,8 @@ pub mod operation;
 /// Transfer manager configuration
 pub mod config;
 
+/// Tower related middleware and components
+pub(crate) mod middleware;
+
 pub use self::client::Client;
 pub use self::config::Config;

aws-s3-transfer-manager/src/middleware.rs

@@ -0,0 +1,6 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+pub(crate) mod retry;

aws-s3-transfer-manager/src/middleware/retry.rs

@@ -0,0 +1,70 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_smithy_types::byte_stream::error::Error as ByteStreamError;
+use futures_util::future;
+use std::sync::Arc;
+use tower::retry::budget::{Budget, TpsBudget};
+
+/// A `tower::retry::Policy` implementation for retrying requests
+#[derive(Debug, Clone)]
+pub(crate) struct RetryPolicy {
+    budget: Arc<TpsBudget>,
+    remaining_attempts: usize,
+}
+
+impl Default for RetryPolicy {
+    fn default() -> Self {
+        Self {
+            budget: Arc::new(TpsBudget::default()),
+            remaining_attempts: 2,
+        }
+    }
+}
+
+fn find_source<'a, E: std::error::Error + 'static>(
+    err: &'a (dyn std::error::Error + 'static),
+) -> Option<&'a E> {
+    let mut next = Some(err);
+    while let Some(err) = next {
+        if let Some(matching_err) = err.downcast_ref::<E>() {
+            return Some(matching_err);
+        }
+        next = err.source();
+    }
+    None
+}
+
+impl<Req, Res, E> tower::retry::Policy<Req, Res, E> for RetryPolicy
+where
+    Req: Clone,
+    E: std::error::Error + 'static,
+{
+    type Future = future::Ready<()>;
+
+    fn retry(&mut self, _req: &mut Req, result: &mut Result<Res, E>) -> Option<Self::Future> {
+        match result {
+            Ok(_) => {
+                self.budget.deposit();
+                None
+            }
+            Err(err) => {
+                // the only type of error we care about at this point is errors that come from
+                // reading the body, all other errors go through the SDK retry implementation
+                // already
+                find_source::<ByteStreamError>(err)?;
+                if self.remaining_attempts == 0 || !self.budget.withdraw() {
+                    return None;
+                }
+                self.remaining_attempts -= 1;
+                Some(future::ready(()))
+            }
+        }
+    }
+
+    fn clone_request(&mut self, req: &Req) -> Option<Req> {
+        Some(req.clone())
+    }
+}

aws-s3-transfer-manager/src/operation/download/service.rs

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 use crate::error;
+use crate::middleware::retry;
 use crate::operation::download::header;
 use crate::operation::download::DownloadContext;
 use aws_smithy_types::body::SdkBody;
@@ -61,6 +62,7 @@ pub(super) fn chunk_service(
 
     ServiceBuilder::new()
         .concurrency_limit(ctx.handle.num_workers())
+        .retry(retry::RetryPolicy::default())
         .service(svc)
 }