Skip to content

Commit 4c2d32a

Browse files
bubankobubanko
committed
fix(batch): make batch failure tracking thread-safe for parallel processing
Replicate the thread-safe `BatchItemFailure` collection implementation from the `processBatchInParallel(event, context)` method to the overloaded variant that additionally accepting an executor. Apply this fix to all three `BatchMessageHandler` implementations. Using a shared, unsynchronized `ArrayList` to collect `BatchItemFailure` instances can lead to lost failure items or cause a worker thread to fail with an `ArrayOutOfBoundsException` due to concurrent modifications corrupting the list's internals state. GitHub issue number: #2712
1 parent 6455f45 commit 4c2d32a

7 files changed

Lines changed: 153 additions & 24 deletions

File tree

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

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

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

17-
import java.util.ArrayList;
1817
import java.util.List;
1918
import java.util.Optional;
2019
import java.util.concurrent.CompletableFuture;
@@ -99,22 +98,22 @@ public StreamsEventResponse processBatchInParallel(DynamodbEvent event, Context
9998
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
10099
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();
101100

102-
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
103-
List<CompletableFuture<Void>> futures = event.getRecords().stream()
104-
.map(eventRecord -> CompletableFuture.runAsync(() -> {
101+
List<CompletableFuture<Optional<StreamsEventResponse.BatchItemFailure>>> futures = event.getRecords().stream()
102+
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
103+
AtomicReference<Optional<StreamsEventResponse.BatchItemFailure>> result = new AtomicReference<>();
105104
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
106105
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
107106
try {
108-
Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
109-
context);
110-
failureOpt.ifPresent(batchItemFailures::add);
107+
result.set(processBatchItem(eventRecord, context));
111108
} finally {
112109
multiThreadMDC.removeThread(Thread.currentThread().getName());
113110
}
114111
});
112+
return result.get();
115113
}, executor))
116114
.collect(Collectors.toList());
117-
futures.forEach(CompletableFuture::join);
115+
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
116+
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).collect(Collectors.toList());
118117
return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
119118
}
120119

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

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

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

17-
import java.util.ArrayList;
1817
import java.util.List;
1918
import java.util.Optional;
2019
import java.util.concurrent.CompletableFuture;
@@ -109,22 +108,22 @@ public StreamsEventResponse processBatchInParallel(KinesisEvent event, Context c
109108
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
110109
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();
111110

112-
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
113-
List<CompletableFuture<Void>> futures = event.getRecords().stream()
114-
.map(eventRecord -> CompletableFuture.runAsync(() -> {
111+
List<CompletableFuture<Optional<StreamsEventResponse.BatchItemFailure>>> futures = event.getRecords().stream()
112+
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
113+
AtomicReference<Optional<StreamsEventResponse.BatchItemFailure>> result = new AtomicReference<>();
115114
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
116115
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
117116
try {
118-
Optional<StreamsEventResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
119-
context);
120-
failureOpt.ifPresent(batchItemFailures::add);
117+
result.set(processBatchItem(eventRecord, context));
121118
} finally {
122119
multiThreadMDC.removeThread(Thread.currentThread().getName());
123120
}
124121
});
122+
return result.get();
125123
}, executor))
126124
.collect(Collectors.toList());
127-
futures.forEach(CompletableFuture::join);
125+
List<StreamsEventResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
126+
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).collect(Collectors.toList());
128127
return StreamsEventResponse.builder().withBatchItemFailures(batchItemFailures).build();
129128
}
130129

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -147,22 +147,22 @@ public SQSBatchResponse processBatchInParallel(SQSEvent event, Context context,
147147
MultiThreadMDC multiThreadMDC = new MultiThreadMDC();
148148
Object capturedSubsegment = XRayTraceEntityPropagator.captureTraceEntity();
149149

150-
List<SQSBatchResponse.BatchItemFailure> batchItemFailures = new ArrayList<>();
151-
List<CompletableFuture<Void>> futures = event.getRecords().stream()
152-
.map(eventRecord -> CompletableFuture.runAsync(() -> {
150+
List<CompletableFuture<Optional<SQSBatchResponse.BatchItemFailure>>> futures = event.getRecords().stream()
151+
.map(eventRecord -> CompletableFuture.supplyAsync(() -> {
152+
AtomicReference<Optional<SQSBatchResponse.BatchItemFailure>> result = new AtomicReference<>();
153153
XRayTraceEntityPropagator.runWithEntity(capturedSubsegment, () -> {
154154
multiThreadMDC.copyMDCToThread(Thread.currentThread().getName());
155155
try {
156-
Optional<SQSBatchResponse.BatchItemFailure> failureOpt = processBatchItem(eventRecord,
157-
context);
158-
failureOpt.ifPresent(batchItemFailures::add);
156+
result.set(processBatchItem(eventRecord, context));
159157
} finally {
160158
multiThreadMDC.removeThread(Thread.currentThread().getName());
161159
}
162160
});
161+
return result.get();
163162
}, executor))
164163
.collect(Collectors.toList());
165-
futures.forEach(CompletableFuture::join);
164+
List<SQSBatchResponse.BatchItemFailure> batchItemFailures = futures.stream().map(
165+
CompletableFuture::join).filter(Optional::isPresent).map(Optional::get).collect(Collectors.toList());
166166

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

‎powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/DdbBatchProcessorTest.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@
1919
import com.amazonaws.services.lambda.runtime.Context;
2020
import com.amazonaws.services.lambda.runtime.events.DynamodbEvent;
2121
import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
22+
import com.amazonaws.services.lambda.runtime.events.models.dynamodb.StreamRecord;
2223
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;
24+
25+
import java.time.Duration;
2326
import java.util.ArrayList;
2427
import java.util.Collections;
2528
import java.util.List;
@@ -28,11 +31,13 @@
2831
import java.util.concurrent.Executors;
2932
import java.util.concurrent.atomic.AtomicBoolean;
3033
import java.util.function.BiConsumer;
34+
import java.util.stream.Collectors;
3135

3236
import org.junit.jupiter.api.AfterEach;
3337
import org.junit.jupiter.params.ParameterizedTest;
3438
import org.mockito.Mock;
3539
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
40+
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;
3641

3742
class DdbBatchProcessorTest {
3843

@@ -88,9 +93,17 @@ private void processMessageInParallelFailsForFixedMessage(DynamodbEvent.Dynamodb
8893
private StreamsEventResponse testParallelBatchExecution(DynamodbEvent event,
8994
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Context> messageHandler,
9095
Executor executor) {
96+
return testParallelBatchExecution(event, messageHandler, null, executor);
97+
}
98+
99+
private StreamsEventResponse testParallelBatchExecution(DynamodbEvent event,
100+
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Context> messageHandler,
101+
BiConsumer<DynamodbEvent.DynamodbStreamRecord, Throwable> failureHandler,
102+
Executor executor) {
91103
// Arrange
92104
BatchMessageHandler<DynamodbEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
93105
.withDynamoDbBatchHandler()
106+
.withFailureHandler(failureHandler)
94107
.buildWithRawMessageHandler(messageHandler);
95108

96109
// Act
@@ -238,4 +251,23 @@ void parallelBatchProcessingWithExecutor_shouldAddMessageToBatchFailure_whenExce
238251
assertThat(threadList).hasSizeGreaterThan(1);
239252
}
240253

254+
@ParameterizedTest
255+
@Event(value = "dynamo_event_big.json", type = DynamodbEvent.class)
256+
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(DynamodbEvent event) {
257+
int nThreads = 8;
258+
List<String> batchItemFailureIds;
259+
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
260+
try {
261+
StreamsEventResponse dynamodbBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
262+
throw new RuntimeException("fake exception");
263+
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
264+
batchItemFailureIds = dynamodbBatchResponse.getBatchItemFailures().stream().map(
265+
StreamsEventResponse.BatchItemFailure::getItemIdentifier).collect(Collectors.toList());
266+
} finally {
267+
executor.shutdown();
268+
}
269+
List<String> sequenceNumbers = event.getRecords().stream().map(DynamodbEvent.DynamodbStreamRecord::getDynamodb)
270+
.map(StreamRecord::getSequenceNumber).collect(Collectors.toList());
271+
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(sequenceNumbers);
272+
}
241273
}

‎powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/KinesisBatchProcessorTest.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import com.amazonaws.services.lambda.runtime.events.KinesisEvent;
2121
import com.amazonaws.services.lambda.runtime.events.StreamsEventResponse;
2222
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;
23+
24+
import java.time.Duration;
2325
import java.util.ArrayList;
2426
import java.util.Collections;
2527
import java.util.List;
@@ -28,12 +30,14 @@
2830
import java.util.concurrent.Executors;
2931
import java.util.concurrent.atomic.AtomicBoolean;
3032
import java.util.function.BiConsumer;
33+
import java.util.stream.Collectors;
3134

3235
import org.junit.jupiter.api.AfterEach;
3336
import org.junit.jupiter.params.ParameterizedTest;
3437
import org.mockito.Mock;
3538
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
3639
import software.amazon.lambda.powertools.batch.model.Product;
40+
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;
3741

3842
class KinesisBatchProcessorTest {
3943

@@ -89,9 +93,17 @@ private void processMessageInParallelFailsForFixedMessage(KinesisEvent.KinesisEv
8993
private StreamsEventResponse testParallelBatchExecution(KinesisEvent event,
9094
BiConsumer<KinesisEvent.KinesisEventRecord, Context> messageHandler,
9195
Executor executor) {
96+
return testParallelBatchExecution(event, messageHandler, null, executor);
97+
}
98+
99+
private StreamsEventResponse testParallelBatchExecution(KinesisEvent event,
100+
BiConsumer<KinesisEvent.KinesisEventRecord, Context> messageHandler,
101+
BiConsumer<KinesisEvent.KinesisEventRecord, Throwable> failureHandler,
102+
Executor executor) {
92103
// Arrange
93104
BatchMessageHandler<KinesisEvent, StreamsEventResponse> handler = new BatchMessageHandlerBuilder()
94105
.withKinesisBatchHandler()
106+
.withFailureHandler(failureHandler)
95107
.buildWithRawMessageHandler(messageHandler);
96108

97109
// Act
@@ -268,4 +280,23 @@ void batchProcessingInParallelWithExecutor_shouldAddMessageToBatchFailure_whenEx
268280
assertThat(threadList).hasSizeGreaterThan(1);
269281
}
270282

283+
@ParameterizedTest
284+
@Event(value = "kinesis_event_big.json", type = KinesisEvent.class)
285+
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(KinesisEvent event) {
286+
int nThreads = 8;
287+
List<String> batchItemFailureIds;
288+
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
289+
try {
290+
StreamsEventResponse kinesisBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
291+
throw new RuntimeException("fake exception");
292+
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
293+
batchItemFailureIds = kinesisBatchResponse.getBatchItemFailures().stream().map(
294+
StreamsEventResponse.BatchItemFailure::getItemIdentifier).collect(Collectors.toList());
295+
} finally {
296+
executor.shutdown();
297+
}
298+
List<String> sequenceNumbers = event.getRecords().stream().map(KinesisEvent.KinesisEventRecord::getKinesis).map(
299+
KinesisEvent.Record::getSequenceNumber).collect(Collectors.toList());
300+
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(sequenceNumbers);
301+
}
271302
}

‎powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/SQSBatchProcessorTest.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
2121
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
2222
import com.amazonaws.services.lambda.runtime.tests.annotations.Event;
23+
24+
import java.time.Duration;
2325
import java.util.ArrayList;
2426
import java.util.Collections;
2527
import java.util.List;
@@ -28,12 +30,14 @@
2830
import java.util.concurrent.Executors;
2931
import java.util.concurrent.atomic.AtomicBoolean;
3032
import java.util.function.BiConsumer;
33+
import java.util.stream.Collectors;
3134

3235
import org.junit.jupiter.api.AfterEach;
3336
import org.junit.jupiter.params.ParameterizedTest;
3437
import org.mockito.Mock;
3538
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
3639
import software.amazon.lambda.powertools.batch.model.Product;
40+
import software.amazon.lambda.powertools.batch.testutils.BatchMessageHandlerTestUtils;
3741

3842
class SQSBatchProcessorTest {
3943
@Mock
@@ -94,9 +98,17 @@ private void processMessageFailsForFixedProduct(Product product, Context context
9498
private SQSBatchResponse testParallelBatchExecution(SQSEvent event,
9599
BiConsumer<SQSEvent.SQSMessage, Context> messageHandler,
96100
Executor executor) {
101+
return testParallelBatchExecution(event, messageHandler, null, executor);
102+
}
103+
104+
private SQSBatchResponse testParallelBatchExecution(SQSEvent event,
105+
BiConsumer<SQSEvent.SQSMessage, Context> messageHandler,
106+
BiConsumer<SQSEvent.SQSMessage, Throwable> failureHandler,
107+
Executor executor) {
97108
// Arrange
98109
BatchMessageHandler<SQSEvent, SQSBatchResponse> handler = new BatchMessageHandlerBuilder()
99110
.withSqsBatchHandler()
111+
.withFailureHandler(failureHandler)
100112
.buildWithRawMessageHandler(messageHandler);
101113

102114
// Act
@@ -276,5 +288,23 @@ void parallelBatchProcessingWithExecutor_shouldAddMessageToBatchFailure_whenExce
276288
assertThat(threadList).hasSizeGreaterThan(1);
277289
}
278290

279-
291+
@ParameterizedTest
292+
@Event(value = "sqs_event_big.json", type = SQSEvent.class)
293+
void parallelBatchProcessing_whenCollectingBatchItemFailures_shouldBeThreadSafe(SQSEvent event) {
294+
int nThreads = 8;
295+
List<String> batchItemFailureIds;
296+
ExecutorService executor = Executors.newFixedThreadPool(nThreads);
297+
try {
298+
SQSBatchResponse sqsBatchResponse = testParallelBatchExecution(event, (message, ctx) -> {
299+
throw new RuntimeException("fake exception");
300+
}, BatchMessageHandlerTestUtils.spinGateFailureHandler(nThreads, Duration.ofSeconds(5)), executor);
301+
batchItemFailureIds = sqsBatchResponse.getBatchItemFailures().stream().map(
302+
SQSBatchResponse.BatchItemFailure::getItemIdentifier).collect(Collectors.toList());
303+
} finally {
304+
executor.shutdown();
305+
}
306+
List<String> messageIds = event.getRecords().stream().map(SQSEvent.SQSMessage::getMessageId)
307+
.collect(Collectors.toList());
308+
assertThat(batchItemFailureIds).containsExactlyInAnyOrderElementsOf(messageIds);
309+
}
280310
}

‎powertools-batch/src/test/java/software/amazon/lambda/powertools/batch/testutils/BatchMessageHandlerTestUtils.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
package software.amazon.lambda.powertools.batch.testutils;
15+
16+
import java.time.Duration;
17+
import java.util.concurrent.atomic.AtomicInteger;
18+
import java.util.function.BiConsumer;
19+
20+
public final class BatchMessageHandlerTestUtils {
21+
22+
private BatchMessageHandlerTestUtils() {
23+
}
24+
25+
public static <T> BiConsumer<T, Throwable> spinGateFailureHandler(int parties, Duration timeout) {
26+
AtomicInteger arrived = new AtomicInteger(0);
27+
return (item, throwable) -> {
28+
arrived.incrementAndGet();
29+
long deadline = System.nanoTime() + timeout.toNanos();
30+
while (arrived.get() < parties) {
31+
if (System.nanoTime() > deadline) {
32+
throw new IllegalStateException("Gate timed out, is the executor smaller than the batch?");
33+
}
34+
Thread.onSpinWait();
35+
}
36+
};
37+
}
38+
}

0 commit comments

Comments
 (0)