Skip to content

Commit ff804c0

Browse files
pankajagrawal16Pankaj Agrawal
pankajagrawal16
and
Pankaj Agrawal
authored
feat(batch-processing): move non retry-able message to DLQ (#500)
* feat(batch-processing): Support for moving non retryable msg to DLQ * chore(performance): Build queue url from queue arn instead of API call * feat(batch-processing): test cases * docs(batch-processing): Support for moving non retryable msg to DLQ Co-authored-by: Pankaj Agrawal <[email protected]>
1 parent da613d3 commit ff804c0

File tree

9 files changed

+834
-70
lines changed

9 files changed

+834
-70
lines changed

‎powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SqsBatch.java

+24-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
* calling {@link SqsMessageHandler#process(SQSMessage)} method for each {@link SQSMessage} in the received {@link SQSEvent}
1818
* </p>
1919
*
20-
* </p>
20+
* <p>
2121
* If any exception is thrown from {@link SqsMessageHandler#process(SQSMessage)} during processing of a messages, Utility
2222
* will take care of deleting all the successful messages from SQS. When one or more single message fails processing due
2323
* to exception thrown from {@link SqsMessageHandler#process(SQSMessage)}, Lambda execution will fail
@@ -32,6 +32,24 @@
3232
* {@link SqsBatch#suppressException()} to true. By default its value is false
3333
* </p>
3434
*
35+
* <p>
36+
* If you want certain exceptions to be treated as permanent failures, i.e. exceptions where the result of retrying will
37+
* always be a failure and want these can be immediately moved to the dead letter queue associated to the source SQS queue,
38+
*
39+
* you can use {@link SqsBatch#nonRetryableExceptions()} to configure such exceptions.
40+
* Make sure function execution role has sqs:GetQueueAttributes permission on source SQS queue and sqs:SendMessage,
41+
* sqs:SendMessageBatch permission for configured DLQ.
42+
*
43+
* If you want such messages to be deleted instead, set {@link SqsBatch#deleteNonRetryableMessageFromQueue()} to true.
44+
* By default its value is false.
45+
*
46+
* If there is no DLQ configured on source SQS queue and {@link SqsBatch#nonRetryableExceptions()} attribute is set, if
47+
* nonRetryableExceptions occurs from {@link SqsMessageHandler}, such exceptions will still be treated as temporary
48+
* exceptions and the message will be moved back to source SQS queue for reprocessing. The same behaviour will occur if
49+
* for some reason the utility is unable to move the message to the DLQ. An example of this could be because the function
50+
* is missing the correct permissions.
51+
* </p>
52+
*
3553
* <pre>
3654
* public class SqsMessageHandler implements RequestHandler<SQSEvent, String> {
3755
*
@@ -51,6 +69,7 @@
5169
*
5270
* ...
5371
* </pre>
72+
* @see <a href="https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-dead-letter-queues.html">Amazon SQS dead-letter queues</a>
5473
*/
5574
@Retention(RetentionPolicy.RUNTIME)
5675
@Target(ElementType.METHOD)
@@ -59,4 +78,8 @@
5978
Class<? extends SqsMessageHandler<Object>> value();
6079

6180
boolean suppressException() default false;
81+
82+
Class<? extends Exception>[] nonRetryableExceptions() default {};
83+
84+
boolean deleteNonRetryableMessageFromQueue() default false;
6285
}

‎powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SqsUtils.java

+266-18
Large diffs are not rendered by default.

‎powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/BatchContext.java

+162-31
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,42 @@
11
package software.amazon.lambda.powertools.sqs.internal;
22

3-
import org.slf4j.Logger;
4-
import org.slf4j.LoggerFactory;
5-
63
import java.util.ArrayList;
4+
import java.util.Arrays;
5+
import java.util.HashMap;
76
import java.util.List;
7+
import java.util.Map;
8+
import java.util.Optional;
89

10+
import com.fasterxml.jackson.core.JsonProcessingException;
11+
import com.fasterxml.jackson.databind.JsonNode;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
14+
import software.amazon.awssdk.core.SdkBytes;
915
import software.amazon.awssdk.services.sqs.SqsClient;
1016
import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
1117
import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
1218
import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
13-
import software.amazon.awssdk.services.sqs.model.GetQueueUrlRequest;
19+
import software.amazon.awssdk.services.sqs.model.GetQueueAttributesRequest;
20+
import software.amazon.awssdk.services.sqs.model.GetQueueAttributesResponse;
21+
import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
22+
import software.amazon.awssdk.services.sqs.model.QueueAttributeName;
23+
import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
24+
import software.amazon.awssdk.services.sqs.model.SendMessageBatchResponse;
1425
import software.amazon.lambda.powertools.sqs.SQSBatchProcessingException;
26+
import software.amazon.lambda.powertools.sqs.SqsUtils;
1527

1628
import static com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
1729
import static java.lang.String.format;
30+
import static java.util.Optional.ofNullable;
1831
import static java.util.stream.Collectors.toList;
1932

2033
public final class BatchContext {
2134
private static final Logger LOG = LoggerFactory.getLogger(BatchContext.class);
35+
private static final Map<String, String> QUEUE_ARN_TO_DLQ_URL_MAPPING = new HashMap<>();
2236

37+
private final Map<SQSMessage, Exception> messageToException = new HashMap<>();
2338
private final List<SQSMessage> success = new ArrayList<>();
24-
private final List<SQSMessage> failures = new ArrayList<>();
25-
private final List<Exception> exceptions = new ArrayList<>();
39+
2640
private final SqsClient client;
2741

2842
public BatchContext(SqsClient client) {
@@ -34,53 +48,170 @@ public void addSuccess(SQSMessage event) {
3448
}
3549

3650
public void addFailure(SQSMessage event, Exception e) {
37-
failures.add(event);
38-
exceptions.add(e);
51+
messageToException.put(event, e);
3952
}
4053

41-
public <T> void processSuccessAndHandleFailed(final List<T> successReturns,
42-
final boolean suppressException) {
54+
@SafeVarargs
55+
public final <T> void processSuccessAndHandleFailed(final List<T> successReturns,
56+
final boolean suppressException,
57+
final boolean deleteNonRetryableMessageFromQueue,
58+
final Class<? extends Exception>... nonRetryableExceptions) {
4359
if (hasFailures()) {
44-
deleteSuccessMessage();
4560

46-
if (suppressException) {
47-
List<String> messageIds = failures.stream().
48-
map(SQSMessage::getMessageId)
49-
.collect(toList());
61+
List<Exception> exceptions = new ArrayList<>();
62+
List<SQSMessage> failedMessages = new ArrayList<>();
63+
Map<SQSMessage, Exception> nonRetryableMessageToException = new HashMap<>();
5064

51-
LOG.debug(format("[%s] records failed processing, but exceptions are suppressed. " +
52-
"Failed messages %s", failures.size(), messageIds));
65+
if (nonRetryableExceptions.length == 0) {
66+
exceptions.addAll(messageToException.values());
67+
failedMessages.addAll(messageToException.keySet());
5368
} else {
54-
throw new SQSBatchProcessingException(exceptions, failures, successReturns);
69+
messageToException.forEach((sqsMessage, exception) -> {
70+
boolean nonRetryableException = isNonRetryableException(exception, nonRetryableExceptions);
71+
72+
if (nonRetryableException) {
73+
nonRetryableMessageToException.put(sqsMessage, exception);
74+
} else {
75+
exceptions.add(exception);
76+
failedMessages.add(sqsMessage);
77+
}
78+
});
79+
}
80+
81+
List<SQSMessage> messagesToBeDeleted = new ArrayList<>(success);
82+
83+
if (!nonRetryableMessageToException.isEmpty() && deleteNonRetryableMessageFromQueue) {
84+
messagesToBeDeleted.addAll(nonRetryableMessageToException.keySet());
85+
} else if (!nonRetryableMessageToException.isEmpty()) {
86+
87+
boolean isMovedToDlq = moveNonRetryableMessagesToDlqIfConfigured(nonRetryableMessageToException);
88+
89+
if (!isMovedToDlq) {
90+
exceptions.addAll(nonRetryableMessageToException.values());
91+
failedMessages.addAll(nonRetryableMessageToException.keySet());
92+
}
5593
}
94+
95+
deleteMessagesFromQueue(messagesToBeDeleted);
96+
97+
processFailedMessages(successReturns, suppressException, exceptions, failedMessages);
98+
}
99+
}
100+
101+
private <T> void processFailedMessages(List<T> successReturns,
102+
boolean suppressException,
103+
List<Exception> exceptions,
104+
List<SQSMessage> failedMessages) {
105+
if (failedMessages.isEmpty()) {
106+
return;
107+
}
108+
109+
if (suppressException) {
110+
List<String> messageIds = failedMessages.stream().
111+
map(SQSMessage::getMessageId)
112+
.collect(toList());
113+
114+
LOG.debug(format("[%s] records failed processing, but exceptions are suppressed. " +
115+
"Failed messages %s", failedMessages.size(), messageIds));
116+
} else {
117+
throw new SQSBatchProcessingException(exceptions, failedMessages, successReturns);
56118
}
57119
}
58120

121+
private boolean isNonRetryableException(Exception exception, Class<? extends Exception>[] nonRetryableExceptions) {
122+
return Arrays.stream(nonRetryableExceptions)
123+
.anyMatch(aClass -> aClass.isInstance(exception));
124+
}
125+
126+
private boolean moveNonRetryableMessagesToDlqIfConfigured(Map<SQSMessage, Exception> nonRetryableMessageToException) {
127+
Optional<String> dlqUrl = fetchDlqUrl(nonRetryableMessageToException);
128+
129+
if (!dlqUrl.isPresent()) {
130+
return false;
131+
}
132+
133+
List<SendMessageBatchRequestEntry> dlqMessages = nonRetryableMessageToException.keySet().stream()
134+
.map(sqsMessage -> {
135+
Map<String, MessageAttributeValue> messageAttributesMap = new HashMap<>();
136+
137+
sqsMessage.getMessageAttributes().forEach((s, messageAttribute) -> {
138+
MessageAttributeValue.Builder builder = MessageAttributeValue.builder();
139+
140+
builder
141+
.dataType(messageAttribute.getDataType())
142+
.stringValue(messageAttribute.getStringValue());
143+
144+
if (null != messageAttribute.getBinaryValue()) {
145+
builder.binaryValue(SdkBytes.fromByteBuffer(messageAttribute.getBinaryValue()));
146+
}
147+
148+
messageAttributesMap.put(s, builder.build());
149+
});
150+
151+
return SendMessageBatchRequestEntry.builder()
152+
.messageBody(sqsMessage.getBody())
153+
.id(sqsMessage.getMessageId())
154+
.messageAttributes(messageAttributesMap)
155+
.build();
156+
})
157+
.collect(toList());
158+
159+
SendMessageBatchResponse sendMessageBatchResponse = client.sendMessageBatch(builder -> builder.queueUrl(dlqUrl.get())
160+
.entries(dlqMessages));
161+
162+
LOG.debug("Response from send batch message to DLQ request {}", sendMessageBatchResponse);
163+
164+
return true;
165+
}
166+
167+
private Optional<String> fetchDlqUrl(Map<SQSMessage, Exception> nonRetryableMessageToException) {
168+
return nonRetryableMessageToException.keySet().stream()
169+
.findFirst()
170+
.map(sqsMessage -> QUEUE_ARN_TO_DLQ_URL_MAPPING.computeIfAbsent(sqsMessage.getEventSourceArn(), sourceArn -> {
171+
String queueUrl = url(sourceArn);
172+
173+
GetQueueAttributesResponse queueAttributes = client.getQueueAttributes(GetQueueAttributesRequest.builder()
174+
.attributeNames(QueueAttributeName.REDRIVE_POLICY)
175+
.queueUrl(queueUrl)
176+
.build());
177+
178+
return ofNullable(queueAttributes.attributes().get(QueueAttributeName.REDRIVE_POLICY))
179+
.map(policy -> {
180+
try {
181+
return SqsUtils.objectMapper().readTree(policy);
182+
} catch (JsonProcessingException e) {
183+
LOG.debug("Unable to parse Re drive policy for queue {}. Even if DLQ exists, failed messages will be send back to main queue.", queueUrl, e);
184+
return null;
185+
}
186+
})
187+
.map(node -> node.get("deadLetterTargetArn"))
188+
.map(JsonNode::asText)
189+
.map(this::url)
190+
.orElse(null);
191+
}));
192+
}
193+
59194
private boolean hasFailures() {
60-
return !failures.isEmpty();
195+
return !messageToException.isEmpty();
61196
}
62197

63-
private void deleteSuccessMessage() {
64-
if (!success.isEmpty()) {
198+
private void deleteMessagesFromQueue(final List<SQSMessage> messages) {
199+
if (!messages.isEmpty()) {
65200
DeleteMessageBatchRequest request = DeleteMessageBatchRequest.builder()
66-
.queueUrl(url())
67-
.entries(success.stream().map(m -> DeleteMessageBatchRequestEntry.builder()
201+
.queueUrl(url(messages.get(0).getEventSourceArn()))
202+
.entries(messages.stream().map(m -> DeleteMessageBatchRequestEntry.builder()
68203
.id(m.getMessageId())
69204
.receiptHandle(m.getReceiptHandle())
70205
.build()).collect(toList()))
71206
.build();
72207

73208
DeleteMessageBatchResponse deleteMessageBatchResponse = client.deleteMessageBatch(request);
74-
LOG.debug(format("Response from delete request %s", deleteMessageBatchResponse));
209+
LOG.debug("Response from delete request {}", deleteMessageBatchResponse);
75210
}
76211
}
77212

78-
private String url() {
79-
String[] arnArray = success.get(0).getEventSourceArn().split(":");
80-
return client.getQueueUrl(GetQueueUrlRequest.builder()
81-
.queueOwnerAWSAccountId(arnArray[4])
82-
.queueName(arnArray[5])
83-
.build())
84-
.queueUrl();
213+
private String url(String queueArn) {
214+
String[] arnArray = queueArn.split(":");
215+
return String.format("https://sqs.%s.amazonaws.com/%s/%s", arnArray[3], arnArray[4], arnArray[5]);
85216
}
86217
}

‎powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/SqsMessageBatchProcessorAspect.java

+5-1
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,11 @@ && placedOnSqsEventRequestHandler(pjp)) {
2929

3030
SQSEvent sqsEvent = (SQSEvent) proceedArgs[0];
3131

32-
batchProcessor(sqsEvent, sqsBatch.suppressException(), sqsBatch.value());
32+
batchProcessor(sqsEvent,
33+
sqsBatch.suppressException(),
34+
sqsBatch.value(),
35+
sqsBatch.deleteNonRetryableMessageFromQueue(),
36+
sqsBatch.nonRetryableExceptions());
3337
}
3438

3539
return pjp.proceed(proceedArgs);

0 commit comments

Comments
 (0)