Bug: BatchMessageHandler implementations can fail or lose failed batch item identifiers when using a custom Executor

Expected Behaviour

When BatchMessageHandler#processBatchInParallel is called with a custom Executor and multiple items fail during concurrent processing, the method should return a response containing a list of all failed item identifiers.

Current Behaviour

The method may fail with java.util.concurrent.CompletionException, caused by java.lang.ArrayIndexOutOfBoundsException, or the returned response may contain an incomplete list of failed batch item identifiers.

The issue is caused by concurrent additions of identifiers to an ArrayList without synchronization.

The same issue is present in all three implementations: SqsBatchMessageHandler, DynamoDbBatchMessageHandler, and KinesisStreamsBatchMessageHandler.

The likelihood of this occurring in production depends on several factors, since the issue requires conflicting concurrent modifications of the batchItemFailures list. A high degree of concurrency (e.g., when using a lightweight-thread-based executor, such as one based on Java virtual threads), larger batches, short-running tasks, and a high failure rate all increase the likelihood of the issue occurring.

Code snippet

package com.example;

import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.events.SQSBatchResponse;
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
import org.junit.jupiter.api.RepeatedTest;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.slf4j.LoggerFactory;
import software.amazon.lambda.powertools.batch.BatchMessageHandlerBuilder;
import software.amazon.lambda.powertools.batch.handler.BatchMessageHandler;
import software.amazon.lambda.powertools.batch.handler.SqsBatchMessageHandler;
import software.amazon.lambda.powertools.batch.internal.MultiThreadMDC;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.BiConsumer;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.BDDMockito.then;
import static org.mockito.BDDMockito.willThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;

@ExtendWith(MockitoExtension.class)
class BatchMessageHandlerTests {

    static {
        // logging slows everything down and reduces the likelihood of race conditions
        ((Logger) LoggerFactory.getLogger(MultiThreadMDC.class)).setLevel(Level.INFO);
        ((Logger) LoggerFactory.getLogger(SqsBatchMessageHandler.class)).setLevel(Level.OFF);
    }

    private static final int TEST_REPETITIONS = 150;
    private static final int BATCH_SIZE = 15;

    // private static final Executor EXECUTOR = Executors.newVirtualThreadPerTaskExecutor();
    private static final Executor EXECUTOR = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("virtual-", 0).factory());
    // private static final Executor EXECUTOR = Executors.newFixedThreadPool(16);

    // private static final Executor EXECUTOR = Executors.newFixedThreadPool(1);

    @Mock
    private Context context;

    @RepeatedTest(TEST_REPETITIONS)
    void testSqsBatchHandler() {
        // using a mock for the handler so we can assert all invocations
        BiConsumer<SQSEvent.SQSMessage, Context> rawMessageHandler = mock();
        BatchMessageHandler<SQSEvent, SQSBatchResponse> handler = new BatchMessageHandlerBuilder().withSqsBatchHandler()
                .buildWithRawMessageHandler(rawMessageHandler);
        RuntimeException exception = new RuntimeException();
        // fail every message
        willThrow(exception).given(rawMessageHandler).accept(any(), any());
        // generate a list of messages
        List<SQSEvent.SQSMessage> messages = IntStream.range(0, BATCH_SIZE)
                .mapToObj(BatchMessageHandlerTests::createSQSEvent).toList();
        SQSEvent event = new SQSEvent();
        event.setRecords(messages);
        SQSBatchResponse result;
        try {
            result = handler.processBatchInParallel(event, this.context, EXECUTOR);
        }
        catch (CompletionException e) {
            if (isMultiThreadMdcIssue(e)) {
                fail("Failure cause by MultiThreadMDC", e);
            }
            throw e;
        }
        // assert that the handler was invoked for every message
        ArgumentCaptor<SQSEvent.SQSMessage> messageCaptor = ArgumentCaptor.forClass(SQSEvent.SQSMessage.class);
        then(rawMessageHandler).should(times(BATCH_SIZE)).accept(messageCaptor.capture(), same(this.context));
        assertThat(messageCaptor.getAllValues()).containsExactlyInAnyOrderElementsOf(messages);
        // assert that the batch item failures contains all (failed) message IDs
        assertThat(result.getBatchItemFailures()).extracting(SQSBatchResponse.BatchItemFailure::getItemIdentifier)
                .containsExactlyInAnyOrderElementsOf(
                        messages.stream().map(BatchMessageHandlerTests::extractMessageId).toList());
    }

    private static SQSEvent.SQSMessage createSQSEvent(int messageId) {
        SQSEvent.SQSMessage message = new SQSEvent.SQSMessage();
        message.setMessageId(String.valueOf(messageId));
        message.setAttributes(Collections.emptyMap());
        return message;
    }

    private static String extractMessageId(SQSEvent.SQSMessage message) {
        return message.getMessageId();
    }

    private static boolean isMultiThreadMdcIssue(CompletionException e) {
        return e.getCause() instanceof ArrayIndexOutOfBoundsException oob && Stream.of(oob.getStackTrace())
                .anyMatch(ste -> ste.toString().contains("MultiThreadMDC.removeThread")
                        || ste.toString().contains("MultiThreadMDC.copyMDCToThread"));
    }
}

Possible Solution

Replicate the pattern used in the other overloaded processBatchInParallel method, which does not accept a custom Executor.

Use an AtomicReference inside the CompletableFuture.supplyAsync supplier, and have the code calling processBatchItem(eventRecord, context) set the AtomicReference instead of adding the result to the shared ArrayList.

The type of the futures variable would change from List<CompletableFuture<Void>> to List<CompletableFuture<Optional<BatchItemFailure>>>, where BatchItemFailure is the concrete return type of the corresponding processBatchItem implementation.

    @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();
        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);
+       List<SQSBatchResponse.BatchItemFailure> batchItemFailures = futures.stream()
+               .map(CompletableFuture::join)
+               .filter(Optional::isPresent)
+               .map(Optional::get)
+               .collect(Collectors.toList());

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

Or simply use a thread-safe data structure, e.g., ConcurrentHashMap.newKeySet().

Steps to Reproduce

Environment:

JDK: Temurin 25.0.4
Gradle: 9.2.0
Powertools Lambda Java Batch: 2.11.0

Reproducing Code

Example code (Reproducer.zip) tests SqsBatchMessageHandler, but similar code can be written to test the other two implementations.

The code uses several strategies to make the issue more likely to occur:

Reduces the log levels of MultiThreadMDC and SqsBatchMessageHandler, since writing logs to the console is slow
Repeats the test hundreds of times to increase the chance of at least one iteration failing
Generates a "large" batch to provide enough work for concurrent processing
Uses a virtual-thread executor to reduce the overhead of creating a large number of threads
Uses named virtual threads to reduce the likelihood of issues stemming from concurrency problems in MultiThreadMDC (to be reported separately)

Running the Reproducer

Run the tests from the project root:

./gradlew clean check

Failure Rate

The issue is intermittent. If the tests do not fail, try increasing the batch size and/or the number of test repetitions.

I observed different results depending on how the tests were run: from the console (./gradlew check), using IntelliJ's Run configuration, or using IntelliJ's Debug configuration.

Enabling logging reduces the likelihood of reproducing the issue, but it can still occur if the tests are run enough times.

When running the checks from the console with TEST_REPETITIONS=100, BATCH_SIZE=10, and logging disabled for both components, I observed a 20% failure rate.

Enabling logging for SqsBatchMessageHandler significantly reduced the failure rate. I had to increase BATCH_SIZE to 200 and TEST_REPETITIONS to 200 to observe a failure rate of approximately 1%.

Other Issues

Tests may also fail with CompletionException / ArrayIndexOutOfBoundsException caused by concurrency issues in MultiThreadMDC. And in some cases, the resulting stack trace is too terse to determine which part of the code failed.

Powertools for AWS Lambda (Java) version

2.11.0

AWS Lambda function runtime

Java 25

Debugging logs

// Test failure where `batchItemFailures::add` does not throw an exception but some failed item identifiers are lost:

java.lang.AssertionError:
Expecting actual:
  [...]
to contain exactly in any order:
  [...]
but could not find the following elements:
  ["55", "89", "123", "147"]


// Test failure where `batchItemFailures::add` throws an exception:

java.lang.ArrayIndexOutOfBoundsException: Index 16 out of bounds for length 15
java.util.concurrent.CompletionException: java.lang.ArrayIndexOutOfBoundsException: Index 16 out of bounds for length 15
	at java.base/java.util.concurrent.CompletableFuture.wrapInCompletionException(CompletableFuture.java:323)
	at java.base/java.util.concurrent.CompletableFuture.encodeThrowable(CompletableFuture.java:359)
	at java.base/java.util.concurrent.CompletableFuture.completeThrowable(CompletableFuture.java:364)
	at java.base/java.util.concurrent.CompletableFuture$AsyncRun.run(CompletableFuture.java:1828)
	at java.base/java.util.concurrent.ThreadPerTaskExecutor$TaskRunner.run(ThreadPerTaskExecutor.java:291)
	at java.base/java.lang.VirtualThread.run(VirtualThread.java:460)
Caused by: java.lang.ArrayIndexOutOfBoundsException: Index 16 out of bounds for length 15
	at java.base/java.util.ArrayList.add(ArrayList.java:485)
	at java.base/java.util.ArrayList.add(ArrayList.java:497)
	at java.base/java.util.Optional.ifPresent(Optional.java:178)
	at software.amazon.lambda.powertools.batch.handler.SqsBatchMessageHandler.lambda$processBatchInParallel$4(SqsBatchMessageHandler.java:158)
	at software.amazon.lambda.powertools.batch.internal.XRayTraceEntityPropagator.runWithEntity(XRayTraceEntityPropagator.java:71)
	at software.amazon.lambda.powertools.batch.handler.SqsBatchMessageHandler.lambda$processBatchInParallel$5(SqsBatchMessageHandler.java:153)
	at java.base/java.util.concurrent.CompletableFuture$AsyncRun.run(CompletableFuture.java:1825)
	... 2 more

Thanks for opening your first issue here! We'll come back to you as soon as we can.
In the meantime, check out our GitHub Discussions: Join the conversation

@bubanko

Hi @bubanko,

thank you for bringing this issue to our attention. Let me investigate and try to reproduce this.

@bubanko

Thanks for the detailed report and the reproducer, @bubanko. I reproduced both issues locally on main and confirmed them as bugs.

1. Lost batch item failures with a custom Executor

The processBatchInParallel(event, context, executor) overload in all three handlers adds failures to a shared ArrayList from several threads at once. I ran 200 failing messages on a 32-thread pool, 200 times per handler. SQS returned an incomplete failure list in 3 runs and Kinesis in 1. A missing ID tells Lambda that the message succeeded, so SQS deletes it and Kinesis or DynamoDB moves the checkpoint past it.

I agree with your fix:

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 {
                    result.set(processBatchItem(eventRecord, context));
                } finally {
                    multiThreadMDC.removeThread(Thread.currentThread().getName());
                }
            });
            return result.get();
        }, executor))
        .collect(Collectors.toList());

List<SQSBatchResponse.BatchItemFailure> batchItemFailures = futures.stream()
        .map(CompletableFuture::join)
        .flatMap(Optional::stream)
        .collect(Collectors.toList());

2. ArrayIndexOutOfBoundsException in MultiThreadMDC

MultiThreadMDC tracks threads in a plain ArrayList that all worker threads share. In my runs, Kinesis failed twice in MultiThreadMDC.removeThread, and the error failed the whole batch. The parallelStream overload uses the same class, so it has the same race. I agree with your proposal to use a thread-safe data structure, something like this:

private final Set<String> mdcAwareThreads = ConcurrentHashMap.newKeySet();

public void copyMDCToThread(String thread) {
    if (mdcAwareThreads.add(thread)) {
        MDC.setContextMap(contextMap);
    }
}

public void removeThread(String thread) {
    mdcAwareThreads.remove(thread);
}

Would you be interested in contributing to this project and opening a PR to address these fixes?

@phipag

Hi @phipag!

Thanks for verifying!

I was actually planning on opening a separate issue for MultiThreadMDC, but maybe it doesn't make sense since it's related and already mentioned here.

I'm willing to work on this; Fixes are straightforward enough. I just need to see what to do with regards to tests.

My plan then would be to open a PR which fixes the MultiThreadMDC first, since it can be tested in isolation. And then open a separate PR for the three BatchMessageHandler implementations.

Makes sense?

@bubanko

Hey @bubanko,

yes, the plan sounds good. I am okay with this. Let's do two isolated PRs.

My proposal is to try if you can create a test that fails consistently (not something relying on randomness) before the fix and passes after the fix. Potentially, if you can build a scenario with artificially interleaved threads you can come up with a reproducible stable test. I believe, for these kind of tests this is is the hard part.

I can also help with validating it at runtime in my AWS account and run the E2E tests in case you don't have an easy way to test in an AWS dev account.

We do something like this using CountDownLatch here:

powertools-lambda-java/powertools-lambda-metadata/src/test/java/software/amazon/lambda/powertools/metadata/LambdaMetadataClientConcurrencyTest.java

Lines 51 to 83 in d767ff5

    
           void get_shouldBeThreadSafe() throws Exception { 
        
               // Given 
        
               LambdaMetadata metadata = new LambdaMetadata("use1-az1"); 
        
               when(mockHttpClient.fetchMetadata()).thenReturn(metadata); 
        
               int threadCount = 10; 
        
               ExecutorService executor = Executors.newFixedThreadPool(threadCount); 
        
               CountDownLatch startLatch = new CountDownLatch(1); 
        
               List<Future<LambdaMetadata>> futures = new ArrayList<>(); 
        
               // When - all threads try to get metadata simultaneously 
        
               for (int i = 0; i < threadCount; i++) { 
        
                   futures.add(executor.submit(() -> { 
        
                       startLatch.await(); 
        
                       return LambdaMetadataClient.get(); 
        
                   })); 
        
               } 
        
               startLatch.countDown(); 
        
               // Then - all threads should get the same instance 
        
               LambdaMetadata firstResult = null; 
        
               for (Future<LambdaMetadata> future : futures) { 
        
                   LambdaMetadata result = future.get(5, TimeUnit.SECONDS); 
        
                   if (firstResult == null) { 
        
                       firstResult = result; 
        
                   } 
        
                   assertThat(result).isSameAs(firstResult); 
        
                   assertThat(result.getAvailabilityZoneId()).isEqualTo("use1-az1"); 
        
               } 
        
               executor.shutdown(); 
        
               executor.awaitTermination(5, TimeUnit.SECONDS); 
        
           }

fix(batch): use thread-safe collection in MultiThreadMDC

@phipag

@phipag,

I would appreciate help with validating at runtime!

I opened a PR for the MultiThreadMDC issue and will work on the BatchMessageHandler implementations next. Hopefully, testing won’t be too much of an issue.

refactor(batch): make MultiThreadMDC methods more testable

fix(batch): use thread-safe collection in `MultiThreadMDC` (#2717)

@bubanko

Thanks @bubanko,

your first PR is merged. Looking forward to the next one. Super good work, congrats! 🥳

@phipag

Hey @phipag,

thanks for the support! 🙂

As for the second PR, logging is turning out to be the biggest obstacle to reproducibility. There are two LOGGER.error statements in the catch block, both executed "just" before the failed message ID is added to the shared ArrayList. And since logging introduces synchronization between threads, it reduces the likelihood of contention. While failure is still possible, it is far less consistent.

My first attempt was to (ab)use the fact that failureHandler is called after these log statements, adding a CyclicBarrier there to release all threads simultaneously. This works, but it’s not entirely reliable — I still get an occasional passing test. Running the test through IntelliJ, either in Run or Debug mode, makes it even less consistent than running it from the terminal. You can find the code here.

My second approach involves disabling logging. However, I didn’t want to do this at the configuration level, as that would disable logging for all BatchMessageHandler implementations across the module's tests. Instead, my PoC extends SimpleLogger to support runtime log-level overrides. This makes the failure very reliable, at least for me, but it does involve additional testing infrastructure code (see the PoC commit here).

This second approach also has the benefit of keeping the build logs manageable when a test fails, since exception stack traces are otherwise printed for thousands of messages.

Let me know what you think of the two options, or if you have any alternative ideas.

@bubanko

Thanks @bubanko, your analysis of the logging matches my own analysis as well. I came to a similar solution with the CyclicBarrier at first and ran a few more tests and simulations.

Eventually, I came to this solution with a spin gate. It releases all threads at the same instant, so it works with logging switched on and needs no custom logger:

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

I applied it to withFailureHandler(spinGate(new AtomicInteger(), BATCH_SIZE)), with a batch of 8 items and a pool of 8 threads, and repeat the batch a few times in the test. On my machine, the test failed in 10 of 10 runs on main. With the fix in place it passed every time.

I suggest we use this instead of the logging override, so we don't have to maintain a custom SLF4J provider. What do you think?

Nice! It didn't occur to me to try a busy wait. If it works consistently I'm definitely for it! Especially if it also works with one of the existing event examples (like sqs_event_big.json).

I'll work on this and if all goes well, I'll open a PR.

Navigation Menu

Repository navigation

Bug: `BatchMessageHandler` implementations can fail or lose failed batch item identifiers when using a custom `Executor` #2712

Expected Behaviour

Current Behaviour

Code snippet

Possible Solution

Steps to Reproduce

Environment:

Reproducing Code

Running the Reproducer

Failure Rate

Other Issues

Powertools for AWS Lambda (Java) version

AWS Lambda function runtime

Debugging logs

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	void get_shouldBeThreadSafe() throws Exception {
	// Given
	LambdaMetadata metadata = new LambdaMetadata("use1-az1");
	when(mockHttpClient.fetchMetadata()).thenReturn(metadata);

	int threadCount = 10;
	ExecutorService executor = Executors.newFixedThreadPool(threadCount);
	CountDownLatch startLatch = new CountDownLatch(1);
	List<Future<LambdaMetadata>> futures = new ArrayList<>();

	// When - all threads try to get metadata simultaneously
	for (int i = 0; i < threadCount; i++) {
	futures.add(executor.submit(() -> {
	startLatch.await();
	return LambdaMetadataClient.get();
	}));
	}
	startLatch.countDown();

	// Then - all threads should get the same instance
	LambdaMetadata firstResult = null;
	for (Future<LambdaMetadata> future : futures) {
	LambdaMetadata result = future.get(5, TimeUnit.SECONDS);
	if (firstResult == null) {
	firstResult = result;
	}
	assertThat(result).isSameAs(firstResult);
	assertThat(result.getAvailabilityZoneId()).isEqualTo("use1-az1");
	}

	executor.shutdown();
	executor.awaitTermination(5, TimeUnit.SECONDS);
	}

Bug: BatchMessageHandler implementations can fail or lose failed batch item identifiers when using a custom Executor #2712

Description

Expected Behaviour

Current Behaviour

Code snippet

Possible Solution

Steps to Reproduce

Environment:

Reproducing Code

Running the Reproducer

Failure Rate

Other Issues

Powertools for AWS Lambda (Java) version

AWS Lambda function runtime

Debugging logs

Activity

powertools-for-aws-oss-automation commented on Oct 25, 2026

phipag commented on Oct 26, 2026

phipag commented on Oct 26, 2026

bubanko commented on Oct 26, 2026

phipag commented on Oct 26, 2026

phipag commented on Oct 26, 2026

bubanko commented on Oct 27, 2026

phipag commented on Oct 29, 2026

bubanko commented on Oct 29, 2026

phipag commented on Oct 29, 2026

bubanko commented on Oct 29, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Bug: `BatchMessageHandler` implementations can fail or lose failed batch item identifiers when using a custom `Executor` #2712