Add a few sample Lambda request handlers useful for benchmarking

Author: robin-aws

src/main/java/com/amazonaws/services/sqs/IdentityFunction.java

@@ -0,0 +1,15 @@
+package com.amazonaws.services.sqs;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+
+/**
+ * A Lambda that just returns its input. Useful for benchmarking purposes.
+ */
+public class IdentityFunction implements RequestHandler<String, String> {
+
+    @Override
+    public String handleRequest(String input, Context context) {
+        return input;
+    }
+}

src/main/java/com/amazonaws/services/sqs/SQSIdentityFunction.java

@@ -0,0 +1,45 @@
+package com.amazonaws.services.sqs;
+
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent.MessageAttribute;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
+import com.amazonaws.services.sqs.model.MessageAttributeValue;
+
+/**
+ * An SQS-triggered Lambda function that just returns the body of the
+ * SQS message. If the message is a request sent by a Requester client,
+ * it sends the response through SQS as well.
+ */
+public class SQSIdentityFunction implements RequestHandler<SQSEvent, String> {
+
+    private static final AmazonSQSResponder sqs = 
+            AmazonSQSResponderClientBuilder.defaultClient();
+    
+    @Override
+    public String handleRequest(SQSEvent input, Context context) {
+        SQSMessage message = input.getRecords().get(0);
+        String messageBody = message.getBody();
+        
+        MessageContent requestMessage = new MessageContent(messageBody, 
+                message.getMessageAttributes().entrySet().stream().collect(
+                        Collectors.toMap(Entry::getKey, e -> toMessageAttributeValue(e.getValue()))));
+        if (sqs.isResponseMessageRequested(requestMessage)) {
+            MessageContent response = new MessageContent(messageBody);
+            sqs.sendResponseMessage(requestMessage, response);
+        }
+        return messageBody;
+    }
+
+    private MessageAttributeValue toMessageAttributeValue(MessageAttribute attribute) {
+        return new MessageAttributeValue().withDataType(attribute.getDataType())
+                                          .withStringValue(attribute.getStringValue())
+                                          .withStringListValues(attribute.getStringListValues())
+                                          .withBinaryValue(attribute.getBinaryValue())
+                                          .withBinaryListValues(attribute.getBinaryListValues());
+    }
+}

src/main/java/com/amazonaws/services/sqs/SQSIdentityFunctionRequester.java

@@ -0,0 +1,36 @@
+package com.amazonaws.services.sqs;
+
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.sqs.model.Message;
+import com.amazonaws.services.sqs.model.SendMessageRequest;
+
+/**
+ * A Lambda that invokes the SQSIdentityFunction using the SQS request client.
+ */
+public class SQSIdentityFunctionRequester implements RequestHandler<String, String> {
+
+    private static final AmazonSQSRequester requester = AmazonSQSRequesterClientBuilder.defaultClient();
+    
+    private static final String SQS_IDENTITY_FUNCTION_QUEUE_URL = System.getenv("SQS_IDENTITY_FUNCTION_QUEUE_URL");
+    
+    @Override
+    public String handleRequest(String input, Context context) {
+        SendMessageRequest sendRequest = new SendMessageRequest()
+            .withQueueUrl(SQS_IDENTITY_FUNCTION_QUEUE_URL)
+            .withMessageBody(input);
+        try {
+            Message response = requester.sendMessageAndGetResponse(sendRequest, 30, TimeUnit.SECONDS);
+            String result = response.getBody();
+            if (!input.equals(result)) {
+                throw new RuntimeException("Incorrect result: " + result);
+            }
+            return result;
+        } catch (TimeoutException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}