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";
+    static final String SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME = "sqsMessagesAccessDenied";
+    static final String SQS_MESSAGE_THROTTLED_METRIC_NAME = "sqsMessagesThrottled";
+    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;
     }
+
+    private 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;
@@ -71,8 +78,10 @@
 import static org.junit.jupiter.params.provider.Arguments.arguments;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -866,6 +875,54 @@ private static String createEventBridgeNotification(final Instant startTime) {
                 "\"reason\":\"PutObject\"}}";
     }
 
+    @Test
+    void run_increments_access_denied_metric_on_403_exception() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.statusCode()).thenReturn(403);
+        when(sqsClient.receiveMessage(any(ReceiveMessageRequest.class))).thenThrow(sqsException);
+        
+        final Counter accessDeniedCounter = mock(Counter.class);
+        when(pluginMetrics.counter(SqsWorker.SQS_MESSAGE_ACCESS_DENIED_METRIC_NAME)).thenReturn(accessDeniedCounter);
+        lenient().when(pluginMetrics.counter(anyString())).thenReturn(mock(Counter.class));
+
+        final SqsWorker sqsWorker = createObjectUnderTest();
+        sqsWorker.run();
+
+        verify(accessDeniedCounter).increment();
+    }
+
+    @Test
+    void run_increments_queue_not_found_metric_on_404_exception() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.statusCode()).thenReturn(404);
+        when(sqsClient.receiveMessage(any(ReceiveMessageRequest.class))).thenThrow(sqsException);
+        
+        final Counter queueNotFoundCounter = mock(Counter.class);
+        when(pluginMetrics.counter(SqsWorker.SQS_QUEUE_NOT_FOUND_METRIC_NAME)).thenReturn(queueNotFoundCounter);
+        lenient().when(pluginMetrics.counter(anyString())).thenReturn(mock(Counter.class));
+
+        final SqsWorker sqsWorker = createObjectUnderTest();
+        sqsWorker.run();
+
+        verify(queueNotFoundCounter).increment();
+    }
+
+    @Test
+    void run_increments_throttled_metric_on_throttling_exception() {
+        final SqsException sqsException = mock(SqsException.class);
+        when(sqsException.isThrottlingException()).thenReturn(true);
+        when(sqsClient.receiveMessage(any(ReceiveMessageRequest.class))).thenThrow(sqsException);
+        
+        final Counter throttledCounter = mock(Counter.class);
+        when(pluginMetrics.counter(SqsWorker.SQS_MESSAGE_THROTTLED_METRIC_NAME)).thenReturn(throttledCounter);
+        lenient().when(pluginMetrics.counter(anyString())).thenReturn(mock(Counter.class));
+
+        final SqsWorker sqsWorker = createObjectUnderTest();
+        sqsWorker.run();
+
+        verify(throttledCounter).increment();
+    }
+
     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\"}," +