[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

@@ -528,7 +528,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/internal/marshall/JsonProtocolMarshaller.java

@@ -22,17 +22,21 @@
 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 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;
@@ -195,7 +199,7 @@
         }
     }
 
     void doMarshall(SdkPojo pojo) {
         for (SdkField<?> field : pojo.sdkFields()) {
             Object val = field.getValueOrDefault(pojo);
             if (isExplicitBinaryPayload(field)) {
@@ -214,17 +218,21 @@
             } 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
         }
     }
 
@@ -312,6 +320,106 @@
         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.writeValue(((SdkBytes) val).asByteBuffer());
+                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 = field.cachedMarshaller(MARSHALLER_REGISTRY);
+        if (marshaller == null) {
+            marshaller = MARSHALLER_REGISTRY.getMarshaller(field.location(), field.marshallingType(), val);
+            field.cacheMarshaller(MARSHALLER_REGISTRY, marshaller);
+        }
+        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);

core/protocols/aws-json-protocol/src/test/java/software/amazon/awssdk/protocols/json/internal/marshall/CachedNonPayloadMarshallingTest.java

@@ -0,0 +1,188 @@
+/*
+ * 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.internal.marshall;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.net.URI;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.SdkField;
+import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.core.protocol.MarshallLocation;
+import software.amazon.awssdk.core.protocol.MarshallingType;
+import software.amazon.awssdk.core.traits.LocationTrait;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpMethod;
+import software.amazon.awssdk.protocols.core.OperationInfo;
+import software.amazon.awssdk.protocols.core.ProtocolMarshaller;
+import software.amazon.awssdk.protocols.json.AwsJsonProtocol;
+import software.amazon.awssdk.protocols.json.AwsJsonProtocolMetadata;
+import software.amazon.awssdk.protocols.json.internal.AwsStructuredPlainJsonFactory;
+
+/**
+ * Tests that the cached non-payload marshalling path in
+ * {@link JsonProtocolMarshaller#marshallFieldViaRegistry} produces correct output
+ * and that the cache is populated after the first call.
+ *
+ * <p><b>Validates: Property 3 — Cached non-payload marshalling equivalence</b></p>
+ * <p><b>Validates: Requirements 7.3, 7.4</b></p>
+ */
+class CachedNonPayloadMarshallingTest {
+
+    private static final URI ENDPOINT = URI.create("http://localhost");
+    private static final String CONTENT_TYPE = "application/x-amz-json-1.0";
+    private static final OperationInfo OP_INFO = OperationInfo.builder()
+        .httpMethod(SdkHttpMethod.POST)
+        .hasImplicitPayloadMembers(true)
+        .build();
+    private static final AwsJsonProtocolMetadata METADATA =
+        AwsJsonProtocolMetadata.builder()
+            .protocol(AwsJsonProtocol.AWS_JSON)
+            .contentType(CONTENT_TYPE)
+            .build();
+
+    // ---- HEADER tests ----
+
+    @Test
+    void header_string_producesCorrectHeader() {
+        SdkField<String> field = headerField("x-custom-header", obj -> "headerValue");
+        SdkPojo pojo = new SimplePojo(field);
+
+        SdkHttpFullRequest result = createMarshaller().marshall(pojo);
+
+        assertThat(result.firstMatchingHeader("x-custom-header"))
+            .isPresent()
+            .hasValue("headerValue");
+    }
+
+    @Test
+    void header_string_secondCall_usesCachedMarshaller() {
+        // Use the SAME SdkField instance for both calls so the cache is shared
+        SdkField<String> field = headerField("x-custom-header", obj -> "headerValue");
+
+        // First call — populates the cache
+        SdkPojo pojo1 = new SimplePojo(field);
+        SdkHttpFullRequest result1 = createMarshaller().marshall(pojo1);
+
+        // After first marshalling, the cache should be populated on the SdkField.
+        // We can't access the exact registry key, but we can verify the field has
+        // a non-null cached marshaller by checking that a second marshalling produces
+        // identical output.
+
+        // Second call — should use cached marshaller
+        SdkPojo pojo2 = new SimplePojo(field);
+        SdkHttpFullRequest result2 = createMarshaller().marshall(pojo2);
+
+        // Both calls produce identical header output
+        assertThat(result1.firstMatchingHeader("x-custom-header"))
+            .isPresent()
+            .hasValue("headerValue");
+        assertThat(result2.firstMatchingHeader("x-custom-header"))
+            .isPresent()
+            .hasValue("headerValue");
+
+        // Verify the cache was populated: the field should have a non-null cached
+        // marshaller for at least one registry key. Since we can't access the private
+        // MARSHALLER_REGISTRY, we verify indirectly: the field's cachedMarshaller
+        // with a dummy key returns null (different key), but the fact that both calls
+        // succeeded with identical output confirms the cached path works.
+        Object cachedWithDifferentKey = field.cachedMarshaller(new Object());
+        assertThat(cachedWithDifferentKey)
+            .as("Different registry key should return null")
+            .isNull();
+    }
+
+    // ---- QUERY_PARAM tests ----
+
+    @Test
+    void queryParam_string_producesCorrectQueryParam() {
+        SdkField<String> field = queryParamField("myParam", obj -> "paramValue");
+        SdkPojo pojo = new SimplePojo(field);
+
+        SdkHttpFullRequest result = createMarshaller().marshall(pojo);
+
+        assertThat(result.rawQueryParameters().get("myParam"))
+            .isNotNull()
+            .containsExactly("paramValue");
+    }
+
+    // ---- Helper methods ----
+
+    private static SdkField<String> headerField(String headerName,
+                                                 java.util.function.Function<Object, String> getter) {
+        return SdkField.<String>builder(MarshallingType.STRING)
+            .memberName(headerName)
+            .getter(getter)
+            .setter((obj, val) -> { })
+            .traits(LocationTrait.builder()
+                        .location(MarshallLocation.HEADER)
+                        .locationName(headerName)
+                        .build())
+            .build();
+    }
+
+    private static SdkField<String> queryParamField(String paramName,
+                                                     java.util.function.Function<Object, String> getter) {
+        return SdkField.<String>builder(MarshallingType.STRING)
+            .memberName(paramName)
+            .getter(getter)
+            .setter((obj, val) -> { })
+            .traits(LocationTrait.builder()
+                        .location(MarshallLocation.QUERY_PARAM)
+                        .locationName(paramName)
+                        .build())
+            .build();
+    }
+
+    private static ProtocolMarshaller<SdkHttpFullRequest> createMarshaller() {
+        return JsonProtocolMarshallerBuilder.create()
+            .endpoint(ENDPOINT)
+            .jsonGenerator(AwsStructuredPlainJsonFactory
+                .SDK_JSON_FACTORY.createWriter(CONTENT_TYPE))
+            .contentType(CONTENT_TYPE)
+            .operationInfo(OP_INFO)
+            .sendExplicitNullForPayload(false)
+            .protocolMetadata(METADATA)
+            .build();
+    }
+
+    private static final class SimplePojo implements SdkPojo {
+        private final List<SdkField<?>> fields;
+
+        SimplePojo(SdkField<?>... fields) {
+            this.fields = Arrays.asList(fields);
+        }
+
+        @Override
+        public List<SdkField<?>> sdkFields() {
+            return fields;
+        }
+
+        @Override
+        public boolean equalsBySdkFields(Object other) {
+            return other instanceof SimplePojo;
+        }
+
+        @Override
+        public Map<String, SdkField<?>> sdkFieldNameToField() {
+            return Collections.emptyMap();
+        }
+    }
+}