feat(batch): add support for batch execution in parallel with custom Executor (#1900)

* feat: add support for batch execution in parallel with custom Executor

* pmd_analyse fix

* add docs

* Update docs/utilities/batch.md

Co-authored-by: Philipp Page <[email protected]>

---------

Co-authored-by: Philipp Page <[email protected]>
Co-authored-by: Philipp Page <[email protected]>

Author: visheshruparelia

docs/utilities/batch.md

@@ -510,6 +510,31 @@ used with SQS FIFO. In that case, an `UnsupportedOperationException` is thrown.
         }
     }
     ```
+=== "Example with SQS (using custom executor)"
+
+    ```java hl_lines="4 10 15"
+    public class SqsBatchHandler implements RequestHandler<SQSEvent, SQSBatchResponse> {
+
+        private final BatchMessageHandler<SQSEvent, SQSBatchResponse> handler;
+        private final ExecutorService executor;
+    
+        public SqsBatchHandler() {
+            handler = new BatchMessageHandlerBuilder()
+                    .withSqsBatchHandler()
+                    .buildWithMessageHandler(this::processMessage, Product.class);
+            executor = Executors.newFixedThreadPool(2);
+        }
+    
+        @Override
+        public SQSBatchResponse handleRequest(SQSEvent sqsEvent, Context context) {
+            return handler.processBatchInParallel(sqsEvent, context, executor);
+        }
+    
+        private void processMessage(Product p, Context c) {
+            // Process the product
+        }
+    }
+    ```
 
 
 ## Handling Messages

examples/powertools-examples-batch/src/main/java/org/demo/batch/dynamo/DynamoDBStreamBatchHandler.java

@@ -25,7 +25,7 @@ public StreamsEventResponse handleRequest(DynamodbEvent ddbEvent, Context contex
         return handler.processBatch(ddbEvent, context);
     }
 
-    private void processMessage(DynamodbEvent.DynamodbStreamRecord dynamodbStreamRecord, Context context) {
+    private void processMessage(DynamodbEvent.DynamodbStreamRecord dynamodbStreamRecord) {
         LOGGER.info("Processing DynamoDB Stream Record" + dynamodbStreamRecord);
     }
 

examples/powertools-examples-batch/src/main/java/org/demo/batch/dynamo/DynamoDBStreamBatchHandlerParallel.java

@@ -0,0 +1,37 @@
+package org.demo.batch.dynamo;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.DynamodbEvent;
+import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.lambda.powertools.batch.BatchMessageHandlerBuilder;
+import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class DynamoDBStreamBatchHandlerParallel implements RequestHandler<DynamodbEvent, StreamsEventResponse> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBStreamBatchHandlerParallel.class);
+    private final BatchMessageHandler<DynamodbEvent, StreamsEventResponse> handler;
+    private final ExecutorService executor;
+
+    public DynamoDBStreamBatchHandlerParallel() {
+        handler = new BatchMessageHandlerBuilder()
+                .withDynamoDbBatchHandler()
+                .buildWithRawMessageHandler(this::processMessage);
+        executor = Executors.newFixedThreadPool(2);
+    }
+
+    @Override
+    public StreamsEventResponse handleRequest(DynamodbEvent ddbEvent, Context context) {
+        return handler.processBatchInParallel(ddbEvent, context, executor);
+    }
+
+    private void processMessage(DynamodbEvent.DynamodbStreamRecord dynamodbStreamRecord) {
+        LOGGER.info("Processing DynamoDB Stream Record" + dynamodbStreamRecord);
+    }
+
+}

examples/powertools-examples-batch/src/main/java/org/demo/batch/kinesis/KinesisBatchHandler.java

@@ -26,7 +26,7 @@ public StreamsEventResponse handleRequest(KinesisEvent kinesisEvent, Context con
         return handler.processBatch(kinesisEvent, context);
     }
 
-    private void processMessage(Product p, Context c) {
+    private void processMessage(Product p) {
         LOGGER.info("Processing product " + p);
     }
 

examples/powertools-examples-batch/src/main/java/org/demo/batch/kinesis/KinesisBatchHandlerParallel.java

@@ -0,0 +1,39 @@
+package org.demo.batch.kinesis;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
+import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
+import org.demo.batch.model.Product;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.lambda.powertools.batch.BatchMessageHandlerBuilder;
+import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class KinesisBatchHandlerParallel implements RequestHandler<KinesisEvent, StreamsEventResponse> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(KinesisBatchHandlerParallel.class);
+    private final BatchMessageHandler<KinesisEvent, StreamsEventResponse> handler;
+    private final ExecutorService executor;
+
+
+    public KinesisBatchHandlerParallel() {
+        handler = new BatchMessageHandlerBuilder()
+                .withKinesisBatchHandler()
+                .buildWithMessageHandler(this::processMessage, Product.class);
+        executor = Executors.newFixedThreadPool(2);
+    }
+
+    @Override
+    public StreamsEventResponse handleRequest(KinesisEvent kinesisEvent, Context context) {
+        return handler.processBatchInParallel(kinesisEvent, context, executor);
+    }
+
+    private void processMessage(Product p) {
+        LOGGER.info("Processing product " + p);
+    }
+
+}

examples/powertools-examples-batch/src/main/java/org/demo/batch/sqs/AbstractSqsBatchHandler.java

@@ -42,11 +42,10 @@ public class AbstractSqsBatchHandler {
     /**
      * Simulate some processing (I/O + S3 put request)
      * @param p deserialized product
-     * @param context Lambda context
      */
     @Logging
     @Tracing
-    protected void processMessage(Product p, Context context) {
+    protected void processMessage(Product p) {
         TracingUtils.putAnnotation("productId", p.getId());
         TracingUtils.putAnnotation("Thread", Thread.currentThread().getName());
         MDC.put("product", String.valueOf(p.getId()));

examples/powertools-examples-batch/src/main/java/org/demo/batch/sqs/SqsBatchHandlerParallel.java

@@ -0,0 +1,37 @@
+package org.demo.batch.sqs;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import org.demo.batch.model.Product;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.lambda.powertools.batch.BatchMessageHandlerBuilder;
+import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
+import software.amazon.lambda.powertools.logging.Logging;
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class SqsBatchHandlerParallel extends AbstractSqsBatchHandler implements RequestHandler<SQSEvent, SQSBatchResponse> {
+    private static final Logger LOGGER = LoggerFactory.getLogger(SqsBatchHandlerParallel.class);
+    private final BatchMessageHandler<SQSEvent, SQSBatchResponse> handler;
+    private final ExecutorService executor;
+
+    public SqsBatchHandlerParallel() {
+        handler = new BatchMessageHandlerBuilder()
+                .withSqsBatchHandler()
+                .buildWithMessageHandler(this::processMessage, Product.class);
+        executor = Executors.newFixedThreadPool(2);
+    }
+
+    @Logging
+    @Tracing
+    @Override
+    public SQSBatchResponse handleRequest(SQSEvent sqsEvent, Context context) {
+        LOGGER.info("Processing batch of {} messages", sqsEvent.getRecords().size());
+        return handler.processBatchInParallel(sqsEvent, context, executor);
+    }
+}

powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/BatchMessageHandler.java

@@ -16,6 +16,9 @@
 
 import com.amazonaws.services.lambda.runtime.Context;
 
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+
 /**
  * The basic interface a batch message handler must meet.
  *
@@ -50,4 +53,14 @@ public interface BatchMessageHandler<E, R> {
      * @return A partial batch response
      */
     R processBatchInParallel(E event, Context context);
+
+
+    /**
+     * Same as {@link #processBatchInParallel(Object, Context)} but with an option to provide custom {@link Executor}
+     * @param event   The Lambda event containing the batch to process
+     * @param context The lambda context
+     * @param executor Custom executor to use for parallel processing
+     * @return A partial batch response
+     */
+    R processBatchInParallel(E event, Context context, Executor executor);
 }

powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/DynamoDbBatchMessageHandler.java

@@ -17,8 +17,12 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.events.DynamodbEvent;
 import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
+
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -66,7 +70,9 @@ public StreamsEventResponse processBatchInParallel(DynamodbEvent event, Context
                 .parallelStream() // Parallel processing
                 .map(eventRecord -> {
                     multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
-                    return processBatchItem(eventRecord, context);
+                    Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord, context);
+                    multiThreadMDC.removeThread(Thread.currentThread().getName());
+                    return failureOpt;
                 })
                 .filter(Optional::isPresent)
                 .map(Optional::get)
@@ -75,6 +81,23 @@ public StreamsEventResponse processBatchInParallel(DynamodbEvent event, Context
         return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
     }
 
+    @Override
+    public StreamsEventResponse processBatchInParallel(DynamodbEvent event, Context context, Executor executor) {
+        MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
+
+        List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
+        List<CompletableFuture<Void>> futures = event.getRecords().stream()
+                .map(eventRecord -> CompletableFuture.runAsync(() -> {
+                    multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
+                    Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord, context);
+                    failureOpt.ifPresent(batchItemFailures::add);
+                    multiThreadMDC.removeThread(Thread.currentThread().getName());
+                }, executor))
+                .collect(Collectors.toList());
+        futures.forEach(CompletableFuture::join);
+        return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
+    }
+
     private Optional<StreamsEventResponse.BatchItemFailure> processBatchItem(DynamodbEvent.DynamodbStreamRecord streamRecord, Context context) {
         try {
             LOGGER.debug("Processing item {}", streamRecord.getEventID());
@@ -86,19 +109,19 @@ private Optional<StreamsEventResponse.BatchItemFailure> processBatchItem(Dynamod
                 this.successHandler.accept(streamRecord);
             }
             return Optional.empty();
-        } catch (Throwable t) {
+        } catch (Exception e) {
             String sequenceNumber = streamRecord.getDynamodb().getSequenceNumber();
             LOGGER.error("Error while processing record with id {}: {}, adding it to batch item failures",
-                    sequenceNumber, t.getMessage());
-            LOGGER.error("Error was", t);
+                    sequenceNumber, e.getMessage());
+            LOGGER.error("Error was", e);
 
             // Report failure if we have a handler
             if (this.failureHandler != null) {
                 // A failing failure handler is no reason to fail the batch
                 try {
-                    this.failureHandler.accept(streamRecord, t);
-                } catch (Throwable t2) {
-                    LOGGER.warn("failureHandler threw handling failure", t2);
+                    this.failureHandler.accept(streamRecord, e);
+                } catch (Exception e2) {
+                    LOGGER.warn("failureHandler threw handling failure", e2);
                 }
             }
             return Optional.of(StreamsEventResponse.BatchItemFailure.builder().withItemIdentifier(sequenceNumber).build());

powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/KinesisStreamsBatchMessageHandler.java

@@ -18,8 +18,12 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
 import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
+
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -77,7 +81,9 @@ public StreamsEventResponse processBatchInParallel(KinesisEvent event, Context c
                 .parallelStream() // Parallel processing
                 .map(eventRecord -> {
                     multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
-                    return processBatchItem(eventRecord, context);
+                    Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord, context);
+                    multiThreadMDC.removeThread(Thread.currentThread().getName());
+                    return failureOpt;
                 })
                 .filter(Optional::isPresent)
                 .map(Optional::get)
@@ -86,6 +92,23 @@ public StreamsEventResponse processBatchInParallel(KinesisEvent event, Context c
         return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
     }
 
+    @Override
+    public StreamsEventResponse processBatchInParallel(KinesisEvent event, Context context, Executor executor) {
+        MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
+
+        List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
+        List<CompletableFuture<Void>> futures = event.getRecords().stream()
+                .map(eventRecord -> CompletableFuture.runAsync(() -> {
+                    multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
+                    Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord, context);
+                    failureOpt.ifPresent(batchItemFailures::add);
+                    multiThreadMDC.removeThread(Thread.currentThread().getName());
+                }, executor))
+                .collect(Collectors.toList());
+        futures.forEach(CompletableFuture::join);
+        return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
+    }
+
     private Optional<StreamsEventResponse.BatchItemFailure> processBatchItem(KinesisEvent.KinesisEventRecord eventRecord, Context context) {
         try {
             LOGGER.debug("Processing item {}", eventRecord.getEventID());
@@ -102,19 +125,19 @@ private Optional<StreamsEventResponse.BatchItemFailure> processBatchItem(Kinesis
                 this.successHandler.accept(eventRecord);
             }
             return Optional.empty();
-        } catch (Throwable t) {
+        } catch (Exception e) {
             String sequenceNumber = eventRecord.getEventID();
             LOGGER.error("Error while processing record with eventID {}: {}, adding it to batch item failures",
-                    sequenceNumber, t.getMessage());
-            LOGGER.error("Error was", t);
+                    sequenceNumber, e.getMessage());
+            LOGGER.error("Error was", e);
 
             // Report failure if we have a handler
             if (this.failureHandler != null) {
                 // A failing failure handler is no reason to fail the batch
                 try {
-                    this.failureHandler.accept(eventRecord, t);
-                } catch (Throwable t2) {
-                    LOGGER.warn("failureHandler threw handling failure", t2);
+                    this.failureHandler.accept(eventRecord, e);
+                } catch (Exception e2) {
+                    LOGGER.warn("failureHandler threw handling failure", e2);
                 }
             }