Skip to content

Commit 998af09

Browse files
landonxjameslandonxjames
authored
Fix S3 ListParts pagination infinite loop (#3679)
## Motivation and Context <!--- Why is this change required? What problem does it solve? --> <!--- If it fixes an open issue, please link to the issue here --> The paginator for the S3 `ListParts` operation could loop forever: awslabs/aws-sdk-rust#1143. This is because most paginated operations use an empty next token as an indication that pagination is exhausted. `ListParts` instead sets the final next token as `0` causing the pagination to loop back to the first page and loop forever. Instead of an empty next token `ListParts` uses `IsTruncated = false` to indicate that pagination has been exhausted. ## Description <!--- Describe your changes in detail --> * Added a new trait `isTruncatedPaginatorTrait` * Add that trait to the S3 `ListPartsOutput` shape * Use the presence of that trait to vary the logic setting the `is_empty` value in the paginator. * If the trait is absent it looks for an empty next token as always * if the trait is present the value is set based on the value of the response's `is_truncated` field ## Testing <!--- Please describe in detail how you tested your changes --> <!--- Include details of your testing environment, and the tests you ran to --> <!--- see how your change affects other areas of the code, etc. --> * Added integration test confirming that pagination terminates when a response contains `is_truncated = false` **Note:** I'm still working on turning this into a model test rather than an S3 specific integration test, but I wanted to get some feedback on the actual fix while I'm figuring that out) ## Checklist <!--- If a checkbox below is not applicable, then please DELETE it rather than leaving it unchecked --> - [x] I have updated `CHANGELOG.next.toml` if I made changes to the smithy-rs codegen or runtime crates - [x] I have updated `CHANGELOG.next.toml` if I made changes to the AWS SDK, generated SDK code, or SDK runtime crates ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._
1 parent 0cbeef3 commit 998af09

File tree

8 files changed

+309
-5
lines changed

8 files changed

+309
-5
lines changed

‎CHANGELOG.next.toml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ references = ["smithy-rs#3675"]
2929
meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
3030
author = "dastrom"
3131

32+
[[aws-sdk-rust]]
33+
message = "Fix S3 ListParts API paginator infinite loop."
34+
references = ["aws-sdk-rust#1143"]
35+
meta = { "breaking" = false, "tada" = false, "bug" = true }
36+
author = "landonxjames"
37+
3238
[[aws-smithy-runtime-api]]
3339
message = "Add conversions from smithy StatusCode to http StatusCode."
3440
references = ["smithy-rs#3637"]

‎aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.DocsRsMe
1010
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
1111
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
1212
import software.amazon.smithy.rustsdk.customize.DisabledAuthDecorator
13+
import software.amazon.smithy.rustsdk.customize.IsTruncatedPaginatorDecorator
1314
import software.amazon.smithy.rustsdk.customize.RemoveDefaultsDecorator
1415
import software.amazon.smithy.rustsdk.customize.apigateway.ApiGatewayDecorator
1516
import software.amazon.smithy.rustsdk.customize.applyDecorators
@@ -70,6 +71,7 @@ val DECORATORS: List<ClientCodegenDecorator> =
7071
S3Decorator(),
7172
S3ExpressDecorator(),
7273
S3ExtendedRequestIdDecorator(),
74+
IsTruncatedPaginatorDecorator(),
7375
),
7476
S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),
7577
STSDecorator().onlyApplyTo("com.amazonaws.sts#AWSSecurityTokenServiceV20110615"),
@@ -78,7 +80,12 @@ val DECORATORS: List<ClientCodegenDecorator> =
7880
TimestreamDecorator().onlyApplyTo("com.amazonaws.timestreamquery#Timestream_20181101"),
7981
// Only build docs-rs for linux to reduce load on docs.rs
8082
listOf(
81-
DocsRsMetadataDecorator(DocsRsMetadataSettings(targets = listOf("x86_64-unknown-linux-gnu"), allFeatures = true)),
83+
DocsRsMetadataDecorator(
84+
DocsRsMetadataSettings(
85+
targets = listOf("x86_64-unknown-linux-gnu"),
86+
allFeatures = true,
87+
),
88+
),
8289
),
8390
).flatten()
8491

‎aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/IsTruncatedPaginatorDecorator.kt

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.rustsdk.customize
7+
8+
import software.amazon.smithy.model.Model
9+
import software.amazon.smithy.model.shapes.ServiceShape
10+
import software.amazon.smithy.model.shapes.Shape
11+
import software.amazon.smithy.model.shapes.ShapeId
12+
import software.amazon.smithy.model.shapes.StructureShape
13+
import software.amazon.smithy.model.transform.ModelTransformer
14+
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
15+
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
16+
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
17+
import software.amazon.smithy.rust.codegen.core.util.letIf
18+
import java.util.logging.Logger
19+
20+
/**
21+
* Decorator for adding isTruncatedPaginator trait
22+
*/
23+
class IsTruncatedPaginatorDecorator : ClientCodegenDecorator {
24+
override val name: String = "IsTruncatedPaginatorDecorator"
25+
override val order: Byte = 0
26+
private val logger: Logger = Logger.getLogger(javaClass.name)
27+
private val operationsWithIsTruncatedPaginator = setOf(ShapeId.from("com.amazonaws.s3#ListPartsOutput"))
28+
29+
override fun transformModel(
30+
service: ServiceShape,
31+
model: Model,
32+
settings: ClientRustSettings,
33+
): Model =
34+
ModelTransformer.create().mapShapes(model) { shape ->
35+
shape.letIf(isInIsTruncatedList(shape)) {
36+
logger.info("Adding IsTruncatedPaginator trait to $it")
37+
(it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
38+
}
39+
}
40+
41+
private fun isInIsTruncatedList(shape: Shape): Boolean {
42+
return shape.isStructureShape && operationsWithIsTruncatedPaginator.contains(shape.id)
43+
}
44+
}

‎aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ class S3Decorator : ClientCodegenDecorator {
155155
)
156156
}
157157
}
158+
158159
else -> {}
159160
}
160161
}

‎aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/customize/IsTruncatedPaginatorTest.kt

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.rustsdk.customize
7+
8+
import org.junit.jupiter.api.Test
9+
import software.amazon.smithy.model.shapes.ShapeId
10+
import software.amazon.smithy.model.shapes.StructureShape
11+
import software.amazon.smithy.model.transform.ModelTransformer
12+
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
13+
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
14+
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
15+
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
16+
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
17+
import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
18+
import software.amazon.smithy.rust.codegen.core.util.letIf
19+
import software.amazon.smithy.rustsdk.AwsRuntimeType
20+
import software.amazon.smithy.rustsdk.awsSdkIntegrationTest
21+
22+
class IsTruncatedPaginatorTest {
23+
private val model =
24+
"""
25+
namespace test
26+
27+
use aws.protocols#restXml
28+
use aws.api#service
29+
use smithy.rules#endpointRuleSet
30+
31+
@restXml
32+
@service(sdkId: "fake")
33+
@endpointRuleSet({
34+
"version": "1.0",
35+
"rules": [{ "type": "endpoint", "conditions": [], "endpoint": { "url": "https://example.com" } }],
36+
"parameters": {
37+
"Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
38+
}
39+
})
40+
service TestService {
41+
operations: [PaginatedList]
42+
}
43+
44+
@readonly
45+
@optionalAuth
46+
@http(uri: "/PaginatedList", method: "POST")
47+
@paginated(inputToken: "nextToken", outputToken: "nextToken",
48+
pageSize: "maxResults", items: "items")
49+
operation PaginatedList {
50+
input: GetFoosInput,
51+
output: GetFoosOutput
52+
}
53+
54+
structure GetFoosInput {
55+
maxResults: Integer,
56+
nextToken: String
57+
}
58+
59+
structure GetFoosOutput {
60+
nextToken: String,
61+
items: StringList,
62+
isTruncated: Boolean,
63+
}
64+
65+
list StringList {
66+
member: String
67+
}
68+
""".asSmithyModel()
69+
70+
@Test
71+
fun `isTruncated paginators work`() {
72+
// Adding IsTruncated trait to the output shape
73+
val modifiedModel =
74+
ModelTransformer.create().mapShapes(model) { shape ->
75+
shape.letIf(shape.isStructureShape && shape.toShapeId() == ShapeId.from("test#GetFoosOutput")) {
76+
(it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
77+
}
78+
}
79+
80+
awsSdkIntegrationTest(modifiedModel) { context, rustCrate ->
81+
val rc = context.runtimeConfig
82+
val moduleName = context.moduleUseName()
83+
rustCrate.integrationTest("is_truncated_paginator") {
84+
rustTemplate(
85+
"""
86+
##![cfg(feature = "test-util")]
87+
88+
use $moduleName::Config;
89+
use $moduleName::Client;
90+
use #{Region};
91+
use aws_smithy_runtime::client::http::test_util::{ReplayEvent, StaticReplayClient};
92+
use aws_smithy_types::body::SdkBody;
93+
94+
fn mk_response(part_marker: u8) -> http::Response<SdkBody> {
95+
let (part_num_marker, next_num_marker, is_truncated) = if part_marker < 3 {
96+
(part_marker, part_marker + 1, true)
97+
} else {
98+
(part_marker, 0, false)
99+
};
100+
let body = format!(
101+
"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n
102+
<GetFoosOutput
103+
xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">
104+
<token>{part_num_marker}</token>
105+
<nextToken>{next_num_marker}</nextToken>
106+
<isTruncated>{is_truncated}</isTruncated>
107+
</GetFoosOutput>"
108+
);
109+
http::Response::builder().body(SdkBody::from(body)).unwrap()
110+
}
111+
112+
fn mk_request() -> http::Request<SdkBody> {
113+
http::Request::builder()
114+
.uri("https://some-test-bucket.s3.us-east-1.amazonaws.com/test.txt?part-number-marker=PartNumberMarker&uploadId=UploadId")
115+
.body(SdkBody::empty())
116+
.unwrap()
117+
}
118+
119+
##[#{tokio}::test]
120+
async fn is_truncated_pagination_does_not_loop() {
121+
let http_client = StaticReplayClient::new(vec![
122+
ReplayEvent::new(mk_request(), mk_response(0)),
123+
ReplayEvent::new(mk_request(), mk_response(1)),
124+
ReplayEvent::new(mk_request(), mk_response(2)),
125+
ReplayEvent::new(mk_request(), mk_response(3)),
126+
//The events below should never be called because the pagination should
127+
//terminate with the event above
128+
ReplayEvent::new(mk_request(), mk_response(0)),
129+
ReplayEvent::new(mk_request(), mk_response(1)),
130+
]);
131+
132+
let config = Config::builder()
133+
.region(Region::new("fake"))
134+
.http_client(http_client.clone())
135+
.with_test_defaults()
136+
.build();
137+
let client = Client::from_conf(config);
138+
139+
let list_parts_res = client
140+
.paginated_list()
141+
.max_results(1)
142+
.into_paginator()
143+
.send()
144+
.collect::<Vec<_>>()
145+
.await;
146+
147+
// Confirm that the pagination stopped calling the http client after the
148+
// first page with is_truncated = false
149+
assert_eq!(list_parts_res.len(), 4)
150+
}
151+
""",
152+
*preludeScope,
153+
"tokio" to CargoDependency.Tokio.toType(),
154+
"Region" to AwsRuntimeType.awsTypes(rc).resolve("region::Region"),
155+
)
156+
}
157+
}
158+
}
159+
}

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/PaginatorGenerator.kt

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import software.amazon.smithy.model.shapes.OperationShape
1111
import software.amazon.smithy.model.traits.IdempotencyTokenTrait
1212
import software.amazon.smithy.model.traits.PaginatedTrait
1313
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
14+
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
1415
import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
1516
import software.amazon.smithy.rust.codegen.core.rustlang.RustType
1617
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
@@ -21,8 +22,10 @@ import software.amazon.smithy.rust.codegen.core.rustlang.writable
2122
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
2223
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
2324
import software.amazon.smithy.rust.codegen.core.smithy.rustType
25+
import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticOutputTrait
2426
import software.amazon.smithy.rust.codegen.core.util.PANIC
2527
import software.amazon.smithy.rust.codegen.core.util.findMemberWithTrait
28+
import software.amazon.smithy.rust.codegen.core.util.getTrait
2629
import software.amazon.smithy.rust.codegen.core.util.hasTrait
2730
import software.amazon.smithy.rust.codegen.core.util.inputShape
2831
import software.amazon.smithy.rust.codegen.core.util.orNull
@@ -71,6 +74,13 @@ class PaginatorGenerator private constructor(
7174
private val outputType = symbolProvider.toSymbol(outputShape)
7275
private val errorType = symbolProvider.symbolForOperationError(operation)
7376

77+
private val isTruncatedPaginator =
78+
codegenContext.model.getShape(outputShape.toShapeId()).orNull().let { shape ->
79+
shape?.getTrait<SyntheticOutputTrait>()?.originalId.let { shapeId ->
80+
codegenContext.model.getShape(shapeId).orNull()?.hasTrait<IsTruncatedPaginatorTrait>() ?: false
81+
}
82+
}
83+
7484
private fun paginatorType(): RuntimeType =
7585
RuntimeType.forInlineFun(
7686
paginatorName,
@@ -89,7 +99,9 @@ class PaginatorGenerator private constructor(
8999
"Error" to errorType,
90100
"Builder" to symbolProvider.symbolForBuilder(operation.inputShape(model)),
91101
// SDK Types
92-
"HttpResponse" to RuntimeType.smithyRuntimeApiClient(runtimeConfig).resolve("client::orchestrator::HttpResponse"),
102+
"HttpResponse" to
103+
RuntimeType.smithyRuntimeApiClient(runtimeConfig)
104+
.resolve("client::orchestrator::HttpResponse"),
93105
"SdkError" to RuntimeType.sdkError(runtimeConfig),
94106
"pagination_stream" to RuntimeType.smithyAsync(runtimeConfig).resolve("future::pagination_stream"),
95107
// External Types
@@ -161,7 +173,7 @@ class PaginatorGenerator private constructor(
161173
let done = match resp {
162174
#{Ok}(ref resp) => {
163175
let new_token = #{output_token}(resp);
164-
let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true);
176+
#{is_empty_setter:W}
165177
if !is_empty && new_token == input.$inputTokenMember.as_ref() && self.stop_on_duplicate_token {
166178
true
167179
} else {
@@ -211,9 +223,37 @@ class PaginatorGenerator private constructor(
211223
"RuntimePlugins" to RuntimeType.runtimePlugins(runtimeConfig),
212224
)
213225
},
226+
"is_empty_setter" to isEmptySetter(),
214227
)
215228
}
216229

230+
/** Generate code to calculate the value of is_empty. For most paginators this
231+
* is indicated by the next token being the empty string. But for paginators
232+
* with the isTruncatedPaginator trait the next token is not necessarily empty.
233+
* (ex: for s3 ListParts the final next token is "0" when pagination is complete,
234+
* causing the paginator to go back to the first page and loop forever)
235+
* In this case we use a false value of isTruncated as the only indicator that
236+
* the pagination is exhausted.
237+
* */
238+
private fun isEmptySetter() =
239+
writable {
240+
if (isTruncatedPaginator) {
241+
rustTemplate(
242+
"""
243+
// Pagination is exhausted when `is_truncated` is false
244+
let is_empty = !resp.is_truncated.unwrap_or(false);
245+
""",
246+
)
247+
} else {
248+
rustTemplate(
249+
"""
250+
// Pagination is exhausted when the next token is an empty string
251+
let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true);
252+
""",
253+
)
254+
}
255+
}
256+
217257
/** Type of the inner item of the paginator */
218258
private fun itemType(): String {
219259
val members = paginationInfo.itemsMemberPath
@@ -280,7 +320,10 @@ class PaginatorGenerator private constructor(
280320
),
281321
"item_type" to
282322
writable {
283-
rustTemplate("#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>", *codegenScope)
323+
rustTemplate(
324+
"#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>",
325+
*codegenScope,
326+
)
284327
},
285328
*codegenScope,
286329
)

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/traits/IsTruncatedPaginatorTrait.kt

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.rust.codegen.client.smithy.traits
7+
8+
import software.amazon.smithy.model.node.Node
9+
import software.amazon.smithy.model.shapes.ShapeId
10+
import software.amazon.smithy.model.traits.AnnotationTrait
11+
12+
/**
13+
* Indicates that an operation should use the IsTruncated field for detecting the end of pagination.
14+
*/
15+
class IsTruncatedPaginatorTrait : AnnotationTrait(ID, Node.objectNode()) {
16+
companion object {
17+
val ID: ShapeId =
18+
ShapeId.from("software.amazon.smithy.rust.codegen.client.smithy.traits#isTruncatedPaginatorTrait")
19+
}
20+
}

0 commit comments

Comments
 (0)