Skip to content

Commit

Permalink
Fix S3 ListParts pagination infinite loop (#3679)
Browse files Browse the repository at this point in the history
## 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._
  • Loading branch information
landonxjames authored Jun 7, 2024
1 parent 0cbeef3 commit 998af09
Show file tree
Hide file tree
Showing 8 changed files with 309 additions and 5 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.next.toml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,12 @@ references = ["smithy-rs#3675"]
meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
author = "dastrom"

[[aws-sdk-rust]]
message = "Fix S3 ListParts API paginator infinite loop."
references = ["aws-sdk-rust#1143"]
meta = { "breaking" = false, "tada" = false, "bug" = true }
author = "landonxjames"

[[aws-smithy-runtime-api]]
message = "Add conversions from smithy StatusCode to http StatusCode."
references = ["smithy-rs#3637"]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.DocsRsMe
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
import software.amazon.smithy.rustsdk.customize.DisabledAuthDecorator
import software.amazon.smithy.rustsdk.customize.IsTruncatedPaginatorDecorator
import software.amazon.smithy.rustsdk.customize.RemoveDefaultsDecorator
import software.amazon.smithy.rustsdk.customize.apigateway.ApiGatewayDecorator
import software.amazon.smithy.rustsdk.customize.applyDecorators
Expand Down Expand Up @@ -70,6 +71,7 @@ val DECORATORS: List<ClientCodegenDecorator> =
S3Decorator(),
S3ExpressDecorator(),
S3ExtendedRequestIdDecorator(),
IsTruncatedPaginatorDecorator(),
),
S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),
STSDecorator().onlyApplyTo("com.amazonaws.sts#AWSSecurityTokenServiceV20110615"),
Expand All @@ -78,7 +80,12 @@ val DECORATORS: List<ClientCodegenDecorator> =
TimestreamDecorator().onlyApplyTo("com.amazonaws.timestreamquery#Timestream_20181101"),
// Only build docs-rs for linux to reduce load on docs.rs
listOf(
DocsRsMetadataDecorator(DocsRsMetadataSettings(targets = listOf("x86_64-unknown-linux-gnu"), allFeatures = true)),
DocsRsMetadataDecorator(
DocsRsMetadataSettings(
targets = listOf("x86_64-unknown-linux-gnu"),
allFeatures = true,
),
),
),
).flatten()

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package software.amazon.smithy.rustsdk.customize

import software.amazon.smithy.model.Model
import software.amazon.smithy.model.shapes.ServiceShape
import software.amazon.smithy.model.shapes.Shape
import software.amazon.smithy.model.shapes.ShapeId
import software.amazon.smithy.model.shapes.StructureShape
import software.amazon.smithy.model.transform.ModelTransformer
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
import software.amazon.smithy.rust.codegen.core.util.letIf
import java.util.logging.Logger

/**
* Decorator for adding isTruncatedPaginator trait
*/
class IsTruncatedPaginatorDecorator : ClientCodegenDecorator {
override val name: String = "IsTruncatedPaginatorDecorator"
override val order: Byte = 0
private val logger: Logger = Logger.getLogger(javaClass.name)
private val operationsWithIsTruncatedPaginator = setOf(ShapeId.from("com.amazonaws.s3#ListPartsOutput"))

override fun transformModel(
service: ServiceShape,
model: Model,
settings: ClientRustSettings,
): Model =
ModelTransformer.create().mapShapes(model) { shape ->
shape.letIf(isInIsTruncatedList(shape)) {
logger.info("Adding IsTruncatedPaginator trait to $it")
(it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
}
}

private fun isInIsTruncatedList(shape: Shape): Boolean {
return shape.isStructureShape && operationsWithIsTruncatedPaginator.contains(shape.id)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ class S3Decorator : ClientCodegenDecorator {
)
}
}

else -> {}
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,159 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package software.amazon.smithy.rustsdk.customize

import org.junit.jupiter.api.Test
import software.amazon.smithy.model.shapes.ShapeId
import software.amazon.smithy.model.shapes.StructureShape
import software.amazon.smithy.model.transform.ModelTransformer
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
import software.amazon.smithy.rust.codegen.core.util.letIf
import software.amazon.smithy.rustsdk.AwsRuntimeType
import software.amazon.smithy.rustsdk.awsSdkIntegrationTest

class IsTruncatedPaginatorTest {
private val model =
"""
namespace test
use aws.protocols#restXml
use aws.api#service
use smithy.rules#endpointRuleSet
@restXml
@service(sdkId: "fake")
@endpointRuleSet({
"version": "1.0",
"rules": [{ "type": "endpoint", "conditions": [], "endpoint": { "url": "https://example.com" } }],
"parameters": {
"Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
}
})
service TestService {
operations: [PaginatedList]
}
@readonly
@optionalAuth
@http(uri: "/PaginatedList", method: "POST")
@paginated(inputToken: "nextToken", outputToken: "nextToken",
pageSize: "maxResults", items: "items")
operation PaginatedList {
input: GetFoosInput,
output: GetFoosOutput
}
structure GetFoosInput {
maxResults: Integer,
nextToken: String
}
structure GetFoosOutput {
nextToken: String,
items: StringList,
isTruncated: Boolean,
}
list StringList {
member: String
}
""".asSmithyModel()

@Test
fun `isTruncated paginators work`() {
// Adding IsTruncated trait to the output shape
val modifiedModel =
ModelTransformer.create().mapShapes(model) { shape ->
shape.letIf(shape.isStructureShape && shape.toShapeId() == ShapeId.from("test#GetFoosOutput")) {
(it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
}
}

awsSdkIntegrationTest(modifiedModel) { context, rustCrate ->
val rc = context.runtimeConfig
val moduleName = context.moduleUseName()
rustCrate.integrationTest("is_truncated_paginator") {
rustTemplate(
"""
##![cfg(feature = "test-util")]
use $moduleName::Config;
use $moduleName::Client;
use #{Region};
use aws_smithy_runtime::client::http::test_util::{ReplayEvent, StaticReplayClient};
use aws_smithy_types::body::SdkBody;
fn mk_response(part_marker: u8) -> http::Response<SdkBody> {
let (part_num_marker, next_num_marker, is_truncated) = if part_marker < 3 {
(part_marker, part_marker + 1, true)
} else {
(part_marker, 0, false)
};
let body = format!(
"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n
<GetFoosOutput
xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">
<token>{part_num_marker}</token>
<nextToken>{next_num_marker}</nextToken>
<isTruncated>{is_truncated}</isTruncated>
</GetFoosOutput>"
);
http::Response::builder().body(SdkBody::from(body)).unwrap()
}
fn mk_request() -> http::Request<SdkBody> {
http::Request::builder()
.uri("https://some-test-bucket.s3.us-east-1.amazonaws.com/test.txt?part-number-marker=PartNumberMarker&uploadId=UploadId")
.body(SdkBody::empty())
.unwrap()
}
##[#{tokio}::test]
async fn is_truncated_pagination_does_not_loop() {
let http_client = StaticReplayClient::new(vec![
ReplayEvent::new(mk_request(), mk_response(0)),
ReplayEvent::new(mk_request(), mk_response(1)),
ReplayEvent::new(mk_request(), mk_response(2)),
ReplayEvent::new(mk_request(), mk_response(3)),
//The events below should never be called because the pagination should
//terminate with the event above
ReplayEvent::new(mk_request(), mk_response(0)),
ReplayEvent::new(mk_request(), mk_response(1)),
]);
let config = Config::builder()
.region(Region::new("fake"))
.http_client(http_client.clone())
.with_test_defaults()
.build();
let client = Client::from_conf(config);
let list_parts_res = client
.paginated_list()
.max_results(1)
.into_paginator()
.send()
.collect::<Vec<_>>()
.await;
// Confirm that the pagination stopped calling the http client after the
// first page with is_truncated = false
assert_eq!(list_parts_res.len(), 4)
}
""",
*preludeScope,
"tokio" to CargoDependency.Tokio.toType(),
"Region" to AwsRuntimeType.awsTypes(rc).resolve("region::Region"),
)
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.model.traits.IdempotencyTokenTrait
import software.amazon.smithy.model.traits.PaginatedTrait
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
import software.amazon.smithy.rust.codegen.core.rustlang.RustType
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
Expand All @@ -21,8 +22,10 @@ import software.amazon.smithy.rust.codegen.core.rustlang.writable
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
import software.amazon.smithy.rust.codegen.core.smithy.rustType
import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticOutputTrait
import software.amazon.smithy.rust.codegen.core.util.PANIC
import software.amazon.smithy.rust.codegen.core.util.findMemberWithTrait
import software.amazon.smithy.rust.codegen.core.util.getTrait
import software.amazon.smithy.rust.codegen.core.util.hasTrait
import software.amazon.smithy.rust.codegen.core.util.inputShape
import software.amazon.smithy.rust.codegen.core.util.orNull
Expand Down Expand Up @@ -71,6 +74,13 @@ class PaginatorGenerator private constructor(
private val outputType = symbolProvider.toSymbol(outputShape)
private val errorType = symbolProvider.symbolForOperationError(operation)

private val isTruncatedPaginator =
codegenContext.model.getShape(outputShape.toShapeId()).orNull().let { shape ->
shape?.getTrait<SyntheticOutputTrait>()?.originalId.let { shapeId ->
codegenContext.model.getShape(shapeId).orNull()?.hasTrait<IsTruncatedPaginatorTrait>() ?: false
}
}

private fun paginatorType(): RuntimeType =
RuntimeType.forInlineFun(
paginatorName,
Expand All @@ -89,7 +99,9 @@ class PaginatorGenerator private constructor(
"Error" to errorType,
"Builder" to symbolProvider.symbolForBuilder(operation.inputShape(model)),
// SDK Types
"HttpResponse" to RuntimeType.smithyRuntimeApiClient(runtimeConfig).resolve("client::orchestrator::HttpResponse"),
"HttpResponse" to
RuntimeType.smithyRuntimeApiClient(runtimeConfig)
.resolve("client::orchestrator::HttpResponse"),
"SdkError" to RuntimeType.sdkError(runtimeConfig),
"pagination_stream" to RuntimeType.smithyAsync(runtimeConfig).resolve("future::pagination_stream"),
// External Types
Expand Down Expand Up @@ -161,7 +173,7 @@ class PaginatorGenerator private constructor(
let done = match resp {
#{Ok}(ref resp) => {
let new_token = #{output_token}(resp);
let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true);
#{is_empty_setter:W}
if !is_empty && new_token == input.$inputTokenMember.as_ref() && self.stop_on_duplicate_token {
true
} else {
Expand Down Expand Up @@ -211,9 +223,37 @@ class PaginatorGenerator private constructor(
"RuntimePlugins" to RuntimeType.runtimePlugins(runtimeConfig),
)
},
"is_empty_setter" to isEmptySetter(),
)
}

/** Generate code to calculate the value of is_empty. For most paginators this
* is indicated by the next token being the empty string. But for paginators
* with the isTruncatedPaginator trait the next token is not necessarily empty.
* (ex: for s3 ListParts the final next token is "0" when pagination is complete,
* causing the paginator to go back to the first page and loop forever)
* In this case we use a false value of isTruncated as the only indicator that
* the pagination is exhausted.
* */
private fun isEmptySetter() =
writable {
if (isTruncatedPaginator) {
rustTemplate(
"""
// Pagination is exhausted when `is_truncated` is false
let is_empty = !resp.is_truncated.unwrap_or(false);
""",
)
} else {
rustTemplate(
"""
// Pagination is exhausted when the next token is an empty string
let is_empty = new_token.map(|token| token.is_empty()).unwrap_or(true);
""",
)
}
}

/** Type of the inner item of the paginator */
private fun itemType(): String {
val members = paginationInfo.itemsMemberPath
Expand Down Expand Up @@ -280,7 +320,10 @@ class PaginatorGenerator private constructor(
),
"item_type" to
writable {
rustTemplate("#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>", *codegenScope)
rustTemplate(
"#{Result}<${itemType()}, #{SdkError}<#{Error}, #{HttpResponse}>>",
*codegenScope,
)
},
*codegenScope,
)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package software.amazon.smithy.rust.codegen.client.smithy.traits

import software.amazon.smithy.model.node.Node
import software.amazon.smithy.model.shapes.ShapeId
import software.amazon.smithy.model.traits.AnnotationTrait

/**
* Indicates that an operation should use the IsTruncated field for detecting the end of pagination.
*/
class IsTruncatedPaginatorTrait : AnnotationTrait(ID, Node.objectNode()) {
companion object {
val ID: ShapeId =
ShapeId.from("software.amazon.smithy.rust.codegen.client.smithy.traits#isTruncatedPaginatorTrait")
}
}
Loading

0 comments on commit 998af09

Please sign in to comment.