Rename/add metrics for ddb source (opensearch-project#3498)

Signed-off-by: Taylor Gray <[email protected]>
daixba · Oct 16, 2023 · 910f451 · 910f451
1 parent 7605763
commit 910f451
Show file tree

Hide file tree

Showing 14 changed files with 122 additions and 62 deletions.
diff --git a/...rce/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java b/...rce/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java
@@ -102,7 +102,7 @@ public void start(Buffer<Record<Event>> buffer) {
         Runnable fileLoaderScheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics);
 
         ShardConsumerFactory consumerFactory = new ShardConsumerFactory(coordinator, dynamoDbStreamsClient, pluginMetrics, shardManager, buffer);
-        Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager);
+        Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics);
 
         // May consider start or shutdown the scheduler on demand
         // Currently, event after the exports are done, the related scheduler will not be shutdown

diff --git a/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java b/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java
@@ -26,8 +26,8 @@ public class ExportRecordConverter extends RecordConverter {
 
     private static final String ITEM_KEY = "Item";
 
-    static final String EXPORT_RECORD_SUCCESS_COUNT = "exportRecordSuccess";
-    static final String EXPORT_RECORD_ERROR_COUNT = "exportRecordErrors";
+    static final String EXPORT_RECORDS_PROCESSED_COUNT = "exportRecordsProcessed";
+    static final String EXPORT_RECORDS_PROCESSING_ERROR_COUNT = "exportRecordProcessingErrors";
 
 
     IonObjectMapper MAPPER = new IonObjectMapper();
@@ -40,8 +40,8 @@ public class ExportRecordConverter extends RecordConverter {
     public ExportRecordConverter(Buffer<Record<Event>> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) {
         super(buffer, tableInfo);
         this.pluginMetrics = pluginMetrics;
-        this.exportRecordSuccessCounter = pluginMetrics.counter(EXPORT_RECORD_SUCCESS_COUNT);
-        this.exportRecordErrorCounter = pluginMetrics.counter(EXPORT_RECORD_ERROR_COUNT);
+        this.exportRecordSuccessCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSED_COUNT);
+        this.exportRecordErrorCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSING_ERROR_COUNT);
 
     }
 

diff --git a/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java b/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java
@@ -6,7 +6,7 @@
 package org.opensearch.dataprepper.plugins.source.dynamodb.converter;
 
 public class MetadataKeyAttributes {
-     static final String COMPOSITE_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "_id";
+     static final String PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "primary_key";
 
      static final String PARTITION_KEY_METADATA_ATTRIBUTE = "partition_key";
 

diff --git a/...in/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/...in/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java
@@ -17,7 +17,7 @@
 import java.util.List;
 import java.util.Map;
 
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.COMPOSITE_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE;
 import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TABLE_NAME_METADATA_ATTRIBUTE;
 import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE;
 import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE;
@@ -86,7 +86,7 @@ public Record<Event> convertToEvent(Map<String, Object> data, Instant eventCreat
         }
 
         eventMetadata.setAttribute(STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(streamEventName));
-        eventMetadata.setAttribute(COMPOSITE_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, getId(data));
+        eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, getId(data));
         eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, getPartitionKey(data));
         eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, getSortKey(data));
 

diff --git a/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java b/...a/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java
@@ -26,8 +26,8 @@ public class StreamRecordConverter extends RecordConverter {
     private static final Logger LOG = LoggerFactory.getLogger(StreamRecordConverter.class);
 
 
-    static final String CHANGE_EVENT_SUCCESS_COUNT = "changeEventSuccess";
-    static final String CHANGE_EVENT_ERROR_COUNT = "changeEventErrors";
+    static final String CHANGE_EVENTS_PROCESSED_COUNT = "changeEventsProcessed";
+    static final String CHANGE_EVENTS_PROCESSING_ERROR_COUNT = "changeEventsProcessingErrors";
 
     private static final ObjectMapper MAPPER = new ObjectMapper();
 
@@ -42,8 +42,8 @@ public class StreamRecordConverter extends RecordConverter {
     public StreamRecordConverter(Buffer<Record<Event>> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) {
         super(buffer, tableInfo);
         this.pluginMetrics = pluginMetrics;
-        this.changeEventSuccessCounter = pluginMetrics.counter(CHANGE_EVENT_SUCCESS_COUNT);
-        this.changeEventErrorCounter = pluginMetrics.counter(CHANGE_EVENT_ERROR_COUNT);
+        this.changeEventSuccessCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSED_COUNT);
+        this.changeEventErrorCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSING_ERROR_COUNT);
     }
 
     @Override

diff --git a/...ain/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/...ain/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java
@@ -22,6 +22,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
 
 
@@ -41,7 +42,8 @@ public class DataFileScheduler implements Runnable {
      */
     private static final int DEFAULT_LEASE_INTERVAL_MILLIS = 15_000;
 
-    static final String EXPORT_FILE_SUCCESS_COUNT = "exportFileSuccess";
+    static final String EXPORT_S3_OBJECTS_PROCESSED_COUNT = "exportS3ObjectsProcessed";
+    static final String ACTIVE_EXPORT_S3_OBJECT_CONSUMERS_GAUGE = "activeExportS3ObjectConsumers";
 
 
     private final EnhancedSourceCoordinator coordinator;
@@ -54,6 +56,7 @@ public class DataFileScheduler implements Runnable {
 
 
     private final Counter exportFileSuccessCounter;
+    private final AtomicLong activeExportS3ObjectConsumersGauge;
 
 
     public DataFileScheduler(EnhancedSourceCoordinator coordinator, DataFileLoaderFactory loaderFactory, PluginMetrics pluginMetrics) {
@@ -64,7 +67,8 @@ public DataFileScheduler(EnhancedSourceCoordinator coordinator, DataFileLoaderFa
 
         executor = Executors.newFixedThreadPool(MAX_JOB_COUNT);
 
-        this.exportFileSuccessCounter = pluginMetrics.counter(EXPORT_FILE_SUCCESS_COUNT);
+        this.exportFileSuccessCounter = pluginMetrics.counter(EXPORT_S3_OBJECTS_PROCESSED_COUNT);
+        this.activeExportS3ObjectConsumersGauge = pluginMetrics.gauge(ACTIVE_EXPORT_S3_OBJECT_CONSUMERS_GAUGE, new AtomicLong());
     }
 
     private void processDataFilePartition(DataFilePartition dataFilePartition) {
@@ -83,13 +87,15 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) {
     public void run() {
         LOG.info("Start running Data File Scheduler");
 
-        while (!Thread.interrupted()) {
+        while (!Thread.currentThread().isInterrupted()) {
             if (numOfWorkers.get() < MAX_JOB_COUNT) {
                 final Optional<EnhancedSourcePartition> sourcePartition = coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE);
 
                 if (sourcePartition.isPresent()) {
+                    activeExportS3ObjectConsumersGauge.incrementAndGet();
                     DataFilePartition dataFilePartition = (DataFilePartition) sourcePartition.get();
                     processDataFilePartition(dataFilePartition);
+                    activeExportS3ObjectConsumersGauge.decrementAndGet();
                 }
             }
             try {

diff --git a/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java
@@ -51,8 +51,8 @@ public class ExportScheduler implements Runnable {
     private static final String FAILED_STATUS = "Failed";
 
     static final String EXPORT_JOB_SUCCESS_COUNT = "exportJobSuccess";
-    static final String EXPORT_JOB_ERROR_COUNT = "exportJobErrors";
-    static final String EXPORT_FILES_TOTAL_COUNT = "exportFilesTotal";
+    static final String EXPORT_JOB_FAILURE_COUNT = "exportJobFailure";
+    static final String EXPORT_S3_OBJECTS_TOTAL_COUNT = "exportS3ObjectsTotal";
     static final String EXPORT_RECORDS_TOTAL_COUNT = "exportRecordsTotal";
 
     private final PluginMetrics pluginMetrics;
@@ -68,9 +68,9 @@ public class ExportScheduler implements Runnable {
     private final ExportTaskManager exportTaskManager;
 
     private final Counter exportJobSuccessCounter;
-    private final Counter exportJobErrorCounter;
+    private final Counter exportJobFailureCounter;
 
-    private final Counter exportFilesTotalCounter;
+    private final Counter exportS3ObjectsTotalCounter;
     private final Counter exportRecordsTotalCounter;
 
     public ExportScheduler(EnhancedSourceCoordinator enhancedSourceCoordinator, DynamoDbClient dynamoDBClient, ManifestFileReader manifestFileReader, PluginMetrics pluginMetrics) {
@@ -83,17 +83,17 @@ public ExportScheduler(EnhancedSourceCoordinator enhancedSourceCoordinator, Dyna
         executor = Executors.newCachedThreadPool();
 
         exportJobSuccessCounter = pluginMetrics.counter(EXPORT_JOB_SUCCESS_COUNT);
-        exportJobErrorCounter = pluginMetrics.counter(EXPORT_JOB_ERROR_COUNT);
-        exportFilesTotalCounter = pluginMetrics.counter(EXPORT_FILES_TOTAL_COUNT);
+        exportJobFailureCounter = pluginMetrics.counter(EXPORT_JOB_FAILURE_COUNT);
+        exportS3ObjectsTotalCounter = pluginMetrics.counter(EXPORT_S3_OBJECTS_TOTAL_COUNT);
         exportRecordsTotalCounter = pluginMetrics.counter(EXPORT_RECORDS_TOTAL_COUNT);
 
 
     }
 
     @Override
     public void run() {
-        LOG.info("Start running Export Scheduler");
-        while (!Thread.interrupted()) {
+        LOG.debug("Start running Export Scheduler");
+        while (!Thread.currentThread().isInterrupted()) {
             // Does not have limit on max leases
             // As most of the time it's just to wait
             final Optional<EnhancedSourcePartition> sourcePartition = enhancedSourceCoordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE);
@@ -186,7 +186,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map<S
             enhancedSourceCoordinator.createPartition(partition);
         });
 
-        exportFilesTotalCounter.increment(totalFiles.get());
+        exportS3ObjectsTotalCounter.increment(totalFiles.get());
         exportRecordsTotalCounter.increment(totalRecords.get());
 
         // Currently, we need to maintain a global state to track the overall progress.
@@ -197,7 +197,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map<S
 
 
     private void closeExportPartitionWithError(ExportPartition exportPartition) {
-        exportJobErrorCounter.increment(1);
+        exportJobFailureCounter.increment(1);
         ExportProgressState exportProgressState = exportPartition.getProgressState().get();
         // Clear current Arn, so that a new export can be submitted.
         exportProgressState.setExportArn(null);

diff --git a/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java b/.../main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
 
+import org.opensearch.dataprepper.metrics.PluginMetrics;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
 import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition;
 import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition;
@@ -18,6 +19,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
 
 /**
@@ -41,20 +43,28 @@ public class StreamScheduler implements Runnable {
      */
     private static final int DELAY_TO_GET_CHILD_SHARDS_MILLIS = 1_500;
 
+    static final String ACTIVE_CHANGE_EVENT_CONSUMERS = "activeChangeEventConsumers";
+
     private final AtomicInteger numOfWorkers = new AtomicInteger(0);
     private final EnhancedSourceCoordinator coordinator;
     private final ShardConsumerFactory consumerFactory;
     private final ExecutorService executor;
     private final ShardManager shardManager;
+    private final PluginMetrics pluginMetrics;
+    private final AtomicLong activeChangeEventConsumers;
 
 
-    public StreamScheduler(final EnhancedSourceCoordinator coordinator, final ShardConsumerFactory consumerFactory, final ShardManager shardManager) {
+    public StreamScheduler(final EnhancedSourceCoordinator coordinator,
+                           final ShardConsumerFactory consumerFactory,
+                           final ShardManager shardManager,
+                           final PluginMetrics pluginMetrics) {
         this.coordinator = coordinator;
         this.shardManager = shardManager;
         this.consumerFactory = consumerFactory;
+        this.pluginMetrics = pluginMetrics;
 
         executor = Executors.newFixedThreadPool(MAX_JOB_COUNT);
-
+        activeChangeEventConsumers = pluginMetrics.gauge(ACTIVE_CHANGE_EVENT_CONSUMERS, new AtomicLong());
     }
 
     private void processStreamPartition(StreamPartition streamPartition) {
@@ -71,13 +81,15 @@ private void processStreamPartition(StreamPartition streamPartition) {
 
     @Override
     public void run() {
-        LOG.info("Start running Stream Scheduler");
-        while (!Thread.interrupted()) {
+        LOG.debug("Stream Scheduler start to run...");
+        while (!Thread.currentThread().isInterrupted()) {
             if (numOfWorkers.get() < MAX_JOB_COUNT) {
                 final Optional<EnhancedSourcePartition> sourcePartition = coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE);
                 if (sourcePartition.isPresent()) {
+                    activeChangeEventConsumers.incrementAndGet();
                     StreamPartition streamPartition = (StreamPartition) sourcePartition.get();
                     processStreamPartition(streamPartition);
+                    activeChangeEventConsumers.decrementAndGet();
                 }
             }
 

diff --git a/...g/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java b/...g/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java
@@ -33,8 +33,8 @@
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORD_ERROR_COUNT;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORD_SUCCESS_COUNT;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSING_ERROR_COUNT;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSED_COUNT;
 
 @ExtendWith(MockitoExtension.class)
 class ExportRecordConverterTest {
@@ -73,8 +73,8 @@ void setup() {
 
         tableInfo = new TableInfo(tableArn, metadata);
 
-        given(pluginMetrics.counter(EXPORT_RECORD_SUCCESS_COUNT)).willReturn(exportRecordSuccess);
-        given(pluginMetrics.counter(EXPORT_RECORD_ERROR_COUNT)).willReturn(exportRecordErrors);
+        given(pluginMetrics.counter(EXPORT_RECORDS_PROCESSED_COUNT)).willReturn(exportRecordSuccess);
+        given(pluginMetrics.counter(EXPORT_RECORDS_PROCESSING_ERROR_COUNT)).willReturn(exportRecordErrors);
 
     }
 

diff --git a/...g/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java b/...g/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java
@@ -38,8 +38,8 @@
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENT_ERROR_COUNT;
-import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENT_SUCCESS_COUNT;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENTS_PROCESSING_ERROR_COUNT;
+import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENTS_PROCESSED_COUNT;
 
 @ExtendWith(MockitoExtension.class)
 class StreamRecordConverterTest {
@@ -77,8 +77,8 @@ void setup() {
 
         tableInfo = new TableInfo(tableArn, metadata);
 
-        given(pluginMetrics.counter(CHANGE_EVENT_SUCCESS_COUNT)).willReturn(changeEventSuccessCounter);
-        given(pluginMetrics.counter(CHANGE_EVENT_ERROR_COUNT)).willReturn(changeEventErrorCounter);
+        given(pluginMetrics.counter(CHANGE_EVENTS_PROCESSED_COUNT)).willReturn(changeEventSuccessCounter);
+        given(pluginMetrics.counter(CHANGE_EVENTS_PROCESSING_ERROR_COUNT)).willReturn(changeEventErrorCounter);
 
     }
 

diff --git a/...st/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java b/...st/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java
@@ -6,6 +6,7 @@
 package org.opensearch.dataprepper.plugins.source.dynamodb.export;
 
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -30,14 +31,21 @@
 import java.util.Random;
 import java.util.StringJoiner;
 import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.zip.GZIPOutputStream;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
+@Disabled
 class DataFileLoaderTest {
 
     @Mock
@@ -128,9 +136,12 @@ void test_run_loadFile_correctly() throws InterruptedException {
                 .checkpointer(checkpointer)
                 .build();
 
-        loader.run();
-        // Run for a while
-        Thread.sleep(1000);
+        ExecutorService executorService = Executors.newSingleThreadExecutor();
+        final Future<?> future = executorService.submit(loader);
+        Thread.sleep(100);
+        executorService.shutdown();
+        future.cancel(true);
+        assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true));
 
         // Should call s3 getObject
         verify(s3Client).getObject(any(GetObjectRequest.class));