Revert to using DynamoDB lease refresher

Signed-off-by: Souvik Bose <[email protected]>
sb2k16 · Jul 18, 2024 · 0ca6cf0 · 0ca6cf0
1 parent e9ac708
commit 0ca6cf0
Showing 1 changed file with 8 additions and 11 deletions.
diff --git a/...source/src/main/java/org/opensearch/dataprepper/plugins/source/kinesis/KinesisSource.java b/...source/src/main/java/org/opensearch/dataprepper/plugins/source/kinesis/KinesisSource.java
@@ -22,6 +22,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.model.BillingMode;
 import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
+import software.amazon.kinesis.common.ConfigsBuilder;
 import software.amazon.kinesis.common.KinesisClientUtil;
 import software.amazon.kinesis.coordinator.Scheduler;
 import software.amazon.kinesis.leases.LeaseManagementConfig;
@@ -146,10 +147,12 @@ public void start(final Buffer<Record<Event>> buffer) {
         final ShardRecordProcessorFactory processorFactory = new KinesisShardRecordProcessorFactory(
                 buffer, codec, sourceConfig, acknowledgementSetManager, pluginMetrics);
 
-        MultiTenantMultiStreamConfigsBuilder configsBuilder = new MultiTenantMultiStreamConfigsBuilder(
-                new KinesisMultiStreamTracker(kinesisClient, sourceConfig, applicationName),
-                applicationName, kinesisClient, dynamoClient, cloudWatchClient, workerIdentifier, processorFactory);
-        configsBuilder.tableName("KinesisDynamoDBLeaseCoordinationTable");
+        ConfigsBuilder configsBuilder =
+                new ConfigsBuilder(
+                    new KinesisMultiStreamTracker(kinesisClient, sourceConfig, applicationName),
+                    applicationName, kinesisClient, dynamoClient, cloudWatchClient,
+                        workerIdentifier, processorFactory
+                ).tableName(applicationName);
 
         sourceConfig.getStreams().forEach(stream -> {
             if (stream.getConsumerStrategy() == KinesisStreamConfig.ConsumerStrategy.POLLING) {
@@ -160,17 +163,11 @@ public void start(final Buffer<Record<Event>> buffer) {
             }
         });
 
-        MultiTenantDynamoDBLeaseManagementFactory multiTenantDynamoDBLeaseManagementFactory =
-                (MultiTenantDynamoDBLeaseManagementFactory) configsBuilder
-                        .leaseManagementConfig()
-                        .leaseManagementFactory(new MultiTenantDynamoDBLeaseSerializer(applicationName), true);
-
         scheduler = new Scheduler(
                 configsBuilder.checkpointConfig(),
                 configsBuilder.coordinatorConfig(),
                 configsBuilder.leaseManagementConfig()
-                        .billingMode(BillingMode.PAY_PER_REQUEST)
-                        .leaseManagementFactory(multiTenantDynamoDBLeaseManagementFactory),
+                        .billingMode(BillingMode.PAY_PER_REQUEST),
                 configsBuilder.lifecycleConfig(),
                 configsBuilder.metricsConfig(),
                 configsBuilder.processorConfig(),