KCL metrics management through config

RashmiRam · RashmiRam · commit 496230d324a8 · 2026-01-21T10:14:12.000+05:30
1.Added "metrics" config to disable/enable sending KCL metrics to cloudwatch
2.Default value is false.

Signed-off-by: RashmiRam &lt;ras.xena@gmail.com&gt;
diff --git a/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java b/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java
@@ -42,6 +42,8 @@
 import software.amazon.kinesis.coordinator.Scheduler;
 import software.amazon.kinesis.exceptions.KinesisClientLibDependencyException;
 import software.amazon.kinesis.exceptions.ThrottlingException;
+import software.amazon.kinesis.metrics.MetricsConfig;
+import software.amazon.kinesis.metrics.NullMetricsFactory;
 import software.amazon.kinesis.processor.ShardRecordProcessorFactory;
 import software.amazon.kinesis.retrieval.RetrievalConfig;
 import software.amazon.kinesis.retrieval.polling.PollingConfig;
@@ -198,14 +200,18 @@ public Scheduler createScheduler(final Buffer<Record<Event>> buffer) {
                             kinesisSourceConfig.getPollingConfig().getIdleTimeBetweenReads().toMillis()));
         }
 
+        MetricsConfig metricsConfig = kinesisSourceConfig.isMetrics()
+                ? configsBuilder.metricsConfig()
+                : configsBuilder.metricsConfig().metricsFactory(new NullMetricsFactory());
+
         return new Scheduler(
                 configsBuilder.checkpointConfig(),
                 configsBuilder.coordinatorConfig()
                         .schedulerInitializationBackoffTimeMillis(kinesisSourceConfig.getInitializationBackoffTime().toMillis())
                         .maxInitializationAttempts(kinesisSourceConfig.getMaxInitializationAttempts()),
                 configsBuilder.leaseManagementConfig().billingMode(BillingMode.PAY_PER_REQUEST),
                 configsBuilder.lifecycleConfig(),
-                configsBuilder.metricsConfig(),
+                metricsConfig,
                 configsBuilder.processorConfig(),
                 retrievalConfig
         );
diff --git a/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/configuration/KinesisSourceConfig.java b/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/configuration/KinesisSourceConfig.java
@@ -79,6 +79,10 @@ public Duration getShardAcknowledgmentTimeout() {
     @Getter
     @JsonProperty("initialization_backoff_time")
     private Duration initializationBackoffTime = DEFAULT_INITIALIZATION_BACKOFF_TIME;
+
+    @Getter
+    @JsonProperty("metrics")
+    private boolean metrics = false;
 }
 
 
diff --git a/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java b/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java
@@ -41,7 +41,9 @@
 import software.amazon.awssdk.services.kinesis.model.StreamDescriptionSummary;
 import software.amazon.kinesis.common.InitialPositionInStream;
 import software.amazon.kinesis.coordinator.Scheduler;
+import software.amazon.kinesis.metrics.CloudWatchMetricsFactory;
 import software.amazon.kinesis.metrics.MetricsLevel;
+import software.amazon.kinesis.metrics.NullMetricsFactory;
 import software.amazon.kinesis.retrieval.polling.PollingConfig;
 
 import java.time.Duration;
@@ -58,6 +60,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -257,6 +260,30 @@ void testCreateScheduler() {
         verify(workerIdentifierGenerator, times(1)).generate();
     }
 
+    @Test
+    void testCreateSchedulerUsesNullMetricsFactoryWhenMetricsDisabled() {
+        when(kinesisSourceConfig.isMetrics()).thenReturn(false);
+        KinesisService kinesisService = new KinesisService(kinesisSourceConfig, kinesisClientFactory, pluginMetrics, pluginFactory,
+                pipelineDescription, acknowledgementSetManager, kinesisLeaseConfigSupplier, workerIdentifierGenerator);
+        Scheduler schedulerObjectUnderTest = kinesisService.createScheduler(buffer);
+
+        assertNotNull(schedulerObjectUnderTest);
+        assertNotNull(schedulerObjectUnderTest.metricsConfig());
+        assertInstanceOf(NullMetricsFactory.class, schedulerObjectUnderTest.metricsConfig().metricsFactory());
+    }
+
+    @Test
+    void testCreateSchedulerUsesCloudWatchMetricsFactoryWhenMetricsEnabled() {
+        when(kinesisSourceConfig.isMetrics()).thenReturn(true);
+        KinesisService kinesisService = new KinesisService(kinesisSourceConfig, kinesisClientFactory, pluginMetrics, pluginFactory,
+                pipelineDescription, acknowledgementSetManager, kinesisLeaseConfigSupplier, workerIdentifierGenerator);
+        Scheduler schedulerObjectUnderTest = kinesisService.createScheduler(buffer);
+
+        assertNotNull(schedulerObjectUnderTest);
+        assertNotNull(schedulerObjectUnderTest.metricsConfig());
+        assertInstanceOf(CloudWatchMetricsFactory.class, schedulerObjectUnderTest.metricsConfig().metricsFactory());
+    }
+
     @Test
     void testCreateSchedulerWithPollingStrategy() {
         when(kinesisSourceConfig.getConsumerStrategy()).thenReturn(ConsumerStrategy.POLLING);
diff --git a/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/configuration/KinesisSourceConfigTest.java b/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/configuration/KinesisSourceConfigTest.java
@@ -46,6 +46,7 @@ public class KinesisSourceConfigTest {
     private static final String PIPELINE_CONFIG_CHECKPOINT_ENABLED = "pipeline_with_checkpoint_enabled.yaml";
     private static final String PIPELINE_CONFIG_STREAM_ARN_ENABLED = "pipeline_with_stream_arn_config.yaml";
     private static final String PIPELINE_CONFIG_STREAM_ARN_CONSUMER_ARN_ENABLED = "pipeline_with_stream_arn_consumer_arn_config.yaml";
+    private static final String PIPELINE_CONFIG_WITH_METRICS_ENABLED = "pipeline_with_metrics_enabled.yaml";
     private static final Duration MINIMAL_CHECKPOINT_INTERVAL = Duration.ofMillis(2 * 60 * 1000); // 2 minute
 
     KinesisSourceConfig kinesisSourceConfig;
@@ -81,6 +82,7 @@ void testSourceConfig() {
         assertEquals(KinesisSourceConfig.DEFAULT_MAX_INITIALIZATION_ATTEMPTS, kinesisSourceConfig.getMaxInitializationAttempts());
         assertEquals(KinesisSourceConfig.DEFAULT_INITIALIZATION_BACKOFF_TIME, kinesisSourceConfig.getInitializationBackoffTime());
         assertTrue(kinesisSourceConfig.isAcknowledgments());
+        assertFalse(kinesisSourceConfig.isMetrics());
         assertEquals(KinesisSourceConfig.DEFAULT_SHARD_ACKNOWLEDGEMENT_TIMEOUT, kinesisSourceConfig.getShardAcknowledgmentTimeout());
         assertThat(kinesisSourceConfig.getAwsAuthenticationConfig(), notNullValue());
         assertEquals(kinesisSourceConfig.getAwsAuthenticationConfig().getAwsRegion(), Region.US_EAST_1);
@@ -114,6 +116,7 @@ void testSourceConfigWithStreamCodec() {
         assertEquals(KinesisSourceConfig.DEFAULT_MAX_INITIALIZATION_ATTEMPTS, kinesisSourceConfig.getMaxInitializationAttempts());
         assertEquals(KinesisSourceConfig.DEFAULT_INITIALIZATION_BACKOFF_TIME, kinesisSourceConfig.getInitializationBackoffTime());
         assertFalse(kinesisSourceConfig.isAcknowledgments());
+        assertFalse(kinesisSourceConfig.isMetrics());
         assertEquals(KinesisSourceConfig.DEFAULT_SHARD_ACKNOWLEDGEMENT_TIMEOUT, kinesisSourceConfig.getShardAcknowledgmentTimeout());
         assertThat(kinesisSourceConfig.getAwsAuthenticationConfig(), notNullValue());
         assertEquals(kinesisSourceConfig.getAwsAuthenticationConfig().getAwsRegion(), Region.US_EAST_1);
@@ -234,4 +237,26 @@ void testSourceConfigWithStreamArnConsumerArn() {
             assertEquals(kinesisStreamConfig.getCheckPointInterval(), MINIMAL_CHECKPOINT_INTERVAL);
         }
     }
+
+    @Test
+    @Tag(PIPELINE_CONFIG_WITH_METRICS_ENABLED)
+    void testSourceConfigWithMetricsEnabled() {
+
+        assertThat(kinesisSourceConfig, notNullValue());
+        assertTrue(kinesisSourceConfig.isMetrics());
+        assertEquals(KinesisSourceConfig.DEFAULT_NUMBER_OF_RECORDS_TO_ACCUMULATE, kinesisSourceConfig.getNumberOfRecordsToAccumulate());
+        assertEquals(KinesisSourceConfig.DEFAULT_TIME_OUT_IN_MILLIS, kinesisSourceConfig.getBufferTimeout());
+        assertFalse(kinesisSourceConfig.isAcknowledgments());
+        assertThat(kinesisSourceConfig.getAwsAuthenticationConfig(), notNullValue());
+        assertEquals(kinesisSourceConfig.getAwsAuthenticationConfig().getAwsRegion(), Region.US_EAST_1);
+        assertEquals(kinesisSourceConfig.getAwsAuthenticationConfig().getAwsStsRoleArn(), "arn:aws:iam::123456789012:role/OSI-PipelineRole");
+
+        List<KinesisStreamConfig> streamConfigs = kinesisSourceConfig.getStreams();
+        assertNotNull(kinesisSourceConfig.getCodec());
+        assertEquals(streamConfigs.size(), 3);
+
+        for (KinesisStreamConfig kinesisStreamConfig: streamConfigs) {
+            assertTrue(kinesisStreamConfig.getName().contains("stream"));
+        }
+    }
 }
diff --git a/data-prepper-plugins/kinesis-source/src/test/resources/pipeline_with_metrics_enabled.yaml b/data-prepper-plugins/kinesis-source/src/test/resources/pipeline_with_metrics_enabled.yaml
@@ -0,0 +1,13 @@
+source:
+  kinesis:
+    streams:
+      - stream_name: "stream-1"
+      - stream_name: "stream-2"
+      - stream_name: "stream-3"
+    codec:
+      ndjson:
+    aws:
+      sts_role_arn: "arn:aws:iam::123456789012:role/OSI-PipelineRole"
+      region: "us-east-1"
+    metrics: true
+

Original file line number	Diff line number	Diff line change
`@@ -79,6 +79,10 @@ public Duration getShardAcknowledgmentTimeout() {`
`79`	`79`	`@Getter`
`80`	`80`	`@JsonProperty("initialization_backoff_time")`
`81`	`81`	`private Duration initializationBackoffTime = DEFAULT_INITIALIZATION_BACKOFF_TIME;`
	`82`	`+`
	`83`	`+ @Getter`
	`84`	`+ @JsonProperty("metrics")`
	`85`	`+ private boolean metrics = false;`
`82`	`86`	`}`
`83`	`87`
`84`	`88`