[da-vinci][server] Guard version-role-change resubscribe against in-flight blob transfer (linkedin#2756)

Author: jingy-li

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java

@@ -793,6 +793,19 @@ public synchronized void subscribePartition(
   void resubscribeForAllPartitions() throws InterruptedException {
     throwIfNotRunning();
     for (PartitionConsumptionState partitionConsumptionState: partitionConsumptionStateMap.values()) {
+      /**
+       * Skip partitions with an in-flight blob transfer. Resubscribing reopens RocksDB on the final
+       * partition directory while blob transfer is still receiving files into the temp directory,
+       * causing the post-transfer rename to fail with "Final partition directory is not empty".
+       * completeBlobTransferAndSubscribe will subscribe once the transfer finishes, using the
+       * already-updated versionRole/workloadType.
+       */
+      if (partitionConsumptionState.isBlobTransferInProgress()) {
+        LOGGER.info(
+            "Skipping version-role-change resubscribe for replica: {} because blob transfer is in progress.",
+            partitionConsumptionState.getReplicaId());
+        continue;
+      }
       /**
        * For completed current version replica, if we resubscribe during version role change, it will be resubscribed to
        * correct high priority pool. We will mark {@link PartitionConsumptionState#hasResubscribedAfterBootstrapAsCurrentVersion}
@@ -5907,6 +5920,20 @@ && getResubscribeRequestQueue().peek() != null) {
             getReplicaId(versionTopic, partition));
         continue;
       }
+
+      /**
+       * Skip partitions with an in-flight blob transfer. Resubscribing reopens RocksDB on the final
+       * partition directory while blob transfer is still receiving files into the temp directory,
+       * causing the post-transfer rename to fail. completeBlobTransferAndSubscribe will subscribe
+       * once the transfer finishes.
+       */
+      if (pcs.isBlobTransferInProgress()) {
+        LOGGER.info(
+            "Skipping lag-based resubscribe for replica: {} because blob transfer is in progress.",
+            pcs.getReplicaId());
+        continue;
+      }
+
       /**
        * As of now, this feature intends to resolve ingestion performance issue introduced by consumer. We will rely on
        * the error reset feature to handle the error replica properly.

clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskTest.java

@@ -270,6 +270,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.PriorityBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -4720,6 +4721,108 @@ public void testResubscribeForCompletedCurrentVersionPartition() throws Interrup
 
   }
 
+  /**
+   * Verifies that {@link StoreIngestionTask#resubscribeForAllPartitions()} skips partitions whose
+   * blob transfer is still in flight. Without this guard, a version-role change (e.g. FUTURE->CURRENT)
+   * would unsubscribe + resubscribe the partition, causing RocksDB to reopen on the final partition
+   * directory while the blob transfer was still streaming into the temp directory -- which breaks
+   * the post-transfer rename and drives the replica to Helix ERROR.
+   */
+  @Test
+  public void testResubscribeForAllPartitionsSkipsBlobTransferInProgress() throws Exception {
+    StoreIngestionTask storeIngestionTask = mock(StoreIngestionTask.class);
+    doCallRealMethod().when(storeIngestionTask).resubscribeForAllPartitions();
+
+    // pcs0: blob transfer in progress -- MUST be skipped.
+    PartitionConsumptionState pcsBlobInProgress = mock(PartitionConsumptionState.class);
+    doReturn(true).when(pcsBlobInProgress).isBlobTransferInProgress();
+    doReturn("store_v1-0").when(pcsBlobInProgress).getReplicaId();
+
+    // pcs1: no blob transfer, not complete -- should be resubscribed.
+    PartitionConsumptionState pcsNotComplete = mock(PartitionConsumptionState.class);
+    doReturn(false).when(pcsNotComplete).isBlobTransferInProgress();
+    doReturn(false).when(pcsNotComplete).isComplete();
+
+    // pcs2: no blob transfer, complete, not yet flipped -- should be resubscribed AND flipped.
+    PartitionConsumptionState pcsCompleteNotFlipped = mock(PartitionConsumptionState.class);
+    doReturn(false).when(pcsCompleteNotFlipped).isBlobTransferInProgress();
+    doReturn(true).when(pcsCompleteNotFlipped).isComplete();
+    doReturn(false).when(pcsCompleteNotFlipped).hasResubscribedAfterBootstrapAsCurrentVersion();
+
+    VeniceConcurrentHashMap<Integer, PartitionConsumptionState> pcsMap = new VeniceConcurrentHashMap<>();
+    pcsMap.put(0, pcsBlobInProgress);
+    pcsMap.put(1, pcsNotComplete);
+    pcsMap.put(2, pcsCompleteNotFlipped);
+
+    // Inject into the protected final partitionConsumptionStateMap field on the mock.
+    Field pcsMapField = StoreIngestionTask.class.getDeclaredField("partitionConsumptionStateMap");
+    pcsMapField.setAccessible(true);
+    pcsMapField.set(storeIngestionTask, pcsMap);
+
+    doReturn(true).when(storeIngestionTask).isCurrentVersion();
+
+    storeIngestionTask.resubscribeForAllPartitions();
+
+    // Blob-transfer-in-progress partition is skipped entirely.
+    verify(storeIngestionTask, never()).resubscribe(pcsBlobInProgress);
+    verify(pcsBlobInProgress, never()).setHasResubscribedAfterBootstrapAsCurrentVersion(anyBoolean());
+
+    // Other partitions are resubscribed as usual.
+    verify(storeIngestionTask, times(1)).resubscribe(pcsNotComplete);
+    verify(storeIngestionTask, times(1)).resubscribe(pcsCompleteNotFlipped);
+    verify(pcsCompleteNotFlipped, times(1)).setHasResubscribedAfterBootstrapAsCurrentVersion(eq(true));
+  }
+
+  /**
+   * Verifies that {@link StoreIngestionTask#maybeProcessResubscribeRequest()} (the lag-based
+   * auto-resubscribe path fed by {@code HeartbeatMonitoringService}) skips partitions whose blob
+   * transfer is still in flight. Same race as the version-role-change path: resubscribing would
+   * reopen RocksDB on the final partition directory mid-transfer and break the post-transfer rename.
+   */
+  @Test
+  public void testMaybeProcessResubscribeRequestSkipsBlobTransferInProgress() throws Exception {
+    StoreIngestionTask storeIngestionTask = mock(StoreIngestionTask.class);
+    doCallRealMethod().when(storeIngestionTask).maybeProcessResubscribeRequest();
+
+    VeniceServerConfig serverConfig = mock(VeniceServerConfig.class);
+    doReturn(serverConfig).when(storeIngestionTask).getServerConfig();
+    doReturn(10).when(serverConfig).getLagBasedReplicaAutoResubscribeMaxReplicaCount();
+    doReturn(60).when(serverConfig).getLagBasedReplicaAutoResubscribeIntervalInSeconds();
+
+    // pcs0: blob transfer in progress -- MUST be skipped.
+    PartitionConsumptionState pcsBlobInProgress = mock(PartitionConsumptionState.class);
+    doReturn(true).when(pcsBlobInProgress).isBlobTransferInProgress();
+    doReturn("store_v1-0").when(pcsBlobInProgress).getReplicaId();
+
+    // pcs1: no blob transfer -- should be resubscribed.
+    PartitionConsumptionState pcsReady = mock(PartitionConsumptionState.class);
+    doReturn(false).when(pcsReady).isBlobTransferInProgress();
+    doReturn(false).when(pcsReady).isErrorReported();
+    doReturn("store_v1-1").when(pcsReady).getReplicaId();
+
+    VeniceConcurrentHashMap<Integer, PartitionConsumptionState> pcsMap = new VeniceConcurrentHashMap<>();
+    pcsMap.put(0, pcsBlobInProgress);
+    pcsMap.put(1, pcsReady);
+    doReturn(pcsMap).when(storeIngestionTask).getPartitionConsumptionStateMap();
+
+    PriorityBlockingQueue<Integer> resubscribeQueue = new PriorityBlockingQueue<>();
+    resubscribeQueue.add(0);
+    resubscribeQueue.add(1);
+    doReturn(resubscribeQueue).when(storeIngestionTask).getResubscribeRequestQueue();
+
+    // Empty previous-resubscribe-time map means no rate-limiting will trigger.
+    doReturn(new VeniceConcurrentHashMap<Integer, Long>()).when(storeIngestionTask)
+        .getPartitionToPreviousResubscribeTimeMap();
+
+    storeIngestionTask.maybeProcessResubscribeRequest();
+
+    // Blob-transfer-in-progress partition is skipped entirely.
+    verify(storeIngestionTask, never()).resubscribe(pcsBlobInProgress);
+
+    // Ready partition is resubscribed as usual.
+    verify(storeIngestionTask, times(1)).resubscribe(pcsReady);
+  }
+
   @Test(dataProvider = "aaConfigProvider")
   public void testWrappedInterruptExceptionDuringGracefulShutdown(AAConfig aaConfig) throws Exception {
     hybridStoreConfig = Optional.of(