From 089c6ab18ef9429f5e808cf4db722a999399c2a6 Mon Sep 17 00:00:00 2001
From: Joshua Kim <20001595+joshua-kim@users.noreply.github.com>
Date: Thu, 16 Jul 2020 16:27:25 -0400
Subject: [PATCH] Making isLeaseTableEmpty use one paginated scan call. (#75)

Co-authored-by: Joshua Kim <kimjos@amazon.com>
---
 .../kinesis/leases/impl/LeaseManager.java      | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java b/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
index 7fe4551a..e5860870 100644
--- a/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
+++ b/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
@@ -240,7 +240,7 @@ public class LeaseManager<T extends Lease> implements ILeaseManager<T> {
      */
     @Override
     public boolean isLeaseTableEmpty() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
-        return list(1).isEmpty();
+        return list(1, 1).isEmpty();
     }
 
     /**
@@ -253,6 +253,20 @@ public class LeaseManager<T extends Lease> implements ILeaseManager<T> {
      * @throws ProvisionedThroughputException if DynamoDB scan fail due to exceeded capacity
      */
     List<T> list(Integer limit) throws DependencyException, InvalidStateException, ProvisionedThroughputException {
+        return list(limit, Integer.MAX_VALUE);
+    }
+
+    /**
+     * List with the given page size, up to a limit of paginated calls.
+     *
+     * @param limit number of items to consider at a time - used by integration tests to force paging.
+     * @param maxPages max number of paginated scan calls.
+     * @return list of leases
+     * @throws InvalidStateException if table does not exist
+     * @throws DependencyException if DynamoDB scan fail in an unexpected way
+     * @throws ProvisionedThroughputException if DynamoDB scan fail due to exceeded capacity
+     */
+    private List<T> list(Integer limit, Integer maxPages) throws InvalidStateException, ProvisionedThroughputException, DependencyException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("Listing leases from table " + table);
         }
@@ -277,7 +291,7 @@ public class LeaseManager<T extends Lease> implements ILeaseManager<T> {
                 }
 
                 Map<String, AttributeValue> lastEvaluatedKey = scanResult.getLastEvaluatedKey();
-                if (lastEvaluatedKey == null) {
+                if (lastEvaluatedKey == null || --maxPages <= 0) {
                     // Signify that we're done.
                     scanResult = null;
                     if (LOG.isDebugEnabled()) {