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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+378 −47
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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) {
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) {
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) {
Copy link
Author

Choose a reason for hiding this comment

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

I think this will be helpful if the same invocation of processBatchInParallel() ends up re-using the thread.

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