Fixing bug for multistream not using list shards with filter. (#66)

Fixing bug for mulitstreaming not using list shards with filter.

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

@@ -125,8 +125,6 @@ public class HierarchicalShardSyncer {
             List<Shard> latestShards, final boolean ignoreUnexpectedChildShards, final MetricsScope scope, final boolean isLeaseTableEmpty)
             throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
 
-        //TODO: Need to add multistream support for this https://sim.amazon.com/issues/KinesisLTR-191
-
         if (!CollectionUtils.isNullOrEmpty(latestShards)) {
             log.debug("{} - Num shards: {}", streamIdentifier, latestShards.size());
         } else {

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

@@ -217,6 +217,19 @@ public interface LeaseRefresher {
      */
     boolean isLeaseTableEmpty() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
 
+    /**
+     * Check (synchronously) if there are any leases in the lease table for a given stream identifier.
+     *
+     * @param streamIdentifier for multi-stream mode. Can be null.
+     * @return true if there are no leases in the lease table
+     *
+     * @throws DependencyException if DynamoDB scan fails in an unexpected way
+     * @throws InvalidStateException if lease table does not exist
+     * @throws ProvisionedThroughputException if DynamoDB scan fails due to lack of capacity
+     */
+    boolean isLeaseTableEmptyForStreamIdentifier(StreamIdentifier streamIdentifier) throws DependencyException,
+            InvalidStateException, ProvisionedThroughputException;
+
     /**
      * Gets the current checkpoint of the shard. This is useful in the resharding use case
      * where we will wait for the parent shard to complete before starting on the records from a child shard.

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

@@ -69,7 +69,7 @@ public class ShardSyncTask implements ConsumerTask {
         try {
             hierarchicalShardSyncer.checkAndCreateLeaseForNewShards(shardDetector, leaseRefresher,
                     initialPosition, scope, ignoreUnexpectedChildShards,
-                    leaseRefresher.isLeaseTableEmpty());
+                    leaseRefresher.isLeaseTableEmptyForStreamIdentifier(shardDetector.streamIdentifier()));
 
             if (shardSyncTaskIdleTimeMillis > 0) {
                 Thread.sleep(shardSyncTaskIdleTimeMillis);

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

@@ -292,7 +292,13 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
     @Override
     public boolean isLeaseTableEmpty()
             throws DependencyException, InvalidStateException, ProvisionedThroughputException {
-        return list(1, 1, null).isEmpty();
+        return isLeaseTableEmptyForStreamIdentifier(null);
+    }
+
+    @Override
+    public boolean isLeaseTableEmptyForStreamIdentifier(StreamIdentifier streamIdentifier)
+            throws DependencyException, ProvisionedThroughputException, InvalidStateException {
+        return list(1, 1, streamIdentifier).isEmpty();
     }
 
     /**

amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/ExceptionThrowingLeaseRefresher.java

@@ -215,6 +215,11 @@ public class ExceptionThrowingLeaseRefresher implements LeaseRefresher {
         return false;
     }
 
+    @Override
+    public boolean isLeaseTableEmptyForStreamIdentifier(StreamIdentifier streamIdentifier) throws DependencyException, InvalidStateException, ProvisionedThroughputException {
+        return false;
+    }
+
     @Override
     public ExtendedSequenceNumber getCheckpoint(final String leaseKey)
             throws ProvisionedThroughputException, InvalidStateException, DependencyException {

amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/HierarchicalShardSyncerTest.java

@@ -335,7 +335,8 @@ public class HierarchicalShardSyncerTest {
         setupMultiStream();
         hierarchicalShardSyncer
                 .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
-                        SCOPE, false, dynamoDBLeaseRefresher.isLeaseTableEmpty());
+                        SCOPE, false, dynamoDBLeaseRefresher.isLeaseTableEmptyForStreamIdentifier(
+                                StreamIdentifier.multiStreamInstance(STREAM_IDENTIFIER)));
 
         final Set<String> expectedShardIds = new HashSet<>(
                 toMultiStreamLeaseList(Arrays.asList("shardId-4", "shardId-8", "shardId-9", "shardId-10")));
@@ -417,8 +418,8 @@ public class HierarchicalShardSyncerTest {
         setupMultiStream();
         hierarchicalShardSyncer
                 .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
-                        latestShards, false, SCOPE,
+                        latestShards, false, SCOPE, dynamoDBLeaseRefresher.isLeaseTableEmptyForStreamIdentifier(
-                        dynamoDBLeaseRefresher.isLeaseTableEmpty());
+                                StreamIdentifier.multiStreamInstance(STREAM_IDENTIFIER)));
 
         final Set<String> expectedShardIds = new HashSet<>(
                 toMultiStreamLeaseList(Arrays.asList("shardId-4", "shardId-8", "shardId-9", "shardId-10")));
@@ -682,7 +683,7 @@ public class HierarchicalShardSyncerTest {
         try {
             hierarchicalShardSyncer.checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher,
                     INITIAL_POSITION_TRIM_HORIZON, SCOPE, false,
-                    dynamoDBLeaseRefresher.isLeaseTableEmpty());
+                    dynamoDBLeaseRefresher.isLeaseTableEmptyForStreamIdentifier(StreamIdentifier.multiStreamInstance(STREAM_IDENTIFIER)));
         } finally {
             verify(shardDetector).listShards();
             verify(dynamoDBLeaseRefresher, never()).listLeases();
@@ -762,7 +763,8 @@ public class HierarchicalShardSyncerTest {
         setupMultiStream();
         hierarchicalShardSyncer
                 .checkAndCreateLeaseForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
-                        SCOPE, true, dynamoDBLeaseRefresher.isLeaseTableEmpty());
+                        SCOPE, true, dynamoDBLeaseRefresher.isLeaseTableEmptyForStreamIdentifier(
+                                StreamIdentifier.multiStreamInstance(STREAM_IDENTIFIER)));
 
         final List<Lease> leases = leaseCaptor.getAllValues();
         final Set<String> leaseKeys = leases.stream().map(Lease::leaseKey).collect(Collectors.toSet());