Update AWS SDK to 1.x (#684)

The main change here is that we no longer differentiate between "other" and `Unhandled` errors. The distinction didn't make sense in the first place and is now handled uniformly as a retry case.
TraceMachina · Mar 3, 2024 · cd78ed2 · cd78ed2
1 parent c5851f9
commit cd78ed2
Show file tree

Hide file tree

Showing 11 changed files with 348 additions and 301 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -47,7 +47,7 @@ scopeguard = "1.2.0"
 serde_json5 = "0.1.0"
 tokio = { version = "1.35.1", features = ["rt-multi-thread", "signal"] }
 tokio-rustls = "0.25.0"
-tonic = { version = "0.10.2", features = ["gzip", "tls"] }
+tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tower = "0.4.13"
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
diff --git a/nativelink-error/Cargo.toml b/nativelink-error/Cargo.toml
@@ -14,4 +14,4 @@ hex = "0.4.3"
 prost = "0.12.3"
 prost-types = "0.12.3"
 tokio = { version = "1.35.1" }
-tonic = { version = "0.10.2", features = ["gzip"] }
+tonic = { version = "0.11.0", features = ["gzip"] }
diff --git a/nativelink-proto/Cargo.toml b/nativelink-proto/Cargo.toml
@@ -21,8 +21,8 @@ doctest = false
 [dependencies]
 prost = "0.12.3"
 prost-types = "0.12.3"
-tonic = { version = "0.10.2", features = ["gzip"] }
+tonic = { version = "0.11.0", features = ["gzip"] }
 
 [dev-dependencies]
 prost-build = "0.12.3"
-tonic-build = "0.10.2"
+tonic-build = "0.11.0"
diff --git a/nativelink-scheduler/Cargo.toml b/nativelink-scheduler/Cargo.toml
@@ -26,7 +26,7 @@ rand = "0.8.5"
 scopeguard = "1.2.0"
 tokio = { version = "1.35.1", features = ["sync", "rt", "parking_lot"] }
 tokio-stream = { version = "0.1.14", features = ["sync"] }
-tonic = { version = "0.10.2", features = ["gzip", "tls"] }
+tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tracing = "0.1.40"
 
 [dev-dependencies]

diff --git a/nativelink-service/Cargo.toml b/nativelink-service/Cargo.toml
@@ -19,7 +19,7 @@ prost = "0.12.3"
 rand = "0.8.5"
 tokio = { version = "1.35.1", features = ["sync", "rt"] }
 tokio-stream = { version = "0.1.14", features = ["sync"] }
-tonic = { version = "0.10.2", features = ["gzip", "tls"] }
+tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tracing = "0.1.40"
 uuid = { version = "1.6.1", features = ["v4"] }
 

diff --git a/nativelink-store/Cargo.toml b/nativelink-store/Cargo.toml
@@ -11,9 +11,9 @@ nativelink-proto = { path = "../nativelink-proto" }
 
 async-lock = "3.3.0"
 async-trait = "0.1.74"
-aws-config = "0.57.2"
-aws-sdk-s3 = { version = "0.35.0" }
-aws-smithy-runtime = { version = "0.57.2" }
+aws-config = "1.1.7"
+aws-sdk-s3 = { version = "1.17.0" }
+aws-smithy-runtime = { version = "1.1.7" }
 bincode = "1.3.3"
 blake3 = "1.5.0"
 byteorder = "1.5.0"
@@ -34,15 +34,15 @@ tempfile = "3.9.0"
 tokio = { version = "1.35.1" }
 tokio-stream = { version = "0.1.14", features = ["fs"] }
 tokio-util = { version = "0.7.10" }
-tonic = { version = "0.10.2", features = ["gzip", "tls"] }
+tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tracing = "0.1.40"
 uuid = { version = "1.6.1", features = ["v4"] }
 
 [dev-dependencies]
 pretty_assertions = "1.4.0"
 memory-stats = "1.1.0"
 once_cell = "1.19.0"
-http = "^0.2.11"
-aws-smithy-types = "0.57.2"
-aws-sdk-s3 = { version = "0.35.0"  }
-aws-smithy-runtime = { version = "0.57.2", features = ["test-util"] }
+http = "1.0.0"
+aws-smithy-types = "1.1.7"
+aws-sdk-s3 = { version = "1.17.0"  }
+aws-smithy-runtime = { version = "1.1.7", features = ["test-util"] }
diff --git a/nativelink-store/src/s3_store.rs b/nativelink-store/src/s3_store.rs
@@ -22,6 +22,7 @@ use std::{cmp, env};
 
 use async_trait::async_trait;
 use aws_config::default_provider::credentials;
+use aws_config::BehaviorVersion;
 use aws_sdk_s3::config::Region;
 use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadOutput;
 use aws_sdk_s3::operation::get_object::GetObjectError;
@@ -149,7 +150,7 @@ impl S3Store {
         let s3_client = {
             let http_client = HyperClientBuilder::new().build(TlsConnector::new(config, jitter_fn.clone()));
             let credential_provider = credentials::default_provider().await;
-            let mut config_builder = aws_config::from_env()
+            let mut config_builder = aws_config::defaults(BehaviorVersion::v2023_11_09())
                 .credentials_provider(credential_provider)
                 .region(Region::new(Cow::Owned(config.region.clone())))
                 .http_client(http_client);
@@ -201,19 +202,26 @@ impl S3Store {
 
                 match result {
                     Ok(head_object_output) => {
-                        let sz = head_object_output.content_length;
-                        Some((RetryResult::Ok(Some(sz as usize)), state))
+                        let Some(length) = head_object_output.content_length else {
+                            return Some((RetryResult::Ok(None), state));
+                        };
+                        if length >= 0 {
+                            return Some((RetryResult::Ok(Some(length as usize)), state));
+                        }
+                        Some((
+                            RetryResult::Err(make_err!(
+                                Code::InvalidArgument,
+                                "Negative content length in S3: {length:?}",
+                            )),
+                            state,
+                        ))
                     }
                     Err(sdk_error) => match sdk_error.into_service_error() {
                         HeadObjectError::NotFound(_) => Some((RetryResult::Ok(None), state)),
-                        HeadObjectError::Unhandled(e) => Some((
-                            RetryResult::Retry(make_err!(Code::Unavailable, "Unhandled HeadObjectError in S3: {e:?}")),
-                            state,
-                        )),
                         other => Some((
-                            RetryResult::Err(make_err!(
+                            RetryResult::Retry(make_err!(
                                 Code::Unavailable,
-                                "Unkown error getting head_object in S3: {other:?}",
+                                "Unhandled HeadObjectError in S3: {other:?}"
                             )),
                             state,
                         )),
@@ -470,20 +478,11 @@ impl Store for S3Store {
                                 writer,
                             ));
                         }
-                        GetObjectError::Unhandled(e) => {
-                            return Some((
-                                RetryResult::Retry(make_err!(
-                                    Code::Unavailable,
-                                    "Unhandled GetObjectError in S3: {e:?}",
-                                )),
-                                writer,
-                            ));
-                        }
                         other => {
                             return Some((
-                                RetryResult::Err(make_err!(
+                                RetryResult::Retry(make_err!(
                                     Code::Unavailable,
-                                    "Unkown error getting result in S3: {other:?}"
+                                    "Unhandled GetObjectError in S3: {other:?}",
                                 )),
                                 writer,
                             ));

diff --git a/nativelink-store/tests/s3_store_test.rs b/nativelink-store/tests/s3_store_test.rs
@@ -16,7 +16,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::time::Duration;
 
-use aws_sdk_s3::config::{Builder, Region};
+use aws_sdk_s3::config::{BehaviorVersion, Builder, Region};
 use aws_smithy_runtime::client::http::test_util::{ReplayEvent, StaticReplayClient};
 use aws_smithy_types::body::SdkBody;
 use bytes::Bytes;
@@ -53,6 +53,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -83,6 +84,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -136,6 +138,7 @@ mod s3_store_tests {
             ),
         ]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -188,6 +191,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client.clone())
             .build();
@@ -224,6 +228,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -268,6 +273,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client.clone())
             .build();
@@ -328,6 +334,7 @@ mod s3_store_tests {
             ),
         ]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -455,6 +462,7 @@ mod s3_store_tests {
             ),
         ]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client.clone())
             .build();
@@ -490,6 +498,7 @@ mod s3_store_tests {
                 .unwrap(),
         )]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client.clone())
             .build();
@@ -536,6 +545,7 @@ mod s3_store_tests {
 
         let mock_client = StaticReplayClient::new(vec![]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();
@@ -578,6 +588,7 @@ mod s3_store_tests {
 
         let mock_client = StaticReplayClient::new(vec![]);
         let test_config = Builder::new()
+            .behavior_version(BehaviorVersion::v2023_11_09())
             .region(Region::from_static(REGION))
             .http_client(mock_client)
             .build();

diff --git a/nativelink-util/Cargo.toml b/nativelink-util/Cargo.toml
@@ -26,7 +26,7 @@ serde = { version = "1.0.193", features = ["derive"] }
 sha2 = "0.10.8"
 tokio = { version = "1.35.1", features = [ "sync", "fs", "rt", "time", "io-util", "macros" ] }
 tokio-util = { version = "0.7.10" }
-tonic = { version = "0.10.2", features = ["tls"] }
+tonic = { version = "0.11.0", features = ["tls"] }
 tracing = "0.1.40"
 
 [dev-dependencies]

diff --git a/nativelink-worker/Cargo.toml b/nativelink-worker/Cargo.toml
@@ -30,7 +30,7 @@ serde_json5 = "0.1.0"
 shlex = "1.2.0"
 tokio = { version = "1.35.1", features = ["sync", "rt", "process"] }
 tokio-stream = { version = "0.1.14", features = ["fs"] }
-tonic = { version = "0.10.2", features = ["gzip", "tls"] }
+tonic = { version = "0.11.0", features = ["gzip", "tls"] }
 tracing = "0.1.40"
 uuid = { version = "1.6.1", features = ["v4"] }