Examples with FunctionWithIndex org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams.FunctionWithIndex used on opensource projects

Example 1 with FunctionWithIndex

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams.FunctionWithIndex in project beam by apache.

the class SqsUnboundedReader method deleteBatch.

/**
 * Delete the provided {@code receiptHandles} from SQS. Blocking until all messages are deleted.
 *
 * <p>CAUTION: May be invoked from a separate thread.
 */
private void deleteBatch(List<String> receiptHandles) throws IOException {
    int retries = 0;
    FunctionWithIndex<String, DeleteMessageBatchRequestEntry> buildEntry = (handle, id) -> DeleteMessageBatchRequestEntry.builder().id(Long.toString(id)).receiptHandle(handle).build();
    Map<String, DeleteMessageBatchRequestEntry> pendingDeletes = mapWithIndex(receiptHandles.stream(), buildEntry).collect(toMap(e -> e.id(), identity()));
    while (!pendingDeletes.isEmpty()) {
        if (retries >= BATCH_OPERATION_MAX_RETIRES) {
            throw new IOException("Failed to delete " + pendingDeletes.size() + " messages after " + retries + " retries");
        }
        DeleteMessageBatchResponse result = sqsClient.deleteMessageBatch(DeleteMessageBatchRequest.builder().queueUrl(queueUrl()).entries(pendingDeletes.values()).build());
        Map<Boolean, Set<String>> failures = result.failed().stream().collect(partitioningBy(this::isHandleInvalid, mapping(e -> e.id(), toSet())));
        // Keep failed IDs only, but discard invalid receipt handles
        pendingDeletes.keySet().retainAll(failures.getOrDefault(FALSE, ImmutableSet.of()));
        int invalidHandles = failures.getOrDefault(TRUE, ImmutableSet.of()).size();
        if (invalidHandles > 0) {
            LOG.warn("Failed to delete {} messages due to expired receipt handles.", invalidHandles);
        }
        retries += 1;
    }
}