refactor: make AWS retry policy inheritable (#1122)

awslabs · Nov 20, 2023 · c710c72 · c710c72
1 parent 094809f
commit c710c72
Show file tree

Hide file tree

Showing 7 changed files with 86 additions and 65 deletions.
diff --git a/.changes/4cbc20ed-789f-4040-b81c-328ecd3d0700.json b/.changes/4cbc20ed-789f-4040-b81c-328ecd3d0700.json
@@ -0,0 +1,8 @@
+{
+    "id": "4cbc20ed-789f-4040-b81c-328ecd3d0700",
+    "type": "feature",
+    "description": "Make the AWS retry policy inheritable",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1055"
+    ]
+}
diff --git a/aws-runtime/aws-http/api/aws-http.api b/aws-runtime/aws-http/api/aws-http.api
@@ -49,7 +49,13 @@ public final class aws/sdk/kotlin/runtime/http/operation/CustomUserAgentMetadata
 	public static final fun getCustomUserAgentMetadata (Laws/smithy/kotlin/runtime/operation/ExecutionContext;)Laws/sdk/kotlin/runtime/http/operation/CustomUserAgentMetadata;
 }
 
-public final class aws/sdk/kotlin/runtime/http/retries/AwsDefaultRetryPolicy : aws/smithy/kotlin/runtime/retries/policy/StandardRetryPolicy {
-	public static final field INSTANCE Laws/sdk/kotlin/runtime/http/retries/AwsDefaultRetryPolicy;
+public class aws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy : aws/smithy/kotlin/runtime/retries/policy/StandardRetryPolicy {
+	public static final field Companion Laws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy$Companion;
+	public fun <init> ()V
+	protected fun evaluateSpecificExceptions (Ljava/lang/Throwable;)Laws/smithy/kotlin/runtime/retries/policy/RetryDirective;
+}
+
+public final class aws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy$Companion {
+	public final fun getDefault ()Laws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy;
 }
 
diff --git a/aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/retries/AwsDefaultRetryPolicy.kt b/aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/retries/AwsDefaultRetryPolicy.kt
diff --git a/aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy.kt b/aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/retries/AwsRetryPolicy.kt
@@ -0,0 +1,62 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package aws.sdk.kotlin.runtime.http.retries
+
+import aws.smithy.kotlin.runtime.ServiceErrorMetadata
+import aws.smithy.kotlin.runtime.ServiceException
+import aws.smithy.kotlin.runtime.http.response.HttpResponse
+import aws.smithy.kotlin.runtime.retries.policy.RetryDirective
+import aws.smithy.kotlin.runtime.retries.policy.RetryErrorType.*
+import aws.smithy.kotlin.runtime.retries.policy.StandardRetryPolicy
+
+public open class AwsRetryPolicy : StandardRetryPolicy() {
+    public companion object {
+        /**
+         * The default [aws.smithy.kotlin.runtime.retries.policy.RetryPolicy] used by AWS service clients
+         */
+        public val Default: AwsRetryPolicy = AwsRetryPolicy()
+
+        internal val knownErrorTypes = mapOf(
+            "BandwidthLimitExceeded" to Throttling,
+            "EC2ThrottledException" to Throttling,
+            "IDPCommunicationError" to Transient,
+            "LimitExceededException" to Throttling,
+            "PriorRequestNotComplete" to Throttling,
+            "ProvisionedThroughputExceededException" to Throttling,
+            "RequestLimitExceeded" to Throttling,
+            "RequestThrottled" to Throttling,
+            "RequestThrottledException" to Throttling,
+            "RequestTimeout" to Transient,
+            "RequestTimeoutException" to Transient,
+            "SlowDown" to Throttling,
+            "ThrottledException" to Throttling,
+            "Throttling" to Throttling,
+            "ThrottlingException" to Throttling,
+            "TooManyRequestsException" to Throttling,
+            "TransactionInProgressException" to Throttling,
+        )
+
+        internal val knownStatusCodes = mapOf(
+            500 to Transient,
+            502 to Transient,
+            503 to Transient,
+            504 to Transient,
+        )
+    }
+
+    override fun evaluateSpecificExceptions(ex: Throwable): RetryDirective? = when (ex) {
+        is ServiceException -> evaluateServiceException(ex)
+        else -> null
+    }
+
+    private fun evaluateServiceException(ex: ServiceException): RetryDirective? = with(ex.sdkErrorMetadata) {
+        (knownErrorTypes[errorCode] ?: knownStatusCodes[statusCode])
+            ?.let { RetryDirective.RetryError(it) }
+    }
+
+    private val ServiceErrorMetadata.statusCode: Int?
+        get() = (protocolResponse as? HttpResponse)?.status?.value
+}
diff --git a/...http/retries/AwsDefaultRetryPolicyTest.kt → ...untime/http/retries/AwsRetryPolicyTest.kt b/...http/retries/AwsDefaultRetryPolicyTest.kt → ...untime/http/retries/AwsRetryPolicyTest.kt
@@ -14,25 +14,25 @@ import aws.smithy.kotlin.runtime.retries.policy.RetryDirective
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
-class AwsDefaultRetryPolicyTest {
+class AwsRetryPolicyTest {
     @Test
     fun testErrorsByErrorCode() {
-        AwsDefaultRetryPolicy.knownErrorTypes.forEach { (errorCode, errorType) ->
+        AwsRetryPolicy.knownErrorTypes.forEach { (errorCode, errorType) ->
             val ex = ServiceException()
             ex.sdkErrorMetadata.attributes[ServiceErrorMetadata.ErrorCode] = errorCode
-            val result = AwsDefaultRetryPolicy.evaluate(Result.failure(ex))
+            val result = AwsRetryPolicy.Default.evaluate(Result.failure(ex))
             assertEquals(RetryDirective.RetryError(errorType), result)
         }
     }
 
     @Test
     fun testErrorsByStatusCode() {
-        AwsDefaultRetryPolicy.knownStatusCodes.forEach { (statusCode, errorType) ->
+        AwsRetryPolicy.knownStatusCodes.forEach { (statusCode, errorType) ->
             val modeledStatusCode = HttpStatusCode.fromValue(statusCode)
             val response = HttpResponse(modeledStatusCode, Headers.Empty, HttpBody.Empty)
             val ex = ServiceException()
             ex.sdkErrorMetadata.attributes[ServiceErrorMetadata.ProtocolResponse] = response
-            val result = AwsDefaultRetryPolicy.evaluate(Result.failure(ex))
+            val result = AwsRetryPolicy.Default.evaluate(Result.failure(ex))
             assertEquals(RetryDirective.RetryError(errorType), result)
         }
     }

diff --git a/codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt b/codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt
@@ -63,7 +63,7 @@ object AwsRuntimeTypes {
         }
 
         object Retries {
-            val AwsDefaultRetryPolicy = symbol("AwsDefaultRetryPolicy", "retries")
+            val AwsRetryPolicy = symbol("AwsRetryPolicy", "retries")
         }
 
         object Middleware {

diff --git a/...-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsServiceConfigIntegration.kt b/...-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsServiceConfigIntegration.kt
@@ -121,8 +121,8 @@ class AwsServiceConfigIntegration : KotlinIntegration {
             .RetryPolicy
             .toBuilder()
             .apply {
-                propertyType = ConfigPropertyType.RequiredWithDefault("AwsDefaultRetryPolicy")
-                additionalImports = listOf(AwsRuntimeTypes.Http.Retries.AwsDefaultRetryPolicy)
+                propertyType = ConfigPropertyType.RequiredWithDefault("AwsRetryPolicy.Default")
+                additionalImports = listOf(AwsRuntimeTypes.Http.Retries.AwsRetryPolicy)
             }
             .build()