Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jul 11, 2025
+425 −69
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6e7e643
feat: add support for batch execution in parallel with custom Executor
visheshruparelia Jun 18, 2025
f06e692
Merge branch 'main' into parallel_execution_custom_executor
phipag Jul 10, 2025
ab4b667
Merge branch 'main' into parallel_execution_custom_executor
phipag Jul 10, 2025
850c155
pmd_analyse fix
visheshruparelia Jul 11, 2025
4d0cd06
add docs
visheshruparelia Jul 11, 2025
c4361a9
Merge remote-tracking branch 'parent/main' into parallel_execution_cu…
visheshruparelia Jul 11, 2025
5758c71
Update docs/utilities/batch.md
visheshruparelia Jul 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions ...xamples-batch/src/main/java/org/demo/batch/dynamo/DynamoDBStreamBatchHandlerParallel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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, Context context) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can remove Context c here please. There is an overload for a message handler without context in AbstractBatchMessageHandlerBuilder.java

    /**
     * Builds a BatchMessageHandler that can be used to process batches, given
     * a user-defined handler to process each item in the batch. This variant
     * takes a function that consumes a raw message and the Lambda context. This
     * is useful for handlers that need access to the entire message object, not
     * just the deserialized contents of the body.
     *
     * @param handler Takes a raw message - the underlying AWS Events Library event - to process.
     *                For instance for SQS this would be an SQSMessage.
     * @return A BatchMessageHandler for processing the batch
     */
    public BatchMessageHandler<E, R> buildWithRawMessageHandler(Consumer<T> handler) {
        return buildWithRawMessageHandler((f, c) -> handler.accept(f));
    }

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The older examples in the package follow the same pattern as the new one. Do you want those changed as well?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this would be awesome if you can update those as well. Ideally, we have no Sonar or pmd_analyze findings before merging the PR. With the exception of the code duplication finding.

Copy link
Contributor Author

@visheshruparelia visheshruparelia Jul 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

850c155

LOGGER.info("Processing DynamoDB Stream Record" + dynamodbStreamRecord);
}

}
39 changes: 39 additions & 0 deletions ...ools-examples-batch/src/main/java/org/demo/batch/kinesis/KinesisBatchHandlerParallel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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, Context c) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can remove Context c here please. There is an overload for a message handler without context in AbstractBatchMessageHandlerBuilder.java

    /**
     * Builds a BatchMessageHandler that can be used to process batches, given
     * a user-defined handler to process each item in the batch. This variant
     * takes a function that consumes the deserialized body of the given message
     * If deserialization fails, it will be treated as
     * failure of the processing of that item in the batch.
     * Note:  If you don't need the Lambda context, use the variant of this function
     * that does not require it.
     *
     * @param handler Processes the deserialized body of the message
     * @return A BatchMessageHandler for processing the batch
     */
    public <M> BatchMessageHandler<E, R> buildWithMessageHandler(Consumer<M> handler, Class<M> messageClass) {
        return buildWithMessageHandler((f, c) -> handler.accept(f), messageClass);
    }

Copy link
Contributor Author

@visheshruparelia visheshruparelia Jul 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

850c155

LOGGER.info("Processing product " + p);
}

}
37 changes: 37 additions & 0 deletions ...s/powertools-examples-batch/src/main/java/org/demo/batch/sqs/SqsBatchHandlerParallel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
}
}
13 changes: 13 additions & 0 deletions ...ch/src/main/java/software/amazon/lambda/powertools/batch/handler/BatchMessageHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down Expand Up @@ -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);
}
26 changes: 25 additions & 1 deletion ...ain/java/software/amazon/lambda/powertools/batch/handler/DynamoDbBatchMessageHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,13 @@
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.concurrent.ExecutorService;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;
Expand Down Expand Up @@ -66,7 +71,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)
Expand All @@ -75,6 +82,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());
Expand Down
25 changes: 24 additions & 1 deletion ...va/software/amazon/lambda/powertools/batch/handler/KinesisStreamsBatchMessageHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand All @@ -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());
Expand Down
36 changes: 34 additions & 2 deletions ...src/main/java/software/amazon/lambda/powertools/batch/handler/SqsBatchMessageHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,20 @@
package software.amazon.lambda.powertools.batch.handler;

import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
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.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.lambda.powertools.batch.internal.MultiThreadMDC;
Expand Down Expand Up @@ -99,7 +104,7 @@ public SQSBatchResponse processBatch(SQSEvent event, Context context) {

@Override
public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context) {
if (!event.getRecords().isEmpty() && event.getRecords().get(0).getAttributes().get(MESSAGE_GROUP_ID_KEY) != null) {
if (isFIFOEnabled(event)) {
throw new UnsupportedOperationException("FIFO queues are not supported in parallel mode, use the processBatch method instead");
}

Expand All @@ -109,7 +114,9 @@ public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context)
.map(sqsMessage -> {

multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
return processBatchItem(sqsMessage, context);
Optional<SQSBatchResponse.BatchItemFailure> failureOpt = processBatchItem(sqsMessage, context);
multiThreadMDC.removeThread(Thread.currentThread().getName());
return failureOpt;
})
.filter(Optional::isPresent)
.map(Optional::get)
Expand All @@ -118,6 +125,27 @@ public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context)
return SQSBatchResponse.builder().withBatchItemFailures(batchItemFailures).build();
}

@Override
public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context, Executor executor) {
if (isFIFOEnabled(event)) {
throw new UnsupportedOperationException("FIFO queues are not supported in parallel mode, use the processBatch method instead");
}

MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
List<SQSBatchResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
List<CompletableFuture<Void>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.runAsync(() -> {
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
Optional<SQSBatchResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord, context);
failureOpt.ifPresent(batchItemFailures::add);
multiThreadMDC.removeThread(Thread.currentThread().getName());
}, executor))
.collect(Collectors.toList());
futures.forEach(CompletableFuture::join);

return SQSBatchResponse.builder().withBatchItemFailures(batchItemFailures).build();
}

private Optional<SQSBatchResponse.BatchItemFailure> processBatchItem(SQSEvent.SQSMessage message, Context context) {
try {
LOGGER.debug("Processing message {}", message.getMessageId());
Expand Down Expand Up @@ -152,4 +180,8 @@ private Optional<SQSBatchResponse.BatchItemFailure> processBatchItem(SQSEvent.SQ
.build());
}
}

private boolean isFIFOEnabled(SQSEvent sqsEvent) {
return !sqsEvent.getRecords().isEmpty() && sqsEvent.getRecords().get(0).getAttributes().get(MESSAGE_GROUP_ID_KEY) != null;
}
}
7 changes: 7 additions & 0 deletions ...-batch/src/main/java/software/amazon/lambda/powertools/batch/internal/MultiThreadMDC.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,4 +44,11 @@ public void copyMDCToThread(String thread) {
mdcAwareThreads.add(thread);
}
}

public void removeThread(String thread) {
if (mdcAwareThreads.contains(thread)) {
LOGGER.debug("Removing thread {}", thread);
mdcAwareThreads.remove(thread);
}
}
}
Loading