Changing ingestion consumption check avoiding race condition. (linked…

…in#635) Previously, when OFFLINE => DROP transition happen, we will check if we can remove the partition from disk only based on partition consumption state exist or not. However, it is also possible that there are some pending state transitions in consume action queue of the StoreIngestionTask. If we drop the partition and there is a pending SUBSCRIBE action in consumption queue, then it will cause data persistence failure due to partition not existed in disk. To avoid this, we add pending queue checking for the partition about to be deleted during the partition consumption check. Co-authored-by: Hao Xu <[email protected]>
m-nagarajan · Oct 12, 2023 · 94c59de · 94c59de
1 parent e4d7a79
commit 94c59de
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/...-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java b/...-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java
@@ -1075,7 +1075,8 @@ public boolean containsRunningConsumption(String topic) {
   public boolean isPartitionConsuming(String topic, int partitionId) {
     try (AutoCloseableLock ignore = topicLockManager.getLockForResource(topic)) {
       StoreIngestionTask ingestionTask = topicNameToIngestionTaskMap.get(topic);
-      return ingestionTask != null && ingestionTask.isRunning() && ingestionTask.isPartitionConsuming(partitionId);
+      return ingestionTask != null && ingestionTask.isRunning()
+          && ingestionTask.isPartitionConsumingOrHasPendingIngestionAction(partitionId);
     }
   }
 

diff --git a/...da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java b/...da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java
@@ -1776,7 +1776,6 @@ protected void processCommonConsumerAction(
         } else {
           LOGGER.info("{} Unsubscribed to: {}", consumerTaskId, topicPartition);
         }
-
         break;
       case RESET_OFFSET:
         resetOffset(partition, topicPartition, false);
@@ -3322,8 +3321,11 @@ public void disableMetricsEmission() {
   /**
    * To check whether the given partition is still consuming message from Kafka
    */
-  public boolean isPartitionConsuming(int userPartition) {
-    return amplificationFactorAdapter.meetsAny(userPartition, partitionConsumptionStateMap::containsKey);
+  public boolean isPartitionConsumingOrHasPendingIngestionAction(int userPartition) {
+    boolean subPartitionConsumptionStateExist =
+        amplificationFactorAdapter.meetsAny(userPartition, partitionConsumptionStateMap::containsKey);
+    boolean pendingPartitionIngestionAction = hasPendingPartitionIngestionAction(userPartition);
+    return pendingPartitionIngestionAction || subPartitionConsumptionStateExist;
   }
 
   /**