From 651bf4d28cc9dea2d911f4284876b745a1d9a6d8 Mon Sep 17 00:00:00 2001
From: Aravinda Kidambi Srinivasan <akidambi@amazon.com>
Date: Mon, 29 Apr 2024 10:22:09 -0700
Subject: [PATCH] Fix a race condition between ShardConsumer shutdown and
 initialization

When Kinesis shards have no data, there can be a race condition where
the shard-end record processing from RecordProcessorThread
interleaves with Scheduler performing initialization.
This leads to ShardConsumer making incorrect state transition
during initialization (moves from PROCESSING -> SHUTTING_DOWN) state
and during shutdown handling it moves from SHUTTING_DOWN -> SHUTDOWN_COMPLETE
without running the ShutdownTask.

This can cause the ShardConsumer to not perform proper shutdown
processing that is required for a child shard processing
to be unblocked. So the child shard could be blocked forever unless the
lease for the parent shard moves to a new worker and that worker does
not run into the race condition.

This patch fixes the race condition as follows:

The intializationComplete invocation is not needed after
needsInitialization has been set to false. Because initializationComplete
is mean to perform initialization in an async manner, but once
its done, the async task is a no-op in happy-path, but it can
perform incorrect state transition during a race condition.
---
 .../amazon/kinesis/lifecycle/ShardConsumer.java        | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/lifecycle/ShardConsumer.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/lifecycle/ShardConsumer.java
index 96261131..16dac30b 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/lifecycle/ShardConsumer.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/lifecycle/ShardConsumer.java
@@ -179,8 +179,18 @@ public class ShardConsumer {
                         // Task rejection during the subscribe() call will not be propagated back as it not executed
                         // in the context of the Scheduler thread. Hence we should not assume the subscription will
                         // always be successful.
+                        // But if subscription was not successful, then it will recover
+                        // during healthCheck which will restart subscription.
+                        // From Shardconsumer point of view, initialization after the below subscribe call
+                        // is complete
                         subscribe();
                         needsInitialization = false;
+                        // Initialization is complete, return now, because we dont need to do
+                        // initializeComplete anymore. ShardConsumer is in ProcessingState now and any further activity
+                        // will be driven by publisher pushing data to subscriber which invokes handleInput
+                        // and that triggers ProcessTask. Scheduler is only meant to do health-checks
+                        // to ensure the consumer is not stuck for any reason and to do shutdown handling.
+                        return;
                     }
                 }
                 stateChangeFuture = initializeComplete();