Added metrics to SQS notification queue for S3 source

Signed-off-by: Manuel Mangas Zurita <mzurita@amazon.com>

Author: mzurita-amz

data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java

@@ -33,6 +33,11 @@
 import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
 import software.amazon.awssdk.services.sqs.model.SqsException;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.services.sqs.model.KmsAccessDeniedException;
+import software.amazon.awssdk.services.sqs.model.KmsNotFoundException;
+import software.amazon.awssdk.services.sqs.model.KmsThrottledException;
+import software.amazon.awssdk.services.sqs.model.QueueDoesNotExistException;
 import software.amazon.awssdk.services.sts.model.StsException;
 
 import java.time.Duration;
@@ -58,6 +63,9 @@ public class SqsWorker implements Runnable {
     static final String SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME = "sqsVisibilityTimeoutChangedCount";
     static final String SQS_VISIBILITY_TIMEOUT_CHANGE_FAILED_COUNT_METRIC_NAME = "sqsVisibilityTimeoutChangeFailedCount";
     static final String ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME = "acknowledgementSetCallbackCounter";
+    public static final String SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME = "sqsMessagesAccessDenied";
+    public static final String SQS_MESSAGE_THROTTLED_METRIC_NAME = "sqsMessagesThrottled";
+    public static final String SQS_QUEUE_NOT_FOUND_METRIC_NAME = "sqsQueueNotFound";
 
     private final S3SourceConfig s3SourceConfig;
     private final SqsClient sqsClient;
@@ -74,6 +82,9 @@ public class SqsWorker implements Runnable {
     private final Counter acknowledgementSetCallbackCounter;
     private final Counter sqsVisibilityTimeoutChangedCount;
     private final Counter sqsVisibilityTimeoutChangeFailedCount;
+    private final Counter sqsMessageAccessDeniedCounter;
+    private final Counter sqsMessageThrottledCounter;
+    private final Counter sqsQueueNotFoundCounter;
     private final Timer sqsMessageDelayTimer;
     private final Backoff standardBackoff;
     private final SqsMessageParser sqsMessageParser;
@@ -111,6 +122,9 @@ public SqsWorker(final AcknowledgementSetManager acknowledgementSetManager,
         acknowledgementSetCallbackCounter = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME);
         sqsVisibilityTimeoutChangedCount = pluginMetrics.counter(SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME);
         sqsVisibilityTimeoutChangeFailedCount = pluginMetrics.counter(SQS_VISIBILITY_TIMEOUT_CHANGE_FAILED_COUNT_METRIC_NAME);
+        sqsMessageAccessDeniedCounter = pluginMetrics.counter(SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME);
+        sqsMessageThrottledCounter = pluginMetrics.counter(SQS_MESSAGE_THROTTLED_METRIC_NAME);
+        sqsQueueNotFoundCounter = pluginMetrics.counter(SQS_QUEUE_NOT_FOUND_METRIC_NAME);
     }
 
     @Override
@@ -164,6 +178,7 @@ private List<Message> getMessagesFromSqs() {
             return messages;
         } catch (final SqsException | StsException e) {
             LOG.error("Error reading from SQS: {}. Retrying with exponential backoff.", e.getMessage());
+            recordSqsException(e);
             sqsReceiveMessagesFailedCounter.increment();
             applyBackoff();
             return Collections.emptyList();
@@ -452,4 +467,21 @@ private int getApproximateReceiveCount(final Message message) {
     void stop() {
         isStopped = true;
     }
+
+    public void recordSqsException(final AwsServiceException e) {
+        // AWS SQS emits some of their exceptions without the matching HTTP code. As we want to generate an aggregate version of
+        // these exceptions, we have to explicitly catch the type alongside the status code for the ones that leverage the status
+        // code (i.e. InvalidAddressException)
+        if (e.statusCode() == 403 ||
+                e instanceof KmsAccessDeniedException) {
+            sqsMessageAccessDeniedCounter.increment();
+        } else if (e.statusCode() == 404 ||
+                e instanceof QueueDoesNotExistException ||
+                e instanceof KmsNotFoundException) {
+            sqsQueueNotFoundCounter.increment();
+        } else if (e.isThrottlingException() ||
+                e instanceof KmsThrottledException) {
+            sqsMessageThrottledCounter.increment();
+        }
+    }
 }

data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java

@@ -23,6 +23,8 @@
 import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
 import org.opensearch.dataprepper.metrics.PluginMetrics;
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
 import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
@@ -45,6 +47,11 @@
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
 import software.amazon.awssdk.services.sqs.model.SqsException;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.services.sqs.model.KmsAccessDeniedException;
+import software.amazon.awssdk.services.sqs.model.KmsNotFoundException;
+import software.amazon.awssdk.services.sqs.model.KmsThrottledException;
+import software.amazon.awssdk.services.sqs.model.QueueDoesNotExistException;
 import software.amazon.awssdk.services.sts.model.StsException;
 
 import java.io.IOException;
@@ -92,6 +99,7 @@
 import static org.opensearch.dataprepper.plugins.source.s3.SqsWorker.SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME;
 
 @ExtendWith(MockitoExtension.class)
+@MockitoSettings(strictness = Strictness.LENIENT)
 class SqsWorkerTest {
     private SqsClient sqsClient;
     private S3Service s3Service;
@@ -866,6 +874,58 @@ private static String createEventBridgeNotification(final Instant startTime) {
                 "\"reason\":\"PutObject\"}}";
     }
 
+    @Test
+    void testRecordSqsException_KmsThrottledException() {
+        testRecordSqsException(mock(KmsThrottledException.class), SqsWorker.SQS_MESSAGE_THROTTLED_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_403Status() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.statusCode()).thenReturn(403);
+        testRecordSqsException(sqsException, SqsWorker.SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_KmsAccessDeniedException() {
+        testRecordSqsException(mock(KmsAccessDeniedException.class), SqsWorker.SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_404Status() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.statusCode()).thenReturn(404);
+        testRecordSqsException(sqsException, SqsWorker.SQS_QUEUE_NOT_FOUND_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_QueueDoesNotExistException() {
+        testRecordSqsException(mock(QueueDoesNotExistException.class), SqsWorker.SQS_QUEUE_NOT_FOUND_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_KmsNotFoundException() {
+        testRecordSqsException(mock(KmsNotFoundException.class), SqsWorker.SQS_QUEUE_NOT_FOUND_METRIC_NAME);
+    }
+
+    @Test
+    void testRecordSqsException_isThrottlingException() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.isThrottlingException()).thenReturn(true);
+        testRecordSqsException(sqsException, SqsWorker.SQS_MESSAGE_THROTTLED_METRIC_NAME);
+    }
+
+    private void testRecordSqsException(AwsServiceException exception, String metricName) {
+        final Counter expectedCounter = mock(Counter.class);
+        when(pluginMetrics.counter(metricName)).thenReturn(expectedCounter);
+        final SqsWorker testWorker = createObjectUnderTest();
+
+        testWorker.recordSqsException(exception);
+
+        verify(expectedCounter).increment();
+        verifyNoMoreInteractions(expectedCounter);
+    }
+
     private static String createSecurityLakeNotification(final Instant startTime) {
         return "{\"source\":\"aws.s3\",\"time\":\"" + startTime + "\",\"account\":\"123456789012\",\"region\":\"ca-central-1\"," +
                 "\"resources\":[\"arn:aws:s3:::example-bucket\"],\"detail\":{\"bucket\":{\"name\":\"example-bucket\"}," +