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 3 commits 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
3 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
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) {

Check failure on line 33 in examples/powertools-examples-batch/src/main/java/org/demo/batch/dynamo/DynamoDBStreamBatchHandlerParallel.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Avoid unused method parameters such as 'context'. 

Reports parameters of methods and constructors that are not referenced them in the method body.
Parameters whose name starts with `ignored` or `unused` are filtered out.

Removing unused formal parameters from public methods could cause a ripple effect through the code base.
Hence, by default, this rule only considers private methods. To include non-private methods, set the
`checkAll` property to `true`.


UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter
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));
    }

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) {

Check failure on line 35 in examples/powertools-examples-batch/src/main/java/org/demo/batch/kinesis/KinesisBatchHandlerParallel.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Avoid unused method parameters such as 'c'. 

Reports parameters of methods and constructors that are not referenced them in the method body.
Parameters whose name starts with `ignored` or `unused` are filtered out.

Removing unused formal parameters from public methods could cause a ripple effect through the code base.
Hence, by default, this rule only considers private methods. To include non-private methods, set the
`checkAll` property to `true`.


UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter
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);
    }

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 @@
.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 @@
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 All @@ -86,7 +110,7 @@
this.successHandler.accept(streamRecord);
}
return Optional.empty();
} catch (Throwable t) {

Check failure on line 113 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/DynamoDbBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
String sequenceNumber = streamRecord.getDynamodb().getSequenceNumber();
LOGGER.error("Error while processing record with id {}: {}, adding it to batch item failures",
sequenceNumber, t.getMessage());
Expand All @@ -97,7 +121,7 @@
// A failing failure handler is no reason to fail the batch
try {
this.failureHandler.accept(streamRecord, t);
} catch (Throwable t2) {

Check failure on line 124 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/DynamoDbBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
LOGGER.warn("failureHandler threw handling failure", t2);
}
}
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 @@
.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());
Copy link
Contributor

Choose a reason for hiding this comment

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

Good catch!!

return failureOpt;
})
.filter(Optional::isPresent)
.map(Optional::get)
Expand All @@ -86,6 +92,23 @@
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 All @@ -102,7 +125,7 @@
this.successHandler.accept(eventRecord);
}
return Optional.empty();
} catch (Throwable t) {

Check failure on line 128 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/KinesisStreamsBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
String sequenceNumber = eventRecord.getEventID();
LOGGER.error("Error while processing record with eventID {}: {}, adding it to batch item failures",
sequenceNumber, t.getMessage());
Expand All @@ -113,7 +136,7 @@
// A failing failure handler is no reason to fail the batch
try {
this.failureHandler.accept(eventRecord, t);
} catch (Throwable t2) {

Check failure on line 139 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/KinesisStreamsBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
LOGGER.warn("failureHandler threw handling failure", t2);
}
}
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 @@

@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 @@
.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 @@
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 All @@ -134,7 +162,7 @@
this.successHandler.accept(message);
}
return Optional.empty();
} catch (Throwable t) {

Check failure on line 165 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/SqsBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
LOGGER.error("Error while processing message with messageId {}: {}, adding it to batch item failures",
message.getMessageId(), t.getMessage());
LOGGER.error("Error was", t);
Expand All @@ -144,7 +172,7 @@
// A failing failure handler is no reason to fail the batch
try {
this.failureHandler.accept(message, t);
} catch (Throwable t2) {

Check failure on line 175 in powertools-batch/src/main/java/software/amazon/lambda/powertools/batch/handler/SqsBatchMessageHandler.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

A catch statement should never catch throwable since it includes errors. 

Catching Throwable errors is not recommended since its scope is very broad. It includes runtime issues such as
OutOfMemoryError that should be exposed and managed separately.

AvoidCatchingThrowable (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#avoidcatchingthrowable
LOGGER.warn("failureHandler threw handling failure", t2);
}
}
Expand All @@ -152,4 +180,8 @@
.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.

Copy link
Contributor

Choose a reason for hiding this comment

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

Yes, this is good. I believe it was a small memory leak because we never removed the thread either.

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