Skip to content

Commit e1c9e42

Browse files
aajtoddaajtodd
authored
feat(rt): implement retries for imds (#404)
1 parent edb5d79 commit e1c9e42

File tree

4 files changed

+107
-11
lines changed

4 files changed

+107
-11
lines changed

‎aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/imds/ImdsClient.kt‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,13 @@ import aws.smithy.kotlin.runtime.client.SdkLogMode
1717
import aws.smithy.kotlin.runtime.http.*
1818
import aws.smithy.kotlin.runtime.http.engine.HttpClientEngine
1919
import aws.smithy.kotlin.runtime.http.middleware.ResolveEndpoint
20+
import aws.smithy.kotlin.runtime.http.middleware.RetryFeature
2021
import aws.smithy.kotlin.runtime.http.operation.*
2122
import aws.smithy.kotlin.runtime.http.response.HttpResponse
2223
import aws.smithy.kotlin.runtime.io.Closeable
2324
import aws.smithy.kotlin.runtime.io.middleware.Phase
2425
import aws.smithy.kotlin.runtime.logging.Logger
26+
import aws.smithy.kotlin.runtime.retries.impl.*
2527
import aws.smithy.kotlin.runtime.time.Clock
2628
import aws.smithy.kotlin.runtime.util.Platform
2729
import aws.smithy.kotlin.runtime.util.PlatformProvider
@@ -86,11 +88,18 @@ public class ImdsClient private constructor(builder: Builder) : InstanceMetadata
8688
UserAgent.create {
8789
staticMetadata = AwsUserAgentMetadata.fromEnvironment(ApiMetadata(SERVICE, "unknown"))
8890
},
91+
RetryFeature.create {
92+
val tokenBucket = StandardRetryTokenBucket(StandardRetryTokenBucketOptions.Default)
93+
val delayProvider = ExponentialBackoffWithJitter(ExponentialBackoffWithJitterOptions.Default)
94+
strategy = StandardRetryStrategy(StandardRetryStrategyOptions.Default, tokenBucket, delayProvider)
95+
policy = ImdsRetryPolicy()
96+
},
97+
// must come after retries
8998
TokenMiddleware.create {
9099
httpClient = this@ImdsClient.httpClient
91100
ttl = tokenTtl
92101
clock = this@ImdsClient.clock
93-
}
102+
},
94103
)
95104

96105
public companion object {

‎aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/imds/ImdsRetryPolicy.kt‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0.
4+
*/
5+
6+
package aws.sdk.kotlin.runtime.config.imds
7+
8+
import aws.smithy.kotlin.runtime.http.HttpStatusCode
9+
import aws.smithy.kotlin.runtime.http.category
10+
import aws.smithy.kotlin.runtime.logging.Logger
11+
import aws.smithy.kotlin.runtime.retries.RetryDirective
12+
import aws.smithy.kotlin.runtime.retries.RetryErrorType
13+
import aws.smithy.kotlin.runtime.retries.RetryPolicy
14+
15+
internal class ImdsRetryPolicy : RetryPolicy<Any?> {
16+
override fun evaluate(result: Result<Any?>): RetryDirective = when {
17+
result.isSuccess -> RetryDirective.TerminateAndSucceed
18+
else -> evaluate(result.exceptionOrNull()!!)
19+
}
20+
21+
private fun evaluate(throwable: Throwable): RetryDirective = when (throwable) {
22+
is EC2MetadataError -> {
23+
val status = HttpStatusCode.fromValue(throwable.statusCode)
24+
when {
25+
status.category() == HttpStatusCode.Category.SERVER_ERROR -> RetryDirective.RetryError(RetryErrorType.ServerSide)
26+
// 401 indicates the token has expired, this is retryable
27+
status == HttpStatusCode.Unauthorized -> RetryDirective.RetryError(RetryErrorType.ServerSide)
28+
else -> {
29+
val logger = Logger.getLogger<ImdsRetryPolicy>()
30+
logger.debug { "Non retryable IMDS error: statusCode=${throwable.statusCode}; ${throwable.message}" }
31+
RetryDirective.TerminateAndFail
32+
}
33+
}
34+
}
35+
else -> RetryDirective.TerminateAndFail
36+
}
37+
}

‎aws-runtime/aws-config/common/src/aws/sdk/kotlin/runtime/config/imds/TokenMiddleware.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ internal class TokenMiddleware(config: Config) : Feature {
5151
}
5252

5353
override fun <I, O> install(operation: SdkHttpOperation<I, O>) {
54-
operation.execution.mutate.intercept { req, next ->
54+
operation.execution.finalize.intercept { req, next ->
5555
val token = cachedToken.getOrLoad { getToken(clock, req).let { ExpiringValue(it, it.expires) } }
5656
req.subject.headers.append(X_AWS_EC2_METADATA_TOKEN, token.value.decodeToString())
5757
next.call(req)

‎aws-runtime/aws-config/common/test/aws/sdk/kotlin/runtime/config/imds/ImdsClientTest.kt‎

Lines changed: 59 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,11 @@ package aws.sdk.kotlin.runtime.config.imds
77

88
import aws.sdk.kotlin.runtime.testing.TestPlatformProvider
99
import aws.sdk.kotlin.runtime.testing.runSuspendTest
10+
import aws.smithy.kotlin.runtime.http.Headers
11+
import aws.smithy.kotlin.runtime.http.HttpBody
12+
import aws.smithy.kotlin.runtime.http.HttpStatusCode
1013
import aws.smithy.kotlin.runtime.http.operation.Endpoint
14+
import aws.smithy.kotlin.runtime.http.response.HttpResponse
1115
import aws.smithy.kotlin.runtime.httptest.buildTestConnection
1216
import aws.smithy.kotlin.runtime.time.Instant
1317
import aws.smithy.kotlin.runtime.time.ManualClock
@@ -139,24 +143,70 @@ class ImdsClientTest {
139143
connection.assertRequests()
140144
}
141145

142-
@Ignore
143146
@Test
144-
fun testRetryHttp500() {
145-
fail("not implemented yet")
147+
fun testRetryHttp500(): Unit = runSuspendTest {
148+
val connection = buildTestConnection {
149+
expect(
150+
tokenRequest("http://169.254.169.254", DEFAULT_TOKEN_TTL_SECONDS),
151+
tokenResponse(DEFAULT_TOKEN_TTL_SECONDS, "TOKEN_A")
152+
)
153+
expect(
154+
imdsRequest("http://169.254.169.254/latest/metadata", "TOKEN_A"),
155+
HttpResponse(HttpStatusCode.InternalServerError, Headers.Empty, HttpBody.Empty)
156+
)
157+
expect(
158+
imdsRequest("http://169.254.169.254/latest/metadata", "TOKEN_A"),
159+
imdsResponse("output 2")
160+
)
161+
}
162+
163+
val client = ImdsClient { engine = connection }
164+
val r1 = client.get("/latest/metadata")
165+
assertEquals("output 2", r1)
166+
connection.assertRequests()
146167
}
147168

148-
@Ignore
149169
@Test
150-
fun testRetryTokenFailure() {
170+
fun testRetryTokenFailure(): Unit = runSuspendTest {
151171
// 500 during token acquisition should be retried
152-
fail("not implemented yet")
172+
val connection = buildTestConnection {
173+
expect(
174+
tokenRequest("http://169.254.169.254", DEFAULT_TOKEN_TTL_SECONDS),
175+
HttpResponse(HttpStatusCode.InternalServerError, Headers.Empty, HttpBody.Empty)
176+
)
177+
expect(
178+
tokenRequest("http://169.254.169.254", DEFAULT_TOKEN_TTL_SECONDS),
179+
tokenResponse(DEFAULT_TOKEN_TTL_SECONDS, "TOKEN_A")
180+
)
181+
expect(
182+
imdsRequest("http://169.254.169.254/latest/metadata", "TOKEN_A"),
183+
imdsResponse("output 2")
184+
)
185+
}
186+
187+
val client = ImdsClient { engine = connection }
188+
val r1 = client.get("/latest/metadata")
189+
assertEquals("output 2", r1)
190+
connection.assertRequests()
153191
}
154192

155-
@Ignore
156193
@Test
157-
fun testNoRetryHttp403() {
194+
fun testNoRetryHttp403(): Unit = runSuspendTest {
158195
// 403 responses from IMDS during token acquisition MUST not be retried
159-
fail("not implemented yet")
196+
val connection = buildTestConnection {
197+
expect(
198+
tokenRequest("http://169.254.169.254", DEFAULT_TOKEN_TTL_SECONDS),
199+
HttpResponse(HttpStatusCode.Forbidden, Headers.Empty, HttpBody.Empty)
200+
)
201+
}
202+
203+
val client = ImdsClient { engine = connection }
204+
val ex = assertFailsWith<EC2MetadataError> {
205+
client.get("/latest/metadata")
206+
}
207+
208+
assertEquals(HttpStatusCode.Forbidden.value, ex.statusCode)
209+
connection.assertRequests()
160210
}
161211

162212
@Test

0 commit comments

Comments
 (0)