add partition timing metrics to LeaderOnlyTokenCrawler (opensearch-project#6299)

add partition timing metrics to LeaderOnlyTokenCrawler

Author: wandna-amazon

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java

@@ -473,9 +473,8 @@ void testStartSubscriptionsRecordsFailureMetrics() {
         // Verify that failure metrics were recorded
         // Note: We skip verifying recordSubscriptionLatency parameter since lambda expressions
         // get compiled to specific classes that are difficult to match in tests
-        // The retry handler calls recordSubscriptionFailure() for each retry attempt.
-        // With 6 content types and multiple retries per content type, we expect multiple failure calls
-        verify(metricsRecorder, times(6)).recordSubscriptionFailure();
+        // The retry handler calls recordSubscriptionFailure() once after all retries are exhausted
+        verify(metricsRecorder, times(1)).recordSubscriptionFailure();
     }
 
     /**

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/DimensionalTimeSliceCrawler.java

@@ -39,9 +39,9 @@ public class DimensionalTimeSliceCrawler implements Crawler<DimensionalTimeSlice
     // delay five minutes for partition creation on latest time duration to ensure the newly generated events are queryable
     // In general, newly generated events become queryable after 30 ~ 120 second
     protected static final long WAIT_SECONDS_BEFORE_PARTITION_CREATION = 300;
-    private static final String DIMENSIONAL_TIME_SLICE_WORKER_PARTITIONS_CREATED = "DimensionalTimeSliceWorkerPartitionsCreated";
-    private static final String WORKER_PARTITION_WAIT_TIME = "WorkerPartitionWaitTime";
-    private static final String WORKER_PARTITION_PROCESS_LATENCY = "WorkerPartitionProcessLatency";
+    private static final String DIMENSIONAL_TIME_SLICE_WORKER_PARTITIONS_CREATED = "dimensionalTimeSliceWorkerPartitionsCreated";
+    private static final String WORKER_PARTITION_WAIT_TIME = "workerPartitionWaitTime";
+    private static final String WORKER_PARTITION_PROCESS_LATENCY = "workerPartitionProcessLatency";
     private static final Duration HOUR_DURATION = Duration.ofHours(1);
 
     private final CrawlerClient client;

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/LeaderOnlyTokenCrawler.java

@@ -39,6 +39,8 @@ public class LeaderOnlyTokenCrawler implements Crawler<PaginationCrawlerWorkerPr
 
     private static final String METRIC_BATCHES_FAILED = "batchesFailed";
     private static final String METRIC_BUFFER_WRITE_TIME = "bufferWriteTime";
+    private static final String WORKER_PARTITION_WAIT_TIME = "workerPartitionWaitTime";
+    private static final String WORKER_PARTITION_PROCESS_LATENCY = "workerPartitionProcessLatency";
     public static final String ACKNOWLEDGEMENT_SET_SUCCESS_METRIC_NAME = "acknowledgementSetSuccesses";
     public static final String ACKNOWLEDGEMENT_SET_FAILURES_METRIC_NAME = "acknowledgementSetFailures";
 
@@ -55,6 +57,8 @@ public class LeaderOnlyTokenCrawler implements Crawler<PaginationCrawlerWorkerPr
     private final Counter acknowledgementSetSuccesses;
     private final Counter acknowledgementSetFailures;
     private final Timer bufferWriteTimer;
+    private final Timer partitionWaitTimeTimer;
+    private final Timer partitionProcessLatencyTimer;
 
     private String lastToken;
     private Duration noAckTimeout;
@@ -67,6 +71,8 @@ public LeaderOnlyTokenCrawler(
         this.crawlingTimer = pluginMetrics.timer("crawlingTime");
         this.batchesFailedCounter = pluginMetrics.counter(METRIC_BATCHES_FAILED);
         this.bufferWriteTimer = pluginMetrics.timer(METRIC_BUFFER_WRITE_TIME);
+        this.partitionWaitTimeTimer = pluginMetrics.timer(WORKER_PARTITION_WAIT_TIME);
+        this.partitionProcessLatencyTimer = pluginMetrics.timer(WORKER_PARTITION_PROCESS_LATENCY);
         this.acknowledgementSetSuccesses = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_SUCCESS_METRIC_NAME);
         this.acknowledgementSetFailures = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_FAILURES_METRIC_NAME);
         this.noAckTimeout = NO_ACK_TIME_OUT_SECONDS;
@@ -123,7 +129,8 @@ public Instant crawl(LeaderPartition leaderPartition,
 
     @Override
     public void executePartition(PaginationCrawlerWorkerProgressState state, Buffer buffer, AcknowledgementSet acknowledgementSet) {
-        client.executePartition(state, buffer, acknowledgementSet);
+        partitionWaitTimeTimer.record(Duration.between(state.getExportStartTime(), Instant.now()));
+        partitionProcessLatencyTimer.record(() -> client.executePartition(state, buffer, acknowledgementSet));
     }
 
     private List<ItemInfo> collectBatch(Iterator<ItemInfo> iterator) {

data-prepper-plugins/saas-source-plugins/source-crawler/src/main/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/TokenPaginationCrawler.java

@@ -33,8 +33,8 @@ public class TokenPaginationCrawler implements Crawler<PaginationCrawlerWorkerPr
     private static final int batchSize = 50;
     private static final String PAGINATION_WORKER_PARTITIONS_CREATED = "paginationWorkerPartitionsCreated";
     private static final String INVALID_PAGINATION_ITEMS = "invalidPaginationItems";
-    private static final String WORKER_PARTITION_WAIT_TIME = "WorkerPartitionWaitTime";
-    private static final String WORKER_PARTITION_PROCESS_LATENCY = "WorkerPartitionProcessLatency";
+    private static final String WORKER_PARTITION_WAIT_TIME = "workerPartitionWaitTime";
+    private static final String WORKER_PARTITION_PROCESS_LATENCY = "workerPartitionProcessLatency";
     private final Timer crawlingTimer;
     private final Timer partitionWaitTimeTimer;
     private final Timer partitionProcessLatencyTimer;

data-prepper-plugins/saas-source-plugins/source-crawler/src/test/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/DimensionalTimeSliceCrawlerTest.java

@@ -75,8 +75,8 @@ public class DimensionalTimeSliceCrawlerTest {
     @BeforeEach
     void setUp() {
         when(pluginMetrics.counter(anyString())).thenReturn(partitionsCreatedCounter);
-        when(pluginMetrics.timer("WorkerPartitionWaitTime")).thenReturn(partitionWaitTimeTimer);
-        when(pluginMetrics.timer("WorkerPartitionProcessLatency")).thenReturn(partitionProcessLatencyTimer);
+        when(pluginMetrics.timer("workerPartitionWaitTime")).thenReturn(partitionWaitTimeTimer);
+        when(pluginMetrics.timer("workerPartitionProcessLatency")).thenReturn(partitionProcessLatencyTimer);
         crawler = new DimensionalTimeSliceCrawler(client, pluginMetrics);
         crawler.initialize(LOG_TYPES);
     }

data-prepper-plugins/saas-source-plugins/source-crawler/src/test/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/LeaderOnlyTokenCrawlerTest.java

@@ -22,6 +22,8 @@
 import org.opensearch.dataprepper.plugins.source.source_crawler.coordination.state.TokenPaginationCrawlerLeaderProgressState;
 import org.opensearch.dataprepper.plugins.source.source_crawler.model.ItemInfo;
 import org.opensearch.dataprepper.plugins.source.source_crawler.model.TestItemInfo;
+import io.micrometer.core.instrument.Timer;
+import io.micrometer.core.instrument.Counter;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -42,6 +44,9 @@
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.internal.verification.VerificationModeFactory.times;
 
 @ExtendWith(MockitoExtension.class)
@@ -62,6 +67,8 @@ class LeaderOnlyTokenCrawlerTest {
     private AcknowledgementSet acknowledgementSet;
     @Mock
     private Buffer<Record<Event>> buffer;
+    @Mock
+    private PaginationCrawlerWorkerProgressState workerState;
 
     private LeaderOnlyTokenCrawler crawler;
     private final PluginMetrics pluginMetrics = PluginMetrics.fromNames("CrawlerTest", "crawler");
@@ -241,6 +248,50 @@ void testAcknowledgmentTimeout() {
         verify(coordinator).createPartition(any());
     }
 
+    @Test
+    public void testExecutePartitionMetrics() {
+        reset(leaderPartition);
+
+        // mock timers and counters
+        Timer mockCrawlingTimer = mock(Timer.class);
+        Timer partitionWaitTimeTimer = mock(Timer.class);
+        Timer partitionProcessLatencyTimer = mock(Timer.class);
+        Timer mockBufferWriteTimer = mock(Timer.class);
+        Counter mockBatchesFailedCounter = mock(Counter.class);
+        Counter mockAcknowledgementSetSuccesses = mock(Counter.class);
+        Counter mockAcknowledgementSetFailures = mock(Counter.class);
+
+        // setup mock plugin metrics
+        PluginMetrics mockPluginMetrics = mock(PluginMetrics.class);
+        when(mockPluginMetrics.timer("crawlingTime")).thenReturn(mockCrawlingTimer);
+        when(mockPluginMetrics.timer("workerPartitionWaitTime")).thenReturn(partitionWaitTimeTimer);
+        when(mockPluginMetrics.timer("workerPartitionProcessLatency")).thenReturn(partitionProcessLatencyTimer);
+        when(mockPluginMetrics.timer("bufferWriteTime")).thenReturn(mockBufferWriteTimer);
+        when(mockPluginMetrics.counter("batchesFailed")).thenReturn(mockBatchesFailedCounter);
+        when(mockPluginMetrics.counter("acknowledgementSetSuccesses")).thenReturn(mockAcknowledgementSetSuccesses);
+        when(mockPluginMetrics.counter("acknowledgementSetFailures")).thenReturn(mockAcknowledgementSetFailures);
+
+        LeaderOnlyTokenCrawler testCrawler = new LeaderOnlyTokenCrawler(client, mockPluginMetrics);
+
+        // test executePartition with metrics
+        when(workerState.getExportStartTime()).thenReturn(Instant.now().minusSeconds(1));
+
+        // make latency timer execute the runnable so client.executePartition() gets called
+        doAnswer(invocation -> {
+            Runnable runnable = invocation.getArgument(0);
+            runnable.run();
+            return null;
+        }).when(partitionProcessLatencyTimer).record(any(Runnable.class));
+        
+        doNothing().when(partitionWaitTimeTimer).record(any(Duration.class));
+
+        testCrawler.executePartition(workerState, buffer, acknowledgementSet);
+
+        // verify metrics are recorded
+        verify(partitionProcessLatencyTimer).record(any(Runnable.class));
+        verify(partitionWaitTimeTimer).record(any(Duration.class));
+        verify(client).executePartition(workerState, buffer, acknowledgementSet);
+    }
 
     private List<ItemInfo> createTestItems(int count) {
         List<ItemInfo> items = new ArrayList<>();
@@ -249,4 +300,4 @@ private List<ItemInfo> createTestItems(int count) {
         }
         return items;
     }
-}
+}

data-prepper-plugins/saas-source-plugins/source-crawler/src/test/java/org/opensearch/dataprepper/plugins/source/source_crawler/base/TokenPaginationCrawlerTest.java

@@ -102,8 +102,8 @@ public void test_metrics_in_crawler() {
         Counter mockInvalidItemsCounter = mock(Counter.class);
 
         when(mockPluginMetrics.timer("crawlingTime")).thenReturn(mockCrawlingTimer);
-        when(mockPluginMetrics.timer("WorkerPartitionWaitTime")).thenReturn(partitionWaitTimeTimer);
-        when(mockPluginMetrics.timer("WorkerPartitionProcessLatency")).thenReturn(partitionProcessLatencyTimer);
+        when(mockPluginMetrics.timer("workerPartitionWaitTime")).thenReturn(partitionWaitTimeTimer);
+        when(mockPluginMetrics.timer("workerPartitionProcessLatency")).thenReturn(partitionProcessLatencyTimer);
         when(mockPluginMetrics.counter("paginationWorkerPartitionsCreated")).thenReturn(mockPartitionsCreatedCounter);
         when(mockPluginMetrics.counter("invalidPaginationItems")).thenReturn(mockInvalidItemsCounter);