Examples with ChangeMessageVisibilityBatchRequestEntry software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry used on opensource projects

Example 1 with ChangeMessageVisibilityBatchRequestEntry

use of software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry in project beam by apache.

the class SqsUnboundedReader method extendBatch.

/**
 * BLOCKING. Extend the SQS visibility timeout for messages in {@code messages} as {@link KV} of
 * message id, receipt handle.
 */
void extendBatch(long nowMsSinceEpoch, List<KV<String, String>> messages, int extensionSec) throws IOException {
    int retries = 0;
    Function<KV<String, String>, ChangeMessageVisibilityBatchRequestEntry> buildEntry = kv -> ChangeMessageVisibilityBatchRequestEntry.builder().visibilityTimeout(extensionSec).id(kv.getKey()).receiptHandle(kv.getValue()).build();
    Map<String, ChangeMessageVisibilityBatchRequestEntry> pendingExtends = messages.stream().collect(toMap(KV::getKey, buildEntry));
    while (!pendingExtends.isEmpty()) {
        if (retries >= BATCH_OPERATION_MAX_RETIRES) {
            throw new IOException("Failed to extend visibility timeout for " + messages.size() + " messages after " + retries + " retries");
        }
        ChangeMessageVisibilityBatchResponse response = sqsClient.changeMessageVisibilityBatch(ChangeMessageVisibilityBatchRequest.builder().queueUrl(queueUrl()).entries(pendingExtends.values()).build());
        Map<Boolean, Set<String>> failures = response.failed().stream().collect(partitioningBy(this::isHandleInvalid, mapping(e -> e.id(), toSet())));
        // Keep failed IDs only, but discard invalid (expired) receipt handles
        pendingExtends.keySet().retainAll(failures.getOrDefault(FALSE, ImmutableSet.of()));
        // redelivery
        if (extensionSec > 0) {
            numExtendedDeadlines.add(nowMsSinceEpoch, response.successful().size());
            Set<String> invalidMsgIds = failures.getOrDefault(TRUE, ImmutableSet.of());
            if (invalidMsgIds.size() > 0) {
                // consider invalid (expired) messages no longer in flight
                numLateDeadlines.add(nowMsSinceEpoch, invalidMsgIds.size());
                for (String msgId : invalidMsgIds) {
                    inFlight.remove(msgId);
                }
                LOG.warn("Failed to extend visibility timeout for {} messages with expired receipt handles.", invalidMsgIds.size());
            }
        }
        retries += 1;
    }
}