feat: add retryStrategy configuration option for waiters (#1183)

smithy-lang · Nov 26, 2024 · 869fb84 · 869fb84
1 parent 159e784
commit 869fb84
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 11 deletions.
diff --git a/.changes/0857b6f0-0444-479f-be6b-06ef71d482a0.json b/.changes/0857b6f0-0444-479f-be6b-06ef71d482a0.json
@@ -0,0 +1,9 @@
+{
+    "id": "0857b6f0-0444-479f-be6b-06ef71d482a0",
+    "type": "feature",
+    "description": "⚠️ **IMPORTANT**: Add `retryStrategy` configuration option for waiters",
+    "issues": [
+        "https://github.com/awslabs/aws-sdk-kotlin/issues/1431"
+    ],
+    "requiresMinorVersionBump": true
+}
diff --git a/...rc/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/WaiterGenerator.kt b/...rc/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/WaiterGenerator.kt
@@ -17,7 +17,7 @@ import java.text.DecimalFormatSymbols
  * Renders the top-level retry strategy for a waiter.
  */
 private fun KotlinWriter.renderRetryStrategy(wi: WaiterInfo, asValName: String) {
-    withBlock("val #L = #T {", "}", asValName, RuntimeTypes.Core.Retries.StandardRetryStrategy) {
+    withBlock("val #L = retryStrategy ?: #T {", "}", asValName, RuntimeTypes.Core.Retries.StandardRetryStrategy) {
         write("maxAttempts = 20")
         write("tokenBucket = #T", RuntimeTypes.Core.Retries.Delay.InfiniteTokenBucket)
         withBlock("delayProvider {", "}") {
@@ -35,18 +35,21 @@ private fun KotlinWriter.renderRetryStrategy(wi: WaiterInfo, asValName: String)
 internal fun KotlinWriter.renderWaiter(wi: WaiterInfo) {
     write("")
     wi.waiter.documentation.ifPresent(::dokka)
-    val inputParameter = if (wi.input.hasAllOptionalMembers) {
-        format("request: #1T = #1T { }", wi.inputSymbol)
+
+    val requestType = if (wi.input.hasAllOptionalMembers) {
+        format("#1T = #1T { }", wi.inputSymbol)
     } else {
-        format("request: #T", wi.inputSymbol)
+        format("#T", wi.inputSymbol)
     }
+
     withBlock(
-        "#L suspend fun #T.#L(#L): #T<#T> {",
+        "#L suspend fun #T.#L(request: #L, retryStrategy: #T? = null): #T<#T> {",
         "}",
         wi.ctx.settings.api.visibility,
         wi.serviceSymbol,
         wi.methodName,
-        inputParameter,
+        requestType,
+        RuntimeTypes.Core.Retries.RetryStrategy,
         RuntimeTypes.Core.Retries.Outcome,
         wi.outputSymbol,
     ) {

diff --git a/...in/software/amazon/smithy/kotlin/codegen/rendering/waiters/ServiceWaitersGeneratorTest.kt b/...in/software/amazon/smithy/kotlin/codegen/rendering/waiters/ServiceWaitersGeneratorTest.kt
@@ -43,7 +43,7 @@ class ServiceWaitersGeneratorTest {
             /**
              * Wait until a foo exists with optional input
              */
-            public suspend fun TestClient.waitUntilFooOptionalExists(request: DescribeFooOptionalRequest = DescribeFooOptionalRequest { }): Outcome<DescribeFooOptionalResponse> {
+            public suspend fun TestClient.waitUntilFooOptionalExists(request: DescribeFooOptionalRequest = DescribeFooOptionalRequest { }, retryStrategy: RetryStrategy? = null): Outcome<DescribeFooOptionalResponse> {
         """.trimIndent()
         val methodFooter = """
                 val policy = AcceptorRetryPolicy(request, acceptors)
@@ -59,7 +59,7 @@ class ServiceWaitersGeneratorTest {
             /**
              * Wait until a foo exists with required input
              */
-            public suspend fun TestClient.waitUntilFooRequiredExists(request: DescribeFooRequiredRequest): Outcome<DescribeFooRequiredResponse> {
+            public suspend fun TestClient.waitUntilFooRequiredExists(request: DescribeFooRequiredRequest, retryStrategy: RetryStrategy? = null): Outcome<DescribeFooRequiredResponse> {
         """.trimIndent()
         listOf(
             generateService("simple-service-with-operation-waiter.smithy"),
@@ -105,7 +105,7 @@ class ServiceWaitersGeneratorTest {
     @Test
     fun testRetryStrategy() {
         val expected = """
-            val strategy = StandardRetryStrategy {
+            val strategy = retryStrategy ?: StandardRetryStrategy {
                 maxAttempts = 20
                 tokenBucket = InfiniteTokenBucket
                 delayProvider {

diff --git a/...est/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/WaiterGeneratorTest.kt b/...est/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/WaiterGeneratorTest.kt
@@ -28,7 +28,7 @@ class WaiterGeneratorTest {
     @Test
     fun testDefaultDelays() {
         val expected = """
-            val strategy = StandardRetryStrategy {
+            val strategy = retryStrategy ?: StandardRetryStrategy {
                 maxAttempts = 20
                 tokenBucket = InfiniteTokenBucket
                 delayProvider {
@@ -45,7 +45,7 @@ class WaiterGeneratorTest {
     @Test
     fun testCustomDelays() {
         val expected = """
-            val strategy = StandardRetryStrategy {
+            val strategy = retryStrategy ?: StandardRetryStrategy {
                 maxAttempts = 20
                 tokenBucket = InfiniteTokenBucket
                 delayProvider {