Skip to content

Commit

Permalink
Handled surface area review comments (#5563)
Browse files Browse the repository at this point in the history 
* Initial version

* Intermediate changes

* Update after internal poll

* ResponseCOnfiguration construction updated

* RequestOverride configuration check added to Bypass batch manager

* Handled review comments
  • Loading branch information
@joviegas
joviegas authored Sep 5, 2024
1 parent b04e5ad commit bff5b97
Show file tree
Hide file tree
Showing 16 changed files with 632 additions and 830 deletions.
429 changes: 148 additions & 281 deletions ...ain/java/software/amazon/awssdk/services/sqs/batchmanager/BatchOverrideConfiguration.java
View file

Large diffs are not rendered by default.

4 changes: 3 additions & 1 deletion ...oftware/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,9 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) {
);

this.receiveMessageBatchManager =
new ReceiveMessageBatchManager(client, scheduledExecutor, builder.overrideConfiguration);
new ReceiveMessageBatchManager(client,
scheduledExecutor,
ResponseBatchConfiguration.builder(builder.overrideConfiguration).build());
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion ...n/java/software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveBatchManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ public CompletableFuture<ReceiveMessageResponse> processRequest(ReceiveMessageRe
}
int numMessages = rq.maxNumberOfMessages() != null ? rq.maxNumberOfMessages() : MAX_SUPPORTED_SQS_RECEIVE_MSG;

return queueAttributesManager.getReceiveMessageTimeout(rq, config.minReceiveWaitTime()).thenCompose(waitTimeMs -> {
return queueAttributesManager.getReceiveMessageTimeout(rq, config.messageMinWaitDuration()).thenCompose(waitTimeMs -> {
CompletableFuture<ReceiveMessageResponse> receiveMessageFuture = new CompletableFuture<>();
receiveQueueBuffer.receiveMessage(receiveMessageFuture, numMessages);
CompletableFuture<ReceiveMessageResponse> timeoutFuture = new CompletableFuture<>();
Expand Down
33 changes: 24 additions & 9 deletions ...software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveMessageBatchManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
import java.util.concurrent.ScheduledExecutorService;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.services.sqs.SqsAsyncClient;
import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
import software.amazon.awssdk.utils.Logger;
Expand All @@ -39,19 +38,21 @@ public class ReceiveMessageBatchManager implements SdkAutoCloseable {

public ReceiveMessageBatchManager(SqsAsyncClient sqsClient,
ScheduledExecutorService executor,
BatchOverrideConfiguration config) {
ResponseBatchConfiguration config) {
this.sqsClient = sqsClient;
this.executor = executor;
this.config = new ResponseBatchConfiguration(config);
this.config = config;


}

public CompletableFuture<ReceiveMessageResponse> batchRequest(ReceiveMessageRequest request) {
if (canBeRetrievedFromQueueBuffer(request)) {
String ineligibleReason = checkBatchingEligibility(request);
if (ineligibleReason == null) {
return receiveBatchManagerMap.computeIfAbsent(generateBatchKey(request), key -> createReceiveBatchManager(request))
.processRequest(request);
} else {
log.debug(() -> "canBeRetrievedFromQueueBuffer failed, so skipping batching for request for Queue with URL: "
+ request.queueUrl());
log.debug(() -> String.format("Batching skipped. Reason: %s", ineligibleReason));
return sqsClient.receiveMessage(request);
}
}
Expand All @@ -77,11 +78,25 @@ public void close() {
receiveBatchManagerMap.values().forEach(ReceiveBatchManager::close);
}

private boolean canBeRetrievedFromQueueBuffer(ReceiveMessageRequest rq) {
return hasCompatibleAttributes(rq) && isBufferingEnabled() && rq.visibilityTimeout() == null;
private String checkBatchingEligibility(ReceiveMessageRequest rq) {
if (!hasCompatibleAttributes(rq)) {
return "Incompatible attributes.";
}
if (rq.visibilityTimeout() != null) {
return "Visibility timeout is set.";
}
if (!isBufferingEnabled()) {
return "Buffering is disabled.";
}
if (rq.overrideConfiguration().isPresent()) {
return "Request has override configurations.";
}
if (rq.waitTimeSeconds() != null && rq.waitTimeSeconds() != 0) {
return "Request has long polling enabled.";
}
return null;
}


private boolean hasCompatibleAttributes(ReceiveMessageRequest rq) {
return !rq.hasAttributeNames()
&& hasCompatibleSystemAttributes(rq)
Expand Down
15 changes: 7 additions & 8 deletions ...in/java/software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveQueueBuffer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@

package software.amazon.awssdk.services.sqs.internal.batchmanager;

import static software.amazon.awssdk.services.sqs.internal.batchmanager.ResponseBatchConfiguration.MAX_SUPPORTED_SQS_RECEIVE_MSG;

import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
Expand Down Expand Up @@ -119,14 +121,11 @@ private void spawnMoreReceiveTasks() {

private int determineDesiredBatches() {
int desiredBatches = Math.max(config.maxDoneReceiveBatches(), 1);

if (config.adaptivePrefetching()) {
int totalRequested = futures.stream()
.mapToInt(FutureRequestWrapper::getRequestedSize)
.sum();
int batchesNeededToFulfillFutures = (int) Math.ceil((float) totalRequested / config.maxBatchItems());
desiredBatches = Math.min(batchesNeededToFulfillFutures, desiredBatches);
}
int totalRequested = futures.stream()
.mapToInt(FutureRequestWrapper::getRequestedSize)
.sum();
int batchesNeededToFulfillFutures = (int) Math.ceil((float) totalRequested / MAX_SUPPORTED_SQS_RECEIVE_MSG);
desiredBatches = Math.min(batchesNeededToFulfillFutures, desiredBatches);

return desiredBatches;
}
Expand Down
3 changes: 0 additions & 3 deletions ...va/software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveSqsMessageHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,9 +87,6 @@ public CompletableFuture<ReceiveSqsMessageHelper> asyncReceiveMessage() {

request.visibilityTimeout(NumericUtils.saturatedCast(this.visibilityTimeout.getSeconds()));

if (config.longPollWaitTimeout() != null) {
request.waitTimeSeconds(NumericUtils.saturatedCast(config.longPollWaitTimeout().getSeconds()));
}
try {
return asyncClient.receiveMessage(request.build())
.handle((response, throwable) -> {
Expand Down
24 changes: 12 additions & 12 deletions .../software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,22 +24,23 @@ public final class RequestBatchConfiguration {

public static final int DEFAULT_MAX_BATCH_ITEMS = 10;
public static final int DEFAULT_MAX_BATCH_BYTES_SIZE = -1;
public static final int DEFAULT_MAX_BATCH_KEYS = 1000;
public static final int DEFAULT_MAX_BATCH_KEYS = 10000;
public static final int DEFAULT_MAX_BUFFER_SIZE = 500;
public static final Duration DEFAULT_MAX_BATCH_OPEN_IN_MS = Duration.ofMillis(200);

private final Integer maxBatchItems;
private final Integer maxBatchKeys;
private final Integer maxBufferSize;
private final Duration maxBatchOpenDuration;
private final Duration sendRequestFrequency;
private final Integer maxBatchBytesSize;

private RequestBatchConfiguration(Builder builder) {

this.maxBatchItems = builder.maxBatchItems != null ? builder.maxBatchItems : DEFAULT_MAX_BATCH_ITEMS;
this.maxBatchKeys = builder.maxBatchKeys != null ? builder.maxBatchKeys : DEFAULT_MAX_BATCH_KEYS;
this.maxBufferSize = builder.maxBufferSize != null ? builder.maxBufferSize : DEFAULT_MAX_BUFFER_SIZE;
this.maxBatchOpenDuration = builder.maxBatchOpenDuration != null ? builder.maxBatchOpenDuration :
this.sendRequestFrequency = builder.sendRequestFrequency != null ?
builder.sendRequestFrequency :
DEFAULT_MAX_BATCH_OPEN_IN_MS;
this.maxBatchBytesSize = builder.maxBatchBytesSize != null ? builder.maxBatchBytesSize : DEFAULT_MAX_BATCH_BYTES_SIZE;

Expand All @@ -52,16 +53,15 @@ public static Builder builder() {
public static Builder builder(BatchOverrideConfiguration configuration) {
if (configuration != null) {
return new Builder()
.maxBatchKeys(configuration.maxBatchKeys())
.maxBatchItems(configuration.maxBatchItems())
.maxBatchOpenDuration(configuration.maxBatchOpenDuration())
.maxBufferSize(configuration.maxBufferSize());
.maxBatchItems(configuration.maxBatchSize())
.sendRequestFrequency(configuration.sendRequestFrequency())
.maxBatchBytesSize(configuration.maxBatchSize());
}
return new Builder();
}

public Duration maxBatchOpenDuration() {
return maxBatchOpenDuration;
public Duration sendRequestFrequency() {
return sendRequestFrequency;
}

public int maxBatchItems() {
Expand All @@ -85,7 +85,7 @@ public static final class Builder {
private Integer maxBatchItems;
private Integer maxBatchKeys;
private Integer maxBufferSize;
private Duration maxBatchOpenDuration;
private Duration sendRequestFrequency;
private Integer maxBatchBytesSize;

private Builder() {
Expand All @@ -106,8 +106,8 @@ public Builder maxBufferSize(Integer maxBufferSize) {
return this;
}

public Builder maxBatchOpenDuration(Duration maxBatchOpenDuration) {
this.maxBatchOpenDuration = maxBatchOpenDuration;
public Builder sendRequestFrequency(Duration sendRequestFrequency) {
this.sendRequestFrequency = sendRequestFrequency;
return this;
}

Expand Down
17 changes: 9 additions & 8 deletions ...n/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,9 @@ public abstract class RequestBatchManager<RequestT, ResponseT, BatchResponseT> {
protected final RequestBatchConfiguration batchConfiguration ;

private final int maxBatchItems;
private final Duration maxBatchOpenDuration;
private final Duration sendRequestFrequency;
private final BatchingMap<RequestT, ResponseT> requestsAndResponsesMaps;
private final ScheduledExecutorService scheduledExecutor;

private final Set<CompletableFuture<BatchResponseT>> pendingBatchResponses ;
private final Set<CompletableFuture<ResponseT>> pendingResponses ;

Expand All @@ -57,7 +56,7 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration,
ScheduledExecutorService scheduledExecutor) {
batchConfiguration = overrideConfiguration;
this.maxBatchItems = batchConfiguration.maxBatchItems();
this.maxBatchOpenDuration = batchConfiguration.maxBatchOpenDuration();
this.sendRequestFrequency = batchConfiguration.sendRequestFrequency();
this.scheduledExecutor = Validate.notNull(scheduledExecutor, "Null scheduledExecutor");
pendingBatchResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());
pendingResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());
Expand All @@ -71,7 +70,6 @@ public CompletableFuture<ResponseT> batchRequest(RequestT request) {

try {
String batchKey = getBatchKey(request);

// Handle potential byte size overflow only if there are request in map and if feature enabled
if (requestsAndResponsesMaps.contains(batchKey) && batchConfiguration.maxBatchBytesSize() > 0) {
Optional.of(requestsAndResponsesMaps.flushableRequestsOnByteLimitBeforeAdd(batchKey, request))
Expand All @@ -81,7 +79,9 @@ public CompletableFuture<ResponseT> batchRequest(RequestT request) {

// Add request and response to the map, scheduling a flush if necessary
requestsAndResponsesMaps.put(batchKey,
() -> scheduleBufferFlush(batchKey, maxBatchOpenDuration.toMillis(), scheduledExecutor),
() -> scheduleBufferFlush(batchKey,
sendRequestFrequency.toMillis(),
scheduledExecutor),
request,
response);

Expand Down Expand Up @@ -109,8 +109,10 @@ private void manualFlushBuffer(String batchKey,
Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableRequests) {
requestsAndResponsesMaps.cancelScheduledFlush(batchKey);
flushBuffer(batchKey, flushableRequests);
requestsAndResponsesMaps.putScheduledFlush(batchKey, scheduleBufferFlush(batchKey, maxBatchOpenDuration.toMillis(),
scheduledExecutor));
requestsAndResponsesMaps.putScheduledFlush(batchKey,
scheduleBufferFlush(batchKey,
sendRequestFrequency.toMillis(),
scheduledExecutor));
}

private void flushBuffer(String batchKey, Map<String, BatchingExecutionContext<RequestT, ResponseT>> flushableRequests) {
Expand Down Expand Up @@ -174,5 +176,4 @@ public void close() {
requestsAndResponsesMaps.clear();
}


}
Loading

0 comments on commit bff5b97

Please sign in to comment.