Fix SqsAsyncBatchManager excessive batch flushing under heavy load (#6374)

[thornhillcody]

• Synchronize replacement of scheduled flush in RequestBatchBuffer to prevent orphaned, uncanceled scheduled flushes

Motivation and Context

This change enables the SqsAsyncBatchManager to make batches of correct size during high concurrency. More info in #6374

Modifications

I chose to synchronize the cancel and assignment of the scheduledFlush in RequestBatchBuffer.putScheduledFlush(). This is the only place this is used and it's an internal API, so I felt comfortable changing the name to make it more clear what it's doing.

Testing

Ran sample app used to reproduce the issue (found in bug report) and verified the average batch sizes were as expected with 500 TPS, 1000 TPS, 5000 TPS.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)

Checklist

• I have read the CONTRIBUTING document
• Local run of mvn install succeeds
• My code follows the code style of this project
• My change requires a change to the Javadoc documentation
• I have updated the Javadoc documentation accordingly
• I have added tests to cover my changes
• All new and existing tests passed
• I have added a changelog entry. Adding a new entry must be accomplished by running the scripts/new-change script and following the instructions. Commit the new file created by the script in .changes/next-release with your changes.
• My change is to implement 1.11 parity feature and I have updated LaunchChangelog

License

• I confirm that this pull request can be released under the Apache 2 license

[joviegas]

Can we add a new Junit Test class testing end to end something like

import static com.github.tomakehurst.wiremock.client.WireMock.*;
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
import static org.assertj.core.api.Assertions.assertThat;

import com.github.tomakehurst.wiremock.junit5.WireMockExtension;
import com.github.tomakehurst.wiremock.verification.LoggedRequest;
import java.net.URI;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.services.sqs.SqsAsyncClient;

public class HighThroughputBatchingTest {

    private static final String QUEUE_URL = "https://sqs.us-east-1.amazonaws.com/123456789012/MyQueue";
    private static final int MESSAGE_COUNT = 5000;
    private static final int CONCURRENT_THREADS = 50;
    private static final int MAX_BATCH_SIZE = 10;

    @RegisterExtension
    static WireMockExtension wireMock = WireMockExtension.newInstance()
                                                         .options(wireMockConfig().dynamicPort())
                                                         .configureStaticDsl(true)
                                                         .build();

    private SqsAsyncClient client;
    private SqsAsyncBatchManager batchManager;

    @BeforeEach
    void setUp() {
        client = SqsAsyncClient.builder()
                               .endpointOverride(URI.create("http://localhost:" + wireMock.getPort()))
                               .checksumValidationEnabled(false)
                               .credentialsProvider(StaticCredentialsProvider.create(
                                   AwsBasicCredentials.create("key", "secret")))
                               .build();

        batchManager = SqsAsyncBatchManager.builder()
                                           .client(client)
                                           .scheduledExecutor(Executors.newScheduledThreadPool(10))
                                           .overrideConfiguration(config -> config
                                               .sendRequestFrequency(Duration.ofMillis(10))
                                               .maxBatchSize(MAX_BATCH_SIZE))
                                           .build();
    }

    @AfterEach
    void tearDown() {
        batchManager.close();
        client.close();
    }

    @Test
    void shouldEfficientlyBatchMessagesUnderHighLoad() throws Exception {
        // Given: SQS returns success for batch requests
        stubFor(post(anyUrl())
                    .willReturn(aResponse()
                                    .withStatus(200)
                                    .withBody("{\"Successful\": []}")));

        // When: Send 5000 messages concurrently using 50 threads
        ExecutorService executor = Executors.newFixedThreadPool(CONCURRENT_THREADS);
        CountDownLatch startSignal = new CountDownLatch(1);

        for (int i = 0; i < MESSAGE_COUNT; i++) {
            final String messageBody = String.valueOf(i);
            executor.submit(() -> {
                try {
                    startSignal.await(); // Wait to start all at once
                    batchManager.sendMessage(builder ->
                                                 builder.queueUrl(QUEUE_URL)
                                                        .messageBody(messageBody));
                } catch (Exception ignored) {
                    // Test will fail on assertions if messages aren't sent
                }
            });
        }

        startSignal.countDown(); // Fire all threads simultaneously
        executor.shutdown();
        executor.awaitTermination(10, TimeUnit.SECONDS);

        // Allow batch manager to complete processing
        Thread.sleep(2000);

        // Then: Verify messages were efficiently batched
        List<LoggedRequest> batchRequests = findAll(postRequestedFor(anyUrl()));

        // Calculate batching metrics
        List<Integer> batchSizes = batchRequests.stream()
                                                .map(req -> req.getBodyAsString().split("\"Id\"").length - 1)
                                                .collect(Collectors.toList());

        double avgBatchSize = batchSizes.stream()
                                        .mapToInt(Integer::intValue)
                                        .average()
                                        .orElse(0);

        double fullBatchRatio = batchSizes.stream()
                                          .filter(size -> size >= 9)
                                          .count() / (double) batchSizes.size();

        // Assert efficient batching
        assertThat(avgBatchSize)
            .as("Average batch size")
            .isGreaterThan(8.0);

        assertThat(fullBatchRatio)
            .as("Ratio of nearly full batches (9-10 messages)")
            .isGreaterThan(0.8);

        assertThat(batchRequests.size())
            .as("Total batch requests for %d messages", MESSAGE_COUNT)
            .isLessThan(MESSAGE_COUNT / 5);
    }

   //please add more test cases with different combination of maxBatchSize(too low 1 and max 10) and sendFrequency(too long and too short  )

}

[thornhillcody]

I used this as a base with some changes to make it work. Added a rate limiter to spread out traffic and a small batch frequency to reduce the amount of time required to run the test. Also did medium/small batch size tests.

[thornhillcody]

[INFO] Running software.amazon.awssdk.services.sqs.batchmanager.BatchingEfficiencyUnderLoadTest
[INFO] Tests run: 3, Failures: 0, Errors: 0, Skipped: 0, Time elapsed: 6.365 s -- in software.amazon.awssdk.services.sqs.batchmanager.BatchingEfficiencyUnderLoadTest

Takes 6 seconds (2 per test) as it is.

[joviegas]

Can we remove the cancelScheduledFlush here since we are already cancelling it in Line 112 ?

[thornhillcody]

I opted to leave it in because you'd potentially run a scheduled flush while the manual flush is happening; I didn't see a downside to leaving the cancel in. Seems minor, I don't feel strongly either way.

[joviegas]

you'd potentially run a scheduled flush while the manual flush is happening;

Yeah , make sense , can we add a comment about this .

Thanks for doing a deep dive on this issue and raising the PR 👍

[thornhillcody]

We can certainly remove the first cancel. flushBuffer shouldn't take a significant amount of time (async client), so canceling after will be fine.

[joviegas]

Can we also ad synchronization here since cancelScheduledFlush() can be independently call

    public void cancelScheduledFlush() {
        synchronized (scheduledFlushLock) {
            scheduledFlush.cancel(false);
        }
    }

[joviegas]

Can we please add some test cases in RequestBatchBufferTest for this fix ?

[joviegas]

Can you please merge latest changes from master, since some codebuild check are failing .

[thornhillcody]

OK, I think I've addressed your comments @joviegas

[sonarqubecloud]

Quality Gate passed

Issues
3 New issues
0 Accepted issues

Measures
0 Security Hotspots
100.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud