Skip to content

fix(batch): make batch failure tracking thread-safe for parallel processing #2735

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
phipag merged 1 commit into from
Nov 4, 2026
+152 −27
Merged

fix(batch): make batch failure tracking thread-safe for parallel processing #2735

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
17 changes: 8 additions & 9 deletions ...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 @@ -14,7 +14,6 @@

package software.amazon.lambda.powertools.batch.handler;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
Expand Down Expand Up @@ -99,22 +98,22 @@ public StreamsEventResponse processBatchInParallel(DynamodbEvent event, Context
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();

List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
List<CompletableFuture<Void>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.runAsync(() -> {
List<CompletableFuture<Optional<StreamsEventResponse.BatchItemFailure>>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
AtomicReference<Optional<StreamsEventResponse.BatchItemFailure>> result = new AtomicReference<>();
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
try {
Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
context);
failureOpt.ifPresent(batchItemFailures::add);
result.set(processBatchItem(eventRecord, context));
} finally {
multiThreadMDC.removeThread(Thread.currentThread().getName());
}
});
return result.get();
}, executor))
.collect(Collectors.toList());
futures.forEach(CompletableFuture::join);
.toList();
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).toList();
return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
}

Expand Down
17 changes: 8 additions & 9 deletions ...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 @@ -14,7 +14,6 @@

package software.amazon.lambda.powertools.batch.handler;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
Expand Down Expand Up @@ -109,22 +108,22 @@ public StreamsEventResponse processBatchInParallel(KinesisEvent event, Context c
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();

List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
List<CompletableFuture<Void>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.runAsync(() -> {
List<CompletableFuture<Optional<StreamsEventResponse.BatchItemFailure>>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
AtomicReference<Optional<StreamsEventResponse.BatchItemFailure>> result = new AtomicReference<>();
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
try {
Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
context);
failureOpt.ifPresent(batchItemFailures::add);
result.set(processBatchItem(eventRecord, context));
} finally {
multiThreadMDC.removeThread(Thread.currentThread().getName());
}
});
return result.get();
}, executor))
.collect(Collectors.toList());
futures.forEach(CompletableFuture::join);
.toList();
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).toList();
return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
}

Expand Down
16 changes: 8 additions & 8 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 @@ -147,22 +147,22 @@ public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context,
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();

List<SQSBatchResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
List<CompletableFuture<Void>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.runAsync(() -> {
List<CompletableFuture<Optional<SQSBatchResponse.BatchItemFailure>>> futures = event.getRecords().stream()
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
AtomicReference<Optional<SQSBatchResponse.BatchItemFailure>> result = new AtomicReference<>();
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
try {
Optional<SQSBatchResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
context);
failureOpt.ifPresent(batchItemFailures::add);
result.set(processBatchItem(eventRecord, context));
} finally {
multiThreadMDC.removeThread(Thread.currentThread().getName());
}
});
return result.get();
}, executor))
.collect(Collectors.toList());
futures.forEach(CompletableFuture::join);
.toList();
List<SQSBatchResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).toList();

return SQSBatchResponse.builder().withBatchItemFailures(batchItemFailures).build();
}
Expand Down
31 changes: 31 additions & 0 deletions ...ls-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,10 @@
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.DynamodbEvent;
import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
import com.amazonaws.services.lambda.runtime.events.models.dynamodb.StreamRecord;
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
Expand All @@ -33,6 +36,7 @@
import org.junit.jupiter.params.ParameterizedTest;
import org.mockito.Mock;
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;

class DdbBatchProcessorTest {

Expand All @@ -52,13 +56,13 @@
System.out.println("Great success, record: " + record + ", context: " + context);
}

private void processMessageFailsForFixedMessage(DynamodbEvent.DynamodbStreamRecord record, Context context) {

Check failure on line 59 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.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`. The same applies to public constructors.

UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter
if (record.getDynamodb().getSequenceNumber().equals("4421584500000000017450439091")) {

Check failure on line 60 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Position literals first in String comparisons 

Position literals first in all String comparisons, if the second argument is null then NullPointerExceptions
can be avoided, they will just return false. Note that switching literal positions for compareTo and
compareToIgnoreCase may change the result, see examples.

Note that compile-time constant strings are treated like literals. This is because they are inlined into
the class file, are necessarily non-null, and therefore cannot cause an NPE at runtime.

LiteralsFirstInComparisons (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#literalsfirstincomparisons
throw new RuntimeException("fake exception");
}
}

private void processMessageInParallelSucceeds(DynamodbEvent.DynamodbStreamRecord record, Context context) {

Check failure on line 65 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.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`. The same applies to public constructors.

UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter

Check failure on line 65 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Avoid unused method parameters such as 'record'. 

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`. The same applies to public constructors.

UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter
String thread = Thread.currentThread().getName();
if (!threadList.contains(thread)) {
threadList.add(thread);
Expand All @@ -70,7 +74,7 @@
}
}

private void processMessageInParallelFailsForFixedMessage(DynamodbEvent.DynamodbStreamRecord record, Context context) {

Check failure on line 77 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.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`. The same applies to public constructors.

UnusedFormalParameter (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#unusedformalparameter
String thread = Thread.currentThread().getName();
if (!threadList.contains(thread)) {
threadList.add(thread);
Expand All @@ -80,7 +84,7 @@
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
if (record.getDynamodb().getSequenceNumber().equals("4421584500000000017450439091")) {

Check failure on line 87 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Position literals first in String comparisons 

Position literals first in all String comparisons, if the second argument is null then NullPointerExceptions
can be avoided, they will just return false. Note that switching literal positions for compareTo and
compareToIgnoreCase may change the result, see examples.

Note that compile-time constant strings are treated like literals. This is because they are inlined into
the class file, are necessarily non-null, and therefore cannot cause an NPE at runtime.

LiteralsFirstInComparisons (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#literalsfirstincomparisons
throw new RuntimeException("fake exception");
}
}
Expand All @@ -88,9 +92,17 @@
private StreamsEventResponse testParallelBatchExecution(DynamodbEvent event,
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Context> messageHandler,
Executor executor) {
return testParallelBatchExecution(event, messageHandler, null, executor);
}

private StreamsEventResponse testParallelBatchExecution(DynamodbEvent event,
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Context> messageHandler,
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Throwable> failureHandler,
Executor executor) {
// Arrange
BatchMessageHandler<DynamodbEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
.withDynamoDbBatchHandler()
.withFailureHandler(failureHandler)
.buildWithRawMessageHandler(messageHandler);

// Act
Expand Down Expand Up @@ -166,7 +178,7 @@
BatchMessageHandler<DynamodbEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
.withDynamoDbBatchHandler()
.withFailureHandler((m, e) -> {
if (m.getDynamodb().getSequenceNumber().equals("4421584500000000017450439091")) {

Check failure on line 181 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Position literals first in String comparisons 

Position literals first in all String comparisons, if the second argument is null then NullPointerExceptions
can be avoided, they will just return false. Note that switching literal positions for compareTo and
compareToIgnoreCase may change the result, see examples.

Note that compile-time constant strings are treated like literals. This is because they are inlined into
the class file, are necessarily non-null, and therefore cannot cause an NPE at runtime.

LiteralsFirstInComparisons (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#literalsfirstincomparisons
wasCalledAndFailed.set(true);
throw new RuntimeException("Success handler throws");
}
Expand All @@ -192,7 +204,7 @@
BatchMessageHandler<DynamodbEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
.withDynamoDbBatchHandler()
.withSuccessHandler((e) -> {
if (e.getDynamodb().getSequenceNumber().equals("4421584500000000017450439091")) {

Check failure on line 207 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Position literals first in String comparisons 

Position literals first in all String comparisons, if the second argument is null then NullPointerExceptions
can be avoided, they will just return false. Note that switching literal positions for compareTo and
compareToIgnoreCase may change the result, see examples.

Note that compile-time constant strings are treated like literals. This is because they are inlined into
the class file, are necessarily non-null, and therefore cannot cause an NPE at runtime.

LiteralsFirstInComparisons (Priority: 1, Ruleset: Best Practices)
https://docs.pmd-code.org/snapshot/pmd_rules_java_bestpractices.html#literalsfirstincomparisons
wasCalledAndFailed.set(true);
throw new RuntimeException("Success handler throws");
}
Expand All @@ -213,7 +225,7 @@
@ParameterizedTest
@Event(value = "dynamo_event_big.json", type = DynamodbEvent.class)
void parallelBatchProcessingWithExecutorSucceedsAndReturns(DynamodbEvent event) {
ExecutorService executor = Executors.newFixedThreadPool(2);

Check failure on line 228 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Ensure that resources like this ExecutorService object are closed after use 

Ensure that resources (like `java.sql.Connection`, `java.sql.Statement`, and `java.sql.ResultSet` objects
and any subtype of `java.lang.AutoCloseable`) are always closed after use.
Failing to do so might result in resource leaks.

Note: It suffices to configure the super type, e.g. `java.lang.AutoCloseable`, so that this rule automatically triggers
on any subtype (e.g. `java.io.FileInputStream`). Additionally specifying `java.sql.Connection` helps in detecting
the types, if the type resolution / auxclasspath is not correctly setup.

Note: Since PMD 6.16.0 the default value for the property `types` contains `java.lang.AutoCloseable` and detects
now cases where the standard `java.io.*Stream` classes are involved. In order to restore the old behaviour,
just remove "AutoCloseable" from the types.

The property `allowedResourceMethodPatterns` can be used to specify method invocation patterns that return
resources which are managed externally and don't need to be closed by the caller. This is useful for
servlet-related streams like `HttpServletRequest.getReader()` or `HttpServletResponse.getWriter()`,
which are managed by the servlet container, and for mocking frameworks, whose mocks are not real
resources even when the mocked type implements `java.lang.AutoCloseable`. The patterns use
InvocationMatcher syntax (e.g., `javax.servlet.ServletRequest#getReader()`).

The defaults cover the servlet API and Mockito's `mock()`. Mocks created by other frameworks
can be allowed by adding their factory methods, for example `org.easymock.EasyMock#createMock(_*)`,
`org.jmock.Mockery#mock(_*)` or `mockit.Mocked#new(_*)`.

Note that `Mockito.spy()` and `Mockito.mockStatic()` are intentionally not among the defaults.
A spy calls through to the real object by default, so either the spy or the object it wraps still
has to be closed; `mockStatic()` returns a `MockedStatic`, which has to be closed as well. A violation
in either case is a true positive.

CloseResource (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#closeresource

StreamsEventResponse dynamodbBatchResponse = testParallelBatchExecution(event, this::processMessageInParallelSucceeds, executor);
executor.shutdown();
Expand All @@ -226,7 +238,7 @@
@ParameterizedTest
@Event(value = "dynamo_event_big.json", type = DynamodbEvent.class)
void parallelBatchProcessingWithExecutor_shouldAddMessageToBatchFailure_whenException_withMessage(DynamodbEvent event) {
ExecutorService executor = Executors.newFixedThreadPool(2);

Check failure on line 241 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java

View workflow job for this annotation

GitHub Actions / pmd_analyse

Ensure that resources like this ExecutorService object are closed after use 

Ensure that resources (like `java.sql.Connection`, `java.sql.Statement`, and `java.sql.ResultSet` objects
and any subtype of `java.lang.AutoCloseable`) are always closed after use.
Failing to do so might result in resource leaks.

Note: It suffices to configure the super type, e.g. `java.lang.AutoCloseable`, so that this rule automatically triggers
on any subtype (e.g. `java.io.FileInputStream`). Additionally specifying `java.sql.Connection` helps in detecting
the types, if the type resolution / auxclasspath is not correctly setup.

Note: Since PMD 6.16.0 the default value for the property `types` contains `java.lang.AutoCloseable` and detects
now cases where the standard `java.io.*Stream` classes are involved. In order to restore the old behaviour,
just remove "AutoCloseable" from the types.

The property `allowedResourceMethodPatterns` can be used to specify method invocation patterns that return
resources which are managed externally and don't need to be closed by the caller. This is useful for
servlet-related streams like `HttpServletRequest.getReader()` or `HttpServletResponse.getWriter()`,
which are managed by the servlet container, and for mocking frameworks, whose mocks are not real
resources even when the mocked type implements `java.lang.AutoCloseable`. The patterns use
InvocationMatcher syntax (e.g., `javax.servlet.ServletRequest#getReader()`).

The defaults cover the servlet API and Mockito's `mock()`. Mocks created by other frameworks
can be allowed by adding their factory methods, for example `org.easymock.EasyMock#createMock(_*)`,
`org.jmock.Mockery#mock(_*)` or `mockit.Mocked#new(_*)`.

Note that `Mockito.spy()` and `Mockito.mockStatic()` are intentionally not among the defaults.
A spy calls through to the real object by default, so either the spy or the object it wraps still
has to be closed; `mockStatic()` returns a `MockedStatic`, which has to be closed as well. A violation
in either case is a true positive.

CloseResource (Priority: 1, Ruleset: Error Prone)
https://docs.pmd-code.org/snapshot/pmd_rules_java_errorprone.html#closeresource

StreamsEventResponse dynamodbBatchResponse = testParallelBatchExecution(event, this::processMessageInParallelFailsForFixedMessage, executor);
executor.shutdown();
Expand All @@ -238,4 +250,23 @@
assertThat(threadList).hasSizeGreaterThan(1);
}

@ParameterizedTest
@Event(value = "dynamo_event_big.json", type = DynamodbEvent.class)
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(DynamodbEvent event) {
int nThreads = 8;
List<String> batchItemFailureIds;
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
try {
StreamsEventResponse dynamodbBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
throw new RuntimeException("fake exception");
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
batchItemFailureIds = dynamodbBatchResponse.getBatchItemFailures().stream().map(
StreamsEventResponse.BatchItemFailure::getItemIdentifier).toList();
} finally {
executor.shutdown();
}
List<String> sequenceNumbers = event.getRecords().stream().map(DynamodbEvent.DynamodbStreamRecord::getDynamodb)
.map(StreamRecord::getSequenceNumber).toList();
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(sequenceNumbers);
}
}
30 changes: 30 additions & 0 deletions ...atch/src/test/java/software/amazon/lambda/powertools/batch/KinesisBatchProcessorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
Expand All @@ -34,6 +36,7 @@
import org.mockito.Mock;
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
import software.amazon.lambda.powertools.batch.model.Product;
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;

class KinesisBatchProcessorTest {

Expand Down Expand Up @@ -89,9 +92,17 @@
private StreamsEventResponse testParallelBatchExecution(KinesisEvent event,
BiConsumer<KinesisEvent.KinesisEventRecord, Context> messageHandler,
Executor executor) {
return testParallelBatchExecution(event, messageHandler, null, executor);
}

private StreamsEventResponse testParallelBatchExecution(KinesisEvent event,
BiConsumer<KinesisEvent.KinesisEventRecord, Context> messageHandler,
BiConsumer<KinesisEvent.KinesisEventRecord, Throwable> failureHandler,
Executor executor) {
// Arrange
BatchMessageHandler<KinesisEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
.withKinesisBatchHandler()
.withFailureHandler(failureHandler)
.buildWithRawMessageHandler(messageHandler);

// Act
Expand All @@ -106,7 +117,7 @@
}

// A handler that throws an exception for _one_ of the deserialized products in the same messages
public void processMessageFailsForFixedProduct(Product product, Context context) {

Check warning on line 120 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/KinesisBatchProcessorTest.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Remove this unused method parameter "context". 


See more on https://sonarcloud.io/project/issues?id=aws-powertools_powertools-lambda-java&issues=AaGM5wd6pr5tEPrTRa3U&open=AaGM5wd6pr5tEPrTRa3U&pullRequest=2735
if (product.getId() == 1234) {
throw new RuntimeException("fake exception");
}
Expand Down Expand Up @@ -268,4 +279,23 @@
assertThat(threadList).hasSizeGreaterThan(1);
}

@ParameterizedTest
@Event(value = "kinesis_event_big.json", type = KinesisEvent.class)
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(KinesisEvent event) {
int nThreads = 8;
List<String> batchItemFailureIds;
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
try {
StreamsEventResponse kinesisBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
throw new RuntimeException("fake exception");
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
batchItemFailureIds = kinesisBatchResponse.getBatchItemFailures().stream().map(
StreamsEventResponse.BatchItemFailure::getItemIdentifier).toList();
} finally {
executor.shutdown();
}
List<String> sequenceNumbers = event.getRecords().stream().map(KinesisEvent.KinesisEventRecord::getKinesis).map(
KinesisEvent.Record::getSequenceNumber).toList();
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(sequenceNumbers);
}
}
30 changes: 29 additions & 1 deletion ...ls-batch/src/test/java/software/amazon/lambda/powertools/batch/SQSBatchProcessorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
Expand All @@ -34,6 +36,7 @@
import org.mockito.Mock;
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
import software.amazon.lambda.powertools.batch.model.Product;
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;

class SQSBatchProcessorTest {
@Mock
Expand Down Expand Up @@ -85,7 +88,7 @@
}

// A handler that throws an exception for _one_ of the deserialized products in the same messages
private void processMessageFailsForFixedProduct(Product product, Context context) {

Check warning on line 91 in powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/SQSBatchProcessorTest.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Remove this unused method parameter "context". 


See more on https://sonarcloud.io/project/issues?id=aws-powertools_powertools-lambda-java&issues=AaGM5wY_pr5tEPrTRa3N&open=AaGM5wY_pr5tEPrTRa3N&pullRequest=2735
if (product.getId() == 12345) {
throw new RuntimeException("fake exception");
}
Expand All @@ -94,9 +97,17 @@
private SQSBatchResponse testParallelBatchExecution(SQSEvent event,
BiConsumer<SQSEvent.SQSMessage, Context> messageHandler,
Executor executor) {
return testParallelBatchExecution(event, messageHandler, null, executor);
}

private SQSBatchResponse testParallelBatchExecution(SQSEvent event,
BiConsumer<SQSEvent.SQSMessage, Context> messageHandler,
BiConsumer<SQSEvent.SQSMessage, Throwable> failureHandler,
Executor executor) {
// Arrange
BatchMessageHandler<SQSEvent, SQSBatchResponse> handler = new BatchMessageHandlerBuilder()
.withSqsBatchHandler()
.withFailureHandler(failureHandler)
.buildWithRawMessageHandler(messageHandler);

// Act
Expand Down Expand Up @@ -276,5 +287,22 @@
assertThat(threadList).hasSizeGreaterThan(1);
}


@ParameterizedTest
@Event(value = "sqs_event_big.json", type = SQSEvent.class)
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(SQSEvent event) {
int nThreads = 8;
List<String> batchItemFailureIds;
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
try {
SQSBatchResponse sqsBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
throw new RuntimeException("fake exception");
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
batchItemFailureIds = sqsBatchResponse.getBatchItemFailures().stream().map(
SQSBatchResponse.BatchItemFailure::getItemIdentifier).toList();
} finally {
executor.shutdown();
}
List<String> messageIds = event.getRecords().stream().map(SQSEvent.SQSMessage::getMessageId).toList();
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(messageIds);
}
}
38 changes: 38 additions & 0 deletions .../java/software/amazon/lambda/powertools/batch/testutils/BatchMessageHandlerTestUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
/*
* Copyright 2023 Amazon.com, Inc. or its affiliates.
* Licensed under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
* http://www.apache.org/licenses/LICENSE-2.0
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package software.amazon.lambda.powertools.batch.testutils;

import java.time.Duration;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;

public final class BatchMessageHandlerTestUtils {

private BatchMessageHandlerTestUtils() {
}

public static <T> BiConsumer<T, Throwable> spinGateFailureHandler(int parties, Duration timeout) {
AtomicInteger arrived = new AtomicInteger(0);
return (item, throwable) -> {
arrived.incrementAndGet();
long deadline = System.nanoTime() + timeout.toNanos();
while (arrived.get() < parties) {
if (System.nanoTime() > deadline) {
throw new IllegalStateException("Gate timed out, is the executor smaller than the batch?");
}
Thread.onSpinWait();
}
};
}
}
Loading