Run multiple instance of scheduler on one JVM

amazon-kinesis-client/src/main/java/software/amazon/kinesis/coordinator/Scheduler.java

@@ -47,6 +47,7 @@
 import software.amazon.kinesis.leases.ShardPrioritization;
 import software.amazon.kinesis.leases.ShardSyncTask;
 import software.amazon.kinesis.leases.ShardSyncTaskManager;
+import software.amazon.kinesis.leases.HierarchicalShardSyncer;
 import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseCoordinator;
 import software.amazon.kinesis.leases.exceptions.LeasingException;
 import software.amazon.kinesis.lifecycle.LifecycleConfig;
@@ -113,6 +114,7 @@ public class Scheduler implements Runnable {
     private final ShardDetector shardDetector;
     private final boolean ignoreUnexpetedChildShards;
     private final AggregatorUtil aggregatorUtil;
+    private final HierarchicalShardSyncer hierarchicalShardSyncer;
 
     // Holds consumers for shards the worker is currently tracking. Key is shard
     // info, value is ShardConsumer.
@@ -193,6 +195,7 @@ public Scheduler(@NonNull final CheckpointConfig checkpointConfig,
         this.shardDetector = this.shardSyncTaskManager.shardDetector();
         this.ignoreUnexpetedChildShards = this.leaseManagementConfig.ignoreUnexpectedChildShards();
         this.aggregatorUtil = this.lifecycleConfig.aggregatorUtil();
+        this.hierarchicalShardSyncer = leaseManagementConfig.hierarchicalShardSyncer();
     }
 
     /**
@@ -236,7 +239,7 @@ private void initialize() {
                 if (!skipShardSyncAtWorkerInitializationIfLeasesExist || leaseRefresher.isLeaseTableEmpty()) {
                     log.info("Syncing Kinesis shard info");
                     ShardSyncTask shardSyncTask = new ShardSyncTask(shardDetector, leaseRefresher, initialPosition,
-                            cleanupLeasesUponShardCompletion, ignoreUnexpetedChildShards, 0L, metricsFactory);
+                            cleanupLeasesUponShardCompletion, ignoreUnexpetedChildShards, 0L, hierarchicalShardSyncer, metricsFactory);
                     result = new MetricsCollectingTaskDecorator(shardSyncTask, metricsFactory).call();
                 } else {
                     log.info("Skipping shard sync per configuration setting (and lease table is not empty)");
@@ -569,8 +572,9 @@ protected ShardConsumer buildConsumer(@NonNull final ShardInfo shardInfo,
                 cleanupLeasesUponShardCompletion,
                 ignoreUnexpetedChildShards,
                 shardDetector,
-                metricsFactory,
-                aggregatorUtil);
+                aggregatorUtil,
+                hierarchicalShardSyncer,
+                metricsFactory);
         return new ShardConsumer(cache, executorService, shardInfo, lifecycleConfig.logWarningForTaskAfterMillis(), argument);
     }
 

amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/ShardSyncer.java → amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/HierarchicalShardSyncer.java

@@ -30,13 +30,12 @@
 
 import org.apache.commons.lang3.StringUtils;
 
-import lombok.AccessLevel;
-import lombok.NoArgsConstructor;
 import lombok.NonNull;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import software.amazon.awssdk.services.kinesis.model.Shard;
 import software.amazon.awssdk.utils.CollectionUtils;
+import software.amazon.kinesis.annotations.KinesisClientInternalApi;
 import software.amazon.kinesis.common.InitialPositionInStream;
 import software.amazon.kinesis.common.InitialPositionInStreamExtended;
 import software.amazon.kinesis.exceptions.internal.KinesisClientLibIOException;
@@ -54,24 +53,27 @@
  * It deletes leases for shards that have been trimmed from Kinesis, or if we've completed processing it
  * and begun processing it's child shards.
  */
-@NoArgsConstructor(access = AccessLevel.PRIVATE)
 @Slf4j
-public class ShardSyncer {
+@KinesisClientInternalApi
+public class HierarchicalShardSyncer {
+
     /**
      * Check and create leases for any new shards (e.g. following a reshard operation). Sync leases with Kinesis shards
      * (e.g. at startup, or when we reach end of a shard).
-     * 
+     *
+     * @param shardDetector
      * @param leaseRefresher
      * @param initialPosition
      * @param cleanupLeasesOfCompletedShards
      * @param ignoreUnexpectedChildShards
+     * @param scope
      * @throws DependencyException
      * @throws InvalidStateException
      * @throws ProvisionedThroughputException
      * @throws KinesisClientLibIOException
      */
     // CHECKSTYLE:OFF CyclomaticComplexity
-    public static synchronized void checkAndCreateLeasesForNewShards(@NonNull final ShardDetector shardDetector,
+    public synchronized void checkAndCreateLeaseForNewShards(@NonNull final ShardDetector shardDetector,
             final LeaseRefresher leaseRefresher, final InitialPositionInStreamExtended initialPosition,
             final boolean cleanupLeasesOfCompletedShards, final boolean ignoreUnexpectedChildShards,
             final MetricsScope scope) throws DependencyException, InvalidStateException,
@@ -152,7 +154,7 @@ private static Set<String> findInconsistentShardIds(final Map<String, Set<String
      * @return ShardIds of child shards (children of the expectedClosedShard)
      * @throws KinesisClientLibIOException
      */
-    static synchronized void assertClosedShardsAreCoveredOrAbsent(final Map<String, Shard> shardIdToShardMap,
+    synchronized void assertClosedShardsAreCoveredOrAbsent(final Map<String, Shard> shardIdToShardMap,
             final Map<String, Set<String>> shardIdToChildShardIdsMap, final Set<String> shardIdsOfClosedShards)
             throws KinesisClientLibIOException {
         final String exceptionMessageSuffix = "This can happen if we constructed the list of shards "
@@ -181,7 +183,7 @@ static synchronized void assertClosedShardsAreCoveredOrAbsent(final Map<String,
         }
     }
 
-    private static synchronized void assertHashRangeOfClosedShardIsCovered(final Shard closedShard,
+    private synchronized void assertHashRangeOfClosedShardIsCovered(final Shard closedShard,
             final Map<String, Shard> shardIdToShardMap, final Set<String> childShardIds)
             throws KinesisClientLibIOException {
         BigInteger minStartingHashKeyOfChildren = null;
@@ -583,7 +585,7 @@ static boolean isCandidateForCleanup(final Lease lease, final Set<String> curren
      * @throws ProvisionedThroughputException
      * @throws KinesisClientLibIOException
      */
-    private static synchronized void cleanupLeasesOfFinishedShards(final Collection<Lease> currentLeases,
+    private synchronized void cleanupLeasesOfFinishedShards(final Collection<Lease> currentLeases,
             final Map<String, Shard> shardIdToShardMap, final Map<String, Set<String>> shardIdToChildShardIdsMap,
             final List<Lease> trackedLeases, final LeaseRefresher leaseRefresher) throws DependencyException,
             InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
@@ -625,7 +627,7 @@ private static synchronized void cleanupLeasesOfFinishedShards(final Collection<
      * @throws InvalidStateException 
      * @throws DependencyException 
      */
-    static synchronized void cleanupLeaseForClosedShard(final String closedShardId, final Set<String> childShardIds,
+    synchronized void cleanupLeaseForClosedShard(final String closedShardId, final Set<String> childShardIds,
             final Map<String, Lease> trackedLeases, final LeaseRefresher leaseRefresher)
             throws DependencyException, InvalidStateException, ProvisionedThroughputException {
         final Lease leaseForClosedShard = trackedLeases.get(closedShardId);

amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/LeaseManagementConfig.java

@@ -220,6 +220,8 @@ static class LeaseManagementThreadPool extends ThreadPoolExecutor {
         }
     };
 
+    private HierarchicalShardSyncer hierarchicalShardSyncer = new HierarchicalShardSyncer();
+
     private LeaseManagementFactory leaseManagementFactory;
 
     public LeaseManagementFactory leaseManagementFactory() {
@@ -246,7 +248,8 @@ public LeaseManagementFactory leaseManagementFactory() {
                     listShardsCacheAllowedAgeInSeconds(),
                     cacheMissWarningModulus(),
                     initialLeaseTableReadCapacity(),
-                    initialLeaseTableWriteCapacity());
+                    initialLeaseTableWriteCapacity(),
+                    hierarchicalShardSyncer());
         }
         return leaseManagementFactory;
     }

amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/ShardSyncTask.java

@@ -48,6 +48,8 @@ public class ShardSyncTask implements ConsumerTask {
     private final boolean ignoreUnexpectedChildShards;
     private final long shardSyncTaskIdleTimeMillis;
     @NonNull
+    private final HierarchicalShardSyncer hierarchicalShardSyncer;
+    @NonNull
     private final MetricsFactory metricsFactory;
 
     private final TaskType taskType = TaskType.SHARDSYNC;
@@ -62,7 +64,7 @@ public TaskResult call() {
         final MetricsScope scope = MetricsUtil.createMetricsWithOperation(metricsFactory, SHARD_SYNC_TASK_OPERATION);
 
         try {
-            ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, leaseRefresher, initialPosition,
+            hierarchicalShardSyncer.checkAndCreateLeaseForNewShards(shardDetector, leaseRefresher, initialPosition,
                     cleanupLeasesUponShardCompletion, ignoreUnexpectedChildShards, scope);
             if (shardSyncTaskIdleTimeMillis > 0) {
                 Thread.sleep(shardSyncTaskIdleTimeMillis);

amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/ShardSyncTaskManager.java

@@ -50,8 +50,68 @@ public class ShardSyncTaskManager {
     @NonNull
     private final ExecutorService executorService;
     @NonNull
+    private final HierarchicalShardSyncer hierarchicalShardSyncer;
+    @NonNull
     private final MetricsFactory metricsFactory;
 
+    /**
+     * Constructor.
+     *
+     * <p>NOTE: This constructor is deprecated and will be removed in a future release.</p>
+     *
+     * @param shardDetector
+     * @param leaseRefresher
+     * @param initialPositionInStream
+     * @param cleanupLeasesUponShardCompletion
+     * @param ignoreUnexpectedChildShards
+     * @param shardSyncIdleTimeMillis
+     * @param executorService
+     * @param metricsFactory
+     */
+    @Deprecated
+    public ShardSyncTaskManager(ShardDetector shardDetector, LeaseRefresher leaseRefresher,
+            InitialPositionInStreamExtended initialPositionInStream, boolean cleanupLeasesUponShardCompletion,
+            boolean ignoreUnexpectedChildShards, long shardSyncIdleTimeMillis, ExecutorService executorService,
+            MetricsFactory metricsFactory) {
+        this.shardDetector = shardDetector;
+        this.leaseRefresher = leaseRefresher;
+        this.initialPositionInStream = initialPositionInStream;
+        this.cleanupLeasesUponShardCompletion = cleanupLeasesUponShardCompletion;
+        this.ignoreUnexpectedChildShards = ignoreUnexpectedChildShards;
+        this.shardSyncIdleTimeMillis = shardSyncIdleTimeMillis;
+        this.executorService = executorService;
+        this.hierarchicalShardSyncer = new HierarchicalShardSyncer();
+        this.metricsFactory = metricsFactory;
+    }
+
+    /**
+     * Constructor.
+     *
+     * @param shardDetector
+     * @param leaseRefresher
+     * @param initialPositionInStream
+     * @param cleanupLeasesUponShardCompletion
+     * @param ignoreUnexpectedChildShards
+     * @param shardSyncIdleTimeMillis
+     * @param executorService
+     * @param hierarchicalShardSyncer
+     * @param metricsFactory
+     */
+    public ShardSyncTaskManager(ShardDetector shardDetector, LeaseRefresher leaseRefresher,
+            InitialPositionInStreamExtended initialPositionInStream, boolean cleanupLeasesUponShardCompletion,
+            boolean ignoreUnexpectedChildShards, long shardSyncIdleTimeMillis, ExecutorService executorService,
+            HierarchicalShardSyncer hierarchicalShardSyncer, MetricsFactory metricsFactory) {
+        this.shardDetector = shardDetector;
+        this.leaseRefresher = leaseRefresher;
+        this.initialPositionInStream = initialPositionInStream;
+        this.cleanupLeasesUponShardCompletion = cleanupLeasesUponShardCompletion;
+        this.ignoreUnexpectedChildShards = ignoreUnexpectedChildShards;
+        this.shardSyncIdleTimeMillis = shardSyncIdleTimeMillis;
+        this.executorService = executorService;
+        this.hierarchicalShardSyncer = hierarchicalShardSyncer;
+        this.metricsFactory = metricsFactory;
+    }
+
     private ConsumerTask currentTask;
     private Future<TaskResult> future;
 
@@ -82,6 +142,7 @@ private synchronized boolean checkAndSubmitNextTask() {
                                     cleanupLeasesUponShardCompletion,
                                     ignoreUnexpectedChildShards,
                                     shardSyncIdleTimeMillis,
+                                    hierarchicalShardSyncer,
                                     metricsFactory),
                             metricsFactory);
             future = executorService.submit(currentTask);

amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseManagementFactory.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
 import software.amazon.kinesis.annotations.KinesisClientInternalApi;
 import software.amazon.kinesis.common.InitialPositionInStreamExtended;
+import software.amazon.kinesis.leases.HierarchicalShardSyncer;
 import software.amazon.kinesis.leases.KinesisShardDetector;
 import software.amazon.kinesis.leases.LeaseCoordinator;
 import software.amazon.kinesis.leases.LeaseManagementFactory;
@@ -50,6 +51,9 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
     private final ExecutorService executorService;
     @NonNull
     private final InitialPositionInStreamExtended initialPositionInStream;
+    @NonNull
+    private final HierarchicalShardSyncer hierarchicalShardSyncer;
+
     private final long failoverTimeMillis;
     private final long epsilonMillis;
     private final int maxLeasesForWorker;
@@ -116,6 +120,8 @@ public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, fi
     /**
      * Constructor.
      *
+     * <p>NOTE: This constructor is deprecated and will be removed in a future release.</p>
+     *
      * @param kinesisClient
      * @param streamName
      * @param dynamoDBClient
@@ -140,6 +146,7 @@ public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, fi
      * @param initialLeaseTableReadCapacity
      * @param initialLeaseTableWriteCapacity
      */
+    @Deprecated
     public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, final String streamName,
             final DynamoDbAsyncClient dynamoDBClient, final String tableName, final String workerIdentifier,
             final ExecutorService executorService, final InitialPositionInStreamExtended initialPositionInStream,
@@ -150,6 +157,54 @@ public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, fi
             final int maxListShardsRetryAttempts, final int maxCacheMissesBeforeReload,
             final long listShardsCacheAllowedAgeInSeconds, final int cacheMissWarningModulus,
             final long initialLeaseTableReadCapacity, final long initialLeaseTableWriteCapacity) {
+        this(kinesisClient, streamName, dynamoDBClient, tableName, workerIdentifier, executorService,
+                initialPositionInStream, failoverTimeMillis, epsilonMillis, maxLeasesForWorker,
+                maxLeasesToStealAtOneTime, maxLeaseRenewalThreads, cleanupLeasesUponShardCompletion,
+                ignoreUnexpectedChildShards, shardSyncIntervalMillis, consistentReads, listShardsBackoffTimeMillis,
+                maxListShardsRetryAttempts, maxCacheMissesBeforeReload, listShardsCacheAllowedAgeInSeconds,
+                cacheMissWarningModulus, initialLeaseTableReadCapacity, initialLeaseTableWriteCapacity,
+                new HierarchicalShardSyncer());
+    }
+
+    /**
+     * Constructor.
+     *
+     * @param kinesisClient
+     * @param streamName
+     * @param dynamoDBClient
+     * @param tableName
+     * @param workerIdentifier
+     * @param executorService
+     * @param initialPositionInStream
+     * @param failoverTimeMillis
+     * @param epsilonMillis
+     * @param maxLeasesForWorker
+     * @param maxLeasesToStealAtOneTime
+     * @param maxLeaseRenewalThreads
+     * @param cleanupLeasesUponShardCompletion
+     * @param ignoreUnexpectedChildShards
+     * @param shardSyncIntervalMillis
+     * @param consistentReads
+     * @param listShardsBackoffTimeMillis
+     * @param maxListShardsRetryAttempts
+     * @param maxCacheMissesBeforeReload
+     * @param listShardsCacheAllowedAgeInSeconds
+     * @param cacheMissWarningModulus
+     * @param initialLeaseTableReadCapacity
+     * @param initialLeaseTableWriteCapacity
+     * @param hierarchicalShardSyncer
+     */
+    public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, final String streamName,
+            final DynamoDbAsyncClient dynamoDBClient, final String tableName, final String workerIdentifier,
+            final ExecutorService executorService, final InitialPositionInStreamExtended initialPositionInStream,
+            final long failoverTimeMillis, final long epsilonMillis, final int maxLeasesForWorker,
+            final int maxLeasesToStealAtOneTime, final int maxLeaseRenewalThreads,
+            final boolean cleanupLeasesUponShardCompletion, final boolean ignoreUnexpectedChildShards,
+            final long shardSyncIntervalMillis, final boolean consistentReads, final long listShardsBackoffTimeMillis,
+            final int maxListShardsRetryAttempts, final int maxCacheMissesBeforeReload,
+            final long listShardsCacheAllowedAgeInSeconds, final int cacheMissWarningModulus,
+            final long initialLeaseTableReadCapacity, final long initialLeaseTableWriteCapacity,
+            final HierarchicalShardSyncer hierarchicalShardSyncer) {
         this.kinesisClient = kinesisClient;
         this.streamName = streamName;
         this.dynamoDBClient = dynamoDBClient;
@@ -173,6 +228,7 @@ public DynamoDBLeaseManagementFactory(final KinesisAsyncClient kinesisClient, fi
         this.cacheMissWarningModulus = cacheMissWarningModulus;
         this.initialLeaseTableReadCapacity = initialLeaseTableReadCapacity;
         this.initialLeaseTableWriteCapacity = initialLeaseTableWriteCapacity;
+        this.hierarchicalShardSyncer = hierarchicalShardSyncer;
     }
 
     @Override
@@ -198,6 +254,7 @@ public ShardSyncTaskManager createShardSyncTaskManager(@NonNull final MetricsFac
                 ignoreUnexpectedChildShards,
                 shardSyncIntervalMillis,
                 executorService,
+                hierarchicalShardSyncer,
                 metricsFactory);
     }