Ensure worker re-initializes after startup failure (#1804)

Author: ehasah-aws

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

@@ -169,6 +169,10 @@ private void initializeStartupComponents(
             final LeaseAssignmentMode assignmentMode,
             final Supplier<? extends LeaderDecider> leaderDeciderSupplier,
             final boolean createLamAndStartMetrics) {
+        if (initialized) {
+            log.info("Startup components already initialized, skipping initializeStartupComponents");
+            return;
+        }
         this.dualMode = dualMode;
         this.currentAssignmentMode = assignmentMode;
 
@@ -299,11 +303,7 @@ private void createGsi(final boolean blockingWait) throws DependencyException {
         if (blockingWait) {
             log.info("Waiting for Lease table GSI creation");
             final long secondsBetweenPolls = 10L;
-            // TODO: there exists an issue where if this timeout is reached the KCL starts up erroneously for a new
-            //   application (Scheduler appears to start but doesn't take any leases and some threads don't start
-            //   properly) where the only recovery is bouncing the worker after the GSI is created.
-            //   Extending the timeout to a long time to ensure GSI is created as a short term fix and will revisit
-            //   on the proper fix in the future
+            // Extending the timeout to a long time to ensure GSI is created.
             final long timeoutSeconds = 3600L;
             final boolean isIndexActive =
                     leaseRefresher.waitUntilLeaseOwnerToLeaseKeyIndexExists(secondsBetweenPolls, timeoutSeconds);

amazon-kinesis-client/src/main/java/software/amazon/kinesis/coordinator/migration/MigrationStateMachineImpl.java

@@ -113,16 +113,19 @@ public void initialize() throws DependencyException {
                             coordinatorStateDAO, clientVersionConfig, random, workerId);
             final SimpleEntry<ClientVersion, MigrationState> dataForInitialization =
                     startingStateInitializer.getInitialState();
-            startingClientVersion = dataForInitialization.getKey();
             startingMigrationState = dataForInitialization.getValue();
 
             // Create and enter the starting state. The state's enter() method
             // writes MigrationState to DDB (except CLIENT_VERSION_INIT which skips DDB)
             // and initializes components. If enter() throws DependencyException, Scheduler
             // retries the whole initialization loop.
             final MigrationClientVersionState startingState =
-                    createMigrationClientVersionState(startingClientVersion, startingMigrationState);
+                    createMigrationClientVersionState(dataForInitialization.getKey(), startingMigrationState);
             startingState.enter(ClientVersion.CLIENT_VERSION_INIT);
+
+            // Only set startingClientVersion after enter() succeeds — this is the guard
+            // that prevents re-initialization from being skipped on retry.
+            startingClientVersion = dataForInitialization.getKey();
             currentMigrationClientVersionState = startingState;
             log.info("MigrationStateMachine initial clientVersion {}", startingClientVersion);
 

amazon-kinesis-client/src/test/java/software/amazon/kinesis/coordinator/DynamicMigrationComponentsInitializerTest.java

@@ -409,6 +409,77 @@ public void testWorkerMetricsReporting() throws Exception {
                 77.0, statsCaptor.getValue().getMetricStats().get("CPU").get(2));
     }
 
+    @Test
+    public void testRetry_3x_gsiFailsThenSucceeds_noDoubleInit() throws Exception {
+        // First call: GSI times out. Second call: GSI succeeds.
+        when(mockLeaseRefresher.waitUntilLeaseOwnerToLeaseKeyIndexExists(anyLong(), anyLong()))
+                .thenReturn(false)
+                .thenReturn(true);
+
+        // First call should throw
+        assertThrows(
+                DependencyException.class,
+                () -> migrationInitializer.initializeClientVersionFor3x(ClientVersion.CLIENT_VERSION_INIT));
+
+        // Second call should succeed
+        migrationInitializer.initializeClientVersionFor3x(ClientVersion.CLIENT_VERSION_INIT);
+
+        // initializeStartupComponents internals should only be called once
+        verify(mockWorkerMetricsManager, Mockito.times(1)).startManager();
+        verify(mockDdbLockBasedLeaderDeciderCreator, Mockito.times(1)).get();
+        verify(mockLamCreator, Mockito.times(1)).apply(any(), any());
+
+        // GSI creation and idempotent operations should be called twice (once per attempt)
+        verify(mockLeaseRefresher, Mockito.times(2)).createLeaseOwnerToLeaseKeyIndexIfNotExists();
+        verify(mockLam).start();
+    }
+
+    @Test
+    public void testRetry_upgradeFrom2x_gsiFailsThenSucceeds_noDoubleInit() throws Exception {
+        when(mockLeaseRefresher.waitUntilLeaseOwnerToLeaseKeyIndexExists(anyLong(), anyLong()))
+                .thenReturn(false)
+                .thenReturn(true);
+
+        // UpgradeFrom2x uses createGsi(false) — non-blocking, but startWorkerMetricsReporting can throw
+        Mockito.doThrow(new DependencyException(new RuntimeException("WorkerMetrics init failed")))
+                .doNothing()
+                .when(mockWorkerMetricsDAO)
+                .initialize();
+
+        assertThrows(
+                DependencyException.class,
+                () -> migrationInitializer.initializeClientVersionForUpgradeFrom2x(ClientVersion.CLIENT_VERSION_INIT));
+
+        migrationInitializer.initializeClientVersionForUpgradeFrom2x(ClientVersion.CLIENT_VERSION_INIT);
+
+        // Non-idempotent startup components called only once
+        verify(mockWorkerMetricsManager, Mockito.times(1)).startManager();
+        verify(mockDeterministicLeaderDeciderCreator, Mockito.times(1)).get();
+        verify(mockLamCreator, Mockito.times(1)).apply(any(), any());
+    }
+
+    @Test
+    public void testRetry_3xWithRollback_metricsFailsThenSucceeds_noDoubleInit() throws Exception {
+        // startWorkerMetricsReporting throws on first call
+        Mockito.doThrow(new DependencyException(new RuntimeException("WorkerMetrics init failed")))
+                .doNothing()
+                .when(mockWorkerMetricsDAO)
+                .initialize();
+
+        assertThrows(
+                DependencyException.class,
+                () -> migrationInitializer.initializeClientVersionFor3xWithRollback(ClientVersion.CLIENT_VERSION_INIT));
+
+        migrationInitializer.initializeClientVersionFor3xWithRollback(ClientVersion.CLIENT_VERSION_INIT);
+
+        // Non-idempotent startup components called only once
+        verify(mockWorkerMetricsManager, Mockito.times(1)).startManager();
+        verify(mockDdbLockBasedLeaderDeciderCreator, Mockito.times(1)).get();
+        verify(mockAdaptiveLeaderDeciderCreator, Mockito.times(1)).get();
+        verify(mockLamCreator, Mockito.times(1)).apply(any(), any());
+        verify(mockLam).start();
+    }
+
     private abstract static class DynamoDBLockBasedLeaderDeciderSupplier
             implements Supplier<DynamoDBLockBasedLeaderDecider> {}
 

amazon-kinesis-client/src/test/java/software/amazon/kinesis/coordinator/migration/MigrationStateMachineTest.java

@@ -41,6 +41,7 @@
 import software.amazon.kinesis.metrics.NullMetricsFactory;
 import software.amazon.kinesis.worker.metricstats.WorkerMetricStatsDAO;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
@@ -347,4 +348,96 @@ private void initiateAndTestSuccessfulUpgrade(final Runnable successfulUpgradeMo
         Assertions.assertEquals(ClientVersion.CLIENT_VERSION_3X, stateMachineUnderTest.getCurrentClientVersion());
         verify(mockInitializer).initializeClientVersionFor3x(ClientVersion.CLIENT_VERSION_3X_WITH_ROLLBACK);
     }
+
+    @Test
+    public void testInitializeRetry_3x_enterFailsThenSucceeds() throws Exception {
+        // First call to initializeClientVersionFor3x throws, simulating GSI timeout
+        Mockito.doThrow(new DependencyException(new RuntimeException("GSI creation timed out")))
+                .doNothing()
+                .when(mockInitializer)
+                .initializeClientVersionFor3x(any());
+
+        final MigrationStateMachineImpl stateMachine = new MigrationStateMachineImpl(
+                nullMetricsFactory,
+                mockTimeProvider,
+                mockCoordinatorStateDAO,
+                mockMigrationStateMachineThreadPool,
+                ClientVersionConfig.CLIENT_VERSION_CONFIG_3X,
+                mockRandom,
+                mockInitializer,
+                WORKER_ID,
+                Duration.ofMinutes(0).getSeconds());
+
+        // First initialize() should throw
+        assertThrows(DependencyException.class, stateMachine::initialize);
+        // startingClientVersion should still be null — retry must not be skipped
+        Assertions.assertNull(stateMachine.getStartingClientVersion());
+
+        // Second initialize() should succeed
+        stateMachine.initialize();
+        Assertions.assertEquals(ClientVersion.CLIENT_VERSION_3X, stateMachine.getStartingClientVersion());
+
+        // initializeClientVersionFor3x should have been called twice (once per attempt)
+        verify(mockInitializer, times(2)).initializeClientVersionFor3x(any());
+    }
+
+    @Test
+    public void testInitializeRetry_upgradeFrom2x_enterFailsThenSucceeds() throws Exception {
+        // First call throws, simulating DDB failure during enter
+        Mockito.doThrow(new DependencyException(new RuntimeException("DDB connection timeout")))
+                .doNothing()
+                .when(mockInitializer)
+                .initializeClientVersionForUpgradeFrom2x(any());
+
+        final MigrationStateMachineImpl stateMachine = new MigrationStateMachineImpl(
+                nullMetricsFactory,
+                mockTimeProvider,
+                mockCoordinatorStateDAO,
+                mockMigrationStateMachineThreadPool,
+                ClientVersionConfig.CLIENT_VERSION_CONFIG_COMPATIBLE_WITH_2X,
+                mockRandom,
+                mockInitializer,
+                WORKER_ID,
+                Duration.ofMinutes(0).getSeconds());
+
+        // First initialize() should throw
+        assertThrows(DependencyException.class, stateMachine::initialize);
+        Assertions.assertNull(stateMachine.getStartingClientVersion());
+
+        // Second initialize() should succeed
+        stateMachine.initialize();
+        Assertions.assertEquals(ClientVersion.CLIENT_VERSION_UPGRADE_FROM_2X, stateMachine.getStartingClientVersion());
+
+        verify(mockInitializer, times(2)).initializeClientVersionForUpgradeFrom2x(any());
+    }
+
+    @Test
+    public void testInitializeRetry_phase1_enterFailsThenSucceeds() throws Exception {
+        // initializeClientVersionForPhase1 doesn't throw checked exceptions,
+        // so use RuntimeException to simulate failure
+        Mockito.doThrow(new RuntimeException("DDB failure"))
+                .doNothing()
+                .when(mockInitializer)
+                .initializeClientVersionForPhase1();
+
+        final MigrationStateMachineImpl stateMachine = new MigrationStateMachineImpl(
+                nullMetricsFactory,
+                mockTimeProvider,
+                mockCoordinatorStateDAO,
+                mockMigrationStateMachineThreadPool,
+                ClientVersionConfig.CLIENT_VERSION_CONFIG_COMPATIBLE_WITH_2X_PHASE1,
+                mockRandom,
+                mockInitializer,
+                WORKER_ID,
+                Duration.ofMinutes(0).getSeconds());
+
+        // First initialize() should throw (wrapped as RuntimeException)
+        Assertions.assertThrows(RuntimeException.class, stateMachine::initialize);
+        Assertions.assertNull(stateMachine.getStartingClientVersion());
+
+        // Second initialize() should succeed
+        stateMachine.initialize();
+
+        verify(mockInitializer, times(2)).initializeClientVersionForPhase1();
+    }
 }