Fix lambda plugin threshold and flaky test (opensearch-project#5523)

Signed-off-by: Srikanth Govindarajan <srigovs@amazon.com>

Author: srikanthjg

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

@@ -63,16 +63,18 @@ private static List<Buffer> createBufferBatches(Collection<Record<Event>> record
 
         LOG.debug("Batch size received to lambda processor: {}", records.size());
         for (Record<Event> record : records) {
-            //check size or time has exceeded threshold
-            if (ThresholdCheck.checkSizeThresholdExceed(currentBufferPerBatch, maxBytes, record)) {
+            //check size has exceeded threshold
+            if (currentBufferPerBatch.getEventCount() > 0 &&
+                    ThresholdCheck.checkSizeThresholdExceed(currentBufferPerBatch, maxBytes, record)) {
                 batchedBuffers.add(currentBufferPerBatch);
                 currentBufferPerBatch = new InMemoryBuffer(keyName, outputCodecContext);
             }
 
             currentBufferPerBatch.addRecord(record);
 
             // After adding, check if the event count threshold is reached.
-            if (ThresholdCheck.checkEventCountThresholdExceeded(currentBufferPerBatch, maxEvents)) {
+            if (currentBufferPerBatch.getEventCount() > 0 &&
+                    ThresholdCheck.checkEventCountThresholdExceeded(currentBufferPerBatch, maxEvents)) {
                 batchedBuffers.add(currentBufferPerBatch);
                 currentBufferPerBatch = new InMemoryBuffer(keyName, outputCodecContext);
             }
@@ -102,8 +104,12 @@ public static Map<Buffer, CompletableFuture<InvokeResponse>> invokeLambdaAndGetF
         for (Buffer buffer : batchedBuffers) {
             InvokeRequest requestPayload = buffer.getRequestPayload(config.getFunctionName(),
                     config.getInvocationType().getAwsLambdaValue());
-            CompletableFuture<InvokeResponse> future = lambdaAsyncClient.invoke(requestPayload);
-            bufferToFutureMap.put(buffer, future);
+            if(requestPayload!=null) {
+                CompletableFuture<InvokeResponse> future = lambdaAsyncClient.invoke(requestPayload);
+                bufferToFutureMap.put(buffer, future);
+            }else{
+                LOG.warn("Request Payload is null, skipping lambda invocation");
+            }
         }
         return bufferToFutureMap;
     }

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

@@ -19,15 +19,24 @@
 public class ThresholdCheck {
 
       public static boolean checkTimeoutExceeded(final Buffer currentBuffer, final Duration maxCollectionDuration) {
+          if (currentBuffer.getEventCount() == 0) {
+              return false;
+          }
         return currentBuffer.getDuration().compareTo(maxCollectionDuration) > 0;
     }
 
      public static boolean checkSizeThresholdExceed(final Buffer currentBuffer, final ByteCount maxBytes, Record<Event> nextRecord) {
+         if (currentBuffer.getEventCount() == 0) {
+             return false;
+         }
         int estimatedRecordSize = estimateRecordSize(nextRecord);
         return (currentBuffer.getSize() + estimatedRecordSize) > maxBytes.getBytes();
     }
 
     public static boolean checkEventCountThresholdExceeded(final Buffer currentBuffer, final int maxEvents) {
+        if (currentBuffer.getEventCount() == 0) {
+            return false;
+        }
         return currentBuffer.getEventCount() >= maxEvents;
     }
 

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

@@ -36,6 +36,7 @@
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.services.lambda.LambdaAsyncClient;
 import software.amazon.awssdk.services.lambda.model.InvokeResponse;
+import software.amazon.awssdk.services.lambda.model.LambdaException;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -201,7 +202,9 @@ public Collection<Record<Event>> doExecute(Collection<Record<Event>> records) {
 
             } catch (Exception e) {
                 LOG.error(NOISY, e.getMessage(), e);
-                if(e.getMessage().contains(EXCEEDING_PAYLOAD_LIMIT_EXCEPTION)){
+                if (e instanceof LambdaException &&
+                        e.getMessage() != null &&
+                        e.getMessage().contains(EXCEEDING_PAYLOAD_LIMIT_EXCEPTION)) {
                     batchExceedingThresholdCounter.increment();
                 }
                 /* fall through */

data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/ThresholdCheckTest.java

@@ -55,6 +55,7 @@ void setUp() {
     void testTimeoutExceededTrue() {
         // Simulate a buffer that has been open for 6 minutes (exceeding the 5-minute limit)
         when(buffer.getDuration()).thenReturn(Duration.ofMinutes(6));
+        when(buffer.getEventCount()).thenReturn(1);
         assertTrue(ThresholdCheck.checkTimeoutExceeded(buffer, maxCollectionDuration),
                 "Expected timeout threshold to be exceeded.");
     }
@@ -63,6 +64,7 @@ void testTimeoutExceededTrue() {
     void testTimeoutExceededFalse() {
         // Simulate a buffer that has been open for 4 minutes (within the limit)
         when(buffer.getDuration()).thenReturn(Duration.ofMinutes(4));
+        when(buffer.getEventCount()).thenReturn(1);
         assertFalse(ThresholdCheck.checkTimeoutExceeded(buffer, maxCollectionDuration),
                 "Expected timeout threshold to NOT be exceeded.");
     }
@@ -71,6 +73,7 @@ void testTimeoutExceededFalse() {
     @Test
     void testSizeThresholdExceedTrue() {
         long maxBytesValue = maxBytes.getBytes();
+        when(buffer.getEventCount()).thenReturn(1);
         when(buffer.getSize()).thenReturn(maxBytesValue - 1);
         // The record's estimated size is 2 bytes (from "{}").
         // So, adding the record yields: (maxBytes - 1) + 2 = maxBytes + 1 (exceeds limit).