Normailze UploadContext with the other OperationContext types (#123)

Author: landonxjames

aws-sdk-s3-transfer-manager/src/operation/upload.rs

@@ -120,13 +120,13 @@ async fn put_object(
         .scheduler
         .acquire_permit(PermitType::Network(NetworkPermitContext {
             payload_size_estimate: content_length as u64,
-            bucket_type: ctx.bucket_type(),
+            bucket_type: ctx.state.bucket_type(),
             direction: TransferDirection::Upload,
         }))
         .await?;
 
     let req = copy_fields_to_put_object_request(
-        &ctx.request,
+        &ctx.state.request,
         ctx.client()
             .put_object()
             .body(body)
@@ -139,8 +139,8 @@ async fn put_object(
         .send()
         .instrument(tracing::info_span!(
             "send-upload-part",
-            bucket = ctx.request.bucket().unwrap_or_default(),
-            key = ctx.request.key().unwrap_or_default()
+            bucket = ctx.state.request.bucket().unwrap_or_default(),
+            key = ctx.state.request.key().unwrap_or_default()
         ))
         .await?;
     Ok(UploadOutputBuilder::from(resp).build()?)
@@ -189,16 +189,16 @@ async fn try_start_mpu_upload(
 }
 
 fn new_context(handle: Arc<crate::client::Handle>, req: UploadInput) -> UploadContext {
-    UploadContext {
+    UploadContext::new(
         handle,
-        bucket_type: BucketType::from_bucket_name(req.bucket().expect("bucket is availabe")),
-        request: Arc::new(req),
-    }
+        BucketType::from_bucket_name(req.bucket().expect("bucket is availabe")),
+        req,
+    )
 }
 
 /// start a new multipart upload by invoking `CreateMultipartUpload`
 async fn start_mpu(ctx: &UploadContext) -> Result<UploadOutputBuilder, crate::error::Error> {
-    let req = ctx.request();
+    let req = ctx.state.request();
     let client = ctx.client();
 
     let req = copy_fields_to_mpu_request(req, client.create_multipart_upload());

aws-sdk-s3-transfer-manager/src/operation/upload/context.rs

@@ -4,28 +4,38 @@
  */
 
 use crate::operation::upload::UploadInput;
+use crate::operation::TransferContext;
 use crate::types::BucketType;
 use std::ops::Deref;
 use std::sync::Arc;
 
+pub(crate) type UploadContext = TransferContext<UploadState>;
+
+impl UploadContext {
+    pub(crate) fn new(
+        handle: Arc<crate::client::Handle>,
+        bucket_type: BucketType,
+        req: UploadInput,
+    ) -> Self {
+        let state = Arc::new(UploadState {
+            request: Arc::new(req),
+            bucket_type,
+        });
+        TransferContext { handle, state }
+    }
+}
+
 /// Internal context used to drive a single Upload operation
 #[derive(Debug, Clone)]
-pub(crate) struct UploadContext {
-    /// reference to client handle used to do actual work
-    pub(crate) handle: Arc<crate::client::Handle>,
+pub(crate) struct UploadState {
     /// the original request (NOTE: the body will have been taken for processing, only the other fields remain)
     pub(crate) request: Arc<UploadInput>,
 
     /// Type of S3 bucket targeted by this operation
     pub(crate) bucket_type: BucketType,
 }
 
-impl UploadContext {
-    /// The S3 client to use for SDK operations
-    pub(crate) fn client(&self) -> &aws_sdk_s3::Client {
-        self.handle.config.client()
-    }
-
+impl UploadState {
     /// The original request (sans the body as it will have been taken for processing)
     pub(crate) fn request(&self) -> &UploadInput {
         self.request.deref()

aws-sdk-s3-transfer-manager/src/operation/upload/handle.rs

@@ -126,6 +126,7 @@ async fn abort_multipart_upload(
     while (tasks.join_next().await).is_some() {}
 
     let abort_policy = ctx
+        .state
         .request
         .failed_multipart_upload_policy
         .clone()
@@ -134,7 +135,7 @@ async fn abort_multipart_upload(
         FailedMultipartUploadPolicy::Retain => Ok(AbortedUpload::default()),
         FailedMultipartUploadPolicy::AbortUpload => {
             let abort_mpu_resp = copy_fields_to_abort_mpu_request(
-                &ctx.request,
+                &ctx.state.request,
                 ctx.client()
                     .abort_multipart_upload()
                     .set_upload_id(Some(mpu_data.upload_id.clone())),
@@ -210,7 +211,7 @@ async fn complete_upload(handle: UploadHandle) -> Result<UploadOutput, crate::er
 
             // complete the multipart upload
             let req = copy_fields_to_complete_mpu_request(
-                &handle.ctx.request,
+                &handle.ctx.state.request,
                 handle
                     .ctx
                     .client()

aws-sdk-s3-transfer-manager/src/operation/upload/service.rs

@@ -41,7 +41,7 @@ impl ProvideNetworkPermitContext for UploadPartRequest {
     fn network_permit_context(&self) -> NetworkPermitContext {
         NetworkPermitContext {
             payload_size_estimate: self.part_data.data.len() as u64,
-            bucket_type: self.ctx.bucket_type(),
+            bucket_type: self.ctx.state.bucket_type(),
             direction: TransferDirection::Upload,
         }
     }
@@ -52,7 +52,7 @@ pub(crate) struct UploadHedgePolicy;
 
 impl Policy<UploadPartRequest> for UploadHedgePolicy {
     fn clone_request(&self, req: &UploadPartRequest) -> Option<UploadPartRequest> {
-        if req.ctx.bucket_type() == BucketType::Standard {
+        if req.ctx.state.bucket_type() == BucketType::Standard {
             Some(req.clone())
         } else {
             None
@@ -71,7 +71,7 @@ async fn upload_part_handler(request: UploadPartRequest) -> Result<CompletedPart
     let part_number = part_data.part_number as i32;
 
     let req = copy_fields_to_upload_part_request(
-        &ctx.request,
+        &ctx.state.request,
         ctx.client()
             .upload_part()
             .set_upload_id(Some(request.upload_id))
@@ -139,8 +139,8 @@ pub(super) fn distribute_work(
     // group all spawned tasks together
     let parent_span_for_all_tasks = tracing::debug_span!(
         parent: None, "upload-tasks", // TODO: for upload_objects, parent should be upload-objects-tasks
-        bucket = ctx.request.bucket().unwrap_or_default(),
-        key = ctx.request.key().unwrap_or_default(),
+        bucket = ctx.state.request.bucket().unwrap_or_default(),
+        key = ctx.state.request.key().unwrap_or_default(),
     );
     parent_span_for_all_tasks.follows_from(tracing::Span::current());
 
@@ -225,6 +225,7 @@ pub(super) async fn read_body(
 mod tests {
     use super::*;
     use crate::client::Handle;
+    use crate::operation::upload::context::UploadState;
     use crate::operation::upload::UploadInput;
     use crate::runtime::scheduler::Scheduler;
     use crate::types::ConcurrencyMode;
@@ -240,14 +241,16 @@ mod tests {
                     config: Config::builder().client(s3_client).build(),
                     scheduler: Scheduler::new(ConcurrencyMode::Explicit(1)),
                 }),
-                request: Arc::new(
-                    UploadInput::builder()
-                        .bucket(bucket_name)
-                        .key("test-key")
-                        .build()
-                        .unwrap(),
-                ),
-                bucket_type: BucketType::from_bucket_name(bucket_name),
+                state: Arc::new(UploadState {
+                    request: Arc::new(
+                        UploadInput::builder()
+                            .bucket(bucket_name)
+                            .key("test-key")
+                            .build()
+                            .unwrap(),
+                    ),
+                    bucket_type: BucketType::from_bucket_name(bucket_name),
+                }),
             },
             part_data: PartData::new(1, Bytes::default()),
             upload_id: "test-id".to_string(),