[WIP] Optimize json marshaller

.changes/next-release/feature-AWSSDKforJavav2-439f346.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Optimized JSON marshalling performance for JSON RPC and REST JSON protocols."
+}

build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml

@@ -530,7 +530,10 @@
          whose NULL marshallers handle null validation. -->
     <Match>
         <Class name="software.amazon.awssdk.protocols.json.internal.marshall.JsonProtocolMarshaller"/>
-        <Method name="doMarshall"/>
+        <Or>
+            <Method name="doMarshall"/>
+            <Method name="marshallFieldViaRegistry"/>
+        </Or>
         <Bug pattern="NP_LOAD_OF_KNOWN_NULL_VALUE"/>
     </Match>
 </FindBugsFilter>

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/ExposedByteArrayOutputStream.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.json;
+
+import java.io.ByteArrayOutputStream;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * A thin subclass of {@link ByteArrayOutputStream} that exposes the internal buffer and count
+ * without copying. This allows {@link SdkJsonGenerator} to create a {@code ContentStreamProvider}
+ * that wraps the buffer directly via {@code ByteArrayInputStream(buf, 0, count)}, avoiding the
+ * contiguous copy that {@link ByteArrayOutputStream#toByteArray()} performs.
+ *
+ * <p>The write path is identical to {@code ByteArrayOutputStream} — no overhead is added.
+ * Only the final "get the bytes" step is optimized.
+ *
+ * <p>This class is not thread-safe.
+ */
+@SdkInternalApi
+final class ExposedByteArrayOutputStream extends ByteArrayOutputStream {
+
+    ExposedByteArrayOutputStream(int size) {
+        super(size);
+    }
+
+    /**
+     * Returns the internal buffer. The valid data is in {@code buf[0..count-1]}.
+     * The returned array may be larger than {@link #size()}; callers must use
+     * {@link #size()} to determine the valid range.
+     *
+     * <p><b>Warning:</b> The returned array is the live internal buffer. Do not modify it,
+     * and do not write to this stream after capturing the reference — the buffer may be
+     * replaced by a larger one on the next write if growth is needed.
+     */
+    byte[] buf() {
+        return buf;
+    }
+}

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/SdkJsonGenerator.java

@@ -15,14 +15,15 @@
 
 package software.amazon.awssdk.protocols.json;
 
-import java.io.ByteArrayOutputStream;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.time.Instant;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonGenerator;
 import software.amazon.awssdk.utils.BinaryUtils;
@@ -39,7 +40,7 @@ public class SdkJsonGenerator implements StructuredJsonGenerator {
      * prevent frequent resizings but small enough to avoid wasted allocations for small requests.
      */
     private static final int DEFAULT_BUFFER_SIZE = 1024;
-    private final ByteArrayOutputStream baos = new ByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
+    private final ExposedByteArrayOutputStream baos = new ExposedByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
     private final JsonGenerator generator;
     private final String contentType;
 
@@ -206,6 +207,16 @@ public StructuredJsonGenerator writeValue(ByteBuffer bytes) {
         return this;
     }
 
+    @Override
+    public StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
+        try {
+            generator.writeBinary(bytes);
+        } catch (IOException e) {
+            throw new JsonGenerationException(e);
+        }
+        return this;
+    }
+
     @Override
     //TODO: This date formatting is coupled to AWS's format. Should generalize it
     public StructuredJsonGenerator writeValue(Instant instant) {
@@ -277,6 +288,28 @@ public byte[] getBytes() {
         return baos.toByteArray();
     }
 
+    /**
+     * Returns the size of the generated content in bytes without copying.
+     */
+    public int contentSize() {
+        close();
+        return baos.size();
+    }
+
+    /**
+     * Returns a {@link ContentStreamProvider} that wraps the internal buffer directly,
+     * avoiding the contiguous copy that {@link #getBytes()} performs via
+     * {@code ByteArrayOutputStream.toByteArray()}. Each call to
+     * {@link ContentStreamProvider#newStream()} creates a fresh {@code ByteArrayInputStream}
+     * over the same buffer for retry safety.
+     */
+    public ContentStreamProvider contentStreamProvider() {
+        close();
+        byte[] buf = baos.buf();
+        int count = baos.size();
+        return () -> new ByteArrayInputStream(buf, 0, count);
+    }
+
     @Override
     public String getContentType() {
         return contentType;

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/StructuredJsonGenerator.java

@@ -169,6 +169,15 @@ default StructuredJsonGenerator writeValue(byte val) {
 
     StructuredJsonGenerator writeValue(ByteBuffer bytes);
 
+    /**
+     * Writes binary data directly from a byte array, avoiding the overhead of wrapping in a
+     * {@link ByteBuffer}. The default implementation wraps the array and delegates to
+     * {@link #writeValue(ByteBuffer)}.
+     */
+    default StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
+        return writeValue(ByteBuffer.wrap(bytes));
+    }
+
     StructuredJsonGenerator writeValue(Instant instant);
 
     StructuredJsonGenerator writeNumber(String number);

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/marshall/JsonProtocolMarshaller.java

@@ -22,22 +22,28 @@
 import static software.amazon.awssdk.http.Header.TRANSFER_ENCODING;
 
 import java.io.ByteArrayInputStream;
+import java.math.BigDecimal;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.Collections;
 import java.util.EnumMap;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.SdkField;
 import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.core.document.Document;
 import software.amazon.awssdk.core.protocol.MarshallLocation;
+import software.amazon.awssdk.core.protocol.MarshallingKnownType;
 import software.amazon.awssdk.core.protocol.MarshallingType;
 import software.amazon.awssdk.core.traits.PayloadTrait;
 import software.amazon.awssdk.core.traits.RequiredTrait;
 import software.amazon.awssdk.core.traits.TimestampFormatTrait;
 import software.amazon.awssdk.core.traits.TraitType;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.protocols.core.InstantToString;
 import software.amazon.awssdk.protocols.core.OperationInfo;
@@ -47,6 +53,7 @@
 import software.amazon.awssdk.protocols.json.AwsJsonProtocol;
 import software.amazon.awssdk.protocols.json.AwsJsonProtocolMetadata;
 import software.amazon.awssdk.protocols.json.BaseAwsJsonProtocolFactory;
+import software.amazon.awssdk.protocols.json.SdkJsonGenerator;
 import software.amazon.awssdk.protocols.json.StructuredJsonGenerator;
 import software.amazon.awssdk.protocols.json.internal.ProtocolFact;
 
@@ -61,6 +68,12 @@ public class JsonProtocolMarshaller implements ProtocolMarshaller<SdkHttpFullReq
 
     private static final JsonMarshallerRegistry MARSHALLER_REGISTRY = createMarshallerRegistry();
 
+    // Caches the resolved marshaller for non-PAYLOAD fields, keyed by SdkField identity.
+    // SdkField instances are static final per generated model class, so identity-based lookup is correct.
+    // ConcurrentHashMap is used for thread safety; the one-time put per SdkField is negligible.
+    private static final ConcurrentHashMap<SdkField<?>, JsonMarshaller<Object>> MARSHALLER_CACHE =
+        new ConcurrentHashMap<>();
+
     private final URI endpoint;
     private final StructuredJsonGenerator jsonGenerator;
     private final SdkHttpFullRequest.Builder request;
@@ -214,17 +227,21 @@ void doMarshall(SdkPojo pojo) {
             } else if (isExplicitPayloadMember(field)) {
                 marshallExplicitJsonPayload(field, val);
             } else if (val != null) {
-                marshallField(field, val);
+                if (field.location() == MarshallLocation.PAYLOAD) {
+                    // HOT PATH: switch-based dispatch, no registry, no interface dispatch
+                    marshallPayloadField(field, val);
+                } else {
+                    // WARM PATH: cached registry lookup + interface dispatch
+                    marshallFieldViaRegistry(field, val);
+                }
             } else if (field.location() != MarshallLocation.PAYLOAD) {
-                // Null payload fields that aren't required are no-op in the marshaller registry.
-                // We short circuit to avoid the registry lookup and dispatch overhead.
-                // Non payload locations (path, header, query) have null marshallers with
-                // different behavior, so they must still go through marshallField.
-                marshallField(field, val);
+                // Null non-payload: must go through registry (null marshallers vary by location)
+                marshallFieldViaRegistry(field, val);
             } else if (field.containsTrait(RequiredTrait.class, TraitType.REQUIRED_TRAIT)) {
                 throw new IllegalArgumentException(
                     String.format("Parameter '%s' must not be null", field.locationName()));
             }
+            // else: null payload field, not required → no-op
         }
     }
 
@@ -273,12 +290,24 @@ private SdkHttpFullRequest finishMarshalling() {
                 jsonGenerator.writeEndObject();
             }
 
-            byte[] content = jsonGenerator.getBytes();
+            if (jsonGenerator instanceof SdkJsonGenerator) {
+                // Optimized path: stream directly from chunked buffers, avoiding a single
+                // contiguous byte[] allocation that can cause G1GC humongous allocations.
+                SdkJsonGenerator sdkGenerator = (SdkJsonGenerator) jsonGenerator;
+                ContentStreamProvider contentProvider = sdkGenerator.contentStreamProvider();
+                request.contentStreamProvider(contentProvider);
+                int contentSize = sdkGenerator.contentSize();
+                if (contentSize > 0) {
+                    request.putHeader(CONTENT_LENGTH, Integer.toString(contentSize));
+                }
+            } else {
+                byte[] content = jsonGenerator.getBytes();
 
-            if (content != null) {
-                request.contentStreamProvider(() -> new ByteArrayInputStream(content));
-                if (content.length > 0) {
-                    request.putHeader(CONTENT_LENGTH, Integer.toString(content.length));
+                if (content != null) {
+                    request.contentStreamProvider(() -> new ByteArrayInputStream(content));
+                    if (content.length > 0) {
+                        request.putHeader(CONTENT_LENGTH, Integer.toString(content.length));
+                    }
                 }
             }
         }
@@ -312,6 +341,103 @@ private SdkHttpFullRequest finishMarshalling() {
         return request.build();
     }
 
+    /**
+     * Marshalls a PAYLOAD-location field using a switch on {@link MarshallingKnownType} instead of
+     * registry lookup and interface dispatch. Each case is a monomorphic call site that the JIT can inline.
+     */
+    @SuppressWarnings("unchecked")
+    private void marshallPayloadField(SdkField<?> field, Object val) {
+        MarshallingKnownType knownType = field.marshallingType().getKnownType();
+        if (knownType == null) {
+            marshallFieldViaRegistry(field, val);
+            return;
+        }
+
+        StructuredJsonGenerator gen = marshallerContext.jsonGenerator();
+        String fieldName = field.locationName();
+
+        switch (knownType) {
+            case STRING:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((String) val);
+                break;
+            case INTEGER:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((int) (Integer) val);
+                break;
+            case LONG:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((long) (Long) val);
+                break;
+            case SHORT:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((short) (Short) val);
+                break;
+            case BYTE:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((byte) (Byte) val);
+                break;
+            case FLOAT:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((float) (Float) val);
+                break;
+            case DOUBLE:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((double) (Double) val);
+                break;
+            case BIG_DECIMAL:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((BigDecimal) val);
+                break;
+            case BOOLEAN:
+                gen.writeFieldName(fieldName);
+                gen.writeValue((boolean) (Boolean) val);
+                break;
+            case INSTANT:
+                // Delegate to existing INSTANT marshaller to preserve TimestampFormatTrait handling.
+                // Note: INSTANT marshaller writes the field name itself.
+                SimpleTypeJsonMarshaller.INSTANT.marshall((Instant) val, marshallerContext,
+                    fieldName, (SdkField<Instant>) field);
+                break;
+            case SDK_BYTES:
+                gen.writeFieldName(fieldName);
+                gen.writeBinaryValue(((SdkBytes) val).asByteArrayUnsafe());
+                break;
+            case SDK_POJO:
+                SimpleTypeJsonMarshaller.SDK_POJO.marshall((SdkPojo) val, marshallerContext,
+                    fieldName, (SdkField<SdkPojo>) field);
+                break;
+            case LIST:
+                SimpleTypeJsonMarshaller.LIST.marshall((List<?>) val, marshallerContext,
+                    fieldName, (SdkField<List<?>>) field);
+                break;
+            case MAP:
+                SimpleTypeJsonMarshaller.MAP.marshall((Map<String, ?>) val, marshallerContext,
+                    fieldName, (SdkField<Map<String, ?>>) field);
+                break;
+            case DOCUMENT:
+                SimpleTypeJsonMarshaller.DOCUMENT.marshall((Document) val, marshallerContext,
+                    fieldName, (SdkField<Document>) field);
+                break;
+            default:
+                // Unknown type — fall back to registry lookup
+                marshallFieldViaRegistry(field, val);
+                break;
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private void marshallFieldViaRegistry(SdkField<?> field, Object val) {
+        if (val == null) {
+            MARSHALLER_REGISTRY.getMarshaller(field.location(), field.marshallingType(), val)
+                               .marshall(val, marshallerContext, field.locationName(), (SdkField<Object>) field);
+            return;
+        }
+        JsonMarshaller<Object> marshaller = MARSHALLER_CACHE.computeIfAbsent(field,
+            f -> MARSHALLER_REGISTRY.getMarshaller(f.location(), f.marshallingType(), val));
+        marshaller.marshall(val, marshallerContext, field.locationName(), (SdkField<Object>) field);
+    }
+
     private void marshallField(SdkField<?> field, Object val) {
         MARSHALLER_REGISTRY.getMarshaller(field.location(), field.marshallingType(), val)
                            .marshall(val, marshallerContext, field.locationName(), (SdkField<Object>) field);