Address multi-threading issues with AWS Lambda Plugin

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/annotations/SingleThread.java

@@ -17,6 +17,6 @@
 
 @Documented
 @Retention(RetentionPolicy.RUNTIME)
-@Target({ElementType.TYPE})
+@Target({ElementType.CONSTRUCTOR, ElementType.TYPE})
 public @interface SingleThread {
 }

data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/common/LambdaCommonHandler.java

@@ -1,48 +1,16 @@
 package org.opensearch.dataprepper.plugins.lambda.common;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
-import org.opensearch.dataprepper.plugins.lambda.common.accumlator.Buffer;
-import org.opensearch.dataprepper.plugins.lambda.common.accumlator.BufferFactory;
 import org.slf4j.Logger;
-import software.amazon.awssdk.services.lambda.LambdaAsyncClient;
+import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.lambda.model.InvokeResponse;
 
-import java.io.IOException;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
 public class LambdaCommonHandler {
-    private final Logger LOG;
-    private final LambdaAsyncClient lambdaAsyncClient;
-    private final String functionName;
-    private final String invocationType;
-    BufferFactory bufferFactory;
-    private final ObjectMapper objectMapper = new ObjectMapper();
+    private static final Logger LOG = LoggerFactory.getLogger(LambdaCommonHandler.class);
 
-    public LambdaCommonHandler(
-            final Logger log,
-            final LambdaAsyncClient lambdaAsyncClient,
-            final String functionName,
-            final String invocationType,
-            BufferFactory bufferFactory){
-        this.LOG = log;
-        this.lambdaAsyncClient = lambdaAsyncClient;
-        this.functionName = functionName;
-        this.invocationType = invocationType;
-        this.bufferFactory = bufferFactory;
-    }
-
-    public Buffer createBuffer(Buffer currentBuffer) {
-        try {
-            LOG.debug("Resetting buffer");
-            currentBuffer = bufferFactory.getBuffer(lambdaAsyncClient, functionName, invocationType);
-            return currentBuffer;
-        } catch (IOException e) {
-            throw new RuntimeException("Failed to reset buffer", e);
-        }
-    }
-
-    public boolean checkStatusCode(InvokeResponse response) {
+    public static boolean checkStatusCode(InvokeResponse response) {
         int statusCode = response.statusCode();
         if (statusCode < 200 || statusCode >= 300) {
             LOG.error("Lambda invocation returned with non-success status code: {}", statusCode);
@@ -51,7 +19,7 @@ public boolean checkStatusCode(InvokeResponse response) {
         return true;
     }
 
-    public void waitForFutures(List<CompletableFuture<Void>> futureList) {
+    public static void waitForFutures(List<CompletableFuture<Void>> futureList) {
         if (!futureList.isEmpty()) {
             try {
                 CompletableFuture.allOf(futureList.toArray(new CompletableFuture[0])).join();
@@ -63,4 +31,4 @@ public void waitForFutures(List<CompletableFuture<Void>> futureList) {
             }
         }
     }
-}
+}

data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/common/accumlator/Buffer.java

@@ -42,11 +42,8 @@ public interface Buffer {
 
     public Long getPayloadRequestSize();
 
-    public Long getPayloadResponseSize();
-
     public Duration stopLatencyWatch();
 
-
     void reset();
 
 }

data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/common/accumlator/InMemoryBuffer.java

@@ -39,7 +39,6 @@ public class InMemoryBuffer implements Buffer {
     private StopWatch lambdaLatencyWatch;
     private long payloadRequestSize;
     private long payloadResponseSize;
-    private boolean isCodecStarted;
     private final List<Record<Event>> records;
 
 
@@ -53,7 +52,6 @@ public InMemoryBuffer(LambdaAsyncClient lambdaAsyncClient, String functionName,
         bufferWatch.start();
         lambdaLatencyWatch = new StopWatch();
         eventCount = 0;
-        isCodecStarted = false;
         payloadRequestSize = 0;
         payloadResponseSize = 0;
     }
@@ -86,7 +84,6 @@ public void reset() {
         eventCount = 0;
         bufferWatch.reset();
         lambdaLatencyWatch.reset();
-        isCodecStarted = false;
         payloadRequestSize = 0;
         payloadResponseSize = 0;
     }
@@ -160,13 +157,10 @@ public Long getPayloadRequestSize() {
         return payloadRequestSize;
     }
 
-    public Long getPayloadResponseSize() {
-        return payloadResponseSize;
-    }
-
     public StopWatch getBufferWatch() {return bufferWatch;}
 
     public StopWatch getLambdaLatencyWatch(){return lambdaLatencyWatch;}
 
+
 }
 

data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/processor/AggregateResponseEventHandlingStrategy.java

@@ -4,14 +4,19 @@
 import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.record.Record;
-import org.opensearch.dataprepper.plugins.lambda.common.accumlator.Buffer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.List;
 
 public class AggregateResponseEventHandlingStrategy implements ResponseEventHandlingStrategy {
 
+    private static final Logger LOG = LoggerFactory.getLogger(AggregateResponseEventHandlingStrategy.class);
+
     @Override
-    public void handleEvents(List<Event> parsedEvents, List<Record<Event>> originalRecords, List<Record<Event>> resultRecords, Buffer flushedBuffer) {
+    public void handleEvents(List<Event> parsedEvents, List<Record<Event>> originalRecords,
+                             List<Record<Event>> resultRecords) {
+
         Event originalEvent = originalRecords.get(0).getData();
         DefaultEventHandle eventHandle = (DefaultEventHandle) originalEvent.getEventHandle();
         AcknowledgementSet originalAcknowledgementSet = eventHandle.getAcknowledgementSet();
@@ -26,5 +31,6 @@ public void handleEvents(List<Event> parsedEvents, List<Record<Event>> originalR
                 originalAcknowledgementSet.add(responseEvent);
             }
         }
+        LOG.info("Successfully handled {} events in Aggregate response strategy", parsedEvents.size());
     }
 }