Search in sources :

Example 1 with ChangeMessageVisibilityBatchResponse

use of software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse 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;
    }
}
Also used : ChangeMessageVisibilityBatchResponse(software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse) ChangeMessageVisibilityBatchRequestEntry(software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry) SENT_TIMESTAMP(software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName.SENT_TIMESTAMP) UnboundedSource(org.apache.beam.sdk.io.UnboundedSource) LoggerFactory(org.slf4j.LoggerFactory) Streams.mapWithIndex(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams.mapWithIndex) ChangeMessageVisibilityBatchRequest(software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest) ImmutableSet(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet) Collectors.toMap(java.util.stream.Collectors.toMap) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Map(java.util.Map) Iterables(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables) ReceiveMessageRequest(software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest) Collectors.toSet(java.util.stream.Collectors.toSet) Min(org.apache.beam.sdk.transforms.Min) Set(java.util.Set) ClientBuilderFactory(org.apache.beam.sdk.io.aws2.common.ClientBuilderFactory) Sum(org.apache.beam.sdk.transforms.Sum) SqsClient(software.amazon.awssdk.services.sqs.SqsClient) MovingFunction(org.apache.beam.sdk.util.MovingFunction) List(java.util.List) Max(org.apache.beam.sdk.transforms.Max) Collections2.transform(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Collections2.transform) DeleteMessageBatchResponse(software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse) AwsOptions(org.apache.beam.sdk.io.aws2.options.AwsOptions) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting) Function.identity(java.util.function.Function.identity) Queue(java.util.Queue) TRUE(java.lang.Boolean.TRUE) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue) Collectors.partitioningBy(java.util.stream.Collectors.partitioningBy) KV(org.apache.beam.sdk.values.KV) Combine(org.apache.beam.sdk.transforms.Combine) Duration(org.joda.time.Duration) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) BucketingFunction(org.apache.beam.sdk.util.BucketingFunction) Function(java.util.function.Function) DeleteMessageBatchRequest(software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) LinkedHashMap(java.util.LinkedHashMap) Message(software.amazon.awssdk.services.sqs.model.Message) Collectors.mapping(java.util.stream.Collectors.mapping) NoSuchElementException(java.util.NoSuchElementException) FALSE(java.lang.Boolean.FALSE) Logger(org.slf4j.Logger) UTF_8(java.nio.charset.StandardCharsets.UTF_8) FunctionWithIndex(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams.FunctionWithIndex) IOException(java.io.IOException) ReceiveMessageResponse(software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse) BatchResultErrorEntry(software.amazon.awssdk.services.sqs.model.BatchResultErrorEntry) Collectors.toList(java.util.stream.Collectors.toList) CheckpointMark(org.apache.beam.sdk.io.UnboundedSource.CheckpointMark) Preconditions.checkState(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) Instant(org.joda.time.Instant) VISIBILITY_TIMEOUT(software.amazon.awssdk.services.sqs.model.QueueAttributeName.VISIBILITY_TIMEOUT) ClientConfiguration(org.apache.beam.sdk.io.aws2.common.ClientConfiguration) Clock(java.time.Clock) ArrayDeque(java.util.ArrayDeque) DeleteMessageBatchRequestEntry(software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry) ImmutableSet(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet) Collectors.toSet(java.util.stream.Collectors.toSet) Set(java.util.Set) HashSet(java.util.HashSet) KV(org.apache.beam.sdk.values.KV) IOException(java.io.IOException) ChangeMessageVisibilityBatchResponse(software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse) ChangeMessageVisibilityBatchRequestEntry(software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean)

Aggregations

IOException (java.io.IOException)1 FALSE (java.lang.Boolean.FALSE)1 TRUE (java.lang.Boolean.TRUE)1 UTF_8 (java.nio.charset.StandardCharsets.UTF_8)1 Clock (java.time.Clock)1 ArrayDeque (java.util.ArrayDeque)1 ArrayList (java.util.ArrayList)1 HashSet (java.util.HashSet)1 LinkedHashMap (java.util.LinkedHashMap)1 List (java.util.List)1 Map (java.util.Map)1 NoSuchElementException (java.util.NoSuchElementException)1 Queue (java.util.Queue)1 Set (java.util.Set)1 ConcurrentLinkedQueue (java.util.concurrent.ConcurrentLinkedQueue)1 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)1 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)1 Function (java.util.function.Function)1 Function.identity (java.util.function.Function.identity)1 Collectors.mapping (java.util.stream.Collectors.mapping)1