Revert "Remove fast unmarshaller dialup code (#6050)"

This reverts commit d8588e1.

Author: davidh44

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java

@@ -81,6 +81,7 @@
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.IdentityProviders;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.protocols.json.internal.unmarshall.SdkClientJsonProtocolAdvancedOption;
 import software.amazon.awssdk.regions.ServiceMetadataAdvancedOption;
 import software.amazon.awssdk.utils.AttributeMap;
 import software.amazon.awssdk.utils.CollectionUtils;
@@ -507,6 +508,13 @@ private MethodSpec finalizeServiceConfigurationMethod() {
                .addCode("    .fipsEnabled(c.get($T.FIPS_ENDPOINT_ENABLED))", AwsClientOption.class)
                .addCode("    .build());");
 
+        if (model.getMetadata().isJsonProtocol()) {
+            if (model.getCustomizationConfig().getEnableFastUnmarshaller()) {
+                builder.addStatement("builder.option($1T.ENABLE_FAST_UNMARSHALLER, true)",
+                                     SdkClientJsonProtocolAdvancedOption.class);
+            }
+        }
+
         if (hasRequestAlgorithmMember(model) || hasResponseAlgorithms(model)) {
             builder.addStatement("$T clientConfig = config", SdkClientConfiguration.class);
 

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

@@ -53,6 +53,7 @@
 import software.amazon.awssdk.protocols.json.internal.unmarshall.JsonProtocolUnmarshaller;
 import software.amazon.awssdk.protocols.json.internal.unmarshall.JsonResponseHandler;
 import software.amazon.awssdk.protocols.json.internal.unmarshall.ProtocolUnmarshallDependencies;
+import software.amazon.awssdk.protocols.json.internal.unmarshall.SdkClientJsonProtocolAdvancedOption;
 
 @SdkProtectedApi
 public abstract class BaseAwsJsonProtocolFactory {
@@ -87,7 +88,16 @@ protected BaseAwsJsonProtocolFactory(Builder<?> builder) {
         this.customErrorCodeFieldName = builder.customErrorCodeFieldName;
         this.hasAwsQueryCompatible = builder.hasAwsQueryCompatible;
         this.clientConfiguration = builder.clientConfiguration;
+        Boolean enableFastUnmarshalling = false;
+        if (clientConfiguration != null) {
+            enableFastUnmarshalling =
+                clientConfiguration.option(SdkClientJsonProtocolAdvancedOption.ENABLE_FAST_UNMARSHALLER);
+            if (enableFastUnmarshalling == null) {
+                enableFastUnmarshalling = false;
+            }
+        }
         this.protocolUnmarshaller = JsonProtocolUnmarshaller.builder()
+                                                            .enableFastUnmarshalling(enableFastUnmarshalling)
                                                             .protocolUnmarshallDependencies(
                                                                 builder.protocolUnmarshallDependencies.get())
                                                             .build();

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/unmarshall/JsonProtocolUnmarshaller.java

@@ -48,6 +48,7 @@
 import software.amazon.awssdk.protocols.json.internal.MarshallerUtil;
 import software.amazon.awssdk.protocols.json.internal.unmarshall.document.DocumentUnmarshaller;
 import software.amazon.awssdk.protocols.jsoncore.JsonNode;
+import software.amazon.awssdk.protocols.jsoncore.JsonNodeParser;
 import software.amazon.awssdk.protocols.jsoncore.JsonValueNodeFactory;
 import software.amazon.awssdk.utils.Lazy;
 import software.amazon.awssdk.utils.builder.Buildable;
@@ -58,24 +59,42 @@
  */
 @SdkInternalApi
 @ThreadSafe
-public final class JsonProtocolUnmarshaller {
+public class JsonProtocolUnmarshaller {
     private static final Lazy<DefaultProtocolUnmarshallDependencies> DEFAULT_DEPENDENCIES =
         new Lazy<>(JsonProtocolUnmarshaller::newProtocolUnmarshallDependencies);
 
     private final JsonUnmarshallerRegistry registry;
     private final JsonUnmarshallingParser unmarshallingParser;
+    private final JsonNodeParser parser;
 
     private JsonProtocolUnmarshaller(Builder builder) {
         ProtocolUnmarshallDependencies dependencies = builder.protocolUnmarshallDependencies;
         this.registry = dependencies.jsonUnmarshallerRegistry();
-        this.unmarshallingParser = JsonUnmarshallingParser.builder()
-                                                          .jsonValueNodeFactory(dependencies.nodeValueFactory())
-                                                          .jsonFactory(dependencies.jsonFactory())
-                                                          .unmarshallerRegistry(dependencies.jsonUnmarshallerRegistry())
-                                                          .defaultTimestampFormat(dependencies.timestampFormats()
-                                                                                              .get(MarshallLocation.PAYLOAD))
+        if (builder.enableFastUnmarshalling) {
+            this.unmarshallingParser = JsonUnmarshallingParser.builder()
+                                                              .jsonValueNodeFactory(dependencies.nodeValueFactory())
+                                                              .jsonFactory(dependencies.jsonFactory())
+                                                              .unmarshallerRegistry(dependencies.jsonUnmarshallerRegistry())
+                                                              .defaultTimestampFormat(dependencies.timestampFormats()
+                                                                                                  .get(MarshallLocation.PAYLOAD))
+
+                                                              .build();
+            this.parser = null;
+        } else {
+            this.unmarshallingParser = null;
+            this.parser = createParser(builder, dependencies);
+        }
+    }
 
-                                                          .build();
+    private JsonNodeParser createParser(Builder builder, ProtocolUnmarshallDependencies dependencies) {
+        if (builder.parser != null) {
+            return builder.parser;
+        }
+        return JsonNodeParser
+            .builder()
+            .jsonFactory(dependencies.jsonFactory())
+            .jsonValueNodeFactory(dependencies.nodeValueFactory())
+            .build();
     }
 
     public static DefaultProtocolUnmarshallDependencies defaultProtocolUnmarshallDependencies() {
@@ -220,6 +239,15 @@ public T unmarshall(JsonUnmarshallerContext context,
 
     public <TypeT extends SdkPojo> TypeT unmarshall(SdkPojo sdkPojo,
                                                     SdkHttpFullResponse response) throws IOException {
+        if (this.unmarshallingParser != null) {
+            return fastUnmarshall(sdkPojo, response);
+        }
+        JsonNode jsonNode = hasJsonPayload(sdkPojo, response) ? parser.parse(response.content().get()) : null;
+        return unmarshall(sdkPojo, response, jsonNode);
+    }
+
+    private <TypeT extends SdkPojo> TypeT fastUnmarshall(SdkPojo sdkPojo,
+                                                    SdkHttpFullResponse response) throws IOException {
         if (!hasJsonPayload(sdkPojo, response)) {
             return unmarshallResponse(sdkPojo, response);
         }
@@ -425,11 +453,23 @@ public static JsonUnmarshallerRegistry timestampFormatRegistryFactory(
      * Builder for {@link JsonProtocolUnmarshaller}.
      */
     public static final class Builder {
+
+        private JsonNodeParser parser;
         private ProtocolUnmarshallDependencies protocolUnmarshallDependencies;
+        private boolean enableFastUnmarshalling = false;
 
         private Builder() {
         }
 
+        /**
+         * @param parser JSON parser to use.
+         * @return This builder for method chaining.
+         */
+        public Builder parser(JsonNodeParser parser) {
+            this.parser = parser;
+            return this;
+        }
+
         /**
          * @param formats The default timestamp formats for each location in the HTTP response.
          * @return This builder for method chaining.
@@ -451,6 +491,15 @@ public Builder protocolUnmarshallDependencies(
             return this;
         }
 
+        /**
+         * @param enableFastUnmarshalling Whether to enable the fast unmarshalling codepath. Default to {@code false}.
+         * @return This builder for method chaining.
+         */
+        public Builder enableFastUnmarshalling(boolean enableFastUnmarshalling) {
+            this.enableFastUnmarshalling = enableFastUnmarshalling;
+            return this;
+        }
+
         /**
          * @return New instance of {@link JsonProtocolUnmarshaller}.
          */

test/architecture-tests/src/test/java/software/amazon/awssdk/archtests/InternalApiBoundaryTest.java

@@ -32,12 +32,14 @@
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.awscore.internal.AwsProtocolMetadata;
 import software.amazon.awssdk.awscore.internal.AwsServiceProtocol;
 import software.amazon.awssdk.core.internal.interceptor.trait.RequestCompression;
 import software.amazon.awssdk.core.internal.util.MetricUtils;
 import software.amazon.awssdk.core.internal.waiters.WaiterAttribute;
 import software.amazon.awssdk.http.auth.aws.internal.signer.util.ChecksumUtil;
+import software.amazon.awssdk.protocols.json.internal.unmarshall.SdkClientJsonProtocolAdvancedOption;
 import software.amazon.awssdk.utils.internal.EnumUtils;
 import software.amazon.awssdk.utils.internal.SystemSettingUtils;
 
@@ -55,7 +57,7 @@ public class InternalApiBoundaryTest {
     private static final Set<Class<?>> ALLOWED_INTERNAL_API_ACROSS_MODULE_SUPPRESSION = new HashSet<>(
         Arrays.asList(WaiterAttribute.class, RequestCompression.class, RequestCompression.Builder.class, EnumUtils.class,
                       AwsServiceProtocol.class, AwsProtocolMetadata.class, MetricUtils.class, SystemSettingUtils.class,
-                      ChecksumUtil.class));
+                      ChecksumUtil.class, SdkClientJsonProtocolAdvancedOption.class));
 
     @Test
     void internalApi_shouldNotUsedAcrossModule() {

test/sdk-benchmarks/src/main/java/software/amazon/awssdk/benchmark/apicall/protocol/JsonCodec.java

@@ -67,6 +67,7 @@ public SdkPojo unmarshall(AwsJsonProtocol protocol, SdkPojo pojo, byte[] bytes)
             JsonProtocolUnmarshaller unmarshaller =
                 JsonProtocolUnmarshaller
                     .builder()
+                    .enableFastUnmarshalling(true)
                     .protocolUnmarshallDependencies(behavior.protocolUnmarshallDependencies())
                     .build();
             SdkHttpFullResponse response = SdkHttpFullResponse

test/sdk-benchmarks/src/main/java/software/amazon/awssdk/benchmark/marshaller/dynamodb/V2DynamoDbAttributeValue.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.protocols.json.AwsJsonProtocol;
 import software.amazon.awssdk.protocols.json.AwsJsonProtocolFactory;
 import software.amazon.awssdk.protocols.json.JsonOperationMetadata;
+import software.amazon.awssdk.protocols.json.internal.unmarshall.SdkClientJsonProtocolAdvancedOption;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.BackupInUseException;
 import software.amazon.awssdk.services.dynamodb.model.BackupNotFoundException;
@@ -79,6 +80,7 @@ public class V2DynamoDbAttributeValue {
         .clientConfiguration(SdkClientConfiguration
                                  .builder()
                                  .option(SdkClientOption.ENDPOINT, URI.create("https://localhost"))
+                                 .option(SdkClientJsonProtocolAdvancedOption.ENABLE_FAST_UNMARSHALLER, true)
                                  .build())
         .defaultServiceExceptionSupplier(DynamoDbException::builder)
         .protocol(AwsJsonProtocol.AWS_JSON)