Set longPoll ops for SQS,SWF,SFN

Add a customization to set the long polling trait for the following
service operations:

 - SQS#ReceiveMessage
 - SFN#GetActivityTask
 - SWF#PollForActivityTask
 - SWF#PollForDecisionTask

Author: dagnir

codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/DefaultCustomizationProcessor.java

@@ -42,7 +42,8 @@ public static CodegenCustomizationProcessor getProcessorFor(
                 new S3RemoveBucketFromUriProcessor(),
                 new S3ControlRemoveAccountIdHostPrefixProcessor(),
                 new ExplicitStringPayloadQueryProtocolProcessor(),
-                new LowercaseShapeValidatorProcessor()
+                new LowercaseShapeValidatorProcessor(),
+                new LongPollingOperationProcessor()
                 );
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/customization/processors/LongPollingOperationProcessor.java

@@ -0,0 +1,71 @@
+/*
+ * 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.codegen.customization.processors;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.codegen.customization.CodegenCustomizationProcessor;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+import software.amazon.awssdk.codegen.model.service.ServiceModel;
+
+// TODO: Remove this when the long polling trait is formalized as a c2j trait.
+/**
+ * Marks specific service operations as having the long polling trait.
+ */
+public class LongPollingOperationProcessor implements CodegenCustomizationProcessor {
+    private static final Logger log = LoggerFactory.getLogger(LongPollingOperationProcessor.class);
+
+    // Note: static mapping instead of exposed via CustomizationConfig to avoid exposing it for wider use unless necessary.
+    private static final Map<String, List<String>> SERVICE_ID_TO_OPERATIONS_MAP;
+
+    static {
+        Map<String, List<String>> serviceIdToOperationsMap = new HashMap<>();
+
+        serviceIdToOperationsMap.put("SQS", Collections.singletonList("ReceiveMessage"));
+        serviceIdToOperationsMap.put("SFN", Collections.singletonList("GetActivityTask"));
+        serviceIdToOperationsMap.put("SWF", Collections.unmodifiableList(Arrays.asList("PollForActivityTask",
+                                                                                       "PollForDecisionTask")));
+
+        SERVICE_ID_TO_OPERATIONS_MAP = Collections.unmodifiableMap(serviceIdToOperationsMap);
+    }
+
+    @Override
+    public void preprocess(ServiceModel serviceModel) {
+        // no-op
+    }
+
+    @Override
+    public void postprocess(IntermediateModel intermediateModel) {
+        String serviceId = intermediateModel.getMetadata().getServiceId();
+        List<String> longPollingOperations = SERVICE_ID_TO_OPERATIONS_MAP.getOrDefault(serviceId, Collections.emptyList());
+
+        for (String longPollingOperation : longPollingOperations) {
+            OperationModel opModel = intermediateModel.getOperation(longPollingOperation);
+            if (opModel != null) {
+                log.info("Setting the longPoll trait for {}#{}", serviceId, longPollingOperation);
+                opModel.setLongPolling(true);
+            } else {
+                log.warn("Did not find operation {}#{}", serviceId, longPollingOperation);
+            }
+        }
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/OperationModel.java

@@ -90,6 +90,8 @@ public class OperationModel extends DocumentationModel {
 
     private boolean unsignedPayload;
 
+    private boolean longPolling;
+
     public String getOperationName() {
         return operationName;
     }
@@ -381,6 +383,14 @@ public void setUnsignedPayload(boolean unsignedPayload) {
         this.unsignedPayload = unsignedPayload;
     }
 
+    public boolean isLongPolling() {
+        return longPolling;
+    }
+
+    public void setLongPolling(boolean longPolling) {
+        this.longPolling = longPolling;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) {
@@ -395,7 +405,8 @@ public boolean equals(Object o) {
                && hasStringMemberAsPayload == that.hasStringMemberAsPayload && isAuthenticated == that.isAuthenticated
                && isPaginated == that.isPaginated && endpointOperation == that.endpointOperation
                && endpointCacheRequired == that.endpointCacheRequired && httpChecksumRequired == that.httpChecksumRequired
-               && unsignedPayload == that.unsignedPayload && Objects.equals(operationName, that.operationName)
+               && unsignedPayload == that.unsignedPayload && longPolling == that.longPolling
+               && Objects.equals(operationName, that.operationName)
                && Objects.equals(serviceProtocol, that.serviceProtocol)
                && Objects.equals(deprecatedMessage, that.deprecatedMessage) && Objects.equals(input, that.input)
                && Objects.equals(returnType, that.returnType) && Objects.equals(exceptions, that.exceptions)
@@ -437,6 +448,7 @@ public int hashCode() {
         result = 31 * result + Objects.hashCode(staticContextParams);
         result = 31 * result + Objects.hashCode(operationContextParams);
         result = 31 * result + Boolean.hashCode(unsignedPayload);
+        result = 31 * result + Boolean.hashCode(longPolling);
         return result;
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumRequiredTrait;
 import software.amazon.awssdk.codegen.poet.client.traits.HttpChecksumTrait;
+import software.amazon.awssdk.codegen.poet.client.traits.LongPollTrait;
 import software.amazon.awssdk.codegen.poet.client.traits.RequestCompressionTrait;
 import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
 import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
@@ -219,6 +220,7 @@ public CodeBlock executionHandler(OperationModel opModel) {
                      .add(hostPrefixExpression(opModel))
                      .add(discoveredEndpoint(opModel))
                      .add(credentialType(opModel, model))
+                     .add(LongPollTrait.executionParamSetter(opModel))
                      .add(".withRequestConfiguration(clientConfiguration)")
                      .add(".withInput($L)\n", opModel.getInput().getVariableName())
                      .add(".withMetricCollector(apiCallMetricCollector)")
@@ -290,6 +292,7 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                .add(".withMarshaller($L)\n", asyncMarshaller(model, opModel, marshaller, protocolFactory))
                .add(asyncRequestBody(opModel))
                .add(fullDuplex(opModel))
+               .add(LongPollTrait.executionParamSetter(opModel))
                .add(hasInitialRequestEvent(opModel, isRestJson))
                .add(".withResponseHandler($L)\n", responseHandlerName(opModel, isRestJson))
                .add(".withErrorResponseHandler(errorResponseHandler)\n")

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/LongPollTrait.java

@@ -0,0 +1,34 @@
+/*
+ * 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.codegen.poet.client.traits;
+
+import com.squareup.javapoet.CodeBlock;
+import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
+
+/**
+ * Helper methods for working with the long poll trait for operations.
+ */
+public final class LongPollTrait {
+    private LongPollTrait() {
+    }
+    public static CodeBlock executionParamSetter(OperationModel operationModel) {
+        if (operationModel.isLongPolling()) {
+            return CodeBlock.of(".withLongPolling(true)");
+        }
+        return CodeBlock.of("");
+    }
+
+}