From 8c3e24f0cb6aa86077353415050a19285bddc820 Mon Sep 17 00:00:00 2001
From: David Yaffe <dayaffe@amazon.com>
Date: Tue, 26 Sep 2023 17:36:24 -0400
Subject: [PATCH] feat: add ability to configure timeout in
 AWSClientConfiguration (#1152)

---
 .../AWSClientConfiguration.swift              | 27 ++++++++++++++-----
 .../AWSClientConfigurationTests.swift         |  9 +++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/Sources/Core/AWSClientRuntime/AWSClientConfiguration.swift b/Sources/Core/AWSClientRuntime/AWSClientConfiguration.swift
index d92d6964f0a..2ec22d16b17 100644
--- a/Sources/Core/AWSClientRuntime/AWSClientConfiguration.swift
+++ b/Sources/Core/AWSClientRuntime/AWSClientConfiguration.swift
@@ -102,6 +102,11 @@ public class AWSClientConfiguration<ServiceSpecificConfiguration: AWSServiceSpec
     /// This structure is custom code-generated for each AWS service.
     public var serviceSpecific: ServiceSpecificConfiguration
 
+    /// The timeout for a request in milliseconds
+    ///
+    /// If none is provided the client will use default values based on the platform.
+    public var connectTimeoutMs: UInt32?
+
     /// Internal designated init
     /// All convenience inits should call this.
     private init(
@@ -115,7 +120,8 @@ public class AWSClientConfiguration<ServiceSpecificConfiguration: AWSServiceSpec
         _ useFIPS: Swift.Bool?,
         _ retryStrategyOptions: RetryStrategyOptions?,
         _ appID: String?,
-        _ awsRetryMode: AWSRetryMode
+        _ awsRetryMode: AWSRetryMode,
+        _ connectTimeoutMs: UInt32? = nil
     ) throws {
         typealias RuntimeConfigType =
             DefaultSDKRuntimeConfiguration<DefaultRetryStrategy, DefaultRetryErrorInfoProvider>
@@ -128,8 +134,13 @@ public class AWSClientConfiguration<ServiceSpecificConfiguration: AWSServiceSpec
         self.useDualStack = useDualStack
         self.useFIPS = useFIPS
         self.clientLogMode = RuntimeConfigType.defaultClientLogMode
+        self.connectTimeoutMs = connectTimeoutMs
         self.httpClientConfiguration = RuntimeConfigType.defaultHttpClientConfiguration
-        self.httpClientEngine = RuntimeConfigType.defaultHttpClientEngine
+        if let connectTimeoutMs = self.connectTimeoutMs {
+            self.httpClientEngine = RuntimeConfigType.httpClientEngineWithTimeout(timeoutMs: connectTimeoutMs)
+        } else {
+            self.httpClientEngine = RuntimeConfigType.defaultHttpClientEngine
+        }
         self.idempotencyTokenGenerator = RuntimeConfigType.defaultIdempotencyTokenGenerator
         self.logger = RuntimeConfigType.defaultLogger(clientName: self.serviceSpecific.clientName)
         self.retryStrategyOptions = retryStrategyOptions ?? RuntimeConfigType.defaultRetryStrategyOptions
@@ -162,7 +173,8 @@ extension AWSClientConfiguration {
         useFIPS: Swift.Bool? = nil,
         retryMode: AWSRetryMode? = nil,
         maxAttempts: Int? = nil,
-        appID: String? = nil
+        appID: String? = nil,
+        connectTimeoutMs: UInt32? = nil
     ) async throws {
         let fileBasedConfig = try await CRTFileBasedConfiguration.makeAsync()
         let resolvedRegionResolver = try regionResolver ?? DefaultRegionResolver { _, _ in fileBasedConfig }
@@ -200,7 +212,8 @@ extension AWSClientConfiguration {
             useFIPS,
             retryStrategyOptions,
             resolvedAppID,
-            resolvedAWSRetryMode
+            resolvedAWSRetryMode,
+            connectTimeoutMs
         )
     }
 
@@ -214,7 +227,8 @@ extension AWSClientConfiguration {
         useFIPS: Swift.Bool? = nil,
         retryMode: AWSRetryMode? = nil,
         maxAttempts: Int? = nil,
-        appID: String? = nil
+        appID: String? = nil,
+        connectTimeoutMs: UInt32? = nil
     ) throws {
         let fileBasedConfig = try CRTFileBasedConfiguration.make()
         let resolvedCredentialsProvider: CredentialsProviding
@@ -245,7 +259,8 @@ extension AWSClientConfiguration {
             useFIPS,
             retryStrategyOptions,
             resolvedAppID,
-            resolvedAWSRetryMode
+            resolvedAWSRetryMode,
+            connectTimeoutMs
         )
     }
 
diff --git a/Tests/Core/AWSClientRuntimeTests/AWSClientConfigurationTests.swift b/Tests/Core/AWSClientRuntimeTests/AWSClientConfigurationTests.swift
index 04e9aca4451..1a6769ab5fb 100644
--- a/Tests/Core/AWSClientRuntimeTests/AWSClientConfigurationTests.swift
+++ b/Tests/Core/AWSClientRuntimeTests/AWSClientConfigurationTests.swift
@@ -63,6 +63,14 @@ class AWSClientConfigurationTests: XCTestCase {
         let subject = try await Subject(region: region, appID: appID)
         XCTAssertEqual(subject.appID, appID)
     }
+
+    // MARK: - Timeout
+
+    func test_sync_configureTimeoutOptionsFromParams() throws {
+        let customTimeout: UInt32 = 10_000
+        let subject = try Subject(region: region, connectTimeoutMs: customTimeout)
+        XCTAssertEqual(subject.connectTimeoutMs, customTimeout)
+    }
 }
 
 struct TestAWSServiceSpecificConfiguration: AWSServiceSpecificConfiguration {
@@ -78,4 +86,5 @@ struct TestAWSServiceSpecificConfiguration: AWSServiceSpecificConfiguration {
 
     var serviceName: String { "TestAWSService" }
     var clientName: String { "TestAWSServiceClient" }
+    var connectTimeoutMs: UInt32 { 10_000 }
 }