From feadd5e043522456e5bd9cdc0c551b01547e26ec Mon Sep 17 00:00:00 2001
From: stair <123031771+stair-aws@users.noreply.github.com>
Date: Wed, 28 Jun 2023 10:36:32 -0400
Subject: [PATCH] Fix NPE on graceful shutdown before DDB `LeaseCoordinator`
 starts. (#1157)

---
 .../dynamodb/DynamoDBLeaseCoordinator.java    |  5 +--
 .../DynamoDBLeaseCoordinatorTest.java         | 32 ++++++++++++++-----
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinator.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinator.java
index 07e9068d..da6d8e07 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinator.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinator.java
@@ -329,8 +329,9 @@ public class DynamoDBLeaseCoordinator implements LeaseCoordinator {
 
     @Override
     public void stopLeaseTaker() {
-        takerFuture.cancel(false);
-
+        if (takerFuture != null) {
+            takerFuture.cancel(false);
+        }
     }
 
     @Override
diff --git a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinatorTest.java b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinatorTest.java
index caa7a6c7..cf1c536b 100644
--- a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinatorTest.java
+++ b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseCoordinatorTest.java
@@ -1,6 +1,5 @@
 package software.amazon.kinesis.leases.dynamodb;
 
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -12,7 +11,7 @@ import software.amazon.kinesis.metrics.MetricsFactory;
 
 import java.util.UUID;
 
-import static org.mockito.Mockito.times;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -51,17 +50,34 @@ public class DynamoDBLeaseCoordinatorTest {
 
         leaseCoordinator.initialize();
 
-        verify(leaseRefresher, times(1)).createLeaseTableIfNotExists();
-        verify(leaseRefresher, times(1)).waitUntilLeaseTableExists(SECONDS_BETWEEN_POLLS, TIMEOUT_SECONDS);
+        verify(leaseRefresher).createLeaseTableIfNotExists();
+        verify(leaseRefresher).waitUntilLeaseTableExists(SECONDS_BETWEEN_POLLS, TIMEOUT_SECONDS);
     }
 
-    @Test
+    @Test(expected = DependencyException.class)
     public void testInitialize_tableCreationFails() throws Exception {
         when(leaseRefresher.createLeaseTableIfNotExists()).thenReturn(false);
         when(leaseRefresher.waitUntilLeaseTableExists(SECONDS_BETWEEN_POLLS, TIMEOUT_SECONDS)).thenReturn(false);
 
-        Assert.assertThrows(DependencyException.class, () -> leaseCoordinator.initialize());
-        verify(leaseRefresher, times(1)).createLeaseTableIfNotExists();
-        verify(leaseRefresher, times(1)).waitUntilLeaseTableExists(SECONDS_BETWEEN_POLLS, TIMEOUT_SECONDS);
+        try {
+            leaseCoordinator.initialize();
+        } finally {
+            verify(leaseRefresher).createLeaseTableIfNotExists();
+            verify(leaseRefresher).waitUntilLeaseTableExists(SECONDS_BETWEEN_POLLS, TIMEOUT_SECONDS);
+        }
     }
+
+    /**
+     * Validates a {@link NullPointerException} is not thrown when the lease taker
+     * is stopped before it starts/exists.
+     *
+     * @see <a href="https://github.com/awslabs/amazon-kinesis-client/issues/745">issue #745</a>
+     * @see <a href="https://github.com/awslabs/amazon-kinesis-client/issues/900">issue #900</a>
+     */
+    @Test
+    public void testStopLeaseTakerBeforeStart() {
+        leaseCoordinator.stopLeaseTaker();
+        assertTrue(leaseCoordinator.getAssignments().isEmpty());
+    }
+
 }