fix: include more info for retry token bucket capacity errors

.changes/3456d00f-1b29-4d88-ab02-045db1b1ebce.json

@@ -0,0 +1,8 @@
+{
+    "id": "3456d00f-1b29-4d88-ab02-045db1b1ebce",
+    "type": "bugfix",
+    "description": "Include more information when retry strategy halts early due to token bucket capacity errors",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1321"
+    ]
+}

runtime/runtime-core/api/runtime-core.api

@@ -1,3 +1,13 @@
+public final class aws/smithy/kotlin/runtime/ClientErrorContext {
+	public fun <init> (Ljava/lang/String;Ljava/lang/String;)V
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getFormatted ()Ljava/lang/String;
+	public final fun getKey ()Ljava/lang/String;
+	public final fun getValue ()Ljava/lang/String;
+	public fun hashCode ()I
+	public fun toString ()Ljava/lang/String;
+}
+
 public class aws/smithy/kotlin/runtime/ClientException : aws/smithy/kotlin/runtime/SdkBaseException {
 	public fun <init> ()V
 	public fun <init> (Ljava/lang/String;)V
@@ -9,11 +19,13 @@ public class aws/smithy/kotlin/runtime/ErrorMetadata {
 	public static final field Companion Laws/smithy/kotlin/runtime/ErrorMetadata$Companion;
 	public fun <init> ()V
 	public final fun getAttributes ()Laws/smithy/kotlin/runtime/collections/MutableAttributes;
+	public final fun getClientContext ()Ljava/util/List;
 	public final fun isRetryable ()Z
 	public final fun isThrottling ()Z
 }
 
 public final class aws/smithy/kotlin/runtime/ErrorMetadata$Companion {
+	public final fun getClientContext ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
 	public final fun getRetryable ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
 	public final fun getThrottlingError ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
 }
@@ -62,7 +74,6 @@ public class aws/smithy/kotlin/runtime/ServiceException : aws/smithy/kotlin/runt
 	public fun <init> (Ljava/lang/String;)V
 	public fun <init> (Ljava/lang/String;Ljava/lang/Throwable;)V
 	public fun <init> (Ljava/lang/Throwable;)V
-	protected fun getDisplayMetadata ()Ljava/util/List;
 	public fun getMessage ()Ljava/lang/String;
 	public synthetic fun getSdkErrorMetadata ()Laws/smithy/kotlin/runtime/ErrorMetadata;
 	public fun getSdkErrorMetadata ()Laws/smithy/kotlin/runtime/ServiceErrorMetadata;
@@ -134,6 +145,7 @@ public final class aws/smithy/kotlin/runtime/collections/AttributesBuilder {
 }
 
 public final class aws/smithy/kotlin/runtime/collections/AttributesKt {
+	public static final fun appendValue (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;Ljava/lang/Object;)V
 	public static final fun attributesOf (Lkotlin/jvm/functions/Function1;)Laws/smithy/kotlin/runtime/collections/Attributes;
 	public static final fun emptyAttributes ()Laws/smithy/kotlin/runtime/collections/Attributes;
 	public static final fun get (Laws/smithy/kotlin/runtime/collections/Attributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;)Ljava/lang/Object;
@@ -1685,6 +1697,7 @@ public abstract class aws/smithy/kotlin/runtime/retries/RetryException : aws/smi
 	public final fun getAttempts ()I
 	public final fun getLastException ()Ljava/lang/Throwable;
 	public final fun getLastResponse ()Ljava/lang/Object;
+	public fun toString ()Ljava/lang/String;
 }
 
 public final class aws/smithy/kotlin/runtime/retries/RetryFailureException : aws/smithy/kotlin/runtime/retries/RetryException {

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/Exceptions.kt

@@ -8,6 +8,41 @@ import aws.smithy.kotlin.runtime.collections.AttributeKey
 import aws.smithy.kotlin.runtime.collections.MutableAttributes
 import aws.smithy.kotlin.runtime.collections.mutableAttributes
 
+/**
+ * Describes additional context about an error which may be useful in client-side debugging. This information will be
+ * included in exception messages. This contrasts with [ErrorMetadata] which is not _necessarily_ included in messages
+ * and not _necessarily_ client-related.
+ * @param key A header or key for the information
+ * @param value A value for the information
+ */
+public class ClientErrorContext(public val key: String, public val value: String) {
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other == null || this::class != other::class) return false
+
+        other as ClientErrorContext
+
+        if (key != other.key) return false
+        if (value != other.value) return false
+
+        return true
+    }
+
+    /**
+     * Gets a formatted representation of this error context suitable for inclusion in a message. This format is
+     * generally `"$key: $value"`.
+     */
+    public val formatted: String = "$key: $value"
+
+    override fun hashCode(): Int {
+        var result = key.hashCode()
+        result = 31 * result + value.hashCode()
+        return result
+    }
+
+    override fun toString(): String = "ClientErrorContext(key='$key', value='$value')"
+}
+
 /**
  * Additional metadata about an error
  */
@@ -16,6 +51,12 @@ public open class ErrorMetadata {
     public val attributes: MutableAttributes = mutableAttributes()
 
     public companion object {
+        /**
+         * Set if there are additional context elements about the error
+         */
+        public val ClientContext: AttributeKey<List<ClientErrorContext>> =
+            AttributeKey("aws.smithy.kotlin#ClientContext")
+
         /**
          * Set if an error is retryable
          */
@@ -32,6 +73,9 @@ public open class ErrorMetadata {
 
     public val isThrottling: Boolean
         get() = attributes.getOrNull(ThrottlingError) ?: false
+
+    public val clientContext: List<ClientErrorContext>
+        get() = attributes.getOrNull(ClientContext).orEmpty()
 }
 
 /**
@@ -156,7 +200,7 @@ public open class ServiceException : SdkBaseException {
 
     public constructor(cause: Throwable?) : super(cause)
 
-    protected open val displayMetadata: List<String>
+    private val displayMetadata: List<String>
         get() = buildList {
             val serviceProvidedMessage = super.message ?: sdkErrorMetadata.errorMessage
             if (serviceProvidedMessage == null) {
@@ -166,7 +210,10 @@ public open class ServiceException : SdkBaseException {
             } else {
                 add(serviceProvidedMessage)
             }
+
             sdkErrorMetadata.requestId?.let { add("Request ID: $it") }
+
+            sdkErrorMetadata.clientContext.mapTo(this@buildList) { it.formatted }
         }
 
     override val message: String

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/collections/Attributes.kt

@@ -115,6 +115,17 @@ public fun MutableAttributes.merge(other: Attributes) {
     }
 }
 
+/**
+ * Appends a value to a list-typed attribute. If the attribute does not exist, it will be created.
+ * @param key The key for the attribute
+ * @param element The element to append to the existing (or new) list value of the attribute
+ */
+public fun <E> MutableAttributes.appendValue(key: AttributeKey<List<E>>, element: E) {
+    val existingList = getOrNull(key).orEmpty()
+    val newList = existingList + element
+    set(key, newList)
+}
+
 private class AttributesImpl constructor(seed: Attributes) : MutableAttributes {
     private val map: MutableMap<AttributeKey<*>, Any> = mutableMapOf()
     constructor() : this(emptyAttributes())

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/retries/Exceptions.kt

@@ -22,7 +22,28 @@ public sealed class RetryException(
     public val attempts: Int,
     public val lastResponse: Any?,
     public val lastException: Throwable?,
-) : ClientException(message, cause)
+) : ClientException(message, cause) {
+    override fun toString(): String = buildString {
+        append(this@RetryException::class.simpleName)
+        append("(")
+
+        append("message=")
+        append(message)
+
+        append(",attempts=")
+        append(attempts)
+
+        if (lastException != null) {
+            append(",lastException=")
+            append(lastException)
+        } else if (lastResponse != null) {
+            append(",lastResponse=")
+            append(lastResponse)
+        }
+
+        append(")")
+    }
+}
 
 /**
  * Indicates that retrying has failed because too many attempts have completed unsuccessfully.

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/retries/StandardRetryStrategy.kt

@@ -5,6 +5,10 @@
 
 package aws.smithy.kotlin.runtime.retries
 
+import aws.smithy.kotlin.runtime.ClientErrorContext
+import aws.smithy.kotlin.runtime.ErrorMetadata
+import aws.smithy.kotlin.runtime.ServiceException
+import aws.smithy.kotlin.runtime.collections.appendValue
 import aws.smithy.kotlin.runtime.retries.delay.*
 import aws.smithy.kotlin.runtime.retries.policy.RetryDirective
 import aws.smithy.kotlin.runtime.retries.policy.RetryPolicy
@@ -129,17 +133,35 @@ public open class StandardRetryStrategy(override val config: Config = Config.def
         }
     }
 
-    private fun <R> throwCapacityExceeded(cause: Throwable, attempt: Int, result: Result<R>?): Nothing =
-        when (val ex = result?.exceptionOrNull()) {
+    private fun <R> throwCapacityExceeded(
+        cause: RetryCapacityExceededException,
+        attempt: Int,
+        result: Result<R>?,
+    ): Nothing {
+        val capacityMessage = buildString {
+            append("Insufficient client capacity to attempt retry, halting on attempt ")
+            append(attempt)
+            append(" of ")
+            append(config.maxAttempts)
+        }
+
+        throw when (val retryableException = result?.exceptionOrNull()) {
             null -> throw TooManyAttemptsException(
-                cause.message!!,
+                capacityMessage,
                 cause,
                 attempt,
                 result?.getOrNull(),
                 result?.exceptionOrNull(),
             )
-            else -> throw ex
+
+            is ServiceException -> retryableException.apply {
+                val addCtx = ClientErrorContext("Early retry termination", capacityMessage)
+                sdkErrorMetadata.attributes.appendValue(ErrorMetadata.ClientContext, addCtx)
+            }
+
+            else -> retryableException
         }
+    }
 
     /**
      * Handles the termination of the retry loop because of a non-retryable failure by throwing a

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/retries/delay/StandardRetryTokenBucket.kt

@@ -60,7 +60,7 @@ public class StandardRetryTokenBucket internal constructor(
             capacity -= size
         } else {
             if (config.useCircuitBreakerMode) {
-                throw RetryCapacityExceededException("Insufficient capacity to attempt another retry")
+                throw RetryCapacityExceededException("Insufficient capacity to attempt retry")
             }
 
             val extraRequiredCapacity = size - capacity

runtime/runtime-core/common/test/aws/smithy/kotlin/runtime/ExceptionsTest.kt

@@ -5,9 +5,14 @@
 package aws.smithy.kotlin.runtime
 
 import aws.smithy.kotlin.runtime.collections.MutableAttributes
+import aws.smithy.kotlin.runtime.collections.appendValue
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
+private const val CTX_KEY_1 = "Color"
+private const val CTX_VALUE_1 = "blue"
+private const val CTX_KEY_2 = "Shape"
+private const val CTX_VALUE_2 = "square"
 private const val ERROR_CODE = "ErrorWithNoMessage"
 private const val METADATA_MESSAGE = "This is a message included in metadata but not the regular response"
 private const val PROTOCOL_RESPONSE_SUMMARY = "HTTP 418 I'm a teapot"
@@ -104,6 +109,35 @@ class ExceptionsTest {
             e.message,
         )
     }
+
+    @Test
+    fun testNoMessageWithClientContext() {
+        val e = FooServiceException {
+            appendValue(ErrorMetadata.ClientContext, ClientErrorContext(CTX_KEY_1, CTX_VALUE_1))
+            appendValue(ErrorMetadata.ClientContext, ClientErrorContext(CTX_KEY_2, CTX_VALUE_2))
+        }
+        assertEquals(
+            buildList {
+                add("Error type: Unknown")
+                add("Protocol response: (empty response)")
+                add("$CTX_KEY_1: $CTX_VALUE_1")
+                add("$CTX_KEY_2: $CTX_VALUE_2")
+            }.joinToString(),
+            e.message,
+        )
+    }
+
+    @Test
+    fun testMessageWithClientContext() {
+        val e = FooServiceException(SERVICE_MESSAGE) {
+            appendValue(ErrorMetadata.ClientContext, ClientErrorContext(CTX_KEY_1, CTX_VALUE_1))
+            appendValue(ErrorMetadata.ClientContext, ClientErrorContext(CTX_KEY_2, CTX_VALUE_2))
+        }
+        assertEquals(
+            "$SERVICE_MESSAGE, $CTX_KEY_1: $CTX_VALUE_1, $CTX_KEY_2: $CTX_VALUE_2",
+            e.message,
+        )
+    }
 }
 
 private class FooServiceException(

runtime/runtime-core/jvm/test/aws/smithy/kotlin/runtime/retries/impl/StandardRetryIntegrationTest.kt

@@ -5,8 +5,8 @@
 
 package aws.smithy.kotlin.runtime.retries.impl
 
+import aws.smithy.kotlin.runtime.ServiceException
 import aws.smithy.kotlin.runtime.retries.StandardRetryStrategy
-import aws.smithy.kotlin.runtime.retries.TooManyAttemptsException
 import aws.smithy.kotlin.runtime.retries.delay.StandardRetryTokenBucket
 import aws.smithy.kotlin.runtime.retries.getOrThrow
 import aws.smithy.kotlin.runtime.retries.policy.RetryDirective
@@ -15,6 +15,7 @@ import aws.smithy.kotlin.runtime.retries.policy.RetryPolicy
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.currentTime
 import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.assertThrows
 import kotlin.test.*
 import kotlin.time.Duration.Companion.milliseconds
 
@@ -38,7 +39,7 @@ class StandardRetryIntegrationTest {
 
             val block = object {
                 var index = 0
-                suspend fun doIt() = tc.responses[index++].response.statusCode
+                suspend fun doIt() = tc.responses[index++].response.getOrThrow()
             }::doIt
 
             val startTimeMs = currentTime
@@ -47,9 +48,29 @@ class StandardRetryIntegrationTest {
 
             val finalState = tc.responses.last().expected
             when (finalState.outcome) {
-                TestOutcome.Success -> assertEquals(200, result.getOrNull()?.getOrThrow(), "Unexpected outcome for $name")
-                TestOutcome.MaxAttemptsExceeded -> assertIs<TooManyAttemptsException>(result.exceptionOrNull())
-                TestOutcome.RetryQuotaExceeded -> assertIs<TooManyAttemptsException>(result.exceptionOrNull())
+                TestOutcome.Success ->
+                    assertEquals(Ok, result.getOrThrow().getOrThrow(), "Unexpected outcome for $name")
+
+                TestOutcome.MaxAttemptsExceeded -> {
+                    val e = assertThrows<HttpCodeException>("Expected exception for $name") {
+                        result.getOrThrow()
+                    }
+
+                    assertEquals(tc.responses.last().response.statusCode, e.code, "Unexpected error code for $name")
+                }
+
+                TestOutcome.RetryQuotaExceeded -> {
+                    val e = assertThrows<HttpCodeException>("Expected exception for $name") {
+                        result.getOrThrow()
+                    }
+
+                    assertEquals(tc.responses.last().response.statusCode, e.code, "Unexpected error code for $name")
+
+                    assertTrue("Expected retry capacity message in exception for $name") {
+                        "Insufficient client capacity to attempt retry" in e.message
+                    }
+                }
+
                 else -> fail("Unexpected outcome for $name: ${finalState.outcome}")
             }
 
@@ -72,10 +93,19 @@ class StandardRetryIntegrationTest {
     }
 }
 
-object IntegrationTestPolicy : RetryPolicy<Int> {
-    override fun evaluate(result: Result<Int>): RetryDirective = when (val code = result.getOrNull()!!) {
-        200 -> RetryDirective.TerminateAndSucceed
-        500, 502 -> RetryDirective.RetryError(RetryErrorType.ServerSide)
-        else -> fail("Unexpected status code: $code")
+object IntegrationTestPolicy : RetryPolicy<Ok> {
+    override fun evaluate(result: Result<Ok>): RetryDirective = when {
+        result.isSuccess -> RetryDirective.TerminateAndSucceed
+        result.isFailure -> RetryDirective.RetryError(RetryErrorType.ServerSide)
+        else -> fail("Unexpected result condition")
     }
 }
+
+data object Ok
+
+class HttpCodeException(val code: Int) : ServiceException()
+
+fun Response.getOrThrow() = when (statusCode) {
+    200 -> Ok
+    else -> throw HttpCodeException(statusCode)
+}