Pre-allocate response body using Content-Length (aws#1565)

Updates SDK API client deserialization to pre-allocate byte slice and string response payloads. This saves us allocations when the response body is greater than 512 bytes. For large responses, like multi-megabyte Lambda Invoke responses, this can be a very significant number of allocations and slice copies because ioutil.ReadAll / io.ReadAll start with 512 bytes and follow append's growth rules: * https://cs.opensource.google/go/go/+/refs/tags/go1.17.6:src/io/io.go;drc=dc289d3dcb59f80b9e23c7e8f237628359d21d92;l=627 * https://cs.opensource.google/go/go/%20/master:src/runtime/slice.go;l=166?q=growslice S3 GetBucketPolicy benchmark: name old time/op new time/op delta GetBucketPolicy-12 12.7ms ± 2% 8.7ms ± 2% -31.60% (p=0.000 n=9+9) name old alloc/op new alloc/op delta GetBucketPolicy-12 83.4MB ± 0% 50.4MB ± 0% -39.60% (p=0.000 n=10+10) name old allocs/op new allocs/op delta GetBucketPolicy-12 256 ± 0% 216 ± 0% -15.49% (p=0.000 n=10+10) Schemas GetCodeBindingSource benchmarks: name old time/op new time/op delta GetCodeBindingSource-12 10.8ms ± 3% 7.4ms ± 4% -31.79% (p=0.000 n=10+10) name old alloc/op new alloc/op delta GetCodeBindingSource-12 70.8MB ± 0% 37.8MB ± 0% -46.64% (p=0.000 n=10+10) name old allocs/op new allocs/op delta GetCodeBindingSource-12 241 ± 0% 200 ± 0% -17.01% (p=0.000 n=10+10) * Add changelog Co-authored-by: Jason Del Ponte <[email protected]>
jrichard8 · Feb 14, 2022 · feeedfd · feeedfd
1 parent bac80ad
commit feeedfd
Show file tree

Hide file tree

Showing 25 changed files with 658 additions and 224 deletions.
diff --git a/.changelog/183e085d88ad43669b697fde5cbf7cc5.json b/.changelog/183e085d88ad43669b697fde5cbf7cc5.json
@@ -0,0 +1,37 @@
+{
+    "id": "183e085d-88ad-4366-9b69-7fde5cbf7cc5",
+    "type": "bugfix",
+    "collapse": true,
+    "description": "Updates SDK API client deserialization to pre-allocate byte slice and string response payloads, [#1565](https://github.com/aws/aws-sdk-go-v2/pull/1565). Thanks to [Tyson Mote](https://github.com/tysonmote) for submitting this PR.",
+    "modules": [
+        "internal/protocoltest/awsrestjson",
+        "internal/protocoltest/restxml",
+        "service/apigateway",
+        "service/apigatewayv2",
+        "service/appconfig",
+        "service/appconfigdata",
+        "service/appsync",
+        "service/cloudfront",
+        "service/codeartifact",
+        "service/codeguruprofiler",
+        "service/dataexchange",
+        "service/ebs",
+        "service/glacier",
+        "service/internal/benchmark",
+        "service/iotdataplane",
+        "service/kinesisvideoarchivedmedia",
+        "service/kinesisvideomedia",
+        "service/lakeformation",
+        "service/lambda",
+        "service/lexruntimeservice",
+        "service/lexruntimev2",
+        "service/location",
+        "service/medialive",
+        "service/mediastoredata",
+        "service/polly",
+        "service/s3",
+        "service/sagemakerruntime",
+        "service/schemas",
+        "service/workmailmessageflow"
+    ]
+}
diff --git a/...odegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java b/...odegen/src/main/java/software/amazon/smithy/aws/go/codegen/RestJsonProtocolGenerator.java
@@ -130,18 +130,18 @@ protected void writeMiddlewarePayloadAsDocumentSerializerDelegator(
                     writer.addUseImports(SmithyGoDependency.SMITHY_JSON);
                     writer.addUseImports(SmithyGoDependency.BYTES);
                     writer.write("""
-                                 jsonEncoder := smithyjson.NewEncoder()
-                                 if err := $L($L, jsonEncoder.Value); err != nil {
-                                     return out, metadata, &smithy.SerializationError{Err: err}
-                                 }
-                                 payload := bytes.NewReader(jsonEncoder.Bytes())""", functionName, s);
+                            jsonEncoder := smithyjson.NewEncoder()
+                            if err := $L($L, jsonEncoder.Value); err != nil {
+                                return out, metadata, &smithy.SerializationError{Err: err}
+                            }
+                            payload := bytes.NewReader(jsonEncoder.Bytes())""", functionName, s);
                     writeSetStream(writer, "payload");
                     if (payloadShape.getType() == ShapeType.STRUCTURE) {
                         writer.openBlock("} else {", "", () -> {
                             writer.write("""
-                                         jsonEncoder := smithyjson.NewEncoder()
-                                         jsonEncoder.Value.Object().Close()
-                                         payload := bytes.NewReader(jsonEncoder.Bytes())""");
+                                    jsonEncoder := smithyjson.NewEncoder()
+                                    jsonEncoder.Value.Object().Close()
+                                    payload := bytes.NewReader(jsonEncoder.Bytes())""");
                             writeSetStream(writer, "payload");
                         });
                     }
@@ -186,7 +186,8 @@ protected void writeMiddlewareDocumentSerializerDelegator(
         writer.write("");
 
         Shape inputShape = ProtocolUtils.expectInput(context.getModel(), operation);
-        String functionName = ProtocolGenerator.getDocumentSerializerFunctionName(inputShape, context.getService(), getProtocolName());
+        String functionName = ProtocolGenerator.getDocumentSerializerFunctionName(inputShape, context.getService(),
+                getProtocolName());
 
         writer.addUseImports(SmithyGoDependency.SMITHY_JSON);
         writer.write("jsonEncoder := smithyjson.NewEncoder()");
@@ -238,7 +239,8 @@ protected void writeMiddlewareDocumentDeserializerDelegator(
 
             // if target shape is of type String or type Blob, then delegate deserializers for explicit payload shapes
             if (payloadShape.isStringShape() || payloadShape.isBlobShape()) {
-                writeMiddlewarePayloadBindingDeserializerDelegator(writer, context.getService(), targetShape);
+                writeMiddlewarePayloadBindingDeserializerDelegator(writer, context.getService(), targetShape,
+                        payloadShape);
                 return;
             }
             // for other payload target types we should deserialize using the appropriate document deserializer
@@ -266,10 +268,17 @@ protected void writeMiddlewareDocumentDeserializerDelegator(
 
     // Writes middleware that delegates to deserializers for shapes that have explicit payload.
     private void writeMiddlewarePayloadBindingDeserializerDelegator(
-            GoWriter writer, ServiceShape service, Shape shape
+            GoWriter writer, ServiceShape service, Shape outputShape, Shape payloadShape
     ) {
-        String deserFuncName = ProtocolGenerator.getDocumentDeserializerFunctionName(shape, service, getProtocolName());
-        writer.write("err = $L(output, response.Body)", deserFuncName);
+        String deserFuncName = ProtocolGenerator.getDocumentDeserializerFunctionName(outputShape, service,
+                getProtocolName());
+
+        String contentLengthParam = "";
+        if (!payloadShape.hasTrait(StreamingTrait.class)) {
+            contentLengthParam = "response.ContentLength";
+        }
+
+        writer.write("err = $L(output, response.Body, $L)", deserFuncName, contentLengthParam);
         writer.openBlock("if err != nil {", "}", () -> {
             writer.addUseImports(SmithyGoDependency.SMITHY);
             writer.write(String.format("return out, metadata, &smithy.DeserializationError{Err:%s}",
@@ -395,7 +404,8 @@ private void writePayloadBindingDeserializer(
         var writer = context.getWriter().get();
         var symbolProvider = context.getSymbolProvider();
         var shapeSymbol = symbolProvider.toSymbol(shape);
-        var funcName = ProtocolGenerator.getDocumentDeserializerFunctionName(shape, context.getService(), getProtocolName());
+        var funcName = ProtocolGenerator.getDocumentDeserializerFunctionName(shape, context.getService(),
+                getProtocolName());
 
         for (var memberShape : new TreeSet<>(shape.members())) {
             if (!filterMemberShapes.test(memberShape)) {
@@ -405,26 +415,40 @@ private void writePayloadBindingDeserializer(
             var memberName = symbolProvider.toMemberName(memberShape);
             var targetShape = context.getModel().expectShape(memberShape.getTarget());
             if (targetShape.isStringShape() || targetShape.isBlobShape()) {
-                writer.openBlock("func $L(v $P, body io.ReadCloser) error {", "}",
-                        funcName, shapeSymbol, () -> {
+
+                String contentLengthParam = "";
+                if (!targetShape.hasTrait(StreamingTrait.class)) {
+                    contentLengthParam = "contentLength int64";
+                }
+
+                writer.openBlock("func $L(v $P, body io.ReadCloser, $L) error {", "}",
+                        funcName, shapeSymbol, contentLengthParam, () -> {
                             writer.openBlock("if v == nil {", "}", () -> {
                                 writer.write("return fmt.Errorf(\"unsupported deserialization of nil %T\", v)");
                             });
                             writer.write("");
 
                             if (!targetShape.hasTrait(StreamingTrait.class)) {
-                                writer.addUseImports(SmithyGoDependency.IOUTIL);
-                                writer.write("bs, err := ioutil.ReadAll(body)");
+                                writer.addUseImports(SmithyGoDependency.BYTES);
+                                writer.write("var buf bytes.Buffer");
+                                writer.openBlock("if contentLength > 0 {", "", () -> {
+                                    writer.write("buf.Grow(int(contentLength))");
+                                    writer.openBlock("} else {", "}", () -> {
+                                        writer.write("buf.Grow(512)");
+                                    });
+                                });
+                                writer.write("_, err := buf.ReadFrom(body)");
                                 writer.write("if err != nil { return err }");
-                                writer.openBlock("if len(bs) > 0 {", "}", () -> {
+                                writer.openBlock("if buf.Len() > 0 {", "}", () -> {
                                     if (targetShape.isBlobShape()) {
-                                        writer.write("v.$L = bs", memberName);
+                                        writer.write("v.$L = buf.Bytes()", memberName);
                                     } else { // string
                                         writer.addUseImports(SmithyGoDependency.SMITHY_PTR);
                                         if (targetShape.hasTrait(EnumTrait.class)) {
-                                            writer.write("v.$L = $T(bs)", memberName, symbolProvider.toSymbol(targetShape));
+                                            writer.write("v.$L = $T(buf.Bytes())", memberName,
+                                                    symbolProvider.toSymbol(targetShape));
                                         } else {
-                                            writer.write("v.$L = ptr.String(string(bs))", memberName);
+                                            writer.write("v.$L = ptr.String(buf.String())", memberName);
                                         }
                                     }
                                 });
@@ -511,7 +535,7 @@ protected void generateEventStreamSerializers(
 
             var hasBindings = targetShape.members().stream()
                     .filter(ms -> ms.getTrait(EventHeaderTrait.class).isPresent()
-                                  || ms.getTrait(EventPayloadTrait.class).isPresent())
+                            || ms.getTrait(EventPayloadTrait.class).isPresent())
                     .findAny();
             if (hasBindings.isPresent()) {
                 var payload = targetShape.members().stream()
@@ -571,13 +595,13 @@ protected void generateEventStreamDeserializers(
                                     payloadTarget, ctx.getService(), ctx.getProtocolName());
                             var ctxWriter = ctx.getWriter().get();
                             ctxWriter.openBlock("if err := $L(&$L, shape); err != nil {", "}", functionName, operand,
-                                            () -> handleDecodeError(ctxWriter))
+                                    () -> handleDecodeError(ctxWriter))
                                     .write("return nil");
                         });
 
                 var hasBindings = targetShape.members().stream()
                         .filter(ms -> ms.getTrait(EventHeaderTrait.class).isPresent()
-                                      || ms.getTrait(EventPayloadTrait.class).isPresent())
+                                || ms.getTrait(EventPayloadTrait.class).isPresent())
                         .findAny();
                 if (hasBindings.isPresent()) {
                     var payload = targetShape.members().stream()