Receive Batch Manager Implementation

joviegas commented 3 months ago

Motivation and Context

This change is required to efficiently manage the receiving of message batches from AWS SQS, ensuring proper handling of message visibility timeout and error scenarios. It solves the problem of handling asynchronous message retrieval and processing, managing visibility timeouts, and handling exceptions during the process.

Modifications

AsyncReceiveMessageBatch Class:
- Forms a ReceiveMessageRequest based on configuration settings.
- Collects messages from an AWS SQS queue.
- Handles exceptions during the process.
- Manages message visibility timeout by tracking the visibility deadline and expiring messages if not processed in time.
- Supports clearing messages in the batch and changing their visibility as needed.
ReceiveMessageCompletableFuture Class:
- Manages the completion of the receive message future with success or failure.
- Starts a wait timer for the receive message request and handles timeout scenarios.
ReceiveQueueBuffer Class:
- Manages the queue buffer for received messages.
- Handles the reception of messages and the spawning of additional receive tasks.
- Manages the clearing of finished tasks and handling of futures for received messages.
- Supports shutdown functionality, ensuring all futures are completed exceptionally upon shutdown.
ResponseBatchConfiguration Class:
- Configures various batch settings such as visibility timeout, long poll settings, and adaptive prefetching.
QueueAttributesManager Class:
- Manages SQS queue attributes.
- Fetches and caches attributes like visibility timeout and receive message wait time.
- Ensures only a single call is made to fetch the attributes, with proper handling of concurrent requests and failures.

Design Details

Asynchronous Handling:
- The AsyncReceiveMessageBatch uses asynchronous handling for receiving messages from SQS. It leverages CompletableFuture to manage asynchronous operations and handle results or exceptions.
Visibility Timeout Management:
- The class tracks the visibility deadline for messages to ensure they are reprocessed if not handled within the specified time. This is managed using visibilityDeadlineNano.
Error Handling:
- The AsyncReceiveMessageBatch and QueueAttributesManager classes handle exceptions by setting the future to complete exceptionally and resetting state where necessary.
Concurrent Task Management:
- ReceiveQueueBuffer uses a concurrent queue to manage finished tasks and inflight receive message batches, ensuring thread-safe operations.
- QueueAttributesManager uses an AtomicReference to manage the cached future for fetching queue attributes, ensuring that only one request is made even under concurrent access.
Adaptive Prefetching:
- The ResponseBatchConfiguration supports adaptive prefetching, adjusting the number of batches based on the load and configuration.

Testing

Unit tests were added for AsyncReceiveMessageBatch to cover:
- Message receiving and handling.
- Visibility timeout management.
- Error handling during message receiving.
Tests for ReceiveQueueBuffer include:
- Handling multiple receive message requests.
- Managing finished tasks and inflight batches.
- Testing shutdown behavior.
QueueAttributesManager tests:
- Fetching and caching queue attributes.
- Handling exceptions during attribute fetching.
- Ensuring concurrent requests are handled correctly and only a single call is made to fetch attributes.