Add read failure metric to 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/S3ObjectPluginMetrics.java

@@ -13,6 +13,7 @@ public class S3ObjectPluginMetrics {
     static final String S3_OBJECTS_SIZE_PROCESSED = "s3ObjectProcessedBytes";
     static final String S3_OBJECTS_FAILED_METRIC_NAME = "s3ObjectsFailed";
     static final String S3_OBJECTS_DELETE_FAILED_METRIC_NAME = "s3ObjectsDeleteFailed";
+    static final String S3_OBJECTS_READ_FAILED_METRIC_NAME = "s3ObjectReadFailed";
     static final String S3_OBJECTS_SUCCEEDED_METRIC_NAME = "s3ObjectsSucceeded";
     static final String S3_OBJECTS_EVENTS = "s3ObjectsEvents";
     static final String S3_OBJECTS_FAILED_NOT_FOUND_METRIC_NAME = "s3ObjectsNotFound";
@@ -33,6 +34,7 @@ public class S3ObjectPluginMetrics {
     private final Counter s3ObjectNoRecordsFound;
 
     private final Counter s3ObjectsDeleteFailed;
+    private final Counter s3ObjectReadFailedCounter;
 
     public S3ObjectPluginMetrics(final PluginMetrics pluginMetrics){
         s3ObjectsFailedCounter = pluginMetrics.counter(S3_OBJECTS_FAILED_METRIC_NAME);
@@ -46,6 +48,7 @@ public S3ObjectPluginMetrics(final PluginMetrics pluginMetrics){
         s3ObjectEventsSummary = pluginMetrics.summary(S3_OBJECTS_EVENTS);
         s3ObjectNoRecordsFound = pluginMetrics.counter(S3_OBJECTS_NO_RECORDS_FOUND);
         s3ObjectsDeleteFailed = pluginMetrics.counter(S3_OBJECTS_DELETE_FAILED_METRIC_NAME);
+        s3ObjectReadFailedCounter = pluginMetrics.counter(S3_OBJECTS_READ_FAILED_METRIC_NAME);
     }
 
     public Counter getS3ObjectsFailedCounter() {
@@ -89,4 +92,6 @@ public Counter getS3ObjectsThrottledCounter() {
     }
 
     public Counter getS3ObjectsDeleteFailed() { return s3ObjectsDeleteFailed; }
+
+    public Counter getS3ObjectReadFailedCounter() { return s3ObjectReadFailedCounter; }
 }

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

@@ -155,10 +155,15 @@ public long consumeS3Object(final S3InputFile inputFile, final S3DataSelection d
             final CompressionOption fileCompressionOption = compressionOption != CompressionOption.AUTOMATIC ?
                     compressionOption : CompressionOption.fromFileName(s3ObjectReference.getKey());
 
-            codec.parse(inputFile, fileCompressionOption.getDecompressionEngine(), record -> {
-                consumer.accept(record, dataSelection);
-            });
-            return inputFile.getLength();
+            try {
+                codec.parse(inputFile, fileCompressionOption.getDecompressionEngine(), record -> {
+                    consumer.accept(record, dataSelection);
+                });
+                return inputFile.getLength();
+            } catch (final Exception e) {
+                s3ObjectPluginMetrics.getS3ObjectReadFailedCounter().increment();
+                throw new S3ReadFailedException(e);
+            }
         }
     }
 

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.source.s3;
+
+/**
+ * Exception thrown when there is a failure reading from S3 objects.
+ * This is used to distinguish actual S3 read failures from other processing failures.
+ */
+public class S3ReadFailedException extends RuntimeException {
+    public S3ReadFailedException(final Throwable cause) {
+        super(cause);
+    }
+}

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

@@ -46,5 +46,6 @@ public void s3ObjectPluginMetricsTest(){
         assertThat(metrics.getS3ObjectsFailedAccessDeniedCounter(),sameInstance(counter));
         assertThat(metrics.getS3ObjectsFailedNotFoundCounter(),sameInstance(counter));
         assertThat(metrics.getS3ObjectsThrottledCounter(),sameInstance(counter));
+        assertThat(metrics.getS3ObjectReadFailedCounter(),sameInstance(counter));
     }
 }

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

@@ -289,11 +289,12 @@ void S3ObjectWorker_with_MetadataOnly_Test() throws Exception {
     void processS3Object_codec_parse_exception() throws Exception {
         when(s3Client.headObject(any(HeadObjectRequest.class))).thenReturn(headObjectResponse);
         when(s3ObjectPluginMetrics.getS3ObjectsFailedCounter()).thenReturn(s3ObjectsFailedCounter);
+        when(s3ObjectPluginMetrics.getS3ObjectReadFailedCounter()).thenReturn(mock(Counter.class));
 
         doThrow(IOException.class).when(codec).parse(any(InputFile.class), any(DecompressionEngine.class), any(Consumer.class));
 
         assertThrows(
-            IOException.class,
+            S3ReadFailedException.class,
             () -> createObjectUnderTest(s3ObjectPluginMetrics).processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null));
 
         final ArgumentCaptor<InputFile> inputFileArgumentCaptor = ArgumentCaptor.forClass(InputFile.class);
@@ -302,6 +303,21 @@ void processS3Object_codec_parse_exception() throws Exception {
         assertThat(actualInputFile, instanceOf(S3InputFile.class));
     }
 
+    @Test
+    void processS3Object_increments_s3ObjectReadFailed_counter_when_codec_parse_fails() throws Exception {
+        when(s3Client.headObject(any(HeadObjectRequest.class))).thenReturn(headObjectResponse);
+        when(s3ObjectPluginMetrics.getS3ObjectsFailedCounter()).thenReturn(s3ObjectsFailedCounter);
+        final Counter s3ObjectReadFailedCounter = mock(Counter.class);
+        when(s3ObjectPluginMetrics.getS3ObjectReadFailedCounter()).thenReturn(s3ObjectReadFailedCounter);
+
+        doThrow(IOException.class).when(codec).parse(any(InputFile.class), any(DecompressionEngine.class), any(Consumer.class));
+
+        assertThrows(S3ReadFailedException.class, 
+            () -> createObjectUnderTest(s3ObjectPluginMetrics).processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null));
+
+        verify(s3ObjectReadFailedCounter, times(1)).increment();
+    }
+
     @Test
     void processS3Object_calls_Codec_parse_with_Consumer_that_adds_to_BufferAccumulator() throws Exception {
         when(s3Client.headObject(any(HeadObjectRequest.class))).thenReturn(headObjectResponse);
@@ -415,6 +431,7 @@ void processS3Object_increments_success_counter_after_parsing_S3_object() throws
     @Test
     void processS3Object_throws_Exception_and_increments_failure_counter_when_unable_to_parse_S3_object() throws IOException {
         when(s3ObjectPluginMetrics.getS3ObjectsFailedCounter()).thenReturn(s3ObjectsFailedCounter);
+        when(s3ObjectPluginMetrics.getS3ObjectReadFailedCounter()).thenReturn(mock(Counter.class));
 
         when(s3Client.headObject(any(HeadObjectRequest.class))).thenReturn(headObjectResponse);
 
@@ -423,13 +440,13 @@ void processS3Object_throws_Exception_and_increments_failure_counter_when_unable
                 .when(codec).parse(any(InputFile.class), any(DecompressionEngine.class), any(Consumer.class));
 
         final S3ObjectHandler objectUnderTest = createObjectUnderTest(s3ObjectPluginMetrics);
-        final IOException actualException = assertThrows(IOException.class, () -> objectUnderTest.processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null));
+        final S3ReadFailedException actualException = assertThrows(S3ReadFailedException.class, () -> objectUnderTest.processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null));
 
-        assertThat(actualException, sameInstance(expectedException));
+        assertThat(actualException.getCause(), sameInstance(expectedException));
 
         verify(s3ObjectsFailedCounter).increment();
         verifyNoInteractions(s3ObjectsSucceededCounter);
-        assertThat(exceptionThrownByCallable, sameInstance(expectedException));
+        assertThat(exceptionThrownByCallable, sameInstance(actualException));
     }
 
     @Test

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

@@ -152,6 +152,7 @@ void setup() {
         lenient().when(selectJsonOption.getType()).thenReturn(JSON_LINES_TYPE);
 
         given(s3ObjectRequest.getS3ObjectPluginMetrics()).willReturn(s3ObjectPluginMetrics);
+        lenient().when(s3ObjectPluginMetrics.getS3ObjectReadFailedCounter()).thenReturn(mock(Counter.class));
         bucketOwnerProvider = mock(BucketOwnerProvider.class);
         given(bucketOwnerProvider.getBucketOwner(any(String.class))).willReturn(Optional.of("my-bucket-1"));
         given(s3ObjectRequest.getBucketOwnerProvider()).willReturn(bucketOwnerProvider);