[da-vinci] [server] Release OFFLINE->STANDBY Latch on Version Demotion to Backup (linkedin#2586)

* release the OFFLINE->STANDBY ingestion latch when a current version is demoted to backup so waitConsumptionCompleted does not block indefinitely and the state transition is left not completed
* add regression and idempotency coverage across state-model, dispatcher, and ingestion-task tests

Author: KaiSernLim

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

@@ -1835,6 +1835,10 @@ protected void refreshIngestionContextIfChanged(Store store) throws InterruptedE
       return;
     }
 
+    if (versionRole == VersionRole.CURRENT && newVersionRole != VersionRole.CURRENT) {
+      stopTrackingCurrentVersionIngestion();
+    }
+
     LOGGER.info(
         "Trigger for version topic: {} due to Previous: version role: {}, workload type: {} "
             + "changed to New: version role: {}, workload type: {}",
@@ -1855,6 +1859,22 @@ protected void refreshIngestionContextIfChanged(Store store) throws InterruptedE
     }
   }
 
+  /**
+   * {@link AbstractPartitionStateModel#onBecomeStandbyFromOffline} only needs to synchronously wait for ingestion
+   * to be completed for current versions. If a current version becomes no longer the current version, it no longer
+   * needs to {@link AbstractPartitionStateModel#waitConsumptionCompleted}. In that case, we stop tracking ingestion
+   * for those partitions and explicitly release their latches.
+   */
+  private void stopTrackingCurrentVersionIngestion() {
+    partitionConsumptionStateMap.values()
+        .stream()
+        .filter(pcs -> pcs.isLatchCreated() && !pcs.isLatchReleased())
+        .forEach(pcs -> {
+          ingestionNotificationDispatcher.reportStopped(pcs);
+          pcs.releaseLatch();
+        });
+  }
+
   private void maybeUnsubscribeCompletedPartitions(Store store) {
     if (hybridStoreConfig.isPresent() || (!serverConfig.isUnsubscribeAfterBatchpushEnabled())) {
       return;

clients/da-vinci-client/src/test/java/com/linkedin/davinci/helix/LeaderFollowerPartitionStateModelTest.java

@@ -17,6 +17,8 @@
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 
 import com.linkedin.davinci.config.VeniceServerConfig;
@@ -40,12 +42,14 @@
 import com.linkedin.venice.meta.VersionStatus;
 import com.linkedin.venice.store.rocksdb.RocksDBUtils;
 import com.linkedin.venice.utils.PropertyBuilder;
+import com.linkedin.venice.utils.TestUtils;
 import com.linkedin.venice.utils.Utils;
 import com.linkedin.venice.utils.VeniceProperties;
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.util.Random;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import org.apache.commons.io.FileUtils;
 import org.apache.helix.NotificationContext;
 import org.apache.helix.model.Message;
@@ -84,7 +88,11 @@ public void setUp() {
     partitionPushStatusAccessorFuture = CompletableFuture.completedFuture(mock(HelixPartitionStatusAccessor.class));
     stateTransitionStats = mock(ParticipantStateTransitionStats.class);
     heartbeatMonitoringService = mock(HeartbeatMonitoringService.class);
-    leaderFollowerPartitionStateModel = new LeaderFollowerPartitionStateModel(
+    leaderFollowerPartitionStateModel = buildModel(notifier);
+  }
+
+  private LeaderFollowerPartitionStateModel buildModel(LeaderFollowerIngestionProgressNotifier notifier) {
+    return new LeaderFollowerPartitionStateModel(
         ingestionBackend,
         storeAndServerConfigs,
         partition,
@@ -359,6 +367,48 @@ public void testNonCurrentVersionSkipsGracefulDrop() {
     assertEquals(resetTimestamp, -1L, "Timestamp should be reset to -1 after transition");
   }
 
+  /**
+   * OFFLINE->STANDBY needs to complete when the current version is demoted to a backup version.
+   */
+  @Test
+  public void testOnStandbyCompletionUponDemotionToBackupVersion() throws Exception {
+    Store store = mock(Store.class);
+    when(store.getCurrentVersion()).thenReturn(storeVersion); // Version starts as current so the latch is created
+    when(store.getBootstrapToOnlineTimeoutInHours()).thenReturn(24); // large timeout for latch await
+    doReturn(store).when(metadataRepo).getStoreOrThrow(anyString());
+    LeaderFollowerIngestionProgressNotifier notifier = new LeaderFollowerIngestionProgressNotifier();
+    LeaderFollowerPartitionStateModel model = buildModel(notifier);
+
+    // Start the OFFLINE -> STANDBY transition in a background thread which should be blocked in
+    // waitConsumptionCompleted, because this is the current version,
+    Message message = mock(Message.class);
+    when(message.getResourceName()).thenReturn(resourceName);
+    NotificationContext context = mock(NotificationContext.class);
+    CompletableFuture<Void> transitionFuture =
+        CompletableFuture.runAsync(() -> model.onBecomeStandbyFromOffline(message, context));
+
+    // Wait until the latch is created, confirming the transition is blocking.
+    TestUtils.waitForNonDeterministicAssertion(
+        5,
+        TimeUnit.SECONDS,
+        () -> assertNotNull(
+            notifier.getIngestionCompleteFlag(resourceName, partition),
+            "Latch must be created for a current-version OFFLINE->STANDBY transition"));
+    assertFalse(transitionFuture.isDone(), "Transition must be blocked waiting for the latch");
+
+    // Simulate CURRENT -> BACKUP version role flip: a newer version has been promoted.
+    when(store.getCurrentVersion()).thenReturn(storeVersion + 1);
+
+    // Simulate StoreIngestionTask.stopTrackingCurrentVersionIngestion() being called
+    notifier.stopped(resourceName, partition, null);
+    TestUtils.waitForNonDeterministicCompletion(5, TimeUnit.SECONDS, transitionFuture::isDone);
+    // Ensure the transition completed successfully (will throw if it completed exceptionally)
+    transitionFuture.join();
+    assertNull(
+        notifier.getIngestionCompleteFlag(resourceName, partition),
+        "Latch should be released/removed after current version is demoted");
+  }
+
   /**
    * Integration test that verifies RocksDB deletion rate limiting is honored during
    * OFFLINE->DROPPED state transition (backup version deletion).

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

@@ -89,6 +89,7 @@
 import com.linkedin.venice.serialization.avro.AvroProtocolDefinition;
 import com.linkedin.venice.serialization.avro.ChunkedValueManifestSerializer;
 import com.linkedin.venice.serialization.avro.InternalAvroSpecificSerializer;
+import com.linkedin.venice.server.VersionRole;
 import com.linkedin.venice.stats.dimensions.VeniceRecordType;
 import com.linkedin.venice.storage.protocol.ChunkedValueManifest;
 import com.linkedin.venice.utils.ByteUtils;
@@ -1411,6 +1412,40 @@ public void testIngestionTimeoutHandling() throws Exception {
     verify(mockHostLevelStats, times(1)).recordIngestionFailure();
   }
 
+  @Test
+  public void testStopTrackingCurrentVersionIngestionOnDemotion() throws Exception {
+    LeaderFollowerStoreIngestionTask storeIngestionTask = mock(LeaderFollowerStoreIngestionTask.class);
+    VeniceServerConfig serverConfig = mock(VeniceServerConfig.class);
+    doReturn(true).when(serverConfig).isResubscriptionTriggeredByVersionIngestionContextChangeEnabled();
+    doReturn(0).when(serverConfig).getResubscriptionCheckIntervalInSeconds();
+    setField(storeIngestionTask, "serverConfig", serverConfig);
+    setField(storeIngestionTask, "versionRole", VersionRole.CURRENT);
+    setVersion(storeIngestionTask, 1);
+    doReturn(true).when(storeIngestionTask).isHybridMode();
+    doCallRealMethod().when(storeIngestionTask).refreshIngestionContextIfChanged(any(Store.class));
+
+    Store store = mock(Store.class);
+    doReturn(5).when(store).getCurrentVersion();
+
+    PartitionConsumptionState pcs = mock(PartitionConsumptionState.class);
+    doReturn(true).when(pcs).isLatchCreated();
+    doReturn(false).when(pcs).isLatchReleased();
+    VeniceConcurrentHashMap<Integer, PartitionConsumptionState> pcsMap = new VeniceConcurrentHashMap<>();
+    pcsMap.put(1, pcs);
+    setField(storeIngestionTask, "partitionConsumptionStateMap", pcsMap);
+
+    IngestionNotificationDispatcher mockDispatcher = mock(IngestionNotificationDispatcher.class);
+    setField(storeIngestionTask, "ingestionNotificationDispatcher", mockDispatcher);
+
+    // First call: versionRole transitions CURRENT -> BACKUP, reportStopped should be called once
+    storeIngestionTask.refreshIngestionContextIfChanged(store);
+    verify(mockDispatcher, times(1)).reportStopped(eq(pcs));
+
+    // Second call: versionRole is already BACKUP, no further latch release
+    storeIngestionTask.refreshIngestionContextIfChanged(store);
+    verify(mockDispatcher, times(1)).reportStopped(eq(pcs));
+  }
+
   private static void addStandbyPcs(Map<Integer, PartitionConsumptionState> pcsMap, int partition, long ageMs) {
     PartitionConsumptionState pcs = mock(PartitionConsumptionState.class);
     pcsMap.put(partition, pcs);