Added support for faster shard reallocation on graceful shutdown

.gitignore

@@ -2,5 +2,7 @@ target/
 AwsCredentials.properties
 .idea
 *.iml
+.sdkmanrc
+.vscode
 *.swp
 .DS_Store

amazon-kinesis-client/src/main/java/software/amazon/kinesis/common/CommonCalculations.java

@@ -28,4 +28,22 @@ public class CommonCalculations {
     public static long getRenewerTakerIntervalMillis(long leaseDurationMillis, long epsilonMillis) {
         return leaseDurationMillis / 3 - epsilonMillis;
     }
+
+    /**
+     * Convenience method for calculating lease taker intervals in milliseconds.
+     *
+     * @param leaseTakerIntervalMillis Current value for interval (from default or overriden).
+     * @param leaseDurationMillis Duration of a lease
+     * @param epsilonMillis Allow for some variance when calculating lease expirations
+     * @return lease taker interval.
+     */
+    public static long getLeaseTakerIntervalMillis(
+      long leaseTakerIntervalMillis, long leaseDurationMillis, long epsilonMillis
+    ) {
+      if (leaseTakerIntervalMillis > 0) {
+        return leaseTakerIntervalMillis;
+      }
+
+      return (leaseDurationMillis + epsilonMillis) * 2;
+    }
 }

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

@@ -953,7 +953,8 @@ protected ShardConsumer buildConsumer(@NonNull final ShardInfo shardInfo,
                 hierarchicalShardSyncerProvider.apply(streamConfig),
                 metricsFactory,
                 leaseCleanupManager,
-                schemaRegistryDecoder
+                schemaRegistryDecoder,
+                leaseManagementConfig.evictLeaseOnShutdown()
             );
         return new ShardConsumer(cache, executorService, shardInfo, lifecycleConfig.logWarningForTaskAfterMillis(),
                 argument, lifecycleConfig.taskExecutionListener(), lifecycleConfig.readTimeoutsToIgnoreBeforeWarning());

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

@@ -44,7 +44,7 @@
 public class Lease {
     /**
      * See javadoc for System.nanoTime - summary:
-     * 
+     *
      * Sometimes System.nanoTime's return values will wrap due to overflow. When they do, the difference between two
      * values will be very large. We will consider leases to be expired if they are more than a year old.
      */
@@ -106,7 +106,7 @@ public class Lease {
 
     /**
      * Copy constructor, used by clone().
-     * 
+     *
      * @param lease lease to copy
      */
     protected Lease(Lease lease) {
@@ -159,7 +159,7 @@ public Set<String> parentShardIds() {
     /**
      * Updates this Lease's mutable, application-specific fields based on the passed-in lease object. Does not update
      * fields that are internal to the leasing library (leaseKey, leaseOwner, leaseCounter).
-     * 
+     *
      * @param lease
      */
     public void update(final Lease lease) {
@@ -190,9 +190,16 @@ public boolean isExpired(long leaseDurationNanos, long asOfNanos) {
         }
     }
 
+    /**
+     * @return true if this lease is unassigned (no assigned owner), false otherwise.
+     */
+    public boolean isUnassigned() {
+        return leaseOwner == null;
+    }
+
     /**
      * Sets lastCounterIncrementNanos
-     * 
+     *
      * @param lastCounterIncrementNanos last renewal in nanoseconds since the epoch
      */
     public void lastCounterIncrementNanos(Long lastCounterIncrementNanos) {
@@ -201,7 +208,7 @@ public void lastCounterIncrementNanos(Long lastCounterIncrementNanos) {
 
     /**
      * Sets concurrencyToken.
-     * 
+     *
      * @param concurrencyToken may not be null
      */
     public void concurrencyToken(@NonNull final UUID concurrencyToken) {
@@ -210,7 +217,7 @@ public void concurrencyToken(@NonNull final UUID concurrencyToken) {
 
     /**
      * Sets leaseKey. LeaseKey is immutable once set.
-     * 
+     *
      * @param leaseKey may not be null.
      */
     public void leaseKey(@NonNull final String leaseKey) {
@@ -222,7 +229,7 @@ public void leaseKey(@NonNull final String leaseKey) {
 
     /**
      * Sets leaseCounter.
-     * 
+     *
      * @param leaseCounter may not be null
      */
     public void leaseCounter(@NonNull final Long leaseCounter) {
@@ -298,7 +305,7 @@ public void hashKeyRange(HashKeyRangeForLease hashKeyRangeForLease) {
 
     /**
      * Sets leaseOwner.
-     * 
+     *
      * @param leaseOwner may be null.
      */
     public void leaseOwner(String leaseOwner) {
@@ -307,11 +314,10 @@ public void leaseOwner(String leaseOwner) {
 
     /**
      * Returns a deep copy of this object. Type-unsafe - there aren't good mechanisms for copy-constructing generics.
-     * 
+     *
      * @return A deep copy of this object.
      */
     public Lease copy() {
         return new Lease(this);
     }
-
 }

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

@@ -227,6 +227,18 @@ public class LeaseManagementConfig {
     private long listShardsCacheAllowedAgeInSeconds = 30;
     private int cacheMissWarningModulus = 250;
 
+    /**
+     * Interval at which the lease taker will execute.
+     * If unspecified, an interval will be calculated based on the lease duration.
+     */
+    private long leaseTakerIntervalMillis = -1L;
+
+    /**
+     * If leases should be evicted or not on shutdown requested.
+     * By default, leases are not evicted.
+     */
+    private boolean evictLeaseOnShutdown = false;
+
     private MetricsFactory metricsFactory = new NullMetricsFactory();
 
     @Deprecated
@@ -329,6 +341,7 @@ public LeaseManagementFactory leaseManagementFactory() {
                     initialPositionInStream(),
                     failoverTimeMillis(),
                     epsilonMillis(),
+                    leaseTakerIntervalMillis,
                     maxLeasesForWorker(),
                     maxLeasesToStealAtOneTime(),
                     maxLeaseRenewalThreads(),
@@ -364,6 +377,7 @@ public LeaseManagementFactory leaseManagementFactory(final LeaseSerializer lease
                     executorService(),
                     failoverTimeMillis(),
                     epsilonMillis(),
+                    leaseTakerIntervalMillis,
                     maxLeasesForWorker(),
                     maxLeasesToStealAtOneTime(),
                     maxLeaseRenewalThreads(),

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

@@ -49,6 +49,7 @@
 import software.amazon.kinesis.metrics.MetricsUtil;
 
 import static software.amazon.kinesis.common.CommonCalculations.getRenewerTakerIntervalMillis;
+import static software.amazon.kinesis.common.CommonCalculations.getLeaseTakerIntervalMillis;
 
 /**
  * LeaseCoordinator abstracts away LeaseTaker and LeaseRenewer from the application code that's using leasing. It owns
@@ -108,12 +109,13 @@ public DynamoDBLeaseCoordinator(final LeaseRefresher leaseRefresher,
                                     final String workerIdentifier,
                                     final long leaseDurationMillis,
                                     final long epsilonMillis,
+                                    final long leaseTakerIntervalMillis,
                                     final int maxLeasesForWorker,
                                     final int maxLeasesToStealAtOneTime,
                                     final int maxLeaseRenewerThreadCount,
                                     final MetricsFactory metricsFactory) {
-        this(leaseRefresher, workerIdentifier, leaseDurationMillis, epsilonMillis, maxLeasesForWorker,
-                maxLeasesToStealAtOneTime, maxLeaseRenewerThreadCount,
+        this(leaseRefresher, workerIdentifier, leaseDurationMillis, epsilonMillis, leaseTakerIntervalMillis,
+                maxLeasesForWorker, maxLeasesToStealAtOneTime, maxLeaseRenewerThreadCount,
                 TableConstants.DEFAULT_INITIAL_LEASE_TABLE_READ_CAPACITY,
                 TableConstants.DEFAULT_INITIAL_LEASE_TABLE_WRITE_CAPACITY, metricsFactory);
     }
@@ -144,6 +146,7 @@ public DynamoDBLeaseCoordinator(final LeaseRefresher leaseRefresher,
                                     final String workerIdentifier,
                                     final long leaseDurationMillis,
                                     final long epsilonMillis,
+                                    final long leaseTakerIntervalMillis,
                                     final int maxLeasesForWorker,
                                     final int maxLeasesToStealAtOneTime,
                                     final int maxLeaseRenewerThreadCount,
@@ -158,7 +161,7 @@ public DynamoDBLeaseCoordinator(final LeaseRefresher leaseRefresher,
         this.leaseRenewer = new DynamoDBLeaseRenewer(
                 leaseRefresher, workerIdentifier, leaseDurationMillis, leaseRenewalThreadpool, metricsFactory);
         this.renewerIntervalMillis = getRenewerTakerIntervalMillis(leaseDurationMillis, epsilonMillis);
-        this.takerIntervalMillis = (leaseDurationMillis + epsilonMillis) * 2;
+        this.takerIntervalMillis = getLeaseTakerIntervalMillis(leaseTakerIntervalMillis, leaseDurationMillis, epsilonMillis);
         if (initialLeaseTableReadCapacity <= 0) {
             throw new IllegalArgumentException("readCapacity should be >= 1");
         }